#include "nvfp4_linear.h"

#include <cuda_bf16.h>
#include <cuda_fp4.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <cublasLt.h>

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <limits>
#include <mutex>
#include <stdexcept>
#include <string>

namespace {

constexpr int kAbiVersion = 1;
constexpr int kFp4BlockElements = 16;
constexpr int kScaleTileOuter = 128;
constexpr int kScaleTileInner = 4;
constexpr int kWarpsPerQuantBlock = 8;
constexpr int kQuantThreads = 32 * kWarpsPerQuantBlock;
constexpr int kReduceThreads = 256;
constexpr int kReduceItemsPerThread = 4;
constexpr int kBiasThreads = 256;
constexpr float kFp4E2M1Max = 6.0f;
constexpr float kFp4TensorScaleMax = 448.0f;
constexpr float kTensorScaleDenominator =
    kFp4E2M1Max * kFp4TensorScaleMax;
constexpr size_t kWorkspaceBytes = 64ull * 1024ull * 1024ull;

thread_local std::string g_last_error;

[[noreturn]] void fail(const std::string& message) {
  throw std::runtime_error(message);
}

#define CUDA_CHECK(expr)                                                     \
  do {                                                                       \
    const cudaError_t status_ = (expr);                                      \
    if (status_ != cudaSuccess) {                                            \
      fail(std::string(#expr) + ": " + cudaGetErrorString(status_));         \
    }                                                                        \
  } while (0)

#define CUBLASLT_CHECK(expr)                                                 \
  do {                                                                       \
    const cublasStatus_t status_ = (expr);                                   \
    if (status_ != CUBLAS_STATUS_SUCCESS) {                                  \
      fail(std::string(#expr) + " failed with cuBLASLt status " +            \
           std::to_string(static_cast<int>(status_)));                       \
    }                                                                        \
  } while (0)

int round_up(int value, int multiple) {
  if (value <= 0 || multiple <= 0 ||
      value > std::numeric_limits<int>::max() - (multiple - 1)) {
    fail("invalid or overflowing round_up arguments");
  }
  return ((value + multiple - 1) / multiple) * multiple;
}

size_t checked_multiply(size_t left, size_t right, const char* label) {
  if (right != 0 && left > std::numeric_limits<size_t>::max() / right) {
    fail(std::string(label) + " size overflows size_t");
  }
  return left * right;
}

size_t round_up_divide(size_t dividend, size_t divisor) {
  return (dividend + divisor - 1) / divisor;
}

struct ScaleLayout {
  int inner_dim = 0;
  int outer_tiles = 0;
  size_t bytes = 0;
};

ScaleLayout make_scale_layout(int rows_k, int outer_columns) {
  if (rows_k <= 0 || rows_k % 16 != 0 || outer_columns <= 0) {
    fail("scale layout requires positive K divisible by 16 and positive M/N");
  }
  ScaleLayout layout;
  layout.inner_dim = round_up(rows_k / kFp4BlockElements, kScaleTileInner);
  layout.outer_tiles =
      (outer_columns + kScaleTileOuter - 1) / kScaleTileOuter;
  layout.bytes = checked_multiply(
      checked_multiply(static_cast<size_t>(layout.outer_tiles),
                       static_cast<size_t>(layout.inner_dim), "scale tensor"),
      static_cast<size_t>(kScaleTileOuter), "scale tensor");
  return layout;
}

size_t packed_weight_bytes_checked(int n, int k) {
  if (n <= 0 || k <= 0 || n % 8 != 0 || k % 32 != 0) {
    fail("NVFP4 weight requires N divisible by 8 and K divisible by 32");
  }
  return checked_multiply(static_cast<size_t>(n),
                          static_cast<size_t>(k), "packed weight") /
         2;
}

size_t host_scale_offset(int outer, int inner_scale, int scale_inner_dim) {
  const int outer_tile = outer / kScaleTileOuter;
  const int local_outer = outer % kScaleTileOuter;
  const int local_inner = inner_scale % kScaleTileInner;
  const int inner_tile_start = inner_scale - local_inner;
  const size_t tile_base =
      static_cast<size_t>(inner_tile_start +
                          outer_tile * scale_inner_dim) *
      kScaleTileOuter;
  return tile_base + static_cast<size_t>(local_outer % 32) * 16 +
         static_cast<size_t>(local_outer / 32) * 4 + local_inner;
}

__device__ __forceinline__ size_t device_scale_offset(
    int outer, int inner_scale, int scale_inner_dim) {
  const int outer_tile = outer / kScaleTileOuter;
  const int local_outer = outer % kScaleTileOuter;
  const int local_inner = inner_scale % kScaleTileInner;
  const int inner_tile_start = inner_scale - local_inner;
  const size_t tile_base =
      static_cast<size_t>(inner_tile_start +
                          outer_tile * scale_inner_dim) *
      kScaleTileOuter;
  return tile_base + static_cast<size_t>(local_outer % 32) * 16 +
         static_cast<size_t>(local_outer / 32) * 4 + local_inner;
}

float host_ue4m3_to_float(uint8_t raw) {
  const __half_raw half_raw = __nv_cvt_fp8_to_halfraw(raw, __NV_E4M3);
  return __half2float(static_cast<__half>(half_raw));
}

__device__ __forceinline__ float device_ue4m3_to_float(uint8_t raw) {
  const __half_raw half_raw = __nv_cvt_fp8_to_halfraw(raw, __NV_E4M3);
  return __half2float(static_cast<__half>(half_raw));
}

__device__ __forceinline__ float tensor_scale_from_amax(float amax) {
  return amax == 0.0f ? 1.0f : amax / kTensorScaleDenominator;
}

__device__ __forceinline__ float block_reduce_max(float value) {
  __shared__ float shared[kReduceThreads];
  shared[threadIdx.x] = value;
  __syncthreads();
  for (int offset = kReduceThreads / 2; offset > 0; offset >>= 1) {
    if (threadIdx.x < offset) {
      shared[threadIdx.x] =
          fmaxf(shared[threadIdx.x], shared[threadIdx.x + offset]);
    }
    __syncthreads();
  }
  return shared[0];
}

__global__ void reduce_abs_max_bf16(
    const __nv_bfloat16* source, float* block_maxima,
    size_t element_count) {
  const size_t block_start =
      static_cast<size_t>(blockIdx.x) * kReduceThreads *
      kReduceItemsPerThread;
  float local_max = 0.0f;
  for (int item = 0; item < kReduceItemsPerThread; ++item) {
    const size_t index =
        block_start + static_cast<size_t>(threadIdx.x) +
        static_cast<size_t>(item) * kReduceThreads;
    if (index < element_count) {
      local_max =
          fmaxf(local_max, fabsf(__bfloat162float(source[index])));
    }
  }
  const float block_max = block_reduce_max(local_max);
  if (threadIdx.x == 0) {
    block_maxima[blockIdx.x] = block_max;
  }
}

__global__ void reduce_max_float(
    const float* source, float* block_maxima, size_t element_count) {
  const size_t block_start =
      static_cast<size_t>(blockIdx.x) * kReduceThreads *
      kReduceItemsPerThread;
  float local_max = 0.0f;
  for (int item = 0; item < kReduceItemsPerThread; ++item) {
    const size_t index =
        block_start + static_cast<size_t>(threadIdx.x) +
        static_cast<size_t>(item) * kReduceThreads;
    if (index < element_count) {
      local_max = fmaxf(local_max, source[index]);
    }
  }
  const float block_max = block_reduce_max(local_max);
  if (threadIdx.x == 0) {
    block_maxima[blockIdx.x] = block_max;
  }
}

__global__ void finalize_activation_scale(
    const float* activation_amax, const float* weight_tensor_scale,
    float* activation_tensor_scale, float* fp4_alpha) {
  if (blockIdx.x == 0 && threadIdx.x == 0) {
    const float activation_scale =
        tensor_scale_from_amax(activation_amax[0]);
    activation_tensor_scale[0] = activation_scale;
    fp4_alpha[0] = weight_tensor_scale[0] * activation_scale;
  }
}

// logical_m may be smaller than padded_m. Padded rows are written as exact
// FP4 zero with zero block scales, so reusable scratch never exposes a prior
// call's tail.
__global__ void dynamic_quantize_activation(
    const __nv_bfloat16* source, uint8_t* destination_fp4,
    uint8_t* destination_scales,
    const float* activation_tensor_scale, int rows_k, int logical_m,
    int scale_inner_dim, uint64_t padded_scale_blocks) {
  const int lane = threadIdx.x & 31;
  const int warp_in_block = threadIdx.x >> 5;
  const uint64_t logical_block =
      static_cast<uint64_t>(blockIdx.x) * kWarpsPerQuantBlock +
      static_cast<uint64_t>(warp_in_block);
  if (logical_block >= padded_scale_blocks) {
    return;
  }

  const int blocks_per_column = rows_k / kFp4BlockElements;
  const int outer =
      static_cast<int>(logical_block / blocks_per_column);
  const int inner_scale = static_cast<int>(
      logical_block -
      static_cast<uint64_t>(outer) * blocks_per_column);
  const int first_row = inner_scale * kFp4BlockElements;
  const size_t column_base = static_cast<size_t>(outer) * rows_k;
  const bool valid_outer = outer < logical_m;
  const float tensor_scale = activation_tensor_scale[0];
  const float inverse_tensor_scale =
      tensor_scale == 0.0f ? 0.0f : 1.0f / tensor_scale;

  float value0 = 0.0f;
  float value1 = 0.0f;
  float magnitude = 0.0f;
  size_t destination_source_index = 0;
  if (lane < kFp4BlockElements / 2) {
    const int row0 = first_row + lane * 2;
    destination_source_index = column_base + row0;
    if (valid_outer) {
      value0 = __bfloat162float(source[destination_source_index]) *
               inverse_tensor_scale;
      value1 = __bfloat162float(source[destination_source_index + 1]) *
               inverse_tensor_scale;
      magnitude = fmaxf(fabsf(value0), fabsf(value1));
    }
  }
  for (int offset = 16; offset > 0; offset >>= 1) {
    magnitude =
        fmaxf(magnitude,
              __shfl_down_sync(0xffffffffU, magnitude, offset));
  }

  float rounded_scale = 0.0f;
  if (lane == 0) {
    const uint8_t scale_raw = __nv_cvt_float_to_fp8(
        magnitude / kFp4E2M1Max, __NV_SATFINITE, __NV_E4M3);
    rounded_scale = device_ue4m3_to_float(scale_raw);
    destination_scales[device_scale_offset(
        outer, inner_scale, scale_inner_dim)] = scale_raw;
  }
  rounded_scale = __shfl_sync(0xffffffffU, rounded_scale, 0);

  if (lane < kFp4BlockElements / 2) {
    const float inverse_scale =
        rounded_scale == 0.0f ? 0.0f : 1.0f / rounded_scale;
    destination_fp4[destination_source_index / 2] =
        __nv_cvt_float2_to_fp4x2(
            make_float2(value0 * inverse_scale,
                        value1 * inverse_scale),
            __NV_E2M1, cudaRoundNearest);
  }
}

__global__ void add_bias_bf16(
    __nv_bfloat16* output, const __nv_bfloat16* bias,
    size_t element_count, int n) {
  const size_t index =
      static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
  if (index < element_count) {
    const float value = __bfloat162float(output[index]);
    const float bias_value = __bfloat162float(bias[index % n]);
    output[index] = __float2bfloat16(value + bias_value);
  }
}

size_t reduction_block_count(size_t element_count) {
  return round_up_divide(
      element_count,
      static_cast<size_t>(kReduceThreads * kReduceItemsPerThread));
}

size_t reduction_scratch_elements(size_t element_count) {
  size_t max_blocks = 1;
  while (element_count > 1) {
    const size_t blocks = reduction_block_count(element_count);
    max_blocks = std::max(max_blocks, blocks);
    element_count = blocks;
  }
  return max_blocks;
}

void enqueue_activation_amax_reduce(
    const __nv_bfloat16* source, size_t element_count,
    float* scratch_a, float* scratch_b, float* destination_amax,
    int max_grid_x, cudaStream_t stream) {
  const float* current_source = nullptr;
  float* current_destination = scratch_a;
  size_t current_count = element_count;
  bool first_stage = true;
  while (true) {
    const size_t blocks = reduction_block_count(current_count);
    if (blocks == 0 || blocks > static_cast<size_t>(max_grid_x)) {
      fail("activation amax reduction exceeds the GPU grid limit");
    }
    if (first_stage) {
      reduce_abs_max_bf16<<<
          static_cast<unsigned int>(blocks), kReduceThreads, 0,
          stream>>>(source, current_destination, current_count);
    } else {
      reduce_max_float<<<
          static_cast<unsigned int>(blocks), kReduceThreads, 0,
          stream>>>(current_source, current_destination, current_count);
    }
    CUDA_CHECK(cudaPeekAtLastError());
    if (blocks == 1) {
      if (current_destination != destination_amax) {
        CUDA_CHECK(cudaMemcpyAsync(
            destination_amax, current_destination, sizeof(float),
            cudaMemcpyDeviceToDevice, stream));
      }
      return;
    }
    current_count = blocks;
    current_source = current_destination;
    current_destination =
        current_destination == scratch_a ? scratch_b : scratch_a;
    first_stage = false;
  }
}

struct DeviceAllocation {
  void* pointer = nullptr;
  size_t bytes = 0;

  ~DeviceAllocation() {
    if (pointer != nullptr) {
      cudaFree(pointer);
    }
  }

  DeviceAllocation() = default;
  DeviceAllocation(const DeviceAllocation&) = delete;
  DeviceAllocation& operator=(const DeviceAllocation&) = delete;
};

struct Context {
  int device = -1;
  int max_grid_x = 0;
  cublasLtHandle_t handle = nullptr;
  DeviceAllocation x_fp4;
  DeviceAllocation x_scales;
  DeviceAllocation reduce_a;
  DeviceAllocation reduce_b;
  DeviceAllocation activation_amax;
  DeviceAllocation activation_scale;
  DeviceAllocation fp4_alpha;
  DeviceAllocation fp4_beta;
  DeviceAllocation workspace;
  uintptr_t bound_stream = 0;
  bool stream_bound = false;
  std::mutex mutex;

  ~Context() {
    if (handle != nullptr) {
      cublasLtDestroy(handle);
    }
  }
};

void allocate_exact(DeviceAllocation* allocation, size_t bytes) {
  if (allocation->pointer != nullptr) {
    CUDA_CHECK(cudaFree(allocation->pointer));
    allocation->pointer = nullptr;
    allocation->bytes = 0;
  }
  if (bytes != 0) {
    CUDA_CHECK(cudaMalloc(&allocation->pointer, bytes));
    allocation->bytes = bytes;
  }
}

void ensure_capacity(
    Context* context, DeviceAllocation* allocation, size_t bytes,
    cudaStream_t stream) {
  if (allocation->bytes >= bytes) {
    return;
  }
  // A growth invalidates scratch pointers. Synchronize the one bound stream
  // before freeing; steady-state forwards do not synchronize.
  if (context->stream_bound) {
    CUDA_CHECK(cudaStreamSynchronize(stream));
  }
  allocate_exact(allocation, bytes);
}

void set_scale_mode(
    cublasLtMatmulDesc_t operation,
    cublasLtMatmulDescAttributes_t attribute) {
  const int32_t mode =
      CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      operation, attribute, &mode, sizeof(mode)));
}

void set_pointer_attribute(
    cublasLtMatmulDesc_t operation,
    cublasLtMatmulDescAttributes_t attribute, const void* pointer) {
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      operation, attribute, &pointer, sizeof(pointer)));
}

struct Descriptors {
  cublasLtMatmulDesc_t operation = nullptr;
  cublasLtMatrixLayout_t w = nullptr;
  cublasLtMatrixLayout_t x = nullptr;
  cublasLtMatrixLayout_t c = nullptr;
  cublasLtMatrixLayout_t d = nullptr;

  ~Descriptors() {
    if (d != nullptr) cublasLtMatrixLayoutDestroy(d);
    if (c != nullptr) cublasLtMatrixLayoutDestroy(c);
    if (x != nullptr) cublasLtMatrixLayoutDestroy(x);
    if (w != nullptr) cublasLtMatrixLayoutDestroy(w);
    if (operation != nullptr) cublasLtMatmulDescDestroy(operation);
  }
};

cublasLtMatmulAlgo_t choose_algo(
    Context* context, const Descriptors& descriptors,
    size_t* workspace_bytes) {
  cublasLtMatmulPreference_t preference = nullptr;
  CUBLASLT_CHECK(cublasLtMatmulPreferenceCreate(&preference));
  CUBLASLT_CHECK(cublasLtMatmulPreferenceSetAttribute(
      preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
      &kWorkspaceBytes, sizeof(kWorkspaceBytes)));

  cublasLtMatmulHeuristicResult_t candidates[16]{};
  int returned = 0;
  const cublasStatus_t status = cublasLtMatmulAlgoGetHeuristic(
      context->handle, descriptors.operation, descriptors.w,
      descriptors.x, descriptors.c, descriptors.d, preference, 16,
      candidates, &returned);
  cublasLtMatmulPreferenceDestroy(preference);
  if (status != CUBLAS_STATUS_SUCCESS || returned == 0) {
    fail("cuBLASLt returned no NVFP4 heuristic");
  }
  for (int index = 0; index < returned; ++index) {
    if (candidates[index].state == CUBLAS_STATUS_SUCCESS) {
      *workspace_bytes = candidates[index].workspaceSize;
      return candidates[index].algo;
    }
  }
  fail("all cuBLASLt NVFP4 heuristics were unsupported");
}

