Skip to content

Migrate QuTLASS kernels to torch Stable ABI #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cleonard530 wants to merge 10 commits into
base: main
Choose a base branch
from
+942 −841
Open

Migrate QuTLASS kernels to torch Stable ABI #12

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
d1a46ae
migrated kernels used in vllm to use the torch stable abi
cleonard530 Jun 23, 2026
a426f34
moved new binding.cpp helper functions to bindings_utils.h to make gi…
cleonard530 Jun 23, 2026
7aa8755
small updates to bindings_utils.h to intent easier to understand.
cleonard530 Jun 24, 2026
7b2601a
very minor formatting fixes
cleonard530 Jul 7, 2026
82163c5
added the functions that I removed from bindings.cpp back, but made t…
cleonard530 Jul 9, 2026
3bb0e69
Guard QUTLASS_MINIMAL_BUILD op defs and restore fusedQuantizeNv for vLLM
cleonard530 Jul 14, 2026
e2257ab
removed temporary fusedQuantizeNv wrapper
cleonard530 Jul 15, 2026
92e4a44
Wire standalone pip build for stable ABI and dispatch via torch.ops._…
cleonard530 Jul 15, 2026
4a8934c
Fix pip editable build: use plain BuildExtension instead of with_options
cleonard530 Jul 15, 2026
9425b22
Remove legacy torch/extension.h from quartet_bwd_sm120.cu
cleonard530 Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
639 changes: 234 additions & 405 deletions qutlass/csrc/bindings.cpp
View file
Open in desktop

Large diffs are not rendered by default.

82 changes: 36 additions & 46 deletions qutlass/csrc/fused_quantize_mx.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,6 @@
* limitations under the License.
*/

#include <ATen/ATen.h>
#include <torch/types.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#ifndef QUTLASS_DISABLE_PYBIND
#include <torch/extension.h>
#endif

#include <iostream>

#include "cutlass/cutlass.h"
Expand Down Expand Up @@ -78,12 +68,12 @@ struct GemmRunner {
GemmRunner() { }

bool run(
torch::Tensor &out,
torch::Tensor &out_sf,
torch::Tensor const&x,
torch::Tensor const&y,
torch::stable::Tensor &out,
torch::stable::Tensor &out_sf,
torch::stable::Tensor const& x,
torch::stable::Tensor const& y,
int32_t M, int32_t N, int32_t K,
torch::Device device)
torch::stable::Device device)
{

using GemmCoord = cutlass::gemm::GemmCoord;
Expand All @@ -93,16 +83,16 @@ struct GemmRunner {
{static_cast<GemmCoord::Index>(M),
static_cast<GemmCoord::Index>(N),
static_cast<GemmCoord::Index>(K)},
{(cutlass::bfloat16_t *)x.data_ptr(), K},
{(cutlass::bfloat16_t *)y.data_ptr(), N},
{(cutlass::float_e2m1_t *)out.data_ptr(), N},
{(cutlass::float_e2m1_t *)out.data_ptr(), N},
{(cutlass::float_ue8m0_t *)out_sf.data_ptr(), M},
{static_cast<const cutlass::bfloat16_t*>(x.const_data_ptr()), K},
{static_cast<const cutlass::bfloat16_t*>(y.const_data_ptr()), N},
{static_cast<cutlass::float_e2m1_t*>(out.mutable_data_ptr()), N},
{static_cast<cutlass::float_e2m1_t*>(out.mutable_data_ptr()), N},
{static_cast<cutlass::float_ue8m0_t*>(out_sf.mutable_data_ptr()), M},
cutlass::bfloat16_t(0) //TODO (later): float
};

const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index());
const torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index());
cudaStream_t stream = get_current_cuda_stream(device.index());

CUTLASS_CHECK(gemmOp.initialize(arguments, nullptr, stream));

Expand All @@ -114,10 +104,10 @@ struct GemmRunner {
};


void fusedQuantizeMxQuest_host(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B)
void fusedQuantizeMxQuest_host(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B)
{
int32_t M = A.numel() / 32;
int32_t N = B.size(1);
Expand All @@ -131,10 +121,10 @@ void fusedQuantizeMxQuest_host(torch::Tensor& D,
bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device());
}

