[WIP] Fix RMS and test MoE for batch invariance [4/n]

[bwasti]

This change gets Qwen/Qwen3-30B-A3B on 1GPU working at bitwise parity across batch sizes.
It also improves the e2e testing logic (makes it more harsh). Please only look at the last commit (it is rebased on #25769)

Purpose

Add RMS kernel directly and override in Python. The change to the csrc version doesn't provide full coverage, but does increase invariance.

Test Plan

VLLM_TEST_MODEL="Qwen/Qwen3-30B-A3B" VLLM_ATTENTION_BACKEND=FLASHINFER VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT=1 HF_HUB_DISABLE_XET=1 pytest -s -v tests/v1/generation/test_batch_invariance.py -k test_logprobs_bitwise_batch_invariance_bs1_vs_bsN

Test Result

Pass

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces significant improvements for batch invariance, particularly for RMSNorm and MoE layers, and strengthens the e2e testing framework. The addition of a batch-invariant Triton kernel for RMSNorm and the corresponding Python overrides are key changes. The updates to the testing logic, making it more rigorous by using randomized batch sizes and prompts, are a great enhancement. The modifications for the FlashInfer backend to ensure deterministic behavior under batch invariance are also well-implemented. However, I've identified a critical issue in the fused_add_rms_norm implementation and a high-severity issue in the MoE softmax kernel that could undermine the goal of batch invariance.

[bwasti]

btw, I'm counting on some RMS test to run here. this shouldn't land until that passes

[yewentao256]

btw, I'm counting on some RMS test to run here. this shouldn't land until that passes

Sounds great, please change the title once it is ready

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @bwasti.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork