test_cutlass.py 20.4 KB
Newer Older
1
2
3
4
"""Tests for cutlass kernels

Run `pytest tests/kernels/test_cutlass.py`.
"""
5
from typing import Type
6
7
8
9

import pytest
import torch

10
from tests.kernels.utils import opcheck
11
from vllm import _custom_ops as ops
12
from vllm.platforms import current_platform
13
from vllm.utils import cdiv
14

15
16
from .utils import baseline_scaled_mm, to_fp8, to_int8

17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
MNK_FACTORS = [
    (1, 256, 128),
    (1, 16384, 1024),
    (1, 24576, 496),
    (16, 256, 496),
    (16, 16384, 128),
    (16, 24576, 4096),
    (32, 8192, 4096),
    (32, 16384, 4096),
    (33, 1024, 1024),
    (33, 8192, 128),
    (64, 2048, 496),
    (64, 16384, 1024),
    (100, 8192, 496),
    (128, 32768, 4096),
    (256, 4096, 4096),
    (512, 256, 1024),
    (512, 8192, 4096),
    (512, 16384, 128),
    (512, 24576, 128),
]

39
40
41
42
CUDA_DEVICES = [
    f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
]

43
44
45
46
47
# -1 means full extent in that dimension
TENSORWISE_GROUP_SHAPE = (-1, -1)
PER_TOKEN_GROUP_SHAPE = (1, -1)
PER_OUT_CH_GROUP_SHAPE = (-1, 1)

48
49
50
capability = current_platform.get_device_capability()
capability = capability[0] * 10 + capability[1]

51

52
53
54
55
def rand_int8(shape: tuple, device: str = "cuda"):
    return to_int8(torch.rand(shape, device=device) * 255 - 128)


56
57
58
59
60
61
62
63
64
65
66
def group_scale_helper(shape, group_shape):
    return [shape[i] if s < 0 else s for i, s in enumerate(group_shape)]


def scale_shape(shape, group_shape):
    assert len(shape) == len(group_shape)
    group_shape = group_scale_helper(shape, group_shape)
    return tuple(
        cdiv(shape[i], group_shape[i]) for i in range(len(group_shape)))


67
68
69
def cutlass_fp8_gemm_helper(m: int,
                            n: int,
                            k: int,
70
71
                            a_scale_group_shape: tuple,
                            b_scale_group_shape: tuple,
72
                            use_bias: bool,
73
74
75
76
77
78
79
                            out_dtype: Type[torch.dtype] = torch.bfloat16,
                            device: str = "cuda"):
    # Test for a cutlass kernel with per-token activation quantization
    # and per-output channel weight quantization.
    a = to_fp8(torch.randn((m, k), device=device))
    b = to_fp8(torch.randn((n, k), device=device).t())

80
81
82
83
84
85
86
87
88
89
    a_scales_shape = scale_shape(a.shape, a_scale_group_shape)
    b_scales_shape = scale_shape(b.shape, b_scale_group_shape)

    scale_a = (torch.randn(a_scales_shape, device=device, dtype=torch.float32))
    scale_b = (torch.randn(b_scales_shape, device=device, dtype=torch.float32))

    # make scales M-major for blockwise quant, doesn't affect 1D scales
    scale_a = scale_a.t().contiguous().t()
    # make scales K-major for blockwise quant, doesn't affect 1D scales
    scale_b = scale_b.t().contiguous().t()
90

91
92
    if use_bias:
        bias = torch.rand((n, ), device=device, dtype=out_dtype) * 10
93
    else:
94
        bias = None
95

96
97
    out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
    baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
98

99
    torch.testing.assert_close(out, baseline, rtol=1e-2, atol=5e-2)
100

101
102
103
    opcheck(torch.ops._C.cutlass_scaled_mm,
            (out, a, b, scale_a, scale_b, bias))

104
105
106
107

def cutlass_int8_gemm_helper(m: int,
                             n: int,
                             k: int,
108
109
                             a_scale_group_shape: tuple,
                             b_scale_group_shape: tuple,
110
                             use_bias: bool,
111
112
113
114
115
116
117
                             out_dtype: Type[torch.dtype] = torch.bfloat16,
                             device: str = "cuda"):
    # Test for a cutlass kernel with per-token activation quantization
    # and per-output channel weight quantization.
    a = to_int8(torch.randn((m, k), device=device) * 5)
    b = to_int8(torch.randn((n, k), device=device).t() * 5)

