From 58d2610bc2cf547bc03e9afcba5b61e74b6f4b2f Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:44:28 +0900 Subject: [PATCH 1/9] feat(rocm): extract GGUF Q8 activation quantizer --- .../kernel/csrc/gguf/quantize_q8_1.cuh | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 python/freetoken/kernel/csrc/gguf/quantize_q8_1.cuh diff --git a/python/freetoken/kernel/csrc/gguf/quantize_q8_1.cuh b/python/freetoken/kernel/csrc/gguf/quantize_q8_1.cuh new file mode 100644 index 00000000..a41ce6cd --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/quantize_q8_1.cuh @@ -0,0 +1,71 @@ +// Extracted from gguf_kernel.cu so ROCm can JIT-compile GGUF operations in +// independent translation units. The implementation is unchanged from the +// vendored sgl-kernel/llama.cpp-derived FreeToken GGUF path. +#pragma once + +#include + +// Q8 activation quantization used by MMVQ/MMQ and grouped MoE kernels. +template +static __global__ void quantize_q8_1( + const scalar_t* __restrict__ x, + void* __restrict__ vy, + const int kx, + const int kx_padded) { + const auto ix = blockDim.x * blockIdx.x + threadIdx.x; + if (ix >= kx_padded) { + return; + } + const auto iy = blockDim.y * blockIdx.y + threadIdx.y; + const int i_padded = iy * kx_padded + ix; + + block_q8_1* y = (block_q8_1*)vy; + + const int ib = i_padded / QK8_1; + const int iqs = i_padded % QK8_1; + + const float xi = ix < kx ? static_cast(x[iy * kx + ix]) : 0.0f; + float amax = fabsf(xi); + float sum = xi; + +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + amax = fmaxf(amax, SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), amax, mask, 32)); + sum += SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), sum, mask, 32); + } + + const float d = amax / 127; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + + y[ib].qs[iqs] = q; + + if (iqs > 0) { + return; + } + + y[ib].ds.x = __float2half(d); + y[ib].ds.y = __float2half(sum); +} + +template +static void quantize_row_q8_1_cuda( + const scalar_t* x, + void* vy, + const int kx, + const int ky, + cudaStream_t stream) { + const int64_t kx_padded = (kx + 512 - 1) / 512 * 512; + const int block_num_x = + (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; + constexpr int MAX_BLOCK_SIZE = 65535; + for (int off = 0; off < ky; off += MAX_BLOCK_SIZE) { + const int num_blocks_y = std::min(ky, off + MAX_BLOCK_SIZE) - off; + const dim3 num_blocks(block_num_x, num_blocks_y, 1); + const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1); + quantize_q8_1<<>>( + &x[off * kx], + (int32_t*)vy + off * (kx_padded / 32 * 9), + kx, + kx_padded); + } +} From 157440b3a03bf089ad0a36034de30e8d48a49b49 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:44:45 +0900 Subject: [PATCH 2/9] feat(rocm): split GGUF dequant JIT translation unit --- .../kernel/csrc/gguf/gguf_dequant_kernel.cu | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 python/freetoken/kernel/csrc/gguf/gguf_dequant_kernel.cu diff --git a/python/freetoken/kernel/csrc/gguf/gguf_dequant_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_dequant_kernel.cu new file mode 100644 index 00000000..3b9d5fc6 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_dequant_kernel.cu @@ -0,0 +1,40 @@ +// ROCm operation-split binding for the vendored GGUF dequant kernels. +// Kernel implementations remain in the sgl-kernel/llama.cpp-derived headers. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "dequantize.cuh" + +torch::Tensor ggml_dequantize( + torch::Tensor W, + int64_t type, + int64_t m, + int64_t n, + std::optional const& dtype) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(W)); + auto dtype_ = dtype.value_or(torch::kFloat16); + auto options = torch::TensorOptions().dtype(dtype_).device(W.device()); + at::Tensor DW = torch::empty({m, n}, options); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { + auto to_cuda = ggml_get_to_cuda(type); + TORCH_CHECK( + to_cuda != nullptr, + "ggml_dequantize: unsupported GGUF quant type ", type, + " (dequant kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/" + "IQ2_XS/IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); + to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream); + }); + + return DW; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_dequantize", &ggml_dequantize, ""); +} From 2f85ce6a1def50a579f00a1c78fe602b13accb92 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:45:16 +0900 Subject: [PATCH 3/9] feat(rocm): split GGUF MMVQ JIT translation unit --- .../kernel/csrc/gguf/gguf_mmvq_kernel.cu | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 python/freetoken/kernel/csrc/gguf/gguf_mmvq_kernel.cu diff --git a/python/freetoken/kernel/csrc/gguf/gguf_mmvq_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_mmvq_kernel.cu new file mode 100644 index 00000000..8ea3f818 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_mmvq_kernel.cu @@ -0,0 +1,102 @@ +// ROCm operation-split binding for GGUF small-batch MMVQ. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "mmvq.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_mul_mat_vec_a8( + torch::Tensor W, + torch::Tensor X, + int64_t type, + int64_t row) { + int col = X.sizes()[1]; + int vecs = X.sizes()[0]; + const int padded = (col + 512 - 1) / 512 * 512; + const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y = torch::empty({vecs, row}, options); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, vecs, stream); + switch (type) { + case 2: + mul_mat_vec_q4_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 3: + mul_mat_vec_q4_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 6: + mul_mat_vec_q5_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 7: + mul_mat_vec_q5_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 8: + mul_mat_vec_q8_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 10: + mul_mat_vec_q2_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 11: + mul_mat_vec_q3_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 12: + mul_mat_vec_q4_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 13: + mul_mat_vec_q5_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 14: + mul_mat_vec_q6_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 16: + mul_mat_vec_iq2_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 17: + mul_mat_vec_iq2_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 18: + mul_mat_vec_iq3_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 19: + mul_mat_vec_iq1_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 20: + mul_mat_vec_iq4_nl_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 21: + mul_mat_vec_iq3_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 22: + mul_mat_vec_iq2_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 23: + mul_mat_vec_iq4_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 29: + mul_mat_vec_iq1_m_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + default: + TORCH_CHECK( + false, + "ggml_mul_mat_vec_a8: unsupported GGUF quant type ", type, + " (MMVQ kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/IQ2_XS/" + "IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); + } + }); + return Y; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_mul_mat_vec_a8", &ggml_mul_mat_vec_a8, ""); +} From 66e1a07269369a7fcaac5c3cf2dac2ab4cf683c6 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:45:36 +0900 Subject: [PATCH 4/9] feat(rocm): split GGUF MMQ JIT translation unit --- .../kernel/csrc/gguf/gguf_mmq_kernel.cu | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 python/freetoken/kernel/csrc/gguf/gguf_mmq_kernel.cu diff --git a/python/freetoken/kernel/csrc/gguf/gguf_mmq_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_mmq_kernel.cu new file mode 100644 index 00000000..666087b0 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_mmq_kernel.cu @@ -0,0 +1,57 @@ +// ROCm operation-split binding for GGUF large-batch MMQ. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "mmq.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_mul_mat_a8( + torch::Tensor W, + torch::Tensor X, + int64_t type, + int64_t row) { + int col = X.sizes()[1]; + int padded = (col + 512 - 1) / 512 * 512; + int batch = X.sizes()[0]; + const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y = torch::empty({batch, row}, options); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X = torch::empty({batch, padded / 32 * 9}, options); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, batch, stream); + using Fn = void (*)(const void*, const void*, scalar_t*, int, int, int, int, int, cudaStream_t); + Fn fn = nullptr; + switch (type) { + case 2: fn = &ggml_mul_mat_q4_0_q8_1_cuda; break; + case 3: fn = &ggml_mul_mat_q4_1_q8_1_cuda; break; + case 6: fn = &ggml_mul_mat_q5_0_q8_1_cuda; break; + case 7: fn = &ggml_mul_mat_q5_1_q8_1_cuda; break; + case 8: fn = &ggml_mul_mat_q8_0_q8_1_cuda; break; + case 10: fn = &ggml_mul_mat_q2_K_q8_1_cuda; break; + case 11: fn = &ggml_mul_mat_q3_K_q8_1_cuda; break; + case 12: fn = &ggml_mul_mat_q4_K_q8_1_cuda; break; + case 13: fn = &ggml_mul_mat_q5_K_q8_1_cuda; break; + case 14: fn = &ggml_mul_mat_q6_K_q8_1_cuda; break; + default: + TORCH_CHECK(false, "ggml_mul_mat_a8: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); + } + fn(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + col, row, batch, padded, row, stream); + }); + return Y; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_mul_mat_a8", &ggml_mul_mat_a8, ""); +} From 6ec1e7d69781e9a26a699aaaee770a915912ed52 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:45:53 +0900 Subject: [PATCH 5/9] feat(rocm): split GGUF MoE-vector JIT translation unit --- .../kernel/csrc/gguf/gguf_moe_vec_kernel.cu | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) create mode 100644 python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu diff --git a/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu new file mode 100644 index 00000000..aab4f7eb --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu @@ -0,0 +1,70 @@ +// ROCm operation-split binding for GGUF routed small-batch MoE vector kernels. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "moe_vec.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_moe_a8_vec( + torch::Tensor X, + torch::Tensor W, + torch::Tensor topk_ids, + int64_t top_k, + int64_t type, + int64_t row, + int64_t tokens) { + int col = X.sizes()[1]; + const int padded = (col + 512 - 1) / 512 * 512; + const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y = torch::zeros({tokens * top_k, row}, options); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); +#define FT_MOE_VEC_CASE(TYPE, FN) \ + case TYPE: \ + FN(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), \ + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); \ + break + switch (type) { + FT_MOE_VEC_CASE(2, moe_vec_q4_0_q8_1_cuda); + FT_MOE_VEC_CASE(3, moe_vec_q4_1_q8_1_cuda); + FT_MOE_VEC_CASE(6, moe_vec_q5_0_q8_1_cuda); + FT_MOE_VEC_CASE(7, moe_vec_q5_1_q8_1_cuda); + FT_MOE_VEC_CASE(8, moe_vec_q8_0_q8_1_cuda); + FT_MOE_VEC_CASE(10, moe_vec_q2_K_q8_1_cuda); + FT_MOE_VEC_CASE(11, moe_vec_q3_K_q8_1_cuda); + FT_MOE_VEC_CASE(12, moe_vec_q4_K_q8_1_cuda); + FT_MOE_VEC_CASE(13, moe_vec_q5_K_q8_1_cuda); + FT_MOE_VEC_CASE(14, moe_vec_q6_K_q8_1_cuda); + FT_MOE_VEC_CASE(16, moe_vec_iq2_xxs_q8_1_cuda); + FT_MOE_VEC_CASE(17, moe_vec_iq2_xs_q8_1_cuda); + FT_MOE_VEC_CASE(18, moe_vec_iq3_xxs_q8_1_cuda); + FT_MOE_VEC_CASE(19, moe_vec_iq1_s_q8_1_cuda); + FT_MOE_VEC_CASE(20, moe_vec_iq4_nl_q8_1_cuda); + FT_MOE_VEC_CASE(21, moe_vec_iq3_s_q8_1_cuda); + FT_MOE_VEC_CASE(22, moe_vec_iq2_s_q8_1_cuda); + FT_MOE_VEC_CASE(23, moe_vec_iq4_xs_q8_1_cuda); + FT_MOE_VEC_CASE(29, moe_vec_iq1_m_q8_1_cuda); + default: + TORCH_CHECK(false, "ggml_moe_a8_vec: unsupported GGUF quant type ", type, + " (MMVQ kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/IQ2_XS/" + "IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); + } +#undef FT_MOE_VEC_CASE + }); + return Y; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_moe_a8_vec", &ggml_moe_a8_vec, ""); +} From 3493d65bdd6770ca0dfca295decba6b62e1cdccf Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:46:15 +0900 Subject: [PATCH 6/9] feat(rocm): split GGUF grouped-MoE JIT translation unit --- .../kernel/csrc/gguf/gguf_moe_kernel.cu | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 python/freetoken/kernel/csrc/gguf/gguf_moe_kernel.cu diff --git a/python/freetoken/kernel/csrc/gguf/gguf_moe_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_moe_kernel.cu new file mode 100644 index 00000000..f562f45e --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_moe_kernel.cu @@ -0,0 +1,87 @@ +// ROCm operation-split binding for GGUF grouped large-batch MoE kernels. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "mmq.cuh" +#include "moe.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_moe_a8( + torch::Tensor X, + torch::Tensor W, + torch::Tensor sorted_token_ids, + torch::Tensor expert_ids, + torch::Tensor num_tokens_post_padded, + int64_t type, + int64_t row, + int64_t top_k, + int64_t tokens) { + int col = X.sizes()[1]; + int padded = (col + 512 - 1) / 512 * 512; + const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y = torch::empty({tokens * top_k, row}, options); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_a8", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); + using Fn = void (*)( + const void*, const void*, scalar_t*, const int*, const int*, const int*, + int, int, int, int, int, int, int, int, cudaStream_t); + Fn fn = nullptr; + switch (type) { + case 2: fn = &ggml_moe_q4_0_q8_1_cuda; break; + case 3: fn = &ggml_moe_q4_1_q8_1_cuda; break; + case 6: fn = &ggml_moe_q5_0_q8_1_cuda; break; + case 7: fn = &ggml_moe_q5_1_q8_1_cuda; break; + case 8: fn = &ggml_moe_q8_0_q8_1_cuda; break; + case 10: fn = &ggml_moe_q2_K_q8_1_cuda; break; + case 11: fn = &ggml_moe_q3_K_q8_1_cuda; break; + case 12: fn = &ggml_moe_q4_K_q8_1_cuda; break; + case 13: fn = &ggml_moe_q5_K_q8_1_cuda; break; + case 14: fn = &ggml_moe_q6_K_q8_1_cuda; break; + default: + TORCH_CHECK(false, "ggml_moe_a8: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); + } + fn(quant_X.data_ptr(), W.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)sorted_token_ids.data_ptr(), (int*)expert_ids.data_ptr(), + (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, tokens, + padded, row, top_k, sorted_token_ids.sizes()[0], stream); + }); + return Y; +} + +int64_t ggml_moe_get_block_size(int64_t type) { + switch (type) { + case 2: return MOE_X_Q4_0; + case 3: return MOE_X_Q4_1; + case 6: return MOE_X_Q5_0; + case 7: return MOE_X_Q5_1; + case 8: return MOE_X_Q8_0; + case 10: return MOE_X_Q2_K; + case 11: return MOE_X_Q3_K; + case 12: return MOE_X_Q4_K; + case 13: return MOE_X_Q5_K; + case 14: return MOE_X_Q6_K; + default: + TORCH_CHECK(false, "ggml_moe_get_block_size: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); + return 0; + } +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_moe_a8", &ggml_moe_a8, ""); + m.def("ggml_moe_get_block_size", &ggml_moe_get_block_size, ""); +} From 16ac89f153fea1879f5be4e9569cbbed0039a9a4 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:46:55 +0900 Subject: [PATCH 7/9] feat(rocm): route GGUF JIT by operation family --- python/freetoken/kernel/gguf.py | 121 ++++++++++++++++++++++++-------- 1 file changed, 92 insertions(+), 29 deletions(-) diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 4d44ae5d..fee3db00 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -1,13 +1,18 @@ -"""Borrowed llama.cpp GGUF dequant/GEMM CUDA kernels, JIT-compiled on first use. +"""Borrowed llama.cpp GGUF dequant/GEMM kernels, JIT-compiled on first use. -The ``.cu``/``.cuh`` under ``csrc/gguf/`` are vendored verbatim from sgl-kernel +The ``.cu``/``.cuh`` under ``csrc/gguf/`` are vendored from sgl-kernel (``csrc/quantization/gguf/``), which are themselves ports of llama.cpp. We compile -them through ``torch.utils.cpp_extension.load`` (the same toolchain sglang/vllm use) -into a torch-op module and expose the handful of ops the GGUF path needs. This is a -separate, torch-native extension that sits alongside FreeToken's tvm-ffi kernels. +them through ``torch.utils.cpp_extension.load`` into torch-op modules and expose +the handful of ops the GGUF path needs. -All ops keep the weight in its native GGUF block layout (packed ``uint8`` rows) and -dequantize *inside* the kernel -- no bf16 copy of the weight is ever materialized. +CUDA keeps the original monolithic translation unit. ROCm uses operation-split +translation units because AMD clang can spend many minutes optimizing the full +all-quant dequant+MMVQ+MMQ+MoE unit even though each operation family compiles in +seconds on RDNA. Splitting by operation preserves all quant coverage without +forcing per-quant JIT modules. + +All ops keep the weight in its native GGUF block layout (packed ``uint8`` rows) +and dequantize inside the kernel -- no bf16 copy of the weight is materialized. """ from __future__ import annotations @@ -21,17 +26,31 @@ import torch _CSRC = pathlib.Path(__file__).parent / "csrc" / "gguf" +_ROCM_OPERATION_SOURCES = { + "dequant": "gguf_dequant_kernel.cu", + "mmvq": "gguf_mmvq_kernel.cu", + "mmq": "gguf_mmq_kernel.cu", + "moe_vec": "gguf_moe_vec_kernel.cu", + "moe": "gguf_moe_kernel.cu", +} + + +def _is_rocm() -> bool: + return getattr(torch.version, "hip", None) is not None def _staged_rocm_sources() -> pathlib.Path: """Copy CUDA sources out of the checkout before PyTorch HIPifies them. ``torch.utils.cpp_extension.load`` writes generated ``*_hip`` sources next to - the input file. Keeping the staging directory under the extension cache makes + the input file. Keeping the staging directory under the extension cache makes the source checkout stay clean while still allowing normal incremental builds. """ cache_root = pathlib.Path( - os.environ.get("TORCH_EXTENSIONS_DIR", pathlib.Path.home() / ".cache" / "torch_extensions") + os.environ.get( + "TORCH_EXTENSIONS_DIR", + pathlib.Path.home() / ".cache" / "torch_extensions", + ) ) digest = hashlib.sha256() digest.update(f"torch={torch.__version__};hip={torch.version.hip}".encode()) @@ -75,21 +94,49 @@ def _c_compiler_for(cxx: str) -> str: return shutil.which(cc) or cc +@functools.cache +def _rocm_module(operation: str): + """Build one all-quant GGUF operation family on ROCm. + + Keeping quant types together avoids a large fleet of JIT extensions while + keeping AMD clang away from the pathological monolithic translation unit. + """ + if operation not in _ROCM_OPERATION_SOURCES: + raise ValueError(f"unknown ROCm GGUF operation: {operation}") + + from freetoken.kernel.utils import _rocm_link_flags + from torch.utils.cpp_extension import load + + csrc = _staged_rocm_sources() + return load( + name=f"freetoken_gguf_rocm_{operation}_kernels", + sources=[str(csrc / _ROCM_OPERATION_SOURCES[operation])], + extra_include_paths=[str(csrc)], + extra_cuda_cflags=[ + "-O3", + "-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_CPP", + ], + extra_ldflags=_rocm_link_flags(), + verbose=True, + ) + + @functools.cache def _module(): + """Build the original monolithic CUDA extension. + + This remains available on ROCm for compatibility/debugging, but public + wrappers route ROCm calls through ``_rocm_module`` instead. + """ from torch.utils.cpp_extension import load - is_rocm = getattr(torch.version, "hip", None) is not None + is_rocm = _is_rocm() extra_cuda_cflags = ["-O3"] extra_ldflags: list[str] = [] if is_rocm: from freetoken.kernel.utils import _rocm_link_flags extra_ldflags = _rocm_link_flags() - # Ubuntu's generic Thrust headers otherwise select the CUDA backend and - # try to include cuda_runtime_api.h. GGUF only reaches Thrust through a - # libtorch complex-number header, so the backend-neutral C++ path is - # sufficient for HIP compilation. extra_cuda_cflags.append("-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_CPP") csrc = _staged_rocm_sources() else: @@ -98,16 +145,11 @@ def _module(): host_cxx = None if is_rocm else _host_compiler() if host_cxx is not None: - # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a - # libtorch/nvcc-compatible compiler. Force (not setdefault): the system - # default (CXX unset -> g++) can be a gcc too new for the torch headers. cxx_path = shutil.which(host_cxx) or host_cxx extra_cuda_cflags += ["-ccbin", cxx_path] os.environ["CXX"] = cxx_path os.environ["CC"] = _c_compiler_for(cxx_path) - # gguf_kernel.cu carries its own PYBIND11_MODULE (appended at the end), so a - # plain `load` of the single source compiles + binds the ggml_* ops. return load( name="freetoken_gguf_kernels", sources=[str(csrc / "gguf_kernel.cu")], @@ -118,28 +160,40 @@ def _module(): ) +def _operation_module(operation: str): + return _rocm_module(operation) if _is_rocm() else _module() + + # ---- thin typed wrappers (signatures mirror sgl_kernel.quantization.gguf) ---- def ggml_dequantize( - weight: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None = None + weight: torch.Tensor, + quant_type: int, + m: int, + n: int, + dtype: torch.dtype | None = None, ) -> torch.Tensor: - """Dequantize a packed GGUF weight ``[m, row_bytes]`` to a dense ``[m, n]`` tensor.""" - return _module().ggml_dequantize(weight, quant_type, m, n, dtype) + """Dequantize a packed GGUF weight ``[m, row_bytes]`` to dense ``[m, n]``.""" + return _operation_module("dequant").ggml_dequantize( + weight, quant_type, m, n, dtype + ) def ggml_mul_mat_vec_a8( weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int ) -> torch.Tensor: """MMVQ: small-batch GEMV with on-the-fly dequant. ``row`` = output features.""" - return _module().ggml_mul_mat_vec_a8(weight, x, quant_type, row) + return _operation_module("mmvq").ggml_mul_mat_vec_a8( + weight, x, quant_type, row + ) def ggml_mul_mat_a8( weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int ) -> torch.Tensor: """MMQ: large-batch quantized matmul. ``row`` = output features.""" - return _module().ggml_mul_mat_a8(weight, x, quant_type, row) + return _operation_module("mmq").ggml_mul_mat_a8(weight, x, quant_type, row) def ggml_moe_a8( @@ -154,9 +208,16 @@ def ggml_moe_a8( tokens: int, ) -> torch.Tensor: """MMQ grouped expert matmul over stacked experts ``weight[E, row, *]``.""" - return _module().ggml_moe_a8( - x, weight, sorted_token_ids, expert_ids, num_tokens_post_padded, - quant_type, row, top_k, tokens, + return _operation_module("moe").ggml_moe_a8( + x, + weight, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type, + row, + top_k, + tokens, ) @@ -170,11 +231,13 @@ def ggml_moe_a8_vec( tokens: int, ) -> torch.Tensor: """MMVQ grouped expert GEMV over stacked experts ``weight[E, row, *]``.""" - return _module().ggml_moe_a8_vec(x, weight, topk_ids, top_k, quant_type, row, tokens) + return _operation_module("moe_vec").ggml_moe_a8_vec( + x, weight, topk_ids, top_k, quant_type, row, tokens + ) def ggml_moe_get_block_size(quant_type: int) -> int: - return _module().ggml_moe_get_block_size(quant_type) + return _operation_module("moe").ggml_moe_get_block_size(quant_type) __all__ = [ From 93cc7ecc6efb1b16b2140e8b045027ced68e4605 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:58:33 +0900 Subject: [PATCH 8/9] fix(rocm): keep GGUF MoE-vector switch HIPify-safe --- .../kernel/csrc/gguf/gguf_moe_vec_kernel.cu | 101 +++++++++++++----- 1 file changed, 76 insertions(+), 25 deletions(-) diff --git a/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu index aab4f7eb..ddf6fa8f 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu @@ -29,37 +29,88 @@ torch::Tensor ggml_moe_a8_vec( DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] { quantize_row_q8_1_cuda( (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); -#define FT_MOE_VEC_CASE(TYPE, FN) \ - case TYPE: \ - FN(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), \ - (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); \ - break switch (type) { - FT_MOE_VEC_CASE(2, moe_vec_q4_0_q8_1_cuda); - FT_MOE_VEC_CASE(3, moe_vec_q4_1_q8_1_cuda); - FT_MOE_VEC_CASE(6, moe_vec_q5_0_q8_1_cuda); - FT_MOE_VEC_CASE(7, moe_vec_q5_1_q8_1_cuda); - FT_MOE_VEC_CASE(8, moe_vec_q8_0_q8_1_cuda); - FT_MOE_VEC_CASE(10, moe_vec_q2_K_q8_1_cuda); - FT_MOE_VEC_CASE(11, moe_vec_q3_K_q8_1_cuda); - FT_MOE_VEC_CASE(12, moe_vec_q4_K_q8_1_cuda); - FT_MOE_VEC_CASE(13, moe_vec_q5_K_q8_1_cuda); - FT_MOE_VEC_CASE(14, moe_vec_q6_K_q8_1_cuda); - FT_MOE_VEC_CASE(16, moe_vec_iq2_xxs_q8_1_cuda); - FT_MOE_VEC_CASE(17, moe_vec_iq2_xs_q8_1_cuda); - FT_MOE_VEC_CASE(18, moe_vec_iq3_xxs_q8_1_cuda); - FT_MOE_VEC_CASE(19, moe_vec_iq1_s_q8_1_cuda); - FT_MOE_VEC_CASE(20, moe_vec_iq4_nl_q8_1_cuda); - FT_MOE_VEC_CASE(21, moe_vec_iq3_s_q8_1_cuda); - FT_MOE_VEC_CASE(22, moe_vec_iq2_s_q8_1_cuda); - FT_MOE_VEC_CASE(23, moe_vec_iq4_xs_q8_1_cuda); - FT_MOE_VEC_CASE(29, moe_vec_iq1_m_q8_1_cuda); + case 2: + moe_vec_q4_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 3: + moe_vec_q4_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 6: + moe_vec_q5_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 7: + moe_vec_q5_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 8: + moe_vec_q8_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 10: + moe_vec_q2_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 11: + moe_vec_q3_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 12: + moe_vec_q4_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 13: + moe_vec_q5_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 14: + moe_vec_q6_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 16: + moe_vec_iq2_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 17: + moe_vec_iq2_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 18: + moe_vec_iq3_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 19: + moe_vec_iq1_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 20: + moe_vec_iq4_nl_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 21: + moe_vec_iq3_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 22: + moe_vec_iq2_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 23: + moe_vec_iq4_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 29: + moe_vec_iq1_m_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; default: TORCH_CHECK(false, "ggml_moe_a8_vec: unsupported GGUF quant type ", type, " (MMVQ kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/IQ2_XS/" "IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); } -#undef FT_MOE_VEC_CASE }); return Y; } From 3ee61fa8b2ca49dfc42f3033afba41823667ec7b Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:50:54 +0900 Subject: [PATCH 9/9] test(rocm): protect GGUF operation-split routing --- .../kernels/test_gguf_rocm_operation_split.py | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 tests/kernels/test_gguf_rocm_operation_split.py diff --git a/tests/kernels/test_gguf_rocm_operation_split.py b/tests/kernels/test_gguf_rocm_operation_split.py new file mode 100644 index 00000000..d3d48b09 --- /dev/null +++ b/tests/kernels/test_gguf_rocm_operation_split.py @@ -0,0 +1,96 @@ +"""Routing contract for ROCm GGUF operation-split JIT modules. + +These tests deliberately do not compile a native extension. Physical gfx1101 coverage +lives in the ROCm validation receipts; this file protects the cheap Python dispatch +contract so a future refactor cannot silently route one public GGUF op back through the +pathological monolithic ROCm translation unit. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.kernel import gguf + + +_EXPECTED_SOURCES = { + "dequant": "gguf_dequant_kernel.cu", + "mmvq": "gguf_mmvq_kernel.cu", + "mmq": "gguf_mmq_kernel.cu", + "moe_vec": "gguf_moe_vec_kernel.cu", + "moe": "gguf_moe_kernel.cu", +} + + +def test_rocm_operation_source_contract(): + assert gguf._ROCM_OPERATION_SOURCES == _EXPECTED_SOURCES + + +def test_operation_module_uses_rocm_family(monkeypatch): + seen = [] + marker = object() + monkeypatch.setattr(gguf, "_is_rocm", lambda: True) + monkeypatch.setattr(gguf, "_rocm_module", lambda operation: seen.append(operation) or marker) + monkeypatch.setattr(gguf, "_module", lambda: pytest.fail("ROCm dispatch reached monolithic module")) + + assert gguf._operation_module("mmvq") is marker + assert seen == ["mmvq"] + + +def test_operation_module_keeps_cuda_monolith(monkeypatch): + marker = object() + monkeypatch.setattr(gguf, "_is_rocm", lambda: False) + monkeypatch.setattr(gguf, "_module", lambda: marker) + monkeypatch.setattr(gguf, "_rocm_module", lambda operation: pytest.fail(f"CUDA reached ROCm split {operation}")) + + assert gguf._operation_module("dequant") is marker + + +def test_rocm_loader_rejects_unknown_family_before_build(): + with pytest.raises(ValueError, match="unknown ROCm GGUF operation"): + gguf._rocm_module("not-an-operation") + + +def test_public_wrappers_route_to_owned_operation_family(monkeypatch): + calls = [] + + def op(name): + def invoke(*args, **kwargs): + calls.append((name, args, kwargs)) + return name + return invoke + + modules = { + family: SimpleNamespace( + ggml_dequantize=op("dequant"), + ggml_mul_mat_vec_a8=op("mmvq"), + ggml_mul_mat_a8=op("mmq"), + ggml_moe_a8_vec=op("moe_vec"), + ggml_moe_a8=op("moe"), + ggml_moe_get_block_size=op("moe_block"), + ) + for family in _EXPECTED_SOURCES + } + requested = [] + + def load_family(family): + requested.append(family) + return modules[family] + + monkeypatch.setattr(gguf, "_operation_module", load_family) + + w = torch.empty(1, dtype=torch.uint8) + x = torch.empty(1) + ids = torch.empty(1, dtype=torch.int32) + + assert gguf.ggml_dequantize(w, 2, 1, 32) == "dequant" + assert gguf.ggml_mul_mat_vec_a8(w, x, 2, 1) == "mmvq" + assert gguf.ggml_mul_mat_a8(w, x, 2, 1) == "mmq" + assert gguf.ggml_moe_a8_vec(x, w, ids, 1, 2, 1, 1) == "moe_vec" + assert gguf.ggml_moe_a8(x, w, ids, ids, ids, 2, 1, 1, 1) == "moe" + assert gguf.ggml_moe_get_block_size(2) == "moe_block" + + assert requested == ["dequant", "mmvq", "mmq", "moe_vec", "moe", "moe"]