Move imports out of line to fix compilation

Signed-off-by: Bram Wasti <bwasti@meta.com>

Author: bwasti

vllm/distributed/device_communicators/all_reduce_utils.py

@@ -19,6 +19,9 @@
 import vllm.envs as envs
 from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from vllm.logger import init_logger
+from vllm.model_executor.layers.batch_invariant import (
+    vllm_kernel_override_batch_invariant,
+)
 from vllm.utils import cuda_device_count_stateless, update_environment_variables
 
 logger = init_logger(__name__)
@@ -70,9 +73,6 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor)
     from vllm.distributed.device_communicators.pynccl_allocator import (
         is_symmetric_memory_enabled,
     )
-    from vllm.model_executor.layers.batch_invariant import (
-        vllm_kernel_override_batch_invariant,
-    )
 
     if vllm_kernel_override_batch_invariant():
         return False

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -1134,11 +1134,8 @@ def fused_topk_bias(
     scores_for_choice = scores.view(
         -1, n_routed_experts
     ) + e_score_correction_bias.unsqueeze(0)
-    # For batch invariance, use sorted=True to ensure deterministic expert selection
-    from vllm.model_executor.layers.batch_invariant import (
-        vllm_kernel_override_batch_invariant,
-    )
 
+    # For batch invariance, use sorted=True to ensure deterministic expert selection
     use_sorted = vllm_kernel_override_batch_invariant()
     topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[1]
     topk_weights = scores.gather(1, topk_indices)
@@ -1201,11 +1198,8 @@ def grouped_topk(
         group_scores = (
             scores.view(num_token, num_expert_group, -1).max(dim=-1).values
         )  # [n, n_group]
-    # For batch invariance, use sorted=True to ensure deterministic expert selection
-    from vllm.model_executor.layers.batch_invariant import (
-        vllm_kernel_override_batch_invariant,
-    )
 
+    # For batch invariance, use sorted=True to ensure deterministic expert selection
     use_sorted = vllm_kernel_override_batch_invariant()
     group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=use_sorted)[
         1

vllm/v1/attention/backends/mla/triton_mla.py

@@ -13,6 +13,9 @@
 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
 from vllm.attention.ops.triton_flash_attention import triton_attention
 from vllm.logger import init_logger
+from vllm.model_executor.layers.batch_invariant import (
+    vllm_kernel_override_batch_invariant,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON
 from vllm.v1.attention.backends.mla.common import (
@@ -159,10 +162,6 @@ def _forward_decode(
         )
         lse = torch.zeros(B, q_num_heads, dtype=q.dtype, device=q.device)
 
-        from vllm.model_executor.layers.batch_invariant import (
-            vllm_kernel_override_batch_invariant,
-        )
-
         # For batch invariance, use only 1 split to ensure deterministic reduction
         num_kv_splits = 1 if vllm_kernel_override_batch_invariant() else 4