mircuda-sys 0.1.0

Private native CUDA integration for mircuda
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

#include "variable_grouped_fp4_sm120.cuh"

namespace mircuda::variable_grouped_fp4 {

__global__ void setup_pair(
    Plan plan,
    const ElementType* left_a, const ElementSF* left_a_scales,
    const ElementType* left_b, const ElementSF* left_b_scales,
    const float* left_alphas,
    const ElementType* right_a, const ElementSF* right_a_scales,
    const ElementType* right_b, const ElementSF* right_b_scales,
    const float* right_alphas,
    const unsigned int* indices, const unsigned int* rows,
    const unsigned int* offsets, ElementC* left_c, ElementC* right_c) {
  const int group = blockIdx.x * blockDim.x + threadIdx.x;
  if (group >= plan.groups) return;
  const int logical_groups = plan.groups / 2;
  const int logical = group % logical_groups;
  const bool right = group >= logical_groups;
  const unsigned int matrix = indices[logical];
  if (matrix >= static_cast<unsigned int>(plan.matrices)) return;
  const int row_count = min(static_cast<int>(rows[logical]), plan.max_rows);
  const size_t offset = offsets[logical];
  const int packed_k = plan.k / 2;
  const int a_scale_stride = ((plan.max_rows + 127) / 128) * 128 *
                             ((plan.k / 16 + 3) / 4) * 4;
  const int b_scale_stride = ((plan.n + 127) / 128) * 128 *
                             ((plan.k / 16 + 3) / 4) * 4;
  const ElementType* a = right ? right_a : left_a;
  const ElementSF* a_scales = right ? right_a_scales : left_a_scales;
  const ElementType* b = right ? right_b : left_b;
  const ElementSF* b_scales = right ? right_b_scales : left_b_scales;
  const float* alphas = right ? right_alphas : left_alphas;
  ElementC* c = right ? right_c : left_c;
  plan.rows[group] = row_count;
  plan.a_ptrs[group] = const_cast<ElementType*>(b) + matrix * plan.n * packed_k;
  plan.b_ptrs[group] = const_cast<ElementType*>(a) + offset * packed_k;
  plan.c_ptrs[group] = c + offset * plan.n;
  plan.a_scale_ptrs[group] =
      const_cast<ElementSF*>(b_scales) + matrix * b_scale_stride;
  plan.b_scale_ptrs[group] =
      const_cast<ElementSF*>(a_scales) + logical * a_scale_stride;
  plan.alpha_ptrs[group] = const_cast<float*>(alphas) + matrix;
  plan.a_layouts[group] = ScaleConfig::tile_atom_to_shape_SFA(
      cute::make_shape(plan.n, row_count, plan.k, 1));
  plan.b_layouts[group] = ScaleConfig::tile_atom_to_shape_SFB(
      cute::make_shape(plan.n, row_count, plan.k, 1));
}

}  // namespace mircuda::variable_grouped_fp4

extern "C" int mircuda_paired_variable_grouped_fp4_create(
    int groups, int matrices, int max_m, int n, int k, void* stream,
    void** output) {
  using namespace mircuda::variable_grouped_fp4;
  if (groups > INT_MAX / 2) return -1;
  Plan* plan = nullptr;
  const int status = create_plan(
      groups * 2, matrices, max_m, n, k, stream, &plan);
  if (status == 0) *output = plan;
  return status;
}

extern "C" int mircuda_paired_variable_grouped_fp4_execute(
    void* raw, void* stream,
    const void* left_a, const void* left_a_scales,
    const void* left_b, const void* left_b_scales, const void* left_alphas,
    const void* right_a, const void* right_a_scales,
    const void* right_b, const void* right_b_scales, const void* right_alphas,
    const unsigned int* indices, const unsigned int* rows,
    const unsigned int* offsets, void* left_c, void* right_c) {
  using namespace mircuda::variable_grouped_fp4;
  if (raw == nullptr || stream == nullptr) return -1;
  auto* plan = static_cast<Plan*>(raw);
  auto cuda_stream = static_cast<cudaStream_t>(stream);
  if (cuda_stream != plan->stream || plan->groups % 2 != 0) return -1;
  constexpr int threads = 256;
  const int blocks = (plan->groups + threads - 1) / threads;
  setup_pair<<<blocks, threads, 0, cuda_stream>>>(*plan,
      static_cast<const ElementType*>(left_a),
      static_cast<const ElementSF*>(left_a_scales),
      static_cast<const ElementType*>(left_b),
      static_cast<const ElementSF*>(left_b_scales),
      static_cast<const float*>(left_alphas),
      static_cast<const ElementType*>(right_a),
      static_cast<const ElementSF*>(right_a_scales),
      static_cast<const ElementType*>(right_b),
      static_cast<const ElementSF*>(right_b_scales),
      static_cast<const float*>(right_alphas), indices, rows, offsets,
      static_cast<ElementC*>(left_c), static_cast<ElementC*>(right_c));
  const int status = static_cast<int>(cudaPeekAtLastError());
  return status == 0 ? execute_plan(plan, cuda_stream) : status;
}

extern "C" void mircuda_paired_variable_grouped_fp4_destroy(void* raw) {
  auto* plan = static_cast<mircuda::variable_grouped_fp4::Plan*>(raw);
  if (plan == nullptr) return;
  mircuda::variable_grouped_fp4::release_plan(plan);
  delete plan;
}