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

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

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <algorithm>
#include <climits>

static void check_blas(cublasStatus_t status) {
    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
                "cuBLAS error: ", static_cast<int>(status));
}

// Retain two BF16 components per FP32 value.
__global__ void split_inputs(
    const float* __restrict__ a,
    const float* __restrict__ b,
    __nv_bfloat16* __restrict__ ah,
    __nv_bfloat16* __restrict__ al,
    __nv_bfloat16* __restrict__ bh,
    __nv_bfloat16* __restrict__ bl,
    int64_t na,
    int64_t nb
) {
    const int64_t count = na + nb;
    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < count;
         i += int64_t(blockDim.x) * gridDim.x) {
        if (i < na) {
            const float x = a[i];
            const __nv_bfloat16 hi = __float2bfloat16_rn(x);
            ah[i] = hi;
            al[i] = __float2bfloat16_rn(x - __bfloat162float(hi));
        } else {
            const int64_t j = i - na;
            const float x = b[j];
            const __nv_bfloat16 hi = __float2bfloat16_rn(x);
            bh[j] = hi;
            bl[j] = __float2bfloat16_rn(x - __bfloat162float(hi));
        }
    }
}

struct MathModeGuard {
    cublasHandle_t handle;
    cublasMath_t previous;

    explicit MathModeGuard(cublasHandle_t h) : handle(h) {
        check_blas(cublasGetMathMode(handle, &previous));
        check_blas(cublasSetMathMode(
            handle,
            static_cast<cublasMath_t>(
                CUBLAS_DEFAULT_MATH |
                CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION)));
    }

    ~MathModeGuard() {
        cublasSetMathMode(handle, previous);
    }
};

torch::Tensor split_gemm(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2,
                "Expected two matrices");
    TORCH_CHECK(A.size(1) == B.size(0),
                "Incompatible matrix dimensions");
    TORCH_CHECK(A.device() == B.device(),
                "Matrices must be on the same device");

    if (!A.is_cuda() ||
        A.scalar_type() != at::kFloat ||
        B.scalar_type() != at::kFloat ||
        A.requires_grad() || B.requires_grad()) {
        return at::matmul(A, B);
    }

    const c10::cuda::CUDAGuard device_guard(A.device());

    const int64_t m64 = A.size(0);
    const int64_t k64 = A.size(1);
    const int64_t n64 = B.size(1);

    if (at::cuda::getCurrentDeviceProperties()->major < 8 ||
        m64 < 128 || n64 < 128 || k64 < 128 ||
        m64 > INT_MAX || n64 > INT_MAX || k64 > INT_MAX) {
        return at::matmul(A, B);
    }

    auto a = A.contiguous();
    auto b = B.contiguous();
    auto out = torch::empty({m64, n64}, A.options());

    const int64_t na = a.numel();
    const int64_t nb = b.numel();
    auto packed = torch::empty(
        {2 * (na + nb)}, A.options().dtype(at::kBFloat16));

    auto* ah = reinterpret_cast<__nv_bfloat16*>(
        packed.data_ptr<at::BFloat16>());
    auto* al = ah + na;
    auto* bh = al + na;
    auto* bl = bh + nb;

    const auto stream = at::cuda::getCurrentCUDAStream();
    const int blocks = static_cast<int>(
        std::min<int64_t>((na + nb + 255) / 256, 8192));

    split_inputs<<<blocks, 256, 0, stream.stream()>>>(
        a.data_ptr<float>(), b.data_ptr<float>(),
        ah, al, bh, bl, na, nb);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const int m = static_cast<int>(m64);
    const int n = static_cast<int>(n64);
    const int k = static_cast<int>(k64);

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    MathModeGuard math_guard(handle);

    const float one = 1.0f;
    const float zero = 0.0f;
    float* c = out.data_ptr<float>();

    // Row-major A @ B is column-major B^T @ A^T.
    auto gemm = [&](const __nv_bfloat16* left,
                    const __nv_bfloat16* right,
                    const float* beta) {
        check_blas(cublasGemmEx(
            handle,
            CUBLAS_OP_N, CUBLAS_OP_N,
            n, m, k,
            &one,
            left, CUDA_R_16BF, n,
            right, CUDA_R_16BF, k,
            beta,
            c, CUDA_R_32F, n,
            CUBLAS_COMPUTE_32F,
            CUBLAS_GEMM_DEFAULT_TENSOR_OP));
    };

    // Accumulate correction products before the dominant product.
    // The omitted low-low product is second order in BF16 rounding.
    gemm(bh, al, &zero);
    gemm(bl, ah, &one);
    gemm(bh, ah, &one);

    return out;
}
"""

_split_gemm = load_inline(
    name="split_bf16_matmul_ext",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["split_gemm"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        return _split_gemm.split_gemm(A, B)
```