def METHOD_NAME(self):
    device = torch.device("cuda")
    dtype = torch.float32
    batch, head_dim, head_num, hidden = 128, 64, 8, 768
    x = torch.randn(batch, head_dim, head_num, hidden, device=device, dtype=dtype)
    dst_dim = 1024
    linear = Linear2D(hidden, dst_dim)
    linear.to(device).to(dtype)
    old_linear = nn.Linear(hidden, dst_dim)
    old_linear.load_state_dict(linear.state_dict())
    old_linear.to(device).to(dtype)
    for i in range(5):  # warmup
        y = linear(x)
        y1 = old_linear(x)
    torch.cuda.synchronize()
    start_event = torch.cuda.Event(enable_timing=True)
    start_event2 = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    end_event2 = torch.cuda.Event(enable_timing=True)
    with torch.profiler.profile(
        activities=[
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.CUDA,
        ],
        with_stack=True,
    ) as prof:
        start_event.record()
        for i in range(100):
            y = linear(x)
        end_event.record()
        torch.cuda.synchronize()
        fused_time = start_event.elapsed_time(end_event)
        start_event2.record()
        for i in range(100):
            y1 = old_linear(x)
        end_event2.record()
        torch.cuda.synchronize()
        origin_time = start_event2.elapsed_time(end_event2)
        torch.cuda.synchronize()
        Y = torch.randn_like(y)
        loss = torch.nn.functional.mse_loss(y, Y)
        loss1 = torch.nn.functional.mse_loss(y1, Y)
        loss.backward()
        loss1.backward()
        torch.testing.assert_close(y, y1, rtol=1e-4, atol=1e-4)
        name2grad = {}
        for name, p in linear.named_parameters():
            name2grad[name] = p.grad
        for name, p in old_linear.named_parameters():
            torch.testing.assert_close(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)
        self.assertLess(fused_time, origin_time)
        print("fused_time", fused_time, "origin_time", origin_time)
    prof.export_chrome_trace("test_2d_linear_gpu.json")