Hugging Face's logo

Kernels:
phanerozoic
/
mamba3

Kernel card Files
xet
 Community
mamba3 / torch-ext /torch_binding.cpp
phanerozoic's picture
phanerozoic
Bring the v1 sources and card up to main
6e06900
Raw
History Blame
18.5 kB
#include <torch/library.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <algorithm>
#include "torch_binding.h"
// Not CHECK: torch/library.h pulls in glog's CHECK and redefining it warns.
#define M3_F32(x) TORCH_CHECK((x).is_cuda() && (x).scalar_type() == at::kFloat && (x).is_contiguous(), \
 #x " must be a contiguous float32 CUDA tensor")
#define STEP_ARGS \
 const void* q, const void* k, const void* v, const void* z, const float* adt, \
 const float* dt, const float* trap, const void* q_bias, const void* k_bias, \
 const float* angles, const void* mimo_v, const void* mimo_o, const void* mimo_z, \
 const void* Dvec, const float* angle_in, float* angle_out, float* S_st, \
 const float* kprev_in, const float* vprev_in, float* kprev_out, float* vprev_out, \
 void* y, int B, int H, int Gqk, int R, int N, int P, int Na, void* stream
void mamba3_mimo_step_launch_f32(STEP_ARGS);
void mamba3_mimo_step_launch_bf16(STEP_ARGS);
static const void* vptr(const c10::optional<at::Tensor>& t, at::ScalarType st) {
 if (!t.has_value() || !t->defined()) return nullptr;
 TORCH_CHECK(t->is_cuda() && t->scalar_type() == st && t->is_contiguous(),
 "optional tensors must be contiguous CUDA tensors of the input dtype");
 return t->data_ptr();
}
void mamba3_cumulative_angles_launch(
 const float* angles, const float* dt, float* out, float* seg_tot,
 int B, int S, int H, int Na, int nseg, void* stream);
int mamba3_cumulative_angles_segments(int B, int S, int H);
torch::Tensor mamba3_cumulative_angles(torch::Tensor angles, torch::Tensor dt)
{
 M3_F32(angles); M3_F32(dt);
 TORCH_CHECK(angles.dim() == 4 && dt.dim() == 3, "angles is (B,S,H,Na), dt is (B,H,S)");
 const int B = (int)angles.size(0), S = (int)angles.size(1);
 const int H = (int)angles.size(2), Na = (int)angles.size(3);
 TORCH_CHECK(dt.size(0) == B && dt.size(1) == H && dt.size(2) == S,
 "dt must be (B,H,S) matching angles");
 auto out = torch::empty_like(angles);
 const int nseg = mamba3_cumulative_angles_segments(B, S, H);
 auto seg_tot = torch::empty({B * H, nseg, Na}, angles.options());
 mamba3_cumulative_angles_launch(
 angles.data_ptr<float>(), dt.data_ptr<float>(), out.data_ptr<float>(),
 seg_tot.data_ptr<float>(), B, S, H, Na, nseg,
 c10::cuda::getCurrentCUDAStream().stream());
 C10_CUDA_KERNEL_LAUNCH_CHECK();
 return out;
}
// Takes and returns no tensors, so it carries no meta kernel and belongs
// outside a compiled region; it exists so a caller, and the dispatch test, can
// ask which path a geometry takes without restating the gates.
std::vector<int64_t> mamba3_fwd_dispatch(
 int64_t mimo_rank, int64_t dstate, int64_t headdim, int64_t chunk_size, bool bf16)
{
 TORCH_CHECK(chunk_size > 0, "chunk_size must be positive");
 int state = 0, scan = 0, tf32 = 0;
 mamba3_fwd_dispatch_paths((int)mimo_rank, (int)dstate, (int)headdim,
 (int)chunk_size, bf16 ? 1 : 0, &state, &scan, &tf32);
 return {state, scan, tf32};
}
torch::Tensor mamba3_mimo_step(
 torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<at::Tensor> z,
 torch::Tensor adt, torch::Tensor dt, torch::Tensor trap,
 torch::Tensor q_bias, torch::Tensor k_bias, torch::Tensor angles,
 torch::Tensor mimo_v, torch::Tensor mimo_o,
 c10::optional<at::Tensor> mimo_z, c10::optional<at::Tensor> D,
 torch::Tensor angle_in, torch::Tensor angle_out, torch::Tensor S_st,
 torch::Tensor kprev_in, torch::Tensor vprev_in,
 torch::Tensor kprev_out, torch::Tensor vprev_out)
{
 const auto sdt = q.scalar_type();
 TORCH_CHECK(sdt == at::kFloat || sdt == at::kBFloat16,
 "inputs must be float32 or bfloat16");
 for (const auto& t : {q, k, v, q_bias, k_bias, mimo_v, mimo_o})
 TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == sdt,
 "all value inputs must be contiguous CUDA tensors of one dtype");
 M3_F32(adt); M3_F32(dt); M3_F32(trap); M3_F32(angles);
 M3_F32(angle_in); M3_F32(angle_out); M3_F32(S_st);
 M3_F32(kprev_in); M3_F32(vprev_in); M3_F32(kprev_out); M3_F32(vprev_out);
