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
[ROCm] Apply FP8 weights padding to values not divisible by 512 bytes on ROCm (#13231)
Gregory Shtrasberg committed
Feb 22, 2025 at 08:54 UTC
c904fdddf68f8d2f622f7505763156d1aaa0e36b
3 files changed
+20
-1
vllm/envs.py
+4
@@ -74,6 +74,7 @@ if TYPE_CHECKING:
74
VLLM_SKIP_P2P_CHECK: bool = False
75
VLLM_DISABLED_KERNELS: List[str] = []
76
VLLM_USE_V1: bool = False
77
+ VLLM_ROCM_FP8_PADDING: bool = True
78
VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
79
VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
80
VLLM_DISABLE_COMPILE_CACHE: bool = False
@@ -507,6 +508,9 @@ environment_variables: Dict[str, Callable[[], Any]] = {
508
"VLLM_USE_V1":
509
lambda: bool(int(os.getenv("VLLM_USE_V1", "0"))),
510
511
+ # Pad the fp8 weights to 256 bytes for ROCm
512
+ "VLLM_ROCM_FP8_PADDING":
513
+ lambda: bool(int(os.getenv("VLLM_ROCM_FP8_PADDING", "1"))),
514
# Divisor for dynamic key scale factor calculation for FP8 KV Cache
515
"K_SCALE_CONSTANT":
516
lambda: int(os.getenv("K_SCALE_CONSTANT", "200")),
vllm/model_executor/layers/quantization/fp8.py
+15
@@ -3,6 +3,7 @@
3
from typing import Any, Callable, Dict, List, Optional
4
5
import torch
6
+import torch.nn.functional as F
7
from torch.nn import Module
8
from torch.nn.parameter import Parameter
9
@@ -251,6 +252,17 @@ class Fp8LinearMethod(LinearMethodBase):
252
else:
253
layer.register_parameter("input_scale", None)
254
255
+ def add_padding_to_weight(self, weight: torch.Tensor) -> torch.Tensor:
256
+ # Pad the weight tensor. This is an optimization on ROCm platform, which
257
+ # can benefit from tensors located far enough from one another in memory
258
+ if (envs.VLLM_ROCM_FP8_PADDING and current_platform.is_rocm()
259
+ and weight.stride(-1) == 1
260
+ and (weight.stride(-2) * weight.element_size()) % 512 == 0):
261
+ num_pad = 256 // weight.element_size()
262
+ weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
263
+ torch.cuda.empty_cache()
264
+ return weight
265
+
266
def process_weights_after_loading(self, layer: Module) -> None:
267
# TODO(rob): refactor block quant into separate class.
268
if self.block_quant:
@@ -264,6 +276,8 @@ class Fp8LinearMethod(LinearMethodBase):
276
weight = layer.weight.data
277
weight_scale_inv = layer.weight_scale_inv.data
278
279
+ weight = self.add_padding_to_weight(weight)
280
+
281
# Torch.compile cannot use Parameter subclasses.
282
layer.weight = Parameter(weight, requires_grad=False)
283
layer.weight_scale_inv = Parameter(weight_scale_inv,
@@ -327,6 +341,7 @@ class Fp8LinearMethod(LinearMethodBase):
341
logical_widths=layer.logical_widths,
342
)
343
344
+ weight = self.add_padding_to_weight(weight)
345
# Update layer with new values.
346
layer.weight = Parameter(weight.t(), requires_grad=False)
347
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1
-1
@@ -494,7 +494,7 @@ def w8a8_block_fp8_matmul(
494
assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
495
M = A.numel() // A.shape[-1]
496
497
- assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
497
+ assert B.ndim == 2 and Bs.ndim == 2
498
N, K = B.shape
499
assert triton.cdiv(N, block_n) == Bs.shape[0]
500
assert triton.cdiv(K, block_k) == Bs.shape[1]