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 9, 21:48 UTC
[Quantization][Deprecation] Remove RTN (#32697)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Jan 21, 2026 at 11:34 UTC
4e31b7f2284b9c5b88e8504895794b9511980345
4 files changed
-730
tests/quantization/test_rtn.py
deleted
-39
@@ -1,39 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# Copyright © 2025, Oracle and/or its affiliates.
4
-"""Tests RTN quantization startup and generation,
5
-doesn't test correctness
6
-"""
7
-
8
-import pytest
9
-
10
-from tests.quantization.utils import is_quant_method_supported
11
-
12
-MODELS = [
13
- "ai21labs/Jamba-tiny-dev", # MoE model
14
-]
15
-
16
-
17
-@pytest.mark.skipif(
18
- not is_quant_method_supported("rtn"),
19
- reason="RTN is not supported on this GPU type.",
20
-)
21
-@pytest.mark.parametrize("model", MODELS)
22
-@pytest.mark.parametrize("dtype", ["bfloat16"])
23
-@pytest.mark.parametrize("max_tokens", [10])
24
-def test_model_rtn_startup(
25
- hf_runner,
26
- vllm_runner,
27
- example_prompts,
28
- model: str,
29
- dtype: str,
30
- max_tokens: int,
31
-) -> None:
32
- with vllm_runner(
33
- model,
34
- enforce_eager=True,
35
- dtype=dtype,
36
- quantization="rtn",
37
- allow_deprecated_quantization=True,
38
- ) as vllm_model:
39
- vllm_model.generate_greedy(example_prompts, max_tokens)
vllm/model_executor/layers/quantization/__init__.py
-4
@@ -31,7 +31,6 @@ QuantizationMethods = Literal[
31
"quark",
32
"moe_wna16",
33
"torchao",
34
- "rtn",
34
"inc",
35
"mxfp4",
36
"petit_nvfp4",
@@ -49,7 +48,6 @@ DEPRECATED_QUANTIZATION_METHODS = [
48
"gptq_bitblas",
49
"experts_int8",
50
"ipex",
52
- "rtn",
51
"petit_nvfp4",
52
]
53
@@ -138,7 +136,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
136
from .mxfp4 import Mxfp4Config
137
from .petit import PetitNvFp4Config
138
from .ptpc_fp8 import PTPCFp8Config
141
- from .rtn import RTNConfig
139
from .torchao import TorchAOConfig
140
141
method_to_config: dict[str, type[QuantizationConfig]] = {
@@ -163,7 +160,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
160
"quark": QuarkConfig,
161
"moe_wna16": MoeWNA16Config,
162
"torchao": TorchAOConfig,
166
- "rtn": RTNConfig,
163
"auto-round": INCConfig,
164
"inc": INCConfig,
165
"mxfp4": Mxfp4Config,
vllm/model_executor/layers/quantization/rtn.py
deleted
-626
@@ -1,626 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# Copyright © 2025, Oracle and/or its affiliates.
4
-
5
-import os
6
-from typing import Any, Optional
7
-
8
-import numpy as np
9
-import torch
10
-from torch.nn.parameter import Parameter
11
-
12
-from vllm.logger import init_logger
13
-from vllm.model_executor.layers.fused_moe import FusedMoERouter
14
-from vllm.model_executor.layers.fused_moe.config import (
15
- FusedMoEConfig,
16
- FusedMoEQuantConfig,
17
-)
18
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
19
-from vllm.model_executor.layers.fused_moe.layer import (
20
- FusedMoE,
21
- FusedMoEMethodBase,
22
-)
23
-from vllm.model_executor.layers.linear import (
24
- LinearBase,
25
- LinearMethodBase,
26
- set_weight_attrs,
27
-)
28
-from vllm.model_executor.layers.quantization import QuantizationMethods
29
-from vllm.model_executor.layers.quantization.base_config import (
30
- QuantizationConfig,
31
- QuantizeMethodBase,
32
-)
33
-from vllm.model_executor.layers.quantization.utils import replace_parameter
34
-from vllm.model_executor.layers.quantization.utils.marlin_utils import (
35
- apply_rtn_marlin_linear,
36
- marlin_make_workspace_new,
37
-)
38
-from vllm.scalar_type import scalar_types
39
-
40
-logger = init_logger(__name__)
41
-"""By default, use 8 bit as target precision, but it can be
42
-overridden by setting the RTN_NUM_BITS envvar
43
-"""
44
-NUM_BITS = os.getenv("RTN_NUM_BITS", "8")
45
-"""By default, use group size of 128 parameters, but it can be
46
-overridden by setting the RTN_GROUP_SIZE envvar
47
-"""
48
-GROUP_SIZE = os.getenv("RTN_GROUP_SIZE", "128")
49
-"""Global Marlin workspace shared by all modules
50
-"""
51
-workspace = None
52
-
53
-
54
-class RTNConfig(QuantizationConfig):
55
- """Config class for RTN."""
56
-
57
- def __init__(
58
- self,
59
- weight_bits: int = int(NUM_BITS),
60
- group_size: int = int(GROUP_SIZE),
61
- ) -> None:
62
- self.weight_bits = weight_bits
63
- self.group_size = group_size
64
-
65
- if self.weight_bits != 4 and self.weight_bits != 8:
66
- raise ValueError(
67
- "Currently, only 4-bit or 8-bit weight quantization is "
68
- f"supported for RTN, but got {self.weight_bits} bits."
69
- )
70
-
71
- self.quant_type = (
72
- scalar_types.uint8b128 if self.weight_bits == 8 else scalar_types.uint4b8
73
- )
74
-
75
- def __repr__(self) -> str:
76
- return (
77
- f"RTNConfig(weight_bits={self.weight_bits}, group_size={self.group_size})"
78
- )
79
-
80
- @classmethod
81
- def get_name(cls) -> QuantizationMethods:
82
- return "rtn"
83
-
84
- @classmethod
85
- def get_supported_act_dtypes(cls) -> list[torch.dtype]:
86
- return [torch.bfloat16, torch.half]
87
-
88
- @classmethod
89
- def get_min_capability(cls) -> int:
90
- return 80
91
-
92
- @classmethod
93
- def get_config_filenames(cls) -> list[str]:
94
- return []
95
-
96
- @classmethod
97
- def from_config(cls, config: dict[str, Any]) -> "RTNConfig":
98
- weight_bits = cls.get_from_keys(config, ["bits"])
99
- group_size = cls.get_from_keys(config, ["group_size"])
100
- return cls(weight_bits, group_size)
101
-
102
- def get_quant_method(
103
- self, layer: torch.nn.Module, prefix: str
104
- ) -> Optional["QuantizeMethodBase"]:
105
- if isinstance(layer, LinearBase):
106
- return RTNLinearMethod(self)
107
- elif isinstance(layer, FusedMoE):
108
- return RTNMoEMethod(self, layer.moe_config)
109
- return None
110
-
111
-
112
-class RTNTensor:
113
- """A wrapper over Tensor that enables quantization on-the-fly by
114
- overloading the copy_ method.
115
- """
116
-
117
- def __init__(
118
- self, data: torch.Tensor, scale: torch.Tensor, quant_config: RTNConfig
119
- ) -> None:
120
- self.data = data
121
- self.scale = scale
122
- self.quant_config = quant_config
123
-
124
- def narrow(self, dim, start, length):
125
- factor = 1 if self.quant_config.weight_bits == 8 else 2
126
- return RTNTensor(
127
- self.data.narrow(dim, start // factor, length // factor),
128
- self.scale.narrow(dim, start, length),
129
- self.quant_config,
130
- )
131
-
132
- def __getitem__(self, key):
133
- return RTNTensor(self.data[key], self.scale[key], self.quant_config)
134
-
135
- @property
136
- def shape(self):
137
- shape = self.data.shape
138
- factor = 1 if self.quant_config.weight_bits == 8 else 2
139
- batch_present = len(shape) == 3
140
- if batch_present:
141
- return torch.Size((shape[0], shape[1] * factor, shape[2]))
142
- else:
143
- return torch.Size((shape[0] * factor, shape[1]))
144
-
145
- def copy_(self, loaded_weight: torch.Tensor) -> None:
146
- qweight, weight_scale = rtn_quantize(
147
- loaded_weight.cuda(),
148
- self.quant_config.weight_bits,
149
- self.quant_config.group_size,
150
- )
151
-
152
- self.data.copy_(qweight)
153
- self.scale.data.copy_(weight_scale)
154
-
155
-
156
-class RTNParameter(Parameter):
157
- """A wrapper over Parameter that returns RTNTensor (a wrapper over Tensor)
158
- when its data is accessed. We need this wrapper for the data loading phase
159
- only, so we can intercept a weight copying function (torch.Tensor.copy_)
160
- and apply quantization on-the-fly.
161
- """
162
-
163
- def __new__(cls, data: torch.Tensor, **kwargs):
164
- return super().__new__(cls, data=data, requires_grad=False)
165
-
166
- def __init__(
167
- self, data: torch.Tensor, scale: torch.Tensor, quant_config: RTNConfig
168
- ) -> None:
169
- self.scale = scale
170
- self.quant_config = quant_config
171
-
172
- @property
173
- def data(self):
174
- return RTNTensor(super().data, self.scale, self.quant_config)
175
-
176
-
177
-class RTNLinearMethod(LinearMethodBase):
178
- """Linear method for RTN.
179
-
180
- Args:
181
- quant_config: The RTN quantization config.
182
- """
183
-
184
- def __init__(self, quant_config: RTNConfig):
185
- self.quant_config = quant_config
186
-
187
- def create_weights(
188
- self,
189
- layer: torch.nn.Module,
190
- input_size_per_partition: int,
191
- output_partition_sizes: list[int],
192
- input_size: int,
193
- output_size: int,
194
- params_dtype: torch.dtype,
195
- **extra_weight_attrs,
196
- ):
197
- output_size_per_partition = sum(output_partition_sizes)
198
- num_groups_per_col = (
199
- input_size_per_partition // self.quant_config.group_size
200
- if self.quant_config.group_size != -1
201
- else 1
202
- )
203
-
204
- scale = Parameter(
205
- torch.empty(
206
- output_size_per_partition, num_groups_per_col, dtype=params_dtype
207
- ),
208
- requires_grad=False,
209
- )
210
- factor = 1 if self.quant_config.weight_bits == 8 else 2
211
-
212
- weight = RTNParameter(
213
- data=torch.empty(
214
- output_size_per_partition // factor,
215
- input_size_per_partition,
216
- dtype=torch.uint8,
217
- ),
218
- scale=scale,
219
- quant_config=self.quant_config,
220
- )
221
-
222
- layer.register_parameter("weight", weight)
223
- set_weight_attrs(
224
- weight,
225
- {
226
- **extra_weight_attrs,
227
- "input_dim": 1,
228
- "output_dim": 0,
229
- },
230
- )
231
-
232
- layer.register_parameter("scale", scale)
233
- layer.output_size_per_partition = output_size_per_partition
234
-
235
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
236
- """Repack weights and scales for Marlin kernels."""
237
- weight_bits = self.quant_config.weight_bits
238
-
239
- weight, scale = repack_weights(layer.weight, layer.scale, weight_bits)
240
-
241
- replace_parameter(layer, "weight", weight)
242
- replace_parameter(layer, "scale", scale)
243
-
244
- init_workspace(layer.weight.device)
245
-
246
- def apply(
247
- self,
248
- layer: torch.nn.Module,
249
- x: torch.Tensor,
250
- bias: torch.Tensor | None = None,
251
- ) -> torch.Tensor:
252
- return apply_rtn_marlin_linear(
253
- input=x,
254
- weight=layer.weight,
255
- weight_scale=layer.scale,
256
- workspace=workspace,
257
- quant_type=self.quant_config.quant_type,
258
- output_size_per_partition=layer.output_size_per_partition,
259
- input_size_per_partition=layer.input_size_per_partition,
260
- bias=bias,
261
- )
262
-
263
-
264
-class RTNMoEMethod(FusedMoEMethodBase):
265
- def __init__(self, quant_config: RTNConfig, moe: FusedMoEConfig):
266
- super().__init__(moe)
267
- self.quant_config = quant_config
268
-
269
- def create_weights(
270
- self,
271
- layer: torch.nn.Module,
272
- num_experts: int,
273
- hidden_size: int,
274
- intermediate_size_per_partition: int,
275
- params_dtype: torch.dtype,
276
- **extra_weight_attrs,
277
- ):
278
- factor = 1 if self.quant_config.weight_bits == 8 else 2
279
-
280
- # Fused gate_up_proj (column parallel)
281
- num_groups_per_col = (
282
- hidden_size // self.quant_config.group_size
283
- if self.quant_config.group_size != -1
284
- else 1
285
- )
286
- w13_scale = Parameter(
287
- torch.empty(
288
- num_experts,
289
- 2 * intermediate_size_per_partition,
290
- num_groups_per_col,
291
- dtype=params_dtype,
292
- ),
293
- requires_grad=False,
294
- )
295
- layer.register_parameter("w13_scale", w13_scale)
296
-
297
- w13_weight = RTNParameter(
298
- data=torch.empty(
299
- num_experts,
300
- 2 * intermediate_size_per_partition // factor,
301
- hidden_size,
302
- dtype=torch.uint8,
303
- ),
304
- scale=w13_scale,
305
- quant_config=self.quant_config,
306
- )
307
- layer.register_parameter("w13_weight", w13_weight)
308
- set_weight_attrs(w13_weight, extra_weight_attrs)
309
-
310
- # down_proj (row parallel)
311
- num_groups_per_col = (
312
- intermediate_size_per_partition // self.quant_config.group_size
313
- if self.quant_config.group_size != -1
314
- else 1
315
- )
316
- w2_scale = Parameter(
317
- torch.zeros(
318
- num_experts, hidden_size, num_groups_per_col, dtype=params_dtype
319
- ),
320
- requires_grad=False,
321
- )
322
- layer.register_parameter("w2_scale", w2_scale)
323
-
324
- w2_weight = RTNParameter(
325
- data=torch.empty(
326
- num_experts,
327
- hidden_size // factor,
328
- intermediate_size_per_partition,
329
- dtype=torch.uint8,
330
- ),
331
- scale=w2_scale,
332
- quant_config=self.quant_config,
333
- )
334
- layer.register_parameter("w2_weight", w2_weight)
335
- set_weight_attrs(w2_weight, extra_weight_attrs)
336
-
337
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
338
- """Repack weights and scales for Marlin kernels."""
339
- weight_bits = self.quant_config.weight_bits
340
-
341
- w13_weight, w13_scale = repack_weights(
342
- layer.w13_weight, layer.w13_scale, weight_bits
343
- )
344
- replace_parameter(layer, "w13_weight", w13_weight)
345
- replace_parameter(layer, "w13_scale", w13_scale)
346
-
347
- w2_weight, w2_scale = repack_weights(
348
- layer.w2_weight, layer.w2_scale, weight_bits
349
- )
350
- replace_parameter(layer, "w2_weight", w2_weight)
351
- replace_parameter(layer, "w2_scale", w2_scale)
352
-
353
- init_workspace(layer.w13_weight.device)
354
-
355
- def get_fused_moe_quant_config(
356
- self, layer: torch.nn.Module
357
- ) -> FusedMoEQuantConfig | None:
358
- return None
359
-
360
- def apply(
361
- self,
362
- layer: FusedMoE,
363
- router: FusedMoERouter,
364
- x: torch.Tensor,
365
- router_logits: torch.Tensor,
366
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
367
- topk_weights, topk_ids = router.select_experts(
368
- hidden_states=x,
369
- router_logits=router_logits,
370
- )
371
-
372
- return fused_marlin_moe(
373
- x,
374
- layer.w13_weight,
375
- layer.w2_weight,
376
- getattr(layer, "w13_bias", None),
377
- getattr(layer, "w2_bias", None),
378
- layer.w13_scale,
379
- layer.w2_scale,
380
- router_logits,
381
- topk_weights,
382
- topk_ids,
383
- quant_type_id=self.quant_config.quant_type.id,
384
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
385
- global_num_experts=layer.global_num_experts,
386
- expert_map=layer.expert_map,
387
- workspace=workspace,
388
- )
389
-
390
-
391
-def rtn_quantize(
392
- tensor: torch.Tensor, num_bits: int, group_size: int
393
-) -> tuple[torch.Tensor, torch.Tensor]:
394
- """Quantize a tensor using per-group static scaling factor.
395
-
396
- Args:
397
- tensor: The input tensor.
398
- num_bits: Target precision for the result (supported values are
399
- 8 or 4).
400
- group_size: Quantization granularity.
401
- If equal to -1, each row in the input tensor is treated
402
- as one group.
403
- """
404
- batch_present = len(tensor.shape) == 3
405
- if not batch_present:
406
- tensor = tensor.unsqueeze(0)
407
-
408
- q_range = 2**num_bits
409
- num_groups = (
410
- tensor.shape[1] * tensor.shape[2] // group_size
411
- if group_size != -1
412
- else tensor.shape[1]
413
- )
414
- """Calculate a scaling factor per input group.
415
- """
416
- input_flat = tensor.reshape(tensor.shape[0], num_groups, -1)
417
- input_min = torch.min(input_flat, dim=2, keepdim=True)[0]
418
- input_max = torch.max(input_flat, dim=2, keepdim=True)[0]
419
- input_max_abs = torch.max(input_min.abs(), input_max.abs())
420
- scale = input_max_abs * 2.0 / (q_range - 1)
421
- """Scale each input group, round to the nearest integer, shift
422
- the range and truncate.
423
- """
424
- scaled_input = input_flat / scale
425
- scaled_input = scaled_input.round()
426
- scaled_input += q_range // 2
427
- scaled_input = scaled_input.clamp(0, q_range - 1)
428
-
429
- scale = scale.reshape(tensor.shape[0], tensor.shape[1], -1).contiguous()
430
- inputs_q = scaled_input.reshape(tensor.shape).to(torch.uint8)
431
- inputs_q = inputs_q.contiguous()
432
-
433
- if num_bits == 4:
434
- """Pack two 4-bit values into each byte.
435
- """
436
- inputs_q = (inputs_q[:, :, 1::2] << 4) | (inputs_q[:, :, ::2] & 0xF)
437
- inputs_q = inputs_q.reshape(
438
- tensor.shape[0], tensor.shape[1] // 2, tensor.shape[2]
439
- )
440
- inputs_q = inputs_q.contiguous()
441
-
442
- if not batch_present:
443
- inputs_q = inputs_q.squeeze(0)
444
- scale = scale.squeeze(0)
445
-
446
- return inputs_q, scale
447
-
448
-
449
-def rtn_dequantize(tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
450
- """Dequantize a tensor using per-group static scaling factors.
451
-
452
- Args:
453
- tensor: The input tensor.
454
- scale: The tensor with per-group scale factors.
455
- """
456
- batch_present = len(tensor.shape) == 3
457
- if not batch_present:
458
- tensor = tensor.unsqueeze(0)
459
- scale = scale.unsqueeze(0)
460
-
461
- num_groups = scale.size(1) * scale.size(2)
462
- batch, input_dim, output_dim = tensor.shape
463
-
464
- num_bits = 8 if input_dim == scale.size(1) else 4
465
- q_range = 2**num_bits
466
- if num_bits == 4:
467
- input_dim *= 2
468
-
469
- data = torch.empty(
470
- (batch, input_dim, output_dim), dtype=scale.dtype, device=tensor.device
471
- )
472
-
473
- if num_bits == 8:
474
- data.copy_(tensor)
475
- data -= q_range // 2
476
- else:
477
- """Unpack two 4-bit values from each byte.
478
- """
479
- tensor = tensor.reshape(batch, input_dim, output_dim // 2)
480
- for i in range(2):
481
- data[:, :, i::2] = ((tensor << 4 * (1 - i)) >> 4).to(
482
- torch.int8
483
- ) - q_range // 2
484
- """Scale each input group with its scaling factor.
485
- """
486
- scale = scale.reshape(batch, num_groups, -1)
487
- data = data.reshape(batch, num_groups, -1)
488
- data = torch.mul(data, scale)
489
-
490
- input_deq = data.reshape((batch, input_dim, output_dim)).contiguous()
491
- if not batch_present:
492
- input_deq = input_deq.squeeze(0)
493
-
494
- return input_deq
495
-
496
-
497
-def _get_perms():
498
- perm = []
499
- for i in range(32):
500
- perm1 = []
501
- col = i // 4
502
- for block in [0, 1]:
503
- for row in [
504
- 2 * (i % 4),
505
- 2 * (i % 4) + 1,
506
- 2 * (i % 4 + 4),
507
- 2 * (i % 4 + 4) + 1,
508
- ]:
509
- perm1.append(16 * row + col + 8 * block)
510
- for j in range(4):
511
- perm.extend([p + 256 * j for p in perm1])
512
-
513
- perm_arr = np.array(perm)
514
- interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
515
- perm_arr = perm_arr.reshape((-1, 8))[:, interleave].ravel()
516
- perm_tensor = torch.from_numpy(perm_arr)
517
- scale_perm = []
518
- for i in range(8):
519
- scale_perm.extend([i + 8 * j for j in range(8)])
520
- scale_perm_single = []
521
- for i in range(4):
522
- scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])
523
- return perm_tensor, scale_perm, scale_perm_single
524
-
525
-
526
-_perm, _scale_perm, _scale_perm_single = _get_perms()
527
-
528
-
529
-def pack_for_marlin(weight, scale, qbits):
530
- batch = weight.shape[0]
531
-
532
- n = weight.size(1)
533
- k = weight.size(2)
534
- groupsize = k // scale.size(2)
535
-
536
- tile = 16
537
- s = scale.permute(0, 2, 1) # transpose
538
- w = weight.permute(0, 2, 1) # transpose
539
- if groupsize != k:
540
- w = w.reshape((batch, -1, groupsize, n))
541
- w = w.permute(0, 2, 1, 3)
542
- w = w.reshape((batch, groupsize, -1))
543
- s = s.reshape((batch, 1, -1))
544
-
545
- if groupsize != k:
546
- w = w.reshape((batch, groupsize, -1, n))
547
- w = w.permute(0, 2, 1, 3)
548
- w = w.reshape((batch, k, n)).contiguous()
549
- s = s.reshape((batch, -1, len(_scale_perm)))[:, :, _scale_perm]
550
- else:
551
- s = s.reshape((batch, -1, len(_scale_perm_single)))[:, :, _scale_perm_single]
552
- s = s.reshape((batch, -1, n)).contiguous()
553
- w = w.reshape((batch, k // tile, tile, n // tile, tile))
554
- w = w.permute((0, 1, 3, 2, 4))
555
- w = w.reshape((batch, k // tile, n * tile))
556
- res = w
557
- res = res.reshape((batch, -1, _perm.numel()))[:, :, _perm].reshape(res.shape)
558
- if qbits == 4:
559
- q = torch.zeros(
560
- (batch, res.shape[1], res.shape[2] // 2), dtype=torch.int8, device=w.device
561
- )
562
- for i in range(2):
563
- q |= res[:, :, i::2] << 4 * i
564
- q = q.reshape(batch, -1, n).contiguous()
565
- else:
566
- q = res.clone()
567
- q[:, :, 2::8] = res[:, :, 4::8]
568
- q[:, :, 3::8] = res[:, :, 5::8]
569
- q[:, :, 4::8] = res[:, :, 2::8]
570
- q[:, :, 5::8] = res[:, :, 3::8]
571
- q = q.reshape(batch, -1, n).to(torch.int8).contiguous()
572
-
573
- return q, s
574
-
575
-
576
-def repack_8bit_into_32bit(input):
577
- output = torch.zeros(
578
- (input.shape[0], input.shape[1], input.shape[2] // 4),
579
- dtype=torch.int32,
580
- device=input.device,
581
- )
582
- for i in range(4):
583
- output |= (input[:, :, i::4] & 0xFF).to(torch.int32) << 8 * i
584
-
585
- return output
586
-
587
-
588
-def repack_weights(qweight, scale, weight_bits):
589
- batch_present = len(qweight.shape) == 3
590
- if not batch_present:
591
- qweight = qweight.unsqueeze(0)
592
- scale = scale.unsqueeze(0)
593
-
594
- if weight_bits == 4:
595
- """Unpack two 4-bit values from each byte.
596
- """
597
- qweight_unpacked = torch.empty(
598
- (qweight.shape[0], qweight.shape[1] * 2, qweight.shape[2]),
599
- dtype=torch.uint8,
600
- device=qweight.device,
601
- )
602
- for i in range(2):
603
- qweight_unpacked[:, :, i::2] = ((qweight << 4 * (1 - i)) >> 4).reshape(
604
- qweight.shape[0], qweight.shape[1] * 2, qweight.shape[2] // 2
605
- )
606
- else:
607
- qweight_unpacked = qweight
608
-
609
- qweight_packed, scale_packed = pack_for_marlin(qweight_unpacked, scale, weight_bits)
610
- """Marlin kernels expect tensors in int32 format in a certain shape
611
- """
612
- qweight_repacked = repack_8bit_into_32bit(qweight_packed.to(torch.uint8))
613
- qweight_reshaped = qweight_repacked.reshape(
614
- qweight.shape[0], qweight.shape[2] // 16, -1
615
- )
616
- if not batch_present:
617
- qweight_reshaped = qweight_reshaped.squeeze(0)
618
- scale_packed = scale_packed.squeeze(0)
619
-
620
- return qweight_reshaped, scale_packed
621
-
622
-
623
-def init_workspace(device):
624
- global workspace
625
- if workspace is None:
626
- workspace = marlin_make_workspace_new(device, 4)
vllm/model_executor/layers/quantization/utils/marlin_utils.py
-61
@@ -650,64 +650,3 @@ def apply_awq_marlin_linear(
650
)
651
652
return output.reshape(out_shape)
653
-
654
-
655
-def apply_rtn_marlin_linear(
656
- input: torch.Tensor,
657
- weight: torch.Tensor,
658
- weight_scale: torch.Tensor,
659
- workspace: torch.Tensor,
660
- quant_type: ScalarType,
661
- output_size_per_partition: int,
662
- input_size_per_partition: int,
663
- input_global_scale: torch.Tensor | None = None,
664
- bias: torch.Tensor | None = None,
665
- use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
666
- input_dtype: torch.dtype | None = None,
667
-) -> torch.Tensor:
668
- reshaped_x = input.reshape(-1, input.shape[-1])
669
- out_shape = input.shape[:-1] + (output_size_per_partition,)
670
-
671
- use_atomic_add = should_use_atomic_add_reduce(
672
- m=reshaped_x.size(0),
673
- n=output_size_per_partition,
674
- k=reshaped_x.size(1),
675
- device=input.device,
676
- dtype=input.dtype,
677
- )
678
-
679
- a_scales = None
680
- if input_dtype == torch.int8:
681
- assert quant_type == scalar_types.uint4b8, (
682
- "W8A8-INT8 is not supported by marlin kernel."
683
- )
684
- reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
685
- a_scales = a_scales * input_global_scale
686
- elif input_dtype == torch.float8_e4m3fn:
687
- assert quant_type == scalar_types.uint4b8, (
688
- "INT8 weight + FP8 activation is not supported."
689
- )
690
- reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
691
-
692
- output = ops.gptq_marlin_gemm(
693
- reshaped_x,
694
- None,
695
- weight,
696
- bias,
697
- weight_scale,
698
- a_scales,
699
- None,
700
- None,
701
- None,
702
- None,
703
- workspace,
704
- quant_type,
705
- size_m=reshaped_x.shape[0],
706
- size_n=output_size_per_partition,
707
- size_k=input_size_per_partition,
708
- use_atomic_add=use_atomic_add,
709
- use_fp32_reduce=use_fp32_reduce,
710
- is_zp_float=False,
711
- )
712
-
713
- return output.reshape(out_shape)