[liuanji]

Summary

1. A precision option for the sum layers (ebc503b3). TensorCircuit.precision (also juice.compile(..., precision = ...)) replaces force_use_bf16 / force_use_fp32:

mode	forward products	backward products	log-likelihood error
"auto" (default)	today's kernels: bf16 from batch 16	round-to-nearest TF32	~4e-3
"tf32"	round-to-nearest TF32, never bf16	round-to-nearest TF32	~5e-4
"fp32"	fp32-level (tf32x3 in the Triton tlmm kernel)	fp32-level (tf32x3)	~1e-6

The mode can be changed at any time, and a recorded CUDA graph is keyed on it. Batches below 16 run fp32 kernels in every mode, as before.

2. Dense sum layers through cuBLAS (79b76e7e). A large dense layer exponentiates its children, multiplies on cuBLAS using views of the parameters, and takes the log. "Large" means block size ≥ 1024 and children × batch ≥ 2^19, for LL without tempering.

• Layers whose blocks share their children (HMMs) exponentiate them once for all blocks.
• Layers with evenly strided child groups (HCLT) run as batched GEMMs, with scratch capped at 64 MB.
• Any other layout keeps the kernels.

The products are TF32 in "auto" and "tf32", and exact fp32 in "fp32". The gate is fixed, never chosen by timing, because the two paths differ numerically.

3. "fp32"'s backward on tensor cores (18f7192f). In "fp32" mode the element- and parameter-flow kernels used to drop tl.dot for broadcast sums on CUDA cores, as force_use_fp32 always did. They now keep the dots and run them in three TF32 passes (input_precision = "tf32x3").

• On one HMM layer, against float64, the flows are 5e-7 from exact. The round-to-nearest TF32 of "auto" and "tf32" is 5e-5 (element flows) and 3e-4 (parameter flows).
• The other modes launch the kernels as before.
• The fp16/TF32 CUDA backward kernels stay off in "fp32".

Breaking change

force_use_bf16 = True / force_use_fp32 = True now raise a TypeError that names the replacement, in TensorCircuit.forward / backward and in SumLayer.forward / backward. Passing False (their old default) is still accepted and ignored.

Results

Whole-circuit timings against main, "auto", RTX PRO 6000. Each figure is the median over processes: two GPUs on different CPU sockets with the arms swapped, plus alternating one-at-a-time runs.

workload	batch	forward	forward + backward
HMM, 4096 latents, 32 positions	128	×2.25	×1.31
HMM, 4096 latents, 32 positions	512	×1.57	×1.12
HMM, 1024 latents, 64 positions	512	×1.10	×1.02
HCLT, 1024 latents, 768 variables, 512 categories	512	×1.20	×1.05
layers the gate does not take (PD, smaller batches)		×1.00	×1.00 (within ~1%)

• Accuracy: dense layers now use TF32 instead of bf16 products. On a 1024-latent HMM layer at batch 128, the log-likelihood error drops from ~4e-3 to ~5e-4.

• Memory: on HCLT at batch 512 the dense path needs 64 MB of scratch, against 3 MB for the kernels, with 6.7 GB resident.

• "fp32" speed: forward+backward is now ×0.64–0.98 of "auto", from 2–50× slower before the third commit:

  workload	batch	"auto"	"fp32", broadcast-sum backward	"fp32", tf32x3 backward
  HMM, 4096 latents, 32 positions	128	13.9 ms	299.6 ms	17.9 ms
  HMM, 4096 latents, 32 positions	512	32.1 ms	1608.9 ms	49.9 ms
  HMM, 1024 latents, 64 positions	512	10.2 ms	153.1 ms	12.3 ms
  HCLT, 1024 latents, 768 variables	512	48.1 ms	1579.4 ms	64.3 ms

Tests

Full suite: 1301 passed, 26 skipped.

• tests/layer/sum_layer_precision_test.py:
  • each mode against a float64 forward, on the Triton path, on one-group dense layers and on batched three-group dense layers;
  • the per-group column-maximum edge cases;
  • the backward's element and parameter flows in each mode against float64, with "fp32" passing tf32x3 to both flow kernels and the other modes passing nothing. Downgrading either kernel's dots to TF32 fails it;
  • the gate, and the plans the dense path makes of made-up layouts;
  • the removed flags;
  • unknown modes.
  • Each of nine mutations of the dense path is caught.
• tests/model/precision_option_test.py:
  • compiling with and setting the mode;
  • every sum layer receiving it in both passes;
  • the removed flags at the circuit;
  • CUDA graphs keyed on it.
• The tests that set force_use_fp32 now set pc.precision.