iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] Remove redundant repeat and copy for dsv4, 1.8% E2E TPOT improvement. (#48137)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Wentao Ye
2026-07-15 00:48:10 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 0bd6b85a1f
commit 442c421e79
3 changed files with 316 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/kernels/mhc/tilelang.py
+107
View File
@@ -300,6 +300,113 @@ def _mhc_pre_tilelang_fake(
return post_mix, comb_mix, layer_input
def mhc_pre_broadcast_tilelang(
residual: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
n_splits: int = 1,
norm_weight: torch.Tensor | None = None,
norm_eps: float = 1e-6,
fn_broadcast: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""First-layer mHC pre for a residual broadcast from ``(T, H)``."""
from vllm.model_executor.kernels.mhc.tilelang_kernels import (
compute_num_split,
mhc_pre_big_fuse_broadcast_with_norm_tilelang,
)
from vllm.utils.math_utils import cdiv
assert norm_weight is not None, "broadcast mHC pre currently requires fused RMSNorm"
assert residual.dtype == torch.bfloat16
assert residual.dim() == 2
assert fn.dtype == torch.float32
assert hc_scale.dtype == torch.float32
assert hc_base.dtype == torch.float32
hidden_size = residual.shape[-1]
hc_mult = fn.shape[1] // hidden_size
hc_mult2 = hc_mult * hc_mult
hc_mult3 = hc_mult * 2 + hc_mult2
assert fn.shape == (hc_mult3, hc_mult * hidden_size)
assert hc_scale.shape == (3,)
assert hc_base.shape == (hc_mult3,)
assert fn_broadcast is not None
assert fn_broadcast.dtype == torch.float32
assert fn_broadcast.shape == (hc_mult3, hidden_size)
if norm_weight.dtype != torch.bfloat16:
norm_weight = norm_weight.to(torch.bfloat16)
if not norm_weight.is_contiguous():
norm_weight = norm_weight.contiguous()
residual_flat = residual
num_tokens = residual.shape[0]
n_splits = compute_num_split(64, hidden_size, cdiv(num_tokens, 64))
residual_out = torch.empty(
num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=residual.device
)
post_mix = torch.empty(
num_tokens, hc_mult, dtype=torch.float32, device=residual.device
)
comb_mix = torch.empty(
num_tokens, hc_mult2, dtype=torch.float32, device=residual.device
)
layer_input = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
)
gemm_out_mul = torch.empty(
n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
)
gemm_out_sqrsum = torch.empty(
n_splits, num_tokens, dtype=torch.float32, device=residual.device
)
from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
tf32_hc_prenorm_gemm(
residual_flat,
fn_broadcast,
gemm_out_mul,
gemm_out_sqrsum,
n_splits,
)
mhc_pre_big_fuse_broadcast_with_norm_tilelang(
gemm_out_mul,
gemm_out_sqrsum,
hc_scale,
hc_base,
residual_flat,
residual_out,
post_mix,
comb_mix,
layer_input,
norm_weight,
hidden_size,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
norm_eps,
n_splits,
hc_mult,
)
return (
residual_out,
post_mix.unsqueeze(-1),
comb_mix.view(num_tokens, hc_mult, hc_mult),
layer_input,
)
def mhc_post_tilelang(
x: torch.Tensor,
residual: torch.Tensor,
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+164
View File
@@ -353,6 +353,170 @@ def mhc_pre_big_fuse_with_norm_tilelang(
T.pdl_trigger()
@tilelang.jit(
pass_configs=pass_configs,
)
def mhc_pre_big_fuse_broadcast_with_norm_tilelang(
gemm_out_mul,
gemm_out_sqrsum,
hc_scale,
hc_base,
residual,
residual_out,
post_mix,
comb_mix,
layer_input,
norm_weight,
hidden_size: int,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
norm_eps: float,
n_splits: int = 1,
hc_mult: int = 4,
gemm_last_dim: int = -1,
):
num_tokens = T.dynamic("num_tokens")
hc_mult3 = hc_mult * (2 + hc_mult)
if gemm_last_dim < 0:
gemm_last_dim = hc_mult3
hidden_block = math.gcd(1024, hidden_size)
gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type]
gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type]
hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type]
hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
residual: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
residual_out: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type]
layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
with T.Kernel(num_tokens, threads=96) as i:
rms = T.alloc_fragment(1, T.float32)
mixes = T.alloc_fragment(hc_mult3, T.float32)
T.clear(mixes)
rms[0] = 0
if ENABLE_PDL:
T.pdl_sync()
for i_split in T.serial(n_splits):
rms[0] += gemm_out_sqrsum[i_split, i]
# rms scale
rms[0] = T.rsqrt(rms[0] / hidden_size + rms_eps)
for j in T.Parallel(hc_mult3):
mixes[j] = 0
for i_split in T.serial(n_splits):
mixes[j] += gemm_out_mul[i_split, i, j]
mixes[j] *= rms[0]
mixes_shared = T.alloc_shared(hc_mult3, T.float32)
T.copy(mixes, mixes_shared)
if T.get_thread_binding() < 32:
cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
for j in T.Parallel(hc_mult):
post_mix[i, j] = (
T.sigmoid(
mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
)
* hc_post_mult_value
)
for j, k in T.Parallel(hc_mult, hc_mult):
cm[j, k] = (
mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
+ hc_base[j * hc_mult + k + hc_mult * 2]
)
row_sum = T.alloc_fragment(hc_mult, T.float32)
col_sum = T.alloc_fragment(hc_mult, T.float32)
row_max = T.alloc_fragment(hc_mult, T.float32)
T.reduce_max(cm, row_max, dim=1)
for j, k in T.Parallel(hc_mult, hc_mult):
cm[j, k] = T.exp(cm[j, k] - row_max[j])
T.reduce_sum(cm, row_sum, dim=1)
for j, k in T.Parallel(hc_mult, hc_mult):
cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
T.reduce_sum(cm, col_sum, dim=0)
for j, k in T.Parallel(hc_mult, hc_mult):
cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
for _ in T.serial(sinkhorn_repeat - 1):
T.reduce_sum(cm, row_sum, dim=1)
for j, k in T.Parallel(hc_mult, hc_mult):
cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
T.reduce_sum(cm, col_sum, dim=0)
for j, k in T.Parallel(hc_mult, hc_mult):
cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
for j, k in T.Parallel(hc_mult, hc_mult):
comb_mix[i, j * hc_mult + k] = cm[j, k]
else:
pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
for j in T.Parallel(hc_mult):
pre_mix_shared[j] = (
T.sigmoid(
mixes_shared[j] * hc_scale[0] + hc_base[j],
)
+ hc_pre_eps
)
output_shared = T.alloc_shared(hidden_size, T.bfloat16)
sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32)
T.clear(sumsq_per_pos)
for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
xs = T.alloc_shared(hidden_block, T.bfloat16)
xl = T.alloc_fragment(hidden_block, T.float32)
T.copy(residual[i, i0_h * hidden_block], xs)
T.copy(xs, xl)
for i_hc, i1_h in T.Parallel(hc_mult, hidden_block):
residual_out[i, i_hc, i0_h * hidden_block + i1_h] = xs[i1_h]
ol = T.alloc_fragment(hidden_block, T.float32)
T.clear(ol)
for i_hc in T.serial(hc_mult):
pre = pre_mix_shared[i_hc]
for i1_h in T.Parallel(hidden_block):
ol[i1_h] += pre * xl[i1_h]
for i1_h in T.Parallel(hidden_block):
sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])
sumsq = T.alloc_fragment(1, T.float32)
T.reduce_sum(sumsq_per_pos, sumsq, dim=0)
rsqrt_norm = T.alloc_fragment(1, T.float32)
rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps)
for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
w_shared = T.alloc_shared(hidden_block, T.bfloat16)
w_local = T.alloc_fragment(hidden_block, T.float32)
T.copy(norm_weight[i0_h * hidden_block], w_shared)
T.copy(w_shared, w_local)
ol = T.alloc_fragment(hidden_block, T.float32)
for i1_h in T.Parallel(hidden_block):
ol[i1_h] = (
output_shared[i0_h * hidden_block + i1_h]
* rsqrt_norm[0]
* w_local[i1_h]
)
T.copy(ol, layer_input[i, i0_h * hidden_block])
if ENABLE_PDL:
T.pdl_trigger()
@tilelang.jit(
pass_configs=pass_configs,
)
vllm/models/deepseek_v4/nvidia/model.py
+45 -15
View File
@@ -22,6 +22,7 @@ from vllm.model_executor.kernels.mhc.tilelang import (
hc_head_fused_kernel_tilelang,
mhc_fused_post_pre_tilelang,
mhc_post_tilelang,
mhc_pre_broadcast_tilelang,
mhc_pre_tilelang,
)
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
@@ -827,6 +828,7 @@ class DeepseekV4DecoderLayer(nn.Module):
),
requires_grad=False,
)
self.hc_attn_fn_broadcast: torch.Tensor | None = None
self.hc_ffn_fn = nn.Parameter(
torch.empty(
(mix_hc, hc_dim),
@@ -876,20 +878,37 @@ class DeepseekV4DecoderLayer(nn.Module):
attn_norm_eps = self.attn_norm.variance_epsilon
if residual is None:
# Run standalone mhc_pre on first layer
residual = x
post_mix, res_mix, x = mhc_pre_tilelang(
x,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
self.rms_norm_eps,
self.hc_eps,
self.hc_eps,
self.hc_post_alpha,
self.hc_sinkhorn_iters,
norm_weight=attn_norm_weight,
norm_eps=attn_norm_eps,
)
if x.dim() == 2:
assert self.hc_attn_fn_broadcast is not None
residual, post_mix, res_mix, x = mhc_pre_broadcast_tilelang(
x,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
self.rms_norm_eps,
self.hc_eps,
self.hc_eps,
self.hc_post_alpha,
self.hc_sinkhorn_iters,
norm_weight=attn_norm_weight,
norm_eps=attn_norm_eps,
fn_broadcast=self.hc_attn_fn_broadcast,
)
else:
residual = x
post_mix, res_mix, x = mhc_pre_tilelang(
x,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
self.rms_norm_eps,
self.hc_eps,
self.hc_eps,
self.hc_post_alpha,
self.hc_sinkhorn_iters,
norm_weight=attn_norm_weight,
norm_eps=attn_norm_eps,
)
else:
residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang(
x,
@@ -1070,7 +1089,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
hidden_states = inputs_embeds
else:
hidden_states = self.embed_input_ids(input_ids)
hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
else:
assert intermediate_tensors is not None
hidden_states = intermediate_tensors["hidden_states"]
@@ -1278,6 +1296,17 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
for layer in islice(self.layers, self.start_layer, self.end_layer):
layer.ffn.finalize_mega_moe_weights()
def finalize_mhc_broadcast_weights(self) -> None:
if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
return
layer = self.layers[self.start_layer]
if isinstance(layer, DeepseekV4DecoderLayer):
layer.hc_attn_fn_broadcast = (
layer.hc_attn_fn.detach()
.view(-1, layer.hc_mult, layer.hidden_size)
.sum(dim=1)
)
def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
if expert_dtype == "fp4":
@@ -1440,6 +1469,7 @@ class DeepseekV4ForCausalLM(
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self.model.finalize_mega_moe_weights()
self.model.finalize_mhc_broadcast_weights()
return loaded_params
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: