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
Add missing rocm_skinny_gemms kernel test to CI (#17060)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Apr 24, 2025 at 08:49 UTC
82e43b2d7e00dfffef4101c3ceb9e7893bf13abb
5 files changed
+62
-65
tests/kernels/quant_utils.py
+60
@@ -87,3 +87,63 @@ def ref_dynamic_per_tensor_fp8_quant(x: torch.tensor) \
87
ref_out = (as_float32_tensor(x) * ref_iscale).clamp(
88
fp8_traits_min, fp8_traits_max).to(FP8_DTYPE)
89
return ref_out, ref_scale.view((1, ))
90
+
91
+
92
+def native_w8a8_block_matmul(A: torch.Tensor, B: torch.Tensor,
93
+ As: torch.Tensor, Bs: torch.Tensor, block_size,
94
+ output_dtype):
95
+ """This function performs matrix multiplication with block-wise
96
+ quantization using native torch.
97
+ It is agnostic to the input data type and can be used for both int8 and
98
+ fp8 data types.
99
+
100
+ It takes two input tensors `A` and `B` (int8) with scales `As` and
101
+ `Bs` (float32).
102
+ The output is returned in the specified `output_dtype`.
103
+ """
104
+ A = A.to(torch.float32)
105
+ B = B.to(torch.float32)
106
+ assert A.shape[-1] == B.shape[-1]
107
+ assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
108
+ assert len(block_size) == 2
109
+ block_n, block_k = block_size[0], block_size[1]
110
+ assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
111
+ assert A.shape[:-1] == As.shape[:-1]
112
+
113
+ M = A.numel() // A.shape[-1]
114
+ N, K = B.shape
115
+ origin_C_shape = A.shape[:-1] + (N, )
116
+ A = A.reshape(M, A.shape[-1])
117
+ As = As.reshape(M, As.shape[-1])
118
+ n_tiles = (N + block_n - 1) // block_n
119
+ k_tiles = (K + block_k - 1) // block_k
120
+ assert n_tiles == Bs.shape[0]
121
+ assert k_tiles == Bs.shape[1]
122
+
123
+ C_shape = (M, N)
124
+ C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
125
+
126
+ A_tiles = [
127
+ A[:, i * block_k:min((i + 1) * block_k, K)] for i in range(k_tiles)
128
+ ]
129
+ B_tiles = [[
130
+ B[
131
+ j * block_n:min((j + 1) * block_n, N),
132
+ i * block_k:min((i + 1) * block_k, K),
133
+ ] for i in range(k_tiles)
134
+ ] for j in range(n_tiles)]
135
+ C_tiles = [
136
+ C[:, j * block_n:min((j + 1) * block_n, N)] for j in range(n_tiles)
137
+ ]
138
+ As_tiles = [As[:, i:i + 1] for i in range(k_tiles)]
139
+
140
+ for i in range(k_tiles):
141
+ for j in range(n_tiles):
142
+ a = A_tiles[i]
143
+ b = B_tiles[j][i]
144
+ c = C_tiles[j]
145
+ s = As_tiles[i] * Bs[j][i]
146
+ c[:, :] += torch.matmul(a, b.t()) * s
147
+
148
+ C = C.reshape(origin_C_shape).to(output_dtype)
149
+ return C
tests/kernels/quantization/test_block_fp8.py
+1
-1
@@ -6,7 +6,7 @@ import itertools
6
import pytest
7
import torch
8
9
-from tests.kernels.utils_block import native_w8a8_block_matmul
9
+from tests.kernels.quant_utils import native_w8a8_block_matmul
10
from vllm.config import VllmConfig, set_current_vllm_config
11
from vllm.model_executor.layers.activation import SiluAndMul
12
from vllm.model_executor.layers.fused_moe import fused_moe
tests/kernels/quantization/test_block_int8.py
+1
-1
@@ -6,7 +6,7 @@ import itertools
6
import pytest
7
import torch
8
9
-from tests.kernels.utils_block import native_w8a8_block_matmul
9
+from tests.kernels.quant_utils import native_w8a8_block_matmul
10
from vllm.config import VllmConfig, set_current_vllm_config
11
from vllm.model_executor.layers.activation import SiluAndMul
12
from vllm.model_executor.layers.fused_moe import fused_moe
tests/kernels/quantization/test_rocm_skinny_gemms.py
renamed
tests/kernels/utils_block.py
deleted
-63
@@ -1,63 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-
3
-import torch
4
-
5
-
6
-def native_w8a8_block_matmul(A: torch.Tensor, B: torch.Tensor,
7
- As: torch.Tensor, Bs: torch.Tensor, block_size,
8
- output_dtype):
9
- """This function performs matrix multiplication with block-wise
10
- quantization using native torch.
11
- It is agnostic to the input data type and can be used for both int8 and
12
- fp8 data types.
13
-
14
- It takes two input tensors `A` and `B` (int8) with scales `As` and
15
- `Bs` (float32).
16
- The output is returned in the specified `output_dtype`.
17
- """
18
- A = A.to(torch.float32)
19
- B = B.to(torch.float32)
20
- assert A.shape[-1] == B.shape[-1]
21
- assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
22
- assert len(block_size) == 2
23
- block_n, block_k = block_size[0], block_size[1]
24
- assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
25
- assert A.shape[:-1] == As.shape[:-1]
26
-
27
- M = A.numel() // A.shape[-1]
28
- N, K = B.shape
29
- origin_C_shape = A.shape[:-1] + (N, )
30
- A = A.reshape(M, A.shape[-1])
31
- As = As.reshape(M, As.shape[-1])
32
- n_tiles = (N + block_n - 1) // block_n
33
- k_tiles = (K + block_k - 1) // block_k
34
- assert n_tiles == Bs.shape[0]
35
- assert k_tiles == Bs.shape[1]
36
-
37
- C_shape = (M, N)
38
- C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
39
-
40
- A_tiles = [
41
- A[:, i * block_k:min((i + 1) * block_k, K)] for i in range(k_tiles)
42
- ]
43
- B_tiles = [[
44
- B[
45
- j * block_n:min((j + 1) * block_n, N),
46
- i * block_k:min((i + 1) * block_k, K),
47
- ] for i in range(k_tiles)
48
- ] for j in range(n_tiles)]
49
- C_tiles = [
50
- C[:, j * block_n:min((j + 1) * block_n, N)] for j in range(n_tiles)
51
- ]
52
- As_tiles = [As[:, i:i + 1] for i in range(k_tiles)]
53
-
54
- for i in range(k_tiles):
55
- for j in range(n_tiles):
56
- a = A_tiles[i]
57
- b = B_tiles[j][i]
58
- c = C_tiles[j]
59
- s = As_tiles[i] * Bs[j][i]
60
- c[:, :] += torch.matmul(a, b.t()) * s
61
-
62
- C = C.reshape(origin_C_shape).to(output_dtype)
63
- return C