118
119
    a_scales_shape = scale_shape(a.shape, a_scale_group_shape)
    b_scales_shape = scale_shape(b.shape, b_scale_group_shape)
120

121
122
    scale_a = (torch.randn(a_scales_shape, device=device, dtype=torch.float32))
    scale_b = (torch.randn(b_scales_shape, device=device, dtype=torch.float32))
123

124
125
    if use_bias:
        bias = torch.rand((n, ), device=device, dtype=out_dtype) * 10
126
    else:
127
128
129
130
        bias = None

    out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
    baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
131

132
    torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
133

134
135
136
    opcheck(torch.ops._C.cutlass_scaled_mm,
            (out, a, b, scale_a, scale_b, bias))

137

138
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
139
140
141
142
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
143
@pytest.mark.parametrize("use_bias", [True, False])
144
@pytest.mark.skipif(not current_platform.has_device_capability(89),
145
                    reason="FP8 is not supported on this GPU type.")
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def test_cutlass_fp8_gemm(m: int, n: int, k: int, a_scale_group_shape,
                          b_scale_group_shape, use_bias: bool):
    cutlass_fp8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape,
                            use_bias)


@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
@pytest.mark.parametrize("a_scale_group_shape,b_scale_group_shape",
                         [((1, 128), (128, 128))])
@pytest.mark.parametrize("use_bias", [False])
@pytest.mark.skipif(not current_platform.has_device_capability(90),
                    reason="FP8 blockwise is not supported on this GPU type.")
def test_cutlass_fp8_blockwise_scale_gemm(m: int, n: int, k: int,
                                          a_scale_group_shape,
                                          b_scale_group_shape, use_bias: bool):
    if k % b_scale_group_shape[0] != 0 or n % b_scale_group_shape[1] != 0:
        return
    if m % a_scale_group_shape[0] != 0 or k % a_scale_group_shape[1] != 0:
        return
    cutlass_fp8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape,
                            use_bias)
167
168


169
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
170
171
172
173
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
174
@pytest.mark.parametrize("use_bias", [True, False])
175
176
177
178
def test_cutlass_int8_gemm(m: int, n: int, k: int, a_scale_group_shape,
                           b_scale_group_shape, use_bias: bool):
    cutlass_int8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape,
                             use_bias)
179
180


181
182
183
184
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
185
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
186
@pytest.mark.parametrize("use_bias", [True, False])
187
188
def test_cutlass_int8_gemm_output_dtype(a_scale_group_shape,
                                        b_scale_group_shape,
189
                                        out_dtype: Type[torch.dtype],
190
                                        use_bias: bool):
191
192
193
    cutlass_int8_gemm_helper(512,
                             512,
                             512,
194
195
                             a_scale_group_shape,
                             b_scale_group_shape,
196
                             use_bias,
197
                             out_dtype=out_dtype)
198
199


200
201
202
203
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
204
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
205
@pytest.mark.parametrize("use_bias", [True, False])
206
@pytest.mark.skipif(not current_platform.has_device_capability(89),
207
                    reason="FP8 is not supported on this GPU type.")
208
209
def test_cutlass_fp8_gemm_output_dtype(a_scale_group_shape,
                                       b_scale_group_shape,
210
                                       out_dtype: Type[torch.dtype],
211
                                       use_bias: bool):
212
213
214
    cutlass_fp8_gemm_helper(512,
                            512,
                            512,
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
                            a_scale_group_shape,
                            b_scale_group_shape,
                            use_bias,
                            out_dtype=out_dtype)


@pytest.mark.parametrize("a_scale_group_shape,b_scale_group_shape",
                         [((1, 128), (128, 128))])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("use_bias", [False])
@pytest.mark.skipif(not current_platform.has_device_capability(90),
                    reason="FP8 blockwise is not supported on this GPU type.")
