```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 triangular_mm_cuda(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>

namespace wmma = nvcuda::wmma;

__global__ void triangular_mm_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int N)
{
    constexpr int TILE = 64;
    constexpr int BK = 32;

    const int tid = threadIdx.x;
    const int row0 = blockIdx.y * TILE;
    const int col0 = blockIdx.x * TILE;

    if (blockIdx.x > blockIdx.y) {
        for (int p = tid; p < TILE * TILE; p += blockDim.x) {
            const int r = row0 + p / TILE;
            const int c = col0 + p % TILE;
            if (r < N && c < N)
                C[(int64_t)r * N + c] = 0.0f;
        }
        return;
    }

    __shared__ __align__(32) float storage[8192];
    float* Ah = storage;
    float* Al = storage + 2048;
    float* Bh = storage + 4096;
    float* Bl = storage + 6144;

    const int kend = min(row0 + TILE, N);

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    const int warp = tid / 32;
    const int wr = (warp / 2) * 32;
    const int wc = (warp % 2) * 32;

    using AFrag = wmma::fragment<
        wmma::matrix_a, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using BFrag = wmma::fragment<
        wmma::matrix_b, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using CFrag = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

    CFrag main_acc[2][2];
    CFrag correction[2][2];

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

    for (int k0 = col0; k0 < kend; k0 += BK) {
        for (int p = tid; p < TILE * BK; p += blockDim.x) {
            const int ar = row0 + p / BK;
            const int ak = k0 + p % BK;
            const int bk = k0 + p / TILE;
            const int bc = col0 + p % TILE;

            const float a =
                (ar < N && ak < N) ? A[(int64_t)ar * N + ak] : 0.0f;
            const float b =
                (bk < N && bc < N) ? B[(int64_t)bk * N + bc] : 0.0f;

            const float ah = wmma::__float_to_tf32(a);
            const float bh = wmma::__float_to_tf32(b);
            Ah[p] = ah;
            Al[p] = wmma::__float_to_tf32(a - ah);
            Bh[p] = bh;
            Bl[p] = wmma::__float_to_tf32(b - bh);
        }
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BK; k += 8) {
            AFrag ah[2], al[2];
            BFrag bh[2], bl[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                wmma::load_matrix_sync(ah[i], Ah + (wr + 16 * i) * BK + k, BK);
                wmma::load_matrix_sync(al[i], Al + (wr + 16 * i) * BK + k, BK);
            }
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::load_matrix_sync(bh[j], Bh + k * TILE + wc + 16 * j, TILE);
                wmma::load_matrix_sync(bl[j], Bl + k * TILE + wc + 16 * j, TILE);
            }

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(
                        main_acc[i][j], ah[i], bh[j], main_acc[i][j]);
                    wmma::mma_sync(
                        correction[i][j], ah[i], bl[j], correction[i][j]);
                    wmma::mma_sync(
                        correction[i][j], al[i], bh[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 < CFrag::num_elements; ++e)
                main_acc[i][j].x[e] += correction[i][j].x[e];

            wmma::store_matrix_sync(
                storage + (wr + 16 * i) * TILE + wc + 16 * j,
                main_acc[i][j], TILE, wmma::mem_row_major);
        }
    }
    __syncthreads();

    for (int p = tid; p < TILE * TILE; p += blockDim.x) {
        const int r = row0 + p / TILE;
        const int c = col0 + p % TILE;
        if (r < N && c < N)
            C[(int64_t)r * N + c] = (r >= c) ? storage[p] : 0.0f;
    }
#else
    // CUDA-core fallback for GPUs without TF32 tensor cores.
    const int tx = tid % 16;
    const int ty = tid / 16;
    float acc[8][4] = {};

    for (int k0 = col0; k0 < kend; k0 += BK) {
        for (int p = tid; p < TILE * BK; p += blockDim.x) {
            const int ar = row0 + p / BK;
            const int ak = k0 + p % BK;
            const int bk = k0 + p / TILE;
            const int bc = col0 + p % TILE;

            Ah[p] = (ar < N && ak < N)
                ? A[(int64_t)ar * N + ak] : 0.0f;
            Bh[p] = (bk < N && bc < N)
                ? B[(int64_t)bk * N + bc] : 0.0f;
        }
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            float bv[4];
            #pragma unroll
            for (int j = 0; j < 4; ++j)
                bv[j] = Bh[k * TILE + tx + 16 * j];

            #pragma unroll
            for (int i = 0; i < 8; ++i) {
                const float av = Ah[(ty + 8 * i) * BK + k];
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    acc[i][j] = fmaf(av, bv[j], acc[i][j]);
            }
        }
        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        const int r = row0 + ty + 8 * i;
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int c = col0 + tx + 16 * j;
            if (r < N && c < N)
                C[(int64_t)r * N + c] = (r >= c) ? acc[i][j] : 0.0f;
        }
    }
#endif
}

torch::Tensor triangular_mm_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Inputs must share a device");
    TORCH_CHECK(A.scalar_type() == at::kFloat &&
                B.scalar_type() == at::kFloat, "Inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 &&
                A.size(0) == A.size(1) && A.sizes() == B.sizes(),
                "Inputs must be equally sized square matrices");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Inputs must be contiguous");

    c10::cuda::CUDAGuard guard(A.device());
    auto C = torch::empty_like(A);
    const int N = static_cast<int>(A.size(0));

    if (N != 0) {
        dim3 grid((N + 63) / 64, (N + 63) / 64);
        triangular_mm_kernel<<<
            grid, 128, 0, at::cuda::getCurrentCUDAStream()
        >>>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), N);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return C;
}
"""

_triangular_mm = load_inline(
    name="triangular_mm_tf32_split_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["triangular_mm_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


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

    def forward(self, A, B):
        if (
            not A.is_cuda
            or not B.is_cuda
            or A.dtype != torch.float32
            or B.dtype != torch.float32
        ):
            return torch.tril(torch.matmul(A, B))
        return _triangular_mm.triangular_mm_cuda(
            A.contiguous(), B.contiguous()
        )
```