A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Transform] Deterministic Hadacore Transforms (#24106)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Kyle Sayers committed
Sep 15, 2025 at 19:59 UTC
a0b26701c91a386e38105816057c1b3e6c1d080f
10 files changed
+979
-43
CMakeLists.txt
+11
@@ -783,6 +783,17 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
783
endif()
784
endif()
785
786
+ # Hadacore kernels
787
+ cuda_archs_loose_intersection(HADACORE_ARCHS "8.0;8.9;9.0" "${CUDA_ARCHS}")
788
+ if(HADACORE_ARCHS)
789
+ set(SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
790
+ set_gencode_flags_for_srcs(
791
+ SRCS "${SRCS}"
792
+ CUDA_ARCHS "${HADACORE_ARCHS}")
793
+ list(APPEND VLLM_EXT_SRC "${SRCS}")
794
+ message(STATUS "Building hadacore")
795
+ endif()
796
+
797
# if CUDA endif
798
endif()
799
csrc/ops.h
+2
@@ -347,6 +347,8 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
347
int64_t open_mem_handle(torch::Tensor& mem_handle);
348
void free_shared_buffer(int64_t buffer);
349
350
+torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace);
351
+
352
#ifdef USE_ROCM
353
fptr_t init_custom_qr(int64_t rank, int64_t world_size,
354
std::optional<int64_t> qr_max_size = std::nullopt);
csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu
new
+817
@@ -0,0 +1,817 @@
1
+// clang-format off
2
+// Adapted from: https://github.com/meta-pytorch/applied-ai/blob/main/kernels/cuda/inference/hadamard_transform/hadamard_transform_cuda.cu
3
+
4
+/***********
5
+Copyright 2024 Meta
6
+
7
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
8
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
9
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
10
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
11
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
12
+***********/
13
+
14
+#include <torch/all.h>
15
+#include <stdint.h>
16
+#include <cuda_runtime.h>
17
+#include <mma.h>
18
+#include <cuda/annotated_ptr>
19
+#include <c10/cuda/CUDAException.h>
20
+
21
+#include <ATen/cuda/CUDAContext.h>
22
+#include <c10/cuda/CUDAGuard.h>
23
+
24
+#include "core/registration.h"
25
+#include "dispatch_utils.h"
26
+
27
+namespace hadacore {
28
+
29
+#ifndef __CUDACC__
30
+#define __launch_bounds__(x,y)
31
+#endif
32
+
33
+#define MAX_WARPS_PER_SM 48
34
+
35
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
36
+
37
+using b16 = uint16_t;
38
+using b32 = uint32_t;
39
+
40
+constexpr int launch_configs_big[7][3] = {
41
+ // default
42
+ {2, 1, 24},
43
+ {2, 2, 16},
44
+ {2, 4, 8},
45
+ {2, 8, 4},
46
+ {2, 16, 3},
47
+ {4, 16, 2},
48
+ {8, 16, 1}
49
+ // // extra coalescing
50
+ // {2, 1, 24},
51
+ // {2, 2, 16},
52
+ // {2, 4, 8},
53
+ // {2, 8, 4},
54
+ // {4, 8, 3},
55
+ // {8, 8, 2},
56
+ // {16, 8, 1}
57
+ // // less coalescing
58
+ // {2, 1, 24},
59
+ // {2, 2, 16},
60
+ // {2, 4, 8},
61
+ // {2, 8, 4},
62
+ // {1, 32, 1},
63
+ // {2, 32, 1},
64
+ // {4, 32, 1}
65
+};
66
+
67
+// a 4x2, b 2x2, c 2x2
68
+template <torch::ScalarType dtype>
69
+__device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32& c0, b32& c1){
70
+ static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16);
71
+ // d, a, b, c
72
+ b32 zero = 0;
73
+ if constexpr(dtype == torch::ScalarType::Half) {
74
+ asm (
75
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
76
+ "{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n\t"
77
+ : "=r"(c0), "=r"(c1) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(zero), "r"(zero)
78
+ );
79
+ } else {
80
+ b32 temp0, temp1, temp2, temp3;
81
+ asm (
82
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
83
+ "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n\t"
84
+ : "=r"(temp0), "=r"(temp1), "=r"(temp2), "=r"(temp3) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(zero), "r"(zero), "r"(zero), "r"(zero)
85
+ );
86
+ asm ("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c0) : "r"(temp1), "r"(temp0));
87
+ asm ("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c1) : "r"(temp3), "r"(temp2));
88
+ }
89
+}
90
+
91
+// a 4x2, b 4x2, c 4x2
92
+template <torch::ScalarType dtype>
93
+__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32 b2, b32 b3, b32& c0, b32& c1, b32& c2, b32& c3){
94
+ mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b0, b1, c0, c1);
95
+ mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b2, b3, c2, c3);
96
+}
97
+
98
+__device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(b32& a0) {
99
+ asm (
100
+ "movmatrix.sync.aligned.m8n8.trans.b16 "
101
+ "%0, %1;\n\t"
102
+ : "=r"(a0) : "r"(a0)
103
+ );
104
+}
105
+
106
+#define p_p(i) ((val_1p[i] & 0x0000FFFF) | val_1p[i] << 16)
107
+#define p_n(i) ((val_1p[i] & 0x0000FFFF) | val_1n[i] << 16)
108
+#define n_p(i) ((val_1n[i] & 0x0000FFFF) | val_1p[i] << 16)
109
+#define n_n(i) ((val_1n[i] & 0x0000FFFF) | val_1n[i] << 16)
110
+
111
+template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::ScalarType dtype>
112
+__global__ void __launch_bounds__(32 * warps_per_block, blocks_per_sm)
113
+// a is column major, b is row major
114
+hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
115
+ static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
116
+
117
+ b32 b_frag_all[num_chunks][4]; // for all chunks, holds matrix fragment (which takes 4 regs of b16x2 * 32 threads)
118
+
119
+ int64_t blockid = blockIdx.x * warps_per_block + threadIdx.x / 32;
120
+ int64_t threadid = threadIdx.x % 32;
121
+ extern __shared__ b32 bfrag_arr[]; // num_chunks * warps_per_block * 128
122
+ int64_t real_num_chunks = ((blockid + 1) * num_chunks) > total_num_chunks ? (total_num_chunks - (blockid * num_chunks)) : num_chunks;
123
+ int64_t diff_num_chunks = real_num_chunks - num_chunks;
124
+
125
+ b32* a_start_ptr = (b32*) (a + blockid * num_chunks * 256); // offset a to where this warp starts
126
+ b32* out_start_ptr = (b32*) (out + blockid * num_chunks * 256);
127
+ b32* a_ptr = a_start_ptr + threadid * 4;
128
+ b32* b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * 128 + threadid * 4;
129
+
130
+ #if (__CUDA_ARCH__ < 900) // SM80, SM89
131
+ uint64_t cache_policy;
132
+ asm volatile(
133
+ "createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n"
134
+ : "=l"(cache_policy)
135
+ );
136
+ #endif
137
+
138
+ #pragma unroll
139
+ for (int64_t k = 0; k < num_chunks; k++) {
140
+ size_t shared_ptr = __cvta_generic_to_shared(b_frag_ptr);
141
+ #if (__CUDA_ARCH__ >= 900) // SM90
142
+ asm volatile(
143
+ "cp.async.cg.shared.global [%0], [%1], 16;\n"
144
+ "cp.async.commit_group;\n"
145
+ :: "l"(shared_ptr), "l"(a_ptr)
146
+ );
147
+ #else // SM80, SM89
148
+ asm volatile(
149
+ "cp.async.cg.shared.global.L2::cache_hint.L2::256B [%0], [%1], 16, %2;\n"
150
+ "cp.async.commit_group;\n"
151
+ :: "l"(shared_ptr), "l"(a_ptr), "l"(cache_policy)
152
+ );
153
+ #endif
154
+
155
+ a_ptr += 128;
156
+ b_frag_ptr += 128;
157
+ }
158
+
159
+ // generate hadamard 16x16 (up to 2 of them)
160
+ constexpr b16 fp16_1p[4] = {0b0011100110101000, 0b0011100000000000, 0b0011010110101000, 0b0011010000000000};
161
+ constexpr b16 fp16_1n[4] = {0b1011100110101000, 0b1011100000000000, 0b1011010110101000, 0b1011010000000000};
162
+ constexpr b16 bf16_1p[4] = {0b0011111100110101, 0b0011111100000000, 0b0011111010110101, 0b0011111010000000};
163
+ constexpr b16 bf16_1n[4] = {0b1011111100110101, 0b1011111100000000, 0b1011111010110101, 0b1011111010000000};
164
+
165
+ #define val_type_1p(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
166
+ #define val_type_1n(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
167
+ constexpr b16 val_1p[4] = {val_type_1p(0), val_type_1p(1), val_type_1p(2), val_type_1p(3)};
168
+ constexpr b16 val_1n[4] = {val_type_1n(0), val_type_1n(1), val_type_1n(2), val_type_1n(3)};
169
+
170
+ constexpr b32 p_p[4] = {p_p(0), p_p(1), p_p(2), p_p(3)};
171
+ constexpr b32 p_n[4] = {p_n(0), p_n(1), p_n(2), p_n(3)};
172
+ constexpr b32 n_p[4] = {n_p(0), n_p(1), n_p(2), n_p(3)};
173
+ constexpr b32 n_n[4] = {n_n(0), n_n(1), n_n(2), n_n(3)};
174
+ const b32 had_16_p1[4][4] = {
175
+ {
176
+ 0b10001000010001000010001000010001,
177
+ 0b00000000000000000000000000000000,
178
+ 0b00000000000000000000000000000000,
179
+ 0b10001000010001000010001000010001
180
+ },
181
+ {
182
+ 0b11001100100010000011001100100010,
183
+ 0b00000000000000000000000000000000,
184
+ 0b00000000000000000000000000000000,
185
+ 0b11001100100010000011001100100010
186
+ },
187
+ {
188
+ 0b11111111101010101100110010011001,
189
+ 0b00000000000000000000000000000000,
190
+ 0b00000000000000000000000000000000,
191
+ 0b11111111101010101100110010011001
192
+ },
193
+ {
194
+ 0b11111111101010101100110010011001,
195
+ 0b11111111101010101100110010011001,
196
+ 0b11111111101010101100110010011001,
197
+ 0b00000000010101010011001101100110
198
+ }
199
+ };
200
+ const b32 had_16_p2[4][4] = {
201
+ {
202
+ 0b10000000010000000010000000010000,
203
+ 0b00000000000000000000000000000000,
204
+ 0b00000000000000000000000000000000,
205
+ 0b10000000010000000010000000010000
206
+ },
207
+ {
208
+ 0b11000000100001000011000000100001,
209
+ 0b00000000000000000000000000000000,
210
+ 0b00000000000000000000000000000000,
211
+ 0b11000000100001000011000000100001
212
+ },
213
+ {
214
+ 0b11110000101001011100001110010110,
215
+ 0b00000000000000000000000000000000,
216
+ 0b00000000000000000000000000000000,
217
+ 0b11110000101001011100001110010110
218
+ },
219
+ {
220
+ 0b11110000101001011100001110010110,
221
+ 0b11110000101001011100001110010110,
222
+ 0b11110000101001011100001110010110,
223
+ 0b00001111010110100011110001101001
224
+ }
225
+ };
226
+ const b32 had_16_mask[3][4] = {
227
+ {
228
+ 0b10001000010001000010001000010001,
229
+ 0b00000000000000000000000000000000,
230
+ 0b00000000000000000000000000000000,
231
+ 0b10001000010001000010001000010001
232
+ },
233
+ {
234
+ 0b11001100110011000011001100110011,
235
+ 0b00000000000000000000000000000000,
236
+ 0b00000000000000000000000000000000,
237
+ 0b11001100110011000011001100110011
238
+ },
239
+ {
240
+ 0b11111111111111111111111111111111,
241
+ 0b00000000000000000000000000000000,
242
+ 0b00000000000000000000000000000000,
243
+ 0b11111111111111111111111111111111
244
+ }
245
+ };
246
+ b32 had_frag[8];
247
+ #pragma unroll
248
+ for (int64_t i = 0; i < 2; i++) {
249
+ int64_t c_log_h = (i == 0) ? MIN(4, log_had_size) : log_had_size % 4;
250
+ #pragma unroll
251
+ for (int64_t j = 0; j < 4; j++) {
252
+ if (c_log_h < 4) {
253
+ bool mask = had_16_mask[c_log_h - 1][j] & (1 << (31 - threadid));
254
+ if (!mask) {
255
+ had_frag[i * 4 + j] = 0;
256
+ continue;
257
+ }
258
+ }
259
+ bool pred1 = had_16_p1[c_log_h - 1][j] & (1 << (31 - threadid));
260
+ bool pred2 = had_16_p2[c_log_h - 1][j] & (1 << (31 - threadid));
261
+ b32 val = pred1 ? (pred2 ? p_p[c_log_h - 1] : p_n[c_log_h - 1]) : (pred2 ? n_p[c_log_h - 1] : n_n[c_log_h - 1]);
262
+ had_frag[i * 4 + j] = val;
263
+ }
264
+ if constexpr(log_had_size <= 4 || log_had_size % 4 == 0) break;
265
+ }
266
+
267
+ // log had size above 8, only used for above 2^8 = 256 size
268
+ constexpr int64_t part8_log_had_size = log_had_size - 8;
269
+
270
+ b32* a_chunk_ptr = a_start_ptr; // first chunk starts at this warp's data starts
271
+ b32* out_chunk_ptr = out_start_ptr;
272
+
273
+ #pragma unroll
274
+ for (int64_t l = 0; l < 2; l++) {
275
+ if constexpr(log_had_size <= 8) { // l == 0 guaranteed, redundant simplified version of else body, to help compiler warnings
276
+ b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * 128;
277
+ } else {
278
+ b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * (l == 0 ? 128 : (128 >> part8_log_had_size));
279
+ }
280
+
281
+ if (l == 1) {
282
+ if constexpr(log_had_size > 8) {
283
+ __syncthreads(); // sync between first and second iterations if above size 256
284
+
285
+ if constexpr(log_had_size >= 12) {
286
+ // sizes 4k and above
287
+
288
+ // a + threadblock offset + warp offset
289
+ // can then index into all chunks owned by this warp
290
+ b32* store = bfrag_arr + (128 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block));
291
+
292
+ #pragma unroll
293
+ for (int64_t j = 0; j < 4; j++) {
294
+ #pragma unroll
295
+ for (int64_t k = 0; k < num_chunks; k++) {
296
+ // here, j represents register, and k represents 8-offset/chunk
297
+ uint64_t real_chunk_num = (num_chunks - (threadid % num_chunks) + k) % num_chunks; // chunk at which you have target thread #'s data
298
+
299
+ int64_t real_thread_id = (threadid / num_chunks) * num_chunks + k; // target thread #
300
+ int64_t chunk_idx = 128 * real_chunk_num; // index due to fetching from another chunk (chunk in which this thread has the target thread's original data)
301
+ int64_t thread_group_idx = (real_thread_id / 4) * 16; // index due to fetching from another group of num_chunk threads (since shuffle is between num_chunk threads)
302
+ int64_t thread_idx = (real_thread_id % 4) * 2; // index due to original thread's position within the group of num_chunk threads
303
+ int64_t reg_idx = (j / 2) * 8 + (j % 2); // index due to target register
304
+ int64_t idx = chunk_idx + thread_group_idx + thread_idx + reg_idx; // final index
305
+
306
+ // fix idx for majorness
307
+ int64_t rowidx = idx % (1 << part8_log_had_size);
308
+ int64_t colidx = idx >> part8_log_had_size;
309
+
310
+ // store[rowidx * 128 + colidx] = data;
311
+ b32 data = store[rowidx * 128 + colidx];
312
+
313
+ // compiler generates excessive instructions, so we manually do the if statement
314
+ #pragma unroll
315
+ for (uint64_t i = 0; i < num_chunks; i++) {
316
+ asm volatile (
317
+ "{\n\t"
318
+ " .reg .pred p0;\n\t"
319
+ " setp.eq.s64 p0, %1, %2;\n\t"
320
+ " @p0 mov.b32 %0, %3;\n\t"
321
+ "}\n\t"
322
+ : "+r"(b_frag_all[i][j]) // Output operand %0
323
+ : "l"(real_chunk_num), "l"(i), "r"(data) // Input operands %1, %2, %3
324
+ );
325
+ }
326
+ }
327
+ }
328
+
329
+ #pragma unroll
330
+ for (int64_t j = 0; j < 4; j++) {
331
+ #pragma unroll
332
+ for (int64_t k = 1; k < num_chunks; k++) {
333
+ int64_t threadid_contig = threadid % num_chunks;
334
+ int64_t threadid_mul = threadid / num_chunks;
335
+ int64_t threadid2 = (threadid_contig + num_chunks - k) % num_chunks + threadid_mul * num_chunks; // thread to give your data to
336
+ b_frag_all[k][j] = __shfl_sync(0xFFFFFFFF, b_frag_all[k][j], threadid2);
337
+ }
338
+ }
339
+ }
340
+ }
341
+ }
342
+
343
+ #pragma unroll
344
+ for (int64_t k = 0; k < num_chunks; k++) {
345
+ if constexpr(enable_mask) {
346
+ if (k >= real_num_chunks)
347
+ break;
348
+ }
349
+ if (l == 0) {
350
+ // bad fix for k not being recognized as a constexpr by compiler
351
+ // asm("cp.async.wait_group %0;\n" :: "n"(num_chunks - k - 1));
352
+ #define SWITCH_WAIT_ASYNC_LOAD_GROUP(i) case i: asm volatile("cp.async.wait_group %0;\n" :: "n"(num_chunks - i - 1)); break;
353
+ if constexpr(enable_mask) {
354
+ switch(k + diff_num_chunks) {
355
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(0)
356
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(1)
357
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(2)
358
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(3)
359
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(4)
360
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(5)
361
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(6)
362
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(7)
363
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(8)
364
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(9)
365
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(10)
366
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(11)
367
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(12)
368
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(13)
369
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(14)
370
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(15)
371
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(16)
372
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(17)
373
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(18)
374
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(19)
375
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(20)
376
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(21)
377
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(22)
378
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(23)
379
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(24)
380
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(25)
381
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(26)
382
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(27)
383
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(28)
384
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(29)
385
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(30)
386
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(31)
387
+ }
388
+ } else {
389
+ switch(k) {
390
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(0)
391
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(1)
392
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(2)
393
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(3)
394
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(4)
395
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(5)
396
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(6)
397
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(7)
398
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(8)
399
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(9)
400
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(10)
401
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(11)
402
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(12)
403
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(13)
404
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(14)
405
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(15)
406
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(16)
407
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(17)
408
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(18)
409
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(19)
410
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(20)
411
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(21)
412
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(22)
413
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(23)
414
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(24)
415
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(25)
416
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(26)
417
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(27)
418
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(28)
419
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(29)
420
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(30)
421
+ SWITCH_WAIT_ASYNC_LOAD_GROUP(31)
422
+ }
423
+ }
424
+ }
425
+
426
+ if (l == 0) {
427
+ // loading for the first iteration
428
+
429
+ // thread 0 loads [t0r0, t16r1, t0r2, t16r3]
430
+ // thread 16 loads [t0r1, t16r0, t0r3, t16r2]
431
+ // allows full coalescing, same for t1/t17, t2/t18, etc.
432
+ #pragma unroll
433
+ for (int64_t j = 0; j < 4; j++) {
434
+ int64_t reg = ((threadid & 16) == 0) ? j : (j / 2 * 2 + (1 - j % 2));
435
+ int64_t real_thread_id = (reg == 0 || reg == 2) ? threadid : (threadid ^ 16);
436
+ int64_t real_row = real_thread_id % 4;
437
+ int64_t real_col = real_thread_id / 4;
438
+ b_frag_all[k][j] = b_frag_ptr[(real_row + (reg % 2) * 4) + (real_col + (j / 2) * 8) * 8];
439
+ }
440
+
441
+ // for t16 swap r0/r1 and r2/r3 to have [t16r0, t0r1, t16r2, t0r3]
442
+ // so registers are in right order, same for t17, t18, etc.
443
+ if ((threadid & 16) != 0) {
444
+ b32 temp = b_frag_all[k][0];
445
+ b_frag_all[k][0] = b_frag_all[k][1];
446
+ b_frag_all[k][1] = temp;
447
+
448
+ temp = b_frag_all[k][2];
449
+ b_frag_all[k][2] = b_frag_all[k][3];
450
+ b_frag_all[k][3] = temp;
451
+ }
452
+
453
+ // t0 and t16 swap r1 and r3 to have their own data,
454
+ // same for t1/t17, t2/18, etc.
455
+ #pragma unroll
456
+ for (int64_t j = 1; j < 4; j += 2) {
457
+ b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], 16);
458
+ }
459
+ } else if constexpr(log_had_size > 8) { // condition is redundant to help compiler warnings
460
+ if constexpr(log_had_size < 12) {
461
+ // sizes 512, 1k, and 2k
462
+
463
+ // for 512:
464
+ // thread 0 loads [t0r0, t0r1, t16r2, t16r3]
465
+ // thread 16 loads [t0r2, t0r3, t16r0, t16r1]
466
+ // same for t1/t17, t2/t18, etc.
467
+ // for 1k and 2k:
468
+ // thread 0 loads [t0r0, t0r1, t1r2, t1r3]
469
+ // thread 1 loads [t0r2, t0r3, t1r0, t1r1]
470
+ // same for t2/t3, t4/t5, etc.
471
+ // allows full coalescing for 512 and 1k, 16x coalescing for 2k
472
+ constexpr int64_t xor_val = log_had_size == 9 ? 16 : 1;
473
+
474
+ #pragma unroll
475
+ for (int64_t j = 0; j < 4; j++) {
476
+ int64_t reg = ((threadid & xor_val) == 0) ? j : (j + 2) % 4;
477
+ int64_t real_thread_id = reg < 2 ? threadid : (threadid ^ xor_val);
478
+ int64_t idx = (real_thread_id / 4 * 16) + (real_thread_id % 4 * 2) + (reg / 2 * 8) + (reg % 2);
479
+ int64_t rowidx = idx % (1 << part8_log_had_size);
480
+ int64_t colidx = idx >> part8_log_had_size;
481
+ b_frag_all[k][j] = b_frag_ptr[rowidx * 128 + colidx];
482
+ }
483
+
484
+ if ((threadid & xor_val) != 0) {
485
+ b32 temp = b_frag_all[k][0];
486
+ b_frag_all[k][0] = b_frag_all[k][2];
487
+ b_frag_all[k][2] = temp;
488
+
489
+ temp = b_frag_all[k][1];
490
+ b_frag_all[k][1] = b_frag_all[k][3];
491
+ b_frag_all[k][3] = temp;
492
+ }
493
+
494
+ #pragma unroll
495
+ for (int64_t j = 2; j < 4; j++) {
496
+ b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], xor_val);
497
+ }
498
+ }
499
+ }
500
+
501
+ if (l == 1) {
502
+ // for second iteration, we load 2 consecutive b16s (1 b32) per register,
503
+ // but tensor core register layout requires 2 b16s that are in the
504
+ // same column/consecutive rows to be in the same register, so do the swap
505
+ b32 f0 = ((b_frag_all[k][1] & 0xFFFF) << 16) | (b_frag_all[k][0] & 0xFFFF);
506
+ b32 f1 = ((b_frag_all[k][3] & 0xFFFF) << 16) | (b_frag_all[k][2] & 0xFFFF);
507
+ b32 f2 = (b_frag_all[k][1] & 0xFFFF0000) | (b_frag_all[k][0] >> 16);
508
+ b32 f3 = (b_frag_all[k][3] & 0xFFFF0000) | (b_frag_all[k][2] >> 16);
509
+ b_frag_all[k][0] = f0;
510
+ b_frag_all[k][1] = f1;
511
+ b_frag_all[k][2] = f2;
512
+ b_frag_all[k][3] = f3;
513
+ }
514
+
515
+ #pragma unroll
516
+ for(int64_t i = 0, remaining_log_had_size = log_had_size - l * 8; i < 2 && remaining_log_had_size > 0; i++) {
517
+ int64_t had_off = ((remaining_log_had_size < 4) && !(log_had_size <= 4 || log_had_size % 4 == 0)) ? 4 : 0;
518
+ mma_m16_n16_k16_b16_b16_b16_noacc<dtype>(had_frag[had_off + 0], had_frag[had_off + 1], had_frag[had_off + 2], had_frag[had_off + 3], b_frag_all[k][0], b_frag_all[k][1], b_frag_all[k][2], b_frag_all[k][3], b_frag_all[k][0], b_frag_all[k][1], b_frag_all[k][2], b_frag_all[k][3]);
519
+
520
+ remaining_log_had_size -= 4;
521
+ if (remaining_log_had_size <= 0 && i == 0) {
522
+ // TODO: consider different storing so no need for transpose
523
+ matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][0]);
524
+ matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][1]);
525
+ matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][2]);
526
+ matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][3]);
527
+ } else {
528
+ // swap and use output directly as b_frag for next iteration as an actually free transpose
529
+ b32 temp = b_frag_all[k][1];
530
+ b_frag_all[k][1] = b_frag_all[k][2];
531
+ b_frag_all[k][2] = temp;
532
+ }
533
+ }
534
+
535
+ if (l == 1) {
536
+ // invert swap from above for second iteration
537
+ b32 f0 = ((b_frag_all[k][2] & 0xFFFF) << 16) | (b_frag_all[k][0] & 0xFFFF);
538
+ b32 f1 = (b_frag_all[k][2] & 0xFFFF0000) | (b_frag_all[k][0] >> 16);
539
+ b32 f2 = ((b_frag_all[k][3] & 0xFFFF) << 16) | (b_frag_all[k][1] & 0xFFFF);
540
+ b32 f3 = (b_frag_all[k][3] & 0xFFFF0000) | (b_frag_all[k][1] >> 16);
541
+ b_frag_all[k][0] = f0;
542
+ b_frag_all[k][1] = f1;
543
+ b_frag_all[k][2] = f2;
544
+ b_frag_all[k][3] = f3;
545
+ }
546
+
547
+ if (l == 0) {
548
+ // inverse of coalesced load for first iteration to store result
549
+ #pragma unroll
550
+ for (int64_t j = 1; j < 4; j += 2) {
551
+ b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], 16);
552
+ }
553
+
554
+ if ((threadid & 16) != 0) {
555
+ b32 temp = b_frag_all[k][0];
556
+ b_frag_all[k][0] = b_frag_all[k][1];
557
+ b_frag_all[k][1] = temp;
558
+
559
+ temp = b_frag_all[k][2];
560
+ b_frag_all[k][2] = b_frag_all[k][3];
561
+ b_frag_all[k][3] = temp;
562
+ }
563
+
564
+ // if only going up to 256 size, store directly back to global memory,
565
+ // otherwise store back to shared memory for next iteration
566
+ b32* store = (log_had_size <= 8) ? out_chunk_ptr : b_frag_ptr;
567
+
568
+ #pragma unroll
569
+ for (int64_t j = 0; j < 4; j++) {
570
+ int64_t reg = ((threadid & 16) == 0) ? j : (j / 2 * 2 + (1 - j % 2));
571
+ int64_t real_thread_id = (reg == 0 || reg == 2) ? threadid : (threadid ^ 16);
572
+ int64_t real_row = real_thread_id % 4;
573
+ int64_t real_col = real_thread_id / 4;
574
+ store[(real_row + (reg % 2) * 4) + (real_col + (reg / 2) * 8) * 8] = b_frag_all[k][j];
575
+ }
576
+ } else if constexpr(log_had_size > 8) { // condition is redundant to help compiler warnings
577
+ if (log_had_size < 12) {
578
+ // inverse of coalesced load for sizes 512, 1k and 2k to store result
579
+ constexpr int xor_val = log_had_size == 9 ? 16 : 1;
580
+ #pragma unroll
581
+ for (int64_t j = 2; j < 4; j++) {
582
+ b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], xor_val);
583
+ }
584
+
585
+ if ((threadid & xor_val) != 0) {
586
+ b32 temp = b_frag_all[k][0];
587
+ b_frag_all[k][0] = b_frag_all[k][2];
588
+ b_frag_all[k][2] = temp;
589
+
590
+ temp = b_frag_all[k][1];
591
+ b_frag_all[k][1] = b_frag_all[k][3];
592
+ b_frag_all[k][3] = temp;
593
+ }
594
+
595
+ b32* store = (b32*)(out + (blockid / warps_per_block) * (num_chunks * warps_per_block) * 256 + (256 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block) + k));
596
+ #pragma unroll
597
+ for (int64_t j = 0; j < 4; j++) {
598
+ int64_t reg = ((threadid & xor_val) == 0) ? j : (j + 2) % 4;
599
+ b32 data = b_frag_all[k][j];
600
+ int64_t real_thread_id = reg < 2 ? threadid : (threadid ^ xor_val);
601
+ int64_t idx = (real_thread_id / 4 * 16) + (real_thread_id % 4 * 2) + (reg / 2 * 8) + (reg % 2);
602
+ int64_t rowidx = idx % (1 << part8_log_had_size);
603
+ int64_t colidx = idx >> part8_log_had_size;
604
+ store[rowidx * 128 + colidx] = data;
605
+ }
606
+ }
607
+ // for size 4k and above, wait to process all chunks so a final store can be performed coalesced
608
+ }
609
+
610
+ a_chunk_ptr += 128; // (only affects first 256 size) move on to next chunk by skipping 256 elements in b16 (= 128 in b32)
611
+ out_chunk_ptr += 128;
612
+ if constexpr(log_had_size > 8) {
613
+ b_frag_ptr += (l == 0 ? 128 : (128 >> part8_log_had_size));
614
+ } else { // else is redundant, simplified version of if body, to help compiler warnings
615
+ b_frag_ptr += 128;
616
+ }
617
+ }
618
+ if (log_had_size <= 8)
619
+ break;
620
+ }
621
+
622
+ if constexpr(log_had_size >= 12) {
623
+ // for sizes 4k and above, perform final coalesced store after processing all chunks
624
+ #pragma unroll
625
+ for (int64_t j = 0; j < 4; j++) {
626
+ #pragma unroll
627
+ for (int64_t k = 1; k < num_chunks; k++) {
628
+ int64_t threadid_contig = threadid % num_chunks;
629
+ int64_t threadid_mul = threadid / num_chunks;
630
+ int64_t threadid2 = (threadid_contig + k) % num_chunks + threadid_mul * num_chunks; // thread to give your data to
631
+ b_frag_all[k][j] = __shfl_sync(0xFFFFFFFF, b_frag_all[k][j], threadid2);
632
+ }
633
+ }
634
+
635
+ // a + threadblock offset + warp offset
636
+ // can then index into all chunks owned by this warp
637
+ b32* store = bfrag_arr + (128 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block));
638
+
639
+ #pragma unroll
640
+ for (int64_t j = 0; j < 4; j++) {
641
+ #pragma unroll
642
+ for (int64_t k = 0; k < num_chunks; k++) {
643
+ // here, j represents register, and k represents 8-offset/chunk
644
+ int64_t real_chunk_num = (num_chunks - (threadid % num_chunks) + k) % num_chunks; // chunk at which you have target thread #'s data
645
+
646
+ // b32 data = b_frag_all[real_chunk_num][j]; // target thread data
647
+ b32 data;
648
+ #pragma unroll
649
+ for (int64_t i = 0; i < num_chunks; i++) {
650
+ if (real_chunk_num == i) data = b_frag_all[i][j];
651
+ }
652
+
653
+ int64_t real_thread_id = (threadid / num_chunks) * num_chunks + k; // target thread #
654
+ int64_t chunk_idx = 128 * real_chunk_num; // index due to fetching from another chunk (chunk in which this thread has the target thread's original data)
655
+ int64_t thread_group_idx = (real_thread_id / 4) * 16; // index due to fetching from another group of num_chunk threads (since shuffle is between num_chunk threads)
656
+ int64_t thread_idx = (real_thread_id % 4) * 2; // index due to original thread's position within the group of num_chunk threads
657
+ int64_t reg_idx = (j / 2) * 8 + (j % 2); // index due to target register
658
+ int64_t idx = chunk_idx + thread_group_idx + thread_idx + reg_idx; // final index
659
+
660
+ // fix idx for majorness
661
+ int64_t rowidx = idx % (1 << part8_log_had_size);
662
+ int64_t colidx = idx >> part8_log_had_size;
663
+
664
+ store[rowidx * 128 + colidx] = data;
665
+ }
666
+ }
667
+
668
+ __syncthreads();
669
+ store = ((b32*) out) + (blockid / warps_per_block) * (num_chunks * warps_per_block) * 128;
670
+ int4* store4 = (int4*) store;
671
+ int4* bfrag_arr4 = (int4*) bfrag_arr;
672
+ // flush smem, simply linearly write to store
673
+ // always divisible by 128*32b, so (32*4)*32b is ok
674
+ #pragma unroll
675
+ for (int64_t warp_off = 0; warp_off < (num_chunks * warps_per_block * 128 / 4); warp_off += 32 * warps_per_block) {
676
+ int64_t total_off = warp_off + threadid + (blockid % warps_per_block) * 32;
677
+ store4[total_off] = bfrag_arr4[total_off];
678
+ }
679
+ }
680
+
681
+}
682
+
683
+constexpr int64_t ceil_div(int64_t a, int64_t b) {
684
+ return (a + b - 1) / b;
685
+}
686
+
687
+template <torch::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
688
+void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaStream_t stream) {
689
+ int64_t shared_size = chunks_per_warp * warps_per_block * 128 * 4;
690
+ dim3 block_size = 32 * warps_per_block;
691
+
692
+ #define CHECK_SHARED_LIM() { \
693
+ if (shared_size > 48 * 1024) { \
694
+ C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
695
+ } \
696
+ } \
697
+
698
+ if constexpr(check_masking) {
699
+ if (num_chunks % (chunks_per_warp * warps_per_block) != 0) {
700
+ dim3 grid_size = ceil_div(ceil_div(num_chunks, chunks_per_warp), warps_per_block);
701
+ auto kernel = hadamard_transform_kernel<chunks_per_warp, warps_per_block, log_had_size, blocks_per_sm, true, dtype>;
702
+ CHECK_SHARED_LIM();
703
+ kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
704
+ } else {
705
+ dim3 grid_size = num_chunks / chunks_per_warp / warps_per_block;
706
+ auto kernel = hadamard_transform_kernel<chunks_per_warp, warps_per_block, log_had_size, blocks_per_sm, false, dtype>;
707
+ CHECK_SHARED_LIM();
708
+ kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
709
+ }
710
+ } else {
711
+ dim3 grid_size = num_chunks / chunks_per_warp / warps_per_block;
712
+ auto kernel = hadamard_transform_kernel<chunks_per_warp, warps_per_block, log_had_size, blocks_per_sm, false, dtype>;
713
+ CHECK_SHARED_LIM();
714
+ kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
715
+ }
716
+
717
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
718
+}
719
+
720
+template <torch::ScalarType dtype>
721
+void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream) {
722
+ int64_t num_chunks = numel / 256; // caller required to ensure divisible by 256
723
+ // for size 256, use (2, 1)
724
+ // for size 32k use (8, 16)
725
+ constexpr int64_t chunks_per_warp_small = 1;// 8;
726
+ constexpr int64_t warps_per_block_small = 1;//2;//16;
727
+ constexpr int64_t blocks_per_sm_small = 24;
728
+ constexpr int64_t chunks_per_warp_large = 2;
729
+ constexpr int64_t warps_per_block_large = 1;
730
+ constexpr int64_t blocks_per_sm_large = 24;
731
+
732
+ b16* a_mat = (b16*) a_mat_ptr;
733
+ b16* out = (b16*) out_ptr;
734
+
735
+ if (numel <= 256) {
736
+ switch (had_size) {
737
+ case (1<<1): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 1, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
738
+ case (1<<2): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 2, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
739
+ case (1<<3): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 3, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
740
+ case (1<<4): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 4, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
741
+ case (1<<5): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 5, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
742
+ case (1<<6): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 6, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
743
+ case (1<<7): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 7, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
744
+ case (1<<8): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 8, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
745
+ }
746
+ } else {
747
+ switch (had_size) {
748
+ case (1<<1): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 1, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
749
+ case (1<<2): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 2, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
750
+ case (1<<3): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 3, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
751
+ case (1<<4): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 4, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
752
+ case (1<<5): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 5, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
753
+ case (1<<6): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 6, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
754
+ case (1<<7): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 7, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
755
+ case (1<<8): run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 8, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
756
+ case (1<<9): run_kernel<dtype, launch_configs_big[0][0], launch_configs_big[0][1], 9 , launch_configs_big[0][2]>(a_mat, out, num_chunks, stream); break;
757
+ case (1<<10): run_kernel<dtype, launch_configs_big[1][0], launch_configs_big[1][1], 10, launch_configs_big[1][2]>(a_mat, out, num_chunks, stream); break;
758
+ case (1<<11): run_kernel<dtype, launch_configs_big[2][0], launch_configs_big[2][1], 11, launch_configs_big[2][2]>(a_mat, out, num_chunks, stream); break;
759
+ case (1<<12): run_kernel<dtype, launch_configs_big[3][0], launch_configs_big[3][1], 12, launch_configs_big[3][2]>(a_mat, out, num_chunks, stream); break;
760
+ case (1<<13): run_kernel<dtype, launch_configs_big[4][0], launch_configs_big[4][1], 13, launch_configs_big[4][2]>(a_mat, out, num_chunks, stream); break;
761
+ case (1<<14): run_kernel<dtype, launch_configs_big[5][0], launch_configs_big[5][1], 14, launch_configs_big[5][2]>(a_mat, out, num_chunks, stream); break;
762
+ case (1<<15): run_kernel<dtype, launch_configs_big[6][0], launch_configs_big[6][1], 15, launch_configs_big[6][2]>(a_mat, out, num_chunks, stream); break;
763
+ }
764
+ }
765
+}
766
+
767
+template void run_fht<torch::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
768
+template void run_fht<torch::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
769
+
770
+} // namespace hadacore
771
+
772
+constexpr bool is_power_of_two(int x) { return x && !(x & (x - 1)); }
773
+
774
+torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace) {
775
+ auto dtype = x.scalar_type();
776
+ TORCH_CHECK(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
777
+ TORCH_CHECK(x.is_cuda());
778
+
779
+ const int had_size = x.size(-1);
780
+ TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
781
+ "Only power of two Hadamard sizes up to 2^15 are supported, got ", had_size);
782
+
783
+ const auto res_shape = x.sizes();
784
+ x = x.reshape({-1, had_size});
785
+
786
+ auto numel = x.numel();
787
+ if (numel % 256 != 0) {
788
+ x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 0, 0, (256 - numel % 256) / had_size}));
789
+ }
790
+
791
+ if (x.stride(-1) != 1) {
792
+ x = x.contiguous();
793
+ }
794
+ torch::Tensor out = inplace ? x : torch::empty_like(x);
795
+
796
+ at::cuda::CUDAGuard device_guard{(char)x.get_device()};
797
+ auto stream = at::cuda::getCurrentCUDAStream().stream();
798
+
799
+ VLLM_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
800
+ auto constexpr SCALAR_TYPE = c10::CppTypeToScalarType<scalar_t>::value;
801
+ hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), x.data_ptr(), x.numel(), had_size, stream);
802
+ });
803
+
804
+ if (numel % 256 != 0) {
805
+ out = out.index({torch::indexing::Slice(0, numel / had_size)});
806
+ }
807
+
808
+ if (inplace && out.data_ptr() != x.data_ptr()) {
809
+ x.copy_(out.view(res_shape));
810
+ return x;
811
+ }
812
+ return out.reshape(res_shape);
813
+}
814
+
815
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
816
+ m.impl("hadacore_transform", &hadacore_transform);
817
+}
csrc/torch_bindings.cpp
+3
@@ -613,6 +613,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
613
"int pad_slot_id) -> ()");
614
ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
615
616
+ // Hadamard transforms
617
+ ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
618
+
619
#ifndef USE_ROCM
620
// Compute per-token-group FP8 quantized tensor and scaling factor.
621
ops.def(
tests/kernels/quantization/test_hadacore.py
new
+25
@@ -0,0 +1,25 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import math
5
+
6
+import pytest
7
+import torch
8
+from compressed_tensors.transform import deterministic_hadamard_matrix
9
+
10
+from vllm import _custom_ops as ops
11
+
12
+
13
+@pytest.mark.parametrize("batch_size", [1, 32])
14
+@pytest.mark.parametrize("hidden_dim", [2**n for n in range(10)])
15
+def test_hadacore(batch_size, hidden_dim, dtype=torch.bfloat16, device="cuda"):
16
+ x = torch.eye(hidden_dim, dtype=dtype, device=device)
17
+ hadamard = deterministic_hadamard_matrix(
18
+ hidden_dim, dtype=torch.float64, device="cuda") / math.sqrt(hidden_dim)
19
+
20
+ y = ops.hadacore_transform(x.clone())
21
+ y_true = (x.to(hadamard.dtype) @ hadamard.T).to(y.dtype)
22
+ assert torch.allclose(y, y_true)
23
+
24
+ y = ops.hadacore_transform(y)
25
+ assert torch.allclose(y, x)
vllm/_custom_ops.py
+24
@@ -2011,3 +2011,27 @@ def onednn_scaled_mm(
2011
input_zp_adj, bias, dnnl_handler.handler)
2012
2013
return output
2014
+
2015
+
2016
+def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
2017
+ """
2018
+ Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)
2019
+ kernels. Note that these kernels exploit the recursive properties of
2020
+ Sylvester Hadamards, and therefore do not require transform weight data
2021
+
2022
+ Note that sylvester hadamard transforms are also symmetric, which means that
2023
+ this function is also applies the (transpose <=> inverse) transform.
2024
+
2025
+ :param x: value to be transformed inplace
2026
+ :param inplace: modify value in place
2027
+ :return: value after transformation
2028
+ """
2029
+ return torch.ops._C.hadacore_transform(x, inplace)
2030
+
2031
+
2032
+if hasattr(torch.ops._C, "hadacore_transform"):
2033
+
2034
+ @register_fake("_C::hadacore_transform")
2035
+ def _hadacore_transform_fake(x: torch.Tensor,
2036
+ inplace: bool) -> torch.Tensor:
2037
+ return torch.empty_like(x) if not inplace else x
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+1
-1
@@ -129,7 +129,7 @@ class CompressedTensorsConfig(QuantizationConfig):
129
# choose transform method
130
if any((input_tfms, output_tfms)):
131
return CompressedTensorsLinearTransformMethod.from_schemes(
132
- quant_method, input_tfms, output_tfms)
132
+ quant_method, quant_scheme, input_tfms, output_tfms)
133
134
else:
135
return quant_method
vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py
+18
-7
@@ -12,6 +12,8 @@ from compressed_tensors.utils import is_match
12
from vllm.model_executor.layers.linear import (WEIGHT_LOADER_V2_SUPPORTED,
13
LinearMethodBase,
14
QKVCrossParallelLinear)
15
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
16
+ CompressedTensorsScheme)
17
from vllm.model_executor.layers.quantization.compressed_tensors.transform.module import ( # noqa: E501
18
HadamardTransform)
19
from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501
@@ -26,14 +28,22 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase):
28
29
@classmethod
30
def from_schemes(
29
- cls, quant_method: LinearMethodBase, input_tfms: dict[int,
30
- TransformTuple],
31
- output_tfms: dict[int, TransformTuple]
31
+ cls,
32
+ quant_method: LinearMethodBase,
33
+ quant_scheme: Optional[CompressedTensorsScheme],
34
+ input_tfms: dict[int, TransformTuple],
35
+ output_tfms: dict[int, TransformTuple],
36
) -> "CompressedTensorsLinearTransformMethod":
37
+ from vllm.model_executor.layers.quantization.compressed_tensors.transform.schemes.linear_qutlass_nvfp4 import ( # noqa: E501
38
+ QutlassNvFP4LinearMethod, is_qutlass_fp4_scheme)
39
+
40
assert input_tfms or output_tfms
41
35
- # TODO (@ksayers): implement QutlassLinearMethodNvFP4
36
- # hadacore and fwht can be selected by Transform module
42
+ if is_qutlass_fp4_scheme(quant_scheme, input_tfms):
43
+ return QutlassNvFP4LinearMethod(quant_method, input_tfms,
44
+ output_tfms)
45
+
46
+ # hadacore or dense gemm is selected by Transform module
47
48
return cls(quant_method, input_tfms, output_tfms)
49
@@ -129,11 +139,12 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase):
139
assert bias is None
140
x = self.quant_method.apply(layer, x, bias)
141
132
- # TODO (@ksayers): Write a triton kernel to do this in parallel
142
+ # In most cases, input transforms are preferred over output transforms
143
+ # (@ksayers): confirm that this is done concurrently
144
if self.output_transform is not None:
145
for part_id, (start, length) in enumerate(self.partition_ranges):
146
x[:, start:start + length] = self.output_transform(
136
- x[:, start:start + length], part_id=part_id)
147
+ x[:, start:start + length].contiguous(), part_id=part_id)
148
149
return x
150
vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py
+47
-29
@@ -2,12 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import math
4
from collections.abc import Hashable
5
-from typing import Callable, Optional
5
+from typing import Callable
6
7
import torch
8
-from compressed_tensors.transform import TransformLocation, TransformScheme
8
+from compressed_tensors.transform import (TransformArgs, TransformLocation,
9
+ TransformScheme)
10
from torch import Tensor
11
12
+import vllm._custom_ops as ops
13
from vllm.distributed.parallel_state import (
14
get_tensor_model_parallel_world_size)
15
from vllm.model_executor.layers.linear import LinearBase
@@ -28,16 +30,12 @@ class HadamardTransform(torch.nn.Module):
30
transforms: dict[int, TransformTuple] # info parsed from transforms config
31
weight: SharedWeightParameter # container for shared tensors
32
31
- kernel: Callable # function used during application
33
scales: dict[int, float] # hadamard scale, usually sqrt(matrix.size(0))
34
34
- def __init__(self,
35
- transforms: dict[int, TransformTuple],
36
- layer: torch.nn.Module,
37
- weight_loader: Callable,
35
+ def __init__(self, transforms: dict[int, TransformTuple],
36
+ layer: torch.nn.Module, weight_loader: Callable,
37
input_size_per_partition: int,
39
- output_partition_sizes: list[int],
40
- kernel: Optional[Callable] = None):
38
+ output_partition_sizes: list[int]):
39
super().__init__()
40
self.transforms = transforms
41
self.scales = {}
@@ -55,7 +53,7 @@ class HadamardTransform(torch.nn.Module):
53
for part_index, (_scheme_name, scheme,
54
args) in self.transforms.items():
55
output_size = output_partition_sizes[part_index]
58
- weight_size = self._get_weight_size(layer, args.location,
56
+ weight_size = self._get_weight_size(layer, scheme, args,
57
input_size, output_size)
58
59
data_key = self._get_data_key(scheme, weight_size)
@@ -69,9 +67,6 @@ class HadamardTransform(torch.nn.Module):
67
# validate that shared tensors and schemes are correct
68
self._validate_input_transforms()
69
72
- # select kernel based on transform schemes
73
- self.kernel = self._infer_kernel() if kernel is None else kernel
74
-
70
def process_weights_after_loading(self):
71
for part_id in self.weight.partitions:
72
data = self.weight.partitions[part_id].data
@@ -90,32 +85,59 @@ class HadamardTransform(torch.nn.Module):
85
if part_id not in self.weight.partitions:
86
return value
87
93
- weight = self.weight.partitions[part_id]
94
- weight = weight if self.transforms[
95
- part_id].args.inverse else weight.T # linear := x(W.T)
96
- scale = self.scales[part_id]
97
- return self.kernel(self, value.to(weight.dtype), weight, None).to(
98
- value.dtype) * scale
88
+ # use hadacore if possible
89
+ if self.transforms[part_id].scheme.type == "hadamard":
90
+ if self.transforms[part_id].scheme.head_dim is not None:
91
+ weight_size = self.transforms[part_id].scheme.head_dim
92
+ value = value.unflatten(-1, (-1, weight_size))
93
+ value = ops.hadacore_transform(value)
94
+ value = value.flatten(-2, -1)
95
+
96
+ return value
97
+
98
+ # sylvester transforms are symmetric, inv => transpose => original
99
+ return ops.hadacore_transform(value)
100
+
101
+ # fall back to dense
102
+ else:
103
+ weight = self.weight.partitions[part_id]
104
+ weight = weight if self.transforms[
105
+ part_id].args.inverse else weight.T # linear := x(W.T)
106
+ scale = self.scales[part_id]
107
+
108
+ if self.transforms[part_id].scheme.head_dim is not None:
109
+ value = value.unflatten(-1, (-1, weight.size(0)))
110
+ value = dispatch_unquantized_gemm()(self, value.to(
111
+ weight.dtype), weight, None).to(value.dtype) * scale
112
+ value = value.flatten(-2, -1)
113
+
114
+ return value
115
+
116
+ return dispatch_unquantized_gemm()(self, value.to(
117
+ weight.dtype), weight, None).to(value.dtype) * scale
118
119
def _get_data_key(self, scheme: TransformScheme,
120
weight_size: int) -> Hashable:
121
return (id(scheme), weight_size)
122
104
- def _get_weight_size(self, layer: torch.nn.Module,
105
- location: TransformLocation, input_size: int,
123
+ def _get_weight_size(self, layer: torch.nn.Module, scheme: TransformScheme,
124
+ args: TransformArgs, input_size: int,
125
output_size: int) -> int:
126
+ if scheme.head_dim is not None:
127
+ return scheme.head_dim
128
+
129
if isinstance(layer, LinearBase):
108
- if location == TransformLocation.INPUT:
130
+ if args.location == TransformLocation.INPUT:
131
return input_size
132
111
- elif location == TransformLocation.OUTPUT:
133
+ elif args.location == TransformLocation.OUTPUT:
134
return output_size
135
136
elif isinstance(layer, VocabParallelEmbedding):
115
- if location == TransformLocation.INPUT:
137
+ if args.location == TransformLocation.INPUT:
138
return output_size
139
118
- elif location == TransformLocation.OUTPUT:
140
+ elif args.location == TransformLocation.OUTPUT:
141
return input_size
142
143
raise ValueError()
@@ -129,7 +151,3 @@ class HadamardTransform(torch.nn.Module):
151
for partition in self.weight.partitions.values():
152
if partition.data.data_ptr() != first_data.data_ptr():
153
raise ValueError("")
132
-
133
- def _infer_kernel(self) -> Callable:
134
- # TODO (@ksayers): use fwht, hadacore
135
- return dispatch_unquantized_gemm()
vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py
+31
-6
@@ -4,18 +4,43 @@ from typing import Optional
4
5
import torch
6
7
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
8
+ CompressedTensorsScheme, CompressedTensorsW4A4Fp4)
9
from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear import ( # noqa: E501
8
- CompressedTensorsLinearTransformMethod)
10
+ CompressedTensorsLinearTransformMethod, TransformTuple)
11
12
+__all__ = ["is_qutlass_fp4_scheme", "QutlassNvFP4LinearMethod"]
13
11
-# Because qutlass fuses hadamard with quantization, it cannot automatically be
12
-# composed with kernels in the way CompressedTensorsLinearTransformMethod does.
13
-# Therefore, a separate scheme must be created for each quantized dtype
14
-class QutlassLinearMethodNvFP4(CompressedTensorsLinearTransformMethod):
14
+
15
+def is_qutlass_fp4_scheme(quant_scheme: Optional[CompressedTensorsScheme],
16
+ input_tfms: dict[int, TransformTuple]) -> bool:
17
+ return isinstance(
18
+ quant_scheme,
19
+ (CompressedTensorsW4A4Fp4, )) and len(input_tfms) == 1 and input_tfms[
20
+ 0].scheme.head_dim == quant_scheme.group_size
21
+
22
+
23
+class QutlassNvFP4LinearMethod(CompressedTensorsLinearTransformMethod):
24
+
25
+ def create_weights(self, layer, input_size_per_partition,
26
+ output_partition_sizes, input_size, output_size,
27
+ params_dtype, **extra_weight_attrs):
28
+ # initializes fp4 qparams
29
+ assert isinstance(layer.scheme, (CompressedTensorsW4A4Fp4, ))
30
+ ret = super().create_weights(layer, input_size_per_partition,
31
+ output_partition_sizes, input_size,
32
+ output_size, params_dtype,
33
+ **extra_weight_attrs)
34
+
35
+ assert self.input_transform is not None
36
+ assert len(self.input_transform.weight) == 1
37
+ assert self.input_transform.weight[0].size(
38
+ 0) == layer.scheme.group_size
39
+
40
+ return ret
41
42
def apply(self,
43
layer: torch.nn.Module,
44
x: torch.Tensor,
45
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
20
- # fused hadamard quant linear method
46
raise NotImplementedError()