[Quant][Kernel] Remove GPTQ Group/Dynamic Activation Ordering (#54809)
Signed-off-by: Roderick-Wu <roderickwu2003@gmail.com> Signed-off-by: Roderick Wu <roderickwu2003@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Roderick Wu committed
Sep 8, 2026 at 15:45 UTC
2c9d68f80bd20af7af1bc539c044a4aebbe52f8e
99 files changed
+1183
-3918
CMakeLists.txt
-1
@@ -436,7 +436,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
436
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
437
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
438
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
439
- "csrc/libtorch_stable/permute_cols.cu"
439
"csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
440
"csrc/libtorch_stable/pos_encoding_kernels.cu"
441
"csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
benchmarks/kernels/benchmark_machete.py
+1
-6
@@ -220,10 +220,8 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
220
bt.w_g_s, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.group_size
221
)
222
223
- sort_indices = torch.empty(0, dtype=torch.int, device=device)
224
- g_idx = torch.empty(0, dtype=torch.int, device=device)
223
w_q = ops.gptq_marlin_repack(
226
- bt.w_q, sort_indices, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits
224
+ bt.w_q, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits
225
)
226
227
if bt.a.dtype.is_floating_point:
@@ -240,14 +238,11 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
238
a_scales=None,
239
global_scale=None,
240
b_zeros=w_zp,
243
- g_idx=g_idx,
244
- perm=sort_indices,
241
workspace=workspace.scratch,
242
b_q_type=bt.wtype,
243
size_m=bt.a.shape[0],
244
size_n=bt.w_ref.shape[1],
245
size_k=bt.w_ref.shape[0],
250
- is_k_full=True,
246
is_zp_float=False,
247
)
248
else:
benchmarks/kernels/benchmark_marlin.py
+33
-85
@@ -32,7 +32,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
32
gptq_pack,
33
gptq_quantize_weights,
34
quantize_weights,
35
- sort_weights,
35
)
36
from vllm.scalar_type import ScalarType, scalar_types
37
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -40,15 +39,10 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
39
DEFAULT_MODELS = ["meta-llama/Llama-2-7b-hf/TP1"]
40
DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
41
43
-ACT_ORDER_OPTS = [False, True]
44
-K_FULL_OPTS = [False, True]
45
-
42
43
def bench_run(
44
results: list[benchmark.Measurement],
45
model: str,
50
- act_order: bool,
51
- is_k_full: bool,
46
quant_type: ScalarType,
47
group_size: int,
48
size_m: int,
@@ -56,38 +50,33 @@ def bench_run(
50
size_n: int,
51
):
52
label = "Quant Matmul"
59
- sub_label = "{}, act={} k_full={}, q={}, g={}, MKN=({}x{}x{})".format(
60
- model, act_order, is_k_full, str(quant_type), group_size, size_m, size_k, size_n
53
+ sub_label = "{}, q={}, g={}, MKN=({}x{}x{})".format(
54
+ model, str(quant_type), group_size, size_m, size_k, size_n
55
)
56
print(f"Testing: {sub_label}")
57
58
a = torch.randn(size_m, size_k).to(torch.half).cuda()
59
b = torch.rand(size_k, size_n).to(torch.half).cuda()
60
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
67
- if act_order and (group_size == -1 or group_size == size_k or has_zp):
68
- return
61
if size_k % group_size != 0:
62
return
63
64
repack_supported = group_size in MARLIN_SUPPORTED_GROUP_SIZES
65
allspark_supported = (
74
- quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES
75
- and group_size == -1
76
- and not act_order
77
- and is_k_full
66
+ quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES and group_size == -1
67
)
68
69
def gen_marlin_params():
70
# Marlin quant
82
- marlin_g_idx = marlin_sort_indices = marlin_zp = marlin_s2 = None
71
+ marlin_zp = marlin_s2 = None
72
if quant_type == scalar_types.float4_e2m1f:
84
- if group_size != 16 or act_order:
73
+ if group_size != 16:
74
return
75
marlin_w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_fp4_like(
76
b.T, group_size
77
)
78
elif quant_type == scalar_types.float8_e4m3fn:
90
- if group_size not in [-1, 128] or act_order:
79
+ if group_size not in [-1, 128]:
80
return
81
marlin_w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(b.T, group_size)
82
elif group_size == 16:
@@ -97,8 +86,8 @@ def bench_run(
86
b, quant_type, group_size
87
)
88
else:
100
- marlin_w_ref, marlin_q_w, marlin_s, marlin_g_idx, marlin_sort_indices, _ = (
101
- marlin_quantize(b, quant_type, group_size, act_order)
89
+ marlin_w_ref, marlin_q_w, marlin_s = marlin_quantize(
90
+ b, quant_type, group_size
91
)
92
return (
93
marlin_w_ref,
@@ -106,25 +95,14 @@ def bench_run(
95
marlin_s,
96
marlin_s2,
97
marlin_zp,
109
- marlin_g_idx,
110
- marlin_sort_indices,
98
)
99
100
def gen_repack_params():
101
q_w_gptq = None
115
- repack_sort_indices = None
102
if repack_supported:
117
- (w_ref, q_w, s, g_idx, rand_perm) = gptq_quantize_weights(
118
- b, quant_type, group_size, act_order
119
- )
103
+ _, q_w, _ = gptq_quantize_weights(b, quant_type, group_size)
104
q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
121
-
122
- # For act_order, sort the "weights" and "g_idx"
123
- # so that group ids are increasing
124
- repack_sort_indices = torch.empty(0, dtype=torch.int, device=b.device)
125
- if act_order:
126
- (q_w, g_idx, repack_sort_indices) = sort_weights(q_w, g_idx)
127
- return q_w_gptq, repack_sort_indices
105
+ return q_w_gptq
106
107
def gen_allspark_params():
108
qw_reorder = s_reorder = zp_reorder = sm_count = sm_version = (
@@ -161,10 +139,8 @@ def bench_run(
139
marlin_s,
140
marlin_s2,
141
marlin_zp,
164
- marlin_g_idx,
165
- marlin_sort_indices,
142
) = gen_marlin_params()
167
- q_w_gptq, repack_sort_indices = gen_repack_params()
143
+ q_w_gptq = gen_repack_params()
144
qw_reorder, s_reorder, zp_reorder, sm_count, sm_version, CUBLAS_M_THRESHOLD = (
145
gen_allspark_params()
146
)
@@ -188,13 +164,9 @@ def bench_run(
164
"marlin_s": marlin_s,
165
"marlin_s2": marlin_s2,
166
"marlin_zp": marlin_zp,
191
- "marlin_g_idx": marlin_g_idx,
192
- "marlin_sort_indices": marlin_sort_indices,
167
"marlin_workspace": marlin_workspace,
194
- "is_k_full": is_k_full,
168
# GPTQ params
169
"q_w_gptq": q_w_gptq,
197
- "repack_sort_indices": repack_sort_indices,
170
# AllSpark W8A16 params
171
"qw_reorder": qw_reorder,
172
"s_reorder": s_reorder,
@@ -226,7 +198,7 @@ def bench_run(
198
199
results.append(
200
benchmark.Timer(
229
- 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
201
+ 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
202
globals=globals,
203
label=label,
204
sub_label=sub_label,
@@ -236,7 +208,7 @@ def bench_run(
208
209
results.append(
210
benchmark.Timer(
239
- 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
211
+ 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
212
globals=globals,
213
label=label,
214
sub_label=sub_label,
@@ -247,7 +219,7 @@ def bench_run(
219
if repack_supported:
220
results.append(
221
benchmark.Timer(
250
- stmt="q_res = gptq_marlin_repack(q_w_gptq, repack_sort_indices, size_k, size_n, quant_type.size_bits)", # noqa: E501
222
+ stmt="q_res = gptq_marlin_repack(q_w_gptq, size_k, size_n, quant_type.size_bits)", # noqa: E501
223
globals=globals,
224
label=label,
225
sub_label=sub_label,
@@ -284,61 +256,39 @@ def main(args):
256
if len(args.limit_n) > 0 and size_n not in args.limit_n:
257
continue
258
287
- for act_order in ACT_ORDER_OPTS:
259
+ for quant_type in query_marlin_supported_quant_types():
260
if (
289
- len(args.limit_act_order) > 0
290
- and act_order not in args.limit_act_order
261
+ len(args.limit_num_bits) > 0
262
+ and quant_type.size_bits not in args.limit_num_bits
263
):
264
continue
265
294
- for is_k_full in K_FULL_OPTS:
266
+ for group_size in (
267
+ MARLIN_SUPPORTED_GROUP_SIZES + FP4_MARLIN_SUPPORTED_GROUP_SIZES
268
+ ):
269
if (
296
- len(args.limit_k_full) > 0
297
- and is_k_full not in args.limit_k_full
270
+ len(args.limit_group_size) > 0
271
+ and group_size not in args.limit_group_size
272
):
273
continue
274
301
- for quant_type in query_marlin_supported_quant_types():
302
- if (
303
- len(args.limit_num_bits) > 0
304
- and quant_type.size_bits not in args.limit_num_bits
305
- ):
306
- continue
307
-
308
- for group_size in (
309
- MARLIN_SUPPORTED_GROUP_SIZES
310
- + FP4_MARLIN_SUPPORTED_GROUP_SIZES
311
- ):
312
- if (
313
- len(args.limit_group_size) > 0
314
- and group_size not in args.limit_group_size
315
- ):
316
- continue
317
-
318
- # For act_order, the group_size must be less than
319
- # size_k
320
- if act_order and (group_size == size_k or group_size == -1):
321
- continue
322
-
323
- for size_m in args.batch_sizes:
324
- bench_run(
325
- results,
326
- model,
327
- act_order,
328
- is_k_full,
329
- quant_type,
330
- group_size,
331
- size_m,
332
- size_k,
333
- size_n,
334
- )
275
+ for size_m in args.batch_sizes:
276
+ bench_run(
277
+ results,
278
+ model,
279
+ quant_type,
280
+ group_size,
281
+ size_m,
282
+ size_k,
283
+ size_n,
284
+ )
285
286
compare = benchmark.Compare(results)
287
compare.print()
288
289
290
# For quick benchmarking use:
341
-# 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
291
+# 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
292
#
293
if __name__ == "__main__":
294
parser = FlexibleArgumentParser(
@@ -358,8 +308,6 @@ if __name__ == "__main__":
308
parser.add_argument("--limit-n", nargs="+", type=int, default=[])
309
parser.add_argument("--limit-group-size", nargs="+", type=int, default=[])
310
parser.add_argument("--limit-num-bits", nargs="+", type=int, default=[])
361
- parser.add_argument("--limit-act-order", nargs="+", type=int, default=[])
362
- parser.add_argument("--limit-k-full", nargs="+", type=int, default=[])
311
312
args = parser.parse_args()
313
main(args)
csrc/cpu/cpu_wna16.cpp
+47
-105
@@ -37,7 +37,7 @@ namespace {
37
using cpu_utils::ISA;
38
using cpu_utils::VecTypeTrait;
39
40
-template <typename scalar_t, ISA isa, bool has_zp, bool use_desc_act>
40
+template <typename scalar_t, ISA isa, bool has_zp>
41
class Dequantizer4b {
42
public:
43
constexpr static int32_t pack_num = 32 / 4;
@@ -47,9 +47,9 @@ class Dequantizer4b {
47
static void dequant(int32_t* __restrict__ q_weight,
48
scalar_t* __restrict__ weight,
49
scalar_t* __restrict__ scales,
50
- int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
51
- const int64_t scales_stride, const int64_t zeros_stride,
52
- const int32_t k_size, const int32_t group_size) {
50
+ int32_t* __restrict__ zeros, const int64_t scales_stride,
51
+ const int64_t zeros_stride, const int32_t k_size,
52
+ const int32_t group_size) {
53
vec_op::FP32Vec16 lut;
54
if constexpr (has_zp) {
55
// AWQ
@@ -81,30 +81,15 @@ class Dequantizer4b {
81
vec_op::FP32Vec16 wb_0(qwb_0, lut);
82
vec_op::FP32Vec16 wb_1(qwb_1, lut);
83
84
- if constexpr (!use_desc_act) {
85
- if (group_counter == 0) {
86
- scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale));
87
- scale_1 = vec_op::FP32Vec16(scale_0);
88
- curr_scale += scales_stride;
84
+ if (group_counter == 0) {
85
+ scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale));
86
+ scale_1 = vec_op::FP32Vec16(scale_0);
87
+ curr_scale += scales_stride;
88
90
- if constexpr (has_zp) {
91
- zero_0 = vec_op::FP32Vec16(*curr_zeros, lut);
92
- zero_1 = vec_op::FP32Vec16(zero_0);
93
- curr_zeros += zeros_stride / 2;
94
- }
95
- }
96
- } else {
97
- int32_t g_idx_0 = g_idx[k_idx];
98
- int32_t g_idx_1 = g_idx[k_idx + 1];
99
- scale_0 = vec_op::FP32Vec16(
100
- scalar_vec_t(curr_scale + g_idx_0 * scales_stride));
101
- scale_1 = vec_op::FP32Vec16(
102
- scalar_vec_t(curr_scale + g_idx_1 * scales_stride));
89
if constexpr (has_zp) {
104
- zero_0 = vec_op::FP32Vec16(*(curr_zeros + g_idx_0 * zeros_stride / 2),
105
- lut);
106
- zero_1 = vec_op::FP32Vec16(*(curr_zeros + g_idx_1 * zeros_stride / 2),
107
- lut);
90
+ zero_0 = vec_op::FP32Vec16(*curr_zeros, lut);
91
+ zero_1 = vec_op::FP32Vec16(zero_0);
92
+ curr_zeros += zeros_stride / 2;
93
}
94
}
95
@@ -130,11 +115,9 @@ class Dequantizer4b {
115
// update
116
curr_q_weight += 2;
117
curr_weight += 32;
133
- if constexpr (!use_desc_act) {
134
- group_counter += 2;
135
- if (group_counter == group_size) {
136
- group_counter = 0;
137
- }
118
+ group_counter += 2;
119
+ if (group_counter == group_size) {
120
+ group_counter = 0;
121
}
122
}
123
}
@@ -145,11 +128,10 @@ template <typename scalar_t, typename dequantizer_t, typename gemm_t>
128
void cpu_gemm_wna16_impl(
129
scalar_t* __restrict__ input, int32_t* __restrict__ q_weight,
130
scalar_t* __restrict__ output, scalar_t* __restrict__ scales,
148
- int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
149
- scalar_t* __restrict__ bias, const int32_t m_size, const int32_t n_size,
150
- const int32_t k_size, const int64_t input_stride,
151
- const int64_t output_stride, const int64_t scales_group_stride,
152
- const int64_t zeros_group_stride, const int32_t group_num,
131
+ int32_t* __restrict__ zeros, scalar_t* __restrict__ bias,
132
+ const int32_t m_size, const int32_t n_size, const int32_t k_size,
133
+ const int64_t input_stride, const int64_t output_stride,
134
+ const int64_t scales_group_stride, const int64_t zeros_group_stride,
135
const int32_t group_size, const int64_t pack_factor) {
136
constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
137
constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
@@ -229,7 +211,7 @@ void cpu_gemm_wna16_impl(
211
int32_t* __restrict__ curr_zeros = zeros + n_start_idx / pack_factor;
212
for (int32_t block_idx = 0; block_idx < n_block_num; ++block_idx) {
213
dequantizer_t::dequant(curr_q_weight, curr_b_buffer, curr_scales,
232
- curr_zeros, g_idx, scales_group_stride,
214
+ curr_zeros, scales_group_stride,
215
zeros_group_stride, k_size, group_size);
216
217
// if (block_idx == 0 && n_start_idx == 0) {
@@ -296,8 +278,7 @@ void cpu_gemm_wna16(
278
const torch::Tensor& scales, // [group_num, N]
279
const std::optional<torch::Tensor>&
280
zeros, // [group_num, N / pack_factor], packed as int32
299
- const std::optional<torch::Tensor>& g_idx, // [K]
300
- const std::optional<torch::Tensor>& bias, // [N]
281
+ const std::optional<torch::Tensor>& bias, // [N]
282
const int64_t pack_factor, const std::string& isa_hint) {
283
using cpu_utils::ISA;
284
TORCH_CHECK_EQ(pack_factor, 8); // only supports 4bits
@@ -314,8 +295,6 @@ void cpu_gemm_wna16(
295
const int64_t output_m_stride = output.stride(0);
296
297
bool has_zp = zeros.has_value();
317
- bool use_desc_act = g_idx.has_value();
318
- TORCH_CHECK(!(has_zp && use_desc_act));
298
299
ISA isa = [&]() {
300
if (isa_hint == "amx") {
@@ -331,109 +310,72 @@ void cpu_gemm_wna16(
310
311
int32_t* zeros_ptr = has_zp ? zeros->data_ptr<int32_t>() : nullptr;
312
const int64_t zeros_group_stride = has_zp ? zeros->stride(0) : 0;
334
- int32_t* g_idx_ptr = use_desc_act ? g_idx->data_ptr<int32_t>() : nullptr;
313
314
VLLM_DISPATCH_16B_TYPES(input.scalar_type(), "cpu_gemm_wna16", [&]() {
315
if (isa == ISA::AMX) {
316
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::AMX, scalar_t>;
317
if (has_zp) {
340
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, true, false>;
318
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, true>;
319
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
320
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
321
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
344
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
345
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
346
- scales_group_stride, zeros_group_stride, group_num, group_size,
347
- pack_factor);
322
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
323
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
324
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
325
return;
326
}
350
- if (use_desc_act) {
351
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, true>;
352
- cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
353
- input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
354
- output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
355
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
356
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
357
- scales_group_stride, zeros_group_stride, group_num, group_size,
358
- pack_factor);
359
- return;
360
- } else {
361
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, false>;
327
+ {
328
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false>;
329
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
330
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
331
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
365
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
366
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
367
- scales_group_stride, zeros_group_stride, group_num, group_size,
368
- pack_factor);
332
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
333
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
334
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
335
return;
336
}
337
} else if (isa == ISA::VEC) {
338
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::VEC, scalar_t>;
339
if (has_zp) {
374
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, true, false>;
340
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, true>;
341
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
342
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
343
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
378
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
379
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
380
- scales_group_stride, zeros_group_stride, group_num, group_size,
381
- pack_factor);
344
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
345
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
346
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
347
return;
348
}
384
- if (use_desc_act) {
385
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, true>;
386
- cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
387
- input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
388
- output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
389
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
390
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
391
- scales_group_stride, zeros_group_stride, group_num, group_size,
392
- pack_factor);
393
- return;
394
- } else {
395
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, false>;
349
+ {
350
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false>;
351
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
352
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
353
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
399
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
400
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
401
- scales_group_stride, zeros_group_stride, group_num, group_size,
402
- pack_factor);
354
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
355
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
356
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
357
return;
358
}
359
} else if (isa == ISA::RVV) {
360
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
361
if (has_zp) {
408
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
362
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true>;
363
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
364
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
365
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
412
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
413
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
414
- scales_group_stride, zeros_group_stride, group_num, group_size,
415
- pack_factor);
366
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
367
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
368
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
369
return;
370
}
418
- if (use_desc_act) {
419
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
420
- cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
421
- input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
422
- output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
423
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
424
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
425
- scales_group_stride, zeros_group_stride, group_num, group_size,
426
- pack_factor);
427
- return;
428
- } else {
429
- using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
371
+ {
372
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false>;
373
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
374
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
375
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
433
- g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
434
- a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
435
- scales_group_stride, zeros_group_stride, group_num, group_size,
436
- pack_factor);
376
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
377
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
378
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
379
return;
380
}
381
}
csrc/cpu/torch_bindings.cpp
+1
-2
@@ -226,7 +226,6 @@ void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
226
void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight,
227
torch::Tensor& output, const torch::Tensor& scales,
228
const std::optional<torch::Tensor>& zeros,
229
- const std::optional<torch::Tensor>& g_idx,
229
const std::optional<torch::Tensor>& bias,
230
const int64_t pack_factor, const std::string& isa_hint);
231
@@ -671,7 +670,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
670
#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__)
671
ops.def(
672
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
674
- "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
673
+ "Tensor scales, Tensor? zeros, Tensor? bias, SymInt "
674
"pack_factor, str isa_hint) -> ()");
675
ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16);
676
#endif
csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py
+2
-2
@@ -75,14 +75,14 @@ QUANT_CONFIGS = [
75
"b_type": "kU4B8",
76
"thread_configs": THREAD_CONFIGS,
77
"thread_m_blocks": THREAD_M_BLOCKS,
78
- "group_blocks": [-1, 0, 2, 4, 8],
78
+ "group_blocks": [-1, 2, 4, 8],
79
},
80
# AWQ-INT8
81
{
82
"b_type": "kU8B128",
83
"thread_configs": THREAD_CONFIGS,
84
"thread_m_blocks": THREAD_M_BLOCKS,
85
- "group_blocks": [-1, 0, 2, 4, 8],
85
+ "group_blocks": [-1, 2, 4, 8],
86
},
87
# FP8
88
{
csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h
+14
-15
@@ -7,21 +7,20 @@
7
#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
8
#include "core/scalar_type.hpp"
9
10
-#define MARLIN_KERNEL_PARAMS \
11
- const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
- int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
- const int4 *__restrict__ b_bias_ptr, \
14
- const float *__restrict__ a_scales_ptr, \
15
- const int4 *__restrict__ scales_ptr, \
16
- const float *__restrict__ global_scale_ptr, \
17
- const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
- const int32_t *__restrict__ sorted_token_ids_ptr, \
19
- const int32_t *__restrict__ expert_ids_ptr, \
20
- const int32_t *__restrict__ num_tokens_past_padded_ptr, \
21
- const float *__restrict__ topk_weights_ptr, int top_k, \
22
- bool mul_topk_weights, int num_groups, int prob_m, int prob_n, \
23
- int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
24
- bool use_fp32_reduce
10
+#define MARLIN_KERNEL_PARAMS \
11
+ const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
+ int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
+ const int4 *__restrict__ b_bias_ptr, \
14
+ const float *__restrict__ a_scales_ptr, \
15
+ const int4 *__restrict__ scales_ptr, \
16
+ const float *__restrict__ global_scale_ptr, \
17
+ const int4 *__restrict__ zp_ptr, \
18
+ const int32_t *__restrict__ sorted_token_ids_ptr, \
19
+ const int32_t *__restrict__ expert_ids_ptr, \
20
+ const int32_t *__restrict__ num_tokens_past_padded_ptr, \
21
+ const float *__restrict__ topk_weights_ptr, int top_k, \
22
+ bool mul_topk_weights, int prob_m, int prob_n, int prob_k, int *locks, \
23
+ bool has_bias, bool use_atomic_add, bool use_fp32_reduce
24
25
namespace MARLIN_NAMESPACE_NAME {
26
template <const vllm::ScalarTypeId a_type_id, // A ScalarType id
csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h
+95
-345
@@ -50,10 +50,9 @@ template <typename scalar_t, // compute dtype, half or nv_float16
50
// only works when thread_m_blocks == 1
51
const int stages, // number of stages for the async global->shared
52
// fetch pipeline
53
- const bool has_act_order, // whether act_order is enabled
54
- const int group_blocks, // number of consecutive 16x16 blocks
55
- // with a separate quantization scale
56
- const bool is_zp_float // is zero point of float16 type?
53
+ const int group_blocks, // number of consecutive 16x16 blocks
54
+ // with a separate quantization scale
55
+ const bool is_zp_float // is zero point of float16 type?
56
>
57
__global__ void Marlin(
58
const int4* __restrict__ A, // fp16 input matrix of shape mxk
@@ -64,14 +63,12 @@ __global__ void Marlin(
63
// (k/groupsize)xn
64
const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
65
// (k/groupsize)x(n/pack_factor)
67
- const int* __restrict__ g_idx, // int32 group indices of shape k
66
const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids
67
const int32_t* __restrict__ expert_ids_ptr, // moe expert ids
68
const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens
69
const float* __restrict__ topk_weights_ptr, // moe top weights
70
int top_k, // num of experts per token
71
bool mul_topk_weights, // mul topk weights or not
74
- int num_groups, // number of scale groups per output channel
72
int prob_m, // batch dimension m
73
int prob_n, // output dimension n
74
int prob_k, // reduction dimension k
@@ -148,29 +145,6 @@ __device__ inline void sub_zp(
145
frag_b[1] = __hsub2(frag_b[1], zp);
146
}
147
151
-// Same as above, but for act_order (each K is multiplied individually)
152
-template <vllm::ScalarTypeId type_id>
153
-__device__ inline void scale4(
154
- typename MarlinScalarType<type_id>::FragB& frag_b,
155
- typename MarlinScalarType<type_id>::FragS& frag_s_1,
156
- typename MarlinScalarType<type_id>::FragS& frag_s_2,
157
- typename MarlinScalarType<type_id>::FragS& frag_s_3,
158
- typename MarlinScalarType<type_id>::FragS& frag_s_4, int i) {
159
- using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
160
- using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
161
-
162
- scalar_t2 s_val_1_2;
163
- s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
164
- s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
165
-
166
- scalar_t2 s_val_3_4;
167
- s_val_3_4.x = reinterpret_cast<scalar_t*>(&frag_s_3)[i];
168
- s_val_3_4.y = reinterpret_cast<scalar_t*>(&frag_s_4)[i];
169
-
170
- frag_b[0] = __hmul2(frag_b[0], s_val_1_2);
171
- frag_b[1] = __hmul2(frag_b[1], s_val_3_4);
172
-}
173
-
148
// Given 2 floats multiply by 2 scales (halves)
149
template <vllm::ScalarTypeId type_id>
150
__device__ inline void scale_float(
@@ -264,15 +238,12 @@ __global__ void Marlin(
238
// 4bit packed zero-points of shape
239
// (k/groupsize, n/pack_factor)
240
const int4* __restrict__ zp_ptr,
267
- // int32 group indices of shape k
268
- const int* __restrict__ g_idx,
241
const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids
242
const int32_t* __restrict__ expert_ids_ptr, // moe expert ids
243
const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens
244
const float* __restrict__ topk_weights_ptr, // moe top weights
245
int top_k, // num of experts per token
246
bool mul_topk_weights, // mul topk weights or not
275
- int num_groups, // number of scale groups per output channel
247
int prob_m, // batch dimension m
248
int prob_n, // output dimension n
249
int prob_k, // reduction dimension k
@@ -369,12 +340,12 @@ __global__ void Marlin(
340
341
float global_scale_f32 = 1.0f;
342
372
- constexpr bool has_act_order = group_blocks == 0;
343
+ static_assert(group_blocks != 0,
344
+ "runtime activation ordering is not supported");
345
346
constexpr int pack_factor = 32 / b_type.size_bits();
347
static_assert(thread_m_blocks == 1 || !m_block_size_8);
376
- const int group_size =
377
- (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
348
+ const int group_size = group_blocks == -1 ? prob_k : group_blocks * 16;
349
const int scales_expert_stride =
350
prob_n * prob_k / group_size / (is_8bit_scale ? 16 : 8);
351
const int zp_expert_stride =
@@ -405,7 +376,7 @@ __global__ void Marlin(
376
377
int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x);
378
408
- if constexpr (!has_act_order && group_blocks != -1) {
379
+ if constexpr (group_blocks != -1) {
380
if (group_blocks >= thread_k_blocks) {
381
// Ensure that the number of tiles in each stripe is a multiple of the
382
// groupsize; this avoids an annoying special case where a stripe starts
@@ -551,9 +522,6 @@ __global__ void Marlin(
522
if constexpr (has_zp) {
523
zp_ptr += (expert_id - old_expert_id) * zp_expert_stride;
524
}
554
- if constexpr (has_act_order) {
555
- g_idx += (expert_id - old_expert_id) * prob_k;
556
- }
525
if (has_bias) {
526
b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride;
527
}
@@ -694,27 +662,17 @@ __global__ void Marlin(
662
b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
663
constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
664
697
- // Scale sizes/strides without act_order
665
+ // Scale sizes/strides
666
int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8);
667
constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8);
668
constexpr int s_tb_groups =
701
- !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
669
+ group_blocks != -1 && group_blocks < thread_k_blocks
670
? thread_k_blocks / group_blocks
671
: 1;
672
constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
673
int s_gl_rd_delta = s_gl_stride;
674
707
- // Scale size/strides with act_order
708
- constexpr int tb_k = 16 * thread_k_blocks;
709
- constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
710
- // constexpr int act_s_row_stride = 1;
711
- // int act_s_col_stride = act_s_row_stride * num_groups;
712
- constexpr int act_s_max_num_groups = 32;
713
- int act_s_col_stride = 1;
714
- int act_s_col_warp_stride = act_s_col_stride * 8;
715
-
675
constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4);
717
- int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
676
677
// Zero-points sizes/strides
678
int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4;
@@ -750,25 +708,16 @@ __global__ void Marlin(
708
auto b_sh_rd = threadIdx.x * b_thread_vecs;
709
b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1));
710
753
- // For act_order
754
- int slice_k_start = tb_k * slice_row;
755
- int slice_k_finish = slice_k_start + tb_k * slice_iters;
756
- int slice_k_start_shared_fetch = slice_k_start;
757
- int slice_n_offset = act_s_col_tb_stride * slice_col;
758
-
759
- // No act_order
711
int s_gl_rd;
761
- if constexpr (!has_act_order) {
762
- if constexpr (group_blocks == -1) {
763
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
764
- } else if constexpr (group_blocks >= thread_k_blocks) {
765
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
766
- s_sh_stride * slice_col + threadIdx.x;
767
- } else {
768
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
769
- threadIdx.x / s_sh_stride) +
770
- s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
771
- }
712
+ if constexpr (group_blocks == -1) {
713
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
714
+ } else if constexpr (group_blocks >= thread_k_blocks) {
715
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
716
+ s_sh_stride * slice_col + threadIdx.x;
717
+ } else {
718
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
719
+ threadIdx.x / s_sh_stride) +
720
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
721
}
722
auto s_sh_wr = threadIdx.x;
723
bool s_sh_wr_pred = threadIdx.x < s_sh_stage;
@@ -885,10 +834,8 @@ __global__ void Marlin(
834
: (sh_size_b_red_min + sh_bias_size);
835
836
int4* sh_bias = sh_new + sh_size_b_red_min;
888
- int4* sh_g_idx = sh_new + sh_b_red_bias_size;
889
- int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
890
- constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
891
- : (stages * s_sh_stage);
837
+ int4* sh_zp = sh_new + sh_b_red_bias_size;
838
+ constexpr int sh_s_size = stages * s_sh_stage;
839
int4* sh_s = sh_zp + (stages * zp_sh_stage);
840
int4* sh_a = sh_s + sh_s_size;
841
@@ -897,9 +844,8 @@ __global__ void Marlin(
844
I4 frag_b_quant[2][b_thread_vecs];
845
FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2];
846
FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2];
900
- FragS frag_s[2][4]; // No act-order
847
+ FragS frag_s[2][4];
848
FragS frag_bias[2][4];
902
- FragS act_frag_s[2][4][4]; // For act-order
849
int frag_qzp[2][num_ints_per_thread]; // Zero-points
850
FragZP frag_zp; // Zero-points in fp16
851
FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ
@@ -929,42 +875,6 @@ __global__ void Marlin(
875
reinterpret_cast<float*>(frag_c)[i] = 0;
876
};
877
932
- int sh_first_group_id = -1;
933
- int sh_num_groups = -1;
934
-
935
- auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
936
- int last_group_id) {
937
- sh_first_group_id = first_group_id;
938
- sh_num_groups = last_group_id - first_group_id + 1;
939
-
940
- if (sh_num_groups > act_s_max_num_groups) {
941
- sh_num_groups = act_s_max_num_groups;
942
- }
943
-
944
- if (sh_first_group_id + sh_num_groups > num_groups) {
945
- sh_num_groups = num_groups - sh_first_group_id;
946
- }
947
-
948
- int row_offset = first_group_id * s_gl_stride;
949
-
950
- if (is_async) {
951
- for (int i = 0; i < sh_num_groups; i++) {
952
- if (threadIdx.x < s_sh_stride) {
953
- cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
954
- &scales_ptr[row_offset + (i * s_gl_stride) +
955
- slice_n_offset + threadIdx.x]);
956
- }
957
- }
958
- } else {
959
- for (int i = 0; i < sh_num_groups; i++) {
960
- if (threadIdx.x < s_sh_stride) {
961
- sh_s[(i * s_sh_stride) + threadIdx.x] =
962
- scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
963
- threadIdx.x];
964
- }
965
- }
966
- }
967
- };
878
// Asynchronously fetch the next A, B and s tile from global to the next
879
// shared memory pipeline location.
880
auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
@@ -995,44 +905,27 @@ __global__ void Marlin(
905
906
b_gl_rd += b_gl_rd_delta_o;
907
998
- if constexpr (has_act_order) {
999
- // Fetch g_idx thread-block portion
1000
- int full_pipe = a_off;
1001
- int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
1002
- if (cur_k < prob_k && cur_k < slice_k_finish) {
1003
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1004
-
1005
- int4 const* cur_g_idx_stage_ptr =
1006
- reinterpret_cast<int4 const*>(&g_idx[cur_k]);
1007
-
1008
- if (threadIdx.x < g_idx_stage) {
1009
- cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
1010
- &cur_g_idx_stage_ptr[threadIdx.x]);
1011
- }
1012
- }
1013
- } else {
1014
- if constexpr (group_blocks != -1) {
1015
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
908
+ if constexpr (group_blocks != -1) {
909
+ int4* sh_s_stage = sh_s + s_sh_stage * pipe;
910
1017
- // Only fetch scales if this tile starts a new group
1018
- if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
1019
- if (s_sh_wr_pred) {
1020
- cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
1021
- }
1022
- s_gl_rd += s_gl_rd_delta * s_tb_groups;
911
+ // Only fetch scales if this tile starts a new group
912
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
913
+ if (s_sh_wr_pred) {
914
+ cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
915
}
916
+ s_gl_rd += s_gl_rd_delta * s_tb_groups;
917
}
918
+ }
919
1026
- if constexpr (has_zp && group_blocks != -1) {
1027
- int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
920
+ if constexpr (has_zp && group_blocks != -1) {
921
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
922
1029
- // Only fetch zero points if this tile starts a new group
1030
- if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
1031
- if (zp_sh_wr_pred) {
1032
- cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
1033
- }
1034
- zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
923
+ // Only fetch zero points if this tile starts a new group
924
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
925
+ if (zp_sh_wr_pred) {
926
+ cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
927
}
928
+ zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
929
}
930
}
931
}
@@ -1080,150 +973,57 @@ __global__ void Marlin(
973
}
974
};
975
1083
- bool is_same_group[stages];
1084
- int same_group_id[stages];
1085
-
1086
- auto init_same_group = [&](int pipe) {
1087
- if constexpr (!has_act_order) {
1088
- return;
1089
- }
1090
-
1091
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1092
- int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
1093
-
1094
- int group_id_1 = sh_g_idx_int_ptr[0];
1095
- int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];
1096
-
1097
- is_same_group[pipe] = group_id_1 == group_id_2;
1098
- same_group_id[pipe] = group_id_1;
1099
- };
1100
-
976
auto fetch_scales_to_registers = [&](int k, int full_pipe) {
977
int pipe = full_pipe % stages;
978
using IT1 = typename std::conditional_t<is_a_8bit, int2, int4>;
979
using IT0 = typename std::conditional_t<is_a_8bit, int, int2>;
980
constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1);
981
1107
- if constexpr (!has_act_order) {
1108
- // No act-order case
1109
- if constexpr (group_blocks == -1) {
1110
- // load only when starting a new slice
1111
- if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
1112
- reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
1113
- reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1114
- }
1115
- } else if constexpr (group_blocks != -1) {
1116
- if constexpr (group_blocks >= thread_k_blocks) {
1117
- constexpr int g = group_blocks / thread_k_blocks;
1118
- if (pipe % g == 0) {
1119
- if (k % b_sh_wr_iters == 0) {
1120
- int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
1121
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1122
- } else {
1123
- reinterpret_cast<int4*>(&frag_s[1])[0] =
1124
- reinterpret_cast<int4*>(&frag_s[0])[0];
1125
- }
1126
- }
1127
- } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
1128
- auto warp_id = threadIdx.x / 32;
1129
- int warp_row = warp_id / tb_n_warps;
1130
-
1131
- int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1132
- int cur_group_id = k_blocks / group_blocks2;
1133
-
1134
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1135
-
1136
- if constexpr (!is_8bit_scale) {
1137
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1138
- sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1139
- } else {
1140
- reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1141
- reinterpret_cast<int2*>(
1142
- sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1143
- }
1144
- } else if (group_blocks >= b_sh_wr_iters) {
1145
- if constexpr (!is_8bit_scale) {
1146
- reinterpret_cast<int4*>(&frag_s[1])[0] =
1147
- reinterpret_cast<int4*>(&frag_s[0])[0];
1148
- } else {
1149
- reinterpret_cast<int2*>(&frag_s[1])[0] =
1150
- reinterpret_cast<int2*>(&frag_s[0])[0];
1151
- }
982
+ if constexpr (group_blocks == -1) {
983
+ // Load only when starting a new slice.
984
+ if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
985
+ reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
986
+ reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
987
+ }
988
+ } else if constexpr (group_blocks >= thread_k_blocks) {
989
+ constexpr int g = group_blocks / thread_k_blocks;
990
+ if (pipe % g == 0) {
991
+ if (k % b_sh_wr_iters == 0) {
992
+ int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
993
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
994
+ } else {
995
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
996
+ reinterpret_cast<int4*>(&frag_s[0])[0];
997
}
998
}
999
+ } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
1000
+ auto warp_id = threadIdx.x / 32;
1001
+ int warp_row = warp_id / tb_n_warps;
1002
1155
- return;
1156
- }
1157
-
1158
- // Act-order case
1159
-
1160
- // Determine K of the "current" thread-block
1161
- int cur_k = slice_k_start + tb_k * full_pipe;
1162
- if (cur_k >= prob_k || cur_k >= slice_k_finish) {
1163
- return;
1164
- }
1165
-
1166
- // Reset (to current thread-block) since we read g_idx portion from the
1167
- // shared memory
1168
- cur_k = 0;
1003
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1004
+ int cur_group_id = k_blocks / group_blocks2;
1005
1170
- // Progress to current iteration
1171
- cur_k += k % b_sh_wr_iters;
1006
+ int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1007
1173
- // Determine "position" inside the thread-block (based on warp and
1174
- // thread-id)
1175
- auto warp_id = threadIdx.x / 32;
1176
- int warp_row = warp_id / tb_n_warps;
1177
- int warp_col = warp_id % tb_n_warps;
1178
-
1179
- cur_k += warp_row * 16 * b_sh_wr_iters;
1180
-
1181
- auto th_id = threadIdx.x % 32;
1182
- cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
1183
-
1184
- int s_col_shift =
1185
- /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +
1186
- (th_id / 4) * act_s_col_stride;
1187
-
1188
- if (is_same_group[pipe]) {
1189
- if (k % 2 == 0) {
1190
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1191
- sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +
1192
- s_col_shift];
1008
+ if constexpr (!is_8bit_scale) {
1009
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1010
+ sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1011
} else {
1194
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1195
- *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));
1012
+ reinterpret_cast<int2*>(&frag_s[k % 2])[0] = reinterpret_cast<int2*>(
1013
+ sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1014
}
1197
-
1198
- for (int i = 1; i < 4; i++) {
1199
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1200
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));
1015
+ } else if (group_blocks >= b_sh_wr_iters) {
1016
+ if constexpr (!is_8bit_scale) {
1017
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
1018
+ reinterpret_cast<int4*>(&frag_s[0])[0];
1019
+ } else {
1020
+ reinterpret_cast<int2*>(&frag_s[1])[0] =
1021
+ reinterpret_cast<int2*>(&frag_s[0])[0];
1022
}
1202
- return;
1203
- }
1204
-
1205
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1206
- int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
1207
-
1208
- constexpr int k_frag_offsets[4] = {0, 1, 8,
1209
- 9}; // Tensor core offsets per thread
1210
-
1211
- #pragma unroll
1212
- for (int i = 0; i < 4; i++) {
1213
- int actual_k = cur_k + k_frag_offsets[i];
1214
-
1215
- int group_id = sh_g_idx_int_ptr[actual_k];
1216
- int rel_group_id = group_id - sh_first_group_id;
1217
-
1218
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1219
- sh_s[rel_group_id * s_sh_stride + s_col_shift];
1023
}
1024
};
1025
1026
auto fetch_zp_to_registers = [&](int k, int full_pipe) {
1224
- // This code does not handle group_blocks == 0,
1225
- // which signifies act_order.
1226
- // has_zp implies AWQ, which doesn't have act_order,
1027
static_assert(!has_zp || group_blocks != 0);
1028
1029
if constexpr (has_zp && !is_zp_float) {
@@ -1313,8 +1113,7 @@ __global__ void Marlin(
1113
constexpr int g =
1114
group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1;
1115
const bool is_new_zp =
1316
- (group_blocks == 0) ||
1317
- ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) &&
1116
+ (group_blocks > 0 && (group_blocks < b_sh_wr_iters || k == 0)) &&
1117
(pipe % g == 0) ||
1118
(group_blocks == -1 && is_first_matmul_in_slice);
1119
if constexpr (has_zp && !is_zp_float) {
@@ -1382,15 +1181,8 @@ __global__ void Marlin(
1181
sub_zp<a_type_id>(frag_b1, frag_zp[j], 1);
1182
}
1183
1385
- // Apply scale to frag_b0
1386
- if constexpr (has_act_order && !is_a_8bit) {
1387
- static_assert(group_blocks != -1);
1388
- scale4<a_type_id>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1389
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1390
- scale4<a_type_id>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1391
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1392
- } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1393
- group_blocks == -1 && !is_a_8bit) {
1184
+ if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1185
+ group_blocks == -1 && !is_a_8bit) {
1186
int idx = (threadIdx.x / 4) % 2;
1187
scalar_t2 s2 = Adtype::nums2num2(
1188
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
@@ -1805,7 +1597,7 @@ __global__ void Marlin(
1597
1598
// For per-column quantization we finally apply the scale here (only for
1599
// 4-bit)
1808
- if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit &&
1600
+ if constexpr (group_blocks == -1 && !is_a_8bit &&
1601
b_type.size_bits() == 4 &&
1602
(has_zp && dequant_skip_flop || !has_zp)) {
1603
c_scalar_t2 tmp_scale = s[0];
@@ -1911,15 +1703,6 @@ __global__ void Marlin(
1703
1704
#pragma unroll
1705
for (int i = 0; i < stages - 1; i++) {
1914
- if (has_act_order && i == 0) {
1915
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1916
- if (last_g_idx >= prob_k) {
1917
- last_g_idx = prob_k - 1;
1918
- }
1919
- fetch_act_order_scales_to_shared(true, g_idx[slice_k_start],
1920
- g_idx[last_g_idx]);
1921
- }
1922
-
1706
if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1707
if (i == 0) {
1708
fetch_col_zp_to_shared();
@@ -1933,14 +1716,10 @@ __global__ void Marlin(
1716
1717
zero_accums();
1718
wait_for_stage();
1936
- init_same_group(0);
1719
fetch_to_registers(0, 0);
1720
fetch_scales_to_registers(0, 0);
1721
fetch_zp_to_registers(0, 0);
1722
a_gl_rd_col += a_gl_rd_delta_o * (stages - 1);
1941
- if constexpr (has_act_order) {
1942
- slice_k_start_shared_fetch += tb_k * (stages - 1);
1943
- }
1723
};
1724
if (slice_iters) {
1725
start_pipes();
@@ -1965,7 +1744,6 @@ __global__ void Marlin(
1744
slice_iters >= stages);
1745
pipe++;
1746
wait_for_stage();
1968
- init_same_group(pipe % stages);
1747
}
1748
1749
if constexpr (!is_a_8bit) {
@@ -1983,25 +1761,6 @@ __global__ void Marlin(
1761
1762
a_gl_rd_col += a_gl_rd_delta_o * stages;
1763
1986
- if constexpr (has_act_order) {
1987
- slice_k_start += tb_k * stages;
1988
-
1989
- if (slice_k_start < prob_k) {
1990
- slice_k_start_shared_fetch += tb_k * stages;
1991
- int first_group_id = g_idx[slice_k_start];
1992
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1993
- if (last_g_idx >= prob_k) {
1994
- last_g_idx = prob_k - 1;
1995
- }
1996
- int last_group_id = g_idx[last_g_idx];
1997
- if (last_group_id >= sh_first_group_id + sh_num_groups) {
1998
- fetch_act_order_scales_to_shared(false, first_group_id,
1999
- last_group_id);
2000
- __syncthreads();
2001
- }
2002
- }
2003
- }
2004
-
1764
// Process results and, if necessary, proceed to the next column slice.
1765
// While this pattern may not be the most readable, other ways of writing
1766
// the loop seemed to noticeably worse performance after compilation.
@@ -2059,7 +1818,7 @@ __global__ void Marlin(
1818
bool last = slice_idx == slice_count - 1;
1819
// For per-column scales, we only fetch them here in the final step before
1820
// write-out
2062
- if constexpr (!has_act_order && group_blocks == -1 &&
1821
+ if constexpr (group_blocks == -1 &&
1822
(has_zp && dequant_skip_flop || !has_zp)) {
1823
if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) {
1824
if (s_sh_wr_pred) {
@@ -2078,7 +1837,7 @@ __global__ void Marlin(
1837
cp_async_fence();
1838
}
1839
2081
- if constexpr (!has_act_order && group_blocks == -1 &&
1840
+ if constexpr (group_blocks == -1 &&
1841
(has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) {
1842
if constexpr (is_a_8bit) {
1843
cp_async_wait<0>();
@@ -2109,7 +1868,7 @@ __global__ void Marlin(
1868
// For 8-bit channelwise, we apply the scale before the global reduction
1869
// that converts the fp32 results to fp16 (so that we avoid possible
1870
// overflow in fp16)
2112
- if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) {
1871
+ if constexpr (group_blocks == -1 && is_a_8bit) {
1872
#pragma unroll
1873
for (int j = 0; j < 2; j++) {
1874
float2 aa[2];
@@ -2131,8 +1890,7 @@ __global__ void Marlin(
1890
}
1891
}
1892
}
2134
- } else if (!has_act_order && group_blocks == -1 &&
2135
- b_type.size_bits() == 8 &&
1893
+ } else if (group_blocks == -1 && b_type.size_bits() == 8 &&
1894
(has_zp && dequant_skip_flop || !has_zp)) {
1895
if (threadIdx.x / 32 < tb_n_warps) {
1896
#pragma unroll
@@ -2202,33 +1960,25 @@ __global__ void Marlin(
1960
b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row;
1961
1962
bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
2205
- // Update slice k/n for scales loading
2206
- if constexpr (has_act_order) {
2207
- slice_k_start = tb_k * slice_row;
2208
- slice_k_finish = slice_k_start + tb_k * slice_iters;
2209
- slice_k_start_shared_fetch = slice_k_start;
2210
- slice_n_offset = act_s_col_tb_stride * slice_col;
1963
+ if constexpr (group_blocks == -1) {
1964
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
1965
+ zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
1966
+ } else if constexpr (group_blocks >= thread_k_blocks) {
1967
+ s_gl_rd =
1968
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
1969
+ s_sh_stride * slice_col + threadIdx.x;
1970
+ zp_gl_rd =
1971
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
1972
+ zp_sh_stride * slice_col + threadIdx.x;
1973
} else {
2212
- if constexpr (group_blocks == -1) {
2213
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
2214
- zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
2215
- } else if constexpr (group_blocks >= thread_k_blocks) {
2216
- s_gl_rd =
2217
- s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2218
- s_sh_stride * slice_col + threadIdx.x;
2219
- zp_gl_rd =
2220
- zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2221
- zp_sh_stride * slice_col + threadIdx.x;
2222
- } else {
2223
- s_gl_rd =
2224
- s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2225
- threadIdx.x / s_sh_stride) +
2226
- s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
2227
- zp_gl_rd =
2228
- zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2229
- threadIdx.x / zp_sh_stride) +
2230
- zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
2231
- }
1974
+ s_gl_rd =
1975
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
1976
+ threadIdx.x / s_sh_stride) +
1977
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
1978
+ zp_gl_rd =
1979
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
1980
+ threadIdx.x / zp_sh_stride) +
1981
+ zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
1982
}
1983
start_pipes();
1984
}
csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu
+78
-270
@@ -45,84 +45,6 @@ __global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
45
46
using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
47
48
-// For a given "a" of size [M,K] performs a permutation of the K columns based
49
-// on the given "perm" indices.
50
-template <int moe_block_size>
51
-__global__ void permute_cols_kernel(
52
- int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
53
- int4* __restrict__ out_int4_ptr,
54
- const int32_t* __restrict__ sorted_token_ids_ptr,
55
- const int32_t* __restrict__ expert_ids_ptr,
56
- const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
57
- int size_k, int top_k) {
58
- int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
59
- int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size);
60
- int32_t block_sorted_ids[moe_block_size];
61
- int block_num_valid_tokens = 0;
62
- int64_t old_expert_id = 0;
63
- int64_t expert_id = 0;
64
- int row_stride = size_k * sizeof(half) / 16;
65
-
66
- auto read_moe_block_data = [&](int block_id) {
67
- block_num_valid_tokens = moe_block_size;
68
- int4* tmp_block_sorted_ids = reinterpret_cast<int4*>(block_sorted_ids);
69
- for (int i = 0; i < moe_block_size / 4; i++) {
70
- tmp_block_sorted_ids[i] =
71
- ((int4*)sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i];
72
- }
73
- for (int i = 0; i < moe_block_size; i++) {
74
- if (block_sorted_ids[i] >= size_m * top_k) {
75
- block_num_valid_tokens = i;
76
- break;
77
- };
78
- }
79
- };
80
-
81
- auto permute_row = [&](int row) {
82
- int iters = size_k / default_threads;
83
- int rest = size_k % default_threads;
84
-
85
- int in_offset = (row / top_k) * row_stride;
86
- int out_offset = row * row_stride;
87
-
88
- half const* a_row_half =
89
- reinterpret_cast<half const*>(a_int4_ptr + in_offset);
90
- half* out_half = reinterpret_cast<half*>(out_int4_ptr + out_offset);
91
-
92
- int base_k = 0;
93
-
94
- for (int i = 0; i < iters; i++) {
95
- auto cur_k = base_k + threadIdx.x;
96
- int src_pos = perm_int_ptr[cur_k];
97
-
98
- out_half[cur_k] = a_row_half[src_pos];
99
-
100
- base_k += default_threads;
101
- }
102
-
103
- if (rest) {
104
- if (threadIdx.x < rest) {
105
- auto cur_k = base_k + threadIdx.x;
106
- int src_pos = perm_int_ptr[cur_k];
107
-
108
- out_half[cur_k] = a_row_half[src_pos];
109
- }
110
- }
111
- };
112
-
113
- for (int index = blockIdx.x; index < num_moe_blocks; index += gridDim.x) {
114
- old_expert_id = expert_id;
115
- int tmp_expert_id = expert_ids_ptr[index];
116
- if (tmp_expert_id == -1) continue;
117
- expert_id = tmp_expert_id;
118
- perm_int_ptr += (expert_id - old_expert_id) * size_k;
119
- read_moe_block_data(index);
120
-
121
- for (int i = 0; i < block_num_valid_tokens; i++)
122
- permute_row(block_sorted_ids[i]);
123
- }
124
-}
125
-
48
typedef struct {
49
int thread_k;
50
int thread_n;
@@ -152,9 +74,7 @@ typedef struct {
74
75
int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
76
int prob_n, int prob_k, int num_bits, int group_size,
155
- bool has_act_order, bool is_k_full, int stages) {
156
- bool cache_scales_chunk = has_act_order && !is_k_full;
157
-
77
+ int stages) {
78
int tb_n = th_config.thread_n;
79
int tb_k = th_config.thread_k;
80
@@ -162,28 +82,17 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
82
int tb_groups;
83
if (group_size == -1) {
84
tb_groups = 1;
165
- } else if (group_size == 0) {
166
- tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size
85
} else {
86
tb_groups = div_ceil(tb_k, group_size);
87
}
88
171
- if (cache_scales_chunk) {
172
- int load_groups =
173
- tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K
174
- load_groups = max(load_groups, 32); // We load at least 32 scale groups
175
- return load_groups * tb_n * 2;
176
- } else {
177
- int tb_scales = tb_groups * tb_n * 2;
178
-
179
- return tb_scales * stages;
180
- }
89
+ int tb_scales = tb_groups * tb_n * 2;
90
+ return tb_scales * stages;
91
}
92
93
int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
94
int thread_m_blocks, int prob_m, int prob_n,
185
- int prob_k, int num_bits, int group_size,
186
- bool has_act_order, bool is_k_full, int has_zp,
95
+ int prob_k, int num_bits, int group_size, int has_zp,
96
int is_zp_float, bool is_a_8bit, int stages) {
97
int pack_factor = 32 / num_bits;
98
@@ -203,10 +112,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
112
(sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size;
113
tmp_size = max(max(sh_b_size, sh_red_size), tmp_size);
114
206
- int sh_s_size =
207
- get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
208
- group_size, has_act_order, is_k_full, stages);
209
- int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0;
115
+ int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k,
116
+ num_bits, group_size, stages);
117
int sh_zp_size = 0;
118
if (has_zp) {
119
if (is_zp_float)
@@ -217,16 +124,15 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
124
sh_zp_size = sh_s_size / 2;
125
}
126
220
- int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size +
221
- sh_g_idx_size + sh_block_meta_size;
127
+ int total_size =
128
+ tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_block_meta_size;
129
130
return total_size;
131
}
132
133
bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
134
int thread_m_blocks, int prob_m, int prob_n, int prob_k,
228
- int num_bits, int group_size, bool has_act_order,
229
- bool is_k_full, int has_zp, int is_zp_float,
135
+ int num_bits, int group_size, int has_zp, int is_zp_float,
136
bool is_a_8bit, int stages, int max_shared_mem) {
137
// Sanity
138
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,
156
}
157
158
// Check that pipeline fits into cache
253
- int cache_size =
254
- get_kernel_cache_size(th_config, m_block_size_8, thread_m_blocks, prob_m,
255
- prob_n, prob_k, num_bits, group_size, has_act_order,
256
- is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
159
+ int cache_size = get_kernel_cache_size(
160
+ th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
161
+ num_bits, group_size, has_zp, is_zp_float, is_a_8bit, stages);
162
return cache_size <= max_shared_mem;
163
}
164
260
-MarlinFuncPtr get_marlin_kernel(
261
- const vllm::ScalarType a_type, const vllm::ScalarType b_type,
262
- const vllm::ScalarType c_type, const vllm::ScalarType s_type,
263
- int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
264
- bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
265
- int threads, bool is_zp_float, int stages) {
165
+MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType a_type,
166
+ const vllm::ScalarType b_type,
167
+ const vllm::ScalarType c_type,
168
+ const vllm::ScalarType s_type,
169
+ int thread_m_blocks, int thread_n_blocks,
170
+ int thread_k_blocks, bool m_block_size_8,
171
+ bool has_zp, int group_blocks, int threads,
172
+ bool is_zp_float, int stages) {
173
int num_bits = b_type.size_bits();
174
auto kernel = MarlinDefault;
175
@@ -275,9 +182,8 @@ exec_config_t determine_exec_config(
182
const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
183
const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
184
int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks,
278
- bool m_block_size_8, int num_bits, int group_size, bool has_act_order,
279
- bool is_k_full, bool has_zp, bool is_zp_float, bool is_a_8bit, int stages,
280
- int max_shared_mem, int sms) {
185
+ bool m_block_size_8, int num_bits, int group_size, bool has_zp,
186
+ bool is_zp_float, bool is_a_8bit, int stages, int max_shared_mem, int sms) {
187
exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
188
thread_config_t* thread_configs = thread_m_blocks > 1
189
? large_batch_thread_configs
@@ -293,27 +199,22 @@ exec_config_t determine_exec_config(
199
thread_config_t th_config = thread_configs[i];
200
201
if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m,
296
- prob_n, prob_k, num_bits, group_size, has_act_order,
297
- is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
202
+ prob_n, prob_k, num_bits, group_size, has_zp,
203
+ is_zp_float, is_a_8bit, stages,
204
max_shared_mem - 512)) {
205
continue;
206
}
207
208
int cache_size = get_kernel_cache_size(
209
th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
304
- num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
305
- is_a_8bit, stages);
210
+ num_bits, group_size, has_zp, is_zp_float, is_a_8bit, stages);
211
307
- int group_blocks = 0;
308
- if (!has_act_order) {
309
- group_blocks = group_size == -1 ? -1 : (group_size / 16);
310
- }
212
+ int group_blocks = group_size == -1 ? -1 : (group_size / 16);
213
312
- auto kernel =
313
- get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
314
- th_config.thread_n / 16, th_config.thread_k / 16,
315
- m_block_size_8, has_act_order, has_zp, group_blocks,
316
- th_config.num_threads, is_zp_float, stages);
214
+ auto kernel = get_marlin_kernel(
215
+ a_type, b_type, c_type, s_type, thread_m_blocks,
216
+ th_config.thread_n / 16, th_config.thread_k / 16, m_block_size_8,
217
+ has_zp, group_blocks, th_config.num_threads, is_zp_float, stages);
218
219
if (kernel == MarlinDefault) continue;
220
@@ -342,15 +243,14 @@ exec_config_t determine_exec_config(
243
}
244
245
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
345
- void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
346
- void* perm, void* a_tmp, void* sorted_token_ids,
347
- void* expert_ids, void* num_tokens_past_padded,
348
- void* topk_weights, int moe_block_size, int num_experts,
349
- int top_k, bool mul_topk_weights, int prob_m, int prob_n,
350
- int prob_k, void* workspace, vllm::ScalarType const& a_type,
246
+ void* a_s, void* b_s, void* g_s, void* zp,
247
+ void* sorted_token_ids, void* expert_ids,
248
+ void* num_tokens_past_padded, void* topk_weights,
249
+ int moe_block_size, int num_experts, int top_k,
250
+ bool mul_topk_weights, int prob_m, int prob_n, int prob_k,
251
+ void* workspace, vllm::ScalarType const& a_type,
252
vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
352
- vllm::ScalarType const& s_type, bool has_bias,
353
- bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
253
+ vllm::ScalarType const& s_type, bool has_bias, bool has_zp,
254
int group_size, int dev, cudaStream_t stream, int thread_k,
255
int thread_n, int sms, int blocks_per_sm, bool use_atomic_add,
256
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,
261
STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [",
262
prob_m, ", ", prob_n, ", ", prob_k, "]");
263
364
- int group_blocks = 0;
365
- if (has_act_order) {
366
- if (is_k_full) {
367
- STD_TORCH_CHECK(group_size != -1);
368
- group_blocks = group_size / 16;
369
- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
370
- " is not divisible by group_blocks = ", group_blocks);
371
- } else {
372
- STD_TORCH_CHECK(group_size == 0);
373
- group_blocks = 0;
374
- }
264
+ int group_blocks;
265
+ if (group_size == -1) {
266
+ group_blocks = -1;
267
} else {
376
- if (group_size == -1) {
377
- group_blocks = -1;
378
- } else {
379
- group_blocks = group_size / 16;
380
- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
381
- " is not divisible by group_blocks = ", group_blocks);
382
- }
268
+ group_blocks = group_size / 16;
269
+ STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
270
+ " is not divisible by group_blocks = ", group_blocks);
271
}
272
273
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,
280
const int4* b_s_ptr = (const int4*)b_s;
281
const float* g_s_ptr = (const float*)g_s;
282
const int4* zp_ptr = (const int4*)zp;
395
- const int* g_idx_ptr = (const int*)g_idx;
396
- const int* perm_ptr = (const int*)perm;
397
- int4* a_tmp_ptr = (int4*)a_tmp;
283
const int32_t* sorted_token_ids_ptr = (const int32_t*)sorted_token_ids;
284
const int32_t* expert_ids_ptr = (const int32_t*)expert_ids;
285
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,
287
const float* topk_weights_ptr = (const float*)topk_weights;
288
int* locks = (int*)workspace;
289
405
- if (has_act_order) {
406
- // Permute A columns
407
- auto kernel = permute_cols_kernel<8>;
408
- if (moe_block_size == 8) {
409
- } else if (moe_block_size == 16)
410
- kernel = permute_cols_kernel<16>;
411
- else if (moe_block_size == 32)
412
- kernel = permute_cols_kernel<32>;
413
- else if (moe_block_size == 48)
414
- kernel = permute_cols_kernel<48>;
415
- else if (moe_block_size == 64)
416
- kernel = permute_cols_kernel<64>;
417
- else
418
- STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
419
-
420
- // avoid ">>>" being formatted to "> > >"
421
- // clang-format off
422
- kernel<<<sms, default_threads, 0, stream>>>(
423
- A_ptr, perm_ptr, a_tmp_ptr, sorted_token_ids_ptr, expert_ids_ptr,
424
- num_tokens_past_padded_ptr, prob_m, prob_k, top_k);
425
- // clang-format on
426
- A_ptr = a_tmp_ptr;
427
- prob_m = prob_m * top_k;
428
- top_k = 1;
429
-
430
- // If we have a full K, then we can run the non-act-order version of Marlin
431
- // (since the weight rows are reordered by increasing group ids, and by
432
- // having a full K, we have full original groups)
433
- if (is_k_full) has_act_order = false;
434
- }
435
-
290
int max_shared_mem = 0;
291
cudaDeviceGetAttribute(&max_shared_mem,
292
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
@@ -476,9 +330,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
330
// Auto config
331
exec_cfg = determine_exec_config(
332
a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts,
479
- top_k, thread_m_blocks, m_block_size_8, num_bits, group_size,
480
- has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
481
- max_shared_mem, sms);
333
+ top_k, thread_m_blocks, m_block_size_8, num_bits, group_size, has_zp,
334
+ is_zp_float, is_a_8bit, stages, max_shared_mem, sms);
335
thread_tfg = exec_cfg.tb_cfg;
336
}
337
@@ -494,36 +347,32 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
347
348
STD_TORCH_CHECK(
349
is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
497
- prob_n, prob_k, num_bits, group_size, has_act_order,
498
- is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
499
- max_shared_mem),
350
+ prob_n, prob_k, num_bits, group_size, has_zp, is_zp_float,
351
+ is_a_8bit, stages, max_shared_mem),
352
"Invalid thread config: thread_m_blocks = ", thread_m_blocks,
353
", thread_k = ", thread_tfg.thread_k,
354
", thread_n = ", thread_tfg.thread_n,
355
", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
356
prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
505
- ", group_size = ", group_size, ", has_act_order = ", has_act_order,
506
- ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
357
+ ", group_size = ", group_size, ", has_zp = ", has_zp,
358
", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
359
509
- int sh_cache_size =
510
- get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
511
- prob_n, prob_k, num_bits, group_size, has_act_order,
512
- is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
360
+ int sh_cache_size = get_kernel_cache_size(
361
+ thread_tfg, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
362
+ num_bits, group_size, has_zp, is_zp_float, is_a_8bit, stages);
363
514
- auto kernel = get_marlin_kernel(
515
- a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
516
- thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
517
- num_threads, is_zp_float, stages);
364
+ auto kernel =
365
+ get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
366
+ thread_n_blocks, thread_k_blocks, m_block_size_8,
367
+ has_zp, group_blocks, num_threads, is_zp_float, stages);
368
369
if (kernel == MarlinDefault) {
520
- STD_TORCH_CHECK(
521
- false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ",
522
- prob_k, "]", ", has_act_order = ", has_act_order,
523
- ", num_groups = ", num_groups, ", group_size = ", group_size,
524
- ", thread_m_blocks = ", thread_m_blocks,
525
- ", thread_n_blocks = ", thread_n_blocks,
526
- ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits);
370
+ STD_TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
371
+ ", ", prob_k, "]", ", group_size = ", group_size,
372
+ ", thread_m_blocks = ", thread_m_blocks,
373
+ ", thread_n_blocks = ", thread_n_blocks,
374
+ ", thread_k_blocks = ", thread_k_blocks,
375
+ ", num_bits = ", num_bits);
376
}
377
378
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
@@ -531,10 +380,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
380
// avoid ">>>" being formatted to "> > >"
381
// clang-format off
382
kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
534
- 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,
383
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr,
384
sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
536
- topk_weights_ptr, top_k, mul_topk_weights, num_groups, prob_m,
537
- prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce);
385
+ topk_weights_ptr, top_k, mul_topk_weights, prob_m, prob_n, prob_k, locks,
386
+ has_bias, use_atomic_add, use_fp32_reduce);
387
// clang-format on
388
}
389
@@ -548,15 +397,13 @@ torch::stable::Tensor moe_wna16_marlin_gemm(
397
std::optional<torch::stable::Tensor> const& a_scales_or_none,
398
std::optional<torch::stable::Tensor> const& global_scale_or_none,
399
std::optional<torch::stable::Tensor> const& b_zeros_or_none,
551
- std::optional<torch::stable::Tensor> const& g_idx_or_none,
552
- std::optional<torch::stable::Tensor> const& perm_or_none,
400
torch::stable::Tensor& workspace, torch::stable::Tensor& sorted_token_ids,
401
torch::stable::Tensor& expert_ids,
402
torch::stable::Tensor& num_tokens_past_padded,
403
torch::stable::Tensor& topk_weights, int64_t moe_block_size, int64_t top_k,
404
bool mul_topk_weights, vllm::ScalarTypeId const& b_type_id, int64_t size_m,
558
- int64_t size_n, int64_t size_k, bool is_k_full, bool use_atomic_add,
559
- bool use_fp32_reduce, bool is_zp_float, int64_t thread_k, int64_t thread_n,
405
+ int64_t size_n, int64_t size_k, bool use_atomic_add, bool use_fp32_reduce,
406
+ bool is_zp_float, int64_t thread_k, int64_t thread_n,
407
int64_t blocks_per_sm) {
408
vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
409
@@ -717,7 +564,7 @@ torch::stable::Tensor moe_wna16_marlin_gemm(
564
c_tmp = torch::stable::new_empty(a, {0}, kFloat);
565
}
566
720
- // Detect groupsize and act_order
567
+ // Detect group size.
568
int num_groups = -1;
569
int group_size = -1;
570
@@ -728,50 +575,13 @@ torch::stable::Tensor moe_wna16_marlin_gemm(
575
" is not size_n = ", size_n);
576
num_groups = b_scales.size(1);
577
731
- torch::stable::Tensor g_idx, perm, a_tmp;
732
- if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
733
- g_idx = g_idx_or_none.value();
734
- perm = perm_or_none.value();
735
-
736
- STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
737
- STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
738
- STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
739
- STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
740
-
741
- // Verify g_idx and perm
742
- STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
743
- (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
744
- "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
745
- " and perm.size(-1) = ", perm.size(-1),
746
- ", where size_k = ", size_k);
747
- } else {
748
- g_idx = torch::stable::new_empty(a, {0}, c_dtype);
749
- perm = torch::stable::new_empty(a, {0}, c_dtype);
750
- a_tmp = torch::stable::new_empty(a, {0}, c_dtype);
751
- }
752
- bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
753
-
754
- if (has_act_order) {
755
- a_tmp = torch::stable::new_empty(a, {size_m * top_k, size_k}, c_dtype);
756
- if (is_k_full) {
757
- STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
758
- STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
759
- ", is not divisible by num_groups = ", num_groups);
760
- group_size = size_k / num_groups;
761
- } else {
762
- group_size = 0;
763
- }
764
-
578
+ if (num_groups > 1) {
579
+ STD_TORCH_CHECK(
580
+ size_k % num_groups == 0, "size_k = ", size_k,
581
+ ", is not divisible by b_scales.size(1) = ", b_scales.size(1));
582
+ group_size = size_k / num_groups;
583
} else {
766
- a_tmp = torch::stable::new_empty(a, {0}, c_dtype);
767
- if (num_groups > 1) {
768
- STD_TORCH_CHECK(
769
- size_k % num_groups == 0, "size_k = ", size_k,
770
- ", is not divisible by b_scales.size(1) = ", b_scales.size(1));
771
- group_size = size_k / num_groups;
772
- } else {
773
- group_size = -1;
774
- }
584
+ group_size = -1;
585
}
586
587
torch::stable::Tensor global_scale;
@@ -880,15 +690,13 @@ torch::stable::Tensor moe_wna16_marlin_gemm(
690
c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(),
691
a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(),
692
global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(),
883
- g_idx.mutable_data_ptr(), perm.mutable_data_ptr(),
884
- a_tmp.mutable_data_ptr(), sorted_token_ids.mutable_data_ptr(),
885
- expert_ids.mutable_data_ptr(), num_tokens_past_padded.mutable_data_ptr(),
693
+ sorted_token_ids.mutable_data_ptr(), expert_ids.mutable_data_ptr(),
694
+ num_tokens_past_padded.mutable_data_ptr(),
695
topk_weights.mutable_data_ptr(), moe_block_size, num_experts, top_k,
696
mul_topk_weights, size_m, size_n, size_k, workspace.mutable_data_ptr(),
888
- a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
889
- has_zp, num_groups, group_size, dev, get_current_cuda_stream(dev),
890
- thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce,
891
- is_zp_float);
697
+ a_type, b_type, c_type, s_type, has_bias, has_zp, group_size, dev,
698
+ get_current_cuda_stream(dev), thread_k, thread_n, sms, blocks_per_sm,
699
+ use_atomic_add, use_fp32_reduce, is_zp_float);
700
701
return c;
702
}
csrc/libtorch_stable/moe/torch_bindings.cpp
+2
-3
@@ -79,13 +79,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
79
"Tensor! b_q_weight, Tensor? b_bias_or_none,"
80
"Tensor! b_scales, Tensor? a_scales, Tensor? global_scale, Tensor? "
81
"b_zeros_or_none,"
82
- "Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace,"
82
+ "Tensor! workspace,"
83
"Tensor sorted_token_ids,"
84
"Tensor! expert_ids, Tensor! num_tokens_past_padded,"
85
"Tensor! topk_weights, int moe_block_size, int top_k, "
86
"bool mul_topk_weights, int b_type_id,"
87
- "int size_m, int size_n, int size_k,"
88
- "bool is_full_k, bool use_atomic_add,"
87
+ "int size_m, int size_n, int size_k, bool use_atomic_add,"
88
"bool use_fp32_reduce, bool is_zp_float,"
89
"int thread_k, int thread_n, int blocks_per_sm) -> Tensor");
90
csrc/libtorch_stable/ops.h
+3
-7
@@ -43,9 +43,6 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
43
int64_t group_size, double eps, double int8_min,
44
double int8_max);
45
46
-torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
47
- torch::stable::Tensor const& perm);
48
-
46
#ifndef USE_ROCM
47
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
48
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,
584
torch::stable::Tensor b_q_weight,
585
torch::stable::Tensor b_gptq_qzeros,
586
torch::stable::Tensor b_gptq_scales,
590
- torch::stable::Tensor b_g_idx, bool use_exllama,
591
- bool use_v2_format, int64_t bit);
587
+ bool use_exllama, bool use_v2_format,
588
+ int64_t bit);
589
593
-void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
594
- int64_t bit);
590
+void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit);
591
592
// Cache ops (shared CUDA/ROCm)
593
void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
csrc/libtorch_stable/permute_cols.cu
deleted
-94
@@ -1,94 +0,0 @@
1
-#include <torch/csrc/stable/library.h>
2
-#include <torch/csrc/stable/tensor.h>
3
-#include <torch/csrc/stable/accelerator.h>
4
-#include <torch/csrc/stable/ops.h>
5
-#include <torch/headeronly/core/ScalarType.h>
6
-
7
-#include <cuda_fp16.h>
8
-
9
-#include "torch_utils.h"
10
-
11
-static constexpr int default_threads = 256;
12
-static constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
13
-
14
-// For a given "a" of size [M,K] performs a permutation of the K columns based
15
-// on the given "perm" indices.
16
-// Currently only supports 16bit types (since we permute half types)
17
-__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
18
- int const* __restrict__ perm_int_ptr,
19
- int4* __restrict__ out_int4_ptr, int size_m,
20
- int size_k, int block_rows) {
21
- int start_row = block_rows * blockIdx.x;
22
- int finish_row = start_row + block_rows;
23
- if (finish_row > size_m) {
24
- finish_row = size_m;
25
- }
26
- int cur_block_rows = std::max(finish_row - start_row, 0);
27
-
28
- int row_stride = size_k * sizeof(half) / 16;
29
-
30
- auto permute_row = [&](int row) {
31
- int iters = size_k / default_threads;
32
- int rest = size_k % default_threads;
33
-
34
- int offset = row * row_stride;
35
-
36
- half const* a_row_half = reinterpret_cast<half const*>(a_int4_ptr + offset);
37
- half* out_half = reinterpret_cast<half*>(out_int4_ptr + offset);
38
-
39
- int base_k = 0;
40
-
41
- for (int i = 0; i < iters; i++) {
42
- int cur_k = base_k + threadIdx.x;
43
- int src_pos = perm_int_ptr[cur_k];
44
-
45
- out_half[cur_k] = a_row_half[src_pos];
46
-
47
- base_k += default_threads;
48
- }
49
-
50
- if (rest) {
51
- if (threadIdx.x < rest) {
52
- int cur_k = base_k + threadIdx.x;
53
- int src_pos = perm_int_ptr[cur_k];
54
-
55
- out_half[cur_k] = a_row_half[src_pos];
56
- }
57
- }
58
- };
59
-
60
- for (int i = 0; i < cur_block_rows; i++) {
61
- int cur_row = start_row + i;
62
- if (cur_row < size_m) {
63
- permute_row(cur_row);
64
- }
65
- }
66
-}
67
-
68
-// More efficient version of A[..., perm]
69
-// taken from gptq_marlin.cu
70
-torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
71
- torch::stable::Tensor const& perm) {
72
- const int32_t dev = A.get_device_index();
73
- const torch::stable::accelerator::DeviceGuard device_guard(dev);
74
- const auto stream = get_current_cuda_stream(dev);
75
-
76
- STD_TORCH_CHECK(
77
- A.scalar_type() == torch::headeronly::ScalarType::Half ||
78
- A.scalar_type() == torch::headeronly::ScalarType::BFloat16,
79
- "Currently only 16bit types are supported");
80
- STD_TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
81
- STD_TORCH_CHECK(A.size(-1) % 8 == 0,
82
- "A columns must be a multiple of 8 (128bits)");
83
- auto A_2d = torch::stable::view(A, {-1, A.size(-1)});
84
-
85
- torch::stable::Tensor D = torch::stable::empty_like(A);
86
- int sms;
87
- cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
88
- int block_rows = div_ceil(A_2d.size(0), sms);
89
- permute_cols_kernel<<<sms, default_threads, 0, stream>>>(
90
- reinterpret_cast<int4 const*>(A_2d.const_data_ptr()),
91
- perm.const_data_ptr<int>(), reinterpret_cast<int4*>(D.mutable_data_ptr()),
92
- A_2d.size(0), A_2d.size(1), block_rows);
93
- return D;
94
-}
\ No newline at end of file
csrc/libtorch_stable/quantization/gptq/q_gemm.cu
+112
-423
@@ -121,7 +121,7 @@ typedef void (*fp_gemm_half_q_half_gptq_kernel)(const half*, const uint32_t*,
121
const uint32_t*, const half*,
122
half*, const int, const int,
123
const int, const int,
124
- const bool, const int*);
124
+ const bool);
125
126
template <bool first_block, int m_count>
127
__global__ void gemm_half_q_half_gptq_4bit_kernel(
@@ -129,7 +129,7 @@ __global__ void gemm_half_q_half_gptq_4bit_kernel(
129
const uint32_t* __restrict__ b_gptq_qzeros,
130
const half* __restrict__ b_gptq_scales, half* __restrict__ c,
131
const int size_m, const int size_n, const int size_k, const int groups,
132
- const bool use_v2_format, const int* __restrict__ b_q_perm) {
132
+ const bool use_v2_format) {
133
MatrixView_half a_(a, size_m, size_k);
134
MatrixView_half_rw c_(c, size_m, size_n);
135
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(
157
const half* a_ptr = a_.item_ptr(offset_m + m, 0);
158
half* block_a_ptr = block_a[m];
159
160
- half a0;
161
- if (b_q_perm)
162
- a0 = a_ptr[b_q_perm[offset_k + t]];
163
- else
164
- a0 = a_ptr[offset_k + t];
165
- block_a_ptr[t] = a0;
160
+ block_a_ptr[t] = a_ptr[offset_k + t];
161
}
162
}
163
@@ -263,7 +258,7 @@ __global__ void gemm_half_q_half_gptq_2bit_kernel(
258
const uint32_t* __restrict__ b_gptq_qzeros,
259
const half* __restrict__ b_gptq_scales, half* __restrict__ c,
260
const int size_m, const int size_n, const int size_k, const int groups,
266
- const bool use_v2_format, const int* __restrict__ b_q_perm) {
261
+ const bool use_v2_format) {
262
MatrixView_half a_(a, size_m, size_k);
263
MatrixView_half_rw c_(c, size_m, size_n);
264
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(
286
const half* a_ptr = a_.item_ptr(offset_m + m, 0);
287
half* block_a_ptr = block_a[m];
288
294
- half a0;
295
- if (b_q_perm)
296
- a0 = a_ptr[b_q_perm[offset_k + t]];
297
- else
298
- a0 = a_ptr[offset_k + t];
299
- block_a_ptr[t] = a0;
289
+ block_a_ptr[t] = a_ptr[offset_k + t];
290
}
291
}
292
@@ -380,7 +370,7 @@ __global__ void gemm_half_q_half_gptq_3bit_kernel(
370
const uint32_t* __restrict__ b_gptq_qzeros,
371
const half* __restrict__ b_gptq_scales, half* __restrict__ c,
372
const int size_m, const int size_n, const int size_k, const int groups,
383
- const bool use_v2_format, const int* __restrict__ b_q_perm) {
373
+ const bool use_v2_format) {
374
MatrixView_half a_(a, size_m, size_k);
375
MatrixView_half_rw c_(c, size_m, size_n);
376
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(
398
const half* a_ptr = a_.item_ptr(offset_m + m, 0);
399
half* block_a_ptr = block_a[m];
400
411
- half a0;
412
- if (b_q_perm)
413
- a0 = a_ptr[b_q_perm[offset_k + t]];
414
- else
415
- a0 = a_ptr[offset_k + t];
416
- block_a_ptr[t] = a0;
401
+ block_a_ptr[t] = a_ptr[offset_k + t];
402
}
403
}
404
@@ -504,7 +489,7 @@ __global__ void gemm_half_q_half_gptq_8bit_kernel(
489
const uint32_t* __restrict__ b_gptq_qzeros,
490
const half* __restrict__ b_gptq_scales, half* __restrict__ c,
491
const int size_m, const int size_n, const int size_k, const int groups,
507
- const bool use_v2_format, const int* __restrict__ b_q_perm) {
492
+ const bool use_v2_format) {
493
MatrixView_half a_(a, size_m, size_k);
494
MatrixView_half_rw c_(c, size_m, size_n);
495
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(
517
const half* a_ptr = a_.item_ptr(offset_m + m, 0);
518
half* block_a_ptr = block_a[m];
519
535
- half a0;
536
- if (b_q_perm)
537
- a0 = a_ptr[b_q_perm[offset_k + t]];
538
- else
539
- a0 = a_ptr[offset_k + t];
540
- block_a_ptr[t] = a0;
520
+ block_a_ptr[t] = a_ptr[offset_k + t];
521
}
522
}
523
@@ -656,10 +636,9 @@ fp_gemm_half_q_half_gptq_kernel pick_gemm_half_q_half_gptq_kernel(
636
637
void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight,
638
const uint32_t* b_gptq_qzeros,
659
- const half* b_gptq_scales, const int* b_q_perm,
660
- half* c, int size_m, int size_n, int size_k,
661
- int m_count, int groups, bool use_v2_format,
662
- int bit) {
639
+ const half* b_gptq_scales, half* c, int size_m,
640
+ int size_n, int size_k, int m_count, int groups,
641
+ bool use_v2_format, int bit) {
642
dim3 blockDim, gridDim;
643
blockDim.x = BLOCK_KN_SIZE;
644
blockDim.y = 1;
@@ -672,13 +651,13 @@ void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight,
651
pick_gemm_half_q_half_gptq_kernel(true, m_count, bit);
652
653
const cudaStream_t stream = get_current_cuda_stream();
675
- kernel<<<gridDim, blockDim, 0, stream>>>(
676
- a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m, size_n, size_k,
677
- groups, use_v2_format, b_q_perm);
654
+ kernel<<<gridDim, blockDim, 0, stream>>>(a, b_q_weight, b_gptq_qzeros,
655
+ b_gptq_scales, c, size_m, size_n,
656
+ size_k, groups, use_v2_format);
657
}
658
659
__global__ void reconstruct_exllama_8bit_kernel(
681
- const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm,
660
+ const uint32_t* __restrict__ b_q_weight,
661
const uint32_t* __restrict__ b_gptq_qzeros,
662
const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
663
const int groups, const bool use_v2_format, half* __restrict__ b) {
@@ -694,14 +673,8 @@ __global__ void reconstruct_exllama_8bit_kernel(
673
674
int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
675
697
- // Preload remapping table
698
- __shared__ int perm[BLOCK_KN_SIZE];
676
auto t = threadIdx.x;
677
701
- if (b_q_perm) {
702
- if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
703
- }
704
-
678
// Column
679
int n = offset_n + t * 4;
680
if (n >= size_n) return;
@@ -722,8 +695,6 @@ __global__ void reconstruct_exllama_8bit_kernel(
695
b_gptq_qzeros_.item4(zeros, group, n);
696
b_gptq_scales_.item4_h2(scales, group, n);
697
725
- __syncthreads();
726
-
698
int k = offset_k;
699
int lk = 0;
700
@@ -752,25 +723,13 @@ __global__ void reconstruct_exllama_8bit_kernel(
723
dequant_8bit_8(load_int4[0].w, load_int4[1].w, dq[3], size_n,
724
zeros[3] + zero_offset);
725
755
- // half* dqh = (half*)dq;
756
- if (b_q_perm) {
757
- for (int j = 0; j < 4; j++) {
758
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
759
- b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]),
760
- __low2half(dq[2][j]), __low2half(dq[3][j]));
761
- b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]),
762
- __high2half(dq[2][j]), __high2half(dq[3][j]));
763
- }
764
- } else {
765
- for (int j = 0; j < 4; j++) {
766
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
767
- b_.set4(offset_k + lk++, n, __low2half(dq[0][j]),
768
- __low2half(dq[1][j]), __low2half(dq[2][j]),
769
- __low2half(dq[3][j]));
770
- b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
771
- __high2half(dq[1][j]), __high2half(dq[2][j]),
772
- __high2half(dq[3][j]));
773
- }
726
+ for (int j = 0; j < 4; j++) {
727
+ for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
728
+ b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
729
+ __low2half(dq[2][j]), __low2half(dq[3][j]));
730
+ b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
731
+ __high2half(dq[1][j]), __high2half(dq[2][j]),
732
+ __high2half(dq[3][j]));
733
}
734
}
735
k += 32;
@@ -778,7 +737,7 @@ __global__ void reconstruct_exllama_8bit_kernel(
737
}
738
739
__global__ void reconstruct_exllama_4bit_kernel(
781
- const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm,
740
+ const uint32_t* __restrict__ b_q_weight,
741
const uint32_t* __restrict__ b_gptq_qzeros,
742
const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
743
const int groups, const bool use_v2_format, half* __restrict__ b) {
@@ -794,14 +753,8 @@ __global__ void reconstruct_exllama_4bit_kernel(
753
754
int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
755
797
- // Preload remapping table
798
- __shared__ int perm[BLOCK_KN_SIZE];
756
auto t = threadIdx.x;
757
801
- if (b_q_perm) {
802
- if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
803
- }
804
-
758
// Column
759
int n = offset_n + t * 4;
760
if (n >= size_n) return;
@@ -828,8 +781,6 @@ __global__ void reconstruct_exllama_4bit_kernel(
781
dequant_4bit_8_prep_zero(zeros[2] + zero_offset, z1z16[2], y1y16[2]);
782
dequant_4bit_8_prep_zero(zeros[3] + zero_offset, z1z16[3], y1y16[3]);
783
831
- __syncthreads();
832
-
784
int k = offset_k;
785
int lk = 0;
786
@@ -860,25 +811,13 @@ __global__ void reconstruct_exllama_4bit_kernel(
811
false);
812
813
b_ptr += size_n;
863
- // half* dqh = (half*)dq;
864
- if (b_q_perm) {
865
- for (int j = 0; j < 4; j++) {
866
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
867
- b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]),
868
- __low2half(dq[2][j]), __low2half(dq[3][j]));
869
- b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]),
870
- __high2half(dq[2][j]), __high2half(dq[3][j]));
871
- }
872
- } else {
873
- for (int j = 0; j < 4; j++) {
874
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
875
- b_.set4(offset_k + lk++, n, __low2half(dq[0][j]),
876
- __low2half(dq[1][j]), __low2half(dq[2][j]),
877
- __low2half(dq[3][j]));
878
- b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
879
- __high2half(dq[1][j]), __high2half(dq[2][j]),
880
- __high2half(dq[3][j]));
881
- }
814
+ for (int j = 0; j < 4; j++) {
815
+ for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
816
+ b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
817
+ __low2half(dq[2][j]), __low2half(dq[3][j]));
818
+ b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
819
+ __high2half(dq[1][j]), __high2half(dq[2][j]),
820
+ __high2half(dq[3][j]));
821
}
822
}
823
k += 32;
@@ -886,7 +825,7 @@ __global__ void reconstruct_exllama_4bit_kernel(
825
}
826
827
__global__ void reconstruct_exllama_3bit_kernel(
889
- const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm,
828
+ const uint32_t* __restrict__ b_q_weight,
829
const uint32_t* __restrict__ b_gptq_qzeros,
830
const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
831
const int groups, const bool use_v2_format, half* __restrict__ b) {
@@ -902,14 +841,8 @@ __global__ void reconstruct_exllama_3bit_kernel(
841
842
int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
843
905
- // Preload remapping table
906
- __shared__ int perm[BLOCK_KN_SIZE];
844
auto t = threadIdx.x;
845
909
- if (b_q_perm) {
910
- if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
911
- }
912
-
846
// Column
847
int n = offset_n + t * 4;
848
if (n >= size_n) return;
@@ -930,8 +863,6 @@ __global__ void reconstruct_exllama_3bit_kernel(
863
b_gptq_qzeros_.item4(zeros, group, n);
864
b_gptq_scales_.item4_h2(scales, group, n);
865
933
- __syncthreads();
934
-
866
int k = offset_k;
867
int lk = 0;
868
@@ -962,24 +893,13 @@ __global__ void reconstruct_exllama_3bit_kernel(
893
dequant_3bit_32(load_int4[0].w, load_int4[1].w, load_int4[2].w, dq[3],
894
size_n, zeros[3] + zero_offset);
895
965
- if (b_q_perm) {
966
- for (int j = 0; j < 16; j++) {
967
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
968
- b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]),
969
- __low2half(dq[2][j]), __low2half(dq[3][j]));
970
- b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]),
971
- __high2half(dq[2][j]), __high2half(dq[3][j]));
972
- }
973
- } else {
974
- for (int j = 0; j < 16; j++) {
975
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
976
- b_.set4(offset_k + lk++, n, __low2half(dq[0][j]),
977
- __low2half(dq[1][j]), __low2half(dq[2][j]),
978
- __low2half(dq[3][j]));
979
- b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
980
- __high2half(dq[1][j]), __high2half(dq[2][j]),
981
- __high2half(dq[3][j]));
982
- }
896
+ for (int j = 0; j < 16; j++) {
897
+ for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
898
+ b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
899
+ __low2half(dq[2][j]), __low2half(dq[3][j]));
900
+ b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
901
+ __high2half(dq[1][j]), __high2half(dq[2][j]),
902
+ __high2half(dq[3][j]));
903
}
904
}
905
k += 32;
@@ -987,7 +907,7 @@ __global__ void reconstruct_exllama_3bit_kernel(
907
}
908
909
__global__ void reconstruct_exllama_2bit_kernel(
990
- const uint32_t* __restrict__ b_q_weight, const int* __restrict__ b_q_perm,
910
+ const uint32_t* __restrict__ b_q_weight,
911
const uint32_t* __restrict__ b_gptq_qzeros,
912
const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
913
const int groups, const bool use_v2_format, half* __restrict__ b) {
@@ -1003,14 +923,8 @@ __global__ void reconstruct_exllama_2bit_kernel(
923
924
int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
925
1006
- // Preload remapping table
1007
- __shared__ int perm[BLOCK_KN_SIZE];
926
auto t = threadIdx.x;
927
1010
- if (b_q_perm) {
1011
- if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
1012
- }
1013
-
928
// Column
929
int n = offset_n + t * 4;
930
if (n >= size_n) return;
@@ -1031,8 +945,6 @@ __global__ void reconstruct_exllama_2bit_kernel(
945
b_gptq_qzeros_.item4(zeros, group, n);
946
b_gptq_scales_.item4_h2(scales, group, n);
947
1034
- __syncthreads();
1035
-
948
int k = offset_k;
949
int lk = 0;
950
@@ -1055,25 +967,13 @@ __global__ void reconstruct_exllama_2bit_kernel(
967
dequant_2bit_16(load_int4.w, dq[3], size_n, zeros[3] + zero_offset);
968
969
b_ptr += size_n;
1058
- // half* dqh = (half*)dq;
1059
- if (b_q_perm) {
1060
- for (int j = 0; j < 8; j++) {
1061
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
1062
- b_.set4(perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]),
1063
- __low2half(dq[2][j]), __low2half(dq[3][j]));
1064
- b_.set4(perm[lk++], n, __high2half(dq[0][j]), __high2half(dq[1][j]),
1065
- __high2half(dq[2][j]), __high2half(dq[3][j]));
1066
- }
1067
- } else {
1068
- for (int j = 0; j < 8; j++) {
1069
- for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
1070
- b_.set4(offset_k + lk++, n, __low2half(dq[0][j]),
1071
- __low2half(dq[1][j]), __low2half(dq[2][j]),
1072
- __low2half(dq[3][j]));
1073
- b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
1074
- __high2half(dq[1][j]), __high2half(dq[2][j]),
1075
- __high2half(dq[3][j]));
1076
- }
970
+ for (int j = 0; j < 8; j++) {
971
+ for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
972
+ b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
973
+ __low2half(dq[2][j]), __low2half(dq[3][j]));
974
+ b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
975
+ __high2half(dq[1][j]), __high2half(dq[2][j]),
976
+ __high2half(dq[3][j]));
977
}
978
}
979
k += 32;
@@ -1082,9 +982,8 @@ __global__ void reconstruct_exllama_2bit_kernel(
982
983
void reconstruct_exllama(const uint32_t* b_q_weight,
984
const uint32_t* b_gptq_qzeros,
1085
- const half* b_gptq_scales, const int* b_q_perm,
1086
- half* out, int height, int width, int groups,
1087
- bool use_v2_format, int bit) {
985
+ const half* b_gptq_scales, half* out, int height,
986
+ int width, int groups, bool use_v2_format, int bit) {
987
dim3 blockDim, gridDim;
988
blockDim.x = BLOCK_KN_SIZE;
989
blockDim.y = 1;
@@ -1102,15 +1001,15 @@ void reconstruct_exllama(const uint32_t* b_q_weight,
1001
1002
const cudaStream_t stream = get_current_cuda_stream();
1003
reconstruct_exllama_kernel<<<gridDim, blockDim, 0, stream>>>(
1105
- b_q_weight, b_q_perm, b_gptq_qzeros, b_gptq_scales, height, width, groups,
1004
+ b_q_weight, b_gptq_qzeros, b_gptq_scales, height, width, groups,
1005
use_v2_format, out);
1006
}
1007
1008
__global__ void gemm_half_q_half_alt_4bit_kernel(
1009
const half2* __restrict__ vec, const uint32_t* __restrict__ mat,
1010
half* __restrict__ mul, const half* __restrict__ scales,
1112
- const uint32_t* __restrict__ zeros, const int* __restrict__ g_idx,
1113
- int batch, int height, int width, bool use_v2_format) {
1011
+ const uint32_t* __restrict__ zeros, int batch, int height, int width,
1012
+ int groups, bool use_v2_format) {
1013
int zero_width = width / 8;
1014
int vec_height = height * 4;
1015
const int blockwidth2 = BLOCK_KN_SIZE / 2;
@@ -1144,6 +1043,7 @@ __global__ void gemm_half_q_half_alt_4bit_kernel(
1043
1044
int i = width * h + w;
1045
int g_h = h * 8;
1046
+ int group_size = vec_height * 2 / groups;
1047
int k = 0;
1048
int z_w = w / 8;
1049
int z_mod = (w % 8) * 4;
@@ -1156,8 +1056,8 @@ __global__ void gemm_half_q_half_alt_4bit_kernel(
1056
half2 scales_tmp[4];
1057
half2 zeros_tmp[4];
1058
for (int tmp_k = 0; tmp_k < 4; tmp_k++) {
1159
- int g = g_idx[g_h + (k + tmp_k) * 2];
1160
- int g2 = g_idx[g_h + (k + tmp_k) * 2 + 1];
1059
+ int g = (g_h + (k + tmp_k) * 2) / group_size;
1060
+ int g2 = (g_h + (k + tmp_k) * 2 + 1) / group_size;
1061
half scale_f = scales[g * width + w];
1062
half scale_f2 = scales[g2 * width + w];
1063
half2 scale = __halves2half2(scale_f, scale_f2);
@@ -1209,8 +1109,8 @@ __global__ void gemm_half_q_half_alt_4bit_kernel(
1109
__global__ void gemm_half_q_half_alt_8bit_kernel(
1110
const half2* __restrict__ vec, const uint32_t* __restrict__ mat,
1111
half* __restrict__ mul, const half* __restrict__ scales,
1212
- const uint32_t* __restrict__ zeros, const int* __restrict__ g_idx,
1213
- int batch, int height, int width, bool use_v2_format) {
1112
+ const uint32_t* __restrict__ zeros, int batch, int height, int width,
1113
+ int groups, bool use_v2_format) {
1114
int zero_width = width / 4;
1115
int vec_height = height * 2;
1116
const int blockwidth2 = BLOCK_KN_SIZE / 2;
@@ -1236,6 +1136,7 @@ __global__ void gemm_half_q_half_alt_8bit_kernel(
1136
1137
int i = width * h + w;
1138
int g_h = h * 4;
1139
+ int group_size = vec_height * 2 / groups;
1140
int k = 0;
1141
int z_w = w / 4;
1142
int z_mod = (w % 4) * 8;
@@ -1248,8 +1149,8 @@ __global__ void gemm_half_q_half_alt_8bit_kernel(
1149
half2 scales_tmp[2];
1150
half2 zeros_tmp[2];
1151
for (int tmp_k = 0; tmp_k < 2; tmp_k++) {
1251
- int g = g_idx[g_h + (k + tmp_k) * 2];
1252
- int g2 = g_idx[g_h + (k + tmp_k) * 2 + 1];
1152
+ int g = (g_h + (k + tmp_k) * 2) / group_size;
1153
+ int g2 = (g_h + (k + tmp_k) * 2 + 1) / group_size;
1154
half scale_f = scales[g * width + w];
1155
half scale_f2 = scales[g2 * width + w];
1156
half2 scale = __halves2half2(scale_f, scale_f2);
@@ -1296,8 +1197,8 @@ __global__ void gemm_half_q_half_alt_8bit_kernel(
1197
1198
void gemm_half_q_half_alt(const half* a, const uint32_t* b_q_weight,
1199
const uint32_t* b_gptq_qzeros,
1299
- const half* b_gptq_scales, const int* b_g_idx,
1300
- half* c, int size_m, int size_n, int size_k,
1200
+ const half* b_gptq_scales, half* c, int size_m,
1201
+ int size_n, int size_k, int groups,
1202
bool use_v2_format, int bit) {
1203
dim3 blockDim, gridDim;
1204
blockDim.x = BLOCK_KN_SIZE;
@@ -1314,16 +1215,15 @@ void gemm_half_q_half_alt(const half* a, const uint32_t* b_q_weight,
1215
1216
const cudaStream_t stream = get_current_cuda_stream();
1217
kernel<<<gridDim, blockDim, 0, stream>>>(
1317
- (const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, b_g_idx,
1318
- size_m, size_k / 32 * bit, size_n, use_v2_format);
1218
+ (const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, size_m,
1219
+ size_k / 32 * bit, size_n, groups, use_v2_format);
1220
}
1221
1222
template <class T, int bit>
1223
__global__ void reconstruct_gptq_kernel(
1224
const uint32_t* __restrict__ w, const half* __restrict__ w_scales,
1324
- const uint32_t* __restrict__ w_zeros, const int* __restrict__ g_idx,
1325
- const int height, const int width, const int group,
1326
- const bool use_v2_format, half* __restrict__ out) {
1225
+ const uint32_t* __restrict__ w_zeros, const int height, const int width,
1226
+ const int num_groups, const bool use_v2_format, half* __restrict__ out) {
1227
// Start of block
1228
1229
auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
@@ -1333,8 +1233,8 @@ __global__ void reconstruct_gptq_kernel(
1233
// Views
1234
1235
MatrixView_half_rw out_(out, height, width);
1336
- MatrixView_half w_scales_(w_scales, group, width);
1337
- T w_zeros_(w_zeros, group, width);
1236
+ MatrixView_half w_scales_(w_scales, num_groups, width);
1237
+ T w_zeros_(w_zeros, num_groups, width);
1238
1239
// GPTQv2 and GPTQv1 handles zero points differently
1240
int zero_offset = use_v2_format ? 0 : 1;
@@ -1344,7 +1244,7 @@ __global__ void reconstruct_gptq_kernel(
1244
1245
#pragma unroll
1246
for (int s = 0; s < 32; s += bit) {
1347
- int group = g_idx[row + s / bit];
1247
+ int group = (row + s / bit) / (height / num_groups);
1248
half w_scale = w_scales_.item(group, column);
1249
uint32_t w_zero = w_zeros_.item(group, column) + zero_offset;
1250
half w_item =
@@ -1357,9 +1257,8 @@ __global__ void reconstruct_gptq_kernel(
1257
1258
__global__ void reconstruct_gptq_3bit_kernel(
1259
const uint32_t* __restrict__ w, const half* __restrict__ w_scales,
1360
- const uint32_t* __restrict__ w_zeros, const int* __restrict__ g_idx,
1361
- const int height, const int width, const int group,
1362
- const bool use_v2_format, half* __restrict__ out) {
1260
+ const uint32_t* __restrict__ w_zeros, const int height, const int width,
1261
+ const int num_groups, const bool use_v2_format, half* __restrict__ out) {
1262
// Start of block
1263
auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1264
auto row = blockIdx.y * 32;
@@ -1368,8 +1267,8 @@ __global__ void reconstruct_gptq_3bit_kernel(
1267
// Views
1268
1269
MatrixView_half_rw out_(out, height, width);
1371
- MatrixView_half w_scales_(w_scales, group, width);
1372
- MatrixView_q3_row w_zeros_(w_zeros, group, width);
1270
+ MatrixView_half w_scales_(w_scales, num_groups, width);
1271
+ MatrixView_q3_row w_zeros_(w_zeros, num_groups, width);
1272
1273
// GPTQv2 and GPTQv1 handles zero points differently
1274
int zero_offset = use_v2_format ? 0 : 1;
@@ -1381,7 +1280,7 @@ __global__ void reconstruct_gptq_3bit_kernel(
1280
1281
#pragma unroll
1282
for (int i = 0; i < 32; i += 1) {
1384
- int group = g_idx[row + i];
1283
+ int group = (row + i) / (height / num_groups);
1284
half w_scale = w_scales_.item(group, column);
1285
uint32_t w_zero = w_zeros_.item(group, column) + zero_offset;
1286
int w_item;
@@ -1402,9 +1301,8 @@ __global__ void reconstruct_gptq_3bit_kernel(
1301
}
1302
1303
void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros,
1405
- const half* b_gptq_scales, const int* b_g_idx, half* out,
1406
- int height, int width, int groups, bool use_v2_format,
1407
- int bit) {
1304
+ const half* b_gptq_scales, half* out, int height,
1305
+ int width, int groups, bool use_v2_format, int bit) {
1306
dim3 blockDim, gridDim;
1307
blockDim.x = BLOCK_KN_SIZE;
1308
blockDim.y = 1;
@@ -1423,17 +1321,16 @@ void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros,
1321
1322
const cudaStream_t stream = get_current_cuda_stream();
1323
kernel<<<gridDim, blockDim, 0, stream>>>(b_q_weight, b_gptq_scales,
1426
- b_gptq_qzeros, b_g_idx, height,
1427
- width, groups, use_v2_format, out);
1324
+ b_gptq_qzeros, height, width, groups,
1325
+ use_v2_format, out);
1326
}
1327
1328
void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a,
1329
const uint32_t* b_q_weight,
1330
const uint32_t* b_gptq_qzeros,
1433
- const half* b_gptq_scales, const int* b_g_idx,
1434
- half* c, half* temp_dq, int size_m, int size_n,
1435
- int size_k, int groups, bool use_exllama,
1436
- bool use_v2_format, int bit) {
1331
+ const half* b_gptq_scales, half* c, half* temp_dq,
1332
+ int size_m, int size_n, int size_k, int groups,
1333
+ bool use_exllama, bool use_v2_format, int bit) {
1334
bool use_reconstruct;
1335
if (use_exllama) {
1336
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,
1343
if (use_reconstruct) {
1344
// Reconstruct FP16 matrix, then cuBLAS
1345
if (use_exllama) {
1449
- reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
1450
- temp_dq, size_k, size_n, groups, use_v2_format, bit);
1346
+ reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, temp_dq,
1347
+ size_k, size_n, groups, use_v2_format, bit);
1348
} else {
1452
- reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
1453
- temp_dq, size_k, size_n, groups, use_v2_format, bit);
1349
+ reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, temp_dq,
1350
+ size_k, size_n, groups, use_v2_format, bit);
1351
}
1352
1353
const half alpha = __float2half(1.0f);
@@ -1464,20 +1361,20 @@ void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a,
1361
int last_chunk_size = size_m - last_chunk;
1362
1363
if (max_chunks) {
1467
- gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales,
1468
- b_g_idx, c, last_chunk, size_n, size_k,
1469
- BLOCK_M_SIZE_MAX, groups, use_v2_format, bit);
1364
+ gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c,
1365
+ last_chunk, size_n, size_k, BLOCK_M_SIZE_MAX,
1366
+ groups, use_v2_format, bit);
1367
}
1368
1369
if (last_chunk_size) {
1370
gemm_half_q_half_cuda_part(
1371
a + last_chunk * size_k, b_q_weight, b_gptq_qzeros, b_gptq_scales,
1475
- b_g_idx, c + last_chunk * size_n, last_chunk_size, size_n, size_k,
1372
+ c + last_chunk * size_n, last_chunk_size, size_n, size_k,
1373
last_chunk_size, groups, use_v2_format, bit);
1374
}
1375
} else {
1479
- gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
1480
- c, size_m, size_n, size_k, use_v2_format, bit);
1376
+ gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m,
1377
+ size_n, size_k, groups, use_v2_format, bit);
1378
}
1379
}
1380
@@ -1533,214 +1430,8 @@ __global__ void shuffle_3bit_kernel(uint32_t* __restrict__ b_q_weight,
1430
}
1431
}
1432
1536
-__global__ void make_sequential_4bit_kernel(const uint32_t* __restrict__ w,
1537
- uint32_t* __restrict__ w_new,
1538
- const int* __restrict__ q_perm,
1539
- const int w_width) {
1540
- const uint64_t* w2 = (uint64_t*)w;
1541
- uint64_t* w_new2 = (uint64_t*)w_new;
1542
- int w2_stride = w_width >> 1;
1543
- auto w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1544
- if (w2_column >= w2_stride) return;
1545
- auto w_new2_row = blockIdx.y;
1546
- int q_perm_idx = w_new2_row << 3;
1547
- uint64_t dst = 0;
1548
-
1549
-#pragma unroll
1550
- for (int i = 0; i < 8; i++) {
1551
- int source_row = q_perm[q_perm_idx++];
1552
-
1553
- int w2_row = source_row >> 3;
1554
- int w2_subrow = source_row & 0x07;
1555
- int w2_row_shift = w2_subrow << 2;
1556
- int wnew2_row_shift = i << 2;
1557
-
1558
- uint64_t src = w2[w2_row * w2_stride + w2_column];
1559
- src >>= w2_row_shift;
1560
- src &= 0x0000000f0000000f;
1561
- src <<= wnew2_row_shift;
1562
- dst |= src;
1563
- }
1564
- w_new2[w_new2_row * w2_stride + w2_column] = dst;
1565
-}
1566
-
1567
-__global__ void make_sequential_2bit_kernel(const uint32_t* __restrict__ w,
1568
- uint32_t* __restrict__ w_new,
1569
- const int* __restrict__ q_perm,
1570
- const int w_width) {
1571
- const uint64_t* w2 = (uint64_t*)w;
1572
- uint64_t* w_new2 = (uint64_t*)w_new;
1573
- int w2_stride = w_width >> 1;
1574
- auto w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1575
- if (w2_column >= w2_stride) return;
1576
- auto w_new2_row = blockIdx.y;
1577
- int q_perm_idx = w_new2_row << 4;
1578
- uint64_t dst = 0;
1579
-
1580
-#pragma unroll
1581
- for (int i = 0; i < 16; i++) {
1582
- int source_row = q_perm[q_perm_idx++];
1583
-
1584
- int w2_row = source_row >> 4;
1585
- int w2_subrow = source_row & 0x0f;
1586
- int w2_row_shift = w2_subrow << 1;
1587
- int wnew2_row_shift = i << 1;
1588
-
1589
- uint64_t src = w2[w2_row * w2_stride + w2_column];
1590
- src >>= w2_row_shift;
1591
- src &= 0x0000000300000003;
1592
- src <<= wnew2_row_shift;
1593
- dst |= src;
1594
- }
1595
- w_new2[w_new2_row * w2_stride + w2_column] = dst;
1596
-}
1597
-
1598
-__global__ void make_sequential_3bit_kernel(const uint32_t* __restrict__ w,
1599
- uint32_t* __restrict__ w_new,
1600
- const int* __restrict__ q_perm,
1601
- const int w_width) {
1602
- auto w_column = THREADS_X * blockIdx.x + threadIdx.x;
1603
- if (w_column >= w_width) return;
1604
- auto w_new_row = blockIdx.y * 3;
1605
- auto q_perm_idx = blockIdx.y << 5;
1606
- uint32_t dst[3] = {0, 0, 0};
1607
-
1608
-#pragma unroll
1609
- for (int i = 0; i < 32; i++) {
1610
- int source_row = q_perm[q_perm_idx++];
1611
- int z_w = (source_row / 32) * 3;
1612
- int z_mod = source_row % 32;
1613
- int z_bit;
1614
-
1615
- if (z_mod != 10) {
1616
- if (z_mod != 21) {
1617
- z_bit = z_mod;
1618
- if (z_bit > 21) {
1619
- z_bit *= 3;
1620
- z_bit -= 64;
1621
- z_w += 2;
1622
- } else if (z_bit > 10) {
1623
- z_bit *= 3;
1624
- z_bit -= 32;
1625
- z_w += 1;
1626
- } else {
1627
- z_bit *= 3;
1628
- }
1629
- } else {
1630
- z_w += 1;
1631
- }
1632
- }
1633
-
1634
- uint64_t src;
1635
- if (z_mod == 10) {
1636
- src = (w[z_w * w_width + w_column] >> 30) |
1637
- ((w[(z_w + 1) * w_width + w_column] << 2) & 0x4);
1638
- } else if (z_mod == 21) {
1639
- src = (w[z_w * w_width + w_column] >> 31) |
1640
- ((w[(z_w + 1) * w_width + w_column] << 1) & 0x6);
1641
- } else {
1642
- src = w[z_w * w_width + w_column];
1643
- src >>= z_bit;
1644
- src &= 0x07;
1645
- }
1646
-
1647
- z_w = 0;
1648
- if (i != 10) {
1649
- if (i != 21) {
1650
- z_bit = i;
1651
- if (z_bit > 21) {
1652
- z_bit *= 3;
1653
- z_bit -= 64;
1654
- z_w += 2;
1655
- } else if (z_bit > 10) {
1656
- z_bit *= 3;
1657
- z_bit -= 32;
1658
- z_w += 1;
1659
- } else {
1660
- z_bit *= 3;
1661
- }
1662
- } else {
1663
- z_w += 1;
1664
- }
1665
- }
1666
- if (i == 10) {
1667
- dst[z_w] |= (src & 0x03) << 30;
1668
- dst[z_w + 1] |= ((src & 0x4) >> 2);
1669
- } else if (i == 21) {
1670
- dst[z_w] |= (src & 0x01) << 31;
1671
- dst[z_w + 1] |= ((src & 0x6) >> 1);
1672
- } else {
1673
- dst[z_w] |= (src << z_bit);
1674
- }
1675
- }
1676
- w_new[w_new_row * w_width + w_column] = dst[0];
1677
- w_new[(w_new_row + 1) * w_width + w_column] = dst[1];
1678
- w_new[(w_new_row + 2) * w_width + w_column] = dst[2];
1679
-}
1680
-
1681
-__global__ void make_sequential_8bit_kernel(const uint32_t* __restrict__ w,
1682
- uint32_t* __restrict__ w_new,
1683
- const int* __restrict__ q_perm,
1684
- const int w_width) {
1685
- const uint64_t* w2 = (uint64_t*)w;
1686
- uint64_t* w_new2 = (uint64_t*)w_new;
1687
- int w2_stride = w_width >> 1;
1688
- auto w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1689
- if (w2_column >= w2_stride) return;
1690
- auto w_new2_row = blockIdx.y;
1691
- int q_perm_idx = w_new2_row << 2;
1692
- uint64_t dst = 0;
1693
-
1694
-#pragma unroll
1695
- for (int i = 0; i < 4; i++) {
1696
- int source_row = q_perm[q_perm_idx++];
1697
-
1698
- int w2_row = source_row >> 2;
1699
- int w2_subrow = source_row & 0x03;
1700
- int w2_row_shift = w2_subrow << 3;
1701
- int wnew2_row_shift = i << 3;
1702
-
1703
- uint64_t src = w2[w2_row * w2_stride + w2_column];
1704
- src >>= w2_row_shift;
1705
- src &= 0x000000ff000000ff;
1706
- src <<= wnew2_row_shift;
1707
- dst |= src;
1708
- }
1709
- w_new2[w_new2_row * w2_stride + w2_column] = dst;
1710
-}
1711
-
1712
-void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height,
1713
- int width, int bit) {
1714
- if (q_perm) {
1715
- uint32_t* new_qweight = NULL;
1716
- cudaMalloc(&new_qweight, height / 32 * bit * width * sizeof(uint32_t));
1717
-
1718
- dim3 blockDim, gridDim;
1719
- blockDim.x = THREADS_X;
1720
- blockDim.y = 1;
1721
- gridDim.x = DIVIDE(width, THREADS_X);
1722
- gridDim.y = height / 32 * bit;
1723
-
1724
- auto kernel = make_sequential_4bit_kernel;
1725
- if (bit == 2) {
1726
- kernel = make_sequential_2bit_kernel;
1727
- } else if (bit == 3) {
1728
- kernel = make_sequential_3bit_kernel;
1729
- gridDim.y = height / 32;
1730
- } else if (bit == 8) {
1731
- kernel = make_sequential_8bit_kernel;
1732
- }
1733
- const cudaStream_t stream = get_current_cuda_stream();
1734
- kernel<<<gridDim, blockDim, 0, stream>>>(q_weight, new_qweight, q_perm,
1735
- width);
1736
- // Replace qweights
1737
- cudaMemcpyAsync(q_weight, new_qweight,
1738
- height / 32 * bit * width * sizeof(uint32_t),
1739
- cudaMemcpyDeviceToDevice);
1740
- // Cleanup
1741
- cudaDeviceSynchronize();
1742
- cudaFree(new_qweight);
1743
- }
1433
+void shuffle_exllama_weight(uint32_t* q_weight, int height, int width,
1434
+ int bit) {
1435
dim3 blockDim, gridDim;
1436
blockDim.x = THREADS_X;
1437
blockDim.y = 1;
@@ -1765,10 +1456,17 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
1456
torch::stable::Tensor b_q_weight,
1457
torch::stable::Tensor b_gptq_qzeros,
1458
torch::stable::Tensor b_gptq_scales,
1768
- torch::stable::Tensor b_g_idx, bool use_exllama,
1769
- bool use_v2_format, int64_t bit) {
1459
+ bool use_exllama, bool use_v2_format,
1460
+ int64_t bit) {
1461
const torch::stable::accelerator::DeviceGuard device_guard(
1462
a.get_device_index());
1463
+ const int64_t num_groups = b_gptq_qzeros.size(0);
1464
+ STD_TORCH_CHECK(num_groups > 0, "GPTQ requires at least one scale group");
1465
+ STD_TORCH_CHECK(b_gptq_scales.size(0) == num_groups,
1466
+ "GPTQ scales and zero-points must have the same number of "
1467
+ "groups");
1468
+ STD_TORCH_CHECK(a.size(1) % num_groups == 0, "Input size K = ", a.size(1),
1469
+ " is not divisible by the number of groups = ", num_groups);
1470
auto c = torch::stable::new_zeros(a, {a.size(0), b_q_weight.size(1)});
1471
auto temp_dq =
1472
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,
1476
get_current_cuda_blas_handle(), (const half*)a.data_ptr(),
1477
(const uint32_t*)b_q_weight.data_ptr(),
1478
(const uint32_t*)b_gptq_qzeros.data_ptr(),
1781
- (const half*)b_gptq_scales.data_ptr(),
1782
- b_g_idx.device().type() == torch::stable::DeviceType::Meta
1783
- ? NULL
1784
- : (const int*)b_g_idx.data_ptr(),
1785
- (half*)c.data_ptr(), (half*)temp_dq.data_ptr(),
1786
- c.size(0), // m
1787
- c.size(1), // n
1788
- a.size(1), // k
1789
- b_gptq_qzeros.size(0), // group number
1790
- use_exllama, use_v2_format, bit);
1479
+ (const half*)b_gptq_scales.data_ptr(), (half*)c.data_ptr(),
1480
+ (half*)temp_dq.data_ptr(),
1481
+ c.size(0), // m
1482
+ c.size(1), // n
1483
+ a.size(1), // k
1484
+ num_groups, use_exllama, use_v2_format, bit);
1485
return c;
1486
}
1487
1794
-void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
1795
- int64_t bit) {
1488
+void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit) {
1489
const torch::stable::accelerator::DeviceGuard device_guard(
1490
q_weight.get_device_index());
1798
- vllm::gptq::shuffle_exllama_weight(
1799
- (uint32_t*)q_weight.data_ptr(),
1800
- q_perm.device().type() == torch::stable::DeviceType::Meta ||
1801
- q_perm.numel() == 0
1802
- ? NULL
1803
- : (int*)q_perm.data_ptr(),
1804
- q_weight.size(0) * 32 / bit, q_weight.size(1), bit);
1491
+ vllm::gptq::shuffle_exllama_weight((uint32_t*)q_weight.data_ptr(),
1492
+ q_weight.size(0) * 32 / bit,
1493
+ q_weight.size(1), bit);
1494
}
csrc/libtorch_stable/quantization/marlin/generate_kernels.py
+2
-2
@@ -75,14 +75,14 @@ QUANT_CONFIGS = [
75
"b_type": "kU4B8",
76
"thread_configs": THREAD_CONFIGS,
77
"thread_m_blocks": THREAD_M_BLOCKS,
78
- "group_blocks": [-1, 0, 2, 4, 8],
78
+ "group_blocks": [-1, 2, 4, 8],
79
},
80
# GPTQ-INT8
81
{
82
"b_type": "kU8B128",
83
"thread_configs": THREAD_CONFIGS,
84
"thread_m_blocks": THREAD_M_BLOCKS,
85
- "group_blocks": [-1, 0, 2, 4, 8],
85
+ "group_blocks": [-1, 2, 4, 8],
86
},
87
# FP8
88
{
csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu
+44
-122
@@ -11,11 +11,9 @@
11
12
namespace marlin {
13
14
-template <int const num_threads, int const num_bits, bool const has_perm,
15
- bool is_a_8bit>
14
+template <int const num_threads, int const num_bits, bool is_a_8bit>
15
__global__ void gptq_marlin_repack_kernel(
17
- uint32_t const* __restrict__ b_q_weight_ptr,
18
- uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
16
+ uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr,
17
int size_k, int size_n) {
18
constexpr int pack_factor = 32 / num_bits;
19
@@ -44,31 +42,14 @@ __global__ void gptq_marlin_repack_kernel(
42
43
extern __shared__ int4 sh[];
44
47
- constexpr int perm_size = target_tile_k_size / 4;
48
-
49
- int4* sh_perm_ptr = sh;
50
- int4* sh_pipe_ptr = sh_perm_ptr;
51
- if constexpr (has_perm) {
52
- sh_pipe_ptr += perm_size;
53
- }
45
+ int4* sh_pipe_ptr = sh;
46
47
constexpr int tile_ints = target_tile_k_size / pack_factor;
48
49
constexpr int stage_n_threads = target_tile_n_size / 4;
58
- constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints;
50
+ constexpr int stage_k_threads = tile_ints;
51
constexpr int stage_size = stage_k_threads * stage_n_threads;
52
61
- auto load_perm_to_shared = [&](int k_tile_id) {
62
- int first_k_int4 = (k_tile_id * target_tile_k_size) / 4;
63
-
64
- int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
65
-
66
- if (threadIdx.x < perm_size) {
67
- sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x];
68
- }
69
- __syncthreads();
70
- };
71
-
53
auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
54
if (n_tile_id >= n_tiles) {
55
cp_async_fence();
@@ -79,36 +60,17 @@ __global__ void gptq_marlin_repack_kernel(
60
61
int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
62
82
- if constexpr (has_perm) {
83
- if (threadIdx.x < stage_size) {
84
- auto k_id = threadIdx.x / stage_n_threads;
85
- auto n_id = threadIdx.x % stage_n_threads;
63
+ if (threadIdx.x < stage_size) {
64
+ auto k_id = threadIdx.x / stage_n_threads;
65
+ auto n_id = threadIdx.x % stage_n_threads;
66
87
- uint32_t const* sh_perm_int_ptr =
88
- reinterpret_cast<uint32_t const*>(sh_perm_ptr);
67
+ int first_k = k_tile_id * target_tile_k_size;
68
+ int first_k_packed = first_k / pack_factor;
69
90
- int src_k = sh_perm_int_ptr[k_id];
91
- int src_k_packed = src_k / pack_factor;
92
-
93
- cp_async4(
94
- &sh_ptr[k_id * stage_n_threads + n_id],
95
- reinterpret_cast<int4 const*>(&(
96
- b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)])));
97
- }
98
-
99
- } else {
100
- if (threadIdx.x < stage_size) {
101
- auto k_id = threadIdx.x / stage_n_threads;
102
- auto n_id = threadIdx.x % stage_n_threads;
103
-
104
- int first_k = k_tile_id * target_tile_k_size;
105
- int first_k_packed = first_k / pack_factor;
106
-
107
- cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
108
- reinterpret_cast<int4 const*>(
109
- &(b_q_weight_ptr[(first_k_packed + k_id) * size_n +
110
- first_n + (n_id * 4)])));
111
- }
70
+ cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
71
+ reinterpret_cast<int4 const*>(
72
+ &(b_q_weight_ptr[(first_k_packed + k_id) * size_n +
73
+ first_n + (n_id * 4)])));
74
}
75
76
cp_async_fence();
@@ -139,56 +101,34 @@ __global__ void gptq_marlin_repack_kernel(
101
int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
102
uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
103
142
- uint32_t* sh_perm_int_ptr = reinterpret_cast<uint32_t*>(sh_perm_ptr);
143
-
104
uint32_t vals[8];
105
146
- if constexpr (has_perm) {
147
- static_assert(!is_a_8bit);
148
- for (int i = 0; i < 4; i++) {
149
- int k_idx = tc_row + tc_offsets[i];
150
-
151
- uint32_t src_k = sh_perm_int_ptr[k_idx];
152
- uint32_t src_k_pos = src_k % pack_factor;
153
-
154
- uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n];
155
- uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask;
156
-
157
- uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8];
158
- uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask;
159
-
160
- vals[i] = b1_cur_val;
161
- vals[4 + i] = b2_cur_val;
162
- }
163
-
164
- } else {
165
- uint32_t b1_vals[tile_ints];
166
- uint32_t b2_vals[tile_ints];
106
+ uint32_t b1_vals[tile_ints];
107
+ uint32_t b2_vals[tile_ints];
108
109
#pragma unroll
169
- for (int i = 0; i < tile_ints; i++) {
170
- if constexpr (is_a_8bit) {
171
- b1_vals[i] =
172
- sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8];
173
- } else {
174
- b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
175
- b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
176
- }
110
+ for (int i = 0; i < tile_ints; i++) {
111
+ if constexpr (is_a_8bit) {
112
+ b1_vals[i] =
113
+ sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8];
114
+ } else {
115
+ b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
116
+ b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
117
}
118
+ }
119
120
#pragma unroll
180
- for (int i = 0; i < 4; i++) {
181
- int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]);
182
- int cur_int = cur_elem / pack_factor;
183
- int cur_pos = cur_elem % pack_factor;
184
-
185
- vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
186
- if constexpr (is_a_8bit)
187
- vals[4 + i] =
188
- (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask;
189
- else
190
- vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
191
- }
121
+ for (int i = 0; i < 4; i++) {
122
+ int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]);
123
+ int cur_int = cur_elem / pack_factor;
124
+ int cur_pos = cur_elem % pack_factor;
125
+
126
+ vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
127
+ if constexpr (is_a_8bit)
128
+ vals[4 + i] =
129
+ (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask;
130
+ else
131
+ vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
132
}
133
134
constexpr int tile_size =
@@ -248,10 +188,6 @@ __global__ void gptq_marlin_repack_kernel(
188
for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
189
int n_tile_id = 0;
190
251
- if constexpr (has_perm) {
252
- load_perm_to_shared(k_tile_id);
253
- }
254
-
191
start_pipes(k_tile_id, n_tile_id);
192
193
while (n_tile_id < n_tiles) {
@@ -269,21 +205,19 @@ __global__ void gptq_marlin_repack_kernel(
205
206
} // namespace marlin
207
272
-#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \
273
- else if (num_bits == NUM_BITS && has_perm == HAS_PERM && \
274
- is_a_8bit == IS_A_8BIT) { \
208
+#define CALL_IF(NUM_BITS, IS_A_8BIT) \
209
+ else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) { \
210
cudaFuncSetAttribute( \
211
marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
277
- HAS_PERM, IS_A_8BIT>, \
212
+ IS_A_8BIT>, \
213
cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
214
marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
280
- HAS_PERM, IS_A_8BIT> \
215
+ IS_A_8BIT> \
216
<<<blocks, marlin::repack_threads, max_shared_mem, stream>>>( \
282
- b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \
217
+ b_q_weight_ptr, out_ptr, size_k, size_n); \
218
}
219
220
torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight,
286
- torch::stable::Tensor& perm,
221
int64_t size_k, int64_t size_n,
222
int64_t num_bits, bool is_a_8bit) {
223
// Verify compatibility with marlin tile of 16x64
@@ -311,11 +245,6 @@ torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight,
245
b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int,
246
"b_q_weight type is not kInt");
247
314
- STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU");
315
- STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
316
- STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int,
317
- "perm type is not at::kInt");
318
-
248
const int32_t device_index = b_q_weight.get_device_index();
249
torch::stable::accelerator::DeviceGuard device_guard(device_index);
250
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,
254
{size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
255
b_q_weight.scalar_type(), std::nullopt, b_q_weight.device());
256
328
- // Detect if there is act_order
329
- bool has_perm = perm.size(0) != 0;
330
-
257
// Get ptrs
258
uint32_t const* b_q_weight_ptr =
259
reinterpret_cast<uint32_t const*>(b_q_weight.const_data_ptr());
334
- uint32_t const* perm_ptr =
335
- reinterpret_cast<uint32_t const*>(perm.const_data_ptr());
260
uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.mutable_data_ptr());
261
262
int blocks;
@@ -345,17 +269,15 @@ torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight,
269
270
if (false) {
271
}
348
- CALL_IF(4, false, false)
349
- CALL_IF(4, true, false)
350
- CALL_IF(8, false, false)
351
- CALL_IF(8, true, false)
272
+ CALL_IF(4, false)
273
+ CALL_IF(8, false)
274
353
- CALL_IF(4, false, true)
354
- CALL_IF(8, false, true)
275
+ CALL_IF(4, true)
276
+ CALL_IF(8, true)
277
278
else {
279
STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
358
- ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
280
+ ", is_a_8bit = ", is_a_8bit);
281
}
282
283
return out;
csrc/libtorch_stable/quantization/marlin/kernel.h
+10
-11
@@ -7,17 +7,16 @@
7
#include "marlin_dtypes.cuh"
8
#include "core/scalar_type.hpp"
9
10
-#define MARLIN_KERNEL_PARAMS \
11
- const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
- int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
- const int4 *__restrict__ b_bias_ptr, \
14
- const float *__restrict__ a_scales_ptr, \
15
- const int4 *__restrict__ scales_ptr, \
16
- const float *__restrict__ global_scale_ptr, \
17
- const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
- int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
19
- bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \
20
- int max_shared_mem
10
+#define MARLIN_KERNEL_PARAMS \
11
+ const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
+ int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
+ const int4 *__restrict__ b_bias_ptr, \
14
+ const float *__restrict__ a_scales_ptr, \
15
+ const int4 *__restrict__ scales_ptr, \
16
+ const float *__restrict__ global_scale_ptr, \
17
+ const int4 *__restrict__ zp_ptr, int prob_m, int prob_n, int prob_k, \
18
+ int lda, int *locks, bool has_bias, bool use_atomic_add, \
19
+ bool use_fp32_reduce, int max_shared_mem
20
21
namespace MARLIN_NAMESPACE_NAME {
22
template <const vllm::ScalarTypeId a_type_id, // A ScalarType id
csrc/libtorch_stable/quantization/marlin/marlin.cu
+68
-233
@@ -47,11 +47,6 @@ using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
47
48
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
49
50
-__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
51
- int const* __restrict__ perm_int_ptr,
52
- int4* __restrict__ out_int4_ptr, int size_m,
53
- int size_k, int lda, int block_rows) {}
54
-
50
} // namespace marlin
51
52
torch::stable::Tensor marlin_gemm(
@@ -62,11 +57,9 @@ torch::stable::Tensor marlin_gemm(
57
std::optional<torch::stable::Tensor> const& a_scales_or_none,
58
std::optional<torch::stable::Tensor> const& global_scale_or_none,
59
std::optional<torch::stable::Tensor> const& b_zeros_or_none,
65
- std::optional<torch::stable::Tensor> const& g_idx_or_none,
66
- std::optional<torch::stable::Tensor> const& perm_or_none,
60
torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id,
68
- int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full,
69
- bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
61
+ int64_t size_m, int64_t size_n, int64_t size_k, bool use_atomic_add,
62
+ bool use_fp32_reduce, bool is_zp_float) {
63
STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
64
"marlin_gemm(..) requires CUDA_ARCH >= 7.5");
65
return torch::stable::empty({1, 1});
@@ -74,62 +67,6 @@ torch::stable::Tensor marlin_gemm(
67
68
#else
69
77
-// For a given "a" of size [M,K] performs a permutation of the K columns based
78
-// on the given "perm" indices.
79
-__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
80
- int const* __restrict__ perm_int_ptr,
81
- int4* __restrict__ out_int4_ptr, int size_m,
82
- int size_k, int lda, int block_rows) {
83
- auto start_row = block_rows * blockIdx.x;
84
- int finish_row = start_row + block_rows;
85
- if (finish_row > size_m) {
86
- finish_row = size_m;
87
- }
88
- int cur_block_rows = finish_row - start_row;
89
-
90
- int input_row_stride = lda * sizeof(half) / 16;
91
- int output_row_stride = size_k * sizeof(half) / 16;
92
-
93
- auto permute_row = [&](int row) {
94
- int iters = size_k / default_threads;
95
- int rest = size_k % default_threads;
96
-
97
- int input_offset = row * input_row_stride;
98
- int output_offset = row * output_row_stride;
99
-
100
- half const* a_row_half =
101
- reinterpret_cast<half const*>(a_int4_ptr + input_offset);
102
- half* out_half = reinterpret_cast<half*>(out_int4_ptr + output_offset);
103
-
104
- int base_k = 0;
105
-
106
- for (int i = 0; i < iters; i++) {
107
- auto cur_k = base_k + threadIdx.x;
108
- int src_pos = perm_int_ptr[cur_k];
109
-
110
- out_half[cur_k] = a_row_half[src_pos];
111
-
112
- base_k += default_threads;
113
- }
114
-
115
- if (rest) {
116
- if (threadIdx.x < rest) {
117
- auto cur_k = base_k + threadIdx.x;
118
- int src_pos = perm_int_ptr[cur_k];
119
-
120
- out_half[cur_k] = a_row_half[src_pos];
121
- }
122
- }
123
- };
124
-
125
- for (int i = 0; i < cur_block_rows; i++) {
126
- int cur_row = start_row + i;
127
- if (cur_row < size_m) {
128
- permute_row(cur_row);
129
- }
130
- }
131
-}
132
-
70
typedef struct {
71
int thread_k;
72
int thread_n;
@@ -159,9 +96,7 @@ typedef struct {
96
97
int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
98
int prob_n, int prob_k, int num_bits, int group_size,
162
- bool has_act_order, bool is_k_full, int stages) {
163
- bool cache_scales_chunk = has_act_order && !is_k_full;
164
-
99
+ int stages) {
100
int tb_n = th_config.thread_n;
101
int tb_k = th_config.thread_k;
102
@@ -169,29 +104,18 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
104
int tb_groups;
105
if (group_size == -1) {
106
tb_groups = 1;
172
- } else if (group_size == 0) {
173
- tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size
107
} else {
108
tb_groups = div_ceil(tb_k, group_size);
109
}
110
178
- if (cache_scales_chunk) {
179
- int load_groups =
180
- tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K
181
- load_groups = max(load_groups, 32); // We load at least 32 scale groups
182
- return load_groups * tb_n * 2;
183
- } else {
184
- int tb_scales = tb_groups * tb_n * 2;
185
-
186
- return tb_scales * stages;
187
- }
111
+ int tb_scales = tb_groups * tb_n * 2;
112
+ return tb_scales * stages;
113
}
114
115
int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
116
int prob_m, int prob_n, int prob_k, int num_bits,
192
- int group_size, bool has_act_order, bool is_k_full,
193
- int has_zp, bool is_zp_float, bool is_a_8bit,
194
- int stages) {
117
+ int group_size, int has_zp, bool is_zp_float,
118
+ bool is_a_8bit, int stages) {
119
int pack_factor = 32 / num_bits;
120
121
// Get B size
@@ -206,10 +130,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
130
(sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size;
131
tmp_size = max(max(sh_b_size, sh_red_size), tmp_size);
132
209
- int sh_s_size =
210
- get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
211
- group_size, has_act_order, is_k_full, stages);
212
- int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0;
133
+ int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k,
134
+ num_bits, group_size, stages);
135
int sh_zp_size = 0;
136
if (has_zp) {
137
if (is_zp_float)
@@ -220,17 +142,15 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
142
sh_zp_size = sh_s_size / 2;
143
}
144
223
- int total_size =
224
- tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size;
145
+ int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size;
146
147
return total_size;
148
}
149
150
bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
151
int prob_m, int prob_n, int prob_k, int num_bits,
231
- int group_size, bool has_act_order, bool is_k_full,
232
- int has_zp, bool is_zp_float, bool is_a_8bit, int stages,
233
- int max_shared_mem) {
152
+ int group_size, int has_zp, bool is_zp_float,
153
+ bool is_a_8bit, int stages, int max_shared_mem) {
154
// Sanity
155
if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
156
th_config.num_threads == -1) {
@@ -255,16 +175,18 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
175
// Check that pipeline fits into cache
176
int cache_size = get_kernel_cache_size(
177
th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size,
258
- has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
178
+ has_zp, is_zp_float, is_a_8bit, stages);
179
return cache_size <= max_shared_mem;
180
}
181
262
-MarlinFuncPtr get_marlin_kernel(
263
- const vllm::ScalarType a_type, const vllm::ScalarType b_type,
264
- const vllm::ScalarType c_type, const vllm::ScalarType s_type,
265
- int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
266
- bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
267
- int threads, bool is_zp_float, int stages) {
182
+MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType a_type,
183
+ const vllm::ScalarType b_type,
184
+ const vllm::ScalarType c_type,
185
+ const vllm::ScalarType s_type,
186
+ int thread_m_blocks, int thread_n_blocks,
187
+ int thread_k_blocks, bool m_block_size_8,
188
+ bool has_zp, int group_blocks, int threads,
189
+ bool is_zp_float, int stages) {
190
int num_bits = b_type.size_bits();
191
auto kernel = MarlinDefault;
192
@@ -277,9 +199,8 @@ exec_config_t determine_exec_config(
199
const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
200
const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
201
int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8,
280
- int num_bits, int group_size, bool has_act_order, bool is_k_full,
281
- bool has_zp, bool is_zp_float, int is_a_8bit, int stages,
282
- int max_shared_mem, int sms) {
202
+ int num_bits, int group_size, bool has_zp, bool is_zp_float, int is_a_8bit,
203
+ int stages, int max_shared_mem, int sms) {
204
exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
205
thread_config_t* thread_configs = thread_m_blocks > 1
206
? large_batch_thread_configs
@@ -293,27 +214,21 @@ exec_config_t determine_exec_config(
214
thread_config_t th_config = thread_configs[i];
215
216
if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k,
296
- num_bits, group_size, has_act_order, is_k_full, has_zp,
297
- is_zp_float, is_a_8bit, stages,
298
- max_shared_mem - 512)) {
217
+ num_bits, group_size, has_zp, is_zp_float, is_a_8bit,
218
+ stages, max_shared_mem - 512)) {
219
continue;
220
}
221
302
- int cache_size = get_kernel_cache_size(th_config, thread_m_blocks, prob_m,
303
- prob_n, prob_k, num_bits, group_size,
304
- has_act_order, is_k_full, has_zp,
305
- is_zp_float, is_a_8bit, stages);
222
+ int cache_size = get_kernel_cache_size(
223
+ th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits,
224
+ group_size, has_zp, is_zp_float, is_a_8bit, stages);
225
307
- int group_blocks = 0;
308
- if (!has_act_order) {
309
- group_blocks = group_size == -1 ? -1 : group_size / 16;
310
- }
226
+ int group_blocks = group_size == -1 ? -1 : group_size / 16;
227
312
- auto kernel =
313
- get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
314
- th_config.thread_n / 16, th_config.thread_k / 16,
315
- m_block_size_8, has_act_order, has_zp, group_blocks,
316
- th_config.num_threads, is_zp_float, stages);
228
+ auto kernel = get_marlin_kernel(
229
+ a_type, b_type, c_type, s_type, thread_m_blocks,
230
+ th_config.thread_n / 16, th_config.thread_k / 16, m_block_size_8,
231
+ has_zp, group_blocks, th_config.num_threads, is_zp_float, stages);
232
233
if (kernel == MarlinDefault) continue;
234
@@ -324,38 +239,25 @@ exec_config_t determine_exec_config(
239
}
240
241
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
327
- void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
328
- void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k,
329
- int lda, void* workspace, vllm::ScalarType const& a_type,
330
- vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
331
- vllm::ScalarType const& s_type, bool has_bias,
332
- bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
333
- int group_size, int dev, cudaStream_t stream, int thread_k_init,
242
+ void* a_s, void* b_s, void* g_s, void* zp, int prob_m,
243
+ int prob_n, int prob_k, int lda, void* workspace,
244
+ vllm::ScalarType const& a_type, vllm::ScalarType const& b_type,
245
+ vllm::ScalarType const& c_type, vllm::ScalarType const& s_type,
246
+ bool has_bias, bool has_zp, int num_groups, int group_size,
247
+ int dev, cudaStream_t stream, int thread_k_init,
248
int thread_n_init, int sms, bool use_atomic_add,
249
bool use_fp32_reduce, bool is_zp_float) {
250
bool is_a_8bit = a_type.size_bits() == 8;
251
STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [",
252
prob_m, ", ", prob_n, ", ", prob_k, "]");
253
340
- int group_blocks = 0;
341
- if (has_act_order) {
342
- if (is_k_full) {
343
- STD_TORCH_CHECK(group_size != -1);
344
- group_blocks = group_size / 16;
345
- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
346
- " is not divisible by group_blocks = ", group_blocks);
347
- } else {
348
- STD_TORCH_CHECK(group_size == 0);
349
- group_blocks = 0;
350
- }
254
+ int group_blocks;
255
+ if (group_size == -1) {
256
+ group_blocks = -1;
257
} else {
352
- if (group_size == -1) {
353
- group_blocks = -1;
354
- } else {
355
- group_blocks = group_size / 16;
356
- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
357
- " is not divisible by group_blocks = ", group_blocks);
358
- }
258
+ group_blocks = group_size / 16;
259
+ STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
260
+ " is not divisible by group_blocks = ", group_blocks);
261
}
262
263
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,
272
const float* g_s_ptr = (const float*)g_s;
273
274
const int4* zp_ptr = (const int4*)zp;
373
- const int* g_idx_ptr = (const int*)g_idx;
374
- const int* perm_ptr = (const int*)perm;
375
- int4* a_tmp_ptr = (int4*)a_tmp;
275
int* locks = (int*)workspace;
276
378
- if (has_act_order) {
379
- // Permute A columns
380
- int block_rows = div_ceil(prob_m, sms);
381
- // avoid ">>>" being formatted to "> > >"
382
- // clang-format off
383
- permute_cols_kernel<<<sms, default_threads, 0, stream>>>(
384
- A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows);
385
- // clang-format on
386
- A_ptr = a_tmp_ptr;
387
- lda = prob_k;
388
-
389
- // If we have a full K, then we can run the non-act-order version of Marlin
390
- // (since the weight rows are reordered by increasing group ids, and by
391
- // having a full K, we have full original groups)
392
- if (is_k_full) has_act_order = false;
393
- }
394
-
277
int max_shared_mem = 0;
278
cudaDeviceGetAttribute(&max_shared_mem,
279
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
@@ -451,18 +333,17 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
333
// Auto config
334
exec_cfg = determine_exec_config(
335
a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k,
454
- thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order,
455
- is_k_full, has_zp, is_zp_float, is_a_8bit, stages, max_shared_mem,
456
- sms);
336
+ thread_m_blocks, m_block_size_8, num_bits, group_size, has_zp,
337
+ is_zp_float, is_a_8bit, stages, max_shared_mem, sms);
338
thread_tfg = exec_cfg.tb_cfg;
339
if (thread_tfg.thread_n != -1) {
340
if (prob_n / thread_tfg.thread_n *
341
div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <=
342
sms) {
343
if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split,
463
- prob_n, prob_k, num_bits, group_size,
464
- has_act_order, is_k_full, has_zp, is_zp_float,
465
- is_a_8bit, stages, max_shared_mem_new)) {
344
+ prob_n, prob_k, num_bits, group_size, has_zp,
345
+ is_zp_float, is_a_8bit, stages,
346
+ max_shared_mem_new)) {
347
thread_tfg = {128, 64, 128};
348
exec_cfg = {1, thread_tfg};
349
}
@@ -487,30 +368,27 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
368
369
STD_TORCH_CHECK(
370
is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n,
490
- prob_k, num_bits, group_size, has_act_order, is_k_full,
491
- has_zp, is_zp_float, is_a_8bit, stages,
492
- max_shared_mem_new),
371
+ prob_k, num_bits, group_size, has_zp, is_zp_float,
372
+ is_a_8bit, stages, max_shared_mem_new),
373
"Invalid thread config: thread_m_blocks = ", thread_m_blocks,
374
", thread_k = ", thread_tfg.thread_k,
375
", thread_n = ", thread_tfg.thread_n,
376
", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m,
377
", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
378
", prob_m_split = ", prob_m_split, ", group_size = ", group_size,
499
- ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
379
", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
380
", stages = ", stages, ", max_shared_mem_new = ", max_shared_mem_new);
381
382
auto kernel = get_marlin_kernel(
383
a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
505
- thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
506
- num_threads, is_zp_float, stages);
384
+ thread_k_blocks, m_block_size_8, has_zp, group_blocks, num_threads,
385
+ is_zp_float, stages);
386
387
if (kernel == MarlinDefault) {
388
STD_TORCH_CHECK(
389
false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ",
511
- prob_k, "]", ", has_act_order = ", has_act_order,
512
- ", num_groups = ", num_groups, ", group_size = ", group_size,
513
- ", prob_m_split = ", prob_m_split,
390
+ prob_k, "]", ", num_groups = ", num_groups,
391
+ ", group_size = ", group_size, ", prob_m_split = ", prob_m_split,
392
", thread_m_blocks = ", thread_m_blocks,
393
", thread_n_blocks = ", thread_n_blocks,
394
", 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,
405
// clang-format off
406
kernel<<<blocks, num_threads, max_shared_mem_new, stream>>>(
407
A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr,
530
- g_idx_ptr, num_groups,
408
prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add,
409
use_fp32_reduce, max_shared_mem_new);
410
// clang-format on
@@ -550,11 +427,9 @@ torch::stable::Tensor marlin_gemm(
427
std::optional<torch::stable::Tensor> const& a_scales_or_none,
428
std::optional<torch::stable::Tensor> const& global_scale_or_none,
429
std::optional<torch::stable::Tensor> const& b_zeros_or_none,
553
- std::optional<torch::stable::Tensor> const& g_idx_or_none,
554
- std::optional<torch::stable::Tensor> const& perm_or_none,
430
torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id,
556
- int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full,
557
- bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
431
+ int64_t size_m, int64_t size_n, int64_t size_k, bool use_atomic_add,
432
+ bool use_fp32_reduce, bool is_zp_float) {
433
vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
434
435
auto c_scalar_type = a.scalar_type();
@@ -722,7 +597,7 @@ torch::stable::Tensor marlin_gemm(
597
std::nullopt, device);
598
}
599
725
- // Detect groupsize and act_order
600
+ // Detect group size.
601
int num_groups = -1;
602
int group_size = -1;
603
@@ -733,51 +608,13 @@ torch::stable::Tensor marlin_gemm(
608
" is not size_n = ", size_n);
609
num_groups = b_scales.size(0);
610
736
- torch::stable::Tensor g_idx, perm, a_tmp;
737
- if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
738
- g_idx = g_idx_or_none.value();
739
- perm = perm_or_none.value();
740
-
741
- STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
742
- STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
743
- STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
744
- STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
745
-
746
- // Verify g_idx and perm
747
- STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
748
- (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
749
- "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
750
- " and perm.size(-1) = ", perm.size(-1),
751
- ", where size_k = ", size_k);
752
- } else {
753
- g_idx = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
754
- perm = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
755
- a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
756
- }
757
- bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
758
-
759
- if (has_act_order) {
760
- a_tmp = torch::stable::empty({size_m, size_k}, c_scalar_type, std::nullopt,
761
- device);
762
- if (is_k_full) {
763
- STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
764
- STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
765
- ", is not divisible by num_groups = ", num_groups);
766
- group_size = size_k / num_groups;
767
- } else {
768
- group_size = 0;
769
- }
770
-
611
+ if (num_groups > 1) {
612
+ STD_TORCH_CHECK(
613
+ size_k % num_groups == 0, "size_k = ", size_k,
614
+ ", is not divisible by b_scales.size(0) = ", b_scales.size(0));
615
+ group_size = size_k / num_groups;
616
} else {
772
- a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
773
- if (num_groups > 1) {
774
- STD_TORCH_CHECK(
775
- size_k % num_groups == 0, "size_k = ", size_k,
776
- ", is not divisible by b_scales.size(0) = ", b_scales.size(0));
777
- group_size = size_k / num_groups;
778
- } else {
779
- group_size = -1;
780
- }
617
+ group_size = -1;
618
}
619
620
torch::stable::Tensor global_scale;
@@ -882,11 +719,9 @@ torch::stable::Tensor marlin_gemm(
719
a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(),
720
c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(),
721
a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(),
885
- global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(),
886
- g_idx.mutable_data_ptr(), perm.mutable_data_ptr(),
887
- a_tmp.mutable_data_ptr(), size_m, size_n, size_k, a.stride(0),
888
- workspace.mutable_data_ptr(), a_type, b_type, c_type, s_type, has_bias,
889
- has_act_order, is_k_full, has_zp, num_groups, group_size, device_index,
722
+ global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), size_m,
723
+ size_n, size_k, a.stride(0), workspace.mutable_data_ptr(), a_type, b_type,
724
+ c_type, s_type, has_bias, has_zp, num_groups, group_size, device_index,
725
get_current_cuda_stream(device_index), thread_k, thread_n, sms,
726
use_atomic_add, use_fp32_reduce, is_zp_float);
727
csrc/libtorch_stable/quantization/marlin/marlin_template.h
+102
-348
@@ -51,10 +51,9 @@ template <typename scalar_t, // compute dtype, half or nv_float16
51
// only works when thread_m_blocks == 1
52
const int stages, // number of stages for the async global->shared
53
// fetch pipeline
54
- const bool has_act_order, // whether act_order is enabled
55
- const int group_blocks, // number of consecutive 16x16 blocks
56
- // with a separate quantization scale
57
- const bool is_zp_float // is zero point of float16 type?
54
+ const int group_blocks, // number of consecutive 16x16 blocks
55
+ // with a separate quantization scale
56
+ const bool is_zp_float // is zero point of float16 type?
57
>
58
__global__ void Marlin(
59
const int4* __restrict__ A, // fp16 input matrix of shape mxk
@@ -63,11 +62,9 @@ __global__ void Marlin(
62
int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
63
const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
64
// (k/groupsize)xn
66
- const int* __restrict__ g_idx, // int32 group indices of shape k
67
- int num_groups, // number of scale groups per output channel
68
- int prob_m, // batch dimension m
69
- int prob_n, // output dimension n
70
- int prob_k, // reduction dimension k
65
+ int prob_m, // batch dimension m
66
+ int prob_n, // output dimension n
67
+ int prob_k, // reduction dimension k
68
int* locks, // extra global storage for barrier synchronization
69
bool use_fp32_reduce // whether to use fp32 global reduce
70
) {}
@@ -140,29 +137,6 @@ __device__ inline void sub_zp(
137
frag_b[1] = __hsub2(frag_b[1], zp);
138
}
139
143
-// Same as above, but for act_order (each K is multiplied individually)
144
-template <vllm::ScalarTypeId type_id>
145
-__device__ inline void scale4(
146
- typename MarlinScalarType<type_id>::FragB& frag_b,
147
- typename MarlinScalarType<type_id>::FragS& frag_s_1,
148
- typename MarlinScalarType<type_id>::FragS& frag_s_2,
149
- typename MarlinScalarType<type_id>::FragS& frag_s_3,
150
- typename MarlinScalarType<type_id>::FragS& frag_s_4, int i) {
151
- using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
152
- using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
153
-
154
- scalar_t2 s_val_1_2;
155
- s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
156
- s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
157
-
158
- scalar_t2 s_val_3_4;
159
- s_val_3_4.x = reinterpret_cast<scalar_t*>(&frag_s_3)[i];
160
- s_val_3_4.y = reinterpret_cast<scalar_t*>(&frag_s_4)[i];
161
-
162
- frag_b[0] = __hmul2(frag_b[0], s_val_1_2);
163
- frag_b[1] = __hmul2(frag_b[1], s_val_3_4);
164
-}
165
-
140
// Given 2 floats multiply by 2 scales (halves)
141
template <vllm::ScalarTypeId type_id>
142
__device__ inline void scale_float(
@@ -256,14 +230,11 @@ __global__ void Marlin(
230
// 4bit packed zero-points of shape
231
// (k/groupsize, n/pack_factor)
232
const int4* __restrict__ zp_ptr,
259
- // int32 group indices of shape k
260
- const int* __restrict__ g_idx,
261
- int num_groups, // number of scale groups per output channel
262
- int prob_m, // batch dimension m
263
- int prob_n, // output dimension n
264
- int prob_k, // reduction dimension k
265
- int lda, // A.stride(0), equal to prob_k is A is contiguous
266
- int* locks, // extra global storage for barrier synchronization
233
+ int prob_m, // batch dimension m
234
+ int prob_n, // output dimension n
235
+ int prob_k, // reduction dimension k
236
+ int lda, // A.stride(0), equal to prob_k is A is contiguous
237
+ int* locks, // extra global storage for barrier synchronization
238
bool has_bias,
239
bool use_atomic_add, // whether to use atomic add to reduce
240
bool use_fp32_reduce, // whether to use fp32 global reduce
@@ -358,7 +329,8 @@ __global__ void Marlin(
329
global_scale_f32 = global_scale_ptr[0];
330
}
331
361
- constexpr bool has_act_order = group_blocks == 0;
332
+ static_assert(group_blocks != 0,
333
+ "runtime activation ordering is not supported");
334
constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
335
336
extern __shared__ int4 sh[];
@@ -391,7 +363,7 @@ __global__ void Marlin(
363
364
int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x);
365
394
- if constexpr (!has_act_order && group_blocks != -1) {
366
+ if constexpr (group_blocks != -1) {
367
if (group_blocks >= thread_k_blocks) {
368
// Ensure that the number of tiles in each stripe is a multiple of the
369
// groupsize; this avoids an annoying special case where a stripe starts
@@ -558,27 +530,17 @@ __global__ void Marlin(
530
b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
531
constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
532
561
- // Scale sizes/strides without act_order
533
+ // Scale sizes/strides
534
int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8);
535
constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8);
536
constexpr int s_tb_groups =
565
- !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
537
+ group_blocks != -1 && group_blocks < thread_k_blocks
538
? thread_k_blocks / group_blocks
539
: 1;
540
constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
541
int s_gl_rd_delta = s_gl_stride;
542
571
- // Scale size/strides with act_order
572
- constexpr int tb_k = 16 * thread_k_blocks;
573
- constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
574
- // constexpr int act_s_row_stride = 1;
575
- // int act_s_col_stride = act_s_row_stride * num_groups;
576
- constexpr int act_s_max_num_groups = 32;
577
- int act_s_col_stride = 1;
578
- int act_s_col_warp_stride = act_s_col_stride * 8;
579
-
543
constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4);
581
- int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
544
545
// Zero-points sizes/strides
546
int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4;
@@ -615,25 +577,16 @@ __global__ void Marlin(
577
auto b_sh_rd = threadIdx.x * b_thread_vecs;
578
b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1));
579
618
- // For act_order
619
- int slice_k_start = tb_k * slice_row;
620
- int slice_k_finish = slice_k_start + tb_k * slice_iters;
621
- int slice_k_start_shared_fetch = slice_k_start;
622
- int slice_n_offset = act_s_col_tb_stride * slice_col;
623
-
624
- // No act_order
580
int s_gl_rd;
626
- if constexpr (!has_act_order) {
627
- if constexpr (group_blocks == -1) {
628
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
629
- } else if constexpr (group_blocks >= thread_k_blocks) {
630
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
631
- s_sh_stride * slice_col + threadIdx.x;
632
- } else {
633
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
634
- threadIdx.x / s_sh_stride) +
635
- s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
636
- }
581
+ if constexpr (group_blocks == -1) {
582
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
583
+ } else if constexpr (group_blocks >= thread_k_blocks) {
584
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
585
+ s_sh_stride * slice_col + threadIdx.x;
586
+ } else {
587
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
588
+ threadIdx.x / s_sh_stride) +
589
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
590
}
591
auto s_sh_wr = threadIdx.x;
592
bool s_sh_wr_pred = threadIdx.x < s_sh_stage;
@@ -757,10 +710,8 @@ __global__ void Marlin(
710
: (sh_size_b_red_min + sh_bias_size);
711
712
int4* sh_bias = sh_new + sh_size_b_red_min;
760
- int4* sh_g_idx = sh_new + sh_b_red_bias_size;
761
- int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
762
- constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
763
- : (stages * s_sh_stage);
713
+ int4* sh_zp = sh_new + sh_b_red_bias_size;
714
+ constexpr int sh_s_size = stages * s_sh_stage;
715
int4* sh_s = sh_zp + (stages * zp_sh_stage);
716
int4* sh_a = sh_s + sh_s_size;
717
@@ -769,9 +720,8 @@ __global__ void Marlin(
720
I4 frag_b_quant[2][b_thread_vecs];
721
FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2];
722
FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2];
772
- FragS frag_s[2][4]; // No act-order
723
+ FragS frag_s[2][4];
724
FragS frag_bias[2][4];
774
- FragS act_frag_s[2][4][4]; // For act-order
725
int frag_qzp[2][num_ints_per_thread]; // Zero-points
726
FragZP frag_zp; // Zero-points in fp16
727
FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ
@@ -801,42 +751,6 @@ __global__ void Marlin(
751
reinterpret_cast<float*>(frag_c)[i] = 0;
752
};
753
804
- int sh_first_group_id = -1;
805
- int sh_num_groups = -1;
806
-
807
- auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
808
- int last_group_id) {
809
- sh_first_group_id = first_group_id;
810
- sh_num_groups = last_group_id - first_group_id + 1;
811
-
812
- if (sh_num_groups > act_s_max_num_groups) {
813
- sh_num_groups = act_s_max_num_groups;
814
- }
815
-
816
- if (sh_first_group_id + sh_num_groups > num_groups) {
817
- sh_num_groups = num_groups - sh_first_group_id;
818
- }
819
-
820
- int row_offset = first_group_id * s_gl_stride;
821
-
822
- if (is_async) {
823
- for (int i = 0; i < sh_num_groups; i++) {
824
- if (threadIdx.x < s_sh_stride) {
825
- cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
826
- &scales_ptr[row_offset + (i * s_gl_stride) +
827
- slice_n_offset + threadIdx.x]);
828
- }
829
- }
830
- } else {
831
- for (int i = 0; i < sh_num_groups; i++) {
832
- if (threadIdx.x < s_sh_stride) {
833
- sh_s[(i * s_sh_stride) + threadIdx.x] =
834
- scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
835
- threadIdx.x];
836
- }
837
- }
838
- }
839
- };
754
// Asynchronously fetch the next A, B and s tile from global to the next
755
// shared memory pipeline location.
756
auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
@@ -862,44 +776,27 @@ __global__ void Marlin(
776
777
b_gl_rd += b_gl_rd_delta_o;
778
865
- if constexpr (has_act_order) {
866
- // Fetch g_idx thread-block portion
867
- int full_pipe = a_off;
868
- int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
869
- if (cur_k < prob_k && cur_k < slice_k_finish) {
870
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
871
-
872
- int4 const* cur_g_idx_stage_ptr =
873
- reinterpret_cast<int4 const*>(&g_idx[cur_k]);
874
-
875
- if (threadIdx.x < g_idx_stage) {
876
- cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
877
- &cur_g_idx_stage_ptr[threadIdx.x]);
878
- }
879
- }
880
- } else {
881
- if constexpr (group_blocks != -1) {
882
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
779
+ if constexpr (group_blocks != -1) {
780
+ int4* sh_s_stage = sh_s + s_sh_stage * pipe;
781
884
- // Only fetch scales if this tile starts a new group
885
- if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
886
- if (s_sh_wr_pred) {
887
- cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
888
- }
889
- s_gl_rd += s_gl_rd_delta * s_tb_groups;
782
+ // Only fetch scales if this tile starts a new group
783
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
784
+ if (s_sh_wr_pred) {
785
+ cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
786
}
787
+ s_gl_rd += s_gl_rd_delta * s_tb_groups;
788
}
789
+ }
790
893
- if constexpr (has_zp && group_blocks != -1) {
894
- int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
791
+ if constexpr (has_zp && group_blocks != -1) {
792
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
793
896
- // Only fetch zero points if this tile starts a new group
897
- if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
898
- if (zp_sh_wr_pred) {
899
- cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
900
- }
901
- zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
794
+ // Only fetch zero points if this tile starts a new group
795
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
796
+ if (zp_sh_wr_pred) {
797
+ cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
798
}
799
+ zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
800
}
801
}
802
}
@@ -947,150 +844,57 @@ __global__ void Marlin(
844
}
845
};
846
950
- bool is_same_group[stages];
951
- int same_group_id[stages];
952
-
953
- auto init_same_group = [&](int pipe) {
954
- if constexpr (!has_act_order) {
955
- return;
956
- }
957
-
958
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
959
- int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
960
-
961
- int group_id_1 = sh_g_idx_int_ptr[0];
962
- int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];
963
-
964
- is_same_group[pipe] = group_id_1 == group_id_2;
965
- same_group_id[pipe] = group_id_1;
966
- };
967
-
847
auto fetch_scales_to_registers = [&](int k, int full_pipe) {
848
int pipe = full_pipe % stages;
849
using IT1 = typename std::conditional_t<is_a_8bit, int2, int4>;
850
using IT0 = typename std::conditional_t<is_a_8bit, int, int2>;
851
constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1);
852
974
- if constexpr (!has_act_order) {
975
- // No act-order case
976
- if constexpr (group_blocks == -1) {
977
- // load only when starting a new slice
978
- if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
979
- reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
980
- reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
981
- }
982
- } else if constexpr (group_blocks != -1) {
983
- if constexpr (group_blocks >= thread_k_blocks) {
984
- constexpr int g = group_blocks / thread_k_blocks;
985
- if (pipe % g == 0) {
986
- if (k % b_sh_wr_iters == 0) {
987
- int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
988
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
989
- } else {
990
- reinterpret_cast<int4*>(&frag_s[1])[0] =
991
- reinterpret_cast<int4*>(&frag_s[0])[0];
992
- }
993
- }
994
- } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
995
- auto warp_id = threadIdx.x / 32;
996
- int warp_row = warp_id / tb_n_warps;
997
-
998
- int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
999
- int cur_group_id = k_blocks / group_blocks2;
1000
-
1001
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1002
-
1003
- if constexpr (!is_8bit_scale) {
1004
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1005
- sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1006
- } else {
1007
- reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1008
- reinterpret_cast<int2*>(
1009
- sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1010
- }
1011
- } else if (group_blocks >= b_sh_wr_iters) {
1012
- if constexpr (!is_8bit_scale) {
1013
- reinterpret_cast<int4*>(&frag_s[1])[0] =
1014
- reinterpret_cast<int4*>(&frag_s[0])[0];
1015
- } else {
1016
- reinterpret_cast<int2*>(&frag_s[1])[0] =
1017
- reinterpret_cast<int2*>(&frag_s[0])[0];
1018
- }
853
+ if constexpr (group_blocks == -1) {
854
+ // Load only when starting a new slice.
855
+ if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
856
+ reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
857
+ reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
858
+ }
859
+ } else if constexpr (group_blocks >= thread_k_blocks) {
860
+ constexpr int g = group_blocks / thread_k_blocks;
861
+ if (pipe % g == 0) {
862
+ if (k % b_sh_wr_iters == 0) {
863
+ int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
864
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
865
+ } else {
866
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
867
+ reinterpret_cast<int4*>(&frag_s[0])[0];
868
}
869
}
870
+ } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
871
+ auto warp_id = threadIdx.x / 32;
872
+ int warp_row = warp_id / tb_n_warps;
873
1022
- return;
1023
- }
1024
-
1025
- // Act-order case
1026
-
1027
- // Determine K of the "current" thread-block
1028
- int cur_k = slice_k_start + tb_k * full_pipe;
1029
- if (cur_k >= prob_k || cur_k >= slice_k_finish) {
1030
- return;
1031
- }
1032
-
1033
- // Reset (to current thread-block) since we read g_idx portion from the
1034
- // shared memory
1035
- cur_k = 0;
874
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
875
+ int cur_group_id = k_blocks / group_blocks2;
876
1037
- // Progress to current iteration
1038
- cur_k += k % b_sh_wr_iters;
877
+ int4* sh_s_stage = sh_s + s_sh_stage * pipe;
878
1040
- // Determine "position" inside the thread-block (based on warp and
1041
- // thread-id)
1042
- auto warp_id = threadIdx.x / 32;
1043
- int warp_row = warp_id / tb_n_warps;
1044
- int warp_col = warp_id % tb_n_warps;
1045
-
1046
- cur_k += warp_row * 16 * b_sh_wr_iters;
1047
-
1048
- auto th_id = threadIdx.x % 32;
1049
- cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
1050
-
1051
- int s_col_shift =
1052
- /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +
1053
- (th_id / 4) * act_s_col_stride;
1054
-
1055
- if (is_same_group[pipe]) {
1056
- if (k % 2 == 0) {
1057
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1058
- sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +
1059
- s_col_shift];
879
+ if constexpr (!is_8bit_scale) {
880
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
881
+ sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
882
} else {
1061
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1062
- *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));
883
+ reinterpret_cast<int2*>(&frag_s[k % 2])[0] = reinterpret_cast<int2*>(
884
+ sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
885
}
1064
-
1065
- for (int i = 1; i < 4; i++) {
1066
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1067
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));
886
+ } else if (group_blocks >= b_sh_wr_iters) {
887
+ if constexpr (!is_8bit_scale) {
888
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
889
+ reinterpret_cast<int4*>(&frag_s[0])[0];
890
+ } else {
891
+ reinterpret_cast<int2*>(&frag_s[1])[0] =
892
+ reinterpret_cast<int2*>(&frag_s[0])[0];
893
}
1069
- return;
1070
- }
1071
-
1072
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1073
- int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
1074
-
1075
- constexpr int k_frag_offsets[4] = {0, 1, 8,
1076
- 9}; // Tensor core offsets per thread
1077
-
1078
- #pragma unroll
1079
- for (int i = 0; i < 4; i++) {
1080
- int actual_k = cur_k + k_frag_offsets[i];
1081
-
1082
- int group_id = sh_g_idx_int_ptr[actual_k];
1083
- int rel_group_id = group_id - sh_first_group_id;
1084
-
1085
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1086
- sh_s[rel_group_id * s_sh_stride + s_col_shift];
894
}
895
};
896
897
auto fetch_zp_to_registers = [&](int k, int full_pipe) {
1091
- // This code does not handle group_blocks == 0,
1092
- // which signifies act_order.
1093
- // has_zp implies AWQ, which doesn't have act_order,
898
static_assert(!has_zp || group_blocks != 0);
899
900
if constexpr (has_zp && !is_zp_float) {
@@ -1180,8 +984,7 @@ __global__ void Marlin(
984
constexpr int g =
985
group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1;
986
const bool is_new_zp =
1183
- (group_blocks == 0) ||
1184
- ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) &&
987
+ (group_blocks > 0 && (group_blocks < b_sh_wr_iters || k == 0)) &&
988
(pipe % g == 0) ||
989
(group_blocks == -1 && is_first_matmul_in_slice);
990
if constexpr (has_zp && !is_zp_float) {
@@ -1249,15 +1052,8 @@ __global__ void Marlin(
1052
sub_zp<a_type_id>(frag_b1, frag_zp[j], 1);
1053
}
1054
1252
- // Apply scale to frag_b0
1253
- if constexpr (has_act_order && !is_a_8bit) {
1254
- static_assert(group_blocks != -1);
1255
- scale4<a_type_id>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1256
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1257
- scale4<a_type_id>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1258
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1259
- } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1260
- group_blocks == -1 && !is_a_8bit) {
1055
+ if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1056
+ group_blocks == -1 && !is_a_8bit) {
1057
int idx = (threadIdx.x / 4) % 2;
1058
scalar_t2 s2 = Adtype::nums2num2(
1059
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
@@ -1661,7 +1457,7 @@ __global__ void Marlin(
1457
1458
// For per-column quantization we finally apply the scale here (only for
1459
// 4-bit)
1664
- if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit &&
1460
+ if constexpr (group_blocks == -1 && !is_a_8bit &&
1461
b_type.size_bits() == 4 &&
1462
(has_zp && dequant_skip_flop || !has_zp)) {
1463
c_scalar_t2 tmp_scale = s[0];
@@ -1750,15 +1546,6 @@ __global__ void Marlin(
1546
1547
#pragma unroll
1548
for (int i = 0; i < stages - 1; i++) {
1753
- if (has_act_order && i == 0) {
1754
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1755
- if (last_g_idx >= prob_k) {
1756
- last_g_idx = prob_k - 1;
1757
- }
1758
- fetch_act_order_scales_to_shared(true, g_idx[slice_k_start],
1759
- g_idx[last_g_idx]);
1760
- }
1761
-
1549
if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1550
if (i == 0) {
1551
fetch_col_zp_to_shared();
@@ -1772,14 +1559,10 @@ __global__ void Marlin(
1559
1560
zero_accums();
1561
wait_for_stage();
1775
- init_same_group(0);
1562
fetch_to_registers(0, 0);
1563
fetch_scales_to_registers(0, 0);
1564
fetch_zp_to_registers(0, 0);
1565
a_gl_rd += a_gl_rd_delta_o * (stages - 1);
1780
- if constexpr (has_act_order) {
1781
- slice_k_start_shared_fetch += tb_k * (stages - 1);
1782
- }
1566
};
1567
if (slice_iters) {
1568
start_pipes();
@@ -1804,7 +1587,6 @@ __global__ void Marlin(
1587
slice_iters >= stages);
1588
pipe++;
1589
wait_for_stage();
1807
- init_same_group(pipe % stages);
1590
}
1591
1592
if constexpr (!is_a_8bit) {
@@ -1822,25 +1604,6 @@ __global__ void Marlin(
1604
1605
a_gl_rd += a_gl_rd_delta_o * stages;
1606
1825
- if constexpr (has_act_order) {
1826
- slice_k_start += tb_k * stages;
1827
-
1828
- if (slice_k_start < prob_k) {
1829
- slice_k_start_shared_fetch += tb_k * stages;
1830
- int first_group_id = g_idx[slice_k_start];
1831
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1832
- if (last_g_idx >= prob_k) {
1833
- last_g_idx = prob_k - 1;
1834
- }
1835
- int last_group_id = g_idx[last_g_idx];
1836
- if (last_group_id >= sh_first_group_id + sh_num_groups) {
1837
- fetch_act_order_scales_to_shared(false, first_group_id,
1838
- last_group_id);
1839
- __syncthreads();
1840
- }
1841
- }
1842
- }
1843
-
1607
// Process results and, if necessary, proceed to the next column slice.
1608
// While this pattern may not be the most readable, other ways of writing
1609
// the loop seemed to noticeably worse performance after compilation.
@@ -1898,7 +1661,7 @@ __global__ void Marlin(
1661
bool last = slice_idx == slice_count - 1;
1662
// For per-column scales, we only fetch them here in the final step before
1663
// write-out
1901
- if constexpr (!has_act_order && group_blocks == -1 &&
1664
+ if constexpr (group_blocks == -1 &&
1665
(has_zp && dequant_skip_flop || !has_zp)) {
1666
if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) {
1667
if (s_sh_wr_pred) {
@@ -1917,7 +1680,7 @@ __global__ void Marlin(
1680
cp_async_fence();
1681
}
1682
1920
- if constexpr (!has_act_order && group_blocks == -1 &&
1683
+ if constexpr (group_blocks == -1 &&
1684
(has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) {
1685
if constexpr (is_a_8bit) {
1686
cp_async_wait<0>();
@@ -1948,7 +1711,7 @@ __global__ void Marlin(
1711
// For 8-bit channelwise, we apply the scale before the global reduction
1712
// that converts the fp32 results to fp16 (so that we avoid possible
1713
// overflow in fp16)
1951
- if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) {
1714
+ if constexpr (group_blocks == -1 && is_a_8bit) {
1715
#pragma unroll
1716
for (int j = 0; j < 2; j++) {
1717
float2 aa[2];
@@ -1970,8 +1733,7 @@ __global__ void Marlin(
1733
}
1734
}
1735
}
1973
- } else if (!has_act_order && group_blocks == -1 &&
1974
- b_type.size_bits() == 8 &&
1736
+ } else if (group_blocks == -1 && b_type.size_bits() == 8 &&
1737
(has_zp && dequant_skip_flop || !has_zp)) {
1738
if (threadIdx.x / 32 < tb_n_warps) {
1739
#pragma unroll
@@ -2042,33 +1804,25 @@ __global__ void Marlin(
1804
b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row;
1805
1806
bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
2045
- // Update slice k/n for scales loading
2046
- if constexpr (has_act_order) {
2047
- slice_k_start = tb_k * slice_row;
2048
- slice_k_finish = slice_k_start + tb_k * slice_iters;
2049
- slice_k_start_shared_fetch = slice_k_start;
2050
- slice_n_offset = act_s_col_tb_stride * slice_col;
1807
+ if constexpr (group_blocks == -1) {
1808
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
1809
+ zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
1810
+ } else if constexpr (group_blocks >= thread_k_blocks) {
1811
+ s_gl_rd =
1812
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
1813
+ s_sh_stride * slice_col + threadIdx.x;
1814
+ zp_gl_rd =
1815
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
1816
+ zp_sh_stride * slice_col + threadIdx.x;
1817
} else {
2052
- if constexpr (group_blocks == -1) {
2053
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
2054
- zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
2055
- } else if constexpr (group_blocks >= thread_k_blocks) {
2056
- s_gl_rd =
2057
- s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2058
- s_sh_stride * slice_col + threadIdx.x;
2059
- zp_gl_rd =
2060
- zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2061
- zp_sh_stride * slice_col + threadIdx.x;
2062
- } else {
2063
- s_gl_rd =
2064
- s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2065
- threadIdx.x / s_sh_stride) +
2066
- s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
2067
- zp_gl_rd =
2068
- zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2069
- threadIdx.x / zp_sh_stride) +
2070
- zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
2071
- }
1818
+ s_gl_rd =
1819
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
1820
+ threadIdx.x / s_sh_stride) +
1821
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
1822
+ zp_gl_rd =
1823
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
1824
+ threadIdx.x / zp_sh_stride) +
1825
+ zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
1826
}
1827
start_pipes();
1828
}
csrc/libtorch_stable/torch_bindings.cpp
+5
-10
@@ -27,8 +27,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
27
"per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
28
"output_s, int group_size, float eps, float int8_min, float int8_max) -> "
29
"()");
30
- ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
31
-
30
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
31
32
#ifndef USE_ROCM
@@ -85,15 +83,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
83
"marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
84
"Tensor? b_bias_or_none,Tensor b_scales, "
85
"Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
88
- "Tensor? "
89
- "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, "
90
- "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
86
+ "Tensor workspace, int b_type_id, "
87
+ "SymInt size_m, SymInt size_n, SymInt size_k, "
88
"bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
89
// conditionally compiled so impl registrations are in source file
90
91
// gptq_marlin repack from GPTQ.
92
ops.def(
96
- "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
93
+ "gptq_marlin_repack(Tensor b_q_weight, "
94
"SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
95
// conditionally compiled so impl registrations are in source file
96
@@ -679,12 +676,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
676
// to prevent the meta function registry.
677
ops.def(
678
"gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
682
- "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
679
+ "Tensor b_gptq_scales, bool use_exllama, bool "
680
"use_v2_format, int bit) "
681
"-> Tensor");
682
683
// Post processing for GPTQ.
687
- ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()");
684
+ ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()");
685
686
// Mamba selective scan kernel
687
ops.def(
@@ -735,8 +732,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
732
ops.impl("per_token_group_quant_int8",
733
TORCH_BOX(&per_token_group_quant_int8));
734
738
- ops.impl("permute_cols", TORCH_BOX(&permute_cols));
739
-
735
#ifndef USE_ROCM
736
// CUTLASS scaled_mm ops
737
ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm));
csrc/rocm/ops.h
+2
-3
@@ -26,12 +26,11 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
26
27
torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
28
torch::Tensor b_qzeros, torch::Tensor b_scales,
29
- torch::Tensor b_g_idx, bool use_v2_format);
29
+ bool use_v2_format);
30
31
torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
32
torch::Tensor b_qzeros,
33
- torch::Tensor b_scales,
34
- torch::Tensor b_g_idx, bool use_v2_format);
33
+ torch::Tensor b_scales, bool use_v2_format);
34
35
void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
36
torch::Tensor b_q_weight, torch::Tensor b_scales,
csrc/rocm/q_gemm_rdna3.cu
+40
-63
@@ -253,11 +253,13 @@ __forceinline__ __device__ void load4_scales(const T* scales_row, int n,
253
// ---------------------------------------------------------------------------
254
255
template <typename T, int M_COUNT>
256
-__global__ void gemm_q4_kernel_rdna3(
257
- const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
258
- const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
259
- T* __restrict__ c, const int size_m, const int size_n, const int size_k,
260
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
256
+__global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
257
+ const uint32_t* __restrict__ b_q_weight,
258
+ const uint32_t* __restrict__ b_qzeros,
259
+ const T* __restrict__ b_scales,
260
+ T* __restrict__ c, const int size_m,
261
+ const int size_n, const int size_k,
262
+ const int groups, const int zero_offset) {
263
const int t = threadIdx.x;
264
const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
265
const int offset_m = blockIdx.y * M_COUNT;
@@ -274,8 +276,8 @@ __global__ void gemm_q4_kernel_rdna3(
276
constexpr int LDS_PAD = 8;
277
__shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD];
278
277
- // Stage A: each thread loads 1 K element per M row into LDS (with optional
278
- // act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map.
279
+ // Stage A: each thread loads 1 K element per M row into LDS.
280
+ // THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map.
281
// For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m
282
// are zero-padded so the dot product contribution is 0 (we then skip the
283
// atomic write for those rows below).
@@ -298,10 +300,7 @@ __global__ void gemm_q4_kernel_rdna3(
300
T av;
301
if (offset_m + m < size_m) {
302
const T* a_row = a + (offset_m + m) * size_k;
301
- if (b_q_perm)
302
- av = a_row[b_q_perm[offset_k + t]];
303
- else
304
- av = a_row[offset_k + t];
303
+ av = a_row[offset_k + t];
304
} else {
305
av = tzero<T>(); // zero-pad invalid M rows
306
}
@@ -314,14 +313,6 @@ __global__ void gemm_q4_kernel_rdna3(
313
// the LDS load above — but here all THREADS_X (=256) threads always do,
314
// regardless of whether their `n` is in bounds.
315
__syncthreads();
317
- } else if (b_q_perm) {
318
- // bf16 M=1 fast path skips LDS, but its global read below is sequential
319
- // and cannot apply act-order. When a permutation is present, stage the
320
- // single A row through LDS (as fp16 / M>1 do) so the read picks it up.
321
- // b_q_perm is block-uniform, so the __syncthreads is non-divergent.
322
- if (offset_k + t < end_k)
323
- block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]];
324
- __syncthreads();
316
}
317
if (n >= size_n) return;
318
@@ -452,14 +443,11 @@ __global__ void gemm_q4_kernel_rdna3(
443
// fp32-aliased union. Storing as uint32 keeps the IR-level type
444
// opaque so the inner v_dot2 cannot be folded to fp32 widening.
445
//
455
- // A is read direct from global (no LDS staging — see USE_LDS_A above),
456
- // except under act-order, where it comes from the permuted LDS copy.
446
+ // A is read direct from global (no LDS staging — see USE_LDS_A above).
447
pack4 a_pack;
448
{
449
const uint32_t* a_words =
460
- b_q_perm
461
- ? reinterpret_cast<const uint32_t*>(&block_a[0][a_off])
462
- : reinterpret_cast<const uint32_t*>(a + offset_k + a_off);
450
+ reinterpret_cast<const uint32_t*>(a + offset_k + a_off);
451
a_pack.u[0] = a_words[0];
452
a_pack.u[1] = a_words[1];
453
a_pack.u[2] = a_words[2];
@@ -623,8 +611,7 @@ __global__ void gemm_q4_kernel_rdna3(
611
template <typename T, int M_COUNT>
612
__global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*,
613
const T*, T*, const int, const int,
626
- const int, const int, const int,
627
- const int*) {}
614
+ const int, const int, const int) {}
615
616
#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
617
@@ -635,17 +622,17 @@ __global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*,
622
template <typename T, int M_COUNT>
623
void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight,
624
const uint32_t* b_qzeros, const T* b_scales,
638
- const int* b_q_perm, T* c, int size_m,
639
- int size_n, int size_k, int groups,
640
- int zero_offset, cudaStream_t stream) {
625
+ T* c, int size_m, int size_n, int size_k,
626
+ int groups, int zero_offset,
627
+ cudaStream_t stream) {
628
dim3 block(THREADS_X);
629
dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
630
(size_m + M_COUNT - 1) / M_COUNT,
631
(size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
632
646
- gemm_q4_kernel_rdna3<T, M_COUNT><<<grid, block, 0, stream>>>(
647
- a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
648
- zero_offset, b_q_perm);
633
+ gemm_q4_kernel_rdna3<T, M_COUNT>
634
+ <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
635
+ size_n, size_k, groups, zero_offset);
636
}
637
638
// 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,
648
// only burn instructions on the last block, never affect correctness.
649
template <typename T>
650
void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
664
- const uint32_t* b_qzeros, const T* b_scales,
665
- const int* b_q_perm, T* c, int size_m, int size_n,
666
- int size_k, int groups, bool use_v2_format,
667
- cudaStream_t stream) {
651
+ const uint32_t* b_qzeros, const T* b_scales, T* c,
652
+ int size_m, int size_n, int size_k, int groups,
653
+ bool use_v2_format, cudaStream_t stream) {
654
const int zero_offset = use_v2_format ? 0 : 1;
655
656
if (size_m == 1) {
671
- launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
672
- c, size_m, size_n, size_k, groups,
673
- zero_offset, stream);
657
+ launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, c,
658
+ size_m, size_n, size_k, groups, zero_offset,
659
+ stream);
660
} else if (size_m <= 3) {
675
- launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
676
- c, size_m, size_n, size_k, groups,
677
- zero_offset, stream);
661
+ launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, c,
662
+ size_m, size_n, size_k, groups, zero_offset,
663
+ stream);
664
} else if (size_m <= 7) {
679
- launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
680
- c, size_m, size_n, size_k, groups,
681
- zero_offset, stream);
665
+ launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, c,
666
+ size_m, size_n, size_k, groups, zero_offset,
667
+ stream);
668
} else {
669
// M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the
670
// WMMA path, but if it falls through we still produce correct output —
671
// just leaving 3-5× of throughput on the table for prefill workloads.
686
- launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
687
- c, size_m, size_n, size_k, groups,
688
- zero_offset, stream);
672
+ launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, c,
673
+ size_m, size_n, size_k, groups, zero_offset,
674
+ stream);
675
}
676
}
677
@@ -701,7 +687,6 @@ void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
687
// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle)
688
// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros)
689
// b_scales [groups, N] half or bfloat16
704
-// b_g_idx [K] or empty int32 (act-order permutation; empty=identity)
690
// use_v2_format bool (true = GPTQv2, no +1 zero offset)
691
//
692
// Output:
@@ -709,17 +694,16 @@ void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
694
695
torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
696
torch::Tensor b_qzeros,
712
- torch::Tensor b_scales,
713
- torch::Tensor b_g_idx, bool use_v2_format);
697
+ torch::Tensor b_scales, bool use_v2_format);
698
699
torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
700
torch::Tensor b_qzeros, torch::Tensor b_scales,
717
- torch::Tensor b_g_idx, bool use_v2_format) {
701
+ bool use_v2_format) {
702
if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 &&
703
b_q_weight.size(1) % 16 == 0 &&
704
((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) ||
705
(a.scalar_type() == torch::kHalf && a.size(0) >= 64))) {
722
- return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx,
706
+ return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales,
707
use_v2_format);
708
}
709
@@ -753,25 +737,18 @@ torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
737
auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
738
at::Tensor c = torch::zeros({size_m, size_n}, opts);
739
756
- const int* g_idx_ptr = nullptr;
757
- if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) {
758
- TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32,
759
- "b_g_idx must be int32");
760
- g_idx_ptr = (const int*)b_g_idx.data_ptr();
761
- }
762
-
740
if (a.scalar_type() == torch::kHalf) {
741
vllm::gptq_rdna3::launch_gemm_q4<half>(
742
(const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
743
(const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
767
- g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups,
768
- use_v2_format, stream);
744
+ (half*)c.data_ptr(), size_m, size_n, size_k, groups, use_v2_format,
745
+ stream);
746
} else {
747
vllm::gptq_rdna3::launch_gemm_q4<vllm::gptq_rdna3::bf16_t>(
748
(const vllm::gptq_rdna3::bf16_t*)a.data_ptr(),
749
(const uint32_t*)b_q_weight.data_ptr(),
750
(const uint32_t*)b_qzeros.data_ptr(),
774
- (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr,
751
+ (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(),
752
(vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups,
753
use_v2_format, stream);
754
}
csrc/rocm/q_gemm_rdna3_wmma.cu
+83
-142
@@ -373,7 +373,7 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w(
373
const T* __restrict__ a, const uint32_t* __restrict__ b_q,
374
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
375
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
376
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
376
+ const int groups, const int zero_offset) {
377
using E = typename WmmaNative<T>::elem;
378
using V16 = typename WmmaNative<T>::v16;
379
@@ -471,25 +471,16 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w(
471
472
if (m_row < size_m) {
473
const T* a_row = a + m_row * size_k;
474
- if (b_q_perm) {
475
- // Permuted (act-order): scattered global reads, no vectorization.
476
- #pragma unroll
477
- for (int i = 0; i < 16; i++) {
478
- T v = a_row[b_q_perm[k_tile + i]];
479
- a_frag[i] = bitcast_elem<T, E>(v);
480
- }
481
- } else {
482
- // Sequential A reads: replace 16 single-element global_load_b16 with
483
- // a bulk 32-byte copy. The AMDGPU backend lowers a memcpy of this
484
- // size + alignment to two `global_load_b128` instructions. size_k is
485
- // a multiple of 16 (TORCH_CHECK above) and k_tile increments by 16,
486
- // so k_tile + 16 is always within bounds — no tail handling needed.
487
- // Note: we memcpy into the whole vector (`&a_frag`) rather than
488
- // `&a_frag[0]`; ext_vector_type element addresses aren't reliably
489
- // valid C pointers across compiler versions.
490
- static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)");
491
- __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
492
- }
474
+ // Sequential A reads: replace 16 single-element global_load_b16 with
475
+ // a bulk 32-byte copy. The AMDGPU backend lowers a memcpy of this
476
+ // size + alignment to two `global_load_b128` instructions. size_k is
477
+ // a multiple of 16 (TORCH_CHECK above) and k_tile increments by 16,
478
+ // so k_tile + 16 is always within bounds — no tail handling needed.
479
+ // Note: we memcpy into the whole vector (`&a_frag`) rather than
480
+ // `&a_frag[0]`; ext_vector_type element addresses aren't reliably
481
+ // valid C pointers across compiler versions.
482
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)");
483
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
484
} else {
485
#pragma unroll
486
for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
@@ -597,25 +588,24 @@ template <typename T>
588
__global__ void gemm_q4_wmma_kernel_16x16_1w(const T*, const uint32_t*,
589
const uint32_t*, const T*, T*,
590
const int, const int, const int,
600
- const int, const int, const int*) {
601
-}
591
+ const int, const int) {}
592
#endif
593
594
template <typename T>
595
void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight,
596
const uint32_t* b_qzeros, const T* b_scales,
607
- const int* b_q_perm, T* c, int size_m,
608
- int size_n, int size_k, int groups,
609
- int zero_offset, cudaStream_t stream) {
597
+ T* c, int size_m, int size_n, int size_k,
598
+ int groups, int zero_offset,
599
+ cudaStream_t stream) {
600
// 1 wave per block (32 lanes), 16x16 C tile per block. gridDim.z splits
601
// K so that more blocks (and therefore more waves) are in flight; with
602
// K_SPLIT > 1 the kernel switches to atomic write-back at the epilogue.
603
const int k_split = compute_wmma_k_split(size_k);
604
dim3 block(32);
605
dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split);
616
- gemm_q4_wmma_kernel_16x16_1w<T><<<grid, block, 0, stream>>>(
617
- a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
618
- zero_offset, b_q_perm);
606
+ gemm_q4_wmma_kernel_16x16_1w<T>
607
+ <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
608
+ size_n, size_k, groups, zero_offset);
609
}
610
611
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -657,7 +647,7 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w(
647
const T* __restrict__ a, const uint32_t* __restrict__ b_q,
648
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
649
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
660
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
650
+ const int groups, const int zero_offset) {
651
using E = typename WmmaNative<T>::elem;
652
using V16 = typename WmmaNative<T>::v16;
653
@@ -757,16 +747,8 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w(
747
V16 a_frag, b_frag;
748
if (m_row < size_m) {
749
const T* a_row = a + m_row * size_k;
760
- if (b_q_perm) {
761
- #pragma unroll
762
- for (int i = 0; i < 16; i++) {
763
- T v = a_row[b_q_perm[k_tile + i]];
764
- a_frag[i] = bitcast_elem<T, E>(v);
765
- }
766
- } else {
767
- static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
768
- __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
769
- }
750
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
751
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
752
} else {
753
#pragma unroll
754
for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
@@ -844,16 +826,15 @@ template <typename T>
826
__global__ void gemm_q4_wmma_kernel_32x16_2w(const T*, const uint32_t*,
827
const uint32_t*, const T*, T*,
828
const int, const int, const int,
847
- const int, const int, const int*) {
848
-}
829
+ const int, const int) {}
830
#endif
831
832
template <typename T>
833
void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight,
834
const uint32_t* b_qzeros, const T* b_scales,
854
- const int* b_q_perm, T* c, int size_m,
855
- int size_n, int size_k, int groups,
856
- int zero_offset, cudaStream_t stream) {
835
+ T* c, int size_m, int size_n, int size_k,
836
+ int groups, int zero_offset,
837
+ cudaStream_t stream) {
838
// Fallback to v1 for size_m < 32. With M-tile=32 the v2 block has 2 waves
839
// working on rows [0..15] and [16..31]; at M < 32 the second wave processes
840
// 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,
843
// max-num-seqs=32 lands at M≈32 steady-state; the M=16 sliver is edge),
844
// but the fallback costs nothing and is the right shape.
845
if (size_m < 32) {
865
- launch_gemm_q4_wmma_16x16_1w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
866
- c, size_m, size_n, size_k, groups,
867
- zero_offset, stream);
846
+ launch_gemm_q4_wmma_16x16_1w<T>(a, b_q_weight, b_qzeros, b_scales, c,
847
+ size_m, size_n, size_k, groups, zero_offset,
848
+ stream);
849
return;
850
}
851
@@ -875,9 +856,9 @@ void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight,
856
const int k_split = compute_wmma_k_split(size_k);
857
dim3 block(64);
858
dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split);
878
- gemm_q4_wmma_kernel_32x16_2w<T><<<grid, block, 0, stream>>>(
879
- a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
880
- zero_offset, b_q_perm);
859
+ gemm_q4_wmma_kernel_32x16_2w<T>
860
+ <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
861
+ size_n, size_k, groups, zero_offset);
862
}
863
864
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -917,7 +898,7 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w(
898
const T* __restrict__ a, const uint32_t* __restrict__ b_q,
899
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
900
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
920
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
901
+ const int groups, const int zero_offset) {
902
using E = typename WmmaNative<T>::elem;
903
using V16 = typename WmmaNative<T>::v16;
904
@@ -1019,16 +1000,8 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w(
1000
V16 a_frag, b_frag;
1001
if (m_row < size_m) {
1002
const T* a_row = a + m_row * size_k;
1022
- if (b_q_perm) {
1023
- #pragma unroll
1024
- for (int i = 0; i < 16; i++) {
1025
- T v = a_row[b_q_perm[k_tile + i]];
1026
- a_frag[i] = bitcast_elem<T, E>(v);
1027
- }
1028
- } else {
1029
- static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1030
- __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1031
- }
1003
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1004
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1005
} else {
1006
#pragma unroll
1007
for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
@@ -1100,21 +1073,20 @@ template <typename T>
1073
__global__ void gemm_q4_wmma_kernel_64x16_4w(const T*, const uint32_t*,
1074
const uint32_t*, const T*, T*,
1075
const int, const int, const int,
1103
- const int, const int, const int*) {
1104
-}
1076
+ const int, const int) {}
1077
#endif
1078
1079
template <typename T>
1080
void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight,
1081
const uint32_t* b_qzeros, const T* b_scales,
1110
- const int* b_q_perm, T* c, int size_m,
1111
- int size_n, int size_k, int groups,
1112
- int zero_offset, cudaStream_t stream) {
1082
+ T* c, int size_m, int size_n, int size_k,
1083
+ int groups, int zero_offset,
1084
+ cudaStream_t stream) {
1085
// Fall back to v2 for M < 64 (would waste 1+ waves on out-of-range rows).
1086
if (size_m < 64) {
1115
- launch_gemm_q4_wmma_32x16_2w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
1116
- c, size_m, size_n, size_k, groups,
1117
- zero_offset, stream);
1087
+ launch_gemm_q4_wmma_32x16_2w<T>(a, b_q_weight, b_qzeros, b_scales, c,
1088
+ size_m, size_n, size_k, groups, zero_offset,
1089
+ stream);
1090
return;
1091
}
1092
@@ -1122,9 +1094,9 @@ void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight,
1094
const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 16);
1095
dim3 block(128);
1096
dim3 grid((size_n + 15) / 16, (size_m + 63) / 64, k_split);
1125
- gemm_q4_wmma_kernel_64x16_4w<T><<<grid, block, 0, stream>>>(
1126
- a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
1127
- zero_offset, b_q_perm);
1097
+ gemm_q4_wmma_kernel_64x16_4w<T>
1098
+ <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
1099
+ size_n, size_k, groups, zero_offset);
1100
}
1101
1102
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -1158,7 +1130,7 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w(
1130
const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1131
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1132
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1161
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1133
+ const int groups, const int zero_offset) {
1134
using E = typename WmmaNative<T>::elem;
1135
using V16 = typename WmmaNative<T>::v16;
1136
@@ -1264,16 +1236,8 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w(
1236
V16 a_frag, b_frag0, b_frag1;
1237
if (m_row < size_m) {
1238
const T* a_row = a + m_row * size_k;
1267
- if (b_q_perm) {
1268
- #pragma unroll
1269
- for (int i = 0; i < 16; i++) {
1270
- T v = a_row[b_q_perm[k_tile + i]];
1271
- a_frag[i] = bitcast_elem<T, E>(v);
1272
- }
1273
- } else {
1274
- static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1275
- __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1276
- }
1239
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1240
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1241
} else {
1242
#pragma unroll
1243
for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
@@ -1352,23 +1316,22 @@ template <typename T>
1316
__global__ void gemm_q4_wmma_kernel_64x32_4w(const T*, const uint32_t*,
1317
const uint32_t*, const T*, T*,
1318
const int, const int, const int,
1355
- const int, const int, const int*) {
1356
-}
1319
+ const int, const int) {}
1320
#endif
1321
1322
template <typename T>
1323
void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight,
1324
const uint32_t* b_qzeros, const T* b_scales,
1362
- const int* b_q_perm, T* c, int size_m,
1363
- int size_n, int size_k, int groups,
1364
- int zero_offset, cudaStream_t stream) {
1325
+ T* c, int size_m, int size_n, int size_k,
1326
+ int groups, int zero_offset,
1327
+ cudaStream_t stream) {
1328
// Fall back to v3 when M < 64 (small-M decode/prefill stays on the
1329
// narrower 64M × 16N path) or when N < 32 (tile would waste a wave on
1330
// out-of-range cols).
1331
if (size_m < 64 || size_n < 32) {
1369
- launch_gemm_q4_wmma_64x16_4w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
1370
- c, size_m, size_n, size_k, groups,
1371
- zero_offset, stream);
1332
+ launch_gemm_q4_wmma_64x16_4w<T>(a, b_q_weight, b_qzeros, b_scales, c,
1333
+ size_m, size_n, size_k, groups, zero_offset,
1334
+ stream);
1335
return;
1336
}
1337
@@ -1376,9 +1339,9 @@ void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight,
1339
const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 32);
1340
dim3 block(128);
1341
dim3 grid((size_n + 31) / 32, (size_m + 63) / 64, k_split);
1379
- gemm_q4_wmma_kernel_64x32_4w<T><<<grid, block, 0, stream>>>(
1380
- a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
1381
- zero_offset, b_q_perm);
1342
+ gemm_q4_wmma_kernel_64x32_4w<T>
1343
+ <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
1344
+ size_n, size_k, groups, zero_offset);
1345
}
1346
1347
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -1412,7 +1375,7 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w(
1375
const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1376
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1377
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1415
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1378
+ const int groups, const int zero_offset) {
1379
using E = typename WmmaNative<T>::elem;
1380
using V16 = typename WmmaNative<T>::v16;
1381
@@ -1510,16 +1473,8 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w(
1473
V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3;
1474
if (m_row < size_m) {
1475
const T* a_row = a + m_row * size_k;
1513
- if (b_q_perm) {
1514
- #pragma unroll
1515
- for (int i = 0; i < 16; i++) {
1516
- T v = a_row[b_q_perm[k_tile + i]];
1517
- a_frag[i] = bitcast_elem<T, E>(v);
1518
- }
1519
- } else {
1520
- static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1521
- __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1522
- }
1476
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1477
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1478
} else {
1479
#pragma unroll
1480
for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
@@ -1611,7 +1566,7 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16(
1566
const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1567
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1568
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1614
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1569
+ const int groups, const int zero_offset) {
1570
using E = typename WmmaNative<T>::elem;
1571
using V16 = typename WmmaNative<T>::v16;
1572
@@ -1712,8 +1667,7 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16(
1667
dequant_into(next_buf, k_next);
1668
}
1669
1715
- // A-load: vectorized 256-bit. b_q_perm branch removed from V7 hot path
1716
- // to eliminate ~450 ISA instructions of dead code (6× icache bloat).
1670
+ // A-load: vectorized 256-bit.
1671
V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3;
1672
if (a_row_ptr) {
1673
__builtin_memcpy(&a_frag, a_row_ptr + k_tile, sizeof(a_frag));
@@ -1810,7 +1764,7 @@ __global__ void gemm_q4_wmma_kernel_128x64_k32(
1764
const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1765
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1766
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1813
- const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1767
+ const int groups, const int zero_offset) {
1768
using E = typename WmmaNative<T>::elem;
1769
using V16 = typename WmmaNative<T>::v16;
1770
@@ -2006,40 +1960,36 @@ template <typename T>
1960
__global__ void gemm_q4_wmma_kernel_64x64_4w(const T*, const uint32_t*,
1961
const uint32_t*, const T*, T*,
1962
const int, const int, const int,
2009
- const int, const int, const int*) {
2010
-}
1963
+ const int, const int) {}
1964
template <typename T>
1965
__global__ void gemm_q4_wmma_kernel_128x64_k16(const T*, const uint32_t*,
1966
const uint32_t*, const T*, T*,
1967
const int, const int, const int,
2015
- const int, const int,
2016
- const int*) {}
1968
+ const int, const int) {}
1969
template <typename T>
1970
__global__ void gemm_q4_wmma_kernel_128x64_k32(const T*, const uint32_t*,
1971
const uint32_t*, const T*, T*,
1972
const int, const int, const int,
2021
- const int, const int,
2022
- const int*) {}
1973
+ const int, const int) {}
1974
#endif
1975
1976
template <typename T>
1977
void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight,
1978
const uint32_t* b_qzeros, const T* b_scales,
2028
- const int* b_q_perm, T* c, int size_m,
2029
- int size_n, int size_k, int groups,
2030
- int zero_offset, cudaStream_t stream) {
1979
+ T* c, int size_m, int size_n, int size_k,
1980
+ int groups, int zero_offset,
1981
+ cudaStream_t stream) {
1982
// Fall back to v4 when N < 64 (would waste 1+ waves on out-of-range cols).
1983
if (size_m < 64 || size_n < 64) {
2033
- launch_gemm_q4_wmma_64x32_4w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
2034
- c, size_m, size_n, size_k, groups,
2035
- zero_offset, stream);
1984
+ launch_gemm_q4_wmma_64x32_4w<T>(a, b_q_weight, b_qzeros, b_scales, c,
1985
+ size_m, size_n, size_k, groups, zero_offset,
1986
+ stream);
1987
return;
1988
}
1989
1990
// V8 (128M × 64N, K=32/iter, 8-wave dequant) when K%32==0 and gs≥32.
2040
- // Falls back to V7 otherwise. V7/V8 read A sequentially, so act-order
2041
- // (b_q_perm != null) must skip them and use v5, which honors the perm.
2042
- if (size_m >= 128 && b_q_perm == nullptr) {
1991
+ // Falls back to V7 otherwise.
1992
+ if (size_m >= 128) {
1993
const int k_split =
1994
compute_wmma_k_split_mn(size_m, size_n, size_k, 128, 64);
1995
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,
1998
if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) {
1999
gemm_q4_wmma_kernel_128x64_k32<T><<<grid, block, 0, stream>>>(
2000
a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
2051
- zero_offset, b_q_perm);
2001
+ zero_offset);
2002
} else {
2003
gemm_q4_wmma_kernel_128x64_k16<T><<<grid, block, 0, stream>>>(
2004
a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
2055
- zero_offset, b_q_perm);
2005
+ zero_offset);
2006
}
2007
return;
2008
}
@@ -2061,9 +2011,9 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight,
2011
const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 64);
2012
dim3 block(128);
2013
dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split);
2064
- gemm_q4_wmma_kernel_64x64_4w<T><<<grid, block, 0, stream>>>(
2065
- a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
2066
- zero_offset, b_q_perm);
2014
+ gemm_q4_wmma_kernel_64x64_4w<T>
2015
+ <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
2016
+ size_n, size_k, groups, zero_offset);
2017
}
2018
2019
} // namespace gptq_rdna3_wmma
@@ -2078,7 +2028,6 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight,
2028
// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle)
2029
// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros)
2030
// b_scales [groups, N] half or bfloat16
2081
-// b_g_idx [K] or empty int32 (act-order permutation; empty=identity)
2031
// use_v2_format bool (true = GPTQv2, no +1 zero offset)
2032
//
2033
// Output:
@@ -2091,8 +2040,7 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight,
2040
2041
torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
2042
torch::Tensor b_qzeros,
2094
- torch::Tensor b_scales,
2095
- torch::Tensor b_g_idx, bool use_v2_format) {
2043
+ torch::Tensor b_scales, bool use_v2_format) {
2044
TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
2045
TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
2046
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,
2076
// negligible (< 1.5% of prefill time on gfx1100).
2077
at::Tensor c = torch::zeros({size_m, size_n}, opts);
2078
2131
- const int* g_idx_ptr = nullptr;
2132
- if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) {
2133
- TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32,
2134
- "b_g_idx must be int32");
2135
- g_idx_ptr = (const int*)b_g_idx.data_ptr();
2136
- }
2137
-
2079
const int zero_offset = use_v2_format ? 0 : 1;
2080
2081
// 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,
2089
vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w<half>(
2090
(const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
2091
(const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
2151
- g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups,
2152
- zero_offset, stream);
2092
+ (half*)c.data_ptr(), size_m, size_n, size_k, groups, zero_offset,
2093
+ stream);
2094
} else {
2095
vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w<
2096
vllm::gptq_rdna3_wmma::bf16_t>(
2097
(const vllm::gptq_rdna3_wmma::bf16_t*)a.data_ptr(),
2098
(const uint32_t*)b_q_weight.data_ptr(),
2099
(const uint32_t*)b_qzeros.data_ptr(),
2159
- (const vllm::gptq_rdna3_wmma::bf16_t*)b_scales.data_ptr(), g_idx_ptr,
2100
+ (const vllm::gptq_rdna3_wmma::bf16_t*)b_scales.data_ptr(),
2101
(vllm::gptq_rdna3_wmma::bf16_t*)c.data_ptr(), size_m, size_n, size_k,
2102
groups, zero_offset, stream);
2103
}
csrc/rocm/torch_bindings.cpp
+2
-2
@@ -56,12 +56,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
56
// W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
57
rocm_ops.def(
58
"gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
59
- "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
59
+ "Tensor b_scales, bool use_v2_format) -> Tensor");
60
rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3);
61
62
rocm_ops.def(
63
"gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
64
- "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
64
+ "Tensor b_scales, bool use_v2_format) -> Tensor");
65
rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
66
67
rocm_ops.def(
tests/compile/fullgraph/test_full_graph.py
+1
-1
@@ -30,7 +30,7 @@ def models_list(*, all: bool = True, keywords: list[str] | None = None):
30
if is_quant_method_supported("gptq_marlin"):
31
TEST_MODELS.append(
32
(
33
- "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ",
33
+ "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
34
{"quantization": "gptq_marlin"},
35
)
36
)
tests/kernels/core/test_permute_cols.py
deleted
-21
@@ -1,21 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-import torch
6
-
7
-from tests.kernels.utils import opcheck
8
-from vllm._custom_ops import permute_cols
9
-
10
-if not hasattr(torch.ops._C, "permute_cols"):
11
- pytest.skip(reason="permute_cols is not supported on ROCm", allow_module_level=True)
12
-
13
-
14
-@pytest.mark.parametrize("shape", [(1, 512), (544, 4096), (67, 8192)])
15
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
16
-def test_permute_cols(shape, dtype):
17
- x = torch.randn(shape, dtype=dtype).cuda()
18
- perm = torch.randperm(x.shape[1]).to(torch.int).cuda()
19
- opcheck(torch.ops._C.permute_cols, (x, perm))
20
- y = permute_cols(x, perm)
21
- torch.testing.assert_close(y, x[:, perm])
tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
+1
-8
@@ -113,9 +113,7 @@ def marlin_quantize_moe_weights(
113
for i in range(e):
114
# Transpose for Marlin: [n, k] → [k, n]
115
w_t = weights_bf16[i].T.contiguous()
116
- _, w_q, w_s, _, _, _ = marlin_quantize(
117
- w_t, scalar_types.uint4b8, group_size, act_order=False
118
- )
116
+ _, w_q, w_s = marlin_quantize(w_t, scalar_types.uint4b8, group_size)
117
weight_list.append(w_q)
118
scale_list.append(w_s)
119
@@ -248,16 +246,11 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group
246
expert_map=None,
247
global_scale1=None,
248
global_scale2=None,
251
- g_idx1=None,
252
- g_idx2=None,
249
input_global_scale1=None,
250
input_global_scale2=None,
255
- sort_indices1=None,
256
- sort_indices2=None,
251
w1_zeros=None,
252
w2_zeros=None,
253
input_dtype=dtype,
260
- is_k_full=True,
254
)
255
256
# Sanity check: manually compute BF16 reference for comparison
tests/kernels/moe/test_moe.py
+13
-85
@@ -178,13 +178,11 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
178
# GPTQ-INT4
179
{
180
"b_type": scalar_types.uint4b8,
181
- "support_act_order": True,
181
"group_blocks": [-1, 2, 4, 8],
182
},
183
# GPTQ-INT8
184
{
185
"b_type": scalar_types.uint8b128,
187
- "support_act_order": True,
186
"group_blocks": [-1, 2, 4, 8],
187
},
188
# FP8
@@ -794,27 +792,22 @@ def test_fused_moe_wn16(
792
793
794
MARLIN_MOE_SCENARIOS = [
797
- # (m, n, k, e, topk, ep_size, act_order, is_k_full)
798
- # No act_order: is_k_full=True matches usual case (marlin_is_k_full).
795
+ # (m, n, k, e, topk, ep_size)
796
# N>=256 required for Marlin kernel thread config for MXFP8.
797
# Single token, small matrices
801
- (1, 128, 256, 5, 2, 1, False, True),
798
+ (1, 128, 256, 5, 2, 1),
799
# Single token, large matrices
803
- (1, 1024, 2048, 5, 2, 1, False, True),
800
+ (1, 1024, 2048, 5, 2, 1),
801
# Unaligned m, small matrices
805
- (133, 256, 256, 5, 2, 1, False, True),
802
+ (133, 256, 256, 5, 2, 1),
803
# Unaligned m, large matrices
807
- (133, 1024, 2048, 12, 3, 1, False, True),
804
+ (133, 1024, 2048, 12, 3, 1),
805
# Aligned batch, small matrices
809
- (128, 256, 256, 5, 2, 1, False, True),
806
+ (128, 256, 256, 5, 2, 1),
807
# Aligned batch, large matrices
811
- (128, 1024, 2048, 12, 3, 1, False, True),
808
+ (128, 1024, 2048, 12, 3, 1),
809
# Expert parallelism
813
- (64, 1024, 2048, 12, 3, 4, False, True),
814
- # Act order with is_k_full=True (no tensor parallelism)
815
- (1, 1024, 2048, 5, 2, 1, True, True),
816
- # Act order with is_k_full=False (tensor parallelism)
817
- (133, 256, 256, 5, 2, 1, True, False),
810
+ (64, 1024, 2048, 12, 3, 4),
811
]
812
813
@@ -832,17 +825,13 @@ def marlin_moe_generate_valid_test_cases():
825
e,
826
topk,
827
ep_size,
835
- act_order,
836
- is_k_full,
828
):
829
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
830
if group_size > 0 and k % group_size != 0:
831
return False
841
- if act_order and group_size in [-1, k, n]:
842
- return False
832
if group_size in [k, n]:
833
return False
845
- if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
834
+ if b_type == scalar_types.float8_e4m3fn and group_size == 32:
835
return False
836
return a_type.size_bits < 16 or a_type is c_type
837
@@ -858,8 +847,6 @@ def marlin_moe_generate_valid_test_cases():
847
)
848
)
849
861
- supports_act_order = quant_test_config.get("support_act_order", False)
862
-
850
for sub_case in inner_combinations:
851
if (
852
sub_case[0] == scalar_types.float8_e4m3fn
@@ -869,10 +856,8 @@ def marlin_moe_generate_valid_test_cases():
856
continue
857
858
for scenario in MARLIN_MOE_SCENARIOS:
872
- m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
873
- if act_order and not supports_act_order:
874
- continue
875
- args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
859
+ m, n, k, e, topk, ep_size = scenario
860
+ args = sub_case + (m, n, k, e, topk, ep_size)
861
if is_valid(*args):
862
cases.append(args)
863
return cases
@@ -885,9 +870,7 @@ class MarlinMoEWeightData:
870
scales: torch.Tensor
871
global_scale: torch.Tensor | None
872
a_scales_factor: torch.Tensor | None
888
- g_idx: torch.Tensor | None
873
zeros: torch.Tensor | None
890
- sort_indices: torch.Tensor | None
874
marlin_bias: torch.Tensor | None
875
876
@staticmethod
@@ -895,15 +878,12 @@ class MarlinMoEWeightData:
878
w: torch.Tensor,
879
quant_type: ScalarType,
880
group_size: int,
898
- act_order: bool | None = None,
881
bias: torch.Tensor | None = None,
882
input_type: ScalarType = None,
883
) -> "MarlinMoEWeightData":
884
assert w.ndim == 3
885
886
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
905
- k = w.shape[-1]
906
-
887
if input_type == scalar_types.int8:
888
input_dtype = torch.int8
889
elif input_type == scalar_types.float8_e4m3fn:
@@ -916,8 +896,6 @@ class MarlinMoEWeightData:
896
scales_l: list[torch.Tensor] = []
897
global_scale_l: list[torch.Tensor] = []
898
zeros_l: list[torch.Tensor] = []
919
- g_idx_l: list[torch.Tensor] = []
920
- sort_indices_l: list[torch.Tensor] = []
899
bias_l: list[torch.Tensor] = []
900
901
for i in range(w.shape[0]):
@@ -959,21 +937,16 @@ class MarlinMoEWeightData:
937
scales_l.append(scales)
938
zeros_l.append(zeros)
939
else:
962
- test_perm = torch.randperm(k)
963
- w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
940
+ w_ref, qweight, scales = marlin_quantize(
941
w[i].transpose(1, 0),
942
quant_type,
943
group_size,
967
- act_order,
968
- test_perm,
944
input_dtype=input_dtype,
945
)
946
947
w_ref_l.append(w_ref.T)
948
qweight_l.append(qweight)
949
scales_l.append(scales)
975
- g_idx_l.append(g_idx)
976
- sort_indices_l.append(sort_indices)
950
951
if bias is not None:
952
bias_l.append(marlin_permute_bias(bias[i]))
@@ -982,9 +955,7 @@ class MarlinMoEWeightData:
955
qweight = stack_and_dev(qweight_l).contiguous()
956
scales = stack_and_dev(scales_l)
957
global_scale = stack_and_dev(global_scale_l) if global_scale_l else None
985
- g_idx = stack_and_dev(g_idx_l) if g_idx_l else None
958
zeros = stack_and_dev(zeros_l) if zeros_l else None
987
- sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
959
marlin_bias = stack_and_dev(bias_l) if bias_l else None
960
961
a_scales_factor = None
@@ -999,19 +970,14 @@ class MarlinMoEWeightData:
970
scales=scales,
971
global_scale=global_scale,
972
a_scales_factor=a_scales_factor,
1002
- g_idx=g_idx,
973
zeros=zeros,
1004
- sort_indices=sort_indices,
974
marlin_bias=marlin_bias,
975
)
976
977
978
@pytest.mark.flaky(reruns=2)
979
@pytest.mark.parametrize(
1011
- (
1012
- "a_type, b_type, c_type, group_blocks,"
1013
- "m, n, k, e, topk, ep_size, act_order, is_k_full"
1014
- ),
980
+ ("a_type, b_type, c_type, group_blocks,m, n, k, e, topk, ep_size"),
981
marlin_moe_generate_valid_test_cases(),
982
)
983
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
@@ -1027,8 +993,6 @@ def test_fused_marlin_moe(
993
e: int,
994
topk: int,
995
ep_size: int,
1030
- act_order: bool,
1031
- is_k_full: bool,
996
):
997
set_random_seed(1)
998
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
@@ -1065,7 +1029,6 @@ def test_fused_marlin_moe(
1029
w=w1,
1030
quant_type=b_type,
1031
group_size=group_size,
1068
- act_order=act_order,
1032
input_type=a_type,
1033
)
1034
@@ -1073,7 +1036,6 @@ def test_fused_marlin_moe(
1036
w=w2,
1037
quant_type=b_type,
1038
group_size=group_size,
1076
- act_order=act_order,
1039
input_type=a_type,
1040
)
1041
@@ -1127,17 +1089,12 @@ def test_fused_marlin_moe(
1089
expert_map=e_map,
1090
global_scale1=w1_data.global_scale,
1091
global_scale2=w2_data.global_scale,
1130
- g_idx1=w1_data.g_idx,
1131
- g_idx2=w2_data.g_idx,
1092
input_global_scale1=w1_data.a_scales_factor,
1093
input_global_scale2=w2_data.a_scales_factor,
1134
- sort_indices1=w1_data.sort_indices,
1135
- sort_indices2=w2_data.sort_indices,
1094
w1_zeros=w1_data.zeros,
1095
w2_zeros=w2_data.zeros,
1096
input_dtype=a_dtype,
1097
quant_type_id=b_type.id,
1140
- is_k_full=is_k_full,
1098
activation_func=instance_activation,
1099
)
1100
@@ -1154,8 +1111,6 @@ def test_fused_marlin_moe_with_bias(m):
1111
e, topk = 32, 4
1112
n, k = 2048, 2048
1113
group_size = 128
1157
- act_order = False
1158
- is_k_full = True
1114
quant_type = scalar_types.uint4b8
1115
dtype = torch.half
1116
@@ -1169,7 +1124,6 @@ def test_fused_marlin_moe_with_bias(m):
1124
w=w1,
1125
quant_type=quant_type,
1126
group_size=group_size,
1172
- act_order=act_order,
1127
bias=b_bias1,
1128
)
1129
@@ -1177,7 +1131,6 @@ def test_fused_marlin_moe_with_bias(m):
1131
w=w2,
1132
quant_type=quant_type,
1133
group_size=group_size,
1180
- act_order=act_order,
1134
bias=b_bias2,
1135
)
1136
@@ -1204,14 +1157,9 @@ def test_fused_marlin_moe_with_bias(m):
1157
expert_map=None,
1158
global_scale1=w1_data.global_scale,
1159
global_scale2=w2_data.global_scale,
1207
- g_idx1=w1_data.g_idx,
1208
- g_idx2=w2_data.g_idx,
1209
- sort_indices1=w1_data.sort_indices,
1210
- sort_indices2=w2_data.sort_indices,
1160
w1_zeros=w1_data.zeros,
1161
w2_zeros=w2_data.zeros,
1162
quant_type_id=quant_type.id,
1214
- is_k_full=is_k_full,
1163
)
1164
1165
torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)
@@ -1235,7 +1183,6 @@ def test_fused_marlin_moe_non_gated(
1183
set_random_seed(42)
1184
1185
group_size = 16 # NVFP4 group size
1238
- is_k_full = True
1186
quant_type = scalar_types.float4_e2m1f
1187
dtype = torch.bfloat16
1188
@@ -1248,14 +1195,12 @@ def test_fused_marlin_moe_non_gated(
1195
w=w1,
1196
quant_type=quant_type,
1197
group_size=group_size,
1251
- act_order=False,
1198
)
1199
1200
w2_data = MarlinMoEWeightData.make(
1201
w=w2,
1202
quant_type=quant_type,
1203
group_size=group_size,
1258
- act_order=False,
1204
)
1205
1206
score = torch.randn((m, e), device="cuda", dtype=dtype)
@@ -1286,14 +1231,9 @@ def test_fused_marlin_moe_non_gated(
1231
expert_map=None,
1232
global_scale1=w1_data.global_scale,
1233
global_scale2=w2_data.global_scale,
1289
- g_idx1=w1_data.g_idx,
1290
- g_idx2=w2_data.g_idx,
1291
- sort_indices1=w1_data.sort_indices,
1292
- sort_indices2=w2_data.sort_indices,
1234
w1_zeros=w1_data.zeros,
1235
w2_zeros=w2_data.zeros,
1236
quant_type_id=quant_type.id,
1296
- is_k_full=is_k_full,
1237
activation=activation,
1238
)
1239
@@ -1612,14 +1552,9 @@ def test_batched_marlin_activation_uses_expert_token_counts(
1552
quant_type_id=0,
1553
a1_gscale=None,
1554
a2_gscale=None,
1615
- w13_g_idx=None,
1616
- w2_g_idx=None,
1617
- w13_g_idx_sort_indices=None,
1618
- w2_g_idx_sort_indices=None,
1555
w1_zp=None,
1556
w2_zp=None,
1557
input_dtype=None,
1622
- is_k_full=True,
1558
activation_config=ApplyMoEActivationConfig(
1559
clamp_limit=3.0, alpha=1.3, beta=0.5
1560
),
@@ -1982,14 +1917,12 @@ def test_batched_fused_marlin_moe(
1917
w=w1,
1918
quant_type=quant_dtype,
1919
group_size=group_size,
1985
- act_order=None,
1920
input_type=input_type,
1921
)
1922
w2_data = MarlinMoEWeightData.make(
1923
w=w2,
1924
quant_type=quant_dtype,
1925
group_size=group_size,
1992
- act_order=None,
1926
input_type=input_type,
1927
)
1928
@@ -2101,14 +2034,9 @@ def test_batched_fused_marlin_moe(
2034
"expert_map": None,
2035
"global_scale1": w1_data.global_scale,
2036
"global_scale2": w2_data.global_scale,
2104
- "g_idx1": w1_data.g_idx,
2105
- "g_idx2": w2_data.g_idx,
2106
- "sort_indices1": w1_data.sort_indices,
2107
- "sort_indices2": w2_data.sort_indices,
2037
"w1_zeros": w1_data.zeros,
2038
"w2_zeros": w2_data.zeros,
2039
"quant_type_id": quant_dtype.id,
2111
- "is_k_full": True,
2040
}
2041
if input_dtype is not None:
2042
kwargs["input_dtype"] = input_dtype
tests/kernels/quantization/test_allspark_gemm.py
+1
-1
@@ -70,7 +70,7 @@ def test_gptq_allspark_gemm_ampere(mnk_factors, group_size, has_zp, dtype):
70
input = rand_data((m, k), dtype=dtype)
71
weight = rand_data((k, n), dtype=dtype)
72
73
- # Quantize (and apply act_order if provided)
73
+ # Quantize weights.
74
w_ref, qw, s, zp = quantize_weights(
75
weight, scalar_types.uint8b128, group_size, has_zp
76
)
tests/kernels/quantization/test_gptq.py
+3
-9
@@ -11,9 +11,8 @@ def test_gptq_shuffle_opcheck():
11
weight = torch.randint(
12
-2000000, 2000000, (1792, 4096), device="cuda", dtype=torch.int32
13
)
14
- perm = torch.empty((0,), device="cuda", dtype=torch.int32)
14
bit = 4
16
- opcheck(torch.ops._C.gptq_shuffle, (weight, perm, bit))
15
+ opcheck(torch.ops._C.gptq_shuffle, (weight, bit))
16
17
18
def test_gptq_gemm_opcheck():
@@ -23,13 +22,8 @@ def test_gptq_gemm_opcheck():
22
)
23
zeros = torch.zeros((32, 768), device="cuda", dtype=torch.int32)
24
scales = torch.rand((32, 6144), device="cuda", dtype=torch.float16)
26
- idx = torch.empty((0,), device="cuda", dtype=torch.int32)
25
use_exllama = True
26
bit = 4
27
# Test both GPTQv1 and GPTQv2 format
30
- opcheck(
31
- torch.ops._C.gptq_gemm, (a, weight, zeros, scales, idx, use_exllama, True, bit)
32
- )
33
- opcheck(
34
- torch.ops._C.gptq_gemm, (a, weight, zeros, scales, idx, use_exllama, False, bit)
35
- )
28
+ opcheck(torch.ops._C.gptq_gemm, (a, weight, zeros, scales, use_exllama, True, bit))
29
+ opcheck(torch.ops._C.gptq_gemm, (a, weight, zeros, scales, use_exllama, False, bit))
tests/kernels/quantization/test_marlin_gemm.py
+14
-71
@@ -17,7 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
17
per_token_quant_int8,
18
)
19
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
20
- marlin_make_empty_g_idx,
20
+ marlin_make_empty,
21
marlin_make_workspace_new,
22
marlin_permute_bias,
23
query_marlin_supported_quant_types,
@@ -40,7 +40,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
gptq_pack,
41
gptq_quantize_weights,
42
quantize_weights,
43
- sort_weights,
43
)
44
from vllm.platforms import current_platform
45
from vllm.scalar_type import scalar_types
@@ -51,8 +50,6 @@ if current_platform.is_rocm():
50
allow_module_level=True,
51
)
52
54
-ACT_ORDER_OPTS = [False, True]
55
-K_FULL_OPTS = [False, True]
53
USE_ATOMIC_ADD_OPTS = [False, True]
54
USE_FP32_REDUCE_OPTS = [True]
55
@@ -80,13 +77,11 @@ DENSE_MARLIN_QUANT_TEST_CONFIGS = [
77
# GPTQ-INT4
78
{
79
"b_type": scalar_types.uint4b8,
83
- "support_act_order": True,
80
"group_blocks": [-1, 2, 4, 8],
81
},
82
# GPTQ-INT8
83
{
84
"b_type": scalar_types.uint8b128,
89
- "support_act_order": True,
85
"group_blocks": [-1, 2, 4, 8],
86
},
87
# FP8
@@ -202,27 +197,15 @@ def test_marlin_int4_fp8_preprocess_awq():
197
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
198
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
199
@pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(False, False))
205
-@pytest.mark.parametrize("act_order", ACT_ORDER_OPTS)
200
@pytest.mark.parametrize("is_a_8bit", [True, False])
201
@pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS)
208
-def test_gptq_marlin_repack(
209
- k_chunk, n_chunk, quant_type, act_order, is_a_8bit, nk_factors
210
-):
202
+def test_gptq_marlin_repack(k_chunk, n_chunk, quant_type, is_a_8bit, nk_factors):
203
n_factor, k_factor = nk_factors
204
205
size_k = k_chunk * k_factor
206
size_n = n_chunk * n_factor
207
group_size = 128
208
217
- # Filter act_order
218
- if act_order:
219
- if group_size == -1:
220
- return
221
- if group_size == size_k:
222
- return
223
- if is_a_8bit:
224
- return
225
-
209
# Normalize group_size
210
if group_size == -1:
211
group_size = size_k
@@ -231,20 +214,12 @@ def test_gptq_marlin_repack(
214
# Create input
215
b_weight = rand_data((size_k, size_n))
216
234
- # Quantize (and apply act_order if provided)
235
- w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights(
236
- b_weight, quant_type, group_size, act_order
237
- )
217
+ # Quantize
218
+ _, q_w, _ = gptq_quantize_weights(b_weight, quant_type, group_size)
219
220
# Pack to GPTQ format
221
q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
222
242
- # For act_order, sort the "weights" and "g_idx" so that group ids are
243
- # increasing
244
- sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device)
245
- if act_order:
246
- q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
247
-
223
# Pack to Marlin format
224
weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit)
225
marlin_q_w_1 = marlin_weights(
@@ -253,12 +228,12 @@ def test_gptq_marlin_repack(
228
229
opcheck(
230
torch.ops._C.gptq_marlin_repack,
256
- (q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit),
231
+ (q_w_gptq, size_k, size_n, quant_type.size_bits, is_a_8bit),
232
)
233
234
# Run Marlin repack GPU kernel
235
marlin_q_w_2 = ops.gptq_marlin_repack(
261
- q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit
236
+ q_w_gptq, size_k, size_n, quant_type.size_bits, is_a_8bit
237
)
238
torch.accelerator.synchronize()
239
@@ -319,8 +294,6 @@ def marlin_generate_valid_test_cases():
294
MNK_FACTORS,
295
MARLIN_N_CHUNKS,
296
MARLIN_K_CHUNKS,
322
- ACT_ORDER_OPTS,
323
- K_FULL_OPTS,
297
USE_ATOMIC_ADD_OPTS,
298
USE_FP32_REDUCE_OPTS,
299
)
@@ -333,8 +306,6 @@ def marlin_generate_valid_test_cases():
306
size_m,
307
size_n,
308
size_k,
336
- act_order,
337
- is_k_full,
309
use_atomic_add,
310
use_fp32_reduce,
311
):
@@ -351,25 +322,18 @@ def marlin_generate_valid_test_cases():
322
if group_size > 0 and size_k % group_size != 0:
323
return False
324
354
- if act_order and group_size in [-1, size_k]:
355
- return False
325
if group_size == size_k:
326
return False
358
- if not act_order and is_k_full:
359
- return False
327
328
return a_type.size_bits < 16 or a_type is c_type
329
330
cases = []
331
for case in all_combinations:
365
- quant_test_config, mnk_factors, n_chunk, k_chunk, act_order, *_ = case
332
+ quant_test_config, mnk_factors, n_chunk, k_chunk, *_ = case
333
size_m = mnk_factors[0]
334
size_n = mnk_factors[1] * n_chunk
335
size_k = mnk_factors[2] * k_chunk
336
370
- if act_order and not quant_test_config.get("support_act_order", False):
371
- continue
372
-
337
f16_types = [scalar_types.float16, scalar_types.bfloat16]
338
inner_combinations = itertools.product(
339
quant_test_config.get("a_type", f16_types),
@@ -398,7 +362,7 @@ def marlin_generate_valid_test_cases():
362
@pytest.mark.parametrize(
363
(
364
"a_type, b_type, c_type, group_blocks,"
401
- "size_m, size_n, size_k, act_order, is_k_full,"
365
+ "size_m, size_n, size_k,"
366
"use_atomic_add, use_fp32_reduce"
367
),
368
marlin_generate_valid_test_cases(),
@@ -411,8 +375,6 @@ def test_marlin_gemm(
375
size_m,
376
size_n,
377
size_k,
414
- act_order,
415
- is_k_full,
378
use_atomic_add,
379
use_fp32_reduce,
380
):
@@ -448,27 +410,21 @@ def test_marlin_gemm(
410
)
411
marlin_s2 = None
412
451
- g_idx = None
452
- sort_indices = None
413
marlin_zp = None
414
elif b_type == scalar_types.float8_e4m3fn:
415
w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(
416
b_weight.T, group_size, input_dtype=a_dtype
417
)
458
- g_idx = None
459
- sort_indices = None
418
marlin_zp = None
419
marlin_s2 = None
420
elif has_zp:
421
w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize(
422
b_weight, b_type, group_size, input_dtype=a_dtype
423
)
466
- g_idx = None
467
- sort_indices = None
424
marlin_s2 = None
425
else:
470
- w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
471
- b_weight, b_type, group_size, act_order, input_dtype=a_dtype
426
+ w_ref, marlin_q_w, marlin_s = marlin_quantize(
427
+ b_weight, b_type, group_size, input_dtype=a_dtype
428
)
429
430
marlin_zp = None
@@ -506,14 +462,11 @@ def test_marlin_gemm(
462
a_scales,
463
marlin_s2,
464
marlin_zp,
509
- g_idx,
510
- sort_indices,
465
workspace,
466
b_type,
467
a_input.shape[0],
468
b_weight.shape[1],
469
a_input.shape[1],
516
- is_k_full=is_k_full,
470
use_atomic_add=use_atomic_add,
471
use_fp32_reduce=use_fp32_reduce,
472
is_zp_float=False,
@@ -535,11 +488,9 @@ def test_marlin_gemm_subset_input():
488
a_input = rand_data((big_m, big_k))[8 : size_m + 8, 8 : size_k + 8]
489
b_weight = rand_data((size_k, size_n))
490
538
- w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
539
- b_weight, quant_type, group_size, False
540
- )
491
+ w_ref, marlin_q_w, marlin_s = marlin_quantize(b_weight, quant_type, group_size)
492
542
- marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
493
+ marlin_zp = marlin_make_empty(marlin_s.device)
494
workspace = marlin_make_workspace_new(a_input.device)
495
496
output = ops.marlin_gemm(
@@ -551,14 +502,11 @@ def test_marlin_gemm_subset_input():
502
None,
503
None,
504
marlin_zp,
554
- g_idx,
555
- sort_indices,
505
workspace,
506
quant_type,
507
a_input.shape[0],
508
b_weight.shape[1],
509
a_input.shape[1],
561
- is_k_full=True,
510
use_atomic_add=False,
511
use_fp32_reduce=True,
512
is_zp_float=False,
@@ -584,11 +532,9 @@ def test_marlin_gemm_with_bias(size_m):
532
533
marlin_bias = marlin_permute_bias(b_bias)
534
587
- w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
588
- b_weight, quant_type, group_size, False
589
- )
535
+ w_ref, marlin_q_w, marlin_s = marlin_quantize(b_weight, quant_type, group_size)
536
591
- marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
537
+ marlin_zp = marlin_make_empty(marlin_s.device)
538
workspace = marlin_make_workspace_new(a_input.device)
539
540
output = ops.marlin_gemm(
@@ -600,14 +546,11 @@ def test_marlin_gemm_with_bias(size_m):
546
None,
547
None,
548
marlin_zp,
603
- g_idx,
604
- sort_indices,
549
workspace,
550
quant_type,
551
a_input.shape[0],
552
b_weight.shape[1],
553
a_input.shape[1],
610
- is_k_full=True,
554
use_atomic_add=False,
555
use_fp32_reduce=True,
556
is_zp_float=False,
tests/kernels/quantization/test_marlin_tile_padding.py
+6
-22
@@ -14,7 +14,7 @@ from vllm import _custom_ops as ops
14
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
15
GPTQ_MARLIN_TILE,
16
apply_gptq_marlin_linear,
17
- marlin_make_empty_g_idx,
17
+ marlin_make_empty,
18
marlin_make_workspace_new,
19
marlin_moe_padded_intermediate,
20
marlin_pad_qweight,
@@ -319,16 +319,13 @@ def test_gptq_marlin_padded_round_trip(shape, group_size):
319
device = torch.device("cuda")
320
321
weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5
322
- w_ref, q_w, s, _, _ = gptq_quantize_weights(
323
- weight, quant_type, group_size, act_order=False
324
- )
322
+ w_ref, q_w, s = gptq_quantize_weights(weight, quant_type, group_size)
323
qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
324
325
padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
326
qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
327
marlin_qweight = ops.gptq_marlin_repack(
328
b_q_weight=qweight,
331
- perm=torch.empty(0, dtype=torch.int, device=device),
329
size_k=padded_k,
330
size_n=padded_n,
331
num_bits=quant_type.size_bits,
@@ -343,14 +340,11 @@ def test_gptq_marlin_padded_round_trip(shape, group_size):
340
input=x,
341
weight=marlin_qweight,
342
weight_scale=marlin_s,
346
- weight_zp=marlin_make_empty_g_idx(device),
347
- g_idx=marlin_make_empty_g_idx(device),
348
- g_idx_sort_indices=marlin_make_empty_g_idx(device),
343
+ weight_zp=marlin_make_empty(device),
344
workspace=marlin_make_workspace_new(device),
345
wtype=quant_type,
346
output_size_per_partition=size_n,
347
input_size_per_partition=size_k,
353
- is_k_full=True,
348
)
349
ref = x @ w_ref
350
@@ -480,7 +474,6 @@ def test_awq_zp_marlin_padded_round_trip(shape):
474
qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
475
marlin_qweight = ops.gptq_marlin_repack(
476
b_q_weight=qweight,
483
- perm=torch.empty(0, dtype=torch.int, device=device),
477
size_k=padded_k,
478
size_n=padded_n,
479
num_bits=quant_type.size_bits,
@@ -503,13 +496,10 @@ def test_awq_zp_marlin_padded_round_trip(shape):
496
weight=marlin_qweight,
497
weight_scale=marlin_s,
498
weight_zp=marlin_zp,
506
- g_idx=marlin_make_empty_g_idx(device),
507
- g_idx_sort_indices=marlin_make_empty_g_idx(device),
499
workspace=marlin_make_workspace_new(device),
500
wtype=quant_type,
501
output_size_per_partition=size_n,
502
input_size_per_partition=size_k,
512
- is_k_full=True,
503
)
504
ref = x @ w_ref
505
@@ -585,9 +575,7 @@ def test_gptq_marlin_moe_padded_round_trip(shape, group_size):
575
576
def quant(w, size_k, size_n):
577
# w is (size_n, size_k); gptq expects (size_k, size_n).
588
- ref, q_w, s, _, _ = gptq_quantize_weights(
589
- w.T, quant_type, group_size, act_order=False
590
- )
578
+ ref, q_w, s = gptq_quantize_weights(w.T, quant_type, group_size)
579
return ref, gptq_pack(q_w, bits, size_k, size_n), s
580
581
w13_qw, w13_s, w13_ref = [], [], []
@@ -616,12 +604,11 @@ def test_gptq_marlin_moe_padded_round_trip(shape, group_size):
604
if group_size > 0:
605
w2_scales = _pad_rows(w2_scales, padded_n // group_size)
606
619
- sort_idx = torch.empty((e, 0), dtype=torch.int32, device=device)
607
marlin_w13 = ops.gptq_marlin_moe_repack(
621
- w13_qweight, sort_idx, w13_qweight.shape[1] * pack, w13_qweight.shape[2], bits
608
+ w13_qweight, w13_qweight.shape[1] * pack, w13_qweight.shape[2], bits
609
)
610
marlin_w2 = ops.gptq_marlin_moe_repack(
624
- w2_qweight, sort_idx, w2_qweight.shape[1] * pack, w2_qweight.shape[2], bits
611
+ w2_qweight, w2_qweight.shape[1] * pack, w2_qweight.shape[2], bits
612
)
613
group_or_pack = group_size if group_size != -1 else pack
614
marlin_w13_s = marlin_moe_permute_scales(
@@ -649,7 +636,6 @@ def test_gptq_marlin_moe_padded_round_trip(shape, group_size):
636
topk_ids,
637
quant_type_id=quant_type.id,
638
global_num_experts=e,
652
- is_k_full=True,
639
)
640
with set_current_vllm_config(VllmConfig()):
641
ref = torch_experts(
@@ -770,7 +756,6 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant):
756
topk_ids,
757
quant_type_id=scalar_types.float8_e4m3fn.id,
758
global_num_experts=e,
773
- is_k_full=True,
759
workspace=layer.workspace,
760
)
761
with set_current_vllm_config(VllmConfig()):
@@ -848,7 +833,6 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):
833
topk_ids,
834
quant_type_id=scalar_types.float8_e4m3fn.id,
835
global_num_experts=e,
851
- is_k_full=True,
836
workspace=layer.workspace,
837
)
838
with set_current_vllm_config(VllmConfig()):
tests/kernels/quantization/test_rdna3_compile_guards.py
-3
@@ -196,7 +196,6 @@ def test_dense_kernel_rejects_on_non_gfx1100():
196
act_type=torch.float16,
197
group_size=128,
198
zero_points=False,
199
- has_g_idx=False,
199
)
200
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
201
assert ok is False, f"RDNA3 dense kernel accepted on non-gfx1100: {reason}"
@@ -431,7 +430,6 @@ class TestDenseKernelSelectionMocked:
430
act_type=torch.float16,
431
group_size=128,
432
zero_points=False,
434
- has_g_idx=False,
433
)
434
ok, _ = RDNA3W4A16LinearKernel.can_implement(config)
435
assert ok is True
@@ -461,7 +459,6 @@ class TestDenseKernelSelectionMocked:
459
act_type=torch.float16,
460
group_size=128,
461
zero_points=False,
464
- has_g_idx=False,
462
)
463
with (
464
patch("vllm.platforms.rocm.on_gfx1100", return_value=False),
tests/kernels/quantization/test_rdna3_moe_w4a16.py
+1
-7
@@ -69,10 +69,9 @@ def _make_packed_weights(E, K, N):
69
packed = torch.zeros(E, K // 8, N, dtype=torch.int32, device=device)
70
for i in range(8):
71
packed |= (w[:, i::8, :] & 0xF) << (i * 4)
72
- g_idx = torch.empty(0, dtype=torch.int32, device=device)
72
for e in range(E):
73
we = packed[e].contiguous()
75
- ops.gptq_shuffle(we, g_idx, 4)
74
+ ops.gptq_shuffle(we, 4)
75
packed[e] = we
76
return packed
77
@@ -113,7 +112,6 @@ def test_fused_moe_w1_matches_dense(
112
w13 = _make_packed_weights(E, K, N_gate_up)
113
w13_s = _make_scales(E, groups, N_gate_up, dtype)
114
w13_z = _make_qzeros(E, groups, N_gate_up)
116
- g_idx = torch.empty(0, dtype=torch.int32, device=device)
115
116
topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
117
si, ei, ntp = moe_align_block_size(topk_ids, block_size_m, E)
@@ -147,7 +145,6 @@ def test_fused_moe_w1_matches_dense(
145
w13[e],
146
w13_z[e],
147
w13_s[e],
150
- g_idx,
148
False,
149
)
150
ref_out[flat] = ref.squeeze()
@@ -243,7 +240,6 @@ def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype):
240
w2 = _make_packed_weights(E, N_inter, hidden)
241
w2_s = _make_scales(E, N_inter // group_size, hidden, dtype)
242
w2_z = _make_qzeros(E, N_inter // group_size, hidden)
246
- g_idx = torch.empty(0, dtype=torch.int32, device=device)
243
244
topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
245
topk_w = torch.softmax(
@@ -300,7 +296,6 @@ def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype):
296
w13[e],
297
w13_z[e],
298
w13_s[e],
303
- g_idx,
299
False,
300
)
301
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):
305
w2[e],
306
w2_z[e],
307
w2_s[e],
313
- g_idx,
308
False,
309
)
310
ref[m_idx] += r2.squeeze() * w
tests/kernels/quantization/test_rdna3_w4a16.py
-2
@@ -173,7 +173,6 @@ def _run_kernel(
173
act_type=dtype,
174
group_size=group_size,
175
zero_points=has_zp,
176
- has_g_idx=False,
176
)
177
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
178
assert ok, f"can_implement rejected a supported config: {reason}"
@@ -184,7 +183,6 @@ def _run_kernel(
183
w_q_param_name="qweight",
184
w_s_param_name="scales",
185
w_zp_param_name="qzeros" if has_zp else None,
187
- w_gidx_param_name=None,
186
)
187
kernel.process_weights_after_loading(layer)
188
return kernel.apply_weights(layer, x_mk, bias=bias)
tests/kernels/quantization/test_rdna3_w4a16_selection.py
-2
@@ -57,7 +57,6 @@ def test_selection_prefers_rdna3(dtype):
57
act_type=dtype,
58
group_size=128,
59
zero_points=False,
60
- has_g_idx=False,
60
)
61
assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel"
62
@@ -83,7 +82,6 @@ def test_can_implement(weight_type, group_size, N, full_k, expected_ok):
82
act_type=torch.float16,
83
group_size=group_size,
84
zero_points=False,
86
- has_g_idx=False,
85
)
86
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
87
assert ok is expected_ok, reason
tests/kernels/quantization/test_rdna_hybrid_w4a16.py
-5
@@ -306,14 +306,12 @@ def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_ini
306
act_type=torch.float16,
307
group_size=G,
308
zero_points=False,
309
- has_g_idx=False,
309
)
310
kernel = RDNAHybridW4A16LinearKernel(
311
config,
312
w_q_param_name="weight_packed",
313
w_s_param_name="weight_scale",
314
w_zp_param_name=None,
316
- w_gidx_param_name=None,
315
)
316
kernel.process_weights_after_loading(layer)
317
@@ -367,14 +365,12 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in
365
act_type=torch.float16,
366
group_size=G,
367
zero_points=True,
370
- has_g_idx=False,
368
)
369
kernel = RDNAHybridW4A16LinearKernel(
370
config,
371
w_q_param_name="weight_packed",
372
w_s_param_name="weight_scale",
373
w_zp_param_name="weight_zero_point",
377
- w_gidx_param_name=None,
374
)
375
kernel.process_weights_after_loading(layer)
376
@@ -414,7 +410,6 @@ def test_hybrid_can_implement_group_size(group_size, expected_ok):
410
act_type=torch.float16,
411
group_size=group_size,
412
zero_points=False,
417
- has_g_idx=False,
413
)
414
ok, _ = RDNAHybridW4A16LinearKernel.can_implement(config)
415
assert ok is expected_ok
tests/kernels/quantization/test_triton_w4a16.py
-6
@@ -273,14 +273,12 @@ def test_triton_w4a16_process_weights_after_loading_repacks_layout():
273
act_type=torch.float16,
274
group_size=G,
275
zero_points=True,
276
- has_g_idx=False,
276
)
277
kernel = TritonW4A16LinearKernel(
278
config,
279
w_q_param_name="weight_packed",
280
w_s_param_name="weight_scale",
281
w_zp_param_name="weight_zero_point",
283
- w_gidx_param_name=None,
282
)
283
284
# Build dummy layer with vLLM parameter wrappers.
@@ -384,14 +382,12 @@ def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout():
382
act_type=torch.float16,
383
group_size=G,
384
zero_points=True,
387
- has_g_idx=False,
385
)
386
kernel = TritonW4A16LinearKernel(
387
config,
388
w_q_param_name="qweight",
389
w_s_param_name="scales",
390
w_zp_param_name="qzeros",
394
- w_gidx_param_name=None,
391
)
392
393
weight_loader = lambda *args, **kwargs: None
@@ -461,14 +457,12 @@ def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch):
457
act_type=torch.float16,
458
group_size=G,
459
zero_points=False,
464
- has_g_idx=False,
460
)
461
kernel = TritonW4A16LinearKernel(
462
config,
463
w_q_param_name="qweight",
464
w_s_param_name="scales",
465
w_zp_param_name="qzeros",
471
- w_gidx_param_name=None,
466
)
467
468
class DummyLayer(torch.nn.Module):
tests/kernels/quantization/test_w4a16_kernel_selection.py
-2
@@ -47,7 +47,6 @@ def test_choose_mp_linear_kernel_uint4b8():
47
act_type=torch.float16,
48
group_size=128,
49
zero_points=False,
50
- has_g_idx=False,
50
)
51
52
kernel_type = choose_mp_linear_kernel(config)
@@ -65,7 +64,6 @@ def test_choose_mp_linear_kernel_uint4_asymmetric():
64
act_type=torch.bfloat16,
65
group_size=64,
66
zero_points=True,
68
- has_g_idx=False,
67
)
68
69
kernel_type = choose_mp_linear_kernel(config)
tests/lora/conftest.py
-5
@@ -270,11 +270,6 @@ def qwen3_woofing_lora_files():
270
return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA")
271
272
273
-@pytest.fixture(scope="session")
274
-def tinyllama_lora_files():
275
- return hf_api().snapshot_download(repo_id="jashing/tinyllama-colorist-lora")
276
-
277
-
273
@pytest.fixture(scope="session")
274
def deepseekv2_lora_files():
275
return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA")
tests/lora/test_quant_model.py
+73
-76
@@ -24,42 +24,51 @@ MODELS: list[ModelWithQuantization]
24
if current_platform.is_rocm():
25
MODELS = [
26
ModelWithQuantization(
27
- model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", quantization="gptq"
27
+ model_path="JunHowie/Qwen3-0.6B-GPTQ-Int4",
28
+ quantization="gptq",
29
),
30
]
31
else:
32
MODELS = [
33
ModelWithQuantization(
33
- model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", quantization="awq"
34
+ model_path="Orion-zhen/Qwen3-0.6B-AWQ", quantization="awq"
35
),
36
ModelWithQuantization(
36
- model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", quantization="gptq"
37
+ model_path="JunHowie/Qwen3-0.6B-GPTQ-Int4",
38
+ quantization="gptq",
39
),
40
]
41
42
43
def do_sample(
42
- llm: vllm.LLM, lora_path: str, lora_id: int, max_tokens: int = 256
44
+ llm: vllm.LLM, lora_path: str | None, lora_id: int, max_tokens: int = 256
45
) -> list[str]:
44
- raw_prompts = [
45
- "Give me an orange-ish brown color",
46
- "Give me a neon pink color",
46
+ messages = [
47
+ [
48
+ {
49
+ "role": "system",
50
+ "content": "Follow the instructions to make animal noises",
51
+ },
52
+ {"role": "user", "content": "Make your favorite animal noise."},
53
+ ],
54
+ [
55
+ {
56
+ "role": "system",
57
+ "content": "You are a cat. Reply only with your sound.",
58
+ },
59
+ {"role": "user", "content": "What do you say?"},
60
+ ],
61
]
48
-
49
- def format_prompt_tuples(prompt):
50
- return f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
51
-
52
- prompts = [format_prompt_tuples(p) for p in raw_prompts]
53
-
54
- sampling_params = vllm.SamplingParams(
55
- temperature=0, max_tokens=max_tokens, stop=["<|im_end|>"]
56
- )
57
- outputs = llm.generate(
58
- prompts,
62
+ sampling_params = vllm.SamplingParams(temperature=0, max_tokens=max_tokens)
63
+ outputs = llm.chat(
64
+ messages,
65
sampling_params,
60
- lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
66
+ chat_template_kwargs={"enable_thinking": False},
67
+ lora_request=(
68
+ LoRARequest(str(lora_id), lora_id, lora_path) if lora_path else None
69
+ ),
70
+ use_tqdm=False,
71
)
62
- # Print the outputs.
72
generated_texts: list[str] = []
73
for output in outputs:
74
prompt = output.prompt
@@ -70,7 +79,7 @@ def do_sample(
79
80
81
@pytest.mark.parametrize("model", MODELS)
73
-def test_quant_model_lora(tinyllama_lora_files, model):
82
+def test_quant_model_lora(qwen3_meowing_lora_files, qwen3_woofing_lora_files, model):
83
llm = vllm.LLM(
84
model=model.model_path,
85
enable_lora=True,
@@ -79,75 +88,59 @@ def test_quant_model_lora(tinyllama_lora_files, model):
88
max_model_len=400,
89
gpu_memory_utilization=0.2, # avoid OOM
90
quantization=model.quantization,
82
- trust_remote_code=True,
91
enable_chunked_prefill=True,
84
- tokenizer=tinyllama_lora_files,
92
)
93
87
- if model.quantization is None:
88
- expected_lora_output = [
89
- "#ff8050",
90
- "#ff8080",
91
- ]
92
- elif model.quantization == "awq":
93
- expected_lora_output = [
94
- "#f07700: A v",
95
- "#f00000: A v",
96
- ]
97
- elif model.quantization == "gptq":
98
- expected_lora_output = [
99
- "#f08800: This is",
100
- "#f07788 \n#",
101
- ]
102
-
103
- def expect_match(output, expected_output):
104
- # HACK: GPTQ lora outputs are just incredibly unstable.
105
- # Assert that the outputs changed.
106
- if model.quantization == "gptq" and expected_output is expected_lora_output:
107
- for i, o in enumerate(output):
108
- assert o.startswith("#"), (
109
- f"Expected example {i} to start with # but got {o}"
110
- )
111
- return
112
- assert output == expected_output
113
-
94
+ loras = [
95
+ (
96
+ qwen3_meowing_lora_files,
97
+ 1,
98
+ ["Meow Meow Meow Meow Meow"] * 2,
99
+ ),
100
+ (
101
+ qwen3_woofing_lora_files,
102
+ 2,
103
+ ["Woof Woof Woof Woof Woof"] * 2,
104
+ ),
105
+ ]
106
max_tokens = 10
115
-
116
- print("lora adapter created")
117
- print("lora 1")
118
- output = do_sample(llm, tinyllama_lora_files, lora_id=1, max_tokens=max_tokens)
119
- expect_match(output, expected_lora_output)
120
-
121
- print("lora 2")
122
- output = do_sample(llm, tinyllama_lora_files, lora_id=2, max_tokens=max_tokens)
123
- expect_match(output, expected_lora_output)
124
-
125
- print("removing lora")
126
-
127
- del llm
128
- cleanup_dist_env_and_memory()
107
+ try:
108
+ base_output = do_sample(llm, None, lora_id=0, max_tokens=max_tokens)
109
+ for lora_path, lora_id, expected_output in loras:
110
+ output = do_sample(
111
+ llm,
112
+ lora_path,
113
+ lora_id=lora_id,
114
+ max_tokens=max_tokens,
115
+ )
116
+ assert output == expected_output
117
+ assert output != base_output
118
+ finally:
119
+ del llm
120
+ cleanup_dist_env_and_memory()
121
122
123
@pytest.mark.parametrize("model", MODELS)
132
-def test_quant_model_tp_equality(tinyllama_lora_files, num_gpus_available, model):
124
+def test_quant_model_tp_equality(qwen3_meowing_lora_files, num_gpus_available, model):
125
if num_gpus_available < 2:
126
pytest.skip(f"Not enough GPUs for tensor parallelism {2}")
135
- if model.quantization == "gptq":
136
- pytest.skip("GPTQ lora outputs are just incredibly unstable")
127
llm_tp1 = vllm.LLM(
128
model=model.model_path,
129
enable_lora=True,
130
max_num_seqs=16,
131
max_loras=4,
132
+ max_model_len=400,
133
gpu_memory_utilization=0.2, # avoid OOM
134
quantization=model.quantization,
144
- trust_remote_code=True,
135
enable_chunked_prefill=True,
136
)
147
- output_tp1 = do_sample(llm_tp1, tinyllama_lora_files, lora_id=1)
148
-
149
- del llm_tp1
150
- cleanup_dist_env_and_memory()
137
+ try:
138
+ output_tp1 = do_sample(
139
+ llm_tp1, qwen3_meowing_lora_files, lora_id=1, max_tokens=10
140
+ )
141
+ finally:
142
+ del llm_tp1
143
+ cleanup_dist_env_and_memory()
144
145
llm_tp2 = vllm.LLM(
146
model=model.model_path,
@@ -155,13 +148,17 @@ def test_quant_model_tp_equality(tinyllama_lora_files, num_gpus_available, model
148
max_num_seqs=16,
149
max_loras=4,
150
tensor_parallel_size=2,
151
+ max_model_len=400,
152
gpu_memory_utilization=0.2, # avoid OOM
153
quantization=model.quantization,
154
enable_chunked_prefill=True,
155
)
162
- output_tp2 = do_sample(llm_tp2, tinyllama_lora_files, lora_id=1)
163
-
164
- del llm_tp2
165
- cleanup_dist_env_and_memory()
156
+ try:
157
+ output_tp2 = do_sample(
158
+ llm_tp2, qwen3_meowing_lora_files, lora_id=1, max_tokens=10
159
+ )
160
+ finally:
161
+ del llm_tp2
162
+ cleanup_dist_env_and_memory()
163
164
assert output_tp1 == output_tp2
tests/model_executor/model_loader/test_reload.py
+8
-103
@@ -272,13 +272,13 @@ def _stub_marlin_ops(monkeypatch):
272
monkeypatch.setattr(
273
ops,
274
"gptq_marlin_repack",
275
- lambda w, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros(
275
+ lambda w, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros(
276
size_k // 16, size_n * 2, dtype=torch.int32
277
),
278
)
279
280
281
-def _make_act_order_marlin_kernel():
281
+def _make_marlin_kernel():
282
from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
283
MarlinLinearKernel,
284
)
@@ -295,20 +295,17 @@ def _make_act_order_marlin_kernel():
295
act_type=torch.float16,
296
group_size=_MARLIN_GROUP_SIZE,
297
zero_points=False,
298
- has_g_idx=True,
298
)
299
kernel.w_q_name = "qweight"
300
kernel.w_s_name = "scales"
301
kernel.w_zp_name = None
303
- kernel.w_gidx_name = "g_idx"
302
return kernel
303
304
307
-def _load_marlin_checkpoint_format_weights(layer, g_idx):
305
+def _load_marlin_checkpoint_format_weights(layer):
306
from vllm.model_executor.parameter import (
307
GroupQuantScaleParameter,
308
PackedvLLMParameter,
311
- RowvLLMParameter,
309
)
310
311
layer.qweight = PackedvLLMParameter(
@@ -327,54 +324,24 @@ def _load_marlin_checkpoint_format_weights(layer, g_idx):
324
output_dim=1,
325
weight_loader=default_weight_loader,
326
)
330
- layer.g_idx = RowvLLMParameter(
331
- data=g_idx.clone(),
332
- input_dim=0,
333
- weight_loader=default_weight_loader,
334
- )
335
-
336
-
337
-def _random_g_idx(generator):
338
- return torch.randint(
339
- 0,
340
- _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE,
341
- (_MARLIN_SIZE_K,),
342
- dtype=torch.int32,
343
- generator=generator,
344
- )
327
328
329
def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init):
348
- """Marlin workspace and act-order sort indices must be recomputed into
349
- the same storage when weights are reloaded (RL weight sync), so device
350
- addresses captured by CUDA graphs remain valid."""
351
- from vllm.model_executor.layers.quantization.utils import marlin_utils
352
-
330
+ """Marlin must reuse the workspace storage captured by CUDA graphs."""
331
_stub_marlin_ops(monkeypatch)
354
- kernel = _make_act_order_marlin_kernel()
355
-
356
- generator = torch.Generator().manual_seed(0)
357
- first_g_idx = _random_g_idx(generator)
358
- second_g_idx = _random_g_idx(generator)
332
+ kernel = _make_marlin_kernel()
333
334
layer = torch.nn.Module()
361
- _load_marlin_checkpoint_format_weights(layer, first_g_idx)
335
+ _load_marlin_checkpoint_format_weights(layer)
336
kernel.process_weights_after_loading(layer)
337
338
workspace_ptr = kernel.workspace.data_ptr()
365
- sort_indices_ptr = layer.g_idx_sort_indices.data_ptr()
339
367
- # Reload: fresh checkpoint-format tensors with a different act-order
368
- _load_marlin_checkpoint_format_weights(layer, second_g_idx)
340
+ _load_marlin_checkpoint_format_weights(layer)
341
kernel.process_weights_after_loading(layer)
342
343
assert kernel.workspace.data_ptr() == workspace_ptr
344
assert torch.all(kernel.workspace == 0)
373
- assert layer.g_idx_sort_indices.data_ptr() == sort_indices_ptr
374
- expected_sort_indices = marlin_utils.marlin_sort_g_idx(second_g_idx)[1]
375
- assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices)
376
- # registered as a Parameter so layerwise reload copy-back preserves it
377
- assert isinstance(layer.g_idx_sort_indices, torch.nn.Parameter)
345
346
347
@pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"])
@@ -394,7 +361,7 @@ def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant):
361
monkeypatch.setattr(
362
ops,
363
"gptq_marlin_repack",
397
- lambda b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros(
364
+ lambda b_q_weight, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros(
365
size_k // 16, size_n * 2, dtype=torch.int32
366
),
367
)
@@ -476,68 +443,6 @@ def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch):
443
)
444
445
479
-def test_marlin_act_order_layerwise_reload_accounting(monkeypatch, dist_init):
480
- """`g_idx_sort_indices` is generated during weight processing and never
481
- loaded from checkpoints. Registering it as a Parameter must not count it
482
- toward `load_numel_total`: reload restores the construction-time tensor
483
- set before sizing, so act-order layers still process during streaming
484
- instead of deferring (and buffering weights) until finalization."""
485
- from vllm.model_executor.layers.quantization.base_config import (
486
- QuantizeMethodBase,
487
- )
488
- from vllm.model_executor.layers.quantization.utils import marlin_utils
489
- from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info
490
-
491
- _stub_marlin_ops(monkeypatch)
492
- kernel = _make_act_order_marlin_kernel()
493
-
494
- class _KernelQuantMethod(QuantizeMethodBase):
495
- def create_weights(self, layer, *args, **kwargs):
496
- raise NotImplementedError
497
-
498
- def apply(self, layer, *args, **kwargs):
499
- raise NotImplementedError
500
-
501
- def process_weights_after_loading(self, layer):
502
- kernel.process_weights_after_loading(layer)
503
-
504
- generator = torch.Generator().manual_seed(0)
505
- layer = torch.nn.Module()
506
- layer.quant_method = _KernelQuantMethod()
507
- _load_marlin_checkpoint_format_weights(layer, _random_g_idx(generator))
508
-
509
- # Metadata is recorded at model construction, before any processing
510
- record_metadata_for_reloading(layer)
511
- checkpoint_numel = sum(t.numel() for t in get_layer_tensors(layer).values())
512
-
513
- kernel.process_weights_after_loading(layer)
514
- sort_indices = layer.g_idx_sort_indices
515
-
516
- initialize_layerwise_reload(layer)
517
- info = get_layerwise_info(layer)
518
- assert info.load_numel_total == checkpoint_numel
519
-
520
- # Stream a new checkpoint; the layer must process as soon as its last
521
- # tensor arrives
522
- new_g_idx = _random_g_idx(generator)
523
- checkpoint = {
524
- "qweight": torch.zeros(_MARLIN_SIZE_K // 8, _MARLIN_SIZE_N, dtype=torch.int32),
525
- "scales": torch.ones(
526
- _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, _MARLIN_SIZE_N, dtype=torch.float16
527
- ),
528
- "g_idx": new_g_idx,
529
- }
530
- for name, weight in checkpoint.items():
531
- param = getattr(layer, name)
532
- param.weight_loader(param, weight)
533
-
534
- assert not info.can_load()
535
- assert not info.loaded_weights
536
- assert layer.g_idx_sort_indices is sort_indices
537
- expected_sort_indices = marlin_utils.marlin_sort_g_idx(new_g_idx)[1]
538
- assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices)
539
-
540
-
446
def test_model_cleanup(dist_init, default_vllm_config):
447
layer = QKVParallelLinear(2, 3, 4)
448
assert layer.weight.weight_loader.__self__ is layer
tests/model_executor/test_weight_utils.py
+11
@@ -160,6 +160,17 @@ class TestMaybeRemapKvScaleName:
160
assert result is None
161
162
163
+def test_checkpoint_weight_mapper_discards_serialized_g_idx():
164
+ from vllm.model_executor.layers.quantization.base_config import (
165
+ QuantizationConfig,
166
+ )
167
+
168
+ mapper = QuantizationConfig.get_checkpoint_weight_mapper()
169
+
170
+ assert mapper._map_name("model.layers.0.g_idx") is None
171
+ assert mapper._map_name("model.layers.0.qweight") == "model.layers.0.qweight"
172
+
173
+
174
class TestKvCacheScaleMapper:
175
"""The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the
176
per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of
tests/models/quantization/test_gptq_marlin.py
+3
-5
@@ -22,12 +22,10 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true"
22
MAX_MODEL_LEN = 1024
23
24
MODELS = [
25
- # act_order==True, group_size=128
26
- ("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", "main"),
27
- # 8-bit, act_order==True, group_size=channelwise
25
+ ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", "main"),
26
+ # desc_act=True is a no-op for channelwise quantization.
27
("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", "gptq-8bit--1g-actorder_True"),
29
- # 4-bit, act_order==True, group_size=128
30
- ("TechxGenus/gemma-1.1-2b-it-GPTQ", "main"),
28
+ ("TechxGenus/gemma-2b-it-GPTQ", "main"),
29
]
30
31
tests/models/transformers/test_backend.py
+1
-1
@@ -243,7 +243,7 @@ def test_distributed(
243
"model, quantization_kwargs",
244
[
245
("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {}),
246
- ("TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", {}),
246
+ ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", {}),
247
],
248
)
249
@pytest.mark.parametrize("max_tokens", [32])
tests/quantization/test_auto_gptq.py
+50
-2
@@ -23,12 +23,13 @@ from vllm.model_executor.layers.quantization.auto_gptq import (
23
AutoGPTQLinearMethod,
24
AutoGPTQMoEMethod,
25
)
26
+from vllm.model_executor.layers.quantization.utils.gptq_utils import override_config
27
from vllm.platforms import current_platform
28
29
PROMPT = "On the surface of Mars, we found"
30
31
MODELS = [
31
- "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ",
32
+ "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
33
]
34
35
@@ -70,6 +71,54 @@ def test_auto_gptq_config_get_name():
71
assert AutoGPTQConfig.get_name() == "auto_gptq"
72
73
74
+@pytest.mark.parametrize(
75
+ "dynamic",
76
+ [
77
+ {},
78
+ {r"+:model\.layers\.0\..*": {"desc_act": True}},
79
+ {
80
+ r"+:model\.layers\.0\..*": {
81
+ "desc_act": True,
82
+ "group_size": 32,
83
+ }
84
+ },
85
+ ],
86
+)
87
+def test_auto_gptq_rejects_group_activation_order(dynamic):
88
+ desc_act = not dynamic
89
+ with pytest.raises(ValueError, match="group activation ordering"):
90
+ AutoGPTQConfig(4, 128, desc_act, True, False, dynamic, {})
91
+
92
+
93
+def test_auto_gptq_normalizes_channelwise_activation_order():
94
+ config = AutoGPTQConfig(4, -1, True, True, False, {}, {})
95
+ assert not config.desc_act
96
+
97
+ config = AutoGPTQConfig(
98
+ 4,
99
+ 128,
100
+ False,
101
+ True,
102
+ False,
103
+ {r"+:model\.layers\.0\..*": {"desc_act": True, "group_size": -1}},
104
+ {},
105
+ )
106
+ override_config(config, "model.layers.0.self_attn.q_proj")
107
+ assert config.group_size == -1
108
+ assert not config.desc_act
109
+
110
+ with pytest.raises(ValueError, match="group activation ordering"):
111
+ AutoGPTQConfig(
112
+ 4,
113
+ -1,
114
+ True,
115
+ True,
116
+ False,
117
+ {r"+:model\.layers\.0\..*": {"group_size": 128}},
118
+ {},
119
+ )
120
+
121
+
122
def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
123
method = object.__new__(AutoGPTQMoEMethod)
124
method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
@@ -84,7 +133,6 @@ def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
133
hidden_size=8,
134
intermediate_size_per_partition=4,
135
params_dtype=torch.float16,
87
- intermediate_size_full=4,
136
weight_loader=lambda *args, **kwargs: None,
137
)
138
tests/quantization/test_auto_round.py
+3
-4
@@ -1537,8 +1537,7 @@ def test_inc_ark_linear_method_xpu_int2_create_weights(monkeypatch) -> None:
1537
assert layer.scales.dtype == torch.bfloat16
1538
assert layer.qzeros.shape == (1, 4)
1539
assert layer.qzeros.dtype == torch.int32
1540
- assert layer.g_idx.shape == (64,)
1541
- assert layer.g_idx.dtype == torch.int32
1540
+ assert not hasattr(layer, "g_idx")
1541
assert layer.in_features == 64
1542
assert layer.out_features == 64
1543
assert layer.params_dtype == torch.bfloat16
@@ -2444,7 +2443,7 @@ def test_calls_kernel_at_threshold(monkeypatch, w4a8_layer) -> None:
2443
w_scale,
2444
w_zp,
2445
group_size,
2447
- g_idx,
2446
+ group_idx,
2447
bias,
2448
) = captured["gemm_args"]
2449
assert quant_x.dtype is torch.int8
@@ -2456,7 +2455,7 @@ def test_calls_kernel_at_threshold(monkeypatch, w4a8_layer) -> None:
2455
assert qweight is layer.qweight
2456
assert w_zp is layer.qzeros
2457
assert group_size == 128
2459
- assert g_idx is None
2458
+ assert group_idx is None
2459
assert bias is None
2460
2461
# The kernel emits fp16; the result is cast back to the activation dtype.
tests/quantization/test_compressed_tensors.py
+43
-40
@@ -84,6 +84,48 @@ def enable_pickle(monkeypatch):
84
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
85
86
87
+def _compressed_tensors_actorder_config(actorder):
88
+ return {
89
+ "format": "pack-quantized",
90
+ "config_groups": {
91
+ "group_0": {
92
+ "targets": ["Linear"],
93
+ "weights": {
94
+ "num_bits": 4,
95
+ "type": "int",
96
+ "strategy": "group",
97
+ "group_size": 128,
98
+ "symmetric": True,
99
+ "dynamic": False,
100
+ "actorder": actorder,
101
+ },
102
+ }
103
+ },
104
+ }
105
+
106
+
107
+@pytest.mark.parametrize(
108
+ "actorder", [ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC, True]
109
+)
110
+def test_compressed_tensors_rejects_runtime_activation_order(actorder):
111
+ with pytest.raises(ValueError, match="group/dynamic activation ordering"):
112
+ CompressedTensorsConfig.from_config(
113
+ _compressed_tensors_actorder_config(actorder)
114
+ )
115
+
116
+
117
+@pytest.mark.parametrize(
118
+ "actorder", [ActivationOrdering.STATIC, ActivationOrdering.WEIGHT, False, None]
119
+)
120
+def test_compressed_tensors_accepts_preordered_weights(actorder):
121
+ config = CompressedTensorsConfig.from_config(
122
+ _compressed_tensors_actorder_config(actorder)
123
+ )
124
+ assert config.target_scheme_map["Linear"]["weights"].actorder == (
125
+ actorder.value if isinstance(actorder, ActivationOrdering) else None
126
+ )
127
+
128
+
129
@pytest.mark.parametrize(
130
"model_args",
131
[
@@ -271,15 +313,6 @@ def test_compressed_tensors_w8a8_dynamic_per_token(
313
None,
314
8,
315
True,
274
- False,
275
- ),
276
- (
277
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
278
- "group",
279
- 128,
280
- 8,
281
- False,
282
- True,
316
),
317
],
318
)
@@ -287,7 +320,7 @@ def test_compressed_tensors_w8a8_dynamic_per_token(
320
not current_platform.is_cuda(), reason="The tests are skipped on non-CUDA platform."
321
)
322
def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
290
- model, strategy, group, pack_factor, symmetric, has_g_idx = wNa16_args
323
+ model, strategy, group, pack_factor, symmetric = wNa16_args
324
with vllm_runner(model, enforce_eager=True) as llm:
325
326
def check_model(model):
@@ -302,7 +335,6 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
335
336
assert qkv_proj.scheme.pack_factor == pack_factor
337
assert qkv_proj.scheme.symmetric == symmetric
305
- assert qkv_proj.scheme.has_g_idx == has_g_idx
338
339
llm.apply_model(check_model)
340
@@ -926,35 +958,6 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner):
958
assert output
959
960
929
-@pytest.mark.parametrize(
930
- "actorder,group_size,part,full,expected",
931
- [
932
- # actorder="group" with real grouping: must load full-K w2 scales and,
933
- # when sharded (part != full), report is_k_full=False.
934
- (ActivationOrdering.GROUP, 32, 64, 128, (True, 128, False)),
935
- # actorder="group" but unsharded (part == full): full scales, k_full.
936
- (ActivationOrdering.GROUP, 32, 128, 128, (True, 128, True)),
937
- # actorder="group" with channel-wise (group_size == -1): no full load.
938
- (ActivationOrdering.GROUP, -1, 64, 128, (False, 64, False)),
939
- # "static"/"weight" reorder at quant time -> shard normally + k_full.
940
- # Regression: static actorder under TP must keep is_k_full=True so the
941
- # Marlin kernel never gets the invalid (group_size=16, is_k_full=0).
942
- ("static", 32, 64, 128, (False, 64, True)),
943
- ("weight", 32, 64, 128, (False, 64, True)),
944
- (None, 32, 64, 128, (False, 64, True)),
945
- ],
946
-)
947
-def test_wna16_moe_w2_scale_sharding(actorder, group_size, part, full, expected):
948
- from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
949
- CompressedTensorsWNA16MoEMethod,
950
- )
951
-
952
- result = CompressedTensorsWNA16MoEMethod._w2_scale_sharding(
953
- actorder, group_size, part, full
954
- )
955
- assert result == expected
956
-
957
-
961
@pytest.mark.parametrize("num_bits", range(2, 9))
962
def test_humming_supports_compressed_tensors_wna16_quant_key(num_bits):
963
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
tests/quantization/test_cpu_wna16.py
+1
-2
@@ -9,8 +9,7 @@ if not current_platform.is_cpu():
9
10
MODELS = [
11
"TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
12
- "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx
13
- "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx
12
+ "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4",
13
"RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp
14
"OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
15
"Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear
tests/quantization/test_cutlass_w4a16.py
+4
-10
@@ -64,7 +64,6 @@ def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points)
64
weight_type=weight_type,
65
group_size=group_size,
66
zero_points=zero_points,
67
- has_g_idx=False,
67
)
68
kernel = choose_mp_linear_kernel(config)
69
assert kernel is MacheteLinearKernel, (
@@ -73,23 +72,21 @@ def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points)
72
73
74
@pytest.mark.parametrize(
76
- "full_shape,part_shape,weight_type,group_size,has_g_idx,expected_reason",
75
+ "full_shape,part_shape,weight_type,group_size,expected_reason",
76
[
78
- ((4096, 4096), (2048, 4096), scalar_types.uint4b8, 128, True, "Act reordering"),
77
(
78
(4096, 4096),
79
(4096, 4096),
80
scalar_types.float6_e3m2f,
81
128,
84
- False,
82
"Quant type",
83
),
87
- ((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, False, "Group size"),
84
+ ((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, "Group size"),
85
],
89
- ids=["partitioned-g_idx", "unsupported-quant-type", "unsupported-group-size"],
86
+ ids=["unsupported-quant-type", "unsupported-group-size"],
87
)
88
def test_machete_rejects_invalid_config(
92
- full_shape, part_shape, weight_type, group_size, has_g_idx, expected_reason
89
+ full_shape, part_shape, weight_type, group_size, expected_reason
90
):
91
"""Verify Machete rejects unsupported configurations."""
92
config = MPLinearLayerConfig(
@@ -99,7 +96,6 @@ def test_machete_rejects_invalid_config(
96
weight_type=weight_type,
97
group_size=group_size,
98
zero_points=False,
102
- has_g_idx=has_g_idx,
99
)
100
can_impl, reason = MacheteLinearKernel.can_implement(config)
101
assert not can_impl
@@ -117,7 +113,6 @@ def test_kernel_selection_with_disabled_machete(monkeypatch):
113
weight_type=scalar_types.uint4b8,
114
group_size=128,
115
zero_points=False,
120
- has_g_idx=False,
116
)
117
kernel = choose_mp_linear_kernel(config)
118
assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled"
@@ -127,7 +122,6 @@ def test_kernel_selection_with_disabled_machete(monkeypatch):
122
"model_name",
123
[
124
"nm-testing/tinyllama-oneshot-w4a16-channel-v2",
130
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
125
],
126
)
127
def test_w4a16_machete_e2e(vllm_runner, model_name):
tests/quantization/test_gptq_dynamic.py
+7
-50
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests whether gptq models with dynamic quantized can be loaded.
3
+"""Tests whether unsupported GPTQ group actorder models are rejected.
4
5
Run `pytest tests/quantization/test_gptq_dynamic.py --forked`.
6
@@ -8,59 +8,16 @@ Note: Only symmetric GPTQ models are supported after consolidation to Marlin.
8
"""
9
10
import pytest
11
-import torch
11
13
-from tests.quantization.utils import load_model_without_vllm_runner
14
-from vllm.model_executor.layers.linear import UnquantizedLinearMethod
15
-from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
16
-from vllm.model_executor.layers.quantization.utils.gptq_utils import (
17
- get_dynamic_override,
18
-)
19
-
20
-PROMPT = "On the surface of Mars, we found"
21
-
22
-# The first layer is quantized using bits=4, group_size=128
23
-# The second layer is quantized using bits=8, group_size=32
24
-# All other layers (layer index >= 2) are not quantized
25
-# Note: Only symmetric models are supported with Marlin kernels
12
MODELS = [
13
"ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head-symTrue",
14
]
15
16
17
@pytest.mark.parametrize("model_id", MODELS)
32
-def test_gptq_with_dynamic(model_id: str, monkeypatch, dist_init, workspace_init):
33
- # `LLM.apply_model` requires pickling a function.
34
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
35
-
36
- linear_method_cls = AutoGPTQLinearMethod
37
-
38
- model, _ = load_model_without_vllm_runner(
39
- model_id,
40
- dtype=torch.float16,
41
- model_config_kwargs={
42
- "max_model_len": 2048,
43
- "hf_overrides": {"num_hidden_layers": 3},
44
- },
45
- )
46
-
47
- for name, submodule in model.named_modules():
48
- if name == "lm_head":
49
- assert isinstance(submodule.quant_method, linear_method_cls)
50
- elif name == "model.layers.0.self_attn.qkv_proj":
51
- assert isinstance(submodule.quant_method, linear_method_cls)
52
- config = submodule.quant_method.quant_config
53
- assert config.weight_bits == 4
54
- assert config.group_size == 128
55
- assert config.desc_act
56
- elif name == "model.layers.1.self_attn.qkv_proj":
57
- assert isinstance(submodule.quant_method, linear_method_cls)
58
- config = submodule.quant_method.quant_config
59
- assert get_dynamic_override(config, layer_name=name, key="bits") == 8
60
- assert get_dynamic_override(config, layer_name=name, key="group_size") == 32
61
- assert not get_dynamic_override(config, layer_name=name, key="desc_act")
62
- elif name in (
63
- "model.layers.2.self_attn.qkv_proj",
64
- "model.layers.2.mlp.gate_up_proj",
65
- ):
66
- assert isinstance(submodule.quant_method, UnquantizedLinearMethod)
18
+def test_gptq_with_dynamic_desc_act_rejected(vllm_runner, model_id: str):
19
+ with (
20
+ pytest.raises(ValueError, match="desc_act=True"),
21
+ vllm_runner(model_id, max_model_len=2048, enforce_eager=True),
22
+ ):
23
+ pass
tests/quantization/test_lm_head.py
+1
-2
@@ -17,8 +17,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
17
PROMPT = "On the surface of Mars, we found"
18
19
MODELS_QUANT = [
20
- ("ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head", True),
21
- ("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", False),
20
+ ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", False),
21
]
22
23
tests/quantization/test_moe_wna16.py
+29
-25
@@ -6,7 +6,6 @@ from types import SimpleNamespace
6
import pytest
7
import torch
8
from compressed_tensors.quantization import (
9
- ActivationOrdering,
9
QuantizationArgs,
10
QuantizationStrategy,
11
QuantizationType,
@@ -33,6 +32,35 @@ def test_map_wna16_backend_supports_triton():
32
assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON
33
34
35
+@pytest.mark.parametrize(
36
+ "config",
37
+ [
38
+ {"desc_act": True, "group_size": 128},
39
+ {
40
+ "desc_act": False,
41
+ "group_size": 128,
42
+ "dynamic": {r"+:model\.layers\.0\..*": {"desc_act": True}},
43
+ },
44
+ ],
45
+)
46
+def test_moe_wna16_rejects_gptq_group_activation_order(config):
47
+ config.update({"quant_method": "gptq", "bits": 4, "sym": True})
48
+ with pytest.raises(ValueError, match="group activation ordering"):
49
+ MoeWNA16Config.from_config(config)
50
+
51
+
52
+def test_moe_wna16_accepts_channelwise_gptq_activation_order():
53
+ config = {
54
+ "quant_method": "gptq",
55
+ "bits": 4,
56
+ "group_size": -1,
57
+ "desc_act": True,
58
+ "sym": True,
59
+ }
60
+ assert MoeWNA16Config.is_moe_wna16_compatible(config)
61
+ MoeWNA16Config.from_config(config)
62
+
63
+
64
@pytest.mark.parametrize(
65
("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"),
66
[
@@ -43,28 +71,6 @@ def test_map_wna16_backend_supports_triton():
71
False,
72
"AutoAWQ weight layout",
73
),
46
- (
47
- WNA16MoEBackend.TRITON,
48
- AutoGPTQConfig(4, 128, True, True, False, {}, {}),
49
- False,
50
- False,
51
- "activation ordering",
52
- ),
53
- (
54
- WNA16MoEBackend.TRITON,
55
- QuantizationArgs(
56
- num_bits=4,
57
- type=QuantizationType.INT,
58
- strategy=QuantizationStrategy.GROUP,
59
- symmetric=True,
60
- dynamic=False,
61
- group_size=128,
62
- actorder=ActivationOrdering.GROUP,
63
- ),
64
- False,
65
- False,
66
- "activation ordering",
67
- ),
74
(
75
WNA16MoEBackend.TRITON,
76
AutoGPTQConfig(4, 128, False, True, False, {}, {}),
@@ -244,7 +250,6 @@ def test_compressed_tensors_wna16_moe_create_weights_uses_ceil_packed_shapes(
250
hidden_size=256,
251
intermediate_size_per_partition=512,
252
params_dtype=torch.float16,
247
- intermediate_size_full=512,
253
)
254
255
packed_hidden = (256 * num_bits + 31) // 32
@@ -296,7 +301,6 @@ def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel():
301
num_experts=2,
302
hidden_size=256,
303
intermediate_size_per_partition=512,
299
- intermediate_size_full=512,
304
params_dtype=torch.bfloat16,
305
)
306
layer.cuda()
tests/weight_loading/models-large.txt
+1
-3
@@ -1,8 +1,6 @@
1
compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-quantized, main
2
compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-channel-quantized, main
3
compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W8A16-quantized, main
4
-compressed-tensors, nm-testing/test-w4a16-mixtral-actorder-group, main
4
gptq, TheBloke/Mixtral-8x7B-v0.1-GPTQ, main
6
-gptq, TheBloke/Mixtral-8x7B-v0.1-GPTQ, gptq-8bit-128g-actorder_True
5
awq_marlin, casperhansen/deepseek-coder-v2-instruct-awq, main
8
-compressed-tensors, RedHatAI/Llama-4-Scout-17B-16E-Instruct-quantized.w4a16, main
\ No newline at end of file
6
+compressed-tensors, RedHatAI/Llama-4-Scout-17B-16E-Instruct-quantized.w4a16, main
tests/weight_loading/models.txt
+3
-4
@@ -1,9 +1,8 @@
1
gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main
2
gptq, TheBloke/Llama-2-7B-GPTQ, main
3
-gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, main
3
+gptq, LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit, main
4
gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True
5
-gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit-32g-actorder_True
6
-gptq, TechxGenus/gemma-1.1-2b-it-GPTQ, main
5
+gptq, TechxGenus/gemma-2b-it-GPTQ, main
6
compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main
7
compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main
8
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
15
#compressed-tensors, mgoin/DeepSeek-Coder-V2-Lite-Instruct-FP8, main
16
awq, casperhansen/mixtral-instruct-awq, main
17
awq_marlin, casperhansen/mixtral-instruct-awq, main
19
-fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
\ No newline at end of file
18
+fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
tests/weight_loading/test_weight_loading.py
+1
-4
@@ -34,10 +34,7 @@ def test_weight_loading(vllm_runner):
34
"""
35
36
# MoE models need fp16.
37
- NEEDS_FP16 = (
38
- QUANTIZATION == "gptq"
39
- or MODEL_NAME == "nm-testing/test-w4a16-mixtral-actorder-group"
40
- )
37
+ NEEDS_FP16 = QUANTIZATION == "gptq"
38
with vllm_runner(
39
model_name=MODEL_NAME,
40
revision=REVISION,
vllm/_custom_ops.py
+5
-46
@@ -613,7 +613,6 @@ def gptq_gemm(
613
b_q_weight: torch.Tensor,
614
b_gptq_qzeros: torch.Tensor,
615
b_gptq_scales: torch.Tensor,
616
- b_g_idx: torch.Tensor,
616
use_exllama: bool,
617
use_v2_format: bool,
618
bit: int,
@@ -623,7 +622,6 @@ def gptq_gemm(
622
b_q_weight,
623
b_gptq_qzeros,
624
b_gptq_scales,
626
- b_g_idx,
625
use_exllama,
626
use_v2_format,
627
bit,
@@ -638,7 +636,6 @@ if hasattr(torch.ops._C, "gptq_gemm"):
636
b_q_weight: torch.Tensor,
637
b_gptq_qzeros: torch.Tensor,
638
b_gptq_scales: torch.Tensor,
641
- b_g_idx: torch.Tensor,
639
use_exllama: bool,
640
use_v2_format: bool,
641
bit: int,
@@ -648,8 +645,8 @@ if hasattr(torch.ops._C, "gptq_gemm"):
645
)
646
647
651
-def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None:
652
- torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
648
+def gptq_shuffle(q_weight: torch.Tensor, bit: int) -> None:
649
+ torch.ops._C.gptq_shuffle(q_weight, bit)
650
651
652
def gptq_gemm_rdna3(
@@ -657,11 +654,10 @@ def gptq_gemm_rdna3(
654
b_q_weight: torch.Tensor,
655
b_qzeros: torch.Tensor,
656
b_scales: torch.Tensor,
660
- b_g_idx: torch.Tensor,
657
use_v2_format: bool,
658
) -> torch.Tensor:
659
return torch.ops._rocm_C.gptq_gemm_rdna3(
664
- a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format
660
+ a, b_q_weight, b_qzeros, b_scales, use_v2_format
661
)
662
663
@@ -673,7 +669,6 @@ if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3
669
b_q_weight: torch.Tensor,
670
b_qzeros: torch.Tensor,
671
b_scales: torch.Tensor,
676
- b_g_idx: torch.Tensor,
672
use_v2_format: bool,
673
) -> torch.Tensor:
674
return torch.empty(
@@ -689,7 +684,6 @@ if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3
684
b_q_weight: torch.Tensor,
685
b_qzeros: torch.Tensor,
686
b_scales: torch.Tensor,
692
- b_g_idx: torch.Tensor,
687
use_v2_format: bool,
688
) -> torch.Tensor:
689
return torch.empty(
@@ -1114,14 +1108,13 @@ def cutlass_mxfp4_moe_mm(
1108
# gptq_marlin
1109
def gptq_marlin_repack(
1110
b_q_weight: torch.Tensor,
1117
- perm: torch.Tensor,
1111
size_k: int,
1112
size_n: int,
1113
num_bits: int,
1114
is_a_8bit: bool = False,
1115
) -> torch.Tensor:
1116
return torch.ops._C.gptq_marlin_repack(
1124
- b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit
1117
+ b_q_weight, size_k, size_n, num_bits, is_a_8bit
1118
)
1119
1120
@@ -1130,7 +1123,6 @@ if hasattr(torch.ops._C, "gptq_marlin_repack"):
1123
@register_fake("_C::gptq_marlin_repack")
1124
def _gptq_marlin_repack_fake(
1125
b_q_weight: torch.Tensor,
1133
- perm: torch.Tensor,
1126
size_k: torch.SymInt,
1127
size_n: torch.SymInt,
1128
num_bits: int,
@@ -1179,7 +1171,6 @@ if hasattr(torch.ops._C, "awq_marlin_repack"):
1171
1172
def gptq_marlin_moe_repack(
1173
b_q_weight: torch.Tensor,
1182
- perm: torch.Tensor,
1174
size_k: int,
1175
size_n: int,
1176
num_bits: int,
@@ -1194,14 +1185,13 @@ def gptq_marlin_moe_repack(
1185
)
1186
for e in range(num_experts):
1187
output[e] = torch.ops._C.gptq_marlin_repack(
1197
- b_q_weight[e], perm[e], size_k, size_n, num_bits, is_a_8bit
1188
+ b_q_weight[e], size_k, size_n, num_bits, is_a_8bit
1189
)
1190
return output
1191
1192
1193
def awq_marlin_moe_repack(
1194
b_q_weight: torch.Tensor,
1204
- perm: torch.Tensor,
1195
size_k: int,
1196
size_n: int,
1197
num_bits: int,
@@ -1238,14 +1228,11 @@ def marlin_gemm(
1228
a_scales: torch.Tensor | None,
1229
global_scale: torch.Tensor | None,
1230
b_zeros: torch.Tensor | None,
1241
- g_idx: torch.Tensor | None,
1242
- perm: torch.Tensor | None,
1231
workspace: torch.Tensor,
1232
b_q_type: ScalarType,
1233
size_m: int,
1234
size_n: int,
1235
size_k: int,
1248
- is_k_full: bool = True,
1236
use_atomic_add: bool = False,
1237
use_fp32_reduce: bool = False,
1238
is_zp_float: bool = False,
@@ -1259,14 +1246,11 @@ def marlin_gemm(
1246
a_scales,
1247
global_scale,
1248
b_zeros,
1262
- g_idx,
1263
- perm,
1249
workspace,
1250
b_q_type.id,
1251
size_m,
1252
size_n,
1253
size_k,
1269
- is_k_full,
1254
use_atomic_add,
1255
use_fp32_reduce,
1256
is_zp_float,
@@ -1285,14 +1269,11 @@ if hasattr(torch.ops._C, "marlin_gemm"):
1269
a_scales: torch.Tensor | None,
1270
global_scale: torch.Tensor | None,
1271
b_zeros: torch.Tensor | None,
1288
- g_idx: torch.Tensor | None,
1289
- perm: torch.Tensor | None,
1272
workspace: torch.Tensor,
1273
b_q_type_id: int,
1274
size_m: torch.SymInt,
1275
size_n: torch.SymInt,
1276
size_k: torch.SymInt,
1295
- is_k_full: bool = True,
1277
use_atomic_add: bool = False,
1278
use_fp32_reduce: bool = False,
1279
is_zp_float: bool = False,
@@ -1540,17 +1521,6 @@ if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b_grouped"):
1521
return torch.empty_like(b, memory_format=torch.contiguous_format)
1522
1523
1543
-def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
1544
- return torch.ops._C.permute_cols(a, perm)
1545
-
1546
-
1547
-if hasattr(torch.ops._C, "permute_cols"):
1548
-
1549
- @register_fake("_C::permute_cols")
1550
- def _permute_cols_fake(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
1551
- return torch.empty_like(a)
1552
-
1553
-
1524
# fp4
1525
def scaled_fp4_quant(
1526
input: torch.Tensor,
@@ -2545,8 +2515,6 @@ def moe_wna16_marlin_gemm(
2515
a_scales: torch.Tensor | None,
2516
global_scale: torch.Tensor | None,
2517
b_qzeros: torch.Tensor | None,
2548
- g_idx: torch.Tensor | None,
2549
- perm: torch.Tensor | None,
2518
workspace: torch.Tensor,
2519
sorted_token_ids: torch.Tensor,
2520
expert_ids: torch.Tensor,
@@ -2559,7 +2527,6 @@ def moe_wna16_marlin_gemm(
2527
size_m: int,
2528
size_n: int,
2529
size_k: int,
2562
- is_k_full: bool,
2530
use_atomic_add: bool,
2531
use_fp32_reduce: bool,
2532
is_zp_float: bool,
@@ -2576,8 +2543,6 @@ def moe_wna16_marlin_gemm(
2543
a_scales,
2544
global_scale,
2545
b_qzeros,
2579
- g_idx,
2580
- perm,
2546
workspace,
2547
sorted_token_ids,
2548
expert_ids,
@@ -2590,7 +2555,6 @@ def moe_wna16_marlin_gemm(
2555
size_m,
2556
size_n,
2557
size_k,
2593
- is_k_full,
2558
use_atomic_add,
2559
use_fp32_reduce,
2560
is_zp_float,
@@ -2612,8 +2576,6 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_
2576
a_scales: torch.Tensor | None,
2577
global_scale: torch.Tensor | None,
2578
b_qzeros: torch.Tensor | None,
2615
- g_idx: torch.Tensor | None,
2616
- perm: torch.Tensor | None,
2579
workspace: torch.Tensor,
2580
sorted_token_ids: torch.Tensor,
2581
expert_ids: torch.Tensor,
@@ -2626,7 +2588,6 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_
2588
size_m: int,
2589
size_n: int,
2590
size_k: int,
2629
- is_k_full: bool,
2591
use_atomic_add: bool,
2592
use_fp32_reduce: bool,
2593
is_zp_float: bool,
@@ -4103,7 +4064,6 @@ def cpu_gemm_wna16(
4064
q_weight: torch.Tensor,
4065
scales: torch.Tensor,
4066
zeros: torch.Tensor | None,
4106
- g_idx: torch.Tensor | None,
4067
bias: torch.Tensor | None,
4068
pack_factor: int,
4069
isa_hint: str,
@@ -4115,7 +4075,6 @@ def cpu_gemm_wna16(
4075
output,
4076
scales,
4077
zeros,
4118
- g_idx,
4078
bias,
4079
pack_factor,
4080
isa_hint,
vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py
+1
-8
@@ -19,7 +19,6 @@ class MPLinearLayerConfig:
19
act_type: torch.dtype
20
group_size: int
21
zero_points: bool
22
- has_g_idx: bool
22
out_type: torch.dtype | None = None
23
24
@@ -40,7 +39,6 @@ class MPLinearKernel(ABC):
39
w_q_param_name: str,
40
w_s_param_name: str,
41
w_zp_param_name: str | None = None,
43
- w_gidx_param_name: str | None = None,
42
) -> None:
43
assert self.can_implement(c)
44
self.config = c
@@ -48,10 +46,7 @@ class MPLinearKernel(ABC):
46
self.w_s_name = w_s_param_name
47
if c.zero_points:
48
assert w_zp_param_name is not None
51
- if c.has_g_idx:
52
- assert w_gidx_param_name is not None
49
self.w_zp_name: str | None = w_zp_param_name
54
- self.w_gidx_name: str | None = w_gidx_param_name
50
51
@abstractmethod
52
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -83,12 +78,10 @@ class MPLinearKernel(ABC):
78
) -> tuple[
79
torch.Tensor, # w_q
80
torch.Tensor, # w_s
86
- torch.Tensor | None, # w_zp,
87
- torch.Tensor | None, # w_gidx
81
+ torch.Tensor | None, # w_zp
82
]:
83
return (
84
getattr(layer, self.w_q_name),
85
getattr(layer, self.w_s_name),
86
getattr(layer, self.w_zp_name or "", None),
93
- getattr(layer, self.w_gidx_name or "", None),
87
)
vllm/model_executor/kernels/linear/mixed_precision/allspark.py
+1
-4
@@ -23,9 +23,6 @@ class AllSparkLinearKernel(MPLinearKernel):
23
24
@classmethod
25
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
26
- if c.has_g_idx:
27
- return False, "Act reordering currently not supported by AllSpark"
28
-
26
if c.zero_points:
27
return False, "Zero points currently not supported by AllSpark"
28
@@ -91,7 +88,7 @@ class AllSparkLinearKernel(MPLinearKernel):
88
) -> torch.Tensor:
89
c = self.config
90
gemm_args = self.gemm_args
94
- w_q, w_s, _, _ = self._get_weight_params(layer)
91
+ w_q, w_s, _ = self._get_weight_params(layer)
92
93
reshaped_x = x.reshape(-1, x.shape[-1])
94
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
vllm/model_executor/kernels/linear/mixed_precision/conch.py
+1
-7
@@ -43,12 +43,6 @@ class ConchLinearKernel(MPLinearKernel):
43
)
44
return False, error_msg
45
46
- if c.has_g_idx:
47
- return (
48
- False,
49
- "Activation reordering (g_idx) is not supported by ConchLinearKernel",
50
- )
51
-
46
if find_spec("conch") is None:
47
error_msg = (
48
"conch-triton-kernels is not installed, please "
@@ -130,7 +124,7 @@ class ConchLinearKernel(MPLinearKernel):
124
) -> torch.Tensor:
125
from conch.ops.quantization.gemm import mixed_precision_gemm
126
133
- w_q, w_s, w_zp, _ = self._get_weight_params(layer)
127
+ w_q, w_s, w_zp = self._get_weight_params(layer)
128
129
# Map channelwise group_size=-1 to the actual input dimension K.
130
# The conch kernel computes stride_mul = block_k / group_size;
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+1
-6
@@ -144,9 +144,6 @@ class CPUWNA16LinearKernel(MPLinearKernel):
144
if (not self.config.zero_points) and (self.w_zp_name is not None):
145
setattr(layer, self.w_zp_name, None)
146
147
- if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
148
- setattr(layer, self.w_gidx_name, None)
149
-
147
weights = getattr(layer, self.w_q_name)
148
# Require GPTQ pack format
149
assert weights.input_dim == weights.packed_dim
@@ -172,7 +169,6 @@ class CPUWNA16LinearKernel(MPLinearKernel):
169
supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV
170
layer.use_w4a8 = (
171
envs.VLLM_CPU_INT4_W4A8
175
- and not self.config.has_g_idx
172
and self.config.act_type == torch.bfloat16
173
and (supports_amx or supports_riscv)
174
)
@@ -189,7 +185,7 @@ class CPUWNA16LinearKernel(MPLinearKernel):
185
x: torch.Tensor,
186
bias: torch.Tensor | None = None,
187
) -> torch.Tensor:
192
- w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
188
+ w_q, w_s, w_zp = self._get_weight_params(layer)
189
if layer.use_w4a8:
190
x = ops.int4_scaled_mm_cpu(
191
x=x,
@@ -204,7 +200,6 @@ class CPUWNA16LinearKernel(MPLinearKernel):
200
q_weight=w_q,
201
scales=w_s,
202
zeros=w_zp,
207
- g_idx=w_gidx,
203
bias=bias,
204
pack_factor=8, # 32 // 4
205
isa_hint=layer.isa_hint,
vllm/model_executor/kernels/linear/mixed_precision/cutlass.py
+1
-4
@@ -39,9 +39,6 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
39
if c.act_type != torch.float8_e4m3fn:
40
return False, "CUTLASS W4A8 only supports FP8 (e4m3) activations"
41
42
- if c.has_g_idx:
43
- return False, "Act reordering not supported by CUTLASS W4A8"
44
-
42
if c.zero_points:
43
return False, "Zero points not supported by CUTLASS W4A8"
44
@@ -109,7 +106,7 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
106
bias: torch.Tensor | None = None,
107
) -> torch.Tensor:
108
c = self.config
112
- w_q, w_s, _, _ = self._get_weight_params(layer)
109
+ w_q, w_s, _ = self._get_weight_params(layer)
110
w_ch_s = layer.weight_chan_scale
111
112
x_2d = x.reshape(-1, x.shape[-1])
vllm/model_executor/kernels/linear/mixed_precision/exllama.py
+3
-31
@@ -32,14 +32,6 @@ class ExllamaLinearKernel(MPLinearKernel):
32
"Exllama is only supported on CUDA and ROCm",
33
)
34
35
- if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
36
- return (
37
- False,
38
- "Act reordering currently not supported by Exllama, "
39
- "when the input features are partitioned across "
40
- "devices",
41
- )
42
-
35
if c.partition_weight_shape[1] % (32 // c.weight_type.size_bits) != 0:
36
return (
37
False,
@@ -111,29 +103,11 @@ class ExllamaLinearKernel(MPLinearKernel):
103
layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
104
)
105
114
- if c.has_g_idx:
115
-
116
- def transform_w_g_idx(x):
117
- # Exllama wants the permutation array instead of the group
118
- # indices
119
- return torch.argsort(x).to(torch.int)
120
-
121
- self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore
122
- else:
123
- self.w_gidx_name = "g_idx"
124
- empty_g_idx = torch.nn.Parameter(
125
- torch.empty((0,), dtype=torch.int, device=device), requires_grad=False
126
- )
127
- setattr(layer, self.w_gidx_name, empty_g_idx)
128
-
106
def transform_w_q(x):
107
assert isinstance(x, BasevLLMParameter)
131
- assert self.w_gidx_name is not None
132
- g_idx = getattr(layer, self.w_gidx_name)
133
-
108
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
109
x_cont = x.data.contiguous()
136
- ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits)
110
+ ops.gptq_shuffle(x_cont, c.weight_type.size_bits)
111
return x_cont
112
113
def transform_w_s(x):
@@ -157,17 +131,15 @@ class ExllamaLinearKernel(MPLinearKernel):
131
x_2d = x.reshape(-1, x.shape[-1])
132
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
133
160
- w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer)
161
-
134
+ w_q, w_s, w_zp = self._get_weight_params(layer)
135
# gptq_gemm supports GPTQv2 format by passing use_v2_format=True.
136
# However, the MPLinearLayerConfig doesn't contain format info.
137
# So hardcode GPTQv1 format here, to keep its behavior unchanged.
138
use_v2_format = False
139
140
assert w_zp is not None, "Zero points are required by Exllama"
168
- assert w_g_idx is not None, "Group index is required by Exllama"
141
output = ops.gptq_gemm(
170
- x_2d, w_q, w_zp, w_s, w_g_idx, True, use_v2_format, c.weight_type.size_bits
142
+ x_2d, w_q, w_zp, w_s, True, use_v2_format, c.weight_type.size_bits
143
)
144
145
if bias is not None:
vllm/model_executor/kernels/linear/mixed_precision/humming.py
-2
@@ -21,8 +21,6 @@ class HummingLinearKernel(MPLinearKernel):
21
return False, "Humming is only supported on CUDA"
22
if not has_humming():
23
return False, "Humming is not installed"
24
- if c.has_g_idx:
25
- return False, "Humming does not support act-order (g_idx)"
24
return True, None
25
26
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/kernels/linear/mixed_precision/machete.py
+1
-34
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from functools import partial
4
5
import torch
6
@@ -12,7 +11,6 @@ from vllm.model_executor.layers.quantization.utils.machete_utils import (
11
query_machete_supported_quant_types,
12
)
13
from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
- pack_quantized_values_into_int32,
14
unpack_quantized_values_into_int32,
15
)
16
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
@@ -35,14 +33,6 @@ class MacheteLinearKernel(MPLinearKernel):
33
if not current_platform.is_device_capability(90):
34
return False, "Machete requires compute capability of 90 (Hopper)"
35
38
- if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
39
- return (
40
- False,
41
- "Act reordering currently not supported by Machete, "
42
- "when the input features are partitioned across "
43
- "devices",
44
- )
45
-
36
if c.weight_type not in query_machete_supported_quant_types(c.zero_points):
37
return (
38
False,
@@ -70,29 +60,9 @@ class MacheteLinearKernel(MPLinearKernel):
60
def process_weights_after_loading(self, layer: torch.nn.Module):
61
c = self.config
62
73
- if c.has_g_idx:
74
- assert self.w_gidx_name is not None
75
- perm = torch.argsort(getattr(layer, self.w_gidx_name)).to(torch.int)
76
-
77
- self.act_perm = lambda x: x[:, perm]
78
- # use `ops.permute_cols` if possible
79
- if (
80
- c.act_type in [torch.float16, torch.bfloat16]
81
- and c.partition_weight_shape[0] % 8 == 0
82
- ):
83
- self.act_perm = partial(ops.permute_cols, perm=perm)
84
-
63
def transform_w_q(x):
64
assert isinstance(x, BasevLLMParameter)
65
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
88
- if c.has_g_idx:
89
- x_unpacked = unpack_quantized_values_into_int32(
90
- x.data, c.weight_type, packed_dim=0
91
- )
92
- x_perm = x_unpacked[perm, :]
93
- x.data = pack_quantized_values_into_int32(
94
- x_perm, c.weight_type, packed_dim=0
95
- )
66
x.data = ops.machete_prepack_B(
67
x.data.t().contiguous().t(),
68
a_type=c.act_type,
@@ -131,14 +101,11 @@ class MacheteLinearKernel(MPLinearKernel):
101
bias: torch.Tensor | None = None,
102
) -> torch.Tensor:
103
c = self.config
134
- w_q, w_s, w_zp, _ = self._get_weight_params(layer)
104
+ w_q, w_s, w_zp = self._get_weight_params(layer)
105
106
x_2d = x.reshape(-1, x.shape[-1])
107
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
108
139
- if c.has_g_idx:
140
- x_2d = self.act_perm(x_2d)
141
-
109
if c.zero_points:
110
assert w_zp is not None
111
else:
vllm/model_executor/kernels/linear/mixed_precision/marlin.py
+5
-48
@@ -8,10 +8,8 @@ from vllm import _custom_ops as ops
8
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
9
MARLIN_SUPPORTED_GROUP_SIZES,
10
apply_gptq_marlin_linear,
11
- check_marlin_supports_shape,
11
marlin_act_int8_process_scales,
13
- marlin_is_k_full,
14
- marlin_make_empty_g_idx,
12
+ marlin_make_empty,
13
marlin_make_workspace_new,
14
marlin_pad_dim,
15
marlin_pad_qweight,
@@ -19,13 +17,11 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
17
marlin_padded_nk,
18
marlin_permute_bias,
19
marlin_permute_scales,
22
- marlin_sort_g_idx,
20
marlin_zero_points,
21
query_marlin_supported_quant_types,
22
unpack_cols,
23
)
24
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
28
-from vllm.model_executor.utils import replace_parameter
25
from vllm.platforms import current_platform
26
from vllm.scalar_type import scalar_types
27
@@ -59,16 +55,6 @@ class MarlinLinearKernel(MPLinearKernel):
55
f"{MARLIN_SUPPORTED_GROUP_SIZES}",
56
)
57
62
- if c.has_g_idx:
63
- # Act-order couples K to the full-model group layout, so tile
64
- # padding is not supported; keep the strict shape check.
65
- return check_marlin_supports_shape(
66
- c.partition_weight_shape[1], # out_features
67
- c.partition_weight_shape[0], # in_features
68
- c.full_weight_shape[0], # in_features
69
- c.group_size,
70
- )
71
-
58
# A group straddling TP ranks cannot be fixed by padding.
59
if (
60
c.group_size != -1
@@ -102,25 +88,16 @@ class MarlinLinearKernel(MPLinearKernel):
88
getattr(layer, self.w_s_name).data * 512
89
)
90
105
- row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0]
106
- self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel)
107
-
91
size_k, size_n = c.partition_weight_shape
109
- if c.has_g_idx:
110
- # Act-order shapes were strictly validated in can_implement.
111
- padded_n, padded_k = size_n, size_k
112
- else:
113
- padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
92
+ padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
93
94
# Allocate marlin workspace, reusing existing storage on reload.
95
self.workspace = marlin_make_workspace_new(
96
device, existing=getattr(self, "workspace", None)
97
)
98
120
- # Default names since marlin requires empty parameters for these,
99
+ # Default name since marlin requires empty parameter for zp,
100
# TODO: remove this requirement from marlin (allow optional tensors)
122
- if self.w_gidx_name is None:
123
- self.w_gidx_name = "g_idx"
101
if self.w_zp_name is None:
102
self.w_zp_name = "w_zp"
103
@@ -131,7 +108,6 @@ class MarlinLinearKernel(MPLinearKernel):
108
marlin_pad_qweight(
109
x.data.contiguous(), size_n, size_k, padded_n, padded_k
110
),
134
- perm=layer.g_idx_sort_indices,
111
size_k=padded_k,
112
size_n=padded_n,
113
num_bits=c.weight_type.size_bits,
@@ -172,18 +148,6 @@ class MarlinLinearKernel(MPLinearKernel):
148
layer.input_global_scale = None
149
return x
150
175
- if c.has_g_idx:
176
- g_idx, g_idx_sort_indices = marlin_sort_g_idx(
177
- getattr(layer, self.w_gidx_name)
178
- )
179
- self._transform_param(layer, self.w_gidx_name, lambda _: g_idx)
180
- replace_parameter(
181
- layer, "g_idx_sort_indices", g_idx_sort_indices, prefer_copy=True
182
- )
183
- else:
184
- setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device))
185
- layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
186
-
151
if c.zero_points:
152
grouped_k = size_k // c.group_size if c.group_size != -1 else 1
153
padded_grouped_k = padded_k // c.group_size if c.group_size != -1 else 1
@@ -211,7 +175,7 @@ class MarlinLinearKernel(MPLinearKernel):
175
),
176
)
177
else:
214
- setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device))
178
+ setattr(layer, self.w_zp_name, marlin_make_empty(device))
179
self._transform_param(layer, self.w_q_name, transform_w_q)
180
self._transform_param(layer, self.w_s_name, transform_w_s)
181
@@ -227,23 +191,16 @@ class MarlinLinearKernel(MPLinearKernel):
191
bias: torch.Tensor | None = None,
192
) -> torch.Tensor:
193
c = self.config
230
- w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
231
-
232
- # `process_weights_after_loading` will ensure w_zp and w_gidx are not
233
- # None for marlin
234
-
194
+ w_q, w_s, w_zp = self._get_weight_params(layer)
195
return apply_gptq_marlin_linear(
196
input=x,
197
weight=w_q,
198
weight_scale=w_s,
199
weight_zp=w_zp, # type: ignore
240
- g_idx=w_gidx, # type: ignore
241
- g_idx_sort_indices=layer.g_idx_sort_indices,
200
workspace=self.workspace,
201
wtype=c.weight_type,
202
input_size_per_partition=c.partition_weight_shape[0],
203
output_size_per_partition=c.partition_weight_shape[1],
246
- is_k_full=self.is_k_full,
204
input_global_scale=getattr(layer, "input_global_scale", None),
205
bias=bias,
206
input_dtype=c.act_type,
vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py
+3
-33
@@ -86,13 +86,6 @@ class RDNA3W4A16LinearKernel(MPLinearKernel):
86
"W4A16 kernel (qzeros packing)",
87
)
88
89
- if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
90
- return (
91
- False,
92
- "Act-order with TP-partitioned input features is not "
93
- "supported by the RDNA3 W4A16 kernel",
94
- )
95
-
89
return True, None
90
91
# ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) -----
@@ -127,34 +120,14 @@ class RDNA3W4A16LinearKernel(MPLinearKernel):
120
layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
121
)
122
130
- # Act-order: convert g_idx to the inverse permutation array exllama
131
- # expects (kernel reads a[perm[k]] instead of using groups indirected
132
- # by g_idx[k]).
133
- if c.has_g_idx:
134
-
135
- def transform_w_g_idx(x):
136
- return torch.argsort(x).to(torch.int)
137
-
138
- self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore
139
- else:
140
- self.w_gidx_name = "g_idx"
141
- empty_g_idx = torch.nn.Parameter(
142
- torch.empty((0,), dtype=torch.int, device=device),
143
- requires_grad=False,
144
- )
145
- setattr(layer, self.w_gidx_name, empty_g_idx)
146
-
123
def transform_w_q(x):
124
assert isinstance(x, BasevLLMParameter)
149
- assert self.w_gidx_name is not None
150
- g_idx = getattr(layer, self.w_gidx_name)
151
-
125
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
126
x_cont = x.data.contiguous()
127
# Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in
128
# the same shuffled int32 layout and uses the (qa & 0x000F000F)
129
# bit-trick on top.
157
- ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits)
130
+ ops.gptq_shuffle(x_cont, c.weight_type.size_bits)
131
return x_cont
132
133
def transform_w_s(x):
@@ -181,12 +154,9 @@ class RDNA3W4A16LinearKernel(MPLinearKernel):
154
x_2d = x.reshape(-1, x.shape[-1])
155
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
156
184
- w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer)
185
-
157
+ w_q, w_s, w_zp = self._get_weight_params(layer)
158
assert w_zp is not None, "Zero points are required by RDNA3 W4A16"
187
- assert w_g_idx is not None, "g_idx tensor (possibly empty) required"
188
-
189
- output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, w_g_idx, False)
159
+ output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, False)
160
161
if bias is not None:
162
output.add_(bias)
vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
+4
-7
@@ -142,13 +142,13 @@ def _triton_w4a16_skinny_fmt_kernel(
142
b = tl.interleave(b, b)
143
b = (b >> shifts_full) & 0xF
144
145
- g_idx = (k_start * BLOCK_K) // group_size
146
- scale_ptrs = scales_ptr + offs_n * num_groups + g_idx
145
+ group_idx = (k_start * BLOCK_K) // group_size
146
+ scale_ptrs = scales_ptr + offs_n * num_groups + group_idx
147
scale_mask = offs_n < N
148
scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0)
149
150
if HAS_ZP:
151
- zp_ptrs = zp_ptr + offs_n * num_groups + g_idx
151
+ zp_ptrs = zp_ptr + offs_n * num_groups + group_idx
152
zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0)
153
b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None]
154
else:
@@ -476,9 +476,6 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel):
476
if c.act_type not in (torch.float16, torch.bfloat16):
477
return False, "requires float16 or bfloat16 activations"
478
479
- if c.has_g_idx:
480
- return False, "does not support g_idx reordering"
481
-
479
gs = c.group_size
480
if gs not in SUPPORTED_GROUP_SIZES:
481
return (
@@ -541,7 +538,7 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel):
538
from vllm.utils.platform_utils import num_compute_units
539
540
c = self.config
544
- w_q, w_s, w_zp, _ = self._get_weight_params(layer)
541
+ w_q, w_s, w_zp = self._get_weight_params(layer)
542
543
x_2d = x.reshape(-1, x.shape[-1])
544
N = w_q.shape[0]
vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+5
-12
@@ -126,10 +126,10 @@ def triton_w4a16_gemm_kernel(
126
b = (b >> shifts) & 0xF
127
128
# ---- Compute scale/zero group row index ----
129
- g_idx = (k_start * BLOCK_K) // group_size
129
+ group_idx = (k_start * BLOCK_K) // group_size
130
131
# ---- Load scales: [BLOCK_N] → broadcast to [BLOCK_K, BLOCK_N] ----
132
- scale_offset = g_idx * N + offs_sn
132
+ scale_offset = group_idx * N + offs_sn
133
scale_mask = offs_sn < N
134
scales = tl.load(scales_ptr + scale_offset, mask=scale_mask, other=1.0)
135
scales = tl.broadcast_to(scales[None, :], (BLOCK_K, BLOCK_N))
@@ -137,7 +137,7 @@ def triton_w4a16_gemm_kernel(
137
# ---- Load / compute zeros ----
138
if HAS_ZP:
139
# Load packed zeros row: [BLOCK_N//8] int32
140
- zero_offset = g_idx * (N // 8) + offs_bn
140
+ zero_offset = group_idx * (N // 8) + offs_bn
141
zero_mask = offs_bn < N // 8
142
z_packed = tl.load(zeros_ptr + zero_offset, mask=zero_mask, other=0)
143
# Unpack to [BLOCK_N] using same interleave+shift pattern
@@ -237,7 +237,7 @@ def _triton_w4a16_gemm_impl(
237
BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
238
239
# The kernel loads scales/zeros for a single group per BLOCK_K tile
240
- # (one g_idx per iteration). If BLOCK_K > group_size, rows at the tail
240
+ # (one group index per iteration). If BLOCK_K > group_size, rows at the tail
241
# of the tile dequantize with the wrong group's scales, silently
242
# corrupting the output. Clamp BLOCK_K to group_size to keep one
243
# scale group per tile.
@@ -340,13 +340,6 @@ class TritonW4A16LinearKernel(MPLinearKernel):
340
"(8 int4 values packed per int32)",
341
)
342
343
- if c.has_g_idx:
344
- return (
345
- False,
346
- "Activation reordering (g_idx) is not supported by "
347
- "TritonW4A16LinearKernel",
348
- )
349
-
343
gs = c.group_size
344
if (
345
gs not in TRITON_W4A16_SUPPORTED_GROUP_SIZES
@@ -471,7 +464,7 @@ class TritonW4A16LinearKernel(MPLinearKernel):
464
self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
465
) -> torch.Tensor:
466
c = self.config
474
- w_q, w_s, w_zp, _ = self._get_weight_params(layer)
467
+ w_q, w_s, w_zp = self._get_weight_params(layer)
468
469
x_2d = x.reshape(-1, x.shape[-1]).contiguous()
470
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+4
-17
@@ -54,10 +54,8 @@ class XPUwNa16LinearKernel(MPLinearKernel):
54
return True, None
55
56
def process_weights_after_loading(self, layer: torch.nn.Module):
57
- # Default names since marlin requires empty parameters for these,
57
+ # Default name since marlin requires empty parameter for zp,
58
# TODO: remove this requirement from marlin (allow optional tensors)
59
- if self.w_gidx_name is None:
60
- self.w_gidx_name = "g_idx"
59
if self.w_zp_name is None:
60
self.w_zp_name = "w_zp"
61
@@ -90,17 +88,6 @@ class XPUwNa16LinearKernel(MPLinearKernel):
88
setattr(
89
layer, self.w_zp_name, Parameter(weight_zero_point, requires_grad=False)
90
)
93
- if self.config.has_g_idx:
94
- setattr(
95
- layer,
96
- self.w_gidx_name,
97
- Parameter(
98
- getattr(layer, self.w_gidx_name).data.t().contiguous(),
99
- requires_grad=False,
100
- ),
101
- )
102
- else:
103
- setattr(layer, self.w_gidx_name, None)
91
92
def apply_weights(
93
self,
@@ -109,7 +96,7 @@ class XPUwNa16LinearKernel(MPLinearKernel):
96
bias: torch.Tensor | None = None,
97
) -> torch.Tensor:
98
reshaped_x = x.reshape(-1, x.shape[-1])
112
- w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
99
+ w_q, w_s, w_zp = self._get_weight_params(layer)
100
out = torch.ops._xpu_C.int4_gemm_w4a16(
101
reshaped_x,
102
w_q.t(),
@@ -117,7 +104,7 @@ class XPUwNa16LinearKernel(MPLinearKernel):
104
w_s,
105
w_zp,
106
self.config.group_size,
120
- w_gidx,
107
+ None, # Retained by the external XPU op ABI.
108
)
109
return out
110
@@ -224,7 +211,7 @@ class XPUW4A8IntLinearKernel(MPLinearKernel):
211
layer.weight_scale,
212
layer.weight_zero_point,
213
self.config.group_size,
227
- None, # g_idx not currently supported
214
+ None, # Retained by the external XPU op ABI.
215
bias,
216
)
217
vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
+1
-12
@@ -52,8 +52,6 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
52
"zentorch_woq_linear} are not registered.",
53
)
54
55
- if c.has_g_idx:
56
- return False, "ZentorchWNA16 does not support activation re-ordering."
55
return True, None
56
57
def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool:
@@ -62,12 +60,6 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
60
Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()``
61
with ``layer`` untouched).
62
"""
65
- if (
66
- self.w_gidx_name is not None
67
- and getattr(layer, self.w_gidx_name, None) is not None
68
- ) or (getattr(self.config, "has_g_idx", False)):
69
- return False
70
-
63
weight_packed = getattr(layer, self.w_q_name, None)
64
weight_scale = getattr(layer, self.w_s_name, None)
65
if weight_packed is None or weight_scale is None:
@@ -112,7 +104,7 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
104
if not self._zentorch_woq_eligible(layer):
105
logger.info_once(
106
"[zen_cpu] ZentorchWNA16 fast path not eligible for this "
115
- "layer (AWQ pack layout, g_idx, or non-int32 storage); "
107
+ "layer (AWQ pack layout or non-int32 storage); "
108
"falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)."
109
)
110
super().process_weights_after_loading(layer)
@@ -121,9 +113,6 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
113
if (not self.config.zero_points) and (self.w_zp_name is not None):
114
setattr(layer, self.w_zp_name, None)
115
124
- if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
125
- setattr(layer, self.w_gidx_name, None)
126
-
116
weight_q = getattr(layer, self.w_q_name)
117
weight_s = getattr(layer, self.w_s_name)
118
weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
-72
@@ -81,10 +81,6 @@ def _fused_marlin_moe(
81
input_global_scale2: torch.Tensor | None = None,
82
global_scale1: torch.Tensor | None = None,
83
global_scale2: torch.Tensor | None = None,
84
- g_idx1: torch.Tensor | None = None,
85
- g_idx2: torch.Tensor | None = None,
86
- sort_indices1: torch.Tensor | None = None,
87
- sort_indices2: torch.Tensor | None = None,
84
w1_zeros: torch.Tensor | None = None,
85
w2_zeros: torch.Tensor | None = None,
86
workspace: torch.Tensor | None = None,
@@ -92,7 +88,6 @@ def _fused_marlin_moe(
88
intermediate_cache2: torch.Tensor | None = None,
89
output: torch.Tensor | None = None,
90
input_dtype: torch.dtype | None = None,
95
- is_k_full: bool = True,
91
activation_config: ApplyMoEActivationConfig | None = None,
92
) -> torch.Tensor:
93
assert hidden_states.ndim == 2
@@ -142,8 +137,6 @@ def _fused_marlin_moe(
137
a_scales1,
138
global_scale1,
139
w1_zeros,
145
- g_idx1,
146
- sort_indices1,
140
workspace,
141
sorted_token_ids,
142
expert_ids,
@@ -156,7 +149,6 @@ def _fused_marlin_moe(
149
size_m=M,
150
size_n=w13_num_shards * N,
151
size_k=K,
159
- is_k_full=is_k_full,
152
use_atomic_add=False,
153
use_fp32_reduce=True,
154
is_zp_float=False,
@@ -209,8 +201,6 @@ def _fused_marlin_moe(
201
a_scales2,
202
global_scale2,
203
w2_zeros,
212
- g_idx2,
213
- sort_indices2,
204
workspace,
205
sorted_token_ids,
206
expert_ids,
@@ -223,7 +213,6 @@ def _fused_marlin_moe(
213
size_m=M * num_topk,
214
size_n=K,
215
size_k=N,
226
- is_k_full=is_k_full,
216
use_atomic_add=False,
217
use_fp32_reduce=True,
218
is_zp_float=False,
@@ -253,16 +242,11 @@ def fused_marlin_moe(
242
input_global_scale2: torch.Tensor | None = None,
243
global_scale1: torch.Tensor | None = None,
244
global_scale2: torch.Tensor | None = None,
256
- g_idx1: torch.Tensor | None = None,
257
- g_idx2: torch.Tensor | None = None,
258
- sort_indices1: torch.Tensor | None = None,
259
- sort_indices2: torch.Tensor | None = None,
245
w1_zeros: torch.Tensor | None = None,
246
w2_zeros: torch.Tensor | None = None,
247
workspace: torch.Tensor | None = None,
248
intermediate_cache13: torch.Tensor | None = None,
249
intermediate_cache2: torch.Tensor | None = None,
265
- is_k_full: bool = True,
250
output: torch.Tensor | None = None,
251
input_dtype: torch.dtype | None = None,
252
activation_config: ApplyMoEActivationConfig | None = None,
@@ -277,12 +261,6 @@ def fused_marlin_moe(
261
- w2 (torch.Tensor): The second set of expert weights.
262
- w1_scale (torch.Tensor): Scale to be used for w1.
263
- w2_scale (torch.Tensor): Scale to be used for w2.
280
- - g_idx1 (torch.Tensor|None): The first set of act_order indices.
281
- - g_idx2 (torch.Tensor|None): The second set of act_order indices.
282
- - sort_indices1 (torch.Tensor|None): The first act_order input
283
- permutation.
284
- - sort_indices2 (torch.Tensor|None): The second act_order input
285
- permutation.
264
- topk_weights (torch.Tensor): Top-k weights.
265
- topk_ids (torch.Tensor): Indices of topk-k elements.
266
- w1_zeros (torch.Tensor|None): Optional zero points to be used for w1.
@@ -372,10 +350,6 @@ def fused_marlin_moe(
350
input_global_scale2=input_global_scale2,
351
global_scale1=global_scale1,
352
global_scale2=global_scale2,
375
- g_idx1=g_idx1,
376
- g_idx2=g_idx2,
377
- sort_indices1=sort_indices1,
378
- sort_indices2=sort_indices2,
353
w1_zeros=w1_zeros,
354
w2_zeros=w2_zeros,
355
workspace=workspace,
@@ -383,7 +357,6 @@ def fused_marlin_moe(
357
intermediate_cache2=intermediate_cache2,
358
output=None,
359
input_dtype=input_dtype,
386
- is_k_full=is_k_full,
360
).view(-1, topk, K)
361
362
if output is None:
@@ -416,16 +389,11 @@ def batched_fused_marlin_moe(
389
input_global_scale2: torch.Tensor | None = None,
390
global_scale1: torch.Tensor | None = None,
391
global_scale2: torch.Tensor | None = None,
419
- g_idx1: torch.Tensor | None = None,
420
- g_idx2: torch.Tensor | None = None,
421
- sort_indices1: torch.Tensor | None = None,
422
- sort_indices2: torch.Tensor | None = None,
392
w1_zeros: torch.Tensor | None = None,
393
w2_zeros: torch.Tensor | None = None,
394
workspace: torch.Tensor | None = None,
395
intermediate_cache13: torch.Tensor | None = None,
396
intermediate_cache2: torch.Tensor | None = None,
428
- is_k_full: bool = True,
397
output: torch.Tensor | None = None,
398
input_dtype: torch.dtype | None = None,
399
activation_func: Callable[..., None] | None = None,
@@ -548,10 +516,6 @@ def batched_fused_marlin_moe(
516
input_global_scale2=input_global_scale2,
517
global_scale1=global_scale1,
518
global_scale2=global_scale2,
551
- g_idx1=g_idx1,
552
- g_idx2=g_idx2,
553
- sort_indices1=sort_indices1,
554
- sort_indices2=sort_indices2,
519
w1_zeros=w1_zeros,
520
w2_zeros=w2_zeros,
521
workspace=workspace,
@@ -559,7 +523,6 @@ def batched_fused_marlin_moe(
523
intermediate_cache2=intermediate_cache2,
524
output=output.view(-1, K) if output is not None else output,
525
input_dtype=input_dtype,
562
- is_k_full=is_k_full,
526
)
527
528
output = output.view(B, BATCH_TOKENS_MAX, K)
@@ -574,11 +537,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
537
quant_config: FusedMoEQuantConfig,
538
max_num_tokens: int | None = None,
539
num_dispatchers: int | None = None,
577
- w13_g_idx: torch.Tensor | None = None,
578
- w2_g_idx: torch.Tensor | None = None,
579
- w13_g_idx_sort_indices: torch.Tensor | None = None,
580
- w2_g_idx_sort_indices: torch.Tensor | None = None,
581
- is_k_full: bool = True,
540
):
541
# TODO (varun) : Enable activation quantization
542
assert (
@@ -588,11 +546,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
546
or quant_config.use_int8_w8a16
547
or quant_config.use_fp8_w8a16
548
), "Supports only {mxfp,nvfp,int}4_w4a16, int8_w8a16 or fp8_w8a16"
591
- self.w13_g_idx = w13_g_idx
592
- self.w2_g_idx = w2_g_idx
593
- self.w13_g_idx_sort_indices = w13_g_idx_sort_indices
594
- self.w2_g_idx_sort_indices = w2_g_idx_sort_indices
595
- self.is_k_full = is_k_full
549
self.input_dtype = get_marlin_input_dtype()
550
551
super().__init__(
@@ -786,11 +739,6 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
739
# output buffer allocation. Please refer to workspace_shapes().
740
intermediate_cache13=workspace2,
741
intermediate_cache2=workspace13,
789
- g_idx1=self.w13_g_idx,
790
- g_idx2=self.w2_g_idx,
791
- sort_indices1=self.w13_g_idx_sort_indices,
792
- sort_indices2=self.w2_g_idx_sort_indices,
793
- is_k_full=self.is_k_full,
742
input_dtype=self.input_dtype,
743
)
744
return
@@ -901,11 +849,6 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
849
output=output,
850
intermediate_cache13=workspace2,
851
intermediate_cache2=workspace13,
904
- g_idx1=self.w13_g_idx,
905
- g_idx2=self.w2_g_idx,
906
- sort_indices1=self.w13_g_idx_sort_indices,
907
- sort_indices2=self.w2_g_idx_sort_indices,
908
- is_k_full=self.is_k_full,
852
input_dtype=self.input_dtype,
853
)
854
@@ -931,22 +874,12 @@ class BatchedMarlinExperts(MarlinExpertsBase):
874
quant_config: FusedMoEQuantConfig,
875
max_num_tokens: int,
876
num_dispatchers: int,
934
- w13_g_idx: torch.Tensor | None = None,
935
- w2_g_idx: torch.Tensor | None = None,
936
- w13_g_idx_sort_indices: torch.Tensor | None = None,
937
- w2_g_idx_sort_indices: torch.Tensor | None = None,
938
- is_k_full: bool = True,
877
):
878
super().__init__(
879
moe_config=moe_config,
880
quant_config=quant_config,
881
max_num_tokens=max_num_tokens,
882
num_dispatchers=num_dispatchers,
945
- w13_g_idx=w13_g_idx,
946
- w2_g_idx=w2_g_idx,
947
- w13_g_idx_sort_indices=w13_g_idx_sort_indices,
948
- w2_g_idx_sort_indices=w2_g_idx_sort_indices,
949
- is_k_full=is_k_full,
883
)
884
885
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
@@ -1038,14 +971,9 @@ class BatchedMarlinExperts(MarlinExpertsBase):
971
output=output,
972
intermediate_cache13=workspace13,
973
intermediate_cache2=workspace2,
1041
- g_idx1=self.w13_g_idx,
1042
- g_idx2=self.w2_g_idx,
1043
- sort_indices1=self.w13_g_idx_sort_indices,
1044
- sort_indices2=self.w2_g_idx_sort_indices,
974
w1_zeros=self.w1_zp,
975
w2_zeros=self.w2_zp,
976
input_dtype=self.input_dtype,
1048
- is_k_full=self.is_k_full,
977
activation_func=activation_func,
978
activation_config=self.activation_config,
979
)
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+5
-150
@@ -22,7 +22,6 @@ from vllm.model_executor.layers.fused_moe.config import (
22
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
23
BatchedMarlinExperts,
24
MarlinExperts,
25
- MarlinExpertsBase,
25
)
26
from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
27
TritonWNA16Experts,
@@ -362,11 +361,6 @@ def make_wna16_moe_kernel(
361
moe_config: FusedMoEConfig,
362
experts_cls: type[mk.FusedMoEExperts],
363
backend: WNA16MoEBackend = WNA16MoEBackend.MARLIN,
365
- is_k_full: bool = False,
366
- w13_g_idx: torch.Tensor | None = None,
367
- w2_g_idx: torch.Tensor | None = None,
368
- w13_g_idx_sort_indices: torch.Tensor | None = None,
369
- w2_g_idx_sort_indices: torch.Tensor | None = None,
364
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
365
) -> mk.FusedMoEKernel:
366
from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -413,15 +407,6 @@ def make_wna16_moe_kernel(
407
logger.info_once("Using %s", experts_cls.__name__, scope="local")
408
409
extra_args: dict[str, Any] = {}
416
- if issubclass(experts_cls, MarlinExpertsBase):
417
- extra_args = {
418
- "w13_g_idx": w13_g_idx,
419
- "w2_g_idx": w2_g_idx,
420
- "w13_g_idx_sort_indices": w13_g_idx_sort_indices,
421
- "w2_g_idx_sort_indices": w2_g_idx_sort_indices,
422
- "is_k_full": is_k_full,
423
- }
424
-
410
if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
411
max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
412
assert max_num_tokens is not None
@@ -450,8 +435,6 @@ def _process_weights_flashinfer(
435
w2_qweight: torch.Tensor,
436
w13_scales: torch.Tensor,
437
w2_scales: torch.Tensor,
453
- w13_g_idx: torch.Tensor,
454
- w2_g_idx: torch.Tensor,
438
w13_bias: torch.Tensor | None = None,
439
w2_bias: torch.Tensor | None = None,
440
) -> tuple[
@@ -459,10 +442,6 @@ def _process_weights_flashinfer(
442
torch.Tensor, # w2_qweight
443
torch.Tensor, # w13_scales
444
torch.Tensor, # w2_scales
462
- torch.Tensor, # w13_g_idx
463
- torch.Tensor, # w2_g_idx
464
- torch.Tensor | None, # w13_g_idx_sort_indices
465
- torch.Tensor | None, # w2_g_idx_sort_indices
445
torch.Tensor | None, # w13_qzeros
446
torch.Tensor | None, # w2_qzeros
447
torch.Tensor | None, # w13_input_global_scale
@@ -475,7 +454,7 @@ def _process_weights_flashinfer(
454
Steps
455
-----
456
1. Transform weights/scales via ``prepare_static_weights_for_trtllm_mxint4_moe``.
478
- 2. Return transformed tensors, passing through g_idx/bias unchanged.
457
+ 2. Return transformed tensors and biases.
458
"""
459
from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
460
prepare_static_weights_for_trtllm_mxint4_moe,
@@ -493,10 +472,6 @@ def _process_weights_flashinfer(
472
dict_weights_mxint4["gemm2_weights"],
473
dict_weights_mxint4["gemm1_scales"],
474
dict_weights_mxint4["gemm2_scales"],
496
- w13_g_idx,
497
- w2_g_idx,
498
- None,
499
- None,
475
None,
476
None,
477
None,
@@ -540,13 +515,10 @@ def _process_weights_marlin(
515
num_bits: int,
516
pack_factor: int,
517
group_size: int,
543
- actorder: str | None,
518
w13_qweight: torch.Tensor,
519
w2_qweight: torch.Tensor,
520
w13_scales: torch.Tensor,
521
w2_scales: torch.Tensor,
548
- w13_g_idx: torch.Tensor,
549
- w2_g_idx: torch.Tensor,
522
w13_qzeros: torch.Tensor | None = None,
523
w2_qzeros: torch.Tensor | None = None,
524
w13_bias: torch.Tensor | None = None,
@@ -556,10 +528,6 @@ def _process_weights_marlin(
528
torch.Tensor, # w2_qweight
529
torch.Tensor, # w13_scales
530
torch.Tensor, # w2_scales
559
- torch.Tensor, # w13_g_idx
560
- torch.Tensor, # w2_g_idx
561
- torch.Tensor, # w13_g_idx_sort_indices
562
- torch.Tensor, # w2_g_idx_sort_indices
531
torch.Tensor | None, # w13_qzeros
532
torch.Tensor | None, # w2_qzeros
533
torch.Tensor | None, # w13_input_global_scale
@@ -573,10 +541,9 @@ def _process_weights_marlin(
541
Steps
542
-----
543
1. Optional FP8 preprocessing of packed weights / scales.
576
- 2. Sort / reset g_idx tensors for act-order handling.
577
- 3. Repack weights via ``gptq_marlin_moe_repack``.
578
- 4. Permute scales (and optionally extract INT8 global scales).
579
- 5. Permute bias tensors.
544
+ 2. Repack weights via ``gptq_marlin_moe_repack``.
545
+ 3. Permute scales (and optionally extract INT8 global scales).
546
+ 4. Permute bias tensors.
547
"""
548
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
549
@@ -584,8 +551,6 @@ def _process_weights_marlin(
551
marlin_w2_qweight: torch.Tensor
552
marlin_w13_scales: torch.Tensor
553
marlin_w2_scales: torch.Tensor
587
- w13_g_idx_sort_indices: torch.Tensor | None = None
588
- w2_g_idx_sort_indices: torch.Tensor | None = None
554
w13_input_global_scale: torch.Tensor | None = None
555
w2_input_global_scale: torch.Tensor | None = None
556
w13_bias_out: torch.Tensor | None = None
@@ -606,13 +571,9 @@ def _process_weights_marlin(
571
572
# --- Pad the intermediate size to a valid Marlin thread tile ---
573
# GPTQ packs along K: w13's N is in the (shard) columns, w2's N in the rows.
609
- # Act-order keeps the strict shape and is never padded.
574
N = layer.intermediate_size_per_partition
575
padded_N = marlin_moe_padded_intermediate(N, group_size)
576
if padded_N != N:
613
- assert actorder != "group", (
614
- "Marlin MoE thread-tile padding is unsupported with act-order"
615
- )
577
marlin_w13_qweight = _pad_w13_shard_cols(marlin_w13_qweight, N, padded_N)
578
marlin_w2_qweight = _pad_rows(marlin_w2_qweight, padded_N // pack_factor)
579
marlin_w13_scales = _pad_w13_shard_cols(marlin_w13_scales, N, padded_N)
@@ -627,44 +588,9 @@ def _process_weights_marlin(
588
if w13_bias is not None:
589
w13_bias = _pad_w13_bias(w13_bias, N, padded_N)
590
630
- # --- Process act_order (g_idx) ---
631
- if actorder == "group":
632
- num_experts = w13_g_idx.shape[0]
633
- w13_g_idx_sort_indices = torch.empty_like(w13_g_idx)
634
- w2_g_idx_sort_indices = torch.empty_like(w2_g_idx)
635
- w13_sorted_g_idx = torch.empty_like(w13_g_idx)
636
- w2_sorted_g_idx = torch.empty_like(w2_g_idx)
637
- for e in range(num_experts):
638
- w13_g_idx_sort_indices[e] = torch.argsort(w13_g_idx[e]).to(torch.int32)
639
- w2_g_idx_sort_indices[e] = torch.argsort(w2_g_idx[e]).to(torch.int32)
640
- w13_sorted_g_idx[e] = w13_g_idx[e][w13_g_idx_sort_indices[e]]
641
- w2_sorted_g_idx[e] = w2_g_idx[e][w2_g_idx_sort_indices[e]]
642
- w13_g_idx = w13_sorted_g_idx
643
- w2_g_idx = w2_sorted_g_idx
644
- else:
645
- num_experts = w13_g_idx.shape[0]
646
- device = w13_g_idx.device
647
- w13_g_idx = torch.nn.Parameter(
648
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
649
- requires_grad=False,
650
- )
651
- w2_g_idx = torch.nn.Parameter(
652
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
653
- requires_grad=False,
654
- )
655
- w13_g_idx_sort_indices = torch.nn.Parameter(
656
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
657
- requires_grad=False,
658
- )
659
- w2_g_idx_sort_indices = torch.nn.Parameter(
660
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
661
- requires_grad=False,
662
- )
663
-
591
# --- Repack weights ---
592
marlin_w13_qweight = ops.gptq_marlin_moe_repack(
593
marlin_w13_qweight,
667
- w13_g_idx_sort_indices,
594
marlin_w13_qweight.shape[1] * pack_factor,
595
marlin_w13_qweight.shape[2],
596
num_bits,
@@ -672,7 +598,6 @@ def _process_weights_marlin(
598
)
599
marlin_w2_qweight = ops.gptq_marlin_moe_repack(
600
marlin_w2_qweight,
675
- w2_g_idx_sort_indices,
601
marlin_w2_qweight.shape[1] * pack_factor,
602
marlin_w2_qweight.shape[2],
603
num_bits,
@@ -734,10 +659,6 @@ def _process_weights_marlin(
659
marlin_w2_qweight,
660
marlin_w13_scales,
661
marlin_w2_scales,
737
- w13_g_idx,
738
- w2_g_idx,
739
- w13_g_idx_sort_indices,
740
- w2_g_idx_sort_indices,
662
w13_qzeros,
663
w2_qzeros,
664
w13_input_global_scale,
@@ -766,10 +687,6 @@ def _process_awq_weights_marlin(
687
torch.Tensor, # w2_qweight
688
torch.Tensor, # w13_scales
689
torch.Tensor, # w2_scales
769
- torch.Tensor | None, # w13_g_idx
770
- torch.Tensor | None, # w2_g_idx
771
- torch.Tensor | None, # w13_g_idx_sort_indices
772
- torch.Tensor | None, # w2_g_idx_sort_indices
690
torch.Tensor | None, # w13_qzeros
691
torch.Tensor | None, # w2_qzeros
692
torch.Tensor | None, # w13_input_global_scale
@@ -782,8 +699,6 @@ def _process_awq_weights_marlin(
699
AWQ checkpoints use a different packing order than GPTQ, so they need
700
AWQ-specific weight repacking and zero-point conversion before Marlin runs.
701
"""
785
- num_experts = w13_qweight.shape[0]
786
- device = w13_qweight.device
702
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
703
w13_input_global_scale: torch.Tensor | None = None
704
w2_input_global_scale: torch.Tensor | None = None
@@ -823,18 +738,8 @@ def _process_awq_weights_marlin(
738
if w13_bias is not None:
739
w13_bias = _pad_w13_bias(w13_bias, N, padded_N)
740
826
- w13_g_idx_sort_indices = torch.nn.Parameter(
827
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
828
- requires_grad=False,
829
- )
830
- w2_g_idx_sort_indices = torch.nn.Parameter(
831
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
832
- requires_grad=False,
833
- )
834
-
741
marlin_w13_qweight = ops.awq_marlin_moe_repack(
742
w13_qweight,
837
- w13_g_idx_sort_indices,
743
size_k=w13_qweight.shape[1],
744
size_n=w13_qweight.shape[2] * pack_factor,
745
num_bits=weight_bits,
@@ -842,7 +747,6 @@ def _process_awq_weights_marlin(
747
)
748
marlin_w2_qweight = ops.awq_marlin_moe_repack(
749
w2_qweight,
845
- w2_g_idx_sort_indices,
750
size_k=w2_qweight.shape[1],
751
size_n=w2_qweight.shape[2] * pack_factor,
752
num_bits=weight_bits,
@@ -898,10 +802,6 @@ def _process_awq_weights_marlin(
802
marlin_w2_qweight,
803
marlin_w13_scales,
804
marlin_w2_scales,
901
- None,
902
- None,
903
- w13_g_idx_sort_indices,
904
- w2_g_idx_sort_indices,
805
marlin_w13_qzeros,
806
marlin_w2_qzeros,
807
w13_input_global_scale,
@@ -917,8 +817,6 @@ def _process_weights_cpu(
817
w2: torch.Tensor,
818
w13_scale: torch.Tensor,
819
w2_scale: torch.Tensor,
920
- w13_g_idx: torch.Tensor | None = None,
921
- w2_g_idx: torch.Tensor | None = None,
820
w13_qzeros: torch.Tensor | None = None,
821
w2_qzeros: torch.Tensor | None = None,
822
w13_bias: torch.Tensor | None = None,
@@ -928,10 +826,6 @@ def _process_weights_cpu(
826
torch.Tensor, # w2_qweight
827
torch.Tensor, # w13_scales
828
torch.Tensor, # w2_scales
931
- torch.Tensor | None, # w13_g_idx
932
- torch.Tensor | None, # w2_g_idx
933
- torch.Tensor | None, # w13_g_idx_sort_indices
934
- torch.Tensor | None, # w2_g_idx_sort_indices
829
torch.Tensor | None, # w13_qzeros
830
torch.Tensor | None, # w2_qzeros
831
torch.Tensor | None, # w13_input_global_scale
@@ -962,7 +856,7 @@ def _process_weights_cpu(
856
if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act:
857
raise NotImplementedError(
858
"CPU WNA16 MoE backend does not support GPTQ with "
965
- "desc_act=True. The fused MoE kernel has no g_idx "
859
+ "desc_act=True. The fused MoE kernel has no activation "
860
"reordering support."
861
)
862
cpu_quant_algo = ops.CPUQuantAlgo.GPTQ
@@ -1009,10 +903,6 @@ def _process_weights_cpu(
903
blocked_w2,
904
blocked_s13,
905
blocked_s2,
1012
- w13_g_idx,
1013
- w2_g_idx,
1014
- None, # w13_g_idx_sort_indices (unused on CPU)
1015
- None, # w2_g_idx_sort_indices (unused on CPU)
906
blocked_z13,
907
blocked_z2,
908
None, # w13_input_global_scale
@@ -1352,10 +1242,6 @@ def _process_weights_emulation_gptq(
1242
w2_bf16, # w2_qweight (now bf16, not int32)
1243
dummy, # w13_scales (unused; nulled out in Int4EmulationTritonExperts)
1244
dummy, # w2_scales (unused)
1355
- None, # w13_g_idx
1356
- None, # w2_g_idx
1357
- None, # w13_g_idx_sort_indices
1358
- None, # w2_g_idx_sort_indices
1245
None, # w13_qzeros
1246
None, # w2_qzeros
1247
None, # w13_input_global_scale
@@ -1413,10 +1299,6 @@ def _process_weights_emulation_awq(
1299
None,
1300
None,
1301
None,
1416
- None,
1417
- None,
1418
- None,
1419
- None,
1302
)
1303
1304
@@ -1429,8 +1311,6 @@ def convert_to_wna16_moe_kernel_format(
1311
w2: torch.Tensor,
1312
w13_scale: torch.Tensor,
1313
w2_scale: torch.Tensor,
1432
- w13_g_idx: torch.Tensor | None = None,
1433
- w2_g_idx: torch.Tensor | None = None,
1314
w13_qzeros: torch.Tensor | None = None,
1315
w2_qzeros: torch.Tensor | None = None,
1316
w13_bias: torch.Tensor | None = None,
@@ -1441,10 +1321,6 @@ def convert_to_wna16_moe_kernel_format(
1321
torch.Tensor, # w2_qweight
1322
torch.Tensor, # w13_scales
1323
torch.Tensor, # w2_scales
1444
- torch.Tensor | None, # w13_g_idx
1445
- torch.Tensor | None, # w2_g_idx
1446
- torch.Tensor | None, # w13_g_idx_sort_indices
1447
- torch.Tensor | None, # w2_g_idx_sort_indices
1324
torch.Tensor | None, # w13_qzeros
1325
torch.Tensor | None, # w2_qzeros
1326
torch.Tensor | None, # w13_input_global_scale
@@ -1509,14 +1385,12 @@ def convert_to_wna16_moe_kernel_format(
1385
num_bits = quant_config.quant_type.size_bits
1386
pack_factor = quant_config.pack_factor
1387
group_size = quant_config.group_size
1512
- actorder = "group" if quant_config.desc_act else None
1388
elif isinstance(quant_config, QuantizationArgs):
1389
num_bits = quant_config.num_bits
1390
pack_factor = 32 // quant_config.num_bits
1391
group_size = (
1392
-1 if quant_config.group_size is None else quant_config.group_size
1393
)
1519
- actorder = quant_config.actorder
1394
else:
1395
raise TypeError(
1396
"Marlin WNA16 MoE backend requires AutoGPTQConfig, AutoAWQConfig or "
@@ -1543,22 +1417,16 @@ def convert_to_wna16_moe_kernel_format(
1417
w2_bias,
1418
)
1419
else:
1546
- if w13_g_idx is None or w2_g_idx is None:
1547
- raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")
1548
-
1420
return _process_weights_marlin(
1421
layer,
1422
input_dtype,
1423
num_bits,
1424
pack_factor,
1425
group_size,
1555
- actorder,
1426
w13,
1427
w2,
1428
w13_scale,
1429
w2_scale,
1560
- w13_g_idx,
1561
- w2_g_idx,
1430
w13_qzeros,
1431
w2_qzeros,
1432
w13_bias,
@@ -1571,8 +1439,6 @@ def convert_to_wna16_moe_kernel_format(
1439
w2,
1440
w13_scale,
1441
w2_scale,
1574
- w13_g_idx,
1575
- w2_g_idx,
1442
w13_qzeros,
1443
w2_qzeros,
1444
w13_bias,
@@ -1584,8 +1450,6 @@ def convert_to_wna16_moe_kernel_format(
1450
w2,
1451
w13_scale,
1452
w2_scale,
1587
- w13_g_idx,
1588
- w2_g_idx,
1453
w13_bias,
1454
w2_bias,
1455
)
@@ -1608,16 +1472,11 @@ def convert_to_wna16_moe_kernel_format(
1472
w13_bias,
1473
w2_bias,
1474
)
1611
- empty = torch.empty((0,), dtype=torch.int32, device=w13.device)
1475
return (
1476
w13_xpu,
1477
w2_xpu,
1478
w13_scale_xpu,
1479
w2_scale_xpu,
1617
- empty, # w13_g_idx
1618
- empty, # w2_g_idx
1619
- empty, # w13_g_idx_sort_indices
1620
- empty, # w2_g_idx_sort_indices
1480
None, # w13_qzeros — sym int4 on XPU has none; kernel does uint4b8→s4
1481
None, # w2_qzeros
1482
None, # w13_input_global_scale
@@ -1712,10 +1571,6 @@ def convert_to_wna16_moe_kernel_format(
1571
w2_uint8,
1572
w13_scale,
1573
w2_scale,
1715
- None,
1716
- None,
1717
- None,
1718
- None,
1574
w13_qzeros,
1575
w2_qzeros,
1576
None,
vllm/model_executor/layers/fused_moe/routed_experts.py
+2
-54
@@ -168,12 +168,6 @@ class RoutedExperts(PluggableLayer):
168
"global_num_experts": moe_config.num_experts,
169
}
170
171
- # need full intermediate size pre-sharding for WNA16 act order
172
- if self._needs_intermediate_size_param(self.quant_method):
173
- moe_quant_params["intermediate_size_full"] = (
174
- self.moe_config.intermediate_size
175
- )
176
-
171
self.quant_method.create_weights(layer=self, **moe_quant_params)
172
173
self.lora_base_layer_prefix = ""
@@ -206,14 +200,6 @@ class RoutedExperts(PluggableLayer):
200
assert isinstance(quant_method, FusedMoEMethodBase)
201
return quant_method
202
209
- # TODO(bnell): make this a method on quant_method
210
- def _needs_intermediate_size_param(self, quant_method: FusedMoEMethodBase) -> bool:
211
- return quant_method.__class__.__name__ in (
212
- "AutoGPTQMoEMethod",
213
- "CompressedTensorsWNA16MoEMethod",
214
- "CompressedTensorsW4A16FlydslMoEMethod",
215
- )
216
-
203
def _ensure_moe_quant_config_init(self):
204
if self.quant_method.moe_quant_config is None:
205
# Note: the moe_quant_config can't be constructed until after
@@ -349,7 +335,6 @@ class RoutedExperts(PluggableLayer):
335
shard_id: str,
336
loaded_weight: torch.Tensor,
337
tp_rank: int,
352
- load_full_w2: bool = False,
338
):
339
"""
340
Load grouped weight scales for group quantization or model weights
@@ -360,17 +345,13 @@ class RoutedExperts(PluggableLayer):
345
shard_id: either w1, w2, or w3
346
loaded_weight: checkpoint weight to load into the param
347
tp_rank: tensor parallel rank
363
- load_full_w2: whether or not the w2 loaded should be sharded.
348
"""
349
if shard_id == "w2":
366
- # In the case where we have actorder/g_idx, we do not partition the
367
- # w2 scales, as indicated by `load_full` argument, for all tp cases
350
self._load_w2(
351
shard_dim=shard_dim,
352
loaded_weight=loaded_weight,
353
expert_data=expert_data,
354
tp_rank=tp_rank,
373
- load_full=load_full_w2,
355
)
356
elif shard_id in ("w1", "w3"):
357
self._load_w13(
@@ -532,13 +513,11 @@ class RoutedExperts(PluggableLayer):
513
shard_dim: int,
514
loaded_weight: torch.Tensor,
515
tp_rank: int,
535
- load_full: bool = False,
516
):
517
# Index the loaded weight for tp sharding.
518
# down_proj: "RowParallel" so tp sharding on input_dim
539
- # Only narrow if the loaded_weight is not a scalar (0-dim tensor)
540
- # and we're not loading the full weight
541
- if not load_full and loaded_weight.ndim > 0:
519
+ # Only narrow if the loaded_weight is not a scalar (0-dim tensor).
520
+ if loaded_weight.ndim > 0:
521
# Same padding fix as _load_w13: use unpadded per-rank size.
522
tp_size = self.moe_config.moe_parallel_config.tp_size
523
loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size
@@ -571,25 +550,6 @@ class RoutedExperts(PluggableLayer):
550
# _to_scalar's reshape(()) would reject the size-2 weight_shape.
551
param_data[expert_id] = loaded_weight
552
574
- def _load_g_idx(
575
- self,
576
- shard_id: str,
577
- expert_data: torch.Tensor,
578
- shard_dim: int,
579
- loaded_weight: torch.Tensor,
580
- tp_rank: int,
581
- ):
582
- if shard_id == "w2":
583
- self._load_w2(
584
- shard_dim=shard_dim,
585
- loaded_weight=loaded_weight,
586
- expert_data=expert_data,
587
- tp_rank=tp_rank,
588
- )
589
- else:
590
- assert shard_id in ("w1", "w3")
591
- expert_data.copy_(loaded_weight)
592
-
553
@overload
554
def weight_loader(
555
self,
@@ -739,17 +699,6 @@ class RoutedExperts(PluggableLayer):
699
)
700
return True if return_success else None
701
742
- # Case g_idx
743
- if "g_idx" in weight_name:
744
- self._load_g_idx(
745
- shard_dim=0,
746
- shard_id=shard_id,
747
- loaded_weight=loaded_weight,
748
- expert_data=expert_data,
749
- tp_rank=self.moe_config.tp_rank,
750
- )
751
- return True if return_success else None
752
-
702
# TODO @dsikka: ModelOpt should follow the proper MoE loading pattern
703
if "ModelOpt" in quant_method_name:
704
# Determine per-tensor weight scale patterns based on variant
@@ -853,7 +802,6 @@ class RoutedExperts(PluggableLayer):
802
loaded_weight=loaded_weight,
803
expert_data=expert_data,
804
tp_rank=self.moe_config.tp_rank,
856
- load_full_w2=getattr(param, "load_full_w2", False),
805
)
806
elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value:
807
self._load_per_tensor_weight_scale(
vllm/model_executor/layers/quantization/auto_awq.py
-23
@@ -439,7 +439,6 @@ class AutoAWQMarlinLinearMethod(LinearMethodBase):
439
act_type=params_dtype if self.input_dtype is None else self.input_dtype,
440
group_size=self.quant_config.group_size,
441
zero_points=self.quant_config.zero_point,
442
- has_g_idx=False,
442
)
443
444
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
@@ -558,11 +557,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
557
}
558
)
559
561
- intermediate_size_full = extra_weight_attrs.pop(
562
- "intermediate_size_full", intermediate_size_per_partition
563
- )
564
- self.is_k_full = intermediate_size_per_partition == intermediate_size_full
565
-
560
w13_qweight = Parameter(
561
torch.empty(
562
num_experts,
@@ -672,10 +666,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
666
w2,
667
w13_scale,
668
w2_scale,
675
- w13_g_idx,
676
- w2_g_idx,
677
- w13_g_idx_sort_indices,
678
- w2_g_idx_sort_indices,
669
w13_qzeros,
670
w2_qzeros,
671
w13_input_global_scale,
@@ -696,14 +686,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
686
687
replace_parameter(layer, "w13_scales", w13_scale)
688
replace_parameter(layer, "w2_scales", w2_scale)
699
- _replace_or_register_parameter(
700
- layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices
701
- )
702
- _replace_or_register_parameter(
703
- layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices
704
- )
705
- _replace_or_register_parameter(layer, "w13_g_idx", w13_g_idx)
706
- _replace_or_register_parameter(layer, "w2_g_idx", w2_g_idx)
689
_replace_or_register_parameter(layer, "w13_qzeros", w13_qzeros)
690
_replace_or_register_parameter(layer, "w2_qzeros", w2_qzeros)
691
_replace_or_register_parameter(
@@ -726,11 +708,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
708
moe_config=self.moe,
709
experts_cls=self.experts_cls,
710
backend=self.wna16_moe_backend,
729
- is_k_full=self.is_k_full,
730
- w13_g_idx=getattr(layer, "w13_g_idx", None),
731
- w2_g_idx=getattr(layer, "w2_g_idx", None),
732
- w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None),
733
- w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None),
711
routing_tables=layer._expert_routing_tables(),
712
)
713
vllm/model_executor/layers/quantization/auto_gptq.py
+7
-92
@@ -40,6 +40,7 @@ from vllm.model_executor.layers.quantization.utils import replace_parameter
40
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
41
get_dynamic_override,
42
get_linear_quant_method,
43
+ normalize_and_validate_gptq_desc_act,
44
override_config,
45
)
46
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
@@ -59,7 +60,6 @@ from vllm.model_executor.parameter import (
60
GroupQuantScaleParameter,
61
PackedColumnParameter,
62
PackedvLLMParameter,
62
- RowvLLMParameter,
63
)
64
from vllm.scalar_type import scalar_types
65
from vllm.transformers_utils.config import get_safetensors_params_metadata
@@ -115,10 +115,7 @@ class AutoGPTQConfig(QuantizationConfig):
115
modules_in_block_to_quantize: list[str] | None = None,
116
) -> None:
117
super().__init__()
118
- if desc_act and group_size == -1:
119
- # In this case, act_order == True is the same as act_order == False
120
- # (since we have only one group per output channel)
121
- desc_act = False
118
+ desc_act = normalize_and_validate_gptq_desc_act(desc_act, group_size, dynamic)
119
120
# GPTQModel use `dynamic` config property to allow per module
121
# quantization config so each module can be individually optimized.
@@ -150,7 +147,7 @@ class AutoGPTQConfig(QuantizationConfig):
147
148
self.pack_factor = 32 // weight_bits # packed into int32
149
self.group_size = group_size
153
- self.desc_act = desc_act
150
+ self.desc_act: bool = desc_act
151
self.lm_head_quantized = lm_head_quantized
152
self.full_config = full_config
153
@@ -244,7 +241,7 @@ class AutoGPTQConfig(QuantizationConfig):
241
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
242
243
if not check_moe_marlin_supports_layer(
247
- layer, self.group_size, allow_tile_padding=not self.desc_act
244
+ layer, self.group_size, allow_tile_padding=True
245
):
246
logger.warning_once(
247
f"Layer '{prefix}' is not supported by GPTQMoeMarlin. "
@@ -348,7 +345,6 @@ class AutoGPTQLinearMethod(LinearMethodBase):
345
act_type=params_dtype if input_dtype is None else input_dtype,
346
group_size=self.quant_config.group_size,
347
zero_points=False,
351
- has_g_idx=self.quant_config.desc_act,
348
)
349
350
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
@@ -365,7 +361,7 @@ class AutoGPTQLinearMethod(LinearMethodBase):
361
362
# Determine sharding
363
if marlin_repeat_scales_on_all_ranks(
368
- self.quant_config.desc_act, self.quant_config.group_size, is_row_parallel
364
+ self.quant_config.group_size, is_row_parallel
365
):
366
# By setting scale_dim == None, weight_loader will
367
# repeat the scales on each GPU in TP>1 case.
@@ -391,16 +387,6 @@ class AutoGPTQLinearMethod(LinearMethodBase):
387
weight_loader=weight_loader,
388
)
389
394
- # Activation order
395
- g_idx = RowvLLMParameter(
396
- data=torch.empty(
397
- input_size_per_partition,
398
- dtype=torch.int32,
399
- ),
400
- input_dim=0,
401
- weight_loader=weight_loader,
402
- )
403
-
390
qzeros_args = {
391
"data": torch.empty(
392
scales_and_zp_size,
@@ -440,7 +426,6 @@ class AutoGPTQLinearMethod(LinearMethodBase):
426
)
427
428
layer.register_parameter("qweight", qweight)
443
- layer.register_parameter("g_idx", g_idx)
429
layer.register_parameter("scales", scales)
430
layer.register_parameter("qzeros", qzeros)
431
@@ -449,7 +434,6 @@ class AutoGPTQLinearMethod(LinearMethodBase):
434
w_q_param_name="qweight",
435
w_s_param_name="scales",
436
w_zp_param_name="qzeros",
452
- w_gidx_param_name="g_idx",
437
)
438
439
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -492,7 +476,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
476
quant_config=self.quant_config,
477
may_have_zp=not self.quant_config.is_sym,
478
may_have_bias=True,
495
- allow_tile_padding=not self.quant_config.desc_act,
479
+ allow_tile_padding=True,
480
)
481
482
def create_weights(
@@ -512,19 +496,9 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
496
"W8A8-INT8 is not supported by marlin kernel."
497
)
498
515
- intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
516
-
517
- self.is_k_full = (not self.quant_config.desc_act) or (
518
- intermediate_size_per_partition == intermediate_size_full
519
- )
520
-
499
if self.quant_config.group_size != -1:
500
scales_size13 = hidden_size // self.quant_config.group_size
523
- w2_scales_size = (
524
- intermediate_size_full
525
- if self.quant_config.desc_act
526
- else intermediate_size_per_partition
527
- )
501
+ w2_scales_size = intermediate_size_per_partition
502
scales_size2 = w2_scales_size // self.quant_config.group_size
503
strategy = FusedMoeWeightScaleSupported.GROUP.value
504
else:
@@ -579,8 +553,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
553
)
554
layer.register_parameter("w2_scales", w2_scales)
555
set_weight_attrs(w2_scales, extra_weight_attrs)
582
- # don't shard the w2 scales when running act order
583
- set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act})
556
# up_proj zero points
557
w13_qzeros = torch.nn.Parameter(
558
torch.empty(
@@ -607,48 +579,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
579
)
580
layer.register_parameter("w2_qzeros", w2_qzeros)
581
set_weight_attrs(w2_qzeros, extra_weight_attrs)
610
- # don't shard the w2 scales when running act order
611
- set_weight_attrs(w2_qzeros, {"load_full_w2": self.quant_config.desc_act})
612
- w13_g_idx = torch.nn.Parameter(
613
- torch.empty(
614
- num_experts,
615
- hidden_size,
616
- dtype=torch.int32,
617
- ),
618
- requires_grad=False,
619
- )
620
- layer.register_parameter("w13_g_idx", w13_g_idx)
621
- set_weight_attrs(w13_g_idx, extra_weight_attrs)
622
- w2_g_idx = torch.nn.Parameter(
623
- torch.empty(
624
- num_experts,
625
- intermediate_size_per_partition,
626
- dtype=torch.int32,
627
- ),
628
- requires_grad=False,
629
- )
630
- layer.register_parameter("w2_g_idx", w2_g_idx)
631
- set_weight_attrs(w2_g_idx, extra_weight_attrs)
632
- w13_g_idx_sort_indices = torch.nn.Parameter(
633
- torch.empty(
634
- num_experts,
635
- hidden_size,
636
- dtype=torch.int32,
637
- ),
638
- requires_grad=False,
639
- )
640
- layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
641
- set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
642
- w2_g_idx_sort_indices = torch.nn.Parameter(
643
- torch.empty(
644
- num_experts,
645
- intermediate_size_per_partition,
646
- dtype=torch.int32,
647
- ),
648
- requires_grad=False,
649
- )
650
- layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
651
- set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
582
583
# Some GPTQ checkpoints contain expert biases even when the model
584
# architecture does not declare them. Zero initialization keeps
@@ -712,8 +642,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
642
w2=layer.w2_qweight,
643
w13_scale=layer.w13_scales,
644
w2_scale=layer.w2_scales,
715
- w13_g_idx=layer.w13_g_idx,
716
- w2_g_idx=layer.w2_g_idx,
645
w13_bias=w13_bias,
646
w2_bias=w2_bias,
647
w13_qzeros=getattr(layer, "w13_qzeros", None),
@@ -730,10 +658,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
658
w2,
659
w13_scale,
660
w2_scale,
733
- w13_g_idx,
734
- w2_g_idx,
735
- w13_g_idx_sort_indices,
736
- w2_g_idx_sort_indices,
661
w13_qzeros,
662
w2_qzeros,
663
w13_input_global_scale,
@@ -746,10 +670,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
670
replace_parameter(layer, "w2_qweight", w2)
671
replace_parameter(layer, "w13_scales", w13_scale)
672
replace_parameter(layer, "w2_scales", w2_scale)
749
- replace_parameter(layer, "w13_g_idx", w13_g_idx)
750
- replace_parameter(layer, "w2_g_idx", w2_g_idx)
751
- replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
752
- replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
673
replace_or_register("w13_input_global_scale", w13_input_global_scale)
674
replace_or_register("w2_input_global_scale", w2_input_global_scale)
675
replace_or_register("w13_bias", w13_bias)
@@ -772,11 +692,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
692
moe_config=self.moe,
693
experts_cls=self.experts_cls,
694
backend=self.wna16_moe_backend,
775
- is_k_full=self.is_k_full,
776
- w13_g_idx=getattr(layer, "w13_g_idx", None),
777
- w2_g_idx=getattr(layer, "w2_g_idx", None),
778
- w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None),
779
- w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None),
695
routing_tables=layer._expert_routing_tables(),
696
)
697
vllm/model_executor/layers/quantization/base_config.py
+7
@@ -238,6 +238,13 @@ class QuantizationConfig(ABC):
238
}
239
return WeightsMapper(orig_to_new_regex=orig_to_new_regex)
240
241
+ @staticmethod
242
+ def get_checkpoint_weight_mapper() -> "WeightsMapper":
243
+ """Discard activation-order metadata unused by supported kernels."""
244
+ from vllm.model_executor.models.utils import WeightsMapper
245
+
246
+ return WeightsMapper(orig_to_new_suffix={".g_idx": None})
247
+
248
def apply_vllm_mapper( # noqa: B027
249
self, hf_to_vllm_mapper: "WeightsMapper"
250
):
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+13
-3
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast
7
import torch
8
from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
9
from compressed_tensors.quantization import (
10
+ ActivationOrdering,
11
QuantizationArgs,
12
QuantizationStrategy,
13
QuantizationType,
@@ -325,9 +326,20 @@ class CompressedTensorsConfig(QuantizationConfig):
326
targets = quant_config.get("targets")
327
for target in targets:
328
target_scheme_map[target] = {}
328
- target_scheme_map[target]["weights"] = QuantizationArgs.model_validate(
329
+ weight_quant = QuantizationArgs.model_validate(
330
quant_config.get("weights")
331
)
332
+ if weight_quant.actorder in (
333
+ ActivationOrdering.GROUP.value,
334
+ ActivationOrdering.DYNAMIC.value,
335
+ ):
336
+ raise ValueError(
337
+ "Compressed-tensors group/dynamic activation ordering "
338
+ f"(actorder={weight_quant.actorder!r}) is no longer supported. "
339
+ "Use a checkpoint with static/weight activation ordering or "
340
+ "without activation ordering."
341
+ )
342
+ target_scheme_map[target]["weights"] = weight_quant
343
344
target_scheme_map[target]["input_activations"] = None
345
target_scheme_map[target]["format"] = quant_config.get("format")
@@ -754,7 +766,6 @@ class CompressedTensorsConfig(QuantizationConfig):
766
strategy=weight_quant.strategy,
767
symmetric=weight_quant.symmetric,
768
group_size=weight_quant.group_size,
757
- actorder=weight_quant.actorder,
769
)
770
771
# Must come before the WNA16 check; standard 4/8-bit weight-only (no
@@ -803,7 +814,6 @@ class CompressedTensorsConfig(QuantizationConfig):
814
strategy=weight_quant.strategy,
815
symmetric=weight_quant.symmetric,
816
group_size=weight_quant.group_size,
806
- actorder=weight_quant.actorder,
817
layer_name=layer_name,
818
)
819
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py
-47
@@ -104,8 +104,6 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod):
104
assert weight_quant.strategy == "group"
105
assert weight_quant.group_size == 32
106
self.group_size = weight_quant.group_size
107
- # grouped actorder isn't supported by this kernel
108
- assert weight_quant.actorder != "group"
107
assert weight_quant.symmetric, (
108
"Only symmetric quantization is supported for MoE. "
109
"Try --moe-backend emulation."
@@ -180,7 +178,6 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod):
178
)
179
layer.register_parameter("w2_weight_scale", w2_scale)
180
set_weight_attrs(w2_scale, extra_weight_attrs)
183
- set_weight_attrs(w2_scale, {"load_full_w2": False})
181
182
w2_weight_shape = torch.nn.Parameter(
183
torch.empty(num_experts, 2), requires_grad=False
@@ -194,50 +191,6 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod):
191
layer.register_parameter("w13_weight_shape", w13_weight_shape)
192
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
193
197
- w13_g_idx = torch.nn.Parameter(
198
- torch.empty(
199
- num_experts,
200
- hidden_size,
201
- dtype=torch.int32,
202
- ),
203
- requires_grad=False,
204
- )
205
- layer.register_parameter("w13_weight_g_idx", w13_g_idx)
206
- set_weight_attrs(w13_g_idx, extra_weight_attrs)
207
-
208
- w2_g_idx = torch.nn.Parameter(
209
- torch.empty(
210
- num_experts,
211
- intermediate_size_per_partition,
212
- dtype=torch.int32,
213
- ),
214
- requires_grad=False,
215
- )
216
- layer.register_parameter("w2_weight_g_idx", w2_g_idx)
217
- set_weight_attrs(w2_g_idx, extra_weight_attrs)
218
-
219
- w13_g_idx_sort_indices = torch.nn.Parameter(
220
- torch.empty(
221
- num_experts,
222
- hidden_size,
223
- dtype=torch.int32,
224
- ),
225
- requires_grad=False,
226
- )
227
- layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
228
- set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
229
-
230
- w2_g_idx_sort_indices = torch.nn.Parameter(
231
- torch.empty(
232
- num_experts,
233
- intermediate_size_per_partition,
234
- dtype=torch.int32,
235
- ),
236
- requires_grad=False,
237
- )
238
- layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
239
- set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
240
-
194
layer.a13_scale = None
195
layer.a2_scale = None
196
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
-1
@@ -47,7 +47,6 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
47
assert self.weight_quant.symmetric, (
48
"Only symmetric quantization is supported for W4A8 MoE"
49
)
50
- assert self.weight_quant.actorder != "group"
50
assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE"
51
52
self.w4a8_backend, self.experts_cls = select_w4a8_moe_backend(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+5
-126
@@ -3,13 +3,10 @@
3
4
import math
5
from fractions import Fraction
6
-from typing import Any
6
7
import torch
8
from compressed_tensors.quantization import (
10
- ActivationOrdering,
9
QuantizationArgs,
12
- QuantizationStrategy,
10
)
11
12
from vllm.logger import init_logger
@@ -75,7 +72,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
72
self.group_size = (
73
-1 if weight_quant.group_size is None else weight_quant.group_size
74
)
78
- self.actorder = weight_quant.actorder
75
76
# Extract quant_type and create weight key for oracle selection
77
self.quant_type = (
@@ -104,11 +100,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
100
101
weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)
102
107
- is_actorder = self.strategy == QuantizationStrategy.GROUP and self.actorder in (
108
- ActivationOrdering.GROUP,
109
- ActivationOrdering.DYNAMIC,
110
- )
111
-
103
# Select WNA16 MoE backend via oracle.
104
self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
105
config=self.moe,
@@ -116,7 +107,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
107
quant_config=self.weight_quant,
108
may_have_zp=not self.symmetric,
109
may_have_bias=False,
119
- allow_tile_padding=not is_actorder,
110
+ allow_tile_padding=True,
111
)
112
113
self.is_marlin = self.wna16_backend in [
@@ -130,14 +121,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
121
122
if self.is_marlin:
123
assert check_moe_marlin_supports_config(
133
- self.moe, self.group_size, allow_tile_padding=not is_actorder
124
+ self.moe, self.group_size, allow_tile_padding=True
125
)
126
self.input_dtype = get_marlin_input_dtype(layer_name)
127
else:
128
# channelwise is not supported by this kernel
129
assert weight_quant.strategy == "group"
139
- # grouped actorder isn't supported by this kernel
140
- assert weight_quant.actorder != "group"
130
131
# Non-Marlin WNA16 always uses bf16/fp16 inputs
132
self.input_dtype = torch.bfloat16
@@ -230,29 +219,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
219
assert shape[1] is not None and shape[2] is not None
220
return shape[0], shape[1], shape[2]
221
233
- @staticmethod
234
- def _w2_scale_sharding(
235
- actorder,
236
- group_size: int,
237
- intermediate_size_per_partition: int,
238
- intermediate_size_full: int,
239
- ) -> tuple[bool, int, bool]:
240
- """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE.
241
-
242
- Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime
243
- and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``.
244
- ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at
245
- quantization time, so scales shard normally per TP rank.
246
- """
247
- load_full_w2 = (actorder == "group") and group_size != -1
248
- w2_scales_size = (
249
- intermediate_size_full if load_full_w2 else intermediate_size_per_partition
250
- )
251
- is_k_full = (actorder != "group") or (
252
- intermediate_size_per_partition == intermediate_size_full
253
- )
254
- return load_full_w2, w2_scales_size, is_k_full
255
-
222
def create_weights(
223
self,
224
layer: torch.nn.Module,
@@ -262,8 +228,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
228
params_dtype: torch.dtype,
229
**extra_weight_attrs,
230
):
265
- intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
266
-
231
# Will transpose the loaded weight along the
232
# intermediate and hidden dim sizes. Will
233
# shard for TP along the transposed dims
@@ -301,13 +265,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
265
layer.register_parameter("w2_weight_packed", w2_weight)
266
set_weight_attrs(w2_weight, extra_weight_attrs)
267
304
- load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding(
305
- self.actorder,
306
- self.group_size,
307
- intermediate_size_per_partition,
308
- intermediate_size_full,
309
- )
310
-
268
if self.strategy == "channel":
269
num_groups_w2 = num_groups_w13 = 1
270
self.group_size = -1
@@ -318,10 +275,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
275
f"({hidden_size}) to be divisible by group_size "
276
f"({self.group_size})."
277
)
321
- if (
322
- not load_full_w2
323
- and intermediate_size_per_partition % self.group_size != 0
324
- ):
278
+ if intermediate_size_per_partition % self.group_size != 0:
279
raise ValueError(
280
"CompressedTensors WNA16 MoE with static group "
281
"scales requires the MoE intermediate size per "
@@ -331,7 +285,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
285
"otherwise cross TP shard boundaries; use a compatible TP "
286
"size or enable expert parallelism."
287
)
334
- num_groups_w2 = w2_scales_size // self.group_size
288
+ num_groups_w2 = intermediate_size_per_partition // self.group_size
289
num_groups_w13 = hidden_size // self.group_size
290
291
if not self.symmetric:
@@ -389,7 +343,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
343
)
344
layer.register_parameter("w2_weight_scale", w2_scale)
345
set_weight_attrs(w2_scale, extra_weight_attrs)
392
- set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})
346
347
if not self.symmetric:
348
w13_zp = torch.nn.Parameter(
@@ -436,50 +389,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
389
layer.register_parameter("w13_weight_shape", w13_weight_shape)
390
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
391
439
- w13_g_idx = torch.nn.Parameter(
440
- torch.empty(
441
- num_experts,
442
- hidden_size,
443
- dtype=torch.int32,
444
- ),
445
- requires_grad=False,
446
- )
447
- layer.register_parameter("w13_weight_g_idx", w13_g_idx)
448
- set_weight_attrs(w13_g_idx, extra_weight_attrs)
449
-
450
- w2_g_idx = torch.nn.Parameter(
451
- torch.empty(
452
- num_experts,
453
- intermediate_size_per_partition,
454
- dtype=torch.int32,
455
- ),
456
- requires_grad=False,
457
- )
458
- layer.register_parameter("w2_weight_g_idx", w2_g_idx)
459
- set_weight_attrs(w2_g_idx, extra_weight_attrs)
460
-
461
- w13_g_idx_sort_indices = torch.nn.Parameter(
462
- torch.empty(
463
- num_experts,
464
- hidden_size,
465
- dtype=torch.int32,
466
- ),
467
- requires_grad=False,
468
- )
469
- layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
470
- set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
471
-
472
- w2_g_idx_sort_indices = torch.nn.Parameter(
473
- torch.empty(
474
- num_experts,
475
- intermediate_size_per_partition,
476
- dtype=torch.int32,
477
- ),
478
- requires_grad=False,
479
- )
480
- layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
481
- set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
482
-
392
layer.a13_scale = None
393
layer.a2_scale = None
394
@@ -488,24 +397,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
397
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
398
assert self.moe_quant_config is not None
399
491
- # Add Marlin-specific arguments
492
- marlin_args: dict[str, Any] = {}
493
- if self.is_marlin:
494
- marlin_args = {
495
- "w13_g_idx": layer.w13_weight_g_idx,
496
- "w2_g_idx": layer.w2_weight_g_idx,
497
- "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices,
498
- "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices,
499
- "is_k_full": self.is_k_full,
500
- }
501
-
400
self.moe_kernel = make_wna16_moe_kernel(
401
moe_quant_config=self.moe_quant_config,
402
moe_config=self.moe,
403
experts_cls=self.experts_cls,
404
backend=self.wna16_backend,
405
routing_tables=layer._expert_routing_tables(),
508
- **marlin_args,
406
)
407
408
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -519,8 +416,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
416
w2=layer.w2_weight_packed,
417
w13_scale=layer.w13_weight_scale,
418
w2_scale=layer.w2_weight_scale,
522
- w13_g_idx=layer.w13_weight_g_idx,
523
- w2_g_idx=layer.w2_weight_g_idx,
419
w13_qzeros=getattr(layer, "w13_weight_zero_point", None),
420
w2_qzeros=getattr(layer, "w2_weight_zero_point", None),
421
)
@@ -534,10 +429,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
429
w2_qweight,
430
w13_scales,
431
w2_scales,
537
- w13_g_idx_processed,
538
- w2_g_idx_processed,
539
- w13_g_idx_sort_indices,
540
- w2_g_idx_sort_indices,
432
w13_qzeros,
433
w2_qzeros,
434
w13_input_global_scale,
@@ -561,19 +452,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
452
replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
453
replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
454
564
- # Marlin-specific parameters (not needed for Flashinfer)
455
if self.is_marlin:
566
- if w13_g_idx_processed is not None:
567
- replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
568
- if w2_g_idx_processed is not None:
569
- replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
570
- if w13_g_idx_sort_indices is not None:
571
- replace_parameter(
572
- layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices
573
- )
574
- if w2_g_idx_sort_indices is not None:
575
- replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
576
-
456
# Register input global scales if present
457
if w13_input_global_scale is not None:
458
layer.register_parameter(
@@ -593,7 +472,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
472
and self.wna16_backend != WNA16MoEBackend.EMULATION
473
):
474
layer.workspace = marlin_make_workspace_new(
596
- layer.w13_weight_g_idx.device,
475
+ layer.w13_weight_packed.device,
476
4,
477
existing=getattr(layer, "workspace", None),
478
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py
+2
-3
@@ -65,15 +65,14 @@ class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod):
65
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
66
device = layer.w13_weight_packed.device
67
num_experts = layer.w13_weight_packed.shape[0]
68
- empty_g_idx = torch.empty(0, dtype=torch.int32, device=device)
68
69
# Shuffle weights in-place per expert (exllama nibble interleave)
70
for e in range(num_experts):
71
w13_e = layer.w13_weight_packed.data[e].contiguous()
73
- ops.gptq_shuffle(w13_e, empty_g_idx, 4)
72
+ ops.gptq_shuffle(w13_e, 4)
73
layer.w13_weight_packed.data[e] = w13_e
74
w2_e = layer.w2_weight_packed.data[e].contiguous()
76
- ops.gptq_shuffle(w2_e, empty_g_idx, 4)
75
+ ops.gptq_shuffle(w2_e, 4)
76
layer.w2_weight_packed.data[e] = w2_e
77
78
# Keep scales as [E, groups, N] in activation dtype
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py
+1
-6
@@ -4,7 +4,6 @@
4
from collections.abc import Callable
5
6
import torch
7
-from compressed_tensors.quantization import ActivationOrdering
7
8
from vllm.distributed.utils import verify_group_size_divides_partition
9
from vllm.logger import init_logger
@@ -43,13 +42,11 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
42
num_bits: int,
43
group_size: int | None = None,
44
symmetric: bool | None = True,
46
- actorder: ActivationOrdering | None = None,
45
):
46
self.pack_factor = 32 // num_bits
47
self.strategy = strategy
48
self.symmetric = symmetric
49
self.group_size = -1 if group_size is None else group_size
52
- self.has_g_idx = actorder == ActivationOrdering.GROUP
50
51
if self.group_size != 128 or self.strategy != "group":
52
raise ValueError(
@@ -92,7 +89,6 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
89
act_type=torch.float8_e4m3fn, # always use fp8(e4m3)
90
group_size=self.group_size,
91
zero_points=not self.symmetric,
95
- has_g_idx=self.has_g_idx,
92
out_type=params_dtype,
93
)
94
@@ -106,7 +102,7 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
102
group_size = self.group_size if self.group_size != -1 else input_size
103
row_parallel = input_size != input_size_per_partition
104
partition_scales = not marlin_repeat_scales_on_all_ranks(
109
- self.has_g_idx, self.group_size, row_parallel
105
+ self.group_size, row_parallel
106
)
107
108
scales_and_zp_size = input_size // group_size
@@ -162,7 +158,6 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
158
w_q_param_name="weight_packed",
159
w_s_param_name="weight_scale",
160
w_zp_param_name="weight_zero_point",
165
- w_gidx_param_name="weight_g_idx",
161
)
162
163
# Checkpoints are serialized in compressed-tensors format, which is
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py
-2
@@ -98,7 +98,6 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme):
98
act_type=params_dtype,
99
group_size=effective_group_size,
100
zero_points=False,
101
- has_g_idx=False,
101
)
102
103
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
@@ -140,7 +139,6 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme):
139
w_q_param_name="weight_packed",
140
w_s_param_name="weight_scale",
141
w_zp_param_name=None,
143
- w_gidx_param_name=None,
142
)
143
144
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
+1
-19
@@ -6,7 +6,6 @@ from collections.abc import Callable
6
from fractions import Fraction
7
8
import torch
9
-from compressed_tensors.quantization import ActivationOrdering
9
10
from vllm.distributed.utils import verify_group_size_divides_partition
11
from vllm.logger import init_logger
@@ -28,7 +27,6 @@ from vllm.model_executor.parameter import (
27
GroupQuantScaleParameter,
28
PackedColumnParameter,
29
PackedvLLMParameter,
31
- RowvLLMParameter,
30
)
31
from vllm.scalar_type import scalar_types
32
@@ -57,7 +55,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
55
num_bits: int,
56
group_size: int | None = None,
57
symmetric: bool | None = True,
60
- actorder: ActivationOrdering | None = None,
58
layer_name: str | None = None,
59
):
60
self.num_bits = num_bits
@@ -65,7 +62,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
62
self.strategy = strategy
63
self.symmetric = symmetric
64
self.group_size = -1 if group_size is None else group_size
68
- self.has_g_idx = actorder == ActivationOrdering.GROUP
65
self.layer_name = layer_name
66
67
if self.group_size == -1 and self.strategy != "channel":
@@ -128,7 +124,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
124
act_type=params_dtype,
125
group_size=self.group_size,
126
zero_points=not self.symmetric,
131
- has_g_idx=self.has_g_idx,
127
)
128
129
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
@@ -146,7 +141,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
141
group_size = self.group_size if self.group_size != -1 else input_size
142
row_parallel = input_size != input_size_per_partition
143
partition_scales = not marlin_repeat_scales_on_all_ranks(
149
- self.has_g_idx, self.group_size, row_parallel
144
+ self.group_size, row_parallel
145
)
146
147
scales_and_zp_size = input_size // group_size
@@ -226,24 +221,11 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
221
if not self.symmetric:
222
layer.register_parameter("weight_zero_point", qzeros)
223
229
- # group index (for activation reordering)
230
- if self.has_g_idx:
231
- weight_g_idx = RowvLLMParameter(
232
- data=torch.empty(
233
- input_size_per_partition,
234
- dtype=torch.int32,
235
- ),
236
- input_dim=0,
237
- weight_loader=weight_loader,
238
- )
239
- layer.register_parameter("weight_g_idx", weight_g_idx)
240
-
224
self.kernel = kernel_type(
225
mp_linear_kernel_config,
226
w_q_param_name="weight_packed",
227
w_s_param_name="weight_scale",
228
w_zp_param_name="weight_zero_point",
246
- w_gidx_param_name="weight_g_idx",
229
)
230
231
# Checkpoints are serialized in compressed-tensors format, which is
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py
+1
-2
@@ -144,7 +144,6 @@ class CompressedTensorsWNA4Int(CompressedTensorsScheme):
144
act_type=params_dtype,
145
group_size=self.group_size,
146
zero_points=False,
147
- has_g_idx=False,
147
)
148
149
kernel_type = choose_mp_linear_kernel(mp_config)
@@ -166,7 +165,7 @@ class CompressedTensorsWNA4Int(CompressedTensorsScheme):
165
group_size = self.group_size if self.group_size != -1 else input_size
166
row_parallel = input_size != input_size_per_partition
167
partition_scales = not marlin_repeat_scales_on_all_ranks(
169
- False, self.group_size, row_parallel
168
+ self.group_size, row_parallel
169
)
170
scales_size = input_size // group_size
171
if partition_scales:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py
+1
-2
@@ -144,7 +144,6 @@ class CompressedTensorsWNA8Int(CompressedTensorsScheme):
144
act_type=params_dtype,
145
group_size=self.group_size,
146
zero_points=False,
147
- has_g_idx=False,
147
)
148
149
kernel_type = choose_mp_linear_kernel(mp_config)
@@ -166,7 +165,7 @@ class CompressedTensorsWNA8Int(CompressedTensorsScheme):
165
group_size = self.group_size if self.group_size != -1 else input_size
166
row_parallel = input_size != input_size_per_partition
167
partition_scales = not marlin_repeat_scales_on_all_ranks(
169
- False, self.group_size, row_parallel
168
+ self.group_size, row_parallel
169
)
170
scales_size = input_size // group_size
171
if partition_scales:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py
+1
-2
@@ -114,7 +114,6 @@ class CompressedTensorsWNA8O8Int(CompressedTensorsScheme):
114
act_type=params_dtype, # activation quant applied externally (SRQ)
115
group_size=self.group_size,
116
zero_points=False,
117
- has_g_idx=False,
117
)
118
self.kernel = choose_mp_linear_kernel(mp_config)(
119
mp_config,
@@ -167,7 +166,7 @@ class CompressedTensorsWNA8O8Int(CompressedTensorsScheme):
166
# Scale: per-output-channel, or per group along the input dim under TP.
167
group_size = self.group_size if self.group_size != -1 else input_size
168
partitioned = not marlin_repeat_scales_on_all_ranks(
170
- False, self.group_size, input_size != input_size_per_partition
169
+ self.group_size, input_size != input_size_per_partition
170
)
171
scales = (input_size_per_partition if partitioned else input_size) // group_size
172
scale_data = torch.empty(out, scales, dtype=params_dtype)
vllm/model_executor/layers/quantization/inc/schemes/inc_w4a8_linear.py
+1
-1
@@ -106,7 +106,7 @@ class INCXPUW4A8LinearMethod(INCXPULinearMethod):
106
layer.scales_fp16,
107
layer.qzeros,
108
self.group_size,
109
- None,
109
+ None, # Retained by the external XPU op ABI.
110
bias,
111
)
112
return out.to(x.dtype).reshape(out_shape)
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py
+1
-12
@@ -14,7 +14,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
14
from vllm.model_executor.parameter import (
15
GroupQuantScaleParameter,
16
PackedvLLMParameter,
17
- RowvLLMParameter,
17
)
18
from vllm.scalar_type import scalar_types
19
@@ -254,16 +253,6 @@ class INCXPULinearBase(INCLinearScheme):
253
layer.register_parameter("scales", scales)
254
layer.register_parameter("qzeros", qzeros)
255
257
- g_idx = RowvLLMParameter(
258
- data=torch.tensor(
259
- [i // self.group_size for i in range(input_size_per_partition)],
260
- dtype=torch.int32,
261
- ),
262
- input_dim=0,
263
- weight_loader=weight_loader,
264
- )
265
- layer.register_parameter("g_idx", g_idx)
266
-
256
def _convert_awq_qweight_to_gptq(self, qw: torch.Tensor) -> torch.Tensor:
257
"""Convert AWQ qweight [K, N // pf] to GPTQ qweight [K // pf, N].
258
@@ -356,7 +345,7 @@ class INCXPULinearMethod(INCXPULinearBase):
345
layer.scales,
346
layer.qzeros,
347
self.group_size,
359
- None,
348
+ None, # Retained by the external XPU op ABI.
349
)
350
return out.reshape(out_shape)
351
vllm/model_executor/layers/quantization/moe_wna16.py
+22
-12
@@ -32,6 +32,9 @@ from vllm.model_executor.layers.quantization.base_config import (
32
QuantizationConfig,
33
QuantizeMethodBase,
34
)
35
+from vllm.model_executor.layers.quantization.utils.gptq_utils import (
36
+ normalize_and_validate_gptq_desc_act,
37
+)
38
from vllm.model_executor.layers.quantization.utils.quant_utils import (
39
INT4_DTYPE,
40
INT8_DTYPE,
@@ -69,7 +72,11 @@ class MoeWNA16Config(QuantizationConfig):
72
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
73
74
if self.linear_quant_method == "gptq":
72
- pass
75
+ normalize_and_validate_gptq_desc_act(
76
+ full_config.get("desc_act", False),
77
+ group_size,
78
+ full_config.get("dynamic") or {},
79
+ )
80
elif self.linear_quant_method in ("awq", "awq_marlin"):
81
capability_tuple = current_platform.get_device_capability()
82
device_capability = (
@@ -148,7 +155,16 @@ class MoeWNA16Config(QuantizationConfig):
155
# Extract data from quant config.
156
quant_method = quant_config.get("quant_method", "").lower()
157
num_bits = quant_config.get("bits")
151
- desc_act = quant_config.get("desc_act")
158
+ desc_act = quant_config.get("desc_act", False)
159
+ try:
160
+ normalize_and_validate_gptq_desc_act(
161
+ desc_act,
162
+ quant_config.get("group_size"),
163
+ quant_config.get("dynamic") or {},
164
+ )
165
+ supported_act_order = True
166
+ except ValueError:
167
+ supported_act_order = False
168
169
capability_tuple = current_platform.get_device_capability()
170
device_capability = (
@@ -159,7 +175,9 @@ class MoeWNA16Config(QuantizationConfig):
175
176
awq_min_capability = AutoAWQConfig.get_min_capability()
177
162
- gptq_compatible = quant_method == "gptq" and not desc_act and num_bits in [4, 8]
178
+ gptq_compatible = (
179
+ quant_method == "gptq" and supported_act_order and num_bits in [4, 8]
180
+ )
181
awq_compatible = (
182
quant_method == "awq"
183
and num_bits == 4
@@ -356,7 +374,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
374
if self.quant_config.linear_quant_method == "gptq":
375
# some param are unused, but we need to init them in order to
376
# load weights
359
- invalid_param_keys = ["w13_g_idx", "w2_g_idx"]
377
+ invalid_param_keys: list[str] = []
378
if not self.quant_config.has_zp:
379
invalid_param_keys += ["w13_qzeros", "w2_qzeros"]
380
for key in invalid_param_keys:
@@ -414,8 +432,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
432
w2=layer.w2_qweight,
433
w13_scale=layer.w13_scales,
434
w2_scale=layer.w2_scales,
417
- w13_g_idx=None,
418
- w2_g_idx=None,
435
w13_qzeros=layer.w13_qzeros if has_zp else None,
436
w2_qzeros=layer.w2_qzeros if has_zp else None,
437
)
@@ -430,10 +446,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
446
w2_qweight,
447
w13_scales,
448
w2_scales,
433
- _,
434
- _,
435
- _,
436
- _,
449
w13_qzeros,
450
w2_qzeros,
451
w13_input_global_scale,
@@ -576,8 +588,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
588
expert_id: int,
589
return_success: bool = False,
590
):
579
- if "g_idx" in weight_name:
580
- return False if return_success else None
591
if not layer.quant_config.has_zp and "qzeros" in weight_name:
592
return False if return_success else None
593
vllm/model_executor/layers/quantization/utils/gptq_utils.py
+34
@@ -20,6 +20,36 @@ else:
20
AutoGPTQConfig = object
21
22
23
+def normalize_and_validate_gptq_desc_act(
24
+ desc_act: bool,
25
+ group_size: int | None,
26
+ dynamic: Mapping[str, Any],
27
+) -> bool:
28
+ base_desc_act = desc_act
29
+ if desc_act and group_size == -1:
30
+ # With one group per output channel, activation ordering is a no-op.
31
+ desc_act = False
32
+ if desc_act:
33
+ raise ValueError(
34
+ "GPTQ group activation ordering (desc_act=True) is no longer "
35
+ "supported. Use a checkpoint with static activation ordering "
36
+ "or desc_act=False."
37
+ )
38
+
39
+ for pattern, overrides in dynamic.items():
40
+ if pattern.startswith("-:") or not isinstance(overrides, Mapping):
41
+ continue
42
+ dynamic_desc_act = overrides.get("desc_act", base_desc_act)
43
+ dynamic_group_size = overrides.get("group_size", group_size)
44
+ if dynamic_desc_act and dynamic_group_size != -1:
45
+ raise ValueError(
46
+ "GPTQ group activation ordering (desc_act=True) is no longer "
47
+ f"supported, but is enabled by dynamic rule {pattern!r}. Use "
48
+ "a checkpoint with static activation ordering or desc_act=False."
49
+ )
50
+ return desc_act
51
+
52
+
53
# Match dynamic rules with module name (prefix) and override quantize
54
# config if module (prefix) matches a rule
55
def override_config(config: AutoGPTQConfig, prefix: str):
@@ -33,6 +63,10 @@ def override_config(config: AutoGPTQConfig, prefix: str):
63
if isinstance(desc_act, bool):
64
config.desc_act = desc_act
65
66
+ if config.desc_act and config.group_size == -1:
67
+ # Activation ordering is a no-op for channelwise quantization.
68
+ config.desc_act = False
69
+
70
config.pack_factor = 32 // config.weight_bits # packed into int32
71
assert isinstance(config, AutoGPTQConfig)
72
is_sym = get_dynamic_override(config, prefix, "sym", config.is_sym)
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+8
-27
@@ -359,11 +359,11 @@ def check_moe_marlin_supports_config(
359
) -> bool:
360
"""Whether the fused MoE Marlin kernel supports ``config``.
361
362
- Callers without act-order may pass ``allow_tile_padding=True``: a
363
- tile-misaligned intermediate size is then zero-padded to a valid thread
364
- tile at weight prep (see marlin_moe_padded_intermediate), so only a group
365
- straddling the padded boundary stays unsupported. hidden_size is the MoE
366
- I/O extent and is never padded. Act-order keeps the strict shape.
362
+ With ``allow_tile_padding=True``, a tile-misaligned intermediate size is
363
+ zero-padded to a valid thread tile at weight prep (see
364
+ marlin_moe_padded_intermediate), so only a group straddling the padded
365
+ boundary stays unsupported. hidden_size is the MoE I/O extent and is never
366
+ padded.
367
"""
368
if current_platform.is_rocm():
369
return False
@@ -433,30 +433,17 @@ def marlin_make_workspace_new(
433
return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False)
434
435
436
-def marlin_is_k_full(act_order: bool, is_row_parallel: bool) -> bool:
437
- return (not act_order) or (act_order and not is_row_parallel)
438
-
439
-
440
-def marlin_repeat_scales_on_all_ranks(
441
- act_order: bool, group_size: int, is_row_parallel: bool
442
-) -> bool:
443
- # Need to repeat scales on every rank if act_ordering or
444
- # channelwise and RowParallelLinear
436
+def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool:
437
is_channelwise = group_size == -1
446
- return act_order or (is_channelwise and is_row_parallel)
438
+ return is_channelwise and is_row_parallel
439
440
449
-def marlin_make_empty_g_idx(device: torch.device) -> torch.Tensor:
441
+def marlin_make_empty(device: torch.device) -> torch.Tensor:
442
return torch.nn.Parameter(
443
torch.empty(0, dtype=torch.int, device=device), requires_grad=False
444
)
445
446
455
-def marlin_sort_g_idx(g_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
456
- g_idx_sort_indices = torch.argsort(g_idx).to(torch.int)
457
- return g_idx[g_idx_sort_indices], g_idx_sort_indices
458
-
459
-
447
def get_scale_perms():
448
scale_perm: list[int] = []
449
for i in range(8):
@@ -700,13 +687,10 @@ def apply_gptq_marlin_linear(
687
weight: torch.Tensor,
688
weight_scale: torch.Tensor,
689
weight_zp: torch.Tensor,
703
- g_idx: torch.Tensor,
704
- g_idx_sort_indices: torch.Tensor,
690
workspace: torch.Tensor,
691
wtype: ScalarType,
692
output_size_per_partition: int,
693
input_size_per_partition: int,
709
- is_k_full: bool,
694
input_global_scale: torch.Tensor | None = None,
695
bias: torch.Tensor | None = None,
696
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
@@ -749,14 +733,11 @@ def apply_gptq_marlin_linear(
733
a_scales,
734
None,
735
weight_zp,
752
- g_idx,
753
- g_idx_sort_indices,
736
workspace,
737
wtype,
738
size_m=reshaped_x.shape[0],
739
size_n=padded_n,
740
size_k=padded_k,
759
- is_k_full=is_k_full,
741
use_atomic_add=use_atomic_add,
742
use_fp32_reduce=use_fp32_reduce,
743
is_zp_float=False,
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+3
-20
@@ -203,8 +203,6 @@ def apply_fp4_marlin_linear(
203
a_scales=a_scales,
204
global_scale=weight_global_scale,
205
b_zeros=None,
206
- g_idx=None,
207
- perm=None,
206
workspace=workspace,
207
b_q_type=scalar_types.float4_e2m1f,
208
size_m=reshaped_x.size(0),
@@ -246,14 +244,12 @@ def prepare_fp4_layer_for_marlin(
244
245
# WEIGHT
246
# Repack weights to marlin format
249
- perm = torch.empty(0, dtype=torch.int, device=device)
247
qweight = layer.weight.view(torch.int32).T.contiguous()
248
qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k)
249
250
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
251
marlin_qweight = ops.gptq_marlin_repack(
252
b_q_weight=qweight,
256
- perm=perm,
253
size_k=padded_k,
254
size_n=padded_n,
255
num_bits=4,
@@ -312,7 +308,6 @@ def _repack_marlin_experts(
308
weight: torch.Tensor,
309
size_n: int,
310
size_k: int,
315
- perm: torch.Tensor,
311
is_a_8bit: bool,
312
) -> torch.Tensor:
313
"""Repack each expert to marlin format into a preallocated output."""
@@ -322,7 +317,6 @@ def _repack_marlin_experts(
317
qweight = weight[i].view(torch.int32).T.contiguous()
318
marlin_qweight = ops.gptq_marlin_repack(
319
b_q_weight=qweight,
325
- perm=perm,
320
size_k=size_k,
321
size_n=size_n,
322
num_bits=4,
@@ -401,7 +395,6 @@ def prepare_nvfp4_moe_layer_for_marlin(
395
layer.workspace = marlin_make_workspace_new(
396
device, 4, existing=getattr(layer, "workspace", None)
397
)
404
- perm = torch.empty(0, dtype=torch.int, device=device)
398
399
# WEIGHT
400
# Repack weights to marlin format
@@ -417,7 +410,7 @@ def prepare_nvfp4_moe_layer_for_marlin(
410
weight = pad_w2(weight, packing=2)
411
size_k = padded_N
412
420
- return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit)
413
+ return _repack_marlin_experts(weight, size_n, size_k, is_a_8bit)
414
415
w13 = repack_weight(w13, "w13")
416
w2 = repack_weight(w2, "w2")
@@ -492,7 +485,6 @@ def prepare_moe_fp4_layer_for_marlin(
485
layer.workspace = marlin_make_workspace_new(
486
device, 4, existing=getattr(layer, "workspace", None)
487
)
495
- perm = torch.empty(0, dtype=torch.int, device=device)
488
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
489
490
# WEIGHT
@@ -506,7 +498,7 @@ def prepare_moe_fp4_layer_for_marlin(
498
499
assert weight.shape == (e, size_n, size_k // 2)
500
509
- weight = _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit)
501
+ weight = _repack_marlin_experts(weight, size_n, size_k, is_a_8bit)
502
weight = torch.nn.Parameter(weight, requires_grad=False)
503
504
setattr(layer, name, weight)
@@ -621,10 +613,8 @@ def prepare_moe_mxfp4_layer_for_marlin(
613
n = w13.shape[1] // 2 # intermediate_size_per_partition
614
k = w13.shape[2] * 2 # hidden_size
615
624
- device = w13.device
616
param_dtype = layer.params_dtype
617
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
627
- perm = torch.empty(0, dtype=torch.int, device=device)
618
619
# WEIGHT: Repack weights to marlin format
620
def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
@@ -635,7 +625,7 @@ def prepare_moe_mxfp4_layer_for_marlin(
625
626
assert weight.shape == (e, size_n, size_k // 2)
627
638
- return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit)
628
+ return _repack_marlin_experts(weight, size_n, size_k, is_a_8bit)
629
630
w13 = repack_weight(w13, "w13")
631
w2 = repack_weight(w2, "w2")
@@ -691,8 +681,6 @@ def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
681
assert not is_a_8bit, "NVFP4 weight + INT8/FP8 activation is not supported."
682
assert group_size > 0
683
size_n, size_k = weight.shape
694
- device = weight.device
695
-
684
scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 6
685
global_scale = scales.max() / 448
686
scales = (scales / global_scale).to(torch.float8_e4m3fn)
@@ -720,7 +708,6 @@ def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
708
709
marlin_qweight = ops.gptq_marlin_repack(
710
b_q_weight=fp4_weight.view(torch.int32).T.contiguous(),
723
- perm=torch.empty(0, dtype=torch.int, device=device),
711
size_k=size_k,
712
size_n=size_n,
713
num_bits=4,
@@ -751,8 +738,6 @@ def rand_marlin_weight_mxfp4_like(weight, group_size, input_dtype=None):
738
739
assert group_size > 0
740
size_n, size_k = weight.shape
754
- device = weight.device
755
-
741
scales = torch.randint(
742
110,
743
120,
@@ -779,11 +764,9 @@ def rand_marlin_weight_mxfp4_like(weight, group_size, input_dtype=None):
764
).view(size_n, size_k)
765
weight_ref = weight_ref * scales.repeat_interleave(group_size, 1).to(weight.dtype)
766
782
- perm = torch.empty(0, dtype=torch.int, device=device)
767
fp4_weight = fp4_weight.view(torch.int32).T.contiguous()
768
marlin_qweight = ops.gptq_marlin_repack(
769
b_q_weight=fp4_weight,
786
- perm=perm,
770
size_k=size_k,
771
size_n=size_n,
772
num_bits=4,
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+1
-16
@@ -89,8 +89,6 @@ def apply_fp8_marlin_linear(
89
a_scales=a_scales,
90
global_scale=None,
91
b_zeros=None,
92
- g_idx=None,
93
- perm=None,
92
workspace=workspace,
93
b_q_type=scalar_types.float8_e4m3fn,
94
size_m=reshaped_x.size(0),
@@ -138,7 +136,6 @@ def prepare_fp8_layer_for_marlin(
136
137
# WEIGHT
138
# Repack weights to marlin format
141
- perm = torch.empty(0, dtype=torch.int, device=device)
139
qweight = pack_fp8_to_int32(layer.weight, size_k_first)
140
if not size_k_first:
141
qweight = qweight.T.contiguous()
@@ -146,7 +143,6 @@ def prepare_fp8_layer_for_marlin(
143
144
marlin_qweight = ops.gptq_marlin_repack(
145
b_q_weight=qweight,
149
- perm=perm,
146
size_k=padded_k,
147
size_n=padded_n,
148
num_bits=8,
@@ -283,7 +279,6 @@ def prepare_fp8_moe_layer_for_marlin(
279
layer.workspace = marlin_make_workspace_new(
280
device, 4, existing=getattr(layer, "workspace", None)
281
)
286
- perm = torch.empty(0, dtype=torch.int, device=device)
282
283
# WEIGHT
284
# Repack weights to marlin format
@@ -295,7 +290,7 @@ def prepare_fp8_moe_layer_for_marlin(
290
qweight = qweight.T.contiguous()
291
292
marlin_qweight = ops.gptq_marlin_repack(
298
- b_q_weight=qweight, perm=perm, size_k=size_k, size_n=size_n, num_bits=8
293
+ b_q_weight=qweight, size_k=size_k, size_n=size_n, num_bits=8
294
)
295
tensor_list.append(marlin_qweight)
296
@@ -434,8 +429,6 @@ def apply_mxfp8_marlin_linear(
429
a_scales=None,
430
global_scale=None,
431
b_zeros=None,
437
- g_idx=None,
438
- perm=None,
432
workspace=workspace,
433
b_q_type=scalar_types.float8_e4m3fn,
434
size_m=reshaped_x.size(0),
@@ -470,14 +463,12 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
463
)
464
465
# WEIGHT - repack FP8 weights to Marlin format
473
- perm = torch.empty(0, dtype=torch.int, device=device)
466
qweight = pack_fp8_to_int32(layer.weight, size_k_first=False)
467
qweight = qweight.T.contiguous()
468
qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k)
469
470
marlin_qweight = ops.gptq_marlin_repack(
471
b_q_weight=qweight,
480
- perm=perm,
472
size_k=padded_k,
473
size_n=padded_n,
474
num_bits=8,
@@ -551,7 +542,6 @@ def prepare_mxfp8_moe_layer_for_marlin(
542
543
device = w13.device
544
param_dtype = torch.get_default_dtype()
554
- perm = torch.empty(0, dtype=torch.int, device=device)
545
546
layer.workspace = marlin_make_workspace_new(
547
device, 4, existing=getattr(layer, "workspace", None)
@@ -571,7 +561,6 @@ def prepare_mxfp8_moe_layer_for_marlin(
561
qweight = qweight.T.contiguous()
562
marlin_qweight = ops.gptq_marlin_repack(
563
b_q_weight=qweight,
574
- perm=perm,
564
size_k=size_k,
565
size_n=size_n,
566
num_bits=8,
@@ -615,8 +604,6 @@ def marlin_quant_fp8_torch(weight, group_size, input_dtype=None):
604
assert input_dtype == torch.float8_e4m3fn
605
606
size_n, size_k = weight.shape
618
- device = weight.device
619
-
607
if group_size != -1:
608
scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 448
609
repeated_scales = scales.repeat_interleave(group_size, 1)
@@ -629,10 +616,8 @@ def marlin_quant_fp8_torch(weight, group_size, input_dtype=None):
616
weight_ref = fp8_weight.to(weight.dtype) * repeated_scales
617
618
packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous()
632
- perm = torch.empty(0, dtype=torch.int, device=device)
619
marlin_qweight = ops.gptq_marlin_repack(
620
b_q_weight=packed_weight,
635
- perm=perm,
621
size_k=size_k,
622
size_n=size_n,
623
num_bits=8,
vllm/model_executor/layers/quantization/utils/marlin_utils_test.py
+2
-14
@@ -13,7 +13,6 @@ from .quant_utils import (
13
get_pack_factor,
14
gptq_quantize_weights,
15
quantize_weights,
16
- sort_weights,
16
)
17
18
@@ -129,8 +128,6 @@ def marlin_quantize(
128
w: torch.Tensor,
129
quant_type: ScalarType,
130
group_size: int,
132
- act_order: bool,
133
- test_perm: torch.Tensor | None = None,
131
input_dtype: torch.dtype | None = None,
132
):
133
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
@@ -143,16 +140,7 @@ def marlin_quantize(
140
group_size = size_k
141
assert group_size <= size_k
142
146
- # Quantize (and apply act_order if provided)
147
- w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights(
148
- w, quant_type, group_size, act_order, test_perm
149
- )
150
-
151
- # For act_order, sort the "weights" and "g_idx" so that group ids are
152
- # increasing
153
- sort_indices = torch.empty(0, dtype=torch.int, device=w.device)
154
- if act_order:
155
- q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
143
+ w_ref, q_w, s = gptq_quantize_weights(w, quant_type, group_size)
144
145
# Reformat to marlin
146
weight_perm = get_weight_perm(num_bits, is_a_8bit)
@@ -166,7 +154,7 @@ def marlin_quantize(
154
marlin_s = marlin_s * 512
155
156
# Create result
169
- res_list = [w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm]
157
+ res_list = [w_ref, marlin_q_w, marlin_s]
158
for i in range(len(res_list)):
159
res_list[i] = res_list[i].to(w.device)
160
vllm/model_executor/layers/quantization/utils/quant_utils.py
+1
-61
@@ -691,36 +691,6 @@ def get_pack_factor(num_bits):
691
return 32 // num_bits
692
693
694
-def permute_rows(
695
- q_w: torch.Tensor,
696
- w_ref: torch.Tensor,
697
- group_size: int,
698
- test_perm: torch.Tensor | None = None,
699
-):
700
- assert q_w.shape == w_ref.shape
701
-
702
- orig_device = q_w.device
703
- k_size, _ = q_w.shape
704
-
705
- g_idx = torch.zeros((k_size,), dtype=torch.int32)
706
- for i in range(k_size):
707
- g_idx[i] = i // group_size
708
-
709
- # Simulate act_order by doing a random permutation on K
710
- rand_perm = test_perm if test_perm is not None else torch.randperm(k_size)
711
-
712
- g_idx = g_idx[rand_perm].contiguous()
713
- q_w = q_w[rand_perm, :].contiguous()
714
- w_ref = w_ref[rand_perm, :].contiguous()
715
-
716
- return (
717
- w_ref.to(device=orig_device),
718
- q_w.to(device=orig_device),
719
- g_idx.to(device=orig_device),
720
- rand_perm.to(device=orig_device),
721
- )
722
-
723
-
694
def quantize_weights(
695
w: torch.Tensor,
696
quant_type: ScalarType,
@@ -824,8 +794,6 @@ def gptq_quantize_weights(
794
w: torch.Tensor,
795
quant_type: ScalarType,
796
group_size: int,
827
- act_order: bool,
828
- test_perm: torch.Tensor | None = None,
797
):
798
size_k, _ = w.shape
799
@@ -838,35 +806,7 @@ def gptq_quantize_weights(
806
)
807
808
w_ref, w_q, w_s, _ = quantize_weights(w, quant_type, group_size)
841
-
842
- # Apply act_order
843
- g_idx = torch.empty(0, dtype=torch.int, device=w.device)
844
- rand_perm = torch.empty(0, dtype=torch.int, device=w.device)
845
- if act_order:
846
- assert group_size < size_k, (
847
- "For act_order, groupsize = {} must be less than size_k = {}".format(
848
- group_size, size_k
849
- )
850
- )
851
-
852
- w_ref, w_q, g_idx, rand_perm = permute_rows(w_q, w_ref, group_size, test_perm)
853
-
854
- return w_ref, w_q, w_s, g_idx, rand_perm
855
-
856
-
857
-def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor):
858
- orig_device = q_w.device
859
-
860
- sort_indices = torch.argsort(g_idx).to(dtype=torch.int32) # Sort based on g_idx
861
-
862
- g_idx = g_idx[sort_indices].contiguous()
863
- q_w = q_w[sort_indices, :].contiguous()
864
-
865
- return (
866
- q_w.to(device=orig_device),
867
- g_idx.to(device=orig_device),
868
- sort_indices.to(device=orig_device),
869
- )
809
+ return w_ref, w_q, w_s
810
811
812
def pack_rows(
vllm/model_executor/layers/vocab_parallel_embedding.py
+1
-2
@@ -463,8 +463,7 @@ class VocabParallelEmbedding(PluggableLayer):
463
output_dim = getattr(param, "output_dim", None)
464
packed_dim = getattr(param, "packed_dim", None)
465
466
- # If parameter does not have output dim, then it should
467
- # be copied onto all gpus (e.g. g_idx for act_order gptq).
466
+ # Parameters without an output dimension are copied onto all GPUs.
467
if output_dim is None:
468
if (
469
loaded_weight.ndim == 0
vllm/model_executor/model_loader/weight_utils.py
-2
@@ -1543,8 +1543,6 @@ def maybe_remap_moe_expert_param_name(
1543
"w2_bias",
1544
"w13_scale",
1545
"w2_scale",
1546
- "w13_g_idx",
1547
- "w2_g_idx",
1546
"w13_qweight",
1547
"w2_qweight",
1548
"w13_qzeros",
vllm/model_executor/models/utils.py
+2
-1
@@ -432,8 +432,9 @@ class AutoWeightsLoader:
432
modules = (self.module, *self.module.children())
433
iterator = (m.quant_config for m in modules if hasattr(m, "quant_config"))
434
if quant_config := next(iterator, None):
435
- # Get mappings and ignore prefixes for KV cache quantization scales
435
+ # Apply mappings for quantization-specific checkpoint tensors.
436
mapper |= quant_config.get_cache_scale_mapper()
437
+ mapper |= quant_config.get_checkpoint_weight_mapper()
438
ignore_unexpected_suffixes = quant_config._ignore_unexpected_suffixes
439
self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes)
440
mapper |= self.REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER