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 12, 22:15 UTC
cleanup at::Tag::needs_fixed_stride_order (#28974)
Signed-off-by: Boyuan Feng <boyuan@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Boyuan Feng committed
Nov 20, 2025 at 02:51 UTC
a903d59ffaffd9160c517fa337b3ab0265a898c3
2 files changed
+20
-51
csrc/cpu/torch_bindings.cpp
+2
-5
@@ -172,7 +172,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
172
// Quantization
173
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
174
defined(__powerpc64__)
175
- at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
175
// Helper function to release oneDNN handlers
176
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
177
&release_dnnl_matmul_handler);
@@ -208,15 +207,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
207
// Compute int8 quantized tensor for given scaling factor.
208
ops.def(
209
"static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
211
- "Tensor? azp) -> ()",
212
- {stride_tag});
210
+ "Tensor? azp) -> ()");
211
ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);
212
213
// Compute int8 quantized tensor and scaling factor
214
ops.def(
215
"dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
218
- "Tensor!? azp) -> ()",
219
- {stride_tag});
216
+ "Tensor!? azp) -> ()");
217
ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
218
&dynamic_scaled_int8_quant);
219
#endif
csrc/torch_bindings.cpp
+18
-46
@@ -20,18 +20,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
20
// vLLM custom ops
21
//
22
23
- // The default behavior in PyTorch 2.6 was changed to "requires_contiguous",
24
- // so we need
25
- // to override this for many GEMMs with the following tag. Otherwise,
26
- // torch.compile will force all input tensors to be contiguous(), which
27
- // will break many custom ops that require column-major weight matrices.
28
- // This was a bug and PyTorch 2.7 has since fixed this.
29
-#if TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 6
30
- #define stride_tag at::Tag::needs_fixed_stride_order
31
-#else
32
- #define stride_tag
33
-#endif
34
-
23
ops.def(
24
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
25
"y_q, Tensor! y_s,"
@@ -241,15 +229,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
229
// Quantized GEMM for AWQ.
230
ops.def(
231
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
244
- "Tensor _zeros, SymInt split_k_iters) -> Tensor",
245
- {stride_tag});
232
+ "Tensor _zeros, SymInt split_k_iters) -> Tensor");
233
ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
234
235
// Dequantization for AWQ.
236
ops.def(
237
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
251
- "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor",
252
- {stride_tag});
238
+ "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
239
ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
240
241
// Note about marlin kernel 'workspace' arguments:
@@ -271,8 +257,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
257
"gptq_marlin_24_gemm(Tensor a, Tensor b_q_weight, Tensor b_meta, "
258
"Tensor b_scales, Tensor workspace, "
259
"int b_q_type, "
274
- "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor",
275
- {stride_tag});
260
+ "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor");
261
// conditionally compiled so impl in source file
262
263
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
@@ -298,8 +283,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
283
" Tensor? channel_scales,"
284
" Tensor? token_scales,"
285
" str? schedule"
301
- ") -> Tensor",
302
- {stride_tag});
286
+ ") -> Tensor");
287
ops.def(
288
"machete_prepack_B("
289
" Tensor B,"
@@ -319,8 +303,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
303
"Tensor b_scales, Tensor? global_scale, Tensor? b_zeros_or_none, Tensor? "
304
"g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_q_type, "
305
"SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
322
- "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor",
323
- {stride_tag});
306
+ "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
307
// conditionally compiled so impl registration is in source file
308
309
// gptq_marlin repack from GPTQ.
@@ -346,8 +329,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
329
" Tensor token_scales,"
330
" ScalarType? out_type,"
331
" str? maybe_schedule"
349
- ") -> Tensor",
350
- {stride_tag});
332
+ ") -> Tensor");
333
// pack scales
334
ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
335
// encode and reorder weight matrix
@@ -394,24 +376,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
376
ops.def(
377
"cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
378
" Tensor block_scale_a, Tensor block_scale_b,"
397
- " Tensor alpha) -> ()",
398
- {stride_tag});
379
+ " Tensor alpha) -> ()");
380
ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
381
382
// cutlass blockwise scaledgroup GEMM
383
ops.def(
384
"cutlass_blockwise_scaled_grouped_mm(Tensor! output, Tensor a, Tensor b, "
385
"Tensor scales_a, Tensor scales_b, "
405
- "Tensor problem_sizes, Tensor expert_offsets) -> ()",
406
- {stride_tag});
386
+ "Tensor problem_sizes, Tensor expert_offsets) -> ()");
387
// conditionally compiled so impl registration is in source file
388
389
// cutlass nvfp4 block scaled group GEMM
390
ops.def(
391
"cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
392
" Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
413
- " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()",
414
- {stride_tag});
393
+ " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
394
// conditionally compiled so impl registration is in source file
395
396
// CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
@@ -419,8 +398,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
398
ops.def(
399
"cutlass_scaled_mm(Tensor! out, Tensor a,"
400
" Tensor b, Tensor a_scales,"
422
- " Tensor b_scales, Tensor? bias) -> ()",
423
- {stride_tag});
401
+ " Tensor b_scales, Tensor? bias) -> ()");
402
ops.impl("cutlass_scaled_mm", torch::kCUDA, &cutlass_scaled_mm);
403
404
// CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
@@ -429,8 +407,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
407
"cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
408
" Tensor b, Tensor a_scales,"
409
" Tensor b_scales, Tensor azp_adj,"
432
- " Tensor? azp, Tensor? bias) -> ()",
433
- {stride_tag});
410
+ " Tensor? azp, Tensor? bias) -> ()");
411
ops.impl("cutlass_scaled_mm_azp", torch::kCUDA, &cutlass_scaled_mm_azp);
412
413
// Check if cutlass scaled_mm is supported for CUDA devices of the given
@@ -449,8 +426,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
426
" Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
427
" Tensor problem_sizes, Tensor a_strides, "
428
" Tensor b_strides, Tensor c_strides, bool per_act_token, "
452
- " bool per_out_ch) -> ()",
453
- {stride_tag});
429
+ " bool per_out_ch) -> ()");
430
ops.impl("cutlass_moe_mm", torch::kCUDA, &cutlass_moe_mm);
431
432
// A function that computes data required to run fused MoE with w8a8 grouped
@@ -464,8 +440,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
440
" Tensor! problem_sizes1, Tensor! problem_sizes2, "
441
" Tensor! input_permutation, "
442
" Tensor! output_permutation, int num_experts, "
467
- " int n, int k, Tensor? blockscale_offsets) -> ()",
468
- {stride_tag});
443
+ " int n, int k, Tensor? blockscale_offsets) -> "
444
+ "()");
445
ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);
446
447
// A function that computes problem sizes for each expert's multiplication
@@ -476,8 +452,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
452
" Tensor! problem_sizes1, "
453
" Tensor! problem_sizes2, "
454
" int num_experts, int n, int k, "
479
- " Tensor? blockscale_offsets) -> ()",
480
- {stride_tag});
455
+ " Tensor? blockscale_offsets) -> ()");
456
ops.impl("get_cutlass_moe_mm_problem_sizes", torch::kCUDA,
457
&get_cutlass_moe_mm_problem_sizes);
458
@@ -492,8 +467,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
467
" Tensor! problem_sizes2, "
468
" Tensor expert_num_tokens, "
469
" int num_local_experts, int padded_m, "
495
- " int n, int k) -> ()",
496
- {stride_tag});
470
+ " int n, int k) -> ()");
471
ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA,
472
&get_cutlass_pplx_moe_mm_data);
473
@@ -517,8 +491,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
491
"cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
492
" Tensor bt_nzs,"
493
" Tensor bt_meta, Tensor a_scales,"
520
- " Tensor b_scales, Tensor? bias) -> ()",
521
- {stride_tag});
494
+ " Tensor b_scales, Tensor? bias) -> ()");
495
ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);
496
497
// CUTLASS sparse matrix compressor
@@ -567,8 +540,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
540
"gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
541
"Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
542
"use_v2_format, int bit) "
570
- "-> Tensor",
571
- {stride_tag});
543
+ "-> Tensor");
544
ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
545
546
// Post processing for GPTQ.