[lora] Remove synchronous .any().item() guard in LoRA MoE prefill path (#25531)
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
co-authored by
Cursor
parent
ca9dc17be4
commit
a24c374f84
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user