diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh index e0a8a9577..ad17a1325 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh @@ -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 SGL_DEVICE device::AlignedVector finalize_vec(const FusionParams& params, uint32_t vid) { using namespace device; + static_assert(std::is_same_v || std::is_same_v, "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 finalize_vec(const FusionParams& p for (uint32_t j = 0; j < kIdxVecs; ++j) { idx[j].load(params.fin_idx, token * kIdxVecs + j); } - AlignedVector weight[kWVecs]; + AlignedVector 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 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& 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 +// 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 __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; + constexpr bool kFinalize = !std::is_void_v; 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(params, vid); } else { ptx::ld_global_16B(vec, params.input, vid); } @@ -796,7 +801,9 @@ struct AllReduceFusionKernel { SymbolicDevice device; device.set_options(); TensorMatcher({P, kNormDim}).with_dtype().with_device(device).verify(gemm2_out); - TensorMatcher({T, K}).with_dtype().with_device(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().with_device(device).verify(expert_weights); TensorMatcher({TK}).with_dtype().with_device(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(gemm2_out.data_ptr()); params.fin_idx = static_cast(permuted_idx.data_ptr()); params.fin_weights = static_cast(expert_weights.data_ptr()); + // stage 1 reads the routing row with one aligned vector load per kMaxVecBytes + CHECK_HOST(reinterpret_cast(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(std::min(num_tokens, kMaxClusters), 1); CHECK_HOST(num_row_clusters < push.num_blocks); + const auto kernel = is_type(expert_weights.dtype()) + ? all_reduce_push_norm_cluster_kernel + : all_reduce_push_norm_cluster_kernel; host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, out.device()) - .enable_pdl(kUsePDL)( - all_reduce_push_norm_cluster_kernel, params); + .enable_pdl(kUsePDL)(kernel, params); } /// Low-SM NVLS pull (+ optional residual): in-place reduce-scatter + diff --git a/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu b/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu index b94c311ae..0e0ea7e71 100644 --- a/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu +++ b/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu @@ -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`` diff --git a/python/sglang/kernels/jit/include/sgl_kernel/math.cuh b/python/sglang/kernels/jit/include/sgl_kernel/math.cuh index f92e3d8b2..92f4a40c7 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/math.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/math.cuh @@ -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(a), b, acc); +} + } // namespace device::math } // namespace sglang diff --git a/python/sglang/kernels/ops/kimi_k3/all_reduce.py b/python/sglang/kernels/ops/kimi_k3/all_reduce.py index f50f95cff..9cff2cd5c 100644 --- a/python/sglang/kernels/ops/kimi_k3/all_reduce.py +++ b/python/sglang/kernels/ops/kimi_k3/all_reduce.py @@ -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, diff --git a/python/sglang/srt/layers/k3_ar_fusion.py b/python/sglang/srt/layers/k3_ar_fusion.py index d5e0a116e..b91774a5c 100644 --- a/python/sglang/srt/layers/k3_ar_fusion.py +++ b/python/sglang/srt/layers/k3_ar_fusion.py @@ -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, diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 9770af47b..b0ede48c8 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -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" diff --git a/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py index 35c239a30..4b04e87b4 100644 --- a/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py +++ b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py @@ -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(