diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index e0007141d5..c3c773d07f 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -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, diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 9fa13041b3..50abac3e01 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -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, ) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 5f4a638c24..ddc2fe0f4b 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -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]]: