Fused complex-M2L Pallas: v1/v2 tuning (on-chip rotation reuse + kill power-of-2 z-core padding)

[TobiBu]

Context

The v0 fused complex-M2L Pallas kernel was re-wired into runtime dispatch in #40 and validated
on A100 (sm_80). Correctness is perfect (relerr ~1e-16 vs the solidfmm reference at orders 2–5,
~1e-17 end-to-end), but as-wired it is a net perf regression, so it ships gated OFF behind
JACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS. This issue tracks the v1/v2 tuning needed to turn the
validated kernel into an actual win.

Evidence (A100, steady-state, JIT'd, f64, best-of-20 ms)

ref = jitted _m2l_complex_batch_kernel (solidfmm reference);
fused = jitted _m2l_complex_batch_kernel_fused_pallas (build solidfmm blocks + kernel — the
runtime dispatch path); kernel-only = jitted m2l_complex_fused_pallas with blocks prebuilt.

order	pairs	ref	fused dispatch	speedup	kernel-only
2	100k	6.48	8.06	0.80×	2.02
3	100k	10.0	11.4	0.88×	2.30
4	100k	14.1	31.5	0.45×	18.9
5	100k	19.2	107	0.18×	91.1

Key reads:

• The pure kernel beats the reference at low order — ~3.2× (p2) / ~4.3× (p3) at 100k pairs.
  The on-chip rotate→z→rotate is sound.
• The win is destroyed twice over by the two levers below.

Lever 1 — stop materialising rotation blocks to HBM

_m2l_complex_batch_kernel_fused_pallas calls the solidfmm adapter to build
blocks_to_z / blocks_from_z as [N, p+1, 2p+1, 2p+1] complex arrays, writes them to HBM,
and the kernel reads them back. That block-build + HBM round-trip is the fused − kernel-only
gap and dominates the dispatch cost (e.g. p3/100k: 11.4 − 2.30 ≈ 9.1 ms of the 11.4). The
reference keeps rotation fused with the matmul and never pays it.

Target: generate the rotation coefficients on-chip / reuse them across pairs instead of
pre-materialising per-pair block matrices — i.e. the planned class-major shared-memory rotation
reuse + fused gather. Pairs that share a rotation class should compute the rotation once in
shared memory rather than each reading a full block from HBM.

Lever 2 — kill the power-of-2 padding waste in the z-core

The Triton lowering pads dims to the next power of 2 (Cp, Bp, mdp, Kp). The dense z-core
operates on [Cp, Cp]:

order	logical C	Cp	wasted
2	9	16	3.2×
3	16	16	1.0×
4	25	32	1.6×
5	36	64	3.2×

This is why even the pure kernel loses at p ≥ 4 (kernel-only 18.9 ms @ p4, 91.1 ms @ p5). The
dense [Cp,Cp] matmul does mostly wasted work on padded lanes.

Target: block-diagonal exploitation — the z-core couples only (n,m)←(k,m) within the
same m, so the operator is block-diagonal by m, not a dense [Cp,Cp]. Exploiting that (or
tighter tiling that doesn't round the whole coefficient axis to a power of 2) removes the padding
blow-up, most impactfully at high order.

Acceptance / done-when

• Fused dispatch (_m2l_complex_batch_kernel_fused_pallas, blocks built in-path) beats the
  solidfmm reference at the production orders/particle counts, not just kernel-only.
• Parity preserved (relerr ≤ 1e-10 f64 vs _m2l_complex_batch_kernel, orders 2–5).
• Re-run the A100 benchmark and update the table; then flip the default (or the fast-lane wiring)
  to use it where it wins.

Refs: #40, jaccpot/pallas/m2l_complex_fused.py (see the v0/v1/v2 note in the module docstring),
docs/phase5_pallas_plan.md.