// 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);
}