From 2c9d68f80bd20af7af1bc539c044a4aebbe52f8e Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Tue, 8 Sep 2026 15:45:12 -0400 Subject: [PATCH] [Quant][Kernel] Remove GPTQ Group/Dynamic Activation Ordering (#54809) Signed-off-by: Roderick-Wu Signed-off-by: Roderick Wu Signed-off-by: mgoin Co-authored-by: mgoin --- CMakeLists.txt | 1 - benchmarks/kernels/benchmark_machete.py | 7 +- benchmarks/kernels/benchmark_marlin.py | 118 ++-- csrc/cpu/cpu_wna16.cpp | 152 ++--- csrc/cpu/torch_bindings.cpp | 3 +- .../moe/marlin_moe_wna16/generate_kernels.py | 4 +- .../moe/marlin_moe_wna16/kernel.h | 29 +- .../moe/marlin_moe_wna16/marlin_template.h | 440 ++++---------- .../moe/marlin_moe_wna16/ops.cu | 348 +++--------- csrc/libtorch_stable/moe/torch_bindings.cpp | 5 +- csrc/libtorch_stable/ops.h | 10 +- csrc/libtorch_stable/permute_cols.cu | 94 --- .../quantization/gptq/q_gemm.cu | 535 ++++-------------- .../quantization/marlin/generate_kernels.py | 4 +- .../quantization/marlin/gptq_marlin_repack.cu | 164 ++---- .../quantization/marlin/kernel.h | 21 +- .../quantization/marlin/marlin.cu | 301 +++------- .../quantization/marlin/marlin_template.h | 450 ++++----------- csrc/libtorch_stable/torch_bindings.cpp | 15 +- csrc/rocm/ops.h | 5 +- csrc/rocm/q_gemm_rdna3.cu | 103 ++-- csrc/rocm/q_gemm_rdna3_wmma.cu | 225 +++----- csrc/rocm/torch_bindings.cpp | 4 +- tests/compile/fullgraph/test_full_graph.py | 2 +- tests/kernels/core/test_permute_cols.py | 21 - .../moe/test_marlin_vs_trtllm_mxint4.py | 9 +- tests/kernels/moe/test_moe.py | 98 +--- .../quantization/test_allspark_gemm.py | 2 +- tests/kernels/quantization/test_gptq.py | 12 +- .../kernels/quantization/test_marlin_gemm.py | 85 +-- .../quantization/test_marlin_tile_padding.py | 28 +- .../quantization/test_rdna3_compile_guards.py | 3 - .../quantization/test_rdna3_moe_w4a16.py | 8 +- .../kernels/quantization/test_rdna3_w4a16.py | 2 - .../test_rdna3_w4a16_selection.py | 2 - .../quantization/test_rdna_hybrid_w4a16.py | 5 - .../kernels/quantization/test_triton_w4a16.py | 6 - .../test_w4a16_kernel_selection.py | 2 - tests/lora/conftest.py | 5 - tests/lora/test_quant_model.py | 149 +++-- .../model_loader/test_reload.py | 111 +--- tests/model_executor/test_weight_utils.py | 11 + tests/models/quantization/test_gptq_marlin.py | 8 +- tests/models/transformers/test_backend.py | 2 +- tests/quantization/test_auto_gptq.py | 52 +- tests/quantization/test_auto_round.py | 7 +- tests/quantization/test_compressed_tensors.py | 83 +-- tests/quantization/test_cpu_wna16.py | 3 +- tests/quantization/test_cutlass_w4a16.py | 14 +- tests/quantization/test_gptq_dynamic.py | 57 +- tests/quantization/test_lm_head.py | 3 +- tests/quantization/test_moe_wna16.py | 54 +- tests/weight_loading/models-large.txt | 4 +- tests/weight_loading/models.txt | 7 +- tests/weight_loading/test_weight_loading.py | 5 +- vllm/_custom_ops.py | 51 +- .../linear/mixed_precision/MPLinearKernel.py | 9 +- .../linear/mixed_precision/allspark.py | 5 +- .../kernels/linear/mixed_precision/conch.py | 8 +- .../kernels/linear/mixed_precision/cpu.py | 7 +- .../kernels/linear/mixed_precision/cutlass.py | 5 +- .../kernels/linear/mixed_precision/exllama.py | 34 +- .../kernels/linear/mixed_precision/humming.py | 2 - .../kernels/linear/mixed_precision/machete.py | 35 +- .../kernels/linear/mixed_precision/marlin.py | 53 +- .../linear/mixed_precision/rdna3_w4a16.py | 36 +- .../mixed_precision/rdna_hybrid_w4a16.py | 11 +- .../linear/mixed_precision/triton_w4a16.py | 17 +- .../kernels/linear/mixed_precision/xpu.py | 21 +- .../linear/mixed_precision/zentorch.py | 13 +- .../layers/fused_moe/experts/marlin_moe.py | 72 --- .../layers/fused_moe/oracle/int_wna16.py | 155 +---- .../layers/fused_moe/routed_experts.py | 56 +- .../layers/quantization/auto_awq.py | 23 - .../layers/quantization/auto_gptq.py | 99 +--- .../layers/quantization/base_config.py | 7 + .../compressed_tensors/compressed_tensors.py | 16 +- .../compressed_tensors_moe_w4a16_flydsl.py | 47 -- .../compressed_tensors_moe_w4a8_fp8.py | 1 - .../compressed_tensors_moe_wna16.py | 131 +---- .../compressed_tensors_moe_wna16_rdna3.py | 5 +- .../schemes/compressed_tensors_w4a8_fp8.py | 7 +- .../schemes/compressed_tensors_w4a8_int.py | 2 - .../schemes/compressed_tensors_wNa16.py | 20 +- .../schemes/compressed_tensors_wNa4.py | 3 +- .../schemes/compressed_tensors_wNa8.py | 3 +- .../schemes/compressed_tensors_wNa8o8.py | 3 +- .../inc/schemes/inc_w4a8_linear.py | 2 +- .../inc/schemes/inc_wna16_linear.py | 13 +- .../layers/quantization/moe_wna16.py | 34 +- .../layers/quantization/utils/gptq_utils.py | 34 ++ .../layers/quantization/utils/marlin_utils.py | 35 +- .../quantization/utils/marlin_utils_fp4.py | 23 +- .../quantization/utils/marlin_utils_fp8.py | 17 +- .../quantization/utils/marlin_utils_test.py | 16 +- .../layers/quantization/utils/quant_utils.py | 62 +- .../layers/vocab_parallel_embedding.py | 3 +- .../model_loader/weight_utils.py | 2 - vllm/model_executor/models/utils.py | 3 +- 99 files changed, 1182 insertions(+), 3917 deletions(-) delete mode 100644 csrc/libtorch_stable/permute_cols.cu delete mode 100644 tests/kernels/core/test_permute_cols.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 91a1b19841..2592c0eef7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -436,7 +436,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" - "csrc/libtorch_stable/permute_cols.cu" "csrc/libtorch_stable/quantization/gptq/q_gemm.cu" "csrc/libtorch_stable/pos_encoding_kernels.cu" "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu" diff --git a/benchmarks/kernels/benchmark_machete.py b/benchmarks/kernels/benchmark_machete.py index 4e6f098665..3ec8b2ecde 100644 --- a/benchmarks/kernels/benchmark_machete.py +++ b/benchmarks/kernels/benchmark_machete.py @@ -220,10 +220,8 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable: bt.w_g_s, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.group_size ) - sort_indices = torch.empty(0, dtype=torch.int, device=device) - g_idx = torch.empty(0, dtype=torch.int, device=device) w_q = ops.gptq_marlin_repack( - bt.w_q, sort_indices, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits + bt.w_q, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits ) if bt.a.dtype.is_floating_point: @@ -240,14 +238,11 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable: a_scales=None, global_scale=None, b_zeros=w_zp, - g_idx=g_idx, - perm=sort_indices, workspace=workspace.scratch, b_q_type=bt.wtype, size_m=bt.a.shape[0], size_n=bt.w_ref.shape[1], size_k=bt.w_ref.shape[0], - is_k_full=True, is_zp_float=False, ) else: diff --git a/benchmarks/kernels/benchmark_marlin.py b/benchmarks/kernels/benchmark_marlin.py index c0019a51cd..751211993a 100644 --- a/benchmarks/kernels/benchmark_marlin.py +++ b/benchmarks/kernels/benchmark_marlin.py @@ -32,7 +32,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( gptq_pack, gptq_quantize_weights, quantize_weights, - sort_weights, ) from vllm.scalar_type import ScalarType, scalar_types from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -40,15 +39,10 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser DEFAULT_MODELS = ["meta-llama/Llama-2-7b-hf/TP1"] DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] -ACT_ORDER_OPTS = [False, True] -K_FULL_OPTS = [False, True] - def bench_run( results: list[benchmark.Measurement], model: str, - act_order: bool, - is_k_full: bool, quant_type: ScalarType, group_size: int, size_m: int, @@ -56,38 +50,33 @@ def bench_run( size_n: int, ): label = "Quant Matmul" - sub_label = "{}, act={} k_full={}, q={}, g={}, MKN=({}x{}x{})".format( - model, act_order, is_k_full, str(quant_type), group_size, size_m, size_k, size_n + sub_label = "{}, q={}, g={}, MKN=({}x{}x{})".format( + model, str(quant_type), group_size, size_m, size_k, size_n ) print(f"Testing: {sub_label}") a = torch.randn(size_m, size_k).to(torch.half).cuda() b = torch.rand(size_k, size_n).to(torch.half).cuda() has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] - if act_order and (group_size == -1 or group_size == size_k or has_zp): - return if size_k % group_size != 0: return repack_supported = group_size in MARLIN_SUPPORTED_GROUP_SIZES allspark_supported = ( - quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES - and group_size == -1 - and not act_order - and is_k_full + quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES and group_size == -1 ) def gen_marlin_params(): # Marlin quant - marlin_g_idx = marlin_sort_indices = marlin_zp = marlin_s2 = None + marlin_zp = marlin_s2 = None if quant_type == scalar_types.float4_e2m1f: - if group_size != 16 or act_order: + if group_size != 16: return marlin_w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_fp4_like( b.T, group_size ) elif quant_type == scalar_types.float8_e4m3fn: - if group_size not in [-1, 128] or act_order: + if group_size not in [-1, 128]: return marlin_w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(b.T, group_size) elif group_size == 16: @@ -97,8 +86,8 @@ def bench_run( b, quant_type, group_size ) else: - marlin_w_ref, marlin_q_w, marlin_s, marlin_g_idx, marlin_sort_indices, _ = ( - marlin_quantize(b, quant_type, group_size, act_order) + marlin_w_ref, marlin_q_w, marlin_s = marlin_quantize( + b, quant_type, group_size ) return ( marlin_w_ref, @@ -106,25 +95,14 @@ def bench_run( marlin_s, marlin_s2, marlin_zp, - marlin_g_idx, - marlin_sort_indices, ) def gen_repack_params(): q_w_gptq = None - repack_sort_indices = None if repack_supported: - (w_ref, q_w, s, g_idx, rand_perm) = gptq_quantize_weights( - b, quant_type, group_size, act_order - ) + _, q_w, _ = gptq_quantize_weights(b, quant_type, group_size) q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n) - - # For act_order, sort the "weights" and "g_idx" - # so that group ids are increasing - repack_sort_indices = torch.empty(0, dtype=torch.int, device=b.device) - if act_order: - (q_w, g_idx, repack_sort_indices) = sort_weights(q_w, g_idx) - return q_w_gptq, repack_sort_indices + return q_w_gptq def gen_allspark_params(): qw_reorder = s_reorder = zp_reorder = sm_count = sm_version = ( @@ -161,10 +139,8 @@ def bench_run( marlin_s, marlin_s2, marlin_zp, - marlin_g_idx, - marlin_sort_indices, ) = gen_marlin_params() - q_w_gptq, repack_sort_indices = gen_repack_params() + q_w_gptq = gen_repack_params() qw_reorder, s_reorder, zp_reorder, sm_count, sm_version, CUBLAS_M_THRESHOLD = ( gen_allspark_params() ) @@ -188,13 +164,9 @@ def bench_run( "marlin_s": marlin_s, "marlin_s2": marlin_s2, "marlin_zp": marlin_zp, - "marlin_g_idx": marlin_g_idx, - "marlin_sort_indices": marlin_sort_indices, "marlin_workspace": marlin_workspace, - "is_k_full": is_k_full, # GPTQ params "q_w_gptq": q_w_gptq, - "repack_sort_indices": repack_sort_indices, # AllSpark W8A16 params "qw_reorder": qw_reorder, "s_reorder": s_reorder, @@ -226,7 +198,7 @@ def bench_run( results.append( benchmark.Timer( - stmt="output = marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501 + stmt="output = marlin_gemm(a, None, marlin_q_w, None, marlin_s, None, marlin_s2, marlin_zp, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, False, False, False)", # noqa: E501 globals=globals, label=label, sub_label=sub_label, @@ -236,7 +208,7 @@ def bench_run( results.append( benchmark.Timer( - stmt="output = marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501 + stmt="output = marlin_gemm(a, None, marlin_q_w, None, marlin_s, None, marlin_s2, marlin_zp, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, False, True, False)", # noqa: E501 globals=globals, label=label, sub_label=sub_label, @@ -247,7 +219,7 @@ def bench_run( if repack_supported: results.append( benchmark.Timer( - stmt="q_res = gptq_marlin_repack(q_w_gptq, repack_sort_indices, size_k, size_n, quant_type.size_bits)", # noqa: E501 + stmt="q_res = gptq_marlin_repack(q_w_gptq, size_k, size_n, quant_type.size_bits)", # noqa: E501 globals=globals, label=label, sub_label=sub_label, @@ -284,61 +256,39 @@ def main(args): if len(args.limit_n) > 0 and size_n not in args.limit_n: continue - for act_order in ACT_ORDER_OPTS: + for quant_type in query_marlin_supported_quant_types(): if ( - len(args.limit_act_order) > 0 - and act_order not in args.limit_act_order + len(args.limit_num_bits) > 0 + and quant_type.size_bits not in args.limit_num_bits ): continue - for is_k_full in K_FULL_OPTS: + for group_size in ( + MARLIN_SUPPORTED_GROUP_SIZES + FP4_MARLIN_SUPPORTED_GROUP_SIZES + ): if ( - len(args.limit_k_full) > 0 - and is_k_full not in args.limit_k_full + len(args.limit_group_size) > 0 + and group_size not in args.limit_group_size ): continue - for quant_type in query_marlin_supported_quant_types(): - if ( - len(args.limit_num_bits) > 0 - and quant_type.size_bits not in args.limit_num_bits - ): - continue - - for group_size in ( - MARLIN_SUPPORTED_GROUP_SIZES - + FP4_MARLIN_SUPPORTED_GROUP_SIZES - ): - if ( - len(args.limit_group_size) > 0 - and group_size not in args.limit_group_size - ): - continue - - # For act_order, the group_size must be less than - # size_k - if act_order and (group_size == size_k or group_size == -1): - continue - - for size_m in args.batch_sizes: - bench_run( - results, - model, - act_order, - is_k_full, - quant_type, - group_size, - size_m, - size_k, - size_n, - ) + for size_m in args.batch_sizes: + bench_run( + results, + model, + quant_type, + group_size, + size_m, + size_k, + size_n, + ) compare = benchmark.Compare(results) compare.print() # For quick benchmarking use: -# python benchmark_marlin.py --batch-sizes 1 16 32 --limit-k 4096 --limit-n 4096 --limit-group-size 128 --limit-num-bits 4 --limit-act-order 0 --limit-k-full 1 # noqa E501 +# python benchmark_marlin.py --batch-sizes 1 16 32 --limit-k 4096 --limit-n 4096 --limit-group-size 128 --limit-num-bits 4 # noqa E501 # if __name__ == "__main__": parser = FlexibleArgumentParser( @@ -358,8 +308,6 @@ if __name__ == "__main__": parser.add_argument("--limit-n", nargs="+", type=int, default=[]) parser.add_argument("--limit-group-size", nargs="+", type=int, default=[]) parser.add_argument("--limit-num-bits", nargs="+", type=int, default=[]) - parser.add_argument("--limit-act-order", nargs="+", type=int, default=[]) - parser.add_argument("--limit-k-full", nargs="+", type=int, default=[]) args = parser.parse_args() main(args) diff --git a/csrc/cpu/cpu_wna16.cpp b/csrc/cpu/cpu_wna16.cpp index ae7aef74c4..d3e42660f2 100644 --- a/csrc/cpu/cpu_wna16.cpp +++ b/csrc/cpu/cpu_wna16.cpp @@ -37,7 +37,7 @@ namespace { using cpu_utils::ISA; using cpu_utils::VecTypeTrait; -template +template class Dequantizer4b { public: constexpr static int32_t pack_num = 32 / 4; @@ -47,9 +47,9 @@ class Dequantizer4b { static void dequant(int32_t* __restrict__ q_weight, scalar_t* __restrict__ weight, scalar_t* __restrict__ scales, - int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx, - const int64_t scales_stride, const int64_t zeros_stride, - const int32_t k_size, const int32_t group_size) { + int32_t* __restrict__ zeros, const int64_t scales_stride, + const int64_t zeros_stride, const int32_t k_size, + const int32_t group_size) { vec_op::FP32Vec16 lut; if constexpr (has_zp) { // AWQ @@ -81,30 +81,15 @@ class Dequantizer4b { vec_op::FP32Vec16 wb_0(qwb_0, lut); vec_op::FP32Vec16 wb_1(qwb_1, lut); - if constexpr (!use_desc_act) { - if (group_counter == 0) { - scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale)); - scale_1 = vec_op::FP32Vec16(scale_0); - curr_scale += scales_stride; + if (group_counter == 0) { + scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale)); + scale_1 = vec_op::FP32Vec16(scale_0); + curr_scale += scales_stride; - if constexpr (has_zp) { - zero_0 = vec_op::FP32Vec16(*curr_zeros, lut); - zero_1 = vec_op::FP32Vec16(zero_0); - curr_zeros += zeros_stride / 2; - } - } - } else { - int32_t g_idx_0 = g_idx[k_idx]; - int32_t g_idx_1 = g_idx[k_idx + 1]; - scale_0 = vec_op::FP32Vec16( - scalar_vec_t(curr_scale + g_idx_0 * scales_stride)); - scale_1 = vec_op::FP32Vec16( - scalar_vec_t(curr_scale + g_idx_1 * scales_stride)); if constexpr (has_zp) { - zero_0 = vec_op::FP32Vec16(*(curr_zeros + g_idx_0 * zeros_stride / 2), - lut); - zero_1 = vec_op::FP32Vec16(*(curr_zeros + g_idx_1 * zeros_stride / 2), - lut); + zero_0 = vec_op::FP32Vec16(*curr_zeros, lut); + zero_1 = vec_op::FP32Vec16(zero_0); + curr_zeros += zeros_stride / 2; } } @@ -130,11 +115,9 @@ class Dequantizer4b { // update curr_q_weight += 2; curr_weight += 32; - if constexpr (!use_desc_act) { - group_counter += 2; - if (group_counter == group_size) { - group_counter = 0; - } + group_counter += 2; + if (group_counter == group_size) { + group_counter = 0; } } } @@ -145,11 +128,10 @@ template void cpu_gemm_wna16_impl( scalar_t* __restrict__ input, int32_t* __restrict__ q_weight, scalar_t* __restrict__ output, scalar_t* __restrict__ scales, - int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx, - scalar_t* __restrict__ bias, const int32_t m_size, const int32_t n_size, - const int32_t k_size, const int64_t input_stride, - const int64_t output_stride, const int64_t scales_group_stride, - const int64_t zeros_group_stride, const int32_t group_num, + int32_t* __restrict__ zeros, scalar_t* __restrict__ bias, + const int32_t m_size, const int32_t n_size, const int32_t k_size, + const int64_t input_stride, const int64_t output_stride, + const int64_t scales_group_stride, const int64_t zeros_group_stride, const int32_t group_size, const int64_t pack_factor) { constexpr int32_t gemm_n_tile_size = gemm_t::NSize; constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; @@ -229,7 +211,7 @@ void cpu_gemm_wna16_impl( int32_t* __restrict__ curr_zeros = zeros + n_start_idx / pack_factor; for (int32_t block_idx = 0; block_idx < n_block_num; ++block_idx) { dequantizer_t::dequant(curr_q_weight, curr_b_buffer, curr_scales, - curr_zeros, g_idx, scales_group_stride, + curr_zeros, scales_group_stride, zeros_group_stride, k_size, group_size); // if (block_idx == 0 && n_start_idx == 0) { @@ -296,8 +278,7 @@ void cpu_gemm_wna16( const torch::Tensor& scales, // [group_num, N] const std::optional& zeros, // [group_num, N / pack_factor], packed as int32 - const std::optional& g_idx, // [K] - const std::optional& bias, // [N] + const std::optional& bias, // [N] const int64_t pack_factor, const std::string& isa_hint) { using cpu_utils::ISA; TORCH_CHECK_EQ(pack_factor, 8); // only supports 4bits @@ -314,8 +295,6 @@ void cpu_gemm_wna16( const int64_t output_m_stride = output.stride(0); bool has_zp = zeros.has_value(); - bool use_desc_act = g_idx.has_value(); - TORCH_CHECK(!(has_zp && use_desc_act)); ISA isa = [&]() { if (isa_hint == "amx") { @@ -331,109 +310,72 @@ void cpu_gemm_wna16( int32_t* zeros_ptr = has_zp ? zeros->data_ptr() : nullptr; const int64_t zeros_group_stride = has_zp ? zeros->stride(0) : 0; - int32_t* g_idx_ptr = use_desc_act ? g_idx->data_ptr() : nullptr; VLLM_DISPATCH_16B_TYPES(input.scalar_type(), "cpu_gemm_wna16", [&]() { if (isa == ISA::AMX) { using gemm_t = cpu_micro_gemm::MicroGemm; if (has_zp) { - using dequantizer_t = Dequantizer4b; + using dequantizer_t = Dequantizer4b; cpu_gemm_wna16_impl( input.data_ptr(), q_weight.data_ptr(), output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); return; } - if (use_desc_act) { - using dequantizer_t = Dequantizer4b; + { + using dequantizer_t = Dequantizer4b; cpu_gemm_wna16_impl( input.data_ptr(), q_weight.data_ptr(), output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); - return; - } else { - using dequantizer_t = Dequantizer4b; - cpu_gemm_wna16_impl( - input.data_ptr(), q_weight.data_ptr(), - output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); return; } } else if (isa == ISA::VEC) { using gemm_t = cpu_micro_gemm::MicroGemm; if (has_zp) { - using dequantizer_t = Dequantizer4b; + using dequantizer_t = Dequantizer4b; cpu_gemm_wna16_impl( input.data_ptr(), q_weight.data_ptr(), output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); return; } - if (use_desc_act) { - using dequantizer_t = Dequantizer4b; + { + using dequantizer_t = Dequantizer4b; cpu_gemm_wna16_impl( input.data_ptr(), q_weight.data_ptr(), output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); - return; - } else { - using dequantizer_t = Dequantizer4b; - cpu_gemm_wna16_impl( - input.data_ptr(), q_weight.data_ptr(), - output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); return; } } else if (isa == ISA::RVV) { using gemm_t = cpu_micro_gemm::MicroGemm; if (has_zp) { - using dequantizer_t = Dequantizer4b; + using dequantizer_t = Dequantizer4b; cpu_gemm_wna16_impl( input.data_ptr(), q_weight.data_ptr(), output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); return; } - if (use_desc_act) { - using dequantizer_t = Dequantizer4b; + { + using dequantizer_t = Dequantizer4b; cpu_gemm_wna16_impl( input.data_ptr(), q_weight.data_ptr(), output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); - return; - } else { - using dequantizer_t = Dequantizer4b; - cpu_gemm_wna16_impl( - input.data_ptr(), q_weight.data_ptr(), - output.data_ptr(), scales.data_ptr(), zeros_ptr, - g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, - a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, - scales_group_stride, zeros_group_stride, group_num, group_size, - pack_factor); + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); return; } } diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 7906ae0e1b..ea3469b30a 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -226,7 +226,6 @@ void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); } void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight, torch::Tensor& output, const torch::Tensor& scales, const std::optional& zeros, - const std::optional& g_idx, const std::optional& bias, const int64_t pack_factor, const std::string& isa_hint); @@ -671,7 +670,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__) ops.def( "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, " - "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt " + "Tensor scales, Tensor? zeros, Tensor? bias, SymInt " "pack_factor, str isa_hint) -> ()"); ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16); #endif diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py b/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py index 64b47b607b..6948a98a37 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py @@ -75,14 +75,14 @@ QUANT_CONFIGS = [ "b_type": "kU4B8", "thread_configs": THREAD_CONFIGS, "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 0, 2, 4, 8], + "group_blocks": [-1, 2, 4, 8], }, # AWQ-INT8 { "b_type": "kU8B128", "thread_configs": THREAD_CONFIGS, "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 0, 2, 4, 8], + "group_blocks": [-1, 2, 4, 8], }, # FP8 { diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h index 783736ab50..25c57f3b2d 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h @@ -7,21 +7,20 @@ #include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" #include "core/scalar_type.hpp" -#define MARLIN_KERNEL_PARAMS \ - const int4 *__restrict__ A, const int4 *__restrict__ B, \ - int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ b_bias_ptr, \ - const float *__restrict__ a_scales_ptr, \ - const int4 *__restrict__ scales_ptr, \ - const float *__restrict__ global_scale_ptr, \ - const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ - const int32_t *__restrict__ sorted_token_ids_ptr, \ - const int32_t *__restrict__ expert_ids_ptr, \ - const int32_t *__restrict__ num_tokens_past_padded_ptr, \ - const float *__restrict__ topk_weights_ptr, int top_k, \ - bool mul_topk_weights, int num_groups, int prob_m, int prob_n, \ - int prob_k, int *locks, bool has_bias, bool use_atomic_add, \ - bool use_fp32_reduce +#define MARLIN_KERNEL_PARAMS \ + const int4 *__restrict__ A, const int4 *__restrict__ B, \ + int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ + const int4 *__restrict__ b_bias_ptr, \ + const float *__restrict__ a_scales_ptr, \ + const int4 *__restrict__ scales_ptr, \ + const float *__restrict__ global_scale_ptr, \ + const int4 *__restrict__ zp_ptr, \ + const int32_t *__restrict__ sorted_token_ids_ptr, \ + const int32_t *__restrict__ expert_ids_ptr, \ + const int32_t *__restrict__ num_tokens_past_padded_ptr, \ + const float *__restrict__ topk_weights_ptr, int top_k, \ + bool mul_topk_weights, int prob_m, int prob_n, int prob_k, int *locks, \ + bool has_bias, bool use_atomic_add, bool use_fp32_reduce namespace MARLIN_NAMESPACE_NAME { template shared // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin( const int4* __restrict__ A, // fp16 input matrix of shape mxk @@ -64,14 +63,12 @@ __global__ void Marlin( // (k/groupsize)xn const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape // (k/groupsize)x(n/pack_factor) - const int* __restrict__ g_idx, // int32 group indices of shape k const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids const int32_t* __restrict__ expert_ids_ptr, // moe expert ids const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens const float* __restrict__ topk_weights_ptr, // moe top weights int top_k, // num of experts per token bool mul_topk_weights, // mul topk weights or not - int num_groups, // number of scale groups per output channel int prob_m, // batch dimension m int prob_n, // output dimension n int prob_k, // reduction dimension k @@ -148,29 +145,6 @@ __device__ inline void sub_zp( frag_b[1] = __hsub2(frag_b[1], zp); } -// Same as above, but for act_order (each K is multiplied individually) -template -__device__ inline void scale4( - typename MarlinScalarType::FragB& frag_b, - typename MarlinScalarType::FragS& frag_s_1, - typename MarlinScalarType::FragS& frag_s_2, - typename MarlinScalarType::FragS& frag_s_3, - typename MarlinScalarType::FragS& frag_s_4, int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - - scalar_t2 s_val_1_2; - s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; - s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; - - scalar_t2 s_val_3_4; - s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; - s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; - - frag_b[0] = __hmul2(frag_b[0], s_val_1_2); - frag_b[1] = __hmul2(frag_b[1], s_val_3_4); -} - // Given 2 floats multiply by 2 scales (halves) template __device__ inline void scale_float( @@ -264,15 +238,12 @@ __global__ void Marlin( // 4bit packed zero-points of shape // (k/groupsize, n/pack_factor) const int4* __restrict__ zp_ptr, - // int32 group indices of shape k - const int* __restrict__ g_idx, const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids const int32_t* __restrict__ expert_ids_ptr, // moe expert ids const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens const float* __restrict__ topk_weights_ptr, // moe top weights int top_k, // num of experts per token bool mul_topk_weights, // mul topk weights or not - int num_groups, // number of scale groups per output channel int prob_m, // batch dimension m int prob_n, // output dimension n int prob_k, // reduction dimension k @@ -369,12 +340,12 @@ __global__ void Marlin( float global_scale_f32 = 1.0f; - constexpr bool has_act_order = group_blocks == 0; + static_assert(group_blocks != 0, + "runtime activation ordering is not supported"); constexpr int pack_factor = 32 / b_type.size_bits(); static_assert(thread_m_blocks == 1 || !m_block_size_8); - const int group_size = - (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups; + const int group_size = group_blocks == -1 ? prob_k : group_blocks * 16; const int scales_expert_stride = prob_n * prob_k / group_size / (is_8bit_scale ? 16 : 8); const int zp_expert_stride = @@ -405,7 +376,7 @@ __global__ void Marlin( int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); - if constexpr (!has_act_order && group_blocks != -1) { + if constexpr (group_blocks != -1) { if (group_blocks >= thread_k_blocks) { // Ensure that the number of tiles in each stripe is a multiple of the // groupsize; this avoids an annoying special case where a stripe starts @@ -551,9 +522,6 @@ __global__ void Marlin( if constexpr (has_zp) { zp_ptr += (expert_id - old_expert_id) * zp_expert_stride; } - if constexpr (has_act_order) { - g_idx += (expert_id - old_expert_id) * prob_k; - } if (has_bias) { b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride; } @@ -694,27 +662,17 @@ __global__ void Marlin( b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - // Scale sizes/strides without act_order + // Scale sizes/strides int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8); constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8); constexpr int s_tb_groups = - !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + group_blocks != -1 && group_blocks < thread_k_blocks ? thread_k_blocks / group_blocks : 1; constexpr int s_sh_stage = s_tb_groups * s_sh_stride; int s_gl_rd_delta = s_gl_stride; - // Scale size/strides with act_order - constexpr int tb_k = 16 * thread_k_blocks; - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; - // constexpr int act_s_row_stride = 1; - // int act_s_col_stride = act_s_row_stride * num_groups; - constexpr int act_s_max_num_groups = 32; - int act_s_col_stride = 1; - int act_s_col_warp_stride = act_s_col_stride * 8; - constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; // Zero-points sizes/strides int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; @@ -750,25 +708,16 @@ __global__ void Marlin( auto b_sh_rd = threadIdx.x * b_thread_vecs; b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); - // For act_order - int slice_k_start = tb_k * slice_row; - int slice_k_finish = slice_k_start + tb_k * slice_iters; - int slice_k_start_shared_fetch = slice_k_start; - int slice_n_offset = act_s_col_tb_stride * slice_col; - - // No act_order int s_gl_rd; - if constexpr (!has_act_order) { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + - s_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + - threadIdx.x / s_sh_stride) + - s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - } + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; } auto s_sh_wr = threadIdx.x; bool s_sh_wr_pred = threadIdx.x < s_sh_stage; @@ -885,10 +834,8 @@ __global__ void Marlin( : (sh_size_b_red_min + sh_bias_size); int4* sh_bias = sh_new + sh_size_b_red_min; - int4* sh_g_idx = sh_new + sh_b_red_bias_size; - int4* sh_zp = sh_g_idx + (stages * g_idx_stage); - constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) - : (stages * s_sh_stage); + int4* sh_zp = sh_new + sh_b_red_bias_size; + constexpr int sh_s_size = stages * s_sh_stage; int4* sh_s = sh_zp + (stages * zp_sh_stage); int4* sh_a = sh_s + sh_s_size; @@ -897,9 +844,8 @@ __global__ void Marlin( I4 frag_b_quant[2][b_thread_vecs]; FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragS frag_s[2][4]; // No act-order + FragS frag_s[2][4]; FragS frag_bias[2][4]; - FragS act_frag_s[2][4][4]; // For act-order int frag_qzp[2][num_ints_per_thread]; // Zero-points FragZP frag_zp; // Zero-points in fp16 FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ @@ -929,42 +875,6 @@ __global__ void Marlin( reinterpret_cast(frag_c)[i] = 0; }; - int sh_first_group_id = -1; - int sh_num_groups = -1; - - auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, - int last_group_id) { - sh_first_group_id = first_group_id; - sh_num_groups = last_group_id - first_group_id + 1; - - if (sh_num_groups > act_s_max_num_groups) { - sh_num_groups = act_s_max_num_groups; - } - - if (sh_first_group_id + sh_num_groups > num_groups) { - sh_num_groups = num_groups - sh_first_group_id; - } - - int row_offset = first_group_id * s_gl_stride; - - if (is_async) { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], - &scales_ptr[row_offset + (i * s_gl_stride) + - slice_n_offset + threadIdx.x]); - } - } - } else { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - sh_s[(i * s_sh_stride) + threadIdx.x] = - scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + - threadIdx.x]; - } - } - } - }; // Asynchronously fetch the next A, B and s tile from global to the next // shared memory pipeline location. auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { @@ -995,44 +905,27 @@ __global__ void Marlin( b_gl_rd += b_gl_rd_delta_o; - if constexpr (has_act_order) { - // Fetch g_idx thread-block portion - int full_pipe = a_off; - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; - if (cur_k < prob_k && cur_k < slice_k_finish) { - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; - int4 const* cur_g_idx_stage_ptr = - reinterpret_cast(&g_idx[cur_k]); - - if (threadIdx.x < g_idx_stage) { - cp_async4_pred(&sh_g_idx_stage[threadIdx.x], - &cur_g_idx_stage_ptr[threadIdx.x]); + // Only fetch scales if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); } + s_gl_rd += s_gl_rd_delta * s_tb_groups; } - } else { - if constexpr (group_blocks != -1) { - int4* sh_s_stage = sh_s + s_sh_stage * pipe; + } - // Only fetch scales if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (s_sh_wr_pred) { - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta * s_tb_groups; - } - } + if constexpr (has_zp && group_blocks != -1) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - if constexpr (has_zp && group_blocks != -1) { - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - // Only fetch zero points if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; + // Only fetch zero points if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (zp_sh_wr_pred) { + cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); } + zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; } } } @@ -1080,150 +973,57 @@ __global__ void Marlin( } }; - bool is_same_group[stages]; - int same_group_id[stages]; - - auto init_same_group = [&](int pipe) { - if constexpr (!has_act_order) { - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - int group_id_1 = sh_g_idx_int_ptr[0]; - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; - - is_same_group[pipe] = group_id_1 == group_id_2; - same_group_id[pipe] = group_id_1; - }; - auto fetch_scales_to_registers = [&](int k, int full_pipe) { int pipe = full_pipe % stages; using IT1 = typename std::conditional_t; using IT0 = typename std::conditional_t; constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); - if constexpr (!has_act_order) { - // No act-order case - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 && dequant_skip_flop) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - } - } else if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0) { - if (k % b_sh_wr_iters == 0) { - int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; - } else { - reinterpret_cast(&frag_s[1])[0] = - reinterpret_cast(&frag_s[0])[0]; - } - } - } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks2; - - int4* sh_s_stage = sh_s + s_sh_stage * pipe; - - if constexpr (!is_8bit_scale) { - reinterpret_cast(&frag_s[k % 2])[0] = - sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; - } else { - reinterpret_cast(&frag_s[k % 2])[0] = - reinterpret_cast( - sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - } else if (group_blocks >= b_sh_wr_iters) { - if constexpr (!is_8bit_scale) { - reinterpret_cast(&frag_s[1])[0] = - reinterpret_cast(&frag_s[0])[0]; - } else { - reinterpret_cast(&frag_s[1])[0] = - reinterpret_cast(&frag_s[0])[0]; - } + if constexpr (group_blocks == -1) { + // Load only when starting a new slice. + if (k == 0 && full_pipe == 0 && dequant_skip_flop) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } else if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0) { + if (k % b_sh_wr_iters == 0) { + int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; } } + } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; - return; - } + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks2; - // Act-order case + int4* sh_s_stage = sh_s + s_sh_stage * pipe; - // Determine K of the "current" thread-block - int cur_k = slice_k_start + tb_k * full_pipe; - if (cur_k >= prob_k || cur_k >= slice_k_finish) { - return; - } - - // Reset (to current thread-block) since we read g_idx portion from the - // shared memory - cur_k = 0; - - // Progress to current iteration - cur_k += k % b_sh_wr_iters; - - // Determine "position" inside the thread-block (based on warp and - // thread-id) - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - int warp_col = warp_id % tb_n_warps; - - cur_k += warp_row * 16 * b_sh_wr_iters; - - auto th_id = threadIdx.x % 32; - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix - - int s_col_shift = - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + - (th_id / 4) * act_s_col_stride; - - if (is_same_group[pipe]) { - if (k % 2 == 0) { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = - sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + - s_col_shift]; + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; } else { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = - *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast( + sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; } - - for (int i = 1; i < 4; i++) { - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } else if (group_blocks >= b_sh_wr_iters) { + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; } - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - constexpr int k_frag_offsets[4] = {0, 1, 8, - 9}; // Tensor core offsets per thread - - #pragma unroll - for (int i = 0; i < 4; i++) { - int actual_k = cur_k + k_frag_offsets[i]; - - int group_id = sh_g_idx_int_ptr[actual_k]; - int rel_group_id = group_id - sh_first_group_id; - - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = - sh_s[rel_group_id * s_sh_stride + s_col_shift]; } }; auto fetch_zp_to_registers = [&](int k, int full_pipe) { - // This code does not handle group_blocks == 0, - // which signifies act_order. - // has_zp implies AWQ, which doesn't have act_order, static_assert(!has_zp || group_blocks != 0); if constexpr (has_zp && !is_zp_float) { @@ -1313,8 +1113,7 @@ __global__ void Marlin( constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; const bool is_new_zp = - (group_blocks == 0) || - ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && + (group_blocks > 0 && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) || (group_blocks == -1 && is_first_matmul_in_slice); if constexpr (has_zp && !is_zp_float) { @@ -1382,15 +1181,8 @@ __global__ void Marlin( sub_zp(frag_b1, frag_zp[j], 1); } - // Apply scale to frag_b0 - if constexpr (has_act_order && !is_a_8bit) { - static_assert(group_blocks != -1); - scale4(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); - scale4(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); - } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && - group_blocks == -1 && !is_a_8bit) { + if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && + group_blocks == -1 && !is_a_8bit) { int idx = (threadIdx.x / 4) % 2; scalar_t2 s2 = Adtype::nums2num2( reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], @@ -1805,7 +1597,7 @@ __global__ void Marlin( // For per-column quantization we finally apply the scale here (only for // 4-bit) - if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit && + if constexpr (group_blocks == -1 && !is_a_8bit && b_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { c_scalar_t2 tmp_scale = s[0]; @@ -1911,15 +1703,6 @@ __global__ void Marlin( #pragma unroll for (int i = 0; i < stages - 1; i++) { - if (has_act_order && i == 0) { - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], - g_idx[last_g_idx]); - } - if constexpr (has_zp && !is_zp_float && group_blocks == -1) { if (i == 0) { fetch_col_zp_to_shared(); @@ -1933,14 +1716,10 @@ __global__ void Marlin( zero_accums(); wait_for_stage(); - init_same_group(0); fetch_to_registers(0, 0); fetch_scales_to_registers(0, 0); fetch_zp_to_registers(0, 0); a_gl_rd_col += a_gl_rd_delta_o * (stages - 1); - if constexpr (has_act_order) { - slice_k_start_shared_fetch += tb_k * (stages - 1); - } }; if (slice_iters) { start_pipes(); @@ -1965,7 +1744,6 @@ __global__ void Marlin( slice_iters >= stages); pipe++; wait_for_stage(); - init_same_group(pipe % stages); } if constexpr (!is_a_8bit) { @@ -1983,25 +1761,6 @@ __global__ void Marlin( a_gl_rd_col += a_gl_rd_delta_o * stages; - if constexpr (has_act_order) { - slice_k_start += tb_k * stages; - - if (slice_k_start < prob_k) { - slice_k_start_shared_fetch += tb_k * stages; - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) { - fetch_act_order_scales_to_shared(false, first_group_id, - last_group_id); - __syncthreads(); - } - } - } - // Process results and, if necessary, proceed to the next column slice. // While this pattern may not be the most readable, other ways of writing // the loop seemed to noticeably worse performance after compilation. @@ -2059,7 +1818,7 @@ __global__ void Marlin( bool last = slice_idx == slice_count - 1; // For per-column scales, we only fetch them here in the final step before // write-out - if constexpr (!has_act_order && group_blocks == -1 && + if constexpr (group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) { if (s_sh_wr_pred) { @@ -2078,7 +1837,7 @@ __global__ void Marlin( cp_async_fence(); } - if constexpr (!has_act_order && group_blocks == -1 && + if constexpr (group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) { if constexpr (is_a_8bit) { cp_async_wait<0>(); @@ -2109,7 +1868,7 @@ __global__ void Marlin( // For 8-bit channelwise, we apply the scale before the global reduction // that converts the fp32 results to fp16 (so that we avoid possible // overflow in fp16) - if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) { + if constexpr (group_blocks == -1 && is_a_8bit) { #pragma unroll for (int j = 0; j < 2; j++) { float2 aa[2]; @@ -2131,8 +1890,7 @@ __global__ void Marlin( } } } - } else if (!has_act_order && group_blocks == -1 && - b_type.size_bits() == 8 && + } else if (group_blocks == -1 && b_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { if (threadIdx.x / 32 < tb_n_warps) { #pragma unroll @@ -2202,33 +1960,25 @@ __global__ void Marlin( b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - // Update slice k/n for scales loading - if constexpr (has_act_order) { - slice_k_start = tb_k * slice_row; - slice_k_finish = slice_k_start + tb_k * slice_iters; - slice_k_start_shared_fetch = slice_k_start; - slice_n_offset = act_s_col_tb_stride * slice_col; + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; } else { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = - s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + - s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = - zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + - zp_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = - s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + - threadIdx.x / s_sh_stride) + - s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - zp_gl_rd = - zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + - threadIdx.x / zp_sh_stride) + - zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; } start_pipes(); } diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu index 177eefa2c6..48de07dc75 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu @@ -45,84 +45,6 @@ __global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -template -__global__ void permute_cols_kernel( - int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr, - int4* __restrict__ out_int4_ptr, - const int32_t* __restrict__ sorted_token_ids_ptr, - const int32_t* __restrict__ expert_ids_ptr, - const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m, - int size_k, int top_k) { - int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; - int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size); - int32_t block_sorted_ids[moe_block_size]; - int block_num_valid_tokens = 0; - int64_t old_expert_id = 0; - int64_t expert_id = 0; - int row_stride = size_k * sizeof(half) / 16; - - auto read_moe_block_data = [&](int block_id) { - block_num_valid_tokens = moe_block_size; - int4* tmp_block_sorted_ids = reinterpret_cast(block_sorted_ids); - for (int i = 0; i < moe_block_size / 4; i++) { - tmp_block_sorted_ids[i] = - ((int4*)sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i]; - } - for (int i = 0; i < moe_block_size; i++) { - if (block_sorted_ids[i] >= size_m * top_k) { - block_num_valid_tokens = i; - break; - }; - } - }; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int in_offset = (row / top_k) * row_stride; - int out_offset = row * row_stride; - - half const* a_row_half = - reinterpret_cast(a_int4_ptr + in_offset); - half* out_half = reinterpret_cast(out_int4_ptr + out_offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int index = blockIdx.x; index < num_moe_blocks; index += gridDim.x) { - old_expert_id = expert_id; - int tmp_expert_id = expert_ids_ptr[index]; - if (tmp_expert_id == -1) continue; - expert_id = tmp_expert_id; - perm_int_ptr += (expert_id - old_expert_id) * size_k; - read_moe_block_data(index); - - for (int i = 0; i < block_num_valid_tokens; i++) - permute_row(block_sorted_ids[i]); - } -} - typedef struct { int thread_k; int thread_n; @@ -152,9 +74,7 @@ typedef struct { int get_scales_cache_size(thread_config_t const& th_config, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, - bool has_act_order, bool is_k_full, int stages) { - bool cache_scales_chunk = has_act_order && !is_k_full; - + int stages) { int tb_n = th_config.thread_n; int tb_k = th_config.thread_k; @@ -162,28 +82,17 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m, int tb_groups; if (group_size == -1) { tb_groups = 1; - } else if (group_size == 0) { - tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size } else { tb_groups = div_ceil(tb_k, group_size); } - if (cache_scales_chunk) { - int load_groups = - tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K - load_groups = max(load_groups, 32); // We load at least 32 scale groups - return load_groups * tb_n * 2; - } else { - int tb_scales = tb_groups * tb_n * 2; - - return tb_scales * stages; - } + int tb_scales = tb_groups * tb_n * 2; + return tb_scales * stages; } int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8, int thread_m_blocks, int prob_m, int prob_n, - int prob_k, int num_bits, int group_size, - bool has_act_order, bool is_k_full, int has_zp, + int prob_k, int num_bits, int group_size, int has_zp, int is_zp_float, bool is_a_8bit, int stages) { int pack_factor = 32 / num_bits; @@ -203,10 +112,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8, (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); - int sh_s_size = - get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, - group_size, has_act_order, is_k_full, stages); - int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0; + int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, + num_bits, group_size, stages); int sh_zp_size = 0; if (has_zp) { if (is_zp_float) @@ -217,16 +124,15 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8, sh_zp_size = sh_s_size / 2; } - int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size + - sh_g_idx_size + sh_block_meta_size; + int total_size = + tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_block_meta_size; return total_size; } bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8, int thread_m_blocks, int prob_m, int prob_n, int prob_k, - int num_bits, int group_size, bool has_act_order, - bool is_k_full, int has_zp, int is_zp_float, + int num_bits, int group_size, int has_zp, int is_zp_float, bool is_a_8bit, int stages, int max_shared_mem) { // Sanity if (th_config.thread_k == -1 || th_config.thread_n == -1 || @@ -250,19 +156,20 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8, } // Check that pipeline fits into cache - int cache_size = - get_kernel_cache_size(th_config, m_block_size_8, thread_m_blocks, prob_m, - prob_n, prob_k, num_bits, group_size, has_act_order, - is_k_full, has_zp, is_zp_float, is_a_8bit, stages); + int cache_size = get_kernel_cache_size( + th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_zp, is_zp_float, is_a_8bit, stages); return cache_size <= max_shared_mem; } -MarlinFuncPtr get_marlin_kernel( - const vllm::ScalarType a_type, const vllm::ScalarType b_type, - const vllm::ScalarType c_type, const vllm::ScalarType s_type, - int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, - bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, - int threads, bool is_zp_float, int stages) { +MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType a_type, + const vllm::ScalarType b_type, + const vllm::ScalarType c_type, + const vllm::ScalarType s_type, + int thread_m_blocks, int thread_n_blocks, + int thread_k_blocks, bool m_block_size_8, + bool has_zp, int group_blocks, int threads, + bool is_zp_float, int stages) { int num_bits = b_type.size_bits(); auto kernel = MarlinDefault; @@ -275,9 +182,8 @@ exec_config_t determine_exec_config( const vllm::ScalarType& a_type, const vllm::ScalarType& b_type, const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m, int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks, - bool m_block_size_8, int num_bits, int group_size, bool has_act_order, - bool is_k_full, bool has_zp, bool is_zp_float, bool is_a_8bit, int stages, - int max_shared_mem, int sms) { + bool m_block_size_8, int num_bits, int group_size, bool has_zp, + bool is_zp_float, bool is_a_8bit, int stages, int max_shared_mem, int sms) { exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; thread_config_t* thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs @@ -293,27 +199,22 @@ exec_config_t determine_exec_config( thread_config_t th_config = thread_configs[i]; if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m, - prob_n, prob_k, num_bits, group_size, has_act_order, - is_k_full, has_zp, is_zp_float, is_a_8bit, stages, + prob_n, prob_k, num_bits, group_size, has_zp, + is_zp_float, is_a_8bit, stages, max_shared_mem - 512)) { continue; } int cache_size = get_kernel_cache_size( th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, - num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, - is_a_8bit, stages); + num_bits, group_size, has_zp, is_zp_float, is_a_8bit, stages); - int group_blocks = 0; - if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : (group_size / 16); - } + int group_blocks = group_size == -1 ? -1 : (group_size / 16); - auto kernel = - get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks, - th_config.thread_n / 16, th_config.thread_k / 16, - m_block_size_8, has_act_order, has_zp, group_blocks, - th_config.num_threads, is_zp_float, stages); + auto kernel = get_marlin_kernel( + a_type, b_type, c_type, s_type, thread_m_blocks, + th_config.thread_n / 16, th_config.thread_k / 16, m_block_size_8, + has_zp, group_blocks, th_config.num_threads, is_zp_float, stages); if (kernel == MarlinDefault) continue; @@ -342,15 +243,14 @@ exec_config_t determine_exec_config( } void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, - void* a_s, void* b_s, void* g_s, void* zp, void* g_idx, - void* perm, void* a_tmp, void* sorted_token_ids, - void* expert_ids, void* num_tokens_past_padded, - void* topk_weights, int moe_block_size, int num_experts, - int top_k, bool mul_topk_weights, int prob_m, int prob_n, - int prob_k, void* workspace, vllm::ScalarType const& a_type, + void* a_s, void* b_s, void* g_s, void* zp, + void* sorted_token_ids, void* expert_ids, + void* num_tokens_past_padded, void* topk_weights, + int moe_block_size, int num_experts, int top_k, + bool mul_topk_weights, int prob_m, int prob_n, int prob_k, + void* workspace, vllm::ScalarType const& a_type, vllm::ScalarType const& b_type, vllm::ScalarType const& c_type, - vllm::ScalarType const& s_type, bool has_bias, - bool has_act_order, bool is_k_full, bool has_zp, int num_groups, + vllm::ScalarType const& s_type, bool has_bias, bool has_zp, int group_size, int dev, cudaStream_t stream, int thread_k, int thread_n, int sms, int blocks_per_sm, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { @@ -361,25 +261,13 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]"); - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { - STD_TORCH_CHECK(group_size != -1); - group_blocks = group_size / 16; - STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } else { - STD_TORCH_CHECK(group_size == 0); - group_blocks = 0; - } + int group_blocks; + if (group_size == -1) { + group_blocks = -1; } else { - if (group_size == -1) { - group_blocks = -1; - } else { - group_blocks = group_size / 16; - STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } + group_blocks = group_size / 16; + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } int num_bits = b_type.size_bits(); @@ -392,9 +280,6 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, const int4* b_s_ptr = (const int4*)b_s; const float* g_s_ptr = (const float*)g_s; const int4* zp_ptr = (const int4*)zp; - const int* g_idx_ptr = (const int*)g_idx; - const int* perm_ptr = (const int*)perm; - int4* a_tmp_ptr = (int4*)a_tmp; const int32_t* sorted_token_ids_ptr = (const int32_t*)sorted_token_ids; const int32_t* expert_ids_ptr = (const int32_t*)expert_ids; const int32_t* num_tokens_past_padded_ptr = @@ -402,37 +287,6 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, const float* topk_weights_ptr = (const float*)topk_weights; int* locks = (int*)workspace; - if (has_act_order) { - // Permute A columns - auto kernel = permute_cols_kernel<8>; - if (moe_block_size == 8) { - } else if (moe_block_size == 16) - kernel = permute_cols_kernel<16>; - else if (moe_block_size == 32) - kernel = permute_cols_kernel<32>; - else if (moe_block_size == 48) - kernel = permute_cols_kernel<48>; - else if (moe_block_size == 64) - kernel = permute_cols_kernel<64>; - else - STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size); - - // avoid ">>>" being formatted to "> > >" - // clang-format off - kernel<<>>( - A_ptr, perm_ptr, a_tmp_ptr, sorted_token_ids_ptr, expert_ids_ptr, - num_tokens_past_padded_ptr, prob_m, prob_k, top_k); - // clang-format on - A_ptr = a_tmp_ptr; - prob_m = prob_m * top_k; - top_k = 1; - - // If we have a full K, then we can run the non-act-order version of Marlin - // (since the weight rows are reordered by increasing group ids, and by - // having a full K, we have full original groups) - if (is_k_full) has_act_order = false; - } - int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); @@ -476,9 +330,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, // Auto config exec_cfg = determine_exec_config( a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts, - top_k, thread_m_blocks, m_block_size_8, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages, - max_shared_mem, sms); + top_k, thread_m_blocks, m_block_size_8, num_bits, group_size, has_zp, + is_zp_float, is_a_8bit, stages, max_shared_mem, sms); thread_tfg = exec_cfg.tb_cfg; } @@ -494,36 +347,32 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, STD_TORCH_CHECK( is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m, - prob_n, prob_k, num_bits, group_size, has_act_order, - is_k_full, has_zp, is_zp_float, is_a_8bit, stages, - max_shared_mem), + prob_n, prob_k, num_bits, group_size, has_zp, is_zp_float, + is_a_8bit, stages, max_shared_mem), "Invalid thread config: thread_m_blocks = ", thread_m_blocks, ", thread_k = ", thread_tfg.thread_k, ", thread_n = ", thread_tfg.thread_n, ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, - ", group_size = ", group_size, ", has_act_order = ", has_act_order, - ", is_k_full = ", is_k_full, ", has_zp = ", has_zp, + ", group_size = ", group_size, ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem); - int sh_cache_size = - get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m, - prob_n, prob_k, num_bits, group_size, has_act_order, - is_k_full, has_zp, is_zp_float, is_a_8bit, stages); + int sh_cache_size = get_kernel_cache_size( + thread_tfg, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_zp, is_zp_float, is_a_8bit, stages); - auto kernel = get_marlin_kernel( - a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks, - thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, - num_threads, is_zp_float, stages); + auto kernel = + get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks, + thread_n_blocks, thread_k_blocks, m_block_size_8, + has_zp, group_blocks, num_threads, is_zp_float, stages); if (kernel == MarlinDefault) { - STD_TORCH_CHECK( - false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", - prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, - ", thread_m_blocks = ", thread_m_blocks, - ", thread_n_blocks = ", thread_n_blocks, - ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits); + STD_TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, + ", ", prob_k, "]", ", group_size = ", group_size, + ", thread_m_blocks = ", thread_m_blocks, + ", thread_n_blocks = ", thread_n_blocks, + ", thread_k_blocks = ", thread_k_blocks, + ", num_bits = ", num_bits); } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -531,10 +380,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, // avoid ">>>" being formatted to "> > >" // clang-format off kernel<<>>( - A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr, + A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, - topk_weights_ptr, top_k, mul_topk_weights, num_groups, prob_m, - prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce); + topk_weights_ptr, top_k, mul_topk_weights, prob_m, prob_n, prob_k, locks, + has_bias, use_atomic_add, use_fp32_reduce); // clang-format on } @@ -548,15 +397,13 @@ torch::stable::Tensor moe_wna16_marlin_gemm( std::optional const& a_scales_or_none, std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::stable::Tensor& workspace, torch::stable::Tensor& sorted_token_ids, torch::stable::Tensor& expert_ids, torch::stable::Tensor& num_tokens_past_padded, torch::stable::Tensor& topk_weights, int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, vllm::ScalarTypeId const& b_type_id, int64_t size_m, - int64_t size_n, int64_t size_k, bool is_k_full, bool use_atomic_add, - bool use_fp32_reduce, bool is_zp_float, int64_t thread_k, int64_t thread_n, + int64_t size_n, int64_t size_k, bool use_atomic_add, bool use_fp32_reduce, + bool is_zp_float, int64_t thread_k, int64_t thread_n, int64_t blocks_per_sm) { vllm::ScalarTypeId a_type_id, c_type_id, s_type_id; @@ -717,7 +564,7 @@ torch::stable::Tensor moe_wna16_marlin_gemm( c_tmp = torch::stable::new_empty(a, {0}, kFloat); } - // Detect groupsize and act_order + // Detect group size. int num_groups = -1; int group_size = -1; @@ -728,50 +575,13 @@ torch::stable::Tensor moe_wna16_marlin_gemm( " is not size_n = ", size_n); num_groups = b_scales.size(1); - torch::stable::Tensor g_idx, perm, a_tmp; - if (g_idx_or_none.has_value() && perm_or_none.has_value()) { - g_idx = g_idx_or_none.value(); - perm = perm_or_none.value(); - - STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); - STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); - STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); - - // Verify g_idx and perm - STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || - (g_idx.size(-1) == size_k && perm.size(-1) == size_k), - "Unexpected g_idx.size(-1) = ", g_idx.size(-1), - " and perm.size(-1) = ", perm.size(-1), - ", where size_k = ", size_k); + if (num_groups > 1) { + STD_TORCH_CHECK( + size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by b_scales.size(1) = ", b_scales.size(1)); + group_size = size_k / num_groups; } else { - g_idx = torch::stable::new_empty(a, {0}, c_dtype); - perm = torch::stable::new_empty(a, {0}, c_dtype); - a_tmp = torch::stable::new_empty(a, {0}, c_dtype); - } - bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0; - - if (has_act_order) { - a_tmp = torch::stable::new_empty(a, {size_m * top_k, size_k}, c_dtype); - if (is_k_full) { - STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); - STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by num_groups = ", num_groups); - group_size = size_k / num_groups; - } else { - group_size = 0; - } - - } else { - a_tmp = torch::stable::new_empty(a, {0}, c_dtype); - if (num_groups > 1) { - STD_TORCH_CHECK( - size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by b_scales.size(1) = ", b_scales.size(1)); - group_size = size_k / num_groups; - } else { - group_size = -1; - } + group_size = -1; } torch::stable::Tensor global_scale; @@ -880,15 +690,13 @@ torch::stable::Tensor moe_wna16_marlin_gemm( c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(), a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(), global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), - g_idx.mutable_data_ptr(), perm.mutable_data_ptr(), - a_tmp.mutable_data_ptr(), sorted_token_ids.mutable_data_ptr(), - expert_ids.mutable_data_ptr(), num_tokens_past_padded.mutable_data_ptr(), + sorted_token_ids.mutable_data_ptr(), expert_ids.mutable_data_ptr(), + num_tokens_past_padded.mutable_data_ptr(), topk_weights.mutable_data_ptr(), moe_block_size, num_experts, top_k, mul_topk_weights, size_m, size_n, size_k, workspace.mutable_data_ptr(), - a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full, - has_zp, num_groups, group_size, dev, get_current_cuda_stream(dev), - thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce, - is_zp_float); + a_type, b_type, c_type, s_type, has_bias, has_zp, group_size, dev, + get_current_cuda_stream(dev), thread_k, thread_n, sms, blocks_per_sm, + use_atomic_add, use_fp32_reduce, is_zp_float); return c; } diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index 800284d70f..c0e966b2be 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -79,13 +79,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { "Tensor! b_q_weight, Tensor? b_bias_or_none," "Tensor! b_scales, Tensor? a_scales, Tensor? global_scale, Tensor? " "b_zeros_or_none," - "Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace," + "Tensor! workspace," "Tensor sorted_token_ids," "Tensor! expert_ids, Tensor! num_tokens_past_padded," "Tensor! topk_weights, int moe_block_size, int top_k, " "bool mul_topk_weights, int b_type_id," - "int size_m, int size_n, int size_k," - "bool is_full_k, bool use_atomic_add," + "int size_m, int size_n, int size_k, bool use_atomic_add," "bool use_fp32_reduce, bool is_zp_float," "int thread_k, int thread_n, int blocks_per_sm) -> Tensor"); diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 1f36f7796c..acf665b07a 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -43,9 +43,6 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input, int64_t group_size, double eps, double int8_min, double int8_max); -torch::stable::Tensor permute_cols(torch::stable::Tensor const& A, - torch::stable::Tensor const& perm); - #ifndef USE_ROCM bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); @@ -587,11 +584,10 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, torch::stable::Tensor b_q_weight, torch::stable::Tensor b_gptq_qzeros, torch::stable::Tensor b_gptq_scales, - torch::stable::Tensor b_g_idx, bool use_exllama, - bool use_v2_format, int64_t bit); + bool use_exllama, bool use_v2_format, + int64_t bit); -void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm, - int64_t bit); +void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit); // Cache ops (shared CUDA/ROCm) void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, diff --git a/csrc/libtorch_stable/permute_cols.cu b/csrc/libtorch_stable/permute_cols.cu deleted file mode 100644 index 3162ac02c0..0000000000 --- a/csrc/libtorch_stable/permute_cols.cu +++ /dev/null @@ -1,94 +0,0 @@ -#include -#include -#include -#include -#include - -#include - -#include "torch_utils.h" - -static constexpr int default_threads = 256; -static constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } - -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -// Currently only supports 16bit types (since we permute half types) -__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, - int const* __restrict__ perm_int_ptr, - int4* __restrict__ out_int4_ptr, int size_m, - int size_k, int block_rows) { - int start_row = block_rows * blockIdx.x; - int finish_row = start_row + block_rows; - if (finish_row > size_m) { - finish_row = size_m; - } - int cur_block_rows = std::max(finish_row - start_row, 0); - - int row_stride = size_k * sizeof(half) / 16; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int offset = row * row_stride; - - half const* a_row_half = reinterpret_cast(a_int4_ptr + offset); - half* out_half = reinterpret_cast(out_int4_ptr + offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - int cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - int cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int i = 0; i < cur_block_rows; i++) { - int cur_row = start_row + i; - if (cur_row < size_m) { - permute_row(cur_row); - } - } -} - -// More efficient version of A[..., perm] -// taken from gptq_marlin.cu -torch::stable::Tensor permute_cols(torch::stable::Tensor const& A, - torch::stable::Tensor const& perm) { - const int32_t dev = A.get_device_index(); - const torch::stable::accelerator::DeviceGuard device_guard(dev); - const auto stream = get_current_cuda_stream(dev); - - STD_TORCH_CHECK( - A.scalar_type() == torch::headeronly::ScalarType::Half || - A.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "Currently only 16bit types are supported"); - STD_TORCH_CHECK(A.is_contiguous(), "A must be contiguous"); - STD_TORCH_CHECK(A.size(-1) % 8 == 0, - "A columns must be a multiple of 8 (128bits)"); - auto A_2d = torch::stable::view(A, {-1, A.size(-1)}); - - torch::stable::Tensor D = torch::stable::empty_like(A); - int sms; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); - int block_rows = div_ceil(A_2d.size(0), sms); - permute_cols_kernel<<>>( - reinterpret_cast(A_2d.const_data_ptr()), - perm.const_data_ptr(), reinterpret_cast(D.mutable_data_ptr()), - A_2d.size(0), A_2d.size(1), block_rows); - return D; -} \ No newline at end of file diff --git a/csrc/libtorch_stable/quantization/gptq/q_gemm.cu b/csrc/libtorch_stable/quantization/gptq/q_gemm.cu index af0b9837b7..6ce6f6a909 100644 --- a/csrc/libtorch_stable/quantization/gptq/q_gemm.cu +++ b/csrc/libtorch_stable/quantization/gptq/q_gemm.cu @@ -121,7 +121,7 @@ typedef void (*fp_gemm_half_q_half_gptq_kernel)(const half*, const uint32_t*, const uint32_t*, const half*, half*, const int, const int, const int, const int, - const bool, const int*); + const bool); template __global__ void gemm_half_q_half_gptq_4bit_kernel( @@ -129,7 +129,7 @@ __global__ void gemm_half_q_half_gptq_4bit_kernel( const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, half* __restrict__ c, const int size_m, const int size_n, const int size_k, const int groups, - const bool use_v2_format, const int* __restrict__ b_q_perm) { + const bool use_v2_format) { MatrixView_half a_(a, size_m, size_k); MatrixView_half_rw c_(c, size_m, size_n); MatrixView_q4_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); @@ -157,12 +157,7 @@ __global__ void gemm_half_q_half_gptq_4bit_kernel( const half* a_ptr = a_.item_ptr(offset_m + m, 0); half* block_a_ptr = block_a[m]; - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; + block_a_ptr[t] = a_ptr[offset_k + t]; } } @@ -263,7 +258,7 @@ __global__ void gemm_half_q_half_gptq_2bit_kernel( const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, half* __restrict__ c, const int size_m, const int size_n, const int size_k, const int groups, - const bool use_v2_format, const int* __restrict__ b_q_perm) { + const bool use_v2_format) { MatrixView_half a_(a, size_m, size_k); MatrixView_half_rw c_(c, size_m, size_n); MatrixView_q2_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); @@ -291,12 +286,7 @@ __global__ void gemm_half_q_half_gptq_2bit_kernel( const half* a_ptr = a_.item_ptr(offset_m + m, 0); half* block_a_ptr = block_a[m]; - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; + block_a_ptr[t] = a_ptr[offset_k + t]; } } @@ -380,7 +370,7 @@ __global__ void gemm_half_q_half_gptq_3bit_kernel( const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, half* __restrict__ c, const int size_m, const int size_n, const int size_k, const int groups, - const bool use_v2_format, const int* __restrict__ b_q_perm) { + const bool use_v2_format) { MatrixView_half a_(a, size_m, size_k); MatrixView_half_rw c_(c, size_m, size_n); MatrixView_q3_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); @@ -408,12 +398,7 @@ __global__ void gemm_half_q_half_gptq_3bit_kernel( const half* a_ptr = a_.item_ptr(offset_m + m, 0); half* block_a_ptr = block_a[m]; - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; + block_a_ptr[t] = a_ptr[offset_k + t]; } } @@ -504,7 +489,7 @@ __global__ void gemm_half_q_half_gptq_8bit_kernel( const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, half* __restrict__ c, const int size_m, const int size_n, const int size_k, const int groups, - const bool use_v2_format, const int* __restrict__ b_q_perm) { + const bool use_v2_format) { MatrixView_half a_(a, size_m, size_k); MatrixView_half_rw c_(c, size_m, size_n); MatrixView_q8_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); @@ -532,12 +517,7 @@ __global__ void gemm_half_q_half_gptq_8bit_kernel( const half* a_ptr = a_.item_ptr(offset_m + m, 0); half* block_a_ptr = block_a[m]; - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; + block_a_ptr[t] = a_ptr[offset_k + t]; } } @@ -656,10 +636,9 @@ fp_gemm_half_q_half_gptq_kernel pick_gemm_half_q_half_gptq_kernel( void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, const int* b_q_perm, - half* c, int size_m, int size_n, int size_k, - int m_count, int groups, bool use_v2_format, - int bit) { + const half* b_gptq_scales, half* c, int size_m, + int size_n, int size_k, int m_count, int groups, + bool use_v2_format, int bit) { dim3 blockDim, gridDim; blockDim.x = BLOCK_KN_SIZE; blockDim.y = 1; @@ -672,13 +651,13 @@ void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight, pick_gemm_half_q_half_gptq_kernel(true, m_count, bit); const cudaStream_t stream = get_current_cuda_stream(); - kernel<<>>( - a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m, size_n, size_k, - groups, use_v2_format, b_q_perm); + kernel<<>>(a, b_q_weight, b_gptq_qzeros, + b_gptq_scales, c, size_m, size_n, + size_k, groups, use_v2_format); } __global__ void reconstruct_exllama_8bit_kernel( - const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm, + const uint32_t* __restrict__ b_q_weight, const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, const int size_k, const int size_n, const int groups, const bool use_v2_format, half* __restrict__ b) { @@ -694,14 +673,8 @@ __global__ void reconstruct_exllama_8bit_kernel( int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; auto t = threadIdx.x; - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - // Column int n = offset_n + t * 4; if (n >= size_n) return; @@ -722,8 +695,6 @@ __global__ void reconstruct_exllama_8bit_kernel( b_gptq_qzeros_.item4(zeros, group, n); b_gptq_scales_.item4_h2(scales, group, n); - __syncthreads(); - int k = offset_k; int lk = 0; @@ -752,25 +723,13 @@ __global__ void reconstruct_exllama_8bit_kernel( dequant_8bit_8(load_int4[0].w, load_int4[1].w, dq[3], size_n, zeros[3] + zero_offset); - // half* dqh = (half*)dq; - if (b_q_perm) { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), - __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]), - __high2half(dq[2][j]), __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), - __low2half(dq[1][j]), __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), - __high2half(dq[1][j]), __high2half(dq[2][j]), - __high2half(dq[3][j])); - } + for (int j = 0; j < 4; j++) { + for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); + b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]), + __low2half(dq[2][j]), __low2half(dq[3][j])); + b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), + __high2half(dq[1][j]), __high2half(dq[2][j]), + __high2half(dq[3][j])); } } k += 32; @@ -778,7 +737,7 @@ __global__ void reconstruct_exllama_8bit_kernel( } __global__ void reconstruct_exllama_4bit_kernel( - const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm, + const uint32_t* __restrict__ b_q_weight, const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, const int size_k, const int size_n, const int groups, const bool use_v2_format, half* __restrict__ b) { @@ -794,14 +753,8 @@ __global__ void reconstruct_exllama_4bit_kernel( int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; auto t = threadIdx.x; - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - // Column int n = offset_n + t * 4; if (n >= size_n) return; @@ -828,8 +781,6 @@ __global__ void reconstruct_exllama_4bit_kernel( dequant_4bit_8_prep_zero(zeros[2] + zero_offset, z1z16[2], y1y16[2]); dequant_4bit_8_prep_zero(zeros[3] + zero_offset, z1z16[3], y1y16[3]); - __syncthreads(); - int k = offset_k; int lk = 0; @@ -860,25 +811,13 @@ __global__ void reconstruct_exllama_4bit_kernel( false); b_ptr += size_n; - // half* dqh = (half*)dq; - if (b_q_perm) { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), - __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]), - __high2half(dq[2][j]), __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), - __low2half(dq[1][j]), __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), - __high2half(dq[1][j]), __high2half(dq[2][j]), - __high2half(dq[3][j])); - } + for (int j = 0; j < 4; j++) { + for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); + b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]), + __low2half(dq[2][j]), __low2half(dq[3][j])); + b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), + __high2half(dq[1][j]), __high2half(dq[2][j]), + __high2half(dq[3][j])); } } k += 32; @@ -886,7 +825,7 @@ __global__ void reconstruct_exllama_4bit_kernel( } __global__ void reconstruct_exllama_3bit_kernel( - const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm, + const uint32_t* __restrict__ b_q_weight, const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, const int size_k, const int size_n, const int groups, const bool use_v2_format, half* __restrict__ b) { @@ -902,14 +841,8 @@ __global__ void reconstruct_exllama_3bit_kernel( int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; auto t = threadIdx.x; - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - // Column int n = offset_n + t * 4; if (n >= size_n) return; @@ -930,8 +863,6 @@ __global__ void reconstruct_exllama_3bit_kernel( b_gptq_qzeros_.item4(zeros, group, n); b_gptq_scales_.item4_h2(scales, group, n); - __syncthreads(); - int k = offset_k; int lk = 0; @@ -962,24 +893,13 @@ __global__ void reconstruct_exllama_3bit_kernel( dequant_3bit_32(load_int4[0].w, load_int4[1].w, load_int4[2].w, dq[3], size_n, zeros[3] + zero_offset); - if (b_q_perm) { - for (int j = 0; j < 16; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), - __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]), - __high2half(dq[2][j]), __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 16; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), - __low2half(dq[1][j]), __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), - __high2half(dq[1][j]), __high2half(dq[2][j]), - __high2half(dq[3][j])); - } + for (int j = 0; j < 16; j++) { + for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); + b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]), + __low2half(dq[2][j]), __low2half(dq[3][j])); + b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), + __high2half(dq[1][j]), __high2half(dq[2][j]), + __high2half(dq[3][j])); } } k += 32; @@ -987,7 +907,7 @@ __global__ void reconstruct_exllama_3bit_kernel( } __global__ void reconstruct_exllama_2bit_kernel( - const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm, + const uint32_t* __restrict__ b_q_weight, const uint32_t* __restrict__ b_gptq_qzeros, const half* __restrict__ b_gptq_scales, const int size_k, const int size_n, const int groups, const bool use_v2_format, half* __restrict__ b) { @@ -1003,14 +923,8 @@ __global__ void reconstruct_exllama_2bit_kernel( int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; auto t = threadIdx.x; - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - // Column int n = offset_n + t * 4; if (n >= size_n) return; @@ -1031,8 +945,6 @@ __global__ void reconstruct_exllama_2bit_kernel( b_gptq_qzeros_.item4(zeros, group, n); b_gptq_scales_.item4_h2(scales, group, n); - __syncthreads(); - int k = offset_k; int lk = 0; @@ -1055,25 +967,13 @@ __global__ void reconstruct_exllama_2bit_kernel( dequant_2bit_16(load_int4.w, dq[3], size_n, zeros[3] + zero_offset); b_ptr += size_n; - // half* dqh = (half*)dq; - if (b_q_perm) { - for (int j = 0; j < 8; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), - __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]), - __high2half(dq[2][j]), __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 8; j++) { - for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), - __low2half(dq[1][j]), __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), - __high2half(dq[1][j]), __high2half(dq[2][j]), - __high2half(dq[3][j])); - } + for (int j = 0; j < 8; j++) { + for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]); + b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]), + __low2half(dq[2][j]), __low2half(dq[3][j])); + b_.set4(offset_k + lk++, n, __high2half(dq[0][j]), + __high2half(dq[1][j]), __high2half(dq[2][j]), + __high2half(dq[3][j])); } } k += 32; @@ -1082,9 +982,8 @@ __global__ void reconstruct_exllama_2bit_kernel( void reconstruct_exllama(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, const int* b_q_perm, - half* out, int height, int width, int groups, - bool use_v2_format, int bit) { + const half* b_gptq_scales, half* out, int height, + int width, int groups, bool use_v2_format, int bit) { dim3 blockDim, gridDim; blockDim.x = BLOCK_KN_SIZE; blockDim.y = 1; @@ -1102,15 +1001,15 @@ void reconstruct_exllama(const uint32_t* b_q_weight, const cudaStream_t stream = get_current_cuda_stream(); reconstruct_exllama_kernel<<>>( - b_q_weight, b_q_perm, b_gptq_qzeros, b_gptq_scales, height, width, groups, + b_q_weight, b_gptq_qzeros, b_gptq_scales, height, width, groups, use_v2_format, out); } __global__ void gemm_half_q_half_alt_4bit_kernel( const half2* __restrict__ vec, const uint32_t* __restrict__ mat, half* __restrict__ mul, const half* __restrict__ scales, - const uint32_t* __restrict__ zeros, const int* __restrict__ g_idx, - int batch, int height, int width, bool use_v2_format) { + const uint32_t* __restrict__ zeros, int batch, int height, int width, + int groups, bool use_v2_format) { int zero_width = width / 8; int vec_height = height * 4; const int blockwidth2 = BLOCK_KN_SIZE / 2; @@ -1144,6 +1043,7 @@ __global__ void gemm_half_q_half_alt_4bit_kernel( int i = width * h + w; int g_h = h * 8; + int group_size = vec_height * 2 / groups; int k = 0; int z_w = w / 8; int z_mod = (w % 8) * 4; @@ -1156,8 +1056,8 @@ __global__ void gemm_half_q_half_alt_4bit_kernel( half2 scales_tmp[4]; half2 zeros_tmp[4]; for (int tmp_k = 0; tmp_k < 4; tmp_k++) { - int g = g_idx[g_h + (k + tmp_k) * 2]; - int g2 = g_idx[g_h + (k + tmp_k) * 2 + 1]; + int g = (g_h + (k + tmp_k) * 2) / group_size; + int g2 = (g_h + (k + tmp_k) * 2 + 1) / group_size; half scale_f = scales[g * width + w]; half scale_f2 = scales[g2 * width + w]; half2 scale = __halves2half2(scale_f, scale_f2); @@ -1209,8 +1109,8 @@ __global__ void gemm_half_q_half_alt_4bit_kernel( __global__ void gemm_half_q_half_alt_8bit_kernel( const half2* __restrict__ vec, const uint32_t* __restrict__ mat, half* __restrict__ mul, const half* __restrict__ scales, - const uint32_t* __restrict__ zeros, const int* __restrict__ g_idx, - int batch, int height, int width, bool use_v2_format) { + const uint32_t* __restrict__ zeros, int batch, int height, int width, + int groups, bool use_v2_format) { int zero_width = width / 4; int vec_height = height * 2; const int blockwidth2 = BLOCK_KN_SIZE / 2; @@ -1236,6 +1136,7 @@ __global__ void gemm_half_q_half_alt_8bit_kernel( int i = width * h + w; int g_h = h * 4; + int group_size = vec_height * 2 / groups; int k = 0; int z_w = w / 4; int z_mod = (w % 4) * 8; @@ -1248,8 +1149,8 @@ __global__ void gemm_half_q_half_alt_8bit_kernel( half2 scales_tmp[2]; half2 zeros_tmp[2]; for (int tmp_k = 0; tmp_k < 2; tmp_k++) { - int g = g_idx[g_h + (k + tmp_k) * 2]; - int g2 = g_idx[g_h + (k + tmp_k) * 2 + 1]; + int g = (g_h + (k + tmp_k) * 2) / group_size; + int g2 = (g_h + (k + tmp_k) * 2 + 1) / group_size; half scale_f = scales[g * width + w]; half scale_f2 = scales[g2 * width + w]; half2 scale = __halves2half2(scale_f, scale_f2); @@ -1296,8 +1197,8 @@ __global__ void gemm_half_q_half_alt_8bit_kernel( void gemm_half_q_half_alt(const half* a, const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, const int* b_g_idx, - half* c, int size_m, int size_n, int size_k, + const half* b_gptq_scales, half* c, int size_m, + int size_n, int size_k, int groups, bool use_v2_format, int bit) { dim3 blockDim, gridDim; blockDim.x = BLOCK_KN_SIZE; @@ -1314,16 +1215,15 @@ void gemm_half_q_half_alt(const half* a, const uint32_t* b_q_weight, const cudaStream_t stream = get_current_cuda_stream(); kernel<<>>( - (const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, b_g_idx, - size_m, size_k / 32 * bit, size_n, use_v2_format); + (const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, size_m, + size_k / 32 * bit, size_n, groups, use_v2_format); } template __global__ void reconstruct_gptq_kernel( const uint32_t* __restrict__ w, const half* __restrict__ w_scales, - const uint32_t* __restrict__ w_zeros, const int* __restrict__ g_idx, - const int height, const int width, const int group, - const bool use_v2_format, half* __restrict__ out) { + const uint32_t* __restrict__ w_zeros, const int height, const int width, + const int num_groups, const bool use_v2_format, half* __restrict__ out) { // Start of block auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x; @@ -1333,8 +1233,8 @@ __global__ void reconstruct_gptq_kernel( // Views MatrixView_half_rw out_(out, height, width); - MatrixView_half w_scales_(w_scales, group, width); - T w_zeros_(w_zeros, group, width); + MatrixView_half w_scales_(w_scales, num_groups, width); + T w_zeros_(w_zeros, num_groups, width); // GPTQv2 and GPTQv1 handles zero points differently int zero_offset = use_v2_format ? 0 : 1; @@ -1344,7 +1244,7 @@ __global__ void reconstruct_gptq_kernel( #pragma unroll for (int s = 0; s < 32; s += bit) { - int group = g_idx[row + s / bit]; + int group = (row + s / bit) / (height / num_groups); half w_scale = w_scales_.item(group, column); uint32_t w_zero = w_zeros_.item(group, column) + zero_offset; half w_item = @@ -1357,9 +1257,8 @@ __global__ void reconstruct_gptq_kernel( __global__ void reconstruct_gptq_3bit_kernel( const uint32_t* __restrict__ w, const half* __restrict__ w_scales, - const uint32_t* __restrict__ w_zeros, const int* __restrict__ g_idx, - const int height, const int width, const int group, - const bool use_v2_format, half* __restrict__ out) { + const uint32_t* __restrict__ w_zeros, const int height, const int width, + const int num_groups, const bool use_v2_format, half* __restrict__ out) { // Start of block auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x; auto row = blockIdx.y * 32; @@ -1368,8 +1267,8 @@ __global__ void reconstruct_gptq_3bit_kernel( // Views MatrixView_half_rw out_(out, height, width); - MatrixView_half w_scales_(w_scales, group, width); - MatrixView_q3_row w_zeros_(w_zeros, group, width); + MatrixView_half w_scales_(w_scales, num_groups, width); + MatrixView_q3_row w_zeros_(w_zeros, num_groups, width); // GPTQv2 and GPTQv1 handles zero points differently int zero_offset = use_v2_format ? 0 : 1; @@ -1381,7 +1280,7 @@ __global__ void reconstruct_gptq_3bit_kernel( #pragma unroll for (int i = 0; i < 32; i += 1) { - int group = g_idx[row + i]; + int group = (row + i) / (height / num_groups); half w_scale = w_scales_.item(group, column); uint32_t w_zero = w_zeros_.item(group, column) + zero_offset; int w_item; @@ -1402,9 +1301,8 @@ __global__ void reconstruct_gptq_3bit_kernel( } void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, const int* b_g_idx, half* out, - int height, int width, int groups, bool use_v2_format, - int bit) { + const half* b_gptq_scales, half* out, int height, + int width, int groups, bool use_v2_format, int bit) { dim3 blockDim, gridDim; blockDim.x = BLOCK_KN_SIZE; blockDim.y = 1; @@ -1423,17 +1321,16 @@ void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros, const cudaStream_t stream = get_current_cuda_stream(); kernel<<>>(b_q_weight, b_gptq_scales, - b_gptq_qzeros, b_g_idx, height, - width, groups, use_v2_format, out); + b_gptq_qzeros, height, width, groups, + use_v2_format, out); } void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a, const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, const int* b_g_idx, - half* c, half* temp_dq, int size_m, int size_n, - int size_k, int groups, bool use_exllama, - bool use_v2_format, int bit) { + const half* b_gptq_scales, half* c, half* temp_dq, + int size_m, int size_n, int size_k, int groups, + bool use_exllama, bool use_v2_format, int bit) { bool use_reconstruct; if (use_exllama) { use_reconstruct = ((bit == 8 && size_m > MAX_Q_GEMM_ROWS_8BIT) || @@ -1446,11 +1343,11 @@ void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a, if (use_reconstruct) { // Reconstruct FP16 matrix, then cuBLAS if (use_exllama) { - reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, - temp_dq, size_k, size_n, groups, use_v2_format, bit); + reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, temp_dq, + size_k, size_n, groups, use_v2_format, bit); } else { - reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, - temp_dq, size_k, size_n, groups, use_v2_format, bit); + reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, temp_dq, + size_k, size_n, groups, use_v2_format, bit); } const half alpha = __float2half(1.0f); @@ -1464,20 +1361,20 @@ void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a, int last_chunk_size = size_m - last_chunk; if (max_chunks) { - gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, - b_g_idx, c, last_chunk, size_n, size_k, - BLOCK_M_SIZE_MAX, groups, use_v2_format, bit); + gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, + last_chunk, size_n, size_k, BLOCK_M_SIZE_MAX, + groups, use_v2_format, bit); } if (last_chunk_size) { gemm_half_q_half_cuda_part( a + last_chunk * size_k, b_q_weight, b_gptq_qzeros, b_gptq_scales, - b_g_idx, c + last_chunk * size_n, last_chunk_size, size_n, size_k, + c + last_chunk * size_n, last_chunk_size, size_n, size_k, last_chunk_size, groups, use_v2_format, bit); } } else { - gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, - c, size_m, size_n, size_k, use_v2_format, bit); + gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m, + size_n, size_k, groups, use_v2_format, bit); } } @@ -1533,214 +1430,8 @@ __global__ void shuffle_3bit_kernel(uint32_t* __restrict__ b_q_weight, } } -__global__ void make_sequential_4bit_kernel(const uint32_t* __restrict__ w, - uint32_t* __restrict__ w_new, - const int* __restrict__ q_perm, - const int w_width) { - const uint64_t* w2 = (uint64_t*)w; - uint64_t* w_new2 = (uint64_t*)w_new; - int w2_stride = w_width >> 1; - auto w2_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w2_column >= w2_stride) return; - auto w_new2_row = blockIdx.y; - int q_perm_idx = w_new2_row << 3; - uint64_t dst = 0; - -#pragma unroll - for (int i = 0; i < 8; i++) { - int source_row = q_perm[q_perm_idx++]; - - int w2_row = source_row >> 3; - int w2_subrow = source_row & 0x07; - int w2_row_shift = w2_subrow << 2; - int wnew2_row_shift = i << 2; - - uint64_t src = w2[w2_row * w2_stride + w2_column]; - src >>= w2_row_shift; - src &= 0x0000000f0000000f; - src <<= wnew2_row_shift; - dst |= src; - } - w_new2[w_new2_row * w2_stride + w2_column] = dst; -} - -__global__ void make_sequential_2bit_kernel(const uint32_t* __restrict__ w, - uint32_t* __restrict__ w_new, - const int* __restrict__ q_perm, - const int w_width) { - const uint64_t* w2 = (uint64_t*)w; - uint64_t* w_new2 = (uint64_t*)w_new; - int w2_stride = w_width >> 1; - auto w2_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w2_column >= w2_stride) return; - auto w_new2_row = blockIdx.y; - int q_perm_idx = w_new2_row << 4; - uint64_t dst = 0; - -#pragma unroll - for (int i = 0; i < 16; i++) { - int source_row = q_perm[q_perm_idx++]; - - int w2_row = source_row >> 4; - int w2_subrow = source_row & 0x0f; - int w2_row_shift = w2_subrow << 1; - int wnew2_row_shift = i << 1; - - uint64_t src = w2[w2_row * w2_stride + w2_column]; - src >>= w2_row_shift; - src &= 0x0000000300000003; - src <<= wnew2_row_shift; - dst |= src; - } - w_new2[w_new2_row * w2_stride + w2_column] = dst; -} - -__global__ void make_sequential_3bit_kernel(const uint32_t* __restrict__ w, - uint32_t* __restrict__ w_new, - const int* __restrict__ q_perm, - const int w_width) { - auto w_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w_column >= w_width) return; - auto w_new_row = blockIdx.y * 3; - auto q_perm_idx = blockIdx.y << 5; - uint32_t dst[3] = {0, 0, 0}; - -#pragma unroll - for (int i = 0; i < 32; i++) { - int source_row = q_perm[q_perm_idx++]; - int z_w = (source_row / 32) * 3; - int z_mod = source_row % 32; - int z_bit; - - if (z_mod != 10) { - if (z_mod != 21) { - z_bit = z_mod; - if (z_bit > 21) { - z_bit *= 3; - z_bit -= 64; - z_w += 2; - } else if (z_bit > 10) { - z_bit *= 3; - z_bit -= 32; - z_w += 1; - } else { - z_bit *= 3; - } - } else { - z_w += 1; - } - } - - uint64_t src; - if (z_mod == 10) { - src = (w[z_w * w_width + w_column] >> 30) | - ((w[(z_w + 1) * w_width + w_column] << 2) & 0x4); - } else if (z_mod == 21) { - src = (w[z_w * w_width + w_column] >> 31) | - ((w[(z_w + 1) * w_width + w_column] << 1) & 0x6); - } else { - src = w[z_w * w_width + w_column]; - src >>= z_bit; - src &= 0x07; - } - - z_w = 0; - if (i != 10) { - if (i != 21) { - z_bit = i; - if (z_bit > 21) { - z_bit *= 3; - z_bit -= 64; - z_w += 2; - } else if (z_bit > 10) { - z_bit *= 3; - z_bit -= 32; - z_w += 1; - } else { - z_bit *= 3; - } - } else { - z_w += 1; - } - } - if (i == 10) { - dst[z_w] |= (src & 0x03) << 30; - dst[z_w + 1] |= ((src & 0x4) >> 2); - } else if (i == 21) { - dst[z_w] |= (src & 0x01) << 31; - dst[z_w + 1] |= ((src & 0x6) >> 1); - } else { - dst[z_w] |= (src << z_bit); - } - } - w_new[w_new_row * w_width + w_column] = dst[0]; - w_new[(w_new_row + 1) * w_width + w_column] = dst[1]; - w_new[(w_new_row + 2) * w_width + w_column] = dst[2]; -} - -__global__ void make_sequential_8bit_kernel(const uint32_t* __restrict__ w, - uint32_t* __restrict__ w_new, - const int* __restrict__ q_perm, - const int w_width) { - const uint64_t* w2 = (uint64_t*)w; - uint64_t* w_new2 = (uint64_t*)w_new; - int w2_stride = w_width >> 1; - auto w2_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w2_column >= w2_stride) return; - auto w_new2_row = blockIdx.y; - int q_perm_idx = w_new2_row << 2; - uint64_t dst = 0; - -#pragma unroll - for (int i = 0; i < 4; i++) { - int source_row = q_perm[q_perm_idx++]; - - int w2_row = source_row >> 2; - int w2_subrow = source_row & 0x03; - int w2_row_shift = w2_subrow << 3; - int wnew2_row_shift = i << 3; - - uint64_t src = w2[w2_row * w2_stride + w2_column]; - src >>= w2_row_shift; - src &= 0x000000ff000000ff; - src <<= wnew2_row_shift; - dst |= src; - } - w_new2[w_new2_row * w2_stride + w2_column] = dst; -} - -void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height, - int width, int bit) { - if (q_perm) { - uint32_t* new_qweight = NULL; - cudaMalloc(&new_qweight, height / 32 * bit * width * sizeof(uint32_t)); - - dim3 blockDim, gridDim; - blockDim.x = THREADS_X; - blockDim.y = 1; - gridDim.x = DIVIDE(width, THREADS_X); - gridDim.y = height / 32 * bit; - - auto kernel = make_sequential_4bit_kernel; - if (bit == 2) { - kernel = make_sequential_2bit_kernel; - } else if (bit == 3) { - kernel = make_sequential_3bit_kernel; - gridDim.y = height / 32; - } else if (bit == 8) { - kernel = make_sequential_8bit_kernel; - } - const cudaStream_t stream = get_current_cuda_stream(); - kernel<<>>(q_weight, new_qweight, q_perm, - width); - // Replace qweights - cudaMemcpyAsync(q_weight, new_qweight, - height / 32 * bit * width * sizeof(uint32_t), - cudaMemcpyDeviceToDevice); - // Cleanup - cudaDeviceSynchronize(); - cudaFree(new_qweight); - } +void shuffle_exllama_weight(uint32_t* q_weight, int height, int width, + int bit) { dim3 blockDim, gridDim; blockDim.x = THREADS_X; blockDim.y = 1; @@ -1765,10 +1456,17 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, torch::stable::Tensor b_q_weight, torch::stable::Tensor b_gptq_qzeros, torch::stable::Tensor b_gptq_scales, - torch::stable::Tensor b_g_idx, bool use_exllama, - bool use_v2_format, int64_t bit) { + bool use_exllama, bool use_v2_format, + int64_t bit) { const torch::stable::accelerator::DeviceGuard device_guard( a.get_device_index()); + const int64_t num_groups = b_gptq_qzeros.size(0); + STD_TORCH_CHECK(num_groups > 0, "GPTQ requires at least one scale group"); + STD_TORCH_CHECK(b_gptq_scales.size(0) == num_groups, + "GPTQ scales and zero-points must have the same number of " + "groups"); + STD_TORCH_CHECK(a.size(1) % num_groups == 0, "Input size K = ", a.size(1), + " is not divisible by the number of groups = ", num_groups); auto c = torch::stable::new_zeros(a, {a.size(0), b_q_weight.size(1)}); auto temp_dq = torch::stable::empty({b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)}, @@ -1778,28 +1476,19 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, get_current_cuda_blas_handle(), (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), (const uint32_t*)b_gptq_qzeros.data_ptr(), - (const half*)b_gptq_scales.data_ptr(), - b_g_idx.device().type() == torch::stable::DeviceType::Meta - ? NULL - : (const int*)b_g_idx.data_ptr(), - (half*)c.data_ptr(), (half*)temp_dq.data_ptr(), - c.size(0), // m - c.size(1), // n - a.size(1), // k - b_gptq_qzeros.size(0), // group number - use_exllama, use_v2_format, bit); + (const half*)b_gptq_scales.data_ptr(), (half*)c.data_ptr(), + (half*)temp_dq.data_ptr(), + c.size(0), // m + c.size(1), // n + a.size(1), // k + num_groups, use_exllama, use_v2_format, bit); return c; } -void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm, - int64_t bit) { +void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit) { const torch::stable::accelerator::DeviceGuard device_guard( q_weight.get_device_index()); - vllm::gptq::shuffle_exllama_weight( - (uint32_t*)q_weight.data_ptr(), - q_perm.device().type() == torch::stable::DeviceType::Meta || - q_perm.numel() == 0 - ? NULL - : (int*)q_perm.data_ptr(), - q_weight.size(0) * 32 / bit, q_weight.size(1), bit); + vllm::gptq::shuffle_exllama_weight((uint32_t*)q_weight.data_ptr(), + q_weight.size(0) * 32 / bit, + q_weight.size(1), bit); } diff --git a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py index 2a03847989..0b908d86d5 100644 --- a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py +++ b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py @@ -75,14 +75,14 @@ QUANT_CONFIGS = [ "b_type": "kU4B8", "thread_configs": THREAD_CONFIGS, "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 0, 2, 4, 8], + "group_blocks": [-1, 2, 4, 8], }, # GPTQ-INT8 { "b_type": "kU8B128", "thread_configs": THREAD_CONFIGS, "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 0, 2, 4, 8], + "group_blocks": [-1, 2, 4, 8], }, # FP8 { diff --git a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu index cafa212bcc..5dab9fe932 100644 --- a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu +++ b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu @@ -11,11 +11,9 @@ namespace marlin { -template +template __global__ void gptq_marlin_repack_kernel( - uint32_t const* __restrict__ b_q_weight_ptr, - uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, + uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { constexpr int pack_factor = 32 / num_bits; @@ -44,31 +42,14 @@ __global__ void gptq_marlin_repack_kernel( extern __shared__ int4 sh[]; - constexpr int perm_size = target_tile_k_size / 4; - - int4* sh_perm_ptr = sh; - int4* sh_pipe_ptr = sh_perm_ptr; - if constexpr (has_perm) { - sh_pipe_ptr += perm_size; - } + int4* sh_pipe_ptr = sh; constexpr int tile_ints = target_tile_k_size / pack_factor; constexpr int stage_n_threads = target_tile_n_size / 4; - constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints; + constexpr int stage_k_threads = tile_ints; constexpr int stage_size = stage_k_threads * stage_n_threads; - auto load_perm_to_shared = [&](int k_tile_id) { - int first_k_int4 = (k_tile_id * target_tile_k_size) / 4; - - int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); - - if (threadIdx.x < perm_size) { - sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; - } - __syncthreads(); - }; - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { if (n_tile_id >= n_tiles) { cp_async_fence(); @@ -79,36 +60,17 @@ __global__ void gptq_marlin_repack_kernel( int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; - if constexpr (has_perm) { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; + if (threadIdx.x < stage_size) { + auto k_id = threadIdx.x / stage_n_threads; + auto n_id = threadIdx.x % stage_n_threads; - uint32_t const* sh_perm_int_ptr = - reinterpret_cast(sh_perm_ptr); + int first_k = k_tile_id * target_tile_k_size; + int first_k_packed = first_k / pack_factor; - int src_k = sh_perm_int_ptr[k_id]; - int src_k_packed = src_k / pack_factor; - - cp_async4( - &sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast(&( - b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); - } - - } else { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * target_tile_k_size; - int first_k_packed = first_k / pack_factor; - - cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast( - &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + - first_n + (n_id * 4)]))); - } + cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast( + &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + + first_n + (n_id * 4)]))); } cp_async_fence(); @@ -139,56 +101,34 @@ __global__ void gptq_marlin_repack_kernel( int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - uint32_t vals[8]; - if constexpr (has_perm) { - static_assert(!is_a_8bit); - for (int i = 0; i < 4; i++) { - int k_idx = tc_row + tc_offsets[i]; - - uint32_t src_k = sh_perm_int_ptr[k_idx]; - uint32_t src_k_pos = src_k % pack_factor; - - uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; - uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; - - uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; - uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; - - vals[i] = b1_cur_val; - vals[4 + i] = b2_cur_val; - } - - } else { - uint32_t b1_vals[tile_ints]; - uint32_t b2_vals[tile_ints]; + uint32_t b1_vals[tile_ints]; + uint32_t b2_vals[tile_ints]; #pragma unroll - for (int i = 0; i < tile_ints; i++) { - if constexpr (is_a_8bit) { - b1_vals[i] = - sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8]; - } else { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; - b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; - } + for (int i = 0; i < tile_ints; i++) { + if constexpr (is_a_8bit) { + b1_vals[i] = + sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8]; + } else { + b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; + b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; } + } #pragma unroll - for (int i = 0; i < 4; i++) { - int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]); - int cur_int = cur_elem / pack_factor; - int cur_pos = cur_elem % pack_factor; + for (int i = 0; i < 4; i++) { + int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]); + int cur_int = cur_elem / pack_factor; + int cur_pos = cur_elem % pack_factor; - vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; - if constexpr (is_a_8bit) - vals[4 + i] = - (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask; - else - vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; - } + vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; + if constexpr (is_a_8bit) + vals[4 + i] = + (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask; + else + vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; } constexpr int tile_size = @@ -248,10 +188,6 @@ __global__ void gptq_marlin_repack_kernel( for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { int n_tile_id = 0; - if constexpr (has_perm) { - load_perm_to_shared(k_tile_id); - } - start_pipes(k_tile_id, n_tile_id); while (n_tile_id < n_tiles) { @@ -269,21 +205,19 @@ __global__ void gptq_marlin_repack_kernel( } // namespace marlin -#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \ - else if (num_bits == NUM_BITS && has_perm == HAS_PERM && \ - is_a_8bit == IS_A_8BIT) { \ +#define CALL_IF(NUM_BITS, IS_A_8BIT) \ + else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) { \ cudaFuncSetAttribute( \ marlin::gptq_marlin_repack_kernel, \ + IS_A_8BIT>, \ cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ marlin::gptq_marlin_repack_kernel \ + IS_A_8BIT> \ <<>>( \ - b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ + b_q_weight_ptr, out_ptr, size_k, size_n); \ } torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, - torch::stable::Tensor& perm, int64_t size_k, int64_t size_n, int64_t num_bits, bool is_a_8bit) { // Verify compatibility with marlin tile of 16x64 @@ -311,11 +245,6 @@ torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, "b_q_weight type is not kInt"); - STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU"); - STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); - STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int, - "perm type is not at::kInt"); - const int32_t device_index = b_q_weight.get_device_index(); torch::stable::accelerator::DeviceGuard device_guard(device_index); const cudaStream_t stream = get_current_cuda_stream(device_index); @@ -325,14 +254,9 @@ torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); - // Detect if there is act_order - bool has_perm = perm.size(0) != 0; - // Get ptrs uint32_t const* b_q_weight_ptr = reinterpret_cast(b_q_weight.const_data_ptr()); - uint32_t const* perm_ptr = - reinterpret_cast(perm.const_data_ptr()); uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); int blocks; @@ -345,17 +269,15 @@ torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, if (false) { } - CALL_IF(4, false, false) - CALL_IF(4, true, false) - CALL_IF(8, false, false) - CALL_IF(8, true, false) + CALL_IF(4, false) + CALL_IF(8, false) - CALL_IF(4, false, true) - CALL_IF(8, false, true) + CALL_IF(4, true) + CALL_IF(8, true) else { STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, - ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); + ", is_a_8bit = ", is_a_8bit); } return out; diff --git a/csrc/libtorch_stable/quantization/marlin/kernel.h b/csrc/libtorch_stable/quantization/marlin/kernel.h index 8c9cec88b6..838ab97903 100644 --- a/csrc/libtorch_stable/quantization/marlin/kernel.h +++ b/csrc/libtorch_stable/quantization/marlin/kernel.h @@ -7,17 +7,16 @@ #include "marlin_dtypes.cuh" #include "core/scalar_type.hpp" -#define MARLIN_KERNEL_PARAMS \ - const int4 *__restrict__ A, const int4 *__restrict__ B, \ - int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ b_bias_ptr, \ - const float *__restrict__ a_scales_ptr, \ - const int4 *__restrict__ scales_ptr, \ - const float *__restrict__ global_scale_ptr, \ - const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ - int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \ - bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \ - int max_shared_mem +#define MARLIN_KERNEL_PARAMS \ + const int4 *__restrict__ A, const int4 *__restrict__ B, \ + int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ + const int4 *__restrict__ b_bias_ptr, \ + const float *__restrict__ a_scales_ptr, \ + const int4 *__restrict__ scales_ptr, \ + const float *__restrict__ global_scale_ptr, \ + const int4 *__restrict__ zp_ptr, int prob_m, int prob_n, int prob_k, \ + int lda, int *locks, bool has_bias, bool use_atomic_add, \ + bool use_fp32_reduce, int max_shared_mem namespace MARLIN_NAMESPACE_NAME { template const& a_scales_or_none, std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id, - int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full, - bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { + int64_t size_m, int64_t size_n, int64_t size_k, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float) { STD_TORCH_CHECK_NOT_IMPLEMENTED(false, "marlin_gemm(..) requires CUDA_ARCH >= 7.5"); return torch::stable::empty({1, 1}); @@ -74,62 +67,6 @@ torch::stable::Tensor marlin_gemm( #else -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, - int const* __restrict__ perm_int_ptr, - int4* __restrict__ out_int4_ptr, int size_m, - int size_k, int lda, int block_rows) { - auto start_row = block_rows * blockIdx.x; - int finish_row = start_row + block_rows; - if (finish_row > size_m) { - finish_row = size_m; - } - int cur_block_rows = finish_row - start_row; - - int input_row_stride = lda * sizeof(half) / 16; - int output_row_stride = size_k * sizeof(half) / 16; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int input_offset = row * input_row_stride; - int output_offset = row * output_row_stride; - - half const* a_row_half = - reinterpret_cast(a_int4_ptr + input_offset); - half* out_half = reinterpret_cast(out_int4_ptr + output_offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int i = 0; i < cur_block_rows; i++) { - int cur_row = start_row + i; - if (cur_row < size_m) { - permute_row(cur_row); - } - } -} - typedef struct { int thread_k; int thread_n; @@ -159,9 +96,7 @@ typedef struct { int get_scales_cache_size(thread_config_t const& th_config, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, - bool has_act_order, bool is_k_full, int stages) { - bool cache_scales_chunk = has_act_order && !is_k_full; - + int stages) { int tb_n = th_config.thread_n; int tb_k = th_config.thread_k; @@ -169,29 +104,18 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m, int tb_groups; if (group_size == -1) { tb_groups = 1; - } else if (group_size == 0) { - tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size } else { tb_groups = div_ceil(tb_k, group_size); } - if (cache_scales_chunk) { - int load_groups = - tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K - load_groups = max(load_groups, 32); // We load at least 32 scale groups - return load_groups * tb_n * 2; - } else { - int tb_scales = tb_groups * tb_n * 2; - - return tb_scales * stages; - } + int tb_scales = tb_groups * tb_n * 2; + return tb_scales * stages; } int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, int prob_m, int prob_n, int prob_k, int num_bits, - int group_size, bool has_act_order, bool is_k_full, - int has_zp, bool is_zp_float, bool is_a_8bit, - int stages) { + int group_size, int has_zp, bool is_zp_float, + bool is_a_8bit, int stages) { int pack_factor = 32 / num_bits; // Get B size @@ -206,10 +130,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); - int sh_s_size = - get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, - group_size, has_act_order, is_k_full, stages); - int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0; + int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, + num_bits, group_size, stages); int sh_zp_size = 0; if (has_zp) { if (is_zp_float) @@ -220,17 +142,15 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, sh_zp_size = sh_s_size / 2; } - int total_size = - tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; + int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size; return total_size; } bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks, int prob_m, int prob_n, int prob_k, int num_bits, - int group_size, bool has_act_order, bool is_k_full, - int has_zp, bool is_zp_float, bool is_a_8bit, int stages, - int max_shared_mem) { + int group_size, int has_zp, bool is_zp_float, + bool is_a_8bit, int stages, int max_shared_mem) { // Sanity if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { @@ -255,16 +175,18 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks, // Check that pipeline fits into cache int cache_size = get_kernel_cache_size( th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages); + has_zp, is_zp_float, is_a_8bit, stages); return cache_size <= max_shared_mem; } -MarlinFuncPtr get_marlin_kernel( - const vllm::ScalarType a_type, const vllm::ScalarType b_type, - const vllm::ScalarType c_type, const vllm::ScalarType s_type, - int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, - bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, - int threads, bool is_zp_float, int stages) { +MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType a_type, + const vllm::ScalarType b_type, + const vllm::ScalarType c_type, + const vllm::ScalarType s_type, + int thread_m_blocks, int thread_n_blocks, + int thread_k_blocks, bool m_block_size_8, + bool has_zp, int group_blocks, int threads, + bool is_zp_float, int stages) { int num_bits = b_type.size_bits(); auto kernel = MarlinDefault; @@ -277,9 +199,8 @@ exec_config_t determine_exec_config( const vllm::ScalarType& a_type, const vllm::ScalarType& b_type, const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m, int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, - int num_bits, int group_size, bool has_act_order, bool is_k_full, - bool has_zp, bool is_zp_float, int is_a_8bit, int stages, - int max_shared_mem, int sms) { + int num_bits, int group_size, bool has_zp, bool is_zp_float, int is_a_8bit, + int stages, int max_shared_mem, int sms) { exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; thread_config_t* thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs @@ -293,27 +214,21 @@ exec_config_t determine_exec_config( thread_config_t th_config = thread_configs[i]; if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k, - num_bits, group_size, has_act_order, is_k_full, has_zp, - is_zp_float, is_a_8bit, stages, - max_shared_mem - 512)) { + num_bits, group_size, has_zp, is_zp_float, is_a_8bit, + stages, max_shared_mem - 512)) { continue; } - int cache_size = get_kernel_cache_size(th_config, thread_m_blocks, prob_m, - prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, - is_zp_float, is_a_8bit, stages); + int cache_size = get_kernel_cache_size( + th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, + group_size, has_zp, is_zp_float, is_a_8bit, stages); - int group_blocks = 0; - if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : group_size / 16; - } + int group_blocks = group_size == -1 ? -1 : group_size / 16; - auto kernel = - get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks, - th_config.thread_n / 16, th_config.thread_k / 16, - m_block_size_8, has_act_order, has_zp, group_blocks, - th_config.num_threads, is_zp_float, stages); + auto kernel = get_marlin_kernel( + a_type, b_type, c_type, s_type, thread_m_blocks, + th_config.thread_n / 16, th_config.thread_k / 16, m_block_size_8, + has_zp, group_blocks, th_config.num_threads, is_zp_float, stages); if (kernel == MarlinDefault) continue; @@ -324,38 +239,25 @@ exec_config_t determine_exec_config( } void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, - void* a_s, void* b_s, void* g_s, void* zp, void* g_idx, - void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k, - int lda, void* workspace, vllm::ScalarType const& a_type, - vllm::ScalarType const& b_type, vllm::ScalarType const& c_type, - vllm::ScalarType const& s_type, bool has_bias, - bool has_act_order, bool is_k_full, bool has_zp, int num_groups, - int group_size, int dev, cudaStream_t stream, int thread_k_init, + void* a_s, void* b_s, void* g_s, void* zp, int prob_m, + int prob_n, int prob_k, int lda, void* workspace, + vllm::ScalarType const& a_type, vllm::ScalarType const& b_type, + vllm::ScalarType const& c_type, vllm::ScalarType const& s_type, + bool has_bias, bool has_zp, int num_groups, int group_size, + int dev, cudaStream_t stream, int thread_k_init, int thread_n_init, int sms, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { bool is_a_8bit = a_type.size_bits() == 8; STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]"); - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { - STD_TORCH_CHECK(group_size != -1); - group_blocks = group_size / 16; - STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } else { - STD_TORCH_CHECK(group_size == 0); - group_blocks = 0; - } + int group_blocks; + if (group_size == -1) { + group_blocks = -1; } else { - if (group_size == -1) { - group_blocks = -1; - } else { - group_blocks = group_size / 16; - STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } + group_blocks = group_size / 16; + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } int num_bits = b_type.size_bits(); @@ -370,28 +272,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, const float* g_s_ptr = (const float*)g_s; const int4* zp_ptr = (const int4*)zp; - const int* g_idx_ptr = (const int*)g_idx; - const int* perm_ptr = (const int*)perm; - int4* a_tmp_ptr = (int4*)a_tmp; int* locks = (int*)workspace; - if (has_act_order) { - // Permute A columns - int block_rows = div_ceil(prob_m, sms); - // avoid ">>>" being formatted to "> > >" - // clang-format off - permute_cols_kernel<<>>( - A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); - // clang-format on - A_ptr = a_tmp_ptr; - lda = prob_k; - - // If we have a full K, then we can run the non-act-order version of Marlin - // (since the weight rows are reordered by increasing group ids, and by - // having a full K, we have full original groups) - if (is_k_full) has_act_order = false; - } - int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); @@ -451,18 +333,17 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, // Auto config exec_cfg = determine_exec_config( a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k, - thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order, - is_k_full, has_zp, is_zp_float, is_a_8bit, stages, max_shared_mem, - sms); + thread_m_blocks, m_block_size_8, num_bits, group_size, has_zp, + is_zp_float, is_a_8bit, stages, max_shared_mem, sms); thread_tfg = exec_cfg.tb_cfg; if (thread_tfg.thread_n != -1) { if (prob_n / thread_tfg.thread_n * div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <= sms) { if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split, - prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, - is_a_8bit, stages, max_shared_mem_new)) { + prob_n, prob_k, num_bits, group_size, has_zp, + is_zp_float, is_a_8bit, stages, + max_shared_mem_new)) { thread_tfg = {128, 64, 128}; exec_cfg = {1, thread_tfg}; } @@ -487,30 +368,27 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, STD_TORCH_CHECK( is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, - prob_k, num_bits, group_size, has_act_order, is_k_full, - has_zp, is_zp_float, is_a_8bit, stages, - max_shared_mem_new), + prob_k, num_bits, group_size, has_zp, is_zp_float, + is_a_8bit, stages, max_shared_mem_new), "Invalid thread config: thread_m_blocks = ", thread_m_blocks, ", thread_k = ", thread_tfg.thread_k, ", thread_n = ", thread_tfg.thread_n, ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, ", prob_m_split = ", prob_m_split, ", group_size = ", group_size, - ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, ", stages = ", stages, ", max_shared_mem_new = ", max_shared_mem_new); auto kernel = get_marlin_kernel( a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks, - thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, - num_threads, is_zp_float, stages); + thread_k_blocks, m_block_size_8, has_zp, group_blocks, num_threads, + is_zp_float, stages); if (kernel == MarlinDefault) { STD_TORCH_CHECK( false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", - prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, - ", prob_m_split = ", prob_m_split, + prob_k, "]", ", num_groups = ", num_groups, + ", group_size = ", group_size, ", prob_m_split = ", prob_m_split, ", thread_m_blocks = ", thread_m_blocks, ", thread_n_blocks = ", thread_n_blocks, ", thread_k_blocks = ", thread_k_blocks, @@ -527,7 +405,6 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, // clang-format off kernel<<>>( A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, - g_idx_ptr, num_groups, prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add, use_fp32_reduce, max_shared_mem_new); // clang-format on @@ -550,11 +427,9 @@ torch::stable::Tensor marlin_gemm( std::optional const& a_scales_or_none, std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id, - int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full, - bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { + int64_t size_m, int64_t size_n, int64_t size_k, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float) { vllm::ScalarTypeId a_type_id, c_type_id, s_type_id; auto c_scalar_type = a.scalar_type(); @@ -722,7 +597,7 @@ torch::stable::Tensor marlin_gemm( std::nullopt, device); } - // Detect groupsize and act_order + // Detect group size. int num_groups = -1; int group_size = -1; @@ -733,51 +608,13 @@ torch::stable::Tensor marlin_gemm( " is not size_n = ", size_n); num_groups = b_scales.size(0); - torch::stable::Tensor g_idx, perm, a_tmp; - if (g_idx_or_none.has_value() && perm_or_none.has_value()) { - g_idx = g_idx_or_none.value(); - perm = perm_or_none.value(); - - STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); - STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); - STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); - - // Verify g_idx and perm - STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || - (g_idx.size(-1) == size_k && perm.size(-1) == size_k), - "Unexpected g_idx.size(-1) = ", g_idx.size(-1), - " and perm.size(-1) = ", perm.size(-1), - ", where size_k = ", size_k); + if (num_groups > 1) { + STD_TORCH_CHECK( + size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); + group_size = size_k / num_groups; } else { - g_idx = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); - perm = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); - a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); - } - bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0; - - if (has_act_order) { - a_tmp = torch::stable::empty({size_m, size_k}, c_scalar_type, std::nullopt, - device); - if (is_k_full) { - STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); - STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by num_groups = ", num_groups); - group_size = size_k / num_groups; - } else { - group_size = 0; - } - - } else { - a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); - if (num_groups > 1) { - STD_TORCH_CHECK( - size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); - group_size = size_k / num_groups; - } else { - group_size = -1; - } + group_size = -1; } torch::stable::Tensor global_scale; @@ -882,11 +719,9 @@ torch::stable::Tensor marlin_gemm( a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(), c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(), a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(), - global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), - g_idx.mutable_data_ptr(), perm.mutable_data_ptr(), - a_tmp.mutable_data_ptr(), size_m, size_n, size_k, a.stride(0), - workspace.mutable_data_ptr(), a_type, b_type, c_type, s_type, has_bias, - has_act_order, is_k_full, has_zp, num_groups, group_size, device_index, + global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), size_m, + size_n, size_k, a.stride(0), workspace.mutable_data_ptr(), a_type, b_type, + c_type, s_type, has_bias, has_zp, num_groups, group_size, device_index, get_current_cuda_stream(device_index), thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float); diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_template.h b/csrc/libtorch_stable/quantization/marlin/marlin_template.h index 32b8f8bdd5..962f15c994 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin_template.h +++ b/csrc/libtorch_stable/quantization/marlin/marlin_template.h @@ -51,10 +51,9 @@ template shared // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin( const int4* __restrict__ A, // fp16 input matrix of shape mxk @@ -63,11 +62,9 @@ __global__ void Marlin( int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape // (k/groupsize)xn - const int* __restrict__ g_idx, // int32 group indices of shape k - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k int* locks, // extra global storage for barrier synchronization bool use_fp32_reduce // whether to use fp32 global reduce ) {} @@ -140,29 +137,6 @@ __device__ inline void sub_zp( frag_b[1] = __hsub2(frag_b[1], zp); } -// Same as above, but for act_order (each K is multiplied individually) -template -__device__ inline void scale4( - typename MarlinScalarType::FragB& frag_b, - typename MarlinScalarType::FragS& frag_s_1, - typename MarlinScalarType::FragS& frag_s_2, - typename MarlinScalarType::FragS& frag_s_3, - typename MarlinScalarType::FragS& frag_s_4, int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - - scalar_t2 s_val_1_2; - s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; - s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; - - scalar_t2 s_val_3_4; - s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; - s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; - - frag_b[0] = __hmul2(frag_b[0], s_val_1_2); - frag_b[1] = __hmul2(frag_b[1], s_val_3_4); -} - // Given 2 floats multiply by 2 scales (halves) template __device__ inline void scale_float( @@ -256,14 +230,11 @@ __global__ void Marlin( // 4bit packed zero-points of shape // (k/groupsize, n/pack_factor) const int4* __restrict__ zp_ptr, - // int32 group indices of shape k - const int* __restrict__ g_idx, - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int lda, // A.stride(0), equal to prob_k is A is contiguous - int* locks, // extra global storage for barrier synchronization + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int lda, // A.stride(0), equal to prob_k is A is contiguous + int* locks, // extra global storage for barrier synchronization bool has_bias, bool use_atomic_add, // whether to use atomic add to reduce bool use_fp32_reduce, // whether to use fp32 global reduce @@ -358,7 +329,8 @@ __global__ void Marlin( global_scale_f32 = global_scale_ptr[0]; } - constexpr bool has_act_order = group_blocks == 0; + static_assert(group_blocks != 0, + "runtime activation ordering is not supported"); constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); extern __shared__ int4 sh[]; @@ -391,7 +363,7 @@ __global__ void Marlin( int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); - if constexpr (!has_act_order && group_blocks != -1) { + if constexpr (group_blocks != -1) { if (group_blocks >= thread_k_blocks) { // Ensure that the number of tiles in each stripe is a multiple of the // groupsize; this avoids an annoying special case where a stripe starts @@ -558,27 +530,17 @@ __global__ void Marlin( b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - // Scale sizes/strides without act_order + // Scale sizes/strides int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8); constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8); constexpr int s_tb_groups = - !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + group_blocks != -1 && group_blocks < thread_k_blocks ? thread_k_blocks / group_blocks : 1; constexpr int s_sh_stage = s_tb_groups * s_sh_stride; int s_gl_rd_delta = s_gl_stride; - // Scale size/strides with act_order - constexpr int tb_k = 16 * thread_k_blocks; - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; - // constexpr int act_s_row_stride = 1; - // int act_s_col_stride = act_s_row_stride * num_groups; - constexpr int act_s_max_num_groups = 32; - int act_s_col_stride = 1; - int act_s_col_warp_stride = act_s_col_stride * 8; - constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; // Zero-points sizes/strides int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; @@ -615,25 +577,16 @@ __global__ void Marlin( auto b_sh_rd = threadIdx.x * b_thread_vecs; b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); - // For act_order - int slice_k_start = tb_k * slice_row; - int slice_k_finish = slice_k_start + tb_k * slice_iters; - int slice_k_start_shared_fetch = slice_k_start; - int slice_n_offset = act_s_col_tb_stride * slice_col; - - // No act_order int s_gl_rd; - if constexpr (!has_act_order) { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + - s_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + - threadIdx.x / s_sh_stride) + - s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - } + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; } auto s_sh_wr = threadIdx.x; bool s_sh_wr_pred = threadIdx.x < s_sh_stage; @@ -757,10 +710,8 @@ __global__ void Marlin( : (sh_size_b_red_min + sh_bias_size); int4* sh_bias = sh_new + sh_size_b_red_min; - int4* sh_g_idx = sh_new + sh_b_red_bias_size; - int4* sh_zp = sh_g_idx + (stages * g_idx_stage); - constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) - : (stages * s_sh_stage); + int4* sh_zp = sh_new + sh_b_red_bias_size; + constexpr int sh_s_size = stages * s_sh_stage; int4* sh_s = sh_zp + (stages * zp_sh_stage); int4* sh_a = sh_s + sh_s_size; @@ -769,9 +720,8 @@ __global__ void Marlin( I4 frag_b_quant[2][b_thread_vecs]; FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragS frag_s[2][4]; // No act-order + FragS frag_s[2][4]; FragS frag_bias[2][4]; - FragS act_frag_s[2][4][4]; // For act-order int frag_qzp[2][num_ints_per_thread]; // Zero-points FragZP frag_zp; // Zero-points in fp16 FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ @@ -801,42 +751,6 @@ __global__ void Marlin( reinterpret_cast(frag_c)[i] = 0; }; - int sh_first_group_id = -1; - int sh_num_groups = -1; - - auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, - int last_group_id) { - sh_first_group_id = first_group_id; - sh_num_groups = last_group_id - first_group_id + 1; - - if (sh_num_groups > act_s_max_num_groups) { - sh_num_groups = act_s_max_num_groups; - } - - if (sh_first_group_id + sh_num_groups > num_groups) { - sh_num_groups = num_groups - sh_first_group_id; - } - - int row_offset = first_group_id * s_gl_stride; - - if (is_async) { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], - &scales_ptr[row_offset + (i * s_gl_stride) + - slice_n_offset + threadIdx.x]); - } - } - } else { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - sh_s[(i * s_sh_stride) + threadIdx.x] = - scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + - threadIdx.x]; - } - } - } - }; // Asynchronously fetch the next A, B and s tile from global to the next // shared memory pipeline location. auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { @@ -862,44 +776,27 @@ __global__ void Marlin( b_gl_rd += b_gl_rd_delta_o; - if constexpr (has_act_order) { - // Fetch g_idx thread-block portion - int full_pipe = a_off; - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; - if (cur_k < prob_k && cur_k < slice_k_finish) { - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; - int4 const* cur_g_idx_stage_ptr = - reinterpret_cast(&g_idx[cur_k]); - - if (threadIdx.x < g_idx_stage) { - cp_async4_pred(&sh_g_idx_stage[threadIdx.x], - &cur_g_idx_stage_ptr[threadIdx.x]); + // Only fetch scales if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); } + s_gl_rd += s_gl_rd_delta * s_tb_groups; } - } else { - if constexpr (group_blocks != -1) { - int4* sh_s_stage = sh_s + s_sh_stage * pipe; + } - // Only fetch scales if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (s_sh_wr_pred) { - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta * s_tb_groups; - } - } + if constexpr (has_zp && group_blocks != -1) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - if constexpr (has_zp && group_blocks != -1) { - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - // Only fetch zero points if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; + // Only fetch zero points if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (zp_sh_wr_pred) { + cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); } + zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; } } } @@ -947,150 +844,57 @@ __global__ void Marlin( } }; - bool is_same_group[stages]; - int same_group_id[stages]; - - auto init_same_group = [&](int pipe) { - if constexpr (!has_act_order) { - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - int group_id_1 = sh_g_idx_int_ptr[0]; - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; - - is_same_group[pipe] = group_id_1 == group_id_2; - same_group_id[pipe] = group_id_1; - }; - auto fetch_scales_to_registers = [&](int k, int full_pipe) { int pipe = full_pipe % stages; using IT1 = typename std::conditional_t; using IT0 = typename std::conditional_t; constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); - if constexpr (!has_act_order) { - // No act-order case - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 && dequant_skip_flop) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - } - } else if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0) { - if (k % b_sh_wr_iters == 0) { - int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; - } else { - reinterpret_cast(&frag_s[1])[0] = - reinterpret_cast(&frag_s[0])[0]; - } - } - } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks2; - - int4* sh_s_stage = sh_s + s_sh_stage * pipe; - - if constexpr (!is_8bit_scale) { - reinterpret_cast(&frag_s[k % 2])[0] = - sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; - } else { - reinterpret_cast(&frag_s[k % 2])[0] = - reinterpret_cast( - sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - } else if (group_blocks >= b_sh_wr_iters) { - if constexpr (!is_8bit_scale) { - reinterpret_cast(&frag_s[1])[0] = - reinterpret_cast(&frag_s[0])[0]; - } else { - reinterpret_cast(&frag_s[1])[0] = - reinterpret_cast(&frag_s[0])[0]; - } + if constexpr (group_blocks == -1) { + // Load only when starting a new slice. + if (k == 0 && full_pipe == 0 && dequant_skip_flop) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } else if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0) { + if (k % b_sh_wr_iters == 0) { + int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; } } + } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; - return; - } + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks2; - // Act-order case + int4* sh_s_stage = sh_s + s_sh_stage * pipe; - // Determine K of the "current" thread-block - int cur_k = slice_k_start + tb_k * full_pipe; - if (cur_k >= prob_k || cur_k >= slice_k_finish) { - return; - } - - // Reset (to current thread-block) since we read g_idx portion from the - // shared memory - cur_k = 0; - - // Progress to current iteration - cur_k += k % b_sh_wr_iters; - - // Determine "position" inside the thread-block (based on warp and - // thread-id) - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - int warp_col = warp_id % tb_n_warps; - - cur_k += warp_row * 16 * b_sh_wr_iters; - - auto th_id = threadIdx.x % 32; - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix - - int s_col_shift = - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + - (th_id / 4) * act_s_col_stride; - - if (is_same_group[pipe]) { - if (k % 2 == 0) { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = - sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + - s_col_shift]; + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; } else { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = - *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast( + sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; } - - for (int i = 1; i < 4; i++) { - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } else if (group_blocks >= b_sh_wr_iters) { + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; } - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - constexpr int k_frag_offsets[4] = {0, 1, 8, - 9}; // Tensor core offsets per thread - - #pragma unroll - for (int i = 0; i < 4; i++) { - int actual_k = cur_k + k_frag_offsets[i]; - - int group_id = sh_g_idx_int_ptr[actual_k]; - int rel_group_id = group_id - sh_first_group_id; - - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = - sh_s[rel_group_id * s_sh_stride + s_col_shift]; } }; auto fetch_zp_to_registers = [&](int k, int full_pipe) { - // This code does not handle group_blocks == 0, - // which signifies act_order. - // has_zp implies AWQ, which doesn't have act_order, static_assert(!has_zp || group_blocks != 0); if constexpr (has_zp && !is_zp_float) { @@ -1180,8 +984,7 @@ __global__ void Marlin( constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; const bool is_new_zp = - (group_blocks == 0) || - ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && + (group_blocks > 0 && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) || (group_blocks == -1 && is_first_matmul_in_slice); if constexpr (has_zp && !is_zp_float) { @@ -1249,15 +1052,8 @@ __global__ void Marlin( sub_zp(frag_b1, frag_zp[j], 1); } - // Apply scale to frag_b0 - if constexpr (has_act_order && !is_a_8bit) { - static_assert(group_blocks != -1); - scale4(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); - scale4(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); - } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && - group_blocks == -1 && !is_a_8bit) { + if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && + group_blocks == -1 && !is_a_8bit) { int idx = (threadIdx.x / 4) % 2; scalar_t2 s2 = Adtype::nums2num2( reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], @@ -1661,7 +1457,7 @@ __global__ void Marlin( // For per-column quantization we finally apply the scale here (only for // 4-bit) - if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit && + if constexpr (group_blocks == -1 && !is_a_8bit && b_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { c_scalar_t2 tmp_scale = s[0]; @@ -1750,15 +1546,6 @@ __global__ void Marlin( #pragma unroll for (int i = 0; i < stages - 1; i++) { - if (has_act_order && i == 0) { - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], - g_idx[last_g_idx]); - } - if constexpr (has_zp && !is_zp_float && group_blocks == -1) { if (i == 0) { fetch_col_zp_to_shared(); @@ -1772,14 +1559,10 @@ __global__ void Marlin( zero_accums(); wait_for_stage(); - init_same_group(0); fetch_to_registers(0, 0); fetch_scales_to_registers(0, 0); fetch_zp_to_registers(0, 0); a_gl_rd += a_gl_rd_delta_o * (stages - 1); - if constexpr (has_act_order) { - slice_k_start_shared_fetch += tb_k * (stages - 1); - } }; if (slice_iters) { start_pipes(); @@ -1804,7 +1587,6 @@ __global__ void Marlin( slice_iters >= stages); pipe++; wait_for_stage(); - init_same_group(pipe % stages); } if constexpr (!is_a_8bit) { @@ -1822,25 +1604,6 @@ __global__ void Marlin( a_gl_rd += a_gl_rd_delta_o * stages; - if constexpr (has_act_order) { - slice_k_start += tb_k * stages; - - if (slice_k_start < prob_k) { - slice_k_start_shared_fetch += tb_k * stages; - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) { - fetch_act_order_scales_to_shared(false, first_group_id, - last_group_id); - __syncthreads(); - } - } - } - // Process results and, if necessary, proceed to the next column slice. // While this pattern may not be the most readable, other ways of writing // the loop seemed to noticeably worse performance after compilation. @@ -1898,7 +1661,7 @@ __global__ void Marlin( bool last = slice_idx == slice_count - 1; // For per-column scales, we only fetch them here in the final step before // write-out - if constexpr (!has_act_order && group_blocks == -1 && + if constexpr (group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) { if (s_sh_wr_pred) { @@ -1917,7 +1680,7 @@ __global__ void Marlin( cp_async_fence(); } - if constexpr (!has_act_order && group_blocks == -1 && + if constexpr (group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) { if constexpr (is_a_8bit) { cp_async_wait<0>(); @@ -1948,7 +1711,7 @@ __global__ void Marlin( // For 8-bit channelwise, we apply the scale before the global reduction // that converts the fp32 results to fp16 (so that we avoid possible // overflow in fp16) - if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) { + if constexpr (group_blocks == -1 && is_a_8bit) { #pragma unroll for (int j = 0; j < 2; j++) { float2 aa[2]; @@ -1970,8 +1733,7 @@ __global__ void Marlin( } } } - } else if (!has_act_order && group_blocks == -1 && - b_type.size_bits() == 8 && + } else if (group_blocks == -1 && b_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { if (threadIdx.x / 32 < tb_n_warps) { #pragma unroll @@ -2042,33 +1804,25 @@ __global__ void Marlin( b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - // Update slice k/n for scales loading - if constexpr (has_act_order) { - slice_k_start = tb_k * slice_row; - slice_k_finish = slice_k_start + tb_k * slice_iters; - slice_k_start_shared_fetch = slice_k_start; - slice_n_offset = act_s_col_tb_stride * slice_col; + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; } else { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = - s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + - s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = - zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + - zp_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = - s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + - threadIdx.x / s_sh_stride) + - s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - zp_gl_rd = - zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + - threadIdx.x / zp_sh_stride) + - zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; } start_pipes(); } diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index e5fa535993..44b5191182 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -27,8 +27,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " "output_s, int group_size, float eps, float int8_min, float int8_max) -> " "()"); - ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor"); - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); #ifndef USE_ROCM @@ -85,15 +83,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " "Tensor? b_bias_or_none,Tensor b_scales, " "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " - "Tensor? " - "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, " - "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, " + "Tensor workspace, int b_type_id, " + "SymInt size_m, SymInt size_n, SymInt size_k, " "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); // conditionally compiled so impl registrations are in source file // gptq_marlin repack from GPTQ. ops.def( - "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, " + "gptq_marlin_repack(Tensor b_q_weight, " "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); // conditionally compiled so impl registrations are in source file @@ -679,12 +676,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // to prevent the meta function registry. ops.def( "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " - "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool " + "Tensor b_gptq_scales, bool use_exllama, bool " "use_v2_format, int bit) " "-> Tensor"); // Post processing for GPTQ. - ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()"); + ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()"); // Mamba selective scan kernel ops.def( @@ -735,8 +732,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("per_token_group_quant_int8", TORCH_BOX(&per_token_group_quant_int8)); - ops.impl("permute_cols", TORCH_BOX(&permute_cols)); - #ifndef USE_ROCM // CUTLASS scaled_mm ops ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index d4b2cd44b3..78751841ca 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -26,12 +26,11 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, torch::Tensor b_scales, - torch::Tensor b_g_idx, bool use_v2_format); + bool use_v2_format); torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, - torch::Tensor b_scales, - torch::Tensor b_g_idx, bool use_v2_format); + torch::Tensor b_scales, bool use_v2_format); void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c, torch::Tensor b_q_weight, torch::Tensor b_scales, diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu index fb178f6d3b..104af66d1c 100644 --- a/csrc/rocm/q_gemm_rdna3.cu +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -253,11 +253,13 @@ __forceinline__ __device__ void load4_scales(const T* scales_row, int n, // --------------------------------------------------------------------------- template -__global__ void gemm_q4_kernel_rdna3( - const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight, - const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, - T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { +__global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, + const uint32_t* __restrict__ b_q_weight, + const uint32_t* __restrict__ b_qzeros, + const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, + const int size_n, const int size_k, + const int groups, const int zero_offset) { const int t = threadIdx.x; const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; const int offset_m = blockIdx.y * M_COUNT; @@ -274,8 +276,8 @@ __global__ void gemm_q4_kernel_rdna3( constexpr int LDS_PAD = 8; __shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD]; - // Stage A: each thread loads 1 K element per M row into LDS (with optional - // act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map. + // Stage A: each thread loads 1 K element per M row into LDS. + // THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map. // For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m // are zero-padded so the dot product contribution is 0 (we then skip the // atomic write for those rows below). @@ -298,10 +300,7 @@ __global__ void gemm_q4_kernel_rdna3( T av; if (offset_m + m < size_m) { const T* a_row = a + (offset_m + m) * size_k; - if (b_q_perm) - av = a_row[b_q_perm[offset_k + t]]; - else - av = a_row[offset_k + t]; + av = a_row[offset_k + t]; } else { av = tzero(); // zero-pad invalid M rows } @@ -314,14 +313,6 @@ __global__ void gemm_q4_kernel_rdna3( // the LDS load above — but here all THREADS_X (=256) threads always do, // regardless of whether their `n` is in bounds. __syncthreads(); - } else if (b_q_perm) { - // bf16 M=1 fast path skips LDS, but its global read below is sequential - // and cannot apply act-order. When a permutation is present, stage the - // single A row through LDS (as fp16 / M>1 do) so the read picks it up. - // b_q_perm is block-uniform, so the __syncthreads is non-divergent. - if (offset_k + t < end_k) - block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]]; - __syncthreads(); } if (n >= size_n) return; @@ -452,14 +443,11 @@ __global__ void gemm_q4_kernel_rdna3( // fp32-aliased union. Storing as uint32 keeps the IR-level type // opaque so the inner v_dot2 cannot be folded to fp32 widening. // - // A is read direct from global (no LDS staging — see USE_LDS_A above), - // except under act-order, where it comes from the permuted LDS copy. + // A is read direct from global (no LDS staging — see USE_LDS_A above). pack4 a_pack; { const uint32_t* a_words = - b_q_perm - ? reinterpret_cast(&block_a[0][a_off]) - : reinterpret_cast(a + offset_k + a_off); + reinterpret_cast(a + offset_k + a_off); a_pack.u[0] = a_words[0]; a_pack.u[1] = a_words[1]; a_pack.u[2] = a_words[2]; @@ -623,8 +611,7 @@ __global__ void gemm_q4_kernel_rdna3( template __global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, - const int, const int, const int, - const int*) {} + const int, const int, const int) {} #endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__ @@ -635,17 +622,17 @@ __global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*, template void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, - const int* b_q_perm, T* c, int size_m, - int size_n, int size_k, int groups, - int zero_offset, cudaStream_t stream) { + T* c, int size_m, int size_n, int size_k, + int groups, int zero_offset, + cudaStream_t stream) { dim3 block(THREADS_X); dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4), (size_m + M_COUNT - 1) / M_COUNT, (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE); - gemm_q4_kernel_rdna3<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + gemm_q4_kernel_rdna3 + <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, + size_n, size_k, groups, zero_offset); } // Dispatch to the largest M_COUNT template that doesn't waste more than @@ -661,31 +648,30 @@ void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight, // only burn instructions on the last block, never affect correctness. template void launch_gemm_q4(const T* a, const uint32_t* b_q_weight, - const uint32_t* b_qzeros, const T* b_scales, - const int* b_q_perm, T* c, int size_m, int size_n, - int size_k, int groups, bool use_v2_format, - cudaStream_t stream) { + const uint32_t* b_qzeros, const T* b_scales, T* c, + int size_m, int size_n, int size_k, int groups, + bool use_v2_format, cudaStream_t stream) { const int zero_offset = use_v2_format ? 0 : 1; if (size_m == 1) { - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); } else if (size_m <= 3) { - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); } else if (size_m <= 7) { - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); } else { // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the // WMMA path, but if it falls through we still produce correct output — // just leaving 3-5× of throughput on the table for prefill workloads. - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); } } @@ -701,7 +687,6 @@ void launch_gemm_q4(const T* a, const uint32_t* b_q_weight, // b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle) // b_qzeros [groups, N/8] uint32 (packed 4-bit zeros) // b_scales [groups, N] half or bfloat16 -// b_g_idx [K] or empty int32 (act-order permutation; empty=identity) // use_v2_format bool (true = GPTQv2, no +1 zero offset) // // Output: @@ -709,17 +694,16 @@ void launch_gemm_q4(const T* a, const uint32_t* b_q_weight, torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, - torch::Tensor b_scales, - torch::Tensor b_g_idx, bool use_v2_format); + torch::Tensor b_scales, bool use_v2_format); torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, torch::Tensor b_scales, - torch::Tensor b_g_idx, bool use_v2_format) { + bool use_v2_format) { if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 && b_q_weight.size(1) % 16 == 0 && ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) || (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) { - return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx, + return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, use_v2_format); } @@ -753,25 +737,18 @@ torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); at::Tensor c = torch::zeros({size_m, size_n}, opts); - const int* g_idx_ptr = nullptr; - if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) { - TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32, - "b_g_idx must be int32"); - g_idx_ptr = (const int*)b_g_idx.data_ptr(); - } - if (a.scalar_type() == torch::kHalf) { vllm::gptq_rdna3::launch_gemm_q4( (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), - g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups, - use_v2_format, stream); + (half*)c.data_ptr(), size_m, size_n, size_k, groups, use_v2_format, + stream); } else { vllm::gptq_rdna3::launch_gemm_q4( (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), (const uint32_t*)b_qzeros.data_ptr(), - (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr, + (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), (vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups, use_v2_format, stream); } diff --git a/csrc/rocm/q_gemm_rdna3_wmma.cu b/csrc/rocm/q_gemm_rdna3_wmma.cu index 966d5df403..3876cd2225 100644 --- a/csrc/rocm/q_gemm_rdna3_wmma.cu +++ b/csrc/rocm/q_gemm_rdna3_wmma.cu @@ -373,7 +373,7 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int groups, const int zero_offset) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -471,25 +471,16 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w( if (m_row < size_m) { const T* a_row = a + m_row * size_k; - if (b_q_perm) { - // Permuted (act-order): scattered global reads, no vectorization. - #pragma unroll - for (int i = 0; i < 16; i++) { - T v = a_row[b_q_perm[k_tile + i]]; - a_frag[i] = bitcast_elem(v); - } - } else { - // Sequential A reads: replace 16 single-element global_load_b16 with - // a bulk 32-byte copy. The AMDGPU backend lowers a memcpy of this - // size + alignment to two `global_load_b128` instructions. size_k is - // a multiple of 16 (TORCH_CHECK above) and k_tile increments by 16, - // so k_tile + 16 is always within bounds — no tail handling needed. - // Note: we memcpy into the whole vector (`&a_frag`) rather than - // `&a_frag[0]`; ext_vector_type element addresses aren't reliably - // valid C pointers across compiler versions. - static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)"); - __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); - } + // Sequential A reads: replace 16 single-element global_load_b16 with + // a bulk 32-byte copy. The AMDGPU backend lowers a memcpy of this + // size + alignment to two `global_load_b128` instructions. size_k is + // a multiple of 16 (TORCH_CHECK above) and k_tile increments by 16, + // so k_tile + 16 is always within bounds — no tail handling needed. + // Note: we memcpy into the whole vector (`&a_frag`) rather than + // `&a_frag[0]`; ext_vector_type element addresses aren't reliably + // valid C pointers across compiler versions. + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); } else { #pragma unroll for (int i = 0; i < 16; i++) a_frag[i] = (E)0; @@ -597,25 +588,24 @@ template __global__ void gemm_q4_wmma_kernel_16x16_1w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int, const int*) { -} + const int, const int) {} #endif template void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, - const int* b_q_perm, T* c, int size_m, - int size_n, int size_k, int groups, - int zero_offset, cudaStream_t stream) { + T* c, int size_m, int size_n, int size_k, + int groups, int zero_offset, + cudaStream_t stream) { // 1 wave per block (32 lanes), 16x16 C tile per block. gridDim.z splits // K so that more blocks (and therefore more waves) are in flight; with // K_SPLIT > 1 the kernel switches to atomic write-back at the epilogue. const int k_split = compute_wmma_k_split(size_k); dim3 block(32); dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split); - gemm_q4_wmma_kernel_16x16_1w<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + gemm_q4_wmma_kernel_16x16_1w + <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, + size_n, size_k, groups, zero_offset); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -657,7 +647,7 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int groups, const int zero_offset) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -757,16 +747,8 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w( V16 a_frag, b_frag; if (m_row < size_m) { const T* a_row = a + m_row * size_k; - if (b_q_perm) { - #pragma unroll - for (int i = 0; i < 16; i++) { - T v = a_row[b_q_perm[k_tile + i]]; - a_frag[i] = bitcast_elem(v); - } - } else { - static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); - __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); - } + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); } else { #pragma unroll for (int i = 0; i < 16; i++) a_frag[i] = (E)0; @@ -844,16 +826,15 @@ template __global__ void gemm_q4_wmma_kernel_32x16_2w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int, const int*) { -} + const int, const int) {} #endif template void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, - const int* b_q_perm, T* c, int size_m, - int size_n, int size_k, int groups, - int zero_offset, cudaStream_t stream) { + T* c, int size_m, int size_n, int size_k, + int groups, int zero_offset, + cudaStream_t stream) { // Fallback to v1 for size_m < 32. With M-tile=32 the v2 block has 2 waves // working on rows [0..15] and [16..31]; at M < 32 the second wave processes // out-of-range M rows (zero-padded a_frag → wmma produces nothing useful) @@ -862,9 +843,9 @@ void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight, // max-num-seqs=32 lands at M≈32 steady-state; the M=16 sliver is edge), // but the fallback costs nothing and is the right shape. if (size_m < 32) { - launch_gemm_q4_wmma_16x16_1w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_wmma_16x16_1w(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); return; } @@ -875,9 +856,9 @@ void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight, const int k_split = compute_wmma_k_split(size_k); dim3 block(64); dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split); - gemm_q4_wmma_kernel_32x16_2w<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + gemm_q4_wmma_kernel_32x16_2w + <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, + size_n, size_k, groups, zero_offset); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -917,7 +898,7 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int groups, const int zero_offset) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1019,16 +1000,8 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w( V16 a_frag, b_frag; if (m_row < size_m) { const T* a_row = a + m_row * size_k; - if (b_q_perm) { - #pragma unroll - for (int i = 0; i < 16; i++) { - T v = a_row[b_q_perm[k_tile + i]]; - a_frag[i] = bitcast_elem(v); - } - } else { - static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); - __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); - } + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); } else { #pragma unroll for (int i = 0; i < 16; i++) a_frag[i] = (E)0; @@ -1100,21 +1073,20 @@ template __global__ void gemm_q4_wmma_kernel_64x16_4w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int, const int*) { -} + const int, const int) {} #endif template void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, - const int* b_q_perm, T* c, int size_m, - int size_n, int size_k, int groups, - int zero_offset, cudaStream_t stream) { + T* c, int size_m, int size_n, int size_k, + int groups, int zero_offset, + cudaStream_t stream) { // Fall back to v2 for M < 64 (would waste 1+ waves on out-of-range rows). if (size_m < 64) { - launch_gemm_q4_wmma_32x16_2w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_wmma_32x16_2w(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); return; } @@ -1122,9 +1094,9 @@ void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight, const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 16); dim3 block(128); dim3 grid((size_n + 15) / 16, (size_m + 63) / 64, k_split); - gemm_q4_wmma_kernel_64x16_4w<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + gemm_q4_wmma_kernel_64x16_4w + <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, + size_n, size_k, groups, zero_offset); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -1158,7 +1130,7 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int groups, const int zero_offset) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1264,16 +1236,8 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w( V16 a_frag, b_frag0, b_frag1; if (m_row < size_m) { const T* a_row = a + m_row * size_k; - if (b_q_perm) { - #pragma unroll - for (int i = 0; i < 16; i++) { - T v = a_row[b_q_perm[k_tile + i]]; - a_frag[i] = bitcast_elem(v); - } - } else { - static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); - __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); - } + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); } else { #pragma unroll for (int i = 0; i < 16; i++) a_frag[i] = (E)0; @@ -1352,23 +1316,22 @@ template __global__ void gemm_q4_wmma_kernel_64x32_4w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int, const int*) { -} + const int, const int) {} #endif template void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, - const int* b_q_perm, T* c, int size_m, - int size_n, int size_k, int groups, - int zero_offset, cudaStream_t stream) { + T* c, int size_m, int size_n, int size_k, + int groups, int zero_offset, + cudaStream_t stream) { // Fall back to v3 when M < 64 (small-M decode/prefill stays on the // narrower 64M × 16N path) or when N < 32 (tile would waste a wave on // out-of-range cols). if (size_m < 64 || size_n < 32) { - launch_gemm_q4_wmma_64x16_4w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_wmma_64x16_4w(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); return; } @@ -1376,9 +1339,9 @@ void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight, const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 32); dim3 block(128); dim3 grid((size_n + 31) / 32, (size_m + 63) / 64, k_split); - gemm_q4_wmma_kernel_64x32_4w<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + gemm_q4_wmma_kernel_64x32_4w + <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, + size_n, size_k, groups, zero_offset); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -1412,7 +1375,7 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int groups, const int zero_offset) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1510,16 +1473,8 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w( V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3; if (m_row < size_m) { const T* a_row = a + m_row * size_k; - if (b_q_perm) { - #pragma unroll - for (int i = 0; i < 16; i++) { - T v = a_row[b_q_perm[k_tile + i]]; - a_frag[i] = bitcast_elem(v); - } - } else { - static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); - __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); - } + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); } else { #pragma unroll for (int i = 0; i < 16; i++) a_frag[i] = (E)0; @@ -1611,7 +1566,7 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int groups, const int zero_offset) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1712,8 +1667,7 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16( dequant_into(next_buf, k_next); } - // A-load: vectorized 256-bit. b_q_perm branch removed from V7 hot path - // to eliminate ~450 ISA instructions of dead code (6× icache bloat). + // A-load: vectorized 256-bit. V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3; if (a_row_ptr) { __builtin_memcpy(&a_frag, a_row_ptr + k_tile, sizeof(a_frag)); @@ -1810,7 +1764,7 @@ __global__ void gemm_q4_wmma_kernel_128x64_k32( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int groups, const int zero_offset) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -2006,40 +1960,36 @@ template __global__ void gemm_q4_wmma_kernel_64x64_4w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int, const int*) { -} + const int, const int) {} template __global__ void gemm_q4_wmma_kernel_128x64_k16(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int, - const int*) {} + const int, const int) {} template __global__ void gemm_q4_wmma_kernel_128x64_k32(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int, - const int*) {} + const int, const int) {} #endif template void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, - const int* b_q_perm, T* c, int size_m, - int size_n, int size_k, int groups, - int zero_offset, cudaStream_t stream) { + T* c, int size_m, int size_n, int size_k, + int groups, int zero_offset, + cudaStream_t stream) { // Fall back to v4 when N < 64 (would waste 1+ waves on out-of-range cols). if (size_m < 64 || size_n < 64) { - launch_gemm_q4_wmma_64x32_4w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, - c, size_m, size_n, size_k, groups, - zero_offset, stream); + launch_gemm_q4_wmma_64x32_4w(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, zero_offset, + stream); return; } // V8 (128M × 64N, K=32/iter, 8-wave dequant) when K%32==0 and gs≥32. - // Falls back to V7 otherwise. V7/V8 read A sequentially, so act-order - // (b_q_perm != null) must skip them and use v5, which honors the perm. - if (size_m >= 128 && b_q_perm == nullptr) { + // Falls back to V7 otherwise. + if (size_m >= 128) { const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 128, 64); const int groupsize = size_k / groups; @@ -2048,11 +1998,11 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) { gemm_q4_wmma_kernel_128x64_k32<<>>( a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + zero_offset); } else { gemm_q4_wmma_kernel_128x64_k16<<>>( a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + zero_offset); } return; } @@ -2061,9 +2011,9 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 64); dim3 block(128); dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split); - gemm_q4_wmma_kernel_64x64_4w<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset, b_q_perm); + gemm_q4_wmma_kernel_64x64_4w + <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, + size_n, size_k, groups, zero_offset); } } // namespace gptq_rdna3_wmma @@ -2078,7 +2028,6 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, // b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle) // b_qzeros [groups, N/8] uint32 (packed 4-bit zeros) // b_scales [groups, N] half or bfloat16 -// b_g_idx [K] or empty int32 (act-order permutation; empty=identity) // use_v2_format bool (true = GPTQv2, no +1 zero offset) // // Output: @@ -2091,8 +2040,7 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, - torch::Tensor b_scales, - torch::Tensor b_g_idx, bool use_v2_format) { + torch::Tensor b_scales, bool use_v2_format) { TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor"); @@ -2128,13 +2076,6 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, // negligible (< 1.5% of prefill time on gfx1100). at::Tensor c = torch::zeros({size_m, size_n}, opts); - const int* g_idx_ptr = nullptr; - if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) { - TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32, - "b_g_idx must be int32"); - g_idx_ptr = (const int*)b_g_idx.data_ptr(); - } - const int zero_offset = use_v2_format ? 0 : 1; // launch_gemm_q4_wmma_64x64_4w dispatches: @@ -2148,15 +2089,15 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w( (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), - g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups, - zero_offset, stream); + (half*)c.data_ptr(), size_m, size_n, size_k, groups, zero_offset, + stream); } else { vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w< vllm::gptq_rdna3_wmma::bf16_t>( (const vllm::gptq_rdna3_wmma::bf16_t*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), (const uint32_t*)b_qzeros.data_ptr(), - (const vllm::gptq_rdna3_wmma::bf16_t*)b_scales.data_ptr(), g_idx_ptr, + (const vllm::gptq_rdna3_wmma::bf16_t*)b_scales.data_ptr(), (vllm::gptq_rdna3_wmma::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups, zero_offset, stream); } diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 92dbed7f02..a8114ce196 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -56,12 +56,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). rocm_ops.def( "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " - "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor"); + "Tensor b_scales, bool use_v2_format) -> Tensor"); rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3); rocm_ops.def( "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " - "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor"); + "Tensor b_scales, bool use_v2_format) -> Tensor"); rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma); rocm_ops.def( diff --git a/tests/compile/fullgraph/test_full_graph.py b/tests/compile/fullgraph/test_full_graph.py index 4f0a78e25c..4ce7d6767f 100644 --- a/tests/compile/fullgraph/test_full_graph.py +++ b/tests/compile/fullgraph/test_full_graph.py @@ -30,7 +30,7 @@ def models_list(*, all: bool = True, keywords: list[str] | None = None): if is_quant_method_supported("gptq_marlin"): TEST_MODELS.append( ( - "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", + "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", {"quantization": "gptq_marlin"}, ) ) diff --git a/tests/kernels/core/test_permute_cols.py b/tests/kernels/core/test_permute_cols.py deleted file mode 100644 index 929db34ae2..0000000000 --- a/tests/kernels/core/test_permute_cols.py +++ /dev/null @@ -1,21 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -from tests.kernels.utils import opcheck -from vllm._custom_ops import permute_cols - -if not hasattr(torch.ops._C, "permute_cols"): - pytest.skip(reason="permute_cols is not supported on ROCm", allow_module_level=True) - - -@pytest.mark.parametrize("shape", [(1, 512), (544, 4096), (67, 8192)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_permute_cols(shape, dtype): - x = torch.randn(shape, dtype=dtype).cuda() - perm = torch.randperm(x.shape[1]).to(torch.int).cuda() - opcheck(torch.ops._C.permute_cols, (x, perm)) - y = permute_cols(x, perm) - torch.testing.assert_close(y, x[:, perm]) diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index 0f80ca5c55..6756d479fe 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -113,9 +113,7 @@ def marlin_quantize_moe_weights( for i in range(e): # Transpose for Marlin: [n, k] → [k, n] w_t = weights_bf16[i].T.contiguous() - _, w_q, w_s, _, _, _ = marlin_quantize( - w_t, scalar_types.uint4b8, group_size, act_order=False - ) + _, w_q, w_s = marlin_quantize(w_t, scalar_types.uint4b8, group_size) weight_list.append(w_q) scale_list.append(w_s) @@ -248,16 +246,11 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group expert_map=None, global_scale1=None, global_scale2=None, - g_idx1=None, - g_idx2=None, input_global_scale1=None, input_global_scale2=None, - sort_indices1=None, - sort_indices2=None, w1_zeros=None, w2_zeros=None, input_dtype=dtype, - is_k_full=True, ) # Sanity check: manually compute BF16 reference for comparison diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index a288e0557a..ea732d6887 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -178,13 +178,11 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [ # GPTQ-INT4 { "b_type": scalar_types.uint4b8, - "support_act_order": True, "group_blocks": [-1, 2, 4, 8], }, # GPTQ-INT8 { "b_type": scalar_types.uint8b128, - "support_act_order": True, "group_blocks": [-1, 2, 4, 8], }, # FP8 @@ -794,27 +792,22 @@ def test_fused_moe_wn16( MARLIN_MOE_SCENARIOS = [ - # (m, n, k, e, topk, ep_size, act_order, is_k_full) - # No act_order: is_k_full=True matches usual case (marlin_is_k_full). + # (m, n, k, e, topk, ep_size) # N>=256 required for Marlin kernel thread config for MXFP8. # Single token, small matrices - (1, 128, 256, 5, 2, 1, False, True), + (1, 128, 256, 5, 2, 1), # Single token, large matrices - (1, 1024, 2048, 5, 2, 1, False, True), + (1, 1024, 2048, 5, 2, 1), # Unaligned m, small matrices - (133, 256, 256, 5, 2, 1, False, True), + (133, 256, 256, 5, 2, 1), # Unaligned m, large matrices - (133, 1024, 2048, 12, 3, 1, False, True), + (133, 1024, 2048, 12, 3, 1), # Aligned batch, small matrices - (128, 256, 256, 5, 2, 1, False, True), + (128, 256, 256, 5, 2, 1), # Aligned batch, large matrices - (128, 1024, 2048, 12, 3, 1, False, True), + (128, 1024, 2048, 12, 3, 1), # Expert parallelism - (64, 1024, 2048, 12, 3, 4, False, True), - # Act order with is_k_full=True (no tensor parallelism) - (1, 1024, 2048, 5, 2, 1, True, True), - # Act order with is_k_full=False (tensor parallelism) - (133, 256, 256, 5, 2, 1, True, False), + (64, 1024, 2048, 12, 3, 4), ] @@ -832,17 +825,13 @@ def marlin_moe_generate_valid_test_cases(): e, topk, ep_size, - act_order, - is_k_full, ): group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 if group_size > 0 and k % group_size != 0: return False - if act_order and group_size in [-1, k, n]: - return False if group_size in [k, n]: return False - if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full: + if b_type == scalar_types.float8_e4m3fn and group_size == 32: return False return a_type.size_bits < 16 or a_type is c_type @@ -858,8 +847,6 @@ def marlin_moe_generate_valid_test_cases(): ) ) - supports_act_order = quant_test_config.get("support_act_order", False) - for sub_case in inner_combinations: if ( sub_case[0] == scalar_types.float8_e4m3fn @@ -869,10 +856,8 @@ def marlin_moe_generate_valid_test_cases(): continue for scenario in MARLIN_MOE_SCENARIOS: - m, n, k, e, topk, ep_size, act_order, is_k_full = scenario - if act_order and not supports_act_order: - continue - args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full) + m, n, k, e, topk, ep_size = scenario + args = sub_case + (m, n, k, e, topk, ep_size) if is_valid(*args): cases.append(args) return cases @@ -885,9 +870,7 @@ class MarlinMoEWeightData: scales: torch.Tensor global_scale: torch.Tensor | None a_scales_factor: torch.Tensor | None - g_idx: torch.Tensor | None zeros: torch.Tensor | None - sort_indices: torch.Tensor | None marlin_bias: torch.Tensor | None @staticmethod @@ -895,15 +878,12 @@ class MarlinMoEWeightData: w: torch.Tensor, quant_type: ScalarType, group_size: int, - act_order: bool | None = None, bias: torch.Tensor | None = None, input_type: ScalarType = None, ) -> "MarlinMoEWeightData": assert w.ndim == 3 has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] - k = w.shape[-1] - if input_type == scalar_types.int8: input_dtype = torch.int8 elif input_type == scalar_types.float8_e4m3fn: @@ -916,8 +896,6 @@ class MarlinMoEWeightData: scales_l: list[torch.Tensor] = [] global_scale_l: list[torch.Tensor] = [] zeros_l: list[torch.Tensor] = [] - g_idx_l: list[torch.Tensor] = [] - sort_indices_l: list[torch.Tensor] = [] bias_l: list[torch.Tensor] = [] for i in range(w.shape[0]): @@ -959,21 +937,16 @@ class MarlinMoEWeightData: scales_l.append(scales) zeros_l.append(zeros) else: - test_perm = torch.randperm(k) - w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize( + w_ref, qweight, scales = marlin_quantize( w[i].transpose(1, 0), quant_type, group_size, - act_order, - test_perm, input_dtype=input_dtype, ) w_ref_l.append(w_ref.T) qweight_l.append(qweight) scales_l.append(scales) - g_idx_l.append(g_idx) - sort_indices_l.append(sort_indices) if bias is not None: bias_l.append(marlin_permute_bias(bias[i])) @@ -982,9 +955,7 @@ class MarlinMoEWeightData: qweight = stack_and_dev(qweight_l).contiguous() scales = stack_and_dev(scales_l) global_scale = stack_and_dev(global_scale_l) if global_scale_l else None - g_idx = stack_and_dev(g_idx_l) if g_idx_l else None zeros = stack_and_dev(zeros_l) if zeros_l else None - sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None marlin_bias = stack_and_dev(bias_l) if bias_l else None a_scales_factor = None @@ -999,19 +970,14 @@ class MarlinMoEWeightData: scales=scales, global_scale=global_scale, a_scales_factor=a_scales_factor, - g_idx=g_idx, zeros=zeros, - sort_indices=sort_indices, marlin_bias=marlin_bias, ) @pytest.mark.flaky(reruns=2) @pytest.mark.parametrize( - ( - "a_type, b_type, c_type, group_blocks," - "m, n, k, e, topk, ep_size, act_order, is_k_full" - ), + ("a_type, b_type, c_type, group_blocks,m, n, k, e, topk, ep_size"), marlin_moe_generate_valid_test_cases(), ) @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") @@ -1027,8 +993,6 @@ def test_fused_marlin_moe( e: int, topk: int, ep_size: int, - act_order: bool, - is_k_full: bool, ): set_random_seed(1) group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 @@ -1065,7 +1029,6 @@ def test_fused_marlin_moe( w=w1, quant_type=b_type, group_size=group_size, - act_order=act_order, input_type=a_type, ) @@ -1073,7 +1036,6 @@ def test_fused_marlin_moe( w=w2, quant_type=b_type, group_size=group_size, - act_order=act_order, input_type=a_type, ) @@ -1127,17 +1089,12 @@ def test_fused_marlin_moe( expert_map=e_map, global_scale1=w1_data.global_scale, global_scale2=w2_data.global_scale, - g_idx1=w1_data.g_idx, - g_idx2=w2_data.g_idx, input_global_scale1=w1_data.a_scales_factor, input_global_scale2=w2_data.a_scales_factor, - sort_indices1=w1_data.sort_indices, - sort_indices2=w2_data.sort_indices, w1_zeros=w1_data.zeros, w2_zeros=w2_data.zeros, input_dtype=a_dtype, quant_type_id=b_type.id, - is_k_full=is_k_full, activation_func=instance_activation, ) @@ -1154,8 +1111,6 @@ def test_fused_marlin_moe_with_bias(m): e, topk = 32, 4 n, k = 2048, 2048 group_size = 128 - act_order = False - is_k_full = True quant_type = scalar_types.uint4b8 dtype = torch.half @@ -1169,7 +1124,6 @@ def test_fused_marlin_moe_with_bias(m): w=w1, quant_type=quant_type, group_size=group_size, - act_order=act_order, bias=b_bias1, ) @@ -1177,7 +1131,6 @@ def test_fused_marlin_moe_with_bias(m): w=w2, quant_type=quant_type, group_size=group_size, - act_order=act_order, bias=b_bias2, ) @@ -1204,14 +1157,9 @@ def test_fused_marlin_moe_with_bias(m): expert_map=None, global_scale1=w1_data.global_scale, global_scale2=w2_data.global_scale, - g_idx1=w1_data.g_idx, - g_idx2=w2_data.g_idx, - sort_indices1=w1_data.sort_indices, - sort_indices2=w2_data.sort_indices, w1_zeros=w1_data.zeros, w2_zeros=w2_data.zeros, quant_type_id=quant_type.id, - is_k_full=is_k_full, ) torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0) @@ -1235,7 +1183,6 @@ def test_fused_marlin_moe_non_gated( set_random_seed(42) group_size = 16 # NVFP4 group size - is_k_full = True quant_type = scalar_types.float4_e2m1f dtype = torch.bfloat16 @@ -1248,14 +1195,12 @@ def test_fused_marlin_moe_non_gated( w=w1, quant_type=quant_type, group_size=group_size, - act_order=False, ) w2_data = MarlinMoEWeightData.make( w=w2, quant_type=quant_type, group_size=group_size, - act_order=False, ) score = torch.randn((m, e), device="cuda", dtype=dtype) @@ -1286,14 +1231,9 @@ def test_fused_marlin_moe_non_gated( expert_map=None, global_scale1=w1_data.global_scale, global_scale2=w2_data.global_scale, - g_idx1=w1_data.g_idx, - g_idx2=w2_data.g_idx, - sort_indices1=w1_data.sort_indices, - sort_indices2=w2_data.sort_indices, w1_zeros=w1_data.zeros, w2_zeros=w2_data.zeros, quant_type_id=quant_type.id, - is_k_full=is_k_full, activation=activation, ) @@ -1612,14 +1552,9 @@ def test_batched_marlin_activation_uses_expert_token_counts( quant_type_id=0, a1_gscale=None, a2_gscale=None, - w13_g_idx=None, - w2_g_idx=None, - w13_g_idx_sort_indices=None, - w2_g_idx_sort_indices=None, w1_zp=None, w2_zp=None, input_dtype=None, - is_k_full=True, activation_config=ApplyMoEActivationConfig( clamp_limit=3.0, alpha=1.3, beta=0.5 ), @@ -1982,14 +1917,12 @@ def test_batched_fused_marlin_moe( w=w1, quant_type=quant_dtype, group_size=group_size, - act_order=None, input_type=input_type, ) w2_data = MarlinMoEWeightData.make( w=w2, quant_type=quant_dtype, group_size=group_size, - act_order=None, input_type=input_type, ) @@ -2101,14 +2034,9 @@ def test_batched_fused_marlin_moe( "expert_map": None, "global_scale1": w1_data.global_scale, "global_scale2": w2_data.global_scale, - "g_idx1": w1_data.g_idx, - "g_idx2": w2_data.g_idx, - "sort_indices1": w1_data.sort_indices, - "sort_indices2": w2_data.sort_indices, "w1_zeros": w1_data.zeros, "w2_zeros": w2_data.zeros, "quant_type_id": quant_dtype.id, - "is_k_full": True, } if input_dtype is not None: kwargs["input_dtype"] = input_dtype diff --git a/tests/kernels/quantization/test_allspark_gemm.py b/tests/kernels/quantization/test_allspark_gemm.py index b6272557ce..7e4ce1afaf 100644 --- a/tests/kernels/quantization/test_allspark_gemm.py +++ b/tests/kernels/quantization/test_allspark_gemm.py @@ -70,7 +70,7 @@ def test_gptq_allspark_gemm_ampere(mnk_factors, group_size, has_zp, dtype): input = rand_data((m, k), dtype=dtype) weight = rand_data((k, n), dtype=dtype) - # Quantize (and apply act_order if provided) + # Quantize weights. w_ref, qw, s, zp = quantize_weights( weight, scalar_types.uint8b128, group_size, has_zp ) diff --git a/tests/kernels/quantization/test_gptq.py b/tests/kernels/quantization/test_gptq.py index 7bc7f97ce7..7795fcc95f 100644 --- a/tests/kernels/quantization/test_gptq.py +++ b/tests/kernels/quantization/test_gptq.py @@ -11,9 +11,8 @@ def test_gptq_shuffle_opcheck(): weight = torch.randint( -2000000, 2000000, (1792, 4096), device="cuda", dtype=torch.int32 ) - perm = torch.empty((0,), device="cuda", dtype=torch.int32) bit = 4 - opcheck(torch.ops._C.gptq_shuffle, (weight, perm, bit)) + opcheck(torch.ops._C.gptq_shuffle, (weight, bit)) def test_gptq_gemm_opcheck(): @@ -23,13 +22,8 @@ def test_gptq_gemm_opcheck(): ) zeros = torch.zeros((32, 768), device="cuda", dtype=torch.int32) scales = torch.rand((32, 6144), device="cuda", dtype=torch.float16) - idx = torch.empty((0,), device="cuda", dtype=torch.int32) use_exllama = True bit = 4 # Test both GPTQv1 and GPTQv2 format - opcheck( - torch.ops._C.gptq_gemm, (a, weight, zeros, scales, idx, use_exllama, True, bit) - ) - opcheck( - torch.ops._C.gptq_gemm, (a, weight, zeros, scales, idx, use_exllama, False, bit) - ) + opcheck(torch.ops._C.gptq_gemm, (a, weight, zeros, scales, use_exllama, True, bit)) + opcheck(torch.ops._C.gptq_gemm, (a, weight, zeros, scales, use_exllama, False, bit)) diff --git a/tests/kernels/quantization/test_marlin_gemm.py b/tests/kernels/quantization/test_marlin_gemm.py index 8b35fab81e..4c290dada6 100644 --- a/tests/kernels/quantization/test_marlin_gemm.py +++ b/tests/kernels/quantization/test_marlin_gemm.py @@ -17,7 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import ( per_token_quant_int8, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - marlin_make_empty_g_idx, + marlin_make_empty, marlin_make_workspace_new, marlin_permute_bias, query_marlin_supported_quant_types, @@ -40,7 +40,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( gptq_pack, gptq_quantize_weights, quantize_weights, - sort_weights, ) from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -51,8 +50,6 @@ if current_platform.is_rocm(): allow_module_level=True, ) -ACT_ORDER_OPTS = [False, True] -K_FULL_OPTS = [False, True] USE_ATOMIC_ADD_OPTS = [False, True] USE_FP32_REDUCE_OPTS = [True] @@ -80,13 +77,11 @@ DENSE_MARLIN_QUANT_TEST_CONFIGS = [ # GPTQ-INT4 { "b_type": scalar_types.uint4b8, - "support_act_order": True, "group_blocks": [-1, 2, 4, 8], }, # GPTQ-INT8 { "b_type": scalar_types.uint8b128, - "support_act_order": True, "group_blocks": [-1, 2, 4, 8], }, # FP8 @@ -202,27 +197,15 @@ def test_marlin_int4_fp8_preprocess_awq(): @pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS) @pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS) @pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(False, False)) -@pytest.mark.parametrize("act_order", ACT_ORDER_OPTS) @pytest.mark.parametrize("is_a_8bit", [True, False]) @pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS) -def test_gptq_marlin_repack( - k_chunk, n_chunk, quant_type, act_order, is_a_8bit, nk_factors -): +def test_gptq_marlin_repack(k_chunk, n_chunk, quant_type, is_a_8bit, nk_factors): n_factor, k_factor = nk_factors size_k = k_chunk * k_factor size_n = n_chunk * n_factor group_size = 128 - # Filter act_order - if act_order: - if group_size == -1: - return - if group_size == size_k: - return - if is_a_8bit: - return - # Normalize group_size if group_size == -1: group_size = size_k @@ -231,20 +214,12 @@ def test_gptq_marlin_repack( # Create input b_weight = rand_data((size_k, size_n)) - # Quantize (and apply act_order if provided) - w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights( - b_weight, quant_type, group_size, act_order - ) + # Quantize + _, q_w, _ = gptq_quantize_weights(b_weight, quant_type, group_size) # Pack to GPTQ format q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n) - # For act_order, sort the "weights" and "g_idx" so that group ids are - # increasing - sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device) - if act_order: - q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) - # Pack to Marlin format weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit) marlin_q_w_1 = marlin_weights( @@ -253,12 +228,12 @@ def test_gptq_marlin_repack( opcheck( torch.ops._C.gptq_marlin_repack, - (q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit), + (q_w_gptq, size_k, size_n, quant_type.size_bits, is_a_8bit), ) # Run Marlin repack GPU kernel marlin_q_w_2 = ops.gptq_marlin_repack( - q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit + q_w_gptq, size_k, size_n, quant_type.size_bits, is_a_8bit ) torch.accelerator.synchronize() @@ -319,8 +294,6 @@ def marlin_generate_valid_test_cases(): MNK_FACTORS, MARLIN_N_CHUNKS, MARLIN_K_CHUNKS, - ACT_ORDER_OPTS, - K_FULL_OPTS, USE_ATOMIC_ADD_OPTS, USE_FP32_REDUCE_OPTS, ) @@ -333,8 +306,6 @@ def marlin_generate_valid_test_cases(): size_m, size_n, size_k, - act_order, - is_k_full, use_atomic_add, use_fp32_reduce, ): @@ -351,25 +322,18 @@ def marlin_generate_valid_test_cases(): if group_size > 0 and size_k % group_size != 0: return False - if act_order and group_size in [-1, size_k]: - return False if group_size == size_k: return False - if not act_order and is_k_full: - return False return a_type.size_bits < 16 or a_type is c_type cases = [] for case in all_combinations: - quant_test_config, mnk_factors, n_chunk, k_chunk, act_order, *_ = case + quant_test_config, mnk_factors, n_chunk, k_chunk, *_ = case size_m = mnk_factors[0] size_n = mnk_factors[1] * n_chunk size_k = mnk_factors[2] * k_chunk - if act_order and not quant_test_config.get("support_act_order", False): - continue - f16_types = [scalar_types.float16, scalar_types.bfloat16] inner_combinations = itertools.product( quant_test_config.get("a_type", f16_types), @@ -398,7 +362,7 @@ def marlin_generate_valid_test_cases(): @pytest.mark.parametrize( ( "a_type, b_type, c_type, group_blocks," - "size_m, size_n, size_k, act_order, is_k_full," + "size_m, size_n, size_k," "use_atomic_add, use_fp32_reduce" ), marlin_generate_valid_test_cases(), @@ -411,8 +375,6 @@ def test_marlin_gemm( size_m, size_n, size_k, - act_order, - is_k_full, use_atomic_add, use_fp32_reduce, ): @@ -448,27 +410,21 @@ def test_marlin_gemm( ) marlin_s2 = None - g_idx = None - sort_indices = None marlin_zp = None elif b_type == scalar_types.float8_e4m3fn: w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch( b_weight.T, group_size, input_dtype=a_dtype ) - g_idx = None - sort_indices = None marlin_zp = None marlin_s2 = None elif has_zp: w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize( b_weight, b_type, group_size, input_dtype=a_dtype ) - g_idx = None - sort_indices = None marlin_s2 = None else: - w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( - b_weight, b_type, group_size, act_order, input_dtype=a_dtype + w_ref, marlin_q_w, marlin_s = marlin_quantize( + b_weight, b_type, group_size, input_dtype=a_dtype ) marlin_zp = None @@ -506,14 +462,11 @@ def test_marlin_gemm( a_scales, marlin_s2, marlin_zp, - g_idx, - sort_indices, workspace, b_type, a_input.shape[0], b_weight.shape[1], a_input.shape[1], - is_k_full=is_k_full, use_atomic_add=use_atomic_add, use_fp32_reduce=use_fp32_reduce, is_zp_float=False, @@ -535,11 +488,9 @@ def test_marlin_gemm_subset_input(): a_input = rand_data((big_m, big_k))[8 : size_m + 8, 8 : size_k + 8] b_weight = rand_data((size_k, size_n)) - w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( - b_weight, quant_type, group_size, False - ) + w_ref, marlin_q_w, marlin_s = marlin_quantize(b_weight, quant_type, group_size) - marlin_zp = marlin_make_empty_g_idx(marlin_s.device) + marlin_zp = marlin_make_empty(marlin_s.device) workspace = marlin_make_workspace_new(a_input.device) output = ops.marlin_gemm( @@ -551,14 +502,11 @@ def test_marlin_gemm_subset_input(): None, None, marlin_zp, - g_idx, - sort_indices, workspace, quant_type, a_input.shape[0], b_weight.shape[1], a_input.shape[1], - is_k_full=True, use_atomic_add=False, use_fp32_reduce=True, is_zp_float=False, @@ -584,11 +532,9 @@ def test_marlin_gemm_with_bias(size_m): marlin_bias = marlin_permute_bias(b_bias) - w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( - b_weight, quant_type, group_size, False - ) + w_ref, marlin_q_w, marlin_s = marlin_quantize(b_weight, quant_type, group_size) - marlin_zp = marlin_make_empty_g_idx(marlin_s.device) + marlin_zp = marlin_make_empty(marlin_s.device) workspace = marlin_make_workspace_new(a_input.device) output = ops.marlin_gemm( @@ -600,14 +546,11 @@ def test_marlin_gemm_with_bias(size_m): None, None, marlin_zp, - g_idx, - sort_indices, workspace, quant_type, a_input.shape[0], b_weight.shape[1], a_input.shape[1], - is_k_full=True, use_atomic_add=False, use_fp32_reduce=True, is_zp_float=False, diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index 981046fa53..027ffbf10c 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -14,7 +14,7 @@ from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils.marlin_utils import ( GPTQ_MARLIN_TILE, apply_gptq_marlin_linear, - marlin_make_empty_g_idx, + marlin_make_empty, marlin_make_workspace_new, marlin_moe_padded_intermediate, marlin_pad_qweight, @@ -319,16 +319,13 @@ def test_gptq_marlin_padded_round_trip(shape, group_size): device = torch.device("cuda") weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5 - w_ref, q_w, s, _, _ = gptq_quantize_weights( - weight, quant_type, group_size, act_order=False - ) + w_ref, q_w, s = gptq_quantize_weights(weight, quant_type, group_size) qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n) padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size) qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k) marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, - perm=torch.empty(0, dtype=torch.int, device=device), size_k=padded_k, size_n=padded_n, num_bits=quant_type.size_bits, @@ -343,14 +340,11 @@ def test_gptq_marlin_padded_round_trip(shape, group_size): input=x, weight=marlin_qweight, weight_scale=marlin_s, - weight_zp=marlin_make_empty_g_idx(device), - g_idx=marlin_make_empty_g_idx(device), - g_idx_sort_indices=marlin_make_empty_g_idx(device), + weight_zp=marlin_make_empty(device), workspace=marlin_make_workspace_new(device), wtype=quant_type, output_size_per_partition=size_n, input_size_per_partition=size_k, - is_k_full=True, ) ref = x @ w_ref @@ -480,7 +474,6 @@ def test_awq_zp_marlin_padded_round_trip(shape): qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k) marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, - perm=torch.empty(0, dtype=torch.int, device=device), size_k=padded_k, size_n=padded_n, num_bits=quant_type.size_bits, @@ -503,13 +496,10 @@ def test_awq_zp_marlin_padded_round_trip(shape): weight=marlin_qweight, weight_scale=marlin_s, weight_zp=marlin_zp, - g_idx=marlin_make_empty_g_idx(device), - g_idx_sort_indices=marlin_make_empty_g_idx(device), workspace=marlin_make_workspace_new(device), wtype=quant_type, output_size_per_partition=size_n, input_size_per_partition=size_k, - is_k_full=True, ) ref = x @ w_ref @@ -585,9 +575,7 @@ def test_gptq_marlin_moe_padded_round_trip(shape, group_size): def quant(w, size_k, size_n): # w is (size_n, size_k); gptq expects (size_k, size_n). - ref, q_w, s, _, _ = gptq_quantize_weights( - w.T, quant_type, group_size, act_order=False - ) + ref, q_w, s = gptq_quantize_weights(w.T, quant_type, group_size) return ref, gptq_pack(q_w, bits, size_k, size_n), s w13_qw, w13_s, w13_ref = [], [], [] @@ -616,12 +604,11 @@ def test_gptq_marlin_moe_padded_round_trip(shape, group_size): if group_size > 0: w2_scales = _pad_rows(w2_scales, padded_n // group_size) - sort_idx = torch.empty((e, 0), dtype=torch.int32, device=device) marlin_w13 = ops.gptq_marlin_moe_repack( - w13_qweight, sort_idx, w13_qweight.shape[1] * pack, w13_qweight.shape[2], bits + w13_qweight, w13_qweight.shape[1] * pack, w13_qweight.shape[2], bits ) marlin_w2 = ops.gptq_marlin_moe_repack( - w2_qweight, sort_idx, w2_qweight.shape[1] * pack, w2_qweight.shape[2], bits + w2_qweight, w2_qweight.shape[1] * pack, w2_qweight.shape[2], bits ) group_or_pack = group_size if group_size != -1 else pack marlin_w13_s = marlin_moe_permute_scales( @@ -649,7 +636,6 @@ def test_gptq_marlin_moe_padded_round_trip(shape, group_size): topk_ids, quant_type_id=quant_type.id, global_num_experts=e, - is_k_full=True, ) with set_current_vllm_config(VllmConfig()): ref = torch_experts( @@ -770,7 +756,6 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant): topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - is_k_full=True, workspace=layer.workspace, ) with set_current_vllm_config(VllmConfig()): @@ -848,7 +833,6 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape): topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - is_k_full=True, workspace=layer.workspace, ) with set_current_vllm_config(VllmConfig()): diff --git a/tests/kernels/quantization/test_rdna3_compile_guards.py b/tests/kernels/quantization/test_rdna3_compile_guards.py index c307bfc3ae..3c1e226d2c 100644 --- a/tests/kernels/quantization/test_rdna3_compile_guards.py +++ b/tests/kernels/quantization/test_rdna3_compile_guards.py @@ -196,7 +196,6 @@ def test_dense_kernel_rejects_on_non_gfx1100(): act_type=torch.float16, group_size=128, zero_points=False, - has_g_idx=False, ) ok, reason = RDNA3W4A16LinearKernel.can_implement(config) assert ok is False, f"RDNA3 dense kernel accepted on non-gfx1100: {reason}" @@ -431,7 +430,6 @@ class TestDenseKernelSelectionMocked: act_type=torch.float16, group_size=128, zero_points=False, - has_g_idx=False, ) ok, _ = RDNA3W4A16LinearKernel.can_implement(config) assert ok is True @@ -461,7 +459,6 @@ class TestDenseKernelSelectionMocked: act_type=torch.float16, group_size=128, zero_points=False, - has_g_idx=False, ) with ( patch("vllm.platforms.rocm.on_gfx1100", return_value=False), diff --git a/tests/kernels/quantization/test_rdna3_moe_w4a16.py b/tests/kernels/quantization/test_rdna3_moe_w4a16.py index 4248251635..6212e6f08c 100644 --- a/tests/kernels/quantization/test_rdna3_moe_w4a16.py +++ b/tests/kernels/quantization/test_rdna3_moe_w4a16.py @@ -69,10 +69,9 @@ def _make_packed_weights(E, K, N): packed = torch.zeros(E, K // 8, N, dtype=torch.int32, device=device) for i in range(8): packed |= (w[:, i::8, :] & 0xF) << (i * 4) - g_idx = torch.empty(0, dtype=torch.int32, device=device) for e in range(E): we = packed[e].contiguous() - ops.gptq_shuffle(we, g_idx, 4) + ops.gptq_shuffle(we, 4) packed[e] = we return packed @@ -113,7 +112,6 @@ def test_fused_moe_w1_matches_dense( w13 = _make_packed_weights(E, K, N_gate_up) w13_s = _make_scales(E, groups, N_gate_up, dtype) w13_z = _make_qzeros(E, groups, N_gate_up) - g_idx = torch.empty(0, dtype=torch.int32, device=device) topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) si, ei, ntp = moe_align_block_size(topk_ids, block_size_m, E) @@ -147,7 +145,6 @@ def test_fused_moe_w1_matches_dense( w13[e], w13_z[e], w13_s[e], - g_idx, False, ) ref_out[flat] = ref.squeeze() @@ -243,7 +240,6 @@ def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype): w2 = _make_packed_weights(E, N_inter, hidden) w2_s = _make_scales(E, N_inter // group_size, hidden, dtype) w2_z = _make_qzeros(E, N_inter // group_size, hidden) - g_idx = torch.empty(0, dtype=torch.int32, device=device) topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) topk_w = torch.softmax( @@ -300,7 +296,6 @@ def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype): w13[e], w13_z[e], w13_s[e], - g_idx, False, ) a = torch.empty(1, N_inter, dtype=dtype, device=device) @@ -310,7 +305,6 @@ def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype): w2[e], w2_z[e], w2_s[e], - g_idx, False, ) ref[m_idx] += r2.squeeze() * w diff --git a/tests/kernels/quantization/test_rdna3_w4a16.py b/tests/kernels/quantization/test_rdna3_w4a16.py index b70a2b9a86..5cecd03944 100644 --- a/tests/kernels/quantization/test_rdna3_w4a16.py +++ b/tests/kernels/quantization/test_rdna3_w4a16.py @@ -173,7 +173,6 @@ def _run_kernel( act_type=dtype, group_size=group_size, zero_points=has_zp, - has_g_idx=False, ) ok, reason = RDNA3W4A16LinearKernel.can_implement(config) assert ok, f"can_implement rejected a supported config: {reason}" @@ -184,7 +183,6 @@ def _run_kernel( w_q_param_name="qweight", w_s_param_name="scales", w_zp_param_name="qzeros" if has_zp else None, - w_gidx_param_name=None, ) kernel.process_weights_after_loading(layer) return kernel.apply_weights(layer, x_mk, bias=bias) diff --git a/tests/kernels/quantization/test_rdna3_w4a16_selection.py b/tests/kernels/quantization/test_rdna3_w4a16_selection.py index b53e166378..49538a25c7 100644 --- a/tests/kernels/quantization/test_rdna3_w4a16_selection.py +++ b/tests/kernels/quantization/test_rdna3_w4a16_selection.py @@ -57,7 +57,6 @@ def test_selection_prefers_rdna3(dtype): act_type=dtype, group_size=128, zero_points=False, - has_g_idx=False, ) assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel" @@ -83,7 +82,6 @@ def test_can_implement(weight_type, group_size, N, full_k, expected_ok): act_type=torch.float16, group_size=group_size, zero_points=False, - has_g_idx=False, ) ok, reason = RDNA3W4A16LinearKernel.can_implement(config) assert ok is expected_ok, reason diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index 0d007aa117..562d1ef2ef 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -306,14 +306,12 @@ def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_ini act_type=torch.float16, group_size=G, zero_points=False, - has_g_idx=False, ) kernel = RDNAHybridW4A16LinearKernel( config, w_q_param_name="weight_packed", w_s_param_name="weight_scale", w_zp_param_name=None, - w_gidx_param_name=None, ) kernel.process_weights_after_loading(layer) @@ -367,14 +365,12 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in act_type=torch.float16, group_size=G, zero_points=True, - has_g_idx=False, ) kernel = RDNAHybridW4A16LinearKernel( config, w_q_param_name="weight_packed", w_s_param_name="weight_scale", w_zp_param_name="weight_zero_point", - w_gidx_param_name=None, ) kernel.process_weights_after_loading(layer) @@ -414,7 +410,6 @@ def test_hybrid_can_implement_group_size(group_size, expected_ok): act_type=torch.float16, group_size=group_size, zero_points=False, - has_g_idx=False, ) ok, _ = RDNAHybridW4A16LinearKernel.can_implement(config) assert ok is expected_ok diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py index dba7f4bd08..bd8bd3cae2 100644 --- a/tests/kernels/quantization/test_triton_w4a16.py +++ b/tests/kernels/quantization/test_triton_w4a16.py @@ -273,14 +273,12 @@ def test_triton_w4a16_process_weights_after_loading_repacks_layout(): act_type=torch.float16, group_size=G, zero_points=True, - has_g_idx=False, ) kernel = TritonW4A16LinearKernel( config, w_q_param_name="weight_packed", w_s_param_name="weight_scale", w_zp_param_name="weight_zero_point", - w_gidx_param_name=None, ) # Build dummy layer with vLLM parameter wrappers. @@ -384,14 +382,12 @@ def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout(): act_type=torch.float16, group_size=G, zero_points=True, - has_g_idx=False, ) kernel = TritonW4A16LinearKernel( config, w_q_param_name="qweight", w_s_param_name="scales", w_zp_param_name="qzeros", - w_gidx_param_name=None, ) weight_loader = lambda *args, **kwargs: None @@ -461,14 +457,12 @@ def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch): act_type=torch.float16, group_size=G, zero_points=False, - has_g_idx=False, ) kernel = TritonW4A16LinearKernel( config, w_q_param_name="qweight", w_s_param_name="scales", w_zp_param_name="qzeros", - w_gidx_param_name=None, ) class DummyLayer(torch.nn.Module): diff --git a/tests/kernels/quantization/test_w4a16_kernel_selection.py b/tests/kernels/quantization/test_w4a16_kernel_selection.py index c9935746e1..55fde4ce63 100644 --- a/tests/kernels/quantization/test_w4a16_kernel_selection.py +++ b/tests/kernels/quantization/test_w4a16_kernel_selection.py @@ -47,7 +47,6 @@ def test_choose_mp_linear_kernel_uint4b8(): act_type=torch.float16, group_size=128, zero_points=False, - has_g_idx=False, ) kernel_type = choose_mp_linear_kernel(config) @@ -65,7 +64,6 @@ def test_choose_mp_linear_kernel_uint4_asymmetric(): act_type=torch.bfloat16, group_size=64, zero_points=True, - has_g_idx=False, ) kernel_type = choose_mp_linear_kernel(config) diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index f14b488a17..2c64fb2245 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -270,11 +270,6 @@ def qwen3_woofing_lora_files(): return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") -@pytest.fixture(scope="session") -def tinyllama_lora_files(): - return hf_api().snapshot_download(repo_id="jashing/tinyllama-colorist-lora") - - @pytest.fixture(scope="session") def deepseekv2_lora_files(): return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA") diff --git a/tests/lora/test_quant_model.py b/tests/lora/test_quant_model.py index 06e1b22ab5..38809385c1 100644 --- a/tests/lora/test_quant_model.py +++ b/tests/lora/test_quant_model.py @@ -24,42 +24,51 @@ MODELS: list[ModelWithQuantization] if current_platform.is_rocm(): MODELS = [ ModelWithQuantization( - model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", quantization="gptq" + model_path="JunHowie/Qwen3-0.6B-GPTQ-Int4", + quantization="gptq", ), ] else: MODELS = [ ModelWithQuantization( - model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", quantization="awq" + model_path="Orion-zhen/Qwen3-0.6B-AWQ", quantization="awq" ), ModelWithQuantization( - model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", quantization="gptq" + model_path="JunHowie/Qwen3-0.6B-GPTQ-Int4", + quantization="gptq", ), ] def do_sample( - llm: vllm.LLM, lora_path: str, lora_id: int, max_tokens: int = 256 + llm: vllm.LLM, lora_path: str | None, lora_id: int, max_tokens: int = 256 ) -> list[str]: - raw_prompts = [ - "Give me an orange-ish brown color", - "Give me a neon pink color", + messages = [ + [ + { + "role": "system", + "content": "Follow the instructions to make animal noises", + }, + {"role": "user", "content": "Make your favorite animal noise."}, + ], + [ + { + "role": "system", + "content": "You are a cat. Reply only with your sound.", + }, + {"role": "user", "content": "What do you say?"}, + ], ] - - def format_prompt_tuples(prompt): - return f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n" - - prompts = [format_prompt_tuples(p) for p in raw_prompts] - - sampling_params = vllm.SamplingParams( - temperature=0, max_tokens=max_tokens, stop=["<|im_end|>"] - ) - outputs = llm.generate( - prompts, + sampling_params = vllm.SamplingParams(temperature=0, max_tokens=max_tokens) + outputs = llm.chat( + messages, sampling_params, - lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None, + chat_template_kwargs={"enable_thinking": False}, + lora_request=( + LoRARequest(str(lora_id), lora_id, lora_path) if lora_path else None + ), + use_tqdm=False, ) - # Print the outputs. generated_texts: list[str] = [] for output in outputs: prompt = output.prompt @@ -70,7 +79,7 @@ def do_sample( @pytest.mark.parametrize("model", MODELS) -def test_quant_model_lora(tinyllama_lora_files, model): +def test_quant_model_lora(qwen3_meowing_lora_files, qwen3_woofing_lora_files, model): llm = vllm.LLM( model=model.model_path, enable_lora=True, @@ -79,75 +88,59 @@ def test_quant_model_lora(tinyllama_lora_files, model): max_model_len=400, gpu_memory_utilization=0.2, # avoid OOM quantization=model.quantization, - trust_remote_code=True, enable_chunked_prefill=True, - tokenizer=tinyllama_lora_files, ) - if model.quantization is None: - expected_lora_output = [ - "#ff8050", - "#ff8080", - ] - elif model.quantization == "awq": - expected_lora_output = [ - "#f07700: A v", - "#f00000: A v", - ] - elif model.quantization == "gptq": - expected_lora_output = [ - "#f08800: This is", - "#f07788 \n#", - ] - - def expect_match(output, expected_output): - # HACK: GPTQ lora outputs are just incredibly unstable. - # Assert that the outputs changed. - if model.quantization == "gptq" and expected_output is expected_lora_output: - for i, o in enumerate(output): - assert o.startswith("#"), ( - f"Expected example {i} to start with # but got {o}" - ) - return - assert output == expected_output - + loras = [ + ( + qwen3_meowing_lora_files, + 1, + ["Meow Meow Meow Meow Meow"] * 2, + ), + ( + qwen3_woofing_lora_files, + 2, + ["Woof Woof Woof Woof Woof"] * 2, + ), + ] max_tokens = 10 - - print("lora adapter created") - print("lora 1") - output = do_sample(llm, tinyllama_lora_files, lora_id=1, max_tokens=max_tokens) - expect_match(output, expected_lora_output) - - print("lora 2") - output = do_sample(llm, tinyllama_lora_files, lora_id=2, max_tokens=max_tokens) - expect_match(output, expected_lora_output) - - print("removing lora") - - del llm - cleanup_dist_env_and_memory() + try: + base_output = do_sample(llm, None, lora_id=0, max_tokens=max_tokens) + for lora_path, lora_id, expected_output in loras: + output = do_sample( + llm, + lora_path, + lora_id=lora_id, + max_tokens=max_tokens, + ) + assert output == expected_output + assert output != base_output + finally: + del llm + cleanup_dist_env_and_memory() @pytest.mark.parametrize("model", MODELS) -def test_quant_model_tp_equality(tinyllama_lora_files, num_gpus_available, model): +def test_quant_model_tp_equality(qwen3_meowing_lora_files, num_gpus_available, model): if num_gpus_available < 2: pytest.skip(f"Not enough GPUs for tensor parallelism {2}") - if model.quantization == "gptq": - pytest.skip("GPTQ lora outputs are just incredibly unstable") llm_tp1 = vllm.LLM( model=model.model_path, enable_lora=True, max_num_seqs=16, max_loras=4, + max_model_len=400, gpu_memory_utilization=0.2, # avoid OOM quantization=model.quantization, - trust_remote_code=True, enable_chunked_prefill=True, ) - output_tp1 = do_sample(llm_tp1, tinyllama_lora_files, lora_id=1) - - del llm_tp1 - cleanup_dist_env_and_memory() + try: + output_tp1 = do_sample( + llm_tp1, qwen3_meowing_lora_files, lora_id=1, max_tokens=10 + ) + finally: + del llm_tp1 + cleanup_dist_env_and_memory() llm_tp2 = vllm.LLM( model=model.model_path, @@ -155,13 +148,17 @@ def test_quant_model_tp_equality(tinyllama_lora_files, num_gpus_available, model max_num_seqs=16, max_loras=4, tensor_parallel_size=2, + max_model_len=400, gpu_memory_utilization=0.2, # avoid OOM quantization=model.quantization, enable_chunked_prefill=True, ) - output_tp2 = do_sample(llm_tp2, tinyllama_lora_files, lora_id=1) - - del llm_tp2 - cleanup_dist_env_and_memory() + try: + output_tp2 = do_sample( + llm_tp2, qwen3_meowing_lora_files, lora_id=1, max_tokens=10 + ) + finally: + del llm_tp2 + cleanup_dist_env_and_memory() assert output_tp1 == output_tp2 diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index b17f2ece3c..f19ec35560 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -272,13 +272,13 @@ def _stub_marlin_ops(monkeypatch): monkeypatch.setattr( ops, "gptq_marlin_repack", - lambda w, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( + lambda w, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( size_k // 16, size_n * 2, dtype=torch.int32 ), ) -def _make_act_order_marlin_kernel(): +def _make_marlin_kernel(): from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( MarlinLinearKernel, ) @@ -295,20 +295,17 @@ def _make_act_order_marlin_kernel(): act_type=torch.float16, group_size=_MARLIN_GROUP_SIZE, zero_points=False, - has_g_idx=True, ) kernel.w_q_name = "qweight" kernel.w_s_name = "scales" kernel.w_zp_name = None - kernel.w_gidx_name = "g_idx" return kernel -def _load_marlin_checkpoint_format_weights(layer, g_idx): +def _load_marlin_checkpoint_format_weights(layer): from vllm.model_executor.parameter import ( GroupQuantScaleParameter, PackedvLLMParameter, - RowvLLMParameter, ) layer.qweight = PackedvLLMParameter( @@ -327,54 +324,24 @@ def _load_marlin_checkpoint_format_weights(layer, g_idx): output_dim=1, weight_loader=default_weight_loader, ) - layer.g_idx = RowvLLMParameter( - data=g_idx.clone(), - input_dim=0, - weight_loader=default_weight_loader, - ) - - -def _random_g_idx(generator): - return torch.randint( - 0, - _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, - (_MARLIN_SIZE_K,), - dtype=torch.int32, - generator=generator, - ) def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init): - """Marlin workspace and act-order sort indices must be recomputed into - the same storage when weights are reloaded (RL weight sync), so device - addresses captured by CUDA graphs remain valid.""" - from vllm.model_executor.layers.quantization.utils import marlin_utils - + """Marlin must reuse the workspace storage captured by CUDA graphs.""" _stub_marlin_ops(monkeypatch) - kernel = _make_act_order_marlin_kernel() - - generator = torch.Generator().manual_seed(0) - first_g_idx = _random_g_idx(generator) - second_g_idx = _random_g_idx(generator) + kernel = _make_marlin_kernel() layer = torch.nn.Module() - _load_marlin_checkpoint_format_weights(layer, first_g_idx) + _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) workspace_ptr = kernel.workspace.data_ptr() - sort_indices_ptr = layer.g_idx_sort_indices.data_ptr() - # Reload: fresh checkpoint-format tensors with a different act-order - _load_marlin_checkpoint_format_weights(layer, second_g_idx) + _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) assert kernel.workspace.data_ptr() == workspace_ptr assert torch.all(kernel.workspace == 0) - assert layer.g_idx_sort_indices.data_ptr() == sort_indices_ptr - expected_sort_indices = marlin_utils.marlin_sort_g_idx(second_g_idx)[1] - assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices) - # registered as a Parameter so layerwise reload copy-back preserves it - assert isinstance(layer.g_idx_sort_indices, torch.nn.Parameter) @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) @@ -394,7 +361,7 @@ def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): monkeypatch.setattr( ops, "gptq_marlin_repack", - lambda b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( + lambda b_q_weight, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( size_k // 16, size_n * 2, dtype=torch.int32 ), ) @@ -476,68 +443,6 @@ def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): ) -def test_marlin_act_order_layerwise_reload_accounting(monkeypatch, dist_init): - """`g_idx_sort_indices` is generated during weight processing and never - loaded from checkpoints. Registering it as a Parameter must not count it - toward `load_numel_total`: reload restores the construction-time tensor - set before sizing, so act-order layers still process during streaming - instead of deferring (and buffering weights) until finalization.""" - from vllm.model_executor.layers.quantization.base_config import ( - QuantizeMethodBase, - ) - from vllm.model_executor.layers.quantization.utils import marlin_utils - from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info - - _stub_marlin_ops(monkeypatch) - kernel = _make_act_order_marlin_kernel() - - class _KernelQuantMethod(QuantizeMethodBase): - def create_weights(self, layer, *args, **kwargs): - raise NotImplementedError - - def apply(self, layer, *args, **kwargs): - raise NotImplementedError - - def process_weights_after_loading(self, layer): - kernel.process_weights_after_loading(layer) - - generator = torch.Generator().manual_seed(0) - layer = torch.nn.Module() - layer.quant_method = _KernelQuantMethod() - _load_marlin_checkpoint_format_weights(layer, _random_g_idx(generator)) - - # Metadata is recorded at model construction, before any processing - record_metadata_for_reloading(layer) - checkpoint_numel = sum(t.numel() for t in get_layer_tensors(layer).values()) - - kernel.process_weights_after_loading(layer) - sort_indices = layer.g_idx_sort_indices - - initialize_layerwise_reload(layer) - info = get_layerwise_info(layer) - assert info.load_numel_total == checkpoint_numel - - # Stream a new checkpoint; the layer must process as soon as its last - # tensor arrives - new_g_idx = _random_g_idx(generator) - checkpoint = { - "qweight": torch.zeros(_MARLIN_SIZE_K // 8, _MARLIN_SIZE_N, dtype=torch.int32), - "scales": torch.ones( - _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, _MARLIN_SIZE_N, dtype=torch.float16 - ), - "g_idx": new_g_idx, - } - for name, weight in checkpoint.items(): - param = getattr(layer, name) - param.weight_loader(param, weight) - - assert not info.can_load() - assert not info.loaded_weights - assert layer.g_idx_sort_indices is sort_indices - expected_sort_indices = marlin_utils.marlin_sort_g_idx(new_g_idx)[1] - assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices) - - def test_model_cleanup(dist_init, default_vllm_config): layer = QKVParallelLinear(2, 3, 4) assert layer.weight.weight_loader.__self__ is layer diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 9e67609b78..e577690c92 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -160,6 +160,17 @@ class TestMaybeRemapKvScaleName: assert result is None +def test_checkpoint_weight_mapper_discards_serialized_g_idx(): + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + + mapper = QuantizationConfig.get_checkpoint_weight_mapper() + + assert mapper._map_name("model.layers.0.g_idx") is None + assert mapper._map_name("model.layers.0.qweight") == "model.layers.0.qweight" + + class TestKvCacheScaleMapper: """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of diff --git a/tests/models/quantization/test_gptq_marlin.py b/tests/models/quantization/test_gptq_marlin.py index de97ed9562..3d34ccb257 100644 --- a/tests/models/quantization/test_gptq_marlin.py +++ b/tests/models/quantization/test_gptq_marlin.py @@ -22,12 +22,10 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true" MAX_MODEL_LEN = 1024 MODELS = [ - # act_order==True, group_size=128 - ("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", "main"), - # 8-bit, act_order==True, group_size=channelwise + ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", "main"), + # desc_act=True is a no-op for channelwise quantization. ("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", "gptq-8bit--1g-actorder_True"), - # 4-bit, act_order==True, group_size=128 - ("TechxGenus/gemma-1.1-2b-it-GPTQ", "main"), + ("TechxGenus/gemma-2b-it-GPTQ", "main"), ] diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index 3014bf05c8..1eaf21b4cf 100644 --- a/tests/models/transformers/test_backend.py +++ b/tests/models/transformers/test_backend.py @@ -243,7 +243,7 @@ def test_distributed( "model, quantization_kwargs", [ ("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {}), - ("TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", {}), + ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", {}), ], ) @pytest.mark.parametrize("max_tokens", [32]) diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index e91093737e..e586155225 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -23,12 +23,13 @@ from vllm.model_executor.layers.quantization.auto_gptq import ( AutoGPTQLinearMethod, AutoGPTQMoEMethod, ) +from vllm.model_executor.layers.quantization.utils.gptq_utils import override_config from vllm.platforms import current_platform PROMPT = "On the surface of Mars, we found" MODELS = [ - "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", + "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", ] @@ -70,6 +71,54 @@ def test_auto_gptq_config_get_name(): assert AutoGPTQConfig.get_name() == "auto_gptq" +@pytest.mark.parametrize( + "dynamic", + [ + {}, + {r"+:model\.layers\.0\..*": {"desc_act": True}}, + { + r"+:model\.layers\.0\..*": { + "desc_act": True, + "group_size": 32, + } + }, + ], +) +def test_auto_gptq_rejects_group_activation_order(dynamic): + desc_act = not dynamic + with pytest.raises(ValueError, match="group activation ordering"): + AutoGPTQConfig(4, 128, desc_act, True, False, dynamic, {}) + + +def test_auto_gptq_normalizes_channelwise_activation_order(): + config = AutoGPTQConfig(4, -1, True, True, False, {}, {}) + assert not config.desc_act + + config = AutoGPTQConfig( + 4, + 128, + False, + True, + False, + {r"+:model\.layers\.0\..*": {"desc_act": True, "group_size": -1}}, + {}, + ) + override_config(config, "model.layers.0.self_attn.q_proj") + assert config.group_size == -1 + assert not config.desc_act + + with pytest.raises(ValueError, match="group activation ordering"): + AutoGPTQConfig( + 4, + -1, + True, + True, + False, + {r"+:model\.layers\.0\..*": {"group_size": 128}}, + {}, + ) + + def test_auto_gptq_moe_creates_zero_initialized_expert_biases(): method = object.__new__(AutoGPTQMoEMethod) method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {}) @@ -84,7 +133,6 @@ def test_auto_gptq_moe_creates_zero_initialized_expert_biases(): hidden_size=8, intermediate_size_per_partition=4, params_dtype=torch.float16, - intermediate_size_full=4, weight_loader=lambda *args, **kwargs: None, ) diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index bfc532126a..2bc463b593 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -1537,8 +1537,7 @@ def test_inc_ark_linear_method_xpu_int2_create_weights(monkeypatch) -> None: assert layer.scales.dtype == torch.bfloat16 assert layer.qzeros.shape == (1, 4) assert layer.qzeros.dtype == torch.int32 - assert layer.g_idx.shape == (64,) - assert layer.g_idx.dtype == torch.int32 + assert not hasattr(layer, "g_idx") assert layer.in_features == 64 assert layer.out_features == 64 assert layer.params_dtype == torch.bfloat16 @@ -2444,7 +2443,7 @@ def test_calls_kernel_at_threshold(monkeypatch, w4a8_layer) -> None: w_scale, w_zp, group_size, - g_idx, + group_idx, bias, ) = captured["gemm_args"] assert quant_x.dtype is torch.int8 @@ -2456,7 +2455,7 @@ def test_calls_kernel_at_threshold(monkeypatch, w4a8_layer) -> None: assert qweight is layer.qweight assert w_zp is layer.qzeros assert group_size == 128 - assert g_idx is None + assert group_idx is None assert bias is None # The kernel emits fp16; the result is cast back to the activation dtype. diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 004983cb19..ef80287e5f 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -84,6 +84,48 @@ def enable_pickle(monkeypatch): monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") +def _compressed_tensors_actorder_config(actorder): + return { + "format": "pack-quantized", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": { + "num_bits": 4, + "type": "int", + "strategy": "group", + "group_size": 128, + "symmetric": True, + "dynamic": False, + "actorder": actorder, + }, + } + }, + } + + +@pytest.mark.parametrize( + "actorder", [ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC, True] +) +def test_compressed_tensors_rejects_runtime_activation_order(actorder): + with pytest.raises(ValueError, match="group/dynamic activation ordering"): + CompressedTensorsConfig.from_config( + _compressed_tensors_actorder_config(actorder) + ) + + +@pytest.mark.parametrize( + "actorder", [ActivationOrdering.STATIC, ActivationOrdering.WEIGHT, False, None] +) +def test_compressed_tensors_accepts_preordered_weights(actorder): + config = CompressedTensorsConfig.from_config( + _compressed_tensors_actorder_config(actorder) + ) + assert config.target_scheme_map["Linear"]["weights"].actorder == ( + actorder.value if isinstance(actorder, ActivationOrdering) else None + ) + + @pytest.mark.parametrize( "model_args", [ @@ -271,15 +313,6 @@ def test_compressed_tensors_w8a8_dynamic_per_token( None, 8, True, - False, - ), - ( - "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder", - "group", - 128, - 8, - False, - True, ), ], ) @@ -287,7 +320,7 @@ def test_compressed_tensors_w8a8_dynamic_per_token( not current_platform.is_cuda(), reason="The tests are skipped on non-CUDA platform." ) def test_compressed_tensors_wNa16(vllm_runner, wNa16_args): - model, strategy, group, pack_factor, symmetric, has_g_idx = wNa16_args + model, strategy, group, pack_factor, symmetric = wNa16_args with vllm_runner(model, enforce_eager=True) as llm: def check_model(model): @@ -302,7 +335,6 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args): assert qkv_proj.scheme.pack_factor == pack_factor assert qkv_proj.scheme.symmetric == symmetric - assert qkv_proj.scheme.has_g_idx == has_g_idx llm.apply_model(check_model) @@ -926,35 +958,6 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner): assert output -@pytest.mark.parametrize( - "actorder,group_size,part,full,expected", - [ - # actorder="group" with real grouping: must load full-K w2 scales and, - # when sharded (part != full), report is_k_full=False. - (ActivationOrdering.GROUP, 32, 64, 128, (True, 128, False)), - # actorder="group" but unsharded (part == full): full scales, k_full. - (ActivationOrdering.GROUP, 32, 128, 128, (True, 128, True)), - # actorder="group" with channel-wise (group_size == -1): no full load. - (ActivationOrdering.GROUP, -1, 64, 128, (False, 64, False)), - # "static"/"weight" reorder at quant time -> shard normally + k_full. - # Regression: static actorder under TP must keep is_k_full=True so the - # Marlin kernel never gets the invalid (group_size=16, is_k_full=0). - ("static", 32, 64, 128, (False, 64, True)), - ("weight", 32, 64, 128, (False, 64, True)), - (None, 32, 64, 128, (False, 64, True)), - ], -) -def test_wna16_moe_w2_scale_sharding(actorder, group_size, part, full, expected): - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 - CompressedTensorsWNA16MoEMethod, - ) - - result = CompressedTensorsWNA16MoEMethod._w2_scale_sharding( - actorder, group_size, part, full - ) - assert result == expected - - @pytest.mark.parametrize("num_bits", range(2, 9)) def test_humming_supports_compressed_tensors_wna16_quant_key(num_bits): from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( diff --git a/tests/quantization/test_cpu_wna16.py b/tests/quantization/test_cpu_wna16.py index db8783c921..ffdc93f454 100644 --- a/tests/quantization/test_cpu_wna16.py +++ b/tests/quantization/test_cpu_wna16.py @@ -9,8 +9,7 @@ if not current_platform.is_cpu(): MODELS = [ "TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ", - "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx - "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx + "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", "RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", "Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear diff --git a/tests/quantization/test_cutlass_w4a16.py b/tests/quantization/test_cutlass_w4a16.py index b0ccb6c68a..8bf0f0a408 100644 --- a/tests/quantization/test_cutlass_w4a16.py +++ b/tests/quantization/test_cutlass_w4a16.py @@ -64,7 +64,6 @@ def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points) weight_type=weight_type, group_size=group_size, zero_points=zero_points, - has_g_idx=False, ) kernel = choose_mp_linear_kernel(config) assert kernel is MacheteLinearKernel, ( @@ -73,23 +72,21 @@ def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points) @pytest.mark.parametrize( - "full_shape,part_shape,weight_type,group_size,has_g_idx,expected_reason", + "full_shape,part_shape,weight_type,group_size,expected_reason", [ - ((4096, 4096), (2048, 4096), scalar_types.uint4b8, 128, True, "Act reordering"), ( (4096, 4096), (4096, 4096), scalar_types.float6_e3m2f, 128, - False, "Quant type", ), - ((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, False, "Group size"), + ((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, "Group size"), ], - ids=["partitioned-g_idx", "unsupported-quant-type", "unsupported-group-size"], + ids=["unsupported-quant-type", "unsupported-group-size"], ) def test_machete_rejects_invalid_config( - full_shape, part_shape, weight_type, group_size, has_g_idx, expected_reason + full_shape, part_shape, weight_type, group_size, expected_reason ): """Verify Machete rejects unsupported configurations.""" config = MPLinearLayerConfig( @@ -99,7 +96,6 @@ def test_machete_rejects_invalid_config( weight_type=weight_type, group_size=group_size, zero_points=False, - has_g_idx=has_g_idx, ) can_impl, reason = MacheteLinearKernel.can_implement(config) assert not can_impl @@ -117,7 +113,6 @@ def test_kernel_selection_with_disabled_machete(monkeypatch): weight_type=scalar_types.uint4b8, group_size=128, zero_points=False, - has_g_idx=False, ) kernel = choose_mp_linear_kernel(config) assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled" @@ -127,7 +122,6 @@ def test_kernel_selection_with_disabled_machete(monkeypatch): "model_name", [ "nm-testing/tinyllama-oneshot-w4a16-channel-v2", - "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder", ], ) def test_w4a16_machete_e2e(vllm_runner, model_name): diff --git a/tests/quantization/test_gptq_dynamic.py b/tests/quantization/test_gptq_dynamic.py index d4da7c4fd9..7a42b769ec 100644 --- a/tests/quantization/test_gptq_dynamic.py +++ b/tests/quantization/test_gptq_dynamic.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests whether gptq models with dynamic quantized can be loaded. +"""Tests whether unsupported GPTQ group actorder models are rejected. Run `pytest tests/quantization/test_gptq_dynamic.py --forked`. @@ -8,59 +8,16 @@ Note: Only symmetric GPTQ models are supported after consolidation to Marlin. """ import pytest -import torch -from tests.quantization.utils import load_model_without_vllm_runner -from vllm.model_executor.layers.linear import UnquantizedLinearMethod -from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod -from vllm.model_executor.layers.quantization.utils.gptq_utils import ( - get_dynamic_override, -) - -PROMPT = "On the surface of Mars, we found" - -# The first layer is quantized using bits=4, group_size=128 -# The second layer is quantized using bits=8, group_size=32 -# All other layers (layer index >= 2) are not quantized -# Note: Only symmetric models are supported with Marlin kernels MODELS = [ "ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head-symTrue", ] @pytest.mark.parametrize("model_id", MODELS) -def test_gptq_with_dynamic(model_id: str, monkeypatch, dist_init, workspace_init): - # `LLM.apply_model` requires pickling a function. - monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - - linear_method_cls = AutoGPTQLinearMethod - - model, _ = load_model_without_vllm_runner( - model_id, - dtype=torch.float16, - model_config_kwargs={ - "max_model_len": 2048, - "hf_overrides": {"num_hidden_layers": 3}, - }, - ) - - for name, submodule in model.named_modules(): - if name == "lm_head": - assert isinstance(submodule.quant_method, linear_method_cls) - elif name == "model.layers.0.self_attn.qkv_proj": - assert isinstance(submodule.quant_method, linear_method_cls) - config = submodule.quant_method.quant_config - assert config.weight_bits == 4 - assert config.group_size == 128 - assert config.desc_act - elif name == "model.layers.1.self_attn.qkv_proj": - assert isinstance(submodule.quant_method, linear_method_cls) - config = submodule.quant_method.quant_config - assert get_dynamic_override(config, layer_name=name, key="bits") == 8 - assert get_dynamic_override(config, layer_name=name, key="group_size") == 32 - assert not get_dynamic_override(config, layer_name=name, key="desc_act") - elif name in ( - "model.layers.2.self_attn.qkv_proj", - "model.layers.2.mlp.gate_up_proj", - ): - assert isinstance(submodule.quant_method, UnquantizedLinearMethod) +def test_gptq_with_dynamic_desc_act_rejected(vllm_runner, model_id: str): + with ( + pytest.raises(ValueError, match="desc_act=True"), + vllm_runner(model_id, max_model_len=2048, enforce_eager=True), + ): + pass diff --git a/tests/quantization/test_lm_head.py b/tests/quantization/test_lm_head.py index a8fa8fd0d2..16f8fd8d38 100644 --- a/tests/quantization/test_lm_head.py +++ b/tests/quantization/test_lm_head.py @@ -17,8 +17,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( PROMPT = "On the surface of Mars, we found" MODELS_QUANT = [ - ("ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head", True), - ("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", False), + ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", False), ] diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index e5be841f30..475ca84f6a 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -6,7 +6,6 @@ from types import SimpleNamespace import pytest import torch from compressed_tensors.quantization import ( - ActivationOrdering, QuantizationArgs, QuantizationStrategy, QuantizationType, @@ -33,6 +32,35 @@ def test_map_wna16_backend_supports_triton(): assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON +@pytest.mark.parametrize( + "config", + [ + {"desc_act": True, "group_size": 128}, + { + "desc_act": False, + "group_size": 128, + "dynamic": {r"+:model\.layers\.0\..*": {"desc_act": True}}, + }, + ], +) +def test_moe_wna16_rejects_gptq_group_activation_order(config): + config.update({"quant_method": "gptq", "bits": 4, "sym": True}) + with pytest.raises(ValueError, match="group activation ordering"): + MoeWNA16Config.from_config(config) + + +def test_moe_wna16_accepts_channelwise_gptq_activation_order(): + config = { + "quant_method": "gptq", + "bits": 4, + "group_size": -1, + "desc_act": True, + "sym": True, + } + assert MoeWNA16Config.is_moe_wna16_compatible(config) + MoeWNA16Config.from_config(config) + + @pytest.mark.parametrize( ("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"), [ @@ -43,28 +71,6 @@ def test_map_wna16_backend_supports_triton(): False, "AutoAWQ weight layout", ), - ( - WNA16MoEBackend.TRITON, - AutoGPTQConfig(4, 128, True, True, False, {}, {}), - False, - False, - "activation ordering", - ), - ( - WNA16MoEBackend.TRITON, - QuantizationArgs( - num_bits=4, - type=QuantizationType.INT, - strategy=QuantizationStrategy.GROUP, - symmetric=True, - dynamic=False, - group_size=128, - actorder=ActivationOrdering.GROUP, - ), - False, - False, - "activation ordering", - ), ( WNA16MoEBackend.TRITON, AutoGPTQConfig(4, 128, False, True, False, {}, {}), @@ -244,7 +250,6 @@ def test_compressed_tensors_wna16_moe_create_weights_uses_ceil_packed_shapes( hidden_size=256, intermediate_size_per_partition=512, params_dtype=torch.float16, - intermediate_size_full=512, ) packed_hidden = (256 * num_bits + 31) // 32 @@ -296,7 +301,6 @@ def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel(): num_experts=2, hidden_size=256, intermediate_size_per_partition=512, - intermediate_size_full=512, params_dtype=torch.bfloat16, ) layer.cuda() diff --git a/tests/weight_loading/models-large.txt b/tests/weight_loading/models-large.txt index e701542569..8b3913cc2d 100644 --- a/tests/weight_loading/models-large.txt +++ b/tests/weight_loading/models-large.txt @@ -1,8 +1,6 @@ compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-quantized, main compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-channel-quantized, main compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W8A16-quantized, main -compressed-tensors, nm-testing/test-w4a16-mixtral-actorder-group, main gptq, TheBloke/Mixtral-8x7B-v0.1-GPTQ, main -gptq, TheBloke/Mixtral-8x7B-v0.1-GPTQ, gptq-8bit-128g-actorder_True awq_marlin, casperhansen/deepseek-coder-v2-instruct-awq, main -compressed-tensors, RedHatAI/Llama-4-Scout-17B-16E-Instruct-quantized.w4a16, main \ No newline at end of file +compressed-tensors, RedHatAI/Llama-4-Scout-17B-16E-Instruct-quantized.w4a16, main diff --git a/tests/weight_loading/models.txt b/tests/weight_loading/models.txt index dd197a33b5..87b9650870 100644 --- a/tests/weight_loading/models.txt +++ b/tests/weight_loading/models.txt @@ -1,9 +1,8 @@ gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main gptq, TheBloke/Llama-2-7B-GPTQ, main -gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, main +gptq, LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit, main gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True -gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit-32g-actorder_True -gptq, TechxGenus/gemma-1.1-2b-it-GPTQ, main +gptq, TechxGenus/gemma-2b-it-GPTQ, main compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2, main @@ -16,4 +15,4 @@ compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main #compressed-tensors, mgoin/DeepSeek-Coder-V2-Lite-Instruct-FP8, main awq, casperhansen/mixtral-instruct-awq, main awq_marlin, casperhansen/mixtral-instruct-awq, main -fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main \ No newline at end of file +fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main diff --git a/tests/weight_loading/test_weight_loading.py b/tests/weight_loading/test_weight_loading.py index ce1835e254..e739164a5f 100644 --- a/tests/weight_loading/test_weight_loading.py +++ b/tests/weight_loading/test_weight_loading.py @@ -34,10 +34,7 @@ def test_weight_loading(vllm_runner): """ # MoE models need fp16. - NEEDS_FP16 = ( - QUANTIZATION == "gptq" - or MODEL_NAME == "nm-testing/test-w4a16-mixtral-actorder-group" - ) + NEEDS_FP16 = QUANTIZATION == "gptq" with vllm_runner( model_name=MODEL_NAME, revision=REVISION, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 93b5eec0f2..5856acc987 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -613,7 +613,6 @@ def gptq_gemm( b_q_weight: torch.Tensor, b_gptq_qzeros: torch.Tensor, b_gptq_scales: torch.Tensor, - b_g_idx: torch.Tensor, use_exllama: bool, use_v2_format: bool, bit: int, @@ -623,7 +622,6 @@ def gptq_gemm( b_q_weight, b_gptq_qzeros, b_gptq_scales, - b_g_idx, use_exllama, use_v2_format, bit, @@ -638,7 +636,6 @@ if hasattr(torch.ops._C, "gptq_gemm"): b_q_weight: torch.Tensor, b_gptq_qzeros: torch.Tensor, b_gptq_scales: torch.Tensor, - b_g_idx: torch.Tensor, use_exllama: bool, use_v2_format: bool, bit: int, @@ -648,8 +645,8 @@ if hasattr(torch.ops._C, "gptq_gemm"): ) -def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None: - torch.ops._C.gptq_shuffle(q_weight, q_perm, bit) +def gptq_shuffle(q_weight: torch.Tensor, bit: int) -> None: + torch.ops._C.gptq_shuffle(q_weight, bit) def gptq_gemm_rdna3( @@ -657,11 +654,10 @@ def gptq_gemm_rdna3( b_q_weight: torch.Tensor, b_qzeros: torch.Tensor, b_scales: torch.Tensor, - b_g_idx: torch.Tensor, use_v2_format: bool, ) -> torch.Tensor: return torch.ops._rocm_C.gptq_gemm_rdna3( - a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format + a, b_q_weight, b_qzeros, b_scales, use_v2_format ) @@ -673,7 +669,6 @@ if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3 b_q_weight: torch.Tensor, b_qzeros: torch.Tensor, b_scales: torch.Tensor, - b_g_idx: torch.Tensor, use_v2_format: bool, ) -> torch.Tensor: return torch.empty( @@ -689,7 +684,6 @@ if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3 b_q_weight: torch.Tensor, b_qzeros: torch.Tensor, b_scales: torch.Tensor, - b_g_idx: torch.Tensor, use_v2_format: bool, ) -> torch.Tensor: return torch.empty( @@ -1114,14 +1108,13 @@ def cutlass_mxfp4_moe_mm( # gptq_marlin def gptq_marlin_repack( b_q_weight: torch.Tensor, - perm: torch.Tensor, size_k: int, size_n: int, num_bits: int, is_a_8bit: bool = False, ) -> torch.Tensor: return torch.ops._C.gptq_marlin_repack( - b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit + b_q_weight, size_k, size_n, num_bits, is_a_8bit ) @@ -1130,7 +1123,6 @@ if hasattr(torch.ops._C, "gptq_marlin_repack"): @register_fake("_C::gptq_marlin_repack") def _gptq_marlin_repack_fake( b_q_weight: torch.Tensor, - perm: torch.Tensor, size_k: torch.SymInt, size_n: torch.SymInt, num_bits: int, @@ -1179,7 +1171,6 @@ if hasattr(torch.ops._C, "awq_marlin_repack"): def gptq_marlin_moe_repack( b_q_weight: torch.Tensor, - perm: torch.Tensor, size_k: int, size_n: int, num_bits: int, @@ -1194,14 +1185,13 @@ def gptq_marlin_moe_repack( ) for e in range(num_experts): output[e] = torch.ops._C.gptq_marlin_repack( - b_q_weight[e], perm[e], size_k, size_n, num_bits, is_a_8bit + b_q_weight[e], size_k, size_n, num_bits, is_a_8bit ) return output def awq_marlin_moe_repack( b_q_weight: torch.Tensor, - perm: torch.Tensor, size_k: int, size_n: int, num_bits: int, @@ -1238,14 +1228,11 @@ def marlin_gemm( a_scales: torch.Tensor | None, global_scale: torch.Tensor | None, b_zeros: torch.Tensor | None, - g_idx: torch.Tensor | None, - perm: torch.Tensor | None, workspace: torch.Tensor, b_q_type: ScalarType, size_m: int, size_n: int, size_k: int, - is_k_full: bool = True, use_atomic_add: bool = False, use_fp32_reduce: bool = False, is_zp_float: bool = False, @@ -1259,14 +1246,11 @@ def marlin_gemm( a_scales, global_scale, b_zeros, - g_idx, - perm, workspace, b_q_type.id, size_m, size_n, size_k, - is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float, @@ -1285,14 +1269,11 @@ if hasattr(torch.ops._C, "marlin_gemm"): a_scales: torch.Tensor | None, global_scale: torch.Tensor | None, b_zeros: torch.Tensor | None, - g_idx: torch.Tensor | None, - perm: torch.Tensor | None, workspace: torch.Tensor, b_q_type_id: int, size_m: torch.SymInt, size_n: torch.SymInt, size_k: torch.SymInt, - is_k_full: bool = True, use_atomic_add: bool = False, use_fp32_reduce: bool = False, is_zp_float: bool = False, @@ -1540,17 +1521,6 @@ if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b_grouped"): return torch.empty_like(b, memory_format=torch.contiguous_format) -def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: - return torch.ops._C.permute_cols(a, perm) - - -if hasattr(torch.ops._C, "permute_cols"): - - @register_fake("_C::permute_cols") - def _permute_cols_fake(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: - return torch.empty_like(a) - - # fp4 def scaled_fp4_quant( input: torch.Tensor, @@ -2545,8 +2515,6 @@ def moe_wna16_marlin_gemm( a_scales: torch.Tensor | None, global_scale: torch.Tensor | None, b_qzeros: torch.Tensor | None, - g_idx: torch.Tensor | None, - perm: torch.Tensor | None, workspace: torch.Tensor, sorted_token_ids: torch.Tensor, expert_ids: torch.Tensor, @@ -2559,7 +2527,6 @@ def moe_wna16_marlin_gemm( size_m: int, size_n: int, size_k: int, - is_k_full: bool, use_atomic_add: bool, use_fp32_reduce: bool, is_zp_float: bool, @@ -2576,8 +2543,6 @@ def moe_wna16_marlin_gemm( a_scales, global_scale, b_qzeros, - g_idx, - perm, workspace, sorted_token_ids, expert_ids, @@ -2590,7 +2555,6 @@ def moe_wna16_marlin_gemm( size_m, size_n, size_k, - is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float, @@ -2612,8 +2576,6 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_ a_scales: torch.Tensor | None, global_scale: torch.Tensor | None, b_qzeros: torch.Tensor | None, - g_idx: torch.Tensor | None, - perm: torch.Tensor | None, workspace: torch.Tensor, sorted_token_ids: torch.Tensor, expert_ids: torch.Tensor, @@ -2626,7 +2588,6 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_ size_m: int, size_n: int, size_k: int, - is_k_full: bool, use_atomic_add: bool, use_fp32_reduce: bool, is_zp_float: bool, @@ -4103,7 +4064,6 @@ def cpu_gemm_wna16( q_weight: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor | None, - g_idx: torch.Tensor | None, bias: torch.Tensor | None, pack_factor: int, isa_hint: str, @@ -4115,7 +4075,6 @@ def cpu_gemm_wna16( output, scales, zeros, - g_idx, bias, pack_factor, isa_hint, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py b/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py index b6fc82e87a..d8ff79057c 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py @@ -19,7 +19,6 @@ class MPLinearLayerConfig: act_type: torch.dtype group_size: int zero_points: bool - has_g_idx: bool out_type: torch.dtype | None = None @@ -40,7 +39,6 @@ class MPLinearKernel(ABC): w_q_param_name: str, w_s_param_name: str, w_zp_param_name: str | None = None, - w_gidx_param_name: str | None = None, ) -> None: assert self.can_implement(c) self.config = c @@ -48,10 +46,7 @@ class MPLinearKernel(ABC): self.w_s_name = w_s_param_name if c.zero_points: assert w_zp_param_name is not None - if c.has_g_idx: - assert w_gidx_param_name is not None self.w_zp_name: str | None = w_zp_param_name - self.w_gidx_name: str | None = w_gidx_param_name @abstractmethod def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -83,12 +78,10 @@ class MPLinearKernel(ABC): ) -> tuple[ torch.Tensor, # w_q torch.Tensor, # w_s - torch.Tensor | None, # w_zp, - torch.Tensor | None, # w_gidx + torch.Tensor | None, # w_zp ]: return ( getattr(layer, self.w_q_name), getattr(layer, self.w_s_name), getattr(layer, self.w_zp_name or "", None), - getattr(layer, self.w_gidx_name or "", None), ) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/allspark.py b/vllm/model_executor/kernels/linear/mixed_precision/allspark.py index 5f31538e40..704e661aed 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/allspark.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/allspark.py @@ -23,9 +23,6 @@ class AllSparkLinearKernel(MPLinearKernel): @classmethod def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: - if c.has_g_idx: - return False, "Act reordering currently not supported by AllSpark" - if c.zero_points: return False, "Zero points currently not supported by AllSpark" @@ -91,7 +88,7 @@ class AllSparkLinearKernel(MPLinearKernel): ) -> torch.Tensor: c = self.config gemm_args = self.gemm_args - w_q, w_s, _, _ = self._get_weight_params(layer) + w_q, w_s, _ = self._get_weight_params(layer) reshaped_x = x.reshape(-1, x.shape[-1]) out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/conch.py b/vllm/model_executor/kernels/linear/mixed_precision/conch.py index c65aa66cd6..7f9910a76a 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/conch.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/conch.py @@ -43,12 +43,6 @@ class ConchLinearKernel(MPLinearKernel): ) return False, error_msg - if c.has_g_idx: - return ( - False, - "Activation reordering (g_idx) is not supported by ConchLinearKernel", - ) - if find_spec("conch") is None: error_msg = ( "conch-triton-kernels is not installed, please " @@ -130,7 +124,7 @@ class ConchLinearKernel(MPLinearKernel): ) -> torch.Tensor: from conch.ops.quantization.gemm import mixed_precision_gemm - w_q, w_s, w_zp, _ = self._get_weight_params(layer) + w_q, w_s, w_zp = self._get_weight_params(layer) # Map channelwise group_size=-1 to the actual input dimension K. # The conch kernel computes stride_mul = block_k / group_size; diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py index c2627668cf..19f5283acb 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py @@ -144,9 +144,6 @@ class CPUWNA16LinearKernel(MPLinearKernel): if (not self.config.zero_points) and (self.w_zp_name is not None): setattr(layer, self.w_zp_name, None) - if (not self.config.has_g_idx) and (self.w_gidx_name is not None): - setattr(layer, self.w_gidx_name, None) - weights = getattr(layer, self.w_q_name) # Require GPTQ pack format assert weights.input_dim == weights.packed_dim @@ -172,7 +169,6 @@ class CPUWNA16LinearKernel(MPLinearKernel): supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV layer.use_w4a8 = ( envs.VLLM_CPU_INT4_W4A8 - and not self.config.has_g_idx and self.config.act_type == torch.bfloat16 and (supports_amx or supports_riscv) ) @@ -189,7 +185,7 @@ class CPUWNA16LinearKernel(MPLinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer) + w_q, w_s, w_zp = self._get_weight_params(layer) if layer.use_w4a8: x = ops.int4_scaled_mm_cpu( x=x, @@ -204,7 +200,6 @@ class CPUWNA16LinearKernel(MPLinearKernel): q_weight=w_q, scales=w_s, zeros=w_zp, - g_idx=w_gidx, bias=bias, pack_factor=8, # 32 // 4 isa_hint=layer.isa_hint, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py b/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py index 184a7f71d7..ea867cafd0 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py @@ -39,9 +39,6 @@ class CutlassW4A8LinearKernel(MPLinearKernel): if c.act_type != torch.float8_e4m3fn: return False, "CUTLASS W4A8 only supports FP8 (e4m3) activations" - if c.has_g_idx: - return False, "Act reordering not supported by CUTLASS W4A8" - if c.zero_points: return False, "Zero points not supported by CUTLASS W4A8" @@ -109,7 +106,7 @@ class CutlassW4A8LinearKernel(MPLinearKernel): bias: torch.Tensor | None = None, ) -> torch.Tensor: c = self.config - w_q, w_s, _, _ = self._get_weight_params(layer) + w_q, w_s, _ = self._get_weight_params(layer) w_ch_s = layer.weight_chan_scale x_2d = x.reshape(-1, x.shape[-1]) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/exllama.py b/vllm/model_executor/kernels/linear/mixed_precision/exllama.py index 3ad43a225f..d50a226636 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/exllama.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/exllama.py @@ -32,14 +32,6 @@ class ExllamaLinearKernel(MPLinearKernel): "Exllama is only supported on CUDA and ROCm", ) - if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]: - return ( - False, - "Act reordering currently not supported by Exllama, " - "when the input features are partitioned across " - "devices", - ) - if c.partition_weight_shape[1] % (32 // c.weight_type.size_bits) != 0: return ( False, @@ -111,29 +103,11 @@ class ExllamaLinearKernel(MPLinearKernel): layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False) ) - if c.has_g_idx: - - def transform_w_g_idx(x): - # Exllama wants the permutation array instead of the group - # indices - return torch.argsort(x).to(torch.int) - - self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore - else: - self.w_gidx_name = "g_idx" - empty_g_idx = torch.nn.Parameter( - torch.empty((0,), dtype=torch.int, device=device), requires_grad=False - ) - setattr(layer, self.w_gidx_name, empty_g_idx) - def transform_w_q(x): assert isinstance(x, BasevLLMParameter) - assert self.w_gidx_name is not None - g_idx = getattr(layer, self.w_gidx_name) - permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) x_cont = x.data.contiguous() - ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits) + ops.gptq_shuffle(x_cont, c.weight_type.size_bits) return x_cont def transform_w_s(x): @@ -157,17 +131,15 @@ class ExllamaLinearKernel(MPLinearKernel): x_2d = x.reshape(-1, x.shape[-1]) out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) - w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer) - + w_q, w_s, w_zp = self._get_weight_params(layer) # gptq_gemm supports GPTQv2 format by passing use_v2_format=True. # However, the MPLinearLayerConfig doesn't contain format info. # So hardcode GPTQv1 format here, to keep its behavior unchanged. use_v2_format = False assert w_zp is not None, "Zero points are required by Exllama" - assert w_g_idx is not None, "Group index is required by Exllama" output = ops.gptq_gemm( - x_2d, w_q, w_zp, w_s, w_g_idx, True, use_v2_format, c.weight_type.size_bits + x_2d, w_q, w_zp, w_s, True, use_v2_format, c.weight_type.size_bits ) if bias is not None: diff --git a/vllm/model_executor/kernels/linear/mixed_precision/humming.py b/vllm/model_executor/kernels/linear/mixed_precision/humming.py index 7d57fa680f..3de8b9ca41 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/humming.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/humming.py @@ -21,8 +21,6 @@ class HummingLinearKernel(MPLinearKernel): return False, "Humming is only supported on CUDA" if not has_humming(): return False, "Humming is not installed" - if c.has_g_idx: - return False, "Humming does not support act-order (g_idx)" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mixed_precision/machete.py b/vllm/model_executor/kernels/linear/mixed_precision/machete.py index 7953ed5e8e..9c09a10dbc 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/machete.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/machete.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import partial import torch @@ -12,7 +11,6 @@ from vllm.model_executor.layers.quantization.utils.machete_utils import ( query_machete_supported_quant_types, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( - pack_quantized_values_into_int32, unpack_quantized_values_into_int32, ) from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ @@ -35,14 +33,6 @@ class MacheteLinearKernel(MPLinearKernel): if not current_platform.is_device_capability(90): return False, "Machete requires compute capability of 90 (Hopper)" - if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]: - return ( - False, - "Act reordering currently not supported by Machete, " - "when the input features are partitioned across " - "devices", - ) - if c.weight_type not in query_machete_supported_quant_types(c.zero_points): return ( False, @@ -70,29 +60,9 @@ class MacheteLinearKernel(MPLinearKernel): def process_weights_after_loading(self, layer: torch.nn.Module): c = self.config - if c.has_g_idx: - assert self.w_gidx_name is not None - perm = torch.argsort(getattr(layer, self.w_gidx_name)).to(torch.int) - - self.act_perm = lambda x: x[:, perm] - # use `ops.permute_cols` if possible - if ( - c.act_type in [torch.float16, torch.bfloat16] - and c.partition_weight_shape[0] % 8 == 0 - ): - self.act_perm = partial(ops.permute_cols, perm=perm) - def transform_w_q(x): assert isinstance(x, BasevLLMParameter) permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) - if c.has_g_idx: - x_unpacked = unpack_quantized_values_into_int32( - x.data, c.weight_type, packed_dim=0 - ) - x_perm = x_unpacked[perm, :] - x.data = pack_quantized_values_into_int32( - x_perm, c.weight_type, packed_dim=0 - ) x.data = ops.machete_prepack_B( x.data.t().contiguous().t(), a_type=c.act_type, @@ -131,14 +101,11 @@ class MacheteLinearKernel(MPLinearKernel): bias: torch.Tensor | None = None, ) -> torch.Tensor: c = self.config - w_q, w_s, w_zp, _ = self._get_weight_params(layer) + w_q, w_s, w_zp = self._get_weight_params(layer) x_2d = x.reshape(-1, x.shape[-1]) out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) - if c.has_g_idx: - x_2d = self.act_perm(x_2d) - if c.zero_points: assert w_zp is not None else: diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py index 5a4f652971..387c33c110 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py @@ -8,10 +8,8 @@ from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils.marlin_utils import ( MARLIN_SUPPORTED_GROUP_SIZES, apply_gptq_marlin_linear, - check_marlin_supports_shape, marlin_act_int8_process_scales, - marlin_is_k_full, - marlin_make_empty_g_idx, + marlin_make_empty, marlin_make_workspace_new, marlin_pad_dim, marlin_pad_qweight, @@ -19,13 +17,11 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_padded_nk, marlin_permute_bias, marlin_permute_scales, - marlin_sort_g_idx, marlin_zero_points, query_marlin_supported_quant_types, unpack_cols, ) from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ -from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -59,16 +55,6 @@ class MarlinLinearKernel(MPLinearKernel): f"{MARLIN_SUPPORTED_GROUP_SIZES}", ) - if c.has_g_idx: - # Act-order couples K to the full-model group layout, so tile - # padding is not supported; keep the strict shape check. - return check_marlin_supports_shape( - c.partition_weight_shape[1], # out_features - c.partition_weight_shape[0], # in_features - c.full_weight_shape[0], # in_features - c.group_size, - ) - # A group straddling TP ranks cannot be fixed by padding. if ( c.group_size != -1 @@ -102,25 +88,16 @@ class MarlinLinearKernel(MPLinearKernel): getattr(layer, self.w_s_name).data * 512 ) - row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0] - self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel) - size_k, size_n = c.partition_weight_shape - if c.has_g_idx: - # Act-order shapes were strictly validated in can_implement. - padded_n, padded_k = size_n, size_k - else: - padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size) + padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size) # Allocate marlin workspace, reusing existing storage on reload. self.workspace = marlin_make_workspace_new( device, existing=getattr(self, "workspace", None) ) - # Default names since marlin requires empty parameters for these, + # Default name since marlin requires empty parameter for zp, # TODO: remove this requirement from marlin (allow optional tensors) - if self.w_gidx_name is None: - self.w_gidx_name = "g_idx" if self.w_zp_name is None: self.w_zp_name = "w_zp" @@ -131,7 +108,6 @@ class MarlinLinearKernel(MPLinearKernel): marlin_pad_qweight( x.data.contiguous(), size_n, size_k, padded_n, padded_k ), - perm=layer.g_idx_sort_indices, size_k=padded_k, size_n=padded_n, num_bits=c.weight_type.size_bits, @@ -172,18 +148,6 @@ class MarlinLinearKernel(MPLinearKernel): layer.input_global_scale = None return x - if c.has_g_idx: - g_idx, g_idx_sort_indices = marlin_sort_g_idx( - getattr(layer, self.w_gidx_name) - ) - self._transform_param(layer, self.w_gidx_name, lambda _: g_idx) - replace_parameter( - layer, "g_idx_sort_indices", g_idx_sort_indices, prefer_copy=True - ) - else: - setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device)) - layer.g_idx_sort_indices = marlin_make_empty_g_idx(device) - if c.zero_points: grouped_k = size_k // c.group_size if c.group_size != -1 else 1 padded_grouped_k = padded_k // c.group_size if c.group_size != -1 else 1 @@ -211,7 +175,7 @@ class MarlinLinearKernel(MPLinearKernel): ), ) else: - setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device)) + setattr(layer, self.w_zp_name, marlin_make_empty(device)) self._transform_param(layer, self.w_q_name, transform_w_q) self._transform_param(layer, self.w_s_name, transform_w_s) @@ -227,23 +191,16 @@ class MarlinLinearKernel(MPLinearKernel): bias: torch.Tensor | None = None, ) -> torch.Tensor: c = self.config - w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer) - - # `process_weights_after_loading` will ensure w_zp and w_gidx are not - # None for marlin - + w_q, w_s, w_zp = self._get_weight_params(layer) return apply_gptq_marlin_linear( input=x, weight=w_q, weight_scale=w_s, weight_zp=w_zp, # type: ignore - g_idx=w_gidx, # type: ignore - g_idx_sort_indices=layer.g_idx_sort_indices, workspace=self.workspace, wtype=c.weight_type, input_size_per_partition=c.partition_weight_shape[0], output_size_per_partition=c.partition_weight_shape[1], - is_k_full=self.is_k_full, input_global_scale=getattr(layer, "input_global_scale", None), bias=bias, input_dtype=c.act_type, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py index 268728f4bf..851422c501 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py @@ -86,13 +86,6 @@ class RDNA3W4A16LinearKernel(MPLinearKernel): "W4A16 kernel (qzeros packing)", ) - if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]: - return ( - False, - "Act-order with TP-partitioned input features is not " - "supported by the RDNA3 W4A16 kernel", - ) - return True, None # ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) ----- @@ -127,34 +120,14 @@ class RDNA3W4A16LinearKernel(MPLinearKernel): layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False) ) - # Act-order: convert g_idx to the inverse permutation array exllama - # expects (kernel reads a[perm[k]] instead of using groups indirected - # by g_idx[k]). - if c.has_g_idx: - - def transform_w_g_idx(x): - return torch.argsort(x).to(torch.int) - - self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore - else: - self.w_gidx_name = "g_idx" - empty_g_idx = torch.nn.Parameter( - torch.empty((0,), dtype=torch.int, device=device), - requires_grad=False, - ) - setattr(layer, self.w_gidx_name, empty_g_idx) - def transform_w_q(x): assert isinstance(x, BasevLLMParameter) - assert self.w_gidx_name is not None - g_idx = getattr(layer, self.w_gidx_name) - permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) x_cont = x.data.contiguous() # Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in # the same shuffled int32 layout and uses the (qa & 0x000F000F) # bit-trick on top. - ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits) + ops.gptq_shuffle(x_cont, c.weight_type.size_bits) return x_cont def transform_w_s(x): @@ -181,12 +154,9 @@ class RDNA3W4A16LinearKernel(MPLinearKernel): x_2d = x.reshape(-1, x.shape[-1]) out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) - w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer) - + w_q, w_s, w_zp = self._get_weight_params(layer) assert w_zp is not None, "Zero points are required by RDNA3 W4A16" - assert w_g_idx is not None, "g_idx tensor (possibly empty) required" - - output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, w_g_idx, False) + output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, False) if bias is not None: output.add_(bias) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index 01c353c554..9964c891b6 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -142,13 +142,13 @@ def _triton_w4a16_skinny_fmt_kernel( b = tl.interleave(b, b) b = (b >> shifts_full) & 0xF - g_idx = (k_start * BLOCK_K) // group_size - scale_ptrs = scales_ptr + offs_n * num_groups + g_idx + group_idx = (k_start * BLOCK_K) // group_size + scale_ptrs = scales_ptr + offs_n * num_groups + group_idx scale_mask = offs_n < N scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0) if HAS_ZP: - zp_ptrs = zp_ptr + offs_n * num_groups + g_idx + zp_ptrs = zp_ptr + offs_n * num_groups + group_idx zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0) b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None] else: @@ -476,9 +476,6 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel): if c.act_type not in (torch.float16, torch.bfloat16): return False, "requires float16 or bfloat16 activations" - if c.has_g_idx: - return False, "does not support g_idx reordering" - gs = c.group_size if gs not in SUPPORTED_GROUP_SIZES: return ( @@ -541,7 +538,7 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel): from vllm.utils.platform_utils import num_compute_units c = self.config - w_q, w_s, w_zp, _ = self._get_weight_params(layer) + w_q, w_s, w_zp = self._get_weight_params(layer) x_2d = x.reshape(-1, x.shape[-1]) N = w_q.shape[0] diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py index 1d1bf04b06..fd6fb033fe 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py @@ -126,10 +126,10 @@ def triton_w4a16_gemm_kernel( b = (b >> shifts) & 0xF # ---- Compute scale/zero group row index ---- - g_idx = (k_start * BLOCK_K) // group_size + group_idx = (k_start * BLOCK_K) // group_size # ---- Load scales: [BLOCK_N] → broadcast to [BLOCK_K, BLOCK_N] ---- - scale_offset = g_idx * N + offs_sn + scale_offset = group_idx * N + offs_sn scale_mask = offs_sn < N scales = tl.load(scales_ptr + scale_offset, mask=scale_mask, other=1.0) scales = tl.broadcast_to(scales[None, :], (BLOCK_K, BLOCK_N)) @@ -137,7 +137,7 @@ def triton_w4a16_gemm_kernel( # ---- Load / compute zeros ---- if HAS_ZP: # Load packed zeros row: [BLOCK_N//8] int32 - zero_offset = g_idx * (N // 8) + offs_bn + zero_offset = group_idx * (N // 8) + offs_bn zero_mask = offs_bn < N // 8 z_packed = tl.load(zeros_ptr + zero_offset, mask=zero_mask, other=0) # Unpack to [BLOCK_N] using same interleave+shift pattern @@ -237,7 +237,7 @@ def _triton_w4a16_gemm_impl( BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32 # The kernel loads scales/zeros for a single group per BLOCK_K tile - # (one g_idx per iteration). If BLOCK_K > group_size, rows at the tail + # (one group index per iteration). If BLOCK_K > group_size, rows at the tail # of the tile dequantize with the wrong group's scales, silently # corrupting the output. Clamp BLOCK_K to group_size to keep one # scale group per tile. @@ -340,13 +340,6 @@ class TritonW4A16LinearKernel(MPLinearKernel): "(8 int4 values packed per int32)", ) - if c.has_g_idx: - return ( - False, - "Activation reordering (g_idx) is not supported by " - "TritonW4A16LinearKernel", - ) - gs = c.group_size if ( gs not in TRITON_W4A16_SUPPORTED_GROUP_SIZES @@ -471,7 +464,7 @@ class TritonW4A16LinearKernel(MPLinearKernel): self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: c = self.config - w_q, w_s, w_zp, _ = self._get_weight_params(layer) + w_q, w_s, w_zp = self._get_weight_params(layer) x_2d = x.reshape(-1, x.shape[-1]).contiguous() out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py index 17900c7505..c6545f312a 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py @@ -54,10 +54,8 @@ class XPUwNa16LinearKernel(MPLinearKernel): return True, None def process_weights_after_loading(self, layer: torch.nn.Module): - # Default names since marlin requires empty parameters for these, + # Default name since marlin requires empty parameter for zp, # TODO: remove this requirement from marlin (allow optional tensors) - if self.w_gidx_name is None: - self.w_gidx_name = "g_idx" if self.w_zp_name is None: self.w_zp_name = "w_zp" @@ -90,17 +88,6 @@ class XPUwNa16LinearKernel(MPLinearKernel): setattr( layer, self.w_zp_name, Parameter(weight_zero_point, requires_grad=False) ) - if self.config.has_g_idx: - setattr( - layer, - self.w_gidx_name, - Parameter( - getattr(layer, self.w_gidx_name).data.t().contiguous(), - requires_grad=False, - ), - ) - else: - setattr(layer, self.w_gidx_name, None) def apply_weights( self, @@ -109,7 +96,7 @@ class XPUwNa16LinearKernel(MPLinearKernel): bias: torch.Tensor | None = None, ) -> torch.Tensor: reshaped_x = x.reshape(-1, x.shape[-1]) - w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer) + w_q, w_s, w_zp = self._get_weight_params(layer) out = torch.ops._xpu_C.int4_gemm_w4a16( reshaped_x, w_q.t(), @@ -117,7 +104,7 @@ class XPUwNa16LinearKernel(MPLinearKernel): w_s, w_zp, self.config.group_size, - w_gidx, + None, # Retained by the external XPU op ABI. ) return out @@ -224,7 +211,7 @@ class XPUW4A8IntLinearKernel(MPLinearKernel): layer.weight_scale, layer.weight_zero_point, self.config.group_size, - None, # g_idx not currently supported + None, # Retained by the external XPU op ABI. bias, ) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py index c3e8b17cc9..90968ab2c2 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py @@ -52,8 +52,6 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel): "zentorch_woq_linear} are not registered.", ) - if c.has_g_idx: - return False, "ZentorchWNA16 does not support activation re-ordering." return True, None def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool: @@ -62,12 +60,6 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel): Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()`` with ``layer`` untouched). """ - if ( - self.w_gidx_name is not None - and getattr(layer, self.w_gidx_name, None) is not None - ) or (getattr(self.config, "has_g_idx", False)): - return False - weight_packed = getattr(layer, self.w_q_name, None) weight_scale = getattr(layer, self.w_s_name, None) if weight_packed is None or weight_scale is None: @@ -112,7 +104,7 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel): if not self._zentorch_woq_eligible(layer): logger.info_once( "[zen_cpu] ZentorchWNA16 fast path not eligible for this " - "layer (AWQ pack layout, g_idx, or non-int32 storage); " + "layer (AWQ pack layout or non-int32 storage); " "falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)." ) super().process_weights_after_loading(layer) @@ -121,9 +113,6 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel): if (not self.config.zero_points) and (self.w_zp_name is not None): setattr(layer, self.w_zp_name, None) - if (not self.config.has_g_idx) and (self.w_gidx_name is not None): - setattr(layer, self.w_gidx_name, None) - weight_q = getattr(layer, self.w_q_name) weight_s = getattr(layer, self.w_s_name) weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 7014e40763..378b6a2314 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -81,10 +81,6 @@ def _fused_marlin_moe( input_global_scale2: torch.Tensor | None = None, global_scale1: torch.Tensor | None = None, global_scale2: torch.Tensor | None = None, - g_idx1: torch.Tensor | None = None, - g_idx2: torch.Tensor | None = None, - sort_indices1: torch.Tensor | None = None, - sort_indices2: torch.Tensor | None = None, w1_zeros: torch.Tensor | None = None, w2_zeros: torch.Tensor | None = None, workspace: torch.Tensor | None = None, @@ -92,7 +88,6 @@ def _fused_marlin_moe( intermediate_cache2: torch.Tensor | None = None, output: torch.Tensor | None = None, input_dtype: torch.dtype | None = None, - is_k_full: bool = True, activation_config: ApplyMoEActivationConfig | None = None, ) -> torch.Tensor: assert hidden_states.ndim == 2 @@ -142,8 +137,6 @@ def _fused_marlin_moe( a_scales1, global_scale1, w1_zeros, - g_idx1, - sort_indices1, workspace, sorted_token_ids, expert_ids, @@ -156,7 +149,6 @@ def _fused_marlin_moe( size_m=M, size_n=w13_num_shards * N, size_k=K, - is_k_full=is_k_full, use_atomic_add=False, use_fp32_reduce=True, is_zp_float=False, @@ -209,8 +201,6 @@ def _fused_marlin_moe( a_scales2, global_scale2, w2_zeros, - g_idx2, - sort_indices2, workspace, sorted_token_ids, expert_ids, @@ -223,7 +213,6 @@ def _fused_marlin_moe( size_m=M * num_topk, size_n=K, size_k=N, - is_k_full=is_k_full, use_atomic_add=False, use_fp32_reduce=True, is_zp_float=False, @@ -253,16 +242,11 @@ def fused_marlin_moe( input_global_scale2: torch.Tensor | None = None, global_scale1: torch.Tensor | None = None, global_scale2: torch.Tensor | None = None, - g_idx1: torch.Tensor | None = None, - g_idx2: torch.Tensor | None = None, - sort_indices1: torch.Tensor | None = None, - sort_indices2: torch.Tensor | None = None, w1_zeros: torch.Tensor | None = None, w2_zeros: torch.Tensor | None = None, workspace: torch.Tensor | None = None, intermediate_cache13: torch.Tensor | None = None, intermediate_cache2: torch.Tensor | None = None, - is_k_full: bool = True, output: torch.Tensor | None = None, input_dtype: torch.dtype | None = None, activation_config: ApplyMoEActivationConfig | None = None, @@ -277,12 +261,6 @@ def fused_marlin_moe( - w2 (torch.Tensor): The second set of expert weights. - w1_scale (torch.Tensor): Scale to be used for w1. - w2_scale (torch.Tensor): Scale to be used for w2. - - g_idx1 (torch.Tensor|None): The first set of act_order indices. - - g_idx2 (torch.Tensor|None): The second set of act_order indices. - - sort_indices1 (torch.Tensor|None): The first act_order input - permutation. - - sort_indices2 (torch.Tensor|None): The second act_order input - permutation. - topk_weights (torch.Tensor): Top-k weights. - topk_ids (torch.Tensor): Indices of topk-k elements. - w1_zeros (torch.Tensor|None): Optional zero points to be used for w1. @@ -372,10 +350,6 @@ def fused_marlin_moe( input_global_scale2=input_global_scale2, global_scale1=global_scale1, global_scale2=global_scale2, - g_idx1=g_idx1, - g_idx2=g_idx2, - sort_indices1=sort_indices1, - sort_indices2=sort_indices2, w1_zeros=w1_zeros, w2_zeros=w2_zeros, workspace=workspace, @@ -383,7 +357,6 @@ def fused_marlin_moe( intermediate_cache2=intermediate_cache2, output=None, input_dtype=input_dtype, - is_k_full=is_k_full, ).view(-1, topk, K) if output is None: @@ -416,16 +389,11 @@ def batched_fused_marlin_moe( input_global_scale2: torch.Tensor | None = None, global_scale1: torch.Tensor | None = None, global_scale2: torch.Tensor | None = None, - g_idx1: torch.Tensor | None = None, - g_idx2: torch.Tensor | None = None, - sort_indices1: torch.Tensor | None = None, - sort_indices2: torch.Tensor | None = None, w1_zeros: torch.Tensor | None = None, w2_zeros: torch.Tensor | None = None, workspace: torch.Tensor | None = None, intermediate_cache13: torch.Tensor | None = None, intermediate_cache2: torch.Tensor | None = None, - is_k_full: bool = True, output: torch.Tensor | None = None, input_dtype: torch.dtype | None = None, activation_func: Callable[..., None] | None = None, @@ -548,10 +516,6 @@ def batched_fused_marlin_moe( input_global_scale2=input_global_scale2, global_scale1=global_scale1, global_scale2=global_scale2, - g_idx1=g_idx1, - g_idx2=g_idx2, - sort_indices1=sort_indices1, - sort_indices2=sort_indices2, w1_zeros=w1_zeros, w2_zeros=w2_zeros, workspace=workspace, @@ -559,7 +523,6 @@ def batched_fused_marlin_moe( intermediate_cache2=intermediate_cache2, output=output.view(-1, K) if output is not None else output, input_dtype=input_dtype, - is_k_full=is_k_full, ) output = output.view(B, BATCH_TOKENS_MAX, K) @@ -574,11 +537,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): quant_config: FusedMoEQuantConfig, max_num_tokens: int | None = None, num_dispatchers: int | None = None, - w13_g_idx: torch.Tensor | None = None, - w2_g_idx: torch.Tensor | None = None, - w13_g_idx_sort_indices: torch.Tensor | None = None, - w2_g_idx_sort_indices: torch.Tensor | None = None, - is_k_full: bool = True, ): # TODO (varun) : Enable activation quantization assert ( @@ -588,11 +546,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): or quant_config.use_int8_w8a16 or quant_config.use_fp8_w8a16 ), "Supports only {mxfp,nvfp,int}4_w4a16, int8_w8a16 or fp8_w8a16" - self.w13_g_idx = w13_g_idx - self.w2_g_idx = w2_g_idx - self.w13_g_idx_sort_indices = w13_g_idx_sort_indices - self.w2_g_idx_sort_indices = w2_g_idx_sort_indices - self.is_k_full = is_k_full self.input_dtype = get_marlin_input_dtype() super().__init__( @@ -786,11 +739,6 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): # output buffer allocation. Please refer to workspace_shapes(). intermediate_cache13=workspace2, intermediate_cache2=workspace13, - g_idx1=self.w13_g_idx, - g_idx2=self.w2_g_idx, - sort_indices1=self.w13_g_idx_sort_indices, - sort_indices2=self.w2_g_idx_sort_indices, - is_k_full=self.is_k_full, input_dtype=self.input_dtype, ) return @@ -901,11 +849,6 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): output=output, intermediate_cache13=workspace2, intermediate_cache2=workspace13, - g_idx1=self.w13_g_idx, - g_idx2=self.w2_g_idx, - sort_indices1=self.w13_g_idx_sort_indices, - sort_indices2=self.w2_g_idx_sort_indices, - is_k_full=self.is_k_full, input_dtype=self.input_dtype, ) @@ -931,22 +874,12 @@ class BatchedMarlinExperts(MarlinExpertsBase): quant_config: FusedMoEQuantConfig, max_num_tokens: int, num_dispatchers: int, - w13_g_idx: torch.Tensor | None = None, - w2_g_idx: torch.Tensor | None = None, - w13_g_idx_sort_indices: torch.Tensor | None = None, - w2_g_idx_sort_indices: torch.Tensor | None = None, - is_k_full: bool = True, ): super().__init__( moe_config=moe_config, quant_config=quant_config, max_num_tokens=max_num_tokens, num_dispatchers=num_dispatchers, - w13_g_idx=w13_g_idx, - w2_g_idx=w2_g_idx, - w13_g_idx_sort_indices=w13_g_idx_sort_indices, - w2_g_idx_sort_indices=w2_g_idx_sort_indices, - is_k_full=is_k_full, ) def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -1038,14 +971,9 @@ class BatchedMarlinExperts(MarlinExpertsBase): output=output, intermediate_cache13=workspace13, intermediate_cache2=workspace2, - g_idx1=self.w13_g_idx, - g_idx2=self.w2_g_idx, - sort_indices1=self.w13_g_idx_sort_indices, - sort_indices2=self.w2_g_idx_sort_indices, w1_zeros=self.w1_zp, w2_zeros=self.w2_zp, input_dtype=self.input_dtype, - is_k_full=self.is_k_full, activation_func=activation_func, activation_config=self.activation_config, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 5d3212ae14..899f9bc3d9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -22,7 +22,6 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( BatchedMarlinExperts, MarlinExperts, - MarlinExpertsBase, ) from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( TritonWNA16Experts, @@ -362,11 +361,6 @@ def make_wna16_moe_kernel( moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], backend: WNA16MoEBackend = WNA16MoEBackend.MARLIN, - is_k_full: bool = False, - w13_g_idx: torch.Tensor | None = None, - w2_g_idx: torch.Tensor | None = None, - w13_g_idx_sort_indices: torch.Tensor | None = None, - w2_g_idx_sort_indices: torch.Tensor | None = None, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -413,15 +407,6 @@ def make_wna16_moe_kernel( logger.info_once("Using %s", experts_cls.__name__, scope="local") extra_args: dict[str, Any] = {} - if issubclass(experts_cls, MarlinExpertsBase): - extra_args = { - "w13_g_idx": w13_g_idx, - "w2_g_idx": w2_g_idx, - "w13_g_idx_sort_indices": w13_g_idx_sort_indices, - "w2_g_idx_sort_indices": w2_g_idx_sort_indices, - "is_k_full": is_k_full, - } - if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() assert max_num_tokens is not None @@ -450,8 +435,6 @@ def _process_weights_flashinfer( w2_qweight: torch.Tensor, w13_scales: torch.Tensor, w2_scales: torch.Tensor, - w13_g_idx: torch.Tensor, - w2_g_idx: torch.Tensor, w13_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, ) -> tuple[ @@ -459,10 +442,6 @@ def _process_weights_flashinfer( torch.Tensor, # w2_qweight torch.Tensor, # w13_scales torch.Tensor, # w2_scales - torch.Tensor, # w13_g_idx - torch.Tensor, # w2_g_idx - torch.Tensor | None, # w13_g_idx_sort_indices - torch.Tensor | None, # w2_g_idx_sort_indices torch.Tensor | None, # w13_qzeros torch.Tensor | None, # w2_qzeros torch.Tensor | None, # w13_input_global_scale @@ -475,7 +454,7 @@ def _process_weights_flashinfer( Steps ----- 1. Transform weights/scales via ``prepare_static_weights_for_trtllm_mxint4_moe``. - 2. Return transformed tensors, passing through g_idx/bias unchanged. + 2. Return transformed tensors and biases. """ from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( prepare_static_weights_for_trtllm_mxint4_moe, @@ -493,10 +472,6 @@ def _process_weights_flashinfer( dict_weights_mxint4["gemm2_weights"], dict_weights_mxint4["gemm1_scales"], dict_weights_mxint4["gemm2_scales"], - w13_g_idx, - w2_g_idx, - None, - None, None, None, None, @@ -540,13 +515,10 @@ def _process_weights_marlin( num_bits: int, pack_factor: int, group_size: int, - actorder: str | None, w13_qweight: torch.Tensor, w2_qweight: torch.Tensor, w13_scales: torch.Tensor, w2_scales: torch.Tensor, - w13_g_idx: torch.Tensor, - w2_g_idx: torch.Tensor, w13_qzeros: torch.Tensor | None = None, w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, @@ -556,10 +528,6 @@ def _process_weights_marlin( torch.Tensor, # w2_qweight torch.Tensor, # w13_scales torch.Tensor, # w2_scales - torch.Tensor, # w13_g_idx - torch.Tensor, # w2_g_idx - torch.Tensor, # w13_g_idx_sort_indices - torch.Tensor, # w2_g_idx_sort_indices torch.Tensor | None, # w13_qzeros torch.Tensor | None, # w2_qzeros torch.Tensor | None, # w13_input_global_scale @@ -573,10 +541,9 @@ def _process_weights_marlin( Steps ----- 1. Optional FP8 preprocessing of packed weights / scales. - 2. Sort / reset g_idx tensors for act-order handling. - 3. Repack weights via ``gptq_marlin_moe_repack``. - 4. Permute scales (and optionally extract INT8 global scales). - 5. Permute bias tensors. + 2. Repack weights via ``gptq_marlin_moe_repack``. + 3. Permute scales (and optionally extract INT8 global scales). + 4. Permute bias tensors. """ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 @@ -584,8 +551,6 @@ def _process_weights_marlin( marlin_w2_qweight: torch.Tensor marlin_w13_scales: torch.Tensor marlin_w2_scales: torch.Tensor - w13_g_idx_sort_indices: torch.Tensor | None = None - w2_g_idx_sort_indices: torch.Tensor | None = None w13_input_global_scale: torch.Tensor | None = None w2_input_global_scale: torch.Tensor | None = None w13_bias_out: torch.Tensor | None = None @@ -606,13 +571,9 @@ def _process_weights_marlin( # --- Pad the intermediate size to a valid Marlin thread tile --- # GPTQ packs along K: w13's N is in the (shard) columns, w2's N in the rows. - # Act-order keeps the strict shape and is never padded. N = layer.intermediate_size_per_partition padded_N = marlin_moe_padded_intermediate(N, group_size) if padded_N != N: - assert actorder != "group", ( - "Marlin MoE thread-tile padding is unsupported with act-order" - ) marlin_w13_qweight = _pad_w13_shard_cols(marlin_w13_qweight, N, padded_N) marlin_w2_qweight = _pad_rows(marlin_w2_qweight, padded_N // pack_factor) marlin_w13_scales = _pad_w13_shard_cols(marlin_w13_scales, N, padded_N) @@ -627,44 +588,9 @@ def _process_weights_marlin( if w13_bias is not None: w13_bias = _pad_w13_bias(w13_bias, N, padded_N) - # --- Process act_order (g_idx) --- - if actorder == "group": - num_experts = w13_g_idx.shape[0] - w13_g_idx_sort_indices = torch.empty_like(w13_g_idx) - w2_g_idx_sort_indices = torch.empty_like(w2_g_idx) - w13_sorted_g_idx = torch.empty_like(w13_g_idx) - w2_sorted_g_idx = torch.empty_like(w2_g_idx) - for e in range(num_experts): - w13_g_idx_sort_indices[e] = torch.argsort(w13_g_idx[e]).to(torch.int32) - w2_g_idx_sort_indices[e] = torch.argsort(w2_g_idx[e]).to(torch.int32) - w13_sorted_g_idx[e] = w13_g_idx[e][w13_g_idx_sort_indices[e]] - w2_sorted_g_idx[e] = w2_g_idx[e][w2_g_idx_sort_indices[e]] - w13_g_idx = w13_sorted_g_idx - w2_g_idx = w2_sorted_g_idx - else: - num_experts = w13_g_idx.shape[0] - device = w13_g_idx.device - w13_g_idx = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - w2_g_idx = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - # --- Repack weights --- marlin_w13_qweight = ops.gptq_marlin_moe_repack( marlin_w13_qweight, - w13_g_idx_sort_indices, marlin_w13_qweight.shape[1] * pack_factor, marlin_w13_qweight.shape[2], num_bits, @@ -672,7 +598,6 @@ def _process_weights_marlin( ) marlin_w2_qweight = ops.gptq_marlin_moe_repack( marlin_w2_qweight, - w2_g_idx_sort_indices, marlin_w2_qweight.shape[1] * pack_factor, marlin_w2_qweight.shape[2], num_bits, @@ -734,10 +659,6 @@ def _process_weights_marlin( marlin_w2_qweight, marlin_w13_scales, marlin_w2_scales, - w13_g_idx, - w2_g_idx, - w13_g_idx_sort_indices, - w2_g_idx_sort_indices, w13_qzeros, w2_qzeros, w13_input_global_scale, @@ -766,10 +687,6 @@ def _process_awq_weights_marlin( torch.Tensor, # w2_qweight torch.Tensor, # w13_scales torch.Tensor, # w2_scales - torch.Tensor | None, # w13_g_idx - torch.Tensor | None, # w2_g_idx - torch.Tensor | None, # w13_g_idx_sort_indices - torch.Tensor | None, # w2_g_idx_sort_indices torch.Tensor | None, # w13_qzeros torch.Tensor | None, # w2_qzeros torch.Tensor | None, # w13_input_global_scale @@ -782,8 +699,6 @@ def _process_awq_weights_marlin( AWQ checkpoints use a different packing order than GPTQ, so they need AWQ-specific weight repacking and zero-point conversion before Marlin runs. """ - num_experts = w13_qweight.shape[0] - device = w13_qweight.device is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 w13_input_global_scale: torch.Tensor | None = None w2_input_global_scale: torch.Tensor | None = None @@ -823,18 +738,8 @@ def _process_awq_weights_marlin( if w13_bias is not None: w13_bias = _pad_w13_bias(w13_bias, N, padded_N) - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - marlin_w13_qweight = ops.awq_marlin_moe_repack( w13_qweight, - w13_g_idx_sort_indices, size_k=w13_qweight.shape[1], size_n=w13_qweight.shape[2] * pack_factor, num_bits=weight_bits, @@ -842,7 +747,6 @@ def _process_awq_weights_marlin( ) marlin_w2_qweight = ops.awq_marlin_moe_repack( w2_qweight, - w2_g_idx_sort_indices, size_k=w2_qweight.shape[1], size_n=w2_qweight.shape[2] * pack_factor, num_bits=weight_bits, @@ -898,10 +802,6 @@ def _process_awq_weights_marlin( marlin_w2_qweight, marlin_w13_scales, marlin_w2_scales, - None, - None, - w13_g_idx_sort_indices, - w2_g_idx_sort_indices, marlin_w13_qzeros, marlin_w2_qzeros, w13_input_global_scale, @@ -917,8 +817,6 @@ def _process_weights_cpu( w2: torch.Tensor, w13_scale: torch.Tensor, w2_scale: torch.Tensor, - w13_g_idx: torch.Tensor | None = None, - w2_g_idx: torch.Tensor | None = None, w13_qzeros: torch.Tensor | None = None, w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, @@ -928,10 +826,6 @@ def _process_weights_cpu( torch.Tensor, # w2_qweight torch.Tensor, # w13_scales torch.Tensor, # w2_scales - torch.Tensor | None, # w13_g_idx - torch.Tensor | None, # w2_g_idx - torch.Tensor | None, # w13_g_idx_sort_indices - torch.Tensor | None, # w2_g_idx_sort_indices torch.Tensor | None, # w13_qzeros torch.Tensor | None, # w2_qzeros torch.Tensor | None, # w13_input_global_scale @@ -962,7 +856,7 @@ def _process_weights_cpu( if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act: raise NotImplementedError( "CPU WNA16 MoE backend does not support GPTQ with " - "desc_act=True. The fused MoE kernel has no g_idx " + "desc_act=True. The fused MoE kernel has no activation " "reordering support." ) cpu_quant_algo = ops.CPUQuantAlgo.GPTQ @@ -1009,10 +903,6 @@ def _process_weights_cpu( blocked_w2, blocked_s13, blocked_s2, - w13_g_idx, - w2_g_idx, - None, # w13_g_idx_sort_indices (unused on CPU) - None, # w2_g_idx_sort_indices (unused on CPU) blocked_z13, blocked_z2, None, # w13_input_global_scale @@ -1352,10 +1242,6 @@ def _process_weights_emulation_gptq( w2_bf16, # w2_qweight (now bf16, not int32) dummy, # w13_scales (unused; nulled out in Int4EmulationTritonExperts) dummy, # w2_scales (unused) - None, # w13_g_idx - None, # w2_g_idx - None, # w13_g_idx_sort_indices - None, # w2_g_idx_sort_indices None, # w13_qzeros None, # w2_qzeros None, # w13_input_global_scale @@ -1413,10 +1299,6 @@ def _process_weights_emulation_awq( None, None, None, - None, - None, - None, - None, ) @@ -1429,8 +1311,6 @@ def convert_to_wna16_moe_kernel_format( w2: torch.Tensor, w13_scale: torch.Tensor, w2_scale: torch.Tensor, - w13_g_idx: torch.Tensor | None = None, - w2_g_idx: torch.Tensor | None = None, w13_qzeros: torch.Tensor | None = None, w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, @@ -1441,10 +1321,6 @@ def convert_to_wna16_moe_kernel_format( torch.Tensor, # w2_qweight torch.Tensor, # w13_scales torch.Tensor, # w2_scales - torch.Tensor | None, # w13_g_idx - torch.Tensor | None, # w2_g_idx - torch.Tensor | None, # w13_g_idx_sort_indices - torch.Tensor | None, # w2_g_idx_sort_indices torch.Tensor | None, # w13_qzeros torch.Tensor | None, # w2_qzeros torch.Tensor | None, # w13_input_global_scale @@ -1509,14 +1385,12 @@ def convert_to_wna16_moe_kernel_format( num_bits = quant_config.quant_type.size_bits pack_factor = quant_config.pack_factor group_size = quant_config.group_size - actorder = "group" if quant_config.desc_act else None elif isinstance(quant_config, QuantizationArgs): num_bits = quant_config.num_bits pack_factor = 32 // quant_config.num_bits group_size = ( -1 if quant_config.group_size is None else quant_config.group_size ) - actorder = quant_config.actorder else: raise TypeError( "Marlin WNA16 MoE backend requires AutoGPTQConfig, AutoAWQConfig or " @@ -1543,22 +1417,16 @@ def convert_to_wna16_moe_kernel_format( w2_bias, ) else: - if w13_g_idx is None or w2_g_idx is None: - raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") - return _process_weights_marlin( layer, input_dtype, num_bits, pack_factor, group_size, - actorder, w13, w2, w13_scale, w2_scale, - w13_g_idx, - w2_g_idx, w13_qzeros, w2_qzeros, w13_bias, @@ -1571,8 +1439,6 @@ def convert_to_wna16_moe_kernel_format( w2, w13_scale, w2_scale, - w13_g_idx, - w2_g_idx, w13_qzeros, w2_qzeros, w13_bias, @@ -1584,8 +1450,6 @@ def convert_to_wna16_moe_kernel_format( w2, w13_scale, w2_scale, - w13_g_idx, - w2_g_idx, w13_bias, w2_bias, ) @@ -1608,16 +1472,11 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) - empty = torch.empty((0,), dtype=torch.int32, device=w13.device) return ( w13_xpu, w2_xpu, w13_scale_xpu, w2_scale_xpu, - empty, # w13_g_idx - empty, # w2_g_idx - empty, # w13_g_idx_sort_indices - empty, # w2_g_idx_sort_indices None, # w13_qzeros — sym int4 on XPU has none; kernel does uint4b8→s4 None, # w2_qzeros None, # w13_input_global_scale @@ -1712,10 +1571,6 @@ def convert_to_wna16_moe_kernel_format( w2_uint8, w13_scale, w2_scale, - None, - None, - None, - None, w13_qzeros, w2_qzeros, None, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index e9a7fc4aa7..5d284bcc80 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -168,12 +168,6 @@ class RoutedExperts(PluggableLayer): "global_num_experts": moe_config.num_experts, } - # need full intermediate size pre-sharding for WNA16 act order - if self._needs_intermediate_size_param(self.quant_method): - moe_quant_params["intermediate_size_full"] = ( - self.moe_config.intermediate_size - ) - self.quant_method.create_weights(layer=self, **moe_quant_params) self.lora_base_layer_prefix = "" @@ -206,14 +200,6 @@ class RoutedExperts(PluggableLayer): assert isinstance(quant_method, FusedMoEMethodBase) return quant_method - # TODO(bnell): make this a method on quant_method - def _needs_intermediate_size_param(self, quant_method: FusedMoEMethodBase) -> bool: - return quant_method.__class__.__name__ in ( - "AutoGPTQMoEMethod", - "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsW4A16FlydslMoEMethod", - ) - def _ensure_moe_quant_config_init(self): if self.quant_method.moe_quant_config is None: # Note: the moe_quant_config can't be constructed until after @@ -349,7 +335,6 @@ class RoutedExperts(PluggableLayer): shard_id: str, loaded_weight: torch.Tensor, tp_rank: int, - load_full_w2: bool = False, ): """ Load grouped weight scales for group quantization or model weights @@ -360,17 +345,13 @@ class RoutedExperts(PluggableLayer): shard_id: either w1, w2, or w3 loaded_weight: checkpoint weight to load into the param tp_rank: tensor parallel rank - load_full_w2: whether or not the w2 loaded should be sharded. """ if shard_id == "w2": - # In the case where we have actorder/g_idx, we do not partition the - # w2 scales, as indicated by `load_full` argument, for all tp cases self._load_w2( shard_dim=shard_dim, loaded_weight=loaded_weight, expert_data=expert_data, tp_rank=tp_rank, - load_full=load_full_w2, ) elif shard_id in ("w1", "w3"): self._load_w13( @@ -532,13 +513,11 @@ class RoutedExperts(PluggableLayer): shard_dim: int, loaded_weight: torch.Tensor, tp_rank: int, - load_full: bool = False, ): # Index the loaded weight for tp sharding. # down_proj: "RowParallel" so tp sharding on input_dim - # Only narrow if the loaded_weight is not a scalar (0-dim tensor) - # and we're not loading the full weight - if not load_full and loaded_weight.ndim > 0: + # Only narrow if the loaded_weight is not a scalar (0-dim tensor). + if loaded_weight.ndim > 0: # Same padding fix as _load_w13: use unpadded per-rank size. tp_size = self.moe_config.moe_parallel_config.tp_size loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size @@ -571,25 +550,6 @@ class RoutedExperts(PluggableLayer): # _to_scalar's reshape(()) would reject the size-2 weight_shape. param_data[expert_id] = loaded_weight - def _load_g_idx( - self, - shard_id: str, - expert_data: torch.Tensor, - shard_dim: int, - loaded_weight: torch.Tensor, - tp_rank: int, - ): - if shard_id == "w2": - self._load_w2( - shard_dim=shard_dim, - loaded_weight=loaded_weight, - expert_data=expert_data, - tp_rank=tp_rank, - ) - else: - assert shard_id in ("w1", "w3") - expert_data.copy_(loaded_weight) - @overload def weight_loader( self, @@ -739,17 +699,6 @@ class RoutedExperts(PluggableLayer): ) return True if return_success else None - # Case g_idx - if "g_idx" in weight_name: - self._load_g_idx( - shard_dim=0, - shard_id=shard_id, - loaded_weight=loaded_weight, - expert_data=expert_data, - tp_rank=self.moe_config.tp_rank, - ) - return True if return_success else None - # TODO @dsikka: ModelOpt should follow the proper MoE loading pattern if "ModelOpt" in quant_method_name: # Determine per-tensor weight scale patterns based on variant @@ -853,7 +802,6 @@ class RoutedExperts(PluggableLayer): loaded_weight=loaded_weight, expert_data=expert_data, tp_rank=self.moe_config.tp_rank, - load_full_w2=getattr(param, "load_full_w2", False), ) elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value: self._load_per_tensor_weight_scale( diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 713bd7bd5e..f8c7269add 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -439,7 +439,6 @@ class AutoAWQMarlinLinearMethod(LinearMethodBase): act_type=params_dtype if self.input_dtype is None else self.input_dtype, group_size=self.quant_config.group_size, zero_points=self.quant_config.zero_point, - has_g_idx=False, ) kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config) @@ -558,11 +557,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): } ) - intermediate_size_full = extra_weight_attrs.pop( - "intermediate_size_full", intermediate_size_per_partition - ) - self.is_k_full = intermediate_size_per_partition == intermediate_size_full - w13_qweight = Parameter( torch.empty( num_experts, @@ -672,10 +666,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): w2, w13_scale, w2_scale, - w13_g_idx, - w2_g_idx, - w13_g_idx_sort_indices, - w2_g_idx_sort_indices, w13_qzeros, w2_qzeros, w13_input_global_scale, @@ -696,14 +686,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): replace_parameter(layer, "w13_scales", w13_scale) replace_parameter(layer, "w2_scales", w2_scale) - _replace_or_register_parameter( - layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices - ) - _replace_or_register_parameter( - layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices - ) - _replace_or_register_parameter(layer, "w13_g_idx", w13_g_idx) - _replace_or_register_parameter(layer, "w2_g_idx", w2_g_idx) _replace_or_register_parameter(layer, "w13_qzeros", w13_qzeros) _replace_or_register_parameter(layer, "w2_qzeros", w2_qzeros) _replace_or_register_parameter( @@ -726,11 +708,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): moe_config=self.moe, experts_cls=self.experts_cls, backend=self.wna16_moe_backend, - is_k_full=self.is_k_full, - w13_g_idx=getattr(layer, "w13_g_idx", None), - w2_g_idx=getattr(layer, "w2_g_idx", None), - w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None), - w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None), routing_tables=layer._expert_routing_tables(), ) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 1e6570abee..893f85b207 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -40,6 +40,7 @@ from vllm.model_executor.layers.quantization.utils import replace_parameter from vllm.model_executor.layers.quantization.utils.gptq_utils import ( get_dynamic_override, get_linear_quant_method, + normalize_and_validate_gptq_desc_act, override_config, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( @@ -59,7 +60,6 @@ from vllm.model_executor.parameter import ( GroupQuantScaleParameter, PackedColumnParameter, PackedvLLMParameter, - RowvLLMParameter, ) from vllm.scalar_type import scalar_types from vllm.transformers_utils.config import get_safetensors_params_metadata @@ -115,10 +115,7 @@ class AutoGPTQConfig(QuantizationConfig): modules_in_block_to_quantize: list[str] | None = None, ) -> None: super().__init__() - if desc_act and group_size == -1: - # In this case, act_order == True is the same as act_order == False - # (since we have only one group per output channel) - desc_act = False + desc_act = normalize_and_validate_gptq_desc_act(desc_act, group_size, dynamic) # GPTQModel use `dynamic` config property to allow per module # quantization config so each module can be individually optimized. @@ -150,7 +147,7 @@ class AutoGPTQConfig(QuantizationConfig): self.pack_factor = 32 // weight_bits # packed into int32 self.group_size = group_size - self.desc_act = desc_act + self.desc_act: bool = desc_act self.lm_head_quantized = lm_head_quantized self.full_config = full_config @@ -244,7 +241,7 @@ class AutoGPTQConfig(QuantizationConfig): from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config if not check_moe_marlin_supports_layer( - layer, self.group_size, allow_tile_padding=not self.desc_act + layer, self.group_size, allow_tile_padding=True ): logger.warning_once( f"Layer '{prefix}' is not supported by GPTQMoeMarlin. " @@ -348,7 +345,6 @@ class AutoGPTQLinearMethod(LinearMethodBase): act_type=params_dtype if input_dtype is None else input_dtype, group_size=self.quant_config.group_size, zero_points=False, - has_g_idx=self.quant_config.desc_act, ) kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config) @@ -365,7 +361,7 @@ class AutoGPTQLinearMethod(LinearMethodBase): # Determine sharding if marlin_repeat_scales_on_all_ranks( - self.quant_config.desc_act, self.quant_config.group_size, is_row_parallel + self.quant_config.group_size, is_row_parallel ): # By setting scale_dim == None, weight_loader will # repeat the scales on each GPU in TP>1 case. @@ -391,16 +387,6 @@ class AutoGPTQLinearMethod(LinearMethodBase): weight_loader=weight_loader, ) - # Activation order - g_idx = RowvLLMParameter( - data=torch.empty( - input_size_per_partition, - dtype=torch.int32, - ), - input_dim=0, - weight_loader=weight_loader, - ) - qzeros_args = { "data": torch.empty( scales_and_zp_size, @@ -440,7 +426,6 @@ class AutoGPTQLinearMethod(LinearMethodBase): ) layer.register_parameter("qweight", qweight) - layer.register_parameter("g_idx", g_idx) layer.register_parameter("scales", scales) layer.register_parameter("qzeros", qzeros) @@ -449,7 +434,6 @@ class AutoGPTQLinearMethod(LinearMethodBase): w_q_param_name="qweight", w_s_param_name="scales", w_zp_param_name="qzeros", - w_gidx_param_name="g_idx", ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -492,7 +476,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): quant_config=self.quant_config, may_have_zp=not self.quant_config.is_sym, may_have_bias=True, - allow_tile_padding=not self.quant_config.desc_act, + allow_tile_padding=True, ) def create_weights( @@ -512,19 +496,9 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): "W8A8-INT8 is not supported by marlin kernel." ) - intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full") - - self.is_k_full = (not self.quant_config.desc_act) or ( - intermediate_size_per_partition == intermediate_size_full - ) - if self.quant_config.group_size != -1: scales_size13 = hidden_size // self.quant_config.group_size - w2_scales_size = ( - intermediate_size_full - if self.quant_config.desc_act - else intermediate_size_per_partition - ) + w2_scales_size = intermediate_size_per_partition scales_size2 = w2_scales_size // self.quant_config.group_size strategy = FusedMoeWeightScaleSupported.GROUP.value else: @@ -579,8 +553,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_scales", w2_scales) set_weight_attrs(w2_scales, extra_weight_attrs) - # don't shard the w2 scales when running act order - set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act}) # up_proj zero points w13_qzeros = torch.nn.Parameter( torch.empty( @@ -607,48 +579,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_qzeros", w2_qzeros) set_weight_attrs(w2_qzeros, extra_weight_attrs) - # don't shard the w2 scales when running act order - set_weight_attrs(w2_qzeros, {"load_full_w2": self.quant_config.desc_act}) - w13_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_g_idx", w13_g_idx) - set_weight_attrs(w13_g_idx, extra_weight_attrs) - w2_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_g_idx", w2_g_idx) - set_weight_attrs(w2_g_idx, extra_weight_attrs) - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) # Some GPTQ checkpoints contain expert biases even when the model # architecture does not declare them. Zero initialization keeps @@ -712,8 +642,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): w2=layer.w2_qweight, w13_scale=layer.w13_scales, w2_scale=layer.w2_scales, - w13_g_idx=layer.w13_g_idx, - w2_g_idx=layer.w2_g_idx, w13_bias=w13_bias, w2_bias=w2_bias, w13_qzeros=getattr(layer, "w13_qzeros", None), @@ -730,10 +658,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): w2, w13_scale, w2_scale, - w13_g_idx, - w2_g_idx, - w13_g_idx_sort_indices, - w2_g_idx_sort_indices, w13_qzeros, w2_qzeros, w13_input_global_scale, @@ -746,10 +670,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): replace_parameter(layer, "w2_qweight", w2) replace_parameter(layer, "w13_scales", w13_scale) replace_parameter(layer, "w2_scales", w2_scale) - replace_parameter(layer, "w13_g_idx", w13_g_idx) - replace_parameter(layer, "w2_g_idx", w2_g_idx) - replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) replace_or_register("w13_input_global_scale", w13_input_global_scale) replace_or_register("w2_input_global_scale", w2_input_global_scale) replace_or_register("w13_bias", w13_bias) @@ -772,11 +692,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): moe_config=self.moe, experts_cls=self.experts_cls, backend=self.wna16_moe_backend, - is_k_full=self.is_k_full, - w13_g_idx=getattr(layer, "w13_g_idx", None), - w2_g_idx=getattr(layer, "w2_g_idx", None), - w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None), - w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None), routing_tables=layer._expert_routing_tables(), ) diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index f83c1ec76f..5f2c7eebd2 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -238,6 +238,13 @@ class QuantizationConfig(ABC): } return WeightsMapper(orig_to_new_regex=orig_to_new_regex) + @staticmethod + def get_checkpoint_weight_mapper() -> "WeightsMapper": + """Discard activation-order metadata unused by supported kernels.""" + from vllm.model_executor.models.utils import WeightsMapper + + return WeightsMapper(orig_to_new_suffix={".g_idx": None}) + def apply_vllm_mapper( # noqa: B027 self, hf_to_vllm_mapper: "WeightsMapper" ): diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 86a01434e5..306a540055 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast import torch from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig from compressed_tensors.quantization import ( + ActivationOrdering, QuantizationArgs, QuantizationStrategy, QuantizationType, @@ -325,9 +326,20 @@ class CompressedTensorsConfig(QuantizationConfig): targets = quant_config.get("targets") for target in targets: target_scheme_map[target] = {} - target_scheme_map[target]["weights"] = QuantizationArgs.model_validate( + weight_quant = QuantizationArgs.model_validate( quant_config.get("weights") ) + if weight_quant.actorder in ( + ActivationOrdering.GROUP.value, + ActivationOrdering.DYNAMIC.value, + ): + raise ValueError( + "Compressed-tensors group/dynamic activation ordering " + f"(actorder={weight_quant.actorder!r}) is no longer supported. " + "Use a checkpoint with static/weight activation ordering or " + "without activation ordering." + ) + target_scheme_map[target]["weights"] = weight_quant target_scheme_map[target]["input_activations"] = None target_scheme_map[target]["format"] = quant_config.get("format") @@ -754,7 +766,6 @@ class CompressedTensorsConfig(QuantizationConfig): strategy=weight_quant.strategy, symmetric=weight_quant.symmetric, group_size=weight_quant.group_size, - actorder=weight_quant.actorder, ) # Must come before the WNA16 check; standard 4/8-bit weight-only (no @@ -803,7 +814,6 @@ class CompressedTensorsConfig(QuantizationConfig): strategy=weight_quant.strategy, symmetric=weight_quant.symmetric, group_size=weight_quant.group_size, - actorder=weight_quant.actorder, layer_name=layer_name, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py index 076da0a081..f879d78a1b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py @@ -104,8 +104,6 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod): assert weight_quant.strategy == "group" assert weight_quant.group_size == 32 self.group_size = weight_quant.group_size - # grouped actorder isn't supported by this kernel - assert weight_quant.actorder != "group" assert weight_quant.symmetric, ( "Only symmetric quantization is supported for MoE. " "Try --moe-backend emulation." @@ -180,7 +178,6 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod): ) layer.register_parameter("w2_weight_scale", w2_scale) set_weight_attrs(w2_scale, extra_weight_attrs) - set_weight_attrs(w2_scale, {"load_full_w2": False}) w2_weight_shape = torch.nn.Parameter( torch.empty(num_experts, 2), requires_grad=False @@ -194,50 +191,6 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod): layer.register_parameter("w13_weight_shape", w13_weight_shape) set_weight_attrs(w13_weight_shape, extra_weight_attrs) - w13_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_g_idx", w13_g_idx) - set_weight_attrs(w13_g_idx, extra_weight_attrs) - - w2_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_g_idx", w2_g_idx) - set_weight_attrs(w2_g_idx, extra_weight_attrs) - - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) - - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) - layer.a13_scale = None layer.a2_scale = None diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py index dcf52464b1..8f98cb4813 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py @@ -47,7 +47,6 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod): assert self.weight_quant.symmetric, ( "Only symmetric quantization is supported for W4A8 MoE" ) - assert self.weight_quant.actorder != "group" assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE" self.w4a8_backend, self.experts_cls = select_w4a8_moe_backend( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index ef8e6fe089..6a3f73e539 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -3,13 +3,10 @@ import math from fractions import Fraction -from typing import Any import torch from compressed_tensors.quantization import ( - ActivationOrdering, QuantizationArgs, - QuantizationStrategy, ) from vllm.logger import init_logger @@ -75,7 +72,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): self.group_size = ( -1 if weight_quant.group_size is None else weight_quant.group_size ) - self.actorder = weight_quant.actorder # Extract quant_type and create weight key for oracle selection self.quant_type = ( @@ -104,11 +100,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric) - is_actorder = self.strategy == QuantizationStrategy.GROUP and self.actorder in ( - ActivationOrdering.GROUP, - ActivationOrdering.DYNAMIC, - ) - # Select WNA16 MoE backend via oracle. self.wna16_backend, self.experts_cls = select_wna16_moe_backend( config=self.moe, @@ -116,7 +107,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): quant_config=self.weight_quant, may_have_zp=not self.symmetric, may_have_bias=False, - allow_tile_padding=not is_actorder, + allow_tile_padding=True, ) self.is_marlin = self.wna16_backend in [ @@ -130,14 +121,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): if self.is_marlin: assert check_moe_marlin_supports_config( - self.moe, self.group_size, allow_tile_padding=not is_actorder + self.moe, self.group_size, allow_tile_padding=True ) self.input_dtype = get_marlin_input_dtype(layer_name) else: # channelwise is not supported by this kernel assert weight_quant.strategy == "group" - # grouped actorder isn't supported by this kernel - assert weight_quant.actorder != "group" # Non-Marlin WNA16 always uses bf16/fp16 inputs self.input_dtype = torch.bfloat16 @@ -230,29 +219,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): assert shape[1] is not None and shape[2] is not None return shape[0], shape[1], shape[2] - @staticmethod - def _w2_scale_sharding( - actorder, - group_size: int, - intermediate_size_per_partition: int, - intermediate_size_full: int, - ) -> tuple[bool, int, bool]: - """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE. - - Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime - and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``. - ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at - quantization time, so scales shard normally per TP rank. - """ - load_full_w2 = (actorder == "group") and group_size != -1 - w2_scales_size = ( - intermediate_size_full if load_full_w2 else intermediate_size_per_partition - ) - is_k_full = (actorder != "group") or ( - intermediate_size_per_partition == intermediate_size_full - ) - return load_full_w2, w2_scales_size, is_k_full - def create_weights( self, layer: torch.nn.Module, @@ -262,8 +228,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): params_dtype: torch.dtype, **extra_weight_attrs, ): - intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full") - # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims @@ -301,13 +265,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): layer.register_parameter("w2_weight_packed", w2_weight) set_weight_attrs(w2_weight, extra_weight_attrs) - load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding( - self.actorder, - self.group_size, - intermediate_size_per_partition, - intermediate_size_full, - ) - if self.strategy == "channel": num_groups_w2 = num_groups_w13 = 1 self.group_size = -1 @@ -318,10 +275,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): f"({hidden_size}) to be divisible by group_size " f"({self.group_size})." ) - if ( - not load_full_w2 - and intermediate_size_per_partition % self.group_size != 0 - ): + if intermediate_size_per_partition % self.group_size != 0: raise ValueError( "CompressedTensors WNA16 MoE with static group " "scales requires the MoE intermediate size per " @@ -331,7 +285,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): "otherwise cross TP shard boundaries; use a compatible TP " "size or enable expert parallelism." ) - num_groups_w2 = w2_scales_size // self.group_size + num_groups_w2 = intermediate_size_per_partition // self.group_size num_groups_w13 = hidden_size // self.group_size if not self.symmetric: @@ -389,7 +343,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): ) layer.register_parameter("w2_weight_scale", w2_scale) set_weight_attrs(w2_scale, extra_weight_attrs) - set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) if not self.symmetric: w13_zp = torch.nn.Parameter( @@ -436,50 +389,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): layer.register_parameter("w13_weight_shape", w13_weight_shape) set_weight_attrs(w13_weight_shape, extra_weight_attrs) - w13_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_g_idx", w13_g_idx) - set_weight_attrs(w13_g_idx, extra_weight_attrs) - - w2_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_g_idx", w2_g_idx) - set_weight_attrs(w2_g_idx, extra_weight_attrs) - - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) - - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) - layer.a13_scale = None layer.a2_scale = None @@ -488,24 +397,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.moe_quant_config is not None - # Add Marlin-specific arguments - marlin_args: dict[str, Any] = {} - if self.is_marlin: - marlin_args = { - "w13_g_idx": layer.w13_weight_g_idx, - "w2_g_idx": layer.w2_weight_g_idx, - "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices, - "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices, - "is_k_full": self.is_k_full, - } - self.moe_kernel = make_wna16_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, backend=self.wna16_backend, routing_tables=layer._expert_routing_tables(), - **marlin_args, ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -519,8 +416,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): w2=layer.w2_weight_packed, w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, w13_qzeros=getattr(layer, "w13_weight_zero_point", None), w2_qzeros=getattr(layer, "w2_weight_zero_point", None), ) @@ -534,10 +429,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): w2_qweight, w13_scales, w2_scales, - w13_g_idx_processed, - w2_g_idx_processed, - w13_g_idx_sort_indices, - w2_g_idx_sort_indices, w13_qzeros, w2_qzeros, w13_input_global_scale, @@ -561,19 +452,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) - # Marlin-specific parameters (not needed for Flashinfer) if self.is_marlin: - if w13_g_idx_processed is not None: - replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) - if w2_g_idx_processed is not None: - replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) - if w13_g_idx_sort_indices is not None: - replace_parameter( - layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices - ) - if w2_g_idx_sort_indices is not None: - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) - # Register input global scales if present if w13_input_global_scale is not None: layer.register_parameter( @@ -593,7 +472,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): and self.wna16_backend != WNA16MoEBackend.EMULATION ): layer.workspace = marlin_make_workspace_new( - layer.w13_weight_g_idx.device, + layer.w13_weight_packed.device, 4, existing=getattr(layer, "workspace", None), ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py index 25cc7b115f..05db6ec3dd 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py @@ -65,15 +65,14 @@ class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: device = layer.w13_weight_packed.device num_experts = layer.w13_weight_packed.shape[0] - empty_g_idx = torch.empty(0, dtype=torch.int32, device=device) # Shuffle weights in-place per expert (exllama nibble interleave) for e in range(num_experts): w13_e = layer.w13_weight_packed.data[e].contiguous() - ops.gptq_shuffle(w13_e, empty_g_idx, 4) + ops.gptq_shuffle(w13_e, 4) layer.w13_weight_packed.data[e] = w13_e w2_e = layer.w2_weight_packed.data[e].contiguous() - ops.gptq_shuffle(w2_e, empty_g_idx, 4) + ops.gptq_shuffle(w2_e, 4) layer.w2_weight_packed.data[e] = w2_e # Keep scales as [E, groups, N] in activation dtype diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py index 2c121c6edf..925a211521 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py @@ -4,7 +4,6 @@ from collections.abc import Callable import torch -from compressed_tensors.quantization import ActivationOrdering from vllm.distributed.utils import verify_group_size_divides_partition from vllm.logger import init_logger @@ -43,13 +42,11 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme): num_bits: int, group_size: int | None = None, symmetric: bool | None = True, - actorder: ActivationOrdering | None = None, ): self.pack_factor = 32 // num_bits self.strategy = strategy self.symmetric = symmetric self.group_size = -1 if group_size is None else group_size - self.has_g_idx = actorder == ActivationOrdering.GROUP if self.group_size != 128 or self.strategy != "group": raise ValueError( @@ -92,7 +89,6 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme): act_type=torch.float8_e4m3fn, # always use fp8(e4m3) group_size=self.group_size, zero_points=not self.symmetric, - has_g_idx=self.has_g_idx, out_type=params_dtype, ) @@ -106,7 +102,7 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme): group_size = self.group_size if self.group_size != -1 else input_size row_parallel = input_size != input_size_per_partition partition_scales = not marlin_repeat_scales_on_all_ranks( - self.has_g_idx, self.group_size, row_parallel + self.group_size, row_parallel ) scales_and_zp_size = input_size // group_size @@ -162,7 +158,6 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme): w_q_param_name="weight_packed", w_s_param_name="weight_scale", w_zp_param_name="weight_zero_point", - w_gidx_param_name="weight_g_idx", ) # Checkpoints are serialized in compressed-tensors format, which is diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py index a26a35f4e3..2e1129fc1e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py @@ -98,7 +98,6 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme): act_type=params_dtype, group_size=effective_group_size, zero_points=False, - has_g_idx=False, ) kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config) @@ -140,7 +139,6 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme): w_q_param_name="weight_packed", w_s_param_name="weight_scale", w_zp_param_name=None, - w_gidx_param_name=None, ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py index 3aa8f85188..42688ab834 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py @@ -6,7 +6,6 @@ from collections.abc import Callable from fractions import Fraction import torch -from compressed_tensors.quantization import ActivationOrdering from vllm.distributed.utils import verify_group_size_divides_partition from vllm.logger import init_logger @@ -28,7 +27,6 @@ from vllm.model_executor.parameter import ( GroupQuantScaleParameter, PackedColumnParameter, PackedvLLMParameter, - RowvLLMParameter, ) from vllm.scalar_type import scalar_types @@ -57,7 +55,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): num_bits: int, group_size: int | None = None, symmetric: bool | None = True, - actorder: ActivationOrdering | None = None, layer_name: str | None = None, ): self.num_bits = num_bits @@ -65,7 +62,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): self.strategy = strategy self.symmetric = symmetric self.group_size = -1 if group_size is None else group_size - self.has_g_idx = actorder == ActivationOrdering.GROUP self.layer_name = layer_name if self.group_size == -1 and self.strategy != "channel": @@ -128,7 +124,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): act_type=params_dtype, group_size=self.group_size, zero_points=not self.symmetric, - has_g_idx=self.has_g_idx, ) kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config) @@ -146,7 +141,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): group_size = self.group_size if self.group_size != -1 else input_size row_parallel = input_size != input_size_per_partition partition_scales = not marlin_repeat_scales_on_all_ranks( - self.has_g_idx, self.group_size, row_parallel + self.group_size, row_parallel ) scales_and_zp_size = input_size // group_size @@ -226,24 +221,11 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): if not self.symmetric: layer.register_parameter("weight_zero_point", qzeros) - # group index (for activation reordering) - if self.has_g_idx: - weight_g_idx = RowvLLMParameter( - data=torch.empty( - input_size_per_partition, - dtype=torch.int32, - ), - input_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight_g_idx", weight_g_idx) - self.kernel = kernel_type( mp_linear_kernel_config, w_q_param_name="weight_packed", w_s_param_name="weight_scale", w_zp_param_name="weight_zero_point", - w_gidx_param_name="weight_g_idx", ) # Checkpoints are serialized in compressed-tensors format, which is diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py index 06715badda..427c92310f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py @@ -144,7 +144,6 @@ class CompressedTensorsWNA4Int(CompressedTensorsScheme): act_type=params_dtype, group_size=self.group_size, zero_points=False, - has_g_idx=False, ) kernel_type = choose_mp_linear_kernel(mp_config) @@ -166,7 +165,7 @@ class CompressedTensorsWNA4Int(CompressedTensorsScheme): group_size = self.group_size if self.group_size != -1 else input_size row_parallel = input_size != input_size_per_partition partition_scales = not marlin_repeat_scales_on_all_ranks( - False, self.group_size, row_parallel + self.group_size, row_parallel ) scales_size = input_size // group_size if partition_scales: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py index 4ca8a1c12c..a14077d899 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py @@ -144,7 +144,6 @@ class CompressedTensorsWNA8Int(CompressedTensorsScheme): act_type=params_dtype, group_size=self.group_size, zero_points=False, - has_g_idx=False, ) kernel_type = choose_mp_linear_kernel(mp_config) @@ -166,7 +165,7 @@ class CompressedTensorsWNA8Int(CompressedTensorsScheme): group_size = self.group_size if self.group_size != -1 else input_size row_parallel = input_size != input_size_per_partition partition_scales = not marlin_repeat_scales_on_all_ranks( - False, self.group_size, row_parallel + self.group_size, row_parallel ) scales_size = input_size // group_size if partition_scales: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py index 3ab02a895a..187b6438ab 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py @@ -114,7 +114,6 @@ class CompressedTensorsWNA8O8Int(CompressedTensorsScheme): act_type=params_dtype, # activation quant applied externally (SRQ) group_size=self.group_size, zero_points=False, - has_g_idx=False, ) self.kernel = choose_mp_linear_kernel(mp_config)( mp_config, @@ -167,7 +166,7 @@ class CompressedTensorsWNA8O8Int(CompressedTensorsScheme): # Scale: per-output-channel, or per group along the input dim under TP. group_size = self.group_size if self.group_size != -1 else input_size partitioned = not marlin_repeat_scales_on_all_ranks( - False, self.group_size, input_size != input_size_per_partition + self.group_size, input_size != input_size_per_partition ) scales = (input_size_per_partition if partitioned else input_size) // group_size scale_data = torch.empty(out, scales, dtype=params_dtype) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_w4a8_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_w4a8_linear.py index 01665d40ad..ceedcfb36b 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_w4a8_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_w4a8_linear.py @@ -106,7 +106,7 @@ class INCXPUW4A8LinearMethod(INCXPULinearMethod): layer.scales_fp16, layer.qzeros, self.group_size, - None, + None, # Retained by the external XPU op ABI. bias, ) return out.to(x.dtype).reshape(out_shape) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py index 936a155be6..c9cacef1c2 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py @@ -14,7 +14,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( from vllm.model_executor.parameter import ( GroupQuantScaleParameter, PackedvLLMParameter, - RowvLLMParameter, ) from vllm.scalar_type import scalar_types @@ -254,16 +253,6 @@ class INCXPULinearBase(INCLinearScheme): layer.register_parameter("scales", scales) layer.register_parameter("qzeros", qzeros) - g_idx = RowvLLMParameter( - data=torch.tensor( - [i // self.group_size for i in range(input_size_per_partition)], - dtype=torch.int32, - ), - input_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("g_idx", g_idx) - def _convert_awq_qweight_to_gptq(self, qw: torch.Tensor) -> torch.Tensor: """Convert AWQ qweight [K, N // pf] to GPTQ qweight [K // pf, N]. @@ -356,7 +345,7 @@ class INCXPULinearMethod(INCXPULinearBase): layer.scales, layer.qzeros, self.group_size, - None, + None, # Retained by the external XPU op ABI. ) return out.reshape(out_shape) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 1c4d123a38..430abdf705 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -32,6 +32,9 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) +from vllm.model_executor.layers.quantization.utils.gptq_utils import ( + normalize_and_validate_gptq_desc_act, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( INT4_DTYPE, INT8_DTYPE, @@ -69,7 +72,11 @@ class MoeWNA16Config(QuantizationConfig): from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig if self.linear_quant_method == "gptq": - pass + normalize_and_validate_gptq_desc_act( + full_config.get("desc_act", False), + group_size, + full_config.get("dynamic") or {}, + ) elif self.linear_quant_method in ("awq", "awq_marlin"): capability_tuple = current_platform.get_device_capability() device_capability = ( @@ -148,7 +155,16 @@ class MoeWNA16Config(QuantizationConfig): # Extract data from quant config. quant_method = quant_config.get("quant_method", "").lower() num_bits = quant_config.get("bits") - desc_act = quant_config.get("desc_act") + desc_act = quant_config.get("desc_act", False) + try: + normalize_and_validate_gptq_desc_act( + desc_act, + quant_config.get("group_size"), + quant_config.get("dynamic") or {}, + ) + supported_act_order = True + except ValueError: + supported_act_order = False capability_tuple = current_platform.get_device_capability() device_capability = ( @@ -159,7 +175,9 @@ class MoeWNA16Config(QuantizationConfig): awq_min_capability = AutoAWQConfig.get_min_capability() - gptq_compatible = quant_method == "gptq" and not desc_act and num_bits in [4, 8] + gptq_compatible = ( + quant_method == "gptq" and supported_act_order and num_bits in [4, 8] + ) awq_compatible = ( quant_method == "awq" and num_bits == 4 @@ -356,7 +374,7 @@ class MoeWNA16Method(FusedMoEMethodBase): if self.quant_config.linear_quant_method == "gptq": # some param are unused, but we need to init them in order to # load weights - invalid_param_keys = ["w13_g_idx", "w2_g_idx"] + invalid_param_keys: list[str] = [] if not self.quant_config.has_zp: invalid_param_keys += ["w13_qzeros", "w2_qzeros"] for key in invalid_param_keys: @@ -414,8 +432,6 @@ class MoeWNA16Method(FusedMoEMethodBase): w2=layer.w2_qweight, w13_scale=layer.w13_scales, w2_scale=layer.w2_scales, - w13_g_idx=None, - w2_g_idx=None, w13_qzeros=layer.w13_qzeros if has_zp else None, w2_qzeros=layer.w2_qzeros if has_zp else None, ) @@ -430,10 +446,6 @@ class MoeWNA16Method(FusedMoEMethodBase): w2_qweight, w13_scales, w2_scales, - _, - _, - _, - _, w13_qzeros, w2_qzeros, w13_input_global_scale, @@ -576,8 +588,6 @@ class MoeWNA16Method(FusedMoEMethodBase): expert_id: int, return_success: bool = False, ): - if "g_idx" in weight_name: - return False if return_success else None if not layer.quant_config.has_zp and "qzeros" in weight_name: return False if return_success else None diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index c73c42f4ff..d134abf8dc 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -20,6 +20,36 @@ else: AutoGPTQConfig = object +def normalize_and_validate_gptq_desc_act( + desc_act: bool, + group_size: int | None, + dynamic: Mapping[str, Any], +) -> bool: + base_desc_act = desc_act + if desc_act and group_size == -1: + # With one group per output channel, activation ordering is a no-op. + desc_act = False + if desc_act: + raise ValueError( + "GPTQ group activation ordering (desc_act=True) is no longer " + "supported. Use a checkpoint with static activation ordering " + "or desc_act=False." + ) + + for pattern, overrides in dynamic.items(): + if pattern.startswith("-:") or not isinstance(overrides, Mapping): + continue + dynamic_desc_act = overrides.get("desc_act", base_desc_act) + dynamic_group_size = overrides.get("group_size", group_size) + if dynamic_desc_act and dynamic_group_size != -1: + raise ValueError( + "GPTQ group activation ordering (desc_act=True) is no longer " + f"supported, but is enabled by dynamic rule {pattern!r}. Use " + "a checkpoint with static activation ordering or desc_act=False." + ) + return desc_act + + # Match dynamic rules with module name (prefix) and override quantize # config if module (prefix) matches a rule def override_config(config: AutoGPTQConfig, prefix: str): @@ -33,6 +63,10 @@ def override_config(config: AutoGPTQConfig, prefix: str): if isinstance(desc_act, bool): config.desc_act = desc_act + if config.desc_act and config.group_size == -1: + # Activation ordering is a no-op for channelwise quantization. + config.desc_act = False + config.pack_factor = 32 // config.weight_bits # packed into int32 assert isinstance(config, AutoGPTQConfig) is_sym = get_dynamic_override(config, prefix, "sym", config.is_sym) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 3658761d6f..22cf18f4a5 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -359,11 +359,11 @@ def check_moe_marlin_supports_config( ) -> bool: """Whether the fused MoE Marlin kernel supports ``config``. - Callers without act-order may pass ``allow_tile_padding=True``: a - tile-misaligned intermediate size is then zero-padded to a valid thread - tile at weight prep (see marlin_moe_padded_intermediate), so only a group - straddling the padded boundary stays unsupported. hidden_size is the MoE - I/O extent and is never padded. Act-order keeps the strict shape. + With ``allow_tile_padding=True``, a tile-misaligned intermediate size is + zero-padded to a valid thread tile at weight prep (see + marlin_moe_padded_intermediate), so only a group straddling the padded + boundary stays unsupported. hidden_size is the MoE I/O extent and is never + padded. """ if current_platform.is_rocm(): return False @@ -433,30 +433,17 @@ def marlin_make_workspace_new( return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False) -def marlin_is_k_full(act_order: bool, is_row_parallel: bool) -> bool: - return (not act_order) or (act_order and not is_row_parallel) - - -def marlin_repeat_scales_on_all_ranks( - act_order: bool, group_size: int, is_row_parallel: bool -) -> bool: - # Need to repeat scales on every rank if act_ordering or - # channelwise and RowParallelLinear +def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool: is_channelwise = group_size == -1 - return act_order or (is_channelwise and is_row_parallel) + return is_channelwise and is_row_parallel -def marlin_make_empty_g_idx(device: torch.device) -> torch.Tensor: +def marlin_make_empty(device: torch.device) -> torch.Tensor: return torch.nn.Parameter( torch.empty(0, dtype=torch.int, device=device), requires_grad=False ) -def marlin_sort_g_idx(g_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - g_idx_sort_indices = torch.argsort(g_idx).to(torch.int) - return g_idx[g_idx_sort_indices], g_idx_sort_indices - - def get_scale_perms(): scale_perm: list[int] = [] for i in range(8): @@ -700,13 +687,10 @@ def apply_gptq_marlin_linear( weight: torch.Tensor, weight_scale: torch.Tensor, weight_zp: torch.Tensor, - g_idx: torch.Tensor, - g_idx_sort_indices: torch.Tensor, workspace: torch.Tensor, wtype: ScalarType, output_size_per_partition: int, input_size_per_partition: int, - is_k_full: bool, input_global_scale: torch.Tensor | None = None, bias: torch.Tensor | None = None, use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, @@ -749,14 +733,11 @@ def apply_gptq_marlin_linear( a_scales, None, weight_zp, - g_idx, - g_idx_sort_indices, workspace, wtype, size_m=reshaped_x.shape[0], size_n=padded_n, size_k=padded_k, - is_k_full=is_k_full, use_atomic_add=use_atomic_add, use_fp32_reduce=use_fp32_reduce, is_zp_float=False, diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 01d28e5e91..3b6990101e 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -203,8 +203,6 @@ def apply_fp4_marlin_linear( a_scales=a_scales, global_scale=weight_global_scale, b_zeros=None, - g_idx=None, - perm=None, workspace=workspace, b_q_type=scalar_types.float4_e2m1f, size_m=reshaped_x.size(0), @@ -246,14 +244,12 @@ def prepare_fp4_layer_for_marlin( # WEIGHT # Repack weights to marlin format - perm = torch.empty(0, dtype=torch.int, device=device) qweight = layer.weight.view(torch.int32).T.contiguous() qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k) is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, - perm=perm, size_k=padded_k, size_n=padded_n, num_bits=4, @@ -312,7 +308,6 @@ def _repack_marlin_experts( weight: torch.Tensor, size_n: int, size_k: int, - perm: torch.Tensor, is_a_8bit: bool, ) -> torch.Tensor: """Repack each expert to marlin format into a preallocated output.""" @@ -322,7 +317,6 @@ def _repack_marlin_experts( qweight = weight[i].view(torch.int32).T.contiguous() marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, - perm=perm, size_k=size_k, size_n=size_n, num_bits=4, @@ -401,7 +395,6 @@ def prepare_nvfp4_moe_layer_for_marlin( layer.workspace = marlin_make_workspace_new( device, 4, existing=getattr(layer, "workspace", None) ) - perm = torch.empty(0, dtype=torch.int, device=device) # WEIGHT # Repack weights to marlin format @@ -417,7 +410,7 @@ def prepare_nvfp4_moe_layer_for_marlin( weight = pad_w2(weight, packing=2) size_k = padded_N - return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit) + return _repack_marlin_experts(weight, size_n, size_k, is_a_8bit) w13 = repack_weight(w13, "w13") w2 = repack_weight(w2, "w2") @@ -492,7 +485,6 @@ def prepare_moe_fp4_layer_for_marlin( layer.workspace = marlin_make_workspace_new( device, 4, existing=getattr(layer, "workspace", None) ) - perm = torch.empty(0, dtype=torch.int, device=device) is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 # WEIGHT @@ -506,7 +498,7 @@ def prepare_moe_fp4_layer_for_marlin( assert weight.shape == (e, size_n, size_k // 2) - weight = _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit) + weight = _repack_marlin_experts(weight, size_n, size_k, is_a_8bit) weight = torch.nn.Parameter(weight, requires_grad=False) setattr(layer, name, weight) @@ -621,10 +613,8 @@ def prepare_moe_mxfp4_layer_for_marlin( n = w13.shape[1] // 2 # intermediate_size_per_partition k = w13.shape[2] * 2 # hidden_size - device = w13.device param_dtype = layer.params_dtype is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 - perm = torch.empty(0, dtype=torch.int, device=device) # WEIGHT: Repack weights to marlin format def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: @@ -635,7 +625,7 @@ def prepare_moe_mxfp4_layer_for_marlin( assert weight.shape == (e, size_n, size_k // 2) - return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit) + return _repack_marlin_experts(weight, size_n, size_k, is_a_8bit) w13 = repack_weight(w13, "w13") w2 = repack_weight(w2, "w2") @@ -691,8 +681,6 @@ def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None): assert not is_a_8bit, "NVFP4 weight + INT8/FP8 activation is not supported." assert group_size > 0 size_n, size_k = weight.shape - device = weight.device - scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 6 global_scale = scales.max() / 448 scales = (scales / global_scale).to(torch.float8_e4m3fn) @@ -720,7 +708,6 @@ def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None): marlin_qweight = ops.gptq_marlin_repack( b_q_weight=fp4_weight.view(torch.int32).T.contiguous(), - perm=torch.empty(0, dtype=torch.int, device=device), size_k=size_k, size_n=size_n, num_bits=4, @@ -751,8 +738,6 @@ def rand_marlin_weight_mxfp4_like(weight, group_size, input_dtype=None): assert group_size > 0 size_n, size_k = weight.shape - device = weight.device - scales = torch.randint( 110, 120, @@ -779,11 +764,9 @@ def rand_marlin_weight_mxfp4_like(weight, group_size, input_dtype=None): ).view(size_n, size_k) weight_ref = weight_ref * scales.repeat_interleave(group_size, 1).to(weight.dtype) - perm = torch.empty(0, dtype=torch.int, device=device) fp4_weight = fp4_weight.view(torch.int32).T.contiguous() marlin_qweight = ops.gptq_marlin_repack( b_q_weight=fp4_weight, - perm=perm, size_k=size_k, size_n=size_n, num_bits=4, diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index aaa1d82090..56841cba17 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -89,8 +89,6 @@ def apply_fp8_marlin_linear( a_scales=a_scales, global_scale=None, b_zeros=None, - g_idx=None, - perm=None, workspace=workspace, b_q_type=scalar_types.float8_e4m3fn, size_m=reshaped_x.size(0), @@ -138,7 +136,6 @@ def prepare_fp8_layer_for_marlin( # WEIGHT # Repack weights to marlin format - perm = torch.empty(0, dtype=torch.int, device=device) qweight = pack_fp8_to_int32(layer.weight, size_k_first) if not size_k_first: qweight = qweight.T.contiguous() @@ -146,7 +143,6 @@ def prepare_fp8_layer_for_marlin( marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, - perm=perm, size_k=padded_k, size_n=padded_n, num_bits=8, @@ -283,7 +279,6 @@ def prepare_fp8_moe_layer_for_marlin( layer.workspace = marlin_make_workspace_new( device, 4, existing=getattr(layer, "workspace", None) ) - perm = torch.empty(0, dtype=torch.int, device=device) # WEIGHT # Repack weights to marlin format @@ -295,7 +290,7 @@ def prepare_fp8_moe_layer_for_marlin( qweight = qweight.T.contiguous() marlin_qweight = ops.gptq_marlin_repack( - b_q_weight=qweight, perm=perm, size_k=size_k, size_n=size_n, num_bits=8 + b_q_weight=qweight, size_k=size_k, size_n=size_n, num_bits=8 ) tensor_list.append(marlin_qweight) @@ -434,8 +429,6 @@ def apply_mxfp8_marlin_linear( a_scales=None, global_scale=None, b_zeros=None, - g_idx=None, - perm=None, workspace=workspace, b_q_type=scalar_types.float8_e4m3fn, size_m=reshaped_x.size(0), @@ -470,14 +463,12 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: ) # WEIGHT - repack FP8 weights to Marlin format - perm = torch.empty(0, dtype=torch.int, device=device) qweight = pack_fp8_to_int32(layer.weight, size_k_first=False) qweight = qweight.T.contiguous() qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k) marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, - perm=perm, size_k=padded_k, size_n=padded_n, num_bits=8, @@ -551,7 +542,6 @@ def prepare_mxfp8_moe_layer_for_marlin( device = w13.device param_dtype = torch.get_default_dtype() - perm = torch.empty(0, dtype=torch.int, device=device) layer.workspace = marlin_make_workspace_new( device, 4, existing=getattr(layer, "workspace", None) @@ -571,7 +561,6 @@ def prepare_mxfp8_moe_layer_for_marlin( qweight = qweight.T.contiguous() marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, - perm=perm, size_k=size_k, size_n=size_n, num_bits=8, @@ -615,8 +604,6 @@ def marlin_quant_fp8_torch(weight, group_size, input_dtype=None): assert input_dtype == torch.float8_e4m3fn size_n, size_k = weight.shape - device = weight.device - if group_size != -1: scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 448 repeated_scales = scales.repeat_interleave(group_size, 1) @@ -629,10 +616,8 @@ def marlin_quant_fp8_torch(weight, group_size, input_dtype=None): weight_ref = fp8_weight.to(weight.dtype) * repeated_scales packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous() - perm = torch.empty(0, dtype=torch.int, device=device) marlin_qweight = ops.gptq_marlin_repack( b_q_weight=packed_weight, - perm=perm, size_k=size_k, size_n=size_n, num_bits=8, diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py index 9162afe03d..d4a7c563b8 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py @@ -13,7 +13,6 @@ from .quant_utils import ( get_pack_factor, gptq_quantize_weights, quantize_weights, - sort_weights, ) @@ -129,8 +128,6 @@ def marlin_quantize( w: torch.Tensor, quant_type: ScalarType, group_size: int, - act_order: bool, - test_perm: torch.Tensor | None = None, input_dtype: torch.dtype | None = None, ): is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 @@ -143,16 +140,7 @@ def marlin_quantize( group_size = size_k assert group_size <= size_k - # Quantize (and apply act_order if provided) - w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights( - w, quant_type, group_size, act_order, test_perm - ) - - # For act_order, sort the "weights" and "g_idx" so that group ids are - # increasing - sort_indices = torch.empty(0, dtype=torch.int, device=w.device) - if act_order: - q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) + w_ref, q_w, s = gptq_quantize_weights(w, quant_type, group_size) # Reformat to marlin weight_perm = get_weight_perm(num_bits, is_a_8bit) @@ -166,7 +154,7 @@ def marlin_quantize( marlin_s = marlin_s * 512 # Create result - res_list = [w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm] + res_list = [w_ref, marlin_q_w, marlin_s] for i in range(len(res_list)): res_list[i] = res_list[i].to(w.device) diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index c200fbbf78..f96146e736 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -691,36 +691,6 @@ def get_pack_factor(num_bits): return 32 // num_bits -def permute_rows( - q_w: torch.Tensor, - w_ref: torch.Tensor, - group_size: int, - test_perm: torch.Tensor | None = None, -): - assert q_w.shape == w_ref.shape - - orig_device = q_w.device - k_size, _ = q_w.shape - - g_idx = torch.zeros((k_size,), dtype=torch.int32) - for i in range(k_size): - g_idx[i] = i // group_size - - # Simulate act_order by doing a random permutation on K - rand_perm = test_perm if test_perm is not None else torch.randperm(k_size) - - g_idx = g_idx[rand_perm].contiguous() - q_w = q_w[rand_perm, :].contiguous() - w_ref = w_ref[rand_perm, :].contiguous() - - return ( - w_ref.to(device=orig_device), - q_w.to(device=orig_device), - g_idx.to(device=orig_device), - rand_perm.to(device=orig_device), - ) - - def quantize_weights( w: torch.Tensor, quant_type: ScalarType, @@ -824,8 +794,6 @@ def gptq_quantize_weights( w: torch.Tensor, quant_type: ScalarType, group_size: int, - act_order: bool, - test_perm: torch.Tensor | None = None, ): size_k, _ = w.shape @@ -838,35 +806,7 @@ def gptq_quantize_weights( ) w_ref, w_q, w_s, _ = quantize_weights(w, quant_type, group_size) - - # Apply act_order - g_idx = torch.empty(0, dtype=torch.int, device=w.device) - rand_perm = torch.empty(0, dtype=torch.int, device=w.device) - if act_order: - assert group_size < size_k, ( - "For act_order, groupsize = {} must be less than size_k = {}".format( - group_size, size_k - ) - ) - - w_ref, w_q, g_idx, rand_perm = permute_rows(w_q, w_ref, group_size, test_perm) - - return w_ref, w_q, w_s, g_idx, rand_perm - - -def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor): - orig_device = q_w.device - - sort_indices = torch.argsort(g_idx).to(dtype=torch.int32) # Sort based on g_idx - - g_idx = g_idx[sort_indices].contiguous() - q_w = q_w[sort_indices, :].contiguous() - - return ( - q_w.to(device=orig_device), - g_idx.to(device=orig_device), - sort_indices.to(device=orig_device), - ) + return w_ref, w_q, w_s def pack_rows( diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index ba901550e8..66cbc65051 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -463,8 +463,7 @@ class VocabParallelEmbedding(PluggableLayer): output_dim = getattr(param, "output_dim", None) packed_dim = getattr(param, "packed_dim", None) - # If parameter does not have output dim, then it should - # be copied onto all gpus (e.g. g_idx for act_order gptq). + # Parameters without an output dimension are copied onto all GPUs. if output_dim is None: if ( loaded_weight.ndim == 0 diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index f6a97961f4..3085df85c6 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1543,8 +1543,6 @@ def maybe_remap_moe_expert_param_name( "w2_bias", "w13_scale", "w2_scale", - "w13_g_idx", - "w2_g_idx", "w13_qweight", "w2_qweight", "w13_qzeros", diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 84bc8a2d4b..6cd21763a0 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -432,8 +432,9 @@ class AutoWeightsLoader: modules = (self.module, *self.module.children()) iterator = (m.quant_config for m in modules if hasattr(m, "quant_config")) if quant_config := next(iterator, None): - # Get mappings and ignore prefixes for KV cache quantization scales + # Apply mappings for quantization-specific checkpoint tensors. mapper |= quant_config.get_cache_scale_mapper() + mapper |= quant_config.get_checkpoint_weight_mapper() ignore_unexpected_suffixes = quant_config._ignore_unexpected_suffixes self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes) mapper |= self.REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER