Skip to content

vllm.model_executor.layers.quantization.utils.w8a8_utils

CUTLASS_BLOCK_FP8_SUPPORTED module-attribute 

CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported()

CUTLASS_FP8_SUPPORTED module-attribute 

CUTLASS_FP8_SUPPORTED = cutlass_fp8_supported()

TORCH_DEVICE_IDENTITY module-attribute 

TORCH_DEVICE_IDENTITY = None

USE_ROWWISE_TORCH_SCALED_MM module-attribute 

USE_ROWWISE_TORCH_SCALED_MM = (
    is_rocm()
    and __version__[0:3] >= "2.7"
    and has_device_capability(94)
)

Fp8LinearOp

This class executes a FP8 linear layer using cutlass if supported and torch.scaled_mm otherwise. It needs to be a class instead of a method so that config can be read in the init method, as reading config is not allowed inside forward.

Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
class Fp8LinearOp:
    """
    This class executes a FP8 linear layer using cutlass if supported and
    torch.scaled_mm otherwise.
    It needs to be a class instead of a method so that config can be read
    in the __init__ method, as reading config is not allowed inside forward.
    """

    def __init__(self,
                 cutlass_fp8_supported: bool = cutlass_fp8_supported(),
                 use_per_token_if_dynamic: bool = False,
                 pad_output: Optional[bool] = None):
        self.cutlass_fp8_supported = cutlass_fp8_supported
        self.use_per_token_if_dynamic = use_per_token_if_dynamic

        # Note: we pad the input because torch._scaled_mm is more performant
        # for matrices with batch dimension > 16.
        # This could change in the future.
        # We also don't pad when using torch.compile,
        # as it breaks with dynamic shapes.
        if pad_output is None:
            config = get_current_vllm_config().compilation_config
            pad_output = config.level < CompilationLevel.PIECEWISE
        self.output_padding = 17 if (
            pad_output and not current_platform.is_rocm()) else None

    def apply(
        self,
        input: torch.Tensor,
        weight: torch.Tensor,
        weight_scale: torch.Tensor,
        out_dtype: Optional[torch.dtype] = None,
        input_scale: Optional[torch.Tensor] = None,
        input_scale_ub: Optional[torch.Tensor] = None,
        bias: Optional[torch.Tensor] = None,
        # TODO(luka) remove this parameter in favor of __init__
        use_per_token_if_dynamic: Optional[bool] = None
    ) -> torch.Tensor:
        # ops.scaled_fp8_quant supports both dynamic and static quant.
        #   If dynamic, layer.input_scale is None and x_scale computed from x.
        #   If static, layer.input_scale is scalar and x_scale is input_scale.

        # View input as 2D matrix for fp8 methods
        input_2d = input.view(-1, input.shape[-1])
        output_shape = [*input.shape[:-1], weight.shape[1]]

        # TODO(luka) this is here because currently MLA only decides this
        #  during the forward method instead of in __init__.
        if use_per_token_if_dynamic is None:
            use_per_token_if_dynamic = self.use_per_token_if_dynamic

        if out_dtype is None:
            out_dtype = input.dtype

        # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
        if self.cutlass_fp8_supported:
            assert input.dtype != current_platform.fp8_dtype(
            ), "FP8 input to cutlass is not currently implemented"
            qinput, x_scale = ops.scaled_fp8_quant(
                input_2d,
                input_scale,
                scale_ub=input_scale_ub,
                use_per_token_if_dynamic=use_per_token_if_dynamic)
        else:
            if input.dtype != current_platform.fp8_dtype():
                # Maybe apply padding to output, see comment in __init__
                qinput, x_scale = ops.scaled_fp8_quant(
                    input_2d,
                    input_scale,
                    num_token_padding=self.output_padding,
                    use_per_token_if_dynamic=use_per_token_if_dynamic)
            else:
                qinput, x_scale = input_2d, input_scale

        per_tensor_weights = (weight_scale.numel() == 1)
        per_tensor_activations = (x_scale.numel() == 1)

        w8a8_scaled_mm_func = dispatch_w8a8_scaled_mm(
            self.cutlass_fp8_supported, per_tensor_weights,
            per_tensor_activations, use_per_token_if_dynamic)

        return w8a8_scaled_mm_func(qinput=qinput,
                                   weight=weight,
                                   out_dtype=out_dtype,
                                   scale_a=x_scale,
                                   scale_b=weight_scale,
                                   bias=bias,
                                   input_2d=input_2d,
                                   output_shape=output_shape)

cutlass_fp8_supported instance-attribute 

