```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 {

constexpr int TILE = 64;
constexpr int BK = 32;
constexpr int THREADS = 128;

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
__device__ __forceinline__ float round_tf32(float x) {
    float y;
    asm("cvt.rn.tf32.f32 %0, %1;" : "=f"(y) : "f"(x));
    return y;
}
#endif

__global__ void triangular_mm_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int n
) {
    const int tid = threadIdx.x;
    const int row_base = blockIdx.y * TILE;
    const int col_base = blockIdx.x * TILE;

    if (blockIdx.x > blockIdx.y) {
        for (int q = tid; q < TILE * TILE; q += THREADS) {
            const int row = row_base + q / TILE;
            const int col = col_base + q % TILE;
            if (row < n && col < n)
                C[(int64_t)row * n + col] = 0.0f;
        }
        return;
    }

    const int k_end = min(row_base + TILE, n);

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    using namespace nvcuda;

    union SharedStorage {
        struct {
            float ah[TILE * BK];
            float al[TILE * BK];
            float bh[BK * TILE];
            float bl[BK * TILE];
        } input;
        float output[TILE * TILE];
    };
    __shared__ __align__(32) SharedStorage smem;

    const int warp = tid / 32;
    const int warp_row = (warp / 2) * 32;
    const int warp_col = (warp % 2) * 32;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[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);
    }

    for (int kb = col_base; kb < k_end; kb += BK) {
        for (int q = tid; q < TILE * BK; q += THREADS) {
            const int ar = row_base + q / BK;
            const int ak = kb + q % BK;
            const float av = (ar < n && ak < n)
                ? A[(int64_t)ar * n + ak] : 0.0f;
            const float ah = round_tf32(av);
            smem.input.ah[q] = ah;
            smem.input.al[q] = round_tf32(av - ah);

            const int bk = kb + q / TILE;
            const int bc = col_base + q % TILE;
            const float bv = (bk < n && bc < n)
                ? B[(int64_t)bk * n + bc] : 0.0f;
            const float bh = round_tf32(bv);
            smem.input.bh[q] = bh;
            smem.input.bl[q] = round_tf32(bv - bh);
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major
            > ah[2], al[2];
            wmma::fragment<
                wmma::matrix_b, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major
            > bh[2], bl[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int offset = (warp_row + i * 16) * BK + kk;
                wmma::load_matrix_sync(ah[i], smem.input.ah + offset, BK);
                wmma::load_matrix_sync(al[i], smem.input.al + offset, BK);
            }
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const int offset = kk * TILE + warp_col + j * 16;
                wmma::load_matrix_sync(bh[j], smem.input.bh + offset, TILE);
                wmma::load_matrix_sync(bl[j], smem.input.bl + offset, TILE);
            }

            // Three products retain substantially more precision than TF32.
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(acc[i][j], al[i], bh[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], ah[i], bl[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], ah[i], bh[j], acc[i][j]);
                }
            }
        }
        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            const int offset =
                (warp_row + i * 16) * TILE + warp_col + j * 16;
            wmma::store_matrix_sync(
                smem.output + offset, acc[i][j], TILE, wmma::mem_row_major);
        }
    }
    __syncthreads();

    for (int q = tid; q < TILE * TILE; q += THREADS) {
        const int row = row_base + q / TILE;
        const int col = col_base + q % TILE;
        if (row < n && col < n) {
            C[(int64_t)row * n + col] =
                col <= row ? smem.output[q] : 0.0f;
        }
    }
#else
    // FP32 implementation for GPUs without TF32 tensor cores.
    __shared__ float sa[TILE * BK];
    __shared__ float sb[BK * TILE];

    const int local_row = tid / 16;
    const int local_col = tid % 16;
    float acc[8][4] = {};

    for (int kb = col_base; kb < k_end; kb += BK) {
        for (int q = tid; q < TILE * BK; q += THREADS) {
            const int ar = row_base + q / BK;
            const int ak = kb + q % BK;
            sa[q] = (ar < n && ak < n)
                ? A[(int64_t)ar * n + ak] : 0.0f;

            const int bk = kb + q / TILE;
            const int bc = col_base + q % TILE;
            sb[q] = (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] = sb[k * TILE + local_col + j * 16];

            #pragma unroll
            for (int i = 0; i < 8; ++i) {
                const float av = sa[(local_row + i * 8) * 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 row = row_base + local_row + i * 8;
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int col = col_base + local_col + j * 16;
            if (row < n && col < n)
                C[(int64_t)row * n + col] =
                    col <= row ? acc[i][j] : 0.0f;
        }
    }
#endif
}

} // namespace

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) && B.sizes() == A.sizes(),
        "Inputs must be equally sized square matrices");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Inputs must be contiguous");

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

    const dim3 grid((n + TILE - 1) / TILE, (n + TILE - 1) / TILE);
    triangular_mm_kernel<<<
        grid, THREADS, 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_decomposed_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()
        )
```