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, 13:30 UTC
[ROCm][Misc] Follow-ups for Skinny Gemms on ROCm. (#17011)
Signed-off-by: charlifu <charlifu@amd.com>
Charlie Fu committed
Apr 26, 2025 at 00:05 UTC
54271bb7661be9a05dde5ac1164f47659add2f6f
4 files changed
+18
-15
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+4
-3
@@ -155,8 +155,9 @@ def rocm_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
155
scale_b: torch.Tensor, bias: torch.Tensor,
156
input_2d: torch.Tensor,
157
output_shape: List) -> torch.Tensor:
158
- if envs.VLLM_ROCM_USE_SKINNY_GEMM and qinput.shape[
159
- 0] == 1 and qinput.shape[1] % 16 == 0:
158
+ from vllm.platforms.rocm import on_mi250_mi300
159
+ if envs.VLLM_ROCM_USE_SKINNY_GEMM and not on_mi250_mi300(
160
+ ) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0:
161
output = ops.wvSplitKQ(weight.t(), qinput, out_dtype, scale_a, scale_b,
162
current_platform.get_cu_count())
163
else:
@@ -371,7 +372,7 @@ class Fp8LinearOp:
372
373
return w8a8_scaled_mm_func(qinput=qinput,
374
weight=weight,
374
- out_dtype=input.dtype,
375
+ out_dtype=out_dtype,
376
scale_a=x_scale,
377
scale_b=weight_scale,
378
bias=bias,
vllm/model_executor/layers/utils.py
+4
-3
@@ -70,8 +70,9 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,
70
def rocm_unquantized_gemm(x: torch.Tensor,
71
weight: torch.Tensor,
72
bias: Optional[torch.Tensor] = None):
73
+ from vllm.platforms.rocm import on_mi250_mi300
74
k = weight.shape[1]
74
- use_skinny = (envs.VLLM_ROCM_USE_SKINNY_GEMM and \
75
+ use_skinny = (envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi250_mi300() and \
76
x.dtype in [torch.float16, torch.bfloat16] \
77
and k % 8 == 0 and bias is None)
78
@@ -83,11 +84,11 @@ def rocm_unquantized_gemm(x: torch.Tensor,
84
m = weight.shape[0]
85
cu_count = current_platform.get_cu_count()
86
86
- if m > 8 and n < 4:
87
+ if m > 8 and 0 < n < 4:
88
out = ops.wvSplitK(weight, x_view, cu_count)
89
return out.view(*x.shape[:-1], weight.shape[0])
90
elif m % 4 == 0 and n == 1 and k <= 8192:
90
- out = ops.LLMM1(weight, x_view, out, 4)
91
+ out = ops.LLMM1(weight, x_view, 4)
92
return out.view(*x.shape[:-1], weight.shape[0])
93
return torch.nn.functional.linear(x, weight, bias)
94
vllm/model_executor/layers/vocab_parallel_embedding.py
+2
-1
@@ -12,6 +12,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,
12
tensor_model_parallel_all_reduce)
13
from vllm.model_executor.layers.quantization.base_config import (
14
QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
15
+from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
16
from vllm.model_executor.parameter import BasevLLMParameter
17
from vllm.model_executor.utils import set_weight_attrs
18
from vllm.platforms import current_platform
@@ -40,7 +41,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
41
layer: torch.nn.Module,
42
x: torch.Tensor,
43
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
43
- return F.linear(x, layer.weight, bias)
44
+ return dispatch_unquantized_gemm()(x, layer.weight, bias)
45
46
def embedding(self, layer: torch.nn.Module,
47
input_: torch.Tensor) -> torch.Tensor:
vllm/platforms/rocm.py
+8
-8
@@ -98,22 +98,22 @@ def device_id_to_physical_device_id(device_id: int) -> int:
98
return device_id
99
100
101
+def on_mi250_mi300() -> bool:
102
+ GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
103
+ return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942"])
104
+
105
+
106
@cache
107
def use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
108
block_size: int, gqa_ratio: int,
109
max_seq_len: int,
110
sliding_window: int) -> bool:
111
107
- GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
108
- ON_NAVI = "gfx1" in GPU_ARCH
109
- ON_MI250_MI300 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942"])
110
-
111
- # rocm custom page attention not support on navi (gfx1*)
112
+ # rocm custom page attention not support on gfx1*
113
# custom paged attn always supported on V0. On V1, requires sliding window
114
# disabled due to observed numerical discrepancy.
114
- return (ON_MI250_MI300 and not ON_NAVI
115
- and (not envs.VLLM_USE_V1 or sliding_window == 0
116
- or sliding_window == (-1, -1))
115
+ return (on_mi250_mi300() and (not envs.VLLM_USE_V1 or sliding_window == 0
116
+ or sliding_window == (-1, -1))
117
and (qtype == torch.half or qtype == torch.bfloat16)
118
and (head_size == 64 or head_size == 128)
119
and (block_size == 16 or block_size == 32)