cutlass_fp8_supported = cutlass_fp8_supported

output_padding instance-attribute 

output_padding = (
    17 if pad_output and not is_rocm() else None
)

use_per_token_if_dynamic instance-attribute 

use_per_token_if_dynamic = use_per_token_if_dynamic

__init__ 

__init__(
    cutlass_fp8_supported: bool = cutlass_fp8_supported(),
    use_per_token_if_dynamic: bool = False,
    pad_output: Optional[bool] = None,
)
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
def __init__(self,
             cutlass_fp8_supported: bool = cutlass_fp8_supported(),
             use_per_token_if_dynamic: bool = False,
             pad_output: Optional[bool] = None):
    self.cutlass_fp8_supported = cutlass_fp8_supported
    self.use_per_token_if_dynamic = use_per_token_if_dynamic

    # Note: we pad the input because torch._scaled_mm is more performant
    # for matrices with batch dimension > 16.
    # This could change in the future.
    # We also don't pad when using torch.compile,
    # as it breaks with dynamic shapes.
    if pad_output is None:
        config = get_current_vllm_config().compilation_config
        pad_output = config.level < CompilationLevel.PIECEWISE
    self.output_padding = 17 if (
        pad_output and not current_platform.is_rocm()) else None

apply 

apply(
    input: Tensor,
    weight: Tensor,
    weight_scale: Tensor,
    out_dtype: Optional[dtype] = None,
    input_scale: Optional[Tensor] = None,
    input_scale_ub: Optional[Tensor] = None,
    bias: Optional[Tensor] = None,
    use_per_token_if_dynamic: Optional[bool] = None,
) -> Tensor
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
def apply(
    self,
    input: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    out_dtype: Optional[torch.dtype] = None,
    input_scale: Optional[torch.Tensor] = None,
    input_scale_ub: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    # TODO(luka) remove this parameter in favor of __init__
    use_per_token_if_dynamic: Optional[bool] = None
) -> torch.Tensor:
    # ops.scaled_fp8_quant supports both dynamic and static quant.
    #   If dynamic, layer.input_scale is None and x_scale computed from x.
    #   If static, layer.input_scale is scalar and x_scale is input_scale.

    # View input as 2D matrix for fp8 methods
    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[1]]

    # TODO(luka) this is here because currently MLA only decides this
    #  during the forward method instead of in __init__.
    if use_per_token_if_dynamic is None:
        use_per_token_if_dynamic = self.use_per_token_if_dynamic

    if out_dtype is None:
        out_dtype = input.dtype

    # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
    if self.cutlass_fp8_supported:
        assert input.dtype != current_platform.fp8_dtype(
        ), "FP8 input to cutlass is not currently implemented"
        qinput, x_scale = ops.scaled_fp8_quant(
            input_2d,
            input_scale,
            scale_ub=input_scale_ub,
            use_per_token_if_dynamic=use_per_token_if_dynamic)
    else:
        if input.dtype != current_platform.fp8_dtype():
            # Maybe apply padding to output, see comment in __init__
            qinput, x_scale = ops.scaled_fp8_quant(
                input_2d,
                input_scale,
                num_token_padding=self.output_padding,
                use_per_token_if_dynamic=use_per_token_if_dynamic)
        else:
            qinput, x_scale = input_2d, input_scale

    per_tensor_weights = (weight_scale.numel() == 1)
    per_tensor_activations = (x_scale.numel() == 1)

    w8a8_scaled_mm_func = dispatch_w8a8_scaled_mm(
        self.cutlass_fp8_supported, per_tensor_weights,
        per_tensor_activations, use_per_token_if_dynamic)

    return w8a8_scaled_mm_func(qinput=qinput,
                               weight=weight,
                               out_dtype=out_dtype,
                               scale_a=x_scale,
                               scale_b=weight_scale,
                               bias=bias,
                               input_2d=input_2d,
                               output_shape=output_shape)

all_close_1d 

all_close_1d(x: Tensor) -> bool
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
78
79
80
def all_close_1d(x: torch.Tensor) -> bool:
    assert len(x.shape) == 1
    return all(torch.allclose(x[0], x[i]) for i in range(x.shape[0]))

convert_to_channelwise 

convert_to_channelwise(
    weight_scale: Tensor, logical_widths: list[int]
) -> tuple[Tensor, Tensor]
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def convert_to_channelwise(
        weight_scale: torch.Tensor,
        logical_widths: list[int]) -> tuple[torch.Tensor, torch.Tensor]:
    # Create channelwise buffer
    weight_scale_channel = torch.empty((sum(logical_widths), 1),
                                       dtype=torch.float32,
                                       device=weight_scale.device)

    # Expand each scale to match the size of each logical matrix.
    start = 0
    for idx, logical_width in enumerate(logical_widths):
        end = start + logical_width
        weight_scale_channel[start:end, :] = weight_scale[idx]
        start = end

    return weight_scale_channel

cutlass_block_fp8_supported 

cutlass_block_fp8_supported() -> bool
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
46
47
48
49
50
51
52
53
def cutlass_block_fp8_supported() -> bool:
    if not current_platform.is_cuda():
        return False

    capability_tuple = current_platform.get_device_capability()
    capability = -1 if capability_tuple is None else capability_tuple.to_int()

    return ops.cutlass_scaled_mm_supports_block_fp8(capability)

cutlass_fp8_supported 

cutlass_fp8_supported() -> bool
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
36
37
38
39
40
41
42
43
def cutlass_fp8_supported() -> bool:
    if not current_platform.is_cuda():
        return False

    capability_tuple = current_platform.get_device_capability()
    capability = -1 if capability_tuple is None else capability_tuple.to_int()

    return ops.cutlass_scaled_mm_supports_fp8(capability)

cutlass_group_gemm_supported 

cutlass_group_gemm_supported() -> bool
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
56
57
58
59
60
61
62
63
def cutlass_group_gemm_supported() -> bool:
    if not current_platform.is_cuda():
        return False

    capability_tuple = current_platform.get_device_capability()
    capability = -1 if capability_tuple is None else capability_tuple.to_int()

    return ops.cutlass_group_gemm_supported(capability)

cutlass_w8a8_scaled_mm 

cutlass_w8a8_scaled_mm(
    *,
    qinput: Tensor,
    weight: Tensor,
    out_dtype: dtype,
    scale_a: Tensor,
    scale_b: Tensor,
    bias: Tensor,
    output_shape: list,
    **kwargs,
) -> Tensor
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
def cutlass_w8a8_scaled_mm(*, qinput: torch.Tensor, weight: torch.Tensor,
                           out_dtype: torch.dtype, scale_a: torch.Tensor,
                           scale_b: torch.Tensor, bias: torch.Tensor,
                           output_shape: list, **kwargs) -> torch.Tensor:

    # Fused GEMM_DQ
    output = ops.cutlass_scaled_mm(qinput,
                                   weight,
                                   out_dtype=out_dtype,
                                   scale_a=scale_a,
                                   scale_b=scale_b,
                                   bias=bias)
    return output.view(*output_shape)

dispatch_w8a8_scaled_mm 

dispatch_w8a8_scaled_mm(
    cutlass_fp8_supported: bool,
    per_tensor_weights: bool,
    per_tensor_activations: bool,
    use_per_token_if_dynamic: Optional[bool],
) -> Callable[..., Tensor]
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
def dispatch_w8a8_scaled_mm(
        cutlass_fp8_supported: bool, per_tensor_weights: bool,
        per_tensor_activations: bool, use_per_token_if_dynamic: Optional[bool]
) -> Callable[..., torch.Tensor]:

    if cutlass_fp8_supported:
        return cutlass_w8a8_scaled_mm
    if per_tensor_weights and per_tensor_activations:
        if current_platform.is_rocm():
            return rocm_per_tensor_w8a8_scaled_mm
        return torch_per_tensor_w8a8_scaled_mm
    # torch.scaled_mm supports per tensor weights + activations only
    # so fallback to naive if per channel or per token
    if (use_per_token_if_dynamic and not per_tensor_weights
            and not per_tensor_activations and USE_ROWWISE_TORCH_SCALED_MM):
        return torch_per_token_w8a8_scaled_mm
    return torch_channelwise_w8a8_scaled_mm

maybe_create_device_identity 

maybe_create_device_identity()
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
130
131
132
133
134
def maybe_create_device_identity():
    # Allocate dummy ones tensor for torch._scaled_mm
    global TORCH_DEVICE_IDENTITY
    if TORCH_DEVICE_IDENTITY is None:
        TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32)

normalize_e4m3fn_to_e4m3fnuz 

