```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Custom CUDA kernel that zeroes the strict upper triangle of a square matrix
# in-place. It replaces torch.tril (which reads + writes the whole matrix)
# with a kernel that only writes the ~N^2/2 elements above the diagonal.
zero_upper_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void zero_upper_triangular_kernel(float* __restrict__ c, const int n) {
    const int row = blockIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col > row && col < n) {
        c[static_cast<long long>(row) * static_cast<long long>(n) + col] = 0.0f;
    }
}

void zero_upper_triangular(torch::Tensor c) {
    TORCH_CHECK(c.is_cuda(), "c must be a CUDA tensor");
    TORCH_CHECK(c.scalar_type() == at::kFloat, "c must be float32");
    TORCH_CHECK(c.is_contiguous(), "c must be contiguous");
    TORCH_CHECK(c.dim() == 2 && c.size(0) == c.size(1), "c must be square");
    const int n = static_cast<int>(c.size(0));
    if (n <= 1) {
        return;
    }
    const dim3 block(256);
    const dim3 grid((n + block.x - 1) / block.x, n);
    zero_upper_triangular_kernel<<<grid, block>>>(c.data_ptr<float>(), n);
}
"""

zero_upper_cpp_source = "void zero_upper_triangular(torch::Tensor c);"

try:
    _zero_upper_mod = load_inline(
        name="lower_triangular_matmul_ext",
        cpp_sources=zero_upper_cpp_source,
        cuda_sources=zero_upper_source,
        functions=["zero_upper_triangular"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
    )
    _zero_upper_fn = _zero_upper_mod.zero_upper_triangular
except Exception:
    _zero_upper_fn = None


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

    Rather than one dense (N x N) @ (N x N) matmul followed by torch.tril,
    the product is evaluated as a sequence of block-GEMMs:

        C[k*B:, :(k+1)*B] += A[k*B:, k*B:(k+1)*B] @ B[k*B:(k+1)*B, :(k+1)*B]

    for k = 0 .. nb-1. Because A and B are lower triangular, every block of
    the product strictly above the diagonal is zero, and k-values outside
    [j, i] contribute zeros to C[i, j]. This performs ~N^3/3 FLOPs instead of
    2*N^3 (~6x fewer) while still using cuBLAS (TF32 tensor cores, which are
    allowed) for each block. A small custom CUDA kernel then zeroes the
    strict upper triangle (replacing torch.tril).
    """

    def __init__(self, block_size: int = 512) -> None:
        super().__init__()
        self.block_size = block_size

    def forward(self, A, B):
        # Allow TF32 tensor-core math for the FP32 block-GEMMs.
        prev_tf32 = torch.backends.cuda.matmul.allow_tf32
        torch.backends.cuda.matmul.allow_tf32 = True
        try:
            n = A.size(0)
            A = A.contiguous()
            B = B.contiguous()
            C = A.new_zeros((n, n))
            bs = self.block_size
            for k in range(0, n, bs):
                k2 = k + bs
                if k2 > n:
                    k2 = n
                # In-place block-GEMM update (strides map directly onto
                # cuBLAS leading dimensions, so no copies are made):
                # C[k:, :k2] += A[k:, k:k2] @ B[k:k2, :k2]
                C[k:, :k2].addmm_(A[k:, k:k2], B[k:k2, :k2])
            if _zero_upper_fn is not None and C.is_cuda and C.dtype == torch.float32:
                _zero_upper_fn(C)
            else:
                torch.tril_(C)
        finally:
            torch.backends.cuda.matmul.allow_tf32 = prev_tf32
        return C
```