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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

#pragma once

#include <cuda_bf16.h>
#include <cuda_runtime_api.h>

#include <cute/tensor.hpp>
#include <cutlass/arch/arch.h>
#include <cutlass/cutlass.h>
#include <cutlass/epilogue/collective/collective_builder.hpp>
#include <cutlass/epilogue/fusion/operations.hpp>
#include <cutlass/gemm/collective/collective_builder.hpp>
#include <cutlass/gemm/device/gemm_universal_adapter.h>
#include <cutlass/gemm/group_array_problem_shape.hpp>
#include <cutlass/gemm/kernel/gemm_universal.hpp>

namespace mircuda::variable_grouped_fp4 {

using namespace cute;
using ProblemShape =
    cutlass::gemm::MoEProblemShape<Shape<int32_t, int32_t, int32_t>>;
using ElementType = cutlass::float_e2m1_t;
using ElementSF = cutlass::float_ue4m3_t;
using ElementA = cutlass::nv_float4_t<ElementType>;
using ElementB = cutlass::nv_float4_t<ElementType>;
using ElementC = cutlass::bfloat16_t;
using ElementAccumulator = float;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using ArchTag = cutlass::arch::Sm120;
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
using ClusterShape = Shape<_1, _1, _1>;
using MmaTileShape = Shape<_128, _128, _128>;
using Fusion = cutlass::epilogue::fusion::LinearCombination<
    ElementC, ElementAccumulator, ElementC, ElementAccumulator>;

static constexpr int AlignmentAB = 32;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;

using CollectiveEpilogue =
    typename cutlass::epilogue::collective::CollectiveBuilder<
        ArchTag, OperatorClass, MmaTileShape, ClusterShape,
        cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
        ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementC, LayoutC*,
        AlignmentC, cutlass::epilogue::collective::EpilogueScheduleAuto,
        Fusion>::CollectiveOp;
using CollectiveMainloop =
    typename cutlass::gemm::collective::CollectiveBuilder<
        ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentAB, ElementB,
        LayoutB*, AlignmentAB, ElementAccumulator, MmaTileShape, ClusterShape,
        cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
            sizeof(typename CollectiveEpilogue::SharedStorage))>,
        cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
    ProblemShape, CollectiveMainloop, CollectiveEpilogue>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using StrideA = GemmKernel::InternalStrideA;
using StrideB = GemmKernel::InternalStrideB;
using StrideC = GemmKernel::InternalStrideC;
using LayoutSFA = CollectiveMainloop::InternalLayoutSFA;
using LayoutSFB = CollectiveMainloop::InternalLayoutSFB;
using ScaleConfig = CollectiveMainloop::Sm1xxBlkScaledConfig;

struct Plan {
  int groups;
  int matrices;
  int max_rows;
  int n;
  int k;
  int device_id;
  int sm_count;
  cudaStream_t stream;
  int32_t* rows;
  ElementType** a_ptrs;
  ElementType** b_ptrs;
  ElementC** c_ptrs;
  ElementSF** a_scale_ptrs;
  ElementSF** b_scale_ptrs;
  float** alpha_ptrs;
  int64_t* a_strides;
  int64_t* b_strides;
  int64_t* c_strides;
  LayoutSFA* a_layouts;
  LayoutSFB* b_layouts;
  Gemm* gemm;
  void* workspace;
  size_t workspace_bytes;
};

GemmKernel::Arguments arguments(Plan* plan);
int allocate_plan(Plan* plan);
int create_plan(int groups, int matrices, int max_m, int n, int k,
                void* stream, Plan** output);
int execute_plan(Plan* plan, cudaStream_t stream);
void release_plan(Plan* plan);

}  // namespace mircuda::variable_grouped_fp4

extern "C" {
int mircuda_variable_grouped_fp4_create(int groups, int matrices, int max_m,
                                        int n, int k, void* stream,
                                        void** output);
int mircuda_variable_grouped_fp4_execute(
    void* plan, void* stream, const void* a, const void* a_scales,
    const void* b, const void* b_scales, const void* alphas,
    const unsigned int* indices, const unsigned int* rows,
    const unsigned int* offsets, void* c);
void mircuda_variable_grouped_fp4_destroy(void* plan);
int mircuda_paired_variable_grouped_fp4_create(
    int groups, int matrices, int max_m, int n, int k, void* stream,
    void** output);
int mircuda_paired_variable_grouped_fp4_execute(
    void* plan, 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);
void mircuda_paired_variable_grouped_fp4_destroy(void* plan);
}