```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_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>

namespace wmma = nvcuda::wmma;

constexpr int TILE = 128;
constexpr int BK = 32;
constexpr int AS = 40;
constexpr int BS = 136;
constexpr int NT = 256;
constexpr int SPLITS = 2;

union SharedStorage {
    struct {
        __half ah[TILE * AS];
        __half al[TILE * AS];
        __half bh[BK * BS];
        __half bl[BK * BS];
    } input;
    float output[8 * 16 * 64];
};

__device__ __forceinline__ void split_store(
    float4 v, __half* high, __half* low
) {
    const __half2 h01 = __floats2half2_rn(v.x, v.y);
    const __half2 h23 = __floats2half2_rn(v.z, v.w);
    const float2 f01 = __half22float2(h01);
    const float2 f23 = __half22float2(h23);

    reinterpret_cast<__half2*>(high)[0] = h01;
    reinterpret_cast<__half2*>(high)[1] = h23;
    reinterpret_cast<__half2*>(low)[0] =
        __floats2half2_rn(v.x - f01.x, v.y - f01.y);
    reinterpret_cast<__half2*>(low)[1] =
        __floats2half2_rn(v.z - f23.x, v.w - f23.y);
}

__device__ __forceinline__ int triangular_band(int index) {
    int band = static_cast<int>(
        (sqrtf(8.f * static_cast<float>(index) + 1.f) - 1.f) * 0.5f);
    while ((band + 1) * (band + 2) / 2 <= index)
        ++band;
    while (band * (band + 1) / 2 > index)
        --band;
    return band;
}

__global__ __launch_bounds__(NT) void triangular_partial_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int n,
    int tiles
) {
    const int tid = threadIdx.x;
    const int task = blockIdx.x;
    const int tile_index = task / SPLITS;
    const int split = task % SPLITS;

    const int band = triangular_band(tile_index);
    const int column_tile = tile_index - band * (band + 1) / 2;
    const int row_tile = column_tile + tiles - 1 - band;
    const int row0 = row_tile * TILE;
    const int col0 = column_tile * TILE;

    const int steps = (row0 + TILE - col0) / BK;
    const int first_step = steps * split / SPLITS;
    const int last_step = steps * (split + 1) / SPLITS;

    __shared__ __align__(32) SharedStorage smem;

    const int warp = tid >> 5;
    const int wr = (warp >> 1) * 32;
    const int wc = (warp & 1) * 64;

    wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc[2][4];

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

    for (int step = first_step; step < last_step; ++step) {
        const int k0 = col0 + step * BK;

        #pragma unroll
        for (int p = tid * 4; p < TILE * BK; p += NT * 4) {
            const int r = p / BK;
            const int k = p % BK;
            const float4 value = *reinterpret_cast<const float4*>(
                A + static_cast<int64_t>(row0 + r) * n + k0 + k);

            split_store(
                value,
                smem.input.ah + r * AS + k,
                smem.input.al + r * AS + k);
        }

        #pragma unroll
        for (int p = tid * 4; p < BK * TILE; p += NT * 4) {
            const int k = p / TILE;
            const int c = p % TILE;
            const float4 value = *reinterpret_cast<const float4*>(
                B + static_cast<int64_t>(k0 + k) * n + col0 + c);

            split_store(
                value,
                smem.input.bh + k * BS + c,
                smem.input.bl + k * BS + c);
        }

        __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];

            #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 < 4; ++j) {
                wmma::fragment<
                    wmma::matrix_b, 16, 16, 16, __half, wmma::row_major
                > bh, bl;

                wmma::load_matrix_sync(
                    bh, smem.input.bh + kk * BS + wc + j * 16, BS);
                wmma::load_matrix_sync(
                    bl, smem.input.bl + kk * BS + wc + j * 16, BS);

                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    wmma::mma_sync(acc[i][j], ah[i], bl, acc[i][j]);
                    wmma::mma_sync(acc[i][j], al[i], bh, acc[i][j]);
                    wmma::mma_sync(acc[i][j], ah[i], bh, acc[i][j]);
                }
            }
        }

        __syncthreads();
    }

    float* destination =
        partial + static_cast<int64_t>(task) * TILE * TILE;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            wmma::store_matrix_sync(
                smem.output + warp * 1024 + j * 16,
                acc[i][j], 64, wmma::mem_row_major);
        }

        __syncthreads();

        #pragma unroll
        for (int p = tid * 4; p < 8192; p += NT * 4) {
            const int w = p / 1024;
            const int local = p % 1024;
            const int row = (w >> 1) * 32 + i * 16 + local / 64;
            const int col = (w & 1) * 64 + local % 64;

            *reinterpret_cast<float4*>(destination + row * TILE + col) =
                *reinterpret_cast<const float4*>(smem.output + p);
        }

        __syncthreads();
    }
}

__global__ void finish_kernel(
    const float* __restrict__ partial,
    float* __restrict__ C,
    int n,
    int tiles
) {
    const int row_tile = blockIdx.y;
    const int column_tile = blockIdx.x;
    const int row0 = row_tile * TILE;
    const int col0 = column_tile * TILE;

    if (column_tile > row_tile) {
        #pragma unroll
        for (int p = threadIdx.x * 4; p < TILE * TILE; p += NT * 4) {
            const int row = row0 + p / TILE;
            const int col = col0 + p % TILE;
            *reinterpret_cast<float4*>(
                C + static_cast<int64_t>(row) * n + col) =
                make_float4(0.f, 0.f, 0.f, 0.f);
        }
        return;
    }

    const int band = tiles - 1 - (row_tile - column_tile);
    const int tile_index = band * (band + 1) / 2 + column_tile;
    const float* first = partial +
        static_cast<int64_t>(tile_index) * SPLITS * TILE * TILE;
    const float* second = first + TILE * TILE;

    #pragma unroll
    for (int p = threadIdx.x * 4; p < TILE * TILE; p += NT * 4) {
        const float4 a = *reinterpret_cast<const float4*>(first + p);
        const float4 b = *reinterpret_cast<const float4*>(second + p);
        float4 value = make_float4(
            a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);

        const int row = row0 + p / TILE;
        const int col = col0 + p % TILE;

        if (row_tile == column_tile) {
            value.x = col <= row ? value.x : 0.f;
            value.y = col + 1 <= row ? value.y : 0.f;
            value.z = col + 2 <= row ? value.z : 0.f;
            value.w = col + 3 <= row ? value.w : 0.f;
        }

        *reinterpret_cast<float4*>(
            C + static_cast<int64_t>(row) * n + col) = value;
    }
}

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) && A.sizes() == B.sizes(),
        "Expected equally sized square matrices");

    const int64_t size = A.size(0);
    TORCH_CHECK(
        size > 0 && size % TILE == 0 && size <= 65536,
        "Unsupported matrix size");

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

    const int n = static_cast<int>(size);
    const int tiles = n / TILE;
    const int triangular_tiles = tiles * (tiles + 1) / 2;

    auto C = torch::empty({size, size}, A.options());
    auto partial = torch::empty(
        {static_cast<int64_t>(triangular_tiles) * SPLITS * TILE * TILE},
        A.options());

    const cudaStream_t stream =
        c10::cuda::getCurrentCUDAStream().stream();

    triangular_partial_kernel<<<triangular_tiles * SPLITS, NT, 0, stream>>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        partial.data_ptr<float>(),
        n,
        tiles);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_kernel<<<dim3(tiles, tiles), NT, 0, stream>>>(
        partial.data_ptr<float>(), C.data_ptr<float>(), n, tiles);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return C;
}
"""

_triangular_ops = load_inline(
    name="triangular_mm_split_reduction_v5",
    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.device == B.device
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.ndim == 2
            and B.ndim == 2
            and A.shape == B.shape
            and A.shape[0] == A.shape[1]
            and 0 < A.shape[0] <= 65536
            and A.shape[0] % 128 == 0
        ):
            return _triangular_ops.triangular_mm_cuda(A, B)
        return torch.tril(torch.matmul(A, B))
```