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 15, 21:00 UTC
[Kernel] Add needs_fixed_stride_order tag to most GEMMs (#14306)
Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>
Tyler Michael Smith committed
Mar 6, 2025 at 17:17 UTC
99b0915d3bcce5beeb812407b354179ee3092b4d
1 file changed
+40
-15
csrc/torch_bindings.cpp
+40
-15
@@ -4,6 +4,7 @@
4
#include "core/registration.h"
5
6
#include <torch/library.h>
7
+#include <torch/version.h>
8
9
// Note on op signatures:
10
// The X_meta signatures are for the meta functions corresponding to op X.
@@ -17,6 +18,15 @@
18
19
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
20
// vLLM custom ops
21
+ //
22
+
23
+ // The default behavior in PyTorch 2.6 is "requires_contiguous", so we need
24
+ // to override this for many GEMMs with the following tag. Otherwise,
25
+ // torch.compile will force all input tensors to be contiguous(), which
26
+ // will break many custom ops that require column-major weight matrices.
27
+ // TODO: remove this for PyTorch 2.8, when the default is planned to switch
28
+ // to match exact eager-mode strides.
29
+ at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
30
31
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
32
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
@@ -163,25 +173,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
173
ops.def(
174
"aqlm_gemm(Tensor input, Tensor codes, Tensor codebooks, "
175
"Tensor scales, int[] codebook_partition_sizes, Tensor? bias) "
166
- "-> Tensor");
176
+ "-> Tensor",
177
+ {stride_tag});
178
ops.impl("aqlm_gemm", torch::kCUDA, &aqlm_gemm);
179
180
// Decompression method for AQLM.
181
ops.def(
182
"aqlm_dequant(Tensor codes, Tensor codebooks, "
172
- "int[] codebook_partition_sizes) -> Tensor");
183
+ "int[] codebook_partition_sizes) -> Tensor",
184
+ {stride_tag});
185
ops.impl("aqlm_dequant", torch::kCUDA, &aqlm_dequant);
186
187
// Quantized GEMM for AWQ.
188
ops.def(
189
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
178
- "Tensor _zeros, SymInt split_k_iters) -> Tensor");
190
+ "Tensor _zeros, SymInt split_k_iters) -> Tensor",
191
+ {stride_tag});
192
ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
193
194
// Dequantization for AWQ.
195
ops.def(
196
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
184
- "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
197
+ "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor",
198
+ {stride_tag});
199
ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
200
201
// Note about marlin kernel 'workspace' arguments:
@@ -202,7 +216,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
216
ops.def(
217
"marlin_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, "
218
"Tensor! workspace, SymInt size_m, SymInt size_n, SymInt size_k) -> "
205
- "Tensor");
219
+ "Tensor",
220
+ {stride_tag});
221
// conditionally compiled so impl in source file
222
223
// Marlin_24 (Sparse) Optimized Quantized GEMM for GPTQ.
@@ -210,7 +225,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
225
"gptq_marlin_24_gemm(Tensor a, Tensor b_q_weight, Tensor b_meta, "
226
"Tensor b_scales, Tensor workspace, "
227
"int b_q_type, "
213
- "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor");
228
+ "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor",
229
+ {stride_tag});
230
// conditionally compiled so impl in source file
231
232
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
@@ -236,7 +252,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
252
" Tensor? channel_scales,"
253
" Tensor? token_scales,"
254
" str? schedule"
239
- ") -> Tensor");
255
+ ") -> Tensor",
256
+ {stride_tag});
257
ops.def(
258
"machete_prepack_B("
259
" Tensor B,"
@@ -255,7 +272,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
272
"Tensor b_zeros, Tensor g_idx, Tensor perm, Tensor workspace, "
273
"int b_q_type, "
274
"SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
258
- "bool has_zp, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
275
+ "bool has_zp, bool use_fp32_reduce, bool is_zp_float) -> Tensor",
276
+ {stride_tag});
277
// conditionally compiled so impl registration is in source file
278
279
// gptq_marlin repack from GPTQ.
@@ -291,7 +309,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
309
ops.def(
310
"fp8_marlin_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, "
311
"Tensor! workspace, int num_bits, SymInt size_m, SymInt size_n, "
294
- "SymInt size_k) -> Tensor");
312
+ "SymInt size_k) -> Tensor",
313
+ {stride_tag});
314
// conditionally compiled so impl registration is in source file
315
316
// marlin_qqq_gemm for QQQ.
@@ -299,14 +318,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
318
"marlin_qqq_gemm(Tensor a, Tensor b_q_weight, "
319
"Tensor s_tok, Tensor s_ch, Tensor s_group, "
320
"Tensor! workspace, SymInt size_m, SymInt size_n, "
302
- "SymInt size_k) -> Tensor");
321
+ "SymInt size_k) -> Tensor",
322
+ {stride_tag});
323
// conditionally compiled so impl registration is in source file
324
325
// CUTLASS nvfp4 block scaled GEMM
326
ops.def(
327
"cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
328
" Tensor block_scale_a, Tensor block_scale_b,"
309
- " Tensor alpha) -> ()");
329
+ " Tensor alpha) -> ()",
330
+ {stride_tag});
331
ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
332
333
// CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
@@ -314,7 +335,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
335
ops.def(
336
"cutlass_scaled_mm(Tensor! out, Tensor a,"
337
" Tensor b, Tensor a_scales,"
317
- " Tensor b_scales, Tensor? bias) -> ()");
338
+ " Tensor b_scales, Tensor? bias) -> ()",
339
+ {stride_tag});
340
ops.impl("cutlass_scaled_mm", torch::kCUDA, &cutlass_scaled_mm);
341
342
// CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
@@ -323,7 +345,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
345
"cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
346
" Tensor b, Tensor a_scales,"
347
" Tensor b_scales, Tensor azp_adj,"
326
- " Tensor? azp, Tensor? bias) -> ()");
348
+ " Tensor? azp, Tensor? bias) -> ()",
349
+ {stride_tag});
350
ops.impl("cutlass_scaled_mm_azp", torch::kCUDA, &cutlass_scaled_mm_azp);
351
352
// Check if cutlass scaled_mm is supported for CUDA devices of the given
@@ -351,7 +374,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
374
"cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
375
" Tensor bt_nzs,"
376
" Tensor bt_meta, Tensor a_scales,"
354
- " Tensor b_scales, Tensor? bias) -> ()");
377
+ " Tensor b_scales, Tensor? bias) -> ()",
378
+ {stride_tag});
379
ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);
380
381
// CUTLASS sparse matrix compressor
@@ -407,7 +431,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
431
ops.def(
432
"gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
433
"Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, int bit) "
410
- "-> Tensor");
434
+ "-> Tensor",
435
+ {stride_tag});
436
ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
437
438
// Post processing for GPTQ.