File size: 2,231 Bytes
ca8b3a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | #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
|