Function test_fp8_block_gemm

test/kernel/deepseekv3_fp8_block_gemm_tuning.py:29–71

(
    M: int,
    N: int,
    K: int,
    block_size: int,
    dtype: torch.dtype,
    test_count: int = 20,
    **run_config,
)

@torch.no_grad()
def test_fp8_block_gemm(
    M: int,
    N: int,
    K: int,
    block_size: int,
    dtype: torch.dtype,
    test_count: int = 20,
    **run_config,
):
    set_seed()

    input_tuples = []
    for _ in range(test_count):
        A = torch.randn((M, K), dtype=torch.float32).cuda().to(torch.float8_e4m3fn)  # Activation
        B = torch.randn((K, N), dtype=torch.float32).cuda().to(torch.float8_e4m3fn)  # Weight
        Ascale = torch.ones((M, (K + block_size - 1) // block_size)).cuda()
        Bscale = torch.ones(((K + block_size - 1) // block_size, (N + block_size - 1) // block_size)).cuda()
        C = torch.randn((M, N), dtype=dtype).cuda()  # weight
        input_tuples.append((A, B, Ascale, Bscale, C))
    w8a8_block_fp8_matmul(A, B, Ascale, Bscale, C, (block_size, block_size), dtype, run_config=run_config)

    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for index in range(test_count):
            A, B, Ascale, Bscale, C = input_tuples[index]
            w8a8_block_fp8_matmul(
                A,
                B,
                Ascale,
                Bscale,
                C,
                (block_size, block_size),
                run_config=run_config,
            )

    graph.replay()
    torch.cuda.synchronize()
    start = time.time()
    graph.replay()
    torch.cuda.synchronize()
    cost_time = (time.time() - start) * 1000
    logger.info(f"fp8 mm {M} {N} {K} block {block_size} cost time: {cost_time} ms")
    return cost_time


def worker(