From 64a783b80cc8c6583027a4c9679fe56b597b8728 Mon Sep 17 00:00:00 2001 From: xindongliu594 Date: Thu, 3 Sep 2026 10:26:05 +0800 Subject: [PATCH 1/2] feat: add reusable GGUF linear backend and FP8 support --- include/infinicore/ops/block_fp8_linear.hpp | 25 ++ include/infinicore/ops/linear_gguf.hpp | 31 ++ include/infiniop.h | 1 + include/infiniop/ops/block_fp8_linear.h | 45 +++ include/infiniop/ops/linear_gguf.h | 48 +++ python/infinicore/__init__.py | 3 + python/infinicore/nn/functional/__init__.py | 2 + .../nn/functional/block_fp8_linear.py | 26 ++ python/infinicore/ops/block_fp8_linear.py | 31 ++ .../allocators/pinnable_block_allocator.cc | 58 ++-- .../ops/block_fp8_linear/block_fp8_linear.cc | 46 +++ .../block_fp8_linear_infiniop.cc | 63 ++++ src/infinicore/ops/linear_gguf/linear_gguf.cc | 47 +++ .../ops/linear_gguf/linear_gguf_infiniop.cc | 68 ++++ src/infinicore/pybind11/ops.hpp | 2 + .../pybind11/ops/block_fp8_linear.hpp | 38 +++ src/infiniop/ops/add_rms_norm/info.h | 38 ++- .../nvidia/add_rms_norm_nvidia.cu | 117 +++++-- .../ops/block_fp8_linear/block_fp8_linear.h | 38 +++ src/infiniop/ops/block_fp8_linear/info.h | 78 +++++ .../nvidia/block_fp8_gemm_sm120.cuh | 217 ++++++++++++ .../nvidia/block_fp8_linear_nvidia.cu | 189 +++++++++++ .../nvidia/block_fp8_linear_nvidia.cuh | 8 + src/infiniop/ops/block_fp8_linear/operator.cc | 98 ++++++ src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu | 129 ++++++- src/infiniop/ops/linear_gguf/ggml_blocks.h | 277 +++++++++++++++ src/infiniop/ops/linear_gguf/info.h | 84 +++++ src/infiniop/ops/linear_gguf/linear_gguf.h | 49 +++ .../nvidia/linear_gguf_dequant.cuh | 176 ++++++++++ .../linear_gguf/nvidia/linear_gguf_gemv.cuh | 315 ++++++++++++++++++ .../linear_gguf/nvidia/linear_gguf_nvidia.cu | 191 +++++++++++ .../linear_gguf/nvidia/linear_gguf_nvidia.cuh | 8 + src/infiniop/ops/linear_gguf/operator.cc | 100 ++++++ xmake.lua | 11 +- xmake/nvidia.lua | 39 ++- 35 files changed, 2628 insertions(+), 68 deletions(-) create mode 100644 include/infinicore/ops/block_fp8_linear.hpp create mode 100644 include/infinicore/ops/linear_gguf.hpp create mode 100644 include/infiniop/ops/block_fp8_linear.h create mode 100644 include/infiniop/ops/linear_gguf.h create mode 100644 python/infinicore/nn/functional/block_fp8_linear.py create mode 100644 python/infinicore/ops/block_fp8_linear.py create mode 100644 src/infinicore/ops/block_fp8_linear/block_fp8_linear.cc create mode 100644 src/infinicore/ops/block_fp8_linear/block_fp8_linear_infiniop.cc create mode 100644 src/infinicore/ops/linear_gguf/linear_gguf.cc create mode 100644 src/infinicore/ops/linear_gguf/linear_gguf_infiniop.cc create mode 100644 src/infinicore/pybind11/ops/block_fp8_linear.hpp create mode 100644 src/infiniop/ops/block_fp8_linear/block_fp8_linear.h create mode 100644 src/infiniop/ops/block_fp8_linear/info.h create mode 100644 src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh create mode 100644 src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu create mode 100644 src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cuh create mode 100644 src/infiniop/ops/block_fp8_linear/operator.cc create mode 100644 src/infiniop/ops/linear_gguf/ggml_blocks.h create mode 100644 src/infiniop/ops/linear_gguf/info.h create mode 100644 src/infiniop/ops/linear_gguf/linear_gguf.h create mode 100644 src/infiniop/ops/linear_gguf/nvidia/linear_gguf_dequant.cuh create mode 100644 src/infiniop/ops/linear_gguf/nvidia/linear_gguf_gemv.cuh create mode 100644 src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cu create mode 100644 src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cuh create mode 100644 src/infiniop/ops/linear_gguf/operator.cc diff --git a/include/infinicore/ops/block_fp8_linear.hpp b/include/infinicore/ops/block_fp8_linear.hpp new file mode 100644 index 000000000..b169c019f --- /dev/null +++ b/include/infinicore/ops/block_fp8_linear.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(BlockFP8Linear, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &); + +Tensor block_fp8_linear(const Tensor &input, + const Tensor &weight, + const Tensor &weight_scale); + +void block_fp8_linear_(Tensor output, + const Tensor &input, + const Tensor &weight, + const Tensor &weight_scale); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/linear_gguf.hpp b/include/infinicore/ops/linear_gguf.hpp new file mode 100644 index 000000000..41245e9e7 --- /dev/null +++ b/include/infinicore/ops/linear_gguf.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +// output = input @ dequant(weight)^T, with the weight left in its GGML block +// form: weight is a contiguous [N, row_bytes] U8 tensor and ggml_type is the +// enum ggml_type id of its blocks. See infiniop/ops/linear_gguf.h for the +// accepted types and the batch limit. +INFINICORE_GRAPH_OP_CLASS(LinearGguf, + Tensor, + const Tensor &, + const Tensor &, + int64_t); + +Tensor linear_gguf(const Tensor &input, + const Tensor &weight, + int64_t ggml_type); + +void linear_gguf_(Tensor output, + const Tensor &input, + const Tensor &weight, + int64_t ggml_type); + +} // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 9f632e27f..1631f465f 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -86,6 +86,7 @@ #include "infiniop/ops/ldexp.h" #include "infiniop/ops/lerp.h" #include "infiniop/ops/linear_mxfp4.h" +#include "infiniop/ops/block_fp8_linear.h" #include "infiniop/ops/log10.h" #include "infiniop/ops/log1p.h" #include "infiniop/ops/log_softmax.h" diff --git a/include/infiniop/ops/block_fp8_linear.h b/include/infiniop/ops/block_fp8_linear.h new file mode 100644 index 000000000..44813d406 --- /dev/null +++ b/include/infiniop/ops/block_fp8_linear.h @@ -0,0 +1,45 @@ +#ifndef __INFINIOP_BLOCK_FP8_LINEAR_API_H__ +#define __INFINIOP_BLOCK_FP8_LINEAR_API_H__ + +#include "../operator_descriptor.h" + +/** + * Block-wise FP8 (E4M3) linear operation with dynamic activation quantization. + * + * output: contiguous [M, N] BF16 + * input: contiguous [M, K] BF16 + * weight: contiguous [N, K] F8 (E4M3), used as logical B [K,N] column-major + * weight_scale: contiguous [ceil(N/128), ceil(K/128)] F32 (dequantization scale) + * + * Internally: BF16 activation -> per-128-group dynamic quant to FP8 E4M3 -> + * SM120 CUTLASS blockwise scaled GEMM -> BF16 output. + * Weight stays as 1 byte/element; no full-weight dequantization. + */ +typedef struct InfiniopDescriptor *infiniopBlockFP8LinearDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBlockFP8LinearDescriptor( + infiniopHandle_t handle, + infiniopBlockFP8LinearDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t weight_scale_desc); + +__INFINI_C __export infiniStatus_t infiniopGetBlockFP8LinearWorkspaceSize( + infiniopBlockFP8LinearDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBlockFP8Linear( + infiniopBlockFP8LinearDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *weight, + const void *weight_scale, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBlockFP8LinearDescriptor( + infiniopBlockFP8LinearDescriptor_t desc); + +#endif diff --git a/include/infiniop/ops/linear_gguf.h b/include/infiniop/ops/linear_gguf.h new file mode 100644 index 000000000..959228cf2 --- /dev/null +++ b/include/infiniop/ops/linear_gguf.h @@ -0,0 +1,48 @@ +#ifndef __INFINIOP_LINEAR_GGUF_API_H__ +#define __INFINIOP_LINEAR_GGUF_API_H__ + +#include "../operator_descriptor.h" +#include + +/** + * Linear over GGML block-quantized weights: + * out[M, N] = a[M, K] @ dequant(weight)[N, K]^T + * + * output: contiguous [M, N] BF16 + * input: contiguous [M, K] BF16 + * weight: contiguous [N, row_bytes] U8 -- the GGML block rows of one tensor, + * packed back to back verbatim, so + * row_bytes == (K / block_elems(ggml_type)) * block_bytes(ggml_type) + * ggml_type: enum ggml_type id of the weight blocks. Supported: 8 (Q8_0), + * 12 (Q4_K), 13 (Q5_K), 14 (Q6_K). Any other id is rejected here. + * + * The weight stays in its quantized form: blocks are decoded inside the kernel + * and accumulated in fp32, so a model loaded this way never materializes a + * dense copy of its weights. + * + * Current backend implements the decode (GEMV) path only, i.e. M must not + * exceed kMaxDecodeM from the NVIDIA kernel header. A larger M returns + * INFINI_STATUS_NOT_IMPLEMENTED rather than silently dequantizing the weight. + */ +typedef struct InfiniopDescriptor *infiniopLinearGgufDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLinearGgufDescriptor(infiniopHandle_t handle, + infiniopLinearGgufDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t w_desc, + int64_t ggml_type); + +__INFINI_C __export infiniStatus_t infiniopGetLinearGgufWorkspaceSize(infiniopLinearGgufDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLinearGguf(infiniopLinearGgufDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *a, + const void *w, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLinearGgufDescriptor(infiniopLinearGgufDescriptor_t desc); + +#endif // __INFINIOP_LINEAR_GGUF_API_H__ diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index 612db614c..8cf6efbb3 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -129,6 +129,7 @@ from infinicore.ops.logical_not import logical_not from infinicore.ops.masked_select import masked_select from infinicore.ops.matmul import matmul +from infinicore.ops.block_fp8_linear import block_fp8_linear, block_fp8_linear_ from infinicore.ops.mha import mha from infinicore.ops.mha_kvcache import mha_kvcache from infinicore.ops.mha_varlen import mha_varlen @@ -276,6 +277,8 @@ "logaddexp", "logaddexp2", "matmul", + "block_fp8_linear", + "block_fp8_linear_", "equal", "mul", "mul_scalar", diff --git a/python/infinicore/nn/functional/__init__.py b/python/infinicore/nn/functional/__init__.py index e84d5a718..ee0bfb8d1 100644 --- a/python/infinicore/nn/functional/__init__.py +++ b/python/infinicore/nn/functional/__init__.py @@ -24,6 +24,7 @@ from .layer_norm import layer_norm from .linear import linear from .linear_mxfp4 import linear_mxfp4 +from .block_fp8_linear import block_fp8_linear from .linear_w8a8i8 import linear_w8a8i8 from .log_softmax import log_softmax from .mamba_selective_scan import mamba_selective_scan @@ -69,6 +70,7 @@ "kimi_delta_attention", "linear", "linear_mxfp4", + "block_fp8_linear", "binary_cross_entropy_with_logits", "random_sample", "adaptive_avg_pool1d", diff --git a/python/infinicore/nn/functional/block_fp8_linear.py b/python/infinicore/nn/functional/block_fp8_linear.py new file mode 100644 index 000000000..a1f35a630 --- /dev/null +++ b/python/infinicore/nn/functional/block_fp8_linear.py @@ -0,0 +1,26 @@ +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def block_fp8_linear( + input: Tensor, + weight: Tensor, + weight_scale: Tensor, + out=None, +) -> Tensor: + if out is None: + return Tensor( + _infinicore.block_fp8_linear( + input._underlying, + weight._underlying, + weight_scale._underlying, + ) + ) + + _infinicore.block_fp8_linear_( + out._underlying, + input._underlying, + weight._underlying, + weight_scale._underlying, + ) + return out diff --git a/python/infinicore/ops/block_fp8_linear.py b/python/infinicore/ops/block_fp8_linear.py new file mode 100644 index 000000000..a85da2567 --- /dev/null +++ b/python/infinicore/ops/block_fp8_linear.py @@ -0,0 +1,31 @@ +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def block_fp8_linear(input, weight, weight_scale): + """Block-FP8 linear: BF16 input x F8 weight + block scale -> BF16 output. + + Args: + input: BF16 tensor [M, K] + weight: F8 tensor [N, K] + weight_scale: F32 tensor [ceil(N/128), ceil(K/128)] + + Returns: + BF16 tensor [M, N] + """ + return Tensor(_infinicore.block_fp8_linear( + input._underlying, weight._underlying, weight_scale._underlying)) + + +def block_fp8_linear_(output, input, weight, weight_scale): + """In-place block-FP8 linear. + + Args: + output: pre-allocated BF16 tensor [M, N] + input: BF16 tensor [M, K] + weight: F8 tensor [N, K] + weight_scale: F32 tensor [ceil(N/128), ceil(K/128)] + """ + _infinicore.block_fp8_linear_( + output._underlying, input._underlying, + weight._underlying, weight_scale._underlying) diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc index 32e5c5e9b..453a70516 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc @@ -47,37 +47,41 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { std::shared_ptr block; - // 1. Try size-class allocation for small/medium - for (auto &cls : size_classes_) { - if (size <= cls.block_size) { - if (!cls.free_blocks.empty()) { - block = cls.free_blocks.back(); - while (block != nullptr && block->in_use) { - cls.free_blocks.pop_back(); - if (cls.free_blocks.empty()) { - block = nullptr; - break; - } + // 1. Try size-class allocation for small only (<=1MB) + // For larger allocations, skip size-class to avoid massive internal fragmentation + // (e.g. 85MB tensor in 128MB block wastes 34% memory) + if (size <= 1 * 1024 * 1024) { + for (auto &cls : size_classes_) { + if (size <= cls.block_size) { + if (!cls.free_blocks.empty()) { block = cls.free_blocks.back(); + while (block != nullptr && block->in_use) { + cls.free_blocks.pop_back(); + if (cls.free_blocks.empty()) { + block = nullptr; + break; + } + block = cls.free_blocks.back(); + } + if (block != nullptr) { + cls.free_blocks.pop_back(); + block->in_use = true; + block->use_count = 1; + return reinterpret_cast(block->ptr); + } } - if (block != nullptr) { - cls.free_blocks.pop_back(); - block->in_use = true; - block->use_count = 1; - return reinterpret_cast(block->ptr); - } - } - // Allocate a new block for this class - block = std::make_shared(); - block->size = cls.block_size; - block->frozen = pinned_mode_; - block->in_use = true; - block->use_count = 1; + // Allocate a new block for this class. + block = std::make_shared(); + block->size = cls.block_size; + block->frozen = pinned_mode_; + block->in_use = true; + block->use_count = 1; - INFINICORE_CHECK_ERROR(infinirtMalloc(&block->ptr, block->size)); + INFINICORE_CHECK_ERROR(infinirtMalloc(&block->ptr, block->size)); - all_blocks_[block->ptr] = block; - return reinterpret_cast(block->ptr); + all_blocks_[block->ptr] = block; + return reinterpret_cast(block->ptr); + } } } diff --git a/src/infinicore/ops/block_fp8_linear/block_fp8_linear.cc b/src/infinicore/ops/block_fp8_linear/block_fp8_linear.cc new file mode 100644 index 000000000..2f17de762 --- /dev/null +++ b/src/infinicore/ops/block_fp8_linear/block_fp8_linear.cc @@ -0,0 +1,46 @@ +#include "infinicore/ops/block_fp8_linear.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlockFP8Linear); + +BlockFP8Linear::BlockFP8Linear(Tensor output, + const Tensor &input, + const Tensor &weight, + const Tensor &weight_scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + output, input, weight, weight_scale); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().getType(), output, input, weight, weight_scale); +} + +void BlockFP8Linear::execute(Tensor output, + const Tensor &input, + const Tensor &weight, + const Tensor &weight_scale) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + BlockFP8Linear, output, input, weight, weight_scale); +} + +Tensor block_fp8_linear(const Tensor &input, + const Tensor &weight, + const Tensor &weight_scale) { + INFINICORE_ASSERT(input->ndim() >= 2); + INFINICORE_ASSERT(weight->ndim() == 2); + auto output_shape = input->shape(); + output_shape.back() = weight->size(0); + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + block_fp8_linear_(output, input, weight, weight_scale); + return output; +} + +void block_fp8_linear_(Tensor output, + const Tensor &input, + const Tensor &weight, + const Tensor &weight_scale) { + BlockFP8Linear::execute(output, input, weight, weight_scale); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/block_fp8_linear/block_fp8_linear_infiniop.cc b/src/infinicore/ops/block_fp8_linear/block_fp8_linear_infiniop.cc new file mode 100644 index 000000000..95bd3bff9 --- /dev/null +++ b/src/infinicore/ops/block_fp8_linear/block_fp8_linear_infiniop.cc @@ -0,0 +1,63 @@ +#include "infinicore/ops/block_fp8_linear.hpp" + +#include "../infiniop_impl.hpp" + +namespace infinicore::op::block_fp8_linear_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, BlockFP8Linear, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace; + graph::GraphTensor output; + graph::GraphTensor input; + graph::GraphTensor weight; + graph::GraphTensor weight_scale; +}; + +void *plan(Tensor output, + const Tensor &input, + const Tensor &weight, + const Tensor &weight_scale) { + size_t seed = hash_combine(output, input, weight, weight_scale); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, + descriptor, + BlockFP8Linear, + seed, + output->desc(), + input->desc(), + weight->desc(), + weight_scale->desc()); + INFINIOP_WORKSPACE_TENSOR(workspace, BlockFP8Linear, descriptor); + + return new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(output), + graph::GraphTensor(input), + graph::GraphTensor(weight), + graph::GraphTensor(weight_scale)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopBlockFP8Linear( + planned->descriptor->desc, + planned->workspace->data(), + planned->workspace->numel(), + planned->output->data(), + planned->input->data(), + planned->weight->data(), + planned->weight_scale->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(BlockFP8Linear, &plan, &run, &cleanup); + +} // namespace infinicore::op::block_fp8_linear_impl::infiniop diff --git a/src/infinicore/ops/linear_gguf/linear_gguf.cc b/src/infinicore/ops/linear_gguf/linear_gguf.cc new file mode 100644 index 000000000..f4d322126 --- /dev/null +++ b/src/infinicore/ops/linear_gguf/linear_gguf.cc @@ -0,0 +1,47 @@ +#include "infinicore/ops/linear_gguf.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(LinearGguf); + +LinearGguf::LinearGguf(Tensor output, + const Tensor &input, + const Tensor &weight, + int64_t ggml_type) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, weight); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().getType(), output, input, weight, ggml_type); +} + +void LinearGguf::execute(Tensor output, + const Tensor &input, + const Tensor &weight, + int64_t ggml_type) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + LinearGguf, output, input, weight, ggml_type); +} + +Tensor linear_gguf(const Tensor &input, + const Tensor &weight, + int64_t ggml_type) { + // The decode kernel takes a single [M, K] batch; folding extra leading dims + // in here is left to the prefill path, so refuse them instead of pretending. + INFINICORE_ASSERT(input->ndim() == 2); + INFINICORE_ASSERT(weight->ndim() == 2); + auto output_shape = input->shape(); + output_shape.back() = weight->size(0); + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + linear_gguf_(output, input, weight, ggml_type); + return output; +} + +void linear_gguf_(Tensor output, + const Tensor &input, + const Tensor &weight, + int64_t ggml_type) { + LinearGguf::execute(output, input, weight, ggml_type); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/linear_gguf/linear_gguf_infiniop.cc b/src/infinicore/ops/linear_gguf/linear_gguf_infiniop.cc new file mode 100644 index 000000000..40eb1d57c --- /dev/null +++ b/src/infinicore/ops/linear_gguf/linear_gguf_infiniop.cc @@ -0,0 +1,68 @@ +#include "infinicore/ops/linear_gguf.hpp" + +// Included directly rather than through , so that adding this op +// needs no edit to the umbrella header that lists every other op. +#include "infiniop/ops/linear_gguf.h" + +#include "../infiniop_impl.hpp" + +namespace infinicore::op::linear_gguf_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, LinearGguf, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace; + graph::GraphTensor output; + graph::GraphTensor input; + graph::GraphTensor weight; + int64_t ggml_type; +}; + +void *plan(Tensor output, + const Tensor &input, + const Tensor &weight, + int64_t ggml_type) { + // The type id is part of the key: the same shapes with different block + // formats are different kernels, so they must not share a descriptor. + size_t seed = hash_combine(output, input, weight, static_cast(ggml_type)); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, + descriptor, + LinearGguf, + seed, + output->desc(), + input->desc(), + weight->desc(), + ggml_type); + INFINIOP_WORKSPACE_TENSOR(workspace, LinearGguf, descriptor); + + return new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(output), + graph::GraphTensor(input), + graph::GraphTensor(weight), + ggml_type}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopLinearGguf( + planned->descriptor->desc, + planned->workspace->data(), + planned->workspace->numel(), + planned->output->data(), + planned->input->data(), + planned->weight->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(LinearGguf, &plan, &run, &cleanup); + +} // namespace infinicore::op::linear_gguf_impl::infiniop diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index e78adf275..0c8aa8edc 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -78,6 +78,7 @@ #include "ops/lerp.hpp" #include "ops/linear.hpp" #include "ops/linear_mxfp4.hpp" +#include "ops/block_fp8_linear.hpp" #include "ops/linear_w8a8i8.hpp" #include "ops/log_softmax.hpp" #include "ops/logaddexp.hpp" @@ -249,6 +250,7 @@ inline void bind(py::module &m) { bind_logaddexp2(m); bind_linear(m); bind_linear_mxfp4(m); + bind_block_fp8_linear(m); bind_logdet(m); bind_matmul(m); bind_mamba_selective_scan(m); diff --git a/src/infinicore/pybind11/ops/block_fp8_linear.hpp b/src/infinicore/pybind11/ops/block_fp8_linear.hpp new file mode 100644 index 000000000..37514f293 --- /dev/null +++ b/src/infinicore/pybind11/ops/block_fp8_linear.hpp @@ -0,0 +1,38 @@ +#pragma once + +#include "infinicore/ops/block_fp8_linear.hpp" + +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline Tensor py_block_fp8_linear(Tensor input, + Tensor weight, + Tensor weight_scale) { + return op::block_fp8_linear(input, weight, weight_scale); +} + +inline void py_block_fp8_linear_(Tensor output, + Tensor input, + Tensor weight, + Tensor weight_scale) { + op::block_fp8_linear_(output, input, weight, weight_scale); +} + +inline void bind_block_fp8_linear(py::module &m) { + m.def("block_fp8_linear", + &ops::py_block_fp8_linear, + py::arg("input"), + py::arg("weight"), + py::arg("weight_scale")); + m.def("block_fp8_linear_", + &ops::py_block_fp8_linear_, + py::arg("output"), + py::arg("input"), + py::arg("weight"), + py::arg("weight_scale")); +} + +} // namespace infinicore::ops diff --git a/src/infiniop/ops/add_rms_norm/info.h b/src/infiniop/ops/add_rms_norm/info.h index 883aed343..63b4b2692 100644 --- a/src/infiniop/ops/add_rms_norm/info.h +++ b/src/infiniop/ops/add_rms_norm/info.h @@ -12,7 +12,9 @@ class AddRMSNormInfo { public: infiniDtype_t wtype; + infiniDtype_t ytype; infiniDtype_t atype; + infiniDtype_t btype; float epsilon; std::vector shape; std::vector y_strides; @@ -32,22 +34,36 @@ class AddRMSNormInfo { infiniopTensorDescriptor_t weight_desc, float epsilon) { - auto atype = y_desc->dtype(); + auto ytype = y_desc->dtype(); + auto atype = a_desc->dtype(); + auto btype = b_desc->dtype(); auto wtype = weight_desc->dtype(); - // Check that all input tensors have the same dtype - if (a_desc->dtype() != atype || b_desc->dtype() != atype) { + const bool mixed_f32_bf16 = + ytype == INFINI_DTYPE_BF16 + && atype == INFINI_DTYPE_F32 + && btype == INFINI_DTYPE_BF16 + && wtype == INFINI_DTYPE_BF16; + const bool mixed_bf16_to_f32 = + ytype == INFINI_DTYPE_F32 + && atype == INFINI_DTYPE_BF16 + && btype == INFINI_DTYPE_BF16 + && wtype == INFINI_DTYPE_BF16; + + if (mixed_f32_bf16 || mixed_bf16_to_f32) { + // Experimental precision boundary used by GGUF decode: preserve the + // F32 linear output through add/RMS reduction, or preserve the final + // BF16 residual sum and norm output in F32 before the LM head. + } else if (atype != ytype || btype != ytype) { return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (atype == INFINI_DTYPE_F16 || atype == INFINI_DTYPE_BF16) { + } else if (ytype == INFINI_DTYPE_F16 || ytype == INFINI_DTYPE_BF16) { // For half-precision types (FP16/BF16), weights can be the same half-precision type or FP32 - if (wtype != atype && wtype != INFINI_DTYPE_F32 && wtype != INFINI_DTYPE_BF16 && wtype != INFINI_DTYPE_F16) { + if (wtype != ytype && wtype != INFINI_DTYPE_F32 && wtype != INFINI_DTYPE_BF16 && wtype != INFINI_DTYPE_F16) { return INFINI_STATUS_BAD_TENSOR_DTYPE; } - } else if (atype == INFINI_DTYPE_F32 || atype == INFINI_DTYPE_F64) { + } else if (ytype == INFINI_DTYPE_F32 || ytype == INFINI_DTYPE_F64) { // For FP32/FP64, activations and weights must be of the same type - if (atype != wtype) { + if (ytype != wtype) { return INFINI_STATUS_BAD_TENSOR_DTYPE; } } else { @@ -100,7 +116,7 @@ class AddRMSNormInfo { if (residual_out_ndim != y_ndim) { return INFINI_STATUS_BAD_TENSOR_SHAPE; } - if (residual_out_desc->dtype() != atype) { + if (residual_out_desc->dtype() != ytype) { return INFINI_STATUS_BAD_TENSOR_DTYPE; } // Check shape matches @@ -115,7 +131,9 @@ class AddRMSNormInfo { AddRMSNormInfo info; info.wtype = wtype; + info.ytype = ytype; info.atype = atype; + info.btype = btype; info.epsilon = epsilon; info.shape = y_desc->shape(); info.y_strides = y_desc->strides(); diff --git a/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cu b/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cu index 652f8adf3..76d17a8f8 100644 --- a/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cu +++ b/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cu @@ -35,6 +35,57 @@ INFINIOP_CUDA_KERNEL add_rmsnormKernel( w, nhead, dim, epsilon); } +template +INFINIOP_CUDA_KERNEL add_rmsnormMixedKernel( + Ty *__restrict__ y, + Ty *__restrict__ residual_out, + ptrdiff_t stride_y_batch, + ptrdiff_t stride_y_nhead, + ptrdiff_t stride_residual_out_batch, + ptrdiff_t stride_residual_out_nhead, + const Ta *__restrict__ a, + ptrdiff_t stride_a_batch, + ptrdiff_t stride_a_nhead, + const Tb *__restrict__ b, + ptrdiff_t stride_b_batch, + ptrdiff_t stride_b_nhead, + const Tweight *__restrict__ w, + size_t nhead, + size_t dim, + float epsilon) { + const size_t batch_idx = blockIdx.x / nhead; + const size_t head_idx = blockIdx.x % nhead; + auto y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead; + auto residual_ptr = residual_out + + batch_idx * stride_residual_out_batch + head_idx * stride_residual_out_nhead; + auto a_ptr = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead; + auto b_ptr = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead; + + Tcompute sum_squared = 0; + for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) { + const Tcompute sum_val = Tcompute(a_ptr[i]) + Tcompute(b_ptr[i]); + residual_ptr[i] = Ty(sum_val); + sum_squared += sum_val * sum_val; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + sum_squared = BlockReduce(temp_storage).Sum(sum_squared); + __shared__ Tcompute rms; + if (threadIdx.x == 0) { + rms = Tcompute(rsqrtf(sum_squared / Tcompute(dim) + epsilon)); + } + __syncthreads(); + + // Recompute the F32 sum instead of reading the BF16 residual_out. This + // keeps the GGUF linear accumulator precision through normalization while + // restoring the model's ordinary BF16 boundary for subsequent layers. + for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) { + const Tcompute sum_val = Tcompute(a_ptr[i]) + Tcompute(b_ptr[i]); + y_ptr[i] = Ty(sum_val * Tcompute(w[i]) * rms); + } +} + namespace op::add_rms_norm::nvidia { struct Descriptor::Opaque { @@ -70,10 +121,10 @@ infiniStatus_t Descriptor::create( template infiniStatus_t launchKernel( uint32_t batch_size, size_t nhead, size_t dim, - void *y, infiniDtype_t atype, ptrdiff_t stride_y_batch, ptrdiff_t stride_y_nhead, + void *y, infiniDtype_t ytype, ptrdiff_t stride_y_batch, ptrdiff_t stride_y_nhead, void *residual_out, ptrdiff_t stride_residual_out_batch, ptrdiff_t stride_residual_out_nhead, - const void *a, ptrdiff_t stride_a_batch, ptrdiff_t stride_a_nhead, - const void *b, ptrdiff_t stride_b_batch, ptrdiff_t stride_b_nhead, + const void *a, infiniDtype_t atype, ptrdiff_t stride_a_batch, ptrdiff_t stride_a_nhead, + const void *b, infiniDtype_t btype, ptrdiff_t stride_b_batch, ptrdiff_t stride_b_nhead, const void *w, infiniDtype_t wtype, float epsilon, cudaStream_t cuda_stream) { @@ -97,19 +148,41 @@ infiniStatus_t launchKernel( dim, \ epsilon) - if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) { + if (ytype == INFINI_DTYPE_BF16 && atype == INFINI_DTYPE_F32 + && btype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { + add_rmsnormMixedKernel + <<>>( + reinterpret_cast<__nv_bfloat16 *>(y), + reinterpret_cast<__nv_bfloat16 *>(residual_out), + stride_y_batch, stride_y_nhead, + stride_residual_out_batch, stride_residual_out_nhead, + reinterpret_cast(a), stride_a_batch, stride_a_nhead, + reinterpret_cast(b), stride_b_batch, stride_b_nhead, + reinterpret_cast(w), nhead, dim, epsilon); + } else if (ytype == INFINI_DTYPE_F32 && atype == INFINI_DTYPE_BF16 + && btype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { + add_rmsnormMixedKernel + <<>>( + reinterpret_cast(y), + reinterpret_cast(residual_out), + stride_y_batch, stride_y_nhead, + stride_residual_out_batch, stride_residual_out_nhead, + reinterpret_cast(a), stride_a_batch, stride_a_nhead, + reinterpret_cast(b), stride_b_batch, stride_b_nhead, + reinterpret_cast(w), nhead, dim, epsilon); + } else if (ytype == INFINI_DTYPE_F16 && atype == ytype && btype == ytype && wtype == INFINI_DTYPE_F16) { LAUNCH_KERNEL(half, half, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) { + } else if (ytype == INFINI_DTYPE_F16 && atype == ytype && btype == ytype && wtype == INFINI_DTYPE_BF16) { LAUNCH_KERNEL(half, __nv_bfloat16, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) { + } else if (ytype == INFINI_DTYPE_F16 && atype == ytype && btype == ytype && wtype == INFINI_DTYPE_F32) { LAUNCH_KERNEL(half, float, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { + } else if (ytype == INFINI_DTYPE_BF16 && atype == ytype && btype == ytype && wtype == INFINI_DTYPE_BF16) { LAUNCH_KERNEL(__nv_bfloat16, __nv_bfloat16, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) { + } else if (ytype == INFINI_DTYPE_BF16 && atype == ytype && btype == ytype && wtype == INFINI_DTYPE_F16) { LAUNCH_KERNEL(__nv_bfloat16, half, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) { + } else if (ytype == INFINI_DTYPE_BF16 && atype == ytype && btype == ytype && wtype == INFINI_DTYPE_F32) { LAUNCH_KERNEL(__nv_bfloat16, float, float); - } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) { + } else if (ytype == INFINI_DTYPE_F32 && atype == ytype && btype == ytype && wtype == INFINI_DTYPE_F32) { LAUNCH_KERNEL(float, float, float); } else { return INFINI_STATUS_BAD_TENSOR_DTYPE; @@ -146,34 +219,34 @@ infiniStatus_t Descriptor::calculate( if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { CHECK_STATUS(launchKernel( batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, + y, _info.ytype, stride_y_batch, stride_y_nhead, residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, + a, _info.atype, stride_a_batch, stride_a_nhead, + b, _info.btype, stride_b_batch, stride_b_nhead, weight, _info.wtype, _info.epsilon, cuda_stream)); } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { CHECK_STATUS(launchKernel( batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, + y, _info.ytype, stride_y_batch, stride_y_nhead, residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, + a, _info.atype, stride_a_batch, stride_a_nhead, + b, _info.btype, stride_b_batch, stride_b_nhead, weight, _info.wtype, _info.epsilon, cuda_stream)); } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) { CHECK_STATUS(launchKernel( batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, + y, _info.ytype, stride_y_batch, stride_y_nhead, residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, + a, _info.atype, stride_a_batch, stride_a_nhead, + b, _info.btype, stride_b_batch, stride_b_nhead, weight, _info.wtype, _info.epsilon, cuda_stream)); } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { CHECK_STATUS(launchKernel( batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, + y, _info.ytype, stride_y_batch, stride_y_nhead, residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, + a, _info.atype, stride_a_batch, stride_a_nhead, + b, _info.btype, stride_b_batch, stride_b_nhead, weight, _info.wtype, _info.epsilon, cuda_stream)); } else { return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/block_fp8_linear/block_fp8_linear.h b/src/infiniop/ops/block_fp8_linear/block_fp8_linear.h new file mode 100644 index 000000000..9b03692da --- /dev/null +++ b/src/infiniop/ops/block_fp8_linear/block_fp8_linear.h @@ -0,0 +1,38 @@ +#ifndef __BLOCK_FP8_LINEAR_H__ +#define __BLOCK_FP8_LINEAR_H__ + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::block_fp8_linear::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + BlockFP8LinearInfo _info; \ + \ + Descriptor(Opaque *opaque, BlockFP8LinearInfo info, \ + infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), _info(info) {} \ + \ + public: \ + ~Descriptor(); \ + size_t workspaceSize() const; \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t output_desc, \ + infiniopTensorDescriptor_t input_desc, \ + infiniopTensorDescriptor_t weight_desc, \ + infiniopTensorDescriptor_t weight_scale_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, \ + void *output, const void *input, \ + const void *weight, const void *weight_scale, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/block_fp8_linear/info.h b/src/infiniop/ops/block_fp8_linear/info.h new file mode 100644 index 000000000..f9e552cc2 --- /dev/null +++ b/src/infiniop/ops/block_fp8_linear/info.h @@ -0,0 +1,78 @@ +#ifndef __BLOCK_FP8_LINEAR_INFO_H__ +#define __BLOCK_FP8_LINEAR_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::block_fp8_linear { + +static constexpr size_t BLOCK_SIZE = 128; + +class BlockFP8LinearInfo { + BlockFP8LinearInfo() = default; + +public: + infiniDtype_t dtype; // activation/output dtype (BF16) + size_t M; + size_t N; + size_t K; + size_t num_out_blocks; // ceil(N/128) + size_t num_in_blocks; // ceil(K/128) + + static utils::Result create( + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t weight_scale_desc) { + CHECK_OR_RETURN(output_desc != nullptr && input_desc != nullptr + && weight_desc != nullptr && weight_scale_desc != nullptr, + INFINI_STATUS_NULL_POINTER); + + const auto dtype = input_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_BF16); + CHECK_OR_RETURN(output_desc->dtype() == dtype, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_DTYPE(weight_desc->dtype(), INFINI_DTYPE_F8); + CHECK_DTYPE(weight_scale_desc->dtype(), INFINI_DTYPE_F32); + + CHECK_OR_RETURN(input_desc->ndim() == 2 + && output_desc->ndim() == 2 + && weight_desc->ndim() == 2 + && weight_scale_desc->ndim() == 2, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(input_desc->isContiguous() && output_desc->isContiguous() + && weight_desc->isContiguous() + && weight_scale_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + + const size_t M = input_desc->dim(0); + const size_t K = input_desc->dim(1); + const size_t N = weight_desc->dim(0); + const size_t K_w = weight_desc->dim(1); + + CHECK_OR_RETURN(K == K_w && K > 0 && N > 0 && M > 0, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(output_desc->dim(0) == M && output_desc->dim(1) == N, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + const size_t num_out_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE; + const size_t num_in_blocks = (K + BLOCK_SIZE - 1) / BLOCK_SIZE; + + CHECK_OR_RETURN(weight_scale_desc->dim(0) == num_out_blocks + && weight_scale_desc->dim(1) == num_in_blocks, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + BlockFP8LinearInfo info; + info.dtype = dtype; + info.M = M; + info.N = N; + info.K = K; + info.num_out_blocks = num_out_blocks; + info.num_in_blocks = num_in_blocks; + return utils::Result(info); + } +}; + +} // namespace op::block_fp8_linear + +#endif diff --git a/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh new file mode 100644 index 000000000..75938ad5c --- /dev/null +++ b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh @@ -0,0 +1,217 @@ +#ifndef __BLOCK_FP8_GEMM_SM120_CUH__ +#define __BLOCK_FP8_GEMM_SM120_CUH__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace op::block_fp8_linear::nvidia::sm120 { + +template +struct enable_sm120_family : Kernel { + template + CUTLASS_DEVICE void operator()(Args&&... args) { +#if defined(__CUDA_ARCH__) + #if (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300) + Kernel::operator()(std::forward(args)...); + #else + printf("BlockFP8Linear: kernel only supports sm120 family.\n"); + asm("trap;"); + #endif +#endif + } +}; + +struct GemmConfig { + using ElementAB = cutlass::float_e4m3_t; + using ElementD = cutlass::bfloat16_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBlockScale = float; + using ElementC = void; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutD = cutlass::layout::RowMajor; + + static constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< + 1, 128, 128, + cute::UMMA::Major::MN, + cute::UMMA::Major::K>; + + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); + + using ArchTag = cutlass::arch::Sm120; + using OperatorClass = cutlass::arch::OpClassTensorOp; + + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; + + using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwiseCooperativeSm120; + using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; + + static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest; + + using DefaultOperation = cutlass::epilogue::fusion::LinearCombination< + ElementD, ElementCompute, ElementC, ElementBlockScale, RoundStyle>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, ElementC, + LayoutD, AlignmentD, + ElementD, LayoutD, AlignmentD, + EpilogueSchedule, DefaultOperation>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementAB, + cute::tuple, + AlignmentA, + ElementAB, + cute::tuple, + AlignmentB, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using KernelType = enable_sm120_family, + CollectiveMainloop, + CollectiveEpilogue>>; + + using GemmOp = cutlass::gemm::device::GemmUniversalAdapter; +}; + +inline size_t get_gemm_workspace_size(size_t M, size_t N, size_t K) { + using GC = GemmConfig; + using GemmOp = GC::GemmOp; + using GemmKernel = typename GemmOp::GemmKernel; + + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideC = typename GemmKernel::StrideC; + + auto a_stride = cutlass::make_cute_packed_stride( + StrideA{}, cute::make_shape((int)M, (int)K, 1)); + auto b_stride = cutlass::make_cute_packed_stride( + StrideB{}, cute::make_shape((int)N, (int)K, 1)); + auto c_stride = cutlass::make_cute_packed_stride( + StrideC{}, cute::make_shape((int)M, (int)N, 1)); + + auto layout_SFA = GC::ScaleConfig::tile_atom_to_shape_SFA( + cute::make_shape((int)M, (int)N, (int)K, 1)); + auto layout_SFB = GC::ScaleConfig::tile_atom_to_shape_SFB( + cute::make_shape((int)M, (int)N, (int)K, 1)); + + typename GemmKernel::MainloopArguments mainloop_args{}; + mainloop_args.ptr_A = nullptr; + mainloop_args.dA = a_stride; + mainloop_args.ptr_B = nullptr; + mainloop_args.dB = b_stride; + mainloop_args.ptr_SFA = nullptr; + mainloop_args.ptr_SFB = nullptr; + mainloop_args.layout_SFA = layout_SFA; + mainloop_args.layout_SFB = layout_SFB; + + auto prob_shape = cute::make_shape((int)M, (int)N, (int)K, 1); + + typename GemmKernel::EpilogueArguments epilogue_args{ + {}, nullptr, c_stride, nullptr, c_stride}; + + cutlass::KernelHardwareInfo hw_info; + typename GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, + prob_shape, mainloop_args, epilogue_args, hw_info, {}}; + + GemmOp gemm_op; + size_t ws_size = gemm_op.get_workspace_size(args); + return ws_size; +} + +inline cutlass::Status run_gemm( + void *output, + const void *a_fp8, + const void *a_scales, + const void *weight, + const void *weight_scales, + size_t M, size_t N, size_t K, + void *cutlass_workspace, + cudaStream_t stream) { + + using GC = GemmConfig; + using GemmOp = GC::GemmOp; + using GemmKernel = typename GemmOp::GemmKernel; + + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideC = typename GemmKernel::StrideC; + + auto a_stride = cutlass::make_cute_packed_stride( + StrideA{}, cute::make_shape((int)M, (int)K, 1)); + auto b_stride = cutlass::make_cute_packed_stride( + StrideB{}, cute::make_shape((int)N, (int)K, 1)); + auto c_stride = cutlass::make_cute_packed_stride( + StrideC{}, cute::make_shape((int)M, (int)N, 1)); + + auto layout_SFA = GC::ScaleConfig::tile_atom_to_shape_SFA( + cute::make_shape((int)M, (int)N, (int)K, 1)); + auto layout_SFB = GC::ScaleConfig::tile_atom_to_shape_SFB( + cute::make_shape((int)M, (int)N, (int)K, 1)); + + auto *a_ptr = static_cast(a_fp8); + auto *b_ptr = static_cast(weight); + auto *sfa_ptr = static_cast(a_scales); + auto *sfb_ptr = static_cast(weight_scales); + auto *d_ptr = static_cast(output); + + typename GemmKernel::MainloopArguments mainloop_args{}; + mainloop_args.ptr_A = a_ptr; + mainloop_args.dA = a_stride; + mainloop_args.ptr_B = b_ptr; + mainloop_args.dB = b_stride; + mainloop_args.ptr_SFA = sfa_ptr; + mainloop_args.ptr_SFB = sfb_ptr; + mainloop_args.layout_SFA = layout_SFA; + mainloop_args.layout_SFB = layout_SFB; + + auto prob_shape = cute::make_shape((int)M, (int)N, (int)K, 1); + + typename GemmKernel::EpilogueArguments epilogue_args{ + {}, d_ptr, c_stride, d_ptr, c_stride}; + + cutlass::KernelHardwareInfo hw_info; + typename GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, + prob_shape, mainloop_args, epilogue_args, hw_info, {}}; + + GemmOp gemm_op; + + auto can_impl = gemm_op.can_implement(args); + if (can_impl != cutlass::Status::kSuccess) { + return can_impl; + } + + size_t ws_size = gemm_op.get_workspace_size(args); + return gemm_op.run(args, cutlass_workspace, stream); +} + +} // namespace op::block_fp8_linear::nvidia::sm120 + +#endif diff --git a/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu new file mode 100644 index 000000000..ac4941e3b --- /dev/null +++ b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu @@ -0,0 +1,189 @@ +#include "block_fp8_linear_nvidia.cuh" + +#include "../../../devices/nvidia/nvidia_handle.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +#ifdef ENABLE_CUTLASS_API +#include "block_fp8_gemm_sm120.cuh" +#endif + +#include +#include +#include +#include + +namespace op::block_fp8_linear::nvidia { + +namespace { + +__global__ void per_token_group_quant_kernel( + const __nv_bfloat16 *__restrict__ input, + __nv_fp8_e4m3 *__restrict__ output, + float *__restrict__ scales, + size_t M, + size_t K, + size_t num_groups) { + const size_t row = blockIdx.x; + if (row >= M) return; + + const size_t tid = threadIdx.x; + const size_t threads_per_group = 8; + const size_t group_id = tid / threads_per_group; + const size_t lane = tid % threads_per_group; + const size_t groups_per_block = blockDim.x / threads_per_group; + + for (size_t g = group_id; g < num_groups; g += groups_per_block) { + size_t start = g * 128; + size_t end = min(start + 128, K); + + float local_max = 0.0f; + for (size_t i = lane; i < 128 && (start + i) < K; i += threads_per_group) { + float val = static_cast(input[row * K + start + i]); + local_max = fmaxf(local_max, fabsf(val)); + } + + for (int offset = 4; offset > 0; offset /= 2) { + local_max = fmaxf(local_max, __shfl_xor_sync(0xff, local_max, offset)); + } + + float amax = local_max; + float scale = fmaxf(amax / 448.0f, 1e-10f); + if (lane == 0) { + scales[row * num_groups + g] = scale; + } + + float inv_scale = 1.0f / scale; + for (size_t i = lane; i < 128 && (start + i) < K; i += threads_per_group) { + float val = static_cast(input[row * K + start + i]); + float q = val * inv_scale; + q = fminf(fmaxf(q, -448.0f), 448.0f); + output[row * K + start + i] = __nv_fp8_e4m3(q); + } + } +} + +void launch_per_token_group_quant( + const __nv_bfloat16 *input, + __nv_fp8_e4m3 *output, + float *scales, + size_t M, + size_t K, + size_t num_groups, + cudaStream_t stream) { + int threads = 128; + int blocks = static_cast(M); + per_token_group_quant_kernel<<>>( + input, output, scales, M, K, num_groups); +} + +} // namespace + +struct Descriptor::Opaque { + std::shared_ptr internal; + size_t workspace_size; +}; + +Descriptor::~Descriptor() { delete _opaque; } + +size_t Descriptor::workspaceSize() const { + return _opaque->workspace_size; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t weight_scale_desc) { + auto info_result = BlockFP8LinearInfo::create( + output_desc, input_desc, weight_desc, weight_scale_desc); + CHECK_RESULT(info_result); + + auto info = info_result.take(); + + auto nvidia_handle = reinterpret_cast(handle); + + const size_t M = info.M; + const size_t K = info.K; + const size_t N = info.N; + const size_t num_in_blocks = info.num_in_blocks; + + auto align256 = [](size_t n) { return (n + 255) & ~255; }; + + size_t a_fp8_size = align256(M * K); + size_t a_scale_size = align256(M * num_in_blocks * 4); + +#ifdef ENABLE_CUTLASS_API + size_t cutlass_ws = sm120::get_gemm_workspace_size(M, N, K); +#else + size_t cutlass_ws = 0; +#endif + cutlass_ws = align256(cutlass_ws); + + size_t total_ws = a_fp8_size + a_scale_size + cutlass_ws; + + *desc_ptr = new Descriptor( + new Opaque{nvidia_handle->internal(), total_ws}, + info, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *output, + const void *input, + const void *weight, + const void *weight_scale, + void *stream) const { + auto cuda_stream = reinterpret_cast(stream); + + const size_t M = _info.M; + const size_t N = _info.N; + const size_t K = _info.K; + const size_t num_in_blocks = _info.num_in_blocks; + + auto align256 = [](size_t n) { return (n + 255) & ~255; }; + + size_t a_fp8_offset = 0; + size_t a_fp8_size = align256(M * K); + size_t a_scale_offset = a_fp8_offset + a_fp8_size; + size_t a_scale_size = align256(M * num_in_blocks * 4); + size_t cutlass_offset = a_scale_offset + a_scale_size; + + auto *a_fp8 = static_cast<__nv_fp8_e4m3 *>(static_cast( + static_cast(workspace) + a_fp8_offset)); + auto *a_scales = static_cast(static_cast( + static_cast(workspace) + a_scale_offset)); + auto *cutlass_ws = static_cast( + static_cast(workspace) + cutlass_offset); + + auto *a_bf16 = reinterpret_cast(input); + auto *w_fp8 = reinterpret_cast(weight); + auto *w_scales = reinterpret_cast(weight_scale); + auto *out_bf16 = reinterpret_cast<__nv_bfloat16 *>(output); + + launch_per_token_group_quant( + a_bf16, a_fp8, a_scales, M, K, num_in_blocks, cuda_stream); + +#ifdef ENABLE_CUTLASS_API + auto status = sm120::run_gemm( + out_bf16, a_fp8, a_scales, w_fp8, w_scales, + M, N, K, cutlass_ws, cuda_stream); + + if (status != cutlass::Status::kSuccess) { + return INFINI_STATUS_INTERNAL_ERROR; + } +#else + return INFINI_STATUS_NOT_IMPLEMENTED; +#endif + + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + return INFINI_STATUS_INTERNAL_ERROR; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::block_fp8_linear::nvidia diff --git a/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cuh b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cuh new file mode 100644 index 000000000..a2a86981f --- /dev/null +++ b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __BLOCK_FP8_LINEAR_NVIDIA_CUH__ +#define __BLOCK_FP8_LINEAR_NVIDIA_CUH__ + +#include "../block_fp8_linear.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/block_fp8_linear/operator.cc b/src/infiniop/ops/block_fp8_linear/operator.cc new file mode 100644 index 000000000..a7a54385c --- /dev/null +++ b/src/infiniop/ops/block_fp8_linear/operator.cc @@ -0,0 +1,98 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/block_fp8_linear.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) +#include "nvidia/block_fp8_linear_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateBlockFP8LinearDescriptor( + infiniopHandle_t handle, + infiniopBlockFP8LinearDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t weight_scale_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::block_fp8_linear::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + output_desc, input_desc, weight_desc, weight_scale_desc) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetBlockFP8LinearWorkspaceSize( + infiniopBlockFP8LinearDescriptor_t desc, + size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc) \ + ->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopBlockFP8Linear( + infiniopBlockFP8LinearDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *weight, + const void *weight_scale, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, output, input, weight, weight_scale, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyBlockFP8LinearDescriptor( + infiniopBlockFP8LinearDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu b/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu index 26bff85fd..0cbd9b510 100644 --- a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu +++ b/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu @@ -3,11 +3,81 @@ #if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) #include #endif +#include namespace op::gemm::nvidia { +namespace { + +constexpr int kMixedSmallNMax = 16; +constexpr int kMixedWarpsPerBlock = 8; + +// C[m, n] = A_bf16[m, k] @ B_f32[k, n]. This is the transposed +// representation used when a row-major LM head is evaluated as W @ hidden^T. +// One warp owns a weight row and reuses each BF16 value across a tile of up to +// 16 hidden columns. grid.y tiles longer prompts without growing registers. +__global__ void mixed_bf16_f32_small_n_kernel( + const __nv_bfloat16 *__restrict__ a, + const float *__restrict__ b, + float *__restrict__ c, + size_t m, size_t n, size_t k, + ptrdiff_t a_row_stride, ptrdiff_t a_col_stride, + ptrdiff_t b_row_stride, ptrdiff_t b_col_stride, + ptrdiff_t c_row_stride, ptrdiff_t c_col_stride, + float alpha, float beta) { + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const size_t row = static_cast(blockIdx.x) * kMixedWarpsPerBlock + + static_cast(warp); + const size_t col_base = static_cast(blockIdx.y) * kMixedSmallNMax; + if (row >= m || col_base >= n) { + return; + } + const size_t remaining = n - col_base; + const size_t tile_n = remaining < static_cast(kMixedSmallNMax) + ? remaining : static_cast(kMixedSmallNMax); + + float acc[kMixedSmallNMax]; +#pragma unroll + for (int j = 0; j < kMixedSmallNMax; ++j) { + acc[j] = 0.0f; + } + for (size_t kk = static_cast(lane); kk < k; kk += 32) { + const float av = __bfloat162float(a[row * a_row_stride + kk * a_col_stride]); +#pragma unroll + for (int j = 0; j < kMixedSmallNMax; ++j) { + if (static_cast(j) >= tile_n) { + break; + } + acc[j] += av * b[kk * b_row_stride + + (col_base + static_cast(j)) * b_col_stride]; + } + } + constexpr unsigned mask = 0xffffffffu; +#pragma unroll + for (int j = 0; j < kMixedSmallNMax; ++j) { + if (static_cast(j) >= tile_n) { + break; + } + float value = acc[j]; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(mask, value, offset); + } + if (lane == 0) { + const size_t idx = row * c_row_stride + + (col_base + static_cast(j)) * c_col_stride; + c[idx] = alpha * value + (beta == 0.0f ? 0.0f : beta * c[idx]); + } + } +} + +} // namespace + struct Descriptor::Opaque { std::shared_ptr internal; + infiniDtype_t a_dtype; + infiniDtype_t b_dtype; #if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) cublasLtHandle_t lt_handle = nullptr; cublasLtMatmulDesc_t lt_desc = nullptr; @@ -140,7 +210,11 @@ infiniStatus_t Descriptor::create( CHECK_RESULT(result); auto info = result.take(); - auto opaque = new Opaque{handle->internal()}; + auto a_dtype = a_desc->dtype(); + auto b_dtype = b_desc->dtype(); + CHECK_DTYPE(a_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + CHECK_DTYPE(b_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + auto opaque = new Opaque{handle->internal(), a_dtype, b_dtype}; #if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) if (dtype == INFINI_DTYPE_BF16 && !opaque->createBf16LtDescriptors(info)) { opaque->destroyLtDescriptors(); @@ -171,9 +245,24 @@ infiniStatus_t Descriptor::calculate( cublasComputeType_t compute_type; #endif + auto cuda_dtype = [](infiniDtype_t dtype) { + switch (dtype) { + case INFINI_DTYPE_F16: + return CUDA_R_16F; + case INFINI_DTYPE_BF16: + return CUDA_R_16BF; + case INFINI_DTYPE_F32: + return CUDA_R_32F; + default: + return CUDA_R_32F; + } + }; + a_type = cuda_dtype(_opaque->a_dtype); + b_type = cuda_dtype(_opaque->b_dtype); + c_type = cuda_dtype(_dtype); + switch (_dtype) { case INFINI_DTYPE_F16: - a_type = b_type = c_type = CUDA_R_16F; #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) compute_type = CUDA_R_32F; #else @@ -181,7 +270,6 @@ infiniStatus_t Descriptor::calculate( #endif break; case INFINI_DTYPE_BF16: - a_type = b_type = c_type = CUDA_R_16BF; #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) compute_type = CUDA_R_32F; #else @@ -189,11 +277,14 @@ infiniStatus_t Descriptor::calculate( #endif break; case INFINI_DTYPE_F32: - a_type = b_type = c_type = CUDA_R_32F; #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) compute_type = CUDA_R_32F; #else - compute_type = CUBLAS_COMPUTE_32F_FAST_TF32; + compute_type = + _opaque->a_dtype == INFINI_DTYPE_F32 + && _opaque->b_dtype == INFINI_DTYPE_F32 + ? CUBLAS_COMPUTE_32F_FAST_TF32 + : CUBLAS_COMPUTE_32F; #endif break; @@ -203,11 +294,39 @@ infiniStatus_t Descriptor::calculate( if (_info.is_transed) { std::swap(a, b); + // Row-major output is evaluated as the transposed product B^T @ A^T. + // Keep the runtime CUDA element types attached to the pointers as well; + // this was invisible for equal-dtype GEMMs but breaks mixed F32/BF16. + std::swap(a_type, b_type); } auto op_a = _info.a_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; auto op_b = _info.b_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; + if (_dtype == INFINI_DTYPE_F32 + && a_type == CUDA_R_16BF && b_type == CUDA_R_32F + && _info.batch == 1) { + const dim3 blocks( + static_cast( + (_info.m + kMixedWarpsPerBlock - 1) / kMixedWarpsPerBlock), + static_cast( + (_info.n + kMixedSmallNMax - 1) / kMixedSmallNMax)); + mixed_bf16_f32_small_n_kernel<<< + blocks, kMixedWarpsPerBlock * 32, 0, + reinterpret_cast(stream)>>>( + reinterpret_cast(a), + reinterpret_cast(b), + reinterpret_cast(c), + _info.m, _info.n, _info.k, + _info.a_matrix.row_stride, _info.a_matrix.col_stride, + _info.b_matrix.row_stride, _info.b_matrix.col_stride, + _info.c_matrix.row_stride, _info.c_matrix.col_stride, + alpha, beta); + return cudaGetLastError() == cudaSuccess + ? INFINI_STATUS_SUCCESS + : INFINI_STATUS_INTERNAL_ERROR; + } + #if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) if (_dtype == INFINI_DTYPE_BF16 && _opaque->lt_handle && _opaque->lt_desc && _opaque->a_layout && _opaque->b_layout && _opaque->c_layout) { diff --git a/src/infiniop/ops/linear_gguf/ggml_blocks.h b/src/infiniop/ops/linear_gguf/ggml_blocks.h new file mode 100644 index 000000000..33218127f --- /dev/null +++ b/src/infiniop/ops/linear_gguf/ggml_blocks.h @@ -0,0 +1,277 @@ +// Device-independent GGUF block decoders for the linear_gguf op. +// +// Scope: the four GGML block types route B stores verbatim, i.e. Q8_0 (8), +// Q4_K (12), Q5_K (13), Q6_K (14). Everything else is rejected upstream, at +// packaging time, so no other decoder lives here. +// +// Numerics: scalar port of dequantize_row_q8_0/q4_K/q5_K/q6_K in llama.cpp +// ggml/src/ggml-quants.c. The floating-point operation order is preserved +// on purpose: the K-quant path forms (d * scale) and (dmin * min) first and +// then touches the quants, so a re-associated expression is not bit-exact and +// the fp32 -> bf16 round would land on a different neighbour. Callers that +// need bit-exact agreement with llama.cpp must go through these functions +// instead of re-deriving the formula. +// +// Alignment: fields are read byte by byte. Route B packs the block rows of one +// GGUF tensor back to back into a single uint8 tensor, so a block only starts +// at 2-byte alignment (Q6_K's half sits at offset 208). Reinterpreting the +// buffer as block_q6_K et al. would be undefined behaviour on some backends. +// +// Portability: no CUDA-only type or intrinsic appears here, and bf16 values +// travel as uint16_t bit patterns, so the same header serves .cu kernels, +// host-side tests and future non-NVIDIA backends. It is the single source of +// truth for block decoding; kernels that decode inline for bandwidth must be +// validated against these functions. +#ifndef __GGML_BLOCKS_H__ +#define __GGML_BLOCKS_H__ + +#include +#include + +#ifdef __CUDACC__ +// host+device so a .cu translation unit can also decode on the host side +// (test drivers, CPU fallback) without a second copy of the logic. +#define GGML_BLK_HOST_DEVICE __host__ __device__ __forceinline__ +#else +#define GGML_BLK_HOST_DEVICE inline +#endif + +namespace ggml_blocks { + +// enum ggml_type values, pinned to gguf.h / GGML_QUANT_SIZES. +enum GgmlType : int32_t { + GGML_TYPE_Q8_0 = 8, + GGML_TYPE_Q4_K = 12, + GGML_TYPE_Q5_K = 13, + GGML_TYPE_Q6_K = 14, +}; + +// Elements per super-block / per Q8_0 block, and bytes per block. +constexpr int32_t QK_K = 256; +constexpr int32_t QK8_0 = 32; +constexpr int32_t K_SCALE_SIZE = 12; +constexpr int32_t SIZE_Q8_0 = 34; // 2 + 32 +constexpr int32_t SIZE_Q4_K = 144; // 4 + 12 + 128 +constexpr int32_t SIZE_Q5_K = 176; // 4 + 12 + 32 + 128 +constexpr int32_t SIZE_Q6_K = 210; // 128 + 64 + 16 + 2 + +// Number of elements one block of `type` decodes to, or -1 if `type` has no +// decoder here. +GGML_BLK_HOST_DEVICE int32_t block_elems(int32_t type) { + return type == GGML_TYPE_Q8_0 ? QK8_0 + : (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || + type == GGML_TYPE_Q6_K + ? QK_K + : -1); +} + +// Bytes one block of `type` occupies in the blob, or -1 if unsupported. +GGML_BLK_HOST_DEVICE int32_t block_bytes(int32_t type) { + switch (type) { + case GGML_TYPE_Q8_0: return SIZE_Q8_0; + case GGML_TYPE_Q4_K: return SIZE_Q4_K; + case GGML_TYPE_Q5_K: return SIZE_Q5_K; + case GGML_TYPE_Q6_K: return SIZE_Q6_K; + default: return -1; + } +} + +GGML_BLK_HOST_DEVICE uint16_t read_u16(const uint8_t *p) { + return (uint16_t)((uint16_t)p[0] | ((uint16_t)p[1] << 8)); +} + +// IEEE binary16 -> binary32, exact, subnormal and NaN safe. Equivalent to +// ggml GGML_FP16_TO_FP32 / __half2float without depending on either. +GGML_BLK_HOST_DEVICE float half_to_float(uint16_t h) { + const uint32_t sign = (uint32_t)(h >> 15) << 31; + const uint32_t exp = (uint32_t)(h >> 10) & 0x1Fu; + const uint32_t mant = (uint32_t)h & 0x3FFu; + uint32_t bits; + if (exp == 0) { + if (mant == 0) { + bits = sign; // +-0 + } else { + // Normalise the subnormal: value == mant * 2^-24. + int e = -14; + uint32_t m = mant; + while ((m & 0x400u) == 0) { + m <<= 1; + --e; + } + bits = sign | ((uint32_t)(e + 127) << 23) | ((m & 0x3FFu) << 13); + } + } else if (exp == 31) { + bits = sign | 0x7F800000u | (mant << 13); // inf / nan + } else { + bits = sign | ((exp + 112) << 23) | (mant << 13); + } + float f; + static_assert(sizeof(f) == 4, "float is assumed to be IEEE binary32"); + std::memcpy(&f, &bits, 4); + return f; +} + +// binary32 -> bf16 bit pattern, round-to-nearest-even, matching +// __float2bfloat16 / torch.Tensor.to(bfloat16) behaviour on finite inputs. +GGML_BLK_HOST_DEVICE uint16_t float_to_bf16(float f) { + uint32_t bits; + static_assert(sizeof(f) == 4, "float is assumed to be IEEE binary32"); + std::memcpy(&bits, &f, 4); + const uint32_t exp = (bits >> 23) & 0xFFu; + if (exp == 0xFFu && (bits & 0x7FFFFFu) != 0) { + // Keep the result a NaN the way the hardware conversion does. + return (uint16_t)((bits >> 16) | 0x0040u); + } + const uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return (uint16_t)((bits + bias) >> 16); +} + +// 6-bit scale / min pair j of a Q4_K or Q5_K super-block, as in ggml-quants.c. +GGML_BLK_HOST_DEVICE void get_scale_min_k4(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] >> 6) << 4); + } +} + +// Field offsets inside one block (see the alignment note above). +constexpr int32_t Q8_0_OFF_D = 0, Q8_0_OFF_QS = 2; +constexpr int32_t QK_OFF_DMIN = 2, QK_OFF_SCALES = 4; +constexpr int32_t Q4K_OFF_QS = 16; +constexpr int32_t Q5K_OFF_QH = 16, Q5K_OFF_QS = 48; +constexpr int32_t Q6K_OFF_QH = 128, Q6K_OFF_SCALES = 192, Q6K_OFF_D = 208; + +// out: QK8_0 fp32 values. +GGML_BLK_HOST_DEVICE void decode_q8_0(const uint8_t *blk, float *out) { + const float d = half_to_float(read_u16(blk + Q8_0_OFF_D)); + const uint8_t *qs = blk + Q8_0_OFF_QS; + for (int32_t j = 0; j < QK8_0; ++j) { + out[j] = (float)(int8_t)qs[j] * d; + } +} + +// out: QK_K fp32 values. +GGML_BLK_HOST_DEVICE void decode_q4_K(const uint8_t *blk, float *out) { + const float d = half_to_float(read_u16(blk)); + const float min = half_to_float(read_u16(blk + QK_OFF_DMIN)); + const uint8_t *scales = blk + QK_OFF_SCALES; + const uint8_t *q = blk + Q4K_OFF_QS; + for (int32_t j = 0, is = 0; j < QK_K; j += 64, is += 2) { + uint8_t sc, m; + get_scale_min_k4(is + 0, scales, sc, m); + const float d1 = d * sc; + const float m1 = min * m; + get_scale_min_k4(is + 1, scales, sc, m); + const float d2 = d * sc; + const float m2 = min * m; + for (int32_t l = 0; l < 32; ++l) out[j + l] = d1 * (float)(q[l] & 0xF) - m1; + for (int32_t l = 0; l < 32; ++l) out[j + 32 + l] = d2 * (float)(q[l] >> 4) - m2; + q += 32; + } +} + +// out: QK_K fp32 values. +GGML_BLK_HOST_DEVICE void decode_q5_K(const uint8_t *blk, float *out) { + const float d = half_to_float(read_u16(blk)); + const float min = half_to_float(read_u16(blk + QK_OFF_DMIN)); + const uint8_t *scales = blk + QK_OFF_SCALES; + const uint8_t *qh = blk + Q5K_OFF_QH; + const uint8_t *ql = blk + Q5K_OFF_QS; + uint8_t u1 = 1, u2 = 2; + for (int32_t j = 0, is = 0; j < QK_K; j += 64, is += 2) { + uint8_t sc, m; + get_scale_min_k4(is + 0, scales, sc, m); + const float d1 = d * sc; + const float m1 = min * m; + get_scale_min_k4(is + 1, scales, sc, m); + const float d2 = d * sc; + const float m2 = min * m; + for (int32_t l = 0; l < 32; ++l) + out[j + l] = d1 * (float)((ql[l] & 0xF) + ((qh[l] & u1) ? 16 : 0)) - m1; + for (int32_t l = 0; l < 32; ++l) + out[j + 32 + l] = d2 * (float)((ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0)) - m2; + ql += 32; + u1 <<= 2; + u2 <<= 2; + } + // qh never advances: the 5th bit of the 64 elements of one iteration lives + // in two different bit positions of the same 32 qh bytes, which is what the + // shifting u1/u2 masks walk instead of a second pointer. +} + +// out: QK_K fp32 values. +GGML_BLK_HOST_DEVICE void decode_q6_K(const uint8_t *blk, float *out) { + const float d = half_to_float(read_u16(blk + Q6K_OFF_D)); + const uint8_t *ql = blk; + const uint8_t *qh = blk + Q6K_OFF_QH; + const int8_t *sc = (const int8_t *)(blk + Q6K_OFF_SCALES); + for (int32_t n = 0, y = 0; n < QK_K; n += 128, y += 128) { + for (int32_t l = 0; l < 32; ++l) { + const int is = l / 16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + out[y + l + 0] = d * sc[is + 0] * q1; + out[y + l + 32] = d * sc[is + 2] * q2; + out[y + l + 64] = d * sc[is + 4] * q3; + out[y + l + 96] = d * sc[is + 6] * q4; + } + ql += 64; + qh += 32; + sc += 8; + } +} + +// Decode `n_blocks` consecutive blocks of `type` into fp32. +// Returns false without writing when `type` has no decoder. +GGML_BLK_HOST_DEVICE bool decode_blocks(int32_t type, const uint8_t *blk, + int64_t n_blocks, float *out) { + const int32_t bytes = block_bytes(type); + const int32_t elems = block_elems(type); + if (bytes < 0) return false; + for (int64_t i = 0; i < n_blocks; ++i) { + float *o = out + i * elems; + switch (type) { + case GGML_TYPE_Q8_0: decode_q8_0(blk + i * bytes, o); break; + case GGML_TYPE_Q4_K: decode_q4_K(blk + i * bytes, o); break; + case GGML_TYPE_Q5_K: decode_q5_K(blk + i * bytes, o); break; + case GGML_TYPE_Q6_K: decode_q6_K(blk + i * bytes, o); break; + default: return false; + } + } + return true; +} + +// Same as decode_blocks, storing bf16 bit patterns; `tmp` must hold +// block_elems(type) fp32 values. +template +GGML_BLK_HOST_DEVICE bool decode_blocks_bf16(int32_t type, const uint8_t *blk, + int64_t n_blocks, uint16_t *out) { + const int32_t bytes = block_bytes(type); + const int32_t elems = block_elems(type); + if (bytes < 0 || elems > MAX_ELEMS) return false; + for (int64_t i = 0; i < n_blocks; ++i) { + float tmp[MAX_ELEMS]; + switch (type) { + case GGML_TYPE_Q8_0: decode_q8_0(blk + i * bytes, tmp); break; + case GGML_TYPE_Q4_K: decode_q4_K(blk + i * bytes, tmp); break; + case GGML_TYPE_Q5_K: decode_q5_K(blk + i * bytes, tmp); break; + case GGML_TYPE_Q6_K: decode_q6_K(blk + i * bytes, tmp); break; + default: return false; + } + uint16_t *o = out + i * elems; + for (int32_t j = 0; j < elems; ++j) o[j] = float_to_bf16(tmp[j]); + } + return true; +} + +#undef GGML_BLK_HOST_DEVICE + +} // namespace ggml_blocks + +#endif diff --git a/src/infiniop/ops/linear_gguf/info.h b/src/infiniop/ops/linear_gguf/info.h new file mode 100644 index 000000000..5802165f6 --- /dev/null +++ b/src/infiniop/ops/linear_gguf/info.h @@ -0,0 +1,84 @@ +#ifndef __LINEAR_GGUF_INFO_H__ +#define __LINEAR_GGUF_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" +#include "ggml_blocks.h" + +#include + +namespace op::linear_gguf { + +class LinearGgufInfo { + LinearGgufInfo() = default; + +public: + size_t m_count; // batch size + size_t n_count; // output features + size_t k_count; // input features, in elements (not bytes) + int32_t ggml_type; // enum ggml_type id of the weight blocks + int32_t block_elems; // elements per block, from ggml_blocks + int32_t block_bytes; // bytes per block, from ggml_blocks + int64_t row_bytes; // packed bytes of one weight row + bool out_is_f32; // experimental decode path: keep GEMV accumulator as F32 + + static utils::Result create( + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t w_desc, + int64_t ggml_type) { + CHECK_OR_RETURN(out_desc != nullptr && a_desc != nullptr && w_desc != nullptr, + INFINI_STATUS_NULL_POINTER); + + CHECK_DTYPE(a_desc->dtype(), INFINI_DTYPE_BF16); + CHECK_DTYPE(out_desc->dtype(), INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + // The weight is block bytes, so its "features" are bytes, not elements. + CHECK_DTYPE(w_desc->dtype(), INFINI_DTYPE_U8); + + CHECK_OR_RETURN(a_desc->ndim() == 2 && out_desc->ndim() == 2 + && w_desc->ndim() == 2, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(a_desc->isContiguous() && out_desc->isContiguous() + && w_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + + CHECK_OR_RETURN(ggml_type >= 0 && ggml_type <= INT32_MAX, + INFINI_STATUS_BAD_PARAM); + const int32_t type = static_cast(ggml_type); + const int32_t elems = ggml_blocks::block_elems(type); + const int32_t bytes = ggml_blocks::block_bytes(type); + // Every type without a decoder reports 0 / -1 here, so it is rejected + // the same way a nonexistent id is. + CHECK_OR_RETURN(elems > 0 && bytes > 0, INFINI_STATUS_BAD_PARAM); + + const size_t m_count = a_desc->dim(0); + const size_t k_count = a_desc->dim(1); + const size_t n_count = w_desc->dim(0); + const int64_t packed = static_cast(w_desc->dim(1)); + + CHECK_OR_RETURN(m_count > 0 && n_count > 0 && k_count > 0, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(out_desc->dim(0) == m_count && out_desc->dim(1) == n_count, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(k_count % static_cast(elems) == 0, + INFINI_STATUS_BAD_TENSOR_SHAPE); + // Exact row size: a row may not carry padding past its last block. + CHECK_OR_RETURN(packed == static_cast(k_count / elems) * bytes, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + LinearGgufInfo info; + info.m_count = m_count; + info.n_count = n_count; + info.k_count = k_count; + info.ggml_type = type; + info.block_elems = elems; + info.block_bytes = bytes; + info.row_bytes = packed; + info.out_is_f32 = out_desc->dtype() == INFINI_DTYPE_F32; + return utils::Result(info); + } +}; + +} // namespace op::linear_gguf + +#endif diff --git a/src/infiniop/ops/linear_gguf/linear_gguf.h b/src/infiniop/ops/linear_gguf/linear_gguf.h new file mode 100644 index 000000000..c77c62284 --- /dev/null +++ b/src/infiniop/ops/linear_gguf/linear_gguf.h @@ -0,0 +1,49 @@ +#ifndef __LINEAR_GGUF_H__ +#define __LINEAR_GGUF_H__ + +#include "../../operator.h" +#include "info.h" + +#include + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::linear_gguf::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + LinearGgufInfo _info; \ + size_t _workspace_size; \ + \ + Descriptor( \ + Opaque *opaque, \ + LinearGgufInfo info, \ + size_t workspace_size, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), \ + _info(info), \ + _workspace_size(workspace_size) {} \ + \ + public: \ + ~Descriptor(); \ + \ + size_t workspaceSize() const { return _workspace_size; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t out_desc, \ + infiniopTensorDescriptor_t a_desc, \ + infiniopTensorDescriptor_t w_desc, \ + int64_t ggml_type); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, \ + void *out, \ + const void *a, const void *w, \ + void *stream) const; \ + }; \ + } + +#endif // __LINEAR_GGUF_H__ diff --git a/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_dequant.cuh b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_dequant.cuh new file mode 100644 index 000000000..715babf26 --- /dev/null +++ b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_dequant.cuh @@ -0,0 +1,176 @@ +// Prefill-path weight decoding for GGUF block-quantized weights, NVIDIA backend. +// +// C[M, N] = A[M, K] @ W[N, K]^T, M > kMaxDecodeM +// +// The register-resident GEMV kernel in linear_gguf_gemv.cuh gives up past +// kMaxDecodeM because it keeps one fp32 accumulator per input row in registers. +// A larger batch wants a real gemm instead, but `W` is still a packed GGUF blob, +// so launch_prefill below expands kPrefillTileN weight rows into a BF16 scratch at +// a time and hands that tile to cublas. Only one tile is ever live -- a full +// dequantized copy of the weight would defeat the memory property that route B +// exists to keep. +// +// Tiling walks N rather than K: `C` is [M, N], so distinct tiles own distinct, +// non-overlapping output columns and every cublas call can use beta = 0. +// Accumulating along K instead would need an fp32 copy of `C` plus a conversion +// pass, for no benefit at these shapes. +// +// Block decoding is not duplicated here either: decode_block and +// BlockTraits come from linear_gguf_gemv.cuh, whose shared source ggml_blocks.h +// was verified bit-exactly against gguf-py on both host and device in stage 3.1. +// +// Status: correctness-first v1 (stage 3.3). The scratch is written and then read +// back once per gemm, so prefill moves roughly 3x the traffic a fused +// decode-in-shared-memory kernel would; closing that gap is stage 6.1 (MMQ-style), +// not a precondition for correct results. +#ifndef __LINEAR_GGUF_NVIDIA_DEQUANT_CUH__ +#define __LINEAR_GGUF_NVIDIA_DEQUANT_CUH__ + +#include +#include +#include + +#include +#include + +#include "linear_gguf_gemv.cuh" + +namespace op::linear_gguf::nvidia { + +// Weight rows decoded per scratch tile. 64 keeps the scratch at +// 64 * K * 2 bytes (1.25 MiB at K = 10240) and hands cublas a [K, 64] operand +// whose leading dimension is still a multiple of 8 elements, as bf16 gemms want. +constexpr int kPrefillTileN = 64; +constexpr int kDequantThreads = 256; + +// Bytes of BF16 scratch that Descriptor::calculate needs for one tile. `k` is in +// elements, not packed bytes. +inline size_t prefill_scratch_bytes(int64_t k) { + return static_cast(kPrefillTileN) * static_cast(k) * sizeof(__nv_bfloat16); +} + +template +__global__ void dequant_tile_kernel(const uint8_t *__restrict__ w, + __nv_bfloat16 *__restrict__ tile, + int64_t n_start, int rows, int k, int64_t row_bytes) { + constexpr int32_t kBytes = BlockTraits::kBytes; + constexpr int32_t kElems = BlockTraits::kElems; + + const int blocks_per_row = k / kElems; + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int r = static_cast(idx / blocks_per_row); + if (r >= rows) return; + const int b = static_cast(idx - static_cast(r) * blocks_per_row); + + float wk[kElems]; + decode_block(w + (n_start + static_cast(r)) * row_bytes + + static_cast(b) * kBytes, + wk); + + // kElems is even for every supported type and b * kElems is therefore a + // 4-byte-aligned element offset, so pairs can be stored as one 32-bit write. + uint32_t *out = reinterpret_cast( + tile + static_cast(r) * k + static_cast(b) * kElems); +#pragma unroll + for (int32_t j = 0; j < kElems; j += 2) { + const __nv_bfloat162 pair = __floats2bfloat162_rn(wk[j], wk[j + 1]); + out[j / 2] = *reinterpret_cast(&pair); + } +} + +// Decode weight rows [n_start, n_start + rows) of `w` into `tile`, which is row +// major with leading dimension `k`. Returns false without launching anything +// when `type` has no decoder or the geometry is not a whole number of blocks per +// row; the caller must not turn that into a dense fallback. +inline bool launch_dequant_tile(int32_t type, const uint8_t *w, __nv_bfloat16 *tile, + int64_t n_start, int rows, int k, int64_t row_bytes, + cudaStream_t stream) { + const int32_t elems = ggml_blocks::block_elems(type); + const int32_t bytes = ggml_blocks::block_bytes(type); + if (rows <= 0 || k <= 0 || elems <= 0 || bytes <= 0) return false; + if (n_start < 0) return false; + if (k % elems != 0) return false; + if (row_bytes < static_cast(k / elems) * bytes) return false; + + const int64_t total = static_cast(rows) * (k / elems); + const unsigned grid = static_cast( + (total + kDequantThreads - 1) / kDequantThreads); + switch (type) { + case ggml_blocks::GGML_TYPE_Q8_0: + dequant_tile_kernel<<>>( + w, tile, n_start, rows, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q4_K: + dequant_tile_kernel<<>>( + w, tile, n_start, rows, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q5_K: + dequant_tile_kernel<<>>( + w, tile, n_start, rows, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q6_K: + dequant_tile_kernel<<>>( + w, tile, n_start, rows, k, row_bytes); + break; + default: + return false; + } + return cudaGetLastError() == cudaSuccess; +} + +// c[m, n] = a[m, k] @ w[n, k]^T for m > kMaxDecodeM, one decoded weight tile at a +// time. `scratch` / `scratch_bytes` come from the op's workspace tensor; the +// function refuses rather than silently allocating, so an undersized workspace is +// never mistaken for a wrong result. +// +// `blas` is the caller's cublas handle: the op borrows one from the device handle +// pool (devices/nvidia/nvidia_common.cu::useCublas) while +// scripts/gguf_routeb_gemv_probe.cu creates its own, so the numerical gate and the +// shipped path run this same composition. +// +// Tiling walks n, so each tile owns disjoint output columns and every gemm has +// beta = 0 -- no fp32 accumulator of `c` and no conversion pass. The operand +// mapping (cublas is column-major, so row-major c[m, n] is [n, m]) is the same one +// ops/scaled_mm/nvidia/int8_gemm_nvidia.cu uses: +// A = tile stored [k, rows], op(A) = A^T -> [rows, k], lda = k +// B = a stored [k, m], op(B) = B -> [k, m], ldb = k +// C = c + n0 stored [n, m], ldc = n +inline bool launch_prefill(cublasHandle_t blas, int32_t type, + const __nv_bfloat16 *a, const uint8_t *w, __nv_bfloat16 *c, + int m, int n, int k, int64_t row_bytes, + void *scratch, size_t scratch_bytes, cudaStream_t stream) { + if (blas == nullptr || scratch == nullptr) return false; + // The regular route uses this composition only for prefill. Strict + // compatibility experiments may deliberately route small-M decode here as + // well, so validate only the actual matrix geometry. + if (m <= 0 || n <= 0 || k <= 0) return false; + if (scratch_bytes < prefill_scratch_bytes(k)) return false; + + auto *tile = reinterpret_cast<__nv_bfloat16 *>(scratch); + const float alpha = 1.0f; + const float beta = 0.0f; + if (cublasSetStream(blas, stream) != CUBLAS_STATUS_SUCCESS) return false; + + for (int n0 = 0; n0 < n; n0 += kPrefillTileN) { + const int rows = std::min(kPrefillTileN, n - n0); + if (!launch_dequant_tile(type, w, tile, n0, rows, k, row_bytes, stream)) return false; + if (cublasGemmEx(blas, + CUBLAS_OP_T, // A = tile^T : [rows, k] + CUBLAS_OP_N, // B = a viewed column-major : [k, m] + rows, m, k, + &alpha, + tile, CUDA_R_16BF, k, + a, CUDA_R_16BF, k, + &beta, + c + n0, CUDA_R_16BF, n, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT) != CUBLAS_STATUS_SUCCESS) { + return false; + } + } + return cudaGetLastError() == cudaSuccess; +} + +} // namespace op::linear_gguf::nvidia + +#endif // __LINEAR_GGUF_NVIDIA_DEQUANT_CUH__ diff --git a/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_gemv.cuh b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_gemv.cuh new file mode 100644 index 000000000..70427cb9a --- /dev/null +++ b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_gemv.cuh @@ -0,0 +1,315 @@ +// Decode-path (small M) GEMV for GGUF block-quantized weights, NVIDIA backend. +// +// C[M, N] = A[M, K] @ W[N, K]^T +// +// A is BF16 row-major. C is BF16 on the normal path, with an opt-in F32 output +// variant for strict-consistency diagnostics. W is the packed GGUF blob: N rows +// of `row_bytes` bytes, each row holding K / block_elems(TYPE) quantizer blocks +// of the given ggml type. Nothing is ever expanded to BF16 in memory -- a block +// is decoded into registers and immediately consumed by the dot product, which +// is the whole point of route B (weights stay at ~6.56 bits per element and the +// kernel is bandwidth bound on W). +// +// Block decoding is *not* duplicated here: it comes from ggml_blocks.h, the same +// single source of truth that was verified bit-exactly against numpy / gguf-py +// for both host and device in stage 3.1. +// +// Layout of the work: one warp owns one output row n, lane l takes quantizer +// blocks l, l+32, ... along K, and each lane accumulates a private fp32 partial +// per input row. A warp-local shuffle reduction then writes C. That keeps the +// weight row contiguous per warp (the access pattern that matters, W being ~99% +// of the traffic) and needs no shared memory for M <= kMaxDecodeM. +// +// Status: correctness-first v1 (stage 3.2). The reduction order is "sum inside +// a block, then across blocks per lane, then across lanes", so results differ +// from a plain fp32 matmul in the last ulp only; the acceptance gate is +// cos_sim > 0.999 against the dense reference (plan 1.2 item 2). Vectorized +// weight loads and a shared-memory A tile are stage 6 work. +#ifndef __LINEAR_GGUF_NVIDIA_GEMV_CUH__ +#define __LINEAR_GGUF_NVIDIA_GEMV_CUH__ + +#include +#include + +#include "../ggml_blocks.h" + +namespace op::linear_gguf::nvidia { + +// Compile-time register capacity. The shipped dispatch still uses M <= 8 by +// default; an opt-in strict-consistency experiment may use the extra slots for +// short prefills (M <= 16) without changing the normal path. +constexpr int kMaxDecodeM = 16; + +template +struct BlockTraits; + +template <> +struct BlockTraits { + static constexpr int32_t kBytes = ggml_blocks::SIZE_Q8_0; + static constexpr int32_t kElems = ggml_blocks::QK8_0; +}; +template <> +struct BlockTraits { + static constexpr int32_t kBytes = ggml_blocks::SIZE_Q4_K; + static constexpr int32_t kElems = ggml_blocks::QK_K; +}; +template <> +struct BlockTraits { + static constexpr int32_t kBytes = ggml_blocks::SIZE_Q5_K; + static constexpr int32_t kElems = ggml_blocks::QK_K; +}; +template <> +struct BlockTraits { + static constexpr int32_t kBytes = ggml_blocks::SIZE_Q6_K; + static constexpr int32_t kElems = ggml_blocks::QK_K; +}; + +// Forwards to the shared decoders; TYPE keeps the branch compile-time so only +// one call survives instantiation. +template +__device__ __forceinline__ void decode_block(const uint8_t *blk, float *out) { + if constexpr (TYPE == ggml_blocks::GGML_TYPE_Q8_0) { + ggml_blocks::decode_q8_0(blk, out); + } else if constexpr (TYPE == ggml_blocks::GGML_TYPE_Q4_K) { + ggml_blocks::decode_q4_K(blk, out); + } else if constexpr (TYPE == ggml_blocks::GGML_TYPE_Q5_K) { + ggml_blocks::decode_q5_K(blk, out); + } else { + ggml_blocks::decode_q6_K(blk, out); + } +} + +__device__ __forceinline__ void store_gemv_value(__nv_bfloat16 *c, int64_t idx, float v) { + c[idx] = __float2bfloat16_rn(v); +} + +__device__ __forceinline__ void store_gemv_value(float *c, int64_t idx, float v) { + c[idx] = v; +} + +// Return the signed 6-bit integer code for element j of one Q6_K super-block. +// The layout follows block_q6_K in llama.cpp: two 128-element halves, with +// low nibbles in ql and two high bits in qh. +__device__ __forceinline__ int q6_k_code(const uint8_t *blk, int j) { + const int half = j >> 7; + const int within = j & 127; + const int quarter = within >> 5; + const int lane = within & 31; + const uint8_t *ql = blk + half * 64; + const uint8_t *qh = blk + ggml_blocks::Q6K_OFF_QH + half * 32; + int low; + if (quarter == 0) { + low = ql[lane] & 0x0f; + } else if (quarter == 1) { + low = ql[lane + 32] & 0x0f; + } else if (quarter == 2) { + low = ql[lane] >> 4; + } else { + low = ql[lane + 32] >> 4; + } + const int high = (qh[lane] >> (2 * quarter)) & 0x03; + return (low | (high << 4)) - 32; +} + +template +__global__ void gemv_decode_kernel(const __nv_bfloat16 *__restrict__ a, + const uint8_t *__restrict__ w, + OutT *__restrict__ c, int m_count, int n_count, + int k, int64_t row_bytes) { + constexpr int32_t kBytes = BlockTraits::kBytes; + constexpr int32_t kElems = BlockTraits::kElems; + + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int warps_per_block = static_cast(blockDim.x) >> 5; + const int n = blockIdx.x * warps_per_block + warp; + if (n >= n_count) return; + + const uint8_t *wrow = w + static_cast(n) * row_bytes; + const int blocks_per_row = k / kElems; + + float acc[kMaxDecodeM]; +#pragma unroll + for (int m = 0; m < kMaxDecodeM; ++m) acc[m] = 0.0f; + + for (int b = lane; b < blocks_per_row; b += 32) { + float wk[kElems]; + decode_block(wrow + static_cast(b) * kBytes, wk); + const __nv_bfloat16 *ab = a + static_cast(b) * kElems; +#pragma unroll + for (int m = 0; m < kMaxDecodeM; ++m) { + if (m >= m_count) break; + const __nv_bfloat16 *am = ab + static_cast(m) * k; + float s = 0.0f; + if constexpr (QUANTIZE_ACTIVATION) { + // llama.cpp's CUDA matvec path quantizes the activation side to + // Q8_1 in groups of 32 before the quantized vec-dot. + if constexpr (TYPE == ggml_blocks::GGML_TYPE_Q6_K) { + // Formula-level Q6_K x Q8_1 path: preserve integer q6/q8 + // dot products and sub-block scales instead of multiplying + // two independently dequantized float vectors. + const uint8_t *blk = wrow + static_cast(b) * kBytes; + const int8_t *scales = reinterpret_cast( + blk + ggml_blocks::Q6K_OFF_SCALES); + const float d6 = ggml_blocks::half_to_float( + ggml_blocks::read_u16(blk + ggml_blocks::Q6K_OFF_D)); + float scaled_integer_sum = 0.0f; +#pragma unroll + for (int base = 0; base < kElems; base += 32) { + float amax = 0.0f; +#pragma unroll + for (int j = 0; j < 32; ++j) { + amax = fmaxf(amax, fabsf(__bfloat162float(__ldg(am + base + j)))); + } + const float d_inv = 127.0f / amax; + const float d8 = 1.0f / d_inv; + int dot0 = 0; + int dot1 = 0; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const float av = __bfloat162float(__ldg(am + base + j)); + const int q8 = amax == 0.0f ? 0 : static_cast(roundf(av * d_inv)); + dot0 += q8 * q6_k_code(blk, base + j); + } +#pragma unroll + for (int j = 16; j < 32; ++j) { + const float av = __bfloat162float(__ldg(am + base + j)); + const int q8 = amax == 0.0f ? 0 : static_cast(roundf(av * d_inv)); + dot1 += q8 * q6_k_code(blk, base + j); + } + const int scale_idx = base / 16; + const int scaled_dot = + static_cast(scales[scale_idx]) * dot0 + + static_cast(scales[scale_idx + 1]) * dot1; + scaled_integer_sum += d8 * static_cast(scaled_dot); + } + s = d6 * scaled_integer_sum; + } else { +#pragma unroll + for (int base = 0; base < kElems; base += 32) { + float amax = 0.0f; +#pragma unroll + for (int j = 0; j < 32; ++j) { + amax = fmaxf(amax, fabsf(__bfloat162float(__ldg(am + base + j)))); + } + const float d_inv = 127.0f / amax; + const float d = 1.0f / d_inv; +#pragma unroll + for (int j = 0; j < 32; ++j) { + const float av = __bfloat162float(__ldg(am + base + j)); + const float q = amax == 0.0f ? 0.0f : roundf(av * d_inv); + s += (d * q) * wk[base + j]; + } + } + } + } else { +#pragma unroll + for (int j = 0; j < kElems; ++j) { + s += __bfloat162float(__ldg(am + j)) * wk[j]; + } + } + acc[m] += s; + } + } + + const unsigned members = 0xffffffffu; +#pragma unroll + for (int m = 0; m < kMaxDecodeM; ++m) { + if (m >= m_count) break; + float v = acc[m]; +#pragma unroll + for (int off = 16; off; off >>= 1) { + v += __shfl_down_sync(members, v, off); + } + if (lane == 0) { + store_gemv_value(c, static_cast(m) * n_count + n, v); + } + } +} + +constexpr int kGemvThreads = 256; + +// Returns false when `type` has no decoder here or the geometry is not a whole +// number of blocks per row, without launching anything. The caller turns that +// into an error: silently falling back to a dense weight would defeat the point +// of route B. +template +inline bool launch_gemv_decode_typed(int32_t type, const __nv_bfloat16 *a, const uint8_t *w, + OutT *c, int m_count, int n_count, int k, + int64_t row_bytes, bool quantize_activation, + cudaStream_t stream) { + if (m_count <= 0 || m_count > kMaxDecodeM || n_count <= 0 || k <= 0) return false; + const int32_t elems = ggml_blocks::block_elems(type); + const int32_t bytes = ggml_blocks::block_bytes(type); + if (elems <= 0 || bytes <= 0) return false; + if (k % elems != 0) return false; + if (row_bytes < static_cast(k / elems) * bytes) return false; + + const int warps_per_block = kGemvThreads >> 5; + const unsigned grid = static_cast((n_count + warps_per_block - 1) / warps_per_block); + if (quantize_activation) { + switch (type) { + case ggml_blocks::GGML_TYPE_Q8_0: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q4_K: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q5_K: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q6_K: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + default: + return false; + } + return cudaGetLastError() == cudaSuccess; + } + + switch (type) { + case ggml_blocks::GGML_TYPE_Q8_0: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q4_K: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q5_K: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + case ggml_blocks::GGML_TYPE_Q6_K: + gemv_decode_kernel<<>>( + a, w, c, m_count, n_count, k, row_bytes); + break; + default: + return false; + } + return cudaGetLastError() == cudaSuccess; +} + +inline bool launch_gemv_decode(int32_t type, const __nv_bfloat16 *a, const uint8_t *w, + __nv_bfloat16 *c, int m_count, int n_count, int k, + int64_t row_bytes, bool quantize_activation, + cudaStream_t stream) { + return launch_gemv_decode_typed(type, a, w, c, m_count, n_count, k, + row_bytes, quantize_activation, stream); +} + +inline bool launch_gemv_decode_f32(int32_t type, const __nv_bfloat16 *a, const uint8_t *w, + float *c, int m_count, int n_count, int k, + int64_t row_bytes, bool quantize_activation, + cudaStream_t stream) { + return launch_gemv_decode_typed(type, a, w, c, m_count, n_count, k, + row_bytes, quantize_activation, stream); +} + +} // namespace op::linear_gguf::nvidia + +#endif // __LINEAR_GGUF_NVIDIA_GEMV_CUH__ diff --git a/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cu b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cu new file mode 100644 index 000000000..74833ef3e --- /dev/null +++ b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cu @@ -0,0 +1,191 @@ +#if defined(ENABLE_NVIDIA_API) + +#include "linear_gguf_nvidia.cuh" + +// Opens namespace op::linear_gguf::nvidia itself, so it has to stay at file +// scope rather than move inside that namespace. +#include "linear_gguf_gemv.cuh" +#include "linear_gguf_dequant.cuh" + +#include "../../../devices/nvidia/nvidia_handle.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +#include +#include + +#include +#include + +namespace op::linear_gguf::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { delete _opaque; } + +namespace { + +bool force_decode_cublas() { + const char *value = std::getenv("INFINI_GGUF_DECODE_CUBLAS"); + return value != nullptr && std::strcmp(value, "0") != 0; +} + +bool strict_small_prefill(int32_t ggml_type) { + const char *value = std::getenv("INFINI_GGUF_STRICT_SMALL_PREFILL"); + if (value == nullptr || std::strcmp(value, "0") == 0) { + return false; + } + // Optional diagnostic filter. When present, only this GGML weight type + // gets the short-prefill register path; e.g. 12=Q4_K, 13=Q5_K, + // 14=Q6_K, 8=Q8_0. Omission preserves the original all-types experiment. + const char *type_value = + std::getenv("INFINI_GGUF_STRICT_SMALL_PREFILL_TYPE"); + return type_value == nullptr || *type_value == '\0' + || std::atoi(type_value) == ggml_type; +} + +size_t strict_small_prefill_limit(int32_t ggml_type) { + constexpr size_t kDefaultDecodeM = 8; + if (!strict_small_prefill(ggml_type)) { + return kDefaultDecodeM; + } + const char *max_m_value = + std::getenv("INFINI_GGUF_STRICT_SMALL_PREFILL_MAX_M"); + if (max_m_value != nullptr && *max_m_value != '\0') { + const int requested = std::atoi(max_m_value); + if (requested > static_cast(kDefaultDecodeM) + && requested <= kMaxDecodeM) { + return static_cast(requested); + } + } + return static_cast(kMaxDecodeM); +} + +bool use_register_gemv(size_t m_count, int32_t ggml_type) { + const size_t limit = strict_small_prefill_limit(ggml_type); + return !force_decode_cublas() && m_count <= limit; +} + +bool quantize_decode_activation(int32_t ggml_type) { + const char *type_value = std::getenv("INFINI_GGUF_DECODE_Q8A_TYPE"); + if (type_value != nullptr && *type_value != '\0') { + return std::atoi(type_value) == ggml_type; + } + const char *value = std::getenv("INFINI_GGUF_DECODE_Q8A"); + return value != nullptr && std::strcmp(value, "0") != 0; +} + +} // namespace + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t w_desc, + int64_t ggml_type) { + auto info_result = LinearGgufInfo::create(out_desc, a_desc, w_desc, ggml_type); + if (!info_result) { + return info_result.status(); + } + auto info = info_result.take(); + + // The register-resident GEMV only covers small batches. A larger M goes to + // the prefill path, which needs one BF16 weight tile as scratch, so ask for it + // here -- the infinicore wrapper sizes its workspace tensor from + // workspaceSize(). M <= kMaxDecodeM keeps a zero-size workspace, leaving the + // decode path exactly as cheap as it was. + size_t workspace_size = 0; + if (!use_register_gemv(info.m_count, info.ggml_type)) { + workspace_size = prefill_scratch_bytes(static_cast(info.k_count)); + } + + auto nvidia_handle = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + new Opaque{nvidia_handle->internal()}, std::move(info), workspace_size, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +namespace { + +// All of the prefill composition (tile decode + gemm) lives in +// linear_gguf_dequant.cuh::launch_prefill, because scripts/gguf_routeb_gemv_probe.cu +// drives exactly that function through its own cublas handle. Keeping it here +// instead would have the numerical gate test a copy of the shipped path, which is +// the mistake stage 3.2 had to undo (log A.8). +infiniStatus_t run_prefill( + const std::shared_ptr &internal, + const LinearGgufInfo &info, + void *workspace, size_t workspace_size, void *out, + const void *a, const void *w, cudaStream_t stream) { + bool ok = false; + CHECK_STATUS(internal->useCublas(stream, [&](cublasHandle_t blas) { + ok = launch_prefill( + blas, info.ggml_type, + reinterpret_cast(a), + reinterpret_cast(w), + reinterpret_cast<__nv_bfloat16 *>(out), + static_cast(info.m_count), static_cast(info.n_count), + static_cast(info.k_count), info.row_bytes, + workspace, workspace_size, stream); + return INFINI_STATUS_SUCCESS; + })); + return ok ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; +} + +} // namespace + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *out, + const void *a, const void *w, + void *stream) const { + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + if (!use_register_gemv(_info.m_count, _info.ggml_type)) { + if (_info.out_is_f32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + return run_prefill(_opaque->internal, _info, workspace, workspace_size, out, a, w, + reinterpret_cast(stream)); + } + + // The kernel decodes weight blocks on the fly; nothing here touches a + // dequantized copy of the weight. + const bool ok = _info.out_is_f32 + ? launch_gemv_decode_f32( + _info.ggml_type, + reinterpret_cast(a), + reinterpret_cast(w), + reinterpret_cast(out), + static_cast(_info.m_count), + static_cast(_info.n_count), + static_cast(_info.k_count), + _info.row_bytes, + quantize_decode_activation(_info.ggml_type), + reinterpret_cast(stream)) + : launch_gemv_decode( + _info.ggml_type, + reinterpret_cast(a), + reinterpret_cast(w), + reinterpret_cast<__nv_bfloat16 *>(out), + static_cast(_info.m_count), + static_cast(_info.n_count), + static_cast(_info.k_count), + _info.row_bytes, + quantize_decode_activation(_info.ggml_type), + reinterpret_cast(stream)); + if (!ok) { + return INFINI_STATUS_INTERNAL_ERROR; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::linear_gguf::nvidia + +#endif // ENABLE_NVIDIA_API diff --git a/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cuh b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cuh new file mode 100644 index 000000000..4ba08ee0c --- /dev/null +++ b/src/infiniop/ops/linear_gguf/nvidia/linear_gguf_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef INFINIOP_NVIDIA_LINEAR_GGUF_OPERATOR_H +#define INFINIOP_NVIDIA_LINEAR_GGUF_OPERATOR_H + +#include "../linear_gguf.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/linear_gguf/operator.cc b/src/infiniop/ops/linear_gguf/operator.cc new file mode 100644 index 000000000..8b3690db7 --- /dev/null +++ b/src/infiniop/ops/linear_gguf/operator.cc @@ -0,0 +1,100 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/linear_gguf.h" + +#if defined ENABLE_NVIDIA_API +#include "nvidia/linear_gguf_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateLinearGgufDescriptor( + infiniopHandle_t handle, + infiniopLinearGgufDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t w_desc, + int64_t ggml_type) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::linear_gguf::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + out_desc, \ + a_desc, \ + w_desc, \ + ggml_type) + + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetLinearGgufWorkspaceSize(infiniopLinearGgufDescriptor_t desc, + size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopLinearGguf( + infiniopLinearGgufDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *a, + const void *w, + void *stream) { + +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, out, a, w, stream) + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CALCULATE +} + +__INFINI_C infiniStatus_t +infiniopDestroyLinearGgufDescriptor(infiniopLinearGgufDescriptor_t desc) { + +#define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS; + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DELETE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef DELETE +} diff --git a/xmake.lua b/xmake.lua index 568410287..2e25f0626 100644 --- a/xmake.lua +++ b/xmake.lua @@ -76,6 +76,15 @@ option_end() if has_config("nv-gpu") then add_defines("ENABLE_NVIDIA_API") + + -- NVIDIA operator translation units include cuDNN through the shared handle. + -- Keep the location environment-driven so local and CI installations work. + local cudnn_root_global = os.getenv("CUDNN_ROOT") or os.getenv("CUDNN_HOME") or os.getenv("CUDNN_PATH") + if cudnn_root_global ~= nil and cudnn_root_global ~= "" then + add_includedirs(cudnn_root_global .. "/include") + add_linkdirs(cudnn_root_global .. "/lib") + end + includes("xmake/nvidia.lua") end @@ -92,7 +101,7 @@ end option("cuda_arch") set_showmenu(true) set_description("Set CUDA GPU architecture (e.g. sm_90)") - set_values("sm_50", "sm_60", "sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90", "sm_90a") + set_values("sm_50", "sm_60", "sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90", "sm_90a", "sm_120") set_category("option") option_end() diff --git a/xmake/nvidia.lua b/xmake/nvidia.lua index 0d2fd0f81..640711c29 100644 --- a/xmake/nvidia.lua +++ b/xmake/nvidia.lua @@ -1,6 +1,7 @@ local CUDNN_ROOT = os.getenv("CUDNN_ROOT") or os.getenv("CUDNN_HOME") or os.getenv("CUDNN_PATH") if CUDNN_ROOT ~= nil then add_includedirs(CUDNN_ROOT .. "/include") + add_linkdirs(CUDNN_ROOT .. "/lib") end local CUTLASS_ROOT = os.getenv("CUTLASS_ROOT") or os.getenv("CUTLASS_HOME") or os.getenv("CUTLASS_PATH") @@ -23,6 +24,8 @@ local function apply_cuda_arch_flags(add_fn) for _, arch in ipairs(arch_opt:split(",")) do arch = arch:trim() if arch ~= "" then + if arch == "sm_90" then arch = "sm_90a" end + if arch == "sm_120" then arch = "sm_120a" end local compute = arch:gsub("sm_", "compute_") add_fn("-gencode=arch=" .. compute .. ",code=" .. arch) end @@ -118,7 +121,30 @@ target("infiniop-nvidia") end -- CUDA arch: explicit --cuda_arch > nvidia-smi auto-detect > native - if not apply_cuda_arch_flags(function(flag) target:add("cuflags", flag) end) then + -- Use cugencodes for devlink support + cuflags for compilation + local function add_arch_flags(add_fn, add_gencode_fn) + local arch_opt_inner = get_config("cuda_arch") + if not arch_opt_inner or type(arch_opt_inner) ~= "string" or arch_opt_inner == "" then + return false + end + for _, arch in ipairs(arch_opt_inner:split(",")) do + arch = arch:trim() + if arch ~= "" then + -- Use 'a' suffix for architectures that need it (sm_90, sm_120) + if arch == "sm_90" then arch = "sm_90a" end + if arch == "sm_120" then arch = "sm_120a" end + local compute = arch:gsub("sm_", "compute_") + add_fn("-gencode=arch=" .. compute .. ",code=" .. arch) + add_gencode_fn(arch) + end + end + return true + end + + if not add_arch_flags( + function(flag) target:add("cuflags", flag) end, + function(arch) target:add("cugencodes", arch) end + ) then local ok, sm_str = os.iorunv("nvidia-smi", {"--query-gpu=compute_cap", "--format=csv,noheader,nounits"}) if ok and sm_str then local major, minor = sm_str:match("(%d+)%.(%d+)") @@ -129,9 +155,12 @@ target("infiniop-nvidia") if sm >= 80 then table.insert(archs, "sm_80") end if sm >= 86 then table.insert(archs, "sm_86") end if sm >= 89 then table.insert(archs, "sm_89") end - -- H100 (sm_90a): use sm_90a for cutlass 3.x if sm == 90 then target:add("cuflags", "-gencode=arch=compute_90a,code=sm_90a") + target:add("cugencodes", "sm_90a") + elseif sm >= 120 then + target:add("cuflags", "-gencode=arch=compute_120a,code=sm_120a") + target:add("cugencodes", "sm_120a") elseif sm > 90 then table.insert(archs, "sm_90") end @@ -141,6 +170,7 @@ target("infiniop-nvidia") for _, arch in ipairs(archs) do local compute = arch:gsub("sm_", "compute_") target:add("cuflags", "-gencode=arch=" .. compute .. ",code=" .. arch) + target:add("cugencodes", arch) end else target:add("cugencodes", "native") @@ -172,7 +202,7 @@ target("infiniop-nvidia") end end - add_cuflags("-Xcompiler=-Wno-error=deprecated-declarations", "-Xcompiler=-Wno-error=unused-function") + add_cuflags("-Xcompiler=-Wno-error=deprecated-declarations", "-Xcompiler=-Wno-error=unused-function", "-Xcompiler=-Wno-error=sign-compare") -- Cutlass: enable I8 Gemm when CUTLASS_ROOT is set if CUTLASS_ROOT ~= nil then @@ -208,8 +238,11 @@ target("infiniop-nvidia") if arch_opt and type(arch_opt) == "string" then for _, arch in ipairs(arch_opt:split(",")) do arch = arch:trim() + if arch == "sm_90" then arch = "sm_90a" end + if arch == "sm_120" then arch = "sm_120a" end local compute = arch:gsub("sm_", "compute_") add_cuflags("-gencode=arch=" .. compute .. ",code=" .. arch) + add_cugencodes(arch) end end From 575c34aeac5bef8a3d47cfd782c73a12e8a62a10 Mon Sep 17 00:00:00 2001 From: xindongliu594 Date: Thu, 10 Sep 2026 22:20:11 +0800 Subject: [PATCH 2/2] fix: add GGUF operator coverage and correct FP8 scale layout --- include/infiniop.h | 3 +- include/infiniop/ops/linear_gguf.h | 9 +- scripts/python_test.py | 2 + .../nvidia/block_fp8_gemm_sm120.cuh | 2 +- .../nvidia/block_fp8_linear_nvidia.cu | 7 + test/infiniop/block_fp8_linear.py | 189 +++++++++++ test/infiniop/linear_gguf.py | 302 ++++++++++++++++++ 7 files changed, 509 insertions(+), 5 deletions(-) create mode 100644 test/infiniop/block_fp8_linear.py create mode 100644 test/infiniop/linear_gguf.py diff --git a/include/infiniop.h b/include/infiniop.h index 1631f465f..11c29b4c3 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -85,8 +85,9 @@ #include "infiniop/ops/layer_norm.h" #include "infiniop/ops/ldexp.h" #include "infiniop/ops/lerp.h" -#include "infiniop/ops/linear_mxfp4.h" #include "infiniop/ops/block_fp8_linear.h" +#include "infiniop/ops/linear_gguf.h" +#include "infiniop/ops/linear_mxfp4.h" #include "infiniop/ops/log10.h" #include "infiniop/ops/log1p.h" #include "infiniop/ops/log_softmax.h" diff --git a/include/infiniop/ops/linear_gguf.h b/include/infiniop/ops/linear_gguf.h index 959228cf2..9838a8872 100644 --- a/include/infiniop/ops/linear_gguf.h +++ b/include/infiniop/ops/linear_gguf.h @@ -20,9 +20,12 @@ * and accumulated in fp32, so a model loaded this way never materializes a * dense copy of its weights. * - * Current backend implements the decode (GEMV) path only, i.e. M must not - * exceed kMaxDecodeM from the NVIDIA kernel header. A larger M returns - * INFINI_STATUS_NOT_IMPLEMENTED rather than silently dequantizing the weight. + * The NVIDIA backend uses a register-resident GEMV for small M and a tiled + * dequantization plus cuBLAS GEMM path for larger M. The latter requires a + * workspace returned by infiniopGetLinearGgufWorkspaceSize. F32 output is + * supported only by the small-M GEMV path; the regular BF16 path supports both + * decode and prefill. Unsupported block types and malformed packed rows are + * rejected instead of silently falling back to a dense weight. */ typedef struct InfiniopDescriptor *infiniopLinearGgufDescriptor_t; diff --git a/scripts/python_test.py b/scripts/python_test.py index 4879da32a..dc0b6ac7d 100644 --- a/scripts/python_test.py +++ b/scripts/python_test.py @@ -21,6 +21,8 @@ def run_tests(args): # "dequantize_gptq.py", "gelu.py", "gemm.py", + "linear_gguf.py", + "block_fp8_linear.py", # "layer_norm.py", "logsoftmax.py", # "lp_norm.py", diff --git a/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh index 75938ad5c..6be2dd5f4 100644 --- a/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh +++ b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_gemm_sm120.cuh @@ -49,7 +49,7 @@ struct GemmConfig { using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< 1, 128, 128, - cute::UMMA::Major::MN, + cute::UMMA::Major::K, cute::UMMA::Major::K>; using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); diff --git a/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu index ac4941e3b..cafb7c67d 100644 --- a/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu +++ b/src/infiniop/ops/block_fp8_linear/nvidia/block_fp8_linear_nvidia.cu @@ -136,6 +136,13 @@ infiniStatus_t Descriptor::calculate( const void *weight, const void *weight_scale, void *stream) const { + if (workspace_size < _opaque->workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (_opaque->workspace_size > 0 && workspace == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + auto cuda_stream = reinterpret_cast(stream); const size_t M = _info.M; diff --git a/test/infiniop/block_fp8_linear.py b/test/infiniop/block_fp8_linear.py new file mode 100644 index 000000000..5efda9fbe --- /dev/null +++ b/test/infiniop/block_fp8_linear.py @@ -0,0 +1,189 @@ +import ctypes +from ctypes import POINTER, c_int32, c_size_t, c_void_p + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceEnum, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopHandle_t, + infiniopOperatorDescriptor_t, + infiniopTensorDescriptor_t, + test_operator, +) + + +LIBINFINIOP.infiniopCreateBlockFP8LinearDescriptor.restype = c_int32 +LIBINFINIOP.infiniopCreateBlockFP8LinearDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, +] +LIBINFINIOP.infiniopGetBlockFP8LinearWorkspaceSize.restype = c_int32 +LIBINFINIOP.infiniopGetBlockFP8LinearWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), +] +LIBINFINIOP.infiniopBlockFP8Linear.restype = c_int32 +LIBINFINIOP.infiniopBlockFP8Linear.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, +] +LIBINFINIOP.infiniopDestroyBlockFP8LinearDescriptor.restype = c_int32 +LIBINFINIOP.infiniopDestroyBlockFP8LinearDescriptor.argtypes = [ + infiniopOperatorDescriptor_t +] + + +_TEST_CASES = [ + (1, 128, 128), + (16, 256, 128), + (64, 256, 256), +] +_TENSOR_DTYPES = [None] +_BLOCK_SIZE = 128 + + +def _reference(input_source, weight_source, weight_scale): + m_count, k_count = input_source.shape + n_count = weight_source.shape[0] + input_blocks = k_count // _BLOCK_SIZE + input_fp32 = input_source.float().reshape( + m_count, input_blocks, _BLOCK_SIZE + ) + activation_scale = input_fp32.abs().amax(dim=-1).div(448.0).clamp_min(1e-10) + activation_fp8 = ( + input_fp32.div(activation_scale.unsqueeze(-1)) + .clamp(-448.0, 448.0) + .to(torch.float8_e4m3fn) + ) + dequant_activation = ( + activation_fp8.float() * activation_scale.unsqueeze(-1) + ).reshape(m_count, k_count) + expanded_weight_scale = weight_scale.repeat_interleave( + _BLOCK_SIZE, dim=0 + ).repeat_interleave(_BLOCK_SIZE, dim=1)[:n_count, :k_count] + dequant_weight = weight_source.float() * expanded_weight_scale + return (dequant_activation @ dequant_weight.transpose(0, 1)).to( + torch.bfloat16 + ) + + +def test(handle, device, m_count, n_count, k_count, _dtype, sync): + if device != InfiniDeviceEnum.NVIDIA: + print(f"Skipping BlockFP8Linear on {InfiniDeviceNames[device]}") + return + if torch.cuda.get_device_capability()[0] < 12: + print("Skipping BlockFP8Linear because it requires SM120") + return + + print( + f"Testing BlockFP8Linear on NVIDIA with " + f"M={m_count}, N={n_count}, K={k_count}" + ) + generator = torch.Generator(device="cpu").manual_seed( + 2000 + m_count + n_count + k_count + ) + input_source = ( + torch.randn((m_count, k_count), generator=generator) * 0.25 + ).to(torch.bfloat16) + weight_source = ( + torch.randn((n_count, k_count), generator=generator) * 0.25 + ).to(torch.float8_e4m3fn) + scale_shape = ( + (n_count + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + (k_count + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + ) + weight_scale_source = ( + torch.rand(scale_shape, generator=generator, dtype=torch.float32) * 0.5 + + 0.5 + ) + expected = _reference(input_source, weight_source, weight_scale_source) + + input_tensor = TestTensor.from_torch( + input_source, InfiniDtype.BF16, device + ) + weight = TestTensor.from_torch(weight_source, InfiniDtype.F8, device) + weight_scale = TestTensor.from_torch( + weight_scale_source, InfiniDtype.F32, device + ) + output = TestTensor( + (m_count, n_count), None, InfiniDtype.BF16, device, mode="zeros" + ) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateBlockFP8LinearDescriptor( + handle, + ctypes.byref(descriptor), + output.descriptor, + input_tensor.descriptor, + weight.descriptor, + weight_scale.descriptor, + ) + ) + for tensor in (output, input_tensor, weight, weight_scale): + tensor.destroy_desc() + + workspace_size = c_size_t(0) + check_error( + LIBINFINIOP.infiniopGetBlockFP8LinearWorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, device) + if m_count == 1: + status = LIBINFINIOP.infiniopBlockFP8Linear( + descriptor, + workspace.data(), + workspace_size.value - 1, + output.data(), + input_tensor.data(), + weight.data(), + weight_scale.data(), + None, + ) + assert status != 0 + + check_error( + LIBINFINIOP.infiniopBlockFP8Linear( + descriptor, + workspace.data(), + workspace_size.value, + output.data(), + input_tensor.data(), + weight.data(), + weight_scale.data(), + None, + ) + ) + if sync is not None: + sync() + + actual = output.actual_tensor().cpu() + torch.testing.assert_close( + actual.float(), expected.float(), atol=0.25, rtol=0.08 + ) + check_error(LIBINFINIOP.infiniopDestroyBlockFP8LinearDescriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + for test_device in get_test_devices(args): + test_operator(test_device, test, _TEST_CASES, _TENSOR_DTYPES) + print("\033[92mBlockFP8Linear test passed!\033[0m") diff --git a/test/infiniop/linear_gguf.py b/test/infiniop/linear_gguf.py new file mode 100644 index 000000000..c5fd89e86 --- /dev/null +++ b/test/infiniop/linear_gguf.py @@ -0,0 +1,302 @@ +import ctypes +import struct +from ctypes import POINTER, c_int32, c_int64, c_size_t, c_void_p + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceEnum, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopHandle_t, + infiniopOperatorDescriptor_t, + infiniopTensorDescriptor_t, + test_operator, +) + + +LIBINFINIOP.infiniopCreateLinearGgufDescriptor.restype = c_int32 +LIBINFINIOP.infiniopCreateLinearGgufDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_int64, +] +LIBINFINIOP.infiniopGetLinearGgufWorkspaceSize.restype = c_int32 +LIBINFINIOP.infiniopGetLinearGgufWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), +] +LIBINFINIOP.infiniopLinearGguf.restype = c_int32 +LIBINFINIOP.infiniopLinearGguf.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, +] +LIBINFINIOP.infiniopDestroyLinearGgufDescriptor.restype = c_int32 +LIBINFINIOP.infiniopDestroyLinearGgufDescriptor.argtypes = [ + infiniopOperatorDescriptor_t +] + + +Q8_0 = 8 +Q4_K = 12 +Q5_K = 13 +Q6_K = 14 +BLOCK_ELEMS = {Q8_0: 32, Q4_K: 256, Q5_K: 256, Q6_K: 256} +BLOCK_BYTES = {Q8_0: 34, Q4_K: 144, Q5_K: 176, Q6_K: 210} +_TEST_CASES = [ + (ggml_type, m, 8, 256) + for ggml_type in (Q8_0, Q4_K, Q5_K, Q6_K) + for m in (1, 8, 17) +] +_TENSOR_DTYPES = [None] + + +def _signed(value): + return value - 256 if value >= 128 else value + + +def _half(data, offset=0): + return struct.unpack_from("> 6) << 4) + minimum = (scales[index + 4] >> 4) | ((scales[index] >> 6) << 4) + return scale, minimum + + +def _decode_block(ggml_type, block): + if ggml_type == Q8_0: + scale = _half(block) + return [_signed(value) * scale for value in block[2:34]] + + if ggml_type in (Q4_K, Q5_K): + scale = _half(block) + minimum = _half(block, 2) + scales = block[4:16] + high = block[16:48] if ggml_type == Q5_K else None + quants = block[48:] if ggml_type == Q5_K else block[16:] + output = [0.0] * 256 + q_offset = 0 + high_mask_1, high_mask_2 = 1, 2 + for base, scale_index in zip(range(0, 256, 64), range(0, 8, 2)): + scale_1, min_1 = _scale_min(scales, scale_index) + scale_2, min_2 = _scale_min(scales, scale_index + 1) + for lane in range(32): + low = quants[q_offset + lane] + q1 = low & 15 + q2 = low >> 4 + if high is not None: + q1 += 16 if high[lane] & high_mask_1 else 0 + q2 += 16 if high[lane] & high_mask_2 else 0 + output[base + lane] = scale * scale_1 * q1 - minimum * min_1 + output[base + 32 + lane] = ( + scale * scale_2 * q2 - minimum * min_2 + ) + q_offset += 32 + high_mask_1 <<= 2 + high_mask_2 <<= 2 + return output + + scale = _half(block, 208) + low = block[:128] + high = block[128:192] + scales = [_signed(value) for value in block[192:208]] + output = [0.0] * 256 + for half_index in range(2): + low_offset = half_index * 64 + high_offset = half_index * 32 + scale_offset = half_index * 8 + out_offset = half_index * 128 + for lane in range(32): + group = lane // 16 + q1 = (low[low_offset + lane] & 15) | ( + ((high[high_offset + lane] >> 0) & 3) << 4 + ) + q2 = (low[low_offset + lane + 32] & 15) | ( + ((high[high_offset + lane] >> 2) & 3) << 4 + ) + q3 = (low[low_offset + lane] >> 4) | ( + ((high[high_offset + lane] >> 4) & 3) << 4 + ) + q4 = (low[low_offset + lane + 32] >> 4) | ( + ((high[high_offset + lane] >> 6) & 3) << 4 + ) + output[out_offset + lane] = ( + scale * scales[scale_offset + group] * (q1 - 32) + ) + output[out_offset + lane + 32] = ( + scale * scales[scale_offset + group + 2] * (q2 - 32) + ) + output[out_offset + lane + 64] = ( + scale * scales[scale_offset + group + 4] * (q3 - 32) + ) + output[out_offset + lane + 96] = ( + scale * scales[scale_offset + group + 6] * (q4 - 32) + ) + return output + + +def _make_block(ggml_type, seed): + if ggml_type == Q8_0: + quant = [((seed * 13 + index * 17) % 255) - 127 for index in range(32)] + return bytearray(struct.pack("