@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f41647ee6b
Fork 0

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, 14:45 UTC

[Kernel] Support W8A8 channel-wise weights and per-token activations in triton fused_moe_kernel (#16366)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Apr 11, 2025 at 11:54 UTC f41647ee6b95b38ceb6b81fc397a2526c0499cae
7 files changed +1229 -158
tests/kernels/test_block_fp8.py
+19 -73
@@ -18,6 +18,8 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18 per_token_group_quant_fp8, w8a8_block_fp8_matmul)
19 from vllm.platforms import current_platform
20
21 +from .utils_block import native_w8a8_block_matmul
22 +
23 dg_available = False
24 try:
25 import deep_gemm
@@ -75,61 +77,6 @@ def native_per_token_group_quant_fp8(x,
77 return x_q, x_s
78
79
78 -def native_w8a8_block_fp8_matmul(A,
79 - B,
80 - As,
81 - Bs,
82 - block_size,
83 - output_dtype=torch.float16):
84 - """Matrix multiplication with block-wise quantization using native torch."""
85 - A = A.to(torch.float32)
86 - B = B.to(torch.float32)
87 - assert A.shape[-1] == B.shape[-1]
88 - assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
89 - assert len(block_size) == 2
90 - block_n, block_k = block_size[0], block_size[1]
91 - assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
92 - assert A.shape[:-1] == As.shape[:-1]
93 -
94 - M = A.numel() // A.shape[-1]
95 - N, K = B.shape
96 - origin_C_shape = A.shape[:-1] + (N, )
97 - A = A.reshape(M, A.shape[-1])
98 - As = As.reshape(M, As.shape[-1])
99 - n_tiles = (N + block_n - 1) // block_n
100 - k_tiles = (K + block_k - 1) // block_k
101 - assert n_tiles == Bs.shape[0]
102 - assert k_tiles == Bs.shape[1]
103 -
104 - C_shape = (M, N)
105 - C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
106 -
107 - A_tiles = [
108 - A[:, i * block_k:min((i + 1) * block_k, K)] for i in range(k_tiles)
109 - ]
110 - B_tiles = [[
111 - B[
112 - j * block_n:min((j + 1) * block_n, N),
113 - i * block_k:min((i + 1) * block_k, K),
114 - ] for i in range(k_tiles)
115 - ] for j in range(n_tiles)]
116 - C_tiles = [
117 - C[:, j * block_n:min((j + 1) * block_n, N)] for j in range(n_tiles)
118 - ]
119 - As_tiles = [As[:, i:i + 1] for i in range(k_tiles)]
120 -
121 - for i in range(k_tiles):
122 - for j in range(n_tiles):
123 - a = A_tiles[i]
124 - b = B_tiles[j][i]
125 - c = C_tiles[j]
126 - s = As_tiles[i] * Bs[j][i]
127 - c[:, :] += torch.matmul(a, b.t()) * s
128 -
129 - C = C.reshape(origin_C_shape).to(output_dtype)
130 - return C
131 -
132 -
80 def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
81 """Fused moe with block-wise quantization using native torch."""
82 B, D = a.shape
@@ -146,22 +93,22 @@ def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
93 for i in range(w1.shape[0]):
94 mask = topk_ids == i
95 if mask.sum():
149 - inter_out = native_w8a8_block_fp8_matmul(a_q[mask],
150 - w1[i],
151 - a_s[mask],
152 - w1_s[i],
153 - block_shape,
154 - output_dtype=a.dtype)
96 + inter_out = native_w8a8_block_matmul(a_q[mask],
97 + w1[i],
98 + a_s[mask],
99 + w1_s[i],
100 + block_shape,
101 + output_dtype=a.dtype)
102 act_out = SiluAndMul().forward_native(inter_out)
103 act_out_q, act_out_s = native_per_token_group_quant_fp8(
104 act_out, block_k)
105 act_out = act_out.to(torch.float32)
159 - out[mask] = native_w8a8_block_fp8_matmul(act_out_q,
160 - w2[i],
161 - act_out_s,
162 - w2_s[i],
163 - block_shape,
164 - output_dtype=a.dtype)
106 + out[mask] = native_w8a8_block_matmul(act_out_q,
107 + w2[i],
108 + act_out_s,
109 + w2_s[i],
110 + block_shape,
111 + output_dtype=a.dtype)
112 return (out.view(B, -1, w2.shape[1]) *
113 topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
114
@@ -215,8 +162,8 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
162 As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
163 Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
164
218 - ref_out = native_w8a8_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size,
219 - out_dtype)
165 + ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size,
166 + out_dtype)
167 out = w8a8_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
168
169 rel_diff = (torch.mean(
@@ -239,8 +186,6 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
186 fp8_info = torch.finfo(torch.float8_e4m3fn)
187 fp8_max, fp8_min = fp8_info.max, fp8_info.min
188
242 - vllm_config = VllmConfig()
243 -
189 a = torch.randn((M, K), dtype=dtype) / 10
190
191 w1_bf16 = (torch.rand(
@@ -266,6 +211,7 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
211 score = torch.randn((M, E), dtype=dtype)
212
213 # Set the context to avoid lots of warning spam.
214 + vllm_config = VllmConfig()
215 with set_current_vllm_config(vllm_config):
216 out = fused_moe(
217 a,
@@ -334,8 +280,8 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
280 As = As_fp8.to(torch.float32)
281 Bs = Bs_fp8.to(torch.float32)
282
337 - ref_out = native_w8a8_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size,
338 - out_dtype)
283 + ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size,
284 + out_dtype)
285
286 # Transpose earlier so that the testing will not trigger transposing kernels
287 As_fp8 = deep_gemm.get_col_major_tma_aligned_tensor(As_fp8)
tests/kernels/test_block_int8.py new
+199
@@ -0,0 +1,199 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_block_int8.py
4 +import itertools
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.config import VllmConfig, set_current_vllm_config
10 +from vllm.model_executor.layers.activation import SiluAndMul
11 +from vllm.model_executor.layers.fused_moe import fused_moe
12 +from vllm.model_executor.layers.quantization.utils.int8_utils import (
13 + w8a8_block_int8_matmul)
14 +from vllm.platforms import current_platform
15 +
16 +from .utils_block import native_w8a8_block_matmul
17 +
18 +if current_platform.get_device_capability() < (7, 0):
19 + pytest.skip("INT8 Triton requires CUDA 7.0 or higher",
20 + allow_module_level=True)
21 +
22 +
23 +# For test
24 +def native_per_token_group_quant_int8(x,
25 + group_size,
26 + eps=1e-10,
27 + dtype=torch.int8):
28 + """Function to perform per-token-group quantization on an input tensor
29 + `x` using native torch.
30 +
31 + It converts the tensor values into int8 values and returns the
32 + quantized tensor along with the scaling factor used for quantization.
33 + """
34 + assert (x.shape[-1] % group_size == 0
35 + ), "the last dimension of `x` cannot be divisible by `group_size`"
36 + assert x.is_contiguous(), "`x` is not contiguous"
37 +
38 + iinfo = torch.iinfo(dtype)
39 + int8_min = iinfo.min
40 + int8_max = iinfo.max
41 +
42 + x_ = x.reshape(x.numel() // group_size, group_size)
43 + # Use float32 for scale calculation for stability
44 + amax = x_.abs().max(dim=-1,
45 + keepdim=True)[0].clamp(min=eps).to(torch.float32)
46 + x_s = amax / int8_max
47 + x_q = (x_.to(torch.float32) / x_s).round().clamp(
48 + min=int8_min, max=int8_max).to(dtype) # Round before clamping
49 + x_q = x_q.reshape(x.shape)
50 + x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size, ))
51 +
52 + return x_q, x_s
53 +
54 +
55 +# For test
56 +def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
57 + """This function performs fused moe with block-wise quantization using
58 + native torch."""
59 + B, D = a.shape
60 + a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
61 + out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
62 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
63 + topk_weight, topk_ids = torch.topk(score, topk)
64 + topk_weight = topk_weight.view(-1)
65 + topk_ids = topk_ids.view(-1)
66 +
67 + _, block_k = block_shape[0], block_shape[1]
68 + a_q, a_s = native_per_token_group_quant_int8(a, block_k)
69 + for i in range(w1.shape[0]):
70 + mask = topk_ids == i
71 + if mask.sum():
72 + inter_out = native_w8a8_block_matmul(a_q[mask],
73 + w1[i],
74 + a_s[mask],
75 + w1_s[i],
76 + block_shape,
77 + output_dtype=a.dtype)
78 + act_out = SiluAndMul().forward_native(inter_out)
79 + act_out_q, act_out_s = native_per_token_group_quant_int8(
80 + act_out, block_k)
81 + act_out = act_out.to(torch.float32)
82 + out[mask] = native_w8a8_block_matmul(act_out_q,
83 + w2[i],
84 + act_out_s,
85 + w2_s[i],
86 + block_shape,
87 + output_dtype=a.dtype)
88 + return (out.view(B, -1, w2.shape[1]) *
89 + topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
90 +
91 +
92 +DTYPES = [torch.half, torch.bfloat16]
93 +M = [1, 33, 64, 222]
94 +N = [128, 1024]
95 +K = [256, 4096]
96 +E = [8, 24]
97 +TOP_KS = [2, 6]
98 +# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
99 +BLOCK_SIZE = [[128, 128]]
100 +SEEDS = [0]
101 +
102 +
103 +@pytest.fixture(autouse=True, scope="module")
104 +def setup_cuda():
105 + """Sets the default CUDA device for all tests in this module."""
106 + torch.set_default_device("cuda")
107 +
108 +
109 +@pytest.mark.parametrize("M,N,K,block_size,out_dtype,seed",
110 + itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS))
111 +@torch.inference_mode()
112 +def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed):
113 + torch.manual_seed(seed)
114 + factor_for_scale = 1e-2
115 + int8_info = torch.iinfo(torch.int8)
116 + int8_max, int8_min = int8_info.max, int8_info.min
117 +
118 + A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * int8_max
119 + A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
120 +
121 + B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * int8_max
122 + B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
123 +
124 + block_n, block_k = block_size[0], block_size[1]
125 + n_tiles = (N + block_n - 1) // block_n
126 + k_tiles = (K + block_k - 1) // block_k
127 +
128 + As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
129 + Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
130 +
131 + ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size,
132 + out_dtype)
133 + out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
134 +
135 + rel_diff = (torch.mean(
136 + torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
137 + torch.mean(torch.abs(ref_out.to(torch.float32))))
138 + assert rel_diff < 0.001
139 +
140 +
141 +@pytest.mark.parametrize(
142 + "M, N, K, E, topk, block_size, dtype, seed",
143 + itertools.product(M, N, K, E, TOP_KS, BLOCK_SIZE, DTYPES, SEEDS))
144 +@torch.inference_mode()
145 +def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
146 + """Tests the fused_moe kernel with W8A8 INT8 block quantization against a
147 + native torch reference."""
148 + torch.manual_seed(seed)
149 + # Use a smaller factor for scale initialization to prevent large
150 + # values/overflow especially when output dtype might be float16
151 + factor_for_scale = 1e-2
152 + int8_info = torch.iinfo(torch.int8)
153 + int8_max, int8_min = int8_info.max, int8_info.min
154 +
155 + a = torch.randn((M, K), dtype=dtype) / 10
156 +
157 + w1_fp32 = (torch.rand(
158 + (E, 2 * N, K), dtype=torch.float32) - 0.5) * 2 * int8_max
159 + w1 = w1_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
160 +
161 + w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2 * int8_max
162 + w2 = w2_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
163 +
164 + block_n, block_k = block_size[0], block_size[1]
165 + n_tiles_w1 = (2 * N + block_n - 1) // block_n
166 + n_tiles_w2 = (K + block_n - 1) // block_n
167 + k_tiles_w1 = (K + block_k - 1) // block_k
168 + k_tiles_w2 = (N + block_k - 1) // block_k
169 +
170 + w1_s = (torch.rand(
171 + (E, n_tiles_w1, k_tiles_w1), dtype=torch.float32) * factor_for_scale)
172 + w2_s = (torch.rand(
173 + (E, n_tiles_w2, k_tiles_w2), dtype=torch.float32) * factor_for_scale)
174 +
175 + score = torch.randn((M, E), dtype=dtype)
176 +
177 + # Set the context to avoid lots of warning spam.
178 + vllm_config = VllmConfig()
179 + with set_current_vllm_config(vllm_config):
180 + out = fused_moe(
181 + a,
182 + w1,
183 + w2,
184 + score,
185 + topk,
186 + renormalize=False,
187 + use_int8_w8a8=True,
188 + w1_scale=w1_s,
189 + w2_scale=w2_s,
190 + block_shape=block_size,
191 + )
192 + ref_out = torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk,
193 + block_size)
194 +
195 + # Check results
196 + rel_diff = (torch.mean(
197 + torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
198 + torch.mean(torch.abs(ref_out.to(torch.float32))))
199 + assert rel_diff < 0.06
tests/kernels/test_int8_kernel.py new
+149
@@ -0,0 +1,149 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_int8_kernel.py
4 +import itertools
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.model_executor.layers.activation import SiluAndMul
10 +from vllm.model_executor.layers.fused_moe import fused_moe
11 +from vllm.model_executor.layers.quantization.utils.int8_utils import (
12 + per_token_quant_int8)
13 +from vllm.platforms import current_platform
14 +
15 +if current_platform.get_device_capability() < (7, 0):
16 + pytest.skip("INT8 Triton requires CUDA 7.0 or higher",
17 + allow_module_level=True)
18 +
19 +
20 +def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
21 + """Matrix multiplication function that supports per-token input
22 + quantization and per-column weight quantization"""
23 + A = A.to(torch.float32)
24 + B = B.to(torch.float32)
25 +
26 + assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
27 + assert B.ndim == 2 and B.is_contiguous(
28 + ), "B must be a 2D contiguous tensor"
29 +
30 + # Reshape input
31 + M = A.numel() // A.shape[-1]
32 + B = B.t() # Transpose weight matrix
33 + N, K = B.shape
34 + origin_C_shape = A.shape[:-1] + (K, )
35 + A = A.reshape(M, N)
36 +
37 + # As is per-token [M, 1], Bs is per-column [1, K]
38 + C = torch.matmul(A, B) # [M, K]
39 + C = As * C * Bs.view(1, -1) # Broadcast per-column scale
40 +
41 + return C.reshape(origin_C_shape).to(output_dtype)
42 +
43 +
44 +def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
45 + """This function performs fused moe with per-column int8 quantization
46 + using native torch."""
47 +
48 + B, D = a.shape
49 + # Perform per-token quantization
50 + a_q, a_s = per_token_quant_int8(a)
51 + # Repeat tokens to match topk
52 + a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
53 + # Also repeat the scale
54 + a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1) # [B*topk, 1]
55 +
56 + out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
57 +
58 + # Calculate routing
59 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
60 + topk_weight, topk_ids = torch.topk(score, topk)
61 + topk_weight = topk_weight.view(-1)
62 + topk_ids = topk_ids.view(-1)
63 + # Process each expert
64 + for i in range(w1.shape[0]):
65 + mask = topk_ids == i
66 + if mask.sum():
67 + # First MLP layer: note that a_s is now per-token
68 + inter_out = native_w8a8_per_token_matmul(a_q[mask],
69 + w1[i],
70 + a_s[mask],
71 + w1_s[i],
72 + output_dtype=a.dtype)
73 + # Activation function
74 + act_out = SiluAndMul().forward_native(inter_out)
75 + # Quantize activation output with per-token
76 + act_out_q, act_out_s = per_token_quant_int8(act_out)
77 +
78 + # Second MLP layer
79 + out[mask] = native_w8a8_per_token_matmul(act_out_q,
80 + w2[i],
81 + act_out_s,
82 + w2_s[i],
83 + output_dtype=a.dtype)
84 + # Apply routing weights and sum
85 + return (out.view(B, -1, w2.shape[1]) *
86 + topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
87 +
88 +
89 +@pytest.fixture(autouse=True, scope="module")
90 +def setup_cuda():
91 + """Sets the default CUDA device for all tests in this module."""
92 + torch.set_default_device("cuda")
93 +
94 +
95 +DTYPES = [torch.half, torch.bfloat16]
96 +M = [1, 33]
97 +N = [128, 1024]
98 +K = [256, 4096]
99 +E = [8]
100 +TOP_KS = [2, 6]
101 +SEEDS = [0]
102 +
103 +
104 +@pytest.mark.parametrize("M, N, K, E, topk, dtype, seed",
105 + itertools.product(M, N, K, E, TOP_KS, DTYPES, SEEDS))
106 +@torch.inference_mode()
107 +def test_w8a8_fp8_fused_moe(M, N, K, E, topk, dtype, seed):
108 + torch.manual_seed(seed)
109 + # Initialize int8 quantization parameters
110 + factor_for_scale = 1e-2
111 + int8_max = 127
112 + int8_min = -128
113 +
114 + # Input tensor
115 + # M * K
116 + a = torch.randn((M, K), dtype=dtype) / 10
117 +
118 + # Generate int8 weights
119 + w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
120 + w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
121 +
122 + w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
123 + w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
124 +
125 + # Generate scale for each column (per-column quantization)
126 + w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
127 + w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
128 + score = torch.randn((M, E), dtype=dtype)
129 +
130 + ref_out = torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk)
131 + out = fused_moe(
132 + a,
133 + w1,
134 + w2,
135 + score,
136 + topk,
137 + renormalize=False,
138 + use_int8_w8a8=True, # Using int8-w8a8
139 + per_channel_quant=True,
140 + w1_scale=w1_s,
141 + w2_scale=w2_s,
142 + block_shape=None, # Not using block quantization
143 + )
144 +
145 + # Check results
146 + rel_diff = (torch.mean(
147 + torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
148 + torch.mean(torch.abs(ref_out.to(torch.float32))))
149 + assert rel_diff < 0.05
tests/kernels/test_triton_moe_ptpc_fp8.py new
+159
@@ -0,0 +1,159 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_triton_moe_channel_fp8_kernel.py
4 +import itertools
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm import _custom_ops as ops
10 +from vllm.model_executor.layers.activation import SiluAndMul
11 +from vllm.model_executor.layers.fused_moe import fused_moe
12 +from vllm.platforms import current_platform
13 +
14 +if current_platform.get_device_capability() < (9, 0):
15 + pytest.skip("FP8 Triton requires CUDA 9.0 or higher",
16 + allow_module_level=True)
17 +
18 +
19 +def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
20 + """Matrix multiplication function that supports per-token input
21 + quantization and per-column weight quantization"""
22 + A = A.to(torch.float32)
23 + B = B.to(torch.float32)
24 +
25 + assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
26 + assert B.ndim == 2 and B.is_contiguous(
27 + ), "B must be a 2D contiguous tensor"
28 +
29 + # Reshape input
30 + M = A.numel() // A.shape[-1]
31 + B = B.t() # Transpose weight matrix
32 + N, K = B.shape
33 + origin_C_shape = A.shape[:-1] + (K, )
34 + A = A.reshape(M, N)
35 +
36 + # As is per-token [M, 1], Bs is per-column [1, K]
37 + C = torch.matmul(A, B) # [M, K]
38 + C = As * C * Bs.view(1, -1) # Broadcast per-column scale
39 +
40 + return C.reshape(origin_C_shape).to(output_dtype)
41 +
42 +
43 +def fp8_mask(a, mask):
44 + dtype = a.dtype
45 + return a.view(torch.int8)[mask].view(dtype)
46 +
47 +
48 +def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
49 + """This function performs fused moe with per-column int8
50 + quantization using native torch."""
51 +
52 + B, D = a.shape
53 + # Perform per-token quantization
54 + a_q, a_s = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True)
55 + # Repeat tokens to match topk
56 + a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
57 + # Also repeat the scale
58 + a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1) # [B*topk, 1]
59 +
60 + out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
61 +
62 + # Calculate routing
63 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
64 + topk_weight, topk_ids = torch.topk(score, topk)
65 + topk_weight = topk_weight.view(-1)
66 + topk_ids = topk_ids.view(-1)
67 + # Process each expert
68 + for i in range(w1.shape[0]):
69 + mask = topk_ids == i
70 + if mask.sum():
71 + # First MLP layer: note that a_s is now per-token
72 + inter_out = native_w8a8_per_token_matmul(
73 + fp8_mask(a_q, mask),
74 + w1[i],
75 + fp8_mask(a_s, mask),
76 + w1_s[i],
77 + output_dtype=a.dtype,
78 + )
79 + # Activation function
80 + act_out = SiluAndMul().forward_native(inter_out)
81 + # Quantize activation output with per-token
82 + act_out_q, act_out_s = ops.scaled_fp8_quant(
83 + act_out, use_per_token_if_dynamic=True)
84 +
85 + # Second MLP layer
86 + out[mask] = native_w8a8_per_token_matmul(act_out_q,
87 + w2[i],
88 + act_out_s,
89 + w2_s[i],
90 + output_dtype=a.dtype)
91 + # Apply routing weights and sum
92 + return (out.view(B, -1, w2.shape[1]) *
93 + topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
94 +
95 +
96 +@pytest.fixture(autouse=True, scope="module")
97 +def setup_cuda():
98 + """Sets the default CUDA device for all tests in this module."""
99 + torch.set_default_device("cuda")
100 +
101 +
102 +DTYPES = [torch.half, torch.bfloat16]
103 +M = [1, 33]
104 +N = [128, 1024]
105 +K = [256, 4096]
106 +E = [8]
107 +TOP_KS = [2, 6]
108 +SEEDS = [0]
109 +
110 +
111 +@pytest.mark.parametrize("M, N, K, E, topk, dtype, seed",
112 + itertools.product(M, N, K, E, TOP_KS, DTYPES, SEEDS))
113 +@torch.inference_mode()
114 +def test_w8a8_fp8_fused_moe(M, N, K, E, topk, dtype, seed):
115 + torch.manual_seed(seed)
116 + # Initialize int8 quantization parameters
117 + factor_for_scale = 1e-2
118 + finfo = torch.finfo(torch.float8_e4m3fn)
119 + fp8_max = finfo.max
120 + fp8_min = finfo.min
121 +
122 + # Input tensor
123 + # M * K
124 + a = torch.randn((M, K), dtype=dtype) / 10
125 +
126 + # Generate int8 weights
127 + w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
128 + w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min,
129 + max=fp8_max).to(torch.float8_e4m3fn)
130 +
131 + w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
132 + w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min,
133 + max=fp8_max).to(torch.float8_e4m3fn)
134 +
135 + # Generate scale for each column (per-column quantization)
136 + w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
137 + w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
138 + score = torch.randn((M, E), dtype=dtype)
139 +
140 + ref_out = torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk)
141 + out = fused_moe(
142 + a,
143 + w1,
144 + w2,
145 + score,
146 + topk,
147 + renormalize=False,
148 + use_fp8_w8a8=True, # using fp8
149 + per_channel_quant=True,
150 + w1_scale=w1_s,
151 + w2_scale=w2_s,
152 + block_shape=None, # Not using block quantization
153 + )
154 +
155 + # Check results
156 + rel_diff = (torch.mean(
157 + torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
158 + torch.mean(torch.abs(ref_out.to(torch.float32))))
159 + assert rel_diff < 0.05
tests/kernels/utils_block.py new
+63
@@ -0,0 +1,63 @@
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
vllm/model_executor/layers/fused_moe/fused_moe.py
+181 -85
@@ -16,7 +16,10 @@ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
16 _valid_deep_gemm, deep_gemm_moe_fp8)
17 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
18 moe_align_block_size)
19 -from vllm.model_executor.layers.fused_moe.utils import _fp8_quantize
19 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
20 + per_token_group_quant_fp8)
21 +from vllm.model_executor.layers.quantization.utils.int8_utils import (
22 + per_token_group_quant_int8, per_token_quant_int8)
23 from vllm.platforms import current_platform
24 from vllm.utils import direct_register_custom_op
25
@@ -251,50 +254,53 @@ def fused_moe_kernel_gptq_awq(
254
255 @triton.jit
256 def fused_moe_kernel(
254 - # Pointers to matrices
255 - a_ptr,
256 - b_ptr,
257 - c_ptr,
258 - a_scale_ptr,
259 - b_scale_ptr,
260 - topk_weights_ptr,
261 - sorted_token_ids_ptr,
262 - expert_ids_ptr,
263 - num_tokens_post_padded_ptr,
264 - # Matrix dimensions
265 - N,
266 - K,
267 - EM,
268 - num_valid_tokens,
269 - # The stride variables represent how much to increase the ptr by when
270 - # moving by 1 element in a particular dimension. E.g. `stride_am` is
271 - # how much to increase `a_ptr` by to get the element one row down
272 - # (A has M rows).
273 - stride_am,
274 - stride_ak,
275 - stride_be,
276 - stride_bk,
277 - stride_bn,
278 - stride_cm,
279 - stride_cn,
280 - stride_asm,
281 - stride_ask,
282 - stride_bse,
283 - stride_bsk,
284 - stride_bsn,
285 - # Block size for block-wise quantization
286 - group_n: tl.constexpr,
287 - group_k: tl.constexpr,
288 - # Meta-parameters
289 - BLOCK_SIZE_M: tl.constexpr,
290 - BLOCK_SIZE_N: tl.constexpr,
291 - BLOCK_SIZE_K: tl.constexpr,
292 - GROUP_SIZE_M: tl.constexpr,
293 - MUL_ROUTED_WEIGHT: tl.constexpr,
294 - top_k: tl.constexpr,
295 - compute_type: tl.constexpr,
296 - use_fp8_w8a8: tl.constexpr,
297 - use_int8_w8a16: tl.constexpr):
257 + # Pointers to matrices
258 + a_ptr,
259 + b_ptr,
260 + c_ptr,
261 + a_scale_ptr,
262 + b_scale_ptr,
263 + topk_weights_ptr,
264 + sorted_token_ids_ptr,
265 + expert_ids_ptr,
266 + num_tokens_post_padded_ptr,
267 + # Matrix dimensions
268 + N,
269 + K,
270 + EM,
271 + num_valid_tokens,
272 + # The stride variables represent how much to increase the ptr by when
273 + # moving by 1 element in a particular dimension. E.g. `stride_am` is
274 + # how much to increase `a_ptr` by to get the element one row down
275 + # (A has M rows).
276 + stride_am,
277 + stride_ak,
278 + stride_be,
279 + stride_bk,
280 + stride_bn,
281 + stride_cm,
282 + stride_cn,
283 + stride_asm,
284 + stride_ask,
285 + stride_bse,
286 + stride_bsk,
287 + stride_bsn,
288 + # Block size for block-wise quantization
289 + group_n: tl.constexpr,
290 + group_k: tl.constexpr,
291 + # Meta-parameters
292 + BLOCK_SIZE_M: tl.constexpr,
293 + BLOCK_SIZE_N: tl.constexpr,
294 + BLOCK_SIZE_K: tl.constexpr,
295 + GROUP_SIZE_M: tl.constexpr,
296 + MUL_ROUTED_WEIGHT: tl.constexpr,
297 + top_k: tl.constexpr,
298 + compute_type: tl.constexpr,
299 + use_fp8_w8a8: tl.constexpr,
300 + use_int8_w8a8: tl.constexpr,
301 + use_int8_w8a16: tl.constexpr,
302 + per_channel_quant: tl.constexpr,
303 +):
304 """
305 Implements the fused computation for a Mixture of Experts (MOE) using
306 token and expert matrices.
@@ -371,12 +377,23 @@ def fused_moe_kernel(
377 None, :] * stride_bsn
378 b_scale = tl.load(b_scale_ptrs)
379
374 - if use_fp8_w8a8:
380 + if use_fp8_w8a8 or use_int8_w8a8:
381 + # block-wise
382 if group_k > 0 and group_n > 0:
383 a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm
384 offs_bsn = offs_bn // group_n
385 b_scale_ptrs = (b_scale_ptr + off_experts * stride_bse +
386 offs_bsn * stride_bsn)
387 + # channel-wise
388 + elif per_channel_quant:
389 + b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bn[
390 + None, :] * stride_bsn
391 + b_scale = tl.load(b_scale_ptrs)
392 + # Load per-token scale for activations
393 + a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm
394 + a_scale = tl.load(a_scale_ptrs, mask=token_mask, other=0.0)[:,
395 + None]
396 + # tensor-wise
397 else:
398 a_scale = tl.load(a_scale_ptr)
399 b_scale = tl.load(b_scale_ptr + off_experts)
@@ -400,7 +417,7 @@ def fused_moe_kernel(
417 # We accumulate along the K dimension.
418 if use_int8_w8a16:
419 accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
403 - elif use_fp8_w8a8:
420 + elif use_fp8_w8a8 or use_int8_w8a8:
421 if group_k > 0 and group_n > 0:
422 k_start = k * BLOCK_SIZE_K
423 offs_ks = k_start // group_k
@@ -412,7 +429,11 @@ def fused_moe_kernel(
429 accumulator += tl.dot(a, b) * a_scale[:,
430 None] * b_scale[None, :]
431 else:
415 - accumulator = tl.dot(a, b, acc=accumulator)
432 + if use_fp8_w8a8:
433 + # acc used to enable fp8_fast_accum
434 + accumulator = tl.dot(a, b, acc=accumulator)
435 + else:
436 + accumulator += tl.dot(a, b)
437 else:
438 accumulator += tl.dot(a, b)
439 # Advance the ptrs to the next K block.
@@ -426,7 +447,7 @@ def fused_moe_kernel(
447 accumulator = accumulator * moe_weight[:, None]
448 if use_int8_w8a16:
449 accumulator = (accumulator * b_scale).to(compute_type)
429 - elif use_fp8_w8a8:
450 + elif use_fp8_w8a8 or use_int8_w8a8:
451 if group_k > 0 and group_n > 0:
452 accumulator = accumulator.to(compute_type)
453 else:
@@ -457,27 +478,15 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
478 config: Dict[str, Any],
479 compute_type: tl.dtype,
480 use_fp8_w8a8: bool,
481 + use_int8_w8a8: bool,
482 use_int8_w8a16: bool,
483 use_int4_w4a16: bool,
484 + per_channel_quant: bool,
485 block_shape: Optional[List[int]] = None) -> None:
486 assert topk_weights is not None or not mul_routed_weight
487 assert topk_weights is None or topk_weights.stride(1) == 1
488 assert sorted_token_ids.stride(0) == 1
489
467 - if use_fp8_w8a8:
468 - assert B_scale is not None
469 - assert (block_shape is None or triton.cdiv(B.shape[-2], block_shape[0])
470 - == B_scale.shape[-2])
471 - assert (block_shape is None or triton.cdiv(B.shape[-1], block_shape[1])
472 - == B_scale.shape[-1])
473 -
474 - elif use_int8_w8a16 or use_int4_w4a16:
475 - assert B_scale is not None
476 - assert block_shape is None or block_shape[0] == 0
477 - else:
478 - assert A_scale is None
479 - assert B_scale is None
480 -
490 M = A.shape[0]
491 num_tokens = M * top_k
492
@@ -604,7 +613,9 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
613 top_k=top_k,
614 compute_type=compute_type,
615 use_fp8_w8a8=use_fp8_w8a8,
616 + use_int8_w8a8=use_int8_w8a8,
617 use_int8_w8a16=use_int8_w8a16,
618 + per_channel_quant=per_channel_quant,
619 BLOCK_SIZE_K=BLOCK_SIZE_K,
620 **config,
621 )
@@ -956,8 +967,10 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
967 activation: str = "silu",
968 apply_router_weight_on_input: bool = False,
969 use_fp8_w8a8: bool = False,
970 + use_int8_w8a8: bool = False,
971 use_int8_w8a16: bool = False,
972 use_int4_w4a16: bool = False,
973 + per_channel_quant: bool = False,
974 global_num_experts: int = -1,
975 expert_map: Optional[torch.Tensor] = None,
976 w1_scale: Optional[torch.Tensor] = None,
@@ -969,9 +982,10 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
982 block_shape: Optional[List[int]] = None) -> None:
983 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
984 activation, apply_router_weight_on_input, use_fp8_w8a8,
972 - use_int8_w8a16, use_int4_w4a16, global_num_experts,
973 - expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
974 - a2_scale, block_shape)
985 + use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,
986 + per_channel_quant, global_num_experts, expert_map,
987 + w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
988 + block_shape)
989
990
991 def inplace_fused_experts_fake(
@@ -983,8 +997,10 @@ def inplace_fused_experts_fake(
997 activation: str = "silu",
998 apply_router_weight_on_input: bool = False,
999 use_fp8_w8a8: bool = False,
1000 + use_int8_w8a8: bool = False,
1001 use_int8_w8a16: bool = False,
1002 use_int4_w4a16: bool = False,
1003 + per_channel_quant: bool = False,
1004 global_num_experts: int = -1,
1005 expert_map: Optional[torch.Tensor] = None,
1006 w1_scale: Optional[torch.Tensor] = None,
@@ -1015,8 +1031,10 @@ def outplace_fused_experts(
1031 activation: str = "silu",
1032 apply_router_weight_on_input: bool = False,
1033 use_fp8_w8a8: bool = False,
1034 + use_int8_w8a8: bool = False,
1035 use_int8_w8a16: bool = False,
1036 use_int4_w4a16: bool = False,
1037 + per_channel_quant: bool = False,
1038 global_num_experts: int = -1,
1039 expert_map: Optional[torch.Tensor] = None,
1040 w1_scale: Optional[torch.Tensor] = None,
@@ -1028,7 +1046,8 @@ def outplace_fused_experts(
1046 block_shape: Optional[List[int]] = None) -> torch.Tensor:
1047 return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
1048 False, activation, apply_router_weight_on_input,
1031 - use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16,
1049 + use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
1050 + use_int4_w4a16, per_channel_quant,
1051 global_num_experts, expert_map, w1_scale,
1052 w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
1053 block_shape)
@@ -1042,8 +1061,10 @@ def outplace_fused_experts_fake(
1061 topk_ids: torch.Tensor,
1062 activation: str = "silu",
1063 use_fp8_w8a8: bool = False,
1064 + use_int8_w8a8: bool = False,
1065 use_int8_w8a16: bool = False,
1066 use_int4_w4a16: bool = False,
1067 + per_channel_quant: bool = False,
1068 global_num_experts: int = -1,
1069 expert_map: Optional[torch.Tensor] = None,
1070 w1_scale: Optional[torch.Tensor] = None,
@@ -1092,8 +1113,10 @@ def fused_experts(hidden_states: torch.Tensor,
1113 activation: str = "silu",
1114 apply_router_weight_on_input: bool = False,
1115 use_fp8_w8a8: bool = False,
1116 + use_int8_w8a8: bool = False,
1117 use_int8_w8a16: bool = False,
1118 use_int4_w4a16: bool = False,
1119 + per_channel_quant: bool = False,
1120 global_num_experts: int = -1,
1121 expert_map: Optional[torch.Tensor] = None,
1122 w1_scale: Optional[torch.Tensor] = None,
@@ -1132,8 +1155,10 @@ def fused_experts(hidden_states: torch.Tensor,
1155 activation=activation,
1156 apply_router_weight_on_input=apply_router_weight_on_input,
1157 use_fp8_w8a8=use_fp8_w8a8,
1158 + use_int8_w8a8=use_int8_w8a8,
1159 use_int8_w8a16=use_int8_w8a16,
1160 use_int4_w4a16=use_int4_w4a16,
1161 + per_channel_quant=per_channel_quant,
1162 global_num_experts=global_num_experts,
1163 expert_map=expert_map,
1164 w1_scale=w1_scale,
@@ -1145,6 +1170,59 @@ def fused_experts(hidden_states: torch.Tensor,
1170 block_shape=block_shape)
1171
1172
1173 +def moe_kernel_prepare_input(
1174 + A: torch.Tensor,
1175 + B: torch.Tensor,
1176 + A_scale: Optional[torch.Tensor],
1177 + B_scale: Optional[torch.Tensor],
1178 + use_fp8_w8a8: bool,
1179 + use_int8_w8a8: bool,
1180 + use_int8_w8a16: bool,
1181 + use_int4_w4a16: bool,
1182 + per_channel_quant: bool,
1183 + block_shape: Optional[List[int]] = None,
1184 +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1185 + if use_fp8_w8a8:
1186 + assert B_scale is not None
1187 + if block_shape is None:
1188 + # If weights are per-channel (per_channel_quant=True), then
1189 + # activations apply per-token quantization. Otherwise, assume
1190 + # activation tensor-wise fp8 quantization, dynamic or static
1191 + A, A_scale = ops.scaled_fp8_quant(
1192 + A, A_scale, use_per_token_if_dynamic=per_channel_quant)
1193 + else:
1194 + # activation block-wise fp8 quantization
1195 + assert len(block_shape) == 2
1196 + _, block_k = block_shape[0], block_shape[1]
1197 + A, A_scale = per_token_group_quant_fp8(A, block_k)
1198 + assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
1199 + # assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
1200 + # assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]
1201 + elif use_int8_w8a8:
1202 + assert B_scale is not None
1203 + if block_shape is None:
1204 + # activation channel-wise int8 quantization
1205 + assert (per_channel_quant
1206 + ), "int8 quantization only supports block or channel-wise"
1207 + A, A_scale = per_token_quant_int8(A)
1208 + else:
1209 + # activation block-wise int8 quantization
1210 + assert len(block_shape) == 2
1211 + _, block_k = block_shape[0], block_shape[1]
1212 + A, A_scale = per_token_group_quant_int8(A, block_k)
1213 + assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
1214 + # assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
1215 + # assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]
1216 + elif use_int8_w8a16 or use_int4_w4a16:
1217 + assert B_scale is not None
1218 + assert block_shape is None or block_shape[0] == 0
1219 + else:
1220 + assert A_scale is None
1221 + assert B_scale is None
1222 +
1223 + return A, A_scale
1224 +
1225 +
1226 def fused_experts_impl(hidden_states: torch.Tensor,
1227 w1: torch.Tensor,
1228 w2: torch.Tensor,
@@ -1154,8 +1232,10 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1232 activation: str = "silu",
1233 apply_router_weight_on_input: bool = False,
1234 use_fp8_w8a8: bool = False,
1235 + use_int8_w8a8: bool = False,
1236 use_int8_w8a16: bool = False,
1237 use_int4_w4a16: bool = False,
1238 + per_channel_quant: bool = False,
1239 global_num_experts: int = -1,
1240 expert_map: Optional[torch.Tensor] = None,
1241 w1_scale: Optional[torch.Tensor] = None,
@@ -1257,14 +1337,17 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1337 curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
1338 curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
1339
1260 - a1q_scale: Optional[torch.Tensor] = None
1261 -
1262 - if use_fp8_w8a8:
1263 - qcurr_hidden_states, a1q_scale = _fp8_quantize(
1264 - curr_hidden_states, a1_scale, block_shape)
1265 - else:
1266 - qcurr_hidden_states = curr_hidden_states
1267 - a1q_scale = a1_scale
1340 + qcurr_hidden_states, qa1_scale = moe_kernel_prepare_input(
1341 + A=curr_hidden_states,
1342 + B=w1,
1343 + A_scale=a1_scale,
1344 + B_scale=w1_scale,
1345 + use_fp8_w8a8=use_fp8_w8a8,
1346 + use_int8_w8a8=use_int8_w8a8,
1347 + use_int8_w8a16=use_int8_w8a16,
1348 + use_int4_w4a16=use_int4_w4a16,
1349 + per_channel_quant=per_channel_quant,
1350 + block_shape=block_shape)
1351
1352 sorted_token_ids, expert_ids, num_tokens_post_padded = (
1353 moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'],
@@ -1273,7 +1356,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1356 invoke_fused_moe_kernel(qcurr_hidden_states,
1357 w1,
1358 intermediate_cache1,
1276 - a1q_scale,
1359 + qa1_scale,
1360 w1_scale,
1361 w1_zp,
1362 curr_topk_weights,
@@ -1285,8 +1368,10 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1368 config,
1369 compute_type=compute_type,
1370 use_fp8_w8a8=use_fp8_w8a8,
1371 + use_int8_w8a8=use_int8_w8a8,
1372 use_int8_w8a16=use_int8_w8a16,
1373 use_int4_w4a16=use_int4_w4a16,
1374 + per_channel_quant=per_channel_quant,
1375 block_shape=block_shape)
1376
1377 if activation == "silu":
@@ -1298,19 +1383,22 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1383 else:
1384 raise ValueError(f"Unsupported FusedMoe activation: {activation}")
1385
1301 - a2q_scale: Optional[torch.Tensor] = None
1302 -
1303 - if use_fp8_w8a8:
1304 - qintermediate_cache2, a2q_scale = _fp8_quantize(
1305 - intermediate_cache2, a2_scale, block_shape)
1306 - else:
1307 - qintermediate_cache2 = intermediate_cache2
1308 - a2q_scale = a2_scale
1386 + qintermediate_cache2, qa2_scale = moe_kernel_prepare_input(
1387 + A=intermediate_cache2,
1388 + B=w2,
1389 + A_scale=a2_scale,
1390 + B_scale=w2_scale,
1391 + use_fp8_w8a8=use_fp8_w8a8,
1392 + use_int8_w8a8=use_int8_w8a8,
1393 + use_int8_w8a16=use_int8_w8a16,
1394 + use_int4_w4a16=use_int4_w4a16,
1395 + per_channel_quant=per_channel_quant,
1396 + block_shape=block_shape)
1397
1398 invoke_fused_moe_kernel(qintermediate_cache2,
1399 w2,
1400 intermediate_cache3,
1313 - a2q_scale,
1401 + qa2_scale,
1402 w2_scale,
1403 w2_zp,
1404 curr_topk_weights,
@@ -1322,8 +1410,10 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1410 config,
1411 compute_type=compute_type,
1412 use_fp8_w8a8=use_fp8_w8a8,
1413 + use_int8_w8a8=use_int8_w8a8,
1414 use_int8_w8a16=use_int8_w8a16,
1415 use_int4_w4a16=use_int4_w4a16,
1416 + per_channel_quant=per_channel_quant,
1417 block_shape=block_shape)
1418
1419 ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.shape),
@@ -1346,8 +1436,10 @@ def fused_moe(
1436 topk_group: Optional[int] = None,
1437 custom_routing_function: Optional[Callable] = None,
1438 use_fp8_w8a8: bool = False,
1439 + use_int8_w8a8: bool = False,
1440 use_int8_w8a16: bool = False,
1441 use_int4_w4a16: bool = False,
1442 + per_channel_quant: bool = False,
1443 global_num_experts: int = -1,
1444 expert_map: Optional[torch.Tensor] = None,
1445 w1_scale: Optional[torch.Tensor] = None,
@@ -1380,6 +1472,8 @@ def fused_moe(
1472 note: Deepseekv2 model uses grouped_topk
1473 - use_fp8_w8a8 (bool): If True, use fp8 arithmetic to compute the inner
1474 products for w1 and w2. Defaults to False.
1475 + - use_int8_w8a8 (bool): If True, use int8 arithmetic to compute the inner
1476 + products for w1 and w2. Defaults to False.
1477 - use_int8_w8a16 (bool): If True, use matmul of int8 weight and bf16/fp16
1478 activation to compute the inner products for w1 and w2.
1479 Defaults to False.
@@ -1426,8 +1520,10 @@ def fused_moe(
1520 inplace=inplace,
1521 activation=activation,
1522 use_fp8_w8a8=use_fp8_w8a8,
1523 + use_int8_w8a8=use_int8_w8a8,
1524 use_int8_w8a16=use_int8_w8a16,
1525 use_int4_w4a16=use_int4_w4a16,
1526 + per_channel_quant=per_channel_quant,
1527 global_num_experts=global_num_experts,
1528 expert_map=expert_map,
1529 w1_scale=w1_scale,
vllm/model_executor/layers/quantization/utils/int8_utils.py new
+459
@@ -0,0 +1,459 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from https://github.com/sgl-project/sglang/blob/4cb53ecd0cffceb6dee5c011a58f65997a86f151/python/sglang/srt/layers/quantization/int8_kernel.py
4 +import functools
5 +import json
6 +import logging
7 +import os
8 +from typing import Any, Dict, List, Optional, Tuple
9 +
10 +import torch
11 +import triton
12 +import triton.language as tl
13 +
14 +from vllm.platforms import current_platform
15 +
16 +logger = logging.getLogger(__name__)
17 +
18 +
19 +def apply_w8a8_block_int8_linear(
20 + input: torch.Tensor,
21 + weight: torch.Tensor,
22 + block_size: List[int],
23 + weight_scale: torch.Tensor,
24 + input_scale: Optional[torch.Tensor] = None,
25 + bias: Optional[torch.Tensor] = None,
26 +) -> torch.Tensor:
27 + assert input_scale is None
28 + # View input as 2D matrix for fp8 methods
29 + input_2d = input.view(-1, input.shape[-1])
30 + output_shape = [*input.shape[:-1], weight.shape[0]]
31 +
32 + q_input, x_scale = per_token_group_quant_int8(input_2d, block_size[1])
33 + output = w8a8_block_int8_matmul(q_input,
34 + weight,
35 + x_scale,
36 + weight_scale,
37 + block_size,
38 + output_dtype=input.dtype)
39 +
40 + if bias is not None:
41 + output = output + bias
42 + return output.to(dtype=input.dtype).view(*output_shape)
43 +
44 +
45 +def input_to_int8(
46 + x: torch.Tensor,
47 + dtype: torch.dtype = torch.int8) -> Tuple[torch.Tensor, torch.Tensor]:
48 + """This function quantizes input values to int8 values with
49 + tensor-wise quantization."""
50 + iinfo = torch.iinfo(dtype)
51 + min_val, max_val = x.aminmax()
52 + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
53 + int8_min, int8_max = iinfo.min, iinfo.max
54 + scale = int8_max / amax
55 + x_scl_sat = (x * scale).clamp(min=int8_min, max=int8_max)
56 + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
57 +
58 +
59 +def block_dequant(
60 + x_q_block: torch.Tensor,
61 + x_s: torch.Tensor,
62 + block_size: List[int],
63 +) -> torch.Tensor:
64 + """This function conducts block-wise dequantization.
65 + The inputs are block-wise quantization tensor `x_q_block`,
66 + block-wise quantization scale and the block size.
67 + The outputs are dequantized tensor.
68 + """
69 + block_n, block_k = block_size[0], block_size[1]
70 + n, k = x_q_block.shape
71 + n_tiles = (n + block_n - 1) // block_n
72 + k_tiles = (k + block_k - 1) // block_k
73 + assert n_tiles == x_s.shape[0]
74 + assert k_tiles == x_s.shape[1]
75 +
76 + x_dq_block = x_q_block.to(torch.float32)
77 +
78 + for i in range(k_tiles):
79 + for j in range(n_tiles):
80 + x_dq_block[
81 + j * block_n:min((j + 1) * block_n, n),
82 + i * block_k:min((i + 1) * block_k, k),
83 + ] *= x_s[j][i]
84 +
85 + return x_dq_block
86 +
87 +
88 +@triton.jit
89 +def _per_token_quant_int8(
90 + x_ptr,
91 + xq_ptr,
92 + scale_ptr,
93 + stride_x,
94 + stride_xq,
95 + N,
96 + BLOCK: tl.constexpr,
97 +):
98 + # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
99 + row_id = tl.program_id(0)
100 +
101 + cols = tl.arange(0, BLOCK)
102 + mask = cols < N
103 +
104 + x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask,
105 + other=0.0).to(tl.float32)
106 + absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
107 + scale_x = absmax / 127
108 + x_q = x * (127 / absmax)
109 + x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8)
110 +
111 + tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)
112 + tl.store(scale_ptr + row_id, scale_x)
113 +
114 +
115 +def per_token_quant_int8(x):
116 + M = x.numel() // x.shape[-1]
117 + N = x.shape[-1]
118 + x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
119 + scales = torch.empty(x.shape[:-1] + (1, ),
120 + device=x.device,
121 + dtype=torch.float32)
122 + BLOCK = triton.next_power_of_2(N)
123 + # heuristics for number of warps
124 + num_warps = min(max(BLOCK // 256, 1), 8)
125 +
126 + assert x.is_contiguous()
127 + _per_token_quant_int8[(M, )](
128 + x,
129 + x_q,
130 + scales,
131 + stride_x=x.stride(-2),
132 + stride_xq=x_q.stride(-2),
133 + N=N,
134 + BLOCK=BLOCK,
135 + num_warps=num_warps,
136 + num_stages=1,
137 + )
138 +
139 + return x_q, scales
140 +
141 +
142 +@triton.jit
143 +def _per_token_group_quant_int8(
144 + # Pointers to inputs and output
145 + y_ptr,
146 + y_q_ptr,
147 + y_s_ptr,
148 + # Stride of input
149 + y_stride,
150 + # Columns of input
151 + N,
152 + # Avoid to divide zero
153 + eps,
154 + # Information for int8
155 + int8_min,
156 + int8_max,
157 + # Meta-parameters
158 + BLOCK: tl.constexpr,
159 +):
160 + """A Triton-accelerated function to perform per-token-group
161 + quantization on a tensor.
162 +
163 + This function converts the tensor values into int8 values.
164 + """
165 + # Map the program id to the row of X and Y it should compute.
166 + g_id = tl.program_id(0)
167 + y_ptr += g_id * y_stride
168 + y_q_ptr += g_id * y_stride
169 + y_s_ptr += g_id
170 +
171 + cols = tl.arange(0, BLOCK) # N <= BLOCK
172 + mask = cols < N
173 +
174 + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
175 + # Quant
176 + _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
177 + y_s = _absmax / int8_max
178 + y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty)
179 +
180 + tl.store(y_q_ptr + cols, y_q, mask=mask)
181 + tl.store(y_s_ptr, y_s)
182 +
183 +
184 +def per_token_group_quant_int8(
185 + x: torch.Tensor,
186 + group_size: int,
187 + eps: float = 1e-10,
188 + dtype: torch.dtype = torch.int8,
189 +) -> Tuple[torch.Tensor, torch.Tensor]:
190 + """Function to perform per-token-group quantization on an input tensor `x`.
191 +
192 + It converts the tensor values into signed int8 values and returns the
193 + quantized tensor along with the scaling factor used for quantization.
194 +
195 + Args:
196 + x: The input tenosr with ndim >= 2.
197 + group_size: The group size used for quantization.
198 + eps: The minimum to avoid dividing zero.
199 + dtype: The dype of output tensor. Note that only `torch.int8`
200 + is supported for now.
201 +
202 + Returns:
203 + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
204 + scaling factor for quantization.
205 + """
206 + assert (x.shape[-1] % group_size == 0
207 + ), "the last dimension of `x` cannot be divisible by `group_size`"
208 + assert x.is_contiguous(), "`x` is not contiguous"
209 +
210 + iinfo = torch.iinfo(dtype)
211 + int8_max = iinfo.max
212 + int8_min = iinfo.min
213 +
214 + x_q = torch.empty_like(x, device=x.device, dtype=dtype)
215 + M = x.numel() // group_size
216 + N = group_size
217 + x_s = torch.empty(
218 + x.shape[:-1] + (x.shape[-1] // group_size, ),
219 + device=x.device,
220 + dtype=torch.float32,
221 + )
222 +
223 + BLOCK = triton.next_power_of_2(N)
224 + # heuristics for number of warps
225 + num_warps = min(max(BLOCK // 256, 1), 8)
226 + num_stages = 1
227 + _per_token_group_quant_int8[(M, )](
228 + x,
229 + x_q,
230 + x_s,
231 + group_size,
232 + N,
233 + eps,
234 + int8_min=int8_min,
235 + int8_max=int8_max,
236 + BLOCK=BLOCK,
237 + num_warps=num_warps,
238 + num_stages=num_stages,
239 + )
240 +
241 + return x_q, x_s
242 +
243 +
244 +@triton.jit
245 +def _w8a8_block_int8_matmul(
246 + # Pointers to inputs and output
247 + A,
248 + B,
249 + C,
250 + As,
251 + Bs,
252 + # Shape for matmul
253 + M,
254 + N,
255 + K,
256 + # Block size for block-wise quantization
257 + group_n,
258 + group_k,
259 + # Stride for inputs and output
260 + stride_am,
261 + stride_ak,
262 + stride_bk,
263 + stride_bn,
264 + stride_cm,
265 + stride_cn,
266 + stride_As_m,
267 + stride_As_k,
268 + stride_Bs_k,
269 + stride_Bs_n,
270 + # Meta-parameters
271 + BLOCK_SIZE_M: tl.constexpr,
272 + BLOCK_SIZE_N: tl.constexpr,
273 + BLOCK_SIZE_K: tl.constexpr,
274 + GROUP_SIZE_M: tl.constexpr,
275 +):
276 + """Triton-accelerated function used to perform linear operations (dot
277 + product) on input tensors `A` and `B` with block-wise quantization, and
278 + store the result in output tensor `C`.
279 + """
280 +
281 + pid = tl.program_id(axis=0)
282 + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
283 + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
284 + num_pid_in_group = GROUP_SIZE_M * num_pid_n
285 + group_id = pid // num_pid_in_group
286 + first_pid_m = group_id * GROUP_SIZE_M
287 + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
288 + pid_m = first_pid_m + (pid % group_size_m)
289 + pid_n = (pid % num_pid_in_group) // group_size_m
290 +
291 + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
292 + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
293 + offs_k = tl.arange(0, BLOCK_SIZE_K)
294 + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
295 + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
296 +
297 + As_ptrs = As + offs_am * stride_As_m
298 + offs_bsn = offs_bn // group_n
299 + Bs_ptrs = Bs + offs_bsn * stride_Bs_n
300 +
301 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
302 + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
303 + a = tl.load(a_ptrs,
304 + mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
305 + other=0.0)
306 + b = tl.load(b_ptrs,
307 + mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
308 + other=0.0)
309 +
310 + k_start = k * BLOCK_SIZE_K
311 + offs_ks = k_start // group_k
312 + a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
313 + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
314 +
315 + accumulator += tl.dot(a, b).to(tl.float32) * a_s[:,
316 + None] * b_s[None, :]
317 + a_ptrs += BLOCK_SIZE_K * stride_ak
318 + b_ptrs += BLOCK_SIZE_K * stride_bk
319 +
320 + if C.dtype.element_ty == tl.bfloat16:
321 + c = accumulator.to(tl.bfloat16)
322 + elif C.dtype.element_ty == tl.float16:
323 + c = accumulator.to(tl.float16)
324 + else:
325 + c = accumulator.to(tl.float32)
326 +
327 + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
328 + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
329 + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
330 + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
331 + tl.store(c_ptrs, c, mask=c_mask)
332 +
333 +
334 +@functools.lru_cache
335 +def get_w8a8_block_int8_configs(N: int, K: int, block_n: int,
336 + block_k: int) -> Optional[Dict[int, Any]]:
337 + """
338 + Return optimized configurations for the w8a8 block fp8 kernel.
339 +
340 + The return value will be a dictionary that maps an irregular grid of
341 + batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the
342 + kernel on a given batch size bs, the closest batch size in the grid should
343 + be picked and the associated configuration chosen to invoke the kernel.
344 + """
345 +
346 + # First look up if an optimized configuration is available in the configs
347 + # directory
348 + device_name = current_platform.get_device_name().replace(" ", "_")
349 + json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" # noqa: E501
350 +
351 + config_file_path = os.path.join(
352 + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
353 + if os.path.exists(config_file_path):
354 + with open(config_file_path) as f:
355 + logger.info(
356 + "Using configuration from %s for W8A8 Block INT8 kernel.",
357 + config_file_path,
358 + )
359 + # If a configuration has been found, return it
360 + return {int(key): val for key, val in json.load(f).items()}
361 +
362 + # If no optimized configuration is available, we will use the default
363 + # configuration
364 + logger.warning(
365 + ("Using default W8A8 Block INT8 kernel config. Performance might "
366 + "be sub-optimal! Config file not found at %s"),
367 + config_file_path,
368 + )
369 + return None
370 +
371 +
372 +def w8a8_block_int8_matmul(
373 + A: torch.Tensor,
374 + B: torch.Tensor,
375 + As: torch.Tensor,
376 + Bs: torch.Tensor,
377 + block_size: List[int],
378 + output_dtype: torch.dtype = torch.float16,
379 +) -> torch.Tensor:
380 + """This function performs matrix multiplication with block-wise
381 + quantization.
382 +
383 + It takes two input tensors `A` and `B` with scales `As` and `Bs`.
384 + The output is returned in the specified `output_dtype`.
385 +
386 + Args:
387 + A: The input tensor, e.g., activation.
388 + B: The input tensor, e.g., weight.
389 + As: The per-token-group quantization scale for `A`.
390 + Bs: The per-block quantization scale for `B`.
391 + block_size: The block size for per-block quantization. It should be
392 + 2-dim, e.g., [128, 128].
393 + output_dytpe: The dtype of the returned tensor.
394 +
395 + Returns:
396 + torch.Tensor: The result of matmul.
397 + """
398 + assert len(block_size) == 2
399 + block_n, block_k = block_size[0], block_size[1]
400 +
401 + assert A.shape[-1] == B.shape[-1]
402 + assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
403 + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
404 + M = A.numel() // A.shape[-1]
405 +
406 + assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
407 + N, K = B.shape
408 + assert triton.cdiv(N, block_n) == Bs.shape[0]
409 + assert triton.cdiv(K, block_k) == Bs.shape[1]
410 +
411 + C_shape = A.shape[:-1] + (N, )
412 + C = A.new_empty(C_shape, dtype=output_dtype)
413 +
414 + configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1])
415 + if configs:
416 + # If an optimal configuration map has been found, look up the
417 + # optimal config
418 + config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
419 + else:
420 + # Default config
421 + # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1]
422 + config = {
423 + "BLOCK_SIZE_M": 64,
424 + "BLOCK_SIZE_N": block_size[0],
425 + "BLOCK_SIZE_K": block_size[1],
426 + "GROUP_SIZE_M": 32,
427 + "num_warps": 4,
428 + "num_stages": 3,
429 + }
430 +
431 + def grid(META):
432 + return (triton.cdiv(M, META["BLOCK_SIZE_M"]) *
433 + triton.cdiv(N, META["BLOCK_SIZE_N"]), )
434 +
435 + _w8a8_block_int8_matmul[grid](
436 + A,
437 + B,
438 + C,
439 + As,
440 + Bs,
441 + M,
442 + N,
443 + K,
444 + block_n,
445 + block_k,
446 + A.stride(-2),
447 + A.stride(-1),
448 + B.stride(1),
449 + B.stride(0),
450 + C.stride(-2),
451 + C.stride(-1),
452 + As.stride(-2),
453 + As.stride(-1),
454 + Bs.stride(1),
455 + Bs.stride(0),
456 + **config,
457 + )
458 +
459 + return C