normalize_e4m3fn_to_e4m3fnuz(
    weight: Tensor,
    weight_scale: Tensor,
    input_scale: Optional[Tensor] = None,
) -> tuple[Tensor, Tensor, Optional[Tensor]]
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
def normalize_e4m3fn_to_e4m3fnuz(
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None
) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
    assert weight.dtype == torch.float8_e4m3fn
    # The bits pattern 10000000(-128) represents zero in e4m3fn
    # but NaN in e4m3fnuz. So here we set it to 0.
    # https://onnx.ai/onnx/technical/float8.html
    weight_as_int8 = weight.view(torch.int8)
    ROCM_FP8_NAN_AS_INT = -128
    weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0
    weight = weight_as_int8.view(torch.float8_e4m3fnuz)

    # For the same bits representation, e4m3fnuz value is half of
    # the e4m3fn value, so we should double the scaling factor to
    # get the same dequantized value.
    # https://onnx.ai/onnx/technical/float8.html
    weight_scale = weight_scale * 2.0
    if input_scale is not None:
        input_scale = input_scale * 2.0
    return weight, weight_scale, input_scale

per_tensor_dequantize 

per_tensor_dequantize(
    tensor: Tensor, inv_scale: Union[float, Tensor]
) -> Tensor
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
70
71
72
73
74
75
def per_tensor_dequantize(
        tensor: torch.Tensor, inv_scale: Union[float,
                                               torch.Tensor]) -> torch.Tensor:
    fake_qweight = tensor.to(torch.float16)
    dq_weight = fake_qweight * inv_scale
    return dq_weight

requantize_with_max_scale 

requantize_with_max_scale(
    weight: Tensor,
    weight_scale: Tensor,
    logical_widths: list[int],
) -> tuple[Tensor, Tensor]
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def requantize_with_max_scale(
        weight: torch.Tensor, weight_scale: torch.Tensor,
        logical_widths: list[int]) -> tuple[torch.Tensor, torch.Tensor]:
    # Max scale to be used for requanitzation.
    max_w_scale = weight_scale.max()

    # QKV / MLP is fused in the on disk checkpoint if any of the
    # weight scales are still set to the default since we initialize
    # N weight scales for N shards but we only load 1 weight scale
    # from disk in this case. Skip requantization in this case (since)
    # we already are quantized with the single scale.
    # * Sample Model: nm-testing/Phi-3-mini-128k-instruct-FP8
    unfused_module_in_checkpoint = (weight_scale[-1]
                                    > torch.finfo(torch.float8_e4m3fn).min)

    # If unfused checkpoint, need requanize with the single scale.
    if unfused_module_in_checkpoint:
        start = 0
        for idx, logical_width in enumerate(logical_widths):
            end = start + logical_width
            weight_dq = per_tensor_dequantize(weight[start:end, :],
                                              weight_scale[idx])
            weight[start:end, :], _ = ops.scaled_fp8_quant(
                weight_dq, max_w_scale)
            start = end

    return max_w_scale, weight

rocm_per_tensor_w8a8_scaled_mm 

rocm_per_tensor_w8a8_scaled_mm(
    *,
    qinput: Tensor,
    weight: Tensor,
    out_dtype: dtype,
    scale_a: Tensor,
    scale_b: Tensor,
    bias: Tensor,
    input_2d: Tensor,
    output_shape: list,
) -> Tensor
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def rocm_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
                                   weight: torch.Tensor,
                                   out_dtype: torch.dtype,
                                   scale_a: torch.Tensor,
                                   scale_b: torch.Tensor, bias: torch.Tensor,
                                   input_2d: torch.Tensor,
                                   output_shape: list) -> torch.Tensor:
    from vllm.platforms.rocm import on_mi3xx
    if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi3xx(
    ) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0:
        output = ops.wvSplitKQ(weight.t(), qinput, out_dtype, scale_a, scale_b,
                               current_platform.get_cu_count())
    else:
        output = torch._scaled_mm(qinput,
                                  weight,
                                  out_dtype=out_dtype,
                                  scale_a=scale_a,
                                  scale_b=scale_b,
                                  bias=bias)

    return torch.narrow(output, 0, 0, input_2d.shape[0]).view(*output_shape)

sparse_cutlass_supported 

sparse_cutlass_supported() -> bool
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
26
27
28
29
30
31
32
33
def sparse_cutlass_supported() -> bool:
    if not current_platform.is_cuda():
        return False

    capability_tuple = current_platform.get_device_capability()
    capability = -1 if capability_tuple is None else capability_tuple.to_int()

    return ops.cutlass_sparse_scaled_mm_supported(capability)

torch_channelwise_w8a8_scaled_mm 

