@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4d0ec37267

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 15, 17:15 UTC

[Quantization][FP8] Adding support for fp8 gemm layer input in fp8 (#14578)

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Gregory Shtrasberg committed Mar 27, 2025 at 22:58 UTC 4d0ec37267afaf988e32174ebc31f24268076491

5 files changed +41 -9

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

@@ -23,6 +23,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

24 def __init__(self, strategy: str, is_static_input_scheme: bool):

25 self.strategy = strategy

26 + self.out_dtype = torch.get_default_dtype()

27 self.is_static_input_scheme = is_static_input_scheme

28 self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)

@@ -143,5 +144,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

144 return self.fp8_linear.apply(input=x,

145 weight=layer.weight,

146 weight_scale=layer.weight_scale,

147 + out_dtype=self.out_dtype,

148 input_scale=layer.input_scale,

149 bias=bias)

vllm/model_executor/layers/quantization/fbgemm_fp8.py

@@ -73,6 +73,7 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):

73 def __init__(self, quant_config: FBGEMMFp8Config):

74 self.quant_config = quant_config

75 self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)

76 + self.out_dtype = torch.get_default_dtype()

78 def create_weights(

79 self,

@@ -161,6 +162,7 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):

162 return self.fp8_linear.apply(input=x,

163 weight=layer.weight,

164 weight_scale=layer.weight_scale,

165 + out_dtype=self.out_dtype,

166 input_scale=None,

167 input_scale_ub=layer.input_scale_ub,

168 bias=bias)

vllm/model_executor/layers/quantization/fp8.py

+17

@@ -116,6 +116,21 @@ class Fp8Config(QuantizationConfig):

116 return Fp8KVCacheMethod(self)

117 return None

118

119 + def get_cache_scale(self, name: str) -> Optional[str]:

120 + """

121 + Check whether the param name matches the format for k/v cache scales

122 + in compressed-tensors. If this is the case, return its equivalent

123 + param name expected by vLLM

124 +

125 + :param name: param name

126 + :return: matching param name for KV cache scale in vLLM

127 + """

128 + if name.endswith(".output_scale") and ".k_proj" in name:

129 + return name.replace(".k_proj.output_scale", ".attn.k_scale")

130 + if name.endswith(".output_scale") and ".v_proj" in name:

131 + return name.replace(".v_proj.output_scale", ".attn.v_scale")

132 + return None

133 +

134

135 class Fp8LinearMethod(LinearMethodBase):

136 """Linear method for FP8.

@@ -138,6 +153,7 @@ class Fp8LinearMethod(LinearMethodBase):

153 def __init__(self, quant_config: Fp8Config):

154 self.quant_config = quant_config

155 self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()

156 + self.out_dtype = torch.get_default_dtype()

157

158 # For GPUs that lack FP8 hardware support, we can leverage the Marlin

159 # kernel for fast weight-only FP8 quantization

@@ -386,6 +402,7 @@ class Fp8LinearMethod(LinearMethodBase):

402 return self.fp8_linear.apply(input=x,

403 weight=layer.weight,

404 weight_scale=layer.weight_scale,

405 + out_dtype=self.out_dtype,

406 input_scale=layer.input_scale,

407 bias=bias)

408

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py

@@ -22,6 +22,7 @@ class QuarkW8A8Fp8(QuarkScheme):

22 self.qscheme = qscheme

23 self.is_static_input_scheme = is_static_input_scheme

24 self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)

25 + self.out_dtype = torch.get_default_dtype()

27 @classmethod

28 def get_min_capability(cls) -> int:

@@ -134,5 +135,6 @@ class QuarkW8A8Fp8(QuarkScheme):

135 return self.fp8_linear.apply(input=x,

136 weight=layer.weight,

137 weight_scale=layer.weight_scale,

138 + out_dtype=self.out_dtype,

139 input_scale=layer.input_scale,

140 bias=bias)

vllm/model_executor/layers/quantization/utils/w8a8_utils.py

+18 -9

@@ -163,6 +163,7 @@ class Fp8LinearOp:

163 input: torch.Tensor,

164 weight: torch.Tensor,

165 weight_scale: torch.Tensor,

166 + out_dtype: Optional[torch.dtype] = None,

167 input_scale: Optional[torch.Tensor] = None,

168 input_scale_ub: Optional[torch.Tensor] = None,

169 bias: Optional[torch.Tensor] = None,

@@ -182,8 +183,13 @@ class Fp8LinearOp:

183 if use_per_token_if_dynamic is None:

184 use_per_token_if_dynamic = self.use_per_token_if_dynamic

185

186 + if out_dtype is None:

187 + out_dtype = input.dtype

188 +

189 # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A

190 if self.cutlass_fp8_supported:

191 + assert input.dtype != current_platform.fp8_dtype(

192 + ), "FP8 input to cutlass is not currently implemented"

193 qinput, x_scale = ops.scaled_fp8_quant(

194 input_2d,

195 input_scale,

@@ -193,7 +199,7 @@ class Fp8LinearOp:

199 # Fused GEMM_DQ

200 output = ops.cutlass_scaled_mm(qinput,

201 weight,

196 - out_dtype=input.dtype,

202 + out_dtype=out_dtype,

203 scale_a=x_scale,

204 scale_b=weight_scale,

205 bias=bias)

@@ -202,12 +208,15 @@ class Fp8LinearOp:

208 # torch.scaled_mm supports per tensor weights + activations only

209 # so fallback to naive if per channel or per token

210 else:

205 - # Maybe apply padding to output, see comment in __init__

206 - qinput, x_scale = ops.scaled_fp8_quant(

207 - input_2d,

208 - input_scale,

209 - num_token_padding=self.output_padding,

210 - use_per_token_if_dynamic=use_per_token_if_dynamic)

211 + if input.dtype != current_platform.fp8_dtype():

212 + # Maybe apply padding to output, see comment in __init__

213 + qinput, x_scale = ops.scaled_fp8_quant(

214 + input_2d,

215 + input_scale,

216 + num_token_padding=self.output_padding,

217 + use_per_token_if_dynamic=use_per_token_if_dynamic)

218 + else:

219 + qinput, x_scale = input_2d, input_scale

220

221 per_tensor_weights = (weight_scale.numel() == 1)

222 per_tensor_activations = (x_scale.numel() == 1)

@@ -216,7 +225,7 @@ class Fp8LinearOp:

225 # Fused GEMM_DQ

226 output = torch._scaled_mm(qinput,

227 weight,

219 - out_dtype=input.dtype,

228 + out_dtype=out_dtype,

229 scale_a=x_scale,

230 scale_b=weight_scale,

231 bias=bias)

@@ -240,7 +249,7 @@ class Fp8LinearOp:

249 # Fused GEMM_DQ Rowwise GEMM

250 output = torch._scaled_mm(qinput,

251 weight,

243 - out_dtype=input.dtype,

252 + out_dtype=out_dtype,

253 scale_a=x_scale,

254 scale_b=weight_scale.t(),

255 bias=bias)