fix(true-on-policy): match SGLang's row-linear k-tile reduction order

[zihaow211]

Problem

_sglang_row_parallel_matmul combines the 128-wide K partials with a plain binary tree
(_fixed_tree_sum_tensors), but SGLang's matmul_tp_inv does not reduce that way. Its
kernel accumulates FIRST_LEVEL_BLOCK partials sequentially into one register
accumulator before carrying into the binary levels above it
(table_value = FIRST_LEVEL_BLOCK if level == 0 else 2), and FIRST_LEVEL_BLOCK comes
from halving the k-tile count while it stays even:

T = K // BLOCK_K
FIRST_LEVEL_BLOCK = T
while FIRST_LEVEL_BLOCK > 2 and FIRST_LEVEL_BLOCK % 2 == 0:
    FIRST_LEVEL_BLOCK //= 2

The two orders coincide only when FIRST_LEVEL_BLOCK == 2, i.e. when T is a power of
two. For Qwen3-4B fc2/down_proj at TP=2 the local K = 4864 gives T = 38, so
FIRST_LEVEL_BLOCK = 19 and SGLang computes (p0 + ... + p18) + (p19 + ... + p37)
while Megatron computes a full binary tree over all 38. Accumulation is in bf16
(acc_dtype = A.dtype unless fp32_accum), so the association order is visible in the
output bits and down_proj diverges from rollout even though every partial product is
identical.

Fix

Add _sglang_kernel_order_sum, which accumulates each first-level block sequentially and
then reduces the blocks with the existing binary tree. The block count is always
2**(LEVEL_K - 1), so the tree is exact and no partial is left unpaired.

_fixed_tree_sum_tensors is left as-is: it is still the right shape for the cross-rank
reduction, where it matches SGLang's tree_all_reduce_sum over a power-of-two world
size. Only the k-tile reduction inside _sglang_row_parallel_matmul changes. The
addition count is unchanged (n - 1 either way); only the association differs.

The existing test_sglang_row_parallel_matmul_uses_fixed_k_block_order could not have
caught this — it compares sglang_reference_matmul(..., row_parallel=True) against
_sglang_row_parallel_matmul, which is the function it dispatches to, so it asserts the
function equals itself. Its K=256 also gives T=2, inside the power-of-two blind spot.
The new tests assert the association order directly.

Validation

Compared against the real torch.ops.tp_inv_ops.matmul_tp_inv triton kernel on 1×B200
(sm_100), torch 2.7 / triton 3.3, bf16, M=128, N=2560.

At T=38 (K=4864, the Qwen3-4B down_proj shard at TP=2):

• before: 234686 / 327680 elements differ, max abs diff 4.0
• after: bitwise identical, 0 / 327680

Sweeping T confirms the power-of-two rule, and that the fix holds everywhere:

T	4	6	8	12	16	19	24	32	38	40	48	64
FIRST_LEVEL_BLOCK	2	3	2	3	2	19	3	2	19	5	3	2
before	ok	x	ok	x	ok	x	x	ok	x	x	x	ok
after	ok	ok	ok	ok	ok	ok	ok	ok	ok	ok	ok	ok

After the fix _sglang_row_parallel_matmul is bitwise identical to matmul_tp_inv for
every T above, for the bias path, and for 3-D inputs. New unit tests pass; they need no
GPU and no SGLang.

Out of scope

_sglang_rollout_partition_row_parallel_matmul also reduces with _fixed_tree_sum_tensors,
but that is correct: it runs only when SGLang's should_use_tp_invariant_row_linear gate
fails, where SGLang itself falls back to one plain GEMM per rollout rank plus the tree
all-reduce. I checked it against that equivalent over 100 configurations
(K ∈ {9728, 2560}, train/rollout TP ∈ {1,2,4}×{8,16}, M ∈ {1, 8, 128, 2048, 8192},
bf16 and fp16) — all bitwise identical, including the non-contiguous slices it hands
cuBLAS. Unchanged here.

Separately, RowParallelLinear adds bias after the TP reduce while SGLang fuses it into
rank 0's GEMM before the reduce, which is not bitwise equivalent in bf16. Qwen3 is
unaffected (down_proj is bias=False, released Qwen3 configs set
attention_bias: false), and o_proj at TP=2 has 16 k-tiles so it is untouched by this
PR. Filed separately.

Fixes radixark/miles#1485

Thanks to @MeouSker77 for the report and the root-cause analysis.