fix(gemm)[1th merge]: BF16 store + mid-split K-tree for simulated TP=2

[frank-2077]

Summary

• GEMM still computes in FP32 and only stores the final result as BF16.
• K is reduced with a mid-split tree: FP32 leaves, BF16 adds at internal nodes.
• Simulated TP=2: left + right is the root add, so it matches TP=1 bitwise.
• Simulated TP=8 left-fold uses a different parenthesization and is expected to diverge. Real TP is not wired yet.

Test plan

• pytest tests/test_det_gemm.py tests/distributed/test_det_gemm_simulated_tp.py -k "not test_target_shapes"

[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: Pro Plus

Run ID: 339b6fc6-dcc7-47d0-960a-330df1a40f5e

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

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[zhangj1an]

Thanks for making the K-reduce order TP-invariant!

Previously the fixed-K order folded tiles by ascending index. Now it becomes a mid-split tree, which matches the fixed tree in @CyberSecurityErial's AllReduce (#310). It looks like this (e.g. suppose there are total 8 tiles, kd = 8):

                [0,8)                  <- internal node (root)
               /     \
         [0,4)         [4,8)           <- internal node
        /    \        /    \
    [0,2)  [2,4)  [4,6)  [6,8)         <- internal node
    /  \    /  \   /  \   /  \
   t0  t1  t2 t3  t4 t5  t6 t7         <- leaf = tile (BK = 32)

So as long as K % (32 * TP) == 0, each rank's K-range is exactly a node of the global tree, and TP = 2, 4, 8 are all bitwise identical to TP = 1.

But I think the accumulator should stay FP32 all the way, with a single downcast at the end. Right now this PR downcasts at every tile (L246 __float2bfloat16) and again at every internal node (L256 bf16_add). For K = 2048 that is 64 leaf (tile) roundings plus 63 internal-node roundings, so every output element goes through 7 chained roundings instead of 1. I simulated the mid-split tree with a BF16 round at every leaf and every internal node, and at M,K,N = 128,2048,2048 it misses the contract thresholds (atol=0.05, rtol=0.02 from rl_engine/kernels/gtest/tolerance_contract.json) on ~8.5% of elements, max abs error 1.53. With FP32 internal nodes it passes with max abs error 0.50. Is this intended?

The AllReduce kernel already accepts FP32 (ordered_add has an add.rn.f32 specialization), so it's fine to keep everything in FP32 and only downcast after the reduction. That path is exactly what det_gemm_fwd_fp32 was for, but this PR changed it to gemm_dispatch(a, b), so it now returns BF16. So I want to clarify.

Otherwise the kernel change looks good to me! Thanks again