[Perf] Remove redundant repeat and copy for dsv4, 1.8% E2E TPOT improvement. (#48137)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
Jul 14, 2026 at 20:48 UTC
442c421e794317c37ab5e15feff1b05a035a5a18
3 files changed
+316
-15
vllm/model_executor/kernels/mhc/tilelang.py
+107
@@ -300,6 +300,113 @@ def _mhc_pre_tilelang_fake(
300
return post_mix, comb_mix, layer_input
301
302
303
+def mhc_pre_broadcast_tilelang(
304
+ residual: torch.Tensor,
305
+ fn: torch.Tensor,
306
+ hc_scale: torch.Tensor,
307
+ hc_base: torch.Tensor,
308
+ rms_eps: float,
309
+ hc_pre_eps: float,
310
+ hc_sinkhorn_eps: float,
311
+ hc_post_mult_value: float,
312
+ sinkhorn_repeat: int,
313
+ n_splits: int = 1,
314
+ norm_weight: torch.Tensor | None = None,
315
+ norm_eps: float = 1e-6,
316
+ fn_broadcast: torch.Tensor | None = None,
317
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
318
+ """First-layer mHC pre for a residual broadcast from ``(T, H)``."""
319
+ from vllm.model_executor.kernels.mhc.tilelang_kernels import (
320
+ compute_num_split,
321
+ mhc_pre_big_fuse_broadcast_with_norm_tilelang,
322
+ )
323
+ from vllm.utils.math_utils import cdiv
324
+
325
+ assert norm_weight is not None, "broadcast mHC pre currently requires fused RMSNorm"
326
+ assert residual.dtype == torch.bfloat16
327
+ assert residual.dim() == 2
328
+ assert fn.dtype == torch.float32
329
+ assert hc_scale.dtype == torch.float32
330
+ assert hc_base.dtype == torch.float32
331
+
332
+ hidden_size = residual.shape[-1]
333
+ hc_mult = fn.shape[1] // hidden_size
334
+ hc_mult2 = hc_mult * hc_mult
335
+ hc_mult3 = hc_mult * 2 + hc_mult2
336
+ assert fn.shape == (hc_mult3, hc_mult * hidden_size)
337
+ assert hc_scale.shape == (3,)
338
+ assert hc_base.shape == (hc_mult3,)
339
+ assert fn_broadcast is not None
340
+ assert fn_broadcast.dtype == torch.float32
341
+ assert fn_broadcast.shape == (hc_mult3, hidden_size)
342
+
343
+ if norm_weight.dtype != torch.bfloat16:
344
+ norm_weight = norm_weight.to(torch.bfloat16)
345
+ if not norm_weight.is_contiguous():
346
+ norm_weight = norm_weight.contiguous()
347
+
348
+ residual_flat = residual
349
+ num_tokens = residual.shape[0]
350
+
351
+ n_splits = compute_num_split(64, hidden_size, cdiv(num_tokens, 64))
352
+
353
+ residual_out = torch.empty(
354
+ num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=residual.device
355
+ )
356
+ post_mix = torch.empty(
357
+ num_tokens, hc_mult, dtype=torch.float32, device=residual.device
358
+ )
359
+ comb_mix = torch.empty(
360
+ num_tokens, hc_mult2, dtype=torch.float32, device=residual.device
361
+ )
362
+ layer_input = torch.empty(
363
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
364
+ )
365
+ gemm_out_mul = torch.empty(
366
+ n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
367
+ )
368
+ gemm_out_sqrsum = torch.empty(
369
+ n_splits, num_tokens, dtype=torch.float32, device=residual.device
370
+ )
371
+
372
+ from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
373
+
374
+ tf32_hc_prenorm_gemm(
375
+ residual_flat,
376
+ fn_broadcast,
377
+ gemm_out_mul,
378
+ gemm_out_sqrsum,
379
+ n_splits,
380
+ )
381
+ mhc_pre_big_fuse_broadcast_with_norm_tilelang(
382
+ gemm_out_mul,
383
+ gemm_out_sqrsum,
384
+ hc_scale,
385
+ hc_base,
386
+ residual_flat,
387
+ residual_out,
388
+ post_mix,
389
+ comb_mix,
390
+ layer_input,
391
+ norm_weight,
392
+ hidden_size,
393
+ rms_eps,
394
+ hc_pre_eps,
395
+ hc_sinkhorn_eps,
396
+ hc_post_mult_value,
397
+ sinkhorn_repeat,
398
+ norm_eps,
399
+ n_splits,
400
+ hc_mult,
401
+ )
402
+ return (
403
+ residual_out,
404
+ post_mix.unsqueeze(-1),
405
+ comb_mix.view(num_tokens, hc_mult, hc_mult),
406
+ layer_input,
407
+ )
408
+
409
+
410
def mhc_post_tilelang(
411
x: torch.Tensor,
412
residual: torch.Tensor,
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+164
@@ -353,6 +353,170 @@ def mhc_pre_big_fuse_with_norm_tilelang(
353
T.pdl_trigger()
354
355
356
+@tilelang.jit(
357
+ pass_configs=pass_configs,
358
+)
359
+def mhc_pre_big_fuse_broadcast_with_norm_tilelang(
360
+ gemm_out_mul,
361
+ gemm_out_sqrsum,
362
+ hc_scale,
363
+ hc_base,
364
+ residual,
365
+ residual_out,
366
+ post_mix,
367
+ comb_mix,
368
+ layer_input,
369
+ norm_weight,
370
+ hidden_size: int,
371
+ rms_eps: float,
372
+ hc_pre_eps: float,
373
+ hc_sinkhorn_eps: float,
374
+ hc_post_mult_value: float,
375
+ sinkhorn_repeat: int,
376
+ norm_eps: float,
377
+ n_splits: int = 1,
378
+ hc_mult: int = 4,
379
+ gemm_last_dim: int = -1,
380
+):
381
+ num_tokens = T.dynamic("num_tokens")
382
+ hc_mult3 = hc_mult * (2 + hc_mult)
383
+ if gemm_last_dim < 0:
384
+ gemm_last_dim = hc_mult3
385
+ hidden_block = math.gcd(1024, hidden_size)
386
+
387
+ gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type]
388
+ gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type]
389
+ hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type]
390
+ hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
391
+ residual: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
392
+ residual_out: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
393
+ post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
394
+ comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type]
395
+ layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
396
+ norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
397
+
398
+ with T.Kernel(num_tokens, threads=96) as i:
399
+ rms = T.alloc_fragment(1, T.float32)
400
+ mixes = T.alloc_fragment(hc_mult3, T.float32)
401
+ T.clear(mixes)
402
+ rms[0] = 0
403
+
404
+ if ENABLE_PDL:
405
+ T.pdl_sync()
406
+
407
+ for i_split in T.serial(n_splits):
408
+ rms[0] += gemm_out_sqrsum[i_split, i]
409
+ # rms scale
410
+ rms[0] = T.rsqrt(rms[0] / hidden_size + rms_eps)
411
+ for j in T.Parallel(hc_mult3):
412
+ mixes[j] = 0
413
+ for i_split in T.serial(n_splits):
414
+ mixes[j] += gemm_out_mul[i_split, i, j]
415
+ mixes[j] *= rms[0]
416
+ mixes_shared = T.alloc_shared(hc_mult3, T.float32)
417
+ T.copy(mixes, mixes_shared)
418
+
419
+ if T.get_thread_binding() < 32:
420
+ cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
421
+ for j in T.Parallel(hc_mult):
422
+ post_mix[i, j] = (
423
+ T.sigmoid(
424
+ mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
425
+ )
426
+ * hc_post_mult_value
427
+ )
428
+ for j, k in T.Parallel(hc_mult, hc_mult):
429
+ cm[j, k] = (
430
+ mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
431
+ + hc_base[j * hc_mult + k + hc_mult * 2]
432
+ )
433
+
434
+ row_sum = T.alloc_fragment(hc_mult, T.float32)
435
+ col_sum = T.alloc_fragment(hc_mult, T.float32)
436
+
437
+ row_max = T.alloc_fragment(hc_mult, T.float32)
438
+ T.reduce_max(cm, row_max, dim=1)
439
+ for j, k in T.Parallel(hc_mult, hc_mult):
440
+ cm[j, k] = T.exp(cm[j, k] - row_max[j])
441
+ T.reduce_sum(cm, row_sum, dim=1)
442
+ for j, k in T.Parallel(hc_mult, hc_mult):
443
+ cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
444
+
445
+ T.reduce_sum(cm, col_sum, dim=0)
446
+ for j, k in T.Parallel(hc_mult, hc_mult):
447
+ cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
448
+
449
+ for _ in T.serial(sinkhorn_repeat - 1):
450
+ T.reduce_sum(cm, row_sum, dim=1)
451
+ for j, k in T.Parallel(hc_mult, hc_mult):
452
+ cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
453
+
454
+ T.reduce_sum(cm, col_sum, dim=0)
455
+ for j, k in T.Parallel(hc_mult, hc_mult):
456
+ cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
457
+
458
+ for j, k in T.Parallel(hc_mult, hc_mult):
459
+ comb_mix[i, j * hc_mult + k] = cm[j, k]
460
+ else:
461
+ pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
462
+ for j in T.Parallel(hc_mult):
463
+ pre_mix_shared[j] = (
464
+ T.sigmoid(
465
+ mixes_shared[j] * hc_scale[0] + hc_base[j],
466
+ )
467
+ + hc_pre_eps
468
+ )
469
+
470
+ output_shared = T.alloc_shared(hidden_size, T.bfloat16)
471
+ sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32)
472
+ T.clear(sumsq_per_pos)
473
+
474
+ for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
475
+ xs = T.alloc_shared(hidden_block, T.bfloat16)
476
+ xl = T.alloc_fragment(hidden_block, T.float32)
477
+ T.copy(residual[i, i0_h * hidden_block], xs)
478
+ T.copy(xs, xl)
479
+
480
+ for i_hc, i1_h in T.Parallel(hc_mult, hidden_block):
481
+ residual_out[i, i_hc, i0_h * hidden_block + i1_h] = xs[i1_h]
482
+
483
+ ol = T.alloc_fragment(hidden_block, T.float32)
484
+ T.clear(ol)
485
+
486
+ for i_hc in T.serial(hc_mult):
487
+ pre = pre_mix_shared[i_hc]
488
+ for i1_h in T.Parallel(hidden_block):
489
+ ol[i1_h] += pre * xl[i1_h]
490
+
491
+ for i1_h in T.Parallel(hidden_block):
492
+ sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
493
+ output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])
494
+
495
+ sumsq = T.alloc_fragment(1, T.float32)
496
+ T.reduce_sum(sumsq_per_pos, sumsq, dim=0)
497
+ rsqrt_norm = T.alloc_fragment(1, T.float32)
498
+ rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps)
499
+
500
+ for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
501
+ w_shared = T.alloc_shared(hidden_block, T.bfloat16)
502
+ w_local = T.alloc_fragment(hidden_block, T.float32)
503
+ T.copy(norm_weight[i0_h * hidden_block], w_shared)
504
+ T.copy(w_shared, w_local)
505
+
506
+ ol = T.alloc_fragment(hidden_block, T.float32)
507
+ for i1_h in T.Parallel(hidden_block):
508
+ ol[i1_h] = (
509
+ output_shared[i0_h * hidden_block + i1_h]
510
+ * rsqrt_norm[0]
511
+ * w_local[i1_h]
512
+ )
513
+
514
+ T.copy(ol, layer_input[i, i0_h * hidden_block])
515
+
516
+ if ENABLE_PDL:
517
+ T.pdl_trigger()
518
+
519
+
520
@tilelang.jit(
521
pass_configs=pass_configs,
522
)
vllm/models/deepseek_v4/nvidia/model.py
+45
-15
@@ -22,6 +22,7 @@ from vllm.model_executor.kernels.mhc.tilelang import (
22
hc_head_fused_kernel_tilelang,
23
mhc_fused_post_pre_tilelang,
24
mhc_post_tilelang,
25
+ mhc_pre_broadcast_tilelang,
26
mhc_pre_tilelang,
27
)
28
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
@@ -827,6 +828,7 @@ class DeepseekV4DecoderLayer(nn.Module):
828
),
829
requires_grad=False,
830
)
831
+ self.hc_attn_fn_broadcast: torch.Tensor | None = None
832
self.hc_ffn_fn = nn.Parameter(
833
torch.empty(
834
(mix_hc, hc_dim),
@@ -876,20 +878,37 @@ class DeepseekV4DecoderLayer(nn.Module):
878
attn_norm_eps = self.attn_norm.variance_epsilon
879
if residual is None:
880
# Run standalone mhc_pre on first layer
879
- residual = x
880
- post_mix, res_mix, x = mhc_pre_tilelang(
881
- x,
882
- self.hc_attn_fn,
883
- self.hc_attn_scale,
884
- self.hc_attn_base,
885
- self.rms_norm_eps,
886
- self.hc_eps,
887
- self.hc_eps,
888
- self.hc_post_alpha,
889
- self.hc_sinkhorn_iters,
890
- norm_weight=attn_norm_weight,
891
- norm_eps=attn_norm_eps,
892
- )
881
+ if x.dim() == 2:
882
+ assert self.hc_attn_fn_broadcast is not None
883
+ residual, post_mix, res_mix, x = mhc_pre_broadcast_tilelang(
884
+ x,
885
+ self.hc_attn_fn,
886
+ self.hc_attn_scale,
887
+ self.hc_attn_base,
888
+ self.rms_norm_eps,
889
+ self.hc_eps,
890
+ self.hc_eps,
891
+ self.hc_post_alpha,
892
+ self.hc_sinkhorn_iters,
893
+ norm_weight=attn_norm_weight,
894
+ norm_eps=attn_norm_eps,
895
+ fn_broadcast=self.hc_attn_fn_broadcast,
896
+ )
897
+ else:
898
+ residual = x
899
+ post_mix, res_mix, x = mhc_pre_tilelang(
900
+ x,
901
+ self.hc_attn_fn,
902
+ self.hc_attn_scale,
903
+ self.hc_attn_base,
904
+ self.rms_norm_eps,
905
+ self.hc_eps,
906
+ self.hc_eps,
907
+ self.hc_post_alpha,
908
+ self.hc_sinkhorn_iters,
909
+ norm_weight=attn_norm_weight,
910
+ norm_eps=attn_norm_eps,
911
+ )
912
else:
913
residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang(
914
x,
@@ -1070,7 +1089,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1089
hidden_states = inputs_embeds
1090
else:
1091
hidden_states = self.embed_input_ids(input_ids)
1073
- hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
1092
else:
1093
assert intermediate_tensors is not None
1094
hidden_states = intermediate_tensors["hidden_states"]
@@ -1278,6 +1296,17 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1296
for layer in islice(self.layers, self.start_layer, self.end_layer):
1297
layer.ffn.finalize_mega_moe_weights()
1298
1299
+ def finalize_mhc_broadcast_weights(self) -> None:
1300
+ if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
1301
+ return
1302
+ layer = self.layers[self.start_layer]
1303
+ if isinstance(layer, DeepseekV4DecoderLayer):
1304
+ layer.hc_attn_fn_broadcast = (
1305
+ layer.hc_attn_fn.detach()
1306
+ .view(-1, layer.hc_mult, layer.hidden_size)
1307
+ .sum(dim=1)
1308
+ )
1309
+
1310
1311
def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1312
if expert_dtype == "fp4":
@@ -1440,6 +1469,7 @@ class DeepseekV4ForCausalLM(
1469
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
1470
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1471
self.model.finalize_mega_moe_weights()
1472
+ self.model.finalize_mhc_broadcast_weights()
1473
return loaded_params
1474
1475
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: