[DSv4][P5-3] shared_grouped_lora_delta: CUDA + Triton, both batch-invariant

[ScottNi0809]

[DSv4][P5-3] shared_grouped_lora_delta — CUDA + Triton, both batch-invariant

Summary

Implements P5-3 (#62), the LoRA delta reused at both fc1 and fc2:

U  = X @ A^T          FP32 accumulate, round to BF16
ΔY = (U @ B^T) * α    FP32 accumulate, FP32 output
dV = dY * α    dU = dV @ B    dB = dV^T @ U    dX = dU @ A    dA = dU^T @ X

No dW — base weights are frozen. dY*α and dU round to BF16 between the
GEMMs (D4); every accumulation is FP32 and the association order is frozen.

Path	Status
LoRADeltaProvider (torch-native)	Benchmark baseline. Not batch-invariant. New.
LoRADeltaCudaProvider	Composed from _C.det_gemm_*. Batch-invariant. New.
LoRADeltaTritonProvider	Composed from TritonDetGemmOp. Batch-invariant. New.
tests/test_p5_lora_delta.py	44 cases over both deterministic backends. New.
benchmarks/benchmark_lora_delta.py	native / Triton / CUDA. New.
scripts/check_p5_lora_delta_invariance.py	Invariance sweep. New.

Each provider overrides only this operator; the other eight stay on the oracle,
so the acceptance command runs end to end. Neither deterministic path adds a
GEMM — both compose the repo's existing fixed-K primitives.

Batch invariance did not require avoiding tensor cores. tl.dot is
deterministic once the tile shape stops depending on M; what matters is pinned
BLOCK sizes (autotune picks per-shape configs), no split-K, and
allow_tf32=False.

Acceptance criteria

#	Criterion	Evidence
1	base-only / LoRA-only / base+LoRA pass the golden vector	check_p5.py PASS, 7/7 cases
2	Fixture proving the base weight is not unpacked	test_tampering_packed_base_bytes_changes_the_output
3	dX/dA/dB match the FP32 reference; no base grad	test_matches_oracle_within_bf16_tolerance, test_base_weights_get_no_gradient
4	LoRA gradients finite and non-zero	test_lora_grads_are_finite_and_nonzero
5	fc1 and fc2 insertion points each have a fixture	test_both_insertion_points[fc1/fc2]

On (2): the LoRA path cannot unpack the base weight — its signature never
receives it, shared_grouped_lora_delta_fwd(x, a, b, alpha), and the test
asserts that. Behaviourally it flips one byte of w1.codes and requires the
output to move (494/3072 elements, up to 2.2e-2), with a control assertion that
an untampered rerun is bitwise identical so the check cannot pass on noise.

python scripts/check_p5.py --provider \
  rl_engine.moe.shared_grouped_lora_delta_provider:LoRADeltaTritonProvider --device cuda
# provider=p5-3-lora-delta-triton profile=triton-det-gemm-fp32-accum device=cuda
# RESULT: PASS (all boundaries byte-equal)

python -m pytest tests/test_p5_lora_delta.py -q   # 44 passed
python -m pytest tests/test_p5_*.py -q            # 27 S0 tests still pass

Batch invariance

scripts/check_p5_lora_delta_invariance.py -q — 100 probes per backend
(5 seeds × 5 sub-batch sizes × 4 shapes, fixture 128×64 up to production
4096×2048), checking every boundary rather than only the final output:

backend	u (GEMM1)	y (GEMM2)	dX (bwd)	rerun
torch-native	100/100	87/100	47/100	ok
CUDA det_gemm	100/100	100/100	100/100	ok
Triton det_gemm	100/100	100/100	100/100	ok

Two things worth knowing about how this probe had to be built:

• GEMM1 cannot carry an invariance test. u is [M, 8] — too narrow for
  cuBLAS to switch strategy, so even torch.matmul is invariant there. Only
  y and the backward dX diverge.
• Small K hides the problem. cuBLAS keeps one strategy across all M until
  the matrix is worth splitting, so mutating a deterministic backend back to
  torch.matmul went undetected at 128×64. Production geometry is now swept
  too. (dA/dB reduce over every row, so a sub-batch result is a different
  quantity rather than a slice; they are covered by the rerun column.)

Agreement with the oracle

Neither backend declares ORACLE_PROFILE:

path	vs oracle	why
Triton	bitwise at fixture width; ~1e-6 at production width	tile-wise vs term-wise FP32 accumulation
CUDA	~5e-3 at all K	BF16 rounding per leaf of det_gemm's mid-split tree

The Triton agreement degrades with K rather than breaking at a threshold —
byte-equal runs out of 12 seeds, at M=24 N=2048:

K	64	128	256	512	1024	2048	3072	4096
byte-equal	12	12	12	12	11	9	7	7

This is a precision-headroom effect. BF16 carries 8 mantissa bits, so a
BF16×BF16 product needs at most 16 and lands in FP32's 24 exactly. While the
partial sums still fit inside that headroom nothing rounds at all, so the
oracle's term-by-term order and the op's BLOCK_K=32 tile order cannot
produce different bits. As K grows the accumulator climbs into the remaining
bits and starts dropping low ones — at which point the parenthesisation
matters, and whether a particular input crosses that line depends on its
values. The gap stays FP32-epsilon sized (~1e-6 relative, 1 element in 192 for
u), never a correctness error.

check_p5.py passes because it only exercises the fixture geometry
(K=128 for fc1, K=64 for fc2). Two tests bracket this instead of overclaiming:
one pins the byte-equality at fixture width, the other pins the bounded
divergence at production width and fails if that divergence ever disappears.

The CUDA path's larger gap buys something the Triton path lacks: a contiguous
half-K shard is one child of det_gemm's tree, so TP=2 reproduces TP=1 bitwise.
That is likely what P5-8 needs, so it is left as is deliberately.

Benchmark

python benchmarks/benchmark_lora_delta.py --direction fwd (also bwd,
both). Median of 5×100 iterations after 20 warmups, ms:

shape	M	K	N	r	native fp32	Triton	CUDA	Triton ovh	CUDA ovh
decode_m1	1	4096	2048	8	0.200	0.248	0.205	1.2x	1.0x
small_m8	8	4096	2048	8	0.176	0.233	0.201	1.3x	1.1x
batch_m64	64	4096	2048	8	0.193	0.344	0.352	1.8x	1.8x
batch_m256	256	4096	2048	8	0.214	0.240	0.718	1.1x	3.4x

Overhead is against native fp32 (cuBLAS with TF32 disabled) — the fair
baseline, since both deterministic paths disable TF32. As in
benchmark_det_gemm.py, this is overhead, not speedup.

The CUDA path degrades with M because det_gemm falls back to its naive scalar
kernel below SM90; an SM90 runner would take the TMA path.

Where the time actually goes: at M=64 the two BF16 GEMMs cost 0.070 ms but
the whole native forward costs 0.222 ms — 68% of the wall clock is the
elementwise .float() / *α / BF16-cast traffic, not the matmuls. Traffic
analysis agrees: at X=[M,4096], N=2048, r=8 the split is ~66% reading X, ~33%
writing Y, and only 0.26% for the U round trip. So fusing away U is not where
the headroom is; folding the scale and cast into the GEMM2 epilogue is, and
that is #58 (Step 10), out of scope here. It also means small-M rows are
launch-bound, which is why the TF32 column can land below FP32 — read the table
as relative overhead, not GEMM throughput.

Provenance

{
  "requested_backend": "p5-3-lora-delta-triton",
  "actual_backend": "triton-det_gemm-composed",
  "numeric_profile": "triton-det-gemm-fp32-accum",
  "reduction": "pinned BLOCK 64/64/32, no split-K, no autotune, allow_tf32=False",
  "oracle_agreement": "bitwise at K<=512; ~1e-6 relative by K=4096"
}
{
  "requested_backend": "p5-3-lora-delta-detgemm",
  "actual_backend": "cuda-det_gemm-composed",
  "numeric_profile": "det-gemm-midsplit-tree",
  "reduction": "det_gemm mid-split k-tree, leaf=32, no split-K, no atomics"
}

Validation environment

RTX A2000 8GB Laptop (Ampere, SM86, cc 8.6) · driver 596.08 · nvcc 12.4.131
(conda nvidia) · PyTorch 2.5.1+cu124 · Triton 3.1.0 · Python 3.10.21 ·
gcc 9.4.0 · Ubuntu 20.04 (WSL2) · rl_engine._C built for sm_86

Not SM90, so the CUDA column reflects det_gemm's scalar fallback rather than
its TMA path.

One question

numeric_profile has no enforcement behind it: the start kit says a backend
without equivalent capability "must register its own profile with an explicit
tolerance", but _compare in check_p5.py only diffs sha256 digests and never
reads the field. P5-3 lands exactly in that gap — neither backend can promise
byte equality at arbitrary K, so both declare their own profile and record the
observed bound in provenance(). Fine as is if the field stays declarative;
happy to propose a tolerance path in _compare if it should be machine-checked.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: cbc67b9d-7d3f-4540-8604-6a2d16859126

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[KJLdefeated]

Well-developed PR, reusing current det_gemm operator is good. Overall clean and clear, just left some comment on minor issues. Happy to merge after resolving the requests! Good Work!

[KJLdefeated]

Here y become fp32, but we can can keep it in bf16. In real system, we would prefer bf16 input and output.

[ScottNi0809]

Hi KJ, thanks for the review!

Makes sense as a production shape — the FP32 output doubles the bytes on a [M, 2*ffn] tensor, and it pairs with your fusion comment below.

The blocker is that I can't do it inside this PR alone: oracle.shared_grouped_lora_delta_fwd returns (y_fp32, u_bf16), and check_p5.py compares sha256 with zero tolerance, so a BF16 return diverges from the golden bytes. I measured the blast radius — switching it changes 7 of the 11 recorded boundaries on base_plus_lora, including swiglu_h and fc2_base, so it moves the golden bytes for P5-2 and P5-4 as well and the manifest needs regenerating.

Two things worth weighing if you do want it:

• The pipeline already casts at the end: y = (y_base + y_lora).to(torch.bfloat16) after summing base and LoRA. Casting inside the operator adds a rounding step before that sum — on base_plus_lora, 276/3072 elements move by up to 3.9e-3 relative.
• The pipeline itself won't break either way; z_base + z_lora promotes to FP32 regardless. Only the byte comparison changes.

Since it touches oracle.py and every sub-issue's golden bytes, it felt like your call rather than something to slip into P5-3.
Happy to send a follow-up switching all three providers plus the oracle in one go if you want the contract changed.

[KJLdefeated]

For speed issue, I think we can fuse two gemm together because LoRA operation have little r rank (little K dimension in gemm), so it can fuse into one kernel to boost the performance. But this is non-blocking optimization, if you think it is hard to do it. It is totally OK.

[ScottNi0809]

You're right that r=8 makes this a natural fusion candidate, and the profile backs it up: at these shapes everything is launch-bound — reading a [1,8] tensor costs about as much as reading [M,4096] — so the win would come from removing a kernel launch rather than from memory traffic. (The U round trip itself is only ~0.26% of the bytes moved.)

One constraint worth noting: u_bf16 is an input to the backward, so a fused kernel still has to write it out — it can save the read-back but not the store.

I did try it. The fused kernel is in my tree but not in this PR, because it lost: 0.79 ms vs 0.25 ms native at M=256, same run. My version was naive (one thread per output element, scalar K loop), so the saved launch was swamped by GEMM1 losing to cuBLAS tensor cores. A competitive version needs tensor cores for the [M,4096] read while keeping the pinned-tile / no-split-K schedule that batch invariance depends on, plus reproducing the oracle's rounding point for u.

Where I think it pays off most is when combined with the deferred epilogue-fusion step (Step 10 in the P5 sequencing): the same kernel would fold * alpha and the BF16 cast into the GEMM2 epilogue, and that is where the real headroom is — at M=64, 68% of the wall clock is elementwise dtype/scale traffic rather than the matmuls.
Happy to take that on as a follow-up if it's useful.

[KJLdefeated]

OK I see, let's keep current state as it now. The fused kernel would be future features.

[KJLdefeated]

LGTM.