def test_cutlass_fp8_blockwise_scale_gemm_dtype(a_scale_group_shape,
                                                b_scale_group_shape,
                                                out_dtype: Type[torch.dtype],
                                                use_bias: bool):
    cutlass_fp8_gemm_helper(512,
                            512,
                            512,
                            a_scale_group_shape,
                            b_scale_group_shape,
236
                            use_bias,
237
                            out_dtype=out_dtype)
238
239


240
241
242
243
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
244
@pytest.mark.parametrize("use_bias", [True, False])
245
@pytest.mark.parametrize("device", CUDA_DEVICES)
246
@pytest.mark.skipif(not current_platform.has_device_capability(89),
247
                    reason="FP8 is not supported on this GPU type.")
248
def test_cutlass_fp8_gemm_devices(a_scale_group_shape, b_scale_group_shape,
249
                                  use_bias: bool, device: str):
250
251
252
    cutlass_fp8_gemm_helper(512, 512, 512, a_scale_group_shape,
                            b_scale_group_shape, use_bias, torch.bfloat16,
                            device)
253
254


255
256
257
258
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
259
@pytest.mark.parametrize("use_bias", [True, False])
260
@pytest.mark.parametrize("device", CUDA_DEVICES)
261
def test_cutlass_int8_gemm_devices(a_scale_group_shape, b_scale_group_shape,
262
                                   use_bias: bool, device: str):
263
264
265
    cutlass_int8_gemm_helper(512,
                             512,
                             512,
266
267
                             a_scale_group_shape,
                             b_scale_group_shape,
268
                             use_bias,
269
270
                             out_dtype=torch.bfloat16,
                             device=device)
271
272
273
274
275
276
277


# For the following two tests:
# N and K correspond to the size of the weight matrix and likely to be multiples
# of a large power of two. In any case, the kernel will have a naive fallback
# when N and K are not divisible by 16. But M is the number of tokens and the
# kernel must handle any M thrown at it.
278
279
280
281
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
282
@pytest.mark.parametrize("use_bias", [True, False])
283
@pytest.mark.skipif(not current_platform.has_device_capability(89),
284
                    reason="FP8 is not supported on this GPU type.")
285
def test_cutlass_fp8_gemm_m_sweep(a_scale_group_shape, b_scale_group_shape,
286
                                  use_bias: bool):
287
288
    for nk in range(32, 128, 32):
        for m in range(1, 128):
289
290
            cutlass_fp8_gemm_helper(m, nk, nk, a_scale_group_shape,
                                    b_scale_group_shape, use_bias)
291
292


