```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 <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>

namespace wmma = nvcuda::wmma;

constexpr int BM = 128;
constexpr int BN = 64;
constexpr int BK = 32;
constexpr int AS = 40;
constexpr int BS = 72;
constexpr int THREADS = 256;

union SharedStorage {
    struct {
        __half ah[BM * AS];
        __half al[BM * AS];
        __half bh[BK * BS];
        __half bl[BK * BS];
    } input;
    float output[BM * BN];
};

__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 row0 = blockIdx.y * BM;
    const int col0 = blockIdx.x * BN;

    if (col0 >= row0 + BM) {
        for (int p = tid; p < BM * BN; p += THREADS) {
            const int row = row0 + p / BN;
            const int col = col0 + p % BN;
            if (row < n && col < n)
                C[static_cast<int64_t>(row) * n + col] = 0.0f;
        }
        return;
    }

    __shared__ __align__(32) SharedStorage smem;

    const int warp = tid / 32;
    const int wr = (warp / 2) * 32;
    const int wc = (warp % 2) * 32;

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

    const int kend = min(row0 + BM, n);

    for (int k0 = col0; k0 < kend; k0 += BK) {
        for (int p = tid; p < BM * BK; p += THREADS) {
            const int r = p / BK;
            const int k = p % BK;
            const int gr = row0 + r;
            const int gk = k0 + k;
            const float value = (gr < n && gk < n)
                ? A[static_cast<int64_t>(gr) * n + gk] : 0.0f;
            const __half high = __float2half_rn(value);
            smem.input.ah[r * AS + k] = high;
            smem.input.al[r * AS + k] =
                __float2half_rn(value - __half2float(high));
        }

        for (int p = tid; p < BK * BN; p += THREADS) {
            const int k = p / BN;
            const int c = p % BN;
            const int gk = k0 + k;
            const int gc = col0 + c;
            const float value = (gk < n && gc < n)
                ? B[static_cast<int64_t>(gk) * n + gc] : 0.0f;
            const __half high = __float2half_rn(value);
            smem.input.bh[k * BS + c] = high;
            smem.input.bl[k * BS + c] =
                __float2half_rn(value - __half2float(high));
        }

        __syncthreads();

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

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                wmma::load_matrix_sync(
                    ah[i], smem.input.ah + (wr + i * 16) * AS + kk, AS);
                wmma::load_matrix_sync(
                    al[i], smem.input.al + (wr + i * 16) * AS + kk, AS);
            }

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::load_matrix_sync(
                    bh[j], smem.input.bh + kk * BS + wc + j * 16, BS);
                wmma::load_matrix_sync(
                    bl[j], smem.input.bl + kk * BS + wc + j * 16, BS);
            }

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(acc[i][j], ah[i], bl[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], al[i], bh[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) {
            wmma::store_matrix_sync(
                smem.output + (wr + i * 16) * BN + wc + j * 16,
                acc[i][j], BN, wmma::mem_row_major);
        }
    }

    __syncthreads();

    for (int p = tid; p < BM * BN; p += THREADS) {
        const int row = row0 + p / BN;
        const int col = col0 + p % BN;
        if (row < n && col < n) {
            C[static_cast<int64_t>(row) * n + col] =
                col <= row ? smem.output[p] : 0.0f;
        }
    }
}

torch::Tensor triangular_mm_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Expected matching devices");
    TORCH_CHECK(
        A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
        "Expected float32 tensors");
    TORCH_CHECK(
        A.dim() == 2 && B.dim() == 2 &&
        A.size(0) == A.size(1) &&
        B.size(0) == A.size(0) && B.size(1) == A.size(1),
        "Expected equally sized square matrices");

    const c10::cuda::CUDAGuard guard(A.device());
    A = A.contiguous();
    B = B.contiguous();

    const int64_t size = A.size(0);
    TORCH_CHECK(size <= int64_t(BM) * 65535, "Matrix is too large");
    auto C = torch::empty({size, size}, A.options());
    if (size == 0)
        return C;

    const int n = static_cast<int>(size);
    triangular_mm_kernel<<<
        dim3((n + BN - 1) / BN, (n + BM - 1) / BM),
        THREADS,
        0,
        c10::cuda::getCurrentCUDAStream().stream()
    >>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        C.data_ptr<float>(),
        n
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_triangular_ops = load_inline(
    name="triangular_mm_tensorcore_compensated_v2",
    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 (
            A.is_cuda
            and B.is_cuda
            and A.dtype == torch.float32
            and B.dtype == torch.float32
        ):
            return _triangular_ops.triangular_mm_cuda(A, B)
        return torch.tril(torch.matmul(A, B))
```