ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: vectorize fp8 quantize and swizzle gemm smem

Browse Source
This commit is contained in:
ViperEkura
2026-08-23 20:31:44 +08:00
parent 2bc4d2b8a8
commit 4b10d3ca37
7 changed files with 387 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/ops.cu
+99 -56
View File
@@ -58,7 +58,7 @@ void check_scale(const torch::Tensor& scale, const torch::Tensor& input,
void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out,
const torch::Tensor& sa, const torch::Tensor& sb,
const torch::Tensor* out_scale, int64_t m, int64_t n,
int64_t k) {
int64_t k, int64_t a_ld, int64_t b_ld) {
p.a_ptr = a;
p.b_ptr = b;
p.out_ptr = out;
@@ -68,9 +68,11 @@ void pack_gemm_params(FP8Params& p, const void* a, const void* b, void* out,
p.bias = nullptr;
p.amax_a = nullptr;
p.amax_b = nullptr;
p.m = m;
p.n = n;
p.k = k;
p.m = static_cast<int>(m);
p.n = static_cast<int>(n);
p.k = static_cast<int>(k);
p.a_ld = static_cast<int>(a_ld);
p.b_ld = static_cast<int>(b_ld);
p.total = 0;
}
@@ -87,7 +89,39 @@ void pack_quantize_params(FP8Params& p, const void* x, void* x8,
p.amax_a = amax ? amax->data_ptr<float>() : nullptr;
p.amax_b = nullptr;
p.m = p.n = p.k = 0;
p.total = total;
p.a_ld = p.b_ld = 0;
p.total = static_cast<int>(total);
}
// ---- GEMM launch dispatch (runtime flags -> compile-time kernel variants) ----
template <FP8Format Fmt, int Variant>
void launch_gemm_variant(const FP8Params& p, cudaStream_t stream) {
static_assert(Variant >= 0 && Variant < 8,
"invalid FP8 GEMM dispatch variant");
constexpr bool out_fp8 = (Variant & 4) != 0;
constexpr bool trans_a = (Variant & 2) != 0;
constexpr bool trans_b = (Variant & 1) != 0;
fp8::launch_fp8_gemm<Fmt, out_fp8, trans_a, trans_b>(p, stream);
}
template <FP8Format Fmt>
void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool out_fp8,
bool trans_a, bool trans_b) {
// Encode the runtime flags as [output FP8, transpose A, transpose B].
const int variant = (static_cast<int>(out_fp8) << 2) |
(static_cast<int>(trans_a) << 1) |
static_cast<int>(trans_b);
switch (variant) {
case 0: launch_gemm_variant<Fmt, 0>(p, stream); break;
case 1: launch_gemm_variant<Fmt, 1>(p, stream); break;
case 2: launch_gemm_variant<Fmt, 2>(p, stream); break;
case 3: launch_gemm_variant<Fmt, 3>(p, stream); break;
case 4: launch_gemm_variant<Fmt, 4>(p, stream); break;
case 5: launch_gemm_variant<Fmt, 5>(p, stream); break;
case 6: launch_gemm_variant<Fmt, 6>(p, stream); break;
case 7: launch_gemm_variant<Fmt, 7>(p, stream); break;
}
}
} // namespace
@@ -127,11 +161,13 @@ std::tuple<torch::Tensor, torch::Tensor> quantize_bf16(torch::Tensor x,
torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
torch::Tensor sb, int64_t out_dtype,
c10::optional<torch::Tensor> out_scale) {
// Pre-quantized FP8 GEMM: out = a @ b^T * (sa * sb), FP32 accumulation.
c10::optional<torch::Tensor> out_scale, int64_t trans_a,
int64_t trans_b) {
// Pre-quantized FP8 GEMM: out = op(a) @ op(b)^T * (sa * sb), FP32 accum.
// trans_a / trans_b select the operand layout (0 = stored [M,K]/[K,N],
// 1 = transposed [K,M]/[N,K]); the default (0/0) is the plain a @ b.
// out_dtype: 0 = BF16 (default), 1 = FP8 E4M3 (requires out_scale, the
// quantization step for the output — mirrors torch._scaled_mm's
// out_dtype / scale_result). Both operands share the same FP8 format.
// output quantization step). Both operands share one format.
TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn ||
a.scalar_type() == torch::kFloat8_e5m2,
@@ -140,7 +176,6 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
"a and b must share the same fp8 format");
TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D");
TORCH_CHECK(a.device() == b.device(), "a and b must be on the same device");
TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch");
check_scale(sa, a, "sa");
check_scale(sb, a, "sb");
check_fp8_device(a);
@@ -149,7 +184,16 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
auto a_c = a.contiguous();
auto b_c = b.contiguous();
int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0);
const bool ta = (trans_a == 1), tb = (trans_b == 1);
// Physical leading dimension = column count of each contiguous buffer.
const int64_t a_ld = a_c.size(1);
const int64_t b_ld = b_c.size(1);
// Logical GEMM shape derived from the layout flags.
const int64_t m = ta ? a_c.size(1) : a_c.size(0);
const int64_t k = ta ? a_c.size(0) : a_c.size(1);
const int64_t n = tb ? b_c.size(0) : b_c.size(1);
const int64_t k2 = tb ? b_c.size(1) : b_c.size(0);
TORCH_CHECK(k == k2, "inner dim mismatch");
const bool out_fp8 = (out_dtype == 1);
TORCH_CHECK(out_dtype == 0 || out_fp8,
"out_dtype must be 0 (bf16) or 1 (fp8 e4m3)");
@@ -165,20 +209,11 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa,
: a_c.options().dtype(torch::kBFloat16));
FP8Params p;
pack_gemm_params(p, a_c.data_ptr(), b_c.data_ptr(), out.data_ptr(), sa, sb,
out_fp8 ? &os : nullptr, m, n, k);
if (a.scalar_type() == torch::kFloat8_e4m3fn) {
if (out_fp8) {
fp8::launch_fp8_gemm<FP8Format::E4M3, true>(p, stream.stream());
} else {
fp8::launch_fp8_gemm<FP8Format::E4M3>(p, stream.stream());
}
} else {
if (out_fp8) {
fp8::launch_fp8_gemm<FP8Format::E5M2, true>(p, stream.stream());
} else {
fp8::launch_fp8_gemm<FP8Format::E5M2>(p, stream.stream());
}
}
out_fp8 ? &os : nullptr, m, n, k, a_ld, b_ld);
if (a.scalar_type() == torch::kFloat8_e4m3fn)
dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), out_fp8, ta, tb);
else
dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), out_fp8, ta, tb);
C10_CUDA_CHECK(cudaGetLastError());
return out;
}
@@ -233,12 +268,16 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8(
quantize(w_c, w8, sw, &amax_w);
FP8Params p;
// Forward is the NT layout: A = x8 [M,K] (a_ld = k), B = w8 [N,K]
// (b_ld = k), out = x @ w^T. No operand transposes needed.
pack_gemm_params(p, x8.data_ptr(), w8.data_ptr(), out.data_ptr(), sx, sw,
nullptr, m, n, k);
nullptr, m, n, k, k, k);
if (fmt) {
fp8::launch_fp8_gemm<FP8Format::E5M2>(p, stream.stream());
fp8::launch_fp8_gemm<FP8Format::E5M2, false, false, true>(
p, stream.stream());
} else {
fp8::launch_fp8_gemm<FP8Format::E4M3>(p, stream.stream());
fp8::launch_fp8_gemm<FP8Format::E4M3, false, false, true>(
p, stream.stream());
}
C10_CUDA_CHECK(cudaGetLastError());
@@ -293,24 +332,19 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
fp8::launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
}
};
// Explicit-transpose backward: the gradient/activation tensors keep their
// natural row-major layout, which the GEMM consumes transposed (W is
// [N,K] but dX contracts over N; x is [M,K] and g is [M,N] for dW), so
// the fp8 operands are transposed once and run through the fast non-trans
// pre-quantized GEMM. g is quantized once (amax_g measured here); its
// transpose is derived from the same g8 so both GEMMs share the value.
auto pq_n = [&](const torch::Tensor& a8, const torch::Tensor& b8,
torch::Tensor& out, const torch::Tensor& sa,
const torch::Tensor& sb, int64_t mm, int64_t nn,
int64_t kk) {
FP8Params gp;
pack_gemm_params(gp, a8.data_ptr(), b8.data_ptr(), out.data_ptr(), sa,
sb, nullptr, mm, nn, kk);
if (fmt) {
fp8::launch_fp8_gemm<FP8Format::E5M2>(gp, stream.stream());
} else {
fp8::launch_fp8_gemm<FP8Format::E4M3>(gp, stream.stream());
}
// Four-layout backward: the gradient and activation tensors keep their
// natural row-major layout, and the kernel reads them transposed where the
// GEMM needs it (TransA / TransB). No torch-level `.transpose().contiguous()`
// copies are required — dX uses g8 [M,N] as A with w8 [N,K] read transposed
// as B; dW uses g8 transposed as A with x8 transposed as B.
// g is quantized once (amax_g measured here); both GEMMs share g8.
auto run_bwd_gemm = [&](const FP8Params& gp, bool trans_a, bool trans_b) {
if (fmt)
dispatch_gemm<FP8Format::E5M2>(gp, stream.stream(), false, trans_a,
trans_b);
else
dispatch_gemm<FP8Format::E4M3>(gp, stream.stream(), false, trans_a,
trans_b);
};
torch::Tensor g8;
@@ -318,21 +352,28 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
g8 = torch::empty({m, n}, f8opt);
quantize(g_c, g8, sg, &amax_g);
}
// dX = g @ W: A = g8 [M,N] natural; B = W^T [K,N] (w8 transposed in fp8).
// dX = g @ w: A = g8 [M,N] (contract over N), B = w8 [N,K] read transposed
// (b[p*b_ld + n] = w[p,n]); out = [M,K], a_ld = N, b_ld = K, contract = N.
if (masks[0]) {
auto w8 = torch::empty({n, k}, f8opt);
quantize(w_c, w8, sw, nullptr);
auto w8T = w8.transpose(0, 1).contiguous(); // [K, N]
auto grad_input_2d = grad_input.reshape({m, k});
pq_n(g8, w8T, grad_input_2d, sg, sw, m, k, n);
FP8Params gp;
pack_gemm_params(gp, g8.data_ptr(), w8.data_ptr(),
grad_input_2d.data_ptr(), sg, sw, nullptr, m, k, n, n,
k);
run_bwd_gemm(gp, false, false);
}
// dW = g^T @ x: A = g^T [N,M] (g8 transposed); B = x^T [K,M].
// dW = g^T @ x: A = g8 [M,N] read transposed (a[p*a_ld + m] = g[p,m]), B =
// x8 [M,K] read transposed (b[p*b_ld + n] = x[p,n]); out = [N,K], a_ld = N,
// b_ld = K, contract = M.
if (masks[1]) {
auto g8T = g8.transpose(0, 1).contiguous(); // [N, M]
auto x8 = torch::empty({m, k}, f8opt);
quantize(x_c, x8, sx, nullptr);
auto x8T = x8.transpose(0, 1).contiguous(); // [K, M]
pq_n(g8T, x8T, grad_weight, sg, sx, n, k, m);
FP8Params gp;
pack_gemm_params(gp, g8.data_ptr(), x8.data_ptr(),
grad_weight.data_ptr(), sg, sx, nullptr, n, k, m, n, k);
run_bwd_gemm(gp, true, false);
}
if (!masks[0] && !masks[1]) {
amax_g.copy_(g_c.abs().amax().to(torch::kFloat32));
@@ -348,9 +389,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
"BF16 to FP8 (E4M3/E5M2) quantize with fused amax; returns (x8, amax)");
m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("sa"),
py::arg("sb"), py::arg("out_dtype") = 0,
py::arg("out_scale") = py::none(),
"Pre-quantized FP8 GEMM: a @ b^T * (sa * sb); out_dtype 0=bf16, "
"1=fp8 e4m3 (requires out_scale)");
py::arg("out_scale") = py::none(), py::arg("trans_a") = 0,
py::arg("trans_b") = 0,
"Pre-quantized FP8 GEMM: op(a) @ op(b)^T * (sa * sb); out_dtype "
"0=bf16, 1=fp8 e4m3 (requires out_scale); trans_a/trans_b select "
"the operand layout (default 0/0 = a@b)");
m.def("linear_forward_fp8", &linear_forward_fp8, py::arg("x"),
py::arg("w"), py::arg("bias"), py::arg("sx"), py::arg("sw"),
py::arg("fmt") = 0,