```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 large_k_mm(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_runtime.h>
#include <mma.h>
#include <algorithm>

namespace wmma = nvcuda::wmma;

__device__ __forceinline__ float round_tf32(float x) {
#if __CUDA_ARCH__ >= 800
    unsigned int bits;
    asm("cvt.rn.tf32.f32 %0, %1;" : "=r"(bits) : "f"(x));
    return __uint_as_float(bits);
#else
    return x;
#endif
}

__global__ void split_k_mm(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int K,
    int tiles_per_split
) {
    constexpr int M = 256;
    constexpr int N = 256;
    constexpr int BM = 64;
    constexpr int BN = 64;
    constexpr int BK = 32;

    const int row_base = blockIdx.y * BM;
    const int col_base = blockIdx.x * BN;
    const int tile_begin = blockIdx.z * tiles_per_split;
    const int tile_end = min(tile_begin + tiles_per_split, K / BK);
    float* output = partial + static_cast<size_t>(blockIdx.z) * M * N;

#if __CUDA_ARCH__ >= 800
    __shared__ __align__(32) float shared_A[BM * BK];
    __shared__ __align__(32) float shared_B[BK * BN];

    const int warp = threadIdx.x / 32;
    const int warp_row = (warp / 2) * 32;
    const int warp_col = (warp % 2) * 32;

    wmma::fragment<
        wmma::accumulator, 16, 16, 8, float
    > accum[2][2];

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

    for (int tile = tile_begin; tile < tile_end; ++tile) {
        const int k_base = tile * BK;

        #pragma unroll
        for (int index = threadIdx.x; index < BM * BK; index += 128) {
            const int row = index / BK;
            const int k = index % BK;
            shared_A[index] = round_tf32(
                A[static_cast<size_t>(row_base + row) * K + k_base + k]
            );
        }

        #pragma unroll
        for (int index = threadIdx.x; index < BK * BN; index += 128) {
            const int k = index / BN;
            const int col = index % BN;
            shared_B[index] = round_tf32(
                B[static_cast<size_t>(k_base + k) * N + col_base + col]
            );
        }

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BK; k += 8) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major
            > a_frag[2];

            wmma::fragment<
                wmma::matrix_b, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major
            > b_frag[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                wmma::load_matrix_sync(
                    a_frag[i],
                    shared_A + (warp_row + i * 16) * BK + k,
                    BK
                );
            }

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::load_matrix_sync(
                    b_frag[j],
                    shared_B + k * BN + warp_col + j * 16,
                    BN
                );
            }

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(
                        accum[i][j], a_frag[i], b_frag[j], accum[i][j]
                    );
                }
            }
        }

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::store_matrix_sync(
                output
                    + (row_base + warp_row + i * 16) * N
                    + col_base + warp_col + j * 16,
                accum[i][j],
                N,
                wmma::mem_row_major
            );
        }
    }
#else
    // FP32 fallback for devices without TF32 tensor cores.
    for (int index = threadIdx.x; index < BM * BN; index += blockDim.x) {
        const int row = row_base + index / BN;
        const int col = col_base + index % BN;
        float sum = 0.0f;

        for (int k = tile_begin * BK; k < tile_end * BK; ++k) {
            sum = fmaf(
                A[static_cast<size_t>(row) * K + k],
                B[static_cast<size_t>(k) * N + col],
                sum
            );
        }

        output[row * N + col] = sum;
    }
#endif
}

__global__ void reduce_splits(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int splits
) {
    constexpr int elements = 256 * 256;
    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= elements) {
        return;
    }

    float sum = 0.0f;
    for (int split = 0; split < splits; ++split) {
        sum += partial[static_cast<size_t>(split) * elements + index];
    }
    output[index] = sum;
}

torch::Tensor large_k_mm(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "CUDA tensors required");
    TORCH_CHECK(A.device() == B.device(), "Devices must match");
    TORCH_CHECK(
        A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
        "Float32 tensors required"
    );
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(), "Contiguous tensors required");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Matrices required");
    TORCH_CHECK(
        A.size(0) == 256 && B.size(1) == 256 && A.size(1) == B.size(0),
        "Expected (256, K) and (K, 256)"
    );
    TORCH_CHECK(
        A.size(1) > 0 && A.size(1) % 32 == 0 && A.size(1) <= 2147483647,
        "K must be a positive multiple of 32 fitting int32"
    );

    const c10::cuda::CUDAGuard guard(A.device());
    const int K = static_cast<int>(A.size(1));
    const int tiles = K / 32;
    const int splits = std::min(128, (tiles + 127) / 128);
    const int tiles_per_split = (tiles + splits - 1) / splits;

    auto partial = torch::empty({splits, 256, 256}, A.options());
    auto output = torch::empty({256, 256}, A.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    split_k_mm<<<dim3(4, 4, splits), 128, 0, stream>>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        partial.data_ptr<float>(),
        K,
        tiles_per_split
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    reduce_splits<<<256, 256, 0, stream>>>(
        partial.data_ptr<float>(),
        output.data_ptr<float>(),
        splits
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="large_k_split_tf32_mm_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["large_k_mm"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


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

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            A.is_cuda
            and B.is_cuda
            and A.device == B.device
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.ndim == 2
            and B.ndim == 2
            and A.shape[0] == 256
            and B.shape[1] == 256
            and A.shape[1] == B.shape[0]
            and 0 < A.shape[1] <= 2147483647
            and A.shape[1] % 32 == 0
            and A.is_contiguous()
            and B.is_contiguous()
            and not (torch.is_grad_enabled() and (A.requires_grad or B.requires_grad))
        ):
            return _extension.large_k_mm(A, B)
        return torch.matmul(A, B)
```