iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Quant][Kernel] Remove GPTQ Group/Dynamic Activation Ordering (#54809)

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