bitnet-cpu v1: kernel-builder source tree

ce8679e verified 26 days ago

20.1 kB

	// Ternary x INT8 GEMM for BitNet b1.58 (W1.58 A8) on CPUs.
	//
	// Operand layout matches phanerozoic/bitnet-tc (CUDA):
	// A: [M, K] int8 activations (per-token absmax quantized)
	// B: [N, K/4] uint8 packed ternary weights, 4 codes per byte.
	// Code encoding: 1->-1, 2->0, 3->+1 (weight = code - 2).
	// scale_act: [M] bf16 per-row activation scale
	// scale_wt: [N] bf16 per-column weight scale
	// out: [M, N] bf16/f32 = (A @ W^T) * scale_act[:,None] * scale_wt[None,:]
	//
	// Compute identity: with u = code - 1 in {0,1,2} (unsigned) and signed a,
	// dot(w, a) = dot(u, a) - sum(a)
	// so the ternary product maps onto unsigned x signed multiply-accumulate
	// instructions (vpdpbusd on AVX-512 VNNI / AVX-VNNI, vpmaddubsw on AVX2)
	// with no saturation risk (u <= 2, \|a\| <= 127). sum(a) is one int32 per
	// activation row, computed once. The ISA path is selected once at runtime
	// from cpuid; unsupported machines fall back to a portable scalar loop.

	#include <ATen/Parallel.h>
	#include <torch/all.h>

	#include <algorithm>
	#include <cmath>
	#include <cstdint>
	#include <cstring>
	#include <memory>

	#if defined(__x86_64__) \|\| defined(_M_X64)
	#define BITNET_X86 1
	#include <cpuid.h>
	#include <immintrin.h>
	#else
	#define BITNET_X86 0
	#endif

	namespace {

	// ---------------------------------------------------------------------------
	// 2-bit unpack: packed byte -> 4 u-codes in {0,1,2} (u = code - 1).
	// ---------------------------------------------------------------------------

	struct U4Lut {
	uint32_t v[256];
	U4Lut() {
	for (int b = 0; b < 256; b++) {
	uint32_t out = 0;
	for (int j = 0; j < 4; j++) {
	uint32_t code = (static_cast<uint32_t>(b) >> (2 * j)) & 3u;
	uint32_t u = code == 0 ? 0u : code - 1u; // code 0 is invalid; map safely
	out \|= u << (8 * j);
	}
	v[b] = out;
	}
	}
	};
	const U4Lut u4_lut;

	inline void unpack_row(const uint8_t* wrow, int64_t k4, uint8_t* dst) {
	uint32_t* d = reinterpret_cast<uint32_t*>(dst);
	for (int64_t i = 0; i < k4; i++) {
	d[i] = u4_lut.v[wrow[i]];
	}
	}

	// ---------------------------------------------------------------------------
	// Runtime ISA dispatch.
	// ---------------------------------------------------------------------------

	enum class Isa { Scalar, Avx2, AvxVnni, Avx512Vnni };

	#if BITNET_X86
	inline uint64_t xgetbv0() {
	uint32_t eax, edx;
	__asm__ volatile("xgetbv" : "=a"(eax), "=d"(edx) : "c"(0));
	return (static_cast<uint64_t>(edx) << 32) \| eax;
	}

	Isa detect_isa() {
	uint32_t eax, ebx, ecx, edx;
	if (!__get_cpuid(1, &eax, &ebx, &ecx, &edx)) return Isa::Scalar;
	const bool osxsave = (ecx >> 27) & 1;
	if (!osxsave) return Isa::Scalar;
	const uint64_t xcr0 = xgetbv0();
	const bool ymm_os = (xcr0 & 0x6) == 0x6;
	const bool zmm_os = (xcr0 & 0xe6) == 0xe6;
	if (!ymm_os) return Isa::Scalar;

	uint32_t max_leaf = __get_cpuid_max(0, nullptr);
	if (max_leaf < 7) return Isa::Scalar;
	uint32_t eax7, ebx7, ecx7, edx7;
	__cpuid_count(7, 0, eax7, ebx7, ecx7, edx7);
	const bool avx2 = (ebx7 >> 5) & 1;
	const bool avx512f = (ebx7 >> 16) & 1;
	const bool avx512bw = (ebx7 >> 30) & 1;
	const bool avx512vl = (ebx7 >> 31) & 1;
	const bool avx512vnni = (ecx7 >> 11) & 1;
	if (zmm_os && avx512f && avx512bw && avx512vl && avx512vnni) return Isa::Avx512Vnni;

	bool avxvnni = false;
	if (eax7 >= 1) {
	uint32_t eax71, ebx71, ecx71, edx71;
	__cpuid_count(7, 1, eax71, ebx71, ecx71, edx71);
	avxvnni = (eax71 >> 4) & 1;
	}
	if (avx2 && avxvnni) return Isa::AvxVnni;
	if (avx2) return Isa::Avx2;
	return Isa::Scalar;
	}
	#else
	Isa detect_isa() { return Isa::Scalar; }
	#endif

	Isa isa() {
	static const Isa v = detect_isa();
	return v;
	}

	// ---------------------------------------------------------------------------
	// Dot kernels: dot(u, a) over K elements, 1 and 4 weight rows at a time.
	// u in {0,1,2} (unsigned), a signed int8. K is a multiple of 32.
	// ---------------------------------------------------------------------------

	int32_t dot1_scalar(const uint8_t* u, const int8_t* a, int64_t K) {
	int32_t acc = 0;
	for (int64_t k = 0; k < K; k++) acc += static_cast<int32_t>(u[k]) * static_cast<int32_t>(a[k]);
	return acc;
	}

	void dot4_scalar(const uint8_t* const* u, const int8_t* a, int64_t K, int32_t* out) {
	for (int j = 0; j < 4; j++) out[j] = dot1_scalar(u[j], a, K);
	}

	#if BITNET_X86

	__attribute__((target("avx2"))) inline int32_t hsum256_i32(__m256i v) {
	__m128i lo = _mm256_castsi256_si128(v);
	__m128i hi = _mm256_extracti128_si256(v, 1);
	__m128i s = _mm_add_epi32(lo, hi);
	s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0x4e));
	s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0xb1));
	return _mm_cvtsi128_si32(s);
	}

	// ---- AVX2 (vpmaddubsw) ----

	__attribute__((target("avx2"))) int32_t dot1_avx2(const uint8_t* u, const int8_t* a, int64_t K) {
	const __m256i ones = _mm256_set1_epi16(1);
	__m256i acc = _mm256_setzero_si256();
	for (int64_t k = 0; k < K; k += 32) {
	__m256i uv = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u + k));
	__m256i av = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(a + k));
	__m256i p16 = _mm256_maddubs_epi16(uv, av);
	acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
	}
	return hsum256_i32(acc);
	}

	__attribute__((target("avx2"))) void dot4_avx2(const uint8_t* const* u, const int8_t* a, int64_t K,
	int32_t* out) {
	const __m256i ones = _mm256_set1_epi16(1);
	__m256i acc0 = _mm256_setzero_si256(), acc1 = _mm256_setzero_si256();
	__m256i acc2 = _mm256_setzero_si256(), acc3 = _mm256_setzero_si256();
	for (int64_t k = 0; k < K; k += 32) {
	__m256i av = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(a + k));
	__m256i u0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[0] + k));
	__m256i u1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[1] + k));
	__m256i u2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[2] + k));
	__m256i u3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[3] + k));
	acc0 = _mm256_add_epi32(acc0, _mm256_madd_epi16(_mm256_maddubs_epi16(u0, av), ones));
	acc1 = _mm256_add_epi32(acc1, _mm256_madd_epi16(_mm256_maddubs_epi16(u1, av), ones));
	acc2 = _mm256_add_epi32(acc2, _mm256_madd_epi16(_mm256_maddubs_epi16(u2, av), ones));
	acc3 = _mm256_add_epi32(acc3, _mm256_madd_epi16(_mm256_maddubs_epi16(u3, av), ones));
	}
	out[0] = hsum256_i32(acc0);
	out[1] = hsum256_i32(acc1);
	out[2] = hsum256_i32(acc2);
	out[3] = hsum256_i32(acc3);
	}

	// ---- AVX-VNNI (256-bit vpdpbusd on client parts without AVX-512) ----

	__attribute__((target("avxvnni"))) int32_t dot1_avxvnni(const uint8_t* u, const int8_t* a,
	int64_t K) {
	__m256i acc = _mm256_setzero_si256();
	for (int64_t k = 0; k < K; k += 32) {
	__m256i uv = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u + k));
	__m256i av = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(a + k));
	acc = _mm256_dpbusd_avx_epi32(acc, uv, av);
	}
	return hsum256_i32(acc);
	}

	__attribute__((target("avxvnni"))) void dot4_avxvnni(const uint8_t* const* u, const int8_t* a,
	int64_t K, int32_t* out) {
	__m256i acc0 = _mm256_setzero_si256(), acc1 = _mm256_setzero_si256();
	__m256i acc2 = _mm256_setzero_si256(), acc3 = _mm256_setzero_si256();
	for (int64_t k = 0; k < K; k += 32) {
	__m256i av = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(a + k));
	acc0 = _mm256_dpbusd_avx_epi32(acc0, _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[0] + k)), av);
	acc1 = _mm256_dpbusd_avx_epi32(acc1, _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[1] + k)), av);
	acc2 = _mm256_dpbusd_avx_epi32(acc2, _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[2] + k)), av);
	acc3 = _mm256_dpbusd_avx_epi32(acc3, _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[3] + k)), av);
	}
	out[0] = hsum256_i32(acc0);
	out[1] = hsum256_i32(acc1);
	out[2] = hsum256_i32(acc2);
	out[3] = hsum256_i32(acc3);
	}

	// ---- AVX-512 VNNI ----

	__attribute__((target("avx512f,avx512bw,avx512vl,avx512vnni"))) int32_t dot1_avx512(
	const uint8_t* u, const int8_t* a, int64_t K) {
	__m512i acc = _mm512_setzero_si512();
	int64_t k = 0;
	for (; k + 64 <= K; k += 64) {
	__m512i uv = _mm512_loadu_si512(u + k);
	__m512i av = _mm512_loadu_si512(a + k);
	acc = _mm512_dpbusd_epi32(acc, uv, av);
	}
	int32_t total = _mm512_reduce_add_epi32(acc);
	if (k < K) { // K % 64 == 32 tail
	__m256i acc2 = _mm256_setzero_si256();
	__m256i uv = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u + k));
	__m256i av = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(a + k));
	acc2 = _mm256_dpbusd_epi32(acc2, uv, av);
	total += hsum256_i32(acc2);
	}
	return total;
	}

	__attribute__((target("avx512f,avx512bw,avx512vl,avx512vnni"))) void dot4_avx512(
	const uint8_t* const* u, const int8_t* a, int64_t K, int32_t* out) {
	__m512i acc0 = _mm512_setzero_si512(), acc1 = _mm512_setzero_si512();
	__m512i acc2 = _mm512_setzero_si512(), acc3 = _mm512_setzero_si512();
	int64_t k = 0;
	for (; k + 64 <= K; k += 64) {
	__m512i av = _mm512_loadu_si512(a + k);
	acc0 = _mm512_dpbusd_epi32(acc0, _mm512_loadu_si512(u[0] + k), av);
	acc1 = _mm512_dpbusd_epi32(acc1, _mm512_loadu_si512(u[1] + k), av);
	acc2 = _mm512_dpbusd_epi32(acc2, _mm512_loadu_si512(u[2] + k), av);
	acc3 = _mm512_dpbusd_epi32(acc3, _mm512_loadu_si512(u[3] + k), av);
	}
	out[0] = _mm512_reduce_add_epi32(acc0);
	out[1] = _mm512_reduce_add_epi32(acc1);
	out[2] = _mm512_reduce_add_epi32(acc2);
	out[3] = _mm512_reduce_add_epi32(acc3);
	if (k < K) {
	__m256i av = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(a + k));
	for (int j = 0; j < 4; j++) {
	__m256i accT = _mm256_setzero_si256();
	__m256i uv = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u[j] + k));
	accT = _mm256_dpbusd_epi32(accT, uv, av);
	out[j] += hsum256_i32(accT);
	}
	}
	}

	#endif // BITNET_X86

	int32_t dot1(const uint8_t* u, const int8_t* a, int64_t K) {
	#if BITNET_X86
	switch (isa()) {
	case Isa::Avx512Vnni: return dot1_avx512(u, a, K);
	case Isa::AvxVnni: return dot1_avxvnni(u, a, K);
	case Isa::Avx2: return dot1_avx2(u, a, K);
	default: break;
	}
	#endif
	return dot1_scalar(u, a, K);
	}

	void dot4(const uint8_t* const* u, const int8_t* a, int64_t K, int32_t* out) {
	#if BITNET_X86
	switch (isa()) {
	case Isa::Avx512Vnni: dot4_avx512(u, a, K, out); return;
	case Isa::AvxVnni: dot4_avxvnni(u, a, K, out); return;
	case Isa::Avx2: dot4_avx2(u, a, K, out); return;
	default: break;
	}
	#endif
	dot4_scalar(u, a, K, out);
	}

	// ---------------------------------------------------------------------------
	// GEMM core: out[m, n] = (dot(u_n, a_m) - act_sum[m]) * sa[m] * sw[n]
	// ---------------------------------------------------------------------------

	template <typename OutT>
	void gemm_core(const int8_t* a, const int32_t* act_sum, const float* sa, int64_t M, int64_t K,
	const uint8_t* w, const float* sw, int64_t N, OutT* out) {
	const int64_t k4 = K / 4;
	const int64_t nblocks = (N + 3) / 4;
	at::parallel_for(0, nblocks, 1, [&](int64_t begin, int64_t end) {
	std::unique_ptr<uint8_t[]> ubuf(new uint8_t[4 * K]);
	uint8_t* uptr[4] = {ubuf.get(), ubuf.get() + K, ubuf.get() + 2 * K, ubuf.get() + 3 * K};
	for (int64_t blk = begin; blk < end; blk++) {
	const int64_t n0 = blk * 4;
	const int rows = static_cast<int>(std::min<int64_t>(4, N - n0));
	for (int j = 0; j < rows; j++) unpack_row(w + (n0 + j) * k4, k4, uptr[j]);
	if (rows == 4) {
	const uint8_t* uc[4] = {uptr[0], uptr[1], uptr[2], uptr[3]};
	int32_t d[4];
	for (int64_t m = 0; m < M; m++) {
	dot4(uc, a + m * K, K, d);
	const float s = sa[m];
	const int32_t asum = act_sum[m];
	OutT* orow = out + m * N + n0;
	for (int j = 0; j < 4; j++) {
	orow[j] = static_cast<OutT>(static_cast<float>(d[j] - asum) * s * sw[n0 + j]);
	}
	}
	} else {
	for (int64_t m = 0; m < M; m++) {
	const float s = sa[m];
	const int32_t asum = act_sum[m];
	OutT* orow = out + m * N + n0;
	for (int j = 0; j < rows; j++) {
	int32_t d = dot1(uptr[j], a + m * K, K);
	orow[j] = static_cast<OutT>(static_cast<float>(d - asum) * s * sw[n0 + j]);
	}
	}
	}
	}
	});
	}

	// ---------------------------------------------------------------------------
	// Quantization helpers.
	// ---------------------------------------------------------------------------

	template <typename InT>
	void quantize_rows(const InT* x, int64_t M, int64_t K, int8_t* q, float* scale_out,
	int32_t* sum_out) {
	at::parallel_for(0, M, 1, [&](int64_t begin, int64_t end) {
	for (int64_t m = begin; m < end; m++) {
	const InT* row = x + m * K;
	float amax = 0.f;
	for (int64_t k = 0; k < K; k++) {
	float v = std::fabs(static_cast<float>(row[k]));
	amax = std::max(amax, v);
	}
	amax = std::max(amax, 1e-5f);
	const float scale = amax / 127.0f;
	const float inv = 127.0f / amax;
	int8_t* qrow = q + m * K;
	int32_t s = 0;
	for (int64_t k = 0; k < K; k++) {
	float v = static_cast<float>(row[k]) * inv;
	int32_t iv = static_cast<int32_t>(std::nearbyint(v));
	iv = std::min(127, std::max(-127, iv));
	qrow[k] = static_cast<int8_t>(iv);
	s += iv;
	}
	scale_out[m] = scale;
	if (sum_out) sum_out[m] = s;
	}
	});
	}

	void row_sums(const int8_t* a, int64_t M, int64_t K, int32_t* sums) {
	at::parallel_for(0, M, 1, [&](int64_t begin, int64_t end) {
	for (int64_t m = begin; m < end; m++) {
	const int8_t* row = a + m * K;
	int32_t s = 0;
	for (int64_t k = 0; k < K; k++) s += row[k];
	sums[m] = s;
	}
	});
	}

	void check_weights(const torch::Tensor& w_packed, int64_t K) {
	TORCH_CHECK(w_packed.device().is_cpu(), "w_packed must be a CPU tensor");
	TORCH_CHECK(w_packed.dtype() == torch::kUInt8, "w_packed must be uint8");
	TORCH_CHECK(w_packed.dim() == 2, "w_packed must be 2D [N, K/4]");
	TORCH_CHECK(w_packed.is_contiguous(), "w_packed must be contiguous");
	TORCH_CHECK(w_packed.size(1) * 4 == K, "w_packed second dim must be K/4");
	}

	std::unique_ptr<float[]> scales_to_float(const torch::Tensor& s) {
	const int64_t n = s.numel();
	std::unique_ptr<float[]> out(new float[n]);
	if (s.dtype() == torch::kBFloat16) {
	const at::BFloat16* p = s.const_data_ptr<at::BFloat16>();
	for (int64_t i = 0; i < n; i++) out[i] = static_cast<float>(p[i]);
	} else if (s.dtype() == torch::kFloat32) {
	const float* p = s.const_data_ptr<float>();
	std::memcpy(out.get(), p, n * sizeof(float));
	} else {
	TORCH_CHECK(false, "scales must be bf16 or f32");
	}
	return out;
	}

	template <typename OutT>
	void run_gemm(const torch::Tensor& out, const int8_t* a, const int32_t* sums, const float* sa,
	int64_t M, int64_t K, const torch::Tensor& w_packed, const float* sw, int64_t N) {
	gemm_core<OutT>(a, sums, sa, M, K, w_packed.const_data_ptr<uint8_t>(), sw, N,
	reinterpret_cast<OutT*>(out.data_ptr()));
	}

	void dispatch_out(const torch::Tensor& out, const int8_t* a, const int32_t* sums, const float* sa,
	int64_t M, int64_t K, const torch::Tensor& w_packed, const float* sw,
	int64_t N) {
	if (out.dtype() == torch::kBFloat16) {
	run_gemm<at::BFloat16>(out, a, sums, sa, M, K, w_packed, sw, N);
	} else if (out.dtype() == torch::kFloat32) {
	run_gemm<float>(out, a, sums, sa, M, K, w_packed, sw, N);
	} else {
	TORCH_CHECK(false, "out must be bf16 or f32");
	}
	}

	} // namespace

	// ---------------------------------------------------------------------------
	// Torch entry points (same schema as phanerozoic/bitnet-tc).
	// ---------------------------------------------------------------------------

	void quantize_act(torch::Tensor& Aq, torch::Tensor& scale, torch::Tensor const& A) {
	TORCH_CHECK(A.device().is_cpu() && Aq.device().is_cpu() && scale.device().is_cpu(),
	"all tensors must be CPU");
	TORCH_CHECK(A.dim() == 2 && A.is_contiguous(), "A must be contiguous 2D");
	TORCH_CHECK(Aq.dtype() == torch::kInt8 && Aq.is_contiguous(), "Aq must be contiguous int8");
	TORCH_CHECK(scale.dtype() == torch::kBFloat16, "scale must be bf16");
	const int64_t M = A.size(0), K = A.size(1);
	TORCH_CHECK(Aq.size(0) == M && Aq.size(1) == K, "Aq must match A");
	TORCH_CHECK(scale.numel() == M, "scale must be [M]");

	std::unique_ptr<float[]> sc(new float[M]);
	if (A.dtype() == torch::kBFloat16) {
	quantize_rows<at::BFloat16>(A.const_data_ptr<at::BFloat16>(), M, K, Aq.data_ptr<int8_t>(),
	sc.get(), nullptr);
	} else if (A.dtype() == torch::kFloat32) {
	quantize_rows<float>(A.const_data_ptr<float>(), M, K, Aq.data_ptr<int8_t>(), sc.get(),
	nullptr);
	} else {
	TORCH_CHECK(false, "A must be bf16 or f32");
	}
	at::BFloat16* sp = scale.data_ptr<at::BFloat16>();
	for (int64_t m = 0; m < M; m++) sp[m] = static_cast<at::BFloat16>(sc[m]);
	}

	void bitnet_gemm(torch::Tensor& out, torch::Tensor const& a_int8, torch::Tensor const& w_packed,
	torch::Tensor const& scale_act, torch::Tensor const& scale_wt,
	c10::optional<torch::Tensor> const& scratch) {
	(void)scratch; // CUDA split-K scratch; unused on CPU
	TORCH_CHECK(a_int8.device().is_cpu() && out.device().is_cpu(), "tensors must be CPU");
	TORCH_CHECK(a_int8.dtype() == torch::kInt8 && a_int8.dim() == 2 && a_int8.is_contiguous(),
	"a_int8 must be contiguous int8 [M, K]");
	const int64_t M = a_int8.size(0), K = a_int8.size(1);
	TORCH_CHECK(K % 32 == 0, "K must be a multiple of 32");
	check_weights(w_packed, K);
	const int64_t N = w_packed.size(0);
	TORCH_CHECK(out.dim() == 2 && out.size(0) == M && out.size(1) == N && out.is_contiguous(),
	"out must be contiguous [M, N]");
	TORCH_CHECK(scale_act.numel() == M, "scale_act must be [M]");
	TORCH_CHECK(scale_wt.numel() == N, "scale_wt must be [N]");

	auto sa = scales_to_float(scale_act);
	auto sw = scales_to_float(scale_wt);
	std::unique_ptr<int32_t[]> sums(new int32_t[M]);
	row_sums(a_int8.const_data_ptr<int8_t>(), M, K, sums.get());
	dispatch_out(out, a_int8.const_data_ptr<int8_t>(), sums.get(), sa.get(), M, K, w_packed,
	sw.get(), N);
	}

	void bitnet_gemv_fused(torch::Tensor& out, torch::Tensor const& a, torch::Tensor const& w_packed,
	torch::Tensor const& scale_wt) {
	TORCH_CHECK(a.device().is_cpu() && out.device().is_cpu(), "tensors must be CPU");
	TORCH_CHECK(a.dim() == 2 && a.is_contiguous(), "a must be contiguous 2D [M, K]");
	const int64_t M = a.size(0), K = a.size(1);
	TORCH_CHECK(M < 16, "fused path expects M < 16");
	TORCH_CHECK(K % 32 == 0, "K must be a multiple of 32");
	check_weights(w_packed, K);
	const int64_t N = w_packed.size(0);
	TORCH_CHECK(out.dim() == 2 && out.size(0) == M && out.size(1) == N && out.is_contiguous(),
	"out must be contiguous [M, N]");
	TORCH_CHECK(scale_wt.numel() == N, "scale_wt must be [N]");

	std::unique_ptr<int8_t[]> q(new int8_t[M * K]);
	std::unique_ptr<float[]> sa(new float[M]);
	std::unique_ptr<int32_t[]> sums(new int32_t[M]);
	if (a.dtype() == torch::kBFloat16) {
	quantize_rows<at::BFloat16>(a.const_data_ptr<at::BFloat16>(), M, K, q.get(), sa.get(),
	sums.get());
	} else if (a.dtype() == torch::kFloat32) {
	quantize_rows<float>(a.const_data_ptr<float>(), M, K, q.get(), sa.get(), sums.get());
	} else {
	TORCH_CHECK(false, "a must be bf16 or f32");
	}
	auto sw = scales_to_float(scale_wt);
	dispatch_out(out, q.get(), sums.get(), sa.get(), M, K, w_packed, sw.get(), N);
	}