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, 22:15 UTC
[BugFix] Make FP8 Linear compatible with torch.compile (#13918)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Feb 26, 2025 at 13:48 UTC
b382a7f28f739f3b120e5495fd029089d0399428
2 files changed
+21
-4
vllm/model_executor/layers/quantization/fp8.py
+1
-4
@@ -369,12 +369,9 @@ class Fp8LinearMethod(LinearMethodBase):
369
size_k=layer.input_size_per_partition,
370
bias=bias)
371
372
- # Note: lazy import to avoid triton import error.
373
- from vllm.model_executor.layers.quantization.utils.fp8_utils import (
374
- apply_w8a8_block_fp8_linear)
372
if self.block_quant:
373
assert self.quant_config.weight_block_size is not None
377
- return apply_w8a8_block_fp8_linear(
374
+ return torch.ops.vllm.apply_w8a8_block_fp8_linear(
375
input=x,
376
weight=layer.weight,
377
block_size=self.quant_config.weight_block_size,
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+20
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
18
CUTLASS_BLOCK_FP8_SUPPORTED, CUTLASS_FP8_SUPPORTED, apply_fp8_linear)
19
from vllm.platforms import current_platform
20
+from vllm.utils import direct_register_custom_op
21
22
logger = init_logger(__name__)
23
@@ -81,6 +82,25 @@ def apply_w8a8_block_fp8_linear(
82
return output.to(dtype=input.dtype).view(*output_shape)
83
84
85
+def apply_w8a8_block_fp8_linear_fake(
86
+ input: torch.Tensor,
87
+ weight: torch.Tensor,
88
+ block_size: List[int],
89
+ weight_scale: torch.Tensor,
90
+ input_scale: Optional[torch.Tensor] = None,
91
+) -> torch.Tensor:
92
+ output_shape = [*input.shape[:-1], weight.shape[0]]
93
+ return torch.empty(output_shape, dtype=input.dtype, device=input.device)
94
+
95
+
96
+direct_register_custom_op(
97
+ op_name="apply_w8a8_block_fp8_linear",
98
+ op_func=apply_w8a8_block_fp8_linear,
99
+ mutates_args=[],
100
+ fake_impl=apply_w8a8_block_fp8_linear_fake,
101
+)
102
+
103
+
104
# Unify the interface between `apply_w8a8_block_fp8_linear` and
105
# `apply_fp8_linear`
106
# NOTE(lucas): this is quite messy, we should think through this more formally