From 46c7feddf4bc5204a8332c87ed14bfab39537318 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:50:43 -0700 Subject: [PATCH 01/40] [Chore] docker: Add Beam CUDA image (1/5) Add a Beam serving image on top of the v0.5.22 release image. It clones the upstream beam branch and adds the pins that branch needs, without changing docker/Dockerfile. Split from #44747. --- docker/beam/beam.Dockerfile | 47 +++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 docker/beam/beam.Dockerfile diff --git a/docker/beam/beam.Dockerfile b/docker/beam/beam.Dockerfile new file mode 100644 index 000000000000..1bb93a209d57 --- /dev/null +++ b/docker/beam/beam.Dockerfile @@ -0,0 +1,47 @@ +# syntax=docker/dockerfile:1 +# Beam serving image (CUDA). Build from the repository root; application source +# is cloned from upstream: +# docker buildx build --platform linux/amd64 --load \ +# -f docker/beam/beam.Dockerfile -t beam:local-test . +FROM lmsysorg/sglang:v0.5.22@sha256:e3a4f06001a29542fc4a8638536d9ffc5eb9c9942b08603c09b19a56eae4a1ae AS base + +SHELL ["/bin/bash", "-o", "pipefail", "-c"] + +# Replace the stock SGLang with the upstream beam branch. Pass a commit SHA as +# BEAM_REF to pin the source of a published image. +ARG BEAM_REF=beam +WORKDIR /sgl-workspace +RUN git init -q /sgl-workspace/sglang-beam \ + && git -C /sgl-workspace/sglang-beam fetch -q --depth 1 \ + https://github.com/sgl-project/sglang.git "${BEAM_REF}" \ + && git -C /sgl-workspace/sglang-beam checkout -q --detach FETCH_HEAD \ + && uv pip uninstall sglang \ + && uv pip install --no-deps -e /sgl-workspace/sglang-beam/python + +# Pins the beam branch adds on top of the v0.5.22 release: OpenAI 2.8.0 for the +# Responses shell types, the cuTile compiler for the Beam router GEMM, and a +# gRPC compiler that supports the Protobuf 7 dependency set. Record the new +# versions so later installs keep them. +RUN sed -i -E '/^(openai|nvidia-cuda-tileiras|grpcio-tools)==/d' /sgl-workspace/constraints.txt \ + && uv pip install -c /sgl-workspace/constraints.txt \ + openai==2.8.0 nvidia-cuda-tileiras==13.4.92 grpcio-tools==1.84.0 \ + && uv pip freeze | grep -E '^(openai|nvidia-cuda-tileiras|grpcio-tools)==' \ + >> /sgl-workspace/constraints.txt \ + && python3 -c "from importlib.metadata import version; assert version('openai') == '2.8.0'; assert version('nvidia-cuda-tileiras') == '13.4.92'; assert version('grpcio-tools') == '1.84.0'" + +# The release image's unbounded prerelease request can resolve to the empty +# ai-dynamo==0.0.0a0 placeholder, so pin a stable release. ai-dynamo-runtime +# caps pydantic below the frozen pin: release only the pydantic pins for this +# solve and record the versions it selects. Layered images that bring their +# own Dynamo can pass INSTALL_DYNAMO=0. +ARG INSTALL_DYNAMO=1 +RUN if [ "${INSTALL_DYNAMO}" = "1" ]; then \ + sed -i -E '/^pydantic(-core)?==/d' /sgl-workspace/constraints.txt \ + && uv pip install -c /sgl-workspace/constraints.txt \ + --index-strategy unsafe-best-match \ + --extra-index-url https://pypi.nvidia.com/ ai-dynamo==1.5.1 inflect \ + && uv pip freeze | grep -E '^pydantic(-core)?==' >> /sgl-workspace/constraints.txt \ + && python3 -c "import dynamo.runtime, inflect"; \ + fi + +WORKDIR /sgl-workspace/sglang-beam From 65f3a8bcc5a05add2e30fe608a5c94d2c862b6dc Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:50:45 -0700 Subject: [PATCH 02/40] [Feat] kernels: Add Beam fused ops, router GEMMs and all-reduce boundary (2/5) Add the Beam fused projection/norm/residual and attention-gate ops, the FP32 router GEMV, BF16x3 and cuTile router GEMMs, the Beam all-reduce boundary kernel, and their kernel tests and benchmarks. Split from #44747. Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Caihua Li Co-authored-by: Florent Gbelidji Co-authored-by: iaroslav-reflection Co-authored-by: Jinze Xue Co-authored-by: mloh Co-authored-by: neilkg <33635204+neilkg@users.noreply.github.com> Co-authored-by: thibautbar --- python/pyproject.toml | 1 + .../jit/csrc/distributed/beam_all_reduce.cuh | 252 ++++++++++++++++++ .../kernels/jit/csrc/gemm/tiny_gemm.cuh | 53 ++-- .../include/sgl_kernel/gemm/dot_product.cuh | 10 + .../kernels/ops/activation/_triton_silu.py | 29 ++ .../kernels/ops/communication/__init__.py | 9 + .../ops/communication/beam_all_reduce.py | 72 +++++ .../kernels/ops/elementwise/__init__.py | 80 +++++- .../kernels/ops/elementwise/beam_fused.py | 96 +++++++ python/sglang/kernels/ops/gemm/__init__.py | 8 +- .../kernels/ops/gemm/router_gemm_cutile.py | 217 +++++++++++++++ python/sglang/kernels/ops/gemm/router_gemv.py | 93 ++++++- python/sglang/kernels/ops/gemm/tiny_gemm.py | 35 ++- .../sglang/kernels/ops/kvcache/zero_pages.py | 9 +- .../sglang/kernels/ops/layernorm/__init__.py | 82 ++++++ .../kernels/ops/layernorm/beam_fused.py | 177 ++++++++++++ python/sglang/kernels/ops/moe/__init__.py | 16 ++ .../kernels/ops/moe/beam_decode_gemv.py | 150 +++++++++++ .../ops/moe/fused_moe_triton_kernels.py | 29 +- .../custom_all_reduce_v2.py | 12 + test/manual/kernels/bench_fp32_router.py | 67 +++++ .../attention/bench_trtllm_mha_mixed.py | 142 ++++++++++ .../elementwise/bench_beam_attention_gate.py | 40 +++ .../gemm/bench_router_gemm_cutile.py | 111 ++++++++ .../bench_beam_sandwich_residual_rmsnorm.py | 32 +++ .../elementwise/test_fused_op_gpu_parity.py | 147 ++++++++++ .../ops/gemm/test_router_gemm_cutile.py | 237 ++++++++++++++++ .../kernels/ops/gemm/test_router_gemv.py | 91 +++++++ .../kernels/ops/gemm/test_tiny_gemm_fp32.py | 97 +++++++ 29 files changed, 2323 insertions(+), 71 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh create mode 100644 python/sglang/kernels/ops/activation/_triton_silu.py create mode 100644 python/sglang/kernels/ops/communication/beam_all_reduce.py create mode 100644 python/sglang/kernels/ops/elementwise/beam_fused.py create mode 100644 python/sglang/kernels/ops/gemm/router_gemm_cutile.py create mode 100644 python/sglang/kernels/ops/layernorm/beam_fused.py create mode 100644 python/sglang/kernels/ops/moe/beam_decode_gemv.py create mode 100644 test/manual/kernels/bench_fp32_router.py create mode 100644 test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py create mode 100644 test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py create mode 100644 test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py create mode 100644 test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py create mode 100644 test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py create mode 100644 test/registered/kernels/ops/gemm/test_router_gemm_cutile.py create mode 100644 test/registered/kernels/ops/gemm/test_router_gemv.py create mode 100644 test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py diff --git a/python/pyproject.toml b/python/pyproject.toml index c1cff46ed9c4..fbb30fb2e434 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -52,6 +52,7 @@ dependencies = [ "ninja", "numba==0.65.1", "numpy", + "nvidia-cuda-tileiras==13.4.92", "nvidia-cutlass-dsl[cu13]==4.8.0", "nvidia-mathdx==25.6.0", "nvidia-ml-py", 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..586b0a2f920f --- /dev/null +++ b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh @@ -0,0 +1,252 @@ +// Copyright 2026 SGLang Team. Licensed under the Apache License, Version 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; +}; + +/// 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; + } + 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 their rank-local payloads. + 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); + } + 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); + 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..7e0b775eea17 100644 --- a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh @@ -13,13 +13,14 @@ #include #include #include +#include #include 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 +33,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 +65,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 +83,12 @@ 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]; + // FP32 weights may have been updated by the preceding graph node. + if constexpr (std::is_same_v) PDLWaitPrimary(); + 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 @@ -94,7 +97,7 @@ TINY_GEMM_KERNEL void tiny_n_gemm_kernel(const TinyGEMMParams params) { } } - PDLWaitPrimary(); + if constexpr (!std::is_same_v) PDLWaitPrimary(); vec_t xv[M][kUnroll]; #pragma unroll @@ -155,7 +158,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 +181,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 +206,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 +228,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 +267,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..1f368e70ca6f --- /dev/null +++ b/python/sglang/kernels/ops/activation/_triton_silu.py @@ -0,0 +1,29 @@ +# Copyright 2026 SGLang Team. Licensed under the Apache License, Version 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): + # Preserve the fused MoE epilogue's __fdividef(x, 1 + __expf(-x)): accurate + # exp/div can move the result across a BF16 rounding boundary. silu_and_mul + # matches it only before SM100, where it is built with --use_fast_math. + # Keep SiLU in FP32; callers apply the multiply before the final BF16 store. + 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 93b52b67af38..78c1efc69cbc 100644 --- a/python/sglang/kernels/ops/communication/__init__.py +++ b/python/sglang/kernels/ops/communication/__init__.py @@ -43,3 +43,12 @@ ) ) 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))}), + ) +) 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..0b63c52f0898 --- /dev/null +++ b/python/sglang/kernels/ops/communication/beam_all_reduce.py @@ -0,0 +1,72 @@ +"""Beam BF16 all-reduce with FP32 sandwich normalization and rounded residuals.""" + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit + + +@cache_once +def _module(world_size: int, hidden_size: int, cluster: int = 1, topk: int = 1): + 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", + ) + ], + extra_cuda_cflags=["--fmad=false"], + ) + + +def _cluster_size(world_size: int, rows: int, topk: int) -> int: + # TP8 benefits from distributing the smallest rows across SMs; at larger + # batches a single CTA per row avoids cluster synchronization overhead. + if world_size == 8: + return 5 if rows <= (16 if topk == 8 else 8) else 1 + return 5 if world_size == 2 and rows == 1 and topk == 8 else 1 + + +def beam_all_reduce_norm( + comm, branch, residual, post_weight, next_weight, eps, scales, *, routed_scale=1.0 +): + """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), + 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 a0c3aecab7b3..d084761c43a0 100644 --- a/python/sglang/kernels/ops/elementwise/__init__.py +++ b/python/sglang/kernels/ops/elementwise/__init__.py @@ -8,10 +8,23 @@ they logically belong to. """ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sglang.kernels.fused_op import BaseFusedOp, register_fused_op from sglang.kernels.registry import register_kernel -from sglang.kernels.spec import CapabilityRequirement, KernelBackend, KernelSpec +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, +) -__all__ = [] +if TYPE_CHECKING: + import torch + +__all__ = ["BeamAttentionGateOp"] for _fn in ("hc_combine_gate", "hc_combine_apply", "hc_combine_apply_norm"): register_kernel( @@ -51,3 +64,66 @@ capabilities=frozenset({CapabilityRequirement.CUDA}), ) ) + + +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}} + format_signature = FormatSignature(supported_dtypes=("float16", "bfloat16")) + + 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): + import torch + + if self._compiled_native is None: + # Register and warm the opaque op before tracing its dispatch wrapper. + self.forward_triton(*args, **kwargs) + self._compiled_native = torch.compile( + self.forward_triton, 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: + import torch + + heads = gate_bias.numel() + multiplier = torch.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: + import torch + + if not attention_output.is_cuda or attention_output.dtype not in ( + torch.float16, + torch.bfloat16, + ): + return self.forward_native(attention_output, gate, gate_bias, head_dim) + from sglang.kernels.ops.elementwise.beam_fused import fused_attention_gate + + return fused_attention_gate(attention_output, gate, gate_bias, head_dim) + + +_BEAM_ATTENTION_GATE = register_fused_op( + BeamAttentionGateOp(), __name__, "_BEAM_ATTENTION_GATE" +) 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..bcdd80757c9a --- /dev/null +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -0,0 +1,96 @@ +# Copyright 2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""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.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) + + +@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") + 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 diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index 47b4a0e4addd..46b8b1ed3b80 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -241,10 +241,10 @@ def forward_aot( target="sglang.kernels.ops.gemm.tiny_gemm:tiny_gemm_bf16", capabilities=_CUDA, format_signature=FormatSignature( - supported_dtypes=("bfloat16",), - description="skinny [m, k] @ [n, k].T with m in [1, max_m]; serves MoE routers and gate projections", + supported_dtypes=("bfloat16", "float32"), + description="BF16 activations, BF16 weights or FP32 weights for N < K; m in [1, max_m]", ), - description="Tiny bf16 GEMM (sglang.kernels.jit, JIT-only).", + description="Tiny BF16-activation GEMM (sglang.kernels.jit, JIT-only).", ) ) register_kernel( @@ -384,7 +384,7 @@ def tiny_gemm_bf16( out_dtype: Optional[torch.dtype] = None, max_m: int = 16, ) -> torch.Tensor: - """Tiny bf16 GEMM ``x[m, k] @ w[n, k].T``, support bf16/fp32 out""" + """BF16 activations, BF16/FP32 output; FP32 weights are supported for N < K.""" impl = get_kernel("gemm.tiny_gemm", KernelBackend.JIT) return impl(x, w, out, out_dtype=out_dtype, max_m=max_m) diff --git a/python/sglang/kernels/ops/gemm/router_gemm_cutile.py b/python/sglang/kernels/ops/gemm/router_gemm_cutile.py new file mode 100644 index 000000000000..ca5121f91add --- /dev/null +++ b/python/sglang/kernels/ops/gemm/router_gemm_cutile.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +"""cuTile router GEMM: BF16 activations × FP32 weights → FP32 logits. + +The caller splits each FP32 weight into three BF16 components before serving. +Main-product K tiles are combined with explicit FP32 rounding for accuracy. +Split-K configurations use a separate reduction kernel. +""" + +import functools + +import cuda.tile as ct +import torch + +from sglang.srt.utils.custom_op import register_custom_op + + +@ct.kernel(num_ctas=1) +def _router_gemm_bf16x3_kernel( + X, + W, + WR, + Y, + BM: ct.Constant[int], + BN: ct.Constant[int], + BK: ct.Constant[int], + SPLIT_K: ct.Constant[int], + EXPERT_MAJOR: ct.Constant[bool], + PDL: ct.Constant[bool], +): + col, row, part = ct.bid(0), ct.bid(1), ct.bid(2) + if PDL: + ct.grid_dependency_control_launch_dependents() + ct.grid_dependency_control_wait() + steps = ct.cdiv(ct.cdiv(X.shape[1], BK), SPLIT_K) + if EXPERT_MAJOR: + total = ct.full((4 * BN, BM), 0, ct.float32) + else: + total = ct.full((BM, BN), 0, ct.float32) + residual = ct.full((BM, 2 * BN), 0, ct.float32) + for step in range(steps): + k = part * steps + step + x = ct.load( + X, + (row, k), + (BM, BK), + padding_mode=ct.PaddingMode.ZERO, + latency=4 if EXPERT_MAJOR or BK == 128 else 8, + ) + if EXPERT_MAJOR: + w = ct.load( + W, (0, col, k), (4, BN, BK), padding_mode=ct.PaddingMode.ZERO, latency=4 + ) + local = ct.mma( + ct.reshape(w, (4 * BN, BK)), + ct.transpose(x), + ct.full((4 * BN, BM), 0, ct.float32), + ) + else: + w0 = ct.load( + W, + (0, col, k), + (1, BN, BK), + padding_mode=ct.PaddingMode.ZERO, + latency=4 if BK == 128 else 8, + ) + wr = ct.load( + WR, + (0, col, k), + (2, BN, BK), + padding_mode=ct.PaddingMode.ZERO, + latency=4 if BK == 128 else 8, + ) + local = ct.mma( + x, + ct.transpose(ct.reshape(w0, (BN, BK))), + ct.full((BM, BN), 0, ct.float32), + ) + residual = ct.mma(x, ct.transpose(ct.reshape(wr, (2 * BN, BK))), residual) + # Round each main-product tile before adding the next tile. + total = ct.add(total, local, rounding_mode=ct.RoundingMode.RN) + if EXPERT_MAJOR: + main = ct.extract(total, (0, 0), (BN, BM)) + mid = ct.extract(total, (1, 0), (BN, BM)) + low = ct.extract(total, (2, 0), (BN, BM)) + value = ct.transpose( + ct.add( + main, + ct.add(mid, low, rounding_mode=ct.RoundingMode.RN), + rounding_mode=ct.RoundingMode.RN, + ) + ) + else: + correction = ct.sum( + ct.reshape(residual, (BM, 2, BN)), axis=1, rounding_mode=ct.RoundingMode.RN + ) + value = ct.add(total, correction, rounding_mode=ct.RoundingMode.RN) + ct.store( + Y, + (part, row, col), + ct.reshape(value, (1, BM, BN)), + allow_tma=not (BM == 256 and SPLIT_K > 1), + ) + + +@ct.kernel +def _split_k_reduce_kernel( + P, Y, SPLIT_K: ct.Constant[int], PDL: ct.Constant[bool] = False +): + if PDL: + ct.grid_dependency_control_wait() + if SPLIT_K == 2: + values = ct.load( + P, + (0, ct.bid(0)), + (SPLIT_K, 1024), + padding_mode=ct.PaddingMode.ZERO, + allow_tma=False, + ) + total = ct.sum(values, axis=0, rounding_mode=ct.RoundingMode.RN) + else: + total = ct.full((1024,), 0, ct.float32) + for part in range(SPLIT_K): + value = ct.load( + P, + (part, ct.bid(0)), + (1, 1024), + padding_mode=ct.PaddingMode.ZERO, + allow_tma=False, + ) + total = ct.add( + total, ct.reshape(value, (1024,)), rounding_mode=ct.RoundingMode.RN + ) + ct.store(Y, (ct.bid(0),), total, allow_tma=False) + + +_GEMM_KERNELS = { + 1: _router_gemm_bf16x3_kernel, + 2: _router_gemm_bf16x3_kernel.replace_hints(num_ctas=2), +} +# (max M, (BM, BN, BK, SPLIT_K, CTAS)); larger M uses the fallback below. +_BLACKWELL_CONFIGS = ( + (64, (32, 16, 256, 4, 1)), + (128, (64, 32, 256, 4, 1)), + (512, (128, 64, 256, 8, 1)), + (1024, (128, 64, 128, 4, 1)), + (2048, (128, 64, 128, 2, 1)), + (4096, (256, 128, 256, 2, 2)), +) +# Use smaller tiles on Hopper to limit register pressure. +_HOPPER_CONFIGS = ( + (256, (64, 64, 128, 8, 1)), + (512, (64, 64, 128, 4, 1)), + (1024, (64, 64, 128, 2, 1)), +) + + +@functools.cache +def _is_hopper(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index)[0] == 9 + + +def _select_launch_config(m: int, device: torch.device) -> tuple: + if m > 128 and _is_hopper(device.index): + configs, fallback = _HOPPER_CONFIGS, (64, 64, 128, 1, 1) + else: + configs, fallback = _BLACKWELL_CONFIGS, (256, 128, 256, 1, 2) + return next((config for limit, config in configs if m <= limit), fallback) + + +def _router_gemm_bf16x3_fake(x: torch.Tensor, parts: torch.Tensor) -> torch.Tensor: + return x.new_empty((x.shape[0], parts.shape[0] // 3), dtype=torch.float32) + + +@register_custom_op(fake_impl=_router_gemm_bf16x3_fake) +def router_gemm_bf16x3_cutile(x: torch.Tensor, parts: torch.Tensor) -> torch.Tensor: + """Multiply BF16 [M, K] by prepared BF16 [3N, K], producing FP32 [M, N].""" + m, k = x.shape + n = parts.shape[0] // 3 + bm, bn, bk, splits, ctas = _select_launch_config(m, x.device) + pdl = bk == 128 and splits == 4 + out = torch.empty((m, n), device=x.device, dtype=torch.float32) + if m == 0 or n == 0: + return out + partial = ( + torch.empty((splits, m, n), device=x.device, dtype=torch.float32) + if splits > 1 + else out.unsqueeze(0) + ) + with torch.cuda.device(x.device): + stream = torch.cuda.current_stream() + ct.launch( + stream, + (ct.cdiv(n, bn), ct.cdiv(m, bm), splits), + _GEMM_KERNELS[ctas], + ( + x, + parts.view(3, n, k), + parts[n:].view(2, n, k), + partial, + bm, + bn, + bk, + splits, + m <= 64, + pdl, + ), + programmatic_dependent_launch=pdl, + ) + if splits > 1: + ct.launch( + stream, + (ct.cdiv(out.numel(), 1024),), + _split_k_reduce_kernel, + (partial.view(splits, -1), out.view(-1), splits, pdl), + programmatic_dependent_launch=pdl, + ) + return out diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index ed4d76217186..7ec5f9e16cd5 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 @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import Dict, Tuple +from typing import Dict, Optional, Tuple import torch import triton @@ -31,13 +31,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; above +# 8 rows router_gemm_bf16x3_cublas is faster. Swept for N=256, K=5120 on H200 and B300. +_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 = torch.bfloat16): + 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 +128,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 +139,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 +161,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 +184,63 @@ def router_gemv(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: num_warps=num_warps, ) return out + + +@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 as (hi, mid, lo) blocks with w == hi + mid + lo. + + Exact for |w| >= 2**-110: each part takes the next 8 significant bits, and the + last residual then stays a normal number. + """ + if out is None: + out = torch.empty( + 3 * w.shape[0], w.shape[1], dtype=torch.bfloat16, device=w.device + ) + if w.is_cuda and w.is_contiguous(): + block = 1024 + _split_fp32_to_bf16x3_kernel[(triton.cdiv(w.numel(), block),)]( + w, out, w.numel(), BLOCK=block + ) + return out + hi = w.to(torch.bfloat16) + residual = w - hi.float() + mid = residual.to(torch.bfloat16) + lo = (residual - mid.float()).to(torch.bfloat16) + return torch.cat([hi, mid, lo], out=out) + + +_split_scratch: Dict[Tuple[torch.device, int, int], torch.Tensor] = {} + + +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. Every bf16 x bf16 product is exact in fp32, but + tensor-core accumulation is not IEEE fp32: the error is a few times that of + an fp32 SGEMM. + """ + n, k = w.shape + key = (w.device, n, k) + if key not in _split_scratch: + _split_scratch[key] = torch.empty( + 3 * n, k, dtype=torch.bfloat16, device=w.device + ) + w3 = split_fp32_to_bf16x3(w, out=_split_scratch[key]) + 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..acb15d943e2d 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 @@ -14,6 +14,7 @@ from sglang.kernels.jit.utils import ( cache_once, + get_jit_cuda_arch, is_arch_support_pdl, load_jit, make_cpp_args, @@ -42,10 +43,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 = torch.bfloat16, ) -> 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( @@ -92,7 +107,13 @@ def _supports_k_variant(n: int, k: int) -> bool: @cache_once -def _default_n_gemm_n_split(n: int, k: int, max_m: int) -> int: +def _default_n_gemm_n_split( + n: int, k: int, max_m: int, weight_dtype: torch.dtype = torch.bfloat16, m: int = 16 +) -> int: + if weight_dtype == torch.float32 and (n, k) == (256, 5120) and max_m <= 16: + # Rotating 49 FP32 router weights under CUDA graph replay. + small_m = 8 if get_jit_cuda_arch().major == 10 else 4 + return 1 if m <= small_m else 2 split_cap = _n_gemm_block_threads(k) // max_m candidates = [d for d in range(1, min(n, split_cap) + 1) if n % d == 0] if not candidates: @@ -137,7 +158,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 +188,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: @@ -179,7 +200,7 @@ def tiny_gemm_bf16( n_split = ( _default_k_gemm_n_split(n, k) if _prefer_k_variant(n, k) - else _default_n_gemm_n_split(n, k, max_m) + else _default_n_gemm_n_split(n, k, max_m, w.dtype, x.shape[0]) ) _tiny_gemm_custom_op(x, w, out, max_m, n_split) return out diff --git a/python/sglang/kernels/ops/kvcache/zero_pages.py b/python/sglang/kernels/ops/kvcache/zero_pages.py index afa3c0071085..2c456b23022e 100644 --- a/python/sglang/kernels/ops/kvcache/zero_pages.py +++ b/python/sglang/kernels/ops/kvcache/zero_pages.py @@ -1,8 +1,8 @@ """Zero whole page envelopes of the unified pool by physical page id. -The pool is viewed as int64 words (the MLA page envelope is always -8-byte-aligned: entry bytes per layer = kv_cache_dim * itemsize, a multiple -of 8), one wide element per lane; grid = (num_pages, page word blocks). +The pool is viewed as int64 words (MHA and MLA page envelopes are +8-byte-aligned), one wide element per lane; grid = (num_pages, page word +blocks). """ from __future__ import annotations @@ -41,6 +41,9 @@ def zero_pages( if m == 0: return assert raw.dtype == torch.uint8, f"expected uint8 pool, got {raw.dtype}" + if not raw.is_cuda: + raw[: num_pages * page_bytes].view(num_pages, page_bytes)[pages] = 0 + return assert page_bytes % 8 == 0, f"page_bytes {page_bytes} not int64-aligned" page_words = page_bytes // 8 words = raw[: num_pages * page_bytes].view(torch.int64) diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index 1e71b01c92ef..e4b8220d2f54 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -439,6 +439,83 @@ def forward_jit( residual.copy_(residual_out) +class BeamSandwichResidualRMSNormOp(BaseFusedOp): + """Post-normalize a branch, combine residuals, then normalize the next input. + + Returns ``(normed, combined)``. Either norm weight may be None; without + ``out_weight``, ``normed`` is None (a non-final pipeline stage). + """ + + op = "layernorm.beam_sandwich_residual_rmsnorm" + capabilities = {KernelBackend.TRITON: _CUDA} + format_signature = FormatSignature(supported_dtypes=_NORM_DTYPES) + + 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): + import torch + + if self._compiled_native is None: + # Register and warm the opaque op before tracing its dispatch wrapper. + self.forward_triton(*args, **kwargs) + self._compiled_native = torch.compile( + self.forward_triton, options={"emulate_precision_casts": True} + ) + return self._compiled_native(*args, **kwargs) + + def forward_native( + self, + branch: torch.Tensor, + residual: torch.Tensor, + sandwich_weight: Optional[torch.Tensor], + out_weight: Optional[torch.Tensor], + eps: float, + scales: tuple[float, float], + ) -> tuple[Optional[torch.Tensor], torch.Tensor]: + if sandwich_weight is not None: + branch = _RMSNORM.forward_native(branch, sandwich_weight, eps) + residual = residual.to(branch.dtype) + combined = branch * scales[0] + residual * scales[1] + out = ( + _RMSNORM.forward_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: Optional[torch.Tensor], + out_weight: Optional[torch.Tensor], + eps: float, + scales: tuple[float, float], + ) -> tuple[Optional[torch.Tensor], torch.Tensor]: + import torch + + if not branch.is_cuda or branch.dtype not in (torch.float16, torch.bfloat16): + return self.forward_native( + branch, residual, sandwich_weight, out_weight, eps, scales + ) + from sglang.kernels.jit.utils import is_arch_support_pdl + from sglang.kernels.ops.layernorm.beam_fused import ( + fused_sandwich_residual_rmsnorm, + ) + + return fused_sandwich_residual_rmsnorm( + branch, + residual, + sandwich_weight, + out_weight, + eps, + *scales, + enable_pdl=is_arch_support_pdl(), + ) + + _RMSNORM = register_fused_op(RMSNormOp(), __name__, "_RMSNORM") _FUSED_ADD_RMSNORM = register_fused_op( FusedAddRMSNormOp(), __name__, "_FUSED_ADD_RMSNORM" @@ -448,6 +525,10 @@ def forward_jit( GemmaFusedAddRMSNormOp(), __name__, "_GEMMA_FUSED_ADD_RMSNORM" ) +_BEAM_SANDWICH_RESIDUAL_RMSNORM = register_fused_op( + BeamSandwichResidualRMSNormOp(), __name__, "_BEAM_SANDWICH_RESIDUAL_RMSNORM" +) + def rmsnorm( input: torch.Tensor, @@ -494,6 +575,7 @@ def gemma_fused_add_rmsnorm( __all__ = [ + "BeamSandwichResidualRMSNormOp", "RMSNormOp", "FusedAddRMSNormOp", "GemmaRMSNormOp", 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..5158657b6144 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -0,0 +1,177 @@ +# Copyright 2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""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.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_SANDWICH: 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. + if HAS_SANDWICH: + 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 + ) + if HAS_SANDWICH: + 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: Optional[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 (None skips it); `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") + 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: Optional[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 + dummy = sandwich_weight if sandwich_weight is not None else out_weight + if dummy is None: + dummy = branch # never dereferenced: both norm flags are off + _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 if sandwich_weight is not None else dummy, + out_weight if out_weight is not None else dummy, + eps=eps, + x_scale=x_scale, + residual_scale=residual_scale, + hidden_dim=hidden_dim, + BLOCK_SIZE=triton.next_power_of_2(hidden_dim), + HAS_SANDWICH=sandwich_weight is not None, + 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 {}), + ) diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index 1f752167492f..22eaaa7967fa 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -29,6 +29,22 @@ ) ) +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))} + ), + format_signature=FormatSignature( + supported_dtypes=("bfloat16",), + description="single-token interleaved Beam MoE with FP32 routing weights", + ), + description="SM103 Beam TP4/EP1 weighted expert outputs; caller owns reduction.", + ) +) + register_kernel( KernelSpec( op="moe.moe_align_block_size", 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..a62a062639b8 --- /dev/null +++ b/python/sglang/kernels/ops/moe/beam_decode_gemv.py @@ -0,0 +1,150 @@ +# Copyright 2026 SGLang Team. Licensed under the Apache License, Version 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 592dcf31978f..f89308d21877 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, @@ -629,36 +630,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/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 15b0bbadaedd..3a802c592ddb 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 @@ -214,6 +214,18 @@ def force_thresholds(heuristic): self._graph_mode_allowed = 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 _init_workspace(self) -> None: """Slice one symmetric-memory allocation into every shared buffer. diff --git a/test/manual/kernels/bench_fp32_router.py b/test/manual/kernels/bench_fp32_router.py new file mode 100644 index 000000000000..a8af4d3665f5 --- /dev/null +++ b/test/manual/kernels/bench_fp32_router.py @@ -0,0 +1,67 @@ +"""Compare cuBLAS and mixed-weight tiny_gemm. + +Run with the candidate Python package on CUDA SM90+: + python test/manual/kernels/bench_fp32_router.py + +Times are microseconds per layer under CUDA graph replay. The baseline includes +BF16-to-FP32 activation conversion; weights stay FP32 for all implementations. +The 245 MiB weight working set exceeds H200/B300 L2. Compilation is untimed. +""" + +import json + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 + + +def bench(fn, weights): + def run_layers(): + for weight in weights: + fn(weight) + + result = marker.do_bench( + run_layers, + warmup_iters=2, + metrics=(0.5,), + memory_output=None, + ) + return result.times[0] * 1e6 / len(weights) + + +def main(): + torch.manual_seed(42) + torch.backends.cuda.matmul.allow_tf32 = False + print( + json.dumps({"gpu": torch.cuda.get_device_name(), "torch": torch.__version__}), + flush=True, + ) + weights = [ + torch.randn(256, 5120, device="cuda", dtype=torch.float32) / 5120**0.5 + for _ in range(49) + ] + for m in range(1, 17): + x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) + base = bench(lambda w: torch.nn.functional.linear(x.float(), w), weights) + ref = x.double() @ weights[0].double().T + out = tiny_gemm_bf16(x, weights[0], out_dtype=torch.float32) + err = (out.double() - ref).abs().max().item() + us = bench(lambda w: tiny_gemm_bf16(x, w, out_dtype=torch.float32), weights) + print( + json.dumps( + dict( + m=m, + variant="tiny", + base_us=base, + us=us, + speedup=base / us, + max_error=err, + ) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py b/test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py new file mode 100644 index 000000000000..8f9a2308d37a --- /dev/null +++ b/test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py @@ -0,0 +1,142 @@ +"""Compare split and unsplit mixed-extend attention on SM100-family GPUs. + +Times forward_extend with CUDA graphs and cold L2, excluding cache flushes, +projection, KV writes and metadata construction. Run from the repository root +with PYTHONPATH=python. +""" + +import argparse +import gc +import json +import unittest + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.utils.common import is_sm100_supported +from sglang.test.kits.attention_unittest.attention_methods.dense_attention import ( + DenseAttentionCase, + build_dense_attention_fixture, +) + + +@torch.no_grad() +def benchmark(args, num_decode): + case = DenseAttentionCase( + name="mixed_extend", + backend="trtllm_mha", + forward_mode=ForwardMode.MIXED, + num_heads=args.num_heads, + num_kv_heads=args.num_kv_heads, + page_size=64, + prefix_lens=(0,) + (args.kv_len - 1,) * num_decode, + extend_lens=(args.prefill_tokens,) + (1,) * num_decode, + ) + fixture = build_dense_attention_fixture( + unittest.TestCase(), + case, + head_dim=args.head_dim, + hidden_size=128, + max_context_len=((max(args.kv_len, args.prefill_tokens) + 63) // 64) * 64, + dtype=torch.bfloat16, + ) + runner = fixture.runner + layer = fixture.actual_module.attn + layer.sliding_window_size = args.window_size + if args.kv_cache_dtype == "fp8": + runner.kv_cache_dtype = torch.float8_e4m3fn + runner.token_to_kv_pool = MHATokenToKVPool( + size=runner.token_to_kv_pool.size, + page_size=64, + dtype=runner.kv_cache_dtype, + head_num=args.num_kv_heads, + head_dim=args.head_dim, + layer_num=1, + device="cuda", + enable_memory_saver=False, + enable_alt_stream=False, + ) + runner.init_kv_index_translator() + fixture.backend = TRTLLMHAAttnBackend(runner) + for req, prefix in enumerate(fixture.prefix_hidden): + if prefix.shape[0]: + _, k, v = fixture.actual_module.project_qkv(prefix) + runner.token_to_kv_pool.set_kv_buffer( + layer, + runner.req_to_token_pool.req_to_token[req, : len(prefix)].long(), + k.view(-1, args.num_kv_heads, args.head_dim), + v.view(-1, args.num_kv_heads, args.head_dim), + ) + backend = fixture.backend + if not is_sm100_supported(): + raise RuntimeError("The decode-tail split requires an SM100-family GPU") + q, k, v = fixture.actual_module.project_qkv(fixture.input_hidden) + batch = fixture.forward_batch + backend.init_forward_metadata(batch) + backend.forward_extend(q, k, v, layer, batch) + + def run(): + return backend.forward_extend(q, None, None, layer, batch, save_kv_cache=False) + + outputs, timings = {}, {} + for split in (False, True): + backend.init_forward_metadata(batch) + if not split: + backend.forward_metadata.mixed_prefill_reqs = -1 + outputs[split] = run().clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + # Flush before each replay, not just before a graph of repeated calls. + timings[split] = ( + marker.do_bench( + graph.replay, + use_cuda_graph=False, + estimated_time_ms=200, + disable_log_bandwidth=True, + ).times[0] + * 1e6 + ) + torch.testing.assert_close(outputs[True], outputs[False], atol=1e-2, rtol=1e-2) + delta = outputs[True].float() - outputs[False].float() + return { + "decode_requests": num_decode, + "prefill_tokens": args.prefill_tokens, + "kv_len": args.kv_len, + "kv_cache_dtype": args.kv_cache_dtype, + "window_size": args.window_size, + "num_heads": args.num_heads, + "num_kv_heads": args.num_kv_heads, + "head_dim": args.head_dim, + "unsplit_us": timings[False], + "split_us": timings[True], + "speedup": timings[False] / timings[True], + "max_abs_diff": delta.abs().max().item(), + "relative_rms_diff": (delta.norm() / outputs[False].float().norm()).item(), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--decode-requests", type=int, nargs="+", default=[0, 1, 32, 128, 256, 511] + ) + parser.add_argument("--prefill-tokens", type=int, default=8192) + parser.add_argument("--kv-len", type=int, default=8192) + parser.add_argument("--kv-cache-dtype", choices=["bf16", "fp8"], default="fp8") + parser.add_argument("--window-size", type=int, default=-1) + parser.add_argument("--num-heads", type=int, default=12) + parser.add_argument("--num-kv-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + args = parser.parse_args() + for num_decode in args.decode_requests: + print(json.dumps(benchmark(args, num_decode)), flush=True) + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py b/test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py new file mode 100644 index 000000000000..8f2da343704e --- /dev/null +++ b/test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py @@ -0,0 +1,40 @@ +"""Beam attention-gate latency at decode and prefill sizes.""" + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE +from sglang.kernels.spec import KernelBackend +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=7, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + + +@marker.parametrize("rows", [1, 16, 64, 128, 1024, 8192], [1, 128]) +@marker.parametrize("packed_gate", [False, True]) +@marker.benchmark("impl", ["triton", "torch"]) +def benchmark(rows: int, packed_gate: bool, impl: str): + x = torch.randn(rows, 6 * 128, dtype=torch.bfloat16, device="cuda") + if packed_gate: + # TP8 QKV+gate widths: Q=768, K=128, V=128, gate=768. + gate_source = torch.randn(rows, 1792, dtype=x.dtype, device=x.device) + else: + gate_source = torch.randn_like(x) + gate_start = gate_source.shape[1] - x.shape[1] + bias = torch.randn(6, dtype=x.dtype, device="cuda") + # The marker rotates clones, and a clone of a strided slice is contiguous; + # rotate the backing tensor and slice the gate inside the timed call. + return marker.do_bench( + lambda x, gate_source, bias: _BEAM_ATTENTION_GATE( + x, gate_source[:, gate_start:], bias, 128, backend=KernelBackend(impl) + ), + input_args=(x, gate_source, bias), + memory_args=(x, gate_source[:, gate_start:], bias), + memory_output="out", + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py b/test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py new file mode 100644 index 000000000000..c651c900907f --- /dev/null +++ b/test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py @@ -0,0 +1,111 @@ +"""Compare router GEMM latency and mean absolute error against FP64. + +Run this file on Hopper or Blackwell. Weight preparation is outside the timed region; +split-K reduction is included. Latencies use rotating inputs and CUDA graphs. +All providers use the same random inputs (seed 0, X std 1, W std 0.02). +""" + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random, get_benchmark_range +from sglang.kernels.ops.gemm.router_gemm_cutile import router_gemm_bf16x3_cutile +from sglang.kernels.ops.gemm.router_gemv import split_fp32_to_bf16x3 +from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=30, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +PROVIDERS = ["SGLang tiny FP32", "Cast + TF32", "cuBLAS BF16x3", "cuTile BF16x3"] + + +def _cast_tf32(x, w): + return x.float() @ w.T + + +def _cublas(x, parts): + n = parts.shape[0] // 3 + return torch.mm(x, parts.T, out_dtype=torch.float32).view(-1, 3, n).sum(1) + + +def _print_table(title, rows, number_format): + print(f"\n### {title}\n") + print("| Tokens | " + " | ".join(PROVIDERS) + " |") + print("| ---: " * (len(PROVIDERS) + 1) + "|") + for m, values in rows: + cells = ["—" if v is None else format(v, number_format) for v in values] + print(f"| {m} | " + " | ".join(cells) + " |") + + +def benchmark(batch_sizes, n, k): + torch.manual_seed(0) + w = create_random(n, k, dtype=torch.float32) * 0.02 + parts = split_fp32_to_bf16x3(w) + latencies, errors = [], [] + old_tf32 = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = True + for m in batch_sizes: + x = create_random(m, k, dtype=torch.bfloat16) + ref = x.double() @ w.double().T + cases = [ + ( + tiny_gemm_bf16, + w, + {"out_dtype": torch.float32, "max_m": 16 if m <= 16 else 32}, + ), + (_cast_tf32, w, {}), + (_cublas, parts, {}), + (router_gemm_bf16x3_cutile, parts, {}), + ] + times, maes = [], [] + for fn, weight, kwargs in cases: + if fn is tiny_gemm_bf16 and m > 32: + times.append(None) + maes.append(None) + continue + result = marker.do_bench( + fn, input_args=(x, weight), input_kwargs=kwargs + ) + out = fn(x, weight, **kwargs) + times.append(result.times[0] * 1e6) + maes.append((out.double() - ref).abs().mean().item()) + latencies.append((m, times)) + errors.append((m, maes)) + finally: + torch.backends.cuda.matmul.allow_tf32 = old_tf32 + print(f"\n## K{k}/N{n}") + _print_table("Latency (µs)", latencies, ".2f") + _print_table("Mean absolute error vs FP64", errors, ".2e") + + +if __name__ == "__main__": + major = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0 + if major not in (9, 10): + print("Hopper or Blackwell required.") + else: + batch_sizes = get_benchmark_range( + [ + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + 16896, + ], + [1, 32, 1024, 16384], + ) + for n, k in get_benchmark_range([(256, 5120), (128, 6144)], [(256, 5120)]): + benchmark(batch_sizes, n, k) diff --git a/test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py b/test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py new file mode 100644 index 000000000000..be744a7ce01f --- /dev/null +++ b/test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py @@ -0,0 +1,32 @@ +"""Beam norm/residual latency at decode and prefill sizes.""" + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.layernorm import _BEAM_SANDWICH_RESIDUAL_RMSNORM +from sglang.kernels.spec import KernelBackend +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + + +@marker.parametrize("rows", [1, 16, 64, 128, 1024, 8192], [1, 128]) +@marker.benchmark("impl", ["triton", "torch"]) +def benchmark(rows: int, impl: str): + x = torch.randn(rows, 5120, dtype=torch.bfloat16, device="cuda") + residual = torch.randn_like(x) + weight = torch.ones(5120, dtype=x.dtype, device="cuda") + return marker.do_bench( + lambda *args: _BEAM_SANDWICH_RESIDUAL_RMSNORM( + *args, backend=KernelBackend(impl) + ), + input_args=(x, residual, weight, weight, 1e-6, (0.25, 0.968)), + memory_args=(x, residual, weight), + memory_output="out", + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py b/test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py new file mode 100644 index 000000000000..70c953e04337 --- /dev/null +++ b/test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py @@ -0,0 +1,147 @@ +"""Every-backend-vs-native parity for BaseFusedOp ops on real GPU (RFC #29630). + +For each reworked fused op, run every backend eligible on this platform and +assert it matches the pure-torch ``forward_native`` reference within dtype +tolerance. New backends are picked up automatically. +""" + +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=13, stage="extra-a", runner_config="1-gpu-small") + +_SKIP = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE} + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _eligible(op): + return [ + b for b in op.available_backends() if b not in _SKIP and op.backend_eligible(b) + ] + + +@pytest.mark.parametrize("rows", [0, 1, 64, 1024]) +def test_beam_attention_gate(rows): + from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE as op + + generator = torch.Generator().manual_seed(5) + bias = torch.randn(12, generator=generator).to(device="cuda", dtype=torch.bfloat16) + attention = torch.randn(rows, 12 * 128, generator=generator).to( + device="cuda", dtype=torch.bfloat16 + ) + gate = ( + torch.randn(rows, 12 * 128, generator=generator).to( + device="cuda", dtype=torch.bfloat16 + ) + * 2 + ) + expected = op.forward_native(attention, gate, bias, 128) + for backend in _eligible(op): + torch.testing.assert_close( + op(attention, gate, bias, 128, backend=backend), expected, atol=0, rtol=0 + ) + + +@pytest.mark.parametrize("rows", [1, 16, 64]) +@pytest.mark.parametrize("tp_size", [4, 8]) +@pytest.mark.parametrize("column_step", [1, 2]) +def test_beam_packed_gate_without_copy(rows, tp_size, column_step): + """A QKV+gate slice must preserve rounding without materializing the gate.""" + from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE as op + from sglang.kernels.ops.elementwise.beam_fused import fused_attention_gate + + q_size, kv_size = 48 * 128 // tp_size, 8 * 128 // tp_size + packed = torch.randn( + rows, + (2 * q_size + 2 * kv_size) * column_step, + device="cuda", + dtype=torch.bfloat16, + ) + gate = packed[:, (q_size + 2 * kv_size) * column_step :: column_step] + attention = torch.randn(rows, q_size, device="cuda", dtype=gate.dtype) + bias = torch.randn(q_size // 128, device="cuda", dtype=gate.dtype) + expected = op.forward_native(attention, gate, bias, 128) + fused_attention_gate(attention, gate, bias, 128) # Compile before profiling. + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU] + ) as prof: + actual = fused_attention_gate(attention, gate, bias, 128) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert not {"aten::clone", "aten::copy_"} & { + event.key for event in prof.key_averages() + } + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + replayed = fused_attention_gate(attention, gate, bias, 128) + packed.normal_() + graph.replay() + torch.testing.assert_close( + replayed, op.forward_native(attention, gate, bias, 128), atol=0, rtol=0 + ) + + +def test_beam_fp16_gate_rounding_boundaries(): + from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE as op + + gate = torch.tensor( + [[0.0068359375, -0.00634765625]], device="cuda", dtype=torch.float16 + ) + attention = torch.ones_like(gate) + bias = torch.zeros(1, device="cuda", dtype=torch.float16) + expected = op.forward_native(attention, gate, bias, 2) + for backend in _eligible(op): + torch.testing.assert_close( + op(attention, gate, bias, 2, backend=backend), expected, atol=0, rtol=0 + ) + + +def test_beam_custom_op_contracts(): + from sglang.kernels.ops.elementwise.beam_fused import fused_attention_gate + + # Sliced rows cover real/fake output strides. + x = torch.randn(8, 5120, device="cuda", dtype=torch.bfloat16)[::2] + bias = torch.zeros(40, device="cuda", dtype=x.dtype) + torch.library.opcheck(fused_attention_gate, (x, x.clone(), bias, 128)) + + +@pytest.mark.parametrize("forced_backend", [False, True]) +def test_beam_compiled_rounding(forced_backend): + """Outer compilation must preserve sigmoid rounding.""" + from torch._inductor import config as inductor_config + + from sglang.kernels.ops.elementwise import BeamAttentionGateOp + + op = BeamAttentionGateOp() + gate = torch.tensor( + [[0.0068359375, -0.00634765625]], device="cuda", dtype=torch.float16 + ) + args = ( + torch.ones_like(gate), + gate, + torch.zeros(1, device="cuda", dtype=gate.dtype), + 2, + ) + expected = op.forward_native(*args) + op(*args) # Warm backend imports before full-graph capture. + op.enter_torch_compile(1) + try: + with inductor_config.patch(compile_threads=1): + actual = ( + op(*args, backend=KernelBackend.TORCH_COMPILE) + if forced_backend + else torch.compile(op, fullgraph=True)(*args) + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + finally: + op.leave_torch_compile() + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__])) diff --git a/test/registered/kernels/ops/gemm/test_router_gemm_cutile.py b/test/registered/kernels/ops/gemm/test_router_gemm_cutile.py new file mode 100644 index 000000000000..2f68ea25b35b --- /dev/null +++ b/test/registered/kernels/ops/gemm/test_router_gemm_cutile.py @@ -0,0 +1,237 @@ +"""Accuracy and CUDA graph tests for the prepared BF16x3 router GEMM.""" + +import sys + +import pytest +import torch + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytest.importorskip("cuda.tile") + +import cuda.tile as ct # noqa: E402 + +from sglang.kernels.ops.gemm.router_gemm_cutile import ( # noqa: E402 + router_gemm_bf16x3_cutile, +) +from sglang.kernels.ops.gemm.router_gemv import split_fp32_to_bf16x3 # noqa: E402 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] not in (9, 10), + reason="Hopper or Blackwell required", +) + + +def _check(out, x, w): + ref = x.double() @ w.double().T + error = out.double() - ref + torch.testing.assert_close(out.double(), ref, atol=8e-6, rtol=3e-6) + assert (error.square().mean() / ref.square().mean()).sqrt().item() < 1e-6 + + +@pytest.mark.parametrize( + "m", + [ + 1, + 2, + 8, + 16, + 32, + 64, + 65, + 128, + 129, + 512, + 513, + 1024, + 1025, + 2048, + 2049, + 4096, + 4097, + 8192, + 16384, + ], +) +def test_router_shape(m): + torch.manual_seed(1) + x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 5120, device="cuda") * 0.02 + out = router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)) + assert out.dtype == torch.float32 + _check(out, x, w) + + +@pytest.mark.parametrize( + "m,k,n", + [ + (2, 137, 160), + (17, 2816, 256), + (129, 5120, 160), + (257, 5128, 160), + (257, 6144, 128), + (513, 5128, 160), + (1025, 137, 160), + (8193, 5120, 256), + ], +) +def test_tail_tiles(m, k, n): + torch.manual_seed(1) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + w = torch.randn(n, k, device="cuda") / k**0.5 + _check(router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)), x, w) + + +def test_strided_input(): + torch.manual_seed(1) + x = torch.randn(129, 5144, device="cuda", dtype=torch.bfloat16)[:, :5128] + w = torch.randn(160, 5128, device="cuda") * 0.02 + _check(router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)), x, w) + + +@pytest.mark.parametrize("m,n", [(0, 256), (32, 0)]) +def test_empty_output(m, n): + x = torch.empty(m, 5120, device="cuda", dtype=torch.bfloat16) + parts = torch.empty(3 * n, 5120, device="cuda", dtype=torch.bfloat16) + out = router_gemm_bf16x3_cutile(x, parts) + assert out.shape == (m, n) + assert out.dtype == torch.float32 + assert out.device == x.device + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two GPUs required") +def test_non_current_device(): + current = torch.cuda.current_device() + with torch.cuda.device((current + 1) % torch.cuda.device_count()): + x = torch.randn(32, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 5120, device="cuda") * 0.02 + parts = split_fp32_to_bf16x3(w) + assert torch.cuda.current_device() == current + _check(router_gemm_bf16x3_cutile(x, parts), x, w) + + +@pytest.mark.parametrize("m", [1, 2, 8, 16, 32, 128, 512, 1024, 2048, 16384]) +@pytest.mark.parametrize("k,distance", [(512, 256), (5120, 320), (5120, 2560)]) +def test_cancellation(m, k, distance, record_property): + torch.manual_seed(1) + blocks = k // (2 * distance) + half = torch.randn(m, blocks, distance, device="cuda", dtype=torch.bfloat16) + x = torch.cat((half, half), dim=-1).reshape(m, k) + main = torch.randn( + 256, blocks, distance, device="cuda", dtype=torch.bfloat16 + ).float() + w = torch.cat((main, -main + torch.randn_like(main) * 1e-4), dim=-1).reshape(256, k) + if k != 512: + w = w / k**0.5 + ref = x.double() @ w.double().T + old_tf32 = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = False + baseline_error = (x.float() @ w.T - ref).abs().max() + finally: + torch.backends.cuda.matmul.allow_tf32 = old_tf32 + error = (router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)) - ref).abs() + record_property("fp32_max_error", baseline_error.item()) + record_property("candidate_max_error", error.max().item()) + record_property("candidate_mae", error.mean().item()) + assert error.mean().item() < 1e-6 + + +@pytest.mark.parametrize("m", [64, 512, 1024, 2048, 4096, 8192]) +def test_cuda_graph_replay(m): + torch.manual_seed(1) + x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 5120, device="cuda") * 0.02 + parts = split_fp32_to_bf16x3(w) + router_gemm_bf16x3_cutile(x, parts) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = router_gemm_bf16x3_cutile(x, parts) + for _ in range(2): + x.normal_() + w.normal_(std=0.02) + split_fp32_to_bf16x3(w, out=parts) + graph.replay() + _check(out, x, w) + + +@ct.kernel +def _pdl_producer(source, output): + ct.grid_dependency_control_launch_dependents() + # Keep input production pending after releasing the dependent GEMM. + delay = ct.astype(ct.load(source, (0,), (32,)), ct.float32) + for _ in range(1024): + delay = ct.sin(delay) + 0.001 + offset = ct.sum(delay, axis=0) / 32 + for block in range(ct.bid(0), ct.cdiv(source.shape[0], 1024), 32): + value = ct.load(source, (block,), (1024,), padding_mode=ct.PaddingMode.ZERO) + ct.store(output, (block,), ct.astype(value + offset, ct.bfloat16)) + + +@pytest.mark.parametrize("m", [512, 1024]) +def test_pdl_producer_graph_replay(m): + torch.manual_seed(1) + source = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) + x = torch.empty_like(source) + w = torch.randn(256, 5120, device="cuda") * 0.02 + parts = split_fp32_to_bf16x3(w) + ct.launch( + torch.cuda.current_stream(), (32,), _pdl_producer, (source.view(-1), x.view(-1)) + ) + torch.cuda.synchronize() + router_gemm_bf16x3_cutile(x, parts) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + ct.launch( + torch.cuda.current_stream(), + (32,), + _pdl_producer, + (source.view(-1), x.view(-1)), + ) + out = router_gemm_bf16x3_cutile(x, parts) + for _ in range(3): + source.normal_() + x.fill_(float("nan")) + graph.replay() + _check(out, x, w) + + +def test_concurrent_graphs(): + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + cases = [] + for stream in streams: + with torch.cuda.stream(stream): + x = torch.randn(1024, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 5120, device="cuda") * 0.02 + parts = split_fp32_to_bf16x3(w) + router_gemm_bf16x3_cutile(x, parts) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = router_gemm_bf16x3_cutile(x, parts) + cases.append((graph, x, w, out)) + for _ in range(3): + for stream, (graph, x, _, _) in zip(streams, cases): + with torch.cuda.stream(stream): + x.normal_() + graph.replay() + torch.cuda.synchronize() + for _, x, w, out in cases: + _check(out, x, w) + + +@pytest.mark.parametrize("m", [32, 1024, 8192]) +def test_torch_compile(m): + torch.manual_seed(1) + x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 5120, device="cuda") * 0.02 + parts = split_fp32_to_bf16x3(w) + compiled = torch.compile(router_gemm_bf16x3_cutile, fullgraph=True) + _check(compiled(x, parts), x, w) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) 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..489e60460ca8 --- /dev/null +++ b/test/registered/kernels/ops/gemm/test_router_gemv.py @@ -0,0 +1,91 @@ +"""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, + split_fp32_to_bf16x3, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +# Beam router shape: 256 experts over a 5120 hidden size. +N, K = 256, 5120 + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, + reason="CUDA required", +) + + +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() + + +@requires_cuda +def test_bf16x3_split_reconstructs_fp32_exactly(): + _, w = _inputs(1) + w = torch.cat([w, w * 1e30, w * 1e-20, -w]) + 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) + # The CUDA kernel must round like the torch reference path. + assert torch.equal(w3.cpu(), split_fp32_to_bf16x3(w.cpu())) + + +@requires_cuda +@pytest.mark.parametrize("m", [1, 2, 4, 8]) +def test_fp32_gate_gemv_matches_fp64(m): + x, w = _inputs(m) + assert router_gemv_supported(x, w) + out = router_gemv(x, w) + assert out.dtype == torch.float32 + # Measured about 5e-7 on H200 and B300; a BF16 or TF32 weight read is ~1e-3. + assert _max_error(out, x, w) < 2e-6 + + +@requires_cuda +def test_fp32_gate_gemv_stops_at_its_row_path(): + x, w = _inputs(9) + assert not router_gemv_supported(x, w) + assert not router_gemv_supported(x[:1], w.bfloat16()) + + +@requires_cuda +@pytest.mark.parametrize("m", [9, 64, 1024, 16384]) +def test_bf16x3_router_gemm_matches_fp64(m): + x, w = _inputs(m) + out = router_gemm_bf16x3_cublas(x, w) + assert out.dtype == torch.float32 + # Measured up to about 5e-5 (tensor-core accumulation); an fp32 SGEMM is about 2e-5. + assert _max_error(out, x, w) < 1e-4 + assert _max_error(F.linear(x.float(), w.bfloat16().float()), x, w) > 1e-4 + + +@requires_cuda +def test_bf16x3_router_gemm_reads_in_place_weight_updates(): + x, w = _inputs(64) + router_gemm_bf16x3_cublas(x, w) + # Online weight updates copy into the gate in place, without a reload hook. + w.copy_(-2 * w) + assert _max_error(router_gemm_bf16x3_cublas(x, w), x, w) < 1e-4 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py b/test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py new file mode 100644 index 000000000000..7d7241659e00 --- /dev/null +++ b/test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py @@ -0,0 +1,97 @@ +"""Numerical and graph-replay contract for full-precision router weights.""" + +import unittest + +import torch + +from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +def _tiny_fp32(x, w): + return tiny_gemm_bf16(x, w, out_dtype=torch.float32) + + +@unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9, + "SM90+ required", +) +class TestTinyGemmFp32(CustomTestCase): + def setUp(self): + torch.manual_seed(42) + + def test_fp64_reference(self): + w = torch.randn(256, 5120, device="cuda") / 5120**0.5 + for m in range(1, 17): + with self.subTest(m=m): + x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) + result = _tiny_fp32(x, w) + self.assertEqual(result.dtype, torch.float32) + reference = (x.double() @ w.double().T).float() + torch.testing.assert_close(result, reference, atol=2e-6, rtol=2e-5) + + def test_other_shapes_and_strided_activations(self): + for n, k in ((128, 1024), (384, 7168)): + with self.subTest(n=n, k=k): + x = torch.randn(3, k + 32, device="cuda", dtype=torch.bfloat16)[:, :k] + w = torch.randn(n, k, device="cuda") / k**0.5 + result = _tiny_fp32(x, w) + torch.testing.assert_close( + result, (x.double() @ w.double().T).float(), atol=2e-6, rtol=2e-5 + ) + + def test_weights_remain_fp32(self): + # All these weights round to 1 in BF16. Preserving their low FP32 bits + # is necessary to distinguish the two highest-scoring experts. + x = torch.zeros(1, 5120, device="cuda", dtype=torch.bfloat16) + x[:, 0] = 1 + w = torch.ones(256, 5120, device="cuda") + w[:, 0] += torch.arange(256, device="cuda") * 2**-20 + result = _tiny_fp32(x, w) + torch.testing.assert_close(result[0], w[:, 0], atol=0, rtol=0) + self.assertEqual(result.argmax().item(), 255) + self.assertEqual(torch.unique(w[:, 0].bfloat16()).numel(), 1) + + def test_graph_observes_input_and_weight_updates(self): + x = torch.randn(8, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 5120, device="cuda") / 5120**0.5 + _tiny_fp32(x, w) + next_x = torch.empty_like(x) + next_w = torch.empty_like(w) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + x.copy_(next_x) + w.copy_(next_w) + result = _tiny_fp32(x, w) + for _ in range(3): + next_x.normal_() + next_w.normal_(std=5120**-0.5) + graph.replay() + torch.testing.assert_close( + result, (x.double() @ w.double().T).float(), atol=2e-6, rtol=2e-5 + ) + + def test_rejects_invalid_layout_and_dtype(self): + x = torch.zeros(1, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(256, 5120, device="cuda") + for bad_x, bad_w in ( + (x.float(), w), + (x, w.half()), + (torch.zeros(1, 10240, device="cuda", dtype=torch.bfloat16)[:, ::2], w), + (x, w.T.contiguous().T), + ( + torch.zeros(5121, device="cuda", dtype=torch.bfloat16)[1:].view( + 1, 5120 + ), + w, + ), + ): + with self.assertRaises(Exception): + _tiny_fp32(bad_x, bad_w) + + +if __name__ == "__main__": + unittest.main() From 3cd22c793de1db52359e7404ce868da842401034 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:44:28 -0700 Subject: [PATCH 03/40] Move the Dynamo and grpcio-tools pins out of the Beam image Beam serving does not use Dynamo, and the v0.5.22 release already resolves grpcio-tools 1.84.0. The pins move to their own PR so they can be dropped without touching this layer. --- docker/beam/beam.Dockerfile | 28 ++++++---------------------- 1 file changed, 6 insertions(+), 22 deletions(-) diff --git a/docker/beam/beam.Dockerfile b/docker/beam/beam.Dockerfile index 1bb93a209d57..5c00cabf4a73 100644 --- a/docker/beam/beam.Dockerfile +++ b/docker/beam/beam.Dockerfile @@ -19,29 +19,13 @@ RUN git init -q /sgl-workspace/sglang-beam \ && uv pip install --no-deps -e /sgl-workspace/sglang-beam/python # Pins the beam branch adds on top of the v0.5.22 release: OpenAI 2.8.0 for the -# Responses shell types, the cuTile compiler for the Beam router GEMM, and a -# gRPC compiler that supports the Protobuf 7 dependency set. Record the new -# versions so later installs keep them. -RUN sed -i -E '/^(openai|nvidia-cuda-tileiras|grpcio-tools)==/d' /sgl-workspace/constraints.txt \ +# Responses shell types and the cuTile compiler for the Beam router GEMM. Record +# the new versions so later installs keep them. +RUN sed -i -E '/^(openai|nvidia-cuda-tileiras)==/d' /sgl-workspace/constraints.txt \ && uv pip install -c /sgl-workspace/constraints.txt \ - openai==2.8.0 nvidia-cuda-tileiras==13.4.92 grpcio-tools==1.84.0 \ - && uv pip freeze | grep -E '^(openai|nvidia-cuda-tileiras|grpcio-tools)==' \ + openai==2.8.0 nvidia-cuda-tileiras==13.4.92 \ + && uv pip freeze | grep -E '^(openai|nvidia-cuda-tileiras)==' \ >> /sgl-workspace/constraints.txt \ - && python3 -c "from importlib.metadata import version; assert version('openai') == '2.8.0'; assert version('nvidia-cuda-tileiras') == '13.4.92'; assert version('grpcio-tools') == '1.84.0'" - -# The release image's unbounded prerelease request can resolve to the empty -# ai-dynamo==0.0.0a0 placeholder, so pin a stable release. ai-dynamo-runtime -# caps pydantic below the frozen pin: release only the pydantic pins for this -# solve and record the versions it selects. Layered images that bring their -# own Dynamo can pass INSTALL_DYNAMO=0. -ARG INSTALL_DYNAMO=1 -RUN if [ "${INSTALL_DYNAMO}" = "1" ]; then \ - sed -i -E '/^pydantic(-core)?==/d' /sgl-workspace/constraints.txt \ - && uv pip install -c /sgl-workspace/constraints.txt \ - --index-strategy unsafe-best-match \ - --extra-index-url https://pypi.nvidia.com/ ai-dynamo==1.5.1 inflect \ - && uv pip freeze | grep -E '^pydantic(-core)?==' >> /sgl-workspace/constraints.txt \ - && python3 -c "import dynamo.runtime, inflect"; \ - fi + && python3 -c "from importlib.metadata import version; assert version('openai') == '2.8.0'; assert version('nvidia-cuda-tileiras') == '13.4.92'" WORKDIR /sgl-workspace/sglang-beam From 1a5baa70b01141143db0232d40400affa01d9fa4 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:44:30 -0700 Subject: [PATCH 04/40] Move the nvidia-cuda-tileiras pin to its own PR --- python/pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index fbb30fb2e434..c1cff46ed9c4 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -52,7 +52,6 @@ dependencies = [ "ninja", "numba==0.65.1", "numpy", - "nvidia-cuda-tileiras==13.4.92", "nvidia-cutlass-dsl[cu13]==4.8.0", "nvidia-mathdx==25.6.0", "nvidia-ml-py", From e5f9360f5107e58c591b455f4cc080a3a12758b8 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:53:06 -0700 Subject: [PATCH 05/40] Move the cuTile router GEMM to its own PR The kernel, its test and its benchmark need the tileiras compiler. They move with the router dispatch that uses them, so this layer needs no new dependency. --- .../kernels/ops/gemm/router_gemm_cutile.py | 217 ---------------- .../gemm/bench_router_gemm_cutile.py | 111 -------- .../ops/gemm/test_router_gemm_cutile.py | 237 ------------------ 3 files changed, 565 deletions(-) delete mode 100644 python/sglang/kernels/ops/gemm/router_gemm_cutile.py delete mode 100644 test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py delete mode 100644 test/registered/kernels/ops/gemm/test_router_gemm_cutile.py diff --git a/python/sglang/kernels/ops/gemm/router_gemm_cutile.py b/python/sglang/kernels/ops/gemm/router_gemm_cutile.py deleted file mode 100644 index ca5121f91add..000000000000 --- a/python/sglang/kernels/ops/gemm/router_gemm_cutile.py +++ /dev/null @@ -1,217 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""cuTile router GEMM: BF16 activations × FP32 weights → FP32 logits. - -The caller splits each FP32 weight into three BF16 components before serving. -Main-product K tiles are combined with explicit FP32 rounding for accuracy. -Split-K configurations use a separate reduction kernel. -""" - -import functools - -import cuda.tile as ct -import torch - -from sglang.srt.utils.custom_op import register_custom_op - - -@ct.kernel(num_ctas=1) -def _router_gemm_bf16x3_kernel( - X, - W, - WR, - Y, - BM: ct.Constant[int], - BN: ct.Constant[int], - BK: ct.Constant[int], - SPLIT_K: ct.Constant[int], - EXPERT_MAJOR: ct.Constant[bool], - PDL: ct.Constant[bool], -): - col, row, part = ct.bid(0), ct.bid(1), ct.bid(2) - if PDL: - ct.grid_dependency_control_launch_dependents() - ct.grid_dependency_control_wait() - steps = ct.cdiv(ct.cdiv(X.shape[1], BK), SPLIT_K) - if EXPERT_MAJOR: - total = ct.full((4 * BN, BM), 0, ct.float32) - else: - total = ct.full((BM, BN), 0, ct.float32) - residual = ct.full((BM, 2 * BN), 0, ct.float32) - for step in range(steps): - k = part * steps + step - x = ct.load( - X, - (row, k), - (BM, BK), - padding_mode=ct.PaddingMode.ZERO, - latency=4 if EXPERT_MAJOR or BK == 128 else 8, - ) - if EXPERT_MAJOR: - w = ct.load( - W, (0, col, k), (4, BN, BK), padding_mode=ct.PaddingMode.ZERO, latency=4 - ) - local = ct.mma( - ct.reshape(w, (4 * BN, BK)), - ct.transpose(x), - ct.full((4 * BN, BM), 0, ct.float32), - ) - else: - w0 = ct.load( - W, - (0, col, k), - (1, BN, BK), - padding_mode=ct.PaddingMode.ZERO, - latency=4 if BK == 128 else 8, - ) - wr = ct.load( - WR, - (0, col, k), - (2, BN, BK), - padding_mode=ct.PaddingMode.ZERO, - latency=4 if BK == 128 else 8, - ) - local = ct.mma( - x, - ct.transpose(ct.reshape(w0, (BN, BK))), - ct.full((BM, BN), 0, ct.float32), - ) - residual = ct.mma(x, ct.transpose(ct.reshape(wr, (2 * BN, BK))), residual) - # Round each main-product tile before adding the next tile. - total = ct.add(total, local, rounding_mode=ct.RoundingMode.RN) - if EXPERT_MAJOR: - main = ct.extract(total, (0, 0), (BN, BM)) - mid = ct.extract(total, (1, 0), (BN, BM)) - low = ct.extract(total, (2, 0), (BN, BM)) - value = ct.transpose( - ct.add( - main, - ct.add(mid, low, rounding_mode=ct.RoundingMode.RN), - rounding_mode=ct.RoundingMode.RN, - ) - ) - else: - correction = ct.sum( - ct.reshape(residual, (BM, 2, BN)), axis=1, rounding_mode=ct.RoundingMode.RN - ) - value = ct.add(total, correction, rounding_mode=ct.RoundingMode.RN) - ct.store( - Y, - (part, row, col), - ct.reshape(value, (1, BM, BN)), - allow_tma=not (BM == 256 and SPLIT_K > 1), - ) - - -@ct.kernel -def _split_k_reduce_kernel( - P, Y, SPLIT_K: ct.Constant[int], PDL: ct.Constant[bool] = False -): - if PDL: - ct.grid_dependency_control_wait() - if SPLIT_K == 2: - values = ct.load( - P, - (0, ct.bid(0)), - (SPLIT_K, 1024), - padding_mode=ct.PaddingMode.ZERO, - allow_tma=False, - ) - total = ct.sum(values, axis=0, rounding_mode=ct.RoundingMode.RN) - else: - total = ct.full((1024,), 0, ct.float32) - for part in range(SPLIT_K): - value = ct.load( - P, - (part, ct.bid(0)), - (1, 1024), - padding_mode=ct.PaddingMode.ZERO, - allow_tma=False, - ) - total = ct.add( - total, ct.reshape(value, (1024,)), rounding_mode=ct.RoundingMode.RN - ) - ct.store(Y, (ct.bid(0),), total, allow_tma=False) - - -_GEMM_KERNELS = { - 1: _router_gemm_bf16x3_kernel, - 2: _router_gemm_bf16x3_kernel.replace_hints(num_ctas=2), -} -# (max M, (BM, BN, BK, SPLIT_K, CTAS)); larger M uses the fallback below. -_BLACKWELL_CONFIGS = ( - (64, (32, 16, 256, 4, 1)), - (128, (64, 32, 256, 4, 1)), - (512, (128, 64, 256, 8, 1)), - (1024, (128, 64, 128, 4, 1)), - (2048, (128, 64, 128, 2, 1)), - (4096, (256, 128, 256, 2, 2)), -) -# Use smaller tiles on Hopper to limit register pressure. -_HOPPER_CONFIGS = ( - (256, (64, 64, 128, 8, 1)), - (512, (64, 64, 128, 4, 1)), - (1024, (64, 64, 128, 2, 1)), -) - - -@functools.cache -def _is_hopper(device_index: int) -> bool: - return torch.cuda.get_device_capability(device_index)[0] == 9 - - -def _select_launch_config(m: int, device: torch.device) -> tuple: - if m > 128 and _is_hopper(device.index): - configs, fallback = _HOPPER_CONFIGS, (64, 64, 128, 1, 1) - else: - configs, fallback = _BLACKWELL_CONFIGS, (256, 128, 256, 1, 2) - return next((config for limit, config in configs if m <= limit), fallback) - - -def _router_gemm_bf16x3_fake(x: torch.Tensor, parts: torch.Tensor) -> torch.Tensor: - return x.new_empty((x.shape[0], parts.shape[0] // 3), dtype=torch.float32) - - -@register_custom_op(fake_impl=_router_gemm_bf16x3_fake) -def router_gemm_bf16x3_cutile(x: torch.Tensor, parts: torch.Tensor) -> torch.Tensor: - """Multiply BF16 [M, K] by prepared BF16 [3N, K], producing FP32 [M, N].""" - m, k = x.shape - n = parts.shape[0] // 3 - bm, bn, bk, splits, ctas = _select_launch_config(m, x.device) - pdl = bk == 128 and splits == 4 - out = torch.empty((m, n), device=x.device, dtype=torch.float32) - if m == 0 or n == 0: - return out - partial = ( - torch.empty((splits, m, n), device=x.device, dtype=torch.float32) - if splits > 1 - else out.unsqueeze(0) - ) - with torch.cuda.device(x.device): - stream = torch.cuda.current_stream() - ct.launch( - stream, - (ct.cdiv(n, bn), ct.cdiv(m, bm), splits), - _GEMM_KERNELS[ctas], - ( - x, - parts.view(3, n, k), - parts[n:].view(2, n, k), - partial, - bm, - bn, - bk, - splits, - m <= 64, - pdl, - ), - programmatic_dependent_launch=pdl, - ) - if splits > 1: - ct.launch( - stream, - (ct.cdiv(out.numel(), 1024),), - _split_k_reduce_kernel, - (partial.view(splits, -1), out.view(-1), splits, pdl), - programmatic_dependent_launch=pdl, - ) - return out diff --git a/test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py b/test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py deleted file mode 100644 index c651c900907f..000000000000 --- a/test/registered/kernels/benchmark/gemm/bench_router_gemm_cutile.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Compare router GEMM latency and mean absolute error against FP64. - -Run this file on Hopper or Blackwell. Weight preparation is outside the timed region; -split-K reduction is included. Latencies use rotating inputs and CUDA graphs. -All providers use the same random inputs (seed 0, X std 1, W std 0.02). -""" - -import torch - -from sglang.kernels.jit.benchmark import marker -from sglang.kernels.jit.benchmark.utils import create_random, get_benchmark_range -from sglang.kernels.ops.gemm.router_gemm_cutile import router_gemm_bf16x3_cutile -from sglang.kernels.ops.gemm.router_gemv import split_fp32_to_bf16x3 -from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=30, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" -) - -PROVIDERS = ["SGLang tiny FP32", "Cast + TF32", "cuBLAS BF16x3", "cuTile BF16x3"] - - -def _cast_tf32(x, w): - return x.float() @ w.T - - -def _cublas(x, parts): - n = parts.shape[0] // 3 - return torch.mm(x, parts.T, out_dtype=torch.float32).view(-1, 3, n).sum(1) - - -def _print_table(title, rows, number_format): - print(f"\n### {title}\n") - print("| Tokens | " + " | ".join(PROVIDERS) + " |") - print("| ---: " * (len(PROVIDERS) + 1) + "|") - for m, values in rows: - cells = ["—" if v is None else format(v, number_format) for v in values] - print(f"| {m} | " + " | ".join(cells) + " |") - - -def benchmark(batch_sizes, n, k): - torch.manual_seed(0) - w = create_random(n, k, dtype=torch.float32) * 0.02 - parts = split_fp32_to_bf16x3(w) - latencies, errors = [], [] - old_tf32 = torch.backends.cuda.matmul.allow_tf32 - try: - torch.backends.cuda.matmul.allow_tf32 = True - for m in batch_sizes: - x = create_random(m, k, dtype=torch.bfloat16) - ref = x.double() @ w.double().T - cases = [ - ( - tiny_gemm_bf16, - w, - {"out_dtype": torch.float32, "max_m": 16 if m <= 16 else 32}, - ), - (_cast_tf32, w, {}), - (_cublas, parts, {}), - (router_gemm_bf16x3_cutile, parts, {}), - ] - times, maes = [], [] - for fn, weight, kwargs in cases: - if fn is tiny_gemm_bf16 and m > 32: - times.append(None) - maes.append(None) - continue - result = marker.do_bench( - fn, input_args=(x, weight), input_kwargs=kwargs - ) - out = fn(x, weight, **kwargs) - times.append(result.times[0] * 1e6) - maes.append((out.double() - ref).abs().mean().item()) - latencies.append((m, times)) - errors.append((m, maes)) - finally: - torch.backends.cuda.matmul.allow_tf32 = old_tf32 - print(f"\n## K{k}/N{n}") - _print_table("Latency (µs)", latencies, ".2f") - _print_table("Mean absolute error vs FP64", errors, ".2e") - - -if __name__ == "__main__": - major = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0 - if major not in (9, 10): - print("Hopper or Blackwell required.") - else: - batch_sizes = get_benchmark_range( - [ - 1, - 2, - 4, - 8, - 16, - 32, - 64, - 128, - 256, - 512, - 1024, - 2048, - 4096, - 8192, - 16384, - 16896, - ], - [1, 32, 1024, 16384], - ) - for n, k in get_benchmark_range([(256, 5120), (128, 6144)], [(256, 5120)]): - benchmark(batch_sizes, n, k) diff --git a/test/registered/kernels/ops/gemm/test_router_gemm_cutile.py b/test/registered/kernels/ops/gemm/test_router_gemm_cutile.py deleted file mode 100644 index 2f68ea25b35b..000000000000 --- a/test/registered/kernels/ops/gemm/test_router_gemm_cutile.py +++ /dev/null @@ -1,237 +0,0 @@ -"""Accuracy and CUDA graph tests for the prepared BF16x3 router GEMM.""" - -import sys - -import pytest -import torch - -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large") - -pytest.importorskip("cuda.tile") - -import cuda.tile as ct # noqa: E402 - -from sglang.kernels.ops.gemm.router_gemm_cutile import ( # noqa: E402 - router_gemm_bf16x3_cutile, -) -from sglang.kernels.ops.gemm.router_gemv import split_fp32_to_bf16x3 # noqa: E402 - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() - or torch.cuda.get_device_capability()[0] not in (9, 10), - reason="Hopper or Blackwell required", -) - - -def _check(out, x, w): - ref = x.double() @ w.double().T - error = out.double() - ref - torch.testing.assert_close(out.double(), ref, atol=8e-6, rtol=3e-6) - assert (error.square().mean() / ref.square().mean()).sqrt().item() < 1e-6 - - -@pytest.mark.parametrize( - "m", - [ - 1, - 2, - 8, - 16, - 32, - 64, - 65, - 128, - 129, - 512, - 513, - 1024, - 1025, - 2048, - 2049, - 4096, - 4097, - 8192, - 16384, - ], -) -def test_router_shape(m): - torch.manual_seed(1) - x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) - w = torch.randn(256, 5120, device="cuda") * 0.02 - out = router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)) - assert out.dtype == torch.float32 - _check(out, x, w) - - -@pytest.mark.parametrize( - "m,k,n", - [ - (2, 137, 160), - (17, 2816, 256), - (129, 5120, 160), - (257, 5128, 160), - (257, 6144, 128), - (513, 5128, 160), - (1025, 137, 160), - (8193, 5120, 256), - ], -) -def test_tail_tiles(m, k, n): - torch.manual_seed(1) - x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) - w = torch.randn(n, k, device="cuda") / k**0.5 - _check(router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)), x, w) - - -def test_strided_input(): - torch.manual_seed(1) - x = torch.randn(129, 5144, device="cuda", dtype=torch.bfloat16)[:, :5128] - w = torch.randn(160, 5128, device="cuda") * 0.02 - _check(router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)), x, w) - - -@pytest.mark.parametrize("m,n", [(0, 256), (32, 0)]) -def test_empty_output(m, n): - x = torch.empty(m, 5120, device="cuda", dtype=torch.bfloat16) - parts = torch.empty(3 * n, 5120, device="cuda", dtype=torch.bfloat16) - out = router_gemm_bf16x3_cutile(x, parts) - assert out.shape == (m, n) - assert out.dtype == torch.float32 - assert out.device == x.device - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two GPUs required") -def test_non_current_device(): - current = torch.cuda.current_device() - with torch.cuda.device((current + 1) % torch.cuda.device_count()): - x = torch.randn(32, 5120, device="cuda", dtype=torch.bfloat16) - w = torch.randn(256, 5120, device="cuda") * 0.02 - parts = split_fp32_to_bf16x3(w) - assert torch.cuda.current_device() == current - _check(router_gemm_bf16x3_cutile(x, parts), x, w) - - -@pytest.mark.parametrize("m", [1, 2, 8, 16, 32, 128, 512, 1024, 2048, 16384]) -@pytest.mark.parametrize("k,distance", [(512, 256), (5120, 320), (5120, 2560)]) -def test_cancellation(m, k, distance, record_property): - torch.manual_seed(1) - blocks = k // (2 * distance) - half = torch.randn(m, blocks, distance, device="cuda", dtype=torch.bfloat16) - x = torch.cat((half, half), dim=-1).reshape(m, k) - main = torch.randn( - 256, blocks, distance, device="cuda", dtype=torch.bfloat16 - ).float() - w = torch.cat((main, -main + torch.randn_like(main) * 1e-4), dim=-1).reshape(256, k) - if k != 512: - w = w / k**0.5 - ref = x.double() @ w.double().T - old_tf32 = torch.backends.cuda.matmul.allow_tf32 - try: - torch.backends.cuda.matmul.allow_tf32 = False - baseline_error = (x.float() @ w.T - ref).abs().max() - finally: - torch.backends.cuda.matmul.allow_tf32 = old_tf32 - error = (router_gemm_bf16x3_cutile(x, split_fp32_to_bf16x3(w)) - ref).abs() - record_property("fp32_max_error", baseline_error.item()) - record_property("candidate_max_error", error.max().item()) - record_property("candidate_mae", error.mean().item()) - assert error.mean().item() < 1e-6 - - -@pytest.mark.parametrize("m", [64, 512, 1024, 2048, 4096, 8192]) -def test_cuda_graph_replay(m): - torch.manual_seed(1) - x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) - w = torch.randn(256, 5120, device="cuda") * 0.02 - parts = split_fp32_to_bf16x3(w) - router_gemm_bf16x3_cutile(x, parts) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = router_gemm_bf16x3_cutile(x, parts) - for _ in range(2): - x.normal_() - w.normal_(std=0.02) - split_fp32_to_bf16x3(w, out=parts) - graph.replay() - _check(out, x, w) - - -@ct.kernel -def _pdl_producer(source, output): - ct.grid_dependency_control_launch_dependents() - # Keep input production pending after releasing the dependent GEMM. - delay = ct.astype(ct.load(source, (0,), (32,)), ct.float32) - for _ in range(1024): - delay = ct.sin(delay) + 0.001 - offset = ct.sum(delay, axis=0) / 32 - for block in range(ct.bid(0), ct.cdiv(source.shape[0], 1024), 32): - value = ct.load(source, (block,), (1024,), padding_mode=ct.PaddingMode.ZERO) - ct.store(output, (block,), ct.astype(value + offset, ct.bfloat16)) - - -@pytest.mark.parametrize("m", [512, 1024]) -def test_pdl_producer_graph_replay(m): - torch.manual_seed(1) - source = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) - x = torch.empty_like(source) - w = torch.randn(256, 5120, device="cuda") * 0.02 - parts = split_fp32_to_bf16x3(w) - ct.launch( - torch.cuda.current_stream(), (32,), _pdl_producer, (source.view(-1), x.view(-1)) - ) - torch.cuda.synchronize() - router_gemm_bf16x3_cutile(x, parts) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - ct.launch( - torch.cuda.current_stream(), - (32,), - _pdl_producer, - (source.view(-1), x.view(-1)), - ) - out = router_gemm_bf16x3_cutile(x, parts) - for _ in range(3): - source.normal_() - x.fill_(float("nan")) - graph.replay() - _check(out, x, w) - - -def test_concurrent_graphs(): - streams = [torch.cuda.Stream(), torch.cuda.Stream()] - cases = [] - for stream in streams: - with torch.cuda.stream(stream): - x = torch.randn(1024, 5120, device="cuda", dtype=torch.bfloat16) - w = torch.randn(256, 5120, device="cuda") * 0.02 - parts = split_fp32_to_bf16x3(w) - router_gemm_bf16x3_cutile(x, parts) - stream.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - out = router_gemm_bf16x3_cutile(x, parts) - cases.append((graph, x, w, out)) - for _ in range(3): - for stream, (graph, x, _, _) in zip(streams, cases): - with torch.cuda.stream(stream): - x.normal_() - graph.replay() - torch.cuda.synchronize() - for _, x, w, out in cases: - _check(out, x, w) - - -@pytest.mark.parametrize("m", [32, 1024, 8192]) -def test_torch_compile(m): - torch.manual_seed(1) - x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) - w = torch.randn(256, 5120, device="cuda") * 0.02 - parts = split_fp32_to_bf16x3(w) - compiled = torch.compile(router_gemm_bf16x3_cutile, fullgraph=True) - _check(compiled(x, parts), x, w) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v"])) From cb50998cec6a97a601e1f44a79ae3e0f6fb76e8d Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:40:56 -0700 Subject: [PATCH 06/40] Drop the Beam image from this PR The image is #45253 on its own; nothing in this PR uses it. --- docker/beam/beam.Dockerfile | 31 ------------------------------- 1 file changed, 31 deletions(-) delete mode 100644 docker/beam/beam.Dockerfile diff --git a/docker/beam/beam.Dockerfile b/docker/beam/beam.Dockerfile deleted file mode 100644 index 5c00cabf4a73..000000000000 --- a/docker/beam/beam.Dockerfile +++ /dev/null @@ -1,31 +0,0 @@ -# syntax=docker/dockerfile:1 -# Beam serving image (CUDA). Build from the repository root; application source -# is cloned from upstream: -# docker buildx build --platform linux/amd64 --load \ -# -f docker/beam/beam.Dockerfile -t beam:local-test . -FROM lmsysorg/sglang:v0.5.22@sha256:e3a4f06001a29542fc4a8638536d9ffc5eb9c9942b08603c09b19a56eae4a1ae AS base - -SHELL ["/bin/bash", "-o", "pipefail", "-c"] - -# Replace the stock SGLang with the upstream beam branch. Pass a commit SHA as -# BEAM_REF to pin the source of a published image. -ARG BEAM_REF=beam -WORKDIR /sgl-workspace -RUN git init -q /sgl-workspace/sglang-beam \ - && git -C /sgl-workspace/sglang-beam fetch -q --depth 1 \ - https://github.com/sgl-project/sglang.git "${BEAM_REF}" \ - && git -C /sgl-workspace/sglang-beam checkout -q --detach FETCH_HEAD \ - && uv pip uninstall sglang \ - && uv pip install --no-deps -e /sgl-workspace/sglang-beam/python - -# Pins the beam branch adds on top of the v0.5.22 release: OpenAI 2.8.0 for the -# Responses shell types and the cuTile compiler for the Beam router GEMM. Record -# the new versions so later installs keep them. -RUN sed -i -E '/^(openai|nvidia-cuda-tileiras)==/d' /sgl-workspace/constraints.txt \ - && uv pip install -c /sgl-workspace/constraints.txt \ - openai==2.8.0 nvidia-cuda-tileiras==13.4.92 \ - && uv pip freeze | grep -E '^(openai|nvidia-cuda-tileiras)==' \ - >> /sgl-workspace/constraints.txt \ - && python3 -c "from importlib.metadata import version; assert version('openai') == '2.8.0'; assert version('nvidia-cuda-tileiras') == '13.4.92'" - -WORKDIR /sgl-workspace/sglang-beam From 82fa70c16de123587f6b07775a10a17b115d72c1 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:15:32 -0700 Subject: [PATCH 07/40] Drop the #42930 pick in zero_pages.py main has the merged #42930 version. --- python/sglang/kernels/ops/kvcache/zero_pages.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/python/sglang/kernels/ops/kvcache/zero_pages.py b/python/sglang/kernels/ops/kvcache/zero_pages.py index 2c456b23022e..afa3c0071085 100644 --- a/python/sglang/kernels/ops/kvcache/zero_pages.py +++ b/python/sglang/kernels/ops/kvcache/zero_pages.py @@ -1,8 +1,8 @@ """Zero whole page envelopes of the unified pool by physical page id. -The pool is viewed as int64 words (MHA and MLA page envelopes are -8-byte-aligned), one wide element per lane; grid = (num_pages, page word -blocks). +The pool is viewed as int64 words (the MLA page envelope is always +8-byte-aligned: entry bytes per layer = kv_cache_dim * itemsize, a multiple +of 8), one wide element per lane; grid = (num_pages, page word blocks). """ from __future__ import annotations @@ -41,9 +41,6 @@ def zero_pages( if m == 0: return assert raw.dtype == torch.uint8, f"expected uint8 pool, got {raw.dtype}" - if not raw.is_cuda: - raw[: num_pages * page_bytes].view(num_pages, page_bytes)[pages] = 0 - return assert page_bytes % 8 == 0, f"page_bytes {page_bytes} not int64-aligned" page_words = page_bytes // 8 words = raw[: num_pages * page_bytes].view(torch.int64) From b80a2442122b0b7a52789a9fda6d6682b2695e0c Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:15:32 -0700 Subject: [PATCH 08/40] Drop the kernels layer's tests and benchmarks --- test/manual/kernels/bench_fp32_router.py | 67 -------- .../attention/bench_trtllm_mha_mixed.py | 142 ----------------- .../elementwise/bench_beam_attention_gate.py | 40 ----- .../bench_beam_sandwich_residual_rmsnorm.py | 32 ---- .../elementwise/test_fused_op_gpu_parity.py | 147 ------------------ .../kernels/ops/gemm/test_router_gemv.py | 91 ----------- .../kernels/ops/gemm/test_tiny_gemm_fp32.py | 97 ------------ 7 files changed, 616 deletions(-) delete mode 100644 test/manual/kernels/bench_fp32_router.py delete mode 100644 test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py delete mode 100644 test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py delete mode 100644 test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py delete mode 100644 test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py delete mode 100644 test/registered/kernels/ops/gemm/test_router_gemv.py delete mode 100644 test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py diff --git a/test/manual/kernels/bench_fp32_router.py b/test/manual/kernels/bench_fp32_router.py deleted file mode 100644 index a8af4d3665f5..000000000000 --- a/test/manual/kernels/bench_fp32_router.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Compare cuBLAS and mixed-weight tiny_gemm. - -Run with the candidate Python package on CUDA SM90+: - python test/manual/kernels/bench_fp32_router.py - -Times are microseconds per layer under CUDA graph replay. The baseline includes -BF16-to-FP32 activation conversion; weights stay FP32 for all implementations. -The 245 MiB weight working set exceeds H200/B300 L2. Compilation is untimed. -""" - -import json - -import torch - -from sglang.kernels.jit.benchmark import marker -from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 - - -def bench(fn, weights): - def run_layers(): - for weight in weights: - fn(weight) - - result = marker.do_bench( - run_layers, - warmup_iters=2, - metrics=(0.5,), - memory_output=None, - ) - return result.times[0] * 1e6 / len(weights) - - -def main(): - torch.manual_seed(42) - torch.backends.cuda.matmul.allow_tf32 = False - print( - json.dumps({"gpu": torch.cuda.get_device_name(), "torch": torch.__version__}), - flush=True, - ) - weights = [ - torch.randn(256, 5120, device="cuda", dtype=torch.float32) / 5120**0.5 - for _ in range(49) - ] - for m in range(1, 17): - x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) - base = bench(lambda w: torch.nn.functional.linear(x.float(), w), weights) - ref = x.double() @ weights[0].double().T - out = tiny_gemm_bf16(x, weights[0], out_dtype=torch.float32) - err = (out.double() - ref).abs().max().item() - us = bench(lambda w: tiny_gemm_bf16(x, w, out_dtype=torch.float32), weights) - print( - json.dumps( - dict( - m=m, - variant="tiny", - base_us=base, - us=us, - speedup=base / us, - max_error=err, - ) - ), - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py b/test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py deleted file mode 100644 index 8f9a2308d37a..000000000000 --- a/test/manual/kernels/benchmark/attention/bench_trtllm_mha_mixed.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Compare split and unsplit mixed-extend attention on SM100-family GPUs. - -Times forward_extend with CUDA graphs and cold L2, excluding cache flushes, -projection, KV writes and metadata construction. Run from the repository root -with PYTHONPATH=python. -""" - -import argparse -import gc -import json -import unittest - -import torch - -from sglang.kernels.jit.benchmark import marker -from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend -from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool -from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.srt.utils.common import is_sm100_supported -from sglang.test.kits.attention_unittest.attention_methods.dense_attention import ( - DenseAttentionCase, - build_dense_attention_fixture, -) - - -@torch.no_grad() -def benchmark(args, num_decode): - case = DenseAttentionCase( - name="mixed_extend", - backend="trtllm_mha", - forward_mode=ForwardMode.MIXED, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - page_size=64, - prefix_lens=(0,) + (args.kv_len - 1,) * num_decode, - extend_lens=(args.prefill_tokens,) + (1,) * num_decode, - ) - fixture = build_dense_attention_fixture( - unittest.TestCase(), - case, - head_dim=args.head_dim, - hidden_size=128, - max_context_len=((max(args.kv_len, args.prefill_tokens) + 63) // 64) * 64, - dtype=torch.bfloat16, - ) - runner = fixture.runner - layer = fixture.actual_module.attn - layer.sliding_window_size = args.window_size - if args.kv_cache_dtype == "fp8": - runner.kv_cache_dtype = torch.float8_e4m3fn - runner.token_to_kv_pool = MHATokenToKVPool( - size=runner.token_to_kv_pool.size, - page_size=64, - dtype=runner.kv_cache_dtype, - head_num=args.num_kv_heads, - head_dim=args.head_dim, - layer_num=1, - device="cuda", - enable_memory_saver=False, - enable_alt_stream=False, - ) - runner.init_kv_index_translator() - fixture.backend = TRTLLMHAAttnBackend(runner) - for req, prefix in enumerate(fixture.prefix_hidden): - if prefix.shape[0]: - _, k, v = fixture.actual_module.project_qkv(prefix) - runner.token_to_kv_pool.set_kv_buffer( - layer, - runner.req_to_token_pool.req_to_token[req, : len(prefix)].long(), - k.view(-1, args.num_kv_heads, args.head_dim), - v.view(-1, args.num_kv_heads, args.head_dim), - ) - backend = fixture.backend - if not is_sm100_supported(): - raise RuntimeError("The decode-tail split requires an SM100-family GPU") - q, k, v = fixture.actual_module.project_qkv(fixture.input_hidden) - batch = fixture.forward_batch - backend.init_forward_metadata(batch) - backend.forward_extend(q, k, v, layer, batch) - - def run(): - return backend.forward_extend(q, None, None, layer, batch, save_kv_cache=False) - - outputs, timings = {}, {} - for split in (False, True): - backend.init_forward_metadata(batch) - if not split: - backend.forward_metadata.mixed_prefill_reqs = -1 - outputs[split] = run().clone() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - # Flush before each replay, not just before a graph of repeated calls. - timings[split] = ( - marker.do_bench( - graph.replay, - use_cuda_graph=False, - estimated_time_ms=200, - disable_log_bandwidth=True, - ).times[0] - * 1e6 - ) - torch.testing.assert_close(outputs[True], outputs[False], atol=1e-2, rtol=1e-2) - delta = outputs[True].float() - outputs[False].float() - return { - "decode_requests": num_decode, - "prefill_tokens": args.prefill_tokens, - "kv_len": args.kv_len, - "kv_cache_dtype": args.kv_cache_dtype, - "window_size": args.window_size, - "num_heads": args.num_heads, - "num_kv_heads": args.num_kv_heads, - "head_dim": args.head_dim, - "unsplit_us": timings[False], - "split_us": timings[True], - "speedup": timings[False] / timings[True], - "max_abs_diff": delta.abs().max().item(), - "relative_rms_diff": (delta.norm() / outputs[False].float().norm()).item(), - } - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--decode-requests", type=int, nargs="+", default=[0, 1, 32, 128, 256, 511] - ) - parser.add_argument("--prefill-tokens", type=int, default=8192) - parser.add_argument("--kv-len", type=int, default=8192) - parser.add_argument("--kv-cache-dtype", choices=["bf16", "fp8"], default="fp8") - parser.add_argument("--window-size", type=int, default=-1) - parser.add_argument("--num-heads", type=int, default=12) - parser.add_argument("--num-kv-heads", type=int, default=2) - parser.add_argument("--head-dim", type=int, default=128) - args = parser.parse_args() - for num_decode in args.decode_requests: - print(json.dumps(benchmark(args, num_decode)), flush=True) - gc.collect() - torch.cuda.empty_cache() - - -if __name__ == "__main__": - main() diff --git a/test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py b/test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py deleted file mode 100644 index 8f2da343704e..000000000000 --- a/test/registered/kernels/benchmark/elementwise/bench_beam_attention_gate.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Beam attention-gate latency at decode and prefill sizes.""" - -import torch - -from sglang.kernels.jit.benchmark import marker -from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE -from sglang.kernels.spec import KernelBackend -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=7, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" -) - - -@marker.parametrize("rows", [1, 16, 64, 128, 1024, 8192], [1, 128]) -@marker.parametrize("packed_gate", [False, True]) -@marker.benchmark("impl", ["triton", "torch"]) -def benchmark(rows: int, packed_gate: bool, impl: str): - x = torch.randn(rows, 6 * 128, dtype=torch.bfloat16, device="cuda") - if packed_gate: - # TP8 QKV+gate widths: Q=768, K=128, V=128, gate=768. - gate_source = torch.randn(rows, 1792, dtype=x.dtype, device=x.device) - else: - gate_source = torch.randn_like(x) - gate_start = gate_source.shape[1] - x.shape[1] - bias = torch.randn(6, dtype=x.dtype, device="cuda") - # The marker rotates clones, and a clone of a strided slice is contiguous; - # rotate the backing tensor and slice the gate inside the timed call. - return marker.do_bench( - lambda x, gate_source, bias: _BEAM_ATTENTION_GATE( - x, gate_source[:, gate_start:], bias, 128, backend=KernelBackend(impl) - ), - input_args=(x, gate_source, bias), - memory_args=(x, gate_source[:, gate_start:], bias), - memory_output="out", - ) - - -if __name__ == "__main__": - benchmark.run() diff --git a/test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py b/test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py deleted file mode 100644 index be744a7ce01f..000000000000 --- a/test/registered/kernels/benchmark/layernorm/bench_beam_sandwich_residual_rmsnorm.py +++ /dev/null @@ -1,32 +0,0 @@ -"""Beam norm/residual latency at decode and prefill sizes.""" - -import torch - -from sglang.kernels.jit.benchmark import marker -from sglang.kernels.ops.layernorm import _BEAM_SANDWICH_RESIDUAL_RMSNORM -from sglang.kernels.spec import KernelBackend -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" -) - - -@marker.parametrize("rows", [1, 16, 64, 128, 1024, 8192], [1, 128]) -@marker.benchmark("impl", ["triton", "torch"]) -def benchmark(rows: int, impl: str): - x = torch.randn(rows, 5120, dtype=torch.bfloat16, device="cuda") - residual = torch.randn_like(x) - weight = torch.ones(5120, dtype=x.dtype, device="cuda") - return marker.do_bench( - lambda *args: _BEAM_SANDWICH_RESIDUAL_RMSNORM( - *args, backend=KernelBackend(impl) - ), - input_args=(x, residual, weight, weight, 1e-6, (0.25, 0.968)), - memory_args=(x, residual, weight), - memory_output="out", - ) - - -if __name__ == "__main__": - benchmark.run() diff --git a/test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py b/test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py deleted file mode 100644 index 70c953e04337..000000000000 --- a/test/registered/kernels/ops/elementwise/test_fused_op_gpu_parity.py +++ /dev/null @@ -1,147 +0,0 @@ -"""Every-backend-vs-native parity for BaseFusedOp ops on real GPU (RFC #29630). - -For each reworked fused op, run every backend eligible on this platform and -assert it matches the pure-torch ``forward_native`` reference within dtype -tolerance. New backends are picked up automatically. -""" - -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=13, stage="extra-a", runner_config="1-gpu-small") - -_SKIP = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE} - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -def _eligible(op): - return [ - b for b in op.available_backends() if b not in _SKIP and op.backend_eligible(b) - ] - - -@pytest.mark.parametrize("rows", [0, 1, 64, 1024]) -def test_beam_attention_gate(rows): - from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE as op - - generator = torch.Generator().manual_seed(5) - bias = torch.randn(12, generator=generator).to(device="cuda", dtype=torch.bfloat16) - attention = torch.randn(rows, 12 * 128, generator=generator).to( - device="cuda", dtype=torch.bfloat16 - ) - gate = ( - torch.randn(rows, 12 * 128, generator=generator).to( - device="cuda", dtype=torch.bfloat16 - ) - * 2 - ) - expected = op.forward_native(attention, gate, bias, 128) - for backend in _eligible(op): - torch.testing.assert_close( - op(attention, gate, bias, 128, backend=backend), expected, atol=0, rtol=0 - ) - - -@pytest.mark.parametrize("rows", [1, 16, 64]) -@pytest.mark.parametrize("tp_size", [4, 8]) -@pytest.mark.parametrize("column_step", [1, 2]) -def test_beam_packed_gate_without_copy(rows, tp_size, column_step): - """A QKV+gate slice must preserve rounding without materializing the gate.""" - from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE as op - from sglang.kernels.ops.elementwise.beam_fused import fused_attention_gate - - q_size, kv_size = 48 * 128 // tp_size, 8 * 128 // tp_size - packed = torch.randn( - rows, - (2 * q_size + 2 * kv_size) * column_step, - device="cuda", - dtype=torch.bfloat16, - ) - gate = packed[:, (q_size + 2 * kv_size) * column_step :: column_step] - attention = torch.randn(rows, q_size, device="cuda", dtype=gate.dtype) - bias = torch.randn(q_size // 128, device="cuda", dtype=gate.dtype) - expected = op.forward_native(attention, gate, bias, 128) - fused_attention_gate(attention, gate, bias, 128) # Compile before profiling. - with torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CPU] - ) as prof: - actual = fused_attention_gate(attention, gate, bias, 128) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - assert not {"aten::clone", "aten::copy_"} & { - event.key for event in prof.key_averages() - } - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - replayed = fused_attention_gate(attention, gate, bias, 128) - packed.normal_() - graph.replay() - torch.testing.assert_close( - replayed, op.forward_native(attention, gate, bias, 128), atol=0, rtol=0 - ) - - -def test_beam_fp16_gate_rounding_boundaries(): - from sglang.kernels.ops.elementwise import _BEAM_ATTENTION_GATE as op - - gate = torch.tensor( - [[0.0068359375, -0.00634765625]], device="cuda", dtype=torch.float16 - ) - attention = torch.ones_like(gate) - bias = torch.zeros(1, device="cuda", dtype=torch.float16) - expected = op.forward_native(attention, gate, bias, 2) - for backend in _eligible(op): - torch.testing.assert_close( - op(attention, gate, bias, 2, backend=backend), expected, atol=0, rtol=0 - ) - - -def test_beam_custom_op_contracts(): - from sglang.kernels.ops.elementwise.beam_fused import fused_attention_gate - - # Sliced rows cover real/fake output strides. - x = torch.randn(8, 5120, device="cuda", dtype=torch.bfloat16)[::2] - bias = torch.zeros(40, device="cuda", dtype=x.dtype) - torch.library.opcheck(fused_attention_gate, (x, x.clone(), bias, 128)) - - -@pytest.mark.parametrize("forced_backend", [False, True]) -def test_beam_compiled_rounding(forced_backend): - """Outer compilation must preserve sigmoid rounding.""" - from torch._inductor import config as inductor_config - - from sglang.kernels.ops.elementwise import BeamAttentionGateOp - - op = BeamAttentionGateOp() - gate = torch.tensor( - [[0.0068359375, -0.00634765625]], device="cuda", dtype=torch.float16 - ) - args = ( - torch.ones_like(gate), - gate, - torch.zeros(1, device="cuda", dtype=gate.dtype), - 2, - ) - expected = op.forward_native(*args) - op(*args) # Warm backend imports before full-graph capture. - op.enter_torch_compile(1) - try: - with inductor_config.patch(compile_threads=1): - actual = ( - op(*args, backend=KernelBackend.TORCH_COMPILE) - if forced_backend - else torch.compile(op, fullgraph=True)(*args) - ) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - finally: - op.leave_torch_compile() - - -if __name__ == "__main__": - import sys - - sys.exit(pytest.main([__file__])) diff --git a/test/registered/kernels/ops/gemm/test_router_gemv.py b/test/registered/kernels/ops/gemm/test_router_gemv.py deleted file mode 100644 index 489e60460ca8..000000000000 --- a/test/registered/kernels/ops/gemm/test_router_gemv.py +++ /dev/null @@ -1,91 +0,0 @@ -"""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, - split_fp32_to_bf16x3, -) -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") - -# Beam router shape: 256 experts over a 5120 hidden size. -N, K = 256, 5120 - -requires_cuda = pytest.mark.skipif( - not torch.cuda.is_available() or torch.version.hip is not None, - reason="CUDA required", -) - - -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() - - -@requires_cuda -def test_bf16x3_split_reconstructs_fp32_exactly(): - _, w = _inputs(1) - w = torch.cat([w, w * 1e30, w * 1e-20, -w]) - 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) - # The CUDA kernel must round like the torch reference path. - assert torch.equal(w3.cpu(), split_fp32_to_bf16x3(w.cpu())) - - -@requires_cuda -@pytest.mark.parametrize("m", [1, 2, 4, 8]) -def test_fp32_gate_gemv_matches_fp64(m): - x, w = _inputs(m) - assert router_gemv_supported(x, w) - out = router_gemv(x, w) - assert out.dtype == torch.float32 - # Measured about 5e-7 on H200 and B300; a BF16 or TF32 weight read is ~1e-3. - assert _max_error(out, x, w) < 2e-6 - - -@requires_cuda -def test_fp32_gate_gemv_stops_at_its_row_path(): - x, w = _inputs(9) - assert not router_gemv_supported(x, w) - assert not router_gemv_supported(x[:1], w.bfloat16()) - - -@requires_cuda -@pytest.mark.parametrize("m", [9, 64, 1024, 16384]) -def test_bf16x3_router_gemm_matches_fp64(m): - x, w = _inputs(m) - out = router_gemm_bf16x3_cublas(x, w) - assert out.dtype == torch.float32 - # Measured up to about 5e-5 (tensor-core accumulation); an fp32 SGEMM is about 2e-5. - assert _max_error(out, x, w) < 1e-4 - assert _max_error(F.linear(x.float(), w.bfloat16().float()), x, w) > 1e-4 - - -@requires_cuda -def test_bf16x3_router_gemm_reads_in_place_weight_updates(): - x, w = _inputs(64) - router_gemm_bf16x3_cublas(x, w) - # Online weight updates copy into the gate in place, without a reload hook. - w.copy_(-2 * w) - assert _max_error(router_gemm_bf16x3_cublas(x, w), x, w) < 1e-4 - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py b/test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py deleted file mode 100644 index 7d7241659e00..000000000000 --- a/test/registered/kernels/ops/gemm/test_tiny_gemm_fp32.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Numerical and graph-replay contract for full-precision router weights.""" - -import unittest - -import torch - -from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.test_utils import CustomTestCase - -register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") - - -def _tiny_fp32(x, w): - return tiny_gemm_bf16(x, w, out_dtype=torch.float32) - - -@unittest.skipUnless( - torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9, - "SM90+ required", -) -class TestTinyGemmFp32(CustomTestCase): - def setUp(self): - torch.manual_seed(42) - - def test_fp64_reference(self): - w = torch.randn(256, 5120, device="cuda") / 5120**0.5 - for m in range(1, 17): - with self.subTest(m=m): - x = torch.randn(m, 5120, device="cuda", dtype=torch.bfloat16) - result = _tiny_fp32(x, w) - self.assertEqual(result.dtype, torch.float32) - reference = (x.double() @ w.double().T).float() - torch.testing.assert_close(result, reference, atol=2e-6, rtol=2e-5) - - def test_other_shapes_and_strided_activations(self): - for n, k in ((128, 1024), (384, 7168)): - with self.subTest(n=n, k=k): - x = torch.randn(3, k + 32, device="cuda", dtype=torch.bfloat16)[:, :k] - w = torch.randn(n, k, device="cuda") / k**0.5 - result = _tiny_fp32(x, w) - torch.testing.assert_close( - result, (x.double() @ w.double().T).float(), atol=2e-6, rtol=2e-5 - ) - - def test_weights_remain_fp32(self): - # All these weights round to 1 in BF16. Preserving their low FP32 bits - # is necessary to distinguish the two highest-scoring experts. - x = torch.zeros(1, 5120, device="cuda", dtype=torch.bfloat16) - x[:, 0] = 1 - w = torch.ones(256, 5120, device="cuda") - w[:, 0] += torch.arange(256, device="cuda") * 2**-20 - result = _tiny_fp32(x, w) - torch.testing.assert_close(result[0], w[:, 0], atol=0, rtol=0) - self.assertEqual(result.argmax().item(), 255) - self.assertEqual(torch.unique(w[:, 0].bfloat16()).numel(), 1) - - def test_graph_observes_input_and_weight_updates(self): - x = torch.randn(8, 5120, device="cuda", dtype=torch.bfloat16) - w = torch.randn(256, 5120, device="cuda") / 5120**0.5 - _tiny_fp32(x, w) - next_x = torch.empty_like(x) - next_w = torch.empty_like(w) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - x.copy_(next_x) - w.copy_(next_w) - result = _tiny_fp32(x, w) - for _ in range(3): - next_x.normal_() - next_w.normal_(std=5120**-0.5) - graph.replay() - torch.testing.assert_close( - result, (x.double() @ w.double().T).float(), atol=2e-6, rtol=2e-5 - ) - - def test_rejects_invalid_layout_and_dtype(self): - x = torch.zeros(1, 5120, device="cuda", dtype=torch.bfloat16) - w = torch.zeros(256, 5120, device="cuda") - for bad_x, bad_w in ( - (x.float(), w), - (x, w.half()), - (torch.zeros(1, 10240, device="cuda", dtype=torch.bfloat16)[:, ::2], w), - (x, w.T.contiguous().T), - ( - torch.zeros(5121, device="cuda", dtype=torch.bfloat16)[1:].view( - 1, 5120 - ), - w, - ), - ): - with self.assertRaises(Exception): - _tiny_fp32(bad_x, bad_w) - - -if __name__ == "__main__": - unittest.main() From a25899cbb3fedb593a9a94ea24f0c3459b49646b Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:15:32 -0700 Subject: [PATCH 09/40] Drop the KernelSpec fields main removed --- python/sglang/kernels/ops/elementwise/beam_fused.py | 3 +-- python/sglang/kernels/ops/layernorm/beam_fused.py | 3 +-- python/sglang/kernels/ops/moe/__init__.py | 5 ----- 3 files changed, 2 insertions(+), 9 deletions(-) diff --git a/python/sglang/kernels/ops/elementwise/beam_fused.py b/python/sglang/kernels/ops/elementwise/beam_fused.py index ace8d742fb31..aa18567186cc 100644 --- a/python/sglang/kernels/ops/elementwise/beam_fused.py +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -21,7 +21,7 @@ from triton.language.extra import libdevice from sglang.kernels.fused_op import BaseFusedOp -from sglang.kernels.spec import CapabilityRequirement, FormatSignature, KernelBackend +from sglang.kernels.spec import CapabilityRequirement, KernelBackend from sglang.srt.utils.custom_op import register_custom_op @@ -103,7 +103,6 @@ class BeamAttentionGateOp(BaseFusedOp): op = "elementwise.beam_attention_gate" capabilities = {KernelBackend.TRITON: {CapabilityRequirement.CUDA}} - format_signature = FormatSignature(supported_dtypes=("float16", "bfloat16")) def _torch_compile_forward(self, num_tokens: int): # Native compilation can remove the required intermediate rounding. diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py index 42f6b0ed7678..2a6c01a23479 100644 --- a/python/sglang/kernels/ops/layernorm/beam_fused.py +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -27,7 +27,7 @@ import triton.language as tl from sglang.kernels.fused_op import BaseFusedOp -from sglang.kernels.spec import CapabilityRequirement, FormatSignature, KernelBackend +from sglang.kernels.spec import CapabilityRequirement, KernelBackend from sglang.srt.utils.custom_op import register_custom_op @@ -195,7 +195,6 @@ class BeamSandwichResidualRMSNormOp(BaseFusedOp): op = "layernorm.beam_sandwich_residual_rmsnorm" capabilities = {KernelBackend.TRITON: {CapabilityRequirement.CUDA}} - format_signature = FormatSignature(supported_dtypes=("float16", "bfloat16")) def _torch_compile_forward(self, num_tokens: int): # Native compilation can remove the required intermediate rounding. diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index 593bd3b5a1f7..15401604add2 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -35,11 +35,6 @@ capabilities=frozenset( {CapabilityRequirement.cuda(min_sm=(10, 3), max_sm=(10, 3))} ), - format_signature=FormatSignature( - supported_dtypes=("bfloat16",), - description="single-token interleaved Beam MoE with FP32 routing weights", - ), - description="SM103 Beam TP4/EP1 weighted expert outputs; caller owns reduction.", ) ) From 7899e54d578316ab630228a696e10dd6c2decd5b Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:15:32 -0700 Subject: [PATCH 10/40] Remove the fused-op instances only the dropped tests used --- python/sglang/kernels/ops/elementwise/beam_fused.py | 3 --- python/sglang/kernels/ops/layernorm/beam_fused.py | 3 --- 2 files changed, 6 deletions(-) diff --git a/python/sglang/kernels/ops/elementwise/beam_fused.py b/python/sglang/kernels/ops/elementwise/beam_fused.py index aa18567186cc..3cecf6a0a02b 100644 --- a/python/sglang/kernels/ops/elementwise/beam_fused.py +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -145,6 +145,3 @@ def forward_triton( ): return self.forward_native(attention_output, gate, gate_bias, head_dim) return fused_attention_gate(attention_output, gate, gate_bias, head_dim) - - -_BEAM_ATTENTION_GATE = BeamAttentionGateOp() diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py index 2a6c01a23479..8afa652b7f01 100644 --- a/python/sglang/kernels/ops/layernorm/beam_fused.py +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -253,6 +253,3 @@ def forward_triton( *scales, enable_pdl=is_arch_support_pdl(), ) - - -_BEAM_SANDWICH_RESIDUAL_RMSNORM = BeamSandwichResidualRMSNormOp() From 2fbb520aca00589239c74eabef2dad553e5c70ca Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:48:49 -0700 Subject: [PATCH 11/40] Register the BF16x3 router helpers in the kernel inventory --- python/sglang/kernels/ops/gemm/__init__.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index 804f0d264190..60198b10ffc8 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -511,3 +511,18 @@ def kimi_k3_tiny_gemm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: capabilities=CUDA_ONLY, ) ) +register_kernel( + KernelSpec( + op="gemm.split_fp32_to_bf16x3", + backend=KernelBackend.TRITON, + target="sglang.kernels.ops.gemm.router_gemv:split_fp32_to_bf16x3", + ) +) +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, + ) +) From 7565cdb6017cb3936b6b990530511a0d25e19247 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:44:54 -0700 Subject: [PATCH 12/40] Allocate the BF16x3 router weight split per call --- python/sglang/kernels/ops/gemm/router_gemv.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index 7ec5f9e16cd5..bb91e96843ed 100644 --- a/python/sglang/kernels/ops/gemm/router_gemv.py +++ b/python/sglang/kernels/ops/gemm/router_gemv.py @@ -225,9 +225,6 @@ def split_fp32_to_bf16x3( return torch.cat([hi, mid, lo], out=out) -_split_scratch: Dict[Tuple[torch.device, int, int], torch.Tensor] = {} - - 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. @@ -236,11 +233,6 @@ def router_gemm_bf16x3_cublas(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: tensor-core accumulation is not IEEE fp32: the error is a few times that of an fp32 SGEMM. """ - n, k = w.shape - key = (w.device, n, k) - if key not in _split_scratch: - _split_scratch[key] = torch.empty( - 3 * n, k, dtype=torch.bfloat16, device=w.device - ) - w3 = split_fp32_to_bf16x3(w, out=_split_scratch[key]) + 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) From faddbb3f1868bd3d79bb0008d92ab3a9c09cfe88 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:45:29 -0700 Subject: [PATCH 13/40] Validate the BF16x3 split's operands and drop its unused torch path --- python/sglang/kernels/ops/gemm/__init__.py | 1 + python/sglang/kernels/ops/gemm/router_gemv.py | 30 ++++++++++--------- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index 53de4b375c1b..983bc0dcde3c 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -530,6 +530,7 @@ def kimi_k3_tiny_gemm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: op="gemm.split_fp32_to_bf16x3", backend=KernelBackend.TRITON, target="sglang.kernels.ops.gemm.router_gemv:split_fp32_to_bf16x3", + capabilities=CUDA_ONLY, ) ) register_kernel( diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index bb91e96843ed..82d854f6ae5c 100644 --- a/python/sglang/kernels/ops/gemm/router_gemv.py +++ b/python/sglang/kernels/ops/gemm/router_gemv.py @@ -208,21 +208,23 @@ def split_fp32_to_bf16x3( Exact for |w| >= 2**-110: each part takes the next 8 significant bits, and the last residual then stays a normal number. """ + n, k = w.shape + if not (w.is_cuda and w.dtype == torch.float32 and w.is_contiguous()): + raise ValueError("w must be a contiguous CUDA fp32 [N, K] tensor") if out is None: - out = torch.empty( - 3 * w.shape[0], w.shape[1], dtype=torch.bfloat16, device=w.device - ) - if w.is_cuda and w.is_contiguous(): - block = 1024 - _split_fp32_to_bf16x3_kernel[(triton.cdiv(w.numel(), block),)]( - w, out, w.numel(), BLOCK=block - ) - return out - hi = w.to(torch.bfloat16) - residual = w - hi.float() - mid = residual.to(torch.bfloat16) - lo = (residual - mid.float()).to(torch.bfloat16) - return torch.cat([hi, mid, lo], out=out) + 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 def router_gemm_bf16x3_cublas(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: From 7a33c19b02ebc751e39c1d0fbf00341282e75373 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:46:48 -0700 Subject: [PATCH 14/40] Move the BF16x3 weight split to the quantization group --- python/sglang/kernels/ops/gemm/__init__.py | 8 --- python/sglang/kernels/ops/gemm/router_gemv.py | 44 +--------------- .../kernels/ops/quantization/__init__.py | 8 +++ .../kernels/ops/quantization/bf16x3_split.py | 51 +++++++++++++++++++ 4 files changed, 61 insertions(+), 50 deletions(-) create mode 100644 python/sglang/kernels/ops/quantization/bf16x3_split.py diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index 983bc0dcde3c..8b1b62a210a6 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -525,14 +525,6 @@ def kimi_k3_tiny_gemm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: capabilities=CUDA_ONLY, ) ) -register_kernel( - KernelSpec( - op="gemm.split_fp32_to_bf16x3", - backend=KernelBackend.TRITON, - target="sglang.kernels.ops.gemm.router_gemv:split_fp32_to_bf16x3", - capabilities=CUDA_ONLY, - ) -) register_kernel( KernelSpec( op="gemm.router_gemm_bf16x3_cublas", diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index 82d854f6ae5c..4b02a43e3e6e 100644 --- a/python/sglang/kernels/ops/gemm/router_gemv.py +++ b/python/sglang/kernels/ops/gemm/router_gemv.py @@ -11,12 +11,13 @@ from __future__ import annotations -from typing import Dict, Optional, Tuple +from typing import Dict, Tuple import torch 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() @@ -186,47 +187,6 @@ def router_gemv(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: return out -@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 as (hi, mid, lo) blocks with w == hi + mid + lo. - - Exact for |w| >= 2**-110: each part takes the next 8 significant bits, and the - last residual then stays a normal number. - """ - n, k = w.shape - if not (w.is_cuda and w.dtype == torch.float32 and w.is_contiguous()): - raise ValueError("w must be a contiguous CUDA fp32 [N, K] tensor") - 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 - - 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. 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..ee2d59b9bec6 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/bf16x3_split.py @@ -0,0 +1,51 @@ +# 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 as (hi, mid, lo) blocks with w == hi + mid + lo. + + Exact for |w| >= 2**-110: each part takes the next 8 significant bits, and the + last residual then stays a normal number. + """ + n, k = w.shape + if not (w.is_cuda and w.dtype == torch.float32 and w.is_contiguous()): + raise ValueError("w must be a contiguous CUDA fp32 [N, K] tensor") + 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 From 6a4c90ac8cde382e2f58d1b9ff4f1c1bc512af14 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:47:29 -0700 Subject: [PATCH 15/40] Correct the fast-math SiLU, BF16x3 and Beam all-reduce comments --- .../sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh | 3 ++- python/sglang/kernels/ops/activation/_triton_silu.py | 6 +++--- python/sglang/kernels/ops/communication/beam_all_reduce.py | 1 + python/sglang/kernels/ops/gemm/router_gemv.py | 5 ++--- python/sglang/kernels/ops/quantization/bf16x3_split.py | 6 +++--- 5 files changed, 11 insertions(+), 10 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh index 586b0a2f920f..b11134e186ff 100644 --- a/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh +++ b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh @@ -106,7 +106,7 @@ __global__ void __launch_bounds__(Hidden / 8 / Cluster) __cluster_dims__(1, Clus 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 their rank-local payloads. + // 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); @@ -174,6 +174,7 @@ __global__ void __launch_bounds__(Hidden / 8 / Cluster) __cluster_dims__(1, Clus 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, diff --git a/python/sglang/kernels/ops/activation/_triton_silu.py b/python/sglang/kernels/ops/activation/_triton_silu.py index 1f368e70ca6f..715b1af7e60c 100644 --- a/python/sglang/kernels/ops/activation/_triton_silu.py +++ b/python/sglang/kernels/ops/activation/_triton_silu.py @@ -7,10 +7,10 @@ @triton.jit def _silu_fast_math(gate): - # Preserve the fused MoE epilogue's __fdividef(x, 1 + __expf(-x)): accurate - # exp/div can move the result across a BF16 rounding boundary. silu_and_mul - # matches it only before SM100, where it is built with --use_fast_math. + # 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", diff --git a/python/sglang/kernels/ops/communication/beam_all_reduce.py b/python/sglang/kernels/ops/communication/beam_all_reduce.py index 0b63c52f0898..449f73b91da4 100644 --- a/python/sglang/kernels/ops/communication/beam_all_reduce.py +++ b/python/sglang/kernels/ops/communication/beam_all_reduce.py @@ -20,6 +20,7 @@ def _module(world_size: int, hidden_size: int, cluster: int = 1, topk: int = 1): 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"], ) diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index 4b02a43e3e6e..19c0a8e5aeba 100644 --- a/python/sglang/kernels/ops/gemm/router_gemv.py +++ b/python/sglang/kernels/ops/gemm/router_gemv.py @@ -191,9 +191,8 @@ 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. Every bf16 x bf16 product is exact in fp32, but - tensor-core accumulation is not IEEE fp32: the error is a few times that of - an fp32 SGEMM. + 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) diff --git a/python/sglang/kernels/ops/quantization/bf16x3_split.py b/python/sglang/kernels/ops/quantization/bf16x3_split.py index ee2d59b9bec6..1e685e9ed466 100644 --- a/python/sglang/kernels/ops/quantization/bf16x3_split.py +++ b/python/sglang/kernels/ops/quantization/bf16x3_split.py @@ -27,10 +27,10 @@ def _split_fp32_to_bf16x3_kernel(w_ptr, out_ptr, numel, BLOCK: tl.constexpr): def split_fp32_to_bf16x3( w: torch.Tensor, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - """[N, K] fp32 -> [3N, K] bf16 as (hi, mid, lo) blocks with w == hi + mid + lo. + """[N, K] fp32 -> [3N, K] bf16 (hi, mid, lo) blocks, each rounding what remains. - Exact for |w| >= 2**-110: each part takes the next 8 significant bits, and the - last residual then stays a normal number. + 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. """ n, k = w.shape if not (w.is_cuda and w.dtype == torch.float32 and w.is_contiguous()): From 86f16b6cb379ac12392e246606c748276319553b Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:48:09 -0700 Subject: [PATCH 16/40] Keep the BF16 tiny GEMM split independent of the token count --- python/sglang/kernels/ops/gemm/tiny_gemm.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/python/sglang/kernels/ops/gemm/tiny_gemm.py b/python/sglang/kernels/ops/gemm/tiny_gemm.py index acb15d943e2d..ad192e463a37 100644 --- a/python/sglang/kernels/ops/gemm/tiny_gemm.py +++ b/python/sglang/kernels/ops/gemm/tiny_gemm.py @@ -107,13 +107,7 @@ def _supports_k_variant(n: int, k: int) -> bool: @cache_once -def _default_n_gemm_n_split( - n: int, k: int, max_m: int, weight_dtype: torch.dtype = torch.bfloat16, m: int = 16 -) -> int: - if weight_dtype == torch.float32 and (n, k) == (256, 5120) and max_m <= 16: - # Rotating 49 FP32 router weights under CUDA graph replay. - small_m = 8 if get_jit_cuda_arch().major == 10 else 4 - return 1 if m <= small_m else 2 +def _default_n_gemm_n_split(n: int, k: int, max_m: int) -> int: split_cap = _n_gemm_block_threads(k) // max_m candidates = [d for d in range(1, min(n, split_cap) + 1) if n % d == 0] if not candidates: @@ -197,10 +191,13 @@ def tiny_gemm_bf16( else: assert out_dtype is None or out_dtype == out.dtype if n_split is None: - n_split = ( - _default_k_gemm_n_split(n, k) - if _prefer_k_variant(n, k) - else _default_n_gemm_n_split(n, k, max_m, w.dtype, x.shape[0]) - ) + if _prefer_k_variant(n, k): + n_split = _default_k_gemm_n_split(n, k) + elif w.dtype == torch.float32 and (n, k) == (256, 5120) and max_m <= 16: + # Rotating 49 FP32 router weights under CUDA graph replay. + small_m = 8 if get_jit_cuda_arch().major == 10 else 4 + n_split = 1 if x.shape[0] <= small_m else 2 + else: + n_split = _default_n_gemm_n_split(n, k, max_m) _tiny_gemm_custom_op(x, w, out, max_m, n_split) return out From 77e25c47cda2a006ecdc185d6d32f16d6151b76d Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:48:47 -0700 Subject: [PATCH 17/40] Drop the Beam fused ops' torch_compile overrides and native fallbacks --- .../sglang/kernels/ops/elementwise/beam_fused.py | 14 -------------- python/sglang/kernels/ops/layernorm/beam_fused.py | 13 ------------- 2 files changed, 27 deletions(-) diff --git a/python/sglang/kernels/ops/elementwise/beam_fused.py b/python/sglang/kernels/ops/elementwise/beam_fused.py index 3cecf6a0a02b..c76e3b541c4d 100644 --- a/python/sglang/kernels/ops/elementwise/beam_fused.py +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -108,15 +108,6 @@ 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: - # Register and warm the opaque op before tracing its dispatch wrapper. - self.forward_triton(*args, **kwargs) - self._compiled_native = torch.compile( - self.forward_triton, options={"emulate_precision_casts": True} - ) - return self._compiled_native(*args, **kwargs) - def forward_native( self, attention_output: torch.Tensor, @@ -139,9 +130,4 @@ def forward_triton( gate_bias: torch.Tensor, head_dim: int, ) -> torch.Tensor: - if not attention_output.is_cuda or attention_output.dtype not in ( - torch.float16, - torch.bfloat16, - ): - return self.forward_native(attention_output, gate, gate_bias, head_dim) return fused_attention_gate(attention_output, gate, gate_bias, head_dim) diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py index 8afa652b7f01..12b2856b10b4 100644 --- a/python/sglang/kernels/ops/layernorm/beam_fused.py +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -200,15 +200,6 @@ 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: - # Register and warm the opaque op before tracing its dispatch wrapper. - self.forward_triton(*args, **kwargs) - self._compiled_native = torch.compile( - self.forward_triton, options={"emulate_precision_casts": True} - ) - return self._compiled_native(*args, **kwargs) - def forward_native( self, branch: torch.Tensor, @@ -238,10 +229,6 @@ def forward_triton( eps: float, scales: tuple[float, float], ) -> tuple[Optional[torch.Tensor], torch.Tensor]: - if not branch.is_cuda or branch.dtype not in (torch.float16, torch.bfloat16): - return self.forward_native( - branch, residual, sandwich_weight, out_weight, eps, scales - ) from sglang.kernels.jit.utils import is_arch_support_pdl return fused_sandwich_residual_rmsnorm( From d347ff8b86634ba6c4a004ebde72c7843b3d6cd5 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:49:29 -0700 Subject: [PATCH 18/40] Require the post-norm weight in the Beam sandwich norm --- .../kernels/ops/layernorm/beam_fused.py | 42 ++++++++----------- 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py index 12b2856b10b4..761e173e2325 100644 --- a/python/sglang/kernels/ops/layernorm/beam_fused.py +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -46,7 +46,6 @@ def _fused_sandwich_residual_rmsnorm_kernel( residual_scale, hidden_dim: tl.constexpr, BLOCK_SIZE: tl.constexpr, - HAS_SANDWICH: tl.constexpr, HAS_OUT_NORM: tl.constexpr, USE_PDL: tl.constexpr, ): @@ -56,8 +55,7 @@ def _fused_sandwich_residual_rmsnorm_kernel( act_dtype = x_ptr.dtype.element_ty # The preceding kernel does not write the weights; load them before the PDL wait. - if HAS_SANDWICH: - w = tl.load(sandwich_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + 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: @@ -67,10 +65,9 @@ def _fused_sandwich_residual_rmsnorm_kernel( x = tl.load(x_ptr + row * x_stride_row + offsets, mask=mask, other=0.0).to( tl.float32 ) - if HAS_SANDWICH: - 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) + 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 @@ -92,7 +89,7 @@ def _fused_sandwich_residual_rmsnorm_kernel( def fused_sandwich_residual_rmsnorm( branch: torch.Tensor, residual: torch.Tensor, - sandwich_weight: Optional[torch.Tensor], + sandwich_weight: torch.Tensor, out_weight: Optional[torch.Tensor], eps: float, x_scale: float, @@ -104,10 +101,10 @@ def fused_sandwich_residual_rmsnorm( `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 (None skips it); `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. + 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( @@ -140,7 +137,7 @@ def fused_sandwich_residual_rmsnorm( def _fused_sandwich_residual_rmsnorm( branch: torch.Tensor, residual: torch.Tensor, - sandwich_weight: Optional[torch.Tensor], + sandwich_weight: torch.Tensor, out_weight: Optional[torch.Tensor], eps: float, x_scale: float, @@ -153,9 +150,6 @@ def _fused_sandwich_residual_rmsnorm( rows = branch.shape[0] if rows == 0: return - dummy = sandwich_weight if sandwich_weight is not None else out_weight - if dummy is None: - dummy = branch # never dereferenced: both norm flags are off _fused_sandwich_residual_rmsnorm_kernel[(rows,)]( out if out is not None else combined, combined, @@ -163,14 +157,13 @@ def _fused_sandwich_residual_rmsnorm( residual, branch.stride(0), residual.stride(0), - sandwich_weight if sandwich_weight is not None else dummy, - out_weight if out_weight is not None else dummy, + 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_SANDWICH=sandwich_weight is not None, HAS_OUT_NORM=out_weight is not None, USE_PDL=enable_pdl, num_warps=8, @@ -189,8 +182,8 @@ def _rmsnorm_native(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch. class BeamSandwichResidualRMSNormOp(BaseFusedOp): """Post-normalize a branch, combine residuals, then normalize the next input. - Returns ``(normed, combined)``. Either norm weight may be None; without - ``out_weight``, ``normed`` is None (a non-final pipeline stage). + Returns ``(normed, combined)``; ``normed`` is None without ``out_weight`` + (a non-final pipeline stage). """ op = "layernorm.beam_sandwich_residual_rmsnorm" @@ -204,13 +197,12 @@ def forward_native( self, branch: torch.Tensor, residual: torch.Tensor, - sandwich_weight: Optional[torch.Tensor], + sandwich_weight: torch.Tensor, out_weight: Optional[torch.Tensor], eps: float, scales: tuple[float, float], ) -> tuple[Optional[torch.Tensor], torch.Tensor]: - if sandwich_weight is not None: - branch = _rmsnorm_native(branch, sandwich_weight, eps) + branch = _rmsnorm_native(branch, sandwich_weight, eps) residual = residual.to(branch.dtype) combined = branch * scales[0] + residual * scales[1] out = ( @@ -224,7 +216,7 @@ def forward_triton( self, branch: torch.Tensor, residual: torch.Tensor, - sandwich_weight: Optional[torch.Tensor], + sandwich_weight: torch.Tensor, out_weight: Optional[torch.Tensor], eps: float, scales: tuple[float, float], From a8e836fb488ef34a5d4c8fd7b2d71d5eabce75be Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:50:12 -0700 Subject: [PATCH 19/40] Use SPDX license headers in the Beam kernel files --- .../jit/csrc/distributed/beam_all_reduce.cuh | 2 +- .../sglang/kernels/ops/activation/_triton_silu.py | 2 +- .../kernels/ops/communication/beam_all_reduce.py | 1 + .../sglang/kernels/ops/elementwise/beam_fused.py | 14 +------------- python/sglang/kernels/ops/layernorm/beam_fused.py | 14 +------------- python/sglang/kernels/ops/moe/beam_decode_gemv.py | 2 +- 6 files changed, 6 insertions(+), 29 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh index b11134e186ff..2c800793c833 100644 --- a/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh +++ b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh @@ -1,4 +1,4 @@ -// Copyright 2026 SGLang Team. Licensed under the Apache License, Version 2.0. +// SPDX-License-Identifier: Apache-2.0 // Beam's BF16 sandwich norm over the existing CustomAllReduceV2 push protocol. #pragma once #include diff --git a/python/sglang/kernels/ops/activation/_triton_silu.py b/python/sglang/kernels/ops/activation/_triton_silu.py index 715b1af7e60c..d881f147fa4c 100644 --- a/python/sglang/kernels/ops/activation/_triton_silu.py +++ b/python/sglang/kernels/ops/activation/_triton_silu.py @@ -1,4 +1,4 @@ -# Copyright 2026 SGLang Team. Licensed under the Apache License, Version 2.0. +# SPDX-License-Identifier: Apache-2.0 """Private Triton activation helpers with CUDA fast-math rounding.""" import triton diff --git a/python/sglang/kernels/ops/communication/beam_all_reduce.py b/python/sglang/kernels/ops/communication/beam_all_reduce.py index 449f73b91da4..64989dbf7827 100644 --- a/python/sglang/kernels/ops/communication/beam_all_reduce.py +++ b/python/sglang/kernels/ops/communication/beam_all_reduce.py @@ -1,3 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 """Beam BF16 all-reduce with FP32 sandwich normalization and rounded residuals.""" import torch diff --git a/python/sglang/kernels/ops/elementwise/beam_fused.py b/python/sglang/kernels/ops/elementwise/beam_fused.py index c76e3b541c4d..9f2fb0f0136e 100644 --- a/python/sglang/kernels/ops/elementwise/beam_fused.py +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -1,16 +1,4 @@ -# Copyright 2026 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== +# SPDX-License-Identifier: Apache-2.0 """Beam attention gate with activation-dtype bias and sigmoid rounding.""" from __future__ import annotations diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py index 761e173e2325..ec606b743f29 100644 --- a/python/sglang/kernels/ops/layernorm/beam_fused.py +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -1,16 +1,4 @@ -# Copyright 2026 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== +# 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 diff --git a/python/sglang/kernels/ops/moe/beam_decode_gemv.py b/python/sglang/kernels/ops/moe/beam_decode_gemv.py index a62a062639b8..a7f73912541d 100644 --- a/python/sglang/kernels/ops/moe/beam_decode_gemv.py +++ b/python/sglang/kernels/ops/moe/beam_decode_gemv.py @@ -1,4 +1,4 @@ -# Copyright 2026 SGLang Team. Licensed under the Apache License, Version 2.0. +# 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 From 3b5abfb551b43eb2320b4e8d45574d035f57a149 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:50:44 -0700 Subject: [PATCH 20/40] Drop unused helper defaults and type the Beam all-reduce entry point --- .../ops/communication/beam_all_reduce.py | 21 ++++++++++++++++--- python/sglang/kernels/ops/gemm/router_gemv.py | 2 +- python/sglang/kernels/ops/gemm/tiny_gemm.py | 2 +- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/python/sglang/kernels/ops/communication/beam_all_reduce.py b/python/sglang/kernels/ops/communication/beam_all_reduce.py index 64989dbf7827..07d24618abe6 100644 --- a/python/sglang/kernels/ops/communication/beam_all_reduce.py +++ b/python/sglang/kernels/ops/communication/beam_all_reduce.py @@ -1,13 +1,20 @@ # 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 = 1, topk: int = 1): +def _module(world_size: int, hidden_size: int, cluster: int, topk: int): return load_jit( "beam_all_reduce_norm", str(world_size), @@ -35,8 +42,16 @@ def _cluster_size(world_size: int, rows: int, topk: int) -> int: def beam_all_reduce_norm( - comm, branch, residual, post_weight, next_weight, eps, scales, *, routed_scale=1.0 -): + 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 diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index 19c0a8e5aeba..67343aba8887 100644 --- a/python/sglang/kernels/ops/gemm/router_gemv.py +++ b/python/sglang/kernels/ops/gemm/router_gemv.py @@ -41,7 +41,7 @@ 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 = torch.bfloat16): +def _config(m: int, w_dtype: torch.dtype): configs = _configs(w_dtype) for max_m, cfg in configs: if m <= max_m: diff --git a/python/sglang/kernels/ops/gemm/tiny_gemm.py b/python/sglang/kernels/ops/gemm/tiny_gemm.py index ad192e463a37..1cc8615bca37 100644 --- a/python/sglang/kernels/ops/gemm/tiny_gemm.py +++ b/python/sglang/kernels/ops/gemm/tiny_gemm.py @@ -48,7 +48,7 @@ def _jit_tiny_gemm_module( max_m: int, n_split: int, out_dtype: torch.dtype, - weight_dtype: torch.dtype = torch.bfloat16, + weight_dtype: torch.dtype, ) -> Module: use_k_variant = _prefer_k_variant(n, k) if weight_dtype not in (torch.bfloat16, torch.float32): From 14e419bb2d58a2c835a1dd2d2b531bbee1045cf0 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:28:40 -0700 Subject: [PATCH 21/40] Keep eager rounding in the Beam fused ops' torch_compile backend --- python/sglang/kernels/ops/elementwise/beam_fused.py | 8 ++++++++ python/sglang/kernels/ops/layernorm/beam_fused.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/python/sglang/kernels/ops/elementwise/beam_fused.py b/python/sglang/kernels/ops/elementwise/beam_fused.py index 9f2fb0f0136e..6e5e8e93d888 100644 --- a/python/sglang/kernels/ops/elementwise/beam_fused.py +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -96,6 +96,14 @@ 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, diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py index ec606b743f29..235084ac00c8 100644 --- a/python/sglang/kernels/ops/layernorm/beam_fused.py +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -181,6 +181,14 @@ 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, From 243466189cabaf503e485cbecd8e82fadef897bb Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:28:49 -0700 Subject: [PATCH 22/40] Reject unsupported dtypes in the Beam Triton fused ops --- python/sglang/kernels/ops/elementwise/beam_fused.py | 2 ++ python/sglang/kernels/ops/layernorm/beam_fused.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/python/sglang/kernels/ops/elementwise/beam_fused.py b/python/sglang/kernels/ops/elementwise/beam_fused.py index 6e5e8e93d888..23da828d5326 100644 --- a/python/sglang/kernels/ops/elementwise/beam_fused.py +++ b/python/sglang/kernels/ops/elementwise/beam_fused.py @@ -62,6 +62,8 @@ def fused_attention_gate( """`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(): diff --git a/python/sglang/kernels/ops/layernorm/beam_fused.py b/python/sglang/kernels/ops/layernorm/beam_fused.py index 235084ac00c8..b4bfdf5926fa 100644 --- a/python/sglang/kernels/ops/layernorm/beam_fused.py +++ b/python/sglang/kernels/ops/layernorm/beam_fused.py @@ -100,6 +100,8 @@ def fused_sandwich_residual_rmsnorm( ) 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()): From 4b1354f3fcc11e15e16c12c6fd96fe246fbcc55d Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:28:56 -0700 Subject: [PATCH 23/40] Check the BF16x3 split weight's rank before unpacking it --- python/sglang/kernels/ops/quantization/bf16x3_split.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/quantization/bf16x3_split.py b/python/sglang/kernels/ops/quantization/bf16x3_split.py index 1e685e9ed466..d931ff8b70cf 100644 --- a/python/sglang/kernels/ops/quantization/bf16x3_split.py +++ b/python/sglang/kernels/ops/quantization/bf16x3_split.py @@ -32,9 +32,11 @@ def split_fp32_to_bf16x3( 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. """ - n, k = w.shape - if not (w.is_cuda and w.dtype == torch.float32 and w.is_contiguous()): + 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 ( From 5ba437b76a33360fcc8b0653796457ac14fa3647 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:08:35 -0700 Subject: [PATCH 24/40] Prefetch FP32 tiny GEMM weights before the PDL wait --- python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh index 7e0b775eea17..c543905f3fc7 100644 --- a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh @@ -13,7 +13,6 @@ #include #include #include -#include #include namespace sglang { @@ -85,8 +84,7 @@ TINY_GEMM_KERNEL void tiny_n_gemm_kernel(const TinyGEMMParams params) { const uint32_t tx = threadIdx.x; const Weight* w_tile = static_cast(params.w) + bx * (N_SPLIT * K); - // FP32 weights may have been updated by the preceding graph node. - if constexpr (std::is_same_v) PDLWaitPrimary(); + // prefetch weight before PDL using weight_vec_t = AlignedVector, kVecSize / 2>; weight_vec_t wv[N_SPLIT][kUnroll]; #pragma unroll @@ -97,7 +95,7 @@ TINY_GEMM_KERNEL void tiny_n_gemm_kernel(const TinyGEMMParams params) { } } - if constexpr (!std::is_same_v) PDLWaitPrimary(); + PDLWaitPrimary(); vec_t xv[M][kUnroll]; #pragma unroll From d104fa31c34d176396abf8a7b51f6641c59656fb Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:08:35 -0700 Subject: [PATCH 25/40] Drop the FP32 router tiny GEMM split override --- python/sglang/kernels/ops/gemm/tiny_gemm.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/python/sglang/kernels/ops/gemm/tiny_gemm.py b/python/sglang/kernels/ops/gemm/tiny_gemm.py index 1cc8615bca37..f028b9135b2a 100644 --- a/python/sglang/kernels/ops/gemm/tiny_gemm.py +++ b/python/sglang/kernels/ops/gemm/tiny_gemm.py @@ -14,7 +14,6 @@ from sglang.kernels.jit.utils import ( cache_once, - get_jit_cuda_arch, is_arch_support_pdl, load_jit, make_cpp_args, @@ -191,13 +190,10 @@ def tiny_gemm_bf16( else: assert out_dtype is None or out_dtype == out.dtype if n_split is None: - if _prefer_k_variant(n, k): - n_split = _default_k_gemm_n_split(n, k) - elif w.dtype == torch.float32 and (n, k) == (256, 5120) and max_m <= 16: - # Rotating 49 FP32 router weights under CUDA graph replay. - small_m = 8 if get_jit_cuda_arch().major == 10 else 4 - n_split = 1 if x.shape[0] <= small_m else 2 - else: - n_split = _default_n_gemm_n_split(n, k, max_m) + n_split = ( + _default_k_gemm_n_split(n, k) + if _prefer_k_variant(n, k) + else _default_n_gemm_n_split(n, k, max_m) + ) _tiny_gemm_custom_op(x, w, out, max_m, n_split) return out From 01a68ebe4cc9784490050389768ea9f44fb32992 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:08:57 -0700 Subject: [PATCH 26/40] Drop the unverified FP32 router GEMV crossover claim --- python/sglang/kernels/ops/gemm/router_gemv.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/gemm/router_gemv.py b/python/sglang/kernels/ops/gemm/router_gemv.py index 67343aba8887..b69eab7a1ee0 100644 --- a/python/sglang/kernels/ops/gemm/router_gemv.py +++ b/python/sglang/kernels/ops/gemm/router_gemv.py @@ -32,8 +32,7 @@ (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; above -# 8 rows router_gemm_bf16x3_cublas is faster. Swept for N=256, K=5120 on H200 and B300. +# FP32 weights on CUDA (FP32 router gates) use the FP32 row path only, up to 8 rows. _CUDA_FP32_CONFIGS = ((8, (1, 16, 128, 8, 8)),) From 0f2792970d877cbb011bce2b5df8c028c932e56f Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:09:52 -0700 Subject: [PATCH 27/40] Launch the Beam all-reduce boundary with one CTA per row --- .../jit/csrc/distributed/beam_all_reduce.cuh | 57 +++++++------------ .../ops/communication/beam_all_reduce.py | 20 +------ 2 files changed, 23 insertions(+), 54 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh index 2c800793c833..10149a9df56a 100644 --- a/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh +++ b/python/sglang/kernels/jit/csrc/distributed/beam_all_reduce.cuh @@ -13,8 +13,6 @@ #include -#include - namespace sglang { namespace beam_ar { using Vec = device::AlignedVector; @@ -30,11 +28,11 @@ struct Params { }; /// Every thread holds eight adjacent BF16 values of one row. -template +template