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

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 16
#define SMPAD 132  // padded shared row length (multiple of 4 floats => 16B alignment kept)

// C = tril(A @ B) for lower-triangular A, B.
// Only tiles on/below the diagonal are computed; k-tiles are restricted to the
// contributing range [col0, row0 + BM - 1] (all other terms are exactly zero).
__global__ void __launch_bounds__(256, 2) tril_mm_tiled_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int N)
{
    const int bx = blockIdx.x;  // column tile
    const int by = blockIdx.y;  // row tile
    const int tx = threadIdx.x; // 0..15
    const int ty = threadIdx.y; // 0..15

    float* Cg = C + (size_t)(by * BM) * N + bx * BN;

    if (bx > by) {  // strictly upper tiles: result is exactly zero
        const float4 z = make_float4(0.f, 0.f, 0.f, 0.f);
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int r0 = ty * 4 + i;
            const int r1 = ty * 4 + 64 + i;
            *reinterpret_cast<float4*>(Cg + (size_t)r0 * N + tx * 4)      = z;
            *reinterpret_cast<float4*>(Cg + (size_t)r0 * N + tx * 4 + 64) = z;
            *reinterpret_cast<float4*>(Cg + (size_t)r1 * N + tx * 4)      = z;
            *reinterpret_cast<float4*>(Cg + (size_t)r1 * N + tx * 4 + 64) = z;
        }
        return;
    }

    __shared__ float As[2][BK][SMPAD];
    __shared__ float Bs[2][BK][SMPAD];

    const int tid   = ty * 16 + tx;
    const int a_row = tid >> 2;  // 0..63
    const int a_vf  = tid & 3;   // 0..3  (BK=16 => 4 float4 per row)
    const int b_row = tid >> 5;  // 0..7
    const int b_vf  = tid & 31;  // 0..31 (BN=128 => 32 float4 per row)

    const float* Ag = A + (size_t)(by * BM + a_row) * N + a_vf * 4;
    const float* Bg = B + (size_t)b_row * N + bx * BN + b_vf * 4;

    const int kt_begin = bx * (BN / BK);
    const int kt_end   = (by + 1) * (BM / BK);

    float acc[8][8];
    #pragma unroll
    for (int i = 0; i < 8; ++i)
        #pragma unroll
        for (int j = 0; j < 8; ++j) acc[i][j] = 0.0f;

    // preload first tile
    {
        const size_t koff = (size_t)kt_begin * BK;
        const float4 a0 = *reinterpret_cast<const float4*>(Ag + koff);
        const float4 a1 = *reinterpret_cast<const float4*>(Ag + (size_t)64 * N + koff);
        const float4 b0 = *reinterpret_cast<const float4*>(Bg + koff * (size_t)N);
        const float4 b1 = *reinterpret_cast<const float4*>(Bg + ((size_t)koff + 8) * (size_t)N);

        As[0][a_vf * 4 + 0][a_row]      = a0.x;
        As[0][a_vf * 4 + 1][a_row]      = a0.y;
        As[0][a_vf * 4 + 2][a_row]      = a0.z;
        As[0][a_vf * 4 + 3][a_row]      = a0.w;
        As[0][a_vf * 4 + 0][a_row + 64] = a1.x;
        As[0][a_vf * 4 + 1][a_row + 64] = a1.y;
        As[0][a_vf * 4 + 2][a_row + 64] = a1.z;
        As[0][a_vf * 4 + 3][a_row + 64] = a1.w;

        *reinterpret_cast<float4*>(&Bs[0][b_row][b_vf * 4])     = b0;
        *reinterpret_cast<float4*>(&Bs[0][b_row + 8][b_vf * 4]) = b1;
    }
    __syncthreads();

    int buf = 0;
    for (int kt = kt_begin; kt < kt_end; ++kt, buf ^= 1) {
        float4 a0, a1, b0, b1;
        const bool has_next = (kt + 1 < kt_end);
        if (has_next) {  // prefetch next tile to registers
            const size_t kn = (size_t)(kt + 1) * BK;
            a0 = *reinterpret_cast<const float4*>(Ag + kn);
            a1 = *reinterpret_cast<const float4*>(Ag + (size_t)64 * N + kn);
            b0 = *reinterpret_cast<const float4*>(Bg + kn * (size_t)N);
            b1 = *reinterpret_cast<const float4*>(Bg + (kn + 8) * (size_t)N);
        }

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            const float4 av0 = *reinterpret_cast<const float4*>(&As[buf][kk][ty * 4]);
            const float4 av1 = *reinterpret_cast<const float4*>(&As[buf][kk][ty * 4 + 64]);
            const float4 bv0 = *reinterpret_cast<const float4*>(&Bs[buf][kk][tx * 4]);
            const float4 bv1 = *reinterpret_cast<const float4*>(&Bs[buf][kk][tx * 4 + 64]);
            const float av[8] = {av0.x, av0.y, av0.z, av0.w, av1.x, av1.y, av1.z, av1.w};
            const float bv[8] = {bv0.x, bv0.y, bv0.z, bv0.w, bv1.x, bv1.y, bv1.z, bv1.w};
            #pragma unroll
            for (int i = 0; i < 8; ++i)
                #pragma unroll
                for (int j = 0; j < 8; ++j)
                    acc[i][j] = fmaf(av[i], bv[j], acc[i][j]);
        }

        if (has_next) {  // store prefetched tile into the other buffer
            const int nxt = buf ^ 1;
            As[nxt][a_vf * 4 + 0][a_row]      = a0.x;
            As[nxt][a_vf * 4 + 1][a_row]      = a0.y;
            As[nxt][a_vf * 4 + 2][a_row]      = a0.z;
            As[nxt][a_vf * 4 + 3][a_row]      = a0.w;
            As[nxt][a_vf * 4 + 0][a_row + 64] = a1.x;
            As[nxt][a_vf * 4 + 1][a_row + 64] = a1.y;
            As[nxt][a_vf * 4 + 2][a_row + 64] = a1.z;
            As[nxt][a_vf * 4 + 3][a_row + 64] = a1.w;

            *reinterpret_cast<float4*>(&Bs[nxt][b_row][b_vf * 4])     = b0;
            *reinterpret_cast<float4*>(&Bs[nxt][b_row + 8][b_vf * 4]) = b1;
        }
        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        const int r0 = ty * 4 + i;
        const int r1 = ty * 4 + 64 + i;
        *reinterpret_cast<float4*>(Cg + (size_t)r0 * N + tx * 4) =
            make_float4(acc[i][0], acc[i][1], acc[i][2], acc[i][3]);
        *reinterpret_cast<float4*>(Cg + (size_t)r0 * N + tx * 4 + 64) =
            make_float4(acc[i][4], acc[i][5], acc[i][6], acc[i][7]);
        *reinterpret_cast<float4*>(Cg + (size_t)r1 * N + tx * 4) =
            make_float4(acc[i + 4][0], acc[i + 4][1], acc[i + 4][2], acc[i + 4][3]);
        *reinterpret_cast<float4*>(Cg + (size_t)r1 * N + tx * 4 + 64) =
            make_float4(acc[i + 4][4], acc[i + 4][5], acc[i + 4][6], acc[i + 4][7]);
    }
}

