Deterministic Bmm Batch Scaling
Remove per-item launch overhead from deterministic batched matrix multiplication.
bmm_batch_invariant launches a separate matrix multiplication for each batch item, and the launch overhead grows with batch size. Replace this loop with a deterministic Triton implementation that handles the batch without making the number of launches grow with it. Internal function and kernel names are up to you.
Keep the existing behavior for 3-D FP16, BF16, and FP32 inputs: a batched call must be bitwise identical to concatenated single-batch calls and agree with torch.bmm at rtol=2e-2, atol=2e-2. Preserve input dtype and device validation, existing shape errors, and out= copy semantics. Keep the existing batch-count behavior: ignore extra right-hand batches, and reject a shorter right-hand batch or an empty left-hand batch; exception types and messages need not stay identical. An output destination may have a different dtype, be on the CPU, or have a broadcast-compatible shape; copy the result into it with the required conversion and return that same object. Incompatible destination shapes must still fail.
On one NVIDIA A100-SXM4-40GB, the BF16 implementation should be at least 2.0x faster than the original per-batch path for (B, M, N, K) = (8, 512, 512, 2560) and at least 1.05x faster for each of (32, 512, 512, 2560) and (8, 1280, 1280, 2560). Correctness and launch scaling must hold across supported inputs, beyond these performance cases.
For the performance comparison, warm up each implementation five times and synchronize after its warmup. Time 20 direct calls per round with CUDA events, without CUDA graph replay, and wait for the end event before reading elapsed time. Divide each round's time by 20, then use the median of five rounds for each implementation to calculate the speedup.