void validate_device_pointer(
    const void* pointer, int expected_device, const char* label) {
  if (pointer == nullptr) {
    fail(std::string(label) + " is null");
  }
  cudaPointerAttributes attributes{};
  CUDA_CHECK(cudaPointerGetAttributes(&attributes, pointer));
  if (attributes.type != cudaMemoryTypeDevice ||
      attributes.device != expected_device) {
    fail(std::string(label) + " is not a CUDA allocation on the context device");
  }
  if ((reinterpret_cast<uintptr_t>(pointer) & 0x0fU) != 0) {
    fail(std::string(label) + " is not at least 16-byte aligned");
  }
}

void forward_impl(
    Context* 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_value) {
  if (logical_m <= 0 || k <= 0 || n <= 0) {
    fail("M, K, and N must be positive");
  }
  const size_t expected_weight_bytes =
      packed_weight_bytes_checked(n, k);
  const size_t expected_scale_bytes = make_scale_layout(k, n).bytes;
  if (packed_weight_bytes != expected_weight_bytes ||
      packed_weight_scale_bytes != expected_scale_bytes) {
    fail("packed weight or scale buffer has the wrong byte size");
  }

  CUDA_CHECK(cudaSetDevice(context->device));
  cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_value);
  if (!context->stream_bound) {
    context->bound_stream = stream_value;
    context->stream_bound = true;
  } else if (context->bound_stream != stream_value) {
    fail("resident NVFP4 context is bound to a different CUDA stream");
  }
  cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
  CUDA_CHECK(cudaStreamIsCapturing(stream, &capture_status));
  if (capture_status != cudaStreamCaptureStatusNone) {
    fail("resident ctypes prototype does not support CUDA graph capture");
  }

  validate_device_pointer(input_bf16, context->device, "input");
  validate_device_pointer(packed_weight, context->device, "packed weight");
  validate_device_pointer(
      packed_weight_scales, context->device, "weight scales");
  validate_device_pointer(
      weight_tensor_scale_f32, context->device, "weight tensor scale");
  validate_device_pointer(output_bf16, context->device, "output");
  if (bias_bf16 != nullptr) {
    validate_device_pointer(bias_bf16, context->device, "bias");
  }

  const int padded_m = round_up(logical_m, 8);
  const size_t input_elements = checked_multiply(
      static_cast<size_t>(logical_m), static_cast<size_t>(k), "input");
  const size_t padded_input_elements = checked_multiply(
      static_cast<size_t>(padded_m), static_cast<size_t>(k),
      "padded input");
  const size_t output_elements = checked_multiply(
      static_cast<size_t>(padded_m), static_cast<size_t>(n), "output");
  const ScaleLayout x_scale_layout = make_scale_layout(k, padded_m);
  const size_t reduce_elements =
      reduction_scratch_elements(input_elements);

  ensure_capacity(
      context, &context->x_fp4, padded_input_elements / 2, stream);
  ensure_capacity(
      context, &context->x_scales, x_scale_layout.bytes, stream);
  ensure_capacity(
      context, &context->reduce_a,
      checked_multiply(reduce_elements, sizeof(float), "reduction"),
      stream);
  ensure_capacity(
      context, &context->reduce_b,
      checked_multiply(reduce_elements, sizeof(float), "reduction"),
      stream);

  // Clear the complete tiled scale allocation, including padding that the
  // logical quantizer does not address.
  CUDA_CHECK(cudaMemsetAsync(
      context->x_scales.pointer, 0, x_scale_layout.bytes, stream));
  enqueue_activation_amax_reduce(
      static_cast<const __nv_bfloat16*>(input_bf16), input_elements,
      static_cast<float*>(context->reduce_a.pointer),
      static_cast<float*>(context->reduce_b.pointer),
      static_cast<float*>(context->activation_amax.pointer),
      context->max_grid_x, stream);
  finalize_activation_scale<<<1, 1, 0, stream>>>(
      static_cast<const float*>(context->activation_amax.pointer),
      static_cast<const float*>(weight_tensor_scale_f32),
      static_cast<float*>(context->activation_scale.pointer),
      static_cast<float*>(context->fp4_alpha.pointer));
  CUDA_CHECK(cudaPeekAtLastError());

  const uint64_t padded_blocks =
      static_cast<uint64_t>(padded_m) *
      static_cast<uint64_t>(k / kFp4BlockElements);
  const uint64_t cuda_blocks =
      (padded_blocks + kWarpsPerQuantBlock - 1) /
      kWarpsPerQuantBlock;
  if (cuda_blocks == 0 ||
      cuda_blocks > static_cast<uint64_t>(context->max_grid_x)) {
    fail("activation quantizer exceeds the GPU grid limit");
  }
  dynamic_quantize_activation<<<
      static_cast<unsigned int>(cuda_blocks), kQuantThreads, 0,
      stream>>>(
      static_cast<const __nv_bfloat16*>(input_bf16),
      static_cast<uint8_t*>(context->x_fp4.pointer),
      static_cast<uint8_t*>(context->x_scales.pointer),
      static_cast<const float*>(context->activation_scale.pointer),
      k, logical_m, x_scale_layout.inner_dim, padded_blocks);
  CUDA_CHECK(cudaPeekAtLastError());

  Descriptors descriptors;
  CUBLASLT_CHECK(cublasLtMatmulDescCreate(
      &descriptors.operation, CUBLAS_COMPUTE_32F, CUDA_R_32F));
  const cublasOperation_t transpose_a = CUBLAS_OP_T;
  const cublasOperation_t transpose_b = CUBLAS_OP_N;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      descriptors.operation, CUBLASLT_MATMUL_DESC_TRANSA,
      &transpose_a, sizeof(transpose_a)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      descriptors.operation, CUBLASLT_MATMUL_DESC_TRANSB,
      &transpose_b, sizeof(transpose_b)));
  const cublasLtPointerMode_t pointer_mode =
      CUBLASLT_POINTER_MODE_DEVICE;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      descriptors.operation, CUBLASLT_MATMUL_DESC_POINTER_MODE,
      &pointer_mode, sizeof(pointer_mode)));
  set_scale_mode(
      descriptors.operation, CUBLASLT_MATMUL_DESC_A_SCALE_MODE);
  set_scale_mode(
      descriptors.operation, CUBLASLT_MATMUL_DESC_B_SCALE_MODE);
  // Refresh pointer attributes for every call. Module buffers may differ
  // between adjacent linears even though the shape is identical.
  set_pointer_attribute(
      descriptors.operation, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
      packed_weight_scales);
  set_pointer_attribute(
      descriptors.operation, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
      context->x_scales.pointer);

  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
      &descriptors.w, CUDA_R_4F_E2M1, k, n, k));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
      &descriptors.x, CUDA_R_4F_E2M1, k, padded_m, k));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
      &descriptors.c, CUDA_R_16BF, n, padded_m, n));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
      &descriptors.d, CUDA_R_16BF, n, padded_m, n));

  size_t selected_workspace_bytes = 0;
  const cublasLtMatmulAlgo_t algorithm =
      choose_algo(context, descriptors, &selected_workspace_bytes);
  if (selected_workspace_bytes > context->workspace.bytes) {
    fail("selected cuBLASLt algorithm exceeds the context workspace");
  }
  CUBLASLT_CHECK(cublasLtMatmul(
      context->handle, descriptors.operation,
      context->fp4_alpha.pointer,
      packed_weight, descriptors.w,
      context->x_fp4.pointer, descriptors.x,
      context->fp4_beta.pointer,
      output_bf16, descriptors.c,
      output_bf16, descriptors.d,
      &algorithm, context->workspace.pointer,
      selected_workspace_bytes, stream));

  if (bias_bf16 != nullptr) {
    const size_t blocks =
        round_up_divide(output_elements,
                        static_cast<size_t>(kBiasThreads));
    if (blocks > static_cast<size_t>(context->max_grid_x)) {
      fail("bias kernel exceeds the GPU grid limit");
    }
    add_bias_bf16<<<
        static_cast<unsigned int>(blocks), kBiasThreads, 0, stream>>>(
        static_cast<__nv_bfloat16*>(output_bf16),
        static_cast<const __nv_bfloat16*>(bias_bf16),
        output_elements, n);
    CUDA_CHECK(cudaPeekAtLastError());
  }
}

