@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 22feac8e95

Fork 0

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 14, 05:30 UTC

[Transform] [Quantization] Add transforms to compressed tensors (#22486)

Kyle Sayers committed Aug 28, 2025 at 02:43 UTC 22feac8e957a2f9787eb721c685269afc15bb3b1

9 files changed +661 -36

tests/conftest.py

+37 -6

@@ -1,10 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import json

4 +import math

5 import os

6 import tempfile

7 from enum import Enum

7 -from typing import Any, Callable, Optional, TypedDict, TypeVar, Union

8 +from typing import Any, Callable, Optional, TypedDict, TypeVar, Union, cast

10 import numpy as np

11 import pytest

@@ -33,6 +34,7 @@ from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,

34 from vllm.logger import init_logger

35 from vllm.outputs import RequestOutput

36 from vllm.sampling_params import BeamSearchParams

37 +from vllm.sequence import Logprob

38 from vllm.transformers_utils.utils import maybe_model_redirect

40 logger = init_logger(__name__)

@@ -602,7 +604,7 @@ class HfRunner:

604 def _hidden_states_to_logprobs(

605 self,

606 hidden_states: tuple[tuple[torch.Tensor, ...], ...],

605 - num_logprobs: int,

607 + num_logprobs: Optional[int],

608 ) -> tuple[list[dict[int, float]], int]:

609 seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)

610 output_len = len(hidden_states)

@@ -630,7 +632,7 @@ class HfRunner:

632 self,

633 prompts: list[str],

634 max_tokens: int,

633 - num_logprobs: int,

635 + num_logprobs: Optional[int],

636 images: Optional[PromptImageInput] = None,

637 audios: Optional[PromptAudioInput] = None,

638 videos: Optional[PromptVideoInput] = None,

@@ -677,7 +679,7 @@ class HfRunner:

679 self,

680 encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],

681 max_tokens: int,

680 - num_logprobs: int,

682 + num_logprobs: Optional[int],

683 images: Optional[PromptImageInput] = None,

684 **kwargs: Any,

685 ) -> list[TokensTextLogprobs]:

@@ -966,7 +968,7 @@ class VllmRunner:

968 self,

969 prompts: list[str],

970 max_tokens: int,

969 - num_logprobs: int,

971 + num_logprobs: Optional[int],

972 num_prompt_logprobs: Optional[int] = None,

973 images: Optional[PromptImageInput] = None,

974 audios: Optional[PromptAudioInput] = None,

@@ -991,11 +993,40 @@ class VllmRunner:

993 videos=videos,

994 **kwargs)

995

996 + def generate_prompt_perplexity(self, prompts: list[str]) -> list[float]:

997 + """

998 + Return the perplexity score associated with generating the prompts

999 +

1000 + :param prompts: list of prompts to score

1001 + :return: perplexity score of each prompt

1002 + """

1003 + outputs = self.generate_greedy_logprobs(prompts,

1004 + max_tokens=1,

1005 + num_logprobs=None,

1006 + num_prompt_logprobs=0)

1007 +

1008 + perplexities = []

1009 + for output in outputs:

1010 + output = cast(TokensTextLogprobsPromptLogprobs, output)

1011 + token_datas = cast(list[Optional[dict[int, Logprob]]], output[3])

1012 + assert token_datas[0] is None

1013 + token_log_probs = []

1014 + for token_data in token_datas[1:]:

1015 + assert token_data is not None

1016 + assert len(token_data) == 1

1017 + token_log_prob = list(token_data.values())[0].logprob

1018 + token_log_probs.append(token_log_prob)

1019 +

1020 + perplexity = math.exp(-sum(token_log_probs) / len(token_log_probs))

1021 + perplexities.append(perplexity)

1022 +

1023 + return perplexities

1024 +

1025 def generate_encoder_decoder_greedy_logprobs(

1026 self,

1027 encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],

1028 max_tokens: int,

998 - num_logprobs: int,

1029 + num_logprobs: Optional[int],

1030 num_prompt_logprobs: Optional[int] = None,

1031 skip_special_tokens: bool = True,

