chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh

@@ -0,0 +1,453 @@
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, Panic, div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cfloat>
+#include <cstdint>
+
+namespace {
+
+// Kimi K2 MoE fused gate, supports NUM_EXPERTS in {256 (MiMo V2 Flash), 384 (Kimi K2)}.
+// Routing (DeepSeek "noaux_tc" with num_expert_group = 1):
+//   1. sigmoid(gate_logit)
+//   2. add per-expert correction bias (ranking only)
+//   3. pick top-k by biased score
+//   4. weights = sigmoid (no bias)
+//   5. optional renorm; routed_scaling_factor folded into renorm (no-op when not renormalizing)
+
+__device__ __forceinline__ float sigmoid_accurate(float x) {
+  return 1.0f / (1.0f + expf(-x));
+}
+
+// Scalar widening: input/bias may arrive as fp32, bf16, or fp16; the kernel math
+// always runs in fp32. Widening bf16/fp16 -> fp32 is exact, so results are
+// bitwise identical to upcasting on the host first (the casts we are removing).
+__device__ __forceinline__ float to_float(float x) {
+  return x;
+}
+__device__ __forceinline__ float to_float(__nv_bfloat16 x) {
+  return __bfloat162float(x);
+}
+__device__ __forceinline__ float to_float(__half x) {
+  return __half2float(x);
+}
+
+// Vectorized load of 4 consecutive elements of type T at vector index `vec_idx`,
+// widened to a float4. fp32 reads a 16B float4; bf16/fp16 read an 8B float2 and
+// expand. Used only by the large-token kernel's lane-strided loads.
+template <typename T>
+struct VecLoader;
+
+template <>
+struct VecLoader<float> {
+  __device__ __forceinline__ static float4 load(const float* base, int vec_idx) {
+    return reinterpret_cast<const float4*>(base)[vec_idx];
+  }
+};
+
+template <>
+struct VecLoader<__nv_bfloat16> {
+  __device__ __forceinline__ static float4 load(const __nv_bfloat16* base, int vec_idx) {
+    float2 raw = reinterpret_cast<const float2*>(base)[vec_idx];  // 4 bf16 = 8 bytes
+    const __nv_bfloat162* packed = reinterpret_cast<const __nv_bfloat162*>(&raw);
+    float2 lo = __bfloat1622float2(packed[0]);
+    float2 hi = __bfloat1622float2(packed[1]);
+    return make_float4(lo.x, lo.y, hi.x, hi.y);
+  }
+};
+
+template <>
+struct VecLoader<__half> {
+  __device__ __forceinline__ static float4 load(const __half* base, int vec_idx) {
+    float2 raw = reinterpret_cast<const float2*>(base)[vec_idx];  // 4 fp16 = 8 bytes
+    const __half2* packed = reinterpret_cast<const __half2*>(&raw);
+    float2 lo = __half22float2(packed[0]);
+    float2 hi = __half22float2(packed[1]);
+    return make_float4(lo.x, lo.y, hi.x, hi.y);
+  }
+};
+
+template <int N>
+struct GateConfig {
+  static_assert(
+      N == 256 || N == 384,
+      "kimi_k2_moe_fused_gate currently only supports "
+      "NUM_EXPERTS == 256 or 384");
+  static constexpr int NUM_EXPERTS = N;
+  static constexpr int WARP_SIZE = 32;
+  static constexpr int WARPS_PER_CTA = 6;  // only used by the large-token kernel
+  static constexpr int VPT = N / 32;       // 8 (256) or 12 (384)
+  static constexpr int VEC_SIZE = 4;
+  static constexpr int VEC_PER_LANE = VPT / VEC_SIZE;   // 2 or 3
+  static constexpr int WARPS_PER_TOKEN_SMALL = N / 32;  // 8 or 12
+  static constexpr int THREADS_PER_BLOCK_SMALL = N;     // 256 or 384
+  static constexpr int SMALL_TOKEN_THRESHOLD = 512;
+  static constexpr int MAX_TOPK = 8;  // must match RuntimeCheck(topk <= 8) at the host launcher
+  static_assert(VPT % VEC_SIZE == 0, "VPT must be a multiple of VEC_SIZE for the float4 vec load");
+};
+
+// Small-token kernel: 1 block per token, NUM_EXPERTS threads (1 thread = 1 expert).
+template <int N, typename InputT, typename BiasT>
+__global__ void kimi_k2_moe_fused_gate_kernel_small_token(
+    const InputT* input,
+    const BiasT* bias,
+    float* output_ptr,
+    int32_t* indices_ptr,
+    int64_t num_rows,
+    int64_t topk,
+    bool renormalize,
+    double routed_scaling_factor,
+    bool apply_routed_scaling_factor_on_output) {
+  using Cfg = GateConfig<N>;
+  constexpr int NUM_EXPERTS = Cfg::NUM_EXPERTS;
+  constexpr int WARP_SIZE = Cfg::WARP_SIZE;
+  constexpr int WARPS_PER_TOKEN_SMALL = Cfg::WARPS_PER_TOKEN_SMALL;
+  constexpr int MAX_TOPK = Cfg::MAX_TOPK;
+
+  int64_t row_idx = blockIdx.x;
+  if (row_idx >= num_rows) return;
+
+  int tid = threadIdx.x;
+  int warp_id = tid / WARP_SIZE;
+  int lane_id = tid % WARP_SIZE;
+
+  // Sigmoid weights (no bias) for final lookup, indexed by expert id.
+  __shared__ float shared_original_scores[NUM_EXPERTS];
+  __shared__ float warp_maxs[WARPS_PER_TOKEN_SMALL];
+  __shared__ int warp_experts[WARPS_PER_TOKEN_SMALL];
+  __shared__ int selected_experts[MAX_TOPK];
+
+  // Keep biased_val in register; mask the winner in-place each iteration to
+  // avoid round-tripping through shared memory.
+  float input_val = to_float(input[row_idx * NUM_EXPERTS + tid]);
+  float bias_val = to_float(bias[tid]);
+  float sigmoid_val = sigmoid_accurate(input_val);
+  float biased_val = sigmoid_val + bias_val;
+  shared_original_scores[tid] = sigmoid_val;
+
+  __syncthreads();
+
+  // Lane 0 of warp 0 accumulates the renorm sum as it picks each winner,
+  // saving a second pass over selected_experts during writeback.
+  float sum_for_renorm = 0.0f;
+
+  for (int k = 0; k < topk; k++) {
+    // Stage 1: per-warp argmax.
+    float warp_max_val = biased_val;
+    int warp_max_expert = tid;
+#pragma unroll
+    for (int offset = 16; offset > 0; offset /= 2) {
+      float other_val = __shfl_down_sync(0xFFFFFFFF, warp_max_val, offset);
+      int other_expert = __shfl_down_sync(0xFFFFFFFF, warp_max_expert, offset);
+      if (other_val > warp_max_val) {
+        warp_max_val = other_val;
+        warp_max_expert = other_expert;
+      }
+    }
+    if (lane_id == 0) {
+      warp_maxs[warp_id] = warp_max_val;
+      warp_experts[warp_id] = warp_max_expert;
+    }
+    __syncthreads();
+
+    // Stage 2: warp 0 merges warp-leaders into a single winner.
+    if (warp_id == 0) {
+      float final_max = (lane_id < WARPS_PER_TOKEN_SMALL) ? warp_maxs[lane_id] : -FLT_MAX;
+      int final_expert = (lane_id < WARPS_PER_TOKEN_SMALL) ? warp_experts[lane_id] : -1;
+#pragma unroll
+      for (int offset = 16; offset > 0; offset /= 2) {
+        float other_val = __shfl_down_sync(0xFFFFFFFF, final_max, offset);
+        int other_expert = __shfl_down_sync(0xFFFFFFFF, final_expert, offset);
+        if (other_val > final_max) {
+          final_max = other_val;
+          final_expert = other_expert;
+        }
+      }
+      if (lane_id == 0) {
+        selected_experts[k] = final_expert;
+        if (renormalize && final_expert >= 0 && final_expert < NUM_EXPERTS) {
+          sum_for_renorm += shared_original_scores[final_expert];
+        }
+      }
+    }
+    __syncthreads();
+
+    int selected = selected_experts[k];
+    if (tid == selected) biased_val = -FLT_MAX;
+  }
+
+  // Lane 0 of warp 0 writes the output. sum_for_renorm was accumulated
+  // during the topk loop, so we just fold it into rcp.
+  if (warp_id == 0 && lane_id == 0) {
+    float rcp = 1.0f;
+    if (renormalize && sum_for_renorm > 0.0f) {
+      rcp = 1.0f / sum_for_renorm;
+      if (apply_routed_scaling_factor_on_output) {
+        rcp *= static_cast<float>(routed_scaling_factor);
+      }
+    }
+
+    for (int k = 0; k < topk; k++) {
+      int expert_id = selected_experts[k];
+      bool valid = (expert_id >= 0 && expert_id < NUM_EXPERTS);
+      output_ptr[row_idx * topk + k] = valid ? shared_original_scores[expert_id] * rcp : 0.0f;
+      indices_ptr[row_idx * topk + k] = valid ? expert_id : 0;
+    }
+  }
+}
+
+// Large-token kernel: 1 warp per token, WARPS_PER_CTA warps per block.
+template <int N, typename InputT, typename BiasT>
+__global__ void kimi_k2_moe_fused_gate_kernel(
+    const InputT* input,
+    const BiasT* bias,
+    float* output_ptr,
+    int32_t* indices_ptr,
+    int64_t num_rows,
+    int64_t topk,
+    bool renormalize,
+    double routed_scaling_factor,
+    bool apply_routed_scaling_factor_on_output) {
+  using Cfg = GateConfig<N>;
+  constexpr int NUM_EXPERTS = Cfg::NUM_EXPERTS;
+  constexpr int WARP_SIZE = Cfg::WARP_SIZE;
+  constexpr int WARPS_PER_CTA = Cfg::WARPS_PER_CTA;
+  constexpr int VEC_SIZE = Cfg::VEC_SIZE;
+  constexpr int VEC_PER_LANE = Cfg::VEC_PER_LANE;
+  constexpr int MAX_TOPK = Cfg::MAX_TOPK;
+
+  int64_t row_idx = blockIdx.x * WARPS_PER_CTA + threadIdx.y;
+  if (row_idx >= num_rows) return;
+
+  int lane_id = threadIdx.x;
+  int warp_id = threadIdx.y;
+
+  __shared__ float shared_scores[NUM_EXPERTS * WARPS_PER_CTA];
+  __shared__ float shared_original_scores[NUM_EXPERTS * WARPS_PER_CTA];
+  float* warp_scores = shared_scores + warp_id * NUM_EXPERTS;
+  float* warp_original_scores = shared_original_scores + warp_id * NUM_EXPERTS;
+  float4* warp_scores_v4 = reinterpret_cast<float4*>(warp_scores);
+  float4* warp_original_scores_v4 = reinterpret_cast<float4*>(warp_original_scores);
+
+  const InputT* input_row = input + row_idx * NUM_EXPERTS;
+
+  // Lane-strided vec_idx (each lane k stores at vec_idx k, k+32, k+64, ...) so each
+  // iteration's STS.128 is lane-contiguous, avoiding shared-mem bank conflicts.
+#pragma unroll
+  for (int i = 0; i < VEC_PER_LANE; i++) {
+    int vec_idx = lane_id + i * WARP_SIZE;
+    float4 input_val = VecLoader<InputT>::load(input_row, vec_idx);
+    float4 bias_val = VecLoader<BiasT>::load(bias, vec_idx);
+
+    float4 sigmoid_v4;
+    float4 biased_v4;
+#pragma unroll
+    for (int j = 0; j < VEC_SIZE; j++) {
+      float inp = ((float*)&input_val)[j];
+      float b = ((float*)&bias_val)[j];
+      float sigmoid_val = sigmoid_accurate(inp);
+      ((float*)&sigmoid_v4)[j] = sigmoid_val;
+      ((float*)&biased_v4)[j] = sigmoid_val + b;
+    }
+    warp_original_scores_v4[vec_idx] = sigmoid_v4;
+    warp_scores_v4[vec_idx] = biased_v4;
+  }
+
+  __syncwarp();
+
+  // Lane 0 records the picked expert ids and accumulates the renorm sum as
+  // it goes; the global write is a single pass after the loop.
+  int top_indices[MAX_TOPK];
+  float sum_for_renorm = 0.0f;
+
+  for (int k = 0; k < topk; k++) {
+    float max_val = -FLT_MAX;
+    int max_expert = -1;
+
+    for (int expert = lane_id; expert < NUM_EXPERTS; expert += WARP_SIZE) {
+      if (warp_scores[expert] > max_val) {
+        max_val = warp_scores[expert];
+        max_expert = expert;
+      }
+    }
+
+    // warp shfl reduce; tie-break by lower expert id
+#pragma unroll
+    for (int offset = 16; offset > 0; offset /= 2) {
+      float other_val = __shfl_down_sync(0xFFFFFFFF, max_val, offset);
+      int other_expert = __shfl_down_sync(0xFFFFFFFF, max_expert, offset);
+      if (other_val > max_val || (other_val == max_val && other_expert < max_expert)) {
+        max_val = other_val;
+        max_expert = other_expert;
+      }
+    }
+
+    if (lane_id == 0) {
+      bool valid = (max_expert >= 0 && max_expert < NUM_EXPERTS);
+      top_indices[k] = valid ? max_expert : -1;
+      if (renormalize && valid) {
+        sum_for_renorm += warp_original_scores[max_expert];
+      }
+      if (valid) warp_scores[max_expert] = -FLT_MAX;
+    }
+    __syncwarp();
+  }
+
+  if (lane_id == 0) {
+    float rcp = 1.0f;
+    if (renormalize && sum_for_renorm > 0.0f) {
+      rcp = 1.0f / sum_for_renorm;
+      if (apply_routed_scaling_factor_on_output) {
+        rcp *= static_cast<float>(routed_scaling_factor);
+      }
+    }
+
+    for (int k = 0; k < topk; k++) {
+      int e = top_indices[k];
+      bool valid = (e >= 0);
+      output_ptr[row_idx * topk + k] = valid ? warp_original_scores[e] * rcp : 0.0f;
+      indices_ptr[row_idx * topk + k] = valid ? e : 0;
+    }
+  }
+}
+
+// Bundles the dtype-agnostic launch parameters so the templated dispatch below
+// only has to thread the typed input/bias pointers.
+struct GateLaunchArgs {
+  float* output;
+  int32_t* indices;
+  int64_t num_rows;
+  int64_t topk;
+  bool renormalize;
+  double routed_scaling_factor;
+  bool apply_routed_scaling_factor_on_output;
+  DLDevice device;
+};
+
+template <int N, typename InputT, typename BiasT>
+void launch_for_n(const InputT* input, const BiasT* bias, const GateLaunchArgs& args) {
+  using namespace host;
+  using Cfg = GateConfig<N>;
+  bool use_small_token_kernel = args.num_rows <= Cfg::SMALL_TOKEN_THRESHOLD;
+
+  if (use_small_token_kernel) {
+    LaunchKernel(
+        static_cast<uint32_t>(args.num_rows), static_cast<uint32_t>(Cfg::THREADS_PER_BLOCK_SMALL), args.device)(
+        kimi_k2_moe_fused_gate_kernel_small_token<N, InputT, BiasT>,
+        input,
+        bias,
+        args.output,
+        args.indices,
+        args.num_rows,
+        args.topk,
+        args.renormalize,
+        args.routed_scaling_factor,
+        args.apply_routed_scaling_factor_on_output);
+  } else {
+    uint32_t num_blocks = div_ceil(args.num_rows, static_cast<int64_t>(Cfg::WARPS_PER_CTA));
+    dim3 block_dim(Cfg::WARP_SIZE, Cfg::WARPS_PER_CTA);
+    LaunchKernel(num_blocks, block_dim, args.device)(
+        kimi_k2_moe_fused_gate_kernel<N, InputT, BiasT>,
+        input,
+        bias,
+        args.output,
+        args.indices,
+        args.num_rows,
+        args.topk,
+        args.renormalize,
+        args.routed_scaling_factor,
+        args.apply_routed_scaling_factor_on_output);
+  }
+}
+
+// input/bias each independently arrive as fp32, bf16, or fp16; widen both to
+// fp32 inside the kernel so the host no longer has to upcast. Dispatch is nested:
+// num_experts -> input dtype -> bias dtype.
+template <int N, typename InputT>
+void dispatch_bias(
+    const InputT* input, const void* bias, const host::SymbolicDType& bias_dtype, const GateLaunchArgs& args) {
+  using namespace host;
+  if (bias_dtype.is_type<float>()) {
+    launch_for_n<N, InputT, float>(input, static_cast<const float*>(bias), args);
+  } else if (bias_dtype.is_type<bf16_t>()) {
+    launch_for_n<N, InputT, bf16_t>(input, static_cast<const bf16_t*>(bias), args);
+  } else {
+    launch_for_n<N, InputT, fp16_t>(input, static_cast<const fp16_t*>(bias), args);
+  }
+}
+
+template <int N>
+void dispatch_input(
+    const void* input,
+    const host::SymbolicDType& input_dtype,
+    const void* bias,
+    const host::SymbolicDType& bias_dtype,
+    const GateLaunchArgs& args) {
+  using namespace host;
+  if (input_dtype.is_type<float>()) {
+    dispatch_bias<N, float>(static_cast<const float*>(input), bias, bias_dtype, args);
+  } else if (input_dtype.is_type<bf16_t>()) {
+    dispatch_bias<N, bf16_t>(static_cast<const bf16_t*>(input), bias, bias_dtype, args);
+  } else {
+    dispatch_bias<N, fp16_t>(static_cast<const fp16_t*>(input), bias, bias_dtype, args);
+  }
+}
+
+struct KimiK2MoEFusedGateKernel {
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView bias,
+      const tvm::ffi::TensorView output,
+      const tvm::ffi::TensorView indices,
+      int64_t topk,
+      bool renormalize,
+      double routed_scaling_factor,
+      bool apply_routed_scaling_factor_on_output) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_rows"};
+    auto E = SymbolicSize{"num_experts"};
+    auto K = SymbolicSize{"topk"};
+    auto input_dtype = SymbolicDType{};
+    auto bias_dtype = SymbolicDType{};
+    auto device = SymbolicDevice{};
+    K.set_value(topk);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, E}).with_dtype<float, bf16_t, fp16_t>(input_dtype).with_device(device).verify(input);
+    TensorMatcher({E}).with_dtype<float, bf16_t, fp16_t>(bias_dtype).with_device(device).verify(bias);
+    TensorMatcher({N, K}).with_dtype<float>().with_device(device).verify(output);
+    TensorMatcher({N, K}).with_dtype<int32_t>().with_device(device).verify(indices);
+
+    const auto num_rows = static_cast<int64_t>(N.unwrap());
+    const auto num_experts = static_cast<int64_t>(E.unwrap());
+
+    RuntimeCheck(topk <= 8, "kimi_k2_moe_fused_gate only supports topk <= 8, got ", topk);
+
+    const GateLaunchArgs args{
+        .output = static_cast<float*>(output.data_ptr()),
+        .indices = static_cast<int32_t*>(indices.data_ptr()),
+        .num_rows = num_rows,
+        .topk = topk,
+        .renormalize = renormalize,
+        .routed_scaling_factor = routed_scaling_factor,
+        .apply_routed_scaling_factor_on_output = apply_routed_scaling_factor_on_output,
+        .device = device.unwrap()};
+
+    switch (num_experts) {
+      case 256:
+        dispatch_input<256>(input.data_ptr(), input_dtype, bias.data_ptr(), bias_dtype, args);
+        break;
+      case 384:
+        dispatch_input<384>(input.data_ptr(), input_dtype, bias.data_ptr(), bias_dtype, args);
+        break;
+      default:
+        Panic("kimi_k2_moe_fused_gate only supports num_experts in {256, 384}, got ", num_experts);
+    }
+  }
+};
+
+}  // namespace

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 <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>
+
+#ifndef WARP_SIZE
+#define WARP_SIZE 32
+#endif
+
+#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
+
+#define VEC_SIZE 4
+using Vec = int4;
+
+inline uint32_t next_pow2(uint32_t x) noexcept {
+  --x;
+  x |= x >> 1;
+  x |= x >> 2;
+  x |= x >> 4;
+  x |= x >> 8;
+  x |= x >> 16;
+  return x + 1;
+}
+
+namespace moe_lora_merged {
+
+__device__ __forceinline__ int warp_exclusive_scan(int v, unsigned mask = 0xffffffffu) {
+  int original = v;
+#pragma unroll
+  for (int offset = 1; offset < WARP_SIZE; offset <<= 1) {
+    int n = __shfl_up_sync(mask, v, offset);
+    if ((threadIdx.x & (WARP_SIZE - 1)) >= offset) v += n;
+  }
+  return v - original;
+}
+
+// Inline mirror of _fused_virtual_topk_ids_kernel (virtual_experts.py). Returns
+// the merged virtual expert id for flat slot `i` (range [-1, virtual_num_experts);
+// -1 is the dropped/masked sentinel). The caller adds +1 to get the histogram
+// bucket (sentinel -> bucket 0), matching the native +1 offset convention.
+template <typename scalar_t>
+__device__ __forceinline__ int compute_virtual_id(
+    const scalar_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ token_lora_mapping,
+    size_t i,
+    int top_k,
+    int num_experts_for_weight,
+    int local_expert_offset,
+    int local_num_experts,
+    bool ep_local,
+    bool shared_outer,
+    bool compact) {
+  int m = static_cast<int>(i) / top_k;
+  int lora_id = token_lora_mapping[m];
+  bool mask_val = lora_id >= 0;
+  int safe_lora = lora_id > 0 ? lora_id : 0;
+
+  int base = shared_outer ? 0 : static_cast<int>(topk_ids[i]);
+  if (ep_local) {
+    bool owned = base >= local_expert_offset && base < local_expert_offset + local_num_experts;
+    base = owned ? base : -1;
+  }
+  if (!mask_val || base < 0) return -1;
+  // compact: dense LOCAL expert id in [0, local_num_experts) so the histogram
+  // spans only local_num_experts buckets instead of the full global virtual
+  // space (337/385 empty under EP). Assumes max_loras==1 (safe_lora shift is 0;
+  // the wrapper guards). expert_ids is converted back to global at write time.
+  if (compact) return base - local_expert_offset;
+  return base + safe_lora * num_experts_for_weight;
+}
+
+template <typename scalar_t>
+__global__ void count_and_sort_expert_tokens_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ token_lora_mapping,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ cumsum_buffer,
+    size_t numel,
+    int top_k,
+    int num_experts_for_weight,
+    int local_expert_offset,
+    int local_num_experts,
+    bool ep_local,
+    bool shared_outer,
+    bool do_skip,
+    bool compact) {
+  const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int vid = compute_virtual_id<scalar_t>(
+        topk_ids,
+        token_lora_mapping,
+        i,
+        top_k,
+        num_experts_for_weight,
+        local_expert_offset,
+        local_num_experts,
+        ep_local,
+        shared_outer,
+        compact);
+    // EP skip: dropped/masked slots (vid < 0) produce no delta on this rank, so
+    // they never need a slot in sorted_token_ids -> skip the global atomicAdd
+    // (kills the sentinel-bucket-0 contention). When do_skip is off they fall
+    // into bucket 0 (old behavior, kept for the bitwise-equivalence guardrail).
+    if (do_skip && vid < 0) continue;
+    int32_t expert_id = vid + 1;
+    int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = i;
+  }
+}
+
+template <typename scalar_t>
+__global__ void moe_align_block_size_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ token_lora_mapping,
+    bool* __restrict__ token_lora_mask,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t num_experts,
+    int32_t block_size,
+    size_t numel,
+    int32_t* __restrict__ cumsum,
+    bool pad_sorted_token_ids,
+    const int32_t scan_size,
+    int32_t max_num_tokens_padded,
+    int top_k,
+    int num_experts_for_weight,
+    int local_expert_offset,
+    int local_num_experts,
+    bool ep_local,
+    bool shared_outer,
+    bool do_skip,
+    bool compact) {
+  // Use a separate thread block to populate sorted_token_ids
+  if (blockIdx.x == 1) {
+    if (pad_sorted_token_ids) {
+      Vec fill_vec;
+      fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
+      int32_t total_vecs = (max_num_tokens_padded + VEC_SIZE - 1) / VEC_SIZE;
+      Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
+      for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
+        out_ptr[i] = fill_vec;
+      }
+    }
+    return;
+  }
+
+  extern __shared__ int32_t smem[];
+  int32_t* shared_counts = smem;                  // [num_experts]
+  int32_t* prefix = shared_counts + num_experts;  // [num_experts + 1]
+  int32_t* scan_buf = prefix + num_experts + 1;   // [scan_size]
+  __shared__ int32_t s_total_tokens_post_pad;
+
+  const size_t tid = threadIdx.x;
+  const size_t stride = blockDim.x;
+
+  if (tid < num_experts) {
+    shared_counts[tid] = 0;
+  }
+
+  __syncthreads();
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int vid = compute_virtual_id<scalar_t>(
+        topk_ids,
+        token_lora_mapping,
+        i,
+        top_k,
+        num_experts_for_weight,
+        local_expert_offset,
+        local_num_experts,
+        ep_local,
+        shared_outer,
+        compact);
+    // EP skip: dropped/masked slots don't increment any bucket (sentinel bucket
+    // 0 stays empty), so they never get a block and never reach count_and_sort.
+    if (!do_skip || vid >= 0) {
+      atomicAdd(&shared_counts[vid + 1], 1);
+    }
+    // token_lora_mask[m] = token_lora_mapping[m] >= 0, written once per row.
+    if (static_cast<int>(i) % top_k == 0) {
+      int m = static_cast<int>(i) / top_k;
+      token_lora_mask[m] = token_lora_mapping[m] >= 0;
+    }
+  }
+
+  __syncthreads();
+
+  int32_t padded_count = 0;
+  if (tid < num_experts) {
+    int32_t count = shared_counts[tid];
+    padded_count = (count + block_size - 1) / block_size * block_size;
+    scan_buf[tid] = padded_count;
+  }
+
+  // Intra warp prefix sum
+  int32_t* warp_sums = scan_buf + scan_size;  // [<= 32]
+  const int warp_id = tid / WARP_SIZE;
+  const int lane_id = tid & (WARP_SIZE - 1);
+  const int num_warps_for_scan = (scan_size + WARP_SIZE - 1) / WARP_SIZE;
+  const int warp_sum = warp_exclusive_scan(padded_count) + padded_count;
+  if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = warp_sum;
+  __syncthreads();
+
+  // warp0 accumulate all the block's prefix sum
+  if (tid < WARP_SIZE) {
+    int val = (tid < num_warps_for_scan) ? warp_sums[tid] : 0;
+    int incl = warp_exclusive_scan(val) + val;
+    warp_sums[tid] = incl;
+  }
+  __syncthreads();
+
+  // Every thread obtains the whole block's sum
+  if (tid == 0) {
+    prefix[num_experts] = warp_sums[num_warps_for_scan - 1];
+    s_total_tokens_post_pad = prefix[num_experts];
+    *total_tokens_post_pad = s_total_tokens_post_pad;
+  }
+  __syncthreads();
+
+  // Fill 0 to scan_buf extended area (tid >= num_expert)
+  if (tid >= num_experts && tid < scan_size) scan_buf[tid] = 0;
+  __syncthreads();
+
+  // Perform 2 level exclusive-prefix-sum to scan_buf
+  int v = (tid < scan_size) ? scan_buf[tid] : 0;
+  int pre = warp_exclusive_scan(v);
+  if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = pre + v;
+  __syncthreads();
+
+  if (warp_id == 0) {
+    int val = (lane_id < num_warps_for_scan) ? warp_sums[lane_id] : 0;
+    warp_sums[lane_id] = warp_exclusive_scan(val);
+  }
+  __syncthreads();
+
+  int offset = warp_sums[warp_id];
+  if (tid < scan_size) scan_buf[tid] = pre + offset;
+  __syncthreads();
+
+  // Write prefix[0..num_experts - 1] and cumsum
+  if (tid < num_experts) prefix[tid] = scan_buf[tid];
+
+  if (tid <= num_experts) {
+    cumsum[tid] = prefix[tid];
+  }
+  // fill expert_ids
+  const int32_t num_blocks = s_total_tokens_post_pad / block_size;
+  for (int32_t i = tid; i < num_blocks; i += stride) {
+    int32_t block_start = i * block_size;
+    int left = 0, right = num_experts;
+    while (left < right) {
+      int mid = (left + right) >> 1;
+      if (prefix[mid] <= block_start) {
+        left = mid + 1;
+      } else {
+        right = mid;
+      }
+    }
+    // compact buckets hold LOCAL expert ids; restore the global id (+offset) so
+    // the downstream GEMM still indexes the global contiguous LoRA weight.
+    expert_ids[i] = left - 2 + (compact ? local_expert_offset : 0);
+  }
+}
+
+// Single-block fused variant: does fill + histogram + scan + expert_ids + scatter
+// in ONE threadblock (one launch), eliminating the separate count_and_sort kernel
+// AND its redundant re-computation of the virtual id (cached in shared `svids`).
+// Only valid for small numel (the scatter is single-block); the wrapper routes
+// large numel (prefill) to the 2-kernel path. do_skip is implied (this is the
+// decode hot path); dropped slots are simply never scattered.
+template <typename scalar_t>
+__global__ void fused_align_scatter_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ token_lora_mapping,
+    bool* __restrict__ token_lora_mask,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t num_experts,
+    int32_t block_size,
+    size_t numel,
+    int32_t* __restrict__ cumsum,
+    const int32_t scan_size,
+    int32_t max_num_tokens_padded,
+    int top_k,
+    int num_experts_for_weight,
+    int local_expert_offset,
+    int local_num_experts,
+    bool ep_local,
+    bool shared_outer,
+    bool do_skip,
+    bool compact) {
+  extern __shared__ int32_t smem[];
+  int32_t* shared_counts = smem;                  // [num_experts]
+  int32_t* prefix = shared_counts + num_experts;  // [num_experts + 1]
+  int32_t* scan_buf = prefix + num_experts + 1;   // [scan_size]
+  int32_t* warp_sums = scan_buf + scan_size;      // [WARP_SIZE]
+  int32_t* cursor = warp_sums + WARP_SIZE;        // [num_experts] scatter cursor
+  int32_t* svids = cursor + num_experts;          // [numel] cached virtual ids
+  __shared__ int32_t s_total_tokens_post_pad;
+
+  const size_t tid = threadIdx.x;
+  const size_t stride = blockDim.x;
+  const int warp_id = tid / WARP_SIZE;
+  const int lane_id = tid & (WARP_SIZE - 1);
+  const int num_warps_for_scan = (scan_size + WARP_SIZE - 1) / WARP_SIZE;
+
+  // Phase 1: fill sorted_token_ids with the `numel` padding sentinel.
+  {
+    Vec fill_vec;
+    fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
+    int32_t total_vecs = (max_num_tokens_padded + VEC_SIZE - 1) / VEC_SIZE;
+    Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
+    for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
+      out_ptr[i] = fill_vec;
+    }
+  }
+  if (tid < num_experts) shared_counts[tid] = 0;
+  __syncthreads();
+
+  // Phase 2: histogram + cache the virtual id per slot + token_lora_mask.
+  for (size_t i = tid; i < numel; i += stride) {
+    int vid = compute_virtual_id<scalar_t>(
+        topk_ids,
+        token_lora_mapping,
+        i,
+        top_k,
+        num_experts_for_weight,
+        local_expert_offset,
+        local_num_experts,
+        ep_local,
+        shared_outer,
+        compact);
+    svids[i] = vid;
+    if (!do_skip || vid >= 0) {
+      atomicAdd(&shared_counts[vid + 1], 1);
+    }
+    if (static_cast<int>(i) % top_k == 0) {
+      int m = static_cast<int>(i) / top_k;
+      token_lora_mask[m] = token_lora_mapping[m] >= 0;
+    }
+  }
+  __syncthreads();
+
+  // Phase 3: padded counts + two-level warp exclusive prefix sum (verbatim).
+  int32_t padded_count = 0;
+  if (tid < num_experts) {
+    int32_t count = shared_counts[tid];
+    padded_count = (count + block_size - 1) / block_size * block_size;
+    scan_buf[tid] = padded_count;
+  }
+  const int warp_sum = warp_exclusive_scan(padded_count) + padded_count;
+  if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = warp_sum;
+  __syncthreads();
+  if (tid < WARP_SIZE) {
+    int val = (tid < num_warps_for_scan) ? warp_sums[tid] : 0;
+    int incl = warp_exclusive_scan(val) + val;
+    warp_sums[tid] = incl;
+  }
+  __syncthreads();
+  if (tid == 0) {
+    prefix[num_experts] = warp_sums[num_warps_for_scan - 1];
+    s_total_tokens_post_pad = prefix[num_experts];
+    *total_tokens_post_pad = s_total_tokens_post_pad;
+  }
+  __syncthreads();
+  if (tid >= num_experts && tid < scan_size) scan_buf[tid] = 0;
+  __syncthreads();
+  int v = (tid < scan_size) ? scan_buf[tid] : 0;
+  int pre = warp_exclusive_scan(v);
+  if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = pre + v;
+  __syncthreads();
+  if (warp_id == 0) {
+    int val = (lane_id < num_warps_for_scan) ? warp_sums[lane_id] : 0;
+    warp_sums[lane_id] = warp_exclusive_scan(val);
+  }
+  __syncthreads();
+  int off = warp_sums[warp_id];
+  if (tid < scan_size) scan_buf[tid] = pre + off;
+  __syncthreads();
+  if (tid < num_experts) prefix[tid] = scan_buf[tid];
+  if (tid <= num_experts) cumsum[tid] = prefix[tid];
+  __syncthreads();
+
+  // Phase 4: expert_ids (binary search per block) + init the scatter cursor.
+  const int32_t num_blocks = s_total_tokens_post_pad / block_size;
+  for (int32_t i = tid; i < num_blocks; i += stride) {
+    int32_t block_start = i * block_size;
+    int left = 0, right = num_experts;
+    while (left < right) {
+      int mid = (left + right) >> 1;
+      if (prefix[mid] <= block_start) {
+        left = mid + 1;
+      } else {
+        right = mid;
+      }
+    }
+    expert_ids[i] = left - 2 + (compact ? local_expert_offset : 0);
+  }
+  if (tid < num_experts) cursor[tid] = prefix[tid];
+  __syncthreads();
+
+  // Phase 5: scatter owned tokens using the cached virtual ids + shared cursor.
+  for (size_t i = tid; i < numel; i += stride) {
+    int vid = svids[i];
+    if (do_skip && vid < 0) continue;
+    int bucket = vid + 1;
+    int pos = atomicAdd(&cursor[bucket], 1);
+    sorted_token_ids[pos] = i;
+  }
+}
+
+}  // namespace moe_lora_merged
+
+namespace {
+
+template <typename scalar_t>
+struct MoeLoraMergedAlignKernel {
+  static void
+  run(tvm::ffi::TensorView topk_ids,
+      tvm::ffi::TensorView token_lora_mapping,
+      tvm::ffi::TensorView token_lora_mask,
+      int64_t num_experts,
+      int64_t block_size,
+      tvm::ffi::TensorView sorted_token_ids,
+      tvm::ffi::TensorView expert_ids,
+      tvm::ffi::TensorView num_tokens_post_pad,
+      tvm::ffi::TensorView cumsum_buffer,
+      bool pad_sorted_token_ids,
+      int64_t top_k,
+      int64_t num_experts_for_weight,
+      int64_t local_expert_offset,
+      int64_t local_num_experts,
+      bool ep_local,
+      bool shared_outer,
+      bool do_skip,
+      bool compact,
+      bool fuse_scatter) {
+    using namespace host;
+
+    auto device = topk_ids.device();
+    const cudaStream_t stream = LaunchKernel::resolve_device(device);
+
+    int threads = 1024;
+    threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+
+    int64_t max_num_tokens_padded = sorted_token_ids.size(0);
+
+    // num_experts here is the bucket count. Non-compact: virtual_num_experts+1
+    // (typically 385). Compact: local_num_experts+1 (typically 49). Both use the
+    // same single-block align path (valid for any bucket count <= 1024 that fits
+    // shared memory); the v2 (>1024) regime keeps the old path via the wrapper.
+    RuntimeCheck(
+        num_experts <= 1024, "moe_lora_merged_align: num_experts (bucket count) must be <= 1024, got ", num_experts);
+    // compact buckets hold LOCAL ids and restore global expert ids as
+    // (left-2+offset). For the sentinel bucket 0 that yields (offset-1), NOT the
+    // -1 the GEMM expects to skip -- only safe when do_skip empties bucket 0.
+    RuntimeCheck(
+        !compact || do_skip, "moe_lora_merged_align: compact requires do_skip (sentinel bucket must be empty)");
+
+    const scalar_t* topk_ids_ptr = static_cast<const scalar_t*>(topk_ids.data_ptr());
+    const int32_t* tlm_ptr = static_cast<const int32_t*>(token_lora_mapping.data_ptr());
+    bool* token_lora_mask_ptr = static_cast<bool*>(token_lora_mask.data_ptr());
+    int32_t* sorted_token_ids_ptr = static_cast<int32_t*>(sorted_token_ids.data_ptr());
+    int32_t* expert_ids_ptr = static_cast<int32_t*>(expert_ids.data_ptr());
+    int32_t* num_tokens_post_pad_ptr = static_cast<int32_t*>(num_tokens_post_pad.data_ptr());
+    int32_t* cumsum_buffer_ptr = static_cast<int32_t*>(cumsum_buffer.data_ptr());
+    size_t numel = topk_ids.numel();
+
+    const size_t scan_size = next_pow2(num_experts);
+
+    if (fuse_scatter) {
+      // One block does fill + histogram + scan + expert_ids + scatter. Extra
+      // shared for the scatter cursor [num_experts] and cached virtual ids [numel].
+      const size_t shmem =
+          (num_experts + (num_experts + 1) + scan_size + WARP_SIZE + num_experts + numel) * sizeof(int32_t);
+      auto fused = moe_lora_merged::fused_align_scatter_kernel<scalar_t>;
+      LaunchKernel(dim3(1), dim3(threads), stream, shmem)(
+          fused,
+          topk_ids_ptr,
+          tlm_ptr,
+          token_lora_mask_ptr,
+          sorted_token_ids_ptr,
+          expert_ids_ptr,
+          num_tokens_post_pad_ptr,
+          (int32_t)num_experts,
+          (int32_t)block_size,
+          numel,
+          cumsum_buffer_ptr,
+          (int32_t)scan_size,
+          (int32_t)max_num_tokens_padded,
+          (int)top_k,
+          (int)num_experts_for_weight,
+          (int)local_expert_offset,
+          (int)local_num_experts,
+          ep_local,
+          shared_outer,
+          do_skip,
+          compact);
+      return;
+    }
+
+    const size_t shared_mem_size = (num_experts + (num_experts + 1) + scan_size + WARP_SIZE) * sizeof(int32_t);
+
+    auto align_kernel = moe_lora_merged::moe_align_block_size_kernel<scalar_t>;
+    LaunchKernel(dim3(2), dim3(threads), stream, shared_mem_size)(
+        align_kernel,
+        topk_ids_ptr,
+        tlm_ptr,
+        token_lora_mask_ptr,
+        sorted_token_ids_ptr,
+        expert_ids_ptr,
+        num_tokens_post_pad_ptr,
+        (int32_t)num_experts,
+        (int32_t)block_size,
+        numel,
+        cumsum_buffer_ptr,
+        pad_sorted_token_ids,
+        (int32_t)scan_size,
+        (int32_t)max_num_tokens_padded,
+        (int)top_k,
+        (int)num_experts_for_weight,
+        (int)local_expert_offset,
+        (int)local_num_experts,
+        ep_local,
+        shared_outer,
+        do_skip,
+        compact);
+
+    const int block_threads = std::min(256, threads);
+    const int num_blocks = (numel + block_threads - 1) / block_threads;
+    const int max_blocks = 65535;
+    const int actual_blocks = std::min(num_blocks, max_blocks);
+
+    auto sort_kernel = moe_lora_merged::count_and_sort_expert_tokens_kernel<scalar_t>;
+    LaunchKernel(dim3(actual_blocks), dim3(block_threads), stream)(
+        sort_kernel,
+        topk_ids_ptr,
+        tlm_ptr,
+        sorted_token_ids_ptr,
+        cumsum_buffer_ptr,
+        numel,
+        (int)top_k,
+        (int)num_experts_for_weight,
+        (int)local_expert_offset,
+        (int)local_num_experts,
+        ep_local,
+        shared_outer,
+        do_skip,
+        compact);
+  }
+};
+
+}  // namespace

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 <sgl_kernel/tensor.h>  // TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // LaunchKernel, fp32_t/fp16_t/bf16_t, is_type
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cfloat>
+#include <cstdint>
+
+namespace {
+
+static constexpr int WARP_SIZE = 32;
+
+#define TSP_MAX(a, b) ((a) > (b) ? (a) : (b))
+#define TSP_MIN(a, b) ((a) < (b) ? (a) : (b))
+
+/// Aligned array type (mirrors the AOT kernel's CUTLASS-free aligned array)
+template <typename T, int N, int Alignment = sizeof(T) * N>
+class alignas(Alignment) AlignedArray {
+  T data[N];
+};
+
+template <typename T>
+__device__ float convert_to_float(T x) {
+  if constexpr (std::is_same_v<T, __half>) {
+    return __half2float(x);
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+    return __bfloat162float(x);
+  } else if constexpr (std::is_same_v<T, float>) {
+    return x;
+  } else {
+    return static_cast<float>(x);
+  }
+}
+
+// Reference pack (bit-identical to jit_kernel/flashinfer_trtllm_moe/topk_pack.py):
+// low 16 bits = bf16(weight) bits (round-to-nearest-even, same as torch/triton
+// `.to(bfloat16)`), high 16 bits = int16 expert id.
+__device__ __forceinline__ int32_t pack_routed(int32_t id, float w) {
+  const uint32_t wbits = static_cast<uint32_t>(__bfloat16_as_ushort(__float2bfloat16(w)));
+  return static_cast<int32_t>((static_cast<uint32_t>(id) << 16) | wbits);
+}
+
+template <typename T, int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmaxPack(
+    const T* input,
+    float* output,
+    const int num_rows,
+    int* indices,
+    int* packed_output,
+    const int32_t* num_token_non_padded,
+    const int k,
+    const bool renormalize) {
+  static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
+  static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS), "NUM_EXPERTS must be power of 2");
+  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2");
+  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+  static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
+  static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
+  static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size");
+
+  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp");
+
+  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row = warp_base_row + thread_row_in_warp;
+  if (thread_row >= num_rows) {
+    return;
+  }
+
+  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+  using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+
+  T row_chunk_temp[VPT];
+  AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk_temp);
+  const AccessType* vec_thread_read_ptr = reinterpret_cast<const AccessType*>(thread_read_ptr);
+#pragma unroll
+  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+  }
+
+  float row_chunk[VPT];
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = convert_to_float<T>(row_chunk_temp[ii]);
+  }
+
+  float thread_max = row_chunk[0];
+#pragma unroll
+  for (int ii = 1; ii < VPT; ++ii) {
+    thread_max = max(thread_max, row_chunk[ii]);
+  }
+
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    thread_max = max(thread_max, __shfl_xor_sync(0xffffffffu, thread_max, mask, THREADS_PER_ROW));
+  }
+
+  float row_sum = 0;
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = expf(row_chunk[ii] - thread_max);
+    row_sum += row_chunk[ii];
+  }
+
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    row_sum += __shfl_xor_sync(0xffffffffu, row_sum, mask, THREADS_PER_ROW);
+  }
+
+  const float reciprocal_row_sum = 1.f / row_sum;
+
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+  }
+
+  int start_col = first_elt_read_by_thread;
+  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+  float row_sum_for_renormalize = 0;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    float max_val = row_chunk[0];
+    int expert = start_col;
+#pragma unroll
+    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG) {
+#pragma unroll
+      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+        if (val > max_val) {
+          max_val = val;
+          expert = col + ii;
+        }
+      }
+    }
+
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+      float other_max = __shfl_xor_sync(0xffffffffu, max_val, mask, THREADS_PER_ROW);
+      int other_expert = __shfl_xor_sync(0xffffffffu, expert, mask, THREADS_PER_ROW);
+      if (other_max > max_val || (other_max == max_val && other_expert < expert)) {
+        max_val = other_max;
+        expert = other_expert;
+      }
+    }
+
+    if (thread_group_idx == 0) {
+      const int idx = k * thread_row + k_idx;
+      output[idx] = max_val;
+      indices[idx] = expert;
+      row_sum_for_renormalize += max_val;
+    }
+
+    if (k_idx + 1 < k) {
+      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+      if (thread_group_idx == thread_to_clear_in_group) {
+        const int offset_for_expert = expert % ELTS_PER_LDG;
+        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
+      }
+    }
+  }
+
+  if (thread_group_idx == 0) {
+    // Fused renormalization (same as the AOT kernel).
+    if (renormalize) {
+      float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+#pragma unroll
+      for (int k_idx = 0; k_idx < k; ++k_idx) {
+        const int idx = k * thread_row + k_idx;
+        output[idx] = output[idx] * row_sum_for_renormalize_inv;
+      }
+    }
+    // Fused routed pack: pack the FINAL (post-renorm) weights. Padded rows
+    // (>= *num_token_non_padded) pack id = -1, mirroring the in-place
+    // `_mask_topk_ids_padded_region` sentinel that the separate pack kernel
+    // would otherwise observe. The plain `indices` output is left unmasked
+    // here exactly like the AOT kernel — the existing python post-process
+    // masks it afterwards; only the packed tensor needs the mask baked in
+    // because it is produced BEFORE that post-process runs.
+    const bool row_padded = (num_token_non_padded != nullptr) && (thread_row >= *num_token_non_padded);
+#pragma unroll
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * thread_row + k_idx;
+      const int32_t id = row_padded ? -1 : indices[idx];
+      packed_output[idx] = pack_routed(id, output[idx]);
+    }
+  }
+}
+
+namespace detail {
+template <typename T, int EXPERTS, int BYTES_PER_LDG>
+struct TopkConstants {
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, "");
+  static constexpr int VECs_PER_THREAD = TSP_MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+  static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+};
+}  // namespace detail
+
+template <typename T, int EXPERTS, int WARPS_PER_TB>
+void launchTopkGatingSoftmaxPack(
+    const T* input,
+    float* output,
+    int* indices,
+    int* packed_output,
+    const int32_t* num_token_non_padded,
+    const int num_rows,
+    const int k,
+    const bool renormalize,
+    DLDevice device) {
+  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+  static constexpr int BYTES_PER_LDG = TSP_MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = detail::TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int VPT = Constants::VPT;
+  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+  dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+  host::LaunchKernel(dim3(num_blocks), block_dim, device)(
+      topkGatingSoftmaxPack<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>,
+      input,
+      output,
+      num_rows,
+      indices,
+      packed_output,
+      num_token_non_padded,
+      k,
+      renormalize);
+}
+
+template <typename T>
+void dispatchExperts(
+    const T* input,
+    float* output,
+    int* indices,
+    int* packed_output,
+    const int32_t* num_token_non_padded,
+    const int num_rows,
+    const int num_experts,
+    const int k,
+    const bool renormalize,
+    DLDevice device) {
+  static constexpr int WARPS_PER_TB = 4;
+#define TSP_LAUNCH(E)                              \
+  launchTopkGatingSoftmaxPack<T, E, WARPS_PER_TB>( \
+      input, output, indices, packed_output, num_token_non_padded, num_rows, k, renormalize, device)
+  switch (num_experts) {
+    case 1:
+      TSP_LAUNCH(1);
+      break;
+    case 2:
+      TSP_LAUNCH(2);
+      break;
+    case 4:
+      TSP_LAUNCH(4);
+      break;
+    case 8:
+      TSP_LAUNCH(8);
+      break;
+    case 16:
+      TSP_LAUNCH(16);
+      break;
+    case 32:
+      TSP_LAUNCH(32);
+      break;
+    case 64:
+      TSP_LAUNCH(64);
+      break;
+    case 128:
+      TSP_LAUNCH(128);
+      break;
+    case 256:
+      TSP_LAUNCH(256);
+      break;
+    case 512:
+      TSP_LAUNCH(512);
+      break;
+    default:
+      host::RuntimeCheck(false, "topk_softmax_pack: num_experts must be a power of 2 in [1, 512], got ", num_experts);
+  }
+#undef TSP_LAUNCH
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Launcher
+// ─────────────────────────────────────────────────────────────────────────────
+void topk_softmax_pack(
+    tvm::ffi::TensorView topk_weights,
+    tvm::ffi::TensorView topk_indices,
+    tvm::ffi::TensorView packed,
+    tvm::ffi::TensorView gating_output,
+    tvm::ffi::Optional<tvm::ffi::TensorView> num_token_non_padded,
+    bool renormalize) {
+  using namespace host;
+
+  SymbolicSize N{"num_tokens"};
+  SymbolicSize E{"num_experts"};
+  SymbolicSize K{"topk"};
+  SymbolicDevice device_;
+  device_.set_options<kDLCUDA>();
+
+  TensorMatcher({N, E}).with_dtype<fp32_t, fp16_t, bf16_t>().with_device<kDLCUDA>(device_).verify(gating_output);
+  TensorMatcher({N, K}).with_dtype<fp32_t>().with_device<kDLCUDA>(device_).verify(topk_weights);
+  TensorMatcher({N, K}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(topk_indices);
+  TensorMatcher({N, K}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(packed);
+
+  const int32_t* ntnp_ptr = nullptr;
+  if (num_token_non_padded.has_value()) {
+    SymbolicSize One{"ntnp_numel"};
+    TensorMatcher({One}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(num_token_non_padded.value());
+    RuntimeCheck(One.unwrap() == 1, "num_token_non_padded must be a 1-element tensor");
+    ntnp_ptr = static_cast<const int32_t*>(num_token_non_padded.value().data_ptr());
+  }
+
+  const int num_tokens = static_cast<int>(N.unwrap());
+  const int num_experts = static_cast<int>(E.unwrap());
+  const int topk = static_cast<int>(K.unwrap());
+  DLDevice device = device_.unwrap();
+
+  RuntimeCheck(topk <= num_experts, "topk must be <= num_experts");
+  if (num_tokens == 0) return;
+
+  auto* weights_ptr = static_cast<float*>(topk_weights.data_ptr());
+  auto* indices_ptr = static_cast<int*>(topk_indices.data_ptr());
+  auto* packed_ptr = static_cast<int*>(packed.data_ptr());
+
+  if (is_type<fp32_t>(gating_output.dtype())) {
+    dispatchExperts<float>(
+        static_cast<const float*>(gating_output.data_ptr()),
+        weights_ptr,
+        indices_ptr,
+        packed_ptr,
+        ntnp_ptr,
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        device);
+  } else if (is_type<fp16_t>(gating_output.dtype())) {
+    dispatchExperts<__half>(
+        static_cast<const __half*>(gating_output.data_ptr()),
+        weights_ptr,
+        indices_ptr,
+        packed_ptr,
+        ntnp_ptr,
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        device);
+  } else {
+    dispatchExperts<__nv_bfloat16>(
+        static_cast<const __nv_bfloat16*>(gating_output.data_ptr()),
+        weights_ptr,
+        indices_ptr,
+        packed_ptr,
+        ntnp_ptr,
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        device);
+  }
+}
+
+}  // namespace