template <typename Function>
int guarded(Function&& function) {
  try {
    g_last_error.clear();
    function();
    return 0;
  } catch (const std::exception& error) {
    g_last_error = error.what();
    return 1;
  } catch (...) {
    g_last_error = "unknown native NVFP4 error";
    return 2;
  }
}

}  // namespace

extern "C" int mage_nvfp4_abi_version(void) {
  return kAbiVersion;
}

extern "C" const char* mage_nvfp4_last_error(void) {
  return g_last_error.c_str();
}

extern "C" size_t mage_nvfp4_packed_weight_bytes(int n, int k) {
  try {
    g_last_error.clear();
    return packed_weight_bytes_checked(n, k);
  } catch (const std::exception& error) {
    g_last_error = error.what();
    return 0;
  }
}

extern "C" size_t mage_nvfp4_weight_scale_bytes(int n, int k) {
  try {
    g_last_error.clear();
    packed_weight_bytes_checked(n, k);
    return make_scale_layout(k, n).bytes;
  } catch (const std::exception& error) {
    g_last_error = error.what();
    return 0;
  }
}

extern "C" 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) {
  return guarded([&]() {
    if (weight_bf16 == nullptr || packed_weight == nullptr ||
        packed_scales == nullptr || tensor_scale == nullptr) {
      fail("weight packer received a null pointer");
    }
    const size_t required_weight =
        packed_weight_bytes_checked(n, k);
    const ScaleLayout scale_layout = make_scale_layout(k, n);
    if (packed_weight_capacity != required_weight ||
        packed_scale_capacity != scale_layout.bytes) {
      fail("weight packer received an incorrectly sized destination");
    }

    const auto* source =
        static_cast<const __nv_bfloat16*>(weight_bf16);
    auto* destination = static_cast<uint8_t*>(packed_weight);
    auto* scales = static_cast<uint8_t*>(packed_scales);
    std::memset(destination, 0, required_weight);
    std::memset(scales, 0, scale_layout.bytes);
    const size_t elements =
        checked_multiply(static_cast<size_t>(n),
                         static_cast<size_t>(k), "weight");
    float global_amax = 0.0f;
    for (size_t index = 0; index < elements; ++index) {
      global_amax = std::max(
          global_amax,
          std::abs(__bfloat162float(source[index])));
    }
    *tensor_scale =
        global_amax == 0.0f
            ? 1.0f
            : global_amax / kTensorScaleDenominator;
    const float inverse_tensor_scale = 1.0f / *tensor_scale;

    // nn.Linear weight [N,K] is physical column-major KxN.
    for (int column = 0; column < n; ++column) {
      for (int block = 0; block < k / kFp4BlockElements; ++block) {
        const int first_row = block * kFp4BlockElements;
        float block_amax = 0.0f;
        for (int lane = 0; lane < kFp4BlockElements; ++lane) {
          const size_t index =
              static_cast<size_t>(first_row + lane) +
              static_cast<size_t>(column) * k;
          block_amax = std::max(
              block_amax,
              std::abs(__bfloat162float(source[index]) *
                       inverse_tensor_scale));
        }
        const uint8_t scale_raw = __nv_cvt_float_to_fp8(
            block_amax / kFp4E2M1Max,
            __NV_SATFINITE, __NV_E4M3);
        const float rounded_scale = host_ue4m3_to_float(scale_raw);
        scales[host_scale_offset(
            column, block, scale_layout.inner_dim)] = scale_raw;
        const float inverse_scale =
            rounded_scale == 0.0f ? 0.0f : 1.0f / rounded_scale;
        for (int pair = 0;
             pair < kFp4BlockElements / 2; ++pair) {
          const int row0 = first_row + pair * 2;
          const size_t source_index =
              static_cast<size_t>(row0) +
              static_cast<size_t>(column) * k;
          const float value0 =
              __bfloat162float(source[source_index]) *
              inverse_tensor_scale;
          const float value1 =
              __bfloat162float(source[source_index + 1]) *
              inverse_tensor_scale;
          destination[source_index / 2] =
              __nv_cvt_float2_to_fp4x2(
                  make_float2(value0 * inverse_scale,
                              value1 * inverse_scale),
                  __NV_E2M1, cudaRoundNearest);
        }
      }
    }
  });
}

