[LoRA] Experimental fast LoRA path with experimental_sgl_trtllm MoE backend for FP8 and NVFP4 models (#27329)
Co-authored-by: fzyzcjy <ch271828n@outlook.com> Co-authored-by: Chunan Zeng <zcnrex@gmail.com> Co-authored-by: Ethan (Yusheng) Su <yushengsu.thu@gmail.com>
This commit is contained in:
Yanbin Jiang
co-authored by
fzyzcjy
Chunan Zeng
Ethan Su
fzyzcjy
Chunan Zeng
Ethan Su
parent
d381ec7997
commit
c9f582a272
@@ -0,0 +1,453 @@
|
||||
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.h> // For RuntimeCheck, Panic, div_ceil
|
||||
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
// Kimi K2 MoE fused gate, supports NUM_EXPERTS in {256 (MiMo V2 Flash), 384 (Kimi K2)}.
|
||||
// Routing (DeepSeek "noaux_tc" with num_expert_group = 1):
|
||||
// 1. sigmoid(gate_logit)
|
||||
// 2. add per-expert correction bias (ranking only)
|
||||
// 3. pick top-k by biased score
|
||||
// 4. weights = sigmoid (no bias)
|
||||
// 5. optional renorm; routed_scaling_factor folded into renorm (no-op when not renormalizing)
|
||||
|
||||
__device__ __forceinline__ float sigmoid_accurate(float x) {
|
||||
return 1.0f / (1.0f + expf(-x));
|
||||
}
|
||||
|
||||
// Scalar widening: input/bias may arrive as fp32, bf16, or fp16; the kernel math
|
||||
// always runs in fp32. Widening bf16/fp16 -> fp32 is exact, so results are
|
||||
// bitwise identical to upcasting on the host first (the casts we are removing).
|
||||
__device__ __forceinline__ float to_float(float x) {
|
||||
return x;
|
||||
}
|
||||
__device__ __forceinline__ float to_float(__nv_bfloat16 x) {
|
||||
return __bfloat162float(x);
|
||||
}
|
||||
__device__ __forceinline__ float to_float(__half x) {
|
||||
return __half2float(x);
|
||||
}
|
||||
|
||||
// Vectorized load of 4 consecutive elements of type T at vector index `vec_idx`,
|
||||
// widened to a float4. fp32 reads a 16B float4; bf16/fp16 read an 8B float2 and
|
||||
// expand. Used only by the large-token kernel's lane-strided loads.
|
||||
template <typename T>
|
||||
struct VecLoader;
|
||||
|
||||
template <>
|
||||
struct VecLoader<float> {
|
||||
__device__ __forceinline__ static float4 load(const float* base, int vec_idx) {
|
||||
return reinterpret_cast<const float4*>(base)[vec_idx];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecLoader<__nv_bfloat16> {
|
||||
__device__ __forceinline__ static float4 load(const __nv_bfloat16* base, int vec_idx) {
|
||||
float2 raw = reinterpret_cast<const float2*>(base)[vec_idx]; // 4 bf16 = 8 bytes
|
||||
const __nv_bfloat162* packed = reinterpret_cast<const __nv_bfloat162*>(&raw);
|
||||
float2 lo = __bfloat1622float2(packed[0]);
|
||||
float2 hi = __bfloat1622float2(packed[1]);
|
||||
return make_float4(lo.x, lo.y, hi.x, hi.y);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecLoader<__half> {
|
||||
__device__ __forceinline__ static float4 load(const __half* base, int vec_idx) {
|
||||
float2 raw = reinterpret_cast<const float2*>(base)[vec_idx]; // 4 fp16 = 8 bytes
|
||||
const __half2* packed = reinterpret_cast<const __half2*>(&raw);
|
||||
float2 lo = __half22float2(packed[0]);
|
||||
float2 hi = __half22float2(packed[1]);
|
||||
return make_float4(lo.x, lo.y, hi.x, hi.y);
|
||||
}
|
||||
};
|
||||
|
||||
template <int N>
|
||||
struct GateConfig {
|
||||
static_assert(
|
||||
N == 256 || N == 384,
|
||||
"kimi_k2_moe_fused_gate currently only supports "
|
||||
"NUM_EXPERTS == 256 or 384");
|
||||
static constexpr int NUM_EXPERTS = N;
|
||||
static constexpr int WARP_SIZE = 32;
|
||||
static constexpr int WARPS_PER_CTA = 6; // only used by the large-token kernel
|
||||
static constexpr int VPT = N / 32; // 8 (256) or 12 (384)
|
||||
static constexpr int VEC_SIZE = 4;
|
||||
static constexpr int VEC_PER_LANE = VPT / VEC_SIZE; // 2 or 3
|
||||
static constexpr int WARPS_PER_TOKEN_SMALL = N / 32; // 8 or 12
|
||||
static constexpr int THREADS_PER_BLOCK_SMALL = N; // 256 or 384
|
||||
static constexpr int SMALL_TOKEN_THRESHOLD = 512;
|
||||
static constexpr int MAX_TOPK = 8; // must match RuntimeCheck(topk <= 8) at the host launcher
|
||||
static_assert(VPT % VEC_SIZE == 0, "VPT must be a multiple of VEC_SIZE for the float4 vec load");
|
||||
};
|
||||
|
||||
// Small-token kernel: 1 block per token, NUM_EXPERTS threads (1 thread = 1 expert).
|
||||
template <int N, typename InputT, typename BiasT>
|
||||
__global__ void kimi_k2_moe_fused_gate_kernel_small_token(
|
||||
const InputT* input,
|
||||
const BiasT* bias,
|
||||
float* output_ptr,
|
||||
int32_t* indices_ptr,
|
||||
int64_t num_rows,
|
||||
int64_t topk,
|
||||
bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
bool apply_routed_scaling_factor_on_output) {
|
||||
using Cfg = GateConfig<N>;
|
||||
constexpr int NUM_EXPERTS = Cfg::NUM_EXPERTS;
|
||||
constexpr int WARP_SIZE = Cfg::WARP_SIZE;
|
||||
constexpr int WARPS_PER_TOKEN_SMALL = Cfg::WARPS_PER_TOKEN_SMALL;
|
||||
constexpr int MAX_TOPK = Cfg::MAX_TOPK;
|
||||
|
||||
int64_t row_idx = blockIdx.x;
|
||||
if (row_idx >= num_rows) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int warp_id = tid / WARP_SIZE;
|
||||
int lane_id = tid % WARP_SIZE;
|
||||
|
||||
// Sigmoid weights (no bias) for final lookup, indexed by expert id.
|
||||
__shared__ float shared_original_scores[NUM_EXPERTS];
|
||||
__shared__ float warp_maxs[WARPS_PER_TOKEN_SMALL];
|
||||
__shared__ int warp_experts[WARPS_PER_TOKEN_SMALL];
|
||||
__shared__ int selected_experts[MAX_TOPK];
|
||||
|
||||
// Keep biased_val in register; mask the winner in-place each iteration to
|
||||
// avoid round-tripping through shared memory.
|
||||
float input_val = to_float(input[row_idx * NUM_EXPERTS + tid]);
|
||||
float bias_val = to_float(bias[tid]);
|
||||
float sigmoid_val = sigmoid_accurate(input_val);
|
||||
float biased_val = sigmoid_val + bias_val;
|
||||
shared_original_scores[tid] = sigmoid_val;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Lane 0 of warp 0 accumulates the renorm sum as it picks each winner,
|
||||
// saving a second pass over selected_experts during writeback.
|
||||
float sum_for_renorm = 0.0f;
|
||||
|
||||
for (int k = 0; k < topk; k++) {
|
||||
// Stage 1: per-warp argmax.
|
||||
float warp_max_val = biased_val;
|
||||
int warp_max_expert = tid;
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
float other_val = __shfl_down_sync(0xFFFFFFFF, warp_max_val, offset);
|
||||
int other_expert = __shfl_down_sync(0xFFFFFFFF, warp_max_expert, offset);
|
||||
if (other_val > warp_max_val) {
|
||||
warp_max_val = other_val;
|
||||
warp_max_expert = other_expert;
|
||||
}
|
||||
}
|
||||
if (lane_id == 0) {
|
||||
warp_maxs[warp_id] = warp_max_val;
|
||||
warp_experts[warp_id] = warp_max_expert;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Stage 2: warp 0 merges warp-leaders into a single winner.
|
||||
if (warp_id == 0) {
|
||||
float final_max = (lane_id < WARPS_PER_TOKEN_SMALL) ? warp_maxs[lane_id] : -FLT_MAX;
|
||||
int final_expert = (lane_id < WARPS_PER_TOKEN_SMALL) ? warp_experts[lane_id] : -1;
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
float other_val = __shfl_down_sync(0xFFFFFFFF, final_max, offset);
|
||||
int other_expert = __shfl_down_sync(0xFFFFFFFF, final_expert, offset);
|
||||
if (other_val > final_max) {
|
||||
final_max = other_val;
|
||||
final_expert = other_expert;
|
||||
}
|
||||
}
|
||||
if (lane_id == 0) {
|
||||
selected_experts[k] = final_expert;
|
||||
if (renormalize && final_expert >= 0 && final_expert < NUM_EXPERTS) {
|
||||
sum_for_renorm += shared_original_scores[final_expert];
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int selected = selected_experts[k];
|
||||
if (tid == selected) biased_val = -FLT_MAX;
|
||||
}
|
||||
|
||||
// Lane 0 of warp 0 writes the output. sum_for_renorm was accumulated
|
||||
// during the topk loop, so we just fold it into rcp.
|
||||
if (warp_id == 0 && lane_id == 0) {
|
||||
float rcp = 1.0f;
|
||||
if (renormalize && sum_for_renorm > 0.0f) {
|
||||
rcp = 1.0f / sum_for_renorm;
|
||||
if (apply_routed_scaling_factor_on_output) {
|
||||
rcp *= static_cast<float>(routed_scaling_factor);
|
||||
}
|
||||
}
|
||||
|
||||
for (int k = 0; k < topk; k++) {
|
||||
int expert_id = selected_experts[k];
|
||||
bool valid = (expert_id >= 0 && expert_id < NUM_EXPERTS);
|
||||
output_ptr[row_idx * topk + k] = valid ? shared_original_scores[expert_id] * rcp : 0.0f;
|
||||
indices_ptr[row_idx * topk + k] = valid ? expert_id : 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Large-token kernel: 1 warp per token, WARPS_PER_CTA warps per block.
|
||||
template <int N, typename InputT, typename BiasT>
|
||||
__global__ void kimi_k2_moe_fused_gate_kernel(
|
||||
const InputT* input,
|
||||
const BiasT* bias,
|
||||
float* output_ptr,
|
||||
int32_t* indices_ptr,
|
||||
int64_t num_rows,
|
||||
int64_t topk,
|
||||
bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
bool apply_routed_scaling_factor_on_output) {
|
||||
using Cfg = GateConfig<N>;
|
||||
constexpr int NUM_EXPERTS = Cfg::NUM_EXPERTS;
|
||||
constexpr int WARP_SIZE = Cfg::WARP_SIZE;
|
||||
constexpr int WARPS_PER_CTA = Cfg::WARPS_PER_CTA;
|
||||
constexpr int VEC_SIZE = Cfg::VEC_SIZE;
|
||||
constexpr int VEC_PER_LANE = Cfg::VEC_PER_LANE;
|
||||
constexpr int MAX_TOPK = Cfg::MAX_TOPK;
|
||||
|
||||
int64_t row_idx = blockIdx.x * WARPS_PER_CTA + threadIdx.y;
|
||||
if (row_idx >= num_rows) return;
|
||||
|
||||
int lane_id = threadIdx.x;
|
||||
int warp_id = threadIdx.y;
|
||||
|
||||
__shared__ float shared_scores[NUM_EXPERTS * WARPS_PER_CTA];
|
||||
__shared__ float shared_original_scores[NUM_EXPERTS * WARPS_PER_CTA];
|
||||
float* warp_scores = shared_scores + warp_id * NUM_EXPERTS;
|
||||
float* warp_original_scores = shared_original_scores + warp_id * NUM_EXPERTS;
|
||||
float4* warp_scores_v4 = reinterpret_cast<float4*>(warp_scores);
|
||||
float4* warp_original_scores_v4 = reinterpret_cast<float4*>(warp_original_scores);
|
||||
|
||||
const InputT* input_row = input + row_idx * NUM_EXPERTS;
|
||||
|
||||
// Lane-strided vec_idx (each lane k stores at vec_idx k, k+32, k+64, ...) so each
|
||||
// iteration's STS.128 is lane-contiguous, avoiding shared-mem bank conflicts.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < VEC_PER_LANE; i++) {
|
||||
int vec_idx = lane_id + i * WARP_SIZE;
|
||||
float4 input_val = VecLoader<InputT>::load(input_row, vec_idx);
|
||||
float4 bias_val = VecLoader<BiasT>::load(bias, vec_idx);
|
||||
|
||||
float4 sigmoid_v4;
|
||||
float4 biased_v4;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC_SIZE; j++) {
|
||||
float inp = ((float*)&input_val)[j];
|
||||
float b = ((float*)&bias_val)[j];
|
||||
float sigmoid_val = sigmoid_accurate(inp);
|
||||
((float*)&sigmoid_v4)[j] = sigmoid_val;
|
||||
((float*)&biased_v4)[j] = sigmoid_val + b;
|
||||
}
|
||||
warp_original_scores_v4[vec_idx] = sigmoid_v4;
|
||||
warp_scores_v4[vec_idx] = biased_v4;
|
||||
}
|
||||
|
||||
__syncwarp();
|
||||
|
||||
// Lane 0 records the picked expert ids and accumulates the renorm sum as
|
||||
// it goes; the global write is a single pass after the loop.
|
||||
int top_indices[MAX_TOPK];
|
||||
float sum_for_renorm = 0.0f;
|
||||
|
||||
for (int k = 0; k < topk; k++) {
|
||||
float max_val = -FLT_MAX;
|
||||
int max_expert = -1;
|
||||
|
||||
for (int expert = lane_id; expert < NUM_EXPERTS; expert += WARP_SIZE) {
|
||||
if (warp_scores[expert] > max_val) {
|
||||
max_val = warp_scores[expert];
|
||||
max_expert = expert;
|
||||
}
|
||||
}
|
||||
|
||||
// warp shfl reduce; tie-break by lower expert id
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
float other_val = __shfl_down_sync(0xFFFFFFFF, max_val, offset);
|
||||
int other_expert = __shfl_down_sync(0xFFFFFFFF, max_expert, offset);
|
||||
if (other_val > max_val || (other_val == max_val && other_expert < max_expert)) {
|
||||
max_val = other_val;
|
||||
max_expert = other_expert;
|
||||
}
|
||||
}
|
||||
|
||||
if (lane_id == 0) {
|
||||
bool valid = (max_expert >= 0 && max_expert < NUM_EXPERTS);
|
||||
top_indices[k] = valid ? max_expert : -1;
|
||||
if (renormalize && valid) {
|
||||
sum_for_renorm += warp_original_scores[max_expert];
|
||||
}
|
||||
if (valid) warp_scores[max_expert] = -FLT_MAX;
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
if (lane_id == 0) {
|
||||
float rcp = 1.0f;
|
||||
if (renormalize && sum_for_renorm > 0.0f) {
|
||||
rcp = 1.0f / sum_for_renorm;
|
||||
if (apply_routed_scaling_factor_on_output) {
|
||||
rcp *= static_cast<float>(routed_scaling_factor);
|
||||
}
|
||||
}
|
||||
|
||||
for (int k = 0; k < topk; k++) {
|
||||
int e = top_indices[k];
|
||||
bool valid = (e >= 0);
|
||||
output_ptr[row_idx * topk + k] = valid ? warp_original_scores[e] * rcp : 0.0f;
|
||||
indices_ptr[row_idx * topk + k] = valid ? e : 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Bundles the dtype-agnostic launch parameters so the templated dispatch below
|
||||
// only has to thread the typed input/bias pointers.
|
||||
struct GateLaunchArgs {
|
||||
float* output;
|
||||
int32_t* indices;
|
||||
int64_t num_rows;
|
||||
int64_t topk;
|
||||
bool renormalize;
|
||||
double routed_scaling_factor;
|
||||
bool apply_routed_scaling_factor_on_output;
|
||||
DLDevice device;
|
||||
};
|
||||
|
||||
template <int N, typename InputT, typename BiasT>
|
||||
void launch_for_n(const InputT* input, const BiasT* bias, const GateLaunchArgs& args) {
|
||||
using namespace host;
|
||||
using Cfg = GateConfig<N>;
|
||||
bool use_small_token_kernel = args.num_rows <= Cfg::SMALL_TOKEN_THRESHOLD;
|
||||
|
||||
if (use_small_token_kernel) {
|
||||
LaunchKernel(
|
||||
static_cast<uint32_t>(args.num_rows), static_cast<uint32_t>(Cfg::THREADS_PER_BLOCK_SMALL), args.device)(
|
||||
kimi_k2_moe_fused_gate_kernel_small_token<N, InputT, BiasT>,
|
||||
input,
|
||||
bias,
|
||||
args.output,
|
||||
args.indices,
|
||||
args.num_rows,
|
||||
args.topk,
|
||||
args.renormalize,
|
||||
args.routed_scaling_factor,
|
||||
args.apply_routed_scaling_factor_on_output);
|
||||
} else {
|
||||
uint32_t num_blocks = div_ceil(args.num_rows, static_cast<int64_t>(Cfg::WARPS_PER_CTA));
|
||||
dim3 block_dim(Cfg::WARP_SIZE, Cfg::WARPS_PER_CTA);
|
||||
LaunchKernel(num_blocks, block_dim, args.device)(
|
||||
kimi_k2_moe_fused_gate_kernel<N, InputT, BiasT>,
|
||||
input,
|
||||
bias,
|
||||
args.output,
|
||||
args.indices,
|
||||
args.num_rows,
|
||||
args.topk,
|
||||
args.renormalize,
|
||||
args.routed_scaling_factor,
|
||||
args.apply_routed_scaling_factor_on_output);
|
||||
}
|
||||
}
|
||||
|
||||
// input/bias each independently arrive as fp32, bf16, or fp16; widen both to
|
||||
// fp32 inside the kernel so the host no longer has to upcast. Dispatch is nested:
|
||||
// num_experts -> input dtype -> bias dtype.
|
||||
template <int N, typename InputT>
|
||||
void dispatch_bias(
|
||||
const InputT* input, const void* bias, const host::SymbolicDType& bias_dtype, const GateLaunchArgs& args) {
|
||||
using namespace host;
|
||||
if (bias_dtype.is_type<float>()) {
|
||||
launch_for_n<N, InputT, float>(input, static_cast<const float*>(bias), args);
|
||||
} else if (bias_dtype.is_type<bf16_t>()) {
|
||||
launch_for_n<N, InputT, bf16_t>(input, static_cast<const bf16_t*>(bias), args);
|
||||
} else {
|
||||
launch_for_n<N, InputT, fp16_t>(input, static_cast<const fp16_t*>(bias), args);
|
||||
}
|
||||
}
|
||||
|
||||
template <int N>
|
||||
void dispatch_input(
|
||||
const void* input,
|
||||
const host::SymbolicDType& input_dtype,
|
||||
const void* bias,
|
||||
const host::SymbolicDType& bias_dtype,
|
||||
const GateLaunchArgs& args) {
|
||||
using namespace host;
|
||||
if (input_dtype.is_type<float>()) {
|
||||
dispatch_bias<N, float>(static_cast<const float*>(input), bias, bias_dtype, args);
|
||||
} else if (input_dtype.is_type<bf16_t>()) {
|
||||
dispatch_bias<N, bf16_t>(static_cast<const bf16_t*>(input), bias, bias_dtype, args);
|
||||
} else {
|
||||
dispatch_bias<N, fp16_t>(static_cast<const fp16_t*>(input), bias, bias_dtype, args);
|
||||
}
|
||||
}
|
||||
|
||||
struct KimiK2MoEFusedGateKernel {
|
||||
static void
|
||||
run(const tvm::ffi::TensorView input,
|
||||
const tvm::ffi::TensorView bias,
|
||||
const tvm::ffi::TensorView output,
|
||||
const tvm::ffi::TensorView indices,
|
||||
int64_t topk,
|
||||
bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
bool apply_routed_scaling_factor_on_output) {
|
||||
using namespace host;
|
||||
|
||||
auto N = SymbolicSize{"num_rows"};
|
||||
auto E = SymbolicSize{"num_experts"};
|
||||
auto K = SymbolicSize{"topk"};
|
||||
auto input_dtype = SymbolicDType{};
|
||||
auto bias_dtype = SymbolicDType{};
|
||||
auto device = SymbolicDevice{};
|
||||
K.set_value(topk);
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, E}).with_dtype<float, bf16_t, fp16_t>(input_dtype).with_device(device).verify(input);
|
||||
TensorMatcher({E}).with_dtype<float, bf16_t, fp16_t>(bias_dtype).with_device(device).verify(bias);
|
||||
TensorMatcher({N, K}).with_dtype<float>().with_device(device).verify(output);
|
||||
TensorMatcher({N, K}).with_dtype<int32_t>().with_device(device).verify(indices);
|
||||
|
||||
const auto num_rows = static_cast<int64_t>(N.unwrap());
|
||||
const auto num_experts = static_cast<int64_t>(E.unwrap());
|
||||
|
||||
RuntimeCheck(topk <= 8, "kimi_k2_moe_fused_gate only supports topk <= 8, got ", topk);
|
||||
|
||||
const GateLaunchArgs args{
|
||||
.output = static_cast<float*>(output.data_ptr()),
|
||||
.indices = static_cast<int32_t*>(indices.data_ptr()),
|
||||
.num_rows = num_rows,
|
||||
.topk = topk,
|
||||
.renormalize = renormalize,
|
||||
.routed_scaling_factor = routed_scaling_factor,
|
||||
.apply_routed_scaling_factor_on_output = apply_routed_scaling_factor_on_output,
|
||||
.device = device.unwrap()};
|
||||
|
||||
switch (num_experts) {
|
||||
case 256:
|
||||
dispatch_input<256>(input.data_ptr(), input_dtype, bias.data_ptr(), bias_dtype, args);
|
||||
break;
|
||||
case 384:
|
||||
dispatch_input<384>(input.data_ptr(), input_dtype, bias.data_ptr(), bias_dtype, args);
|
||||
break;
|
||||
default:
|
||||
Panic("kimi_k2_moe_fused_gate only supports num_experts in {256, 384}, got ", num_experts);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -0,0 +1,589 @@
|
||||
/* Copyright 2025 SGLang Team. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
// LoRA merged-virtual-expert routing align, fused with the virtual-expert id
|
||||
// computation. Replaces the (_fused_virtual_topk_ids triton kernel + native
|
||||
// moe_align_block_size) two-launch pair on the `--lora-use-virtual-experts`
|
||||
// path: the align/scatter kernels read the RAW topk_ids + token_lora_mapping
|
||||
// and compute the merged virtual id inline (mirrors _fused_virtual_topk_ids),
|
||||
// so virtual_topk_ids is never materialized to global memory.
|
||||
//
|
||||
// Commit 1 scope: pure fusion (inline virtual id), NO EP skip. Output is
|
||||
// bucket-for-bucket equivalent to the old path (dropped/-1 tokens still land in
|
||||
// the sentinel bucket 0), so it can be asserted equal to the old kernels.
|
||||
// Only the `64 < num_buckets <= 1024` branch is implemented here; other expert
|
||||
// counts keep the old path (handled by the Python dispatcher).
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#ifndef WARP_SIZE
|
||||
#define WARP_SIZE 32
|
||||
#endif
|
||||
|
||||
#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
|
||||
|
||||
#define VEC_SIZE 4
|
||||
using Vec = int4;
|
||||
|
||||
inline uint32_t next_pow2(uint32_t x) noexcept {
|
||||
--x;
|
||||
x |= x >> 1;
|
||||
x |= x >> 2;
|
||||
x |= x >> 4;
|
||||
x |= x >> 8;
|
||||
x |= x >> 16;
|
||||
return x + 1;
|
||||
}
|
||||
|
||||
namespace moe_lora_merged {
|
||||
|
||||
__device__ __forceinline__ int warp_exclusive_scan(int v, unsigned mask = 0xffffffffu) {
|
||||
int original = v;
|
||||
#pragma unroll
|
||||
for (int offset = 1; offset < WARP_SIZE; offset <<= 1) {
|
||||
int n = __shfl_up_sync(mask, v, offset);
|
||||
if ((threadIdx.x & (WARP_SIZE - 1)) >= offset) v += n;
|
||||
}
|
||||
return v - original;
|
||||
}
|
||||
|
||||
// Inline mirror of _fused_virtual_topk_ids_kernel (virtual_experts.py). Returns
|
||||
// the merged virtual expert id for flat slot `i` (range [-1, virtual_num_experts);
|
||||
// -1 is the dropped/masked sentinel). The caller adds +1 to get the histogram
|
||||
// bucket (sentinel -> bucket 0), matching the native +1 offset convention.
|
||||
template <typename scalar_t>
|
||||
__device__ __forceinline__ int compute_virtual_id(
|
||||
const scalar_t* __restrict__ topk_ids,
|
||||
const int32_t* __restrict__ token_lora_mapping,
|
||||
size_t i,
|
||||
int top_k,
|
||||
int num_experts_for_weight,
|
||||
int local_expert_offset,
|
||||
int local_num_experts,
|
||||
bool ep_local,
|
||||
bool shared_outer,
|
||||
bool compact) {
|
||||
int m = static_cast<int>(i) / top_k;
|
||||
int lora_id = token_lora_mapping[m];
|
||||
bool mask_val = lora_id >= 0;
|
||||
int safe_lora = lora_id > 0 ? lora_id : 0;
|
||||
|
||||
int base = shared_outer ? 0 : static_cast<int>(topk_ids[i]);
|
||||
if (ep_local) {
|
||||
bool owned = base >= local_expert_offset && base < local_expert_offset + local_num_experts;
|
||||
base = owned ? base : -1;
|
||||
}
|
||||
if (!mask_val || base < 0) return -1;
|
||||
// compact: dense LOCAL expert id in [0, local_num_experts) so the histogram
|
||||
// spans only local_num_experts buckets instead of the full global virtual
|
||||
// space (337/385 empty under EP). Assumes max_loras==1 (safe_lora shift is 0;
|
||||
// the wrapper guards). expert_ids is converted back to global at write time.
|
||||
if (compact) return base - local_expert_offset;
|
||||
return base + safe_lora * num_experts_for_weight;
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void count_and_sort_expert_tokens_kernel(
|
||||
const scalar_t* __restrict__ topk_ids,
|
||||
const int32_t* __restrict__ token_lora_mapping,
|
||||
int32_t* __restrict__ sorted_token_ids,
|
||||
int32_t* __restrict__ cumsum_buffer,
|
||||
size_t numel,
|
||||
int top_k,
|
||||
int num_experts_for_weight,
|
||||
int local_expert_offset,
|
||||
int local_num_experts,
|
||||
bool ep_local,
|
||||
bool shared_outer,
|
||||
bool do_skip,
|
||||
bool compact) {
|
||||
const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const size_t stride = blockDim.x * gridDim.x;
|
||||
|
||||
for (size_t i = tid; i < numel; i += stride) {
|
||||
int vid = compute_virtual_id<scalar_t>(
|
||||
topk_ids,
|
||||
token_lora_mapping,
|
||||
i,
|
||||
top_k,
|
||||
num_experts_for_weight,
|
||||
local_expert_offset,
|
||||
local_num_experts,
|
||||
ep_local,
|
||||
shared_outer,
|
||||
compact);
|
||||
// EP skip: dropped/masked slots (vid < 0) produce no delta on this rank, so
|
||||
// they never need a slot in sorted_token_ids -> skip the global atomicAdd
|
||||
// (kills the sentinel-bucket-0 contention). When do_skip is off they fall
|
||||
// into bucket 0 (old behavior, kept for the bitwise-equivalence guardrail).
|
||||
if (do_skip && vid < 0) continue;
|
||||
int32_t expert_id = vid + 1;
|
||||
int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
|
||||
sorted_token_ids[rank_post_pad] = i;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void moe_align_block_size_kernel(
|
||||
const scalar_t* __restrict__ topk_ids,
|
||||
const int32_t* __restrict__ token_lora_mapping,
|
||||
bool* __restrict__ token_lora_mask,
|
||||
int32_t* __restrict__ sorted_token_ids,
|
||||
int32_t* __restrict__ expert_ids,
|
||||
int32_t* __restrict__ total_tokens_post_pad,
|
||||
int32_t num_experts,
|
||||
int32_t block_size,
|
||||
size_t numel,
|
||||
int32_t* __restrict__ cumsum,
|
||||
bool pad_sorted_token_ids,
|
||||
const int32_t scan_size,
|
||||
int32_t max_num_tokens_padded,
|
||||
int top_k,
|
||||
int num_experts_for_weight,
|
||||
int local_expert_offset,
|
||||
int local_num_experts,
|
||||
bool ep_local,
|
||||
bool shared_outer,
|
||||
bool do_skip,
|
||||
bool compact) {
|
||||
// Use a separate thread block to populate sorted_token_ids
|
||||
if (blockIdx.x == 1) {
|
||||
if (pad_sorted_token_ids) {
|
||||
Vec fill_vec;
|
||||
fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
|
||||
int32_t total_vecs = (max_num_tokens_padded + VEC_SIZE - 1) / VEC_SIZE;
|
||||
Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
|
||||
for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
|
||||
out_ptr[i] = fill_vec;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
extern __shared__ int32_t smem[];
|
||||
int32_t* shared_counts = smem; // [num_experts]
|
||||
int32_t* prefix = shared_counts + num_experts; // [num_experts + 1]
|
||||
int32_t* scan_buf = prefix + num_experts + 1; // [scan_size]
|
||||
__shared__ int32_t s_total_tokens_post_pad;
|
||||
|
||||
const size_t tid = threadIdx.x;
|
||||
const size_t stride = blockDim.x;
|
||||
|
||||
if (tid < num_experts) {
|
||||
shared_counts[tid] = 0;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (size_t i = tid; i < numel; i += stride) {
|
||||
int vid = compute_virtual_id<scalar_t>(
|
||||
topk_ids,
|
||||
token_lora_mapping,
|
||||
i,
|
||||
top_k,
|
||||
num_experts_for_weight,
|
||||
local_expert_offset,
|
||||
local_num_experts,
|
||||
ep_local,
|
||||
shared_outer,
|
||||
compact);
|
||||
// EP skip: dropped/masked slots don't increment any bucket (sentinel bucket
|
||||
// 0 stays empty), so they never get a block and never reach count_and_sort.
|
||||
if (!do_skip || vid >= 0) {
|
||||
atomicAdd(&shared_counts[vid + 1], 1);
|
||||
}
|
||||
// token_lora_mask[m] = token_lora_mapping[m] >= 0, written once per row.
|
||||
if (static_cast<int>(i) % top_k == 0) {
|
||||
int m = static_cast<int>(i) / top_k;
|
||||
token_lora_mask[m] = token_lora_mapping[m] >= 0;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
int32_t padded_count = 0;
|
||||
if (tid < num_experts) {
|
||||
int32_t count = shared_counts[tid];
|
||||
padded_count = (count + block_size - 1) / block_size * block_size;
|
||||
scan_buf[tid] = padded_count;
|
||||
}
|
||||
|
||||
// Intra warp prefix sum
|
||||
int32_t* warp_sums = scan_buf + scan_size; // [<= 32]
|
||||
const int warp_id = tid / WARP_SIZE;
|
||||
const int lane_id = tid & (WARP_SIZE - 1);
|
||||
const int num_warps_for_scan = (scan_size + WARP_SIZE - 1) / WARP_SIZE;
|
||||
const int warp_sum = warp_exclusive_scan(padded_count) + padded_count;
|
||||
if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = warp_sum;
|
||||
__syncthreads();
|
||||
|
||||
// warp0 accumulate all the block's prefix sum
|
||||
if (tid < WARP_SIZE) {
|
||||
int val = (tid < num_warps_for_scan) ? warp_sums[tid] : 0;
|
||||
int incl = warp_exclusive_scan(val) + val;
|
||||
warp_sums[tid] = incl;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Every thread obtains the whole block's sum
|
||||
if (tid == 0) {
|
||||
prefix[num_experts] = warp_sums[num_warps_for_scan - 1];
|
||||
s_total_tokens_post_pad = prefix[num_experts];
|
||||
*total_tokens_post_pad = s_total_tokens_post_pad;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Fill 0 to scan_buf extended area (tid >= num_expert)
|
||||
if (tid >= num_experts && tid < scan_size) scan_buf[tid] = 0;
|
||||
__syncthreads();
|
||||
|
||||
// Perform 2 level exclusive-prefix-sum to scan_buf
|
||||
int v = (tid < scan_size) ? scan_buf[tid] : 0;
|
||||
int pre = warp_exclusive_scan(v);
|
||||
if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = pre + v;
|
||||
__syncthreads();
|
||||
|
||||
if (warp_id == 0) {
|
||||
int val = (lane_id < num_warps_for_scan) ? warp_sums[lane_id] : 0;
|
||||
warp_sums[lane_id] = warp_exclusive_scan(val);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int offset = warp_sums[warp_id];
|
||||
if (tid < scan_size) scan_buf[tid] = pre + offset;
|
||||
__syncthreads();
|
||||
|
||||
// Write prefix[0..num_experts - 1] and cumsum
|
||||
if (tid < num_experts) prefix[tid] = scan_buf[tid];
|
||||
|
||||
if (tid <= num_experts) {
|
||||
cumsum[tid] = prefix[tid];
|
||||
}
|
||||
// fill expert_ids
|
||||
const int32_t num_blocks = s_total_tokens_post_pad / block_size;
|
||||
for (int32_t i = tid; i < num_blocks; i += stride) {
|
||||
int32_t block_start = i * block_size;
|
||||
int left = 0, right = num_experts;
|
||||
while (left < right) {
|
||||
int mid = (left + right) >> 1;
|
||||
if (prefix[mid] <= block_start) {
|
||||
left = mid + 1;
|
||||
} else {
|
||||
right = mid;
|
||||
}
|
||||
}
|
||||
// compact buckets hold LOCAL expert ids; restore the global id (+offset) so
|
||||
// the downstream GEMM still indexes the global contiguous LoRA weight.
|
||||
expert_ids[i] = left - 2 + (compact ? local_expert_offset : 0);
|
||||
}
|
||||
}
|
||||
|
||||
// Single-block fused variant: does fill + histogram + scan + expert_ids + scatter
|
||||
// in ONE threadblock (one launch), eliminating the separate count_and_sort kernel
|
||||
// AND its redundant re-computation of the virtual id (cached in shared `svids`).
|
||||
// Only valid for small numel (the scatter is single-block); the wrapper routes
|
||||
// large numel (prefill) to the 2-kernel path. do_skip is implied (this is the
|
||||
// decode hot path); dropped slots are simply never scattered.
|
||||
template <typename scalar_t>
|
||||
__global__ void fused_align_scatter_kernel(
|
||||
const scalar_t* __restrict__ topk_ids,
|
||||
const int32_t* __restrict__ token_lora_mapping,
|
||||
bool* __restrict__ token_lora_mask,
|
||||
int32_t* __restrict__ sorted_token_ids,
|
||||
int32_t* __restrict__ expert_ids,
|
||||
int32_t* __restrict__ total_tokens_post_pad,
|
||||
int32_t num_experts,
|
||||
int32_t block_size,
|
||||
size_t numel,
|
||||
int32_t* __restrict__ cumsum,
|
||||
const int32_t scan_size,
|
||||
int32_t max_num_tokens_padded,
|
||||
int top_k,
|
||||
int num_experts_for_weight,
|
||||
int local_expert_offset,
|
||||
int local_num_experts,
|
||||
bool ep_local,
|
||||
bool shared_outer,
|
||||
bool do_skip,
|
||||
bool compact) {
|
||||
extern __shared__ int32_t smem[];
|
||||
int32_t* shared_counts = smem; // [num_experts]
|
||||
int32_t* prefix = shared_counts + num_experts; // [num_experts + 1]
|
||||
int32_t* scan_buf = prefix + num_experts + 1; // [scan_size]
|
||||
int32_t* warp_sums = scan_buf + scan_size; // [WARP_SIZE]
|
||||
int32_t* cursor = warp_sums + WARP_SIZE; // [num_experts] scatter cursor
|
||||
int32_t* svids = cursor + num_experts; // [numel] cached virtual ids
|
||||
__shared__ int32_t s_total_tokens_post_pad;
|
||||
|
||||
const size_t tid = threadIdx.x;
|
||||
const size_t stride = blockDim.x;
|
||||
const int warp_id = tid / WARP_SIZE;
|
||||
const int lane_id = tid & (WARP_SIZE - 1);
|
||||
const int num_warps_for_scan = (scan_size + WARP_SIZE - 1) / WARP_SIZE;
|
||||
|
||||
// Phase 1: fill sorted_token_ids with the `numel` padding sentinel.
|
||||
{
|
||||
Vec fill_vec;
|
||||
fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
|
||||
int32_t total_vecs = (max_num_tokens_padded + VEC_SIZE - 1) / VEC_SIZE;
|
||||
Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
|
||||
for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
|
||||
out_ptr[i] = fill_vec;
|
||||
}
|
||||
}
|
||||
if (tid < num_experts) shared_counts[tid] = 0;
|
||||
__syncthreads();
|
||||
|
||||
// Phase 2: histogram + cache the virtual id per slot + token_lora_mask.
|
||||
for (size_t i = tid; i < numel; i += stride) {
|
||||
int vid = compute_virtual_id<scalar_t>(
|
||||
topk_ids,
|
||||
token_lora_mapping,
|
||||
i,
|
||||
top_k,
|
||||
num_experts_for_weight,
|
||||
local_expert_offset,
|
||||
local_num_experts,
|
||||
ep_local,
|
||||
shared_outer,
|
||||
compact);
|
||||
svids[i] = vid;
|
||||
if (!do_skip || vid >= 0) {
|
||||
atomicAdd(&shared_counts[vid + 1], 1);
|
||||
}
|
||||
if (static_cast<int>(i) % top_k == 0) {
|
||||
int m = static_cast<int>(i) / top_k;
|
||||
token_lora_mask[m] = token_lora_mapping[m] >= 0;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Phase 3: padded counts + two-level warp exclusive prefix sum (verbatim).
|
||||
int32_t padded_count = 0;
|
||||
if (tid < num_experts) {
|
||||
int32_t count = shared_counts[tid];
|
||||
padded_count = (count + block_size - 1) / block_size * block_size;
|
||||
scan_buf[tid] = padded_count;
|
||||
}
|
||||
const int warp_sum = warp_exclusive_scan(padded_count) + padded_count;
|
||||
if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = warp_sum;
|
||||
__syncthreads();
|
||||
if (tid < WARP_SIZE) {
|
||||
int val = (tid < num_warps_for_scan) ? warp_sums[tid] : 0;
|
||||
int incl = warp_exclusive_scan(val) + val;
|
||||
warp_sums[tid] = incl;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
prefix[num_experts] = warp_sums[num_warps_for_scan - 1];
|
||||
s_total_tokens_post_pad = prefix[num_experts];
|
||||
*total_tokens_post_pad = s_total_tokens_post_pad;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid >= num_experts && tid < scan_size) scan_buf[tid] = 0;
|
||||
__syncthreads();
|
||||
int v = (tid < scan_size) ? scan_buf[tid] : 0;
|
||||
int pre = warp_exclusive_scan(v);
|
||||
if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = pre + v;
|
||||
__syncthreads();
|
||||
if (warp_id == 0) {
|
||||
int val = (lane_id < num_warps_for_scan) ? warp_sums[lane_id] : 0;
|
||||
warp_sums[lane_id] = warp_exclusive_scan(val);
|
||||
}
|
||||
__syncthreads();
|
||||
int off = warp_sums[warp_id];
|
||||
if (tid < scan_size) scan_buf[tid] = pre + off;
|
||||
__syncthreads();
|
||||
if (tid < num_experts) prefix[tid] = scan_buf[tid];
|
||||
if (tid <= num_experts) cumsum[tid] = prefix[tid];
|
||||
__syncthreads();
|
||||
|
||||
// Phase 4: expert_ids (binary search per block) + init the scatter cursor.
|
||||
const int32_t num_blocks = s_total_tokens_post_pad / block_size;
|
||||
for (int32_t i = tid; i < num_blocks; i += stride) {
|
||||
int32_t block_start = i * block_size;
|
||||
int left = 0, right = num_experts;
|
||||
while (left < right) {
|
||||
int mid = (left + right) >> 1;
|
||||
if (prefix[mid] <= block_start) {
|
||||
left = mid + 1;
|
||||
} else {
|
||||
right = mid;
|
||||
}
|
||||
}
|
||||
expert_ids[i] = left - 2 + (compact ? local_expert_offset : 0);
|
||||
}
|
||||
if (tid < num_experts) cursor[tid] = prefix[tid];
|
||||
__syncthreads();
|
||||
|
||||
// Phase 5: scatter owned tokens using the cached virtual ids + shared cursor.
|
||||
for (size_t i = tid; i < numel; i += stride) {
|
||||
int vid = svids[i];
|
||||
if (do_skip && vid < 0) continue;
|
||||
int bucket = vid + 1;
|
||||
int pos = atomicAdd(&cursor[bucket], 1);
|
||||
sorted_token_ids[pos] = i;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace moe_lora_merged
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename scalar_t>
|
||||
struct MoeLoraMergedAlignKernel {
|
||||
static void
|
||||
run(tvm::ffi::TensorView topk_ids,
|
||||
tvm::ffi::TensorView token_lora_mapping,
|
||||
tvm::ffi::TensorView token_lora_mask,
|
||||
int64_t num_experts,
|
||||
int64_t block_size,
|
||||
tvm::ffi::TensorView sorted_token_ids,
|
||||
tvm::ffi::TensorView expert_ids,
|
||||
tvm::ffi::TensorView num_tokens_post_pad,
|
||||
tvm::ffi::TensorView cumsum_buffer,
|
||||
bool pad_sorted_token_ids,
|
||||
int64_t top_k,
|
||||
int64_t num_experts_for_weight,
|
||||
int64_t local_expert_offset,
|
||||
int64_t local_num_experts,
|
||||
bool ep_local,
|
||||
bool shared_outer,
|
||||
bool do_skip,
|
||||
bool compact,
|
||||
bool fuse_scatter) {
|
||||
using namespace host;
|
||||
|
||||
auto device = topk_ids.device();
|
||||
const cudaStream_t stream = LaunchKernel::resolve_device(device);
|
||||
|
||||
int threads = 1024;
|
||||
threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
|
||||
|
||||
int64_t max_num_tokens_padded = sorted_token_ids.size(0);
|
||||
|
||||
// num_experts here is the bucket count. Non-compact: virtual_num_experts+1
|
||||
// (typically 385). Compact: local_num_experts+1 (typically 49). Both use the
|
||||
// same single-block align path (valid for any bucket count <= 1024 that fits
|
||||
// shared memory); the v2 (>1024) regime keeps the old path via the wrapper.
|
||||
RuntimeCheck(
|
||||
num_experts <= 1024, "moe_lora_merged_align: num_experts (bucket count) must be <= 1024, got ", num_experts);
|
||||
// compact buckets hold LOCAL ids and restore global expert ids as
|
||||
// (left-2+offset). For the sentinel bucket 0 that yields (offset-1), NOT the
|
||||
// -1 the GEMM expects to skip -- only safe when do_skip empties bucket 0.
|
||||
RuntimeCheck(
|
||||
!compact || do_skip, "moe_lora_merged_align: compact requires do_skip (sentinel bucket must be empty)");
|
||||
|
||||
const scalar_t* topk_ids_ptr = static_cast<const scalar_t*>(topk_ids.data_ptr());
|
||||
const int32_t* tlm_ptr = static_cast<const int32_t*>(token_lora_mapping.data_ptr());
|
||||
bool* token_lora_mask_ptr = static_cast<bool*>(token_lora_mask.data_ptr());
|
||||
int32_t* sorted_token_ids_ptr = static_cast<int32_t*>(sorted_token_ids.data_ptr());
|
||||
int32_t* expert_ids_ptr = static_cast<int32_t*>(expert_ids.data_ptr());
|
||||
int32_t* num_tokens_post_pad_ptr = static_cast<int32_t*>(num_tokens_post_pad.data_ptr());
|
||||
int32_t* cumsum_buffer_ptr = static_cast<int32_t*>(cumsum_buffer.data_ptr());
|
||||
size_t numel = topk_ids.numel();
|
||||
|
||||
const size_t scan_size = next_pow2(num_experts);
|
||||
|
||||
if (fuse_scatter) {
|
||||
// One block does fill + histogram + scan + expert_ids + scatter. Extra
|
||||
// shared for the scatter cursor [num_experts] and cached virtual ids [numel].
|
||||
const size_t shmem =
|
||||
(num_experts + (num_experts + 1) + scan_size + WARP_SIZE + num_experts + numel) * sizeof(int32_t);
|
||||
auto fused = moe_lora_merged::fused_align_scatter_kernel<scalar_t>;
|
||||
LaunchKernel(dim3(1), dim3(threads), stream, shmem)(
|
||||
fused,
|
||||
topk_ids_ptr,
|
||||
tlm_ptr,
|
||||
token_lora_mask_ptr,
|
||||
sorted_token_ids_ptr,
|
||||
expert_ids_ptr,
|
||||
num_tokens_post_pad_ptr,
|
||||
(int32_t)num_experts,
|
||||
(int32_t)block_size,
|
||||
numel,
|
||||
cumsum_buffer_ptr,
|
||||
(int32_t)scan_size,
|
||||
(int32_t)max_num_tokens_padded,
|
||||
(int)top_k,
|
||||
(int)num_experts_for_weight,
|
||||
(int)local_expert_offset,
|
||||
(int)local_num_experts,
|
||||
ep_local,
|
||||
shared_outer,
|
||||
do_skip,
|
||||
compact);
|
||||
return;
|
||||
}
|
||||
|
||||
const size_t shared_mem_size = (num_experts + (num_experts + 1) + scan_size + WARP_SIZE) * sizeof(int32_t);
|
||||
|
||||
auto align_kernel = moe_lora_merged::moe_align_block_size_kernel<scalar_t>;
|
||||
LaunchKernel(dim3(2), dim3(threads), stream, shared_mem_size)(
|
||||
align_kernel,
|
||||
topk_ids_ptr,
|
||||
tlm_ptr,
|
||||
token_lora_mask_ptr,
|
||||
sorted_token_ids_ptr,
|
||||
expert_ids_ptr,
|
||||
num_tokens_post_pad_ptr,
|
||||
(int32_t)num_experts,
|
||||
(int32_t)block_size,
|
||||
numel,
|
||||
cumsum_buffer_ptr,
|
||||
pad_sorted_token_ids,
|
||||
(int32_t)scan_size,
|
||||
(int32_t)max_num_tokens_padded,
|
||||
(int)top_k,
|
||||
(int)num_experts_for_weight,
|
||||
(int)local_expert_offset,
|
||||
(int)local_num_experts,
|
||||
ep_local,
|
||||
shared_outer,
|
||||
do_skip,
|
||||
compact);
|
||||
|
||||
const int block_threads = std::min(256, threads);
|
||||
const int num_blocks = (numel + block_threads - 1) / block_threads;
|
||||
const int max_blocks = 65535;
|
||||
const int actual_blocks = std::min(num_blocks, max_blocks);
|
||||
|
||||
auto sort_kernel = moe_lora_merged::count_and_sort_expert_tokens_kernel<scalar_t>;
|
||||
LaunchKernel(dim3(actual_blocks), dim3(block_threads), stream)(
|
||||
sort_kernel,
|
||||
topk_ids_ptr,
|
||||
tlm_ptr,
|
||||
sorted_token_ids_ptr,
|
||||
cumsum_buffer_ptr,
|
||||
numel,
|
||||
(int)top_k,
|
||||
(int)num_experts_for_weight,
|
||||
(int)local_expert_offset,
|
||||
(int)local_num_experts,
|
||||
ep_local,
|
||||
shared_outer,
|
||||
do_skip,
|
||||
compact);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -0,0 +1,412 @@
|
||||
/*
|
||||
* Fused top-k gating softmax WITH routed-pack output.
|
||||
*
|
||||
* JIT port of the power-of-2 fast path of sgl-kernel's
|
||||
* csrc/moe/moe_topk_softmax_kernels.cu (`topkGatingSoftmax`, itself adapted from
|
||||
* vLLM v0.7.3 / TensorRT-LLM v0.7.1, Apache-2.0), extended with a third output:
|
||||
* the FlashInfer routed-MoE packed format
|
||||
*
|
||||
* packed[idx] = (topk_id << 16) | bf16_bits(topk_weight)
|
||||
*
|
||||
* computed in the kernel epilogue AFTER renormalization — bit-identical to the
|
||||
* standalone `fused_pack_topk` triton kernel applied to the (post-processed)
|
||||
* topk_ids/topk_weights, including the padded-region mask: rows at or beyond
|
||||
* `num_token_non_padded` pack id = -1 (the `_mask_topk_ids_padded_region`
|
||||
* sentinel), matching what the separate pack would produce after the mask.
|
||||
* This removes the per-MoE-layer `_pack_topk_kernel` launch from the decode
|
||||
* critical path entirely (fusion instead of stream overlap).
|
||||
*
|
||||
* Scope intentionally narrowed vs the AOT kernel (callers fall back to the AOT
|
||||
* topk_softmax + separate pack otherwise):
|
||||
* - power-of-2 num_experts in [1, 512] only (no cub workspace fallback)
|
||||
* - no softcapping / correction bias (the Qwen3-MoE softmax path uses neither)
|
||||
*/
|
||||
#include <sgl_kernel/tensor.h> // TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.h> // RuntimeCheck
|
||||
|
||||
#include <sgl_kernel/utils.cuh> // LaunchKernel, fp32_t/fp16_t/bf16_t, is_type
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
static constexpr int WARP_SIZE = 32;
|
||||
|
||||
#define TSP_MAX(a, b) ((a) > (b) ? (a) : (b))
|
||||
#define TSP_MIN(a, b) ((a) < (b) ? (a) : (b))
|
||||
|
||||
/// Aligned array type (mirrors the AOT kernel's CUTLASS-free aligned array)
|
||||
template <typename T, int N, int Alignment = sizeof(T) * N>
|
||||
class alignas(Alignment) AlignedArray {
|
||||
T data[N];
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__device__ float convert_to_float(T x) {
|
||||
if constexpr (std::is_same_v<T, __half>) {
|
||||
return __half2float(x);
|
||||
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
|
||||
return __bfloat162float(x);
|
||||
} else if constexpr (std::is_same_v<T, float>) {
|
||||
return x;
|
||||
} else {
|
||||
return static_cast<float>(x);
|
||||
}
|
||||
}
|
||||
|
||||
// Reference pack (bit-identical to jit_kernel/flashinfer_trtllm_moe/topk_pack.py):
|
||||
// low 16 bits = bf16(weight) bits (round-to-nearest-even, same as torch/triton
|
||||
// `.to(bfloat16)`), high 16 bits = int16 expert id.
|
||||
__device__ __forceinline__ int32_t pack_routed(int32_t id, float w) {
|
||||
const uint32_t wbits = static_cast<uint32_t>(__bfloat16_as_ushort(__float2bfloat16(w)));
|
||||
return static_cast<int32_t>((static_cast<uint32_t>(id) << 16) | wbits);
|
||||
}
|
||||
|
||||
template <typename T, int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG>
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmaxPack(
|
||||
const T* input,
|
||||
float* output,
|
||||
const int num_rows,
|
||||
int* indices,
|
||||
int* packed_output,
|
||||
const int32_t* num_token_non_padded,
|
||||
const int k,
|
||||
const bool renormalize) {
|
||||
static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
|
||||
static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS), "NUM_EXPERTS must be power of 2");
|
||||
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2");
|
||||
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
|
||||
|
||||
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
|
||||
static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
|
||||
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
|
||||
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
|
||||
|
||||
static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
|
||||
static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
|
||||
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
|
||||
static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size");
|
||||
|
||||
static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
|
||||
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
|
||||
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
|
||||
|
||||
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp");
|
||||
|
||||
const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
|
||||
const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
|
||||
const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
|
||||
const int thread_row = warp_base_row + thread_row_in_warp;
|
||||
if (thread_row >= num_rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
|
||||
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
|
||||
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
|
||||
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
|
||||
|
||||
using AccessType = AlignedArray<T, ELTS_PER_LDG>;
|
||||
|
||||
T row_chunk_temp[VPT];
|
||||
AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk_temp);
|
||||
const AccessType* vec_thread_read_ptr = reinterpret_cast<const AccessType*>(thread_read_ptr);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
|
||||
}
|
||||
|
||||
float row_chunk[VPT];
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
row_chunk[ii] = convert_to_float<T>(row_chunk_temp[ii]);
|
||||
}
|
||||
|
||||
float thread_max = row_chunk[0];
|
||||
#pragma unroll
|
||||
for (int ii = 1; ii < VPT; ++ii) {
|
||||
thread_max = max(thread_max, row_chunk[ii]);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
|
||||
thread_max = max(thread_max, __shfl_xor_sync(0xffffffffu, thread_max, mask, THREADS_PER_ROW));
|
||||
}
|
||||
|
||||
float row_sum = 0;
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
row_chunk[ii] = expf(row_chunk[ii] - thread_max);
|
||||
row_sum += row_chunk[ii];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
|
||||
row_sum += __shfl_xor_sync(0xffffffffu, row_sum, mask, THREADS_PER_ROW);
|
||||
}
|
||||
|
||||
const float reciprocal_row_sum = 1.f / row_sum;
|
||||
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
|
||||
}
|
||||
|
||||
int start_col = first_elt_read_by_thread;
|
||||
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
|
||||
|
||||
float row_sum_for_renormalize = 0;
|
||||
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
float max_val = row_chunk[0];
|
||||
int expert = start_col;
|
||||
#pragma unroll
|
||||
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
|
||||
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
|
||||
if (val > max_val) {
|
||||
max_val = val;
|
||||
expert = col + ii;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
|
||||
float other_max = __shfl_xor_sync(0xffffffffu, max_val, mask, THREADS_PER_ROW);
|
||||
int other_expert = __shfl_xor_sync(0xffffffffu, expert, mask, THREADS_PER_ROW);
|
||||
if (other_max > max_val || (other_max == max_val && other_expert < expert)) {
|
||||
max_val = other_max;
|
||||
expert = other_expert;
|
||||
}
|
||||
}
|
||||
|
||||
if (thread_group_idx == 0) {
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = max_val;
|
||||
indices[idx] = expert;
|
||||
row_sum_for_renormalize += max_val;
|
||||
}
|
||||
|
||||
if (k_idx + 1 < k) {
|
||||
const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
|
||||
const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
|
||||
if (thread_group_idx == thread_to_clear_in_group) {
|
||||
const int offset_for_expert = expert % ELTS_PER_LDG;
|
||||
row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (thread_group_idx == 0) {
|
||||
// Fused renormalization (same as the AOT kernel).
|
||||
if (renormalize) {
|
||||
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = output[idx] * row_sum_for_renormalize_inv;
|
||||
}
|
||||
}
|
||||
// Fused routed pack: pack the FINAL (post-renorm) weights. Padded rows
|
||||
// (>= *num_token_non_padded) pack id = -1, mirroring the in-place
|
||||
// `_mask_topk_ids_padded_region` sentinel that the separate pack kernel
|
||||
// would otherwise observe. The plain `indices` output is left unmasked
|
||||
// here exactly like the AOT kernel — the existing python post-process
|
||||
// masks it afterwards; only the packed tensor needs the mask baked in
|
||||
// because it is produced BEFORE that post-process runs.
|
||||
const bool row_padded = (num_token_non_padded != nullptr) && (thread_row >= *num_token_non_padded);
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int idx = k * thread_row + k_idx;
|
||||
const int32_t id = row_padded ? -1 : indices[idx];
|
||||
packed_output[idx] = pack_routed(id, output[idx]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
namespace detail {
|
||||
template <typename T, int EXPERTS, int BYTES_PER_LDG>
|
||||
struct TopkConstants {
|
||||
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
|
||||
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, "");
|
||||
static constexpr int VECs_PER_THREAD = TSP_MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
|
||||
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
|
||||
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
|
||||
static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
|
||||
};
|
||||
} // namespace detail
|
||||
|
||||
template <typename T, int EXPERTS, int WARPS_PER_TB>
|
||||
void launchTopkGatingSoftmaxPack(
|
||||
const T* input,
|
||||
float* output,
|
||||
int* indices,
|
||||
int* packed_output,
|
||||
const int32_t* num_token_non_padded,
|
||||
const int num_rows,
|
||||
const int k,
|
||||
const bool renormalize,
|
||||
DLDevice device) {
|
||||
static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
|
||||
static constexpr int BYTES_PER_LDG = TSP_MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
|
||||
using Constants = detail::TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
|
||||
static constexpr int VPT = Constants::VPT;
|
||||
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
|
||||
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
|
||||
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
|
||||
|
||||
dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
|
||||
host::LaunchKernel(dim3(num_blocks), block_dim, device)(
|
||||
topkGatingSoftmaxPack<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>,
|
||||
input,
|
||||
output,
|
||||
num_rows,
|
||||
indices,
|
||||
packed_output,
|
||||
num_token_non_padded,
|
||||
k,
|
||||
renormalize);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void dispatchExperts(
|
||||
const T* input,
|
||||
float* output,
|
||||
int* indices,
|
||||
int* packed_output,
|
||||
const int32_t* num_token_non_padded,
|
||||
const int num_rows,
|
||||
const int num_experts,
|
||||
const int k,
|
||||
const bool renormalize,
|
||||
DLDevice device) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
#define TSP_LAUNCH(E) \
|
||||
launchTopkGatingSoftmaxPack<T, E, WARPS_PER_TB>( \
|
||||
input, output, indices, packed_output, num_token_non_padded, num_rows, k, renormalize, device)
|
||||
switch (num_experts) {
|
||||
case 1:
|
||||
TSP_LAUNCH(1);
|
||||
break;
|
||||
case 2:
|
||||
TSP_LAUNCH(2);
|
||||
break;
|
||||
case 4:
|
||||
TSP_LAUNCH(4);
|
||||
break;
|
||||
case 8:
|
||||
TSP_LAUNCH(8);
|
||||
break;
|
||||
case 16:
|
||||
TSP_LAUNCH(16);
|
||||
break;
|
||||
case 32:
|
||||
TSP_LAUNCH(32);
|
||||
break;
|
||||
case 64:
|
||||
TSP_LAUNCH(64);
|
||||
break;
|
||||
case 128:
|
||||
TSP_LAUNCH(128);
|
||||
break;
|
||||
case 256:
|
||||
TSP_LAUNCH(256);
|
||||
break;
|
||||
case 512:
|
||||
TSP_LAUNCH(512);
|
||||
break;
|
||||
default:
|
||||
host::RuntimeCheck(false, "topk_softmax_pack: num_experts must be a power of 2 in [1, 512], got ", num_experts);
|
||||
}
|
||||
#undef TSP_LAUNCH
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Launcher
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
void topk_softmax_pack(
|
||||
tvm::ffi::TensorView topk_weights,
|
||||
tvm::ffi::TensorView topk_indices,
|
||||
tvm::ffi::TensorView packed,
|
||||
tvm::ffi::TensorView gating_output,
|
||||
tvm::ffi::Optional<tvm::ffi::TensorView> num_token_non_padded,
|
||||
bool renormalize) {
|
||||
using namespace host;
|
||||
|
||||
SymbolicSize N{"num_tokens"};
|
||||
SymbolicSize E{"num_experts"};
|
||||
SymbolicSize K{"topk"};
|
||||
SymbolicDevice device_;
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, E}).with_dtype<fp32_t, fp16_t, bf16_t>().with_device<kDLCUDA>(device_).verify(gating_output);
|
||||
TensorMatcher({N, K}).with_dtype<fp32_t>().with_device<kDLCUDA>(device_).verify(topk_weights);
|
||||
TensorMatcher({N, K}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(topk_indices);
|
||||
TensorMatcher({N, K}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(packed);
|
||||
|
||||
const int32_t* ntnp_ptr = nullptr;
|
||||
if (num_token_non_padded.has_value()) {
|
||||
SymbolicSize One{"ntnp_numel"};
|
||||
TensorMatcher({One}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(num_token_non_padded.value());
|
||||
RuntimeCheck(One.unwrap() == 1, "num_token_non_padded must be a 1-element tensor");
|
||||
ntnp_ptr = static_cast<const int32_t*>(num_token_non_padded.value().data_ptr());
|
||||
}
|
||||
|
||||
const int num_tokens = static_cast<int>(N.unwrap());
|
||||
const int num_experts = static_cast<int>(E.unwrap());
|
||||
const int topk = static_cast<int>(K.unwrap());
|
||||
DLDevice device = device_.unwrap();
|
||||
|
||||
RuntimeCheck(topk <= num_experts, "topk must be <= num_experts");
|
||||
if (num_tokens == 0) return;
|
||||
|
||||
auto* weights_ptr = static_cast<float*>(topk_weights.data_ptr());
|
||||
auto* indices_ptr = static_cast<int*>(topk_indices.data_ptr());
|
||||
auto* packed_ptr = static_cast<int*>(packed.data_ptr());
|
||||
|
||||
if (is_type<fp32_t>(gating_output.dtype())) {
|
||||
dispatchExperts<float>(
|
||||
static_cast<const float*>(gating_output.data_ptr()),
|
||||
weights_ptr,
|
||||
indices_ptr,
|
||||
packed_ptr,
|
||||
ntnp_ptr,
|
||||
num_tokens,
|
||||
num_experts,
|
||||
topk,
|
||||
renormalize,
|
||||
device);
|
||||
} else if (is_type<fp16_t>(gating_output.dtype())) {
|
||||
dispatchExperts<__half>(
|
||||
static_cast<const __half*>(gating_output.data_ptr()),
|
||||
weights_ptr,
|
||||
indices_ptr,
|
||||
packed_ptr,
|
||||
ntnp_ptr,
|
||||
num_tokens,
|
||||
num_experts,
|
||||
topk,
|
||||
renormalize,
|
||||
device);
|
||||
} else {
|
||||
dispatchExperts<__nv_bfloat16>(
|
||||
static_cast<const __nv_bfloat16*>(gating_output.data_ptr()),
|
||||
weights_ptr,
|
||||
indices_ptr,
|
||||
packed_ptr,
|
||||
ntnp_ptr,
|
||||
num_tokens,
|
||||
num_experts,
|
||||
topk,
|
||||
renormalize,
|
||||
device);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
@@ -0,0 +1,20 @@
|
||||
FlashInfer TRTLLM MoE Overlay
|
||||
=============================
|
||||
|
||||
This directory contains only the editable overlay files used by SGLang's
|
||||
`experimental_sgl_trtllm` SM100 TRTLLM fused MoE backend. Unmodified FlashInfer
|
||||
and TRTLLM sources are compiled from the installed `flashinfer` package at JIT
|
||||
time.
|
||||
|
||||
Local overlay source:
|
||||
|
||||
- `data/csrc/trtllm_fused_moe_kernel_launcher.cu`
|
||||
- `data/csrc/trtllm_fused_moe_runner.cu`
|
||||
- `data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu`
|
||||
- `data/include/flashinfer/trtllm/fused_moe/DevKernel.h`
|
||||
- `data/include/flashinfer/trtllm/fused_moe/runner.h`
|
||||
|
||||
The backend still depends on the installed `flashinfer` and `flashinfer_cubin`
|
||||
packages for the rest of the FlashInfer/TRTLLM JIT source tree and TRTLLM-Gen
|
||||
BMM cubin artifacts. The local include directory is passed before FlashInfer's
|
||||
installed include directory so these overlay headers shadow the originals.
|
||||
@@ -0,0 +1,17 @@
|
||||
from sglang.jit_kernel.trtllm_lora_temp.core import (
|
||||
trtllm_fp4_block_scale_moe_lora_finalize,
|
||||
trtllm_fp4_block_scale_routed_moe_lora,
|
||||
trtllm_fp8_block_scale_moe,
|
||||
trtllm_fp8_block_scale_moe_lora_finalize,
|
||||
trtllm_fp8_block_scale_routed_moe,
|
||||
trtllm_fp8_block_scale_routed_moe_lora,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"trtllm_fp4_block_scale_moe_lora_finalize",
|
||||
"trtllm_fp4_block_scale_routed_moe_lora",
|
||||
"trtllm_fp8_block_scale_moe_lora_finalize",
|
||||
"trtllm_fp8_block_scale_moe",
|
||||
"trtllm_fp8_block_scale_routed_moe",
|
||||
"trtllm_fp8_block_scale_routed_moe_lora",
|
||||
]
|
||||
@@ -0,0 +1,443 @@
|
||||
import functools
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_sgl_trtllm_moe_sm100_module():
|
||||
import flashinfer.fused_moe.core as fi_core
|
||||
|
||||
from sglang.jit_kernel.trtllm_lora_temp.jit import (
|
||||
gen_sgl_trtllm_gen_fused_moe_sm100_module,
|
||||
)
|
||||
|
||||
original_gen = fi_core.gen_trtllm_gen_fused_moe_sm100_module
|
||||
fi_core.gen_trtllm_gen_fused_moe_sm100_module = (
|
||||
gen_sgl_trtllm_gen_fused_moe_sm100_module
|
||||
)
|
||||
try:
|
||||
fi_core.get_trtllm_moe_sm100_module.cache_clear()
|
||||
return fi_core.get_trtllm_moe_sm100_module()
|
||||
finally:
|
||||
fi_core.gen_trtllm_gen_fused_moe_sm100_module = original_gen
|
||||
fi_core.get_trtllm_moe_sm100_module.cache_clear()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_sgl_trtllm_moe_sm100_raw_module():
|
||||
from flashinfer.fused_moe.core import setup_cubin_loader
|
||||
|
||||
from sglang.jit_kernel.trtllm_lora_temp.jit import (
|
||||
gen_sgl_trtllm_gen_fused_moe_sm100_module,
|
||||
)
|
||||
|
||||
module = gen_sgl_trtllm_gen_fused_moe_sm100_module()
|
||||
moe_op = module.build_and_load()
|
||||
setup_cubin_loader(str(module.get_library_path()))
|
||||
return moe_op
|
||||
|
||||
|
||||
def _validate_routing_replay_out(
|
||||
routing_replay_out: Optional[torch.Tensor], top_k: int
|
||||
) -> None:
|
||||
if routing_replay_out is None:
|
||||
return
|
||||
assert routing_replay_out.dim() == 2
|
||||
assert routing_replay_out.shape[1] == top_k
|
||||
assert routing_replay_out.dtype == torch.int16
|
||||
assert routing_replay_out.is_cuda
|
||||
assert routing_replay_out.is_contiguous()
|
||||
|
||||
|
||||
def trtllm_fp8_block_scale_moe(
|
||||
routing_logits: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
n_group: Optional[int],
|
||||
topk_group: Optional[int],
|
||||
intermediate_size: int,
|
||||
local_expert_offset: int,
|
||||
local_num_experts: int,
|
||||
routed_scaling_factor: Optional[float],
|
||||
routing_method_type: int = 0,
|
||||
use_shuffled_weight: bool = False,
|
||||
weight_layout: int = 0,
|
||||
do_finalize: bool = True,
|
||||
enable_pdl: Optional[bool] = None,
|
||||
tune_max_num_tokens: int = 8192,
|
||||
fp8_quantization_type=None,
|
||||
activation_type: Optional[int] = None,
|
||||
norm_topk_prob: bool = True,
|
||||
routing_replay_out: Optional[torch.Tensor] = None,
|
||||
) -> Union[List[torch.Tensor], torch.Tensor]:
|
||||
from flashinfer.fused_moe.core import ActivationType, Fp8QuantizationType
|
||||
|
||||
_validate_routing_replay_out(routing_replay_out, top_k)
|
||||
|
||||
if fp8_quantization_type is None:
|
||||
fp8_quantization_type = Fp8QuantizationType.DeepSeekFp8
|
||||
if activation_type is None:
|
||||
activation_type = ActivationType.Swiglu.value
|
||||
|
||||
output = torch.empty(
|
||||
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
|
||||
)
|
||||
result = get_sgl_trtllm_moe_sm100_module().trtllm_fp8_block_scale_moe(
|
||||
routing_logits,
|
||||
None,
|
||||
None,
|
||||
routing_bias,
|
||||
hidden_states,
|
||||
hidden_states_scale,
|
||||
gemm1_weights,
|
||||
gemm1_weights_scale,
|
||||
gemm2_weights,
|
||||
gemm2_weights_scale,
|
||||
output,
|
||||
num_experts,
|
||||
top_k,
|
||||
n_group,
|
||||
topk_group,
|
||||
intermediate_size,
|
||||
local_expert_offset,
|
||||
local_num_experts,
|
||||
routed_scaling_factor,
|
||||
routing_method_type,
|
||||
use_shuffled_weight,
|
||||
weight_layout,
|
||||
do_finalize,
|
||||
enable_pdl,
|
||||
tune_max_num_tokens,
|
||||
fp8_quantization_type,
|
||||
activation_type,
|
||||
norm_topk_prob,
|
||||
routing_replay_out,
|
||||
)
|
||||
|
||||
return result[0] if do_finalize else result
|
||||
|
||||
|
||||
def trtllm_fp8_block_scale_routed_moe(
|
||||
topk_ids: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
n_group: Optional[int],
|
||||
topk_group: Optional[int],
|
||||
intermediate_size: int,
|
||||
local_expert_offset: int,
|
||||
local_num_experts: int,
|
||||
routed_scaling_factor: Optional[float],
|
||||
routing_method_type: int = 0,
|
||||
use_shuffled_weight: bool = False,
|
||||
weight_layout: int = 0,
|
||||
do_finalize: bool = True,
|
||||
enable_pdl: Optional[bool] = None,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
tune_max_num_tokens: int = 8192,
|
||||
fp8_quantization_type=None,
|
||||
activation_type: Optional[int] = None,
|
||||
) -> Union[List[torch.Tensor], torch.Tensor]:
|
||||
from flashinfer.fused_moe.core import ActivationType, Fp8QuantizationType
|
||||
|
||||
if fp8_quantization_type is None:
|
||||
fp8_quantization_type = Fp8QuantizationType.DeepSeekFp8
|
||||
if activation_type is None:
|
||||
activation_type = ActivationType.Swiglu.value
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
|
||||
)
|
||||
|
||||
result = get_sgl_trtllm_moe_sm100_module().trtllm_fp8_block_scale_moe(
|
||||
None,
|
||||
topk_ids,
|
||||
None,
|
||||
routing_bias,
|
||||
hidden_states,
|
||||
hidden_states_scale,
|
||||
gemm1_weights,
|
||||
gemm1_weights_scale,
|
||||
gemm2_weights,
|
||||
gemm2_weights_scale,
|
||||
output,
|
||||
num_experts,
|
||||
top_k,
|
||||
n_group,
|
||||
topk_group,
|
||||
intermediate_size,
|
||||
local_expert_offset,
|
||||
local_num_experts,
|
||||
routed_scaling_factor,
|
||||
routing_method_type,
|
||||
use_shuffled_weight,
|
||||
weight_layout,
|
||||
do_finalize,
|
||||
enable_pdl,
|
||||
tune_max_num_tokens,
|
||||
fp8_quantization_type,
|
||||
activation_type,
|
||||
True,
|
||||
)
|
||||
|
||||
return result[0] if do_finalize else result
|
||||
|
||||
|
||||
def trtllm_fp8_block_scale_routed_moe_lora(
|
||||
topk_ids: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
gate_up_lora_delta: torch.Tensor,
|
||||
activation_lora_input: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
n_group: Optional[int],
|
||||
topk_group: Optional[int],
|
||||
intermediate_size: int,
|
||||
local_expert_offset: int,
|
||||
local_num_experts: int,
|
||||
routed_scaling_factor: Optional[float],
|
||||
routing_method_type: int = 0,
|
||||
use_shuffled_weight: bool = False,
|
||||
weight_layout: int = 0,
|
||||
do_finalize: bool = True,
|
||||
enable_pdl: Optional[bool] = None,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
tune_max_num_tokens: int = 8192,
|
||||
fp8_quantization_type=None,
|
||||
activation_type: Optional[int] = None,
|
||||
lora_ready_event: int = 0,
|
||||
gemm2_done_event: int = 0,
|
||||
) -> Union[List[torch.Tensor], torch.Tensor]:
|
||||
from flashinfer.fused_moe.core import ActivationType, Fp8QuantizationType
|
||||
from flashinfer.utils import device_support_pdl
|
||||
|
||||
if fp8_quantization_type is None:
|
||||
fp8_quantization_type = Fp8QuantizationType.DeepSeekFp8
|
||||
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()
|
||||
empty_expert_weights = torch.empty(
|
||||
(0,), dtype=torch.bfloat16, device=hidden_states.device
|
||||
)
|
||||
|
||||
result = get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp8_block_scale_moe_lora(
|
||||
None,
|
||||
topk_ids,
|
||||
empty_expert_weights,
|
||||
routing_bias,
|
||||
hidden_states,
|
||||
hidden_states_scale,
|
||||
gemm1_weights,
|
||||
gemm1_weights_scale,
|
||||
gemm2_weights,
|
||||
gemm2_weights_scale,
|
||||
output,
|
||||
num_experts,
|
||||
top_k,
|
||||
n_group,
|
||||
topk_group,
|
||||
intermediate_size,
|
||||
local_expert_offset,
|
||||
local_num_experts,
|
||||
routed_scaling_factor,
|
||||
routing_method_type,
|
||||
use_shuffled_weight,
|
||||
weight_layout,
|
||||
do_finalize,
|
||||
enable_pdl,
|
||||
[-1, -1],
|
||||
fp8_quantization_type,
|
||||
activation_type,
|
||||
True,
|
||||
None,
|
||||
gate_up_lora_delta,
|
||||
activation_lora_input,
|
||||
lora_ready_event,
|
||||
gemm2_done_event,
|
||||
)
|
||||
|
||||
return output if do_finalize else result
|
||||
|
||||
|
||||
def trtllm_fp8_block_scale_moe_lora_finalize(
|
||||
gemm2_output: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
expanded_idx_to_permuted_idx: torch.Tensor,
|
||||
down_lora_delta: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
routed_scaling_factor: Optional[float],
|
||||
) -> torch.Tensor:
|
||||
get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp8_block_scale_moe_lora_finalize(
|
||||
gemm2_output,
|
||||
expert_weights,
|
||||
expanded_idx_to_permuted_idx,
|
||||
down_lora_delta,
|
||||
output,
|
||||
routed_scaling_factor,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def trtllm_fp4_block_scale_routed_moe_lora(
|
||||
topk_ids: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: Optional[torch.Tensor],
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
output1_scales_scalar: torch.Tensor,
|
||||
output1_scales_gate_scalar: torch.Tensor,
|
||||
output2_scales_scalar: 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,
|
||||
act_type: Optional[int] = None,
|
||||
norm_topk_prob: bool = True,
|
||||
lora_ready_event: int = 0,
|
||||
gemm2_done_event: int = 0,
|
||||
) -> Union[List[torch.Tensor], torch.Tensor]:
|
||||
"""NVFP4 sibling of :func:`trtllm_fp8_block_scale_routed_moe_lora`.
|
||||
|
||||
Decomposed (unfused-activation) MoE-LoRA on the flashinfer-trtllm backend:
|
||||
routing -> gather -> gate_up grouped GEMM (raw 2*inter) -> activation that
|
||||
adds ``gate_up_lora_delta`` pre-SwiGLU and writes ``activation_lora_input``
|
||||
-> NvFP4 quant -> down grouped GEMM -> finalize. The down-LoRA is merged into
|
||||
the output afterwards by the dispatch (virtual-experts) or via the finalize.
|
||||
|
||||
``hidden_states`` is bf16 ``[num_tokens, hidden]`` (path 3: the op permutes then
|
||||
NvFP4-quantizes it internally, ``hidden_states_scale=None``) or, legacy, packed NvFP4
|
||||
(uint8 ``[num_tokens, hidden//2]``) with ``hidden_states_scale`` the fp8-e4m3 block
|
||||
scale ``[num_tokens, hidden//16]``.
|
||||
"""
|
||||
from flashinfer.fused_moe.core import ActivationType
|
||||
from flashinfer.utils import device_support_pdl
|
||||
|
||||
if act_type is None:
|
||||
act_type = ActivationType.Swiglu.value
|
||||
if enable_pdl is None:
|
||||
enable_pdl = device_support_pdl(hidden_states.device)
|
||||
|
||||
num_tokens = hidden_states.shape[0]
|
||||
hidden_size = (
|
||||
hidden_states.shape[1] * 2
|
||||
if hidden_states.dtype == torch.uint8
|
||||
else hidden_states.shape[1]
|
||||
)
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
(num_tokens, hidden_size), dtype=torch.bfloat16, device=hidden_states.device
|
||||
)
|
||||
|
||||
assert gate_up_lora_delta.is_contiguous()
|
||||
assert activation_lora_input.is_contiguous()
|
||||
empty_expert_weights = torch.empty(
|
||||
(0,), dtype=torch.bfloat16, device=hidden_states.device
|
||||
)
|
||||
|
||||
result = get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp4_block_scale_moe_lora(
|
||||
None, # routing_logits (precomputed routing via packed topk_ids)
|
||||
topk_ids, # expert_indices (packed: (expert_id<<16)|weight_bf16.view(int16))
|
||||
empty_expert_weights,
|
||||
routing_bias,
|
||||
hidden_states,
|
||||
hidden_states_scale,
|
||||
gemm1_weights,
|
||||
gemm1_weights_scale,
|
||||
None, # gemm1_bias
|
||||
None, # gemm1_alpha
|
||||
None, # gemm1_beta
|
||||
None, # gemm1_clamp_limit
|
||||
gemm2_weights,
|
||||
gemm2_weights_scale,
|
||||
None, # gemm2_bias
|
||||
output1_scales_scalar,
|
||||
output1_scales_gate_scalar,
|
||||
output2_scales_scalar,
|
||||
None, # per_token_scales (the down-GEMM act scale is produced internally)
|
||||
num_experts,
|
||||
top_k,
|
||||
None, # n_group
|
||||
None, # topk_group
|
||||
intermediate_size,
|
||||
local_expert_offset,
|
||||
local_num_experts,
|
||||
routed_scaling_factor,
|
||||
routing_method_type,
|
||||
do_finalize,
|
||||
enable_pdl,
|
||||
act_type,
|
||||
output,
|
||||
[-1, -1], # config_index (autotuner)
|
||||
norm_topk_prob,
|
||||
None, # routing_replay_out
|
||||
gate_up_lora_delta,
|
||||
activation_lora_input,
|
||||
lora_ready_event,
|
||||
bool(lora_envs.SGLANG_OPT_FUSED_PERMUTE_QUANT.get()),
|
||||
gemm2_done_event,
|
||||
)
|
||||
|
||||
return output if do_finalize else result
|
||||
|
||||
|
||||
def trtllm_fp4_block_scale_moe_lora_finalize(
|
||||
gemm2_output: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
expanded_idx_to_permuted_idx: torch.Tensor,
|
||||
down_lora_delta: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
routed_scaling_factor: Optional[float],
|
||||
) -> torch.Tensor:
|
||||
"""NvFP4 analog of :func:`trtllm_fp8_block_scale_moe_lora_finalize` — combines
|
||||
the permuted (bf16) GEMM2 output by expert weight and merges the down-LoRA
|
||||
delta into the per-token output (non-virtual-experts hook path)."""
|
||||
get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp4_block_scale_moe_lora_finalize(
|
||||
gemm2_output,
|
||||
expert_weights,
|
||||
expanded_idx_to_permuted_idx,
|
||||
down_lora_delta,
|
||||
output,
|
||||
routed_scaling_factor,
|
||||
)
|
||||
return output
|
||||
@@ -0,0 +1,230 @@
|
||||
// Fused SwiGLU+LoRA activation -> NVFP4 per-token quant for the FP4 MoE LoRA path.
|
||||
//
|
||||
// Modeled on tensorrt_llm::kernels::nvfp4QuantAndPerTokenScaleKernel (flashinfer
|
||||
// quantization.cuh): one block per expanded row, per-token amax via cub::BlockReduce,
|
||||
// cvt_warp_fp16_to_fp4 for the e4m3 block scale + e2m1 packing + swizzled SF layout.
|
||||
// The ONLY change vs that kernel is the pass-1 input: instead of reading activated_bf16
|
||||
// from gmem, it reads the interleaved gate/up GEMM1 output + the LoRA delta and computes
|
||||
// silu(up)*gate on the fly, rounds to bf16 (matching the standalone activation kernel's
|
||||
// bf16 output exactly), caches in smem, and also writes activation_lora_input. The down
|
||||
// GEMM input (activated_bf16) is therefore never materialized to HBM.
|
||||
//
|
||||
// Because the activated value is rounded to bf16 before quantization (same as the separate
|
||||
// activation kernel) and silu matches, the fp4 / SF / per_token_sf / activation_lora_input
|
||||
// outputs are BITWISE-identical to the unfused activation -> quant#2 chain.
|
||||
#pragma once
|
||||
|
||||
#include <cub/block/block_reduce.cuh>
|
||||
#include <cuda/functional>
|
||||
|
||||
#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
|
||||
#include <cstdint>
|
||||
#include <cuda_bf16.h>
|
||||
#include <optional>
|
||||
|
||||
namespace flashinfer {
|
||||
namespace sgl_fused_act_quant {
|
||||
|
||||
namespace tk = tensorrt_llm::kernels;
|
||||
|
||||
// Same silu as moe::dev::activation (trtllm_fused_moe_dev_kernel.cu:55): x / (1 + exp(-x)).
|
||||
inline __device__ float fused_silu(float x) {
|
||||
return x / (1.0f + expf(-x));
|
||||
}
|
||||
|
||||
// One block per expanded row. gateUp is the column-interleaved GEMM1 output (g0,u0,g1,u1,...)
|
||||
// indexed by permutedIdx; loraDelta is the contiguous [gate|up] delta indexed by expandedIdx.
|
||||
template <uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT, bool DISABLE_FP4_FAST_MATH>
|
||||
__global__ void fusedActivationQuantKernel(
|
||||
int m, // numTokens * topK (number of expanded rows)
|
||||
int innerHalf, // inter == n (output width per row); must be a multiple of 16
|
||||
int innerDim, // gate_up_n == 2 * innerHalf
|
||||
__nv_bfloat16 const* __restrict__ gateUp, // interleaved gate/up, [.., innerDim] by permutedIdx
|
||||
__nv_bfloat16 const* __restrict__ loraDelta, // [.., innerDim] by expandedIdx, may be null
|
||||
__nv_bfloat16* __restrict__ loraInputOut, // [.., innerHalf] by expandedIdx, may be null
|
||||
int32_t const* __restrict__ expandedIdxToPermutedIdx,
|
||||
float globalScaleInv,
|
||||
uint8_t* __restrict__ weightOutput, // fp4 [.., innerHalf/2] by permutedIdx
|
||||
uint8_t* __restrict__ scaleOutput, // swizzled e4m3 SF
|
||||
float* __restrict__ perTokenScaleOutput) {
|
||||
constexpr int SF_VEC_SIZE = 16;
|
||||
using InType = tk::PackedVec<__nv_bfloat16, SF_VEC_SIZE>; // 16 bf16 == 8 __nv_bfloat162
|
||||
using PackedFp4Type = uint64_t; // SF_VEC_SIZE == 16
|
||||
|
||||
int const expandedIdx = blockIdx.x;
|
||||
if (expandedIdx >= m) return;
|
||||
int const permutedIdx = expandedIdxToPermutedIdx[expandedIdx];
|
||||
int const num_vecs_per_row = innerHalf / SF_VEC_SIZE;
|
||||
int64_t const liBaseRow = (int64_t)expandedIdx * innerHalf;
|
||||
|
||||
// Padding row: the separate activation kernel writes 0 to activation_lora_input and skips
|
||||
// the quant outputs. Mirror that, then return.
|
||||
if (permutedIdx < 0) {
|
||||
if (loraInputOut != nullptr) {
|
||||
InType z;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < SF_VEC_SIZE / 2; ++i)
|
||||
z.elts[i] = __float2bfloat162_rn(0.0f);
|
||||
for (int vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
*reinterpret_cast<InType*>(&loraInputOut[liBaseRow + (int64_t)vecIdx * SF_VEC_SIZE]) = z;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
int64_t const permBase = (int64_t)permutedIdx * innerDim; // gate_up row (interleaved)
|
||||
int64_t const expBase = (int64_t)expandedIdx * innerDim; // delta row (contiguous gate|up)
|
||||
(void)DISABLE_FP4_FAST_MATH;
|
||||
|
||||
// 1 SF block (16 outputs) per thread, held in registers across the amax barrier (no smem cache):
|
||||
// requires num_vecs_per_row <= BLOCK_SIZE (inter=2048 -> 128 == BLOCK_SIZE). With
|
||||
// CVT_ELTS_PER_THREAD == SF_VEC_SIZE the cvt needs no cross-thread shuffle, so masking is safe.
|
||||
int const vecIdx = threadIdx.x;
|
||||
bool const active = vecIdx < num_vecs_per_row;
|
||||
|
||||
InType vec;
|
||||
float localAmax = 0.f;
|
||||
if (active) {
|
||||
int const h0 = vecIdx * SF_VEC_SIZE;
|
||||
__nv_bfloat16 const* g = gateUp + permBase + (int64_t)2 * h0; // 32 interleaved bf16
|
||||
__nv_bfloat16 const* dlo = loraDelta + expBase + h0; // silu-arg delta (lower half)
|
||||
__nv_bfloat16 const* dhi = loraDelta + expBase + innerHalf + h0; // multiplier delta (upper half)
|
||||
__nv_bfloat162 amax2 = __float2bfloat162_rn(0.0f);
|
||||
union {
|
||||
int4 v[4];
|
||||
__nv_bfloat16 b[32];
|
||||
} gu;
|
||||
union {
|
||||
int4 v[2];
|
||||
__nv_bfloat16 b[16];
|
||||
} dl, dh;
|
||||
int4 const* gp = reinterpret_cast<int4 const*>(g);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
gu.v[k] = gp[k];
|
||||
if (loraDelta != nullptr) {
|
||||
int4 const* dlp = reinterpret_cast<int4 const*>(dlo);
|
||||
int4 const* dhp = reinterpret_cast<int4 const*>(dhi);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2; ++k) {
|
||||
dl.v[k] = dlp[k];
|
||||
dh.v[k] = dhp[k];
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < SF_VEC_SIZE / 2; ++i) { // 8 bf162 = 16 output elements
|
||||
int const j0 = 2 * i, j1 = 2 * i + 1;
|
||||
float even0 = (float)gu.b[2 * j0], odd0 = (float)gu.b[2 * j0 + 1];
|
||||
float even1 = (float)gu.b[2 * j1], odd1 = (float)gu.b[2 * j1 + 1];
|
||||
float a0 = odd0, b0 = even0, a1 = odd1, b1 = even1;
|
||||
if (loraDelta != nullptr) {
|
||||
a0 += (float)dl.b[j0];
|
||||
b0 += (float)dh.b[j0];
|
||||
a1 += (float)dl.b[j1];
|
||||
b1 += (float)dh.b[j1];
|
||||
}
|
||||
float act0 = fused_silu(a0) * b0;
|
||||
float act1 = fused_silu(a1) * b1;
|
||||
__nv_bfloat162 e = __float22bfloat162_rn(make_float2(act0, act1));
|
||||
vec.elts[i] = e;
|
||||
amax2 = __hmax2(amax2, __habs2(e));
|
||||
}
|
||||
localAmax = (float)__hmax(amax2.x, amax2.y);
|
||||
if (loraInputOut != nullptr) {
|
||||
*reinterpret_cast<InType*>(&loraInputOut[liBaseRow + h0]) = vec;
|
||||
}
|
||||
}
|
||||
|
||||
// ---- per-token scale: blockReduce amax, broadcast via smem (no gmem round-trip) ----
|
||||
using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
__shared__ float sScale;
|
||||
float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
|
||||
if (threadIdx.x == 0) {
|
||||
float const pts = globalAmax * globalScaleInv;
|
||||
perTokenScaleOutput[permutedIdx] = pts;
|
||||
sScale = pts;
|
||||
}
|
||||
__syncthreads();
|
||||
float const globalEncodeScale = tk::reciprocal_approximate_ftz(sScale);
|
||||
|
||||
// ---- quantize from registers (cvt computes the per-16 e4m3 block scale internally) ----
|
||||
if (active) {
|
||||
uint8_t fp8Scale;
|
||||
// 5 template args on this flashinfer build: Type, SF_VEC_SIZE, CVT_ELTS_PER_THREAD,
|
||||
// UE8M0_SF=false, TE_EXACT_NVFP4=false (the default nvfp4 quant path).
|
||||
auto fp4Vals = tk::cvt_warp_fp16_to_fp4<__nv_bfloat16, SF_VEC_SIZE, SF_VEC_SIZE, false, false>(
|
||||
vec, globalEncodeScale, &fp8Scale);
|
||||
int64_t const vecOffset = (int64_t)permutedIdx * num_vecs_per_row + vecIdx;
|
||||
reinterpret_cast<PackedFp4Type*>(weightOutput)[vecOffset] = fp4Vals;
|
||||
|
||||
// Match nvfp4QuantAndPerTokenScaleKernel exactly (it passes the kernel's `m` as numRows).
|
||||
int64_t sfOffset;
|
||||
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);
|
||||
} else {
|
||||
sfOffset = tk::get_sf_out_offset_8x4(std::nullopt, permutedIdx, vecIdx, m, num_vecs_per_row);
|
||||
}
|
||||
scaleOutput[sfOffset] = fp8Scale;
|
||||
}
|
||||
}
|
||||
|
||||
// Host launch: globalScaleInv = 1/448/6. BLOCK_SIZE must be >= innerHalf/16 (one SF block/thread).
|
||||
inline void launchFusedActivationQuant(
|
||||
int m,
|
||||
int innerHalf,
|
||||
int innerDim,
|
||||
__nv_bfloat16 const* gateUp,
|
||||
__nv_bfloat16 const* loraDelta,
|
||||
__nv_bfloat16* loraInputOut,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
float globalScaleInv,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput,
|
||||
tensorrt_llm::QuantizationSFLayout sfLayout,
|
||||
bool disableFp4FastMath,
|
||||
cudaStream_t stream) {
|
||||
constexpr uint32_t BLOCK_SIZE = 128; // == innerHalf/16 for inter=2048 (one SF block per thread)
|
||||
dim3 const grid(m), block(BLOCK_SIZE);
|
||||
|
||||
auto launch = [&](auto layoutTag, auto fastMathTag) {
|
||||
fusedActivationQuantKernel<BLOCK_SIZE, decltype(layoutTag)::value, decltype(fastMathTag)::value>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m,
|
||||
innerHalf,
|
||||
innerDim,
|
||||
gateUp,
|
||||
loraDelta,
|
||||
loraInputOut,
|
||||
expandedIdxToPermutedIdx,
|
||||
globalScaleInv,
|
||||
weightOutput,
|
||||
scaleOutput,
|
||||
perTokenScaleOutput);
|
||||
};
|
||||
auto withFastMath = [&](auto layoutTag) {
|
||||
if (disableFp4FastMath) {
|
||||
launch(layoutTag, std::integral_constant<bool, true>{});
|
||||
} else {
|
||||
launch(layoutTag, std::integral_constant<bool, false>{});
|
||||
}
|
||||
};
|
||||
if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
|
||||
withFastMath(
|
||||
std::integral_constant<
|
||||
tensorrt_llm::QuantizationSFLayout,
|
||||
tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
|
||||
} else if (sfLayout == tensorrt_llm::QuantizationSFLayout::LINEAR) {
|
||||
withFastMath(
|
||||
std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::LINEAR>{});
|
||||
} else {
|
||||
withFastMath(
|
||||
std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace sgl_fused_act_quant
|
||||
} // namespace flashinfer
|
||||
+1137
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,307 @@
|
||||
// Fused permute + NvFP4-per-token-quant for the FP4 MoE-LoRA gate_up path.
|
||||
//
|
||||
// Background (decode bs64, EP8): the plain path runs `permuteKernel` (gather bf16 hidden into the
|
||||
// padded [max_padded, hidden] permuted buffer) then `nvfp4QuantAndPerTokenScaleKernel` over ALL
|
||||
// max_padded rows. At decode only num_tokens*top_k of those rows are real (the rest are padding),
|
||||
// so both kernels waste ~6x of their work on padding, and the bf16 permuted buffer is a full
|
||||
// HBM round-trip (written by permute, read back by quant).
|
||||
//
|
||||
// This kernel fuses the two: it reads the UN-permuted hidden, NvFP4-quantizes each (token,expert)
|
||||
// pair's row, and scatter-writes fp4 + swizzled block-sf + per-token-sf directly to that pair's
|
||||
// permuted position. It iterates only the num_tokens*top_k real pairs (skipping pad), and never
|
||||
// materializes the bf16 permuted buffer.
|
||||
//
|
||||
// It mirrors `nvfp4QuantAndPerTokenScaleKernel` (quantization.cuh) — same amax, same per-token-scale
|
||||
// recipe, same `cvt_warp_fp16_to_fp4`, same swizzled-sf offset (`get_sf_out_offset_8x4`) — with the
|
||||
// single `rowIdx` split into a READ row (the unpermuted source token) and a WRITE row (the permuted
|
||||
// destination). For the valid rows the result is BITWISE-identical to the plain permute->quant chain
|
||||
// (the chain's quant reads permuted_hidden[writeRow], filled by permute from hidden[readRow]; we
|
||||
// read hidden[readRow] directly), verified by the bench's fused-vs-old guard.
|
||||
//
|
||||
// PER-TOKEN-SCALE BRANCH: uses TE_EXACT (globalEncodeScale = __fdiv_rn(globalScale, globalAmax),
|
||||
// stored scale = 1/globalEncodeScale, cvt TE_EXACT_NVFP4=true). This matches the installed
|
||||
// flashinfer 0.6.11.post1, whose DISPATCH macro hard-codes the bf16 kernel to TE_EXACT_NVFP4=true
|
||||
// (quantization.cu DISPATCH_NVP4_QUANT_AND_PER_TOKEN_SCALE_KERNEL). NOTE: if a future flashinfer
|
||||
// reverts the bf16 path to the fast-math branch (reciprocal_approximate / TE_EXACT=false), or if
|
||||
// FLASHINFER_NVFP4_4OVER6 is enabled, this fused kernel would diverge — re-validate the bench's
|
||||
// fused-vs-old bitwise guard against the deployed flashinfer before trusting it there.
|
||||
//
|
||||
// Two variants (both kept, selectable, for cross-scenario perf comparison):
|
||||
// - no-dedup: grid over the num_tokens*top_k pairs; each block re-reads+re-quantizes its source
|
||||
// token and writes 1 destination (more blocks -> better occupancy at tiny decode sizes).
|
||||
// - dedup: grid over num_tokens; each block reads+quantizes its token once and scatter-writes
|
||||
// to all of that token's (valid) permuted destinations (no redundant quant, fewer blocks).
|
||||
//
|
||||
// Helpers are pulled from quantization_utils.cuh (cvt_warp_fp16_to_fp4 / get_sf_out_offset_* /
|
||||
// PackedVec / reciprocal_approximate_ftz) rather than quantization.cuh, because the latter pulls in
|
||||
// nv_internal/.../common/cudaUtils.h, which ODR-conflicts with the flashinfer/trtllm/common twin
|
||||
// already in trtllm_fused_moe_kernel_launcher.cu's TU. loadPackedVec lives in quantization.cuh, so
|
||||
// we do a direct aligned PackedVec load instead.
|
||||
#pragma once
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
|
||||
#include <cstdint>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <optional>
|
||||
#include <type_traits>
|
||||
|
||||
namespace sgl_fused_permute_quant {
|
||||
|
||||
namespace tk = tensorrt_llm::kernels;
|
||||
|
||||
// Quantize source row `readRow` of `input` (unpermuted) and write fp4 + block-sf + per-token-sf to
|
||||
// destination row `writeRow` of the permuted outputs. `numRowsSf` is the SF buffer's row count
|
||||
// (= max_padded), matching the plain quant's `m` arg to get_sf_out_offset_*.
|
||||
template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
|
||||
__device__ __forceinline__ void fused_quant_one_row(
|
||||
uint32_t n,
|
||||
T const* input,
|
||||
int readRow,
|
||||
int writeRow,
|
||||
int numRowsSf,
|
||||
float globalScaleInv,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput) {
|
||||
constexpr int SF_VEC_SIZE = 16;
|
||||
constexpr int ELTS_PER_THREAD = 16;
|
||||
using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
|
||||
using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
|
||||
uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
|
||||
uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
|
||||
InType const* inBase = reinterpret_cast<InType const*>(input);
|
||||
|
||||
// ---- pass 1: per-row amax over the (unpermuted) source row ----
|
||||
float localAmax = 0.f;
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
|
||||
std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
|
||||
a = __hmax2(a, __habs2(vec_in.elts[i]));
|
||||
}
|
||||
localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
|
||||
}
|
||||
|
||||
using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
|
||||
|
||||
// ---- per-token scale (TE_EXACT branch — production instantiates TE_EXACT_NVFP4=true for bf16,
|
||||
// quantization.cu:247): globalEncodeScale = globalScale/globalAmax (exact __fdiv_rn), stored
|
||||
// per-token scale = 1/globalEncodeScale. __shared__ scalar replaces the gmem round-trip
|
||||
// (bit-identical: an fp32 store->load doesn't change the value). ----
|
||||
__shared__ float sEncodeScale;
|
||||
if (threadIdx.x == 0) {
|
||||
float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
|
||||
float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
|
||||
sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
float const globalEncodeScale = sEncodeScale;
|
||||
float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
|
||||
if (threadIdx.x == 0) perTokenScaleOutput[writeRow] = perTokenScale;
|
||||
|
||||
// ---- pass 2: quantize + scatter-write to the permuted destination ----
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
|
||||
uint8_t fp8Scale;
|
||||
auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
|
||||
T,
|
||||
SF_VEC_SIZE,
|
||||
ELTS_PER_THREAD,
|
||||
/*UE8M0_SF=*/false,
|
||||
/*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
|
||||
reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
|
||||
fp4Vals;
|
||||
|
||||
int64_t sfOffset;
|
||||
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);
|
||||
} else {
|
||||
sfOffset = tk::get_sf_out_offset_8x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
|
||||
}
|
||||
scaleOutput[sfOffset] = fp8Scale;
|
||||
}
|
||||
}
|
||||
|
||||
// no-dedup: grid.x = num_tokens*top_k (one block per (token,expert) pair).
|
||||
template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
|
||||
__global__ void fusedPermuteNvfp4QuantKernel(
|
||||
uint32_t numPairs,
|
||||
uint32_t n,
|
||||
uint32_t topK,
|
||||
int numRowsSf,
|
||||
T const* input,
|
||||
float globalScaleInv,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput) {
|
||||
uint32_t const expandedIdx = blockIdx.x;
|
||||
if (expandedIdx >= numPairs) return;
|
||||
int const writeRow = expandedIdxToPermutedIdx[expandedIdx];
|
||||
if (writeRow < 0) return;
|
||||
int const readRow = static_cast<int>(expandedIdx / topK);
|
||||
fused_quant_one_row<T, BLOCK_SIZE, SF_LAYOUT>(
|
||||
n, input, readRow, writeRow, numRowsSf, globalScaleInv, weightOutput, scaleOutput, perTokenScaleOutput);
|
||||
}
|
||||
|
||||
// dedup: grid.x = num_tokens (one block per source token, scatter to its top_k destinations).
|
||||
template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
|
||||
__global__ void fusedPermuteNvfp4QuantDedupKernel(
|
||||
uint32_t numTokens,
|
||||
uint32_t n,
|
||||
uint32_t topK,
|
||||
int numRowsSf,
|
||||
T const* input,
|
||||
float globalScaleInv,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput) {
|
||||
constexpr int SF_VEC_SIZE = 16;
|
||||
constexpr int ELTS_PER_THREAD = 16;
|
||||
using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
|
||||
using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
|
||||
uint32_t const token = blockIdx.x;
|
||||
if (token >= numTokens) return;
|
||||
uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
|
||||
uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
|
||||
InType const* inBase = reinterpret_cast<InType const*>(input);
|
||||
|
||||
// pass 1: amax over the source token row (read once).
|
||||
float localAmax = 0.f;
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
|
||||
std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
|
||||
a = __hmax2(a, __habs2(vec_in.elts[i]));
|
||||
}
|
||||
localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
|
||||
}
|
||||
using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
|
||||
|
||||
// TE_EXACT per-token scale (matches production; see fused_quant_one_row).
|
||||
__shared__ float sEncodeScale;
|
||||
if (threadIdx.x == 0) {
|
||||
float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
|
||||
float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
|
||||
sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
float const globalEncodeScale = sEncodeScale;
|
||||
float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
|
||||
|
||||
// per-token scale -> each (valid) destination (top_k small; first top_k threads write).
|
||||
if (threadIdx.x < topK) {
|
||||
int const writeRow = expandedIdxToPermutedIdx[token * topK + threadIdx.x];
|
||||
if (writeRow >= 0) perTokenScaleOutput[writeRow] = perTokenScale;
|
||||
}
|
||||
|
||||
// pass 2: quantize each vec once, scatter to all valid destinations.
|
||||
for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
|
||||
InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
|
||||
uint8_t fp8Scale;
|
||||
auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
|
||||
T,
|
||||
SF_VEC_SIZE,
|
||||
ELTS_PER_THREAD,
|
||||
/*UE8M0_SF=*/false,
|
||||
/*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
|
||||
#pragma unroll 1
|
||||
for (uint32_t k = 0; k < topK; ++k) {
|
||||
int const writeRow = expandedIdxToPermutedIdx[token * topK + k];
|
||||
if (writeRow < 0) continue;
|
||||
reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
|
||||
fp4Vals;
|
||||
int64_t sfOffset;
|
||||
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);
|
||||
} else {
|
||||
sfOffset = tk::get_sf_out_offset_8x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
|
||||
}
|
||||
scaleOutput[sfOffset] = fp8Scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Launcher. `dedup` picks the variant. `n` (= hidden) must be a multiple of 16. `numRowsSf` is the
|
||||
// SF buffer's row count (max_padded).
|
||||
template <typename T>
|
||||
void invokeFusedPermuteNvfp4Quant(
|
||||
uint32_t numTokens,
|
||||
uint32_t topK,
|
||||
uint32_t n,
|
||||
int numRowsSf,
|
||||
T const* input,
|
||||
float globalScaleInv,
|
||||
int32_t const* expandedIdxToPermutedIdx,
|
||||
uint8_t* weightOutput,
|
||||
uint8_t* scaleOutput,
|
||||
float* perTokenScaleOutput,
|
||||
tensorrt_llm::QuantizationSFLayout sfLayout,
|
||||
bool dedup,
|
||||
cudaStream_t stream) {
|
||||
// [opt] Occupancy tuning (ncu: kernel is occupancy-bound, not DRAM-bound — DRAM <1%, achieved
|
||||
// occupancy was 19.8% no-dedup / 5.5% dedup at BLOCK_SIZE=128). The dedup variant launches only
|
||||
// num_tokens CTAs (=64 at decode bs64), so it is the most CTA-starved; widening the block raises
|
||||
// threads/CTA and hides the per-row amax-reduction + scatter latency. Decode bs64 dedup sweep:
|
||||
// 128 -> 5.52us, 256 -> 4.09us, 512 -> 3.71us. 512 is the chosen default (the prod path uses
|
||||
// dedup). (7168/16 = 448 vecs/row, so >448 threads idle on the tail, but the win dominates.)
|
||||
constexpr uint32_t BLOCK_SIZE = 512;
|
||||
dim3 const block(BLOCK_SIZE);
|
||||
|
||||
auto dispatch = [&](auto layoutTag) {
|
||||
constexpr tensorrt_llm::QuantizationSFLayout LAYOUT = decltype(layoutTag)::value;
|
||||
if (dedup) {
|
||||
dim3 const grid(numTokens);
|
||||
fusedPermuteNvfp4QuantDedupKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
|
||||
numTokens,
|
||||
n,
|
||||
topK,
|
||||
numRowsSf,
|
||||
input,
|
||||
globalScaleInv,
|
||||
expandedIdxToPermutedIdx,
|
||||
weightOutput,
|
||||
scaleOutput,
|
||||
perTokenScaleOutput);
|
||||
} else {
|
||||
dim3 const grid(numTokens * topK);
|
||||
fusedPermuteNvfp4QuantKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
|
||||
numTokens * topK,
|
||||
n,
|
||||
topK,
|
||||
numRowsSf,
|
||||
input,
|
||||
globalScaleInv,
|
||||
expandedIdxToPermutedIdx,
|
||||
weightOutput,
|
||||
scaleOutput,
|
||||
perTokenScaleOutput);
|
||||
}
|
||||
};
|
||||
|
||||
if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
|
||||
dispatch(
|
||||
std::integral_constant<
|
||||
tensorrt_llm::QuantizationSFLayout,
|
||||
tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
|
||||
} else {
|
||||
dispatch(
|
||||
std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace sgl_fused_permute_quant
|
||||
+3920
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+472
@@ -0,0 +1,472 @@
|
||||
/*
|
||||
* Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
|
||||
#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h"
|
||||
#include <cuda.h>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
// #include <cuda_runtime_api.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cutlass/numeric_size.h>
|
||||
#include <cutlass/numeric_types.h>
|
||||
|
||||
#include "flashinfer/exception.h"
|
||||
// #include <tensorrt_llm/common/assert.h>
|
||||
#include "flashinfer/trtllm/common/cudaUtils.h"
|
||||
#include "tensorrt_llm/common/logger.h"
|
||||
|
||||
namespace moe::dev {
|
||||
|
||||
#define CHECK_CUDA_ERROR(cmd) \
|
||||
do { \
|
||||
cudaError_t e = cmd; \
|
||||
if (e != cudaSuccess) { \
|
||||
std::cout << "CUDA error in " << __FILE__ << ":" << __LINE__ << " executing '" << #cmd \
|
||||
<< "': " << cudaGetErrorString(e); \
|
||||
} \
|
||||
FLASHINFER_CHECK(e == cudaSuccess, "Got CUDA error. See above for details."); \
|
||||
} while (0)
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#define LAUNCH_ESC(...) __VA_ARGS__
|
||||
|
||||
#define LAUNCH_PDL(data, coopLaunch, types, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
cudaLaunchConfig_t config{}; \
|
||||
config.gridDim = numBlocks; \
|
||||
config.blockDim = numThreads; \
|
||||
config.dynamicSmemBytes = smemSize; \
|
||||
config.stream = (cudaStream_t)stream; \
|
||||
\
|
||||
cudaLaunchAttribute attributes[2] = {}; \
|
||||
attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
|
||||
attributes[0].val.programmaticStreamSerializationAllowed = int(data.mUsePdl); \
|
||||
attributes[1].id = cudaLaunchAttributeCooperative; \
|
||||
attributes[1].val.cooperative = int(coopLaunch); \
|
||||
config.attrs = attributes; \
|
||||
config.numAttrs = 2; \
|
||||
if (data.mUsePdl) { \
|
||||
auto params = KernelParams<types, true>::setKernelParams(data); \
|
||||
auto kernelTyped = kernel<KernelParams<types, true>>; \
|
||||
if (smemSize > 48 * 1024) \
|
||||
CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
|
||||
CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params)); \
|
||||
} else { \
|
||||
auto params = KernelParams<types, false>::setKernelParams(data); \
|
||||
auto kernelTyped = kernel<KernelParams<types, false>>; \
|
||||
if (smemSize > 48 * 1024) \
|
||||
CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
|
||||
CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params)); \
|
||||
}
|
||||
|
||||
#define LAUNCH(data, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.mDtypeElt == tg::Dtype::Fp16) { \
|
||||
LAUNCH_PDL(data, false, cutlass::half_t, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3) { \
|
||||
LAUNCH_PDL(data, false, cutlass::float_e4m3_t, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL(data, false, cutlass::bfloat16_t, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported dtypeElt"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_NUM_TOKENS_PER_CTA(data, type, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
if (numTokensPerCta == 4) { \
|
||||
LAUNCH_PDL(data, false, LAUNCH_ESC(type, 4), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (numTokensPerCta == 2) { \
|
||||
LAUNCH_PDL(data, false, LAUNCH_ESC(type, 2), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (numTokensPerCta == 1) { \
|
||||
LAUNCH_PDL(data, false, LAUNCH_ESC(type, 1), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported numTokensPerCta"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_ACTIVATION(data, kernel, numTokensPerCta, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.mDtypeElt == tg::Dtype::Fp16) { \
|
||||
LAUNCH_NUM_TOKENS_PER_CTA( \
|
||||
data, cutlass::half_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3) { \
|
||||
LAUNCH_NUM_TOKENS_PER_CTA( \
|
||||
data, cutlass::float_e4m3_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_NUM_TOKENS_PER_CTA( \
|
||||
data, cutlass::bfloat16_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported dtypeElt"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_EXPW(data, kernel, topK, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Fp32) { \
|
||||
LAUNCH_PDL( \
|
||||
data, false, LAUNCH_ESC(cutlass::half_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Fp32) { \
|
||||
LAUNCH_PDL( \
|
||||
data, false, LAUNCH_ESC(cutlass::float_e4m3_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Fp32) { \
|
||||
LAUNCH_PDL( \
|
||||
data, false, LAUNCH_ESC(cutlass::bfloat16_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL( \
|
||||
data, \
|
||||
false, \
|
||||
LAUNCH_ESC(cutlass::half_t, cutlass::bfloat16_t, topK), \
|
||||
kernel, \
|
||||
numBlocks, \
|
||||
numThreads, \
|
||||
smemSize, \
|
||||
stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL( \
|
||||
data, \
|
||||
false, \
|
||||
LAUNCH_ESC(cutlass::float_e4m3_t, cutlass::bfloat16_t, topK), \
|
||||
kernel, \
|
||||
numBlocks, \
|
||||
numThreads, \
|
||||
smemSize, \
|
||||
stream); \
|
||||
} else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \
|
||||
LAUNCH_PDL( \
|
||||
data, \
|
||||
false, \
|
||||
LAUNCH_ESC(cutlass::bfloat16_t, cutlass::bfloat16_t, topK), \
|
||||
kernel, \
|
||||
numBlocks, \
|
||||
numThreads, \
|
||||
smemSize, \
|
||||
stream); \
|
||||
} else { \
|
||||
FLASHINFER_WARN("Unsupported pair"); \
|
||||
}
|
||||
|
||||
#define LAUNCH_TOPK_EXPW(data, kernel, numBlocks, numThreads, smemSize, stream) \
|
||||
if (data.topK % 4 == 0) { \
|
||||
LAUNCH_EXPW(data, kernel, 4, numBlocks, numThreads, smemSize, stream); \
|
||||
} else if (data.topK % 2 == 0) { \
|
||||
LAUNCH_EXPW(data, kernel, 2, numBlocks, numThreads, smemSize, stream); \
|
||||
} else { \
|
||||
LAUNCH_EXPW(data, kernel, 1, numBlocks, numThreads, smemSize, stream); \
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// NOTE: Old routing-specific macros (LAUNCH_TILEN, LAUNCH_ROUTING_LLAMA4,
|
||||
// LAUNCH_ROUTING_DEEPSEEK_*, LAUNCH_ROUTING_WITH_NUM_EXPERTS) have been moved to
|
||||
// RoutingDevKernel.h which uses the new template signature with runtime isPow2/UsePdl.
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace activation {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
tg::Dtype mDtypeElt{tg::Dtype::Fp16};
|
||||
bool mUsePdl{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
|
||||
void* inPtr;
|
||||
void* outPtr;
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
|
||||
cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
|
||||
|
||||
// When true, inPtr holds the column-interleaved gate/up GEMM1 output
|
||||
// (g0,u0,g1,u1,...) and the kernel de-interleaves on read (x1=col 2k, x2=col 2k+1);
|
||||
// when false, inPtr is the contiguous [gate | up] layout. Default false keeps the
|
||||
// FP8 / non-LoRA semantics; the FP4 LoRA path sets it true to fuse the standalone
|
||||
// de-interleave kernel into this activation read.
|
||||
bool interleavedGateUpInput = false;
|
||||
|
||||
int32_t innerDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
|
||||
// Bench/tuning knob: when > 0, overrides the activation launch grid.x (default
|
||||
// innerDim/128). The kernel's grid-stride hidden-dim loop makes any grid.x
|
||||
// produce bitwise-identical output, so a smaller grid.x removes empty blocks and
|
||||
// gives each thread a longer strip (more in-flight loads) without changing math.
|
||||
int32_t actGridXOverride = 0;
|
||||
|
||||
// Activation kernel variant: 0 = scalar activationKernel, 1 = vectorized
|
||||
// activationKernelOpt (128-bit gate/up + 64-bit delta/store, 4 pairs/thread).
|
||||
// Variant 1 applies only to the bf16 interleaved path with (innerDim/2)%4==0;
|
||||
// run() falls back to the scalar kernel otherwise. Output is bitwise-identical.
|
||||
int32_t actOptMode = 0;
|
||||
};
|
||||
|
||||
template <typename Type_, int32_t NumTokensPerCta_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
static constexpr int32_t NumTokensPerCta = NumTokensPerCta_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
Type const* inPtr;
|
||||
Type* outPtr;
|
||||
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
|
||||
cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
|
||||
|
||||
bool interleavedGateUpInput = false;
|
||||
|
||||
int32_t innerDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inPtr = (Type*)data.inPtr;
|
||||
params.outPtr = (Type*)data.outPtr;
|
||||
params.inDqSfsPtr = data.inDqSfsPtr;
|
||||
params.outDqSfsPtr = data.outDqSfsPtr;
|
||||
params.gateUpLoraDeltaPtr = data.gateUpLoraDeltaPtr;
|
||||
params.activationLoraInputOutPtr = data.activationLoraInputOutPtr;
|
||||
params.interleavedGateUpInput = data.interleavedGateUpInput;
|
||||
|
||||
params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
|
||||
|
||||
params.innerDim = data.innerDim;
|
||||
params.numTokens = data.numTokens;
|
||||
params.topK = data.topK;
|
||||
params.totalNumPaddedTokens = data.totalNumPaddedTokens;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
} // namespace activation
|
||||
|
||||
namespace convertsf {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
bool mUsePdl{false};
|
||||
|
||||
void* inSfPtr = nullptr;
|
||||
void* outSfPtr = nullptr;
|
||||
int32_t hiddenDimSf;
|
||||
int32_t numTokens;
|
||||
tg::SfLayout sfLayoutSrc;
|
||||
tg::SfLayout sfLayoutDst;
|
||||
};
|
||||
|
||||
template <typename Type_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
void const* inSfPtr = nullptr;
|
||||
void* outSfPtr = nullptr;
|
||||
int32_t hiddenDimSf;
|
||||
int32_t numTokens;
|
||||
tg::SfLayout sfLayoutSrc;
|
||||
tg::SfLayout sfLayoutDst;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inSfPtr = data.inSfPtr;
|
||||
params.outSfPtr = data.outSfPtr;
|
||||
params.hiddenDimSf = data.hiddenDimSf;
|
||||
params.numTokens = data.numTokens;
|
||||
params.sfLayoutSrc = data.sfLayoutSrc;
|
||||
params.sfLayoutDst = data.sfLayoutDst;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace convertsf
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace permute {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
tg::Dtype mDtypeElt{tg::Dtype::Fp16};
|
||||
bool mUsePdl{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
|
||||
void* inPtr;
|
||||
void* outPtr;
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
int32_t hiddenDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
};
|
||||
|
||||
template <typename Type_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
Type const* inPtr;
|
||||
Type* outPtr;
|
||||
float const* inDqSfsPtr;
|
||||
float* outDqSfsPtr;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
int32_t hiddenDim;
|
||||
int32_t numTokens;
|
||||
int32_t topK;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
bool useDeepSeekFp8;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inPtr = (Type*)data.inPtr;
|
||||
params.outPtr = (Type*)data.outPtr;
|
||||
params.inDqSfsPtr = data.inDqSfsPtr;
|
||||
params.outDqSfsPtr = data.outDqSfsPtr;
|
||||
params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
|
||||
params.hiddenDim = data.hiddenDim;
|
||||
params.numTokens = data.numTokens;
|
||||
params.topK = data.topK;
|
||||
params.totalNumPaddedTokens = data.totalNumPaddedTokens;
|
||||
params.useDeepSeekFp8 = data.mUseDeepSeekFp8;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace permute
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace finalize {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace tg = batchedGemm::trtllm::gen;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Data {
|
||||
tg::Dtype mDtypeElt{tg::Dtype::Fp16};
|
||||
tg::Dtype mDtypeExpW{tg::Dtype::Bfloat16};
|
||||
bool mUsePdl{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
|
||||
void* inPtr;
|
||||
void* outPtr;
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
|
||||
void* expertWeightsPtr;
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t numTokens;
|
||||
int32_t numExperts;
|
||||
int32_t topK;
|
||||
// Hidden dimension output of MoE block. It is not padded.
|
||||
int32_t hiddenDim;
|
||||
// Hidden dimension output of FC2. It might be padded.
|
||||
int32_t hiddenDimPadded;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
};
|
||||
|
||||
template <typename Type_, typename TypeExpW_, int TopKUnrollFactor_, bool UsePdl_>
|
||||
struct KernelParams {
|
||||
using Type = Type_;
|
||||
using TypeExpW = TypeExpW_;
|
||||
static constexpr int TopKUnrollFactor = TopKUnrollFactor_;
|
||||
static constexpr bool UsePdl = UsePdl_;
|
||||
|
||||
Type const* inPtr;
|
||||
TypeExpW const* expertWeightsPtr;
|
||||
Type* outPtr;
|
||||
|
||||
float* inDqSfsPtr = nullptr;
|
||||
float* outDqSfsPtr = nullptr;
|
||||
|
||||
int32_t* expandedIdxToPermutedIdx;
|
||||
|
||||
int32_t hiddenDim;
|
||||
int32_t hiddenDimPadded;
|
||||
int32_t numTokens;
|
||||
int32_t numExperts;
|
||||
int32_t topK;
|
||||
int32_t const* totalNumPaddedTokens;
|
||||
|
||||
static KernelParams setKernelParams(Data const& data) {
|
||||
KernelParams params;
|
||||
|
||||
params.inPtr = (Type*)data.inPtr;
|
||||
params.expertWeightsPtr = (TypeExpW*)data.expertWeightsPtr;
|
||||
params.outPtr = (Type*)data.outPtr;
|
||||
params.inDqSfsPtr = data.inDqSfsPtr;
|
||||
params.outDqSfsPtr = data.outDqSfsPtr;
|
||||
|
||||
params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
|
||||
|
||||
params.hiddenDim = data.hiddenDim;
|
||||
params.hiddenDimPadded = data.hiddenDimPadded;
|
||||
params.numTokens = data.numTokens;
|
||||
params.numExperts = data.numExperts;
|
||||
params.topK = data.topK;
|
||||
params.totalNumPaddedTokens = data.totalNumPaddedTokens;
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
void run(Data const& data, void* stream);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace finalize
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace moe::dev
|
||||
+585
@@ -0,0 +1,585 @@
|
||||
/*
|
||||
* Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "DevKernel.h"
|
||||
#include "flashinfer/trtllm/fused_moe/RoutingKernel.h"
|
||||
#include <string>
|
||||
// #include "flashinfer/trtllm/common/cudaDriverWrapper.h"
|
||||
#include "flashinfer/trtllm/batched_gemm/KernelRunner.h"
|
||||
#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
|
||||
#include "flashinfer/trtllm/common/cudaUtils.h"
|
||||
|
||||
namespace tensorrt_llm {
|
||||
namespace kernels {
|
||||
namespace trtllmgen_moe {
|
||||
|
||||
namespace MoE {
|
||||
class Runner;
|
||||
} // namespace MoE
|
||||
|
||||
namespace Routing {
|
||||
|
||||
// The type of method in top-K routing, for use in torch custom op
|
||||
// Please keep this in sync with the counterpart defined in
|
||||
// flashinfer/fused_moe/core.py
|
||||
enum class RoutingMethodType : int64_t {
|
||||
// Default: Softmax -> TopK
|
||||
Default = 0,
|
||||
// Renormalize: TopK -> Softmax
|
||||
Renormalize = 1,
|
||||
// DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups -> Top8 experts from the
|
||||
// Top4 groups
|
||||
DeepSeekV3 = 2,
|
||||
// Llama4: Top1 -> Sigmoid
|
||||
Llama4 = 3,
|
||||
// RenormalizeNaive: Softmax -> TopK -> Renormalize
|
||||
RenormalizeNaive = 4,
|
||||
// TopK only (no softmax)
|
||||
TopK = 5,
|
||||
// SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K weights)
|
||||
SigmoidRenorm = 6,
|
||||
// MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize (routeScale=1.0, epsilon=1e-20)
|
||||
MiniMax2 = 7,
|
||||
// Sigmoid: Sigmoid -> TopK (no renormalization)
|
||||
Sigmoid = 8,
|
||||
// Unspecified
|
||||
Unspecified = 9,
|
||||
};
|
||||
|
||||
inline int32_t maybeGetMinTokenCount(int32_t numPaddedTokens, int32_t hiddenSize, int32_t dtypeSizeBits) {
|
||||
// Pad so total size exceeds 128KiB for performance reasons
|
||||
int32_t minNumTokensRequired = common::divUp(128 * 1024 * 8, hiddenSize * dtypeSizeBits);
|
||||
return std::max(numPaddedTokens, minNumTokensRequired);
|
||||
}
|
||||
|
||||
inline std::string serializeMoeRoutingMethodType(RoutingMethodType routingMethodType) {
|
||||
switch (routingMethodType) {
|
||||
case RoutingMethodType::Default:
|
||||
return "Default";
|
||||
case RoutingMethodType::Renormalize:
|
||||
return "Renormalize";
|
||||
case RoutingMethodType::DeepSeekV3:
|
||||
return "DeepSeekV3";
|
||||
case RoutingMethodType::Llama4:
|
||||
return "Llama4";
|
||||
case RoutingMethodType::RenormalizeNaive:
|
||||
return "RenormalizeNaive";
|
||||
case RoutingMethodType::TopK:
|
||||
return "TopK";
|
||||
case RoutingMethodType::SigmoidRenorm:
|
||||
return "SigmoidRenorm";
|
||||
case RoutingMethodType::MiniMax2:
|
||||
return "MiniMax2";
|
||||
case RoutingMethodType::Sigmoid:
|
||||
return "Sigmoid";
|
||||
default:
|
||||
return "InvalidRountingMethod"; // TODO throw error
|
||||
};
|
||||
}
|
||||
|
||||
inline int32_t getMaxNumCtasInBatchDim(int32_t numTokens, int32_t topK, int32_t numExperts, int32_t tileTokensDim) {
|
||||
// For MoE, mNumTokens != 0 and the number of CTAs is known only at runtime.
|
||||
// We launch maximally possible number of CTAs and use ptrNumNonExitingCtas to determine
|
||||
// the actual number of CTAs to run.
|
||||
|
||||
// Initialize number of tokens with the number of expanded tokens after routing.
|
||||
int32_t numRemainingTokens = numTokens * topK;
|
||||
int32_t maxNumCtasInBatchDim = 0;
|
||||
// First, distribute one token each expert until token depletion to maximize CTA tile count.
|
||||
int32_t numExpertsFilled = std::min(numExperts, numRemainingTokens);
|
||||
maxNumCtasInBatchDim += numExpertsFilled;
|
||||
numRemainingTokens -= numExpertsFilled;
|
||||
// Next, greedily pour all remaining tokens to one expert to maximize CTA tile count.
|
||||
// E.g., at this point tokens over 4 experts are [1, 1, 1, 1], and we have 4 tokens left.
|
||||
// If each CTA handles 4 tokens/expert, the greedy strategy is to pour all remaining tokens
|
||||
// to any one expert to get to the 5th CTA tile. Otherwise, we can only get 4 tiles in total.
|
||||
//
|
||||
// Another way to reason about this is to pour the remaining tokens into buckets of some fixed
|
||||
// capacity. These buckets, if full, can then be attributed to any expert; it does not have to
|
||||
// belong to the same expert every time.
|
||||
if (numRemainingTokens > 0) {
|
||||
// For every tileTokenDim tokens, we add an extra CTA tile in the token dimension.
|
||||
// The number of CTA tiles is given by divDown(numRemainingTokens, tokenTileDim).
|
||||
maxNumCtasInBatchDim += (numRemainingTokens / tileTokensDim);
|
||||
}
|
||||
return maxNumCtasInBatchDim;
|
||||
}
|
||||
|
||||
inline int32_t
|
||||
getMaxPermutedPaddedCount(int32_t numTokens, int32_t expertsPerToken, int32_t numExperts, int32_t padding) {
|
||||
int32_t maxCtas = getMaxNumCtasInBatchDim(numTokens, expertsPerToken, numExperts, padding);
|
||||
return maxCtas * padding;
|
||||
}
|
||||
|
||||
class Runner {
|
||||
public:
|
||||
explicit Runner();
|
||||
|
||||
explicit Runner(int32_t tileTokensDim);
|
||||
|
||||
void
|
||||
run(void* routingLogits,
|
||||
void* routingBias,
|
||||
int32_t numTokens,
|
||||
int32_t numExperts,
|
||||
int32_t topK,
|
||||
int32_t nGroups,
|
||||
int32_t topkGroups,
|
||||
int32_t localExpertOffset,
|
||||
int32_t localNumExperts,
|
||||
float routedScalingFactor,
|
||||
int32_t* routingExpertIndexes,
|
||||
int32_t* expertCountHistogram,
|
||||
int32_t* permutedIdxSize,
|
||||
int32_t* expandedIdxToPermutedIdx,
|
||||
int32_t* permutedIdxToExpandedIdx,
|
||||
int32_t* permutedIdxToTokenIdx,
|
||||
void* expertWeights,
|
||||
int32_t* numTokensPerExpert,
|
||||
int32_t* ctaIdxXyToBatchIdx,
|
||||
int32_t* ctaIdxXyToMnLimit,
|
||||
int32_t* numNonExitingCtas,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeElt,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeBias,
|
||||
bool useRoutingScalesOnInput,
|
||||
bool useDeepSeekFp8,
|
||||
RoutingMethodType routingMethodType,
|
||||
cudaStream_t stream,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeLogits,
|
||||
bool normTopkProb = true,
|
||||
int16_t* routing_replay_out = nullptr);
|
||||
|
||||
private:
|
||||
friend class MoE::Runner;
|
||||
int32_t mTileTokensDim{8};
|
||||
};
|
||||
} // namespace Routing
|
||||
|
||||
namespace MoE {
|
||||
// The type of activation function
|
||||
// Please keep this in sync with the counterpart defined in flashinfer/flashinfer/fused_moe/core.py
|
||||
enum class ActivationType : int64_t {
|
||||
Gelu = 0,
|
||||
Relu = 1,
|
||||
Silu = 2,
|
||||
Swiglu = 3,
|
||||
Geglu = 4,
|
||||
SwigluBias = 5,
|
||||
Relu2 = 6,
|
||||
Identity = 7,
|
||||
InvalidType = 8, // Must be last
|
||||
};
|
||||
|
||||
inline std::string serializeActivationType(ActivationType activationType) {
|
||||
switch (activationType) {
|
||||
case ActivationType::Gelu:
|
||||
return "Gelu";
|
||||
case ActivationType::Relu:
|
||||
return "Relu";
|
||||
case ActivationType::Silu:
|
||||
return "Silu";
|
||||
case ActivationType::Swiglu:
|
||||
return "Swiglu";
|
||||
case ActivationType::Geglu:
|
||||
return "Geglu";
|
||||
case ActivationType::SwigluBias:
|
||||
return "SwigluBias";
|
||||
case ActivationType::Relu2:
|
||||
return "Relu2";
|
||||
case ActivationType::Identity:
|
||||
return "Identity";
|
||||
default:
|
||||
return "InvalidActivationType"; // TODO throw error
|
||||
};
|
||||
}
|
||||
|
||||
inline bool isGatedActivation(ActivationType activationType) {
|
||||
return activationType == ActivationType::Swiglu || activationType == ActivationType::Geglu ||
|
||||
activationType == ActivationType::SwigluBias;
|
||||
}
|
||||
|
||||
} // namespace MoE
|
||||
|
||||
namespace PermuteGemm1 {
|
||||
class Runner {
|
||||
public:
|
||||
explicit Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeAct,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeWeights,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeOutput,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim,
|
||||
MoE::ActivationType activationType,
|
||||
bool useShuffledMatrix,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout,
|
||||
bool usePerTokenScaling,
|
||||
bool usePerChannelScaling,
|
||||
bool forceUnfusedAct = false);
|
||||
|
||||
size_t getWorkspaceSizeInBytes(
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t configIndex) const;
|
||||
|
||||
[[nodiscard]] int32_t getDefaultValidConfigIndex(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] bool isValidConfigIndex(
|
||||
int32_t configIndex,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
|
||||
|
||||
void
|
||||
run(void* hiddenState,
|
||||
void* hiddenStateScale,
|
||||
void* weight,
|
||||
void* weightScale,
|
||||
void* perTokenScales,
|
||||
void* perChannelScales,
|
||||
float* outputScalesScalar,
|
||||
float* outputScalesGateScalar,
|
||||
float* ptrBias,
|
||||
float* ptrGatedActAlpha,
|
||||
float* ptrGatedActBeta,
|
||||
float* ptrClampLimit,
|
||||
void* output,
|
||||
void* outputScale,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t* permutedIdxToTokenIdx,
|
||||
int32_t* ptrNumNonExitingCtas,
|
||||
int32_t* ptrTotalNumPaddedTokens,
|
||||
int32_t* ptrCtaIdxXyToBatchIdx,
|
||||
int32_t* ptrCtaIdxXyToMnLimit,
|
||||
void* bmm1Workspace,
|
||||
bool useRoutingScalesOnInput,
|
||||
int device,
|
||||
cudaStream_t stream,
|
||||
int32_t configIndex,
|
||||
bool enable_pdl);
|
||||
|
||||
private:
|
||||
friend class MoE::Runner;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeAct;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeWeights;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeOutput;
|
||||
int32_t mTileTokensDim;
|
||||
tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
|
||||
tensorrt_llm::kernels::trtllmgen_moe::MoE::ActivationType mActType;
|
||||
};
|
||||
} // namespace PermuteGemm1
|
||||
|
||||
namespace Gemm2 {
|
||||
class Runner {
|
||||
public:
|
||||
explicit Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeAct,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeWeights,
|
||||
batchedGemm::trtllm::gen::Dtype outputDtype,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim,
|
||||
bool useShuffledMatrix,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout,
|
||||
bool usePerTokenScaling,
|
||||
bool usePerChannelScaling);
|
||||
|
||||
size_t getWorkspaceSizeInBytes(
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t configIndex) const;
|
||||
|
||||
[[nodiscard]] int32_t getDefaultValidConfigIndex(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] bool isValidConfigIndex(
|
||||
int32_t configIndex,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
|
||||
|
||||
void
|
||||
run(void* permutedHiddenState,
|
||||
void* permutedHiddenStateScale,
|
||||
void* weight,
|
||||
void* weightScale,
|
||||
void* perTokenScales,
|
||||
void* perChannelScales,
|
||||
float* outputScalesScalar,
|
||||
float* ptrBias,
|
||||
void* output,
|
||||
void* outputScale,
|
||||
int32_t topK,
|
||||
int32_t hiddenSize,
|
||||
int32_t intermediateSize,
|
||||
int32_t numExperts,
|
||||
int32_t numTokens,
|
||||
int32_t* ptrNumNonExitingCtas,
|
||||
int32_t* ptrTotalNumPaddedTokens,
|
||||
int32_t* ptrCtaIdxXyToBatchIdx,
|
||||
int32_t* ptrCtaIdxXyToMnLimit,
|
||||
void* bmm2Workspace,
|
||||
int device,
|
||||
cudaStream_t stream,
|
||||
int32_t configIndex,
|
||||
bool enable_pdl);
|
||||
|
||||
private:
|
||||
friend class MoE::Runner;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeAct;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeWeights;
|
||||
batchedGemm::trtllm::gen::Dtype mDtypeOut;
|
||||
int32_t mTileTokensDim;
|
||||
tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
|
||||
};
|
||||
} // namespace Gemm2
|
||||
|
||||
namespace MoE {
|
||||
namespace btg = batchedGemm::trtllm::gen;
|
||||
|
||||
struct MoERunnerArgs {
|
||||
void* routing_logits = nullptr; // [num_tokens, num_experts] in float, generated after
|
||||
// gemm(hidden_state, routing_weights)
|
||||
void* routing_bias = nullptr; // [num_experts] in bfloat16 for now = mDtypeExpW
|
||||
void* hidden_states = nullptr; // [num_tokens, hidden_size] in fp8 = mDtypeElt
|
||||
// [hidden_size/128, num_tokens] in float for e4m3 DS recipe
|
||||
// and [num_tokens, hidden_size/16] in float for e2m1
|
||||
void* hidden_states_scale = nullptr;
|
||||
|
||||
// Gemm input:
|
||||
void* gemm1_weights = nullptr;
|
||||
void* gemm1_weights_scale = nullptr;
|
||||
void* gemm2_weights = nullptr;
|
||||
void* gemm2_weights_scale = nullptr;
|
||||
|
||||
float* gemm1_bias = nullptr;
|
||||
float* gemm1_alpha = nullptr;
|
||||
float* gemm1_beta = nullptr;
|
||||
float* gemm1_clamp_limit = nullptr;
|
||||
float* gemm2_bias = nullptr;
|
||||
|
||||
ActivationType activation_type = ActivationType::Swiglu;
|
||||
|
||||
int32_t num_tokens{0};
|
||||
int32_t num_experts{0};
|
||||
// Hidden dimension input of MoE block. It might be padded.
|
||||
int32_t hidden_size{0};
|
||||
// Hidden dimension output of MoE block. It is not padded.
|
||||
// If not provided it is the same as hidden_size.
|
||||
std::optional<int32_t> hidden_size_output;
|
||||
// TODO: only compiled routing kernel supports top_k = 8
|
||||
int32_t top_k{0};
|
||||
int32_t n_group{0};
|
||||
// TODO: only compiled routing kernel supports topk_group = 4
|
||||
int32_t topk_group{0};
|
||||
float routed_scaling_factor{0.0f};
|
||||
int32_t intermediate_size{0};
|
||||
int32_t local_expert_offset{0};
|
||||
int32_t local_num_experts{0};
|
||||
// TODO: support other types
|
||||
btg::Dtype mDtypeElt{btg::Dtype::Void};
|
||||
btg::Dtype mDtypeExpW{btg::Dtype::Bfloat16};
|
||||
btg::Dtype mDtypeOut{btg::Dtype::Bfloat16};
|
||||
|
||||
// Apply routing scale factors to input activations
|
||||
bool mUseRoutingScalesOnInput{false};
|
||||
bool mUseDeepSeekFp8{false};
|
||||
float* output1_scales_scalar = nullptr;
|
||||
float* output1_scales_gate_scalar = nullptr;
|
||||
float* output2_scales_scalar = nullptr;
|
||||
|
||||
// Optional LoRA bridge buffers used by the copied SGLang TRTLLM FP8 path.
|
||||
// gate_up_lora_delta: [num_tokens * top_k, 2 * intermediate_size], bf16,
|
||||
// in FlashInfer gate/up order (up first, gate second).
|
||||
// activation_lora_input: [num_tokens * top_k, intermediate_size], bf16,
|
||||
// populated with the post-activation intermediate for down-proj LoRA.
|
||||
void* gate_up_lora_delta = nullptr;
|
||||
void* activation_lora_input = nullptr;
|
||||
|
||||
// Optional CUDA event (cudaEvent_t) recorded on the LoRA side stream. When set, the
|
||||
// runner waits on it right before the activation kernel (which consumes
|
||||
// gate_up_lora_delta), so permute+GEMM1 overlap the side-stream LoRA shrink/expand
|
||||
// instead of joining before the whole MoE op. nullptr = no wait (serial behavior).
|
||||
void* lora_ready_event = nullptr;
|
||||
|
||||
// Down-LoRA/finalize overlap: optional CUDA event (cudaEvent_t) the runner records on the
|
||||
// MoE stream right after GEMM2 (the base down GEMM), before finalize. The LoRA side stream
|
||||
// waits on it to run the down-proj LoRA shrink/expand concurrent with the finalize kernel.
|
||||
// nullptr = no record (serial behavior).
|
||||
void* gemm2_done_event = nullptr;
|
||||
|
||||
// Output:
|
||||
void* output = nullptr;
|
||||
float* output_scale = nullptr;
|
||||
|
||||
// finalize
|
||||
bool do_finalize{true};
|
||||
};
|
||||
|
||||
struct MoEWorkspace {
|
||||
// Routing intermediate outputs:
|
||||
int32_t* routing_expert_indexes = nullptr;
|
||||
int32_t* permuted_idx_size = nullptr;
|
||||
int32_t* total_num_padded_tokens = nullptr; // TODO: duplicate of permuted_idx_size
|
||||
int32_t total_max_padded_tokens{0};
|
||||
|
||||
int32_t* expanded_idx_to_permuted_idx = nullptr;
|
||||
int32_t* permuted_idx_to_expanded_idx = nullptr;
|
||||
int32_t* permuted_idx_to_token_idx = nullptr;
|
||||
|
||||
// consumed by finalize kernel
|
||||
void* expert_weights = nullptr; // [num_tokens, top_k] in bfloat16 = mDtypeExpW
|
||||
// consumed by permuteGemm1 kernel
|
||||
void* token_scales = nullptr;
|
||||
// consumed by Gemm2 kernel
|
||||
void* token_scales_fc2 = nullptr;
|
||||
|
||||
int32_t* cta_idx_xy_to_batch_idx = nullptr;
|
||||
int32_t* cta_idx_xy_to_mn_limit = nullptr;
|
||||
int32_t* num_non_exiting_ctas = nullptr;
|
||||
|
||||
void* hidden_states_scale_linear = nullptr;
|
||||
|
||||
// Permute intermediate outputs:
|
||||
void* permuted_hidden_states = nullptr;
|
||||
float* permuted_hidden_states_scale = nullptr;
|
||||
|
||||
// Gemm1 intermediate outputs:
|
||||
int32_t ProjUpTileN{0};
|
||||
void* gemm1_output = nullptr;
|
||||
float* gemm1_output_scale = nullptr;
|
||||
|
||||
// Activation intermediate outputs:
|
||||
void* activation_output = nullptr;
|
||||
float* activation_output_scale = nullptr;
|
||||
// Unfused FP4 LoRA: bf16 [max_padded_tokens, intermediate_size] activation output written by the
|
||||
// standalone activation kernel (gate_up LoRA added pre-SwiGLU), then NvFP4-quantized for GEMM2.
|
||||
void* activated_lora_bf16 = nullptr;
|
||||
|
||||
// Gemm2 intermediate outputs:
|
||||
void* gemm2_output = nullptr;
|
||||
float* gemm2_output_scale = nullptr;
|
||||
|
||||
// Finalize intermediate outputs (placeholder not used)
|
||||
void* finalize_output = nullptr;
|
||||
float* finalize_output_scale = nullptr;
|
||||
|
||||
// FC1 workspace:
|
||||
void* bmm1_workspace = nullptr;
|
||||
|
||||
// FC2 workspace:
|
||||
void* bmm2_workspace = nullptr;
|
||||
};
|
||||
|
||||
// Config indices to be used with Batched GEMM runners
|
||||
struct MoEConfig {
|
||||
int64_t gemm1Config;
|
||||
int64_t gemm2Config;
|
||||
};
|
||||
|
||||
class Runner {
|
||||
public:
|
||||
// FIXME: tileTokensDim is hardcoded for now
|
||||
Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeAct,
|
||||
batchedGemm::trtllm::gen::Dtype dtypeWeights,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim = 8,
|
||||
ActivationType activationType = ActivationType::Swiglu,
|
||||
bool useShuffledMatrix = false,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
|
||||
bool usePerTokenScalingGemm1 = false,
|
||||
bool usePerTokenScalingGemm2 = false,
|
||||
bool usePerChannelScalingGemm1 = false,
|
||||
bool usePerChannelScalingGemm2 = false,
|
||||
bool unfuseActForLora = false);
|
||||
Runner(
|
||||
batchedGemm::trtllm::gen::Dtype dtypeElt,
|
||||
bool useDeepSeekFp8,
|
||||
int tileTokensDim = 8,
|
||||
bool useShuffledMatrix = false,
|
||||
batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
|
||||
bool usePerTokenScalingGemm1 = false,
|
||||
bool usePerTokenScalingGemm2 = false,
|
||||
bool usePerChannelScalingGemm1 = false,
|
||||
bool usePerChannelScalingGemm2 = false);
|
||||
|
||||
void
|
||||
run(MoERunnerArgs const& args,
|
||||
MoEWorkspace const& workspace,
|
||||
int device,
|
||||
cudaStream_t stream,
|
||||
int64_t configIndex,
|
||||
bool enable_pdl);
|
||||
|
||||
[[nodiscard]] std::tuple<int32_t, int32_t>
|
||||
getWorkspaceSizeInBytes(MoERunnerArgs const& args, int64_t configIndex) const;
|
||||
|
||||
[[nodiscard]] std::vector<int64_t> getValidConfigIndices(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
|
||||
|
||||
[[nodiscard]] int64_t getDefaultValidConfigIndex(
|
||||
int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
|
||||
|
||||
private:
|
||||
void setOpsData(
|
||||
MoERunnerArgs const& args,
|
||||
MoEWorkspace const& workspace,
|
||||
moe::dev::convertsf::Data& convertSfData,
|
||||
moe::dev::activation::Data& activationData,
|
||||
moe::dev::finalize::Data& finalizeData);
|
||||
|
||||
private:
|
||||
bool mUsePerTokenScalingGemm1;
|
||||
bool mUsePerTokenScalingGemm2;
|
||||
bool mUsePerChannelScalingGemm1;
|
||||
bool mUsePerChannelScalingGemm2;
|
||||
// When true (FP4 LoRA path), GEMM1 emits the raw gate_up projection (fusedAct=false) so the
|
||||
// standalone activation kernel can inject the gate_up LoRA delta pre-SwiGLU and capture the
|
||||
// post-activation input for the down LoRA — mirroring the DeepSeek-FP8 unfused activation.
|
||||
bool mUnfuseActForLora;
|
||||
PermuteGemm1::Runner mPermuteGemm1;
|
||||
Gemm2::Runner mGemm2;
|
||||
|
||||
// This will be the cartesian product of the passing configs for gemm1 and gemm2
|
||||
// This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2
|
||||
// separately
|
||||
std::vector<MoEConfig> mPassingConfigs;
|
||||
};
|
||||
} // namespace MoE
|
||||
|
||||
} // namespace trtllmgen_moe
|
||||
} // namespace kernels
|
||||
} // namespace tensorrt_llm
|
||||
@@ -0,0 +1,101 @@
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _data_dir() -> Path:
|
||||
return Path(__file__).resolve().parent / "data"
|
||||
|
||||
|
||||
def gen_sgl_trtllm_gen_fused_moe_sm100_module():
|
||||
import flashinfer
|
||||
from flashinfer.artifacts import ArtifactPath, CheckSumHash
|
||||
from flashinfer.jit import env as jit_env
|
||||
from flashinfer.jit.core import current_compilation_context, gen_jit_spec
|
||||
from flashinfer.jit.cubin_loader import (
|
||||
ensure_symlink,
|
||||
get_artifact,
|
||||
get_meta_hash,
|
||||
verify_symlinked_headers,
|
||||
)
|
||||
from flashinfer.jit.fused_moe import BMM_EXPORT_HEADERS
|
||||
|
||||
overlay_data_dir = _data_dir()
|
||||
overlay_csrc_dir = overlay_data_dir / "csrc"
|
||||
overlay_include_dir = overlay_data_dir / "include"
|
||||
flashinfer_data_dir = Path(flashinfer.__file__).resolve().parent / "data"
|
||||
flashinfer_csrc_dir = flashinfer_data_dir / "csrc"
|
||||
flashinfer_include_dir = flashinfer_data_dir / "include"
|
||||
|
||||
include_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/include"
|
||||
header_name = "flashinferMetaInfo"
|
||||
checksum_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt"
|
||||
checksum = get_artifact(checksum_path, CheckSumHash.TRTLLM_GEN_BMM)
|
||||
assert checksum, f"Failed to get checksums.txt from {checksum_path}"
|
||||
meta_hash = get_meta_hash(checksum)
|
||||
|
||||
metainfo = get_artifact(f"{include_path}/{header_name}.h", meta_hash)
|
||||
assert metainfo, f"{header_name}.h not found"
|
||||
|
||||
bmm_export_path = f"{include_path}/trtllmGen_bmm_export"
|
||||
for header in BMM_EXPORT_HEADERS:
|
||||
h = get_artifact(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header))
|
||||
assert h, f"{header} not found"
|
||||
|
||||
symlink_path = (
|
||||
jit_env.FLASHINFER_CUBIN_DIR
|
||||
/ "flashinfer"
|
||||
/ "trtllm"
|
||||
/ "batched_gemm"
|
||||
/ "trtllmGen_bmm_export"
|
||||
)
|
||||
ensure_symlink(symlink_path, jit_env.FLASHINFER_CUBIN_DIR / bmm_export_path)
|
||||
verify_symlinked_headers(symlink_path, BMM_EXPORT_HEADERS, checksum)
|
||||
|
||||
nvcc_flags = current_compilation_context.get_nvcc_flags_list(
|
||||
supported_major_versions=[10, 12]
|
||||
)
|
||||
|
||||
return gen_jit_spec(
|
||||
"sgl_fused_moe_trtllm_sm100",
|
||||
[
|
||||
flashinfer_csrc_dir / "nv_internal/cpp/kernels/quantization.cu",
|
||||
flashinfer_csrc_dir / "nv_internal/cpp/common/envUtils.cpp",
|
||||
flashinfer_csrc_dir / "nv_internal/cpp/common/logger.cpp",
|
||||
flashinfer_csrc_dir / "nv_internal/cpp/common/stringUtils.cpp",
|
||||
flashinfer_csrc_dir / "nv_internal/cpp/common/tllmException.cpp",
|
||||
flashinfer_csrc_dir / "nv_internal/cpp/common/memoryUtils.cu",
|
||||
overlay_csrc_dir / "trtllm_fused_moe_kernel_launcher.cu",
|
||||
overlay_csrc_dir / "trtllm_fused_moe_runner.cu",
|
||||
flashinfer_csrc_dir
|
||||
/ "fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu",
|
||||
flashinfer_csrc_dir
|
||||
/ "fused_moe/trtllm_backend/trtllm_fused_moe_routing_llama4.cu",
|
||||
flashinfer_csrc_dir
|
||||
/ "fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu",
|
||||
flashinfer_csrc_dir
|
||||
/ "fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu",
|
||||
overlay_csrc_dir
|
||||
/ "fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu",
|
||||
flashinfer_csrc_dir / "trtllm_batched_gemm_runner.cu",
|
||||
],
|
||||
extra_cuda_cflags=[
|
||||
"-DTLLM_GEN_EXPORT_INTERFACE",
|
||||
"-DTLLM_GEN_EXPORT_FLASHINFER",
|
||||
"-DTLLM_ENABLE_CUDA",
|
||||
"-DENABLE_BF16",
|
||||
"-DENABLE_FP8",
|
||||
"-DENABLE_FP4",
|
||||
"-DCUTLASS_ENABLE_GDC_FOR_SM100=1",
|
||||
f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{ArtifactPath.TRTLLM_GEN_BMM}\\"',
|
||||
]
|
||||
+ nvcc_flags,
|
||||
extra_include_paths=[
|
||||
overlay_include_dir,
|
||||
overlay_csrc_dir,
|
||||
flashinfer_include_dir,
|
||||
flashinfer_csrc_dir,
|
||||
flashinfer_csrc_dir / "nv_internal",
|
||||
flashinfer_csrc_dir / "nv_internal/include",
|
||||
jit_env.FLASHINFER_CUBIN_DIR,
|
||||
jit_env.FLASHINFER_CUBIN_DIR / include_path,
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,70 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_kimi_k2_moe_fused_gate_module() -> Module:
|
||||
return load_jit(
|
||||
"kimi_k2_moe_fused_gate",
|
||||
cuda_files=["trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh"],
|
||||
cuda_wrappers=[
|
||||
("kimi_k2_moe_fused_gate", "KimiK2MoEFusedGateKernel::run"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def kimi_k2_moe_fused_gate(
|
||||
input: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
routed_scaling_factor: float | None = 1.0,
|
||||
apply_routed_scaling_factor_on_output: bool = False,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Kimi K2 MoE fused gate (num_expert_group=1, DeepSeek noaux_tc routing).
|
||||
|
||||
Supports num_experts in {256, 384} and topk <= 8. input and bias are CUDA
|
||||
tensors of float32, bfloat16, or float16 (dtypes may differ between the two);
|
||||
they are widened to fp32 inside the kernel, so callers no longer need to
|
||||
upcast bf16/fp16 router logits or correction bias on the host. Returns
|
||||
(output_weights, expert_indices).
|
||||
"""
|
||||
_supported = (torch.float32, torch.bfloat16, torch.float16)
|
||||
assert (
|
||||
input.dtype in _supported
|
||||
), f"input must be float32/bfloat16/float16, got {input.dtype}"
|
||||
assert (
|
||||
bias.dtype in _supported
|
||||
), f"bias must be float32/bfloat16/float16, got {bias.dtype}"
|
||||
assert input.ndim == 2, "input must be 2D"
|
||||
assert bias.ndim == 1, "bias must be 1D"
|
||||
assert input.size(1) == bias.size(0), "input and bias must have same num_experts"
|
||||
|
||||
num_rows = input.size(0)
|
||||
device = input.device
|
||||
|
||||
output = torch.empty(num_rows, topk, dtype=torch.float32, device=device)
|
||||
indices = torch.empty(num_rows, topk, dtype=torch.int32, device=device)
|
||||
|
||||
module = _jit_kimi_k2_moe_fused_gate_module()
|
||||
module.kimi_k2_moe_fused_gate(
|
||||
input,
|
||||
bias,
|
||||
output,
|
||||
indices,
|
||||
topk,
|
||||
renormalize,
|
||||
float(routed_scaling_factor) if routed_scaling_factor is not None else 1.0,
|
||||
apply_routed_scaling_factor_on_output,
|
||||
)
|
||||
|
||||
return output, indices
|
||||
@@ -0,0 +1,144 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_module(dtype: torch.dtype) -> Module:
|
||||
args = make_cpp_args(dtype)
|
||||
return load_jit(
|
||||
"moe_lora_merged_align",
|
||||
*args,
|
||||
cuda_files=["trtllm_lora_temp/moe_lora_merged_align_kernel.cu"],
|
||||
cuda_wrappers=[
|
||||
("moe_lora_merged_align", f"MoeLoraMergedAlignKernel<{args}>::run"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def supports_merged_align(virtual_num_experts: int) -> bool:
|
||||
"""Commit-1 kernel only implements the (64, 1024] bucket-count branch.
|
||||
|
||||
The bucket count is virtual_num_experts + 1 (the +1 sentinel bucket). Other
|
||||
regimes (small-batch <=64, v2 >1024) keep the old path."""
|
||||
num_buckets = virtual_num_experts + 1
|
||||
return 64 < num_buckets <= 1024
|
||||
|
||||
|
||||
def moe_lora_merged_align(
|
||||
topk_ids: torch.Tensor,
|
||||
token_lora_mapping: torch.Tensor,
|
||||
num_experts: int,
|
||||
shared_outer: bool,
|
||||
max_loras: int,
|
||||
block_size: int,
|
||||
local_expert_offset: int = 0,
|
||||
local_num_experts: Optional[int] = None,
|
||||
do_skip: bool = True,
|
||||
compact: bool = False,
|
||||
fuse_scatter: Optional[bool] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
|
||||
"""Fused replacement for (_fused_virtual_topk_ids + _align_block_size) on the
|
||||
merged-virtual-expert LoRA path.
|
||||
|
||||
Reads raw topk_ids + token_lora_mapping, computes the merged virtual id
|
||||
inline (mirrors _fused_virtual_topk_ids), and aligns to block_size.
|
||||
|
||||
Returns (sorted_token_ids, expert_ids, num_tokens_post_padded,
|
||||
token_lora_mask, virtual_num_experts).
|
||||
"""
|
||||
device = topk_ids.device
|
||||
flat_topk_ids = topk_ids.reshape(-1)
|
||||
if flat_topk_ids.dtype == torch.int64:
|
||||
flat_topk_ids = flat_topk_ids.to(torch.int32)
|
||||
M, top_k = topk_ids.shape
|
||||
numel = M * top_k
|
||||
|
||||
num_experts_for_weight = 1 if shared_outer else num_experts
|
||||
virtual_num_experts = num_experts_for_weight * max_loras
|
||||
ep_local = (
|
||||
(not shared_outer)
|
||||
and (local_num_experts is not None)
|
||||
and (local_num_experts < num_experts_for_weight)
|
||||
)
|
||||
|
||||
# compact: histogram over only the rank's local experts (dense LOCAL ids)
|
||||
# instead of the full global virtual space. Valid only for the single-adapter
|
||||
# EP per-expert path (safe_lora shift is 0; owned ids are a contiguous window
|
||||
# remappable by a single -offset). expert_ids is restored to global in-kernel.
|
||||
compact_eff = compact and ep_local and max_loras == 1 and not shared_outer
|
||||
bucket_experts = local_num_experts if compact_eff else virtual_num_experts
|
||||
|
||||
# fuse_scatter: do the whole align+scatter in one threadblock (one launch).
|
||||
# Only for small numel (the scatter is single-block); large numel (prefill)
|
||||
# keeps the 2-kernel multi-block path. Default auto by numel.
|
||||
fuse_eff = (numel <= 2048) if fuse_scatter is None else fuse_scatter
|
||||
if fuse_eff:
|
||||
# The fused kernel's dynamic shared memory must fit the 48KB default limit
|
||||
# (no cudaFuncSetAttribute opt-in). Layout matches the kernel exactly:
|
||||
# shared_counts + prefix + scan_buf + warp_sums + cursor + svids.
|
||||
nb = bucket_experts + 1 # bucket count (the kernel's num_experts)
|
||||
scan_size = 1 << (nb - 1).bit_length() if nb > 1 else 1
|
||||
fused_shmem = (nb + (nb + 1) + scan_size + 32 + nb + numel) * 4
|
||||
if fused_shmem > 47 * 1024:
|
||||
fuse_eff = False # too big -> fall back to the 2-kernel path
|
||||
|
||||
# Allocation mirrors moe_align_block_size.py (the kernel uses a +1 sentinel
|
||||
# bucket, so the padded buffers are sized with bucket_experts + 1).
|
||||
if numel < bucket_experts + 1:
|
||||
max_num_tokens_padded = numel * block_size
|
||||
else:
|
||||
max_num_tokens_padded = numel + (bucket_experts + 1) * (block_size - 1)
|
||||
max_num_m_blocks = (max_num_tokens_padded + block_size - 1) // block_size
|
||||
|
||||
# The align kernel's block-1 fill writes sorted_token_ids with vectorized int4
|
||||
# stores; the last store can spill up to 3 int32 past the logical end. Pad the
|
||||
# standalone allocation to a multiple of VEC_SIZE (4) so the spill stays in
|
||||
# bounds (matches _align_block_size_jit's _A4). block_size=16 is already a
|
||||
# multiple of 4 in production; this guards non-multiple-of-4 block sizes.
|
||||
sorted_alloc = (max_num_tokens_padded + 3) & ~3
|
||||
sorted_token_ids = torch.empty((sorted_alloc,), dtype=torch.int32, device=device)
|
||||
expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device=device)
|
||||
num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device=device)
|
||||
# No memset: the align kernel writes cumsum[0..num_buckets] before
|
||||
# count_and_sort reads it (same as moe_align_block_size.py's empty cumsum).
|
||||
cumsum_buffer = torch.empty((bucket_experts + 2,), dtype=torch.int32, device=device)
|
||||
token_lora_mask = torch.empty((M,), dtype=torch.bool, device=device)
|
||||
|
||||
module = _jit_module(flat_topk_ids.dtype)
|
||||
module.moe_lora_merged_align(
|
||||
flat_topk_ids,
|
||||
token_lora_mapping,
|
||||
token_lora_mask,
|
||||
bucket_experts + 1, # bucket count (the +1 sentinel)
|
||||
block_size,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_pad,
|
||||
cumsum_buffer,
|
||||
True, # pad_sorted_token_ids
|
||||
top_k,
|
||||
num_experts_for_weight,
|
||||
local_expert_offset,
|
||||
local_num_experts if local_num_experts is not None else 0,
|
||||
ep_local,
|
||||
shared_outer,
|
||||
do_skip,
|
||||
compact_eff,
|
||||
fuse_eff,
|
||||
)
|
||||
|
||||
return (
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_pad,
|
||||
token_lora_mask,
|
||||
virtual_num_experts,
|
||||
)
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Fused pack for the trtllm routed-MoE topk format.
|
||||
|
||||
The trtllm routed MoE consumes top-k routing as a single int32 per (token, slot):
|
||||
``PackedScoreIdx`` = ``(expert_id << 16) | bf16_weight_bits`` (little-endian: low 16
|
||||
bits = bf16 weight, high 16 bits = int16 expert id).
|
||||
|
||||
The torch reference builds this with a cluster of ~4 elementwise ops (cast, cast, view,
|
||||
bitshift, or). On the decode LoRA path these run on the main CUDA stream between
|
||||
``per_token_group_quant_fp8`` and the trtllm MoE op. This kernel collapses them into a
|
||||
single Triton launch. Bit-identical to the torch reference: weights are softmax/renormalized
|
||||
(>= 0 -> bf16 sign bit is 0, so the low 16 bits never collide with the id field and masking
|
||||
== torch's sign-extend-then-or); expert ids are small (< num_experts).
|
||||
"""
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _pack_topk_kernel(ids_ptr, w_ptr, out_ptr, numel, BLOCK: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
offs = pid * BLOCK + tl.arange(0, BLOCK)
|
||||
mask = offs < numel
|
||||
ids = tl.load(ids_ptr + offs, mask=mask) # int32
|
||||
w = tl.load(w_ptr + offs, mask=mask) # float32
|
||||
wb = w.to(tl.bfloat16)
|
||||
wbits = wb.to(tl.int16, bitcast=True).to(tl.int32) & 0xFFFF
|
||||
packed = (ids << 16) | wbits
|
||||
tl.store(out_ptr + offs, packed, mask=mask)
|
||||
|
||||
|
||||
def fused_pack_topk(topk_ids: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor:
|
||||
"""Single-launch replacement for the elementwise routed-MoE topk pack.
|
||||
|
||||
Returns int32 ``[*, top_k]`` packed tensor, bit-identical to the torch reference.
|
||||
"""
|
||||
if topk_ids.dtype != torch.int32:
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
topk_ids = topk_ids.contiguous()
|
||||
if topk_weights.dtype != torch.float32:
|
||||
topk_weights = topk_weights.to(torch.float32)
|
||||
topk_weights = topk_weights.contiguous()
|
||||
out = torch.empty_like(topk_ids, dtype=torch.int32)
|
||||
numel = out.numel()
|
||||
if numel == 0:
|
||||
return out
|
||||
BLOCK = 1024
|
||||
grid = (triton.cdiv(numel, BLOCK),)
|
||||
_pack_topk_kernel[grid](topk_ids, topk_weights, out, numel, BLOCK=BLOCK)
|
||||
return out
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Fused top-k gating softmax with routed-pack output (JIT).
|
||||
|
||||
JIT port of sgl-kernel's AOT ``topk_softmax`` power-of-2 fast path
|
||||
(``topkGatingSoftmax``) extended with a third output: the FlashInfer routed-MoE
|
||||
packed format ``(topk_id << 16) | bf16_bits(topk_weight)`` computed in the
|
||||
kernel epilogue after renormalization — bit-identical to running the standalone
|
||||
``fused_pack_topk`` triton kernel on the post-processed topk_ids/topk_weights
|
||||
(including the ``_mask_topk_ids_padded_region`` id=-1 sentinel for rows at or
|
||||
beyond ``num_token_non_padded``). This removes the per-MoE-layer
|
||||
``_pack_topk_kernel`` launch from the decode critical path.
|
||||
|
||||
Scope (callers must fall back to the AOT ``topk_softmax`` + separate pack
|
||||
otherwise): power-of-2 ``num_experts`` in [1, 512]; no softcapping or
|
||||
correction bias (the Qwen3-MoE softmax routing uses neither).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_topk_softmax_pack_module() -> Module:
|
||||
return load_jit(
|
||||
"topk_softmax_pack",
|
||||
cuda_files=["trtllm_lora_temp/topk_softmax_pack.cuh"],
|
||||
cuda_wrappers=[("topk_softmax_pack", "topk_softmax_pack")],
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["topk_weights", "topk_indices", "packed"])
|
||||
def _jit_topk_softmax_pack_op(
|
||||
topk_weights: torch.Tensor,
|
||||
topk_indices: torch.Tensor,
|
||||
packed: torch.Tensor,
|
||||
gating_output: torch.Tensor,
|
||||
num_token_non_padded: Optional[torch.Tensor],
|
||||
renormalize: bool,
|
||||
) -> None:
|
||||
module = _jit_topk_softmax_pack_module()
|
||||
module.topk_softmax_pack(
|
||||
topk_weights,
|
||||
topk_indices,
|
||||
packed,
|
||||
gating_output,
|
||||
num_token_non_padded,
|
||||
renormalize,
|
||||
)
|
||||
|
||||
|
||||
def topk_softmax_pack(
|
||||
topk_weights: torch.Tensor,
|
||||
topk_indices: torch.Tensor,
|
||||
packed: torch.Tensor,
|
||||
gating_output: torch.Tensor,
|
||||
renormalize: bool,
|
||||
num_token_non_padded: Optional[torch.Tensor] = None,
|
||||
) -> None:
|
||||
"""Drop-in for the AOT ``topk_softmax`` that ALSO writes ``packed``.
|
||||
|
||||
``packed`` is int32 ``[num_tokens, topk]``, ``(id << 16) | bf16_bits(w)``
|
||||
with the final (renormalized) weights; rows >= ``num_token_non_padded``
|
||||
pack id = -1 (the padded-region sentinel). ``topk_weights``/``topk_indices``
|
||||
are written exactly like the AOT kernel (indices NOT masked here — the
|
||||
regular python post-process handles them).
|
||||
"""
|
||||
assert gating_output.dim() == 2
|
||||
num_experts = gating_output.shape[-1]
|
||||
assert num_experts & (num_experts - 1) == 0 and num_experts <= 512, (
|
||||
"topk_softmax_pack supports power-of-2 num_experts in [1, 512] only; "
|
||||
"fall back to topk_softmax + fused_pack_topk"
|
||||
)
|
||||
if gating_output.shape[0] == 0:
|
||||
return
|
||||
_jit_topk_softmax_pack_op(
|
||||
topk_weights,
|
||||
topk_indices,
|
||||
packed,
|
||||
gating_output.contiguous(),
|
||||
num_token_non_padded,
|
||||
renormalize,
|
||||
)
|
||||
@@ -421,6 +421,9 @@ class Envs:
|
||||
SGLANG_NPU_FORWARD_NATIVE_GEMMA_RMS_NORM = EnvBool(False)
|
||||
# Delay all-gather after qlora for better performance for Deepseek v3.2
|
||||
SGLANG_USE_AG_AFTER_QLORA = EnvBool(False)
|
||||
# Master switch for the experimental TRT-LLM LoRA fast path; when OFF (default) every
|
||||
# fine-grained opt switch reads False, keeping non-experimental paths byte-identical.
|
||||
SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False)
|
||||
# Quantize x to int8 in the dispatch operator
|
||||
DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False) # This argument is deprecated
|
||||
SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1)
|
||||
|
||||
@@ -38,6 +38,8 @@ from sglang.srt.utils.common import (
|
||||
next_power_of_2,
|
||||
)
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
logger = __import__("logging").getLogger(__name__)
|
||||
|
||||
|
||||
@@ -1216,3 +1218,9 @@ def fused_experts_none_to_flashinfer_trtllm_routed(
|
||||
raise TypeError(
|
||||
f"Unexpected quant_info type for flashinfer_trtllm_routed: {type(quant_info)}"
|
||||
)
|
||||
|
||||
|
||||
# Register the experimental experimental_sgl_trtllm MoE fused-func (MoeRunner needs it at
|
||||
# build time even for LoRA); gated by the master switch so the upstream path is untouched.
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp import sgl_backend # noqa: E402,F401
|
||||
|
||||
@@ -379,7 +379,9 @@ def fused_moe_kernel(
|
||||
filter_expert: tl.constexpr,
|
||||
swap_ab: tl.constexpr,
|
||||
FUSE_ADD_TO_OUTPUT: tl.constexpr,
|
||||
MASK_OUTPUT: tl.constexpr,
|
||||
FUSE_SUM_ALL_REDUCE: tl.constexpr,
|
||||
LORA_PRESERVE_BASE: tl.constexpr,
|
||||
ROUTER_TOPK: tl.constexpr,
|
||||
):
|
||||
"""
|
||||
@@ -440,11 +442,9 @@ def fused_moe_kernel(
|
||||
off_experts = off_experts_i32.to(tl.int64)
|
||||
|
||||
if filter_expert and off_experts == -1:
|
||||
# -----------------------------------------------------------
|
||||
# Write back zeros to the output when the expert is not
|
||||
# in the current expert parallel rank.
|
||||
if not FUSE_ADD_TO_OUTPUT:
|
||||
# skip the zero-write to preserve existing values.
|
||||
if not FUSE_ADD_TO_OUTPUT and not (FUSE_SUM_ALL_REDUCE and LORA_PRESERVE_BASE):
|
||||
# Write zeros only when this kernel owns the full output; the experimental LoRA
|
||||
# add path (LORA_PRESERVE_BASE) keeps the base output from the prior MoE kernel.
|
||||
write_zeros_to_output(
|
||||
c_ptr,
|
||||
stride_cm,
|
||||
@@ -616,6 +616,18 @@ def fused_moe_kernel(
|
||||
c_mask = token_mask[:, None] & add_mask[:, None] & (offs_cn[None, :] < N)
|
||||
existing = tl.load(c_ptrs, mask=c_mask, other=0.0)
|
||||
tl.store(c_ptrs, existing + accumulator, mask=c_mask)
|
||||
# ===== TO BE REFACTORED ====
|
||||
elif MASK_OUTPUT:
|
||||
# Store a fresh output while zeroing rows whose request has no active LoRA.
|
||||
offs_token_out = offs_token // ROUTER_TOPK
|
||||
output_mask = tl.load(
|
||||
add_mask_ptr + offs_token_out, mask=token_mask, other=False
|
||||
)
|
||||
c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
|
||||
c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
|
||||
accumulator = tl.where(output_mask[:, None], accumulator, 0.0)
|
||||
tl.store(c_ptrs, accumulator, mask=c_mask)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
elif FUSE_SUM_ALL_REDUCE:
|
||||
offs_token_out = offs_token // ROUTER_TOPK
|
||||
c_ptrs = (
|
||||
@@ -731,6 +743,8 @@ def invoke_fused_moe_kernel(
|
||||
router_topk: int = 1,
|
||||
fuse_add_to_output: bool = False,
|
||||
add_output_mask: Optional[torch.Tensor] = None,
|
||||
mask_output: bool = False,
|
||||
lora_preserve_base: bool = False,
|
||||
) -> None:
|
||||
assert topk_weights.stride(1) == 1
|
||||
assert sorted_token_ids.stride(0) == 1
|
||||
@@ -807,6 +821,18 @@ def invoke_fused_moe_kernel(
|
||||
assert (
|
||||
add_output_mask is not None
|
||||
), "add_output_mask required when fuse_add_to_output=True"
|
||||
# ===== TO BE REFACTORED ====
|
||||
if mask_output:
|
||||
assert (
|
||||
not fuse_add_to_output
|
||||
), "mask_output and fuse_add_to_output are mutually exclusive"
|
||||
assert (
|
||||
not fuse_sum_all_reduce
|
||||
), "mask_output and fuse_sum_all_reduce are mutually exclusive"
|
||||
assert (
|
||||
add_output_mask is not None
|
||||
), "add_output_mask required when mask_output=True"
|
||||
# ===== END TO BE REFACTORED ====
|
||||
|
||||
if (
|
||||
(use_int8_w8a16 or use_int4_w4a16)
|
||||
@@ -924,6 +950,8 @@ def invoke_fused_moe_kernel(
|
||||
filter_expert=filter_expert,
|
||||
swap_ab=swap_ab,
|
||||
FUSE_ADD_TO_OUTPUT=fuse_add_to_output,
|
||||
MASK_OUTPUT=mask_output,
|
||||
LORA_PRESERVE_BASE=lora_preserve_base,
|
||||
FUSE_SUM_ALL_REDUCE=fuse_sum_all_reduce,
|
||||
ROUTER_TOPK=router_topk,
|
||||
**config,
|
||||
|
||||
@@ -5,8 +5,11 @@ from typing import Tuple
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_musa, is_xpu
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
_is_hip = is_hip()
|
||||
_is_xpu = is_xpu()
|
||||
@@ -74,14 +77,37 @@ def moe_align_block_size(
|
||||
(num_experts + 2,), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
|
||||
sgl_moe_align_block_size(
|
||||
topk_ids,
|
||||
num_experts + 1,
|
||||
block_size,
|
||||
sorted_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_pad,
|
||||
cumsum_buffer,
|
||||
True,
|
||||
)
|
||||
# ===== TO BE REFACTORED ====
|
||||
use_jit_align = False
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
|
||||
use_jit_align = lora_envs.SGLANG_OPT_USE_JIT_KERNEL_MOE_ALIGN.get()
|
||||
if use_jit_align:
|
||||
from sglang.jit_kernel.moe_align import (
|
||||
moe_align_block_size as jit_moe_align_block_size,
|
||||
)
|
||||
|
||||
jit_moe_align_block_size(
|
||||
topk_ids,
|
||||
num_experts + 1,
|
||||
block_size,
|
||||
sorted_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_pad,
|
||||
cumsum_buffer,
|
||||
True,
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
else:
|
||||
sgl_moe_align_block_size(
|
||||
topk_ids,
|
||||
num_experts + 1,
|
||||
block_size,
|
||||
sorted_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_pad,
|
||||
cumsum_buffer,
|
||||
True,
|
||||
)
|
||||
return sorted_ids, expert_ids, num_tokens_post_pad
|
||||
|
||||
@@ -100,6 +100,7 @@ class StandardDispatcher(BaseDispatcher):
|
||||
backend.is_flashinfer_cutlass()
|
||||
or backend.is_flashinfer_cutedsl()
|
||||
or backend.is_flashinfer_trtllm()
|
||||
or backend.is_experimental_sgl_trtllm()
|
||||
or backend.is_flashinfer_trtllm_routed()
|
||||
or self.enable_flashinfer_mxfp4_moe
|
||||
)
|
||||
|
||||
@@ -113,6 +113,8 @@ from sglang.srt.utils import (
|
||||
)
|
||||
from sglang.srt.utils.patch_torch import register_fake_if_exists
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.quantization import QuantizationConfig
|
||||
|
||||
@@ -220,6 +222,13 @@ class TopKOutputChecker:
|
||||
|
||||
@staticmethod
|
||||
def format_is_standard(topk_output: TopKOutput) -> TypeGuard[StandardTopKOutput]:
|
||||
# ===== TO BE REFACTORED ====
|
||||
# The experimental fused topk+pack carrier only exists under the master switch.
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
return isinstance(
|
||||
topk_output, (StandardTopKOutput, StandardTopKOutputPacked)
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
return isinstance(topk_output, StandardTopKOutput)
|
||||
|
||||
@staticmethod
|
||||
@@ -261,6 +270,25 @@ class StandardTopKOutput(NamedTuple):
|
||||
return TopKOutputFormat.STANDARD
|
||||
|
||||
|
||||
# ===== TO BE REFACTORED ====
|
||||
# Experimental fused topk+pack (SGLANG_OPT_LORA_FUSED_TOPK_PACK) carrier: the FlashInfer
|
||||
# routed-MoE packed topk produced fused in the gating kernel. Kept a SEPARATE type rather
|
||||
# than a 4th StandardTopKOutput field so the OSS `a, b, _ = topk_output` 3-tuple unpack
|
||||
# stays valid; only the gated experimental MoE dispatch reads .packed_topk_ids (getattr).
|
||||
class StandardTopKOutputPacked(NamedTuple):
|
||||
topk_weights: torch.Tensor
|
||||
topk_ids: torch.Tensor
|
||||
router_logits: torch.Tensor
|
||||
packed_topk_ids: torch.Tensor
|
||||
|
||||
@property
|
||||
def format(self) -> TopKOutputFormat:
|
||||
return TopKOutputFormat.STANDARD
|
||||
|
||||
|
||||
# ===== END TO BE REFACTORED ====
|
||||
|
||||
|
||||
class TritonKernelTopKOutput(NamedTuple):
|
||||
"""Triton kernel top-k output format."""
|
||||
|
||||
@@ -434,6 +462,20 @@ class TopK(MultiPlatformOp):
|
||||
output_format = self.topk_config.output_format
|
||||
elif get_moe_runner_backend().is_triton_kernels():
|
||||
output_format = TopKOutputFormat.TRITON_KERNEL
|
||||
# ===== TO BE REFACTORED ====
|
||||
elif get_moe_runner_backend().is_experimental_sgl_trtllm():
|
||||
try:
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
|
||||
use_standard_for_lora = bool(get_global_server_args().enable_lora)
|
||||
except ValueError:
|
||||
use_standard_for_lora = False
|
||||
output_format = (
|
||||
TopKOutputFormat.STANDARD
|
||||
if use_standard_for_lora
|
||||
else TopKOutputFormat.BYPASSED
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
elif get_moe_runner_backend().is_flashinfer_trtllm() or (
|
||||
get_moe_runner_backend().is_flashinfer_mxfp4() and not self.is_fp4_experts
|
||||
):
|
||||
@@ -672,6 +714,8 @@ def fused_topk(
|
||||
renormalize: bool,
|
||||
correction_bias: Optional[torch.Tensor] = None,
|
||||
scoring_func: str = "softmax",
|
||||
packed_out: Optional[torch.Tensor] = None,
|
||||
num_token_non_padded: Optional[torch.Tensor] = None,
|
||||
):
|
||||
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
||||
|
||||
@@ -694,6 +738,23 @@ def fused_topk(
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
)
|
||||
# ===== TO BE REFACTORED ====
|
||||
elif packed_out is not None:
|
||||
# Fused gating + routed pack (SGLANG_OPT_LORA_FUSED_TOPK_PACK): one JIT kernel
|
||||
# writes topk_weights/topk_ids AND the FlashInfer packed topk in one launch.
|
||||
from sglang.jit_kernel.trtllm_lora_temp.topk_softmax_pack import (
|
||||
topk_softmax_pack,
|
||||
)
|
||||
|
||||
topk_softmax_pack(
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
packed_out,
|
||||
gating_output,
|
||||
renormalize,
|
||||
num_token_non_padded=num_token_non_padded,
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
else:
|
||||
topk_softmax(
|
||||
topk_weights,
|
||||
@@ -1206,6 +1267,30 @@ def biased_grouped_topk_gpu(
|
||||
# Use optimized path for Kimi K2 (384 experts with num_expert_group=1)
|
||||
num_experts = gating_output.shape[1]
|
||||
if _is_cuda and num_experts == 384 and num_expert_group == 1:
|
||||
# ===== TO BE REFACTORED ====
|
||||
_use_jit_bf16_gate = False
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
|
||||
_use_jit_bf16_gate = (
|
||||
lora_envs.SGLANG_OPT_USE_JIT_KERNEL_KIMI_GATE.get()
|
||||
and lora_envs.SGLANG_OPT_KIMI_GATE_BF16_INPUT.get()
|
||||
)
|
||||
if _use_jit_bf16_gate:
|
||||
from sglang.jit_kernel.trtllm_lora_temp.kimi_k2_moe_fused_gate import (
|
||||
kimi_k2_moe_fused_gate as _kimi_k2_moe_fused_gate,
|
||||
)
|
||||
|
||||
# bf16 pass-through: skip the two host-side fp32 upcast kernels.
|
||||
return _kimi_k2_moe_fused_gate(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk=topk,
|
||||
renormalize=renormalize,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
return kimi_k2_moe_fused_gate(
|
||||
gating_output.to(dtype=torch.float32),
|
||||
correction_bias,
|
||||
@@ -1471,6 +1556,9 @@ def select_experts(
|
||||
|
||||
scoring_func = topk_config.scoring_func
|
||||
|
||||
# Set by the fused-gating+pack branch below; None everywhere else.
|
||||
packed_topk = None
|
||||
|
||||
(
|
||||
router_logits,
|
||||
correction_bias,
|
||||
@@ -1562,7 +1650,43 @@ def select_experts(
|
||||
renormalize=renormalize,
|
||||
)
|
||||
else:
|
||||
# Fused gating + routed pack (SGLANG_OPT_LORA_FUSED_TOPK_PACK): only on the plain
|
||||
# CUDA softmax path with no EPLB remap / shared experts / bias / routing overrides.
|
||||
_fused_topk_pack = False
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
|
||||
_fused_topk_pack = lora_envs.SGLANG_OPT_LORA_FUSED_TOPK_PACK.get()
|
||||
if (
|
||||
_fused_topk_pack
|
||||
and _is_cuda
|
||||
and not _use_aiter
|
||||
and scoring_func == "softmax"
|
||||
and correction_bias is None
|
||||
and expert_location_dispatch_info is None
|
||||
and num_fused_shared_experts == 0
|
||||
and not envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get()
|
||||
and not envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get()
|
||||
):
|
||||
num_experts = router_logits.shape[-1]
|
||||
if num_experts & (num_experts - 1) == 0 and num_experts <= 512:
|
||||
packed_topk = torch.empty(
|
||||
(hidden_states.shape[0], top_k),
|
||||
dtype=torch.int32,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
# Qwen3MOE uses fused_topk
|
||||
_fused_topk_kwargs = {}
|
||||
# ===== TO BE REFACTORED ====
|
||||
# Only the experimental fused topk+pack passes packed_out/num_token_non_padded;
|
||||
# the default call keeps the upstream signature (fused_topk_cpu lacks these).
|
||||
if packed_topk is not None:
|
||||
_fused_topk_kwargs = dict(
|
||||
packed_out=packed_topk,
|
||||
num_token_non_padded=num_token_non_padded,
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
topk_weights, topk_ids = fused_topk(
|
||||
hidden_states=hidden_states,
|
||||
gating_output=router_logits,
|
||||
@@ -1570,6 +1694,7 @@ def select_experts(
|
||||
renormalize=renormalize,
|
||||
correction_bias=correction_bias,
|
||||
scoring_func=scoring_func,
|
||||
**_fused_topk_kwargs,
|
||||
)
|
||||
else:
|
||||
assert (
|
||||
@@ -1633,6 +1758,12 @@ def select_experts(
|
||||
|
||||
get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
|
||||
|
||||
# ===== TO BE REFACTORED ====
|
||||
if packed_topk is not None:
|
||||
return StandardTopKOutputPacked(
|
||||
topk_weights, topk_ids, router_logits, packed_topk
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
return StandardTopKOutput(topk_weights, topk_ids, router_logits)
|
||||
|
||||
|
||||
|
||||
@@ -91,6 +91,7 @@ class MoeRunnerBackend(Enum):
|
||||
TRITON = "triton"
|
||||
TRITON_KERNELS = "triton_kernel"
|
||||
FLASHINFER_TRTLLM = "flashinfer_trtllm"
|
||||
EXPERIMENTAL_SGL_TRTLLM = "experimental_sgl_trtllm"
|
||||
FLASHINFER_TRTLLM_ROUTED = "flashinfer_trtllm_routed"
|
||||
FLASHINFER_CUTLASS = "flashinfer_cutlass"
|
||||
FLASHINFER_MXFP4 = "flashinfer_mxfp4"
|
||||
@@ -112,7 +113,15 @@ class MoeRunnerBackend(Enum):
|
||||
return self == MoeRunnerBackend.TRITON_KERNELS
|
||||
|
||||
def is_flashinfer_trtllm(self):
|
||||
return self == MoeRunnerBackend.FLASHINFER_TRTLLM
|
||||
# experimental_sgl_trtllm shares the TRT-LLM FP8 kernels + layout, so it inherits
|
||||
# trtllm weight-prep here; divergent sites check is_experimental_sgl_trtllm() first.
|
||||
return self in (
|
||||
MoeRunnerBackend.FLASHINFER_TRTLLM,
|
||||
MoeRunnerBackend.EXPERIMENTAL_SGL_TRTLLM,
|
||||
)
|
||||
|
||||
def is_experimental_sgl_trtllm(self):
|
||||
return self == MoeRunnerBackend.EXPERIMENTAL_SGL_TRTLLM
|
||||
|
||||
def is_flashinfer_trtllm_routed(self):
|
||||
return self == MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED
|
||||
|
||||
@@ -10,6 +10,7 @@ from sglang.srt.distributed import (
|
||||
tensor_model_parallel_all_gather,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
@@ -26,6 +27,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
|
||||
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
|
||||
from sglang.srt.lora.utils import LoRABatchInfo, get_lm_head_lora_b_shard_size
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
|
||||
class BaseLayerWithLoRA(nn.Module):
|
||||
def __init__(
|
||||
@@ -882,6 +885,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
|
||||
self.lora_use_virtual_experts: bool = False
|
||||
self.quant_method = base_layer.quant_method
|
||||
self.moe_runner_config = base_layer.moe_runner_config
|
||||
self.dispatcher = base_layer.dispatcher
|
||||
self.num_local_experts = base_layer.num_local_experts
|
||||
self.should_fuse_routed_scaling_factor_in_topk = (
|
||||
base_layer.should_fuse_routed_scaling_factor_in_topk
|
||||
)
|
||||
|
||||
self.tp_size = getattr(base_layer, "moe_tp_size", 1)
|
||||
self.tp_rank = getattr(base_layer, "moe_tp_rank", 0)
|
||||
@@ -909,6 +917,17 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
|
||||
else:
|
||||
runner_backend = MoeRunnerBackend.TRITON
|
||||
|
||||
# ===== TO BE REFACTORED ====
|
||||
self._lora_runner_backend = runner_backend
|
||||
if runner_backend.is_experimental_sgl_trtllm():
|
||||
from sglang.srt.lora.trtllm_lora_temp.lora_layer import (
|
||||
init_experimental_sgl_trtllm_lora,
|
||||
)
|
||||
|
||||
init_experimental_sgl_trtllm_lora(self, base_layer)
|
||||
return
|
||||
# ===== END TO BE REFACTORED ====
|
||||
|
||||
self._lora_runner = MoeRunner(
|
||||
runner_backend,
|
||||
base_layer.moe_runner_config,
|
||||
@@ -965,6 +984,17 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
|
||||
# the Python weight_indices list, no GPU sync needed.
|
||||
has_active_lora = bool(getattr(batch_info, "has_active_lora", False))
|
||||
|
||||
if self._lora_runner_backend.is_experimental_sgl_trtllm():
|
||||
# Per-rank (local) expert count the LoRA buffers are indexed by, so
|
||||
# virtual-experts indexing matches the buffers under EP.
|
||||
num_experts = (
|
||||
self.down_lora_a_weights.shape[1]
|
||||
if self.down_lora_a_weights is not None
|
||||
else self.base_layer.num_local_experts
|
||||
)
|
||||
else:
|
||||
num_experts = self.base_layer.num_experts
|
||||
|
||||
return LoRAInfo(
|
||||
gate_up_lora_a_weights=self.gate_up_lora_a_weights,
|
||||
gate_up_lora_b_weights=self.gate_up_lora_b_weights,
|
||||
@@ -976,7 +1006,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
|
||||
adapter_enabled=moe_lora_info.adapter_enabled,
|
||||
token_lora_mapping=moe_lora_info.token_lora_mapping,
|
||||
max_lora_rank=max_lora_rank,
|
||||
num_experts=self.base_layer.num_experts,
|
||||
num_experts=num_experts,
|
||||
has_active_lora=has_active_lora,
|
||||
experts_shared_outer_loras=self.experts_shared_outer_loras,
|
||||
cg_buffers=cg_buffers,
|
||||
@@ -1022,10 +1052,20 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
|
||||
# Use pre-computed quant info (doesn't change so not sure why we need to pass it in every time)
|
||||
quant_info = self._quant_info
|
||||
|
||||
# Run the only lora moe runner (Triton)
|
||||
combine_input = self._lora_runner.run(
|
||||
dispatch_output, quant_info, lora_info=lora_info
|
||||
)
|
||||
# ===== TO BE REFACTORED ====
|
||||
if self._lora_runner_backend.is_experimental_sgl_trtllm():
|
||||
from sglang.srt.lora.trtllm_lora_temp.lora_layer import (
|
||||
dispatch_experimental_sgl_trtllm_lora,
|
||||
)
|
||||
|
||||
combine_input = dispatch_experimental_sgl_trtllm_lora(
|
||||
dispatch_output, quant_info, base_layer, lora_info
|
||||
)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
else:
|
||||
combine_input = self._lora_runner.run(
|
||||
dispatch_output, quant_info, lora_info=lora_info
|
||||
)
|
||||
|
||||
final_hidden_states = base_layer.dispatcher.combine(combine_input=combine_input)
|
||||
|
||||
@@ -1157,3 +1197,14 @@ def get_lora_layer(
|
||||
ret = lora_layer_type(layer, lora_backend)
|
||||
return ret
|
||||
raise Exception(f"No corresponding LoRA layer supported for {type(layer)}.")
|
||||
|
||||
|
||||
# ===== TO BE REFACTORED ====
|
||||
# Experimental two-stream LoRA overlap; installed only under the master switch, else no-op.
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp import ( # noqa: E402
|
||||
install_two_stream_overrides as _install_lora_two_stream,
|
||||
)
|
||||
|
||||
_install_lora_two_stream()
|
||||
# ===== END TO BE REFACTORED ====
|
||||
|
||||
@@ -21,6 +21,7 @@ from typing import Dict, Iterable, List, Optional
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.load_config import LoadConfig
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
|
||||
from sglang.srt.layers.utils import get_layer_id
|
||||
from sglang.srt.layers.vocab_parallel_embedding import (
|
||||
@@ -47,6 +48,8 @@ from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import replace_submodule
|
||||
from sglang.srt.utils.hf_transformers_utils import AutoConfig
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -121,6 +124,17 @@ class LoRAManager:
|
||||
num_tokens_per_bs=num_tokens_per_bs,
|
||||
)
|
||||
|
||||
# ===== TO BE REFACTORED ====
|
||||
# Pre-create the experimental LoRA two-stream side stream now (gated) so the
|
||||
# torch.cuda.Stream() call never lands inside a cuda-graph capture region.
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp import (
|
||||
init_lora_two_stream_resources,
|
||||
)
|
||||
|
||||
init_lora_two_stream_resources(self.device)
|
||||
# ===== END TO BE REFACTORED ====
|
||||
|
||||
def init_cuda_graph_moe_buffers(
|
||||
self, max_bs: int, max_loras: int, compute_dtype, moe_layer
|
||||
):
|
||||
|
||||
@@ -23,6 +23,7 @@ from sglang.srt.distributed import (
|
||||
get_moe_tensor_parallel_world_size,
|
||||
get_pp_group,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.lora.eviction_policy import get_eviction_policy
|
||||
from sglang.srt.lora.layers import BaseLayerWithLoRA
|
||||
from sglang.srt.lora.lora import LoRAAdapter
|
||||
@@ -43,6 +44,8 @@ from sglang.srt.lora.utils import (
|
||||
from sglang.srt.utils import is_pin_memory_available
|
||||
from sglang.srt.utils.hf_transformers_utils import AutoConfig
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -119,6 +122,7 @@ def _moe_runner_keeps_global_expert_ids() -> bool:
|
||||
return (
|
||||
b.is_flashinfer_cutlass()
|
||||
or b.is_flashinfer_cutedsl()
|
||||
or b.is_experimental_sgl_trtllm()
|
||||
or b.is_flashinfer_trtllm_routed()
|
||||
)
|
||||
except Exception: # pragma: no cover - backend not initialized
|
||||
@@ -1179,6 +1183,10 @@ class LoRAMemoryPool:
|
||||
weights,
|
||||
)
|
||||
load_lora_weight_tensor(buffer_view, weights)
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
# Zero beyond loaded rank: the experimental dense LoRA-B kernel
|
||||
# contracts over the full padded max_rank, so the tail must be clean.
|
||||
target_buffer[buffer_id, :, lora_rank:].zero_()
|
||||
|
||||
if lora_adapter.embedding_layers:
|
||||
org_vocab_size = self.base_hf_config.vocab_size
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Two-stream LoRA overlap (O1 + O7 + O8 + O9) — installed as a monkey-patch.
|
||||
|
||||
Activates when env ``SGLANG_LORA_TWO_STREAM=1``. Triggered exactly once via
|
||||
:func:`install_two_stream_overrides` (called at end of ``sglang/srt/lora/layers.py``).
|
||||
|
||||
When enabled, these call sites are redirected to side-stream-overlapped versions
|
||||
defined entirely in this package:
|
||||
|
||||
* ``QKVParallelLinearWithLoRA.forward`` → :mod:`.attention.qkv_proj_lora_forward`
|
||||
* ``RowParallelLinearWithLoRA.forward`` → :mod:`.attention.row_parallel_lora_forward`
|
||||
* ``MergedColumnParallelLinearWithLoRA.forward`` → :mod:`.merged_column.merged_column_lora_forward`
|
||||
* ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora`` →
|
||||
:mod:`.moe_overlap.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream`
|
||||
|
||||
When disabled (env unset), ``install_two_stream_overrides`` is a no-op and all
|
||||
the original functions / methods in ``sglang/srt/lora/layers.py`` and
|
||||
``sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py`` run unchanged.
|
||||
|
||||
Per-batch gating still happens inside the patched callables — they fall back
|
||||
to the saved-original implementation for non-decode batches (token count above
|
||||
``SGLANG_TWO_STREAM_MAX_TOKENS`` default 256), so prefill stays on the serial
|
||||
path even with the patch installed.
|
||||
"""
|
||||
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
|
||||
|
||||
def is_two_stream_active(x: torch.Tensor) -> bool:
|
||||
"""Per-batch gate (two-stream now always-on). True iff batch is decode-shaped (<= SGLANG_TWO_STREAM_MAX_TOKENS)."""
|
||||
return x.shape[0] <= lora_envs.SGLANG_TWO_STREAM_MAX_TOKENS.get()
|
||||
|
||||
|
||||
_LORA_SIDE_STREAM: Optional[torch.cuda.Stream] = None
|
||||
|
||||
|
||||
def get_lora_side_stream() -> torch.cuda.Stream:
|
||||
"""Lazily allocate a single shared LoRA side stream.
|
||||
|
||||
Within one decode layer the three sites (qkv → attn → o_proj → moe_gate_up)
|
||||
run sequentially, so one stream suffices and avoids extra graph-capture
|
||||
nodes from per-site streams.
|
||||
"""
|
||||
global _LORA_SIDE_STREAM
|
||||
if _LORA_SIDE_STREAM is None:
|
||||
_LORA_SIDE_STREAM = torch.cuda.Stream()
|
||||
return _LORA_SIDE_STREAM
|
||||
|
||||
|
||||
def init_lora_two_stream_resources(device: Optional[torch.device] = None) -> None:
|
||||
"""Eagerly create the side stream before cuda-graph capture begins.
|
||||
|
||||
``torch.cuda.Stream()`` is a driver call that must not run inside a
|
||||
cuda-graph capture region. Since :func:`get_lora_side_stream` is otherwise
|
||||
lazy, the first eligible decode forward would create it — which can fall
|
||||
inside capture if warmup didn't happen to exercise a two-stream batch.
|
||||
Calling this from a pre-capture hook pins creation to init/warmup on the
|
||||
correct device.
|
||||
"""
|
||||
if device is not None:
|
||||
with torch.cuda.device(device):
|
||||
get_lora_side_stream()
|
||||
else:
|
||||
get_lora_side_stream()
|
||||
|
||||
|
||||
def lora_overlap_alloc_stream() -> Optional[torch.cuda.Stream]:
|
||||
"""Stream to allocate side-stream LoRA-shrink OUTPUT buffers on, or None for default behavior.
|
||||
|
||||
A buffer allocated *inside* ``with torch.cuda.stream(side)`` is tagged to the side stream, so the
|
||||
caching allocator may free/reuse it on the side stream's schedule — before the MAIN stream (the real
|
||||
consumer, via the LoRA-B expand) is done. Under cuda-graph replay that's a premature-reuse WAR ->
|
||||
qwen3.5 mamba decode garbage. With ``SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC`` this returns the MAIN
|
||||
stream so the op allocates the output on the consumer stream (like the MoE O1 ``gate_up_delta``),
|
||||
making a single shared side stream graph-safe. Call on the MAIN stream BEFORE forking to the side.
|
||||
"""
|
||||
if lora_envs.SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC.get():
|
||||
return torch.cuda.current_stream()
|
||||
return None
|
||||
|
||||
|
||||
# References to the original implementations, captured at install time so the
|
||||
# patched callables can defer to them for non-decode batches.
|
||||
_ORIGINAL_QKV_FORWARD: Optional[Callable] = None
|
||||
_ORIGINAL_ROW_FORWARD: Optional[Callable] = None
|
||||
_ORIGINAL_MERGED_FORWARD: Optional[Callable] = None
|
||||
_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
|
||||
_INSTALLED: bool = False
|
||||
|
||||
|
||||
def get_original_qkv_forward() -> Callable:
|
||||
return _ORIGINAL_QKV_FORWARD
|
||||
|
||||
|
||||
def get_original_row_forward() -> Callable:
|
||||
return _ORIGINAL_ROW_FORWARD
|
||||
|
||||
|
||||
def get_original_merged_column_forward() -> Callable:
|
||||
return _ORIGINAL_MERGED_FORWARD
|
||||
|
||||
|
||||
def get_original_column_forward() -> Callable:
|
||||
return _ORIGINAL_COLUMN_FORWARD
|
||||
|
||||
|
||||
def get_original_replicated_forward() -> Callable:
|
||||
return _ORIGINAL_REPLICATED_FORWARD
|
||||
|
||||
|
||||
def get_original_moe_lora_func() -> Callable:
|
||||
return _ORIGINAL_MOE_LORA_FUNC
|
||||
|
||||
|
||||
def get_original_fp4_moe_lora_func() -> Callable:
|
||||
return _ORIGINAL_FP4_MOE_LORA_FUNC
|
||||
|
||||
|
||||
def install_two_stream_overrides() -> None:
|
||||
"""Install the side-stream overlapped overrides if ``SGLANG_LORA_TWO_STREAM=1``.
|
||||
|
||||
Idempotent: subsequent calls are a no-op. Patches:
|
||||
|
||||
1. ``QKVParallelLinearWithLoRA.forward`` (O7 — qkv LoRA shrink overlap)
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
if _INSTALLED:
|
||||
return
|
||||
|
||||
from sglang.srt.lora.layers import (
|
||||
ColumnParallelLinearWithLoRA,
|
||||
MergedColumnParallelLinearWithLoRA,
|
||||
QKVParallelLinearWithLoRA,
|
||||
ReplicatedLinearWithLoRA,
|
||||
RowParallelLinearWithLoRA,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp.attention import (
|
||||
column_parallel_lora_forward,
|
||||
qkv_proj_lora_forward,
|
||||
replicated_lora_forward,
|
||||
row_parallel_lora_forward,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp.merged_column import (
|
||||
merged_column_lora_forward,
|
||||
)
|
||||
|
||||
# Capture all originals before patching: QKV / MergedColumn subclass
|
||||
# ColumnParallel, so the plain-Column O10 patch must not clobber the
|
||||
# subclasses' own (3-/2-slice) forwards captured here as their fallbacks.
|
||||
_ORIGINAL_QKV_FORWARD = QKVParallelLinearWithLoRA.forward
|
||||
_ORIGINAL_ROW_FORWARD = RowParallelLinearWithLoRA.forward
|
||||
_ORIGINAL_MERGED_FORWARD = MergedColumnParallelLinearWithLoRA.forward
|
||||
_ORIGINAL_COLUMN_FORWARD = ColumnParallelLinearWithLoRA.forward
|
||||
_ORIGINAL_REPLICATED_FORWARD = ReplicatedLinearWithLoRA.forward
|
||||
QKVParallelLinearWithLoRA.forward = qkv_proj_lora_forward
|
||||
RowParallelLinearWithLoRA.forward = row_parallel_lora_forward
|
||||
MergedColumnParallelLinearWithLoRA.forward = merged_column_lora_forward
|
||||
# O10 (MLA q_b_proj / kv_b_proj) + O11 (MLA fused_qkv_a_proj_with_mqa).
|
||||
ColumnParallelLinearWithLoRA.forward = column_parallel_lora_forward
|
||||
ReplicatedLinearWithLoRA.forward = replicated_lora_forward
|
||||
|
||||
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_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.
|
||||
_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
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
_INSTALLED = True
|
||||
|
||||
|
||||
__all__ = [
|
||||
"is_two_stream_active",
|
||||
"get_lora_side_stream",
|
||||
"init_lora_two_stream_resources",
|
||||
"get_original_qkv_forward",
|
||||
"get_original_row_forward",
|
||||
"get_original_merged_column_forward",
|
||||
"get_original_column_forward",
|
||||
"get_original_replicated_forward",
|
||||
"get_original_moe_lora_func",
|
||||
"get_original_fp4_moe_lora_func",
|
||||
"install_two_stream_overrides",
|
||||
]
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Two-stream attention LoRA forward implementations (O7 + O8).
|
||||
|
||||
These are monkey-patched onto :class:`QKVParallelLinearWithLoRA` and
|
||||
:class:`RowParallelLinearWithLoRA` by
|
||||
:func:`sglang.srt.lora.trtllm_lora_temp.install_two_stream_overrides` when
|
||||
``SGLANG_LORA_TWO_STREAM=1``. The saved-original forward methods are
|
||||
preserved and called for batches where two-stream isn't active.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
split_tensor_along_last_dim,
|
||||
tensor_model_parallel_all_gather,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp import (
|
||||
get_lora_side_stream,
|
||||
get_original_column_forward,
|
||||
get_original_qkv_forward,
|
||||
get_original_replicated_forward,
|
||||
get_original_row_forward,
|
||||
is_two_stream_active,
|
||||
lora_overlap_alloc_stream,
|
||||
)
|
||||
|
||||
|
||||
def qkv_proj_lora_forward(self, input_: torch.Tensor):
|
||||
"""O7 — side-stream LoRA-A shrink ‖ base qkv_proj GEMM.
|
||||
|
||||
The shrink reads ``input_`` and the LoRA-A weights — same input as the
|
||||
base GEMM, no write conflict. The expand needs the shrink intermediate
|
||||
AND base_output, so it runs after the rejoin on the main stream.
|
||||
"""
|
||||
if not self.set_lora or not is_two_stream_active(input_):
|
||||
return get_original_qkv_forward()(self, input_)
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
|
||||
qkv_lora_b_fwd,
|
||||
sgemm_lora_a_fwd,
|
||||
)
|
||||
|
||||
bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
|
||||
side_stream = get_lora_side_stream()
|
||||
# sgemm_info is host-side (LoRABatchInfo); compute once, share both calls.
|
||||
sgemm_info = self.lora_backend._sgemm_info()
|
||||
|
||||
_alloc = lora_overlap_alloc_stream() # capture MAIN stream here (before the fork)
|
||||
side_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(side_stream):
|
||||
shrink_intermediate = sgemm_lora_a_fwd(
|
||||
input_, self.A_buffer_qkv, sgemm_info, stack_num=3, out_alloc_stream=_alloc
|
||||
)
|
||||
|
||||
# Base qkv_proj GEMM on main, concurrent with the side-stream shrink.
|
||||
output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_, bias)
|
||||
|
||||
# Rejoin: expand reads both side-produced shrink_intermediate and base_output.
|
||||
torch.cuda.current_stream().wait_stream(side_stream)
|
||||
output_parallel = qkv_lora_b_fwd(
|
||||
shrink_intermediate,
|
||||
self.B_buffer_qkv,
|
||||
sgemm_info,
|
||||
self.output_offset,
|
||||
self.max_qkv_out_dim,
|
||||
output_parallel,
|
||||
n_slices=3,
|
||||
)
|
||||
|
||||
if self.base_layer.gather_output:
|
||||
output = tensor_model_parallel_all_gather(output_parallel)
|
||||
else:
|
||||
output = output_parallel
|
||||
output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
|
||||
return output, output_bias
|
||||
|
||||
|
||||
def row_parallel_lora_forward(
|
||||
self, input_: torch.Tensor, skip_all_reduce: bool = False, forward_batch=None
|
||||
):
|
||||
"""O8 — side-stream LoRA-A shrink ‖ base row-parallel (o_proj) GEMM.
|
||||
|
||||
Mirrors O7 but the row-parallel context adds: input split per TP rank
|
||||
(when not already parallel), bias on rank 0 only, optional cross-rank
|
||||
all-reduce on both base output and lora_a intermediate when reducing.
|
||||
|
||||
Falls back to the saved-original :meth:`forward` for non-decode batches
|
||||
or when LoRA isn't set on this layer.
|
||||
"""
|
||||
# We need ``input_parallel`` to gate the per-batch decode check (its
|
||||
# token-count drives the threshold, not the unsplit ``input_``).
|
||||
if self.base_layer.input_is_parallel:
|
||||
input_parallel = input_
|
||||
else:
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
splitted_input = split_tensor_along_last_dim(
|
||||
input_, num_partitions=self.base_layer.tp_size
|
||||
)
|
||||
input_parallel = splitted_input[tp_rank].contiguous()
|
||||
|
||||
if not self.set_lora or not is_two_stream_active(input_parallel):
|
||||
return get_original_row_forward()(self, input_, skip_all_reduce, forward_batch)
|
||||
|
||||
bias_ = (
|
||||
None
|
||||
if (self.base_layer.tp_rank > 0 or self.base_layer.skip_bias_add)
|
||||
else self.base_layer.bias
|
||||
)
|
||||
|
||||
side_stream = get_lora_side_stream()
|
||||
_alloc = lora_overlap_alloc_stream() # capture MAIN stream here (before the fork)
|
||||
side_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(side_stream):
|
||||
lora_a_output = self.lora_backend.run_lora_a_sgemm(
|
||||
input_parallel, self.A_buffer, out_alloc_stream=_alloc
|
||||
)
|
||||
|
||||
# Base row-parallel GEMM on main, concurrent with the side-stream shrink.
|
||||
output_parallel = self.base_layer.quant_method.apply(
|
||||
self.base_layer, input_parallel, bias=bias_
|
||||
)
|
||||
|
||||
torch.cuda.current_stream().wait_stream(side_stream)
|
||||
|
||||
should_reduce = (
|
||||
self.base_layer.reduce_results
|
||||
and self.base_layer.tp_size > 1
|
||||
and not skip_all_reduce
|
||||
)
|
||||
|
||||
if should_reduce:
|
||||
output_ = tensor_model_parallel_all_reduce(output_parallel)
|
||||
lora_a_output = tensor_model_parallel_all_reduce(lora_a_output)
|
||||
output_ = self.lora_backend.run_lora_b_sgemm(
|
||||
x=lora_a_output,
|
||||
weights=self.B_buffer,
|
||||
output_offset=self.output_offset,
|
||||
output_offset_cpu=self.output_offset_cpu,
|
||||
base_output=output_,
|
||||
)
|
||||
else:
|
||||
# Two-stream already produced lora_a_output on the side stream; finish
|
||||
# the LoRA with just the expand atomic-add against output_parallel.
|
||||
output_parallel = self.lora_backend.run_lora_b_sgemm(
|
||||
x=lora_a_output,
|
||||
weights=self.B_buffer,
|
||||
output_offset=self.output_offset,
|
||||
output_offset_cpu=self.output_offset_cpu,
|
||||
base_output=output_parallel,
|
||||
)
|
||||
output_ = output_parallel
|
||||
|
||||
output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
|
||||
return output_, output_bias
|
||||
|
||||
|
||||
def column_parallel_lora_forward(self, input_: torch.Tensor):
|
||||
"""O10 — side-stream LoRA-A shrink ‖ base ColumnParallel GEMM.
|
||||
|
||||
Covers DeepSeek/Kimi MLA's ``q_b_proj`` / ``kv_b_proj`` (plain
|
||||
:class:`ColumnParallelLinearWithLoRA` — the merged/QKV subclasses keep their
|
||||
own O9/O7 overrides). The shrink reads ``input_`` (same as the base GEMM, no
|
||||
write conflict); the expand needs the shrink intermediate AND base_output,
|
||||
so it runs after the rejoin. Byte-identical to the saved-original forward
|
||||
for non-decode batches or when LoRA isn't set on this layer.
|
||||
"""
|
||||
if not self.set_lora or not is_two_stream_active(input_):
|
||||
return get_original_column_forward()(self, input_)
|
||||
|
||||
bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
|
||||
side_stream = get_lora_side_stream()
|
||||
|
||||
_alloc = lora_overlap_alloc_stream() # capture MAIN stream here (before the fork)
|
||||
side_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(side_stream):
|
||||
lora_a_output = self.lora_backend.run_lora_a_sgemm(
|
||||
input_, self.A_buffer, out_alloc_stream=_alloc
|
||||
)
|
||||
|
||||
# Base ColumnParallel GEMM on main, concurrent with the side-stream shrink.
|
||||
output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_, bias)
|
||||
|
||||
# Rejoin: expand reads both the side-produced shrink and base_output.
|
||||
torch.cuda.current_stream().wait_stream(side_stream)
|
||||
output_parallel = self.lora_backend.run_lora_b_sgemm(
|
||||
x=lora_a_output,
|
||||
weights=self.B_buffer,
|
||||
output_offset=self.output_offset,
|
||||
output_offset_cpu=self.output_offset_cpu,
|
||||
base_output=output_parallel,
|
||||
)
|
||||
|
||||
if self.base_layer.gather_output:
|
||||
output = tensor_model_parallel_all_gather(output_parallel)
|
||||
else:
|
||||
output = output_parallel
|
||||
output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
|
||||
return output, output_bias
|
||||
|
||||
|
||||
def replicated_lora_forward(self, x: torch.Tensor):
|
||||
"""O11 — side-stream LoRA-A shrink ‖ base ReplicatedLinear GEMM.
|
||||
|
||||
Covers DeepSeek/Kimi MLA's ``fused_qkv_a_proj_with_mqa``
|
||||
(:class:`ReplicatedLinearWithLoRA`, no TP sharding). Handles both the
|
||||
single-projection (``first_output_dim == 0``) and the fused q_a+kv_a
|
||||
(``> 0``, ``n_slices=2``) cases, splitting the same A-shrink / B-expand the
|
||||
backend's ``run_qkv_lora`` composes internally — A on the side stream, B on
|
||||
the main after the rejoin. Falls back to the saved-original otherwise.
|
||||
"""
|
||||
if not self.set_lora or not is_two_stream_active(x):
|
||||
return get_original_replicated_forward()(self, x)
|
||||
|
||||
bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
|
||||
side_stream = get_lora_side_stream()
|
||||
first_dim = self.first_output_dim
|
||||
|
||||
_alloc = lora_overlap_alloc_stream() # capture MAIN stream here (before the fork)
|
||||
side_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(side_stream):
|
||||
lora_a_output = self.lora_backend.run_lora_a_sgemm(
|
||||
x,
|
||||
self.A_buffer,
|
||||
stack_num=(2 if first_dim > 0 else 1),
|
||||
out_alloc_stream=_alloc,
|
||||
)
|
||||
|
||||
# Base ReplicatedLinear GEMM on main, concurrent with the side-stream shrink.
|
||||
output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
|
||||
|
||||
torch.cuda.current_stream().wait_stream(side_stream)
|
||||
if first_dim == 0:
|
||||
output = self.lora_backend.run_lora_b_sgemm(
|
||||
x=lora_a_output,
|
||||
weights=self.B_buffer,
|
||||
output_offset=self._output_offset,
|
||||
base_output=output,
|
||||
)
|
||||
else:
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import qkv_lora_b_fwd
|
||||
|
||||
output = qkv_lora_b_fwd(
|
||||
lora_a_output,
|
||||
self.B_buffer,
|
||||
self.lora_backend._sgemm_info(),
|
||||
self._output_offset,
|
||||
self._max_out_dim,
|
||||
output,
|
||||
n_slices=2,
|
||||
)
|
||||
output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
|
||||
return output, output_bias
|
||||
@@ -0,0 +1,230 @@
|
||||
"""LoRA correction for absorbed-MLA ``kv_b_proj``.
|
||||
|
||||
The absorbed-MLA path in ``DeepseekV2AttentionMLA`` bypasses
|
||||
``kv_b_proj.forward()`` and folds the K/V contribution into two BMMs against
|
||||
the pre-computed ``w_kc`` / ``w_vc`` weights, so a standard
|
||||
``ColumnParallelLinearWithLoRA`` wrapper would never see the activations and
|
||||
the LoRA delta would silently be dropped. These helpers inject the missing
|
||||
delta on top of the absorbed intermediates via the SGMM-style Triton kernels
|
||||
in ``triton_ops/kv_b_lora_absorbed.py``.
|
||||
|
||||
Used from ``deepseek_common/attention_forward_methods/forward_mla.py``. Call
|
||||
sites should gate the call with :func:`is_kv_b_lora_active` so non-LoRA
|
||||
forwards take a single ``getattr`` and skip the helper entirely.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
# The four step kernels live in triton_ops; importing it pulls the LoRA kernel
|
||||
# modules (and specialized_expand) into the process. They are only ever reached
|
||||
# after _get_state returns a non-None state (a kv_b LoRA adapter is wrapped), so
|
||||
# defer the import to that success path: a no-LoRA forward never imports it here.
|
||||
step_a_q_fwd = step_a_v_fwd = step_b_q_fwd = step_b_v_fwd = None
|
||||
|
||||
|
||||
def _ensure_step_kernels() -> None:
|
||||
global step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, step_b_v_fwd
|
||||
if step_a_q_fwd is None:
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_a_q_fwd as _aq
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_a_v_fwd as _av
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_b_q_fwd as _bq
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import step_b_v_fwd as _bv
|
||||
|
||||
step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, step_b_v_fwd = _aq, _av, _bq, _bv
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
|
||||
|
||||
|
||||
def is_kv_b_lora_active(attn_module: "DeepseekV2AttentionMLA") -> bool:
|
||||
"""Cheap precondition check used at call sites in the attention forward
|
||||
to skip the entire LoRA-correction path when no ``kv_b_proj`` adapter is
|
||||
wrapped on this module (the common case)."""
|
||||
return getattr(attn_module.kv_b_proj, "set_lora", False)
|
||||
|
||||
|
||||
def _get_state(
|
||||
attn_module: "DeepseekV2AttentionMLA",
|
||||
) -> Optional[Tuple[torch.Tensor, torch.Tensor, "LoRABatchInfo"]]:
|
||||
if not is_kv_b_lora_active(attn_module):
|
||||
return None
|
||||
if not hasattr(attn_module.kv_b_proj, "A_buffer"):
|
||||
return None
|
||||
lora_backend = attn_module.kv_b_proj.lora_backend
|
||||
if not hasattr(lora_backend, "batch_info"):
|
||||
return None
|
||||
batch_info = lora_backend.batch_info
|
||||
if batch_info is None:
|
||||
return None
|
||||
|
||||
# Triton backend exposes _sgemm_info() to group decode-shape repeats of
|
||||
# the same adapter; csgmv-style backends just expose batch_info directly.
|
||||
sgemm_info = getattr(lora_backend, "_sgemm_info", None)
|
||||
if callable(sgemm_info):
|
||||
batch_info = sgemm_info()
|
||||
# Non-None state ⇒ a kv_b adapter is active here; load the step kernels now
|
||||
# (cached after the first active forward). No-LoRA forwards return above and
|
||||
# never import triton_ops.
|
||||
_ensure_step_kernels()
|
||||
return attn_module.kv_b_proj.A_buffer, attn_module.kv_b_proj.B_buffer, batch_info
|
||||
|
||||
|
||||
def apply_q_correction(
|
||||
attn_module: "DeepseekV2AttentionMLA",
|
||||
q_nope: torch.Tensor,
|
||||
q_nope_out: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""LoRA correction for the absorbed ``q_nope @ w_kc`` path.
|
||||
|
||||
Computes ``q_nope_out += q_nope @ B_kc @ A * scaling`` per token, per
|
||||
active LoRA slot via two SGMM-style Triton kernels. Factored along the
|
||||
LoRA-A/B boundary so we never materialise ``B @ A`` (~268M FMAs per layer
|
||||
per slot in the naive implementation)::
|
||||
|
||||
step A_q : ``(S,H,qk_nope) @ B_kc[slot, h] (qk_nope, rank) -> (S,H,rank)``
|
||||
step B_q : ``(S,H,rank) @ A[slot] (rank, kv_lora_rank) -> += q_nope_out``
|
||||
"""
|
||||
state = _get_state(attn_module)
|
||||
if state is None:
|
||||
return q_nope_out
|
||||
A_buf, B_buf, batch_info = state
|
||||
|
||||
full_K_per_head = attn_module.qk_nope_head_dim + attn_module.v_head_dim
|
||||
q_lora_a = step_a_q_fwd(q_nope, B_buf, batch_info, full_K_per_head)
|
||||
return step_b_q_fwd(q_lora_a, A_buf, batch_info, q_nope_out)
|
||||
|
||||
|
||||
def apply_v_correction(
|
||||
attn_module: "DeepseekV2AttentionMLA",
|
||||
attn_output: torch.Tensor,
|
||||
attn_bmm_flat: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""LoRA correction for the absorbed ``attn_output @ w_vc`` path.
|
||||
|
||||
Computes ``attn_bmm_flat += attn_output @ A.T @ B_vc.T * scaling`` per
|
||||
token, per active LoRA slot. ``attn_bmm_flat`` is the flat
|
||||
``(S, H*v_head_dim)`` view of the absorbed BMM result; we pass strides
|
||||
matching the implicit ``(S, H, v_head_dim)`` layout to step B_v.
|
||||
"""
|
||||
state = _get_state(attn_module)
|
||||
if state is None:
|
||||
return attn_bmm_flat
|
||||
A_buf, B_buf, batch_info = state
|
||||
|
||||
attn_lora_a = step_a_v_fwd(attn_output, A_buf, batch_info)
|
||||
base_view = attn_bmm_flat.view(
|
||||
-1, attn_module.num_local_heads, attn_module.v_head_dim
|
||||
)
|
||||
step_b_v_fwd(
|
||||
attn_lora_a,
|
||||
B_buf,
|
||||
batch_info,
|
||||
base_view,
|
||||
attn_module.qk_nope_head_dim,
|
||||
attn_module.v_head_dim,
|
||||
)
|
||||
return attn_bmm_flat
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Two-stream overlap (O12) for the absorbed kv_b correction.
|
||||
#
|
||||
# Each correction factors into an input-only A-step (reads q_nope / attn_output,
|
||||
# independent of the absorbed bmm output) and a B-step that adds into that bmm
|
||||
# output. ``*_prepare`` forks the A-step onto the shared LoRA side stream so it
|
||||
# overlaps the main-stream ``q_nope @ w_kc`` / ``attn_output @ w_vc`` bmm;
|
||||
# ``*_apply`` rejoins and runs the B-step.
|
||||
#
|
||||
# Gated by ``SGLANG_LORA_TWO_STREAM`` (decode batches only) via
|
||||
# ``is_two_stream_active``. When inactive, ``*_prepare`` returns None and
|
||||
# ``*_apply`` falls back to the serial ``apply_*_correction`` (or a no-op when no
|
||||
# kv_b adapter is wrapped), so the deepseek call sites stay byte-identical with
|
||||
# two-stream off. Same fork/join (``wait_stream``) idiom as the O7/O8 attention
|
||||
# overrides — cuda-graph-capture safe.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _kv_b_two_stream_state(attn_module, x):
|
||||
from sglang.srt.lora.trtllm_lora_temp import (
|
||||
get_lora_side_stream,
|
||||
is_two_stream_active,
|
||||
)
|
||||
|
||||
if not is_two_stream_active(x):
|
||||
return None
|
||||
state = _get_state(attn_module)
|
||||
if state is None:
|
||||
return None
|
||||
A_buf, B_buf, batch_info = state
|
||||
return A_buf, B_buf, batch_info, get_lora_side_stream()
|
||||
|
||||
|
||||
def kv_b_lora_q_prepare(attn_module, q_nope):
|
||||
"""Fork the q-correction A-step onto the side stream (``step_a_q`` reads only
|
||||
``q_nope``) so it overlaps the main-stream ``q_nope @ w_kc`` bmm. Returns a
|
||||
handle for :func:`kv_b_lora_q_apply`, or None when two-stream is inactive."""
|
||||
st = _kv_b_two_stream_state(attn_module, q_nope)
|
||||
if st is None:
|
||||
return None
|
||||
A_buf, B_buf, batch_info, side_stream = st
|
||||
full_K_per_head = attn_module.qk_nope_head_dim + attn_module.v_head_dim
|
||||
side_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(side_stream):
|
||||
q_lora_a = step_a_q_fwd(q_nope, B_buf, batch_info, full_K_per_head)
|
||||
return q_lora_a, A_buf, batch_info, side_stream
|
||||
|
||||
|
||||
def kv_b_lora_q_apply(attn_module, q_nope, q_nope_out, handle):
|
||||
"""Finish the q-correction: two-stream (rejoin + B-step) when ``handle`` is
|
||||
set, else the serial correction, else a no-op. Single call replacing the
|
||||
``if is_kv_b_lora_active: apply_q_correction`` at the call site."""
|
||||
if handle is not None:
|
||||
q_lora_a, A_buf, batch_info, side_stream = handle
|
||||
torch.cuda.current_stream().wait_stream(side_stream)
|
||||
return step_b_q_fwd(q_lora_a, A_buf, batch_info, q_nope_out)
|
||||
if is_kv_b_lora_active(attn_module):
|
||||
return apply_q_correction(attn_module, q_nope, q_nope_out)
|
||||
return q_nope_out
|
||||
|
||||
|
||||
def kv_b_lora_v_prepare(attn_module, attn_output):
|
||||
"""Fork the v-correction A-step onto the side stream (``step_a_v`` reads only
|
||||
``attn_output``) so it overlaps the main-stream ``attn_output @ w_vc`` bmm.
|
||||
Returns a handle for :func:`kv_b_lora_v_apply`, or None when inactive."""
|
||||
st = _kv_b_two_stream_state(attn_module, attn_output)
|
||||
if st is None:
|
||||
return None
|
||||
A_buf, B_buf, batch_info, side_stream = st
|
||||
side_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(side_stream):
|
||||
attn_lora_a = step_a_v_fwd(attn_output, A_buf, batch_info)
|
||||
return attn_lora_a, B_buf, batch_info, side_stream
|
||||
|
||||
|
||||
def kv_b_lora_v_apply(attn_module, attn_output, attn_bmm_flat, handle):
|
||||
"""Finish the v-correction: two-stream (rejoin + B-step) when ``handle`` is
|
||||
set, else the serial correction, else a no-op."""
|
||||
if handle is not None:
|
||||
attn_lora_a, B_buf, batch_info, side_stream = handle
|
||||
torch.cuda.current_stream().wait_stream(side_stream)
|
||||
base_view = attn_bmm_flat.view(
|
||||
-1, attn_module.num_local_heads, attn_module.v_head_dim
|
||||
)
|
||||
step_b_v_fwd(
|
||||
attn_lora_a,
|
||||
B_buf,
|
||||
batch_info,
|
||||
base_view,
|
||||
attn_module.qk_nope_head_dim,
|
||||
attn_module.v_head_dim,
|
||||
)
|
||||
return attn_bmm_flat
|
||||
if is_kv_b_lora_active(attn_module):
|
||||
return apply_v_correction(attn_module, attn_output, attn_bmm_flat)
|
||||
return attn_bmm_flat
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Local env registry for the experimental TRT-LLM LoRA fast path.
|
||||
|
||||
Every flag here is gated by the single global master switch
|
||||
``SGLANG_EXPERIMENTAL_LORA_OPTI`` (defined in ``sglang.srt.environ``). When the master
|
||||
switch is OFF (the default), every flag reads ``False`` (its default is
|
||||
suppressed), so the no-LoRA path, other MoE backends, and the default
|
||||
(non-experimental) LoRA path are byte-identical to upstream.
|
||||
|
||||
Keeping these flags out of the global ``Envs`` class is deliberate: the only
|
||||
sglang-global addition for this feature is ``SGLANG_EXPERIMENTAL_LORA_OPTI``; all the
|
||||
fine-grained opt switches live here, next to the code that consumes them.
|
||||
|
||||
Default policy (applies only when ``SGLANG_EXPERIMENTAL_LORA_OPTI=1``):
|
||||
* **common** flags — used by BOTH the qwen3.5 (FP8) and kimi (NVFP4) configs —
|
||||
default ``True`` so they need not be repeated on every launch command.
|
||||
* **non-shared** flags default ``False`` and must be set explicitly in the
|
||||
launch environment for the model that needs them.
|
||||
|
||||
C++-getenv flags (read via ``getenv`` in the JIT launcher, not in Python) are
|
||||
listed at the bottom for documentation only; set them in the launch env.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
|
||||
def experimental_lora_enabled() -> bool:
|
||||
"""Master gate. All flags below are forced off unless this is set."""
|
||||
return envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
|
||||
_TRUE = {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
class _GatedBool:
|
||||
def __init__(self, name: str, default: bool):
|
||||
self._name = name
|
||||
self._default = default
|
||||
|
||||
def get(self) -> bool:
|
||||
if not experimental_lora_enabled():
|
||||
return False
|
||||
raw = os.environ.get(self._name)
|
||||
if raw is None:
|
||||
return self._default
|
||||
return raw.strip().lower() in _TRUE
|
||||
|
||||
|
||||
class _GatedInt:
|
||||
def __init__(self, name: str, default: int):
|
||||
self._name = name
|
||||
self._default = default
|
||||
|
||||
def get(self) -> int:
|
||||
# Only consulted on the experimental path; return the default otherwise.
|
||||
if not experimental_lora_enabled():
|
||||
return self._default
|
||||
raw = os.environ.get(self._name)
|
||||
return int(raw) if raw is not None else self._default
|
||||
|
||||
|
||||
class _LoraEnvs:
|
||||
# ---- common (qwen3.5 ∩ kimi): default True when experimental is on ----
|
||||
SGLANG_ENABLE_LORA_SHRINK_SPLIT_K = _GatedBool(
|
||||
"SGLANG_ENABLE_LORA_SHRINK_SPLIT_K", True
|
||||
)
|
||||
SGLANG_OPT_LORA_FUSED_MERGED_ALIGN = _GatedBool(
|
||||
"SGLANG_OPT_LORA_FUSED_MERGED_ALIGN", True
|
||||
)
|
||||
SGLANG_OPT_LORA_FUSED_TOPK_PACK = _GatedBool(
|
||||
"SGLANG_OPT_LORA_FUSED_TOPK_PACK", True
|
||||
)
|
||||
SGLANG_OPT_LORA_QKV_B_STORE = _GatedBool("SGLANG_OPT_LORA_QKV_B_STORE", 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.
|
||||
SGLANG_ENABLE_LORA_MOE_GATEUP_GATED_SPLIT = _GatedBool(
|
||||
"SGLANG_ENABLE_LORA_MOE_GATEUP_GATED_SPLIT", True
|
||||
)
|
||||
# feed bf16 router logits straight to the JIT kimi gate (bitwise-identical).
|
||||
SGLANG_OPT_KIMI_GATE_BF16_INPUT = _GatedBool(
|
||||
"SGLANG_OPT_KIMI_GATE_BF16_INPUT", True
|
||||
)
|
||||
|
||||
# ---- non-shared: default False, set explicitly in the launch env ----
|
||||
# kimi (NVFP4):
|
||||
SGLANG_OPT_USE_JIT_KERNEL_KIMI_GATE = _GatedBool(
|
||||
"SGLANG_OPT_USE_JIT_KERNEL_KIMI_GATE", False
|
||||
)
|
||||
SGLANG_OPT_USE_JIT_KERNEL_MOE_ALIGN = _GatedBool(
|
||||
"SGLANG_OPT_USE_JIT_KERNEL_MOE_ALIGN", False
|
||||
)
|
||||
# qwen3.5 (FP8):
|
||||
SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC = _GatedBool(
|
||||
"SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC", False
|
||||
)
|
||||
# (SGLANG_OPT_LORA_DOWN_FINALIZE_OVERLAP removed: net-neutral + base/decode-corruption hazard; serial down-LoRA only.)
|
||||
SGLANG_OPT_LORA_SHARED_ADD_OVERLAP = _GatedBool(
|
||||
"SGLANG_OPT_LORA_SHARED_ADD_OVERLAP", False
|
||||
)
|
||||
SGLANG_OPT_LORA_CUBLAS = _GatedBool("SGLANG_OPT_LORA_CUBLAS", False)
|
||||
SGLANG_OPT_LORA_CUBLAS_A = _GatedBool("SGLANG_OPT_LORA_CUBLAS_A", False)
|
||||
SGLANG_OPT_LORA_CUBLAS_B = _GatedBool("SGLANG_OPT_LORA_CUBLAS_B", False)
|
||||
SGLANG_OPT_LORA_CUBLAS_GATE_UP = _GatedBool("SGLANG_OPT_LORA_CUBLAS_GATE_UP", False)
|
||||
SGLANG_OPT_LORA_CUBLAS_QKV = _GatedBool("SGLANG_OPT_LORA_CUBLAS_QKV", False)
|
||||
SGLANG_OPT_LORA_CUBLAS_KV_B = _GatedBool("SGLANG_OPT_LORA_CUBLAS_KV_B", False)
|
||||
# diagnostics / tuning:
|
||||
SGLANG_OPT_LORA_SHRINK_TUNE = _GatedBool("SGLANG_OPT_LORA_SHRINK_TUNE", False)
|
||||
|
||||
# kimi NVFP4 permute+quant fuse — read in jit_kernel/trtllm_lora_temp/core.py (Python) to pass
|
||||
# a bool to the kernel, AND C++-side via getenv in the launcher. Default off (kimi-only).
|
||||
SGLANG_OPT_FUSED_PERMUTE_QUANT = _GatedBool("SGLANG_OPT_FUSED_PERMUTE_QUANT", False)
|
||||
|
||||
# ---- integer knob ----
|
||||
# decode two-stream token ceiling (consulted only on the experimental path).
|
||||
SGLANG_TWO_STREAM_MAX_TOKENS = _GatedInt("SGLANG_TWO_STREAM_MAX_TOKENS", 256)
|
||||
|
||||
|
||||
lora_envs = _LoraEnvs()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# C++-getenv-only flags (read via getenv in jit_kernel .cu launchers, NOT in Python).
|
||||
# Set them in the launch env on the model that needs them; default off:
|
||||
# SGLANG_OPT_FUSED_MOE_ACTIVATION_QUANT_FUSE (kimi NVFP4 act+down-quant fuse)
|
||||
# SGLANG_OPT_FUSED_MOE_ACTIVATION_VEC (kimi NVFP4 vectorized activation)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -0,0 +1,203 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
|
||||
def _fake_fp8_block_scale_moe(
|
||||
routing_logits: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
n_group: Optional[int],
|
||||
topk_group: Optional[int],
|
||||
intermediate_size: int,
|
||||
local_expert_offset: int,
|
||||
local_num_experts: int,
|
||||
routed_scaling_factor: Optional[float],
|
||||
routing_method_type: int = 0,
|
||||
use_shuffled_weight: bool = False,
|
||||
weight_layout: int = 0,
|
||||
enable_pdl: Optional[bool] = None,
|
||||
tune_max_num_tokens: int = 8192,
|
||||
fp8_quantization_type: Optional[int] = None,
|
||||
activation_type: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty(
|
||||
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(fake_impl=_fake_fp8_block_scale_moe)
|
||||
def sgl_trtllm_fp8_block_scale_moe_wrapper(
|
||||
routing_logits: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
n_group: Optional[int],
|
||||
topk_group: Optional[int],
|
||||
intermediate_size: int,
|
||||
local_expert_offset: int,
|
||||
local_num_experts: int,
|
||||
routed_scaling_factor: Optional[float],
|
||||
routing_method_type: int = 0,
|
||||
use_shuffled_weight: bool = False,
|
||||
weight_layout: int = 0,
|
||||
enable_pdl: Optional[bool] = None,
|
||||
tune_max_num_tokens: int = 8192,
|
||||
fp8_quantization_type: Optional[int] = None,
|
||||
activation_type: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
try:
|
||||
from flashinfer.fused_moe import Fp8QuantizationType
|
||||
from flashinfer.fused_moe.core import ActivationType
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"experimental_sgl_trtllm requires flashinfer-python to provide "
|
||||
"TRTLLM enums and cubin-loader utilities."
|
||||
) from e
|
||||
|
||||
from sglang.jit_kernel.trtllm_lora_temp import trtllm_fp8_block_scale_moe
|
||||
|
||||
kwargs = {
|
||||
"routing_logits": routing_logits,
|
||||
"routing_bias": routing_bias,
|
||||
"hidden_states": hidden_states,
|
||||
"hidden_states_scale": hidden_states_scale,
|
||||
"gemm1_weights": gemm1_weights,
|
||||
"gemm1_weights_scale": gemm1_weights_scale,
|
||||
"gemm2_weights": gemm2_weights,
|
||||
"gemm2_weights_scale": gemm2_weights_scale,
|
||||
"num_experts": num_experts,
|
||||
"top_k": top_k,
|
||||
"n_group": n_group,
|
||||
"topk_group": topk_group,
|
||||
"intermediate_size": intermediate_size,
|
||||
"local_expert_offset": local_expert_offset,
|
||||
"local_num_experts": local_num_experts,
|
||||
"routed_scaling_factor": routed_scaling_factor,
|
||||
"routing_method_type": routing_method_type,
|
||||
"use_shuffled_weight": use_shuffled_weight,
|
||||
"weight_layout": weight_layout,
|
||||
"enable_pdl": enable_pdl,
|
||||
"tune_max_num_tokens": tune_max_num_tokens,
|
||||
}
|
||||
if fp8_quantization_type is not None:
|
||||
kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
|
||||
if activation_type is not None:
|
||||
kwargs["activation_type"] = ActivationType(activation_type)
|
||||
|
||||
return trtllm_fp8_block_scale_moe(**kwargs)
|
||||
|
||||
|
||||
def _fake_fp8_block_scale_routed_moe(
|
||||
topk_ids: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
n_group: Optional[int],
|
||||
topk_group: Optional[int],
|
||||
intermediate_size: int,
|
||||
local_expert_offset: int,
|
||||
local_num_experts: int,
|
||||
routed_scaling_factor: Optional[float],
|
||||
routing_method_type: int = 0,
|
||||
use_shuffled_weight: bool = False,
|
||||
weight_layout: int = 0,
|
||||
enable_pdl: Optional[bool] = None,
|
||||
tune_max_num_tokens: int = 8192,
|
||||
fp8_quantization_type: Optional[int] = None,
|
||||
activation_type: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty(
|
||||
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(fake_impl=_fake_fp8_block_scale_routed_moe)
|
||||
def sgl_trtllm_fp8_block_scale_routed_moe_wrapper(
|
||||
topk_ids: torch.Tensor,
|
||||
routing_bias: Optional[torch.Tensor],
|
||||
hidden_states: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm1_weights_scale: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
gemm2_weights_scale: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
n_group: Optional[int],
|
||||
topk_group: Optional[int],
|
||||
intermediate_size: int,
|
||||
local_expert_offset: int,
|
||||
local_num_experts: int,
|
||||
routed_scaling_factor: Optional[float],
|
||||
routing_method_type: int = 0,
|
||||
use_shuffled_weight: bool = False,
|
||||
weight_layout: int = 0,
|
||||
enable_pdl: Optional[bool] = None,
|
||||
tune_max_num_tokens: int = 8192,
|
||||
fp8_quantization_type: Optional[int] = None,
|
||||
activation_type: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
try:
|
||||
from flashinfer.fused_moe import Fp8QuantizationType
|
||||
from flashinfer.fused_moe.core import ActivationType
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"experimental_sgl_trtllm requires flashinfer-python to provide "
|
||||
"TRTLLM enums and cubin-loader utilities."
|
||||
) from e
|
||||
|
||||
from sglang.jit_kernel.trtllm_lora_temp import (
|
||||
trtllm_fp8_block_scale_routed_moe,
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"topk_ids": topk_ids,
|
||||
"routing_bias": routing_bias,
|
||||
"hidden_states": hidden_states,
|
||||
"hidden_states_scale": hidden_states_scale,
|
||||
"gemm1_weights": gemm1_weights,
|
||||
"gemm1_weights_scale": gemm1_weights_scale,
|
||||
"gemm2_weights": gemm2_weights,
|
||||
"gemm2_weights_scale": gemm2_weights_scale,
|
||||
"num_experts": num_experts,
|
||||
"top_k": top_k,
|
||||
"n_group": n_group,
|
||||
"topk_group": topk_group,
|
||||
"intermediate_size": intermediate_size,
|
||||
"local_expert_offset": local_expert_offset,
|
||||
"local_num_experts": local_num_experts,
|
||||
"routed_scaling_factor": routed_scaling_factor,
|
||||
"routing_method_type": routing_method_type,
|
||||
"use_shuffled_weight": use_shuffled_weight,
|
||||
"weight_layout": weight_layout,
|
||||
"enable_pdl": enable_pdl,
|
||||
"tune_max_num_tokens": tune_max_num_tokens,
|
||||
}
|
||||
if fp8_quantization_type is not None:
|
||||
kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
|
||||
if activation_type is not None:
|
||||
kwargs["activation_type"] = ActivationType(activation_type)
|
||||
|
||||
return trtllm_fp8_block_scale_routed_moe(**kwargs)
|
||||
@@ -0,0 +1,464 @@
|
||||
"""experimental_sgl_trtllm MoE LoRA dispatch (original single-stream).
|
||||
|
||||
This is the LoRA-enabled fused-experts path added by the trtllm-lora work — it
|
||||
was originally a function in ``layers/moe/moe_runner/flashinfer_trtllm.py`` and
|
||||
is now hosted here so that file holds only a re-export. The function name
|
||||
remains ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora`` for
|
||||
import-site stability.
|
||||
|
||||
When ``SGLANG_LORA_TWO_STREAM=1`` is set, this is the function the
|
||||
``install_two_stream_overrides()`` monkey-patch swaps for the side-stream
|
||||
version in :mod:`sglang.srt.lora.trtllm_lora_temp.moe_overlap`. Otherwise it runs as
|
||||
the active path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
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.quantization.fp8_kernel import per_token_group_quant_fp8
|
||||
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 (
|
||||
FlashInferTrtllmFp4MoeQuantInfo,
|
||||
FlashInferTrtllmFp8MoeQuantInfo,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher import (
|
||||
StandardCombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
|
||||
|
||||
def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
quant_info: "FlashInferTrtllmFp8MoeQuantInfo",
|
||||
runner_config: "MoeRunnerConfig",
|
||||
lora_info,
|
||||
) -> "StandardCombineInput":
|
||||
from flashinfer.fused_moe import Fp8QuantizationType
|
||||
|
||||
from sglang.jit_kernel.trtllm_lora_temp import (
|
||||
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.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.lora_moe_runners import build_lora_hooks
|
||||
from sglang.srt.lora.trtllm_lora_temp.sgl_fp8_moe import (
|
||||
fused_experts_fp8_sgl,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import (
|
||||
maybe_overlap_staged_shared_add,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
|
||||
merged_experts_fused_moe_lora_add,
|
||||
)
|
||||
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
|
||||
|
||||
assert runner_config.activation == "silu" and runner_config.is_gated, (
|
||||
"experimental_sgl_trtllm LoRA currently supports the gated SwiGLU FP8 "
|
||||
"Qwen path only."
|
||||
)
|
||||
assert quant_info.block_quant and not quant_info.use_mxfp8, (
|
||||
"experimental_sgl_trtllm LoRA currently supports DeepSeekFp8 block-quant "
|
||||
"checkpoints only."
|
||||
)
|
||||
assert quant_info.weight_block_k is not None
|
||||
assert quant_info.w13_weight_scale_inv is not None
|
||||
assert quant_info.w2_weight_scale_inv is not None
|
||||
|
||||
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
|
||||
|
||||
if not get_is_capture_mode() and not lora_info.has_active_lora:
|
||||
return fused_experts_fp8_sgl(
|
||||
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
|
||||
)
|
||||
if use_virtual_lora_store:
|
||||
hooks = None
|
||||
token_lora_mapping = lora_info.token_lora_mapping
|
||||
fused_lora_routing_cache: dict = {}
|
||||
else:
|
||||
hooks = build_lora_hooks(hidden_states, lora_info, topk_ids)
|
||||
token_lora_mapping = None
|
||||
fused_lora_routing_cache = {}
|
||||
|
||||
# Fuse the per-token scale transpose into the quant kernel (column-major scales) so the
|
||||
# `.t()` is a free view -> drops the standalone ~2us transpose+copy. Byte/shape-identical.
|
||||
a_q, a_sf = per_token_group_quant_fp8(
|
||||
hidden_states, quant_info.weight_block_k, column_major_scales=True
|
||||
)
|
||||
a_sf_t = a_sf.t()
|
||||
|
||||
# EP-aware LoRA: under MoE EP each rank computes the delta only for the experts it
|
||||
# owns (passed via local_expert_offset/local_num_experts below). gate_up_delta stays
|
||||
# new_empty even though non-owned [token, k] slots are then left unwritten -- the
|
||||
# trtllm MoE is itself EP-aware, so those slots never feed the all-reduced output.
|
||||
gate_up_delta_shape = (
|
||||
hidden_states.shape[0],
|
||||
runner_config.top_k,
|
||||
quant_info.w13_weight.shape[1],
|
||||
)
|
||||
gate_up_delta = (
|
||||
hidden_states.new_empty(gate_up_delta_shape)
|
||||
if use_virtual_lora_store
|
||||
else hidden_states.new_zeros(gate_up_delta_shape)
|
||||
)
|
||||
if use_virtual_lora_store:
|
||||
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=quant_info.local_num_experts,
|
||||
)
|
||||
elif hooks.after_gate_up is not None:
|
||||
hooks.after_gate_up(hidden_states, gate_up_delta, topk_weights, topk_ids)
|
||||
|
||||
activation_lora_input = torch.empty(
|
||||
(hidden_states.shape[0], runner_config.top_k, quant_info.intermediate_size),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
# SGLANG_OPT_LORA_FUSED_TOPK_PACK: the routed pack may already have been produced
|
||||
# fused inside the gating kernel (StandardTopKOutput.packed_topk_ids) — including
|
||||
# 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(
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
)
|
||||
|
||||
direct_down_output = None
|
||||
if use_virtual_lora_store:
|
||||
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,
|
||||
)
|
||||
|
||||
moe_result = trtllm_fp8_block_scale_routed_moe_lora(
|
||||
topk_ids=packed_topk_ids,
|
||||
routing_bias=None,
|
||||
hidden_states=a_q,
|
||||
hidden_states_scale=a_sf_t,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
|
||||
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,
|
||||
n_group=None,
|
||||
topk_group=None,
|
||||
intermediate_size=quant_info.intermediate_size,
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_experts,
|
||||
routed_scaling_factor=(
|
||||
runner_config.routed_scaling_factor
|
||||
if runner_config.routed_scaling_factor is not None
|
||||
else 1.0
|
||||
),
|
||||
routing_method_type=(
|
||||
RoutingMethodType.TopK
|
||||
if quant_info.routing_method_type == RoutingMethodType.DeepSeekV3
|
||||
else quant_info.routing_method_type
|
||||
),
|
||||
use_shuffled_weight=False,
|
||||
do_finalize=use_virtual_lora_store,
|
||||
output=(
|
||||
direct_down_output
|
||||
if direct_down_output is not None
|
||||
else torch.empty_like(hidden_states)
|
||||
),
|
||||
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
|
||||
fp8_quantization_type=Fp8QuantizationType.DeepSeekFp8,
|
||||
activation_type=quant_info.activation_type,
|
||||
)
|
||||
if use_virtual_lora_store:
|
||||
output = moe_result
|
||||
# Shared-add overlap: the trtllm op above already finalized `output`, so the
|
||||
# staged shared-expert add (if any) can run on the main stream concurrent with
|
||||
# the down-LoRA shrink below; the expand waits on it via expand_wait_event.
|
||||
shared_add_done = maybe_overlap_staged_shared_add(output)
|
||||
merged_experts_fused_moe_lora_add(
|
||||
output=output,
|
||||
hidden_states=activation_lora_input.view(-1, quant_info.intermediate_size),
|
||||
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=quant_info.local_num_experts,
|
||||
expand_wait_event=shared_add_done,
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
gemm2_output, expert_weights, expanded_idx_to_permuted_idx = moe_result
|
||||
|
||||
down_delta_shape = (
|
||||
hidden_states.shape[0],
|
||||
runner_config.top_k,
|
||||
hidden_states.shape[1],
|
||||
)
|
||||
down_delta = (
|
||||
hidden_states.new_empty(down_delta_shape)
|
||||
if use_virtual_lora_store
|
||||
else hidden_states.new_zeros(down_delta_shape)
|
||||
)
|
||||
if use_virtual_lora_store:
|
||||
merged_experts_fused_moe_lora_add(
|
||||
output=down_delta,
|
||||
hidden_states=activation_lora_input.view(-1, quant_info.intermediate_size),
|
||||
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,
|
||||
)
|
||||
elif hooks.after_down is not None:
|
||||
hooks.after_down(
|
||||
activation_lora_input.view(-1, quant_info.intermediate_size),
|
||||
down_delta,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
)
|
||||
|
||||
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
|
||||
output = torch.empty(
|
||||
hidden_states.shape[0],
|
||||
hidden_states.shape[1],
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
output = trtllm_fp8_block_scale_moe_lora_finalize(
|
||||
gemm2_output=gemm2_output,
|
||||
expert_weights=expert_weights,
|
||||
expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
|
||||
down_lora_delta=down_delta,
|
||||
output=output,
|
||||
routed_scaling_factor=(
|
||||
runner_config.routed_scaling_factor
|
||||
if runner_config.routed_scaling_factor is not None
|
||||
else 1.0
|
||||
),
|
||||
)
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
quant_info: "FlashInferTrtllmFp4MoeQuantInfo",
|
||||
runner_config: "MoeRunnerConfig",
|
||||
lora_info,
|
||||
) -> "StandardCombineInput":
|
||||
"""NVFP4 sibling of ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``.
|
||||
|
||||
Decomposed (unfused-activation) MoE-LoRA: routing -> gather -> gate_up grouped
|
||||
GEMM (raw 2*inter) -> activation that adds ``gate_up_lora_delta`` pre-SwiGLU and
|
||||
captures ``activation_lora_input`` -> NvFP4 quant -> down grouped GEMM -> finalize,
|
||||
then the virtual-experts down-LoRA is merged into the output. Single-stream
|
||||
version; ``moe_overlap.py`` provides the two-stream variant.
|
||||
"""
|
||||
from sglang.jit_kernel.trtllm_lora_temp import (
|
||||
trtllm_fp4_block_scale_routed_moe_lora,
|
||||
)
|
||||
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
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
|
||||
merged_experts_fused_moe_lora_add,
|
||||
)
|
||||
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
|
||||
|
||||
assert (
|
||||
runner_config.activation == "silu" and runner_config.is_gated
|
||||
), "experimental_sgl_trtllm NVFP4 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) FP4 path.
|
||||
if not get_is_capture_mode() and not lora_info.has_active_lora:
|
||||
return fused_experts_none_to_flashinfer_trtllm_fp4(
|
||||
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, "NVFP4 trtllm LoRA requires virtual-experts."
|
||||
token_lora_mapping = lora_info.token_lora_mapping
|
||||
fused_lora_routing_cache: dict = {}
|
||||
|
||||
inter = quant_info.intermediate_size_per_partition
|
||||
|
||||
# Path 3: feed bf16 hidden DIRECTLY to the op (no python pre-quant). The op permutes the
|
||||
# bf16 hidden (moe::dev::permute, token->expert order) then NvFP4-quantizes ONCE with the
|
||||
# 1/(448*6) global + per-token scale — eliminating the dequant->permute->requant round-trip
|
||||
# (and its magnitude bug) that pre-quantized fp4 input forced. Requires the layer to be in
|
||||
# per-token-activation mode (SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1) so that
|
||||
# g1_scale_c == g1_alphas and g2_alphas == w2_weight_scale_2, making the decomposed
|
||||
# gate_up(g1_alphas)/SwiGLU/down(g2_alphas) scale composition match the plain fused path.
|
||||
|
||||
gate_up_delta_shape = (
|
||||
hidden_states.shape[0],
|
||||
runner_config.top_k,
|
||||
quant_info.w13_weight.shape[1],
|
||||
)
|
||||
# Gated gate_up LoRA delta: a single merged_experts call on the full stacked [gate_A; up_A]
|
||||
# lora_a (rank 2r) and [gate_B; up_B] lora_b (rank r), via the rank-specialized direct expand
|
||||
# (use_direct_expand_add, rank <= 64). EP args scope the delta to this rank's experts, matching
|
||||
# the EP-aware trtllm MoE base.
|
||||
gate_up_delta = hidden_states.new_empty(gate_up_delta_shape)
|
||||
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=quant_info.local_num_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 = _pack_topk_for_flashinfer_routed(
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
)
|
||||
|
||||
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_fp4_block_scale_routed_moe_lora(
|
||||
topk_ids=packed_topk_ids,
|
||||
routing_bias=None,
|
||||
hidden_states=hidden_states,
|
||||
hidden_states_scale=None,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn),
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn),
|
||||
output1_scales_scalar=quant_info.g1_scale_c,
|
||||
output1_scales_gate_scalar=quant_info.g1_alphas,
|
||||
output2_scales_scalar=quant_info.g2_alphas,
|
||||
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=quant_info.local_num_experts,
|
||||
routed_scaling_factor=(
|
||||
runner_config.routed_scaling_factor
|
||||
if runner_config.routed_scaling_factor is not None
|
||||
else 1.0
|
||||
),
|
||||
routing_method_type=quant_info.routing_method_type,
|
||||
do_finalize=True,
|
||||
output=direct_down_output,
|
||||
)
|
||||
|
||||
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,
|
||||
# EP-aware: scope the down delta to this rank's experts, matching the gate_up
|
||||
# call above and the FP8 down call. Harmless at EP=1 (local==global, Kimi today);
|
||||
# required for correctness if MoE-EP is turned on later (otherwise non-owned
|
||||
# experts' deltas get over-counted by the fuse_sum_all_reduce).
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_experts,
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
@@ -0,0 +1,152 @@
|
||||
"""experimental_sgl_trtllm-specific bits of ``FusedMoEWithLoRA``.
|
||||
|
||||
This file holds the two trtllm-specific code blocks that used to be inlined
|
||||
inside the ``FusedMoEWithLoRA`` class in ``lora/layers.py``:
|
||||
|
||||
- :func:`init_experimental_sgl_trtllm_lora` — builds the FP8 block-scale
|
||||
``FlashInferTrtllmFp8MoeQuantInfo`` and stores it on the layer instance.
|
||||
Called from ``FusedMoEWithLoRA.__init__`` when the runner backend is the
|
||||
experimental_sgl_trtllm MoE.
|
||||
- :func:`dispatch_experimental_sgl_trtllm_lora` — dispatches the LoRA fused
|
||||
experts call. Called from ``FusedMoEWithLoRA.run`` for the same backend.
|
||||
|
||||
Keeping them here means ``lora/layers.py`` only has tiny ``if backend == ...:
|
||||
init(self, base_layer)`` / ``dispatch(...)`` injection points for the new
|
||||
trtllm path instead of ~70 lines of inlined logic.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
|
||||
def init_experimental_sgl_trtllm_lora(layer, base_layer) -> None:
|
||||
"""Build and store the trtllm FP8 LoRA quant info on the layer.
|
||||
|
||||
Sets ``layer._lora_runner = None`` (trtllm path doesn't use ``MoeRunner``)
|
||||
and ``layer._quant_info`` to a fully-populated
|
||||
``FlashInferTrtllmFp8MoeQuantInfo`` (or ``FlashInferTrtllmFp4MoeQuantInfo`` for
|
||||
NVFP4 / modelopt checkpoints like Kimi-K2.5-NVFP4).
|
||||
"""
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
FlashInferTrtllmFp8MoeQuantInfo,
|
||||
get_activation_type,
|
||||
)
|
||||
from sglang.srt.layers.moe.utils import RoutingMethodType
|
||||
|
||||
# ---- NVFP4 (modelopt) path ----
|
||||
# The fp4 weight loader sets ``g1_scale_c`` on the FusedMoE layer (see
|
||||
# ModelOptNvFp4FusedMoEMethod.apply). Mirror the non-LoRA construction in
|
||||
# modelopt_quant.py (~L2099) so the fp4 LoRA dispatch gets the same payload.
|
||||
# NOTE(w13 layout): the decomposed op runs the gate_up projection as a *non-gated*
|
||||
# Gemm2-style GEMM and lets the activation kernel do the SwiGLU split (silu(first)*second),
|
||||
# so w13 must be [Gate, Up] with shuffle_matrix_a but WITHOUT reorder_rows_for_gated_act_gemm.
|
||||
# The TRTLLM fp4 path uses load_up_proj_weight_first=False (=> [Gate, Up]); verify the
|
||||
# processed w13 layout against this at e2e (acc gate) and re-prep if mismatched.
|
||||
if hasattr(base_layer, "g1_scale_c"):
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
FlashInferTrtllmFp4MoeQuantInfo,
|
||||
)
|
||||
|
||||
layer._lora_runner = None
|
||||
layer._quant_info = FlashInferTrtllmFp4MoeQuantInfo(
|
||||
w13_weight=base_layer.w13_weight.data,
|
||||
w2_weight=base_layer.w2_weight.data,
|
||||
w13_weight_scale=base_layer.w13_weight_scale.data,
|
||||
w2_weight_scale=base_layer.w2_weight_scale.data,
|
||||
g1_scale_c=base_layer.g1_scale_c.data,
|
||||
g1_alphas=base_layer.g1_alphas.data,
|
||||
g2_alphas=base_layer.g2_alphas.data,
|
||||
w13_input_scale_quant=base_layer.w13_input_scale_quant,
|
||||
global_num_experts=int(base_layer.num_experts),
|
||||
local_expert_offset=int(base_layer.moe_ep_rank)
|
||||
* int(base_layer.num_local_experts),
|
||||
local_num_experts=int(base_layer.num_local_experts),
|
||||
intermediate_size_per_partition=int(
|
||||
base_layer.intermediate_size_per_partition
|
||||
),
|
||||
routing_method_type=int(
|
||||
getattr(base_layer, "routing_method_type", None)
|
||||
or RoutingMethodType.DeepSeekV3
|
||||
),
|
||||
)
|
||||
return
|
||||
|
||||
quant_method = base_layer.quant_method
|
||||
quant_config = getattr(quant_method, "quant_config", None)
|
||||
weight_block_size = getattr(quant_config, "weight_block_size", 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))
|
||||
assert getattr(
|
||||
quant_method, "block_quant", False
|
||||
), "experimental_sgl_trtllm LoRA currently requires FP8 block quant."
|
||||
assert (
|
||||
not use_mxfp8
|
||||
), "experimental_sgl_trtllm LoRA currently targets the non-MX FP8 Qwen path."
|
||||
assert (
|
||||
weight_block_size is not None
|
||||
), "experimental_sgl_trtllm LoRA needs the FP8 weight block size."
|
||||
w13_weight_scale = getattr(base_layer, "w13_weight_scale_inv", None)
|
||||
if w13_weight_scale is None:
|
||||
w13_weight_scale = getattr(base_layer, "w13_weight_scale", None)
|
||||
w2_weight_scale = getattr(base_layer, "w2_weight_scale_inv", None)
|
||||
if w2_weight_scale is None:
|
||||
w2_weight_scale = getattr(base_layer, "w2_weight_scale", None)
|
||||
assert w13_weight_scale is not None and w2_weight_scale is not None
|
||||
|
||||
layer._lora_runner = None
|
||||
layer._quant_info = FlashInferTrtllmFp8MoeQuantInfo(
|
||||
w13_weight=base_layer.w13_weight,
|
||||
w2_weight=base_layer.w2_weight,
|
||||
global_num_experts=int(base_layer.num_experts),
|
||||
local_expert_offset=int(base_layer.moe_ep_rank)
|
||||
* int(base_layer.num_local_experts),
|
||||
local_num_experts=int(base_layer.num_local_experts),
|
||||
intermediate_size=base_layer.w2_weight.shape[2],
|
||||
routing_method_type=int(
|
||||
getattr(base_layer, "routing_method_type", None)
|
||||
or RoutingMethodType.DeepSeekV3
|
||||
),
|
||||
block_quant=True,
|
||||
use_mxfp8=False,
|
||||
weight_block_k=weight_block_size[1],
|
||||
w13_weight_scale_inv=w13_weight_scale,
|
||||
w2_weight_scale_inv=w2_weight_scale,
|
||||
activation_type=get_activation_type(
|
||||
base_layer.moe_runner_config.activation,
|
||||
is_gated=base_layer.moe_runner_config.is_gated,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def dispatch_experimental_sgl_trtllm_lora(
|
||||
dispatch_output, quant_info, base_layer, lora_info
|
||||
) -> "StandardCombineInput":
|
||||
"""Call the trtllm fused-experts LoRA function for a single layer.
|
||||
|
||||
Looked up at call time so the install-time monkey-patch in
|
||||
:mod:`sglang.srt.lora.trtllm_lora_temp` (the two-stream override) takes effect.
|
||||
"""
|
||||
import sglang.srt.lora.trtllm_lora_temp.lora_dispatch as ft
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
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.
|
||||
if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo):
|
||||
fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_fp4_lora
|
||||
else:
|
||||
fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora
|
||||
|
||||
return fused_fn(
|
||||
dispatch_output,
|
||||
quant_info,
|
||||
base_layer.moe_runner_config,
|
||||
lora_info,
|
||||
)
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Two-stream MergedColumnParallelLinear LoRA forward (O9).
|
||||
|
||||
Monkey-patched onto :class:`MergedColumnParallelLinearWithLoRA` by
|
||||
:func:`sglang.srt.lora.trtllm_lora_temp.install_two_stream_overrides` when
|
||||
``SGLANG_LORA_TWO_STREAM=1``. Covers the merged-column LoRA modules not handled
|
||||
by O7 (QKV) / O8 (o_proj) / O1 (MoE experts):
|
||||
|
||||
* Qwen3.5 mamba ``in_proj_qkvz`` (a MergedColumnParallelLinear, every mamba layer)
|
||||
* dense ``gate_up_proj`` MLP layers (e.g. Qwen3-VL non-expert MLP)
|
||||
|
||||
Same shape as O7: the LoRA-A shrink reads ``input_`` (same input as the base
|
||||
GEMM, no write conflict) and runs on the side stream concurrent with the base
|
||||
merged-column GEMM on the main stream; the LoRA-B expand needs both the shrink
|
||||
output and base_output, so it runs after the rejoin on the main stream. The
|
||||
expand mirrors ``MergedColumnParallelLinearWithLoRA.apply_lora`` — gate_up vs
|
||||
general n-slice.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.distributed import tensor_model_parallel_all_gather
|
||||
from sglang.srt.lora.trtllm_lora_temp import (
|
||||
get_lora_side_stream,
|
||||
get_original_merged_column_forward,
|
||||
is_two_stream_active,
|
||||
lora_overlap_alloc_stream,
|
||||
)
|
||||
|
||||
|
||||
def merged_column_lora_forward(self, input_: torch.Tensor):
|
||||
"""O9 — side-stream LoRA-A shrink ‖ base merged-column GEMM."""
|
||||
if not self.set_lora or not is_two_stream_active(input_):
|
||||
return get_original_merged_column_forward()(self, input_)
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
|
||||
gate_up_lora_b_fwd,
|
||||
qkv_lora_b_fwd,
|
||||
sgemm_lora_a_fwd,
|
||||
)
|
||||
|
||||
bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
|
||||
side_stream = get_lora_side_stream()
|
||||
# sgemm_info is host-side (LoRABatchInfo); compute once, share both calls.
|
||||
sgemm_info = self.lora_backend._sgemm_info()
|
||||
lora_n_slices = self._get_lora_n_slices()
|
||||
use_gate_up = lora_n_slices == 2 and self.use_gate_up_lora
|
||||
|
||||
# Shrink on side stream, concurrent with the base merged-column GEMM on main.
|
||||
_alloc = lora_overlap_alloc_stream() # capture MAIN stream here (before the fork)
|
||||
side_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(side_stream):
|
||||
shrink_intermediate = sgemm_lora_a_fwd(
|
||||
input_,
|
||||
self.A_buffer,
|
||||
sgemm_info,
|
||||
stack_num=lora_n_slices,
|
||||
out_alloc_stream=_alloc,
|
||||
)
|
||||
|
||||
output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_, bias)
|
||||
|
||||
# Rejoin: expand reads both side-produced shrink_intermediate and base_output.
|
||||
torch.cuda.current_stream().wait_stream(side_stream)
|
||||
if use_gate_up:
|
||||
output_dim = self.B_buffer.shape[-2] // 2
|
||||
output_parallel = gate_up_lora_b_fwd(
|
||||
shrink_intermediate,
|
||||
self.B_buffer,
|
||||
sgemm_info,
|
||||
output_dim,
|
||||
output_parallel,
|
||||
)
|
||||
else:
|
||||
output_parallel = qkv_lora_b_fwd(
|
||||
shrink_intermediate,
|
||||
self.B_buffer,
|
||||
sgemm_info,
|
||||
self.output_offset,
|
||||
self.max_out_dim,
|
||||
output_parallel,
|
||||
n_slices=lora_n_slices,
|
||||
)
|
||||
|
||||
if self.base_layer.gather_output:
|
||||
output = tensor_model_parallel_all_gather(output_parallel)
|
||||
else:
|
||||
output = output_parallel
|
||||
output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
|
||||
return output, output_bias
|
||||
@@ -0,0 +1,524 @@
|
||||
"""Two-stream MoE LoRA dispatch (O1).
|
||||
|
||||
Monkey-patches ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora`` in
|
||||
``layers/moe/moe_runner/flashinfer_trtllm.py`` (when
|
||||
``SGLANG_LORA_TWO_STREAM=1``) so the gate_up LoRA shrink+expand runs on a
|
||||
side stream concurrent with the main-stream FP8 quant.
|
||||
|
||||
Batches that don't qualify for two-stream (prefill / non-virtual-lora /
|
||||
batch without active LoRA) fall through to the saved-original function so
|
||||
their behavior is byte-identical to the unpatched code path.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp import (
|
||||
get_lora_side_stream,
|
||||
get_original_fp4_moe_lora_func,
|
||||
get_original_moe_lora_func,
|
||||
is_two_stream_active,
|
||||
)
|
||||
|
||||
# GEMM1-LoRA overlap: keep LoRA-ready events recorded during cuda-graph capture alive so the
|
||||
# captured cross-stream wait (resolved inside the trtllm op before activation) isn't torn down
|
||||
# before graph instantiation. Only appended while capturing; eager runs rely on CUDA's
|
||||
# deferred cudaEventDestroy.
|
||||
_LORA_OVERLAP_EVENTS: list = []
|
||||
|
||||
|
||||
def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
|
||||
dispatch_output,
|
||||
quant_info,
|
||||
runner_config,
|
||||
lora_info,
|
||||
):
|
||||
"""Drop-in replacement for the like-named function in flashinfer_trtllm.py.
|
||||
|
||||
Two-stream fast path: only fires when the batch is decode-shaped AND uses
|
||||
virtual-experts LoRA. Everything else delegates to the original function.
|
||||
"""
|
||||
hidden_states = dispatch_output.hidden_states
|
||||
|
||||
use_virtual_lora_store = bool(
|
||||
lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
|
||||
)
|
||||
# Two-stream requires virtual-experts LoRA AND a decode-shaped batch.
|
||||
# Fall back to the original implementation for anything else (prefill,
|
||||
# non-virtual LoRA, non-LoRA capture, etc.).
|
||||
if not (use_virtual_lora_store and is_two_stream_active(hidden_states)):
|
||||
return get_original_moe_lora_func()(
|
||||
dispatch_output, quant_info, runner_config, lora_info
|
||||
)
|
||||
|
||||
# ---- two-stream fast path ----
|
||||
from flashinfer.fused_moe import Fp8QuantizationType
|
||||
|
||||
from sglang.jit_kernel.trtllm_lora_temp import (
|
||||
trtllm_fp8_block_scale_routed_moe_lora,
|
||||
)
|
||||
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
|
||||
from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8
|
||||
from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import (
|
||||
maybe_overlap_staged_shared_add,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
|
||||
merged_experts_fused_moe_lora_add,
|
||||
)
|
||||
from sglang.srt.utils.common import next_power_of_2
|
||||
|
||||
assert runner_config.activation == "silu" and runner_config.is_gated, (
|
||||
"experimental_sgl_trtllm LoRA currently supports the gated SwiGLU FP8 "
|
||||
"Qwen path only."
|
||||
)
|
||||
assert quant_info.block_quant and not quant_info.use_mxfp8, (
|
||||
"experimental_sgl_trtllm LoRA currently supports DeepSeekFp8 block-quant "
|
||||
"checkpoints only."
|
||||
)
|
||||
assert quant_info.weight_block_k is not None
|
||||
assert quant_info.w13_weight_scale_inv is not None
|
||||
assert quant_info.w2_weight_scale_inv is not None
|
||||
|
||||
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 = {}
|
||||
|
||||
side_stream = get_lora_side_stream()
|
||||
|
||||
# EP-aware LoRA: under MoE EP each rank computes the delta only for its owned experts
|
||||
# (passed via local_expert_offset/local_num_experts below). gate_up_delta stays
|
||||
# new_empty even though non-owned [token, k] slots are then left unwritten -- the
|
||||
# trtllm MoE is itself EP-aware, so those slots never feed the all-reduced output.
|
||||
gate_up_delta_shape = (
|
||||
hidden_states.shape[0],
|
||||
runner_config.top_k,
|
||||
quant_info.w13_weight.shape[1],
|
||||
)
|
||||
gate_up_delta = hidden_states.new_empty(gate_up_delta_shape)
|
||||
|
||||
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=quant_info.local_num_experts,
|
||||
)
|
||||
|
||||
# GEMM1-LoRA overlap: fire the gate_up LoRA on the side stream + record an event; the
|
||||
# trtllm op waits on it right before activation (the only consumer of gate_up_delta), so
|
||||
# permute+GEMM1 overlap the side-stream LoRA shrink/expand instead of joining before the
|
||||
# whole op.
|
||||
lora_event = torch.cuda.Event()
|
||||
|
||||
# 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())
|
||||
with torch.cuda.stream(side_stream):
|
||||
_run_gate_up_lora()
|
||||
lora_event.record()
|
||||
|
||||
# Fuse the per-token scale transpose into the quant kernel: column-major scales make
|
||||
# the `.t()` a free view, dropping the standalone ~2us transpose+copy. The trtllm MoE
|
||||
# kernel wants the [K, M]-contiguous scale, which `.t()` of the column-major buffer is
|
||||
# exactly -- byte/shape-identical to the old `a_sf.t().contiguous()`.
|
||||
a_q, a_sf = per_token_group_quant_fp8(
|
||||
hidden_states, quant_info.weight_block_k, column_major_scales=True
|
||||
)
|
||||
a_sf_t = a_sf.t()
|
||||
|
||||
activation_lora_input = torch.empty(
|
||||
(hidden_states.shape[0], runner_config.top_k, quant_info.intermediate_size),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
# SGLANG_OPT_LORA_FUSED_TOPK_PACK: the routed pack may already have been produced
|
||||
# fused inside the gating kernel (StandardTopKOutput.packed_topk_ids) — including
|
||||
# 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(
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# No pre-op join: the trtllm op waits on lora_event right before its activation kernel,
|
||||
# so permute+GEMM1 run concurrent with the side-stream LoRA. Keep the event alive through
|
||||
# cuda-graph capture so the captured cross-stream wait isn't torn down before instantiation.
|
||||
if torch.cuda.is_current_stream_capturing():
|
||||
_LORA_OVERLAP_EVENTS.append(lora_event)
|
||||
lora_ready_handle = lora_event.cuda_event
|
||||
|
||||
# Down-LoRA/finalize overlap (env-gated): the op records gemm2_done_event right after the
|
||||
# base down GEMM (before finalize); the side stream waits on it and runs ONLY the down-proj
|
||||
# LoRA shrink (gemm A) + routing prep concurrent with finalizeKernel. The expand-add
|
||||
# (gemm B) atomic-adds into `output` -- which finalize WRITES concurrently -- so it stays
|
||||
# on the main stream after the op (post-finalize), exactly like the serial path.
|
||||
# DISABLED: the down/finalize overlap is bench-verified net-neutral-to-negative AND
|
||||
# corrupts the base/decode path — the captured gemm2_done cross-stream event under
|
||||
# cuda-graph replay perturbs no-active-LoRA (base) requests (qwen base gsm8k 0.81 -> 0.56
|
||||
# with it on; bisect-confirmed). The serial down-LoRA path below is used unconditionally.
|
||||
down_overlap = False
|
||||
gemm2_done_handle = 0
|
||||
if down_overlap:
|
||||
gemm2_done_event = torch.cuda.Event()
|
||||
# Materialize the underlying cudaEvent (torch creates it lazily on first record) so
|
||||
# .cuda_event is a real handle; the op re-records it after GEMM2.
|
||||
gemm2_done_event.record()
|
||||
gemm2_done_handle = gemm2_done_event.cuda_event
|
||||
|
||||
moe_result = trtllm_fp8_block_scale_routed_moe_lora(
|
||||
topk_ids=packed_topk_ids,
|
||||
routing_bias=None,
|
||||
hidden_states=a_q,
|
||||
hidden_states_scale=a_sf_t,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
|
||||
gate_up_lora_delta=gate_up_delta,
|
||||
activation_lora_input=activation_lora_input,
|
||||
lora_ready_event=lora_ready_handle,
|
||||
num_experts=quant_info.global_num_experts,
|
||||
top_k=runner_config.top_k,
|
||||
n_group=None,
|
||||
topk_group=None,
|
||||
intermediate_size=quant_info.intermediate_size,
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_experts,
|
||||
routed_scaling_factor=(
|
||||
runner_config.routed_scaling_factor
|
||||
if runner_config.routed_scaling_factor is not None
|
||||
else 1.0
|
||||
),
|
||||
routing_method_type=(
|
||||
RoutingMethodType.TopK
|
||||
if quant_info.routing_method_type == RoutingMethodType.DeepSeekV3
|
||||
else quant_info.routing_method_type
|
||||
),
|
||||
use_shuffled_weight=False,
|
||||
do_finalize=True,
|
||||
output=direct_down_output,
|
||||
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
|
||||
fp8_quantization_type=Fp8QuantizationType.DeepSeekFp8,
|
||||
activation_type=quant_info.activation_type,
|
||||
gemm2_done_event=gemm2_done_handle,
|
||||
)
|
||||
|
||||
output = moe_result
|
||||
|
||||
def _run_down_lora(
|
||||
out, stage="all", intermediate_buffer=None, expand_wait_event=None
|
||||
):
|
||||
return merged_experts_fused_moe_lora_add(
|
||||
output=out,
|
||||
hidden_states=activation_lora_input.view(-1, quant_info.intermediate_size),
|
||||
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=quant_info.local_num_experts,
|
||||
stage=stage,
|
||||
intermediate_buffer=intermediate_buffer,
|
||||
expand_wait_event=expand_wait_event,
|
||||
)
|
||||
|
||||
# Shared-add overlap: the trtllm op above already finalized `output`, so the
|
||||
# staged shared-expert add (if any) can run on the producer (main) stream
|
||||
# concurrent with the down-LoRA shrink below; the expand waits on it via
|
||||
# expand_wait_event before atomic-adding into the same buffer.
|
||||
shared_add_done = maybe_overlap_staged_shared_add(output)
|
||||
|
||||
if down_overlap:
|
||||
# Fork at "base down GEMM done": ONLY the shrink (gemm A) + routing prep run on the
|
||||
# side stream, concurrent with the main-stream finalizeKernel. The expand-add (gemm B)
|
||||
# joins back on the MAIN stream after the op -- i.e. strictly after finalize wrote
|
||||
# `output` -- and atomic-adds into it exactly like the serial path: same kernels,
|
||||
# same buffers, identical numerics; the shrink just starts earlier.
|
||||
# The shrink intermediate is allocated HERE (main = consumer stream of the expand),
|
||||
# per the SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC lesson on side-stream allocations.
|
||||
down_intermediate = hidden_states.new_empty(
|
||||
(
|
||||
hidden_states.shape[0],
|
||||
topk_ids.shape[1],
|
||||
lora_info.down_lora_a_weights.shape[2],
|
||||
)
|
||||
)
|
||||
side_stream.wait_event(gemm2_done_event)
|
||||
shrink_done_event = torch.cuda.Event()
|
||||
with torch.cuda.stream(side_stream):
|
||||
_run_down_lora(
|
||||
output, stage="shrink", intermediate_buffer=down_intermediate
|
||||
)
|
||||
shrink_done_event.record()
|
||||
if torch.cuda.is_current_stream_capturing():
|
||||
_LORA_OVERLAP_EVENTS.append(gemm2_done_event)
|
||||
_LORA_OVERLAP_EVENTS.append(shrink_done_event)
|
||||
torch.cuda.current_stream().wait_event(shrink_done_event)
|
||||
_run_down_lora(
|
||||
output,
|
||||
stage="expand",
|
||||
intermediate_buffer=down_intermediate,
|
||||
expand_wait_event=shared_add_done,
|
||||
)
|
||||
else:
|
||||
_run_down_lora(output, expand_wait_event=shared_add_done)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
|
||||
dispatch_output,
|
||||
quant_info,
|
||||
runner_config,
|
||||
lora_info,
|
||||
):
|
||||
"""Two-stream NVFP4 sibling of the FP8 two-stream MoE LoRA dispatch.
|
||||
|
||||
Fires only for virtual-experts LoRA + decode-shaped batches; everything else
|
||||
delegates to the saved-original single-stream FP4 dispatch (byte-identical).
|
||||
"""
|
||||
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_fp4_moe_lora_func()(
|
||||
dispatch_output, quant_info, runner_config, lora_info
|
||||
)
|
||||
|
||||
# ---- two-stream fast path ----
|
||||
from sglang.jit_kernel.trtllm_lora_temp import (
|
||||
trtllm_fp4_block_scale_routed_moe_lora,
|
||||
)
|
||||
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 (
|
||||
merged_experts_fused_moe_lora_add,
|
||||
)
|
||||
|
||||
assert (
|
||||
runner_config.activation == "silu" and runner_config.is_gated
|
||||
), "experimental_sgl_trtllm NVFP4 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 = {}
|
||||
|
||||
# Down-proj LoRA runs serially on the main stream (after the trtllm op) by default. The old
|
||||
# side-stream down-overlap was removed: bench-verified net-neutral-to-negative (the extra
|
||||
# side-stream all-reduce cancels any overlap gain), AND its act_ready_event cross-stream sync
|
||||
# corrupted decode state under sustained heavy LoRA load (cuda-graph replay -> persistent
|
||||
# garbage). SGLANG_OPT_LORA_DOWN_FINALIZE_OVERLAP=1 re-introduces a more conservative variant:
|
||||
# fork at gemm2_done (not act_ready), and only the SHRINK (gemm A) overlaps the finalize
|
||||
# kernel -- the expand-add (gemm B) stays on the MAIN stream post-finalize (it writes the
|
||||
# same `output` finalize writes), so its kernels/numerics match the serial path exactly.
|
||||
inter = quant_info.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, quant_info.w13_weight.shape[1])
|
||||
)
|
||||
|
||||
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=quant_info.local_num_experts,
|
||||
)
|
||||
|
||||
# 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).
|
||||
lora_event = torch.cuda.Event()
|
||||
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 = _pack_topk_for_flashinfer_routed(
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
)
|
||||
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 FP4 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
|
||||
|
||||
# Down-LoRA/finalize overlap (env-gated, see the comment above): record gemm2_done inside
|
||||
# the op; the side stream runs only the down-LoRA shrink + routing prep concurrent with
|
||||
# finalize; the expand-add joins back on the main stream after the op.
|
||||
# DISABLED: the down/finalize overlap is bench-verified net-neutral-to-negative AND
|
||||
# corrupts the base/decode path — the captured gemm2_done cross-stream event under
|
||||
# cuda-graph replay perturbs no-active-LoRA (base) requests (qwen base gsm8k 0.81 -> 0.56
|
||||
# with it on; bisect-confirmed). The serial down-LoRA path below is used unconditionally.
|
||||
down_overlap = False
|
||||
gemm2_done_handle = 0
|
||||
if down_overlap:
|
||||
gemm2_done_event = torch.cuda.Event()
|
||||
# Materialize the underlying cudaEvent (torch creates it lazily on first record) so
|
||||
# .cuda_event is a real handle; the op re-records it after the down GEMM.
|
||||
gemm2_done_event.record()
|
||||
gemm2_done_handle = gemm2_done_event.cuda_event
|
||||
|
||||
output = trtllm_fp4_block_scale_routed_moe_lora(
|
||||
topk_ids=packed_topk_ids,
|
||||
routing_bias=None,
|
||||
hidden_states=hidden_states,
|
||||
hidden_states_scale=None,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn),
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn),
|
||||
output1_scales_scalar=quant_info.g1_scale_c,
|
||||
output1_scales_gate_scalar=quant_info.g1_alphas,
|
||||
output2_scales_scalar=quant_info.g2_alphas,
|
||||
gate_up_lora_delta=gate_up_delta,
|
||||
activation_lora_input=activation_lora_input,
|
||||
lora_ready_event=lora_ready_handle,
|
||||
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=quant_info.local_num_experts,
|
||||
routed_scaling_factor=(
|
||||
runner_config.routed_scaling_factor
|
||||
if runner_config.routed_scaling_factor is not None
|
||||
else 1.0
|
||||
),
|
||||
routing_method_type=quant_info.routing_method_type,
|
||||
do_finalize=True,
|
||||
output=direct_down_output,
|
||||
gemm2_done_event=gemm2_done_handle,
|
||||
)
|
||||
|
||||
def _run_down_lora(out, stage="all", intermediate_buffer=None):
|
||||
return merged_experts_fused_moe_lora_add(
|
||||
output=out,
|
||||
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=quant_info.local_num_experts,
|
||||
stage=stage,
|
||||
intermediate_buffer=intermediate_buffer,
|
||||
)
|
||||
|
||||
if down_overlap:
|
||||
# Fork at "base down GEMM done": shrink (gemm A) + routing prep on the side stream
|
||||
# concurrent with the main-stream finalize; the expand-add (gemm B) joins back on the
|
||||
# MAIN stream after the op (strictly post-finalize) -- serial-path kernels/numerics.
|
||||
# Intermediate allocated on main (= consumer stream of the expand), per the
|
||||
# SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC lesson.
|
||||
down_intermediate = hidden_states.new_empty(
|
||||
(
|
||||
hidden_states.shape[0],
|
||||
topk_ids.shape[1],
|
||||
lora_info.down_lora_a_weights.shape[2],
|
||||
)
|
||||
)
|
||||
side_stream.wait_event(gemm2_done_event)
|
||||
shrink_done_event = torch.cuda.Event()
|
||||
with torch.cuda.stream(side_stream):
|
||||
_run_down_lora(
|
||||
output, stage="shrink", intermediate_buffer=down_intermediate
|
||||
)
|
||||
shrink_done_event.record()
|
||||
if torch.cuda.is_current_stream_capturing():
|
||||
_LORA_OVERLAP_EVENTS.append(gemm2_done_event)
|
||||
_LORA_OVERLAP_EVENTS.append(shrink_done_event)
|
||||
torch.cuda.current_stream().wait_event(shrink_done_event)
|
||||
_run_down_lora(output, stage="expand", intermediate_buffer=down_intermediate)
|
||||
else:
|
||||
_run_down_lora(output)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Registers the ``experimental_sgl_trtllm`` MoE fused-func.
|
||||
|
||||
``MoeRunner.__init__`` requires a registered fused-func at CONSTRUCTION time even
|
||||
for the LoRA case, because LoRA is attached *after* the MoE layer is built (so
|
||||
``lora_enabled`` is False inside ``MoeRunner.__init__``). At run time the runner
|
||||
skips this for the LoRA path; the no-LoRA path delegates entirely to the upstream
|
||||
flashinfer_trtllm dispatch (all quant types), so no-LoRA is identical to the stock backend.
|
||||
|
||||
Registration fires at model-build time via a one-line import of this module in
|
||||
``moe_runner/flashinfer_trtllm.py`` (the module already imported there for the
|
||||
trtllm weight-prep). Keeping the dispatch body here keeps that file otherwise
|
||||
pristine; the sgl FP8 LoRA dispatch lives in ``sgl_fp8_moe.py`` (used only by the LoRA path).
|
||||
"""
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import register_fused_func
|
||||
|
||||
|
||||
@register_fused_func("none", "experimental_sgl_trtllm")
|
||||
def fused_experts_none_to_experimental_sgl_trtllm(
|
||||
dispatch_output, quant_info, runner_config
|
||||
):
|
||||
# No-LoRA on the experimental_sgl_trtllm backend == upstream flashinfer_trtllm for EVERY
|
||||
# quant type (FP8 / NVFP4 / bf16). When LoRA is disabled the runner calls this fused-func,
|
||||
# so delegating entirely to upstream keeps the no-LoRA path byte-identical to the stock
|
||||
# backend. The new sgl kernels (sgl_fp8_moe, trtllm_*_routed_moe_lora) run ONLY on the LoRA
|
||||
# dispatch (lora_dispatch.py), never here.
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
fused_experts_none_to_flashinfer_trtllm,
|
||||
)
|
||||
|
||||
return fused_experts_none_to_flashinfer_trtllm(
|
||||
dispatch_output, quant_info, runner_config
|
||||
)
|
||||
@@ -0,0 +1,248 @@
|
||||
"""Copy of upstream flashinfer-trtllm FP8 MoE dispatch, wired to experimental_sgl_trtllm_moe
|
||||
block-scale wrappers (LoRA-capable) so moe_runner/flashinfer_trtllm.py stays pristine. Body is
|
||||
verbatim from upstream; helper imports are call-time (cycle-safe); two FP8 wrappers shadowed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
FlashInferTrtllmFp8MoeQuantInfo,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import (
|
||||
StandardCombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
|
||||
|
||||
def fused_experts_fp8_sgl(
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
quant_info: FlashInferTrtllmFp8MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
use_routed_topk: bool = False,
|
||||
) -> StandardCombineInput:
|
||||
# Lazy (call-time) imports so this module never triggers the flashinfer_trtllm
|
||||
# <-> quantization import cycle at load time.
|
||||
from flashinfer.fused_moe import Fp8QuantizationType
|
||||
|
||||
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,
|
||||
per_token_group_quant_fp8,
|
||||
scaled_fp8_quant,
|
||||
trtllm_fp8_per_tensor_scale_moe_wrapper,
|
||||
use_symmetric_memory,
|
||||
)
|
||||
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.experimental_sgl_trtllm_moe import (
|
||||
sgl_trtllm_fp8_block_scale_moe_wrapper as trtllm_fp8_block_scale_moe_wrapper,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp.experimental_sgl_trtllm_moe import (
|
||||
sgl_trtllm_fp8_block_scale_routed_moe_wrapper as trtllm_fp8_block_scale_routed_moe_wrapper,
|
||||
)
|
||||
|
||||
_SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"}
|
||||
assert runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, (
|
||||
f"Only {_SUPPORTED_FP8_ACTIVATIONS} are supported for FP8 MoE, "
|
||||
f"got '{runner_config.activation}'."
|
||||
)
|
||||
assert not runner_config.no_combine, "no_combine is not supported for flashinfer."
|
||||
|
||||
hidden_states = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
if TopKOutputChecker.format_is_bypassed(topk_output):
|
||||
router_logits = topk_output.router_logits
|
||||
topk_config = topk_output.topk_config
|
||||
correction_bias = (
|
||||
None
|
||||
if topk_config.correction_bias is None
|
||||
else topk_config.correction_bias.to(hidden_states.dtype)
|
||||
)
|
||||
else:
|
||||
router_logits = None
|
||||
topk_config = None
|
||||
correction_bias = None
|
||||
|
||||
routing_method_type = quant_info.routing_method_type
|
||||
fp8_quantization_type = (
|
||||
Fp8QuantizationType.MxFp8
|
||||
if quant_info.use_mxfp8
|
||||
else Fp8QuantizationType.DeepSeekFp8
|
||||
)
|
||||
use_shuffled_weight = quant_info.use_mxfp8
|
||||
|
||||
if quant_info.block_quant:
|
||||
assert quant_info.weight_block_k is not None
|
||||
assert quant_info.w13_weight_scale_inv is not None
|
||||
assert quant_info.w2_weight_scale_inv is not None
|
||||
|
||||
if quant_info.use_mxfp8:
|
||||
assert quant_info.weight_block_k == 32
|
||||
from flashinfer import mxfp8_quantize
|
||||
|
||||
a_q, a_sf = mxfp8_quantize(hidden_states, False)
|
||||
# FlashInfer TRT-LLM MxFP8 expects token-major activation scales:
|
||||
# [num_tokens, hidden_size // 32] (no transpose).
|
||||
a_sf_t = a_sf.view(torch.uint8).reshape(hidden_states.shape[0], -1)
|
||||
else:
|
||||
a_q, a_sf = per_token_group_quant_fp8(
|
||||
hidden_states, quant_info.weight_block_k
|
||||
)
|
||||
a_sf_t = a_sf.t().contiguous()
|
||||
|
||||
# Allocate output inside symmetric memory context
|
||||
with use_symmetric_memory(
|
||||
get_tp_group(), disabled=not is_allocation_symmetric()
|
||||
):
|
||||
symm_output = torch.empty(
|
||||
hidden_states.shape[0],
|
||||
hidden_states.shape[1],
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
# Move kernel call outside context manager to avoid graph breaks
|
||||
# during torch.compile for piecewise cuda graph.
|
||||
# Use custom op wrapper for torch.compile compatibility.
|
||||
if use_routed_topk:
|
||||
assert (
|
||||
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(
|
||||
topk_ids=topk_output.topk_ids,
|
||||
topk_weights=topk_output.topk_weights,
|
||||
)
|
||||
|
||||
output = trtllm_fp8_block_scale_routed_moe_wrapper(
|
||||
topk_ids=packed_topk_ids,
|
||||
routing_bias=None,
|
||||
hidden_states=a_q,
|
||||
hidden_states_scale=a_sf_t,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
|
||||
num_experts=quant_info.global_num_experts,
|
||||
top_k=runner_config.top_k,
|
||||
n_group=None,
|
||||
topk_group=None,
|
||||
intermediate_size=quant_info.intermediate_size,
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_experts,
|
||||
routed_scaling_factor=(
|
||||
runner_config.routed_scaling_factor
|
||||
if runner_config.routed_scaling_factor is not None
|
||||
else 1.0
|
||||
),
|
||||
routing_method_type=(
|
||||
RoutingMethodType.TopK
|
||||
if routing_method_type == RoutingMethodType.DeepSeekV3
|
||||
else routing_method_type
|
||||
),
|
||||
use_shuffled_weight=use_shuffled_weight,
|
||||
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
|
||||
fp8_quantization_type=int(fp8_quantization_type),
|
||||
activation_type=quant_info.activation_type,
|
||||
)
|
||||
else:
|
||||
assert TopKOutputChecker.format_is_bypassed(topk_output)
|
||||
|
||||
output = trtllm_fp8_block_scale_moe_wrapper(
|
||||
routing_logits=router_logits,
|
||||
routing_bias=correction_bias,
|
||||
hidden_states=a_q,
|
||||
hidden_states_scale=a_sf_t,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
|
||||
num_experts=quant_info.global_num_experts,
|
||||
top_k=topk_config.top_k,
|
||||
n_group=topk_config.num_expert_group,
|
||||
topk_group=topk_config.topk_group,
|
||||
intermediate_size=quant_info.intermediate_size,
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_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,
|
||||
use_shuffled_weight=use_shuffled_weight,
|
||||
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
|
||||
fp8_quantization_type=int(fp8_quantization_type),
|
||||
activation_type=quant_info.activation_type,
|
||||
)
|
||||
# TODO: Once https://github.com/flashinfer-ai/flashinfer/issues/2703 is fixed, pass output to moe kernel and remove this copy.
|
||||
symm_output.copy_(output)
|
||||
output = symm_output
|
||||
else:
|
||||
assert TopKOutputChecker.format_is_bypassed(topk_output)
|
||||
assert quant_info.w13_input_scale is not None
|
||||
assert quant_info.output1_scales_scalar is not None
|
||||
assert quant_info.output1_scales_gate_scalar is not None
|
||||
assert quant_info.output2_scales_scalar is not None
|
||||
|
||||
a_q, _ = scaled_fp8_quant(hidden_states, quant_info.w13_input_scale)
|
||||
routing_bias_cast = (
|
||||
None if correction_bias is None else correction_bias.to(torch.bfloat16)
|
||||
)
|
||||
|
||||
# Allocate output inside symmetric memory context
|
||||
with use_symmetric_memory(
|
||||
get_tp_group(), disabled=not is_allocation_symmetric()
|
||||
):
|
||||
symm_output = torch.empty(
|
||||
hidden_states.shape[0],
|
||||
hidden_states.shape[1],
|
||||
dtype=torch.bfloat16,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
# Move kernel call outside context manager to avoid graph breaks
|
||||
# during torch.compile for piecewise cuda graph.
|
||||
# Use custom op wrapper for torch.compile compatibility.
|
||||
|
||||
router_logits = router_logits.to(torch.bfloat16)
|
||||
|
||||
output = trtllm_fp8_per_tensor_scale_moe_wrapper(
|
||||
routing_logits=router_logits,
|
||||
routing_bias=routing_bias_cast,
|
||||
hidden_states=a_q,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
output1_scales_scalar=quant_info.output1_scales_scalar,
|
||||
output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar,
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
output2_scales_scalar=quant_info.output2_scales_scalar,
|
||||
num_experts=quant_info.global_num_experts,
|
||||
top_k=topk_config.top_k,
|
||||
n_group=topk_config.num_expert_group,
|
||||
topk_group=topk_config.topk_group,
|
||||
intermediate_size=int(quant_info.w2_weight.shape[2]),
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_experts,
|
||||
routed_scaling_factor=(
|
||||
runner_config.routed_scaling_factor
|
||||
if runner_config.routed_scaling_factor is not None
|
||||
else 1.0
|
||||
),
|
||||
use_routing_scales_on_input=False,
|
||||
routing_method_type=routing_method_type,
|
||||
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
|
||||
activation_type=quant_info.activation_type,
|
||||
)
|
||||
symm_output.copy_(output)
|
||||
output = symm_output
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
@@ -0,0 +1,116 @@
|
||||
"""Overlap the MoE shared-expert add with the down-LoRA shrink (gemm A).
|
||||
|
||||
In the Qwen dual-stream decode path (``Qwen2MoeSparseMoeBlock.forward_normal_dual_stream``)
|
||||
the shared expert runs on the main stream while the routed experts + trtllm LoRA
|
||||
MoE run on the alt stream; ``final_hidden_states += shared_output`` then runs on
|
||||
the main stream only after the WHOLE alt-stream chain (base MoE -> finalize ->
|
||||
down-LoRA shrink -> down-LoRA expand) joins — putting a ~2us elementwise add at
|
||||
the tail of every MoE layer's critical path.
|
||||
|
||||
The add's real dependency is only the base-MoE finalize (the trtllm op with
|
||||
``do_finalize=True`` writing the output buffer): the down-LoRA shrink writes a
|
||||
separate intermediate, and the down-LoRA expand atomic-adds into the same output
|
||||
buffer, so addition order is commutative — the only hard constraint is that the
|
||||
non-atomic shared add and the expand must not run CONCURRENTLY on the buffer.
|
||||
(The cross-rank ``tensor_model_parallel_all_reduce`` happens after all of this
|
||||
at the model layer, unchanged.)
|
||||
|
||||
Protocol (gated by ``SGLANG_OPT_LORA_SHARED_ADD_OVERLAP``):
|
||||
|
||||
1. The model layer stages ``(shared_output, producer_stream)`` via
|
||||
:func:`stage_shared_expert_add` after computing the shared expert and before
|
||||
forking the routed experts to the alt stream.
|
||||
2. The LoRA dispatch, right after the trtllm op returns (finalize done), calls
|
||||
:func:`maybe_overlap_staged_shared_add(output)`: it records ``base_ready`` on
|
||||
the current (alt) stream, enqueues ``wait(base_ready); output += shared;
|
||||
record(add_done)`` on the producer (main) stream — main-stream program order
|
||||
already guarantees ``shared_output`` is ready there — and returns ``add_done``.
|
||||
3. The down-LoRA ``merged_experts_fused_moe_lora_add`` waits on ``add_done``
|
||||
right before launching the expand kernel, so the shrink (+ stage-B routing)
|
||||
overlaps the add and the expand never races it.
|
||||
4. After the dual-stream join the model layer calls
|
||||
:func:`unstage_shared_expert_add`; if the dispatch never consumed the staging
|
||||
(prefill / non-virtual-store / fallback paths) it gets the tensor back and
|
||||
performs the original add itself — byte-identical fallback behavior.
|
||||
|
||||
The state is a single slot: MoE layers run sequentially within one scheduler
|
||||
process, and the stage/consume pair lives within a single layer forward.
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
|
||||
_PENDING: Optional[Tuple[torch.Tensor, torch.cuda.Stream]] = None
|
||||
|
||||
# Keep events recorded during cuda-graph capture alive so the captured
|
||||
# cross-stream waits aren't torn down before graph instantiation (same pattern
|
||||
# as moe_overlap._LORA_OVERLAP_EVENTS). Eager runs rely on deferred destroy.
|
||||
_SHARED_ADD_EVENTS: list = []
|
||||
|
||||
|
||||
def shared_add_overlap_enabled() -> bool:
|
||||
return lora_envs.SGLANG_OPT_LORA_SHARED_ADD_OVERLAP.get()
|
||||
|
||||
|
||||
def stage_shared_expert_add(
|
||||
shared_output: torch.Tensor, producer_stream: torch.cuda.Stream
|
||||
) -> None:
|
||||
"""Stage the shared-expert output for the LoRA dispatch to add.
|
||||
|
||||
``producer_stream`` is the stream ``shared_output`` was computed on (the
|
||||
main stream); the overlapped add is enqueued there so its data dependency
|
||||
on the shared expert is carried by stream program order.
|
||||
"""
|
||||
global _PENDING
|
||||
_PENDING = (shared_output, producer_stream)
|
||||
|
||||
|
||||
def unstage_shared_expert_add() -> Optional[torch.Tensor]:
|
||||
"""Reclaim a staged-but-unconsumed shared add (fallback paths).
|
||||
|
||||
Returns the staged tensor if the dispatch did NOT consume it (the model
|
||||
layer must then do the add itself), or None if it was consumed (the add is
|
||||
already enqueued on the producer stream).
|
||||
"""
|
||||
global _PENDING
|
||||
if _PENDING is None:
|
||||
return None
|
||||
shared_output, _ = _PENDING
|
||||
_PENDING = None
|
||||
return shared_output
|
||||
|
||||
|
||||
def maybe_overlap_staged_shared_add(output: torch.Tensor) -> Optional[torch.cuda.Event]:
|
||||
"""Enqueue the staged shared-expert add overlapped with the down-LoRA shrink.
|
||||
|
||||
Call from the LoRA dispatch right after the base-MoE finalize has been
|
||||
enqueued on the current stream with ``output`` fully written. Returns the
|
||||
``add_done`` event the down-LoRA expand must wait on before atomic-adding
|
||||
into ``output``, or None when nothing was staged.
|
||||
"""
|
||||
global _PENDING
|
||||
if _PENDING is None:
|
||||
return None
|
||||
shared_output, producer_stream = _PENDING
|
||||
current_stream = torch.cuda.current_stream()
|
||||
if producer_stream == current_stream:
|
||||
# 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
|
||||
_PENDING = None
|
||||
|
||||
base_ready = torch.cuda.Event()
|
||||
base_ready.record(current_stream)
|
||||
add_done = torch.cuda.Event()
|
||||
with torch.cuda.stream(producer_stream):
|
||||
producer_stream.wait_event(base_ready)
|
||||
output.add_(shared_output)
|
||||
add_done.record(producer_stream)
|
||||
|
||||
if torch.cuda.is_current_stream_capturing():
|
||||
_SHARED_ADD_EVENTS.append(base_ready)
|
||||
_SHARED_ADD_EVENTS.append(add_done)
|
||||
return add_done
|
||||
@@ -0,0 +1,234 @@
|
||||
"""Rank-specialized LoRA-B expand for virtual-expert LoRA.
|
||||
|
||||
The kernel here was originally a chunk in ``lora/triton_ops/virtual_experts.py``.
|
||||
It is rank-specialized: the ``R`` dimension (LoRA rank) is a triton
|
||||
``constexpr``, so each rank value used at runtime gets its own JIT-compiled
|
||||
specialization (R=16, R=32, R=64 are all supported up to the ``R <= 64`` assert,
|
||||
with no perf interaction between them — each gets its own kernel).
|
||||
|
||||
Called from :mod:`sglang.srt.lora.triton_ops.virtual_experts` when
|
||||
``use_direct_expand_add=True`` (the trtllm-lora path uses this when
|
||||
``max_lora_rank <= 64``); the generic ``invoke_fused_moe_kernel`` is used
|
||||
when that flag is False (incl. ranks above 64).
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _moe_lora_expand_add_kernel(
|
||||
# Pointers
|
||||
a_ptr, # [num_tokens * top_k, rank]
|
||||
b_ptr, # [num_virtual_experts, N, rank]
|
||||
c_ptr, # [num_tokens, N]
|
||||
topk_weights_ptr,
|
||||
sorted_token_ids_ptr,
|
||||
expert_ids_ptr,
|
||||
num_tokens_post_padded_ptr,
|
||||
# Dimensions
|
||||
N,
|
||||
R: tl.constexpr,
|
||||
num_valid_tokens,
|
||||
# Strides
|
||||
stride_am,
|
||||
stride_ar,
|
||||
stride_be,
|
||||
stride_bn,
|
||||
stride_br,
|
||||
stride_cm,
|
||||
stride_cn,
|
||||
# Constexprs
|
||||
router_topk: tl.constexpr,
|
||||
MUL_ROUTED_WEIGHT: tl.constexpr,
|
||||
FUSE_SUM_ALL_REDUCE: tl.constexpr,
|
||||
BLOCK_SIZE_M: tl.constexpr,
|
||||
BLOCK_SIZE_N: tl.constexpr,
|
||||
BLOCK_SIZE_R: tl.constexpr,
|
||||
GROUP_SIZE_M: tl.constexpr,
|
||||
GATED_A_HALF: tl.constexpr,
|
||||
):
|
||||
"""Rank-specialized LoRA-B expand for virtual-expert LoRA.
|
||||
|
||||
``GATED_A_HALF`` > 0 enables the gate/up split for a gated (SwiGLU) gate_up
|
||||
LoRA: the intermediate (A) has ``2*R`` columns (gate-shrink ``[0:R]`` then
|
||||
up-shrink ``[R:2R]``) and the output has ``2*GATED_A_HALF`` columns (gate
|
||||
then up). Output tiles in the up half (column >= ``GATED_A_HALF``) read the
|
||||
up-shrink columns ``[R:2R]`` instead of ``[0:R]``. ``GATED_A_HALF`` must be
|
||||
a multiple of ``BLOCK_SIZE_N`` so no tile straddles the gate/up boundary.
|
||||
``GATED_A_HALF == 0`` is the non-gated path (read ``[0:R]`` for all tiles).
|
||||
"""
|
||||
pid = tl.program_id(0)
|
||||
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
|
||||
num_pid_m = tl.cdiv(num_tokens_post_padded, BLOCK_SIZE_M)
|
||||
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
|
||||
|
||||
num_pid_in_group = GROUP_SIZE_M * num_pid_n
|
||||
group_id = pid // num_pid_in_group
|
||||
first_pid_m = group_id * GROUP_SIZE_M
|
||||
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
|
||||
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
|
||||
pid_n = (pid % num_pid_in_group) // group_size_m
|
||||
|
||||
if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
|
||||
return
|
||||
|
||||
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
|
||||
offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
|
||||
token_mask = offs_token < num_valid_tokens
|
||||
|
||||
off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
|
||||
if off_expert == -1:
|
||||
if not FUSE_SUM_ALL_REDUCE:
|
||||
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
|
||||
c_ptrs = (
|
||||
c_ptr + offs_token[:, None] * stride_cm + offs_n[None, :] * stride_cn
|
||||
)
|
||||
c_mask = token_mask[:, None] & (offs_n[None, :] < N)
|
||||
zeros = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=c_ptr.dtype.element_ty)
|
||||
tl.store(c_ptrs, zeros, mask=c_mask)
|
||||
return
|
||||
|
||||
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
|
||||
offs_r = tl.arange(0, BLOCK_SIZE_R).to(tl.int64)
|
||||
rank_mask = offs_r < R
|
||||
|
||||
# Gated gate_up split: the up-half output tiles read up-shrink (A columns [R:2R]); gate-half
|
||||
# tiles read gate-shrink (A columns [0:R]). GATED_A_HALF == 0 -> always read [0:R] (non-gated).
|
||||
a_col = offs_r
|
||||
if GATED_A_HALF > 0:
|
||||
a_col = offs_r + tl.where(pid_n * BLOCK_SIZE_N >= GATED_A_HALF, R, 0)
|
||||
|
||||
a = tl.load(
|
||||
a_ptr + offs_token[:, None] * stride_am + a_col[None, :] * stride_ar,
|
||||
mask=token_mask[:, None] & rank_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
b = tl.load(
|
||||
b_ptr
|
||||
+ off_expert * stride_be
|
||||
+ offs_n[None, :] * stride_bn
|
||||
+ offs_r[:, None] * stride_br,
|
||||
mask=(offs_n[None, :] < N) & rank_mask[:, None],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
accumulator = tl.dot(a, b, out_dtype=tl.float32)
|
||||
if MUL_ROUTED_WEIGHT:
|
||||
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
|
||||
accumulator *= moe_weight[:, None]
|
||||
|
||||
if FUSE_SUM_ALL_REDUCE:
|
||||
offs_token_out = offs_token // router_topk
|
||||
else:
|
||||
offs_token_out = offs_token
|
||||
c_ptrs = c_ptr + offs_token_out[:, None] * stride_cm + offs_n[None, :] * stride_cn
|
||||
c_mask = token_mask[:, None] & (offs_n[None, :] < N)
|
||||
if FUSE_SUM_ALL_REDUCE:
|
||||
tl.atomic_add(c_ptrs, accumulator.to(c_ptr.dtype.element_ty), mask=c_mask)
|
||||
else:
|
||||
tl.store(c_ptrs, accumulator.to(c_ptr.dtype.element_ty), mask=c_mask)
|
||||
|
||||
|
||||
def _invoke_moe_lora_expand_add(
|
||||
intermediate: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
sorted_token_ids: torch.Tensor,
|
||||
expert_ids: torch.Tensor,
|
||||
num_tokens_post_padded: torch.Tensor,
|
||||
config: "dict[str, Any]",
|
||||
mul_routed_weight: bool,
|
||||
fuse_sum_all_reduce: bool,
|
||||
force_block_size_n: "int | None" = None,
|
||||
) -> None:
|
||||
"""Launch the rank-specialized LoRA-B expand kernel.
|
||||
|
||||
``R`` (= ``weight.shape[2]``) up to 64 is supported. ``BLOCK_SIZE_R`` is
|
||||
set to ``next_power_of_2(R)`` so each rank value pairs with the smallest
|
||||
tile that covers it (R=16 → BSR=16, R=32 → BSR=32, R=64 → BSR=64).
|
||||
Triton compiles a separate specialization per (R, BLOCK_SIZE_R) combo
|
||||
so different ranks don't interfere with each other's perf.
|
||||
"""
|
||||
N = weight.shape[1]
|
||||
R = weight.shape[2]
|
||||
assert R <= 64, f"direct LoRA expand/add expects rank <= 64, got {R}"
|
||||
|
||||
block_size_m = config["BLOCK_SIZE_M"]
|
||||
# BLOCK_SIZE_N defaults to 128 when N % 128 == 0 (a good N-divisible tile that also keeps
|
||||
# the gated gate_up split aligned). ``force_block_size_n`` lets a tuner/bench override it;
|
||||
# down-proj has no gate/up boundary, so any divisor of N is valid there.
|
||||
if force_block_size_n is not None:
|
||||
block_size_n = force_block_size_n
|
||||
else:
|
||||
block_size_n = 128 if N % 128 == 0 else config["BLOCK_SIZE_N"]
|
||||
group_size_m = config.get("GROUP_SIZE_M", 1)
|
||||
block_size_r = triton.next_power_of_2(R)
|
||||
|
||||
# gate_up LoRA: the shrink stacks gate_A and up_A, so the intermediate has 2*R columns
|
||||
# ([0:R] = gate-shrink x@gate_A^T, [R:2R] = up-shrink x@up_A^T). The up output half
|
||||
# (column >= N/2) must contract the up-shrink [R:2R], not gate_A's [0:R]. Reading [0:R]
|
||||
# for both halves (the previous hardcode) computed the up delta from gate_A and dropped
|
||||
# up_A -- wrong whenever gate_A != up_A (the normal independently-trained gate/up case;
|
||||
# verified >100% rel error vs a PEFT reference on the real Qwen3.5 adapter). The earlier
|
||||
# "vs cutlass" justification for reading [0:R] was unreliable (the cutlass reference shared
|
||||
# the same bug). Detect the gated layout from the intermediate width and split in-kernel.
|
||||
inter_width = intermediate.shape[1]
|
||||
assert inter_width in (R, 2 * R), (
|
||||
f"LoRA expand intermediate width must be R ({R}, non-gated) or 2*R "
|
||||
f"({2 * R}, gated gate_up), got {inter_width}"
|
||||
)
|
||||
# Lazy import to avoid the trtllm_moe <-> triton_ops package import cycle at load time.
|
||||
|
||||
gated = inter_width == 2 * R
|
||||
use_gated_split = (
|
||||
gated and lora_envs.SGLANG_ENABLE_LORA_MOE_GATEUP_GATED_SPLIT.get()
|
||||
)
|
||||
gated_a_half = (N // 2) if use_gated_split else 0
|
||||
if use_gated_split:
|
||||
assert N % 2 == 0 and (N // 2) % block_size_n == 0, (
|
||||
f"gated gate_up split needs N/2 ({N // 2}) divisible by BLOCK_SIZE_N "
|
||||
f"({block_size_n})"
|
||||
)
|
||||
|
||||
grid = (
|
||||
triton.cdiv(sorted_token_ids.shape[0], block_size_m)
|
||||
* triton.cdiv(N, block_size_n),
|
||||
)
|
||||
|
||||
_moe_lora_expand_add_kernel[grid](
|
||||
intermediate,
|
||||
weight,
|
||||
output,
|
||||
topk_weights,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
N,
|
||||
R,
|
||||
topk_ids.numel(),
|
||||
intermediate.stride(0),
|
||||
intermediate.stride(1),
|
||||
weight.stride(0),
|
||||
weight.stride(1),
|
||||
weight.stride(2),
|
||||
output.stride(-2),
|
||||
output.stride(-1),
|
||||
router_topk=topk_ids.shape[1],
|
||||
MUL_ROUTED_WEIGHT=mul_routed_weight,
|
||||
FUSE_SUM_ALL_REDUCE=fuse_sum_all_reduce,
|
||||
BLOCK_SIZE_M=block_size_m,
|
||||
BLOCK_SIZE_N=block_size_n,
|
||||
BLOCK_SIZE_R=block_size_r,
|
||||
GROUP_SIZE_M=group_size_m,
|
||||
GATED_A_HALF=gated_a_half,
|
||||
num_warps=config.get("num_warps", 4),
|
||||
num_stages=1,
|
||||
)
|
||||
@@ -0,0 +1,32 @@
|
||||
"""Experimental TRT-LLM LoRA kernels (copies of the upstream triton_ops kernels
|
||||
with the SGLANG_EXPERIMENTAL_LORA_OPTI optimizations).
|
||||
|
||||
These are forked from ``sglang.srt.lora.triton_ops`` so the upstream kernels stay
|
||||
byte-pristine; only the experimental forwards / dispatch in this package import
|
||||
from here. The opt branches inside each kernel are still gated by ``lora_envs``
|
||||
(master-gated by SGLANG_EXPERIMENTAL_LORA_OPTI).
|
||||
"""
|
||||
|
||||
from .gate_up_lora_b import gate_up_lora_b_fwd
|
||||
from .kv_b_lora_absorbed import (
|
||||
step_a_q_fwd,
|
||||
step_a_v_fwd,
|
||||
step_b_q_fwd,
|
||||
step_b_v_fwd,
|
||||
)
|
||||
from .qkv_lora_b import qkv_lora_b_fwd
|
||||
from .sgemm_lora_a import sgemm_lora_a_fwd
|
||||
from .sgemm_lora_b import sgemm_lora_b_fwd
|
||||
from .virtual_experts import merged_experts_fused_moe_lora_add
|
||||
|
||||
__all__ = [
|
||||
"gate_up_lora_b_fwd",
|
||||
"qkv_lora_b_fwd",
|
||||
"sgemm_lora_a_fwd",
|
||||
"sgemm_lora_b_fwd",
|
||||
"merged_experts_fused_moe_lora_add",
|
||||
"step_a_q_fwd",
|
||||
"step_a_v_fwd",
|
||||
"step_b_q_fwd",
|
||||
"step_b_v_fwd",
|
||||
]
|
||||
@@ -0,0 +1,262 @@
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
|
||||
_resolve_token_positions,
|
||||
get_pdl_launch_metadata,
|
||||
)
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
# Minimum total_tokens * rank for the single-adapter cuBLAS path; below this
|
||||
# the Triton kernel is faster (crossover measured at output_dim=1536/GPU:
|
||||
# cuBLAS wins rank64 from S>=256 and rank16 only from S>=2048).
|
||||
_CUBLAS_MIN_S_RANK = 16384
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gate_up_lora_b_kernel(
|
||||
# Pointers to matrices
|
||||
x,
|
||||
weights,
|
||||
output,
|
||||
# Parameters of size
|
||||
K, # K = R
|
||||
output_dim,
|
||||
# Strides
|
||||
x_stride_0,
|
||||
x_stride_1,
|
||||
w_stride_0,
|
||||
w_stride_1,
|
||||
w_stride_2,
|
||||
output_stride_0,
|
||||
output_stride_1,
|
||||
# Information on sequence lengths,ranks and weight id
|
||||
seg_lens,
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
sorted_token_ids,
|
||||
# Meta parameters
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
# For fused output scaling
|
||||
scalings,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
):
|
||||
"""
|
||||
This kernel packs 2 sgemms (gate/up) into a single kernel. The multiplication
|
||||
results are accumulated into the output tensor.
|
||||
|
||||
When a sequence's rank is 0, the kernel is essentially a no-op, following
|
||||
the convention in pytorch where the product of two matrices of shape (m, 0)
|
||||
and (0, n) is an all-zero matrix of shape (m, n).
|
||||
|
||||
Args:
|
||||
x (Tensor): The input tensor, which is the result of the LoRA A projection.
|
||||
Shape: (s, 2 * K), where s is the sum of all sequence lengths in the
|
||||
batch and K is the maximum LoRA rank.
|
||||
weights (Tensor): The LoRA B weights for all adapters.
|
||||
Shape: (num_lora, 2 * output_dim, K).
|
||||
output (Tensor): The output tensor where the result is stored.
|
||||
Shape: (s, 2 * output_dim).
|
||||
"""
|
||||
# output_dim >> K
|
||||
|
||||
# Current block computes sequence with batch_id,
|
||||
# which starts from row seg_start of x with length seg_len.
|
||||
# gate_up_id decides which of gate or up (0: gate, 1: up)
|
||||
batch_id = tl.program_id(axis=2)
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
rank = tl.load(lora_ranks + w_index)
|
||||
|
||||
# If rank is 0, this kernel is a no-op.
|
||||
if rank == 0:
|
||||
return
|
||||
|
||||
gate_up_id = tl.program_id(axis=1)
|
||||
pid = tl.program_id(axis=0)
|
||||
seg_len = tl.load(seg_lens + batch_id)
|
||||
if seg_len == 0:
|
||||
return
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
n_start = gate_up_id * output_dim # offset on output dim
|
||||
scaling = tl.load(scalings + w_index)
|
||||
|
||||
# Adjust K (rank) according to the specific LoRA adapter
|
||||
K = tl.minimum(K, rank)
|
||||
|
||||
# The tile in output matrix will have (pid_s, pid_n) as id
|
||||
num_pid_n = tl.cdiv(output_dim, BLOCK_N)
|
||||
pid_s = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
if pid_s * BLOCK_S >= seg_len:
|
||||
return
|
||||
|
||||
# Create pointers for the first block of x and weights
|
||||
# The pointers will be advanced as we move in the K direction
|
||||
# and accumulate
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = tl.arange(0, BLOCK_K)
|
||||
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
x_ptrs = (
|
||||
x
|
||||
+ (gate_up_id * K) * x_stride_1
|
||||
+ (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
|
||||
)
|
||||
w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
|
||||
k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
|
||||
)
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
# Iterate to compute the block in output matrix
|
||||
partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
|
||||
for k in range(0, tl.cdiv(K, BLOCK_K)):
|
||||
x_tile = tl.load(
|
||||
x_ptrs,
|
||||
mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K),
|
||||
other=0.0,
|
||||
)
|
||||
w_tile = tl.load(
|
||||
w_ptrs,
|
||||
mask=(k_offset[:, None] < K - k * BLOCK_K)
|
||||
& (n_offset[None, :] < output_dim),
|
||||
other=0.0,
|
||||
)
|
||||
partial_sum += tl.dot(
|
||||
x_tile.to(w_tile.dtype), w_tile
|
||||
) # cast fused: split-K returns fp32, plain path bf16 (no-op)
|
||||
|
||||
x_ptrs += BLOCK_K * x_stride_1
|
||||
w_ptrs += BLOCK_K * w_stride_2
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
# Store result to output matrix
|
||||
partial_sum *= scaling
|
||||
partial_sum = partial_sum.to(x.dtype.element_ty)
|
||||
output_ptr = (
|
||||
output
|
||||
+ n_start * output_stride_1
|
||||
+ (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1)
|
||||
)
|
||||
output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < output_dim)
|
||||
partial_sum += tl.load(output_ptr, mask=output_mask)
|
||||
tl.store(output_ptr, partial_sum, mask=output_mask)
|
||||
|
||||
|
||||
def _gate_up_lora_b_cublas(
|
||||
x: torch.Tensor,
|
||||
gate_up_lora_b: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
output_dim: int,
|
||||
base_output: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Single-adapter dense path: one cuBLAS addmm_ per gate/up slice.
|
||||
|
||||
The LoRA-A output is rank-packed (slice i at columns [i*rank, (i+1)*rank)),
|
||||
matching the Triton kernel's K = min(K, rank) slice stride. Slices are
|
||||
disjoint output regions, so in-place addmm_ writes never collide.
|
||||
"""
|
||||
r = gate_up_lora_b.shape[-1]
|
||||
if base_output is None:
|
||||
base_output = torch.zeros(
|
||||
(x.shape[0], 2 * output_dim), device=x.device, dtype=x.dtype
|
||||
)
|
||||
w = gate_up_lora_b[0]
|
||||
x_scaled = x[:, : 2 * r] * batch_info.scalings[0]
|
||||
for i in range(2):
|
||||
lo, hi = i * output_dim, (i + 1) * output_dim
|
||||
base_output[:, lo:hi].addmm_(x_scaled[:, i * r : (i + 1) * r], w[lo:hi, :r].t())
|
||||
return base_output
|
||||
|
||||
|
||||
def gate_up_lora_b_fwd(
|
||||
x: torch.Tensor,
|
||||
gate_up_lora_b: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
output_dim: int,
|
||||
base_output: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
# x: (s, 2 * r)
|
||||
# gate_up_lora_b: (num_lora, 2 * output_dim, r)
|
||||
# output: (s, 2 * output_dim)
|
||||
|
||||
# Compute lora_output with shape (s, output_dim) as follows:
|
||||
# lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :])
|
||||
# lora_output[:, output_dim:]
|
||||
# = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :])
|
||||
|
||||
# Get dims
|
||||
s = x.shape[0]
|
||||
input_dim = x.shape[1]
|
||||
r = gate_up_lora_b.shape[-1]
|
||||
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:
|
||||
return _gate_up_lora_b_cublas(
|
||||
x, gate_up_lora_b, batch_info, output_dim, base_output
|
||||
)
|
||||
|
||||
BLOCK_S = 16
|
||||
BLOCK_R = 16
|
||||
BLOCK_OUT = 64
|
||||
|
||||
grid_b = (
|
||||
triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT),
|
||||
2, # this dimension decides current block computes on gate or up proj
|
||||
batch_info.bs,
|
||||
)
|
||||
|
||||
if base_output is None:
|
||||
output = torch.zeros((s, 2 * output_dim), device=x.device, dtype=x.dtype)
|
||||
else:
|
||||
output = base_output
|
||||
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
_gate_up_lora_b_kernel[grid_b](
|
||||
x,
|
||||
gate_up_lora_b,
|
||||
output,
|
||||
r,
|
||||
output_dim,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
gate_up_lora_b.stride(0),
|
||||
gate_up_lora_b.stride(1),
|
||||
gate_up_lora_b.stride(2),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
batch_info.seg_lens,
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
batch_info.permutation,
|
||||
sorted_by_adapter,
|
||||
BLOCK_S,
|
||||
BLOCK_OUT,
|
||||
BLOCK_R,
|
||||
batch_info.scalings,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,31 @@
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.jit_kernel.utils import is_arch_support_pdl
|
||||
|
||||
|
||||
def get_pdl_launch_metadata() -> tuple[bool, dict]:
|
||||
"""Return (ENABLE_PDL constexpr value, extra launch kwargs) for LoRA kernels.
|
||||
|
||||
``launch_pdl`` is NVIDIA-only Triton launch metadata; the HIP backend
|
||||
rejects unknown kwargs, so it is only included when PDL is supported.
|
||||
"""
|
||||
enable_pdl = is_arch_support_pdl()
|
||||
return enable_pdl, ({"launch_pdl": True} if enable_pdl else {})
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER: tl.constexpr
|
||||
):
|
||||
"""Map logical segment offsets to physical token positions.
|
||||
|
||||
When SORTED_BY_ADAPTER is True, segments are grouped by adapter and
|
||||
sorted_token_ids provides the indirection to the original token rows.
|
||||
When False, tokens are already contiguous starting at seg_start.
|
||||
"""
|
||||
if SORTED_BY_ADAPTER:
|
||||
return tl.load(
|
||||
sorted_token_ids + seg_start + s_offset, mask=s_offset < seg_len
|
||||
).to(tl.int64)
|
||||
return (seg_start + s_offset).to(tl.int64)
|
||||
@@ -0,0 +1,953 @@
|
||||
"""Triton kernels for absorbed-MLA ``kv_b_proj`` LoRA correction.
|
||||
|
||||
The absorbed-MLA path bypasses ``kv_b_proj.forward()`` and folds the K/V
|
||||
sides as plain BMMs ``q_nope @ w_kc`` and ``attn_output @ w_vc``. When a
|
||||
LoRA adapter is active on ``kv_b_proj`` we add the LoRA delta to
|
||||
``q_nope_out`` / ``attn_bmm_output`` manually.
|
||||
|
||||
Using the standard LoRA factored math we *never* materialize ``B @ A``:
|
||||
|
||||
q_correction = q_nope @ B_kc @ A * scaling # K-side
|
||||
v_correction = attn_output @ A.T @ B_vc.T * scaling # V-side
|
||||
|
||||
where ``A: (slot, rank, kv_lora_rank)`` is the LoRA-A of ``kv_b_proj``
|
||||
(shared across heads) and ``B: (slot, num_heads*(qk_nope+v_head_dim), rank)``
|
||||
is the LoRA-B; ``B_kc`` / ``B_vc`` are its K-half / V-half slices.
|
||||
|
||||
Four kernels split the math along the factorization boundary, all using
|
||||
the SGMM idiom from ``sgemm_lora_a`` / ``qkv_lora_b`` and the segment-indptr
|
||||
routing used by ``chunked_sgmv_*``:
|
||||
|
||||
* ``step_a_q_fwd``: per-head per-slot SGMM, ``(S,H,qk_nope) -> (S,H,rank)``
|
||||
* ``step_b_q_fwd``: shared-A per-slot SGMM, scaled+accumulated,
|
||||
``(S,H,rank) -> (S,H,kv_lora_rank)``
|
||||
* ``step_a_v_fwd``: shared-A.T per-slot SGMM, ``(S,H,kv_lora_rank) -> (S,H,rank)``
|
||||
* ``step_b_v_fwd``: per-head per-slot SGMM with V-half of B, transposed,
|
||||
scaled+accumulated, ``(S,H,rank) -> (S,H,v_head_dim)``
|
||||
|
||||
Grid axes for each kernel:
|
||||
axis 0 : output tile in (S, N) -- tile_id = pid_s * num_pid_n + pid_n
|
||||
axis 1 : head_id -- per-head weight slice
|
||||
axis 2 : batch_id (segment / request) -- per-slot weight routing via weight_indices
|
||||
|
||||
Per-segment routing: each program derives its segment length from
|
||||
``seg_indptr[segment_id + 1] - seg_indptr[segment_id]``, loads
|
||||
``weight_indices[segment_id]`` once, and uses that slot's slice of the LoRA
|
||||
weight stack. When ``permutation`` is present, rows are routed through it,
|
||||
matching the csgmv backend's adapter-grouped chunks. No Python loops over slots
|
||||
or heads.
|
||||
|
||||
The math also stays in the input dtype (no fp32 round-trip) -- the
|
||||
contraction dim ``rank`` is small (typically 16-64), so bf16 accumulation
|
||||
over it is acceptable. ``tl.dot`` itself uses fp32 accumulation internally.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
|
||||
_resolve_token_positions,
|
||||
get_pdl_launch_metadata,
|
||||
)
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Block sizes -- chosen per-kernel from the natural shape of each step.
|
||||
#
|
||||
# The factored math gives the four kernels these contraction (K) and output
|
||||
# (N) ranges (for Kimi-K2.5: rank=16-32, qk_nope=v_head_dim=128, kv_lora_rank=512):
|
||||
#
|
||||
# K (contraction) N (output)
|
||||
# step_a_q qk_nope (~128) rank (~16-32)
|
||||
# step_b_q rank (~16-32) kv_lora_rank (~512)
|
||||
# step_a_v kv_lora_rank (~512) rank (~16-32)
|
||||
# step_b_v rank (~16-32) v_head_dim (~128)
|
||||
#
|
||||
# So the "step_a_*" kernels want a large BLOCK_K (to keep loop iters small)
|
||||
# and a small BLOCK_N (matched to rank to avoid wasted tile lanes), while
|
||||
# the "step_b_*" kernels are the inverse. Kernels aren't autotuned -- the
|
||||
# decode-shape workload is too small to benefit and the sweep surface is
|
||||
# wide.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BLOCK_S = 16
|
||||
|
||||
|
||||
def _num_segments(batch_info: LoRABatchInfo) -> int:
|
||||
return batch_info.num_segments or batch_info.bs
|
||||
|
||||
|
||||
def _max_segment_len(batch_info: LoRABatchInfo) -> int:
|
||||
if batch_info.max_len is not None:
|
||||
return batch_info.max_len
|
||||
if batch_info.seg_lens is not None:
|
||||
return int(batch_info.seg_lens.max().item())
|
||||
raise ValueError("LoRA batch_info must provide max_len or seg_lens.")
|
||||
|
||||
|
||||
def _segment_grid_size(batch_info: LoRABatchInfo, num_segments: int) -> int:
|
||||
return batch_info.bs if batch_info.use_cuda_graph else num_segments
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kernel 1 -- Step A_q: per-head per-slot SGMM, reads K-half of B
|
||||
#
|
||||
# q_lora_a[t, h, r] = sum_{i<qk_nope} q_nope[t, h, i] * B[slot, h*FULL_K + i, r]
|
||||
#
|
||||
# x : (S, H, qk_nope)
|
||||
# w (B) : (num_lora, H*FULL_K, rank) -- FULL_K = qk_nope + v_head_dim
|
||||
# out : (S, H, rank) -- fresh allocation, no accumulate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit(do_not_specialize=["num_segments"])
|
||||
def _step_a_q_kernel(
|
||||
x,
|
||||
w,
|
||||
out,
|
||||
# dims
|
||||
S,
|
||||
H_FULL_K, # H * (qk_nope + v_head_dim), the row-stride landmark
|
||||
K, # qk_nope (contraction)
|
||||
N, # rank (output)
|
||||
# strides
|
||||
x_stride_s,
|
||||
x_stride_h,
|
||||
x_stride_k,
|
||||
w_stride_l,
|
||||
w_stride_n,
|
||||
w_stride_k,
|
||||
out_stride_s,
|
||||
out_stride_h,
|
||||
out_stride_n,
|
||||
# batch info
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
sorted_token_ids,
|
||||
num_segments,
|
||||
# meta
|
||||
FULL_K: tl.constexpr, # per-head row stride in B (qk_nope + v_head_dim)
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
K_DIV: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
):
|
||||
batch_id = tl.program_id(axis=2)
|
||||
head_id = tl.program_id(axis=1)
|
||||
pid = tl.program_id(axis=0)
|
||||
|
||||
if batch_id >= num_segments:
|
||||
return
|
||||
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
cur_rank = tl.load(lora_ranks + w_index)
|
||||
if cur_rank == 0:
|
||||
return
|
||||
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
seg_end = tl.load(seg_indptr + batch_id + 1)
|
||||
seg_len = seg_end - seg_start
|
||||
if seg_len == 0:
|
||||
return
|
||||
|
||||
# Truncate output N to this slot's rank (allows mixed-rank batches).
|
||||
N_eff = tl.minimum(N, cur_rank)
|
||||
|
||||
num_pid_n = tl.cdiv(N_eff, BLOCK_N)
|
||||
pid_s = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
if pid_s * BLOCK_S >= seg_len:
|
||||
return
|
||||
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = tl.arange(0, BLOCK_K)
|
||||
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
|
||||
# Clamp masked-lane indices into the valid range so pointer arithmetic
|
||||
# stays in-bounds even before the load mask drops the values.
|
||||
row_mask = s_offset < seg_len
|
||||
safe_row = tl.minimum(s_physical, S - 1)
|
||||
safe_n = tl.minimum(n_offset, N_eff - 1)
|
||||
|
||||
head_row_base = (
|
||||
head_id * FULL_K
|
||||
) # row offset for this head's K-half (i in [0, qk_nope))
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
|
||||
for k_block in range(0, tl.cdiv(K, BLOCK_K)):
|
||||
cur_k = k_block * BLOCK_K + k_offset
|
||||
k_mask = cur_k < K
|
||||
safe_k = cur_k if K_DIV else tl.minimum(cur_k, K - 1)
|
||||
|
||||
# x[s, h, k]
|
||||
x_tile = tl.load(
|
||||
x
|
||||
+ safe_row[:, None] * x_stride_s
|
||||
+ head_id * x_stride_h
|
||||
+ safe_k[None, :] * x_stride_k,
|
||||
mask=row_mask[:, None] & k_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# B[slot, h*FULL_K + i, r]: row dim of B carries i (= GEMM K),
|
||||
# column dim carries r (= GEMM N).
|
||||
w_tile = tl.load(
|
||||
w
|
||||
+ w_index * w_stride_l
|
||||
+ (head_row_base + safe_k[:, None]) * w_stride_n
|
||||
+ safe_n[None, :] * w_stride_k,
|
||||
mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
partial_sum += tl.dot(x_tile, w_tile)
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
partial_sum = partial_sum.to(x.dtype.element_ty)
|
||||
out_offs = (
|
||||
safe_row[:, None] * out_stride_s
|
||||
+ head_id * out_stride_h
|
||||
+ safe_n[None, :] * out_stride_n
|
||||
)
|
||||
out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
|
||||
tl.store(out + out_offs, partial_sum, mask=out_mask)
|
||||
|
||||
|
||||
def step_a_q_fwd(
|
||||
q_nope: torch.Tensor,
|
||||
B_buf: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
full_K_per_head: int,
|
||||
) -> torch.Tensor:
|
||||
"""Step A of the q-side correction.
|
||||
|
||||
Args:
|
||||
q_nope: ``(S, H, qk_nope)``, the absorbed-MLA q intermediate.
|
||||
B_buf: ``(num_lora, H*full_K_per_head, rank)`` from the LoRA pool.
|
||||
batch_info: standard ``LoRABatchInfo``.
|
||||
full_K_per_head: ``qk_nope + v_head_dim``, the row stride per head in B.
|
||||
|
||||
Returns:
|
||||
``(S, H, rank)`` -- per-token, per-head low-rank intermediate, ready for step B_q.
|
||||
"""
|
||||
S, H, qk_nope_dim = q_nope.shape
|
||||
rank = B_buf.shape[-1]
|
||||
|
||||
if (
|
||||
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
|
||||
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
|
||||
):
|
||||
# (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, :]
|
||||
return torch.bmm(q_nope.transpose(0, 1), w_kc).transpose(0, 1)
|
||||
|
||||
block_n = triton.next_power_of_2(rank) # output N == rank -> one tile
|
||||
out = torch.empty((S, H, rank), device=q_nope.device, dtype=q_nope.dtype)
|
||||
num_segments = _num_segments(batch_info)
|
||||
max_segment_len = _max_segment_len(batch_info)
|
||||
segment_grid = _segment_grid_size(batch_info, num_segments)
|
||||
_STEP_A_Q_BLOCK_K = 128
|
||||
|
||||
grid = (
|
||||
triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, block_n),
|
||||
H,
|
||||
segment_grid,
|
||||
)
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
|
||||
_step_a_q_kernel[grid](
|
||||
q_nope,
|
||||
B_buf,
|
||||
out,
|
||||
S,
|
||||
H * full_K_per_head,
|
||||
qk_nope_dim,
|
||||
rank,
|
||||
q_nope.stride(0),
|
||||
q_nope.stride(1),
|
||||
q_nope.stride(2),
|
||||
B_buf.stride(0),
|
||||
B_buf.stride(1),
|
||||
B_buf.stride(2),
|
||||
out.stride(0),
|
||||
out.stride(1),
|
||||
out.stride(2),
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
batch_info.permutation,
|
||||
num_segments,
|
||||
FULL_K=full_K_per_head,
|
||||
SORTED_BY_ADAPTER=sorted_by_adapter,
|
||||
K_DIV=(qk_nope_dim % _STEP_A_Q_BLOCK_K == 0),
|
||||
BLOCK_S=_BLOCK_S,
|
||||
BLOCK_N=block_n,
|
||||
BLOCK_K=_STEP_A_Q_BLOCK_K,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kernel 2 -- Step B_q: shared-A per-slot SGMM, scaled + accumulated
|
||||
#
|
||||
# base[t, h, k] += sum_r x[t, h, r] * A[slot, r, k] * scaling
|
||||
#
|
||||
# x : (S, H, rank)
|
||||
# w (A) : (num_lora, rank, kv_lora_rank)
|
||||
# base : (S, H, kv_lora_rank), updated in-place (accumulated)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit(do_not_specialize=["num_segments"])
|
||||
def _step_b_q_kernel(
|
||||
x,
|
||||
w,
|
||||
base,
|
||||
# dims
|
||||
S,
|
||||
K, # rank (contraction)
|
||||
N, # kv_lora_rank (output)
|
||||
# strides
|
||||
x_stride_s,
|
||||
x_stride_h,
|
||||
x_stride_k,
|
||||
w_stride_l,
|
||||
w_stride_k,
|
||||
w_stride_n,
|
||||
b_stride_s,
|
||||
b_stride_h,
|
||||
b_stride_n,
|
||||
# batch info
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
sorted_token_ids,
|
||||
scalings,
|
||||
num_segments,
|
||||
# meta
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
N_DIV: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
):
|
||||
batch_id = tl.program_id(axis=2)
|
||||
head_id = tl.program_id(axis=1)
|
||||
pid = tl.program_id(axis=0)
|
||||
|
||||
if batch_id >= num_segments:
|
||||
return
|
||||
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
cur_rank = tl.load(lora_ranks + w_index)
|
||||
if cur_rank == 0:
|
||||
return
|
||||
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
seg_end = tl.load(seg_indptr + batch_id + 1)
|
||||
seg_len = seg_end - seg_start
|
||||
if seg_len == 0:
|
||||
return
|
||||
scaling = tl.load(scalings + w_index)
|
||||
|
||||
# Truncate contraction K to this slot's rank.
|
||||
K_eff = tl.minimum(K, cur_rank)
|
||||
|
||||
num_pid_n = tl.cdiv(N, BLOCK_N)
|
||||
pid_s = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
if pid_s * BLOCK_S >= seg_len:
|
||||
return
|
||||
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = tl.arange(0, BLOCK_K)
|
||||
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
|
||||
row_mask = s_offset < seg_len
|
||||
safe_row = tl.minimum(s_physical, S - 1)
|
||||
n_mask = n_offset[None, :] < N
|
||||
safe_n = n_offset if N_DIV else tl.minimum(n_offset, N - 1)
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
|
||||
for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
|
||||
cur_k = k_block * BLOCK_K + k_offset
|
||||
k_mask = cur_k < K_eff
|
||||
safe_k = tl.minimum(cur_k, K_eff - 1)
|
||||
|
||||
# x[s, h, k] (k iterates over rank)
|
||||
x_tile = tl.load(
|
||||
x
|
||||
+ safe_row[:, None] * x_stride_s
|
||||
+ head_id * x_stride_h
|
||||
+ safe_k[None, :] * x_stride_k,
|
||||
mask=row_mask[:, None] & k_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# A[slot, k, n]: read k along contraction, n along output.
|
||||
w_tile = tl.load(
|
||||
w
|
||||
+ w_index * w_stride_l
|
||||
+ safe_k[:, None] * w_stride_k
|
||||
+ safe_n[None, :] * w_stride_n,
|
||||
mask=k_mask[:, None] & n_mask,
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
partial_sum += tl.dot(x_tile, w_tile)
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
partial_sum *= scaling
|
||||
partial_sum = partial_sum.to(x.dtype.element_ty)
|
||||
|
||||
# Accumulate into base[s, h, n].
|
||||
base_offs = (
|
||||
safe_row[:, None] * b_stride_s
|
||||
+ head_id * b_stride_h
|
||||
+ safe_n[None, :] * b_stride_n
|
||||
)
|
||||
out_mask = row_mask[:, None] & n_mask
|
||||
tl.atomic_add(base + base_offs, partial_sum, mask=out_mask, sem="relaxed")
|
||||
|
||||
|
||||
def step_b_q_fwd(
|
||||
q_lora_a: torch.Tensor,
|
||||
A_buf: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
base_output: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Step B of the q-side correction, accumulating into ``base_output``.
|
||||
|
||||
Args:
|
||||
q_lora_a: ``(S, H, rank)`` from step A_q.
|
||||
A_buf: ``(num_lora, rank, kv_lora_rank)`` from the LoRA pool.
|
||||
batch_info: standard ``LoRABatchInfo``.
|
||||
base_output: ``(S, H, kv_lora_rank)``, modified in-place
|
||||
(the absorbed ``q_nope @ w_kc`` result).
|
||||
|
||||
Returns:
|
||||
``base_output`` (same object, mutated).
|
||||
"""
|
||||
S, H, rank = q_lora_a.shape
|
||||
kv_lora_rank = A_buf.shape[-1]
|
||||
|
||||
if (
|
||||
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
|
||||
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
|
||||
):
|
||||
# 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.
|
||||
base2d = x2d = None
|
||||
if base_output.is_contiguous():
|
||||
base2d = base_output.view(-1, kv_lora_rank)
|
||||
x2d = q_lora_a[..., :rank].reshape(-1, rank)
|
||||
elif base_output.transpose(0, 1).is_contiguous():
|
||||
base2d = base_output.transpose(0, 1).view(-1, kv_lora_rank)
|
||||
x2d = q_lora_a[..., :rank].transpose(0, 1).reshape(-1, rank)
|
||||
if base2d is not None:
|
||||
base2d.addmm_(x2d, A_buf[0, :, :], alpha=batch_info.scalings[0])
|
||||
return base_output
|
||||
|
||||
num_segments = _num_segments(batch_info)
|
||||
max_segment_len = _max_segment_len(batch_info)
|
||||
segment_grid = _segment_grid_size(batch_info, num_segments)
|
||||
_STEP_B_Q_BLOCK_N = 128
|
||||
_STEP_B_BLOCK_K = 16
|
||||
|
||||
grid = (
|
||||
triton.cdiv(max_segment_len, _BLOCK_S)
|
||||
* triton.cdiv(kv_lora_rank, _STEP_B_Q_BLOCK_N),
|
||||
H,
|
||||
segment_grid,
|
||||
)
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
|
||||
_step_b_q_kernel[grid](
|
||||
q_lora_a,
|
||||
A_buf,
|
||||
base_output,
|
||||
S,
|
||||
rank,
|
||||
kv_lora_rank,
|
||||
q_lora_a.stride(0),
|
||||
q_lora_a.stride(1),
|
||||
q_lora_a.stride(2),
|
||||
A_buf.stride(0),
|
||||
A_buf.stride(1),
|
||||
A_buf.stride(2),
|
||||
base_output.stride(0),
|
||||
base_output.stride(1),
|
||||
base_output.stride(2),
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
batch_info.permutation,
|
||||
batch_info.scalings,
|
||||
num_segments,
|
||||
SORTED_BY_ADAPTER=sorted_by_adapter,
|
||||
N_DIV=(kv_lora_rank % _STEP_B_Q_BLOCK_N == 0),
|
||||
BLOCK_S=_BLOCK_S,
|
||||
BLOCK_N=_STEP_B_Q_BLOCK_N,
|
||||
BLOCK_K=_STEP_B_BLOCK_K,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
return base_output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kernel 3 -- Step A_v: shared-A.T per-slot SGMM (no scaling, fresh output)
|
||||
#
|
||||
# attn_lora_a[t, h, r] = sum_k attn_output[t, h, k] * A[slot, r, k]
|
||||
#
|
||||
# x : (S, H, kv_lora_rank)
|
||||
# w (A) : (num_lora, rank, kv_lora_rank) -- accessed transposed vs step B_q
|
||||
# out : (S, H, rank), fresh allocation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit(do_not_specialize=["num_segments"])
|
||||
def _step_a_v_kernel(
|
||||
x,
|
||||
w,
|
||||
out,
|
||||
# dims
|
||||
S,
|
||||
K, # kv_lora_rank (contraction)
|
||||
N, # rank (output)
|
||||
# strides
|
||||
x_stride_s,
|
||||
x_stride_h,
|
||||
x_stride_k,
|
||||
w_stride_l,
|
||||
w_stride_n, # A's "rank" axis (= GEMM N)
|
||||
w_stride_k, # A's "kv_lora_rank" axis (= GEMM K)
|
||||
out_stride_s,
|
||||
out_stride_h,
|
||||
out_stride_n,
|
||||
# batch info
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
sorted_token_ids,
|
||||
num_segments,
|
||||
# meta
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
K_DIV: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
):
|
||||
batch_id = tl.program_id(axis=2)
|
||||
head_id = tl.program_id(axis=1)
|
||||
pid = tl.program_id(axis=0)
|
||||
|
||||
if batch_id >= num_segments:
|
||||
return
|
||||
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
cur_rank = tl.load(lora_ranks + w_index)
|
||||
if cur_rank == 0:
|
||||
return
|
||||
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
seg_end = tl.load(seg_indptr + batch_id + 1)
|
||||
seg_len = seg_end - seg_start
|
||||
if seg_len == 0:
|
||||
return
|
||||
|
||||
# Truncate output N to this slot's rank.
|
||||
N_eff = tl.minimum(N, cur_rank)
|
||||
|
||||
num_pid_n = tl.cdiv(N_eff, BLOCK_N)
|
||||
pid_s = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
if pid_s * BLOCK_S >= seg_len:
|
||||
return
|
||||
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = tl.arange(0, BLOCK_K)
|
||||
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
|
||||
row_mask = s_offset < seg_len
|
||||
safe_row = tl.minimum(s_physical, S - 1)
|
||||
safe_n = tl.minimum(n_offset, N_eff - 1)
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
|
||||
for k_block in range(0, tl.cdiv(K, BLOCK_K)):
|
||||
cur_k = k_block * BLOCK_K + k_offset
|
||||
k_mask = cur_k < K
|
||||
safe_k = cur_k if K_DIV else tl.minimum(cur_k, K - 1)
|
||||
|
||||
# x[s, h, k]
|
||||
x_tile = tl.load(
|
||||
x
|
||||
+ safe_row[:, None] * x_stride_s
|
||||
+ head_id * x_stride_h
|
||||
+ safe_k[None, :] * x_stride_k,
|
||||
mask=row_mask[:, None] & k_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# A[slot, r, k] -- here we want each (k, r) so we read along k
|
||||
# (inner / contraction) and produce r as output. Stride access:
|
||||
# the row dim is r (= GEMM N), column dim is k (= GEMM K).
|
||||
w_tile = tl.load(
|
||||
w
|
||||
+ w_index * w_stride_l
|
||||
+ safe_k[:, None] * w_stride_k
|
||||
+ safe_n[None, :] * w_stride_n,
|
||||
mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
partial_sum += tl.dot(x_tile, w_tile)
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
partial_sum = partial_sum.to(x.dtype.element_ty)
|
||||
out_offs = (
|
||||
safe_row[:, None] * out_stride_s
|
||||
+ head_id * out_stride_h
|
||||
+ safe_n[None, :] * out_stride_n
|
||||
)
|
||||
out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
|
||||
tl.store(out + out_offs, partial_sum, mask=out_mask)
|
||||
|
||||
|
||||
def step_a_v_fwd(
|
||||
attn_output: torch.Tensor,
|
||||
A_buf: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
) -> torch.Tensor:
|
||||
"""Step A of the v-side correction.
|
||||
|
||||
Args:
|
||||
attn_output: ``(S, H, kv_lora_rank)``, the post-attention intermediate.
|
||||
A_buf: ``(num_lora, rank, kv_lora_rank)``.
|
||||
batch_info: standard ``LoRABatchInfo``.
|
||||
|
||||
Returns:
|
||||
``(S, H, rank)`` -- per-token, per-head low-rank intermediate for step B_v.
|
||||
"""
|
||||
S, H, kv_lora_rank = attn_output.shape
|
||||
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():
|
||||
return torch.mm(
|
||||
attn_output.view(-1, kv_lora_rank), A_buf[0, :rank, :].t()
|
||||
).view(S, H, rank)
|
||||
|
||||
block_n = triton.next_power_of_2(rank)
|
||||
out = torch.empty((S, H, rank), device=attn_output.device, dtype=attn_output.dtype)
|
||||
num_segments = _num_segments(batch_info)
|
||||
max_segment_len = _max_segment_len(batch_info)
|
||||
segment_grid = _segment_grid_size(batch_info, num_segments)
|
||||
_STEP_A_V_BLOCK_K = 256
|
||||
|
||||
grid = (
|
||||
triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, block_n),
|
||||
H,
|
||||
segment_grid,
|
||||
)
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
|
||||
_step_a_v_kernel[grid](
|
||||
attn_output,
|
||||
A_buf,
|
||||
out,
|
||||
S,
|
||||
kv_lora_rank,
|
||||
rank,
|
||||
attn_output.stride(0),
|
||||
attn_output.stride(1),
|
||||
attn_output.stride(2),
|
||||
A_buf.stride(0),
|
||||
A_buf.stride(1),
|
||||
A_buf.stride(2),
|
||||
out.stride(0),
|
||||
out.stride(1),
|
||||
out.stride(2),
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
batch_info.permutation,
|
||||
num_segments,
|
||||
SORTED_BY_ADAPTER=sorted_by_adapter,
|
||||
K_DIV=(kv_lora_rank % _STEP_A_V_BLOCK_K == 0),
|
||||
BLOCK_S=_BLOCK_S,
|
||||
BLOCK_N=block_n,
|
||||
BLOCK_K=_STEP_A_V_BLOCK_K,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kernel 4 -- Step B_v: per-head per-slot SGMM with V-half of B (transposed),
|
||||
# scaled + accumulated
|
||||
#
|
||||
# base[t, h, j] += sum_r x[t, h, r] * B[slot, h*FULL_K + qk_nope + j, r] * scaling
|
||||
#
|
||||
# x : (S, H, rank)
|
||||
# w (B) : (num_lora, H*FULL_K, rank), V-half slice via offset
|
||||
# base : (S, H, v_head_dim), updated in-place (accumulated)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit(do_not_specialize=["num_segments"])
|
||||
def _step_b_v_kernel(
|
||||
x,
|
||||
w,
|
||||
base,
|
||||
# dims
|
||||
S,
|
||||
K, # rank (contraction)
|
||||
N, # v_head_dim (output)
|
||||
# strides
|
||||
x_stride_s,
|
||||
x_stride_h,
|
||||
x_stride_k,
|
||||
w_stride_l,
|
||||
w_stride_n, # B's row dim (h*FULL_K + j) -- this is GEMM N
|
||||
w_stride_k, # B's rank dim -- this is GEMM K
|
||||
b_stride_s,
|
||||
b_stride_h,
|
||||
b_stride_n,
|
||||
# batch info
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
sorted_token_ids,
|
||||
scalings,
|
||||
num_segments,
|
||||
# meta
|
||||
FULL_K: tl.constexpr, # qk_nope + v_head_dim
|
||||
QK_NOPE_OFFSET: tl.constexpr, # offset of V-half within each head's row block
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
N_DIV: tl.constexpr, # N % BLOCK_N == 0 -> drop safe_n (keep the store coalesced)
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
):
|
||||
batch_id = tl.program_id(axis=2)
|
||||
head_id = tl.program_id(axis=1)
|
||||
pid = tl.program_id(axis=0)
|
||||
|
||||
if batch_id >= num_segments:
|
||||
return
|
||||
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
cur_rank = tl.load(lora_ranks + w_index)
|
||||
if cur_rank == 0:
|
||||
return
|
||||
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
seg_end = tl.load(seg_indptr + batch_id + 1)
|
||||
seg_len = seg_end - seg_start
|
||||
if seg_len == 0:
|
||||
return
|
||||
scaling = tl.load(scalings + w_index)
|
||||
|
||||
K_eff = tl.minimum(K, cur_rank)
|
||||
|
||||
num_pid_n = tl.cdiv(N, BLOCK_N)
|
||||
pid_s = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
if pid_s * BLOCK_S >= seg_len:
|
||||
return
|
||||
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = tl.arange(0, BLOCK_K)
|
||||
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
|
||||
row_mask = s_offset < seg_len
|
||||
safe_row = tl.minimum(s_physical, S - 1)
|
||||
n_mask = n_offset[None, :] < N
|
||||
safe_n = n_offset if N_DIV else tl.minimum(n_offset, N - 1)
|
||||
|
||||
# V-half row base for this head: h*FULL_K + qk_nope
|
||||
head_row_base = head_id * FULL_K + QK_NOPE_OFFSET
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
|
||||
for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
|
||||
cur_k = k_block * BLOCK_K + k_offset
|
||||
k_mask = cur_k < K_eff
|
||||
safe_k = tl.minimum(cur_k, K_eff - 1)
|
||||
|
||||
# x[s, h, k]
|
||||
x_tile = tl.load(
|
||||
x
|
||||
+ safe_row[:, None] * x_stride_s
|
||||
+ head_id * x_stride_h
|
||||
+ safe_k[None, :] * x_stride_k,
|
||||
mask=row_mask[:, None] & k_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# B[slot, h*FULL_K + qk_nope + j, r] -- row dim is j (= GEMM N),
|
||||
# column dim is r (= GEMM K). Transposed access vs step A_q.
|
||||
w_tile = tl.load(
|
||||
w
|
||||
+ w_index * w_stride_l
|
||||
+ safe_k[:, None] * w_stride_k
|
||||
+ (head_row_base + safe_n[None, :]) * w_stride_n,
|
||||
mask=k_mask[:, None] & n_mask,
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
partial_sum += tl.dot(x_tile, w_tile)
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
partial_sum *= scaling
|
||||
partial_sum = partial_sum.to(x.dtype.element_ty)
|
||||
|
||||
base_offs = (
|
||||
safe_row[:, None] * b_stride_s
|
||||
+ head_id * b_stride_h
|
||||
+ safe_n[None, :] * b_stride_n
|
||||
)
|
||||
out_mask = row_mask[:, None] & n_mask
|
||||
tl.atomic_add(base + base_offs, partial_sum, mask=out_mask, sem="relaxed")
|
||||
|
||||
|
||||
def step_b_v_fwd(
|
||||
attn_lora_a: torch.Tensor,
|
||||
B_buf: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
base_output: torch.Tensor,
|
||||
qk_nope_head_dim: int,
|
||||
v_head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""Step B of the v-side correction, accumulating into ``base_output``.
|
||||
|
||||
Args:
|
||||
attn_lora_a: ``(S, H, rank)`` from step A_v.
|
||||
B_buf: ``(num_lora, H*(qk_nope+v_head_dim), rank)``.
|
||||
batch_info: standard ``LoRABatchInfo``.
|
||||
base_output: ``(S, H, v_head_dim)``, modified in-place
|
||||
(the absorbed ``attn_output @ w_vc`` result).
|
||||
qk_nope_head_dim: offset of V-half within each head's row block of B.
|
||||
v_head_dim: output feature dim per head.
|
||||
|
||||
Returns:
|
||||
``base_output`` (same object, mutated).
|
||||
"""
|
||||
S, H, rank = attn_lora_a.shape
|
||||
full_K_per_head = qk_nope_head_dim + v_head_dim
|
||||
num_segments = _num_segments(batch_info)
|
||||
max_segment_len = _max_segment_len(batch_info)
|
||||
segment_grid = _segment_grid_size(batch_info, num_segments)
|
||||
_STEP_B_V_BLOCK_N = 64
|
||||
_STEP_B_BLOCK_K = 16
|
||||
|
||||
grid = (
|
||||
triton.cdiv(max_segment_len, _BLOCK_S)
|
||||
* triton.cdiv(v_head_dim, _STEP_B_V_BLOCK_N),
|
||||
H,
|
||||
segment_grid,
|
||||
)
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
|
||||
_step_b_v_kernel[grid](
|
||||
attn_lora_a,
|
||||
B_buf,
|
||||
base_output,
|
||||
S,
|
||||
rank,
|
||||
v_head_dim,
|
||||
attn_lora_a.stride(0),
|
||||
attn_lora_a.stride(1),
|
||||
attn_lora_a.stride(2),
|
||||
B_buf.stride(0),
|
||||
B_buf.stride(1),
|
||||
B_buf.stride(2),
|
||||
base_output.stride(0),
|
||||
base_output.stride(1),
|
||||
base_output.stride(2),
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
batch_info.permutation,
|
||||
batch_info.scalings,
|
||||
num_segments,
|
||||
FULL_K=full_K_per_head,
|
||||
QK_NOPE_OFFSET=qk_nope_head_dim,
|
||||
SORTED_BY_ADAPTER=sorted_by_adapter,
|
||||
N_DIV=(v_head_dim % _STEP_B_V_BLOCK_N == 0),
|
||||
BLOCK_S=_BLOCK_S,
|
||||
BLOCK_N=_STEP_B_V_BLOCK_N,
|
||||
BLOCK_K=_STEP_B_BLOCK_K,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
return base_output
|
||||
@@ -0,0 +1,295 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
|
||||
_resolve_token_positions,
|
||||
get_pdl_launch_metadata,
|
||||
)
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
# Minimum max_len (longest segment) for the single-adapter cuBLAS path; below
|
||||
# this the Triton kernel is faster (measured crossover at the smallest
|
||||
# realistic N=768, where cuBLAS is weakest).
|
||||
_CUBLAS_MIN_MAX_LEN = 8
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _qkv_lora_b_kernel(
|
||||
# Pointers to matrices
|
||||
x,
|
||||
weights,
|
||||
output,
|
||||
# Parameters of size
|
||||
K, # K = R
|
||||
max_qkv_out_dim, # max(output_q_dim, output_kv_dim)
|
||||
# Strides
|
||||
x_stride_0,
|
||||
x_stride_1,
|
||||
w_stride_0,
|
||||
w_stride_1,
|
||||
w_stride_2,
|
||||
output_stride_0,
|
||||
output_stride_1,
|
||||
# Information on sequence lengths and weight id
|
||||
seg_lens,
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
# Offsets of q/k/v slice on output dimension
|
||||
n_offs,
|
||||
sorted_token_ids,
|
||||
# Meta parameters
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
# For fused output scaling
|
||||
scalings,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
STORE_WRITEBACK: tl.constexpr = False,
|
||||
):
|
||||
"""
|
||||
This kernel packs 3 sgemms (q/k/v) into a single kernel. The multiplication
|
||||
results are accumulated into the output tensor.
|
||||
|
||||
When a sequence's rank is 0, the kernel is essentially a no-op, following
|
||||
the convention in pytorch where the product of two matrices of shape (m, 0)
|
||||
and (0, n) is an all-zero matrix of shape (m, n).
|
||||
|
||||
Args:
|
||||
x (Tensor): The input tensor, which is the result of the LoRA A projection.
|
||||
Shape: (s, 3 * K), where s is the sum of all sequence lengths in the
|
||||
batch and K is the maximum LoRA rank. The second dimension is partitioned
|
||||
for Q, K, and V.
|
||||
weights (Tensor): The LoRA B weights for all adapters.
|
||||
Shape: (num_lora, N_Q + 2 * N_KV, K).
|
||||
output (Tensor): The output tensor where the result is stored.
|
||||
Shape: (s, N_Q + 2 * N_KV).
|
||||
"""
|
||||
|
||||
# Current block computes sequence with batch_id,
|
||||
# which starts from row seg_start of x with length seg_len.
|
||||
# qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v)
|
||||
batch_id = tl.program_id(axis=2)
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
rank = tl.load(lora_ranks + w_index)
|
||||
|
||||
# If rank is 0, this kernel is a no-op.
|
||||
if rank == 0:
|
||||
return
|
||||
|
||||
qkv_id = tl.program_id(axis=1)
|
||||
pid = tl.program_id(axis=0)
|
||||
seg_len = tl.load(seg_lens + batch_id)
|
||||
if seg_len == 0:
|
||||
return
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
n_start = tl.load(n_offs + qkv_id)
|
||||
n_size = tl.load(n_offs + qkv_id + 1) - n_start
|
||||
scaling = tl.load(scalings + w_index)
|
||||
# Adjust K (rank) according to the specific LoRA adapter.
|
||||
K = tl.minimum(K, rank)
|
||||
|
||||
# The tile in output matrix will have (pid_s, pid_n) as id
|
||||
num_pid_n = tl.cdiv(max_qkv_out_dim, BLOCK_N)
|
||||
pid_s = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
if pid_s * BLOCK_S >= seg_len:
|
||||
return
|
||||
|
||||
# Create pointers for the first block of x and weights[batch_id][n_start: n_end][:]
|
||||
# The pointers will be advanced as we move in the K direction
|
||||
# and accumulate
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = tl.arange(0, BLOCK_K)
|
||||
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
x_ptrs = (
|
||||
x
|
||||
+ (qkv_id * K) * x_stride_1
|
||||
+ (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
|
||||
)
|
||||
w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
|
||||
k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
|
||||
)
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
n_mask = n_offset[None, :] < n_size
|
||||
x_tile = tl.load(
|
||||
x_ptrs,
|
||||
mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K),
|
||||
other=0.0,
|
||||
)
|
||||
w_tile = tl.load(
|
||||
w_ptrs,
|
||||
mask=(k_offset[:, None] < K) & n_mask,
|
||||
other=0.0,
|
||||
)
|
||||
# cast fused: the split-K shrink returns fp32, plain path bf16 (no-op)
|
||||
partial_sum = tl.dot(x_tile.to(w_tile.dtype), w_tile)
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
# Store result to output matrix (cast to the OUTPUT dtype: x may be the fp32
|
||||
# split-K shrink accumulator while base_output is bf16)
|
||||
partial_sum *= scaling
|
||||
partial_sum = partial_sum.to(output.dtype.element_ty)
|
||||
output_ptr = (
|
||||
output
|
||||
+ n_start * output_stride_1
|
||||
+ (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1)
|
||||
)
|
||||
output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < n_size)
|
||||
if STORE_WRITEBACK:
|
||||
# The expand-add output tiles are disjoint across all programs in this launch
|
||||
# (distinct s-rows / n-cols / slice / segment), so each element is RMW'd by
|
||||
# exactly one program -- a plain read-add-write is correct and avoids the bf16
|
||||
# narrow-tile atomic (the dominant decode cost). base_output is a same-stream
|
||||
# data dependency (base GEMM before apply_lora), not a concurrent writer.
|
||||
partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
|
||||
tl.store(output_ptr, partial_sum, mask=output_mask)
|
||||
else:
|
||||
tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
|
||||
|
||||
|
||||
def _qkv_lora_b_cublas(
|
||||
x: torch.Tensor,
|
||||
qkv_lora_b: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
output_offset_cpu: torch.Tensor,
|
||||
base_output: Optional[torch.Tensor],
|
||||
n_slices: int,
|
||||
) -> torch.Tensor:
|
||||
"""Single-adapter dense path: one cuBLAS addmm_ per q/k/v slice.
|
||||
|
||||
The LoRA-A output is rank-packed (slice i at columns [i*rank, (i+1)*rank)),
|
||||
matching the Triton kernel's K = min(K, rank) slice stride. Slice offsets
|
||||
come from the pinned CPU copy (no GPU sync); slices are disjoint output
|
||||
regions, so in-place addmm_ writes never collide.
|
||||
"""
|
||||
r = qkv_lora_b.shape[-1]
|
||||
if base_output is None:
|
||||
base_output = torch.zeros(
|
||||
(x.shape[0], qkv_lora_b.shape[-2]), device=x.device, dtype=x.dtype
|
||||
)
|
||||
w = qkv_lora_b[0]
|
||||
x_scaled = x[:, : n_slices * r] * batch_info.scalings[0]
|
||||
offsets = output_offset_cpu.tolist()
|
||||
for i in range(n_slices):
|
||||
lo, hi = offsets[i], offsets[i + 1]
|
||||
base_output[:, lo:hi].addmm_(x_scaled[:, i * r : (i + 1) * r], w[lo:hi, :r].t())
|
||||
return base_output
|
||||
|
||||
|
||||
def qkv_lora_b_fwd(
|
||||
x: torch.Tensor,
|
||||
qkv_lora_b: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
output_offset: torch.Tensor,
|
||||
max_qkv_out_dim: int,
|
||||
base_output: torch.Tensor = None,
|
||||
n_slices: int = 3,
|
||||
output_offset_cpu: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
# x: (s, n_slices * r)
|
||||
# qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
|
||||
# output_offset = [0, output_dim_q, output_dim_q + output_dim_kv,
|
||||
# output_dim_q + 2 * output_dim_kv] (length n_slices + 1)
|
||||
# max_qkv_out_dim = max(output_dim_q, output_dim_kv)
|
||||
# output: (s, output_dim_q + 2 * output_dim_kv)
|
||||
|
||||
# Compute lora_output with shape (s, output_dim) as follows:
|
||||
# lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :])
|
||||
# lora_output[:, output_dim_q: output_dim_q + output_dim_kv]
|
||||
# = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :])
|
||||
# lora_output[:, output_dim_q + output_dim_kv: ]
|
||||
# = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :])
|
||||
|
||||
# Get dims
|
||||
s = x.shape[0]
|
||||
input_dim = x.shape[1]
|
||||
r = qkv_lora_b.shape[-1]
|
||||
output_dim = qkv_lora_b.shape[-2]
|
||||
assert input_dim == n_slices * r
|
||||
assert output_offset.shape[0] == n_slices + 1
|
||||
|
||||
if (
|
||||
output_offset_cpu is not None
|
||||
and (
|
||||
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
|
||||
or lora_envs.SGLANG_OPT_LORA_CUBLAS_QKV.get()
|
||||
)
|
||||
and batch_info.max_len >= _CUBLAS_MIN_MAX_LEN
|
||||
):
|
||||
return _qkv_lora_b_cublas(
|
||||
x, qkv_lora_b, batch_info, output_offset_cpu, base_output, n_slices
|
||||
)
|
||||
|
||||
BLOCK_S = 16
|
||||
BLOCK_R = triton.next_power_of_2(r)
|
||||
# BLOCK_OUT stays 64: with the 1-adapter cuBLAS dispatch the Triton path
|
||||
# only runs for decode-sized batches, where 128 halves the grid (96->48
|
||||
# programs on Kimi r16 bs64) and slows the kernel ~60% (11.4->18.5us, B200).
|
||||
# Re-swept for the store path on GB200: 32 vs 64 is within noise (one preset
|
||||
# marginally each way), so the single value is kept for both writebacks.
|
||||
BLOCK_OUT = 64
|
||||
|
||||
grid_b = (
|
||||
triton.cdiv(batch_info.max_len, BLOCK_S)
|
||||
* triton.cdiv(max_qkv_out_dim, BLOCK_OUT),
|
||||
n_slices,
|
||||
batch_info.bs,
|
||||
)
|
||||
|
||||
if base_output is None:
|
||||
output = torch.zeros((s, output_dim), device=x.device, dtype=x.dtype)
|
||||
else:
|
||||
output = base_output
|
||||
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
store_writeback = lora_envs.SGLANG_OPT_LORA_QKV_B_STORE.get()
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
_qkv_lora_b_kernel[grid_b](
|
||||
x,
|
||||
qkv_lora_b,
|
||||
output,
|
||||
r,
|
||||
max_qkv_out_dim,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
qkv_lora_b.stride(0),
|
||||
qkv_lora_b.stride(1),
|
||||
qkv_lora_b.stride(2),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
batch_info.seg_lens,
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
output_offset,
|
||||
batch_info.permutation,
|
||||
sorted_by_adapter,
|
||||
BLOCK_S,
|
||||
BLOCK_OUT,
|
||||
BLOCK_R,
|
||||
batch_info.scalings,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
STORE_WRITEBACK=store_writeback,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,267 @@
|
||||
import functools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
|
||||
_resolve_token_positions,
|
||||
get_pdl_launch_metadata,
|
||||
)
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _sgemm_lora_a_kernel(
|
||||
# Pointers to matrices
|
||||
x,
|
||||
weights,
|
||||
output,
|
||||
# Matrix dimensions
|
||||
N, # stack_num * r
|
||||
K, # input_dim
|
||||
stack_num,
|
||||
# Strides
|
||||
x_stride_0,
|
||||
x_stride_1,
|
||||
w_stride_0,
|
||||
w_stride_1,
|
||||
w_stride_2,
|
||||
output_stride_0,
|
||||
output_stride_1,
|
||||
# Information on sequence lengths,ranks and weight id
|
||||
seg_lens,
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
sorted_token_ids,
|
||||
# Meta parameters
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
SPLIT_K: tl.constexpr = 1,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
):
|
||||
"""
|
||||
Computes a segmented batched matrix multiplication for the LoRA A matrix.
|
||||
|
||||
The kernel ensures that output[seg_start:seg_start + seg_len, :rank * stack_num]
|
||||
stores the product of the input `x` and the LoRA weights for the corresponding
|
||||
sequence. This implies that when rank is 0, the kernel is essentially a no-op,
|
||||
as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty).
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s`
|
||||
is the sum of all sequence lengths in the batch.
|
||||
weights (torch.Tensor): The LoRA 'A' weights for all available adapters,
|
||||
with shape `(num_lora, N, K)`.
|
||||
output (torch.Tensor): The output tensor of shape `(s, N)`.
|
||||
"""
|
||||
|
||||
# Current block computes sequence with batch_id,
|
||||
# which starts from row seg_start of x with length seg_len
|
||||
batch_id = tl.program_id(axis=1)
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
rank = tl.load(lora_ranks + w_index)
|
||||
|
||||
# If rank is 0, this kernel becomes a no-op as the output is always trivially correct.
|
||||
if rank == 0:
|
||||
return
|
||||
|
||||
pid = tl.program_id(axis=0)
|
||||
# Fold the split-K factor out of axis-0 (SPLIT_K == 1 -> pid_sk == 0, pid_tile == pid).
|
||||
pid_sk = pid % SPLIT_K
|
||||
pid_tile = pid // SPLIT_K
|
||||
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
seg_len = tl.load(seg_lens + batch_id)
|
||||
if seg_len == 0:
|
||||
return
|
||||
|
||||
# Adjust N (stack_num * max_rank) to this adapter's actual rank.
|
||||
N = tl.minimum(N, rank * stack_num)
|
||||
|
||||
# The tile in output matrix will have (pid_s, pid_n) as id
|
||||
num_pid_n = tl.cdiv(N, BLOCK_N)
|
||||
pid_s = pid_tile // num_pid_n
|
||||
pid_n = pid_tile % num_pid_n
|
||||
if pid_s * BLOCK_S >= seg_len:
|
||||
return
|
||||
|
||||
# Create pointers for the first block of x and weights[batch_id]
|
||||
# The pointers will be advanced as we move in the K direction
|
||||
# and accumulate.
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
x_ptrs = x + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
|
||||
w_ptrs = (weights + w_index * w_stride_0) + (
|
||||
k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
|
||||
)
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
# Iterate to compute the block in output matrix
|
||||
partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
|
||||
for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
|
||||
k_remaining = K - k * (BLOCK_K * SPLIT_K)
|
||||
x_tile = tl.load(
|
||||
x_ptrs,
|
||||
mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < k_remaining),
|
||||
other=0.0,
|
||||
)
|
||||
w_tile = tl.load(
|
||||
w_ptrs,
|
||||
mask=(k_offset[:, None] < k_remaining) & (n_offset[None, :] < N),
|
||||
other=0.0,
|
||||
)
|
||||
partial_sum += tl.dot(x_tile, w_tile)
|
||||
|
||||
x_ptrs += BLOCK_K * SPLIT_K * x_stride_1
|
||||
w_ptrs += BLOCK_K * SPLIT_K * w_stride_2
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
# Store result to output matrix
|
||||
output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < N)
|
||||
output_ptr = output + (
|
||||
s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
|
||||
)
|
||||
if SPLIT_K == 1:
|
||||
tl.store(output_ptr, partial_sum.to(output.dtype.element_ty), mask=output_mask)
|
||||
else:
|
||||
tl.atomic_add(
|
||||
output_ptr,
|
||||
partial_sum.to(output.dtype.element_ty),
|
||||
mask=output_mask,
|
||||
sem="relaxed",
|
||||
)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=None)
|
||||
def _num_sms(device_index: int) -> int:
|
||||
return torch.cuda.get_device_properties(device_index).multi_processor_count
|
||||
|
||||
|
||||
def sgemm_lora_a_fwd(
|
||||
x: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
stack_num: int = 1,
|
||||
out_alloc_stream=None,
|
||||
) -> torch.Tensor:
|
||||
# x: (s, input_dim)
|
||||
# weights: (num_lora, stack_num * r, input_dim)
|
||||
# output: (s, stack_num * r)
|
||||
# stack_num: run_qkv_lora: 3, run_gate_up_lora: 2
|
||||
# when called by run_qkv_lora, the weights.shape[-2] will be 3 * r
|
||||
# input_dim is much larger than r
|
||||
|
||||
assert x.is_contiguous()
|
||||
assert weights.is_contiguous()
|
||||
assert len(x.shape) == 2
|
||||
assert len(weights.shape) == 3
|
||||
|
||||
S = x.shape[0]
|
||||
R = weights.shape[-2]
|
||||
K = weights.shape[-1]
|
||||
assert x.shape[-1] == K
|
||||
|
||||
if (
|
||||
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
|
||||
or lora_envs.SGLANG_OPT_LORA_CUBLAS_A.get()
|
||||
):
|
||||
# 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
|
||||
# allocate explicitly and matmul into it.
|
||||
if out_alloc_stream is not None:
|
||||
with torch.cuda.stream(out_alloc_stream):
|
||||
output = torch.empty((S, R), device=x.device, dtype=x.dtype)
|
||||
else:
|
||||
output = torch.empty((S, R), device=x.device, dtype=x.dtype)
|
||||
return torch.matmul(x, weights[0].transpose(-2, -1), out=output)
|
||||
|
||||
# Block shapes
|
||||
BLOCK_S = 16
|
||||
BLOCK_K = 256
|
||||
BLOCK_R = triton.next_power_of_2(R)
|
||||
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
|
||||
num_s_tiles = triton.cdiv(batch_info.max_len, BLOCK_S)
|
||||
split_k = 1
|
||||
if lora_envs.SGLANG_ENABLE_LORA_SHRINK_SPLIT_K.get() and x.is_cuda:
|
||||
num_k_tiles = triton.cdiv(K, BLOCK_K)
|
||||
base_grid = batch_info.bs * num_s_tiles
|
||||
num_sms = _num_sms(x.device.index)
|
||||
if base_grid < num_sms and num_k_tiles >= 16:
|
||||
split_k = max(1, min(2 * num_sms // base_grid, num_k_tiles, 16))
|
||||
|
||||
launch_kwargs = {}
|
||||
if split_k > 1:
|
||||
# out_alloc_stream (SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC) is intentionally NOT honored here:
|
||||
# torch.zeros launches its memset on the alloc stream, which would race the side-stream
|
||||
# shrink without extra ordering. No current config exercises split-K together with the
|
||||
# two-stream main-alloc overlap (qwen3.5 leaves split-K off; kimi is single-stream-coherent).
|
||||
output = torch.zeros((S, R), device=x.device, dtype=torch.float32)
|
||||
launch_kwargs = {
|
||||
"num_warps": 2 if split_k <= 4 else 4,
|
||||
"num_stages": 3,
|
||||
}
|
||||
elif out_alloc_stream is not None:
|
||||
# Allocate the output on the MAIN (consumer) stream when requested, so the caching allocator
|
||||
# frees/reuses it on the consumer's schedule, not the side stream's (cuda-graph WAR — see
|
||||
# lora_overlap_alloc_stream / SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC).
|
||||
with torch.cuda.stream(out_alloc_stream):
|
||||
output = torch.empty((S, R), device=x.device, dtype=x.dtype)
|
||||
else:
|
||||
output = torch.empty((S, R), device=x.device, dtype=x.dtype)
|
||||
|
||||
grid = (
|
||||
num_s_tiles * split_k,
|
||||
batch_info.bs,
|
||||
)
|
||||
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
_sgemm_lora_a_kernel[grid](
|
||||
x,
|
||||
weights,
|
||||
output,
|
||||
R,
|
||||
K,
|
||||
stack_num,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
weights.stride(0),
|
||||
weights.stride(1),
|
||||
weights.stride(2),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
batch_info.seg_lens,
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
batch_info.permutation,
|
||||
sorted_by_adapter,
|
||||
BLOCK_S,
|
||||
BLOCK_R,
|
||||
BLOCK_K,
|
||||
split_k,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
**launch_kwargs,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
# split_k>1 returns the fp32 accumulator directly; the LoRA-B expand casts x to the weight dtype
|
||||
# on-load (fused), dropping the standalone fp32->bf16 copy kernel. split_k==1 already returns x.dtype.
|
||||
return output
|
||||
@@ -0,0 +1,219 @@
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops.gate_up_lora_b import (
|
||||
_CUBLAS_MIN_S_RANK,
|
||||
)
|
||||
from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
|
||||
_resolve_token_positions,
|
||||
get_pdl_launch_metadata,
|
||||
)
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
|
||||
def _sgemm_lora_b_cublas(
|
||||
x: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
base_output: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Single-adapter dense path: one cuBLAS addmm_ over the full output.
|
||||
|
||||
Mirrors the Triton kernel exactly: single slice, K = max_r (the kernel
|
||||
reads the full K; the loader zero-pads the weight tail beyond the
|
||||
adapter's rank), scaling fused via a pre-scaled x.
|
||||
"""
|
||||
if base_output is None:
|
||||
base_output = torch.zeros(
|
||||
(x.shape[0], weights.shape[-2]), device=x.device, dtype=x.dtype
|
||||
)
|
||||
w = weights[batch_info.weight_indices[0]]
|
||||
x_scaled = x * batch_info.scalings[0]
|
||||
base_output.addmm_(x_scaled, w.t())
|
||||
return base_output
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _sgemm_lora_b_kernel(
|
||||
# Pointers to matrices
|
||||
x,
|
||||
weights,
|
||||
output,
|
||||
# Matrix dimensions
|
||||
N, # output_dim
|
||||
K, # r
|
||||
# Strides
|
||||
x_stride_0,
|
||||
x_stride_1,
|
||||
w_stride_0,
|
||||
w_stride_1,
|
||||
w_stride_2,
|
||||
output_stride_0,
|
||||
output_stride_1,
|
||||
# Information on sequence lengths and weight id
|
||||
seg_lens,
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
lora_ranks,
|
||||
sorted_token_ids,
|
||||
# Meta parameters
|
||||
SORTED_BY_ADAPTER: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
# For fused output scaling
|
||||
scalings,
|
||||
ENABLE_PDL: tl.constexpr = False,
|
||||
):
|
||||
"""
|
||||
Computes a segmented batched matrix multiplication for the LoRA B matrix
|
||||
and adds the result to the output in-place.
|
||||
|
||||
When a sequence's rank is 0, the kernel is essentially a no-op, following
|
||||
the convention in pytorch where the product of two matrices of shape (m, 0)
|
||||
and (0, n) is an all-zero matrix of shape (m, n).
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The intermediate tensor from the LoRA 'A' multiplication,
|
||||
of shape `(s, K)`, where `s` is the total number of tokens.
|
||||
weights (torch.Tensor): The LoRA 'B' weights for all available adapters,
|
||||
with shape `(num_lora, N, K)`.
|
||||
output (torch.Tensor): The output tensor of shape `(s, N)`. This can be
|
||||
the base model's output for a fused add operation.
|
||||
"""
|
||||
|
||||
pid_s = tl.program_id(axis=0)
|
||||
pid_n = tl.program_id(axis=1)
|
||||
batch_id = tl.program_id(axis=2)
|
||||
w_index = tl.load(weight_indices + batch_id)
|
||||
rank = tl.load(lora_ranks + w_index)
|
||||
|
||||
# If rank is 0, this kernel is a no-op.
|
||||
if rank == 0:
|
||||
return
|
||||
|
||||
seg_len = tl.load(seg_lens + batch_id)
|
||||
if pid_s * BLOCK_S >= seg_len: # also covers seg_len == 0
|
||||
return
|
||||
seg_start = tl.load(seg_indptr + batch_id)
|
||||
scaling = tl.load(scalings + w_index)
|
||||
|
||||
s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
|
||||
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
|
||||
k_offset = tl.arange(0, BLOCK_K)
|
||||
s_physical = _resolve_token_positions(
|
||||
sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
|
||||
)
|
||||
x_ptrs = x + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
|
||||
w_ptrs = (weights + w_index * w_stride_0) + (
|
||||
k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
|
||||
)
|
||||
|
||||
# GDC wait: ensure the prior kernel (producer of x) has fully completed
|
||||
# before consuming its output.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
n_mask = n_offset[None, :] < N
|
||||
output_ptr = output + (
|
||||
s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
|
||||
)
|
||||
output_mask = (s_offset[:, None] < seg_len) & n_mask
|
||||
|
||||
x_tile = tl.load(
|
||||
x_ptrs,
|
||||
mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K),
|
||||
other=0.0,
|
||||
)
|
||||
w_tile = tl.load(
|
||||
w_ptrs,
|
||||
mask=(k_offset[:, None] < K) & n_mask,
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# cast fused: the split-K shrink returns fp32, plain path bf16 (no-op)
|
||||
partial_sum = tl.dot(x_tile.to(w_tile.dtype), w_tile) * scaling
|
||||
|
||||
# All input reads are done; hint the runtime to launch the dependent kernel.
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
# Store result to output matrix (cast to the OUTPUT dtype: x may be the fp32
|
||||
# split-K shrink accumulator while base_output is bf16)
|
||||
partial_sum = partial_sum.to(output.dtype.element_ty)
|
||||
tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
|
||||
|
||||
|
||||
def sgemm_lora_b_fwd(
|
||||
x: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
base_output: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
# x: (s, max_r)
|
||||
# weights: (num_lora, output_dim, max_r)
|
||||
# output: (s, output_dim)
|
||||
# output_dim is much larger than max_r
|
||||
|
||||
assert x.is_contiguous()
|
||||
assert weights.is_contiguous()
|
||||
assert len(x.shape) == 2
|
||||
assert len(weights.shape) == 3
|
||||
|
||||
S = x.shape[0]
|
||||
N = weights.shape[-2]
|
||||
R = weights.shape[-1]
|
||||
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:
|
||||
return _sgemm_lora_b_cublas(x, weights, batch_info, base_output)
|
||||
# Block shapes
|
||||
BLOCK_S = 16
|
||||
BLOCK_R = triton.next_power_of_2(R)
|
||||
BLOCK_N = 256
|
||||
|
||||
grid = (
|
||||
triton.cdiv(batch_info.max_len, BLOCK_S),
|
||||
triton.cdiv(N, BLOCK_N),
|
||||
batch_info.bs,
|
||||
)
|
||||
|
||||
if base_output is None:
|
||||
output = torch.zeros((S, N), device=x.device, dtype=x.dtype)
|
||||
else:
|
||||
output = base_output
|
||||
|
||||
sorted_by_adapter = batch_info.permutation is not None
|
||||
enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
|
||||
_sgemm_lora_b_kernel[grid](
|
||||
x,
|
||||
weights,
|
||||
output,
|
||||
N,
|
||||
R,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
weights.stride(0),
|
||||
weights.stride(1),
|
||||
weights.stride(2),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
batch_info.seg_lens,
|
||||
batch_info.seg_indptr,
|
||||
batch_info.weight_indices,
|
||||
batch_info.lora_ranks,
|
||||
batch_info.permutation,
|
||||
sorted_by_adapter,
|
||||
BLOCK_S,
|
||||
BLOCK_N,
|
||||
BLOCK_R,
|
||||
batch_info.scalings,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
return output
|
||||
File diff suppressed because it is too large
Load Diff
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional
|
||||
import torch
|
||||
|
||||
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers import deep_gemm_wrapper
|
||||
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
|
||||
from sglang.srt.layers.communicator import get_attn_tp_context
|
||||
@@ -45,6 +46,8 @@ from sglang.srt.state_capturer.indexer_topk import (
|
||||
)
|
||||
from sglang.srt.utils import BumpAllocator
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
|
||||
|
||||
@@ -284,6 +287,15 @@ class DeepseekMLAForwardMixin:
|
||||
q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
|
||||
k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)
|
||||
|
||||
_kvb_q = None
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
# Fork the kv_b q-correction A-step onto the LoRA side stream to overlap the bmm.
|
||||
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
|
||||
kv_b_lora_q_prepare,
|
||||
)
|
||||
|
||||
_kvb_q = kv_b_lora_q_prepare(self, q_nope)
|
||||
|
||||
if self.use_deep_gemm_bmm:
|
||||
(
|
||||
q_nope_val,
|
||||
@@ -367,7 +379,13 @@ class DeepseekMLAForwardMixin:
|
||||
q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
|
||||
|
||||
q_nope_out = q_nope_out.transpose(0, 1)
|
||||
if is_kv_b_lora_active(self):
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
|
||||
kv_b_lora_q_apply,
|
||||
)
|
||||
|
||||
q_nope_out = kv_b_lora_q_apply(self, q_nope, q_nope_out, _kvb_q)
|
||||
elif is_kv_b_lora_active(self):
|
||||
q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out)
|
||||
|
||||
skip_rope_for_dsa_tilelang_fused = self._skip_rope_for_dsa_tilelang_fused()
|
||||
@@ -548,6 +566,15 @@ class DeepseekMLAForwardMixin:
|
||||
)
|
||||
attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
|
||||
|
||||
_kvb_v = None
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
# Fork the kv_b v-correction A-step onto the LoRA side stream to overlap the bmm.
|
||||
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
|
||||
kv_b_lora_v_prepare,
|
||||
)
|
||||
|
||||
_kvb_v = kv_b_lora_v_prepare(self, attn_output)
|
||||
|
||||
if self.use_deep_gemm_bmm:
|
||||
(
|
||||
attn_output_val,
|
||||
@@ -686,7 +713,15 @@ class DeepseekMLAForwardMixin:
|
||||
-1, self.num_local_heads, self.v_head_dim
|
||||
).transpose(0, 1),
|
||||
)
|
||||
if is_kv_b_lora_active(self):
|
||||
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
|
||||
kv_b_lora_v_apply,
|
||||
)
|
||||
|
||||
attn_bmm_output = kv_b_lora_v_apply(
|
||||
self, attn_output, attn_bmm_output, _kvb_v
|
||||
)
|
||||
elif is_kv_b_lora_active(self):
|
||||
attn_bmm_output = apply_kv_b_lora_v_correction(
|
||||
self, attn_output, attn_bmm_output
|
||||
)
|
||||
|
||||
@@ -112,6 +112,8 @@ if is_npu():
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils.hf_transformers_utils import get_rope_config
|
||||
|
||||
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
@@ -468,11 +470,31 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
|
||||
self.alt_stream.wait_stream(current_stream)
|
||||
shared_output = self._forward_shared_experts(hidden_states.clone())
|
||||
|
||||
# ===== TO BE REFACTORED ====
|
||||
# Shared-add overlap (SGLANG_OPT_LORA_SHARED_ADD_OVERLAP): hand the add to the LoRA
|
||||
# MoE dispatch so it overlaps the down-LoRA shrink on the alt stream.
|
||||
staged = False
|
||||
if shared_output is not None and _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
||||
from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import (
|
||||
shared_add_overlap_enabled,
|
||||
stage_shared_expert_add,
|
||||
unstage_shared_expert_add,
|
||||
)
|
||||
|
||||
if shared_add_overlap_enabled():
|
||||
stage_shared_expert_add(shared_output, current_stream)
|
||||
staged = True
|
||||
# ===== END TO BE REFACTORED ====
|
||||
|
||||
with torch.cuda.stream(self.alt_stream):
|
||||
router_output = self._forward_router_experts(hidden_states)
|
||||
|
||||
current_stream.wait_stream(self.alt_stream)
|
||||
|
||||
if staged and unstage_shared_expert_add() is None:
|
||||
# The dispatch consumed the staging (add already enqueued); skip the caller's add.
|
||||
shared_output = None
|
||||
|
||||
return router_output, shared_output
|
||||
|
||||
def forward(
|
||||
|
||||
@@ -223,6 +223,7 @@ MOE_RUNNER_BACKEND_CHOICES = [
|
||||
"triton",
|
||||
"triton_kernel",
|
||||
"flashinfer_trtllm",
|
||||
"experimental_sgl_trtllm",
|
||||
"flashinfer_trtllm_routed",
|
||||
"flashinfer_cutlass",
|
||||
"flashinfer_mxfp4",
|
||||
@@ -3344,7 +3345,7 @@ class ServerArgs:
|
||||
"FlashInfer CuteDSL MoE is enabled. --disable-shared-experts-fusion is automatically set."
|
||||
)
|
||||
|
||||
if self.moe_runner_backend == "flashinfer_trtllm":
|
||||
if self.moe_runner_backend in ["flashinfer_trtllm", "experimental_sgl_trtllm"]:
|
||||
assert self.quantization in [
|
||||
"modelopt_fp4",
|
||||
"fp8",
|
||||
|
||||
Reference in New Issue
Block a user