Delete CUTLASS FP8 blockwise for SM90 and SM100, move SM120 to JIT and add SwapAB (#30438)
Co-authored-by: Brayden Zhong <brayden.zhong@radixark.ai> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Co-authored-by: root <root@sgl-b300-inference.datacrunch.io> Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
co-authored by
Brayden Zhong
Claude Sonnet 5
root
Brayden Zhong
Brayden Zhong
Claude Sonnet 5
root
Brayden Zhong
parent
c124bec99d
commit
7431f35fd8
@@ -0,0 +1,25 @@
|
||||
/* Copyright 2026 SGLang Team. All Rights Reserved.
|
||||
|
||||
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.
|
||||
==============================================================================*/
|
||||
|
||||
#include "fp8_blockwise_scaled_mm_sm120.cuh"
|
||||
|
||||
void fp8_blockwise_scaled_mm(
|
||||
tvm::ffi::TensorView out,
|
||||
tvm::ffi::TensorView mat_a,
|
||||
tvm::ffi::TensorView mat_b,
|
||||
tvm::ffi::TensorView scales_a,
|
||||
tvm::ffi::TensorView scales_b) {
|
||||
fp8_blockwise_scaled_mm_sm120(out, mat_a, mat_b, scales_a, scales_b);
|
||||
}
|
||||
@@ -0,0 +1,502 @@
|
||||
/* Copyright 2026 SGLang Team. All Rights Reserved.
|
||||
|
||||
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.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/runtime.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
using namespace host;
|
||||
|
||||
// clang-format off
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/detail/blockwise_scale_layout.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
// clang-format on
|
||||
|
||||
#define CUTLASS_CHECK(status) \
|
||||
{ \
|
||||
cutlass::Status error = status; \
|
||||
RuntimeCheck(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
|
||||
}
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
|
||||
|
||||
template <
|
||||
typename OutType,
|
||||
typename MmaTileShape,
|
||||
typename PerSmTileShape,
|
||||
typename EpilogueTileShape,
|
||||
typename ScalesPerTile,
|
||||
int TileSizeM_ = 128,
|
||||
class ClusterShape = Shape<_1, _1, _1>>
|
||||
void launch_sm120_fp8_blockwise_scaled_mm(
|
||||
tvm::ffi::TensorView out,
|
||||
tvm::ffi::TensorView a,
|
||||
tvm::ffi::TensorView b,
|
||||
tvm::ffi::TensorView scales_a,
|
||||
tvm::ffi::TensorView scales_b,
|
||||
cudaStream_t stream) {
|
||||
using ElementBlockScale = float;
|
||||
|
||||
// A matrix configuration
|
||||
using ElementA = cutlass::float_e4m3_t; // Element type for A matrix operand
|
||||
using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand
|
||||
constexpr int AlignmentA =
|
||||
128 / cutlass::sizeof_bits<ElementA>::value; // Memory access granularity/alignment of A matrix in units of
|
||||
// elements (up to 16 bytes)
|
||||
|
||||
// B matrix configuration
|
||||
using ElementB = cutlass::float_e4m3_t; // Element type for B matrix operand
|
||||
using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand
|
||||
constexpr int AlignmentB =
|
||||
128 / cutlass::sizeof_bits<ElementB>::value; // Memory access granularity/alignment of B matrix in units of
|
||||
// elements (up to 16 bytes)
|
||||
|
||||
// C/D matrix configuration
|
||||
using ElementD = OutType; // Element type for D matrix operand
|
||||
using ElementC = void; // Element type for C matrix operand
|
||||
using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand
|
||||
using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand
|
||||
constexpr int AlignmentD =
|
||||
128 / cutlass::sizeof_bits<ElementD>::value; // Memory access granularity/alignment of C matrix in units of
|
||||
// elements (up to 16 bytes)
|
||||
constexpr int AlignmentC =
|
||||
AlignmentD; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
|
||||
|
||||
// Kernel functional config
|
||||
using ElementAccumulator = float; // Element type for internal accumulation
|
||||
using ArchTag = cutlass::arch::Sm120; // Tag indicating the minimum SM that supports the intended feature
|
||||
using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag - changed from OpClassBlockScaledTensorOp
|
||||
|
||||
static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{});
|
||||
static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile;
|
||||
static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{});
|
||||
static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{});
|
||||
|
||||
using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig<
|
||||
ScaleGranularityM,
|
||||
ScaleGranularityN,
|
||||
ScaleGranularityK,
|
||||
cute::UMMA::Major::MN,
|
||||
cute::UMMA::Major::K>;
|
||||
// FP8 Block-wise scaling configuration
|
||||
using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand
|
||||
using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand
|
||||
|
||||
constexpr bool kCanUsePingpong = (64 % ScaleGranularityM == 0);
|
||||
|
||||
int m = a.size(0);
|
||||
int k = a.size(1);
|
||||
int n = b.size(1);
|
||||
|
||||
auto a_ptr = static_cast<ElementA*>(a.data_ptr());
|
||||
auto b_ptr = static_cast<ElementB*>(b.data_ptr());
|
||||
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
|
||||
|
||||
auto scales_a_ptr = static_cast<ElementBlockScale*>(scales_a.data_ptr());
|
||||
auto scales_b_ptr = static_cast<ElementBlockScale*>(scales_b.data_ptr());
|
||||
|
||||
LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
|
||||
LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
|
||||
|
||||
auto run_gemm = [&](auto tag) -> cutlass::Status {
|
||||
using GemmKernel = decltype(tag);
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
Gemm gemm_op;
|
||||
|
||||
using StrideA = typename GemmKernel::StrideA;
|
||||
using StrideB = typename GemmKernel::StrideB;
|
||||
using StrideC = typename GemmKernel::StrideD;
|
||||
|
||||
StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1));
|
||||
StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
|
||||
StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1));
|
||||
|
||||
typename GemmKernel::MainloopArguments mainloop_args{
|
||||
a_ptr, stride_a, b_ptr, stride_b, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB};
|
||||
|
||||
typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c};
|
||||
epilogue_args.thread.alpha = 1.0f;
|
||||
|
||||
typename Gemm::Arguments args = {
|
||||
cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
{m, n, k, 1},
|
||||
mainloop_args,
|
||||
epilogue_args,
|
||||
};
|
||||
|
||||
auto can_implement = gemm_op.can_implement(args);
|
||||
if (can_implement != cutlass::Status::kSuccess) {
|
||||
return can_implement;
|
||||
}
|
||||
|
||||
size_t workspace_size = gemm_op.get_workspace_size(args);
|
||||
auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device());
|
||||
void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr();
|
||||
|
||||
auto init_status = gemm_op.initialize(args, workspace, stream);
|
||||
if (init_status != cutlass::Status::kSuccess) {
|
||||
return init_status;
|
||||
}
|
||||
|
||||
return gemm_op.run(stream);
|
||||
};
|
||||
|
||||
using CooperativeCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
ArchTag,
|
||||
OperatorClass,
|
||||
PerSmTileShape,
|
||||
ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator,
|
||||
ElementC,
|
||||
LayoutCTag,
|
||||
AlignmentC,
|
||||
ElementD,
|
||||
LayoutDTag,
|
||||
AlignmentD,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
|
||||
|
||||
using CooperativeStageCount = cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
|
||||
sizeof(typename CooperativeCollectiveEpilogue::SharedStorage))>;
|
||||
|
||||
using CooperativeCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag,
|
||||
OperatorClass,
|
||||
ElementA,
|
||||
cute::tuple<LayoutATag, LayoutSFA>,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
cute::tuple<LayoutBTag, LayoutSFB>,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
MmaTileShape,
|
||||
ClusterShape,
|
||||
CooperativeStageCount,
|
||||
cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp;
|
||||
|
||||
using CooperativeGemmKernelStreamK = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>,
|
||||
CooperativeCollectiveMainloop,
|
||||
CooperativeCollectiveEpilogue,
|
||||
cutlass::gemm::StreamKScheduler>;
|
||||
using CooperativeGemmKernelVoid = cutlass::gemm::kernel::
|
||||
GemmUniversal<Shape<int, int, int, int>, CooperativeCollectiveMainloop, CooperativeCollectiveEpilogue, void>;
|
||||
|
||||
auto run_cooperative = [&]() -> cutlass::Status {
|
||||
static const uint32_t kNumSM = host::runtime::get_sm_count(a.device().device_id);
|
||||
constexpr int kTileM = size<0>(MmaTileShape{});
|
||||
constexpr int kTileN = size<1>(MmaTileShape{});
|
||||
uint64_t tiles = static_cast<uint64_t>((m + kTileM - 1) / kTileM) * ((n + kTileN - 1) / kTileN);
|
||||
uint32_t last_wave = static_cast<uint32_t>(tiles % kNumSM);
|
||||
if (last_wave == 0) last_wave = kNumSM;
|
||||
float waste = 1.0f - static_cast<float>(last_wave) / static_cast<float>(kNumSM);
|
||||
return (waste > 0.5f) ? run_gemm(CooperativeGemmKernelStreamK{}) : run_gemm(CooperativeGemmKernelVoid{});
|
||||
};
|
||||
|
||||
cutlass::Status status = cutlass::Status::kSuccess;
|
||||
if constexpr (kCanUsePingpong) {
|
||||
using PingpongMmaTileShape_MNK = Shape<_64, _128, _128>;
|
||||
using PingpongCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
ArchTag,
|
||||
OperatorClass,
|
||||
PerSmTileShape,
|
||||
ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator,
|
||||
ElementC,
|
||||
LayoutCTag,
|
||||
AlignmentC,
|
||||
ElementD,
|
||||
LayoutDTag,
|
||||
AlignmentD,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
|
||||
|
||||
using PingpongStageCount = cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
|
||||
sizeof(typename PingpongCollectiveEpilogue::SharedStorage))>;
|
||||
|
||||
using PingpongCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag,
|
||||
OperatorClass,
|
||||
ElementA,
|
||||
cute::tuple<LayoutATag, LayoutSFA>,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
cute::tuple<LayoutBTag, LayoutSFB>,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
PingpongMmaTileShape_MNK,
|
||||
ClusterShape,
|
||||
PingpongStageCount,
|
||||
cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120>::CollectiveOp;
|
||||
|
||||
using PingpongGemmKernel = cutlass::gemm::kernel::
|
||||
GemmUniversal<Shape<int, int, int, int>, PingpongCollectiveMainloop, PingpongCollectiveEpilogue, void>;
|
||||
|
||||
if (m <= 64) {
|
||||
status = run_gemm(PingpongGemmKernel{});
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
status = run_cooperative();
|
||||
}
|
||||
} else {
|
||||
status = run_cooperative();
|
||||
}
|
||||
} else {
|
||||
status = run_cooperative();
|
||||
}
|
||||
|
||||
CUTLASS_CHECK(status);
|
||||
}
|
||||
|
||||
// Transposed GEMM D^T = Wgemm(weight, activation): puts tokens on the N axis.
|
||||
template <
|
||||
typename OutType,
|
||||
typename MmaTileShape,
|
||||
typename PerSmTileShape,
|
||||
typename EpilogueTileShape,
|
||||
typename ScalesPerTile,
|
||||
class ClusterShape = Shape<_1, _1, _1>>
|
||||
void launch_sm120_fp8_blockwise_scaled_mm_swapab(
|
||||
tvm::ffi::TensorView out,
|
||||
tvm::ffi::TensorView a,
|
||||
tvm::ffi::TensorView b,
|
||||
tvm::ffi::TensorView scales_a,
|
||||
tvm::ffi::TensorView scales_b,
|
||||
cudaStream_t stream) {
|
||||
using ElementBlockScale = float;
|
||||
|
||||
using ElementA = cutlass::float_e4m3_t; // A' = weight
|
||||
using LayoutATag = cutlass::layout::RowMajor; // weight [N, K] is row-major
|
||||
constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
|
||||
|
||||
using ElementB = cutlass::float_e4m3_t; // B' = activation
|
||||
using LayoutBTag = cutlass::layout::ColumnMajor; // activation as [K, M] column-major
|
||||
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
|
||||
|
||||
using ElementD = OutType;
|
||||
using ElementC = void;
|
||||
using LayoutCTag = cutlass::layout::ColumnMajor; // D' = out^T is column-major
|
||||
using LayoutDTag = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
|
||||
constexpr int AlignmentC = AlignmentD;
|
||||
|
||||
using ElementAccumulator = float;
|
||||
using ArchTag = cutlass::arch::Sm120;
|
||||
using OperatorClass = cutlass::arch::OpClassTensorOp;
|
||||
|
||||
static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{});
|
||||
static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile;
|
||||
static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{});
|
||||
static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{});
|
||||
|
||||
// Operands are swapped, so the scale majors swap relative to the non-swap path:
|
||||
// SFA (weight) is K-major; SFB (per-token activation) is MN-major.
|
||||
using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig<
|
||||
ScaleGranularityM,
|
||||
ScaleGranularityN,
|
||||
ScaleGranularityK,
|
||||
cute::UMMA::Major::K,
|
||||
cute::UMMA::Major::MN>;
|
||||
using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
|
||||
using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
|
||||
|
||||
int m = a.size(0); // original tokens -> swapped N'
|
||||
int k = a.size(1);
|
||||
int n = b.size(1); // original weight cols -> swapped M'
|
||||
|
||||
auto weight_ptr = static_cast<ElementA*>(b.data_ptr());
|
||||
auto act_ptr = static_cast<ElementB*>(a.data_ptr());
|
||||
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
|
||||
auto weight_scale_ptr = static_cast<ElementBlockScale*>(scales_b.data_ptr());
|
||||
auto act_scale_ptr = static_cast<ElementBlockScale*>(scales_a.data_ptr());
|
||||
|
||||
// Swapped problem shape (M', N', K) = (n, m, k).
|
||||
LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1));
|
||||
LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1));
|
||||
|
||||
auto run_gemm = [&](auto tag) -> cutlass::Status {
|
||||
using GemmKernel = decltype(tag);
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
Gemm gemm_op;
|
||||
|
||||
using StrideA = typename GemmKernel::StrideA;
|
||||
using StrideB = typename GemmKernel::StrideB;
|
||||
using StrideC = 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));
|
||||
StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(n, m, 1));
|
||||
|
||||
typename GemmKernel::MainloopArguments mainloop_args{
|
||||
weight_ptr, stride_a, act_ptr, stride_b, weight_scale_ptr, layout_SFA, act_scale_ptr, layout_SFB};
|
||||
|
||||
typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c};
|
||||
epilogue_args.thread.alpha = 1.0f;
|
||||
|
||||
typename Gemm::Arguments args = {
|
||||
cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
{n, m, k, 1},
|
||||
mainloop_args,
|
||||
epilogue_args,
|
||||
};
|
||||
|
||||
auto can_implement = gemm_op.can_implement(args);
|
||||
if (can_implement != cutlass::Status::kSuccess) {
|
||||
return can_implement;
|
||||
}
|
||||
|
||||
size_t workspace_size = gemm_op.get_workspace_size(args);
|
||||
auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device());
|
||||
void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr();
|
||||
|
||||
auto init_status = gemm_op.initialize(args, workspace, stream);
|
||||
if (init_status != cutlass::Status::kSuccess) {
|
||||
return init_status;
|
||||
}
|
||||
|
||||
return gemm_op.run(stream);
|
||||
};
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
ArchTag,
|
||||
OperatorClass,
|
||||
PerSmTileShape,
|
||||
ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator,
|
||||
ElementC,
|
||||
LayoutCTag,
|
||||
AlignmentC,
|
||||
ElementD,
|
||||
LayoutDTag,
|
||||
AlignmentD,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
|
||||
|
||||
using StageCount = cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage))>;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag,
|
||||
OperatorClass,
|
||||
ElementA,
|
||||
cute::tuple<LayoutATag, LayoutSFA>,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
cute::tuple<LayoutBTag, LayoutSFB>,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
MmaTileShape,
|
||||
ClusterShape,
|
||||
StageCount,
|
||||
cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp;
|
||||
|
||||
using GemmKernel =
|
||||
cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
|
||||
|
||||
CUTLASS_CHECK(run_gemm(GemmKernel{}));
|
||||
}
|
||||
|
||||
// swapAB (tile N=32) beats the non-swap 128x128 path for M<=64 or M%4!=0
|
||||
// (cold-L2 CUPTI benchmarks, up to ~1.2x); tile N=16 is unsupported by the
|
||||
// SM120 blockwise collective (needs EPI_TILE_N=32 | CTA_N and B LDSM N>=32).
|
||||
template <typename OutType>
|
||||
void sm120_fp8_blockwise_dispatch_shape(
|
||||
tvm::ffi::TensorView out,
|
||||
tvm::ffi::TensorView a,
|
||||
tvm::ffi::TensorView b,
|
||||
tvm::ffi::TensorView scales_a,
|
||||
tvm::ffi::TensorView scales_b,
|
||||
cudaStream_t stream) {
|
||||
const int m = a.size(0);
|
||||
using EpilogueTileShape = Shape<_128, _64>;
|
||||
if (m <= 64 || (m % 4 != 0)) {
|
||||
launch_sm120_fp8_blockwise_scaled_mm_swapab<
|
||||
OutType,
|
||||
Shape<_128, _32, _128>,
|
||||
Shape<_128, _32, _128>,
|
||||
EpilogueTileShape,
|
||||
Shape<_1, _32, _1>>(out, a, b, scales_a, scales_b, stream);
|
||||
return;
|
||||
}
|
||||
|
||||
using MmaTileShape = Shape<_128, _128, _128>;
|
||||
using PerSmTileShape = Shape<_128, _128, _128>;
|
||||
using ScalesPerTile = Shape<_128, _1, _1>;
|
||||
launch_sm120_fp8_blockwise_scaled_mm<OutType, MmaTileShape, PerSmTileShape, EpilogueTileShape, ScalesPerTile>(
|
||||
out, a, b, scales_a, scales_b, stream);
|
||||
}
|
||||
|
||||
inline void fp8_blockwise_scaled_mm_sm120(
|
||||
tvm::ffi::TensorView out,
|
||||
tvm::ffi::TensorView mat_a,
|
||||
tvm::ffi::TensorView mat_b,
|
||||
tvm::ffi::TensorView scales_a,
|
||||
tvm::ffi::TensorView scales_b) {
|
||||
RuntimeCheck(mat_a.device().device_type == kDLCUDA, "mat_a must be a CUDA tensor");
|
||||
RuntimeCheck(mat_b.device().device_type == kDLCUDA, "mat_b must be a CUDA tensor");
|
||||
|
||||
RuntimeCheck(mat_a.dim() == 2, "mat_a must be a 2D tensor");
|
||||
RuntimeCheck(mat_b.dim() == 2, "mat_b must be a 2D tensor");
|
||||
RuntimeCheck(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
|
||||
RuntimeCheck(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
|
||||
RuntimeCheck(mat_a.size(1) == mat_b.size(0), "mat_a and mat_b shapes cannot be multiplied");
|
||||
|
||||
RuntimeCheck(
|
||||
(mat_a.size(1) * (mat_a.dtype().bits / 8)) % 16 == 0, "mat_a must be multiple of 16 bytes for memory alignment");
|
||||
RuntimeCheck(
|
||||
(mat_b.size(0) * (mat_b.dtype().bits / 8)) % 16 == 0, "mat_b must be multiple of 16 bytes for memory alignment");
|
||||
RuntimeCheck(host::is_type<fp8_e4m3_t>(mat_a.dtype()), "mat_a must be Float8_e4m3fn");
|
||||
RuntimeCheck(host::is_type<fp8_e4m3_t>(mat_b.dtype()), "mat_b must be Float8_e4m3fn");
|
||||
|
||||
RuntimeCheck(mat_a.size(0) == scales_a.size(0), "size of scales_a is not matched");
|
||||
RuntimeCheck(mat_a.size(1) / 128 == scales_a.size(1), "size of scales_a is not matched");
|
||||
RuntimeCheck(mat_b.size(0) / 128 == scales_b.size(0), "size of scales_b is not matched");
|
||||
RuntimeCheck(mat_b.size(1) / 128 == scales_b.size(1), "size of scales_b is not matched");
|
||||
RuntimeCheck(host::is_type<float>(scales_a.dtype()), "scales_a must be Float32");
|
||||
RuntimeCheck(host::is_type<float>(scales_b.dtype()), "scales_b must be Float32");
|
||||
|
||||
RuntimeCheck(
|
||||
(out.size(1) * (out.dtype().bits / 8)) % 16 == 0, "out must be multiple of 16 bytes for memory alignment");
|
||||
|
||||
const cudaStream_t stream = LaunchKernel::resolve_device(mat_a.device());
|
||||
|
||||
if (host::is_type<bf16_t>(out.dtype())) {
|
||||
sm120_fp8_blockwise_dispatch_shape<cutlass::bfloat16_t>(out, mat_a, mat_b, scales_a, scales_b, stream);
|
||||
} else if (host::is_type<fp16_t>(out.dtype())) {
|
||||
sm120_fp8_blockwise_dispatch_shape<cutlass::half_t>(out, mat_a, mat_b, scales_a, scales_b, stream);
|
||||
} else {
|
||||
Panic("out_dtype must be Half or BFloat16");
|
||||
}
|
||||
}
|
||||
|
||||
#endif // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
|
||||
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.common import is_sm120_supported
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
def _fp8_blockwise_cuda_flags() -> list[str]:
|
||||
return [
|
||||
"-DNDEBUG",
|
||||
"-DCUTE_USE_PACKED_TUPLE=1",
|
||||
"-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
|
||||
"-DCUTLASS_VERSIONS_GENERATED",
|
||||
"-DCUTLASS_TEST_LEVEL=0",
|
||||
"-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1",
|
||||
"-DCUTLASS_DEBUG_TRACE_LEVEL=0",
|
||||
"--expt-relaxed-constexpr",
|
||||
"--expt-extended-lambda",
|
||||
]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _fp8_blockwise_arch_env():
|
||||
if not is_sm120_supported():
|
||||
raise RuntimeError(
|
||||
"fp8_blockwise_scaled_mm JIT kernel requires SM120 (Blackwell)."
|
||||
)
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
# sm_*a target (e.g. sm_120a) required, not plain sm_120.
|
||||
with override_jit_cuda_arch(major, minor, suffix="a"):
|
||||
yield
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_fp8_blockwise_module() -> Module:
|
||||
"""Compile and cache the SM120 fp8 blockwise GEMM module (handles fp16 + bf16)."""
|
||||
with _fp8_blockwise_arch_env():
|
||||
return load_jit(
|
||||
"fp8_blockwise_scaled_mm",
|
||||
cuda_files=["gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh"],
|
||||
cuda_wrappers=[
|
||||
("fp8_blockwise_scaled_mm", "fp8_blockwise_scaled_mm"),
|
||||
],
|
||||
extra_dependencies=["cutlass"],
|
||||
extra_cuda_cflags=_fp8_blockwise_cuda_flags(),
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(
|
||||
op_name="fp8_blockwise_scaled_mm",
|
||||
mutates_args=["out"],
|
||||
)
|
||||
def _fp8_blockwise_scaled_mm_custom_op(
|
||||
out: torch.Tensor,
|
||||
mat_a: torch.Tensor,
|
||||
mat_b: torch.Tensor,
|
||||
scales_a: torch.Tensor,
|
||||
scales_b: torch.Tensor,
|
||||
) -> None:
|
||||
module = _jit_fp8_blockwise_module()
|
||||
module.fp8_blockwise_scaled_mm(out, mat_a, mat_b, scales_a, scales_b)
|
||||
|
||||
|
||||
@debug_kernel_api
|
||||
def fp8_blockwise_scaled_mm(
|
||||
mat_a: torch.Tensor,
|
||||
mat_b: torch.Tensor,
|
||||
scales_a: torch.Tensor,
|
||||
scales_b: torch.Tensor,
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""FP8 e4m3 block-wise scaled matmul on SM120."""
|
||||
assert out_dtype in (
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
), f"out_dtype must be Half or BFloat16, got {out_dtype}"
|
||||
|
||||
out = torch.empty(
|
||||
(mat_a.shape[0], mat_b.shape[1]),
|
||||
dtype=out_dtype,
|
||||
device=mat_a.device,
|
||||
)
|
||||
_fp8_blockwise_scaled_mm_custom_op(out, mat_a, mat_b, scales_a, scales_b)
|
||||
return out
|
||||
@@ -15,6 +15,7 @@
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <sgl_kernel/ffi.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
@@ -238,6 +239,21 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) {
|
||||
return RuntimeDeviceCheck(::cudaGetLastError(), location);
|
||||
}
|
||||
|
||||
inline int getSMVersion(int device_id) {
|
||||
int sm_major = 0;
|
||||
int sm_minor = 0;
|
||||
RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
|
||||
RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
|
||||
return sm_major * 10 + sm_minor;
|
||||
}
|
||||
|
||||
inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor {
|
||||
if (required_bytes == 0) return {};
|
||||
DLDataType u8 = {kDLUInt, 8, 1};
|
||||
int64_t shape[] = {static_cast<int64_t>(required_bytes)};
|
||||
return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device);
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Kernel launcher with automatic stream resolution and PDL support.
|
||||
*
|
||||
|
||||
@@ -158,8 +158,9 @@ if _use_aiter:
|
||||
|
||||
|
||||
if _is_cuda:
|
||||
from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm
|
||||
from sgl_kernel import fp8_scaled_mm
|
||||
|
||||
from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
|
||||
from sglang.srt.utils.patch_torch import register_fake_if_exists
|
||||
|
||||
@register_fake_if_exists("sgl_kernel::fp8_scaled_mm")
|
||||
@@ -169,13 +170,6 @@ if _is_cuda:
|
||||
N = mat_b.shape[-1]
|
||||
return mat_a.new_empty((M, N), dtype=out_dtype)
|
||||
|
||||
@register_fake_if_exists("sgl_kernel::fp8_blockwise_scaled_mm")
|
||||
def _fp8_blockwise_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype):
|
||||
# mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N].
|
||||
M = mat_a.shape[-2]
|
||||
N = mat_b.shape[-1]
|
||||
return mat_a.new_empty((M, N), dtype=out_dtype)
|
||||
|
||||
|
||||
use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL")
|
||||
|
||||
@@ -274,11 +268,6 @@ class Fp8GemmRunnerBackend(Enum):
|
||||
FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None
|
||||
|
||||
|
||||
def _check_cutlass_block_fp8_hardware_support() -> bool:
|
||||
"""Return True if CUTLASS block FP8 is supported (Hopper or newer with CUDA 12.0+)."""
|
||||
return is_sm90_supported() or is_blackwell_supported()
|
||||
|
||||
|
||||
if is_blackwell_supported() and is_flashinfer_available():
|
||||
from flashinfer import SfLayout
|
||||
from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8
|
||||
@@ -541,11 +530,10 @@ def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable:
|
||||
return flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback
|
||||
|
||||
elif backend.is_cutlass():
|
||||
if not _check_cutlass_block_fp8_hardware_support():
|
||||
if not is_sm120_supported():
|
||||
raise RuntimeError(
|
||||
"CUTLASS block FP8 requested via --fp8-gemm-backend=cutlass, "
|
||||
"but hardware does not support it. CUTLASS block FP8 requires "
|
||||
"Hopper (SM90+) GPUs with CUDA 12.0+."
|
||||
"--fp8-gemm-backend=cutlass is deprecated on this hardware. "
|
||||
"Please switch to DeepGEMM or FlashInfer TRTLLM on SM90/SM100."
|
||||
)
|
||||
return cutlass_w8a8_block_fp8_linear_with_fallback
|
||||
|
||||
@@ -579,7 +567,7 @@ def _dispatch_auto_backend() -> Callable:
|
||||
# Priority order for auto selection:
|
||||
# 1. DeepGEMM (if enabled and available)
|
||||
# 2. FlashInfer TRTLLM (if Blackwell GPU and FlashInfer available)
|
||||
# 3. CUTLASS (if Hopper+ GPU and CUDA 12.0+)
|
||||
# 3. CUTLASS (if SM120 GPU and CUDA 12.8+)
|
||||
# 4. AITER (if AMD GPU with AITER enabled)
|
||||
# 5. Triton (fallback)
|
||||
|
||||
@@ -587,7 +575,7 @@ def _dispatch_auto_backend() -> Callable:
|
||||
return deepgemm_w8a8_block_fp8_linear_with_fallback
|
||||
elif is_blackwell_supported() and is_flashinfer_available():
|
||||
return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback
|
||||
elif _check_cutlass_block_fp8_hardware_support():
|
||||
elif is_sm120_supported():
|
||||
return cutlass_w8a8_block_fp8_linear_with_fallback
|
||||
elif _use_aiter:
|
||||
return aiter_w8a8_block_fp8_linear
|
||||
@@ -601,8 +589,7 @@ def initialize_fp8_gemm_config(server_args: ServerArgs) -> None:
|
||||
|
||||
backend = server_args.fp8_gemm_runner_backend
|
||||
if backend == "auto" and is_sm120_supported():
|
||||
# TODO(brayden): Verify if CUTLASS can be set by default once SwapAB is supported
|
||||
backend = "triton"
|
||||
backend = "cutlass"
|
||||
|
||||
backend = Fp8GemmRunnerBackend(backend)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user