From 88ed4273ebb5fa1936701e50a6b39bbe9fc3a087 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 6 Nov 2026 23:39:35 +0000 Subject: [PATCH 1/5] [CUDA] Add GatherPacked, MatMulPacked, and MoEPacked contrib ops for block-packed weights Add com.microsoft::GatherPacked (CPU and CUDA), com.microsoft::MatMulPacked (CUDA), and com.microsoft::MoEPacked (CUDA) contrib operators that consume 2D/3D block-packed quantized weight blobs (such as GGUF IQ1/IQ2/IQ3/IQ4/Q2_0/Q4/Q5/Q6_K/Q8_0) directly without offline repacking or split scale tensors. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cpu/cpu_contrib_kernels.cc | 12 + .../cpu/quantization/gather_packed.cc | 104 ++ .../cpu/quantization/gguf_quant_tables.h | 1037 +++++++++++++ .../contrib_ops/cuda/cuda_contrib_kernels.cc | 20 + .../contrib_ops/cuda/moe/moe_packed.cc | 135 ++ .../contrib_ops/cuda/moe/moe_packed.cu | 1306 +++++++++++++++++ onnxruntime/contrib_ops/cuda/moe/moe_packed.h | 70 + .../cuda/quantization/gather_packed.cc | 82 ++ .../cuda/quantization/gather_packed.cu | 103 ++ .../cuda/quantization/gather_packed.h | 41 + .../cuda/quantization/gguf_quant_cuda.cuh | 690 +++++++++ .../cuda/quantization/matmul_packed.cc | 90 ++ .../cuda/quantization/matmul_packed.cu | 376 +++++ .../cuda/quantization/matmul_packed.h | 48 + .../core/graph/contrib_ops/contrib_defs.cc | 86 ++ onnxruntime/core/graph/contrib_ops/ms_opset.h | 6 + .../transformers/test_gguf_packed_ops.py | 229 +++ 17 files changed, 4435 insertions(+) create mode 100644 onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc create mode 100644 onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h create mode 100644 onnxruntime/contrib_ops/cuda/moe/moe_packed.cc create mode 100644 onnxruntime/contrib_ops/cuda/moe/moe_packed.cu create mode 100644 onnxruntime/contrib_ops/cuda/moe/moe_packed.h create mode 100644 onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc create mode 100644 onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu create mode 100644 onnxruntime/contrib_ops/cuda/quantization/gather_packed.h create mode 100644 onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh create mode 100644 onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc create mode 100644 onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu create mode 100644 onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h create mode 100644 onnxruntime/test/python/transformers/test_gguf_packed_ops.py diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc index 3bbb8f6bc8c57..3d510ac1ee609 100644 --- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc @@ -78,6 +78,12 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulBnb4); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, int32_t, GatherBlockQuantized); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, int64_t, GatherBlockQuantized); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, int32_t, GatherPacked); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, int64_t, GatherPacked); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, int32_t, GatherPacked); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, GatherPacked); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, BFloat16, int32_t, GatherPacked); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, BFloat16, int64_t, GatherPacked); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int32_t, GatherBlockQuantized); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int64_t, GatherBlockQuantized); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Int4x2, int32_t, GatherBlockQuantized); @@ -427,6 +433,12 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc b/onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc new file mode 100644 index 0000000000000..e2ad697442256 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include + +#include "contrib_ops/cpu/quantization/gguf_quant_tables.h" +#include "core/common/common.h" +#include "core/common/float16.h" +#include "core/framework/op_kernel.h" +#include "core/framework/tensor_shape.h" +#include "core/platform/threadpool.h" + +namespace onnxruntime { +namespace contrib { + +template +class GatherPacked final : public OpKernel { + public: + explicit GatherPacked(const OpKernelInfo& info) : OpKernel(info) { + ORT_ENFORCE(info.GetAttr("K", &k_).IsOK() && k_ > 0, "GatherPacked requires K > 0"); + ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK(), "GatherPacked requires packed_type"); + row_bytes_ = gguf_quant::GgmlRowBytes(static_cast(ggml_type_), k_); + ORT_ENFORCE(row_bytes_ > 0, "GatherPacked: unsupported packed_type=", ggml_type_, " or invalid K=", k_); + } + + Status Compute(OpKernelContext* context) const override { + const Tensor* data = context->Input(0); + const Tensor* indices = context->Input(1); + const auto& data_shape = data->Shape(); + ORT_RETURN_IF_NOT(data_shape.NumDimensions() == 2, "GatherPacked data must be 2D [vocab_size, row_bytes]"); + const int64_t vocab_size = data_shape[0]; + ORT_RETURN_IF_NOT(data_shape[1] == row_bytes_, + "GatherPacked data.shape[1]=", data_shape[1], " does not match expected row_bytes=", row_bytes_); + + const auto& indices_shape = indices->Shape(); + TensorShapeVector out_dims = indices_shape.AsShapeVector(); + out_dims.push_back(k_); + Tensor* output = context->Output(0, TensorShape(out_dims)); + + const int64_t num_rows = indices_shape.Size(); + if (num_rows == 0) return Status::OK(); + + const uint8_t* data_ptr = data->Data(); + const Tind* idx_ptr = indices->Data(); + T2* out_ptr = output->MutableData(); + + auto* tp = context->GetOperatorThreadPool(); + const int type_id = static_cast(ggml_type_); + const int64_t k = k_; + const int64_t rb = row_bytes_; + + concurrency::ThreadPool::TryBatchParallelFor( + tp, static_cast(num_rows), + [&](std::ptrdiff_t r) { + const int64_t raw_id = static_cast(idx_ptr[r]); + const int64_t tok = (raw_id >= 0 && raw_id < vocab_size) ? raw_id : 0; + const uint8_t* src_row = data_ptr + tok * rb; + T2* dst_row = out_ptr + static_cast(r) * k; + if constexpr (std::is_same_v) { + gguf_quant::DequantizeGgufRowCpu(src_row, type_id, k, dst_row); + } else { + std::vector tmp(static_cast(k)); + gguf_quant::DequantizeGgufRowCpu(src_row, type_id, k, tmp.data()); + for (int64_t j = 0; j < k; ++j) { + dst_row[j] = T2(tmp[j]); + } + } + }, + 0); + + return Status::OK(); + } + + private: + int64_t k_{0}; + int64_t ggml_type_{-1}; + int64_t row_bytes_{0}; +}; + +#define REGISTER_GATHER_PACKED_CPU(T2, Tind) \ + ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ + GatherPacked, \ + kMSDomain, 1, \ + T2, Tind, \ + kCpuExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ + GatherPacked); + +REGISTER_GATHER_PACKED_CPU(float, int32_t) +REGISTER_GATHER_PACKED_CPU(float, int64_t) +REGISTER_GATHER_PACKED_CPU(MLFloat16, int32_t) +REGISTER_GATHER_PACKED_CPU(MLFloat16, int64_t) +REGISTER_GATHER_PACKED_CPU(BFloat16, int32_t) +REGISTER_GATHER_PACKED_CPU(BFloat16, int64_t) + +#undef REGISTER_GATHER_PACKED_CPU + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h b/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h new file mode 100644 index 0000000000000..896dec2a42285 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h @@ -0,0 +1,1037 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include "core/common/common.h" +#include "core/common/float16.h" + +namespace onnxruntime { +namespace contrib { +namespace gguf_quant { + +enum GgmlType : int { + kGgmlTypeF32 = 0, + kGgmlTypeF16 = 1, + kGgmlTypeQ4_0 = 2, + kGgmlTypeQ4_1 = 3, + kGgmlTypeQ5_0 = 6, + kGgmlTypeQ5_1 = 7, + kGgmlTypeQ8_0 = 8, + kGgmlTypeQ8_1 = 9, + kGgmlTypeQ2_K = 10, + kGgmlTypeQ3_K = 11, + kGgmlTypeQ4_K = 12, + kGgmlTypeQ5_K = 13, + kGgmlTypeQ6_K = 14, + kGgmlTypeQ8_K = 15, + kGgmlTypeIQ2_XXS = 16, + kGgmlTypeIQ2_XS = 17, + kGgmlTypeIQ3_XXS = 18, + kGgmlTypeIQ1_S = 19, + kGgmlTypeIQ4_NL = 20, + kGgmlTypeIQ3_S = 21, + kGgmlTypeIQ2_S = 22, + kGgmlTypeIQ4_XS = 23, + kGgmlTypeIQ1_M = 29, + kGgmlTypeBF16 = 30, + kGgmlTypeMXFP4 = 39, + kGgmlTypeNVFP4 = 40, + kGgmlTypeQ1_0 = 41, + kGgmlTypeQ2_0 = 42, +}; + +inline bool GgmlTypeBlockInfo(int ggml_type, int64_t& block_elems, int64_t& block_bytes) { + switch (ggml_type) { + case kGgmlTypeF32: block_elems = 1; block_bytes = 4; return true; + case kGgmlTypeF16: block_elems = 1; block_bytes = 2; return true; + case kGgmlTypeBF16: block_elems = 1; block_bytes = 2; return true; + case kGgmlTypeQ4_0: block_elems = 32; block_bytes = 18; return true; + case kGgmlTypeQ4_1: block_elems = 32; block_bytes = 20; return true; + case kGgmlTypeQ5_0: block_elems = 32; block_bytes = 22; return true; + case kGgmlTypeQ5_1: block_elems = 32; block_bytes = 24; return true; + case kGgmlTypeQ8_0: block_elems = 32; block_bytes = 34; return true; + case kGgmlTypeQ8_1: block_elems = 32; block_bytes = 36; return true; + case kGgmlTypeQ2_K: block_elems = 256; block_bytes = 84; return true; + case kGgmlTypeQ3_K: block_elems = 256; block_bytes = 110; return true; + case kGgmlTypeQ4_K: block_elems = 256; block_bytes = 144; return true; + case kGgmlTypeQ5_K: block_elems = 256; block_bytes = 176; return true; + case kGgmlTypeQ6_K: block_elems = 256; block_bytes = 210; return true; + case kGgmlTypeQ8_K: block_elems = 256; block_bytes = 292; return true; + case kGgmlTypeIQ2_XXS: block_elems = 256; block_bytes = 66; return true; + case kGgmlTypeIQ2_XS: block_elems = 256; block_bytes = 74; return true; + case kGgmlTypeIQ3_XXS: block_elems = 256; block_bytes = 98; return true; + case kGgmlTypeIQ1_S: block_elems = 256; block_bytes = 50; return true; + case kGgmlTypeIQ4_NL: block_elems = 32; block_bytes = 18; return true; + case kGgmlTypeIQ3_S: block_elems = 256; block_bytes = 110; return true; + case kGgmlTypeIQ2_S: block_elems = 256; block_bytes = 82; return true; + case kGgmlTypeIQ4_XS: block_elems = 256; block_bytes = 136; return true; + case kGgmlTypeIQ1_M: block_elems = 256; block_bytes = 56; return true; + case kGgmlTypeMXFP4: block_elems = 32; block_bytes = 17; return true; + case kGgmlTypeNVFP4: block_elems = 64; block_bytes = 36; return true; + case kGgmlTypeQ1_0: block_elems = 128; block_bytes = 18; return true; + case kGgmlTypeQ2_0: block_elems = 64; block_bytes = 18; return true; + default: return false; + } +} + +inline int64_t GgmlRowBytes(int ggml_type, int64_t K) { + int64_t block_elems = 0, block_bytes = 0; + if (!GgmlTypeBlockInfo(ggml_type, block_elems, block_bytes) || K < 0 || (K % block_elems) != 0) { + return -1; + } + return (K / block_elems) * block_bytes; +} + +#ifndef GGML_TABLE_BEGIN +#if defined(__CUDA_ARCH__) +#define GGML_TABLE_BEGIN(type, name, size) static const __device__ __align__(16) type name[size] = { +#else +#define GGML_TABLE_BEGIN(type, name, size) static constexpr type name[size] = { +#endif +#define GGML_TABLE_END() }; +#endif + +GGML_TABLE_BEGIN(uint8_t, kmask_iq2xs, 8) + 1, 2, 4, 8, 16, 32, 64, 128 +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint8_t, ksigns_iq2xs, 128) + 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, + 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, + 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, + 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, + 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, + 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, + 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, + 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint64_t, ksigns64, 128) + 0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, + 0xff00000000ff0000, 0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff, + 0xff000000ff000000, 0x00000000ff0000ff, 0x00000000ff00ff00, 0xff000000ff00ffff, + 0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00, 0x00000000ffffffff, + 0xff0000ff00000000, 0x000000ff000000ff, 0x000000ff0000ff00, 0xff0000ff0000ffff, + 0x000000ff00ff0000, 0xff0000ff00ff00ff, 0xff0000ff00ffff00, 0x000000ff00ffffff, + 0x000000ffff000000, 0xff0000ffff0000ff, 0xff0000ffff00ff00, 0x000000ffff00ffff, + 0xff0000ffffff0000, 0x000000ffffff00ff, 0x000000ffffffff00, 0xff0000ffffffffff, + 0xff00ff0000000000, 0x0000ff00000000ff, 0x0000ff000000ff00, 0xff00ff000000ffff, + 0x0000ff0000ff0000, 0xff00ff0000ff00ff, 0xff00ff0000ffff00, 0x0000ff0000ffffff, + 0x0000ff00ff000000, 0xff00ff00ff0000ff, 0xff00ff00ff00ff00, 0x0000ff00ff00ffff, + 0xff00ff00ffff0000, 0x0000ff00ffff00ff, 0x0000ff00ffffff00, 0xff00ff00ffffffff, + 0x0000ffff00000000, 0xff00ffff000000ff, 0xff00ffff0000ff00, 0x0000ffff0000ffff, + 0xff00ffff00ff0000, 0x0000ffff00ff00ff, 0x0000ffff00ffff00, 0xff00ffff00ffffff, + 0xff00ffffff000000, 0x0000ffffff0000ff, 0x0000ffffff00ff00, 0xff00ffffff00ffff, + 0x0000ffffffff0000, 0xff00ffffffff00ff, 0xff00ffffffffff00, 0x0000ffffffffffff, + 0xffff000000000000, 0x00ff0000000000ff, 0x00ff00000000ff00, 0xffff00000000ffff, + 0x00ff000000ff0000, 0xffff000000ff00ff, 0xffff000000ffff00, 0x00ff000000ffffff, + 0x00ff0000ff000000, 0xffff0000ff0000ff, 0xffff0000ff00ff00, 0x00ff0000ff00ffff, + 0xffff0000ffff0000, 0x00ff0000ffff00ff, 0x00ff0000ffffff00, 0xffff0000ffffffff, + 0x00ff00ff00000000, 0xffff00ff000000ff, 0xffff00ff0000ff00, 0x00ff00ff0000ffff, + 0xffff00ff00ff0000, 0x00ff00ff00ff00ff, 0x00ff00ff00ffff00, 0xffff00ff00ffffff, + 0xffff00ffff000000, 0x00ff00ffff0000ff, 0x00ff00ffff00ff00, 0xffff00ffff00ffff, + 0x00ff00ffffff0000, 0xffff00ffffff00ff, 0xffff00ffffffff00, 0x00ff00ffffffffff, + 0x00ffff0000000000, 0xffffff00000000ff, 0xffffff000000ff00, 0x00ffff000000ffff, + 0xffffff0000ff0000, 0x00ffff0000ff00ff, 0x00ffff0000ffff00, 0xffffff0000ffffff, + 0xffffff00ff000000, 0x00ffff00ff0000ff, 0x00ffff00ff00ff00, 0xffffff00ff00ffff, + 0x00ffff00ffff0000, 0xffffff00ffff00ff, 0xffffff00ffffff00, 0x00ffff00ffffffff, + 0xffffffff00000000, 0x00ffffff000000ff, 0x00ffffff0000ff00, 0xffffffff0000ffff, + 0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff, + 0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff, + 0xffffffffffff0000, 0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff, +GGML_TABLE_END() + + +GGML_TABLE_BEGIN(uint64_t, iq2xxs_grid, 256) + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, + 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, + 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, + 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, + 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, + 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, + 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, + 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, + 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, + 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, + 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, + 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, + 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, + 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, + 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, + 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, + 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, + 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, + 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, + 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, + 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, + 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, + 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, + 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, + 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, + 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, + 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, + 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, + 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, + 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, + 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, + 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, + 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, + 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, + 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, + 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, + 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, + 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, + 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, + 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, + 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, + 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, + 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, + 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, + 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, + 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, + 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, + 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, + 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, + 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, + 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, + 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, + 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, + 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, + 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, + 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, + 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, + 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint64_t, iq2xs_grid, 512) + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, + 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, + 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, + 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, + 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, + 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, + 0x080808082b191908, 0x080808082b192b19, 0x080808082b2b0808, 0x0808081908080819, + 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, + 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, + 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, + 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, + 0x08080819192b0808, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908, + 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, + 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, + 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, + 0x0808082b2b080808, 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, + 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, + 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, + 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, + 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, + 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, + 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908190819, + 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, 0x0808191919081908, + 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, + 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, + 0x0808192b1908082b, 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, + 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808082b2b, 0x08082b0808190819, + 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, + 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, + 0x08082b1908190808, 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, + 0x08082b2b08080808, 0x08082b2b082b0808, 0x08082b2b082b2b08, 0x08082b2b2b19192b, + 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, + 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, + 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, + 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, + 0x0819080819191908, 0x08190808192b0808, 0x08190808192b2b2b, 0x081908082b080819, + 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, 0x081908190808082b, + 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, + 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, + 0x081908192b080808, 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, + 0x0819082b08081908, 0x0819082b0808192b, 0x0819082b08190808, 0x0819082b19080808, + 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, + 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, + 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, + 0x08191908192b1908, 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, + 0x0819191908190808, 0x0819191919080808, 0x0819192b08080808, 0x0819192b08191908, + 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, + 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, + 0x08192b2b2b2b2b19, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, + 0x082b080808082b08, 0x082b080808082b2b, 0x082b080808190819, 0x082b080808191908, + 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808, + 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, + 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, + 0x082b082b08080808, 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, + 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b1908082b2b19, + 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, 0x082b19191919082b, + 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, + 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, + 0x082b2b0819191919, 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, + 0x082b2b192b190808, 0x082b2b2b08082b08, 0x082b2b2b082b0808, 0x082b2b2b2b08082b, + 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, + 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, + 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, + 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, + 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x19080808192b0808, + 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, + 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, + 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, + 0x1908081919081908, 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, + 0x190808192b2b082b, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, + 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, 0x1908190808080808, + 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, + 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, + 0x1908190819081908, 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, + 0x1908191908080819, 0x1908191908081908, 0x1908191908190808, 0x19081919082b1908, + 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, 0x1908192b08082b2b, + 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, + 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, + 0x19082b08192b082b, 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, + 0x19082b1919190808, 0x19082b19192b2b19, 0x19082b2b08081908, 0x1919080808080808, + 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, + 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, + 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, + 0x1919081908081908, 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, + 0x191908191908082b, 0x1919082b08080808, 0x1919082b19081908, 0x1919082b2b2b2b2b, + 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, 0x19191908082b0819, + 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, + 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, + 0x1919192b082b0819, 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, + 0x19192b0808191908, 0x19192b0819080819, 0x19192b0819190808, 0x19192b082b192b19, + 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, 0x19192b2b2b081919, + 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, + 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, + 0x192b081908080808, 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, + 0x192b190808080808, 0x192b19080819192b, 0x192b191908190808, 0x192b191919080808, + 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, 0x192b2b08192b2b2b, + 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, + 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, + 0x2b08080808190819, 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, + 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808082b080808, + 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, 0x2b08081908080819, + 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, + 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, + 0x2b08082b2b080808, 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, + 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b0819080819082b, + 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b0819082b082b19, + 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, + 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, + 0x2b082b0819192b2b, 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, + 0x2b082b190808192b, 0x2b082b2b082b082b, 0x2b082b2b2b080808, 0x2b082b2b2b082b08, + 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, 0x2b19080808081908, + 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, + 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, + 0x2b19082b2b082b19, 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, + 0x2b19190819190808, 0x2b19190819192b08, 0x2b191919082b2b19, 0x2b1919192b190808, + 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, 0x2b192b082b2b192b, + 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, + 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, + 0x2b2b0808082b2b2b, 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, + 0x2b2b08192b2b192b, 0x2b2b082b08080808, 0x2b2b082b0808082b, 0x2b2b082b08082b08, + 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, 0x2b2b190819080808, + 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, + 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, + 0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, + 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint64_t, iq2s_grid, 1024) + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, + 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, + 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, + 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, + 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, + 0x08080808192b2b19, 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, + 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b2b0808, + 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, 0x0808081908081908, + 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, + 0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, + 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, + 0x0808081919190819, 0x0808081919191908, 0x080808191919192b, 0x0808081919192b19, + 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, 0x080808192b080819, + 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919, + 0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, + 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, + 0x0808082b082b0808, 0x0808082b082b2b2b, 0x0808082b19080819, 0x0808082b19081908, + 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, 0x0808082b19191919, + 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908, + 0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, + 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, + 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x08081908082b192b, + 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, + 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908, + 0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, + 0x08081908192b1919, 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, + 0x080819082b082b19, 0x080819082b190808, 0x080819082b191919, 0x080819082b192b08, + 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, + 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819, + 0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, + 0x08081919082b1919, 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, + 0x080819191908192b, 0x0808191919082b19, 0x0808191919190808, 0x080819191919082b, + 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, 0x08081919192b1908, + 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08, + 0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, + 0x0808192b08081908, 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, + 0x0808192b08191919, 0x0808192b19080808, 0x0808192b19081919, 0x0808192b19082b08, + 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, 0x0808192b2b080819, + 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b, + 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, + 0x08082b080819192b, 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, + 0x08082b08082b2b2b, 0x08082b0819080819, 0x08082b0819081908, 0x08082b081908192b, + 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, 0x08082b0819191919, + 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808, + 0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, + 0x08082b1908081908, 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, + 0x08082b1908192b08, 0x08082b19082b0819, 0x08082b1919080808, 0x08082b1919081919, + 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, 0x08082b19192b0808, + 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819, + 0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, + 0x08082b2b19190808, 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, + 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, + 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, + 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, + 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, + 0x0819080819192b19, 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, + 0x08190808192b2b08, 0x081908082b080819, 0x081908082b081908, 0x081908082b08192b, + 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, 0x081908082b2b0819, + 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, + 0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, + 0x081908190819192b, 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, + 0x08190819082b1919, 0x08190819082b2b08, 0x0819081919080819, 0x0819081919081908, + 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, 0x081908191919082b, + 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908, + 0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, + 0x081908192b190819, 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, + 0x0819082b08082b19, 0x0819082b08190808, 0x0819082b08191919, 0x0819082b082b0819, + 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, 0x0819082b19190819, + 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808, + 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, + 0x0819190808190819, 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, + 0x08191908082b0808, 0x08191908082b1919, 0x08191908082b2b08, 0x0819190819080819, + 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, 0x0819190819190808, + 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819, + 0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, + 0x081919082b082b08, 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, + 0x0819191908080819, 0x0819191908081908, 0x081919190808192b, 0x0819191908082b19, + 0x0819191908190808, 0x081919190819082b, 0x0819191908191919, 0x0819191908192b08, + 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b, + 0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, + 0x08191919192b0808, 0x081919192b080819, 0x081919192b081908, 0x081919192b190808, + 0x0819192b08080808, 0x0819192b08081919, 0x0819192b08082b08, 0x0819192b08190819, + 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, 0x0819192b19081908, + 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, + 0x08192b0808191919, 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, + 0x08192b081908082b, 0x08192b0819081919, 0x08192b0819082b08, 0x08192b0819190819, + 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, 0x08192b082b081908, + 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08, + 0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, + 0x08192b1919081908, 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, + 0x08192b2b08081908, 0x08192b2b08190808, 0x08192b2b19080808, 0x08192b2b1919192b, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, + 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19, + 0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, + 0x082b080819081908, 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, + 0x082b0808192b1908, 0x082b08082b080808, 0x082b08082b082b2b, 0x082b08082b191908, + 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808, + 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808, + 0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, + 0x082b0819192b0808, 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, + 0x082b082b08080808, 0x082b082b08082b2b, 0x082b082b082b082b, 0x082b082b082b2b08, + 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, 0x082b082b2b082b08, + 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, + 0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, + 0x082b190808192b08, 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, + 0x082b19081908082b, 0x082b190819081919, 0x082b190819082b08, 0x082b190819190819, + 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, 0x082b19082b081908, + 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08, + 0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, + 0x082b191919081908, 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, + 0x082b192b08080819, 0x082b192b08081908, 0x082b192b08190808, 0x082b192b19080808, + 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, 0x082b2b0808190819, + 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808, + 0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, + 0x082b2b1908190808, 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, + 0x082b2b2b192b1908, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, + 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, + 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b, + 0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, + 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, + 0x1908080819190819, 0x1908080819191908, 0x190808081919192b, 0x1908080819192b19, + 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, 0x190808082b080819, + 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08, + 0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, + 0x1908081908081919, 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, + 0x190808190819192b, 0x1908081908192b19, 0x19080819082b0808, 0x19080819082b082b, + 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, 0x190808191908192b, + 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919, + 0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, + 0x190808192b08082b, 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, + 0x190808192b191908, 0x190808192b2b0808, 0x1908082b08080819, 0x1908082b08081908, + 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, 0x1908082b08192b08, + 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08, + 0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, + 0x1908082b2b081908, 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, + 0x1908190808082b08, 0x1908190808082b2b, 0x1908190808190819, 0x1908190808191908, + 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, 0x19081908082b082b, + 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908, + 0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, + 0x1908190819191919, 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, + 0x190819082b080808, 0x190819082b08082b, 0x190819082b081919, 0x190819082b082b08, + 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, 0x1908191908080819, + 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808, + 0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, + 0x19081919082b1908, 0x1908191919080808, 0x190819191908082b, 0x1908191919081919, + 0x1908191919082b08, 0x1908191919190819, 0x1908191919191908, 0x19081919192b0808, + 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, 0x190819192b190808, + 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08, + 0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, + 0x1908192b19081908, 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, + 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808082b19, 0x19082b0808190808, + 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, 0x19082b08082b0819, + 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919, + 0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, + 0x19082b082b081908, 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, + 0x19082b1908081919, 0x19082b1908082b08, 0x19082b1908190819, 0x19082b1908191908, + 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, 0x19082b1919190808, + 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908, + 0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, + 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, + 0x191908080819192b, 0x1919080808192b19, 0x19190808082b0808, 0x19190808082b082b, + 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908, + 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b, + 0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, + 0x191908082b080808, 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, + 0x191908082b190819, 0x191908082b191908, 0x1919081908080819, 0x1919081908081908, + 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, 0x191908190819082b, + 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908, + 0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, + 0x1919081919190819, 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, + 0x191908192b081908, 0x191908192b190808, 0x1919082b08080808, 0x1919082b08081919, + 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, 0x1919082b082b0808, + 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19, + 0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, + 0x1919190808082b19, 0x1919190808190808, 0x191919080819082b, 0x1919190808191919, + 0x1919190808192b08, 0x19191908082b0819, 0x19191908082b1908, 0x1919190819080808, + 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, 0x1919190819190819, + 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908, + 0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, + 0x1919191908082b08, 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, + 0x1919191919080819, 0x1919191919081908, 0x1919191919190808, 0x191919192b080808, + 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, 0x1919192b082b192b, + 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919, + 0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, + 0x19192b0819080819, 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, + 0x19192b082b080808, 0x19192b1908080819, 0x19192b1908081908, 0x19192b1908190808, + 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, 0x19192b2b2b081919, + 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b, + 0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, + 0x192b0808082b0819, 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, + 0x192b080819082b08, 0x192b080819190819, 0x192b080819191908, 0x192b0808192b0808, + 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, 0x192b08190808082b, + 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908, + 0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, + 0x192b08192b080808, 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, + 0x192b082b19080808, 0x192b082b1919192b, 0x192b082b2b2b0819, 0x192b190808080808, + 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, 0x192b190808191908, + 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808, + 0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, + 0x192b191919080808, 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, + 0x192b192b08080808, 0x192b192b2b191908, 0x192b2b0808080819, 0x192b2b0808081908, + 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, 0x192b2b1908080808, + 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808, + 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, + 0x2b08080808191908, 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, + 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808081919082b, + 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, 0x2b0808082b080808, + 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819, + 0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, + 0x2b08081908191919, 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, + 0x2b08081919080808, 0x2b0808191908082b, 0x2b08081919081919, 0x2b08081919082b08, + 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, 0x2b0808192b081908, + 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919, + 0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, + 0x2b08082b19081908, 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, + 0x2b0819080808192b, 0x2b08190808082b19, 0x2b08190808190808, 0x2b0819080819082b, + 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, 0x2b08190819080808, + 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819, + 0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, + 0x2b0819082b190808, 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, + 0x2b08191908082b08, 0x2b08191908190819, 0x2b08191908191908, 0x2b081919082b0808, + 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, 0x2b0819192b080808, + 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808, + 0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, + 0x2b082b0808190819, 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, + 0x2b082b0819190808, 0x2b082b082b2b082b, 0x2b082b1908080819, 0x2b082b1908081908, + 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, 0x2b082b2b19192b08, + 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819, + 0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, + 0x2b19080808191919, 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, + 0x2b1908081908082b, 0x2b19080819081919, 0x2b19080819082b08, 0x2b19080819190819, + 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, 0x2b1908082b081908, + 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819, + 0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, + 0x2b19081919192b2b, 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, + 0x2b19082b19080808, 0x2b19082b2b2b192b, 0x2b19190808080808, 0x2b1919080808082b, + 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, 0x2b19190808191908, + 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808, + 0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, + 0x2b19191908190808, 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, + 0x2b19192b08080808, 0x2b19192b1908192b, 0x2b19192b192b1908, 0x2b192b0808080819, + 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, 0x2b192b0819080808, + 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b, + 0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, + 0x2b2b080808191908, 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, + 0x2b2b080819081908, 0x2b2b080819190808, 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, + 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, 0x2b2b082b08082b2b, + 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b, + 0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, + 0x2b2b190808081908, 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, + 0x2b2b19082b2b1908, 0x2b2b191908080808, 0x2b2b191908192b19, 0x2b2b192b19190819, + 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, 0x2b2b2b1919191908, + 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, + 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint32_t, iq3xxs_grid, 256) + 0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, + 0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, + 0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, + 0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, + 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, + 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, 0x04243e1c, 0x04243e2c, 0x042c040c, + 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, 0x043e0c24, 0x043e0c34, + 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, 0x0c04141c, + 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c, + 0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, + 0x0c143e14, 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, + 0x0c24042c, 0x0c242c04, 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, + 0x0c3e2404, 0x14040404, 0x14040414, 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, + 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, + 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, 0x14140414, 0x14140c0c, 0x14140c3e, + 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, 0x141c0c04, 0x141c0c24, + 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, 0x142c3e24, + 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c, + 0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, + 0x1c0c2424, 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, + 0x1c1c0c0c, 0x1c1c1c1c, 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, + 0x1c2c2c2c, 0x1c340c24, 0x1c341c34, 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, + 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, + 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, 0x24242c0c, 0x24243424, 0x242c142c, + 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, 0x2c040c14, 0x2c04240c, + 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, 0x2c143e14, + 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c, + 0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, + 0x340c340c, 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, + 0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, + 0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, + 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint32_t, iq3s_grid, 512) + 0x01010101, 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305, + 0x01010309, 0x0101030d, 0x01010501, 0x01010503, 0x0101050b, 0x01010707, 0x01010901, 0x01010905, + 0x0101090b, 0x0101090f, 0x01010b03, 0x01010b07, 0x01010d01, 0x01010d05, 0x01010f03, 0x01010f09, + 0x01010f0f, 0x01030101, 0x01030103, 0x01030105, 0x01030109, 0x01030301, 0x01030303, 0x0103030b, + 0x01030501, 0x01030507, 0x0103050f, 0x01030703, 0x0103070b, 0x01030909, 0x01030d03, 0x01030d0b, + 0x01030f05, 0x01050101, 0x01050103, 0x0105010b, 0x0105010f, 0x01050301, 0x01050307, 0x0105030d, + 0x01050503, 0x0105050b, 0x01050701, 0x01050709, 0x01050905, 0x0105090b, 0x0105090f, 0x01050b03, + 0x01050b07, 0x01050f01, 0x01050f07, 0x01070107, 0x01070303, 0x0107030b, 0x01070501, 0x01070505, + 0x01070703, 0x01070707, 0x0107070d, 0x01070909, 0x01070b01, 0x01070b05, 0x01070d0f, 0x01070f03, + 0x01070f0b, 0x01090101, 0x01090307, 0x0109030f, 0x01090503, 0x01090509, 0x01090705, 0x01090901, + 0x01090907, 0x01090b03, 0x01090f01, 0x010b0105, 0x010b0109, 0x010b0501, 0x010b0505, 0x010b050d, + 0x010b0707, 0x010b0903, 0x010b090b, 0x010b090f, 0x010b0d0d, 0x010b0f07, 0x010d010d, 0x010d0303, + 0x010d0307, 0x010d0703, 0x010d0b05, 0x010d0f03, 0x010f0101, 0x010f0105, 0x010f0109, 0x010f0501, + 0x010f0505, 0x010f050d, 0x010f0707, 0x010f0b01, 0x010f0b09, 0x03010101, 0x03010103, 0x03010105, + 0x03010109, 0x03010301, 0x03010303, 0x03010307, 0x0301030b, 0x0301030f, 0x03010501, 0x03010505, + 0x03010703, 0x03010709, 0x0301070d, 0x03010b09, 0x03010b0d, 0x03010d03, 0x03010f05, 0x03030101, + 0x03030103, 0x03030107, 0x0303010d, 0x03030301, 0x03030309, 0x03030503, 0x03030701, 0x03030707, + 0x03030903, 0x03030b01, 0x03030b05, 0x03030f01, 0x03030f0d, 0x03050101, 0x03050305, 0x0305030b, + 0x0305030f, 0x03050501, 0x03050509, 0x03050705, 0x03050901, 0x03050907, 0x03050b0b, 0x03050d01, + 0x03050f05, 0x03070103, 0x03070109, 0x0307010f, 0x03070301, 0x03070307, 0x03070503, 0x0307050f, + 0x03070701, 0x03070709, 0x03070903, 0x03070d05, 0x03070f01, 0x03090107, 0x0309010b, 0x03090305, + 0x03090309, 0x03090703, 0x03090707, 0x03090905, 0x0309090d, 0x03090b01, 0x03090b09, 0x030b0103, + 0x030b0301, 0x030b0307, 0x030b0503, 0x030b0701, 0x030b0705, 0x030b0b03, 0x030d0501, 0x030d0509, + 0x030d050f, 0x030d0909, 0x030d090d, 0x030f0103, 0x030f0107, 0x030f0301, 0x030f0305, 0x030f0503, + 0x030f070b, 0x030f0903, 0x030f0d05, 0x030f0f01, 0x05010101, 0x05010103, 0x05010107, 0x0501010b, + 0x0501010f, 0x05010301, 0x05010305, 0x05010309, 0x0501030d, 0x05010503, 0x05010507, 0x0501050f, + 0x05010701, 0x05010705, 0x05010903, 0x05010907, 0x0501090b, 0x05010b01, 0x05010b05, 0x05010d0f, + 0x05010f01, 0x05010f07, 0x05010f0b, 0x05030101, 0x05030105, 0x05030301, 0x05030307, 0x0503030f, + 0x05030505, 0x0503050b, 0x05030703, 0x05030709, 0x05030905, 0x05030b03, 0x05050103, 0x05050109, + 0x0505010f, 0x05050503, 0x05050507, 0x05050701, 0x0505070f, 0x05050903, 0x05050b07, 0x05050b0f, + 0x05050f03, 0x05050f09, 0x05070101, 0x05070105, 0x0507010b, 0x05070303, 0x05070505, 0x05070509, + 0x05070703, 0x05070707, 0x05070905, 0x05070b01, 0x05070d0d, 0x05090103, 0x0509010f, 0x05090501, + 0x05090507, 0x05090705, 0x0509070b, 0x05090903, 0x05090f05, 0x05090f0b, 0x050b0109, 0x050b0303, + 0x050b0505, 0x050b070f, 0x050b0901, 0x050b0b07, 0x050b0f01, 0x050d0101, 0x050d0105, 0x050d010f, + 0x050d0503, 0x050d0b0b, 0x050d0d03, 0x050f010b, 0x050f0303, 0x050f050d, 0x050f0701, 0x050f0907, + 0x050f0b01, 0x07010105, 0x07010303, 0x07010307, 0x0701030b, 0x0701030f, 0x07010505, 0x07010703, + 0x07010707, 0x0701070b, 0x07010905, 0x07010909, 0x0701090f, 0x07010b03, 0x07010d07, 0x07010f03, + 0x07030103, 0x07030107, 0x0703010b, 0x07030309, 0x07030503, 0x07030507, 0x07030901, 0x07030d01, + 0x07030f05, 0x07030f0d, 0x07050101, 0x07050305, 0x07050501, 0x07050705, 0x07050709, 0x07050b01, + 0x07070103, 0x07070301, 0x07070309, 0x07070503, 0x07070507, 0x0707050f, 0x07070701, 0x07070903, + 0x07070907, 0x0707090f, 0x07070b0b, 0x07070f07, 0x07090107, 0x07090303, 0x0709030d, 0x07090505, + 0x07090703, 0x07090b05, 0x07090d01, 0x07090d09, 0x070b0103, 0x070b0301, 0x070b0305, 0x070b050b, + 0x070b0705, 0x070b0909, 0x070b0b0d, 0x070b0f07, 0x070d030d, 0x070d0903, 0x070f0103, 0x070f0107, + 0x070f0501, 0x070f0505, 0x070f070b, 0x09010101, 0x09010109, 0x09010305, 0x09010501, 0x09010509, + 0x0901050f, 0x09010705, 0x09010903, 0x09010b01, 0x09010f01, 0x09030105, 0x0903010f, 0x09030303, + 0x09030307, 0x09030505, 0x09030701, 0x0903070b, 0x09030907, 0x09030b03, 0x09030b0b, 0x09050103, + 0x09050107, 0x09050301, 0x0905030b, 0x09050503, 0x09050707, 0x09050901, 0x09050b0f, 0x09050d05, + 0x09050f01, 0x09070109, 0x09070303, 0x09070307, 0x09070501, 0x09070505, 0x09070703, 0x0907070b, + 0x09090101, 0x09090105, 0x09090509, 0x0909070f, 0x09090901, 0x09090f03, 0x090b010b, 0x090b010f, + 0x090b0503, 0x090b0d05, 0x090d0307, 0x090d0709, 0x090d0d01, 0x090f0301, 0x090f030b, 0x090f0701, + 0x090f0907, 0x090f0b03, 0x0b010105, 0x0b010301, 0x0b010309, 0x0b010505, 0x0b010901, 0x0b010909, + 0x0b01090f, 0x0b010b05, 0x0b010d0d, 0x0b010f09, 0x0b030103, 0x0b030107, 0x0b03010b, 0x0b030305, + 0x0b030503, 0x0b030705, 0x0b030f05, 0x0b050101, 0x0b050303, 0x0b050507, 0x0b050701, 0x0b05070d, + 0x0b050b07, 0x0b070105, 0x0b07010f, 0x0b070301, 0x0b07050f, 0x0b070909, 0x0b070b03, 0x0b070d0b, + 0x0b070f07, 0x0b090103, 0x0b090109, 0x0b090501, 0x0b090705, 0x0b09090d, 0x0b0b0305, 0x0b0b050d, + 0x0b0b0b03, 0x0b0b0b07, 0x0b0d0905, 0x0b0f0105, 0x0b0f0109, 0x0b0f0505, 0x0d010303, 0x0d010307, + 0x0d01030b, 0x0d010703, 0x0d010707, 0x0d010d01, 0x0d030101, 0x0d030501, 0x0d03050f, 0x0d030d09, + 0x0d050305, 0x0d050709, 0x0d050905, 0x0d050b0b, 0x0d050d05, 0x0d050f01, 0x0d070101, 0x0d070309, + 0x0d070503, 0x0d070901, 0x0d09050b, 0x0d090907, 0x0d090d05, 0x0d0b0101, 0x0d0b0107, 0x0d0b0709, + 0x0d0b0d01, 0x0d0d010b, 0x0d0d0901, 0x0d0f0303, 0x0d0f0307, 0x0f010101, 0x0f010109, 0x0f01010f, + 0x0f010501, 0x0f010505, 0x0f01070d, 0x0f010901, 0x0f010b09, 0x0f010d05, 0x0f030105, 0x0f030303, + 0x0f030509, 0x0f030907, 0x0f03090b, 0x0f050103, 0x0f050109, 0x0f050301, 0x0f05030d, 0x0f050503, + 0x0f050701, 0x0f050b03, 0x0f070105, 0x0f070705, 0x0f07070b, 0x0f070b07, 0x0f090103, 0x0f09010b, + 0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101, +GGML_TABLE_END() + +// TODO: fix name to kvalues_iq4_nl +GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16) + -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113, +GGML_TABLE_END() + +// e2m1 values (doubled), shared by MXFP4 and NVFP4 +// ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf +GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16) + 0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12, +GGML_TABLE_END() +#define kvalues_mxfp4 kvalues_fp4 + +#undef GGML_TABLE_BEGIN +#undef GGML_TABLE_END + +#if !defined(__CUDA_ARCH__) + +inline float Fp16BitsToFloatCpu(uint16_t bits) { + return MLFloat16::FromBits(bits).ToFloat(); +} + +inline float Bf16BitsToFloatCpu(uint16_t bits) { + uint32_t u32 = static_cast(bits) << 16; + float f = 0.0f; + std::memcpy(&f, &u32, sizeof(f)); + return f; +} + +inline uint16_t ReadU16Le(const uint8_t* p) { + uint16_t v = 0; + std::memcpy(&v, p, sizeof(v)); + return v; +} + +inline uint32_t ReadU32Le(const uint8_t* p) { + uint32_t v = 0; + std::memcpy(&v, p, sizeof(v)); + return v; +} + +inline void GetScaleMinK4Cpu(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = q[j] & 63; + m = q[j + 4] & 63; + } else { + d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); + m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +inline bool DequantizeGgufRowCpu(const uint8_t* src, int ggml_type, int64_t K, float* dst) { + switch (ggml_type) { + case kGgmlTypeF32: { + std::memcpy(dst, src, static_cast(K) * sizeof(float)); + return true; + } + case kGgmlTypeF16: { + for (int64_t i = 0; i < K; ++i) { + dst[i] = Fp16BitsToFloatCpu(ReadU16Le(src + i * 2)); + } + return true; + } + case kGgmlTypeBF16: { + for (int64_t i = 0; i < K; ++i) { + dst[i] = Bf16BitsToFloatCpu(ReadU16Le(src + i * 2)); + } + return true; + } + case kGgmlTypeQ8_0: { + const int64_t nb = K / 32; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 34; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const int8_t* qs = reinterpret_cast(blk + 2); + for (int j = 0; j < 32; ++j) { + dst[i * 32 + j] = d * static_cast(qs[j]); + } + } + return true; + } + case kGgmlTypeQ4_0: { + const int64_t nb = K / 32; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 18; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint8_t* qs = blk + 2; + for (int j = 0; j < 16; ++j) { + dst[i * 32 + j] = static_cast((qs[j] & 0xF) - 8) * d; + dst[i * 32 + j + 16] = static_cast((qs[j] >> 4) - 8) * d; + } + } + return true; + } + case kGgmlTypeQ5_0: { + const int64_t nb = K / 32; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 22; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint32_t qh = ReadU32Le(blk + 2); + const uint8_t* qs = blk + 6; + for (int j = 0; j < 16; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12))) & 0x10; + const int32_t x0 = ((qs[j] & 0x0F) | xh_0) - 16; + const int32_t x1 = ((qs[j] >> 4) | xh_1) - 16; + dst[i * 32 + j] = static_cast(x0) * d; + dst[i * 32 + j + 16] = static_cast(x1) * d; + } + } + return true; + } + case kGgmlTypeQ2_0: { + const int64_t nb = K / 64; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 18; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint8_t* qs = blk + 2; + for (int j = 0; j < 64; ++j) { + const int q = (qs[j >> 2] >> ((j & 3) * 2)) & 0x3; + dst[i * 64 + j] = static_cast(q - 1) * d; + } + } + return true; + } + case kGgmlTypeIQ4_NL: { + const int64_t nb = K / 32; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 18; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint8_t* qs = blk + 2; + for (int j = 0; j < 16; ++j) { + dst[i * 32 + j] = d * static_cast(kvalues_iq4nl[qs[j] & 0xF]); + dst[i * 32 + j + 16] = d * static_cast(kvalues_iq4nl[qs[j] >> 4]); + } + } + return true; + } + case kGgmlTypeIQ4_XS: { + const int64_t nb = K / 256; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 136; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint16_t scales_h = ReadU16Le(blk + 2); + const uint8_t* scales_l = blk + 4; + const uint8_t* qs = blk + 8; + float* y = dst + i * 256; + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((scales_l[ib / 2] >> (4 * (ib % 2))) & 0xF) | (((scales_h >> (2 * ib)) & 3) << 4); + const float dl = d * static_cast(ls - 32); + for (int j = 0; j < 16; ++j) { + y[j + 0] = dl * static_cast(kvalues_iq4nl[qs[j] & 0xF]); + y[j + 16] = dl * static_cast(kvalues_iq4nl[qs[j] >> 4]); + } + y += 32; + qs += 16; + } + } + return true; + } + case kGgmlTypeQ5_K: { + const int64_t nb = K / 256; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 176; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const float min = Fp16BitsToFloatCpu(ReadU16Le(blk + 2)); + const uint8_t* scales = blk + 4; + const uint8_t* qh = blk + 16; + const uint8_t* ql = blk + 48; + float* y = dst + i * 256; + int is = 0; + uint8_t u1 = 1, u2 = 2; + for (int j = 0; j < 256; j += 64) { + uint8_t sc, m; + GetScaleMinK4Cpu(is + 0, scales, sc, m); + const float d1 = d * sc; + const float m1 = min * m; + GetScaleMinK4Cpu(is + 1, scales, sc, m); + const float d2 = d * sc; + const float m2 = min * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2; + ql += 32; + is += 2; + u1 <<= 2; + u2 <<= 2; + } + } + return true; + } + case kGgmlTypeQ6_K: { + const int64_t nb = K / 256; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 210; + const uint8_t* ql = blk; + const uint8_t* qh = blk + 128; + const int8_t* sc = reinterpret_cast(blk + 192); + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk + 208)); + float* y = dst + i * 256; + for (int n = 0; n < 256; n += 128) { + for (int l = 0; l < 32; ++l) { + const int is = l / 16; + const int8_t q1 = static_cast((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = static_cast((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = static_cast((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = static_cast((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + y[l + 0] = d * sc[is + 0] * q1; + y[l + 32] = d * sc[is + 2] * q2; + y[l + 64] = d * sc[is + 4] * q3; + y[l + 96] = d * sc[is + 6] * q4; + } + y += 128; + ql += 64; + qh += 32; + sc += 8; + } + } + return true; + } + case kGgmlTypeIQ2_XXS: { + const int64_t nb = K / 256; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 66; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint8_t* qs = blk + 2; + float* y = dst + i * 256; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const uint8_t* aux8 = qs + 8 * ib32; + const uint32_t aux32 = ReadU32Le(aux8 + 4); + const float db = d * (0.5f + (aux32 >> 28)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint8_t* grid = reinterpret_cast(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * l)) & 127]; + for (int j = 0; j < 8; ++j) { + y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + } + } + return true; + } + case kGgmlTypeIQ2_XS: { + const int64_t nb = K / 256; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 74; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint8_t* qs = blk + 2; + const uint8_t* scales = blk + 66; + float* y = dst + i * 256; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const float db0 = d * (0.5f + (scales[ib32] & 0xf)) * 0.25f; + const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint16_t qval = ReadU16Le(qs + (4 * ib32 + l) * 2); + const uint8_t* grid = reinterpret_cast(iq2xs_grid + (qval & 511)); + const uint8_t signs = ksigns_iq2xs[qval >> 9]; + const float db = (l < 2) ? db0 : db1; + for (int j = 0; j < 8; ++j) { + y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + } + } + return true; + } + case kGgmlTypeIQ2_S: { + const int64_t nb = K / 256; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 82; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint8_t* qs = blk + 2; + const uint8_t* signs = qs + 32; + const uint8_t* qh = blk + 66; + const uint8_t* scales = blk + 74; + float* y = dst + i * 256; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const float db0 = d * (0.5f + (scales[ib32] & 0xf)) * 0.25f; + const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const float dl = (l < 2) ? db0 : db1; + const uint8_t* grid = + reinterpret_cast(iq2s_grid + (qs[l] | ((qh[ib32] << (8 - 2 * l)) & 0x300))); + for (int j = 0; j < 8; ++j) { + y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + qs += 4; + signs += 4; + } + } + return true; + } + case kGgmlTypeIQ3_XXS: { + const int64_t nb = K / 256; + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* blk = src + i * 98; + const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)); + const uint8_t* qs = blk + 2; + const uint8_t* scales_and_signs = qs + 64; + float* y = dst + i * 256; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const uint32_t aux32 = ReadU32Le(scales_and_signs + 4 * ib32); + const float db = d * (0.5f + (aux32 >> 28)) * 0.5f; + for (int l = 0; l < 4; ++l) { + const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * l)) & 127]; + const uint8_t* grid1 = reinterpret_cast(iq3xxs_grid + qs[2 * l + 0]); + const uint8_t* grid2 = reinterpret_cast(iq3xxs_grid + qs[2 * l + 1]); + for (int j = 0; j < 4; ++j) { + y[j + 0] = db * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + } + } + return true; + } + default: + return false; + } +} + +#endif // !defined(__CUDA_ARCH__) + +} // namespace gguf_quant +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index c9249b7304446..404d518532f77 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -166,6 +166,16 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, SwiGLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, SwiGLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, RouterGateUpProj); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, RouterGateUpProj); +class CUDA_MS_OP_TWO_TYPED_CLASS_NAME(1, float, int32_t, GatherPacked); +class CUDA_MS_OP_TWO_TYPED_CLASS_NAME(1, float, int64_t, GatherPacked); +class CUDA_MS_OP_TWO_TYPED_CLASS_NAME(1, MLFloat16, int32_t, GatherPacked); +class CUDA_MS_OP_TWO_TYPED_CLASS_NAME(1, MLFloat16, int64_t, GatherPacked); +class CUDA_MS_OP_TWO_TYPED_CLASS_NAME(1, BFloat16, int32_t, GatherPacked); +class CUDA_MS_OP_TWO_TYPED_CLASS_NAME(1, BFloat16, int64_t, GatherPacked); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MatMulPacked); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MatMulPacked); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoEPacked); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoEPacked); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MRotaryEmbedding); @@ -541,6 +551,16 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc new file mode 100644 index 0000000000000..8a88c864ba786 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc @@ -0,0 +1,135 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/moe/moe_packed.h" + +#include +#include + +#include "contrib_ops/cpu/quantization/gguf_quant_tables.h" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define REGISTER_MOE_PACKED_CUDA(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MoEPacked, \ + kMSDomain, 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()),\ + MoEPacked); + +REGISTER_MOE_PACKED_CUDA(MLFloat16) +REGISTER_MOE_PACKED_CUDA(BFloat16) + +#undef REGISTER_MOE_PACKED_CUDA + +template +MoEPacked::MoEPacked(const OpKernelInfo& info) : CudaKernel(info) { + ORT_ENFORCE(info.GetAttr("k", &k_).IsOK() && k_ > 0 && k_ <= 16, + "MoEPacked requires k in 1..16"); + normalize_routing_weights_ = info.GetAttrOrDefault("normalize_routing_weights", 1); + activation_type_ = info.GetAttrOrDefault("activation_type", "swiglu"); + ORT_ENFORCE(activation_type_ == "swiglu", "MoEPacked only supports activation_type='swiglu'"); + ORT_ENFORCE(info.GetAttr("gate_type", &gate_type_).IsOK(), "MoEPacked requires gate_type"); + ORT_ENFORCE(info.GetAttr("up_type", &up_type_).IsOK(), "MoEPacked requires up_type"); + ORT_ENFORCE(info.GetAttr("down_type", &down_type_).IsOK(), "MoEPacked requires down_type"); +} + +template +Status MoEPacked::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + const Tensor* router_probs = context->Input(1); + const Tensor* gate_w = context->Input(2); + const Tensor* up_w = context->Input(3); + const Tensor* down_w = context->Input(4); + + const auto& in_shape = input->Shape(); + const auto& rp_shape = router_probs->Shape(); + const auto& gw_shape = gate_w->Shape(); + const auto& uw_shape = up_w->Shape(); + const auto& dw_shape = down_w->Shape(); + + ORT_RETURN_IF_NOT(in_shape.NumDimensions() >= 2, "MoEPacked input must have rank >= 2"); + ORT_RETURN_IF_NOT(rp_shape.NumDimensions() == 2, "MoEPacked router_probs must be 2D [num_rows, num_experts]"); + ORT_RETURN_IF_NOT(gw_shape.NumDimensions() == 3 && uw_shape.NumDimensions() == 3 && dw_shape.NumDimensions() == 3, + "MoEPacked expert weights must be 3D [num_experts, out_dim, row_bytes]"); + + const size_t last_axis = in_shape.NumDimensions() - 1; + const int64_t hidden_size = in_shape[last_axis]; + const int64_t num_rows = in_shape.SizeToDimension(last_axis); + const int64_t num_experts = rp_shape[1]; + const int64_t inter_size = gw_shape[1]; + + ORT_RETURN_IF_NOT(rp_shape[0] == num_rows, "MoEPacked router_probs row count mismatch"); + ORT_RETURN_IF_NOT(num_experts > 0 && num_experts <= 512, "MoEPacked supports up to 512 experts"); + ORT_RETURN_IF_NOT(gw_shape[0] == num_experts && uw_shape[0] == num_experts && dw_shape[0] == num_experts, + "MoEPacked expert count mismatch"); + ORT_RETURN_IF_NOT(uw_shape[1] == inter_size && dw_shape[1] == hidden_size, + "MoEPacked expert inner/outer dimension mismatch"); + ORT_RETURN_IF_NOT((hidden_size % 32) == 0 && (inter_size % 32) == 0, + "MoEPacked requires hidden_size and inter_size to be multiples of 32"); + + const int64_t exp_gate_rb = gguf_quant::GgmlRowBytes(static_cast(gate_type_), hidden_size); + const int64_t exp_up_rb = gguf_quant::GgmlRowBytes(static_cast(up_type_), hidden_size); + const int64_t exp_down_rb = gguf_quant::GgmlRowBytes(static_cast(down_type_), inter_size); + ORT_RETURN_IF_NOT(gw_shape[2] == exp_gate_rb && uw_shape[2] == exp_up_rb && dw_shape[2] == exp_down_rb, + "MoEPacked expert weight row_bytes mismatch"); + + Tensor* output = context->Output(0, in_shape); + if (num_rows == 0) return Status::OK(); + + // Process in chunks of at most 2048 rows so workspace stays compact during large prefills. + constexpr int64_t kMaxChunkRows = 2048; + const int64_t max_chunk = std::min(num_rows, kMaxChunkRows); + + MoEPackedParams max_params{}; + max_params.num_rows = static_cast(max_chunk); + max_params.num_experts = static_cast(num_experts); + max_params.hidden_size = static_cast(hidden_size); + max_params.inter_size = static_cast(inter_size); + max_params.k = static_cast(k_); + max_params.normalize_routing_weights = (normalize_routing_weights_ != 0); + max_params.gate_type = static_cast(gate_type_); + max_params.up_type = static_cast(up_type_); + max_params.down_type = static_cast(down_type_); + max_params.gate_row_bytes = exp_gate_rb; + max_params.up_row_bytes = exp_up_rb; + max_params.down_row_bytes = exp_down_rb; + + const size_t ws_bytes = GetMoEPackedWorkspaceBytes(max_params, sizeof(T)); + auto workspace = GetScratchBuffer(ws_bytes, GetComputeStream(context)); + constexpr bool is_bf16 = std::is_same_v; + + const uint8_t* in_bytes = static_cast(input->DataRaw()); + const uint8_t* rp_bytes = static_cast(router_probs->DataRaw()); + uint8_t* out_bytes = static_cast(output->MutableDataRaw()); + + for (int64_t r0 = 0; r0 < num_rows; r0 += kMaxChunkRows) { + const int64_t chunk_rows = std::min(kMaxChunkRows, num_rows - r0); + MoEPackedParams chunk_params = max_params; + chunk_params.num_rows = static_cast(chunk_rows); + ORT_RETURN_IF_ERROR(LaunchMoEPackedCuda( + Stream(context), + is_bf16, + in_bytes + static_cast(r0 * hidden_size) * sizeof(T), + rp_bytes + static_cast(r0 * num_experts) * sizeof(T), + gate_w->Data(), + up_w->Data(), + down_w->Data(), + chunk_params, + workspace.get(), + out_bytes + static_cast(r0 * hidden_size) * sizeof(T))); + } + + return Status::OK(); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu new file mode 100644 index 0000000000000..002ebe2fbc1e7 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu @@ -0,0 +1,1306 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/moe/moe_packed.h" + +#include +#include + +#include "contrib_ops/cuda/moe/qmoe_fused_route.cuh" +#include "contrib_ops/cuda/quantization/gguf_quant_cuda.cuh" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +namespace { + +constexpr int kMoEPackedMaxExperts = 512; +constexpr int kMoEPackedMaxK = 16; + +// Softmax + top-k over 512 experts for arbitrary num_rows (8 warps per CTA, 1 warp per row). +// Uses the exact same warp reduction and tie-breaking key as FusedSoftmaxTopKMapsKernel. +template +__global__ __launch_bounds__(256) void MoEPackedSoftmaxTopK512Kernel( + const CudaT* __restrict__ logits, + float* __restrict__ topk_scales, + int* __restrict__ topk_indices, + int num_rows, + int k, + bool normalize_scales) { + constexpr int num_experts = 512; + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int row = static_cast(blockIdx.x) * 8 + warp; + if (row >= num_rows) return; + + const CudaT* row_logits = logits + static_cast(row) * num_experts; + float scores[16]; + uint32_t keys[16]; + float local_max = onnxruntime::cuda::topk::kNegativeInfinity; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const int base = (lane + 32 * w) * 4; + float v[4]; + if ((reinterpret_cast(row_logits + base) & 7u) == 0) { + FusedRouteLoad4(row_logits + base, v); + } else { +#pragma unroll + for (int j = 0; j < 4; ++j) { + v[j] = gguf_cuda::CudaTypeToFloat(row_logits[base + j]); + } + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + scores[w * 4 + j] = v[j] == 0.0f ? 0.0f : v[j]; + keys[w * 4 + j] = FusedRouteKey(v[j], base + j, true); + local_max = fmaxf(local_max, v[j]); + } + } + +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const float other = __shfl_xor_sync(0xFFFFFFFFu, local_max, offset); + local_max = (local_max < other) ? other : local_max; + } + const float max_val = local_max; + + uint32_t winners[kMoEPackedMaxK]; +#pragma unroll + for (int r = 0; r < kMoEPackedMaxK; ++r) { + if (r < k) { + uint32_t best = keys[0]; +#pragma unroll + for (int j = 1; j < 16; ++j) { + best = max(best, keys[j]); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + winners[r] = __reduce_max_sync(0xFFFFFFFFu, best); +#else +#pragma unroll + for (int offset = 16; offset >= 1; offset >>= 1) { + best = max(best, __shfl_xor_sync(0xFFFFFFFFu, best, offset)); + } + winners[r] = best; +#endif +#pragma unroll + for (int j = 0; j < 16; ++j) { + keys[j] = keys[j] == winners[r] ? 0u : keys[j]; + } + } else { + winners[r] = 0u; + } + } + + float part[4]; +#pragma unroll + for (int w = 0; w < 4; ++w) { + float sum = 0.0f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + sum += expf(scores[w * 4 + j] - max_val); + } + part[w] = sum; + } +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { +#pragma unroll + for (int w = 0; w < 4; ++w) { + part[w] += __shfl_xor_sync(0xFFFFFFFFu, part[w], offset); + } + } + const float block_sum = ((__shfl_sync(0xFFFFFFFFu, part[0], 0) + __shfl_sync(0xFFFFFFFFu, part[1], 0)) + + __shfl_sync(0xFFFFFFFFu, part[2], 0)) + + __shfl_sync(0xFFFFFFFFu, part[3], 0); + const float inv_sum = onnxruntime::cuda::topk::SafeInvSum(block_sum); + + uint32_t selected = 0; + float prob = 0.0f; + float scale_sum = 0.0f; +#pragma unroll + for (int r = 0; r < kMoEPackedMaxK; ++r) { + if (r < k) { + const float p = onnxruntime::cuda::topk::SoftmaxScale(FusedRouteScore(winners[r]), max_val, inv_sum); + scale_sum += p; + if (lane == r) { + selected = winners[r]; + prob = p; + } + } + } + const float denom = onnxruntime::cuda::topk::TopKNormalizeDenom(normalize_scales, scale_sum); + if (lane < k) { + const int expert = 0x1FFF - static_cast(selected & 0x1FFFu); + topk_scales[static_cast(row) * k + lane] = normalize_scales ? (prob / denom) : prob; + topk_indices[static_cast(row) * k + lane] = expert; + } +} + +// Generic Softmax + TopK for num_experts != 512 (e.g. unit tests with smaller E). +template +__global__ __launch_bounds__(256) void MoEPackedSoftmaxTopKGenericKernel( + const CudaT* __restrict__ logits, + float* __restrict__ topk_scales, + int* __restrict__ topk_indices, + int num_rows, + int num_experts, + int k, + bool normalize_scales) { + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int row = static_cast(blockIdx.x) * 8 + warp; + if (row >= num_rows) return; + + const CudaT* row_logits = logits + static_cast(row) * num_experts; + float local_max = onnxruntime::cuda::topk::kNegativeInfinity; + for (int e = lane; e < num_experts; e += 32) { + const float v = gguf_cuda::CudaTypeToFloat(row_logits[e]); + local_max = fmaxf(local_max, v); + } +#pragma unroll + for (int offset = 16; offset >= 1; offset >>= 1) { + local_max = fmaxf(local_max, __shfl_xor_sync(0xFFFFFFFFu, local_max, offset)); + } + const float max_val = local_max; + + float local_sum = 0.0f; + for (int e = lane; e < num_experts; e += 32) { + const float v = gguf_cuda::CudaTypeToFloat(row_logits[e]); + local_sum += expf(v - max_val); + } +#pragma unroll + for (int offset = 16; offset >= 1; offset >>= 1) { + local_sum += __shfl_xor_sync(0xFFFFFFFFu, local_sum, offset); + } + const float inv_sum = onnxruntime::cuda::topk::SafeInvSum(local_sum); + + uint32_t winners[kMoEPackedMaxK]; + for (int r = 0; r < k; ++r) { + uint32_t best = 0u; + for (int e = lane; e < num_experts; e += 32) { + const float v = gguf_cuda::CudaTypeToFloat(row_logits[e]); + uint32_t key = FusedRouteKey(v, e, true); + for (int prev = 0; prev < r; ++prev) { + if (key == winners[prev]) key = 0u; + } + best = max(best, key); + } +#pragma unroll + for (int offset = 16; offset >= 1; offset >>= 1) { + best = max(best, __shfl_xor_sync(0xFFFFFFFFu, best, offset)); + } + winners[r] = best; + } + + float scale_sum = 0.0f; + uint32_t selected = 0u; + float prob = 0.0f; + for (int r = 0; r < k; ++r) { + const float p = onnxruntime::cuda::topk::SoftmaxScale(FusedRouteScore(winners[r]), max_val, inv_sum); + scale_sum += p; + if (lane == r) { + selected = winners[r]; + prob = p; + } + } + const float denom = onnxruntime::cuda::topk::TopKNormalizeDenom(normalize_scales, scale_sum); + if (lane < k) { + const int expert = 0x1FFF - static_cast(selected & 0x1FFFu); + topk_scales[static_cast(row) * k + lane] = normalize_scales ? (prob / denom) : prob; + topk_indices[static_cast(row) * k + lane] = expert; + } +} + +// Parallel single-CTA bucket & tile builder: groups total_routes (route_id = 0..total_routes-1) +// by expert (0..num_experts-1) into tiles of up to kTileRoutes routes sharing the same expert. +template +__global__ __launch_bounds__(512) void BuildMoEPackedTilesKernel( + const int* __restrict__ topk_indices, + int total_routes, + int num_experts, + int* __restrict__ sorted_routes, + MoEPackedTile* __restrict__ tiles, + int* __restrict__ d_num_tiles) { + __shared__ int s_counts[kMoEPackedMaxExperts]; + __shared__ int s_cursor[kMoEPackedMaxExperts]; + __shared__ int s_warp_routes[16]; + __shared__ int s_warp_tiles[16]; + + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + const int lane = tid & 31; + + s_counts[tid] = 0; + __syncthreads(); + + for (int r = tid; r < total_routes; r += 512) { + const int e = topk_indices[r]; + if (static_cast(e) < static_cast(num_experts)) { + atomicAdd(&s_counts[e], 1); + } + } + __syncthreads(); + + const int c = (tid < num_experts) ? s_counts[tid] : 0; + const int n_tiles = (c + kTileRoutes - 1) / kTileRoutes; + + int r_scan = c; + int t_scan = n_tiles; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const int r_up = __shfl_up_sync(0xFFFFFFFFu, r_scan, offset); + const int t_up = __shfl_up_sync(0xFFFFFFFFu, t_scan, offset); + if (lane >= offset) { + r_scan += r_up; + t_scan += t_up; + } + } + if (lane == 31) { + s_warp_routes[warp] = r_scan; + s_warp_tiles[warp] = t_scan; + } + __syncthreads(); + + if (warp == 0) { + int wr = (lane < 16) ? s_warp_routes[lane] : 0; + int wt = (lane < 16) ? s_warp_tiles[lane] : 0; +#pragma unroll + for (int offset = 1; offset < 16; offset <<= 1) { + const int wr_up = __shfl_up_sync(0xFFFFFFFFu, wr, offset); + const int wt_up = __shfl_up_sync(0xFFFFFFFFu, wt, offset); + if (lane >= offset) { + wr += wr_up; + wt += wt_up; + } + } + if (lane < 16) { + s_warp_routes[lane] = wr; + s_warp_tiles[lane] = wt; + } + if (lane == 15) { + *d_num_tiles = wt; + } + } + __syncthreads(); + + const int warp_r_base = (warp > 0) ? s_warp_routes[warp - 1] : 0; + const int warp_t_base = (warp > 0) ? s_warp_tiles[warp - 1] : 0; + const int exc_r = warp_r_base + r_scan - c; + const int exc_t = warp_t_base + t_scan - n_tiles; + if (tid < num_experts) { + s_cursor[tid] = exc_r; + } + __syncthreads(); + + for (int r = tid; r < total_routes; r += 512) { + const int e = topk_indices[r]; + if (static_cast(e) < static_cast(num_experts)) { + const int pos = atomicAdd(&s_cursor[e], 1); + sorted_routes[pos] = r; + } + } + + if (tid < num_experts && n_tiles > 0) { + for (int t = 0; t < n_tiles; ++t) { + MoEPackedTile tile{}; + tile.expert = tid; + tile.start_route = exc_r + t * kTileRoutes; + const int rem = c - t * kTileRoutes; + tile.count = rem < kTileRoutes ? rem : kTileRoutes; + tile.pad = 0; + tiles[exc_t + t] = tile; + } + } +} + +// Decode Warp-K-Split Tensor-Core BF16 FC1 + SwiGLU kernel (inspired by bolt_ops.cu MmaGemvKernel): +// KW=4 warps (128 threads) per CTA split K across the 4 warps (only 20 iterations/warp for K=2560, +// zero __syncthreads() in the K-loop) while computing MT*8 = 16 inter_size channels (32 gate+up rows). +template +__global__ __launch_bounds__(KW * 32) void MoEPackedFc1DecodeWarpKSplitKernel( + const __nv_bfloat16* __restrict__ input, + const uint8_t* __restrict__ gate_weights, + const uint8_t* __restrict__ up_weights, + const int* __restrict__ sorted_routes, + const MoEPackedTile* __restrict__ tiles, + const int* __restrict__ d_num_tiles, + MoEPackedParams p, + __nv_bfloat16* __restrict__ fc1_out) { + const int tile_idx = static_cast(blockIdx.y); + if (tile_idx >= __ldg(d_num_tiles)) return; + + constexpr int kThreads = KW * 32; + using namespace onnxruntime::contrib::gguf_quant; + constexpr int kLutU64Count = (kGateUpType == kGgmlTypeIQ2_S) + ? 1024 + : ((kGateUpType == kGgmlTypeIQ2_XS) + ? 512 + : ((kGateUpType == kGgmlTypeIQ2_XXS) + ? 256 + : ((kGateUpType == kGgmlTypeIQ3_XXS) ? 128 : 1))); + + __shared__ uint64_t s_lut[kLutU64Count]; + __shared__ float s_red[KW][MT * 16][8]; + __shared__ int s_routes[8]; + __shared__ int s_token_rows[8]; + + const MoEPackedTile tile = tiles[tile_idx]; + const int count = tile.count; + const int expert = tile.expert; + + const int tid = static_cast(threadIdx.x); + if (tid < 8) { + const int r = (tid < count) ? __ldg(&sorted_routes[tile.start_route + tid]) : 0; + s_routes[tid] = r; + s_token_rows[tid] = r / p.k; + } + if constexpr (kGateUpType == kGgmlTypeIQ2_XXS) { + for (int i = tid; i < 256; i += kThreads) s_lut[i] = __ldg(iq2xxs_grid + i); + } else if constexpr (kGateUpType == kGgmlTypeIQ2_XS) { + for (int i = tid; i < 512; i += kThreads) s_lut[i] = __ldg(iq2xs_grid + i); + } else if constexpr (kGateUpType == kGgmlTypeIQ2_S) { + for (int i = tid; i < 1024; i += kThreads) s_lut[i] = __ldg(iq2s_grid + i); + } else if constexpr (kGateUpType == kGgmlTypeIQ3_XXS) { + uint32_t* s_lut32 = reinterpret_cast(s_lut); + for (int i = tid; i < 256; i += kThreads) s_lut32[i] = __ldg(iq3xxs_grid + i); + } + __syncthreads(); + + const int warp = tid >> 5; + const int lane = tid & 31; + const int g = lane >> 2; // 0..7 + const int t = lane & 3; // 0..3 + const int cta_ch_base = static_cast(blockIdx.x) * (MT * 8); + + const uint8_t* gate_ptrs[MT]; + const uint8_t* up_ptrs[MT]; + bool valid_rows[MT]; + float acc[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + const int j_out = cta_ch_base + m * 8 + g; + valid_rows[m] = (j_out < p.inter_size); + const int64_t exp_row_idx = + static_cast(expert) * p.inter_size + (valid_rows[m] ? j_out : 0); + gate_ptrs[m] = gate_weights + exp_row_idx * p.gate_row_bytes; + up_ptrs[m] = up_weights + exp_row_idx * p.up_row_bytes; +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc[m][i] = 0.0f; + } + } + + const bool x_valid = (g < count); + const __nv_bfloat16* x_row = + input + static_cast(x_valid ? s_token_rows[g] : 0) * p.hidden_size; + + const int num_sb32 = p.hidden_size >> 5; + for (int sb32_idx = warp; sb32_idx < num_sb32; sb32_idx += KW) { + const uint4 b_vec = x_valid + ? __ldg(reinterpret_cast(x_row + sb32_idx * 32 + t * 8)) + : make_uint4(0u, 0u, 0u, 0u); +#pragma unroll + for (int m = 0; m < MT; ++m) { + if (valid_rows[m]) { + float g_vals[8]; + float u_vals[8]; + gguf_cuda::DecodeGgufSubBlock8SmemLutTpl( + gate_ptrs[m], p.gate_type, sb32_idx, t, s_lut, g_vals); + gguf_cuda::DecodeGgufSubBlock8SmemLutTpl( + up_ptrs[m], p.up_type, sb32_idx, t, s_lut, u_vals); + const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(g_vals); + const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(u_vals); + gguf_cuda::MmaBf16(acc[m], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); + gguf_cuda::MmaBf16(acc[m], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); + } + } + } + +#pragma unroll + for (int m = 0; m < MT; ++m) { + s_red[warp][m * 16 + g][2 * t + 0] = acc[m][0]; + s_red[warp][m * 16 + g][2 * t + 1] = acc[m][1]; + s_red[warp][m * 16 + g + 8][2 * t + 0] = acc[m][2]; + s_red[warp][m * 16 + g + 8][2 * t + 1] = acc[m][3]; + } + __syncthreads(); + + for (int idx = tid; idx < MT * 8 * 8; idx += kThreads) { + const int col = idx / (MT * 8); + const int ch = idx - col * (MT * 8); + const int j_out = cta_ch_base + ch; + if (col < count && j_out < p.inter_size) { + const int m = ch >> 3; + const int g_ch = ch & 7; + float g_raw = 0.0f; + float u_raw = 0.0f; +#pragma unroll + for (int w = 0; w < KW; ++w) { + g_raw += s_red[w][m * 16 + g_ch][col]; + u_raw += s_red[w][m * 16 + g_ch + 8][col]; + } + const float g_val = __bfloat162float(__float2bfloat16_rn(g_raw)); + const float u_val = __bfloat162float(__float2bfloat16_rn(u_raw)); + const float swiglu = u_val * (g_val / (1.0f + expf(-g_val))); + fc1_out[static_cast(s_routes[col]) * p.inter_size + j_out] = + __float2bfloat16_rn(swiglu); + } + } +} + +// Tensor-Core BF16 FC1 + SwiGLU kernel (mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32). +// Each warp (32 threads) computes 8 inter_size channels (8 gate rows + 8 up rows = 16 weight rows) +// across up to kMaxSubTiles * 8 routes (8 routes in decode, 32 routes in prefill). +// 8 warps per CTA = 64 inter_size channels per CTA. +template +__global__ __launch_bounds__((kMaxSubTiles == 1) ? 128 : 256) void MoEPackedFc1SwiGLUBf16TcKernel( + const __nv_bfloat16* __restrict__ input, + const uint8_t* __restrict__ gate_weights, + const uint8_t* __restrict__ up_weights, + const int* __restrict__ sorted_routes, + const MoEPackedTile* __restrict__ tiles, + const int* __restrict__ d_num_tiles, + MoEPackedParams p, + __nv_bfloat16* __restrict__ fc1_out) { + const int tile_idx = static_cast(blockIdx.y); + if (tile_idx >= __ldg(d_num_tiles)) return; + + constexpr int kWarps = (kMaxSubTiles == 1) ? 4 : 8; + constexpr int kThreads = kWarps * 32; + constexpr int kChunkK = 512; + constexpr int kChunkU4 = kChunkK / 8; // 64 + constexpr int kRowStrideU4 = kChunkU4 + 4; // 68 (bank-conflict-free for LDS.128) + constexpr int kMaxRoutes = kMaxSubTiles * 8; // 8 or 32 + using namespace onnxruntime::contrib::gguf_quant; + constexpr int kLutU64Count = (kGateUpType == kGgmlTypeIQ2_S) + ? 1024 + : ((kGateUpType == kGgmlTypeIQ2_XS) + ? 512 + : ((kGateUpType == kGgmlTypeIQ2_XXS) + ? 256 + : ((kGateUpType == kGgmlTypeIQ3_XXS) ? 128 : 1))); + + __shared__ uint4 s_x_u4[kMaxRoutes * kRowStrideU4]; + __shared__ uint64_t s_lut[kLutU64Count]; + __shared__ int s_routes[kMaxRoutes]; + __shared__ int s_token_rows[kMaxRoutes]; + + const MoEPackedTile tile = tiles[tile_idx]; + const int count = tile.count; + const int expert = tile.expert; + const int num_subtiles = (count + 7) >> 3; + const int padded_routes = num_subtiles * 8; + + const int tid = static_cast(threadIdx.x); + if (tid < kMaxRoutes) { + const int r = (tid < count) ? __ldg(&sorted_routes[tile.start_route + tid]) : 0; + s_routes[tid] = r; + s_token_rows[tid] = r / p.k; + } + if constexpr (kGateUpType == kGgmlTypeIQ2_XXS) { + for (int i = tid; i < 256; i += kThreads) { + s_lut[i] = __ldg(iq2xxs_grid + i); + } + } else if constexpr (kGateUpType == kGgmlTypeIQ2_XS) { + for (int i = tid; i < 512; i += kThreads) { + s_lut[i] = __ldg(iq2xs_grid + i); + } + } else if constexpr (kGateUpType == kGgmlTypeIQ2_S) { + for (int i = tid; i < 1024; i += kThreads) { + s_lut[i] = __ldg(iq2s_grid + i); + } + } else if constexpr (kGateUpType == kGgmlTypeIQ3_XXS) { + uint32_t* s_lut32 = reinterpret_cast(s_lut); + for (int i = tid; i < 256; i += kThreads) { + s_lut32[i] = __ldg(iq3xxs_grid + i); + } + } + __syncthreads(); + + const int warp = tid >> 5; + const int lane = tid & 31; + const int g = lane >> 2; // 0..7 + const int t = lane & 3; // 0..3 + const int j_out = (static_cast(blockIdx.x) * kWarps + warp) * 8 + g; + const bool valid_row = (j_out < p.inter_size); + + const int64_t exp_row_idx = static_cast(expert) * p.inter_size + (valid_row ? j_out : 0); + const uint8_t* gate_row_ptr = gate_weights + exp_row_idx * p.gate_row_bytes; + const uint8_t* up_row_ptr = up_weights + exp_row_idx * p.up_row_bytes; + + float acc[kMaxSubTiles][4]; +#pragma unroll + for (int st = 0; st < kMaxSubTiles; ++st) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc[st][i] = 0.0f; + } + } + + for (int k_base = 0; k_base < p.hidden_size; k_base += kChunkK) { + const int rem_k = p.hidden_size - k_base; + const int actual_chunk_k = rem_k < kChunkK ? rem_k : kChunkK; + const int actual_chunk_u4 = actual_chunk_k >> 3; + const int total_u4 = padded_routes * actual_chunk_u4; + + if (actual_chunk_u4 == 64) { + for (int idx = tid; idx < total_u4; idx += kThreads) { + const int m = idx >> 6; + const int u4_col = idx & 63; + uint4 val = make_uint4(0u, 0u, 0u, 0u); + if (m < count) { + const int token_row = s_token_rows[m]; + const __nv_bfloat16* src_ptr = + input + static_cast(token_row) * p.hidden_size + k_base + (u4_col << 3); + val = __ldg(reinterpret_cast(src_ptr)); + } + s_x_u4[m * kRowStrideU4 + u4_col] = val; + } + } else { + for (int idx = tid; idx < total_u4; idx += kThreads) { + const int m = idx / actual_chunk_u4; + const int u4_col = idx - m * actual_chunk_u4; + uint4 val = make_uint4(0u, 0u, 0u, 0u); + if (m < count) { + const int token_row = s_token_rows[m]; + const __nv_bfloat16* src_ptr = + input + static_cast(token_row) * p.hidden_size + k_base + (u4_col << 3); + val = __ldg(reinterpret_cast(src_ptr)); + } + s_x_u4[m * kRowStrideU4 + u4_col] = val; + } + } + __syncthreads(); + + if (valid_row) { + const int sb32_base = k_base >> 5; + const int num_sb32 = actual_chunk_k >> 5; + for (int sb_local = 0; sb_local < num_sb32; ++sb_local) { + const int sb32_idx = sb32_base + sb_local; + float g_vals[8]; + float u_vals[8]; + gguf_cuda::DecodeGgufSubBlock8SmemLutTpl(gate_row_ptr, p.gate_type, sb32_idx, t, s_lut, g_vals); + gguf_cuda::DecodeGgufSubBlock8SmemLutTpl(up_row_ptr, p.up_type, sb32_idx, t, s_lut, u_vals); + const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(g_vals); + const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(u_vals); +#pragma unroll + for (int st = 0; st < kMaxSubTiles; ++st) { + if (st < num_subtiles) { + const uint4 b_vec = s_x_u4[(st * 8 + g) * kRowStrideU4 + sb_local * 4 + t]; + gguf_cuda::MmaBf16(acc[st], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); + gguf_cuda::MmaBf16(acc[st], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); + } + } + } + } + if (k_base + kChunkK < p.hidden_size) { + __syncthreads(); + } + } + + if (valid_row) { +#pragma unroll + for (int st = 0; st < kMaxSubTiles; ++st) { + if (st < num_subtiles) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int m = st * 8 + 2 * t + h; + if (m < count) { + const float g_raw = (h == 0) ? acc[st][0] : acc[st][1]; + const float u_raw = (h == 0) ? acc[st][2] : acc[st][3]; + const float g_val = __bfloat162float(__float2bfloat16_rn(g_raw)); + const float u_val = __bfloat162float(__float2bfloat16_rn(u_raw)); + const float swiglu = u_val * (g_val / (1.0f + expf(-g_val))); + const int route = s_routes[m]; + fc1_out[static_cast(route) * p.inter_size + j_out] = __float2bfloat16_rn(swiglu); + } + } + } + } + } +} + +// Tensor-Core BF16 FC2 kernel (mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32). +// Each warp (32 threads) computes 16 hidden_size channels across up to kMaxSubTiles * 8 routes. +// 8 warps per CTA = 128 hidden_size channels per CTA. +template +__global__ __launch_bounds__(256) void MoEPackedFc2Bf16TcKernel( + const __nv_bfloat16* __restrict__ fc1_out, + const uint8_t* __restrict__ down_weights, + const int* __restrict__ sorted_routes, + const MoEPackedTile* __restrict__ tiles, + const int* __restrict__ d_num_tiles, + MoEPackedParams p, + __nv_bfloat16* __restrict__ fc2_out) { + const int tile_idx = static_cast(blockIdx.y); + if (tile_idx >= __ldg(d_num_tiles)) return; + + constexpr int kChunkK = 640; + constexpr int kChunkU4 = kChunkK / 8; // 80 + constexpr int kRowStrideU4 = kChunkU4 + 4; // 84 (bank-conflict-free for LDS.128) + constexpr int kMaxRoutes = kMaxSubTiles * 8; // 8 or 32 + + __shared__ uint4 s_act_u4[kMaxRoutes * kRowStrideU4]; + __shared__ int s_routes[kMaxRoutes]; + + const MoEPackedTile tile = tiles[tile_idx]; + const int count = tile.count; + const int expert = tile.expert; + const int num_subtiles = (count + 7) >> 3; + const int padded_routes = num_subtiles * 8; + + const int tid = static_cast(threadIdx.x); + if (tid < kMaxRoutes) { + s_routes[tid] = (tid < count) ? __ldg(&sorted_routes[tile.start_route + tid]) : 0; + } + __syncthreads(); + + const int warp = tid >> 5; + const int lane = tid & 31; + const int g = lane >> 2; // 0..7 + const int t = lane & 3; // 0..3 + const int n_out_base = (static_cast(blockIdx.x) * 8 + warp) * 16; + const int r_lo = n_out_base + 2 * g; + const int r_hi = r_lo + 1; + const bool valid_lo = (r_lo < p.hidden_size); + const bool valid_hi = (r_hi < p.hidden_size); + + const int64_t exp_base = static_cast(expert) * p.hidden_size; + const uint8_t* down_row_lo = down_weights + (exp_base + (valid_lo ? r_lo : 0)) * p.down_row_bytes; + const uint8_t* down_row_hi = down_weights + (exp_base + (valid_hi ? r_hi : 0)) * p.down_row_bytes; + + float acc[kMaxSubTiles][4]; +#pragma unroll + for (int st = 0; st < kMaxSubTiles; ++st) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc[st][i] = 0.0f; + } + } + + for (int k_base = 0; k_base < p.inter_size; k_base += kChunkK) { + const int rem_k = p.inter_size - k_base; + const int actual_chunk_k = rem_k < kChunkK ? rem_k : kChunkK; + const int actual_chunk_u4 = actual_chunk_k >> 3; + const int total_u4 = padded_routes * actual_chunk_u4; + + for (int idx = tid; idx < total_u4; idx += 256) { + const int m = idx / actual_chunk_u4; + const int u4_col = idx - m * actual_chunk_u4; + uint4 val = make_uint4(0u, 0u, 0u, 0u); + if (m < count) { + const int route = s_routes[m]; + const __nv_bfloat16* src_ptr = + fc1_out + static_cast(route) * p.inter_size + k_base + (u4_col << 3); + val = __ldg(reinterpret_cast(src_ptr)); + } + s_act_u4[m * kRowStrideU4 + u4_col] = val; + } + __syncthreads(); + + if (valid_lo) { + const int sb32_base = k_base >> 5; + const int num_sb32 = actual_chunk_k >> 5; + for (int sb_local = 0; sb_local < num_sb32; ++sb_local) { + const int sb32_idx = sb32_base + sb_local; + float d_lo[8]; + float d_hi[8]; + gguf_cuda::DecodeGgufSubBlock8Tpl(down_row_lo, p.down_type, sb32_idx, t, d_lo); + if (valid_hi) { + gguf_cuda::DecodeGgufSubBlock8Tpl(down_row_hi, p.down_type, sb32_idx, t, d_hi); + } else { +#pragma unroll + for (int j = 0; j < 8; ++j) d_hi[j] = 0.0f; + } + const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(d_lo); + const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(d_hi); +#pragma unroll + for (int st = 0; st < kMaxSubTiles; ++st) { + if (st < num_subtiles) { + const uint4 b_vec = s_act_u4[(st * 8 + g) * kRowStrideU4 + sb_local * 4 + t]; + gguf_cuda::MmaBf16(acc[st], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); + gguf_cuda::MmaBf16(acc[st], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); + } + } + } + } + if (k_base + kChunkK < p.inter_size) { + __syncthreads(); + } + } + + if (valid_lo) { +#pragma unroll + for (int st = 0; st < kMaxSubTiles; ++st) { + if (st < num_subtiles) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int m = st * 8 + 2 * t + h; + if (m < count) { + const int route = s_routes[m]; + const float v_lo = (h == 0) ? acc[st][0] : acc[st][1]; + const float v_hi = (h == 0) ? acc[st][2] : acc[st][3]; + __nv_bfloat16* dst = fc2_out + static_cast(route) * p.hidden_size + r_lo; + if (valid_hi && ((p.hidden_size & 1) == 0)) { + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v_lo, v_hi); + } else { + dst[0] = __float2bfloat16_rn(v_lo); + if (valid_hi) dst[1] = __float2bfloat16_rn(v_hi); + } + } + } + } + } + } +} + +// Fallback SIMT FC1 + SwiGLU kernel (used for FP16). +template +__global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel( + const CudaT* __restrict__ input, + const uint8_t* __restrict__ gate_weights, + const uint8_t* __restrict__ up_weights, + const int* __restrict__ sorted_routes, + const MoEPackedTile* __restrict__ tiles, + const int* __restrict__ d_num_tiles, + MoEPackedParams p, + CudaT* __restrict__ fc1_out) { + const int tile_idx = static_cast(blockIdx.y); + if (tile_idx >= *d_num_tiles) return; + + extern __shared__ char smem_raw[]; + CudaT* s_x = reinterpret_cast(smem_raw); + + const MoEPackedTile tile = tiles[tile_idx]; + const int count = tile.count < 8 ? tile.count : 8; + const int expert = tile.expert; + + const int total_elems = count * p.hidden_size; + for (int idx = static_cast(threadIdx.x); idx < total_elems; idx += 256) { + const int m = idx / p.hidden_size; + const int col = idx - m * p.hidden_size; + const int token_row = sorted_routes[tile.start_route + m] / p.k; + s_x[idx] = input[static_cast(token_row) * p.hidden_size + col]; + } + __syncthreads(); + + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int j_out = static_cast(blockIdx.x) * 8 + warp; + if (j_out >= p.inter_size) return; + + const int64_t exp_row_idx = static_cast(expert) * p.inter_size + j_out; + const uint8_t* gate_row_ptr = gate_weights + exp_row_idx * p.gate_row_bytes; + const uint8_t* up_row_ptr = up_weights + exp_row_idx * p.up_row_bytes; + + float gate_acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float up_acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + + const int total_halves = p.hidden_size / 16; + for (int h_idx = lane; h_idx < total_halves; h_idx += 32) { + const int sb32_idx = h_idx >> 1; + const int half16 = h_idx & 1; + float g_vals[16]; + float u_vals[16]; + gguf_cuda::DecodeGgufSubBlock32(gate_row_ptr, p.gate_type, sb32_idx, half16, g_vals); + gguf_cuda::DecodeGgufSubBlock32(up_row_ptr, p.up_type, sb32_idx, half16, u_vals); + + const int k_offset = h_idx * 16; + for (int m = 0; m < count; ++m) { + const CudaT* xm = s_x + m * p.hidden_size + k_offset; +#pragma unroll + for (int c = 0; c < 16; ++c) { + const float xv = gguf_cuda::CudaTypeToFloat(xm[c]); + gate_acc[m] += g_vals[c] * xv; + up_acc[m] += u_vals[c] * xv; + } + } + } + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int m = 0; m < 8; ++m) { + gate_acc[m] += __shfl_down_sync(0xFFFFFFFFu, gate_acc[m], offset); + up_acc[m] += __shfl_down_sync(0xFFFFFFFFu, up_acc[m], offset); + } + } + + if (lane == 0) { + for (int m = 0; m < count; ++m) { + const float g = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(gate_acc[m])); + const float u = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(up_acc[m])); + const float swiglu = u * (g / (1.0f + expf(-g))); + fc1_out[static_cast(sorted_routes[tile.start_route + m]) * p.inter_size + j_out] = + gguf_cuda::FloatToCudaType(swiglu); + } + } +} + +// Fallback SIMT FC2 kernel (used for FP16). +template +__global__ __launch_bounds__(256) void MoEPackedFc2SimtKernel( + const CudaT* __restrict__ fc1_out, + const uint8_t* __restrict__ down_weights, + const int* __restrict__ sorted_routes, + const MoEPackedTile* __restrict__ tiles, + const int* __restrict__ d_num_tiles, + MoEPackedParams p, + CudaT* __restrict__ fc2_out) { + const int tile_idx = static_cast(blockIdx.y); + if (tile_idx >= *d_num_tiles) return; + + extern __shared__ char smem_raw[]; + CudaT* s_act = reinterpret_cast(smem_raw); + + const MoEPackedTile tile = tiles[tile_idx]; + const int count = tile.count < 8 ? tile.count : 8; + const int expert = tile.expert; + + const int total_elems = count * p.inter_size; + for (int idx = static_cast(threadIdx.x); idx < total_elems; idx += 256) { + const int m = idx / p.inter_size; + const int col = idx - m * p.inter_size; + s_act[idx] = fc1_out[static_cast(sorted_routes[tile.start_route + m]) * p.inter_size + col]; + } + __syncthreads(); + + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int n_out = static_cast(blockIdx.x) * 8 + warp; + if (n_out >= p.hidden_size) return; + + const int64_t exp_row_idx = static_cast(expert) * p.hidden_size + n_out; + const uint8_t* down_row_ptr = down_weights + exp_row_idx * p.down_row_bytes; + + float down_acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + + const int total_halves = p.inter_size / 16; + for (int h_idx = lane; h_idx < total_halves; h_idx += 32) { + const int sb32_idx = h_idx >> 1; + const int half16 = h_idx & 1; + float d_vals[16]; + gguf_cuda::DecodeGgufSubBlock32(down_row_ptr, p.down_type, sb32_idx, half16, d_vals); + + const int k_offset = h_idx * 16; + for (int m = 0; m < count; ++m) { + const CudaT* am = s_act + m * p.inter_size + k_offset; +#pragma unroll + for (int c = 0; c < 16; ++c) { + down_acc[m] += d_vals[c] * gguf_cuda::CudaTypeToFloat(am[c]); + } + } + } + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int m = 0; m < 8; ++m) { + down_acc[m] += __shfl_down_sync(0xFFFFFFFFu, down_acc[m], offset); + } + } + + if (lane == 0) { + for (int m = 0; m < count; ++m) { + fc2_out[static_cast(sorted_routes[tile.start_route + m]) * p.hidden_size + n_out] = + gguf_cuda::FloatToCudaType(down_acc[m]); + } + } +} + +// Vectorized 128-bit Finalize: accumulates the k routed expert outputs per token in rank order 0..k-1. +template +__global__ __launch_bounds__(256) void MoEPackedFinalizeVec8Kernel( + const CudaT* __restrict__ fc2_out, + const float* __restrict__ topk_scales, + int num_rows, + int hidden_size, + int k, + CudaT* __restrict__ output) { + const int row = static_cast(blockIdx.y); + const int col8 = (static_cast(blockIdx.x) * 256 + static_cast(threadIdx.x)) * 8; + if (row >= num_rows || col8 >= hidden_size) return; + + const float* scales = topk_scales + static_cast(row) * k; + const CudaT* row_fc2 = fc2_out + static_cast(row) * k * hidden_size + col8; + float sum[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + for (int r = 0; r < k; ++r) { + const float sc = __ldg(scales + r); + const uint4 v_u4 = __ldg(reinterpret_cast(row_fc2 + static_cast(r) * hidden_size)); + const CudaT* vals = reinterpret_cast(&v_u4); +#pragma unroll + for (int j = 0; j < 8; ++j) { + sum[j] += sc * gguf_cuda::CudaTypeToFloat(vals[j]); + } + } + uint4 out_u4; + CudaT* out_vals = reinterpret_cast(&out_u4); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out_vals[j] = gguf_cuda::FloatToCudaType(sum[j]); + } + *reinterpret_cast(output + static_cast(row) * hidden_size + col8) = out_u4; +} + +template +__global__ __launch_bounds__(256) void MoEPackedFinalizeKernel( + const CudaT* __restrict__ fc2_out, + const float* __restrict__ topk_scales, + int num_rows, + int hidden_size, + int k, + CudaT* __restrict__ output) { + const int row = static_cast(blockIdx.y); + const int col = static_cast(blockIdx.x) * 256 + static_cast(threadIdx.x); + if (row >= num_rows || col >= hidden_size) return; + + const float* scales = topk_scales + static_cast(row) * k; + const CudaT* row_fc2 = fc2_out + static_cast(row) * k * hidden_size + col; + float sum = 0.0f; + for (int r = 0; r < k; ++r) { + sum += scales[r] * gguf_cuda::CudaTypeToFloat(row_fc2[static_cast(r) * hidden_size]); + } + output[static_cast(row) * hidden_size + col] = gguf_cuda::FloatToCudaType(sum); +} + +size_t AlignUpWorkspace(size_t v) { + return (v + 255u) & ~size_t{255u}; +} + +void LaunchMoEPackedFc1DecodeBf16TcDispatch( + cudaStream_t stream, + dim3 grid, + const __nv_bfloat16* input, + const uint8_t* gate_weights, + const uint8_t* up_weights, + const int* sorted_routes, + const MoEPackedTile* tiles, + const int* d_num_tiles, + const MoEPackedParams& params, + __nv_bfloat16* fc1_out) { + using namespace onnxruntime::contrib::gguf_quant; + constexpr int KW = 4; + constexpr int MT = 2; + constexpr unsigned int kBlock = KW * 32; + if (params.gate_type == params.up_type) { + switch (params.gate_type) { + case kGgmlTypeIQ2_XXS: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ2_XS: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ2_S: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ3_XXS: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ4_XS: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ4_NL: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeQ2_0: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeQ5_0: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeQ8_0: + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + default: + break; + } + } + MoEPackedFc1DecodeWarpKSplitKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); +} + +template +void LaunchMoEPackedFc1Bf16TcDispatch( + cudaStream_t stream, + dim3 grid, + const __nv_bfloat16* input, + const uint8_t* gate_weights, + const uint8_t* up_weights, + const int* sorted_routes, + const MoEPackedTile* tiles, + const int* d_num_tiles, + const MoEPackedParams& params, + __nv_bfloat16* fc1_out) { + using namespace onnxruntime::contrib::gguf_quant; + constexpr unsigned int kBlock = (kMaxSubTiles == 1) ? 128u : 256u; + if (params.gate_type == params.up_type) { + switch (params.gate_type) { + case kGgmlTypeIQ2_XXS: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ2_XS: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ2_S: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ3_XXS: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ4_XS: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeIQ4_NL: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeQ2_0: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeQ5_0: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + case kGgmlTypeQ8_0: + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + return; + default: + break; + } + } + MoEPackedFc1SwiGLUBf16TcKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); +} + +template +void LaunchMoEPackedFc2Bf16TcDispatch( + cudaStream_t stream, + dim3 grid, + const __nv_bfloat16* fc1_out, + const uint8_t* down_weights, + const int* sorted_routes, + const MoEPackedTile* tiles, + const int* d_num_tiles, + const MoEPackedParams& params, + __nv_bfloat16* fc2_out) { + using namespace onnxruntime::contrib::gguf_quant; + switch (params.down_type) { + case kGgmlTypeQ2_0: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeIQ4_NL: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeQ5_0: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeIQ2_XXS: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeIQ2_XS: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeIQ2_S: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeIQ3_XXS: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeIQ4_XS: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + case kGgmlTypeQ8_0: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + default: + MoEPackedFc2Bf16TcKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + return; + } +} + +template +Status DispatchMoEPackedCuda( + cudaStream_t stream, + const CudaT* input, + const CudaT* router_probs, + const uint8_t* gate_weights, + const uint8_t* up_weights, + const uint8_t* down_weights, + const MoEPackedParams& params, + void* workspace, + CudaT* output) { + if (params.num_rows == 0) return Status::OK(); + const int total_routes = params.num_rows * params.k; + constexpr bool kIsBf16 = std::is_same::value; + const bool use_prefill_tiles = kIsBf16 && (total_routes > 128); + const int max_tiles = use_prefill_tiles + ? ((total_routes + 31) / 32 + params.num_experts) + : ((total_routes <= 128) ? total_routes : ((total_routes + 7) / 8 + params.num_experts)); + + uint8_t* ws_ptr = static_cast(workspace); + float* topk_scales = reinterpret_cast(ws_ptr); + ws_ptr += AlignUpWorkspace(static_cast(total_routes) * sizeof(float)); + + int* topk_indices = reinterpret_cast(ws_ptr); + ws_ptr += AlignUpWorkspace(static_cast(total_routes) * sizeof(int)); + + int* sorted_routes = reinterpret_cast(ws_ptr); + ws_ptr += AlignUpWorkspace(static_cast(total_routes) * sizeof(int)); + + MoEPackedTile* tiles = reinterpret_cast(ws_ptr); + ws_ptr += AlignUpWorkspace(static_cast(max_tiles) * sizeof(MoEPackedTile)); + + int* d_num_tiles = reinterpret_cast(ws_ptr); + ws_ptr += AlignUpWorkspace(sizeof(int)); + + CudaT* fc1_out = reinterpret_cast(ws_ptr); + ws_ptr += AlignUpWorkspace(static_cast(total_routes) * params.inter_size * sizeof(CudaT)); + + CudaT* fc2_out = reinterpret_cast(ws_ptr); + + // 1. Softmax + TopK + const dim3 route_grid(static_cast((params.num_rows + 7) / 8)); + if (params.num_experts == 512) { + MoEPackedSoftmaxTopK512Kernel<<>>( + router_probs, topk_scales, topk_indices, params.num_rows, params.k, params.normalize_routing_weights); + } else { + MoEPackedSoftmaxTopKGenericKernel<<>>( + router_probs, topk_scales, topk_indices, params.num_rows, params.num_experts, params.k, + params.normalize_routing_weights); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + + // 2. Build expert-grouped tiles (8 routes/tile in decode, 32 routes/tile in prefill) + if (use_prefill_tiles) { + BuildMoEPackedTilesKernel<32><<<1, 512, 0, stream>>>( + topk_indices, total_routes, params.num_experts, sorted_routes, tiles, d_num_tiles); + } else { + BuildMoEPackedTilesKernel<8><<<1, 512, 0, stream>>>( + topk_indices, total_routes, params.num_experts, sorted_routes, tiles, d_num_tiles); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + + if constexpr (kIsBf16) { + // 3. Tensor-Core FC1 (64 inter_size channels/CTA in prefill, 16 inter_size channels/CTA Warp-K-Split in decode) + if (use_prefill_tiles) { + const dim3 fc1_grid(static_cast((params.inter_size + 63) / 64), + static_cast(max_tiles)); + LaunchMoEPackedFc1Bf16TcDispatch<4>( + stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + } else { + const dim3 fc1_grid(static_cast((params.inter_size + 15) / 16), + static_cast(max_tiles)); + LaunchMoEPackedFc1DecodeBf16TcDispatch( + stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + + // 4. Tensor-Core FC2 (128 hidden_size channels per CTA) + const dim3 fc2_grid(static_cast((params.hidden_size + 127) / 128), + static_cast(max_tiles)); + if (use_prefill_tiles) { + LaunchMoEPackedFc2Bf16TcDispatch<4>( + stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + } else { + LaunchMoEPackedFc2Bf16TcDispatch<1>( + stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + } else { + const size_t fc1_smem = static_cast(8) * params.hidden_size * sizeof(CudaT); + const dim3 fc1_grid(static_cast((params.inter_size + 7) / 8), static_cast(max_tiles)); + MoEPackedFc1SwiGLUSimtKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + + const size_t fc2_smem = static_cast(8) * params.inter_size * sizeof(CudaT); + const dim3 fc2_grid(static_cast((params.hidden_size + 7) / 8), static_cast(max_tiles)); + MoEPackedFc2SimtKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + } + + // 5. Finalize (128-bit vectorized when hidden_size % 8 == 0) + if ((params.hidden_size & 7) == 0) { + const int vec_cols = params.hidden_size >> 3; + const dim3 fin_grid(static_cast((vec_cols + 255) / 256), + static_cast(params.num_rows)); + MoEPackedFinalizeVec8Kernel<<>>( + fc2_out, topk_scales, params.num_rows, params.hidden_size, params.k, output); + } else { + const dim3 fin_grid(static_cast((params.hidden_size + 255) / 256), + static_cast(params.num_rows)); + MoEPackedFinalizeKernel<<>>( + fc2_out, topk_scales, params.num_rows, params.hidden_size, params.k, output); + } + return CUDA_CALL(cudaGetLastError()); +} + +} // namespace + +size_t GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_bytes) { + const size_t total_routes = static_cast(params.num_rows) * params.k; + const size_t max_tiles = + (total_routes <= 128) ? total_routes : ((total_routes + 7) / 8 + static_cast(params.num_experts)); + size_t bytes = 0; + bytes += AlignUpWorkspace(total_routes * sizeof(float)); + bytes += AlignUpWorkspace(total_routes * sizeof(int)); + bytes += AlignUpWorkspace(total_routes * sizeof(int)); + bytes += AlignUpWorkspace(max_tiles * sizeof(MoEPackedTile)); + bytes += AlignUpWorkspace(sizeof(int)); + bytes += AlignUpWorkspace(total_routes * static_cast(params.inter_size) * elem_bytes); + bytes += AlignUpWorkspace(total_routes * static_cast(params.hidden_size) * elem_bytes); + return bytes; +} + +Status LaunchMoEPackedCuda( + cudaStream_t stream, + bool is_bf16, + const void* input, + const void* router_probs, + const uint8_t* gate_weights, + const uint8_t* up_weights, + const uint8_t* down_weights, + const MoEPackedParams& params, + void* workspace, + void* output) { + if (is_bf16) { + return DispatchMoEPackedCuda<__nv_bfloat16>( + stream, static_cast(input), static_cast(router_probs), + gate_weights, up_weights, down_weights, params, workspace, static_cast<__nv_bfloat16*>(output)); + } + return DispatchMoEPackedCuda( + stream, static_cast(input), static_cast(router_probs), + gate_weights, up_weights, down_weights, params, workspace, static_cast(output)); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.h b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h new file mode 100644 index 0000000000000..6685651a2e84f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +struct alignas(16) MoEPackedTile { + int expert; + int start_route; + int count; + int pad; +}; + +struct MoEPackedParams { + int num_rows{0}; + int num_experts{0}; + int hidden_size{0}; + int inter_size{0}; + int k{0}; + bool normalize_routing_weights{true}; + int gate_type{0}; + int up_type{0}; + int down_type{0}; + int64_t gate_row_bytes{0}; + int64_t up_row_bytes{0}; + int64_t down_row_bytes{0}; +}; + +size_t GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_bytes); + +Status LaunchMoEPackedCuda( + cudaStream_t stream, + bool is_bf16, + const void* input, + const void* router_probs, + const uint8_t* gate_weights, + const uint8_t* up_weights, + const uint8_t* down_weights, + const MoEPackedParams& params, + void* workspace, + void* output); + +template +class MoEPacked final : public CudaKernel { + public: + explicit MoEPacked(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t k_{0}; + int64_t normalize_routing_weights_{1}; + std::string activation_type_; + int64_t gate_type_{-1}; + int64_t up_type_{-1}; + int64_t down_type_{-1}; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc new file mode 100644 index 0000000000000..0fd7c1d1f0613 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/gather_packed.h" + +#include + +#include "contrib_ops/cpu/quantization/gguf_quant_tables.h" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define REGISTER_GATHER_PACKED_CUDA(T2, Tind) \ + ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ + GatherPacked, \ + kMSDomain, 1, \ + T2, Tind, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ + GatherPacked); + +REGISTER_GATHER_PACKED_CUDA(float, int32_t) +REGISTER_GATHER_PACKED_CUDA(float, int64_t) +REGISTER_GATHER_PACKED_CUDA(MLFloat16, int32_t) +REGISTER_GATHER_PACKED_CUDA(MLFloat16, int64_t) +REGISTER_GATHER_PACKED_CUDA(BFloat16, int32_t) +REGISTER_GATHER_PACKED_CUDA(BFloat16, int64_t) + +#undef REGISTER_GATHER_PACKED_CUDA + +template +GatherPacked::GatherPacked(const OpKernelInfo& info) : CudaKernel(info) { + ORT_ENFORCE(info.GetAttr("K", &k_).IsOK() && k_ > 0, "GatherPacked requires K > 0"); + ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK(), "GatherPacked requires packed_type"); + row_bytes_ = gguf_quant::GgmlRowBytes(static_cast(ggml_type_), k_); + ORT_ENFORCE(row_bytes_ > 0, "GatherPacked: unsupported packed_type=", ggml_type_, " or invalid K=", k_); +} + +template +Status GatherPacked::ComputeInternal(OpKernelContext* context) const { + const Tensor* data = context->Input(0); + const Tensor* indices = context->Input(1); + const auto& data_shape = data->Shape(); + ORT_RETURN_IF_NOT(data_shape.NumDimensions() == 2, "GatherPacked data must be 2D [vocab_size, row_bytes]"); + const int64_t vocab_size = data_shape[0]; + ORT_RETURN_IF_NOT(data_shape[1] == row_bytes_, + "GatherPacked data.shape[1]=", data_shape[1], " does not match expected row_bytes=", row_bytes_); + + const auto& indices_shape = indices->Shape(); + TensorShapeVector out_dims = indices_shape.AsShapeVector(); + out_dims.push_back(k_); + Tensor* output = context->Output(0, TensorShape(out_dims)); + + const int64_t num_rows = indices_shape.Size(); + if (num_rows == 0) return Status::OK(); + + constexpr int out_dtype = + std::is_same_v ? 16 : (std::is_same_v ? 10 : 1); + constexpr bool is_int64 = std::is_same_v; + + return LaunchGatherPackedCuda( + Stream(context), + data->Data(), + indices->DataRaw(), + is_int64, + output->MutableDataRaw(), + out_dtype, + num_rows, + vocab_size, + k_, + row_bytes_, + static_cast(ggml_type_)); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu new file mode 100644 index 0000000000000..f897f25967b03 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu @@ -0,0 +1,103 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/gather_packed.h" +#include "contrib_ops/cuda/quantization/gguf_quant_cuda.cuh" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +namespace { + +template +__global__ __launch_bounds__(256) void GatherPackedKernel( + const uint8_t* __restrict__ data, + const Tind* __restrict__ indices, + CudaT* __restrict__ output, + int64_t num_rows, + int64_t vocab_size, + int k, + int64_t row_bytes, + int ggml_type) { + const int64_t r = static_cast(blockIdx.x); + if (r >= num_rows) return; + const int64_t raw_id = static_cast(indices[r]); + const int64_t tok = (raw_id >= 0 && raw_id < vocab_size) ? raw_id : 0; + const uint8_t* row_ptr = data + tok * row_bytes; + CudaT* dst_row = output + r * static_cast(k); + + const int total_quarters = k / 8; + for (int q_idx = threadIdx.x; q_idx < total_quarters; q_idx += blockDim.x) { + float vals[8]; + gguf_cuda::DecodeGgufSubBlock8(row_ptr, ggml_type, q_idx >> 2, q_idx & 3, vals); + CudaT* dst8 = dst_row + q_idx * 8; +#pragma unroll + for (int j = 0; j < 8; ++j) { + dst8[j] = gguf_cuda::FloatToCudaType(vals[j]); + } + } +} + +template +Status DispatchGatherPackedByIndex( + cudaStream_t stream, + const uint8_t* data, + const void* indices, + bool indices_are_int64, + CudaT* output, + int64_t num_rows, + int64_t vocab_size, + int64_t k, + int64_t row_bytes, + int ggml_type) { + if (num_rows == 0) return Status::OK(); + ORT_RETURN_IF_NOT((k % 32) == 0, "GatherPacked CUDA requires K to be a multiple of 32"); + const dim3 grid(static_cast(num_rows)); + const dim3 block(256); + if (indices_are_int64) { + GatherPackedKernel<<>>( + data, static_cast(indices), output, num_rows, vocab_size, static_cast(k), row_bytes, + ggml_type); + } else { + GatherPackedKernel<<>>( + data, static_cast(indices), output, num_rows, vocab_size, static_cast(k), row_bytes, + ggml_type); + } + return CUDA_CALL(cudaGetLastError()); +} + +} // namespace + +Status LaunchGatherPackedCuda( + cudaStream_t stream, + const uint8_t* data, + const void* indices, + bool indices_are_int64, + void* output, + int output_dtype, + int64_t num_rows, + int64_t vocab_size, + int64_t k, + int64_t row_bytes, + int ggml_type) { + if (output_dtype == 16) { + return DispatchGatherPackedByIndex<__nv_bfloat16>( + stream, data, indices, indices_are_int64, static_cast<__nv_bfloat16*>(output), num_rows, vocab_size, k, + row_bytes, ggml_type); + } else if (output_dtype == 10) { + return DispatchGatherPackedByIndex( + stream, data, indices, indices_are_int64, static_cast(output), num_rows, vocab_size, k, row_bytes, + ggml_type); + } else if (output_dtype == 1) { + return DispatchGatherPackedByIndex( + stream, data, indices, indices_are_int64, static_cast(output), num_rows, vocab_size, k, row_bytes, + ggml_type); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported GatherPacked output_dtype=", output_dtype); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.h b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.h new file mode 100644 index 0000000000000..109815ec48825 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.h @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +Status LaunchGatherPackedCuda( + cudaStream_t stream, + const uint8_t* data, + const void* indices, + bool indices_are_int64, + void* output, + int output_dtype, + int64_t num_rows, + int64_t vocab_size, + int64_t k, + int64_t row_bytes, + int ggml_type); + +template +class GatherPacked final : public CudaKernel { + public: + explicit GatherPacked(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t k_{0}; + int64_t ggml_type_{-1}; + int64_t row_bytes_{0}; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh b/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh new file mode 100644 index 0000000000000..552b972bac4ee --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh @@ -0,0 +1,690 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "contrib_ops/cpu/quantization/gguf_quant_tables.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace gguf_cuda { + +using namespace onnxruntime::contrib::gguf_quant; + +__device__ __forceinline__ uint16_t ReadU16Dev(const uint8_t* p) { + return __ldg(reinterpret_cast(p)); +} + +__device__ __forceinline__ uint32_t ReadU32Dev(const uint8_t* p) { + const uint32_t lo = __ldg(reinterpret_cast(p)); + const uint32_t hi = __ldg(reinterpret_cast(p + 2)); + return __byte_perm(lo, hi, 0x5410u); +} + +__device__ __forceinline__ float Fp16ToFloatDev(uint16_t bits) { + return __half2float(__ushort_as_half(bits)); +} + +__device__ __forceinline__ float Bf16ToFloatDev(uint16_t bits) { + return __uint_as_float(static_cast(bits) << 16); +} + +__device__ __forceinline__ void GetScaleMinK4Dev(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = __ldg(q + j) & 63; + m = __ldg(q + j + 4) & 63; + } else { + d = (__ldg(q + j + 4) & 0xF) | ((__ldg(q + j - 4) >> 6) << 4); + m = (__ldg(q + j + 4) >> 4) | ((__ldg(q + j - 0) >> 6) << 4); + } +} + +__device__ __forceinline__ void MmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +#endif +} + +__device__ __forceinline__ uint4 PackFloat8ToBf16x2(const float (&v)[8]) { + uint4 res; + *reinterpret_cast<__nv_bfloat162*>(&res.x) = __floats2bfloat162_rn(v[0], v[1]); + *reinterpret_cast<__nv_bfloat162*>(&res.y) = __floats2bfloat162_rn(v[2], v[3]); + *reinterpret_cast<__nv_bfloat162*>(&res.z) = __floats2bfloat162_rn(v[4], v[5]); + *reinterpret_cast<__nv_bfloat162*>(&res.w) = __floats2bfloat162_rn(v[6], v[7]); + return res; +} + +// Register-only computation of ksigns_iq2xs[x7] (7-bit index -> 8-bit even-parity sign mask). +__device__ __forceinline__ uint8_t Iq2xsSignsDev(uint32_t x7) { + const uint32_t s = x7 & 127u; + return static_cast(s | ((__popc(s) & 1u) << 7)); +} + +// Register-only lookup of kvalues_iq4nl[n & 0xF] via hardware prmt.b32 (__byte_perm). +__device__ __forceinline__ float Iq4NlValueDev(uint32_t n4) { + const uint32_t n = n4 & 0xFu; + const uint32_t lo = __byte_perm(0xbfad9881u, 0xf6eaddcfu, n); + const uint32_t hi = __byte_perm(0x26190d01u, 0x71594535u, n); + return static_cast(static_cast((n & 8u) ? hi : lo)); +} + +// Decodes 8 contiguous elements [sb32_idx * 32 + l * 8 .. + 7] (where l in 0..3) of a GGUF weight row. +template +__device__ __forceinline__ void DecodeGgufSubBlock8Impl( + const uint8_t* row_ptr, int sb32_idx, int l, float (&out)[8]) { + if constexpr (kCompileType == kGgmlTypeF32) { + const float* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + l * 8; +#pragma unroll + for (int j = 0; j < 8; ++j) out[j] = __ldg(p + j); + } else if constexpr (kCompileType == kGgmlTypeF16) { + const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + l * 8; +#pragma unroll + for (int j = 0; j < 8; ++j) out[j] = Fp16ToFloatDev(__ldg(p + j)); + } else if constexpr (kCompileType == kGgmlTypeBF16) { + const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + l * 8; +#pragma unroll + for (int j = 0; j < 8; ++j) out[j] = Bf16ToFloatDev(__ldg(p + j)); + } else if constexpr (kCompileType == kGgmlTypeQ8_0) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 34; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint32_t q0 = ReadU32Dev(blk + 2 + l * 8); + const uint32_t q1 = ReadU32Dev(blk + 6 + l * 8); +#pragma unroll + for (int j = 0; j < 4; ++j) { + out[j] = d * static_cast(static_cast((q0 >> (8 * j)) & 0xFFu)); + out[j + 4] = d * static_cast(static_cast((q1 >> (8 * j)) & 0xFFu)); + } + } else if constexpr (kCompileType == kGgmlTypeQ4_0) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 18; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const int shift = (l >> 1) * 4; + const int off = (l & 1) * 8; + const uint32_t q0 = ReadU32Dev(blk + 2 + off); + const uint32_t q1 = ReadU32Dev(blk + 6 + off); +#pragma unroll + for (int j = 0; j < 4; ++j) { + out[j] = static_cast(static_cast((q0 >> (8 * j + shift)) & 0xFu) - 8) * d; + out[j + 4] = static_cast(static_cast((q1 >> (8 * j + shift)) & 0xFu) - 8) * d; + } + } else if constexpr (kCompileType == kGgmlTypeQ5_0) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 22; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint32_t qh = ReadU32Dev(blk + 2) >> (l * 8); + const int shift = (l >> 1) * 4; + const int off = (l & 1) * 8; + const uint32_t q0 = ReadU32Dev(blk + 6 + off); + const uint32_t q1 = ReadU32Dev(blk + 10 + off); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int q5_lo = static_cast(((q0 >> (8 * j + shift)) & 0xFu) | (((qh >> j) & 1u) << 4)); + const int q5_hi = static_cast(((q1 >> (8 * j + shift)) & 0xFu) | (((qh >> (j + 4)) & 1u) << 4)); + out[j] = static_cast(q5_lo - 16) * d; + out[j + 4] = static_cast(q5_hi - 16) * d; + } + } else if constexpr (kCompileType == kGgmlTypeQ2_0) { + const int blk64 = sb32_idx >> 1; + const int sub32 = sb32_idx & 1; + const uint8_t* blk = row_ptr + static_cast(blk64) * 18; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint16_t q16 = ReadU16Dev(blk + 2 + sub32 * 8 + l * 2); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int q = static_cast((q16 >> (2 * j)) & 0x3u); + out[j] = static_cast(q - 1) * d; + } + } else if constexpr (kCompileType == kGgmlTypeIQ4_NL) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 18; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const int shift = (l >> 1) * 4; + const int off = (l & 1) * 8; + const uint32_t q0 = ReadU32Dev(blk + 2 + off); + const uint32_t q1 = ReadU32Dev(blk + 6 + off); +#pragma unroll + for (int j = 0; j < 4; ++j) { + out[j] = d * Iq4NlValueDev(q0 >> (8 * j + shift)); + out[j + 4] = d * Iq4NlValueDev(q1 >> (8 * j + shift)); + } + } else if constexpr (kCompileType == kGgmlTypeIQ4_XS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 136; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint16_t scales_h = ReadU16Dev(blk + 2); + const uint8_t sc_l = __ldg(blk + 4 + (ib >> 1)); + const int ls = ((sc_l >> (4 * (ib & 1))) & 0xF) | (((scales_h >> (2 * ib)) & 3) << 4); + const float dl = d * static_cast(ls - 32); + const int shift = (l >> 1) * 4; + const int off = (l & 1) * 8; + const uint32_t q0 = ReadU32Dev(blk + 8 + ib * 16 + off); + const uint32_t q1 = ReadU32Dev(blk + 12 + ib * 16 + off); +#pragma unroll + for (int j = 0; j < 4; ++j) { + out[j] = dl * Iq4NlValueDev(q0 >> (8 * j + shift)); + out[j + 4] = dl * Iq4NlValueDev(q1 >> (8 * j + shift)); + } + } else if constexpr (kCompileType == kGgmlTypeQ5_K) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const int half16 = l >> 1; + const int off = (l & 1) * 8; + const uint8_t* blk = row_ptr + static_cast(sb256) * 176; + const float dall = Fp16ToFloatDev(ReadU16Dev(blk)); + const float dmin = Fp16ToFloatDev(ReadU16Dev(blk + 2)); + uint8_t sc, m; + GetScaleMinK4Dev(ib, blk + 4, sc, m); + const float d_sub = dall * static_cast(sc); + const float m_sub = dmin * static_cast(m); + const uint32_t qh0 = ReadU32Dev(blk + 16 + half16 * 16 + off); + const uint32_t qh1 = ReadU32Dev(blk + 20 + half16 * 16 + off); + const uint32_t ql0 = ReadU32Dev(blk + 48 + (ib >> 1) * 32 + half16 * 16 + off); + const uint32_t ql1 = ReadU32Dev(blk + 52 + (ib >> 1) * 32 + half16 * 16 + off); + const int shift = (ib & 1) * 4; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int q0 = static_cast(((ql0 >> (8 * j + shift)) & 0xFu) | (((qh0 >> (8 * j + ib)) & 1u) << 4)); + const int q1 = static_cast(((ql1 >> (8 * j + shift)) & 0xFu) | (((qh1 >> (8 * j + ib)) & 1u) << 4)); + out[j] = d_sub * static_cast(q0) - m_sub; + out[j + 4] = d_sub * static_cast(q1) - m_sub; + } + } else if constexpr (kCompileType == kGgmlTypeQ6_K) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const int half16 = l >> 1; + const int off = (l & 1) * 8; + const uint8_t* blk = row_ptr + static_cast(sb256) * 210; + const float d = Fp16ToFloatDev(ReadU16Dev(blk + 208)); + const int ip = ib >> 2; + const int q_grp = ib & 3; + const int8_t sc = static_cast(__ldg(blk + 192 + 8 * ip + 2 * q_grp + half16)); + const float dl = d * static_cast(sc); + const uint32_t ql0 = ReadU32Dev(blk + 64 * ip + 32 * (q_grp & 1) + half16 * 16 + off); + const uint32_t ql1 = ReadU32Dev(blk + 64 * ip + 32 * (q_grp & 1) + half16 * 16 + off + 4); + const uint32_t qh0 = ReadU32Dev(blk + 128 + 32 * ip + half16 * 16 + off); + const uint32_t qh1 = ReadU32Dev(blk + 128 + 32 * ip + half16 * 16 + off + 4); + const int ql_shift = (q_grp >> 1) * 4; + const int qh_shift = q_grp * 2; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int q0 = + static_cast(((ql0 >> (8 * j + ql_shift)) & 0xFu) | (((qh0 >> (8 * j + qh_shift)) & 3u) << 4)) - 32; + const int q1 = + static_cast(((ql1 >> (8 * j + ql_shift)) & 0xFu) | (((qh1 >> (8 * j + qh_shift)) & 3u) << 4)) - 32; + out[j] = dl * static_cast(q0); + out[j + 4] = dl * static_cast(q1); + } + } else if constexpr (kCompileType == kGgmlTypeIQ2_XXS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 66; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint8_t* aux8 = blk + 2 + 8 * ib; + const uint32_t aux32 = ReadU32Dev(aux8 + 4); + const float db = d * (0.5f + static_cast(aux32 >> 28)) * 0.25f; + const uint64_t g64 = __ldg(iq2xxs_grid + __ldg(aux8 + l)); + const uint8_t* grid = reinterpret_cast(&g64); + const uint8_t signs = Iq2xsSignsDev(aux32 >> (7 * l)); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = db * static_cast(grid[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); + } + } else if constexpr (kCompileType == kGgmlTypeIQ2_XS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 74; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint8_t sc_byte = __ldg(blk + 66 + ib); + const float db = d * (0.5f + static_cast((sc_byte >> (4 * (l >> 1))) & 0xF)) * 0.25f; + const uint16_t q2 = ReadU16Dev(blk + 2 + (4 * ib + l) * 2); + const uint64_t g64 = __ldg(iq2xs_grid + (q2 & 511)); + const uint8_t* grid = reinterpret_cast(&g64); + const uint8_t signs = Iq2xsSignsDev(q2 >> 9); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = db * static_cast(grid[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); + } + } else if constexpr (kCompileType == kGgmlTypeIQ2_S) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 82; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint8_t sc_byte = __ldg(blk + 74 + ib); + const float db = d * (0.5f + static_cast((sc_byte >> (4 * (l >> 1))) & 0xF)) * 0.25f; + const uint8_t qh_ib = __ldg(blk + 66 + ib); + const int g_idx = + static_cast(__ldg(blk + 2 + 4 * ib + l)) | ((static_cast(qh_ib) << (8 - 2 * l)) & 0x300); + const uint64_t g64 = __ldg(iq2s_grid + g_idx); + const uint8_t* grid = reinterpret_cast(&g64); + const uint8_t signs = __ldg(blk + 34 + 4 * ib + l); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = db * static_cast(grid[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); + } + } else if constexpr (kCompileType == kGgmlTypeIQ3_XXS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 98; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint32_t aux32 = ReadU32Dev(blk + 66 + 4 * ib); + const float db = d * (0.5f + static_cast(aux32 >> 28)) * 0.5f; + const uint8_t signs = Iq2xsSignsDev(aux32 >> (7 * l)); + const uint16_t q2 = ReadU16Dev(blk + 2 + 8 * ib + 2 * l); + const uint32_t g1 = __ldg(iq3xxs_grid + (q2 & 0xFFu)); + const uint32_t g2 = __ldg(iq3xxs_grid + (q2 >> 8)); + const uint8_t* grid1 = reinterpret_cast(&g1); + const uint8_t* grid2 = reinterpret_cast(&g2); +#pragma unroll + for (int j = 0; j < 4; ++j) { + out[j] = db * static_cast(grid1[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); + out[j + 4] = db * static_cast(grid2[j]) * ((signs & (1u << (j + 4))) ? -1.0f : 1.0f); + } + } else { +#pragma unroll + for (int j = 0; j < 8; ++j) out[j] = 0.0f; + } +} + +__device__ __forceinline__ void DecodeGgufSubBlock8( + const uint8_t* row_ptr, int ggml_type, int sb32_idx, int l, float (&out)[8]) { + switch (ggml_type) { + case kGgmlTypeF32: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeF16: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeBF16: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeQ8_0: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeQ4_0: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeQ5_0: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeQ2_0: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeIQ4_NL: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeIQ4_XS: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeQ5_K: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeQ6_K: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeIQ2_XXS: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeIQ2_XS: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeIQ2_S: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + case kGgmlTypeIQ3_XXS: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; + default: +#pragma unroll + for (int j = 0; j < 8; ++j) out[j] = 0.0f; + return; + } +} + +template +__device__ __forceinline__ void DecodeGgufSubBlock8Tpl( + const uint8_t* row_ptr, int runtime_type, int sb32_idx, int l, float (&out)[8]) { + if constexpr (kCompileType >= 0) { + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + } else { + DecodeGgufSubBlock8(row_ptr, runtime_type, sb32_idx, l, out); + } +} + +template +__device__ __forceinline__ void DecodeGgufSubBlock8SmemLutTpl( + const uint8_t* row_ptr, + int runtime_type, + int sb32_idx, + int l, + const void* s_lut, + float (&out)[8]) { + if constexpr (kCompileType == kGgmlTypeIQ2_XXS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 66; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint8_t* aux8 = blk + 2 + 8 * ib; + const uint32_t aux32 = ReadU32Dev(aux8 + 4); + const float db = d * (0.5f + static_cast(aux32 >> 28)) * 0.25f; + const uint64_t g64 = static_cast(s_lut)[__ldg(aux8 + l)]; + const uint8_t* grid = reinterpret_cast(&g64); + const uint8_t signs = Iq2xsSignsDev(aux32 >> (7 * l)); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid[j]); + } + } else if constexpr (kCompileType == kGgmlTypeIQ2_XS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 74; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint8_t sc_byte = __ldg(blk + 66 + ib); + const float db = d * (0.5f + static_cast((sc_byte >> (4 * (l >> 1))) & 0xF)) * 0.25f; + const uint16_t q2 = ReadU16Dev(blk + 2 + (4 * ib + l) * 2); + const uint64_t g64 = static_cast(s_lut)[q2 & 511]; + const uint8_t* grid = reinterpret_cast(&g64); + const uint8_t signs = Iq2xsSignsDev(q2 >> 9); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid[j]); + } + } else if constexpr (kCompileType == kGgmlTypeIQ2_S) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 82; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint8_t sc_byte = __ldg(blk + 74 + ib); + const float db = d * (0.5f + static_cast((sc_byte >> (4 * (l >> 1))) & 0xF)) * 0.25f; + const uint8_t qh_ib = __ldg(blk + 66 + ib); + const int g_idx = + static_cast(__ldg(blk + 2 + 4 * ib + l)) | ((static_cast(qh_ib) << (8 - 2 * l)) & 0x300); + const uint64_t g64 = static_cast(s_lut)[g_idx]; + const uint8_t* grid = reinterpret_cast(&g64); + const uint8_t signs = __ldg(blk + 34 + 4 * ib + l); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid[j]); + } + } else if constexpr (kCompileType == kGgmlTypeIQ3_XXS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 98; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint32_t aux32 = ReadU32Dev(blk + 66 + 4 * ib); + const float db = d * (0.5f + static_cast(aux32 >> 28)) * 0.5f; + const uint8_t signs = Iq2xsSignsDev(aux32 >> (7 * l)); + const uint16_t q2 = ReadU16Dev(blk + 2 + 8 * ib + 2 * l); + const uint32_t g1 = static_cast(s_lut)[q2 & 0xFFu]; + const uint32_t g2 = static_cast(s_lut)[q2 >> 8]; + const uint8_t* grid1 = reinterpret_cast(&g1); + const uint8_t* grid2 = reinterpret_cast(&g2); +#pragma unroll + for (int j = 0; j < 4; ++j) { + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid1[j]); + out[j + 4] = ((signs & (1u << (j + 4))) ? -db : db) * static_cast(grid2[j]); + } + } else { + DecodeGgufSubBlock8Tpl(row_ptr, runtime_type, sb32_idx, l, out); + } +} + +// Decodes 16 contiguous elements [sb32_idx * 32 + half16 * 16 .. + 15] of a GGUF weight row +// with a single scale computation per 16 elements and vectorized 32-bit quant loads. +template +__device__ __forceinline__ void DecodeGgufSubBlock16Impl( + const uint8_t* row_ptr, int sb32_idx, int half16, float (&out)[16]) { + if constexpr (kCompileType == kGgmlTypeF32) { + const float* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + half16 * 16; +#pragma unroll + for (int j = 0; j < 16; ++j) out[j] = __ldg(p + j); + } else if constexpr (kCompileType == kGgmlTypeF16) { + const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + half16 * 16; +#pragma unroll + for (int j = 0; j < 16; ++j) out[j] = Fp16ToFloatDev(__ldg(p + j)); + } else if constexpr (kCompileType == kGgmlTypeBF16) { + const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + half16 * 16; +#pragma unroll + for (int j = 0; j < 16; ++j) out[j] = Bf16ToFloatDev(__ldg(p + j)); + } else if constexpr (kCompileType == kGgmlTypeQ8_0) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 34; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint8_t* qs = blk + 2 + half16 * 16; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t qw = ReadU32Dev(qs + w * 4); +#pragma unroll + for (int b = 0; b < 4; ++b) { + out[w * 4 + b] = d * static_cast(static_cast((qw >> (8 * b)) & 0xFFu)); + } + } + } else if constexpr (kCompileType == kGgmlTypeQ4_0) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 18; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const int shift = half16 * 4; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t qw = ReadU32Dev(blk + 2 + w * 4); +#pragma unroll + for (int b = 0; b < 4; ++b) { + out[w * 4 + b] = static_cast(static_cast((qw >> (8 * b + shift)) & 0xFu) - 8) * d; + } + } + } else if constexpr (kCompileType == kGgmlTypeQ5_0) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 22; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint32_t qh = ReadU32Dev(blk + 2) >> (half16 * 16); + const int shift = half16 * 4; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t qw = ReadU32Dev(blk + 6 + w * 4); +#pragma unroll + for (int b = 0; b < 4; ++b) { + const int idx = w * 4 + b; + const int q5 = static_cast(((qw >> (8 * b + shift)) & 0xFu) | (((qh >> idx) & 1u) << 4)); + out[idx] = static_cast(q5 - 16) * d; + } + } + } else if constexpr (kCompileType == kGgmlTypeQ2_0) { + const int blk64 = sb32_idx >> 1; + const int sub32 = sb32_idx & 1; + const uint8_t* blk = row_ptr + static_cast(blk64) * 18; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint32_t q32 = ReadU32Dev(blk + 2 + sub32 * 8 + half16 * 4); +#pragma unroll + for (int j = 0; j < 16; ++j) { + out[j] = static_cast(static_cast((q32 >> (2 * j)) & 0x3u) - 1) * d; + } + } else if constexpr (kCompileType == kGgmlTypeIQ4_NL) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 18; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const int shift = half16 * 4; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t qw = ReadU32Dev(blk + 2 + w * 4); +#pragma unroll + for (int b = 0; b < 4; ++b) { + out[w * 4 + b] = d * Iq4NlValueDev(qw >> (8 * b + shift)); + } + } + } else if constexpr (kCompileType == kGgmlTypeIQ4_XS) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 136; + const float d = Fp16ToFloatDev(ReadU16Dev(blk)); + const uint16_t scales_h = ReadU16Dev(blk + 2); + const uint8_t sc_l = __ldg(blk + 4 + (ib >> 1)); + const int ls = ((sc_l >> (4 * (ib & 1))) & 0xF) | (((scales_h >> (2 * ib)) & 3) << 4); + const float dl = d * static_cast(ls - 32); + const int shift = half16 * 4; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t qw = ReadU32Dev(blk + 8 + ib * 16 + w * 4); +#pragma unroll + for (int b = 0; b < 4; ++b) { + out[w * 4 + b] = dl * Iq4NlValueDev(qw >> (8 * b + shift)); + } + } + } else if constexpr (kCompileType == kGgmlTypeQ5_K) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 176; + const float dall = Fp16ToFloatDev(ReadU16Dev(blk)); + const float dmin = Fp16ToFloatDev(ReadU16Dev(blk + 2)); + uint8_t sc, m; + GetScaleMinK4Dev(ib, blk + 4, sc, m); + const float d_sub = dall * static_cast(sc); + const float m_sub = dmin * static_cast(m); + const uint8_t* qh = blk + 16 + half16 * 16; + const uint8_t* ql = blk + 48 + (ib >> 1) * 32 + half16 * 16; + const int shift = (ib & 1) * 4; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t qlw = ReadU32Dev(ql + w * 4); + const uint32_t qhw = ReadU32Dev(qh + w * 4); +#pragma unroll + for (int b = 0; b < 4; ++b) { + const int q = static_cast(((qlw >> (8 * b + shift)) & 0xFu) | (((qhw >> (8 * b + ib)) & 1u) << 4)); + out[w * 4 + b] = d_sub * static_cast(q) - m_sub; + } + } + } else if constexpr (kCompileType == kGgmlTypeQ6_K) { + const int sb256 = sb32_idx >> 3; + const int ib = sb32_idx & 7; + const uint8_t* blk = row_ptr + static_cast(sb256) * 210; + const float d = Fp16ToFloatDev(ReadU16Dev(blk + 208)); + const int ip = ib >> 2; + const int q_grp = ib & 3; + const int8_t sc = static_cast(__ldg(blk + 192 + 8 * ip + 2 * q_grp + half16)); + const float dl = d * static_cast(sc); + const uint8_t* ql = blk + 64 * ip + 32 * (q_grp & 1) + half16 * 16; + const uint8_t* qh = blk + 128 + 32 * ip + half16 * 16; + const int ql_shift = (q_grp >> 1) * 4; + const int qh_shift = q_grp * 2; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t qlw = ReadU32Dev(ql + w * 4); + const uint32_t qhw = ReadU32Dev(qh + w * 4); +#pragma unroll + for (int b = 0; b < 4; ++b) { + const int q = + static_cast(((qlw >> (8 * b + ql_shift)) & 0xFu) | (((qhw >> (8 * b + qh_shift)) & 3u) << 4)) - 32; + out[w * 4 + b] = dl * static_cast(q); + } + } + } else { + float lo[8]; + float hi[8]; + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, 2 * half16 + 0, lo); + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, 2 * half16 + 1, hi); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = lo[j]; + out[j + 8] = hi[j]; + } + } +} + +// Decodes 16 contiguous elements [sb32_idx * 32 + half16 * 16 .. + 15] of a GGUF weight row. +__device__ __forceinline__ void DecodeGgufSubBlock32( + const uint8_t* row_ptr, int ggml_type, int sb32_idx, int half16, float (&out)[16]) { + switch (ggml_type) { + case kGgmlTypeF32: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeF16: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeBF16: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeQ8_0: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeQ4_0: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeQ5_0: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeQ2_0: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeIQ4_NL: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeIQ4_XS: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeQ5_K: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeQ6_K: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeIQ2_XXS: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeIQ2_XS: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeIQ2_S: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + case kGgmlTypeIQ3_XXS: + DecodeGgufSubBlock16Impl(row_ptr, sb32_idx, half16, out); + return; + default: +#pragma unroll + for (int j = 0; j < 16; ++j) out[j] = 0.0f; + return; + } +} + +template +__device__ __forceinline__ CudaT FloatToCudaType(float v); + +template <> +__device__ __forceinline__ float FloatToCudaType(float v) { + return v; +} + +template <> +__device__ __forceinline__ half FloatToCudaType(float v) { + return __float2half_rn(v); +} + +template <> +__device__ __forceinline__ __nv_bfloat16 FloatToCudaType<__nv_bfloat16>(float v) { + return __float2bfloat16_rn(v); +} + +template +__device__ __forceinline__ float CudaTypeToFloat(CudaT v); + +template <> +__device__ __forceinline__ float CudaTypeToFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ float CudaTypeToFloat(half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float CudaTypeToFloat<__nv_bfloat16>(__nv_bfloat16 v) { + return __bfloat162float(v); +} + +} // namespace gguf_cuda +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc new file mode 100644 index 0000000000000..0d7e971759542 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_packed.h" + +#include + +#include "contrib_ops/cpu/quantization/gguf_quant_tables.h" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define REGISTER_MATMUL_PACKED_CUDA(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MatMulPacked, \ + kMSDomain, 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()),\ + MatMulPacked); + +REGISTER_MATMUL_PACKED_CUDA(MLFloat16) +REGISTER_MATMUL_PACKED_CUDA(BFloat16) + +#undef REGISTER_MATMUL_PACKED_CUDA + +template +MatMulPacked::MatMulPacked(const OpKernelInfo& info) : CudaKernel(info) { + ORT_ENFORCE(info.GetAttr("K", &k_).IsOK() && k_ > 0, "MatMulPacked requires K > 0"); + ORT_ENFORCE(info.GetAttr("N", &n_).IsOK() && n_ > 0, "MatMulPacked requires N > 0"); + ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK() && ggml_type_ >= 0, + "MatMulPacked requires packed_type >= 0"); +} + +template +Status MatMulPacked::ComputeInternal(OpKernelContext* context) const { + const Tensor* a = context->Input(0); + const Tensor* b = context->Input(1); + ORT_RETURN_IF_NOT(a != nullptr && b != nullptr, "MatMulPacked requires inputs A and B"); + + const auto& a_shape = a->Shape(); + ORT_RETURN_IF_NOT(a_shape.NumDimensions() >= 1, "MatMulPacked input A must have rank >= 1"); + const size_t last_axis = a_shape.NumDimensions() - 1; + ORT_RETURN_IF_NOT(a_shape[last_axis] == k_, "MatMulPacked A last dimension must equal K"); + + const auto& b_shape = b->Shape(); + ORT_RETURN_IF_NOT(b_shape.NumDimensions() == 2, "MatMulPacked weight B must be 2D [N, row_bytes]"); + ORT_RETURN_IF_NOT(b_shape[0] == n_, "MatMulPacked B dim 0 (", b_shape[0], ") != N (", n_, ")"); + const int64_t expected_rb = gguf_quant::GgmlRowBytes(static_cast(ggml_type_), k_); + ORT_RETURN_IF_NOT(expected_rb > 0 && b_shape[1] == expected_rb, + "MatMulPacked weight B row_bytes=", b_shape[1], " != expected=", expected_rb); + + MatMulPackedParams params{}; + params.b_data = b->Data(); + params.k = static_cast(k_); + params.n = static_cast(n_); + params.row_bytes = expected_rb; + params.ggml_type = static_cast(ggml_type_); + + TensorShapeVector y_dims = a_shape.AsShapeVector(); + y_dims[last_axis] = n_; + Tensor* y = context->Output(0, TensorShape(y_dims)); + + const int64_t m = a_shape.SizeToDimension(last_axis); + if (m == 0) return Status::OK(); + + constexpr bool is_bf16 = std::is_same_v; + IAllocatorUniquePtr scratch_w; + if (m > 8) { + scratch_w = GetScratchBuffer(static_cast(n_ * k_), GetComputeStream(context)); + } + + return LaunchMatMulPackedCuda( + GetCublasHandle(context), + Stream(context), + is_bf16, + a->DataRaw(), + params, + scratch_w ? scratch_w.get() : nullptr, + y->MutableDataRaw(), + m); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu new file mode 100644 index 0000000000000..c1df32db6f813 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu @@ -0,0 +1,376 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_packed.h" +#include "contrib_ops/cuda/quantization/gguf_quant_cuda.cuh" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +namespace { + +// Warp-K-Split Tensor-Core BF16 GEMV/GEMM kernel for M in 1..8 on large-N projections (e.g. lm_head). +// KW warps (KW*32 threads) split K across the KW warps with zero __syncthreads() in the K-loop, +// while computing MT*16 output columns per CTA. +template +__global__ __launch_bounds__(KW * 32) void MatMulPackedBf16TcWarpKSplitKernel( + const __nv_bfloat16* __restrict__ a, + MatMulPackedParams p, + __nv_bfloat16* __restrict__ y, + int m_rt) { + constexpr int kThreads = KW * 32; + __shared__ float s_red[KW][MT * 16][8]; + + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + const int lane = tid & 31; + const int g = lane >> 2; // 0..7 + const int t = lane & 3; // 0..3 + + const int n_out_base = static_cast(blockIdx.x) * (MT * 16); + const uint8_t* row_ptrs_lo[MT]; + const uint8_t* row_ptrs_hi[MT]; + bool valid_lo[MT]; + bool valid_hi[MT]; + float acc[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + const int r_lo = n_out_base + m * 16 + g; + const int r_hi = r_lo + 8; + valid_lo[m] = (r_lo < p.n); + valid_hi[m] = (r_hi < p.n); + row_ptrs_lo[m] = p.b_data + static_cast(valid_lo[m] ? r_lo : 0) * p.row_bytes; + row_ptrs_hi[m] = p.b_data + static_cast(valid_hi[m] ? r_hi : 0) * p.row_bytes; +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc[m][i] = 0.0f; + } + } + + const bool a_valid = (g < m_rt); + const __nv_bfloat16* a_row = a + static_cast(a_valid ? g : 0) * p.k; + + const int num_sb32 = p.k >> 5; + for (int sb32_idx = warp; sb32_idx < num_sb32; sb32_idx += KW) { + const uint4 b_vec = a_valid + ? __ldg(reinterpret_cast(a_row + sb32_idx * 32 + t * 8)) + : make_uint4(0u, 0u, 0u, 0u); +#pragma unroll + for (int m = 0; m < MT; ++m) { + if (valid_lo[m]) { + float w_lo[8]; + float w_hi[8]; + gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_lo[m], p.ggml_type, sb32_idx, t, w_lo); + if (valid_hi[m]) { + gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_hi[m], p.ggml_type, sb32_idx, t, w_hi); + } else { +#pragma unroll + for (int j = 0; j < 8; ++j) w_hi[j] = 0.0f; + } + const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(w_lo); + const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(w_hi); + gguf_cuda::MmaBf16(acc[m], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); + gguf_cuda::MmaBf16(acc[m], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); + } + } + } + +#pragma unroll + for (int m = 0; m < MT; ++m) { + s_red[warp][m * 16 + g][2 * t + 0] = acc[m][0]; + s_red[warp][m * 16 + g][2 * t + 1] = acc[m][1]; + s_red[warp][m * 16 + g + 8][2 * t + 0] = acc[m][2]; + s_red[warp][m * 16 + g + 8][2 * t + 1] = acc[m][3]; + } + __syncthreads(); + + for (int idx = tid; idx < MT * 16 * 8; idx += kThreads) { + const int row_m = idx / (MT * 16); + const int rl = idx - row_m * (MT * 16); + const int col = n_out_base + rl; + if (row_m < m_rt && col < p.n) { + float sum = 0.0f; +#pragma unroll + for (int w = 0; w < KW; ++w) { + sum += s_red[w][rl][row_m]; + } + y[static_cast(row_m) * p.n + col] = __float2bfloat16_rn(sum); + } + } +} + +template +__global__ __launch_bounds__(256) void MatMulPackedGemvKernel( + const CudaT* __restrict__ a, + MatMulPackedParams p, + CudaT* __restrict__ y, + int m_rt) { + constexpr int kWarpsPerBlock = 8; + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int n_out = static_cast(blockIdx.x) * kWarpsPerBlock + warp; + if (n_out >= p.n) return; + + const int M = M_FIXED > 0 ? M_FIXED : m_rt; + const uint8_t* row_ptr = p.b_data + static_cast(n_out) * p.row_bytes; + + float acc[M_FIXED > 0 ? M_FIXED : 8]; +#pragma unroll + for (int m = 0; m < (M_FIXED > 0 ? M_FIXED : 8); ++m) { + acc[m] = 0.0f; + } + + const int total_quarters = p.k >> 3; +#pragma unroll 1 + for (int q_idx = lane; q_idx < total_quarters; q_idx += 32) { + const int sb32_idx = q_idx >> 2; + const int l = q_idx & 3; + float w_vals[8]; + gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptr, p.ggml_type, sb32_idx, l, w_vals); + const int k_offset = sb32_idx * 32 + (l << 3); + + if constexpr (M_FIXED > 0) { +#pragma unroll + for (int m = 0; m < M_FIXED; ++m) { + const CudaT* a_ptr = a + static_cast(m) * p.k + k_offset; + const uint4 v0 = __ldg(reinterpret_cast(a_ptr)); + const CudaT* a0 = reinterpret_cast(&v0); +#pragma unroll + for (int j = 0; j < 8; ++j) { + acc[m] += w_vals[j] * gguf_cuda::CudaTypeToFloat(a0[j]); + } + } + } else { + for (int m = 0; m < M; ++m) { + const CudaT* a_ptr = a + static_cast(m) * p.k + k_offset; + const uint4 v0 = __ldg(reinterpret_cast(a_ptr)); + const CudaT* a0 = reinterpret_cast(&v0); +#pragma unroll + for (int j = 0; j < 8; ++j) { + acc[m] += w_vals[j] * gguf_cuda::CudaTypeToFloat(a0[j]); + } + } + } + } + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int m = 0; m < (M_FIXED > 0 ? M_FIXED : 8); ++m) { + acc[m] += __shfl_down_sync(0xFFFFFFFFu, acc[m], offset); + } + } + + if (lane == 0) { + if constexpr (M_FIXED > 0) { +#pragma unroll + for (int m = 0; m < M_FIXED; ++m) { + y[static_cast(m) * p.n + n_out] = gguf_cuda::FloatToCudaType(acc[m]); + } + } else { + for (int m = 0; m < M; ++m) { + y[static_cast(m) * p.n + n_out] = gguf_cuda::FloatToCudaType(acc[m]); + } + } + } +} + +template +__global__ __launch_bounds__(256) void MatMulPackedDequantKernel( + MatMulPackedParams p, + CudaT* __restrict__ w_out) { + constexpr int kWarpsPerBlock = 8; + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int n_out = static_cast(blockIdx.x) * kWarpsPerBlock + warp; + if (n_out >= p.n) return; + + const uint8_t* row_ptr = p.b_data + static_cast(n_out) * p.row_bytes; + CudaT* dst_row = w_out + static_cast(n_out) * p.k; + + const int total_quarters = p.k / 8; + for (int q_idx = lane; q_idx < total_quarters; q_idx += 32) { + const int sb32_idx = q_idx >> 2; + const int l = q_idx & 3; + float w_vals[8]; + gguf_cuda::DecodeGgufSubBlock8(row_ptr, p.ggml_type, sb32_idx, l, w_vals); + CudaT* dst8 = dst_row + sb32_idx * 32 + l * 8; + if constexpr (std::is_same::value) { + *reinterpret_cast(dst8) = gguf_cuda::PackFloat8ToBf16x2(w_vals); + } else { + uint4 out_u4; + CudaT* out_elems = reinterpret_cast(&out_u4); +#pragma unroll + for (int j = 0; j < 8; ++j) { + out_elems[j] = gguf_cuda::FloatToCudaType(w_vals[j]); + } + *reinterpret_cast(dst8) = out_u4; + } + } +} + +template +void LaunchMatMulPackedGemvByM( + cudaStream_t stream, + dim3 grid, + const CudaT* a, + const MatMulPackedParams& params, + CudaT* y, + int m) { + const dim3 block(256); + switch (m) { + case 1: + MatMulPackedGemvKernel<<>>(a, params, y, 1); + return; + case 2: + MatMulPackedGemvKernel<<>>(a, params, y, 2); + return; + case 4: + MatMulPackedGemvKernel<<>>(a, params, y, 4); + return; + case 7: + MatMulPackedGemvKernel<<>>(a, params, y, 7); + return; + case 8: + MatMulPackedGemvKernel<<>>(a, params, y, 8); + return; + default: + MatMulPackedGemvKernel<<>>(a, params, y, m); + return; + } +} + +template +void LaunchMatMulPackedBf16TcWarpKSplit( + cudaStream_t stream, + const __nv_bfloat16* a, + const MatMulPackedParams& params, + __nv_bfloat16* y, + int m_int) { + using namespace onnxruntime::contrib::gguf_quant; + constexpr int kRowsPerBlock = MT * 16; + const dim3 tc_grid(static_cast((params.n + kRowsPerBlock - 1) / kRowsPerBlock)); + constexpr unsigned int kBlock = KW * 32; + switch (params.ggml_type) { + case kGgmlTypeQ8_0: + MatMulPackedBf16TcWarpKSplitKernel<<>>(a, params, y, m_int); + return; + case kGgmlTypeQ6_K: + MatMulPackedBf16TcWarpKSplitKernel<<>>(a, params, y, m_int); + return; + case kGgmlTypeQ5_K: + MatMulPackedBf16TcWarpKSplitKernel<<>>(a, params, y, m_int); + return; + case kGgmlTypeIQ4_XS: + MatMulPackedBf16TcWarpKSplitKernel<<>>(a, params, y, m_int); + return; + default: + MatMulPackedBf16TcWarpKSplitKernel<<>>(a, params, y, m_int); + return; + } +} + +template +Status DispatchMatMulPackedCuda( + cublasHandle_t cublas, + cudaStream_t stream, + bool is_bf16, + const CudaT* a, + const MatMulPackedParams& params, + CudaT* scratch_w, + CudaT* y, + int64_t m) { + if (m == 0) return Status::OK(); + ORT_RETURN_IF_NOT((params.k % 32) == 0, "MatMulPacked requires K to be a multiple of 32"); + + if (m <= 8) { + using namespace onnxruntime::contrib::gguf_quant; + const int m_int = static_cast(m); + + if constexpr (std::is_same::value) { + if (params.n >= 65536) { + LaunchMatMulPackedBf16TcWarpKSplit<4, 2>(stream, a, params, y, m_int); + return CUDA_CALL(cudaGetLastError()); + } + } + + constexpr int kWarpsPerBlock = 8; + const dim3 grid(static_cast((params.n + kWarpsPerBlock - 1) / kWarpsPerBlock)); + + switch (params.ggml_type) { + case kGgmlTypeQ8_0: + LaunchMatMulPackedGemvByM(stream, grid, a, params, y, m_int); + return CUDA_CALL(cudaGetLastError()); + case kGgmlTypeQ6_K: + LaunchMatMulPackedGemvByM(stream, grid, a, params, y, m_int); + return CUDA_CALL(cudaGetLastError()); + case kGgmlTypeQ5_K: + LaunchMatMulPackedGemvByM(stream, grid, a, params, y, m_int); + return CUDA_CALL(cudaGetLastError()); + case kGgmlTypeIQ4_XS: + LaunchMatMulPackedGemvByM(stream, grid, a, params, y, m_int); + return CUDA_CALL(cudaGetLastError()); + default: + LaunchMatMulPackedGemvByM(stream, grid, a, params, y, m_int); + return CUDA_CALL(cudaGetLastError()); + } + } + + ORT_RETURN_IF_NOT(scratch_w != nullptr, "MatMulPacked prefill requires non-null scratch_w"); + constexpr int kWarpsPerBlock = 8; + const dim3 grid(static_cast((params.n + kWarpsPerBlock - 1) / kWarpsPerBlock)); + const dim3 block(256); + MatMulPackedDequantKernel<<>>(params, scratch_w); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + + const float alpha = 1.0f; + const float beta = 0.0f; + const cudaDataType_t data_type = is_bf16 ? CUDA_R_16BF : CUDA_R_16F; + return CUBLAS_CALL(cublasGemmEx( + cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + params.n, + static_cast(m), + params.k, + &alpha, + scratch_w, + data_type, + params.k, + a, + data_type, + params.k, + &beta, + y, + data_type, + params.n, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +} + +} // namespace + +Status LaunchMatMulPackedCuda( + cublasHandle_t cublas, + cudaStream_t stream, + bool is_bf16, + const void* a, + const MatMulPackedParams& params, + void* scratch_w, + void* y, + int64_t m) { + if (is_bf16) { + return DispatchMatMulPackedCuda<__nv_bfloat16>( + cublas, stream, true, static_cast(a), params, + static_cast<__nv_bfloat16*>(scratch_w), static_cast<__nv_bfloat16*>(y), m); + } + return DispatchMatMulPackedCuda( + cublas, stream, false, static_cast(a), params, + static_cast(scratch_w), static_cast(y), m); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h new file mode 100644 index 0000000000000..b4760c5a796ab --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +struct MatMulPackedParams { + const uint8_t* b_data{nullptr}; + int k{0}; + int n{0}; + int64_t row_bytes{0}; + int ggml_type{0}; +}; + +Status LaunchMatMulPackedCuda( + cublasHandle_t cublas, + cudaStream_t stream, + bool is_bf16, + const void* a, + const MatMulPackedParams& params, + void* scratch_w, + void* y, + int64_t m); + +template +class MatMulPacked final : public CudaKernel { + public: + explicit MatMulPacked(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t k_{0}; + int64_t n_{0}; + int64_t ggml_type_{-1}; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index f98390a7b5555..f08447115762a 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -992,6 +992,92 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } })); +ONNX_MS_OPERATOR_SET_SCHEMA( + GatherPacked, 1, + OpSchema() + .SetDoc("Gather rows from a 2D block-packed embedding table and dequantize on the fly.") + .Attr("K", "Logical embedding dimension.", AttributeProto::INT) + .Attr("packed_type", + "Packed quantization format enum (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).", + AttributeProto::INT) + .Attr("output_dtype", "ONNX TensorProto DataType for output (default 16 = BFLOAT16).", + AttributeProto::INT, static_cast(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16)) + .Input(0, "data", "2D packed weight blob [vocab_size, row_bytes].", "T1") + .Input(1, "indices", "Token indices of arbitrary rank.", "Tind") + .Output(0, "output", "Dequantized gathered embeddings [..., K].", "T2") + .TypeConstraint("T1", {"tensor(uint8)"}, "Packed byte storage.") + .TypeConstraint("T2", {"tensor(bfloat16)", "tensor(float16)", "tensor(float)"}, "Output floating-point type.") + .TypeConstraint("Tind", {"tensor(int32)", "tensor(int64)"}, "Index type.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const int64_t k = getAttribute(ctx, "K", 0); + if (k <= 0) fail_shape_inference("GatherPacked K must be positive."); + const int64_t out_dt = getAttribute( + ctx, "output_dtype", static_cast(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16)); + ctx.getOutputType(0)->mutable_tensor_type()->set_elem_type(static_cast(out_dt)); + if (!hasInputShape(ctx, 1)) return; + const auto& indices_shape = getInputShape(ctx, 1); + auto* out_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape(); + for (int i = 0; i < indices_shape.dim_size(); ++i) { + *out_shape->add_dim() = indices_shape.dim(i); + } + out_shape->add_dim()->set_dim_value(k); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + MatMulPacked, 1, + OpSchema() + .SetDoc("Matrix multiplication Y = A @ B^T where B is a 2D block-packed weight tensor.") + .Attr("K", "Reduction dimension.", AttributeProto::INT) + .Attr("N", "Output dimension.", AttributeProto::INT) + .Attr("packed_type", + "Packed quantization format enum for B (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).", + AttributeProto::INT) + .Input(0, "A", "Input activations [..., K].", "T1") + .Input(1, "B", "Packed weight blob [N, row_bytes].", "T2") + .Output(0, "Y", "Output activations [..., N].", "T1") + .TypeConstraint("T1", {"tensor(bfloat16)", "tensor(float16)"}, "Activation type.") + .TypeConstraint("T2", {"tensor(uint8)"}, "Packed byte storage.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const int64_t n = getAttribute(ctx, "N", 0); + if (n <= 0) fail_shape_inference("MatMulPacked N must be positive."); + if (!hasInputShape(ctx, 0)) return; + const auto& a_shape = getInputShape(ctx, 0); + if (a_shape.dim_size() < 1) fail_shape_inference("MatMulPacked A must have rank >= 1."); + auto* y_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape(); + for (int i = 0; i + 1 < a_shape.dim_size(); ++i) { + *y_shape->add_dim() = a_shape.dim(i); + } + y_shape->add_dim()->set_dim_value(n); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + MoEPacked, 1, + OpSchema() + .SetDoc("Mixture-of-Experts layer consuming 3D block-packed gate, up, and down expert weight blobs.") + .Attr("k", "Number of selected experts per token.", AttributeProto::INT) + .Attr("normalize_routing_weights", "Whether to normalize top-k routing probabilities.", + AttributeProto::INT, static_cast(1)) + .Attr("activation_type", "Expert activation function ('swiglu').", + AttributeProto::STRING, std::string("swiglu")) + .Attr("gate_type", "Packed quantization format enum of gate_experts_weights.", AttributeProto::INT) + .Attr("up_type", "Packed quantization format enum of up_experts_weights.", AttributeProto::INT) + .Attr("down_type", "Packed quantization format enum of down_experts_weights.", AttributeProto::INT) + .Input(0, "input", "Input hidden states [..., hidden_size].", "T") + .Input(1, "router_probs", "Router logits [num_rows, num_experts].", "T") + .Input(2, "gate_experts_weights", "3D packed gate weight blob [num_experts, inter_size, gate_row_bytes].", "T2") + .Input(3, "up_experts_weights", "3D packed up weight blob [num_experts, inter_size, up_row_bytes].", "T2") + .Input(4, "down_experts_weights", "3D packed down weight blob [num_experts, hidden_size, down_row_bytes].", "T2") + .Output(0, "output", "Output hidden states [..., hidden_size].", "T") + .TypeConstraint("T", {"tensor(bfloat16)", "tensor(float16)"}, "Activation type.") + .TypeConstraint("T2", {"tensor(uint8)"}, "Packed byte storage.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (hasInputShape(ctx, 0)) { + *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape() = getInputShape(ctx, 0); + } + })); + ONNX_MS_OPERATOR_SET_SCHEMA(BiasDropout, 1, OpSchema() .SetDoc( diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index f471540ce68b6..7966c0049f434 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -124,6 +124,9 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SampleOp); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Sampling); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SwiGLU); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RouterGateUpProj); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatherPacked); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MatMulPacked); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MoEPacked); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipGroupNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipLayerNormalization); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipSimplifiedLayerNormalization); @@ -266,6 +269,9 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/python/transformers/test_gguf_packed_ops.py b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py new file mode 100644 index 0000000000000..fb345292d6dca --- /dev/null +++ b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py @@ -0,0 +1,229 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import unittest + +import numpy as np +from onnx import TensorProto, helper + +import onnxruntime as ort + +GGUF_TYPES_INFO = [ + (0, 1, 4), + (1, 1, 2), + (2, 32, 18), + (6, 32, 22), + (8, 32, 34), + (13, 256, 176), + (14, 256, 210), + (16, 256, 66), + (17, 256, 74), + (18, 256, 98), + (20, 32, 18), + (22, 256, 82), + (23, 256, 136), + (30, 1, 2), + (42, 64, 18), +] + + +def _f32_to_bf16_u16(arr: np.ndarray) -> np.ndarray: + u32 = np.ascontiguousarray(arr, dtype=np.float32).view(np.uint32) + return (((u32 + 0x7FFF + ((u32 >> 16) & 1)) >> 16) & 0xFFFF).astype(np.uint16) + + +def _bf16_u16_to_f32(arr_u16: np.ndarray) -> np.ndarray: + return (arr_u16.astype(np.uint32) << 16).view(np.float32) + + +def _make_synthetic_gguf_rows(rng: np.random.Generator, rows: int, k: int, gtype: int, bel: int, bbytes: int) -> np.ndarray: + row_bytes = (k // bel) * bbytes + if gtype == 0: + return (rng.standard_normal((rows, k), dtype=np.float32) * 0.1).view(np.uint8).reshape(rows, row_bytes) + if gtype == 1: + return (rng.standard_normal((rows, k)) * 0.1).astype(np.float16).view(np.uint8).reshape(rows, row_bytes) + if gtype == 30: + return _f32_to_bf16_u16(rng.standard_normal((rows, k)) * 0.1).view(np.uint8).reshape(rows, row_bytes) + + raw = rng.integers(0, 256, size=(rows, row_bytes), dtype=np.uint8) + for r in range(rows): + for b in range(k // bel): + off = b * bbytes + (208 if gtype == 14 else 0) + raw[r, off] = 0x00 + raw[r, off + 1] = 0x20 # fp16 0.0078125 + if gtype == 13: + raw[r, off + 2] = 0x00 + raw[r, off + 3] = 0x1C + return raw + + +@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDAExecutionProvider is required") +class TestGgufPackedOps(unittest.TestCase): + def test_gather_packed_cpu_vs_cuda_all_types(self): + rng = np.random.default_rng(42) + k = 256 + vocab = 16 + indices = np.arange(vocab, dtype=np.int64) + for gtype, bel, bbytes in GGUF_TYPES_INFO: + row_bytes = (k // bel) * bbytes + w_raw = _make_synthetic_gguf_rows(rng, vocab, k, gtype, bel, bbytes) + node = helper.make_node( + "GatherPacked", + ["W", "I"], + ["Y"], + domain="com.microsoft", + K=k, + packed_type=gtype, + output_dtype=int(TensorProto.FLOAT), + ) + graph = helper.make_graph( + [node], + "gather_packed_test", + [ + helper.make_tensor_value_info("W", TensorProto.UINT8, [vocab, row_bytes]), + helper.make_tensor_value_info("I", TensorProto.INT64, [vocab]), + ], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [vocab, k])], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)] + ) + sess_cpu = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + sess_cuda = ort.InferenceSession(model.SerializeToString(), providers=["CUDAExecutionProvider"]) + y_cpu = sess_cpu.run(None, {"W": w_raw, "I": indices})[0] + y_cuda = sess_cuda.run(None, {"W": w_raw, "I": indices})[0] + np.testing.assert_allclose(y_cpu, y_cuda, rtol=0.0, atol=0.0, err_msg=f"Mismatch for ggml_type={gtype}") + + def test_matmul_packed_cuda_all_types(self): + rng = np.random.default_rng(123) + k = 256 + n = 32 + indices = np.arange(n, dtype=np.int64) + for gtype, bel, bbytes in GGUF_TYPES_INFO: + row_bytes = (k // bel) * bbytes + w_raw = _make_synthetic_gguf_rows(rng, n, k, gtype, bel, bbytes) + + # Dequantize reference via CPU GatherPacked + g_node = helper.make_node( + "GatherPacked", ["W", "I"], ["W_f32"], domain="com.microsoft", K=k, packed_type=gtype, output_dtype=1 + ) + g_model = helper.make_model( + helper.make_graph( + [g_node], + "deq", + [ + helper.make_tensor_value_info("W", TensorProto.UINT8, [n, row_bytes]), + helper.make_tensor_value_info("I", TensorProto.INT64, [n]), + ], + [helper.make_tensor_value_info("W_f32", TensorProto.FLOAT, [n, k])], + ), + opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)], + ) + w_f32 = ort.InferenceSession(g_model.SerializeToString(), providers=["CPUExecutionProvider"]).run( + None, {"W": w_raw, "I": indices} + )[0] + + # Build MatMulPacked model with Cast(FLOAT -> BFLOAT16) -> MatMulPacked -> Cast(BFLOAT16 -> FLOAT) + nodes = [ + helper.make_node("Cast", ["A_f32"], ["A_bf16"], to=int(TensorProto.BFLOAT16)), + helper.make_node( + "MatMulPacked", + ["A_bf16", "B"], + ["Y_bf16"], + domain="com.microsoft", + K=k, + N=n, + packed_type=gtype, + ), + helper.make_node("Cast", ["Y_bf16"], ["Y_f32"], to=int(TensorProto.FLOAT)), + ] + mm_model = helper.make_model( + helper.make_graph( + nodes, + "mm_packed", + [ + helper.make_tensor_value_info("A_f32", TensorProto.FLOAT, ["M", k]), + helper.make_tensor_value_info("B", TensorProto.UINT8, [n, row_bytes]), + ], + [helper.make_tensor_value_info("Y_f32", TensorProto.FLOAT, ["M", n])], + ), + opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)], + ) + w_bf16 = _bf16_u16_to_f32(_f32_to_bf16_u16(w_f32)) + sess = ort.InferenceSession(mm_model.SerializeToString(), providers=["CUDAExecutionProvider"]) + for m in (1, 7, 8, 32): + a_f32 = _bf16_u16_to_f32(_f32_to_bf16_u16(rng.standard_normal((m, k), dtype=np.float32) * 0.2)) + y_got = sess.run(None, {"A_f32": a_f32, "B": w_raw})[0] + w_target = w_f32 if m <= 8 else w_bf16 + y_ref = _bf16_u16_to_f32(_f32_to_bf16_u16(a_f32 @ w_target.T)) + np.testing.assert_allclose( + y_got, y_ref, rtol=1e-2, atol=1e-2, err_msg=f"MatMulPacked mismatch for ggml_type={gtype}, M={m}" + ) + + def test_moe_packed_cuda_decode_and_prefill(self): + rng = np.random.default_rng(999) + num_experts = 16 + hidden_size = 256 + inter_size = 64 + top_k = 4 + gate_type, gate_bel, gate_bb = 16, 256, 66 # IQ2_XXS + up_type, up_bel, up_bb = 16, 256, 66 + down_type, down_bel, down_bb = 42, 64, 18 # Q2_0 + + gate_rb = (hidden_size // gate_bel) * gate_bb + up_rb = (hidden_size // up_bel) * up_bb + down_rb = (inter_size // down_bel) * down_bb + + gate_raw = _make_synthetic_gguf_rows( + rng, num_experts * inter_size, hidden_size, gate_type, gate_bel, gate_bb + ).reshape(num_experts, inter_size, gate_rb) + up_raw = _make_synthetic_gguf_rows( + rng, num_experts * inter_size, hidden_size, up_type, up_bel, up_bb + ).reshape(num_experts, inter_size, up_rb) + down_raw = _make_synthetic_gguf_rows( + rng, num_experts * hidden_size, inter_size, down_type, down_bel, down_bb + ).reshape(num_experts, hidden_size, down_rb) + + nodes = [ + helper.make_node("Cast", ["X_f32"], ["X_bf16"], to=int(TensorProto.BFLOAT16)), + helper.make_node("Cast", ["R_f32"], ["R_bf16"], to=int(TensorProto.BFLOAT16)), + helper.make_node( + "MoEPacked", + ["X_bf16", "R_bf16", "G", "U", "D"], + ["Y_bf16"], + domain="com.microsoft", + k=top_k, + normalize_routing_weights=1, + activation_type="swiglu", + gate_type=gate_type, + up_type=up_type, + down_type=down_type, + ), + helper.make_node("Cast", ["Y_bf16"], ["Y_f32"], to=int(TensorProto.FLOAT)), + ] + moe_model = helper.make_model( + helper.make_graph( + nodes, + "moe_packed", + [ + helper.make_tensor_value_info("X_f32", TensorProto.FLOAT, ["M", hidden_size]), + helper.make_tensor_value_info("R_f32", TensorProto.FLOAT, ["M", num_experts]), + helper.make_tensor_value_info("G", TensorProto.UINT8, [num_experts, inter_size, gate_rb]), + helper.make_tensor_value_info("U", TensorProto.UINT8, [num_experts, inter_size, up_rb]), + helper.make_tensor_value_info("D", TensorProto.UINT8, [num_experts, hidden_size, down_rb]), + ], + [helper.make_tensor_value_info("Y_f32", TensorProto.FLOAT, ["M", hidden_size])], + ), + opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)], + ) + sess = ort.InferenceSession(moe_model.SerializeToString(), providers=["CUDAExecutionProvider"]) + for m in (4, 48): + x_f32 = rng.standard_normal((m, hidden_size), dtype=np.float32) * 0.2 + r_f32 = rng.standard_normal((m, num_experts), dtype=np.float32) + y_got = sess.run(None, {"X_f32": x_f32, "R_f32": r_f32, "G": gate_raw, "U": up_raw, "D": down_raw})[0] + self.assertEqual(y_got.shape, (m, hidden_size)) + self.assertTrue(np.all(np.isfinite(y_got))) + + +if __name__ == "__main__": + unittest.main() From db3e02c47c1c4968a2ee5b13e5a0915822d43e6a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 7 Nov 2026 06:10:58 +0000 Subject: [PATCH 2/5] Fix packed operator validation and CUDA edge cases --- docs/ContribOperators.md | 152 +++++++ docs/OperatorKernels.md | 4 + .../cpu/quantization/gather_packed.cc | 47 +- .../cpu/quantization/gguf_quant_tables.h | 302 +++++++++---- .../contrib_ops/cuda/moe/moe_packed.cc | 29 +- .../contrib_ops/cuda/moe/moe_packed.cu | 69 ++- .../cuda/quantization/gather_packed.cc | 22 +- .../cuda/quantization/gather_packed.cu | 17 +- .../cuda/quantization/matmul_packed.cc | 21 +- .../cuda/quantization/matmul_packed.cu | 22 +- .../transformers/test_gguf_packed_ops.py | 414 ++++++++++++------ 11 files changed, 792 insertions(+), 307 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 102e6d0653d13..f9dca18163c84 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -48,6 +48,7 @@ Do not modify directly.* * com.microsoft.GatedRelativePositionBias * com.microsoft.GatherBlockQuantized * com.microsoft.GatherND + * com.microsoft.GatherPacked * com.microsoft.Gelu * com.microsoft.GemmFastGelu * com.microsoft.GemmFloat8 @@ -75,8 +76,10 @@ Do not modify directly.* * com.microsoft.MatMulNBitsMlp * com.microsoft.MatMulNBitsQLoRA * com.microsoft.MatMulNBitsQkv + * com.microsoft.MatMulPacked * com.microsoft.MaxpoolWithMask * com.microsoft.MoE + * com.microsoft.MoEPacked * com.microsoft.MulInteger * com.microsoft.MultiHeadAttention * com.microsoft.MurmurHash3 @@ -2946,6 +2949,53 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.GatherPacked** + + Gather rows from a 2D block-packed embedding table and dequantize on the fly. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
K : int (required)
+
Logical embedding dimension.
+
output_dtype : int
+
ONNX TensorProto DataType for output (default 16 = BFLOAT16).
+
packed_type : int (required)
+
Packed quantization format enum (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).
+
+ +#### Inputs + +
+
data : T1
+
2D packed weight blob [vocab_size, row_bytes].
+
indices : Tind
+
Token indices of arbitrary rank.
+
+ +#### Outputs + +
+
output : T2
+
Dequantized gathered embeddings [..., K].
+
+ +#### Type Constraints + +
+
T1 : tensor(uint8)
+
Packed byte storage.
+
T2 : tensor(bfloat16), tensor(float16), tensor(float)
+
Output floating-point type.
+
Tind : tensor(int32), tensor(int64)
+
Index type.
+
+ + ### **com.microsoft.Gelu** Gaussian Error Linear Unit. @@ -4671,6 +4721,51 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.MatMulPacked** + + Matrix multiplication Y = A @ B^T where B is a 2D block-packed weight tensor. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
K : int (required)
+
Reduction dimension.
+
N : int (required)
+
Output dimension.
+
packed_type : int (required)
+
Packed quantization format enum for B (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).
+
+ +#### Inputs + +
+
A : T1
+
Input activations [..., K].
+
B : T2
+
Packed weight blob [N, row_bytes].
+
+ +#### Outputs + +
+
Y : T1
+
Output activations [..., N].
+
+ +#### Type Constraints + +
+
T1 : tensor(bfloat16), tensor(float16)
+
Activation type.
+
T2 : tensor(uint8)
+
Packed byte storage.
+
+ + ### **com.microsoft.MaxpoolWithMask** For internal use. @@ -4807,6 +4902,63 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.MoEPacked** + + Mixture-of-Experts layer consuming 3D block-packed gate, up, and down expert weight blobs. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
activation_type : string
+
Expert activation function ('swiglu').
+
down_type : int (required)
+
Packed quantization format enum of down_experts_weights.
+
gate_type : int (required)
+
Packed quantization format enum of gate_experts_weights.
+
k : int (required)
+
Number of selected experts per token.
+
normalize_routing_weights : int
+
Whether to normalize top-k routing probabilities.
+
up_type : int (required)
+
Packed quantization format enum of up_experts_weights.
+
+ +#### Inputs + +
+
input : T
+
Input hidden states [..., hidden_size].
+
router_probs : T
+
Router logits [num_rows, num_experts].
+
gate_experts_weights : T2
+
3D packed gate weight blob [num_experts, inter_size, gate_row_bytes].
+
up_experts_weights : T2
+
3D packed up weight blob [num_experts, inter_size, up_row_bytes].
+
down_experts_weights : T2
+
3D packed down weight blob [num_experts, hidden_size, down_row_bytes].
+
+ +#### Outputs + +
+
output : T
+
Output hidden states [..., hidden_size].
+
+ +#### Type Constraints + +
+
T : tensor(bfloat16), tensor(float16)
+
Activation type.
+
T2 : tensor(uint8)
+
Packed byte storage.
+
+ + ### **com.microsoft.MulInteger** Performs element-wise binary quantized multiplication (with Numpy-style broadcasting support). diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 473b66692a80d..244e760267d4e 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -605,6 +605,7 @@ The **OpSet Version** column uses the following notation: |GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |GatherBlockQuantized|*in* data:**T1**
*in* indices:**Tind**
*in* scales:**T2**
*in* zero_points:**T1**
*out* output:**T2**|1+|**T1** = tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| |GatherND|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| +|GatherPacked|*in* data:**T1**
*in* indices:**Tind**
*out* output:**T2**|1+|**T1** = tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| |Gelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |GreedySearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*in* attention_mask:**I**
*out* sequences:**I**|1+|**T** = tensor(float)| |GridSample|*in* X:**T1**
*in* Grid:**T1**
*out* Y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(float)| @@ -1131,6 +1132,7 @@ The **OpSet Version** column uses the following notation: |GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |GatedRelativePositionBias|*in* query_layer:**T**
*in* query_bias:**T**
*in* rel_pos:**T**
*in* weight:**T**
*in* bias:**T**
*in* eco_a:**T**
*in* token_offset:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |GatherBlockQuantized|*in* data:**T1**
*in* indices:**Tind**
*in* scales:**T2**
*in* zero_points:**T1**
*out* output:**T2**|1+|**T1** = tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| +|GatherPacked|*in* data:**T1**
*in* indices:**Tind**
*out* output:**T2**|1+|**T1** = tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| |Gelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |GemmFloat8|*in* A:**TA**
*in* B:**TB**
*in* C:**TC**
*in* scaleA:**TS**
*in* scaleB:**TS**
*in* scaleY:**TS**
*out* Y:**TR**|1+|**TA** = tensor(bfloat16), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e5m2)
**TB** = tensor(bfloat16), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e5m2)
**TC** = tensor(bfloat16), tensor(float), tensor(float16)
**TR** = tensor(bfloat16), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e5m2)
**TS** = tensor(float)| |GemmaRotaryEmbedding|*in* emb:**U**
*in* q:**T**
*in* q_rot:**T**
*in* k:**T**
*in* k_rot:**T**
*out* output1:**T**
*out* output2:**T**|1+|**T** = tensor(float16)
**U** = tensor(float)| @@ -1151,7 +1153,9 @@ The **OpSet Version** column uses the following notation: |MatMulBlockQuantizedFp8Weight|*in* A:**T**
*in* B:**T1**
*in* b_scale:**T2**
*in* a_scale:**T2**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn)
**T2** = tensor(float)| |MatMulBnb4|*in* A:**T1**
*in* B:**T2**
*in* absmax:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)| |MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8)| +|MatMulPacked|*in* A:**T1**
*in* B:**T2**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)| |MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|MoEPacked|*in* input:**T**
*in* router_probs:**T**
*in* gate_experts_weights:**T2**
*in* up_experts_weights:**T2**
*in* down_experts_weights:**T2**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)| |MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**QK** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)| |NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| diff --git a/onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc b/onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc index e2ad697442256..d864cef00b246 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc @@ -1,13 +1,13 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include -#include -#include +#include #include "contrib_ops/cpu/quantization/gguf_quant_tables.h" #include "core/common/common.h" #include "core/common/float16.h" +#include "core/common/inlined_containers.h" +#include "core/common/safeint.h" #include "core/framework/op_kernel.h" #include "core/framework/tensor_shape.h" #include "core/platform/threadpool.h" @@ -21,6 +21,7 @@ class GatherPacked final : public OpKernel { explicit GatherPacked(const OpKernelInfo& info) : OpKernel(info) { ORT_ENFORCE(info.GetAttr("K", &k_).IsOK() && k_ > 0, "GatherPacked requires K > 0"); ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK(), "GatherPacked requires packed_type"); + ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(ggml_type_), "GatherPacked: unsupported packed_type=", ggml_type_); row_bytes_ = gguf_quant::GgmlRowBytes(static_cast(ggml_type_), k_); ORT_ENFORCE(row_bytes_ > 0, "GatherPacked: unsupported packed_type=", ggml_type_, " or invalid K=", k_); } @@ -41,6 +42,7 @@ class GatherPacked final : public OpKernel { const int64_t num_rows = indices_shape.Size(); if (num_rows == 0) return Status::OK(); + ORT_RETURN_IF_NOT(vocab_size > 0, "GatherPacked requires a nonempty vocabulary for nonempty indices"); const uint8_t* data_ptr = data->Data(); const Tind* idx_ptr = indices->Data(); @@ -48,22 +50,23 @@ class GatherPacked final : public OpKernel { auto* tp = context->GetOperatorThreadPool(); const int type_id = static_cast(ggml_type_); - const int64_t k = k_; - const int64_t rb = row_bytes_; + const size_t k = SafeInt(k_); + const size_t rb = SafeInt(row_bytes_); + const std::ptrdiff_t rows = SafeInt(num_rows); concurrency::ThreadPool::TryBatchParallelFor( - tp, static_cast(num_rows), + tp, rows, [&](std::ptrdiff_t r) { const int64_t raw_id = static_cast(idx_ptr[r]); const int64_t tok = (raw_id >= 0 && raw_id < vocab_size) ? raw_id : 0; - const uint8_t* src_row = data_ptr + tok * rb; - T2* dst_row = out_ptr + static_cast(r) * k; + const uint8_t* src_row = data_ptr + static_cast(tok) * rb; + T2* dst_row = out_ptr + static_cast(r) * k; if constexpr (std::is_same_v) { - gguf_quant::DequantizeGgufRowCpu(src_row, type_id, k, dst_row); + ORT_ENFORCE(gguf_quant::DequantizeGgufRowCpu(src_row, type_id, k_, dst_row)); } else { - std::vector tmp(static_cast(k)); - gguf_quant::DequantizeGgufRowCpu(src_row, type_id, k, tmp.data()); - for (int64_t j = 0; j < k; ++j) { + InlinedVector tmp(k); + ORT_ENFORCE(gguf_quant::DequantizeGgufRowCpu(src_row, type_id, k_, tmp.data())); + for (size_t j = 0; j < k; ++j) { dst_row[j] = T2(tmp[j]); } } @@ -79,16 +82,16 @@ class GatherPacked final : public OpKernel { int64_t row_bytes_{0}; }; -#define REGISTER_GATHER_PACKED_CPU(T2, Tind) \ - ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ - GatherPacked, \ - kMSDomain, 1, \ - T2, Tind, \ - kCpuExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ +#define REGISTER_GATHER_PACKED_CPU(T2, Tind) \ + ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ + GatherPacked, \ + kMSDomain, 1, \ + T2, Tind, \ + kCpuExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ GatherPacked); REGISTER_GATHER_PACKED_CPU(float, int32_t) diff --git a/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h b/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h index 896dec2a42285..e35de46b7634f 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h +++ b/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h @@ -14,67 +14,176 @@ namespace contrib { namespace gguf_quant { enum GgmlType : int { - kGgmlTypeF32 = 0, - kGgmlTypeF16 = 1, - kGgmlTypeQ4_0 = 2, - kGgmlTypeQ4_1 = 3, - kGgmlTypeQ5_0 = 6, - kGgmlTypeQ5_1 = 7, - kGgmlTypeQ8_0 = 8, - kGgmlTypeQ8_1 = 9, - kGgmlTypeQ2_K = 10, - kGgmlTypeQ3_K = 11, - kGgmlTypeQ4_K = 12, - kGgmlTypeQ5_K = 13, - kGgmlTypeQ6_K = 14, - kGgmlTypeQ8_K = 15, + kGgmlTypeF32 = 0, + kGgmlTypeF16 = 1, + kGgmlTypeQ4_0 = 2, + kGgmlTypeQ4_1 = 3, + kGgmlTypeQ5_0 = 6, + kGgmlTypeQ5_1 = 7, + kGgmlTypeQ8_0 = 8, + kGgmlTypeQ8_1 = 9, + kGgmlTypeQ2_K = 10, + kGgmlTypeQ3_K = 11, + kGgmlTypeQ4_K = 12, + kGgmlTypeQ5_K = 13, + kGgmlTypeQ6_K = 14, + kGgmlTypeQ8_K = 15, kGgmlTypeIQ2_XXS = 16, - kGgmlTypeIQ2_XS = 17, + kGgmlTypeIQ2_XS = 17, kGgmlTypeIQ3_XXS = 18, - kGgmlTypeIQ1_S = 19, - kGgmlTypeIQ4_NL = 20, - kGgmlTypeIQ3_S = 21, - kGgmlTypeIQ2_S = 22, - kGgmlTypeIQ4_XS = 23, - kGgmlTypeIQ1_M = 29, - kGgmlTypeBF16 = 30, - kGgmlTypeMXFP4 = 39, - kGgmlTypeNVFP4 = 40, - kGgmlTypeQ1_0 = 41, - kGgmlTypeQ2_0 = 42, + kGgmlTypeIQ1_S = 19, + kGgmlTypeIQ4_NL = 20, + kGgmlTypeIQ3_S = 21, + kGgmlTypeIQ2_S = 22, + kGgmlTypeIQ4_XS = 23, + kGgmlTypeIQ1_M = 29, + kGgmlTypeBF16 = 30, + kGgmlTypeMXFP4 = 39, + kGgmlTypeNVFP4 = 40, + kGgmlTypeQ1_0 = 41, + kGgmlTypeQ2_0 = 42, }; +// Formats implemented by both the CPU row decoder and the CUDA sub-block decoders. +inline bool IsGgmlTypeSupported(int64_t ggml_type) { + switch (ggml_type) { + case kGgmlTypeF32: + case kGgmlTypeF16: + case kGgmlTypeBF16: + case kGgmlTypeQ4_0: + case kGgmlTypeQ5_0: + case kGgmlTypeQ8_0: + case kGgmlTypeQ5_K: + case kGgmlTypeQ6_K: + case kGgmlTypeIQ2_XXS: + case kGgmlTypeIQ2_XS: + case kGgmlTypeIQ3_XXS: + case kGgmlTypeIQ4_NL: + case kGgmlTypeIQ2_S: + case kGgmlTypeIQ4_XS: + case kGgmlTypeQ2_0: + return true; + default: + return false; + } +} + inline bool GgmlTypeBlockInfo(int ggml_type, int64_t& block_elems, int64_t& block_bytes) { switch (ggml_type) { - case kGgmlTypeF32: block_elems = 1; block_bytes = 4; return true; - case kGgmlTypeF16: block_elems = 1; block_bytes = 2; return true; - case kGgmlTypeBF16: block_elems = 1; block_bytes = 2; return true; - case kGgmlTypeQ4_0: block_elems = 32; block_bytes = 18; return true; - case kGgmlTypeQ4_1: block_elems = 32; block_bytes = 20; return true; - case kGgmlTypeQ5_0: block_elems = 32; block_bytes = 22; return true; - case kGgmlTypeQ5_1: block_elems = 32; block_bytes = 24; return true; - case kGgmlTypeQ8_0: block_elems = 32; block_bytes = 34; return true; - case kGgmlTypeQ8_1: block_elems = 32; block_bytes = 36; return true; - case kGgmlTypeQ2_K: block_elems = 256; block_bytes = 84; return true; - case kGgmlTypeQ3_K: block_elems = 256; block_bytes = 110; return true; - case kGgmlTypeQ4_K: block_elems = 256; block_bytes = 144; return true; - case kGgmlTypeQ5_K: block_elems = 256; block_bytes = 176; return true; - case kGgmlTypeQ6_K: block_elems = 256; block_bytes = 210; return true; - case kGgmlTypeQ8_K: block_elems = 256; block_bytes = 292; return true; - case kGgmlTypeIQ2_XXS: block_elems = 256; block_bytes = 66; return true; - case kGgmlTypeIQ2_XS: block_elems = 256; block_bytes = 74; return true; - case kGgmlTypeIQ3_XXS: block_elems = 256; block_bytes = 98; return true; - case kGgmlTypeIQ1_S: block_elems = 256; block_bytes = 50; return true; - case kGgmlTypeIQ4_NL: block_elems = 32; block_bytes = 18; return true; - case kGgmlTypeIQ3_S: block_elems = 256; block_bytes = 110; return true; - case kGgmlTypeIQ2_S: block_elems = 256; block_bytes = 82; return true; - case kGgmlTypeIQ4_XS: block_elems = 256; block_bytes = 136; return true; - case kGgmlTypeIQ1_M: block_elems = 256; block_bytes = 56; return true; - case kGgmlTypeMXFP4: block_elems = 32; block_bytes = 17; return true; - case kGgmlTypeNVFP4: block_elems = 64; block_bytes = 36; return true; - case kGgmlTypeQ1_0: block_elems = 128; block_bytes = 18; return true; - case kGgmlTypeQ2_0: block_elems = 64; block_bytes = 18; return true; - default: return false; + case kGgmlTypeF32: + block_elems = 1; + block_bytes = 4; + return true; + case kGgmlTypeF16: + block_elems = 1; + block_bytes = 2; + return true; + case kGgmlTypeBF16: + block_elems = 1; + block_bytes = 2; + return true; + case kGgmlTypeQ4_0: + block_elems = 32; + block_bytes = 18; + return true; + case kGgmlTypeQ4_1: + block_elems = 32; + block_bytes = 20; + return true; + case kGgmlTypeQ5_0: + block_elems = 32; + block_bytes = 22; + return true; + case kGgmlTypeQ5_1: + block_elems = 32; + block_bytes = 24; + return true; + case kGgmlTypeQ8_0: + block_elems = 32; + block_bytes = 34; + return true; + case kGgmlTypeQ8_1: + block_elems = 32; + block_bytes = 36; + return true; + case kGgmlTypeQ2_K: + block_elems = 256; + block_bytes = 84; + return true; + case kGgmlTypeQ3_K: + block_elems = 256; + block_bytes = 110; + return true; + case kGgmlTypeQ4_K: + block_elems = 256; + block_bytes = 144; + return true; + case kGgmlTypeQ5_K: + block_elems = 256; + block_bytes = 176; + return true; + case kGgmlTypeQ6_K: + block_elems = 256; + block_bytes = 210; + return true; + case kGgmlTypeQ8_K: + block_elems = 256; + block_bytes = 292; + return true; + case kGgmlTypeIQ2_XXS: + block_elems = 256; + block_bytes = 66; + return true; + case kGgmlTypeIQ2_XS: + block_elems = 256; + block_bytes = 74; + return true; + case kGgmlTypeIQ3_XXS: + block_elems = 256; + block_bytes = 98; + return true; + case kGgmlTypeIQ1_S: + block_elems = 256; + block_bytes = 50; + return true; + case kGgmlTypeIQ4_NL: + block_elems = 32; + block_bytes = 18; + return true; + case kGgmlTypeIQ3_S: + block_elems = 256; + block_bytes = 110; + return true; + case kGgmlTypeIQ2_S: + block_elems = 256; + block_bytes = 82; + return true; + case kGgmlTypeIQ4_XS: + block_elems = 256; + block_bytes = 136; + return true; + case kGgmlTypeIQ1_M: + block_elems = 256; + block_bytes = 56; + return true; + case kGgmlTypeMXFP4: + block_elems = 32; + block_bytes = 17; + return true; + case kGgmlTypeNVFP4: + block_elems = 64; + block_bytes = 36; + return true; + case kGgmlTypeQ1_0: + block_elems = 128; + block_bytes = 18; + return true; + case kGgmlTypeQ2_0: + block_elems = 64; + block_bytes = 18; + return true; + default: + return false; } } @@ -92,26 +201,27 @@ inline int64_t GgmlRowBytes(int ggml_type, int64_t K) { #else #define GGML_TABLE_BEGIN(type, name, size) static constexpr type name[size] = { #endif -#define GGML_TABLE_END() }; +#define GGML_TABLE_END() \ + } \ + ; #endif GGML_TABLE_BEGIN(uint8_t, kmask_iq2xs, 8) - 1, 2, 4, 8, 16, 32, 64, 128 -GGML_TABLE_END() +1, 2, 4, 8, 16, 32, 64, 128 GGML_TABLE_END() -GGML_TABLE_BEGIN(uint8_t, ksigns_iq2xs, 128) - 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, - 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, - 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, - 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, - 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, - 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, - 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, + GGML_TABLE_BEGIN(uint8_t, ksigns_iq2xs, 128) 0, + 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, + 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, + 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, + 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, + 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, + 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, + 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, -GGML_TABLE_END() + GGML_TABLE_END() -GGML_TABLE_BEGIN(uint64_t, ksigns64, 128) - 0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, + GGML_TABLE_BEGIN(uint64_t, ksigns64, 128) 0x0000000000000000, + 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, 0xff00000000ff0000, 0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff, 0xff000000ff000000, 0x00000000ff0000ff, 0x00000000ff00ff00, 0xff000000ff00ffff, 0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00, 0x00000000ffffffff, @@ -143,11 +253,10 @@ GGML_TABLE_BEGIN(uint64_t, ksigns64, 128) 0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff, 0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff, 0xffffffffffff0000, 0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff, -GGML_TABLE_END() - + GGML_TABLE_END() -GGML_TABLE_BEGIN(uint64_t, iq2xxs_grid, 256) - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + GGML_TABLE_BEGIN(uint64_t, iq2xxs_grid, 256) 0x0808080808080808, + 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, @@ -211,10 +320,10 @@ GGML_TABLE_BEGIN(uint64_t, iq2xxs_grid, 256) 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, -GGML_TABLE_END() + GGML_TABLE_END() -GGML_TABLE_BEGIN(uint64_t, iq2xs_grid, 512) - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + GGML_TABLE_BEGIN(uint64_t, iq2xs_grid, 512) 0x0808080808080808, + 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, @@ -342,10 +451,10 @@ GGML_TABLE_BEGIN(uint64_t, iq2xs_grid, 512) 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, 0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b, -GGML_TABLE_END() + GGML_TABLE_END() -GGML_TABLE_BEGIN(uint64_t, iq2s_grid, 1024) - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + GGML_TABLE_BEGIN(uint64_t, iq2s_grid, 1024) 0x0808080808080808, + 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, @@ -601,10 +710,10 @@ GGML_TABLE_BEGIN(uint64_t, iq2s_grid, 1024) 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, 0x2b2b2b1919191908, 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b, -GGML_TABLE_END() + GGML_TABLE_END() -GGML_TABLE_BEGIN(uint32_t, iq3xxs_grid, 256) - 0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, + GGML_TABLE_BEGIN(uint32_t, iq3xxs_grid, 256) 0x04040404, + 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, 0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, 0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, 0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, @@ -636,10 +745,10 @@ GGML_TABLE_BEGIN(uint32_t, iq3xxs_grid, 256) 0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, 0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, -GGML_TABLE_END() + GGML_TABLE_END() -GGML_TABLE_BEGIN(uint32_t, iq3s_grid, 512) - 0x01010101, 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305, + GGML_TABLE_BEGIN(uint32_t, iq3s_grid, 512) 0x01010101, + 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305, 0x01010309, 0x0101030d, 0x01010501, 0x01010503, 0x0101050b, 0x01010707, 0x01010901, 0x01010905, 0x0101090b, 0x0101090f, 0x01010b03, 0x01010b07, 0x01010d01, 0x01010d05, 0x01010f03, 0x01010f09, 0x01010f0f, 0x01030101, 0x01030103, 0x01030105, 0x01030109, 0x01030301, 0x01030303, 0x0103030b, @@ -703,18 +812,19 @@ GGML_TABLE_BEGIN(uint32_t, iq3s_grid, 512) 0x0f030509, 0x0f030907, 0x0f03090b, 0x0f050103, 0x0f050109, 0x0f050301, 0x0f05030d, 0x0f050503, 0x0f050701, 0x0f050b03, 0x0f070105, 0x0f070705, 0x0f07070b, 0x0f070b07, 0x0f090103, 0x0f09010b, 0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101, -GGML_TABLE_END() + GGML_TABLE_END() -// TODO: fix name to kvalues_iq4_nl -GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16) - -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113, -GGML_TABLE_END() + // TODO: fix name to kvalues_iq4_nl + GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16) - + 127, + -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113, + GGML_TABLE_END() -// e2m1 values (doubled), shared by MXFP4 and NVFP4 -// ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf -GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16) - 0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12, -GGML_TABLE_END() + // e2m1 values (doubled), shared by MXFP4 and NVFP4 + // ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf + GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16) 0, + 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12, + GGML_TABLE_END() #define kvalues_mxfp4 kvalues_fp4 #undef GGML_TABLE_BEGIN @@ -722,7 +832,7 @@ GGML_TABLE_END() #if !defined(__CUDA_ARCH__) -inline float Fp16BitsToFloatCpu(uint16_t bits) { + inline float Fp16BitsToFloatCpu(uint16_t bits) { return MLFloat16::FromBits(bits).ToFloat(); } diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc index 8a88c864ba786..7a2553bac4429 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc @@ -13,15 +13,15 @@ namespace onnxruntime { namespace contrib { namespace cuda { -#define REGISTER_MOE_PACKED_CUDA(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - MoEPacked, \ - kMSDomain, 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()),\ +#define REGISTER_MOE_PACKED_CUDA(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MoEPacked, \ + kMSDomain, 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ MoEPacked); REGISTER_MOE_PACKED_CUDA(MLFloat16) @@ -39,6 +39,9 @@ MoEPacked::MoEPacked(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(info.GetAttr("gate_type", &gate_type_).IsOK(), "MoEPacked requires gate_type"); ORT_ENFORCE(info.GetAttr("up_type", &up_type_).IsOK(), "MoEPacked requires up_type"); ORT_ENFORCE(info.GetAttr("down_type", &down_type_).IsOK(), "MoEPacked requires down_type"); + ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(gate_type_), "MoEPacked: unsupported gate_type=", gate_type_); + ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(up_type_), "MoEPacked: unsupported up_type=", up_type_); + ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(down_type_), "MoEPacked: unsupported down_type=", down_type_); } template @@ -68,6 +71,7 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(rp_shape[0] == num_rows, "MoEPacked router_probs row count mismatch"); ORT_RETURN_IF_NOT(num_experts > 0 && num_experts <= 512, "MoEPacked supports up to 512 experts"); + ORT_RETURN_IF_NOT(k_ <= num_experts, "MoEPacked requires k <= num_experts"); ORT_RETURN_IF_NOT(gw_shape[0] == num_experts && uw_shape[0] == num_experts && dw_shape[0] == num_experts, "MoEPacked expert count mismatch"); ORT_RETURN_IF_NOT(uw_shape[1] == inter_size && dw_shape[1] == hidden_size, @@ -78,11 +82,15 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { const int64_t exp_gate_rb = gguf_quant::GgmlRowBytes(static_cast(gate_type_), hidden_size); const int64_t exp_up_rb = gguf_quant::GgmlRowBytes(static_cast(up_type_), hidden_size); const int64_t exp_down_rb = gguf_quant::GgmlRowBytes(static_cast(down_type_), inter_size); - ORT_RETURN_IF_NOT(gw_shape[2] == exp_gate_rb && uw_shape[2] == exp_up_rb && dw_shape[2] == exp_down_rb, + ORT_RETURN_IF_NOT(exp_gate_rb > 0 && exp_up_rb > 0 && exp_down_rb > 0 && + gw_shape[2] == exp_gate_rb && uw_shape[2] == exp_up_rb && dw_shape[2] == exp_down_rb, "MoEPacked expert weight row_bytes mismatch"); Tensor* output = context->Output(0, in_shape); if (num_rows == 0) return Status::OK(); + constexpr bool is_bf16 = std::is_same_v; + ORT_RETURN_IF_NOT(!is_bf16 || GetDeviceProp().major >= 8, + "MoEPacked BF16 requires compute capability 8.0 or higher"); // Process in chunks of at most 2048 rows so workspace stays compact during large prefills. constexpr int64_t kMaxChunkRows = 2048; @@ -104,7 +112,6 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { const size_t ws_bytes = GetMoEPackedWorkspaceBytes(max_params, sizeof(T)); auto workspace = GetScratchBuffer(ws_bytes, GetComputeStream(context)); - constexpr bool is_bf16 = std::is_same_v; const uint8_t* in_bytes = static_cast(input->DataRaw()); const uint8_t* rp_bytes = static_cast(router_probs->DataRaw()); diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu index 002ebe2fbc1e7..92672b730be20 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu @@ -756,7 +756,7 @@ __global__ __launch_bounds__(256) void MoEPackedFc2Bf16TcKernel( } // Fallback SIMT FC1 + SwiGLU kernel (used for FP16). -template +template __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel( const CudaT* __restrict__ input, const uint8_t* __restrict__ gate_weights, @@ -776,14 +776,16 @@ __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel( const int count = tile.count < 8 ? tile.count : 8; const int expert = tile.expert; - const int total_elems = count * p.hidden_size; - for (int idx = static_cast(threadIdx.x); idx < total_elems; idx += 256) { - const int m = idx / p.hidden_size; - const int col = idx - m * p.hidden_size; - const int token_row = sorted_routes[tile.start_route + m] / p.k; - s_x[idx] = input[static_cast(token_row) * p.hidden_size + col]; + if constexpr (kUseShared) { + const int total_elems = count * p.hidden_size; + for (int idx = static_cast(threadIdx.x); idx < total_elems; idx += 256) { + const int m = idx / p.hidden_size; + const int col = idx - m * p.hidden_size; + const int token_row = sorted_routes[tile.start_route + m] / p.k; + s_x[idx] = input[static_cast(token_row) * p.hidden_size + col]; + } + __syncthreads(); } - __syncthreads(); const int warp = static_cast(threadIdx.x) >> 5; const int lane = static_cast(threadIdx.x) & 31; @@ -808,7 +810,13 @@ __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel( const int k_offset = h_idx * 16; for (int m = 0; m < count; ++m) { - const CudaT* xm = s_x + m * p.hidden_size + k_offset; + const CudaT* xm; + if constexpr (kUseShared) { + xm = s_x + m * p.hidden_size + k_offset; + } else { + const int token_row = sorted_routes[tile.start_route + m] / p.k; + xm = input + static_cast(token_row) * p.hidden_size + k_offset; + } #pragma unroll for (int c = 0; c < 16; ++c) { const float xv = gguf_cuda::CudaTypeToFloat(xm[c]); @@ -839,7 +847,7 @@ __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel( } // Fallback SIMT FC2 kernel (used for FP16). -template +template __global__ __launch_bounds__(256) void MoEPackedFc2SimtKernel( const CudaT* __restrict__ fc1_out, const uint8_t* __restrict__ down_weights, @@ -858,13 +866,15 @@ __global__ __launch_bounds__(256) void MoEPackedFc2SimtKernel( const int count = tile.count < 8 ? tile.count : 8; const int expert = tile.expert; - const int total_elems = count * p.inter_size; - for (int idx = static_cast(threadIdx.x); idx < total_elems; idx += 256) { - const int m = idx / p.inter_size; - const int col = idx - m * p.inter_size; - s_act[idx] = fc1_out[static_cast(sorted_routes[tile.start_route + m]) * p.inter_size + col]; + if constexpr (kUseShared) { + const int total_elems = count * p.inter_size; + for (int idx = static_cast(threadIdx.x); idx < total_elems; idx += 256) { + const int m = idx / p.inter_size; + const int col = idx - m * p.inter_size; + s_act[idx] = fc1_out[static_cast(sorted_routes[tile.start_route + m]) * p.inter_size + col]; + } + __syncthreads(); } - __syncthreads(); const int warp = static_cast(threadIdx.x) >> 5; const int lane = static_cast(threadIdx.x) & 31; @@ -885,7 +895,12 @@ __global__ __launch_bounds__(256) void MoEPackedFc2SimtKernel( const int k_offset = h_idx * 16; for (int m = 0; m < count; ++m) { - const CudaT* am = s_act + m * p.inter_size + k_offset; + const CudaT* am; + if constexpr (kUseShared) { + am = s_act + m * p.inter_size + k_offset; + } else { + am = fc1_out + static_cast(sorted_routes[tile.start_route + m]) * p.inter_size + k_offset; + } #pragma unroll for (int c = 0; c < 16; ++c) { down_acc[m] += d_vals[c] * gguf_cuda::CudaTypeToFloat(am[c]); @@ -1234,16 +1249,28 @@ Status DispatchMoEPackedCuda( } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); } else { + // Keep staging within CUDA's default shared-memory limit; larger shapes read activations directly. + constexpr size_t kMaxStagingBytes = 48 * 1024; const size_t fc1_smem = static_cast(8) * params.hidden_size * sizeof(CudaT); const dim3 fc1_grid(static_cast((params.inter_size + 7) / 8), static_cast(max_tiles)); - MoEPackedFc1SwiGLUSimtKernel<<>>( - input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + if (fc1_smem <= kMaxStagingBytes) { + MoEPackedFc1SwiGLUSimtKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + } else { + MoEPackedFc1SwiGLUSimtKernel<<>>( + input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); const size_t fc2_smem = static_cast(8) * params.inter_size * sizeof(CudaT); const dim3 fc2_grid(static_cast((params.hidden_size + 7) / 8), static_cast(max_tiles)); - MoEPackedFc2SimtKernel<<>>( - fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + if (fc2_smem <= kMaxStagingBytes) { + MoEPackedFc2SimtKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + } else { + MoEPackedFc2SimtKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); } diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc index 0fd7c1d1f0613..75fd583d33d60 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc @@ -12,16 +12,16 @@ namespace onnxruntime { namespace contrib { namespace cuda { -#define REGISTER_GATHER_PACKED_CUDA(T2, Tind) \ - ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ - GatherPacked, \ - kMSDomain, 1, \ - T2, Tind, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ +#define REGISTER_GATHER_PACKED_CUDA(T2, Tind) \ + ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ + GatherPacked, \ + kMSDomain, 1, \ + T2, Tind, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ GatherPacked); REGISTER_GATHER_PACKED_CUDA(float, int32_t) @@ -37,6 +37,7 @@ template GatherPacked::GatherPacked(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(info.GetAttr("K", &k_).IsOK() && k_ > 0, "GatherPacked requires K > 0"); ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK(), "GatherPacked requires packed_type"); + ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(ggml_type_), "GatherPacked: unsupported packed_type=", ggml_type_); row_bytes_ = gguf_quant::GgmlRowBytes(static_cast(ggml_type_), k_); ORT_ENFORCE(row_bytes_ > 0, "GatherPacked: unsupported packed_type=", ggml_type_, " or invalid K=", k_); } @@ -58,6 +59,7 @@ Status GatherPacked::ComputeInternal(OpKernelContext* context) const { const int64_t num_rows = indices_shape.Size(); if (num_rows == 0) return Status::OK(); + ORT_RETURN_IF_NOT(vocab_size > 0, "GatherPacked requires a nonempty vocabulary for nonempty indices"); constexpr int out_dtype = std::is_same_v ? 16 : (std::is_same_v ? 10 : 1); diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu index f897f25967b03..9435dc23c1ece 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu @@ -28,6 +28,22 @@ __global__ __launch_bounds__(256) void GatherPackedKernel( const uint8_t* row_ptr = data + tok * row_bytes; CudaT* dst_row = output + r * static_cast(k); + // Unquantized rows can end partway through a sub-block. + if ((k % 32) != 0) { + for (int col = threadIdx.x; col < k; col += blockDim.x) { + float value; + if (ggml_type == gguf_quant::kGgmlTypeF32) { + value = __ldg(reinterpret_cast(row_ptr) + col); + } else { + const uint16_t bits = __ldg(reinterpret_cast(row_ptr) + col); + value = ggml_type == gguf_quant::kGgmlTypeF16 ? gguf_cuda::Fp16ToFloatDev(bits) + : gguf_cuda::Bf16ToFloatDev(bits); + } + dst_row[col] = gguf_cuda::FloatToCudaType(value); + } + return; + } + const int total_quarters = k / 8; for (int q_idx = threadIdx.x; q_idx < total_quarters; q_idx += blockDim.x) { float vals[8]; @@ -53,7 +69,6 @@ Status DispatchGatherPackedByIndex( int64_t row_bytes, int ggml_type) { if (num_rows == 0) return Status::OK(); - ORT_RETURN_IF_NOT((k % 32) == 0, "GatherPacked CUDA requires K to be a multiple of 32"); const dim3 grid(static_cast(num_rows)); const dim3 block(256); if (indices_are_int64) { diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc index 0d7e971759542..817b19073c9eb 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc @@ -12,15 +12,15 @@ namespace onnxruntime { namespace contrib { namespace cuda { -#define REGISTER_MATMUL_PACKED_CUDA(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - MatMulPacked, \ - kMSDomain, 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()),\ +#define REGISTER_MATMUL_PACKED_CUDA(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MatMulPacked, \ + kMSDomain, 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ MatMulPacked); REGISTER_MATMUL_PACKED_CUDA(MLFloat16) @@ -34,6 +34,7 @@ MatMulPacked::MatMulPacked(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(info.GetAttr("N", &n_).IsOK() && n_ > 0, "MatMulPacked requires N > 0"); ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK() && ggml_type_ >= 0, "MatMulPacked requires packed_type >= 0"); + ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(ggml_type_), "MatMulPacked: unsupported packed_type=", ggml_type_); } template @@ -69,6 +70,8 @@ Status MatMulPacked::ComputeInternal(OpKernelContext* context) const { if (m == 0) return Status::OK(); constexpr bool is_bf16 = std::is_same_v; + ORT_RETURN_IF_NOT(!is_bf16 || GetDeviceProp().major >= 8, + "MatMulPacked BF16 requires compute capability 8.0 or higher"); IAllocatorUniquePtr scratch_w; if (m > 8) { scratch_w = GetScratchBuffer(static_cast(n_ * k_), GetComputeStream(context)); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu index c1df32db6f813..3ff0eeef331fc 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu @@ -59,21 +59,19 @@ __global__ __launch_bounds__(KW * 32) void MatMulPackedBf16TcWarpKSplitKernel( : make_uint4(0u, 0u, 0u, 0u); #pragma unroll for (int m = 0; m < MT; ++m) { + float w_lo[8] = {}; + float w_hi[8] = {}; if (valid_lo[m]) { - float w_lo[8]; - float w_hi[8]; gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_lo[m], p.ggml_type, sb32_idx, t, w_lo); - if (valid_hi[m]) { - gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_hi[m], p.ggml_type, sb32_idx, t, w_hi); - } else { -#pragma unroll - for (int j = 0; j < 8; ++j) w_hi[j] = 0.0f; - } - const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(w_lo); - const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(w_hi); - gguf_cuda::MmaBf16(acc[m], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); - gguf_cuda::MmaBf16(acc[m], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); } + if (valid_hi[m]) { + gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_hi[m], p.ggml_type, sb32_idx, t, w_hi); + } + // mma.sync requires every lane to participate, including partial output tiles. + const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(w_lo); + const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(w_hi); + gguf_cuda::MmaBf16(acc[m], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); + gguf_cuda::MmaBf16(acc[m], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); } } diff --git a/onnxruntime/test/python/transformers/test_gguf_packed_ops.py b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py index fb345292d6dca..46d5f5e29181e 100644 --- a/onnxruntime/test/python/transformers/test_gguf_packed_ops.py +++ b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py @@ -36,7 +36,139 @@ def _bf16_u16_to_f32(arr_u16: np.ndarray) -> np.ndarray: return (arr_u16.astype(np.uint32) << 16).view(np.float32) -def _make_synthetic_gguf_rows(rng: np.random.Generator, rows: int, k: int, gtype: int, bel: int, bbytes: int) -> np.ndarray: +def _round_to_dtype(arr, dtype): + if dtype == TensorProto.BFLOAT16: + return _bf16_u16_to_f32(_f32_to_bf16_u16(arr)) + return arr.astype(np.float16).astype(np.float32) + + +def _session(model, provider): + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + if provider == "CUDAExecutionProvider": + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + session = ort.InferenceSession(model.SerializeToString(), sess_options=options, providers=[provider]) + session.disable_fallback() + if provider not in session.get_providers(): + raise RuntimeError(f"Required provider {provider} was not initialized") + return session + + +def _model(nodes, inputs, outputs): + return helper.make_model( + helper.make_graph(nodes, "packed_ops", inputs, outputs), + opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)], + ) + + +def _gather_model(vocab, k, row_bytes, gtype, dtype=TensorProto.FLOAT, index_type=TensorProto.INT64): + return _model( + [ + helper.make_node( + "GatherPacked", + ["W", "I"], + ["Y_typed"], + domain="com.microsoft", + K=k, + packed_type=gtype, + output_dtype=int(dtype), + ), + helper.make_node("Cast", ["Y_typed"], ["Y"], to=int(TensorProto.FLOAT)), + ], + [ + helper.make_tensor_value_info("W", TensorProto.UINT8, [vocab, row_bytes]), + helper.make_tensor_value_info("I", index_type, ["rows"]), + ], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, ["rows", k])], + ) + + +def _dequantize_weights(raw, k, gtype): + rows = raw.reshape(-1, raw.shape[-1]) + session = _session(_gather_model(rows.shape[0], k, rows.shape[1], gtype), "CPUExecutionProvider") + result = session.run(None, {"W": rows, "I": np.arange(rows.shape[0], dtype=np.int64)})[0] + return result.reshape(*raw.shape[:-1], k) + + +def _matmul_model(k, n, row_bytes, gtype): + return _model( + [ + helper.make_node("Cast", ["A_f32"], ["A_bf16"], to=int(TensorProto.BFLOAT16)), + helper.make_node( + "MatMulPacked", + ["A_bf16", "B"], + ["Y_bf16"], + domain="com.microsoft", + K=k, + N=n, + packed_type=gtype, + ), + helper.make_node("Cast", ["Y_bf16"], ["Y_f32"], to=int(TensorProto.FLOAT)), + ], + [ + helper.make_tensor_value_info("A_f32", TensorProto.FLOAT, ["M", k]), + helper.make_tensor_value_info("B", TensorProto.UINT8, [n, row_bytes]), + ], + [helper.make_tensor_value_info("Y_f32", TensorProto.FLOAT, ["M", n])], + ) + + +def _moe_model(num_experts, hidden_size, inter_size, top_k, dtype, types, row_bytes, normalize): + return _model( + [ + helper.make_node("Cast", ["X_f32"], ["X"], to=int(dtype)), + helper.make_node("Cast", ["R_f32"], ["R"], to=int(dtype)), + helper.make_node( + "MoEPacked", + ["X", "R", "G", "U", "D"], + ["Y"], + domain="com.microsoft", + k=top_k, + normalize_routing_weights=int(normalize), + activation_type="swiglu", + gate_type=types[0], + up_type=types[1], + down_type=types[2], + ), + helper.make_node("Cast", ["Y"], ["Y_f32"], to=int(TensorProto.FLOAT)), + ], + [ + helper.make_tensor_value_info("X_f32", TensorProto.FLOAT, ["M", hidden_size]), + helper.make_tensor_value_info("R_f32", TensorProto.FLOAT, ["M", num_experts]), + helper.make_tensor_value_info("G", TensorProto.UINT8, [num_experts, inter_size, row_bytes[0]]), + helper.make_tensor_value_info("U", TensorProto.UINT8, [num_experts, inter_size, row_bytes[1]]), + helper.make_tensor_value_info("D", TensorProto.UINT8, [num_experts, hidden_size, row_bytes[2]]), + ], + [helper.make_tensor_value_info("Y_f32", TensorProto.FLOAT, ["M", hidden_size])], + ) + + +def _moe_reference(x, logits, gate, up, down, top_k, dtype, normalize): + x = _round_to_dtype(x, dtype) + logits = _round_to_dtype(logits, dtype) + if dtype == TensorProto.BFLOAT16: + gate, up, down = (_round_to_dtype(w, dtype) for w in (gate, up, down)) + probabilities = np.exp(logits - logits.max(axis=1, keepdims=True)) + probabilities /= probabilities.sum(axis=1, keepdims=True) + # Equal logits route to the lower expert index. + experts = np.argsort(-logits, axis=1, kind="stable")[:, :top_k] + scales = np.take_along_axis(probabilities, experts, axis=1) + if normalize: + scales /= scales.sum(axis=1, keepdims=True) + result = np.zeros_like(x) + for row in range(x.shape[0]): + for rank, expert in enumerate(experts[row]): + g = _round_to_dtype(x[row] @ gate[expert].T, dtype) + u = _round_to_dtype(x[row] @ up[expert].T, dtype) + activation = _round_to_dtype(u * (g / (1.0 + np.exp(-g))), dtype) + expert_output = _round_to_dtype(activation @ down[expert].T, dtype) + result[row] += scales[row, rank] * expert_output + return _round_to_dtype(result, dtype) + + +def _make_synthetic_gguf_rows( + rng: np.random.Generator, rows: int, k: int, gtype: int, bel: int, bbytes: int +) -> np.ndarray: row_bytes = (k // bel) * bbytes if gtype == 0: return (rng.standard_normal((rows, k), dtype=np.float32) * 0.1).view(np.uint8).reshape(rows, row_bytes) @@ -57,6 +189,47 @@ def _make_synthetic_gguf_rows(rng: np.random.Generator, rows: int, k: int, gtype return raw +class TestGatherPackedValidation(unittest.TestCase): + def test_empty_vocabulary(self): + for provider in ("CPUExecutionProvider", "CUDAExecutionProvider"): + if provider not in ort.get_available_providers(): + continue + with self.subTest(provider=provider): + session = _session(_gather_model(0, 32, 128, 0), provider) + weights = np.empty((0, 128), dtype=np.uint8) + empty = session.run(None, {"W": weights, "I": np.array([], dtype=np.int64)})[0] + self.assertEqual(empty.shape, (0, 32)) + with self.assertRaisesRegex(Exception, "nonempty vocabulary"): + session.run(None, {"W": weights, "I": np.array([0], dtype=np.int64)}) + + def test_reject_unsupported_formats(self): + unsupported = [ + (3, 20), + (7, 24), + (9, 36), + (10, 84), + (11, 110), + (12, 144), + (15, 292), + (19, 50), + (21, 110), + (29, 56), + (39, 17), + (40, 36), + (41, 18), + (1 << 32, 4), + ] + for provider in ("CPUExecutionProvider", "CUDAExecutionProvider"): + if provider not in ort.get_available_providers(): + continue + for gtype, row_bytes in unsupported: + with ( + self.subTest(provider=provider, gtype=gtype), + self.assertRaisesRegex(Exception, "unsupported packed_type"), + ): + _session(_gather_model(1, 256, row_bytes, gtype), provider) + + @unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDAExecutionProvider is required") class TestGgufPackedOps(unittest.TestCase): def test_gather_packed_cpu_vs_cuda_all_types(self): @@ -88,141 +261,132 @@ def test_gather_packed_cpu_vs_cuda_all_types(self): model = helper.make_model( graph, opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)] ) - sess_cpu = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) - sess_cuda = ort.InferenceSession(model.SerializeToString(), providers=["CUDAExecutionProvider"]) + sess_cpu = _session(model, "CPUExecutionProvider") + sess_cuda = _session(model, "CUDAExecutionProvider") y_cpu = sess_cpu.run(None, {"W": w_raw, "I": indices})[0] y_cuda = sess_cuda.run(None, {"W": w_raw, "I": indices})[0] np.testing.assert_allclose(y_cpu, y_cuda, rtol=0.0, atol=0.0, err_msg=f"Mismatch for ggml_type={gtype}") + def test_gather_packed_unquantized_short_rows(self): + rng = np.random.default_rng(43) + for gtype, bbytes in ((0, 4), (1, 2), (30, 2)): + for k in (1, 8, 33): + raw = _make_synthetic_gguf_rows(rng, 3, k, gtype, 1, bbytes) + for dtype in (TensorProto.FLOAT, TensorProto.FLOAT16, TensorProto.BFLOAT16): + for index_type, numpy_type in ((TensorProto.INT32, np.int32), (TensorProto.INT64, np.int64)): + with self.subTest(gtype=gtype, k=k, dtype=dtype, index_type=index_type): + model = _gather_model(3, k, k * bbytes, gtype, dtype, index_type) + inputs = {"W": raw, "I": np.array([2, 0, 1, -1, 3], dtype=numpy_type)} + cpu = _session(model, "CPUExecutionProvider").run(None, inputs)[0] + cuda = _session(model, "CUDAExecutionProvider").run(None, inputs)[0] + np.testing.assert_array_equal(cuda, cpu) + def test_matmul_packed_cuda_all_types(self): rng = np.random.default_rng(123) - k = 256 - n = 32 - indices = np.arange(n, dtype=np.int64) + k, n = 256, 32 for gtype, bel, bbytes in GGUF_TYPES_INFO: - row_bytes = (k // bel) * bbytes - w_raw = _make_synthetic_gguf_rows(rng, n, k, gtype, bel, bbytes) - - # Dequantize reference via CPU GatherPacked - g_node = helper.make_node( - "GatherPacked", ["W", "I"], ["W_f32"], domain="com.microsoft", K=k, packed_type=gtype, output_dtype=1 - ) - g_model = helper.make_model( - helper.make_graph( - [g_node], - "deq", - [ - helper.make_tensor_value_info("W", TensorProto.UINT8, [n, row_bytes]), - helper.make_tensor_value_info("I", TensorProto.INT64, [n]), - ], - [helper.make_tensor_value_info("W_f32", TensorProto.FLOAT, [n, k])], - ), - opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)], - ) - w_f32 = ort.InferenceSession(g_model.SerializeToString(), providers=["CPUExecutionProvider"]).run( - None, {"W": w_raw, "I": indices} - )[0] - - # Build MatMulPacked model with Cast(FLOAT -> BFLOAT16) -> MatMulPacked -> Cast(BFLOAT16 -> FLOAT) - nodes = [ - helper.make_node("Cast", ["A_f32"], ["A_bf16"], to=int(TensorProto.BFLOAT16)), - helper.make_node( - "MatMulPacked", - ["A_bf16", "B"], - ["Y_bf16"], - domain="com.microsoft", - K=k, - N=n, - packed_type=gtype, - ), - helper.make_node("Cast", ["Y_bf16"], ["Y_f32"], to=int(TensorProto.FLOAT)), - ] - mm_model = helper.make_model( - helper.make_graph( - nodes, - "mm_packed", - [ - helper.make_tensor_value_info("A_f32", TensorProto.FLOAT, ["M", k]), - helper.make_tensor_value_info("B", TensorProto.UINT8, [n, row_bytes]), - ], - [helper.make_tensor_value_info("Y_f32", TensorProto.FLOAT, ["M", n])], - ), - opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)], - ) - w_bf16 = _bf16_u16_to_f32(_f32_to_bf16_u16(w_f32)) - sess = ort.InferenceSession(mm_model.SerializeToString(), providers=["CUDAExecutionProvider"]) + raw = _make_synthetic_gguf_rows(rng, n, k, gtype, bel, bbytes) + weights = _dequantize_weights(raw, k, gtype) + session = _session(_matmul_model(k, n, raw.shape[1], gtype), "CUDAExecutionProvider") for m in (1, 7, 8, 32): - a_f32 = _bf16_u16_to_f32(_f32_to_bf16_u16(rng.standard_normal((m, k), dtype=np.float32) * 0.2)) - y_got = sess.run(None, {"A_f32": a_f32, "B": w_raw})[0] - w_target = w_f32 if m <= 8 else w_bf16 - y_ref = _bf16_u16_to_f32(_f32_to_bf16_u16(a_f32 @ w_target.T)) - np.testing.assert_allclose( - y_got, y_ref, rtol=1e-2, atol=1e-2, err_msg=f"MatMulPacked mismatch for ggml_type={gtype}, M={m}" - ) + with self.subTest(gtype=gtype, m=m): + a = _round_to_dtype(rng.standard_normal((m, k), dtype=np.float32) * 0.2, TensorProto.BFLOAT16) + got = session.run(None, {"A_f32": a, "B": raw})[0] + target = weights if m <= 8 else _round_to_dtype(weights, TensorProto.BFLOAT16) + ref = _round_to_dtype(a @ target.T, TensorProto.BFLOAT16) + np.testing.assert_allclose(got, ref, rtol=1e-2, atol=1e-2) - def test_moe_packed_cuda_decode_and_prefill(self): + def test_matmul_packed_partial_tensor_core_tile(self): + rng = np.random.default_rng(124) + k, n = 32, 65537 + # F32 values must be rounded to BF16 before the Tensor Core multiplication. + raw = _make_synthetic_gguf_rows(rng, n, k, 0, 1, 4) + weights = _round_to_dtype(raw.view(np.float32).reshape(n, k), TensorProto.BFLOAT16) + session = _session(_matmul_model(k, n, raw.shape[1], 0), "CUDAExecutionProvider") + for m in (1, 7, 8): + with self.subTest(m=m): + a = _round_to_dtype(rng.standard_normal((m, k), dtype=np.float32), TensorProto.BFLOAT16) + got = session.run(None, {"A_f32": a, "B": raw})[0] + ref = _round_to_dtype(a @ weights.T, TensorProto.BFLOAT16) + np.testing.assert_allclose(got, ref, rtol=1e-2, atol=1e-2) + + def test_reject_unsupported_matmul_and_moe_formats(self): + with self.assertRaisesRegex(Exception, "unsupported packed_type"): + _session(_matmul_model(32, 8, 20, 3), "CUDAExecutionProvider") + for index, name in enumerate(("gate", "up", "down")): + with self.subTest(weight=name): + types = [0, 0, 0] + row_bytes = [128, 128, 128] + types[index], row_bytes[index] = 3, 20 + model = _moe_model(2, 32, 32, 1, TensorProto.FLOAT16, types, row_bytes, True) + with self.assertRaisesRegex(Exception, f"unsupported {name}_type"): + _session(model, "CUDAExecutionProvider") + + def test_moe_packed_reject_top_k_above_expert_count(self): + model = _moe_model(2, 32, 32, 4, TensorProto.FLOAT16, (0, 0, 0), (128, 128, 128), True) + session = _session(model, "CUDAExecutionProvider") + inputs = { + "X_f32": np.zeros((1, 32), dtype=np.float32), + "R_f32": np.zeros((1, 2), dtype=np.float32), + "G": np.zeros((2, 32, 128), dtype=np.uint8), + "U": np.zeros((2, 32, 128), dtype=np.uint8), + "D": np.zeros((2, 32, 128), dtype=np.uint8), + } + with self.assertRaisesRegex(Exception, "k <= num_experts"): + session.run(None, inputs) + + def _check_moe(self, num_experts, hidden_size, inter_size, dtype, formats, row_counts, normalize): rng = np.random.default_rng(999) - num_experts = 16 - hidden_size = 256 - inter_size = 64 - top_k = 4 - gate_type, gate_bel, gate_bb = 16, 256, 66 # IQ2_XXS - up_type, up_bel, up_bb = 16, 256, 66 - down_type, down_bel, down_bb = 42, 64, 18 # Q2_0 - - gate_rb = (hidden_size // gate_bel) * gate_bb - up_rb = (hidden_size // up_bel) * up_bb - down_rb = (inter_size // down_bel) * down_bb - - gate_raw = _make_synthetic_gguf_rows( - rng, num_experts * inter_size, hidden_size, gate_type, gate_bel, gate_bb - ).reshape(num_experts, inter_size, gate_rb) - up_raw = _make_synthetic_gguf_rows( - rng, num_experts * inter_size, hidden_size, up_type, up_bel, up_bb - ).reshape(num_experts, inter_size, up_rb) - down_raw = _make_synthetic_gguf_rows( - rng, num_experts * hidden_size, inter_size, down_type, down_bel, down_bb - ).reshape(num_experts, hidden_size, down_rb) - - nodes = [ - helper.make_node("Cast", ["X_f32"], ["X_bf16"], to=int(TensorProto.BFLOAT16)), - helper.make_node("Cast", ["R_f32"], ["R_bf16"], to=int(TensorProto.BFLOAT16)), - helper.make_node( - "MoEPacked", - ["X_bf16", "R_bf16", "G", "U", "D"], - ["Y_bf16"], - domain="com.microsoft", - k=top_k, - normalize_routing_weights=1, - activation_type="swiglu", - gate_type=gate_type, - up_type=up_type, - down_type=down_type, - ), - helper.make_node("Cast", ["Y_bf16"], ["Y_f32"], to=int(TensorProto.FLOAT)), - ] - moe_model = helper.make_model( - helper.make_graph( - nodes, - "moe_packed", - [ - helper.make_tensor_value_info("X_f32", TensorProto.FLOAT, ["M", hidden_size]), - helper.make_tensor_value_info("R_f32", TensorProto.FLOAT, ["M", num_experts]), - helper.make_tensor_value_info("G", TensorProto.UINT8, [num_experts, inter_size, gate_rb]), - helper.make_tensor_value_info("U", TensorProto.UINT8, [num_experts, inter_size, up_rb]), - helper.make_tensor_value_info("D", TensorProto.UINT8, [num_experts, hidden_size, down_rb]), - ], - [helper.make_tensor_value_info("Y_f32", TensorProto.FLOAT, ["M", hidden_size])], - ), - opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)], + top_k = min(4, num_experts) + raw_weights = [] + weights = [] + for out_size, k, (gtype, bel, bbytes) in zip( + (inter_size, inter_size, hidden_size), (hidden_size, hidden_size, inter_size), formats, strict=True + ): + raw = _make_synthetic_gguf_rows(rng, num_experts * out_size, k, gtype, bel, bbytes) + raw = raw.reshape(num_experts, out_size, -1) + raw_weights.append(raw) + weights.append(_dequantize_weights(raw, k, gtype)) + model = _moe_model( + num_experts, + hidden_size, + inter_size, + top_k, + dtype, + [f[0] for f in formats], + [w.shape[-1] for w in raw_weights], + normalize, ) - sess = ort.InferenceSession(moe_model.SerializeToString(), providers=["CUDAExecutionProvider"]) - for m in (4, 48): - x_f32 = rng.standard_normal((m, hidden_size), dtype=np.float32) * 0.2 - r_f32 = rng.standard_normal((m, num_experts), dtype=np.float32) - y_got = sess.run(None, {"X_f32": x_f32, "R_f32": r_f32, "G": gate_raw, "U": up_raw, "D": down_raw})[0] - self.assertEqual(y_got.shape, (m, hidden_size)) - self.assertTrue(np.all(np.isfinite(y_got))) + session = _session(model, "CUDAExecutionProvider") + for m in row_counts: + with self.subTest( + experts=num_experts, hidden=hidden_size, inter=inter_size, dtype=dtype, m=m, normalize=normalize + ): + x = rng.standard_normal((m, hidden_size), dtype=np.float32) * 0.2 + logits = rng.standard_normal((m, num_experts), dtype=np.float32) + # Concentrate routes so BF16 prefill fills multiple eight-route subtiles. + logits[:, :8] += 4.0 + logits[0] = 0.0 # Exercise tie-breaking in both routing kernels. + got = session.run( + None, {"X_f32": x, "R_f32": logits, "G": raw_weights[0], "U": raw_weights[1], "D": raw_weights[2]} + )[0] + ref = _moe_reference(x, logits, *weights, top_k, dtype, normalize) + self.assertEqual(got.shape, (m, hidden_size)) + self.assertTrue(np.all(np.isfinite(got))) + self.assertGreater(np.max(np.abs(ref)), 1e-2) + np.testing.assert_allclose(got, ref, rtol=2e-2, atol=2e-3) + + def test_moe_packed_cuda_decode_and_prefill(self): + formats = ((16, 256, 66), (16, 256, 66), (42, 64, 18)) + for dtype in (TensorProto.FLOAT16, TensorProto.BFLOAT16): + for num_experts in (16, 512): + for normalize in (False, True): + self._check_moe(num_experts, 256, 64, dtype, formats, (4, 48), normalize) + + def test_moe_packed_fp16_large_activations(self): + formats = ((1, 1, 2),) * 3 + for hidden_size, inter_size in ((4096, 64), (256, 4096)): + self._check_moe(2, hidden_size, inter_size, TensorProto.FLOAT16, formats, (4, 48), True) if __name__ == "__main__": From 0051ca2e8b75572a50874c9139a5eedbde35a2ad Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 8 Nov 2026 00:55:09 +0000 Subject: [PATCH 3/5] fix(cuda): preserve packed GEMM FP32 reductions Disallow output-type split-K reductions in MatMulPacked prefill so FP32 compute cannot silently round intermediate sums to BF16 or FP16. This addresses the Windows A10 packed-operator parity failures while retaining Tensor Core execution. Preserve the caller's cuBLAS math mode, propagate mode-update failures, and add FP16/BF16 cancellation coverage with a documented prefill numerical contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 1 + .../cuda/quantization/matmul_packed.cu | 15 ++++++++++- .../core/graph/contrib_ops/contrib_defs.cc | 3 ++- .../transformers/test_gguf_packed_ops.py | 25 +++++++++++++++---- 4 files changed, 37 insertions(+), 7 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index f9dca18163c84..94343a07bd798 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4724,6 +4724,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.MatMulPacked** Matrix multiplication Y = A @ B^T where B is a 2D block-packed weight tensor. + CUDA prefill dequantizes B to the activation type and uses float32 accumulation and reductions. #### Version diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu index 3ff0eeef331fc..f8b4b483f511c 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu @@ -323,10 +323,19 @@ Status DispatchMatMulPackedCuda( MatMulPackedDequantKernel<<>>(params, scratch_w); ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + // FP32 compute alone still permits FP16/BF16 split-K reductions. + cublasMath_t previous_math_mode = CUBLAS_DEFAULT_MATH; + CUBLAS_RETURN_IF_ERROR(cublasGetMathMode(cublas, &previous_math_mode)); + const auto math_mode = static_cast( + previous_math_mode | CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION); + if (math_mode != previous_math_mode) { + CUBLAS_RETURN_IF_ERROR(cublasSetMathMode(cublas, math_mode)); + } + const float alpha = 1.0f; const float beta = 0.0f; const cudaDataType_t data_type = is_bf16 ? CUDA_R_16BF : CUDA_R_16F; - return CUBLAS_CALL(cublasGemmEx( + const auto gemm_status = CUBLAS_CALL(cublasGemmEx( cublas, CUBLAS_OP_T, CUBLAS_OP_N, @@ -346,6 +355,10 @@ Status DispatchMatMulPackedCuda( params.n, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + if (math_mode != previous_math_mode) { + CUBLAS_RETURN_IF_ERROR(cublasSetMathMode(cublas, previous_math_mode)); + } + return gemm_status; } } // namespace diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index f08447115762a..9183ae5452d6d 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1026,7 +1026,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( ONNX_MS_OPERATOR_SET_SCHEMA( MatMulPacked, 1, OpSchema() - .SetDoc("Matrix multiplication Y = A @ B^T where B is a 2D block-packed weight tensor.") + .SetDoc("Matrix multiplication Y = A @ B^T where B is a 2D block-packed weight tensor.\n" + "CUDA prefill dequantizes B to the activation type and uses float32 accumulation and reductions.") .Attr("K", "Reduction dimension.", AttributeProto::INT) .Attr("N", "Output dimension.", AttributeProto::INT) .Attr("packed_type", diff --git a/onnxruntime/test/python/transformers/test_gguf_packed_ops.py b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py index 46d5f5e29181e..7595dd05d4e4c 100644 --- a/onnxruntime/test/python/transformers/test_gguf_packed_ops.py +++ b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py @@ -90,20 +90,20 @@ def _dequantize_weights(raw, k, gtype): return result.reshape(*raw.shape[:-1], k) -def _matmul_model(k, n, row_bytes, gtype): +def _matmul_model(k, n, row_bytes, gtype, dtype=TensorProto.BFLOAT16): return _model( [ - helper.make_node("Cast", ["A_f32"], ["A_bf16"], to=int(TensorProto.BFLOAT16)), + helper.make_node("Cast", ["A_f32"], ["A_typed"], to=int(dtype)), helper.make_node( "MatMulPacked", - ["A_bf16", "B"], - ["Y_bf16"], + ["A_typed", "B"], + ["Y_typed"], domain="com.microsoft", K=k, N=n, packed_type=gtype, ), - helper.make_node("Cast", ["Y_bf16"], ["Y_f32"], to=int(TensorProto.FLOAT)), + helper.make_node("Cast", ["Y_typed"], ["Y_f32"], to=int(TensorProto.FLOAT)), ], [ helper.make_tensor_value_info("A_f32", TensorProto.FLOAT, ["M", k]), @@ -296,6 +296,21 @@ def test_matmul_packed_cuda_all_types(self): ref = _round_to_dtype(a @ target.T, TensorProto.BFLOAT16) np.testing.assert_allclose(got, ref, rtol=1e-2, atol=1e-2) + def test_matmul_packed_prefill_fp32_reduction(self): + k, n = 4096, 32 + weights = np.ones((n, k), dtype=np.float32) + weights[:, -1] = 1.125 + raw = weights.view(np.uint8).reshape(n, k * 4) + for dtype in (TensorProto.BFLOAT16, TensorProto.FLOAT16): + session = _session(_matmul_model(k, n, raw.shape[1], 0, dtype), "CUDAExecutionProvider") + for m in (9, 32): + with self.subTest(dtype=dtype, m=m): + a = np.ones((m, k), dtype=np.float32) + a[:, k // 2 :] = -1 + got = session.run(None, {"A_f32": a, "B": raw})[0] + # Preserve the residual after cancellation of large partial sums. + np.testing.assert_array_equal(got, np.full((m, n), -0.125, dtype=np.float32)) + def test_matmul_packed_partial_tensor_core_tile(self): rng = np.random.default_rng(124) k, n = 32, 65537 From 3206eefe905531e31568cc297688ec050ac71615 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 8 Nov 2026 09:56:03 +0000 Subject: [PATCH 4/5] fix(cuda): address packed op review findings Validate dimensions before narrowing, support unaligned packed and activation buffers, and bound MatMulPacked prefill scratch. Check compiled BF16 capability, preserve endian-safe decoding and ggml notices, and extend CPU/CUDA regression coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ThirdPartyNotices.txt | 26 ++ docs/ContribOperators.md | 10 +- .../cpu/quantization/gguf_quant_tables.h | 69 +++- .../contrib_ops/cuda/moe/moe_packed.cc | 61 +++- .../contrib_ops/cuda/moe/moe_packed.cu | 96 +++-- onnxruntime/contrib_ops/cuda/moe/moe_packed.h | 11 +- .../cuda/quantization/gather_packed.cc | 25 +- .../cuda/quantization/gather_packed.cu | 8 +- .../cuda/quantization/gather_packed.h | 7 + .../cuda/quantization/gguf_quant_cuda.cuh | 81 +++-- .../cuda/quantization/matmul_packed.cc | 46 ++- .../cuda/quantization/matmul_packed.cu | 86 +++-- .../cuda/quantization/matmul_packed.h | 11 + .../core/graph/contrib_ops/contrib_defs.cc | 16 +- .../transformers/test_gguf_packed_ops.py | 335 +++++++++++++++++- 15 files changed, 730 insertions(+), 158 deletions(-) diff --git a/ThirdPartyNotices.txt b/ThirdPartyNotices.txt index 409183748a915..cbbf4cef81e49 100644 --- a/ThirdPartyNotices.txt +++ b/ThirdPartyNotices.txt @@ -19,6 +19,32 @@ General Public License. _____ +ggml-org/ggml + +MIT License + +Copyright (c) 2023-2026 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. + +_____ + protocolbuffers/protobuf Copyright 2008 Google Inc. All rights reserved. diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 94343a07bd798..a9c48a74a56a5 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2961,7 +2961,7 @@ This version of the operator has been available since version 1 of the 'com.micr
K : int (required)
-
Logical embedding dimension.
+
Logical embedding dimension. Must be a positive multiple of the selected packed format's block size.
output_dtype : int
ONNX TensorProto DataType for output (default 16 = BFLOAT16).
packed_type : int (required)
@@ -2972,9 +2972,9 @@ This version of the operator has been available since version 1 of the 'com.micr
data : T1
-
2D packed weight blob [vocab_size, row_bytes].
+
2D little-endian packed weight blob [vocab_size, row_bytes].
indices : Tind
-
Token indices of arbitrary rank.
+
Token indices of arbitrary rank. Negative indices and indices >= vocab_size select row 0.
#### Outputs @@ -4724,7 +4724,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.MatMulPacked** Matrix multiplication Y = A @ B^T where B is a 2D block-packed weight tensor. - CUDA prefill dequantizes B to the activation type and uses float32 accumulation and reductions. + CUDA prefill dequantizes output-dimension tiles of B to the activation type using bounded scratch memory and uses float32 accumulation and reductions. #### Version @@ -4734,7 +4734,7 @@ This version of the operator has been available since version 1 of the 'com.micr
K : int (required)
-
Reduction dimension.
+
Reduction dimension. Must be a whole number of packed blocks; CUDA requires a positive multiple of 32.
N : int (required)
Output dimension.
packed_type : int (required)
diff --git a/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h b/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h index e35de46b7634f..0cd8e4581f083 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h +++ b/onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h @@ -1,12 +1,37 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +// Quantization lookup tables are from ggml (https://github.com/ggml-org/ggml). +// +// MIT License +// +// Copyright (c) 2023-2026 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. + #pragma once #include #include #include #include "core/common/common.h" +#include "core/common/endian.h" #include "core/common/float16.h" namespace onnxruntime { @@ -844,15 +869,12 @@ inline float Bf16BitsToFloatCpu(uint16_t bits) { } inline uint16_t ReadU16Le(const uint8_t* p) { - uint16_t v = 0; - std::memcpy(&v, p, sizeof(v)); - return v; + return static_cast(p[0]) | (static_cast(p[1]) << 8); } inline uint32_t ReadU32Le(const uint8_t* p) { - uint32_t v = 0; - std::memcpy(&v, p, sizeof(v)); - return v; + return static_cast(p[0]) | (static_cast(p[1]) << 8) | + (static_cast(p[2]) << 16) | (static_cast(p[3]) << 24); } inline void GetScaleMinK4Cpu(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { @@ -868,7 +890,14 @@ inline void GetScaleMinK4Cpu(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { inline bool DequantizeGgufRowCpu(const uint8_t* src, int ggml_type, int64_t K, float* dst) { switch (ggml_type) { case kGgmlTypeF32: { - std::memcpy(dst, src, static_cast(K) * sizeof(float)); + if constexpr (endian::native == endian::little) { + std::memcpy(dst, src, static_cast(K) * sizeof(float)); + } else { + for (int64_t i = 0; i < K; ++i) { + const uint32_t bits = ReadU32Le(src + i * sizeof(float)); + std::memcpy(dst + i, &bits, sizeof(float)); + } + } return true; } case kGgmlTypeF16: { @@ -1045,10 +1074,11 @@ inline bool DequantizeGgufRowCpu(const uint8_t* src, int ggml_type, int64_t K, f const uint32_t aux32 = ReadU32Le(aux8 + 4); const float db = d * (0.5f + (aux32 >> 28)) * 0.25f; for (int l = 0; l < 4; ++l) { - const uint8_t* grid = reinterpret_cast(iq2xxs_grid + aux8[l]); + const uint64_t grid = iq2xxs_grid[aux8[l]]; const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * l)) & 127]; for (int j = 0; j < 8; ++j) { - y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + y[j] = db * static_cast((grid >> (8 * j)) & 0xff) * + (signs & kmask_iq2xs[j] ? -1.f : 1.f); } y += 8; } @@ -1069,11 +1099,12 @@ inline bool DequantizeGgufRowCpu(const uint8_t* src, int ggml_type, int64_t K, f const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; for (int l = 0; l < 4; ++l) { const uint16_t qval = ReadU16Le(qs + (4 * ib32 + l) * 2); - const uint8_t* grid = reinterpret_cast(iq2xs_grid + (qval & 511)); + const uint64_t grid = iq2xs_grid[qval & 511]; const uint8_t signs = ksigns_iq2xs[qval >> 9]; const float db = (l < 2) ? db0 : db1; for (int j = 0; j < 8; ++j) { - y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + y[j] = db * static_cast((grid >> (8 * j)) & 0xff) * + (signs & kmask_iq2xs[j] ? -1.f : 1.f); } y += 8; } @@ -1096,10 +1127,10 @@ inline bool DequantizeGgufRowCpu(const uint8_t* src, int ggml_type, int64_t K, f const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; for (int l = 0; l < 4; ++l) { const float dl = (l < 2) ? db0 : db1; - const uint8_t* grid = - reinterpret_cast(iq2s_grid + (qs[l] | ((qh[ib32] << (8 - 2 * l)) & 0x300))); + const uint64_t grid = iq2s_grid[qs[l] | ((qh[ib32] << (8 - 2 * l)) & 0x300)]; for (int j = 0; j < 8; ++j) { - y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); + y[j] = dl * static_cast((grid >> (8 * j)) & 0xff) * + (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); } y += 8; } @@ -1122,11 +1153,13 @@ inline bool DequantizeGgufRowCpu(const uint8_t* src, int ggml_type, int64_t K, f const float db = d * (0.5f + (aux32 >> 28)) * 0.5f; for (int l = 0; l < 4; ++l) { const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * l)) & 127]; - const uint8_t* grid1 = reinterpret_cast(iq3xxs_grid + qs[2 * l + 0]); - const uint8_t* grid2 = reinterpret_cast(iq3xxs_grid + qs[2 * l + 1]); + const uint32_t grid1 = iq3xxs_grid[qs[2 * l + 0]]; + const uint32_t grid2 = iq3xxs_grid[qs[2 * l + 1]]; for (int j = 0; j < 4; ++j) { - y[j + 0] = db * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); - y[j + 4] = db * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); + y[j + 0] = db * static_cast((grid1 >> (8 * j)) & 0xff) * + (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db * static_cast((grid2 >> (8 * j)) & 0xff) * + (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); } y += 8; } diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc index 7a2553bac4429..ec40344492a2c 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc @@ -7,12 +7,25 @@ #include #include "contrib_ops/cpu/quantization/gguf_quant_tables.h" +#include "core/common/safeint.h" +#include "core/framework/allocator.h" #include "core/providers/cuda/cuda_common.h" namespace onnxruntime { namespace contrib { namespace cuda { +namespace { +bool PackedTensorBytesFit(int64_t rows, int64_t columns, size_t element_bytes) { + size_t row_bytes = 0; + size_t bytes = 0; + return rows >= 0 && columns >= 0 && + IAllocator::CalcMemSizeForArray(static_cast(columns), element_bytes, &row_bytes) && + IAllocator::CalcMemSizeForArray(static_cast(rows), row_bytes, &bytes) && + bytes <= static_cast(std::numeric_limits::max()); +} +} // namespace + #define REGISTER_MOE_PACKED_CUDA(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ MoEPacked, \ @@ -42,6 +55,10 @@ MoEPacked::MoEPacked(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(gate_type_), "MoEPacked: unsupported gate_type=", gate_type_); ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(up_type_), "MoEPacked: unsupported up_type=", up_type_); ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(down_type_), "MoEPacked: unsupported down_type=", down_type_); + if constexpr (std::is_same_v) { + ORT_ENFORCE(GetDeviceProp().major >= 8, "MoEPacked BF16 requires compute capability 8.0 or higher"); + ORT_THROW_IF_ERROR(CheckMoEPackedBf16KernelSupport()); + } } template @@ -69,6 +86,10 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { const int64_t num_experts = rp_shape[1]; const int64_t inter_size = gw_shape[1]; + ORT_RETURN_IF_NOT(hidden_size > 0 && hidden_size <= kMoEPackedMaxDimension && + inter_size > 0 && inter_size <= kMoEPackedMaxDimension, + "MoEPacked requires hidden_size and inter_size in 1..", kMoEPackedMaxDimension); + ORT_RETURN_IF_NOT(num_rows >= 0, "MoEPacked requires a nonnegative row count"); ORT_RETURN_IF_NOT(rp_shape[0] == num_rows, "MoEPacked router_probs row count mismatch"); ORT_RETURN_IF_NOT(num_experts > 0 && num_experts <= 512, "MoEPacked supports up to 512 experts"); ORT_RETURN_IF_NOT(k_ <= num_experts, "MoEPacked requires k <= num_experts"); @@ -85,16 +106,23 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(exp_gate_rb > 0 && exp_up_rb > 0 && exp_down_rb > 0 && gw_shape[2] == exp_gate_rb && uw_shape[2] == exp_up_rb && dw_shape[2] == exp_down_rb, "MoEPacked expert weight row_bytes mismatch"); - - Tensor* output = context->Output(0, in_shape); - if (num_rows == 0) return Status::OK(); constexpr bool is_bf16 = std::is_same_v; - ORT_RETURN_IF_NOT(!is_bf16 || GetDeviceProp().major >= 8, - "MoEPacked BF16 requires compute capability 8.0 or higher"); + ORT_RETURN_IF_NOT(PackedTensorBytesFit(num_rows, hidden_size, sizeof(T)) && + PackedTensorBytesFit(num_rows, num_experts, sizeof(T)) && + PackedTensorBytesFit(num_experts * inter_size, exp_gate_rb, 1) && + PackedTensorBytesFit(num_experts * inter_size, exp_up_rb, 1) && + PackedTensorBytesFit(num_experts * hidden_size, exp_down_rb, 1), + "MoEPacked tensor byte size exceeds the supported address range"); // Process in chunks of at most 2048 rows so workspace stays compact during large prefills. - constexpr int64_t kMaxChunkRows = 2048; - const int64_t max_chunk = std::min(num_rows, kMaxChunkRows); + const int64_t max_chunk = std::min(num_rows, kMoEPackedMaxChunkRows); + const int64_t max_routes = max_chunk * k_; + const int64_t max_tiles = max_routes <= 128 ? max_routes : (max_routes + 7) / 8 + num_experts; + ORT_RETURN_IF_NOT((hidden_size + 7) / 8 <= GetDeviceProp().maxGridSize[0] && + (inter_size + 7) / 8 <= GetDeviceProp().maxGridSize[0] && + max_tiles <= GetDeviceProp().maxGridSize[1] && + max_chunk <= GetDeviceProp().maxGridSize[1], + "MoEPacked dimensions exceed the CUDA grid limits"); MoEPackedParams max_params{}; max_params.num_rows = static_cast(max_chunk); @@ -103,6 +131,7 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { max_params.inter_size = static_cast(inter_size); max_params.k = static_cast(k_); max_params.normalize_routing_weights = (normalize_routing_weights_ != 0); + max_params.bf16_tc_supported = is_bf16; max_params.gate_type = static_cast(gate_type_); max_params.up_type = static_cast(up_type_); max_params.down_type = static_cast(down_type_); @@ -110,28 +139,34 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { max_params.up_row_bytes = exp_up_rb; max_params.down_row_bytes = exp_down_rb; - const size_t ws_bytes = GetMoEPackedWorkspaceBytes(max_params, sizeof(T)); + size_t ws_bytes = 0; + ORT_RETURN_IF_ERROR(GetMoEPackedWorkspaceBytes(max_params, sizeof(T), ws_bytes)); + Tensor* output = context->Output(0, in_shape); + if (num_rows == 0) return Status::OK(); auto workspace = GetScratchBuffer(ws_bytes, GetComputeStream(context)); const uint8_t* in_bytes = static_cast(input->DataRaw()); const uint8_t* rp_bytes = static_cast(router_probs->DataRaw()); uint8_t* out_bytes = static_cast(output->MutableDataRaw()); - for (int64_t r0 = 0; r0 < num_rows; r0 += kMaxChunkRows) { - const int64_t chunk_rows = std::min(kMaxChunkRows, num_rows - r0); + for (int64_t r0 = 0; r0 < num_rows;) { + const int64_t chunk_rows = std::min(kMoEPackedMaxChunkRows, num_rows - r0); MoEPackedParams chunk_params = max_params; chunk_params.num_rows = static_cast(chunk_rows); + const size_t input_offset = SafeInt(r0) * hidden_size * sizeof(T); + const size_t router_offset = SafeInt(r0) * num_experts * sizeof(T); ORT_RETURN_IF_ERROR(LaunchMoEPackedCuda( Stream(context), is_bf16, - in_bytes + static_cast(r0 * hidden_size) * sizeof(T), - rp_bytes + static_cast(r0 * num_experts) * sizeof(T), + in_bytes + input_offset, + rp_bytes + router_offset, gate_w->Data(), up_w->Data(), down_w->Data(), chunk_params, workspace.get(), - out_bytes + static_cast(r0 * hidden_size) * sizeof(T))); + out_bytes + input_offset)); + r0 += chunk_rows; } return Status::OK(); diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu index 92672b730be20..1b511ff86e297 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu @@ -5,9 +5,12 @@ #include #include +#include #include "contrib_ops/cuda/moe/qmoe_fused_route.cuh" #include "contrib_ops/cuda/quantization/gguf_quant_cuda.cuh" +#include "core/common/safeint.h" +#include "core/framework/allocator.h" #include "core/providers/cuda/cuda_common.h" namespace onnxruntime { @@ -398,7 +401,7 @@ __global__ __launch_bounds__(KW * 32) void MoEPackedFc1DecodeWarpKSplitKernel( const int num_sb32 = p.hidden_size >> 5; for (int sb32_idx = warp; sb32_idx < num_sb32; sb32_idx += KW) { const uint4 b_vec = x_valid - ? __ldg(reinterpret_cast(x_row + sb32_idx * 32 + t * 8)) + ? gguf_cuda::LoadU4Dev(x_row + sb32_idx * 32 + t * 8) : make_uint4(0u, 0u, 0u, 0u); #pragma unroll for (int m = 0; m < MT; ++m) { @@ -553,7 +556,7 @@ __global__ __launch_bounds__((kMaxSubTiles == 1) ? 128 : 256) void MoEPackedFc1S const int token_row = s_token_rows[m]; const __nv_bfloat16* src_ptr = input + static_cast(token_row) * p.hidden_size + k_base + (u4_col << 3); - val = __ldg(reinterpret_cast(src_ptr)); + val = gguf_cuda::LoadU4Dev(src_ptr); } s_x_u4[m * kRowStrideU4 + u4_col] = val; } @@ -566,7 +569,7 @@ __global__ __launch_bounds__((kMaxSubTiles == 1) ? 128 : 256) void MoEPackedFc1S const int token_row = s_token_rows[m]; const __nv_bfloat16* src_ptr = input + static_cast(token_row) * p.hidden_size + k_base + (u4_col << 3); - val = __ldg(reinterpret_cast(src_ptr)); + val = gguf_cuda::LoadU4Dev(src_ptr); } s_x_u4[m * kRowStrideU4 + u4_col] = val; } @@ -979,8 +982,11 @@ __global__ __launch_bounds__(256) void MoEPackedFinalizeKernel( output[static_cast(row) * hidden_size + col] = gguf_cuda::FloatToCudaType(sum); } -size_t AlignUpWorkspace(size_t v) { - return (v + 255u) & ~size_t{255u}; +size_t AlignUpWorkspace(size_t count, size_t element_bytes = 1) { + size_t bytes = 0; + ORT_ENFORCE(IAllocator::CalcMemSizeForArrayWithAlignment(count, element_bytes, 256, &bytes), + "MoEPacked workspace byte size overflow"); + return bytes; } void LaunchMoEPackedFc1DecodeBf16TcDispatch( @@ -1171,9 +1177,18 @@ Status DispatchMoEPackedCuda( const MoEPackedParams& params, void* workspace, CudaT* output) { + ORT_RETURN_IF_NOT(params.num_rows >= 0 && params.num_rows <= kMoEPackedMaxChunkRows && + params.hidden_size > 0 && params.hidden_size <= kMoEPackedMaxDimension && + params.inter_size > 0 && params.inter_size <= kMoEPackedMaxDimension && + params.hidden_size % 32 == 0 && params.inter_size % 32 == 0 && + params.num_experts > 0 && params.num_experts <= kMoEPackedMaxExperts && + params.k > 0 && params.k <= kMoEPackedMaxK && params.k <= params.num_experts, + "MoEPacked parameters exceed the supported CUDA index range"); if (params.num_rows == 0) return Status::OK(); const int total_routes = params.num_rows * params.k; constexpr bool kIsBf16 = std::is_same::value; + ORT_RETURN_IF_NOT(!kIsBf16 || params.bf16_tc_supported, + "MoEPacked BF16 requires Tensor Core kernels compiled for compute_80 or higher"); const bool use_prefill_tiles = kIsBf16 && (total_routes > 128); const int max_tiles = use_prefill_tiles ? ((total_routes + 31) / 32 + params.num_experts) @@ -1181,22 +1196,22 @@ Status DispatchMoEPackedCuda( uint8_t* ws_ptr = static_cast(workspace); float* topk_scales = reinterpret_cast(ws_ptr); - ws_ptr += AlignUpWorkspace(static_cast(total_routes) * sizeof(float)); + ws_ptr += AlignUpWorkspace(total_routes, sizeof(float)); int* topk_indices = reinterpret_cast(ws_ptr); - ws_ptr += AlignUpWorkspace(static_cast(total_routes) * sizeof(int)); + ws_ptr += AlignUpWorkspace(total_routes, sizeof(int)); int* sorted_routes = reinterpret_cast(ws_ptr); - ws_ptr += AlignUpWorkspace(static_cast(total_routes) * sizeof(int)); + ws_ptr += AlignUpWorkspace(total_routes, sizeof(int)); MoEPackedTile* tiles = reinterpret_cast(ws_ptr); - ws_ptr += AlignUpWorkspace(static_cast(max_tiles) * sizeof(MoEPackedTile)); + ws_ptr += AlignUpWorkspace(max_tiles, sizeof(MoEPackedTile)); int* d_num_tiles = reinterpret_cast(ws_ptr); ws_ptr += AlignUpWorkspace(sizeof(int)); CudaT* fc1_out = reinterpret_cast(ws_ptr); - ws_ptr += AlignUpWorkspace(static_cast(total_routes) * params.inter_size * sizeof(CudaT)); + ws_ptr += AlignUpWorkspace(SafeInt(total_routes) * params.inter_size, sizeof(CudaT)); CudaT* fc2_out = reinterpret_cast(ws_ptr); @@ -1251,7 +1266,7 @@ Status DispatchMoEPackedCuda( } else { // Keep staging within CUDA's default shared-memory limit; larger shapes read activations directly. constexpr size_t kMaxStagingBytes = 48 * 1024; - const size_t fc1_smem = static_cast(8) * params.hidden_size * sizeof(CudaT); + const size_t fc1_smem = SafeInt(8) * params.hidden_size * sizeof(CudaT); const dim3 fc1_grid(static_cast((params.inter_size + 7) / 8), static_cast(max_tiles)); if (fc1_smem <= kMaxStagingBytes) { MoEPackedFc1SwiGLUSimtKernel<<>>( @@ -1262,7 +1277,7 @@ Status DispatchMoEPackedCuda( } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - const size_t fc2_smem = static_cast(8) * params.inter_size * sizeof(CudaT); + const size_t fc2_smem = SafeInt(8) * params.inter_size * sizeof(CudaT); const dim3 fc2_grid(static_cast((params.hidden_size + 7) / 8), static_cast(max_tiles)); if (fc2_smem <= kMaxStagingBytes) { MoEPackedFc2SimtKernel<<>>( @@ -1274,8 +1289,8 @@ Status DispatchMoEPackedCuda( ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); } - // 5. Finalize (128-bit vectorized when hidden_size % 8 == 0) - if ((params.hidden_size & 7) == 0) { + // IO-bound output may only be element-aligned. + if ((params.hidden_size & 7) == 0 && (reinterpret_cast(output) & 15u) == 0) { const int vec_cols = params.hidden_size >> 3; const dim3 fin_grid(static_cast((vec_cols + 255) / 256), static_cast(params.num_rows)); @@ -1292,19 +1307,52 @@ Status DispatchMoEPackedCuda( } // namespace -size_t GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_bytes) { - const size_t total_routes = static_cast(params.num_rows) * params.k; +Status GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_bytes, size_t& workspace_bytes) { + ORT_RETURN_IF_NOT(params.num_rows >= 0 && params.num_rows <= kMoEPackedMaxChunkRows && + params.k > 0 && params.k <= kMoEPackedMaxK && + params.num_experts > 0 && params.num_experts <= kMoEPackedMaxExperts && + params.hidden_size > 0 && params.hidden_size <= kMoEPackedMaxDimension && + params.inter_size > 0 && params.inter_size <= kMoEPackedMaxDimension && + elem_bytes == 2, + "MoEPacked workspace parameters exceed the supported range"); + if (params.num_rows == 0) { + workspace_bytes = 0; + return Status::OK(); + } + const size_t total_routes = SafeInt(params.num_rows) * params.k; const size_t max_tiles = (total_routes <= 128) ? total_routes : ((total_routes + 7) / 8 + static_cast(params.num_experts)); size_t bytes = 0; - bytes += AlignUpWorkspace(total_routes * sizeof(float)); - bytes += AlignUpWorkspace(total_routes * sizeof(int)); - bytes += AlignUpWorkspace(total_routes * sizeof(int)); - bytes += AlignUpWorkspace(max_tiles * sizeof(MoEPackedTile)); - bytes += AlignUpWorkspace(sizeof(int)); - bytes += AlignUpWorkspace(total_routes * static_cast(params.inter_size) * elem_bytes); - bytes += AlignUpWorkspace(total_routes * static_cast(params.hidden_size) * elem_bytes); - return bytes; + const auto add_buffer = [&bytes](size_t count, size_t element_bytes) { + size_t buffer_bytes = 0; + if (!IAllocator::CalcMemSizeForArrayWithAlignment(count, element_bytes, 256, &buffer_bytes) || + buffer_bytes > static_cast(std::numeric_limits::max()) - bytes) { + return false; + } + bytes += buffer_bytes; + return true; + }; + ORT_RETURN_IF_NOT(add_buffer(total_routes, sizeof(float)) && + add_buffer(total_routes, sizeof(int)) && + add_buffer(total_routes, sizeof(int)) && + add_buffer(max_tiles, sizeof(MoEPackedTile)) && + add_buffer(1, sizeof(int)) && + add_buffer(SafeInt(total_routes) * params.inter_size, elem_bytes) && + add_buffer(SafeInt(total_routes) * params.hidden_size, elem_bytes), + "MoEPacked workspace byte size exceeds the supported address range"); + workspace_bytes = bytes; + return Status::OK(); +} + +Status CheckMoEPackedBf16KernelSupport() { + cudaFuncAttributes attributes{}; + // All BF16 MMA instantiations in this translation unit use the same virtual architecture. + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, MoEPackedFc1DecodeWarpKSplitKernel<4, 2, -1>)); + // binaryVersion can be >= 80 after JIT of compute_75 PTX, whose MMA body is still unsupported. + ORT_RETURN_IF_NOT(attributes.ptxVersion >= 80, + "MoEPacked BF16 requires Tensor Core kernels compiled for compute_80 or higher; got compute_", + attributes.ptxVersion); + return Status::OK(); } Status LaunchMoEPackedCuda( diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.h b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h index 6685651a2e84f..ca95f4c4ba757 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h @@ -3,7 +3,9 @@ #pragma once +#include #include +#include #include #include "core/providers/cuda/cuda_kernel.h" @@ -14,6 +16,10 @@ namespace cuda { using namespace onnxruntime::cuda; +inline constexpr int kMoEPackedMaxChunkRows = 2048; +// Covers 512/640-element K-loop increments and rounded 2048-element finalize tiles. +inline constexpr int64_t kMoEPackedMaxDimension = std::numeric_limits::max() - 2047; + struct alignas(16) MoEPackedTile { int expert; int start_route; @@ -28,6 +34,7 @@ struct MoEPackedParams { int inter_size{0}; int k{0}; bool normalize_routing_weights{true}; + bool bf16_tc_supported{false}; int gate_type{0}; int up_type{0}; int down_type{0}; @@ -36,7 +43,9 @@ struct MoEPackedParams { int64_t down_row_bytes{0}; }; -size_t GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_bytes); +Status GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_bytes, size_t& workspace_bytes); + +Status CheckMoEPackedBf16KernelSupport(); Status LaunchMoEPackedCuda( cudaStream_t stream, diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc index 75fd583d33d60..33a3f9fd7e7c9 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cc @@ -6,12 +6,24 @@ #include #include "contrib_ops/cpu/quantization/gguf_quant_tables.h" +#include "core/framework/allocator.h" #include "core/providers/cuda/cuda_common.h" namespace onnxruntime { namespace contrib { namespace cuda { +namespace { +bool PackedTensorBytesFit(int64_t rows, int64_t columns, size_t element_bytes) { + size_t row_bytes = 0; + size_t bytes = 0; + return rows >= 0 && columns >= 0 && + IAllocator::CalcMemSizeForArray(static_cast(columns), element_bytes, &row_bytes) && + IAllocator::CalcMemSizeForArray(static_cast(rows), row_bytes, &bytes) && + bytes <= static_cast(std::numeric_limits::max()); +} +} // namespace + #define REGISTER_GATHER_PACKED_CUDA(T2, Tind) \ ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ GatherPacked, \ @@ -36,6 +48,7 @@ REGISTER_GATHER_PACKED_CUDA(BFloat16, int64_t) template GatherPacked::GatherPacked(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(info.GetAttr("K", &k_).IsOK() && k_ > 0, "GatherPacked requires K > 0"); + ORT_ENFORCE(k_ <= kGatherPackedMaxK, "GatherPacked requires K <= ", kGatherPackedMaxK); ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK(), "GatherPacked requires packed_type"); ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(ggml_type_), "GatherPacked: unsupported packed_type=", ggml_type_); row_bytes_ = gguf_quant::GgmlRowBytes(static_cast(ggml_type_), k_); @@ -53,13 +66,21 @@ Status GatherPacked::ComputeInternal(OpKernelContext* context) const { "GatherPacked data.shape[1]=", data_shape[1], " does not match expected row_bytes=", row_bytes_); const auto& indices_shape = indices->Shape(); + const int64_t num_rows = indices_shape.Size(); + ORT_RETURN_IF_NOT(num_rows >= 0 && num_rows <= kGatherPackedMaxRows && + num_rows <= GetDeviceProp().maxGridSize[0], + "GatherPacked index count exceeds the CUDA grid limit"); + ORT_RETURN_IF_NOT(PackedTensorBytesFit(num_rows, k_, sizeof(T2)) && + PackedTensorBytesFit(vocab_size, row_bytes_, 1), + "GatherPacked tensor byte size exceeds the supported address range"); + ORT_RETURN_IF_NOT(num_rows == 0 || vocab_size > 0, + "GatherPacked requires a nonempty vocabulary for nonempty indices"); + TensorShapeVector out_dims = indices_shape.AsShapeVector(); out_dims.push_back(k_); Tensor* output = context->Output(0, TensorShape(out_dims)); - const int64_t num_rows = indices_shape.Size(); if (num_rows == 0) return Status::OK(); - ORT_RETURN_IF_NOT(vocab_size > 0, "GatherPacked requires a nonempty vocabulary for nonempty indices"); constexpr int out_dtype = std::is_same_v ? 16 : (std::is_same_v ? 10 : 1); diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu index 9435dc23c1ece..e4ead19787ee9 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.cu @@ -33,9 +33,9 @@ __global__ __launch_bounds__(256) void GatherPackedKernel( for (int col = threadIdx.x; col < k; col += blockDim.x) { float value; if (ggml_type == gguf_quant::kGgmlTypeF32) { - value = __ldg(reinterpret_cast(row_ptr) + col); + value = __uint_as_float(gguf_cuda::ReadU32Dev(row_ptr + static_cast(col) * 4)); } else { - const uint16_t bits = __ldg(reinterpret_cast(row_ptr) + col); + const uint16_t bits = gguf_cuda::ReadU16Dev(row_ptr + static_cast(col) * 2); value = ggml_type == gguf_quant::kGgmlTypeF16 ? gguf_cuda::Fp16ToFloatDev(bits) : gguf_cuda::Bf16ToFloatDev(bits); } @@ -68,7 +68,11 @@ Status DispatchGatherPackedByIndex( int64_t k, int64_t row_bytes, int ggml_type) { + ORT_RETURN_IF_NOT(k > 0 && k <= kGatherPackedMaxK, "GatherPacked requires K in 1..", kGatherPackedMaxK); + ORT_RETURN_IF_NOT(num_rows >= 0 && num_rows <= kGatherPackedMaxRows, + "GatherPacked index count exceeds the CUDA grid limit"); if (num_rows == 0) return Status::OK(); + ORT_RETURN_IF_NOT(vocab_size > 0 && row_bytes > 0, "GatherPacked requires nonempty packed rows"); const dim3 grid(static_cast(num_rows)); const dim3 block(256); if (indices_are_int64) { diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.h b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.h index 109815ec48825..9c6e71208c644 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_packed.h +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_packed.h @@ -3,6 +3,9 @@ #pragma once +#include +#include + #include "core/providers/cuda/cuda_kernel.h" namespace onnxruntime { @@ -11,6 +14,10 @@ namespace cuda { using namespace onnxruntime::cuda; +// Leave room for the scalar kernel's final 256-element loop increment. +inline constexpr int64_t kGatherPackedMaxK = std::numeric_limits::max() - 256; +inline constexpr int64_t kGatherPackedMaxRows = std::numeric_limits::max(); + Status LaunchGatherPackedCuda( cudaStream_t stream, const uint8_t* data, diff --git a/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh b/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh index 552b972bac4ee..f074710e6855e 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh @@ -18,13 +18,26 @@ namespace gguf_cuda { using namespace onnxruntime::contrib::gguf_quant; __device__ __forceinline__ uint16_t ReadU16Dev(const uint8_t* p) { - return __ldg(reinterpret_cast(p)); + if ((reinterpret_cast(p) & 1u) == 0) { + return __ldg(reinterpret_cast(p)); + } + return static_cast(__ldg(p)) | (static_cast(__ldg(p + 1)) << 8); } __device__ __forceinline__ uint32_t ReadU32Dev(const uint8_t* p) { - const uint32_t lo = __ldg(reinterpret_cast(p)); - const uint32_t hi = __ldg(reinterpret_cast(p + 2)); - return __byte_perm(lo, hi, 0x5410u); + if ((reinterpret_cast(p) & 3u) == 0) { + return __ldg(reinterpret_cast(p)); + } + return static_cast(ReadU16Dev(p)) | (static_cast(ReadU16Dev(p + 2)) << 16); +} + +// Tensor rows may be element-aligned without being aligned for a 128-bit load. +__device__ __forceinline__ uint4 LoadU4Dev(const void* p) { + if ((reinterpret_cast(p) & 15u) == 0) { + return __ldg(reinterpret_cast(p)); + } + const auto* bytes = static_cast(p); + return make_uint4(ReadU32Dev(bytes), ReadU32Dev(bytes + 4), ReadU32Dev(bytes + 8), ReadU32Dev(bytes + 12)); } __device__ __forceinline__ float Fp16ToFloatDev(uint16_t bits) { @@ -83,17 +96,17 @@ template __device__ __forceinline__ void DecodeGgufSubBlock8Impl( const uint8_t* row_ptr, int sb32_idx, int l, float (&out)[8]) { if constexpr (kCompileType == kGgmlTypeF32) { - const float* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + l * 8; + const uint8_t* p = row_ptr + (static_cast(sb32_idx) * 32 + l * 8) * 4; #pragma unroll - for (int j = 0; j < 8; ++j) out[j] = __ldg(p + j); + for (int j = 0; j < 8; ++j) out[j] = __uint_as_float(ReadU32Dev(p + j * 4)); } else if constexpr (kCompileType == kGgmlTypeF16) { - const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + l * 8; + const uint8_t* p = row_ptr + (static_cast(sb32_idx) * 32 + l * 8) * 2; #pragma unroll - for (int j = 0; j < 8; ++j) out[j] = Fp16ToFloatDev(__ldg(p + j)); + for (int j = 0; j < 8; ++j) out[j] = Fp16ToFloatDev(ReadU16Dev(p + j * 2)); } else if constexpr (kCompileType == kGgmlTypeBF16) { - const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + l * 8; + const uint8_t* p = row_ptr + (static_cast(sb32_idx) * 32 + l * 8) * 2; #pragma unroll - for (int j = 0; j < 8; ++j) out[j] = Bf16ToFloatDev(__ldg(p + j)); + for (int j = 0; j < 8; ++j) out[j] = Bf16ToFloatDev(ReadU16Dev(p + j * 2)); } else if constexpr (kCompileType == kGgmlTypeQ8_0) { const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 34; const float d = Fp16ToFloatDev(ReadU16Dev(blk)); @@ -231,11 +244,11 @@ __device__ __forceinline__ void DecodeGgufSubBlock8Impl( const uint32_t aux32 = ReadU32Dev(aux8 + 4); const float db = d * (0.5f + static_cast(aux32 >> 28)) * 0.25f; const uint64_t g64 = __ldg(iq2xxs_grid + __ldg(aux8 + l)); - const uint8_t* grid = reinterpret_cast(&g64); const uint8_t signs = Iq2xsSignsDev(aux32 >> (7 * l)); #pragma unroll for (int j = 0; j < 8; ++j) { - out[j] = db * static_cast(grid[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); + out[j] = db * static_cast(static_cast(g64 >> (8 * j))) * + ((signs & (1u << j)) ? -1.0f : 1.0f); } } else if constexpr (kCompileType == kGgmlTypeIQ2_XS) { const int sb256 = sb32_idx >> 3; @@ -246,11 +259,11 @@ __device__ __forceinline__ void DecodeGgufSubBlock8Impl( const float db = d * (0.5f + static_cast((sc_byte >> (4 * (l >> 1))) & 0xF)) * 0.25f; const uint16_t q2 = ReadU16Dev(blk + 2 + (4 * ib + l) * 2); const uint64_t g64 = __ldg(iq2xs_grid + (q2 & 511)); - const uint8_t* grid = reinterpret_cast(&g64); const uint8_t signs = Iq2xsSignsDev(q2 >> 9); #pragma unroll for (int j = 0; j < 8; ++j) { - out[j] = db * static_cast(grid[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); + out[j] = db * static_cast(static_cast(g64 >> (8 * j))) * + ((signs & (1u << j)) ? -1.0f : 1.0f); } } else if constexpr (kCompileType == kGgmlTypeIQ2_S) { const int sb256 = sb32_idx >> 3; @@ -263,11 +276,11 @@ __device__ __forceinline__ void DecodeGgufSubBlock8Impl( const int g_idx = static_cast(__ldg(blk + 2 + 4 * ib + l)) | ((static_cast(qh_ib) << (8 - 2 * l)) & 0x300); const uint64_t g64 = __ldg(iq2s_grid + g_idx); - const uint8_t* grid = reinterpret_cast(&g64); const uint8_t signs = __ldg(blk + 34 + 4 * ib + l); #pragma unroll for (int j = 0; j < 8; ++j) { - out[j] = db * static_cast(grid[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); + out[j] = db * static_cast(static_cast(g64 >> (8 * j))) * + ((signs & (1u << j)) ? -1.0f : 1.0f); } } else if constexpr (kCompileType == kGgmlTypeIQ3_XXS) { const int sb256 = sb32_idx >> 3; @@ -280,12 +293,12 @@ __device__ __forceinline__ void DecodeGgufSubBlock8Impl( const uint16_t q2 = ReadU16Dev(blk + 2 + 8 * ib + 2 * l); const uint32_t g1 = __ldg(iq3xxs_grid + (q2 & 0xFFu)); const uint32_t g2 = __ldg(iq3xxs_grid + (q2 >> 8)); - const uint8_t* grid1 = reinterpret_cast(&g1); - const uint8_t* grid2 = reinterpret_cast(&g2); #pragma unroll for (int j = 0; j < 4; ++j) { - out[j] = db * static_cast(grid1[j]) * ((signs & (1u << j)) ? -1.0f : 1.0f); - out[j + 4] = db * static_cast(grid2[j]) * ((signs & (1u << (j + 4))) ? -1.0f : 1.0f); + out[j] = db * static_cast(static_cast(g1 >> (8 * j))) * + ((signs & (1u << j)) ? -1.0f : 1.0f); + out[j + 4] = db * static_cast(static_cast(g2 >> (8 * j))) * + ((signs & (1u << (j + 4))) ? -1.0f : 1.0f); } } else { #pragma unroll @@ -375,11 +388,10 @@ __device__ __forceinline__ void DecodeGgufSubBlock8SmemLutTpl( const uint32_t aux32 = ReadU32Dev(aux8 + 4); const float db = d * (0.5f + static_cast(aux32 >> 28)) * 0.25f; const uint64_t g64 = static_cast(s_lut)[__ldg(aux8 + l)]; - const uint8_t* grid = reinterpret_cast(&g64); const uint8_t signs = Iq2xsSignsDev(aux32 >> (7 * l)); #pragma unroll for (int j = 0; j < 8; ++j) { - out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid[j]); + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(static_cast(g64 >> (8 * j))); } } else if constexpr (kCompileType == kGgmlTypeIQ2_XS) { const int sb256 = sb32_idx >> 3; @@ -390,11 +402,10 @@ __device__ __forceinline__ void DecodeGgufSubBlock8SmemLutTpl( const float db = d * (0.5f + static_cast((sc_byte >> (4 * (l >> 1))) & 0xF)) * 0.25f; const uint16_t q2 = ReadU16Dev(blk + 2 + (4 * ib + l) * 2); const uint64_t g64 = static_cast(s_lut)[q2 & 511]; - const uint8_t* grid = reinterpret_cast(&g64); const uint8_t signs = Iq2xsSignsDev(q2 >> 9); #pragma unroll for (int j = 0; j < 8; ++j) { - out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid[j]); + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(static_cast(g64 >> (8 * j))); } } else if constexpr (kCompileType == kGgmlTypeIQ2_S) { const int sb256 = sb32_idx >> 3; @@ -407,11 +418,10 @@ __device__ __forceinline__ void DecodeGgufSubBlock8SmemLutTpl( const int g_idx = static_cast(__ldg(blk + 2 + 4 * ib + l)) | ((static_cast(qh_ib) << (8 - 2 * l)) & 0x300); const uint64_t g64 = static_cast(s_lut)[g_idx]; - const uint8_t* grid = reinterpret_cast(&g64); const uint8_t signs = __ldg(blk + 34 + 4 * ib + l); #pragma unroll for (int j = 0; j < 8; ++j) { - out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid[j]); + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(static_cast(g64 >> (8 * j))); } } else if constexpr (kCompileType == kGgmlTypeIQ3_XXS) { const int sb256 = sb32_idx >> 3; @@ -424,12 +434,11 @@ __device__ __forceinline__ void DecodeGgufSubBlock8SmemLutTpl( const uint16_t q2 = ReadU16Dev(blk + 2 + 8 * ib + 2 * l); const uint32_t g1 = static_cast(s_lut)[q2 & 0xFFu]; const uint32_t g2 = static_cast(s_lut)[q2 >> 8]; - const uint8_t* grid1 = reinterpret_cast(&g1); - const uint8_t* grid2 = reinterpret_cast(&g2); #pragma unroll for (int j = 0; j < 4; ++j) { - out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(grid1[j]); - out[j + 4] = ((signs & (1u << (j + 4))) ? -db : db) * static_cast(grid2[j]); + out[j] = ((signs & (1u << j)) ? -db : db) * static_cast(static_cast(g1 >> (8 * j))); + out[j + 4] = ((signs & (1u << (j + 4))) ? -db : db) * + static_cast(static_cast(g2 >> (8 * j))); } } else { DecodeGgufSubBlock8Tpl(row_ptr, runtime_type, sb32_idx, l, out); @@ -442,17 +451,17 @@ template __device__ __forceinline__ void DecodeGgufSubBlock16Impl( const uint8_t* row_ptr, int sb32_idx, int half16, float (&out)[16]) { if constexpr (kCompileType == kGgmlTypeF32) { - const float* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + half16 * 16; + const uint8_t* p = row_ptr + (static_cast(sb32_idx) * 32 + half16 * 16) * 4; #pragma unroll - for (int j = 0; j < 16; ++j) out[j] = __ldg(p + j); + for (int j = 0; j < 16; ++j) out[j] = __uint_as_float(ReadU32Dev(p + j * 4)); } else if constexpr (kCompileType == kGgmlTypeF16) { - const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + half16 * 16; + const uint8_t* p = row_ptr + (static_cast(sb32_idx) * 32 + half16 * 16) * 2; #pragma unroll - for (int j = 0; j < 16; ++j) out[j] = Fp16ToFloatDev(__ldg(p + j)); + for (int j = 0; j < 16; ++j) out[j] = Fp16ToFloatDev(ReadU16Dev(p + j * 2)); } else if constexpr (kCompileType == kGgmlTypeBF16) { - const uint16_t* p = reinterpret_cast(row_ptr) + sb32_idx * 32 + half16 * 16; + const uint8_t* p = row_ptr + (static_cast(sb32_idx) * 32 + half16 * 16) * 2; #pragma unroll - for (int j = 0; j < 16; ++j) out[j] = Bf16ToFloatDev(__ldg(p + j)); + for (int j = 0; j < 16; ++j) out[j] = Bf16ToFloatDev(ReadU16Dev(p + j * 2)); } else if constexpr (kCompileType == kGgmlTypeQ8_0) { const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 34; const float d = Fp16ToFloatDev(ReadU16Dev(blk)); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc index 817b19073c9eb..b31f6f0f37a79 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc @@ -3,15 +3,29 @@ #include "contrib_ops/cuda/quantization/matmul_packed.h" +#include #include #include "contrib_ops/cpu/quantization/gguf_quant_tables.h" +#include "core/common/safeint.h" +#include "core/framework/allocator.h" #include "core/providers/cuda/cuda_common.h" namespace onnxruntime { namespace contrib { namespace cuda { +namespace { +bool PackedTensorBytesFit(int64_t rows, int64_t columns, size_t element_bytes) { + size_t row_bytes = 0; + size_t bytes = 0; + return rows >= 0 && columns >= 0 && + IAllocator::CalcMemSizeForArray(static_cast(columns), element_bytes, &row_bytes) && + IAllocator::CalcMemSizeForArray(static_cast(rows), row_bytes, &bytes) && + bytes <= static_cast(std::numeric_limits::max()); +} +} // namespace + #define REGISTER_MATMUL_PACKED_CUDA(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ MatMulPacked, \ @@ -32,9 +46,18 @@ template MatMulPacked::MatMulPacked(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(info.GetAttr("K", &k_).IsOK() && k_ > 0, "MatMulPacked requires K > 0"); ORT_ENFORCE(info.GetAttr("N", &n_).IsOK() && n_ > 0, "MatMulPacked requires N > 0"); + ORT_ENFORCE(k_ <= kMatMulPackedMaxDimension && n_ <= kMatMulPackedMaxDimension, + "MatMulPacked requires K and N <= ", kMatMulPackedMaxDimension); + ORT_ENFORCE(k_ % 32 == 0, "MatMulPacked requires K to be a multiple of 32"); ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK() && ggml_type_ >= 0, "MatMulPacked requires packed_type >= 0"); ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(ggml_type_), "MatMulPacked: unsupported packed_type=", ggml_type_); + if constexpr (std::is_same_v) { + ORT_ENFORCE(GetDeviceProp().major >= 8, "MatMulPacked BF16 requires compute capability 8.0 or higher"); + if (n_ >= kMatMulPackedBf16TcMinN) { + ORT_THROW_IF_ERROR(GetMatMulPackedBf16KernelSupport(bf16_tc_supported_)); + } + } } template @@ -55,26 +78,39 @@ Status MatMulPacked::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(expected_rb > 0 && b_shape[1] == expected_rb, "MatMulPacked weight B row_bytes=", b_shape[1], " != expected=", expected_rb); + const int64_t m = a_shape.SizeToDimension(last_axis); + ORT_RETURN_IF_NOT(m >= 0 && m <= std::numeric_limits::max(), + "MatMulPacked requires M in 0..", std::numeric_limits::max()); + ORT_RETURN_IF_NOT((n_ + 7) / 8 <= GetDeviceProp().maxGridSize[0], + "MatMulPacked N exceeds the CUDA grid limit"); + const int64_t scratch_rows = std::min(n_, kMatMulPackedPrefillNTile); + ORT_RETURN_IF_NOT(PackedTensorBytesFit(m, k_, sizeof(T)) && + PackedTensorBytesFit(n_, expected_rb, 1) && + PackedTensorBytesFit(m, n_, sizeof(T)) && + (m <= 8 || PackedTensorBytesFit(scratch_rows, k_, sizeof(T))), + "MatMulPacked tensor byte size exceeds the supported address range"); + constexpr bool is_bf16 = std::is_same_v; + ORT_RETURN_IF_NOT(!is_bf16 || m == 0 || m > 8 || n_ < kMatMulPackedBf16TcMinN || bf16_tc_supported_, + "MatMulPacked BF16 Tensor Core decode requires kernels compiled for compute_80 or higher"); + MatMulPackedParams params{}; params.b_data = b->Data(); params.k = static_cast(k_); params.n = static_cast(n_); params.row_bytes = expected_rb; params.ggml_type = static_cast(ggml_type_); + params.bf16_tc_supported = bf16_tc_supported_; TensorShapeVector y_dims = a_shape.AsShapeVector(); y_dims[last_axis] = n_; Tensor* y = context->Output(0, TensorShape(y_dims)); - const int64_t m = a_shape.SizeToDimension(last_axis); if (m == 0) return Status::OK(); - constexpr bool is_bf16 = std::is_same_v; - ORT_RETURN_IF_NOT(!is_bf16 || GetDeviceProp().major >= 8, - "MatMulPacked BF16 requires compute capability 8.0 or higher"); IAllocatorUniquePtr scratch_w; if (m > 8) { - scratch_w = GetScratchBuffer(static_cast(n_ * k_), GetComputeStream(context)); + const size_t scratch_count = SafeInt(scratch_rows) * k_; + scratch_w = GetScratchBuffer(scratch_count, GetComputeStream(context)); } return LaunchMatMulPackedCuda( diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu index f8b4b483f511c..492c05d97a396 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu @@ -2,7 +2,12 @@ // Licensed under the MIT License. #include "contrib_ops/cuda/quantization/matmul_packed.h" + +#include +#include + #include "contrib_ops/cuda/quantization/gguf_quant_cuda.cuh" +#include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" namespace onnxruntime { @@ -55,7 +60,7 @@ __global__ __launch_bounds__(KW * 32) void MatMulPackedBf16TcWarpKSplitKernel( const int num_sb32 = p.k >> 5; for (int sb32_idx = warp; sb32_idx < num_sb32; sb32_idx += KW) { const uint4 b_vec = a_valid - ? __ldg(reinterpret_cast(a_row + sb32_idx * 32 + t * 8)) + ? gguf_cuda::LoadU4Dev(a_row + sb32_idx * 32 + t * 8) : make_uint4(0u, 0u, 0u, 0u); #pragma unroll for (int m = 0; m < MT; ++m) { @@ -133,7 +138,7 @@ __global__ __launch_bounds__(256) void MatMulPackedGemvKernel( #pragma unroll for (int m = 0; m < M_FIXED; ++m) { const CudaT* a_ptr = a + static_cast(m) * p.k + k_offset; - const uint4 v0 = __ldg(reinterpret_cast(a_ptr)); + const uint4 v0 = gguf_cuda::LoadU4Dev(a_ptr); const CudaT* a0 = reinterpret_cast(&v0); #pragma unroll for (int j = 0; j < 8; ++j) { @@ -143,7 +148,7 @@ __global__ __launch_bounds__(256) void MatMulPackedGemvKernel( } else { for (int m = 0; m < M; ++m) { const CudaT* a_ptr = a + static_cast(m) * p.k + k_offset; - const uint4 v0 = __ldg(reinterpret_cast(a_ptr)); + const uint4 v0 = gguf_cuda::LoadU4Dev(a_ptr); const CudaT* a0 = reinterpret_cast(&v0); #pragma unroll for (int j = 0; j < 8; ++j) { @@ -280,6 +285,10 @@ Status DispatchMatMulPackedCuda( CudaT* scratch_w, CudaT* y, int64_t m) { + ORT_RETURN_IF_NOT(params.k > 0 && params.k <= kMatMulPackedMaxDimension && + params.n > 0 && params.n <= kMatMulPackedMaxDimension && + m >= 0 && m <= std::numeric_limits::max(), + "MatMulPacked dimensions exceed the supported CUDA index range"); if (m == 0) return Status::OK(); ORT_RETURN_IF_NOT((params.k % 32) == 0, "MatMulPacked requires K to be a multiple of 32"); @@ -288,7 +297,9 @@ Status DispatchMatMulPackedCuda( const int m_int = static_cast(m); if constexpr (std::is_same::value) { - if (params.n >= 65536) { + if (params.n >= kMatMulPackedBf16TcMinN) { + ORT_RETURN_IF_NOT(params.bf16_tc_supported, + "MatMulPacked BF16 Tensor Core decode requires kernels compiled for compute_80 or higher"); LaunchMatMulPackedBf16TcWarpKSplit<4, 2>(stream, a, params, y, m_int); return CUDA_CALL(cudaGetLastError()); } @@ -317,11 +328,6 @@ Status DispatchMatMulPackedCuda( } ORT_RETURN_IF_NOT(scratch_w != nullptr, "MatMulPacked prefill requires non-null scratch_w"); - constexpr int kWarpsPerBlock = 8; - const dim3 grid(static_cast((params.n + kWarpsPerBlock - 1) / kWarpsPerBlock)); - const dim3 block(256); - MatMulPackedDequantKernel<<>>(params, scratch_w); - ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); // FP32 compute alone still permits FP16/BF16 split-K reductions. cublasMath_t previous_math_mode = CUBLAS_DEFAULT_MATH; @@ -335,26 +341,40 @@ Status DispatchMatMulPackedCuda( const float alpha = 1.0f; const float beta = 0.0f; const cudaDataType_t data_type = is_bf16 ? CUDA_R_16BF : CUDA_R_16F; - const auto gemm_status = CUBLAS_CALL(cublasGemmEx( - cublas, - CUBLAS_OP_T, - CUBLAS_OP_N, - params.n, - static_cast(m), - params.k, - &alpha, - scratch_w, - data_type, - params.k, - a, - data_type, - params.k, - &beta, - y, - data_type, - params.n, - CUBLAS_COMPUTE_32F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + Status gemm_status = Status::OK(); + for (int n0 = 0; n0 < params.n;) { + MatMulPackedParams tile_params = params; + tile_params.n = std::min(kMatMulPackedPrefillNTile, params.n - n0); + tile_params.b_data = params.b_data + static_cast(SafeInt(n0) * params.row_bytes); + const dim3 grid(static_cast((tile_params.n + 7) / 8)); + MatMulPackedDequantKernel<<>>(tile_params, scratch_w); + gemm_status = CUDA_CALL(cudaGetLastError()); + if (!gemm_status.IsOK()) break; + + // Column-major [N, M] is row-major [M, N]; retain the full N leading dimension for every tile. + gemm_status = CUBLAS_CALL(cublasGemmEx( + cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + tile_params.n, + static_cast(m), + params.k, + &alpha, + scratch_w, + data_type, + params.k, + a, + data_type, + params.k, + &beta, + y + n0, + data_type, + params.n, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + if (!gemm_status.IsOK()) break; + n0 += tile_params.n; + } if (math_mode != previous_math_mode) { CUBLAS_RETURN_IF_ERROR(cublasSetMathMode(cublas, previous_math_mode)); } @@ -363,6 +383,14 @@ Status DispatchMatMulPackedCuda( } // namespace +Status GetMatMulPackedBf16KernelSupport(bool& supported) { + cudaFuncAttributes attributes{}; + // All BF16 MMA instantiations in this translation unit use the same virtual architecture. + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, MatMulPackedBf16TcWarpKSplitKernel<4, 2, -1>)); + supported = attributes.ptxVersion >= 80; + return Status::OK(); +} + Status LaunchMatMulPackedCuda( cublasHandle_t cublas, cudaStream_t stream, diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h index b4760c5a796ab..c6bbe3f95e241 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h @@ -3,7 +3,9 @@ #pragma once +#include #include +#include #include "core/providers/cuda/cuda_kernel.h" @@ -13,14 +15,22 @@ namespace cuda { using namespace onnxruntime::cuda; +// Leave room for the largest (32-column) decode tile. +inline constexpr int64_t kMatMulPackedMaxDimension = std::numeric_limits::max() - 31; +inline constexpr int kMatMulPackedPrefillNTile = 4096; +inline constexpr int kMatMulPackedBf16TcMinN = 65536; + struct MatMulPackedParams { const uint8_t* b_data{nullptr}; int k{0}; int n{0}; int64_t row_bytes{0}; int ggml_type{0}; + bool bf16_tc_supported{false}; }; +Status GetMatMulPackedBf16KernelSupport(bool& supported); + Status LaunchMatMulPackedCuda( cublasHandle_t cublas, cudaStream_t stream, @@ -41,6 +51,7 @@ class MatMulPacked final : public CudaKernel { int64_t k_{0}; int64_t n_{0}; int64_t ggml_type_{-1}; + bool bf16_tc_supported_{false}; }; } // namespace cuda diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 9183ae5452d6d..1367ecf0db55b 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -996,14 +996,17 @@ ONNX_MS_OPERATOR_SET_SCHEMA( GatherPacked, 1, OpSchema() .SetDoc("Gather rows from a 2D block-packed embedding table and dequantize on the fly.") - .Attr("K", "Logical embedding dimension.", AttributeProto::INT) + .Attr("K", + "Logical embedding dimension. Must be a positive multiple of the selected packed format's block size.", + AttributeProto::INT) .Attr("packed_type", "Packed quantization format enum (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).", AttributeProto::INT) .Attr("output_dtype", "ONNX TensorProto DataType for output (default 16 = BFLOAT16).", AttributeProto::INT, static_cast(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16)) - .Input(0, "data", "2D packed weight blob [vocab_size, row_bytes].", "T1") - .Input(1, "indices", "Token indices of arbitrary rank.", "Tind") + .Input(0, "data", "2D little-endian packed weight blob [vocab_size, row_bytes].", "T1") + .Input(1, "indices", + "Token indices of arbitrary rank. Negative indices and indices >= vocab_size select row 0.", "Tind") .Output(0, "output", "Dequantized gathered embeddings [..., K].", "T2") .TypeConstraint("T1", {"tensor(uint8)"}, "Packed byte storage.") .TypeConstraint("T2", {"tensor(bfloat16)", "tensor(float16)", "tensor(float)"}, "Output floating-point type.") @@ -1027,8 +1030,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( MatMulPacked, 1, OpSchema() .SetDoc("Matrix multiplication Y = A @ B^T where B is a 2D block-packed weight tensor.\n" - "CUDA prefill dequantizes B to the activation type and uses float32 accumulation and reductions.") - .Attr("K", "Reduction dimension.", AttributeProto::INT) + "CUDA prefill dequantizes output-dimension tiles of B to the activation type using bounded scratch " + "memory and uses float32 accumulation and reductions.") + .Attr("K", + "Reduction dimension. Must be a whole number of packed blocks; CUDA requires a positive multiple of 32.", + AttributeProto::INT) .Attr("N", "Output dimension.", AttributeProto::INT) .Attr("packed_type", "Packed quantization format enum for B (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).", diff --git a/onnxruntime/test/python/transformers/test_gguf_packed_ops.py b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py index 7595dd05d4e4c..4cbd4e256d332 100644 --- a/onnxruntime/test/python/transformers/test_gguf_packed_ops.py +++ b/onnxruntime/test/python/transformers/test_gguf_packed_ops.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. +import ctypes +import sys import unittest import numpy as np @@ -42,6 +44,54 @@ def _round_to_dtype(arr, dtype): return arr.astype(np.float16).astype(np.float32) +def _cuda_bf16_supported(): + driver = ctypes.WinDLL("nvcuda.dll") if sys.platform == "win32" else ctypes.CDLL("libcuda.so.1") + + def check(status): + if status != 0: + raise RuntimeError(f"Cannot query CUDA compute capability (CUDA driver error {status})") + + device, major, minor = ctypes.c_int(), ctypes.c_int(), ctypes.c_int() + check(driver.cuInit(0)) + check(driver.cuDeviceGet(ctypes.byref(device), 0)) + check(driver.cuDeviceComputeCapability(ctypes.byref(major), ctypes.byref(minor), device)) + return major.value >= 8 + + +def _activation_storage(values, dtype): + if dtype == TensorProto.BFLOAT16: + return _f32_to_bf16_u16(values) + return np.ascontiguousarray(values, dtype=np.float16) + + +def _run_cuda_bound(session, inputs, output_shape, output_dtype, output_offset=2): + binding = session.io_binding() + owners = [] + for name, (dtype, values, offset) in inputs.items(): + storage = np.concatenate( + (np.zeros(offset, dtype=np.uint8), np.ascontiguousarray(values).view(np.uint8).ravel()) + ) + owner = ort.OrtValue.ortvalue_from_numpy(storage, "cuda", 0) + owners.append(owner) + binding.bind_input(name, "cuda", 0, dtype, values.shape, owner.data_ptr() + offset) + item_size = 4 if output_dtype == TensorProto.FLOAT else 2 + output_storage = ort.OrtValue.ortvalue_from_numpy( + np.zeros(output_offset + int(np.prod(output_shape)) * item_size, dtype=np.uint8), "cuda", 0 + ) + binding.bind_output("Y", "cuda", 0, output_dtype, output_shape, output_storage.data_ptr() + output_offset) + binding.synchronize_inputs() + session.run_with_iobinding(binding) + binding.synchronize_outputs() + raw = output_storage.numpy()[output_offset:] + if output_dtype == TensorProto.BFLOAT16: + return _bf16_u16_to_f32(raw.view(np.uint16)).reshape(output_shape) + return ( + raw.view(np.float32 if output_dtype == TensorProto.FLOAT else np.float16) + .reshape(output_shape) + .astype(np.float32) + ) + + def _session(model, provider): options = ort.SessionOptions() options.intra_op_num_threads = 1 @@ -167,7 +217,7 @@ def _moe_reference(x, logits, gate, up, down, top_k, dtype, normalize): def _make_synthetic_gguf_rows( - rng: np.random.Generator, rows: int, k: int, gtype: int, bel: int, bbytes: int + rng: np.random.Generator, rows: int, k: int, gtype: int, bel: int, bbytes: int, quant_scale: float = 2**-7 ) -> np.ndarray: row_bytes = (k // bel) * bbytes if gtype == 0: @@ -178,18 +228,83 @@ def _make_synthetic_gguf_rows( return _f32_to_bf16_u16(rng.standard_normal((rows, k)) * 0.1).view(np.uint8).reshape(rows, row_bytes) raw = rng.integers(0, 256, size=(rows, row_bytes), dtype=np.uint8) + scale_bytes = np.array([quant_scale], dtype=" 8 else weights + ref = _round_to_dtype(_round_to_dtype(a, dtype) @ target.T, dtype) + np.testing.assert_allclose(got, ref, rtol=1e-2, atol=1e-2) + + def test_moe_packed_unaligned_buffers(self): + rng = np.random.default_rng(1000) + num_experts, hidden_size, inter_size = 2, 256, 256 + for dtype in self.activation_dtypes: + for formats in (((1, 1, 2),) * 3, ((17, 256, 74), (17, 256, 74), (8, 32, 34))): + raw_weights, weights = [], [] + for out_size, k, (gtype, bel, bbytes) in zip( + (inter_size, inter_size, hidden_size), (hidden_size, hidden_size, inter_size), formats, strict=True + ): + raw = _make_synthetic_gguf_rows(rng, num_experts * out_size, k, gtype, bel, bbytes) + raw = raw.reshape(num_experts, out_size, -1) + raw_weights.append(raw) + weights.append(_dequantize_weights(raw, k, gtype)) + model = _model( + [ + helper.make_node( + "MoEPacked", + ["X", "R", "G", "U", "D"], + ["Y"], + domain="com.microsoft", + k=2, + gate_type=formats[0][0], + up_type=formats[1][0], + down_type=formats[2][0], + ) + ], + [ + helper.make_tensor_value_info("X", dtype, ["M", hidden_size]), + helper.make_tensor_value_info("R", dtype, ["M", num_experts]), + *[ + helper.make_tensor_value_info(name, TensorProto.UINT8, list(raw.shape)) + for name, raw in zip(("G", "U", "D"), raw_weights, strict=True) + ], + ], + [helper.make_tensor_value_info("Y", dtype, ["M", hidden_size])], + ) + session = _session(model, "CUDAExecutionProvider") + for m in (4, 48): + with self.subTest(dtype=dtype, formats=formats, m=m): + x = rng.standard_normal((m, hidden_size), dtype=np.float32) * 0.2 + logits = rng.standard_normal((m, num_experts), dtype=np.float32) + inputs = { + "X": (dtype, _activation_storage(x, dtype), 2), + "R": (dtype, _activation_storage(logits, dtype), 2), + } + inputs.update( + { + name: (TensorProto.UINT8, raw, 1) + for name, raw in zip(("G", "U", "D"), raw_weights, strict=True) + } + ) + got = _run_cuda_bound(session, inputs, (m, hidden_size), dtype) + ref = _moe_reference(x, logits, *weights, 2, dtype, True) + np.testing.assert_allclose(got, ref, rtol=2e-2, atol=2e-3) + if __name__ == "__main__": unittest.main() From c46a1daaee04ec13747b232bd394d7083f89ccdc Mon Sep 17 00:00:00 2001 From: rui-ren Date: Sun, 8 Nov 2026 21:58:00 +0000 Subject: [PATCH 5/5] Consolidate Bolt stack for Microsoft benchmark testing Start from current PR #33953 (2d9d48d4) and consolidate the upper three PRs #33962, #33963, and #33969 into a standalone test branch. Preserve the latest packed-op validation, alignment, BF16 capability, and bounded prefill scratch fixes alongside Bolt scaling, routing, precision, and Mamba snapshots. Cover scaled prefill across N tiles. Only the new benchmark branch is changed; the original stack remains untouched. CUDA/host/schema/test translation units and formatting pass. Runtime tests are unqualified because the borrowed overlay support runtime has incompatible symbols. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 77d657ce-b25d-4ca6-9727-5653117637e0 Signed-off-by: rui-ren --- docs/ContribOperators.md | 131 +++++- docs/OperatorKernels.md | 5 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 2 + .../contrib_ops/cuda/moe/moe_packed.cc | 134 ++++-- .../contrib_ops/cuda/moe/moe_packed.cu | 383 ++++++++++++++-- onnxruntime/contrib_ops/cuda/moe/moe_packed.h | 11 + .../cuda/quantization/gguf_quant_cuda.cuh | 15 + .../cuda/quantization/matmul_packed.cc | 46 +- .../cuda/quantization/matmul_packed.cu | 210 +++++++-- .../cuda/quantization/matmul_packed.h | 4 + .../cuda/sequence/FusedVerification.md | 82 ++++ .../cuda/sequence/fused_state_space.cc | 102 +++++ .../cuda/sequence/fused_state_space_impl.cu | 230 ++++++++++ .../cuda/sequence/fused_state_space_impl.h | 23 + .../core/graph/contrib_ops/contrib_defs.cc | 42 +- .../contrib_ops/fused_state_space_defs.h | 93 ++++ .../cuda_kernels/fused_state_space_test.cc | 173 ++++++++ .../contrib_ops/fused_state_space_op_test.cc | 420 ++++++++++++++++++ .../matmul_packed_precision_test.cc | 186 ++++++++ .../contrib_ops/moe_packed_routing_test.cc | 234 ++++++++++ 20 files changed, 2403 insertions(+), 123 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/sequence/FusedVerification.md create mode 100644 onnxruntime/contrib_ops/cuda/sequence/fused_state_space.cc create mode 100644 onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.cu create mode 100644 onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.h create mode 100644 onnxruntime/core/graph/contrib_ops/fused_state_space_defs.h create mode 100644 onnxruntime/test/contrib_ops/cuda_kernels/fused_state_space_test.cc create mode 100644 onnxruntime/test/contrib_ops/fused_state_space_op_test.cc create mode 100644 onnxruntime/test/contrib_ops/matmul_packed_precision_test.cc create mode 100644 onnxruntime/test/contrib_ops/moe_packed_routing_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 9fd8627459c0e..11587a9b573ce 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -66,6 +66,7 @@ Do not modify directly.* * com.microsoft.LinearAttentionGate * com.microsoft.LongformerAttention * com.microsoft.MRotaryEmbedding + * com.microsoft.Mamba2Verify * com.microsoft.MatMulBlockQuantizedFp4Weight * com.microsoft.MatMulBlockQuantizedFp8Weight * com.microsoft.MatMulBnb4 @@ -4057,6 +4058,94 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.Mamba2Verify** + + Fused causal convolution, selective state-space recurrence and gated grouped RMS normalization. + Local provisional contract: head_dim=64, d_state=128, 1<=R<=16. + projection[R, 2*di+2*n_groups*d_state+heads] packs [gate, xBC, dt], di=heads*head_dim. + state_row is a device int32[1] selecting the starting snapshot (0<=state_row
+ +#### Inputs + +
+
projection : T
+
Packed projection [R, width].
+
conv_state : T
+
Snapshots [slots, di+2*n_groups*d_state, 3].
+
ssm_state : F
+
Snapshots [slots, heads, d_state, head_dim].
+
conv_weight : T
+
Depthwise weights [channels,4] or [channels,1,4].
+
conv_bias : T
+
Depthwise convolution bias [channels].
+
dt_bias : F
+
Time-step bias [heads].
+
A_negative : F
+
Negative decay coefficient [heads], not A_log.
+
D : F
+
Skip coefficient [heads].
+
norm_weight : F
+
Normalization weight [di].
+
state_row : I
+
Starting snapshot index int32[1], on device.
+
+ +#### Outputs + +
+
Y : T
+
Normalized mixer result [R, di].
+
conv_snapshots : T
+
Post-token convolution snapshots, same shape as conv_state.
+
ssm_snapshots : F
+
Post-token recurrence snapshots, same shape as ssm_state.
+
xbc : T
+
BF16 convolution SiLU result [R, channels].
+
readout : F
+
FP32 recurrence readout before gating [R, di].
+
+ +#### Type Constraints + +
+
T : tensor(bfloat16)
+
BF16 activation and convolution parameters.
+
F : tensor(float)
+
FP32 recurrence and normalization.
+
I : tensor(int32)
+
Device snapshot selector.
+
+ + ### **com.microsoft.MatMulBlockQuantizedFp4Weight** Weight-only NVFP4 (E2M1) matrix multiplication. @@ -4747,17 +4836,21 @@ This version of the operator has been available since version 1 of the 'com.micr
Reduction dimension. Must be a whole number of packed blocks; CUDA requires a positive multiple of 32.
N : int (required)
Output dimension.
+
dequantize_to_activation_dtype : int
+
Round complete dequantized weights to A's dtype before decode accumulation. Default preserves the existing decode path; prefill already dequantizes to A's dtype.
packed_type : int (required)
Packed quantization format enum for B (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).
-#### Inputs +#### Inputs (2 - 3)
A : T1
Input activations [..., K].
B : T2
Packed weight blob [N, row_bytes].
+
global_scale (optional) : Tscale
+
Optional scalar FP32 multiplier for NVFP4, applied to FP32 accumulation before the output cast.
#### Outputs @@ -4774,6 +4867,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Activation type.
T2 : tensor(uint8)
Packed byte storage.
+
Tscale : tensor(float)
+
Global scale type.
@@ -4978,7 +5073,7 @@ This version of the operator has been available since version 1 of the 'com.micr
activation_type : string
-
Expert activation function ('swiglu').
+
Expert activation function ('swiglu' or 'relu_squared').
down_type : int (required)
Packed quantization format enum of down_experts_weights.
gate_type : int (required)
@@ -4987,23 +5082,37 @@ This version of the operator has been available since version 1 of the 'com.micr
Number of selected experts per token.
normalize_routing_weights : int
Whether to normalize top-k routing probabilities.
-
up_type : int (required)
-
Packed quantization format enum of up_experts_weights.
+
router_scale : float
+
Finite post-normalization scale for projected sigmoid routing.
+
routing_function : string
+
Internal routing: softmax logits or projected FP32 sigmoid with selection-only bias.
+
up_type : int
+
Packed quantization format enum of up_experts_weights; required for swiglu.
-#### Inputs +#### Inputs (5 - 10)
input : T
Input hidden states [..., hidden_size].
-
router_probs : T
-
Router logits [num_rows, num_experts].
+
router_probs (optional) : T
+
Router logits [num_rows, num_experts]; omit when supplying preselected routes.
gate_experts_weights : T2
3D packed gate weight blob [num_experts, inter_size, gate_row_bytes].
-
up_experts_weights : T2
-
3D packed up weight blob [num_experts, inter_size, up_row_bytes].
+
up_experts_weights (optional) : T2
+
3D packed up weight blob [num_experts, inter_size, up_row_bytes]; omit for relu_squared.
down_experts_weights : T2
3D packed down weight blob [num_experts, hidden_size, down_row_bytes].
+
expert_indices (optional) : Tind
+
Preselected distinct expert IDs [num_rows, k], k <= 32. Requires expert_weights, no router_probs and normalize_routing_weights=0. Invalid IDs poison only their token row.
+
expert_weights (optional) : Tweight
+
Preselected FP32 weights [num_rows, k], preserved without normalization and mixed in supplied order.
+
router_input (optional) : T
+
Projected sigmoid routing input [num_rows, router_width]; omit other routing inputs. Requires relu_squared, routing_function='sigmoid' and normalize_routing_weights=1.
+
router_weight (optional) : T
+
Sigmoid router weight [num_experts, router_width].
+
router_bias (optional) : Tweight
+
FP32 correction bias [num_experts], used only for expert selection.
#### Outputs @@ -5020,6 +5129,10 @@ This version of the operator has been available since version 1 of the 'com.micr
Activation type.
T2 : tensor(uint8)
Packed byte storage.
+
Tind : tensor(int32)
+
Preselected expert index type.
+
Tweight : tensor(float)
+
Preselected routing weight type.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 01358613ee6a5..56507718a106c 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1152,14 +1152,15 @@ The **OpSet Version** column uses the following notation: |LinearAttentionGate|*in* a:**T**
*in* dt_bias:**TF**
*in* decay_scale:**TF**
*in* b:**T**
*out* decay:**T**
*out* beta:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TF** = tensor(float)| |LongformerAttention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask:**T**
*in* global_weight:**T**
*in* global_bias:**T**
*in* global:**G**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |MRotaryEmbedding|*in* input:**T**
*in* position_ids:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**|1+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|Mamba2Verify|*in* projection:**T**
*in* conv_state:**T**
*in* ssm_state:**F**
*in* conv_weight:**T**
*in* conv_bias:**T**
*in* dt_bias:**F**
*in* A_negative:**F**
*in* D:**F**
*in* norm_weight:**F**
*in* state_row:**I**
*out* Y:**T**
*out* conv_snapshots:**T**
*out* ssm_snapshots:**F**
*out* xbc:**T**
*out* readout:**F**|1+|**F** = tensor(float)
**I** = tensor(int32)
**T** = tensor(bfloat16)| |MatMulBlockQuantizedFp4Weight|*in* A:**T**
*in* B:**T1**
*in* weight_scale:**T2**
*in* weight_scale_2:**T3**
*in* input_scale:**T3**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(uint8)
**T2** = tensor(uint8)
**T3** = tensor(float)| |MatMulBlockQuantizedFp8Weight|*in* A:**T**
*in* B:**T1**
*in* b_scale:**T2**
*in* a_scale:**T2**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn)
**T2** = tensor(float)| |MatMulBnb4|*in* A:**T1**
*in* B:**T2**
*in* absmax:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)| |MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8)| -|MatMulPacked|*in* A:**T1**
*in* B:**T2**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)| +|MatMulPacked|*in* A:**T1**
*in* B:**T2**
*in* global_scale:**Tscale**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)
**Tscale** = tensor(float)| |MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |MoEGroupRouter|*in* logits:**T**
*in* expert_bias:**TF**
*out* router_probs:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TF** = tensor(float)| -|MoEPacked|*in* input:**T**
*in* router_probs:**T**
*in* gate_experts_weights:**T2**
*in* up_experts_weights:**T2**
*in* down_experts_weights:**T2**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)| +|MoEPacked|*in* input:**T**
*in* router_probs:**T**
*in* gate_experts_weights:**T2**
*in* up_experts_weights:**T2**
*in* down_experts_weights:**T2**
*in* expert_indices:**Tind**
*in* expert_weights:**Tweight**
*in* router_input:**T**
*in* router_weight:**T**
*in* router_bias:**Tweight**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)
**Tind** = tensor(int32)
**Tweight** = tensor(float)| |MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**QK** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)| |NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index dfeb0a37d4b88..b2f5525af966a 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -176,6 +176,7 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MatMulPacked); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MatMulPacked); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoEPacked); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoEPacked); +class CUDA_MS_OP_CLASS_NAME(1, Mamba2Verify); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MRotaryEmbedding); @@ -564,6 +565,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc index ec40344492a2c..593ef1023b71f 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc @@ -4,6 +4,8 @@ #include "contrib_ops/cuda/moe/moe_packed.h" #include +#include +#include #include #include "contrib_ops/cpu/quantization/gguf_quant_tables.h" @@ -26,15 +28,17 @@ bool PackedTensorBytesFit(int64_t rows, int64_t columns, size_t element_bytes) { } } // namespace -#define REGISTER_MOE_PACKED_CUDA(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - MoEPacked, \ - kMSDomain, 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ +#define REGISTER_MOE_PACKED_CUDA(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MoEPacked, \ + kMSDomain, 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("Tweight", DataTypeImpl::GetTensorType()), \ MoEPacked); REGISTER_MOE_PACKED_CUDA(MLFloat16) @@ -44,21 +48,34 @@ REGISTER_MOE_PACKED_CUDA(BFloat16) template MoEPacked::MoEPacked(const OpKernelInfo& info) : CudaKernel(info) { - ORT_ENFORCE(info.GetAttr("k", &k_).IsOK() && k_ > 0 && k_ <= 16, - "MoEPacked requires k in 1..16"); + ORT_ENFORCE(info.GetAttr("k", &k_).IsOK() && k_ > 0 && k_ <= 32, + "MoEPacked requires k in 1..32; internally routed mode requires k <= 16"); normalize_routing_weights_ = info.GetAttrOrDefault("normalize_routing_weights", 1); activation_type_ = info.GetAttrOrDefault("activation_type", "swiglu"); - ORT_ENFORCE(activation_type_ == "swiglu", "MoEPacked only supports activation_type='swiglu'"); + routing_function_ = info.GetAttrOrDefault("routing_function", "softmax"); + router_scale_ = info.GetAttrOrDefault("router_scale", 1.0f); + ORT_ENFORCE(routing_function_ == "softmax" || routing_function_ == "sigmoid", + "MoEPacked routing_function must be softmax or sigmoid"); + ORT_ENFORCE(std::isfinite(router_scale_) && + (routing_function_ == "sigmoid" || router_scale_ == 1.0f), + "MoEPacked router_scale must be finite and non-default scaling requires sigmoid routing"); + ORT_ENFORCE(activation_type_ == "swiglu" || activation_type_ == "relu_squared", + "MoEPacked supports activation_type='swiglu' or 'relu_squared'"); ORT_ENFORCE(info.GetAttr("gate_type", &gate_type_).IsOK(), "MoEPacked requires gate_type"); - ORT_ENFORCE(info.GetAttr("up_type", &up_type_).IsOK(), "MoEPacked requires up_type"); + up_type_ = info.GetAttrOrDefault("up_type", -1); + ORT_ENFORCE(activation_type_ != "swiglu" || up_type_ >= 0, "MoEPacked swiglu requires up_type"); ORT_ENFORCE(info.GetAttr("down_type", &down_type_).IsOK(), "MoEPacked requires down_type"); ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(gate_type_), "MoEPacked: unsupported gate_type=", gate_type_); - ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(up_type_), "MoEPacked: unsupported up_type=", up_type_); + ORT_ENFORCE(activation_type_ != "swiglu" || gguf_quant::IsGgmlTypeSupported(up_type_), + "MoEPacked: unsupported up_type=", up_type_); ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(down_type_), "MoEPacked: unsupported down_type=", down_type_); if constexpr (std::is_same_v) { ORT_ENFORCE(GetDeviceProp().major >= 8, "MoEPacked BF16 requires compute capability 8.0 or higher"); ORT_THROW_IF_ERROR(CheckMoEPackedBf16KernelSupport()); } + ORT_ENFORCE(activation_type_ != "relu_squared" || + (gate_type_ == gguf_quant::kGgmlTypeIQ2_XXS && down_type_ == gguf_quant::kGgmlTypeIQ2_XXS), + "MoEPacked relu_squared currently requires IQ2_XXS gate and down weights"); } template @@ -68,50 +85,94 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { const Tensor* gate_w = context->Input(2); const Tensor* up_w = context->Input(3); const Tensor* down_w = context->Input(4); + const Tensor* indices = context->InputCount() > 5 ? context->Input(5) : nullptr; + const Tensor* weights = context->InputCount() > 6 ? context->Input(6) : nullptr; + const Tensor* router_input = context->InputCount() > 7 ? context->Input(7) : nullptr; + const Tensor* router_weight = context->InputCount() > 8 ? context->Input(8) : nullptr; + const Tensor* router_bias = context->InputCount() > 9 ? context->Input(9) : nullptr; + const bool sigmoid_routing = routing_function_ == "sigmoid"; + const bool preselected = indices || weights; + const bool relu_squared = activation_type_ == "relu_squared"; + ORT_RETURN_IF_NOT(input && gate_w && down_w, "MoEPacked requires input, gate and down weights"); + ORT_RETURN_IF_NOT(sigmoid_routing + ? (router_input && router_weight && router_bias && !router_probs && !preselected && + normalize_routing_weights_ == 1 && relu_squared) + : (!router_input && !router_weight && !router_bias), + "MoEPacked sigmoid routing requires router_input, router_weight, router_bias, " + "relu_squared and normalize_routing_weights=1, with no other routing inputs"); + ORT_RETURN_IF_NOT(sigmoid_routing || + (preselected ? (indices && weights && !router_probs && normalize_routing_weights_ == 0) + : (router_probs && k_ <= 16)), + "MoEPacked preselected routing requires both indices and weights, no router_probs and " + "normalize_routing_weights=0; internal routing requires router_probs and k <= 16"); + ORT_RETURN_IF_NOT(relu_squared ? !up_w : up_w != nullptr, + "MoEPacked swiglu requires up weights; relu_squared requires omitted up weights"); const auto& in_shape = input->Shape(); - const auto& rp_shape = router_probs->Shape(); const auto& gw_shape = gate_w->Shape(); - const auto& uw_shape = up_w->Shape(); const auto& dw_shape = down_w->Shape(); ORT_RETURN_IF_NOT(in_shape.NumDimensions() >= 2, "MoEPacked input must have rank >= 2"); - ORT_RETURN_IF_NOT(rp_shape.NumDimensions() == 2, "MoEPacked router_probs must be 2D [num_rows, num_experts]"); - ORT_RETURN_IF_NOT(gw_shape.NumDimensions() == 3 && uw_shape.NumDimensions() == 3 && dw_shape.NumDimensions() == 3, + ORT_RETURN_IF_NOT(gw_shape.NumDimensions() == 3 && dw_shape.NumDimensions() == 3 && + (!up_w || up_w->Shape().NumDimensions() == 3), "MoEPacked expert weights must be 3D [num_experts, out_dim, row_bytes]"); const size_t last_axis = in_shape.NumDimensions() - 1; const int64_t hidden_size = in_shape[last_axis]; const int64_t num_rows = in_shape.SizeToDimension(last_axis); - const int64_t num_experts = rp_shape[1]; + const int64_t num_experts = gw_shape[0]; const int64_t inter_size = gw_shape[1]; ORT_RETURN_IF_NOT(hidden_size > 0 && hidden_size <= kMoEPackedMaxDimension && inter_size > 0 && inter_size <= kMoEPackedMaxDimension, "MoEPacked requires hidden_size and inter_size in 1..", kMoEPackedMaxDimension); ORT_RETURN_IF_NOT(num_rows >= 0, "MoEPacked requires a nonnegative row count"); - ORT_RETURN_IF_NOT(rp_shape[0] == num_rows, "MoEPacked router_probs row count mismatch"); ORT_RETURN_IF_NOT(num_experts > 0 && num_experts <= 512, "MoEPacked supports up to 512 experts"); ORT_RETURN_IF_NOT(k_ <= num_experts, "MoEPacked requires k <= num_experts"); - ORT_RETURN_IF_NOT(gw_shape[0] == num_experts && uw_shape[0] == num_experts && dw_shape[0] == num_experts, + int64_t router_width = 0; + if (sigmoid_routing) { + ORT_RETURN_IF_NOT(router_input->Shape().NumDimensions() == 2 && + router_input->Shape()[0] == num_rows, + "MoEPacked sigmoid router_input must be [num_rows, router_width]"); + router_width = router_input->Shape()[1]; + ORT_RETURN_IF_NOT(router_width > 0 && router_width <= kMoEPackedMaxDimension && + router_width % 256 == 0 && + router_weight->Shape() == TensorShape({num_experts, router_width}) && + router_bias->Shape() == TensorShape({num_experts}), + "MoEPacked sigmoid routing requires positive int32 width divisible by 256, " + "router_weight[num_experts,width] and router_bias[num_experts]"); + } else if (preselected) { + ORT_RETURN_IF_NOT(indices->Shape() == TensorShape({num_rows, k_}) && weights->Shape() == indices->Shape(), + "MoEPacked preselected routes require indices and weights [num_rows, k]"); + } else { + ORT_RETURN_IF_NOT(router_probs->Shape() == TensorShape({num_rows, num_experts}), + "MoEPacked router_probs must be [num_rows, num_experts]"); + } + ORT_RETURN_IF_NOT((!up_w || up_w->Shape()[0] == num_experts) && dw_shape[0] == num_experts, "MoEPacked expert count mismatch"); - ORT_RETURN_IF_NOT(uw_shape[1] == inter_size && dw_shape[1] == hidden_size, + ORT_RETURN_IF_NOT((!up_w || up_w->Shape()[1] == inter_size) && dw_shape[1] == hidden_size, "MoEPacked expert inner/outer dimension mismatch"); - ORT_RETURN_IF_NOT((hidden_size % 32) == 0 && (inter_size % 32) == 0, + ORT_RETURN_IF_NOT(hidden_size > 0 && inter_size > 0 && + (hidden_size % 32) == 0 && (inter_size % 32) == 0, "MoEPacked requires hidden_size and inter_size to be multiples of 32"); const int64_t exp_gate_rb = gguf_quant::GgmlRowBytes(static_cast(gate_type_), hidden_size); - const int64_t exp_up_rb = gguf_quant::GgmlRowBytes(static_cast(up_type_), hidden_size); + const int64_t exp_up_rb = relu_squared ? exp_gate_rb : gguf_quant::GgmlRowBytes(static_cast(up_type_), hidden_size); const int64_t exp_down_rb = gguf_quant::GgmlRowBytes(static_cast(down_type_), inter_size); ORT_RETURN_IF_NOT(exp_gate_rb > 0 && exp_up_rb > 0 && exp_down_rb > 0 && - gw_shape[2] == exp_gate_rb && uw_shape[2] == exp_up_rb && dw_shape[2] == exp_down_rb, + gw_shape[2] == exp_gate_rb && (!up_w || up_w->Shape()[2] == exp_up_rb) && + dw_shape[2] == exp_down_rb, "MoEPacked expert weight row_bytes mismatch"); constexpr bool is_bf16 = std::is_same_v; ORT_RETURN_IF_NOT(PackedTensorBytesFit(num_rows, hidden_size, sizeof(T)) && PackedTensorBytesFit(num_rows, num_experts, sizeof(T)) && PackedTensorBytesFit(num_experts * inter_size, exp_gate_rb, 1) && PackedTensorBytesFit(num_experts * inter_size, exp_up_rb, 1) && - PackedTensorBytesFit(num_experts * hidden_size, exp_down_rb, 1), + PackedTensorBytesFit(num_experts * hidden_size, exp_down_rb, 1) && + (!preselected || (PackedTensorBytesFit(num_rows, k_, sizeof(int32_t)) && + PackedTensorBytesFit(num_rows, k_, sizeof(float)))) && + (!sigmoid_routing || (PackedTensorBytesFit(num_rows, router_width, sizeof(T)) && + PackedTensorBytesFit(num_experts, router_width, sizeof(T)))), "MoEPacked tensor byte size exceeds the supported address range"); // Process in chunks of at most 2048 rows so workspace stays compact during large prefills. @@ -132,8 +193,12 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { max_params.k = static_cast(k_); max_params.normalize_routing_weights = (normalize_routing_weights_ != 0); max_params.bf16_tc_supported = is_bf16; + max_params.relu_squared = relu_squared; + max_params.sigmoid_routing = sigmoid_routing; + max_params.router_width = static_cast(router_width); + max_params.router_scale = router_scale_; max_params.gate_type = static_cast(gate_type_); - max_params.up_type = static_cast(up_type_); + max_params.up_type = static_cast(relu_squared ? gate_type_ : up_type_); max_params.down_type = static_cast(down_type_); max_params.gate_row_bytes = exp_gate_rb; max_params.up_row_bytes = exp_up_rb; @@ -146,7 +211,7 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { auto workspace = GetScratchBuffer(ws_bytes, GetComputeStream(context)); const uint8_t* in_bytes = static_cast(input->DataRaw()); - const uint8_t* rp_bytes = static_cast(router_probs->DataRaw()); + const uint8_t* rp_bytes = router_probs ? static_cast(router_probs->DataRaw()) : nullptr; uint8_t* out_bytes = static_cast(output->MutableDataRaw()); for (int64_t r0 = 0; r0 < num_rows;) { @@ -155,13 +220,22 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const { chunk_params.num_rows = static_cast(chunk_rows); const size_t input_offset = SafeInt(r0) * hidden_size * sizeof(T); const size_t router_offset = SafeInt(r0) * num_experts * sizeof(T); + const size_t route_offset = SafeInt(r0) * k_; + chunk_params.expert_indices = indices ? indices->Data() + route_offset : nullptr; + chunk_params.expert_weights = weights ? weights->Data() + route_offset : nullptr; + chunk_params.router_input = router_input + ? static_cast(router_input->DataRaw()) + + static_cast(SafeInt(r0) * router_width * sizeof(T)) + : nullptr; + chunk_params.router_weight = router_weight ? router_weight->DataRaw() : nullptr; + chunk_params.router_bias = router_bias ? router_bias->Data() : nullptr; ORT_RETURN_IF_ERROR(LaunchMoEPackedCuda( Stream(context), is_bf16, in_bytes + input_offset, - rp_bytes + router_offset, + rp_bytes ? rp_bytes + router_offset : nullptr, gate_w->Data(), - up_w->Data(), + up_w ? up_w->Data() : gate_w->Data(), down_w->Data(), chunk_params, workspace.get(), diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu index 1b511ff86e297..e626515e58606 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu @@ -20,7 +20,134 @@ namespace cuda { namespace { constexpr int kMoEPackedMaxExperts = 512; -constexpr int kMoEPackedMaxK = 16; +constexpr int kMoEPackedMaxK = 32; + +template +__device__ __forceinline__ uint4 MoEPackedLoadRouter8(const CudaT* input) { + if ((reinterpret_cast(input) & 15u) == 0) { + return __ldg(reinterpret_cast(input)); + } + uint4 result; + CudaT* values = reinterpret_cast(&result); +#pragma unroll + for (int j = 0; j < 8; ++j) values[j] = input[j]; + return result; +} + +template +__global__ void MoEPackedSigmoidScoresKernel( + const CudaT* input, const CudaT* weight, int width, int experts, float* scores) { + const int row = blockIdx.x; + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int expert = blockIdx.y * 8 + warp; + if (expert >= experts) return; + const CudaT* x = input + static_cast(row) * width; + const CudaT* w = weight + static_cast(expert) * width; + float sum = 0.0f; + for (int offset = lane * 8; offset < width; offset += 256) { + const uint4 xv = MoEPackedLoadRouter8(x + offset); + const uint4 wv = MoEPackedLoadRouter8(w + offset); + const CudaT* xp = reinterpret_cast(&xv); + const CudaT* wp = reinterpret_cast(&wv); +#pragma unroll + for (int j = 0; j < 8; ++j) { + sum += gguf_cuda::CudaTypeToFloat(xp[j]) * gguf_cuda::CudaTypeToFloat(wp[j]); + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + sum += __shfl_down_sync(0xffffffffu, sum, offset); + } + if (lane == 0) scores[static_cast(row) * experts + expert] = 1.0f / (1.0f + expf(-sum)); +} + +__device__ __forceinline__ uint32_t MoEPackedSigmoidScoreRank(float value) { + const uint32_t bits = __float_as_uint(value); + const uint32_t magnitude = bits & 0x7fffffffu; + if (magnitude > 0x7f800000u) return 1u; + if (magnitude == 0) return 0x80000000u; + return (bits & 0x80000000u) ? ~bits : (bits ^ 0x80000000u); +} + +template +__device__ __forceinline__ void MoEPackedSelectSigmoidRoutes( + Rank (&choices)[16], const float (&scores)[16], int k, float scale, Rank empty, + int32_t* indices, float* weights) { + const int lane = threadIdx.x; + int selected = -1; + float selected_score = 0.0f; + float sum = 0.0f; + for (int rank = 0; rank < k; ++rank) { + Rank best = empty; + float best_score = 0.0f; + int best_expert = -1; +#pragma unroll + for (int slot = 0; slot < 16; ++slot) { + if (choices[slot] > best) { + best = choices[slot]; + best_score = scores[slot]; + best_expert = lane + 32 * slot; + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + const Rank other = __shfl_xor_sync(0xffffffffu, best, offset); + const int other_expert = __shfl_xor_sync(0xffffffffu, best_expert, offset); + if (other > best || (other == best && other_expert >= 0 && + (best_expert < 0 || other_expert < best_expert))) { + best = other; + best_expert = other_expert; + } + } + const int winner_lane = best_expert & 31; + const int winner_slot = best_expert >> 5; + if (best_expert >= 0 && winner_lane == lane) { +#pragma unroll + for (int slot = 0; slot < 16; ++slot) { + if (slot == winner_slot) choices[slot] = empty; + } + } + const float score = best_expert >= 0 ? __shfl_sync(0xffffffffu, best_score, winner_lane) : 0.0f; + if (lane == rank) { + selected = best_expert; + selected_score = score; + } + // Preserve selected-rank FP32 normalization and post-normalization scaling. + sum += score; + } + if (lane < k) { + indices[lane] = selected; + weights[lane] = selected_score / (sum + 1e-20f) * scale; + } +} + +__global__ void MoEPackedSigmoidTopKKernel( + const float* probabilities, const float* bias, int experts, int k, float scale, + int32_t* indices, float* weights) { + const int row = blockIdx.x; + const int lane = threadIdx.x; + float choices[16]; + float scores[16]; + bool finite = true; +#pragma unroll + for (int slot = 0; slot < 16; ++slot) { + const int expert = lane + 32 * slot; + scores[slot] = expert < experts ? probabilities[static_cast(row) * experts + expert] : 0.0f; + choices[slot] = expert < experts ? scores[slot] + bias[expert] : -INFINITY; + finite = finite && (expert >= experts || isfinite(choices[slot])); + } + indices += static_cast(row) * k; + weights += static_cast(row) * k; + if (__all_sync(0xffffffffu, finite)) { + MoEPackedSelectSigmoidRoutes(choices, scores, k, scale, -INFINITY, indices, weights); + return; + } + uint32_t ranks[16]; +#pragma unroll + for (int slot = 0; slot < 16; ++slot) { + ranks[slot] = lane + 32 * slot < experts ? MoEPackedSigmoidScoreRank(choices[slot]) : 0u; + } + MoEPackedSelectSigmoidRoutes(ranks, scores, k, scale, 0u, indices, weights); +} // Softmax + top-k over 512 experts for arbitrary num_rows (8 warps per CTA, 1 warp per row). // Uses the exact same warp reduction and tie-breaking key as FusedSoftmaxTopKMapsKernel. @@ -407,11 +534,13 @@ __global__ __launch_bounds__(KW * 32) void MoEPackedFc1DecodeWarpKSplitKernel( for (int m = 0; m < MT; ++m) { if (valid_rows[m]) { float g_vals[8]; - float u_vals[8]; + float u_vals[8] = {}; gguf_cuda::DecodeGgufSubBlock8SmemLutTpl( gate_ptrs[m], p.gate_type, sb32_idx, t, s_lut, g_vals); - gguf_cuda::DecodeGgufSubBlock8SmemLutTpl( - up_ptrs[m], p.up_type, sb32_idx, t, s_lut, u_vals); + if (!p.relu_squared) { + gguf_cuda::DecodeGgufSubBlock8SmemLutTpl( + up_ptrs[m], p.up_type, sb32_idx, t, s_lut, u_vals); + } const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(g_vals); const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(u_vals); gguf_cuda::MmaBf16(acc[m], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); @@ -445,9 +574,10 @@ __global__ __launch_bounds__(KW * 32) void MoEPackedFc1DecodeWarpKSplitKernel( } const float g_val = __bfloat162float(__float2bfloat16_rn(g_raw)); const float u_val = __bfloat162float(__float2bfloat16_rn(u_raw)); - const float swiglu = u_val * (g_val / (1.0f + expf(-g_val))); + const float positive = g_val < 0.f ? 0.f : g_val; + const float activation = p.relu_squared ? positive * positive : u_val * (g_val / (1.0f + expf(-g_val))); fc1_out[static_cast(s_routes[col]) * p.inter_size + j_out] = - __float2bfloat16_rn(swiglu); + __float2bfloat16_rn(activation); } } } @@ -582,9 +712,11 @@ __global__ __launch_bounds__((kMaxSubTiles == 1) ? 128 : 256) void MoEPackedFc1S for (int sb_local = 0; sb_local < num_sb32; ++sb_local) { const int sb32_idx = sb32_base + sb_local; float g_vals[8]; - float u_vals[8]; + float u_vals[8] = {}; gguf_cuda::DecodeGgufSubBlock8SmemLutTpl(gate_row_ptr, p.gate_type, sb32_idx, t, s_lut, g_vals); - gguf_cuda::DecodeGgufSubBlock8SmemLutTpl(up_row_ptr, p.up_type, sb32_idx, t, s_lut, u_vals); + if (!p.relu_squared) { + gguf_cuda::DecodeGgufSubBlock8SmemLutTpl(up_row_ptr, p.up_type, sb32_idx, t, s_lut, u_vals); + } const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(g_vals); const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(u_vals); #pragma unroll @@ -614,9 +746,10 @@ __global__ __launch_bounds__((kMaxSubTiles == 1) ? 128 : 256) void MoEPackedFc1S const float u_raw = (h == 0) ? acc[st][2] : acc[st][3]; const float g_val = __bfloat162float(__float2bfloat16_rn(g_raw)); const float u_val = __bfloat162float(__float2bfloat16_rn(u_raw)); - const float swiglu = u_val * (g_val / (1.0f + expf(-g_val))); + const float positive = g_val < 0.f ? 0.f : g_val; + const float activation = p.relu_squared ? positive * positive : u_val * (g_val / (1.0f + expf(-g_val))); const int route = s_routes[m]; - fc1_out[static_cast(route) * p.inter_size + j_out] = __float2bfloat16_rn(swiglu); + fc1_out[static_cast(route) * p.inter_size + j_out] = __float2bfloat16_rn(activation); } } } @@ -807,9 +940,17 @@ __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel( const int sb32_idx = h_idx >> 1; const int half16 = h_idx & 1; float g_vals[16]; - float u_vals[16]; + float u_vals[16] = {}; gguf_cuda::DecodeGgufSubBlock32(gate_row_ptr, p.gate_type, sb32_idx, half16, g_vals); - gguf_cuda::DecodeGgufSubBlock32(up_row_ptr, p.up_type, sb32_idx, half16, u_vals); + if (!p.relu_squared) { + gguf_cuda::DecodeGgufSubBlock32(up_row_ptr, p.up_type, sb32_idx, half16, u_vals); + } + if (p.relu_squared) { +#pragma unroll + for (int c = 0; c < 16; ++c) { + g_vals[c] = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(g_vals[c])); + } + } const int k_offset = h_idx * 16; for (int m = 0; m < count; ++m) { @@ -842,9 +983,10 @@ __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel( for (int m = 0; m < count; ++m) { const float g = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(gate_acc[m])); const float u = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(up_acc[m])); - const float swiglu = u * (g / (1.0f + expf(-g))); + const float positive = g < 0.f ? 0.f : g; + const float activation = p.relu_squared ? positive * positive : u * (g / (1.0f + expf(-g))); fc1_out[static_cast(sorted_routes[tile.start_route + m]) * p.inter_size + j_out] = - gguf_cuda::FloatToCudaType(swiglu); + gguf_cuda::FloatToCudaType(activation); } } } @@ -895,6 +1037,12 @@ __global__ __launch_bounds__(256) void MoEPackedFc2SimtKernel( const int half16 = h_idx & 1; float d_vals[16]; gguf_cuda::DecodeGgufSubBlock32(down_row_ptr, p.down_type, sb32_idx, half16, d_vals); + if (p.relu_squared) { +#pragma unroll + for (int c = 0; c < 16; ++c) { + d_vals[c] = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(d_vals[c])); + } + } const int k_offset = h_idx * 16; for (int m = 0; m < count; ++m) { @@ -989,6 +1137,21 @@ size_t AlignUpWorkspace(size_t count, size_t element_bytes = 1) { return bytes; } +__global__ void MoEPackedPreselectedRoutesKernel( + const int32_t* indices, const float* weights, int* output_indices, float* output_weights, + int count, int k, int num_experts) { + const int route = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (route >= count) return; + const int expert = indices[route]; + bool invalid = expert < 0 || expert >= num_experts || !isfinite(weights[route]); + for (int previous = (route / k) * k; previous < route; ++previous) { + invalid = invalid || indices[previous] == expert; + } + output_indices[route] = invalid ? 0 : expert; + // Preserve valid FP32 weights; NaN marks a malformed token without unsafe expert indexing. + output_weights[route] = invalid ? __int_as_float(0x7fffffff) : weights[route]; +} + void LaunchMoEPackedFc1DecodeBf16TcDispatch( cudaStream_t stream, dim3 grid, @@ -1166,6 +1329,118 @@ void LaunchMoEPackedFc2Bf16TcDispatch( } } +template +__global__ __launch_bounds__(KW * 32) void MoEPackedReluIq2ContiguousKKernel( + const __nv_bfloat16* input, + const uint8_t* weights, + const int* sorted_routes, + const MoEPackedTile* tiles, + const int* d_num_tiles, + MoEPackedParams p, + __nv_bfloat16* output) { + const int tile_index = static_cast(blockIdx.y); + if (tile_index >= __ldg(d_num_tiles)) return; + const MoEPackedTile tile = tiles[tile_index]; + const int first = static_cast(blockIdx.z) * 8; + const int remaining = tile.count - first; + if (remaining <= 0) return; + const int count = remaining < 8 ? remaining : 8; + constexpr int MT = 2; + __shared__ int routes[8]; + __shared__ float reduction[KW][MT * 16][8]; + const int tid = static_cast(threadIdx.x); + if (tid < 8) { + routes[tid] = tid < count ? __ldg(sorted_routes + tile.start_route + first + tid) : 0; + } + __syncthreads(); + const int warp = tid >> 5; + const int lane = tid & 31; + const int g = lane >> 2; + const int t = lane & 3; + const int n = kFc1 ? p.inter_size : p.hidden_size; + const int k = kFc1 ? p.hidden_size : p.inter_size; + const int64_t row_bytes = kFc1 ? p.gate_row_bytes : p.down_row_bytes; + const int base = static_cast(blockIdx.x) * (MT * 16); + const int input_row = kFc1 ? routes[g] / p.k : routes[g]; + const auto* x_row = input + static_cast(input_row) * k; + float acc[MT][4] = {}; + // IQ2 operands and activation fragments share contiguous 64-K partitions. + // FC1 and FC2 retain their four- and eight-warp FP32 reduction boundaries. + for (int chunk = warp; chunk < k / 64; chunk += KW) { + const int offset = chunk * 64 + t * 16; + const uint4 x0 = g < count ? MoEPackedLoadRouter8(x_row + offset) : make_uint4(0, 0, 0, 0); + const uint4 x1 = g < count ? MoEPackedLoadRouter8(x_row + offset + 8) : make_uint4(0, 0, 0, 0); + const uint32_t x_words[8] = {x0.x, x0.y, x0.z, x0.w, x1.x, x1.y, x1.z, x1.w}; +#pragma unroll + for (int m = 0; m < MT; ++m) { + const int low_row = base + m * 16 + g; + const int high_row = low_row + 8; + float low_values[16] = {}; + float high_values[16] = {}; + if (low_row < n) { + const auto* row = weights + (static_cast(tile.expert) * n + low_row) * row_bytes; + gguf_cuda::DecodeGgufSubBlock16Impl( + row, offset / 32, (offset / 16) & 1, low_values); + } + if (high_row < n) { + const auto* row = weights + (static_cast(tile.expert) * n + high_row) * row_bytes; + gguf_cuda::DecodeGgufSubBlock16Impl( + row, offset / 32, (offset / 16) & 1, high_values); + } + float e0[8], e1[8], o0[8], o1[8]; +#pragma unroll + for (int s = 0; s < 4; ++s) { + e0[2 * s] = low_values[4 * s]; + e0[2 * s + 1] = low_values[4 * s + 1]; + e1[2 * s] = low_values[4 * s + 2]; + e1[2 * s + 1] = low_values[4 * s + 3]; + o0[2 * s] = high_values[4 * s]; + o0[2 * s + 1] = high_values[4 * s + 1]; + o1[2 * s] = high_values[4 * s + 2]; + o1[2 * s + 1] = high_values[4 * s + 3]; + } + const uint4 e_low = gguf_cuda::PackFloat8ToBf16x2(e0); + const uint4 e_high = gguf_cuda::PackFloat8ToBf16x2(e1); + const uint4 o_low = gguf_cuda::PackFloat8ToBf16x2(o0); + const uint4 o_high = gguf_cuda::PackFloat8ToBf16x2(o1); + const uint32_t el[4] = {e_low.x, e_low.y, e_low.z, e_low.w}; + const uint32_t eh[4] = {e_high.x, e_high.y, e_high.z, e_high.w}; + const uint32_t ol[4] = {o_low.x, o_low.y, o_low.z, o_low.w}; + const uint32_t oh[4] = {o_high.x, o_high.y, o_high.z, o_high.w}; +#pragma unroll + for (int s = 0; s < 4; ++s) { + gguf_cuda::MmaBf16(acc[m], el[s], ol[s], eh[s], oh[s], x_words[2 * s], x_words[2 * s + 1]); + } + } + } +#pragma unroll + for (int m = 0; m < MT; ++m) { + reduction[warp][m * 16 + g][2 * t] = acc[m][0]; + reduction[warp][m * 16 + g][2 * t + 1] = acc[m][1]; + reduction[warp][m * 16 + g + 8][2 * t] = acc[m][2]; + reduction[warp][m * 16 + g + 8][2 * t + 1] = acc[m][3]; + } + __syncthreads(); + for (int index = tid; index < MT * 16 * 8; index += KW * 32) { + const int row = index / 8; + const int column = index & 7; + if (column < count && base + row < n) { + float sum = 0.f; +#pragma unroll + for (int w = 0; w < KW; ++w) sum += reduction[w][row][column]; + const auto rounded = __float2bfloat16_rn(sum); + if constexpr (kFc1) { + const float value = __bfloat162float(rounded); + const float positive = value < 0.f ? 0.f : value; + output[static_cast(routes[column]) * n + base + row] = + __float2bfloat16_rn(positive * positive); + } else { + output[static_cast(routes[column]) * n + base + row] = rounded; + } + } + } +} + template Status DispatchMoEPackedCuda( cudaStream_t stream, @@ -1182,7 +1457,8 @@ Status DispatchMoEPackedCuda( params.inter_size > 0 && params.inter_size <= kMoEPackedMaxDimension && params.hidden_size % 32 == 0 && params.inter_size % 32 == 0 && params.num_experts > 0 && params.num_experts <= kMoEPackedMaxExperts && - params.k > 0 && params.k <= kMoEPackedMaxK && params.k <= params.num_experts, + params.k > 0 && params.k <= kMoEPackedMaxK && params.k <= params.num_experts && + (params.expert_indices || params.sigmoid_routing || params.k <= 16), "MoEPacked parameters exceed the supported CUDA index range"); if (params.num_rows == 0) return Status::OK(); const int total_routes = params.num_rows * params.k; @@ -1217,7 +1493,20 @@ Status DispatchMoEPackedCuda( // 1. Softmax + TopK const dim3 route_grid(static_cast((params.num_rows + 7) / 8)); - if (params.num_experts == 512) { + if (params.sigmoid_routing) { + float* scores = reinterpret_cast( + ws_ptr + AlignUpWorkspace(SafeInt(total_routes) * params.hidden_size, sizeof(CudaT))); + MoEPackedSigmoidScoresKernel<<>>( + static_cast(params.router_input), static_cast(params.router_weight), + params.router_width, params.num_experts, scores); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + MoEPackedSigmoidTopKKernel<<>>( + scores, params.router_bias, params.num_experts, params.k, params.router_scale, topk_indices, topk_scales); + } else if (params.expert_indices) { + MoEPackedPreselectedRoutesKernel<<((total_routes + 255) / 256), 256, 0, stream>>>( + params.expert_indices, params.expert_weights, topk_indices, topk_scales, + total_routes, params.k, params.num_experts); + } else if (params.num_experts == 512) { MoEPackedSoftmaxTopK512Kernel<<>>( router_probs, topk_scales, topk_indices, params.num_rows, params.k, params.normalize_routing_weights); } else { @@ -1238,31 +1527,46 @@ Status DispatchMoEPackedCuda( ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); if constexpr (kIsBf16) { - // 3. Tensor-Core FC1 (64 inter_size channels/CTA in prefill, 16 inter_size channels/CTA Warp-K-Split in decode) - if (use_prefill_tiles) { - const dim3 fc1_grid(static_cast((params.inter_size + 63) / 64), - static_cast(max_tiles)); - LaunchMoEPackedFc1Bf16TcDispatch<4>( - stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + if (params.relu_squared && params.gate_type == gguf_quant::kGgmlTypeIQ2_XXS && + params.down_type == gguf_quant::kGgmlTypeIQ2_XXS) { + const unsigned int subtiles = use_prefill_tiles ? 4u : 1u; + const dim3 fc1_grid(static_cast((params.inter_size + 31) / 32), + static_cast(max_tiles), subtiles); + MoEPackedReluIq2ContiguousKKernel<<>>( + input, gate_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + const dim3 fc2_grid(static_cast((params.hidden_size + 31) / 32), + static_cast(max_tiles), subtiles); + MoEPackedReluIq2ContiguousKKernel<<>>( + fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); } else { - const dim3 fc1_grid(static_cast((params.inter_size + 15) / 16), - static_cast(max_tiles)); - LaunchMoEPackedFc1DecodeBf16TcDispatch( - stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); - } - ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + // 3. Tensor-Core FC1 (64 inter_size channels/CTA in prefill, 16 inter_size channels/CTA Warp-K-Split in decode) + if (use_prefill_tiles) { + const dim3 fc1_grid(static_cast((params.inter_size + 63) / 64), + static_cast(max_tiles)); + LaunchMoEPackedFc1Bf16TcDispatch<4>( + stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + } else { + const dim3 fc1_grid(static_cast((params.inter_size + 15) / 16), + static_cast(max_tiles)); + LaunchMoEPackedFc1DecodeBf16TcDispatch( + stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - // 4. Tensor-Core FC2 (128 hidden_size channels per CTA) - const dim3 fc2_grid(static_cast((params.hidden_size + 127) / 128), - static_cast(max_tiles)); - if (use_prefill_tiles) { - LaunchMoEPackedFc2Bf16TcDispatch<4>( - stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); - } else { - LaunchMoEPackedFc2Bf16TcDispatch<1>( - stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + // 4. Tensor-Core FC2 (128 hidden_size channels per CTA) + const dim3 fc2_grid(static_cast((params.hidden_size + 127) / 128), + static_cast(max_tiles)); + if (use_prefill_tiles) { + LaunchMoEPackedFc2Bf16TcDispatch<4>( + stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + } else { + LaunchMoEPackedFc2Bf16TcDispatch<1>( + stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); } - ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); } else { // Keep staging within CUDA's default shared-memory limit; larger shapes read activations directly. constexpr size_t kMaxStagingBytes = 48 * 1024; @@ -1340,6 +1644,9 @@ Status GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_byt add_buffer(SafeInt(total_routes) * params.inter_size, elem_bytes) && add_buffer(SafeInt(total_routes) * params.hidden_size, elem_bytes), "MoEPacked workspace byte size exceeds the supported address range"); + ORT_RETURN_IF_NOT(!params.sigmoid_routing || + add_buffer(SafeInt(params.num_rows) * params.num_experts, sizeof(float)), + "MoEPacked sigmoid workspace byte size exceeds the supported address range"); workspace_bytes = bytes; return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.h b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h index ca95f4c4ba757..7e752d8879494 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h @@ -35,6 +35,15 @@ struct MoEPackedParams { int k{0}; bool normalize_routing_weights{true}; bool bf16_tc_supported{false}; + bool relu_squared{false}; + bool sigmoid_routing{false}; + int router_width{0}; + float router_scale{1.0f}; + const void* router_input{nullptr}; + const void* router_weight{nullptr}; + const float* router_bias{nullptr}; + const int32_t* expert_indices{nullptr}; + const float* expert_weights{nullptr}; int gate_type{0}; int up_type{0}; int down_type{0}; @@ -69,6 +78,8 @@ class MoEPacked final : public CudaKernel { int64_t k_{0}; int64_t normalize_routing_weights_{1}; std::string activation_type_; + std::string routing_function_; + float router_scale_{1.0f}; int64_t gate_type_{-1}; int64_t up_type_{-1}; int64_t down_type_{-1}; diff --git a/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh b/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh index f074710e6855e..4afdfb08b6596 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -117,6 +118,17 @@ __device__ __forceinline__ void DecodeGgufSubBlock8Impl( out[j] = d * static_cast(static_cast((q0 >> (8 * j)) & 0xFFu)); out[j + 4] = d * static_cast(static_cast((q1 >> (8 * j)) & 0xFFu)); } + } else if constexpr (kCompileType == kGgmlTypeNVFP4) { + const uint8_t* blk = row_ptr + static_cast(sb32_idx >> 1) * 36; + const int group = 2 * (sb32_idx & 1) + (l >> 1); + const uint8_t scale = __ldg(blk + group); + const float d = scale == 127 ? 0.0f : __half2float(__half(__nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(scale), __NV_E4M3))); + const uint8_t* codes = blk + 4 + group * 8; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int code = (__ldg(codes + j) >> (4 * (l & 1))) & 15; + out[j] = d * (0.5f * static_cast(kvalues_fp4[code])); + } } else if constexpr (kCompileType == kGgmlTypeQ4_0) { const uint8_t* blk = row_ptr + static_cast(sb32_idx) * 18; const float d = Fp16ToFloatDev(ReadU16Dev(blk)); @@ -321,6 +333,9 @@ __device__ __forceinline__ void DecodeGgufSubBlock8( case kGgmlTypeQ8_0: DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); return; + case kGgmlTypeNVFP4: + DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); + return; case kGgmlTypeQ4_0: DecodeGgufSubBlock8Impl(row_ptr, sb32_idx, l, out); return; diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc index b31f6f0f37a79..10ffdada329b4 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cc @@ -26,15 +26,16 @@ bool PackedTensorBytesFit(int64_t rows, int64_t columns, size_t element_bytes) { } } // namespace -#define REGISTER_MATMUL_PACKED_CUDA(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - MatMulPacked, \ - kMSDomain, 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ +#define REGISTER_MATMUL_PACKED_CUDA(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MatMulPacked, \ + kMSDomain, 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("Tscale", DataTypeImpl::GetTensorType()), \ MatMulPacked); REGISTER_MATMUL_PACKED_CUDA(MLFloat16) @@ -51,10 +52,14 @@ MatMulPacked::MatMulPacked(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(k_ % 32 == 0, "MatMulPacked requires K to be a multiple of 32"); ORT_ENFORCE(info.GetAttr("packed_type", &ggml_type_).IsOK() && ggml_type_ >= 0, "MatMulPacked requires packed_type >= 0"); - ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(ggml_type_), "MatMulPacked: unsupported packed_type=", ggml_type_); + ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(ggml_type_) || ggml_type_ == gguf_quant::kGgmlTypeNVFP4, + "MatMulPacked: unsupported packed_type=", ggml_type_); + const int64_t round_weights = info.GetAttrOrDefault("dequantize_to_activation_dtype", 0); + ORT_ENFORCE(round_weights == 0 || round_weights == 1, "dequantize_to_activation_dtype must be 0 or 1"); + dequantize_to_activation_dtype_ = round_weights != 0; if constexpr (std::is_same_v) { ORT_ENFORCE(GetDeviceProp().major >= 8, "MatMulPacked BF16 requires compute capability 8.0 or higher"); - if (n_ >= kMatMulPackedBf16TcMinN) { + if (n_ >= kMatMulPackedBf16TcMinN || dequantize_to_activation_dtype_) { ORT_THROW_IF_ERROR(GetMatMulPackedBf16KernelSupport(bf16_tc_supported_)); } } @@ -64,7 +69,13 @@ template Status MatMulPacked::ComputeInternal(OpKernelContext* context) const { const Tensor* a = context->Input(0); const Tensor* b = context->Input(1); + const Tensor* global_scale = context->InputCount() > 2 ? context->Input(2) : nullptr; ORT_RETURN_IF_NOT(a != nullptr && b != nullptr, "MatMulPacked requires inputs A and B"); + ORT_RETURN_IF_NOT(!global_scale || + (ggml_type_ == gguf_quant::kGgmlTypeNVFP4 && + (global_scale->Shape().NumDimensions() == 0 || + global_scale->Shape() == TensorShape({1}))), + "MatMulPacked global_scale requires NVFP4 and a scalar or [1] shape"); const auto& a_shape = a->Shape(); ORT_RETURN_IF_NOT(a_shape.NumDimensions() >= 1, "MatMulPacked input A must have rank >= 1"); @@ -87,14 +98,18 @@ Status MatMulPacked::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(PackedTensorBytesFit(m, k_, sizeof(T)) && PackedTensorBytesFit(n_, expected_rb, 1) && PackedTensorBytesFit(m, n_, sizeof(T)) && - (m <= 8 || PackedTensorBytesFit(scratch_rows, k_, sizeof(T))), + (m <= 8 || PackedTensorBytesFit(scratch_rows, k_, sizeof(T))) && + (m <= 8 || !global_scale || PackedTensorBytesFit(m, n_, sizeof(float))), "MatMulPacked tensor byte size exceeds the supported address range"); constexpr bool is_bf16 = std::is_same_v; - ORT_RETURN_IF_NOT(!is_bf16 || m == 0 || m > 8 || n_ < kMatMulPackedBf16TcMinN || bf16_tc_supported_, + ORT_RETURN_IF_NOT(!is_bf16 || m == 0 || m > 8 || + (n_ < kMatMulPackedBf16TcMinN && !dequantize_to_activation_dtype_) || bf16_tc_supported_, "MatMulPacked BF16 Tensor Core decode requires kernels compiled for compute_80 or higher"); MatMulPackedParams params{}; params.b_data = b->Data(); + params.global_scale = global_scale ? global_scale->Data() : nullptr; + params.dequantize_to_activation_dtype = dequantize_to_activation_dtype_; params.k = static_cast(k_); params.n = static_cast(n_); params.row_bytes = expected_rb; @@ -108,9 +123,13 @@ Status MatMulPacked::ComputeInternal(OpKernelContext* context) const { if (m == 0) return Status::OK(); IAllocatorUniquePtr scratch_w; + IAllocatorUniquePtr scratch_y; if (m > 8) { const size_t scratch_count = SafeInt(scratch_rows) * k_; scratch_w = GetScratchBuffer(scratch_count, GetComputeStream(context)); + if (global_scale) { + scratch_y = GetScratchBuffer(SafeInt(m) * n_, GetComputeStream(context)); + } } return LaunchMatMulPackedCuda( @@ -120,6 +139,7 @@ Status MatMulPacked::ComputeInternal(OpKernelContext* context) const { a->DataRaw(), params, scratch_w ? scratch_w.get() : nullptr, + scratch_y ? scratch_y.get() : nullptr, y->MutableDataRaw(), m); } diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu index 492c05d97a396..c77dfd0ffa837 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu @@ -19,7 +19,7 @@ namespace { // Warp-K-Split Tensor-Core BF16 GEMV/GEMM kernel for M in 1..8 on large-N projections (e.g. lm_head). // KW warps (KW*32 threads) split K across the KW warps with zero __syncthreads() in the K-loop, // while computing MT*16 output columns per CTA. -template +template __global__ __launch_bounds__(KW * 32) void MatMulPackedBf16TcWarpKSplitKernel( const __nv_bfloat16* __restrict__ a, MatMulPackedParams p, @@ -57,26 +57,109 @@ __global__ __launch_bounds__(KW * 32) void MatMulPackedBf16TcWarpKSplitKernel( const bool a_valid = (g < m_rt); const __nv_bfloat16* a_row = a + static_cast(a_valid ? g : 0) * p.k; - const int num_sb32 = p.k >> 5; - for (int sb32_idx = warp; sb32_idx < num_sb32; sb32_idx += KW) { - const uint4 b_vec = a_valid - ? gguf_cuda::LoadU4Dev(a_row + sb32_idx * 32 + t * 8) - : make_uint4(0u, 0u, 0u, 0u); + if constexpr (kChunk != 32) { + // Keep contiguous 64/128-K partitions and their MMA operand order together; + // interleaving 32-K partitions changes cancellation before the warp reduction. + constexpr int kHalves = kChunk / 64; + for (int chunk = warp; chunk < p.k / kChunk; chunk += KW) { #pragma unroll - for (int m = 0; m < MT; ++m) { - float w_lo[8] = {}; - float w_hi[8] = {}; - if (valid_lo[m]) { - gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_lo[m], p.ggml_type, sb32_idx, t, w_lo); + for (int half = 0; half < kHalves; ++half) { + const int offset = chunk * kChunk + t * (kChunk / 4) + half * 16; + uint4 x0 = make_uint4(0u, 0u, 0u, 0u); + uint4 x1 = x0; + if (a_valid) { + const auto* input = a_row + offset; + if ((reinterpret_cast(input) & 15u) == 0) { + x0 = __ldg(reinterpret_cast(input)); + x1 = __ldg(reinterpret_cast(input + 8)); + } else { + const auto* bits = reinterpret_cast(input); + auto* low = reinterpret_cast(&x0); + auto* high = reinterpret_cast(&x1); +#pragma unroll + for (int i = 0; i < 4; ++i) { + low[i] = __ldg(bits + 2 * i) | (static_cast(__ldg(bits + 2 * i + 1)) << 16); + high[i] = __ldg(bits + 8 + 2 * i) | (static_cast(__ldg(bits + 9 + 2 * i)) << 16); + } + } + } + const uint32_t x_words[8] = {x0.x, x0.y, x0.z, x0.w, x1.x, x1.y, x1.z, x1.w}; +#pragma unroll + for (int m = 0; m < MT; ++m) { + float w_lo[16] = {}; + float w_hi[16] = {}; + if (valid_lo[m]) { + gguf_cuda::DecodeGgufSubBlock16Impl( + row_ptrs_lo[m], offset / 32, (offset / 16) & 1, w_lo); + } + if (valid_hi[m]) { + gguf_cuda::DecodeGgufSubBlock16Impl( + row_ptrs_hi[m], offset / 32, (offset / 16) & 1, w_hi); + } + float e0[8], e1[8], o0[8], o1[8]; +#pragma unroll + for (int s = 0; s < 4; ++s) { + constexpr bool adjacent = kCompileType == gguf_quant::kGgmlTypeIQ2_XXS; + const int i0 = adjacent ? 4 * s : s; + const int i1 = adjacent ? i0 + 1 : s + 4; + const int i2 = adjacent ? i0 + 2 : s + 8; + const int i3 = adjacent ? i0 + 3 : s + 12; + e0[2 * s] = w_lo[i0]; + e0[2 * s + 1] = w_lo[i1]; + e1[2 * s] = w_lo[i2]; + e1[2 * s + 1] = w_lo[i3]; + o0[2 * s] = w_hi[i0]; + o0[2 * s + 1] = w_hi[i1]; + o1[2 * s] = w_hi[i2]; + o1[2 * s + 1] = w_hi[i3]; + } + const uint4 e_lo = gguf_cuda::PackFloat8ToBf16x2(e0); + const uint4 e_hi = gguf_cuda::PackFloat8ToBf16x2(e1); + const uint4 o_lo = gguf_cuda::PackFloat8ToBf16x2(o0); + const uint4 o_hi = gguf_cuda::PackFloat8ToBf16x2(o1); + const uint32_t e_low[4] = {e_lo.x, e_lo.y, e_lo.z, e_lo.w}; + const uint32_t e_high[4] = {e_hi.x, e_hi.y, e_hi.z, e_hi.w}; + const uint32_t o_low[4] = {o_lo.x, o_lo.y, o_lo.z, o_lo.w}; + const uint32_t o_high[4] = {o_hi.x, o_hi.y, o_hi.z, o_hi.w}; +#pragma unroll + for (int s = 0; s < 4; ++s) { + uint32_t b0, b1; + if constexpr (kCompileType == gguf_quant::kGgmlTypeIQ2_XXS) { + b0 = x_words[2 * s]; + b1 = x_words[2 * s + 1]; + } else { + const int word = s / 2; + const int selector = (s & 1) ? 0x7632 : 0x5410; + b0 = __byte_perm(x_words[word], x_words[word + 2], selector); + b1 = __byte_perm(x_words[word + 4], x_words[word + 6], selector); + } + gguf_cuda::MmaBf16(acc[m], e_low[s], o_low[s], e_high[s], o_high[s], b0, b1); + } + } } - if (valid_hi[m]) { - gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_hi[m], p.ggml_type, sb32_idx, t, w_hi); + } + } else { + const int num_sb32 = p.k >> 5; + for (int sb32_idx = warp; sb32_idx < num_sb32; sb32_idx += KW) { + const uint4 b_vec = a_valid + ? gguf_cuda::LoadU4Dev(a_row + sb32_idx * 32 + t * 8) + : make_uint4(0u, 0u, 0u, 0u); +#pragma unroll + for (int m = 0; m < MT; ++m) { + float w_lo[8] = {}; + float w_hi[8] = {}; + if (valid_lo[m]) { + gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_lo[m], p.ggml_type, sb32_idx, t, w_lo); + } + if (valid_hi[m]) { + gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptrs_hi[m], p.ggml_type, sb32_idx, t, w_hi); + } + // Warp-synchronous MMA requires every lane, including padded output rows. + const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(w_lo); + const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(w_hi); + gguf_cuda::MmaBf16(acc[m], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); + gguf_cuda::MmaBf16(acc[m], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); } - // mma.sync requires every lane to participate, including partial output tiles. - const uint4 e_vec = gguf_cuda::PackFloat8ToBf16x2(w_lo); - const uint4 o_vec = gguf_cuda::PackFloat8ToBf16x2(w_hi); - gguf_cuda::MmaBf16(acc[m], e_vec.x, o_vec.x, e_vec.y, o_vec.y, b_vec.x, b_vec.y); - gguf_cuda::MmaBf16(acc[m], e_vec.z, o_vec.z, e_vec.w, o_vec.w, b_vec.z, b_vec.w); } } @@ -99,7 +182,8 @@ __global__ __launch_bounds__(KW * 32) void MatMulPackedBf16TcWarpKSplitKernel( for (int w = 0; w < KW; ++w) { sum += s_red[w][rl][row_m]; } - y[static_cast(row_m) * p.n + col] = __float2bfloat16_rn(sum); + y[static_cast(row_m) * p.n + col] = + __float2bfloat16_rn(p.global_scale ? sum * *p.global_scale : sum); } } } @@ -132,6 +216,13 @@ __global__ __launch_bounds__(256) void MatMulPackedGemvKernel( const int l = q_idx & 3; float w_vals[8]; gguf_cuda::DecodeGgufSubBlock8Tpl(row_ptr, p.ggml_type, sb32_idx, l, w_vals); + if (p.dequantize_to_activation_dtype) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + w_vals[j] = gguf_cuda::CudaTypeToFloat( + gguf_cuda::FloatToCudaType(w_vals[j])); + } + } const int k_offset = sb32_idx * 32 + (l << 3); if constexpr (M_FIXED > 0) { @@ -170,11 +261,13 @@ __global__ __launch_bounds__(256) void MatMulPackedGemvKernel( if constexpr (M_FIXED > 0) { #pragma unroll for (int m = 0; m < M_FIXED; ++m) { - y[static_cast(m) * p.n + n_out] = gguf_cuda::FloatToCudaType(acc[m]); + y[static_cast(m) * p.n + n_out] = + gguf_cuda::FloatToCudaType(p.global_scale ? acc[m] * *p.global_scale : acc[m]); } } else { for (int m = 0; m < M; ++m) { - y[static_cast(m) * p.n + n_out] = gguf_cuda::FloatToCudaType(acc[m]); + y[static_cast(m) * p.n + n_out] = + gguf_cuda::FloatToCudaType(p.global_scale ? acc[m] * *p.global_scale : acc[m]); } } } @@ -245,6 +338,15 @@ void LaunchMatMulPackedGemvByM( } } +template +__global__ void MatMulPackedScaleOutputKernel( + const float* input, const float* global_scale, CudaT* output, int64_t count) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + output[index] = gguf_cuda::FloatToCudaType(input[index] * *global_scale); + } +} + template void LaunchMatMulPackedBf16TcWarpKSplit( cudaStream_t stream, @@ -275,6 +377,31 @@ void LaunchMatMulPackedBf16TcWarpKSplit( } } +template +void LaunchMatMulPackedBf16ContiguousK( + cudaStream_t stream, + const __nv_bfloat16* a, + const MatMulPackedParams& params, + __nv_bfloat16* y, + int rows) { + const dim3 grid(static_cast((params.n + MT * 16 - 1) / (MT * 16))); + using namespace onnxruntime::contrib::gguf_quant; + switch (params.ggml_type) { + case kGgmlTypeQ8_0: + MatMulPackedBf16TcWarpKSplitKernel + <<>>(a, params, y, rows); + return; + case kGgmlTypeIQ2_XXS: + MatMulPackedBf16TcWarpKSplitKernel + <<>>(a, params, y, rows); + return; + case kGgmlTypeNVFP4: + MatMulPackedBf16TcWarpKSplitKernel + <<>>(a, params, y, rows); + return; + } +} + template Status DispatchMatMulPackedCuda( cublasHandle_t cublas, @@ -283,6 +410,7 @@ Status DispatchMatMulPackedCuda( const CudaT* a, const MatMulPackedParams& params, CudaT* scratch_w, + float* scratch_y, CudaT* y, int64_t m) { ORT_RETURN_IF_NOT(params.k > 0 && params.k <= kMatMulPackedMaxDimension && @@ -297,9 +425,28 @@ Status DispatchMatMulPackedCuda( const int m_int = static_cast(m); if constexpr (std::is_same::value) { - if (params.n >= kMatMulPackedBf16TcMinN) { + if (params.n >= kMatMulPackedBf16TcMinN || params.dequantize_to_activation_dtype) { ORT_RETURN_IF_NOT(params.bf16_tc_supported, "MatMulPacked BF16 Tensor Core decode requires kernels compiled for compute_80 or higher"); + } + if (params.dequantize_to_activation_dtype && params.k % 64 == 0 && + (params.ggml_type == kGgmlTypeQ8_0 || params.ggml_type == kGgmlTypeIQ2_XXS || + params.ggml_type == kGgmlTypeNVFP4)) { + const bool wide = params.ggml_type == kGgmlTypeNVFP4 && params.k % 128 == 0; + if (params.n >= 8192) { + if (wide) + LaunchMatMulPackedBf16ContiguousK<4, 2, 128>(stream, a, params, y, m_int); + else + LaunchMatMulPackedBf16ContiguousK<4, 2, 64>(stream, a, params, y, m_int); + } else { + if (wide) + LaunchMatMulPackedBf16ContiguousK<16, 1, 128>(stream, a, params, y, m_int); + else + LaunchMatMulPackedBf16ContiguousK<16, 1, 64>(stream, a, params, y, m_int); + } + return CUDA_CALL(cudaGetLastError()); + } + if (params.n >= kMatMulPackedBf16TcMinN || params.dequantize_to_activation_dtype) { LaunchMatMulPackedBf16TcWarpKSplit<4, 2>(stream, a, params, y, m_int); return CUDA_CALL(cudaGetLastError()); } @@ -328,6 +475,7 @@ Status DispatchMatMulPackedCuda( } ORT_RETURN_IF_NOT(scratch_w != nullptr, "MatMulPacked prefill requires non-null scratch_w"); + ORT_RETURN_IF_NOT(!params.global_scale || scratch_y, "Scaled prefill requires FP32 output workspace"); // FP32 compute alone still permits FP16/BF16 split-K reductions. cublasMath_t previous_math_mode = CUBLAS_DEFAULT_MATH; @@ -367,8 +515,8 @@ Status DispatchMatMulPackedCuda( data_type, params.k, &beta, - y + n0, - data_type, + params.global_scale ? static_cast(scratch_y + n0) : static_cast(y + n0), + params.global_scale ? CUDA_R_32F : data_type, params.n, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); @@ -378,7 +526,14 @@ Status DispatchMatMulPackedCuda( if (math_mode != previous_math_mode) { CUBLAS_RETURN_IF_ERROR(cublasSetMathMode(cublas, previous_math_mode)); } - return gemm_status; + ORT_RETURN_IF_ERROR(gemm_status); + if (params.global_scale) { + const int64_t count = m * params.n; + MatMulPackedScaleOutputKernel<<((count + 255) / 256), 256, 0, stream>>>( + scratch_y, params.global_scale, y, count); + return CUDA_CALL(cudaGetLastError()); + } + return Status::OK(); } } // namespace @@ -398,16 +553,17 @@ Status LaunchMatMulPackedCuda( const void* a, const MatMulPackedParams& params, void* scratch_w, + float* scratch_y, void* y, int64_t m) { if (is_bf16) { return DispatchMatMulPackedCuda<__nv_bfloat16>( cublas, stream, true, static_cast(a), params, - static_cast<__nv_bfloat16*>(scratch_w), static_cast<__nv_bfloat16*>(y), m); + static_cast<__nv_bfloat16*>(scratch_w), scratch_y, static_cast<__nv_bfloat16*>(y), m); } return DispatchMatMulPackedCuda( cublas, stream, false, static_cast(a), params, - static_cast(scratch_w), static_cast(y), m); + static_cast(scratch_w), scratch_y, static_cast(y), m); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h index c6bbe3f95e241..ebf97cc00a224 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_packed.h @@ -22,6 +22,8 @@ inline constexpr int kMatMulPackedBf16TcMinN = 65536; struct MatMulPackedParams { const uint8_t* b_data{nullptr}; + const float* global_scale{nullptr}; + bool dequantize_to_activation_dtype{false}; int k{0}; int n{0}; int64_t row_bytes{0}; @@ -38,6 +40,7 @@ Status LaunchMatMulPackedCuda( const void* a, const MatMulPackedParams& params, void* scratch_w, + float* scratch_y, void* y, int64_t m); @@ -52,6 +55,7 @@ class MatMulPacked final : public CudaKernel { int64_t n_{0}; int64_t ggml_type_{-1}; bool bf16_tc_supported_{false}; + bool dequantize_to_activation_dtype_{false}; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/sequence/FusedVerification.md b/onnxruntime/contrib_ops/cuda/sequence/FusedVerification.md new file mode 100644 index 0000000000000..95df2cc511734 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sequence/FusedVerification.md @@ -0,0 +1,82 @@ +# Snapshot-aware Mamba recurrence + +`com.microsoft::Mamba2Verify`, version 1, fuses causal convolution, selective +state-space recurrence and gated grouped RMS normalization. It is a built-in +CUDA contrib operator, with no custom-op plugin or external registration library. +This layer adds no projection, routing, MoE or squared-ReLU operators. + +## Contract + +The provisional implementation requires `head_dim=64`, `d_state=128`, +`heads % n_groups == 0` and 1 through 16 token rows. Eight rows use the preserved +specialized dispatch. Longer prefills must retain graph/runtime chunking; +this is not an arbitrary-length prefill kernel. + +The ten inputs and five outputs retain their frozen order: + +| Input | Type and layout | +|---|---| +| `projection` | BF16 `[R, 2*di + 2*n_groups*128 + heads]`, packing gate, xBC, dt | +| `conv_state` | BF16 `[slots, di + 2*n_groups*128, 3]` | +| `ssm_state` | FP32 `[slots, heads, 128, 64]` | +| `conv_weight` | BF16 `[channels, 4]` or `[channels, 1, 4]` | +| `conv_bias` | BF16 `[channels]` | +| `dt_bias`, `A_negative`, `D` | FP32 `[heads]`; decay is negative A, not A_log | +| `norm_weight` | FP32 `[di]`, where `di = heads*64` | +| `state_row` | Device INT32 `[1]`, selecting an input snapshot | + +Outputs are BF16 `Y[R, di]`, BF16 convolution snapshots, FP32 SSM snapshots, +BF16 `xbc[R, channels]` and FP32 pre-gating `readout[R, di]`. + +Each token writes its post-token state to the corresponding output row. +The caller selects the accepted/final snapshot on its next invocation; +the operator performs no accepted-length truncation or hidden state mutation. +Corresponding state inputs and outputs may alias. Trailing slots `[R, slots)` +retain their input values, including nonaliased outputs. + +An out-of-range device selector produces NaNs in computed rows without reading +either input state; trailing slots remain preserved. This is a non-finite output +policy, not a synchronous ORT Status error. Element-aligned convolution weights +use scalar loads when vector alignment is unavailable. + +The supplied ORT CUDA stream is used throughout. There are no host selector +reads, device-to-host copies, per-call CUDA allocations or synchronization. +Nonaliased state preservation uses capture-compatible device-to-device copies. + +## Provenance and precision + +The Mamba implementation is retained from the source-exact operator in #33940, +commit `491dc043952066267b2212ba979f0325da318573`, without recurrence, +dispatch or normalization arithmetic changes. Unused routing and squared-ReLU +helpers are omitted. + +Its frozen source is `export-nemotron-pr-20261030/src/bolt_ops.cu`, SHA256 +`1f587d9eacc356805c37d19e94e7aa4465dd50302576b2ca4b9579d7c5ae8ba5`. +Convolution SiLU rounds to BF16; recurrence and ordered readout retain FP32. +Grouped normalization gates, reduces and multiplies FP32 gamma before the final +BF16 output rounding. This describes the frozen Bolt contract, not a claim of +independent HuggingFace model parity. + +## Integration and tests + +Schema and CUDA registration are in `contrib_defs.cc` and +`cuda_contrib_kernels.cc`. Existing CMake globs collect the new source and provider +test files; regenerate CMake when adding them. CUDA internal tests require +`onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS`. + +Run the following on SM80+ and verify that every selected test executes: + +```text +FusedStateSpaceTest.* +FusedStateSpaceKernelTest.Mamba* +``` + +Provider tests cover all five outputs, 1/3/8/16 rows, nonzero snapshot selection, +1024-wide normalization, multi-group nonuniform states, rank-2 convolution +weights, each/both device-state aliases, invalid selectors and trailing/no-trailing +snapshot storage. Internal tests cover the production normalization launch with +65536 groups and offset convolution weights at one/eight rows. + +Component tests do not establish full-model readiness. Decoder parity, +generation, speculative state correctness, accuracy and performance require +separate qualification of the actual loaded runtime and model package. diff --git a/onnxruntime/contrib_ops/cuda/sequence/fused_state_space.cc b/onnxruntime/contrib_ops/cuda/sequence/fused_state_space.cc new file mode 100644 index 0000000000000..6f9b69e8f1df3 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sequence/fused_state_space.cc @@ -0,0 +1,102 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/narrow.h" +#include "core/providers/cuda/cuda_kernel.h" +#include "contrib_ops/cuda/sequence/fused_state_space_impl.h" + +#include +#include + +namespace onnxruntime::contrib::cuda { + +using onnxruntime::cuda::CudaKernel; + +class Mamba2Verify final : public CudaKernel { + public: + explicit Mamba2Verify(const OpKernelInfo& info) : CudaKernel(info) { + heads_ = narrow(info.GetAttrOrDefault("heads", 0)); + groups_ = narrow(info.GetAttrOrDefault("n_groups", 0)); + ORT_ENFORCE(heads_ > 0 && groups_ > 0 && heads_ % groups_ == 0 && + heads_ <= std::numeric_limits::max() / 1024); + ORT_ENFORCE(info.GetAttrOrDefault("head_dim", 64) == 64 && + info.GetAttrOrDefault("d_state", 128) == 128, + "Mamba2Verify currently requires head_dim=64, d_state=128"); + epsilon_ = info.GetAttrOrDefault("epsilon", 1e-5f); + time_step_min_ = info.GetAttrOrDefault("time_step_min", 0.001f); + ORT_ENFORCE(std::isfinite(epsilon_) && epsilon_ > 0 && + std::isfinite(time_step_min_) && time_step_min_ >= 0); + } + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Mamba2Verify); + + Status ComputeInternal(OpKernelContext* context) const override { + const auto* projection = context->Input(0); + const auto* conv = context->Input(1); + const auto* ssm = context->Input(2); + const int di = heads_ * 64; + const int channels = di + 2 * groups_ * 128; + ORT_RETURN_IF_NOT(projection->Shape().NumDimensions() == 2, "projection must have rank 2"); + const int64_t rows64 = projection->Shape()[0]; + ORT_RETURN_IF_NOT(rows64 >= 1 && rows64 <= 16 && + projection->Shape()[1] == di + channels + heads_, + "Invalid projection shape"); + const int rows = static_cast(rows64); + ORT_RETURN_IF_NOT(conv->Shape().NumDimensions() == 3 && conv->Shape()[0] >= rows && + conv->Shape()[1] == channels && conv->Shape()[2] == 3, + "Invalid convolution state shape"); + const int64_t slots = conv->Shape()[0]; + ORT_RETURN_IF_NOT(ssm->Shape() == TensorShape({slots, heads_, 128, 64}), "Invalid SSM state shape"); + const auto& conv_weight_shape = context->Input(3)->Shape(); + ORT_RETURN_IF_NOT((conv_weight_shape == TensorShape({channels, 4}) || + conv_weight_shape == TensorShape({channels, 1, 4})) && + context->Input(4)->Shape() == TensorShape({channels}), + "Invalid convolution parameters"); + for (int i = 5; i <= 7; ++i) { + ORT_RETURN_IF_NOT(context->Input(i)->Shape() == TensorShape({heads_}), "Invalid head parameter"); + } + ORT_RETURN_IF_NOT(context->Input(8)->Shape() == TensorShape({di}) && + context->Input(9)->Shape() == TensorShape({1}), + "Invalid norm weight or state_row"); + auto* y = context->Output(0, {rows, di}); + auto* conv_out = context->Output(1, conv->Shape()); + auto* ssm_out = context->Output(2, ssm->Shape()); + auto* xbc = context->Output(3, {rows, channels}); + auto* readout = context->Output(4, {rows, di}); + // Unwritten snapshot slots retain their input values, including when ORT does not alias buffers. + if (slots > rows && conv_out->MutableDataRaw() != conv->DataRaw()) { + const size_t offset = conv->SizeInBytes() / static_cast(slots) * rows; + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(static_cast(conv_out->MutableDataRaw()) + offset, + static_cast(conv->DataRaw()) + offset, + conv->SizeInBytes() - offset, cudaMemcpyDeviceToDevice, Stream(context))); + } + if (slots > rows && ssm_out->MutableDataRaw() != ssm->DataRaw()) { + const size_t offset = ssm->SizeInBytes() / static_cast(slots) * rows; + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(static_cast(ssm_out->MutableDataRaw()) + offset, + static_cast(ssm->DataRaw()) + offset, + ssm->SizeInBytes() - offset, cudaMemcpyDeviceToDevice, Stream(context))); + } + CUDA_RETURN_IF_ERROR(LaunchMamba2Verify( + Stream(context), rows, slots, heads_, groups_, epsilon_, time_step_min_, projection->DataRaw(), conv->DataRaw(), + ssm->Data(), context->Input(3)->DataRaw(), context->Input(4)->DataRaw(), + context->Input(5)->Data(), context->Input(6)->Data(), + context->Input(7)->Data(), context->Input(8)->Data(), + context->Input(9)->Data(), y->MutableDataRaw(), conv_out->MutableDataRaw(), + ssm_out->MutableData(), xbc->MutableDataRaw(), readout->MutableData())); + return Status::OK(); + } + + private: + int heads_, groups_; + float epsilon_, time_step_min_; +}; + +ONNX_OPERATOR_KERNEL_EX( + Mamba2Verify, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()) + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .MayInplace(1, 1) + .MayInplace(2, 2), + Mamba2Verify); +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.cu b/onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.cu new file mode 100644 index 0000000000000..4809bd1f7949a --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.cu @@ -0,0 +1,230 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/sequence/fused_state_space_impl.h" +#include + +namespace onnxruntime::contrib::cuda { +namespace { +using bf16 = __nv_bfloat16; +constexpr int kMaxRows = 16; +__device__ __forceinline__ float B2F(bf16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float Silu(float v) { return v / (1.0f + expf(-v)); } + +// Preserve the frozen fused path's FP32 ordering and BF16 convolution boundary. +template +__global__ void MambaConvKernel(const bf16* proj, int R_rt, int proj_stride, int di, int conv_dim, + const bf16* conv_state_in, const bf16* conv_w, const bf16* conv_b, bf16* conv_state_out, + bf16* xbc_act, const int32_t* sel, int64_t slots) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= conv_dim) return; + const int R = R_FIXED > 0 ? R_FIXED : R_rt; + if (*sel < 0 || static_cast(*sel) >= slots) { + for (int r = 0; r < R; ++r) { + xbc_act[static_cast(r) * conv_dim + c] = __float2bfloat16(NAN); + for (int k = 0; k < 3; ++k) + conv_state_out[(static_cast(r) * conv_dim + c) * 3 + k] = __float2bfloat16(NAN); + } + return; + } + conv_state_in += static_cast(*sel) * conv_dim * 3; + float inp[3 + (R_FIXED > 0 ? R_FIXED : kMaxRows)]; + inp[0] = B2F(conv_state_in[c * 3]); + inp[1] = B2F(conv_state_in[c * 3 + 1]); + inp[2] = B2F(conv_state_in[c * 3 + 2]); + const bf16* weight = conv_w + static_cast(c) * 4; + uint2 wv; + if ((reinterpret_cast(weight) & (alignof(uint2) - 1)) == 0) { + wv = __ldg(reinterpret_cast(weight)); + } else { + bf16* values = reinterpret_cast(&wv); +#pragma unroll + for (int i = 0; i < 4; ++i) values[i] = weight[i]; + } + const bf16* wp = reinterpret_cast(&wv); + const float w0 = B2F(wp[0]), w1 = B2F(wp[1]), w2 = B2F(wp[2]), w3 = B2F(wp[3]); + const float bias = B2F(conv_b[c]); + for (int r = 0; r < R; ++r) inp[3 + r] = B2F(proj[static_cast(r) * proj_stride + di + c]); + for (int r = 0; r < R; ++r) { + float acc = bias + w0 * inp[r] + w1 * inp[r + 1] + w2 * inp[r + 2] + w3 * inp[r + 3]; + xbc_act[static_cast(r) * conv_dim + c] = __float2bfloat16(Silu(acc)); + } + for (int r = 0; r < R; ++r) { + bf16* out = conv_state_out + (static_cast(r) * conv_dim + c) * 3; + out[0] = __float2bfloat16(inp[r + 1]); + out[1] = __float2bfloat16(inp[r + 2]); + out[2] = __float2bfloat16(inp[r + 3]); + } +} + +template +__global__ void MambaSsmKernel(const bf16* proj, int R_rt, int proj_stride, int di, int conv_dim, int heads, + int heads_per_group, const bf16* xbc_act, const float* ssm_in, const float* dt_bias, + const float* a_neg, const float* d_skip, float tmin, float* ssm_out, float* ypre, + const int32_t* sel, int64_t slots) { + constexpr int DS = 128, HD = 64, NH = DS / QP; + constexpr int MAX_R = R_FIXED > 0 ? R_FIXED : kMaxRows; + const int h = blockIdx.x, g = h / heads_per_group; + const int tid = threadIdx.x, p = tid % HD, half = tid / HD; + const int R = R_FIXED > 0 ? R_FIXED : R_rt; + // Uniform across the block and before either state read; safe during CUDA capture. + if (*sel < 0 || static_cast(*sel) >= slots) { + for (int r = 0; r < R; ++r) { + for (int i = tid; i < DS * HD; i += blockDim.x) + ssm_out[(static_cast(r) * heads + h) * DS * HD + i] = NAN; + for (int i = tid; i < HD; i += blockDim.x) ypre[static_cast(r) * di + h * HD + i] = NAN; + } + return; + } + __shared__ float sB[MAX_R][DS], sC[MAX_R][DS], sPart[MAX_R][QP][HD], sDt[MAX_R]; + const int b_off = di + g * DS; + const int c_off = di + (heads / heads_per_group) * DS + g * DS; + for (int i = tid; i < R * DS; i += blockDim.x) { + const int r = i / DS, n = i % DS; + sB[r][n] = B2F(xbc_act[static_cast(r) * conv_dim + b_off + n]); + sC[r][n] = B2F(xbc_act[static_cast(r) * conv_dim + c_off + n]); + } + if (tid < R) { + const float raw = B2F(proj[static_cast(tid) * proj_stride + di + conv_dim + h]) + dt_bias[h]; + float dt = raw > 20.0f ? raw : log1pf(expf(raw)); + sDt[tid] = fmaxf(dt, tmin); + } + __syncthreads(); + float st[NH]; + const int64_t row_stride = static_cast(heads) * DS * HD; + const float* sp = ssm_in + static_cast(*sel) * row_stride + + (static_cast(h) * DS + half * NH) * HD + p; +#pragma unroll + for (int n = 0; n < NH; ++n) st[n] = sp[static_cast(n) * HD]; + const float A = a_neg[h], D = d_skip[h]; + for (int r = 0; r < R; ++r) { + const float dt = sDt[r]; + const float dA = expf(A * dt); + const float x = B2F(xbc_act[static_cast(r) * conv_dim + h * HD + p]); + const float dtx = dt * x; + float acc = 0.f; +#pragma unroll + for (int n = 0; n < NH; ++n) { + const int nn = half * NH + n; + st[n] = st[n] * dA + dtx * sB[r][nn]; + acc += sC[r][nn] * st[n]; + } + sPart[r][half][p] = acc; + float* out = ssm_out + r * row_stride + (static_cast(h) * DS + half * NH) * HD + p; +#pragma unroll + for (int n = 0; n < NH; ++n) out[static_cast(n) * HD] = st[n]; + } + __syncthreads(); + for (int i = tid; i < R * HD; i += blockDim.x) { + const int r = i / HD, pp = i % HD; + const float x = B2F(xbc_act[static_cast(r) * conv_dim + h * HD + pp]); + float sum = D * x; +#pragma unroll + for (int q = 0; q < QP; ++q) sum += sPart[r][q][pp]; + ypre[static_cast(r) * di + h * HD + pp] = sum; + } +} + +__global__ void MambaNormKernel(const bf16* proj, int proj_stride, int di, int group_size, float eps, + const float* ypre, const float* norm_w, bf16* y) { + const int r = blockIdx.y, g = blockIdx.x; + const int64_t base = static_cast(r) * di + static_cast(g) * group_size; + const bf16* gate = proj + static_cast(r) * proj_stride + static_cast(g) * group_size; + __shared__ float red[32]; + if (group_size == 1024) { + const int tid = threadIdx.x; + float v[4], ss = 0.f; +#pragma unroll + for (int k = 0; k < 4; ++k) { + const int i = tid + k * 256; + const float gt = B2F(__ldg(gate + i)); + v[k] = __ldg(ypre + base + i) * (gt / (1.0f + expf(-gt))); + ss += v[k] * v[k]; + } + for (int o = 16; o > 0; o >>= 1) ss += __shfl_down_sync(0xffffffff, ss, o); + if ((tid & 31) == 0) red[tid >> 5] = ss; + __syncthreads(); + if (tid < 32) { + float w = tid < 8 ? red[tid] : 0.f; + for (int o = 16; o > 0; o >>= 1) w += __shfl_down_sync(0xffffffff, w, o); + if (tid == 0) red[0] = w; + } + __syncthreads(); + const float rs = rsqrtf(red[0] / group_size + eps); +#pragma unroll + for (int k = 0; k < 4; ++k) { + const int i = tid + k * 256; + y[base + i] = __float2bfloat16(v[k] * rs * __ldg(norm_w + g * 1024 + i)); + } + return; + } + float ss = 0.f; + for (int i = threadIdx.x; i < group_size; i += blockDim.x) { + const float gt = B2F(gate[i]); + const float v = ypre[base + i] * (gt / (1.0f + expf(-gt))); + ss += v * v; + } + for (int o = 16; o > 0; o >>= 1) ss += __shfl_down_sync(0xffffffff, ss, o); + if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = ss; + __syncthreads(); + if (threadIdx.x < 32) { + float v = threadIdx.x < 8 ? red[threadIdx.x] : 0.f; + for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o); + if (threadIdx.x == 0) red[0] = v; + } + __syncthreads(); + const float rs = rsqrtf(red[0] / group_size + eps); + for (int i = threadIdx.x; i < group_size; i += blockDim.x) { + const float gt = B2F(gate[i]); + const float v = ypre[base + i] * (gt / (1.0f + expf(-gt))); + y[base + i] = __float2bfloat16(v * rs * norm_w[g * group_size + i]); + } +} + +} // namespace + +cudaError_t LaunchMamba2VerifyNormalization(cudaStream_t stream, int rows, int width, int groups, int projection_stride, + float epsilon, const void* projection, const float* readout, + const float* norm_weight, void* output) { + MambaNormKernel<<>>(static_cast(projection), projection_stride, + width, width / groups, epsilon, readout, norm_weight, + static_cast(output)); + return cudaGetLastError(); +} + +cudaError_t LaunchMamba2Verify(cudaStream_t stream, int rows, int64_t slots, int heads, int groups, float epsilon, + float time_step_min, const void* projection, const void* conv_state, + const float* ssm_state, const void* conv_weight, const void* conv_bias, + const float* dt_bias, const float* a_negative, const float* d_skip, + const float* norm_weight, const int32_t* state_row, void* output, void* conv_snapshots, + float* ssm_snapshots, void* xbc, float* readout) { + const int di = heads * 64, channels = di + 2 * groups * 128, stride = di + channels + heads; + const auto* proj = static_cast(projection); + const auto* conv_in = static_cast(conv_state); + const auto* cw = static_cast(conv_weight); + const auto* cb = static_cast(conv_bias); + auto* conv_out = static_cast(conv_snapshots); + auto* act = static_cast(xbc); + if (rows == 8) { + MambaConvKernel<8><<<(channels + 63) / 64, 64, 0, stream>>>(proj, rows, stride, di, channels, conv_in, cw, cb, + conv_out, act, state_row, slots); + MambaSsmKernel<4, 8><<>>(proj, rows, stride, di, channels, heads, heads / groups, act, + ssm_state, dt_bias, a_negative, d_skip, time_step_min, + ssm_snapshots, readout, state_row, slots); + } else { + MambaConvKernel<<<(channels + 127) / 128, 128, 0, stream>>>(proj, rows, stride, di, channels, conv_in, cw, cb, + conv_out, act, state_row, slots); + if (rows > 1) { + MambaSsmKernel<4><<>>(proj, rows, stride, di, channels, heads, heads / groups, act, + ssm_state, dt_bias, a_negative, d_skip, time_step_min, ssm_snapshots, + readout, state_row, slots); + } else { + MambaSsmKernel<2><<>>(proj, rows, stride, di, channels, heads, heads / groups, act, + ssm_state, dt_bias, a_negative, d_skip, time_step_min, ssm_snapshots, + readout, state_row, slots); + } + } + return LaunchMamba2VerifyNormalization(stream, rows, di, groups, stride, epsilon, proj, readout, norm_weight, output); +} + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.h b/onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.h new file mode 100644 index 0000000000000..8321eeabddc87 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include + +namespace onnxruntime::contrib::cuda { + +// Shared production normalization launch; also permits small normalization-only kernel regressions. +cudaError_t LaunchMamba2VerifyNormalization(cudaStream_t stream, int rows, int width, int groups, int projection_stride, + float epsilon, const void* projection, const float* readout, + const float* norm_weight, void* output); + +// State snapshots use [rows, heads, state_size, head_size], not [rows, heads, head_size, state_size]. +// An invalid device selector produces NaN computed rows without reading either input state. +cudaError_t LaunchMamba2Verify(cudaStream_t stream, int rows, int64_t slots, int heads, int groups, float epsilon, + float time_step_min, const void* projection, const void* conv_state, + const float* ssm_state, const void* conv_weight, const void* conv_bias, + const float* dt_bias, const float* a_negative, const float* d_skip, + const float* norm_weight, const int32_t* state_row, void* output, void* conv_snapshots, + float* ssm_snapshots, void* xbc, float* readout); +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 936ba066de4c7..0387be64763b5 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -13,6 +13,7 @@ #include "core/graph/constants.h" #include "core/graph/contrib_ops/attn_lstm_schema_defs.h" #include "core/graph/contrib_ops/range_schema_defs.h" +#include "core/graph/contrib_ops/fused_state_space_defs.h" #include "core/graph/op.h" #include "core/mlas/inc/mlas.h" #include "core/mlas/inc/mlas_q4.h" @@ -1039,11 +1040,19 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("packed_type", "Packed quantization format enum for B (0..99 map to GGUF ggml_type; >=100 reserved for custom formats).", AttributeProto::INT) + .Attr("dequantize_to_activation_dtype", + "Round complete dequantized weights to A's dtype before decode accumulation. Default preserves " + "the existing decode path; prefill already dequantizes to A's dtype.", + AttributeProto::INT, static_cast(0)) .Input(0, "A", "Input activations [..., K].", "T1") .Input(1, "B", "Packed weight blob [N, row_bytes].", "T2") + .Input(2, "global_scale", + "Optional scalar FP32 multiplier for NVFP4, applied to FP32 accumulation before the output cast.", + "Tscale", OpSchema::Optional) .Output(0, "Y", "Output activations [..., N].", "T1") .TypeConstraint("T1", {"tensor(bfloat16)", "tensor(float16)"}, "Activation type.") .TypeConstraint("T2", {"tensor(uint8)"}, "Packed byte storage.") + .TypeConstraint("Tscale", {"tensor(float)"}, "Global scale type.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); const int64_t n = getAttribute(ctx, "N", 0); @@ -1065,19 +1074,43 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("k", "Number of selected experts per token.", AttributeProto::INT) .Attr("normalize_routing_weights", "Whether to normalize top-k routing probabilities.", AttributeProto::INT, static_cast(1)) - .Attr("activation_type", "Expert activation function ('swiglu').", + .Attr("activation_type", "Expert activation function ('swiglu' or 'relu_squared').", AttributeProto::STRING, std::string("swiglu")) + .Attr("routing_function", "Internal routing: softmax logits or projected FP32 sigmoid with selection-only bias.", + AttributeProto::STRING, std::string("softmax")) + .Attr("router_scale", "Finite post-normalization scale for projected sigmoid routing.", + AttributeProto::FLOAT, 1.0f) .Attr("gate_type", "Packed quantization format enum of gate_experts_weights.", AttributeProto::INT) - .Attr("up_type", "Packed quantization format enum of up_experts_weights.", AttributeProto::INT) + .Attr("up_type", "Packed quantization format enum of up_experts_weights; required for swiglu.", + AttributeProto::INT, static_cast(-1)) .Attr("down_type", "Packed quantization format enum of down_experts_weights.", AttributeProto::INT) .Input(0, "input", "Input hidden states [..., hidden_size].", "T") - .Input(1, "router_probs", "Router logits [num_rows, num_experts].", "T") + .Input(1, "router_probs", "Router logits [num_rows, num_experts]; omit when supplying preselected routes.", + "T", OpSchema::Optional) .Input(2, "gate_experts_weights", "3D packed gate weight blob [num_experts, inter_size, gate_row_bytes].", "T2") - .Input(3, "up_experts_weights", "3D packed up weight blob [num_experts, inter_size, up_row_bytes].", "T2") + .Input(3, "up_experts_weights", + "3D packed up weight blob [num_experts, inter_size, up_row_bytes]; omit for relu_squared.", + "T2", OpSchema::Optional) .Input(4, "down_experts_weights", "3D packed down weight blob [num_experts, hidden_size, down_row_bytes].", "T2") + .Input(5, "expert_indices", + "Preselected distinct expert IDs [num_rows, k], k <= 32. Requires expert_weights, no router_probs " + "and normalize_routing_weights=0. Invalid IDs poison only their token row.", + "Tind", OpSchema::Optional) + .Input(6, "expert_weights", + "Preselected FP32 weights [num_rows, k], preserved without normalization and mixed in supplied order.", + "Tweight", OpSchema::Optional) + .Input(7, "router_input", + "Projected sigmoid routing input [num_rows, router_width]; omit other routing inputs. " + "Requires relu_squared, routing_function='sigmoid' and normalize_routing_weights=1.", + "T", OpSchema::Optional) + .Input(8, "router_weight", "Sigmoid router weight [num_experts, router_width].", "T", OpSchema::Optional) + .Input(9, "router_bias", "FP32 correction bias [num_experts], used only for expert selection.", + "Tweight", OpSchema::Optional) .Output(0, "output", "Output hidden states [..., hidden_size].", "T") .TypeConstraint("T", {"tensor(bfloat16)", "tensor(float16)"}, "Activation type.") .TypeConstraint("T2", {"tensor(uint8)"}, "Packed byte storage.") + .TypeConstraint("Tind", {"tensor(int32)"}, "Preselected expert index type.") + .TypeConstraint("Tweight", {"tensor(float)"}, "Preselected routing weight type.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); if (hasInputShape(ctx, 0)) { @@ -3683,6 +3716,7 @@ static void MatmulWithQuantWeightShapeInference(ONNX_NAMESPACE::InferenceContext } void RegisterContribSchemas() { + ONNX_CONTRIB_OPERATOR_SCHEMA_ELSEWHERE(Mamba2Verify, Mamba2VerifySchema); ONNX_CONTRIB_OPERATOR_SCHEMA_ELSEWHERE(AttnLSTM, RegisterAttnLSTMContribOpSchema); ONNX_CONTRIB_OPERATOR_SCHEMA_ELSEWHERE(Range, RegisterRangeOpSchema); diff --git a/onnxruntime/core/graph/contrib_ops/fused_state_space_defs.h b/onnxruntime/core/graph/contrib_ops/fused_state_space_defs.h new file mode 100644 index 0000000000000..7fda99b3bf390 --- /dev/null +++ b/onnxruntime/core/graph/contrib_ops/fused_state_space_defs.h @@ -0,0 +1,93 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include "onnx/defs/schema.h" +#include "onnx/defs/shape_inference.h" +#include "core/graph/constants.h" +#include + +namespace onnxruntime::contrib { + +inline ONNX_NAMESPACE::OpSchema Mamba2VerifySchema(ONNX_NAMESPACE::OpSchema schema) { + using namespace ONNX_NAMESPACE; + return schema.SetDomain(kMSDomain).SinceVersion(1).SetDoc(R"DOC( +Fused causal convolution, selective state-space recurrence and gated grouped RMS normalization. +Local provisional contract: head_dim=64, d_state=128, 1<=R<=16. +projection[R, 2*di+2*n_groups*d_state+heads] packs [gate, xBC, dt], di=heads*head_dim. +state_row is a device int32[1] selecting the starting snapshot (0<=state_row