torch_channelwise_w8a8_scaled_mm(
    *,
    qinput: Tensor,
    weight: Tensor,
    out_dtype: dtype,
    scale_a: Tensor,
    scale_b: Tensor,
    bias: Tensor,
    input_2d: Tensor,
    output_shape: list,
    **kwargs,
) -> Tensor
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def torch_channelwise_w8a8_scaled_mm(*, qinput: torch.Tensor,
                                     weight: torch.Tensor,
                                     out_dtype: torch.dtype,
                                     scale_a: torch.Tensor,
                                     scale_b: torch.Tensor, bias: torch.Tensor,
                                     input_2d: torch.Tensor,
                                     output_shape: list,
                                     **kwargs) -> torch.Tensor:
    # Use unfused DQ due to limitations with scaled_mm

    # Symmetric quantized GEMM by definition computes the following:
    #   C = (s_x * X) (s_w * W) + bias
    # This is equivalent to dequantizing the weights and activations
    # before applying a GEMM.
    #
    # In order to compute quantized operands, a quantized kernel
    # will rewrite the above like so:
    #   C = s_w * s_x * (X * W) + bias
    #
    # For the scaled_mm fallback case, we break this down, since it
    # does not support s_w being a vector.

    # GEMM
    # This computes C = (X * W).
    # Output in fp32 to allow subsequent ops to happen in-place
    output = torch._scaled_mm(qinput,
                              weight,
                              scale_a=TORCH_DEVICE_IDENTITY,
                              scale_b=TORCH_DEVICE_IDENTITY,
                              out_dtype=torch.float32)
    # A fix for discrepancy in scaled_mm which returns tuple
    # for torch < 2.5 and a single value in torch >= 2.5
    if type(output) is tuple and len(output) == 2:
        output = output[0]
    # Unpad (undo num_token_padding)
    output = torch.narrow(output, 0, 0, input_2d.shape[0])
    x_scale = torch.narrow(scale_a, 0, 0, input_2d.shape[0])

    # DQ
    # C = sw * sx * (X * W) + bias
    output = output * x_scale * scale_b.t()
    if bias is not None:
        output = output + bias
    return output.to(out_dtype).view(*output_shape)

torch_per_tensor_w8a8_scaled_mm 

torch_per_tensor_w8a8_scaled_mm(
    *,
    qinput: Tensor,
    weight: Tensor,
    out_dtype: dtype,
    scale_a: Tensor,
    scale_b: Tensor,
    bias: Tensor,
    input_2d: Tensor,
    output_shape: list,
) -> Tensor
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def torch_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
                                    weight: torch.Tensor,
                                    out_dtype: torch.dtype,
                                    scale_a: torch.Tensor,
                                    scale_b: torch.Tensor, bias: torch.Tensor,
                                    input_2d: torch.Tensor,
                                    output_shape: list) -> torch.Tensor:
    output = torch._scaled_mm(qinput,
                              weight,
                              out_dtype=out_dtype,
                              scale_a=scale_a,
                              scale_b=scale_b,
                              bias=bias)
    # A fix for discrepancy in scaled_mm which returns tuple
    # for torch < 2.5 and a single value in torch >= 2.5
    if type(output) is tuple and len(output) == 2:
        output = output[0]

    return torch.narrow(output, 0, 0, input_2d.shape[0]).view(*output_shape)

torch_per_token_w8a8_scaled_mm 

torch_per_token_w8a8_scaled_mm(
    *,
    qinput: Tensor,
    weight: Tensor,
    out_dtype: dtype,
    scale_a: Tensor,
    scale_b: Tensor,
    bias: Tensor,
    input_2d: Tensor,
    output_shape: list,
) -> Tensor
Source code in vllm/model_executor/layers/quantization/utils/w8a8_utils.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
def torch_per_token_w8a8_scaled_mm(*, qinput: torch.Tensor,
                                   weight: torch.Tensor,
                                   out_dtype: torch.dtype,
                                   scale_a: torch.Tensor,
                                   scale_b: torch.Tensor, bias: torch.Tensor,
                                   input_2d: torch.Tensor,
                                   output_shape: list) -> torch.Tensor:
    # Note: Callers of this function should check USE_ROWWISE_TORCH_SCALED_MM
    #  when using it.
    #  For now it has only been validated on ROCm platform.
    #  fp8 rowwise scaling in torch._scaled_mm is introduced in
    #  https://github.com/pytorch/pytorch/pull/144432 using
    #  hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above.
    #
    #  For CUDA platform please validate if the torch._scaled_mm supports
    #  rowwise scaled GEMM before using it

    # Fused GEMM_DQ Rowwise GEMM
    output = torch._scaled_mm(qinput,
                              weight,
                              out_dtype=out_dtype,
                              scale_a=scale_a,
                              scale_b=scale_b.t(),
                              bias=bias)

    output = torch.narrow(output, 0, 0, input_2d.shape[0])
    output = output.view(*output_shape)
    return output