1032 ) -> Union[list[TokensTextLogprobs],

tests/quantization/test_compressed_tensors.py

+22

@@ -719,3 +719,25 @@ def test_compressed_tensors_w4a8_fp8(vllm_runner, args):

719 output = llm.generate_greedy("Hello my name is", max_tokens=20)

720 print(output)

721 assert output

722 +

723 +

724 +@pytest.mark.skipif(not current_platform.is_cuda(),

725 + reason="This test is skipped on non-CUDA platform.")

726 +@pytest.mark.parametrize("model,prompt,exp_perplexity", [

727 + (

728 + "nm-testing/Llama-3.2-1B-Instruct-spinquantR1R2R4-w4a16",

729 + "Flat is better than nested.\nSparse is better than dense.",

730 + 150.0,

731 + ),

732 + (

733 + "nm-testing/Llama-3.2-1B-Instruct-quip-w4a16",

734 + "Flat is better than nested.\nSparse is better than dense.",

735 + 150.0,

736 + ),

737 +])

738 +def test_compressed_tensors_transforms_perplexity(vllm_runner, model, prompt,

739 + exp_perplexity):

740 + with vllm_runner(model, enforce_eager=True) as llm:

741 + perplexity = llm.generate_prompt_perplexity([prompt])[0]

742 + print(perplexity)

743 + assert perplexity <= exp_perplexity

\ No newline at end of file

vllm/model_executor/layers/linear.py

+15 -1

@@ -35,6 +35,7 @@ logger = init_logger(__name__)

36 WEIGHT_LOADER_V2_SUPPORTED = [

37 "CompressedTensorsLinearMethod",

38 + "CompressedTensorsLinearTransformMethod",

39 "BitBLASLinearMethod",

40 "GPTQBitBLASLinearMethod",

41 "AWQMarlinLinearMethod",

@@ -199,6 +200,7 @@ class UnquantizedLinearMethod(LinearMethodBase):

200 set_weight_attrs(weight, extra_weight_attrs)

201

202 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

203 + # special postprocessing for CPU SGL

204 if current_platform.is_cpu() and envs.VLLM_CPU_SGL_KERNEL:

205 from vllm.model_executor.layers.utils import check_cpu_sgl_kernel

206 N, K = layer.weight.size()

@@ -1470,7 +1472,7 @@ class QKVCrossParallelLinear(LinearBase):

1472 self.bias = torch.nn.Parameter()

1473 set_weight_attrs(self.bias, {

1474 "output_dim": 0,

1473 - "weight_loader": self.weight_loader,

1475 + "weight_loader": self.weight_loader_v1,

1476 })

1477 else:

1478 self.bias = None

@@ -1580,6 +1582,18 @@ class QKVCrossParallelLinear(LinearBase):

1582 k, v = kv_enc.split(self.kv_size, dim=-1)

1583 return q, k, v

1584

1585 + def weight_loader_v1(self,

1586 + param: torch.nn.Parameter,

1587 + loaded_weight: torch.Tensor,

1588 + loaded_shard_id: Optional[str] = None):

1589 + # just like all other parameters, does not yet

1590 + # support loading bias with weight_loader_v2

1591 + layer = (self.q_proj_decoder

1592 + if loaded_shard_id == "q" else self.kv_proj_encoder)

1593 + target_param = self.select_proj_params(layer, param)

1594 + shard_id_args = (loaded_shard_id, ) if loaded_shard_id != "q" else ()

1595 + layer.weight_loader(target_param, loaded_weight, *shard_id_args)

1596 +

1597 def weight_loader(self,

1598 param: torch.nn.Parameter,

1599 loaded_weight: torch.Tensor,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

+37 -15

@@ -11,6 +11,7 @@ from compressed_tensors.config import (CompressionFormat,

11 from compressed_tensors.quantization import (QuantizationArgs,

12 QuantizationStrategy,

13 QuantizationType)

14 +from compressed_tensors.transform import TransformConfig

15 from pydantic import BaseModel

17 import vllm.envs as envs

@@ -30,6 +31,8 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (

31 CompressedTensorsW4A16Fp4, CompressedTensorsW4A16Sparse24,

32 CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,

33 CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)

34 +from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear import ( # noqa: E501

35 + CompressedTensorsLinearTransformMethod, get_linear_transform_schemes)

36 from vllm.model_executor.layers.quantization.compressed_tensors.utils import (

37 find_matched_target, is_activation_quantization_format,

38 should_ignore_layer)

@@ -60,6 +63,7 @@ class CompressedTensorsConfig(QuantizationConfig):

63 sparsity_ignore_list: list[str],

64 kv_cache_scheme: Optional[dict[str, Any]] = None,

65 config: Optional[dict[str, Any]] = None,

66 + transform_config: Optional[TransformConfig] = None,

67 ):

68 super().__init__()

69 self.ignore = ignore

@@ -71,6 +75,12 @@ class CompressedTensorsConfig(QuantizationConfig):

75 self.sparsity_ignore_list = sparsity_ignore_list

76 self.config = config

78 + if transform_config is not None:

79 + self.transform_config = TransformConfig.model_validate(

80 + transform_config)

81 + else:

82 + self.transform_config = None

83 +

84 def get_linear_method(self) -> "CompressedTensorsLinearMethod":

85 return CompressedTensorsLinearMethod(self)

@@ -103,18 +113,27 @@ class CompressedTensorsConfig(QuantizationConfig):

113 ) -> Optional["QuantizeMethodBase"]:

114 from vllm.attention.layer import Attention # Avoid circular import

115

106 - # Check if the layer is skipped for quantization.

107 - # TODO (@robertgshaw2): support module names

108 - if should_ignore_layer(prefix,

109 - ignore=self.ignore,

110 - fused_mapping=self.packed_modules_mapping):

111 - return UnquantizedLinearMethod()

116 if isinstance(layer, LinearBase):

113 - scheme = self.get_scheme(layer=layer, layer_name=prefix)

114 - if scheme is None:

115 - return UnquantizedLinearMethod()

116 - layer.scheme = scheme

117 - return CompressedTensorsLinearMethod(self)

117 + # collect schemes

118 + quant_scheme = self.get_scheme(layer=layer, layer_name=prefix)

119 + input_tfms, output_tfms = get_linear_transform_schemes(

120 + layer, prefix, self.transform_config,

121 + self.packed_modules_mapping)

122 +

123 + # choose quantization method

124 + quant_method: LinearMethodBase = UnquantizedLinearMethod()

125 + if quant_scheme is not None:

126 + layer.scheme = quant_scheme

127 + quant_method = CompressedTensorsLinearMethod(self)

128 +

129 + # choose transform method

130 + if any((input_tfms, output_tfms)):

131 + return CompressedTensorsLinearTransformMethod.from_schemes(

132 + quant_method, input_tfms, output_tfms)

133 +

134 + else:

135 + return quant_method

136 +

137 if isinstance(layer, Attention):

138 return CompressedTensorsKVCacheMethod(self)

139 if isinstance(layer, FusedMoE):

@@ -129,6 +148,7 @@ class CompressedTensorsConfig(QuantizationConfig):

148 config=config)

149 sparsity_scheme_map, sparsity_ignore_list = cls._parse_sparsity_config(

150 config=config)

151 + transform_config = config.get("transform_config")

152

153 return cls(

154 target_scheme_map=target_scheme_map,

@@ -137,6 +157,7 @@ class CompressedTensorsConfig(QuantizationConfig):

157 sparsity_scheme_map=sparsity_scheme_map,

158 sparsity_ignore_list=sparsity_ignore_list,

159 config=config,

160 + transform_config=transform_config,

161 )

162

163 @classmethod

@@ -537,9 +558,11 @@ class CompressedTensorsConfig(QuantizationConfig):

558

559 # Find the "target" in the compressed-tensors config

560 # that our layer conforms to.

540 - # TODO (@robertgshaw): add compressed-tensors as dep

541 - # so we do not have to re-write these functions

542 - # need to make accelerate optional in ct to do this

561 + # TODO (@kylesayrs): support ignore module names with ct matching utils

562 + if should_ignore_layer(layer_name,

563 + ignore=self.ignore,

564 + fused_mapping=self.packed_modules_mapping):

565 + return None

566

567 # Will be empty for models with only sparsity

568 weight_quant = input_quant = None

@@ -722,7 +745,6 @@ class CompressedTensorsLinearMethod(LinearMethodBase):

745 layer input. See LinearMethodBase for param details

746

747 """

725 -

748 scheme = layer.scheme

749 if scheme is None:

750 raise ValueError("A scheme must be defined for each layer")

vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py new

+227

@@ -0,0 +1,227 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Generator

4 +from itertools import accumulate

5 +from typing import Callable, Optional

6 +

7 +import torch

8 +from compressed_tensors.transform import (TransformArgs, TransformConfig,

9 + TransformLocation, TransformScheme)

10 +from compressed_tensors.utils import is_match

11 +

12 +from vllm.model_executor.layers.linear import (WEIGHT_LOADER_V2_SUPPORTED,

13 + LinearMethodBase,

14 + QKVCrossParallelLinear)

15 +from vllm.model_executor.layers.quantization.compressed_tensors.transform.module import ( # noqa: E501

16 + HadamardTransform)

17 +from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501

18 + TransformTuple)

19 +

20 +

21 +class CompressedTensorsLinearTransformMethod(LinearMethodBase):

22 + """

23 + Wraps `CompressedTensorsLinearMethod` or `UnquantizedLinearMethod` and adds

24 + input and output transforms to either side of the original apply method

25 + """

26 +

27 + @classmethod

28 + def from_schemes(

29 + cls, quant_method: LinearMethodBase, input_tfms: dict[int,

30 + TransformTuple],

31 + output_tfms: dict[int, TransformTuple]

32 + ) -> "CompressedTensorsLinearTransformMethod":

33 + assert input_tfms or output_tfms

34 +

35 + # TODO (@ksayers): implement QutlassLinearMethodNvFP4

36 + # hadacore and fwht can be selected by Transform module

37 +

38 + return cls(quant_method, input_tfms, output_tfms)

39 +

40 + def __init__(self, quant_method: LinearMethodBase,

41 + input_tfms: dict[int, TransformTuple],

42 + output_tfms: dict[int, TransformTuple]):

43 + self.quant_method = quant_method

44 + self.input_tfms = input_tfms

45 + self.output_tfms = output_tfms

46 +

47 + self.input_transform: Optional[HadamardTransform] = None

48 + self.output_transform: Optional[HadamardTransform] = None

49 +

50 + def create_weights(self, layer: torch.nn.Module,

51 + input_size_per_partition: int,

52 + output_partition_sizes: list[int], input_size: int,

53 + output_size: int, params_dtype: torch.dtype,

54 + **extra_weight_attrs):

55 +

56 + # get weight loader for transforms

57 + weight_loader: Callable = extra_weight_attrs.get(

58 + "weight_loader") # type: ignore[assignment]

59 +

60 + # HACK: UnquantizedLinearMethod does not support weight loader v2, but

61 + # transforms (specifically SharedWeightParameter) requires

62 + # weight loader v2. Until UnquantizedLinearMethod supports v2, we must

63 + # hack around this by getting weight loader v1 so ULM can load correctly

64 + quant_method_name = self.quant_method.__class__.__name__

65 + if quant_method_name not in WEIGHT_LOADER_V2_SUPPORTED:

66 + if isinstance(layer, QKVCrossParallelLinear):

67 + weight_loader_v1 = layer.weight_loader_v1

68 + else:

69 + weight_loader_v1 = layer.weight_loader

70 + extra_weight_attrs["weight_loader"] = weight_loader_v1

71 +

72 + self.quant_method.create_weights(

73 + layer=layer,

74 + input_size_per_partition=input_size_per_partition,

75 + output_partition_sizes=output_partition_sizes,

76 + input_size=input_size,

77 + output_size=output_size,

78 + params_dtype=params_dtype,

79 + **extra_weight_attrs)

80 +

81 + # validate schemes

82 + num_partitions = len(output_partition_sizes)

83 + self._validate_tfm_schemes(num_partitions)

84 +

85 + # create submodules for weight loading

86 + if len(self.input_tfms) > 0:

87 + scheme_name = list(self.input_tfms.values())[0].scheme_name

88 + location = list(self.input_tfms.values())[0].args.location

89 + transform_name = f"{scheme_name}_{location}"

90 +

91 + transform = HadamardTransform(self.input_tfms, layer,

92 + weight_loader,

93 + input_size_per_partition,

94 + output_partition_sizes)

95 + layer.register_module(transform_name, transform)

96 + self.input_transform = transform

97 +

98 + if len(self.output_tfms) > 0:

99 + scheme_name = list(self.output_tfms.values())[0].scheme_name

100 + location = list(self.output_tfms.values())[0].args.location

101 + transform_name = f"{scheme_name}_{location}"

102 +

103 + transform = HadamardTransform(self.output_tfms, layer,

104 + weight_loader,

105 + input_size_per_partition,

106 + output_partition_sizes)

107 + layer.register_module(transform_name, transform)

108 + self.output_transform = transform

109 +

110 + # compute partition ranges for slicing activations

111 + starts = [0] + list(accumulate(output_partition_sizes))[:-1]

112 + self.partition_ranges = list(zip(starts, output_partition_sizes))

113 +

114 + def process_weights_after_loading(self, layer):

115 + self.quant_method.process_weights_after_loading(layer)

116 +

117 + for submodule in layer.children():

118 + if isinstance(submodule, HadamardTransform):

119 + submodule.process_weights_after_loading()

120 +

121 + def apply(self,

122 + layer: torch.nn.Module,

123 + x: torch.Tensor,

124 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:

125 +

126 + if self.input_transform is not None:

127 + x = self.input_transform(x)

128 +

129 + assert bias is None

130 + x = self.quant_method.apply(layer, x, bias)

131 +

132 + # TODO (@ksayers): Write a triton kernel to do this in parallel

133 + if self.output_transform is not None:

134 + for part_id, (start, length) in enumerate(self.partition_ranges):

135 + x[:, start:start + length] = self.output_transform(

136 + x[:, start:start + length], part_id=part_id)

137 +

138 + return x

139 +

140 + def _validate_tfm_schemes(self, num_partitions: int):

141 + if len(self.input_tfms) > 0:

142 + if 0 not in self.input_tfms:

143 + raise ValueError("Must have same input")

144 +

145 + for part_index in range(num_partitions):

146 + if self.input_tfms[part_index] != self.input_tfms[0]:

147 + raise ValueError("Must have same input")

148 +

149 + if len(self.output_tfms) > 0:

150 + scheme_name = list(self.output_tfms.values())[0].scheme_name

151 + location = list(self.output_tfms.values())[0].args.location

152 +

153 + for tfm in self.output_tfms.values():

154 + if tfm.scheme_name != scheme_name:

155 + raise ValueError("Must have same scheme name")

156 + if tfm.args.location != location:

157 + raise ValueError("Must have same location")

158 +

159 + return self.input_tfms, self.output_tfms

160 +

161 +

162 +def get_linear_transform_schemes(

163 + layer: torch.nn.Module, layer_name: str,

164 + transform_config: Optional[TransformConfig],

165 + packed_modules_mapping: dict[str, list[str]]

166 +) -> tuple[dict[int, TransformTuple], dict[

167 + int, TransformTuple]]: # [input_transform, [output_transform, ...]]

168 + # there can only be one transform input scheme per (fused) module

169 + input_tfms = {}

170 + output_tfms = {}

171 +

172 + partition_names = get_layer_partition_names(layer_name,

173 + packed_modules_mapping)

174 +

175 + for scheme_name, scheme, args in get_schemes_args(transform_config):

176 + for part_index, part_name in enumerate(partition_names):

177 + if is_match(part_name, layer, args.targets,

178 + args.ignore) and args.is_online():

179 + if args.location == TransformLocation.INPUT:

180 + input_tfms[part_index] = TransformTuple(

181 + scheme_name, scheme, args)

182 +

183 + elif args.location == TransformLocation.OUTPUT:

184 + output_tfms[part_index] = TransformTuple(

185 + scheme_name, scheme, args)

186 +

187 + else:

188 + raise ValueError(f"Cannot apply `{args.location}` "

189 + f"transform to `{layer_name}`")

190 +

191 + return (input_tfms, output_tfms)

192 +

193 +

194 +def get_schemes_args(

195 + transform_config: Optional[TransformConfig]

196 +) -> Generator[tuple[str, TransformScheme, TransformArgs]]:

197 + if transform_config is None:

198 + return

199 +

200 + for scheme_name, scheme in transform_config.config_groups.items():

201 + for args in scheme.apply:

202 + yield (scheme_name, scheme, args)

203 +

204 +

205 +def get_layer_partition_names(

206 + layer_name: str, packed_modules_mapping: dict[str,

207 + list[str]]) -> list[str]:

208 + """

209 + Get all partition names associated with this layer.

210 + Names are returned in order of their partition indices.

211 +

212 + ```python

213 + mapping = {"gate_up_proj", "gate_proj", "up_proj"}

214 +

215 + assert get_layer_partition_names(

216 + "mlp.gate_up_proj", mapping) == ["gate_proj", "up_proj"]

217 + assert get_layer_partition_names(

218 + "mlp.down_proj", mapping) == ["down_proj"]

219 + """

220 + for fused_suffix, part_suffixes in packed_modules_mapping.items():

221 + if layer_name.endswith(fused_suffix):

222 + return [

223 + layer_name.removesuffix(fused_suffix) + part_suffix

224 + for part_suffix in part_suffixes

225 + ]

226 +

227 + return [layer_name]

vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py new

+135

@@ -0,0 +1,135 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import math

4 +from collections.abc import Hashable

5 +from typing import Callable, Optional

6 +

7 +import torch

8 +from compressed_tensors.transform import TransformLocation, TransformScheme

9 +from torch import Tensor

10 +

11 +from vllm.distributed.parallel_state import (

12 + get_tensor_model_parallel_world_size)

13 +from vllm.model_executor.layers.linear import LinearBase

14 +from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501

15 + TransformTuple)

16 +from vllm.model_executor.layers.utils import dispatch_unquantized_gemm

17 +from vllm.model_executor.layers.vocab_parallel_embedding import (

18 + VocabParallelEmbedding)

19 +from vllm.model_executor.parameter import SharedWeightParameter

20 +

21 +

22 +class HadamardTransform(torch.nn.Module):

23 + """

24 + Class which handles weight loading, postprocessing, and application of

25 + transforms. Meant to be used with `CompressedTensorsLinearTransformMethod`

26 + and attention transforms method (not implemented yet)

27 + """

28 + transforms: dict[int, TransformTuple] # info parsed from transforms config

29 + weight: SharedWeightParameter # container for shared tensors

30 +

31 + kernel: Callable # function used during application

32 + scales: dict[int, float] # hadamard scale, usually sqrt(matrix.size(0))

33 +

34 + def __init__(self,

35 + transforms: dict[int, TransformTuple],

36 + layer: torch.nn.Module,

37 + weight_loader: Callable,

38 + input_size_per_partition: int,

39 + output_partition_sizes: list[int],

40 + kernel: Optional[Callable] = None):

41 + super().__init__()

42 + self.transforms = transforms

43 + self.scales = {}

44 +

45 + if get_tensor_model_parallel_world_size() > 1:

46 + raise NotImplementedError("Online transforms with tensor "

47 + "parallelism is not supported")

48 +

49 + # Similar to row/col parallel params, but tensors are separate

50 + # to allow for loading with shared memory

51 + self.weight = SharedWeightParameter(weight_loader=weight_loader)

52 +

53 + # create shared partition data for each partition of the original weight

54 + input_size = input_size_per_partition

55 + for part_index, (_scheme_name, scheme,

56 + args) in self.transforms.items():

57 + output_size = output_partition_sizes[part_index]

58 + weight_size = self._get_weight_size(layer, args.location,

59 + input_size, output_size)

60 +

61 + data_key = self._get_data_key(scheme, weight_size)

62 + self.weight.add_partition(

63 + part_index,

64 + data_key,

65 + size=(weight_size, weight_size),

66 + dtype=scheme.precision,

67 + )

68 +

69 + # validate that shared tensors and schemes are correct

70 + self._validate_input_transforms()

71 +

72 + # select kernel based on transform schemes

73 + self.kernel = self._infer_kernel() if kernel is None else kernel

74 +

75 + def process_weights_after_loading(self):

76 + for part_id in self.weight.partitions:

77 + data = self.weight.partitions[part_id].data

78 +

79 + # required by torch.compile

80 + self.weight.process_weights_after_loading()

81 +

82 + # precompute scale as a runtime multiply, not division

83 + # do not fold into weight in order to utilize FWHT

84 + self.scales[part_id] = 1 / math.sqrt(data.size(0))

85 +

86 + # FUTURE: avoid runtime tranpose by processing weights

87 + # prior to apply

88 +

89 + def forward(self, value: Tensor, part_id: int = 0) -> Tensor:

90 + if part_id not in self.weight.partitions:

91 + return value

92 +

93 + weight = self.weight.partitions[part_id]

94 + weight = weight if self.transforms[

95 + part_id].args.inverse else weight.T # linear := x(W.T)

96 + scale = self.scales[part_id]

97 + return self.kernel(self, value.to(weight.dtype), weight, None).to(

98 + value.dtype) * scale

99 +

100 + def _get_data_key(self, scheme: TransformScheme,

101 + weight_size: int) -> Hashable:

102 + return (id(scheme), weight_size)

103 +

104 + def _get_weight_size(self, layer: torch.nn.Module,

105 + location: TransformLocation, input_size: int,

106 + output_size: int) -> int:

107 + if isinstance(layer, LinearBase):

108 + if location == TransformLocation.INPUT:

109 + return input_size

110 +

111 + elif location == TransformLocation.OUTPUT:

112 + return output_size

113 +

114 + elif isinstance(layer, VocabParallelEmbedding):

115 + if location == TransformLocation.INPUT:

116 + return output_size

117 +

118 + elif location == TransformLocation.OUTPUT:

119 + return input_size

120 +

121 + raise ValueError()

122 +

123 + def _validate_input_transforms(self):

124 + assert len(self.transforms) > 0

125 + location = list(self.transforms.values())[0].args.location

126 +

127 + if location == TransformLocation.INPUT:

128 + first_data = self.weight.partitions[0].data

129 + for partition in self.weight.partitions.values():

130 + if partition.data.data_ptr() != first_data.data_ptr():

131 + raise ValueError("")

132 +

133 + def _infer_kernel(self) -> Callable:

134 + # TODO (@ksayers): use fwht, hadacore

135 + return dispatch_unquantized_gemm()

vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py new

+21

@@ -0,0 +1,21 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Optional

4 +

5 +import torch

6 +

7 +from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear import ( # noqa: E501

8 + CompressedTensorsLinearTransformMethod)

9 +

10 +

11 +# Because qutlass fuses hadamard with quantization, it cannot automatically be

12 +# composed with kernels in the way CompressedTensorsLinearTransformMethod does.

13 +# Therefore, a separate scheme must be created for each quantized dtype

14 +class QutlassLinearMethodNvFP4(CompressedTensorsLinearTransformMethod):

15 +

16 + def apply(self,

17 + layer: torch.nn.Module,

18 + x: torch.Tensor,

19 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:

20 + # fused hadamard quant linear method

21 + raise NotImplementedError()

vllm/model_executor/layers/quantization/compressed_tensors/transform/utils.py new

+13

@@ -0,0 +1,13 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import NamedTuple

4 +

5 +from compressed_tensors.transform import TransformArgs, TransformScheme

6 +

7 +__all__ = ["TransformTuple"]

8 +

9 +

10 +class TransformTuple(NamedTuple):

11 + scheme_name: str

12 + scheme: TransformScheme

13 + args: TransformArgs

vllm/model_executor/parameter.py

+154 -14

@@ -1,13 +1,16 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +from collections.abc import Hashable

5 from fractions import Fraction

6 from typing import Callable, Optional, Union

7 +from weakref import WeakValueDictionary

9 import torch

10 from torch.nn import Parameter

10 -from vllm.distributed import get_tensor_model_parallel_rank

12 +from vllm.distributed import (get_tensor_model_parallel_rank,

13 + get_tensor_model_parallel_world_size)

14 from vllm.logger import init_logger

15 from vllm.model_executor.utils import _make_synced_weight_loader

@@ -27,7 +30,7 @@ class BasevLLMParameter(Parameter):

30 into the parameter when the provided weight loader is called.

31 """

30 - def __new__(cls, data: torch.Tensor, **kwargs):

33 + def __new__(cls, data: Optional[torch.Tensor], **kwargs):

35 return super().__new__(cls, data=data, requires_grad=False)

@@ -81,6 +84,17 @@ class BasevLLMParameter(Parameter):

84 def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):

85 self._assert_and_load(loaded_weight)

87 + def _shard_id_as_int(self, shard_id: Union[str, int]) -> int:

88 + if isinstance(shard_id, int):

89 + return shard_id

90 +

91 + # if not int, assume shard_id for qkv

92 + # map to int and return

93 + qkv_idxs = {"q": 0, "k": 1, "v": 2}

94 + assert isinstance(shard_id, str)

95 + assert shard_id in qkv_idxs

96 + return qkv_idxs[shard_id]

97 +

99 class _ColumnvLLMParameter(BasevLLMParameter):

100 """

@@ -113,6 +127,7 @@ class _ColumnvLLMParameter(BasevLLMParameter):

127

128 shard_offset = kwargs.get("shard_offset")

129 shard_size = kwargs.get("shard_size")

130 + # TODO: move these to PackedColumnParameter and PackedvLLMParameter

131 if isinstance(

132 self,

133 (PackedColumnParameter,

@@ -137,6 +152,7 @@ class _ColumnvLLMParameter(BasevLLMParameter):

152 shard_id = kwargs.get("shard_id")

153 num_heads = kwargs.get("num_heads")

154

155 + # TODO: move these to PackedColumnParameter and PackedvLLMParameter

156 if isinstance(

157 self,

158 (PackedColumnParameter,

@@ -224,19 +240,8 @@ class PerTensorScaleParameter(BasevLLMParameter):

240 """

241

242 def __init__(self, **kwargs):

227 - self.qkv_idxs = {"q": 0, "k": 1, "v": 2}

243 super().__init__(**kwargs)

244

230 - def _shard_id_as_int(self, shard_id: Union[str, int]) -> int:

231 - if isinstance(shard_id, int):

232 - return shard_id

233 -

234 - # if not int, assume shard_id for qkv

235 - # map to int and return

236 - assert isinstance(shard_id, str)

237 - assert shard_id in self.qkv_idxs

238 - return self.qkv_idxs[shard_id]

239 -

245 # For row parallel layers, no sharding needed

246 # load weight into parameter as is

247 def load_row_parallel_weight(self, *args, **kwargs):

@@ -373,6 +378,141 @@ class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):

378 pass

379

380

381 +class SharedWeightParameter(BasevLLMParameter):

382 + """

383 + Parameter for weights with many shared tensors across a model

384 +

385 + For example, when applying transforms to the "gate" and "up" partitions of

386 + `MergedColumnParallelLinear`, the transform weights must stay separate

387 + tensors in order to allow for tensor memory sharing between layers.

388 + """

389 + # global registry for sharing tensors based on passed `data_key`

390 + # this dict holds weaksrefs to avoid memory leak after model cleanup

391 + tensors_registry: WeakValueDictionary = WeakValueDictionary()

392 +

393 + # local container for strong references to shared tensors

394 + # this set compensates for the fact that torch.nn.Parameter

395 + # and Parameter subclasses do not hold reliable references to tensors

396 + local_tensors: set[torch.Tensor]

397 +

398 + # dictionary mapping partition indices to associated parameters

399 + partitions: dict[int, Union[ModelWeightParameter, Parameter]]

400 +

401 + def __new__(cls, **kwargs):

402 + return super().__new__(cls, data=None, **kwargs)

403 +

404 + def __init__(self, input_dim: int = 1, output_dim: int = 0, **kwargs):

405 + weight_loader: Callable = kwargs.get(

406 + "weight_loader") # type: ignore[assignment]

407 + super().__init__(data=None, weight_loader=weight_loader)

408 +

409 + self.local_tensors = set()

410 + self.partitions = {}

411 + self.kwargs = {

412 + "input_dim": input_dim,

413 + "output_dim": output_dim,

414 + "weight_loader": self._fake_weight_loader

415 + }

416 +

417 + self.tp_rank = get_tensor_model_parallel_rank()

418 + self.tp_size = get_tensor_model_parallel_world_size()

419 +

420 + if self.tp_size > 1:

421 + raise NotImplementedError(f"{self.__class__.__name__} does not "

422 + "currently support tensor parallelism")

423 +

424 + def add_partition(self, index: int, data_key: Hashable, *args, **kwargs):

425 + """

426 + Add a partition to the weight parameter. Partitions whose `data_key`

427 + is the same will share tensor data

428 +

429 + :param index: index of partition to add

430 + :param data_key: hashable key used to key shared tensors

431 + :param *args: arguments for `torch.empty`

432 + :param **kwargs: keyword arguments for `torch.empty`

433 + """

434 + # load (shared) tensor using `data_key`

435 + if data_key not in self.tensors_registry:

436 + data = torch.empty(*args, **kwargs)

437 + self.tensors_registry[data_key] = data

438 + else:

439 + data = self.tensors_registry[data_key]

440 +

441 + # create associated model parameter

442 + self.partitions[index] = ModelWeightParameter(

443 + data=data, **self.kwargs) # type: ignore[arg-type]

444 +

445 + # hold local reference, since ModelWeightParameter does not

446 + # see https://github.com/pytorch/pytorch/issues/75932

447 + self.local_tensors.add(data)

448 +

449 + def load_column_parallel_weight(self, loaded_weight: torch.Tensor):

450 + assert len(self.partitions) == 1 and 0 in self.partitions

451 + partition = self.partitions[0]

452 +

453 + ModelWeightParameter.load_column_parallel_weight(

454 + partition, loaded_weight)

455 +

456 + def load_row_parallel_weight(self, loaded_weight: torch.Tensor):

457 + assert len(self.partitions) == 1 and 0 in self.partitions

458 + partition = self.partitions[0]

459 +

460 + ModelWeightParameter.load_row_parallel_weight(partition, loaded_weight)

461 +

462 + def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):

463 + partition_id = kwargs.pop("shard_id")

464 + partition_id = self._shard_id_as_int(partition_id)

465 + partition = self.partitions[partition_id]

466 +

467 + input_dim = self.kwargs.get("input_dim")

468 + shard_size = partition.data.size(input_dim) // self.tp_size

469 + shard_offset = self.tp_rank * shard_size

470 +

471 + ModelWeightParameter.load_merged_column_weight(

472 + partition,

473 + loaded_weight,

474 + shard_offset=shard_offset,

475 + shard_size=shard_size)

476 +

477 + def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):

478 + partition_id = self._shard_id_as_int(kwargs.pop("shard_id"))

479 + partition = self.partitions[partition_id]

480 +

481 + input_dim = self.kwargs.get("input_dim")

482 + shard_size = partition.data.size(input_dim) // self.tp_size

483 + shard_offset = self.tp_rank * shard_size

484 + shard_id = "q" # fake first partition

485 + num_heads = kwargs.get("num_heads")

486 +

487 + ModelWeightParameter.load_qkv_weight(

488 + partition,

489 + loaded_weight,

490 + shard_offset=shard_offset,

491 + shard_size=shard_size,

492 + shard_id=shard_id,

493 + num_heads=num_heads,

494 + )

495 +

496 + def process_weights_after_loading(self):

497 + for key in self.partitions:

498 + self.partitions[key] = torch.nn.Parameter(

499 + data=self.partitions[key].data, requires_grad=False)

500 +

501 + @property

502 + def data(self):

503 + raise ValueError("Accessing `data` of a "

504 + "`PartitionedModelWeightParameter` is not allowed. "

505 + "Instead, use `get_partition` to get the weight of "

506 + "the particular partition you want to access")

507 +

508 + def _fake_weight_loader(self, param: BasevLLMParameter,

509 + loaded_weight: torch.Tensor,

510 + loaded_weight_shard_id: Optional[Union[str, int]]):

511 + raise ValueError("When loading partition weights of "

512 + f"{self.__class__.__name__}, use methods provided by "

513 + f"{self.__class__.__name__}, not partition loader")

514 +

515 +

516 def permute_param_layout_(param: BasevLLMParameter, input_dim: int,

517 output_dim: int, **kwargs) -> BasevLLMParameter:

518 """

@@ -456,4 +596,4 @@ def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor,

596 shard_offset=shard_offset,

597 bitblas_tile_size=bitblas_tile_size)

598

459 - return shard_size, shard_offset

\ No newline at end of file

599 + return shard_size, shard_offset