extern "C" int mage_nvfp4_create_context(
    int cuda_device, void** context) {
  return guarded([&]() {
    if (context == nullptr) {
      fail("context output pointer is null");
    }
    *context = nullptr;
    int device_count = 0;
    CUDA_CHECK(cudaGetDeviceCount(&device_count));
    if (cuda_device < 0 || cuda_device >= device_count) {
      fail("invalid CUDA device index");
    }
    CUDA_CHECK(cudaSetDevice(cuda_device));
    cudaDeviceProp properties{};
    CUDA_CHECK(cudaGetDeviceProperties(&properties, cuda_device));
    if (properties.major != 12 || properties.minor != 0) {
      fail("resident NVFP4 prototype requires sm_120");
    }

    Context* created = new Context();
    try {
      created->device = cuda_device;
      created->max_grid_x = properties.maxGridSize[0];
      CUBLASLT_CHECK(cublasLtCreate(&created->handle));
      allocate_exact(&created->activation_amax, sizeof(float));
      allocate_exact(&created->activation_scale, sizeof(float));
      allocate_exact(&created->fp4_alpha, sizeof(float));
      allocate_exact(&created->fp4_beta, sizeof(float));
      allocate_exact(&created->workspace, kWorkspaceBytes);
      const float zero = 0.0f;
      CUDA_CHECK(cudaMemcpy(
          created->fp4_beta.pointer, &zero, sizeof(zero),
          cudaMemcpyHostToDevice));
      *context = created;
    } catch (...) {
      delete created;
      throw;
    }
  });
}