void fusedQuantizeMxAbsMax_host(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B)
void fusedQuantizeMxAbsMax_host(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B)
{
int32_t M = A.numel() / 32;
int32_t N = B.size(1);
Expand All @@ -148,10 +138,10 @@ void fusedQuantizeMxAbsMax_host(torch::Tensor& D,
bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device());
}

void fusedQuantizeMxQuestHad64_host(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B)
void fusedQuantizeMxQuestHad64_host(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B)
{
int32_t M = A.numel() / 64;
int32_t N = B.size(1);
Expand All @@ -165,10 +155,10 @@ void fusedQuantizeMxQuestHad64_host(torch::Tensor& D,
bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device());
}

void fusedQuantizeMxAbsMaxHad64_host(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B)
void fusedQuantizeMxAbsMaxHad64_host(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B)
{
int32_t M = A.numel() / 64;
int32_t N = B.size(1);
Expand All @@ -182,10 +172,10 @@ void fusedQuantizeMxAbsMaxHad64_host(torch::Tensor& D,
bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device());
}

void fusedQuantizeMxQuestHad128_host(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B)
void fusedQuantizeMxQuestHad128_host(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B)
{
int32_t M = A.numel() / 128;
int32_t N = B.size(1);
Expand All @@ -199,10 +189,10 @@ void fusedQuantizeMxQuestHad128_host(torch::Tensor& D,
bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device());
}

void fusedQuantizeMxAbsMaxHad128_host(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B)
void fusedQuantizeMxAbsMaxHad128_host(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B)
{
int32_t M = A.numel() / 128;
int32_t N = B.size(1);
Expand Down
48 changes: 19 additions & 29 deletions qutlass/csrc/fused_quantize_mx_mask.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,6 @@
* limitations under the License.
*/

#include <ATen/ATen.h>
#include <torch/types.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#ifndef QUTLASS_DISABLE_PYBIND
#include <torch/extension.h>
#endif

#include <iostream>

#include "cutlass/cutlass.h"
Expand Down Expand Up @@ -76,13 +66,13 @@ struct GemmRunner {
GemmRunner() { }

bool run(
torch::Tensor &out,
torch::Tensor &out_sf,
torch::Tensor &out_mask,
torch::Tensor const&x,
torch::Tensor const&y,
torch::stable::Tensor &out,
torch::stable::Tensor &out_sf,
torch::stable::Tensor &out_mask,
torch::stable::Tensor const& x,
torch::stable::Tensor const& y,
int32_t M, int32_t N, int32_t K,
torch::Device device)
torch::stable::Device device)
{

using GemmCoord = cutlass::gemm::GemmCoord;
Expand All @@ -92,17 +82,17 @@ struct GemmRunner {
{static_cast<GemmCoord::Index>(M),
static_cast<GemmCoord::Index>(N),
static_cast<GemmCoord::Index>(K)},
{(cutlass::bfloat16_t *)x.data_ptr(), K},
{(cutlass::bfloat16_t *)y.data_ptr(), N},
{(cutlass::float_e2m1_t *)out.data_ptr(), N},
{(cutlass::float_e2m1_t *)out.data_ptr(), N},
{(cutlass::float_ue8m0_t *)out_sf.data_ptr(), M},
{(uint8_t *)out_mask.data_ptr(), M}, //FIXME: bfloat16_t
{static_cast<const cutlass::bfloat16_t*>(x.const_data_ptr()), K},
{static_cast<const cutlass::bfloat16_t*>(y.const_data_ptr()), N},
{static_cast<cutlass::float_e2m1_t*>(out.mutable_data_ptr()), N},
{static_cast<cutlass::float_e2m1_t*>(out.mutable_data_ptr()), N},
{static_cast<cutlass::float_ue8m0_t*>(out_sf.mutable_data_ptr()), M},
{static_cast<uint8_t*>(out_mask.mutable_data_ptr()), M}, //FIXME: bfloat16_t
cutlass::bfloat16_t(0) //TODO (later): float
};

const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index());
const torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index());
cudaStream_t stream = get_current_cuda_stream(device.index());

CUTLASS_CHECK(gemmOp.initialize(arguments, nullptr, stream));

Expand All @@ -113,11 +103,11 @@ struct GemmRunner {

};

