Fix - moe lora support

[yushengsu-thu]

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[Copilot]

Pull request overview

Updates MoE LoRA support by adjusting the per-expert Triton kernel implementation and aligning MoE LoRA A-buffer shapes with stacked module requirements.

Changes:

• Update the per-expert MoE LoRA Triton kernel to handle stacked gate_up_proj LoRA weights by computing gate/up halves separately.
• Change MoE LoRA A buffer allocation to size the rank dimension based on the specific MoE module stacking factor.
• Add commented alternative model/adapter options in the MoE LoRA logprob comparison test.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

File	Description
python/sglang/srt/lora/triton_ops/per_expert_lora_moe.py	Reworks the Triton kernel logic for MoE per-expert LoRA, including special handling for stacked gate_up_proj.
python/sglang/srt/lora/mem_pool.py	Adjusts MoE LoRA A-buffer rank dimension sizing to account for per-module stacking factor.
test/registered/lora/test_lora_hf_sgl_logprob_diff.py	Adds commented-out alternative model and LoRA adapter entries for the MoE comparison test.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The file now contains a large, fully commented-out duplicate implementation (old kernel + forward) at the top of the module. This significantly bloats the source and makes future maintenance/review harder. Please remove the commented-out block and rely on git history if the old version needs to be referenced.

Suggested change

	# """Per-expert LoRA computation kernel for MoE layers."""
	# import torch
	# import triton
	# import triton.language as tl
	# @triton.jit
	# def _per_expert_lora_kernel(
	# # Input/Output pointers
	# hidden_states_ptr, # [num_total_tokens, input_dim]
	# lora_a_weights_ptr, # [num_loras, num_experts, max_rank, input_dim]
	# lora_b_weights_ptr, # [num_loras, num_experts, output_dim, max_rank]
	# output_ptr, # [num_total_tokens, output_dim] - base output (modified in-place)
	# lora_output_ptr, # [num_total_tokens, output_dim] - separate LoRA-only output
	# # Dispatch info (length = num_dispatched)
	# token_ids_ptr, # [num_dispatched] -> index into hidden/output
	# expert_ids_ptr, # [num_dispatched]
	# lora_ids_ptr, # [num_dispatched]
	# # Dimensions
	# input_dim: tl.constexpr,
	# output_dim: tl.constexpr,
	# max_rank: tl.constexpr,
	# num_experts: tl.constexpr,
	# num_dispatched,

[Copilot]

In the gate_up_proj path, the half (gate vs up) is chosen based on out_start (slice-level). If half_output_dim is not an exact multiple of BLOCK_SIZE, the boundary slice will contain indices from both halves, but this code will compute only one half and write zeros for the other half, producing incorrect outputs for those positions. Please make the half-selection elementwise based on out_offs (or restructure the grid so slices never straddle the boundary, e.g., add a separate grid dimension for {gate,up} and slice within half_output_dim).

