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, 20:45 UTC
[Kernel] Remove NormGateLinear (#43554)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
May 25, 2026 at 17:49 UTC
d4004455d2357985830af10e432709b42c820455
16 files changed
+49
-836
CMakeLists.txt
-10
@@ -1240,16 +1240,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1240
CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
1241
list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
1242
message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
1243
-
1244
- # DeepSeek V4 fused RMSNorm + router GEMV - same arch gating as DSV3.
1245
- set(DSV4_NORM_ROUTER_GEMM_SRC
1246
- "csrc/moe/dsv4_norm_router_gemm_entry.cu"
1247
- "csrc/moe/dsv4_norm_router_gemm_kernel.cu")
1248
- set_gencode_flags_for_srcs(
1249
- SRCS "${DSV4_NORM_ROUTER_GEMM_SRC}"
1250
- CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
1251
- list(APPEND VLLM_MOE_EXT_SRC "${DSV4_NORM_ROUTER_GEMM_SRC}")
1252
- message(STATUS "Building DSV4 norm+router GEMV kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
1243
else()
1244
message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
1245
" (requires SM90+ and CUDA >= 12.0)")
benchmarks/kernels/benchmark_norm_router_gemm.py
deleted
-183
@@ -1,183 +0,0 @@
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/dsv3_router_gemm_bf16_out.cu
+1
-1
@@ -182,7 +182,7 @@ void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
182
config.stream = stream;
183
cudaLaunchAttribute attrs[1];
184
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
185
- attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
185
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
186
config.numAttrs = 1;
187
config.attrs = attrs;
188
cudaLaunchKernelEx(
csrc/moe/dsv3_router_gemm_float_out.cu
+1
-1
@@ -182,7 +182,7 @@ void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
182
config.stream = stream;
183
cudaLaunchAttribute attrs[1];
184
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
185
- attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
185
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
186
config.numAttrs = 1;
187
config.attrs = attrs;
188
cudaLaunchKernelEx(
csrc/moe/dsv3_router_gemm_utils.h
-12
@@ -29,15 +29,3 @@ inline int getSMVersion() {
29
auto* props = at::cuda::getCurrentDeviceProperties();
30
return props->major * 10 + props->minor;
31
}
32
-
33
-inline bool getEnvEnablePDL() {
34
- static std::once_flag flag;
35
- static bool enablePDL = false;
36
- std::call_once(flag, [&]() {
37
- if (getSMVersion() >= 90) {
38
- const char* env = std::getenv("TRTLLM_ENABLE_PDL");
39
- enablePDL = env && env[0] == '1' && env[1] == '\0';
40
- }
41
- });
42
- return enablePDL;
43
-}
csrc/moe/dsv4_norm_router_gemm.h
deleted
-30
@@ -1,30 +0,0 @@
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
deleted
-130
@@ -1,130 +0,0 @@
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
deleted
-249
@@ -1,249 +0,0 @@
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,12 +75,4 @@ 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);
78
#endif
csrc/moe/torch_bindings.cpp
-6
@@ -125,12 +125,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
125
// DeepSeek V3 optimized router GEMM for SM90+
126
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
127
// conditionally compiled so impl registration is in source file
128
-
129
- // DeepSeek V4 fused RMSNorm + router GEMV for SM90+
130
- m.def(
131
- "dsv4_norm_router_gemm(Tensor! logits, Tensor! normed_x, Tensor x, "
132
- "Tensor norm_weight, Tensor gate_weight, float eps) -> ()");
133
- // conditionally compiled so impl registration is in source file
128
#endif
129
}
130
vllm/_custom_ops.py
-30
@@ -2367,36 +2367,6 @@ def dsv3_router_gemm(
2367
return output
2368
2369
2370
-def dsv4_norm_router_gemm(
2371
- x: torch.Tensor,
2372
- norm_weight: torch.Tensor,
2373
- gate_weight: torch.Tensor,
2374
- eps: float,
2375
-) -> tuple[torch.Tensor, torch.Tensor]:
2376
- """Fused RMSNorm + router GEMV for DeepSeek V4.
2377
-
2378
- Returns ``(normed_x, router_logits)`` where
2379
- normed_x[m,k] = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
2380
- router_logits[m,n] = sum_k(normed_x[m,k] * gate_weight[n,k])
2381
-
2382
- DSV4-specific constraints (caller must check before dispatching here):
2383
- - x, norm_weight, gate_weight all bf16 contiguous
2384
- - x.shape == [num_tokens, 7168] with num_tokens in [1, 16]
2385
- - gate_weight.shape == [num_experts, 7168] with num_experts in {256, 384}
2386
- - SM 9.x or 10.x device
2387
-
2388
- Logits output is fp32 (hard-coded by DSV4 router).
2389
- """
2390
- num_tokens, hidden = x.shape
2391
- num_experts = gate_weight.shape[0]
2392
- normed_x = torch.empty_like(x)
2393
- logits = torch.empty(num_tokens, num_experts, device=x.device, dtype=torch.float32)
2394
- torch.ops._moe_C.dsv4_norm_router_gemm(
2395
- logits, normed_x, x, norm_weight, gate_weight, float(eps)
2396
- )
2397
- return normed_x, logits
2398
-
2399
-
2370
def topk_softmax(
2371
topk_weights: torch.Tensor,
2372
topk_ids: torch.Tensor,
vllm/model_executor/kernels/mhc/tilelang.py
+3
-3
@@ -170,7 +170,7 @@ def _mhc_pre_tilelang_fake(
170
sinkhorn_repeat: int,
171
n_splits: int = 1,
172
norm_weight: torch.Tensor | None = None,
173
- norm_eps: float = 0.0,
173
+ norm_eps: float = 1e-6,
174
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
175
hc_mult = residual.shape[-2]
176
hidden_size = residual.shape[-1]
@@ -238,7 +238,7 @@ def mhc_fused_post_pre_tilelang(
238
n_splits: int = 1,
239
tile_n: int = 1,
240
norm_weight: torch.Tensor | None = None,
241
- norm_eps: float = 0.0,
241
+ norm_eps: float = 1e-6,
242
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
243
"""
244
Run one MHC post block followed by the next MHC pre block.
@@ -450,7 +450,7 @@ def _mhc_fused_post_pre_tilelang_fake(
450
n_splits: int = 1,
451
tile_n: int = 1,
452
norm_weight: torch.Tensor | None = None,
453
- norm_eps: float = 0.0,
453
+ norm_eps: float = 1e-6,
454
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
455
hc_mult = residual.shape[-2]
456
hidden_size = residual.shape[-1]
vllm/model_executor/layers/fused_moe/router/norm_gate_linear.py
deleted
-114
@@ -1,114 +0,0 @@
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/models/deepseek_v4/amd/model.py
+42
-46
@@ -18,13 +18,10 @@ from vllm.distributed import (
18
)
19
from vllm.forward_context import get_forward_context
20
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
21
-from vllm.model_executor.layers.fused_moe import FusedMoE
21
+from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
22
from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
23
fused_topk_bias,
24
)
25
-from vllm.model_executor.layers.fused_moe.router.norm_gate_linear import (
26
- NormGateLinear,
27
-)
25
from vllm.model_executor.layers.layernorm import RMSNorm
26
from vllm.model_executor.layers.linear import (
27
ColumnParallelLinear,
@@ -655,23 +652,23 @@ class DeepseekV4MoE(nn.Module):
652
"deep_gemm_mega_moe for this checkpoint."
653
)
654
658
- # Fused RMSNorm + gate: owns both ffn_norm and the gate matmul.
659
- self.norm_gate = NormGateLinear(
660
- hidden_size=config.hidden_size,
661
- num_experts=config.n_routed_experts,
662
- rms_eps=config.rms_norm_eps,
663
- prefix=f"{prefix}.norm_gate",
655
+ self.gate = GateLinear(
656
+ input_size=config.hidden_size,
657
+ output_size=config.n_routed_experts,
658
+ bias=False,
659
+ out_dtype=torch.float32,
660
+ prefix=f"{prefix}.gate",
661
)
665
- # Routing-side tensors live on ``norm_gate`` directly (not on the
666
- # inner gate); they are initialized to None in NormGatedLinear and
667
- # populated below depending on the MoE variant.
662
+
663
+ self.gate.e_score_correction_bias = None
664
+ self.gate.tid2eid = None
665
is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
666
self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
667
if is_hash_moe:
668
# hash MoE doesn't use e_score_correction_bias
669
# Use randint instead of empty to avoid garbage values causing
670
# invalid memory access in dummy mode (--load-format="dummy")
674
- self.norm_gate.tid2eid = nn.Parameter(
671
+ self.gate.tid2eid = nn.Parameter(
672
torch.randint(
673
0,
674
config.n_routed_experts,
@@ -681,7 +678,7 @@ class DeepseekV4MoE(nn.Module):
678
requires_grad=False,
679
)
680
elif getattr(config, "topk_method", None) == "noaux_tc":
684
- self.norm_gate.e_score_correction_bias = nn.Parameter(
681
+ self.gate.e_score_correction_bias = nn.Parameter(
682
torch.empty(config.n_routed_experts, dtype=torch.float32),
683
requires_grad=False,
684
)
@@ -744,9 +741,10 @@ class DeepseekV4MoE(nn.Module):
741
self.n_local_experts = config.n_routed_experts // self.tp_size
742
self.experts_start_idx = self.tp_rank * self.n_local_experts
743
self.experts_end_idx = self.experts_start_idx + self.n_local_experts
747
- # We don't pass `gate` into FusedMoE
744
+
745
self.experts = FusedMoE(
746
shared_experts=self.shared_experts,
747
+ gate=self.gate,
748
num_experts=config.n_routed_experts,
749
top_k=config.num_experts_per_tok,
750
hidden_size=config.hidden_size,
@@ -756,8 +754,8 @@ class DeepseekV4MoE(nn.Module):
754
prefix=f"{prefix}.experts",
755
scoring_func=self.scoring_func,
756
routed_scaling_factor=self.routed_scaling_factor,
759
- e_score_correction_bias=self.norm_gate.e_score_correction_bias,
760
- hash_indices_table=self.norm_gate.tid2eid,
757
+ e_score_correction_bias=self.gate.e_score_correction_bias,
758
+ hash_indices_table=self.gate.tid2eid,
759
swiglu_limit=self.swiglu_limit,
760
router_logits_dtype=torch.float32,
761
)
@@ -765,40 +763,40 @@ class DeepseekV4MoE(nn.Module):
763
def forward(
764
self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
765
) -> torch.Tensor:
768
- if self.norm_gate.tid2eid is not None and input_ids is None:
766
+ if self.gate.tid2eid is not None and input_ids is None:
767
raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
768
769
if not self.use_mega_moe:
770
return self._forward_fused_moe(hidden_states, input_ids)
771
772
org_shape = hidden_states.shape
775
- normed_x, router_logits = self.norm_gate(hidden_states)
773
+ router_logits, _ = self.gate(hidden_states)
774
topk_weights, topk_ids = fused_topk_bias(
777
- hidden_states=normed_x,
775
+ hidden_states=hidden_states,
776
gating_output=router_logits,
777
scoring_func=self.scoring_func,
780
- e_score_correction_bias=self.norm_gate.e_score_correction_bias.data
781
- if self.norm_gate.e_score_correction_bias is not None
778
+ e_score_correction_bias=self.gate.e_score_correction_bias.data
779
+ if self.gate.e_score_correction_bias is not None
780
else None,
781
topk=self.n_activated_experts,
782
renormalize=self.renormalize,
783
indices_type=self.hash_indices_dtype,
784
input_tokens=input_ids,
787
- hash_indices_table=self.norm_gate.tid2eid,
785
+ hash_indices_table=self.gate.tid2eid,
786
routed_scaling_factor=self.routed_scaling_factor,
787
)
788
activation_clamp = (
789
float(self.swiglu_limit) if self.swiglu_limit is not None else None
790
)
791
final_hidden_states = self.experts(
794
- normed_x,
792
+ hidden_states,
793
topk_weights,
794
topk_ids,
795
activation_clamp=activation_clamp,
796
)
797
798
if self.shared_experts is not None:
801
- shared_output = self.shared_experts(normed_x)
799
+ shared_output = self.shared_experts(hidden_states)
800
final_hidden_states += shared_output
801
802
return final_hidden_states.view(org_shape)
@@ -806,14 +804,21 @@ class DeepseekV4MoE(nn.Module):
804
def _forward_fused_moe(
805
self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
806
) -> torch.Tensor:
809
- assert not self.experts.is_internal_router
807
org_shape = hidden_states.shape
811
- normed_x, router_logits = self.norm_gate(hidden_states)
812
- final_hidden_states = self.experts(
813
- hidden_states=normed_x,
814
- router_logits=router_logits,
815
- input_ids=input_ids,
816
- )
808
+ if self.experts.is_internal_router:
809
+ # In this case, the gate/router runs inside the FusedMoE class
810
+ final_hidden_states = self.experts(
811
+ hidden_states=hidden_states,
812
+ router_logits=hidden_states,
813
+ input_ids=input_ids,
814
+ )
815
+ else:
816
+ router_logits, _ = self.gate(hidden_states)
817
+ final_hidden_states = self.experts(
818
+ hidden_states=hidden_states,
819
+ router_logits=router_logits,
820
+ input_ids=input_ids,
821
+ )
822
823
return final_hidden_states.view(org_shape)
824
@@ -1017,8 +1022,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1022
self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
1023
1024
self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1020
- # ``ffn_norm`` is owned by ``self.ffn.norm_gate`` (fused with the
1021
- # router gate matmul); see ``NormGatedLinear``.
1025
+ self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1026
self.hc_mult = config.hc_mult
1027
self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
1028
self.hc_eps = config.hc_eps
@@ -1148,8 +1152,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1152
self.hc_post_alpha,
1153
self.hc_sinkhorn_iters,
1154
)
1151
- # ffn_norm is now folded into self.ffn.norm_gate; ffn() takes
1152
- # the pre-norm activation directly.
1155
+ x = self.ffn_norm(x)
1156
x = self.ffn(x, input_ids)
1157
return x, residual, post_mix, res_mix
1158
@@ -1176,8 +1179,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1179
x, post, comb = self.hc_pre(
1180
x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
1181
)
1179
- # ffn_norm is now folded into self.ffn.norm_gate; ffn() takes
1180
- # the pre-norm activation directly.
1182
+ x = self.ffn_norm(x)
1183
x = self.ffn(x, input_ids)
1184
x = self.hc_post(x, residual, post, comb)
1185
return x, None, None, None
@@ -1527,13 +1529,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1529
orig_to_new_suffix={
1530
"head.weight": "lm_head.weight",
1531
"embed.weight": "embed_tokens.weight",
1530
- # Pre-MoE norm + gate are now owned by ``DeepseekV4MoE.norm_gate``
1531
- # (see NormGatedLinear).
1532
- ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
1533
- ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
1534
- ".ffn.gate.bias": ".ffn.norm_gate.e_score_correction_bias",
1535
- # Hash MoE table also moved off the inner gate.
1536
- ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
1532
+ ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
1533
},
1534
orig_to_new_substr={
1535
".attn.compressor.": ".attn.mla_attn.compressor.",
vllm/models/deepseek_v4/amd/mtp.py
+1
-11
@@ -292,11 +292,6 @@ class DeepSeekV4MTP(nn.Module):
292
".emb.tok_emb.weight": ".embed_tokens.weight",
293
".head.weight": ".shared_head.head.weight",
294
".norm.weight": ".shared_head.norm.weight",
295
- # Pre-MoE norm + gate are now owned by
296
- # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear).
297
- ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
298
- ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
299
- ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
295
}
296
297
def _remap_weight_name(name: str) -> str:
@@ -444,12 +439,7 @@ class DeepSeekV4MTP(nn.Module):
439
".shared_experts.w2", ".shared_experts.down_proj"
440
)
441
if name.endswith(".ffn.gate.bias"):
447
- # ``e_score_correction_bias`` lives on
448
- # ``norm_gate`` directly (not on the inner gate).
449
- name = name.replace(
450
- ".ffn.gate.bias",
451
- ".ffn.norm_gate.e_score_correction_bias",
452
- )
442
+ name = name.replace(".bias", ".e_score_correction_bias")
443
param = params_dict[name]
444
weight_loader = getattr(
445
param, "weight_loader", default_weight_loader
vllm/models/deepseek_v4/nvidia/model.py
+1
-2
@@ -1181,8 +1181,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1181
norm_weight=ffn_norm_weight,
1182
norm_eps=ffn_norm_eps,
1183
)
1184
- # ffn_norm is fused into mhc_fused_post_pre above; ffn() takes the
1185
- # already-normed activation directly.
1184
+
1185
x = self.ffn(x, input_ids)
1186
return x, residual, post_mix, res_mix
1187