From 0137a1e736a63cc3bc7c81f7084dc3f049d34b15 Mon Sep 17 00:00:00 2001 From: Hao Liu Date: Mon, 24 Aug 2026 10:17:38 +0800 Subject: [PATCH] [CUDA] Add GGUF MMQ kernels for I-quant types Add dense and batched MoE MMQ kernels for eight supported I-quant GGUF formats, with conservative routing and compatibility fallback for older sglang-kernel wheels. Fixes #35019 Signed-off-by: Hao Liu --- .../sglang/kernels/aot/THIRDPARTYNOTICES.txt | 25 + .../kernels/aot/benchmark/bench_gguf_mmq.py | 142 +++ .../kernels/aot/csrc/common_extension.cc | 7 +- .../aot/csrc/quantization/gguf/gguf_kernel.cu | 324 +++++ .../aot/csrc/quantization/gguf/mmq.cuh | 606 +++++++++ .../aot/csrc/quantization/gguf/moe.cuh | 1067 +++++++++++++++- .../aot/csrc/quantization/gguf/moe_vec.cuh | 2 +- .../aot/csrc/quantization/gguf/vecdotq.cuh | 1096 ++++++++++++++++- .../kernels/aot/include/sgl_kernel_ops.h | 2 + .../kernels/aot/python/sgl_kernel/__init__.py | 1 + .../sgl_kernel/quantization/__init__.py | 1 + .../python/sgl_kernel/quantization/gguf.py | 7 + python/sglang/kernels/aot/tests/test_gguf.py | 295 ++++- python/sglang/srt/layers/quantization/gguf.py | 186 ++- .../quantization/test_gguf_mmq_types.py | 447 +++++++ 15 files changed, 4127 insertions(+), 81 deletions(-) create mode 100644 python/sglang/kernels/aot/benchmark/bench_gguf_mmq.py create mode 100644 test/registered/unit/layers/quantization/test_gguf_mmq_types.py diff --git a/python/sglang/kernels/aot/THIRDPARTYNOTICES.txt b/python/sglang/kernels/aot/THIRDPARTYNOTICES.txt index 2a81411ae9e7..cb26429ce67d 100644 --- a/python/sglang/kernels/aot/THIRDPARTYNOTICES.txt +++ b/python/sglang/kernels/aot/THIRDPARTYNOTICES.txt @@ -429,6 +429,31 @@ Notice for NVIDIA/TensorRT-LLM See the License for the specific language governing permissions and limitations under the License. +Notice for ggml-org/llama.cpp +----------------------------- + +MIT License + +Copyright (c) 2023-2024 The ggml authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + Notice for deepseek-ai/DeepGEMM ------------------------------- diff --git a/python/sglang/kernels/aot/benchmark/bench_gguf_mmq.py b/python/sglang/kernels/aot/benchmark/bench_gguf_mmq.py new file mode 100644 index 000000000000..244c4812b788 --- /dev/null +++ b/python/sglang/kernels/aot/benchmark/bench_gguf_mmq.py @@ -0,0 +1,142 @@ +"""Benchmark GGUF I-quant dense MMQ against dequantize plus matmul.""" + +import argparse +from pathlib import Path + +import torch +import triton.testing +from gguf import GGMLQuantizationType, GGUFReader +from huggingface_hub import hf_hub_download +from sgl_kernel import ggml_dequantize, ggml_mul_mat_a8 + +from sglang.utils import is_in_ci + +REPO_ID = "Isotr0py/test-gguf-sample" +REVISION = "d82b8773934ef260d8d8a896a7c197bc69a0fac1" +IQ_MMQ_TYPES = ( + "IQ1_S", + "IQ2_XXS", + "IQ2_XS", + "IQ2_S", + "IQ3_XXS", + "IQ3_S", + "IQ4_NL", + "IQ4_XS", +) + + +def parse_csv(value: str) -> list[str]: + return [item.strip() for item in value.split(",") if item.strip()] + + +def parse_int_csv(value: str) -> list[int]: + return [int(item) for item in parse_csv(value)] + + +def parse_shape(value: str) -> tuple[int, int]: + rows, columns = value.lower().split("x", maxsplit=1) + return int(rows), int(columns) + + +def load_quantized_weight( + quant_type: GGMLQuantizationType, + rows: int, + columns: int, + *, + local_files_only: bool, +) -> torch.Tensor: + filename = f"Quant_{quant_type.name}_{columns}.gguf" + path = hf_hub_download( + repo_id=REPO_ID, + filename=filename, + revision=REVISION, + local_files_only=local_files_only, + ) + matches = [] + for tensor in GGUFReader(Path(path)).tensors: + shape = tuple(int(item) for item in tensor.name.rsplit("_", 1)[-1].split("x")) + if shape == (rows, columns): + matches.append(tensor) + if len(matches) != 1: + raise RuntimeError( + f"expected one {quant_type.name} tensor with shape " + f"{rows}x{columns}, found {len(matches)}" + ) + return torch.tensor(matches[0].data, dtype=torch.uint8, device="cuda") + + +def benchmark_case( + quant_type: GGMLQuantizationType, + rows: int, + columns: int, + batch_size: int, + *, + local_files_only: bool, +) -> tuple[float, float, float]: + torch.manual_seed(0) + qweight = load_quantized_weight( + quant_type, + rows, + columns, + local_files_only=local_files_only, + ) + x = torch.randn((batch_size, columns), dtype=torch.bfloat16, device="cuda") + + def run_mmq() -> torch.Tensor: + return ggml_mul_mat_a8(qweight, x, quant_type, rows) + + def run_fallback() -> torch.Tensor: + weight = ggml_dequantize(qweight, quant_type, rows, columns, torch.bfloat16) + return x @ weight.T + + torch.testing.assert_close(run_mmq(), run_fallback(), atol=1.5, rtol=0.1) + quantiles = [0.5, 0.2, 0.8] + mmq_ms, _, _ = triton.testing.do_bench_cudagraph(run_mmq, quantiles=quantiles) + fallback_ms, _, _ = triton.testing.do_bench_cudagraph( + run_fallback, quantiles=quantiles + ) + return mmq_ms, fallback_ms, fallback_ms / mmq_ms + + +def main() -> None: + is_ci = is_in_ci() + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--quant-types", + type=parse_csv, + default=["IQ3_S"] if is_ci else list(IQ_MMQ_TYPES), + ) + parser.add_argument( + "--batch-sizes", + type=parse_int_csv, + default=[16] if is_ci else [9, 16, 17, 24, 32, 64], + ) + parser.add_argument("--shape", type=parse_shape, default=(18944, 1024)) + parser.add_argument("--local-files-only", action="store_true") + args = parser.parse_args() + + unsupported = set(args.quant_types) - set(IQ_MMQ_TYPES) + if unsupported: + parser.error(f"unsupported I-quant types: {sorted(unsupported)}") + + rows, columns = args.shape + print("| quant type | M | MMQ (ms) | dequant + matmul (ms) | speedup |") + print("|---|---:|---:|---:|---:|") + for name in args.quant_types: + quant_type = GGMLQuantizationType[name] + for batch_size in args.batch_sizes: + mmq_ms, fallback_ms, speedup = benchmark_case( + quant_type, + rows, + columns, + batch_size, + local_files_only=args.local_files_only, + ) + print( + f"| {name} | {batch_size} | {mmq_ms:.4f} | " + f"{fallback_ms:.4f} | {speedup:.2f}x |" + ) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/kernels/aot/csrc/common_extension.cc b/python/sglang/kernels/aot/csrc/common_extension.cc index 211bbdc055ed..49b2ea1130cf 100644 --- a/python/sglang/kernels/aot/csrc/common_extension.cc +++ b/python/sglang/kernels/aot/csrc/common_extension.cc @@ -433,7 +433,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.impl("ggml_moe_a8_vec", torch::kCUDA, &ggml_moe_a8_vec); m.def("ggml_moe_get_block_size(int type) -> int"); - m.impl("ggml_moe_get_block_size", torch::kCUDA, &ggml_moe_get_block_size); + // This scalar-only helper has no Tensor argument for backend dispatch. + m.impl("ggml_moe_get_block_size", &ggml_moe_get_block_size); + + m.def("ggml_supports_iq_mmq() -> bool"); + // This scalar-only capability query has no Tensor argument for dispatch. + m.impl("ggml_supports_iq_mmq", &ggml_supports_iq_mmq); /* * From csrc/mamba diff --git a/python/sglang/kernels/aot/csrc/quantization/gguf/gguf_kernel.cu b/python/sglang/kernels/aot/csrc/quantization/gguf/gguf_kernel.cu index 1a4955f517a0..27bfc0a85743 100644 --- a/python/sglang/kernels/aot/csrc/quantization/gguf/gguf_kernel.cu +++ b/python/sglang/kernels/aot/csrc/quantization/gguf/gguf_kernel.cu @@ -18,6 +18,49 @@ #include "moe_vec.cuh" // clang-format off +#if !defined(USE_MUSA) +static int64_t ggml_mmq_k_alignment(const int64_t type) { + switch (type) { + case 8: + return 128; + case 10: + case 11: + return 512; + case 2: + case 3: + case 6: + case 7: + case 12: + case 13: + case 14: + case 16: + case 17: + case 18: + case 19: + case 20: + case 21: + case 22: + case 23: + return 256; + default: + return 0; + } +} + +static void check_ggml_mmq_input_size(const int64_t type, const int64_t input_size) { + const int64_t alignment = ggml_mmq_k_alignment(type); + TORCH_CHECK(alignment > 0, "Unsupported GGUF MMQ quantization type: ", type); + TORCH_CHECK( + input_size % alignment == 0, + "GGUF MMQ quantization type ", + type, + " requires an input size divisible by ", + alignment, + ", got ", + input_size); +} +#endif + // Q8 gemv template static __global__ void @@ -85,6 +128,7 @@ torch::Tensor ggml_dequantize( DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { auto to_cuda = ggml_get_to_cuda(type); + TORCH_CHECK(to_cuda != nullptr, "Unsupported GGUF dequantization type: ", type); to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream); }); @@ -184,6 +228,8 @@ torch::Tensor ggml_mul_mat_vec_a8( mul_mat_vec_iq1_m_q8_1_cuda( (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); break; + default: + TORCH_CHECK(false, "Unsupported GGUF MMVQ quantization type: ", type); } }); return Y; @@ -195,6 +241,9 @@ torch::Tensor ggml_mul_mat_a8( int64_t type, int64_t row) { int col = X.sizes()[1]; +#if !defined(USE_MUSA) + check_ggml_mmq_input_size(type, col); +#endif int padded = (col + 512 - 1) / 512 * 512; int batch = X.sizes()[0]; const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); @@ -327,6 +376,104 @@ torch::Tensor ggml_mul_mat_a8( row, stream); break; +#if !defined(USE_MUSA) + case 16: + ggml_mul_mat_iq2_xxs_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; + case 17: + ggml_mul_mat_iq2_xs_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; + case 18: + ggml_mul_mat_iq3_xxs_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; + case 19: + ggml_mul_mat_iq1_s_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; + case 20: + ggml_mul_mat_iq4_nl_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; + case 21: + ggml_mul_mat_iq3_s_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; + case 22: + ggml_mul_mat_iq2_s_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; + case 23: + ggml_mul_mat_iq4_xs_q8_1_cuda( + (void*)W.data_ptr(), + (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), + col, + row, + batch, + padded, + row, + stream); + break; +#endif } }); return Y; @@ -343,6 +490,9 @@ torch::Tensor ggml_moe_a8( int64_t top_k, int64_t tokens) { int col = X.sizes()[1]; +#if !defined(USE_MUSA) + check_ggml_mmq_input_size(type, col); +#endif 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()); @@ -533,6 +683,152 @@ torch::Tensor ggml_moe_a8( sorted_token_ids.sizes()[0], stream); break; +#if !defined(USE_MUSA) + case 16: + ggml_moe_iq2_xxs_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; + case 17: + ggml_moe_iq2_xs_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; + case 18: + ggml_moe_iq3_xxs_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; + case 19: + ggml_moe_iq1_s_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; + case 20: + ggml_moe_iq4_nl_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; + case 21: + ggml_moe_iq3_s_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; + case 22: + ggml_moe_iq2_s_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; + case 23: + ggml_moe_iq4_xs_q8_1_cuda( + (void*)quant_X.data_ptr(), + (void*)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); + break; +#endif } }); return Y; @@ -804,6 +1100,8 @@ torch::Tensor ggml_moe_a8_vec( quant_X.stride(0), stream); break; + default: + TORCH_CHECK(false, "Unsupported GGUF MoE vector quantization type: ", type); } }); return Y; @@ -831,6 +1129,32 @@ int64_t ggml_moe_get_block_size(int64_t type) { return MOE_X_Q5_K; case 14: return MOE_X_Q6_K; +#if !defined(USE_MUSA) + case 16: + return MOE_X_IQ2_XXS; + case 17: + return MOE_X_IQ2_XS; + case 18: + return MOE_X_IQ3_XXS; + case 19: + return MOE_X_IQ1_S; + case 20: + return MOE_X_IQ4_NL; + case 21: + return MOE_X_IQ3_S; + case 22: + return MOE_X_IQ2_S; + case 23: + return MOE_X_IQ4_XS; +#endif } return 0; } + +bool ggml_supports_iq_mmq() { +#if defined(USE_MUSA) + return false; +#else + return true; +#endif +} diff --git a/python/sglang/kernels/aot/csrc/quantization/gguf/mmq.cuh b/python/sglang/kernels/aot/csrc/quantization/gguf/mmq.cuh index 9838c14f0118..cc3e007536d7 100644 --- a/python/sglang/kernels/aot/csrc/quantization/gguf/mmq.cuh +++ b/python/sglang/kernels/aot/csrc/quantization/gguf/mmq.cuh @@ -879,3 +879,609 @@ static void ggml_mul_mat_q6_K_q8_1_cuda( <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); } } + +// Adapted from +// https://github.com/aimbit-ni/vllm/commit/037e1d547c15313aa7e2bb5fc83390ab0c857314. +#if !defined(USE_MUSA) + +#if defined(USE_ROCM) +#define MMQ_X_IQ4_NL 64 +#define MMQ_Y_IQ4_NL 128 +#define NWARPS_IQ4_NL 8 +#else +#define MMQ_X_IQ4_NL 4 +#define MMQ_Y_IQ4_NL 32 +#define NWARPS_IQ4_NL 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ4_NL, 2) +#endif + mul_mat_iq4_nl( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ4_NL; + const int mmq_y = MMQ_Y_IQ4_NL; + const int nwarps = NWARPS_IQ4_NL; + + mul_mat_q< + scalar_t, + QK4_NL, + QR4_NL, + QI4_NL, + true, + block_iq4_nl, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq4_nl, + load_tiles_iq4_nl, + VDR_IQ4_NL_Q8_1_MMQ, + vec_dot_iq4_nl_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq4_nl_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ4_NL; + const int mmq_y = MMQ_Y_IQ4_NL; + const int nwarps = NWARPS_IQ4_NL; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq4_nl + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq4_nl + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_IQ4_XS 64 +#define MMQ_Y_IQ4_XS 128 +#define NWARPS_IQ4_XS 8 +#else +#define MMQ_X_IQ4_XS 4 +#define MMQ_Y_IQ4_XS 32 +#define NWARPS_IQ4_XS 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ4_XS, 2) +#endif + mul_mat_iq4_xs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ4_XS; + const int mmq_y = MMQ_Y_IQ4_XS; + const int nwarps = NWARPS_IQ4_XS; + + mul_mat_q< + scalar_t, + QK_K, + QR_IQ4_XS_MMQ, + QI_IQ4_XS_MMQ, + true, + block_iq4_xs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq4_xs, + load_tiles_iq4_xs, + VDR_IQ4_XS_Q8_1_MMQ, + vec_dot_iq4_xs_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq4_xs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ4_XS; + const int mmq_y = MMQ_Y_IQ4_XS; + const int nwarps = NWARPS_IQ4_XS; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq4_xs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq4_xs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_IQ3_S 64 +#define MMQ_Y_IQ3_S 128 +#define NWARPS_IQ3_S 8 +#else +#define MMQ_X_IQ3_S 4 +#define MMQ_Y_IQ3_S 32 +#define NWARPS_IQ3_S 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ3_S, 2) +#endif + mul_mat_iq3_s( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ3_S; + const int mmq_y = MMQ_Y_IQ3_S; + const int nwarps = NWARPS_IQ3_S; + + mul_mat_q< + scalar_t, + QK_K, + QR_IQ3_S_MMQ, + QI_IQ3_S_MMQ, + true, + block_iq3_s, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq3_s, + load_tiles_iq3_s, + VDR_IQ3_S_Q8_1_MMQ, + vec_dot_iq3_s_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq3_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ3_S; + const int mmq_y = MMQ_Y_IQ3_S; + const int nwarps = NWARPS_IQ3_S; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq3_s + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq3_s + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_IQ3_XXS 64 +#define MMQ_Y_IQ3_XXS 128 +#define NWARPS_IQ3_XXS 8 +#else +#define MMQ_X_IQ3_XXS 4 +#define MMQ_Y_IQ3_XXS 32 +#define NWARPS_IQ3_XXS 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ3_XXS, 2) +#endif + mul_mat_iq3_xxs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ3_XXS; + const int mmq_y = MMQ_Y_IQ3_XXS; + const int nwarps = NWARPS_IQ3_XXS; + + mul_mat_q< + scalar_t, + QK_K, + QR_IQ3_XXS_MMQ, + QI_IQ3_XXS_MMQ, + true, + block_iq3_xxs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq3_xxs, + load_tiles_iq3_xxs, + VDR_IQ3_XXS_Q8_1_MMQ, + vec_dot_iq3_xxs_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq3_xxs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ3_XXS; + const int mmq_y = MMQ_Y_IQ3_XXS; + const int nwarps = NWARPS_IQ3_XXS; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq3_xxs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq3_xxs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_IQ2_XXS 64 +#define MMQ_Y_IQ2_XXS 128 +#define NWARPS_IQ2_XXS 8 +#else +#define MMQ_X_IQ2_XXS 4 +#define MMQ_Y_IQ2_XXS 32 +#define NWARPS_IQ2_XXS 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ2_XXS, 2) +#endif + mul_mat_iq2_xxs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ2_XXS; + const int mmq_y = MMQ_Y_IQ2_XXS; + const int nwarps = NWARPS_IQ2_XXS; + + mul_mat_q< + scalar_t, + QK_K, + QR_IQ2_XXS_MMQ, + QI_IQ2_XXS_MMQ, + true, + block_iq2_xxs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq2_xxs, + load_tiles_iq2_xxs, + VDR_IQ2_XXS_Q8_1_MMQ, + vec_dot_iq2_xxs_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq2_xxs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ2_XXS; + const int mmq_y = MMQ_Y_IQ2_XXS; + const int nwarps = NWARPS_IQ2_XXS; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq2_xxs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq2_xxs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_IQ2_XS 64 +#define MMQ_Y_IQ2_XS 128 +#define NWARPS_IQ2_XS 8 +#else +#define MMQ_X_IQ2_XS 4 +#define MMQ_Y_IQ2_XS 32 +#define NWARPS_IQ2_XS 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ2_XS, 2) +#endif + mul_mat_iq2_xs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ2_XS; + const int mmq_y = MMQ_Y_IQ2_XS; + const int nwarps = NWARPS_IQ2_XS; + + mul_mat_q< + scalar_t, + QK_K, + QR_IQ2_XS_MMQ, + QI_IQ2_XS_MMQ, + true, + block_iq2_xs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq2_xs, + load_tiles_iq2_xs, + VDR_IQ2_XS_Q8_1_MMQ, + vec_dot_iq2_xs_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq2_xs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ2_XS; + const int mmq_y = MMQ_Y_IQ2_XS; + const int nwarps = NWARPS_IQ2_XS; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq2_xs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq2_xs + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_IQ2_S 64 +#define MMQ_Y_IQ2_S 128 +#define NWARPS_IQ2_S 8 +#else +#define MMQ_X_IQ2_S 4 +#define MMQ_Y_IQ2_S 32 +#define NWARPS_IQ2_S 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ2_S, 2) +#endif + mul_mat_iq2_s( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ2_S; + const int mmq_y = MMQ_Y_IQ2_S; + const int nwarps = NWARPS_IQ2_S; + + mul_mat_q< + scalar_t, + QK_K, + QR_IQ2_S_MMQ, + QI_IQ2_S_MMQ, + true, + block_iq2_s, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq2_s, + load_tiles_iq2_s, + VDR_IQ2_S_Q8_1_MMQ, + vec_dot_iq2_s_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq2_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ2_S; + const int mmq_y = MMQ_Y_IQ2_S; + const int nwarps = NWARPS_IQ2_S; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq2_s + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq2_s + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_IQ1_S 64 +#define MMQ_Y_IQ1_S 128 +#define NWARPS_IQ1_S 8 +#else +#define MMQ_X_IQ1_S 4 +#define MMQ_Y_IQ1_S 32 +#define NWARPS_IQ1_S 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ1_S, 2) +#endif + mul_mat_iq1_s( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_IQ1_S; + const int mmq_y = MMQ_Y_IQ1_S; + const int nwarps = NWARPS_IQ1_S; + + mul_mat_q< + scalar_t, + QK_K, + QR_IQ1_S_MMQ, + QI_IQ1_S_MMQ, + true, + block_iq1_s, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq1_s, + load_tiles_iq1_s, + VDR_IQ1_S_Q8_1_MMQ, + vec_dot_iq1_s_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_iq1_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + cudaStream_t stream) { + const int mmq_x = MMQ_X_IQ1_S; + const int mmq_y = MMQ_Y_IQ1_S; + const int nwarps = NWARPS_IQ1_S; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + mul_mat_iq1_s + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + mul_mat_iq1_s + <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#endif // !defined(USE_MUSA) diff --git a/python/sglang/kernels/aot/csrc/quantization/gguf/moe.cuh b/python/sglang/kernels/aot/csrc/quantization/gguf/moe.cuh index 91e434fa6080..cb0761e268f5 100644 --- a/python/sglang/kernels/aot/csrc/quantization/gguf/moe.cuh +++ b/python/sglang/kernels/aot/csrc/quantization/gguf/moe.cuh @@ -25,7 +25,7 @@ static __device__ __forceinline__ void moe_q( const int* __restrict__ sorted_token_ids, const int* __restrict__ expert_ids, const int* __restrict__ num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -43,16 +43,34 @@ static __device__ __forceinline__ void moe_q( const auto col_dst_0 = blockIdx.y * mmq_x; + if (col_dst_0 >= num_tokens_post_padded[0]) return; + int token_offs[mmq_x / nwarps]; for (int i = 0; i < mmq_x; i += nwarps) { token_offs[i / nwarps] = sorted_token_ids[col_dst_0 + threadIdx.y + i]; } const int exp_idx = expert_ids[blockIdx.y]; - if (exp_idx > 255 || exp_idx < 0) return; - if (blockIdx.y * mmq_x > num_tokens_post_padded[0]) return; + if (exp_idx > 255 || exp_idx < 0) { +#pragma unroll + for (int j = 0; j < mmq_x; j += nwarps) { + const int col_dst = token_offs[j / nwarps]; + if (col_dst >= ncols_dst) { + continue; + } + +#pragma unroll + for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { + const int row_dst = row_dst_0 + threadIdx.x + i; + if (row_dst < nrows_dst) { + dst[col_dst * nrows_dst + row_dst] = scalar_t(0); + } + } + } + return; + } - const block_q_t* x = (const block_q_t*)((char*)vx + exp_idx * exp_stride); + const block_q_t* x = (const block_q_t*)((char*)vx + static_cast(exp_idx) * exp_stride); const block_q8_1* y = (const block_q8_1*)(vy); int* tile_x_ql = nullptr; @@ -96,14 +114,19 @@ static __device__ __forceinline__ void moe_q( } } - if (threadIdx.x < n_per_r / QK8_1) { + // Load every destination column's q8 scale. token_offs is local to one + // warp, so indexing it by a global tile column is invalid when + // mmq_x > nwarps and out of bounds when mmq_x == nwarps. +#pragma unroll + for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) { + const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE_GGUF / QI8_1)) % mmq_x; const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1); - const int col_y_eff = token_offs[threadIdx.y] / top_k; + const int col_y_eff = sorted_token_ids[col_dst_0 + ids] / top_k; const int block_x = ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby; if (col_y_eff < ncols_y && block_x < blocks_per_col_y) { const half2* dsi_src = &y[col_y_eff * blocks_per_col_y + block_x].ds; - half2* dsi_dst = &tile_y_ds[threadIdx.y * (WARP_SIZE_GGUF / QI8_1) + kby]; + half2* dsi_dst = &tile_y_ds[ids * (WARP_SIZE_GGUF / QI8_1) + kby]; if (need_sum) { *dsi_dst = *dsi_src; @@ -170,7 +193,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -218,7 +241,7 @@ static void ggml_moe_q4_0_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -293,7 +316,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -341,7 +364,7 @@ static void ggml_moe_q4_1_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -416,7 +439,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -464,7 +487,7 @@ static void ggml_moe_q5_0_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -539,7 +562,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -587,7 +610,7 @@ static void ggml_moe_q5_1_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -662,7 +685,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -710,7 +733,7 @@ static void ggml_moe_q8_0_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -785,7 +808,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -833,7 +856,7 @@ static void ggml_moe_q2_K_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -908,7 +931,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -956,7 +979,7 @@ static void ggml_moe_q3_K_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -1031,7 +1054,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -1079,7 +1102,7 @@ static void ggml_moe_q4_K_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -1154,7 +1177,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -1202,7 +1225,7 @@ static void ggml_moe_q5_K_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -1277,7 +1300,7 @@ __launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2) const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -1325,7 +1348,7 @@ static void ggml_moe_q6_K_q8_1_cuda( const int* sorted_token_ids, const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, + const int64_t exp_stride, const int ncols_x, const int nrows_x, const int ncols_y, @@ -1377,3 +1400,993 @@ static void ggml_moe_q6_K_q8_1_cuda( top_k); } } + +// Adapted from +// https://github.com/aimbit-ni/vllm/commit/037e1d547c15313aa7e2bb5fc83390ab0c857314. +#if !defined(USE_MUSA) + +#if defined(USE_ROCM) +#define MOE_X_IQ4_NL 64 +#define MOE_Y_IQ4_NL 128 +#define NWARPS_IQ4_NL 8 +#else +#define MOE_X_IQ4_NL 4 +#define MOE_Y_IQ4_NL 32 +#define NWARPS_IQ4_NL 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ4_NL, 2) +#endif + moe_iq4_nl( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ4_NL; + const int mmq_y = MOE_Y_IQ4_NL; + const int nwarps = NWARPS_IQ4_NL; + + moe_q< + scalar_t, + QK4_NL, + QR4_NL, + QI4_NL, + true, + block_iq4_nl, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq4_nl, + load_tiles_iq4_nl, + VDR_IQ4_NL_Q8_1_MMQ, + vec_dot_iq4_nl_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq4_nl_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ4_NL; + const int mmq_y = MOE_Y_IQ4_NL; + const int nwarps = NWARPS_IQ4_NL; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq4_nl<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq4_nl<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_IQ4_XS 64 +#define MOE_Y_IQ4_XS 128 +#define NWARPS_IQ4_XS 8 +#else +#define MOE_X_IQ4_XS 4 +#define MOE_Y_IQ4_XS 32 +#define NWARPS_IQ4_XS 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ4_XS, 2) +#endif + moe_iq4_xs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ4_XS; + const int mmq_y = MOE_Y_IQ4_XS; + const int nwarps = NWARPS_IQ4_XS; + + moe_q< + scalar_t, + QK_K, + QR_IQ4_XS_MMQ, + QI_IQ4_XS_MMQ, + true, + block_iq4_xs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq4_xs, + load_tiles_iq4_xs, + VDR_IQ4_XS_Q8_1_MMQ, + vec_dot_iq4_xs_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq4_xs_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ4_XS; + const int mmq_y = MOE_Y_IQ4_XS; + const int nwarps = NWARPS_IQ4_XS; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq4_xs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq4_xs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_IQ3_S 64 +#define MOE_Y_IQ3_S 128 +#define NWARPS_IQ3_S_MOE 8 +#else +#define MOE_X_IQ3_S 4 +#define MOE_Y_IQ3_S 32 +#define NWARPS_IQ3_S_MOE 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ3_S_MOE, 2) +#endif + moe_iq3_s( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ3_S; + const int mmq_y = MOE_Y_IQ3_S; + const int nwarps = NWARPS_IQ3_S_MOE; + + moe_q< + scalar_t, + QK_K, + QR_IQ3_S_MMQ, + QI_IQ3_S_MMQ, + true, + block_iq3_s, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq3_s, + load_tiles_iq3_s, + VDR_IQ3_S_Q8_1_MMQ, + vec_dot_iq3_s_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq3_s_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ3_S; + const int mmq_y = MOE_Y_IQ3_S; + const int nwarps = NWARPS_IQ3_S_MOE; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq3_s<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq3_s<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_IQ3_XXS 64 +#define MOE_Y_IQ3_XXS 128 +#define NWARPS_IQ3_XXS_MOE 8 +#else +#define MOE_X_IQ3_XXS 4 +#define MOE_Y_IQ3_XXS 32 +#define NWARPS_IQ3_XXS_MOE 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ3_XXS_MOE, 2) +#endif + moe_iq3_xxs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ3_XXS; + const int mmq_y = MOE_Y_IQ3_XXS; + const int nwarps = NWARPS_IQ3_XXS_MOE; + + moe_q< + scalar_t, + QK_K, + QR_IQ3_XXS_MMQ, + QI_IQ3_XXS_MMQ, + true, + block_iq3_xxs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq3_xxs, + load_tiles_iq3_xxs, + VDR_IQ3_XXS_Q8_1_MMQ, + vec_dot_iq3_xxs_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq3_xxs_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ3_XXS; + const int mmq_y = MOE_Y_IQ3_XXS; + const int nwarps = NWARPS_IQ3_XXS_MOE; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq3_xxs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq3_xxs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_IQ2_XXS 64 +#define MOE_Y_IQ2_XXS 128 +#define NWARPS_IQ2_XXS_MOE 8 +#else +#define MOE_X_IQ2_XXS 4 +#define MOE_Y_IQ2_XXS 32 +#define NWARPS_IQ2_XXS_MOE 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ2_XXS_MOE, 2) +#endif + moe_iq2_xxs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ2_XXS; + const int mmq_y = MOE_Y_IQ2_XXS; + const int nwarps = NWARPS_IQ2_XXS_MOE; + + moe_q< + scalar_t, + QK_K, + QR_IQ2_XXS_MMQ, + QI_IQ2_XXS_MMQ, + true, + block_iq2_xxs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq2_xxs, + load_tiles_iq2_xxs, + VDR_IQ2_XXS_Q8_1_MMQ, + vec_dot_iq2_xxs_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq2_xxs_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ2_XXS; + const int mmq_y = MOE_Y_IQ2_XXS; + const int nwarps = NWARPS_IQ2_XXS_MOE; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq2_xxs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq2_xxs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_IQ2_XS 64 +#define MOE_Y_IQ2_XS 128 +#define NWARPS_IQ2_XS_MOE 8 +#else +#define MOE_X_IQ2_XS 4 +#define MOE_Y_IQ2_XS 32 +#define NWARPS_IQ2_XS_MOE 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ2_XS_MOE, 2) +#endif + moe_iq2_xs( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ2_XS; + const int mmq_y = MOE_Y_IQ2_XS; + const int nwarps = NWARPS_IQ2_XS_MOE; + + moe_q< + scalar_t, + QK_K, + QR_IQ2_XS_MMQ, + QI_IQ2_XS_MMQ, + true, + block_iq2_xs, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq2_xs, + load_tiles_iq2_xs, + VDR_IQ2_XS_Q8_1_MMQ, + vec_dot_iq2_xs_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq2_xs_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ2_XS; + const int mmq_y = MOE_Y_IQ2_XS; + const int nwarps = NWARPS_IQ2_XS_MOE; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq2_xs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq2_xs<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_IQ2_S 64 +#define MOE_Y_IQ2_S 128 +#define NWARPS_IQ2_S_MOE 8 +#else +#define MOE_X_IQ2_S 4 +#define MOE_Y_IQ2_S 32 +#define NWARPS_IQ2_S_MOE 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ2_S_MOE, 2) +#endif + moe_iq2_s( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ2_S; + const int mmq_y = MOE_Y_IQ2_S; + const int nwarps = NWARPS_IQ2_S_MOE; + + moe_q< + scalar_t, + QK_K, + QR_IQ2_S_MMQ, + QI_IQ2_S_MMQ, + true, + block_iq2_s, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq2_s, + load_tiles_iq2_s, + VDR_IQ2_S_Q8_1_MMQ, + vec_dot_iq2_s_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq2_s_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ2_S; + const int mmq_y = MOE_Y_IQ2_S; + const int nwarps = NWARPS_IQ2_S_MOE; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq2_s<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq2_s<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_IQ1_S 64 +#define MOE_Y_IQ1_S 128 +#define NWARPS_IQ1_S_MOE 8 +#else +#define MOE_X_IQ1_S 4 +#define MOE_Y_IQ1_S 32 +#define NWARPS_IQ1_S_MOE 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_IQ1_S_MOE, 2) +#endif + moe_iq1_s( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_IQ1_S; + const int mmq_y = MOE_Y_IQ1_S; + const int nwarps = NWARPS_IQ1_S_MOE; + + moe_q< + scalar_t, + QK_K, + QR_IQ1_S_MMQ, + QI_IQ1_S_MMQ, + true, + block_iq1_s, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_iq1_s, + load_tiles_iq1_s, + VDR_IQ1_S_Q8_1_MMQ, + vec_dot_iq1_s_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_iq1_s_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int64_t exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + cudaStream_t stream) { + const int mmq_x = MOE_X_IQ1_S; + const int mmq_y = MOE_Y_IQ1_S; + const int nwarps = NWARPS_IQ1_S_MOE; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = tokens_post_padded / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + moe_iq1_s<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + moe_iq1_s<<>>( + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#endif // !defined(USE_MUSA) diff --git a/python/sglang/kernels/aot/csrc/quantization/gguf/moe_vec.cuh b/python/sglang/kernels/aot/csrc/quantization/gguf/moe_vec.cuh index 8cef9e080a22..b8b37dcba9b8 100644 --- a/python/sglang/kernels/aot/csrc/quantization/gguf/moe_vec.cuh +++ b/python/sglang/kernels/aot/csrc/quantization/gguf/moe_vec.cuh @@ -17,7 +17,7 @@ static __global__ void moe_vec_q( const auto token = blockIdx.z / topk; const auto expert = (topk_ids)[blockIdx.z]; - if (row >= nrows) { + if (row >= nrows || expert < 0) { return; } diff --git a/python/sglang/kernels/aot/csrc/quantization/gguf/vecdotq.cuh b/python/sglang/kernels/aot/csrc/quantization/gguf/vecdotq.cuh index 08b4cd269c7b..a86694af0f42 100644 --- a/python/sglang/kernels/aot/csrc/quantization/gguf/vecdotq.cuh +++ b/python/sglang/kernels/aot/csrc/quantization/gguf/vecdotq.cuh @@ -5,10 +5,10 @@ static __device__ __forceinline__ int get_int_b2(const void* x, const int& i32) { const uint16_t* x16 = (const uint16_t*)x; // assume at least 2 byte alignment - int x32 = x16[2 * i32 + 0] << 0; - x32 |= x16[2 * i32 + 1] << 16; + uint32_t x32 = uint32_t(x16[2 * i32 + 0]); + x32 |= uint32_t(x16[2 * i32 + 1]) << 16; - return x32; + return static_cast(x32); } static __device__ __forceinline__ int get_int_b4(const void* x, const int& i32) { @@ -17,18 +17,16 @@ static __device__ __forceinline__ int get_int_b4(const void* x, const int& i32) static __device__ __forceinline__ int get_int_from_int8(const int8_t* x8, const int& i32) { const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32); // assume at least 2 byte alignment - int x32 = 0; - x32 |= x16[0] << 0; - x32 |= x16[1] << 16; - return x32; + uint32_t x32 = uint32_t(x16[0]); + x32 |= uint32_t(x16[1]) << 16; + return static_cast(x32); } static __device__ __forceinline__ int get_int_from_uint8(const uint8_t* x8, const int& i32) { const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32); // assume at least 2 byte alignment - int x32 = 0; - x32 |= x16[0] << 0; - x32 |= x16[1] << 16; - return x32; + uint32_t x32 = uint32_t(x16[0]); + x32 |= uint32_t(x16[1]) << 16; + return static_cast(x32); } static __device__ __forceinline__ int get_int_from_int8_aligned(const int8_t* x8, const int& i32) { @@ -1759,7 +1757,7 @@ vec_dot_iq2_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict_ const uint16_t* q2 = bq2->qs + 4 * ib32; const uint8_t* aux8 = (const uint8_t*)q2; const int8_t* q8 = bq8_1[ib32].qs; - uint32_t aux32 = q2[2] | (q2[3] << 16); + uint32_t aux32 = uint32_t(q2[2]) | (uint32_t(q2[3]) << 16); int sumi = 0; for (int l = 0; l < 4; ++l) { const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[l]); @@ -1853,7 +1851,7 @@ vec_dot_iq3_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict_ const uint8_t* q3 = bq2->qs + 8 * ib32; const uint16_t* gas = (const uint16_t*)(bq2->qs + QK_K / 4) + 2 * ib32; const int8_t* q8 = bq8_1[ib32].qs; - uint32_t aux32 = gas[0] | (gas[1] << 16); + uint32_t aux32 = uint32_t(gas[0]) | (uint32_t(gas[1]) << 16); int sumi = 0; for (int l = 0; l < 4; ++l) { const uint32_t* grid1 = iq3xxs_grid + q3[2 * l + 0]; @@ -1982,11 +1980,11 @@ get_int_from_table_16(const uint32_t& q4, const uint8_t* values, int& val1, int& aux32 = q4 & 0x0f0f0f0f; uint16_t v1 = values[q8[0]] | (values[q8[1]] << 8); uint16_t v2 = values[q8[2]] | (values[q8[3]] << 8); - val1 = v1 | (v2 << 16); + val1 = static_cast(uint32_t(v1) | (uint32_t(v2) << 16)); aux32 = (q4 >> 4) & 0x0f0f0f0f; v1 = values[q8[0]] | (values[q8[1]] << 8); v2 = values[q8[2]] | (values[q8[3]] << 8); - val2 = v1 | (v2 << 16); + val2 = static_cast(uint32_t(v1) | (uint32_t(v2) << 16)); } static __device__ __forceinline__ float @@ -2003,7 +2001,7 @@ vec_dot_iq4_nl_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ int v1, v2; int sumi1 = 0, sumi2 = 0; for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) { - const uint32_t aux = q4[2 * l] | (q4[2 * l + 1] << 16); + const uint32_t aux = uint32_t(q4[2 * l]) | (uint32_t(q4[2 * l + 1]) << 16); get_int_from_table_16(aux, values, v1, v2); sumi1 = __dp4a(v1, q8[l + 0], sumi1); sumi2 = __dp4a(v2, q8[l + 4], sumi2); @@ -2035,3 +2033,1069 @@ vec_dot_iq4_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ return d * (sumi1 + sumi2); #endif } + +// I-quant MMQ tile operations adapted from +// https://github.com/aimbit-ni/vllm/commit/037e1d547c15313aa7e2bb5fc83390ab0c857314 +// and https://github.com/ggml-org/llama.cpp/pull/8495. +#if !defined(USE_MUSA) +static __device__ __forceinline__ int iq_mmq_dp4a(const int a, const int b, int c) { +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 610) || defined(USE_ROCM) + return __dp4a(a, b, c); +#else + const int8_t* va = reinterpret_cast(&a); + const int8_t* vb = reinterpret_cast(&b); +#pragma unroll + for (int i = 0; i < 4; ++i) { + c += va[i] * vb[i]; + } + return c; +#endif +} + +// ========================= IQ4_NL MMQ ========================= +// IQ4_NL is structurally identical to Q4_0 (QK=32, QR=2, QI=4) but uses +// a 16-entry lookup table (kvalues_iq4nl) to convert 4-bit indices to signed int8. +// Tiles store raw packed nibble data (same as Q4_0); the lookup is applied in vec_dot. +// need_sum=true keeps the q8_1 (d, sum) tile layout; signed lookup values use only d. + +#define VDR_IQ4_NL_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq4_nl(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI4_NL) + mmq_y / QI4_NL]; + + *x_ql = tile_x_qs; + *x_dm = (half2*)tile_x_d; +} + +template +static __device__ __forceinline__ void load_tiles_iq4_nl( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI4_NL; + const int kqsx = k % QI4_NL; + + const block_iq4_nl* bx0 = (const block_iq4_nl*)vx; + float* x_dmf = (float*)x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq4_nl* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_NL; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_NL) { + int i = i0 + i_offset * QI4_NL + k / blocks_per_tile_x_row; + if (need_check) { + i = min(i, i_max); + } + const block_iq4_nl* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI4_NL) + i / QI4_NL + kbxd] = __half2float(bxi->d); + } +} + +static __device__ __forceinline__ float vec_dot_iq4_nl_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_qh; + (void)x_sc; + + const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); + const float* x_dmf = (const float*)x_dm; + + int u[2 * VDR_IQ4_NL_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_IQ4_NL_Q8_1_MMQ; ++l) { + u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; + u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_NL) % WARP_SIZE_GGUF]; + } + + int sumi = 0; + + // IQ4_NL lookup table values (same as kvalues_iq4nl in ggml-common.h). + const int8_t iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + +#pragma unroll + for (int l = 0; l < VDR_IQ4_NL_Q8_1_MMQ; ++l) { + int v1, v2; + get_int_from_table_16(x_ql[i * (WARP_SIZE_GGUF + 1) + k + l], (const uint8_t*)iq4nl, v1, v2); + sumi = iq_mmq_dp4a(v1, u[2 * l + 0], sumi); + sumi = iq_mmq_dp4a(v2, u[2 * l + 1], sumi); + } + + // need_sum=true: y_ds is stored as half2 (d, sum). We only need the d component + // since IQ4_NL values are already signed (no unsigned-to-signed bias correction). + const float2 ds8 = __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return x_dmf[i * (WARP_SIZE_GGUF / QI4_NL) + i / QI4_NL + k / QI4_NL] * ds8.x * sumi; +} + +// ========================= IQ4_XS MMQ ========================= +// IQ4_XS is a 256-element super-block (QK_K=256) with 8 sub-blocks of 32 elements. +// Each sub-block has a 6-bit scale: 4 low bits from scales_l + 2 high bits from scales_h. +// For MMQ, we use Q4_K-like tiling: QR=2, QI=32 (blocks_per_warp=1). +// This ensures all 128 bytes of qs are loaded into shared memory per step. +// The MMVQ constants (QR4_XS=8, QI4_XS=8) are NOT used for MMQ. + +#define QR_IQ4_XS_MMQ 2 +#define QI_IQ4_XS_MMQ (QK_K / (4 * QR_IQ4_XS_MMQ)) +#define VDR_IQ4_XS_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq4_xs(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI_IQ4_XS_MMQ) + mmq_y / QI_IQ4_XS_MMQ]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_iq4_xs( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI_IQ4_XS_MMQ; + const int kqsx = k % QI_IQ4_XS_MMQ; + + const block_iq4_xs* bx0 = (const block_iq4_xs*)vx; + float* x_dmf = (float*)x_dm; + + // Load qs data (packed nibbles, 128 bytes per block) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq4_xs* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); + } + + // Load global d scale (1 float per block) + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI_IQ4_XS_MMQ; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI_IQ4_XS_MMQ) { + int i = (i0 + i_offset * QI_IQ4_XS_MMQ + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq4_xs* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI_IQ4_XS_MMQ) + i / QI_IQ4_XS_MMQ + kbxd] = __half2float(bxi->d); + } + + // Load sub-block scales (8 per block, decoded to int8, packed 4 per int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq4_xs* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / (QI_IQ4_XS_MMQ / 8); + + const int ksc = k % (WARP_SIZE_GGUF / 8); + + // Pack 4 decoded sub-block scales per int32 (only need ksc=0,1 for 8 scales) + if (ksc < 2) { + int scales_packed = 0; +#pragma unroll + for (int sb = 0; sb < 4; sb++) { + const int ib32 = ksc * 4 + sb; + const int8_t ls = + ((bxi->scales_l[ib32 / 2] >> (4 * (ib32 % 2))) & 0xf) | (((bxi->scales_h >> (2 * ib32)) & 3) << 4); + ((int8_t*)&scales_packed)[sb] = ls - 32; + } + x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = scales_packed; + } + } +} + +static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_qh; + + const float* x_dmf = (const float*)x_dm; + + // IQ4_NL lookup table values (same as kvalues_iq4nl in ggml-common.h). + const int8_t iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + + // Sub-block index: each sub-block = 4 int32 of packed nibble data = 32 elements + const int ib32 = k / 4; + const int8_t* sc_bytes = (const int8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ib32 / 4]; + const float sub_scale = (float)sc_bytes[ib32 % 4]; + + const int index_y = j * WARP_SIZE_GGUF + (QR_IQ4_XS_MMQ * k) % WARP_SIZE_GGUF; + + int sumi = 0; + +#pragma unroll + for (int l = 0; l < VDR_IQ4_XS_Q8_1_MMQ; ++l) { + int v1, v2; + get_int_from_table_16(x_ql[i * (WARP_SIZE_GGUF + 1) + k + l], (const uint8_t*)iq4nl, v1, v2); + sumi = iq_mmq_dp4a(v1, y_qs[index_y + l], sumi); + sumi = iq_mmq_dp4a(v2, y_qs[index_y + l + 4], sumi); + } + + // need_sum=true: y_ds is stored as half2 (d, sum). We only need the d component + // since IQ4_NL values are already signed (no unsigned-to-signed bias correction). + const float2 ds8 = + __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (QR_IQ4_XS_MMQ * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return x_dmf[i * (WARP_SIZE_GGUF / QI_IQ4_XS_MMQ) + i / QI_IQ4_XS_MMQ] * sub_scale * ds8.x * sumi; +} + +// ========================= IQ3_S MMQ ========================= +// IQ3_S: QK_K=256, qs[64], qh[8], signs[32], scales[4], d(half). +// Custom QR=2, QI=32 for MMQ (1 block/warp, like IQ4_XS_MMQ). +// tile_x_ql[0..15]: qs (64 bytes), tile_x_ql[16..23]: signs (32 bytes). +// tile_x_qh: qh (8 bytes). tile_x_sc: scales (4 bytes). tile_x_dm: d. +// Grid lookup (iq3xs_grid[512]) in vec_dot. + +#define QR_IQ3_S_MMQ 2 +#define QI_IQ3_S_MMQ (QK_K / (4 * QR_IQ3_S_MMQ)) +#define VDR_IQ3_S_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq3_s(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI_IQ3_S_MMQ) + mmq_y / QI_IQ3_S_MMQ]; + __shared__ int tile_x_qh[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; + *x_qh = tile_x_qh; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_iq3_s( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI_IQ3_S_MMQ; // 0 (1 block per warp) + const int kqsx = k % QI_IQ3_S_MMQ; // 0..31 + + const block_iq3_s* bx0 = (const block_iq3_s*)vx; + float* x_dmf = (float*)x_dm; + + // Load qs (64 bytes = 16 int32, slots 0..15) and signs (32 bytes = 8 int32, slots 16..23) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq3_s* bxi = bx0 + i * blocks_per_row + kbx; + if (kqsx < 16) { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + } else if (kqsx < 24) { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->signs, kqsx - 16); + } else { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = 0; + } + } + + // Load d scale + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI_IQ3_S_MMQ; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI_IQ3_S_MMQ) { + int i = (i0 + i_offset * QI_IQ3_S_MMQ + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq3_s* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI_IQ3_S_MMQ) + i / QI_IQ3_S_MMQ + kbxd] = __half2float(bxi->d); + } + + // Load qh (8 bytes = 2 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { + int i = (i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4)) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const int kqh = k % (WARP_SIZE_GGUF / 4); + if (kqh < 2) { + const block_iq3_s* bxi = bx0 + i * blocks_per_row + kbx; + x_qh[i * (WARP_SIZE_GGUF / 4) + i / 4 + kqh] = get_int_from_uint8(bxi->qh, kqh); + } + } + + // Load scales (4 bytes = 1 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const int ksc = k % (WARP_SIZE_GGUF / 8); + if (ksc < 1) { + const block_iq3_s* bxi = bx0 + i * blocks_per_row + kbx; + x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = get_int_from_uint8(bxi->scales, ksc); + } + } +} + +static __device__ __forceinline__ float vec_dot_iq3_s_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const float* x_dmf = (const float*)x_dm; + const int ib32 = k / 4; // sub-block index (same for all VDR iterations since k is multiple of 4) + + // Read qs base for this ib32 (8 bytes of grid indices) + const uint8_t* qs = (const uint8_t*)&x_ql[i * (WARP_SIZE_GGUF + 1) + ib32 * 2]; + // Read qh byte for this ib32 + const uint8_t* qh_all = (const uint8_t*)&x_qh[i * (WARP_SIZE_GGUF / 4) + i / 4]; + const uint8_t qh_val = qh_all[ib32]; + // Signs base + const uint8_t* all_signs = (const uint8_t*)&x_ql[i * (WARP_SIZE_GGUF + 1) + 16]; + // Per-sub-block scale + const uint8_t* sc_bytes = (const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8]; + const float sub_scale = 0.5f + ((sc_bytes[ib32 / 2] >> (4 * (ib32 % 2))) & 0xf); + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_IQ3_S_Q8_1_MMQ; ++l) { + const int sub = l; // k%4=0 always, so (k+l)%4 = l + + // Grid lookup: 9-bit index + const int grid1_val = iq3xs_grid[qs[2 * sub + 0] | ((qh_val << (8 - 2 * sub)) & 256)]; + const int grid2_val = iq3xs_grid[qs[2 * sub + 1] | ((qh_val << (7 - 2 * sub)) & 256)]; + + // Sign extraction + const uint8_t sign_byte = all_signs[4 * ib32 + sub]; + const uint32_t signs0 = __vcmpeq4(((sign_byte & 0xf) * 0x01010101) & 0x08040201, 0x08040201); + const uint32_t signs1 = __vcmpeq4(((sign_byte >> 4) * 0x01010101) & 0x08040201, 0x08040201); + const int grid_l = __vsub4(grid1_val ^ signs0, signs0); + const int grid_h = __vsub4(grid2_val ^ signs1, signs1); + + const int index_y = j * WARP_SIZE_GGUF + (QR_IQ3_S_MMQ * (k + l)) % WARP_SIZE_GGUF; + sumi = iq_mmq_dp4a(grid_l, y_qs[index_y + 0], sumi); + sumi = iq_mmq_dp4a(grid_h, y_qs[index_y + 1], sumi); + } + + const float2 ds8 = + __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (QR_IQ3_S_MMQ * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return x_dmf[i * (WARP_SIZE_GGUF / QI_IQ3_S_MMQ) + i / QI_IQ3_S_MMQ] * sub_scale * 0.5f * ds8.x * sumi; +} + +// ========================= IQ3_XXS MMQ ========================= +// IQ3_XXS: QK_K=256, qs[3*(QK_K/8)=96 bytes]: first 64 = grid indices, last 32 = gas (signs+scales). +// Grid: iq3xxs_grid[256] (uint32). Sign via ksigns64. Scale from gas aux bits. +// Same custom QR=2, QI=32 approach. + +#define QR_IQ3_XXS_MMQ 2 +#define QI_IQ3_XXS_MMQ (QK_K / (4 * QR_IQ3_XXS_MMQ)) +#define VDR_IQ3_XXS_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq3_xxs(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + // tile_x_ql: slots 0..15 = grid indices (64 bytes), slots 16..23 = gas data (32 bytes) + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI_IQ3_XXS_MMQ) + mmq_y / QI_IQ3_XXS_MMQ]; + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; +} + +template +static __device__ __forceinline__ void load_tiles_iq3_xxs( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + (void)x_qh; + (void)x_sc; + + const int kbx = k / QI_IQ3_XXS_MMQ; + const int kqsx = k % QI_IQ3_XXS_MMQ; + + const block_iq3_xxs* bx0 = (const block_iq3_xxs*)vx; + float* x_dmf = (float*)x_dm; + + // Load grid indices (64 bytes, slots 0..15) and gas data (32 bytes, slots 16..23) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq3_xxs* bxi = bx0 + i * blocks_per_row + kbx; + if (kqsx < 16) { + // Grid indices: first 64 bytes of qs + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + } else if (kqsx < 24) { + // Gas data (signs+scales): bytes 64..95 of qs + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs + QK_K / 4, kqsx - 16); + } else { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = 0; + } + } + + // Load d scale + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI_IQ3_XXS_MMQ; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI_IQ3_XXS_MMQ) { + int i = (i0 + i_offset * QI_IQ3_XXS_MMQ + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq3_xxs* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI_IQ3_XXS_MMQ) + i / QI_IQ3_XXS_MMQ + kbxd] = __half2float(bxi->d); + } +} + +static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_qh; + (void)x_sc; + + const float* x_dmf = (const float*)x_dm; + + const int ib32 = k / 4; + + // Grid indices for this ib32 + const uint8_t* q3 = (const uint8_t*)&x_ql[i * (WARP_SIZE_GGUF + 1) + ib32 * 2]; + + // Gas data for this ib32 (signs + scale packed in 4 bytes per ib32) + const uint16_t* gas = (const uint16_t*)&x_ql[i * (WARP_SIZE_GGUF + 1) + 16 + ib32]; + uint32_t aux32 = uint32_t(gas[0]) | (uint32_t(gas[1]) << 16); + const float sub_scale = 0.5f + (aux32 >> 28); + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_IQ3_XXS_Q8_1_MMQ; ++l) { + const int sub = l; + + // Extract signs for this sub-group + uint32_t signs_val = aux32 >> (7 * sub); + const uint32_t* signs = (const uint32_t*)(ksigns64 + (signs_val & 127)); + + // Grid lookup + const uint32_t* grid1 = iq3xxs_grid + q3[2 * sub + 0]; + const uint32_t* grid2 = iq3xxs_grid + q3[2 * sub + 1]; + + // Apply signs + const int grid_l = __vsub4(grid1[0] ^ signs[0], signs[0]); + const int grid_h = __vsub4(grid2[0] ^ signs[1], signs[1]); + + // Q8 data + const int index_y = j * WARP_SIZE_GGUF + (QR_IQ3_XXS_MMQ * (k + l)) % WARP_SIZE_GGUF; + sumi = iq_mmq_dp4a(grid_l, y_qs[index_y + 0], sumi); + sumi = iq_mmq_dp4a(grid_h, y_qs[index_y + 1], sumi); + } + + const float2 ds8 = + __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (QR_IQ3_XXS_MMQ * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return x_dmf[i * (WARP_SIZE_GGUF / QI_IQ3_XXS_MMQ) + i / QI_IQ3_XXS_MMQ] * sub_scale * 0.5f * ds8.x * sumi; +} + +// ========================= IQ2_XXS MMQ ========================= +// IQ2_XXS: QK_K=256, qs[QK_K/8=32 uint16]. Grid: iq2xxs_grid[256] (uint64). +// Per ib32: 4 uint16 from qs. First 2 = grid indices (8-bit each). Last 2 = aux (signs+scale). + +#define QR_IQ2_XXS_MMQ 2 +#define QI_IQ2_XXS_MMQ (QK_K / (4 * QR_IQ2_XXS_MMQ)) +#define VDR_IQ2_XXS_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq2_xxs(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + // qs data: 32 uint16 = 64 bytes = 16 int32 in slots 0..15 + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI_IQ2_XXS_MMQ) + mmq_y / QI_IQ2_XXS_MMQ]; + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; +} + +template +static __device__ __forceinline__ void load_tiles_iq2_xxs( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + (void)x_qh; + (void)x_sc; + + const int kbx = k / QI_IQ2_XXS_MMQ; + const int kqsx = k % QI_IQ2_XXS_MMQ; + + const block_iq2_xxs* bx0 = (const block_iq2_xxs*)vx; + float* x_dmf = (float*)x_dm; + + // Load qs data (32 uint16 = 64 bytes = 16 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq2_xxs* bxi = bx0 + i * blocks_per_row + kbx; + if (kqsx < 16) { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_b2(bxi->qs, kqsx); + } else { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = 0; + } + } + + // Load d scale + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI_IQ2_XXS_MMQ; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI_IQ2_XXS_MMQ) { + int i = (i0 + i_offset * QI_IQ2_XXS_MMQ + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq2_xxs* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI_IQ2_XXS_MMQ) + i / QI_IQ2_XXS_MMQ + kbxd] = __half2float(bxi->d); + } +} + +static __device__ __forceinline__ float vec_dot_iq2_xxs_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_qh; + (void)x_sc; + + const float* x_dmf = (const float*)x_dm; + + const int ib32 = k / 4; + + // Read the 4 uint16 for this ib32 + const uint16_t* q2 = (const uint16_t*)&x_ql[i * (WARP_SIZE_GGUF + 1) + ib32 * 2]; + const uint8_t* aux8 = (const uint8_t*)q2; + uint32_t aux32 = uint32_t(q2[2]) | (uint32_t(q2[3]) << 16); + const float sub_scale = 0.5f + (aux32 >> 28); + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_IQ2_XXS_Q8_1_MMQ; ++l) { + const int sub = l; + const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[sub]); + const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * sub)) & 127]; + + const int index_y = j * WARP_SIZE_GGUF + (QR_IQ2_XXS_MMQ * (k + l)) % WARP_SIZE_GGUF; + const int8_t* q8 = (const int8_t*)&y_qs[index_y]; + + for (int jj = 0; jj < 8; ++jj) { + sumi += q8[jj] * grid[jj] * (signs & kmask_iq2xs[jj] ? -1 : 1); + } + } + + const float2 ds8 = + __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (QR_IQ2_XXS_MMQ * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return x_dmf[i * (WARP_SIZE_GGUF / QI_IQ2_XXS_MMQ) + i / QI_IQ2_XXS_MMQ] * sub_scale * 0.25f * ds8.x * sumi; +} + +// ========================= IQ2_XS MMQ ========================= +// IQ2_XS: QK_K=256, qs[QK_K/8=32 uint16], scales[QK_K/32=8 uint8]. Grid: iq2xs_grid[512] (uint64). +// Per ib32: 4 uint16 from qs. Each uint16: low 9 bits = grid index, upper 7 = signs. + +#define QR_IQ2_XS_MMQ 2 +#define QI_IQ2_XS_MMQ (QK_K / (4 * QR_IQ2_XS_MMQ)) +#define VDR_IQ2_XS_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq2_xs(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI_IQ2_XS_MMQ) + mmq_y / QI_IQ2_XS_MMQ]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_iq2_xs( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + (void)x_qh; + + const int kbx = k / QI_IQ2_XS_MMQ; + const int kqsx = k % QI_IQ2_XS_MMQ; + + const block_iq2_xs* bx0 = (const block_iq2_xs*)vx; + float* x_dmf = (float*)x_dm; + + // Load qs data (32 uint16 = 64 bytes = 16 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq2_xs* bxi = bx0 + i * blocks_per_row + kbx; + if (kqsx < 16) { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_b2(bxi->qs, kqsx); + } else { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = 0; + } + } + + // Load d scale + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI_IQ2_XS_MMQ; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI_IQ2_XS_MMQ) { + int i = (i0 + i_offset * QI_IQ2_XS_MMQ + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq2_xs* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI_IQ2_XS_MMQ) + i / QI_IQ2_XS_MMQ + kbxd] = __half2float(bxi->d); + } + + // Load scales (8 bytes = 2 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const int ksc = k % (WARP_SIZE_GGUF / 8); + if (ksc < 2) { + const block_iq2_xs* bxi = bx0 + i * blocks_per_row + kbx; + x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = get_int_from_uint8(bxi->scales, ksc); + } + } +} + +static __device__ __forceinline__ float vec_dot_iq2_xs_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_qh; + + const float* x_dmf = (const float*)x_dm; + + const int ib32 = k / 4; + + // Read the 4 uint16 for this ib32 + const uint16_t* q2 = (const uint16_t*)&x_ql[i * (WARP_SIZE_GGUF + 1) + ib32 * 2]; + + // Per-sub-block scales + const uint8_t* sc_bytes = (const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8]; + + float sumf = 0; +#pragma unroll + for (int l = 0; l < VDR_IQ2_XS_Q8_1_MMQ; ++l) { + const int sub = l; + + // Grid lookup: 9-bit index, 7-bit signs + const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[sub] & 511)); + const uint8_t signs = ksigns_iq2xs[q2[sub] >> 9]; + + const int index_y = j * WARP_SIZE_GGUF + (QR_IQ2_XS_MMQ * (k + l)) % WARP_SIZE_GGUF; + const int8_t* q8 = (const int8_t*)&y_qs[index_y]; + + int sumi = 0; + for (int jj = 0; jj < 8; ++jj) { + sumi += q8[jj] * grid[jj] * (signs & kmask_iq2xs[jj] ? -1 : 1); + } + + const uint8_t ls = (sub < 2) ? (sc_bytes[ib32] & 0xf) : (sc_bytes[ib32] >> 4); + sumf += (0.5f + ls) * sumi; + } + + const float2 ds8 = + __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (QR_IQ2_XS_MMQ * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return x_dmf[i * (WARP_SIZE_GGUF / QI_IQ2_XS_MMQ) + i / QI_IQ2_XS_MMQ] * 0.25f * ds8.x * sumf; +} + +// ========================= IQ2_S MMQ ========================= +// IQ2_S: QK_K=256, qs[QK_K/4=64], qh[QK_K/32=8], scales[QK_K/32=8]. Grid: iq2s_grid[1024] (uint64). +// Per ib32: qs[4] + qh[1] + signs[4] + scales[1]. +// 10-bit index: 8 from qs + 2 from qh. Signs from dedicated array at qs+64. + +#define QR_IQ2_S_MMQ 2 +#define QI_IQ2_S_MMQ (QK_K / (4 * QR_IQ2_S_MMQ)) +#define VDR_IQ2_S_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq2_s(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + // slots 0..15: qs (64 bytes), slots 16..23: signs (32 bytes) + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI_IQ2_S_MMQ) + mmq_y / QI_IQ2_S_MMQ]; + __shared__ int tile_x_qh[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; + *x_qh = tile_x_qh; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_iq2_s( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI_IQ2_S_MMQ; + const int kqsx = k % QI_IQ2_S_MMQ; + + const block_iq2_s* bx0 = (const block_iq2_s*)vx; + float* x_dmf = (float*)x_dm; + + // Load qs (64 bytes, slots 0..15) and signs (32 bytes at qs+64, slots 16..23) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq2_s* bxi = bx0 + i * blocks_per_row + kbx; + if (kqsx < 16) { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + } else if (kqsx < 24) { + // MMVQ stores grid indices in qs[0..31] and signs in qs[32..63]. + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs + QK_K / 8, kqsx - 16); + } else { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = 0; + } + } + + // Load d + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI_IQ2_S_MMQ; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI_IQ2_S_MMQ) { + int i = (i0 + i_offset * QI_IQ2_S_MMQ + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq2_s* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI_IQ2_S_MMQ) + i / QI_IQ2_S_MMQ + kbxd] = __half2float(bxi->d); + } + + // Load qh (8 bytes = 2 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { + int i = (i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4)) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const int kqh = k % (WARP_SIZE_GGUF / 4); + if (kqh < 2) { + const block_iq2_s* bxi = bx0 + i * blocks_per_row + kbx; + x_qh[i * (WARP_SIZE_GGUF / 4) + i / 4 + kqh] = get_int_from_uint8(bxi->qh, kqh); + } + } + + // Load scales (8 bytes = 2 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const int ksc = k % (WARP_SIZE_GGUF / 8); + if (ksc < 2) { + const block_iq2_s* bxi = bx0 + i * blocks_per_row + kbx; + x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = get_int_from_uint8(bxi->scales, ksc); + } + } +} + +static __device__ __forceinline__ float vec_dot_iq2_s_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const float* x_dmf = (const float*)x_dm; + + const int ib32 = k / 4; + + // Grid indices from qs[0..31] (stored in tile slots 0..7 as int32) + const uint8_t* qs = (const uint8_t*)&x_ql[i * (WARP_SIZE_GGUF + 1)]; + + // qh for the 10-bit index + const uint8_t* qh_all = (const uint8_t*)&x_qh[i * (WARP_SIZE_GGUF / 4) + i / 4]; + const uint8_t qh_val = qh_all[ib32]; + + // Signs from qs[32..63] (stored in tile slots 16..23) + const uint8_t* sign_bytes = (const uint8_t*)&x_ql[i * (WARP_SIZE_GGUF + 1) + 16]; + + // Per-sub-block scales + const uint8_t* sc_bytes = (const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8]; + + float sumf = 0; +#pragma unroll + for (int l = 0; l < VDR_IQ2_S_Q8_1_MMQ; ++l) { + const int sub = l; + + // 10-bit grid index + const uint32_t* grid = (const uint32_t*)(iq2s_grid + (qs[4 * ib32 + sub] | ((qh_val << (8 - 2 * sub)) & 0x300))); + + const uint8_t sign_byte = sign_bytes[4 * ib32 + sub]; + uint32_t signs0 = __vcmpeq4(((sign_byte & 0xf) * 0x01010101) & 0x08040201, 0x08040201); + uint32_t signs1 = __vcmpeq4(((sign_byte >> 4) * 0x01010101) & 0x08040201, 0x08040201); + const int grid_l = __vsub4(grid[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid[1] ^ signs1, signs1); + + const int index_y = j * WARP_SIZE_GGUF + (QR_IQ2_S_MMQ * (k + l)) % WARP_SIZE_GGUF; + + int sumi = 0; + sumi = iq_mmq_dp4a(grid_l, y_qs[index_y + 0], sumi); + sumi = iq_mmq_dp4a(grid_h, y_qs[index_y + 1], sumi); + + const uint8_t ls = (sub < 2) ? (sc_bytes[ib32] & 0xf) : (sc_bytes[ib32] >> 4); + sumf += (0.5f + ls) * sumi; + } + + const float2 ds8 = + __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (QR_IQ2_S_MMQ * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return x_dmf[i * (WARP_SIZE_GGUF / QI_IQ2_S_MMQ) + i / QI_IQ2_S_MMQ] * 0.25f * ds8.x * sumf; +} + +// ========================= IQ1_S MMQ ========================= +// IQ1_S: QK_K=256, qs[QK_K/8=32], qh[QK_K/32=8 uint16]. Grid: iq1s_grid_gpu[2048] (uint64). +// 11-bit grid index (8 from qs + 3 from qh). need_sum=true (delta correction). +// Grid values are packed nibbles: (grid>>0)&0x0F0F0F0F, (grid>>4)&0x0F0F0F0F. + +#define QR_IQ1_S_MMQ 2 +#define QI_IQ1_S_MMQ (QK_K / (4 * QR_IQ1_S_MMQ)) +#define VDR_IQ1_S_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ void allocate_tiles_iq1_s(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + // slots 0..7: qs (32 bytes = 8 int32) + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI_IQ1_S_MMQ) + mmq_y / QI_IQ1_S_MMQ]; + __shared__ int tile_x_qh[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4]; // qh: 16 bytes = 4 int32 + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; + *x_qh = tile_x_qh; +} + +template +static __device__ __forceinline__ void load_tiles_iq1_s( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + (void)x_sc; + + const int kbx = k / QI_IQ1_S_MMQ; + const int kqsx = k % QI_IQ1_S_MMQ; + + const block_iq1_s* bx0 = (const block_iq1_s*)vx; + float* x_dmf = (float*)x_dm; + + // Load qs (32 bytes = 8 int32, slots 0..7) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_iq1_s* bxi = bx0 + i * blocks_per_row + kbx; + if (kqsx < 8) { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_b2(bxi->qs, kqsx); + } else { + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = 0; + } + } + + // Load d scale + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI_IQ1_S_MMQ; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI_IQ1_S_MMQ) { + int i = (i0 + i_offset * QI_IQ1_S_MMQ + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_iq1_s* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI_IQ1_S_MMQ) + i / QI_IQ1_S_MMQ + kbxd] = __half2float(bxi->d); + } + + // Load qh (8 uint16 = 16 bytes = 4 int32) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { + int i = (i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4)) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const int kqh = k % (WARP_SIZE_GGUF / 4); + if (kqh < 4) { + const block_iq1_s* bxi = bx0 + i * blocks_per_row + kbx; + x_qh[i * (WARP_SIZE_GGUF / 4) + i / 4 + kqh] = get_int_b2(bxi->qh, kqh); + } + } +} + +static __device__ __forceinline__ float vec_dot_iq1_s_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_sc; + + const float* x_dmf = (const float*)x_dm; + + const int iqs = k / 4; + + // Read qs (4 bytes per iqs group) + const int qs_packed = x_ql[i * (WARP_SIZE_GGUF + 1) + iqs]; + const uint8_t* qs = (const uint8_t*)&qs_packed; + + // Read qh for this iqs group + const uint16_t* qh_all = (const uint16_t*)&x_qh[i * (WARP_SIZE_GGUF / 4) + i / 4]; + const int qh = qh_all[iqs]; + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_IQ1_S_Q8_1_MMQ; ++l) { + const int sub = l; + + // 11-bit grid index: qs[sub] | (3 bits from qh << 8) + const int grid = iq1s_grid_gpu[qs[sub] | (((qh >> 3 * sub) & 0x07) << 8)]; + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + + const int index_y = j * WARP_SIZE_GGUF + (QR_IQ1_S_MMQ * (k + l)) % WARP_SIZE_GGUF; + sumi = iq_mmq_dp4a(grid0, y_qs[index_y + 0], sumi); + sumi = iq_mmq_dp4a(grid1, y_qs[index_y + 1], sumi); + } + + const float d1q = x_dmf[i * (WARP_SIZE_GGUF / QI_IQ1_S_MMQ) + i / QI_IQ1_S_MMQ] * (((qh >> 11) & 0x0E) + 1); + const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f * IQ1S_DELTA / 0x8000); + + // need_sum=true: y_ds stores half2(d, sum) + const float2 ds = + __half22float2(y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (QR_IQ1_S_MMQ * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); + + return d1q * (ds.x * sumi + ds.y * delta); +} +#endif // !defined(USE_MUSA) diff --git a/python/sglang/kernels/aot/include/sgl_kernel_ops.h b/python/sglang/kernels/aot/include/sgl_kernel_ops.h index 4ee916228de9..bf37554fc084 100644 --- a/python/sglang/kernels/aot/include/sgl_kernel_ops.h +++ b/python/sglang/kernels/aot/include/sgl_kernel_ops.h @@ -750,6 +750,8 @@ torch::Tensor ggml_moe_a8_vec( int64_t ggml_moe_get_block_size(int64_t type); +bool ggml_supports_iq_mmq(); + /* * From csrc/spatial */ diff --git a/python/sglang/kernels/aot/python/sgl_kernel/__init__.py b/python/sglang/kernels/aot/python/sgl_kernel/__init__.py index f725f62c2ed3..cd952382dead 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/__init__.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/__init__.py @@ -103,6 +103,7 @@ ggml_moe_get_block_size, ggml_mul_mat_a8, ggml_mul_mat_vec_a8, + ggml_supports_iq_mmq, ) from sgl_kernel.sampling import ( top_k_renorm_prob, diff --git a/python/sglang/kernels/aot/python/sgl_kernel/quantization/__init__.py b/python/sglang/kernels/aot/python/sgl_kernel/quantization/__init__.py index 0ad12c64e4ac..32a5070d6d98 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/quantization/__init__.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/quantization/__init__.py @@ -5,4 +5,5 @@ ggml_moe_get_block_size, ggml_mul_mat_a8, ggml_mul_mat_vec_a8, + ggml_supports_iq_mmq, ) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/quantization/gguf.py b/python/sglang/kernels/aot/python/sgl_kernel/quantization/gguf.py index 994bc21d09de..316241d71157 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/quantization/gguf.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/quantization/gguf.py @@ -60,3 +60,10 @@ def ggml_moe_a8_vec( def ggml_moe_get_block_size(type: int) -> int: return torch.ops.sgl_kernel.ggml_moe_get_block_size.default(type) + + +def ggml_supports_iq_mmq() -> bool: + try: + return bool(torch.ops.sgl_kernel.ggml_supports_iq_mmq.default()) + except (AttributeError, RuntimeError): + return False diff --git a/python/sglang/kernels/aot/tests/test_gguf.py b/python/sglang/kernels/aot/tests/test_gguf.py index 2f1f0935f5d6..1ce39008f8fb 100644 --- a/python/sglang/kernels/aot/tests/test_gguf.py +++ b/python/sglang/kernels/aot/tests/test_gguf.py @@ -7,7 +7,13 @@ import numpy as np import pytest import torch -from gguf import GGMLQuantizationType, GGUFReader, ReaderTensor, dequantize +from gguf import ( + GGML_QUANT_SIZES, + GGMLQuantizationType, + GGUFReader, + ReaderTensor, + dequantize, +) from huggingface_hub import snapshot_download from sgl_kernel import ( ggml_dequantize, @@ -16,10 +22,31 @@ ggml_moe_get_block_size, ggml_mul_mat_a8, ggml_mul_mat_vec_a8, + ggml_supports_iq_mmq, + moe_align_block_size, + silu_and_mul, ) -GGUF_SAMPLE = snapshot_download("Isotr0py/test-gguf-sample") -GGUF_SAMPLE_MOE = snapshot_download("SzymonOzog/test-gguf-moe-sample") +GGUF_SAMPLE = snapshot_download( + "Isotr0py/test-gguf-sample", + revision="d82b8773934ef260d8d8a896a7c197bc69a0fac1", + allow_patterns="*.gguf", +) +GGUF_SAMPLE_MOE = snapshot_download( + "SzymonOzog/test-gguf-moe-sample", + revision="2b77eb27ae2ea4b3f68cf3042961509a9f5847b8", + allow_patterns=[ + "Quant_IQ1_S_512.gguf", + "Quant_IQ2_XXS_512.gguf", + "Quant_IQ2_XS_512.gguf", + "Quant_IQ2_S_512.gguf", + "Quant_IQ3_XXS_512.gguf", + "Quant_IQ3_S_512.gguf", + "Quant_IQ4_NL_512.gguf", + "Quant_IQ4_XS_512.gguf", + "Quant_Q5_0_512.gguf", + ], +) def get_gguf_sample_tensors( @@ -50,6 +77,7 @@ def get_gguf_MoE_tensors( # i-matrix GGMLQuantizationType.IQ1_M, GGMLQuantizationType.IQ1_S, + GGMLQuantizationType.IQ2_XXS, GGMLQuantizationType.IQ2_S, GGMLQuantizationType.IQ2_XS, GGMLQuantizationType.IQ3_S, @@ -67,6 +95,63 @@ def get_gguf_MoE_tensors( GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0, ] +MMQ_IMATRIX_QUANT_TYPES = [ + GGMLQuantizationType.IQ1_S, + GGMLQuantizationType.IQ2_XXS, + GGMLQuantizationType.IQ2_XS, + GGMLQuantizationType.IQ2_S, + GGMLQuantizationType.IQ3_XXS, + GGMLQuantizationType.IQ3_S, + GGMLQuantizationType.IQ4_NL, + GGMLQuantizationType.IQ4_XS, +] +MMQ_MOE_TEST_TYPES = MMQ_IMATRIX_QUANT_TYPES + [GGMLQuantizationType.Q5_0] +MMQ_MOE_TEST_TYPE_PAIRS = [ + (quant_type, quant_type) for quant_type in MMQ_MOE_TEST_TYPES +] + [ + (GGMLQuantizationType.IQ3_S, GGMLQuantizationType.Q5_0), + (GGMLQuantizationType.Q5_0, GGMLQuantizationType.IQ3_S), +] +MMQ_K_ALIGNMENTS = { + GGMLQuantizationType.Q4_0: 256, + GGMLQuantizationType.Q5_0: 256, + GGMLQuantizationType.Q8_0: 128, + GGMLQuantizationType.Q2_K: 512, + GGMLQuantizationType.Q3_K: 512, + GGMLQuantizationType.Q4_K: 256, + GGMLQuantizationType.Q5_K: 256, + GGMLQuantizationType.Q6_K: 256, + **{quant_type: 256 for quant_type in MMQ_IMATRIX_QUANT_TYPES}, +} + + +def _moe_align( + topk_ids: torch.Tensor, block_size: int, num_experts: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1) + sorted_token_ids = torch.empty( + max_num_tokens_padded, dtype=torch.int32, device=topk_ids.device + ) + expert_ids = torch.empty( + (max_num_tokens_padded + block_size - 1) // block_size, + dtype=torch.int32, + device=topk_ids.device, + ) + num_tokens_post_padded = torch.empty(1, dtype=torch.int32, device=topk_ids.device) + cumsum_buffer = torch.empty( + num_experts + 2, dtype=torch.int32, device=topk_ids.device + ) + moe_align_block_size( + topk_ids, + num_experts + 1, + block_size, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + cumsum_buffer, + True, + ) + return sorted_token_ids, expert_ids, num_tokens_post_padded @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) @@ -122,7 +207,8 @@ def test_mmvq(hidden_size: int, dtype: torch.dtype, quant_type: GGMLQuantization @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize( "quant_type", - [ + MMQ_IMATRIX_QUANT_TYPES + + [ # k-quants GGMLQuantizationType.Q2_K, GGMLQuantizationType.Q3_K, @@ -142,7 +228,7 @@ def test_mmq( dtype: torch.dtype, quant_type: GGMLQuantizationType, ): - + torch.manual_seed(0) tensors = get_gguf_sample_tensors(hidden_size, quant_type) x = torch.rand((num_tokens, hidden_size), dtype=dtype, device="cuda") for tensor in tensors: @@ -152,7 +238,16 @@ def test_mmq( ref_output = x @ weight.T qweight = torch.tensor(tensor.data, device="cuda") + alignment = MMQ_K_ALIGNMENTS[quant_type] + if hidden_size % alignment != 0: + with pytest.raises(RuntimeError, match="requires an input size"): + ggml_mul_mat_a8(qweight, x, quant_type, qweight.shape[0]) + continue output = ggml_mul_mat_a8(qweight, x, quant_type, qweight.shape[0]) + if quant_type in MMQ_IMATRIX_QUANT_TYPES: + assert torch.isfinite(output).all() + torch.testing.assert_close(output, ref_output, atol=1.5, rtol=1e-1) + continue atols = {torch.half: 1, torch.bfloat16: 1.5, torch.float: 1.2} # test matrix has inputs centered around 0 and lower precision from # bfloat16 tends to accumulate and can greatly inflate rtol @@ -163,5 +258,195 @@ def test_mmq( ) +def _reference_moe( + x: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, +) -> torch.Tensor: + output = torch.empty((x.shape[0], w2.shape[1]), dtype=x.dtype, device=x.device) + for token_idx in range(x.shape[0]): + expert_indices = topk_ids[token_idx].long() + gate_up = torch.einsum("h,enh->en", x[token_idx], w13[expert_indices]) + activated = silu_and_mul(gate_up) + expert_output = torch.einsum("ei,eoi->eo", activated, w2[expert_indices]) + output[token_idx] = torch.sum( + expert_output * topk_weights[token_idx, :, None], dim=0 + ) + return output + + +def test_moe_block_size_dispatch(): + assert ggml_moe_get_block_size(GGMLQuantizationType.Q5_0) > 0 + + +def test_iq_mmq_capability_dispatch(): + assert ggml_supports_iq_mmq() + + +def test_q8_1_is_rejected_by_unsupported_low_level_paths(): + qweight = torch.empty((1, 32), dtype=torch.uint8, device="cuda") + x = torch.empty((1, 256), dtype=torch.bfloat16, device="cuda") + + with pytest.raises(RuntimeError, match="Unsupported GGUF dequantization type"): + ggml_dequantize(qweight, GGMLQuantizationType.Q8_1, 1, 256, torch.bfloat16) + with pytest.raises(RuntimeError, match="Unsupported GGUF MMQ quantization type"): + ggml_mul_mat_a8(qweight, x, GGMLQuantizationType.Q8_1, 1) + with pytest.raises(RuntimeError, match="Unsupported GGUF MMVQ quantization type"): + ggml_mul_mat_vec_a8(qweight, x, GGMLQuantizationType.Q8_1, 1) + + topk_ids = torch.zeros((1, 1), dtype=torch.int32, device="cuda") + with pytest.raises( + RuntimeError, match="Unsupported GGUF MoE vector quantization type" + ): + ggml_moe_a8_vec( + x, qweight.unsqueeze(0), topk_ids, 1, GGMLQuantizationType.Q8_1, 1, 1 + ) + + +@pytest.mark.parametrize( + "quant_type,hidden_size", + [ + (GGMLQuantizationType.Q5_0, 32), + (GGMLQuantizationType.Q8_0, 96), + (GGMLQuantizationType.Q2_K, 256), + (GGMLQuantizationType.Q3_K, 256), + (GGMLQuantizationType.IQ4_NL, 32), + (GGMLQuantizationType.IQ4_NL, 96), + ], +) +def test_mmq_rejects_unaligned_k(quant_type: GGMLQuantizationType, hidden_size: int): + torch.manual_seed(0) + tensor = get_gguf_sample_tensors(256, quant_type)[0] + block_size, type_size = GGML_QUANT_SIZES[quant_type] + packed_cols = hidden_size // block_size * type_size + qweight = torch.tensor(tensor.data[:, :packed_cols], device="cuda") + x = torch.rand((9, hidden_size), dtype=torch.bfloat16, device="cuda") + + error = "requires an input size divisible by" + with pytest.raises(RuntimeError, match=error): + ggml_mul_mat_a8(qweight, x, quant_type, qweight.shape[0]) + + sorted_token_ids = torch.zeros(1, dtype=torch.int32, device="cuda") + expert_ids = torch.zeros(1, dtype=torch.int32, device="cuda") + num_tokens_post_padded = torch.ones(1, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match=error): + ggml_moe_a8( + x[:1], + qweight.unsqueeze(0), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type, + qweight.shape[0], + 1, + 1, + ) + + +def test_moe_mmq_invalid_expert_writes_zero(): + torch.manual_seed(0) + quant_type = GGMLQuantizationType.IQ3_S + tensor = get_gguf_MoE_tensors(512, quant_type)[0] + qweight = torch.tensor(tensor.data, device="cuda") + weight = dequantize(tensor.data, quant_type) + x = torch.rand((1, weight.shape[-1]), dtype=torch.bfloat16, device="cuda") + sorted_token_ids = torch.tensor([0, 1, 1, 1], dtype=torch.int32, device="cuda") + expert_ids = torch.tensor([-1], dtype=torch.int32, device="cuda") + num_tokens_post_padded = torch.tensor([4], dtype=torch.int32, device="cuda") + + output = ggml_moe_a8( + x, + qweight, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type, + weight.shape[-2], + 1, + 1, + ) + torch.testing.assert_close(output, torch.zeros_like(output), atol=0, rtol=0) + + vector_output = ggml_moe_a8_vec( + x, + qweight, + torch.tensor([[-1]], dtype=torch.int32, device="cuda"), + 1, + quant_type, + weight.shape[-2], + 1, + ) + torch.testing.assert_close( + vector_output, torch.zeros_like(vector_output), atol=0, rtol=0 + ) + + +@pytest.mark.parametrize("quant_type,quant_type2", MMQ_MOE_TEST_TYPE_PAIRS) +@torch.inference_mode() +def test_moe_mmq(quant_type: GGMLQuantizationType, quant_type2: GGMLQuantizationType): + torch.manual_seed(0) + num_tokens = 65 if quant_type == quant_type2 == GGMLQuantizationType.Q5_0 else 128 + top_k = 4 + w13_tensor = get_gguf_MoE_tensors(512, quant_type)[0] + w2_tensor = get_gguf_MoE_tensors(512, quant_type2)[1] + w13 = torch.tensor(w13_tensor.data, device="cuda") + w2 = torch.tensor(w2_tensor.data, device="cuda") + w13_dequant = torch.tensor( + dequantize(w13_tensor.data, quant_type), device="cuda" + ).to(torch.bfloat16) + w2_dequant = torch.tensor( + dequantize(w2_tensor.data, quant_type2), device="cuda" + ).to(torch.bfloat16) + + x = torch.rand( + (num_tokens, w13_dequant.shape[-1]), + dtype=torch.bfloat16, + device="cuda", + ) + topk_weights = torch.rand((num_tokens, top_k), dtype=torch.bfloat16, device="cuda") + topk_ids = torch.randint( + 0, + w13.shape[0], + (num_tokens, top_k), + dtype=torch.int32, + device="cuda", + ) + + block_size = ggml_moe_get_block_size(quant_type) + sorted_token_ids, expert_ids, num_tokens_post_padded = _moe_align( + topk_ids, block_size, w13.shape[0] + ) + output = ggml_moe_a8( + x, + w13, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type, + w13.shape[1], + top_k, + num_tokens, + ) + output = silu_and_mul(output) + output = ggml_moe_a8( + output, + w2, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type2, + w2.shape[1], + 1, + num_tokens * top_k, + ) + output = output.reshape(num_tokens, top_k, w2.shape[1]) + output = torch.sum(output * topk_weights[:, :, None], dim=1) + + ref_output = _reference_moe(x, w13_dequant, w2_dequant, topk_weights, topk_ids) + torch.testing.assert_close(output, ref_output, atol=1, rtol=1e-1) + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index 98784000d017..aa78a21f6956 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -39,9 +39,10 @@ _is_xpu = is_xpu() _is_musa = is_musa() _is_npu = is_npu() +_supports_iq_mmq = False if _is_cuda: - from sgl_kernel import moe_align_block_size, moe_sum + from sgl_kernel import moe_sum from sgl_kernel.quantization import ( ggml_dequantize, ggml_moe_a8, @@ -51,7 +52,17 @@ ggml_mul_mat_vec_a8, ) + try: + from sgl_kernel.quantization import ggml_supports_iq_mmq + except ImportError: + pass + else: + _supports_iq_mmq = ggml_supports_iq_mmq() + from sglang.kernels.ops.activation.activation import gelu_and_mul, silu_and_mul + from sglang.srt.layers.moe.moe_runner.triton_utils.moe_align_block_size import ( + moe_align_block_size, + ) elif _is_musa: from sgl_kernel import gelu_and_mul, moe_align_block_size, moe_sum, silu_and_mul from sgl_kernel.quantization import ( @@ -141,7 +152,6 @@ def is_layer_skipped_gguf(prefix: str, modules_to_not_convert: list[str]): WeightType.Q5_0, WeightType.Q5_1, WeightType.Q8_0, - WeightType.Q8_1, } KQUANT_TYPES = { WeightType.Q2_K, @@ -161,12 +171,130 @@ def is_layer_skipped_gguf(prefix: str, modules_to_not_convert: list[str]): WeightType.IQ4_XS, WeightType.IQ4_NL, } -# TODO(Isotr0py): Currently, we don't have MMQ kernel for I-Matrix quantization. -# Consolidate DEQUANT_TYPES, MMVQ_QUANT_TYPES and MMQ_QUANT_TYPES after we add -# MMQ kernel for I-Matrix quantization. +MMQ_IMATRIX_QUANT_TYPES = { + WeightType.IQ1_S, + WeightType.IQ2_XXS, + WeightType.IQ2_XS, + WeightType.IQ2_S, + WeightType.IQ3_XXS, + WeightType.IQ3_S, + WeightType.IQ4_XS, + WeightType.IQ4_NL, +} DEQUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES MMVQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES -MMQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES + + +def _build_mmq_quant_types(supports_iq_mmq: bool) -> set[WeightType]: + quant_types = STANDARD_QUANT_TYPES | KQUANT_TYPES + if supports_iq_mmq: + quant_types |= MMQ_IMATRIX_QUANT_TYPES + return quant_types + + +MMQ_QUANT_TYPES = _build_mmq_quant_types(_is_cuda and _supports_iq_mmq) + +IQ_MMQ_MAX_BATCH_SIZE = 16 +IQ_MMQ_NARROW_MAX_BATCH_SIZE = 16 +IQ_MMQ_NARROW_MAX_OUTPUT_ROWS = 5120 +IQ_MOE_MMQ_MIN_TOKENS = 128 +IQ_MOE_MMQ_MIN_ASSIGNMENTS_PER_EXPERT = 2 +IQ_MOE_MMQ_MAX_EXPERTS = 256 +IQ_MOE_MMQ_BLOCK_SIZE = 4 +CUDA_MAX_GRID_Y = 65535 +MMQ_K_ALIGNMENTS = { + WeightType.Q4_0: 256, + WeightType.Q4_1: 256, + WeightType.Q5_0: 256, + WeightType.Q5_1: 256, + WeightType.Q8_0: 128, + WeightType.Q2_K: 512, + WeightType.Q3_K: 512, + WeightType.Q4_K: 256, + WeightType.Q5_K: 256, + WeightType.Q6_K: 256, + **{qweight_type: 256 for qweight_type in MMQ_IMATRIX_QUANT_TYPES}, +} + + +def _is_iq_mmq_batch_size(num_tokens: int, num_output_rows: int) -> bool: + max_batch_size = ( + IQ_MMQ_NARROW_MAX_BATCH_SIZE + if num_output_rows <= IQ_MMQ_NARROW_MAX_OUTPUT_ROWS + else IQ_MMQ_MAX_BATCH_SIZE + ) + return num_tokens <= max_batch_size + + +def _is_iq_moe_mmq_shape(num_tokens: int, num_experts: int, top_k: int) -> bool: + # moe_align_block_size allocates this conservative capacity, and the IQ + # kernels launch one four-assignment block along grid.y. + max_num_tokens_padded = num_tokens * top_k + (num_experts + 1) * ( + IQ_MOE_MMQ_BLOCK_SIZE - 1 + ) + return ( + num_tokens >= IQ_MOE_MMQ_MIN_TOKENS + and num_experts <= IQ_MOE_MMQ_MAX_EXPERTS + and num_tokens * top_k >= IQ_MOE_MMQ_MIN_ASSIGNMENTS_PER_EXPERT * num_experts + and max_num_tokens_padded <= CUDA_MAX_GRID_Y * IQ_MOE_MMQ_BLOCK_SIZE + ) + + +def _is_mmq_input_size(qweight_type: int, input_size: int) -> bool: + alignment = MMQ_K_ALIGNMENTS.get(qweight_type) + return alignment is not None and input_size % alignment == 0 + + +def _select_gguf_matmul_kernel( + num_tokens: int, num_output_rows: int, qweight_type: int, input_size: int +) -> str: + if qweight_type in UNQUANTIZED_TYPES: + return "unquantized" + + if qweight_type in IMATRIX_QUANT_TYPES: + mmvq_safe = 8 if num_output_rows > 5120 else 16 + else: + mmvq_safe = 2 if num_output_rows > 5120 else 6 + + if num_tokens <= mmvq_safe and qweight_type in MMVQ_QUANT_TYPES: + return "mmvq" + if ( + qweight_type in MMQ_QUANT_TYPES + and ( + qweight_type not in MMQ_IMATRIX_QUANT_TYPES + or _is_iq_mmq_batch_size(num_tokens, num_output_rows) + ) + and _is_mmq_input_size(qweight_type, input_size) + ): + return "mmq" + if qweight_type in DEQUANT_TYPES: + return "dequantize" + return "unsupported" + + +def _should_use_moe_mmq( + qweight_type: int, + qweight_type2: int, + num_tokens: int, + num_experts: int, + top_k: int, + input_size: int, + intermediate_size: int, +) -> bool: + if qweight_type not in MMQ_QUANT_TYPES or qweight_type2 not in MMQ_QUANT_TYPES: + return False + if not _is_mmq_input_size(qweight_type, input_size) or not _is_mmq_input_size( + qweight_type2, intermediate_size + ): + return False + + uses_iq_mmq = ( + qweight_type in MMQ_IMATRIX_QUANT_TYPES + or qweight_type2 in MMQ_IMATRIX_QUANT_TYPES + ) + if not uses_iq_mmq: + return num_tokens > 64 + return _is_iq_moe_mmq_shape(num_tokens, num_experts, top_k) def dequantize_gguf_weight( @@ -181,25 +309,21 @@ def dequantize_gguf_weight( def fused_mul_mat_gguf( x: torch.Tensor, qweight: torch.Tensor, qweight_type: int ) -> torch.Tensor: - if qweight_type in IMATRIX_QUANT_TYPES: - mmvq_safe = 8 if qweight.shape[0] > 5120 else 16 - else: - mmvq_safe = 2 if qweight.shape[0] > 5120 else 6 # HACK: when doing chunked prefill we don't generate output tokens # so input to logits generator is empty which causes invalid parameter if x.shape[0] == 0: return torch.empty(x.shape[0], qweight.shape[0], dtype=x.dtype, device=x.device) - # there is no need to call any kernel for fp16/bf16 - if qweight_type in UNQUANTIZED_TYPES: + + kernel = _select_gguf_matmul_kernel( + x.shape[0], qweight.shape[0], qweight_type, x.shape[1] + ) + if kernel == "unquantized": return x @ qweight.T - # enable MMVQ in contiguous batching with batch_size=1 - if x.shape[0] <= mmvq_safe and qweight_type in MMVQ_QUANT_TYPES: + if kernel == "mmvq": y = ggml_mul_mat_vec_a8(qweight, x, qweight_type, qweight.shape[0]) - # Use MMQ Kernel if it's available (standard + k-quants) - elif qweight_type in MMQ_QUANT_TYPES: + elif kernel == "mmq": y = ggml_mul_mat_a8(qweight, x, qweight_type, qweight.shape[0]) - # If there is no available MMQ kernel, fallback to dequantize - elif qweight_type in DEQUANT_TYPES: + elif kernel == "dequantize": weight = dequantize_gguf_weight(qweight, qweight_type, x.dtype) y = x @ weight.T else: @@ -229,19 +353,23 @@ def act(x: torch.Tensor): raise ValueError(f"Unsupported activation: {activation}") out_hidden_states = torch.empty_like(x) - # unless we decent expert reuse we are better off running moe_vec kernel - if ( - qweight_type2 in MMQ_QUANT_TYPES - and qweight_type in MMQ_QUANT_TYPES - and x.shape[0] > 64 + num_tokens, _ = x.shape + num_experts, N, _ = w1.shape + top_k = topk_ids.shape[1] + # The batched kernel needs enough work per expert to amortize alignment. + if _should_use_moe_mmq( + qweight_type, + qweight_type2, + num_tokens, + num_experts, + top_k, + x.shape[1], + N // 2, ): - num_tokens, _ = x.shape - E, N, _ = w1.shape - top_k = topk_ids.shape[1] BLOCK_SIZE = ggml_moe_get_block_size(qweight_type) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( - topk_ids, BLOCK_SIZE, E + topk_ids, BLOCK_SIZE, num_experts ) out = ggml_moe_a8( x, @@ -272,10 +400,6 @@ def act(x: torch.Tensor): # TODO(FlamingoPg): maybe we can use moe_sum_reduce here? moe_sum(out, out_hidden_states) elif qweight_type2 in MMVQ_QUANT_TYPES and qweight_type in MMVQ_QUANT_TYPES: - num_tokens, _ = x.shape - E, N, _ = w1.shape - top_k = topk_ids.shape[1] - out = ggml_moe_a8_vec(x, w1, topk_ids, top_k, qweight_type, N, num_tokens) out = act(out) diff --git a/test/registered/unit/layers/quantization/test_gguf_mmq_types.py b/test/registered/unit/layers/quantization/test_gguf_mmq_types.py new file mode 100644 index 000000000000..4ac177bd82dd --- /dev/null +++ b/test/registered/unit/layers/quantization/test_gguf_mmq_types.py @@ -0,0 +1,447 @@ +"""Registration guardrails for GGUF I-quant MMQ kernels.""" + +import unittest +from unittest.mock import patch + +import torch +from gguf import GGMLQuantizationType as WeightType + +from sglang.srt.layers.quantization import gguf as gguf_quant +from sglang.srt.layers.quantization.gguf import ( + DEQUANT_TYPES, + KQUANT_TYPES, + MMQ_IMATRIX_QUANT_TYPES, + MMQ_K_ALIGNMENTS, + MMQ_QUANT_TYPES, + MMVQ_QUANT_TYPES, + STANDARD_QUANT_TYPES, + _build_mmq_quant_types, + _is_iq_mmq_batch_size, + _is_iq_moe_mmq_shape, + _select_gguf_matmul_kernel, + _should_use_moe_mmq, +) +from sglang.srt.utils import is_cuda +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +EXPECTED_MMQ_IMATRIX_TYPES = { + WeightType.IQ1_S, + WeightType.IQ2_XXS, + WeightType.IQ2_XS, + WeightType.IQ2_S, + WeightType.IQ3_XXS, + WeightType.IQ3_S, + WeightType.IQ4_NL, + WeightType.IQ4_XS, +} +CUDA_MMQ_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES | EXPECTED_MMQ_IMATRIX_TYPES + + +class TestGGUFMMQTypes(CustomTestCase): + def test_only_upstream_supported_iq_types_are_registered(self): + self.assertEqual(MMQ_IMATRIX_QUANT_TYPES, EXPECTED_MMQ_IMATRIX_TYPES) + self.assertNotIn(WeightType.IQ1_M, MMQ_IMATRIX_QUANT_TYPES) + self.assertNotIn(WeightType.Q8_1, STANDARD_QUANT_TYPES) + self.assertNotIn(WeightType.Q8_1, DEQUANT_TYPES) + self.assertNotIn(WeightType.Q8_1, MMVQ_QUANT_TYPES) + self.assertNotIn(WeightType.Q8_1, MMQ_QUANT_TYPES) + + def test_iq_mmq_requires_new_kernel_capability(self): + without_capability = _build_mmq_quant_types(False) + with_capability = _build_mmq_quant_types(True) + self.assertTrue(EXPECTED_MMQ_IMATRIX_TYPES.isdisjoint(without_capability)) + self.assertTrue(EXPECTED_MMQ_IMATRIX_TYPES <= with_capability) + + if is_cuda() and gguf_quant._supports_iq_mmq: + self.assertTrue(EXPECTED_MMQ_IMATRIX_TYPES <= MMQ_QUANT_TYPES) + else: + self.assertTrue(EXPECTED_MMQ_IMATRIX_TYPES.isdisjoint(MMQ_QUANT_TYPES)) + + def test_old_kernel_capability_keeps_iq_fallbacks(self): + without_capability = _build_mmq_quant_types(False) + with patch.object(gguf_quant, "MMQ_QUANT_TYPES", without_capability): + for qweight_type in EXPECTED_MMQ_IMATRIX_TYPES: + self.assertEqual( + _select_gguf_matmul_kernel(9, 18944, qweight_type, 1024), + "dequantize", + ) + self.assertEqual( + _select_gguf_matmul_kernel(32, 18944, qweight_type, 1024), + "dequantize", + ) + self.assertFalse( + _should_use_moe_mmq( + qweight_type, + WeightType.Q5_0, + 128, + 256, + 4, + 1024, + 512, + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.Q5_0, + qweight_type, + 128, + 256, + 4, + 1024, + 512, + ) + ) + + self.assertEqual( + _select_gguf_matmul_kernel(7, 18944, WeightType.Q5_0, 1024), + "mmq", + ) + self.assertTrue( + _should_use_moe_mmq( + WeightType.Q5_0, + WeightType.Q2_K, + 65, + 8, + 2, + 1024, + 512, + ) + ) + + def test_dense_kernel_selection(self): + with patch.object(gguf_quant, "MMQ_QUANT_TYPES", CUDA_MMQ_TYPES): + # Wide I-quant matrices use MMVQ through 8, MMQ through 16. + self.assertEqual( + _select_gguf_matmul_kernel(8, 5121, WeightType.IQ3_S, 1024), + "mmvq", + ) + self.assertEqual( + _select_gguf_matmul_kernel(9, 5121, WeightType.IQ3_S, 1024), + "mmq", + ) + self.assertEqual( + _select_gguf_matmul_kernel(16, 5121, WeightType.IQ3_S, 1024), + "mmq", + ) + self.assertEqual( + _select_gguf_matmul_kernel(17, 5121, WeightType.IQ3_S, 1024), + "dequantize", + ) + + # Narrow I-quant matrices stay on MMVQ through 16, then dequantize. + self.assertEqual( + _select_gguf_matmul_kernel(16, 5120, WeightType.IQ3_S, 1024), + "mmvq", + ) + self.assertEqual( + _select_gguf_matmul_kernel(17, 5120, WeightType.IQ3_S, 1024), + "dequantize", + ) + self.assertEqual( + _select_gguf_matmul_kernel(32, 5120, WeightType.IQ3_S, 1024), + "dequantize", + ) + self.assertEqual( + _select_gguf_matmul_kernel(33, 5120, WeightType.IQ3_S, 1024), + "dequantize", + ) + + # Existing standard and K-quant routing remains unchanged for + # aligned model dimensions. + self.assertEqual( + _select_gguf_matmul_kernel(2, 5121, WeightType.Q5_0, 1024), + "mmvq", + ) + self.assertEqual( + _select_gguf_matmul_kernel(3, 5121, WeightType.Q5_0, 1024), + "mmq", + ) + self.assertEqual( + _select_gguf_matmul_kernel(64, 5121, WeightType.Q2_K, 1024), + "mmq", + ) + + # Every MMQ loader consumes a complete K tile. + for qweight_type, unaligned, aligned in ( + (WeightType.Q5_0, 32, 256), + (WeightType.Q8_0, 96, 128), + (WeightType.Q2_K, 256, 512), + (WeightType.Q3_K, 256, 512), + (WeightType.IQ4_NL, 96, 256), + ): + num_tokens = 9 if qweight_type in MMQ_IMATRIX_QUANT_TYPES else 7 + self.assertEqual( + _select_gguf_matmul_kernel( + num_tokens, 5121, qweight_type, unaligned + ), + "dequantize", + ) + self.assertEqual( + _select_gguf_matmul_kernel(num_tokens, 5121, qweight_type, aligned), + "mmq", + ) + + # IQ1_M intentionally keeps the dequantization fallback. + self.assertEqual( + _select_gguf_matmul_kernel(8, 5121, WeightType.IQ1_M, 1024), + "mmvq", + ) + self.assertEqual( + _select_gguf_matmul_kernel(9, 5121, WeightType.IQ1_M, 1024), + "dequantize", + ) + + self.assertEqual( + _select_gguf_matmul_kernel(1, 5121, WeightType.Q8_1, 32), + "unsupported", + ) + + self.assertEqual(MMQ_K_ALIGNMENTS[WeightType.Q8_0], 128) + self.assertEqual(MMQ_K_ALIGNMENTS[WeightType.Q2_K], 512) + self.assertEqual(MMQ_K_ALIGNMENTS[WeightType.IQ4_NL], 256) + + self.assertTrue(_is_iq_mmq_batch_size(16, 5121)) + self.assertFalse(_is_iq_mmq_batch_size(17, 5121)) + self.assertTrue(_is_iq_mmq_batch_size(16, 5120)) + self.assertFalse(_is_iq_mmq_batch_size(17, 5120)) + + def test_dense_selector_terminal_paths(self): + self.assertEqual( + _select_gguf_matmul_kernel(1, 1024, WeightType.F16, 1024), + "unquantized", + ) + self.assertEqual(_select_gguf_matmul_kernel(1, 1024, -1, 1024), "unsupported") + + def test_empty_dense_input_bypasses_selector(self): + x = torch.empty((0, 4), dtype=torch.bfloat16) + qweight = torch.empty((3, 4), dtype=torch.uint8) + + with patch.object(gguf_quant, "_select_gguf_matmul_kernel") as selector: + output = gguf_quant.fused_mul_mat_gguf(x, qweight, -1) + + selector.assert_not_called() + self.assertEqual(output.shape, (0, 3)) + self.assertEqual(output.dtype, x.dtype) + + def test_moe_iq_mmq_requires_reuse(self): + self.assertFalse(_is_iq_moe_mmq_shape(127, 8, 8)) + self.assertFalse(_is_iq_moe_mmq_shape(128, 256, 2)) + self.assertTrue(_is_iq_moe_mmq_shape(128, 256, 4)) + self.assertTrue(_is_iq_moe_mmq_shape(256, 256, 2)) + self.assertFalse(_is_iq_moe_mmq_shape(256, 257, 4)) + + max_tokens = ( + gguf_quant.CUDA_MAX_GRID_Y * gguf_quant.IQ_MOE_MMQ_BLOCK_SIZE + - (256 + 1) * (gguf_quant.IQ_MOE_MMQ_BLOCK_SIZE - 1) + ) // 4 + self.assertTrue(_is_iq_moe_mmq_shape(max_tokens, 256, 4)) + self.assertFalse(_is_iq_moe_mmq_shape(max_tokens + 1, 256, 4)) + + def test_moe_kernel_selection(self): + with patch.object(gguf_quant, "MMQ_QUANT_TYPES", CUDA_MMQ_TYPES): + # Existing standard/K kernels retain their 64/65 boundary. + self.assertFalse( + _should_use_moe_mmq( + WeightType.Q5_0, WeightType.Q2_K, 64, 8, 2, 1024, 512 + ) + ) + self.assertTrue( + _should_use_moe_mmq( + WeightType.Q5_0, WeightType.Q2_K, 65, 8, 2, 1024, 512 + ) + ) + + # Any I-quant weight selects the conservative I-quant gate. + self.assertFalse( + _should_use_moe_mmq( + WeightType.Q5_0, WeightType.IQ3_S, 127, 8, 8, 1024, 512 + ) + ) + self.assertTrue( + _should_use_moe_mmq( + WeightType.IQ3_S, WeightType.Q5_0, 128, 256, 4, 1024, 512 + ) + ) + self.assertTrue( + _should_use_moe_mmq( + WeightType.Q5_0, WeightType.IQ3_S, 128, 256, 4, 1024, 512 + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.IQ3_S, WeightType.Q5_0, 256, 257, 4, 1024, 512 + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.IQ1_M, WeightType.Q5_0, 256, 8, 8, 1024, 512 + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.Q5_0, WeightType.IQ1_M, 256, 8, 8, 1024, 512 + ) + ) + + self.assertFalse( + _should_use_moe_mmq( + WeightType.Q8_1, WeightType.IQ3_S, 128, 256, 4, 1024, 512 + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.IQ3_S, WeightType.Q8_1, 128, 256, 4, 1024, 512 + ) + ) + + # Both projections must satisfy their type-specific K alignment. + self.assertFalse( + _should_use_moe_mmq( + WeightType.IQ4_NL, WeightType.Q5_0, 128, 256, 4, 96, 512 + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.Q5_0, WeightType.IQ4_NL, 128, 256, 4, 1024, 96 + ) + ) + self.assertTrue( + _should_use_moe_mmq( + WeightType.IQ4_NL, + WeightType.IQ4_NL, + 128, + 256, + 4, + 256, + 512, + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.Q2_K, WeightType.IQ3_S, 128, 256, 4, 256, 512 + ) + ) + self.assertFalse( + _should_use_moe_mmq( + WeightType.IQ3_S, WeightType.Q8_0, 128, 256, 4, 256, 96 + ) + ) + + def test_fused_moe_uses_high_level_align_wrapper(self): + num_tokens, hidden_size, num_experts, top_k = 128, 256, 2, 2 + x = torch.ones((num_tokens, hidden_size), dtype=torch.bfloat16) + w1 = torch.empty((num_experts, 512, 16), dtype=torch.uint8) + w2 = torch.empty((num_experts, hidden_size, 16), dtype=torch.uint8) + topk_ids = torch.zeros((num_tokens, top_k), dtype=torch.int32) + topk_weights = torch.full((num_tokens, top_k), 0.5, dtype=x.dtype) + sorted_token_ids = torch.arange(num_tokens * top_k, dtype=torch.int32) + expert_ids = torch.zeros((num_tokens * top_k // 4,), dtype=torch.int32) + num_tokens_post_padded = torch.tensor([num_tokens * top_k], dtype=torch.int32) + first_projection = torch.ones((num_tokens * top_k, 512), dtype=x.dtype) + activated = torch.ones((num_tokens * top_k, 256), dtype=x.dtype) + second_projection = torch.ones((num_tokens * top_k, hidden_size), dtype=x.dtype) + + def reduce_experts(inp, output): + output.copy_(inp.sum(dim=1)) + + with ( + patch.object(gguf_quant, "MMQ_QUANT_TYPES", CUDA_MMQ_TYPES), + patch.object( + gguf_quant, "ggml_moe_get_block_size", return_value=4, create=True + ), + patch.object( + gguf_quant, + "moe_align_block_size", + return_value=( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + ), + create=True, + ) as align, + patch.object( + gguf_quant, + "ggml_moe_a8", + side_effect=(first_projection, second_projection), + create=True, + ) as moe_a8, + patch.object( + gguf_quant, "silu_and_mul", return_value=activated, create=True + ), + patch.object( + gguf_quant, "moe_sum", side_effect=reduce_experts, create=True + ), + ): + output = gguf_quant.fused_moe_gguf( + x, + w1, + w2, + topk_weights, + topk_ids, + WeightType.IQ3_S, + WeightType.Q5_0, + "silu", + ) + + align_args = align.call_args.args + self.assertIs(align_args[0], topk_ids) + self.assertEqual(align_args[1:], (4, num_experts)) + self.assertEqual(moe_a8.call_count, 2) + self.assertEqual(moe_a8.call_args_list[0].args[5], WeightType.IQ3_S) + self.assertEqual(moe_a8.call_args_list[1].args[5], WeightType.Q5_0) + torch.testing.assert_close(output, torch.ones_like(output)) + + def test_fused_moe_uses_vector_fallback_below_iq_threshold(self): + num_tokens, hidden_size, num_experts, top_k = 127, 256, 2, 2 + x = torch.ones((num_tokens, hidden_size), dtype=torch.bfloat16) + w1 = torch.empty((num_experts, 512, 16), dtype=torch.uint8) + w2 = torch.empty((num_experts, hidden_size, 16), dtype=torch.uint8) + topk_ids = torch.zeros((num_tokens, top_k), dtype=torch.int32) + topk_weights = torch.full((num_tokens, top_k), 0.5, dtype=x.dtype) + first_projection = torch.ones((num_tokens * top_k, 512), dtype=x.dtype) + activated = torch.ones((num_tokens * top_k, 256), dtype=x.dtype) + second_projection = torch.ones((num_tokens * top_k, hidden_size), dtype=x.dtype) + + def reduce_experts(inp, output): + output.copy_(inp.sum(dim=1)) + + with ( + patch.object(gguf_quant, "MMQ_QUANT_TYPES", CUDA_MMQ_TYPES), + patch.object(gguf_quant, "moe_align_block_size", create=True) as align, + patch.object(gguf_quant, "ggml_moe_a8", create=True) as moe_a8, + patch.object( + gguf_quant, + "ggml_moe_a8_vec", + side_effect=(first_projection, second_projection), + create=True, + ) as moe_a8_vec, + patch.object( + gguf_quant, "silu_and_mul", return_value=activated, create=True + ), + patch.object( + gguf_quant, "moe_sum", side_effect=reduce_experts, create=True + ), + ): + output = gguf_quant.fused_moe_gguf( + x, + w1, + w2, + topk_weights, + topk_ids, + WeightType.IQ3_S, + WeightType.Q5_0, + "silu", + ) + + align.assert_not_called() + moe_a8.assert_not_called() + self.assertEqual(moe_a8_vec.call_count, 2) + torch.testing.assert_close(output, torch.ones_like(output)) + + +if __name__ == "__main__": + unittest.main()