From 41148496b066235114828bfb4bd476c0c59b68d9 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Tue, 30 Jun 2026 11:57:17 -0700 Subject: [PATCH 1/3] pybind->TORCH_LIBRARY; get ABI stable with CPython --- csrc/api/api.cpp | 23 ++++++++++++------- csrc/api/common.h | 2 +- csrc/api/dense_decode.h | 10 ++++---- csrc/api/sparse_decode.h | 8 +++---- csrc/api/sparse_fwd.h | 4 ++-- .../kerutils/supplemental/torch_tensors.h | 2 +- csrc/sm100/prefill/dense/common/utils.hpp | 1 - .../prefill/dense/fmha_cutlass_bwd_sm100.cu | 2 +- .../prefill/dense/fmha_cutlass_fwd_sm100.cu | 4 ++-- csrc/sm100/prefill/dense/interface.h | 4 ++-- flash_mla/__init__.py | 11 +++++++++ flash_mla/flash_mla_interface.py | 4 +++- setup.py | 4 ++++ 13 files changed, 51 insertions(+), 28 deletions(-) diff --git a/csrc/api/api.cpp b/csrc/api/api.cpp index f43f2a09..ca0d2269 100644 --- a/csrc/api/api.cpp +++ b/csrc/api/api.cpp @@ -1,15 +1,22 @@ -#include +#include #include "sparse_fwd.h" #include "sparse_decode.h" #include "dense_decode.h" #include "dense_fwd.h" -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "FlashMLA"; - m.def("sparse_decode_fwd", &sparse_attn_decode_interface); - m.def("dense_decode_fwd", &dense_attn_decode_interface); - m.def("sparse_prefill_fwd", &sparse_attn_prefill_interface); - m.def("dense_prefill_fwd", &FMHACutlassSM100FwdRun); - m.def("dense_prefill_bwd", &FMHACutlassSM100BwdRun); +TORCH_LIBRARY(flash_mla, m) { + m.def("sparse_decode_fwd(Tensor q, Tensor kv, Tensor indices, Tensor? topk_length, Tensor? attn_sink, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits, Tensor? extra_kv, Tensor? extra_indices, Tensor? extra_topk_length, int d_v, float sm_scale) -> (Tensor, Tensor, Tensor(a)?, Tensor(b)?)"); + m.def("dense_decode_fwd(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float softmax_scale, bool is_causal, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits) -> (Tensor, Tensor, Tensor(a)?, Tensor(b)?)"); + m.def("sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length) -> Tensor[]"); + m.def("dense_prefill_fwd(Tensor workspace_buffer, Tensor q, Tensor k, Tensor v, Tensor cumulative_seqlen_q, Tensor cumulative_seqlen_kv, Tensor(a!) o, Tensor(b!) lse, int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) -> ()"); + m.def("dense_prefill_bwd(Tensor(a!) workspace_buffer, Tensor d_o, Tensor q, Tensor k, Tensor v, Tensor o, Tensor lse, Tensor cumulative_seqlen_q, Tensor cumulative_seqlen_kv, Tensor(b!) dq, Tensor(c!) dk, Tensor(d!) dv, int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) -> ()"); +} + +TORCH_LIBRARY_IMPL(flash_mla, CUDA, m) { + m.impl("sparse_decode_fwd", &sparse_attn_decode_interface); + m.impl("dense_decode_fwd", &dense_attn_decode_interface); + m.impl("sparse_prefill_fwd", &sparse_attn_prefill_interface); + m.impl("dense_prefill_fwd", &FMHACutlassSM100FwdRun); + m.impl("dense_prefill_bwd", &FMHACutlassSM100BwdRun); } diff --git a/csrc/api/common.h b/csrc/api/common.h index 2c930ed9..d35328d7 100644 --- a/csrc/api/common.h +++ b/csrc/api/common.h @@ -2,7 +2,7 @@ #include -#include +#include #include #include #include diff --git a/csrc/api/dense_decode.h b/csrc/api/dense_decode.h index 7df178a6..7877bf02 100644 --- a/csrc/api/dense_decode.h +++ b/csrc/api/dense_decode.h @@ -12,15 +12,15 @@ static std::tuple, std::optional> dense_attn_decode_interface( - at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size + at::Tensor q, // batch_size x seqlen_q x num_heads x head_size const at::Tensor &kcache, // num_blocks x page_block_size x num_heads_k x head_size (when is_fp8 is False) or num_blocks x num_heads_k x (page_block_size*656) (when is_fp8 is True) - const int head_size_v, + const int64_t head_size_v, const at::Tensor &seqlens_k, // batch_size const at::Tensor &block_table, // batch_size x max_num_blocks_per_seq - const float softmax_scale, + const double softmax_scale, bool is_causal, - std::optional &tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) - std::optional &num_splits // batch_size + 1 + std::optional tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) + std::optional num_splits // batch_size + 1 ) { // Check arch Arch arch = Arch(); diff --git a/csrc/api/sparse_decode.h b/csrc/api/sparse_decode.h index 6df5c841..adef4834 100644 --- a/csrc/api/sparse_decode.h +++ b/csrc/api/sparse_decode.h @@ -187,13 +187,13 @@ sparse_attn_decode_interface( const at::Tensor &indices, // [b, s_q, topk] const std::optional &topk_length, // [b, s_q] const std::optional &attn_sink, // [h_q] - std::optional &tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) - std::optional &num_splits, // batch_size + 1 + std::optional tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) + std::optional num_splits, // batch_size + 1 const std::optional &extra_kv, const std::optional &extra_indices, const std::optional &extra_topk_length, - int d_v, - float sm_scale + int64_t d_v, + double sm_scale ) { using bf16 = cutlass::bfloat16_t; diff --git a/csrc/api/sparse_fwd.h b/csrc/api/sparse_fwd.h index 66d7111e..5157d1cc 100644 --- a/csrc/api/sparse_fwd.h +++ b/csrc/api/sparse_fwd.h @@ -102,8 +102,8 @@ static std::vector sparse_attn_prefill_interface( const at::Tensor &q, const at::Tensor &kv, const at::Tensor &indices, - float sm_scale, - int d_v, + double sm_scale, + int64_t d_v, const std::optional &attn_sink, const std::optional &topk_length ) { diff --git a/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h b/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h index 5b4e564c..4d798bfb 100644 --- a/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h +++ b/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h @@ -2,7 +2,7 @@ #include -#include +#include // libtorch C++ API without pybind (Py_LIMITED_API-safe) #include "kerutils/common/common.h" diff --git a/csrc/sm100/prefill/dense/common/utils.hpp b/csrc/sm100/prefill/dense/common/utils.hpp index fdaeff08..1231f5fe 100644 --- a/csrc/sm100/prefill/dense/common/utils.hpp +++ b/csrc/sm100/prefill/dense/common/utils.hpp @@ -1,6 +1,5 @@ #pragma once -#include #include "cutlass/numeric_types.h" #include "helper.h" diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu index 54d85db8..b67ea9e4 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu +++ b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu @@ -30,7 +30,7 @@ void FMHACutlassSM100BwdRun(at::Tensor workspace_buffer, at::Tensor d_o, at::Ten at::Tensor v, at::Tensor o, at::Tensor lse, at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, at::Tensor dq, at::Tensor dk, at::Tensor dv, - int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) { + int64_t mask_mode_code, double softmax_scale, int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_varlen) { const c10::cuda::OptionalCUDAGuard device_guard(q.device()); diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu index ab66f0fd..6a4ef4ff 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu +++ b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu @@ -31,8 +31,8 @@ void call_run_fmha_fwd([[maybe_unused]] Mask mask, [[maybe_unused]] Varlen is_va void FMHACutlassSM100FwdRun(at::Tensor workspace_buffer, at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, at::Tensor o, at::Tensor lse, - int mask_mode_code, float sm_scale, int max_seqlen_q, - int max_seqlen_kv, bool is_varlen) { + int64_t mask_mode_code, double sm_scale, int64_t max_seqlen_q, + int64_t max_seqlen_kv, bool is_varlen) { const c10::cuda::OptionalCUDAGuard device_guard(q.device()); CHECK(q.scalar_type() == k.scalar_type()); auto scalar_type_in = q.scalar_type(); diff --git a/csrc/sm100/prefill/dense/interface.h b/csrc/sm100/prefill/dense/interface.h index 80ef2bca..0ddc553a 100644 --- a/csrc/sm100/prefill/dense/interface.h +++ b/csrc/sm100/prefill/dense/interface.h @@ -5,10 +5,10 @@ void FMHACutlassSM100FwdRun(at::Tensor workspace_buffer, at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, at::Tensor o, at::Tensor lse, - int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen); + int64_t mask_mode_code, double softmax_scale, int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_varlen); void FMHACutlassSM100BwdRun(at::Tensor workspace_buffer, at::Tensor d_o, at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor o, at::Tensor lse, at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, at::Tensor dq, at::Tensor dk, at::Tensor dv, - int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen); + int64_t mask_mode_code, double softmax_scale, int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_varlen); diff --git a/flash_mla/__init__.py b/flash_mla/__init__.py index 02a8bbab..41837732 100644 --- a/flash_mla/__init__.py +++ b/flash_mla/__init__.py @@ -1,5 +1,16 @@ __version__ = "1.0.0" +import torch +from pathlib import Path + +# Load the compiled extension so its static TORCH_LIBRARY registrations run. +# The operators are then available as torch.ops.flash_mla.. The glob matches +# both the Python-version-specific name (cuda.cpython-*.so) and the abi3 name +# (cuda.abi3.so) produced when Py_LIMITED_API is enabled. +_so_files = list(Path(__file__).parent.glob("cuda*.so")) +assert len(_so_files) == 1, f"Expected one cuda*.so file, found {_so_files}" +torch.ops.load_library(_so_files[0]) + from flash_mla.flash_mla_interface import ( get_mla_metadata, flash_mla_with_kvcache, diff --git a/flash_mla/flash_mla_interface.py b/flash_mla/flash_mla_interface.py index a3740b0f..147bc4a2 100644 --- a/flash_mla/flash_mla_interface.py +++ b/flash_mla/flash_mla_interface.py @@ -3,7 +3,9 @@ import torch -import flash_mla.cuda as flash_mla_cuda +# torch.ops.flash_mla is the op namespace registered by the _C extension (loaded in +# __init__.py). Alias it to the old pybind module name so the call sites below are unchanged. +flash_mla_cuda = torch.ops.flash_mla @dataclasses.dataclass class FlashMLASchedMeta: diff --git a/setup.py b/setup.py index 513b4355..2e4a8a0b 100644 --- a/setup.py +++ b/setup.py @@ -130,6 +130,9 @@ def get_nvcc_thread_args(): Path(this_dir) / "csrc" / "cutlass" / "include", Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", ], + # Build against CPython's Limited API (abi3) so one wheel works across + # multiple CPython versions, which is possible now that pybind11 is gone + py_limited_api=True, ) ) @@ -148,4 +151,5 @@ def get_nvcc_thread_args(): packages=find_packages(include=['flash_mla']), ext_modules=ext_modules, cmdclass={"build_ext": BuildExtension}, + options={"bdist_wheel": {"py_limited_api": "cp310"}}, ) From 3911d005cc61172d6a11194ec2c26d68884a021e Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Thu, 2 Jul 2026 08:24:36 -0700 Subject: [PATCH 2/3] Make FlashMLA ABI stable with torch --- csrc/api/api.cpp | 16 +-- csrc/api/common.h | 66 ++++++--- csrc/api/dense_decode.h | 131 +++++++++--------- csrc/api/sparse_decode.h | 119 ++++++++-------- csrc/api/sparse_fwd.h | 49 ++++--- .../kerutils/supplemental/cuda_stream.h | 19 +++ .../kerutils/supplemental/torch_tensors.h | 29 ++-- .../prefill/dense/fmha_cutlass_bwd_sm100.cu | 25 ++-- .../prefill/dense/fmha_cutlass_bwd_sm100.cuh | 48 ++++--- .../prefill/dense/fmha_cutlass_fwd_sm100.cu | 25 ++-- .../prefill/dense/fmha_cutlass_fwd_sm100.cuh | 39 +++--- csrc/sm100/prefill/dense/interface.h | 16 +-- setup.py | 8 +- 13 files changed, 331 insertions(+), 259 deletions(-) create mode 100644 csrc/kerutils/include/kerutils/supplemental/cuda_stream.h diff --git a/csrc/api/api.cpp b/csrc/api/api.cpp index ca0d2269..918c639b 100644 --- a/csrc/api/api.cpp +++ b/csrc/api/api.cpp @@ -1,11 +1,11 @@ -#include +#include #include "sparse_fwd.h" #include "sparse_decode.h" #include "dense_decode.h" #include "dense_fwd.h" -TORCH_LIBRARY(flash_mla, m) { +STABLE_TORCH_LIBRARY(flash_mla, m) { m.def("sparse_decode_fwd(Tensor q, Tensor kv, Tensor indices, Tensor? topk_length, Tensor? attn_sink, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits, Tensor? extra_kv, Tensor? extra_indices, Tensor? extra_topk_length, int d_v, float sm_scale) -> (Tensor, Tensor, Tensor(a)?, Tensor(b)?)"); m.def("dense_decode_fwd(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float softmax_scale, bool is_causal, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits) -> (Tensor, Tensor, Tensor(a)?, Tensor(b)?)"); m.def("sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length) -> Tensor[]"); @@ -13,10 +13,10 @@ TORCH_LIBRARY(flash_mla, m) { m.def("dense_prefill_bwd(Tensor(a!) workspace_buffer, Tensor d_o, Tensor q, Tensor k, Tensor v, Tensor o, Tensor lse, Tensor cumulative_seqlen_q, Tensor cumulative_seqlen_kv, Tensor(b!) dq, Tensor(c!) dk, Tensor(d!) dv, int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) -> ()"); } -TORCH_LIBRARY_IMPL(flash_mla, CUDA, m) { - m.impl("sparse_decode_fwd", &sparse_attn_decode_interface); - m.impl("dense_decode_fwd", &dense_attn_decode_interface); - m.impl("sparse_prefill_fwd", &sparse_attn_prefill_interface); - m.impl("dense_prefill_fwd", &FMHACutlassSM100FwdRun); - m.impl("dense_prefill_bwd", &FMHACutlassSM100BwdRun); +STABLE_TORCH_LIBRARY_IMPL(flash_mla, CUDA, m) { + m.impl("sparse_decode_fwd", TORCH_BOX(&sparse_attn_decode_interface)); + m.impl("dense_decode_fwd", TORCH_BOX(&dense_attn_decode_interface)); + m.impl("sparse_prefill_fwd", TORCH_BOX(&sparse_attn_prefill_interface)); + m.impl("dense_prefill_fwd", TORCH_BOX(&FMHACutlassSM100FwdRun)); + m.impl("dense_prefill_bwd", TORCH_BOX(&FMHACutlassSM100BwdRun)); } diff --git a/csrc/api/common.h b/csrc/api/common.h index d35328d7..07520738 100644 --- a/csrc/api/common.h +++ b/csrc/api/common.h @@ -1,20 +1,51 @@ #pragma once #include +#include + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include -#include -#include -#include #include +#include #include +using torch::stable::Tensor; +using torch::headeronly::ScalarType; + +// Re-exported from kerutils so existing call sites can use it unqualified. +using kerutils::get_current_cuda_stream; + static constexpr float LOG_2_E = 1.44269504f; -// Instantiation for tensor.data_ptr() -template<> -inline cutlass::bfloat16_t* at::TensorBase::data_ptr() const { - return reinterpret_cast(this->data_ptr()); +// Helper to access device properties (SM count, compute capability) via the CUDA +// Runtime directly and cache per device index. +inline const cudaDeviceProp &get_cached_device_prop() { + int device_idx = static_cast(torch::stable::accelerator::getCurrentDeviceIndex()); + constexpr int kMaxDevices = 16; + static cudaDeviceProp props[kMaxDevices]; + static bool inited[kMaxDevices] = {false}; + if (device_idx < 0 || device_idx >= kMaxDevices) { + // Uncached fallback for unexpectedly large device indices. + static thread_local cudaDeviceProp tmp; + cudaGetDeviceProperties(&tmp, device_idx); + return tmp; + } + if (!inited[device_idx]) { + cudaGetDeviceProperties(&props[device_idx], device_idx); + inited[device_idx] = true; + } + return props[device_idx]; } // A struct that holds the architecture information of the current GPU. @@ -22,10 +53,10 @@ struct Arch { int major; int minor; int num_sms; - cudaDeviceProp* device_prop; + const cudaDeviceProp* device_prop; Arch() { - device_prop = at::cuda::getCurrentDeviceProperties(); + device_prop = &get_cached_device_prop(); major = device_prop->major; minor = device_prop->minor; num_sms = device_prop->multiProcessorCount; @@ -43,7 +74,7 @@ struct Arch { // Convert int64_t stride to int32_t, with overflow check. inline int int64_stride_to_int(int64_t orig_stride) { if (orig_stride > std::numeric_limits::max()) { - TORCH_CHECK(false, "[FlashMLA] Stride exceeds int32 limit: ", orig_stride); + STD_TORCH_CHECK(false, "[FlashMLA] Stride exceeds int32 limit: ", orig_stride); } return static_cast(orig_stride); } @@ -57,7 +88,7 @@ inline int int64_stride_to_int(int64_t orig_stride) { static constexpr int CONSTEXPR_NAME = 64; \ return __VA_ARGS__(); \ } else { \ - TORCH_CHECK(false, "Unsupported num_heads_q: ", NUM_HEADS); \ + STD_TORCH_CHECK(false, "Unsupported num_heads_q: ", NUM_HEADS); \ } \ } (); @@ -70,7 +101,7 @@ inline int int64_stride_to_int(int64_t orig_stride) { static constexpr int CONSTEXPR_NAME = 512; \ return __VA_ARGS__(); \ } else { \ - TORCH_CHECK(false, "Unsupported head_dim_qk: ", HEAD_DIM); \ + STD_TORCH_CHECK(false, "Unsupported head_dim_qk: ", HEAD_DIM); \ } \ } (); @@ -94,7 +125,7 @@ inline int int64_stride_to_int(int64_t orig_stride) { static constexpr ModelType CONSTEXPR_NAME = ModelType::MODEL1; \ return __VA_ARGS__(); \ } else { \ - TORCH_CHECK(false, "Unsupported model type: ", (int)MODEL_TYPE); \ + STD_TORCH_CHECK(false, "Unsupported model type: ", (int)MODEL_TYPE); \ } \ } (); @@ -120,7 +151,7 @@ constexpr auto get_static_enum_name(){ }; } -template +template static constexpr std::size_t get_enum_max(){ constexpr T value = static_cast(N); if constexpr (get_static_enum_name().find(")") == std::string_view::npos) @@ -133,8 +164,8 @@ template requires std::is_enum_v static constexpr std::string get_dynamic_enum_name(T value){ constexpr std::size_t num = get_enum_max(); constexpr auto names = [](std::index_sequence){ - return std::array{ - get_static_enum_name(Is)>()... + return std::array{ + get_static_enum_name(Is)>()... }; }(std::make_index_sequence{}); return (std::string)names[static_cast(value)]; @@ -219,7 +250,7 @@ class ImplBase { Arch cur_gpu_arch = Arch(); fprintf(stderr, "Current GPU: %s, SM %d.%d with %d SMs\n", cur_gpu_arch.device_prop->name, cur_gpu_arch.major, cur_gpu_arch.minor, cur_gpu_arch.num_sms); fprintf(stderr, "This means that the dispatcher has chosen an implementation that does not support all required features. Maybe there is a bug in the dispatcher, or you have requested an invalid combination of features.\n"); - TORCH_CHECK(false, "The chosen implementation does not support all required features. See message above for details."); + STD_TORCH_CHECK(false, "The chosen implementation does not support all required features. See message above for details."); } } @@ -228,4 +259,3 @@ class ImplBase { run_(params, required_features); } }; - diff --git a/csrc/api/dense_decode.h b/csrc/api/dense_decode.h index 7877bf02..db8770cc 100644 --- a/csrc/api/dense_decode.h +++ b/csrc/api/dense_decode.h @@ -10,31 +10,31 @@ #include "smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h" #include "smxx/decode/combine/combine.h" -static std::tuple, std::optional> +static std::tuple, std::optional> dense_attn_decode_interface( - at::Tensor q, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor &kcache, // num_blocks x page_block_size x num_heads_k x head_size (when is_fp8 is False) or num_blocks x num_heads_k x (page_block_size*656) (when is_fp8 is True) + Tensor q, // batch_size x seqlen_q x num_heads x head_size + const Tensor &kcache, // num_blocks x page_block_size x num_heads_k x head_size (when is_fp8 is False) or num_blocks x num_heads_k x (page_block_size*656) (when is_fp8 is True) const int64_t head_size_v, - const at::Tensor &seqlens_k, // batch_size - const at::Tensor &block_table, // batch_size x max_num_blocks_per_seq + const Tensor &seqlens_k, // batch_size + const Tensor &block_table, // batch_size x max_num_blocks_per_seq const double softmax_scale, bool is_causal, - std::optional tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) - std::optional num_splits // batch_size + 1 + std::optional tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) + std::optional num_splits // batch_size + 1 ) { // Check arch Arch arch = Arch(); if (!arch.is_sm90a()) { - TORCH_CHECK(false, "Dense decode MLA is only supported on SM90a architecture"); + STD_TORCH_CHECK(false, "Dense decode MLA is only supported on SM90a architecture"); } // Check data types - auto q_dtype = q.dtype(); - TORCH_CHECK(q_dtype == torch::kBFloat16 || q_dtype == torch::kHalf); - - TORCH_CHECK(kcache.dtype() == q_dtype, "query and key must have the same dtype"); - TORCH_CHECK(seqlens_k.dtype() == torch::kInt32, "seqlens_k must have dtype int32"); - TORCH_CHECK(block_table.dtype() == torch::kInt32, "block_table must have dtype torch.int32"); + auto q_dtype = q.scalar_type(); + STD_TORCH_CHECK(q_dtype == ScalarType::BFloat16 || q_dtype == ScalarType::Half); + + STD_TORCH_CHECK(kcache.scalar_type() == q_dtype, "query and key must have the same dtype"); + STD_TORCH_CHECK(seqlens_k.scalar_type() == ScalarType::Int, "seqlens_k must have dtype int32"); + STD_TORCH_CHECK(block_table.scalar_type() == ScalarType::Int, "block_table must have dtype torch.int32"); // Check device KU_CHECK_DEVICE(q); @@ -45,36 +45,38 @@ dense_attn_decode_interface( KU_CHECK_DEVICE(num_splits); // Check layout - TORCH_CHECK(q.stride(-1) == 1, "q must have contiguous last dimension"); - TORCH_CHECK(kcache.stride(-1) == 1, "kcache must have contiguous last dimension"); + STD_TORCH_CHECK(q.stride(-1) == 1, "q must have contiguous last dimension"); + STD_TORCH_CHECK(kcache.stride(-1) == 1, "kcache must have contiguous last dimension"); KU_CHECK_CONTIGUOUS(seqlens_k); - TORCH_CHECK(block_table.stride(-1) == 1, "block_table must have contiguous last dimension"); + STD_TORCH_CHECK(block_table.stride(-1) == 1, "block_table must have contiguous last dimension"); KU_CHECK_CONTIGUOUS(tile_scheduler_metadata); KU_CHECK_CONTIGUOUS(num_splits); - const auto sizes = q.sizes(); - const int batch_size = sizes[0]; - const int seqlen_q_ori = sizes[1]; - const int num_heads_q = sizes[2]; - const int head_size_k = sizes[3]; - TORCH_CHECK(head_size_k == 576 || head_size_k == 512, "Only head_size_k == 576 or 512 is supported"); - TORCH_CHECK(head_size_v == 512, "Only head_size_v == 576 is supported"); - + const int batch_size = q.size(0); + const int seqlen_q_ori = q.size(1); + const int num_heads_q = q.size(2); + const int head_size_k = q.size(3); + STD_TORCH_CHECK(head_size_k == 576 || head_size_k == 512, "Only head_size_k == 576 or 512 is supported"); + STD_TORCH_CHECK(head_size_v == 512, "Only head_size_v == 576 is supported"); + const int max_num_blocks_per_seq = block_table.size(1); const int num_blocks = kcache.size(0); const int page_block_size = kcache.size(1); const int num_heads_k = kcache.size(2); - TORCH_CHECK(page_block_size == 64, "Currently page_block_size must be 64"); - TORCH_CHECK(batch_size > 0, "batch size must be positive"); - TORCH_CHECK(num_heads_q % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query"); + STD_TORCH_CHECK(page_block_size == 64, "Currently page_block_size must be 64"); + STD_TORCH_CHECK(batch_size > 0, "batch size must be positive"); + STD_TORCH_CHECK(num_heads_q % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query"); if (seqlen_q_ori == 1) { is_causal = false; } const int num_q_heads_per_hk = num_heads_q / num_heads_k; const int q_seq_per_hk = seqlen_q_ori * num_q_heads_per_hk; const int num_heads = num_heads_k; - q = q.view({batch_size, seqlen_q_ori, num_heads_k, num_q_heads_per_hk, head_size_k}).transpose(2, 3) - .reshape({batch_size, q_seq_per_hk, num_heads, head_size_k}); + q = torch::stable::reshape( + torch::stable::transpose( + torch::stable::view(q, {batch_size, seqlen_q_ori, num_heads_k, num_q_heads_per_hk, head_size_k}), + 2, 3), + {batch_size, q_seq_per_hk, num_heads, head_size_k}); int num_sm_parts = std::max(arch.num_sms / num_heads_k / cutlass::ceil_div(seqlen_q_ori*num_heads_q/num_heads_k, 64), 1); KU_CHECK_SHAPE(q, batch_size, q_seq_per_hk, num_heads, head_size_k); @@ -84,17 +86,16 @@ dense_attn_decode_interface( KU_CHECK_SHAPE(tile_scheduler_metadata, num_sm_parts, DecodingSchedMetaSize/sizeof(int)); KU_CHECK_SHAPE(num_splits, batch_size+1); - at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index()); - auto opts = q.options(); - at::Tensor out = torch::empty({batch_size, num_heads, q_seq_per_hk, head_size_v}, opts); - at::Tensor lse = torch::empty({batch_size, num_heads, q_seq_per_hk}, opts.dtype(at::kFloat)); + Tensor out = torch::stable::new_empty(q, {batch_size, num_heads, q_seq_per_hk, head_size_v}); + Tensor lse = torch::stable::new_empty(q, {batch_size, num_heads, q_seq_per_hk}, ScalarType::Float); KU_CHECK_CONTIGUOUS(out); KU_CHECK_CONTIGUOUS(lse); if (!tile_scheduler_metadata.has_value()) { - tile_scheduler_metadata = torch::empty({num_sm_parts, sizeof(DecodingSchedMeta)/4}, opts.dtype(torch::kInt32)); - num_splits = torch::empty({batch_size+1}, opts.dtype(torch::kInt32)); + tile_scheduler_metadata = torch::stable::new_empty(q, {num_sm_parts, sizeof(DecodingSchedMeta)/4}, ScalarType::Int); + num_splits = torch::stable::new_empty(q, {batch_size+1}, ScalarType::Int); KU_CHECK_CONTIGUOUS(tile_scheduler_metadata); KU_CHECK_CONTIGUOUS(num_splits); @@ -104,16 +105,16 @@ dense_attn_decode_interface( 5, -1, -1, nullptr, nullptr, - seqlens_k.data_ptr(), + const_cast(seqlens_k.const_data_ptr()), (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr(), - num_splits->data_ptr(), + num_splits->mutable_data_ptr(), num_sm_parts, - at::cuda::getCurrentCUDAStream().stream() + get_current_cuda_stream(q) }; smxx::decode::run_get_decoding_sched_meta_kernel(get_sched_meta_params); } else { - KU_CHECK_DTYPE(tile_scheduler_metadata, torch::kInt32); - KU_CHECK_DTYPE(num_splits, torch::kInt32); + KU_CHECK_DTYPE(tile_scheduler_metadata, ScalarType::Int); + KU_CHECK_DTYPE(num_splits, ScalarType::Int); KU_CHECK_DEVICE(tile_scheduler_metadata); KU_CHECK_DEVICE(num_splits); KU_CHECK_CONTIGUOUS(tile_scheduler_metadata); @@ -127,7 +128,7 @@ dense_attn_decode_interface( params.b = batch_size; params.s_q = seqlen_q_ori; params.q_seq_per_hk = q_seq_per_hk; - params.seqlens_k_ptr = seqlens_k.data_ptr(); + params.seqlens_k_ptr = const_cast(seqlens_k.const_data_ptr()); params.h_q = num_heads_q; params.h_k = num_heads_k; params.num_blocks = num_blocks; @@ -141,7 +142,7 @@ dense_attn_decode_interface( params.q_ptr = q.data_ptr(); params.k_ptr = kcache.data_ptr(); params.o_ptr = out.data_ptr(); - params.softmax_lse_ptr = lse.data_ptr(); + params.softmax_lse_ptr = lse.mutable_data_ptr(); // All stride are in elements, not bytes. params.q_batch_stride = q.stride(0); params.k_batch_stride = kcache.stride(0); @@ -153,35 +154,35 @@ dense_attn_decode_interface( params.k_head_stride = kcache.stride(2); params.o_head_stride = out.stride(1); - params.block_table = block_table.data_ptr(); + params.block_table = const_cast(block_table.const_data_ptr()); params.block_table_batch_stride = block_table.stride(0); params.page_block_size = page_block_size; params.tile_scheduler_metadata_ptr = (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr(); params.num_sm_parts = num_sm_parts; - params.num_splits_ptr = num_splits->data_ptr(); + params.num_splits_ptr = num_splits->mutable_data_ptr(); const int total_num_splits = batch_size + params.num_sm_parts; - at::Tensor lse_accum = torch::empty({total_num_splits, num_heads, q_seq_per_hk}, opts.dtype(at::kFloat)); - at::Tensor out_accum = torch::empty({total_num_splits, num_heads, q_seq_per_hk, head_size_v}, opts.dtype(at::kFloat)); + Tensor lse_accum = torch::stable::new_empty(q, {total_num_splits, num_heads, q_seq_per_hk}, ScalarType::Float); + Tensor out_accum = torch::stable::new_empty(q, {total_num_splits, num_heads, q_seq_per_hk, head_size_v}, ScalarType::Float); KU_CHECK_CONTIGUOUS(lse_accum); KU_CHECK_CONTIGUOUS(out_accum); params.total_num_splits = total_num_splits; - params.softmax_lseaccum_ptr = lse_accum.data_ptr(); - params.oaccum_ptr = out_accum.data_ptr(); + params.softmax_lseaccum_ptr = lse_accum.mutable_data_ptr(); + params.oaccum_ptr = out_accum.mutable_data_ptr(); - params.stream = at::cuda::getCurrentCUDAStream().stream(); + params.stream = get_current_cuda_stream(q); - if (q_dtype == torch::kBFloat16) { + if (q_dtype == ScalarType::BFloat16) { sm90::run_flash_splitkv_mla_kernel(params); - } else if (q_dtype == torch::kHalf) { + } else if (q_dtype == ScalarType::Half) { #ifdef FLASH_MLA_DISABLE_FP16 - TORCH_CHECK(false, "FlashMLA is compiled with -DFLASH_MLA_DISABLE_FP16. Please remove this flag from your environment and re-compile FlashMLA."); + STD_TORCH_CHECK(false, "FlashMLA is compiled with -DFLASH_MLA_DISABLE_FP16. Please remove this flag from your environment and re-compile FlashMLA."); #else sm90::run_flash_splitkv_mla_kernel(params); #endif } else { - TORCH_CHECK(false, "Unsupported dtype for dense MLA on SM90"); + STD_TORCH_CHECK(false, "Unsupported dtype for dense MLA on SM90"); } CombineParams combine_params = { @@ -203,23 +204,29 @@ dense_attn_decode_interface( params.num_sm_parts, nullptr, - at::cuda::getCurrentCUDAStream().stream() + get_current_cuda_stream(q) }; - if (q_dtype == torch::kBFloat16) { + if (q_dtype == ScalarType::BFloat16) { smxx::decode::run_flash_mla_combine_kernel(combine_params); - } else if (q_dtype == torch::kHalf) { + } else if (q_dtype == ScalarType::Half) { #ifndef FLASH_MLA_DISABLE_FP16 smxx::decode::run_flash_mla_combine_kernel(combine_params); #endif } else { - TORCH_CHECK(false, "Unsupported tensor dtype for query"); + STD_TORCH_CHECK(false, "Unsupported tensor dtype for query"); } - out = out.view({batch_size, num_heads_k, seqlen_q_ori, num_q_heads_per_hk, head_size_v}).transpose(1, 2) - .reshape({batch_size, seqlen_q_ori, num_heads_q, head_size_v}); - lse = lse.view({batch_size, num_heads_k, seqlen_q_ori, num_q_heads_per_hk}).transpose(2, 3) - .reshape({batch_size, num_heads_q, seqlen_q_ori}); + out = torch::stable::reshape( + torch::stable::transpose( + torch::stable::view(out, {batch_size, num_heads_k, seqlen_q_ori, num_q_heads_per_hk, head_size_v}), + 1, 2), + {batch_size, seqlen_q_ori, num_heads_q, head_size_v}); + lse = torch::stable::reshape( + torch::stable::transpose( + torch::stable::view(lse, {batch_size, num_heads_k, seqlen_q_ori, num_q_heads_per_hk}), + 2, 3), + {batch_size, num_heads_q, seqlen_q_ori}); return {out, lse, tile_scheduler_metadata, num_splits}; } diff --git a/csrc/api/sparse_decode.h b/csrc/api/sparse_decode.h index adef4834..3d9fea09 100644 --- a/csrc/api/sparse_decode.h +++ b/csrc/api/sparse_decode.h @@ -180,18 +180,18 @@ class Decode_Sm100_Head128_Impl : public DecodeImplBase { } }; -static std::tuple, std::optional> +static std::tuple, std::optional> sparse_attn_decode_interface( - const at::Tensor &q, // [b, s_q, h_q, d_qk] - const at::Tensor &kv, // [num_blocks, page_block_size, h_k, d_qk] - const at::Tensor &indices, // [b, s_q, topk] - const std::optional &topk_length, // [b, s_q] - const std::optional &attn_sink, // [h_q] - std::optional tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) - std::optional num_splits, // batch_size + 1 - const std::optional &extra_kv, - const std::optional &extra_indices, - const std::optional &extra_topk_length, + const Tensor &q, // [b, s_q, h_q, d_qk] + const Tensor &kv, // [num_blocks, page_block_size, h_k, d_qk] + const Tensor &indices, // [b, s_q, topk] + const std::optional &topk_length, // [b, s_q] + const std::optional &attn_sink, // [h_q] + std::optional tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) + std::optional num_splits, // batch_size + 1 + const std::optional &extra_kv, + const std::optional &extra_indices, + const std::optional &extra_topk_length, int64_t d_v, double sm_scale ) { @@ -228,19 +228,19 @@ sparse_attn_decode_interface( } // metadata sanity check - TORCH_CHECK(b > 0); - TORCH_CHECK(s_q > 0); - TORCH_CHECK(h_q > 0); - TORCH_CHECK(h_kv == 1, "Currently only MQA (i.e. h_kv == 1) is supported for sparse decoding"); - TORCH_CHECK(d_qk == 576 || d_qk == 512, "Only head_size_k == 576 or 512 is supported for sparse decoding"); - TORCH_CHECK(d_v == 512, "Only head_size_v == 512 is supported for sparse decoding"); - TORCH_CHECK(topk > 0); + STD_TORCH_CHECK(b > 0); + STD_TORCH_CHECK(s_q > 0); + STD_TORCH_CHECK(h_q > 0); + STD_TORCH_CHECK(h_kv == 1, "Currently only MQA (i.e. h_kv == 1) is supported for sparse decoding"); + STD_TORCH_CHECK(d_qk == 576 || d_qk == 512, "Only head_size_k == 576 or 512 is supported for sparse decoding"); + STD_TORCH_CHECK(d_v == 512, "Only head_size_v == 512 is supported for sparse decoding"); + STD_TORCH_CHECK(topk > 0); if (have_extra_kcache) { - TORCH_CHECK(extra_indices.has_value(), "extra_indices_in_kvcache must be provided when extra_kcache is provided for sparse attention"); + STD_TORCH_CHECK(extra_indices.has_value(), "extra_indices_in_kvcache must be provided when extra_kcache is provided for sparse attention"); } else { - TORCH_CHECK(!extra_indices.has_value(), "extra_indices_in_kvcache must not be provided when extra_k_cache is not provided"); - TORCH_CHECK(!extra_topk_length.has_value(), "extra_topk_length must not be provided when extra_k_cache is not provided"); + STD_TORCH_CHECK(!extra_indices.has_value(), "extra_indices_in_kvcache must not be provided when extra_k_cache is not provided"); + STD_TORCH_CHECK(!extra_topk_length.has_value(), "extra_topk_length must not be provided when extra_k_cache is not provided"); } // Check device @@ -256,18 +256,18 @@ sparse_attn_decode_interface( KU_CHECK_DEVICE(extra_topk_length); // Check data type - KU_CHECK_DTYPE(q, torch::kBFloat16); - TORCH_CHECK(kv.dtype() == torch::kFloat8_e4m3fn || kv.dtype() == torch::kInt8 || kv.dtype() == torch::kUInt8, "key must have dtype fp8_e4m3fn, int8 or uint8"); + KU_CHECK_DTYPE(q, ScalarType::BFloat16); + STD_TORCH_CHECK(kv.scalar_type() == ScalarType::Float8_e4m3fn || kv.scalar_type() == ScalarType::Char || kv.scalar_type() == ScalarType::Byte, "key must have dtype fp8_e4m3fn, int8 or uint8"); if (extra_kv.has_value()) { - TORCH_CHECK(extra_kv->dtype() == torch::kFloat8_e4m3fn || extra_kv->dtype() == torch::kInt8 || extra_kv->dtype() == torch::kUInt8, "extra k cache must have dtype fp8_e4m3fn, int8 or uint8"); + STD_TORCH_CHECK(extra_kv->scalar_type() == ScalarType::Float8_e4m3fn || extra_kv->scalar_type() == ScalarType::Char || extra_kv->scalar_type() == ScalarType::Byte, "extra k cache must have dtype fp8_e4m3fn, int8 or uint8"); } - KU_CHECK_DTYPE(indices, torch::kInt32); - KU_CHECK_DTYPE(topk_length, torch::kInt32); - KU_CHECK_DTYPE(attn_sink, torch::kFloat32); - KU_CHECK_DTYPE(tile_scheduler_metadata, torch::kInt32); - KU_CHECK_DTYPE(num_splits, torch::kInt32); - KU_CHECK_DTYPE(extra_indices, torch::kInt32); - KU_CHECK_DTYPE(extra_topk_length, torch::kInt32); + KU_CHECK_DTYPE(indices, ScalarType::Int); + KU_CHECK_DTYPE(topk_length, ScalarType::Int); + KU_CHECK_DTYPE(attn_sink, ScalarType::Float); + KU_CHECK_DTYPE(tile_scheduler_metadata, ScalarType::Int); + KU_CHECK_DTYPE(num_splits, ScalarType::Int); + KU_CHECK_DTYPE(extra_indices, ScalarType::Int); + KU_CHECK_DTYPE(extra_topk_length, ScalarType::Int); // Check layout KU_CHECK_LAST_DIM_CONTIGUOUS(q); @@ -294,13 +294,13 @@ sparse_attn_decode_interface( // MODEL1 style bytes_per_token = 448 + 64*2 + (448/64)*1 + 1; } else { - TORCH_CHECK(false, "Unsupported head sizes for is_fp8_kvcache == True"); + STD_TORCH_CHECK(false, "Unsupported head sizes for is_fp8_kvcache == True"); } KU_CHECK_SHAPE(kv, num_blocks, page_block_size, h_kv, bytes_per_token); KU_CHECK_SHAPE(extra_kv, extra_num_blocks, extra_page_block_size, h_kv, bytes_per_token); - TORCH_CHECK(kv.stride(1) == bytes_per_token, "The whole block must be contiguous when is_fp8_cache is True for kv cache"); + STD_TORCH_CHECK(kv.stride(1) == bytes_per_token, "The whole block must be contiguous when is_fp8_cache is True for kv cache"); if (extra_kv.has_value()) { - TORCH_CHECK(extra_kv->stride(1) == bytes_per_token, "The whole block must be contiguous when is_fp8_cache is True for extra kv cache"); + STD_TORCH_CHECK(extra_kv->stride(1) == bytes_per_token, "The whole block must be contiguous when is_fp8_cache is True for extra kv cache"); } } KU_CHECK_SHAPE(indices, b, s_q, topk); @@ -309,11 +309,10 @@ sparse_attn_decode_interface( KU_CHECK_SHAPE(extra_indices, b, s_q, extra_topk); KU_CHECK_SHAPE(extra_topk_length, b); - at::cuda::CUDAGuard device_guard{(char)q.get_device()}; - auto opts = q.options(); + torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index()); - at::Tensor out = torch::empty({b, s_q, h_q, d_v}, opts); - at::Tensor lse = torch::empty({b, s_q, h_q}, opts.dtype(at::kFloat)); + Tensor out = torch::stable::new_empty(q, {b, s_q, h_q, d_v}); + Tensor lse = torch::stable::new_empty(q, {b, s_q, h_q}, ScalarType::Float); ModelType model_type; if (d_qk == 576) { @@ -321,7 +320,7 @@ sparse_attn_decode_interface( } else if (d_qk == 512) { model_type = ModelType::MODEL1; } else { - TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); + STD_TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); } std::vector features; @@ -330,21 +329,21 @@ sparse_attn_decode_interface( } else if (h_q == 128) { features.push_back(DecodeFeatures::HEAD_128); } else { - TORCH_CHECK(false, "Unsupported h_q: ", h_q); + STD_TORCH_CHECK(false, "Unsupported h_q: ", h_q); } if (d_qk == 576) { features.push_back(DecodeFeatures::HEAD_DIM_576); } else if (d_qk == 512) { features.push_back(DecodeFeatures::HEAD_DIM_512); } else { - TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); + STD_TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); } if (model_type == ModelType::V32) { features.push_back(DecodeFeatures::V32_KVCACHE_FORMAT); } else if (model_type == ModelType::MODEL1) { features.push_back(DecodeFeatures::MODEL1_KVCACHE_FORMAT); } else { - TORCH_CHECK(false, "Unsupported model type: ", (int)model_type); + STD_TORCH_CHECK(false, "Unsupported model type: ", (int)model_type); } if (have_attn_sink) { features.push_back(DecodeFeatures::ATTN_SINK); @@ -369,15 +368,15 @@ sparse_attn_decode_interface( } else if (d_qk == 512) { impl = new Decode_Sm100_Head128_Impl(); } else { - TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); + STD_TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); } } else { - TORCH_CHECK(false, "Unsupported h_q: ", h_q); + STD_TORCH_CHECK(false, "Unsupported h_q: ", h_q); } } else if (arch.is_sm90a()) { impl = new Decode_Sm90_Impl(); } else { - TORCH_CHECK(false, "Unsupported architecture for sparse decode fwd"); + STD_TORCH_CHECK(false, "Unsupported architecture for sparse decode fwd"); } DecodeImplMeta impl_meta = impl->get_meta(h_q, s_q); @@ -411,14 +410,14 @@ sparse_attn_decode_interface( have_extra_kcache ? int64_stride_to_int(extra_kv->stride(1)) : 0, have_extra_kcache ? int64_stride_to_int(extra_indices->stride(0)) : 0, have_extra_kcache ? int64_stride_to_int(extra_indices->stride(1)) : 0, - at::cuda::getCurrentCUDAStream().stream() + get_current_cuda_stream(q) }; // Get MLA metadata if necessary - at::Tensor o_accum, lse_accum; + Tensor o_accum, lse_accum; if (!tile_scheduler_metadata.has_value()) { - tile_scheduler_metadata = torch::empty({impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/4}, opts.dtype(torch::kInt32)); - num_splits = torch::empty({b+1}, opts.dtype(torch::kInt32)); + tile_scheduler_metadata = torch::stable::new_empty(q, {impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/4}, ScalarType::Int); + num_splits = torch::stable::new_empty(q, {b+1}, ScalarType::Int); KU_CHECK_CONTIGUOUS(tile_scheduler_metadata); KU_CHECK_CONTIGUOUS(num_splits); @@ -432,33 +431,33 @@ sparse_attn_decode_interface( ku::get_optional_tensor_ptr(extra_topk_length), nullptr, (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr(), - num_splits->data_ptr(), + num_splits->mutable_data_ptr(), impl_meta.num_sm_parts, - at::cuda::getCurrentCUDAStream().stream() + get_current_cuda_stream(q) }; smxx::decode::run_get_decoding_sched_meta_kernel(get_sched_meta_params); } // Stick the metadata pointers to `params` KU_CHECK_DEVICE(tile_scheduler_metadata); KU_CHECK_DEVICE(num_splits); - KU_CHECK_DTYPE(tile_scheduler_metadata, torch::kInt32); - KU_CHECK_DTYPE(num_splits, torch::kInt32); + KU_CHECK_DTYPE(tile_scheduler_metadata, ScalarType::Int); + KU_CHECK_DTYPE(num_splits, ScalarType::Int); KU_CHECK_CONTIGUOUS(tile_scheduler_metadata); KU_CHECK_CONTIGUOUS(num_splits); KU_CHECK_SHAPE(tile_scheduler_metadata, impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/sizeof(int)); KU_CHECK_SHAPE(num_splits, b+1); params.tile_scheduler_metadata_ptr = (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr(); - params.num_splits_ptr = num_splits->data_ptr(); + params.num_splits_ptr = num_splits->mutable_data_ptr(); params.num_sm_parts = impl_meta.num_sm_parts; // Allocate intermediate buffers for split-KV const int total_num_splits = b + impl_meta.num_sm_parts; - lse_accum = torch::empty({total_num_splits, s_q, h_q}, opts.dtype(at::kFloat)); - o_accum = torch::empty({total_num_splits, s_q, h_q, d_v}, opts.dtype(at::kFloat)); + lse_accum = torch::stable::new_empty(q, {total_num_splits, s_q, h_q}, ScalarType::Float); + o_accum = torch::stable::new_empty(q, {total_num_splits, s_q, h_q, d_v}, ScalarType::Float); KU_CHECK_CONTIGUOUS(lse_accum); KU_CHECK_CONTIGUOUS(o_accum); - params.lse_accum = lse_accum.data_ptr(); - params.o_accum = o_accum.data_ptr(); + params.lse_accum = lse_accum.mutable_data_ptr(); + params.o_accum = o_accum.mutable_data_ptr(); params.stride_lse_accum_split = int64_stride_to_int(lse_accum.stride(0)); params.stride_lse_accum_s_q = int64_stride_to_int(lse_accum.stride(1)); params.stride_o_accum_split = int64_stride_to_int(o_accum.stride(0)); @@ -485,11 +484,11 @@ sparse_attn_decode_interface( params.num_sm_parts, ku::get_optional_tensor_ptr(attn_sink), - at::cuda::getCurrentCUDAStream().stream() + get_current_cuda_stream(q) }; smxx::decode::run_flash_mla_combine_kernel(combine_params); delete impl; - return {out, lse.transpose(1, 2), tile_scheduler_metadata, num_splits}; + return {out, torch::stable::transpose(lse, 1, 2), tile_scheduler_metadata, num_splits}; } diff --git a/csrc/api/sparse_fwd.h b/csrc/api/sparse_fwd.h index 5157d1cc..ea146a11 100644 --- a/csrc/api/sparse_fwd.h +++ b/csrc/api/sparse_fwd.h @@ -98,21 +98,21 @@ class Fwd_Sm100_Head128_Small_TopK_Impl : public FwdImplBase { } }; -static std::vector sparse_attn_prefill_interface( - const at::Tensor &q, - const at::Tensor &kv, - const at::Tensor &indices, +static std::vector sparse_attn_prefill_interface( + const Tensor &q, + const Tensor &kv, + const Tensor &indices, double sm_scale, int64_t d_v, - const std::optional &attn_sink, - const std::optional &topk_length + const std::optional &attn_sink, + const std::optional &topk_length ) { using bf16 = cutlass::bfloat16_t; Arch arch = Arch(); bool is_sm90a = arch.is_sm90a(); bool is_sm100f = arch.is_sm100f(); - TORCH_CHECK(is_sm90a || is_sm100f, "Sparse Attention Forward Kernel is only supported on SM90a and SM100f architectures."); + STD_TORCH_CHECK(is_sm90a || is_sm100f, "Sparse Attention Forward Kernel is only supported on SM90a and SM100f architectures."); KU_CHECK_NDIM(q, 3); KU_CHECK_NDIM(kv, 3); @@ -128,8 +128,8 @@ static std::vector sparse_attn_prefill_interface( int topk = indices.size(2); bool have_topk_length = topk_length.has_value(); - TORCH_CHECK(d_qk == 576 || d_qk == 512, "Invalid d_qk: ", d_qk); - TORCH_CHECK(d_v == 512, "Invalid d_v", d_v); + STD_TORCH_CHECK(d_qk == 576 || d_qk == 512, "Invalid d_qk: ", d_qk); + STD_TORCH_CHECK(d_v == 512, "Invalid d_v", d_v); KU_CHECK_DEVICE(q); KU_CHECK_DEVICE(kv); @@ -137,11 +137,11 @@ static std::vector sparse_attn_prefill_interface( KU_CHECK_DEVICE(attn_sink); KU_CHECK_DEVICE(topk_length); - KU_CHECK_DTYPE(q, torch::kBFloat16); - KU_CHECK_DTYPE(kv, torch::kBFloat16); - KU_CHECK_DTYPE(indices, torch::kInt32); - KU_CHECK_DTYPE(attn_sink, torch::kFloat32); - KU_CHECK_DTYPE(topk_length, torch::kInt32); + KU_CHECK_DTYPE(q, ScalarType::BFloat16); + KU_CHECK_DTYPE(kv, ScalarType::BFloat16); + KU_CHECK_DTYPE(indices, ScalarType::Int); + KU_CHECK_DTYPE(attn_sink, ScalarType::Float); + KU_CHECK_DTYPE(topk_length, ScalarType::Int); KU_CHECK_SHAPE(q, s_q, h_q, d_qk); KU_CHECK_SHAPE(kv, s_kv, h_kv, d_qk); @@ -156,12 +156,11 @@ static std::vector sparse_attn_prefill_interface( KU_CHECK_LAST_DIM_CONTIGUOUS(topk_length); // Allocate results and buffers - at::cuda::CUDAGuard device_guard{(char)q.get_device()}; - auto opts = q.options(); - - at::Tensor out = torch::empty({s_q, h_q, d_v}, opts); - at::Tensor lse = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); - at::Tensor max_logits = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); + torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index()); + + Tensor out = torch::stable::new_empty(q, {s_q, h_q, d_v}); + Tensor lse = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float); + Tensor max_logits = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float); KU_CHECK_CONTIGUOUS(out); KU_CHECK_CONTIGUOUS(lse); KU_CHECK_CONTIGUOUS(max_logits); @@ -185,7 +184,7 @@ static std::vector sparse_attn_prefill_interface( (float*)lse.data_ptr(), arch.num_sms, - at::cuda::getCurrentCUDAStream().stream() + get_current_cuda_stream(q) }; std::vector required_features; @@ -194,14 +193,14 @@ static std::vector sparse_attn_prefill_interface( } else if (h_q == 128) { required_features.push_back(FwdFeatures::HEAD_128); } else { - TORCH_CHECK(false, "Unsupported h_q: ", h_q); + STD_TORCH_CHECK(false, "Unsupported h_q: ", h_q); } if (d_qk == 576) { required_features.push_back(FwdFeatures::HEAD_DIM_576); } else if (d_qk == 512) { required_features.push_back(FwdFeatures::HEAD_DIM_512); } else { - TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); + STD_TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); } if (attn_sink.has_value()) { required_features.push_back(FwdFeatures::ATTN_SINK); @@ -233,10 +232,10 @@ static std::vector sparse_attn_prefill_interface( regular_impl.run(params, required_features); } } else { - TORCH_CHECK(false, "Unsupported h_q: ", h_q); + STD_TORCH_CHECK(false, "Unsupported h_q: ", h_q); } } else { - TORCH_CHECK(false, "Unsupported architecture"); + STD_TORCH_CHECK(false, "Unsupported architecture"); } return {out, max_logits, lse}; diff --git a/csrc/kerutils/include/kerutils/supplemental/cuda_stream.h b/csrc/kerutils/include/kerutils/supplemental/cuda_stream.h new file mode 100644 index 00000000..d15c9e49 --- /dev/null +++ b/csrc/kerutils/include/kerutils/supplemental/cuda_stream.h @@ -0,0 +1,19 @@ +#pragma once + +#include + +#include +#include +#include + +namespace kerutils { + +// Helper to return the current CUDA stream for the device that `t` lives on, as +// a raw cudaStream_t. Semantically equivalent to at::cuda::getCurrentCUDAStream(). +inline cudaStream_t get_current_cuda_stream(const torch::stable::Tensor &t) { + void *stream_ptr = nullptr; + TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(t.get_device_index(), &stream_ptr)); + return static_cast(stream_ptr); +} + +} diff --git a/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h b/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h index 4d798bfb..d354092d 100644 --- a/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h +++ b/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h @@ -1,8 +1,11 @@ #pragma once #include +#include -#include // libtorch C++ API without pybind (Py_LIMITED_API-safe) +#include +#include +#include #include "kerutils/common/common.h" @@ -12,8 +15,8 @@ namespace kerutils { // If tensor_or_opt is a tensor, check_fn is applied directly // If tensor_or_opt is an optional tensor, check_fn is applied only when the optional has value template -static inline bool _check_optional_tensor(const T& tensor_or_opt, const std::function& check_fn) { - if constexpr (std::is_same::value) { +static inline bool _check_optional_tensor(const T& tensor_or_opt, const std::function& check_fn) { + if constexpr (std::is_same::value) { return check_fn(tensor_or_opt); } else { if (tensor_or_opt.has_value()) { @@ -25,10 +28,10 @@ static inline bool _check_optional_tensor(const T& tensor_or_opt, const std::fun } // Get the pointer of the given tensor -// Return (PtrT*)tensor.data_ptr() if the tensor has a backend storage, nullptr otherwise +// Return (PtrT*)tensor.data_ptr() if the tensor is defined, nullptr otherwise template -static inline PtrT* get_tensor_ptr(const at::Tensor& tensor) { - if (tensor.has_storage()) { +static inline PtrT* get_tensor_ptr(const torch::stable::Tensor& tensor) { + if (tensor.defined()) { return (PtrT*)tensor.data_ptr(); } else { return nullptr; @@ -39,7 +42,7 @@ static inline PtrT* get_tensor_ptr(const at::Tensor& tensor) { // Return (PtrT*)tensor.data_ptr() if tensor_or_opt has value and points to a valid tensor, return nullptr otherwise template static inline PtrT* get_optional_tensor_ptr(const T& tensor_or_opt) { - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { return get_tensor_ptr(tensor_or_opt); } else { if (tensor_or_opt.has_value()) { @@ -53,19 +56,19 @@ static inline PtrT* get_optional_tensor_ptr(const T& tensor_or_opt) { } // Check whether the given tensor (or optional) is on cuda -#define KU_CHECK_DEVICE(tensor) TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const at::Tensor& t) { return t.is_cuda(); }), #tensor " must be on CUDA") +#define KU_CHECK_DEVICE(tensor) STD_TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const torch::stable::Tensor& t) { return t.is_cuda(); }), #tensor " must be on CUDA") // Check whether the given tensor (or optional) has the given number of dimensions -#define KU_CHECK_NDIM(tensor, ndim) TORCH_CHECK(ku::_check_optional_tensor(tensor, [&](const at::Tensor& t) { return t.dim() == (ndim); }), #tensor " must have " #ndim " dimensions") +#define KU_CHECK_NDIM(tensor, ndim) STD_TORCH_CHECK(ku::_check_optional_tensor(tensor, [&](const torch::stable::Tensor& t) { return t.dim() == (ndim); }), #tensor " must have " #ndim " dimensions") // Check whether the given tensor (or optional) has the given shape -#define KU_CHECK_SHAPE(tensor, ...) TORCH_CHECK(ku::_check_optional_tensor(tensor, [&](const at::Tensor& t) { return t.sizes() == torch::IntArrayRef({__VA_ARGS__}); }), #tensor " must have shape (" #__VA_ARGS__ ")") +#define KU_CHECK_SHAPE(tensor, ...) STD_TORCH_CHECK(ku::_check_optional_tensor(tensor, [&](const torch::stable::Tensor& t) { return t.sizes().equals({__VA_ARGS__}); }), #tensor " must have shape (" #__VA_ARGS__ ")") // Check whether the given tensor (or optional) is contiguous -#define KU_CHECK_CONTIGUOUS(tensor) TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const at::Tensor& t) { return t.is_contiguous(); }), #tensor " must be contiguous") +#define KU_CHECK_CONTIGUOUS(tensor) STD_TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const torch::stable::Tensor& t) { return t.is_contiguous(); }), #tensor " must be contiguous") // Check whether the last dimention of the given tensor (or optional) -#define KU_CHECK_LAST_DIM_CONTIGUOUS(tensor) TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const at::Tensor& t) { return t.size(-1) == 1 || t.stride(-1) == 1; }), #tensor " must have contiguous last dimension") +#define KU_CHECK_LAST_DIM_CONTIGUOUS(tensor) STD_TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const torch::stable::Tensor& t) { return t.size(-1) == 1 || t.stride(-1) == 1; }), #tensor " must have contiguous last dimension") // Check whether the given tensor (or optional) has the specified dtype -#define KU_CHECK_DTYPE(tensor, target_dtype) TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const at::Tensor& t) { return t.dtype() == (target_dtype); }), #tensor " must have dtype " #target_dtype) +#define KU_CHECK_DTYPE(tensor, target_dtype) STD_TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const torch::stable::Tensor& t) { return t.scalar_type() == (target_dtype); }), #tensor " must have dtype " #target_dtype) diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu index b67ea9e4..1499c846 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu +++ b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu @@ -1,7 +1,8 @@ #include "interface.h" -#include -#include +#include +#include +#include #include #include "common/mask.cuh" #include "common/utils.hpp" @@ -11,10 +12,10 @@ template void call_run_fmha_bwd([[maybe_unused]] Mask mask, [[maybe_unused]] Varlen is_varlen, [[maybe_unused]] Element in, [[maybe_unused]] ElementOut out, [[maybe_unused]] Mla mla, - at::Tensor workspace_buffer, at::Tensor d_o, at::Tensor q, at::Tensor k, - at::Tensor v, at::Tensor o, at::Tensor lse, - at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, - at::Tensor dq, at::Tensor dk, at::Tensor dv, + torch::stable::Tensor workspace_buffer, torch::stable::Tensor d_o, torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor v, torch::stable::Tensor o, torch::stable::Tensor lse, + torch::stable::Tensor cumulative_seqlen_q, torch::stable::Tensor cumulative_seqlen_kv, + torch::stable::Tensor dq, torch::stable::Tensor dk, torch::stable::Tensor dv, float softmax_scale, int max_seqlen_q, int total_seqlen_kv) { static constexpr bool IsVarlen = std::is_same_v; static constexpr bool IsMla = std::is_same_v; @@ -26,13 +27,13 @@ void call_run_fmha_bwd([[maybe_unused]] Mask mask, [[maybe_unused]] Varlen is_va } -void FMHACutlassSM100BwdRun(at::Tensor workspace_buffer, at::Tensor d_o, at::Tensor q, at::Tensor k, - at::Tensor v, at::Tensor o, at::Tensor lse, - at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, - at::Tensor dq, at::Tensor dk, at::Tensor dv, +void FMHACutlassSM100BwdRun(torch::stable::Tensor workspace_buffer, torch::stable::Tensor d_o, torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor v, torch::stable::Tensor o, torch::stable::Tensor lse, + torch::stable::Tensor cumulative_seqlen_q, torch::stable::Tensor cumulative_seqlen_kv, + torch::stable::Tensor dq, torch::stable::Tensor dk, torch::stable::Tensor dv, int64_t mask_mode_code, double softmax_scale, int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_varlen) { - const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index()); int head_dim_qk = q.size(-1); int head_dim_vo = v.size(-1); @@ -40,7 +41,7 @@ void FMHACutlassSM100BwdRun(at::Tensor workspace_buffer, at::Tensor d_o, at::Ten auto scalar_type_in = q.scalar_type(); auto scalar_type_out = o.scalar_type(); - if(scalar_type_in == at::ScalarType::BFloat16 && scalar_type_out == at::ScalarType::BFloat16) { + if(scalar_type_in == torch::headeronly::ScalarType::BFloat16 && scalar_type_out == torch::headeronly::ScalarType::BFloat16) { using Element = cutlass::bfloat16_t; using ElementOut = cutlass::bfloat16_t; diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh index 086a4687..29a3abb9 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh +++ b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh @@ -49,8 +49,12 @@ #include "collective/fmha_fusion.hpp" #include "device/fmha_device_bwd.hpp" -#include -#include +#include +#include +#include +#include + +#include using namespace cute; using namespace cutlass::fmha::kernel; @@ -93,13 +97,13 @@ struct BwdRunner { using StrideDV = TensorStride; // Seq DVO (H B) using StrideDO = TensorStride; - static void run(at::Tensor workspace_buffer, at::Tensor d_o, at::Tensor q, at::Tensor k, - at::Tensor v, at::Tensor o, at::Tensor lse, - at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, - at::Tensor dq, at::Tensor dk, at::Tensor dv, + static void run(torch::stable::Tensor workspace_buffer, torch::stable::Tensor d_o, torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor v, torch::stable::Tensor o, torch::stable::Tensor lse, + torch::stable::Tensor cumulative_seqlen_q, torch::stable::Tensor cumulative_seqlen_kv, + torch::stable::Tensor dq, torch::stable::Tensor dk, torch::stable::Tensor dv, float softmax_scale, int max_seqlen_q, int max_seqlen_kv) { - const at::cuda::CUDAGuard device_guard{(char)q.get_device()}; - const int device_id = q.get_device(); + const torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index()); + const int device_id = q.get_device_index(); cutlass::KernelHardwareInfo hw_info; hw_info.device_id =device_id; @@ -142,15 +146,15 @@ struct BwdRunner { int dk_stride0 = dk.stride(0), dk_stride1 = dk.stride(1), dk_stride2 = dk.stride(2); int dv_stride0 = dv.stride(0), dv_stride1 = dv.stride(1), dv_stride2 = dv.stride(2); int do_stride0 = d_o.stride(0), do_stride1 = d_o.stride(1), do_stride2 = d_o.stride(2); - TORCH_CHECK(q_stride2 == 1); - TORCH_CHECK(k_stride2 == 1); - TORCH_CHECK(v_stride2 == 1); - TORCH_CHECK(o_stride2 == 1); - TORCH_CHECK(lse_stride0 == 1); - TORCH_CHECK(dq_stride2 == 1); - TORCH_CHECK(dk_stride2 == 1); - TORCH_CHECK(dv_stride2 == 1); - TORCH_CHECK(do_stride2 == 1); + STD_TORCH_CHECK(q_stride2 == 1); + STD_TORCH_CHECK(k_stride2 == 1); + STD_TORCH_CHECK(v_stride2 == 1); + STD_TORCH_CHECK(o_stride2 == 1); + STD_TORCH_CHECK(lse_stride0 == 1); + STD_TORCH_CHECK(dq_stride2 == 1); + STD_TORCH_CHECK(dk_stride2 == 1); + STD_TORCH_CHECK(dv_stride2 == 1); + STD_TORCH_CHECK(do_stride2 == 1); StrideQ stride_Q = make_stride(q_stride0, _1{}, make_stride(q_stride1, B == 1 ? 0 : q_stride0*Q)); StrideK stride_K = make_stride(k_stride0, _1{}, make_stride(k_stride1, B == 1 ? 0 : k_stride0*K)); @@ -184,17 +188,17 @@ struct BwdRunner { CUTLASS_CHECK(op.can_implement(arguments)); CUTLASS_CHECK(op.initialize(arguments, workspace_ptr)); - CUTLASS_CHECK(op.run(at::cuda::getCurrentCUDAStream())); + CUTLASS_CHECK(op.run(kerutils::get_current_cuda_stream(q))); } }; template -void run_fmha_bwd(at::Tensor workspace_buffer, at::Tensor d_o, at::Tensor q, at::Tensor k, - at::Tensor v, at::Tensor o, at::Tensor lse, - at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, - at::Tensor dq, at::Tensor dk, at::Tensor dv, +void run_fmha_bwd(torch::stable::Tensor workspace_buffer, torch::stable::Tensor d_o, torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor v, torch::stable::Tensor o, torch::stable::Tensor lse, + torch::stable::Tensor cumulative_seqlen_q, torch::stable::Tensor cumulative_seqlen_kv, + torch::stable::Tensor dq, torch::stable::Tensor dk, torch::stable::Tensor dv, float softmax_scale, int max_seqlen_q, int total_seqlen_kv) { BwdRunner::run(workspace_buffer, d_o, q, k, v, o, lse, cumulative_seqlen_q, cumulative_seqlen_kv, diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu index 6a4ef4ff..2ca4a18d 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu +++ b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu @@ -1,7 +1,8 @@ #include "interface.h" -#include -#include +#include +#include +#include #include #include "common/mask.cuh" @@ -12,9 +13,9 @@ template void call_run_fmha_fwd([[maybe_unused]] Mask mask, [[maybe_unused]] Varlen is_varlen, [[maybe_unused]] Element in, [[maybe_unused]] ElementOut out, - [[maybe_unused]] Mla mla, at::Tensor workspace_buffer, at::Tensor q, - at::Tensor k, at::Tensor v, at::Tensor cumulative_seqlen_q, - at::Tensor cumulative_seqlen_kv, at::Tensor o, at::Tensor lse, + [[maybe_unused]] Mla mla, torch::stable::Tensor workspace_buffer, torch::stable::Tensor q, + torch::stable::Tensor k, torch::stable::Tensor v, torch::stable::Tensor cumulative_seqlen_q, + torch::stable::Tensor cumulative_seqlen_kv, torch::stable::Tensor o, torch::stable::Tensor lse, float softmax_scale, int max_seqlen_q, int max_seqlen_kv) { static constexpr bool IsVarlen = std::is_same_v; static constexpr bool IsMla = std::is_same_v; @@ -28,21 +29,21 @@ void call_run_fmha_fwd([[maybe_unused]] Mask mask, [[maybe_unused]] Varlen is_va softmax_scale, max_seqlen_q, max_seqlen_kv); } -void FMHACutlassSM100FwdRun(at::Tensor workspace_buffer, at::Tensor q, at::Tensor k, - at::Tensor v, at::Tensor cumulative_seqlen_q, - at::Tensor cumulative_seqlen_kv, at::Tensor o, at::Tensor lse, +void FMHACutlassSM100FwdRun(torch::stable::Tensor workspace_buffer, torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor v, torch::stable::Tensor cumulative_seqlen_q, + torch::stable::Tensor cumulative_seqlen_kv, torch::stable::Tensor o, torch::stable::Tensor lse, int64_t mask_mode_code, double sm_scale, int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_varlen) { - const c10::cuda::OptionalCUDAGuard device_guard(q.device()); - CHECK(q.scalar_type() == k.scalar_type()); + const torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index()); + STD_TORCH_CHECK(q.scalar_type() == k.scalar_type()); auto scalar_type_in = q.scalar_type(); auto scalar_type_out = o.scalar_type(); int head_dim_qk = q.size(-1); int head_dim_vo = v.size(-1); MaskMode mask_mode = static_cast(mask_mode_code); - if (scalar_type_in == at::ScalarType::BFloat16 && - scalar_type_out == at::ScalarType::BFloat16) { + if (scalar_type_in == torch::headeronly::ScalarType::BFloat16 && + scalar_type_out == torch::headeronly::ScalarType::BFloat16) { using Element = cutlass::bfloat16_t; using ElementOut = cutlass::bfloat16_t; diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh index 58adf3c6..a0f70bd0 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh +++ b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh @@ -12,9 +12,12 @@ #include "kernel/fmha_tile_scheduler.hpp" #include "kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp" -#include -#include -#include +#include +#include +#include +#include + +#include using namespace cute; using namespace cutlass::fmha::collective; @@ -220,10 +223,10 @@ struct FwdRunner { } template - void run(const Options &options, const cutlass::KernelHardwareInfo &hw_info, at::Tensor q, - at::Tensor k, at::Tensor v, at::Tensor o, at::Tensor lse, float scale_softmax, - at::Tensor workspace, at::Tensor cumulative_seqlen_q, - at::Tensor cumulative_seqlen_kv, int max_seqlen_q, int max_seqlen_kv) { + void run(const Options &options, const cutlass::KernelHardwareInfo &hw_info, torch::stable::Tensor q, + torch::stable::Tensor k, torch::stable::Tensor v, torch::stable::Tensor o, torch::stable::Tensor lse, float scale_softmax, + torch::stable::Tensor workspace, torch::stable::Tensor cumulative_seqlen_q, + torch::stable::Tensor cumulative_seqlen_kv, int max_seqlen_q, int max_seqlen_kv) { int total_seqlen_q = q.size(0); int total_seqlen_kv = k.size(0); @@ -244,11 +247,11 @@ struct FwdRunner { int v_stride0 = v.stride(0), v_stride1 = v.stride(1), v_stride2 = v.stride(2); int o_stride0 = o.stride(0), o_stride1 = o.stride(1), o_stride2 = o.stride(2); int lse_stride0 = lse.stride(0), lse_stride1 = lse.stride(1); - TORCH_CHECK(q_stride2 == 1); - TORCH_CHECK(k_stride2 == 1); - TORCH_CHECK(v_stride2 == 1); - TORCH_CHECK(o_stride2 == 1); - TORCH_CHECK(lse_stride0 == 1); + STD_TORCH_CHECK(q_stride2 == 1); + STD_TORCH_CHECK(k_stride2 == 1); + STD_TORCH_CHECK(v_stride2 == 1); + STD_TORCH_CHECK(o_stride2 == 1); + STD_TORCH_CHECK(lse_stride0 == 1); stride_Q = make_stride(q_stride0, _1{}, make_stride(make_stride(q_stride1, H_Q * q_stride1), SQ * q_stride0)); stride_O = make_stride(o_stride0, _1{}, make_stride(make_stride(o_stride1, H_Q * o_stride1), SQ * o_stride0)); @@ -279,18 +282,18 @@ struct FwdRunner { CUTLASS_CHECK(op.can_implement(arguments)); CUTLASS_CHECK(op.initialize(arguments, nullptr)); - CUTLASS_CHECK(op.run(at::cuda::getCurrentCUDAStream())); + CUTLASS_CHECK(op.run(kerutils::get_current_cuda_stream(q))); } }; template -void run_fmha_fwd(at::Tensor workspace, at::Tensor q, at::Tensor k, at::Tensor v, - at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, at::Tensor o, - at::Tensor lse, float scale_softmax, int max_seqlen_q, int max_seqlen_kv) { +void run_fmha_fwd(torch::stable::Tensor workspace, torch::stable::Tensor q, torch::stable::Tensor k, torch::stable::Tensor v, + torch::stable::Tensor cumulative_seqlen_q, torch::stable::Tensor cumulative_seqlen_kv, torch::stable::Tensor o, + torch::stable::Tensor lse, float scale_softmax, int max_seqlen_q, int max_seqlen_kv) { - const at::cuda::CUDAGuard device_guard{(char)q.get_device()}; - const int device_id = q.get_device(); + const torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index()); + const int device_id = q.get_device_index(); cutlass::KernelHardwareInfo hw_info; hw_info.device_id = device_id; diff --git a/csrc/sm100/prefill/dense/interface.h b/csrc/sm100/prefill/dense/interface.h index 0ddc553a..7ff8246e 100644 --- a/csrc/sm100/prefill/dense/interface.h +++ b/csrc/sm100/prefill/dense/interface.h @@ -1,14 +1,14 @@ #pragma once -#include +#include -void FMHACutlassSM100FwdRun(at::Tensor workspace_buffer, at::Tensor q, at::Tensor k, at::Tensor v, - at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, - at::Tensor o, at::Tensor lse, +void FMHACutlassSM100FwdRun(torch::stable::Tensor workspace_buffer, torch::stable::Tensor q, torch::stable::Tensor k, torch::stable::Tensor v, + torch::stable::Tensor cumulative_seqlen_q, torch::stable::Tensor cumulative_seqlen_kv, + torch::stable::Tensor o, torch::stable::Tensor lse, int64_t mask_mode_code, double softmax_scale, int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_varlen); -void FMHACutlassSM100BwdRun(at::Tensor workspace_buffer, at::Tensor d_o, at::Tensor q, at::Tensor k, - at::Tensor v, at::Tensor o, at::Tensor lse, - at::Tensor cumulative_seqlen_q, at::Tensor cumulative_seqlen_kv, - at::Tensor dq, at::Tensor dk, at::Tensor dv, +void FMHACutlassSM100BwdRun(torch::stable::Tensor workspace_buffer, torch::stable::Tensor d_o, torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor v, torch::stable::Tensor o, torch::stable::Tensor lse, + torch::stable::Tensor cumulative_seqlen_q, torch::stable::Tensor cumulative_seqlen_kv, + torch::stable::Tensor dq, torch::stable::Tensor dk, torch::stable::Tensor dv, int64_t mask_mode_code, double softmax_scale, int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_varlen); diff --git a/setup.py b/setup.py index 2e4a8a0b..4cc6b422 100644 --- a/setup.py +++ b/setup.py @@ -17,7 +17,13 @@ def is_flag_set(flag: str) -> bool: return os.getenv(flag, "FALSE").lower() in ["true", "1", "y", "yes"] def get_features_args(): - features_args = [] + # ABI-stable flags (always on): TORCH_TARGET_VERSION pins the minimum runtime + # PyTorch version and bans unstable ATen/c10/torch headers at compile time; + # USE_CUDA exposes aoti_torch_get_current_cuda_stream from the shim. + features_args = [ + "-DTORCH_TARGET_VERSION=0x020a000000000000", # PyTorch >= 2.10 at runtime + "-DUSE_CUDA", + ] if is_flag_set("FLASH_MLA_DISABLE_FP16"): features_args.append("-DFLASH_MLA_DISABLE_FP16") return features_args From 1638811f52338290aaab7d53c051d6d0eec847d5 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Thu, 2 Jul 2026 10:11:38 -0700 Subject: [PATCH 3/3] Add a README edit describing the fork and build instructions --- README.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/README.md b/README.md index 1945725f..a17ffc52 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,31 @@ # FlashMLA +This fork is a LibTorch and CPython ABI stable version of the original FlashMLA. This means that you can build one wheel that suffices across Python 3.10+ and torch 2.10+ at runtime with a `pip install`. You can build and install it directly from GitHub; the cutlass submodule will be fetched automatically during the build: + +```zsh +pip install -v --no-build-isolation "git+https://github.com/janeyx99/FlashMLA-ABI-Stable.git" +``` + +**Build prerequisites** +- PyTorch **2.11+** + - Example command for cuda 13.0: `pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cu130` + - Note that the build time requirement is intentionally 2.11 to access nicer UX features for the stable ABI. However at runtime, the wheel will work with torch 2.10+. + - `--no-build-isolation` is required so the build can import your existing torch. +- A CUDA toolkit with `nvcc` **12.9+** for the SM100 (Blackwell) kernels. With `nvcc` 12.8 you must disable SM100: prefix the command with `FLASH_MLA_DISABLE_SM100=1`. + +### Install from a local clone + +```zsh +git clone https://github.com/janeyx99/FlashMLA-ABI-Stable.git +cd FlashMLA-ABI-Stable +git submodule update --init --recursive +pip install -v --no-build-isolation . +``` + +Below is the original README with all the deets -- credits to the original authors from DeepSeek. + + + ## Introduction FlashMLA is DeepSeek's library of optimized attention kernels, powering the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) and [DeepSeek-V3.2-Exp](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp) models. This repository contains the following implementations: