runtime/nvfp4_linear.h

#pragma once

#include <stddef.h>
#include <stdint.h>

#if defined(_WIN32)
#define MAGE_NVFP4_API __declspec(dllexport)
#else
#define MAGE_NVFP4_API __attribute__((visibility("default")))
#endif

#ifdef __cplusplus
extern "C" {
#endif

// The ABI is inference-only. All functions return zero on success and a
// non-zero value on failure unless their return type documents otherwise.
MAGE_NVFP4_API int mage_nvfp4_abi_version(void);
MAGE_NVFP4_API const char* mage_nvfp4_last_error(void);

MAGE_NVFP4_API size_t mage_nvfp4_packed_weight_bytes(int n, int k);
MAGE_NVFP4_API size_t mage_nvfp4_weight_scale_bytes(int n, int k);

// Pack a contiguous CPU BF16 nn.Linear weight in logical row-major [N, K]
// order into cuBLASLt E2M1 payload and tiled VEC16 UE4M3 block scales.
MAGE_NVFP4_API int mage_nvfp4_pack_weight_bf16(
    const void* weight_bf16, int n, int k,
    void* packed_weight, size_t packed_weight_capacity,
    void* packed_scales, size_t packed_scale_capacity,
    float* tensor_scale);

// A context owns one cuBLASLt handle and reusable activation/reduction/
// workspace buffers. The first forward binds it to one CUDA stream. This
// initial prototype deliberately rejects a different stream and is not
// re-entrant.
MAGE_NVFP4_API int mage_nvfp4_create_context(
    int cuda_device, void** context);
MAGE_NVFP4_API int mage_nvfp4_destroy_context(void* context);
MAGE_NVFP4_API size_t mage_nvfp4_context_reserved_bytes(
    const void* context);

// Run Y[M,N] = X[M,K] * W[N,K]^T + bias.
//
// X, packed weight/scales, tensor scale, bias, and output are CUDA pointers on
// context's device. X is contiguous BF16 [logical_m, K]. Output is contiguous
// BF16 [round_up(logical_m, 8), N]. The caller slices the padded rows.
// bias_bf16 may be null. stream is a cudaStream_t passed as uintptr_t.
MAGE_NVFP4_API int mage_nvfp4_linear_forward(
    void* context,
    const void* input_bf16,
    const void* packed_weight,
    size_t packed_weight_bytes,
    const void* packed_weight_scales,
    size_t packed_weight_scale_bytes,
    const void* weight_tensor_scale_f32,
    const void* bias_bf16,
    void* output_bf16,
    int logical_m,
    int k,
    int n,
    uintptr_t stream);

#ifdef __cplusplus
}
#endif