[lora] Remove synchronous .any().item() guard in LoRA MoE prefill path (#25531)

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Ethan (Yusheng) Su
2026-05-21 23:58:57 +08:00
committed by GitHub
co-authored by Cursor
parent ca9dc17be4
commit a24c374f84
2 changed files with 14 additions and 18 deletions
+5
View File
@@ -980,6 +980,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
)
req_to_lora = wi
# Single source of truth: lora_manager precomputes this per-batch from
# the Python weight_indices list, no GPU sync needed.
has_active_lora = bool(getattr(batch_info, "has_active_lora", False))
return LoRAInfo(
gate_up_lora_a_weights=self.gate_up_lora_a_weights,
gate_up_lora_b_weights=self.gate_up_lora_b_weights,
@@ -989,6 +993,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
req_to_lora=req_to_lora,
lora_ranks=lora_ranks,
adapter_enabled=adapter_enabled,
has_active_lora=has_active_lora,
max_lora_rank=max_lora_rank,
num_experts=self.base_layer.num_experts,
experts_shared_outer_loras=self.experts_shared_outer_loras,
+9 -18
View File
@@ -179,6 +179,7 @@ class LoRAInfo:
max_lora_rank: int # Maximum LoRA rank across all adapters
num_experts: int
has_active_lora: bool = True
experts_shared_outer_loras: bool = False
cg_buffers: dict | None = None
@@ -335,24 +336,9 @@ def _add_lora_gate_up_delta(
merged_experts_fused_moe_lora_add,
)
if get_is_capture_mode():
# During CUDA graph capture, always enter the LoRA path so that
# the LoRA kernels are recorded in the graph. adapter_enabled is
# all-zeros during capture, so the Triton kernel early-exits per
# program (zero overhead). During replay the tensor is updated
# in-place with the real adapter mask before graph.replay().
has_active_lora = True
else:
num_loras = len(lora_info.lora_ranks)
has_active_lora = (
(
lora_info.adapter_enabled[:num_loras]
* (lora_info.lora_ranks > 0).to(lora_info.adapter_enabled.dtype)
)
.any()
.item()
)
if not has_active_lora or lora_info is None or lora_info.max_lora_rank == 0:
if lora_info is None or lora_info.max_lora_rank == 0:
return
if not get_is_capture_mode() and not lora_info.has_active_lora:
return
M, top_k, gate_up_dim = intermediate_cache.shape
@@ -520,6 +506,11 @@ def build_lora_hooks(
"""
if lora_info is None or lora_info.max_lora_rank == 0:
return LoRAHooks()
# Skip alignment/mapping work entirely when the batch has no active adapter.
# During CUDA graph capture we still need to record the kernels into the
# graph (adapter_enabled is all-zero, kernels early-exit on GPU).
if not get_is_capture_mode() and not lora_info.has_active_lora:
return LoRAHooks()
# Compute alignment / mapping (once, shared by both hooks)
token_lora_mapping: torch.Tensor | None = None