const int B = (int)q.size(0), R = (int)q.size(1), Gqk = (int)q.size(2), N = (int)q.size(3);
 const int H = (int)v.size(1), P = (int)v.size(2), Na = (int)angles.size(2);
TORCH_CHECK(N % 2 == 0, "state dimension N must be even");
 TORCH_CHECK(H % Gqk == 0, "nheads must be divisible by the query-group count");
 TORCH_CHECK(Na <= N / 2, "angle count must not exceed N/2");
 TORCH_CHECK(R <= 8, "mimo_rank above 8 is not supported");
 TORCH_CHECK(S_st.size(2) == P && S_st.size(3) == N, "state shape does not match inputs");
auto y = torch::empty({B, H, P}, v.options());
 auto fn = (sdt == at::kFloat) ? mamba3_mimo_step_launch_f32 : mamba3_mimo_step_launch_bf16;
 fn(q.data_ptr(), k.data_ptr(), v.data_ptr(), vptr(z, sdt),
 adt.data_ptr<float>(), dt.data_ptr<float>(), trap.data_ptr<float>(),
 q_bias.data_ptr(), k_bias.data_ptr(), angles.data_ptr<float>(),
 mimo_v.data_ptr(), mimo_o.data_ptr(), vptr(mimo_z, sdt), vptr(D, sdt),
 angle_in.data_ptr<float>(), angle_out.data_ptr<float>(), S_st.data_ptr<float>(),
 kprev_in.data_ptr<float>(), vprev_in.data_ptr<float>(),
 kprev_out.data_ptr<float>(), vprev_out.data_ptr<float>(), y.data_ptr(),
 B, H, Gqk, R, N, P, Na, c10::cuda::getCurrentCUDAStream().stream());
 C10_CUDA_KERNEL_LAUNCH_CHECK();
 return y;
}
#define FWD_ARGS \
 const void* q, const void* k, const void* v, const void* z, const void* q_bias, \
 const void* k_bias, const void* mimo_v, const void* mimo_o, const void* mimo_z, \
 const void* Dvec, const void* angles, const float* dt, const float* trap, \
 const float* dA_cs, const float* dA_cs_rev, const void* norm_w, float norm_eps, \
 int fused_norm, void* kv_ws, void* out, \
 int B, int S, int H, int Gqk, int R, int N, int P, int Na, int C, void* stream
void mamba3_mimo_fwd_launch_f32(FWD_ARGS);
void mamba3_mimo_fwd_launch_bf16(FWD_ARGS);
torch::Tensor mamba3_mimo_fwd(
 torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<at::Tensor> z,
 torch::Tensor q_bias, torch::Tensor k_bias, torch::Tensor mimo_v, torch::Tensor mimo_o,
 c10::optional<at::Tensor> mimo_z, c10::optional<at::Tensor> D,
 torch::Tensor angles, torch::Tensor dt, torch::Tensor trap,
 torch::Tensor dA_cs, torch::Tensor dA_cs_rev, int64_t chunk_size,
 c10::optional<at::Tensor> norm_weight, double norm_eps)
{
 const auto sdt = q.scalar_type();
 TORCH_CHECK(sdt == at::kFloat || sdt == at::kBFloat16,
 "inputs must be float32 or bfloat16");
 for (const auto& t : {q, k, v, q_bias, k_bias, mimo_v, mimo_o})
 TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == sdt,
 "all value inputs must be contiguous CUDA tensors of one dtype");
 // The schedule and the accumulated rotation angles stay float32 in both paths.
 M3_F32(dt); M3_F32(trap); M3_F32(dA_cs); M3_F32(dA_cs_rev); M3_F32(angles);
const int B = (int)q.size(0), S = (int)q.size(1), R = (int)q.size(2);
 const int Gqk = (int)q.size(3), N = (int)q.size(4);
 const int H = (int)v.size(2), P = (int)v.size(3), Na = (int)angles.size(3);
 const int C = (int)chunk_size;
 TORCH_CHECK(C > 0, "chunk_size must be positive");
 const int Nc = (S + C - 1) / C;
