diff --git a/python/sglang/jit_kernel/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh b/python/sglang/jit_kernel/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh new file mode 100644 index 000000000..ff90efb6d --- /dev/null +++ b/python/sglang/jit_kernel/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh @@ -0,0 +1,453 @@ +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck, Panic, div_ceil + +#include // For LaunchKernel + +#include + +#include +#include + +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 +struct VecLoader; + +template <> +struct VecLoader { + __device__ __forceinline__ static float4 load(const float* base, int vec_idx) { + return reinterpret_cast(base)[vec_idx]; + } +}; + +template <> +struct VecLoader<__nv_bfloat16> { + __device__ __forceinline__ static float4 load(const __nv_bfloat16* base, int vec_idx) { + float2 raw = reinterpret_cast(base)[vec_idx]; // 4 bf16 = 8 bytes + const __nv_bfloat162* packed = reinterpret_cast(&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(base)[vec_idx]; // 4 fp16 = 8 bytes + const __half2* packed = reinterpret_cast(&raw); + float2 lo = __half22float2(packed[0]); + float2 hi = __half22float2(packed[1]); + return make_float4(lo.x, lo.y, hi.x, hi.y); + } +}; + +template +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 +__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; + 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(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 +__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; + 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(warp_scores); + float4* warp_original_scores_v4 = reinterpret_cast(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::load(input_row, vec_idx); + float4 bias_val = VecLoader::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(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 +void launch_for_n(const InputT* input, const BiasT* bias, const GateLaunchArgs& args) { + using namespace host; + using Cfg = GateConfig; + bool use_small_token_kernel = args.num_rows <= Cfg::SMALL_TOKEN_THRESHOLD; + + if (use_small_token_kernel) { + LaunchKernel( + static_cast(args.num_rows), static_cast(Cfg::THREADS_PER_BLOCK_SMALL), args.device)( + kimi_k2_moe_fused_gate_kernel_small_token, + 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(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, + 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 +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()) { + launch_for_n(input, static_cast(bias), args); + } else if (bias_dtype.is_type()) { + launch_for_n(input, static_cast(bias), args); + } else { + launch_for_n(input, static_cast(bias), args); + } +} + +template +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()) { + dispatch_bias(static_cast(input), bias, bias_dtype, args); + } else if (input_dtype.is_type()) { + dispatch_bias(static_cast(input), bias, bias_dtype, args); + } else { + dispatch_bias(static_cast(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(); + + TensorMatcher({N, E}).with_dtype(input_dtype).with_device(device).verify(input); + TensorMatcher({E}).with_dtype(bias_dtype).with_device(device).verify(bias); + TensorMatcher({N, K}).with_dtype().with_device(device).verify(output); + TensorMatcher({N, K}).with_dtype().with_device(device).verify(indices); + + const auto num_rows = static_cast(N.unwrap()); + const auto num_experts = static_cast(E.unwrap()); + + RuntimeCheck(topk <= 8, "kimi_k2_moe_fused_gate only supports topk <= 8, got ", topk); + + const GateLaunchArgs args{ + .output = static_cast(output.data_ptr()), + .indices = static_cast(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 diff --git a/python/sglang/jit_kernel/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu b/python/sglang/jit_kernel/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu new file mode 100644 index 000000000..e2f672088 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu @@ -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 +#include + +#include + +#include + +#include + +#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 +__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(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(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 +__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( + 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 +__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(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( + 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(i) % top_k == 0) { + int m = static_cast(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 +__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(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( + 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(i) % top_k == 0) { + int m = static_cast(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 +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(topk_ids.data_ptr()); + const int32_t* tlm_ptr = static_cast(token_lora_mapping.data_ptr()); + bool* token_lora_mask_ptr = static_cast(token_lora_mask.data_ptr()); + int32_t* sorted_token_ids_ptr = static_cast(sorted_token_ids.data_ptr()); + int32_t* expert_ids_ptr = static_cast(expert_ids.data_ptr()); + int32_t* num_tokens_post_pad_ptr = static_cast(num_tokens_post_pad.data_ptr()); + int32_t* cumsum_buffer_ptr = static_cast(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; + 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; + 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; + 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 diff --git a/python/sglang/jit_kernel/csrc/trtllm_lora_temp/topk_softmax_pack.cuh b/python/sglang/jit_kernel/csrc/trtllm_lora_temp/topk_softmax_pack.cuh new file mode 100644 index 000000000..7a85f612a --- /dev/null +++ b/python/sglang/jit_kernel/csrc/trtllm_lora_temp/topk_softmax_pack.cuh @@ -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 // TensorMatcher, SymbolicSize, SymbolicDevice +#include // RuntimeCheck + +#include // LaunchKernel, fp32_t/fp16_t/bf16_t, is_type + +#include +#include + +#include +#include + +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 +class alignas(Alignment) AlignedArray { + T data[N]; +}; + +template +__device__ float convert_to_float(T x) { + if constexpr (std::is_same_v) { + return __half2float(x); + } else if constexpr (std::is_same_v) { + return __bfloat162float(x); + } else if constexpr (std::is_same_v) { + return x; + } else { + return static_cast(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(__bfloat16_as_ushort(__float2bfloat16(w))); + return static_cast((static_cast(id) << 16) | wbits); +} + +template +__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 row_chunk_temp[VPT]; + AccessType* row_chunk_vec_ptr = reinterpret_cast(&row_chunk_temp); + const AccessType* vec_thread_read_ptr = reinterpret_cast(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(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 +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 +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; + 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, + input, + output, + num_rows, + indices, + packed_output, + num_token_non_padded, + k, + renormalize); +} + +template +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( \ + 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 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(); + + TensorMatcher({N, E}).with_dtype().with_device(device_).verify(gating_output); + TensorMatcher({N, K}).with_dtype().with_device(device_).verify(topk_weights); + TensorMatcher({N, K}).with_dtype().with_device(device_).verify(topk_indices); + TensorMatcher({N, K}).with_dtype().with_device(device_).verify(packed); + + const int32_t* ntnp_ptr = nullptr; + if (num_token_non_padded.has_value()) { + SymbolicSize One{"ntnp_numel"}; + TensorMatcher({One}).with_dtype().with_device(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(num_token_non_padded.value().data_ptr()); + } + + const int num_tokens = static_cast(N.unwrap()); + const int num_experts = static_cast(E.unwrap()); + const int topk = static_cast(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(topk_weights.data_ptr()); + auto* indices_ptr = static_cast(topk_indices.data_ptr()); + auto* packed_ptr = static_cast(packed.data_ptr()); + + if (is_type(gating_output.dtype())) { + dispatchExperts( + static_cast(gating_output.data_ptr()), + weights_ptr, + indices_ptr, + packed_ptr, + ntnp_ptr, + num_tokens, + num_experts, + topk, + renormalize, + device); + } else if (is_type(gating_output.dtype())) { + dispatchExperts<__half>( + static_cast(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(gating_output.data_ptr()), + weights_ptr, + indices_ptr, + packed_ptr, + ntnp_ptr, + num_tokens, + num_experts, + topk, + renormalize, + device); + } +} + +} // namespace diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/SOURCE.md b/python/sglang/jit_kernel/trtllm_lora_temp/SOURCE.md new file mode 100644 index 000000000..dbe24b686 --- /dev/null +++ b/python/sglang/jit_kernel/trtllm_lora_temp/SOURCE.md @@ -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. diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py b/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py new file mode 100644 index 000000000..5269f90d3 --- /dev/null +++ b/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py @@ -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", +] diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/core.py b/python/sglang/jit_kernel/trtllm_lora_temp/core.py new file mode 100644 index 000000000..c300f92a8 --- /dev/null +++ b/python/sglang/jit_kernel/trtllm_lora_temp/core.py @@ -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 diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh new file mode 100644 index 000000000..10dc14213 --- /dev/null +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh @@ -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 +#include + +#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh" +#include +#include +#include + +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 +__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(&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(g); +#pragma unroll + for (int k = 0; k < 4; ++k) + gu.v[k] = gp[k]; + if (loraDelta != nullptr) { + int4 const* dlp = reinterpret_cast(dlo); + int4 const* dhp = reinterpret_cast(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(&loraInputOut[liBaseRow + h0]) = vec; + } + } + + // ---- per-token scale: blockReduce amax, broadcast via smem (no gmem round-trip) ---- + using BlockReduce = cub::BlockReduce; + __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(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 + <<>>( + m, + innerHalf, + innerDim, + gateUp, + loraDelta, + loraInputOut, + expandedIdxToPermutedIdx, + globalScaleInv, + weightOutput, + scaleOutput, + perTokenScaleOutput); + }; + auto withFastMath = [&](auto layoutTag) { + if (disableFp4FastMath) { + launch(layoutTag, std::integral_constant{}); + } else { + launch(layoutTag, std::integral_constant{}); + } + }; + 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{}); + } else { + withFastMath( + std::integral_constant{}); + } +} + +} // namespace sgl_fused_act_quant +} // namespace flashinfer diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu new file mode 100644 index 000000000..3954f32c0 --- /dev/null +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu @@ -0,0 +1,1137 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "flashinfer/exception.h" +#include "flashinfer/trtllm/fused_moe/DevKernel.h" +#include "flashinfer/utils.cuh" +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper function for array conversion +template +__host__ __device__ constexpr static U arrayConvert(T const& input) { + cutlass::NumericArrayConverter converter; + return converter(input); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace moe::dev { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace activation { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float silu(float x) { + return x / (1.0f + expf(-x)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void activationKernel(KernelParams params) { + using Type = typename KernelParams::Type; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // immediately trigger the secondary kernel when using PDL, then wait on primary + if constexpr (KernelParams::UsePdl) { + cudaTriggerProgrammaticLaunchCompletion(); + cudaGridDependencySynchronize(); + } +#endif + + for (int tokenIdx = blockIdx.z; tokenIdx < params.numTokens; tokenIdx += gridDim.z) { + // Look over experts per token + for (int k = blockIdx.y; k < params.topK; k += gridDim.y) { + int const expandedIdx = tokenIdx * params.topK + k; + int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; + + // Loop over hidden dim + for (int hiddenIdx = threadIdx.x + blockDim.x * blockIdx.x; hiddenIdx < params.innerDim / 2; + hiddenIdx += blockDim.x * gridDim.x) { + if (permutedIdx == -1) { + if (params.activationLoraInputOutPtr != nullptr) { + int64_t const activationIdx = (int64_t)expandedIdx * (params.innerDim / 2) + hiddenIdx; + params.activationLoraInputOutPtr[activationIdx] = cutlass::bfloat16_t(0.0f); + } + continue; + } + + // Use int64_t to avoid overflow when permutedIdx * innerDim > INT32_MAX + int64_t const permBase = (int64_t)permutedIdx * params.innerDim; + + // Contiguous input: gate = col hiddenIdx, up = col innerDim/2 + hiddenIdx. + // Interleaved input (fused de-interleave): gate = col 2*hiddenIdx, up = col 2*hiddenIdx+1. + float x1, x2; + if (params.interleavedGateUpInput) { + x1 = (float)params.inPtr[permBase + 2 * hiddenIdx]; + x2 = (float)params.inPtr[permBase + 2 * hiddenIdx + 1]; + } else { + x1 = (float)params.inPtr[permBase + hiddenIdx]; + x2 = (float)params.inPtr[permBase + hiddenIdx + params.innerDim / 2]; + } + if (params.gateUpLoraDeltaPtr != nullptr) { + int64_t const loraBaseIdx = (int64_t)expandedIdx * params.innerDim + hiddenIdx; + x1 += static_cast(params.gateUpLoraDeltaPtr[loraBaseIdx + params.innerDim / 2]); + x2 += static_cast(params.gateUpLoraDeltaPtr[loraBaseIdx]); + } + + float act = silu(x2); + Type out = (Type)(act * x1); + if (params.activationLoraInputOutPtr != nullptr) { + int64_t const activationIdx = (int64_t)expandedIdx * (params.innerDim / 2) + hiddenIdx; + params.activationLoraInputOutPtr[activationIdx] = static_cast(act * x1); + } + + int64_t const outIdx = (int64_t)permutedIdx * (params.innerDim / 2) + hiddenIdx; + params.outPtr[outIdx] = out; + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Vectorized bf16 SwiGLU+LoRA activation for the FP4-LoRA path (interleaved gate/up input). +// Each thread processes 4 consecutive hidden pairs per step: one 128-bit load for the +// interleaved gate/up, two 64-bit loads for the LoRA delta, and 64-bit stores. This raises +// memory-level parallelism toward the HBM bandwidth roofline (vs the scalar activationKernel, +// which does 1 element/thread with per-element scalar loads). Numerically identical to +// activationKernel (same per-element float math + silu), so the testbed asserts bitwise +// equality. Requires interleaved bf16 input and (innerDim/2) % 4 == 0 (run() guards this). +__global__ void activationKernelOpt( + cutlass::bfloat16_t const* __restrict__ inPtr, // interleaved gate/up GEMM1 output + cutlass::bfloat16_t* __restrict__ outPtr, // activated [.., innerDim/2] + cutlass::bfloat16_t const* __restrict__ gateUpLoraDeltaPtr, // may be null + cutlass::bfloat16_t* __restrict__ activationLoraInputOutPtr, // may be null + int const* __restrict__ expandedIdxToPermutedIdx, + int innerDim, + int numTokens, + int topK) { + int const innerHalf = innerDim / 2; + for (int tokenIdx = blockIdx.z; tokenIdx < numTokens; tokenIdx += gridDim.z) { + for (int k = blockIdx.y; k < topK; k += gridDim.y) { + int const expandedIdx = tokenIdx * topK + k; + int const permutedIdx = expandedIdxToPermutedIdx[expandedIdx]; + for (int h = (threadIdx.x + blockDim.x * blockIdx.x) * 4; h < innerHalf; h += blockDim.x * gridDim.x * 4) { + int64_t const liBase = (int64_t)expandedIdx * innerHalf + h; + if (permutedIdx == -1) { + if (activationLoraInputOutPtr != nullptr) { + *reinterpret_cast(&activationLoraInputOutPtr[liBase]) = make_int2(0, 0); + } + continue; + } + + // gate/up: 8 interleaved bf16 (g0,u0,...,g3,u3) via one 128-bit load. + int64_t const permBase = (int64_t)permutedIdx * innerDim + 2 * h; + int4 const rawGU = *reinterpret_cast(&inPtr[permBase]); + cutlass::bfloat16_t const* gu = reinterpret_cast(&rawGU); + + float gate[4], up[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + gate[j] = (float)gu[2 * j]; // even col 2k -> x1 (gate) + up[j] = (float)gu[2 * j + 1]; // odd col 2k+1 -> x2 (up) + } + + if (gateUpLoraDeltaPtr != nullptr) { + // delta is contiguous [gate | up] per expandedIdx row: up += delta[h], gate += delta[h+innerHalf]. + int64_t const dBase = (int64_t)expandedIdx * innerDim + h; + int2 const rawUp = *reinterpret_cast(&gateUpLoraDeltaPtr[dBase]); + int2 const rawGate = *reinterpret_cast(&gateUpLoraDeltaPtr[dBase + innerHalf]); + cutlass::bfloat16_t const* dUp = reinterpret_cast(&rawUp); + cutlass::bfloat16_t const* dGate = reinterpret_cast(&rawGate); +#pragma unroll + for (int j = 0; j < 4; ++j) { + up[j] += (float)dUp[j]; + gate[j] += (float)dGate[j]; + } + } + + __align__(8) cutlass::bfloat16_t res[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + res[j] = (cutlass::bfloat16_t)(silu(up[j]) * gate[j]); + } + int2 const packed = *reinterpret_cast(res); + + int64_t const outBase = (int64_t)permutedIdx * innerHalf + h; + *reinterpret_cast(&outPtr[outBase]) = packed; + if (activationLoraInputOutPtr != nullptr) { + *reinterpret_cast(&activationLoraInputOutPtr[liBase]) = packed; + } + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Float4Max { + __device__ __forceinline__ float4 operator()(float4 const& a, float4 const& b) const { + float4 result; + result.x = fmaxf(a.x, b.x); + result.y = fmaxf(a.y, b.y); + result.z = fmaxf(a.z, b.z); + result.w = fmaxf(a.w, b.w); + return result; + } +}; + +struct Float2Max { + __device__ __forceinline__ float2 operator()(float2 const& a, float2 const& b) const { + float2 result; + result.x = fmaxf(a.x, b.x); + result.y = fmaxf(a.y, b.y); + return result; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__device__ __forceinline__ VecType packedTypeFromArray(float data[size]) { + return {}; +} + +template <> +__device__ __forceinline__ float4 packedTypeFromArray(float data[4]) { + float4 result; + result.x = data[0]; + result.y = data[1]; + result.z = data[2]; + result.w = data[3]; + return result; +} + +template <> +__device__ __forceinline__ float2 packedTypeFromArray(float data[2]) { + float2 result; + result.x = data[0]; + result.y = data[1]; + return result; +} + +template <> +__device__ __forceinline__ float packedTypeFromArray(float data[1]) { + return data[0]; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__device__ __forceinline__ cutlass::Array arrayFromPackedType(PackedType data) { + return cutlass::Array{}; +} + +template <> +__device__ __forceinline__ cutlass::Array arrayFromPackedType(float4 data) { + return cutlass::Array{data.x, data.y, data.z, data.w}; +} + +template <> +__device__ __forceinline__ cutlass::Array arrayFromPackedType(float2 data) { + return cutlass::Array{data.x, data.y}; +} + +template <> +__device__ __forceinline__ cutlass::Array arrayFromPackedType(float data) { + return cutlass::Array{data}; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct KernelTraits; + +template <> +struct KernelTraits<4> { + using MaxOp = Float4Max; + using PackedType = float4; +}; + +template <> +struct KernelTraits<2> { + using MaxOp = Float2Max; + using PackedType = float2; +}; + +template <> +struct KernelTraits<1> { + using MaxOp = cuda::maximum<>; + using PackedType = float; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +constexpr int DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA = 128; + +template +__global__ void activationDeepSeekKernel(KernelParams params) { + using Type = typename KernelParams::Type; + int32_t constexpr NumTokensPerCta = KernelParams::NumTokensPerCta; + using KernelTraits = KernelTraits; + using MaxOp = typename KernelTraits::MaxOp; + using PackedType = typename KernelTraits::PackedType; + using BlockReduce = cub::BlockReduce; + + __shared__ float s_scaleOutArr[NumTokensPerCta]; + __shared__ typename BlockReduce::TempStorage tempStorage; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // immediately trigger the secondary kernel when using PDL, then wait on primary + if constexpr (KernelParams::UsePdl) { + cudaTriggerProgrammaticLaunchCompletion(); + cudaGridDependencySynchronize(); + } +#endif + + // The largest (finite) value that can be represented using E4m3. + float constexpr E4m3MaxVal{448.f}; + + int const totalNumPaddedTokens = params.totalNumPaddedTokens[0]; + // Loop over tokens + float scale1Arr[NumTokensPerCta]; + float scale2Arr[NumTokensPerCta]; + float dataX1Arr[NumTokensPerCta]; + float dataX2Arr[NumTokensPerCta]; + float outArr[NumTokensPerCta]; + float absOutArr[NumTokensPerCta]; + int permutedIdxArr[NumTokensPerCta]; + + // Loop over tokens + for (int k = blockIdx.z; k < params.topK; k += gridDim.z) { + for (int tokenCtaIdx = blockIdx.y * NumTokensPerCta; tokenCtaIdx < params.numTokens; + tokenCtaIdx += gridDim.y * NumTokensPerCta) { + for (int hiddenIdx = threadIdx.x + blockDim.x * blockIdx.x; hiddenIdx < params.innerDim / 2; + hiddenIdx += blockDim.x * gridDim.x) { +#pragma unroll + for (int tokenInCtaIdx = 0; tokenInCtaIdx < NumTokensPerCta; tokenInCtaIdx++) { + scale1Arr[tokenInCtaIdx] = 0.0f; + scale2Arr[tokenInCtaIdx] = 0.0f; + dataX1Arr[tokenInCtaIdx] = 0.0f; + dataX2Arr[tokenInCtaIdx] = 0.0f; + outArr[tokenInCtaIdx] = 0.0f; + absOutArr[tokenInCtaIdx] = 0.0f; + } +#pragma unroll + for (int tokenInCtaIdx = 0; tokenInCtaIdx < NumTokensPerCta; tokenInCtaIdx++) { + int const tokenIdx = tokenCtaIdx + tokenInCtaIdx; + if (tokenIdx >= params.numTokens) { + break; + } + + int const expandedIdx = tokenIdx * params.topK + k; + int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; + permutedIdxArr[tokenInCtaIdx] = permutedIdx; + if (permutedIdx == -1) { + continue; + } + + // Process blocks for this CTA + // Use int64_t to avoid overflow when permutedIdx * innerDim > INT32_MAX + int64_t const baseIdx = (int64_t)permutedIdx * params.innerDim + hiddenIdx; + + int64_t const scale1Idx = (int64_t)permutedIdx + (int64_t)totalNumPaddedTokens * (hiddenIdx / 128); + int64_t const scale2Idx = + (int64_t)permutedIdx + (int64_t)totalNumPaddedTokens * ((hiddenIdx / 128) + (params.innerDim / 2 / 128)); + + scale1Arr[tokenInCtaIdx] = params.inDqSfsPtr[scale1Idx]; + scale2Arr[tokenInCtaIdx] = params.inDqSfsPtr[scale2Idx]; + dataX1Arr[tokenInCtaIdx] = static_cast(params.inPtr[baseIdx]); + dataX2Arr[tokenInCtaIdx] = static_cast(params.inPtr[baseIdx + params.innerDim / 2]); + } + +#pragma unroll + for (int tokenInCtaIdx = 0; tokenInCtaIdx < NumTokensPerCta; tokenInCtaIdx++) { + float x1 = scale1Arr[tokenInCtaIdx] * dataX1Arr[tokenInCtaIdx]; + float x2 = scale2Arr[tokenInCtaIdx] * dataX2Arr[tokenInCtaIdx]; + auto const tokenIdx = tokenCtaIdx + tokenInCtaIdx; + if (params.gateUpLoraDeltaPtr != nullptr && tokenIdx < params.numTokens) { + int const expandedIdx = tokenIdx * params.topK + k; + int64_t const loraBaseIdx = (int64_t)expandedIdx * params.innerDim + hiddenIdx; + x1 += static_cast(params.gateUpLoraDeltaPtr[loraBaseIdx + params.innerDim / 2]); + x2 += static_cast(params.gateUpLoraDeltaPtr[loraBaseIdx]); + } + float act = silu(x2); + float out = act * x1; + outArr[tokenInCtaIdx] = out; + absOutArr[tokenInCtaIdx] = fabsf(out); + } + + auto absOutPacked = packedTypeFromArray(absOutArr); + auto aMaxPacked = BlockReduce(tempStorage).Reduce(absOutPacked, MaxOp{}); + auto aMaxArr = arrayFromPackedType(aMaxPacked); + +#pragma unroll + for (int tokenInCtaIdx = 0; tokenInCtaIdx < NumTokensPerCta; tokenInCtaIdx++) { + if (threadIdx.x == 0) { + auto const tokenIdx = tokenCtaIdx + tokenInCtaIdx; + if (tokenIdx >= params.numTokens) { + break; + } + int const permutedIdx = permutedIdxArr[tokenInCtaIdx]; + if (permutedIdx == -1) { + continue; + } + // Make sure the scale is strictly positive to avoid division by zero in case the + // maximum is zero. + float scaleOut = fmaxf(aMaxArr[tokenInCtaIdx] / E4m3MaxVal, std::numeric_limits::min()); + s_scaleOutArr[tokenInCtaIdx] = scaleOut; + int64_t const scaleOut_idx = + (int64_t)permutedIdxArr[tokenInCtaIdx] + (int64_t)totalNumPaddedTokens * (hiddenIdx / 128); + params.outDqSfsPtr[scaleOut_idx] = scaleOut; + } + } + __syncthreads(); + +#pragma unroll + for (int tokenInCtaIdx = 0; tokenInCtaIdx < NumTokensPerCta; tokenInCtaIdx++) { + auto const tokenIdx = tokenCtaIdx + tokenInCtaIdx; + if (tokenIdx >= params.numTokens) { + break; + } + int const permutedIdx = permutedIdxArr[tokenInCtaIdx]; + if (permutedIdx == -1) { + if (params.activationLoraInputOutPtr != nullptr) { + int const expandedIdx = tokenIdx * params.topK + k; + int64_t const activationIdx = (int64_t)expandedIdx * (params.innerDim / 2) + hiddenIdx; + params.activationLoraInputOutPtr[activationIdx] = cutlass::bfloat16_t(0.0f); + } + continue; + } + float const scaleOut = s_scaleOutArr[tokenInCtaIdx]; + int64_t const outIdx = (int64_t)permutedIdx * (params.innerDim / 2) + hiddenIdx; + params.outPtr[outIdx] = static_cast(outArr[tokenInCtaIdx] / scaleOut); + if (params.activationLoraInputOutPtr != nullptr) { + int const expandedIdx = tokenIdx * params.topK + k; + int64_t const activationIdx = (int64_t)expandedIdx * (params.innerDim / 2) + hiddenIdx; + params.activationLoraInputOutPtr[activationIdx] = static_cast(outArr[tokenInCtaIdx]); + } + } + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void run(Data const& data, void* stream) { + if (data.mDtypeElt == tg::Dtype::E2m1) { + // Note: this should be unreachable because the options are checked beforehand. + // E2m1 requires using higher-precision intermediate data (bf16). + FLASHINFER_CHECK(false, "Activation with E2m1_t isn't supported."); + return; + } + + if (data.mUseDeepSeekFp8) { + constexpr int NUM_ELTS_PER_LOAD = 1; + constexpr int NUM_ELTS_PER_SF = 128; + + int device{-1}; + cudaGetDevice(&device); + int numSms = 0; + cudaDeviceGetAttribute(&numSms, cudaDevAttrMultiProcessorCount, device); + + // Output dimension is innerDim / 2, and each scale block is 128 elements + int const outputDim = data.innerDim / 2; + int const numScaleBlocks = (outputDim + NUM_ELTS_PER_SF - 1) / NUM_ELTS_PER_SF; + int const gridSizeX = (numScaleBlocks + NUM_ELTS_PER_LOAD - 1) / NUM_ELTS_PER_LOAD; + + auto numCtas = gridSizeX * data.numTokens * data.topK; + // FIXME: This is heruistic based on very short benchmark. + int numTokensPerCta = 1; + if (numCtas > numSms * 32) { + numTokensPerCta = 4; + } else if (numCtas > numSms * 4) { + numTokensPerCta = 2; + } else { + numTokensPerCta = 1; + } + + int const gridSizeY = std::min(8192, (data.numTokens + numTokensPerCta - 1) / numTokensPerCta); + + const dim3 grid(gridSizeX, gridSizeY, data.topK); + + LAUNCH_ACTIVATION( + data, activationDeepSeekKernel, numTokensPerCta, grid, DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA, 0, stream); + } else if ( + data.actOptMode == 1 && data.mDtypeElt == tg::Dtype::Bfloat16 && data.interleavedGateUpInput && + (data.innerDim / 2) % 4 == 0) { + int const numThreads = 256; + int const innerHalf = data.innerDim / 2; + int const defaultGx = (innerHalf / 4 + numThreads - 1) / numThreads; + int const gridX = data.actGridXOverride > 0 ? data.actGridXOverride : defaultGx; + const dim3 grid(gridX, data.topK, std::min(8192, data.numTokens)); + + activationKernelOpt<<>>( + static_cast(data.inPtr), + static_cast(data.outPtr), + data.gateUpLoraDeltaPtr, + data.activationLoraInputOutPtr, + data.expandedIdxToPermutedIdx, + data.innerDim, + data.numTokens, + data.topK); + } else { + int const numThreads = 256; + int const gridX = data.actGridXOverride > 0 ? data.actGridXOverride : (data.innerDim / 128); + const dim3 grid(gridX, data.topK, std::min(8192, data.numTokens)); + + LAUNCH_ACTIVATION(data, activationKernel, 1, grid, numThreads, 0, stream); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace activation + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace convertsf { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +namespace dev { +// Compute the offset that corresponds to (dataRowIdx, dataBlkColIdx) in the SF tensor where +// dataRowIdx and dataBlkColIdx are the respective indices of the row and the block of 16 elts +// from the K dim in the tensor of data. +inline __device__ int64_t getSfOffset(int32_t dataRowIdx, int32_t dataBlkColIdx, int32_t numDataBlksPerRow) { + // The number of rows of SF per block. + static int32_t constexpr NumRowsPerSfBlock = 128; + // The number of cols of SF per block. + static int32_t constexpr NumColsPerSfBlock = 4; + // The size of each SF block. + static int32_t constexpr NumBytesPerSfBlock = NumRowsPerSfBlock * NumColsPerSfBlock; + + // The number of rows of data per SF block. + static int32_t constexpr NumDataRowsPerSfBlock = NumRowsPerSfBlock; + // The number of cols of blocks of data per SF block. + static int32_t constexpr NumDataBlkColsPerSfBlock = NumColsPerSfBlock; + + // The row of the SF block in the SF tensor. + int sfBlkRowIdx = dataRowIdx / NumDataRowsPerSfBlock; + // The col of the SF block in the SF tensor. + int sfBlkColIdx = dataBlkColIdx / NumDataBlkColsPerSfBlock; + // The blocks are stored row-major in the tensor of scaling factors. + int sfBlkIdx = sfBlkRowIdx * numDataBlksPerRow / NumDataBlkColsPerSfBlock + sfBlkColIdx; + + // Find the row in the SF block. + int sfRowIdx = (dataRowIdx % 32) * 4 + (dataRowIdx % NumDataRowsPerSfBlock) / 32; + // Find the col in the SF block. + int sfColIdx = (dataBlkColIdx % 4); + + // Compute the offset in bytes. + return sfBlkIdx * NumBytesPerSfBlock + sfRowIdx * NumColsPerSfBlock + sfColIdx; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Given the GMEM address of an output element, compute the offset of the corresponding scaling +// factor in the SF tensor. Optionally, a startTokenIndex can be provided if the first token is not +// the start token in the SF tensor. This is useful when inflight batching is enabled in TRT-LLM, +// where the context and generation output are stored as one output tensor. In this case, the +// generation output may not start with zero offset in the SF output tensor. +template +inline __device__ int64_t getSfOffset(int64_t gmemOffsetInBytes, int32_t hiddenDim, int32_t startTokenIdx = 0) { + // The number of elements per sf. + int32_t constexpr NumEltsPerSf = 16; + // The GMEM offset of the output element. + int64_t gmemOffset = gmemOffsetInBytes * 8 /*bits*/ / NumBitsPerElt; + // The row/col indices of the corresponding SF element. + int32_t sfRowIdx = gmemOffset / hiddenDim + startTokenIdx; + int32_t sfColIdx = (gmemOffset % hiddenDim) / NumEltsPerSf; + // Compute the SF offset. + return getSfOffset(sfRowIdx, sfColIdx, hiddenDim / NumEltsPerSf); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// TODO(tizheng): Refactor to track gmem offset instead of doing pointer subtraction. +template +inline __device__ int64_t +getSfOffset(void const* gmemOutPtr, void const* gmemBasePtr, int32_t hiddenDim, int32_t startTokenIdx = 0) { + return getSfOffset( + reinterpret_cast(gmemOutPtr) - reinterpret_cast(gmemBasePtr), hiddenDim, startTokenIdx); +} + +} // namespace dev + +// TODO: it would be nice to move some of that logic to Fp4Utils.h +template +inline __device__ int32_t getSfOffset(int32_t dataRowIdx, int32_t dataBlkColIdx, int32_t numDataBlksPerRow) { + if constexpr (Layout == tg::SfLayout::Linear) { + return numDataBlksPerRow * dataRowIdx + dataBlkColIdx; + } else if constexpr (Layout == tg::SfLayout::R128c4) { + return static_cast(dev::getSfOffset(dataRowIdx, dataBlkColIdx, numDataBlksPerRow)); + } else if constexpr (Layout == tg::SfLayout::R8c4 || Layout == tg::SfLayout::R8c16) { + static int32_t constexpr NumRowsPerSfBlock = 8; + static int32_t constexpr NumColsPerSfBlock = (Layout == tg::SfLayout::R8c4) ? 4 : 16; + static int32_t constexpr NumBytesPerSfBlock = NumRowsPerSfBlock * NumColsPerSfBlock; + int sfBlkRowIdx = dataRowIdx / NumRowsPerSfBlock; + int sfBlkColIdx = dataBlkColIdx / NumColsPerSfBlock; + int sfBlkIdx = sfBlkRowIdx * numDataBlksPerRow / NumColsPerSfBlock + sfBlkColIdx; + int sfRowIdx = dataRowIdx % NumRowsPerSfBlock; + int sfColIdx = dataBlkColIdx % NumColsPerSfBlock; + return sfBlkIdx * NumBytesPerSfBlock + sfRowIdx * NumColsPerSfBlock + sfColIdx; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__device__ void convertSfCommon(KernelParams params) { + // Note: it's assumed that the number of scaling factors per row is a multiple of 4. + constexpr int VecSize = 4; + using VecType = uint32_t; + static_assert(sizeof(VecType) == VecSize); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Immediately trigger the secondary kernel when using PDL, then wait on primary. + if constexpr (KernelParams::UsePdl) { + cudaTriggerProgrammaticLaunchCompletion(); + cudaGridDependencySynchronize(); + } +#endif + + // TODO: consider optimizing if used in production. + // This is a naive kernel. It's not doing coalesced loads. + + int const numSfPerRow = params.hiddenDimSf; + + for (int tokenIdx = blockIdx.y; tokenIdx < params.numTokens; tokenIdx += gridDim.y) { + for (int hiddenSfVecIdx = threadIdx.x + blockDim.x * blockIdx.x; hiddenSfVecIdx < numSfPerRow / VecSize; + hiddenSfVecIdx += blockDim.x * gridDim.x) { + // Index of the first SF in the vector. + int const hiddenSfIdx = VecSize * hiddenSfVecIdx; + + // Load scale factors. + int sfIdxIn = getSfOffset(tokenIdx, hiddenSfIdx, numSfPerRow); + const VecType sfVec = reinterpret_cast(params.inSfPtr)[sfIdxIn / VecSize]; + + // Store scale factors. + int const sfIdxOut = getSfOffset(tokenIdx, hiddenSfIdx, numSfPerRow); + reinterpret_cast(params.outSfPtr)[sfIdxOut / VecSize] = sfVec; + } + } +} + +#define CONVERT_FP4_SF_KERNEL(LayoutSrc, LayoutDst) \ + template \ + __global__ void convertSf##LayoutSrc##To##LayoutDst##Kernel(KernelParams params) { \ + convertSfCommon(params); \ + } +// We only need a conversion to the linear layout. +CONVERT_FP4_SF_KERNEL(R128c4, Linear); +CONVERT_FP4_SF_KERNEL(R8c4, Linear); +CONVERT_FP4_SF_KERNEL(R8c16, Linear); +#undef CONVERT_FP4_SF_KERNEL + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void run(Data const& data, void* stream) { + constexpr int VecSize = 4; + int const numThreads = 128; + int const numBlocksX = (data.hiddenDimSf / VecSize - 1 + numThreads) / numThreads; + int const numBlocksY = std::min(8192, data.numTokens); + dim3 numBlocks(numBlocksX, numBlocksY); +#define CONVERT_FP4_SF_LAUNCH(LayoutSrc, LayoutDst) \ + if (data.sfLayoutSrc == tg::SfLayout::LayoutSrc && data.sfLayoutDst == tg::SfLayout::LayoutDst) { \ + LAUNCH_PDL( \ + data, \ + false, \ + cutlass::float_e4m3_t, \ + convertSf##LayoutSrc##To##LayoutDst##Kernel, \ + numBlocks, \ + numThreads, \ + 0, \ + stream); \ + return; \ + } + CONVERT_FP4_SF_LAUNCH(R128c4, Linear); + CONVERT_FP4_SF_LAUNCH(R8c4, Linear); + CONVERT_FP4_SF_LAUNCH(R8c16, Linear); +#undef CONVERT_FP4_SF_LAUNCH +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace convertsf + +namespace permute { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void permuteKernel(KernelParams params) { + using Type = typename KernelParams::Type; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // immediately trigger the secondary kernel when using PDL, then wait on primary + if constexpr (KernelParams::UsePdl) { + cudaTriggerProgrammaticLaunchCompletion(); + cudaGridDependencySynchronize(); + } +#endif + + for (int tokenIdx = blockIdx.y; tokenIdx < params.numTokens; tokenIdx += gridDim.y) { + // Loop over hidden dim + for (int hiddenIdx = threadIdx.x + blockDim.x * blockIdx.x; hiddenIdx < params.hiddenDim; + hiddenIdx += blockDim.x * gridDim.x) { + // Load chunk of token into registers + const Type data = params.inPtr[tokenIdx * params.hiddenDim + hiddenIdx]; + + // Write to topK places + for (int k = 0; k < params.topK; k++) { + int const expandedIdx = tokenIdx * params.topK + k; + int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; + params.outPtr[permutedIdx * params.hiddenDim + hiddenIdx] = data; + } + } + if (params.useDeepSeekFp8) { + for (int scaleIdx = threadIdx.x + blockDim.x * blockIdx.x; scaleIdx < params.hiddenDim / 128; + scaleIdx += blockDim.x * gridDim.x) { + for (int k = 0; k < params.topK; k++) { + int const expandedIdx = tokenIdx * params.topK + k; + int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; + + int const idx_in = tokenIdx + params.numTokens * scaleIdx; + int const idx_out = permutedIdx + params.totalNumPaddedTokens[0] * scaleIdx; + + params.outDqSfsPtr[idx_out] = params.inDqSfsPtr[idx_in]; + } + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void run(Data const& data, void* stream) { + int const numThreads = 256; + int const numBlocksX = (data.hiddenDim - 1 + numThreads) / numThreads; + int const numBlocksY = std::min(8192, data.numTokens); + dim3 numBlocks(numBlocksX, numBlocksY); + + LAUNCH(data, permuteKernel, numBlocks, numThreads, 0, stream); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace permute + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace finalize { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void finalizeKernel(KernelParams params) { + using Type = typename KernelParams::Type; + using TypeExpW = typename KernelParams::TypeExpW; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // wait on primary kernel when using PDL + if constexpr (KernelParams::UsePdl) { + cudaGridDependencySynchronize(); + } +#endif + + for (int tokenIdx = blockIdx.y; tokenIdx < params.numTokens; tokenIdx += gridDim.y) { + // Loop over hidden dim + for (int hiddenIdx = threadIdx.x + blockDim.x * blockIdx.x; hiddenIdx < params.hiddenDim; + hiddenIdx += blockDim.x * gridDim.x) { + // Accumulate chunk of token into registers + float data = 0.0F; + + // Write to topK places + for (int k = 0; k < params.topK; k++) { + int const expandedIdx = tokenIdx * params.topK + k; + int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; + + if (permutedIdx == -1) { + continue; + } + + if (params.expertWeightsPtr != nullptr) { + TypeExpW const scale = params.expertWeightsPtr[expandedIdx]; + data += float{scale} * float{params.inPtr[permutedIdx * params.hiddenDimPadded + hiddenIdx]}; + } else { + data += float{params.inPtr[permutedIdx * params.hiddenDimPadded + hiddenIdx]}; + } + } + + params.outPtr[tokenIdx * params.hiddenDim + hiddenIdx] = static_cast(data); + } + } +} + +constexpr static int FINALIZE_THREADS_PER_BLOCK = 256; + +__device__ float4 vectorizedLoadPtx(float4 const* ptr) { + float4 ret; + asm volatile("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) + : "l"(ptr)); + return ret; +} + +// Final kernel to unpermute and scale +// This kernel unpermutes the original data, does the k-way reduction and performs the final skip +// connection. +//////////////////////////////////////////////////////////////////////////////////////////////////// + +constexpr int MaxTopK = 64; + +typedef struct __CUDA_ALIGN__(4) { + cutlass::bfloat16_t array[2]; +} bfloat16_2; + +typedef struct __CUDA_ALIGN__(8) { + cutlass::bfloat16_t array[4]; +} bfloat16_4; + +typedef struct __CUDA_ALIGN__(8) { + half array[4]; +} half_4; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct ScaleTraitsStruct; + +template <> +struct ScaleTraitsStruct<1, cutlass::bfloat16_t> { + using PackedType = cutlass::bfloat16_t; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<2, cutlass::bfloat16_t> { + using PackedType = bfloat16_2; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<4, cutlass::bfloat16_t> { + using PackedType = bfloat16_4; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<1, float> { + using PackedType = float; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<2, float> { + using PackedType = float2; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<4, float> { + using PackedType = float4; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<1, half> { + using PackedType = half; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<2, half> { + using PackedType = half2; + using ArrayType = cutlass::Array; +}; + +template <> +struct ScaleTraitsStruct<4, half> { + using PackedType = half_4; + using ArrayType = cutlass::Array; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct FinalizeTraits; + +template +struct FinalizeTraits<1, TypeExpW_> { + using IdxPackedType = int; + using IdxArrayType = cutlass::Array; + using ScaleTraits = ScaleTraitsStruct<1, TypeExpW_>; + using ScalePackedType = typename ScaleTraits::PackedType; + using ScaleArrayType = typename ScaleTraits::ArrayType; +}; + +template +struct FinalizeTraits<2, TypeExpW_> { + using IdxPackedType = int2; + using IdxArrayType = cutlass::Array; + using ScaleTraits = ScaleTraitsStruct<2, TypeExpW_>; + using ScalePackedType = typename ScaleTraits::PackedType; + using ScaleArrayType = typename ScaleTraits::ArrayType; +}; + +template +struct FinalizeTraits<4, TypeExpW_> { + using IdxPackedType = int4; + using IdxArrayType = cutlass::Array; + using ScaleTraits = ScaleTraitsStruct<4, TypeExpW_>; + using ScalePackedType = typename ScaleTraits::PackedType; + using ScaleArrayType = typename ScaleTraits::ArrayType; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void finalizeKernelVecLoad(KernelParams params) { + using Type = typename KernelParams::Type; + using TypeExpW = typename KernelParams::TypeExpW; + int constexpr TopKUnrollFactor = KernelParams::TopKUnrollFactor; + + static_assert( + TopKUnrollFactor == 1 || TopKUnrollFactor == 2 || TopKUnrollFactor == 4, "TopKUnrollFactor must be 1, 2, or 4"); + using FinalizeTraits = FinalizeTraits; + using IdxPackedType = typename FinalizeTraits::IdxPackedType; + using IdxArrayType = typename FinalizeTraits::IdxArrayType; + using ScalePackedType = typename FinalizeTraits::ScalePackedType; + using ScaleArrayType = typename FinalizeTraits::ScaleArrayType; + + int const hiddenDimPaddedBits = params.hiddenDimPadded * cutlass::sizeof_bits::value; + int const hiddenDimBits = params.hiddenDim * cutlass::sizeof_bits::value; + assert(hiddenDimPaddedBits % 128 == 0); + assert(hiddenDimBits % 128 == 0); + + // Load 128-bits per thread, according to the smallest data type we read/write + constexpr int64_t FINALIZE_ELEM_PER_THREAD = 128 / cutlass::sizeof_bits::value; + using InputElem = cutlass::Array; + using OutputElem = cutlass::Array; + using ComputeElem = cutlass::Array; + + int64_t const tokenIdx = blockIdx.x; + int64_t const startOffset = threadIdx.x; + int64_t const stride = FINALIZE_THREADS_PER_BLOCK; + int64_t const numElemsInPaddedCol = params.hiddenDimPadded / FINALIZE_ELEM_PER_THREAD; + int64_t const numElemsInCol = params.hiddenDim / FINALIZE_ELEM_PER_THREAD; + bool const useScale = params.expertWeightsPtr != nullptr; + + __shared__ ScalePackedType scaleArrSmem[MaxTopK / TopKUnrollFactor]; + __shared__ IdxPackedType permutedIdxArrSmem[MaxTopK / TopKUnrollFactor]; + + for (int kChunkIdx = threadIdx.x; kChunkIdx < params.topK / TopKUnrollFactor; kChunkIdx += blockDim.x) { + int const expandedIdx = tokenIdx * params.topK + kChunkIdx * TopKUnrollFactor; + auto permutedIdxPacked = + reinterpret_cast(params.expandedIdxToPermutedIdx)[expandedIdx / TopKUnrollFactor]; + auto scalePacked = + useScale ? reinterpret_cast(params.expertWeightsPtr)[expandedIdx / TopKUnrollFactor] + : ScalePackedType{TypeExpW(1.f)}; + + scaleArrSmem[kChunkIdx] = scalePacked; + permutedIdxArrSmem[kChunkIdx] = permutedIdxPacked; + } + + auto const offset = tokenIdx * params.hiddenDim; + Type* outputPtr = params.outPtr + offset; + auto* outElemPtr = reinterpret_cast(outputPtr); + auto const* inElemPtr = reinterpret_cast(params.inPtr); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // wait on primary kernel when using PDL + if constexpr (KernelParams::UsePdl) { + cudaGridDependencySynchronize(); + } +#endif + __syncthreads(); + + for (int elemIndex = startOffset; elemIndex < numElemsInCol; elemIndex += stride) { + ComputeElem threadOutput; + threadOutput.fill(0); + for (int kChunkIdx = 0; kChunkIdx < params.topK / TopKUnrollFactor; kChunkIdx++) { + auto permutedIdxArr = *reinterpret_cast(&permutedIdxArrSmem[kChunkIdx]); + InputElem inputElemArr[TopKUnrollFactor]; +#pragma unroll + for (int ki = 0; ki < TopKUnrollFactor; ++ki) { + auto const permutedIdx = permutedIdxArr[ki]; + if (permutedIdx == -1) { + continue; + } + + auto const* inputPermutedPtr = inElemPtr + permutedIdx * numElemsInPaddedCol; + + float4 input = vectorizedLoadPtx(reinterpret_cast(&inputPermutedPtr[elemIndex])); + inputElemArr[ki] = *reinterpret_cast(&input); + } + auto scaleArr = *reinterpret_cast(&scaleArrSmem[kChunkIdx]); + auto const scaleFloatArr = arrayConvert>(scaleArr); + +#pragma unroll + for (int ki = 0; ki < TopKUnrollFactor; ++ki) { + auto const permutedIdx = permutedIdxArr[ki]; + if (permutedIdx == -1) { + continue; + } + auto scale = useScale ? scaleFloatArr[ki] : 1.0f; + ComputeElem expertResult = arrayConvert(inputElemArr[ki]); + threadOutput = threadOutput + scale * expertResult; + } + } + OutputElem outputElem = arrayConvert(threadOutput); + outElemPtr[elemIndex] = outputElem; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void finalizeDeepSeekKernel(KernelParams params) { + using Type = typename KernelParams::Type; + using BlockReduce = cub::BlockReduce; + + __shared__ float s_scaleOut; + __shared__ typename BlockReduce::TempStorage temp_storage; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // wait on primary kernel when using PDL + if constexpr (KernelParams::UsePdl) { + cudaGridDependencySynchronize(); + } +#endif + + for (int tokenIdx = blockIdx.y; tokenIdx < params.numTokens; tokenIdx += gridDim.y) { + // Loop over hidden dim + for (int hiddenIdx = threadIdx.x + blockDim.x * blockIdx.x; hiddenIdx < params.hiddenDim; + hiddenIdx += blockDim.x * gridDim.x) { + // Accumulate chunk of token into registers + float acc = 0.0f; + + for (int k = 0; k < params.topK; k++) { + int const expandedIdx = tokenIdx * params.topK + k; + int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; + if (permutedIdx == -1) { + continue; + } + int const totalNumPaddedTokens = params.totalNumPaddedTokens[0]; + int const scaleIdx = permutedIdx + totalNumPaddedTokens * (hiddenIdx / 128); + float const blockScale = params.inDqSfsPtr ? params.inDqSfsPtr[scaleIdx] : 1; + + float const expertProb = (float)params.expertWeightsPtr[tokenIdx * params.topK + k]; + + float const scale = expertProb * blockScale; + acc += scale * static_cast(params.inPtr[permutedIdx * params.hiddenDimPadded + hiddenIdx]); + } + + // The largest (finite) value that can be represented using E4m3. + float constexpr E4m3MaxVal{448.f}; + + // Compute the absolute max + float aMax = BlockReduce(temp_storage).Reduce(fabsf(acc), cuda::maximum<>{}); + + if (threadIdx.x == 0) { + if (params.outDqSfsPtr) { + s_scaleOut = aMax / E4m3MaxVal; + int const scaleOut_idx = tokenIdx + hiddenIdx / 128 * params.numTokens; + params.outDqSfsPtr[scaleOut_idx] = aMax / E4m3MaxVal; + } else { + s_scaleOut = 1.0f; + } + } + __syncthreads(); + float const scaleOut = s_scaleOut; + __syncthreads(); + params.outPtr[tokenIdx * params.hiddenDim + hiddenIdx] = (Type)(acc / scaleOut); + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +void run(Data const& data, void* stream) { + if (data.mUseDeepSeekFp8) { + int const numThreads = 128; + int const numBlocksX = (data.hiddenDim - 1 + numThreads) / numThreads; + // Capped at rather arbitrary 8192 to avoid gridDim exceeding 65535 specified by CUDA. + int const numBlocksY = std::min(8192, data.numTokens); + dim3 numBlocks(numBlocksX, numBlocksY); + + LAUNCH_TOPK_EXPW(data, finalizeDeepSeekKernel, numBlocks, numThreads, 0, stream); + } else { + int const numThreads = 256; + int const numBlocksX = (data.hiddenDim - 1 + numThreads) / numThreads; + // Capped at rather arbitrary 8192 to avoid gridDim exceeding 65535 specified by CUDA. + int const numBlocksY = std::min(8192, data.numTokens); + + if (numBlocksX * numBlocksY < 1184) { + // The number 1184 comes from 148 * 8, where 148 is the number of SMs (Streaming + // Multiprocessors) in the Blackwell architecture, and the value 8 means that each Streaming + // Multiprocessor (SM) can hold up to 8 blocks for this kernel. This limitation is intended to + // ensure that when the number of waves is greater than 1, we choose to use the kernel with + // vectorized loading. + dim3 numBlocks(numBlocksX, numBlocksY); + LAUNCH_TOPK_EXPW(data, finalizeKernel, numBlocks, numThreads, 0, stream); + } else { + FLASHINFER_CHECK( + data.topK <= MaxTopK, + "Finalize kernel with vectorized loading is not supported for this TopK value: %d", + data.topK); + LAUNCH_TOPK_EXPW( + data, + finalizeKernelVecLoad, + /*numBlocks=*/data.numTokens, + /*numThreads=*/FINALIZE_THREADS_PER_BLOCK, + 0, + stream); + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace finalize + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace moe::dev diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh new file mode 100644 index 000000000..7b4bc0de7 --- /dev/null +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh @@ -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 + +#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh" +#include +#include +#include +#include +#include + +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 +__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; + using PackedFp4Type = std::conditional_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(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(readRow) * num_vecs_per_row + vecIdx]; + std::remove_reference_t 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(__hmax(a.x, a.y))); + } + + using BlockReduce = cub::BlockReduce; + __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(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(weightOutput)[static_cast(writeRow) * num_vecs_per_row + vecIdx] = + fp4Vals; + + int64_t sfOffset; + if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) { + sfOffset = static_cast(writeRow) * num_sf_vecs_per_row + vecIdx; + } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) { + sfOffset = tk::get_sf_out_offset_128x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row); + } 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 +__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(expandedIdx / topK); + fused_quant_one_row( + 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 +__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; + using PackedFp4Type = std::conditional_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(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(token) * num_vecs_per_row + vecIdx]; + std::remove_reference_t 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(__hmax(a.x, a.y))); + } + using BlockReduce = cub::BlockReduce; + __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(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(weightOutput)[static_cast(writeRow) * num_vecs_per_row + vecIdx] = + fp4Vals; + int64_t sfOffset; + if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) { + sfOffset = static_cast(writeRow) * num_sf_vecs_per_row + vecIdx; + } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) { + sfOffset = tk::get_sf_out_offset_128x4(std::nullopt, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row); + } 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 +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<<>>( + numTokens, + n, + topK, + numRowsSf, + input, + globalScaleInv, + expandedIdxToPermutedIdx, + weightOutput, + scaleOutput, + perTokenScaleOutput); + } else { + dim3 const grid(numTokens * topK); + fusedPermuteNvfp4QuantKernel<<>>( + 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{}); + } +} + +} // namespace sgl_fused_permute_quant diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu new file mode 100644 index 000000000..a3006324d --- /dev/null +++ b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -0,0 +1,3920 @@ +/* + * 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. + */ +#include + +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h" +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" +#include "flashinfer/trtllm/fused_moe/DevKernel.h" +#include "flashinfer/trtllm/fused_moe/RoutingKernel.h" +#include "flashinfer/trtllm/fused_moe/runner.h" +#include "fused_activation_quant.cuh" +#include "fused_permute_quant.cuh" // fused permute+nvfp4-quant (gate_up de-pad), used by bench_fused_permute_quant +#include "nv_internal/tensorrt_llm/kernels/quantization.h" +#include "nv_internal/tensorrt_llm/thop/utils.h" +#include "tvm_ffi_utils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flashinfer { + +namespace btg = batchedGemm::trtllm::gen; +using tensorrt_llm::kernels::trtllmgen_moe::MoE::ActivationType; +using tensorrt_llm::kernels::trtllmgen_moe::Routing::RoutingMethodType; +using tvm::ffi::Array; +using tvm::ffi::Optional; + +// Validate routing_replay_out tensor properties. +// NOTE: dim0 >= num_tokens is intentionally NOT checked — with CUDA graphs the buffer +// is pre-allocated at maximum batch size and reused across steps with varying num_tokens. +static void validate_routing_replay_out(TensorView const& replay, TensorView const& hidden_states, int64_t top_k) { + TVM_FFI_ICHECK(replay.device().device_type == kDLCUDA) << "routing_replay_out must be a CUDA tensor"; + TVM_FFI_ICHECK(replay.device().device_id == hidden_states.device().device_id) + << "routing_replay_out must be on the same device as hidden_states"; + TVM_FFI_ICHECK(replay.ndim() == 2) << "routing_replay_out must be 2D [num_tokens, top_k]"; + TVM_FFI_ICHECK(replay.size(1) == top_k) << "routing_replay_out dim1 must equal top_k"; + TVM_FFI_ICHECK((replay.dtype() == DLDataType{kDLInt, 16, 1})) << "routing_replay_out must be int16 dtype"; + TVM_FFI_ICHECK(replay.IsContiguous()) << "routing_replay_out must be contiguous (packed row-major)"; +} + +enum class Fp8QuantizationType { + NoneFp8, + DeepSeekFp8, + MxFp8, + PerTensorFp8, +}; + +inline std::string fp8QuantizationTypeToString(Fp8QuantizationType quantization_type) { + switch (quantization_type) { + default: + case Fp8QuantizationType::NoneFp8: + return "NoneFp8"; + case Fp8QuantizationType::DeepSeekFp8: + return "DeepSeekFp8"; + case Fp8QuantizationType::MxFp8: + return "MxFp8"; + case Fp8QuantizationType::PerTensorFp8: + return "PerTensorFp8"; + } +} + +inline ActivationType validateAndCastActivationType(int64_t act_type) { + TVM_FFI_ICHECK(act_type >= 0 && act_type < static_cast(ActivationType::InvalidType)) + << "Invalid activation type: " << act_type; + return static_cast(act_type); +} + +// Utility function to compute the next power of two +inline int32_t nextPowerOfTwo(float value) { + int32_t n = static_cast(std::ceil(value)); + if (n <= 1) return 1; + + // If n is already a power of 2, return it + if ((n & (n - 1)) == 0) return n; + + // Find the next power of 2 + n--; + n |= n >> 1; + n |= n >> 2; + n |= n >> 4; + n |= n >> 8; + n |= n >> 16; + n++; + + return n; +} + +std::set computeSelectedTileN( + std::vector const& supported_tile_nums, + int64_t const num_tokens, + int64_t const top_k, + int64_t const num_local_experts) { + TVM_FFI_ICHECK(!supported_tile_nums.empty()) << "supported_tile_nums must not be empty."; + float const avg_tokens_per_expert = static_cast(num_tokens * top_k) / num_local_experts; + // NOTE: This differs from Python AutoTuner bucketing: + // - AutoTuner maps raw num_tokens with last_positive_power_of_2 (round-down). + // - Here we map derived avg_tokens_per_expert and use nextPowerOfTwo (round-up). + // Because they round different quantities in different directions, cache bucket and runtime + // tile candidates can diverge; launcher-side tactic resolution handles that mismatch. + // assume supported_tile_nums is sorted + int32_t tile_tokens_dim = + std::clamp(nextPowerOfTwo(avg_tokens_per_expert), supported_tile_nums.front(), supported_tile_nums.back()); + auto it = std::find(supported_tile_nums.begin(), supported_tile_nums.end(), tile_tokens_dim); + FLASHINFER_CHECK( + it != supported_tile_nums.end(), + "computeSelectedTileN expected exact tile ", + tile_tokens_dim, + " in supported_tile_nums (size=", + supported_tile_nums.size(), + "). Please keep supported_tile_nums as a dense power-of-2 ladder for this launcher."); + + // Candidate tile set centered on the heuristic tile. + // This function returns nearby candidates (not a single final tile): + // center, +1, +2, and -1 neighbors when available. + // Final tile choice is made later (autotuner-provided tile if valid, otherwise fallback policy). + std::set selected_tile_nums; + selected_tile_nums.insert(tile_tokens_dim); + if (std::next(it) != supported_tile_nums.end()) { + selected_tile_nums.insert(*std::next(it)); + if (std::next(std::next(it)) != supported_tile_nums.end()) { + selected_tile_nums.insert(*std::next(std::next(it))); + } + } + if (it != supported_tile_nums.begin()) { + selected_tile_nums.insert(*std::prev(it)); + } + + return selected_tile_nums; +} + +int64_t selectDefaultTileN( + std::vector const& supported_tile_nums, + int64_t const num_tokens, + int64_t const top_k, + int64_t const num_local_experts) { + auto selected = computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); + TVM_FFI_ICHECK(!selected.empty()) << "No selected tile_N candidates for current MoE input."; + return *selected.begin(); +} + +// Resolve the (tile_N, config) pair passed from Python side, applying fallback logic +// when tile_N is -1. +std::pair resolveMoeTileAndConfig( + Array const& config_index, + std::vector const& supported_tile_nums, + int64_t const num_tokens, + int64_t const top_k, + int64_t const num_local_experts) { + // Python side convention: tactic is [tile_N, config] + TVM_FFI_ICHECK(config_index.size() == 2) + << "Invalid tactic, expected to be [tile_N, config], but got array of size " << config_index.size(); + const int64_t tile_N = config_index[0]; + const int64_t config = config_index[1]; + + if (tile_N == -1 || config == -1) { + // Use fallback tactic + auto const default_tile_N = selectDefaultTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); + return {default_tile_N, -1}; + } + + return {tile_N, config}; +} + +class FusedMoeLauncher { + protected: + Optional routing_logits; + Optional routing_bias; + TensorView hidden_states; + TensorView gemm1_weights; + Optional output1_scales_scalar; + Optional output1_scales_gate_scalar; + TensorView gemm2_weights; + Optional output2_scales_scalar; + Optional per_token_scales; + Tensor per_token_scales_fc2; + + int64_t tile_tokens_dim{}; + int64_t routing_method_type{}; + bool use_shuffled_weight{}; + batchedGemm::gemm::MatrixLayout weight_layout{batchedGemm::gemm::MatrixLayout::MajorK}; + + std::tuple device_version; + std::unique_ptr args; + tensorrt_llm::kernels::trtllmgen_moe::MoE::MoEWorkspace workspace; + + btg::Dtype mDtypeAct{btg::Dtype::Bfloat16}; + btg::Dtype mDtypeWeights{btg::Dtype::Bfloat16}; + btg::Dtype mRoutingBiasDtype{btg::Dtype::Bfloat16}; // Dtype for expert weights in routing, based on routing bias + btg::Dtype mRoutingLogitsDtype{btg::Dtype::Bfloat16}; + bool norm_topk_prob{true}; + ActivationType activation_type{ActivationType::Swiglu}; + btg::Dtype mDtypeScore{btg::Dtype::Bfloat16}; + + // Optional routing replay output: [num_tokens, top_k] int16 tensor + Optional routing_replay_out; + + int64_t intermediate_size_factor{2}; + + public: + // Constructor that initializes all TensorView members + FusedMoeLauncher( + const Optional& routing_logits, + const Optional& routing_bias, + const TensorView& hidden_states, + const TensorView& gemm1_weights, + const Optional& output1_scales_scalar, + const Optional& output1_scales_gate_scalar, + const TensorView& gemm2_weights, + const Optional& output2_scales_scalar, + const Optional& per_token_scales) + : routing_logits(routing_logits), + routing_bias(routing_bias), + hidden_states(hidden_states), + gemm1_weights(gemm1_weights), + output1_scales_scalar(output1_scales_scalar), + output1_scales_gate_scalar(output1_scales_gate_scalar), + gemm2_weights(gemm2_weights), + output2_scales_scalar(output2_scales_scalar), + per_token_scales(per_token_scales), + tile_tokens_dim{}, + routing_method_type{}, + use_shuffled_weight{}, + weight_layout{batchedGemm::gemm::MatrixLayout::MajorK}, + mDtypeAct{btg::Dtype::Bfloat16}, + mDtypeWeights{btg::Dtype::Bfloat16}, + activation_type{ActivationType::Swiglu}, + intermediate_size_factor{2} {} + + public: + void set_routing_replay_out(const Optional& replay_out) { + routing_replay_out = replay_out; + } + + protected: + // Initialize common data necessary for later. + // May throw exception from TVM_FFI_ICHECK. + void init_common( + std::unique_ptr&& args, + int64_t tile_tokens_dim, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + ActivationType activation_type, + bool norm_topk_prob = true); + + // Routing logits [num_tokens, num_experts] + void check_routing_logits() const { + if (routing_logits.has_value()) { + // Check shape + TVM_FFI_ICHECK_EQ(routing_logits.value().ndim(), 2) << "routing_logits must be 2D."; + TVM_FFI_ICHECK_EQ(routing_logits.value().size(0), hidden_states.size(0)) + << "routing_logits and hidden_states must have the same number of tokens."; + TVM_FFI_ICHECK_EQ(routing_logits.value().size(1), args->num_experts) + << "routing_logits dim1 must match num_experts."; + + // Check dtype + TVM_FFI_ICHECK(routing_logits.value().dtype() == dl_float32 || routing_logits.value().dtype() == dl_bfloat16) + << "routing_logits must be float or bfloat16."; + } + } + + // Routing bias [num_experts] + void check_routing_bias_shape() const { + if (routing_bias.has_value()) { + TVM_FFI_ICHECK_EQ(routing_bias.value().ndim(), 1) << "routing_bias must be 1D."; + TVM_FFI_ICHECK_EQ(routing_bias.value().size(0), args->num_experts) << "routing_bias has incorrect shape."; + } + } + + // Hidden states [num_tokens, hidden_size] + void check_hidden_states_shape() const { + TVM_FFI_ICHECK_EQ(hidden_states.ndim(), 2) << "hidden_states must be 2D."; + TVM_FFI_ICHECK_EQ(hidden_states.size(1), args->intermediate_size) << "hidden_states has incorrect shape."; + } + + // GEMM1 or GEMM2 weights [num_experts, M, K] or [num_experts, K/block_k, M, block_k] + void check_weights_shape(std::string which_weights) const { + TensorView weights = (which_weights == "gemm1") ? gemm1_weights : gemm2_weights; + if (which_weights != "gemm1" && which_weights != "gemm2") { + TVM_FFI_LOG_AND_THROW(InternalError) << "Internal error: which_weights = " << which_weights; + } + + int64_t Mn = 0, K = 0; + if (weight_layout == batchedGemm::gemm::MatrixLayout::MajorK) { + // MajorK [num_experts, M, K] + Mn = weights.size(1); + K = weights.size(2); + } else if (weight_layout == batchedGemm::gemm::MatrixLayout::BlockMajorK) { + // BlockMajorK [num_experts, K/block_k, M, block_k] + Mn = weights.size(2); + int64_t block_k = weights.size(3); + K = weights.size(1) * block_k; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported weight_layout: " << (int)weight_layout; + } + if (which_weights == "gemm1") { + // Gated MoE activations (e.g. Swiglu/Geglu) pack gate+up projections in GEMM1, + // so Mn = 2 * intermediate_size and must be even. + if (intermediate_size_factor == 2) { + TVM_FFI_ICHECK_EQ(Mn % 2, 0) << which_weights << " weights Mn dimension must be even."; + } + // Non-gated activations (e.g. Relu2) use a single projection in GEMM1, + // so Mn = intermediate_size. This check covers both gated and non-gated cases. + TVM_FFI_ICHECK_EQ(args->intermediate_size * intermediate_size_factor, Mn) + << "intermediate_size has incorrect shape."; + TVM_FFI_ICHECK_EQ(K, hidden_states.size(1)) + << which_weights << " weights K dimension must be equal to hidden_size."; + } else if (which_weights == "gemm2") { + // GEMM2 always consumes the post-activation hidden of size intermediate_size. + TVM_FFI_ICHECK_EQ(K, args->intermediate_size) + << which_weights << " weights K dimension must be equal to intermediate_size."; + } + } + + void check_routing_common() const { + TVM_FFI_ICHECK(args->top_k > 0 && args->top_k <= args->num_experts) << "top_k must be between 1 and num_experts"; + TVM_FFI_ICHECK(args->local_num_experts > 0 && args->local_num_experts <= args->num_experts) + << "local_num_experts must be between 1 and num_experts"; + TVM_FFI_ICHECK( + args->local_expert_offset >= 0 && args->local_expert_offset + args->local_num_experts <= args->num_experts) + << "expert offset and count must be within valid range"; + + check_routing_logits(); + + if (routing_bias.has_value()) { + check_routing_bias_shape(); + } + } + + // Routing phase workspace tensors (allocated in prepare_routing() or prepare_routing_common()) + Tensor num_tokens_per_expert; + Tensor total_num_padded_tokens; + Tensor expanded_idx_to_permuted_idx; + Tensor permuted_idx_to_token_idx; + Tensor expert_weights; + Tensor expert_indexes; + Tensor expert_count_histogram; + Tensor cta_idx_xy_to_batch_idx; + Tensor cta_idx_xy_to_mn_limit; + Tensor num_non_exiting_ctas; + + void prepare_routing_common() { + // Allocate routing phase workspace tensors + num_tokens_per_expert = alloc_tensor({args->num_experts}, dl_int32, hidden_states.device()); + int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxPermutedPaddedCount( + args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); + + total_num_padded_tokens = alloc_tensor({1}, dl_int32, hidden_states.device()); + + expanded_idx_to_permuted_idx = alloc_tensor({args->num_tokens * args->top_k}, dl_int32, hidden_states.device()); + + permuted_idx_to_token_idx = alloc_tensor({max_num_padded_tokens}, dl_int32, hidden_states.device()); + + expert_indexes = alloc_tensor({args->num_tokens, args->top_k}, dl_int32, hidden_states.device()); + + // expert_weights allocation should be done by derived class since data type could vary + + int64_t const size_of_expert_count_histogram = std::max(args->num_experts * 2, 256 * 2); + expert_count_histogram = alloc_tensor( + {size_of_expert_count_histogram}, + dl_int32, // 256 is the max number of threads per block + // and max number of experts + hidden_states.device()); + + int32_t max_num_ctas = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxNumCtasInBatchDim( + args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); + + cta_idx_xy_to_batch_idx = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); + + cta_idx_xy_to_mn_limit = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); + + num_non_exiting_ctas = alloc_tensor({1}, dl_int32, hidden_states.device()); + + workspace.total_num_padded_tokens = static_cast(total_num_padded_tokens.data_ptr()); + workspace.total_max_padded_tokens = max_num_padded_tokens; + workspace.ProjUpTileN = tile_tokens_dim; + workspace.routing_expert_indexes = static_cast(expert_indexes.data_ptr()); + workspace.permuted_idx_size = static_cast(total_num_padded_tokens.data_ptr()); + workspace.expanded_idx_to_permuted_idx = static_cast(expanded_idx_to_permuted_idx.data_ptr()); + workspace.permuted_idx_to_token_idx = static_cast(permuted_idx_to_token_idx.data_ptr()); + // workspace.expert_weights will be set by derived class after expert_weights allocation + workspace.cta_idx_xy_to_batch_idx = static_cast(cta_idx_xy_to_batch_idx.data_ptr()); + workspace.cta_idx_xy_to_mn_limit = static_cast(cta_idx_xy_to_mn_limit.data_ptr()); + workspace.num_non_exiting_ctas = static_cast(num_non_exiting_ctas.data_ptr()); + + // Set dtype of score based on actual routing_logits dtype + if (routing_logits.has_value()) { + if (routing_logits.value().dtype() == dl_float32) { + mDtypeScore = btg::Dtype::Fp32; + } else { + mDtypeScore = btg::Dtype::Bfloat16; + } + } + } + + void check_moe_common() const { + // Hidden states [num_tokens, hidden_size] + TVM_FFI_ICHECK_EQ(hidden_states.ndim(), 2) << "hidden_states must be 2D."; + } + + // MoE computation phase workspace tensors (allocated in prepare_moe() or prepare_moe_common()) + Tensor gemm1_output; + Tensor activation_output; + Tensor gemm2_output; + Tensor workspace_fc1; + Tensor workspace_fc2; + Tensor output; + int64_t moe_tactic{-1}; + std::unique_ptr moe_runner; + + void prepare_moe_common(int64_t& moe_tactic) { + using RunnerType = tensorrt_llm::kernels::trtllmgen_moe::MoE::Runner; + bool usePerTokenScalingGemm1 = + per_token_scales.has_value() || + static_cast(this->routing_method_type) == RoutingMethodType::Llama4; + bool usePerTokenScalingGemm2 = per_token_scales.has_value() && this->mDtypeAct != btg::Dtype::Bfloat16; + // For FP8 block-scale (E4m3 activations, E4m3 weights) with DeepSeek FP8, use the + // weights-only Runner constructor to match the original kernel path and numerics. + if (this->mDtypeAct == btg::Dtype::E4m3 && this->mDtypeWeights == btg::Dtype::E4m3 && args->mUseDeepSeekFp8) { + moe_runner = std::make_unique( + this->mDtypeWeights, + args->mUseDeepSeekFp8, + (int32_t)tile_tokens_dim, + this->use_shuffled_weight, + this->weight_layout, + usePerTokenScalingGemm1, + usePerTokenScalingGemm2, + false, + false); + } else { + moe_runner = std::make_unique( + this->mDtypeAct, + this->mDtypeWeights, + args->mUseDeepSeekFp8, + (int32_t)tile_tokens_dim, + this->activation_type, + this->use_shuffled_weight, + this->weight_layout, + usePerTokenScalingGemm1, + usePerTokenScalingGemm2); + } + + if (moe_tactic == -1) { + moe_tactic = moe_runner->getDefaultValidConfigIndex( + args->top_k, args->hidden_size, args->intermediate_size, args->local_num_experts, args->num_tokens); + } + auto valid_cfgs = moe_runner->getValidConfigIndices( + args->top_k, args->hidden_size, args->intermediate_size, args->local_num_experts, args->num_tokens); + auto valid_it = std::find(valid_cfgs.begin(), valid_cfgs.end(), moe_tactic); + FLASHINFER_CHECK( + valid_it != valid_cfgs.end(), + "Invalid MoE tactic ", + moe_tactic, + " for tile_N=", + tile_tokens_dim, + ". Number of valid tactics for this tile is ", + valid_cfgs.size(), + ". This often indicates a stale or mismatched autotuner cache entry."); + this->moe_tactic = moe_tactic; + + auto workspace_sizes = moe_runner->getWorkspaceSizeInBytes(*args, moe_tactic); + workspace_fc1 = alloc_tensor({std::get<0>(workspace_sizes)}, dl_int8, hidden_states.device()); + workspace_fc2 = alloc_tensor({std::get<1>(workspace_sizes)}, dl_int8, hidden_states.device()); + workspace.bmm1_workspace = workspace_fc1.data_ptr(); + workspace.bmm2_workspace = workspace_fc2.data_ptr(); + } + + public: + virtual void check_routing() const = 0; + virtual void prepare_routing() = 0; + virtual void check_moe() const = 0; + virtual void prepare_moe(int64_t& moe_tactic) = 0; + + // Main entry point for all the executions. + // Do initializations prior to calling this as the initializations are different for bf16, fp8 and + // fp4. The executions are non-blocking by default. + virtual Array + run(int64_t moe_tactic, + bool enable_pdl = true, + bool use_routing_scales_on_input = false, + bool use_deep_seek_fp8 = false) { + check_routing(); + prepare_routing(); + + // Execute routing + tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim); + cudaStream_t routing_stream = get_stream(hidden_states.device()); + + int16_t* replay_ptr = nullptr; + if (routing_replay_out.has_value()) { + replay_ptr = reinterpret_cast(routing_replay_out.value().data_ptr()); + } + + routing_runner.run( + args->routing_logits, + args->routing_bias, + args->num_tokens, + args->num_experts, + args->top_k, + args->n_group, + args->topk_group, + args->local_expert_offset, + args->local_num_experts, + args->routed_scaling_factor, + workspace.routing_expert_indexes, + static_cast(expert_count_histogram.data_ptr()), + static_cast(total_num_padded_tokens.data_ptr()), + static_cast(expanded_idx_to_permuted_idx.data_ptr()), + nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, + static_cast(permuted_idx_to_token_idx.data_ptr()), + workspace.expert_weights, + static_cast(num_tokens_per_expert.data_ptr()), + static_cast(cta_idx_xy_to_batch_idx.data_ptr()), + static_cast(cta_idx_xy_to_mn_limit.data_ptr()), + static_cast(num_non_exiting_ctas.data_ptr()), + args->mDtypeElt, + mRoutingBiasDtype, + use_routing_scales_on_input, + use_deep_seek_fp8, + static_cast(routing_method_type), + routing_stream, + mRoutingLogitsDtype, + norm_topk_prob, + replay_ptr); + + check_moe(); + prepare_moe(moe_tactic); + + cudaStream_t moe_stream = get_stream(hidden_states.device()); + moe_runner->run(*args, workspace, hidden_states.device().device_id, moe_stream, moe_tactic, enable_pdl); + + if (args->do_finalize) { + return {output}; + } + return {gemm2_output, FusedMoeLauncher::expert_weights, expanded_idx_to_permuted_idx}; + } +}; + +void FusedMoeLauncher::init_common( + std::unique_ptr&& args, + int64_t tile_tokens_dim, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + ActivationType activation_type, + bool norm_topk_prob) { + // Check devicearchitecture: Blackwell (SM 10.x) required + auto device = hidden_states.device().device_id; + int major = 0, minor = 0; + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device); + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device); + TVM_FFI_ICHECK(major == 10 || major == 12) + << "MoE kernel requires SM 10.x or SM 12.x architecture. Current device has SM " << major << minor; + this->device_version = std::make_tuple(major, minor); + + args->routing_logits = routing_logits.has_value() ? routing_logits.value().data_ptr() : nullptr; + args->routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr; + args->hidden_states = hidden_states.data_ptr(); + args->gemm1_weights = gemm1_weights.data_ptr(); + args->gemm2_weights = gemm2_weights.data_ptr(); + + this->args = std::move(args); + this->tile_tokens_dim = tile_tokens_dim; + this->routing_method_type = routing_method_type; + this->use_shuffled_weight = use_shuffled_weight; + TVM_FFI_ICHECK(0 <= weight_layout && weight_layout <= 2) << "the value of weight_layout is not recognized"; + this->weight_layout = static_cast(weight_layout); + this->activation_type = activation_type; + this->intermediate_size_factor = isGatedActivation(activation_type) ? 2 : 1; + this->norm_topk_prob = norm_topk_prob; +} + +class Bf16MoeLauncher : public FusedMoeLauncher { + public: + static constexpr std::array mSupportedTileNums = {8, 16, 32, 64, 128}; + + Bf16MoeLauncher( + Optional const& routing_logits, + Optional const& routing_bias, + TensorView const& expert_indices, + TensorView const& expert_weights, + TensorView const& hidden_states, + TensorView const& gemm1_weights, + TensorView const& gemm2_weights) + : FusedMoeLauncher( + routing_logits, + routing_bias, + hidden_states, + gemm1_weights, + Optional(), + Optional(), + gemm2_weights, + Optional(), + Optional()), + expert_indices(expert_indices), + expert_weights(expert_weights) {} + + void init( + std::unique_ptr&& args, + int64_t tile_tokens_dim, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + ActivationType activation_type, + bool norm_topk_prob = true) { + // Do base class init and perform common checks + FusedMoeLauncher::init_common( + std::move(args), + tile_tokens_dim, + routing_method_type, + use_shuffled_weight, + weight_layout, + activation_type, + norm_topk_prob); + } + + void check_routing() const override { + FusedMoeLauncher::check_routing_common(); + if (expert_indices.ndim() == 2 && expert_indices.size(0) > 0) { + // Pre-computed routing: expert_indices is a packed tensor + // Format: (expert_id << 16) | (weight_bf16.view(int16)) + TVM_FFI_ICHECK_EQ(expert_indices.ndim(), 2) << "expert_indices must be 2D."; + TVM_FFI_ICHECK_EQ(expert_indices.size(0), hidden_states.size(0)) + << "expert_indices and hidden_states must have same number of tokens."; + TVM_FFI_ICHECK_EQ(expert_indices.size(1), args->top_k) << "expert_indices dim1 must match top_k."; + TVM_FFI_ICHECK_EQ(expert_indices.dtype(), dl_int32) << "expert_indices must be int32."; + } + + // TODO n_group, topk_group validation? + } + + void prepare_routing() override { + FusedMoeLauncher::prepare_routing_common(); + + args->mDtypeElt = btg::Dtype::Bfloat16; + args->mUseDeepSeekFp8 = false; + + // Set expert weights dtype based on routing bias + auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; + mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; + + auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; + mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; + + // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr + bool has_precomputed_indices = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; + if (has_precomputed_indices) { + // Use expert_indices directly + workspace.routing_expert_indexes = static_cast(const_cast(expert_indices.data_ptr())); + } + bool has_precomputed_weights = expert_weights.ndim() == 2 && expert_weights.size(0) > 0; + if (has_precomputed_weights) { + workspace.expert_weights = const_cast(expert_weights.data_ptr()); + } else { + auto ew_dtype = mDtypeScore == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; + FusedMoeLauncher::expert_weights = + alloc_tensor({args->num_tokens, args->top_k}, ew_dtype, hidden_states.device()); + workspace.expert_weights = FusedMoeLauncher::expert_weights.data_ptr(); + } + } + + void check_moe() const override { + FusedMoeLauncher::check_moe_common(); + + TVM_FFI_ICHECK(weight_layout == batchedGemm::gemm::MatrixLayout::BlockMajorK) + << "BF16 Moe: weight_layout must be BlockMajorK"; + check_weights_shape("gemm1"); + check_weights_shape("gemm2"); + + TVM_FFI_ICHECK_EQ(args->intermediate_size % 128, 0) << "the second dimension of weights must be a multiple of 128."; + } + + void prepare_moe(int64_t& moe_tactic) override { + FusedMoeLauncher::prepare_moe_common(moe_tactic); + + int32_t max_num_padded_tokens = workspace.total_max_padded_tokens; + gemm1_output = alloc_tensor({max_num_padded_tokens, args->intermediate_size}, dl_bfloat16, hidden_states.device()); + activation_output = + alloc_tensor({max_num_padded_tokens, args->intermediate_size}, dl_bfloat16, hidden_states.device()); + gemm2_output = alloc_tensor({max_num_padded_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); + + workspace.hidden_states_scale_linear = nullptr; + workspace.gemm1_output = gemm1_output.data_ptr(); + workspace.gemm1_output_scale = nullptr; + workspace.activation_output = activation_output.data_ptr(); + workspace.activation_output_scale = nullptr; + workspace.gemm2_output = gemm2_output.data_ptr(); + workspace.gemm2_output_scale = nullptr; + + if (args->output == nullptr) { + output = alloc_tensor({args->num_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); + args->output = output.data_ptr(); + } + args->output_scale = nullptr; + } + + static Array> getValidConfigs( + int64_t top_k, + int64_t hidden_size, + int64_t intermediate_size, + int64_t num_local_experts, + int64_t num_tokens, + int64_t act_type, + bool use_shuffled_weight, + int64_t weight_layout) { + Array> valid_configs; + + std::vector supported_tile_nums(mSupportedTileNums.begin(), mSupportedTileNums.end()); + std::set selected_tile_nums = + computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); + + for (int32_t tile_N : selected_tile_nums) { + auto moe_runner = std::make_unique( + btg::Dtype::Bfloat16, // dtype_act + btg::Dtype::Bfloat16, // dtype_weights + false, // useDeepSeekFp8 + tile_N, + static_cast(act_type), + use_shuffled_weight, + static_cast(weight_layout)); + + auto cfgs = + moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); + + for (auto cfg : cfgs) { + valid_configs.push_back({tile_N, cfg}); + } + } + + return valid_configs; + } + + private: + TensorView expert_weights; + TensorView expert_indices; +}; + +class Fp8PerTensorLauncher : public FusedMoeLauncher { + public: + static constexpr std::array mSupportedTileNums = {8, 16, 32, 64, 128}; + + // Constructor that passes TensorView parameters to base constructor + Fp8PerTensorLauncher( + TensorView const& routing_logits, + Optional const& routing_bias, + TensorView const& hidden_states, + TensorView const& gemm1_weights, + TensorView const& output1_scales_scalar, + TensorView const& output1_scales_gate_scalar, + TensorView const& gemm2_weights, + TensorView const& output2_scales_scalar) + : FusedMoeLauncher( + Optional(routing_logits), + routing_bias, + hidden_states, + gemm1_weights, + Optional(output1_scales_scalar), + Optional(output1_scales_gate_scalar), + gemm2_weights, + Optional(output2_scales_scalar), + Optional()), + use_routing_scales_on_input(false) {} + + void init( + std::unique_ptr&& args, + int64_t tile_tokens_dim, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + bool use_routing_scales_on_input_param, + ActivationType activation_type, + bool norm_topk_prob = true) { + this->use_routing_scales_on_input = use_routing_scales_on_input_param; + + auto dtype = hidden_states.dtype(); + if (dtype == dl_float16) { + mDtypeAct = btg::Dtype::Fp16; + } else if (dtype == dl_bfloat16) { + mDtypeAct = btg::Dtype::Bfloat16; + } else if (dtype == dl_float8_e4m3fn) { + mDtypeAct = btg::Dtype::E4m3; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for FP8 MoE."; + } + mDtypeWeights = btg::Dtype::E4m3; + + FusedMoeLauncher::init_common( + std::move(args), + tile_tokens_dim, + routing_method_type, + use_shuffled_weight, + weight_layout, + activation_type, + norm_topk_prob); + } + + void check_routing() const override { + FusedMoeLauncher::check_routing_common(); + } + + void prepare_routing() override { + FusedMoeLauncher::prepare_routing_common(); + + auto dtype = hidden_states.dtype(); + if (dtype == dl_float16) { + args->mDtypeElt = btg::Dtype::Fp16; + } else if (dtype == dl_bfloat16) { + args->mDtypeElt = btg::Dtype::Bfloat16; + } else if (dtype == dl_float8_e4m3fn) { + args->mDtypeElt = btg::Dtype::E4m3; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; + } + + args->mDtypeOut = btg::Dtype::Bfloat16; + args->mUseDeepSeekFp8 = false; + + auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; + mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; + + auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; + mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; + + auto expert_weights_dtype = mRoutingLogitsDtype == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; + expert_weights = alloc_tensor({args->num_tokens, args->top_k}, expert_weights_dtype, hidden_states.device()); + + workspace.expert_weights = expert_weights.data_ptr(); + if (static_cast(routing_method_type) == RoutingMethodType::Llama4) { + workspace.token_scales = expert_weights.data_ptr(); // Consumed by permuteGemm1 kernel + } + } + + void check_moe() const override { + FusedMoeLauncher::check_moe_common(); + + TVM_FFI_ICHECK(output1_scales_scalar.has_value()) << "output1_scales_scalar is required for FP8 MoE"; + TVM_FFI_ICHECK_EQ(output1_scales_scalar.value().dtype(), dl_float32) << "output1_scales_scalar must be float."; + TVM_FFI_ICHECK_EQ(output1_scales_scalar.value().ndim(), 1) << "output1_scales_scalar must be 1D."; + TVM_FFI_ICHECK_EQ(output1_scales_scalar.value().size(0), args->local_num_experts) + << "output1_scales_scalar has incorrect dim 0."; + + TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()) << "output1_scales_gate_scalar is required for FP8 MoE"; + TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.value().dtype(), dl_float32) + << "output1_scales_gate_scalar must be float."; + TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.value().ndim(), 1) << "output1_scales_gate_scalar must be 1D."; + TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.value().size(0), args->local_num_experts) + << "output1_scales_gate_scalar has incorrect dim 0."; + + TVM_FFI_ICHECK(output2_scales_scalar.has_value()) << "output2_scales_scalar is required for FP8 MoE"; + TVM_FFI_ICHECK_EQ(output2_scales_scalar.value().dtype(), dl_float32) << "output2_scales_scalar must be float."; + TVM_FFI_ICHECK_EQ(output2_scales_scalar.value().ndim(), 1) << "output2_scales_scalar must be 1D."; + TVM_FFI_ICHECK_EQ(output2_scales_scalar.value().size(0), args->local_num_experts) + << "output2_scales_scalar has incorrect dim 0."; + + TVM_FFI_ICHECK( + hidden_states.dtype() == dl_float8_e4m3fn || hidden_states.dtype() == dl_float16 || + hidden_states.dtype() == dl_bfloat16) + << "FP8 MoE: hidden_states must be float8_e4m3fn, float16, or bfloat16."; + TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm1_weights must be float8_e4m3fn."; + TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm2_weights must be float8_e4m3fn."; + } + + void prepare_moe(int64_t& moe_tactic) override { + FusedMoeLauncher::prepare_moe_common(moe_tactic); + + int32_t max_num_padded_tokens_gemm1 = workspace.total_max_padded_tokens + args->num_experts; + int32_t max_num_padded_tokens_gemm2 = workspace.total_max_padded_tokens; + + gemm1_output = + alloc_tensor({max_num_padded_tokens_gemm1, 2 * args->intermediate_size}, dl_uint8, hidden_states.device()); + gemm1_output_scale = alloc_tensor( + {2 * args->intermediate_size / 128, max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); + + activation_output = + alloc_tensor({max_num_padded_tokens_gemm1, args->intermediate_size}, dl_uint8, hidden_states.device()); + activation_output_scale = + alloc_tensor({args->intermediate_size / 128, max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); + + gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); + + workspace.hidden_states_scale_linear = nullptr; + workspace.gemm1_output = gemm1_output.data_ptr(); + workspace.gemm1_output_scale = static_cast(gemm1_output_scale.data_ptr()); + workspace.activation_output = activation_output.data_ptr(); + workspace.activation_output_scale = static_cast(activation_output_scale.data_ptr()); + workspace.gemm2_output = gemm2_output.data_ptr(); + workspace.gemm2_output_scale = nullptr; + + if (args->output == nullptr) { + output = alloc_tensor({args->num_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); + args->output = output.data_ptr(); + } + args->output_scale = nullptr; + + // Set scale pointers + TVM_FFI_ICHECK(output1_scales_scalar.has_value()); + TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()); + TVM_FFI_ICHECK(output2_scales_scalar.has_value()); + + args->output1_scales_scalar = static_cast(output1_scales_scalar.value().data_ptr()); + args->output1_scales_gate_scalar = static_cast(output1_scales_gate_scalar.value().data_ptr()); + args->output2_scales_scalar = static_cast(output2_scales_scalar.value().data_ptr()); + } + + private: + bool use_routing_scales_on_input; + Tensor gemm1_output_scale; + Tensor activation_output_scale; + + public: + static Array> getValidConfigs( + int64_t top_k, + int64_t hidden_size, + int64_t intermediate_size, + int64_t num_local_experts, + int64_t num_tokens, + int64_t act_type, + bool use_shuffled_weight, + int64_t weight_layout, + btg::Dtype dtype_act, + btg::Dtype dtype_weights) { + Array> valid_configs; + + std::vector supported_tile_nums(mSupportedTileNums.begin(), mSupportedTileNums.end()); + std::set selected_tile_nums = + computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); + + for (int32_t tile_N : selected_tile_nums) { + auto moe_runner = std::make_unique( + dtype_act, + dtype_weights, + false, // useDeepSeekFp8 + tile_N, + static_cast(act_type), + use_shuffled_weight, + static_cast(weight_layout), + true, // usePerTokenScalingGemm1. always true for per-tensor fp8 due to llama4 routing + false, + false, + false); + + auto cfgs = + moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); + + for (auto cfg : cfgs) { + valid_configs.push_back({tile_N, cfg}); + } + } + + return valid_configs; + } +}; + +class Fp8BlockScaleLauncher : public FusedMoeLauncher { + public: + static constexpr std::array mBaseSupportedTileNums = {8, 16, 32, 64, 128}; + + static std::vector getSupportedTileNums(Fp8QuantizationType quantization_type) { + std::vector tiles(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); + if (quantization_type == Fp8QuantizationType::MxFp8) { + tiles.push_back(256); + } + return tiles; + } + + Fp8BlockScaleLauncher( + Optional const& routing_logits, + Optional const& routing_bias, + TensorView const& hidden_states, + TensorView const& hidden_states_scale, + TensorView const& gemm1_weights, + TensorView const& gemm1_weights_scale, + TensorView const& gemm2_weights, + TensorView const& gemm2_weights_scale, + TensorView const& expert_indices, + TensorView const& expert_weights, + Fp8QuantizationType quantization_type, + Optional const& gate_up_lora_delta = Optional(), + Optional const& activation_lora_input = Optional()) + : FusedMoeLauncher( + routing_logits, + routing_bias, + hidden_states, + gemm1_weights, + Optional(), + Optional(), + gemm2_weights, + Optional(), + Optional()), + hidden_states_scale(hidden_states_scale), + gemm1_weights_scale(gemm1_weights_scale), + gemm2_weights_scale(gemm2_weights_scale), + expert_indices(expert_indices), + expert_weights(expert_weights), + gate_up_lora_delta(gate_up_lora_delta), + activation_lora_input(activation_lora_input), + quantization_type(quantization_type) {} + + void init( + std::unique_ptr&& args, + int64_t tile_tokens_dim, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + ActivationType activation_type, + bool norm_topk_prob = true) { + if (quantization_type == Fp8QuantizationType::MxFp8) { + mDtypeAct = btg::Dtype::MxE4m3; + mDtypeWeights = btg::Dtype::MxE4m3; + } else { + mDtypeAct = btg::Dtype::E4m3; + mDtypeWeights = btg::Dtype::E4m3; + } + + auto dtype = hidden_states.dtype(); + if (dtype == dl_float16) { + args->mDtypeElt = btg::Dtype::Fp16; + } else if (dtype == dl_bfloat16) { + args->mDtypeElt = btg::Dtype::Bfloat16; + } else if (dtype == dl_float8_e4m3fn) { + args->mDtypeElt = btg::Dtype::E4m3; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; + } + + // Output is always bfloat16 for FP8 block scale + args->mDtypeOut = btg::Dtype::Bfloat16; + + FusedMoeLauncher::init_common( + std::move(args), + tile_tokens_dim, + routing_method_type, + use_shuffled_weight, + weight_layout, + activation_type, + norm_topk_prob); + } + + void check_routing() const override { + // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr + if (expert_indices.ndim() == 2 && expert_indices.size(0) > 0) { + // Pre-computed routing: expert_indices is a packed tensor + // Format: (expert_id << 16) | (weight_bf16.view(int16)) + TVM_FFI_ICHECK_EQ(expert_indices.ndim(), 2) << "expert_indices must be 2D."; + TVM_FFI_ICHECK_EQ(expert_indices.size(0), hidden_states.size(0)) + << "expert_indices and hidden_states must have same number of tokens."; + TVM_FFI_ICHECK_EQ(expert_indices.size(1), args->top_k) << "expert_indices dim1 must match top_k."; + TVM_FFI_ICHECK_EQ(expert_indices.dtype(), dl_int32) << "expert_indices must be int32."; + } + + FusedMoeLauncher::check_routing_common(); + + if (static_cast(routing_method_type) != RoutingMethodType::DeepSeekV3) { + TVM_FFI_ICHECK(args->n_group <= 1) << "Current routing kernel (no groups) only supports n_group <= 1"; + TVM_FFI_ICHECK(args->topk_group <= 1) << "Current routing kernel (no groups) only supports topk_group <= 1"; + } + + if (static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3) { + TVM_FFI_ICHECK(args->n_group != 0) << "n_group should not be zero for DeepSeekV3 routing"; + TVM_FFI_ICHECK(args->topk_group != 0) << "if n_group is given, topk_group must be given"; + TVM_FFI_ICHECK_EQ(args->num_experts % args->n_group, 0) << "num_experts must be divisible by n_group"; + // DeepSeekV3 routing supports top_k up to: + // - 8 when num_experts <= 384 (NumKimiK2Experts) + // - 22 when num_experts > 384 (NumNemotronExperts path) + // Keep this in sync with LAUNCH_ROUTING_DEEPSEEK in trtllm_fused_moe_routing_deepseek.cu. + constexpr int32_t kNumKimiK2Experts = 384; // same as in trtllm_fused_moe_routing_deepseek.cu + int32_t max_supported_top_k = args->num_experts <= kNumKimiK2Experts ? 8 : 22; + TVM_FFI_ICHECK(args->top_k <= max_supported_top_k && args->top_k > 0) + << "Current routing kernel (with groups) only supports top_k<=" << max_supported_top_k + << " && top_k>0 for num_experts=" << args->num_experts << "."; + TVM_FFI_ICHECK(args->topk_group <= 4 && args->topk_group > 0) + << "Current routing kernel only (with groups) supports topk_group<=4 && topk_group > 0."; + TVM_FFI_ICHECK_LE(args->topk_group, args->n_group) << "n_group must not be smaller than topk_group."; + TVM_FFI_ICHECK_LT(args->top_k, (args->topk_group * args->num_experts / args->n_group)) + << "top_k must be less than total number of experts in selected groups"; + } else if ( + static_cast(routing_method_type) == RoutingMethodType::Renormalize || + static_cast(routing_method_type) == RoutingMethodType::RenormalizeNaive || + static_cast(routing_method_type) == RoutingMethodType::SigmoidRenorm || + static_cast(routing_method_type) == RoutingMethodType::Sigmoid) { + TVM_FFI_ICHECK(args->top_k <= 32 && args->top_k > 0) + << "Current routing kernel (no groups) only supports top_k<=32 && top_k>0."; + } else if (static_cast(routing_method_type) == RoutingMethodType::Llama4) { + TVM_FFI_ICHECK_EQ(args->top_k, 1) << "Current routing kernel (no groups, Llama4) only supports top_k=1."; + } + + TVM_FFI_ICHECK_EQ(args->num_experts % 4, 0) << "Routing kernel expects that num_experts must be divisible by 4"; + TVM_FFI_ICHECK_GT(args->num_experts, args->top_k) << "num_experts must be greater than top_k"; + TVM_FFI_ICHECK_LE(args->local_num_experts + args->local_expert_offset, args->num_experts) + << "num_experts must be greater or equal to local_num_experts + local_expert_offset"; + } + + void prepare_routing() override { + FusedMoeLauncher::prepare_routing_common(); + + auto dtype = hidden_states.dtype(); + if (dtype == dl_float16) { + args->mDtypeElt = btg::Dtype::Fp16; + } else if (dtype == dl_bfloat16) { + args->mDtypeElt = btg::Dtype::Bfloat16; + } else if (dtype == dl_float8_e4m3fn) { + args->mDtypeElt = btg::Dtype::E4m3; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; + } + + args->mUseDeepSeekFp8 = quantization_type == Fp8QuantizationType::DeepSeekFp8; + // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr + bool has_precomputed_indices = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; + if (has_precomputed_indices) { + // Use expert_indices directly + workspace.routing_expert_indexes = static_cast(const_cast(expert_indices.data_ptr())); + } else { + // Use routing_logits directly + args->routing_logits = static_cast(routing_logits.value().data_ptr()); + } + // Set expert weights dtype based on routing bias + auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; + mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; + + auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; + mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; + + // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr + bool has_precomputed_weights = expert_weights.ndim() == 2 && expert_weights.size(0) > 0; + if (!has_precomputed_weights) { + auto ew_dtype = mDtypeScore == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; + FusedMoeLauncher::expert_weights = + alloc_tensor({args->num_tokens, args->top_k}, ew_dtype, hidden_states.device()); + workspace.expert_weights = FusedMoeLauncher::expert_weights.data_ptr(); + } else { + workspace.expert_weights = const_cast(expert_weights.data_ptr()); + } + } + + void check_moe() const override { + FusedMoeLauncher::check_moe_common(); + + TVM_FFI_ICHECK_EQ(hidden_states.dtype(), dl_float8_e4m3fn) << "hidden_states must be fp8."; + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_float32) << "hidden_states_scale must be float."; + TVM_FFI_ICHECK_EQ(hidden_states_scale.ndim(), 2) << "hidden_states_scale must be 2D."; + TVM_FFI_ICHECK_EQ(hidden_states_scale.size(0), hidden_states.size(1) / 128) + << "hidden_states_scale dim0 must match hidden_states dim1 / 128."; + TVM_FFI_ICHECK_EQ(hidden_states_scale.size(1), args->num_tokens) + << "hidden_states_scale dim1 must match num_tokens."; + } else if (quantization_type == Fp8QuantizationType::MxFp8) { + TVM_FFI_CHECK( + weight_layout == batchedGemm::gemm::MatrixLayout::MajorK, "weight_layout must be MajorK for MxFp8."); + TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_uint8); + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; + } + + TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "gemm1_weights must be fp8."; + TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "gemm2_weights must be fp8."; + + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_float32) << "gemm1_weights_scale must be float."; + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.ndim(), 3) << "gemm1_weights_scale must be 3D."; + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.size(0), args->local_num_experts) + << "gemm1_weights_scale has incorrect shape."; + TVM_FFI_ICHECK_EQ(args->intermediate_size % 128, 0) << "intermediate_size must be a multiple of 128."; + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.size(1), intermediate_size_factor * args->intermediate_size / 128) + << "gemm1_weights_scale has incorrect shape."; + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.size(2), args->hidden_size / 128) + << "gemm1_weights_scale has incorrect shape."; + } else if (quantization_type == Fp8QuantizationType::MxFp8) { + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_uint8) << "gemm1_weights_scale must be uint8."; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; + } + + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_float32) << "gemm2_weights_scale must be float."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.ndim(), 3) << "gemm2_weights_scale must be 3D."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.size(0), args->local_num_experts) + << "gemm2_weights_scale has incorrect shape."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.size(1), args->hidden_size / 128) + << "gemm2_weights_scale has incorrect shape."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.size(2), args->intermediate_size / 128) + << "gemm2_weights_scale has incorrect shape."; + } else if (quantization_type == Fp8QuantizationType::MxFp8) { + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_uint8) << "gemm2_weights_scale must be uint8."; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; + } + + check_weights_shape("gemm1"); + check_weights_shape("gemm2"); + + if (gate_up_lora_delta.has_value()) { + TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().dtype(), dl_bfloat16) << "gate_up_lora_delta must be bf16."; + TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().ndim(), 3) + << "gate_up_lora_delta must be [num_tokens, top_k, 2 * intermediate_size]."; + TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().size(0), args->num_tokens); + TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().size(1), args->top_k); + TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().size(2), args->intermediate_size * intermediate_size_factor); + } + if (activation_lora_input.has_value()) { + TVM_FFI_ICHECK_EQ(activation_lora_input.value().dtype(), dl_bfloat16) << "activation_lora_input must be bf16."; + TVM_FFI_ICHECK_EQ(activation_lora_input.value().ndim(), 3) + << "activation_lora_input must be [num_tokens, top_k, intermediate_size]."; + TVM_FFI_ICHECK_EQ(activation_lora_input.value().size(0), args->num_tokens); + TVM_FFI_ICHECK_EQ(activation_lora_input.value().size(1), args->top_k); + TVM_FFI_ICHECK_EQ(activation_lora_input.value().size(2), args->intermediate_size); + } + + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + TVM_FFI_ICHECK_EQ(args->intermediate_size % 128, 0) << "intermediate_size must be a multiple of 128."; + } + } + + void prepare_moe(int64_t& moe_tactic) override { + FusedMoeLauncher::prepare_moe_common(moe_tactic); + + // Calculate max_num_padded_tokens for gemm1 and gemm2 using maybeGetMinTokenCount + int32_t max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + workspace.total_max_padded_tokens, args->intermediate_size, btg::dtypeGetNumBits(args->mDtypeElt)); + int32_t max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + workspace.total_max_padded_tokens, args->hidden_size, btg::dtypeGetNumBits(args->mDtypeOut)); + + gemm1_output = alloc_tensor( + {max_num_padded_tokens_gemm1, intermediate_size_factor * args->intermediate_size}, + dl_uint8, + hidden_states.device()); + + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + gemm1_output_scale = alloc_tensor( + {intermediate_size_factor * args->intermediate_size / 128, workspace.total_max_padded_tokens}, + dl_float32, + hidden_states.device()); + } else if (quantization_type == Fp8QuantizationType::MxFp8) { + // MxFP8 fuses the activation so no need for intermediate_size_factor + int64_t sf_size = + tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens_gemm1, args->intermediate_size / 32); + gemm1_output_scale = alloc_tensor({sf_size}, dl_uint8, hidden_states.device()); + } + + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + activation_output = + alloc_tensor({max_num_padded_tokens_gemm1, args->intermediate_size}, dl_uint8, hidden_states.device()); + activation_output_scale = alloc_tensor( + {args->intermediate_size / 128, max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); + } + + gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); + + workspace.hidden_states_scale_linear = nullptr; + workspace.gemm1_output = gemm1_output.data_ptr(); + workspace.gemm1_output_scale = static_cast(gemm1_output_scale.data_ptr()); + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + workspace.activation_output = activation_output.data_ptr(); + workspace.activation_output_scale = static_cast(activation_output_scale.data_ptr()); + } + workspace.gemm2_output = gemm2_output.data_ptr(); + workspace.gemm2_output_scale = nullptr; + + if (args->output == nullptr) { + output = alloc_tensor({args->num_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); + args->output = output.data_ptr(); + } + args->output_scale = nullptr; + + args->hidden_states_scale = static_cast(hidden_states_scale.data_ptr()); + args->gemm1_weights_scale = static_cast(gemm1_weights_scale.data_ptr()); + args->gemm2_weights_scale = static_cast(gemm2_weights_scale.data_ptr()); + args->gate_up_lora_delta = gate_up_lora_delta.has_value() ? gate_up_lora_delta.value().data_ptr() : nullptr; + args->activation_lora_input = + activation_lora_input.has_value() ? activation_lora_input.value().data_ptr() : nullptr; + } + + private: + TensorView hidden_states_scale; + TensorView gemm1_weights_scale; + TensorView gemm2_weights_scale; + Tensor gemm1_output_scale; + Tensor activation_output_scale; + TensorView expert_indices; + TensorView expert_weights; + Optional gate_up_lora_delta; + Optional activation_lora_input; + Fp8QuantizationType quantization_type; + + public: + // Override to handle pre-computed routing + Array + run(int64_t moe_tactic, + bool enable_pdl = true, + bool use_routing_scales_on_input = false, + bool use_deep_seek_fp8 = false) override { + check_routing(); + prepare_routing(); + + cudaStream_t routing_stream = get_stream(hidden_states.device()); + tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim); + + // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr + bool use_precomputed = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; + // When using pre-computed routing, pass nullptr as routing_logits to tell the + // routing runner to use the pre-computed expert indices from workspace.routing_expert_indexes + int16_t* replay_ptr = nullptr; + if (routing_replay_out.has_value()) { + replay_ptr = reinterpret_cast(routing_replay_out.value().data_ptr()); + } + + routing_runner.run( + use_precomputed ? nullptr : args->routing_logits, + args->routing_bias, + args->num_tokens, + args->num_experts, + args->top_k, + args->n_group, + args->topk_group, + args->local_expert_offset, + args->local_num_experts, + args->routed_scaling_factor, + workspace.routing_expert_indexes, + static_cast(expert_count_histogram.data_ptr()), + static_cast(total_num_padded_tokens.data_ptr()), + static_cast(expanded_idx_to_permuted_idx.data_ptr()), + nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, + static_cast(permuted_idx_to_token_idx.data_ptr()), + workspace.expert_weights, + static_cast(num_tokens_per_expert.data_ptr()), + static_cast(cta_idx_xy_to_batch_idx.data_ptr()), + static_cast(cta_idx_xy_to_mn_limit.data_ptr()), + static_cast(num_non_exiting_ctas.data_ptr()), + args->mDtypeElt, + mRoutingBiasDtype, + use_routing_scales_on_input, + use_deep_seek_fp8, + static_cast(routing_method_type), + routing_stream, + mRoutingLogitsDtype, + norm_topk_prob, + replay_ptr); + + check_moe(); + prepare_moe(moe_tactic); + + cudaStream_t moe_stream = get_stream(hidden_states.device()); + moe_runner->run(*args, workspace, hidden_states.device().device_id, moe_stream, moe_tactic, enable_pdl); + + if (args->do_finalize) { + return {output}; + } + return {gemm2_output, FusedMoeLauncher::expert_weights, expanded_idx_to_permuted_idx}; + } + + static Array> getValidConfigs( + int64_t top_k, + int64_t hidden_size, + int64_t intermediate_size, + int64_t num_local_experts, + int64_t num_tokens, + bool use_shuffled_weight, + int64_t weight_layout, + btg::Dtype dtype_act, + btg::Dtype dtype_weights, + Fp8QuantizationType quantization_type, + int64_t act_type) { + Array> valid_configs; + auto activation_type = validateAndCastActivationType(act_type); + + auto supported_tile_nums = getSupportedTileNums(quantization_type); + std::set selected_tile_nums = + computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); + + for (int32_t tile_N : selected_tile_nums) { + std::unique_ptr moe_runner; + // Keep getValidConfigs constructor path aligned with runtime prepare_moe_common(). + // This branch is for DeepSeek FP8 (E4m3 activations + E4m3 weights). + if (quantization_type == Fp8QuantizationType::DeepSeekFp8 && dtype_act == btg::Dtype::E4m3 && + dtype_weights == btg::Dtype::E4m3) { + TVM_FFI_ICHECK(static_cast(activation_type) == static_cast(ActivationType::Swiglu)) + << "DeepSeekFp8 only supports ActivationType::Swiglu, got " << static_cast(activation_type) << "."; + moe_runner = std::make_unique( + dtype_weights, + true /* useDeepSeekFp8 */, + tile_N, + use_shuffled_weight, + static_cast(weight_layout)); + } else { + // Under current trtllm_get_valid_moe_configs() dispatch rules, this else-path is + // reached only by FP8 block-scale MXFP8 (dtype_act=dtype_weights=MxE4m3). + moe_runner = std::make_unique( + dtype_act, // dtypeAct + dtype_weights, // dtypeWeights + quantization_type == Fp8QuantizationType::DeepSeekFp8, // useDeepSeekFp8 + tile_N, + activation_type, + use_shuffled_weight, + static_cast(weight_layout)); + } + + auto cfgs = + moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); + + for (auto cfg : cfgs) { + valid_configs.push_back({tile_N, cfg}); + } + } + + return valid_configs; + } +}; + +class MxInt4BlockScaleLauncher : public FusedMoeLauncher { + public: + static constexpr std::array mSupportedTileNums = {8, 16, 32, 64, 128}; + + MxInt4BlockScaleLauncher( + TensorView const& routing_logits, + Optional const& routing_bias, + TensorView const& hidden_states, + TensorView const& gemm1_weights, + TensorView const& gemm1_weights_scale, + Optional const& gemm1_alpha, + Optional const& gemm1_beta, + Optional const& gemm1_clamp_limit, + TensorView const& gemm2_weights, + TensorView const& gemm2_weights_scale) + : FusedMoeLauncher( + Optional(routing_logits), + routing_bias, + hidden_states, + gemm1_weights, + Optional(), + Optional(), + gemm2_weights, + Optional(), + Optional()), + gemm1_weights_scale(gemm1_weights_scale), + gemm2_weights_scale(gemm2_weights_scale) {} + + void init( + std::unique_ptr&& args, + int64_t tile_tokens_dim, + int64_t routing_method_type, + bool norm_topk_prob = true) { + // currently only support mxint4 x bf16 + auto dtype = hidden_states.dtype(); + if (dtype == dl_bfloat16) { + args->mDtypeElt = btg::Dtype::Bfloat16; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; + } + args->mDtypeOut = btg::Dtype::Bfloat16; + + mDtypeAct = btg::Dtype::Bfloat16; + mDtypeWeights = btg::Dtype::MxInt4; + + FusedMoeLauncher::init_common( + std::move(args), + tile_tokens_dim, + routing_method_type, + /*use_shuffled_weight=*/true, + static_cast(batchedGemm::gemm::MatrixLayout::BlockMajorK), + ActivationType::Swiglu, + norm_topk_prob); + } + + void check_routing() const override { + FusedMoeLauncher::check_routing_common(); + } + + void prepare_routing() override { + FusedMoeLauncher::prepare_routing_common(); + + args->mDtypeElt = mDtypeAct; + args->mUseDeepSeekFp8 = false; + // Set expert weights dtype based on routing bias + auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; + mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; + + auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; + mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; + + auto expert_weights_dtype = mRoutingLogitsDtype == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; + expert_weights = alloc_tensor({args->num_tokens, args->top_k}, expert_weights_dtype, hidden_states.device()); + + workspace.expert_weights = expert_weights.data_ptr(); + } + + void check_moe() const override { + TVM_FFI_ICHECK(mDtypeAct == btg::Dtype::Bfloat16) << "Only Bfloat16 is supported by MxInt4 block scale MoE"; + + TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_uint8) << "gemm1_weights must be uint8."; + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_bfloat16) << "gemm1_weights_scale must be bf16."; + TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_uint8) << "gemm2_weights must be uint8."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_bfloat16) << "gemm2_weights_scale must be bf16."; + } + + void prepare_moe(int64_t& moe_tactic) override { + args->hidden_states = hidden_states.data_ptr(); + args->hidden_states_scale = nullptr; + args->gemm1_weights = gemm1_weights.data_ptr(); + args->gemm1_weights_scale = gemm1_weights_scale.data_ptr(); + args->gemm1_alpha = gemm1_alpha.has_value() ? static_cast(gemm1_alpha.value().data_ptr()) : nullptr; + args->gemm1_beta = gemm1_beta.has_value() ? static_cast(gemm1_beta.value().data_ptr()) : nullptr; + args->gemm1_clamp_limit = + gemm1_clamp_limit.has_value() ? static_cast(gemm1_clamp_limit.value().data_ptr()) : nullptr; + args->gemm2_weights = gemm2_weights.data_ptr(); + args->gemm2_weights_scale = gemm2_weights_scale.data_ptr(); + args->output1_scales_scalar = nullptr; + args->output1_scales_gate_scalar = nullptr; + args->output2_scales_scalar = nullptr; + + FusedMoeLauncher::prepare_moe_common(moe_tactic); + + max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + workspace.total_max_padded_tokens, args->intermediate_size, btg::dtypeGetNumBits(mDtypeAct)); + max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + workspace.total_max_padded_tokens, + args->hidden_size, + btg::dtypeGetNumBits(btg::Dtype::Bfloat16)); // Output is always BF16 + + auto const gemm1_output_hidden = args->intermediate_size; + gemm1_output = + alloc_tensor({max_num_padded_tokens_gemm1, gemm1_output_hidden}, dl_bfloat16, hidden_states.device()); + + // Allocate gemm2_output + gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); + + // Setup workspace pointers + workspace.hidden_states_scale_linear = nullptr; // MxInt4 doesn't use linear scale + workspace.gemm1_output = gemm1_output.data_ptr(); + workspace.gemm1_output_scale = nullptr; + // Note: activation_output and activation_output_scale are set by the base class + // prepare_moe_common() when gated activation is used + workspace.gemm2_output = gemm2_output.data_ptr(); + workspace.gemm2_output_scale = nullptr; + } + + private: + TensorView gemm1_weights_scale; + Optional gemm1_alpha; + Optional gemm1_beta; + Optional gemm1_clamp_limit; + TensorView gemm2_weights_scale; + int32_t max_num_padded_tokens_gemm1{}; + int32_t max_num_padded_tokens_gemm2{}; + + public: + static Array> getValidConfigs( + int64_t top_k, int64_t hidden_size, int64_t intermediate_size, int64_t num_local_experts, int64_t num_tokens) { + Array> valid_configs; + + std::vector tile_sizes(mSupportedTileNums.begin(), mSupportedTileNums.end()); + std::set selected_tile_nums = computeSelectedTileN(tile_sizes, num_tokens, top_k, num_local_experts); + + for (int32_t tile_N : selected_tile_nums) { + auto moe_runner = std::make_unique( + btg::Dtype::Bfloat16, + btg::Dtype::MxInt4, + false, // useDeepSeekFp8 + tile_N, + ActivationType::Swiglu, + /*useShuffledMatrix*/ true, + batchedGemm::gemm::MatrixLayout::BlockMajorK); + + auto cfgs = + moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); + + for (auto cfg : cfgs) { + valid_configs.push_back({tile_N, cfg}); + } + } + + return valid_configs; + } +}; + +class FP4BlockScaleLauncher : public FusedMoeLauncher { + public: + static constexpr std::array mBaseSupportedTileNums = {8, 16, 32, 64}; + + static std::vector getSupportedTileNums(btg::Dtype dtype_act) { + std::vector tiles(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); + if (dtype_act != btg::Dtype::Bfloat16) { + tiles.push_back(128); + tiles.push_back(256); + } + return tiles; + } + + FP4BlockScaleLauncher( + Optional const& routing_logits, + Optional const& routing_bias, + TensorView const& hidden_states, + Optional const& hidden_states_scale, + TensorView const& gemm1_weights, + TensorView const& gemm1_weights_scale, + Optional const& gemm1_bias, + Optional const& gemm1_alpha, + Optional const& gemm1_beta, + Optional const& gemm1_clamp_limit, + TensorView const& gemm2_weights, + TensorView const& gemm2_weights_scale, + Optional const& gemm2_bias, + Optional const& output1_scales_scalar, + Optional const& output1_scales_gate_scalar, + Optional const& output2_scales_scalar, + Optional const& per_token_scales, + TensorView const& expert_indices, + TensorView const& expert_weights) + : FusedMoeLauncher( + routing_logits, + routing_bias, + hidden_states, + gemm1_weights, + output1_scales_scalar, + output1_scales_gate_scalar, + gemm2_weights, + output2_scales_scalar, + per_token_scales), + hidden_states_scale(hidden_states_scale), + gemm1_weights_scale(gemm1_weights_scale), + gemm1_bias(gemm1_bias), + gemm1_alpha(gemm1_alpha), + gemm1_beta(gemm1_beta), + gemm1_clamp_limit(gemm1_clamp_limit), + gemm2_weights_scale(gemm2_weights_scale), + gemm2_bias(gemm2_bias), + expert_indices(expert_indices), + expert_weights(expert_weights) {} + + void init( + std::unique_ptr&& args, + int64_t tile_tokens_dim, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + ActivationType activation_type, + btg::Dtype dtype_act, + btg::Dtype dtype_weights, + bool norm_topk_prob = true) { + // Set data types + args->mDtypeElt = dtype_act; + args->mDtypeOut = btg::Dtype::Bfloat16; // Output is always BF16 for FP4 + args->mUseDeepSeekFp8 = false; // FP4 doesn't use DeepSeek FP8 + + mDtypeAct = dtype_act; + mDtypeWeights = dtype_weights; + + FusedMoeLauncher::init_common( + std::move(args), + tile_tokens_dim, + routing_method_type, + use_shuffled_weight, + weight_layout, + activation_type, + norm_topk_prob); + } + + void check_routing() const override { + // First call base class common routing checks + FusedMoeLauncher::check_routing_common(); + } + + void prepare_routing() override { + num_tokens_per_expert = alloc_tensor({args->num_experts}, dl_int32, hidden_states.device()); + int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxPermutedPaddedCount( + args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); + + total_num_padded_tokens = alloc_tensor({1}, dl_int32, hidden_states.device()); + expanded_idx_to_permuted_idx = alloc_tensor({args->num_tokens * args->top_k}, dl_int32, hidden_states.device()); + permuted_idx_to_token_idx = alloc_tensor({max_num_padded_tokens}, dl_int32, hidden_states.device()); + + int64_t const size_of_expert_count_histogram = std::max(args->num_experts * 2, 256 * 2); + expert_count_histogram = alloc_tensor({size_of_expert_count_histogram}, dl_int32, hidden_states.device()); + + int32_t max_num_ctas = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxNumCtasInBatchDim( + args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); + cta_idx_xy_to_batch_idx = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); + cta_idx_xy_to_mn_limit = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); + num_non_exiting_ctas = alloc_tensor({1}, dl_int32, hidden_states.device()); + + workspace.total_num_padded_tokens = static_cast(total_num_padded_tokens.data_ptr()); + workspace.total_max_padded_tokens = max_num_padded_tokens; + workspace.ProjUpTileN = tile_tokens_dim; + workspace.routing_expert_indexes = static_cast(const_cast(expert_indices.data_ptr())); + workspace.expert_weights = const_cast(expert_weights.data_ptr()); + workspace.permuted_idx_size = static_cast(total_num_padded_tokens.data_ptr()); + workspace.expanded_idx_to_permuted_idx = static_cast(expanded_idx_to_permuted_idx.data_ptr()); + workspace.permuted_idx_to_token_idx = static_cast(permuted_idx_to_token_idx.data_ptr()); + workspace.cta_idx_xy_to_batch_idx = static_cast(cta_idx_xy_to_batch_idx.data_ptr()); + workspace.cta_idx_xy_to_mn_limit = static_cast(cta_idx_xy_to_mn_limit.data_ptr()); + workspace.num_non_exiting_ctas = static_cast(num_non_exiting_ctas.data_ptr()); + + args->mDtypeElt = mDtypeAct; + auto routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; + mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; + + auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; + mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; + } + + void check_moe() const override { + TVM_FFI_ICHECK( + mDtypeAct == btg::Dtype::E2m1 || mDtypeAct == btg::Dtype::Bfloat16 || mDtypeAct == btg::Dtype::E4m3 || + mDtypeAct == btg::Dtype::MxE4m3) + << "Only E2m1, Bfloat16, MxE4m3 and E4m3 are supported by Fp4 block scale MoE"; + + if (mDtypeAct == btg::Dtype::E2m1) { + TVM_FFI_ICHECK(mDtypeWeights == btg::Dtype::E2m1) + << "Only E2m1 and MxE2m1 are supported by block scale MoE with E2m1 activation"; + TVM_FFI_ICHECK(hidden_states_scale.has_value()) << "hidden_states_scale is required for E2m1 activation"; + TVM_FFI_ICHECK(output1_scales_scalar.has_value()) << "output1_scales_scalar is required for E2m1 activation"; + TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()) + << "output1_scales_gate_scalar is required for E2m1 activation"; + TVM_FFI_ICHECK(output2_scales_scalar.has_value()) << "output2_scales_scalar is required for E2m1 activation"; + } else if (mDtypeAct == btg::Dtype::Bfloat16 || mDtypeAct == btg::Dtype::E4m3 || mDtypeAct == btg::Dtype::MxE4m3) { + TVM_FFI_ICHECK(mDtypeWeights == btg::Dtype::MxE2m1) + << "Only MxE2m1 weights are supported by block scale MoE with Bfloat16, E4m3 or " + "MxE4m3 activation"; + } + + if (mDtypeAct == btg::Dtype::E4m3) { + TVM_FFI_ICHECK(output1_scales_scalar.has_value()) << "output1_scales_scalar is required for E4m3 activation"; + TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()) + << "output1_scales_gate_scalar is required for E4m3 activation"; + TVM_FFI_ICHECK(output2_scales_scalar.has_value()) << "output2_scales_scalar is required for E4m3 activation"; + } + + TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_uint8) << "gemm1_weights must be byte."; + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_float8_e4m3fn) << "gemm1_weights_scale must be fp8."; + TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_uint8) << "gemm2_weights must be byte."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_float8_e4m3fn) << "gemm2_weights_scale must be fp8."; + } + + void prepare_moe(int64_t& moe_tactic) override { + args->hidden_states = hidden_states.data_ptr(); + args->hidden_states_scale = hidden_states_scale.has_value() ? hidden_states_scale.value().data_ptr() : nullptr; + args->gemm1_weights = gemm1_weights.data_ptr(); + args->gemm1_weights_scale = gemm1_weights_scale.data_ptr(); + args->gemm1_bias = gemm1_bias.has_value() ? static_cast(gemm1_bias.value().data_ptr()) : nullptr; + args->gemm1_alpha = gemm1_alpha.has_value() ? static_cast(gemm1_alpha.value().data_ptr()) : nullptr; + args->gemm1_beta = gemm1_beta.has_value() ? static_cast(gemm1_beta.value().data_ptr()) : nullptr; + args->gemm1_clamp_limit = + gemm1_clamp_limit.has_value() ? static_cast(gemm1_clamp_limit.value().data_ptr()) : nullptr; + args->gemm2_weights = gemm2_weights.data_ptr(); + args->gemm2_weights_scale = gemm2_weights_scale.data_ptr(); + args->gemm2_bias = gemm2_bias.has_value() ? static_cast(gemm2_bias.value().data_ptr()) : nullptr; + args->output1_scales_scalar = + output1_scales_scalar.has_value() ? static_cast(output1_scales_scalar.value().data_ptr()) : nullptr; + args->output1_scales_gate_scalar = output1_scales_gate_scalar.has_value() + ? static_cast(output1_scales_gate_scalar.value().data_ptr()) + : nullptr; + args->output2_scales_scalar = + output2_scales_scalar.has_value() ? static_cast(output2_scales_scalar.value().data_ptr()) : nullptr; + + FusedMoeLauncher::prepare_moe_common(moe_tactic); + + auto const sf_vec_size = mDtypeWeights == btg::Dtype::MxE2m1 ? 32 : 16; + + max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + workspace.total_max_padded_tokens, args->intermediate_size, btg::dtypeGetNumBits(mDtypeAct)); + max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + workspace.total_max_padded_tokens, + args->hidden_size, + btg::dtypeGetNumBits(btg::Dtype::Bfloat16)); // Output is always BF16 + + auto const gemm1_output_hidden = + mDtypeAct == btg::Dtype::E2m1 ? args->intermediate_size / 2 : args->intermediate_size; + if (mDtypeAct == btg::Dtype::E2m1 || mDtypeAct == btg::Dtype::MxE4m3) { + int64_t sf_size = + tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens_gemm1, args->intermediate_size / sf_vec_size); + gemm1_output_scale = alloc_tensor({sf_size}, dl_uint8, hidden_states.device()); + } + if (!per_token_scales.has_value()) { + gemm1_output = alloc_tensor( + {max_num_padded_tokens_gemm1, gemm1_output_hidden}, + mDtypeAct == btg::Dtype::Bfloat16 ? dl_bfloat16 : dl_uint8, + hidden_states.device()); + } else { // FC1 output is Bfloat16 + TVM_FFI_ICHECK(mDtypeAct == btg::Dtype::E2m1) + << "NvFP4 MoE: currently only support NvFP4 x NvFP4 when using per-token scaling."; + // When per-token scales are used, the FC1 output is always BF16 and will be quantized + gemm1_output = + alloc_tensor({max_num_padded_tokens_gemm1, args->intermediate_size}, dl_bfloat16, hidden_states.device()); + activation_output = + alloc_tensor({max_num_padded_tokens_gemm1, gemm1_output_hidden}, dl_uint8, hidden_states.device()); + per_token_scales_fc2 = alloc_tensor({max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); + } + + // Allocate gemm2_output + gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); + + // Setup workspace pointers + workspace.hidden_states_scale_linear = nullptr; // FP4 doesn't use linear scale + workspace.gemm1_output = gemm1_output.data_ptr(); + workspace.gemm1_output_scale = + gemm1_output_scale.has_value() ? static_cast(gemm1_output_scale.value().data_ptr()) : nullptr; + if (per_token_scales.has_value()) { + workspace.token_scales = per_token_scales.value().data_ptr(); + workspace.activation_output = activation_output.data_ptr(); + workspace.activation_output_scale = workspace.gemm1_output_scale; + workspace.token_scales_fc2 = per_token_scales_fc2.data_ptr(); + } + workspace.gemm2_output = gemm2_output.data_ptr(); + workspace.gemm2_output_scale = nullptr; + } + + private: + Optional hidden_states_scale; + TensorView gemm1_weights_scale; + Optional gemm1_bias; + Optional gemm1_alpha; + Optional gemm1_beta; + Optional gemm1_clamp_limit; + TensorView gemm2_weights_scale; + Optional gemm2_bias; + int32_t max_num_padded_tokens_gemm1{}; + int32_t max_num_padded_tokens_gemm2{}; + Optional gemm1_output_scale; + TensorView expert_indices; + TensorView expert_weights; + + public: + Array + run(int64_t moe_tactic, + bool enable_pdl = true, + bool use_routing_scales_on_input = false, + bool use_deep_seek_fp8 = false) override { + check_routing(); + prepare_routing(); + + // Execute routing + tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim); + cudaStream_t routing_stream = get_stream(hidden_states.device()); + + int16_t* replay_ptr = nullptr; + if (routing_replay_out.has_value()) { + replay_ptr = reinterpret_cast(routing_replay_out.value().data_ptr()); + } + + routing_runner.run( + args->routing_logits, + args->routing_bias, + args->num_tokens, + args->num_experts, + args->top_k, + args->n_group, + args->topk_group, + args->local_expert_offset, + args->local_num_experts, + args->routed_scaling_factor, + static_cast(expert_indices.data_ptr()), + static_cast(expert_count_histogram.data_ptr()), + static_cast(total_num_padded_tokens.data_ptr()), + static_cast(expanded_idx_to_permuted_idx.data_ptr()), + nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, + static_cast(permuted_idx_to_token_idx.data_ptr()), + expert_weights.data_ptr(), + static_cast(num_tokens_per_expert.data_ptr()), + static_cast(cta_idx_xy_to_batch_idx.data_ptr()), + static_cast(cta_idx_xy_to_mn_limit.data_ptr()), + static_cast(num_non_exiting_ctas.data_ptr()), + args->mDtypeElt, + mRoutingBiasDtype, + use_routing_scales_on_input, + use_deep_seek_fp8, + static_cast(routing_method_type), + routing_stream, + mRoutingLogitsDtype, + norm_topk_prob, + replay_ptr); + + check_moe(); + prepare_moe(moe_tactic); + + cudaStream_t moe_stream = get_stream(hidden_states.device()); + moe_runner->run(*args, workspace, hidden_states.device().device_id, moe_stream, moe_tactic, enable_pdl); + + // Match original FP4 behavior for return values + if (args->do_finalize) { + return {output}; + } + return {gemm2_output, FusedMoeLauncher::expert_weights, expanded_idx_to_permuted_idx}; + } + + static Array> getValidConfigs( + int64_t top_k, + int64_t hidden_size, + int64_t intermediate_size, + int64_t num_local_experts, + int64_t num_tokens, + int64_t act_type, + btg::Dtype dtype_act, + btg::Dtype dtype_weights, + bool use_per_token_scaling) { + Array> valid_configs; + + std::vector tile_sizes = getSupportedTileNums(dtype_act); + std::set selected_tile_nums = computeSelectedTileN(tile_sizes, num_tokens, top_k, num_local_experts); + + for (int32_t tile_N : selected_tile_nums) { + auto moe_runner = std::make_unique( + dtype_act, + dtype_weights, + false, // useDeepSeekFp8 + tile_N, + static_cast(act_type), + /*useShuffledMatrix*/ true, + /*weight_layout*/ batchedGemm::gemm::MatrixLayout::MajorK, + // NOTE(siyuan): currently FP4 MoE always apply per-token scaling to both FC1 and FC2. + /*usePerTokenScalingGemm1*/ use_per_token_scaling, + /*usePerTokenScalingGemm2*/ use_per_token_scaling, + false, + false); + + auto cfgs = + moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); + + for (auto cfg : cfgs) { + valid_configs.push_back({tile_N, cfg}); + } + } + + return valid_configs; + } +}; + +Array trtllm_bf16_moe( + Optional const& routing_logits, + Optional const& routing_bias, + TensorView const& expert_indices, + TensorView const& expert_weights, + TensorView const& hidden_states, + TensorView const& gemm1_weights, + TensorView const& gemm2_weights, + TensorView output, + int64_t num_experts, + int64_t top_k, + Optional n_group, + Optional topk_group, + int64_t intermediate_size, + int64_t local_expert_offset, + int64_t local_num_experts, + Optional routed_scaling_factor, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + bool do_finalize, + bool enable_pdl, + Array moe_tactic, + int64_t activation_type, + bool norm_topk_prob, + Optional routing_replay_out) { + // Just some basic type validation first and leave more checks to the launcher + if (routing_logits.has_value()) { + TVM_FFI_ICHECK(routing_logits.value().dtype() == dl_float32 || routing_logits.value().dtype() == dl_bfloat16) + << "BF16 MoE: routing_logits must be bfloat16 or float."; + } + TVM_FFI_ICHECK_EQ(hidden_states.dtype(), dl_bfloat16) << "BF16 MoE: hidden_states must be bfloat16."; + TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_bfloat16) << "BF16 MoE: gemm1_weights must be bfloat16."; + TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_bfloat16) << "BF16 MoE: gemm2_weights must be bfloat16."; + + if (routing_replay_out.has_value()) { + validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); + } + + auto const num_tokens = hidden_states.size(0); + auto const hidden_size = hidden_states.size(1); + auto const activation = validateAndCastActivationType(activation_type); + + // Calculate supported tile sizes + std::vector mSupportedTileN( + Bf16MoeLauncher::mSupportedTileNums.begin(), Bf16MoeLauncher::mSupportedTileNums.end()); + // Build launchers for ALL supported tiles (not just the computeSelectedTileN subset) + // so that autotuner-cached tactics always find their tile_N in the map. + // Launcher creation is cheap (no GPU allocation until run()), so this is safe. + + // Create a map of launchers for each tile size + std::unordered_map> launchers_map; + + for (int32_t curr_tile_N : mSupportedTileN) { + // Create MoE arguments for this launcher + auto args = std::make_unique(); + args->num_tokens = num_tokens; + args->num_experts = num_experts; + args->hidden_size = hidden_size; + args->hidden_size_output = args->hidden_size; + args->top_k = top_k; + args->n_group = n_group.value_or(0); + args->topk_group = topk_group.value_or(0); + args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); + args->local_expert_offset = local_expert_offset; + args->local_num_experts = local_num_experts; + args->intermediate_size = intermediate_size; + args->do_finalize = do_finalize; + args->output = output.data_ptr(); + args->output_scale = nullptr; + + // Create and initialize launcher for this tile size + auto launcher = std::make_unique( + routing_logits, routing_bias, expert_indices, expert_weights, hidden_states, gemm1_weights, gemm2_weights); + launcher->init( + std::move(args), + curr_tile_N, + routing_method_type, + use_shuffled_weight, + weight_layout, + activation, + norm_topk_prob); + launcher->set_routing_replay_out(routing_replay_out); + + launchers_map[curr_tile_N] = std::move(launcher); + } + + auto const [tile_N, config] = + resolveMoeTileAndConfig(moe_tactic, mSupportedTileN, num_tokens, top_k, local_num_experts); + + // Get the launcher for the selected tile_N + auto launcher_it = launchers_map.find(static_cast(tile_N)); + FLASHINFER_CHECK(launcher_it != launchers_map.end(), "Internal error: missing BF16 MoE launcher for tile_N=", tile_N); + auto& selected_launcher = launcher_it->second; + + // Run the launcher - it will create its own runner internally + return selected_launcher->run(config, enable_pdl); +} + +Array trtllm_fp8_per_tensor_scale_moe( + TensorView routing_logits, + Optional routing_bias, + TensorView hidden_states, + TensorView gemm1_weights, + TensorView output1_scales_scalar, + TensorView output1_scales_gate_scalar, + TensorView gemm2_weights, + TensorView output2_scales_scalar, + TensorView output, + int64_t num_experts, + int64_t top_k, + Optional n_group, + Optional topk_group, + int64_t intermediate_size, + int64_t local_expert_offset, + int64_t local_num_experts, + Optional routed_scaling_factor, + bool use_routing_scales_on_input, + int64_t routing_method_type, + bool do_finalize, + bool enable_pdl, + Array config_index, + int64_t activation_type, + bool norm_topk_prob, + Optional routing_replay_out) { + // Basic type validation + auto dtype = hidden_states.dtype(); + auto activation = validateAndCastActivationType(activation_type); + + TVM_FFI_ICHECK(dtype == dl_float8_e4m3fn || dtype == dl_float16 || dtype == dl_bfloat16) + << "FP8 MoE: hidden_states must be float8_e4m3fn, float16, or bfloat16."; + TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm1_weights must be float8_e4m3fn."; + TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm2_weights must be float8_e4m3fn."; + TVM_FFI_ICHECK_EQ(output1_scales_scalar.dtype(), dl_float32) << "FP8 MoE: output1_scales_scalar must be float32."; + TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.dtype(), dl_float32) + << "FP8 MoE: output1_scales_gate_scalar must be float32."; + TVM_FFI_ICHECK_EQ(output2_scales_scalar.dtype(), dl_float32) << "FP8 MoE: output2_scales_scalar must be float32."; + + if (routing_replay_out.has_value()) { + validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); + } + + auto const num_tokens = hidden_states.size(0); + auto const hidden_size = hidden_states.size(1); + + // Use default values that match the original function behavior + bool use_shuffled_weight = true; // Original uses /*useShuffledMatrix*/ true + int64_t weight_layout = 0; // Default to MajorK + + // Calculate supported tile sizes + std::vector mSupportedTileN( + Fp8PerTensorLauncher::mSupportedTileNums.begin(), Fp8PerTensorLauncher::mSupportedTileNums.end()); + // Build launchers for ALL supported tiles so autotuner-cached tactics always find their tile_N. + + // Create a map of launchers for each tile size + std::unordered_map> launchers_map; + + for (int32_t curr_tile_N : mSupportedTileN) { + // Create MoE arguments for this launcher + auto args = std::make_unique(); + args->num_tokens = num_tokens; + args->num_experts = num_experts; + args->hidden_size = hidden_size; + args->hidden_size_output = args->hidden_size; + args->top_k = top_k; + args->n_group = n_group.value_or(0); + args->topk_group = topk_group.value_or(0); + args->local_expert_offset = local_expert_offset; + args->local_num_experts = local_num_experts; + args->intermediate_size = intermediate_size; + args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); + args->do_finalize = do_finalize; + args->output = output.data_ptr(); + args->output_scale = nullptr; + + // Create and initialize launcher for this tile size + auto launcher = std::make_unique( + routing_logits, + routing_bias, + hidden_states, + gemm1_weights, + output1_scales_scalar, + output1_scales_gate_scalar, + gemm2_weights, + output2_scales_scalar); + launcher->init( + std::move(args), + curr_tile_N, + routing_method_type, + use_shuffled_weight, + weight_layout, + use_routing_scales_on_input, + activation, + norm_topk_prob); + launcher->set_routing_replay_out(routing_replay_out); + + launchers_map[curr_tile_N] = std::move(launcher); + } + + auto const [tile_N, config] = + resolveMoeTileAndConfig(config_index, mSupportedTileN, num_tokens, top_k, local_num_experts); + + // Get the launcher for the selected tile_N + auto launcher_it = launchers_map.find(static_cast(tile_N)); + FLASHINFER_CHECK( + launcher_it != launchers_map.end(), "Internal error: missing FP8 per-tensor MoE launcher for tile_N=", tile_N); + auto& selected_launcher = launcher_it->second; + + // Run the launcher - it will create its own runner internally + return selected_launcher->run(config, enable_pdl, use_routing_scales_on_input); +} + +Array trtllm_fp8_block_scale_moe_impl( + Optional routing_logits, + TensorView expert_indices, + TensorView expert_weights, + Optional routing_bias, + TensorView hidden_states, + TensorView hidden_states_scale, + TensorView gemm1_weights, + TensorView gemm1_weights_scale, + TensorView gemm2_weights, + TensorView gemm2_weights_scale, + TensorView output, + int64_t num_experts, + int64_t top_k, + Optional n_group, + Optional topk_group, + int64_t intermediate_size, + int64_t local_expert_offset, + int64_t local_num_experts, + Optional routed_scaling_factor, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + bool do_finalize, + bool enable_pdl, + Array config_index, + Fp8QuantizationType quantization_type, + int64_t act_type, + bool norm_topk_prob, + Optional routing_replay_out, + Optional gate_up_lora_delta, + Optional activation_lora_input, + int64_t lora_ready_event = 0, + int64_t gemm2_done_event = 0) { + auto activation_type = validateAndCastActivationType(act_type); + // DeepSeekFp8 currently uses a TRTLLM runner that hardwires Swiglu activation semantics. + // Fail for any other activation to avoid silently running incorrect activation behavior. + if (quantization_type == Fp8QuantizationType::DeepSeekFp8 && activation_type != ActivationType::Swiglu) { + TVM_FFI_LOG_AND_THROW(NotImplementedError) + << "DeepSeekFp8 only supports ActivationType::Swiglu in this runner path. " + << "Received activation_type=" << static_cast(activation_type); + } + + // Basic type validation + auto dtype = hidden_states.dtype(); + + // Either routing_logits or expert_indices must be provided + // expert_indices is a packed tensor: (expert_id << 16) | (weight_bf16.view(int16)) + bool use_routing_logits = routing_logits.has_value(); + // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr + bool use_precomputed_routing = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; + + TVM_FFI_ICHECK(use_routing_logits || use_precomputed_routing) + << "Either routing_logits or expert_indices must be provided."; + + (void)use_routing_logits; + TVM_FFI_ICHECK(dtype == dl_float16 || dtype == dl_bfloat16 || dtype == dl_float8_e4m3fn) + << "FP8 block scale MoE: hidden_states must be fp16, bf16, or fp8."; + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_float32) + << "FP8 block scale MoE: hidden_states_scale must be float32."; + } else if (quantization_type == Fp8QuantizationType::MxFp8) { + TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_uint8) + << "FP8 block scale MoE: hidden_states_scale must be uint8."; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; + } + TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "FP8 block scale MoE: gemm1_weights must be fp8."; + TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "FP8 block scale MoE: gemm2_weights must be fp8."; + if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_float32) + << "FP8 block scale MoE: gemm1_weights_scale must be float32."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_float32) + << "FP8 block scale MoE: gemm2_weights_scale must be float32."; + } else if (quantization_type == Fp8QuantizationType::MxFp8) { + TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_uint8) + << "FP8 block scale MoE: gemm1_weights_scale must be uint8."; + TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_uint8) + << "FP8 block scale MoE: gemm2_weights_scale must be uint8."; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; + } + + if (quantization_type == Fp8QuantizationType::MxFp8) { + TVM_FFI_ICHECK(use_shuffled_weight) << "use_shuffled_weight must be true for MxFp8."; + TVM_FFI_ICHECK(weight_layout == 0) << "weight_layout must be 0 for MxFp8."; + } + + if (routing_replay_out.has_value()) { + validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); + } + + auto const num_tokens = hidden_states.size(0); + auto const hidden_size = hidden_states.size(1); + + auto supported_tile_nums = Fp8BlockScaleLauncher::getSupportedTileNums(quantization_type); + // Build launchers for ALL supported tiles so autotuner-cached tactics always find their tile_N. + + // Create a map of launchers for each tile size + std::unordered_map> launchers_map; + + for (int32_t curr_tile_N : supported_tile_nums) { + // Create MoE arguments for this launcher + auto args = std::make_unique(); + args->num_tokens = num_tokens; + args->num_experts = num_experts; + args->hidden_size = hidden_size; + args->hidden_size_output = args->hidden_size; + args->top_k = top_k; + args->n_group = n_group.value_or(0); + args->topk_group = topk_group.value_or(0); + args->local_expert_offset = local_expert_offset; + args->local_num_experts = local_num_experts; + args->intermediate_size = intermediate_size; + args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); + args->do_finalize = do_finalize; + args->output = output.data_ptr(); + args->output_scale = nullptr; + // GEMM1-LoRA overlap: cudaEvent_t handle (recorded on the LoRA side stream) the runner + // waits on right before activation; 0 = no wait (serial path). + args->lora_ready_event = reinterpret_cast(lora_ready_event); + // Down-LoRA/finalize overlap: cudaEvent_t handle the runner records right after GEMM2 + // (before finalize) so the LoRA side stream can overlap the down-proj LoRA with + // finalize; 0 = no record (serial path). + args->gemm2_done_event = reinterpret_cast(gemm2_done_event); + + // Create and initialize launcher for this tile size + auto launcher = std::make_unique( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm2_weights, + gemm2_weights_scale, + expert_indices, + expert_weights, + quantization_type, + gate_up_lora_delta, + activation_lora_input); + launcher->init( + std::move(args), + curr_tile_N, + routing_method_type, + use_shuffled_weight, + weight_layout, + activation_type, + norm_topk_prob); + launcher->set_routing_replay_out(routing_replay_out); + + launchers_map[curr_tile_N] = std::move(launcher); + } + + auto const [tile_N, config] = + resolveMoeTileAndConfig(config_index, supported_tile_nums, num_tokens, top_k, local_num_experts); + + // Get the launcher for the selected tile_N + auto launcher_it = launchers_map.find(static_cast(tile_N)); + FLASHINFER_CHECK( + launcher_it != launchers_map.end(), "Internal error: missing FP8 block-scale MoE launcher for tile_N=", tile_N); + auto& selected_launcher = launcher_it->second; + + // Run the launcher with DeepSeek FP8 enabled - it will create its own runner internally + return selected_launcher->run( + config, + enable_pdl, + false /* use_routing_scales_on_input */, + quantization_type == Fp8QuantizationType::DeepSeekFp8 /* use_deep_seek_fp8 */); +} + +Array trtllm_fp8_block_scale_moe( + Optional routing_logits, + TensorView expert_indices, + TensorView expert_weights, + Optional routing_bias, + TensorView hidden_states, + TensorView hidden_states_scale, + TensorView gemm1_weights, + TensorView gemm1_weights_scale, + TensorView gemm2_weights, + TensorView gemm2_weights_scale, + TensorView output, + int64_t num_experts, + int64_t top_k, + Optional n_group, + Optional topk_group, + int64_t intermediate_size, + int64_t local_expert_offset, + int64_t local_num_experts, + Optional routed_scaling_factor, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + bool do_finalize, + bool enable_pdl, + Array config_index, + Fp8QuantizationType quantization_type, + int64_t act_type, + bool norm_topk_prob, + Optional routing_replay_out) { + return trtllm_fp8_block_scale_moe_impl( + routing_logits, + expert_indices, + 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, + config_index, + quantization_type, + act_type, + norm_topk_prob, + routing_replay_out, + Optional(), + Optional()); +} + +Array sgl_trtllm_fp8_block_scale_moe_lora( + Optional routing_logits, + TensorView expert_indices, + TensorView expert_weights, + Optional routing_bias, + TensorView hidden_states, + TensorView hidden_states_scale, + TensorView gemm1_weights, + TensorView gemm1_weights_scale, + TensorView gemm2_weights, + TensorView gemm2_weights_scale, + TensorView output, + int64_t num_experts, + int64_t top_k, + Optional n_group, + Optional topk_group, + int64_t intermediate_size, + int64_t local_expert_offset, + int64_t local_num_experts, + Optional routed_scaling_factor, + int64_t routing_method_type, + bool use_shuffled_weight, + int64_t weight_layout, + bool do_finalize, + bool enable_pdl, + Array config_index, + Fp8QuantizationType quantization_type, + int64_t act_type, + bool norm_topk_prob, + Optional routing_replay_out, + TensorView gate_up_lora_delta, + TensorView activation_lora_input, + int64_t lora_ready_event, + int64_t gemm2_done_event) { + if (quantization_type != Fp8QuantizationType::DeepSeekFp8) { + TVM_FFI_LOG_AND_THROW(NotImplementedError) + << "sgl_trtllm_fp8_block_scale_moe_lora currently supports DeepSeekFp8 only."; + } + return trtllm_fp8_block_scale_moe_impl( + routing_logits, + expert_indices, + 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, + config_index, + quantization_type, + act_type, + norm_topk_prob, + routing_replay_out, + Optional(gate_up_lora_delta), + Optional(activation_lora_input), + lora_ready_event, + gemm2_done_event); +} + +__global__ void sgl_trtllm_fp8_block_scale_moe_lora_finalize_kernel( + cutlass::bfloat16_t const* __restrict__ gemm2_output, + cutlass::bfloat16_t const* __restrict__ expert_weights, + int32_t const* __restrict__ expanded_idx_to_permuted_idx, + cutlass::bfloat16_t const* __restrict__ down_lora_delta, + cutlass::bfloat16_t* __restrict__ output, + int64_t num_tokens, + int64_t top_k, + int64_t hidden_size, + int64_t hidden_size_padded, + float routed_scaling_factor) { + for (int64_t token_idx = blockIdx.y; token_idx < num_tokens; token_idx += gridDim.y) { + for (int64_t hidden_idx = threadIdx.x + blockDim.x * blockIdx.x; hidden_idx < hidden_size; + hidden_idx += blockDim.x * gridDim.x) { + float acc = 0.0f; + float lora_acc = 0.0f; + for (int64_t k = 0; k < top_k; ++k) { + int64_t const expanded_idx = token_idx * top_k + k; + int32_t const permuted_idx = expanded_idx_to_permuted_idx[expanded_idx]; + if (permuted_idx != -1) { + float const expert_prob = static_cast(expert_weights[token_idx * top_k + k]); + acc += expert_prob * static_cast(gemm2_output[permuted_idx * hidden_size_padded + hidden_idx]); + } + lora_acc += static_cast(down_lora_delta[expanded_idx * hidden_size + hidden_idx]); + } + output[token_idx * hidden_size + hidden_idx] = + static_cast(acc + routed_scaling_factor * lora_acc); + } + } +} + +void sgl_trtllm_fp8_block_scale_moe_lora_finalize( + TensorView gemm2_output, + TensorView expert_weights, + TensorView expanded_idx_to_permuted_idx, + TensorView down_lora_delta, + TensorView output, + Optional routed_scaling_factor) { + TVM_FFI_ICHECK_EQ(gemm2_output.dtype(), dl_bfloat16) << "gemm2_output must be bfloat16."; + TVM_FFI_ICHECK_EQ(expert_weights.dtype(), dl_bfloat16) << "expert_weights must be bfloat16."; + TVM_FFI_ICHECK((expanded_idx_to_permuted_idx.dtype() == DLDataType{kDLInt, 32, 1})) + << "expanded_idx_to_permuted_idx must be int32."; + TVM_FFI_ICHECK_EQ(down_lora_delta.dtype(), dl_bfloat16) << "down_lora_delta must be bfloat16."; + TVM_FFI_ICHECK_EQ(output.dtype(), dl_bfloat16) << "output must be bfloat16."; + TVM_FFI_ICHECK_EQ(gemm2_output.ndim(), 2) << "gemm2_output must be 2D."; + TVM_FFI_ICHECK_EQ(expert_weights.ndim(), 2) << "expert_weights must be 2D."; + TVM_FFI_ICHECK_EQ(expanded_idx_to_permuted_idx.ndim(), 1) << "expanded_idx_to_permuted_idx must be 1D."; + TVM_FFI_ICHECK_EQ(down_lora_delta.ndim(), 3) << "down_lora_delta must be 3D."; + TVM_FFI_ICHECK_EQ(output.ndim(), 2) << "output must be 2D."; + TVM_FFI_ICHECK(gemm2_output.IsContiguous()) << "gemm2_output must be contiguous."; + TVM_FFI_ICHECK(expert_weights.IsContiguous()) << "expert_weights must be contiguous."; + TVM_FFI_ICHECK(expanded_idx_to_permuted_idx.IsContiguous()) << "expanded_idx_to_permuted_idx must be contiguous."; + TVM_FFI_ICHECK(down_lora_delta.IsContiguous()) << "down_lora_delta must be contiguous."; + TVM_FFI_ICHECK(output.IsContiguous()) << "output must be contiguous."; + + int64_t const num_tokens = output.size(0); + int64_t const hidden_size = output.size(1); + int64_t const top_k = down_lora_delta.size(1); + TVM_FFI_ICHECK_EQ(expert_weights.size(0), num_tokens) << "expert_weights dim0 must equal num_tokens."; + TVM_FFI_ICHECK_EQ(expert_weights.size(1), top_k) << "expert_weights dim1 must equal top_k."; + TVM_FFI_ICHECK_EQ(down_lora_delta.size(0), num_tokens) << "down_lora_delta dim0 must equal num_tokens."; + TVM_FFI_ICHECK_EQ(down_lora_delta.size(2), hidden_size) << "down_lora_delta dim2 must equal hidden_size."; + TVM_FFI_ICHECK_EQ(expanded_idx_to_permuted_idx.size(0), num_tokens * top_k) + << "expanded_idx_to_permuted_idx size must equal num_tokens * top_k."; + TVM_FFI_ICHECK(gemm2_output.size(1) >= hidden_size) + << "gemm2_output hidden dimension is smaller than output hidden dimension."; + + int const num_threads = 128; + int const num_blocks_x = (hidden_size + num_threads - 1) / num_threads; + int const num_blocks_y = std::min(8192, num_tokens); + dim3 grid(num_blocks_x, num_blocks_y); + cudaStream_t stream = get_stream(output.device()); + sgl_trtllm_fp8_block_scale_moe_lora_finalize_kernel<<>>( + static_cast(gemm2_output.data_ptr()), + static_cast(expert_weights.data_ptr()), + static_cast(expanded_idx_to_permuted_idx.data_ptr()), + static_cast(down_lora_delta.data_ptr()), + static_cast(output.data_ptr()), + num_tokens, + top_k, + hidden_size, + gemm2_output.size(1), + static_cast(routed_scaling_factor.value_or(1.0))); + auto err = cudaGetLastError(); + FLASHINFER_CHECK(err == cudaSuccess, cudaGetErrorString(err)); +} + +Array trtllm_fp4_block_scale_moe( + Optional routing_logits, + TensorView expert_indices, + TensorView expert_weights, + Optional routing_bias, + TensorView hidden_states, + Optional hidden_states_scale, + TensorView gemm1_weights, + TensorView gemm1_weights_scale, + Optional gemm1_bias, + Optional gemm1_alpha, + Optional gemm1_beta, + Optional gemm1_clamp_limit, + TensorView gemm2_weights, + TensorView gemm2_weights_scale, + Optional gemm2_bias, + Optional output1_scales_scalar, + Optional output1_scales_gate_scalar, + Optional output2_scales_scalar, + Optional per_token_scales, + int64_t num_experts, + int64_t top_k, + Optional n_group, + Optional topk_group, + int64_t intermediate_size, + int64_t local_expert_offset, + int64_t local_num_experts, + Optional routed_scaling_factor, + int64_t routing_method_type, + bool do_finalize, + bool enable_pdl, + int64_t act_type, + TensorView output, + Array config_index, + bool norm_topk_prob, + Optional routing_replay_out) { + // Determine data types based on input format + int const num_tokens = hidden_states.size(0); + int hidden_size = hidden_states.size(1); + if (hidden_states.dtype() == dl_uint8) hidden_size *= 2; + + int64_t hidden_states_scale_vec_size = -1; + if (hidden_states_scale.has_value()) { + hidden_states_scale_vec_size = + (static_cast(num_tokens) * hidden_size) / hidden_states_scale.value().numel(); + } + int64_t intermediate_size_factor = isGatedActivation(static_cast(act_type)) ? 2 : 1; + int64_t logical_scale_count = + static_cast(local_num_experts) * intermediate_size * intermediate_size_factor * hidden_size; + int64_t weight_scale_vec_size_raw = logical_scale_count / gemm1_weights_scale.numel(); + + // Snap to nearest valid sf_vec_size (16 or 32). + // The raw value may be slightly smaller than the true vec_size because + // block_scale_interleave pads scale columns to a multiple of 4, inflating numel(). + int64_t weight_scale_vec_size = weight_scale_vec_size_raw > 16 ? 32 : 16; + + // Round-trip validation: the unpadded scale count must not exceed actual numel + // (padding only adds elements, never removes them). + int64_t expected_unpadded = logical_scale_count / weight_scale_vec_size; + TVM_FFI_ICHECK(gemm1_weights_scale.numel() >= expected_unpadded) + << "weight scale tensor too small: numel=" << gemm1_weights_scale.numel() << " but expected at least " + << expected_unpadded << " for sf_vec_size=" << weight_scale_vec_size; + + auto mDtypeWeights = weight_scale_vec_size == 16 ? btg::Dtype::E2m1 : btg::Dtype::MxE2m1; + + if (routing_bias.has_value()) { + TVM_FFI_ICHECK(routing_bias.value().dtype() == dl_bfloat16 || routing_bias.value().dtype() == dl_float32) + << "routing_bias must be bfloat16 or float."; + + TVM_FFI_ICHECK_EQ(routing_bias.value().ndim(), 1) << "routing_bias must be 1D."; + TVM_FFI_ICHECK_EQ(routing_bias.value().size(0), num_experts) << "routing_bias has incorrect shape."; + } + + if (routing_replay_out.has_value()) { + validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); + } + + // Determine activation type + TVM_FFI_ICHECK(gemm1_weights.dtype() == dl_uint8 && gemm2_weights.dtype() == dl_uint8) + << "weights must be fp4 packed in uint8."; + TVM_FFI_ICHECK( + hidden_states.dtype() == dl_uint8 || hidden_states.dtype() == dl_bfloat16 || + hidden_states.dtype() == dl_float8_e4m3fn) + << "hidden_states must be bf16, fp8 or uint8 (packed fp4)."; + + auto mDtypeAct = btg::Dtype::Bfloat16; + if (hidden_states.dtype() == dl_uint8) { + TVM_FFI_ICHECK(hidden_states_scale.has_value() && hidden_states_scale.value().dtype() == dl_float8_e4m3fn) + << "hidden_states_scale must be provided for fp4 activation."; + if (hidden_states_scale_vec_size == 16) { + mDtypeAct = btg::Dtype::E2m1; + } else if (hidden_states_scale_vec_size == 32) { + mDtypeAct = btg::Dtype::MxE2m1; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported hidden state scale shape."; + } + } else if (hidden_states.dtype() == dl_float8_e4m3fn) { + if (hidden_states_scale.has_value()) { + if (hidden_states_scale_vec_size == 32) { + mDtypeAct = btg::Dtype::MxE4m3; + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported hidden state scale shape."; + } + } else { + mDtypeAct = btg::Dtype::E4m3; + } + } + + // Determine supported tile sizes + std::vector mSupportedTileN = FP4BlockScaleLauncher::getSupportedTileNums(mDtypeAct); + // Build launchers for ALL supported tiles so autotuner-cached tactics always find their tile_N. + + // Create a map of launchers for each tile size + std::unordered_map> launchers_map; + + for (int32_t curr_tile_N : mSupportedTileN) { + // Create MoE arguments for this launcher + auto args = std::make_unique(); + args->num_tokens = num_tokens; + args->num_experts = num_experts; + // For E2m1, hidden_size is already multiplied by 2 above, so use it directly + args->hidden_size = hidden_size; + args->hidden_size_output = output.size(1); + args->top_k = top_k; + args->n_group = n_group.value_or(0); + args->topk_group = topk_group.value_or(0); + args->local_expert_offset = local_expert_offset; + args->local_num_experts = local_num_experts; + args->intermediate_size = intermediate_size; + args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); + args->do_finalize = do_finalize; + args->output = output.data_ptr(); + args->output_scale = nullptr; + + // Create and initialize launcher for this tile size + auto launcher = std::make_unique( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm1_bias, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + gemm2_weights, + gemm2_weights_scale, + gemm2_bias, + output1_scales_scalar, + output1_scales_gate_scalar, + output2_scales_scalar, + per_token_scales, + expert_indices, + expert_weights); + launcher->init( + std::move(args), + curr_tile_N, + routing_method_type, + /*use_shuffled_weight=*/true, + /*weight_layout=*/0, + static_cast(act_type), + mDtypeAct, + mDtypeWeights, + norm_topk_prob); + launcher->set_routing_replay_out(routing_replay_out); + + launchers_map[curr_tile_N] = std::move(launcher); + } + + auto const [tile_N, config] = + resolveMoeTileAndConfig(config_index, mSupportedTileN, num_tokens, top_k, local_num_experts); + + // Get the launcher for the selected tile_N + auto launcher_it = launchers_map.find(static_cast(tile_N)); + FLASHINFER_CHECK( + launcher_it != launchers_map.end(), "Internal error: missing FP4 block-scale MoE launcher for tile_N=", tile_N); + auto& selected_launcher = launcher_it->second; + + // Run the launcher - it will create its own runner internally + return selected_launcher->run(config, enable_pdl); +} + +// =========================================================================== +// NVFP4 MoE LoRA (decomposed / unfused-activation) — FP4 sibling of the FP8 +// trtllm-lora op. The standard NVFP4 path fuses SwiGLU into GEMM1, which leaves +// no seam to inject the gate_up LoRA delta pre-activation. We therefore run the +// MoE as a hand-wired pipeline that mirrors what MoE::Runner::run does for the +// DeepSeek-FP8 + per-token-NvFP4 path, but with the gate_up projection executed +// as a raw (no-activation) grouped GEMM via Gemm2::Runner so the standalone, +// LoRA-aware activation kernel can run between the two GEMMs: +// +// gather (permute bf16) -> NvFP4 quant -> gate_up GEMM (K=hidden, N=2*inter, +// raw bf16 out) -> activation (adds gate_up_lora_delta pre-SwiGLU, writes +// activation_lora_input) -> NvFP4 quant -> down GEMM (K=inter, N=hidden) -> +// finalize. +// +// The hidden states are supplied as bf16 (path 3: the dispatch feeds bf16 and this op permutes +// then NvFP4-quantizes internally, with globalScaleInv = 1/448/6 + per-token scaling, matching +// SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION). No fp4-input dequant round-trip. +// =========================================================================== + +// Decomposed NvFP4 MoE-LoRA launcher. Reuses FusedMoeLauncher's routing-phase +// workspace allocation/bookkeeping (via prepare_routing-style setup) but owns +// the MoE compute pipeline. +class FP4BlockScaleLoraLauncher { + public: + // Match the plain FP4 E2m1 path's tile ladder (FP4BlockScaleLauncher::getSupportedTileNums for + // non-bf16 act). Large prefills (high avg tokens/expert) need 128/256; capping at 64 makes + // selectDefaultTileN pick a tile too small for the Gemm2 cubin to have a valid config at that + // token count -> "Failed to initialize the TMA descriptor / illegal memory access". + static constexpr std::array mBaseSupportedTileNums = {8, 16, 32, 64, 128, 256}; + + static std::vector getSupportedTileNums() { + return std::vector(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); + } + + FP4BlockScaleLoraLauncher( + TensorView const& expert_indices, + TensorView const& expert_weights, + Optional const& routing_bias, + TensorView const& hidden_states, + Optional const& hidden_states_scale, + TensorView const& gemm1_weights, + TensorView const& gemm1_weights_scale, + TensorView const& gemm2_weights, + TensorView const& gemm2_weights_scale, + Optional const& output1_scales_scalar, + Optional const& output1_scales_gate_scalar, + Optional const& output2_scales_scalar, + TensorView const& gate_up_lora_delta, + TensorView const& activation_lora_input, + TensorView const& output, + int64_t lora_ready_event, + int64_t gemm2_done_event) + : expert_indices_(expert_indices), + expert_weights_(expert_weights), + 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), + output1_scales_scalar_(output1_scales_scalar), + output1_scales_gate_scalar_(output1_scales_gate_scalar), + output2_scales_scalar_(output2_scales_scalar), + gate_up_lora_delta_(gate_up_lora_delta), + activation_lora_input_(activation_lora_input), + output_(output), + lora_ready_event_(lora_ready_event), + gemm2_done_event_(gemm2_done_event) {} + + // Returns {output} when do_finalize, else {gemm2_output, expert_weights, + // expanded_idx_to_permuted_idx} for a downstream finalize kernel. + Array + run(int64_t num_experts, + int64_t top_k, + int64_t intermediate_size, + int64_t local_expert_offset, + int64_t local_num_experts, + double routed_scaling_factor, + int64_t routing_method_type, + int64_t tile_tokens_dim, + bool norm_topk_prob, + bool do_finalize, + bool enable_pdl, + bool use_fused_permute_quant) { + namespace moe_ns = tensorrt_llm::kernels::trtllmgen_moe; + auto device = hidden_states_.device(); + int dev_id = device.device_id; + cudaStream_t stream = get_stream(device); + + int64_t const num_tokens = hidden_states_.size(0); + int64_t const hidden_size = + hidden_states_.dtype() == dl_uint8 ? hidden_states_.size(1) * 2 : hidden_states_.size(1); + int64_t const inter = intermediate_size; + int64_t const gate_up_n = 2 * inter; // gated SwiGLU + + // ---- 1) routing (precomputed packed topk) ---- + Tensor num_tokens_per_expert = alloc_tensor({num_experts}, dl_int32, device); + int32_t max_num_padded_tokens = + moe_ns::Routing::getMaxPermutedPaddedCount(num_tokens, top_k, num_experts, tile_tokens_dim); + Tensor total_num_padded_tokens = alloc_tensor({1}, dl_int32, device); + Tensor expanded_idx_to_permuted_idx = alloc_tensor({num_tokens * top_k}, dl_int32, device); + Tensor permuted_idx_to_token_idx = alloc_tensor({max_num_padded_tokens}, dl_int32, device); + int64_t const hist_size = std::max(num_experts * 2, 256 * 2); + Tensor expert_count_histogram = alloc_tensor({hist_size}, dl_int32, device); + int32_t max_num_ctas = moe_ns::Routing::getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts, tile_tokens_dim); + Tensor cta_idx_xy_to_batch_idx = alloc_tensor({max_num_ctas}, dl_int32, device); + Tensor cta_idx_xy_to_mn_limit = alloc_tensor({max_num_ctas}, dl_int32, device); + Tensor num_non_exiting_ctas = alloc_tensor({1}, dl_int32, device); + + auto routing_bias_dtype = routing_bias_.has_value() ? routing_bias_.value().dtype() : dl_bfloat16; + btg::Dtype mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; + + // The wrapper passes an empty placeholder for expert_weights; the routing runner + // writes the unpacked per-(token,k) weights here. Allocate it ourselves (mirrors + // Fp8BlockScaleLauncher::prepare_routing when has_precomputed_weights is false). + // If the caller did pass a real expert_weights tensor, copy it into the allocation + // afterwards is unnecessary; we just compute into our own buffer for a clean Tensor + // return type. The bf16 routing-weight values are identical either way. + auto ew_dtype = mRoutingBiasDtype == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; + Tensor expert_weights_alloc = alloc_tensor({num_tokens, top_k}, ew_dtype, device); + void* expert_weights_ptr = expert_weights_alloc.data_ptr(); + + moe_ns::Routing::Runner routing_runner(tile_tokens_dim); + routing_runner.run( + /*routing_logits=*/nullptr, + routing_bias_.has_value() ? routing_bias_.value().data_ptr() : nullptr, + num_tokens, + num_experts, + top_k, + /*n_group=*/0, + /*topk_group=*/0, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + static_cast(const_cast(expert_indices_.data_ptr())), + static_cast(expert_count_histogram.data_ptr()), + static_cast(total_num_padded_tokens.data_ptr()), + static_cast(expanded_idx_to_permuted_idx.data_ptr()), + /*permuted_idx_to_expanded_idx=*/nullptr, + static_cast(permuted_idx_to_token_idx.data_ptr()), + expert_weights_ptr, + static_cast(num_tokens_per_expert.data_ptr()), + static_cast(cta_idx_xy_to_batch_idx.data_ptr()), + static_cast(cta_idx_xy_to_mn_limit.data_ptr()), + static_cast(num_non_exiting_ctas.data_ptr()), + btg::Dtype::Bfloat16, + mRoutingBiasDtype, + /*useRoutingScalesOnInput=*/false, + /*useDeepSeekFp8=*/false, + static_cast(routing_method_type), + stream, + btg::Dtype::Bfloat16, + norm_topk_prob, + /*routing_replay_out=*/nullptr); + + // ---- 2) hidden as bf16 (path 3: dispatch feeds bf16; the op quantizes internally) ---- + TVM_FFI_ICHECK(hidden_states_.dtype() == dl_bfloat16) + << "fp4 LoRA (path 3) requires bf16 hidden_states; the dispatch feeds bf16 and the op " + "permutes+NvFP4-quantizes internally (no python pre-quant / dequant round-trip)."; + void* hidden_bf16_ptr = hidden_states_.data_ptr(); + + int64_t const tile = tile_tokens_dim; + // gate_up GEMM act operand (permuted fp4 hidden + scales). Declared in run() scope because the + // gate_up GEMM (step 5) consumes them; the LARGE [max_padded, hidden] bf16 gather buffer + // (permuted_hidden_bf16, ~4 GB at a 32K-token prefill) lives only inside the block below, so it + // frees right after the quant -- before the equally-large [max_padded, hidden] gemm2_output is + // allocated (step 8), letting the caching allocator reuse its block (halves the op's peak). + auto gu_sfLayout = tile >= 128 ? tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4 + : tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4; + int64_t const hidden_sf_size = tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens, hidden_size / 16); + Tensor hidden_fp4 = alloc_tensor({max_num_padded_tokens, hidden_size / 2}, dl_uint8, device); + Tensor hidden_fp4_sf = alloc_tensor({hidden_sf_size}, dl_uint8, device); + Tensor hidden_per_token_sf = alloc_tensor({max_num_padded_tokens}, dl_float32, device); + if (use_fused_permute_quant && tile < 128) { + // Invariants the fused kernel relies on (review hardening): hidden must be a multiple of the + // 16-wide PackedVec load, and top_k must fit the dedup per-token-scale write (threadIdx