void fusedQuantizeMxQuestWithMask_host(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor& D_mask,
torch::Tensor const& A,
torch::Tensor const& B)
void fusedQuantizeMxQuestWithMask_host(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor& D_mask,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B)
{
int32_t M = A.numel() / 32;
int32_t N = B.size(1);
Expand Down
58 changes: 24 additions & 34 deletions qutlass/csrc/fused_quantize_mx_sm100.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,6 @@
* limitations under the License.
*/

#include <ATen/ATen.h>
#include <torch/types.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#ifndef QUTLASS_DISABLE_PYBIND
#include <torch/extension.h>
#endif

#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
Expand Down Expand Up @@ -128,11 +118,11 @@ constexpr bool IsBlockScaleSupported = FusionOp::IsBlockScaleSupported;
using SfdOutputCfg = cutlass::detail::Sm1xxBlockScaledOutputConfig<OutputSFVectorSize>;
using LayoutSFD = typename SfdOutputCfg::LayoutSF;

typename Gemm::Arguments args_from_options(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B,
torch::Tensor const& global_scale,
typename Gemm::Arguments args_from_options(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B,
torch::stable::Tensor const& global_scale,
int32_t M, int32_t N, int32_t K)
{
using ElementA = typename Gemm::ElementA;
Expand All @@ -155,30 +145,30 @@ typename Gemm::Arguments args_from_options(torch::Tensor& D,
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{
static_cast<ElementA const*>(A.data_ptr()), stride_A,
static_cast<ElementB const*>(B.data_ptr()), stride_B},
static_cast<ElementA const*>(A.const_data_ptr()), stride_A,
static_cast<ElementB const*>(B.const_data_ptr()), stride_B},
{
{1.f, 0.f},
nullptr, stride_C,
static_cast<ElementD*>(D.data_ptr()), stride_D
static_cast<ElementD*>(D.mutable_data_ptr()), stride_D
}
};

if constexpr (IsBlockScaleSupported) {
arguments.epilogue.thread.block_scale_factor_ptr = static_cast<cutlass::float_ue8m0_t*>(D_sf.data_ptr());
arguments.epilogue.thread.norm_constant_ptr = static_cast<ElementAccumulator const*>(global_scale.data_ptr());;
arguments.epilogue.thread.block_scale_factor_ptr = static_cast<cutlass::float_ue8m0_t*>(D_sf.mutable_data_ptr());
arguments.epilogue.thread.norm_constant_ptr = static_cast<ElementAccumulator const*>(global_scale.const_data_ptr());
}

return arguments;
}

void runGemm(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B,
torch::Tensor const& global_scale,
void runGemm(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B,
torch::stable::Tensor const& global_scale,
int32_t M, int32_t N, int32_t K,
torch::Device device)
torch::stable::Device device)
{
Gemm gemm;

Expand All @@ -189,8 +179,8 @@ void runGemm(torch::Tensor& D,

cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);

const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index());
const torch::stable::accelerator::DeviceGuard device_guard(A.get_device_index());
cudaStream_t stream = get_current_cuda_stream(device.index());

CUTLASS_CHECK(gemm.can_implement(arguments));

Expand All @@ -199,11 +189,11 @@ void runGemm(torch::Tensor& D,
CUTLASS_CHECK(gemm.run(arguments, workspace.get(), stream));
}

void fusedQuantizeMxAbsMax_host_sm100(torch::Tensor& D,
torch::Tensor& D_sf,
torch::Tensor const& A,
torch::Tensor const& B,
torch::Tensor const& global_scale)
void fusedQuantizeMxAbsMax_host_sm100(torch::stable::Tensor& D,
torch::stable::Tensor& D_sf,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B,
torch::stable::Tensor const& global_scale)
{
#if TARGET_CUDA_ARCH == 100
int32_t M = A.numel() / 128;
Expand All @@ -212,7 +202,7 @@ void fusedQuantizeMxAbsMax_host_sm100(torch::Tensor& D,

runGemm(D, D_sf, A, B, global_scale, M, N, K, A.device());
#else
TORCH_CHECK(false, "Unsupported CUDA arch");
STD_TORCH_CHECK(false, "Unsupported CUDA arch");
#endif
}

Expand Down
Loading