TORCH_CHECK(N % 2 == 0, "state dimension N must be even");
 TORCH_CHECK(H % Gqk == 0, "nheads must be divisible by the query-group count");
 TORCH_CHECK(Na <= N / 2, "angle count must not exceed N/2");
 // The tensor-core scan is instantiated for a fixed set of ranks; above it the
 // dispatch has no kernel to fall back to that carries the rank at runtime.
 TORCH_CHECK(R <= 8, "mimo_rank above 8 is not supported");
const size_t smem = mamba3_fwd_smem_bytes(R, N, P, C, sdt == at::kBFloat16 ? 1 : 0);
 const size_t cap = (size_t)mamba3_max_smem_optin();
 TORCH_CHECK(smem <= cap,
 "chunk_size ", C, " with mimo_rank ", R, ", headdim ", P, " and state ", N,
 " needs ", smem / 1024, " KB of shared memory per block, above this "
 "device's ", cap / 1024, " KB opt-in limit");
const int fused = (norm_weight.has_value() && norm_weight->defined()) ? 1 : 0;
 auto out = torch::empty({B, S, H, P}, v.options());
 // The state workspace carries the value dtype. chunk_scan casts the states
 // to bfloat16 the moment it reads them, so bf16 storage loses nothing there
 // and halves the traffic of the scan that walks this buffer.
 auto kv_ws = torch::empty({B, H, Nc, N, P}, v.options());
auto fn = (sdt == at::kFloat) ? mamba3_mimo_fwd_launch_f32 : mamba3_mimo_fwd_launch_bf16;
 fn(q.data_ptr(), k.data_ptr(), v.data_ptr(), vptr(z, sdt), q_bias.data_ptr(),
 k_bias.data_ptr(), mimo_v.data_ptr(), mimo_o.data_ptr(), vptr(mimo_z, sdt), vptr(D, sdt),
 angles.data_ptr(), dt.data_ptr<float>(), trap.data_ptr<float>(),
 dA_cs.data_ptr<float>(), dA_cs_rev.data_ptr<float>(), vptr(norm_weight, sdt),
 (float)norm_eps, fused, kv_ws.data_ptr(), out.data_ptr(),
 B, S, H, Gqk, R, N, P, Na, C, c10::cuda::getCurrentCUDAStream().stream());
 C10_CUDA_KERNEL_LAUNCH_CHECK();
 return out;
}
#define BWD_ARGS \
 const void* q, const void* k, const void* v, const void* z, const void* q_bias, \
 const void* k_bias, const void* mimo_v, const void* mimo_o, const void* mimo_z, \
 const void* Dvec, const float* angles, const float* dt, const float* trap, \
 const float* dA_cs, const float* dA_cs_rev, const float* states, const void* norm_w, \
 float norm_eps, int fused_norm, const void* dy, float* dO_ws, float* dSt_ws, float* ws, \
 float* dq, float* dk, float* dv, float* dz, float* dq_bias, float* dk_bias, \
 float* dmimo_v, float* dmimo_o, float* dmimo_z, float* dD, float* dnorm_w, \
 float* d_dA_cs, float* d_dA_cs_rev, float* ddt, float* dtrap, float* dang, \
 int B, int S, int H, int Gqk, int R, int N, int P, int Na, int C, void* stream
void mamba3_mimo_bwd_launch_f32(BWD_ARGS);
void mamba3_mimo_bwd_launch_bf16(BWD_ARGS);
void mamba3_mimo_states_launch_f32(
 const void* k, const void* v, const void* k_bias, const void* mimo_v,
 const float* angles, const float* dt, const float* trap,
 const float* dA_cs, const float* dA_cs_rev, float* kv_ws,
 int B, int S, int H, int Gqk, int R, int N, int P, int Na, int C, void* stream);
std::vector<at::Tensor> mamba3_mimo_bwd(
 torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<at::Tensor> z,
 torch::Tensor q_bias, torch::Tensor k_bias, torch::Tensor mimo_v, torch::Tensor mimo_o,
 c10::optional<at::Tensor> mimo_z, c10::optional<at::Tensor> D,
 torch::Tensor angles, torch::Tensor dt, torch::Tensor trap,
 torch::Tensor dA_cs, torch::Tensor dA_cs_rev, int64_t chunk_size,
 c10::optional<at::Tensor> norm_weight, double norm_eps, torch::Tensor dy)
{
 const auto sdt = q.scalar_type();
 TORCH_CHECK(sdt == at::kFloat || sdt == at::kBFloat16,
 "backward inputs must be float32 or bfloat16");
 // The kernels index these with contiguous strides; a strided view would be
 // read as if it were packed and would return wrong gradients silently.
 for (const auto& t : {q, k, v, q_bias, k_bias, mimo_v, mimo_o})
 TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == sdt,
 "all value inputs must be contiguous CUDA tensors of one dtype");
 M3_F32(dt); M3_F32(trap); M3_F32(dA_cs); M3_F32(dA_cs_rev); M3_F32(angles);
