"wrappers/vscode:/vscode.git/clone" did not exist on "7a60e4343e3e7157380bac641396a5250549deb9"
flatmm_a8w8_blockscale_asm_pybind.cu 347 Bytes
Newer Older
Xiaowei.zhang's avatar
Xiaowei.zhang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
// SPDX-License-Identifier: MIT
 
#include "asm_flatmm_a8w8_blockscale.h"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("flatmm_a8w8_blockscale_asm", &flatmm_a8w8_blockscale_asm, "flatmm_a8w8_blockscale_asm", 
        py::arg("XQ"),
        py::arg("WQ"),
        py::arg("x_scale"),
        py::arg("w_scale"),
        py::arg("Out"));
}