A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 05:45 UTC
[ROCm][Refactor] Enable AWQMarlinConfig on ROCm to use choose_mp_linear_kernel (#36505)
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Matthias Gehre committed
Mar 23, 2026 at 08:36 UTC
410d30089310c3e6bc385230f2e1bb2ccdc72edc
3 files changed
+143
-102
vllm/model_executor/kernels/linear/mixed_precision/conch.py
+2
@@ -113,6 +113,8 @@ class ConchLinearKernel(MPLinearKernel):
113
self._transform_param(layer, self.w_s_name, transform_w_s)
114
if self.config.zero_points:
115
self._transform_param(layer, self.w_zp_name, transform_w_zp)
116
+ elif self.w_zp_name is not None:
117
+ layer.register_parameter(self.w_zp_name, None)
118
119
def apply_weights(
120
self,
vllm/model_executor/layers/quantization/awq_marlin.py
+133
-93
@@ -10,6 +10,10 @@ from torch.nn import Parameter
10
import vllm.model_executor.layers.fused_moe # noqa
11
from vllm import _custom_ops as ops
12
from vllm.logger import init_logger
13
+from vllm.model_executor.kernels.linear import (
14
+ MPLinearLayerConfig,
15
+ choose_mp_linear_kernel,
16
+)
17
from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEConfig,
19
FusedMoEQuantConfig,
@@ -34,21 +38,16 @@ from vllm.model_executor.layers.quantization.base_config import (
38
)
39
from vllm.model_executor.layers.quantization.utils import replace_parameter
40
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
37
- apply_awq_marlin_linear,
38
- awq_to_marlin_zero_points,
41
check_marlin_supported,
42
check_marlin_supports_layer,
43
check_moe_marlin_supports_layer,
44
get_marlin_input_dtype,
45
marlin_act_int8_process_scales,
44
- marlin_make_empty_g_idx,
46
marlin_make_workspace_new,
47
marlin_moe_permute_scales,
48
marlin_permute_bias,
48
- marlin_permute_scales,
49
moe_awq_to_marlin_zero_points,
50
verify_marlin_supported,
51
- verify_marlin_supports_shape,
51
)
52
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
53
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
@@ -63,6 +62,90 @@ if TYPE_CHECKING:
62
63
logger = init_logger(__name__)
64
65
+# AWQ uses a non-standard packing order within int32 values.
66
+# For 4-bit: standard order stores values at bit positions [0,4,8,12,16,20,24,28]
67
+# for indices [0,1,2,3,4,5,6,7], while AWQ stores them for indices
68
+# [0,4,1,5,2,6,3,7]. This permutation reverses that ordering.
69
+_REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
70
+
71
+
72
+def _convert_awq_to_standard_format(
73
+ layer: torch.nn.Module,
74
+ w_q_name: str,
75
+ w_zp_name: str,
76
+ size_bits: int,
77
+) -> None:
78
+ """Convert AWQ weight and zero-point tensors to standard GPTQ-like format.
79
+
80
+ AWQ packs qweight along the output dim with a non-standard bit order.
81
+ This converts to standard bit order and repacks qweight along the input
82
+ dim, matching the format expected by the MPLinearKernel framework.
83
+ """
84
+ pack_factor = 32 // size_bits
85
+ mask = (1 << size_bits) - 1
86
+ device = getattr(layer, w_q_name).device
87
+ reverse_order = torch.tensor(
88
+ _REVERSE_AWQ_PACK_ORDER, dtype=torch.long, device=device
89
+ )
90
+ shifts = torch.arange(0, 32, size_bits, dtype=torch.int32, device=device)
91
+
92
+ # --- Convert qweight: (K, N // pack) packed_dim=1 → (K // pack, N) packed_dim=0
93
+ qw = getattr(layer, w_q_name).data
94
+ K, N_packed = qw.shape
95
+ N = N_packed * pack_factor
96
+
97
+ # Unpack int32 → individual values, fix AWQ ordering
98
+ unpacked = (qw.unsqueeze(-1) >> shifts) & mask # (K, N_packed, pack_factor)
99
+ unpacked = unpacked[:, :, reverse_order]
100
+ unpacked = unpacked.reshape(K, N) # (K, N)
101
+
102
+ # Repack along input dim (dim 0)
103
+ unpacked = unpacked.reshape(K // pack_factor, pack_factor, N)
104
+ new_qw = (unpacked.to(torch.int32) << shifts[None, :, None]).sum(
105
+ dim=1, dtype=torch.int32
106
+ )
107
+
108
+ def _noop_loader(*args, **kwargs):
109
+ pass
110
+
111
+ new_param = PackedvLLMParameter(
112
+ data=new_qw.contiguous(),
113
+ input_dim=0,
114
+ output_dim=1,
115
+ packed_dim=0,
116
+ packed_factor=pack_factor,
117
+ weight_loader=_noop_loader,
118
+ )
119
+ setattr(layer, w_q_name, new_param)
120
+
121
+ # --- Convert qzeros: fix AWQ bit ordering and repack
122
+ # AWQ qzeros: (G, N // pack) packed along dim 1, AWQ bit order
123
+ # Target: (N // pack, G) packed along dim 0, standard bit order
124
+ # This matches the CompressedTensors layout expected by the kernels.
125
+ qz = getattr(layer, w_zp_name).data
126
+ G, _ = qz.shape
127
+
128
+ unpacked_zp = (qz.unsqueeze(-1) >> shifts) & mask # (G, N_packed, pack_factor)
129
+ unpacked_zp = unpacked_zp[:, :, reverse_order]
130
+ unpacked_zp = unpacked_zp.reshape(G, N) # (G, N) individual values
131
+
132
+ # Transpose and repack along dim 0 (output dim)
133
+ unpacked_zp = unpacked_zp.T # (N, G)
134
+ unpacked_zp = unpacked_zp.reshape(N // pack_factor, pack_factor, G)
135
+ new_qz = (unpacked_zp.to(torch.int32) << shifts[None, :, None]).sum(
136
+ dim=1, dtype=torch.int32
137
+ )
138
+
139
+ new_zp_param = PackedvLLMParameter(
140
+ data=new_qz.contiguous(),
141
+ output_dim=0,
142
+ input_dim=1,
143
+ packed_dim=0,
144
+ packed_factor=pack_factor,
145
+ weight_loader=_noop_loader,
146
+ )
147
+ setattr(layer, w_zp_name, new_zp_param)
148
+
149
150
class AWQMarlinConfig(QuantizationConfig):
151
"""Config class for AWQ Marlin"""
@@ -226,7 +309,7 @@ class AWQMarlinConfig(QuantizationConfig):
309
group_size = quant_config.get("group_size")
310
zero_point = quant_config.get("zero_point")
311
229
- if not current_platform.is_cuda():
312
+ if not (current_platform.is_cuda_alike() or current_platform.is_cpu()):
313
return False
314
315
if quant_method != "awq":
@@ -268,15 +351,26 @@ class AWQMarlinConfig(QuantizationConfig):
351
class AWQMarlinLinearMethod(LinearMethodBase):
352
"""Linear method for AWQ Marlin.
353
354
+ Uses choose_mp_linear_kernel to select the best available kernel
355
+ (Conch, Exllama, or Marlin) for the current platform.
356
+
357
Args:
358
quant_config: The AWQ Marlin quantization config.
359
"""
360
361
+ _kernel_backends_being_used: set[str] = set()
362
+
363
def __init__(self, quant_config: AWQMarlinConfig) -> None:
364
self.quant_config = quant_config
365
self.quant_type = scalar_types.uint4
366
self.input_dtype = None
367
368
+ verify_marlin_supported(
369
+ quant_type=self.quant_config.quant_type,
370
+ group_size=self.quant_config.group_size,
371
+ has_zp=self.quant_config.zero_point,
372
+ )
373
+
374
def create_weights(
375
self,
376
layer: torch.nn.Module,
@@ -287,23 +381,35 @@ class AWQMarlinLinearMethod(LinearMethodBase):
381
params_dtype: torch.dtype,
382
**extra_weight_attrs,
383
) -> None:
290
- del output_size
384
output_size_per_partition = sum(output_partition_sizes)
385
weight_loader = extra_weight_attrs.get("weight_loader")
386
294
- # Normalize group_size
387
if self.quant_config.group_size != -1:
388
group_size = self.quant_config.group_size
389
else:
390
group_size = input_size
391
300
- verify_marlin_supports_shape(
301
- output_size_per_partition=output_size_per_partition,
302
- input_size_per_partition=input_size_per_partition,
303
- input_size=input_size,
304
- group_size=group_size,
392
+ mp_linear_kernel_config = MPLinearLayerConfig(
393
+ full_weight_shape=(input_size, output_size),
394
+ partition_weight_shape=(
395
+ input_size_per_partition,
396
+ output_size_per_partition,
397
+ ),
398
+ weight_type=self.quant_config.quant_type,
399
+ act_type=params_dtype if self.input_dtype is None else self.input_dtype,
400
+ group_size=self.quant_config.group_size,
401
+ zero_points=self.quant_config.zero_point,
402
+ has_g_idx=False,
403
)
404
405
+ kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
406
+
407
+ if kernel_type.__name__ not in self._kernel_backends_being_used:
408
+ logger.info("Using %s for AWQMarlinLinearMethod", kernel_type.__name__)
409
+ self._kernel_backends_being_used.add(kernel_type.__name__)
410
+
411
+ # Weights are loaded in AWQ checkpoint format (packed along output dim).
412
+ # Conversion to GPTQ-like format happens in process_weights_after_loading.
413
qweight = PackedvLLMParameter(
414
data=torch.empty(
415
input_size_per_partition,
@@ -318,7 +424,6 @@ class AWQMarlinLinearMethod(LinearMethodBase):
424
)
425
426
num_groups = input_size_per_partition // group_size
321
- layer.num_groups = num_groups
427
428
qzeros = PackedvLLMParameter(
429
data=torch.empty(
@@ -348,73 +453,22 @@ class AWQMarlinLinearMethod(LinearMethodBase):
453
layer.register_parameter("qzeros", qzeros)
454
layer.register_parameter("scales", scales)
455
351
- layer.input_size_per_partition = input_size_per_partition
352
- layer.output_size_per_partition = output_size_per_partition
353
- layer.num_groups = num_groups
354
-
355
- # TODO: Update this docs
356
- # Checkpoints are serialized in AutoAWQ format, which is different from the
357
- # marlin format. This function is called after the weights are loaded.
358
- # Here, we handle the repacking
359
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
360
- device = layer.qweight.device
361
- layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
362
- layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
363
- layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
364
-
365
- # Allocate marlin workspace
366
- layer.workspace = marlin_make_workspace_new(device)
367
-
368
- is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
369
-
370
- if self.input_dtype == torch.float8_e4m3fn:
371
- ops.marlin_int4_fp8_preprocess(layer.qweight, layer.qzeros, inplace=True)
372
- layer.scales.data = layer.scales.data * 512
373
-
374
- # Repack weights from AWQ format to marlin format.
375
- marlin_qweight = ops.awq_marlin_repack(
376
- layer.qweight,
377
- size_k=layer.input_size_per_partition,
378
- size_n=layer.output_size_per_partition,
379
- num_bits=self.quant_config.quant_type.size_bits,
380
- is_a_8bit=is_a_8bit,
456
+ self.kernel = kernel_type(
457
+ mp_linear_kernel_config,
458
+ w_q_param_name="qweight",
459
+ w_s_param_name="scales",
460
+ w_zp_param_name="qzeros",
461
)
382
- replace_parameter(layer, "qweight", marlin_qweight)
462
384
- # Permute scales from AWQ format to marlin format.
385
- marlin_scales = marlin_permute_scales(
386
- layer.scales,
387
- size_k=layer.input_size_per_partition,
388
- size_n=layer.output_size_per_partition,
389
- group_size=self.quant_config.group_size,
390
- is_a_8bit=is_a_8bit,
391
- )
392
- if self.input_dtype == torch.int8 and layer.num_groups > 1:
393
- marlin_scales, input_global_scale = marlin_act_int8_process_scales(
394
- marlin_scales
395
- )
396
- layer.register_parameter(
397
- "input_global_scale", Parameter(input_global_scale, requires_grad=False)
398
- )
399
-
400
- replace_parameter(layer, "scales", marlin_scales)
401
-
402
- # Permute zero-points from AWQ format to marlin format.
403
- marlin_zp = awq_to_marlin_zero_points(
404
- layer.qzeros,
405
- size_k=layer.num_groups,
406
- size_n=layer.output_size_per_partition,
407
- num_bits=self.quant_config.quant_type.size_bits,
408
- is_a_8bit=is_a_8bit,
463
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
464
+ # AWQ checkpoints use a non-standard packing order and pack qweight
465
+ # along the output dimension. Convert to the standard format
466
+ # (GPTQ-like: standard bit order, qweight packed along input dim)
467
+ # before handing off to the kernel.
468
+ _convert_awq_to_standard_format(
469
+ layer, "qweight", "qzeros", self.quant_config.quant_type.size_bits
470
)
410
- replace_parameter(layer, "qzeros", marlin_zp)
411
-
412
- # Not-used
413
- layer.g_idx = marlin_make_empty_g_idx(device)
414
- layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
415
-
416
- if hasattr(layer, "bias") and layer.bias is not None:
417
- layer.bias.data = marlin_permute_bias(layer.bias)
471
+ self.kernel.process_weights_after_loading(layer)
472
473
def apply(
474
self,
@@ -422,21 +476,7 @@ class AWQMarlinLinearMethod(LinearMethodBase):
476
x: torch.Tensor,
477
bias: torch.Tensor | None = None,
478
) -> torch.Tensor:
425
- return apply_awq_marlin_linear(
426
- input=x,
427
- weight=layer.qweight,
428
- weight_scale=layer.scales,
429
- weight_zp=layer.qzeros,
430
- g_idx=layer.g_idx,
431
- g_idx_sort_indices=layer.g_idx_sort_indices,
432
- workspace=layer.workspace,
433
- quant_type=self.quant_config.quant_type,
434
- output_size_per_partition=layer.output_size_per_partition,
435
- input_size_per_partition=layer.input_size_per_partition,
436
- input_global_scale=getattr(layer, "input_global_scale", None),
437
- bias=bias,
438
- input_dtype=self.input_dtype,
439
- )
479
+ return self.kernel.apply_weights(layer, x, bias)
480
481
482
class AWQMarlinMoEMethod(FusedMoEMethodBase):
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+8
-9
@@ -46,14 +46,15 @@ def query_marlin_supported_quant_types(
46
if current_platform.is_cpu():
47
return _query_cpu_marlin_supported_quant_types(has_zp, include_fp_type)
48
49
- if device_capability is None:
50
- capability_tuple = current_platform.get_device_capability()
51
- device_capability = (
52
- -1 if capability_tuple is None else capability_tuple.to_int()
53
- )
49
+ if not current_platform.is_rocm():
50
+ if device_capability is None:
51
+ capability_tuple = current_platform.get_device_capability()
52
+ device_capability = (
53
+ -1 if capability_tuple is None else capability_tuple.to_int()
54
+ )
55
55
- if device_capability < 75:
56
- return []
56
+ if device_capability < 75:
57
+ return []
58
59
# - has_zp is True: return quant_types that has zero points
60
# - has_zp is False: return quant_types that has not zero points
@@ -210,8 +211,6 @@ def check_marlin_supports_shape(
211
212
213
def check_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool:
213
- if current_platform.is_rocm():
214
- return False
214
output_size_per_partition = (
215
getattr(layer, "output_size_per_partition", None) or layer.output_size
216
)