const int B = (int)q.size(0), S = (int)q.size(1), R = (int)q.size(2);
 const int Gqk = (int)q.size(3), N = (int)q.size(4);
 const int H = (int)v.size(2), P = (int)v.size(3), Na = (int)angles.size(3);
 TORCH_CHECK(chunk_size > 0, "chunk_size must be positive");
 TORCH_CHECK(R <= 8, "mimo_rank above 8 is not supported");
 const int C = (int)chunk_size, Nc = (S + C - 1) / C, CR = C * R;
// The backward has no row-block streamed path, so its tiles are resident and
 // its budget is tighter than the forward's at the same geometry.
 const size_t smem = std::max(mamba3_bwd_smem_bytes(R, N, P, C),
 mamba3_states_smem_bytes(R, N, P, C));
 const size_t cap = (size_t)mamba3_max_smem_optin();
 TORCH_CHECK(smem <= cap,
 "backward with chunk_size ", C, ", mimo_rank ", R, ", headdim ", P,
 " and state ", N, " needs ", smem / 1024, " KB of shared memory per "
 "block, above this device's ", cap / 1024, " KB opt-in limit; reduce "
 "chunk_size * mimo_rank or headdim");
auto f32 = v.options().dtype(at::kFloat);
 // The state recompute runs in float32 regardless; it feeds the inter term.
 auto kf = k.to(at::kFloat).contiguous(), vf = v.to(at::kFloat).contiguous();
 auto kbf = k_bias.to(at::kFloat).contiguous(), mvf = mimo_v.to(at::kFloat).contiguous();
 auto states = torch::empty({B, H, Nc, N, P}, f32);
 mamba3_mimo_states_launch_f32(
 kf.data_ptr(), vf.data_ptr(), kbf.data_ptr(), mvf.data_ptr(),
 angles.data_ptr<float>(), dt.data_ptr<float>(), trap.data_ptr<float>(),
 dA_cs.data_ptr<float>(), dA_cs_rev.data_ptr<float>(), states.data_ptr<float>(),
 B, S, H, Gqk, R, N, P, Na, C, c10::cuda::getCurrentCUDAStream().stream());
 C10_CUDA_KERNEL_LAUNCH_CHECK();
auto dO_ws = torch::empty({B, H, Nc, CR, P}, f32);
 auto dSt_ws = torch::empty({B, H, Nc, N, P}, f32);
 // Per-block slots for the parameter gradients, reduced in block order so the
 // result does not depend on scheduling.
 const int64_t nblk = (int64_t)B * Nc;
 const int64_t gy = ((int64_t)N * P + 255) / 256;
 const int64_t nws = nblk * ((int64_t)3 * H * R * P + 2 * H * R * N + H * P + H)
 + 2LL * B * H * S + (int64_t)B * H * Nc * gy;
 auto ws = torch::empty({nws}, f32);
auto dq = torch::zeros_like(q, f32), dk = torch::zeros_like(k, f32);
 auto dv = torch::zeros_like(v, f32), dz = torch::zeros_like(v, f32);
 auto dqb = torch::zeros_like(q_bias, f32), dkb = torch::zeros_like(k_bias, f32);
 auto dmv = torch::zeros_like(mimo_v, f32), dmo = torch::zeros_like(mimo_o, f32);
 auto dmz = torch::zeros_like(mimo_v, f32);
 auto dD = torch::zeros({H}, f32), dnw = torch::zeros({H, P}, f32);
 auto dacs = torch::zeros_like(dA_cs), dacsr = torch::zeros_like(dA_cs_rev);
 auto ddt = torch::zeros_like(dt), dtrap = torch::zeros_like(trap);
 auto dang = torch::zeros_like(angles);
