```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_source = r"""
#include <torch/extension.h>
torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <mma.h>

namespace {

constexpr int BM = 128;
constexpr int BN = 64;
constexpr int BK = 32;
constexpr int THREADS = 256;

__global__ void matmul_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
#if __CUDA_ARCH__ >= 800
    namespace wmma = nvcuda::wmma;

    __shared__ __align__(32) __nv_bfloat16 ah[BM * BK];
    __shared__ __align__(32) __nv_bfloat16 al[BM * BK];
    __shared__ __align__(32) __nv_bfloat16 bh[BK * BN];
    __shared__ __align__(32) __nv_bfloat16 bl[BK * BN];
    __shared__ __align__(32) float edge[8 * 16 * 16];

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int lane = tid % 32;
    const int row_base = blockIdx.y * BM;
    const int col_base = blockIdx.x * BN;
    const int warp_row = (warp / 2) * 32;
    const int warp_col = (warp % 2) * 32;

    using AFrag = wmma::fragment<
        wmma::matrix_a, 16, 16, 16, __nv_bfloat16, wmma::row_major>;
    using BFrag = wmma::fragment<
        wmma::matrix_b, 16, 16, 16, __nv_bfloat16, wmma::row_major>;
    using AccFrag = wmma::fragment<
        wmma::accumulator, 16, 16, 16, float>;

    AccFrag acc[2][2];
    AccFrag correction[2][2];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::fill_fragment(acc[i][j], 0.0f);
            wmma::fill_fragment(correction[i][j], 0.0f);
        }
    }

    for (int kb = 0; kb < K; kb += BK) {
        #pragma unroll
        for (int p = tid; p < BM * BK; p += THREADS) {
            const int r = row_base + p / BK;
            const int k = kb + p % BK;
            const float x = (r < M && k < K)
                ? A[static_cast<int64_t>(r) * K + k] : 0.0f;
            const __nv_bfloat16 hi = __float2bfloat16_rn(x);
            ah[p] = hi;
            al[p] = __float2bfloat16_rn(x - __bfloat162float(hi));
        }

        #pragma unroll
        for (int p = tid; p < BK * BN; p += THREADS) {
            const int k = kb + p / BN;
            const int c = col_base + p % BN;
            const float x = (k < K && c < N)
                ? B[static_cast<int64_t>(k) * N + c] : 0.0f;
            const __nv_bfloat16 hi = __float2bfloat16_rn(x);
            bh[p] = hi;
            bl[p] = __float2bfloat16_rn(x - __bfloat162float(hi));
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 16) {
            AFrag a_hi[2], a_lo[2];
            BFrag b_hi[2], b_lo[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int offset = (warp_row + i * 16) * BK + kk;
                wmma::load_matrix_sync(a_hi[i], ah + offset, BK);
                wmma::load_matrix_sync(a_lo[i], al + offset, BK);
            }

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const int offset = kk * BN + warp_col + j * 16;
                wmma::load_matrix_sync(b_hi[j], bh + offset, BN);
                wmma::load_matrix_sync(b_lo[j], bl + offset, BN);
            }

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(
                        acc[i][j], a_hi[i], b_hi[j], acc[i][j]);
                    wmma::mma_sync(
                        correction[i][j], a_lo[i], b_hi[j],
                        correction[i][j]);
                    wmma::mma_sync(
                        correction[i][j], a_hi[i], b_lo[j],
                        correction[i][j]);
                }
            }
        }

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int e = 0; e < acc[i][j].num_elements; ++e) {
                acc[i][j].x[e] += correction[i][j].x[e];
            }

            const int r = row_base + warp_row + i * 16;
            const int c = col_base + warp_col + j * 16;

            if (r + 16 <= M && c + 16 <= N && N % 4 == 0) {
                wmma::store_matrix_sync(
                    C + static_cast<int64_t>(r) * N + c,
                    acc[i][j], N, wmma::mem_row_major);
            } else {
                float* tile = edge + warp * 256;
                wmma::store_matrix_sync(
                    tile, acc[i][j], 16, wmma::mem_row_major);
                __syncwarp();

                for (int p = lane; p < 256; p += 32) {
                    const int rr = r + p / 16;
                    const int cc = c + p % 16;
                    if (rr < M && cc < N) {
                        C[static_cast<int64_t>(rr) * N + cc] = tile[p];
                    }
                }
                __syncwarp();
            }
        }
    }
#endif
}

} // namespace

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Device mismatch");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 &&
                B.scalar_type() == torch::kFloat32,
                "Expected float32 tensors");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Expected matrices");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible dimensions");

    c10::cuda::CUDAGuard guard(A.device());
    A = A.contiguous();
    B = B.contiguous();

    const int M = static_cast<int>(A.size(0));
    const int K = static_cast<int>(A.size(1));
    const int N = static_cast<int>(B.size(1));

    auto C = torch::empty({M, N}, A.options());
    if (M == 0 || N == 0) {
        return C;
    }

    const dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
    const auto stream = c10::cuda::getCurrentCUDAStream(A.get_device());

    matmul_kernel<<<grid, THREADS, 0, stream.stream()>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
        M, N, K);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="compensated_bf16_matmul_v1",
            cpp_sources=_cpp_source,
            cuda_sources=_cuda_source,
            functions=["matmul_cuda"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            verbose=False,
        )
    return _extension


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self._capabilities = {}

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            not A.is_cuda
            or not B.is_cuda
            or A.device != B.device
            or A.dtype != torch.float32
            or B.dtype != torch.float32
            or A.ndim != 2
            or B.ndim != 2
            or A.requires_grad
            or B.requires_grad
        ):
            return torch.matmul(A, B)

        device = A.device.index
        if device not in self._capabilities:
            self._capabilities[device] = torch.cuda.get_device_capability(device)

        if self._capabilities[device][0] < 8:
            return torch.matmul(A, B)

        return _get_extension().matmul_cuda(A, B)
```