minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[DSV4.1] Reduce mHC, metadata and small-batch router overhead (#39704)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-09-19 19:51:21 +08:00
committed by GitHub
parent d1acbe0746
commit 7fac84b639
15 changed files with 1039 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/mhc_post_combine_norm_prefill.cuh
+190
View File
@@ -0,0 +1,190 @@
// Prefill mHC post/combine/RMSNorm with the original BF16 intermediates.
// Stage the collapsed row in shared memory to bound register use while preserving
// the original Triton prefill normalization reduction and PTX arithmetic.
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace sglang {
struct MhcPostCombineNormPrefillParams {
const bf16_t *x, *residual, *weight;
bf16_t *residual_out, *output;
const float *post, *comb, *pre;
float eps;
};
template <int Threads>
__global__ __launch_bounds__(Threads) void mhc_post_combine_norm_prefill_kernel(
const __grid_constant__ MhcPostCombineNormPrefillParams p) {
using namespace device;
using V = AlignedVector<bf16x2_t, 4>;
__shared__ __align__(16) bf16_t collapsed[5120];
__shared__ float warp_sums[4];
__shared__ float inv_rms;
const int tid = threadIdx.x;
const int lane = tid % 32;
const int64_t row = blockIdx.x;
const float coeff = lane < 16 ? p.comb[row * 16 + lane]
: lane < 20 ? p.post[row * 4 + lane - 16]
: lane < 24 ? p.pre[row * 4 + lane - 20]
: 0.f;
// Vectorize loads but keep only one 8-element tile live at a time.
#pragma unroll 1
for (int vid = tid; vid < 640; vid += Threads) {
V x, residual[4];
x.load(p.x + row * 5120, vid);
#pragma unroll
for (int j = 0; j < 4; ++j)
residual[j].load(p.residual + row * 20480 + j * 5120, vid);
float combined[8] = {};
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float post = __shfl_sync(0xffffffff, coeff, 16 + i);
const float pre = __shfl_sync(0xffffffff, coeff, 20 + i);
const float c0 = __shfl_sync(0xffffffff, coeff, i);
const float c1 = __shfl_sync(0xffffffff, coeff, 4 + i);
const float c2 = __shfl_sync(0xffffffff, coeff, 8 + i);
const float c3 = __shfl_sync(0xffffffff, coeff, 12 + i);
V updated;
#pragma unroll
for (int k = 0; k < 4; ++k) {
const auto xx = cast<fp32x2_t>(x[k]);
const auto r0 = cast<fp32x2_t>(residual[0][k]);
const auto r1 = cast<fp32x2_t>(residual[1][k]);
const auto r2 = cast<fp32x2_t>(residual[2][k]);
const auto r3 = cast<fp32x2_t>(residual[3][k]);
float a = __fmaf_rn(post, xx.x, __fmul_rn(c0, r0.x));
float b = __fmaf_rn(post, xx.y, __fmul_rn(c0, r0.y));
a = __fmaf_rn(c1, r1.x, a);
b = __fmaf_rn(c1, r1.y, b);
a = __fmaf_rn(c2, r2.x, a);
b = __fmaf_rn(c2, r2.y, b);
a = __fmaf_rn(c3, r3.x, a);
b = __fmaf_rn(c3, r3.y, b);
updated[k] = cast<bf16x2_t>(fp32x2_t{a, b});
const auto rounded = cast<fp32x2_t>(updated[k]);
combined[2 * k] = __fmaf_rn(pre, rounded.x, combined[2 * k]);
combined[2 * k + 1] = __fmaf_rn(pre, rounded.y, combined[2 * k + 1]);
}
updated.store(p.residual_out + row * 20480 + i * 5120, vid);
}
V rounded;
#pragma unroll
for (int k = 0; k < 4; ++k)
rounded[k] = cast<bf16x2_t>(fp32x2_t{combined[2 * k], combined[2 * k + 1]});
rounded.store(collapsed, vid);
}
__syncthreads();
// Match the actual prefill Triton layout: 8 consecutive elements/thread,
// 128 threads, then stride 1024. The PTX folds each thread's elements in
// order, followed by XOR 16,8,4,2,1 and a four-warp XOR 2,1 reduction.
if (tid < 128) {
float sum = 0.f;
#pragma unroll
for (int group = 0; group < 5; ++group) {
V v;
v.load(collapsed + group * 1024, tid);
#pragma unroll
for (int k = 0; k < 4; ++k) {
const auto value = cast<fp32x2_t>(v[k]);
if (group == 0 && k == 0)
sum = __fadd_rn(__fmul_rn(value.y, value.y), __fmul_rn(value.x, value.x));
else {
sum = __fmaf_rn(value.x, value.x, sum);
sum = __fadd_rn(__fmul_rn(value.y, value.y), sum);
}
}
}
#pragma unroll
for (int offset = 16; offset; offset >>= 1)
sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, offset));
if (lane == 0) warp_sums[tid / 32] = sum;
}
__syncthreads();
if (tid < 32) {
float sum = lane < 4 ? warp_sums[lane] : 0.f;
sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 2));
sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 1));
if (lane == 0) {
float mean, inverse;
asm("div.full.f32 %0, %1, %2;" : "=f"(mean) : "f"(sum), "f"(5120.f));
// Match Triton's unqualified PTX add: ptxas may contract the
// constant division's reciprocal multiply with this epsilon add.
asm("add.f32 %0, %1, %2;" : "=f"(mean) : "f"(mean), "f"(p.eps));
asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(inverse) : "f"(mean));
inv_rms = inverse;
}
}
__syncthreads();
#pragma unroll 1
for (int vid = tid; vid < 640; vid += Threads) {
V v, w, out;
v.load(collapsed, vid);
w.load(p.weight, vid);
#pragma unroll
for (int k = 0; k < 4; ++k) {
const auto value = cast<fp32x2_t>(v[k]);
const auto weight = cast<fp32x2_t>(w[k]);
out[k] = cast<bf16x2_t>(
fp32x2_t{__fmul_rn(__fmul_rn(value.x, inv_rms), weight.x), __fmul_rn(__fmul_rn(value.y, inv_rms), weight.y)});
}
out.store(p.output + row * 5120, vid);
}
}
template <int Threads>
struct MhcPostCombineNormPrefill {
static void
run(tvm::ffi::TensorView x,
tvm::ffi::TensorView residual,
tvm::ffi::TensorView post,
tvm::ffi::TensorView comb,
tvm::ffi::TensorView pre,
tvm::ffi::TensorView weight,
tvm::ffi::TensorView residual_out,
tvm::ffi::TensorView output,
float eps) {
using namespace host;
auto m = SymbolicSize{"num_tokens"};
auto dev = SymbolicDevice{};
dev.set_options<kDLCUDA>();
TensorMatcher({m, 5120}).with_dtype<bf16_t>().with_device(dev).verify(x).verify(output);
TensorMatcher({m, 4, 5120}).with_dtype<bf16_t>().with_device(dev).verify(residual).verify(residual_out);
TensorMatcher({m, 4}).with_dtype<float>().with_device(dev).verify(post).verify(pre);
TensorMatcher({m, 4, 4}).with_dtype<float>().with_device(dev).verify(comb);
TensorMatcher({5120}).with_dtype<bf16_t>().with_device(dev).verify(weight);
CHECK_HOST(m.unwrap() >= 4096 && m.unwrap() <= 65536) << "prefill rows must be in [4096, 65536]";
for (const auto* ptr :
{x.data_ptr(), residual.data_ptr(), weight.data_ptr(), residual_out.data_ptr(), output.data_ptr()}) {
CHECK_HOST(reinterpret_cast<uintptr_t>(ptr) % 16 == 0) << "BF16 pointers must be 16-byte aligned";
}
CHECK_HOST(
residual_out.data_ptr() != residual.data_ptr() && residual_out.data_ptr() != x.data_ptr() &&
output.data_ptr() != residual.data_ptr() && output.data_ptr() != x.data_ptr() &&
output.data_ptr() != residual_out.data_ptr())
<< "outputs must not alias inputs or each other";
const MhcPostCombineNormPrefillParams params{
static_cast<const bf16_t*>(x.data_ptr()),
static_cast<const bf16_t*>(residual.data_ptr()),
static_cast<const bf16_t*>(weight.data_ptr()),
static_cast<bf16_t*>(residual_out.data_ptr()),
static_cast<bf16_t*>(output.data_ptr()),
static_cast<const float*>(post.data_ptr()),
static_cast<const float*>(comb.data_ptr()),
static_cast<const float*>(pre.data_ptr()),
eps};
LaunchKernel(dim3(m.unwrap()), Threads, dev.unwrap()).launch(mhc_post_combine_norm_prefill_kernel<Threads>, params);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/all_reduce_fusion.cuh
+49 -5
View File
@@ -286,7 +286,8 @@ template <
bool kNorm,
typename WeightT,
bool kMhc = false,
bool kQuant = false>
bool kQuant = false,
bool kCollapse = false>
__global__ __launch_bounds__(RowClusterTrait<kHiddenDim, kClusterSize>::kBlockSize)
__cluster_dims__(1, kClusterSize, 1) void moe_finalize_all_reduce_kernel(
const __grid_constant__ MoeFinalizeAllReduceParams<kWorldSize, WeightT> params) {
@@ -377,7 +378,14 @@ __global__ __launch_bounds__(RowClusterTrait<kHiddenDim, kClusterSize>::kBlockSi
if constexpr (!kNorm) {
const auto red = reduce_vec(vec);
ptx::st_global_16B(red, params.out, vid);
if constexpr (kMhc) mhc_post_vec<kHiddenDim>(params, red, row_idx, hvec);
if constexpr (kMhc) {
if constexpr (kCollapse) {
const auto combined = mhc_post_vec<kHiddenDim, true>(params, red, row_idx, hvec);
ptx::st_global_16B(combined, params.normalized, vid);
} else {
mhc_post_vec<kHiddenDim>(params, red, row_idx, hvec);
}
}
// ensure epoch is consumed, so flipping it won't lead to error
barrier_cluster_wait();
} else {
@@ -447,7 +455,8 @@ template <
bool kUsePDL,
typename WeightT,
bool kMhc = false,
bool kQuant = false>
bool kQuant = false,
bool kCollapse = false>
struct MoeFinalizeAllReduceKernel {
private:
static_assert(std::is_same_v<WeightT, bf16_t> || std::is_same_v<WeightT, fp32_t>);
@@ -466,7 +475,8 @@ struct MoeFinalizeAllReduceKernel {
kNorm,
WeightT,
kMhc,
kQuant>;
kQuant,
kCollapse>;
public:
/// out = [allreduce over ranks of] finalize(gemm2_out, idx, weights) [+ shared] [-> RMSNorm(norm_weight, eps)].
@@ -528,6 +538,40 @@ struct MoeFinalizeAllReduceKernel {
comb);
}
// Keep the original RMSNorm as a separate kernel, while reusing the
// BF16-rounded post values for the next sublayer's pre-combine.
static void run_mhc_combine(
CommunicatorRef ref,
TensorView out,
TensorView gemm2_out,
TensorView permuted_idx,
TensorView expert_weights,
std::optional<TensorView> shared_output,
TensorView mhc_out,
TensorView residual,
TensorView post,
TensorView comb,
TensorView pre,
TensorView combined) {
static_assert(kMhc && kCollapse && !kQuant);
run_impl(
ref,
out,
gemm2_out,
permuted_idx,
expert_weights,
shared_output,
std::nullopt,
0.0,
false,
mhc_out,
residual,
post,
comb,
pre,
combined);
}
static void run_mhc_norm(
CommunicatorRef ref,
TensorView out,
@@ -673,7 +717,7 @@ struct MoeFinalizeAllReduceKernel {
}
if constexpr (kMhc) {
static_assert(kHiddenDim == 5120);
if (norm_weight.has_value()) {
if (norm_weight.has_value() || kCollapse) {
TensorMatcher({T, 4}).with_dtype<fp32_t>().with_device<kDLCUDA>(device).verify(pre.value());
TensorMatcher({T, kHiddenDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(normalized.value());
}
python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py
+65 -1
View File
@@ -6,6 +6,7 @@ import triton.language as tl
from triton.language.extra import libdevice
from sglang.kernels.ops.attention.dsv4.torch_quant import FP4_AMAX_FLOOR
from sglang.srt.runtime_context import get_platform
INDEX_HEAD_DIM = 128
# One index-K slot: 64 packed e2m1 bytes and four ue8m0 block exponents.
@@ -132,6 +133,56 @@ def _quantize_fp4_indexer_kernel(
tl.store(x_fp4 + token_id * (BLOCK_N // 2) + pair_offsets, packed)
@triton.jit
def _quantize_fp4_indexer_rows(
x,
x_fp4,
x_sf,
M,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
GROUP_N: tl.constexpr,
RNE: tl.constexpr,
):
tl.static_assert(BLOCK_N == 128 and GROUP_N == 32)
# Each reduction covers one scale group. Keep its values for packing,
# avoiding four masked full-row reductions and a second input load.
group = tl.program_id(0) * BLOCK_M * 4 + tl.arange(0, BLOCK_M * 4)
offs = tl.arange(0, GROUP_N)
values = tl.load(
x + group[:, None].to(tl.int64) * GROUP_N + offs[None, :],
group[:, None] < M * 4,
0,
).to(tl.float32)
amax = tl.max(tl.abs(values), axis=1)
exp = _ceil_ue8m0_exp(tl.maximum(amax / 6.0, 1.0e-4))
scale = (exp << 23).to(tl.float32, bitcast=True)
v0, v1 = tl.split(
tl.reshape(values / scale[:, None], (BLOCK_M * 4, GROUP_N // 2, 2))
)
if RNE:
code0 = _fp4_e2m1_code_rne(v0)
code1 = _fp4_e2m1_code_rne(v1)
else:
code0 = _fp4_e2m1_code(v0)
code1 = _fp4_e2m1_code(v1)
packed = (code0 & 0x0F) | ((code1 & 0x0F) << 4)
tl.store(
x_fp4
+ group[:, None].to(tl.int64) * (GROUP_N // 2)
+ tl.arange(0, GROUP_N // 2)[None, :],
packed,
group[:, None] < M * 4,
)
# The four exponents occupy disjoint bytes, so integer sum packs them.
shifts = tl.arange(0, 4) * 8
packed_sf = tl.sum(
tl.reshape(exp.to(tl.uint32), (BLOCK_M, 4)) << shifts[None, :], axis=1
)
token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
tl.store(x_sf + token_id, packed_sf.to(tl.int32), token_id < M)
@triton.jit
def _store_fp4_index_k_cache_kernel(
k_fp4,
@@ -169,7 +220,20 @@ def quantize_fp4_indexer_tensor(
x = x.contiguous().view(-1, x.shape[-1])
x_fp4 = torch.empty((x.shape[0], 64), device=x.device, dtype=torch.int8)
x_sf = torch.empty((x.shape[0],), device=x.device, dtype=torch.int32)
if x.shape[0] > 0:
if x.shape[0] >= 4096 and get_platform().is_blackwell:
# Independent rows share a CTA to avoid one block per 128 values.
_quantize_fp4_indexer_rows[(triton.cdiv(x.shape[0], 8),)](
x,
x_fp4,
x_sf,
x.shape[0],
8,
BLOCK_N=128,
GROUP_N=32,
RNE=rne,
num_warps=4,
)
elif x.shape[0] > 0:
_quantize_fp4_indexer_kernel[(x.shape[0],)](
x,
x_fp4,
python/sglang/kernels/ops/communication/all_reduce_mhc_combine.py
+121
View File
@@ -0,0 +1,121 @@
"""Finalize/all-reduce + post + combine, retaining standalone RMSNorm."""
from typing import Optional
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernels.ops.communication.all_reduce_fusion import (
default_cluster_size,
get_registered_comm,
)
from sglang.srt.utils.custom_op import register_custom_op
@cache_once
def _module(world_size, top_k, cluster_size, weight_dtype):
args = make_cpp_args(
world_size,
5120,
top_k,
cluster_size,
is_arch_support_pdl(),
weight_dtype,
True,
False,
True,
)
return load_jit(
"moe_finalize_all_reduce_mhc_combine",
*args,
cuda_files=["distributed/all_reduce_fusion.cuh"],
cuda_wrappers=[("run", f"MoeFinalizeAllReduceKernel<{args}>::run_mhc_combine")],
)
@register_custom_op(mutates_args=["out", "mhc_out", "combined"])
def _moe_finalize_all_reduce_mhc_combine(
world_size: int,
top_k: int,
cluster_size: int,
out: torch.Tensor,
mhc_out: torch.Tensor,
combined: torch.Tensor,
gemm2: torch.Tensor,
idx: torch.Tensor,
weights: torch.Tensor,
shared: Optional[torch.Tensor],
residual: torch.Tensor,
post: torch.Tensor,
comb: torch.Tensor,
pre: torch.Tensor,
) -> None:
comm = get_registered_comm(world_size)
assert comm is not None
_module(world_size, top_k, cluster_size, weights.dtype).run(
comm,
out,
gemm2,
idx,
weights,
shared,
mhc_out,
residual,
post,
comb,
pre,
combined,
)
def moe_finalize_all_reduce_mhc_combine(
gemm2,
idx,
weights,
top_k,
shared,
residual,
post,
comb,
pre,
*,
world_size,
cluster_size=None,
):
out = torch.empty(
(weights.shape[0], 5120), dtype=torch.bfloat16, device=gemm2.device
)
mhc_out = torch.empty_like(residual)
combined = torch.empty_like(out)
if weights.shape[0]:
_moe_finalize_all_reduce_mhc_combine(
world_size,
top_k,
cluster_size or default_cluster_size(5120),
out,
mhc_out,
combined,
gemm2,
idx,
weights,
shared,
residual,
post,
comb,
pre,
)
return out, mhc_out, combined
def all_reduce_mhc_combine(x, residual, post, comb, pre, *, world_size):
from sglang.kernels.ops.communication.all_reduce_mhc import _identity_routing
idx, weights = _identity_routing(x.shape[0], x.device)
return moe_finalize_all_reduce_mhc_combine(
x, idx, weights, 1, None, residual, post, comb, pre, world_size=world_size
)
python/sglang/kernels/ops/layernorm/mhc_post_combine.py
+85
View File
@@ -0,0 +1,85 @@
"""HC=4 post-mix and pre-combine with the original BF16 intermediates."""
import torch
import triton
import triton.language as tl
@triton.jit
def _mhc_post_combine(X, R, P, C, A, RO, Y, H: tl.constexpr, B: tl.constexpr):
row = tl.program_id(0)
h = tl.program_id(1) * B + tl.arange(0, B)
mask = h < H
x = tl.load(X + row * H + h, mask, 0).to(tl.float32)
r0 = tl.load(R + (row * 4 + 0) * H + h, mask, 0).to(tl.float32)
r1 = tl.load(R + (row * 4 + 1) * H + h, mask, 0).to(tl.float32)
r2 = tl.load(R + (row * 4 + 2) * H + h, mask, 0).to(tl.float32)
r3 = tl.load(R + (row * 4 + 3) * H + h, mask, 0).to(tl.float32)
collapsed = tl.full((B,), 0, tl.float32)
for i in tl.static_range(4):
post = tl.load(P + row * 4 + i)
c0 = tl.load(C + row * 16 + i)
c1 = tl.load(C + row * 16 + 4 + i)
c2 = tl.load(C + row * 16 + 8 + i)
c3 = tl.load(C + row * 16 + 12 + i)
pre = tl.load(A + row * 4 + i)
# Match mhc_post_split_h's contraction order and BF16 store.
mixed = tl.fma(post, x, c0 * r0)
mixed = tl.fma(c1, r1, mixed)
mixed = tl.fma(c2, r2, mixed)
mixed = tl.fma(c3, r3, mixed)
rounded = mixed.to(tl.bfloat16)
tl.store(RO + (row * 4 + i) * H + h, rounded, mask)
# Match hc_combine's sequential FP32 accumulation, then BF16 store.
collapsed = tl.fma(pre, rounded.to(tl.float32), collapsed)
tl.store(Y + row * H + h, collapsed, mask)
def mhc_post_combine(x, residual, post, comb, pre):
"""Return (updated HC streams, collapsed BF16 input) out of place."""
assert x.dtype == residual.dtype == torch.bfloat16
assert post.dtype == comb.dtype == pre.dtype == torch.float32
assert x.ndim == 2 and x.shape[1] == 5120
assert residual.shape == (x.shape[0], 4, x.shape[1])
assert post.shape == pre.shape == (x.shape[0], 4)
assert comb.shape == (x.shape[0], 4, 4)
assert all(t.is_contiguous() for t in (x, residual, post, comb, pre))
updated = torch.empty_like(residual)
combined = torch.empty_like(x)
if x.shape[0]:
block = 1024 if x.shape[0] <= 192 or x.shape[0] >= 4096 else 512
_mhc_post_combine[(x.shape[0], triton.cdiv(x.shape[1], block))](
x,
residual,
post,
comb,
pre,
updated,
combined,
H=x.shape[1],
B=block,
num_warps=4,
enable_fp_fusion=False,
)
return updated, combined
@triton.jit
def _hc_norm_prefill(X, W, Y, EPS: tl.constexpr):
row = tl.program_id(0).to(tl.int64)
h = tl.arange(0, 8192)
value = tl.load(X + row * 5120 + h, h < 5120, 0).to(tl.float32)
inv_rms = tl.rsqrt(tl.sum(value * value, 0) / 5120 + EPS)
weight = tl.load(W + h, h < 5120, 0).to(tl.float32)
tl.store(Y + row * 5120 + h, value * inv_rms * weight, h < 5120)
def hc_norm_prefill(combined, weight, eps):
"""Keep the original fused prefill norm's arithmetic on collapsed input."""
assert 4096 <= combined.shape[0] <= 65536 and combined.shape[1] == 5120
assert weight.shape == (5120,)
assert combined.dtype == weight.dtype == torch.bfloat16
assert combined.is_contiguous() and weight.is_contiguous()
output = torch.empty_like(combined)
_hc_norm_prefill[(combined.shape[0],)](combined, weight, output, eps, num_warps=4)
return output
python/sglang/kernels/ops/layernorm/mhc_post_combine_norm_prefill.py
+38
View File
@@ -0,0 +1,38 @@
"""Prefill post/combine/RMSNorm with the original BF16 boundaries and norm tree."""
import torch
from sglang.kernels.jit.utils import cache_once, load_jit
from sglang.srt.utils.custom_op import register_custom_op
@cache_once
def _module():
return load_jit(
"mhc_post_combine_norm_prefill",
cuda_files=["deepseek_v4/mhc_post_combine_norm_prefill.cuh"],
cuda_wrappers=[("run", "MhcPostCombineNormPrefill<128>::run")],
)
@register_custom_op(mutates_args=["updated", "normalized"])
def _mhc_post_combine_norm_prefill(
x: torch.Tensor,
residual: torch.Tensor,
post: torch.Tensor,
comb: torch.Tensor,
pre: torch.Tensor,
weight: torch.Tensor,
updated: torch.Tensor,
normalized: torch.Tensor,
eps: float,
) -> None:
_module().run(x, residual, post, comb, pre, weight, updated, normalized, eps)
def mhc_post_combine_norm_prefill(x, residual, post, comb, pre, weight, eps):
updated, normalized = torch.empty_like(residual), torch.empty_like(x)
_mhc_post_combine_norm_prefill(
x, residual, post, comb, pre, weight, updated, normalized, eps
)
return updated, normalized
python/sglang/kernels/ops/moe/moe_fused_gate.py
+12 -1
View File
@@ -8,7 +8,12 @@ import triton
import triton.language as tl
from triton.language.extra import libdevice
from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl, load_jit
from sglang.kernels.jit.utils import (
cache_once,
get_jit_cuda_arch,
is_arch_support_pdl,
load_jit,
)
from sglang.kernels.kernel_api_logging import debug_kernel_api
from sglang.kernels.ops.moe import moe_route_radix
@@ -461,6 +466,12 @@ def moe_fused_gate(
num_warps = 1 if BLOCK_N <= 512 else 4
grid = (triton.cdiv(M, BLOCK_M),)
use_pdl = is_arch_support_pdl()
if use_pdl and scoring_func_int == 1 and N == 384 and K == 6 and M <= 8:
# On SM103, early-launching the small DSV4.1 target router increases
# latency when it overlaps with mHC/shared-expert work. Use ordinary
# stream dependencies; keep PDL for the draft router and larger batches.
arch = get_jit_cuda_arch()
use_pdl = (arch.major, arch.minor) != (10, 3)
extra = {"launch_pdl": True} if use_pdl else {}
# Dynamo cannot analyze the kernel (PDL inline asm), so it writes back every
# pointer arg; aliasing an output as an unused arg's fallback clobbers it.
python/sglang/kernels/ops/speculative/dspark/dspark_accept.py
+4 -1
View File
@@ -676,7 +676,10 @@ def _row_argmax(logits: torch.Tensor, fused: bool = False) -> torch.Tensor:
and logits.dim() == 2
and logits.dtype == torch.float32
and logits.stride(1) == 1
and logits.shape[0] <= 64
and (
logits.shape[0] <= 64
or (logits.shape[0] <= 384 and logits.shape[1] >= 65536)
)
and logits.shape[1] >= 4096
):
from sglang.kernels.ops.speculative.row_argmax import row_argmax
python/sglang/kernels/ops/speculative/row_argmax.py
+53 -1
View File
@@ -51,11 +51,63 @@ def _argmax_final_kernel(INV, INI, OUT, SPLITS: tl.constexpr, BLOCK: tl.constexp
_SPLITS = 64
@triton.jit
def _argmax_pair(av, ai, bv, bi):
# Torch gives NaNs priority and picks the first index for all ties.
an, bn = av != av, bv != bv
take_a = (an & ~bn) | ((an == bn) & ((av > bv) | (((av == bv) | an) & (ai < bi))))
return tl.where(take_a, av, bv), tl.where(take_a, ai, bi)
@triton.jit
def _medium_argmax_partial_kernel(
X,
PV,
PI,
N: tl.constexpr,
SX: tl.constexpr,
SPLITS: tl.constexpr,
BLOCK: tl.constexpr,
):
row, part = tl.program_id(0), tl.program_id(1)
ix = part * BLOCK + tl.arange(0, BLOCK)
valid = ix < N
v = tl.load(X + row * SX + ix, valid, float("-inf"))
i = tl.where(valid, ix, N)
best_v, best_i = tl.reduce((v, i), 0, _argmax_pair)
tl.store(PV + row * SPLITS + part, best_v)
tl.store(PI + row * SPLITS + part, best_i)
@triton.jit
def _medium_argmax_final_kernel(PV, PI, OUT, SPLITS: tl.constexpr, BLOCK: tl.constexpr):
row = tl.program_id(0)
part = tl.arange(0, BLOCK)
v = tl.load(PV + row * SPLITS + part, part < SPLITS, float("-inf"))
i = tl.load(PI + row * SPLITS + part, part < SPLITS, 0x7FFFFFFF)
_, index = tl.reduce((v, i), 0, _argmax_pair)
tl.store(OUT + row, index.to(tl.int64))
def row_argmax(x: torch.Tensor) -> torch.Tensor:
"""``x.argmax(dim=-1)`` for a 2D FP32 tensor with few rows and a wide vocab."""
"""``x.argmax(dim=-1)`` for FP32 speculative logits with a wide vocab."""
assert x.dim() == 2 and x.dtype == torch.float32 and x.stride(1) == 1
rows, n = x.shape
out = torch.empty((rows,), dtype=torch.int64, device=x.device)
if rows > 64:
# Whole aligned tiles avoid the loop and unaligned partition starts.
# Larger batches use wider tiles to limit the total number of CTAs.
block = 4096 if rows <= 256 else 8192
splits = triton.cdiv(n, block)
pv = torch.empty((rows, splits), dtype=torch.float32, device=x.device)
pi = torch.empty((rows, splits), dtype=torch.int32, device=x.device)
_medium_argmax_partial_kernel[(rows, splits)](
x, pv, pi, n, x.stride(0), splits, block, num_warps=4
)
_medium_argmax_final_kernel[(rows,)](
pv, pi, out, splits, triton.next_power_of_2(splits), num_warps=1
)
return out
pv = torch.empty((rows, _SPLITS), dtype=torch.float32, device=x.device)
pi = torch.empty((rows, _SPLITS), dtype=torch.int32, device=x.device)
_argmax_partial_kernel[(rows, _SPLITS)](
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+1 -1
View File
@@ -4294,7 +4294,7 @@ class DeepseekV4AttnBackend(
small_metadata = (
not is_prefill
and seq_lens_casual.is_cuda
and 0 < seq_lens_casual.numel() <= 8
and 0 < seq_lens_casual.numel() <= 384
and out_loc.numel() == seq_lens_casual.numel()
and self.low_ratios == (1, 2)
and set(self.present_ratios) == {1, 2}
python/sglang/srt/layers/moe/mhc_post_fusion.py
+9
View File
@@ -20,10 +20,19 @@ class MhcPostFusion:
norm_eps: float = 0.0
normalized: Optional[torch.Tensor] = None
quantized: Optional[tuple[torch.Tensor, torch.Tensor]] = None
combine_only: bool = False
combined: Optional[torch.Tensor] = None
overlap_only: bool = False
record_stats: Optional[
Callable[[], tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
] = None
def start_stats_before_all_reduce(self):
if self.overlap_only and self.record_stats is not None:
assert self.stats_stream is not None
self.stats_stream.wait_stream(torch.cuda.current_stream())
self.materialize_stats()
def materialize_stats(self):
# Record after the main parent; graph replay must keep the join on the
# main stream.
python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py
+37 -8
View File
@@ -416,7 +416,10 @@ class Mxfp4FlashinferTrtllmMoEMethod:
# triple instead of the finalized [T, hidden] tensor.
defer_finalize = is_deferred_finalize_enabled()
symm_output = None
if not defer_finalize:
# Preserve the ordinary output shape in the medium-batch autotuner
# cache key. The deferred ABI ignores this allocation and returns the
# expanded GEMM output for the separate fused finalize epilogue.
if not defer_finalize or 96 < num_tokens <= 384:
with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
@@ -519,11 +522,17 @@ def maybe_fuse_routed_scale_and_shared_add(
# Fused finalize + shared add + TP all-reduce
_fused_finalize_all_reduce_world_size: Optional[int] = None
_fused_finalize_all_reduce_probed = False
_fused_finalize_all_reduce_comm = None
def _fused_finalize_all_reduce_comm_world_size() -> Optional[int]:
"""Reserve a separate push plane for up to 384 rows of fused MoE output."""
global _fused_finalize_all_reduce_world_size, _fused_finalize_all_reduce_probed
global _fused_finalize_all_reduce_comm
if not _fused_finalize_all_reduce_probed:
# The eager warmup must initialize peer workspaces before capture.
if torch.cuda.is_current_stream_capturing():
return None
_fused_finalize_all_reduce_probed = True
from sglang.kernels.ops.communication import all_reduce_fusion
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
@@ -532,8 +541,24 @@ def _fused_finalize_all_reduce_comm_world_size() -> Optional[int]:
ca_comm = get_parallel().tp_group.ca_comm
if isinstance(ca_comm, CustomAllReduceV2) and not ca_comm.disabled:
all_reduce_fusion.register_comm(ca_comm.obj)
_fused_finalize_all_reduce_world_size = ca_comm.world_size
fused_comm = ca_comm
if ca_comm.world_size == 4:
from sglang.kernels.ops.communication.mp import register_comm_cleanup
fused_comm = CustomAllReduceV2(
ca_comm.group,
ca_comm.device,
max_pull_size=0,
max_pull_blocks=0,
max_push_size=4 * 1024 * 1024,
max_push_blocks=512,
)
register_comm_cleanup(fused_comm)
if fused_comm.disabled:
return None
_fused_finalize_all_reduce_comm = fused_comm
all_reduce_fusion.register_comm(fused_comm.obj)
_fused_finalize_all_reduce_world_size = fused_comm.world_size
else:
log_info_on_rank0(
logger,
@@ -556,6 +581,11 @@ def should_use_fuse_finalize_all_reduce(
return False
if num_tokens <= 0:
return False
if num_tokens > 96:
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
if is_batch_invariant_mode_enabled():
return False
from sglang.kernels.ops.communication import all_reduce_fusion
if not all_reduce_fusion.valid_cluster_sizes(hidden_dim):
@@ -563,9 +593,8 @@ def should_use_fuse_finalize_all_reduce(
tp_group = get_parallel().tp_group
if _fused_finalize_all_reduce_comm_world_size() != tp_group.world_size:
return False
# one push phase counter per row (the plane has num_sm of them)
if num_tokens > tp_group.ca_comm.config.num_push_blocks:
comm = _fused_finalize_all_reduce_comm
# Each token row owns a push phase counter on the epilogue's plane.
if num_tokens > comm.config.num_push_blocks:
return False
return all_reduce_fusion.fits_push_slot(
tp_group.ca_comm.max_push_size, num_tokens, hidden_dim
)
return all_reduce_fusion.fits_push_slot(comm.max_push_size, num_tokens, hidden_dim)
python/sglang/srt/models/deepseek_v2.py
+32 -4
View File
@@ -552,9 +552,9 @@ class MoEGate(nn.Module):
return logits
# 96 rows of 5120 bf16 fit the 1 MiB CustomAllReduceV2 push slot the whole
# [T, hidden] view is staged through.
_FUSED_FINALIZE_ALL_REDUCE_MAX_TOKENS = 96
# The dedicated 4 MiB push slot fits 384 rows of 5120 BF16 values.
# The dispatch gate also checks slot capacity and available counters.
_FUSED_FINALIZE_ALL_REDUCE_MAX_TOKENS = 384
class DeepseekV2MoE(nn.Module):
@@ -1129,7 +1129,17 @@ class DeepseekV2MoE(nn.Module):
mhc.post,
mhc.comb,
)
if mhc.norm_weight is not None:
if mhc.combine_only:
from sglang.kernels.ops.communication.all_reduce_mhc_combine import (
moe_finalize_all_reduce_mhc_combine,
)
final_hidden_states, mhc.output, mhc.combined = (
moe_finalize_all_reduce_mhc_combine(
*args, mhc.pre, world_size=self.tp_size
)
)
elif mhc.norm_weight is not None:
from sglang.kernels.ops.communication.all_reduce_mhc import (
moe_finalize_all_reduce_mhc_quant,
)
@@ -1175,6 +1185,16 @@ class DeepseekV2MoE(nn.Module):
)
if not all_reduce_done:
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True
):
from sglang.srt.layers.moe.mhc_post_fusion import (
current_mhc_post_fusion,
)
mhc = current_mhc_post_fusion()
if mhc is not None:
mhc.start_stats_before_all_reduce()
final_hidden_states = post_experts_all_reduce(final_hidden_states)
# TP1 shared experts are replicated, so add them after all-reduce to
# avoid summing the same shared output once per TP rank.
@@ -1321,6 +1341,14 @@ class DeepseekV2MoE(nn.Module):
self.routed_scaling_factor,
)
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True,
):
from sglang.srt.layers.moe.mhc_post_fusion import current_mhc_post_fusion
mhc = current_mhc_post_fusion()
if mhc is not None:
mhc.start_stats_before_all_reduce()
final_hidden_states = post_experts_all_reduce(final_hidden_states)
# TP1 shared experts are replicated, so add them after all-reduce to
# avoid summing the same shared output once per TP rank.
python/sglang/srt/models/deepseek_v4.py
+260 -37
View File
@@ -2525,7 +2525,10 @@ class MQALayer(MqaAttentionBase):
o if isinstance(o, Mxfp8SwizzledInput) else o.flatten(1),
skip_all_reduce=mhc is not None,
)
if mhc is not None:
if mhc is not None and mhc.overlap_only:
mhc.start_stats_before_all_reduce()
o = attn_tp_all_reduce(o)
elif mhc is not None:
from sglang.kernels.ops.communication.all_reduce_mhc import (
all_reduce_mhc_norm,
)
@@ -2533,16 +2536,30 @@ class MQALayer(MqaAttentionBase):
mhc.materialize_stats()
if mhc.stats_stream is not None:
torch.cuda.current_stream().wait_stream(mhc.stats_stream)
o, mhc.output, mhc.normalized = all_reduce_mhc_norm(
o,
mhc.residual,
mhc.post,
mhc.comb,
mhc.pre,
mhc.norm_weight,
mhc.norm_eps,
world_size=self.attn_tp_size,
)
if mhc.combine_only:
from sglang.kernels.ops.communication.all_reduce_mhc_combine import (
all_reduce_mhc_combine,
)
o, mhc.output, mhc.combined = all_reduce_mhc_combine(
o,
mhc.residual,
mhc.post,
mhc.comb,
mhc.pre,
world_size=self.attn_tp_size,
)
else:
o, mhc.output, mhc.normalized = all_reduce_mhc_norm(
o,
mhc.residual,
mhc.post,
mhc.comb,
mhc.pre,
mhc.norm_weight,
mhc.norm_eps,
world_size=self.attn_tp_size,
)
if self.attn_tp_size > 1 and self.attn_tp_size < get_parallel().tp_size:
o = attn_tp_all_reduce(o)
@@ -3180,6 +3197,7 @@ class DeepseekV4DecoderLayer(nn.Module):
quantized: Optional[list] = None,
normalized: Optional[torch.Tensor] = None,
precomputed: Optional[tuple] = None,
combined: Optional[torch.Tensor] = None,
) -> torch.Tensor:
from sglang.kernels.ops.layernorm.mhc import hc_combine
@@ -3195,8 +3213,23 @@ class DeepseekV4DecoderLayer(nn.Module):
quantized.append(precomputed[1])
return precomputed[0]
if normalized is not None:
assert not quantize
# Prefill projections still quantize the BF16 input themselves;
# the optional fused-quantization list stays empty for this case.
assert not quantize or 4096 <= x.shape[0] <= 65536
return normalized
if combined is not None:
if (
4096 <= combined.shape[0] <= 65536
and norm.weight.dtype == torch.bfloat16
and not norm.cast_x_before_out_mul
and norm.variance_size_override is None
):
from sglang.kernels.ops.layernorm.mhc_post_combine import (
hc_norm_prefill,
)
return hc_norm_prefill(combined, norm.weight, norm.variance_epsilon)
return norm(combined)
if apply_pre is None:
return norm(x[:, 0, :].contiguous())
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
@@ -3378,7 +3411,24 @@ class DeepseekV4DecoderLayer(nn.Module):
)
def _get_hc_stats_stream(self, hidden_states, forward_batch):
# Every branch joins this stream before hc_post reads the coefficients.
# Prefill stats share one model-wide stream. Start them immediately
# before the sublayer's all-reduce, after its compute has completed.
if (
self.config.model_type == "deepseek_v41"
and hidden_states.is_cuda
and get_platform().is_blackwell
and forward_batch.forward_mode.is_extend_without_speculative()
and 4096 <= hidden_states.shape[0] <= 65536
and get_parallel().attn_dp_size == 1
and not get_forward().sp_active
and not self.dsa_enable_prefill_cp
):
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
if not is_batch_invariant_mode_enabled():
return self.hc_stats_stream
# Verify batches can also compute coefficients beside the
# sublayer; each branch joins before hc_post reads those coefficients.
return (
self.hc_stats_stream
if (
@@ -3392,6 +3442,65 @@ class DeepseekV4DecoderLayer(nn.Module):
else None
)
def _hc_post_with_combine(
self, x, residual, post, comb, pre, forward_batch, norm=None
):
"""Return updated HC streams and optional combined/normalized inputs."""
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
if (
self.config.model_type == "deepseek_v41"
and x.is_cuda
and get_platform().is_blackwell
and (
(
128 <= x.shape[0] <= 384
and (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
)
)
or (
4096 <= x.shape[0] <= 65536
and forward_batch.forward_mode.is_extend_without_speculative()
and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get()
and not envs.SGLANG_OPT_USE_FLASHINFER_MHC.get()
)
)
and x.shape[1] == 5120
and self.hc_mult == 4
and x.dtype == residual.dtype == torch.bfloat16
and post.dtype == comb.dtype == pre.dtype == torch.float32
and all(t.is_contiguous() for t in (x, residual, post, comb, pre))
and get_parallel().attn_dp_size == 1
and not get_forward().sp_active
and not self.dsa_enable_prefill_cp
and not is_batch_invariant_mode_enabled()
):
if (
x.shape[0] >= 4096
and norm is not None
and not norm.cast_x_before_out_mul
and norm.variance_size_override is None
and norm.weight.dtype == torch.bfloat16
and norm.weight.shape == (5120,)
and norm.weight.is_contiguous()
and all(t.data_ptr() % 16 == 0 for t in (x, residual, norm.weight))
):
from sglang.kernels.ops.layernorm.mhc_post_combine_norm_prefill import (
mhc_post_combine_norm_prefill,
)
updated, normalized = mhc_post_combine_norm_prefill(
x, residual, post, comb, pre, norm.weight, norm.variance_epsilon
)
return updated, None, normalized
from sglang.kernels.ops.layernorm.mhc_post_combine import mhc_post_combine
updated, combined = mhc_post_combine(x, residual, post, comb, pre)
return updated, combined, None
return self.hc_post(x, residual, post, comb), None, None
def forward_hc_pre_from_prev(
self,
positions: torch.Tensor,
@@ -3403,11 +3512,20 @@ class DeepseekV4DecoderLayer(nn.Module):
precomputed_attn: Optional[tuple] = None,
next_norm: Optional[RMSNorm] = None,
next_input: Optional[list] = None,
combined_attn: Optional[torch.Tensor] = None,
normalized_attn: Optional[torch.Tensor] = None,
next_combined: Optional[list] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Layer forward where each sublayer consumes the previous sublayer's
pre-mix. Returns (hidden_states, ffn_pre)."""
from functools import partial
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
from sglang.srt.layers.moe.mhc_post_fusion import (
MhcPostFusion,
use_mhc_post_fusion,
)
stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch)
residual = hidden_states
attn_quantized: Optional[list] = (
@@ -3428,13 +3546,23 @@ class DeepseekV4DecoderLayer(nn.Module):
stats_stream=stats_stream,
quantized=attn_quantized,
precomputed=precomputed_attn,
combined=combined_attn,
normalized=normalized_attn,
)
prefill_overlap = (
stats_stream is not None
and forward_batch.forward_mode.is_extend_without_speculative()
)
medium_verify = 128 <= x.shape[0] <= 384 and (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
)
attn_mhc = None
if (
self.config.model_type == "deepseek_v41"
and x.is_cuda
and get_platform().is_blackwell
and 0 < x.shape[0] <= 8
and (0 < x.shape[0] <= 8 or medium_verify)
and x.shape[1] == 5120
and self.hc_mult == 4
and x.dtype == residual.dtype == torch.bfloat16
@@ -3451,16 +3579,18 @@ class DeepseekV4DecoderLayer(nn.Module):
from sglang.kernels.ops.communication.all_reduce_fusion import (
get_registered_comm,
)
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
from sglang.srt.layers.moe.mhc_post_fusion import (
MhcPostFusion,
use_mhc_post_fusion,
)
if (
not is_batch_invariant_mode_enabled()
and get_registered_comm(self.self_attn.attn_tp_size) is not None
):
comm_ready = get_registered_comm(self.self_attn.attn_tp_size) is not None
if medium_verify and not is_batch_invariant_mode_enabled():
from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import (
_fused_finalize_all_reduce_comm_world_size,
)
comm_ready = (
_fused_finalize_all_reduce_comm_world_size()
== self.self_attn.attn_tp_size
)
if not is_batch_invariant_mode_enabled() and comm_ready:
attn_mhc = MhcPostFusion(
residual,
None,
@@ -3469,7 +3599,21 @@ class DeepseekV4DecoderLayer(nn.Module):
record_stats=attn_stats,
norm_weight=self.post_attention_layernorm.weight,
norm_eps=self.post_attention_layernorm.variance_epsilon,
combine_only=medium_verify,
)
if (
prefill_overlap
and get_parallel().tp_size == self.self_attn.attn_tp_size == 4
and self.self_attn.wo_b.reduce_results
):
attn_mhc = MhcPostFusion(
residual,
None,
None,
stats_stream,
overlap_only=True,
record_stats=attn_stats,
)
context = (
use_mhc_post_fusion(attn_mhc) if attn_mhc is not None else nullcontext()
)
@@ -3480,15 +3624,32 @@ class DeepseekV4DecoderLayer(nn.Module):
forward_batch=forward_batch,
x_quant=attn_quantized[0] if attn_quantized else None,
)
ffn_combined = None
ffn_normalized = None
if attn_mhc is not None:
attn_mhc.materialize_stats()
if attn_mhc is not None and attn_mhc.output is not None:
attn_pre = attn_mhc.pre
hidden_states = attn_mhc.output
ffn_combined = attn_mhc.combined
ffn_normalized = attn_mhc.normalized
else:
attn_pre, attn_post, attn_comb = attn_stats()
attn_pre, attn_post, attn_comb = (
(attn_mhc.pre, attn_mhc.post, attn_mhc.comb)
if attn_mhc is not None
else attn_stats()
)
if stats_stream is not None:
torch.cuda.current_stream().wait_stream(stats_stream)
hidden_states = self.hc_post(x, residual, attn_post, attn_comb)
hidden_states, ffn_combined, ffn_normalized = self._hc_post_with_combine(
x,
residual,
attn_post,
attn_comb,
attn_pre,
forward_batch,
norm=self.post_attention_layernorm,
)
residual = hidden_states
ffn_stats = partial(
@@ -3504,38 +3665,52 @@ class DeepseekV4DecoderLayer(nn.Module):
apply_pre=attn_pre,
norm=self.post_attention_layernorm,
stats_stream=stats_stream,
normalized=attn_mhc.normalized if attn_mhc is not None else None,
normalized=ffn_normalized,
combined=ffn_combined,
)
mhc = None
if (
self.config.model_type == "deepseek_v41"
and x.is_cuda
and get_platform().is_blackwell
and 0 < x.shape[0] <= 8
and (0 < x.shape[0] <= 8 or (medium_verify and next_combined is not None))
and x.shape[1] == 5120
and self.hc_mult == 4
and x.dtype == residual.dtype == torch.bfloat16
and residual.is_contiguous()
and get_parallel().attn_dp_size == 1
and get_moe_a2a_backend().is_none()
and not get_forward().sp_active
and not self.dsa_enable_prefill_cp
and not self.mlp._shared_expert_tp1
and self.mlp.tp_size == 4
and (not medium_verify or not is_batch_invariant_mode_enabled())
):
from sglang.srt.layers.moe.mhc_post_fusion import (
MhcPostFusion,
use_mhc_post_fusion,
)
mhc = MhcPostFusion(
residual, None, None, stats_stream, record_stats=ffn_stats
residual,
None,
None,
stats_stream,
record_stats=ffn_stats,
combine_only=medium_verify,
)
if next_norm is not None:
mhc.norm_weight = next_norm.weight
mhc.norm_eps = next_norm.variance_epsilon
context = use_mhc_post_fusion(mhc)
else:
context = nullcontext()
if (
prefill_overlap
and self.mlp.tp_size == 4
and get_moe_a2a_backend().is_none()
):
mhc = MhcPostFusion(
residual,
None,
None,
stats_stream,
overlap_only=True,
record_stats=ffn_stats,
)
context = use_mhc_post_fusion(mhc) if mhc is not None else nullcontext()
with context:
x = self._run_moe_ffn_dp_sync(
x, forward_batch, input_ids=input_ids, input_ids_global=input_ids_global
@@ -3549,10 +3724,25 @@ class DeepseekV4DecoderLayer(nn.Module):
hidden_states = mhc.output
if next_input is not None and mhc.quantized is not None:
next_input.append((mhc.normalized, Mxfp8SwizzledInput(*mhc.quantized)))
if next_combined is not None and mhc.combined is not None:
next_combined.append((mhc.combined, None))
else:
if stats_stream is not None:
torch.cuda.current_stream().wait_stream(stats_stream)
hidden_states = self.hc_post(x, residual, ffn_post, ffn_comb)
if next_combined is not None:
hidden_states, combined, normalized = self._hc_post_with_combine(
x,
residual,
ffn_post,
ffn_comb,
ffn_pre,
forward_batch,
norm=next_norm,
)
if combined is not None or normalized is not None:
next_combined.append((combined, normalized))
else:
hidden_states = self.hc_post(x, residual, ffn_post, ffn_comb)
return hidden_states, ffn_pre
def _run_moe_ffn_dp_sync(
@@ -4243,8 +4433,12 @@ class DeepseekV4Model(nn.Module):
saved_full = None
prev_pre = None
precomputed_attn = None
combined_attn = None
normalized_attn = None
for i in range(self.start_layer, self.end_layer):
if tail is not None and i == self.late_layer_start:
combined_attn = None
normalized_attn = None
# Decode reaches back at most SWA_WINDOW positions.
saved_full = attn_backend.enter_late_layer_tail(forward_batch)
hidden_states, prev_pre, input_ids, input_ids_global = (
@@ -4259,6 +4453,8 @@ class DeepseekV4Model(nn.Module):
engram = self.layers[i].engram
if engram is not None:
precomputed_attn = None
combined_attn = None
normalized_attn = None
before_engram = hidden_states
hidden_states = engram(
hidden_states,
@@ -4288,6 +4484,27 @@ class DeepseekV4Model(nn.Module):
)
next_norm = None
next_input = []
# The next layer can consume a collapsed input only if no Engram
# or row selection changes the residual between the two layers.
next_combined = (
[]
if (
self.config.model_type == "deepseek_v41"
and (
128 <= hidden_states.shape[0] <= 384
or (
4096 <= hidden_states.shape[0] <= 65536
and forward_batch.forward_mode.is_extend_without_speculative()
)
)
and i + 1 < self.end_layer
and tail is None
and self.layers[i + 1].engram is None
)
else None
)
if next_combined is not None and hidden_states.shape[0] >= 4096:
next_norm = self.layers[i + 1].input_layernorm
if (
self.config.model_type == "deepseek_v41"
and i + 1 < self.end_layer
@@ -4328,8 +4545,14 @@ class DeepseekV4Model(nn.Module):
precomputed_attn=precomputed_attn,
next_norm=next_norm,
next_input=next_input,
combined_attn=combined_attn,
normalized_attn=normalized_attn,
next_combined=next_combined,
)
precomputed_attn = next_input[0] if next_input else None
combined_attn, normalized_attn = (
next_combined[0] if next_combined else (None, None)
)
if saved_full is not None:
attn_backend.exit_late_layer_tail(saved_full, forward_batch)
return hidden_states, prev_pre, tail