feat: add reusable GGUF linear backend and FP8 support

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

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 <cstdint>
+
+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

include/infiniop.h

@@ -85,6 +85,8 @@
 #include "infiniop/ops/layer_norm.h"
 #include "infiniop/ops/ldexp.h"
 #include "infiniop/ops/lerp.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/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

include/infiniop/ops/linear_gguf.h

@@ -0,0 +1,51 @@
+#ifndef __INFINIOP_LINEAR_GGUF_API_H__
+#define __INFINIOP_LINEAR_GGUF_API_H__
+
+#include "../operator_descriptor.h"
+#include <cstdint>
+
+/**
+ * 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.
+ *
+ * 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;
+
+__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__

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",

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",

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

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)

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",

src/infinicore/context/allocators/pinnable_block_allocator.cc

@@ -47,37 +47,41 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) {
 
     std::shared_ptr<Block> 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<std::byte *>(block->ptr);
+                    }
                 }
-                if (block != nullptr) {
-                    cls.free_blocks.pop_back();
-                    block->in_use = true;
-                    block->use_count = 1;
-                    return reinterpret_cast<std::byte *>(block->ptr);
-                }
-            }
-            // Allocate a new block for this class
-            block = std::make_shared<Block>();
-            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>();
+                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<std::byte *>(block->ptr);
+                all_blocks_[block->ptr] = block;
+                return reinterpret_cast<std::byte *>(block->ptr);
+            }
         }
     }
 

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