asm_flatmm_a8w8_blockscale.h

#pragma once
// SPDX-License-Identifier: MIT
 
#include <torch/extension.h>

torch::Tensor flatmm_a8w8_blockscale_asm(
    torch::Tensor &XQ,      // [M, K]
    torch::Tensor &WQ,      // [N, K] -> [N/128, K*128]
    torch::Tensor &x_scale, // [K/128, M]
    torch::Tensor &w_scale, // [K/128, N/128]
    torch::Tensor &out      // Out:[M, N] fp16
);