Suggested change

	half_output_dim = output_dim // 2 # inter_dim
	# Determine which half this slice belongs to (gate or up)
	# Use tl.where to avoid type mismatch between branches
	is_up_half = out_start >= half_output_dim
	# For UP part: local_out_offs = out_offs - half_output_dim, a_row_offset = rank
	# For GATE part: local_out_offs = out_offs, a_row_offset = 0
	# Use multiplication to get consistent types: rank * 1 or rank * 0
	a_row_offset = rank * tl.where(is_up_half, 1, 0)
	b_row_offset = half_output_dim * tl.where(is_up_half, 1, 0)
	local_out_offs = out_offs - half_output_dim * tl.where(is_up_half, 1, 0)
	local_out_mask = (local_out_offs >= 0) & (local_out_offs < half_output_dim)
	global_out_mask = out_offs < output_dim
	# Stage 1: Compute intermediate for the appropriate half (gate or up)
	# Read from the correct half of stacked A matrix
	intermediate = tl.zeros((max_rank,), dtype=tl.float32)
	NUM_INPUT_TILES = (input_dim + BLOCK_SIZE - 1) // BLOCK_SIZE
	for input_tile_idx in range(NUM_INPUT_TILES):
	input_start = input_tile_idx * BLOCK_SIZE
	input_offs = input_start + tl.arange(0, BLOCK_SIZE)
	input_mask = input_offs < input_dim
	h_vals = tl.load(hidden_ptr + input_offs, mask=input_mask, other=0.0).to(tl.float32)
	# Read from the appropriate half of A (offset by a_row_offset)
	a_ptrs = (
	lora_a_base
	+ (r_offs[:, None] + a_row_offset) * lora_a_stride_rank
	+ input_offs[None, :] * lora_a_stride_input
	)
	a_vals = tl.load(a_ptrs, mask=rank_mask[:, None] & input_mask[None, :], other=0.0).to(tl.float32)
	intermediate += tl.sum(a_vals * h_vals[None, :], axis=1)
	# Stage 2: Compute output using the appropriate half of B
	# B is indexed using local_out_offs (0 to inter_dim-1 for either gate or up)
	# We add b_row_offset to access the correct half in stacked B
	b_ptrs = (
	lora_b_base
	+ (local_out_offs[None, :] + b_row_offset) * lora_b_stride_output
	+ r_offs[:, None] * lora_b_stride_rank
	)
	b_vals = tl.load(b_ptrs, mask=rank_mask[:, None] & local_out_mask[None, :], other=0.0).to(tl.float32)
	out_vals = tl.sum(b_vals * intermediate[:, None], axis=0)
	half_output_dim = output_dim // 2 # inter_dim
	# Elementwise half selection based on out_offs
	# gate half: out_offs < half_output_dim
	# up half: out_offs >= half_output_dim
	is_up_half = out_offs >= half_output_dim
	# Local offsets within each half
	local_out_offs_gate = out_offs
	local_out_offs_up = out_offs - half_output_dim
	gate_out_mask = (local_out_offs_gate >= 0) & (local_out_offs_gate < half_output_dim)
	up_out_mask = (local_out_offs_up >= 0) & (local_out_offs_up < half_output_dim)
	global_out_mask = out_offs < output_dim
	# Stage 1: Compute intermediates for BOTH halves (gate and up)
	# gate_inter uses A rows [0 .. rank-1], up_inter uses A rows [rank .. 2*rank-1]
	gate_inter = tl.zeros((max_rank,), dtype=tl.float32)
	up_inter = tl.zeros((max_rank,), dtype=tl.float32)
	NUM_INPUT_TILES = (input_dim + BLOCK_SIZE - 1) // BLOCK_SIZE
	for input_tile_idx in range(NUM_INPUT_TILES):
	input_start = input_tile_idx * BLOCK_SIZE
	input_offs = input_start + tl.arange(0, BLOCK_SIZE)
	input_mask = input_offs < input_dim
	h_vals = tl.load(hidden_ptr + input_offs, mask=input_mask, other=0.0).to(tl.float32)
	# Gate half: rows [0 .. rank-1]
	a_gate_ptrs = (
	lora_a_base
	+ r_offs[:, None] * lora_a_stride_rank
	+ input_offs[None, :] * lora_a_stride_input
	)
	a_gate_vals = tl.load(
	a_gate_ptrs,
	mask=rank_mask[:, None] & input_mask[None, :],
	other=0.0,
	).to(tl.float32)
	gate_inter += tl.sum(a_gate_vals * h_vals[None, :], axis=1)
	# Up half: rows [rank .. 2*rank-1]
	a_up_ptrs = (
	lora_a_base
	+ (r_offs[:, None] + rank) * lora_a_stride_rank
	+ input_offs[None, :] * lora_a_stride_input
	)
	a_up_vals = tl.load(
	a_up_ptrs,
	mask=rank_mask[:, None] & input_mask[None, :],
	other=0.0,
	).to(tl.float32)
	up_inter += tl.sum(a_up_vals * h_vals[None, :], axis=1)
	# Stage 2: Compute outputs for both halves using B
	# Gate half: B rows [0 .. half_output_dim-1]
	b_gate_ptrs = (
	lora_b_base
	+ local_out_offs_gate[None, :] * lora_b_stride_output
	+ r_offs[:, None] * lora_b_stride_rank
	)
	b_gate_vals = tl.load(
	b_gate_ptrs,
	mask=rank_mask[:, None] & gate_out_mask[None, :],
	other=0.0,
	).to(tl.float32)
	gate_out_vals = tl.sum(b_gate_vals * gate_inter[:, None], axis=0)
	# Up half: B rows [half_output_dim .. 2*half_output_dim-1]
	b_up_ptrs = (
	lora_b_base
	+ (local_out_offs_up[None, :] + half_output_dim) * lora_b_stride_output
	+ r_offs[:, None] * lora_b_stride_rank
	)
	b_up_vals = tl.load(
	b_up_ptrs,
	mask=rank_mask[:, None] & up_out_mask[None, :],
	other=0.0,
	).to(tl.float32)
	up_out_vals = tl.sum(b_up_vals * up_inter[:, None], axis=0)
	# Elementwise select gate vs up outputs based on out_offs
	out_vals = tl.where(is_up_half, up_out_vals, gate_out_vals)
	# Zero out positions beyond output_dim (for safety/clarity)
	out_vals = tl.where(global_out_mask, out_vals, 0.0)

[Copilot]

per_expert_lora_forward hard-codes max_rank = 256. This will silently produce incorrect results if any loaded adapter rank exceeds 256 (or exceeds 128 for gate_up_proj where A is stacked to 2*rank). Consider deriving max_rank from the allocated weight tensor shape (e.g., lora_a_weights.shape[2]) and passing it as the tl.constexpr max_rank, or enforcing/validating an explicit upper bound against the configured max_lora_rank and failing fast when exceeded.

[Copilot]

This change leaves a sizable commented-out alternative implementation/return block in get_lora_A_shape, and also includes a redundant nested if self.is_moe_module(module_name) inside an already-true MoE branch. Please remove the commented-out block and simplify the control flow (compute c = get_stacked_multiply(module_name) once) to keep the allocator logic clear.

[Copilot]

These commented-out alternative model/LoRA paths add noise to a registered test and can become stale over time. If you want to keep alternative configs, prefer documenting them in the PR description or a developer doc; otherwise please remove the commented lines to keep the test focused and deterministic.

Suggested change

	# model_path = "Qwen/Qwen3-30B-A3B"
	lora_paths = ["sai-lakkshmii/Qwen1.5-MoE-A2.7B-squad-lora-latest"]
	# lora_paths = ["debaterhub/sentence-selection-orpo-lora"]
	lora_paths = ["sai-lakkshmii/Qwen1.5-MoE-A2.7B-squad-lora-latest"]

[Copilot]

Variable global_out_mask is not used.

Suggested change

global_out_mask = out_offs < output_dim