diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/mhc_post_combine_norm_prefill.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/mhc_post_combine_norm_prefill.cuh new file mode 100644 index 000000000..56631833d --- /dev/null +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/mhc_post_combine_norm_prefill.cuh @@ -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 +#include + +#include +#include +#include +#include + +#include + +#include + +namespace sglang { +struct MhcPostCombineNormPrefillParams { + const bf16_t *x, *residual, *weight; + bf16_t *residual_out, *output; + const float *post, *comb, *pre; + float eps; +}; + +template +__global__ __launch_bounds__(Threads) void mhc_post_combine_norm_prefill_kernel( + const __grid_constant__ MhcPostCombineNormPrefillParams p) { + using namespace device; + using V = AlignedVector; + __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(x[k]); + const auto r0 = cast(residual[0][k]); + const auto r1 = cast(residual[1][k]); + const auto r2 = cast(residual[2][k]); + const auto r3 = cast(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(fp32x2_t{a, b}); + const auto rounded = cast(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(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(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(v[k]); + const auto weight = cast(w[k]); + out[k] = cast( + 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 +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(); + TensorMatcher({m, 5120}).with_dtype().with_device(dev).verify(x).verify(output); + TensorMatcher({m, 4, 5120}).with_dtype().with_device(dev).verify(residual).verify(residual_out); + TensorMatcher({m, 4}).with_dtype().with_device(dev).verify(post).verify(pre); + TensorMatcher({m, 4, 4}).with_dtype().with_device(dev).verify(comb); + TensorMatcher({5120}).with_dtype().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(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(x.data_ptr()), + static_cast(residual.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(residual_out.data_ptr()), + static_cast(output.data_ptr()), + static_cast(post.data_ptr()), + static_cast(comb.data_ptr()), + static_cast(pre.data_ptr()), + eps}; + LaunchKernel(dim3(m.unwrap()), Threads, dev.unwrap()).launch(mhc_post_combine_norm_prefill_kernel, params); + } +}; +} // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/distributed/all_reduce_fusion.cuh b/python/sglang/kernels/jit/csrc/distributed/all_reduce_fusion.cuh index c23127a07..175d2beb8 100644 --- a/python/sglang/kernels/jit/csrc/distributed/all_reduce_fusion.cuh +++ b/python/sglang/kernels/jit/csrc/distributed/all_reduce_fusion.cuh @@ -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::kBlockSize) __cluster_dims__(1, kClusterSize, 1) void moe_finalize_all_reduce_kernel( const __grid_constant__ MoeFinalizeAllReduceParams params) { @@ -377,7 +378,14 @@ __global__ __launch_bounds__(RowClusterTrait::kBlockSi if constexpr (!kNorm) { const auto red = reduce_vec(vec); ptx::st_global_16B(red, params.out, vid); - if constexpr (kMhc) mhc_post_vec(params, red, row_idx, hvec); + if constexpr (kMhc) { + if constexpr (kCollapse) { + const auto combined = mhc_post_vec(params, red, row_idx, hvec); + ptx::st_global_16B(combined, params.normalized, vid); + } else { + mhc_post_vec(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 || std::is_same_v); @@ -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 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().with_device(device).verify(pre.value()); TensorMatcher({T, kHiddenDim}).with_dtype().with_device(device).verify(normalized.value()); } diff --git a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py index 9cac53e38..50572a5fd 100644 --- a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py +++ b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py @@ -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, diff --git a/python/sglang/kernels/ops/communication/all_reduce_mhc_combine.py b/python/sglang/kernels/ops/communication/all_reduce_mhc_combine.py new file mode 100644 index 000000000..df4410650 --- /dev/null +++ b/python/sglang/kernels/ops/communication/all_reduce_mhc_combine.py @@ -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 + ) diff --git a/python/sglang/kernels/ops/layernorm/mhc_post_combine.py b/python/sglang/kernels/ops/layernorm/mhc_post_combine.py new file mode 100644 index 000000000..4509d6e36 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/mhc_post_combine.py @@ -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 diff --git a/python/sglang/kernels/ops/layernorm/mhc_post_combine_norm_prefill.py b/python/sglang/kernels/ops/layernorm/mhc_post_combine_norm_prefill.py new file mode 100644 index 000000000..174b042d7 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/mhc_post_combine_norm_prefill.py @@ -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 diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index b7c9fde82..2a7e17f64 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -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. diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py index 785a9ddd3..963f98ad3 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py @@ -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 diff --git a/python/sglang/kernels/ops/speculative/row_argmax.py b/python/sglang/kernels/ops/speculative/row_argmax.py index 6dbc6c5c0..21adafc2e 100644 --- a/python/sglang/kernels/ops/speculative/row_argmax.py +++ b/python/sglang/kernels/ops/speculative/row_argmax.py @@ -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)]( diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 91fcd94f5..86bb2e580 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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} diff --git a/python/sglang/srt/layers/moe/mhc_post_fusion.py b/python/sglang/srt/layers/moe/mhc_post_fusion.py index 64b90d32e..3d3ff204c 100644 --- a/python/sglang/srt/layers/moe/mhc_post_fusion.py +++ b/python/sglang/srt/layers/moe/mhc_post_fusion.py @@ -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. diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index f90160145..7e5f844b7 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -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) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 67243d3d7..3ca0915da 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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. diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 717931280..a3b8b6102 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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 diff --git a/test/registered/kernels/ops/attention/test_fp4_indexer.py b/test/registered/kernels/ops/attention/test_fp4_indexer.py index 45639a098..3a8933056 100644 --- a/test/registered/kernels/ops/attention/test_fp4_indexer.py +++ b/test/registered/kernels/ops/attention/test_fp4_indexer.py @@ -116,6 +116,89 @@ def test_quantize_fp4_indexer_tensor(num_tokens: int) -> None: torch.testing.assert_close(x_sf, ref_sf) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="Vectorized prefill dispatch targets Blackwell", +) +@pytest.mark.parametrize("rows", [4096, 4097, 16384, 524288]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rne", [False, True]) +@pytest.mark.parametrize("strided", [False, True]) +def test_prefill_quantization_matches_single_row_and_replay(rows, dtype, rne, strided): + from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( + _quantize_fp4_indexer_kernel, + ) + + x = torch.randn(rows, 256 if strided else 128, device="cuda", dtype=dtype) + if strided: + x = x[:, ::2] + + def reference(): + q = torch.empty(rows, 64, device="cuda", dtype=torch.int8) + sf = torch.empty(rows, device="cuda", dtype=torch.int32) + _quantize_fp4_indexer_kernel[(rows,)]( + x.contiguous(), + q, + sf, + BLOCK_N=128, + GROUP_N=32, + RNE=rne, + ) + return q, sf + + for _ in range(3): + quantize_fp4_indexer_tensor(x, rne) + reference() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = quantize_fp4_indexer_tensor(x, rne) + boundaries = torch.tensor( + [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, 6.0], device="cuda", dtype=dtype + ) + for scale in (0.0, 1e-6, 1.0, 1e3): + x.normal_().mul_(scale) + x[:2] = boundaries.repeat(16) + x[1].neg_() + graph.replay() + for a, b in zip(actual, reference()): + assert torch.equal(a, b) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="Vectorized prefill dispatch targets Blackwell", +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rne", [False, True]) +def test_prefill_quantization_nonfinite_group_replay(dtype, rne): + from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( + _quantize_fp4_indexer_kernel, + ) + + rows = 4097 # Also exercise the final partial CTA. + x = torch.randn(rows, HEAD_DIM, device="cuda", dtype=dtype) + expected = ( + torch.empty(rows, FP4_DIM, device="cuda", dtype=torch.int8), + torch.empty(rows, device="cuda", dtype=torch.int32), + ) + for _ in range(3): + quantize_fp4_indexer_tensor(x, rne) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = quantize_fp4_indexer_tensor(x, rne) + for value in (float("inf"), float("-inf"), float("nan"), 0.0): + x.normal_() + x[:, 0] = value + x[:, 32:64] = value + x[-1, 96:] = value + graph.replay() + _quantize_fp4_indexer_kernel[(rows,)]( + x, *expected, BLOCK_N=HEAD_DIM, GROUP_N=GROUP_SIZE, RNE=rne + ) + for a, b in zip(actual, expected): + assert torch.equal(a, b) + + @pytest.mark.parametrize("num_tokens", [1, 16, 96]) def test_fp4_index_cache_store_layout(num_tokens: int) -> None: torch.manual_seed(num_tokens)