chore: import upstream snapshot with attribution

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 <cub/block/block_reduce.cuh>
+#include <cuda/functional>
+
+#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <optional>
+
+namespace flashinfer {
+namespace sgl_fused_act_quant {
+
+namespace tk = tensorrt_llm::kernels;
+
+// Same silu as moe::dev::activation (trtllm_fused_moe_dev_kernel.cu:55): x / (1 + exp(-x)).
+inline __device__ float fused_silu(float x) {
+  return x / (1.0f + expf(-x));
+}
+
+// One block per expanded row. gateUp is the column-interleaved GEMM1 output (g0,u0,g1,u1,...)
+// indexed by permutedIdx; loraDelta is the contiguous [gate|up] delta indexed by expandedIdx.
+template <uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT, bool DISABLE_FP4_FAST_MATH>
+__global__ void fusedActivationQuantKernel(
+    int m,                                        // numTokens * topK (number of expanded rows)
+    int innerHalf,                                // inter == n (output width per row); must be a multiple of 16
+    int innerDim,                                 // gate_up_n == 2 * innerHalf
+    __nv_bfloat16 const* __restrict__ gateUp,     // interleaved gate/up, [.., innerDim] by permutedIdx
+    __nv_bfloat16 const* __restrict__ loraDelta,  // [.., innerDim] by expandedIdx, may be null
+    __nv_bfloat16* __restrict__ loraInputOut,     // [.., innerHalf] by expandedIdx, may be null
+    int32_t const* __restrict__ expandedIdxToPermutedIdx,
+    float globalScaleInv,
+    uint8_t* __restrict__ weightOutput,  // fp4 [.., innerHalf/2] by permutedIdx
+    uint8_t* __restrict__ scaleOutput,   // swizzled e4m3 SF
+    float* __restrict__ perTokenScaleOutput) {
+  constexpr int SF_VEC_SIZE = 16;
+  using InType = tk::PackedVec<__nv_bfloat16, SF_VEC_SIZE>;  // 16 bf16 == 8 __nv_bfloat162
+  using PackedFp4Type = uint64_t;                            // SF_VEC_SIZE == 16
+
+  int const expandedIdx = blockIdx.x;
+  if (expandedIdx >= m) return;
+  int const permutedIdx = expandedIdxToPermutedIdx[expandedIdx];
+  int const num_vecs_per_row = innerHalf / SF_VEC_SIZE;
+  int64_t const liBaseRow = (int64_t)expandedIdx * innerHalf;
+
+  // Padding row: the separate activation kernel writes 0 to activation_lora_input and skips
+  // the quant outputs. Mirror that, then return.
+  if (permutedIdx < 0) {
+    if (loraInputOut != nullptr) {
+      InType z;
+#pragma unroll
+      for (int i = 0; i < SF_VEC_SIZE / 2; ++i)
+        z.elts[i] = __float2bfloat162_rn(0.0f);
+      for (int vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+        *reinterpret_cast<InType*>(&loraInputOut[liBaseRow + (int64_t)vecIdx * SF_VEC_SIZE]) = z;
+      }
+    }
+    return;
+  }
+
+  int64_t const permBase = (int64_t)permutedIdx * innerDim;  // gate_up row (interleaved)
+  int64_t const expBase = (int64_t)expandedIdx * innerDim;   // delta row (contiguous gate|up)
+  (void)DISABLE_FP4_FAST_MATH;
+
+  // 1 SF block (16 outputs) per thread, held in registers across the amax barrier (no smem cache):
+  // requires num_vecs_per_row <= BLOCK_SIZE (inter=2048 -> 128 == BLOCK_SIZE). With
+  // CVT_ELTS_PER_THREAD == SF_VEC_SIZE the cvt needs no cross-thread shuffle, so masking is safe.
+  int const vecIdx = threadIdx.x;
+  bool const active = vecIdx < num_vecs_per_row;
+
+  InType vec;
+  float localAmax = 0.f;
+  if (active) {
+    int const h0 = vecIdx * SF_VEC_SIZE;
+    __nv_bfloat16 const* g = gateUp + permBase + (int64_t)2 * h0;     // 32 interleaved bf16
+    __nv_bfloat16 const* dlo = loraDelta + expBase + h0;              // silu-arg delta (lower half)
+    __nv_bfloat16 const* dhi = loraDelta + expBase + innerHalf + h0;  // multiplier delta (upper half)
+    __nv_bfloat162 amax2 = __float2bfloat162_rn(0.0f);
+    union {
+      int4 v[4];
+      __nv_bfloat16 b[32];
+    } gu;
+    union {
+      int4 v[2];
+      __nv_bfloat16 b[16];
+    } dl, dh;
+    int4 const* gp = reinterpret_cast<int4 const*>(g);
+#pragma unroll
+    for (int k = 0; k < 4; ++k)
+      gu.v[k] = gp[k];
+    if (loraDelta != nullptr) {
+      int4 const* dlp = reinterpret_cast<int4 const*>(dlo);
+      int4 const* dhp = reinterpret_cast<int4 const*>(dhi);
+#pragma unroll
+      for (int k = 0; k < 2; ++k) {
+        dl.v[k] = dlp[k];
+        dh.v[k] = dhp[k];
+      }
+    }
+#pragma unroll
+    for (int i = 0; i < SF_VEC_SIZE / 2; ++i) {  // 8 bf162 = 16 output elements
+      int const j0 = 2 * i, j1 = 2 * i + 1;
+      float even0 = (float)gu.b[2 * j0], odd0 = (float)gu.b[2 * j0 + 1];
+      float even1 = (float)gu.b[2 * j1], odd1 = (float)gu.b[2 * j1 + 1];
+      float a0 = odd0, b0 = even0, a1 = odd1, b1 = even1;
+      if (loraDelta != nullptr) {
+        a0 += (float)dl.b[j0];
+        b0 += (float)dh.b[j0];
+        a1 += (float)dl.b[j1];
+        b1 += (float)dh.b[j1];
+      }
+      float act0 = fused_silu(a0) * b0;
+      float act1 = fused_silu(a1) * b1;
+      __nv_bfloat162 e = __float22bfloat162_rn(make_float2(act0, act1));
+      vec.elts[i] = e;
+      amax2 = __hmax2(amax2, __habs2(e));
+    }
+    localAmax = (float)__hmax(amax2.x, amax2.y);
+    if (loraInputOut != nullptr) {
+      *reinterpret_cast<InType*>(&loraInputOut[liBaseRow + h0]) = vec;
+    }
+  }
+
+  // ---- per-token scale: blockReduce amax, broadcast via smem (no gmem round-trip) ----
+  using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+  __shared__ typename BlockReduce::TempStorage tempStorage;
+  __shared__ float sScale;
+  float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
+  if (threadIdx.x == 0) {
+    float const pts = globalAmax * globalScaleInv;
+    perTokenScaleOutput[permutedIdx] = pts;
+    sScale = pts;
+  }
+  __syncthreads();
+  float const globalEncodeScale = tk::reciprocal_approximate_ftz(sScale);
+
+  // ---- quantize from registers (cvt computes the per-16 e4m3 block scale internally) ----
+  if (active) {
+    uint8_t fp8Scale;
+    // 5 template args on this flashinfer build: Type, SF_VEC_SIZE, CVT_ELTS_PER_THREAD,
+    // UE8M0_SF=false, TE_EXACT_NVFP4=false (the default nvfp4 quant path).
+    auto fp4Vals = tk::cvt_warp_fp16_to_fp4<__nv_bfloat16, SF_VEC_SIZE, SF_VEC_SIZE, false, false>(
+        vec, globalEncodeScale, &fp8Scale);
+    int64_t const vecOffset = (int64_t)permutedIdx * num_vecs_per_row + vecIdx;
+    reinterpret_cast<PackedFp4Type*>(weightOutput)[vecOffset] = fp4Vals;
+
+    // Match nvfp4QuantAndPerTokenScaleKernel exactly (it passes the kernel's `m` as numRows).
+    int64_t sfOffset;
+    if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+      sfOffset = vecOffset;
+    } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+      sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
+    } else {
+      sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
+    }
+    scaleOutput[sfOffset] = fp8Scale;
+  }
+}
+
+// Host launch: globalScaleInv = 1/448/6. BLOCK_SIZE must be >= innerHalf/16 (one SF block/thread).
+inline void launchFusedActivationQuant(
+    int m,
+    int innerHalf,
+    int innerDim,
+    __nv_bfloat16 const* gateUp,
+    __nv_bfloat16 const* loraDelta,
+    __nv_bfloat16* loraInputOut,
+    int32_t const* expandedIdxToPermutedIdx,
+    float globalScaleInv,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput,
+    tensorrt_llm::QuantizationSFLayout sfLayout,
+    bool disableFp4FastMath,
+    cudaStream_t stream) {
+  constexpr uint32_t BLOCK_SIZE = 128;  // == innerHalf/16 for inter=2048 (one SF block per thread)
+  dim3 const grid(m), block(BLOCK_SIZE);
+
+  auto launch = [&](auto layoutTag, auto fastMathTag) {
+    fusedActivationQuantKernel<BLOCK_SIZE, decltype(layoutTag)::value, decltype(fastMathTag)::value>
+        <<<grid, block, 0, stream>>>(
+            m,
+            innerHalf,
+            innerDim,
+            gateUp,
+            loraDelta,
+            loraInputOut,
+            expandedIdxToPermutedIdx,
+            globalScaleInv,
+            weightOutput,
+            scaleOutput,
+            perTokenScaleOutput);
+  };
+  auto withFastMath = [&](auto layoutTag) {
+    if (disableFp4FastMath) {
+      launch(layoutTag, std::integral_constant<bool, true>{});
+    } else {
+      launch(layoutTag, std::integral_constant<bool, false>{});
+    }
+  };
+  if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+    withFastMath(
+        std::integral_constant<
+            tensorrt_llm::QuantizationSFLayout,
+            tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
+  } else if (sfLayout == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+    withFastMath(
+        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::LINEAR>{});
+  } else {
+    withFastMath(
+        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
+  }
+}
+
+}  // namespace sgl_fused_act_quant
+}  // namespace flashinfer

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 <cub/cub.cuh>
+
+#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <optional>
+#include <type_traits>
+
+namespace sgl_fused_permute_quant {
+
+namespace tk = tensorrt_llm::kernels;
+
+// Quantize source row `readRow` of `input` (unpermuted) and write fp4 + block-sf + per-token-sf to
+// destination row `writeRow` of the permuted outputs. `numRowsSf` is the SF buffer's row count
+// (= max_padded), matching the plain quant's `m` arg to get_sf_out_offset_*.
+template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
+__device__ __forceinline__ void fused_quant_one_row(
+    uint32_t n,
+    T const* input,
+    int readRow,
+    int writeRow,
+    int numRowsSf,
+    float globalScaleInv,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput) {
+  constexpr int SF_VEC_SIZE = 16;
+  constexpr int ELTS_PER_THREAD = 16;
+  using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
+  using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
+  uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
+  uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
+  InType const* inBase = reinterpret_cast<InType const*>(input);
+
+  // ---- pass 1: per-row amax over the (unpermuted) source row ----
+  float localAmax = 0.f;
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
+    std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
+#pragma unroll
+    for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
+      a = __hmax2(a, __habs2(vec_in.elts[i]));
+    }
+    localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
+  }
+
+  using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+  __shared__ typename BlockReduce::TempStorage tempStorage;
+  float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
+
+  // ---- per-token scale (TE_EXACT branch — production instantiates TE_EXACT_NVFP4=true for bf16,
+  //      quantization.cu:247): globalEncodeScale = globalScale/globalAmax (exact __fdiv_rn), stored
+  //      per-token scale = 1/globalEncodeScale. __shared__ scalar replaces the gmem round-trip
+  //      (bit-identical: an fp32 store->load doesn't change the value). ----
+  __shared__ float sEncodeScale;
+  if (threadIdx.x == 0) {
+    float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
+    float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
+    sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
+  }
+  __syncthreads();
+  float const globalEncodeScale = sEncodeScale;
+  float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
+  if (threadIdx.x == 0) perTokenScaleOutput[writeRow] = perTokenScale;
+
+  // ---- pass 2: quantize + scatter-write to the permuted destination ----
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
+    uint8_t fp8Scale;
+    auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
+        T,
+        SF_VEC_SIZE,
+        ELTS_PER_THREAD,
+        /*UE8M0_SF=*/false,
+        /*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
+    reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
+        fp4Vals;
+
+    int64_t sfOffset;
+    if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+      sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
+    } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+      sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+    } else {
+      sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+    }
+    scaleOutput[sfOffset] = fp8Scale;
+  }
+}
+
+// no-dedup: grid.x = num_tokens*top_k (one block per (token,expert) pair).
+template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
+__global__ void fusedPermuteNvfp4QuantKernel(
+    uint32_t numPairs,
+    uint32_t n,
+    uint32_t topK,
+    int numRowsSf,
+    T const* input,
+    float globalScaleInv,
+    int32_t const* expandedIdxToPermutedIdx,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput) {
+  uint32_t const expandedIdx = blockIdx.x;
+  if (expandedIdx >= numPairs) return;
+  int const writeRow = expandedIdxToPermutedIdx[expandedIdx];
+  if (writeRow < 0) return;
+  int const readRow = static_cast<int>(expandedIdx / topK);
+  fused_quant_one_row<T, BLOCK_SIZE, SF_LAYOUT>(
+      n, input, readRow, writeRow, numRowsSf, globalScaleInv, weightOutput, scaleOutput, perTokenScaleOutput);
+}
+
+// dedup: grid.x = num_tokens (one block per source token, scatter to its top_k destinations).
+template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
+__global__ void fusedPermuteNvfp4QuantDedupKernel(
+    uint32_t numTokens,
+    uint32_t n,
+    uint32_t topK,
+    int numRowsSf,
+    T const* input,
+    float globalScaleInv,
+    int32_t const* expandedIdxToPermutedIdx,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput) {
+  constexpr int SF_VEC_SIZE = 16;
+  constexpr int ELTS_PER_THREAD = 16;
+  using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
+  using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
+  uint32_t const token = blockIdx.x;
+  if (token >= numTokens) return;
+  uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
+  uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
+  InType const* inBase = reinterpret_cast<InType const*>(input);
+
+  // pass 1: amax over the source token row (read once).
+  float localAmax = 0.f;
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
+    std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
+#pragma unroll
+    for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
+      a = __hmax2(a, __habs2(vec_in.elts[i]));
+    }
+    localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
+  }
+  using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+  __shared__ typename BlockReduce::TempStorage tempStorage;
+  float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
+
+  // TE_EXACT per-token scale (matches production; see fused_quant_one_row).
+  __shared__ float sEncodeScale;
+  if (threadIdx.x == 0) {
+    float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
+    float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
+    sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
+  }
+  __syncthreads();
+  float const globalEncodeScale = sEncodeScale;
+  float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
+
+  // per-token scale -> each (valid) destination (top_k small; first top_k threads write).
+  if (threadIdx.x < topK) {
+    int const writeRow = expandedIdxToPermutedIdx[token * topK + threadIdx.x];
+    if (writeRow >= 0) perTokenScaleOutput[writeRow] = perTokenScale;
+  }
+
+  // pass 2: quantize each vec once, scatter to all valid destinations.
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
+    uint8_t fp8Scale;
+    auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
+        T,
+        SF_VEC_SIZE,
+        ELTS_PER_THREAD,
+        /*UE8M0_SF=*/false,
+        /*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
+#pragma unroll 1
+    for (uint32_t k = 0; k < topK; ++k) {
+      int const writeRow = expandedIdxToPermutedIdx[token * topK + k];
+      if (writeRow < 0) continue;
+      reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
+          fp4Vals;
+      int64_t sfOffset;
+      if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+        sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
+      } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+        sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+      } else {
+        sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+      }
+      scaleOutput[sfOffset] = fp8Scale;
+    }
+  }
+}
+
+// Launcher. `dedup` picks the variant. `n` (= hidden) must be a multiple of 16. `numRowsSf` is the
+// SF buffer's row count (max_padded).
+template <typename T>
+void invokeFusedPermuteNvfp4Quant(
+    uint32_t numTokens,
+    uint32_t topK,
+    uint32_t n,
+    int numRowsSf,
+    T const* input,
+    float globalScaleInv,
+    int32_t const* expandedIdxToPermutedIdx,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput,
+    tensorrt_llm::QuantizationSFLayout sfLayout,
+    bool dedup,
+    cudaStream_t stream) {
+  // [opt] Occupancy tuning (ncu: kernel is occupancy-bound, not DRAM-bound — DRAM <1%, achieved
+  // occupancy was 19.8% no-dedup / 5.5% dedup at BLOCK_SIZE=128). The dedup variant launches only
+  // num_tokens CTAs (=64 at decode bs64), so it is the most CTA-starved; widening the block raises
+  // threads/CTA and hides the per-row amax-reduction + scatter latency. Decode bs64 dedup sweep:
+  // 128 -> 5.52us, 256 -> 4.09us, 512 -> 3.71us. 512 is the chosen default (the prod path uses
+  // dedup). (7168/16 = 448 vecs/row, so >448 threads idle on the tail, but the win dominates.)
+  constexpr uint32_t BLOCK_SIZE = 512;
+  dim3 const block(BLOCK_SIZE);
+
+  auto dispatch = [&](auto layoutTag) {
+    constexpr tensorrt_llm::QuantizationSFLayout LAYOUT = decltype(layoutTag)::value;
+    if (dedup) {
+      dim3 const grid(numTokens);
+      fusedPermuteNvfp4QuantDedupKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
+          numTokens,
+          n,
+          topK,
+          numRowsSf,
+          input,
+          globalScaleInv,
+          expandedIdxToPermutedIdx,
+          weightOutput,
+          scaleOutput,
+          perTokenScaleOutput);
+    } else {
+      dim3 const grid(numTokens * topK);
+      fusedPermuteNvfp4QuantKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
+          numTokens * topK,
+          n,
+          topK,
+          numRowsSf,
+          input,
+          globalScaleInv,
+          expandedIdxToPermutedIdx,
+          weightOutput,
+          scaleOutput,
+          perTokenScaleOutput);
+    }
+  };
+
+  if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+    dispatch(
+        std::integral_constant<
+            tensorrt_llm::QuantizationSFLayout,
+            tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
+  } else {
+    dispatch(
+        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
+  }
+}
+
+}  // namespace sgl_fused_permute_quant

python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/DevKernel.h

@@ -0,0 +1,472 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
+#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h"
+#include <cuda.h>
+#include <iostream>
+#include <string>
+// #include <cuda_runtime_api.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_size.h>
+#include <cutlass/numeric_types.h>
+
+#include "flashinfer/exception.h"
+// #include <tensorrt_llm/common/assert.h>
+#include "flashinfer/trtllm/common/cudaUtils.h"
+#include "tensorrt_llm/common/logger.h"
+
+namespace moe::dev {
+
+#define CHECK_CUDA_ERROR(cmd)                                                                \
+  do {                                                                                       \
+    cudaError_t e = cmd;                                                                     \
+    if (e != cudaSuccess) {                                                                  \
+      std::cout << "CUDA error in " << __FILE__ << ":" << __LINE__ << " executing '" << #cmd \
+                << "': " << cudaGetErrorString(e);                                           \
+    }                                                                                        \
+    FLASHINFER_CHECK(e == cudaSuccess, "Got CUDA error. See above for details.");            \
+  } while (0)
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define LAUNCH_ESC(...) __VA_ARGS__
+
+#define LAUNCH_PDL(data, coopLaunch, types, kernel, numBlocks, numThreads, smemSize, stream)                      \
+  cudaLaunchConfig_t config{};                                                                                    \
+  config.gridDim = numBlocks;                                                                                     \
+  config.blockDim = numThreads;                                                                                   \
+  config.dynamicSmemBytes = smemSize;                                                                             \
+  config.stream = (cudaStream_t)stream;                                                                           \
+                                                                                                                  \
+  cudaLaunchAttribute attributes[2] = {};                                                                         \
+  attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;                                          \
+  attributes[0].val.programmaticStreamSerializationAllowed = int(data.mUsePdl);                                   \
+  attributes[1].id = cudaLaunchAttributeCooperative;                                                              \
+  attributes[1].val.cooperative = int(coopLaunch);                                                                \
+  config.attrs = attributes;                                                                                      \
+  config.numAttrs = 2;                                                                                            \
+  if (data.mUsePdl) {                                                                                             \
+    auto params = KernelParams<types, true>::setKernelParams(data);                                               \
+    auto kernelTyped = kernel<KernelParams<types, true>>;                                                         \
+    if (smemSize > 48 * 1024)                                                                                     \
+      CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
+    CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params));                                           \
+  } else {                                                                                                        \
+    auto params = KernelParams<types, false>::setKernelParams(data);                                              \
+    auto kernelTyped = kernel<KernelParams<types, false>>;                                                        \
+    if (smemSize > 48 * 1024)                                                                                     \
+      CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
+    CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params));                                           \
+  }
+
+#define LAUNCH(data, kernel, numBlocks, numThreads, smemSize, stream)                                \
+  if (data.mDtypeElt == tg::Dtype::Fp16) {                                                           \
+    LAUNCH_PDL(data, false, cutlass::half_t, kernel, numBlocks, numThreads, smemSize, stream);       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3) {                                                    \
+    LAUNCH_PDL(data, false, cutlass::float_e4m3_t, kernel, numBlocks, numThreads, smemSize, stream); \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16) {                                                \
+    LAUNCH_PDL(data, false, cutlass::bfloat16_t, kernel, numBlocks, numThreads, smemSize, stream);   \
+  } else {                                                                                           \
+    FLASHINFER_WARN("Unsupported dtypeElt");                                                         \
+  }
+
+#define LAUNCH_NUM_TOKENS_PER_CTA(data, type, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream) \
+  if (numTokensPerCta == 4) {                                                                                   \
+    LAUNCH_PDL(data, false, LAUNCH_ESC(type, 4), kernel, numBlocks, numThreads, smemSize, stream);              \
+  } else if (numTokensPerCta == 2) {                                                                            \
+    LAUNCH_PDL(data, false, LAUNCH_ESC(type, 2), kernel, numBlocks, numThreads, smemSize, stream);              \
+  } else if (numTokensPerCta == 1) {                                                                            \
+    LAUNCH_PDL(data, false, LAUNCH_ESC(type, 1), kernel, numBlocks, numThreads, smemSize, stream);              \
+  } else {                                                                                                      \
+    FLASHINFER_WARN("Unsupported numTokensPerCta");                                                             \
+  }
+
+#define LAUNCH_ACTIVATION(data, kernel, numTokensPerCta, numBlocks, numThreads, smemSize, stream)       \
+  if (data.mDtypeElt == tg::Dtype::Fp16) {                                                              \
+    LAUNCH_NUM_TOKENS_PER_CTA(                                                                          \
+        data, cutlass::half_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream);       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3) {                                                       \
+    LAUNCH_NUM_TOKENS_PER_CTA(                                                                          \
+        data, cutlass::float_e4m3_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16) {                                                   \
+    LAUNCH_NUM_TOKENS_PER_CTA(                                                                          \
+        data, cutlass::bfloat16_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream);   \
+  } else {                                                                                              \
+    FLASHINFER_WARN("Unsupported dtypeElt");                                                            \
+  }
+
+#define LAUNCH_EXPW(data, kernel, topK, numBlocks, numThreads, smemSize, stream)                                       \
+  if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Fp32) {                                       \
+    LAUNCH_PDL(                                                                                                        \
+        data, false, LAUNCH_ESC(cutlass::half_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream);       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Fp32) {                                \
+    LAUNCH_PDL(                                                                                                        \
+        data, false, LAUNCH_ESC(cutlass::float_e4m3_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Fp32) {                            \
+    LAUNCH_PDL(                                                                                                        \
+        data, false, LAUNCH_ESC(cutlass::bfloat16_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream);   \
+  } else if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Bfloat16) {                            \
+    LAUNCH_PDL(                                                                                                        \
+        data,                                                                                                          \
+        false,                                                                                                         \
+        LAUNCH_ESC(cutlass::half_t, cutlass::bfloat16_t, topK),                                                        \
+        kernel,                                                                                                        \
+        numBlocks,                                                                                                     \
+        numThreads,                                                                                                    \
+        smemSize,                                                                                                      \
+        stream);                                                                                                       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Bfloat16) {                            \
+    LAUNCH_PDL(                                                                                                        \
+        data,                                                                                                          \
+        false,                                                                                                         \
+        LAUNCH_ESC(cutlass::float_e4m3_t, cutlass::bfloat16_t, topK),                                                  \
+        kernel,                                                                                                        \
+        numBlocks,                                                                                                     \
+        numThreads,                                                                                                    \
+        smemSize,                                                                                                      \
+        stream);                                                                                                       \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Bfloat16) {                        \
+    LAUNCH_PDL(                                                                                                        \
+        data,                                                                                                          \
+        false,                                                                                                         \
+        LAUNCH_ESC(cutlass::bfloat16_t, cutlass::bfloat16_t, topK),                                                    \
+        kernel,                                                                                                        \
+        numBlocks,                                                                                                     \
+        numThreads,                                                                                                    \
+        smemSize,                                                                                                      \
+        stream);                                                                                                       \
+  } else {                                                                                                             \
+    FLASHINFER_WARN("Unsupported pair");                                                                               \
+  }
+
+#define LAUNCH_TOPK_EXPW(data, kernel, numBlocks, numThreads, smemSize, stream) \
+  if (data.topK % 4 == 0) {                                                     \
+    LAUNCH_EXPW(data, kernel, 4, numBlocks, numThreads, smemSize, stream);      \
+  } else if (data.topK % 2 == 0) {                                              \
+    LAUNCH_EXPW(data, kernel, 2, numBlocks, numThreads, smemSize, stream);      \
+  } else {                                                                      \
+    LAUNCH_EXPW(data, kernel, 1, numBlocks, numThreads, smemSize, stream);      \
+  }
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// NOTE: Old routing-specific macros (LAUNCH_TILEN, LAUNCH_ROUTING_LLAMA4,
+// LAUNCH_ROUTING_DEEPSEEK_*, LAUNCH_ROUTING_WITH_NUM_EXPERTS) have been moved to
+// RoutingDevKernel.h which uses the new template signature with runtime isPow2/UsePdl.
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace activation {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  tg::Dtype mDtypeElt{tg::Dtype::Fp16};
+  bool mUsePdl{false};
+  bool mUseDeepSeekFp8{false};
+
+  void* inPtr;
+  void* outPtr;
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+  cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
+  cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
+
+  // When true, inPtr holds the column-interleaved gate/up GEMM1 output
+  // (g0,u0,g1,u1,...) and the kernel de-interleaves on read (x1=col 2k, x2=col 2k+1);
+  // when false, inPtr is the contiguous [gate | up] layout. Default false keeps the
+  // FP8 / non-LoRA semantics; the FP4 LoRA path sets it true to fuse the standalone
+  // de-interleave kernel into this activation read.
+  bool interleavedGateUpInput = false;
+
+  int32_t innerDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t const* totalNumPaddedTokens;
+
+  // Bench/tuning knob: when > 0, overrides the activation launch grid.x (default
+  // innerDim/128). The kernel's grid-stride hidden-dim loop makes any grid.x
+  // produce bitwise-identical output, so a smaller grid.x removes empty blocks and
+  // gives each thread a longer strip (more in-flight loads) without changing math.
+  int32_t actGridXOverride = 0;
+
+  // Activation kernel variant: 0 = scalar activationKernel, 1 = vectorized
+  // activationKernelOpt (128-bit gate/up + 64-bit delta/store, 4 pairs/thread).
+  // Variant 1 applies only to the bf16 interleaved path with (innerDim/2)%4==0;
+  // run() falls back to the scalar kernel otherwise. Output is bitwise-identical.
+  int32_t actOptMode = 0;
+};
+
+template <typename Type_, int32_t NumTokensPerCta_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  static constexpr int32_t NumTokensPerCta = NumTokensPerCta_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  Type const* inPtr;
+  Type* outPtr;
+
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+  cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
+  cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
+
+  bool interleavedGateUpInput = false;
+
+  int32_t innerDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t const* totalNumPaddedTokens;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inPtr = (Type*)data.inPtr;
+    params.outPtr = (Type*)data.outPtr;
+    params.inDqSfsPtr = data.inDqSfsPtr;
+    params.outDqSfsPtr = data.outDqSfsPtr;
+    params.gateUpLoraDeltaPtr = data.gateUpLoraDeltaPtr;
+    params.activationLoraInputOutPtr = data.activationLoraInputOutPtr;
+    params.interleavedGateUpInput = data.interleavedGateUpInput;
+
+    params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
+
+    params.innerDim = data.innerDim;
+    params.numTokens = data.numTokens;
+    params.topK = data.topK;
+    params.totalNumPaddedTokens = data.totalNumPaddedTokens;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+}  // namespace activation
+
+namespace convertsf {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  bool mUsePdl{false};
+
+  void* inSfPtr = nullptr;
+  void* outSfPtr = nullptr;
+  int32_t hiddenDimSf;
+  int32_t numTokens;
+  tg::SfLayout sfLayoutSrc;
+  tg::SfLayout sfLayoutDst;
+};
+
+template <typename Type_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  void const* inSfPtr = nullptr;
+  void* outSfPtr = nullptr;
+  int32_t hiddenDimSf;
+  int32_t numTokens;
+  tg::SfLayout sfLayoutSrc;
+  tg::SfLayout sfLayoutDst;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inSfPtr = data.inSfPtr;
+    params.outSfPtr = data.outSfPtr;
+    params.hiddenDimSf = data.hiddenDimSf;
+    params.numTokens = data.numTokens;
+    params.sfLayoutSrc = data.sfLayoutSrc;
+    params.sfLayoutDst = data.sfLayoutDst;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace convertsf
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace permute {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  tg::Dtype mDtypeElt{tg::Dtype::Fp16};
+  bool mUsePdl{false};
+  bool mUseDeepSeekFp8{false};
+
+  void* inPtr;
+  void* outPtr;
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+  int32_t* expandedIdxToPermutedIdx;
+  int32_t hiddenDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t const* totalNumPaddedTokens;
+};
+
+template <typename Type_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  Type const* inPtr;
+  Type* outPtr;
+  float const* inDqSfsPtr;
+  float* outDqSfsPtr;
+  int32_t* expandedIdxToPermutedIdx;
+  int32_t hiddenDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t const* totalNumPaddedTokens;
+  bool useDeepSeekFp8;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inPtr = (Type*)data.inPtr;
+    params.outPtr = (Type*)data.outPtr;
+    params.inDqSfsPtr = data.inDqSfsPtr;
+    params.outDqSfsPtr = data.outDqSfsPtr;
+    params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
+    params.hiddenDim = data.hiddenDim;
+    params.numTokens = data.numTokens;
+    params.topK = data.topK;
+    params.totalNumPaddedTokens = data.totalNumPaddedTokens;
+    params.useDeepSeekFp8 = data.mUseDeepSeekFp8;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace permute
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace finalize {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  tg::Dtype mDtypeElt{tg::Dtype::Fp16};
+  tg::Dtype mDtypeExpW{tg::Dtype::Bfloat16};
+  bool mUsePdl{false};
+  bool mUseDeepSeekFp8{false};
+
+  void* inPtr;
+  void* outPtr;
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+
+  void* expertWeightsPtr;
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t numTokens;
+  int32_t numExperts;
+  int32_t topK;
+  // Hidden dimension output of MoE block. It is not padded.
+  int32_t hiddenDim;
+  // Hidden dimension output of FC2. It might be padded.
+  int32_t hiddenDimPadded;
+  int32_t const* totalNumPaddedTokens;
+};
+
+template <typename Type_, typename TypeExpW_, int TopKUnrollFactor_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  using TypeExpW = TypeExpW_;
+  static constexpr int TopKUnrollFactor = TopKUnrollFactor_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  Type const* inPtr;
+  TypeExpW const* expertWeightsPtr;
+  Type* outPtr;
+
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t hiddenDim;
+  int32_t hiddenDimPadded;
+  int32_t numTokens;
+  int32_t numExperts;
+  int32_t topK;
+  int32_t const* totalNumPaddedTokens;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inPtr = (Type*)data.inPtr;
+    params.expertWeightsPtr = (TypeExpW*)data.expertWeightsPtr;
+    params.outPtr = (Type*)data.outPtr;
+    params.inDqSfsPtr = data.inDqSfsPtr;
+    params.outDqSfsPtr = data.outDqSfsPtr;
+
+    params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
+
+    params.hiddenDim = data.hiddenDim;
+    params.hiddenDimPadded = data.hiddenDimPadded;
+    params.numTokens = data.numTokens;
+    params.numExperts = data.numExperts;
+    params.topK = data.topK;
+    params.totalNumPaddedTokens = data.totalNumPaddedTokens;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace finalize
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace moe::dev

python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h

@@ -0,0 +1,586 @@
+/*
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "DevKernel.h"
+#include "flashinfer/trtllm/fused_moe/RoutingKernel.h"
+#include <string>
+// #include "flashinfer/trtllm/common/cudaDriverWrapper.h"
+#include "flashinfer/trtllm/batched_gemm/KernelRunner.h"
+#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
+#include "flashinfer/trtllm/common/cudaUtils.h"
+
+namespace tensorrt_llm {
+namespace kernels {
+namespace trtllmgen_moe {
+
+namespace MoE {
+class Runner;
+}  // namespace MoE
+
+namespace Routing {
+
+// The type of method in top-K routing, for use in torch custom op
+// Please keep this in sync with the counterpart defined in
+// flashinfer/fused_moe/core.py
+enum class RoutingMethodType : int64_t {
+  // Default: Softmax -> TopK
+  Default = 0,
+  // Renormalize: TopK -> Softmax
+  Renormalize = 1,
+  // DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups -> Top8 experts from the
+  // Top4 groups
+  DeepSeekV3 = 2,
+  // Llama4: Top1 -> Sigmoid
+  Llama4 = 3,
+  // RenormalizeNaive: Softmax -> TopK -> Renormalize
+  RenormalizeNaive = 4,
+  // TopK only (no softmax)
+  TopK = 5,
+  // SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K weights)
+  SigmoidRenorm = 6,
+  // MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize (routeScale=1.0, epsilon=1e-20)
+  MiniMax2 = 7,
+  // Sigmoid: Sigmoid -> TopK (no renormalization)
+  Sigmoid = 8,
+  // Unspecified
+  Unspecified = 9,
+};
+
+inline int32_t maybeGetMinTokenCount(int32_t numPaddedTokens, int32_t hiddenSize, int32_t dtypeSizeBits) {
+  // Pad so total size exceeds 128KiB for performance reasons
+  int32_t minNumTokensRequired = common::divUp(128 * 1024 * 8, hiddenSize * dtypeSizeBits);
+  return std::max(numPaddedTokens, minNumTokensRequired);
+}
+
+inline std::string serializeMoeRoutingMethodType(RoutingMethodType routingMethodType) {
+  switch (routingMethodType) {
+    case RoutingMethodType::Default:
+      return "Default";
+    case RoutingMethodType::Renormalize:
+      return "Renormalize";
+    case RoutingMethodType::DeepSeekV3:
+      return "DeepSeekV3";
+    case RoutingMethodType::Llama4:
+      return "Llama4";
+    case RoutingMethodType::RenormalizeNaive:
+      return "RenormalizeNaive";
+    case RoutingMethodType::TopK:
+      return "TopK";
+    case RoutingMethodType::SigmoidRenorm:
+      return "SigmoidRenorm";
+    case RoutingMethodType::MiniMax2:
+      return "MiniMax2";
+    case RoutingMethodType::Sigmoid:
+      return "Sigmoid";
+    default:
+      return "InvalidRountingMethod";  // TODO throw error
+  };
+}
+
+inline int32_t getMaxNumCtasInBatchDim(int32_t numTokens, int32_t topK, int32_t numExperts, int32_t tileTokensDim) {
+  // For MoE, mNumTokens != 0 and the number of CTAs is known only at runtime.
+  // We launch maximally possible number of CTAs and use ptrNumNonExitingCtas to determine
+  // the actual number of CTAs to run.
+
+  // Initialize number of tokens with the number of expanded tokens after routing.
+  int32_t numRemainingTokens = numTokens * topK;
+  int32_t maxNumCtasInBatchDim = 0;
+  // First, distribute one token each expert until token depletion to maximize CTA tile count.
+  int32_t numExpertsFilled = std::min(numExperts, numRemainingTokens);
+  maxNumCtasInBatchDim += numExpertsFilled;
+  numRemainingTokens -= numExpertsFilled;
+  // Next, greedily pour all remaining tokens to one expert to maximize CTA tile count.
+  // E.g., at this point tokens over 4 experts are [1, 1, 1, 1], and we have 4 tokens left.
+  // If each CTA handles 4 tokens/expert, the greedy strategy is to pour all remaining tokens
+  // to any one expert to get to the 5th CTA tile. Otherwise, we can only get 4 tiles in total.
+  //
+  // Another way to reason about this is to pour the remaining tokens into buckets of some fixed
+  // capacity. These buckets, if full, can then be attributed to any expert; it does not have to
+  // belong to the same expert every time.
+  if (numRemainingTokens > 0) {
+    // For every tileTokenDim tokens, we add an extra CTA tile in the token dimension.
+    // The number of CTA tiles is given by divDown(numRemainingTokens, tokenTileDim).
+    maxNumCtasInBatchDim += (numRemainingTokens / tileTokensDim);
+  }
+  return maxNumCtasInBatchDim;
+}
+
+inline int32_t
+getMaxPermutedPaddedCount(int32_t numTokens, int32_t expertsPerToken, int32_t numExperts, int32_t padding) {
+  int32_t maxCtas = getMaxNumCtasInBatchDim(numTokens, expertsPerToken, numExperts, padding);
+  return maxCtas * padding;
+}
+
+class Runner {
+ public:
+  explicit Runner();
+
+  explicit Runner(int32_t tileTokensDim);
+
+  void
+  run(void* routingLogits,
+      void* routingBias,
+      int32_t numTokens,
+      int32_t numExperts,
+      int32_t topK,
+      int32_t nGroups,
+      int32_t topkGroups,
+      int32_t localExpertOffset,
+      int32_t localNumExperts,
+      float routedScalingFactor,
+      int32_t* routingExpertIndexes,
+      int32_t* expertCountHistogram,
+      int32_t* permutedIdxSize,
+      int32_t* expandedIdxToPermutedIdx,
+      int32_t* permutedIdxToExpandedIdx,
+      int32_t* permutedIdxToTokenIdx,
+      int32_t* expertIds,
+      void* expertWeights,
+      int32_t* numTokensPerExpert,
+      int32_t* ctaIdxXyToBatchIdx,
+      int32_t* ctaIdxXyToMnLimit,
+      int32_t* numNonExitingCtas,
+      batchedGemm::trtllm::gen::Dtype dtypeElt,
+      batchedGemm::trtllm::gen::Dtype dtypeBias,
+      bool useRoutingScalesOnInput,
+      bool useDeepSeekFp8,
+      RoutingMethodType routingMethodType,
+      cudaStream_t stream,
+      batchedGemm::trtllm::gen::Dtype dtypeLogits,
+      bool normTopkProb = true,
+      int16_t* routing_replay_out = nullptr);
+
+ private:
+  friend class MoE::Runner;
+  int32_t mTileTokensDim{8};
+};
+}  // namespace Routing
+
+namespace MoE {
+// The type of activation function
+// Please keep this in sync with the counterpart defined in flashinfer/flashinfer/fused_moe/core.py
+enum class ActivationType : int64_t {
+  Gelu = 0,
+  Relu = 1,
+  Silu = 2,
+  Swiglu = 3,
+  Geglu = 4,
+  SwigluBias = 5,
+  Relu2 = 6,
+  Identity = 7,
+  InvalidType = 8,  // Must be last
+};
+
+inline std::string serializeActivationType(ActivationType activationType) {
+  switch (activationType) {
+    case ActivationType::Gelu:
+      return "Gelu";
+    case ActivationType::Relu:
+      return "Relu";
+    case ActivationType::Silu:
+      return "Silu";
+    case ActivationType::Swiglu:
+      return "Swiglu";
+    case ActivationType::Geglu:
+      return "Geglu";
+    case ActivationType::SwigluBias:
+      return "SwigluBias";
+    case ActivationType::Relu2:
+      return "Relu2";
+    case ActivationType::Identity:
+      return "Identity";
+    default:
+      return "InvalidActivationType";  // TODO throw error
+  };
+}
+
+inline bool isGatedActivation(ActivationType activationType) {
+  return activationType == ActivationType::Swiglu || activationType == ActivationType::Geglu ||
+         activationType == ActivationType::SwigluBias;
+}
+
+}  // namespace MoE
+
+namespace PermuteGemm1 {
+class Runner {
+ public:
+  explicit Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeAct,
+      batchedGemm::trtllm::gen::Dtype dtypeWeights,
+      batchedGemm::trtllm::gen::Dtype dtypeOutput,
+      bool useDeepSeekFp8,
+      int tileTokensDim,
+      MoE::ActivationType activationType,
+      bool useShuffledMatrix,
+      batchedGemm::gemm::MatrixLayout weight_layout,
+      bool usePerTokenScaling,
+      bool usePerChannelScaling,
+      bool forceUnfusedAct = false);
+
+  size_t getWorkspaceSizeInBytes(
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t configIndex) const;
+
+  [[nodiscard]] int32_t getDefaultValidConfigIndex(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
+
+  [[nodiscard]] bool isValidConfigIndex(
+      int32_t configIndex,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens) const;
+
+  [[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
+
+  void
+  run(void* hiddenState,
+      void* hiddenStateScale,
+      void* weight,
+      void* weightScale,
+      void* perTokenScales,
+      void* perChannelScales,
+      float* outputScalesScalar,
+      float* outputScalesGateScalar,
+      float* ptrBias,
+      float* ptrGatedActAlpha,
+      float* ptrGatedActBeta,
+      float* ptrClampLimit,
+      void* output,
+      void* outputScale,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t* permutedIdxToTokenIdx,
+      int32_t* ptrNumNonExitingCtas,
+      int32_t* ptrTotalNumPaddedTokens,
+      int32_t* ptrCtaIdxXyToBatchIdx,
+      int32_t* ptrCtaIdxXyToMnLimit,
+      void* bmm1Workspace,
+      bool useRoutingScalesOnInput,
+      int device,
+      cudaStream_t stream,
+      int32_t configIndex,
+      bool enable_pdl);
+
+ private:
+  friend class MoE::Runner;
+  batchedGemm::trtllm::gen::Dtype mDtypeAct;
+  batchedGemm::trtllm::gen::Dtype mDtypeWeights;
+  batchedGemm::trtllm::gen::Dtype mDtypeOutput;
+  int32_t mTileTokensDim;
+  tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
+  tensorrt_llm::kernels::trtllmgen_moe::MoE::ActivationType mActType;
+};
+}  // namespace PermuteGemm1
+
+namespace Gemm2 {
+class Runner {
+ public:
+  explicit Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeAct,
+      batchedGemm::trtllm::gen::Dtype dtypeWeights,
+      batchedGemm::trtllm::gen::Dtype outputDtype,
+      bool useDeepSeekFp8,
+      int tileTokensDim,
+      bool useShuffledMatrix,
+      batchedGemm::gemm::MatrixLayout weight_layout,
+      bool usePerTokenScaling,
+      bool usePerChannelScaling);
+
+  size_t getWorkspaceSizeInBytes(
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t configIndex) const;
+
+  [[nodiscard]] int32_t getDefaultValidConfigIndex(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
+
+  [[nodiscard]] bool isValidConfigIndex(
+      int32_t configIndex,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens) const;
+
+  [[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
+
+  void
+  run(void* permutedHiddenState,
+      void* permutedHiddenStateScale,
+      void* weight,
+      void* weightScale,
+      void* perTokenScales,
+      void* perChannelScales,
+      float* outputScalesScalar,
+      float* ptrBias,
+      void* output,
+      void* outputScale,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t* ptrNumNonExitingCtas,
+      int32_t* ptrTotalNumPaddedTokens,
+      int32_t* ptrCtaIdxXyToBatchIdx,
+      int32_t* ptrCtaIdxXyToMnLimit,
+      void* bmm2Workspace,
+      int device,
+      cudaStream_t stream,
+      int32_t configIndex,
+      bool enable_pdl);
+
+ private:
+  friend class MoE::Runner;
+  batchedGemm::trtllm::gen::Dtype mDtypeAct;
+  batchedGemm::trtllm::gen::Dtype mDtypeWeights;
+  batchedGemm::trtllm::gen::Dtype mDtypeOut;
+  int32_t mTileTokensDim;
+  tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
+};
+}  // namespace Gemm2
+
+namespace MoE {
+namespace btg = batchedGemm::trtllm::gen;
+
+struct MoERunnerArgs {
+  void* routing_logits = nullptr;  // [num_tokens, num_experts] in float, generated after
+                                   // gemm(hidden_state, routing_weights)
+  void* routing_bias = nullptr;    // [num_experts] in bfloat16 for now = mDtypeExpW
+  void* hidden_states = nullptr;   // [num_tokens, hidden_size] in fp8 = mDtypeElt
+  // [hidden_size/128, num_tokens] in float for e4m3 DS recipe
+  // and [num_tokens, hidden_size/16] in float for e2m1
+  void* hidden_states_scale = nullptr;
+
+  // Gemm input:
+  void* gemm1_weights = nullptr;
+  void* gemm1_weights_scale = nullptr;
+  void* gemm2_weights = nullptr;
+  void* gemm2_weights_scale = nullptr;
+
+  float* gemm1_bias = nullptr;
+  float* gemm1_alpha = nullptr;
+  float* gemm1_beta = nullptr;
+  float* gemm1_clamp_limit = nullptr;
+  float* gemm2_bias = nullptr;
+
+  ActivationType activation_type = ActivationType::Swiglu;
+
+  int32_t num_tokens{0};
+  int32_t num_experts{0};
+  // Hidden dimension input of MoE block. It might be padded.
+  int32_t hidden_size{0};
+  // Hidden dimension output of MoE block. It is not padded.
+  // If not provided it is the same as hidden_size.
+  std::optional<int32_t> hidden_size_output;
+  // TODO: only compiled routing kernel supports top_k = 8
+  int32_t top_k{0};
+  int32_t n_group{0};
+  // TODO: only compiled routing kernel supports topk_group = 4
+  int32_t topk_group{0};
+  float routed_scaling_factor{0.0f};
+  int32_t intermediate_size{0};
+  int32_t local_expert_offset{0};
+  int32_t local_num_experts{0};
+  // TODO: support other types
+  btg::Dtype mDtypeElt{btg::Dtype::Void};
+  btg::Dtype mDtypeExpW{btg::Dtype::Bfloat16};
+  btg::Dtype mDtypeOut{btg::Dtype::Bfloat16};
+
+  // Apply routing scale factors to input activations
+  bool mUseRoutingScalesOnInput{false};
+  bool mUseDeepSeekFp8{false};
+  float* output1_scales_scalar = nullptr;
+  float* output1_scales_gate_scalar = nullptr;
+  float* output2_scales_scalar = nullptr;
+
+  // Optional LoRA bridge buffers used by the copied SGLang TRTLLM FP8 path.
+  // gate_up_lora_delta: [num_tokens * top_k, 2 * intermediate_size], bf16,
+  // in FlashInfer gate/up order (up first, gate second).
+  // activation_lora_input: [num_tokens * top_k, intermediate_size], bf16,
+  // populated with the post-activation intermediate for down-proj LoRA.
+  void* gate_up_lora_delta = nullptr;
+  void* activation_lora_input = nullptr;
+
+  // Optional CUDA event (cudaEvent_t) recorded on the LoRA side stream. When set, the
+  // runner waits on it right before the activation kernel (which consumes
+  // gate_up_lora_delta), so permute+GEMM1 overlap the side-stream LoRA shrink/expand
+  // instead of joining before the whole MoE op. nullptr = no wait (serial behavior).
+  void* lora_ready_event = nullptr;
+
+  // Down-LoRA/finalize overlap: optional CUDA event (cudaEvent_t) the runner records on the
+  // MoE stream right after GEMM2 (the base down GEMM), before finalize. The LoRA side stream
+  // waits on it to run the down-proj LoRA shrink/expand concurrent with the finalize kernel.
+  // nullptr = no record (serial behavior).
+  void* gemm2_done_event = nullptr;
+
+  // Output:
+  void* output = nullptr;
+  float* output_scale = nullptr;
+
+  // finalize
+  bool do_finalize{true};
+};
+
+struct MoEWorkspace {
+  // Routing intermediate outputs:
+  int32_t* routing_expert_indexes = nullptr;
+  int32_t* permuted_idx_size = nullptr;
+  int32_t* total_num_padded_tokens = nullptr;  // TODO: duplicate of permuted_idx_size
+  int32_t total_max_padded_tokens{0};
+
+  int32_t* expanded_idx_to_permuted_idx = nullptr;
+  int32_t* permuted_idx_to_expanded_idx = nullptr;
+  int32_t* permuted_idx_to_token_idx = nullptr;
+
+  // consumed by finalize kernel
+  void* expert_weights = nullptr;  // [num_tokens, top_k] in bfloat16 = mDtypeExpW
+  // consumed by permuteGemm1 kernel
+  void* token_scales = nullptr;
+  // consumed by Gemm2 kernel
+  void* token_scales_fc2 = nullptr;
+
+  int32_t* cta_idx_xy_to_batch_idx = nullptr;
+  int32_t* cta_idx_xy_to_mn_limit = nullptr;
+  int32_t* num_non_exiting_ctas = nullptr;
+
+  void* hidden_states_scale_linear = nullptr;
+
+  // Permute intermediate outputs:
+  void* permuted_hidden_states = nullptr;
+  float* permuted_hidden_states_scale = nullptr;
+
+  // Gemm1 intermediate outputs:
+  int32_t ProjUpTileN{0};
+  void* gemm1_output = nullptr;
+  float* gemm1_output_scale = nullptr;
+
+  // Activation intermediate outputs:
+  void* activation_output = nullptr;
+  float* activation_output_scale = nullptr;
+  // Unfused FP4 LoRA: bf16 [max_padded_tokens, intermediate_size] activation output written by the
+  // standalone activation kernel (gate_up LoRA added pre-SwiGLU), then NvFP4-quantized for GEMM2.
+  void* activated_lora_bf16 = nullptr;
+
+  // Gemm2 intermediate outputs:
+  void* gemm2_output = nullptr;
+  float* gemm2_output_scale = nullptr;
+
+  // Finalize intermediate outputs (placeholder not used)
+  void* finalize_output = nullptr;
+  float* finalize_output_scale = nullptr;
+
+  // FC1 workspace:
+  void* bmm1_workspace = nullptr;
+
+  // FC2 workspace:
+  void* bmm2_workspace = nullptr;
+};
+
+// Config indices to be used with Batched GEMM runners
+struct MoEConfig {
+  int64_t gemm1Config;
+  int64_t gemm2Config;
+};
+
+class Runner {
+ public:
+  // FIXME: tileTokensDim is hardcoded for now
+  Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeAct,
+      batchedGemm::trtllm::gen::Dtype dtypeWeights,
+      bool useDeepSeekFp8,
+      int tileTokensDim = 8,
+      ActivationType activationType = ActivationType::Swiglu,
+      bool useShuffledMatrix = false,
+      batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
+      bool usePerTokenScalingGemm1 = false,
+      bool usePerTokenScalingGemm2 = false,
+      bool usePerChannelScalingGemm1 = false,
+      bool usePerChannelScalingGemm2 = false,
+      bool unfuseActForLora = false);
+  Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeElt,
+      bool useDeepSeekFp8,
+      int tileTokensDim = 8,
+      bool useShuffledMatrix = false,
+      batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
+      bool usePerTokenScalingGemm1 = false,
+      bool usePerTokenScalingGemm2 = false,
+      bool usePerChannelScalingGemm1 = false,
+      bool usePerChannelScalingGemm2 = false);
+
+  void
+  run(MoERunnerArgs const& args,
+      MoEWorkspace const& workspace,
+      int device,
+      cudaStream_t stream,
+      int64_t configIndex,
+      bool enable_pdl);
+
+  [[nodiscard]] std::tuple<int32_t, int32_t>
+  getWorkspaceSizeInBytes(MoERunnerArgs const& args, int64_t configIndex) const;
+
+  [[nodiscard]] std::vector<int64_t> getValidConfigIndices(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
+
+  [[nodiscard]] int64_t getDefaultValidConfigIndex(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
+
+ private:
+  void setOpsData(
+      MoERunnerArgs const& args,
+      MoEWorkspace const& workspace,
+      moe::dev::convertsf::Data& convertSfData,
+      moe::dev::activation::Data& activationData,
+      moe::dev::finalize::Data& finalizeData);
+
+ private:
+  bool mUsePerTokenScalingGemm1;
+  bool mUsePerTokenScalingGemm2;
+  bool mUsePerChannelScalingGemm1;
+  bool mUsePerChannelScalingGemm2;
+  // When true (FP4 LoRA path), GEMM1 emits the raw gate_up projection (fusedAct=false) so the
+  // standalone activation kernel can inject the gate_up LoRA delta pre-SwiGLU and capture the
+  // post-activation input for the down LoRA — mirroring the DeepSeek-FP8 unfused activation.
+  bool mUnfuseActForLora;
+  PermuteGemm1::Runner mPermuteGemm1;
+  Gemm2::Runner mGemm2;
+
+  // This will be the cartesian product of the passing configs for gemm1 and gemm2
+  // This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2
+  // separately
+  std::vector<MoEConfig> mPassingConfigs;
+};
+}  // namespace MoE
+
+}  // namespace trtllmgen_moe
+}  // namespace kernels
+}  // namespace tensorrt_llm