const int fused = (norm_weight.has_value() && norm_weight->defined()) ? 1 : 0;
 auto bwd = (sdt == at::kFloat) ? mamba3_mimo_bwd_launch_f32 : mamba3_mimo_bwd_launch_bf16;
 bwd(q.data_ptr(), k.data_ptr(), v.data_ptr(), vptr(z, sdt), q_bias.data_ptr(),
 k_bias.data_ptr(), mimo_v.data_ptr(), mimo_o.data_ptr(), vptr(mimo_z, sdt),
 vptr(D, sdt), angles.data_ptr<float>(), dt.data_ptr<float>(),
 trap.data_ptr<float>(), dA_cs.data_ptr<float>(), dA_cs_rev.data_ptr<float>(),
 states.data_ptr<float>(), vptr(norm_weight, sdt), (float)norm_eps, fused,
 dy.contiguous().data_ptr(), dO_ws.data_ptr<float>(), dSt_ws.data_ptr<float>(),
 ws.data_ptr<float>(),
 dq.data_ptr<float>(), dk.data_ptr<float>(), dv.data_ptr<float>(),
 dz.data_ptr<float>(), dqb.data_ptr<float>(), dkb.data_ptr<float>(),
 dmv.data_ptr<float>(), dmo.data_ptr<float>(), dmz.data_ptr<float>(),
 dD.data_ptr<float>(), dnw.data_ptr<float>(), dacs.data_ptr<float>(),
 dacsr.data_ptr<float>(), ddt.data_ptr<float>(), dtrap.data_ptr<float>(),
 dang.data_ptr<float>(),
 B, S, H, Gqk, R, N, P, Na, C, c10::cuda::getCurrentCUDAStream().stream());
 C10_CUDA_KERNEL_LAUNCH_CHECK();
return {dq, dk, dv, dz, dqb, dkb, dmv, dmo, dmz, dD, dnw, dacs, dacsr, ddt, dtrap, dang};
}
// registration.h comes from the kernel builder; a local JIT build has neither
// it nor TORCH_LIBRARY_EXPAND, and reaches the ops through torch.ops instead.
#if defined(__has_include) && __has_include("registration.h")
# include "registration.h"
# define MAMBA3_HAVE_REGISTRATION 1
#else
# define TORCH_LIBRARY_EXPAND(NAME, MOD) TORCH_LIBRARY(NAME, MOD)
# define CUDA_KERNEL 1
#endif
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 // Defined with its implementation: no tensor arguments, so there is no
 // dispatch key to select on and no meta kernel to register.
 ops.def("mamba3_fwd_dispatch(int mimo_rank, int dstate, int headdim, "
 "int chunk_size, bool bf16) -> int[]", &mamba3_fwd_dispatch);
 ops.def("mamba3_cumulative_angles(Tensor angles, Tensor dt) -> Tensor");
 ops.def("mamba3_mimo_step(Tensor q, Tensor k, Tensor v, Tensor? z, Tensor adt, "
 "Tensor dt, Tensor trap, Tensor q_bias, Tensor k_bias, Tensor angles, "
 "Tensor mimo_v, Tensor mimo_o, Tensor? mimo_z, Tensor? D, "
 "Tensor angle_in, Tensor(a!) angle_out, Tensor(b!) S_st, "
 "Tensor kprev_in, Tensor vprev_in, Tensor(c!) kprev_out, "
 "Tensor(d!) vprev_out) -> Tensor");
 ops.def("mamba3_mimo_fwd(Tensor q, Tensor k, Tensor v, Tensor? z, Tensor q_bias, "
 "Tensor k_bias, Tensor mimo_v, Tensor mimo_o, Tensor? mimo_z, Tensor? D, "
 "Tensor angles, Tensor dt, Tensor trap, Tensor dA_cs, Tensor dA_cs_rev, "
 "int chunk_size, Tensor? norm_weight, float norm_eps) -> Tensor");
ops.def("mamba3_mimo_bwd(Tensor q, Tensor k, Tensor v, Tensor? z, Tensor q_bias, "
 "Tensor k_bias, Tensor mimo_v, Tensor mimo_o, Tensor? mimo_z, Tensor? D, "
 "Tensor angles, Tensor dt, Tensor trap, Tensor dA_cs, Tensor dA_cs_rev, "
 "int chunk_size, Tensor? norm_weight, float norm_eps, Tensor dy) -> Tensor[]");
#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
 ops.impl("mamba3_cumulative_angles", torch::kCUDA, &mamba3_cumulative_angles);
 ops.impl("mamba3_mimo_step", torch::kCUDA, &mamba3_mimo_step);
 ops.impl("mamba3_mimo_fwd", torch::kCUDA, &mamba3_mimo_fwd);
 ops.impl("mamba3_mimo_bwd", torch::kCUDA, &mamba3_mimo_bwd);
#endif
}
#ifdef MAMBA3_HAVE_REGISTRATION
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
#endif