A high-throughput and memory-efficient inference and serving engine for LLMs
[11a/n] Migrate Marlin kernels to torch stable ABI (#45176)
Signed-off-by: Chris Leonard <chleonar@redhat.com>
Chris Leonard committed
Jun 12, 2026 at 00:22 UTC
7021be66e8c351fa819fd07b4053e946f11c1147
19 files changed
+627
-573
CMakeLists.txt
+139
-139
@@ -358,145 +358,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
358
SRCS "${VLLM_EXT_SRC}"
359
CUDA_ARCHS "${CUDA_ARCHS}")
360
361
- # Only build Marlin kernels if we are building for at least some compatible archs.
362
- # Keep building Marlin for 9.0 as there are some group sizes and shapes that
363
- # are not supported by Machete yet.
364
-
365
- # marlin arches for fp16 output
366
- # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
367
- # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
368
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
369
- cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
370
- else()
371
- cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
372
- endif()
373
- # marlin has limited support for turing
374
- cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
375
- # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
376
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
377
- cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
378
- else()
379
- cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
380
- endif()
381
- # marlin arches for fp8 input
382
- # - sm80 doesn't support fp8 computation
383
- # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
384
- # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
385
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
386
- cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
387
- else()
388
- cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
389
- endif()
390
- # marlin arches for other files
391
- cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
392
-
393
- if (MARLIN_OTHER_ARCHS)
394
-
395
- #
396
- # For the Marlin kernels we automatically generate sources for various
397
- # preselected input type pairs and schedules.
398
- # Generate sources:
399
- set(MARLIN_GEN_SCRIPT
400
- ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py)
401
- file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
402
- list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
403
- set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
404
-
405
- message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
406
- message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
407
-
408
- if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
409
- OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
410
- execute_process(
411
- COMMAND ${CMAKE_COMMAND} -E env
412
- PYTHONPATH=$ENV{PYTHONPATH}
413
- ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
414
- RESULT_VARIABLE marlin_generation_result
415
- OUTPUT_VARIABLE marlin_generation_result
416
- OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
417
- ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
418
- )
419
-
420
- if (NOT marlin_generation_result EQUAL 0)
421
- message(FATAL_ERROR "Marlin generation failed."
422
- " Result: \"${marlin_generation_result}\""
423
- "\nCheck the log for details: "
424
- "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
425
- else()
426
- set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
427
- CACHE STRING "Last run Marlin generate script hash and arch" FORCE)
428
- message(STATUS "Marlin generation completed successfully.")
429
- endif()
430
- else()
431
- message(STATUS "Marlin generation script has not changed, skipping generation.")
432
- endif()
433
-
434
- if (MARLIN_ARCHS)
435
- file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu")
436
- set_gencode_flags_for_srcs(
437
- SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
438
- CUDA_ARCHS "${MARLIN_ARCHS}")
439
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
440
- set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
441
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
442
- endif()
443
- list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
444
-
445
- file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu")
446
- set_gencode_flags_for_srcs(
447
- SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
448
- CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
449
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
450
- set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
451
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
452
- endif()
453
- list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
454
- endif()
455
-
456
- if (MARLIN_SM75_ARCHS)
457
- file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu")
458
- set_gencode_flags_for_srcs(
459
- SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
460
- CUDA_ARCHS "${MARLIN_SM75_ARCHS}")
461
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
462
- set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC}
463
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
464
- endif()
465
- list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
466
- endif()
467
-
468
- if (MARLIN_FP8_ARCHS)
469
- file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu")
470
- set_gencode_flags_for_srcs(
471
- SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
472
- CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
473
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
474
- set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC}
475
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
476
- endif()
477
- list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC})
478
- endif()
479
-
480
- set(MARLIN_SRCS
481
- "csrc/quantization/marlin/marlin.cu"
482
- "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu"
483
- "csrc/quantization/marlin/gptq_marlin_repack.cu"
484
- "csrc/quantization/marlin/awq_marlin_repack.cu")
485
- set_gencode_flags_for_srcs(
486
- SRCS "${MARLIN_SRCS}"
487
- CUDA_ARCHS "${MARLIN_OTHER_ARCHS}")
488
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
489
- set_source_files_properties(${MARLIN_SRCS}
490
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
491
- endif()
492
- list(APPEND VLLM_EXT_SRC "${MARLIN_SRCS}")
493
-
494
- message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}")
495
- else()
496
- message(STATUS "Not building Marlin kernels as no compatible archs found"
497
- " in CUDA target architectures")
498
- endif()
499
-
361
# Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
362
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
363
cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
@@ -676,6 +537,145 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
537
SRCS "${VLLM_STABLE_EXT_SRC}"
538
CUDA_ARCHS "${CUDA_ARCHS}")
539
540
+ # Only build Marlin kernels if we are building for at least some compatible archs.
541
+ # Keep building Marlin for 9.0 as there are some group sizes and shapes that
542
+ # are not supported by Machete yet.
543
+
544
+ # marlin arches for fp16 output
545
+ # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
546
+ # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
547
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
548
+ cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
549
+ else()
550
+ cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
551
+ endif()
552
+ # marlin has limited support for turing
553
+ cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
554
+ # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
555
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
556
+ cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
557
+ else()
558
+ cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
559
+ endif()
560
+ # marlin arches for fp8 input
561
+ # - sm80 doesn't support fp8 computation
562
+ # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
563
+ # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
564
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
565
+ cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
566
+ else()
567
+ cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
568
+ endif()
569
+ # marlin arches for other files
570
+ cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
571
+
572
+ if (MARLIN_OTHER_ARCHS)
573
+
574
+ #
575
+ # For the Marlin kernels we automatically generate sources for various
576
+ # preselected input type pairs and schedules.
577
+ # Generate sources:
578
+ set(MARLIN_GEN_SCRIPT
579
+ ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py)
580
+ file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
581
+ list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
582
+ set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
583
+
584
+ message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
585
+ message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
586
+
587
+ if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
588
+ OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
589
+ execute_process(
590
+ COMMAND ${CMAKE_COMMAND} -E env
591
+ PYTHONPATH=$ENV{PYTHONPATH}
592
+ ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
593
+ RESULT_VARIABLE marlin_generation_result
594
+ OUTPUT_VARIABLE marlin_generation_result
595
+ OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
596
+ ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
597
+ )
598
+
599
+ if (NOT marlin_generation_result EQUAL 0)
600
+ message(FATAL_ERROR "Marlin generation failed."
601
+ " Result: \"${marlin_generation_result}\""
602
+ "\nCheck the log for details: "
603
+ "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
604
+ else()
605
+ set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
606
+ CACHE STRING "Last run Marlin generate script hash and arch" FORCE)
607
+ message(STATUS "Marlin generation completed successfully.")
608
+ endif()
609
+ else()
610
+ message(STATUS "Marlin generation script has not changed, skipping generation.")
611
+ endif()
612
+
613
+ if (MARLIN_ARCHS)
614
+ file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu")
615
+ set_gencode_flags_for_srcs(
616
+ SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
617
+ CUDA_ARCHS "${MARLIN_ARCHS}")
618
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
619
+ set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
620
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
621
+ endif()
622
+ list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
623
+
624
+ file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu")
625
+ set_gencode_flags_for_srcs(
626
+ SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
627
+ CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
628
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
629
+ set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
630
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
631
+ endif()
632
+ list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
633
+ endif()
634
+
635
+ if (MARLIN_SM75_ARCHS)
636
+ file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu")
637
+ set_gencode_flags_for_srcs(
638
+ SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
639
+ CUDA_ARCHS "${MARLIN_SM75_ARCHS}")
640
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
641
+ set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC}
642
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
643
+ endif()
644
+ list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
645
+ endif()
646
+
647
+ if (MARLIN_FP8_ARCHS)
648
+ file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu")
649
+ set_gencode_flags_for_srcs(
650
+ SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
651
+ CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
652
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
653
+ set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC}
654
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
655
+ endif()
656
+ list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC})
657
+ endif()
658
+
659
+ set(MARLIN_SRCS
660
+ "csrc/libtorch_stable/quantization/marlin/marlin.cu"
661
+ "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu"
662
+ "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu"
663
+ "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu")
664
+ set_gencode_flags_for_srcs(
665
+ SRCS "${MARLIN_SRCS}"
666
+ CUDA_ARCHS "${MARLIN_OTHER_ARCHS}")
667
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
668
+ set_source_files_properties(${MARLIN_SRCS}
669
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
670
+ endif()
671
+ list(APPEND VLLM_STABLE_EXT_SRC "${MARLIN_SRCS}")
672
+
673
+ message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}")
674
+ else()
675
+ message(STATUS "Not building Marlin kernels as no compatible archs found"
676
+ " in CUDA target architectures")
677
+ endif()
678
+
679
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
680
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
681
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h
+2
-2
@@ -3,8 +3,8 @@
3
#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
4
#endif
5
6
-#include "quantization/marlin/marlin.cuh"
7
-#include "quantization/marlin/marlin_dtypes.cuh"
6
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
7
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
8
#include "core/scalar_type.hpp"
9
10
#define MARLIN_KERNEL_PARAMS \
csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h
+4
-4
@@ -23,10 +23,10 @@
23
#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
24
#endif
25
26
-#include "quantization/marlin/marlin.cuh"
27
-#include "quantization/marlin/marlin_dtypes.cuh"
28
-#include "quantization/marlin/dequant.h"
29
-#include "quantization/marlin/marlin_mma.h"
26
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
27
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
28
+#include "libtorch_stable/quantization/marlin/dequant.h"
29
+#include "libtorch_stable/quantization/marlin/marlin_mma.h"
30
#include "core/scalar_type.hpp"
31
32
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh
+1
-1
@@ -6,7 +6,7 @@
6
7
#include <iostream>
8
9
-#include "quantization/marlin/marlin_dtypes.cuh"
9
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
10
using marlin::MarlinScalarType2;
11
12
namespace allspark {
csrc/libtorch_stable/quantization/marlin/.gitignore
renamed
csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu
renamed
+43
-37
@@ -1,6 +1,13 @@
1
#include "marlin.cuh"
2
3
-#include "core/registration.h"
3
+#include <torch/csrc/stable/accelerator.h>
4
+#include <torch/csrc/stable/library.h>
5
+#include <torch/csrc/stable/ops.h>
6
+#include <torch/csrc/stable/tensor.h>
7
+#include <torch/headeronly/core/ScalarType.h>
8
+#include <torch/headeronly/util/Exception.h>
9
+
10
+#include "libtorch_stable/torch_utils.h"
11
12
namespace marlin {
13
@@ -218,56 +225,55 @@ __global__ void awq_marlin_repack_kernel(
225
b_q_weight_ptr, out_ptr, size_k, size_n); \
226
}
227
221
-torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
222
- int64_t size_n, int64_t num_bits,
223
- bool is_a_8bit) {
228
+torch::stable::Tensor awq_marlin_repack(torch::stable::Tensor& b_q_weight,
229
+ int64_t size_k, int64_t size_n,
230
+ int64_t num_bits, bool is_a_8bit) {
231
// Verify compatibility with marlin tile of 16x64
225
- TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
226
- " is not divisible by tile_k_size = ", marlin::tile_k_size);
227
- TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
228
- " is not divisible by tile_n_size = ", marlin::tile_n_size);
232
+ STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
233
+ " is not divisible by tile_k_size = ", marlin::tile_k_size);
234
+ STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
235
+ " is not divisible by tile_n_size = ", marlin::tile_n_size);
236
230
- TORCH_CHECK(num_bits == 4 || num_bits == 8,
231
- "num_bits must be 4 or 8. Got = ", num_bits);
237
+ STD_TORCH_CHECK(num_bits == 4 || num_bits == 8,
238
+ "num_bits must be 4 or 8. Got = ", num_bits);
239
int const pack_factor = 32 / num_bits;
240
241
// Verify B
235
- TORCH_CHECK(b_q_weight.size(0) == size_k,
236
- "b_q_weight.size(0) = ", b_q_weight.size(0),
237
- " is not size_k = ", size_k);
238
- TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1),
239
- "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
240
- ", size_n = ", size_n, ", pack_factor = ", pack_factor);
242
+ STD_TORCH_CHECK(b_q_weight.size(0) == size_k,
243
+ "b_q_weight.size(0) = ", b_q_weight.size(0),
244
+ " is not size_k = ", size_k);
245
+ STD_TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1),
246
+ "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
247
+ ", size_n = ", size_n, ", pack_factor = ", pack_factor);
248
249
// Verify device and strides
243
- TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
244
- TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
245
- TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt");
250
+ STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU");
251
+ STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
252
+ STD_TORCH_CHECK(
253
+ b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int,
254
+ "b_q_weight type is not kInt");
255
+
256
+ const int32_t device_index = b_q_weight.get_device_index();
257
+ torch::stable::accelerator::DeviceGuard device_guard(device_index);
258
+ const cudaStream_t stream = get_current_cuda_stream(device_index);
259
260
// Alloc buffers
248
- const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
249
- auto options = torch::TensorOptions()
250
- .dtype(b_q_weight.dtype())
251
- .device(b_q_weight.device());
252
- torch::Tensor out = torch::empty(
261
+ torch::stable::Tensor out = torch::stable::empty(
262
{size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
254
- options);
263
+ b_q_weight.scalar_type(), std::nullopt, b_q_weight.device());
264
265
// Get ptrs
266
uint32_t const* b_q_weight_ptr =
258
- reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
259
- uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
267
+ reinterpret_cast<uint32_t const*>(b_q_weight.const_data_ptr());
268
+ uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.mutable_data_ptr());
269
261
- // Get dev info
262
- int dev = b_q_weight.get_device();
263
- cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
270
int blocks;
265
- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
271
+ cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index);
272
273
int max_shared_mem = 0;
274
cudaDeviceGetAttribute(&max_shared_mem,
269
- cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
270
- TORCH_CHECK(max_shared_mem > 0);
275
+ cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index);
276
+ STD_TORCH_CHECK(max_shared_mem > 0);
277
278
if (false) {
279
}
@@ -276,13 +282,13 @@ torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
282
CALL_IF(4, true)
283
CALL_IF(8, true)
284
else {
279
- TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
280
- ", is_a_8bit = ", is_a_8bit);
285
+ STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
286
+ ", is_a_8bit = ", is_a_8bit);
287
}
288
289
return out;
290
}
291
286
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
287
- m.impl("awq_marlin_repack", &awq_marlin_repack);
292
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
293
+ m.impl("awq_marlin_repack", TORCH_BOX(&awq_marlin_repack));
294
}
csrc/libtorch_stable/quantization/marlin/dequant.h
renamed
csrc/libtorch_stable/quantization/marlin/generate_kernels.py
renamed
+1
-1
@@ -303,7 +303,7 @@ def generate_new_kernels():
303
if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
304
kernel_selector_str += (
305
"else if (a_type == vllm::kFE4M3fn)\n"
306
- " TORCH_CHECK(false, "
306
+ " STD_TORCH_CHECK(false, "
307
'"marlin kernel with fp8 activation is not built.");'
308
)
309
csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu
renamed
+50
-41
@@ -1,6 +1,13 @@
1
#include "marlin.cuh"
2
3
-#include "core/registration.h"
3
+#include <torch/csrc/stable/accelerator.h>
4
+#include <torch/csrc/stable/library.h>
5
+#include <torch/csrc/stable/ops.h>
6
+#include <torch/csrc/stable/tensor.h>
7
+#include <torch/headeronly/core/ScalarType.h>
8
+#include <torch/headeronly/util/Exception.h>
9
+
10
+#include "libtorch_stable/torch_utils.h"
11
12
namespace marlin {
13
@@ -275,64 +282,66 @@ __global__ void gptq_marlin_repack_kernel(
282
b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \
283
}
284
278
-torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
279
- int64_t size_k, int64_t size_n,
280
- int64_t num_bits, bool is_a_8bit) {
285
+torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight,
286
+ torch::stable::Tensor& perm,
287
+ int64_t size_k, int64_t size_n,
288
+ int64_t num_bits, bool is_a_8bit) {
289
// Verify compatibility with marlin tile of 16x64
282
- TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
283
- " is not divisible by tile_k_size = ", marlin::tile_k_size);
284
- TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
285
- " is not divisible by tile_n_size = ", marlin::tile_n_size);
290
+ STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
291
+ " is not divisible by tile_k_size = ", marlin::tile_k_size);
292
+ STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
293
+ " is not divisible by tile_n_size = ", marlin::tile_n_size);
294
287
- TORCH_CHECK(num_bits == 4 || num_bits == 8,
288
- "num_bits must be 4 or 8. Got = ", num_bits);
295
+ STD_TORCH_CHECK(num_bits == 4 || num_bits == 8,
296
+ "num_bits must be 4 or 8. Got = ", num_bits);
297
int const pack_factor = 32 / num_bits;
298
299
// Verify B
292
- TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
293
- "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
294
- ", size_k = ", size_k, ", pack_factor = ", pack_factor);
295
- TORCH_CHECK(b_q_weight.size(1) == size_n,
296
- "b_q_weight.size(1) = ", b_q_weight.size(1),
297
- " is not size_n = ", size_n);
300
+ STD_TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
301
+ "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
302
+ ", size_k = ", size_k, ", pack_factor = ", pack_factor);
303
+ STD_TORCH_CHECK(b_q_weight.size(1) == size_n,
304
+ "b_q_weight.size(1) = ", b_q_weight.size(1),
305
+ " is not size_n = ", size_n);
306
307
// Verify device and strides
300
- TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
301
- TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
302
- TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt");
308
+ STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU");
309
+ STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
310
+ STD_TORCH_CHECK(
311
+ b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int,
312
+ "b_q_weight type is not kInt");
313
+
314
+ STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU");
315
+ STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
316
+ STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int,
317
+ "perm type is not at::kInt");
318
304
- TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
305
- TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
306
- TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt");
319
+ const int32_t device_index = b_q_weight.get_device_index();
320
+ torch::stable::accelerator::DeviceGuard device_guard(device_index);
321
+ const cudaStream_t stream = get_current_cuda_stream(device_index);
322
323
// Alloc buffers
309
- const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
310
- auto options = torch::TensorOptions()
311
- .dtype(b_q_weight.dtype())
312
- .device(b_q_weight.device());
313
- torch::Tensor out = torch::empty(
324
+ torch::stable::Tensor out = torch::stable::empty(
325
{size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
315
- options);
326
+ b_q_weight.scalar_type(), std::nullopt, b_q_weight.device());
327
328
// Detect if there is act_order
329
bool has_perm = perm.size(0) != 0;
330
331
// Get ptrs
332
uint32_t const* b_q_weight_ptr =
322
- reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
323
- uint32_t const* perm_ptr = reinterpret_cast<uint32_t const*>(perm.data_ptr());
324
- uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
333
+ reinterpret_cast<uint32_t const*>(b_q_weight.const_data_ptr());
334
+ uint32_t const* perm_ptr =
335
+ reinterpret_cast<uint32_t const*>(perm.const_data_ptr());
336
+ uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.mutable_data_ptr());
337
326
- // Get dev info
327
- int dev = b_q_weight.get_device();
328
- cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
338
int blocks;
330
- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
339
+ cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index);
340
341
int max_shared_mem = 0;
342
cudaDeviceGetAttribute(&max_shared_mem,
334
- cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
335
- TORCH_CHECK(max_shared_mem > 0);
343
+ cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index);
344
+ STD_TORCH_CHECK(max_shared_mem > 0);
345
346
if (false) {
347
}
@@ -345,13 +354,13 @@ torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
354
CALL_IF(8, false, true)
355
356
else {
348
- TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
349
- ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
357
+ STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
358
+ ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
359
}
360
361
return out;
362
}
363
355
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
356
- m.impl("gptq_marlin_repack", &gptq_marlin_repack);
364
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
365
+ m.impl("gptq_marlin_repack", TORCH_BOX(&gptq_marlin_repack));
366
}
csrc/libtorch_stable/quantization/marlin/kernel.h
renamed
csrc/libtorch_stable/quantization/marlin/marlin.cu
renamed
+240
-205
@@ -24,7 +24,15 @@
24
#endif
25
26
#include "kernel.h"
27
-#include "core/registration.h"
27
+
28
+#include <torch/csrc/stable/accelerator.h>
29
+#include <torch/csrc/stable/library.h>
30
+#include <torch/csrc/stable/ops.h>
31
+#include <torch/csrc/stable/tensor.h>
32
+#include <torch/headeronly/core/ScalarType.h>
33
+#include <torch/headeronly/util/Exception.h>
34
+
35
+#include "libtorch_stable/torch_utils.h"
36
37
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
38
static_assert(std::is_same<scalar_t, half>::value || \
@@ -46,19 +54,22 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
54
55
} // namespace marlin
56
49
-torch::Tensor marlin_gemm(
50
- torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
51
- torch::Tensor& b_q_weight,
52
- std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
53
- std::optional<torch::Tensor> const& b_zeros_or_none,
54
- std::optional<torch::Tensor> const& g_idx_or_none,
55
- std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
56
- vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
57
- int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
58
- bool is_zp_float) {
59
- TORCH_CHECK_NOT_IMPLEMENTED(false,
60
- "marlin_gemm(..) requires CUDA_ARCH >= 7.5");
61
- return torch::empty({1, 1});
57
+torch::stable::Tensor marlin_gemm(
58
+ torch::stable::Tensor& a, std::optional<torch::stable::Tensor> c_or_none,
59
+ torch::stable::Tensor& b_q_weight,
60
+ std::optional<torch::stable::Tensor> const& b_bias_or_none,
61
+ torch::stable::Tensor& b_scales,
62
+ std::optional<torch::stable::Tensor> const& a_scales_or_none,
63
+ std::optional<torch::stable::Tensor> const& global_scale_or_none,
64
+ std::optional<torch::stable::Tensor> const& b_zeros_or_none,
65
+ std::optional<torch::stable::Tensor> const& g_idx_or_none,
66
+ std::optional<torch::stable::Tensor> const& perm_or_none,
67
+ torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id,
68
+ int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full,
69
+ bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
70
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
71
+ "marlin_gemm(..) requires CUDA_ARCH >= 7.5");
72
+ return torch::stable::empty({1, 1});
73
}
74
75
#else
@@ -323,18 +334,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
334
int thread_n_init, int sms, bool use_atomic_add,
335
bool use_fp32_reduce, bool is_zp_float) {
336
bool is_a_8bit = a_type.size_bits() == 8;
326
- TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
327
- ", ", prob_n, ", ", prob_k, "]");
337
+ STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [",
338
+ prob_m, ", ", prob_n, ", ", prob_k, "]");
339
340
int group_blocks = 0;
341
if (has_act_order) {
342
if (is_k_full) {
332
- TORCH_CHECK(group_size != -1);
343
+ STD_TORCH_CHECK(group_size != -1);
344
group_blocks = group_size / 16;
334
- TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
335
- " is not divisible by group_blocks = ", group_blocks);
345
+ STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
346
+ " is not divisible by group_blocks = ", group_blocks);
347
} else {
337
- TORCH_CHECK(group_size == 0);
348
+ STD_TORCH_CHECK(group_size == 0);
349
group_blocks = 0;
350
}
351
} else {
@@ -342,8 +353,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
353
group_blocks = -1;
354
} else {
355
group_blocks = group_size / 16;
345
- TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
346
- " is not divisible by group_blocks = ", group_blocks);
356
+ STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
357
+ " is not divisible by group_blocks = ", group_blocks);
358
}
359
}
360
@@ -384,25 +395,25 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
395
int max_shared_mem = 0;
396
cudaDeviceGetAttribute(&max_shared_mem,
397
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
387
- TORCH_CHECK(max_shared_mem > 0);
398
+ STD_TORCH_CHECK(max_shared_mem > 0);
399
400
int major_capability, minor_capability;
401
cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
402
dev);
403
cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
404
dev);
394
- TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
395
- "marlin kernel only support Turing or newer GPUs.");
405
+ STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
406
+ "marlin kernel only support Turing or newer GPUs.");
407
int stages = 4;
408
if (major_capability == 7 && minor_capability == 5) {
409
stages = 2;
399
- TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
400
- "Turing only support FP16 or INT8 activation.");
410
+ STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
411
+ "Turing only support FP16 or INT8 activation.");
412
}
413
if (a_type == vllm::kFE4M3fn) {
403
- TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
404
- "FP8 only support Ada Lovelace or newer GPUs.");
405
- TORCH_CHECK(
414
+ STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
415
+ "FP8 only support Ada Lovelace or newer GPUs.");
416
+ STD_TORCH_CHECK(
417
major_capability * 10 + minor_capability == 89 ||
418
major_capability == 12,
419
"Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
@@ -432,10 +443,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
443
if (thread_k != -1 && thread_n != -1) {
444
thread_tfg = thread_config_t{thread_k, thread_n, default_threads};
445
exec_cfg = exec_config_t{1, thread_tfg};
435
- TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
436
- " is not divisible by thread_n = ", thread_n);
437
- TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
438
- " is not divisible by thread_k = ", thread_k);
446
+ STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
447
+ " is not divisible by thread_n = ", thread_n);
448
+ STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
449
+ " is not divisible by thread_k = ", thread_k);
450
} else {
451
// Auto config
452
exec_cfg = determine_exec_config(
@@ -474,7 +485,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
485
int thread_k_blocks = thread_k / 16;
486
int thread_n_blocks = thread_n / 16;
487
477
- TORCH_CHECK(
488
+ STD_TORCH_CHECK(
489
is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n,
490
prob_k, num_bits, group_size, has_act_order, is_k_full,
491
has_zp, is_zp_float, is_a_8bit, stages,
@@ -495,14 +506,15 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
506
num_threads, is_zp_float, stages);
507
508
if (kernel == MarlinDefault) {
498
- TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
499
- ", ", prob_k, "]", ", has_act_order = ", has_act_order,
500
- ", num_groups = ", num_groups, ", group_size = ", group_size,
501
- ", prob_m_split = ", prob_m_split,
502
- ", thread_m_blocks = ", thread_m_blocks,
503
- ", thread_n_blocks = ", thread_n_blocks,
504
- ", thread_k_blocks = ", thread_k_blocks,
505
- ", num_threads = ", num_threads, ", num_bits = ", num_bits);
509
+ STD_TORCH_CHECK(
510
+ false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ",
511
+ prob_k, "]", ", has_act_order = ", has_act_order,
512
+ ", num_groups = ", num_groups, ", group_size = ", group_size,
513
+ ", prob_m_split = ", prob_m_split,
514
+ ", thread_m_blocks = ", thread_m_blocks,
515
+ ", thread_n_blocks = ", thread_n_blocks,
516
+ ", thread_k_blocks = ", thread_k_blocks,
517
+ ", num_threads = ", num_threads, ", num_bits = ", num_bits);
518
}
519
520
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
@@ -530,71 +542,76 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
542
543
} // namespace marlin
544
533
-torch::Tensor marlin_gemm(
534
- torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
535
- torch::Tensor& b_q_weight,
536
- std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
537
- std::optional<torch::Tensor> const& a_scales_or_none,
538
- std::optional<torch::Tensor> const& global_scale_or_none,
539
- std::optional<torch::Tensor> const& b_zeros_or_none,
540
- std::optional<torch::Tensor> const& g_idx_or_none,
541
- std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
542
- vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
543
- int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
544
- bool is_zp_float) {
545
+torch::stable::Tensor marlin_gemm(
546
+ torch::stable::Tensor& a, std::optional<torch::stable::Tensor> c_or_none,
547
+ torch::stable::Tensor& b_q_weight,
548
+ std::optional<torch::stable::Tensor> const& b_bias_or_none,
549
+ torch::stable::Tensor& b_scales,
550
+ std::optional<torch::stable::Tensor> const& a_scales_or_none,
551
+ std::optional<torch::stable::Tensor> const& global_scale_or_none,
552
+ std::optional<torch::stable::Tensor> const& b_zeros_or_none,
553
+ std::optional<torch::stable::Tensor> const& g_idx_or_none,
554
+ std::optional<torch::stable::Tensor> const& perm_or_none,
555
+ torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id,
556
+ int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full,
557
+ bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
558
vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
559
547
- auto c_dtype = a.dtype();
548
- if (a.scalar_type() == at::ScalarType::Half) {
560
+ auto c_scalar_type = a.scalar_type();
561
+ if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
562
a_type_id = vllm::kFloat16.id();
563
c_type_id = vllm::kFloat16.id();
551
- } else if (a.scalar_type() == at::ScalarType::BFloat16) {
564
+ } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
565
a_type_id = vllm::kBFloat16.id();
566
c_type_id = vllm::kBFloat16.id();
567
} else {
555
- c_dtype = b_scales.dtype();
556
- if (b_scales.scalar_type() == at::ScalarType::Half) {
568
+ c_scalar_type = b_scales.scalar_type();
569
+ if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
570
c_type_id = vllm::kFloat16.id();
558
- } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
571
+ } else if (b_scales.scalar_type() ==
572
+ torch::headeronly::ScalarType::BFloat16) {
573
c_type_id = vllm::kBFloat16.id();
574
} else {
575
c_type_id = vllm::kBFloat16.id();
576
563
- TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
564
- torch::Tensor c = c_or_none.value();
565
- c_dtype = c.dtype();
577
+ STD_TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
578
+ torch::stable::Tensor c = c_or_none.value();
579
+ c_scalar_type = c.scalar_type();
580
567
- if (c.scalar_type() == at::ScalarType::Half) {
581
+ if (c.scalar_type() == torch::headeronly::ScalarType::Half) {
582
c_type_id = vllm::kFloat16.id();
569
- } else if (c.scalar_type() == at::ScalarType::BFloat16) {
583
+ } else if (c.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
584
c_type_id = vllm::kBFloat16.id();
585
} else {
572
- TORCH_CHECK(false, "unsupported c dtype");
586
+ STD_TORCH_CHECK(false, "unsupported c dtype");
587
}
588
}
589
576
- if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
590
+ if (a.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn) {
591
a_type_id = vllm::kFE4M3fn.id();
578
- } else if (a.scalar_type() == at::ScalarType::Char) {
592
+ } else if (a.scalar_type() == torch::headeronly::ScalarType::Char) {
593
a_type_id = vllm::kS8.id();
594
} else {
581
- TORCH_CHECK(false, "unsupported `a` scalar_type");
595
+ STD_TORCH_CHECK(false, "unsupported `a` scalar_type");
596
}
597
}
598
599
s_type_id = c_type_id;
600
if (b_type_id == vllm::kFE2M1f.id()) {
587
- if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
601
+ if (b_scales.scalar_type() ==
602
+ torch::headeronly::ScalarType::Float8_e4m3fn) {
603
s_type_id = vllm::kFE4M3fn.id();
589
- } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
604
+ } else if (b_scales.scalar_type() ==
605
+ torch::headeronly::ScalarType::Float8_e8m0fnu) {
606
s_type_id = vllm::kFE8M0fnu.id();
607
} else {
592
- TORCH_CHECK(false,
593
- "When b_type = float4_e2m1f, b_scale scalar type must be",
594
- "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
608
+ STD_TORCH_CHECK(
609
+ false, "When b_type = float4_e2m1f, b_scale scalar type must be",
610
+ "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
611
}
612
} else if (b_type_id == vllm::kFE4M3fn.id() &&
597
- b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
613
+ b_scales.scalar_type() ==
614
+ torch::headeronly::ScalarType::Float8_e8m0fnu) {
615
s_type_id = vllm::kFE8M0fnu.id();
616
}
617
@@ -606,54 +623,58 @@ torch::Tensor marlin_gemm(
623
int pack_factor = 32 / b_type.size_bits();
624
625
// Verify A
609
- TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
610
- ", size_m = ", size_m);
611
- TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1),
612
- ", size_k = ", size_k);
626
+ STD_TORCH_CHECK(a.size(0) == size_m,
627
+ "Shape mismatch: a.size(0) = ", a.size(0),
628
+ ", size_m = ", size_m);
629
+ STD_TORCH_CHECK(a.size(1) == size_k,
630
+ "Shape mismatch: a.size(1) = ", a.size(1),
631
+ ", size_k = ", size_k);
632
633
// Verify B
615
- TORCH_CHECK(
634
+ STD_TORCH_CHECK(
635
size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k,
636
" is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
618
- TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0),
619
- "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
620
- ", size_k = ", size_k,
621
- ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
622
- TORCH_CHECK(
637
+ STD_TORCH_CHECK(
638
+ (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0),
639
+ "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
640
+ ", size_k = ", size_k,
641
+ ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
642
+ STD_TORCH_CHECK(
643
b_q_weight.size(1) % MARLIN_NAMESPACE_NAME::tile_size == 0,
644
"b_q_weight.size(1) = ", b_q_weight.size(1),
645
" is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
646
int actual_size_n =
647
(b_q_weight.size(1) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor;
628
- TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
629
- ", actual_size_n = ", actual_size_n);
648
+ STD_TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
649
+ ", actual_size_n = ", actual_size_n);
650
651
// Verify device and strides
632
- TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
633
- TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1");
652
+ STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
653
+ STD_TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1");
654
// We use int4 (16 bytes) to load A, so A must aligned to 16 bytes
635
- TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8");
636
- TORCH_CHECK(((uint64_t)a.data_ptr()) % 16 == 0, "A must aligned to 16 bytes");
655
+ STD_TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8");
656
+ STD_TORCH_CHECK(((uint64_t)a.const_data_ptr()) % 16 == 0,
657
+ "A must aligned to 16 bytes");
658
638
- TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
639
- TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
659
+ STD_TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
660
+ STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
661
641
- TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
642
- TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
662
+ STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
663
+ STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
664
644
- torch::Tensor a_scales;
645
- auto options = torch::TensorOptions().dtype(c_dtype).device(a.device());
646
- auto options_fp32 =
647
- torch::TensorOptions().dtype(at::kFloat).device(a.device());
665
+ torch::stable::Tensor a_scales;
666
+ const auto device = a.device();
667
668
if (a_scales_or_none.has_value()) {
669
a_scales = a_scales_or_none.value();
651
- TORCH_CHECK(a_type.size_bits() == 8,
652
- "a_scales can only be used for 8bit activation.");
670
+ STD_TORCH_CHECK(a_type.size_bits() == 8,
671
+ "a_scales can only be used for 8bit activation.");
672
} else {
654
- a_scales = torch::empty({0}, options_fp32);
655
- TORCH_CHECK(a_type.size_bits() != 8,
656
- "the a_scales parameter must be passed for 8bit activation.");
673
+ a_scales = torch::stable::empty({0}, torch::headeronly::ScalarType::Float,
674
+ std::nullopt, device);
675
+ STD_TORCH_CHECK(
676
+ a_type.size_bits() != 8,
677
+ "the a_scales parameter must be passed for 8bit activation.");
678
}
679
680
// thread_k: `k` size of a thread_tile in `weights` (can usually be left as
@@ -664,84 +685,93 @@ torch::Tensor marlin_gemm(
685
int thread_n = -1;
686
// sms: number of SMs to use for the kernel
687
int sms = -1;
667
- cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
688
+ const int32_t device_index = a.get_device_index();
689
+ cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_index);
690
691
// Alloc buffers
670
- const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
671
- torch::Tensor c;
692
+ torch::stable::accelerator::DeviceGuard device_guard(device_index);
693
+ torch::stable::Tensor c;
694
if (c_or_none.has_value()) {
695
c = c_or_none.value();
674
- TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
675
- TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
676
- TORCH_CHECK(c.size(0) == size_m, "Shape mismatch: c.size(0) = ", c.size(0),
677
- ", size_m = ", size_m);
678
- TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1),
679
- ", size_n = ", size_n);
696
+ STD_TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
697
+ STD_TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
698
+ STD_TORCH_CHECK(c.size(0) == size_m,
699
+ "Shape mismatch: c.size(0) = ", c.size(0),
700
+ ", size_m = ", size_m);
701
+ STD_TORCH_CHECK(c.size(1) == size_n,
702
+ "Shape mismatch: c.size(1) = ", c.size(1),
703
+ ", size_n = ", size_n);
704
} else {
681
- c = torch::empty({size_m, size_n}, options);
705
+ c = torch::stable::empty({size_m, size_n}, c_scalar_type, std::nullopt,
706
+ device);
707
}
708
if (size_m == 0) return c;
709
710
// Alloc C tmp buffer that is going to be used for the global reduce
686
- torch::Tensor c_tmp;
711
+ torch::stable::Tensor c_tmp;
712
if (use_fp32_reduce) {
713
int max_m_block_size = (size_m + 16 - 1) / 16 * 16;
714
max_m_block_size = min(max_m_block_size, 64);
715
int max_c_tmp_size =
716
sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n;
692
- c_tmp = torch::empty({max_c_tmp_size}, options_fp32);
717
+ c_tmp = torch::stable::empty({max_c_tmp_size},
718
+ torch::headeronly::ScalarType::Float,
719
+ std::nullopt, device);
720
} else {
694
- c_tmp = torch::empty({0}, options_fp32);
721
+ c_tmp = torch::stable::empty({0}, torch::headeronly::ScalarType::Float,
722
+ std::nullopt, device);
723
}
724
725
// Detect groupsize and act_order
726
int num_groups = -1;
727
int group_size = -1;
728
701
- int rank = b_scales.sizes().size();
702
- TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2");
703
- TORCH_CHECK(b_scales.size(1) == size_n, "b_scales dim 1 = ", b_scales.size(1),
704
- " is not size_n = ", size_n);
729
+ int rank = b_scales.dim();
730
+ STD_TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2");
731
+ STD_TORCH_CHECK(b_scales.size(1) == size_n,
732
+ "b_scales dim 1 = ", b_scales.size(1),
733
+ " is not size_n = ", size_n);
734
num_groups = b_scales.size(0);
735
707
- torch::Tensor g_idx, perm, a_tmp;
736
+ torch::stable::Tensor g_idx, perm, a_tmp;
737
if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
738
g_idx = g_idx_or_none.value();
739
perm = perm_or_none.value();
740
712
- TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
713
- TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
714
- TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
715
- TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
741
+ STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
742
+ STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
743
+ STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
744
+ STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
745
746
// Verify g_idx and perm
718
- TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
719
- (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
720
- "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
721
- " and perm.size(-1) = ", perm.size(-1),
722
- ", where size_k = ", size_k);
747
+ STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
748
+ (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
749
+ "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
750
+ " and perm.size(-1) = ", perm.size(-1),
751
+ ", where size_k = ", size_k);
752
} else {
724
- g_idx = torch::empty({0}, options);
725
- perm = torch::empty({0}, options);
726
- a_tmp = torch::empty({0}, options);
753
+ g_idx = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
754
+ perm = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
755
+ a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
756
}
757
bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
758
759
if (has_act_order) {
731
- a_tmp = torch::empty({size_m, size_k}, options);
760
+ a_tmp = torch::stable::empty({size_m, size_k}, c_scalar_type, std::nullopt,
761
+ device);
762
if (is_k_full) {
733
- TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
734
- TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
735
- ", is not divisible by num_groups = ", num_groups);
763
+ STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
764
+ STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
765
+ ", is not divisible by num_groups = ", num_groups);
766
group_size = size_k / num_groups;
767
} else {
768
group_size = 0;
769
}
770
771
} else {
742
- a_tmp = torch::empty({0}, options);
772
+ a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
773
if (num_groups > 1) {
744
- TORCH_CHECK(
774
+ STD_TORCH_CHECK(
775
size_k % num_groups == 0, "size_k = ", size_k,
776
", is not divisible by b_scales.size(0) = ", b_scales.size(0));
777
group_size = size_k / num_groups;
@@ -750,109 +780,114 @@ torch::Tensor marlin_gemm(
780
}
781
}
782
753
- torch::Tensor global_scale;
783
+ torch::stable::Tensor global_scale;
784
if (global_scale_or_none.has_value()) {
785
global_scale = global_scale_or_none.value();
756
- TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
757
- "global_scale can only be used for nvfp4 format.");
786
+ STD_TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
787
+ "global_scale can only be used for nvfp4 format.");
788
} else {
759
- global_scale = torch::empty({0}, options_fp32);
760
- TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
761
- "the global_scale parameter must be passed for nvfp4 format.");
789
+ global_scale = torch::stable::empty(
790
+ {0}, torch::headeronly::ScalarType::Float, std::nullopt, device);
791
+ STD_TORCH_CHECK(
792
+ !(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
793
+ "the global_scale parameter must be passed for nvfp4 format.");
794
}
795
796
bool has_bias = b_bias_or_none.has_value();
765
- torch::Tensor b_bias;
797
+ torch::stable::Tensor b_bias;
798
if (has_bias) {
799
b_bias = b_bias_or_none.value();
768
- TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
769
- TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
770
- TORCH_CHECK(b_bias.size(0) == size_n, "b_bias.size(0) != size_n");
771
- TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1");
800
+ STD_TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
801
+ STD_TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
802
+ STD_TORCH_CHECK(b_bias.size(0) == size_n, "b_bias.size(0) != size_n");
803
+ STD_TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1");
804
} else {
773
- b_bias = torch::empty({0}, options);
805
+ b_bias = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
806
}
807
776
- torch::Tensor b_zeros;
808
+ torch::stable::Tensor b_zeros;
809
if (b_zeros_or_none.has_value()) {
810
b_zeros = b_zeros_or_none.value();
779
- TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
780
- TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
811
+ STD_TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
812
+ STD_TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
813
} else {
782
- b_zeros = torch::empty({0}, options);
814
+ b_zeros = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
815
}
816
bool has_zp = b_zeros.size(-1) > 0;
817
if (has_zp) {
786
- TORCH_CHECK(
818
+ STD_TORCH_CHECK(
819
b_type == vllm::kU4 || b_type == vllm::kU8,
820
"b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str());
821
} else {
790
- TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
791
- b_type == vllm::kS4 || b_type == vllm::kS8 ||
792
- b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
793
- "b_type must be uint4b8, uint8b128, int4, int8, "
794
- "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
795
- b_type.str());
822
+ STD_TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
823
+ b_type == vllm::kS4 || b_type == vllm::kS8 ||
824
+ b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
825
+ "b_type must be uint4b8, uint8b128, int4, int8, "
826
+ "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
827
+ b_type.str());
828
}
829
830
if (has_zp && is_zp_float) {
799
- TORCH_CHECK(a.scalar_type() == at::ScalarType::Half,
800
- "Computation type must be float16 (half) when using float zero "
801
- "points.");
831
+ STD_TORCH_CHECK(
832
+ a.scalar_type() == torch::headeronly::ScalarType::Half,
833
+ "Computation type must be float16 (half) when using float zero "
834
+ "points.");
835
}
836
837
// Verify b_zeros
838
if (has_zp) {
806
- int rank = b_zeros.sizes().size();
807
- TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2");
839
+ int rank = b_zeros.dim();
840
+ STD_TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2");
841
if (is_zp_float) {
809
- TORCH_CHECK(b_zeros.size(1) == size_n,
810
- "b_zeros dim 1 = ", b_zeros.size(1),
811
- " is not size_n = ", size_n);
812
- TORCH_CHECK(num_groups == b_zeros.size(0),
813
- "b_zeros dim 0 = ", b_zeros.size(0),
814
- " is not num_groups = ", num_groups);
815
- TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
842
+ STD_TORCH_CHECK(b_zeros.size(1) == size_n,
843
+ "b_zeros dim 1 = ", b_zeros.size(1),
844
+ " is not size_n = ", size_n);
845
+ STD_TORCH_CHECK(num_groups == b_zeros.size(0),
846
+ "b_zeros dim 0 = ", b_zeros.size(0),
847
+ " is not num_groups = ", num_groups);
848
+ STD_TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
849
} else {
817
- TORCH_CHECK(b_zeros.size(0) == num_groups,
818
- "b_zeros dim 0 = ", b_zeros.size(0),
819
- " is not num_groups = ", num_groups);
820
- TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor,
821
- "b_zeros dim 1 = ", b_zeros.size(1),
822
- " is not size_n / pack_factor = ", size_n / pack_factor);
850
+ STD_TORCH_CHECK(b_zeros.size(0) == num_groups,
851
+ "b_zeros dim 0 = ", b_zeros.size(0),
852
+ " is not num_groups = ", num_groups);
853
+ STD_TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor,
854
+ "b_zeros dim 1 = ", b_zeros.size(1),
855
+ " is not size_n / pack_factor = ", size_n / pack_factor);
856
}
857
}
858
859
// Verify workspace size
827
- TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
828
- "size_n = ", size_n, ", is not divisible by min_thread_n = ",
829
- MARLIN_NAMESPACE_NAME::min_thread_n);
860
+ STD_TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
861
+ "size_n = ", size_n, ", is not divisible by min_thread_n = ",
862
+ MARLIN_NAMESPACE_NAME::min_thread_n);
863
864
int min_workspace_size = sms;
832
- TORCH_CHECK(workspace.numel() >= min_workspace_size,
833
- "workspace.numel = ", workspace.numel(),
834
- " is below min_workspace_size = ", min_workspace_size);
835
-
836
- int dev = a.get_device();
837
-
838
- TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
839
- "scalar type of a_scales must be float");
840
- TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
841
- "scalar type of global_scale must be float");
865
+ STD_TORCH_CHECK(workspace.numel() >= min_workspace_size,
866
+ "workspace.numel = ", workspace.numel(),
867
+ " is below min_workspace_size = ", min_workspace_size);
868
+
869
+ STD_TORCH_CHECK(
870
+ a_scales.scalar_type() == torch::headeronly::ScalarType::Float,
871
+ "scalar type of a_scales must be float");
872
+ STD_TORCH_CHECK(
873
+ global_scale.scalar_type() == torch::headeronly::ScalarType::Float,
874
+ "scalar type of global_scale must be float");
875
if (a_type.size_bits() == 16) {
843
- TORCH_CHECK(
876
+ STD_TORCH_CHECK(
877
a.scalar_type() == c.scalar_type(),
878
"scalar type of a must be the same with c for 16 bit activation");
879
}
880
881
marlin::marlin_mm(
849
- a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(),
850
- b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(),
851
- global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(),
852
- perm.data_ptr(), a_tmp.data_ptr(), size_m, size_n, size_k, a.stride(0),
853
- workspace.data_ptr(), a_type, b_type, c_type, s_type, has_bias,
854
- has_act_order, is_k_full, has_zp, num_groups, group_size, dev,
855
- at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
882
+ a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(),
883
+ c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(),
884
+ a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(),
885
+ global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(),
886
+ g_idx.mutable_data_ptr(), perm.mutable_data_ptr(),
887
+ a_tmp.mutable_data_ptr(), size_m, size_n, size_k, a.stride(0),
888
+ workspace.mutable_data_ptr(), a_type, b_type, c_type, s_type, has_bias,
889
+ has_act_order, is_k_full, has_zp, num_groups, group_size, device_index,
890
+ get_current_cuda_stream(device_index), thread_k, thread_n, sms,
891
use_atomic_add, use_fp32_reduce, is_zp_float);
892
893
return c;
@@ -860,6 +895,6 @@ torch::Tensor marlin_gemm(
895
896
#endif
897
863
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
864
- m.impl("marlin_gemm", &marlin_gemm);
898
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
899
+ m.impl("marlin_gemm", TORCH_BOX(&marlin_gemm));
900
}
csrc/libtorch_stable/quantization/marlin/marlin.cuh
renamed
-8
@@ -2,14 +2,6 @@
2
3
#ifndef _marlin_cuh
4
#define _marlin_cuh
5
- // These torch headers are only needed by non-stable callers (e.g. ops.cu).
6
- // Guard them so that stable ABI targets can still include marlin.cuh
7
- // for Vec, constants, and cp_async helpers without pulling in torch/all.h.
8
- #ifndef TORCH_TARGET_VERSION
9
- #include <torch/all.h>
10
- #include <ATen/cuda/CUDAContext.h>
11
- #include <c10/cuda/CUDAGuard.h>
12
- #endif
5
#include <cuda.h>
6
#include <cuda_fp16.h>
7
#include <cuda_runtime.h>
csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh
renamed
csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu
new
+118
@@ -0,0 +1,118 @@
1
+
2
+#include "marlin.cuh"
3
+
4
+#include <torch/csrc/stable/accelerator.h>
5
+#include <torch/csrc/stable/library.h>
6
+#include <torch/csrc/stable/ops.h>
7
+#include <torch/csrc/stable/tensor.h>
8
+#include <torch/headeronly/core/ScalarType.h>
9
+#include <torch/headeronly/util/Exception.h>
10
+
11
+#include "libtorch_stable/torch_utils.h"
12
+
13
+// for only non-zp format (like gptq)
14
+__global__ void marlin_int4_fp8_preprocess_kernel_without_zp(
15
+ // qweight: (size_k * size_n // 8,)
16
+ const int32_t* __restrict__ qweight,
17
+ // output: same shape with qweight
18
+ int32_t* __restrict__ output) {
19
+ int32_t val = qweight[blockIdx.x * 32 + threadIdx.x];
20
+ int32_t new_val = 0;
21
+
22
+#pragma unroll
23
+ for (int32_t i = 0; i < 8; i++) {
24
+ int32_t single_val = val & 0xF;
25
+ single_val = single_val >= 8 ? single_val - 8 : 15 - single_val;
26
+ new_val |= single_val << (i * 4);
27
+ val >>= 4;
28
+ }
29
+
30
+ output[blockIdx.x * 32 + threadIdx.x] = new_val;
31
+}
32
+
33
+// for awq format only (with zp and with awq weight layout)
34
+__global__ void marlin_int4_fp8_preprocess_kernel_awq(
35
+ // AWQ qweight: (size_k, size_n // 8)
36
+ const int32_t* __restrict__ qweight,
37
+ // output: same shape with qweight
38
+ int32_t* __restrict__ output,
39
+ // AWQ zeros: (size_k // group_size, size_n // 8)
40
+ const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
41
+ int32_t group_size) {
42
+ int32_t val =
43
+ qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
44
+ int32_t zero =
45
+ qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
46
+ blockIdx.y];
47
+ int32_t new_val = 0;
48
+
49
+#pragma unroll
50
+ for (int32_t i = 0; i < 8; i++) {
51
+ int32_t single_val = val & 0xF;
52
+ int32_t single_zero = zero & 0xF;
53
+
54
+ single_val =
55
+ single_val >= single_zero ? single_val - single_zero : 15 - single_val;
56
+ new_val |= single_val << (i * 4);
57
+ val >>= 4;
58
+ zero >>= 4;
59
+ }
60
+
61
+ output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
62
+}
63
+
64
+torch::stable::Tensor marlin_int4_fp8_preprocess(
65
+ torch::stable::Tensor& qweight,
66
+ std::optional<torch::stable::Tensor> qzeros_or_none, bool inplace) {
67
+ STD_TORCH_CHECK(qweight.is_cuda(), "qweight is not on GPU");
68
+ STD_TORCH_CHECK(qweight.scalar_type() == torch::headeronly::ScalarType::Int,
69
+ "qweight.dtype != torch.int32");
70
+
71
+ const int32_t device_index = qweight.get_device_index();
72
+ torch::stable::accelerator::DeviceGuard device_guard(device_index);
73
+ const cudaStream_t stream = get_current_cuda_stream(device_index);
74
+
75
+ torch::stable::Tensor output =
76
+ inplace ? qweight : torch::stable::empty_like(qweight);
77
+
78
+ if (!qzeros_or_none.has_value()) {
79
+ STD_TORCH_CHECK(qweight.numel() * 8 % 256 == 0,
80
+ "qweight.numel() * 8 % 256 != 0");
81
+
82
+ int blocks = qweight.numel() * 8 / 256;
83
+ marlin_int4_fp8_preprocess_kernel_without_zp<<<blocks, 32, 0, stream>>>(
84
+ reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
85
+ reinterpret_cast<int32_t*>(output.mutable_data_ptr()));
86
+ } else {
87
+ int32_t size_k = qweight.size(0);
88
+ int32_t size_n = qweight.size(1) * 8;
89
+ torch::stable::Tensor qzeros = qzeros_or_none.value();
90
+
91
+ STD_TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0");
92
+ STD_TORCH_CHECK(qzeros.is_cuda(), "qzeros is not on GPU");
93
+ STD_TORCH_CHECK(qzeros.scalar_type() == torch::headeronly::ScalarType::Int,
94
+ "qweight.dtype != torch.int32");
95
+ STD_TORCH_CHECK(qzeros.get_device_index() == device_index,
96
+ "qzeros is not on the same device with qweight");
97
+
98
+ int32_t group_size = qweight.size(0) / qzeros.size(0);
99
+ STD_TORCH_CHECK(qweight.size(1) == qzeros.size(1),
100
+ "qweight.size(1) != qzeros.size(1)");
101
+ STD_TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0,
102
+ "qweight.size(0) % qzeros.size(0) != 0");
103
+ STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
104
+
105
+ dim3 blocks(size_k / 32, size_n / 8);
106
+ marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
107
+ reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
108
+ reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
109
+ reinterpret_cast<const int32_t*>(qzeros.const_data_ptr()), size_n,
110
+ size_k, group_size);
111
+ }
112
+
113
+ return output;
114
+}
115
+
116
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
117
+ m.impl("marlin_int4_fp8_preprocess", TORCH_BOX(&marlin_int4_fp8_preprocess));
118
+}
csrc/libtorch_stable/quantization/marlin/marlin_mma.h
renamed
csrc/libtorch_stable/quantization/marlin/marlin_template.h
renamed
csrc/libtorch_stable/torch_bindings.cpp
+29
@@ -33,6 +33,35 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
33
34
// TODO: Remove this once ROCm upgrade to torch 2.11.
35
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
36
+
37
+ // Marlin GEMM
38
+ ops.def(
39
+ "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
40
+ "Tensor? b_bias_or_none,Tensor b_scales, "
41
+ "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
42
+ "Tensor? "
43
+ "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, "
44
+ "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
45
+ "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
46
+ // conditionally compiled so impl registrations are in source file
47
+
48
+ // gptq_marlin repack from GPTQ.
49
+ ops.def(
50
+ "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
51
+ "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
52
+ // conditionally compiled so impl registrations are in source file
53
+
54
+ // awq_marlin repack from AWQ.
55
+ ops.def(
56
+ "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
57
+ "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
58
+ // conditionally compiled so impl registrations are in source file
59
+
60
+ // preprocess W-int4A-fp8 weight for marlin kernel
61
+ ops.def(
62
+ "marlin_int4_fp8_preprocess(Tensor qweight, "
63
+ "Tensor? qzeros_or_none, bool inplace) -> Tensor");
64
+ // conditionally compiled so impl registrations are in source file
65
#endif
66
67
#ifndef USE_ROCM
csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu
deleted
-106
@@ -1,106 +0,0 @@
1
-
2
-
3
-#include "marlin.cuh"
4
-
5
-#include "core/registration.h"
6
-
7
-// for only non-zp format (like gptq)
8
-__global__ void marlin_int4_fp8_preprocess_kernel_without_zp(
9
- // qweight: (size_k * size_n // 8,)
10
- const int32_t* __restrict__ qweight,
11
- // output: same shape with qweight
12
- int32_t* __restrict__ output) {
13
- int32_t val = qweight[blockIdx.x * 32 + threadIdx.x];
14
- int32_t new_val = 0;
15
-
16
-#pragma unroll
17
- for (int32_t i = 0; i < 8; i++) {
18
- int32_t single_val = val & 0xF;
19
- single_val = single_val >= 8 ? single_val - 8 : 15 - single_val;
20
- new_val |= single_val << (i * 4);
21
- val >>= 4;
22
- }
23
-
24
- output[blockIdx.x * 32 + threadIdx.x] = new_val;
25
-}
26
-
27
-// for awq format only (with zp and with awq weight layout)
28
-__global__ void marlin_int4_fp8_preprocess_kernel_awq(
29
- // AWQ qweight: (size_k, size_n // 8)
30
- const int32_t* __restrict__ qweight,
31
- // output: same shape with qweight
32
- int32_t* __restrict__ output,
33
- // AWQ zeros: (size_k // group_size, size_n // 8)
34
- const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
35
- int32_t group_size) {
36
- int32_t val =
37
- qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
38
- int32_t zero =
39
- qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
40
- blockIdx.y];
41
- int32_t new_val = 0;
42
-
43
-#pragma unroll
44
- for (int32_t i = 0; i < 8; i++) {
45
- int32_t single_val = val & 0xF;
46
- int32_t single_zero = zero & 0xF;
47
-
48
- single_val =
49
- single_val >= single_zero ? single_val - single_zero : 15 - single_val;
50
- new_val |= single_val << (i * 4);
51
- val >>= 4;
52
- zero >>= 4;
53
- }
54
-
55
- output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
56
-}
57
-
58
-torch::Tensor marlin_int4_fp8_preprocess(
59
- torch::Tensor& qweight, std::optional<torch::Tensor> qzeros_or_none,
60
- bool inplace) {
61
- TORCH_CHECK(qweight.device().is_cuda(), "qweight is not on GPU");
62
- TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Int,
63
- "qweight.dtype != torch.int32");
64
-
65
- const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight));
66
-
67
- torch::Tensor output = inplace ? qweight : torch::empty_like(qweight);
68
-
69
- if (!qzeros_or_none.has_value()) {
70
- TORCH_CHECK(qweight.numel() * 8 % 256 == 0,
71
- "qweight.numel() * 8 % 256 != 0");
72
-
73
- int blocks = qweight.numel() * 8 / 256;
74
- marlin_int4_fp8_preprocess_kernel_without_zp<<<blocks, 32>>>(
75
- (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr());
76
- } else {
77
- int32_t size_k = qweight.size(0);
78
- int32_t size_n = qweight.size(1) * 8;
79
- torch::Tensor qzeros = qzeros_or_none.value();
80
-
81
- TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0");
82
- TORCH_CHECK(qzeros.device().is_cuda(), "qzeros is not on GPU");
83
- TORCH_CHECK(qzeros.scalar_type() == at::ScalarType::Int,
84
- "qweight.dtype != torch.int32");
85
- TORCH_CHECK(device_of(qweight) == device_of(qzeros),
86
- "qzeros is not on the same device with qweight");
87
-
88
- int32_t group_size = qweight.size(0) / qzeros.size(0);
89
- TORCH_CHECK(qweight.size(1) == qzeros.size(1),
90
- "qweight.size(1) != qzeros.size(1)");
91
- TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0,
92
- "qweight.size(0) % qzeros.size(0) != 0");
93
- TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
94
-
95
- dim3 blocks(size_k / 32, size_n / 8);
96
- marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32>>>(
97
- (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr(),
98
- (const int32_t*)qzeros.data_ptr(), size_n, size_k, group_size);
99
- }
100
-
101
- return output;
102
-}
103
-
104
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
105
- m.impl("marlin_int4_fp8_preprocess", &marlin_int4_fp8_preprocess);
106
-}
csrc/torch_bindings.cpp
-29
@@ -101,35 +101,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
101
") -> Tensor");
102
// conditionally compiled so impl registration is in source file
103
104
- // Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4).
105
- ops.def(
106
- "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
107
- "Tensor? b_bias_or_none,Tensor b_scales, "
108
- "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
109
- "Tensor? "
110
- "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, "
111
- "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
112
- "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
113
- // conditionally compiled so impl registration is in source file
114
-
115
- // gptq_marlin repack from GPTQ.
116
- ops.def(
117
- "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
118
- "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
119
- // conditionally compiled so impl registrations are in source file
120
-
121
- // awq_marlin repack from AWQ.
122
- ops.def(
123
- "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
124
- "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
125
- // conditionally compiled so impl registrations are in source file
126
-
127
- // preprocess W-int4A-fp8 weight for marlin kernel
128
- ops.def(
129
- "marlin_int4_fp8_preprocess(Tensor qweight, "
130
- "Tensor? qzeros_or_none, bool inplace) -> Tensor");
131
- // conditionally compiled so impl registrations are in source file
132
-
104
#endif
105
}
106