Expose Batch Invariant Kernels

[wdykas]

Using batch invariant kernels has become common for many post training workloads to get 0 log prob mismatch between training and inference. We can monkey patch around this in TE but it would be very useful if TE exposed efficient batch invariant kernels for common operations like grouped gemm and regular gemm. This would be very useful for our team megatron inference and megatron RL for larger experimentation and a better interface.

A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]

Describe the solution you'd like

batch invariant mode in TE

Describe alternatives you've considered

I am currently monkey patching general grouped gemm, rmsnorm, and general gemm with my own inefficient kernels.

Additional context

Add any other context or screenshots about the feature request here.

[0z5a]

@sbhavani @wdykas I’d like to help upstream the missing self-contained Batch Invariant Kernel support here.

I saw #2977 already depends on Megatron’s BIK implementation for its TP-invariant path, so I don’t want to duplicate its TP orchestration. I’d instead start with one self-contained TE primitive, preferably regular GEMM, behind an opt-in batch-invariant mode.

For the first slice I’d validate:

• identical sample output across different batch compositions;
• different sample positions / padding layouts;
• forward and backward behavior;
• multi-L20 TP=2/4 where applicable;
• performance relative to the current TE path.

If regular GEMM is the preferred first primitive, I can start there and leave Grouped GEMM / RMSNorm as follow-ups. If you’d rather establish the common BIK interface first, I’m also happy to scope the first PR around that.

Would that be a useful slice for me to take?

[0z5a]

I took the regular-GEMM slice I proposed and have measurements. Short version: the current BF16 forward is not batch-invariant on L20, and it is not the split-accumulator flag — it is cuBLASLt picking the kernel from the whole problem shape. A self-contained path can make it invariant, at a small cost.

Environment: L20 (cc 8.9), TE main @ 9233c73a built with NVTE_CUDA_ARCHS=89, torch 2.11.0+cu128, CUDA 12.8, cuDNN 9.19. BF16, N=K=2048, eval, no dropout.

1. Current behaviour — one anchor row placed at different offsets in batches of different M, output bytes compared exactly:

M	1	2	17	33	128	257	1024
anchor row bitwise == M=1 result	yes	no	no	yes	no	no	no

5 of 7 M values disagree. The disagreement is independent of the row's position within the batch (pos 0 / middle / last give the same answer), so it is driven by M, not by where the row sits.

2. It is not use_split_accumulator. Forcing it off and on produced the identical 5/7 pattern. The reduction order is coming from the heuristic's choice of algorithm/workspace, which is a function of the full problem shape. (NVTE_EXT_MARGIN_SM was left at its default.)

3. A self-contained path

I added an opt-in Triton kernel: each output tile accumulates over K in a fixed BLOCK_K loop, and the tile grid derives from tile coordinates only, so nothing inside the kernel depends on the total M. Result:

check	result
7 M values × 3 row positions	18/18 bitwise equal (max ULP 0)
row permutation (anchor moved to a different row index in a shuffled batch)	bitwise equal
non-tile-aligned M=1000	bitwise equal
accuracy vs FP64	max abs 2.34e-4, same as the cuBLAS path
unsupported combinations (fp32, K mismatch, 3-D input)	rejected, not silently falling back

4. Cost (eager, warmup before timing)

M	N	K	default	batch-invariant	ratio
1	4096	4096	0.046 ms	0.081 ms	0.57×
8	4096	4096	0.050 ms	0.074 ms	0.67×
128	4096	4096	0.084 ms	0.104 ms	0.81×
1024	4096	4096	0.348 ms	0.379 ms	0.92×
4096	4096	4096	1.251 ms	1.345 ms	0.93×
4096	11008	4096	3.237 ms	3.585 ms	0.90×
4096	4096	11008	3.612 ms	3.647 ms	0.99×
8192	4096	4096	2.486 ms	2.686 ms	0.93×

Median 0.91×: 2.5–10% slower at realistic shapes, 1.4–1.8× at very small M where launch overhead dominates. I would not hide that behind "more numerically stable".

5. Scope and what I am not claiming

• BF16 only, Y = X @ W.T, contiguous 2-D operands, no bias, no quantization, no split-K.
• default path untouched; opt-in only.
• tested with eval/no-dropout and the same weights, so this is batch-composition invariance of one GEMM, not of a model. A whole-model logit-invariance claim would additionally need every other op in the graph to have the property.
• not claiming cross-TP-degree invariance, and not claiming cross-architecture/cuBLAS-version invariance — L20 cannot demonstrate either.
• the new path is verified standalone; I have not wired it into te.Linear's dispatch yet. So today it is a validated primitive, not a drop-in "TE Linear is now batch-invariant".

Question before I send anything

The commit in #2977 deliberately moved the FP32 TP-reduce out as a separable feature, and the reviewer there asked for TE's behaviour to be self-contained rather than depending on Megatron's kernels. Before I write a PR I would like to know which of these you want:

1. this primitive behind an explicit opt-in flag, with TE owning the kernel and its tests (my current shape), or
2. a common BIK interface first, with the GEMM as one implementation behind it, or
3. something that reuses an existing numeric-mode abstraction in TE rather than adding a parallel one.

I have not opened a PR; I would rather agree the interface first than add a second numeric mode that later has to be reconciled.

[0z5a]

Opened #3556 — opt-in batch-invariant BF16 GEMM with a fixed tile geometry and reduction order. 18/18 row slices are bitwise identical to the full-batch result (the general path differed on 5/7 splits on L20 and 17/18 on 5090); median cost 0.91×. Not wired into te.Linear's dispatch, so existing users are unaffected.