feat(linear-logp): add batch-invariant fused SM90 operator

[ThreeMonth03]

The SM90 fused linear_logp from #122 chooses its split-V schedule from the number of row blocks and the GPU SM count. This improves occupancy, but changing the batch size can also change the vocabulary partition and online-softmax merge order.

As a result, the existing path is deterministic for a fixed shape but does not guarantee batch invariance: the same token row may produce different output bits when evaluated alone, moved within a batch, or evaluated in different chunks. This matters when rollout and training regroup the same tokens but compare their log probabilities exactly.

Method

This PR adds a separate forward-only batch_invariant_linear_logp operator for SM90. It reuses the fused TMA + tensor-core kernel while making the reduction topology independent of the token-row count N:

• fixed hidden-dimension traversal with no split-K reduction;
• split-V derived only from V;
• fixed ascending merge of the partial online-softmax states.

The operator therefore preserves identical output bits across batch sizes, row positions, neighboring rows, and batch chunking without materializing the [N, V] logits tensor.

The initial scope is intentionally limited to single-card Hopper, BF16 hidden states and weights, and forward execution. Tensor parallelism, backward, and fallback backends are not included.

Verification

Tests cover FP32 reference accuracy and exact batch invariance across row positions, batch sizes, chunk sizes, vocabulary boundaries, layouts, and CUDA streams. The code compiles for CUDA 12.4 and 13.1, and the non-H100 test and lint suites pass.

The real SM90 launch and bitwise checks still require H100 GPU CI, so this PR remains a draft.

Related to #122. Potentially overlaps with the broader strict mode in #336.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

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.}