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

[Kimi-K3] Accept fp32 routing weights in the fused MoE finalize (#38612)

Browse Source 
Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
This commit is contained in:
Mohammad Miadh Angkad
2026-09-09 01:33:39 -07:00
committed by GitHub
co-authored by Mohammad Angkad
parent d10ebdd0cb
commit 72d5c5bb73
7 changed files with 55 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh
+26 -14
View File
@@ -78,7 +78,7 @@ struct FusionParams {
// deferred finalize path); `input` is then output-only ([T, kNormDim])
const uint8_t* fin_gemm2; // [P, kNormDim] bf16, permuted rows
const uint8_t* fin_idx; // [T * kFinTopK] int32, -1 = dropped slot
const uint8_t* fin_weights; // [T, kFinTopK] bf16
const uint8_t* fin_weights; // [T, kFinTopK] bf16 or fp32
};
// The *_norm variants view the input as rows of the K3 latent width (3584
@@ -161,13 +161,16 @@ constexpr uint32_t kFinTopK = 16;
// One 16B vector of the deferred MoE finalize (latent width fixed to kNormDim):
// local[t] = sum_k fin_weights[t, k] * fin_gemm2[fin_idx[t*16 + k]]
// All 16 gathers issue before the FMA chain; threads of the same token
// broadcast-load the same routing rows.
// broadcast-load the same routing rows. `W` is the routing-weight dtype
// (bf16 or fp32; see finalize_push_norm).
template <typename W>
SGL_DEVICE device::AlignedVector<bf16x2_t, 4> finalize_vec(const FusionParams& params, uint32_t vid) {
using namespace device;
static_assert(std::is_same_v<W, bf16_t> || std::is_same_v<W, fp32_t>, "unsupported routing-weight dtype");
constexpr uint32_t kIdxVecSize = kMaxVecBytes / sizeof(int32_t);
constexpr uint32_t kWVecSize = kMaxVecBytes / sizeof(bf16_t);
constexpr uint32_t kWVecSize = kMaxVecBytes / sizeof(W);
constexpr uint32_t kIdxVecs = kFinTopK / kIdxVecSize; // 2 on SM100+
constexpr uint32_t kWVecs = kFinTopK / kWVecSize; // 1 on SM100+
constexpr uint32_t kWVecs = kFinTopK / kWVecSize; // 1 (bf16) / 2 (fp32) on SM100+
const uint32_t token = vid / kNormRowVecs;
const uint32_t hvec = vid % kNormRowVecs;
@@ -177,7 +180,7 @@ SGL_DEVICE device::AlignedVector<bf16x2_t, 4> finalize_vec(const FusionParams& p
for (uint32_t j = 0; j < kIdxVecs; ++j) {
idx[j].load(params.fin_idx, token * kIdxVecs + j);
}
AlignedVector<bf16_t, kWVecSize> weight[kWVecs];
AlignedVector<W, kWVecSize> weight[kWVecs];
#pragma unroll
for (uint32_t j = 0; j < kWVecs; ++j) {
weight[j].load(params.fin_weights, token * kWVecs + j);
@@ -197,7 +200,7 @@ SGL_DEVICE device::AlignedVector<bf16x2_t, 4> finalize_vec(const FusionParams& p
for (uint32_t k = 0; k < kFinTopK; ++k) {
const int32_t row = idx[k / kIdxVecSize][k % kIdxVecSize];
if (row < 0) continue;
const bf16_t w_k = weight[k / kWVecSize][k % kWVecSize];
const W w_k = weight[k / kWVecSize][k % kWVecSize];
#pragma unroll
for (uint32_t i = 0; i < 8; ++i) {
acc[i] = device::math::fma_f32_bf16(in[k][i], w_k, acc[i]);
@@ -234,15 +237,17 @@ SGL_DEVICE float reduce_sqr(device::AlignedVector<T2, N>& out_vec, device::Align
return sum_eq;
}
// kFinalize: stage 1 computes the deferred MoE finalize per vector instead of
// reading a staged input tensor; `input` is then output-only. The host sets
// num_norm_rows to the full row count (every reduced row is normed).
template <uint32_t kWorldSize, uint32_t kClusterSize, bool kUsePDL, bool kFinalize = false>
// FinWeight != void selects the finalize variant: stage 1 computes the
// deferred MoE finalize per vector (with FinWeight-typed routing weights)
// instead of reading a staged input tensor; `input` is then output-only. The
// host sets num_norm_rows to the full row count (every reduced row is normed).
template <uint32_t kWorldSize, uint32_t kClusterSize, bool kUsePDL, typename FinWeight = void>
__global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClusterSize, 1, 1) //
void all_reduce_push_norm_cluster_kernel(const __grid_constant__ FusionParams params) {
namespace cg = cooperative_groups;
using namespace device;
using vec_t = AlignedVector<bf16x2_t, 4>;
constexpr bool kFinalize = !std::is_void_v<FinWeight>;
constexpr uint32_t kBlockSize = kNormRowVecs / kClusterSize;
constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
@@ -289,7 +294,7 @@ __global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClus
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec;
if constexpr (kFinalize) {
vec = finalize_vec(params, vid);
vec = finalize_vec<FinWeight>(params, vid);
} else {
ptx::ld_global_16B(vec, params.input, vid);
}
@@ -796,7 +801,9 @@ struct AllReduceFusionKernel {
SymbolicDevice device;
device.set_options<kDLCUDA>();
TensorMatcher({P, kNormDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(gemm2_out);
TensorMatcher({T, K}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(expert_weights);
// trtllm-gen returns the routing weights at the dtype the routing form
// carries -- fp32 for unpacked, bf16 for packed -- consumed as given
TensorMatcher({T, K}).with_dtype<fp32_t, bf16_t>().with_device<kDLCUDA>(device).verify(expert_weights);
TensorMatcher({TK}).with_dtype<int32_t>().with_device<kDLCUDA>(device).verify(permuted_idx);
CHECK_HOST(K.unwrap() == kFinTopK) << "finalize_push_norm is specialized for top_k = " << kFinTopK;
@@ -805,13 +812,18 @@ struct AllReduceFusionKernel {
params.fin_gemm2 = static_cast<const uint8_t*>(gemm2_out.data_ptr());
params.fin_idx = static_cast<const uint8_t*>(permuted_idx.data_ptr());
params.fin_weights = static_cast<const uint8_t*>(expert_weights.data_ptr());
// stage 1 reads the routing row with one aligned vector load per kMaxVecBytes
CHECK_HOST(reinterpret_cast<uintptr_t>(params.fin_weights) % device::kMaxVecBytes == 0)
<< "expert_weights must be " << device::kMaxVecBytes << "B aligned";
constexpr uint32_t kMaxClusters = 96;
const auto num_row_clusters = std::max<uint32_t>(std::min(num_tokens, kMaxClusters), 1);
CHECK_HOST(num_row_clusters < push.num_blocks);
const auto kernel = is_type<fp32_t>(expert_weights.dtype())
? all_reduce_push_norm_cluster_kernel<kWorldSize, kClusterSize, kUsePDL, fp32_t>
: all_reduce_push_norm_cluster_kernel<kWorldSize, kClusterSize, kUsePDL, bf16_t>;
host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, out.device())
.enable_pdl(kUsePDL)(
all_reduce_push_norm_cluster_kernel<kWorldSize, kClusterSize, kUsePDL, /*kFinalize=*/true>, params);
.enable_pdl(kUsePDL)(kernel, params);
}
/// Low-SM NVLS pull (+ optional residual): in-place reduce-scatter +
python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu
+2 -3
View File
@@ -36,9 +36,8 @@
* a clean PDL handoff.
*
* Expert-weight dtype is templated on ``TypeExpW`` so we accept both bf16
* and fp32 topk weights. The trtllm deferred-finalize path always feeds bf16
* (the trtllm-gen routing kernel emits bf16 for every routing method); fp32
* is kept for callers that produce topk weights in fp32.
* and fp32 topk weights: the trtllm deferred finalize returns bf16 for packed
* routing and the caller's own fp32 weights for unpacked routing.
*
* Expert-weight scale convention: in our target backends
* (flashinfer trtllm nvfp4 + unquantized), ``apply_routed_scaling_factor_on_output``
python/sglang/kernels/jit/include/sgl_kernel/math.cuh
+6
View File
@@ -106,6 +106,12 @@ SGL_DEVICE float fma_f32_bf16(bf16_t a, bf16_t b, float acc) {
#endif
}
// bf16 x fp32 -> fp32 fused multiply-add: same one-rounding contract as the
// overload above (the bf16 -> f32 convert is exact).
SGL_DEVICE float fma_f32_bf16(bf16_t a, float b, float acc) {
return fmaf(cast<fp32_t>(a), b, acc);
}
} // namespace device::math
} // namespace sglang
python/sglang/kernels/ops/kimi_k3/all_reduce.py
+2 -1
View File
@@ -281,7 +281,8 @@ def finalize_all_reduce_push_norm(
(``sum_k expert_weights[t, k] * gemm2_out[idx[t*16 + k]]``, -1 slots
skipped) is computed during the multicast staging pass from the
trtllm-gen deferred-finalize triple (``do_finalize=False``) and never
materializes in global memory. top_k is fixed to 16 (K3)."""
materializes in global memory. ``expert_weights`` is fp32 or bf16,
consumed as given; top_k is fixed to 16 (K3)."""
_finalize_push_norm_op(
world_size,
out,
python/sglang/srt/layers/k3_ar_fusion.py
-2
View File
@@ -371,8 +371,6 @@ def finalize_all_reduce_push_norm(
state = _get_state()
assert state is not None
if expert_weights.dtype != torch.bfloat16:
expert_weights = expert_weights.to(torch.bfloat16)
return mod.finalize_all_reduce_push_norm(
state.world_size,
out,
python/sglang/srt/models/kimi_k3.py
+4 -3
View File
@@ -552,9 +552,10 @@ class KimiK3MoE(nn.Module):
# Defer the trtllm-gen finalize (top-k weighted unpermute) out of the
# MoE op and fuse it into the push all-reduce's staging pass
# (k3_ar_fusion.finalize_all_reduce_push_norm): the rank-local latent
# never materializes. Only the situ packed-routing trtllm-gen path
# serves the deferral; sizes beyond the push window fall back to the
# in-op finalize at runtime (finalize_push_fits).
# never materializes. Only the situ trtllm-gen path serves the
# deferral, on either routing form (packed ids on a fused route+quant
# hit, unpacked fp32 weights otherwise); sizes beyond the push window
# fall back to the in-op finalize at runtime (finalize_push_fits).
self._defer_moe_finalize = (
get_moe_runner_backend().is_flashinfer_mxfp4()
and config.hidden_act == "situ"
test/registered/kernels/ops/kimi_k3/test_ar_fusion.py
+15 -6
View File
@@ -248,10 +248,13 @@ def test_ar_fusion_push_norm(num_tokens: int, rows_per_token: int):
FIN_TOPK = 16
def _build_permuted_layout(num_tokens: int, seed: int):
def _build_permuted_layout(
num_tokens: int, seed: int, w_dtype: torch.dtype = torch.float32
):
"""trtllm-gen permuted gemm2 layout (rows grouped by expert, per-expert
tile padding). Deterministic on CPU: idx/weights are identical on every
rank (TP semantics — same routing), gemm2 values are per-rank."""
rank (TP semantics — same routing), gemm2 values are per-rank.
``w_dtype`` is the routing-weight dtype the deferred finalize hands back."""
num_experts, tile = 896, 8
gen = torch.Generator(device="cpu").manual_seed(seed)
topk_ids = torch.stack(
@@ -268,7 +271,7 @@ def _build_permuted_layout(num_tokens: int, seed: int):
for i, e in enumerate(topk_ids.flatten().tolist()):
idx[i] = bases[e] + fill[e]
fill[e] += 1
weights = torch.rand(num_tokens, FIN_TOPK, generator=gen).to(torch.bfloat16)
weights = torch.rand(num_tokens, FIN_TOPK, generator=gen).to(w_dtype)
num_rows = int(padded.sum())
g = torch.Generator(device="cpu").manual_seed(seed * 31 + dist.get_rank())
gemm2 = (torch.randn(num_rows, NORM_DIM, generator=g) * 2).to(torch.bfloat16)
@@ -296,13 +299,14 @@ def _finalize_norm_ref(gemm2, idx, weights, norm_w, eps: float) -> torch.Tensor:
return (total * factor * norm_w.float()).to(torch.bfloat16)
@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("bs", PUSH_BS)
@torch.inference_mode()
def test_ar_fusion_finalize_push_norm(bs: int):
def test_ar_fusion_finalize_push_norm(bs: int, w_dtype: torch.dtype):
comm = _init_comm()
world = comm.world_size
eps = 1e-6
gemm2, idx, weights = _build_permuted_layout(bs, seed=bs + 23)
gemm2, idx, weights = _build_permuted_layout(bs, seed=bs + 23, w_dtype=w_dtype)
g = torch.Generator(device="cpu").manual_seed(77)
norm_w = (torch.rand(NORM_DIM, generator=g) + 0.5).to(torch.bfloat16).to(_device())
ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps)
@@ -325,7 +329,12 @@ def test_ar_fusion_finalize_push_norm_stress():
norm_w = (torch.rand(NORM_DIM, generator=g) + 0.5).to(torch.bfloat16).to(_device())
for it in range(12):
bs = (1, 8, 32)[it % 3]
gemm2, idx, weights = _build_permuted_layout(bs, seed=9000 + it)
# alternate the routing-weight precision: both kernel instantiations
# share the one push workspace
w_dtype = (torch.float32, torch.bfloat16)[it % 2]
gemm2, idx, weights = _build_permuted_layout(
bs, seed=9000 + it, w_dtype=w_dtype
)
ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps)
out = torch.empty(bs, NORM_DIM, dtype=torch.bfloat16, device=_device())
all_reduce.finalize_all_reduce_push_norm(