@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 208 lines 6.05 KB
Raw
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
4 // clang-format off
5
6 #include "common.h"
7 #include "gemm.h"
8 #include "vec.h"
9
10 namespace {
11
12 template <typename scalar_t, typename packed_t>
13 void bmm_kernel_impl(
14 scalar_t* __restrict__ out,
15 const scalar_t* __restrict__ mat1,
16 const packed_t* __restrict__ mat2,
17 int64_t B,
18 int64_t M,
19 int64_t N,
20 int64_t K,
21 int64_t mat1_strideB,
22 int64_t mat1_strideM,
23 int64_t out_strideB,
24 int64_t out_strideM,
25 float scale = 0.f) {
26 constexpr int64_t BLOCK_M = block_size_m();
27 constexpr int64_t BLOCK_N = block_size_n();
28 const int64_t MB = div_up(M, BLOCK_M);
29 const int64_t NB = div_up(N, BLOCK_N);
30
31 // mat2 contiguous in [B, N, K]
32 int64_t mat2_strideB = N * K;
33 int64_t mat2_strideN = K;
34
35 const bool use_brgemm = can_use_brgemm<scalar_t>(M);
36
37 // parallel on [B, MB, NB]
38 at::parallel_for(0, B * MB * NB, 0, [&](int64_t begin, int64_t end) {
39 int64_t bs{0}, mb{0}, nb{0};
40 data_index_init(begin, bs, B, mb, MB, nb, NB);
41
42 // for brgemm, use float32 for accumulate
43 alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
44
45 for (int i = begin; i < end; ++i) {
46 UNUSED(i);
47 int mb_start = mb * BLOCK_M;
48 int mb_size = std::min(M - mb_start, BLOCK_M);
49 int nb_start = nb * BLOCK_N;
50 int nb_size = std::min(N - nb_start, BLOCK_N);
51
52 tinygemm_kernel<scalar_t>(
53 /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM,
54 /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */,
55 /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start,
56 /* Ctmp*/ Ctmp,
57 /* M */ mb_size,
58 /* N */ nb_size,
59 /* K */ K,
60 /* lda */ mat1_strideM,
61 /* ldb */ nb_size,
62 /* ldc */ out_strideM,
63 /* brg */ use_brgemm);
64
65 // move to the next index
66 data_index_step(bs, B, mb, MB, nb, NB);
67 }
68
69 if (use_brgemm) {
70 at::native::cpublas::brgemm_release();
71 }
72 });
73 }
74
75 template <>
76 void bmm_kernel_impl(
77 at::BFloat16* __restrict__ out,
78 const at::BFloat16* __restrict__ mat1,
79 const at::Float8_e4m3fn* __restrict__ mat2,
80 int64_t B,
81 int64_t M,
82 int64_t N,
83 int64_t K,
84 int64_t mat1_strideB,
85 int64_t mat1_strideM,
86 int64_t out_strideB,
87 int64_t out_strideM,
88 float scale) {
89 constexpr int64_t BLOCK_M = block_size_m();
90 constexpr int64_t BLOCK_N = block_size_n();
91 const int64_t MB = div_up(M, BLOCK_M);
92 const int64_t NB = div_up(N, BLOCK_N);
93
94 // mat2 contiguous in [B, N, K]
95 int64_t mat2_strideB = N * K;
96 int64_t mat2_strideN = K;
97
98 const bool use_brgemm = can_use_brgemm<at::BFloat16>(M);
99
100 // parallel on [B, MB, NB]
101 parallel_2d(B * MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
102 // for brgemm, use float32 for accumulate
103 alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
104 // for brgemm when mat2 is float8_e4m3
105 alignas(64) at::BFloat16 Btmp[BLOCK_N * BLOCK_K];
106
107 loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
108 int64_t bs = mb / MB;
109 int64_t mb_start = (mb % MB) * BLOCK_M;
110 int64_t mb_size = std::min(M - mb_start, BLOCK_M);
111 int64_t nb_start = nb * BLOCK_N;
112 int64_t nb_size = std::min(N - nb_start, BLOCK_N);
113
114 tinygemm_kernel(
115 /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM,
116 /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */,
117 /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start,
118 /* Btmp*/ Btmp,
119 /* Ctmp*/ Ctmp,
120 /*scale*/ scale,
121 /* M */ mb_size,
122 /* N */ nb_size,
123 /* K */ K,
124 /* lda */ mat1_strideM,
125 /* ldb */ nb_size,
126 /* ldc */ out_strideM,
127 /* brg */ use_brgemm);
128 });
129
130 if (use_brgemm) {
131 at::native::cpublas::brgemm_release();
132 }
133 });
134 }
135
136 } // anonymous namespace
137
138 // mat1 : [B, M, K]
139 // mat2 : [B, N, K] or [B, OC, IC]
140 // out : [B, M, N]
141 // scale: [] 0-dim tensor for per tensor quant
142 //
143 void bmm_cpu(
144 at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale) {
145 auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
146
147 // input and out could be non-contiguous
148 // weight needs to be contiguous in [OC, IC] order
149 CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
150 CHECK_LAST_DIM_CONTIGUOUS_INPUT(out);
151 CHECK_INPUT(mat2);
152 CHECK_DIM(3, out);
153 CHECK_DIM(3, mat1);
154 CHECK_DIM(3, mat2);
155
156 int64_t B = mat1.size(0);
157 int64_t M = mat1.size(1);
158 int64_t N = mat2.size(1);
159 int64_t K = mat1.size(2);
160
161 const bool use_fp8_w8a16 = scale.has_value();
162 TORCH_CHECK(N % 32 == 0, "tinygemm requires N to be 32x.");
163
164 int64_t mat1_strideB = mat1.stride(0);
165 int64_t mat1_strideM = mat1.stride(1);
166 int64_t out_strideB = out.stride(0);
167 int64_t out_strideM = out.stride(1);
168
169 // check shapes
170 TORCH_CHECK(mat2.size(0) == B && mat2.size(2) == K, "bmm: mat2 shape mismatch!");
171 TORCH_CHECK(out.size(0) == B && out.size(1) == M, "bmm: out shape mismatch!");
172 if (!use_fp8_w8a16) {
173 AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] {
174 bmm_kernel_impl<scalar_t, scalar_t>(
175 out.data_ptr<scalar_t>(),
176 mat1.data_ptr<scalar_t>(),
177 packed_w.data_ptr<scalar_t>(),
178 B,
179 M,
180 N,
181 K,
182 mat1_strideB,
183 mat1_strideM,
184 out_strideB,
185 out_strideM);
186 });
187 } else { // fp8 bmm
188 float scale_val = 0.f;
189
190 auto scale_tensor = scale.value();
191 TORCH_CHECK(scale_tensor.ndimension() == 0, "bmm: expect scale to be 0-dim tensor.");
192 scale_val = scale_tensor.item<float>();
193
194 bmm_kernel_impl<at::BFloat16, at::Float8_e4m3fn>(
195 out.data_ptr<at::BFloat16>(),
196 mat1.data_ptr<at::BFloat16>(),
197 packed_w.data_ptr<at::Float8_e4m3fn>(),
198 B,
199 M,
200 N,
201 K,
202 mat1_strideB,
203 mat1_strideM,
204 out_strideB,
205 out_strideM,
206 scale_val);
207 }
208 }