From ccf9fe6590ee7437005d8353c3c67d2dc4d25fcb Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 5 Sep 2026 22:27:06 +0800 Subject: [PATCH] [Kernel] Add KDA FP8 skinny GEMM for SM120 (#38082) Co-authored-by: Waterpine --- python/sglang/kernels/kda_kernels/README.md | 3 +- .../csrc/diffusion/causal_conv3d_cat_pad.cuh | 2 +- .../csrc/diffusion/ltx2_qknorm_split_rope.cuh | 2 +- .../csrc/diffusion/norm_scale_shift.cuh | 2 +- .../csrc/diffusion/residual_gate_add.cuh | 2 +- .../csrc/gemm/sm120_fp8_skinny_gemm.cuh | 459 ++++++++++++++++++ .../kda_kernels/qwen3x_nvfp4_gemm_sm120.py | 2 - .../sglang/kernels/kda_kernels/sm120_fp8.py | 151 ++++++ .../sm120_fp8_skinny_gemm_sm120.py | 90 ++++ python/sglang/kernels/ops/gemm/__init__.py | 34 ++ .../srt/layers/quantization/modelopt_quant.py | 49 +- .../gemm/bench_kda_fp8_skinny_gemm.py | 149 ++++++ .../kernels/ops/gemm/test_sm120_fp8_linear.py | 175 +++++++ .../ops/layernorm/test_kernels_namespace.py | 1 + 14 files changed, 1082 insertions(+), 39 deletions(-) create mode 100644 python/sglang/kernels/kda_kernels/csrc/gemm/sm120_fp8_skinny_gemm.cuh create mode 100644 python/sglang/kernels/kda_kernels/sm120_fp8.py create mode 100644 python/sglang/kernels/kda_kernels/sm120_fp8_skinny_gemm_sm120.py create mode 100644 test/registered/kernels/benchmark/gemm/bench_kda_fp8_skinny_gemm.py create mode 100644 test/registered/kernels/ops/gemm/test_sm120_fp8_linear.py diff --git a/python/sglang/kernels/kda_kernels/README.md b/python/sglang/kernels/kda_kernels/README.md index 6f1b2fdff..b1b892272 100644 --- a/python/sglang/kernels/kda_kernels/README.md +++ b/python/sglang/kernels/kda_kernels/README.md @@ -13,7 +13,8 @@ load Triton, CUTLASS, or compile a JIT extension. | Kernel family | Implementation | Provenance | |---|---|---| -| Qwen3.x ModelOpt NVFP4 GEMM on SM120 | `qwen3x_nvfp4_gemm_sm120.py` | [BBuf/KDA-Pilot#195](https://github.com/BBuf/KDA-Pilot/pull/195) at `516c976cee824a236679adf6eb525275a0a9a120` | +| Qwen3.x ModelOpt NVFP4 GEMM on SM120 | `qwen3x_nvfp4_gemm_sm120.py` | [sgl-project/sglang#36865](https://github.com/sgl-project/sglang/pull/36865), merge commit `c593527f33` | +| ModelOpt static per-tensor FP8 small-batch dispatch on SM12x | `sm120_fp8.py`, `sm120_fp8_skinny_gemm_sm120.py`, `csrc/gemm/sm120_fp8_skinny_gemm.cuh` | [sgl-project/sglang#38082](https://github.com/sgl-project/sglang/pull/38082) | | Qwen-Image norm / residual-norm scale-shift | `norm_scale_shift_jit.py` | [sgl-project/sglang#27392](https://github.com/sgl-project/sglang/pull/27392), merge commit `26e1d4d847` | | Cosmos3 causal Conv3D cat-pad | `causal_conv3d_cat_pad_jit.py` | [sgl-project/sglang#29281](https://github.com/sgl-project/sglang/pull/29281), merge commit `5996b54bd3` | | Diffusion residual-gate add | `residual_gate_add_jit.py` | [sgl-project/sglang#29361](https://github.com/sgl-project/sglang/pull/29361), merge commit `495f13fa12` | diff --git a/python/sglang/kernels/kda_kernels/csrc/diffusion/causal_conv3d_cat_pad.cuh b/python/sglang/kernels/kda_kernels/csrc/diffusion/causal_conv3d_cat_pad.cuh index dccfb6a7c..5c07ca2ef 100644 --- a/python/sglang/kernels/kda_kernels/csrc/diffusion/causal_conv3d_cat_pad.cuh +++ b/python/sglang/kernels/kda_kernels/csrc/diffusion/causal_conv3d_cat_pad.cuh @@ -1,4 +1,4 @@ -// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #29281. +// KDA provenance: SGLang PR #29281. // Native CUDA fast path for Cosmos3 VAE causal-Conv3D cat/pad copy. // // The op writes the output of: diff --git a/python/sglang/kernels/kda_kernels/csrc/diffusion/ltx2_qknorm_split_rope.cuh b/python/sglang/kernels/kda_kernels/csrc/diffusion/ltx2_qknorm_split_rope.cuh index 038cdee58..63f321a2a 100644 --- a/python/sglang/kernels/kda_kernels/csrc/diffusion/ltx2_qknorm_split_rope.cuh +++ b/python/sglang/kernels/kda_kernels/csrc/diffusion/ltx2_qknorm_split_rope.cuh @@ -1,4 +1,4 @@ -// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #29708. +// KDA provenance: SGLang PR #29708. // CUDA fast path for LTX2 Q/K RMSNorm + split RoPE. // // Developed with MIT HAN Lab Kernel Design Agents: diff --git a/python/sglang/kernels/kda_kernels/csrc/diffusion/norm_scale_shift.cuh b/python/sglang/kernels/kda_kernels/csrc/diffusion/norm_scale_shift.cuh index 86a797080..550fa3036 100644 --- a/python/sglang/kernels/kda_kernels/csrc/diffusion/norm_scale_shift.cuh +++ b/python/sglang/kernels/kda_kernels/csrc/diffusion/norm_scale_shift.cuh @@ -1,4 +1,4 @@ -// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #27392. +// KDA provenance: SGLang PR #27392. // Minimal native-CUDA fast path for generic bf16 hidden=3072 norm-scale-shift. // // Supported shape family: diff --git a/python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh b/python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh index 31ad61ea8..bd74391ef 100644 --- a/python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh +++ b/python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh @@ -1,4 +1,4 @@ -// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #29361. +// KDA provenance: SGLang PR #29361. // CUDA fast path for bit-exact diffusion residual-gate updates: // out = residual + update * gate diff --git a/python/sglang/kernels/kda_kernels/csrc/gemm/sm120_fp8_skinny_gemm.cuh b/python/sglang/kernels/kda_kernels/csrc/gemm/sm120_fp8_skinny_gemm.cuh new file mode 100644 index 000000000..70e58d550 --- /dev/null +++ b/python/sglang/kernels/kda_kernels/csrc/gemm/sm120_fp8_skinny_gemm.cuh @@ -0,0 +1,459 @@ +// Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +// KDA provenance: this kernel was automatically optimized by the Humanize2 +// workflow (https://github.com/PolyArch/humanize) and Kernel Design Agents +// (https://github.com/mit-han-lab/kernel-design-agents). + +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: + +// 1. Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. + +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. + +// 3. Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software without +// specific prior written permission. + +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. + +// This file integrates the SM120 FP8 skinny GEMM with SGLang's JIT kernel and +// KDA backend interfaces. + +#pragma once + +#include +#include + +#include +#include + +#include + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/threadblock/default_thread_map_tensor_op.h" +#include "cutlass/epilogue/threadblock/fusion/visitors.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/util/packed_stride.hpp" +#include + +namespace kda_sm120_fp8_skinny { + +using namespace cute; +using namespace sglang; + +using ElementInput = cutlass::float_e4m3_t; +using ElementOutput = cutlass::bfloat16_t; +using ElementAccumulator = float; +using ArchTag = cutlass::arch::Sm120; +using OperatorClass = cutlass::arch::OpClassTensorOp; + +using ClusterShape = Shape<_1, _1, _1>; +using LayoutA = cutlass::layout::RowMajor; +using LayoutB = cutlass::layout::ColumnMajor; +using LayoutC = cutlass::layout::ColumnMajor; +using LayoutD = cutlass::layout::ColumnMajor; +constexpr int AlignmentInput = 16; +constexpr int AlignmentOutput = 8; + +using Accum = cutlass::epilogue::fusion::Sm90AccFetch; +using OutputScale = cutlass::epilogue::fusion::Sm90ScalarBroadcast; +using Multiply = cutlass::epilogue::fusion:: + Sm90Compute; +using ScaleOutput = cutlass::epilogue::fusion::Sm90EVT; + +template +CUTLASS_HOST_DEVICE constexpr auto select_b_smem_layout() { + if constexpr (SwizzleBytes == 128) { + return cute::UMMA::Layout_K_SW128_Atom{}; + } else if constexpr (SwizzleBytes == 64) { + return cute::UMMA::Layout_K_SW64_Atom{}; + } else if constexpr (SwizzleBytes == 32) { + return cute::UMMA::Layout_K_SW32_Atom{}; + } else { + static_assert(SwizzleBytes == 0, "unsupported B shared-memory swizzle"); + return cute::UMMA::Layout_K_INTER_Atom{}; + } +} + +// FlashInfer's cuBLAS backend combines the static activation and weight scales +// before applying them to the accumulator. Use the same single FP32 multiply +// before BF16 conversion to minimize speculative target-logit drift. +class OneScaleLinearCombination { + public: + using ElementAccumulator = float; + using ElementCompute = float; + using ElementOutput = kda_sm120_fp8_skinny::ElementOutput; + using ElementC = ElementOutput; + using ElementD = ElementOutput; + static constexpr int kCount = 1; + using FragmentAccumulator = cutlass::Array; + using FragmentOutput = cutlass::Array; + + struct Params { + const float* scale = nullptr; + }; + + CUTLASS_HOST_DEVICE + explicit OneScaleLinearCombination(const Params& params) : scale_(*params.scale) {} + + CUTLASS_HOST_DEVICE bool is_source_needed() const { + return false; + } + + CUTLASS_HOST_DEVICE FragmentOutput operator()(const FragmentAccumulator& accumulator) const { + cutlass::NumericConverter convert; + FragmentOutput output; + output[0] = convert(accumulator[0] * scale_); + return output; + } + + CUTLASS_HOST_DEVICE ElementD operator()(ElementAccumulator accumulator) const { + cutlass::NumericConverter convert; + return convert(accumulator * scale_); + } + + CUTLASS_HOST_DEVICE ElementD operator()(ElementAccumulator accumulator, ElementC) const { + return (*this)(accumulator); + } + + private: + ElementCompute scale_; +}; + +// D.T = W @ A.T moves the tiny token dimension to the narrow N axis. A +// 16-wide token tile minimizes padded MMA work for every qualified M <= 9. The generic +// warp-specialized schedule needs only one MMA consumer group for this skinny +// tile; unlike ping-pong, it compiles to 256 threads without register spills. +template < + int TileM, + int TileK, + int MainloopStages = 0, + bool UseCacheHints = (TileM == 32 && TileK == 512), + int BSmemSwizzleBytes = 128, + int TmaLoopUnroll = 0> +struct Fp8Gemm { + using TileShape = Shape, _16, Int>; + // Preserve the tuned mainloop's stage count by reserving the same carveout + // used by the former TMA epilogue. Only the executed epilogue changes. + using StagingEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + TileShape, + ClusterShape, + Shape, _16>, + ElementAccumulator, + float, + void, + LayoutC, + AlignmentOutput, + ElementOutput, + LayoutD, + AlignmentOutput, + cutlass::epilogue::collective::EpilogueScheduleAuto, + ScaleOutput>::CollectiveOp; + using DirectEpilogue = cutlass::epilogue::collective::DefaultEpilogue< + void, + cutlass::detail::TagToStrideC_t, + cutlass::detail::TagToStrideC_t, + OneScaleLinearCombination, + cutlass::gemm::EpilogueDefault>; + using CollectiveEpilogue = cutlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter; + using AutoStageCount = cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename StagingEpilogue::SharedStorage))>; + using StageCount = + cute::conditional_t>; + // The SM120 builder currently exposes ping-pong/cooperative schedules only. + // Use it as a type/layout factory, then rebuild the collective with the + // generic one-consumer-group dispatch policy below. + using BuiltMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementInput, + LayoutA, + AlignmentInput, + ElementInput, + LayoutB, + AlignmentInput, + ElementAccumulator, + TileShape, + ClusterShape, + StageCount, + cutlass::gemm::KernelTmaWarpSpecializedPingpong>::CollectiveOp; + // CUTLASS's generic SM120 builder selects ldmatrix.x4 for both operands. + // A 16-wide B tile contains only enough values for ldmatrix.x2, so retain + // every generated policy/layout while narrowing just that shared-to-register + // copy atom. + using MainloopDispatch = cutlass::gemm::MainloopSm120TmaWarpSpecialized< + BuiltMainloop::DispatchPolicy::Stages, + 2, + ClusterShape, + cutlass::gemm::KernelTmaWarpSpecialized>; + using SmemLayoutAtomB = decltype(select_b_smem_layout()); + using BaseCollectiveMainloop = cutlass::gemm::collective::CollectiveMma< + MainloopDispatch, + typename BuiltMainloop::TileShape, + typename BuiltMainloop::ElementA, + typename BuiltMainloop::StrideA, + typename BuiltMainloop::ElementB, + typename BuiltMainloop::StrideB, + typename BuiltMainloop::TiledMma, + typename BuiltMainloop::GmemTiledCopyA, + typename BuiltMainloop::SmemLayoutAtomA, + typename BuiltMainloop::SmemCopyAtomA, + typename BuiltMainloop::TransformA, + typename BuiltMainloop::GmemTiledCopyB, + SmemLayoutAtomB, + cute::Copy_Atom, + typename BuiltMainloop::TransformB>; + // The generic GemmUniversal kernel calls the legacy mma overload. SM120's + // collective also accepts a block coordinate for block-scaled variants; this + // scalar-scale kernel does not use it, so bridge the two interfaces. + struct CollectiveMainloop : BaseCollectiveMainloop { + template + CUTLASS_DEVICE void load( + typename BaseCollectiveMainloop::Params const& mainloop_params, + typename BaseCollectiveMainloop::MainloopPipeline pipeline, + typename BaseCollectiveMainloop::PipelineState smem_pipe_write, + cute::tuple const& load_inputs, + BlockCoord const& block_coord, + KTileIterator k_tile_iter, + int k_tile_count, + int thread_idx, + uint32_t block_rank_in_cluster, + typename BaseCollectiveMainloop::TensorStorage& shared_tensors) { + if constexpr (!UseCacheHints) { + BaseCollectiveMainloop::load( + mainloop_params, + pipeline, + smem_pipe_write, + load_inputs, + block_coord, + k_tile_iter, + k_tile_count, + thread_idx, + block_rank_in_cluster, + shared_tensors); + return; + } + + if (cute::elect_one_sync()) { + auto shared_a = cute::make_tensor( + cute::make_smem_ptr(shared_tensors.smem_A.data()), typename BaseCollectiveMainloop::SmemLayoutA{}); + auto shared_b = cute::make_tensor( + cute::make_smem_ptr(shared_tensors.smem_B.data()), typename BaseCollectiveMainloop::SmemLayoutB{}); + + // This kernel always uses a 1x1 cluster, so the TMA slice is fixed. + auto block_tma_a = mainloop_params.tma_load_a.get_slice(0); + auto block_tma_b = mainloop_params.tma_load_b.get_slice(0); + auto [m_coord, n_coord, k_coord, l_coord] = block_coord; + auto global_a = cute::get<0>(load_inputs)(cute::_, cute::_, m_coord, cute::_, l_coord); + auto global_b = cute::get<1>(load_inputs)(cute::_, cute::_, n_coord, cute::_, l_coord); + auto tiled_global_a = block_tma_a.partition_S(global_a); + auto tiled_shared_a = block_tma_a.partition_D(shared_a); + auto tiled_global_b = block_tma_b.partition_S(global_b); + auto tiled_shared_b = block_tma_b.partition_D(shared_b); + + if constexpr (TmaLoopUnroll == 0) { + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + pipeline.producer_acquire(smem_pipe_write); + using Barrier = typename BaseCollectiveMainloop::MainloopPipeline::ProducerBarrierType; + Barrier* barrier = pipeline.producer_get_barrier(smem_pipe_write); + const int stage = smem_pipe_write.index(); + // The wide projection streams 80 MiB of weights exactly once. + // Mark that traffic evict-first so it does not displace the + // freshly quantized 45 KiB activation reused by every output tile. + cute::copy( + mainloop_params.tma_load_a.with(*barrier, 0, cute::TMA::CacheHintSm90::EVICT_FIRST), + tiled_global_a(cute::_, cute::_, cute::_, *k_tile_iter), + tiled_shared_a(cute::_, cute::_, cute::_, stage)); + cute::copy( + mainloop_params.tma_load_b.with(*barrier, 0, cute::TMA::CacheHintSm90::EVICT_NORMAL), + tiled_global_b(cute::_, cute::_, cute::_, *k_tile_iter), + tiled_shared_b(cute::_, cute::_, cute::_, stage)); + ++k_tile_iter; + ++smem_pipe_write; + } + } else { + auto issue_stage = [&] { + pipeline.producer_acquire(smem_pipe_write); + using Barrier = typename BaseCollectiveMainloop::MainloopPipeline::ProducerBarrierType; + Barrier* barrier = pipeline.producer_get_barrier(smem_pipe_write); + const int stage = smem_pipe_write.index(); + cute::copy( + mainloop_params.tma_load_a.with(*barrier, 0, cute::TMA::CacheHintSm90::EVICT_FIRST), + tiled_global_a(cute::_, cute::_, cute::_, *k_tile_iter), + tiled_shared_a(cute::_, cute::_, cute::_, stage)); + cute::copy( + mainloop_params.tma_load_b.with(*barrier, 0, cute::TMA::CacheHintSm90::EVICT_NORMAL), + tiled_global_b(cute::_, cute::_, cute::_, *k_tile_iter), + tiled_shared_b(cute::_, cute::_, cute::_, stage)); + ++k_tile_iter; + ++smem_pipe_write; + }; + + if constexpr (TmaLoopUnroll == 2) { +#pragma unroll 2 + for (; k_tile_count > 0; --k_tile_count) { + issue_stage(); + } + } else if constexpr (TmaLoopUnroll == 5) { +#pragma unroll 5 + for (; k_tile_count > 0; --k_tile_count) { + issue_stage(); + } + } else { +#pragma unroll 10 + for (; k_tile_count > 0; --k_tile_count) { + issue_stage(); + } + } + } + } + } + + template + CUTLASS_DEVICE void + mma(typename BaseCollectiveMainloop::MainloopPipeline pipeline, + typename BaseCollectiveMainloop::PipelineState state, + FrgTensorC& accumulators, + int k_tile_count, + int thread_idx, + typename BaseCollectiveMainloop::TensorStorage& shared_tensors, + typename BaseCollectiveMainloop::Params const& params) { + cute::Underscore unused_block_coord; + BaseCollectiveMainloop::mma( + pipeline, state, accumulators, k_tile_count, thread_idx, shared_tensors, params, unused_block_coord); + } + }; + using Kernel = + cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue, void>; + using Device = cutlass::gemm::device::GemmUniversalAdapter; +}; + +template < + int TileM, + int TileK, + int MainloopStages = 0, + bool UseCacheHints = (TileM == 32 && TileK == 512), + int BSmemSwizzleBytes = 128, + int TmaLoopUnroll = 0> +void run( + void* output, + const void* activation, + const void* weight, + const void* output_scale, + int n, + int k, + cudaStream_t stream, + int m = 9) { + using GemmKernel = + typename Fp8Gemm::Kernel; + using Gemm = typename Fp8Gemm::Device; + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideD = typename GemmKernel::StrideD; + + StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(n, k, 1)); + StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(m, k, 1)); + StrideD stride_d = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(n, m, 1)); + + typename GemmKernel::MainloopArguments mainloop_args{ + static_cast(weight), stride_a, static_cast(activation), stride_b}; + auto* result = static_cast(output); + typename GemmKernel::EpilogueArguments epilogue_args{ + {static_cast(output_scale)}, nullptr, stride_d, result, stride_d}; + + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, {n, m, k, 1}, mainloop_args, epilogue_args}; + + Gemm gemm; + CHECK_HOST(gemm.can_implement(arguments) == cutlass::Status::kSuccess) + << "KDA SM120 FP8 GEMM cannot implement M=" << m << ", N=" << n << ", K=" << k; + CHECK_HOST(gemm.initialize(arguments, nullptr, stream) == cutlass::Status::kSuccess) + << "KDA SM120 FP8 GEMM initialization failed"; + CHECK_HOST(gemm.run(stream, nullptr, false) == cutlass::Status::kSuccess) << "KDA SM120 FP8 GEMM launch failed"; +} + +void launch_gemm( + void* output, + const void* activation, + const void* weight, + const void* output_scale, + int m, + int n, + int k, + cudaStream_t stream) { + if (n >= 16384) { + kda_sm120_fp8_skinny::run<32, 512>(output, activation, weight, output_scale, n, k, stream, m); + } else if (n >= 8192) { + kda_sm120_fp8_skinny::run<64, 256>(output, activation, weight, output_scale, n, k, stream, m); + } else { + kda_sm120_fp8_skinny::run<32, 256>(output, activation, weight, output_scale, n, k, stream, m); + } +} + +} // namespace kda_sm120_fp8_skinny + +namespace sglang { + +/** + * \brief SM120 per-tensor FP8 skinny GEMM for the KDA backend. + * + * SGLang performs its reference static quantization before this entry so the + * full path remains within BF16 tolerance of the FlashInfer fallback. + */ +struct KdaSm120Fp8SkinnyGemm { + static void run_quantized( + const tvm::ffi::TensorView input, + const tvm::ffi::TensorView weight, + const tvm::ffi::TensorView output_scale, + const tvm::ffi::TensorView output) { + using namespace host; + + auto M = SymbolicSize{"M"}; + auto K = SymbolicSize{"K"}; + auto N = SymbolicSize{"N"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({M, K}).with_dtype().with_device(device).verify(input); + TensorMatcher({K, N}).with_strides({1, K}).with_dtype().with_device(device).verify(weight); + TensorMatcher({}).with_dtype().with_device(device).verify(output_scale); + TensorMatcher({M, N}).with_dtype().with_device(device).verify(output); + + const int m = static_cast(M.unwrap()); + const int k = static_cast(K.unwrap()); + const int n = static_cast(N.unwrap()); + CHECK_HOST(m == 1 || m == 2 || m == 4 || m == 8 || m == 9) << "M must be one of {1, 2, 4, 8, 9}, got " << m; + CHECK_HOST(k >= 256 && k % 256 == 0) << "K must be a positive multiple of 256, got " << k; + CHECK_HOST(n >= 32 && n % 32 == 0) << "N must be a positive multiple of 32, got " << n; + + const cudaStream_t stream = LaunchKernel::resolve_device(device.unwrap()); + kda_sm120_fp8_skinny::launch_gemm( + output.data_ptr(), input.data_ptr(), weight.data_ptr(), output_scale.data_ptr(), m, n, k, stream); + RuntimeDeviceCheck(); + } +}; + +} // namespace sglang diff --git a/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py b/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py index 45aaa9375..eadf82cbf 100644 --- a/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py +++ b/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py @@ -4,8 +4,6 @@ # KDA provenance: this kernel was automatically optimized by the Humanize2 # workflow (https://github.com/PolyArch/humanize) and Kernel Design Agents # (https://github.com/mit-han-lab/kernel-design-agents). -# Source: https://github.com/BBuf/KDA-Pilot/pull/195 @ -# 516c976cee824a236679adf6eb525275a0a9a120. # Redistribution and use in source and binary forms, with or without diff --git a/python/sglang/kernels/kda_kernels/sm120_fp8.py b/python/sglang/kernels/kda_kernels/sm120_fp8.py new file mode 100644 index 000000000..fcc497e64 --- /dev/null +++ b/python/sglang/kernels/kda_kernels/sm120_fp8.py @@ -0,0 +1,151 @@ +"""Adaptive SM12x dispatch for small-M static per-tensor FP8 linear.""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + import torch + + +_SUPPORTED_KDA_M = (2, 4, 8, 9) + +# The streaming GEMV is faster for supported M=1 shapes, so the CUTLASS path +# keeps only the oversized M=1 projection that GEMV cannot serve. M>=2 shapes +# are enabled only after BF16 comparison, cold-L2 benchmarking, and model-level +# E2E validation on RTX PRO 6000 Blackwell. +_KDA_M_BY_PROJECTION = { + (5120, 8192): _SUPPORTED_KDA_M, + (5120, 16384): (2, 4, 8), + (6144, 5120): _SUPPORTED_KDA_M, + (4096, 4096): _SUPPORTED_KDA_M, + (5120, 7168): _SUPPORTED_KDA_M, + (5120, 5120): (2, 4, 8, 9), + (5120, 34816): (1, 2, 4, 8, 9), +} +_KDA_PRODUCTION_SHAPES = frozenset( + (m, k, n) + for (k, n), supported_m in _KDA_M_BY_PROJECTION.items() + for m in supported_m +) + + +@functools.cache +def _device_capability(device: torch.device) -> tuple[int, int]: + import torch + + return torch.cuda.get_device_capability(device) + + +@functools.cache +def _has_kda_runtime() -> bool: + try: + from sglang.kernels.jit.utils.deps import get_cutlass_include_paths + + get_cutlass_include_paths() + return True + except (ModuleNotFoundError, RuntimeError): + return False + + +def _supports_common( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: Optional[torch.Tensor], + output_scale: Optional[torch.Tensor], + bias: Optional[torch.Tensor], +) -> bool: + import torch + + if ( + input.device.type != "cuda" + or input.ndim != 2 + or input.dtype != torch.bfloat16 + or not input.is_contiguous() + or input_scale is None + or output_scale is None + or bias is not None + ): + return False + + m, k = input.shape + if weight.ndim != 2: + return False + n = weight.shape[1] + if ( + weight.dtype != torch.float8_e4m3fn + or weight.shape[0] != k + or weight.stride() != (1, k) + ): + return False + if ( + input_scale.dtype != torch.float32 + or input_scale.numel() != 1 + or not input_scale.is_contiguous() + or output_scale.dtype != torch.float32 + or output_scale.numel() != 1 + or not output_scale.is_contiguous() + ): + return False + if any(t.device != input.device for t in (weight, input_scale, output_scale)): + return False + return True + + +def try_sm120_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: Optional[torch.Tensor], + output_scale: Optional[torch.Tensor], + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor | None: + """Run the best qualified SM12x FP8 small-M kernel, or return ``None``.""" + if not _supports_common( + input, + weight, + input_scale, + output_scale, + bias, + ): + return None + + assert input_scale is not None + assert output_scale is not None + capability = _device_capability(input.device) + if capability[0] != 12: + return None + + m, k = input.shape + n = weight.shape[1] + use_native_gemv = False + if m == 1: + from sglang.kernels.ops.gemm.sm120_fp8_gemv import use_sm120_fp8_gemv + + use_native_gemv = use_sm120_fp8_gemv(m, n, k) + + use_kda_gemm = ( + not use_native_gemv + and capability == (12, 0) + and (m, k, n) in _KDA_PRODUCTION_SHAPES + and _has_kda_runtime() + ) + if not use_native_gemv and not use_kda_gemm: + return None + + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + + quantized, _ = static_quant_fp8(input, input_scale, repeat_scale=False) + if use_native_gemv: + from sglang.kernels.ops.gemm.sm120_fp8_gemv import sm120_fp8_gemv + + return sm120_fp8_gemv(quantized, weight.t(), output_scale.reshape(1)) + + from sglang.kernels.kda_kernels.sm120_fp8_skinny_gemm_sm120 import ( + _run_sm120_fp8_skinny_gemm_quantized, + ) + + return _run_sm120_fp8_skinny_gemm_quantized(quantized, weight, output_scale) + + +__all__ = ["try_sm120_fp8_linear"] diff --git a/python/sglang/kernels/kda_kernels/sm120_fp8_skinny_gemm_sm120.py b/python/sglang/kernels/kda_kernels/sm120_fp8_skinny_gemm_sm120.py new file mode 100644 index 000000000..d459f1c93 --- /dev/null +++ b/python/sglang/kernels/kda_kernels/sm120_fp8_skinny_gemm_sm120.py @@ -0,0 +1,90 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# KDA provenance: this kernel was automatically optimized by the Humanize2 +# workflow (https://github.com/PolyArch/humanize) and Kernel Design Agents +# (https://github.com/mit-han-lab/kernel-design-agents). + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# This file integrates the SM120 FP8 skinny GEMM with SGLang's JIT kernel and +# KDA backend interfaces. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sglang.kernels.jit.utils import cache_once, load_jit +from sglang.kernels.kda_kernels import _cuda_source + +if TYPE_CHECKING: + import torch + from tvm_ffi.module import Module + + +@cache_once +def _jit_sm120_fp8_skinny_module() -> Module: + import torch + + if torch.cuda.get_device_capability() != (12, 0): + raise RuntimeError("KDA FP8 skinny GEMM requires CUDA SM120") + return load_jit( + "kda_sm120_fp8_skinny_gemm", + cuda_files=[_cuda_source("gemm/sm120_fp8_skinny_gemm.cuh")], + cuda_wrappers=[ + ("run_quantized", "KdaSm120Fp8SkinnyGemm::run_quantized"), + ], + extra_cuda_cflags=[ + "-O3", + "-DCUTLASS_ENABLE_GDC_FOR_SM100", + "--expt-relaxed-constexpr", + "-static-global-template-stub=false", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + ], + extra_dependencies=["cutlass"], + ) + + +def _run_sm120_fp8_skinny_gemm_quantized( + input: torch.Tensor, + weight: torch.Tensor, + output_scale: torch.Tensor, +) -> torch.Tensor: + """Run the low-level FP8-input entry used for kernel qualification.""" + import torch + + output = torch.empty( + (input.shape[0], weight.shape[1]), + dtype=torch.bfloat16, + device=input.device, + ) + _jit_sm120_fp8_skinny_module().run_quantized( + input, weight, output_scale.reshape(()), output + ) + return output diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index c148e2d1f..4e66ef727 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: _CUDA = frozenset({CapabilityRequirement.CUDA}) _SM90 = frozenset({CapabilityRequirement.cuda(min_sm=(9, 0), max_sm=(9, 0))}) _SM120 = frozenset({CapabilityRequirement.cuda(min_sm=(12, 0), max_sm=(12, 0))}) +_SM12X = frozenset({CapabilityRequirement.cuda(min_sm=(12, 0), max_sm=(12, 9))}) _KDA_PACKAGE = "sglang.kernels.kda_kernels" @@ -230,6 +231,21 @@ register_kernel( ), ) ) +register_kernel( + KernelSpec( + op="gemm.sm120_fp8_linear", + backend=KernelBackend.KDA, + target=f"{_KDA_PACKAGE}.sm120_fp8:try_sm120_fp8_linear", + capabilities=_SM12X, + format_signature=FormatSignature( + supported_dtypes=("bfloat16", "float8_e4m3fn", "float32"), + description=( + "Adaptive static per-tensor FP8 linear for small batches on SM12x" + ), + ), + description=("SM12x streaming GEMV plus KDA CUTLASS skinny GEMM dispatch."), + ) +) def fp8_scaled_mm( @@ -297,6 +313,23 @@ def try_qwen3x_nvfp4_gemm( ) +def try_sm120_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: Optional[torch.Tensor], + output_scale: Optional[torch.Tensor], + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor | None: + """Run the best validated SM12x FP8 small-batch path, or return ``None``.""" + return get_kernel("gemm.sm120_fp8_linear", KernelBackend.KDA)( + input, + weight, + input_scale, + output_scale, + bias, + ) + + __all__ = [ "Fp8ScaledMMOp", "bmm_fp8", @@ -304,6 +337,7 @@ __all__ = [ "fp8_scaled_mm", "tiny_gemm_bf16", "try_qwen3x_nvfp4_gemm", + "try_sm120_fp8_linear", ] diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index f0fbf9275..ec7a6bf4f 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -4,7 +4,6 @@ from __future__ import annotations import logging -import os from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import regex as re @@ -536,14 +535,9 @@ class ModelOptFp8LinearMethod(LinearMethodBase): self.use_marlin = ( envs.SGLANG_FORCE_FP8_MARLIN.get() or can_auto_enable_marlin_fp8() ) - # SM120 decode fast path: cuBLAS serves M=1 fp8 GEMMs with SM89 tiles - # at 50-70% DRAM bandwidth for mid-sized N; a streaming GEMV recovers - # the gap. Kill switch: SGLANG_DISABLE_SM120_FP8_GEMV=1. - self.use_sm120_gemv = ( - is_cuda() - and torch.cuda.get_device_capability()[0] == 12 - and os.environ.get("SGLANG_DISABLE_SM120_FP8_GEMV", "0") != "1" - ) + # The SM12x facade selects the best qualified small-M FP8 kernel. + cuda_capability = torch.cuda.get_device_capability() if is_cuda() else None + self.use_sm120_fp8 = cuda_capability is not None and cuda_capability[0] == 12 def create_weights( self, @@ -614,12 +608,12 @@ class ModelOptFp8LinearMethod(LinearMethodBase): layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) if ( - self.use_sm120_gemv + self.use_sm120_fp8 and layer.weight_scale.numel() == 1 and layer.input_scale.numel() == 1 ): - # Combined GEMM epilogue scale for the SM120 M=1 GEMV fast path. - layer.sm120_gemv_alpha = ( + # Precompute the combined epilogue scale for the SM12x facade. + layer.sm120_fp8_alpha = ( (layer.input_scale.float() * layer.weight_scale.float()) .reshape(1) .contiguous() @@ -646,26 +640,18 @@ class ModelOptFp8LinearMethod(LinearMethodBase): size_k=layer.input_size_per_partition, bias=bias, ) - if ( - self.use_sm120_gemv - and bias is None - and x.dim() == 2 - and x.shape[0] == 1 - and hasattr(layer, "sm120_gemv_alpha") - ): - from sglang.kernels.ops.gemm.sm120_fp8_gemv import ( - sm120_fp8_gemv, - use_sm120_fp8_gemv, + if self.use_sm120_fp8: + from sglang.kernels.ops.gemm import try_sm120_fp8_linear + + output = try_sm120_fp8_linear( + x, + layer.weight, + layer.input_scale, + getattr(layer, "sm120_fp8_alpha", None), + bias, ) - - # layer.weight is the [K, N] transposed view of an [N, K]-contiguous - # buffer, so .t() recovers the row-major weight the GEMV streams. - w = layer.weight.t() - if use_sm120_fp8_gemv(1, w.shape[0], w.shape[1]) and w.is_contiguous(): - from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 - - qinput, _ = static_quant_fp8(x, layer.input_scale, repeat_scale=False) - return sm120_fp8_gemv(qinput, w, layer.sm120_gemv_alpha) + if output is not None: + return output if layer.use_flashinfer_bmm: return apply_fp8_linear_bmm_flashinfer( input=x, @@ -2900,7 +2886,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): layer: FusedMoE, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - # Note: dispatch_output may be a DeepEPLLDispatchOutput (no topk_output # attribute -- topk_ids/topk_weights live directly on the dispatch # tuple). Defer per-attribute access to the branches that actually diff --git a/test/registered/kernels/benchmark/gemm/bench_kda_fp8_skinny_gemm.py b/test/registered/kernels/benchmark/gemm/bench_kda_fp8_skinny_gemm.py new file mode 100644 index 000000000..cd8af5a86 --- /dev/null +++ b/test/registered/kernels/benchmark/gemm/bench_kda_fp8_skinny_gemm.py @@ -0,0 +1,149 @@ +"""Cross-model benchmark for the KDA SM120 FP8 skinny GEMM.""" + +from __future__ import annotations + +import sys + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.kda_kernels.sm120_fp8_skinny_gemm_sm120 import ( + _run_sm120_fp8_skinny_gemm_quantized, +) +from sglang.kernels.ops.gemm.sm120_fp8_gemv import sm120_fp8_gemv +from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + apply_fp8_linear_bmm_flashinfer, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=90, stage="base-b", runner_config="1-gpu-small") + + +def _make_inputs(m: int, n: int, k: int): + input = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + weight = ( + torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + .mul_(0.25) + .to(torch.float8_e4m3fn) + .t() + ) + weight_scale = torch.tensor(0.02, dtype=torch.float32, device="cuda") + input_scale = torch.tensor(0.025, dtype=torch.float32, device="cuda") + output_scale = input_scale * weight_scale + return input, weight, weight_scale, input_scale, output_scale + + +def _torch_impl(input, weight, weight_scale, input_scale, output_scale): + del output_scale + return apply_fp8_linear( + input, + weight, + weight_scale, + input_scale, + cutlass_fp8_supported=False, + pad_output=False, + ) + + +def _kda_impl(input, weight, weight_scale, input_scale, output_scale): + del weight_scale + qinput, _ = static_quant_fp8(input, input_scale, repeat_scale=False) + return _run_sm120_fp8_skinny_gemm_quantized(qinput, weight, output_scale) + + +def _flashinfer_impl(input, weight, weight_scale, input_scale, output_scale): + del output_scale + return apply_fp8_linear_bmm_flashinfer(input, weight, weight_scale, input_scale) + + +def _native_gemv_impl(input, weight, weight_scale, input_scale, output_scale): + del weight_scale + qinput, _ = static_quant_fp8(input, input_scale, repeat_scale=False) + return sm120_fp8_gemv(qinput, weight.t(), output_scale.reshape(1)) + + +FN_MAP = { + "kda": _kda_impl, + "torch": _torch_impl, + "flashinfer": _flashinfer_impl, + "native": _native_gemv_impl, +} + + +PROJECTIONS = [ + ("Qwen3.8-27B/gdn-in", 16384, 5120), + ("Qwen3.8-27B/attn-qkv", 8192, 5120), + ("Qwen3.8-27B/out", 5120, 6144), + ("Qwen3-8B+Llama-3.1-8B/qkv", 6144, 4096), + ("Qwen3-8B+Llama-3.1-8B/o", 4096, 4096), + ("Qwen3-8B/gate-up", 24576, 4096), + ("Qwen3-8B/down", 4096, 12288), + ("Qwen3-14B/qkv", 7168, 5120), + ("Qwen3-14B/o", 5120, 5120), + ("Qwen3-14B/gate-up", 34816, 5120), + ("Qwen3-14B/down", 5120, 17408), + ("Llama-3.1-8B/gate-up", 28672, 4096), + ("Llama-3.1-8B/down", 4096, 14336), + ("Nemotron-3-Super/mamba-in", 18560, 4096), + ("Nemotron-3-Super/mamba-out", 4096, 8192), + ("Nemotron-3-Super/shared-up", 5376, 4096), + ("Nemotron-3-Super/shared-down", 4096, 5376), +] +M_VALUES = (1, 2, 4, 8, 9) +MODEL_FP8_CASES = [(model, m, n, k) for m in M_VALUES for model, n, k in PROJECTIONS] +NATIVE_M1_CASES = [ + ("Qwen3.8-27B/attn-qkv", 8192, 5120), + ("Qwen3.8-27B/out", 5120, 6144), + ("Qwen3-8B+Llama-3.1-8B/o", 4096, 4096), + ("Qwen3-14B/qkv", 7168, 5120), + ("Nemotron-3-Super/shared-up", 5376, 4096), +] + + +# The full sweep covers decode/verify M values and representative per-tensor +# FP8 projections from several model families. +@marker.parametrize( + "model,m,n,k", + MODEL_FP8_CASES, + [("Qwen3.8-27B/attn-qkv", 9, 8192, 5120)], +) +@marker.benchmark("provider", ["kda", "torch", "flashinfer"]) +def benchmark(model: str, m: int, n: int, k: int, provider: str): + del model + args = _make_inputs(m, n, k) + return marker.do_bench( + FN_MAP[provider], + input_args=args, + graph_clone_args=(0, 1, 2, 3), + disable_log_bandwidth=True, + ) + + +@marker.parametrize( + "model,n,k", + NATIVE_M1_CASES, + [("Qwen3.8-27B/attn-qkv", 8192, 5120)], +) +@marker.benchmark("provider", ["kda", "native"]) +def benchmark_m1_native(model: str, n: int, k: int, provider: str): + """Compare M=1 with SGLang's existing SM120 GEMV before dispatching.""" + del model + args = _make_inputs(1, n, k) + return marker.do_bench( + FN_MAP[provider], + input_args=args, + graph_clone_args=(0, 1, 2, 3), + disable_log_bandwidth=True, + ) + + +if __name__ == "__main__": + if not ( + torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 0) + ): + print("[skip] KDA FP8 skinny GEMM benchmark requires CUDA SM120") + sys.exit(0) + benchmark.run() + benchmark_m1_native.run() diff --git a/test/registered/kernels/ops/gemm/test_sm120_fp8_linear.py b/test/registered/kernels/ops/gemm/test_sm120_fp8_linear.py new file mode 100644 index 000000000..96cd27dbf --- /dev/null +++ b/test/registered/kernels/ops/gemm/test_sm120_fp8_linear.py @@ -0,0 +1,175 @@ +"""Correctness and dispatch tests for SM12x small-M FP8 linear.""" + +from __future__ import annotations + +import sys + +import pytest +import torch + +from sglang.kernels.ops.gemm import try_sm120_fp8_linear +from sglang.srt.layers.quantization.fp8_utils import apply_fp8_linear_bmm_flashinfer +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=180, stage="base-b", runner_config="1-gpu-small") + +if not (torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 0)): + pytest.skip( + "SM120 FP8 linear dispatch requires CUDA SM120", allow_module_level=True + ) + + +ALL_M = (1, 2, 4, 8, 9) + +# M=1 uses the streaming GEMV whenever its broad shape gate accepts the call; +# qualified M>=2 cases and the oversized M=1 gate/up projection use KDA. +SUPPORTED_CONFIGS = [ + ("Qwen3.8-27B", "attn-qkv", 8192, 5120, ALL_M), + ("Qwen3.8-27B", "gdn-in", 16384, 5120, (1, 2, 4, 8)), + ("Qwen3.8-27B", "out", 5120, 6144, ALL_M), + ("Qwen3-8B/Llama-3.1-8B", "qkv", 6144, 4096, (1,)), + ("Qwen3-8B/Llama-3.1-8B", "o", 4096, 4096, ALL_M), + ("Qwen3-14B", "qkv", 7168, 5120, ALL_M), + ("Qwen3-14B", "o", 5120, 5120, ALL_M), + ("Qwen3-14B", "gate-up", 34816, 5120, ALL_M), + ("Nemotron-3-Super", "shared-up", 5376, 4096, (1,)), +] +SUPPORTED_SHAPES = [ + pytest.param(m, n, k, id=f"{model}-{op}-m{m}".lower()) + for model, op, n, k, supported_m in SUPPORTED_CONFIGS + for m in supported_m +] + +# These M>1 shapes are numerically valid but missed either the KDA accuracy or +# performance gate. Their M=1 variants may still use the streaming GEMV. +UNSUPPORTED_CONFIGS = [ + pytest.param((3,), 8192, 5120, id="unsupported-m"), + pytest.param((9,), 16384, 5120, id="qwen38-gdn-in"), + pytest.param((2, 4, 8, 9), 5376, 4096, id="nemotron-shared-up"), + pytest.param((2, 4, 8, 9), 4096, 5376, id="nemotron-shared-down"), + pytest.param((8,), 6144, 4096, id="qwen3-8b-qkv"), +] + + +def _make_inputs(m: int, n: int, k: int, seed: int = 0): + torch.manual_seed(seed) + input_scale = torch.tensor(0.025, dtype=torch.float32, device="cuda") + weight_scale = torch.tensor(0.02, dtype=torch.float32, device="cuda") + input = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + weight = ( + torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + .mul_(32) + .to(torch.float8_e4m3fn) + .t() + ) + return input, weight, weight_scale, input_scale + + +def _reference(args): + return apply_fp8_linear_bmm_flashinfer(*args) + + +def _assert_matches_flashinfer(actual, expected): + # CUTLASS and cuBLAS may accumulate in a different order. On SM120 the + # observed differences are sparse and stay within standard BF16 tolerance. + torch.testing.assert_close(actual, expected) + + +def _run_sm120(args, *, m=None, vector_scales=False, bias=None): + input, weight, weight_scale, input_scale = args + if m is not None: + input = input[:m] + output_scale = input_scale * weight_scale + if vector_scales: + input_scale = input_scale.reshape(1) + output_scale = output_scale.reshape(1) + return try_sm120_fp8_linear( + input, + weight, + input_scale, + output_scale, + bias, + ) + + +@pytest.mark.parametrize("seed", [0, 1, 7]) +@pytest.mark.parametrize("m,n,k", SUPPORTED_SHAPES) +def test_supported_shapes_match_flashinfer(m: int, n: int, k: int, seed: int): + args = _make_inputs(m, n, k, seed) + expected = _reference(args) + actual = _run_sm120(args) + assert actual is not None + _assert_matches_flashinfer(actual, expected) + + +@pytest.mark.parametrize("m", ALL_M) +def test_cuda_graph_replay_uses_current_input(m: int): + args = _make_inputs(m, 8192, 5120) + input, weight, weight_scale, input_scale = args + output_scale = input_scale * weight_scale + + # Compile the selected provider before capture; JIT compilation is not + # CUDA Graph safe and this test must also work when selected in isolation. + warmup = _run_sm120(args) + assert warmup is not None + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = try_sm120_fp8_linear(input, weight, input_scale, output_scale) + assert actual is not None + + # A replay must consume the new activation rather than the values present + # during capture. This also exercises the quantize-to-GEMM dependency. + torch.manual_seed(17) + input.copy_(torch.randn_like(input).mul_(8)) + graph.replay() + expected = _reference(args) + _assert_matches_flashinfer(actual, expected) + + +@pytest.mark.parametrize("m", (2, 4, 8)) +def test_saturated_real_activation_range_matches_flashinfer(m: int): + args = _make_inputs(m, 16384, 5120, seed=11) + input = args[0] + # Static ModelOpt scales can expose both saturation and FP8 rounding + # boundaries in live GDN activations. This distribution reproduces the + # class of mismatch that the former fused quantizer caused in E2E decode. + input.mul_(32) + input[:, :8] = torch.tensor( + [-32.0, -11.25, -11.0, -0.013, 0.013, 11.0, 11.25, 32.0], + dtype=torch.bfloat16, + device="cuda", + ) + expected = _reference(args) + actual = _run_sm120(args) + assert actual is not None + _assert_matches_flashinfer(actual, expected) + + +def test_scalar_and_vector_scale_layouts_dispatch(): + args = _make_inputs(8, 8192, 5120) + expected = _run_sm120(args) + assert expected is not None + + # Some callers retain per-tensor scales as one-element vectors. The Python + # facade normalizes both layouts to the scalar TVM-FFI contract. + vector_scale_output = _run_sm120(args, vector_scales=True) + torch.testing.assert_close(vector_scale_output, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("m_values,n,k", UNSUPPORTED_CONFIGS) +def test_unsupported_shapes_fall_back(m_values, n: int, k: int): + args = _make_inputs(max(m_values), n, k) + for m in m_values: + assert _run_sm120(args, m=m) is None + + +def test_bias_falls_back(): + args = _make_inputs(8, 8192, 5120) + bias = torch.zeros(args[1].shape[1], dtype=torch.bfloat16, device="cuda") + assert _run_sm120(args, bias=bias) is None + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py index fc3f9eddd..60abc2700 100644 --- a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py +++ b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py @@ -44,6 +44,7 @@ EXPECTED = { "diffusion.flux2_qkv_epilogue": {"KDA"}, "diffusion.flux2_token_cat_fp8": {"KDA"}, "gemm.qwen3x_nvfp4": {"KDA"}, + "gemm.sm120_fp8_linear": {"KDA"}, } _CPU = PlatformInfo(device_type="cpu")