diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh index 10dc14213..911079923 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh @@ -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; } diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh index 7b4bc0de7..4fed01297 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh @@ -122,9 +122,9 @@ __device__ __forceinline__ void fused_quant_one_row( if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) { sfOffset = static_cast(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(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; } 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 a3006324d..5d1287e1d 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 @@ -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(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(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(routing_replay_out.value().data_ptr()); @@ -527,6 +540,7 @@ class FusedMoeLauncher { static_cast(expanded_idx_to_permuted_idx.data_ptr()), nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, static_cast(permuted_idx_to_token_idx.data_ptr()), + expert_ids_param, workspace.expert_weights, static_cast(num_tokens_per_expert.data_ptr()), static_cast(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(routing_replay_out.value().data_ptr()); @@ -1340,6 +1355,7 @@ class Fp8BlockScaleLauncher : public FusedMoeLauncher { static_cast(expanded_idx_to_permuted_idx.data_ptr()), nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, static_cast(permuted_idx_to_token_idx.data_ptr()), + nullptr, // expertIds - FP8 doesn't support UnpackedPrecomputed mode workspace.expert_weights, static_cast(num_tokens_per_expert.data_ptr()), static_cast(cta_idx_xy_to_batch_idx.data_ptr()), @@ -1605,6 +1621,7 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { } FP4BlockScaleLauncher( + RoutingInputMode routing_input_mode, Optional const& routing_logits, Optional const& routing_bias, TensorView const& hidden_states, @@ -1622,8 +1639,8 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { Optional const& output1_scales_gate_scalar, Optional const& output2_scales_scalar, Optional 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&& args, @@ -1699,8 +1717,8 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { workspace.total_num_padded_tokens = static_cast(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(const_cast(expert_indices.data_ptr())); - workspace.expert_weights = const_cast(expert_weights.data_ptr()); + workspace.routing_expert_indexes = static_cast(const_cast(topk_ids.data_ptr())); + workspace.expert_weights = const_cast(topk_weights.data_ptr()); workspace.permuted_idx_size = static_cast(total_num_padded_tokens.data_ptr()); workspace.expanded_idx_to_permuted_idx = static_cast(expanded_idx_to_permuted_idx.data_ptr()); workspace.permuted_idx_to_token_idx = static_cast(permuted_idx_to_token_idx.data_ptr()); @@ -1823,6 +1841,7 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { } private: + RoutingInputMode routing_input_mode_; Optional hidden_states_scale; TensorView gemm1_weights_scale; Optional 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 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 @@ -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(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(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(expert_indices.data_ptr()), + static_cast(topk_ids.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()), - nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, + nullptr /*permuted_idx_to_expanded_idx*/, static_cast(permuted_idx_to_token_idx.data_ptr()), - expert_weights.data_ptr(), + expert_ids_param, + expert_weights_param, 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()), @@ -2551,9 +2595,10 @@ void sgl_trtllm_fp8_block_scale_moe_lora_finalize( } Array trtllm_fp4_block_scale_moe( + int64_t routing_input_mode, Optional routing_logits, - TensorView expert_indices, - TensorView expert_weights, + TensorView topk_ids, + TensorView topk_weights, Optional routing_bias, TensorView hidden_states, Optional hidden_states_scale, @@ -2686,6 +2731,7 @@ Array trtllm_fp4_block_scale_moe( // Create and initialize launcher for this tile size auto launcher = std::make_unique( + static_cast(routing_input_mode), routing_logits, routing_bias, hidden_states, @@ -2703,8 +2749,8 @@ Array 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(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()), diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu index b42ddc460..e02e51d21 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu @@ -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 @@ -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, diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h b/python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h index d4bf64d58..1d78d81cb 100644 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h @@ -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, 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 5cd3b0c66..ee088c517 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py @@ -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, ) 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 ef02217d3..9b3934057 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py @@ -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, ) diff --git a/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py b/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py index c7321243a..c2914339c 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py @@ -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, ) diff --git a/python/sglang/srt/lora/trtllm_lora_temp/shared_add_overlap.py b/python/sglang/srt/lora/trtllm_lora_temp/shared_add_overlap.py index f7018c153..27fad9366 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/shared_add_overlap.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/shared_add_overlap.py @@ -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() diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py index c212115d2..b8a76de26 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py @@ -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 ) diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py index 1f6adca7c..a69d7b189 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py @@ -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) diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py index 0186ac730..628d1fbb0 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py @@ -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 diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py index a7a3a0085..c109fcf3f 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py @@ -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 diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py index c0c34306f..abcc49d96 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py @@ -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 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 f3bb6d53e..d3452c122 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 @@ -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(