293
294
295
296
@pytest.mark.parametrize("a_scale_group_shape",
                         [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
@pytest.mark.parametrize("b_scale_group_shape",
                         [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE])
297
@pytest.mark.parametrize("use_bias", [True, False])
298
def test_cutlass_int8_gemm_m_sweep(a_scale_group_shape, b_scale_group_shape,
299
                                   use_bias: bool):
300
301
    for nk in range(32, 128, 32):
        for m in range(1, 128):
302
303
            cutlass_int8_gemm_helper(m, nk, nk, a_scale_group_shape,
                                     b_scale_group_shape, use_bias)
304
305


306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
@pytest.mark.parametrize("m", [32, 64, 128])
@pytest.mark.parametrize("n", [16, 32, 64])
@pytest.mark.parametrize("k", [64, 128, 256])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
@pytest.mark.skip
def test_cutlass_int8_azp_bias_fold(m: int, n: int, k: int,
                                    out_dtype: torch.dtype):
    # Currently, the test is failing because folding azp into
    # 16-bit bias loses too much precision
    scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
    scale_b = torch.randn((1, n), device="cuda", dtype=torch.float32) / 10

    aq_i8 = rand_int8((m, k))
    bq_i8 = rand_int8((n, k)).t()

    aq_i32 = aq_i8.to(dtype=torch.int32)
    bq_i32 = bq_i8.to(dtype=torch.int32)

    aq_f32 = aq_i8.to(dtype=torch.float32)
    bq_f32 = bq_i8.to(dtype=torch.float32)

    b_dq = scale_b * bq_f32

    azp_a = torch.rand((1, ), device="cuda", dtype=torch.float32) * 10 + 1.5
    azp_aq_i8 = (azp_a / scale_a).to(dtype=torch.int8)
    azp_a = azp_aq_i8.to(dtype=torch.float32) * scale_a  # correct for rounding

    a_dq = scale_a * (aq_i32 + azp_aq_i8).to(dtype=torch.float32)
334
    torch.testing.assert_close(a_dq, scale_a * aq_f32 + azp_a)
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352

    baseline_dq = torch.mm(a_dq, b_dq).to(out_dtype)

    J = torch.ones((1, k), device="cuda", dtype=torch.float32)
    azp_bias = (azp_a * scale_b * (J @ bq_f32)).to(out_dtype)
    assert azp_bias.shape == (1, n)
    assert azp_bias[0, :].shape == (n, )

    baseline_q = (scale_a.to(device='cpu') * scale_b.to(device='cpu') * (
        (aq_i32 + azp_aq_i8).to(device='cpu') @ bq_i32.to(device='cpu'))).to(
            dtype=out_dtype, device='cuda')

    out = ops.cutlass_scaled_mm(aq_i8,
                                bq_i8,
                                scale_a,
                                scale_b,
                                out_dtype=out_dtype,
                                bias=azp_bias[0, :])
353
354
    torch.testing.assert_close(out, baseline_dq, rtol=1e-2, atol=1e0)
    torch.testing.assert_close(out, baseline_q, rtol=1e-2, atol=1e0)
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383


@pytest.mark.parametrize("m", [32, 64, 128])
@pytest.mark.parametrize("n", [16, 32, 64])
@pytest.mark.parametrize("k", [64, 128, 256])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("use_bias", [True, False])
@pytest.mark.parametrize("azp_per_token", [True, False])
def test_cutlass_int8_azp(m: int, n: int, k: int, out_dtype: torch.dtype,
                          use_bias: bool, azp_per_token: bool):
    m_azp = m if azp_per_token else 1
    scale_a = torch.randn((m_azp, 1), device="cuda", dtype=torch.float32) / 10
    scale_b = torch.randn((1, n), device="cuda", dtype=torch.float32) / 10

    aq_i8 = rand_int8((m, k))
    aq_i32 = aq_i8.to(dtype=torch.int32)
    aq_f32 = aq_i8.to(dtype=torch.float32)

    bq_i8 = rand_int8((n, k)).t()
    bq_i32 = bq_i8.to(dtype=torch.int32)
    bq_f32 = bq_i8.to(dtype=torch.float32)
    b_dq = scale_b * bq_f32

    azp_a = torch.rand(
        (m_azp, 1), device="cuda", dtype=torch.float32) * 10 + 1.5
    azp_aq_i8 = (azp_a / scale_a).to(dtype=torch.int8)
    azp_a = azp_aq_i8.to(dtype=torch.float32) * scale_a  # correct for rounding

    a_dq = scale_a * (aq_i32 - azp_aq_i8).to(dtype=torch.float32)
384
385
386
387
    torch.testing.assert_close(a_dq,
                               scale_a * aq_f32 - azp_a,
                               rtol=1e-4,
                               atol=1e-3)
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419

    if use_bias:
        bias = torch.rand((1, n), device="cuda", dtype=out_dtype) * 10 + 2.5
    else:
        bias = torch.zeros((1, n), device="cuda", dtype=out_dtype)

    baseline_dq = (torch.mm(a_dq, b_dq) + bias).to(out_dtype)

    # int32 mm not supported on CUDA
    a_noazp_i32_cpu = (aq_i32 - azp_aq_i8).to(device='cpu')
    cq = (a_noazp_i32_cpu @ bq_i32.to(device='cpu')).to(device='cuda')
    baseline_q = (scale_a * scale_b * cq + bias).to(dtype=out_dtype)

    # Hadamard is just the sum of the cols
    azp_adj_i32 = bq_i32.sum(dim=0, keepdim=True, dtype=torch.int32)
    azp_i32 = azp_aq_i8.to(dtype=torch.int32)
    func_bias = bias if use_bias else None

    if azp_per_token:
        out = ops.cutlass_scaled_mm_azp(aq_i8, bq_i8, scale_a, scale_b,
                                        out_dtype, azp_adj_i32, azp_i32,
                                        func_bias)
    else:
        azp_with_adj_i32 = azp_i32 * azp_adj_i32
        out = ops.cutlass_scaled_mm_azp(aq_i8, bq_i8, scale_a, scale_b,
                                        out_dtype, azp_with_adj_i32, None,
                                        func_bias)

    # bfloat16 precision is 7-bit mantissa -> 2^-8 ~ 0.4%
    # float16 precision is 10-bit mantissa -> 2^-11 ~ 0.05%
    rtol = 1e-2 if out_dtype == torch.bfloat16 else 1e-3
    atol = 1e-3
420
421
    torch.testing.assert_close(out, baseline_dq, rtol=rtol, atol=atol)
    torch.testing.assert_close(out, baseline_q, rtol=rtol, atol=atol)
422

423
424
425
426
427
428
429
430
431
    if azp_per_token:
        opcheck(torch.ops._C.cutlass_scaled_mm_azp,
                (out, aq_i8, bq_i8, scale_a, scale_b, azp_adj_i32, azp_i32,
                 func_bias))
    else:
        opcheck(torch.ops._C.cutlass_scaled_mm_azp,
                (out, aq_i8, bq_i8, scale_a, scale_b, azp_with_adj_i32, None,
                 func_bias))

432

433
434
435
436
437
438
439
440
441
442
443
444
445
# Test working with a subset of A and B
def test_cutlass_subset():
    big_m, big_n, big_k = 1024, 1024, 1024
    m, n, k = 512, 512, 512

    whole_a = to_int8(torch.randn((big_m, big_k), device="cuda") * 5)
    whole_b = to_int8(torch.randn((big_n, big_k), device="cuda").t() * 5)
    a = whole_a[0:m, 0:k]
    b = whole_b[0:k, 0:n]

    scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
    scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10

446
447
448
449
450
    out = ops.cutlass_scaled_mm(a,
                                b,
                                scale_a,
                                scale_b,
                                out_dtype=torch.bfloat16)
451
452
453
454
455
    baseline = baseline_scaled_mm(a,
                                  b,
                                  scale_a,
                                  scale_b,
                                  out_dtype=torch.bfloat16)
456

457
    torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
458
459
460
461
462
463
464
465
466
467
468
469
470


# Test to make sure cuda graphs work
class CutlassLayer(torch.nn.Module):

    def __init__(self, b, scale_a, scale_b, out_dtype):
        super().__init__()
        self.b = b
        self.scale_a = scale_a
        self.scale_b = scale_b
        self.out_dtype = out_dtype

    def forward(self, a):
471
472
        return ops.cutlass_scaled_mm(a, self.b, self.scale_a, self.scale_b,
                                     self.out_dtype)
473
474


475
476
477
@pytest.mark.parametrize("per_act_token", [True, False])
@pytest.mark.parametrize("per_out_ch", [True, False])
def test_cutlass_cuda_graph(per_act_token: bool, per_out_ch: bool):
478
479
480
481
482
    m, n, k = 512, 512, 512

    a = to_int8(torch.randn((m, k), device="cuda"))
    b = to_int8(torch.randn((n, k), device="cuda").t())

483
484
485
486
487
488
489
    m_a_scales = m if per_act_token else 1
    n_b_scales = n if per_out_ch else 1

    scale_a = (torch.randn(
        (m_a_scales, 1), device="cuda", dtype=torch.float32) / 10)
    scale_b = (torch.randn(
        (1, n_b_scales), device="cuda", dtype=torch.float32) / 10)
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504

    # Construct a trivial model with a single layer that calls a CUTLASS kernel
    model = CutlassLayer(b, scale_a, scale_b, torch.bfloat16)

    # Run the model with a cuda graph
    stream = torch.cuda.Stream()
    with torch.cuda.stream(stream):
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            out = model(a)
    out.zero_()
    g.replay()

    baseline = torch.mm(scale_a * a.to(dtype=torch.float32),
                        scale_b * b.to(dtype=torch.float32)).to(torch.bfloat16)
505
    torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
506
507
508
509


def test_cutlass_support_opcheck():
    opcheck(torch.ops._C.cutlass_scaled_mm_supports_fp8, (capability, ))