Test nit

Signed-off-by: Luka Govedič <lgovedic@redhat.com>

tests/kernels/ir/test_layernorm.py

@@ -182,7 +182,7 @@ class TestFusedAddRMSNorm:
 
     def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
         x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
-            num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
+            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
         )
         out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps)
 

vllm/config/vllm.py

@@ -94,13 +94,14 @@ IS_DENSE = False
 
 def enable_norm_fusion(cfg: "VllmConfig") -> bool:
     """Enable if either RMS norm or quant FP8 custom op is active;
-    otherwise Inductor handles fusion."""
+    otherwise Inductor handles fusion. Also disable for batch invariant
+    as the custom fused kernels are not currently batch invariant."""
 
     return (
         cfg.compilation_config.is_custom_op_enabled("rms_norm")
         or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
         or cfg.kernel_config.ir_op_priority.rms_norm[0] != "native"
-    )
+    ) and not envs.VLLM_BATCH_INVARIANT
 
 
 def enable_act_fusion(cfg: "VllmConfig") -> bool:
@@ -117,7 +118,9 @@ def enable_act_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
-    """Enable if TP > 1 and Hopper/Blackwell and flashinfer installed."""
+    """Enable if TP > 1 and Hopper/Blackwell and flashinfer installed.
+    Disable if batch invariance is enabled.
+    """
     from vllm.platforms import current_platform
     from vllm.utils.flashinfer import has_flashinfer
 
@@ -136,6 +139,7 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
             current_platform.is_device_capability_family(100)
             or current_platform.is_device_capability(90)
         )
+        and not envs.VLLM_BATCH_INVARIANT
     )
 
 

vllm/ir/op.py

@@ -333,7 +333,8 @@ class IrOp:
         Context manager to set the dispatch priority for implementations for this op.
         """
         assert all(p in self.impls for p in priority), (
-            "All providers in priority must be registered implementations."
+            f"All providers in priority must be registered implementations, missing "
+            f"{','.join(p for p in priority if p not in self.impls)}"
         )
 
         def filter_priority_impls(p_list: list[str]) -> list[IrOpImpl]: