```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Custom CUDA kernel: C = A @ B^T using TF32 tensor cores via the WMMA API.
# A is (M, K) row-major, B is (N, K) row-major, C is (M, N) row-major.
# Both operands are streamed through shared memory with fully coalesced
# float4 global loads; each warp computes a 64x64 output tile with
# 16x16x8 TF32 MMA fragments (FP32 accumulation).
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 32
#define PAD 8
#define LDS (BK + PAD)
#define WM 64
#define WN 64
#define FRAGS_M (WM / 16)
#define FRAGS_N (WN / 16)
#define NWARPS 4
#define THREADS (NWARPS * 32)

__global__ void __launch_bounds__(THREADS)
gemm_nt_tf32_kernel(const float* __restrict__ A,
                    const float* __restrict__ B,
                    float* __restrict__ C,
                    int M, int N, int K)
{
    __shared__ float As[BM][LDS];
    __shared__ float Bs[BN][LDS];

    const int block_m = blockIdx.y * BM;
    const int block_n = blockIdx.x * BN;
    const int tid = threadIdx.x;

#if __CUDA_ARCH__ >= 800
    const int warp_id = tid >> 5;
    const int lane = tid & 31;
    const int warp_m = (warp_id >> 1) * WM;
    const int warp_n = (warp_id & 1) * WN;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[FRAGS_M][FRAGS_N];
    #pragma unroll
    for (int i = 0; i < FRAGS_M; ++i)
        #pragma unroll
        for (int j = 0; j < FRAGS_N; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    const bool vec4 = ((K & 3) == 0);

    for (int k0 = 0; k0 < K; k0 += BK) {
        // ---- Stage A tile (BM x BK) into shared memory ----
        #pragma unroll
        for (int t = 0; t < (BM * BK / 4) / THREADS; ++t) {
            const int idx = tid + t * THREADS;
            const int row = idx / (BK / 4);
            const int col = (idx % (BK / 4)) * 4;
            const int gr = block_m + row;
            const int gk = k0 + col;
            float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
            if (vec4 && gr < M && (gk + 3) < K) {
                v = *reinterpret_cast<const float4*>(A + (size_t)gr * K + gk);
            } else if (gr < M) {
                float* pv = reinterpret_cast<float*>(&v);
                #pragma unroll
                for (int u = 0; u < 4; ++u)
                    if (gk + u < K) pv[u] = A[(size_t)gr * K + gk + u];
            }
            *reinterpret_cast<float4*>(&As[row][col]) = v;
        }
        // ---- Stage B tile (BN x BK) into shared memory ----
        #pragma unroll
        for (int t = 0; t < (BN * BK / 4) / THREADS; ++t) {
            const int idx = tid + t * THREADS;
            const int row = idx / (BK / 4);
            const int col = (idx % (BK / 4)) * 4;
            const int gr = block_n + row;
            const int gk = k0 + col;
            float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
            if (vec4 && gr < N && (gk + 3) < K) {
                v = *reinterpret_cast<const float4*>(B + (size_t)gr * K + gk);
            } else if (gr < N) {
                float* pv = reinterpret_cast<float*>(&v);
                #pragma unroll
                for (int u = 0; u < 4; ++u)
                    if (gk + u < K) pv[u] = B[(size_t)gr * K + gk + u];
            }
            *reinterpret_cast<float4*>(&Bs[row][col]) = v;
        }
        __syncthreads();

        // ---- Tensor-core MMA on the staged tiles ----
        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag[FRAGS_M];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::col_major> b_frag[FRAGS_N];
            #pragma unroll
            for (int i = 0; i < FRAGS_M; ++i)
                wmma::load_matrix_sync(a_frag[i], &As[warp_m + i * 16][kk], LDS);
            #pragma unroll
            for (int j = 0; j < FRAGS_N; ++j)
                wmma::load_matrix_sync(b_frag[j], &Bs[warp_n + j * 16][kk], LDS);
            #pragma unroll
            for (int i = 0; i < FRAGS_M; ++i)
                #pragma unroll
                for (int j = 0; j < FRAGS_N; ++j)
                    wmma::mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]);
        }
        __syncthreads();
    }

    // ---- Write results ----
    if (block_m + BM <= M && block_n + BN <= N) {
        #pragma unroll
        for (int i = 0; i < FRAGS_M; ++i)
            #pragma unroll
            for (int j = 0; j < FRAGS_N; ++j)
                wmma::store_matrix_sync(
                    &C[(size_t)(block_m + warp_m + i * 16) * N + block_n + warp_n + j * 16],
                    acc[i][j], N, wmma::mem_row_major);
    } else {
        // Edge tiles: stage through shared memory and store with bounds checks.
        float* stage = &As[0][0];
        for (int w = 0; w < NWARPS; ++w) {
            if (warp_id == w) {
                #pragma unroll
                for (int i = 0; i < FRAGS_M; ++i)
                    #pragma unroll
                    for (int j = 0; j < FRAGS_N; ++j)
                        wmma::store_matrix_sync(stage + (i * 16) * WN + j * 16,
                                                acc[i][j], WN, wmma::mem_row_major);
            }
            __syncthreads();
            if (warp_id == w) {
                for (int idx = lane; idx < WM * WN; idx += 32) {
                    const int r = idx / WN;
                    const int c = idx % WN;
                    const int gm = block_m + warp_m + r;
                    const int gn = block_n + warp_n + c;
                    if (gm < M && gn < N) C[(size_t)gm * N + gn] = stage[idx];
                }
            }
            __syncthreads();
        }
    }
#else
    // Fallback for pre-Ampere GPUs: simple global-memory GEMM.
    for (int t = 0; t < (BM * BN) / THREADS; ++t) {
        const int idx = tid + t * THREADS;
        const int r = block_m + idx / BN;
        const int c = block_n + idx % BN;
        if (r < M && c < N) {
            float sum = 0.0f;
            for (int k = 0; k < K; ++k)
                sum += A[(size_t)r * K + k] * B[(size_t)c * K + k];
            C[(size_t)r * N + c] = sum;
        }
    }
#endif
}

torch::Tensor gemm_nt_tf32(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D");
    TORCH_CHECK(A.size(1) == B.size(1), "inner (K) dimensions must match");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);
    const int64_t N = Bc.size(0);

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

    dim3 grid((unsigned)((N + BN - 1) / BN), (unsigned)((M + BM - 1) / BM));
    dim3 block(THREADS);
    auto stream = at::cuda::getCurrentCUDAStream();
    gemm_nt_tf32_kernel<<<grid, block, 0, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, (int)K);
    return C;
}
"""

cpp_source = "torch::Tensor gemm_nt_tf32(torch::Tensor A, torch::Tensor B);"

# Compile the inline CUDA extension
_gemm_ext = load_inline(
    name="gemm_nt_tf32_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["gemm_nt_tf32"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A @ B^T with a custom TF32
    tensor-core CUDA kernel (FP32 storage and FP32 accumulation).
    """
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (A.is_cuda and B.is_cuda and A.dtype == torch.float32
                and B.dtype == torch.float32 and A.dim() == 2 and B.dim() == 2):
            return _gemm_ext.gemm_nt_tf32(A, B)
        return torch.matmul(A, B.T)
```