[vLLM IR] 2/N batch-invariant-aware dispatching and rms_norm

[ProExpertProg]

Purpose

Test Plan

Test Result

---

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 the concept of batch-invariance for IR ops, which is a good step towards more efficient dispatching. A new has_reduction flag is added to IrOp to control the default batch_invariant status of its implementations. However, I've found a critical issue in how the batch_invariant flag is set for native PyTorch implementations. It's hardcoded to True, which is incorrect for reduction operations and inconsistent with how other implementations are handled. This could lead to incorrect behavior and numerical results when the dispatcher is implemented. My review includes a suggestion to fix this inconsistency.

[gemini-code-assist]

The native implementation is being hardcoded as batch_invariant=True. This seems incorrect for reduction operations and is inconsistent with the handling of other implementations.

For an op registered with has_reduction=True, its native PyTorch implementation (e.g., torch.sum) is generally not batch-invariant. Applying it to a batched tensor would typically reduce across the batch dimension, which is not the desired behavior for batch processing.

The logic for determining batch_invariant for non-native implementations (which defaults to not self.has_reduction) is sound and should also be applied to the native implementation for consistency. This ensures that all implementations for a reduction op are correctly marked as non-batch-invariant by default.

        self.impls["native"] = IrOpImpl(
            self,
            "native",
            native_impl,
            # always supported
            supported=True,
            supports_args=None,
            batch_invariant=not self.has_reduction,
        )