fix(energy): j-sub-chunking + NRT fallback for large-shape DF-MP2

[scttfrdmn]

Summary

• j-sub-chunking (_J_CHUNK_MAX_B_BYTES = 256 MB): splits the j-loop in _nki_batched_pair_energy_chunked_impl so each _j_batched_kernel call receives B:(j_chunk × nvir_pad × naux_pad) ≤ 256 MB, preventing neuronx-cc from OOMing on large shapes where the full B:(96, 768, 2304) ≈ 680 MB exhausted compiler RAM.
• Proactive NRT fallback (_NRT_INNER_OP_LIMIT = 300): when N_A × N_B × N_K > 300, the compiled NEFF cannot be loaded by the NeuronCore NRT hardware regardless of j_chunk size (confirmed j_chunk = 1, 16, 32 all fail). Falls back to _torch_batched_pair_energy immediately, skipping ~30 min of wasted NEFF compilation.
• Large-shape bench results (first complete run, trn1.2xlarge): chol/half/metric use NKI; energy uses PyTorch CPU fallback. Cold 174.5 s, warm 80.5 s, E = −43.51 Ha, ~0.25 TFLOPS warm.

Large-shape timing (2026-04-24, trn1.2xlarge, neuronxcc 2.24.5133)

Step	Cold	Warm
Cholesky (NKI)	38.5 s	12.5 s
Half-transform (NKI)	96.5 s	30.4 s
Metric (NKI)	3.8 s	2.0 s
Energy (PyTorch fallback)	35.6 s	35.5 s
Total	174.5 s	80.5 s

Energy cold ≈ warm (35.5 s) confirms PyTorch CPU (no NKI cache benefit). NKI energy for large shape requires kernel redesign: N_A × N_B × N_K = 648 exceeds NRT limit (~192).

Test plan

• pytest tests/ -m "not neuron" -q — 58 passed
• test_j_sub_chunked_agrees_with_full_j and test_j_sub_chunked_matches_torch_ref
• Medium bench (warm ~120 s): unchanged
• Large bench: cold 174.5 s, warm 80.5 s

[codecov]

Codecov Report

❌ Patch coverage is 5.55556% with 34 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
trnblas/nki/dispatch.py	5.55%	34 Missing ⚠️

📢 Thoughts on this report? Let us know!