extern "C" int mage_nvfp4_destroy_context(void* context) {
  return guarded([&]() {
    if (context == nullptr) {
      return;
    }
    auto* typed = static_cast<Context*>(context);
    {
      std::lock_guard<std::mutex> lock(typed->mutex);
      CUDA_CHECK(cudaSetDevice(typed->device));
      if (typed->stream_bound) {
        CUDA_CHECK(cudaStreamSynchronize(
            reinterpret_cast<cudaStream_t>(typed->bound_stream)));
      }
    }
    delete typed;
  });
}

extern "C" size_t mage_nvfp4_context_reserved_bytes(
    const void* context) {
  if (context == nullptr) {
    return 0;
  }
  const auto* typed = static_cast<const Context*>(context);
  return typed->x_fp4.bytes + typed->x_scales.bytes +
         typed->reduce_a.bytes + typed->reduce_b.bytes +
         typed->activation_amax.bytes +
         typed->activation_scale.bytes +
         typed->fp4_alpha.bytes + typed->fp4_beta.bytes +
         typed->workspace.bytes;
}

extern "C" 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) {
  return guarded([&]() {
    if (context == nullptr) {
      fail("context is null");
    }
    auto* typed = static_cast<Context*>(context);
    std::lock_guard<std::mutex> lock(typed->mutex);
    forward_impl(
        typed, input_bf16, packed_weight, packed_weight_bytes,
        packed_weight_scales, packed_weight_scale_bytes,
        weight_tensor_scale_f32, bias_bf16, output_bf16,
        logical_m, k, n, stream);
  });
}