[small][batch invariance] Rename the env and internal flags to simplify usage #26855
Conversation
There was a problem hiding this comment.
Code Review
This pull request primarily renames an environment variable and an internal function related to batch invariance to simplify their usage. However, it also introduces a substantial set of changes to properly implement and enforce this batch invariance feature across the codebase. This includes adding deterministic custom kernels for operations like RMS norm and matrix multiplication, disabling potentially non-deterministic optimizations when the feature is active, and significantly expanding the test suite to validate the deterministic behavior. The changes are well-structured and align with the goal of ensuring bit-for-bit determinism, which is valuable for debugging and testing. The new tests for the RMS norm implementation are particularly thorough. I have one suggestion to improve the robustness of the new deterministic matrix multiplication function.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
be130ac to
bc96863
Compare
35cb4b4 to
7e7bc78
Compare
Signed-off-by: Bram Wasti <bwasti@meta.com>
Environment variable VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT -> VLLM_BATCH_INVARIANT
Internal flag: vllm_kernel_override_batch_invariant() -> vllm_is_batch_invariant()
Purpose
Simplify usage across the stack.
Test Plan
Rebuild for the C++ component:
pytest -s -v tests/v1/generation/test_batch_invariance.py -k test_logprobs_bitwise_batch_invariance_bs1_vs_bsNTest Result
Pass.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.