diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py b/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py index 5269f90d3..5b2b4c1e2 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py +++ b/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py @@ -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", diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/core.py b/python/sglang/jit_kernel/trtllm_lora_temp/core.py index c300f92a8..4606f1864 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/core.py +++ b/python/sglang/jit_kernel/trtllm_lora_temp/core.py @@ -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], diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu index 3954f32c0..49148c775 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu @@ -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; } } diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu index 5d1287e1d..ee75dd4f2 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -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 mBaseSupportedTileNums = {8, 16, 32, 64, 128}; + + static std::vector getSupportedTileNums() { + return std::vector(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); + } + + Bf16LoraLauncher( + TensorView const& expert_indices, + Optional 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 + 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(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(const_cast(expert_indices_.data_ptr())), + static_cast(expert_count_histogram.data_ptr()), + static_cast(total_num_padded_tokens.data_ptr()), + static_cast(expanded_idx_to_permuted_idx.data_ptr()), + /*permuted_idx_to_expanded_idx=*/nullptr, + static_cast(permuted_idx_to_token_idx.data_ptr()), + /*expertIds=*/nullptr, + expert_weights_ptr, + static_cast(num_tokens_per_expert.data_ptr()), + static_cast(cta_idx_xy_to_batch_idx.data_ptr()), + static_cast(cta_idx_xy_to_mn_limit.data_ptr()), + static_cast(num_non_exiting_ctas.data_ptr()), + btg::Dtype::Bfloat16, + mRoutingBiasDtype, + /*useRoutingScalesOnInput=*/false, + /*useDeepSeekFp8=*/false, + static_cast(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(expanded_idx_to_permuted_idx.data_ptr()); + permData.hiddenDim = hidden_size; + permData.numTokens = num_tokens; + permData.topK = top_k; + permData.totalNumPaddedTokens = static_cast(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(num_non_exiting_ctas.data_ptr()), + static_cast(total_num_padded_tokens.data_ptr()), + static_cast(cta_idx_xy_to_batch_idx.data_ptr()), + static_cast(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(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(gate_up_lora_delta_.data_ptr()); + actData.activationLoraInputOutPtr = static_cast(activation_lora_input_.data_ptr()); + actData.innerDim = gate_up_n; + actData.numTokens = num_tokens; + actData.topK = top_k; + actData.expandedIdxToPermutedIdx = static_cast(expanded_idx_to_permuted_idx.data_ptr()); + actData.totalNumPaddedTokens = static_cast(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(num_non_exiting_ctas.data_ptr()), + static_cast(total_num_padded_tokens.data_ptr()), + static_cast(cta_idx_xy_to_batch_idx.data_ptr()), + static_cast(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(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(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(total_num_padded_tokens.data_ptr()); + moe::dev::finalize::run(finData, stream); + } + sync_check_cuda_error(stream); + return Array(); + } + + private: + TensorView expert_indices_; + Optional 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 sgl_trtllm_bf16_routed_moe_lora( + TensorView expert_indices, + Optional 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 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 trtllm_mxint4_block_scale_moe( TensorView routing_logits, Optional 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); diff --git a/python/sglang/srt/lora/trtllm_lora_temp/__init__.py b/python/sglang/srt/lora/trtllm_lora_temp/__init__.py index ec07ddf5a..cedc793d0 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/__init__.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/__init__.py @@ -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", ] diff --git a/python/sglang/srt/lora/trtllm_lora_temp/environ.py b/python/sglang/srt/lora/trtllm_lora_temp/environ.py index d84805d3a..48f91b2b4 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/environ.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/environ.py @@ -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. diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py index ee088c517..cebed4c06 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py @@ -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, diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py index b1c4c0840..b8258078e 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py @@ -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 diff --git a/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py b/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py index 9b3934057..e0699c60f 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py @@ -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) diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py index d3452c122..bf184f2be 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py @@ -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,