// Generic fallback (any N): one thread per output element.
__global__ void tril_mm_naive_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int N)
{
    const int j = blockIdx.x * blockDim.x + threadIdx.x;
    const int i = blockIdx.y * blockDim.y + threadIdx.y;
    if (i < N && j < N) {
        if (j > i) {
            C[(size_t)i * N + j] = 0.0f;
            return;
        }
        float s = 0.0f;
        for (int k = j; k <= i; ++k)
            s = fmaf(A[(size_t)i * N + k], B[(size_t)k * N + j], s);
        C[(size_t)i * N + j] = s;
    }
}

torch::Tensor tril_matmul_cuda(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D");
    TORCH_CHECK(A.size(0) == A.size(1) && B.size(0) == B.size(1) && A.size(0) == B.size(0),
                "A and B must be square matrices of the same size");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const int N = (int)Ac.size(0);
    auto C = at::empty_like(Ac);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (N % 128 == 0) {
        const int nb = N / 128;
        dim3 grid(nb, nb);
        dim3 block(16, 16);
        tril_mm_tiled_kernel<<<grid, block, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(), N);
    } else {
        dim3 block(16, 16);
        dim3 grid((N + 15) / 16, (N + 15) / 16);
        tril_mm_naive_kernel<<<grid, block, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(), N);
    }
    return C;
}
"""

cpp_src = "torch::Tensor tril_matmul_cuda(torch::Tensor A, torch::Tensor B);"

tri_mm = load_inline(
    name="tril_matmul_ext",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["tril_matmul_cuda"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized lower-triangular matrix multiplication.

    Uses a custom CUDA kernel that:
      * skips all tiles strictly above the diagonal (writes exact zeros there),
      * restricts the k-loop to the contributing range only (~N^3/6 MACs instead of N^3),
      * fuses the trailing tril() into the matmul,
      * uses 128x128 tiles, 8x8 register blocking, float4 vectorized global/shared
        access, padded shared memory to avoid bank conflicts, and double buffering.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._tri_mm = tri_mm

    def forward(self, A, B):
        if (
            A.is_cuda
            and B.is_cuda
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.dim() == 2
            and A.shape == B.shape
            and A.size(0) == A.size(1)
        ):
            return self._tri_mm.tril_matmul_cuda(A, B)
        return torch.tril(torch.matmul(A, B))
```