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

[LoRA] Fix experimental fast-path multi-adapter correctness + flashinfer 0.6.12 compatibility (#28091)

Browse Source
This commit is contained in:
Yanbin Jiang
2026-06-19 16:20:19 -07:00
committed by GitHub
parent 364bf976be
commit 6b945c16f4
15 changed files with 229 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh
+2 -2
View File
@@ -163,9 +163,9 @@ __global__ void fusedActivationQuantKernel(
if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
sfOffset = vecOffset;
} else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
sfOffset = tk::get_sf_out_offset_128x4(std::nullopt, permutedIdx, vecIdx, m, num_vecs_per_row);
sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
} else {
sfOffset = tk::get_sf_out_offset_8x4(std::nullopt, permutedIdx, vecIdx, m, num_vecs_per_row);
sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
}
scaleOutput[sfOffset] = fp8Scale;
}
python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh
+4 -4
View File
@@ -122,9 +122,9 @@ __device__ __forceinline__ void fused_quant_one_row(
if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
} else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
sfOffset = tk::get_sf_out_offset_128x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
} else {
sfOffset = tk::get_sf_out_offset_8x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
}
scaleOutput[sfOffset] = fp8Scale;
}
@@ -227,9 +227,9 @@ __global__ void fusedPermuteNvfp4QuantDedupKernel(
if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
} else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
sfOffset = tk::get_sf_out_offset_128x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
} else {
sfOffset = tk::get_sf_out_offset_8x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
}
scaleOutput[sfOffset] = fp8Scale;
}
python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu
+65 -18
View File
@@ -44,6 +44,12 @@ using tensorrt_llm::kernels::trtllmgen_moe::Routing::RoutingMethodType;
using tvm::ffi::Array;
using tvm::ffi::Optional;
enum class RoutingInputMode {
FromLogits, // Mode 1: Compute routing from logits
PackedPrecomputed, // Mode 2: Pre-computed with packed (score << 16 | id) format
UnpackedPrecomputed // Mode 3: Pre-computed with separate topk_ids and topk_weights
};
// Validate routing_replay_out tensor properties.
// NOTE: dim0 >= num_tokens is intentionally NOT checked — with CUDA graphs the buffer
// is pre-allocated at maximum batch size and reused across steps with varying num_tokens.
@@ -429,10 +435,14 @@ class FusedMoeLauncher {
void prepare_moe_common(int64_t& moe_tactic) {
using RunnerType = tensorrt_llm::kernels::trtllmgen_moe::MoE::Runner;
// FIXME(siyuan): check llama4 routing after the fp4 FC1 kernels with bf16 scale factors were
// generated
bool usePerTokenScalingGemm1 =
per_token_scales.has_value() ||
static_cast<RoutingMethodType>(this->routing_method_type) == RoutingMethodType::Llama4;
bool usePerTokenScalingGemm2 = per_token_scales.has_value() && this->mDtypeAct != btg::Dtype::Bfloat16;
per_token_scales.has_value() /* ||
static_cast<RoutingMethodType>(this->routing_method_type) == RoutingMethodType::Llama4*/
;
// FIXME(siyuan): currently only nvfp4 x nvfp4 uses per-token scaling in both FC1 and FC2
bool usePerTokenScalingGemm2 = per_token_scales.has_value() && mDtypeAct == btg::Dtype::E2m1;
// For FP8 block-scale (E4m3 activations, E4m3 weights) with DeepSeek FP8, use the
// weights-only Runner constructor to match the original kernel path and numerics.
if (this->mDtypeAct == btg::Dtype::E4m3 && this->mDtypeWeights == btg::Dtype::E4m3 && args->mUseDeepSeekFp8) {
@@ -505,6 +515,9 @@ class FusedMoeLauncher {
tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim);
cudaStream_t routing_stream = get_stream(hidden_states.device());
// This base class only supports Mode 1 (FromLogits) - compute routing from logits
int32_t* expert_ids_param = nullptr;
int16_t* replay_ptr = nullptr;
if (routing_replay_out.has_value()) {
replay_ptr = reinterpret_cast<int16_t*>(routing_replay_out.value().data_ptr());
@@ -527,6 +540,7 @@ class FusedMoeLauncher {
static_cast<int*>(expanded_idx_to_permuted_idx.data_ptr()),
nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/,
static_cast<int*>(permuted_idx_to_token_idx.data_ptr()),
expert_ids_param,
workspace.expert_weights,
static_cast<int*>(num_tokens_per_expert.data_ptr()),
static_cast<int*>(cta_idx_xy_to_batch_idx.data_ptr()),
@@ -1318,6 +1332,7 @@ class Fp8BlockScaleLauncher : public FusedMoeLauncher {
bool use_precomputed = expert_indices.ndim() == 2 && expert_indices.size(0) > 0;
// When using pre-computed routing, pass nullptr as routing_logits to tell the
// routing runner to use the pre-computed expert indices from workspace.routing_expert_indexes
// FP8 only supports Mode 1 (FromLogits) and Mode 2 (PackedPrecomputed), so expertIds is nullptr
int16_t* replay_ptr = nullptr;
if (routing_replay_out.has_value()) {
replay_ptr = reinterpret_cast<int16_t*>(routing_replay_out.value().data_ptr());
@@ -1340,6 +1355,7 @@ class Fp8BlockScaleLauncher : public FusedMoeLauncher {
static_cast<int*>(expanded_idx_to_permuted_idx.data_ptr()),
nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/,
static_cast<int*>(permuted_idx_to_token_idx.data_ptr()),
nullptr, // expertIds - FP8 doesn't support UnpackedPrecomputed mode
workspace.expert_weights,
static_cast<int*>(num_tokens_per_expert.data_ptr()),
static_cast<int*>(cta_idx_xy_to_batch_idx.data_ptr()),
@@ -1605,6 +1621,7 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
}
FP4BlockScaleLauncher(
RoutingInputMode routing_input_mode,
Optional<TensorView> const& routing_logits,
Optional<TensorView> const& routing_bias,
TensorView const& hidden_states,
@@ -1622,8 +1639,8 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
Optional<TensorView> const& output1_scales_gate_scalar,
Optional<TensorView> const& output2_scales_scalar,
Optional<TensorView> const& per_token_scales,
TensorView const& expert_indices,
TensorView const& expert_weights)
TensorView const& topk_ids,
TensorView const& topk_weights)
: FusedMoeLauncher(
routing_logits,
routing_bias,
@@ -1634,6 +1651,7 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
gemm2_weights,
output2_scales_scalar,
per_token_scales),
routing_input_mode_(routing_input_mode),
hidden_states_scale(hidden_states_scale),
gemm1_weights_scale(gemm1_weights_scale),
gemm1_bias(gemm1_bias),
@@ -1642,8 +1660,8 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
gemm1_clamp_limit(gemm1_clamp_limit),
gemm2_weights_scale(gemm2_weights_scale),
gemm2_bias(gemm2_bias),
expert_indices(expert_indices),
expert_weights(expert_weights) {}
topk_ids(topk_ids),
topk_weights(topk_weights) {}
void init(
std::unique_ptr<tensorrt_llm::kernels::trtllmgen_moe::MoE::MoERunnerArgs>&& args,
@@ -1699,8 +1717,8 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
workspace.total_num_padded_tokens = static_cast<int*>(total_num_padded_tokens.data_ptr());
workspace.total_max_padded_tokens = max_num_padded_tokens;
workspace.ProjUpTileN = tile_tokens_dim;
workspace.routing_expert_indexes = static_cast<int*>(const_cast<void*>(expert_indices.data_ptr()));
workspace.expert_weights = const_cast<void*>(expert_weights.data_ptr());
workspace.routing_expert_indexes = static_cast<int*>(const_cast<void*>(topk_ids.data_ptr()));
workspace.expert_weights = const_cast<void*>(topk_weights.data_ptr());
workspace.permuted_idx_size = static_cast<int*>(total_num_padded_tokens.data_ptr());
workspace.expanded_idx_to_permuted_idx = static_cast<int*>(expanded_idx_to_permuted_idx.data_ptr());
workspace.permuted_idx_to_token_idx = static_cast<int*>(permuted_idx_to_token_idx.data_ptr());
@@ -1823,6 +1841,7 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
}
private:
RoutingInputMode routing_input_mode_;
Optional<TensorView> hidden_states_scale;
TensorView gemm1_weights_scale;
Optional<TensorView> gemm1_bias;
@@ -1834,8 +1853,8 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
int32_t max_num_padded_tokens_gemm1{};
int32_t max_num_padded_tokens_gemm2{};
Optional<Tensor> gemm1_output_scale;
TensorView expert_indices;
TensorView expert_weights;
TensorView topk_ids; // [num_tokens, top_k] - pre-computed or output top-k expert indices
TensorView topk_weights; // [num_tokens, top_k] - pre-computed or output top-k routing weights
public:
Array<Tensor>
@@ -1850,6 +1869,30 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim);
cudaStream_t routing_stream = get_stream(hidden_states.device());
// Set routing kernel parameters based on mode (see RoutingInputMode enum for documentation)
int32_t* expert_ids_param = nullptr; // INPUT: pre-computed expert IDs (Mode 3 only)
void* expert_weights_param = nullptr; // INPUT or OUTPUT depending on mode
switch (routing_input_mode_) {
case RoutingInputMode::FromLogits:
// Mode 1: Kernel computes routing, writes weights to expert_weights_param (OUTPUT)
expert_ids_param = nullptr;
expert_weights_param = topk_weights.data_ptr();
break;
case RoutingInputMode::PackedPrecomputed:
// Mode 2: Kernel unpacks from topk_ids, writes weights to expert_weights_param (OUTPUT)
expert_ids_param = nullptr;
expert_weights_param = topk_weights.data_ptr();
break;
case RoutingInputMode::UnpackedPrecomputed:
// Mode 3: Both are INPUTS, kernel uses them directly
expert_ids_param = static_cast<int32_t*>(topk_ids.data_ptr());
expert_weights_param = topk_weights.data_ptr();
break;
}
int16_t* replay_ptr = nullptr;
if (routing_replay_out.has_value()) {
replay_ptr = reinterpret_cast<int16_t*>(routing_replay_out.value().data_ptr());
@@ -1866,13 +1909,14 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
args->local_expert_offset,
args->local_num_experts,
args->routed_scaling_factor,
static_cast<int*>(expert_indices.data_ptr()),
static_cast<int*>(topk_ids.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()),
nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/,
nullptr /*permuted_idx_to_expanded_idx*/,
static_cast<int*>(permuted_idx_to_token_idx.data_ptr()),
expert_weights.data_ptr(),
expert_ids_param,
expert_weights_param,
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()),
@@ -2551,9 +2595,10 @@ void sgl_trtllm_fp8_block_scale_moe_lora_finalize(
}
Array<Tensor> trtllm_fp4_block_scale_moe(
int64_t routing_input_mode,
Optional<TensorView> routing_logits,
TensorView expert_indices,
TensorView expert_weights,
TensorView topk_ids,
TensorView topk_weights,
Optional<TensorView> routing_bias,
TensorView hidden_states,
Optional<TensorView> hidden_states_scale,
@@ -2686,6 +2731,7 @@ Array<Tensor> trtllm_fp4_block_scale_moe(
// Create and initialize launcher for this tile size
auto launcher = std::make_unique<FP4BlockScaleLauncher>(
static_cast<RoutingInputMode>(routing_input_mode),
routing_logits,
routing_bias,
hidden_states,
@@ -2703,8 +2749,8 @@ Array<Tensor> trtllm_fp4_block_scale_moe(
output1_scales_gate_scalar,
output2_scales_scalar,
per_token_scales,
expert_indices,
expert_weights);
topk_ids,
topk_weights);
launcher->init(
std::move(args),
curr_tile_N,
@@ -2874,6 +2920,7 @@ class FP4BlockScaleLoraLauncher {
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()),
python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu
+19 -7
View File
@@ -21,6 +21,7 @@
#include "flashinfer/trtllm/fused_moe/DevKernel.h"
#include "flashinfer/trtllm/fused_moe/RoutingKernel.h"
#include "flashinfer/trtllm/fused_moe/runner.h"
#include "tensorrt_llm/common/envUtils.h"
#include "tensorrt_llm/kernels/quantization.h"
#include <iostream>
@@ -66,6 +67,7 @@ void Runner::run(
int32_t* expandedIdxToPermutedIdx,
int32_t* permutedIdxToExpandedIdx,
int32_t* permutedIdxToTokenIdx,
int32_t* expertIds,
void* expertWeights,
int32_t* numTokensPerExpert,
int32_t* ctaIdxXyToBatchIdx,
@@ -95,7 +97,8 @@ void Runner::run(
routingData.mDtypeBias = dtypeBias;
routingData.mRouteScale = routedScalingFactor;
routingData.mPtrScores = routingLogits;
routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr;
routingData.mPtrTopKIds = expertIds;
routingData.mPtrTopKPacked = routingExpertIndexes;
routingData.mPtrExpertCounts = expertCountHistogram;
routingData.mPtrPermutedIdxSize = permutedIdxSize;
@@ -135,7 +138,8 @@ void Runner::run(
routingData.mRouteScale = 1.0f;
routingData.mSumEpsilon = 1e-20f;
routingData.mPtrScores = routingLogits;
routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr;
routingData.mPtrTopKIds = expertIds;
routingData.mPtrTopKPacked = routingExpertIndexes;
routingData.mPtrExpertCounts = expertCountHistogram;
routingData.mPtrPermutedIdxSize = permutedIdxSize;
@@ -183,7 +187,9 @@ void Runner::run(
// input:
routingData.mPtrRoutingBias = routingBias;
routingData.mPtrScores = routingLogits; // type-erased; InputT selected by forceFloatInput
// Pre-computed routing support: when expertIds is provided, use it directly
routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr;
routingData.mPtrTopKIds = expertIds;
routingData.mNumTokens = numTokens;
routingData.mNumExperts = numExperts;
routingData.mNumExpertGroups = nGroup;
@@ -222,7 +228,9 @@ void Runner::run(
routingData.mPtrNumNonExitingCtas = numNonExitingCtas;
// input:
routingData.mPtrScores = routingLogits;
// Pre-computed routing support: when expertIds is provided, use it directly
routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr;
routingData.mPtrTopKIds = expertIds;
routingData.mNumTokens = numTokens;
routingData.mNumExperts = numExperts;
routingData.mTopK = topK;
@@ -282,7 +290,9 @@ void Runner::run(
routingData.mPostprocessType = RoutingPostprocessType::None;
}
routingData.mPtrScores = routingLogits;
// Pre-computed routing support: when expertIds is provided, use it directly
routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr;
routingData.mPtrTopKIds = expertIds;
//
// Outputs
@@ -1033,8 +1043,10 @@ void Runner::run(
auto sfLayout =
mGemm2.mTileTokensDim >= 128 ? QuantizationSFLayout::SWIZZLED_128x4 : QuantizationSFLayout::SWIZZLED_8x4;
// TODO(siyuan): should this value be exposed?
float globalScaleInv = 1.f / 448.f / 6.f;
float globalScaleInv = 1.f / (448.f * 6.f);
if (tensorrt_llm::common::getEnvNVFP4Use4Over6() && tensorrt_llm::common::getEnvNVFP44Over6E4M3Use256()) {
globalScaleInv = 1.f / (256.f * 6.f);
}
invokeNvfp4QuantAndPerTokenScale<__nv_bfloat16>(
args.num_tokens * args.top_k,
args.intermediate_size,
python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h
+1
View File
@@ -149,6 +149,7 @@ class Runner {
int32_t* expandedIdxToPermutedIdx,
int32_t* permutedIdxToExpandedIdx,
int32_t* permutedIdxToTokenIdx,
int32_t* expertIds,
void* expertWeights,
int32_t* numTokensPerExpert,
int32_t* ctaIdxXyToBatchIdx,
python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py
+4 -6
View File
@@ -50,9 +50,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
trtllm_fp8_block_scale_moe_lora_finalize,
trtllm_fp8_block_scale_routed_moe_lora,
)
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
_pack_topk_for_flashinfer_routed,
)
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
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
@@ -160,7 +158,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
# the padded-region id=-1 mask. Fall back to the separate pack otherwise.
packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
if packed_topk_ids is None:
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
@@ -318,8 +316,8 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
from sglang.jit_kernel.trtllm_lora_temp import (
trtllm_fp4_block_scale_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 (
_pack_topk_for_flashinfer_routed,
fused_experts_none_to_flashinfer_trtllm_fp4,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
@@ -397,7 +395,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
device=hidden_states.device,
)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py
+60 -8
View File
@@ -56,14 +56,12 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
from sglang.jit_kernel.trtllm_lora_temp import (
trtllm_fp8_block_scale_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 (
_pack_topk_for_flashinfer_routed,
)
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
@@ -127,6 +125,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
intermediate_buffer=gate_up_lora_intermediate,
)
# GEMM1-LoRA overlap: fire the gate_up LoRA on the side stream + record an event; the
@@ -135,6 +134,34 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
# whole op.
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.
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=quant_info.local_num_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],
)
)
# O1 fork — gate_up shrink/expand on side stream concurrent with the main-stream
# per-token-group FP8 quant + the trtllm op's permute+GEMM1 below.
side_stream.wait_stream(torch.cuda.current_stream())
@@ -162,7 +189,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
# the padded-region id=-1 mask. Fall back to the separate pack otherwise.
packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
if packed_topk_ids is None:
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
@@ -333,14 +360,12 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
from sglang.jit_kernel.trtllm_lora_temp import (
trtllm_fp4_block_scale_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 (
_pack_topk_for_flashinfer_routed,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
@@ -391,8 +416,35 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
intermediate_buffer=gate_up_lora_intermediate,
)
# Hoist every side-chain allocation onto the MAIN stream (cuda-graph
# allocator safety -- see the "routing" stage in virtual_experts.py).
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=quant_info.local_num_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],
)
)
# O1-fp4 fork: gate_up shrink/expand on the side stream, concurrent with the
# FP4 op's permute + gate_up GEMM1 below. The op waits on lora_event right
# before its activation kernel (the only consumer of gate_up_delta).
@@ -407,7 +459,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
dtype=hidden_states.dtype,
device=hidden_states.device,
)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py
+2 -2
View File
@@ -30,8 +30,8 @@ def fused_experts_fp8_sgl(
# <-> quantization import cycle at load time.
from flashinfer.fused_moe import Fp8QuantizationType
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
_pack_topk_for_flashinfer_routed,
get_tp_group,
is_allocation_symmetric,
next_power_of_2,
@@ -118,7 +118,7 @@ def fused_experts_fp8_sgl(
runner_config.top_k is not None
), "runner_config.top_k is required for flashinfer_trtllm_routed."
assert TopKOutputChecker.format_is_standard(topk_output)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_output.topk_ids,
topk_weights=topk_output.topk_weights,
)
python/sglang/srt/lora/trtllm_lora_temp/shared_add_overlap.py
+7
View File
@@ -100,6 +100,13 @@ def maybe_overlap_staged_shared_add(output: torch.Tensor) -> Optional[torch.cuda
# Single-stream caller: nothing to overlap. Leave the staging in place
# so the model layer reclaims it and does the add as before.
return None
if torch.cuda.is_current_stream_capturing():
# The cross-stream producer-stream add_ (ordered via base_ready/add_done
# events) is NOT cuda-graph-capture-safe: it corrupts `output` on replay.
# Fall back to the serial caller-side add -- leave the staging so the model
# layer reclaims it via unstage_shared_expert_add and adds shared_output
# after current_stream.wait_stream(alt_stream).
return None
_PENDING = None
base_ready = torch.cuda.Event()
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py
+7 -3
View File
@@ -208,9 +208,13 @@ def gate_up_lora_b_fwd(
assert input_dim == 2 * r
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_GATE_UP.get()
) and s * r >= _CUBLAS_MIN_S_RANK:
(
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_GATE_UP.get()
)
and s * r >= _CUBLAS_MIN_S_RANK
and gate_up_lora_b.shape[0] == 1
): # single-adapter fast path: only valid with one resident slot
return _gate_up_lora_b_cublas(
x, gate_up_lora_b, batch_info, output_dim, base_output
)
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py
+13 -5
View File
@@ -259,7 +259,9 @@ def step_a_q_fwd(
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
):
) and B_buf.shape[
0
] == 1: # single-adapter fast path: only valid with one resident slot
# (S,H,r) view of a (H,S,r)-contiguous bmm result; step_b_q's dense
# path flattens in (h,s) order, so the chain needs no copies.
w_kc = B_buf[0].view(H, full_K_per_head, -1)[:, :qk_nope_dim, :]
@@ -474,7 +476,9 @@ def step_b_q_fwd(
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
):
) and A_buf.shape[
0
] == 1: # single-adapter fast path: only valid with one resident slot
# Flatten (S,H) in whichever order base_output's storage allows
# without a copy (the absorbed q path passes a transpose view of a
# (H,S,kv)-contiguous bmm result). x is small; reshape may copy it.
@@ -688,9 +692,13 @@ def step_a_v_fwd(
rank = A_buf.shape[1]
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
) and attn_output.is_contiguous():
(
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
)
and attn_output.is_contiguous()
and A_buf.shape[0] == 1
): # single-adapter fast path: only valid with one resident slot
return torch.mm(
attn_output.view(-1, kv_lora_rank), A_buf[0, :rank, :].t()
).view(S, H, rank)
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py
+2
View File
@@ -234,6 +234,8 @@ def qkv_lora_b_fwd(
or lora_envs.SGLANG_OPT_LORA_CUBLAS_QKV.get()
)
and batch_info.max_len >= _CUBLAS_MIN_MAX_LEN
and qkv_lora_b.shape[0]
== 1 # single-adapter fast path: only valid with one resident slot
):
return _qkv_lora_b_cublas(
x, qkv_lora_b, batch_info, output_offset_cpu, base_output, n_slices
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py
+3 -1
View File
@@ -180,7 +180,9 @@ def sgemm_lora_a_fwd(
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_A.get()
):
) and weights.shape[
0
] == 1: # single-adapter fast path: only valid with one resident slot
# Honor out_alloc_stream like the Triton path below: under SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC
# the shrink output must be allocated on the MAIN (consumer) stream so the caching allocator
# frees/reuses it on the consumer's schedule (cuda-graph WAR). F.linear has no out=, so
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py
+7 -3
View File
@@ -168,9 +168,13 @@ def sgemm_lora_b_fwd(
assert x.shape[-1] == R
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_B.get()
) and S * R >= _CUBLAS_MIN_S_RANK:
(
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_B.get()
)
and S * R >= _CUBLAS_MIN_S_RANK
and weights.shape[0] == 1
): # single-adapter fast path: only valid with one resident slot
return _sgemm_lora_b_cublas(x, weights, batch_info, base_output)
# Block shapes
BLOCK_S = 16
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py
+33 -1
View File
@@ -859,13 +859,45 @@ def _merged_experts_fused_moe_lora_add_impl(
invoke_fused_moe_kernel,
)
assert stage in ("all", "shrink", "expand"), f"invalid stage {stage!r}"
assert stage in (
"all",
"shrink",
"expand",
"routing",
), f"invalid stage {stage!r}"
lora_a_virtual = _merge_lora_expert_weight(lora_a)
lora_b_virtual = _merge_lora_expert_weight(lora_b)
num_experts_a = lora_a.shape[1]
num_experts_b = lora_b.shape[1]
b_stage_config = _get_stage_config(lora_b_virtual, 1)
if stage == "routing":
# Pre-warm the routing cache on the CALLER'S (main) stream so the
# side-stream chain performs no allocations. Tensors allocated inside a
# side-stream context during cuda-graph capture can be pool-reused by
# later allocations on other streams with no cross-stream guard (the
# allocator's stream tracking is disabled while capturing), corrupting
# replays. Routing needs only topk_ids + token_lora_mapping, which are
# both ready before the side-stream fork, so it can run on main.
a_cfg = _get_shrink_stage_config(lora_a_virtual, token_lora_mapping.shape[0])
if lora_envs.SGLANG_OPT_LORA_SHRINK_TUNE.get():
a_cfg = {**a_cfg, "BLOCK_SIZE_M": 16}
_get_routing(
topk_ids,
token_lora_mapping,
num_experts_a,
experts_shared_outer_loras_a,
a_cfg["BLOCK_SIZE_M"],
)
_get_routing(
topk_ids,
token_lora_mapping,
num_experts_b,
experts_shared_outer_loras_b,
b_stage_config["BLOCK_SIZE_M"],
)
return None
intermediate = intermediate_buffer
if stage != "expand":
a_stage_config = _get_shrink_stage_config(