Skip to content

Conversation

bwasti
Copy link
Contributor

@bwasti bwasti commented Oct 2, 2025

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.

@bwasti bwasti requested a review from mgoin as a code owner October 2, 2025 21:23
@mergify mergify bot added the v1 label Oct 2, 2025
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Contributor Author

bwasti commented Oct 2, 2025

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

@bwasti bwasti force-pushed the det_moe branch 3 times, most recently from f867388 to 648ce25 Compare October 2, 2025 22:40
Signed-off-by: Bram Wasti <[email protected]>
Copy link
Member

@yewentao256 yewentao256 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@yewentao256 yewentao256 changed the title Fix RMS and test MoE for batch invariance [4/n] [WIP] Fix RMS and test MoE for batch invariance [4/n] Oct 3, 2025
Copy link

mergify bot commented Oct 4, 2025

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

@mergify mergify bot added the needs-rebase label Oct 4, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
Status: In Progress
Development

Successfully merging this pull request may close these issues.

2 participants