feat(cutile): chunk FLCE with Triton sizing and cuBLAS GEMMs

[vaibhavjindal]

Summary

Add chunked cuTile FLCE on Hopper SM90 and Blackwell SM100, using cuBLAS for all matrix multiplications and matching Triton's current chunking, precision, and forward-allocation baseline. This PR is based on main after #1454 and adds no public chunk_size argument.

Matched baseline

• Both backends call the same private chunk helper: inc = ceil(V / (16 * H)), chunk = min(N, next_pow2(ceil(N / inc))).
• cuTile allocates fresh logits for each chunk with input_chunk @ weight.T; CE overwrites that tensor with dZ. Previous/current chunks can overlap through live gradient/transpose views, as in Triton.
• dX uses a temporary dZ @ weight result copied into a zero-initialized gradient buffer, not a direct GEMM output write.
• dW uses a zero-initialized accumulator and addmm on every chunk, including the first. The meaningful transposed-gradient operand is held in the loop just as in Triton.
• accum_dtype=None or BF16 uses a BF16 running dW buffer; explicit FP32 uses an FP32 buffer and casts the completed gradient to BF16 at the end of forward. Mean normalization is applied in FP32 to dZ before the BF16 store and gradient GEMMs. Backward applies only upstream scaling, out-of-place for retained-graph safety.
• Fresh GEMMs explicitly preserve the supported BF16 operand dtype under ambient FP16 autocast. This preserves the previous out-based GEMMs' precision behavior without bringing back their allocation optimization.
• The CE kernel still needs a small, chunk-bounded statistics/counter scratch allocation specific to its reduction algorithm.

Deliberately deferred: reusable logits storage, direct dX output writes, and first-chunk initialization shortcuts will be introduced together for Triton and cuTile in a separate optimization PR. There are no artificial padding allocations or unused references to force matching memory numbers.

Memory comparison

Reference shape: N=8192, H=4096, V=128256, BF16 inputs and weights, 4096-row default chunks on both sides. Full forward + backward, mean reduction, upstream 1.0. Values are warmed incremental peak allocated MiB above inputs, not total or reserved memory.

Accumulator mode	Triton	cuTile
Default BF16	3102.04 MiB	3103.05 MiB
Explicit FP32	4104.04 MiB	4105.05 MiB

The remaining approximately 1 MiB is cuTile's CE statistics/counter scratch, not a different logits or gradient-buffer allocation strategy. Actual allocation traces confirm two overlapping logits chunks and the same temporary dX pattern in both implementations. Older measurements showing a cuTile-only workspace-memory advantage describe the superseded allocation implementation and should not be used for this revision.

Scope and compatibility

BF16 input/weight, mean/sum, and None/BF16/FP32 accumulation are supported. CE reductions need not be bitwise identical to Triton and optional features such as bias, class weights, z-loss, label smoothing, softcap and token scaling remain unsupported on this cuTile path. The original exported unchunked raw-dZ/7-argument backward helpers retain their compatibility behavior; normal training uses the aligned chunked path. No custom cuTile GEMM remains.

Testing Done

• Hardware: NVIDIA B200, SM100, driver 580.105.08; torch 2.11.0.1+cu130, Triton 3.6.0, cuda-tile 1.5.0.
• Existing and new relevant FLCE/backend suites: 719 passed, 2 skipped. Physical Hopper and model-convergence runs were not performed.
• Coverage includes actual GEMM/allocation traces, global mean normalization, accumulator dtype, default/shared chunk geometry, tails, independent gradient masks, ignored chunks, retained backward, legacy helpers, older-Torch fallback and BF16 operands under ambient FP16 autocast.
• Independent review found the autocast regression introduced by fresh matmuls; it was fixed and is covered by six exact-parity regression cases.
• Local make checkstyle passes.
• GitHub Checkstyle passes on 5f1f7b9.
• Full-repository make test and make test-convergence were not run; the relevant suites above were run.

The architecture-neutral test file is test/cutile/test_fused_linear_cross_entropy.py. User-local setup.py pins are not included.