@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0a65d46628
Fork 0

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 3, 16:30 UTC

[DSV4] Fuse norm and router for low latency scenario (#41263)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: jeejeelee <jeejeelee@verda-b300-05.datacrunch.io> Co-authored-by: jeejeelee <jeejeelee@verda-b300-05.datacrunch.io> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Jee Jee Li committed May 14, 2026 at 20:11 UTC 0a65d46628fdc7d8e83c1111b46a078219d7d983
11 files changed +815 -43
CMakeLists.txt
+10
@@ -1222,6 +1222,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1222 CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
1223 list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
1224 message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
1225 +
1226 + # DeepSeek V4 fused RMSNorm + router GEMV - same arch gating as DSV3.
1227 + set(DSV4_NORM_ROUTER_GEMM_SRC
1228 + "csrc/moe/dsv4_norm_router_gemm_entry.cu"
1229 + "csrc/moe/dsv4_norm_router_gemm_kernel.cu")
1230 + set_gencode_flags_for_srcs(
1231 + SRCS "${DSV4_NORM_ROUTER_GEMM_SRC}"
1232 + CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
1233 + list(APPEND VLLM_MOE_EXT_SRC "${DSV4_NORM_ROUTER_GEMM_SRC}")
1234 + message(STATUS "Building DSV4 norm+router GEMV kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
1235 else()
1236 message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
1237 " (requires SM90+ and CUDA >= 12.0)")
benchmarks/kernels/benchmark_norm_router_gemm.py new
+183
@@ -0,0 +1,183 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark and correctness check for ``ops.dsv4_norm_router_gemm``.
4 +
5 +Two implementations are compared:
6 +
7 + 1. ``unfused`` — ``vllm_ops.rms_norm`` then ``ops.dsv3_router_gemm``,
8 + i.e. the current vLLM hot path (two kernel launches).
9 + 2. ``fused`` — ``ops.dsv4_norm_router_gemm``, the new single-kernel
10 + fused path.
11 +
12 +Both produce ``(normed_x: bf16, router_logits: fp32)``. The correctness
13 +check verifies that ``fused`` and ``unfused`` agree to within ~1 bf16
14 +ULP — that is the precision floor for this op.
15 +"""
16 +
17 +import argparse
18 +
19 +import torch
20 +
21 +from vllm import _custom_ops as vllm_ops
22 +from vllm.triton_utils import triton
23 +
24 +# The fused dsv4_norm_router_gemm kernel is templated only for DSV4-Pro
25 +# (hidden_size=7168, num_experts=384). Other shapes fall back to the
26 +# unfused path on the Python side (NormGatedLinear), so benchmark only
27 +# the configuration that the fused kernel actually targets.
28 +HIDDEN_SIZE = 7168
29 +NUM_EXPERTS_CHOICES = (384,)
30 +RMS_EPS = 1e-6
31 +
32 +
33 +def unfused_norm_router_gemm(
34 + x: torch.Tensor,
35 + norm_weight: torch.Tensor,
36 + gate_weight: torch.Tensor,
37 + eps: float,
38 +) -> tuple[torch.Tensor, torch.Tensor]:
39 + # Call ``_C::rms_norm`` directly (mirroring ``_dsv4_pro_norm_gate``'s
40 + # fallback path) so the benchmarked baseline doesn't inherit any
41 + # Python wrapper overhead or risk falling through to the native
42 + # eager-primitive ``RMSNorm.forward_native`` path.
43 + normed = torch.empty_like(x)
44 + torch.ops._C.rms_norm(normed, x, norm_weight, eps)
45 + logits = vllm_ops.dsv3_router_gemm(normed, gate_weight, torch.float32)
46 + return normed, logits
47 +
48 +
49 +def fused_norm_router_gemm(
50 + x: torch.Tensor,
51 + norm_weight: torch.Tensor,
52 + gate_weight: torch.Tensor,
53 + eps: float,
54 +) -> tuple[torch.Tensor, torch.Tensor]:
55 + return vllm_ops.dsv4_norm_router_gemm(x, norm_weight, gate_weight, eps)
56 +
57 +
58 +def _make_inputs(num_tokens: int, num_experts: int, hidden_size: int, seed: int = 0):
59 + torch.manual_seed(seed)
60 + device = "cuda"
61 + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
62 + norm_w = torch.randn(hidden_size, dtype=torch.bfloat16, device=device)
63 + gate_w = torch.randn(num_experts, hidden_size, dtype=torch.bfloat16, device=device)
64 + # Down-scale gate_w so the GEMV output stays in a representable range.
65 + gate_w = gate_w / float(hidden_size) ** 0.5
66 + norm_w = (norm_w * 0.1) + 1.0
67 + return x, norm_w, gate_w
68 +
69 +
70 +def calculate_diff(
71 + num_tokens: int,
72 + num_experts: int,
73 + hidden_size: int = HIDDEN_SIZE,
74 + normed_atol: float = 2e-3,
75 + logits_atol: float = 1e-2,
76 + rtol: float = 1e-2,
77 +) -> None:
78 + x, norm_w, gate_w = _make_inputs(num_tokens, num_experts, hidden_size)
79 +
80 + normed_unfused, logits_unfused = unfused_norm_router_gemm(
81 + x.clone(), norm_w, gate_w, RMS_EPS
82 + )
83 + normed_fused, logits_fused = fused_norm_router_gemm(
84 + x.clone(), norm_w, gate_w, RMS_EPS
85 + )
86 +
87 + def _max_abs(a, b):
88 + return (a.float() - b.float()).abs().max().item()
89 +
90 + print(f"\n=== M={num_tokens} E={num_experts} H={hidden_size} ===")
91 + print(f"normed_x |fused - unfused| = {_max_abs(normed_fused, normed_unfused):.3e}")
92 + print(f"logits |fused - unfused| = {_max_abs(logits_fused, logits_unfused):.3e}")
93 +
94 + ok_normed = torch.allclose(
95 + normed_fused.float(),
96 + normed_unfused.float(),
97 + atol=normed_atol,
98 + rtol=rtol,
99 + )
100 + ok_logits = torch.allclose(
101 + logits_fused.float(),
102 + logits_unfused.float(),
103 + atol=logits_atol,
104 + rtol=rtol,
105 + )
106 + if ok_normed and ok_logits:
107 + print(
108 + f"OK fused vs unfused within "
109 + f"normed_atol={normed_atol:.0e} logits_atol={logits_atol:.0e} "
110 + f"rtol={rtol:.0e}"
111 + )
112 + else:
113 + print(
114 + f"FAIL normed_ok={ok_normed} logits_ok={ok_logits}; "
115 + f"see max-abs values above"
116 + )
117 +
118 +
119 +def get_benchmark():
120 + # Only num_tokens varies (DSV4-Pro hard-codes E=384); single-axis
121 + # sweep yields a clean line plot with M on the x-axis.
122 + num_experts = NUM_EXPERTS_CHOICES[0]
123 +
124 + @triton.testing.perf_report(
125 + triton.testing.Benchmark(
126 + x_names=["num_tokens"],
127 + x_vals=list(range(1, 17)),
128 + line_arg="provider",
129 + line_vals=["unfused", "fused"],
130 + line_names=["unfused (rms+dsv3)", "fused (dsv4)"],
131 + styles=[("green", "-"), ("red", "-")],
132 + ylabel="us",
133 + plot_name=f"norm-router-gemm-E{num_experts}-H{HIDDEN_SIZE}",
134 + args={},
135 + )
136 + )
137 + def benchmark(num_tokens, provider):
138 + x, norm_w, gate_w = _make_inputs(num_tokens, num_experts, HIDDEN_SIZE)
139 +
140 + quantiles = [0.5, 0.2, 0.8]
141 + if provider == "unfused":
142 + fn = lambda: unfused_norm_router_gemm( # noqa: E731
143 + x, norm_w, gate_w, RMS_EPS
144 + )
145 + else:
146 + fn = lambda: fused_norm_router_gemm( # noqa: E731
147 + x, norm_w, gate_w, RMS_EPS
148 + )
149 +
150 + ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
151 + return 1000 * ms, 1000 * max_ms, 1000 * min_ms
152 +
153 + return benchmark
154 +
155 +
156 +def main() -> None:
157 + parser = argparse.ArgumentParser()
158 + parser.add_argument(
159 + "--save-path",
160 + type=str,
161 + default="./configs/norm_router_gemm/",
162 + )
163 + parser.add_argument(
164 + "--skip-bench",
165 + action="store_true",
166 + help="Run only the correctness check, not the perf sweep.",
167 + )
168 + args = parser.parse_args()
169 +
170 + # Correctness sweep over the full fast-path range M=1..16.
171 + for m in range(1, 17):
172 + for e in NUM_EXPERTS_CHOICES:
173 + calculate_diff(num_tokens=m, num_experts=e, hidden_size=HIDDEN_SIZE)
174 +
175 + if args.skip_bench:
176 + return
177 +
178 + benchmark = get_benchmark()
179 + benchmark.run(print_data=True, save_path=args.save_path)
180 +
181 +
182 +if __name__ == "__main__":
183 + main()
csrc/moe/dsv4_norm_router_gemm.h new
+30
@@ -0,0 +1,30 @@
1 +/*
2 + * Fused RMSNorm + router GEMV for DeepSeek V4.
3 + *
4 + * Computes in a single kernel:
5 + * normed_x[m,k] = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
6 + * router_logits[m,n] = sum_k(normed_x[m,k] * gate_weight[n,k])
7 + *
8 + * The GEMV body mirrors the algorithm in csrc/moe/dsv3_router_gemm_*.cu
9 + * (warp butterfly + smem cross-warp reduction, fp32 accumulation, PDL on
10 + * SM90+). Blocks 0..kNumTokens-1 each materialize one token's normed_x
11 + * row to global memory using the algebraic identity
12 + * logits[m,n] = rsqrt[m] * sum_k(x[m,k] * nw[k] * gw[n,k])
13 + * which lets every block produce its column of logits before normed_x
14 + * exists in gmem.
15 + *
16 + * Logits output is fp32 only — DeepSeek V4 router gate is hard-coded to
17 + * fp32 (vllm/model_executor/models/deepseek_v4.py:749).
18 + */
19 +
20 +#pragma once
21 +
22 +#include <cuda_bf16.h>
23 +#include <cuda_runtime.h>
24 +
25 +#include "dsv3_router_gemm_utils.h"
26 +
27 +template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
28 +void invokeNormRouterGemm(float* logits, __nv_bfloat16* normed_x, T const* x,
29 + T const* norm_weight, T const* gate_weight, float eps,
30 + cudaStream_t stream);
csrc/moe/dsv4_norm_router_gemm_entry.cu new
+130
@@ -0,0 +1,130 @@
1 +/*
2 + * TORCH op entry for the fused RMSNorm + router GEMV kernel
3 + * (DeepSeek V4 Pro). This op is DSV4-Pro-specific: the kernel is
4 + * instantiated only for ``num_experts == 384`` and ``hidden_dim ==
5 + * 7168``. Other configurations (e.g. DSV4-Flash with H=4096) must
6 + * fall back to the unfused ``rms_norm`` + ``dsv3_router_gemm`` path.
7 + */
8 +
9 +#include <ATen/ATen.h>
10 +#include <ATen/cuda/CUDAContext.h>
11 +#include <torch/all.h>
12 +
13 +#include <cuda_bf16.h>
14 +#include <cuda_runtime.h>
15 +
16 +#include "core/registration.h"
17 +#include "dsv4_norm_router_gemm.h"
18 +
19 +namespace {
20 +
21 +// DSV4-Pro hard-coded shape constants. Renamed from the earlier
22 +// ``kKimiK2NumExperts`` to avoid the misleading impression that this
23 +// kernel targets Kimi K2 — 384 happens to match Kimi K2's gate but the
24 +// intent here is DSV4-Pro.
25 +constexpr int kDsv4NumExperts = 384;
26 +constexpr int kDsv4HiddenDim = 7168;
27 +
28 +template <int kBegin, int kEnd>
29 +struct LoopUnroller {
30 + static void unroll(int num_tokens, float* logits, __nv_bfloat16* normed_x,
31 + __nv_bfloat16 const* x, __nv_bfloat16 const* norm_weight,
32 + __nv_bfloat16 const* gate_weight, float eps,
33 + cudaStream_t stream) {
34 + if (num_tokens == kBegin) {
35 + invokeNormRouterGemm<__nv_bfloat16, kBegin, kDsv4NumExperts,
36 + kDsv4HiddenDim>(logits, normed_x, x, norm_weight,
37 + gate_weight, eps, stream);
38 + } else {
39 + LoopUnroller<kBegin + 1, kEnd>::unroll(num_tokens, logits, normed_x, x,
40 + norm_weight, gate_weight, eps,
41 + stream);
42 + }
43 + }
44 +};
45 +
46 +template <int kEnd>
47 +struct LoopUnroller<kEnd, kEnd> {
48 + static void unroll(int num_tokens, float* logits, __nv_bfloat16* normed_x,
49 + __nv_bfloat16 const* x, __nv_bfloat16 const* norm_weight,
50 + __nv_bfloat16 const* gate_weight, float eps,
51 + cudaStream_t stream) {
52 + if (num_tokens == kEnd) {
53 + invokeNormRouterGemm<__nv_bfloat16, kEnd, kDsv4NumExperts,
54 + kDsv4HiddenDim>(logits, normed_x, x, norm_weight,
55 + gate_weight, eps, stream);
56 + } else {
57 + throw std::invalid_argument(
58 + "Invalid num_tokens, only supports 1 to 16 for "
59 + "dsv4_norm_router_gemm");
60 + }
61 + }
62 +};
63 +
64 +} // namespace
65 +
66 +void dsv4_norm_router_gemm(at::Tensor& logits, // [num_tokens, E] fp32
67 + at::Tensor& normed_x, // [num_tokens, H] bf16
68 + at::Tensor const& x, // [num_tokens, H] bf16
69 + at::Tensor const& norm_weight, // [H] bf16
70 + at::Tensor const& gate_weight, // [E, H] bf16
71 + double eps) {
72 + TORCH_CHECK(x.dim() == 2 && norm_weight.dim() == 1 && gate_weight.dim() == 2,
73 + "x must be 2D, norm_weight 1D, gate_weight 2D");
74 + TORCH_CHECK(logits.dim() == 2 && normed_x.dim() == 2,
75 + "logits and normed_x must be 2D");
76 +
77 + int const num_tokens = x.size(0);
78 + int const hidden_dim = x.size(1);
79 + int const num_experts = gate_weight.size(0);
80 +
81 + TORCH_CHECK(hidden_dim == kDsv4HiddenDim,
82 + "Expected hidden_dim=", kDsv4HiddenDim,
83 + " (DSV4-Pro), but got hidden_dim=", hidden_dim);
84 + TORCH_CHECK(gate_weight.size(1) == hidden_dim,
85 + "gate_weight.shape[1] must equal x.shape[1]");
86 + TORCH_CHECK(norm_weight.size(0) == hidden_dim,
87 + "norm_weight.shape[0] must equal x.shape[1]");
88 + TORCH_CHECK(num_experts == kDsv4NumExperts,
89 + "Expected num_experts=", kDsv4NumExperts,
90 + " (DSV4-Pro), but got num_experts=", num_experts);
91 + TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
92 + "num_tokens must be in [1, 16] for dsv4_norm_router_gemm");
93 +
94 + TORCH_CHECK(x.dtype() == at::kBFloat16, "x must be bf16");
95 + TORCH_CHECK(norm_weight.dtype() == at::kBFloat16, "norm_weight must be bf16");
96 + TORCH_CHECK(gate_weight.dtype() == at::kBFloat16, "gate_weight must be bf16");
97 + TORCH_CHECK(normed_x.dtype() == at::kBFloat16, "normed_x must be bf16");
98 + TORCH_CHECK(logits.dtype() == at::kFloat,
99 + "logits must be float32 (DSV4 router output is hard-coded fp32)");
100 +
101 + TORCH_CHECK(normed_x.size(0) == num_tokens && normed_x.size(1) == hidden_dim,
102 + "normed_x must be [num_tokens, hidden_dim]");
103 + TORCH_CHECK(logits.size(0) == num_tokens && logits.size(1) == num_experts,
104 + "logits must be [num_tokens, num_experts]");
105 +
106 + TORCH_CHECK(x.is_contiguous() && norm_weight.is_contiguous() &&
107 + gate_weight.is_contiguous() && normed_x.is_contiguous() &&
108 + logits.is_contiguous(),
109 + "all tensors must be contiguous");
110 +
111 + auto const sm = getSMVersion();
112 + TORCH_CHECK(sm >= 90 && sm <= 103,
113 + "dsv4_norm_router_gemm requires SM_90 <= CUDA ARCH <= SM_103");
114 +
115 + cudaStream_t const stream = at::cuda::getCurrentCUDAStream();
116 +
117 + auto* logits_ptr = reinterpret_cast<float*>(logits.mutable_data_ptr());
118 + auto* nx_ptr = reinterpret_cast<__nv_bfloat16*>(normed_x.mutable_data_ptr());
119 + auto* x_ptr = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr());
120 + auto* nw_ptr = reinterpret_cast<__nv_bfloat16 const*>(norm_weight.data_ptr());
121 + auto* gw_ptr = reinterpret_cast<__nv_bfloat16 const*>(gate_weight.data_ptr());
122 + float const eps_f = static_cast<float>(eps);
123 +
124 + LoopUnroller<1, 16>::unroll(num_tokens, logits_ptr, nx_ptr, x_ptr, nw_ptr,
125 + gw_ptr, eps_f, stream);
126 +}
127 +
128 +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
129 + m.impl("dsv4_norm_router_gemm", &dsv4_norm_router_gemm);
130 +}
csrc/moe/dsv4_norm_router_gemm_kernel.cu new
+249
@@ -0,0 +1,249 @@
1 +/*
2 + * Fused RMSNorm + router GEMV for DeepSeek V4 (logits are fp32; bf16
3 + * output is unsupported because DSV4 hard-codes fp32 logits). See
4 + * dsv4_norm_router_gemm.h for the math.
5 + *
6 + * The GEMV body mirrors csrc/moe/dsv3_router_gemm_float_out.cu (warp
7 + * butterfly reduction + smem cross-warp reduction, fp32 accumulation,
8 + * 128-thread block, PDL on SM90+). RMSNorm is folded into the same
9 + * pass via the identity
10 + * logits[m,n] = rsqrt[m] * sum_k(x[m,k] * nw[k] * gw[n,k])
11 + * so x is read exactly once per block during the GEMV phase. Blocks
12 + * 0..kNumTokens-1 each materialize one row of normed_x for downstream
13 + * experts / shared_experts to consume.
14 + */
15 +
16 +#include <ATen/ATen.h>
17 +#include <ATen/cuda/CUDAContext.h>
18 +
19 +#include <cuda_bf16.h>
20 +#include <cuda_runtime.h>
21 +
22 +#include "dsv4_norm_router_gemm.h"
23 +
24 +namespace {
25 +
26 +// Convert 8 bf16 values packed in uint4 into 8 floats. Mirrors the helper
27 +// in dsv3_router_gemm_float_out.cu (kept local so the dsv3 file stays
28 +// untouched).
29 +template <int VPT>
30 +__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
31 + float* dst) {
32 + __nv_bfloat16* bf16_ptr =
33 + reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
34 +#pragma unroll
35 + for (int i = 0; i < VPT; i++) {
36 + dst[i] = __bfloat162float(bf16_ptr[i]);
37 + }
38 +}
39 +
40 +template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
41 + int kHiddenDim>
42 +__global__ __launch_bounds__(128, 1) void norm_router_gemm_kernel(
43 + float* __restrict__ logits, __nv_bfloat16* __restrict__ normed_x,
44 + T const* __restrict__ x, T const* __restrict__ norm_weight,
45 + T const* __restrict__ gate_weight, float eps) {
46 + static_assert(kBlockSize == 128, "kernel assumes blockDim.x == 128");
47 + static_assert(kHiddenDim % (VPT * kBlockSize) == 0,
48 + "kHiddenDim must be a multiple of VPT * kBlockSize");
49 +
50 + int const n_idx = blockIdx.x;
51 + int const tid = threadIdx.x;
52 + constexpr int kWarpSize = 32;
53 + constexpr int kNumWarps = kBlockSize / kWarpSize;
54 + constexpr int k_elems_per_iter = VPT * kBlockSize;
55 + constexpr int k_iterations = kHiddenDim / k_elems_per_iter;
56 +
57 + T const* gw_col = gate_weight + n_idx * kHiddenDim;
58 +
59 + // Per-thread accumulators — fp32 throughout, matching dsv3 / layernorm.
60 + float partial[kNumTokens] = {};
61 + float ss[kNumTokens] = {};
62 +
63 + // Cross-warp reduction scratch.
64 + __shared__ float sm_partial[kNumTokens][kNumWarps];
65 + __shared__ float sm_ss[kNumTokens][kNumWarps];
66 + __shared__ float s_rsqrt[kNumTokens];
67 +
68 + int k_bases[k_iterations];
69 +#pragma unroll
70 + for (int ki = 0; ki < k_iterations; ki++) {
71 + k_bases[ki] = ki * k_elems_per_iter + tid * VPT;
72 + }
73 +
74 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
75 + asm volatile("griddepcontrol.wait;");
76 +#endif
77 +
78 + // ---- Phase 1: single pass over x, accumulate partial GEMV and ss. ----
79 +#pragma unroll
80 + for (int ki = 0; ki < k_iterations; ki++) {
81 + int const k_base = k_bases[ki];
82 +
83 + uint4 nw_vec = *reinterpret_cast<uint4 const*>(norm_weight + k_base);
84 + float nw_f[VPT];
85 + bf16_uint4_to_float8<VPT>(nw_vec, nw_f);
86 +
87 + uint4 b_vec = *reinterpret_cast<uint4 const*>(gw_col + k_base);
88 + float b_f[VPT];
89 + bf16_uint4_to_float8<VPT>(b_vec, b_f);
90 +
91 +#pragma unroll
92 + for (int m = 0; m < kNumTokens; m++) {
93 + uint4 a_vec =
94 + *reinterpret_cast<uint4 const*>(x + m * kHiddenDim + k_base);
95 + float a_f[VPT];
96 + bf16_uint4_to_float8<VPT>(a_vec, a_f);
97 +
98 +#pragma unroll
99 + for (int k = 0; k < VPT; k++) {
100 + float a = a_f[k];
101 + ss[m] += a * a;
102 + partial[m] += a * nw_f[k] * b_f[k];
103 + }
104 + }
105 + }
106 +
107 + // ---- Phase 2: warp butterfly reduction for both ss[] and partial[]. ----
108 + int const warpId = tid / kWarpSize;
109 + int const laneId = tid % kWarpSize;
110 +
111 +#pragma unroll
112 + for (int m = 0; m < kNumTokens; m++) {
113 + float p = partial[m];
114 + float s = ss[m];
115 +
116 + p += __shfl_xor_sync(0xffffffff, p, 16);
117 + s += __shfl_xor_sync(0xffffffff, s, 16);
118 + p += __shfl_xor_sync(0xffffffff, p, 8);
119 + s += __shfl_xor_sync(0xffffffff, s, 8);
120 + p += __shfl_xor_sync(0xffffffff, p, 4);
121 + s += __shfl_xor_sync(0xffffffff, s, 4);
122 + p += __shfl_xor_sync(0xffffffff, p, 2);
123 + s += __shfl_xor_sync(0xffffffff, s, 2);
124 + p += __shfl_xor_sync(0xffffffff, p, 1);
125 + s += __shfl_xor_sync(0xffffffff, s, 1);
126 +
127 + if (laneId == 0) {
128 + sm_partial[m][warpId] = p;
129 + sm_ss[m][warpId] = s;
130 + }
131 + }
132 +
133 + __syncthreads();
134 +
135 + // ---- Phase 3: tid 0 finalises the reduction, writes logits, stashes
136 + // rsqrt[m] in smem for phase 4. ----
137 + if (tid == 0) {
138 +#pragma unroll
139 + for (int m = 0; m < kNumTokens; m++) {
140 + float p_sum = 0.0f;
141 + float s_sum = 0.0f;
142 +#pragma unroll
143 + for (int w = 0; w < kNumWarps; w++) {
144 + p_sum += sm_partial[m][w];
145 + s_sum += sm_ss[m][w];
146 + }
147 + // Order matches layernorm_kernels.cu: rsqrtf(variance / H + eps).
148 + // Use division (not multiply-by-reciprocal) to avoid an extra ULP
149 + // mismatch with the reference RMSNorm.
150 + float rs = rsqrtf(s_sum / static_cast<float>(kHiddenDim) + eps);
151 + s_rsqrt[m] = rs;
152 + logits[m * kNumExperts + n_idx] = p_sum * rs;
153 + }
154 + }
155 +
156 + __syncthreads();
157 +
158 + // ---- Phase 4: spread normed_x writes across blocks 0..kNumTokens-1.
159 + // Each writer block handles exactly one token row,
160 + // avoiding the long tail of block 0 doing all M rows.
161 + // Every block has every token's rsqrt[] in s_rsqrt
162 + // already (computed independently in phase 3), so no
163 + // cross-block synchronization is required. ----
164 + if (n_idx < kNumTokens) {
165 + int const m_writer = n_idx;
166 + float const rs = s_rsqrt[m_writer];
167 + __nv_bfloat16 const* x_row = x + m_writer * kHiddenDim;
168 + __nv_bfloat16* normed_row = normed_x + m_writer * kHiddenDim;
169 +
170 +#pragma unroll
171 + for (int ki = 0; ki < k_iterations; ki++) {
172 + int const k_base = k_bases[ki];
173 +
174 + uint4 nw_vec = *reinterpret_cast<uint4 const*>(norm_weight + k_base);
175 + float nw_f[VPT];
176 + bf16_uint4_to_float8<VPT>(nw_vec, nw_f);
177 +
178 + uint4 a_vec = *reinterpret_cast<uint4 const*>(x_row + k_base);
179 + float a_f[VPT];
180 + bf16_uint4_to_float8<VPT>(a_vec, a_f);
181 +
182 + uint4 normed_vec;
183 + __nv_bfloat16* np = reinterpret_cast<__nv_bfloat16*>(&normed_vec);
184 +#pragma unroll
185 + for (int k = 0; k < VPT; k++) {
186 + np[k] = __float2bfloat16(a_f[k] * rs * nw_f[k]);
187 + }
188 + *reinterpret_cast<uint4*>(normed_row + k_base) = normed_vec;
189 + }
190 + }
191 +
192 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
193 + asm volatile("griddepcontrol.launch_dependents;");
194 +#endif
195 +}
196 +
197 +} // namespace
198 +
199 +template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
200 +void invokeNormRouterGemm(float* logits, __nv_bfloat16* normed_x, T const* x,
201 + T const* norm_weight, T const* gate_weight, float eps,
202 + cudaStream_t stream) {
203 + constexpr int VPT = 16 / sizeof(T);
204 + constexpr int kBlockSize = 128;
205 +
206 + cudaLaunchConfig_t config;
207 + config.gridDim = kNumExperts;
208 + config.blockDim = kBlockSize;
209 + config.dynamicSmemBytes = 0;
210 + config.stream = stream;
211 +
212 + cudaLaunchAttribute attrs[1];
213 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
214 + attrs[0].val.programmaticStreamSerializationAllowed = 1;
215 + config.numAttrs = 1;
216 + config.attrs = attrs;
217 +
218 + cudaLaunchKernelEx(&config,
219 + norm_router_gemm_kernel<T, kBlockSize, VPT, kNumTokens,
220 + kNumExperts, kHiddenDim>,
221 + logits, normed_x, x, norm_weight, gate_weight, eps);
222 +}
223 +
224 +// Template instantiations — DSV4-Pro is the only supported configuration:
225 +// num_experts=384, hidden_dim=7168. Other shapes (e.g. DSV4-Flash with
226 +// hidden_dim=4096) fall back to the unfused path on the Python side.
227 +#define INSTANTIATE(M) \
228 + template void invokeNormRouterGemm<__nv_bfloat16, M, 384, 7168>( \
229 + float*, __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, \
230 + __nv_bfloat16 const*, float, cudaStream_t);
231 +
232 +INSTANTIATE(1)
233 +INSTANTIATE(2)
234 +INSTANTIATE(3)
235 +INSTANTIATE(4)
236 +INSTANTIATE(5)
237 +INSTANTIATE(6)
238 +INSTANTIATE(7)
239 +INSTANTIATE(8)
240 +INSTANTIATE(9)
241 +INSTANTIATE(10)
242 +INSTANTIATE(11)
243 +INSTANTIATE(12)
244 +INSTANTIATE(13)
245 +INSTANTIATE(14)
246 +INSTANTIATE(15)
247 +INSTANTIATE(16)
248 +
249 +#undef INSTANTIATE
csrc/moe/moe_ops.h
+8
@@ -75,4 +75,12 @@ void shuffle_rows(const torch::Tensor& input_tensor,
75 // Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
76 void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
77 const torch::Tensor& mat_b);
78 +
79 +// Fused RMSNorm + router GEMV for DeepSeek V4. Produces both:
80 +// normed_x[m,k] = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
81 +// logits[m,n] = sum_k(normed_x[m,k] * gate_weight[n,k])
82 +// in a single kernel launch. Same dim/dtype constraints as dsv3_router_gemm.
83 +void dsv4_norm_router_gemm(at::Tensor& logits, at::Tensor& normed_x,
84 + at::Tensor const& x, at::Tensor const& norm_weight,
85 + at::Tensor const& gate_weight, double eps);
86 #endif
csrc/moe/torch_bindings.cpp
+6
@@ -135,6 +135,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
135 // DeepSeek V3 optimized router GEMM for SM90+
136 m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
137 // conditionally compiled so impl registration is in source file
138 +
139 + // DeepSeek V4 fused RMSNorm + router GEMV for SM90+
140 + m.def(
141 + "dsv4_norm_router_gemm(Tensor! logits, Tensor! normed_x, Tensor x, "
142 + "Tensor norm_weight, Tensor gate_weight, float eps) -> ()");
143 + // conditionally compiled so impl registration is in source file
144 #endif
145 }
146
vllm/_custom_ops.py
+30
@@ -2465,6 +2465,36 @@ def dsv3_router_gemm(
2465 return output
2466
2467
2468 +def dsv4_norm_router_gemm(
2469 + x: torch.Tensor,
2470 + norm_weight: torch.Tensor,
2471 + gate_weight: torch.Tensor,
2472 + eps: float,
2473 +) -> tuple[torch.Tensor, torch.Tensor]:
2474 + """Fused RMSNorm + router GEMV for DeepSeek V4.
2475 +
2476 + Returns ``(normed_x, router_logits)`` where
2477 + normed_x[m,k] = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
2478 + router_logits[m,n] = sum_k(normed_x[m,k] * gate_weight[n,k])
2479 +
2480 + DSV4-specific constraints (caller must check before dispatching here):
2481 + - x, norm_weight, gate_weight all bf16 contiguous
2482 + - x.shape == [num_tokens, 7168] with num_tokens in [1, 16]
2483 + - gate_weight.shape == [num_experts, 7168] with num_experts in {256, 384}
2484 + - SM 9.x or 10.x device
2485 +
2486 + Logits output is fp32 (hard-coded by DSV4 router).
2487 + """
2488 + num_tokens, hidden = x.shape
2489 + num_experts = gate_weight.shape[0]
2490 + normed_x = torch.empty_like(x)
2491 + logits = torch.empty(num_tokens, num_experts, device=x.device, dtype=torch.float32)
2492 + torch.ops._moe_C.dsv4_norm_router_gemm(
2493 + logits, normed_x, x, norm_weight, gate_weight, float(eps)
2494 + )
2495 + return normed_x, logits
2496 +
2497 +
2498 def topk_softmax(
2499 topk_weights: torch.Tensor,
2500 topk_ids: torch.Tensor,
vllm/model_executor/layers/fused_moe/router/norm_gate_linear.py new
+114
@@ -0,0 +1,114 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused RMSNorm + GateLinear for DeepSeek V4 MoE routing."""
4 +
5 +import torch
6 +from torch import nn
7 +
8 +import vllm._custom_ops as ops
9 +from vllm.model_executor.custom_op import PluggableLayer
10 +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
11 +from vllm.model_executor.layers.layernorm import RMSNorm
12 +from vllm.utils.torch_utils import direct_register_custom_op
13 +
14 +DSV4_PRO_NUM_EXPERTS = 384
15 +DSV4_PRO_HIDDEN_SIZE = 7168
16 +DSV4_PRO_MAX_NUM_TOKENS = 16
17 +
18 +
19 +def _dsv4_pro_norm_gate(
20 + x: torch.Tensor,
21 + norm_weight: torch.Tensor,
22 + gate_weight: torch.Tensor,
23 + rms_eps: float,
24 +) -> tuple[torch.Tensor, torch.Tensor]:
25 + """Runtime dispatcher: fused ``dsv4_norm_router_gemm`` (M<=16) vs the
26 + unfused ``rms_norm + dsv3_router_gemm`` fallback (M>16).
27 +
28 + """
29 + if x.shape[0] <= DSV4_PRO_MAX_NUM_TOKENS:
30 + return ops.dsv4_norm_router_gemm(x, norm_weight, gate_weight, rms_eps)
31 +
32 + normed = torch.empty_like(x)
33 + # Call `_C::rms_norm` here to avoid select the path of native rms
34 + torch.ops._C.rms_norm(normed, x, norm_weight, rms_eps)
35 + logits = torch.mm(normed, gate_weight.t(), out_dtype=torch.float32)
36 + return normed, logits
37 +
38 +
39 +def _dsv4_pro_norm_gate_fake(
40 + x: torch.Tensor,
41 + norm_weight: torch.Tensor,
42 + gate_weight: torch.Tensor,
43 + rms_eps: float,
44 +) -> tuple[torch.Tensor, torch.Tensor]:
45 + num_tokens = x.shape[0]
46 + num_experts = gate_weight.shape[0]
47 + return (
48 + torch.empty_like(x),
49 + torch.empty(num_tokens, num_experts, dtype=torch.float32, device=x.device),
50 + )
51 +
52 +
53 +direct_register_custom_op(
54 + op_name="dsv4_pro_norm_gate",
55 + op_func=_dsv4_pro_norm_gate,
56 + mutates_args=[],
57 + fake_impl=_dsv4_pro_norm_gate_fake,
58 +)
59 +
60 +
61 +@PluggableLayer.register("norm_gated_linear")
62 +class NormGateLinear(nn.Module):
63 + """RMSNorm + GateLinear, fused on DSV4-Pro only."""
64 +
65 + def __init__(
66 + self,
67 + hidden_size: int,
68 + num_experts: int,
69 + rms_eps: float = 1e-6,
70 + params_dtype: torch.dtype | None = None,
71 + prefix: str = "",
72 + ) -> None:
73 + super().__init__()
74 + self.hidden_size = hidden_size
75 + self.num_experts = num_experts
76 + self.rms_eps = rms_eps
77 +
78 + self.norm = RMSNorm(hidden_size, eps=rms_eps, dtype=params_dtype)
79 + self.gate = GateLinear(
80 + hidden_size,
81 + num_experts,
82 + bias=False,
83 + out_dtype=torch.float32, # DSV4 router output is fp32
84 + params_dtype=params_dtype,
85 + prefix=f"{prefix}.gate" if prefix else "gate",
86 + )
87 +
88 + self.e_score_correction_bias = None
89 + self.tid2eid = None
90 +
91 + self._fused_kernel_supported = (
92 + hidden_size == DSV4_PRO_HIDDEN_SIZE
93 + and num_experts == DSV4_PRO_NUM_EXPERTS
94 + and self.gate.allow_dsv3_router_gemm # cuda platform
95 + )
96 +
97 + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
98 + if self._fused_kernel_supported:
99 + assert x.shape[1] == DSV4_PRO_HIDDEN_SIZE
100 + assert self.gate.weight.shape == (
101 + DSV4_PRO_NUM_EXPERTS,
102 + DSV4_PRO_HIDDEN_SIZE,
103 + )
104 + # This must be wrapped in a custom op because our torch.compile integration
105 + # does not support runtime dispatching on num_tokens.
106 + return torch.ops.vllm.dsv4_pro_norm_gate(
107 + x, self.norm.weight, self.gate.weight, self.rms_eps
108 + )
109 +
110 + # Non-Pro fallback (e.g. DSV4-Flash with hidden_size=4096):
111 +
112 + normed_x = self.norm(x)
113 + logits, _ = self.gate(normed_x)
114 + return normed_x, logits
vllm/model_executor/models/deepseek_v4.py
+44 -42
@@ -23,11 +23,14 @@ from vllm.model_executor.layers.deepseek_v4_attention import (
23 DeepseekV4MLAModules,
24 DeepseekV4MultiHeadLatentAttentionWrapper,
25 )
26 -from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
26 +from vllm.model_executor.layers.fused_moe import FusedMoE
27 from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
28 from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
29 fused_topk_bias,
30 )
31 +from vllm.model_executor.layers.fused_moe.router.norm_gate_linear import (
32 + NormGateLinear,
33 +)
34 from vllm.model_executor.layers.layernorm import RMSNorm
35 from vllm.model_executor.layers.linear import (
36 ColumnParallelLinear,
@@ -755,23 +758,23 @@ class DeepseekV4MoE(nn.Module):
758 "deep_gemm_mega_moe for this checkpoint."
759 )
760
758 - self.gate = GateLinear(
759 - config.hidden_size,
760 - config.n_routed_experts,
761 - out_dtype=torch.float32,
762 - bias=False,
763 - prefix=f"{prefix}.gate",
761 + # Fused RMSNorm + gate: owns both ffn_norm and the gate matmul.
762 + self.norm_gate = NormGateLinear(
763 + hidden_size=config.hidden_size,
764 + num_experts=config.n_routed_experts,
765 + rms_eps=config.rms_norm_eps,
766 + prefix=f"{prefix}.norm_gate",
767 )
765 - self.gate.e_score_correction_bias = None
766 - self.gate.tid2eid = None
768 + # Routing-side tensors live on ``norm_gate`` directly (not on the
769 + # inner gate); they are initialized to None in NormGatedLinear and
770 + # populated below depending on the MoE variant.
771 is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
772 self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
769 -
773 if is_hash_moe:
774 # hash MoE doesn't use e_score_correction_bias
775 # Use randint instead of empty to avoid garbage values causing
776 # invalid memory access in dummy mode (--load-format="dummy")
774 - self.gate.tid2eid = nn.Parameter(
777 + self.norm_gate.tid2eid = nn.Parameter(
778 torch.randint(
779 0,
780 config.n_routed_experts,
@@ -781,7 +784,7 @@ class DeepseekV4MoE(nn.Module):
784 requires_grad=False,
785 )
786 elif getattr(config, "topk_method", None) == "noaux_tc":
784 - self.gate.e_score_correction_bias = nn.Parameter(
787 + self.norm_gate.e_score_correction_bias = nn.Parameter(
788 torch.empty(config.n_routed_experts, dtype=torch.float32),
789 requires_grad=False,
790 )
@@ -844,10 +847,9 @@ class DeepseekV4MoE(nn.Module):
847 self.n_local_experts = config.n_routed_experts // self.tp_size
848 self.experts_start_idx = self.tp_rank * self.n_local_experts
849 self.experts_end_idx = self.experts_start_idx + self.n_local_experts
847 -
850 + # We don't pass `gate` into FusedMoE
851 self.experts = FusedMoE(
852 shared_experts=self.shared_experts,
850 - gate=self.gate,
853 num_experts=config.n_routed_experts,
854 top_k=config.num_experts_per_tok,
855 hidden_size=config.hidden_size,
@@ -857,8 +859,8 @@ class DeepseekV4MoE(nn.Module):
859 prefix=f"{prefix}.experts",
860 scoring_func=self.scoring_func,
861 routed_scaling_factor=self.routed_scaling_factor,
860 - e_score_correction_bias=self.gate.e_score_correction_bias,
861 - hash_indices_table=self.gate.tid2eid,
862 + e_score_correction_bias=self.norm_gate.e_score_correction_bias,
863 + hash_indices_table=self.norm_gate.tid2eid,
864 swiglu_limit=self.swiglu_limit,
865 router_logits_dtype=torch.float32,
866 )
@@ -866,40 +868,40 @@ class DeepseekV4MoE(nn.Module):
868 def forward(
869 self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
870 ) -> torch.Tensor:
869 - if self.gate.tid2eid is not None and input_ids is None:
871 + if self.norm_gate.tid2eid is not None and input_ids is None:
872 raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
873
874 if not self.use_mega_moe:
875 return self._forward_fused_moe(hidden_states, input_ids)
876
877 org_shape = hidden_states.shape
876 - router_logits, _ = self.gate(hidden_states)
878 + normed_x, router_logits = self.norm_gate(hidden_states)
879 topk_weights, topk_ids = fused_topk_bias(
878 - hidden_states=hidden_states,
880 + hidden_states=normed_x,
881 gating_output=router_logits,
882 scoring_func=self.scoring_func,
881 - e_score_correction_bias=self.gate.e_score_correction_bias.data
882 - if self.gate.e_score_correction_bias is not None
883 + e_score_correction_bias=self.norm_gate.e_score_correction_bias.data
884 + if self.norm_gate.e_score_correction_bias is not None
885 else None,
886 topk=self.n_activated_experts,
887 renormalize=self.renormalize,
888 indices_type=self.hash_indices_dtype,
889 input_tokens=input_ids,
888 - hash_indices_table=self.gate.tid2eid,
890 + hash_indices_table=self.norm_gate.tid2eid,
891 routed_scaling_factor=self.routed_scaling_factor,
892 )
893 activation_clamp = (
894 float(self.swiglu_limit) if self.swiglu_limit is not None else None
895 )
896 final_hidden_states = self.experts(
895 - hidden_states,
897 + normed_x,
898 topk_weights,
899 topk_ids,
900 activation_clamp=activation_clamp,
901 )
902
903 if self.shared_experts is not None:
902 - shared_output = self.shared_experts(hidden_states)
904 + shared_output = self.shared_experts(normed_x)
905 final_hidden_states += shared_output
906
907 return final_hidden_states.view(org_shape)
@@ -907,21 +909,14 @@ class DeepseekV4MoE(nn.Module):
909 def _forward_fused_moe(
910 self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
911 ) -> torch.Tensor:
912 + assert not self.experts.is_internal_router
913 org_shape = hidden_states.shape
911 - if self.experts.is_internal_router:
912 - # In this case, the gate/router runs inside the FusedMoE class
913 - final_hidden_states = self.experts(
914 - hidden_states=hidden_states,
915 - router_logits=hidden_states,
916 - input_ids=input_ids,
917 - )
918 - else:
919 - router_logits, _ = self.gate(hidden_states)
920 - final_hidden_states = self.experts(
921 - hidden_states=hidden_states,
922 - router_logits=router_logits,
923 - input_ids=input_ids,
924 - )
914 + normed_x, router_logits = self.norm_gate(hidden_states)
915 + final_hidden_states = self.experts(
916 + hidden_states=normed_x,
917 + router_logits=router_logits,
918 + input_ids=input_ids,
919 + )
920
921 return final_hidden_states.view(org_shape)
922
@@ -1125,7 +1120,8 @@ class DeepseekV4DecoderLayer(nn.Module):
1120 self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
1121
1122 self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1128 - self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1123 + # ``ffn_norm`` is owned by ``self.ffn.norm_gate`` (fused with the
1124 + # router gate matmul); see ``NormGatedLinear``.
1125 self.hc_mult = config.hc_mult
1126 self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
1127 self.hc_eps = config.hc_eps
@@ -1255,8 +1251,8 @@ class DeepseekV4DecoderLayer(nn.Module):
1251 self.hc_post_alpha,
1252 self.hc_sinkhorn_iters,
1253 )
1258 -
1259 - x = self.ffn_norm(x)
1254 + # ffn_norm is now folded into self.ffn.norm_gate; ffn() takes
1255 + # the pre-norm activation directly.
1256 x = self.ffn(x, input_ids)
1257 return x, residual, post_mix, res_mix
1258
@@ -1629,7 +1625,13 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1625 orig_to_new_suffix={
1626 "head.weight": "lm_head.weight",
1627 "embed.weight": "embed_tokens.weight",
1632 - ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
1628 + # Pre-MoE norm + gate are now owned by ``DeepseekV4MoE.norm_gate``
1629 + # (see NormGatedLinear).
1630 + ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
1631 + ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
1632 + ".ffn.gate.bias": ".ffn.norm_gate.e_score_correction_bias",
1633 + # Hash MoE table also moved off the inner gate.
1634 + ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
1635 },
1636 orig_to_new_substr={
1637 ".attn.compressor.": ".attn.mla_attn.compressor.",
vllm/model_executor/models/deepseek_v4_mtp.py
+11 -1
@@ -290,6 +290,11 @@ class DeepSeekV4MTP(nn.Module):
290 ".emb.tok_emb.weight": ".embed_tokens.weight",
291 ".head.weight": ".shared_head.head.weight",
292 ".norm.weight": ".shared_head.norm.weight",
293 + # Pre-MoE norm + gate are now owned by
294 + # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear).
295 + ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
296 + ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
297 + ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
298 }
299
300 def _remap_weight_name(name: str) -> str:
@@ -437,7 +442,12 @@ class DeepSeekV4MTP(nn.Module):
442 ".shared_experts.w2", ".shared_experts.down_proj"
443 )
444 if name.endswith(".ffn.gate.bias"):
440 - name = name.replace(".bias", ".e_score_correction_bias")
445 + # ``e_score_correction_bias`` lives on
446 + # ``norm_gate`` directly (not on the inner gate).
447 + name = name.replace(
448 + ".ffn.gate.bias",
449 + ".ffn.norm_gate.e_score_correction_bias",
450 + )
451 param = params_dict[name]
452 weight_loader = getattr(
453 param, "weight_loader", default_weight_loader