diff --git a/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh new file mode 100644 index 000000000000..c828803d8a02 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh @@ -0,0 +1,263 @@ +// SPDX-License-Identifier: Apache-2.0 +// Beam's BF16 sandwich norm over the existing CustomAllReduceV2 push protocol. +#pragma once +#include + +#include +#include +#include +#include + +#include +#include + +#include + +#include + +namespace sglang { +namespace beam_ar { +using Vec = device::AlignedVector; +using Lamport = device::distributed::LamportTrait; + +template +struct Params { + const bf16_t *input, *residual, *post_weight, *next_weight; + bf16_t *combined, *normalized; + float eps, x_scale, residual_scale, routed_scale; + uint32_t rows, rank, counters; + device::distributed::PushWorkSpace workspace; +}; + +SGL_DEVICE void cluster_arrive_relaxed() { + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} +SGL_DEVICE void cluster_wait() { + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); +} + +/// Every thread holds eight adjacent BF16 values of one row. +template +SGL_DEVICE float inv_rms(const Vec& v, float eps) { + using namespace device; + constexpr uint32_t warps = Hidden / 8 / Cluster / 32; + __shared__ float partial[Cluster][warps]; + float sum = 0.0f; +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const auto f = cast(v[j]); + sum += f.x * f.x; + sum += f.y * f.y; + } + sum = warp::reduce_sum(sum); + const auto lane = threadIdx.x % 32; + if constexpr (Cluster > 1) { + const auto cluster = cooperative_groups::this_cluster(); + if (lane < Cluster) *cluster.map_shared_rank(&partial[blockIdx.y][threadIdx.x / 32], lane) = sum; + cluster.sync(); + } else { + if (lane == 0) partial[0][threadIdx.x / 32] = sum; + __syncthreads(); + } + sum = lane < Cluster * warps ? partial[lane / warps][lane % warps] : 0.0f; + sum = warp::reduce_sum(sum); + if constexpr (Cluster > 1) + cooperative_groups::this_cluster().sync(); + else + __syncthreads(); + return rsqrtf(sum / float(Hidden) + eps); +} + +template +__global__ void __launch_bounds__(Hidden / 8 / Cluster) __cluster_dims__(1, Cluster, 1) + kernel(const __grid_constant__ Params p) { + using namespace device; + namespace ptx = device::ptx; + device::PDLWaitPrimary(); + const auto row = blockIdx.x; + const auto tx = blockIdx.y * blockDim.x + threadIdx.x; + const auto epoch = distributed::PushEpoch{p.workspace}; + if (row == p.rows) { + // Keep every counter in phase when interleaving with generic all-reduces. + __syncthreads(); + if (blockIdx.y == 0) epoch.unsafe_flip_range(row, p.counters); + device::PDLTriggerSecondary(); + return; + } + // inv_rms writes peers' shared memory, which is valid only once every CTA in the cluster has started. + if constexpr (Cluster > 1) cluster_arrive_relaxed(); + const auto idx = row * (Hidden / 8) + tx; + Vec value; + if constexpr (TopK == 1) { + value.load(p.input, idx); + } else { + fp32x2_t sum[4]{}; +#pragma unroll + for (uint32_t k = 0; k < TopK; ++k) { + Vec expert; + expert.load(p.input, (row * TopK + k) * (Hidden / 8) + tx); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + auto v = cast(expert[j]); + sum[j].x += v.x; + sum[j].y += v.y; + } + } +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) + value[j] = cast(fp32x2_t{sum[j].x * p.routed_scale, sum[j].y * p.routed_scale}); + } + Lamport::clear_pos_zero(value.data()); + if (p.workspace.mc_workspace != nullptr) { + ptx::st_multimem_16B(value, p.workspace.mc_workspace + epoch.slot_offset(p.rank), idx); + } else { +#pragma unroll + for (uint32_t rank = 0; rank < World; ++rank) + ptx::st_relaxed_16B(value, epoch.slot_ptr(rank, p.rank), idx); + } + // Fetch epilogue inputs while the peers publish; asm loads stay ahead of the poll. + Vec post, residual; + ptx::ld_relaxed_16B(post, p.post_weight, tx); + ptx::ld_relaxed_16B(residual, p.residual, idx); + Vec peers[World]; + bool pending; + do { + pending = false; +#pragma unroll + for (uint32_t rank = 0; rank < World; ++rank) { + ptx::ld_relaxed_16B(peers[rank], epoch.slot_ptr(p.rank, rank), idx); + pending |= Lamport::has_pos_zero(peers[rank].data()); + } + } while (pending); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + fp32x2_t sum{0.0f, 0.0f}; +#pragma unroll + for (uint32_t rank = 0; rank < World; ++rank) { + const auto x = cast(peers[rank][j]); + sum.x += x.x; + sum.y += x.y; + } + value[j] = cast(sum); + } + if constexpr (Cluster > 1) cluster_wait(); + const float post_rstd = inv_rms(value, p.eps); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + auto v = cast(value[j]); + const auto w = cast(post[j]); + v = cast(cast(fp32x2_t{v.x * post_rstd * w.x, v.y * post_rstd * w.y})); + v = cast(cast(fp32x2_t{v.x * p.x_scale, v.y * p.x_scale})); + auto r = cast(residual[j]); + r = cast(cast(fp32x2_t{r.x * p.residual_scale, r.y * p.residual_scale})); + value[j] = cast(fp32x2_t{v.x + r.x, v.y + r.y}); + } + ptx::st_global_16B(value, p.combined, idx); + if constexpr (NextNorm) { + const float rstd = inv_rms(value, p.eps); + Vec weight; + weight.load(p.next_weight, tx); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const auto v = cast(value[j]); + const auto w = cast(weight[j]); + value[j] = cast(fp32x2_t{v.x * rstd * w.x, v.y * rstd * w.y}); + } + ptx::st_global_16B(value, p.normalized, idx); + } + Vec zero; + Lamport::fill_pos_zero(zero.data()); +#pragma unroll + for (uint32_t rank = 0; rank < World; ++rank) + ptx::st_global_16B(zero, epoch.slot_ptr(p.rank, rank), idx); + // A successor must wait for completion before consuming this plane again. + device::PDLTriggerSecondary(); + if constexpr (Cluster > 1) + cooperative_groups::this_cluster().sync(); + else + __syncthreads(); + if (blockIdx.y == 0) epoch.flip(); +} +} // namespace beam_ar + +/// Validate and launch a BF16 all-reduce + sandwich/residual/next norm. +template +struct BeamAllReduceNorm { + static_assert(World == 2 || World == 4 || World == 8); + // Whole warps per CTA, and at most 32 in a cluster so inv_rms reduces them in one warp. + static_assert(Hidden % (Cluster * 256) == 0 && Hidden <= 8192); + static void + run(host::distributed::CommunicatorRef ref, + tvm::ffi::TensorView input, + tvm::ffi::TensorView residual, + tvm::ffi::TensorView post, + std::optional next, + tvm::ffi::TensorView combined, + std::optional normalized, + double eps, + double x_scale, + double residual_scale, + double routed_scale) { + using namespace host; + const auto& push = ref.get()->get_push_obj(); + CHECK_HOST(push.world_size == World); + SymbolicSize rows{"rows"}; + SymbolicDevice device; + const auto matrix = [&](tvm::ffi::TensorView tensor) { + TensorMatcher({rows, Hidden}) + .with_strides({Hidden, 1}) + .with_dtype() + .with_device(device) + .ensure_alignment(16) + .verify(tensor); + }; + const auto weight = [&](tvm::ffi::TensorView tensor) { + TensorMatcher({Hidden}) + .with_strides({1}) + .with_dtype() + .with_device(device) + .ensure_alignment(16) + .verify(tensor); + }; + if constexpr (TopK == 1) + matrix(input); + else + TensorMatcher({rows, TopK, Hidden}) + .with_strides({TopK * Hidden, Hidden, 1}) + .with_dtype() + .with_device(device) + .ensure_alignment(16) + .verify(input); + matrix(residual); + matrix(combined); + weight(post); + CHECK_HOST(next.has_value() == normalized.has_value()); + if (next) { + weight(*next); + matrix(*normalized); + } + const auto n = rows.unwrap(); + const int64_t bytes = n * Hidden * sizeof(bf16_t); + CHECK_HOST(n > 0 && n <= push.num_blocks && bytes <= push.slot_bytes); + const auto p = beam_ar::Params{ + .input = static_cast(input.data_ptr()), + .residual = static_cast(residual.data_ptr()), + .post_weight = static_cast(post.data_ptr()), + .next_weight = next ? static_cast(next->data_ptr()) : nullptr, + .combined = static_cast(combined.data_ptr()), + .normalized = normalized ? static_cast(normalized->data_ptr()) : nullptr, + .eps = static_cast(eps), + .x_scale = static_cast(x_scale), + .residual_scale = static_cast(residual_scale), + .routed_scale = static_cast(routed_scale), + .rows = static_cast(n), + .rank = push.rank, + .counters = push.num_blocks, + .workspace = push.get_workspace(bytes)}; + auto kernel = next ? beam_ar::kernel + : beam_ar::kernel; + LaunchKernel(dim3(n + (n < push.num_blocks), Cluster), Hidden / 8 / Cluster, input.device()) + .enable_pdl(true)(kernel, p); + } +}; +} // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh index d4db025da6fe..c543905f3fc7 100644 --- a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh @@ -19,7 +19,7 @@ namespace sglang { constexpr uint32_t kMaxBlockThreads = 1024; -template +template struct GEMMTraitN { static constexpr uint32_t N = N_; static constexpr uint32_t K = K_; @@ -32,9 +32,9 @@ struct GEMMTraitN { static constexpr uint32_t kNumBlocks = N / N_SPLIT; static constexpr uint32_t kBlockSize = (K * sizeof(bf16_t)) / kBytes; // always use the largest possible vector - static constexpr uint32_t kVecSize = device::kMaxVecBytes / sizeof(bf16_t); + static constexpr uint32_t kVecSize = device::kMaxVecBytes / sizeof(Weight); // may reduce block size for less thread usage - static constexpr uint32_t kUnroll = kBytes / device::kMaxVecBytes; + static constexpr uint32_t kUnroll = kBytes / (kVecSize * sizeof(bf16_t)); static_assert(kBlockSize % device::kWarpThreads == 0, "block size must be divisible by warp size"); static_assert(kBlockSize <= kMaxBlockThreads, "block size exceeds the maximum block size"); }; @@ -64,11 +64,11 @@ struct GEMMTraitK { struct TinyGEMMParams { void* __restrict__ out; const bf16_t* __restrict__ x; - const bf16_t* __restrict__ w; + const void* __restrict__ w; int64_t stride_x; }; -template +template TINY_GEMM_KERNEL void tiny_n_gemm_kernel(const TinyGEMMParams params) { using namespace device; constexpr uint32_t N = Trait::N; @@ -82,10 +82,11 @@ TINY_GEMM_KERNEL void tiny_n_gemm_kernel(const TinyGEMMParams params) { const uint32_t bx = blockIdx.x; const uint32_t tx = threadIdx.x; - const bf16_t* w_tile = params.w + bx * (N_SPLIT * K); + const Weight* w_tile = static_cast(params.w) + bx * (N_SPLIT * K); // prefetch weight before PDL - vec_t wv[N_SPLIT][kUnroll]; + using weight_vec_t = AlignedVector, kVecSize / 2>; + weight_vec_t wv[N_SPLIT][kUnroll]; #pragma unroll for (uint32_t n = 0; n < N_SPLIT; ++n) { #pragma unroll @@ -155,7 +156,7 @@ TINY_GEMM_KERNEL void tiny_k_gemm_kernel(const TinyGEMMParams params) { const uint32_t tx = threadIdx.x; const uint32_t n_idx = bx * N_SPLIT + tx / kNumKLanes; const uint32_t work_id = tx % kNumKLanes; - const bf16_t* w_tile = params.w + n_idx * K; + const bf16_t* w_tile = static_cast(params.w) + n_idx * K; // Weight prefetch: address is input-independent, load before the PDL wait. vec_t wv; @@ -178,14 +179,21 @@ TINY_GEMM_KERNEL void tiny_k_gemm_kernel(const TinyGEMMParams params) { PDLTriggerSecondary(); } -template +template < + uint32_t N, + uint32_t K, + uint32_t kMaxM, + uint32_t N_SPLIT, + typename OutT, + bool kUsePDL, + typename Weight = bf16_t> struct TinyNGemmKernel { - using Trait = GEMMTraitN; + using Trait = GEMMTraitN; using KernelFn = void (*)(TinyGEMMParams); template static constexpr auto make_table(std::index_sequence) { - return std::array{nullptr, tiny_n_gemm_kernel...}; + return std::array{nullptr, tiny_n_gemm_kernel...}; } static constexpr auto kTable = make_table(std::make_index_sequence{}); @@ -196,8 +204,17 @@ struct TinyNGemmKernel { auto M = SymbolicSize{"num_tokens"}; auto device = SymbolicDevice{}; device.set_options(); - TensorMatcher({M, K}).with_strides({-1, 1}).with_dtype().with_device(device).verify(x); - TensorMatcher({N, K}).with_dtype().with_device(device).verify(w); + TensorMatcher({M, K}) + .with_strides({-1, 1}) + .with_dtype() + .with_device(device) + .ensure_alignment(Trait::kVecSize * sizeof(bf16_t)) + .verify(x); + TensorMatcher({N, K}) + .with_dtype() + .with_device(device) + .ensure_alignment(Trait::kVecSize * sizeof(Weight)) + .verify(w); TensorMatcher({M, N}).with_dtype().with_device(device).verify(out); const auto num_tokens = static_cast(M.unwrap()); if (num_tokens == 0) return; @@ -209,7 +226,7 @@ struct TinyNGemmKernel { const auto params = TinyGEMMParams{ .out = out.data_ptr(), .x = static_cast(x.data_ptr()), - .w = static_cast(w.data_ptr()), + .w = w.data_ptr(), .stride_x = static_cast(x.stride(0)), }; LaunchKernel(Trait::kNumBlocks, Trait::kBlockSize, device.unwrap()) // @@ -248,7 +265,7 @@ struct TinyKGemmKernel { const auto params = TinyGEMMParams{ .out = out.data_ptr(), .x = static_cast(x.data_ptr()), - .w = static_cast(w.data_ptr()), + .w = w.data_ptr(), .stride_x = static_cast(x.stride(0)), }; LaunchKernel(Trait::kNumBlocks, Trait::kBlockSize, device.unwrap()) // diff --git a/python/sglang/kernels/jit/include/sgl_kernel/gemm/dot_product.cuh b/python/sglang/kernels/jit/include/sgl_kernel/gemm/dot_product.cuh index 140980cdb934..bb393cc6eb5a 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/gemm/dot_product.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/gemm/dot_product.cuh @@ -35,6 +35,16 @@ SGL_DEVICE void dot_product_vec(AlignedVector a, AlignedVector +SGL_DEVICE void dot_product_vec(AlignedVector a, AlignedVector b, float& acc) { +#pragma unroll + for (uint32_t i = 0; i < N; ++i) { + acc = math::fma_f32_bf16(a[i].x, b[i].x, acc); + acc = math::fma_f32_bf16(a[i].y, b[i].y, acc); + } +} + } // namespace device } // namespace sglang diff --git a/python/sglang/kernels/ops/activation/_triton_silu.py b/python/sglang/kernels/ops/activation/_triton_silu.py new file mode 100644 index 000000000000..d881f147fa4c --- /dev/null +++ b/python/sglang/kernels/ops/activation/_triton_silu.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Private Triton activation helpers with CUDA fast-math rounding.""" + +import triton +import triton.language as tl + + +@triton.jit +def _silu_fast_math(gate): + # Matches CUDA silu_and_mul, whose --use_fast_math build computes + # __fdividef(x, 1 + __expf(-x)); accurate exp/div can move BF16 results. + # Keep SiLU in FP32; callers apply the multiply before the final BF16 store. + # __expf(-x) == ex2.approx(x * -log2(e)); 0fBFB8AA3B is -log2(e). + exp_neg_gate = tl.inline_asm_elementwise( + "{ mul.ftz.f32 $0, $1, 0fBFB8AA3B; ex2.approx.ftz.f32 $0, $0; }", + "=f,f", + [gate], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + return tl.inline_asm_elementwise( + "div.approx.ftz.f32 $0, $1, $2;", + "=f,f,f", + [gate, 1.0 + exp_neg_gate], + dtype=tl.float32, + is_pure=True, + pack=1, + ) diff --git a/python/sglang/kernels/ops/communication/__init__.py b/python/sglang/kernels/ops/communication/__init__.py index cb2b2411e9e1..8e6e79a69442 100644 --- a/python/sglang/kernels/ops/communication/__init__.py +++ b/python/sglang/kernels/ops/communication/__init__.py @@ -50,6 +50,15 @@ ) del _mod, _fn +register_kernel( + KernelSpec( + op="communication.beam_all_reduce_norm", + backend=KernelBackend.JIT, + target="sglang.kernels.ops.communication.beam_all_reduce:beam_all_reduce_norm", + capabilities=frozenset({CapabilityRequirement.cuda(min_sm=(9, 0))}), + ) +) + # Muse Spark collectives. The MoE finalize op (with its optional fused TP # all-reduce) is built per layer: inventoried through its class methods. register_kernel( diff --git a/python/sglang/kernels/ops/communication/beam_all_reduce.py b/python/sglang/kernels/ops/communication/beam_all_reduce.py new file mode 100644 index 000000000000..32e8dc0cba2f --- /dev/null +++ b/python/sglang/kernels/ops/communication/beam_all_reduce.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Beam BF16 all-reduce with FP32 sandwich normalization and rounded residuals.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit + +if TYPE_CHECKING: + from sglang.kernels.ops.communication.all_reduce import Communicator + + +@cache_once +def _module(world_size: int, hidden_size: int, cluster: int, topk: int): + return load_jit( + "beam_all_reduce_norm", + str(world_size), + str(hidden_size), + str(cluster), + str(topk), + cuda_files=["distributed/beam_all_reduce.cuh"], + cuda_wrappers=[ + ( + "run", + f"BeamAllReduceNorm<{world_size}, {hidden_size}, {cluster}, {topk}>::run", + ) + ], + # Round every multiply and add separately, as the unfused norm does. + extra_cuda_cflags=["--fmad=false"], + ) + + +def _cluster_size(world_size: int, rows: int) -> int: + # Measured on 8xH200: 5-CTA clusters are faster at TP8 up to 16 rows; one CTA + # per row is faster at TP2 and TP4. + return 5 if world_size == 8 and rows <= 16 else 1 + + +def beam_all_reduce_norm( + comm: Communicator, + branch: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor, + next_weight: Optional[torch.Tensor], + eps: float, + scales: tuple[float, float], + *, + routed_scale: float = 1.0, +) -> tuple[Optional[torch.Tensor], torch.Tensor]: + """Return `(normalized, combined)` using a caller-owned V2 communicator. + + All ranks must call with identical shapes/epilogue and serialize use of the + communicator. Inputs and weights are contiguous BF16 CUDA tensors; the + complete row view must fit the communicator's push plane and counter array. + A three-dimensional branch contains eight already routing-weighted expert + outputs per token. Their FP32 sum is multiplied by routed_scale and rounded + to BF16 before rank reduction. A two-dimensional branch is already combined. + The post-norm, each scaled residual term, their sum, and the optional next + norm retain separate BF16 rounding boundaries. + """ + topk = branch.shape[1] if branch.ndim == 3 else 1 + if branch.ndim not in (2, 3) or branch.shape[-1] != 5120 or topk not in (1, 8): + raise ValueError("Beam fusion requires [tokens, 5120] or [tokens, 8, 5120]") + combined = torch.empty_like(residual) + normalized = torch.empty_like(residual) if next_weight is not None else None + if branch.shape[0]: + _module( + comm.world_size, + branch.shape[-1], + _cluster_size(comm.world_size, branch.shape[0]), + topk, + ).run( + comm, + branch, + residual, + post_weight, + next_weight, + combined, + normalized, + eps, + *scales, + routed_scale, + ) + return normalized, combined diff --git a/python/sglang/kernels/ops/elementwise/__init__.py b/python/sglang/kernels/ops/elementwise/__init__.py index da2e40a3041c..e67a457849dc 100644 --- a/python/sglang/kernels/ops/elementwise/__init__.py +++ b/python/sglang/kernels/ops/elementwise/__init__.py @@ -78,3 +78,17 @@ capabilities=CUDA_ONLY, ) ) + + +# Beam fused op, built per layer: inventoried through its class methods. +for _bk in ("torch", "triton"): + register_kernel( + KernelSpec( + op="elementwise.beam_attention_gate", + backend=KernelBackend(_bk), + target="sglang.kernels.ops.elementwise.beam_fused:BeamAttentionGateOp." + + ("forward_native" if _bk == "torch" else f"forward_{_bk}"), + capabilities=frozenset() if _bk == "torch" else CUDA_ONLY, + ) + ) +del _bk diff --git a/python/sglang/kernels/ops/elementwise/beam_fused.py b/python/sglang/kernels/ops/elementwise/beam_fused.py new file mode 100644 index 000000000000..75256b19f4ad --- /dev/null +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Beam attention gate with activation-dtype bias and sigmoid rounding.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +from sglang.kernels.fused_op import BaseFusedOp +from sglang.kernels.spec import CapabilityRequirement, KernelBackend +from sglang.srt.utils.custom_op import register_custom_op + + +@triton.jit +def _fused_attention_gate_kernel( + out_ptr, + attn_ptr, + gate_ptr, + bias_ptr, + hidden_dim: tl.constexpr, + gate_row_stride: tl.constexpr, + gate_col_stride: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_H: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + block = tl.program_id(1) + offsets = block * BLOCK_H + tl.arange(0, BLOCK_H) + mask = offsets < hidden_dim + head = offsets // HEAD_DIM + act_dtype = attn_ptr.dtype.element_ty + + off = row * hidden_dim + offsets + attn = tl.load(attn_ptr + off, mask=mask, other=0.0).to(tl.float32) + gate_off = row * gate_row_stride + offsets * gate_col_stride + g = tl.load(gate_ptr + gate_off, mask=mask, other=0.0).to(tl.float32) + b = tl.load(bias_ptr + head, mask=mask, other=0.0).to(tl.float32) + # As in eager, gate + bias, its sigmoid, and the product round to act_dtype. + g = (g + b).to(act_dtype).to(tl.float32) + if act_dtype == tl.float16: + # Approximate exp/div can cross an FP16 rounding boundary near zero. + m = tl.div_rn(1.0, 1.0 + libdevice.exp(-g)) + else: + m = tl.sigmoid(g) + m = m.to(act_dtype).to(tl.float32) + tl.store(out_ptr + off, (attn * m).to(act_dtype), mask=mask) + + +def _fake_attention_gate(attention_output, gate, gate_bias, head_dim): + return torch.empty_like(attention_output, memory_format=torch.contiguous_format) + + +@torch.compiler.disable +def _eager_sigmoid(x: torch.Tensor) -> torch.Tensor: + # Inductor's sigmoid can round FP16 one ULP away from ATen's. + return torch.sigmoid(x) + + +@register_custom_op(fake_impl=_fake_attention_gate) +def fused_attention_gate( + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + head_dim: int, +) -> torch.Tensor: + """`attention_output * sigmoid(gate + gate_bias[head])` for 2-D `[tokens, heads*head_dim]` inputs.""" + if attention_output.shape != gate.shape or attention_output.ndim != 2: + raise ValueError("attention_output and gate must be matching 2-D tensors") + if attention_output.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("attention_output must be fp16, bf16 or fp32") + attention_output = attention_output.contiguous() + hidden_dim = attention_output.shape[1] + if gate_bias.numel() * head_dim != hidden_dim or not gate_bias.is_contiguous(): + raise ValueError("gate_bias must hold one contiguous entry per local head") + out = torch.empty_like(attention_output) + rows = attention_output.shape[0] + if rows == 0: + return out + block_h = 1024 + _fused_attention_gate_kernel[(rows, triton.cdiv(hidden_dim, block_h))]( + out, + attention_output, + gate, + gate_bias, + hidden_dim=hidden_dim, + gate_row_stride=gate.stride(0), + gate_col_stride=gate.stride(1), + HEAD_DIM=head_dim, + BLOCK_H=block_h, + num_warps=4, + ) + return out + + +class BeamAttentionGateOp(BaseFusedOp): + """Apply a per-head bias and sigmoid gate to flat attention output.""" + + op = "elementwise.beam_attention_gate" + capabilities = {KernelBackend.TRITON: {CapabilityRequirement.CUDA}} + + def _torch_compile_forward(self, num_tokens: int): + # Native compilation can remove the required intermediate rounding. + return None + + def forward_torch_compile(self, *args, **kwargs): + if self._compiled_native is None: + # Without emulated casts, Inductor drops eager's intermediate rounding. + self._compiled_native = torch.compile( + self.forward_native, options={"emulate_precision_casts": True} + ) + return self._compiled_native(*args, **kwargs) + + def forward_native( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + head_dim: int, + ) -> torch.Tensor: + heads = gate_bias.numel() + multiplier = _eager_sigmoid( + gate.reshape(-1, heads, head_dim) + gate_bias[None, :, None] + ) + return (attention_output.reshape(-1, heads, head_dim) * multiplier).reshape_as( + attention_output + ) + + def forward_triton( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + head_dim: int, + ) -> torch.Tensor: + return fused_attention_gate(attention_output, gate, gate_bias, head_dim) diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index bd5dcaa1eb63..9a0eff8a5a9f 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -533,3 +533,19 @@ def kimi_k3_tiny_gemm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: capabilities=_SM10X, ) ) +register_kernel( + KernelSpec( + op="gemm.router_gemv", + backend=KernelBackend.TRITON, + target="sglang.kernels.ops.gemm.router_gemv:router_gemv", + capabilities=frozenset({CapabilityRequirement.CUDA, CapabilityRequirement.HIP}), + ) +) +register_kernel( + KernelSpec( + op="gemm.router_gemm_bf16x3_cublas", + backend=KernelBackend.TORCH, + target="sglang.kernels.ops.gemm.router_gemv:router_gemm_bf16x3_cublas", + capabilities=CUDA_ONLY, + ) +) diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index ed4d76217186..906c77319445 100644 --- a/python/sglang/kernels/ops/gemm/router_gemv.py +++ b/python/sglang/kernels/ops/gemm/router_gemv.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Skinny router-logits GEMV: [M, K] bf16 x [N, K] bf16 -> [M, N] fp32. +"""Skinny router-logits GEMV: [M, K] bf16 x [N, K] bf16 or fp32 -> [M, N] fp32. MoE router gates are tiny (N ~ 128 experts, K ~ hidden) but hipblaslt's solutions for M<=8, N=128 run at ~0.1 TB/s on gfx950 (~12us for 1.6MB). This @@ -17,6 +17,7 @@ import triton import triton.language as tl +from sglang.kernels.ops.quantization.bf16x3_split import split_fp32_to_bf16x3 from sglang.srt.utils import is_gfx95_supported, is_hip _is_hip = is_hip() @@ -31,13 +32,21 @@ (32, (16, 32, 512, 8, 4)), (_MAX_M, (32, 32, 512, 8, 4)), ) +# FP32 weights on CUDA (FP32 router gates) use the FP32 row path only, up to 8 rows, +# a conservative cap: for N=256, K=5120 on GB300 it beats bf16x3 through 14 rows. +_CUDA_FP32_CONFIGS = ((8, (1, 16, 128, 8, 8)),) -def _config(m: int): - for max_m, cfg in _CONFIGS: +def _configs(w_dtype: torch.dtype): + return _CUDA_FP32_CONFIGS if w_dtype == torch.float32 else _CONFIGS + + +def _config(m: int, w_dtype: torch.dtype): + configs = _configs(w_dtype) + for max_m, cfg in configs: if m <= max_m: return cfg - return _CONFIGS[-1][1] + return configs[-1][1] @triton.jit @@ -120,7 +129,7 @@ def _router_gemv_kernel( def _get_scratch(device: torch.device, n: int) -> Tuple[torch.Tensor, torch.Tensor]: key = (device, n) if key not in _scratch: - max_split_k = max(cfg[3] for _, cfg in _CONFIGS) + max_split_k = max(cfg[3] for _, cfg in _CONFIGS + _CUDA_FP32_CONFIGS) _scratch[key] = ( torch.empty(max_split_k * _MAX_M * n, dtype=torch.float32, device=device), torch.zeros(n * _MAX_M // 16, dtype=torch.int32, device=device), @@ -131,18 +140,21 @@ def _get_scratch(device: torch.device, n: int) -> Tuple[torch.Tensor, torch.Tens def router_gemv_supported(x: torch.Tensor, w: torch.Tensor) -> bool: m, k = x.shape n = w.shape[0] - _, block_n, block_k, _, _ = _config(m) + _, block_n, block_k, _, _ = _config(m, w.dtype) + # BF16 gates are tuned for gfx950 only; FP32 gates for CUDA only. + platform_ok = ( + _is_gfx95_supported if w.dtype == torch.bfloat16 else not _is_hip + ) and w.dtype in (torch.bfloat16, torch.float32) return ( - _is_gfx95_supported + platform_ok and x.device.type == "cuda" and w.device == x.device and x.dtype == torch.bfloat16 - and w.dtype == torch.bfloat16 and x.stride(1) == 1 and w.stride(1) == 1 and n % block_n == 0 and k % block_k == 0 - and m <= _MAX_M + and m <= _configs(w.dtype)[-1][0] ) @@ -150,7 +162,7 @@ def router_gemv(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: """Router logits in fp32; caller guards with router_gemv_supported().""" m, k = x.shape n = w.shape[0] - block_m, block_n, block_k, split_k, num_warps = _config(m) + block_m, block_n, block_k, split_k, num_warps = _config(m, w.dtype) partials, counter = _get_scratch(x.device, n) out = torch.empty(m, n, device=x.device, dtype=torch.float32) grid = (n // block_n, triton.cdiv(m, block_m), split_k) @@ -173,3 +185,15 @@ def router_gemv(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: num_warps=num_warps, ) return out + + +def router_gemm_bf16x3_cublas(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """[M, K] bf16 x [N, K] fp32 router logits on BF16 tensor cores. + + The fp32 weight is split per call, so in-place weight updates take effect + without any refresh. Each bf16 x bf16 product is exact within fp32's exponent + range, but tensor-core accumulation is not IEEE fp32 summation. + """ + n = w.shape[0] + w3 = split_fp32_to_bf16x3(w) + return torch.mm(x, w3.t(), out_dtype=torch.float32).view(-1, 3, n).sum(dim=1) diff --git a/python/sglang/kernels/ops/gemm/tiny_gemm.py b/python/sglang/kernels/ops/gemm/tiny_gemm.py index 1b429732f62c..f028b9135b2a 100644 --- a/python/sglang/kernels/ops/gemm/tiny_gemm.py +++ b/python/sglang/kernels/ops/gemm/tiny_gemm.py @@ -1,4 +1,4 @@ -"""Tiny bf16 GEMM for skinny ``x[m, k] @ w[n, k].T`` with a handful of rows. +"""Tiny BF16-activation GEMM for skinny ``x[m, k] @ w[n, k].T`` with a handful of rows. Two kernels behind one entry point, picked by whichever of N / K is the tiny dimension: the N variant spreads K across a block and walks N, the K variant @@ -42,10 +42,24 @@ def _prefer_k_variant(n: int, k: int) -> bool: @cache_once def _jit_tiny_gemm_module( - n: int, k: int, max_m: int, n_split: int, out_dtype: torch.dtype + n: int, + k: int, + max_m: int, + n_split: int, + out_dtype: torch.dtype, + weight_dtype: torch.dtype, ) -> Module: use_k_variant = _prefer_k_variant(n, k) - args = make_cpp_args(n, k, max_m, n_split, out_dtype, is_arch_support_pdl()) + if weight_dtype not in (torch.bfloat16, torch.float32): + raise ValueError("tiny_gemm requires BF16 or FP32 weights") + if use_k_variant and weight_dtype != torch.bfloat16: + raise ValueError( + "FP32 weights require the tiny-N variant (N >= K is unsupported)" + ) + weight_args = () if use_k_variant else (weight_dtype,) + args = make_cpp_args( + n, k, max_m, n_split, out_dtype, is_arch_support_pdl(), *weight_args + ) name = "tiny_k_gemm" if use_k_variant else "tiny_n_gemm" kernel = "TinyKGemmKernel" if use_k_variant else "TinyNGemmKernel" return load_jit( @@ -137,7 +151,7 @@ def _tiny_gemm_custom_op( n_split: int, ) -> None: n, k = w.shape - module = _jit_tiny_gemm_module(n, k, max_m, n_split, out.dtype) + module = _jit_tiny_gemm_module(n, k, max_m, n_split, out.dtype, w.dtype) module.run(x, w, out) @@ -167,7 +181,7 @@ def tiny_gemm_bf16( rather than falling back. :param x: Shape [m, k], must be bf16_t, `m <= max_m` - :param w: Shape [n, k], must be bf16_t + :param w: Shape [n, k], BF16 or (for N < K) FP32; never downcast """ n, k = w.shape if out is None: diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index f9d04b0b83d1..a5d51a823a3c 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -83,8 +83,8 @@ ) ) -# Muse Spark fused norms. The per-layer ops are inventoried through their class -# methods; the NVFP4 quant op through its module-level instance. +# Muse Spark and Beam fused norms. The per-layer ops are inventoried through +# their class methods; the NVFP4 quant op through its module-level instance. for _op, _target, _backends in ( ( "muse_spark_norm_gate", @@ -101,6 +101,11 @@ "muse_spark_rmsnorm_nvfp4_quant:_FUSED_NORM_FP4_QUANT", ("torch", "jit", "helion"), ), + ( + "beam_sandwich_residual_rmsnorm", + "beam_fused:BeamSandwichResidualRMSNormOp", + ("torch", "triton"), + ), ): for _bk in _backends: register_kernel( diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py new file mode 100644 index 000000000000..b4bfdf5926fa --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Beam sandwich RMSNorm, depth-scaled residual, and next RMSNorm in one launch. + +The post-norm, both scaled operands, and the combined stream round to the +activation dtype. Norm reduction differences can be amplified by residual +cancellation, so the result is not guaranteed within one ULP. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.kernels.fused_op import BaseFusedOp +from sglang.kernels.spec import CapabilityRequirement, KernelBackend +from sglang.srt.utils.custom_op import register_custom_op + + +@triton.jit +def _fused_sandwich_residual_rmsnorm_kernel( + out_ptr, + combined_ptr, + x_ptr, + residual_ptr, + x_stride_row, + residual_stride_row, + sandwich_weight_ptr, + out_weight_ptr, + eps, + x_scale, + residual_scale, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_OUT_NORM: tl.constexpr, + USE_PDL: tl.constexpr, +): + row = tl.program_id(axis=0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + act_dtype = x_ptr.dtype.element_ty + + # The preceding kernel does not write the weights; load them before the PDL wait. + w = tl.load(sandwich_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + if HAS_OUT_NORM: + w2 = tl.load(out_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + x = tl.load(x_ptr + row * x_stride_row + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + rstd = tl.rsqrt(tl.sum(x * x, axis=0) / hidden_dim + eps) + # The standalone RMSNorm writes the activation dtype before the combine. + x = (x * rstd * w).to(act_dtype).to(tl.float32) + + residual = tl.load( + residual_ptr + row * residual_stride_row + offsets, mask=mask, other=0.0 + ) + # `depth_scaled_residual` casts the residual to the branch dtype first. + residual = residual.to(act_dtype).to(tl.float32) + xs = (x * x_scale).to(act_dtype).to(tl.float32) + rs = (residual * residual_scale).to(act_dtype).to(tl.float32) + combined = (xs + rs).to(act_dtype) + out_offsets = row * hidden_dim + offsets + tl.store(combined_ptr + out_offsets, combined, mask=mask) + + if HAS_OUT_NORM: + c = combined.to(tl.float32) + rstd2 = tl.rsqrt(tl.sum(c * c, axis=0) / hidden_dim + eps) + tl.store(out_ptr + out_offsets, (c * rstd2 * w2).to(act_dtype), mask=mask) + + +def fused_sandwich_residual_rmsnorm( + branch: torch.Tensor, + residual: torch.Tensor, + sandwich_weight: torch.Tensor, + out_weight: Optional[torch.Tensor], + eps: float, + x_scale: float, + residual_scale: float, + *, + enable_pdl: bool = False, +) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + """Return `(next_norm(combined), combined)`; the first is None without `out_weight`. + + `branch` and `residual` are 2-D `[tokens, hidden]` in the activation dtype. + `sandwich_weight` is the post-norm weight applied to `branch` before the + combine; `out_weight` is the following norm's weight (None returns only the + combined stream, for the last layer of a non-final pipeline stage). + `enable_pdl` launches with programmatic dependent launch; the caller must + check that the device supports it. + """ + if branch.ndim != 2 or residual.shape != branch.shape: + raise ValueError( + f"expected matching 2-D tensors; got {tuple(branch.shape)} and {tuple(residual.shape)}" + ) + if branch.stride(1) != 1 or residual.stride(1) != 1: + raise ValueError("branch and residual must be contiguous along hidden") + if branch.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("branch must be fp16, bf16 or fp32") + hidden_dim = branch.shape[1] + for w in (sandwich_weight, out_weight): + if w is not None and (w.shape != (hidden_dim,) or not w.is_contiguous()): + raise ValueError("norm weights must be contiguous [hidden]") + combined = torch.empty_like(branch, memory_format=torch.contiguous_format) + out = torch.empty_like(combined) if out_weight is not None else None + _fused_sandwich_residual_rmsnorm( + branch, + residual, + sandwich_weight, + out_weight, + eps, + x_scale, + residual_scale, + out, + combined, + enable_pdl, + ) + return out, combined + + +@register_custom_op(mutates_args=["out", "combined"]) +def _fused_sandwich_residual_rmsnorm( + branch: torch.Tensor, + residual: torch.Tensor, + sandwich_weight: torch.Tensor, + out_weight: Optional[torch.Tensor], + eps: float, + x_scale: float, + residual_scale: float, + out: Optional[torch.Tensor], + combined: torch.Tensor, + enable_pdl: bool, +) -> None: + hidden_dim = branch.shape[1] + rows = branch.shape[0] + if rows == 0: + return + _fused_sandwich_residual_rmsnorm_kernel[(rows,)]( + out if out is not None else combined, + combined, + branch, + residual, + branch.stride(0), + residual.stride(0), + sandwich_weight, + out_weight if out_weight is not None else sandwich_weight, + eps=eps, + x_scale=x_scale, + residual_scale=residual_scale, + hidden_dim=hidden_dim, + BLOCK_SIZE=triton.next_power_of_2(hidden_dim), + HAS_OUT_NORM=out_weight is not None, + USE_PDL=enable_pdl, + num_warps=8, + enable_fp_fusion=False, + **({"launch_pdl": True} if enable_pdl else {}), + ) + + +def _rmsnorm_native(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Pure-torch RMSNorm reference: ``(x / RMS(x)) * weight`` in FP32.""" + y = x.to(torch.float32) + y = y * torch.rsqrt(y.pow(2).mean(dim=-1, keepdim=True) + eps) + return (y * weight).to(x.dtype) + + +class BeamSandwichResidualRMSNormOp(BaseFusedOp): + """Post-normalize a branch, combine residuals, then normalize the next input. + + Returns ``(normed, combined)``; ``normed`` is None without ``out_weight`` + (a non-final pipeline stage). + """ + + op = "layernorm.beam_sandwich_residual_rmsnorm" + capabilities = {KernelBackend.TRITON: {CapabilityRequirement.CUDA}} + + def _torch_compile_forward(self, num_tokens: int): + # Native compilation can remove the required intermediate rounding. + return None + + def forward_torch_compile(self, *args, **kwargs): + if self._compiled_native is None: + # Without emulated casts, Inductor drops eager's intermediate rounding. + self._compiled_native = torch.compile( + self.forward_native, options={"emulate_precision_casts": True} + ) + return self._compiled_native(*args, **kwargs) + + def forward_native( + self, + branch: torch.Tensor, + residual: torch.Tensor, + sandwich_weight: torch.Tensor, + out_weight: Optional[torch.Tensor], + eps: float, + scales: tuple[float, float], + ) -> tuple[Optional[torch.Tensor], torch.Tensor]: + branch = _rmsnorm_native(branch, sandwich_weight, eps) + residual = residual.to(branch.dtype) + combined = branch * scales[0] + residual * scales[1] + out = ( + _rmsnorm_native(combined, out_weight, eps) + if out_weight is not None + else None + ) + return out, combined + + def forward_triton( + self, + branch: torch.Tensor, + residual: torch.Tensor, + sandwich_weight: torch.Tensor, + out_weight: Optional[torch.Tensor], + eps: float, + scales: tuple[float, float], + ) -> tuple[Optional[torch.Tensor], torch.Tensor]: + from sglang.kernels.jit.utils import is_arch_support_pdl + + return fused_sandwich_residual_rmsnorm( + branch, + residual, + sandwich_weight, + out_weight, + eps, + *scales, + enable_pdl=is_arch_support_pdl(), + ) diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index d519c41efb7e..0a03c3f7425c 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -27,6 +27,17 @@ ) ) +register_kernel( + KernelSpec( + op="moe.beam_decode_gemv", + backend=KernelBackend.TRITON, + target="sglang.kernels.ops.moe.beam_decode_gemv:beam_decode_gemv", + capabilities=frozenset( + {CapabilityRequirement.cuda(min_sm=(10, 3), max_sm=(10, 3))} + ), + ) +) + register_kernel( KernelSpec( op="moe.bf16_mega_gate", diff --git a/python/sglang/kernels/ops/moe/beam_decode_gemv.py b/python/sglang/kernels/ops/moe/beam_decode_gemv.py new file mode 100644 index 000000000000..a7f73912541d --- /dev/null +++ b/python/sglang/kernels/ops/moe/beam_decode_gemv.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SM103 single-token Beam MoE with weighted expert outputs. + +Gate/up and activation retain BF16 rounding boundaries. The down projection +multiplies FP32 routing coefficients before rounding to BF16. The caller owns +expert reduction, routed scaling and TP communication. Vector reductions use +a different FP32 summation order from the tensor-core reference. +""" + +import torch +import triton +import triton.language as tl + +from sglang.kernels.ops.activation._triton_silu import _silu_fast_math + + +@triton.jit +def _gate_up_gemv_kernel( + HiddenStates, + Weight, + TopkIds, + Intermediate, + HIDDEN_SIZE: tl.constexpr, + INTERMEDIATE_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + route = tl.program_id(0) + n = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + k = tl.arange(0, BLOCK_K) + expert = tl.load(TopkIds + route).to(tl.int64) + accumulator = tl.full((BLOCK_N, BLOCK_K), 0, tl.float32) + for start in range(tl.cdiv(HIDDEN_SIZE, BLOCK_K)): + kk = start * BLOCK_K + k + w = tl.load( + Weight + + expert * (2 * INTERMEDIATE_SIZE * HIDDEN_SIZE) + + n[:, None] * HIDDEN_SIZE + + kk[None, :], + (n[:, None] < 2 * INTERMEDIATE_SIZE) & (kk[None, :] < HIDDEN_SIZE), + 0, + ).to(tl.float32) + x = tl.load(HiddenStates + kk, kk < HIDDEN_SIZE, 0).to(tl.float32) + accumulator = accumulator + w * x[None, :] + y = tl.sum(accumulator, 1).to(tl.bfloat16).to(tl.float32) + gate, up = tl.split(tl.reshape(y, (BLOCK_N // 2, 2))) + silu = _silu_fast_math(gate) + column = tl.program_id(1) * (BLOCK_N // 2) + tl.arange(0, BLOCK_N // 2) + tl.store( + Intermediate + route * INTERMEDIATE_SIZE + column, + silu * up, + column < INTERMEDIATE_SIZE, + ) + + +@triton.jit +def _down_gemv_kernel( + Intermediate, + Weight, + TopkIds, + TopkWeights, + Output, + HIDDEN_SIZE: tl.constexpr, + INTERMEDIATE_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + route = tl.program_id(0) + n = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + k = tl.arange(0, BLOCK_K) + expert = tl.load(TopkIds + route).to(tl.int64) + w = tl.load( + Weight + + expert * (HIDDEN_SIZE * INTERMEDIATE_SIZE) + + n[:, None] * INTERMEDIATE_SIZE + + k[None, :], + (n[:, None] < HIDDEN_SIZE) & (k[None, :] < INTERMEDIATE_SIZE), + 0, + ).to(tl.float32) + a = tl.load( + Intermediate + route * INTERMEDIATE_SIZE + k, k < INTERMEDIATE_SIZE, 0 + ).to(tl.float32) + y = tl.sum(w * a[None, :], 1) * tl.load(TopkWeights + route) + tl.store(Output + route * HIDDEN_SIZE + n, y, n < HIDDEN_SIZE) + + +def beam_decode_gemv( + hidden_states: torch.Tensor, + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, +) -> torch.Tensor: + """Return weighted BF16 expert outputs for Beam's single-token TP4 path. + + Weights use the interleaved gate/up layout. Expert reduction, routed scaling + and TP communication remain with the caller. + """ + assert hidden_states.shape == (1, 5120) + assert w13_weight.shape == (256, 1280, 5120) + assert w2_weight.shape == (256, 5120, 640) + assert topk_ids.shape == topk_weights.shape == (1, 8) + assert hidden_states.dtype == w13_weight.dtype == w2_weight.dtype == torch.bfloat16 + assert topk_weights.dtype == torch.float32 + assert topk_ids.dtype in (torch.int32, torch.int64) + assert hidden_states.is_cuda + assert all( + t.device == hidden_states.device and t.is_contiguous() + for t in (hidden_states, w13_weight, w2_weight, topk_weights, topk_ids) + ) + + hidden_size = hidden_states.shape[1] + top_k = topk_ids.shape[1] + intermediate_size = w2_weight.shape[2] + intermediate = torch.empty( + (top_k, intermediate_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + output = torch.empty( + (1, top_k, hidden_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + _gate_up_gemv_kernel[(top_k, triton.cdiv(2 * intermediate_size, 8))]( + hidden_states, + w13_weight, + topk_ids, + intermediate, + hidden_size, + intermediate_size, + 8, + 256, + num_warps=4, + enable_fp_fusion=False, + ) + _down_gemv_kernel[(top_k, triton.cdiv(hidden_size, 8))]( + intermediate, + w2_weight, + topk_ids, + topk_weights, + output, + hidden_size, + intermediate_size, + 8, + triton.next_power_of_2(intermediate_size), + num_warps=4, + enable_fp_fusion=False, + ) + return output diff --git a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py index babf74af8e2c..b943ea6247f1 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -9,6 +9,7 @@ import triton.language as tl from sglang.kernels.jit.utils import is_arch_support_pdl +from sglang.kernels.ops.activation._triton_silu import _silu_fast_math from sglang.kernels.ops.quantization.fp8_kernel import ( per_token_group_quant_fp8, scaled_fp8_quant, @@ -647,36 +648,10 @@ def fused_moe_kernel( # tile; silu(gate) * up is applied in-register and only the # half-width activation is stored, eliminating intermediate_cache1 # and the standalone activation launch. - # - # The asm below is the bit-parity contract with the `silu_and_mul` this - # replaces, not an optimization: that kernel uses the fast-math - # intrinsics `__fdividef(x, 1 + __expf(-x))`, while Triton's operators - # lower to accurate expf and IEEE `div.rn`. The 1-2 ULP gap is enough to - # flip the stored bf16 (8 mantissa bits) on many inputs. The final - # multiply needs no asm -- plain fp32 multiply matches `mul.ftz.f32` - # except on denormals. Silu stays fp32 until the store, as in the - # reference; rounding it to bf16 first double-rounds and diverges. acc_pairs = tl.reshape(accumulator, (BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2)) gate_b, up_b = tl.split(acc_pairs) gate_f = gate_b.to(tl.float32) - # __expf(-x) == ex2.approx(x * -log2(e)); folding the negation into - # the constant (0fBFB8AA3B == -log2(e)) flips the sign exactly. - exp_neg = tl.inline_asm_elementwise( - "{ mul.ftz.f32 $0, $1, 0fBFB8AA3B; ex2.approx.ftz.f32 $0, $0; }", - "=f,f", - [gate_f], - dtype=tl.float32, - is_pure=True, - pack=1, - ) - silu_f = tl.inline_asm_elementwise( - "div.approx.ftz.f32 $0, $1, $2;", - "=f,f,f", - [gate_f, 1.0 + exp_neg], - dtype=tl.float32, - is_pure=True, - pack=1, - ) + silu_f = _silu_fast_math(gate_f) out_act = (silu_f * up_b.to(tl.float32)).to(compute_type) offs_half = pid_n * (BLOCK_SIZE_N // 2) + tl.arange(0, BLOCK_SIZE_N // 2) if c_sorted: diff --git a/python/sglang/kernels/ops/quantization/__init__.py b/python/sglang/kernels/ops/quantization/__init__.py index 67dd4d56b0c0..90c590b99a93 100644 --- a/python/sglang/kernels/ops/quantization/__init__.py +++ b/python/sglang/kernels/ops/quantization/__init__.py @@ -142,3 +142,11 @@ def sgl_per_token_group_quant_8bit( target="sglang.kernels.ops.quantization.act_quant_triton:act_quant", ) ) +register_kernel( + KernelSpec( + op="quantization.split_fp32_to_bf16x3", + backend=KernelBackend.TRITON, + target="sglang.kernels.ops.quantization.bf16x3_split:split_fp32_to_bf16x3", + capabilities=CUDA_ONLY, + ) +) diff --git a/python/sglang/kernels/ops/quantization/bf16x3_split.py b/python/sglang/kernels/ops/quantization/bf16x3_split.py new file mode 100644 index 000000000000..d931ff8b70cf --- /dev/null +++ b/python/sglang/kernels/ops/quantization/bf16x3_split.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Split FP32 weights into three BF16 parts for BF16x3 tensor-core GEMMs.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _split_fp32_to_bf16x3_kernel(w_ptr, out_ptr, numel, BLOCK: tl.constexpr): + offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offs < numel + w = tl.load(w_ptr + offs, mask=mask) + hi = w.to(tl.bfloat16) + residual = w - hi.to(tl.float32) + mid = residual.to(tl.bfloat16) + lo = (residual - mid.to(tl.float32)).to(tl.bfloat16) + tl.store(out_ptr + offs, hi, mask=mask) + tl.store(out_ptr + numel + offs, mid, mask=mask) + tl.store(out_ptr + 2 * numel + offs, lo, mask=mask) + + +def split_fp32_to_bf16x3( + w: torch.Tensor, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + """[N, K] fp32 -> [3N, K] bf16 (hi, mid, lo) blocks, each rounding what remains. + + hi + mid + lo == w for 2**-110 <= |w| < 2**128 - 2**119: below that the last + remainder drops under BF16's subnormals, above it hi rounds to infinity. + """ + if not ( + w.is_cuda and w.ndim == 2 and w.dtype == torch.float32 and w.is_contiguous() + ): + raise ValueError("w must be a contiguous CUDA fp32 [N, K] tensor") + n, k = w.shape + if out is None: + out = torch.empty(3 * n, k, dtype=torch.bfloat16, device=w.device) + elif not ( + out.shape == (3 * n, k) + and out.dtype == torch.bfloat16 + and out.device == w.device + and out.is_contiguous() + ): + raise ValueError("out must be a contiguous bf16 [3N, K] tensor on w's device") + block = 1024 + _split_fp32_to_bf16x3_kernel[(triton.cdiv(w.numel(), block),)]( + w, out, w.numel(), BLOCK=block + ) + return out diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 164af847d42a..9ceed02e40ec 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -223,6 +223,18 @@ def force_thresholds(heuristic): self._capture_started = False self.disabled = False + def supports_push(self, *, nbytes: int, num_blocks: int) -> bool: + """Whether the push plane can hold a payload and its synchronization blocks. + + This checks allocated capacity, not algorithm-selection thresholds or + tensor eligibility. Disabled communicators have no usable capacity. + """ + return ( + not self.disabled + and 0 < nbytes <= self.max_push_size + and 0 < num_blocks <= self._push_counter.numel() + ) + def reserve_push_capacity(self, nbytes: int) -> None: """Grow shared storage without changing AR dispatch or counter capacity. diff --git a/test/manual/kernels/moe/beam_decode_gemv_probes/README.md b/test/manual/kernels/moe/beam_decode_gemv_probes/README.md new file mode 100644 index 000000000000..c7782db385c9 --- /dev/null +++ b/test/manual/kernels/moe/beam_decode_gemv_probes/README.md @@ -0,0 +1,62 @@ +# Beam SM103 decode GEMV: manual correctness and latency checks + +`sglang.kernels.ops.moe.beam_decode_gemv` serves one BF16 token through +Beam's routed experts on SM103 (GB300). These files are manual checks, not CI +tests. Run them from the repository root with this tree's `python/` importable, +on one SM103 GPU, correctness first: + +```bash +python test/manual/kernels/moe/test_beam_decode_gemv.py +python test/manual/kernels/moe/beam_decode_gemv_probes/bench_decode_gemv.py +``` + +## Shapes and numerics + +Both checks use Beam's TP4 expert shard: 256 experts, hidden size 5120, +intermediate size 640, top-k 8, one contiguous BF16 token and FP32 routing +weights. `w13` is `[256, 1280, 5120]` with gate and up rows interleaved (row +`2j` is gate `j`, row `2j + 1` is up `j`); `w2` is `[256, 5120, 640]`. No +checkpoint, server or download is needed. + +`../test_beam_decode_gemv.py` compares the kernel's routing-weighted +per-expert outputs with an FP64 reference that keeps the kernel's BF16 rounding +points (gate/up, activation, weighted down projection), at +`rtol=1.6e-2, atol=1e-3`. It covers int32 and int64 expert ids eagerly and five +CUDA-graph replays whose inputs change between replays. + +## Latency probe + +`bench_decode_gemv.py` times two paths that return the same `[1, 5120]` BF16 +sum over the 8 routed experts: + +- `fused_moe`: the existing Triton `fused_experts_impl` with interleaved + SwiGLU (`fuse_swiglu_interleaved=True`), `routed_scaling_factor=1.0` and its + own expert combine; +- `decode_gemv`: `beam_decode_gemv` followed by an explicit FP32 expert sum + rounded to BF16. + +Both include the expert reduction. Neither includes TP communication or the +fused all-reduce boundary, so this is not a measurement of the model's +deferred-sum integration. + +Protocol: + +- each path is warmed up and captured in its own CUDA graph before timing; +- every timed replay follows an untimed write of at least 5x the device's L2 + size, because the 8 routed experts' weights (about 157 MB) are comparable to + L2 and would otherwise stay hot between replays; +- each run is the median of a fixed number of single replays timed with CUDA + events (`--repeats`, default 20); +- runs alternate `fused_moe, decode_gemv, decode_gemv, fused_moe` in at least + 5 blocks (`--blocks`, default 6). Every block's pair, the medians and the + median delta are printed with the GPU, SM count, torch and Triton versions. + +The probe prints the maximum absolute difference between the two summed +outputs for information only; correctness is the test above. + +## Status + +Both scripts are checked statically. The correctness test passed on one GB300. +The latency probe has not been run, so no speedup is claimed. The opt-in model +dispatch for this kernel lives in the Beam model part and moves with it when +that part is restacked. diff --git a/test/manual/kernels/moe/beam_decode_gemv_probes/bench_decode_gemv.py b/test/manual/kernels/moe/beam_decode_gemv_probes/bench_decode_gemv.py new file mode 100644 index 000000000000..e8b1d209ab60 --- /dev/null +++ b/test/manual/kernels/moe/beam_decode_gemv_probes/bench_decode_gemv.py @@ -0,0 +1,196 @@ +"""Single-token latency of the Beam decode GEMV against the Triton fused MoE. + +Manual and not CI-registered. Run the correctness test first, then this probe, +from the repository root on one SM103 GPU (GB300): + + python test/manual/kernels/moe/test_beam_decode_gemv.py + python test/manual/kernels/moe/beam_decode_gemv_probes/bench_decode_gemv.py + +Both timed paths return the same [1, 5120] BF16 sum over the 8 routed experts: +`fused_experts_impl` with interleaved SwiGLU and its own expert combine, and +`beam_decode_gemv` followed by an explicit FP32 expert sum rounded to BF16. +Neither path includes TP communication or the fused all-reduce boundary. +""" + +import argparse +import json +import os +import statistics + +import torch +import triton + +from sglang.kernels.ops.moe.beam_decode_gemv import beam_decode_gemv + +# Beam's experts, hidden size, TP4 intermediate size and top-k. +E, H, I, TOPK = 256, 5120, 640, 8 + + +def _init_runtime(): + # Same single-rank setup as test_fused_swiglu_epilogue.py: fused_experts_impl + # reads global server args and allocates its output under the TP group. + from sglang.srt.distributed.parallel_state import ( + init_distributed_environment, + initialize_model_parallel, + ) + from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler + from sglang.test.test_utils import publish_build_topology + + set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) + for key, value in ( + ("MASTER_ADDR", "127.0.0.1"), + ("MASTER_PORT", "29642"), + ("RANK", "0"), + ("WORLD_SIZE", "1"), + ("LOCAL_RANK", "0"), + ): + os.environ.setdefault(key, value) + init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="gloo") + publish_build_topology(tp_size=1, ep_size=1, pp_size=1) + initialize_model_parallel(backend="gloo") + + +def _shutdown_runtime(): + from sglang.srt.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + ) + + destroy_model_parallel() + destroy_distributed_environment() + + +def _capture(fn): + for _ in range(3): # Compile and autotune before capture. + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = fn() + graph.replay() + torch.cuda.synchronize() + return graph, out + + +def _time_us(graph, flush, repeats): + """Median of `repeats` single replays, each after an untimed L2 flush.""" + events = [] + for _ in range(repeats): + flush.zero_() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + graph.replay() + end.record() + events.append((start, end)) + torch.cuda.synchronize() + return statistics.median(s.elapsed_time(e) * 1e3 for s, e in events) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--blocks", type=int, default=6, help="ABBA blocks (>= 5)") + parser.add_argument("--repeats", type=int, default=20, help="replays per run") + parser.add_argument("--flush-l2-multiple", type=int, default=5) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + if args.blocks < 5 or args.flush_l2_multiple < 5: + parser.error("use at least 5 ABBA blocks and a flush of at least 5x L2") + + props = torch.cuda.get_device_properties(0) + capability = torch.cuda.get_device_capability(0) + if capability != (10, 3): + raise SystemExit(f"beam_decode_gemv targets SM103; this GPU is {capability}") + + _init_runtime() + try: + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( + fused_experts_impl, + ) + + torch.manual_seed(args.seed) + # Interleaved gate/up rows: row 2j is gate j and row 2j + 1 is up j. + w13 = (torch.randn(E, 2 * I, H, device="cuda") / H**0.5).bfloat16() + w2 = (torch.randn(E, H, I, device="cuda") / I**0.5).bfloat16() + hidden = torch.randn(1, H, device="cuda").bfloat16() + ids = torch.randperm(E, device="cuda")[:TOPK].view(1, TOPK).int() + weights = torch.softmax(torch.randn(1, TOPK, device="cuda"), -1) + + def fused(): + return fused_experts_impl( + hidden, + w13, + w2, + weights, + ids, + inplace=False, + # As fused_experts passes when every expert is local (EP=1). + filter_expert=False, + routed_scaling_factor=1.0, + fuse_swiglu_interleaved=True, + ) + + def gemv(): + out = beam_decode_gemv(hidden, w13, w2, weights, ids) + return out.sum(dim=1, dtype=torch.float32).bfloat16() + + graphs = {"fused_moe": _capture(fused), "decode_gemv": _capture(gemv)} + l2_bytes = props.L2_cache_size + flush = torch.empty( + args.flush_l2_multiple * l2_bytes // 4, dtype=torch.int32, device="cuda" + ) + print( + json.dumps( + { + "gpu": props.name, + "capability": capability, + "sm_count": props.multi_processor_count, + "torch": torch.__version__, + "triton": triton.__version__, + "l2_bytes": l2_bytes, + "flush_bytes": flush.numel() * 4, + # Only the routed experts' weights are read. + "routed_weight_bytes": TOPK * 3 * I * H * 2, + "max_abs_output_diff": ( + graphs["fused_moe"][1].float() + - graphs["decode_gemv"][1].float() + ) + .abs() + .max() + .item(), + } + ), + flush=True, + ) + + pairs = [] + for block in range(args.blocks): + times = {name: [] for name in graphs} + for name in ("fused_moe", "decode_gemv", "decode_gemv", "fused_moe"): + times[name].append(_time_us(graphs[name][0], flush, args.repeats)) + pair = {name: statistics.mean(t) for name, t in times.items()} + pair["delta_us"] = pair["decode_gemv"] - pair["fused_moe"] + pairs.append(pair) + print(json.dumps({"block": block, **pair}), flush=True) + print( + json.dumps( + { + "blocks": args.blocks, + "repeats": args.repeats, + "median_fused_moe_us": statistics.median( + p["fused_moe"] for p in pairs + ), + "median_decode_gemv_us": statistics.median( + p["decode_gemv"] for p in pairs + ), + "median_delta_us": statistics.median(p["delta_us"] for p in pairs), + } + ), + flush=True, + ) + finally: + _shutdown_runtime() + + +if __name__ == "__main__": + main() diff --git a/test/manual/kernels/moe/test_beam_decode_gemv.py b/test/manual/kernels/moe/test_beam_decode_gemv.py new file mode 100644 index 000000000000..2f1af20f61d5 --- /dev/null +++ b/test/manual/kernels/moe/test_beam_decode_gemv.py @@ -0,0 +1,72 @@ +"""Manual correctness test for the SM103 single-token Beam MoE decode GEMV. + +Run on one SM103 GPU (GB300): ``python test/manual/kernels/moe/test_beam_decode_gemv.py``. + +Compares the weighted per-expert outputs with an independent FP64 reference that +keeps the kernel's BF16 rounding points (gate/up, activation, weighted down +projection), eagerly and through CUDA-graph replays whose inputs change between +replays. +""" + +import sys + +import pytest +import torch + +from sglang.kernels.ops.moe.beam_decode_gemv import beam_decode_gemv + +# Beam's experts, hidden size and TP4 intermediate size. +E, H, I = 256, 5120, 640 +TOL = dict(rtol=1.6e-2, atol=1e-3) + + +@pytest.fixture(scope="module") +def weights(): + torch.manual_seed(0) + # Interleaved gate/up rows: row 2j is gate j and row 2j + 1 is up j. + w13 = (torch.randn(E, 2 * I, H, device="cuda") / H**0.5).bfloat16() + w2 = (torch.randn(E, H, I, device="cuda") / I**0.5).bfloat16() + return w13, w2 + + +def _reference(h, ids, tw, w13, w2): + out = [] + for r in range(8): + e = int(ids[0, r]) + y = (w13[e].double() @ h[0].double()).bfloat16().float() + gate, up = y[0::2], y[1::2] + act = (gate / (1 + torch.exp(-gate)) * up).bfloat16().double() + out.append(((w2[e].double() @ act).float() * tw[0, r]).bfloat16()) + return torch.stack(out)[None] + + +def _refill(seed, h, ids, tw): + g = torch.Generator(device="cuda").manual_seed(seed) + h.copy_(torch.randn(1, H, device="cuda", generator=g)) + ids.copy_(torch.randperm(E, device="cuda", generator=g)[:8].view(1, 8)) + tw.copy_(torch.softmax(torch.randn(1, 8, device="cuda", generator=g), -1)) + + +def test_matches_fp64_reference_under_graph_replay(weights): + w13, w2 = weights + h = torch.empty(1, H, device="cuda", dtype=torch.bfloat16) + ids = torch.empty(1, 8, device="cuda", dtype=torch.int64) + tw = torch.empty(1, 8, device="cuda") + _refill(0, h, ids, tw) + expected = _reference(h, ids, tw, w13, w2) + torch.testing.assert_close(beam_decode_gemv(h, w13, w2, tw, ids), expected, **TOL) + torch.testing.assert_close( + beam_decode_gemv(h, w13, w2, tw, ids.int()), expected, **TOL + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = beam_decode_gemv(h, w13, w2, tw, ids) + for seed in range(1, 6): + _refill(seed, h, ids, tw) + graph.replay() + torch.testing.assert_close(out, _reference(h, ids, tw, w13, w2), **TOL) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/activation/test_fused_op_gpu_parity.py b/test/registered/kernels/ops/activation/test_fused_op_gpu_parity.py index 6a709ddb18b5..ba434f569c6a 100644 --- a/test/registered/kernels/ops/activation/test_fused_op_gpu_parity.py +++ b/test/registered/kernels/ops/activation/test_fused_op_gpu_parity.py @@ -5,13 +5,15 @@ tolerance. New backends are picked up automatically. """ +import math + import pytest import torch from sglang.kernels.spec import KernelBackend from sglang.test.ci.ci_register import register_cuda_ci -register_cuda_ci(est_time=4, stage="extra-a", runner_config="1-gpu-small") +register_cuda_ci(est_time=15, stage="extra-a", runner_config="1-gpu-small") _SKIP = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE} @@ -45,6 +47,101 @@ def test_silu_and_mul(dtype, shape): _close(op.forward(x, backend=b), ref, dtype, f"{op.op} {b.value}") +def _exact(got, ref, msg): + torch.testing.assert_close(got, ref, atol=0, rtol=0, msg=msg) + + +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_beam_attention_gate(dtype, with_bias): + """Gate + bias, its sigmoid and the product round like eager, bit for bit.""" + from sglang.kernels.ops.elementwise.beam_fused import BeamAttentionGateOp + + op = BeamAttentionGateOp() + # 256 consecutive values per sign around 7/1024, where sigmoid lands next to + # 0.5 and a non-eager sigmoid (approximate exp/div, Inductor's) moves FP16. + center = torch.tensor(7 / 1024, dtype=dtype).view(torch.int16) + gate = (center + torch.arange(-128, 128, dtype=torch.int16)).view(dtype) + gate = torch.cat([gate, -gate]).repeat(2, 1).cuda() + generator = torch.Generator().manual_seed(5) + # A unit row exposes the sigmoid's rounding; a random row checks the product. + attention = torch.stack([torch.ones(512), torch.randn(512, generator=generator)]) + bias = torch.randn(4, generator=generator) if with_bias else torch.zeros(4) + attention, bias = attention.to("cuda", dtype), bias.to("cuda", dtype) + expected = op.forward_native(attention, gate, bias, 128) + for b in _eligible(op): + _exact( + op(attention, gate, bias, 128, backend=b), expected, f"{op.op} {b.value}" + ) + + +def _bf16_ulps(got, ref): + """Largest distance in representable BF16 steps, counted across zero.""" + + def ordered(t): + bits = t.view(torch.int16).int() + return torch.where(bits < 0, -(bits & 0x7FFF), bits) + + return (ordered(got) - ordered(ref)).abs().max().item() + + +@pytest.mark.parametrize("with_next", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_beam_sandwich_residual_rmsnorm(dtype, with_next): + """BF16 stays within one ULP of eager; FP16 cancellation is bounded, not one ULP.""" + from sglang.kernels.ops.layernorm.beam_fused import ( + BeamSandwichResidualRMSNormOp, + _rmsnorm_native, + ) + + op = BeamSandwichResidualRMSNormOp() + # The recorded 64-row fixture draws the gate inputs first; keeping those draws + # keeps its FP16 residual-cancellation points. + torch.manual_seed(64) + torch.randn(64, 2 * 5120 + 2048, device="cuda", dtype=dtype) + torch.randn(40, device="cuda", dtype=dtype) + x = torch.randn(64, 5120, device="cuda", dtype=dtype) + r = torch.randn_like(x) + w1 = torch.randn(5120, device="cuda", dtype=dtype) + w2 = torch.randn_like(w1) + args = (x, r, w1, w2 if with_next else None, 1e-6, (0.3, 0.9)) + expected_out, expected_combined = op.forward_native(*args) + for b in _eligible(op): + out, combined = op(*args, backend=b) + msg = f"{op.op} {b.value}" + if not with_next: + assert out is None, msg + if dtype == torch.bfloat16: + for got, ref in ((combined, expected_combined), (out, expected_out)): + if ref is None: + continue + torch.testing.assert_close(got, ref, atol=0, rtol=2**-7, msg=msg) + assert _bf16_ulps(got, ref) <= 1, msg + else: + fp16_tol = dict(atol=2**-11, rtol=2**-10, msg=msg) + torch.testing.assert_close(combined, expected_combined, **fp16_tol) + if with_next: + # The next norm reads the combined stream, so bound it on the + # kernel's own combined output rather than eager's. + ref_out = _rmsnorm_native(combined, w2, 1e-6) + torch.testing.assert_close(out, ref_out, **fp16_tol) + + +def test_beam_residual_cancellation(): + """Each scaled term rounds before the sum, so exactly opposite terms cancel.""" + from sglang.kernels.ops.layernorm.beam_fused import BeamSandwichResidualRMSNormOp + + op = BeamSandwichResidualRMSNormOp() + branch = torch.ones(1, 5120, device="cuda", dtype=torch.bfloat16) + post = torch.full((5120,), 5, device="cuda", dtype=branch.dtype) + scales = (1 / math.sqrt(26), 5 / math.sqrt(26)) + args = (branch, -branch, post, torch.ones_like(post), 1e-6, scales) + expected = op.forward_native(*args) + assert torch.count_nonzero(expected[1]).item() == 0 + for b in _eligible(op): + _exact(op(*args, backend=b), expected, f"{op.op} {b.value}") + + if __name__ == "__main__": import sys diff --git a/test/registered/kernels/ops/communication/test_beam_all_reduce.py b/test/registered/kernels/ops/communication/test_beam_all_reduce.py new file mode 100644 index 000000000000..6ff9d4087542 --- /dev/null +++ b/test/registered/kernels/ops/communication/test_beam_all_reduce.py @@ -0,0 +1,210 @@ +"""Beam all-reduce boundary numerics and shared-plane CUDA graph replay.""" + +import os +import unittest + +import torch +import torch.distributed as dist + +import sglang.srt.distributed.parallel_state as ps +from sglang.kernels.ops.communication.all_reduce import ( + AllReduceAlgo, + Communicator, + PushPlane, +) +from sglang.kernels.ops.communication.beam_all_reduce import ( + _cluster_size, + _module, + beam_all_reduce_norm, +) +from sglang.srt.distributed.comm_cleanup import register_comm_cleanup +from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, +) +from sglang.srt.model_executor.cuda_graph_config import CudaGraphConfig +from sglang.srt.runtime_context import get_context, get_parallel +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=180, stage="extra-b", runner_config="8-gpu-h200") + +ROWS = (1, 8, 16, 32, 64, 128) + + +def _precompile(num_gpus): + for world in num_gpus: + for cluster in {_cluster_size(world, rows) for rows in ROWS}: + for topk in (1, 8): + _module(world, 5120, cluster, topk) + + +class TestBeamAllReduce(CustomTestCase): + @classmethod + def setUpClass(cls): + override = get_context().override_server_args( + cuda_graph_config=CudaGraphConfig() + ) + override.install() + cls.addClassCleanup(override.restore) + # Global rank for the TP layout, local rank for the device (2-node TP8). + rank = int(os.environ["RANK"]) + local = int(os.environ["LOCAL_RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(local) + dist.init_process_group(backend="gloo") + ps._WORLD = coord = ps.init_world_group(list(range(world)), local, "nccl") + get_parallel().override_permanently( + world_group=coord, + tp_group=coord, + attn_tp_group=coord, + tp_size=world, + tp_rank=rank, + attn_tp_size=world, + attn_tp_rank=rank, + moe_tp_size=world, + moe_tp_rank=rank, + ) + cls.group = coord.device_group + cls.comm = CustomAllReduceV2( + coord.cpu_group, + torch.device("cuda", local), + max_push_size=2 * 1024 * 1024, + max_push_blocks=128, + ) + if cls.comm.disabled: + raise RuntimeError( + "Beam collective test requires an NVLink V2 communicator" + ) + register_comm_cleanup(cls.comm) + # Exercise the shared phase counters, regardless of platform tuning. + cls.comm.override_algo = AllReduceAlgo.ONE_SHOT_PUSH + # The same workspaces and counters without a multicast address take the + # per-peer store path; the multicast plane is the control where it exists. + unicast = Communicator( + push=PushPlane( + cls.comm.rank, + cls.comm.world_size, + workspaces=cls.comm._push_workspaces, + counter=cls.comm._push_counter.view(-1, 1).view(torch.uint8), + ), + pull=cls.comm.obj.pull, + ) + cls.planes = [("unicast", unicast)] + if cls.comm.has_multicast: + cls.planes.append(("multicast", cls.comm.obj)) + + @classmethod + def tearDownClass(cls): + cls.planes = [] # Drop the extra pointer holder before the workspaces. + cls.comm.close() + dist.destroy_process_group() + + @staticmethod + def norm(x, w): + x = x.float() + return ( + x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-6) * w.float() + ).bfloat16() + + @staticmethod + def local_sum(x): + if x.ndim == 2: + return x + summed = torch.zeros_like(x[:, 0], dtype=torch.float32) + for expert in range(8): + summed += x[:, expert].float() + return (summed * 2.5).bfloat16() + + def reference(self, x, residual, post, next_weight): + x = self.local_sum(x) + peers = [torch.empty_like(x) for _ in range(dist.get_world_size())] + dist.all_gather(peers, x, group=self.group) + reduced = torch.zeros_like(x, dtype=torch.float32) + for peer in peers: + reduced += peer.float() + reduced = reduced.bfloat16() + combined = self.norm(reduced, post) * 0.13 + residual * 0.97 + return ( + (None if next_weight is None else self.norm(combined, next_weight)), + combined, + reduced, + ) + + def test_rounding_graph_replay_and_interleaved_collectives(self): + if dist.get_world_size() == 8: + # Both TP8 launch shapes (5-CTA clusters and one CTA per row) must run. + self.assertEqual({_cluster_size(8, rows) for rows in ROWS}, {1, 5}) + owned = self.comm.obj + try: + for name, plane in self.planes: + # The interleaved generic all-reduce pushes through the same plane. + self.comm.obj = plane + try: + self._check_plane() + except AssertionError as error: + raise AssertionError(f"{name} push plane: {error}") from error + finally: + self.comm.obj = owned + + def _check_plane(self): + for topk in (1, 8): + for rows in ROWS: + torch.manual_seed(42 + rows) + shape = (rows, 5120) if topk == 1 else (rows, 8, 5120) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + x *= dist.get_rank() + 1 + # Consecutive zero pairs exercise the 32-bit Lamport marker. + x[..., :16] = 0 + residual = torch.randn(rows, 5120, device="cuda", dtype=torch.bfloat16) + post = torch.randn(5120, device="cuda", dtype=torch.bfloat16) + weight = torch.randn_like(post) + for next_weight in (None, weight): + + def fused(): + return beam_all_reduce_norm( + self.comm.obj, + x, + residual, + post, + next_weight, + 1e-6, + (0.13, 0.97), + routed_scale=2.5, + ) + + fused() # Compile before capture. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + normed, combined = fused() + reduced = self.comm.custom_all_reduce(self.local_sum(x)) + normed2, combined2 = fused() + # Reuse graph storage with different payloads and several epochs. + for change in (1.0, -0.5, 2.0): + x.mul_(change) + for _ in range(5): + graph.replay() + expected_norm, expected, expected_reduced = self.reference( + x, residual, post, next_weight + ) + torch.testing.assert_close( + combined, expected, atol=0.016, rtol=0.016 + ) + torch.testing.assert_close(combined2, combined, atol=0, rtol=0) + torch.testing.assert_close( + reduced, expected_reduced, atol=0, rtol=0 + ) + if next_weight is not None: + torch.testing.assert_close( + normed, expected_norm, atol=0.032, rtol=0.016 + ) + torch.testing.assert_close(normed2, normed, atol=0, rtol=0) + + +if __name__ == "__main__": + if os.environ.get("_IS_TEST_MULTIGPU_SGLANG_JIT_KERNEL"): + unittest.main() + else: + multigpu_pytest_main( + __name__, __file__, num_gpus=(2, 4, 8), pre_launch_fn=_precompile + ) diff --git a/test/registered/kernels/ops/gemm/test_router_gemv.py b/test/registered/kernels/ops/gemm/test_router_gemv.py new file mode 100644 index 000000000000..a1b5747429d5 --- /dev/null +++ b/test/registered/kernels/ops/gemm/test_router_gemv.py @@ -0,0 +1,74 @@ +"""Tests for the router-logits GEMV and the BF16x3 router GEMM with FP32 gates.""" + +import sys + +import pytest +import torch +import torch.nn.functional as F + +from sglang.kernels.ops.gemm.router_gemv import ( + router_gemm_bf16x3_cublas, + router_gemv, + router_gemv_supported, +) +from sglang.kernels.ops.quantization.bf16x3_split import split_fp32_to_bf16x3 +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, + reason="CUDA required", +) + +# Beam router shape: 256 experts over a 5120 hidden size. +N, K = 256, 5120 + + +def _inputs(m: int): + generator = torch.Generator(device="cuda").manual_seed(m) + x = torch.randn(m, K, device="cuda", generator=generator).bfloat16() + w = 0.02 * torch.randn(N, K, device="cuda", generator=generator) + return x, w + + +def _max_error(out: torch.Tensor, x: torch.Tensor, w: torch.Tensor) -> float: + return (out.double() - x.double() @ w.double().t()).abs().max().item() + + +def test_bf16x3_split_reconstructs_fp32_exactly(): + _, w = _inputs(1) + # Ends of the exact range 2**-110 <= |w| < 2**128 - 2**119, full mantissas. + edges = torch.tensor( + [ + 2.0**-110 * (1 + 2.0**-23), + 2.0**-110 * (2 - 2.0**-23), + 2.0**128 - 2.0**119 - 2.0**104, + ], + device="cuda", + ).repeat(K)[:K] + w = torch.cat([w, w * 1e30, w * 1e-20, -w, edges[None], -edges[None]]) + w3 = split_fp32_to_bf16x3(w) + n = w.shape[0] + assert w3.dtype == torch.bfloat16 and w3.shape == (3 * n, K) + hi, mid, lo = (part.float() for part in w3.split(n)) + assert torch.equal(hi + mid + lo, w) + + +@pytest.mark.parametrize("m", [1, 8]) +def test_fp32_gate_gemv_matches_fp64(m): + x, w = _inputs(m) + assert router_gemv_supported(x, w) + assert _max_error(router_gemv(x, w), x, w) < 2e-6 + + +@pytest.mark.parametrize("m", [9, 64]) +def test_bf16x3_router_gemm_matches_fp64(m): + x, w = _inputs(m) + assert _max_error(router_gemm_bf16x3_cublas(x, w), x, w) < 1e-4 + # A BF16 read of the FP32 gate misses the same bound. + assert _max_error(F.linear(x.float(), w.bfloat16().float()), x, w) > 1e-4 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/gemm/test_tiny_gemm.py b/test/registered/kernels/ops/gemm/test_tiny_gemm.py index 4d4f2730be6a..f80c483c56d7 100644 --- a/test/registered/kernels/ops/gemm/test_tiny_gemm.py +++ b/test/registered/kernels/ops/gemm/test_tiny_gemm.py @@ -13,7 +13,7 @@ from sglang.kernels.ops.gemm.tiny_gemm import can_use_tiny_gemm, tiny_gemm_bf16 from sglang.test.ci.ci_register import register_cuda_ci -register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_cuda_ci(est_time=300, stage="nightly", runner_config="1-gpu-large") # One kernel is built per m in [1, MAX_M], so hold max_m fixed across the sweep: @@ -54,5 +54,21 @@ def test_tiny_gemm(n, k, num_tokens, out_dtype): torch.testing.assert_close(out.float(), ref.float(), atol=1e-2, rtol=1e-2) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 16]) +def test_tiny_gemm_fp32_weights(num_tokens): + """FP32 router weights are read in FP32: a BF16 read misses FP64 by ~1e-3.""" + if is_hip_runtime() or get_jit_cuda_arch().major < 9: + pytest.skip("SM90+ required") + + torch.manual_seed(42) + # Beam's router gate: 256 experts over a 5120 hidden size. + w = torch.randn(256, 5120, device="cuda") / 5120**0.5 + x = torch.randn(num_tokens, 5120, dtype=torch.bfloat16, device="cuda") + out = tiny_gemm_bf16(x, w, out_dtype=torch.float32, max_m=MAX_M) + ref = (x.double() @ w.double().T).float() + torch.testing.assert_close(out, ref, atol=2e-6, rtol=2e-5) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py b/test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py index ea790dd0200b..19db00594666 100644 --- a/test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py +++ b/test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py @@ -1,6 +1,7 @@ from unittest.mock import Mock import pytest +import torch from sglang.srt.distributed.device_communicators import custom_all_reduce_v2 from sglang.test.ci.ci_register import register_cpu_ci @@ -88,5 +89,28 @@ def is_vmm_backed(device): intra_node_capability.assert_not_called() +@pytest.mark.parametrize( + ("disabled", "has_multicast", "nbytes", "num_blocks", "expected"), + [ + (False, True, 1024, 4, True), + (False, True, 1025, 4, False), + (False, True, 1024, 5, False), + (False, True, 0, 1, False), + (False, False, 1024, 4, True), + (True, None, 1024, 4, False), + ], +) +def test_supports_push(disabled, has_multicast, nbytes, num_blocks, expected): + """Push is offered within the plane's inclusive capacity, with or without multicast.""" + comm = object.__new__(custom_all_reduce_v2.CustomAllReduceV2) + comm.disabled = disabled + # A disabled communicator never builds its planes, so these stay unset. + if not disabled: + comm.has_multicast = has_multicast + comm.max_push_size = 1024 + comm._push_counter = torch.zeros(4) + assert comm.supports_push(nbytes=nbytes, num_blocks=num_blocks) is expected + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/kernels/test_kernels_namespace.py b/test/registered/unit/kernels/test_kernels_namespace.py index fc34d43e2cc2..9822b8bc04c8 100644 --- a/test/registered/unit/kernels/test_kernels_namespace.py +++ b/test/registered/unit/kernels/test_kernels_namespace.py @@ -249,6 +249,14 @@ def test_kimi_k3_kernels_are_inventoried_by_operator(op, backend, device): assert not K.capabilities_satisfied(spec.capabilities, _CPU) +def test_router_gemv_is_inventoried_for_cuda_and_hip(): + """CUDA serves FP32 gates and ROCm gfx950 BF16 gates through one Triton GEMV.""" + spec = K.select_kernel("gemm.router_gemv", backend=KernelBackend.TRITON) + for platform in (_SM90, _HIP): + assert K.capabilities_satisfied(spec.capabilities, platform) + assert not K.capabilities_satisfied(spec.capabilities, _CPU) + + def _readme() -> str: return (Path(K.__file__).resolve().parent / "README.md").read_text()