From 2b769d37a41d2b96b28d8bf1168113c8246dd222 Mon Sep 17 00:00:00 2001 From: "Ethan (Yusheng) Su" Date: Tue, 5 May 2026 03:11:28 +0200 Subject: [PATCH] (2/n - prefill optimize)perf(lora): remove GPU-CPU sync barrier (.item()) in MoE LoRA path and remove duplicate code (#24246) Co-authored-by: Cursor --- python/sglang/srt/lora/lora_moe_runners.py | 129 +-------------------- 1 file changed, 4 insertions(+), 125 deletions(-) diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 3a08cb782..b3f1389b5 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -34,7 +34,6 @@ from sglang.srt.utils import is_cuda, is_hip, is_xpu, next_power_of_2 _is_cuda = is_cuda() _is_hip = is_hip() -_is_hip = is_hip() _is_xpu = is_xpu() if _is_cuda or _is_hip or _is_xpu: @@ -64,112 +63,6 @@ def _get_moe_lora_block_config(max_lora_rank: int) -> dict: _SPARSITY_FACTOR = 8 -def _naive_moe_lora_align_block_size( - topk_ids: torch.Tensor, - seg_indptr: torch.Tensor, - req_to_lora: torch.Tensor, - num_experts: int, - block_size_m: int, - max_loras: int, - max_num_tokens_padded: int, - max_num_m_blocks: int, - adapter_enabled: torch.Tensor, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Construct LoRA token-expert alignment on CPU for small batches. - - When the number of tokens is very small, the overhead of launching the - CUDA-based moe_lora_align_block_size kernel exceeds the actual - computation. This function builds the same data structures using simple - Python loops on CPU and transfers the result to GPU in one shot. - """ - M, top_k = topk_ids.shape - num_valid_tokens = M * top_k - - sorted_token_ids = torch.full( - (max_loras * max_num_tokens_padded,), - num_valid_tokens, - dtype=torch.int32, - ) - expert_ids_out = torch.full((max_loras * max_num_m_blocks,), -1, dtype=torch.int32) - num_tokens_post_padded = torch.zeros(max_loras, dtype=torch.int32) - - seg_indptr_list = seg_indptr.cpu().tolist() - req_to_lora_list = req_to_lora.cpu().tolist() - topk_ids_list = topk_ids.cpu().tolist() - adapter_enabled_list = adapter_enabled.cpu().tolist() - - for lora_id in range(max_loras): - if not adapter_enabled_list[lora_id]: - continue - - pairs: list[tuple[int, int]] = [] - for seg_idx in range(len(seg_indptr_list) - 1): - if req_to_lora_list[seg_idx] != lora_id: - continue - start = seg_indptr_list[seg_idx] - end = seg_indptr_list[seg_idx + 1] - for m in range(start, end): - for k in range(top_k): - pairs.append((topk_ids_list[m][k], m * top_k + k)) - - if not pairs: - continue - - pairs.sort() - - base_t = lora_id * max_num_tokens_padded - base_e = lora_id * max_num_m_blocks - pos = 0 - block_idx = 0 - i = 0 - while i < len(pairs): - cur_expert = pairs[i][0] - group_start = pos - while i < len(pairs) and pairs[i][0] == cur_expert: - sorted_token_ids[base_t + pos] = pairs[i][1] - pos += 1 - i += 1 - group_len = pos - group_start - padded_len = ((group_len + block_size_m - 1) // block_size_m) * block_size_m - num_blocks = padded_len // block_size_m - for b in range(num_blocks): - expert_ids_out[base_e + block_idx + b] = cur_expert - block_idx += num_blocks - pos = group_start + padded_len - - num_tokens_post_padded[lora_id] = pos - - return ( - sorted_token_ids.to(device), - expert_ids_out.to(device), - num_tokens_post_padded.to(device), - ) - - -def _get_moe_lora_block_config(max_lora_rank: int) -> dict: - """Compute rank-aware block sizes for MoE LoRA kernels. - - Shrink: output dim is the rank -> cap BLOCK_SIZE_N to avoid waste. - Expand: input dim is the rank -> cap BLOCK_SIZE_K similarly. - """ - if max_lora_rank <= 0: - rank_pow2 = 64 - else: - rank_pow2 = next_power_of_2(max_lora_rank) - - shrink_n = min(64, rank_pow2) - expand_k = max(16, min(64, rank_pow2)) - - return { - "shrink_block_size_n": shrink_n, - "expand_block_size_k": expand_k, - } - - -_SPARSITY_FACTOR = 8 - - def _naive_moe_lora_align_block_size( topk_ids: torch.Tensor, seg_indptr: torch.Tensor, @@ -288,7 +181,6 @@ class LoRAInfo: num_experts: int experts_shared_outer_loras: bool = False cg_buffers: dict | None = None - cg_buffers: dict | None = None fully_sharded: bool = False tp_size: int = 1 @@ -447,18 +339,10 @@ def _add_lora_gate_up_delta( from sglang.srt.model_executor.cuda_graph_runner import get_capture_lora_variant # Record LoRA kernels for lora graph; skip for nolora graph. - has_active_lora = get_capture_lora_variant() != "nolora" - 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 get_capture_lora_variant() == "nolora": + return + + if lora_info is None or lora_info.max_lora_rank == 0: return M, top_k, gate_up_dim = intermediate_cache.shape @@ -565,11 +449,6 @@ def _add_lora_down_delta( if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts: down_lora_b = down_lora_b.expand(-1, lora_info.num_experts, -1, -1) - if lora_info.fully_sharded and lora_info.tp_size > 1: - shard_size = lora_info.hidden_size // lora_info.tp_size - offset = shard_size * lora_info.tp_rank - else: - offset = 0 if lora_info.fully_sharded and lora_info.tp_size > 1: shard_size = lora_info.hidden_size // lora_info.tp_size offset = shard_size * lora_info.tp_rank