minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[LoRA] BF16 support + EP cuda-graph crash fix for experimental_sgl_trtllm MoE-LoRA (#28953)

Browse Source 
Co-authored-by: Yusheng Su <yushengsu.thu@gmail.com>
This commit is contained in:
Yanbin Jiang
2026-06-24 21:16:50 -07:00
committed by GitHub
co-authored by Yusheng Su
parent a9270250c3
commit 9fd6d0ec99
10 changed files with 946 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/trtllm_lora_temp/__init__.py
+2
View File
@@ -1,4 +1,5 @@
from sglang.jit_kernel.trtllm_lora_temp.core import (
trtllm_bf16_routed_moe_lora,
trtllm_fp4_block_scale_moe_lora_finalize,
trtllm_fp4_block_scale_routed_moe_lora,
trtllm_fp8_block_scale_moe,
@@ -8,6 +9,7 @@ from sglang.jit_kernel.trtllm_lora_temp.core import (
)
__all__ = [
"trtllm_bf16_routed_moe_lora",
"trtllm_fp4_block_scale_moe_lora_finalize",
"trtllm_fp4_block_scale_routed_moe_lora",
"trtllm_fp8_block_scale_moe_lora_finalize",
python/sglang/jit_kernel/trtllm_lora_temp/core.py
+71
View File
@@ -308,6 +308,77 @@ def trtllm_fp8_block_scale_moe_lora_finalize(
return output
def trtllm_bf16_routed_moe_lora(
topk_ids: torch.Tensor,
routing_bias: Optional[torch.Tensor],
hidden_states: torch.Tensor,
gemm1_weights: torch.Tensor,
gemm2_weights: torch.Tensor,
gate_up_lora_delta: torch.Tensor,
activation_lora_input: torch.Tensor,
num_experts: int,
top_k: int,
intermediate_size: int,
local_expert_offset: int,
local_num_experts: int,
routed_scaling_factor: Optional[float],
routing_method_type: int = 0,
do_finalize: bool = True,
enable_pdl: Optional[bool] = None,
output: Optional[torch.Tensor] = None,
activation_type: Optional[int] = None,
lora_ready_event: int = 0,
gemm2_done_event: int = 0,
) -> Union[List[torch.Tensor], torch.Tensor]:
"""BF16 MoE-LoRA: decomposed trtllm pipeline (permute -> raw gate_up GEMM ->
LoRA-aware activation -> down GEMM), bf16 end-to-end (no quantization).
Weights are the SAME prepared bf16 tensors the plain trtllm_bf16 path uses
(shuffled + BlockMajorK). With do_finalize=False returns
(gemm2_output, expert_weights, expanded_idx_to_permuted_idx) for the
python-side down-LoRA delta + trtllm_fp8_block_scale_moe_lora_finalize
(which is pure bf16 and shared across fp8/fp4/bf16)."""
from flashinfer.fused_moe.core import ActivationType
from flashinfer.utils import device_support_pdl
if activation_type is None:
activation_type = ActivationType.Swiglu.value
if enable_pdl is None:
enable_pdl = device_support_pdl(hidden_states.device)
if output is None:
output = torch.empty(
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
)
assert gate_up_lora_delta.is_contiguous()
assert activation_lora_input.is_contiguous()
result = get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_bf16_routed_moe_lora(
topk_ids,
routing_bias,
hidden_states,
gemm1_weights,
gemm2_weights,
num_experts,
top_k,
intermediate_size,
local_expert_offset,
local_num_experts,
routed_scaling_factor,
routing_method_type,
do_finalize,
enable_pdl,
activation_type,
output,
True,
gate_up_lora_delta,
activation_lora_input,
lora_ready_event,
gemm2_done_event,
)
return output if do_finalize else result
def trtllm_fp4_block_scale_routed_moe_lora(
topk_ids: torch.Tensor,
routing_bias: Optional[torch.Tensor],
python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu
+7
View File
@@ -720,6 +720,13 @@ __global__ void permuteKernel(KernelParams params) {
for (int k = 0; k < params.topK; k++) {
int const expandedIdx = tokenIdx * params.topK + k;
int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx];
// Skip EP-unrouted (token, k) slots: under expert parallelism the routing emits
// permutedIdx == -1 for slots whose expert is not on this rank. The write was
// previously unguarded, so at prefill scale the negative index OOBs the output
// buffer (illegal memory access). moe::dev::finalize already skips permutedIdx == -1.
if (permutedIdx < 0) {
continue;
}
params.outPtr[permutedIdx * params.hiddenDim + hiddenIdx] = data;
}
}
python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu
+413
View File
@@ -3450,6 +3450,418 @@ void sgl_trtllm_fp4_block_scale_moe_lora_finalize(
FLASHINFER_CHECK(err == cudaSuccess, cudaGetErrorString(err));
}
// ===========================================================================
// BF16 MoE LoRA (decomposed / unfused-activation) — bf16 sibling of the FP4
// trtllm-lora op (FP4BlockScaleLoraLauncher above). The standard BF16 path
// fuses SwiGLU into GEMM1 (runner.cu: fusedAct = !useDeepSeekFp8), which
// leaves no seam to inject the gate_up LoRA delta pre-activation. This is the
// FP4 decomposed pipeline with the two NvFP4 quant stages REMOVED — every
// stage is bf16 end-to-end:
//
// routing -> permute (gather bf16) -> gate_up GEMM (raw Gemm2::Runner,
// Bf16 x Bf16 -> Bf16, K=hidden, N=2*inter, gated-INTERLEAVED out) ->
// moe::dev::activation (de-interleaves on read, adds gate_up_lora_delta
// pre-SwiGLU, writes activation_lora_input for the down-proj LoRA, bf16
// out) -> down GEMM (Gemm2::Runner, K=inter, N=hidden) -> finalize (or
// return {gemm2_output, expert_weights, idx} for the python-side down-delta
// + sgl_trtllm_fp8_block_scale_moe_lora_finalize, which is pure bf16).
//
// Weights are the SAME prepared bf16 tensors the plain trtllm_bf16_moe path
// consumes (shuffled epilogue_tile_m=128 + BlockMajorK block_k=128, see
// sglang unquant.py process_weights_after_loading) — layout untouched.
// FP8/FP4 paths are untouched by this addition.
// ===========================================================================
class Bf16LoraLauncher {
public:
// Match the plain BF16 path's tile ladder (Bf16MoeLauncher::mSupportedTileNums).
static constexpr std::array<int32_t, 5> mBaseSupportedTileNums = {8, 16, 32, 64, 128};
static std::vector<int32_t> getSupportedTileNums() {
return std::vector<int32_t>(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end());
}
Bf16LoraLauncher(
TensorView const& expert_indices,
Optional<TensorView> const& routing_bias,
TensorView const& hidden_states,
TensorView const& gemm1_weights,
TensorView const& gemm2_weights,
TensorView const& gate_up_lora_delta,
TensorView const& activation_lora_input,
TensorView const& output,
int64_t lora_ready_event,
int64_t gemm2_done_event)
: expert_indices_(expert_indices),
routing_bias_(routing_bias),
hidden_states_(hidden_states),
gemm1_weights_(gemm1_weights),
gemm2_weights_(gemm2_weights),
gate_up_lora_delta_(gate_up_lora_delta),
activation_lora_input_(activation_lora_input),
output_(output),
lora_ready_event_(lora_ready_event),
gemm2_done_event_(gemm2_done_event) {}
// Returns {output} when do_finalize, else {gemm2_output, expert_weights,
// expanded_idx_to_permuted_idx} for a downstream finalize kernel.
Array<Tensor>
run(int64_t num_experts,
int64_t top_k,
int64_t intermediate_size,
int64_t local_expert_offset,
int64_t local_num_experts,
double routed_scaling_factor,
int64_t routing_method_type,
int64_t tile_tokens_dim,
bool norm_topk_prob,
bool do_finalize,
bool enable_pdl) {
namespace moe_ns = tensorrt_llm::kernels::trtllmgen_moe;
auto device = hidden_states_.device();
int dev_id = device.device_id;
cudaStream_t stream = get_stream(device);
int64_t const num_tokens = hidden_states_.size(0);
int64_t const hidden_size = hidden_states_.size(1);
int64_t const inter = intermediate_size;
int64_t const gate_up_n = 2 * inter; // gated SwiGLU
// ---- 1) routing (precomputed packed topk) — identical to the FP4 lora path ----
Tensor num_tokens_per_expert = alloc_tensor({num_experts}, dl_int32, device);
int32_t max_num_padded_tokens =
moe_ns::Routing::getMaxPermutedPaddedCount(num_tokens, top_k, num_experts, tile_tokens_dim);
Tensor total_num_padded_tokens = alloc_tensor({1}, dl_int32, device);
Tensor expanded_idx_to_permuted_idx = alloc_tensor({num_tokens * top_k}, dl_int32, device);
Tensor permuted_idx_to_token_idx = alloc_tensor({max_num_padded_tokens}, dl_int32, device);
int64_t const hist_size = std::max<int64_t>(num_experts * 2, 256 * 2);
Tensor expert_count_histogram = alloc_tensor({hist_size}, dl_int32, device);
int32_t max_num_ctas = moe_ns::Routing::getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts, tile_tokens_dim);
Tensor cta_idx_xy_to_batch_idx = alloc_tensor({max_num_ctas}, dl_int32, device);
Tensor cta_idx_xy_to_mn_limit = alloc_tensor({max_num_ctas}, dl_int32, device);
Tensor num_non_exiting_ctas = alloc_tensor({1}, dl_int32, device);
auto routing_bias_dtype = routing_bias_.has_value() ? routing_bias_.value().dtype() : dl_bfloat16;
btg::Dtype mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32;
auto ew_dtype = mRoutingBiasDtype == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16;
Tensor expert_weights_alloc = alloc_tensor({num_tokens, top_k}, ew_dtype, device);
void* expert_weights_ptr = expert_weights_alloc.data_ptr();
moe_ns::Routing::Runner routing_runner(tile_tokens_dim);
routing_runner.run(
/*routing_logits=*/nullptr,
routing_bias_.has_value() ? routing_bias_.value().data_ptr() : nullptr,
num_tokens,
num_experts,
top_k,
/*n_group=*/0,
/*topk_group=*/0,
local_expert_offset,
local_num_experts,
routed_scaling_factor,
static_cast<int*>(const_cast<void*>(expert_indices_.data_ptr())),
static_cast<int*>(expert_count_histogram.data_ptr()),
static_cast<int*>(total_num_padded_tokens.data_ptr()),
static_cast<int*>(expanded_idx_to_permuted_idx.data_ptr()),
/*permuted_idx_to_expanded_idx=*/nullptr,
static_cast<int*>(permuted_idx_to_token_idx.data_ptr()),
/*expertIds=*/nullptr,
expert_weights_ptr,
static_cast<int*>(num_tokens_per_expert.data_ptr()),
static_cast<int*>(cta_idx_xy_to_batch_idx.data_ptr()),
static_cast<int*>(cta_idx_xy_to_mn_limit.data_ptr()),
static_cast<int*>(num_non_exiting_ctas.data_ptr()),
btg::Dtype::Bfloat16,
mRoutingBiasDtype,
/*useRoutingScalesOnInput=*/false,
/*useDeepSeekFp8=*/false,
static_cast<RoutingMethodType>(routing_method_type),
stream,
btg::Dtype::Bfloat16,
norm_topk_prob,
/*routing_replay_out=*/nullptr);
int64_t const tile = tile_tokens_dim;
// ---- 2+3) permute (gather) bf16 hidden + gate_up GEMM ----
// The permuted gather buffer ([max_padded, hidden] bf16) lives only inside this block: it
// frees right after the gate_up GEMM is enqueued (stream-ordered caching allocator), before
// the equally-large gemm2_output is allocated — same peak-memory trick as the FP4 path.
Tensor gate_up_bf16 = alloc_tensor({max_num_padded_tokens, gate_up_n}, dl_bfloat16, device);
{
Tensor permuted_hidden_bf16 = alloc_tensor({max_num_padded_tokens, hidden_size}, dl_bfloat16, device);
// Padding rows intentionally left uninitialized — never reach a valid output (see the FP4
// path's rationale; finalize gathers only real tokens via the index map).
{
moe::dev::permute::Data permData;
permData.mDtypeElt = btg::Dtype::Bfloat16;
permData.mUsePdl = false;
permData.mUseDeepSeekFp8 = false;
permData.inPtr = hidden_states_.data_ptr();
permData.outPtr = permuted_hidden_bf16.data_ptr();
permData.inDqSfsPtr = nullptr;
permData.outDqSfsPtr = nullptr;
permData.expandedIdxToPermutedIdx = static_cast<int*>(expanded_idx_to_permuted_idx.data_ptr());
permData.hiddenDim = hidden_size;
permData.numTokens = num_tokens;
permData.topK = top_k;
permData.totalNumPaddedTokens = static_cast<int*>(total_num_padded_tokens.data_ptr());
moe::dev::permute::run(permData, stream);
}
// gate_up GEMM: raw Gemm2::Runner(Bf16,Bf16,Bf16, K=hidden, N=2*inter). The gated w13
// weight makes the output columns pairwise-INTERLEAVED (g0,u0,g1,u1,...) — the activation
// below de-interleaves on read (interleavedGateUpInput), same as the FP4 path. Weights are
// the plain bf16 path's prepared tensors: shuffled + BlockMajorK (NOT MajorK like FP4's
// packed-fp4 weights).
moe_ns::Gemm2::Runner gemm_gate_up(
btg::Dtype::Bfloat16,
btg::Dtype::Bfloat16,
btg::Dtype::Bfloat16,
/*useDeepSeekFp8=*/false,
(int)tile,
/*useShuffledMatrix=*/true,
batchedGemm::gemm::MatrixLayout::BlockMajorK,
/*usePerTokenScaling=*/false,
/*usePerChannelScaling=*/false);
int64_t cfg =
gemm_gate_up.getDefaultValidConfigIndex(top_k, hidden_size, gate_up_n, local_num_experts, num_tokens);
size_t ws =
gemm_gate_up.getWorkspaceSizeInBytes(top_k, hidden_size, gate_up_n, local_num_experts, num_tokens, cfg);
Tensor gemm_ws = alloc_tensor({(int64_t)ws}, dl_int8, device);
// Gemm2::Runner semantics: hiddenSize is the OUTPUT N dim, intermediateSize is the K dim.
gemm_gate_up.run(
permuted_hidden_bf16.data_ptr(),
/*hiddenStateScale=*/nullptr,
gemm1_weights_.data_ptr(),
/*weightScale=*/nullptr,
/*perTokenScales=*/nullptr,
/*perChannelScales=*/nullptr,
/*scaleC=*/nullptr,
/*ptrBias=*/nullptr,
gate_up_bf16.data_ptr(),
/*outputScale=*/nullptr,
top_k,
/*hiddenSize(N)=*/gate_up_n,
/*intermediateSize(K)=*/hidden_size,
local_num_experts,
num_tokens,
static_cast<int*>(num_non_exiting_ctas.data_ptr()),
static_cast<int*>(total_num_padded_tokens.data_ptr()),
static_cast<int*>(cta_idx_xy_to_batch_idx.data_ptr()),
static_cast<int*>(cta_idx_xy_to_mn_limit.data_ptr()),
gemm_ws.data_ptr(),
dev_id,
stream,
(int)cfg,
enable_pdl);
} // permuted_hidden_bf16 frees here -> its block is reused by gemm2_output below.
// GEMM1-LoRA overlap: wait the side-stream LoRA event that produced the gate_up_lora_delta
// consumed by the activation below, so the gate_up GEMM above overlaps the side-stream
// LoRA shrink/expand. No-op (0) on the single-stream path.
if (lora_ready_event_ != 0) {
cudaStreamWaitEvent(stream, reinterpret_cast<cudaEvent_t>(lora_ready_event_), 0);
}
// ---- 4) activation (SwiGLU + gate_up LoRA delta, bf16 -> bf16, NO quant) ----
// Same generic activation kernel the FP4 lora path uses in its non-fused mode: bf16 in,
// de-interleave on read, add gate_up_lora_delta pre-SwiGLU, capture activation_lora_input
// (by EXPANDED idx, for the python-side down-proj LoRA shrink), bf16 out (by permuted idx).
auto envFlag = [](char const* name) {
char const* e = std::getenv(name);
return e != nullptr && (e[0] == '1' || e[0] == 't' || e[0] == 'T' || e[0] == 'y' || e[0] == 'Y');
};
static int const actOptMode = envFlag("SGLANG_OPT_FUSED_MOE_ACTIVATION_VEC") ? 1 : 0;
Tensor activated_bf16 = alloc_tensor({max_num_padded_tokens, inter}, dl_bfloat16, device);
{
moe::dev::activation::Data actData;
actData.mDtypeElt = btg::Dtype::Bfloat16;
actData.mUsePdl = false;
actData.mUseDeepSeekFp8 = false;
actData.inPtr = gate_up_bf16.data_ptr();
actData.interleavedGateUpInput = true;
actData.outPtr = activated_bf16.data_ptr();
actData.inDqSfsPtr = nullptr;
actData.outDqSfsPtr = nullptr;
actData.gateUpLoraDeltaPtr = static_cast<cutlass::bfloat16_t const*>(gate_up_lora_delta_.data_ptr());
actData.activationLoraInputOutPtr = static_cast<cutlass::bfloat16_t*>(activation_lora_input_.data_ptr());
actData.innerDim = gate_up_n;
actData.numTokens = num_tokens;
actData.topK = top_k;
actData.expandedIdxToPermutedIdx = static_cast<int*>(expanded_idx_to_permuted_idx.data_ptr());
actData.totalNumPaddedTokens = static_cast<int*>(total_num_padded_tokens.data_ptr());
actData.actOptMode = actOptMode;
moe::dev::activation::run(actData, stream);
}
// ---- 5) down GEMM: Gemm2::Runner(Bf16,Bf16,Bf16, K=inter, N=hidden) ----
Tensor gemm2_output = alloc_tensor({max_num_padded_tokens, hidden_size}, dl_bfloat16, device);
{
moe_ns::Gemm2::Runner gemm_down(
btg::Dtype::Bfloat16,
btg::Dtype::Bfloat16,
btg::Dtype::Bfloat16,
/*useDeepSeekFp8=*/false,
(int)tile,
/*useShuffledMatrix=*/true,
batchedGemm::gemm::MatrixLayout::BlockMajorK,
/*usePerTokenScaling=*/false,
/*usePerChannelScaling=*/false);
int64_t cfg = gemm_down.getDefaultValidConfigIndex(top_k, hidden_size, inter, local_num_experts, num_tokens);
size_t ws = gemm_down.getWorkspaceSizeInBytes(top_k, hidden_size, inter, local_num_experts, num_tokens, cfg);
Tensor gemm_ws = alloc_tensor({(int64_t)ws}, dl_int8, device);
gemm_down.run(
activated_bf16.data_ptr(),
/*hiddenStateScale=*/nullptr,
gemm2_weights_.data_ptr(),
/*weightScale=*/nullptr,
/*perTokenScales=*/nullptr,
/*perChannelScales=*/nullptr,
/*scaleC=*/nullptr,
/*ptrBias=*/nullptr,
gemm2_output.data_ptr(),
/*outputScale=*/nullptr,
top_k,
/*hiddenSize(N)=*/hidden_size,
/*intermediateSize(K)=*/inter,
local_num_experts,
num_tokens,
static_cast<int*>(num_non_exiting_ctas.data_ptr()),
static_cast<int*>(total_num_padded_tokens.data_ptr()),
static_cast<int*>(cta_idx_xy_to_batch_idx.data_ptr()),
static_cast<int*>(cta_idx_xy_to_mn_limit.data_ptr()),
gemm_ws.data_ptr(),
dev_id,
stream,
(int)cfg,
enable_pdl);
}
// Down-LoRA/finalize overlap: signal "base down GEMM done" so the LoRA side stream can
// start the down-proj LoRA shrink/expand concurrent with the finalize below. 0 = no-op.
if (gemm2_done_event_ != 0) {
cudaEventRecord(reinterpret_cast<cudaEvent_t>(gemm2_done_event_), stream);
}
if (!do_finalize) {
return {gemm2_output, expert_weights_alloc, expanded_idx_to_permuted_idx};
}
// ---- 6) finalize (combine by expert weight) -> output [num_tokens, hidden] ----
{
moe::dev::finalize::Data finData;
finData.mDtypeElt = btg::Dtype::Bfloat16;
finData.mDtypeExpW = mRoutingBiasDtype;
finData.mUsePdl = false;
finData.mUseDeepSeekFp8 = false;
finData.inPtr = gemm2_output.data_ptr();
finData.outPtr = output_.data_ptr();
finData.inDqSfsPtr = nullptr;
finData.outDqSfsPtr = nullptr;
finData.expertWeightsPtr = expert_weights_ptr;
finData.expandedIdxToPermutedIdx = static_cast<int*>(expanded_idx_to_permuted_idx.data_ptr());
finData.numTokens = num_tokens;
finData.numExperts = num_experts;
finData.topK = top_k;
finData.hiddenDim = hidden_size;
finData.hiddenDimPadded = hidden_size;
finData.totalNumPaddedTokens = static_cast<int*>(total_num_padded_tokens.data_ptr());
moe::dev::finalize::run(finData, stream);
}
sync_check_cuda_error(stream);
return Array<Tensor>();
}
private:
TensorView expert_indices_;
Optional<TensorView> routing_bias_;
TensorView hidden_states_;
TensorView gemm1_weights_;
TensorView gemm2_weights_;
TensorView gate_up_lora_delta_;
TensorView activation_lora_input_;
TensorView output_;
int64_t lora_ready_event_ = 0;
int64_t gemm2_done_event_ = 0;
};
Array<Tensor> sgl_trtllm_bf16_routed_moe_lora(
TensorView expert_indices,
Optional<TensorView> routing_bias,
TensorView hidden_states,
TensorView gemm1_weights,
TensorView gemm2_weights,
int64_t num_experts,
int64_t top_k,
int64_t intermediate_size,
int64_t local_expert_offset,
int64_t local_num_experts,
Optional<double> routed_scaling_factor,
int64_t routing_method_type,
bool do_finalize,
bool enable_pdl,
int64_t act_type,
TensorView output,
bool norm_topk_prob,
TensorView gate_up_lora_delta,
TensorView activation_lora_input,
int64_t lora_ready_event,
int64_t gemm2_done_event) {
auto activation_type = validateAndCastActivationType(act_type);
TVM_FFI_ICHECK(isGatedActivation(activation_type))
<< "sgl_trtllm_bf16_routed_moe_lora currently supports gated (SwiGLU) activation only.";
// Precomputed routing is required (packed topk_ids).
TVM_FFI_ICHECK(expert_indices.ndim() == 2 && expert_indices.size(0) == hidden_states.size(0))
<< "bf16 LoRA requires precomputed packed expert_indices [num_tokens, top_k].";
TVM_FFI_ICHECK_EQ(expert_indices.dtype(), dl_int32) << "expert_indices must be int32.";
TVM_FFI_ICHECK_EQ(hidden_states.dtype(), dl_bfloat16) << "bf16 LoRA: hidden_states must be bf16.";
TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_bfloat16) << "bf16 LoRA: gemm1_weights must be bf16.";
TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_bfloat16) << "bf16 LoRA: gemm2_weights must be bf16.";
TVM_FFI_ICHECK_EQ(intermediate_size % 128, 0)
<< "bf16 LoRA: intermediate_size must be a multiple of 128 (BlockMajorK weights).";
TVM_FFI_ICHECK_EQ(gate_up_lora_delta.dtype(), dl_bfloat16) << "gate_up_lora_delta must be bf16.";
TVM_FFI_ICHECK_EQ(gate_up_lora_delta.ndim(), 3)
<< "gate_up_lora_delta must be [num_tokens, top_k, 2*intermediate_size].";
TVM_FFI_ICHECK_EQ(gate_up_lora_delta.size(2), 2 * intermediate_size);
TVM_FFI_ICHECK_EQ(activation_lora_input.dtype(), dl_bfloat16) << "activation_lora_input must be bf16.";
TVM_FFI_ICHECK_EQ(activation_lora_input.ndim(), 3)
<< "activation_lora_input must be [num_tokens, top_k, intermediate_size].";
TVM_FFI_ICHECK_EQ(activation_lora_input.size(2), intermediate_size);
TVM_FFI_ICHECK(gate_up_lora_delta.IsContiguous() && activation_lora_input.IsContiguous())
<< "lora bridge buffers must be contiguous.";
int64_t const num_tokens = hidden_states.size(0);
int64_t const tile_tokens_dim =
selectDefaultTileN(Bf16LoraLauncher::getSupportedTileNums(), num_tokens, top_k, local_num_experts);
Bf16LoraLauncher launcher(
expert_indices,
routing_bias,
hidden_states,
gemm1_weights,
gemm2_weights,
gate_up_lora_delta,
activation_lora_input,
output,
lora_ready_event,
gemm2_done_event);
return launcher.run(
num_experts,
top_k,
intermediate_size,
local_expert_offset,
local_num_experts,
routed_scaling_factor.value_or(1.0),
routing_method_type,
tile_tokens_dim,
norm_topk_prob,
do_finalize,
enable_pdl);
}
Array<Tensor> trtllm_mxint4_block_scale_moe(
TensorView routing_logits,
Optional<TensorView> routing_bias,
@@ -3951,6 +4363,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(sgl_trtllm_fp8_block_scale_moe_lora, sgl_trtllm_fp
TVM_FFI_DLL_EXPORT_TYPED_FUNC(
sgl_trtllm_fp8_block_scale_moe_lora_finalize, sgl_trtllm_fp8_block_scale_moe_lora_finalize);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(trtllm_fp4_block_scale_moe, trtllm_fp4_block_scale_moe);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(sgl_trtllm_bf16_routed_moe_lora, sgl_trtllm_bf16_routed_moe_lora);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(sgl_trtllm_fp4_block_scale_moe_lora, sgl_trtllm_fp4_block_scale_moe_lora);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(
sgl_trtllm_fp4_block_scale_moe_lora_finalize, sgl_trtllm_fp4_block_scale_moe_lora_finalize);
python/sglang/srt/lora/trtllm_lora_temp/__init__.py
+20 -5
View File
@@ -92,6 +92,7 @@ _ORIGINAL_COLUMN_FORWARD: Optional[Callable] = None
_ORIGINAL_REPLICATED_FORWARD: Optional[Callable] = None
_ORIGINAL_MOE_LORA_FUNC: Optional[Callable] = None
_ORIGINAL_FP4_MOE_LORA_FUNC: Optional[Callable] = None
_ORIGINAL_BF16_MOE_LORA_FUNC: Optional[Callable] = None
_INSTALLED: bool = False
@@ -123,6 +124,10 @@ def get_original_fp4_moe_lora_func() -> Callable:
return _ORIGINAL_FP4_MOE_LORA_FUNC
def get_original_bf16_moe_lora_func() -> Callable:
return _ORIGINAL_BF16_MOE_LORA_FUNC
def install_two_stream_overrides() -> None:
"""Install the side-stream overlapped overrides if ``SGLANG_LORA_TWO_STREAM=1``.
@@ -132,14 +137,15 @@ def install_two_stream_overrides() -> None:
2. ``RowParallelLinearWithLoRA.forward`` (O8 — o_proj LoRA shrink overlap)
3. ``MergedColumnParallelLinearWithLoRA.forward`` (O9 — merged-column LoRA
shrink overlap: dense gate_up + mamba in_proj_qkvz)
4. ``flashinfer_trtllm.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``
(O1 — MoE gate_up LoRA overlap)
4. ``lora_dispatch.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``
(O1 — MoE gate_up LoRA overlap), plus its fp4 (O1-fp4) and bf16
(O1-bf16) siblings
The saved originals are exposed via :func:`get_original_qkv_forward`,
:func:`get_original_row_forward`, :func:`get_original_moe_lora_func` so the
new versions can fall back when their per-batch gate says single-stream.
"""
global _INSTALLED, _ORIGINAL_QKV_FORWARD, _ORIGINAL_ROW_FORWARD, _ORIGINAL_MERGED_FORWARD, _ORIGINAL_COLUMN_FORWARD, _ORIGINAL_REPLICATED_FORWARD, _ORIGINAL_MOE_LORA_FUNC, _ORIGINAL_FP4_MOE_LORA_FUNC
global _INSTALLED, _ORIGINAL_QKV_FORWARD, _ORIGINAL_ROW_FORWARD, _ORIGINAL_MERGED_FORWARD, _ORIGINAL_COLUMN_FORWARD, _ORIGINAL_REPLICATED_FORWARD, _ORIGINAL_MOE_LORA_FUNC, _ORIGINAL_FP4_MOE_LORA_FUNC, _ORIGINAL_BF16_MOE_LORA_FUNC
if _INSTALLED:
return
@@ -178,22 +184,30 @@ def install_two_stream_overrides() -> None:
import sglang.srt.lora.trtllm_lora_temp.lora_dispatch as ft
from sglang.srt.lora.trtllm_lora_temp.moe_overlap import (
fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream,
fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream,
fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream,
)
# O1 (FP8 Qwen) + O1-fp4 (NVFP4 Kimi): MoE gate_up LoRA overlap. Each patched
# fn falls back to its saved single-stream original for non-decode batches.
# O1 (FP8 Qwen) + O1-fp4 (NVFP4 Kimi) + O1-bf16 (unquantized Qwen): MoE gate_up
# LoRA overlap. Each patched fn falls back to its saved single-stream original
# for non-decode batches.
_ORIGINAL_MOE_LORA_FUNC = ft.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora
_ORIGINAL_FP4_MOE_LORA_FUNC = (
ft.fused_experts_none_to_experimental_sgl_trtllm_fp4_lora
)
_ORIGINAL_BF16_MOE_LORA_FUNC = (
ft.fused_experts_none_to_experimental_sgl_trtllm_bf16_lora
)
ft.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora = (
fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream
)
ft.fused_experts_none_to_experimental_sgl_trtllm_fp4_lora = (
fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream
)
ft.fused_experts_none_to_experimental_sgl_trtllm_bf16_lora = (
fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream
)
_INSTALLED = True
@@ -209,5 +223,6 @@ __all__ = [
"get_original_replicated_forward",
"get_original_moe_lora_func",
"get_original_fp4_moe_lora_func",
"get_original_bf16_moe_lora_func",
"install_two_stream_overrides",
]
python/sglang/srt/lora/trtllm_lora_temp/environ.py
+8
View File
@@ -72,6 +72,14 @@ class _LoraEnvs:
"SGLANG_OPT_LORA_FUSED_TOPK_PACK", True
)
SGLANG_OPT_LORA_QKV_B_STORE = _GatedBool("SGLANG_OPT_LORA_QKV_B_STORE", True)
# F1-①: prefill routing reuse — unify the A (shrink) stage's routing BLOCK_SIZE_M with
# the B stage's at prefill (>=512 tokens) so the per-layer routing_cache key matches
# across stages and the Triton align/sort runs once per layer-forward instead of once
# per stage (4x at prefill). Dtype-agnostic (the chain is shared by fp8/nvfp4/bf16).
# Decode (<512) keeps the opt1 fused merged-align path and its tuned shrink block.
SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE = _GatedBool(
"SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE", True
)
# ---- correctness fixes: on by default when experimental ----
# gate_up gated-split fix (up_A shrink for the up half); set =0 only to A/B bisect.
python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py
+152
View File
@@ -29,6 +29,7 @@ from sglang.srt.utils.common import next_power_of_2
if TYPE_CHECKING:
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
FlashInferTrtllmBf16MoeQuantInfo,
FlashInferTrtllmFp4MoeQuantInfo,
FlashInferTrtllmFp8MoeQuantInfo,
)
@@ -299,6 +300,157 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
return StandardCombineInput(hidden_states=output)
def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora(
dispatch_output: StandardDispatchOutput,
quant_info: FlashInferTrtllmBf16MoeQuantInfo,
runner_config: MoeRunnerConfig,
lora_info,
) -> StandardCombineInput:
"""BF16 sibling of ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``.
Decomposed (unfused-activation) MoE-LoRA, bf16 end-to-end (no quantization):
routing -> gather -> gate_up grouped GEMM (raw 2*inter, bf16) -> activation that
adds ``gate_up_lora_delta`` pre-SwiGLU and captures ``activation_lora_input`` ->
down grouped GEMM -> finalize, then the virtual-experts down-LoRA is merged into
the output. Single-stream version (no two-stream overlap yet — phase 2).
"""
from sglang.jit_kernel.trtllm_lora_temp import trtllm_bf16_routed_moe_lora
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
fused_experts_none_to_flashinfer_trtllm_bf16,
get_activation_type,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
merged_experts_fused_moe_lora_add,
)
from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
assert (
runner_config.activation == "silu" and runner_config.is_gated
), "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only."
hidden_states = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
assert TopKOutputChecker.format_is_standard(topk_output)
assert runner_config.top_k is not None
# No active LoRA in a non-capture decode -> plain (fast) bf16 path.
if not get_is_capture_mode() and not lora_info.has_active_lora:
return fused_experts_none_to_flashinfer_trtllm_bf16(
dispatch_output, quant_info, runner_config, use_routed_topk=True
)
topk_ids = topk_output.topk_ids
topk_weights = topk_output.topk_weights
use_virtual_lora_store = bool(
lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
)
assert use_virtual_lora_store, "BF16 trtllm LoRA requires virtual-experts."
token_lora_mapping = lora_info.token_lora_mapping
fused_lora_routing_cache: dict = {}
inter = runner_config.intermediate_size_per_partition
# Gated gate_up LoRA delta (same shape/semantics as the fp8/fp4 paths). EP args scope
# the delta to this rank's experts, matching the EP-aware trtllm MoE base.
gate_up_delta = hidden_states.new_empty(
(hidden_states.shape[0], runner_config.top_k, 2 * inter)
)
merged_experts_fused_moe_lora_add(
output=gate_up_delta,
hidden_states=hidden_states,
lora_a=lora_info.gate_up_lora_a_weights,
lora_b=lora_info.gate_up_lora_b_weights,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=False,
experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
experts_shared_outer_loras_b=False,
routing_cache=fused_lora_routing_cache,
fuse_add_to_output=False,
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=runner_config.num_local_experts,
)
activation_lora_input = torch.empty(
(hidden_states.shape[0], runner_config.top_k, inter),
dtype=hidden_states.dtype,
device=hidden_states.device,
)
packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
if packed_topk_ids is None:
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
routing_method_type = runner_config.routing_method_type
if routing_method_type is None:
routing_method_type = RoutingMethodType.Default
elif routing_method_type == RoutingMethodType.DeepSeekV3:
routing_method_type = RoutingMethodType.TopK
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
direct_down_output = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
dtype=hidden_states.dtype,
device=hidden_states.device,
)
output = trtllm_bf16_routed_moe_lora(
topk_ids=packed_topk_ids,
routing_bias=None,
hidden_states=hidden_states,
gemm1_weights=quant_info.gemm1_weights,
gemm2_weights=quant_info.gemm2_weights,
gate_up_lora_delta=gate_up_delta,
activation_lora_input=activation_lora_input,
num_experts=quant_info.global_num_experts,
top_k=runner_config.top_k,
intermediate_size=inter,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=runner_config.num_local_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
routing_method_type=routing_method_type,
do_finalize=True,
output=direct_down_output,
activation_type=get_activation_type(
runner_config.activation, is_gated=runner_config.is_gated
),
)
merged_experts_fused_moe_lora_add(
output=output,
hidden_states=activation_lora_input.view(-1, inter),
lora_a=lora_info.down_lora_a_weights,
lora_b=lora_info.down_lora_b_weights,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=True,
experts_shared_outer_loras_a=False,
experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
routing_cache=fused_lora_routing_cache,
fuse_add_to_output=False,
fuse_sum_all_reduce=True,
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=runner_config.num_local_experts,
)
return StandardCombineInput(hidden_states=output)
def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
dispatch_output: StandardDispatchOutput,
quant_info: FlashInferTrtllmFp4MoeQuantInfo,
python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py
+37 -1
View File
@@ -81,6 +81,39 @@ def init_experimental_sgl_trtllm_lora(layer, base_layer) -> None:
if weight_block_size is None:
weight_block_size = getattr(quant_method, "weight_block_size", None)
use_mxfp8 = bool(getattr(quant_config, "use_mxfp8", False))
# ---- BF16 (unquantized) path ----
# No quant_config / block scales => the checkpoint is bf16. The bf16 LoRA dispatch
# runs the decomposed trtllm pipeline (sgl_trtllm_bf16_routed_moe_lora): permute ->
# raw gate_up GEMM -> LoRA-aware activation -> down GEMM, all bf16 — using the SAME
# prepared w13/w2 tensors (shuffled + BlockMajorK) the plain trtllm_bf16 path consumes.
# intermediate_size / local_num_experts / routing_method_type come from
# base_layer.moe_runner_config at dispatch time (the bf16 quant-info is minimal).
if quant_config is None and not getattr(quant_method, "block_quant", False):
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
FlashInferTrtllmBf16MoeQuantInfo,
)
layer._lora_runner = None
layer._quant_info = FlashInferTrtllmBf16MoeQuantInfo(
gemm1_weights=base_layer.w13_weight.data,
gemm2_weights=base_layer.w2_weight.data,
global_num_experts=int(base_layer.num_experts),
local_expert_offset=int(base_layer.moe_ep_rank)
* int(base_layer.num_local_experts),
)
# Expose w13_weight/w2_weight on the bf16 quant-info so the backend-agnostic
# cuda-graph MoE buffer init (BaseLoRABackend.init_cuda_graph_moe_buffers) reads
# expert dims uniformly with the FP8/FP4 quant-infos (which name them w13/w2).
# The bf16 BlockMajorK weights are 4-D [E, N, K/128, 128]; collapsing the inner
# dims to a 3-D [E, N, K] view (free for contiguous weights) makes the upstream
# `E, N, _ = w13_weight.shape` dim-extraction work without touching base_backend.
_g1 = layer._quant_info.gemm1_weights
_g2 = layer._quant_info.gemm2_weights
layer._quant_info.w13_weight = _g1.reshape(_g1.shape[0], _g1.shape[1], -1)
layer._quant_info.w2_weight = _g2.reshape(_g2.shape[0], _g2.shape[1], -1)
return
assert getattr(
quant_method, "block_quant", False
), "experimental_sgl_trtllm LoRA currently requires FP8 block quant."
@@ -133,14 +166,17 @@ def dispatch_experimental_sgl_trtllm_lora(
"""
import sglang.srt.lora.trtllm_lora_temp.lora_dispatch as ft
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
FlashInferTrtllmBf16MoeQuantInfo,
FlashInferTrtllmFp4MoeQuantInfo,
)
# Resolve the fused-experts fn on the module at CALL TIME so the install-time
# two-stream monkey-patch (sglang.srt.lora.trtllm_lora_temp) takes effect. Route by
# quant dtype: NVFP4 -> fp4 LoRA op, else the FP8 path.
# quant dtype: NVFP4 -> fp4 LoRA op, BF16 (unquantized) -> bf16 LoRA op, else FP8.
if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo):
fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_fp4_lora
elif isinstance(quant_info, FlashInferTrtllmBf16MoeQuantInfo):
fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_bf16_lora
else:
fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora
python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py
+210
View File
@@ -14,6 +14,7 @@ import torch
from sglang.srt.lora.trtllm_lora_temp import (
get_lora_side_stream,
get_original_bf16_moe_lora_func,
get_original_fp4_moe_lora_func,
get_original_moe_lora_func,
is_two_stream_active,
@@ -574,3 +575,212 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
else:
_run_down_lora(output)
return StandardCombineInput(hidden_states=output)
def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream(
dispatch_output,
quant_info,
runner_config,
lora_info,
):
"""Two-stream BF16 sibling of the FP8/FP4 two-stream MoE LoRA dispatches.
O1-bf16 fork: the gate_up LoRA shrink/expand runs on the side stream
concurrent with the bf16 op's routing + permute + gate_up GEMM; the op
waits on ``lora_ready_event`` right before its activation kernel (the only
consumer of ``gate_up_delta``). Fires only for virtual-experts LoRA +
decode-shaped batches; everything else delegates to the saved-original
single-stream bf16 dispatch (byte-identical). Down-LoRA stays serial on
the main stream — the down/finalize overlap was bench-verified
net-neutral-to-negative on the FP8/FP4 paths and corrupted the base
decode path under cuda-graph replay (see the comment in the FP4 variant).
"""
hidden_states = dispatch_output.hidden_states
use_virtual_lora_store = bool(
lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
)
if not (use_virtual_lora_store and is_two_stream_active(hidden_states)):
return get_original_bf16_moe_lora_func()(
dispatch_output, quant_info, runner_config, lora_info
)
# ---- two-stream fast path ----
from sglang.jit_kernel.trtllm_lora_temp import trtllm_bf16_routed_moe_lora
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
get_activation_type,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
merged_experts_fused_moe_lora_add,
)
assert (
runner_config.activation == "silu" and runner_config.is_gated
), "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only."
topk_output = dispatch_output.topk_output
assert TopKOutputChecker.format_is_standard(topk_output)
assert runner_config.top_k is not None
topk_ids = topk_output.topk_ids
topk_weights = topk_output.topk_weights
token_lora_mapping = lora_info.token_lora_mapping
fused_lora_routing_cache: dict = {}
inter = runner_config.intermediate_size_per_partition
side_stream = get_lora_side_stream()
gate_up_delta = hidden_states.new_empty(
(hidden_states.shape[0], runner_config.top_k, 2 * inter)
)
def _run_gate_up_lora():
merged_experts_fused_moe_lora_add(
output=gate_up_delta,
hidden_states=hidden_states,
lora_a=lora_info.gate_up_lora_a_weights,
lora_b=lora_info.gate_up_lora_b_weights,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=False,
experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
experts_shared_outer_loras_b=False,
routing_cache=fused_lora_routing_cache,
fuse_add_to_output=False,
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=runner_config.num_local_experts,
intermediate_buffer=gate_up_lora_intermediate,
)
# O1-bf16 fork: gate_up shrink/expand on the side stream, concurrent with the
# bf16 op's routing + permute + gate_up GEMM below. The op waits on lora_event
# right before its activation kernel (the only consumer of gate_up_delta).
lora_event = torch.cuda.Event()
# Hoist every side-chain allocation onto the MAIN stream (cuda-graph allocator
# safety -- see the "routing" stage in virtual_experts.py): pre-warm the routing
# cache and pre-allocate the shrink intermediate here, so the side-stream block
# below launches kernels only. Without this, the routing tensors + shrink
# intermediate get allocated inside the side-stream context during cuda-graph
# capture, where cross-stream tracking is off -> pool blocks reused with no graph
# edge -> '!!!!' decode corruption at max-loras>=2 (matches the fp8/fp4 fix).
merged_experts_fused_moe_lora_add(
output=gate_up_delta,
hidden_states=hidden_states,
lora_a=lora_info.gate_up_lora_a_weights,
lora_b=lora_info.gate_up_lora_b_weights,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=False,
experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
experts_shared_outer_loras_b=False,
routing_cache=fused_lora_routing_cache,
stage="routing",
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=runner_config.num_local_experts,
)
gate_up_lora_intermediate = hidden_states.new_empty(
(
hidden_states.shape[0],
topk_ids.shape[1],
lora_info.gate_up_lora_a_weights.shape[2],
)
)
side_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(side_stream):
_run_gate_up_lora()
lora_event.record()
activation_lora_input = torch.empty(
(hidden_states.shape[0], runner_config.top_k, inter),
dtype=hidden_states.dtype,
device=hidden_states.device,
)
packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
if packed_topk_ids is None:
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
routing_method_type = runner_config.routing_method_type
if routing_method_type is None:
routing_method_type = RoutingMethodType.Default
elif routing_method_type == RoutingMethodType.DeepSeekV3:
routing_method_type = RoutingMethodType.TopK
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
direct_down_output = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
dtype=hidden_states.dtype,
device=hidden_states.device,
)
# Keep the event alive through cuda-graph capture so the captured wait inside
# the bf16 op isn't torn down before instantiation (eager relies on deferred destroy).
if torch.cuda.is_current_stream_capturing():
_LORA_OVERLAP_EVENTS.append(lora_event)
lora_ready_handle = lora_event.cuda_event
output = trtllm_bf16_routed_moe_lora(
topk_ids=packed_topk_ids,
routing_bias=None,
hidden_states=hidden_states,
gemm1_weights=quant_info.gemm1_weights,
gemm2_weights=quant_info.gemm2_weights,
gate_up_lora_delta=gate_up_delta,
activation_lora_input=activation_lora_input,
num_experts=quant_info.global_num_experts,
top_k=runner_config.top_k,
intermediate_size=inter,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=runner_config.num_local_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
routing_method_type=routing_method_type,
do_finalize=True,
output=direct_down_output,
activation_type=get_activation_type(
runner_config.activation, is_gated=runner_config.is_gated
),
lora_ready_event=lora_ready_handle,
# Down-LoRA/finalize overlap intentionally NOT wired (gemm2_done_event=0):
# bench-verified net-neutral-to-negative on FP8/FP4 and corrupts the base
# decode path under cuda-graph replay. Serial down-LoRA below.
gemm2_done_event=0,
)
merged_experts_fused_moe_lora_add(
output=output,
hidden_states=activation_lora_input.view(-1, inter),
lora_a=lora_info.down_lora_a_weights,
lora_b=lora_info.down_lora_b_weights,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=True,
experts_shared_outer_loras_a=False,
experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
routing_cache=fused_lora_routing_cache,
fuse_add_to_output=False,
fuse_sum_all_reduce=True,
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=runner_config.num_local_experts,
)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py
+26 -5
View File
@@ -761,15 +761,22 @@ def _merged_experts_fused_moe_lora_add_impl(
# Fused LoRA-local align: one kernel does inline virtual id + EP skip +
# compact (local experts) + single-block scatter, replacing the 3-kernel
# (_fused_virtual_topk_ids + moe_align + count_and_sort) pipeline. Only the
# supported per-expert single-adapter EP path; everything else falls back.
# (_fused_virtual_topk_ids + moe_align + count_and_sort) pipeline. Two
# single-adapter (max_loras==1) regimes are fused; everything else falls back:
# - per-expert EP path (ep_local): compact local-expert histogram.
# - shared-outer path (shared_outer): lora-id routing (compute_virtual_id
# uses base=0; the kernel + launcher already size num_experts_for_weight=1
# and have no bucket-count blocker, so it just needs compact=False — compact
# + shared_outer would mis-map the id as base-offset). This is the opt1
# align/sort fusion: shared-outer used to fall through to the unfused
# _fused_virtual_topk_ids + moe_align_block_size_small_batch pair (~10.2us/
# layer at decode bs16); now it takes the single fused launch.
# Decode-only: the fused kernel's single-block scatter targets the small
# decode batch; prefill (>= 512 tokens) keeps the multi-block old path.
if (
lora_envs.SGLANG_OPT_LORA_FUSED_MERGED_ALIGN.get()
and max_loras == 1
and not shared_outer
and ep_local
and (shared_outer or ep_local)
and topk_ids.shape[0] < 512
):
from sglang.jit_kernel.trtllm_lora_temp.moe_lora_merged_align import (
@@ -792,7 +799,9 @@ def _merged_experts_fused_moe_lora_add_impl(
local_expert_offset,
local_num_experts,
do_skip=True,
compact=True,
# compact local-expert histogram is only valid for the per-expert EP
# path; shared_outer routes by lora id (base=0) so it must stay global.
compact=not shared_outer,
)
result = (
sorted_token_ids,
@@ -912,6 +921,18 @@ def _merged_experts_fused_moe_lora_add_impl(
"num_warps": 4,
"num_stages": 4,
}
# F1-① prefill routing reuse: the A stage routes with BLOCK_SIZE_M 32 at prefill
# but the B stage with the tuned fused-moe config (typically 64), so the
# (num_experts, shared_outer, block_size) routing_cache key never matches across
# stages and the align/sort pipeline reruns per stage (4x/layer at prefill).
# Matching the A stage's routing block to the B stage's collapses them to one
# align/sort per layer-forward. Decode (<512 tokens) keeps the opt1 fused
# merged-align path and its tuned shrink block untouched.
if (
lora_envs.SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE.get()
and token_lora_mapping.shape[0] >= 512
):
a_stage_config["BLOCK_SIZE_M"] = b_stage_config["BLOCK_SIZE_M"]
(
sorted_token_ids,
expert_ids,