Function test_kernel

test/kernel/moe_sum_reduce_tuning_bf16.py:30–68

(
    m: int,
    topk_num: int,
    hidden_dim: int,
    dtype: torch.dtype,
    test_count: int,
    **config,
)

@torch.no_grad()
def test_kernel(
    m: int,
    topk_num: int,
    hidden_dim: int,
    dtype: torch.dtype,
    test_count: int,
    **config,
):
    set_seed()
    input_tuples = []

    input = torch.randn((m, topk_num, hidden_dim), device="cuda", dtype=dtype) / 10
    output = torch.randn((m, hidden_dim), device="cuda", dtype=dtype)

    for _ in range(test_count):
        input_tuples.append((input.clone(), output.clone()))

    # warm_up
    moe_sum_reduce(input, output, run_config=config)

    graph = torch.cuda.CUDAGraph()

    with torch.cuda.graph(graph):
        for index in range(test_count):
            input, output = input_tuples[index]
            moe_sum_reduce(input, output, run_config=config)

    graph.replay()

    torch.cuda.synchronize()
    start = time.time()
    graph.replay()
    torch.cuda.synchronize()

    cost_time = (time.time() - start) * 1000

    logger.info(str(config))
    logger.info(f"bf16 {m} cost time: {cost_time} ms")
    return cost_time


def worker(