chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/gemm/awq_dequantize.cuh

@@ -0,0 +1,227 @@
+// Adapted from
+// https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L350
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+namespace device::awq {
+
+template <int lut>
+__device__ inline int lop3(int a, int b, int c) {
+  int res;
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut));
+  return res;
+}
+
+__device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750
+  uint4 result;
+
+  uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+  uint32_t const i4s = reinterpret_cast<uint32_t const&>(source);
+
+  // First, we extract the i4s and construct an intermediate fp16 number.
+  static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
+  static constexpr uint32_t BOTTOM_MASK = 0x000f000f;
+  static constexpr uint32_t TOP_MASK = 0x00f000f0;
+  static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400;
+
+  // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW
+  // dependency if we issue immediately before required.
+  const uint32_t top_i4s = i4s >> 8;
+  // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+               : "=r"(h[0])
+               : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
+  // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+               : "=r"(h[1])
+               : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
+  // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+               : "=r"(h[2])
+               : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
+  // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+               : "=r"(h[3])
+               : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
+
+  // This is the half2 {1024, 1024} represented as an integer.
+  static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400;
+  // This is the half2 {1 / 16, 1 / 16} represented as an integer.
+  static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00;
+  // This is the half2 {-64, -64} represented as an integer.
+  static constexpr uint32_t NEG_64 = 0xd400d400;
+
+  // Finally, we construct the output numbers.
+  // Convert elt_01
+  asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM));
+  // Convert elt_23
+  asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64));
+  // Convert elt_45
+  asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM));
+  // Convert elt_67
+  asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64));
+
+  return result;
+#else
+  assert(false);
+  return {};
+#endif
+}
+
+__device__ uint4 dequantize_s4_to_bf16x2(uint32_t const& source) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  uint4 result;
+  uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+  uint32_t const i4s = source;
+
+  // Define masks and constants
+  static constexpr uint32_t MASK = 0x000f000f;
+  static constexpr uint32_t EX = 0x43004300;
+  static constexpr uint32_t MUL = 0x3F803F80;
+  static constexpr uint32_t ADD = 0xC300C300;
+
+  int lo0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s, MASK, EX);
+  int hi0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 4, MASK, EX);
+  int lo1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 8, MASK, EX);
+  int hi1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 12, MASK, EX);
+
+  nv_bfloat162* res = reinterpret_cast<nv_bfloat162*>(h);
+  res[0] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&lo0),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+  res[1] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&hi0),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+  res[2] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&lo1),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+  res[3] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&hi1),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+
+  return result;
+#else
+  assert(false);
+  return {};
+#endif
+}
+
+template <typename OutputT>
+__global__ void __launch_bounds__(256) dequantize_weights(
+    int* __restrict__ qweight,
+    OutputT* __restrict__ scales,
+    int* __restrict__ qzeros,
+    OutputT* __restrict__ output,
+    int group_size,
+    int qweight_cols,
+    int qweight_rows) {
+  int col = blockIdx.x * blockDim.x + threadIdx.x;
+  int row = blockIdx.y * blockDim.y + threadIdx.y;
+  if (col >= qweight_cols || row >= qweight_rows) return;
+
+  int group_idx = row / group_size;
+  int scale_offset = 8 * col + group_idx * qweight_cols * 8;
+  uint4 loaded_scale = *(uint4*)(scales + scale_offset);
+
+  // Handle different data types
+  if constexpr (std::is_same<OutputT, half>::value) {
+    // FP16 path
+    uint4 zeros = dequantize_s4_to_fp16x2(qzeros[col + group_idx * qweight_cols]);
+    uint4 weight_fp16 = dequantize_s4_to_fp16x2(qweight[col + row * qweight_cols]);
+
+    // Use PTX assembly for FP16 operations
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(zeros.x));
+    asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(loaded_scale.x));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(zeros.y));
+    asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(loaded_scale.y));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(zeros.z));
+    asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(loaded_scale.z));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(zeros.w));
+    asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(loaded_scale.w));
+
+    OutputT* output_ptr = output + 8 * col + 8 * row * qweight_cols;
+    *(uint4*)output_ptr = weight_fp16;
+  } else if constexpr (std::is_same<OutputT, __nv_bfloat16>::value) {
+    uint4 weight_raw = dequantize_s4_to_bf16x2(qweight[col + row * qweight_cols]);
+    uint4 zero_raw = dequantize_s4_to_bf16x2(qzeros[col + group_idx * qweight_cols]);
+    uint4 scale_raw = *reinterpret_cast<uint4*>(scales + scale_offset);
+
+    // Vectorized processing (each uint4 contains 4 nv_bfloat162)
+    nv_bfloat162* weight_vec = reinterpret_cast<nv_bfloat162*>(&weight_raw);
+    nv_bfloat162* zero_vec = reinterpret_cast<nv_bfloat162*>(&zero_raw);
+    nv_bfloat162* scale_vec = reinterpret_cast<nv_bfloat162*>(&scale_raw);
+
+// Single instruction dual-channel operation
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {  // uint4 = 4 * nv_bfloat162
+      weight_vec[i] = __hmul2(__hsub2(weight_vec[i], zero_vec[i]), scale_vec[i]);
+    }
+
+    // Directly store to OutputT array (guaranteed contiguous memory)
+    OutputT* output_ptr = output + 8 * col + row * qweight_cols * 8;
+    static_assert(sizeof(uint4) == 8 * sizeof(OutputT), "Memory layout mismatch");
+    *reinterpret_cast<uint4*>(output_ptr) = weight_raw;
+  }
+}
+
+}  // namespace device::awq
+
+// Host wrapper
+template <typename OutputT>
+void awq_dequantize(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView qweight,
+    tvm::ffi::TensorView scales,
+    tvm::ffi::TensorView qzeros) {
+  using namespace host;
+
+  int64_t qweight_rows = qweight.size(0);
+  int64_t qweight_cols = qweight.size(1);
+  int64_t scales_rows = scales.size(0);
+
+  // Validate tensors
+  SymbolicDevice cuda_device;
+  cuda_device.set_options<kDLCUDA>();
+
+  TensorMatcher({qweight_rows, qweight_cols}).with_dtype<int32_t>().with_device(cuda_device).verify(qweight);
+  TensorMatcher({scales_rows, qweight_cols * 8}).with_dtype<OutputT>().with_device(cuda_device).verify(scales);
+  TensorMatcher({scales_rows, qweight_cols}).with_dtype<int32_t>().with_device(cuda_device).verify(qzeros);
+  TensorMatcher({qweight_rows, qweight_cols * 8}).with_dtype<OutputT>().with_device(cuda_device).verify(output);
+
+  // Get device and stream
+  auto device = cuda_device.unwrap();
+  auto stream = LaunchKernel::resolve_device(device);
+
+  int group_size = static_cast<int>(qweight_rows / scales_rows);
+  int x_num_threads = 16;
+  int y_num_threads = 16;
+  int x_blocks = (static_cast<int>(qweight_cols) + x_num_threads - 1) / x_num_threads;
+  int y_blocks = (static_cast<int>(qweight_rows) + y_num_threads - 1) / y_num_threads;
+
+  dim3 num_blocks(x_blocks, y_blocks);
+  dim3 threads_per_block(x_num_threads, y_num_threads);
+
+  // Get pointers
+  auto* qweight_ptr = reinterpret_cast<int*>(qweight.data_ptr());
+  auto* scales_ptr = reinterpret_cast<OutputT*>(scales.data_ptr());
+  auto* qzeros_ptr = reinterpret_cast<int*>(qzeros.data_ptr());
+  auto* output_ptr = reinterpret_cast<OutputT*>(output.data_ptr());
+
+  LaunchKernel(num_blocks, threads_per_block, stream)(
+      device::awq::dequantize_weights<OutputT>,
+      qweight_ptr,
+      scales_ptr,
+      qzeros_ptr,
+      output_ptr,
+      group_size,
+      static_cast<int>(qweight_cols),
+      static_cast<int>(qweight_rows));
+}

python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh

@@ -0,0 +1,631 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu
+ *
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2021, NAVER Corp.  Authored by CLOVA.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>
+#include <cuda_bf16.h>
+#include <utility>
+
+namespace {
+
+using bf16_t = __nv_bfloat16;
+
+__device__ void hmma_16_8_16_f32acc_bf16ab(
+    float (&d_reg)[4], const bf16_t (&a_reg)[8], const bf16_t (&b_reg)[4], float const (&c_reg)[4]) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t a0 = *reinterpret_cast<uint32_t const*>(a_reg + 0);
+  uint32_t a1 = *reinterpret_cast<uint32_t const*>(a_reg + 2);
+  uint32_t a2 = *reinterpret_cast<uint32_t const*>(a_reg + 4);
+  uint32_t a3 = *reinterpret_cast<uint32_t const*>(a_reg + 6);
+  uint32_t b0 = *reinterpret_cast<uint32_t const*>(b_reg + 0);
+  uint32_t b1 = *reinterpret_cast<uint32_t const*>(b_reg + 2);
+  asm volatile(
+      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+      "{%0,  %1,  %2,  %3},"
+      "{%4,  %5,  %6,  %7},"
+      "{%8,  %9},"
+      "{%10, %11, %12, %13};\n"
+      : "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3])
+      : "r"(a0),
+        "r"(a1),
+        "r"(a2),
+        "r"(a3),
+        "r"(b0),
+        "r"(b1),
+        "f"(d_reg[0]),
+        "f"(d_reg[1]),
+        "f"(d_reg[2]),
+        "f"(d_reg[3]));
+#endif
+}
+
+extern "C" {
+__device__ uint32_t __nvvm_get_smem_pointer(void*);
+}
+
+__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  if (pred) {
+    uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr);
+    asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(smemPtrAsUint32), "l"(gPtr), "n"(16));
+  }
+#endif
+}
+
+__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+               : "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3])
+               : "r"(__nvvm_get_smem_pointer(smem_ptr)));
+#endif
+}
+
+template <class Type>
+__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) {
+  uint32_t row_idx = *reinterpret_cast<uint32_t*>(&row_idx_);
+  uint32_t col_idx = *reinterpret_cast<uint32_t*>(&col_idx_);
+  row_idx = row_idx % 8;
+  row_idx = row_idx * (16 / sizeof(Type));
+  col_idx = col_idx ^ row_idx;
+  return *reinterpret_cast<int*>(&col_idx);
+}
+
+__device__ void initialize_barrier(uint64_t* smem_barrier, int thread_count = 1) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr), "r"(thread_count));
+#endif
+}
+
+__device__ void wait_barrier(uint64_t* smem_barrier, int phase_bit) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile(
+      "{\n"
+      ".reg .pred                P1;\n"
+      "LAB_WAIT:\n"
+      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
+      "@P1                       bra DONE;\n"
+      "bra                   LAB_WAIT;\n"
+      "DONE:\n"
+      "}\n" ::"r"(smem_int_ptr),
+      "r"(phase_bit));
+#endif
+}
+
+__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t wait_complete;
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr);
+  asm volatile(
+      "{\n\t"
+      ".reg .pred P1; \n\t"
+      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
+      "selp.b32 %0, 1, 0, P1; \n\t"
+      "}"
+      : "=r"(wait_complete)
+      : "r"(smem_int_ptr), "r"(phase_bit));
+  return static_cast<bool>(wait_complete);
+#endif
+  return false;
+}
+
+__device__ void arrive_barrier(uint64_t* smem_barrier) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile(
+      "{\n"
+      ".reg .b64 state; \n"
+      "mbarrier.arrive.shared::cta.b64   state, [%0];\n"
+      "}\n" ::"r"(smem_int_ptr));
+#endif
+}
+
+__device__ void ldgsts_arrive(uint64_t* smem_barrier) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];" : : "r"(smem_int_ptr));
+#endif
+}
+
+template <int gemm_k, int tile_m, int tile_k, int stage_cnt>
+struct GmemLoaderA {
+  static constexpr int elem_bytes = 2;
+  static constexpr int vec_bytes = 16;
+  static constexpr int vec_elems = vec_bytes / elem_bytes;
+  static constexpr int thread_cnt = 64;
+  static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0);
+  static constexpr int a_inst_cnt_per_iter = (tile_m * tile_k) / (vec_elems * thread_cnt);
+  static_assert(gemm_k % tile_k == 0);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+
+  static constexpr int mma_warp_cnt = 4;
+  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
+  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
+
+ private:
+  __device__ int k_project(int tile_k_idx) {
+    return (tile_k_idx / per_mma_warp_k * k_each_chunk) + (tile_k_idx % per_mma_warp_k);
+  }
+
+ public:
+  __device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_, uint64_t* smem_barrier_)
+      : gmem_a(gmem_a_local_), smem_a(smem_a_), smem_barrier(smem_barrier_), local_tid(threadIdx.x % thread_cnt) {}
+
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+#pragma unroll
+    for (int i = 0; i < a_inst_cnt_per_iter; i++) {
+      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+      int m_idx = linear_idx / tile_k;
+      int k_idx = linear_idx % tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
+      a_smem_offsets[i] = m_idx * tile_k + k_idx;
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+#pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      if (need_wait) {
+        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
+      }
+      int next_stage_idx = stage_idx + 1;
+      int next_phase_bit = next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
+      if (loop_idx != k_iter_cnt - 1) {
+        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2, next_phase_bit);
+      }
+
+#pragma unroll
+      for (int i = 0; i < a_inst_cnt_per_iter; i++) {
+        int smem_offset = a_smem_offsets[i];
+        bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset;
+        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+        int m_idx = linear_idx / tile_k;
+        int k_idx = linear_idx % tile_k;
+        int gmem_offset = m_idx * gemm_k + k_project(k_idx);
+        bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset;
+        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true);
+      }
+      ldgsts_arrive(smem_barrier + stage_idx * 2);
+
+      stage_idx = next_stage_idx;
+      phase_bit = next_phase_bit;
+      gmem_a += per_mma_warp_k;
+    }
+#endif
+  }
+
+  bf16_t const* gmem_a;
+  bf16_t* smem_a;
+  uint64_t* smem_barrier;
+  int local_tid;
+  int stage_idx = 0;
+  int phase_bit = 1;
+  bool need_wait = true;
+
+  int a_smem_offsets[a_inst_cnt_per_iter];
+};
+
+template <int gemm_k, int tile_n, int tile_k, int stage_cnt>
+struct GmemLoaderB {
+  static constexpr int elem_bytes = 2;
+  static constexpr int vec_bytes = 16;
+  static constexpr int vec_elems = vec_bytes / elem_bytes;
+  static constexpr int thread_cnt = 64;
+  static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0);
+  static constexpr int b_inst_cnt_per_iter = (tile_n * tile_k) / (vec_elems * thread_cnt);
+  static_assert(gemm_k % tile_k == 0);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+
+  static constexpr int mma_warp_cnt = 4;
+  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
+  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
+
+ private:
+  __device__ int k_project(int tile_k_idx) {
+    return (tile_k_idx / per_mma_warp_k * k_each_chunk) + (tile_k_idx % per_mma_warp_k);
+  }
+
+ public:
+  __device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_, uint64_t* smem_barrier_, int gemm_n_)
+      : gmem_b(gmem_b_local_),
+        smem_b(smem_b_),
+        smem_barrier(smem_barrier_),
+        gemm_n(gemm_n_),
+        local_tid(threadIdx.x % thread_cnt) {}
+
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+#pragma unroll
+    for (int i = 0; i < b_inst_cnt_per_iter; i++) {
+      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+      int n_idx = linear_idx / tile_k;
+      int k_idx = linear_idx % tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
+      b_smem_offsets[i] = n_idx * tile_k + k_idx;
+      preds[i] = n_idx < gemm_n;
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    cudaGridDependencySynchronize();
+#pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      if (need_wait) {
+        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
+      }
+      int next_stage_idx = stage_idx + 1;
+      int next_phase_bit = next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
+      if (loop_idx != k_iter_cnt - 1) {
+        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2, next_phase_bit);
+      }
+#pragma unroll
+      for (int i = 0; i < b_inst_cnt_per_iter; i++) {
+        int smem_offset = b_smem_offsets[i];
+        bf16_t* smem_ptr_this_iter = smem_b + stage_idx * tile_n * tile_k + smem_offset;
+        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+        int n_idx = linear_idx / tile_k;
+        int k_idx = linear_idx % tile_k;
+        int gmem_offset = n_idx * gemm_k + k_project(k_idx);
+        bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
+        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
+      }
+      ldgsts_arrive(smem_barrier + stage_idx * 2);
+
+      stage_idx = next_stage_idx;
+      phase_bit = next_phase_bit;
+      gmem_b += per_mma_warp_k;
+    }
+#endif
+  }
+
+  bf16_t const* gmem_b;
+  bf16_t* smem_b;
+  uint64_t* smem_barrier;
+  int gemm_n;
+  int local_tid;
+  int stage_idx = 0;
+  int phase_bit = 1;
+  bool need_wait = true;
+
+  int b_smem_offsets[b_inst_cnt_per_iter];
+  uint32_t preds[b_inst_cnt_per_iter];
+};
+
+template <int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k, int stage_cnt>
+struct MmaComputer {
+  static constexpr int elem_bytes = 2;
+  static constexpr int thread_cnt = 128;
+  static_assert(gemm_k % tile_k == 0);
+  static_assert(tile_k % (thread_cnt / 32) == 0);
+  static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+  static constexpr int k_phase_cnt = per_warp_tile_k / 16;
+  static constexpr int m_iter_cnt = (tile_m + 15) / 16;
+  static constexpr int n_iter_cnt = (tile_n + 7) / 8;
+  static_assert(m_iter_cnt == 1);
+  static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
+
+  __device__ MmaComputer(
+      bf16_t* gmem_c_local_, bf16_t* smem_a_, bf16_t* smem_b_, uint64_t* smem_barrier_, int warp_idx_, int gemm_n_)
+      : gmem_c(gmem_c_local_),
+        smem_a(smem_a_),
+        smem_b(smem_b_),
+        smem_barrier(smem_barrier_),
+        warp_idx(warp_idx_ - (thread_cnt / 32)),
+        gemm_n(gemm_n_) {}
+
+ private:
+  __device__ constexpr int internal_b_atom_func(int tid) {
+    if constexpr (tile_n < 8) {
+      return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n;
+    } else {
+      return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8));
+    }
+  }
+
+ public:
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+#pragma unroll
+    for (int i = 0; i < k_phase_cnt; i++) {
+      int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
+      int m_idx = linear_idx % tile_m;
+      int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
+      a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
+    }
+#pragma unroll
+    for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+#pragma unroll
+      for (int i = 0; i < k_phase_cnt; i += 2) {
+        int linear_idx = internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8;
+        int n_idx = linear_idx % tile_n;
+        int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k;
+        k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
+        b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx;
+      }
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+#pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
+
+#pragma unroll
+      for (int i = 0; i < k_phase_cnt; i++) {
+        int smem_offset = a_smem_offsets[0][i];
+        bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset;
+        ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
+      }
+
+#pragma unroll
+      for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+#pragma unroll
+        for (int i = 0; i < k_phase_cnt; i += 2) {
+          int smem_offset = b_smem_offsets[n_iter_idx][i];
+          bf16_t* smem_ptr_this_iter = smem_b + stage_idx * tile_n * tile_k + smem_offset;
+          ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(b_reg[n_iter_idx][i]));
+        }
+      }
+
+#pragma unroll
+      for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
+#pragma unroll
+        for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+          hmma_16_8_16_f32acc_bf16ab(
+              acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx], b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
+        }
+      }
+      ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
+      stage_idx += 1;
+      phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
+    }
+#endif
+  }
+
+  __device__ void epi() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
+    constexpr int thread_m = 2;
+    constexpr int thread_n = 2 * n_iter_cnt;
+    constexpr int cta_mma_n = n_iter_cnt * 8;
+    float acc_reg_reorg[thread_m][thread_n];
+
+    for (int i = 0; i < thread_m; i++) {
+      for (int j = 0; j < thread_n; j++) {
+        acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
+      }
+    }
+
+    float* smem_c = reinterpret_cast<float*>(smem_a);
+    auto smem_c_index_func = [&](int m_idx, int n_idx) {
+      int group_rows = 32 / cta_mma_n;
+      int group_cnt = 2;
+      return (m_idx % group_rows * cta_mma_n) + (m_idx / group_rows * (32 + group_cnt)) + n_idx;
+    };
+    constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2);
+
+#pragma unroll
+    for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
+#pragma unroll
+      for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
+        int m_idx = (lane_idx / 4) + m_idx_thread * 8;
+        int n_idx = ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
+        smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] = acc_reg_reorg[m_idx_thread][n_idx_thread];
+      }
+    }
+    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
+
+    if (warp_idx == 0) {
+      constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32;
+      float acc_final[final_acc_reg_cnt]{};
+
+#pragma unroll
+      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
+        int linear_idx = reg_idx * 32 + lane_idx;
+        int m_idx = linear_idx % tile_m;
+        int n_idx = linear_idx / tile_m;
+        acc_final[reg_idx] += smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] +
+                              smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] +
+                              smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] +
+                              smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c];
+      }
+
+#pragma unroll
+      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
+        int linear_idx = reg_idx * 32 + lane_idx;
+        int m_idx = linear_idx % tile_m;
+        int n_idx = linear_idx / tile_m;
+        if (m_idx < tile_m && n_idx < gemm_n) {
+          gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx];
+        }
+      }
+    }
+#endif
+  }
+
+  bf16_t* gmem_c;
+  bf16_t* smem_a;
+  bf16_t* smem_b;
+  uint64_t* smem_barrier;
+  int warp_idx;
+  int gemm_n;
+  int stage_idx = 0;
+  int phase_bit = 0;
+  int lane_idx = threadIdx.x % 32;
+  int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k;
+
+  int a_smem_offsets[m_iter_cnt][k_phase_cnt];
+  int b_smem_offsets[n_iter_cnt][k_phase_cnt];
+
+  bf16_t a_reg[m_iter_cnt][k_phase_cnt][8];
+  bf16_t b_reg[n_iter_cnt][k_phase_cnt][4];
+  float acc_reg[m_iter_cnt][n_iter_cnt][4]{};
+};
+
+template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k, int stage_cnt>
+__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
+    bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  constexpr int load_thread_cnt = 128;
+  constexpr int compute_thread_cnt = 128;
+  constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt;
+  (void)thread_cnt;
+  static_assert(gemm_m % 16 == 0);
+  static_assert(gemm_k % tile_k == 0);
+  static_assert(gemm_m % tile_m == 0);
+  static_assert(tile_k == 128 || tile_k == 256 || tile_k == 512 || tile_k == 1024);
+  static_assert(tile_m == 16);
+  constexpr int g2s_vec_bytes = 16;
+  constexpr int a_elem_bytes = 2;
+  constexpr int b_elem_bytes = 2;
+  static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k * stage_cnt <= 225 * 1024);
+  static_assert((tile_m * tile_k * a_elem_bytes) % (load_thread_cnt * g2s_vec_bytes) == 0);
+  static_assert((tile_n * tile_k * b_elem_bytes) % (load_thread_cnt * g2s_vec_bytes) == 0);
+
+  extern __shared__ char smem[];
+  uint64_t* smem_barrier = reinterpret_cast<uint64_t*>(smem);
+  bf16_t* smem_a = reinterpret_cast<bf16_t*>(smem + (stage_cnt * 8 * 2 + 1024) / 1024 * 1024);
+  bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt;
+
+  int cta_m_idx = tile_m * blockIdx.x;
+  int cta_n_idx = tile_n * blockIdx.y;
+  bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
+  bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
+  bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
+
+  int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+
+  if (warp_idx == 4) {
+    for (int i = 0; i < stage_cnt; i++) {
+      initialize_barrier(smem_barrier + i * 2 + 0, load_thread_cnt);
+      initialize_barrier(smem_barrier + i * 2 + 1, compute_thread_cnt);
+    }
+  }
+  __syncthreads();
+
+  if (warp_idx < 2) {
+    GmemLoaderA<gemm_k, tile_m, tile_k, stage_cnt> a_loader(gmem_a_local, smem_a, smem_barrier);
+    a_loader.prepare();
+    a_loader.issue_mainloop();
+  } else if (warp_idx < 4) {
+    GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(gmem_b_local, smem_b, smem_barrier, gemm_n);
+    b_loader.prepare();
+    b_loader.issue_mainloop();
+  } else {
+    MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
+        gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
+    mma_computer.prepare();
+    mma_computer.issue_mainloop();
+    mma_computer.epi();
+  }
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+template <typename T, int kHdIn, int kHdOut, int kTileN, bool kUsePDL>
+void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, DLDevice device) {
+  constexpr int gemm_m = kHdOut;  // 2112
+  int const gemm_n = num_tokens;  // 16
+  constexpr int gemm_k = kHdIn;   // 7168
+  constexpr int batch_size = 1;
+  std::swap(mat_a, mat_b);
+  constexpr int tile_m = 16;
+  constexpr int tile_n = kTileN;                        // 8 or 16
+  constexpr int tile_k = std::max(256, 1024 / tile_n);  // 256
+#if defined(SGL_CUDA_ARCH) && SGL_CUDA_ARCH >= 1200
+  constexpr int smem_stage_budget = 96 * 1024;
+#else
+  constexpr int smem_stage_budget = 192 * 1024;
+#endif
+  constexpr int max_stage_cnt = smem_stage_budget / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
+  constexpr int k_iter_cnt = gemm_k / tile_k;
+  constexpr int stage_cnt = k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt;
+  int cta_m_cnt = gemm_m / tile_m;
+  int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n;
+  constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024;  // 4096
+  constexpr int smem_bytes = ((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) / 1024 * 1024;
+
+  dim3 grid(cta_m_cnt, cta_n_cnt, 1);
+  dim3 block_size(256);
+
+  auto kernel = fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt>;
+  if (smem_bytes >= (48 * 1024)) {
+    host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));
+  }
+  host::LaunchKernel(grid, block_size, device, smem_bytes).enable_pdl(kUsePDL)(kernel, output, mat_a, mat_b, gemm_n);
+}
+
+template <int kHdIn, int kHdOut, bool kUsePDL>
+struct DSV3FusedAGemmKernel {
+  static void
+  run(const tvm::ffi::TensorView mat_a, const tvm::ffi::TensorView mat_b, const tvm::ffi::TensorView output) {
+    using namespace host;
+
+    auto M = SymbolicSize{"num_tokens"};
+    auto K = SymbolicSize{"hd_in"};
+    auto N = SymbolicSize{"hd_out"};
+    auto device = SymbolicDevice{};
+    K.set_value(kHdIn);
+    N.set_value(kHdOut);
+    device.set_options<kDLCUDA>();
+
+    // mat_a: [num_tokens, hd_in] row-major; output: [num_tokens, hd_out] row-major;
+    // mat_b: [hd_in, hd_out] column-major (weight.T), i.e. strides {1, hd_in}.
+    TensorMatcher({M, K}).with_dtype<bf16_t>().with_device(device).verify(mat_a);
+    TensorMatcher({K, N}).with_dtype<bf16_t>().with_device(device).with_strides({1, kHdIn}).verify(mat_b);
+    TensorMatcher({M, N}).with_dtype<bf16_t>().with_device(device).verify(output);
+
+    const int num_tokens = static_cast<int>(M.unwrap());
+    RuntimeCheck(
+        num_tokens >= 1 && num_tokens <= 16, "dsv3_fused_a_gemm: num_tokens must be in [1, 16], got ", num_tokens);
+
+    const DLDevice dev = device.unwrap();
+    auto* out_ptr = static_cast<bf16_t*>(output.data_ptr());
+    auto* a_ptr = static_cast<bf16_t const*>(mat_a.data_ptr());
+    auto* b_ptr = static_cast<bf16_t const*>(mat_b.data_ptr());
+
+    if (num_tokens <= 8) {
+      invokeFusedAGemm<bf16_t, kHdIn, kHdOut, 8, kUsePDL>(out_ptr, a_ptr, b_ptr, num_tokens, dev);
+    } else {
+      invokeFusedAGemm<bf16_t, kHdIn, kHdOut, 16, kUsePDL>(out_ptr, a_ptr, b_ptr, num_tokens, dev);
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/gemm/dsv3_router_gemm.cuh

@@ -0,0 +1,184 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
+ *
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <type_traits>
+
+namespace {
+
+using namespace device;
+
+static constexpr int kDefaultNumExperts = 256;
+static constexpr int kKimiK2NumExperts = 384;
+static constexpr int kDefaultHiddenDim = 7168;
+
+// kOutFloat: true = float32 output, false = bfloat16 output
+template <
+    typename T,
+    typename OutT,
+    int kBlockSize,
+    int VPT,
+    int kNumTokens,
+    int kNumExperts,
+    int kHiddenDim,
+    bool kUsePDL>
+__global__ __launch_bounds__(kBlockSize, 1) void router_gemm_kernel(OutT* out, T const* mat_a, T const* mat_b) {
+  constexpr int kWarpSize = 32;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;
+  constexpr int kElemsPerKIter = VPT * kBlockSize;
+  static_assert(kHiddenDim % kElemsPerKIter == 0, "hidden_dim must be divisible by one K iteration");
+  constexpr int kIters = kHiddenDim / kElemsPerKIter;
+  // Padding to avoid shared memory bank conflicts when kNumTokens > 8
+  constexpr int kSmReductionPad = (kNumTokens > 8) ? 1 : 0;
+  static_assert(kSmReductionPad == 0 || kSmReductionPad == 1, "kSmReductionPad only supports 0 or 1");
+
+  int const n_idx = blockIdx.x;
+  int const tid = threadIdx.x;
+  int const warp_id = tid / kWarpSize;
+  int const lane_id = tid % kWarpSize;
+
+  float acc[kNumTokens] = {};
+  __shared__ float sm_reduction[kNumTokens][kNumWarps + kSmReductionPad];
+
+  T const* b_col = mat_b + n_idx * kHiddenDim;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  int k_base = tid * VPT;
+#pragma unroll
+  for (int ki = 0; ki < kIters; ++ki, k_base += kElemsPerKIter) {
+    AlignedVector<bf16_t, VPT> b_vec;
+    b_vec.load(b_col + k_base);
+#pragma unroll
+    for (int m_idx = 0; m_idx < kNumTokens; ++m_idx) {
+      AlignedVector<bf16_t, VPT> a_vec;
+      a_vec.load(mat_a + m_idx * kHiddenDim + k_base);
+#pragma unroll
+      for (int k = 0; k < VPT; ++k) {
+        acc[m_idx] += cast<float>(a_vec[k]) * cast<float>(b_vec[k]);
+      }
+    }
+  }
+
+#pragma unroll
+  for (int m_idx = 0; m_idx < kNumTokens; ++m_idx) {
+    float sum = warp::reduce_sum(acc[m_idx]);
+    if (lane_id == 0) {
+      sm_reduction[m_idx][warp_id] = sum;
+    }
+  }
+
+  __syncthreads();
+
+  if (warp_id == 0 && lane_id < kNumTokens) {
+    float final_sum = 0.0f;
+#pragma unroll
+    for (int w = 0; w < kNumWarps; ++w) {
+      final_sum += sm_reduction[lane_id][w];
+    }
+    out[lane_id * kNumExperts + n_idx] = cast<OutT>(final_sum);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <typename T, typename OutT, int kNumTokens, int kNumExperts, int kHiddenDim, bool kUsePDL>
+void invokeRouterGemm(OutT* output, T const* mat_a, T const* mat_b, DLDevice device) {
+  constexpr int VPT = 16 / sizeof(T);
+  constexpr int kBlockSize = 128;
+  constexpr auto kernel = router_gemm_kernel<T, OutT, kBlockSize, VPT, kNumTokens, kNumExperts, kHiddenDim, kUsePDL>;
+  host::LaunchKernel(kNumExperts, kBlockSize, device).enable_pdl(kUsePDL)(kernel, output, mat_a, mat_b);
+}
+
+// Dispatch runtime num_tokens to compile-time template parameter [kBegin, kEnd]
+template <int kBegin, int kEnd, typename OutT, int kNumExperts, int kHiddenDim, bool kUsePDL>
+struct RouterGemmDispatcher {
+  static void run(int num_tokens, OutT* output, bf16_t const* mat_a, bf16_t const* mat_b, DLDevice device) {
+    if (num_tokens == kBegin) {
+      invokeRouterGemm<bf16_t, OutT, kBegin, kNumExperts, kHiddenDim, kUsePDL>(output, mat_a, mat_b, device);
+    } else {
+      RouterGemmDispatcher<kBegin + 1, kEnd, OutT, kNumExperts, kHiddenDim, kUsePDL>::run(
+          num_tokens, output, mat_a, mat_b, device);
+    }
+  }
+};
+
+// Base case: kBegin == kEnd
+template <int kEnd, typename OutT, int kNumExperts, int kHiddenDim, bool kUsePDL>
+struct RouterGemmDispatcher<kEnd, kEnd, OutT, kNumExperts, kHiddenDim, kUsePDL> {
+  static void run(int num_tokens, OutT* output, bf16_t const* mat_a, bf16_t const* mat_b, DLDevice device) {
+    if (num_tokens == kEnd) {
+      invokeRouterGemm<bf16_t, OutT, kEnd, kNumExperts, kHiddenDim, kUsePDL>(output, mat_a, mat_b, device);
+    } else {
+      host::panic({}, "dsv3_router_gemm: num_tokens must be between 1 and 16, got ", num_tokens);
+    }
+  }
+};
+
+// kNumExperts: compile-time 256 or 384
+// kHiddenDim: compile-time hidden dim, any multiple of one K iteration (1024)
+// kUsePDL: compile-time bool (true on SM90+)
+// kOutFloat: compile-time bool (true = float32 output, false = bfloat16 output)
+template <int kNumExperts, int kHiddenDim, bool kUsePDL, bool kOutFloat>
+struct DSV3RouterGemmKernel {
+  static_assert(
+      kNumExperts == kDefaultNumExperts || kNumExperts == kKimiK2NumExperts,
+      "required num_experts == 256 or num_experts == 384");
+
+  using OutT = std::conditional_t<kOutFloat, fp32_t, bf16_t>;
+
+  static void
+  run(const tvm::ffi::TensorView mat_a, const tvm::ffi::TensorView mat_b, const tvm::ffi::TensorView output) {
+    using namespace host;
+
+    auto M = SymbolicSize{"num_tokens"};
+    auto K = SymbolicSize{"hidden_dim"};
+    auto N = SymbolicSize{"num_experts"};
+    auto device = SymbolicDevice{};
+    K.set_value(kHiddenDim);
+    N.set_value(kNumExperts);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M, K}).with_dtype<bf16_t>().with_device(device).verify(mat_a);
+    TensorMatcher({N, K}).with_dtype<bf16_t>().with_device(device).verify(mat_b);
+    TensorMatcher({M, N}).with_dtype<OutT>().with_device(device).verify(output);
+
+    const int num_tokens = static_cast<int>(M.unwrap());
+
+    RouterGemmDispatcher<1, 16, OutT, kNumExperts, kHiddenDim, kUsePDL>::run(
+        num_tokens,
+        static_cast<OutT*>(output.data_ptr()),
+        static_cast<bf16_t const*>(mat_a.data_ptr()),
+        static_cast<bf16_t const*>(mat_b.data_ptr()),
+        device.unwrap());
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh

@@ -0,0 +1,251 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include "marlin.cuh"
+
+namespace device::marlin {
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+template <int const num_threads, int const num_bits>
+__global__ void awq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) {
+  return;
+}
+#else
+
+template <int const num_threads, int const num_bits>
+__global__ void awq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) {
+  constexpr int pack_factor = 32 / num_bits;
+
+  int k_tiles = size_k / tile_k_size;
+  int n_tiles = size_n / tile_n_size;
+  int block_k_tiles = div_ceil(k_tiles, (int)gridDim.x);
+
+  auto start_k_tile = blockIdx.x * block_k_tiles;
+  if (start_k_tile >= k_tiles) {
+    return;
+  }
+
+  int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles);
+
+  // Wait until the next thread tile has been loaded to shared memory.
+  auto wait_for_stage = [&]() {
+    // We only have `stages - 2` active fetches since we are double buffering
+    // and can only issue the next fetch when it is guaranteed that the previous
+    // shared memory load is fully complete (as it may otherwise be
+    // overwritten).
+    cp_async_wait<repack_stages - 2>();
+    __syncthreads();
+  };
+
+  extern __shared__ int4 sh[];
+
+  constexpr int tile_n_ints = tile_n_size / pack_factor;
+
+  constexpr int stage_n_threads = tile_n_ints / 4;
+  constexpr int stage_k_threads = tile_k_size;
+  constexpr int stage_size = stage_k_threads * stage_n_threads;
+
+  auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      cp_async_fence();
+      return;
+    }
+
+    int first_n = n_tile_id * tile_n_size;
+    int first_n_packed = first_n / pack_factor;
+
+    int4* sh_ptr = sh + stage_size * pipe;
+
+    if (threadIdx.x < stage_size) {
+      auto k_id = threadIdx.x / stage_n_threads;
+      auto n_id = threadIdx.x % stage_n_threads;
+
+      int first_k = k_tile_id * tile_k_size;
+
+      cp_async4(
+          &sh_ptr[k_id * stage_n_threads + n_id],
+          reinterpret_cast<int4 const*>(
+              &(b_q_weight_ptr[(first_k + k_id) * (size_n / pack_factor) + first_n_packed + (n_id * 4)])));
+    }
+
+    cp_async_fence();
+  };
+
+  auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      return;
+    }
+
+    auto warp_id = threadIdx.x / 32;
+    auto th_id = threadIdx.x % 32;
+
+    if (warp_id >= 4) {
+      return;
+    }
+
+    int tc_col = th_id / 4;
+    int tc_row = (th_id % 4) * 2;
+
+    constexpr int tc_offsets[4] = {0, 1, 8, 9};
+
+    int cur_n = warp_id * 16 + tc_col;
+    int cur_n_packed = cur_n / pack_factor;
+    int cur_n_pos = cur_n % pack_factor;
+
+    constexpr int sh_stride = tile_n_ints;
+    constexpr uint32_t mask = (1 << num_bits) - 1;
+
+    int4* sh_stage_ptr = sh + stage_size * pipe;
+    uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
+
+    // Undo interleaving
+    int cur_n_pos_unpacked;
+    if constexpr (num_bits == 4) {
+      constexpr int undo_pack[8] = {0, 4, 1, 5, 2, 6, 3, 7};
+      cur_n_pos_unpacked = undo_pack[cur_n_pos];
+    } else {
+      constexpr int undo_pack[4] = {0, 2, 1, 3};
+      cur_n_pos_unpacked = undo_pack[cur_n_pos];
+    }
+
+    uint32_t vals[8];
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+      int cur_elem = tc_row + tc_offsets[i];
+
+      int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem];
+      int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) + sh_stride * cur_elem];
+
+      vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask;
+      vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask;
+    }
+
+    constexpr int tile_size_val = tile_k_size * tile_n_size / pack_factor;
+    int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size_val;
+
+    // Result of:
+    // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
+    if constexpr (num_bits == 4) {
+      constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
+
+      uint32_t res = 0;
+#pragma unroll
+      for (int i = 0; i < 8; i++) {
+        res |= vals[pack_idx[i]] << (i * 4);
+      }
+
+      out_ptr[out_offset + th_id * 4 + warp_id] = res;
+
+    } else {
+      constexpr int pack_idx[4] = {0, 2, 1, 3};
+
+      uint32_t res1 = 0;
+      uint32_t res2 = 0;
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        res1 |= vals[pack_idx[i]] << (i * 8);
+        res2 |= vals[4 + pack_idx[i]] << (i * 8);
+      }
+
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2;
+    }
+  };
+
+  auto start_pipes = [&](int k_tile_id, int n_tile_id) {
+#pragma unroll
+    for (int pipe = 0; pipe < repack_stages - 1; pipe++) {
+      fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe);
+    }
+
+    wait_for_stage();
+  };
+#pragma unroll
+  for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
+    int n_tile_id = 0;
+
+    start_pipes(k_tile_id, n_tile_id);
+
+    while (n_tile_id < n_tiles) {
+#pragma unroll
+      for (int pipe = 0; pipe < repack_stages; pipe++) {
+        fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1);
+        repack_tile(pipe, k_tile_id, n_tile_id + pipe);
+        wait_for_stage();
+      }
+      n_tile_id += repack_stages;
+    }
+  }
+}
+#endif
+
+}  // namespace device::marlin
+
+// Host wrapper
+void awq_marlin_repack(
+    tvm::ffi::TensorView out, tvm::ffi::TensorView b_q_weight, int64_t size_k, int64_t size_n, int64_t num_bits) {
+  using namespace host;
+  using namespace device::marlin;
+
+  // Validate alignment
+  RuntimeCheck(size_k % tile_k_size == 0, "size_k = ", size_k, " is not divisible by tile_k_size = ", tile_k_size);
+  RuntimeCheck(size_n % tile_n_size == 0, "size_n = ", size_n, " is not divisible by tile_n_size = ", tile_n_size);
+  RuntimeCheck(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits);
+
+  int const pack_factor = 32 / num_bits;
+
+  // Validate tensors
+  SymbolicDevice cuda_device;
+  cuda_device.set_options<kDLCUDA>();
+
+  TensorMatcher({size_k, size_n / pack_factor}).with_dtype<int32_t>().with_device(cuda_device).verify(b_q_weight);
+
+  TensorMatcher({size_k / tile_size, size_n * tile_size / pack_factor})
+      .with_dtype<int32_t>()
+      .with_device(cuda_device)
+      .verify(out);
+
+  // Get device and stream
+  auto device = cuda_device.unwrap();
+  auto stream = LaunchKernel::resolve_device(device);
+
+  // Get pointers
+  auto* b_q_weight_ptr = reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
+  auto* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
+
+  // Get device attributes
+  int blocks = 0;
+  cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device.device_id);
+
+  int max_shared_mem = 0;
+  cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device.device_id);
+  RuntimeCheck(max_shared_mem > 0, "max_shared_mem must be > 0");
+
+  // Dispatch based on num_bits
+  if (num_bits == 4) {
+    cudaFuncSetAttribute(
+        awq_marlin_repack_kernel<repack_threads, 4>, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);
+    LaunchKernel(blocks, repack_threads, stream, max_shared_mem)(
+        awq_marlin_repack_kernel<repack_threads, 4>,
+        b_q_weight_ptr,
+        out_ptr,
+        static_cast<int>(size_k),
+        static_cast<int>(size_n));
+  } else if (num_bits == 8) {
+    cudaFuncSetAttribute(
+        awq_marlin_repack_kernel<repack_threads, 8>, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);
+    LaunchKernel(blocks, repack_threads, stream, max_shared_mem)(
+        awq_marlin_repack_kernel<repack_threads, 8>,
+        b_q_weight_ptr,
+        out_ptr,
+        static_cast<int>(size_k),
+        static_cast<int>(size_n));
+  } else {
+    RuntimeCheck(false, "Unsupported repack config: num_bits = ", num_bits);
+  }
+}

python/sglang/jit_kernel/csrc/gemm/marlin/dequant.h

@@ -0,0 +1,504 @@
+/*
+Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16)
+
+The process of fast dequantization can be summarized as a combination
+of bitwise operations and floating-point computations:
+
+weight =>(bit_op / bitwise operations)=>
+f16_value =>(flop / floating-point computation)=>
+dequantized_weight
+
+Since the dequantized weights typically require subtracting the zero point and
+applying a scale factor, the floating-point computation step can be fused with
+the zero-point subtraction and scaling operations.
+
+The following are the parts that need to be modified for the fused operation
+of zero-point subtraction and scaling.
+
+## INT4 => FP16/BF16 or INT8 => FP16
+
+The floating-point computation is `__hsub2`
+
+If has zero points:
+
+    flop(bit_op(weight)) - flop(bit_op(zp))
+  = sub(bit_op(weight), bias) - sub(bit_op(zp), bias)
+  = bit_op(weight) - bit_op(zp)
+
+so we don't need additional modification.
+
+If has float zero points:
+
+    flop(bit_op(weight)) - fzp
+  = sub(bit_op(weight), bias) - fzp
+  = bit_op(weight) - (fzp + bias)
+
+where the `fzp + bias` can be computed at weight loading. But this
+may have accuracy issue, so we should not use this in most cases.
+
+If has not zero points:
+
+    scale(flop(bit_op(weight)))
+  = scale(sub(bit_op(weight), bias))
+  = scale(bit_op(weight)) - scale(bias)
+  = fma(bit_op(weight), scale_factor, scale(bias))
+
+where the `scale(bias)` can be cached. But this may have accuracy issue,
+so we should not use this in most cases.
+
+
+## INT8 => BF16
+
+INT8 => BF16 is a special case, it use byte_perm instead of flop.
+We cannot fused byte_perm with scaling.
+
+
+## FP4/FP8 => FP16/BF16
+
+    scale(flop(bit_op(weight)))
+  = scale(mul(bit_op(weight), multiplier))
+  = mul(bit_op(weight), scale_factor * multiplier)
+
+where `scale_factor * multiplier` can be computed at weight loading.
+
+*/
+
+#include "marlin_dtypes.cuh"
+
+namespace device::marlin {
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+// Lookup-table based 3-input logical operation; explicitly used for
+// dequantization as the compiler does not seem to automatically recognize it in
+// all cases.
+template <int lut>
+__device__ inline int lop3(int a, int b, int c) {
+  int res;
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut));
+  return res;
+}
+
+// Constructs destination register by taking bytes from 2 sources (based on
+// mask)
+template <int start_byte, int mask>
+__device__ inline uint32_t prmt(uint32_t a) {
+  uint32_t res;
+  asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(res) : "r"(a), "n"(start_byte), "n"(mask));
+  return res;
+}
+
+template <typename scalar_t2, host::ScalarTypeId w_type_id, bool skip_flop = false>
+__device__ inline void dequant(int q, scalar_t2* frag_b);
+
+//
+// Efficiently dequantize 4bit values packed in an int32 value into a full
+// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,
+// with some small changes:
+// - FP16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287
+// - BF16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
+//
+template <>
+__device__ inline void dequant<half2, host::kU4B8.id(), true>(int q, half2* frag_b) {
+  const int MASK = 0x000f000f;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  q >>= 4;
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+
+  frag_b[0] = *reinterpret_cast<half2*>(&lo);
+  frag_b[1] = *reinterpret_cast<half2*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kU4B8.id(), false>(int q, half2* frag_b) {
+  const int LO = 0x000f000f;
+  const int HI = 0x00f000f0;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
+  // clang-format on
+  // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
+  // directly into `SUB` and `ADD`.
+  const int SUB = 0x64086408;
+  const int MUL = 0x2c002c00;
+  const int ADD = 0xd480d480;
+  frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo), *reinterpret_cast<const half2*>(&SUB));
+  frag_b[1] = __hfma2(
+      *reinterpret_cast<half2*>(&hi), *reinterpret_cast<const half2*>(&MUL), *reinterpret_cast<const half2*>(&ADD));
+}
+
+template <>
+__device__ inline void dequant<half2, host::kU4.id(), true>(int q, half2* frag_b) {
+  dequant<half2, host::kU4B8.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kU4.id(), false>(int q, half2* frag_b) {
+  const int LO = 0x000f000f;
+  const int HI = 0x00f000f0;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
+  // clang-format on
+  // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
+  // directly into `SUB` and `ADD`.
+  const int SUB = 0x64006400;
+  const int MUL = 0x2c002c00;
+  const int ADD = 0xd400d400;
+  frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo), *reinterpret_cast<const half2*>(&SUB));
+  frag_b[1] = __hfma2(
+      *reinterpret_cast<half2*>(&hi), *reinterpret_cast<const half2*>(&MUL), *reinterpret_cast<const half2*>(&ADD));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kU4B8.id(), true>(int q, nv_bfloat162* frag_b) {
+  static constexpr uint32_t MASK = 0x000f000f;
+  static constexpr uint32_t EX = 0x43004300;
+
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  q >>= 4;
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  // clang-format on
+
+  frag_b[0] = *reinterpret_cast<nv_bfloat162*>(&lo);
+  frag_b[1] = *reinterpret_cast<nv_bfloat162*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kU4B8.id(), false>(int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, host::kU4B8.id(), true>(q, frag_b);
+
+  static constexpr uint32_t SUB = 0x43084308;
+
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kU4.id(), true>(int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, host::kU4B8.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kU4.id(), false>(int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, host::kU4.id(), true>(q, frag_b);
+
+  static constexpr uint32_t SUB = 0x43004300;
+
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+}
+
+//
+// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or
+// bf16 Reference:
+// - FP16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85
+// - BF16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
+//
+template <>
+__device__ inline void dequant<half2, host::kU8B128.id(), true>(int q, half2* frag_b) {
+  static constexpr uint32_t mask_for_elt_01 = 0x5250;
+  static constexpr uint32_t mask_for_elt_23 = 0x5351;
+  static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
+
+  uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
+  uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
+
+  frag_b[0] = *reinterpret_cast<half2*>(&lo);
+  frag_b[1] = *reinterpret_cast<half2*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kU8B128.id(), false>(int q, half2* frag_b) {
+  dequant<half2, host::kU8B128.id(), true>(q, frag_b);
+
+  static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+}
+
+template <>
+__device__ inline void dequant<half2, host::kU8.id(), true>(int q, half2* frag_b) {
+  dequant<half2, host::kU8B128.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kU8.id(), false>(int q, half2* frag_b) {
+  dequant<half2, host::kU8.id(), true>(q, frag_b);
+
+  static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kU8B128.id(), false>(int q, nv_bfloat162* frag_b) {
+  float fp32_intermediates[4];
+  uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
+
+  static constexpr uint32_t fp32_base = 0x4B000000;
+  fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
+  fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
+  fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
+  fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
+
+  fp32_intermediates[0] -= 8388736.f;
+  fp32_intermediates[1] -= 8388736.f;
+  fp32_intermediates[2] -= 8388736.f;
+  fp32_intermediates[3] -= 8388736.f;
+
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
+  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kU8.id(), false>(int q, nv_bfloat162* frag_b) {
+  float fp32_intermediates[4];
+  uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
+
+  static constexpr uint32_t fp32_base = 0x4B000000;
+  fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
+  fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
+  fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
+  fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
+
+  fp32_intermediates[0] -= 8388608.f;
+  fp32_intermediates[1] -= 8388608.f;
+  fp32_intermediates[2] -= 8388608.f;
+  fp32_intermediates[3] -= 8388608.f;
+
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
+  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kFE4M3fn.id(), true>(int q, half2* frag_b) {
+  // Constants for FP8 (E4M3) and FP16 formats
+  constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
+  constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kFE4M3fn.id(), false>(int q, half2* frag_b) {
+  dequant<half2, host::kFE4M3fn.id(), true>(q, frag_b);
+
+  // Constants for FP8 (E4M3) and FP16 formats
+  constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
+  const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kFE4M3fn.id(), true>(int q, nv_bfloat162* frag_b) {
+  // Constants for FP8 (E4M3) and BF16 formats
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
+
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to BF16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kFE4M3fn.id(), false>(int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, host::kFE4M3fn.id(), true>(q, frag_b);
+
+  // Constants for FP8 (E4M3) and BF16 formats
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
+  // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
+  // position
+  constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
+  const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
+
+  // Convert to bfloat162 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kFE2M1f.id(), true>(int q, half2* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
+  constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70007000;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<half2, host::kFE2M1f.id(), false>(int q, half2* frag_b) {
+  dequant<half2, host::kFE2M1f.id(), true>(q, frag_b);
+
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
+  const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kFE2M1f.id(), true>(int q, nv_bfloat162* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70007000;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, host::kFE2M1f.id(), false>(int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, host::kFE2M1f.id(), true>(q, frag_b);
+
+  // Constants for FP4 (E2M1) and BF16 formats
+  constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
+  // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
+  // position
+  constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
+  const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <typename scalar_t2>
+__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b);
+
+template <>
+__device__ inline void dequant_fp8_scales<half2>(int q, half2* frag_b) {
+  int Out1 = (q & 0xFF00FF00) >> 1;
+  ;
+  q <<= 8;
+  int Out2 = (q & 0xFF00FF00) >> 1;
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+};
+
+template <>
+__device__ inline void dequant_fp8_scales<nv_bfloat162>(int q, nv_bfloat162* frag_b) {
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to BF16 format
+  int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+};
+
+// New version with s_type_id parameter for marlin_moe_wna16_v2
+template <typename scalar_t2, host::ScalarTypeId s_type_id>
+__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b);
+
+template <>
+__device__ inline void dequant_fp8_scales<half2, host::kFE4M3fn.id()>(int q, half2* frag_b) {
+  int Out1 = (q & 0xFF00FF00) >> 1;
+  ;
+  q <<= 8;
+  int Out2 = (q & 0xFF00FF00) >> 1;
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+};
+
+template <>
+__device__ inline void dequant_fp8_scales<nv_bfloat162, host::kFE4M3fn.id()>(int q, nv_bfloat162* frag_b) {
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to BF16 format
+  int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant_fp8_scales<nv_bfloat162, host::kFE8M0fnu.id()>(int q, nv_bfloat162* frag_b) {
+  // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16,
+  // but we assume that such a extreme value would not occur in real models.
+  int Out1 = (q & 0xFF00FF00) >> 1;
+  q <<= 7;
+  int Out2 = q & 0x7F807F80;
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+#endif
+
+}  // namespace device::marlin

python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh

@@ -0,0 +1,362 @@
+/*
+ * Modified by Neural Magic
+ * Copyright (C) Marlin.2024 Elias Frantar
+ *
+ * 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.
+ */
+
+/*
+ * Adapted from https://github.com/IST-DASLab/marlin
+ */
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include "marlin.cuh"
+
+namespace device::marlin {
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+template <int const num_threads, int const num_bits, bool const has_perm>
+__global__ void gptq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr,
+    uint32_t const* __restrict__ perm_ptr,
+    uint32_t* __restrict__ out_ptr,
+    int size_k,
+    int size_n) {
+  return;
+}
+#else
+template <int const num_threads, int const num_bits, bool const has_perm>
+__global__ void gptq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr,
+    uint32_t const* __restrict__ perm_ptr,
+    uint32_t* __restrict__ out_ptr,
+    int size_k,
+    int size_n) {
+  constexpr int pack_factor = 32 / num_bits;
+
+  int k_tiles = size_k / tile_k_size;
+  int n_tiles = size_n / tile_n_size;
+  int block_k_tiles = div_ceil(k_tiles, gridDim.x);
+
+  auto start_k_tile = blockIdx.x * block_k_tiles;
+  if (start_k_tile >= k_tiles) {
+    return;
+  }
+
+  int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles);
+
+  // Wait until the next thread tile has been loaded to shared memory.
+  auto wait_for_stage = [&]() {
+    // We only have `stages - 2` active fetches since we are double buffering
+    // and can only issue the next fetch when it is guaranteed that the previous
+    // shared memory load is fully complete (as it may otherwise be
+    // overwritten).
+    cp_async_wait<repack_stages - 2>();
+    __syncthreads();
+  };
+
+  extern __shared__ int4 sh[];
+
+  constexpr int perm_size = tile_k_size / 4;
+
+  int4* sh_perm_ptr = sh;
+  int4* sh_pipe_ptr = sh_perm_ptr;
+  if constexpr (has_perm) {
+    sh_pipe_ptr += perm_size;
+  }
+
+  constexpr int tile_ints = tile_k_size / pack_factor;
+
+  constexpr int stage_n_threads = tile_n_size / 4;
+  constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints;
+  constexpr int stage_size = stage_k_threads * stage_n_threads;
+
+  auto load_perm_to_shared = [&](int k_tile_id) {
+    int first_k_int4 = (k_tile_id * tile_k_size) / 4;
+
+    int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
+
+    if (threadIdx.x < perm_size) {
+      sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x];
+    }
+    __syncthreads();
+  };
+
+  auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      cp_async_fence();
+      return;
+    }
+
+    int first_n = n_tile_id * tile_n_size;
+
+    int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
+
+    if constexpr (has_perm) {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        uint32_t const* sh_perm_int_ptr = reinterpret_cast<uint32_t const*>(sh_perm_ptr);
+
+        int src_k = sh_perm_int_ptr[k_id];
+        int src_k_packed = src_k / pack_factor;
+
+        cp_async4(
+            &sh_ptr[k_id * stage_n_threads + n_id],
+            reinterpret_cast<int4 const*>(&(b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)])));
+      }
+
+    } else {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        int first_k = k_tile_id * tile_k_size;
+        int first_k_packed = first_k / pack_factor;
+
+        cp_async4(
+            &sh_ptr[k_id * stage_n_threads + n_id],
+            reinterpret_cast<int4 const*>(&(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)])));
+      }
+    }
+
+    cp_async_fence();
+  };
+
+  auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      return;
+    }
+
+    auto warp_id = threadIdx.x / 32;
+    auto th_id = threadIdx.x % 32;
+
+    if (warp_id >= 4) {
+      return;
+    }
+
+    int tc_col = th_id / 4;
+    int tc_row = (th_id % 4) * 2;
+
+    constexpr int tc_offsets[4] = {0, 1, 8, 9};
+
+    int cur_n = warp_id * 16 + tc_col;
+
+    constexpr int sh_stride = 64;
+    constexpr uint32_t mask = (1 << num_bits) - 1;
+
+    int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
+    uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
+
+    uint32_t* sh_perm_int_ptr = reinterpret_cast<uint32_t*>(sh_perm_ptr);
+
+    uint32_t vals[8];
+
+    if constexpr (has_perm) {
+      for (int i = 0; i < 4; i++) {
+        int k_idx = tc_row + tc_offsets[i];
+
+        uint32_t src_k = sh_perm_int_ptr[k_idx];
+        uint32_t src_k_pos = src_k % pack_factor;
+
+        uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n];
+        uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask;
+
+        uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8];
+        uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask;
+
+        vals[i] = b1_cur_val;
+        vals[4 + i] = b2_cur_val;
+      }
+
+    } else {
+      uint32_t b1_vals[tile_ints];
+      uint32_t b2_vals[tile_ints];
+
+#pragma unroll
+      for (int i = 0; i < tile_ints; i++) {
+        b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
+        b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
+      }
+
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        int cur_elem = tc_row + tc_offsets[i];
+        int cur_int = cur_elem / pack_factor;
+        int cur_pos = cur_elem % pack_factor;
+
+        vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+        vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+      }
+    }
+
+    constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
+    int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
+
+    // Result of:
+    // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
+    if constexpr (num_bits == 4) {
+      constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
+
+      uint32_t res = 0;
+#pragma unroll
+      for (int i = 0; i < 8; i++) {
+        res |= vals[pack_idx[i]] << (i * 4);
+      }
+
+      out_ptr[out_offset + th_id * 4 + warp_id] = res;
+
+    } else {
+      constexpr int pack_idx[4] = {0, 2, 1, 3};
+
+      uint32_t res1 = 0;
+      uint32_t res2 = 0;
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        res1 |= vals[pack_idx[i]] << (i * 8);
+        res2 |= vals[4 + pack_idx[i]] << (i * 8);
+      }
+
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2;
+    }
+  };
+
+  auto start_pipes = [&](int k_tile_id, int n_tile_id) {
+#pragma unroll
+    for (int pipe = 0; pipe < repack_stages - 1; pipe++) {
+      fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe);
+    }
+
+    wait_for_stage();
+  };
+#pragma unroll
+  for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
+    int n_tile_id = 0;
+
+    if constexpr (has_perm) {
+      load_perm_to_shared(k_tile_id);
+    }
+
+    start_pipes(k_tile_id, n_tile_id);
+
+    while (n_tile_id < n_tiles) {
+#pragma unroll
+      for (int pipe = 0; pipe < repack_stages; pipe++) {
+        fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1);
+        repack_tile(pipe, k_tile_id, n_tile_id + pipe);
+        wait_for_stage();
+      }
+      n_tile_id += repack_stages;
+    }
+  }
+}
+#endif
+
+}  // namespace device::marlin
+
+#define CALL_IF_REPACK(NUM_BITS, HAS_PERM)                                                                        \
+  else if (num_bits == NUM_BITS && has_perm == HAS_PERM) {                                                        \
+    host::RuntimeDeviceCheck(cudaFuncSetAttribute(                                                                \
+        device::marlin::gptq_marlin_repack_kernel<device::marlin::repack_threads, NUM_BITS, HAS_PERM>,            \
+        cudaFuncAttributeMaxDynamicSharedMemorySize,                                                              \
+        max_shared_mem));                                                                                         \
+    host::LaunchKernel(blocks, device::marlin::repack_threads, stream, static_cast<std::size_t>(max_shared_mem))( \
+        device::marlin::gptq_marlin_repack_kernel<device::marlin::repack_threads, NUM_BITS, HAS_PERM>,            \
+        b_q_weight_ptr,                                                                                           \
+        perm_ptr,                                                                                                 \
+        out_ptr,                                                                                                  \
+        size_k,                                                                                                   \
+        size_n);                                                                                                  \
+  }
+
+void gptq_marlin_repack(
+    tvm::ffi::TensorView b_q_weight,
+    tvm::ffi::TensorView perm,
+    tvm::ffi::TensorView out,
+    int64_t size_k,
+    int64_t size_n,
+    int64_t num_bits) {
+  using namespace host;
+
+  // Validate num_bits
+  RuntimeCheck(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits);
+  int const pack_factor = 32 / static_cast<int>(num_bits);
+
+  // Validate size alignment
+  RuntimeCheck(
+      size_k % device::marlin::tile_k_size == 0,
+      "size_k = ",
+      size_k,
+      " is not divisible by tile_k_size = ",
+      device::marlin::tile_k_size);
+  RuntimeCheck(
+      size_n % device::marlin::tile_n_size == 0,
+      "size_n = ",
+      size_n,
+      " is not divisible by tile_n_size = ",
+      device::marlin::tile_n_size);
+
+  // Validate b_q_weight
+  auto bqw_dim0 = SymbolicSize{"bqw_dim0"};
+  auto bqw_dim1 = SymbolicSize{"bqw_dim1"};
+  bqw_dim0.set_value(size_k / pack_factor);
+  bqw_dim1.set_value(size_n);
+  auto device_ = SymbolicDevice{};
+  device_.set_options<kDLCUDA>();
+  TensorMatcher({bqw_dim0, bqw_dim1}).with_dtype<int32_t>().with_device(device_).verify(b_q_weight);
+
+  // Validate out
+  auto out_dim0 = SymbolicSize{"out_dim0"};
+  auto out_dim1 = SymbolicSize{"out_dim1"};
+  out_dim0.set_value(size_k / device::marlin::tile_size);
+  out_dim1.set_value(size_n * device::marlin::tile_size / pack_factor);
+  TensorMatcher({out_dim0, out_dim1}).with_dtype<int32_t>().with_device(device_).verify(out);
+
+  // Detect if there is act_order
+  bool has_perm = perm.size(0) != 0;
+
+  // Get ptrs
+  uint32_t const* b_q_weight_ptr = reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
+  uint32_t const* perm_ptr = reinterpret_cast<uint32_t const*>(perm.data_ptr());
+  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
+
+  // Get dev info
+  DLDevice dl_device = device_.unwrap();
+  int dev = dl_device.device_id;
+  cudaStream_t stream = LaunchKernel::resolve_device(dl_device);
+  int blocks;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev));
+
+  int max_shared_mem = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev));
+  RuntimeCheck(max_shared_mem > 0, "max_shared_mem must be > 0");
+
+  if (false) {
+  }
+  CALL_IF_REPACK(4, false)
+  CALL_IF_REPACK(4, true)
+  CALL_IF_REPACK(8, false)
+  CALL_IF_REPACK(8, true)
+  else {
+    Panic("Unsupported repack config: num_bits = ", num_bits, ", has_perm = ", has_perm);
+  }
+}
+
+#undef CALL_IF_REPACK

python/sglang/jit_kernel/csrc/gemm/marlin/kernel.h

@@ -0,0 +1,33 @@
+
+#include <sgl_kernel/scalar_type.hpp>
+
+#include "marlin.cuh"
+#include "marlin_dtypes.cuh"
+
+#define MARLIN_KERNEL_PARAMS                                                                                         \
+  const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp,            \
+      const int4 *__restrict__ scales_ptr, const uint16_t *__restrict__ scale2_ptr, const int4 *__restrict__ zp_ptr, \
+      const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks,        \
+      bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem
+
+namespace device::marlin {
+template <
+    typename scalar_t,                   // compute dtype, half or nv_float16
+    const host::ScalarTypeId w_type_id,  // weight ScalarType id
+    const int threads,                   // number of threads in a threadblock
+    const int thread_m_blocks,           // number of 16x16 blocks in the m
+                                         // dimension (batchsize) of the
+                                         // threadblock
+    const int thread_n_blocks,           // same for n dimension (output)
+    const int thread_k_blocks,           // same for k dimension (reduction)
+    const bool m_block_size_8,           // whether m_block_size == 8
+                                         // only works when thread_m_blocks == 1
+    const int stages,                    // number of stages for the async global->shared
+                                         // fetch pipeline
+    const int group_blocks,              // number of consecutive 16x16 blocks
+                                         // with a separate quantization scale
+    const bool is_zp_float               // is zero point of float16 type?
+    >
+__global__ void Marlin(MARLIN_KERNEL_PARAMS);
+
+}  // namespace device::marlin

python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh

@@ -0,0 +1,83 @@
+#pragma once
+
+#include <sgl_kernel/utils.cuh>
+
+#include <iostream>
+
+namespace device::marlin {
+// Marlin params
+
+// 8 warps are a good choice since every SM has 4 schedulers and having more
+// than 1 warp per schedule allows some more latency hiding. At the same time,
+// we want relatively few warps to have many registers per warp and small tiles.
+static constexpr int default_threads = 256;
+
+static constexpr int pipe_stages = 4;  // 4 pipeline stages fit into shared memory
+
+static constexpr int min_thread_n = 64;
+static constexpr int min_thread_k = 64;
+static constexpr int max_thread_n = 256;
+
+static constexpr int tile_size = 16;
+static constexpr int max_par = 16;
+
+// Repack params
+static constexpr int repack_stages = 8;
+
+static constexpr int repack_threads = 256;
+
+static constexpr int tile_k_size = tile_size;
+static constexpr int tile_n_size = tile_k_size * 4;
+
+// Helpers
+template <typename T, int n>
+struct Vec {
+  T elems[n];
+  __device__ T& operator[](int i) {
+    return elems[i];
+  }
+};
+
+using I4 = Vec<int, 4>;
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+// No support for async
+#else
+
+__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, bool pred = true) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.cg.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem),
+      "l"(glob_ptr),
+      "n"(BYTES));
+}
+
+__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   cp.async.cg.shared.global [%0], [%1], %2;\n"
+      "}\n" ::"r"(smem),
+      "l"(glob_ptr),
+      "n"(BYTES));
+}
+
+__device__ inline void cp_async_fence() {
+  asm volatile("cp.async.commit_group;\n" ::);
+}
+
+template <int n>
+__device__ inline void cp_async_wait() {
+  asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
+}
+
+#endif
+
+}  // namespace device::marlin

python/sglang/jit_kernel/csrc/gemm/marlin/marlin_dtypes.cuh

@@ -0,0 +1,77 @@
+#ifndef _data_types_cuh
+#define _data_types_cuh
+#include <sgl_kernel/utils.cuh>
+
+#include "marlin.cuh"
+
+namespace device::marlin {
+
+template <typename scalar_t>
+class ScalarType {};
+
+template <>
+class ScalarType<fp16_t> {
+ public:
+  using scalar_t = fp16_t;
+  using scalar_t2 = fp16x2_t;
+
+  // Matrix fragments for tensor core instructions; their precise layout is
+  // documented here:
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type
+  using FragA = Vec<fp16x2_t, 4>;
+  using FragB = Vec<fp16x2_t, 2>;
+  using FragC = Vec<float, 4>;
+  using FragS = Vec<fp16x2_t, 1>;
+  using FragZP = Vec<fp16x2_t, 4>;
+
+  static __device__ float inline num2float(const fp16_t x) {
+    return __half2float(x);
+  }
+
+  static __device__ fp16x2_t inline num2num2(const fp16_t x) {
+    return __half2half2(x);
+  }
+
+  static __device__ fp16x2_t inline nums2num2(const fp16_t x1, const fp16_t x2) {
+    return __halves2half2(x1, x2);
+  }
+
+  static __host__ __device__ fp16_t inline float2num(const float x) {
+    return __float2half(x);
+  }
+};
+
+template <>
+class ScalarType<bf16_t> {
+ public:
+  using scalar_t = bf16_t;
+  using scalar_t2 = bf16x2_t;
+
+  using FragA = Vec<bf16x2_t, 4>;
+  using FragB = Vec<bf16x2_t, 2>;
+  using FragC = Vec<float, 4>;
+  using FragS = Vec<bf16x2_t, 1>;
+  using FragZP = Vec<bf16x2_t, 4>;
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+  static __device__ float inline num2float(const bf16_t x) {
+    return __bfloat162float(x);
+  }
+
+  static __device__ bf16x2_t inline num2num2(const bf16_t x) {
+    return __bfloat162bfloat162(x);
+  }
+
+  static __device__ bf16x2_t inline nums2num2(const bf16_t x1, const bf16_t x2) {
+    return __halves2bfloat162(x1, x2);
+  }
+
+  static __host__ __device__ bf16_t inline float2num(const float x) {
+    return __float2bfloat16(x);
+  }
+#endif
+};
+
+}  // namespace device::marlin
+
+#endif

python/sglang/jit_kernel/csrc/gemm/marlin_moe/kernel.h

@@ -0,0 +1,37 @@
+
+#include <sgl_kernel/scalar_type.hpp>
+
+#include "../marlin/marlin.cuh"
+#include "../marlin/marlin_dtypes.cuh"
+
+#define MARLIN_KERNEL_PARAMS                                                                                         \
+  const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp,            \
+      const int4 *__restrict__ b_bias_ptr, const int4 *__restrict__ scales_ptr,                                      \
+      const uint16_t *__restrict__ scale2_ptr, const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx,       \
+      const int32_t *__restrict__ sorted_token_ids_ptr, const int32_t *__restrict__ expert_ids_ptr,                  \
+      const int32_t *__restrict__ num_tokens_past_padded_ptr, const float *__restrict__ topk_weights_ptr, int top_k, \
+      bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, int prob_n, int prob_k, int *locks,             \
+      bool has_bias, bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem
+
+namespace device::marlin_moe {
+template <
+    typename scalar_t,                   // compute dtype, half or nv_float16
+    const host::ScalarTypeId w_type_id,  // weight ScalarType id
+    const host::ScalarTypeId s_type_id,  // weight scale ScalarType id
+    const int threads,                   // number of threads in a threadblock
+    const int thread_m_blocks,           // number of 16x16 blocks in the m
+                                         // dimension (batchsize) of the
+                                         // threadblock
+    const int thread_n_blocks,           // same for n dimension (output)
+    const int thread_k_blocks,           // same for k dimension (reduction)
+    const bool m_block_size_8,           // whether m_block_size == 8
+                                         // only works when thread_m_blocks == 1
+    const int stages,                    // number of stages for the async global->shared
+                                         // fetch pipeline
+    const int group_blocks,              // number of consecutive 16x16 blocks
+                                         // with a separate quantization scale
+    const bool is_zp_float               // is zero point of float16 type?
+    >
+__global__ void Marlin(MARLIN_KERNEL_PARAMS);
+
+}  // namespace device::marlin_moe

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh

@@ -0,0 +1,806 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include "nvfp4_quant.cuh"
+#include <cuda_runtime.h>
+#include <cuda_runtime_api.h>
+
+using namespace host;
+
+// Quantizes the provided PackedVec into the uint32_t output
+template <class Type, bool UE8M0_SF = false>
+SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  // Get absolute maximum values among the local 8 values.
+  auto localMax = __habs2(vec.elts[0]);
+
+// Local maximum value.
+#pragma unroll
+  for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
+    localMax = __hmax2(localMax, __habs2(vec.elts[i]));
+  }
+
+  // Get the absolute maximum among all 16 values (two threads).
+  localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
+  // Get the final absolute maximum values.
+  float vecMax = float(__hmax(localMax.x, localMax.y));
+
+  // Get the SF (max value of the vector / max value of e2m1).
+  // maximum value of e2m1 = 6.0.
+  // TODO: use half as compute data type.
+  float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
+  // 8 bits representation of the SF.
+  uint8_t fp8SFVal;
+  // Write the SF to global memory (STG.8).
+  if constexpr (UE8M0_SF) {
+    // Extract the 8 exponent bits from float32.
+    // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
+    uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
+    fp8SFVal = tmp & 0xff;
+    // Convert back to fp32.
+    reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
+  } else {
+    // Here SFValue is always positive, so E4M3 is the same as UE4M3.
+    __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
+    reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
+    // Convert back to fp32.
+    SFValue = float(tmp);
+  }
+  // Get the output scale.
+  // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
+  //                       reciprocal(SFScaleVal))
+  float outputScale =
+      SFValue != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f;
+
+  if (SFout) {
+    // Write the SF to global memory (STG.8).
+    *SFout = fp8SFVal;
+  }
+
+  // Convert the input to float.
+  float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
+
+#pragma unroll
+  for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
+    fp2Vals[i] = device::cast<float2>(vec.elts[i]);
+    fp2Vals[i].x *= outputScale;
+    fp2Vals[i].y *= outputScale;
+  }
+
+  // Convert to e2m1 values.
+  uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
+
+  // Write the e2m1 values to global memory.
+  return e2m1Vec;
+#else
+  return 0;
+#endif
+}
+
+SGL_DEVICE float silu(const float& val) {
+  return val / (1.0f + __expf(-val));
+}
+
+template <class Type>
+SGL_DEVICE void silu_and_mul(PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
+  float2 x[CVT_FP4_ELTS_PER_THREAD / 2];
+  float2 y[CVT_FP4_ELTS_PER_THREAD / 2];
+
+#pragma unroll
+  for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
+    x[i] = device::cast<float2>(x_vec.elts[i]);
+    y[i] = device::cast<float2>(y_vec.elts[i]);
+    x[i].x = silu(x[i].x) * y[i].x;
+    x[i].y = silu(x[i].y) * y[i].y;
+    x_vec.elts[i] = device::cast<packed_t<Type>>(x[i]);
+  }
+}
+
+// Use UE4M3 by default.
+template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
+__global__ void
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+__launch_bounds__(512, 4) cvt_fp16_to_fp4(
+#else
+cvt_fp16_to_fp4(
+#endif
+    int32_t numRows,
+    int32_t numCols,
+    Type const* in,
+    float const* SFScale,
+    uint32_t* out,
+    uint32_t* SFout,
+    uint32_t* input_offset_by_experts,
+    uint32_t* output_scale_offset_by_experts,
+    int32_t* mask,
+    int n_experts,
+    bool low_latency,
+    bool use_silu_and_mul) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  using PackedVec = PackedVec<Type>;
+  static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
+  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched.");
+
+  // Input tensor row/col loops.
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
+  bool use_mask = mask != nullptr;
+  // When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated).
+  int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow;
+
+  // Each global thread processes one element
+  for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) {
+    // Calculate which row and column this global thread should process
+    int rowIdx = globalIdx / colsPerRow;
+    int colIdx = globalIdx % colsPerRow;
+
+    // Find index within the experts using different strategies based on expert
+    // count
+    int rowIdx_in_expert = 0;
+    int expert_idx = 0;
+
+    if constexpr (SMALL_NUM_EXPERTS) {
+      for (int i = 0; i < n_experts; i++) {
+        uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
+        uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
+        if (rowIdx >= current_offset && rowIdx < next_offset) {
+          rowIdx_in_expert = rowIdx - current_offset;
+          expert_idx = i;
+          break;
+        }
+      }
+    } else {
+      // Load input offsets into registers first, then do the computation.
+      // Local array size set to 17 because of register limit.
+      uint32_t local_offsets[17];
+      for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
+        *reinterpret_cast<int4*>(local_offsets) =
+            __ldca(reinterpret_cast<const int4*>(&input_offset_by_experts[chunk_start]));
+        *reinterpret_cast<int4*>(local_offsets + 4) =
+            __ldca(reinterpret_cast<const int4*>(&input_offset_by_experts[chunk_start + 4]));
+        *reinterpret_cast<int4*>(local_offsets + 8) =
+            __ldca(reinterpret_cast<const int4*>(&input_offset_by_experts[chunk_start + 8]));
+        *reinterpret_cast<int4*>(local_offsets + 12) =
+            __ldca(reinterpret_cast<const int4*>(&input_offset_by_experts[chunk_start + 12]));
+        local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
+
+// Check against the 16 loaded offsets
+#pragma unroll
+        for (int i = 0; i < 16; i++) {
+          if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
+            rowIdx_in_expert = rowIdx - local_offsets[i];
+            expert_idx = chunk_start + i;
+            break;
+          }
+        }
+      }
+    }
+
+    // Early exit when using masks.
+    if (use_mask && rowIdx_in_expert >= mask[expert_idx]) {
+      continue;
+    }
+
+    int64_t inOffset = rowIdx * actualColsPerRow + colIdx;
+    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
+    if (use_mask || use_silu_and_mul) {
+      PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
+      silu_and_mul(in_vec, in_vec_mul);
+    }
+
+    // Get the output tensor offset.
+    // Same as inOffset because 8 elements are packed into one uint32_t.
+    int64_t outOffset = rowIdx * colsPerRow + colIdx;
+    auto& out_pos = out[outOffset];
+
+    // Get the global scaling factor, which will be applied to the SF.
+    // Note SFScale is the same as next GEMM's alpha, which is
+    // (448.f / (Alpha_A / 6.f)).
+    float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
+
+    int factor = CVT_FP4_SF_VEC_SIZE * 4;
+    // The actual output_scales dim is computed from the padded numCols.
+    int32_t numCols_padded = (numCols + factor - 1) / factor * factor;
+    int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4;
+    uint32_t* SFout_in_expert = SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout;
+
+    auto sf_out = cvt_quant_to_fp4_get_sf_out_offset<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(
+        rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
+
+    out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
+  }
+#endif
+}
+
+// Use UE4M3 by default.
+template <class Type, bool UE8M0_SF = false>
+__global__ void
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+__launch_bounds__(512, 4) cvt_fp16_to_fp4_expert(
+#else
+cvt_fp16_to_fp4_expert(
+#endif
+    int32_t numRows,
+    int32_t numCols,
+    Type const* in,
+    float const* SFScale,
+    uint32_t* out,
+    uint32_t* SFout,
+    int32_t* mask,
+    bool use_silu_and_mul,
+    int n_experts) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  using PackedVec = PackedVec<Type>;
+  static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
+  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched.");
+
+  // Input tensor row/col loops.
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = (gridDim.x * blockDim.x) / n_experts;
+  int remainder = (gridDim.x * blockDim.x) % n_experts;
+  int expert_idx;
+  int tid_in_expert;
+  int actual_stride;
+  if (remainder > 0) {
+    int bound = remainder * (stride + 1);
+    if (tid < bound) {
+      expert_idx = tid / (stride + 1);
+      tid_in_expert = tid % (stride + 1);
+      actual_stride = stride + 1;
+    } else {
+      expert_idx = remainder + (tid - bound) / stride;
+      tid_in_expert = (tid - bound) % stride;
+      actual_stride = stride;
+    }
+  } else {
+    expert_idx = tid / stride;
+    tid_in_expert = tid % stride;
+    actual_stride = stride;
+  }
+  int m = numRows / n_experts;
+  int padded_m = (m + (128 - 1)) / 128 * 128;
+
+  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
+  // TODO(kaixih@nvidia): For now, we assume mask is used together with
+  // silu_and_mal. Maybe we want a more general behavior of mask later. In the
+  // silu case, the input last dim doubles.
+  bool use_mask = mask != nullptr;
+  int actualColsPerRow = use_silu_and_mul ? colsPerRow * 2 : colsPerRow;
+
+  // Each global thread processes one element
+  for (int globalIdx = tid_in_expert + expert_idx * m * colsPerRow; globalIdx < (expert_idx + 1) * m * colsPerRow;
+       globalIdx += actual_stride) {
+    // Calculate which row and column this global thread should process
+    int rowIdx = globalIdx / colsPerRow;
+    int colIdx = globalIdx % colsPerRow;
+
+    // Find index within the experts
+    int rowIdx_in_expert = rowIdx - expert_idx * m;
+
+    // Early exit when using masks.
+    if (use_mask && rowIdx_in_expert >= mask[expert_idx]) {
+      break;
+    }
+
+    int64_t inOffset = rowIdx * actualColsPerRow + colIdx;
+    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
+    if (use_silu_and_mul) {
+      PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
+      silu_and_mul(in_vec, in_vec_mul);
+    }
+
+    // Get the output tensor offset.
+    // Same as inOffset because 8 elements are packed into one uint32_t.
+    int64_t outOffset = rowIdx * colsPerRow + colIdx;
+    auto& out_pos = out[outOffset];
+
+    // Get the global scaling factor, which will be applied to the SF.
+    // Note SFScale is the same as next GEMM's alpha, which is
+    // (448.f / (Alpha_A / 6.f)).
+    float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
+
+    int factor = CVT_FP4_SF_VEC_SIZE * 4;
+    // The actual output_scales dim is computed from the padded numCols.
+    int32_t numCols_padded = (numCols + factor - 1) / factor * factor;
+    int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4;
+    uint32_t* SFout_in_expert = SFout + expert_idx * padded_m * numCols_SFout;
+
+    auto sf_out = cvt_quant_to_fp4_get_sf_out_offset<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(
+        rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
+
+    out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
+  }
+#endif
+}
+
+// Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
+template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
+__global__ void
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+__launch_bounds__(1024, 4) cvt_fp16_to_fp4(
+#else
+cvt_fp16_to_fp4(
+#endif
+    int32_t numRows,
+    int32_t numCols,
+    Type const* in,
+    float const* SFScale,
+    uint32_t* out,
+    uint32_t* SFout,
+    uint32_t* input_offset_by_experts,
+    uint32_t* output_scale_offset_by_experts,
+    int32_t* mask,
+    int n_experts,
+    bool use_silu_and_mul) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  using PackedVec = PackedVec<Type>;
+  static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
+  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched.");
+  extern __shared__ uint32_t shared_input_offsets[];
+
+  // Load input offsets into shared memory.
+  // If n_experts is larger than 4, use vectorized int4 to save instructions.
+  // If n_experts is smaller than 4, read directly.
+  if constexpr (SMALL_NUM_EXPERTS) {
+    for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
+      shared_input_offsets[i] = input_offset_by_experts[i];
+    }
+  } else {
+    for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
+      *reinterpret_cast<int4*>(&shared_input_offsets[i]) = *reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
+    }
+    if (threadIdx.x == 0) {
+      shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
+    }
+  }
+
+  __syncthreads();
+
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
+  bool use_mask = mask != nullptr;
+  // When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated).
+  int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow;
+
+  // Each global thread processes one element
+  for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) {
+    // Calculate which row and column this global thread should process
+    int rowIdx = globalIdx / colsPerRow;
+    int colIdx = globalIdx % colsPerRow;
+
+    // Find expert using binary search for better performance with large m_topk
+    int rowIdx_in_expert = 0;
+    int expert_idx = 0;
+
+    // Binary search through experts using shared memory
+    int left = 0, right = n_experts - 1;
+    while (left <= right) {
+      int mid = (left + right) / 2;
+      // Get offsets: shared_input_offsets[i] corresponds to
+      // input_offset_by_experts[i]
+      uint32_t mid_offset = shared_input_offsets[mid];
+      uint32_t next_offset = shared_input_offsets[mid + 1];
+
+      if (rowIdx >= mid_offset && rowIdx < next_offset) {
+        rowIdx_in_expert = rowIdx - mid_offset;
+        expert_idx = mid;
+        break;
+      } else if (rowIdx < mid_offset) {
+        right = mid - 1;
+      } else {
+        left = mid + 1;
+      }
+    }
+
+    if (use_mask && rowIdx_in_expert >= mask[expert_idx]) {
+      continue;
+    }
+
+    int64_t inOffset = rowIdx * actualColsPerRow + colIdx;
+
+    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
+    if (use_mask || use_silu_and_mul) {
+      PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
+      silu_and_mul(in_vec, in_vec_mul);
+    }
+
+    int64_t outOffset = rowIdx * colsPerRow + colIdx;
+    auto& out_pos = out[outOffset];
+
+    float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
+
+    int factor = CVT_FP4_SF_VEC_SIZE * 4;
+    int32_t numCols_padded = (numCols + factor - 1) / factor * factor;
+    int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4;
+    uint32_t* SFout_in_expert = SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout;
+
+    auto sf_out = cvt_quant_to_fp4_get_sf_out_offset<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(
+        rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
+
+    out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
+  }
+#endif
+}
+
+template <typename T>
+void quant_impl(
+    void* output,
+    void* output_scale,
+    void* input,
+    void* input_global_scale,
+    void* input_offset_by_experts,
+    void* output_scale_offset_by_experts,
+    void* mask,
+    bool use_silu_and_mul,
+    int m_topk,
+    int k,
+    int n_experts,
+    cudaStream_t stream) {
+  // TODO: this multiProcessorCount should be cached.
+  int device;
+  cudaGetDevice(&device);
+  int multiProcessorCount;
+  cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, device);
+
+  // Grid, Block size.
+  // Each thread converts 8 values.
+  int const workSizePerRow = k / ELTS_PER_THREAD;
+  int const totalWorkSize = m_topk * workSizePerRow;
+  dim3 block(std::min(workSizePerRow, 512));
+  // Get number of blocks per SM (assume we can fully utilize the SM).
+  int const numBlocksPerSM = 2048 / block.x;
+  dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x), multiProcessorCount * numBlocksPerSM));
+  while (grid.x <= multiProcessorCount && block.x > 64) {
+    grid.x *= 2;
+    block.x = (block.x + 1) / 2;
+  }
+
+  // TODO(kaixih@nvidia): Should relax this to allow any grid size.
+  if (mask != nullptr) {
+    grid.x = (grid.x + n_experts - 1) / n_experts * n_experts;
+    cvt_fp16_to_fp4_expert<T, false><<<grid, block, 0, stream>>>(
+        m_topk,
+        k,
+        reinterpret_cast<T*>(input),
+        reinterpret_cast<float*>(input_global_scale),
+        reinterpret_cast<uint32_t*>(output),
+        reinterpret_cast<uint32_t*>(output_scale),
+        reinterpret_cast<int32_t*>(mask),
+        use_silu_and_mul,
+        n_experts);
+    return;
+  }
+
+  int const blockRepeat = (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
+  if (blockRepeat > 1) {
+    size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
+    if (n_experts >= 4) {
+      cvt_fp16_to_fp4<T, false, false><<<grid, block, shared_mem_size, stream>>>(
+          m_topk,
+          k,
+          reinterpret_cast<T*>(input),
+          reinterpret_cast<float*>(input_global_scale),
+          reinterpret_cast<uint32_t*>(output),
+          reinterpret_cast<uint32_t*>(output_scale),
+          reinterpret_cast<uint32_t*>(input_offset_by_experts),
+          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+          reinterpret_cast<int32_t*>(mask),
+          n_experts,
+          use_silu_and_mul);
+    } else {
+      cvt_fp16_to_fp4<T, false, true><<<grid, block, shared_mem_size, stream>>>(
+          m_topk,
+          k,
+          reinterpret_cast<T*>(input),
+          reinterpret_cast<float*>(input_global_scale),
+          reinterpret_cast<uint32_t*>(output),
+          reinterpret_cast<uint32_t*>(output_scale),
+          reinterpret_cast<uint32_t*>(input_offset_by_experts),
+          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+          reinterpret_cast<int32_t*>(mask),
+          n_experts,
+          use_silu_and_mul);
+    }
+  } else {
+    if (n_experts >= 16) {
+      cvt_fp16_to_fp4<T, false, false><<<grid, block, 0, stream>>>(
+          m_topk,
+          k,
+          reinterpret_cast<T*>(input),
+          reinterpret_cast<float*>(input_global_scale),
+          reinterpret_cast<uint32_t*>(output),
+          reinterpret_cast<uint32_t*>(output_scale),
+          reinterpret_cast<uint32_t*>(input_offset_by_experts),
+          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+          reinterpret_cast<int32_t*>(mask),
+          n_experts,
+          /* bool low_latency */ true,
+          use_silu_and_mul);
+    } else {
+      cvt_fp16_to_fp4<T, false, true><<<grid, block, 0, stream>>>(
+          m_topk,
+          k,
+          reinterpret_cast<T*>(input),
+          reinterpret_cast<float*>(input_global_scale),
+          reinterpret_cast<uint32_t*>(output),
+          reinterpret_cast<uint32_t*>(output_scale),
+          reinterpret_cast<uint32_t*>(input_offset_by_experts),
+          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+          reinterpret_cast<int32_t*>(mask),
+          n_experts,
+          /* bool low_latency */ true,
+          use_silu_and_mul);
+    }
+  }
+}
+
+inline int getSMVersion(int device_id) {
+  int sm_major = 0;
+  int sm_minor = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
+  return sm_major * 10 + sm_minor;
+}
+
+void scaled_fp4_experts_quant_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts) {
+  auto MTopK = SymbolicSize{"m_topk"};
+  auto K = SymbolicSize{"k"};
+  auto OutputCols = SymbolicSize{"output_cols"};
+  auto OutputScaleRows = SymbolicSize{"output_scale_rows"};
+  auto OutputScaleCols = SymbolicSize{"output_scale_cols"};
+  auto NExperts = SymbolicSize{"n_experts"};
+  auto OffsetSize = SymbolicSize{"offset_size"};
+  auto device = SymbolicDevice{};
+
+  TensorMatcher({MTopK, K})  //
+      .with_dtype<fp16_t, bf16_t>()
+      .template with_device<kDLCUDA>(device)
+      .verify(input);
+  TensorMatcher({MTopK, OutputCols})  //
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(output);
+  TensorMatcher({OutputScaleRows, OutputScaleCols})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(output_scale);
+  TensorMatcher({NExperts})  //
+      .with_dtype<float>()
+      .with_device(device)
+      .verify(input_global_scale);
+  TensorMatcher({OffsetSize})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(input_offset_by_experts)
+      .verify(output_scale_offset_by_experts);
+
+  const int device_id = input.device().device_id;
+  RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+");
+
+  const int BLOCK_SIZE = 16;
+  const auto m_topk = static_cast<int>(MTopK.unwrap());
+  const auto k = static_cast<int>(K.unwrap());
+  RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
+  const auto n_experts = static_cast<int>(NExperts.unwrap());
+  const auto offset_size = static_cast<int>(OffsetSize.unwrap());
+  RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch");
+  RuntimeCheck(static_cast<int>(OutputCols.unwrap()) == k / 2, "output second dim mismatch");
+  const int scales_k = k / BLOCK_SIZE;
+  const int padded_k = (scales_k + 3) / 4 * 4;
+  RuntimeCheck(static_cast<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
+
+  const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
+  if (host::is_type<fp16_t>(input.dtype())) {
+    quant_impl<half>(
+        output.data_ptr(),
+        output_scale.data_ptr(),
+        input.data_ptr(),
+        input_global_scale.data_ptr(),
+        input_offset_by_experts.data_ptr(),
+        output_scale_offset_by_experts.data_ptr(),
+        nullptr,  // mask
+        false,    // use_silu_and_mul
+        m_topk,
+        k,
+        n_experts,
+        stream);
+  } else {
+    quant_impl<__nv_bfloat16>(
+        output.data_ptr(),
+        output_scale.data_ptr(),
+        input.data_ptr(),
+        input_global_scale.data_ptr(),
+        input_offset_by_experts.data_ptr(),
+        output_scale_offset_by_experts.data_ptr(),
+        nullptr,  // mask
+        false,    // use_silu_and_mul
+        m_topk,
+        k,
+        n_experts,
+        stream);
+  }
+}
+
+void silu_and_mul_scaled_fp4_experts_quant_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView mask,
+    bool use_silu_and_mul) {
+  auto MTopK = SymbolicSize{"m_topk"};
+  auto KBy2 = SymbolicSize{"k_by_2"};
+  auto OutputCols = SymbolicSize{"output_cols"};
+  auto OutputScaleRows = SymbolicSize{"output_scale_rows"};
+  auto OutputScaleCols = SymbolicSize{"output_scale_cols"};
+  auto NExperts = SymbolicSize{"n_experts"};
+  auto device = SymbolicDevice{};
+
+  TensorMatcher({MTopK, KBy2})  //
+      .with_dtype<fp16_t, bf16_t>()
+      .template with_device<kDLCUDA>(device)
+      .verify(input);
+  TensorMatcher({MTopK, OutputCols})  //
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(output);
+  TensorMatcher({OutputScaleRows, OutputScaleCols})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(output_scale);
+  TensorMatcher({NExperts})  //
+      .with_dtype<float>()
+      .with_device(device)
+      .verify(input_global_scale);
+  TensorMatcher({NExperts})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(mask);
+
+  const int device_id = input.device().device_id;
+  RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+");
+
+  const int BLOCK_SIZE = 16;
+  const auto m_topk = static_cast<int>(MTopK.unwrap());
+  const auto k_by_2 = static_cast<int>(KBy2.unwrap());
+  int k = k_by_2;
+  if (use_silu_and_mul) {
+    RuntimeCheck(k_by_2 % 2 == 0, "k must be a multiple of 2");
+    k = k_by_2 / 2;
+  }
+  const auto n_experts = static_cast<int>(NExperts.unwrap());
+  RuntimeCheck(static_cast<int>(OutputCols.unwrap()) == k / 2, "output second dim mismatch");
+  const int scales_k = k / BLOCK_SIZE;
+  const int padded_k = (scales_k + 3) / 4 * 4;
+  RuntimeCheck(static_cast<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
+
+  const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
+  if (host::is_type<fp16_t>(input.dtype())) {
+    quant_impl<half>(
+        output.data_ptr(),
+        output_scale.data_ptr(),
+        input.data_ptr(),
+        input_global_scale.data_ptr(),
+        nullptr,  // input_offset_by_experts
+        nullptr,  // output_scale_offset_by_experts
+        mask.data_ptr(),
+        use_silu_and_mul,
+        m_topk,
+        k,
+        n_experts,
+        stream);
+  } else {
+    quant_impl<__nv_bfloat16>(
+        output.data_ptr(),
+        output_scale.data_ptr(),
+        input.data_ptr(),
+        input_global_scale.data_ptr(),
+        nullptr,  // input_offset_by_experts
+        nullptr,  // output_scale_offset_by_experts
+        mask.data_ptr(),
+        use_silu_and_mul,
+        m_topk,
+        k,
+        n_experts,
+        stream);
+  }
+}
+
+void silu_and_mul_scaled_fp4_experts_quant_packed_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts) {
+  auto MTopK = SymbolicSize{"m_topk"};
+  auto KBy2 = SymbolicSize{"k_by_2"};
+  auto OutputCols = SymbolicSize{"output_cols"};
+  auto OutputScaleRows = SymbolicSize{"output_scale_rows"};
+  auto OutputScaleCols = SymbolicSize{"output_scale_cols"};
+  auto NExperts = SymbolicSize{"n_experts"};
+  auto OffsetSize = SymbolicSize{"offset_size"};
+  auto device = SymbolicDevice{};
+
+  TensorMatcher({MTopK, KBy2})  //
+      .with_dtype<fp16_t, bf16_t>()
+      .template with_device<kDLCUDA>(device)
+      .verify(input);
+  TensorMatcher({MTopK, OutputCols})  //
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(output);
+  TensorMatcher({OutputScaleRows, OutputScaleCols})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(output_scale);
+  TensorMatcher({NExperts})  //
+      .with_dtype<float>()
+      .with_device(device)
+      .verify(input_global_scale);
+  TensorMatcher({OffsetSize})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(input_offset_by_experts)
+      .verify(output_scale_offset_by_experts);
+
+  const int device_id = input.device().device_id;
+  RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+");
+
+  const int BLOCK_SIZE = 16;
+  const auto m_topk = static_cast<int>(MTopK.unwrap());
+  const auto k_by_2 = static_cast<int>(KBy2.unwrap());
+  // Input last dim is 2*k (gate+up concatenated). The kernel does SiLU(gate)*up
+  // then FP4-quantizes the k-dim result.
+  RuntimeCheck(k_by_2 % 2 == 0, "input last dim must be even (2*k)");
+  const int k = k_by_2 / 2;
+  RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
+  const auto n_experts = static_cast<int>(NExperts.unwrap());
+  const auto offset_size = static_cast<int>(OffsetSize.unwrap());
+  RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch");
+  RuntimeCheck(static_cast<int>(OutputCols.unwrap()) == k / 2, "output second dim mismatch");
+  const int scales_k = k / BLOCK_SIZE;
+  const int padded_k = (scales_k + 3) / 4 * 4;
+  RuntimeCheck(static_cast<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
+
+  const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
+  if (host::is_type<fp16_t>(input.dtype())) {
+    quant_impl<half>(
+        output.data_ptr(),
+        output_scale.data_ptr(),
+        input.data_ptr(),
+        input_global_scale.data_ptr(),
+        input_offset_by_experts.data_ptr(),
+        output_scale_offset_by_experts.data_ptr(),
+        nullptr,  // mask
+        true,     // use_silu_and_mul
+        m_topk,
+        k,
+        n_experts,
+        stream);
+  } else {
+    quant_impl<__nv_bfloat16>(
+        output.data_ptr(),
+        output_scale.data_ptr(),
+        input.data_ptr(),
+        input_global_scale.data_ptr(),
+        input_offset_by_experts.data_ptr(),
+        output_scale_offset_by_experts.data_ptr(),
+        nullptr,  // mask
+        true,     // use_silu_and_mul
+        m_topk,
+        k,
+        n_experts,
+        stream);
+  }
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh

@@ -0,0 +1,160 @@
+/* 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.
+==============================================================================*/
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <cutlass/arch/config.h>
+
+#include <cuda.h>
+#include <cuda_fp8.h>
+
+#define ELTS_PER_THREAD 8
+
+constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
+constexpr int CVT_FP4_SF_VEC_SIZE = 16;
+
+// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
+SGL_DEVICE uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
+  // PTX instructions used here requires >= sm100f.
+#if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \
+    (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000))
+  uint32_t val;
+  asm volatile(
+      "{\n"
+      ".reg .b8 byte0;\n"
+      ".reg .b8 byte1;\n"
+      ".reg .b8 byte2;\n"
+      ".reg .b8 byte3;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte0, %2, %1;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte1, %4, %3;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte2, %6, %5;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte3, %8, %7;\n"
+      "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
+      "}"
+      : "=r"(val)
+      : "f"(array[0]),
+        "f"(array[1]),
+        "f"(array[2]),
+        "f"(array[3]),
+        "f"(array[4]),
+        "f"(array[5]),
+        "f"(array[6]),
+        "f"(array[7]));
+  return val;
+#else
+  printf("fp32_vec_to_e2m1 is not supported on this architecture\n");
+  __trap();
+  return 0;
+#endif
+}
+
+// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
+SGL_DEVICE uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
+  // PTX instructions used here requires >= sm100f.
+#if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \
+    (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000))
+  uint32_t val;
+  asm volatile(
+      "{\n"
+      ".reg .b8 byte0;\n"
+      ".reg .b8 byte1;\n"
+      ".reg .b8 byte2;\n"
+      ".reg .b8 byte3;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte0, %2, %1;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte1, %4, %3;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte2, %6, %5;\n"
+      "cvt.rn.satfinite.e2m1x2.f32   byte3, %8, %7;\n"
+      "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
+      "}"
+      : "=r"(val)
+      : "f"(array[0].x),
+        "f"(array[0].y),
+        "f"(array[1].x),
+        "f"(array[1].y),
+        "f"(array[2].x),
+        "f"(array[2].y),
+        "f"(array[3].x),
+        "f"(array[3].y));
+  return val;
+#else
+  printf("fp32_vec_to_e2m1 is not supported on this architecture\n");
+  __trap();
+  return 0;
+#endif
+}
+
+// Fast reciprocal.
+SGL_DEVICE float reciprocal_approximate_ftz(float a) {
+  float b;
+  asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
+  return b;
+}
+
+template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
+SGL_DEVICE uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx, int numCols, SFType* SFout) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || CVT_FP4_NUM_THREADS_PER_SF == 2);
+
+  // One pair of threads write one SF to global memory.
+  // TODO: stage through smem for packed STG.32
+  // is it better than STG.8 from 4 threads ?
+  if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
+    // SF vector index (16 elements share one SF in the K dimension).
+    int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
+    int32_t mIdx = rowIdx;
+
+    // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
+    // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
+
+    int32_t mTileIdx = mIdx / (32 * 4);
+    // SF vector size 16.
+    int factor = CVT_FP4_SF_VEC_SIZE * 4;
+    int32_t numKTiles = (numCols + factor - 1) / factor;
+    int64_t mTileStride = numKTiles * 32 * 4 * 4;
+
+    int32_t kTileIdx = (kIdx / 4);
+    int64_t kTileStride = 32 * 4 * 4;
+
+    // M tile layout [32, 4] is column-major.
+    int32_t outerMIdx = (mIdx % 32);
+    int64_t outerMStride = 4 * 4;
+
+    int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
+    int64_t innerMStride = 4;
+
+    int32_t innerKIdx = (kIdx % 4);
+    int64_t innerKStride = 1;
+
+    // Compute the global offset.
+    int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride + outerMIdx * outerMStride +
+                       innerMIdx * innerMStride + innerKIdx * innerKStride;
+
+    return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
+  }
+#endif
+  return nullptr;
+}
+
+// Define a 16 bytes packed data type.
+template <class Type>
+struct PackedVec {
+  packed_t<Type> elts[4];
+};
+
+template <>
+struct PackedVec<__nv_fp8_e4m3> {
+  __nv_fp8x2_e4m3 elts[8];
+};

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh

@@ -0,0 +1,87 @@
+/* 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.
+==============================================================================*/
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+void scaled_fp4_quant_sm100a_sm120a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView output_sf,
+    tvm::ffi::TensorView input_sf);
+
+void scaled_fp4_experts_quant_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts);
+
+void silu_and_mul_scaled_fp4_experts_quant_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView mask,
+    bool use_silu_and_mul);
+
+void silu_and_mul_scaled_fp4_experts_quant_packed_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts);
+
+void scaled_fp4_quant(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView output_sf,
+    tvm::ffi::TensorView input_sf) {
+  scaled_fp4_quant_sm100a_sm120a(output, input, output_sf, input_sf);
+}
+
+void scaled_fp4_experts_quant(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts) {
+  scaled_fp4_experts_quant_sm100a(
+      output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts);
+}
+
+void silu_and_mul_scaled_fp4_experts_quant(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView mask,
+    bool use_silu_and_mul) {
+  silu_and_mul_scaled_fp4_experts_quant_sm100a(output, output_scale, input, input_global_scale, mask, use_silu_and_mul);
+}
+
+void silu_and_mul_scaled_fp4_experts_quant_packed(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts) {
+  silu_and_mul_scaled_fp4_experts_quant_packed_sm100a(
+      output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts);
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh

@@ -0,0 +1,241 @@
+/* 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.
+==============================================================================*/
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include "nvfp4_quant.cuh"
+#include <cuda_runtime.h>
+#include <cuda_runtime_api.h>
+
+using namespace host;
+
+// Quantizes the provided PackedVec into the uint32_t output
+template <class Type, bool UE8M0_SF = false>
+SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  // Get absolute maximum values among the local 8 values.
+  auto localMax = __habs2(vec.elts[0]);
+
+// Local maximum value.
+#pragma unroll
+  for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
+    localMax = __hmax2(localMax, __habs2(vec.elts[i]));
+  }
+
+  // Get the absolute maximum among all 16 values (two threads).
+  localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
+  // Get the final absolute maximum values.
+  float vecMax = float(__hmax(localMax.x, localMax.y));
+
+  // Get the SF (max value of the vector / max value of e2m1).
+  // maximum value of e2m1 = 6.0.
+  // TODO: use half as compute data type.
+  float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
+  // 8 bits representation of the SF.
+  uint8_t fp8SFVal;
+  // Write the SF to global memory (STG.8).
+  if constexpr (UE8M0_SF) {
+    __nv_fp8_e8m0 tmp;
+    tmp.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf);
+    SFValue = static_cast<float>(tmp);
+    fp8SFVal = tmp.__x;
+  } else {
+    // Here SFValue is always positive, so E4M3 is the same as UE4M3.
+    __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
+    fp8SFVal = tmp.__x;
+    SFValue = static_cast<float>(tmp);
+  }
+  // Get the output scale.
+  // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
+  //                       reciprocal(SFScaleVal))
+  float outputScale =
+      SFValue != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f;
+
+  if (SFout) {
+    // Write the SF to global memory (STG.8).
+    *SFout = fp8SFVal;
+  }
+
+  // Convert the input to float.
+  float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
+
+#pragma unroll
+  for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
+    if constexpr (std::is_same_v<Type, half>) {
+      fp2Vals[i] = __half22float2(vec.elts[i]);
+    } else {
+      fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
+    }
+    fp2Vals[i].x *= outputScale;
+    fp2Vals[i].y *= outputScale;
+  }
+
+  // Convert to e2m1 values.
+  uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
+
+  // Write the e2m1 values to global memory.
+  return e2m1Vec;
+#else
+  return 0;
+#endif
+}
+
+// Use UE4M3 by default.
+template <class Type, bool UE8M0_SF = false>
+__global__ void
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+__launch_bounds__(512, 4) cvt_fp16_to_fp4(
+#else
+cvt_fp16_to_fp4(
+#endif
+    int32_t numRows, int32_t numCols, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  using PackedVec = PackedVec<Type>;
+  static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
+  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched.");
+
+  // Get the global scaling factor, which will be applied to the SF.
+  // Note SFScale is the same as next GEMM's alpha, which is
+  // (448.f / (Alpha_A / 6.f)).
+  float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0];
+
+  // Input tensor row/col loops.
+  for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
+    for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD; colIdx += blockDim.x) {
+      int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
+      PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
+      // Get the output tensor offset.
+      // Same as inOffset because 8 elements are packed into one uint32_t.
+      int64_t outOffset = inOffset;
+      auto& out_pos = out[outOffset];
+
+      auto sf_out =
+          cvt_quant_to_fp4_get_sf_out_offset<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(rowIdx, colIdx, numCols, SFout);
+
+      out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
+    }
+  }
+#endif
+}
+
+template <typename T>
+void invokeFP4Quantization(
+    int m,
+    int n,
+    T const* input,
+    float const* SFScale,
+    int64_t* output,
+    int32_t* SFOuput,
+    bool useUE8M0,
+    int multiProcessorCount,
+    cudaStream_t stream) {
+  // Grid, Block size.
+  // Each thread converts 8 values.
+  dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
+  // Get number of blocks per SM (assume we can fully utilize the SM).
+  int const numBlocksPerSM = 2048 / block.x;
+  dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
+
+  // Launch the cvt kernel.
+  if (useUE8M0) {
+    cvt_fp16_to_fp4<T, true><<<grid, block, 0, stream>>>(
+        m, n, input, SFScale, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(SFOuput));
+  } else {
+    cvt_fp16_to_fp4<T, false><<<grid, block, 0, stream>>>(
+        m, n, input, SFScale, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(SFOuput));
+  }
+}
+
+// Instantiate the function.
+template void invokeFP4Quantization(
+    int m,
+    int n,
+    half const* input,
+    float const* SFScale,
+    int64_t* output,
+    int32_t* SFOuput,
+    bool useUE8M0,
+    int multiProcessorCount,
+    cudaStream_t stream);
+
+template void invokeFP4Quantization(
+    int m,
+    int n,
+    __nv_bfloat16 const* input,
+    float const* SFScale,
+    int64_t* output,
+    int32_t* SFOuput,
+    bool useUE8M0,
+    int multiProcessorCount,
+    cudaStream_t stream);
+
+inline int getSMVersion(int device_id) {
+  int sm_major = 0;
+  int sm_minor = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
+  return sm_major * 10 + sm_minor;
+}
+
+void scaled_fp4_quant_sm100a_sm120a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView output_sf,
+    tvm::ffi::TensorView input_sf) {
+  RuntimeCheck(input.device().device_type == kDLCUDA, "input must be a CUDA tensor");
+  RuntimeCheck(output.device() == input.device(), "output and input must be on same device");
+  RuntimeCheck(output_sf.device() == input.device(), "output_sf and input must be on same device");
+  RuntimeCheck(input_sf.device() == input.device(), "input_sf and input must be on same device");
+  RuntimeCheck(input.dim() == 2, "input must be a 2D tensor");
+  RuntimeCheck(output.dim() == 2, "output must be a 2D tensor");
+  RuntimeCheck(output_sf.dim() == 2, "output_sf must be a 2D tensor");
+  RuntimeCheck(input_sf.numel() == 1, "input_sf must have exactly one element");
+  RuntimeCheck(host::is_type<uint8_t>(output.dtype()), "output must be uint8");
+  RuntimeCheck(host::is_type<int32_t>(output_sf.dtype()), "output_sf must be int32");
+  RuntimeCheck(host::is_type<float>(input_sf.dtype()), "input_sf must be float32");
+  RuntimeCheck(
+      host::is_type<fp16_t>(input.dtype()) || host::is_type<bf16_t>(input.dtype()), "input dtype must be fp16 or bf16");
+
+  const int device_id = input.device().device_id;
+  const auto sm_version = getSMVersion(device_id);
+  RuntimeCheck(sm_version >= 100, "fp4_quant is only supported on sm100+");
+
+  const int32_t m = static_cast<int32_t>(input.size(0));
+  const int32_t n = static_cast<int32_t>(input.size(1));
+
+  RuntimeCheck(output.size(0) == m, "output row size mismatch");
+  RuntimeCheck(output.size(1) == n / 2, "output column size mismatch");
+  RuntimeCheck(n % 16 == 0, "The N dimension must be multiple of 16.");
+
+  const int multiProcessorCount = static_cast<int>(runtime::get_sm_count(device_id));
+
+  auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
+  auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
+  auto output_ptr = static_cast<int64_t*>(output.data_ptr());
+  const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
+
+  constexpr bool useUE8M0 = false;
+  if (host::is_type<fp16_t>(input.dtype())) {
+    auto input_ptr = reinterpret_cast<half const*>(input.data_ptr());
+    invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream);
+  } else {
+    auto input_ptr = reinterpret_cast<__nv_bfloat16 const*>(input.data_ptr());
+    invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream);
+  }
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_common.cuh

@@ -0,0 +1,66 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#pragma once
+
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime.h>
+
+using namespace host;
+
+// clang-format off
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/util/packed_stride.hpp"
+// clang-format on
+
+#define CUTLASS_CHECK(status)                                                        \
+  {                                                                                  \
+    cutlass::Status error = status;                                                  \
+    RuntimeCheck(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
+  }
+
+using namespace cute;
+
+inline uint32_t next_pow_2(uint32_t x) noexcept {
+  if (x <= 1) return 1;
+  return 1u << (32 - __builtin_clz(x - 1));
+}
+
+inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor {
+  if (required_bytes == 0) return {};
+  DLDataType u8 = {kDLUInt, 8, 1};
+  int64_t shape[] = {static_cast<int64_t>(required_bytes)};
+  return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device);
+}
+
+inline int getSMVersion(int device_id) {
+  int sm_major = 0;
+  int sm_minor = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
+  return sm_major * 10 + sm_minor;
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh

@@ -0,0 +1,34 @@
+/* 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.
+==============================================================================*/
+
+#include <sgl_kernel/tensor.h>
+
+void cutlass_scaled_fp4_mm_sm100a_sm120a(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha);
+
+void cutlass_scaled_fp4_mm(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha) {
+  cutlass_scaled_fp4_mm_sm100a_sm120a(D, A, B, A_sf, B_sf, alpha);
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh

@@ -0,0 +1,146 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#include "nvfp4_scaled_mm_common.cuh"
+#include "nvfp4_scaled_mm_sm100.cuh"
+#include "nvfp4_scaled_mm_sm120.cuh"
+
+void cutlass_scaled_fp4_mm_sm100a_sm120a(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha) {
+  RuntimeCheck(A.device().device_type == kDLCUDA, "a must be a CUDA tensor");
+  RuntimeCheck(B.device().device_type == kDLCUDA, "b must be a CUDA tensor");
+  RuntimeCheck(A_sf.device().device_type == kDLCUDA, "scale_a must be a CUDA tensor");
+  RuntimeCheck(B_sf.device().device_type == kDLCUDA, "scale_b must be a CUDA tensor");
+  RuntimeCheck(alpha.device().device_type == kDLCUDA, "alpha must be a CUDA tensor");
+  RuntimeCheck(D.device().device_type == kDLCUDA, "out must be a CUDA tensor");
+
+  RuntimeCheck(A.device() == B.device(), "a and b must be on same device");
+  RuntimeCheck(A.device() == A_sf.device(), "a and scale_a must be on same device");
+  RuntimeCheck(A.device() == B_sf.device(), "a and scale_b must be on same device");
+  RuntimeCheck(A.device() == alpha.device(), "a and alpha must be on same device");
+  RuntimeCheck(A.device() == D.device(), "a and out must be on same device");
+
+  RuntimeCheck(A.is_contiguous(), "a must be contiguous");
+  RuntimeCheck(B.is_contiguous(), "b must be contiguous");
+  RuntimeCheck(A_sf.is_contiguous(), "scale_a must be contiguous");
+  RuntimeCheck(B_sf.is_contiguous(), "scale_b must be contiguous");
+  RuntimeCheck(alpha.is_contiguous(), "alpha must be contiguous");
+  RuntimeCheck(D.is_contiguous(), "out must be contiguous");
+
+  RuntimeCheck(host::is_type<uint8_t>(A.dtype()), "a must be uint8");
+  RuntimeCheck(host::is_type<uint8_t>(B.dtype()), "b must be uint8");
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(A_sf.dtype()), "scale_a must be float8_e4m3fn");
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(B_sf.dtype()), "scale_b must be float8_e4m3fn");
+  RuntimeCheck(host::is_type<float>(alpha.dtype()), "alpha must be float32");
+
+  RuntimeCheck(A.dim() == 2, "a must be a matrix");
+  RuntimeCheck(B.dim() == 2, "b must be a matrix");
+  RuntimeCheck(A_sf.dim() == 2, "scale_a must be a matrix");
+  RuntimeCheck(B_sf.dim() == 2, "scale_b must be a matrix");
+  RuntimeCheck(alpha.numel() == 1, "alpha must have exactly one element");
+
+  RuntimeCheck(
+      A.size(1) == B.size(1),
+      "a and b shapes cannot be multiplied (",
+      A.size(0),
+      "x",
+      A.size(1),
+      " and ",
+      B.size(0),
+      "x",
+      B.size(1),
+      ")");
+
+  const auto m = static_cast<int64_t>(A.size(0));
+  const auto n = static_cast<int64_t>(B.size(0));
+  const auto k = static_cast<int64_t>(A.size(1) * 2);
+
+  RuntimeCheck(D.dim() == 2, "out must be 2D");
+  RuntimeCheck(D.size(0) == m, "out first dim must equal m");
+  RuntimeCheck(D.size(1) == n, "out second dim must equal n");
+
+  constexpr int alignment = 32;
+  RuntimeCheck(k % alignment == 0, "Expected k to be divisible by ", alignment, ", but got k: ", k);
+  RuntimeCheck(n % alignment == 0, "Expected n to be divisible by ", alignment, ", but got n: ", n);
+
+  auto round_up = [](int64_t x, int64_t y) { return (x + y - 1) / y * y; };
+  const int64_t rounded_m = round_up(m, 128);
+  const int64_t rounded_n = round_up(n, 128);
+  const int64_t rounded_k = round_up(k / 16, 4);
+
+  RuntimeCheck(
+      A_sf.size(1) == B_sf.size(1),
+      "scale_a and scale_b shapes cannot be multiplied (",
+      A_sf.size(0),
+      "x",
+      A_sf.size(1),
+      " and ",
+      B_sf.size(0),
+      "x",
+      B_sf.size(1),
+      ")");
+  RuntimeCheck(
+      A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k,
+      "scale_a must be padded/swizzled to shape (",
+      rounded_m,
+      "x",
+      rounded_k,
+      "), got (",
+      A_sf.size(0),
+      "x",
+      A_sf.size(1),
+      ")");
+  RuntimeCheck(
+      B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k,
+      "scale_b must be padded/swizzled to shape (",
+      rounded_n,
+      "x",
+      rounded_k,
+      "), got (",
+      B_sf.size(0),
+      "x",
+      B_sf.size(1),
+      ")");
+
+  const cudaStream_t stream = LaunchKernel::resolve_device(A.device());
+  const int sm_version = getSMVersion(A.device().device_id);
+
+  if (sm_version >= 120) {
+    if (host::is_type<fp16_t>(D.dtype())) {
+      cutlass_fp4_f16_gemm_dispatch_sm120(
+          D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
+    } else if (host::is_type<bf16_t>(D.dtype())) {
+      cutlass_fp4_bf16_gemm_dispatch_sm120(
+          D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
+    } else {
+      Panic("Unsupported output data type of nvfp4 mm sm120");
+    }
+  } else {
+    if (host::is_type<fp16_t>(D.dtype())) {
+      cutlassFp4GemmDispatchSm100<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+    } else if (host::is_type<bf16_t>(D.dtype())) {
+      cutlassFp4GemmDispatchSm100<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+    } else if (host::is_type<float>(D.dtype())) {
+      cutlassFp4GemmDispatchSm100<float>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+    } else {
+      Panic("Unsupported output data type of nvfp4 mm");
+    }
+  }
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm100.cuh

@@ -0,0 +1,305 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#pragma once
+
+#include "nvfp4_scaled_mm_common.cuh"
+
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+
+// Config(half_t/bfloat16_t) for M <= 128
+template <typename T>
+struct KernelConfigM128 {
+  using OutputType = T;
+  using MmaTileShape = Shape<_128, _256, _256>;
+  using ClusterShape = Shape<int, int, _1>;
+  using EpilogueTile = Shape<_128, _64>;  // Avoid register spilling
+  using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm;
+  using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+template <typename T>
+const dim3 KernelConfigM128<T>::preferred_cluster(1, 4, 1);
+template <typename T>
+const dim3 KernelConfigM128<T>::fallback_cluster(1, 2, 1);
+
+// Config(half_t/bfloat16_t) for M <= 256
+template <typename T>
+struct KernelConfigM256 {
+  using OutputType = T;
+  using MmaTileShape = Shape<_256, _256, _256>;
+  using ClusterShape = Shape<int, int, _1>;
+  using EpilogueTile = Shape<_128, _64>;  // Avoid register spilling
+  using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm;
+  using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+template <typename T>
+const dim3 KernelConfigM256<T>::preferred_cluster(2, 4, 1);
+template <typename T>
+const dim3 KernelConfigM256<T>::fallback_cluster(2, 1, 1);
+
+// Config(half_t/bfloat16_t) for 256 < M <= 1024
+template <typename T>
+struct KernelConfigDefault {
+  using OutputType = T;
+  using MmaTileShape = Shape<_256, _256, _256>;
+  using ClusterShape = Shape<int, int, _1>;
+  using EpilogueTile = Shape<_128, _64>;  // Avoid register spilling
+  using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm;
+  using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+template <typename T>
+const dim3 KernelConfigDefault<T>::preferred_cluster(2, 4, 1);
+template <typename T>
+const dim3 KernelConfigDefault<T>::fallback_cluster(2, 1, 1);
+
+// Config(half_t/bfloat16_t) for M > 1024: 1x4 cluster reduces M-tail waste.
+template <typename T>
+struct KernelConfigLargeM {
+  using OutputType = T;
+  using MmaTileShape = Shape<_256, _256, _256>;
+  using ClusterShape = Shape<int, int, _1>;
+  using EpilogueTile = Shape<_128, _64>;
+  using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm;
+  using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+template <typename T>
+const dim3 KernelConfigLargeM<T>::preferred_cluster(1, 4, 1);
+template <typename T>
+const dim3 KernelConfigLargeM<T>::fallback_cluster(1, 2, 1);
+
+struct KernelConfigFp32 {
+  using OutputType = float;
+  using MmaTileShape = Shape<_128, _128, _256>;
+  using ClusterShape = Shape<int, int, _1>;
+  using EpilogueTile = cutlass::epilogue::collective::EpilogueTileAuto;
+  using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm;
+  using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+const dim3 KernelConfigFp32::preferred_cluster = dim3(1, 4, 1);
+const dim3 KernelConfigFp32::fallback_cluster = dim3(1, 2, 1);
+
+template <typename KernelConfig>
+struct Fp4GemmSm100 {
+  using Config = KernelConfig;
+  using OutputType = typename KernelConfig::OutputType;
+
+  using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+  using LayoutATag = cutlass::layout::RowMajor;
+  static constexpr int AlignmentA = 32;
+
+  using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+  using LayoutBTag = cutlass::layout::ColumnMajor;
+  static constexpr int AlignmentB = 32;
+
+  using ElementD = OutputType;
+  using ElementC = OutputType;
+  using LayoutCTag = cutlass::layout::RowMajor;
+  using LayoutDTag = cutlass::layout::RowMajor;
+  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+
+  using ElementAccumulator = float;
+  using ArchTag = cutlass::arch::Sm100;
+  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
+
+  using MmaTileShape = typename KernelConfig::MmaTileShape;
+  using ClusterShape = typename KernelConfig::ClusterShape;
+  using EpilogueTile = typename KernelConfig::EpilogueTile;
+  using EpilogueSchedule = typename KernelConfig::EpilogueSchedule;
+  using MainloopSchedule = typename KernelConfig::MainloopSchedule;
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      MmaTileShape,
+      ClusterShape,
+      EpilogueTile,
+      ElementAccumulator,
+      ElementAccumulator,
+      void,
+      LayoutCTag,
+      AlignmentC,
+      ElementD,
+      LayoutDTag,
+      AlignmentD,
+      EpilogueSchedule,
+      cutlass::epilogue::fusion::LinearCombination<ElementD, float, void, float>>::CollectiveOp;
+
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      ElementA,
+      LayoutATag,
+      AlignmentA,
+      ElementB,
+      LayoutBTag,
+      AlignmentB,
+      ElementAccumulator,
+      MmaTileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+          sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      MainloopSchedule>::CollectiveOp;
+
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{}));
+  using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using LayoutB = decltype(cute::make_layout(make_shape(0, 0, 0), StrideB{}));
+  using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using LayoutC = decltype(cute::make_layout(make_shape(0, 0, 0), StrideC{}));
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+  using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{}));
+};
+
+template <typename T>
+typename T::Gemm::Arguments args_from_options(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int64_t M,
+    int64_t N,
+    int64_t K) {
+  using ElementA = typename T::Gemm::ElementA;
+  using ElementB = typename T::Gemm::ElementB;
+  using ElementSFA = cutlass::float_ue4m3_t;
+  using ElementSFB = cutlass::float_ue4m3_t;
+  using ElementD = typename T::Gemm::ElementD;
+  using ElementCompute = float;
+  using StrideA = typename T::StrideA;
+  using StrideB = typename T::StrideB;
+  using StrideD = typename T::StrideD;
+  using Sm1xxBlkScaledConfig = typename T::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
+
+  int m = static_cast<int>(M);
+  int n = static_cast<int>(N);
+  int k = static_cast<int>(K);
+  auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1});
+  auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1});
+  auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1});
+
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, 1));
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, 1));
+
+  typename T::Gemm::Arguments arguments{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {m, n, k, 1},
+      {// Mainloop arguments
+       static_cast<ElementA const*>(A.data_ptr()),
+       stride_A,
+       static_cast<ElementB const*>(B.data_ptr()),
+       stride_B,
+       static_cast<ElementSFA const*>(A_sf.data_ptr()),
+       layout_SFA,
+       static_cast<ElementSFB const*>(B_sf.data_ptr()),
+       layout_SFB},
+      {     // Epilogue arguments
+       {},  // epilogue.thread
+       nullptr,
+       stride_D,
+       static_cast<ElementD*>(D.data_ptr()),
+       stride_D}};
+  auto& fusion_args = arguments.epilogue.thread;
+  fusion_args.alpha_ptr = static_cast<ElementCompute const*>(alpha.data_ptr());
+  using KernelConfig = typename T::Config;
+  arguments.hw_info.cluster_shape = KernelConfig::preferred_cluster;
+  arguments.hw_info.cluster_shape_fallback = KernelConfig::fallback_cluster;
+  return arguments;
+}
+
+template <typename T>
+void runGemm(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int64_t m,
+    int64_t n,
+    int64_t k,
+    cudaStream_t stream) {
+  typename T::Gemm gemm;
+  auto arguments = args_from_options<T>(D, A, B, A_sf, B_sf, alpha, m, n, k);
+
+  size_t workspace_size = T::Gemm::get_workspace_size(arguments);
+  auto workspace_tensor = alloc_workspace_tensor(workspace_size, A.device());
+  void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr();
+
+  CUTLASS_CHECK(gemm.can_implement(arguments));
+
+  CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream));
+
+  CUTLASS_CHECK(gemm.run(arguments, workspace, stream));
+}
+
+template <typename OutType>
+void cutlassFp4GemmDispatchSm100(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int64_t m,
+    int64_t n,
+    int64_t k,
+    cudaStream_t stream) {
+  if (m <= 128) {
+    runGemm<Fp4GemmSm100<KernelConfigM128<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else if (m <= 256) {
+    runGemm<Fp4GemmSm100<KernelConfigM256<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else if (m <= 1024) {
+    // m in (256, 1024]: 2x4 cluster balances SM occupancy and data reuse
+    runGemm<Fp4GemmSm100<KernelConfigDefault<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else {
+    // m in (1024, inf): 1x4 cluster eliminates M-tail waste for FLUX-class shapes
+    runGemm<Fp4GemmSm100<KernelConfigLargeM<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  }
+}
+
+template <>
+void cutlassFp4GemmDispatchSm100<float>(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int64_t m,
+    int64_t n,
+    int64_t k,
+    cudaStream_t stream) {
+  runGemm<Fp4GemmSm100<KernelConfigFp32>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+}
+
+#endif  // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm120.cuh

@@ -0,0 +1,228 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#pragma once
+
+#include "nvfp4_scaled_mm_common.cuh"
+
+#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
+
+struct sm120_fp4_config_small_m {
+  using ClusterShape = Shape<_1, _1, _1>;
+  using MmaTileShape = Shape<_128, _128, _256>;
+  using PerSmTileShape_MNK = Shape<_128, _128, _256>;
+};
+
+struct sm120_fp4_config_M256 {
+  using ClusterShape = Shape<_1, _1, _1>;
+  using MmaTileShape = Shape<_128, _128, _128>;
+  using PerSmTileShape_MNK = Shape<_128, _128, _128>;
+};
+
+struct sm120_fp4_config_default {
+  using ClusterShape = Shape<_1, _1, _1>;
+  using MmaTileShape = Shape<_256, _128, _128>;
+  using PerSmTileShape_MNK = Shape<_256, _128, _128>;
+};
+
+template <typename Config, typename OutType>
+struct Fp4GemmSm120 {
+  using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+  using LayoutATag = cutlass::layout::RowMajor;
+  static constexpr int AlignmentA = 32;
+
+  using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+  using LayoutBTag = cutlass::layout::ColumnMajor;
+  static constexpr int AlignmentB = 32;
+
+  using ElementD = OutType;
+  using ElementC = OutType;
+  using LayoutCTag = cutlass::layout::RowMajor;
+  using LayoutDTag = cutlass::layout::RowMajor;
+  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+
+  using ElementAccumulator = float;
+  using ArchTag = cutlass::arch::Sm120;
+  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
+
+  using MmaTileShape = typename Config::MmaTileShape;
+  using ClusterShape = typename Config::ClusterShape;
+  using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK;
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      PerSmTileShape_MNK,
+      ClusterShape,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      ElementAccumulator,
+      ElementAccumulator,
+      void,
+      LayoutCTag,
+      AlignmentC,
+      ElementD,
+      LayoutDTag,
+      AlignmentD,
+      cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      ElementA,
+      LayoutATag,
+      AlignmentA,
+      ElementB,
+      LayoutBTag,
+      AlignmentB,
+      ElementAccumulator,
+      MmaTileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+          sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
+
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+};
+
+template <typename Gemm>
+typename Gemm::Arguments args_from_options_sm120(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int M,
+    int N,
+    int K) {
+  using ElementA = typename Gemm::ElementA;
+  using ElementB = typename Gemm::ElementB;
+  using ElementD = typename Gemm::ElementD;
+  using ElementSFA = cutlass::float_ue4m3_t;
+  using ElementSFB = cutlass::float_ue4m3_t;
+  using ElementCompute = float;
+
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
+
+  auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
+  auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
+  auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
+
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
+
+  typename Gemm::Arguments arguments{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {static_cast<ElementA const*>(A.data_ptr()),
+       stride_A,
+       static_cast<ElementB const*>(B.data_ptr()),
+       stride_B,
+       static_cast<ElementSFA const*>(A_sf.data_ptr()),
+       layout_SFA,
+       static_cast<ElementSFB const*>(B_sf.data_ptr()),
+       layout_SFB},
+      {{}, nullptr, stride_D, static_cast<ElementD*>(D.data_ptr()), stride_D}};
+  auto& fusion_args = arguments.epilogue.thread;
+  fusion_args.alpha_ptr = static_cast<ElementCompute const*>(alpha.data_ptr());
+
+  return arguments;
+}
+
+template <typename Gemm>
+void runGemmSm120(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int M,
+    int N,
+    int K,
+    cudaStream_t stream) {
+  Gemm gemm;
+
+  auto arguments = args_from_options_sm120<Gemm>(D, A, B, A_sf, B_sf, alpha, M, N, K);
+
+  size_t workspace_size = Gemm::get_workspace_size(arguments);
+  auto workspace_tensor = alloc_workspace_tensor(workspace_size, A.device());
+  void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr();
+
+  CUTLASS_CHECK(gemm.can_implement(arguments));
+
+  CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream));
+
+  CUTLASS_CHECK(gemm.run(arguments, workspace, stream));
+}
+
+void cutlass_fp4_bf16_gemm_dispatch_sm120(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int m,
+    int n,
+    int k,
+    cudaStream_t stream) {
+  uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
+  if (mp2 <= 32) {
+    runGemmSm120<Fp4GemmSm120<sm120_fp4_config_small_m, cutlass::bfloat16_t>::Gemm>(
+        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else if (mp2 <= 256) {
+    runGemmSm120<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::bfloat16_t>::Gemm>(
+        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else {
+    runGemmSm120<Fp4GemmSm120<sm120_fp4_config_default, cutlass::bfloat16_t>::Gemm>(
+        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  }
+}
+
+void cutlass_fp4_f16_gemm_dispatch_sm120(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha,
+    int m,
+    int n,
+    int k,
+    cudaStream_t stream) {
+  uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
+  if (mp2 <= 32) {
+    runGemmSm120<Fp4GemmSm120<sm120_fp4_config_small_m, cutlass::half_t>::Gemm>(
+        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else if (mp2 <= 256) {
+    runGemmSm120<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::half_t>::Gemm>(
+        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else {
+    runGemmSm120<Fp4GemmSm120<sm120_fp4_config_default, cutlass::half_t>::Gemm>(
+        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  }
+}
+
+#endif  // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)

python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh

@@ -0,0 +1,139 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/atomic.cuh>
+#include <sgl_kernel/cta.cuh>
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+constexpr size_t kBlockSize = 256;
+
+// each warp will handle 512B data
+template <typename T>
+__global__ void
+per_tensor_absmax_kernel(const T* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) {
+  using namespace device;
+  constexpr uint32_t VEC_SIZE = 16 / sizeof(T);
+
+  const int64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+
+  float max_value = 0.0f;
+  if (gid * VEC_SIZE + VEC_SIZE <= num_elements) {
+    using vec_t = AlignedVector<T, VEC_SIZE>;
+    const auto gmem_in = tile::Memory<vec_t>::thread();
+    const auto input_vec = gmem_in.load(input, gid);
+#pragma unroll
+    for (uint32_t i = 0; i < VEC_SIZE; ++i) {
+      const float value = static_cast<float>(input_vec[i]);
+      max_value = math::max(max_value, math::abs(value));
+    }
+  } else if (gid * VEC_SIZE < num_elements) {
+    [[unlikely]];  // poorly aligned case, do not optimize
+    const auto remainder = num_elements - gid * VEC_SIZE;
+    for (uint32_t i = 0; i < remainder; ++i) {
+      const float value = static_cast<float>(input[gid * VEC_SIZE + i]);
+      max_value = math::max(max_value, math::abs(value));
+    }
+  }
+
+  // reduce within block and then atomic reduce between blocks
+  __shared__ float smem[kWarpThreads];
+  cta::reduce_max(max_value, smem);
+  if (threadIdx.x == 0) {
+    const auto max_value = smem[0];
+    atomic::max(output_s, max_value / math::FP8_E4M3_MAX);
+  }
+}
+
+[[maybe_unused]]
+SGL_DEVICE float fp8_e4m3_clip(float val) {
+  namespace math = device::math;
+  return math::max(math::min(val, math::FP8_E4M3_MAX), -math::FP8_E4M3_MAX);
+}
+
+template <typename T, typename DST_DTYPE>
+__global__ void per_tensor_quant_fp8_kernel(
+    const T* __restrict__ input,
+    DST_DTYPE* __restrict__ output,
+    const float* __restrict__ scale,
+    const int64_t num_elements) {
+  using namespace device;
+  constexpr uint32_t VEC_SIZE = 16 / sizeof(T);
+
+  const int64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+  const float scale_val = 1.0f / (*scale);
+
+  if (gid * VEC_SIZE + VEC_SIZE <= num_elements) {
+    using input_vec_t = AlignedVector<T, VEC_SIZE>;
+    using output_vec_t = AlignedVector<DST_DTYPE, VEC_SIZE>;
+    const auto gmem_in = tile::Memory<input_vec_t>::thread();
+    const auto gmem_out = tile::Memory<output_vec_t>::thread();
+    const auto input_vec = gmem_in.load(input, gid);
+    output_vec_t output_vec;
+#pragma unroll
+    for (uint32_t i = 0; i < VEC_SIZE; ++i) {
+      const float value = fp8_e4m3_clip(static_cast<float>(input_vec[i]) * scale_val);
+      output_vec[i] = static_cast<DST_DTYPE>(value);
+    }
+    gmem_out.store(output, output_vec, gid);
+  } else if (gid * VEC_SIZE < num_elements) {
+    [[unlikely]];  // poorly aligned case, do not optimize
+    const auto remainder = num_elements - gid * VEC_SIZE;
+    for (uint32_t i = 0; i < remainder; ++i) {
+      const float value = fp8_e4m3_clip(static_cast<float>(input[gid * VEC_SIZE + i]) * scale_val);
+      output[gid * VEC_SIZE + i] = static_cast<DST_DTYPE>(value);
+    }
+  }
+}
+
+template <bool kIsStatic, typename DType>
+void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
+  using namespace host;
+
+  auto device = SymbolicDevice{};
+  auto N = SymbolicSize{"num_elements"};
+  device.set_options<kDLCUDA>();
+
+  TensorMatcher({N})  //
+      .with_dtype<DType>()
+      .with_device(device)
+      .verify(input);
+  TensorMatcher({N})  //
+      .with_dtype<fp8_e4m3_t>()
+      .with_device(device)
+      .verify(output_q);
+  TensorMatcher({1})  //
+      .with_dtype<float>()
+      .with_device(device)
+      .verify(output_s);
+
+  const auto num_elements = N.unwrap();
+
+  constexpr size_t kElementsPerBlock = kBlockSize * (16 / sizeof(DType));
+  const uint32_t num_blocks = div_ceil(num_elements, kElementsPerBlock);
+
+  if constexpr (!kIsStatic) {
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
+        per_tensor_absmax_kernel<DType>,
+        static_cast<const DType*>(input.data_ptr()),
+        static_cast<float*>(output_s.data_ptr()),
+        static_cast<int64_t>(num_elements));
+  }
+
+  LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
+      per_tensor_quant_fp8_kernel<DType, fp8_e4m3_t>,
+      static_cast<const DType*>(input.data_ptr()),
+      static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
+      static_cast<const float*>(output_s.data_ptr()),
+      static_cast<int64_t>(num_elements));
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh

@@ -0,0 +1,261 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
+
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::inv_scale_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+// Optimized per-token-group quant to FP8-e4m3 (or int8), with optional
+// column-major UE8M0 (int32-packed) scale. Memory-bound rewrite of the AOT
+// `per_token_group_quant_8bit` (sgl-kernel): the previous JIT clone read the
+// input TWICE and used 16 threads/group with only group/8 active. This version:
+//   * loads each group once into registers (single 128-bit load per thread),
+//   * uses exactly kGroupSize/kVec threads per group (no idle lanes),
+//   * sub-warp shuffle-reduces the absmax over those lanes (no shared memory),
+//   * launches warp-aligned, ~256-thread blocks for high occupancy / latency
+//     hiding (the AOT kernel's 1-warp blocks left HBM ~80% idle at prefill).
+// The UE8M0 path reuses the dsv4 cast_to_ue8m0/inv_scale_ue8m0 primitives and is
+// byte-identical to `sgl_per_token_group_quant_8bit_v2` (both ceil-round the
+// scale and store the biased exponent byte). ROCm portability comes from the
+// portable `warp::reduce_max<kThreadsPerGroup>` used directly in the kernel
+// (gfx942-safe; no separate GroupReduceMax helper needed).
+
+template <bool kUE8M0>
+using scale_packed_t_t = std::conditional_t<kUE8M0, uint32_t, float>;
+template <bool kUE8M0>
+using scale_element_t_t = std::conditional_t<kUE8M0, uint8_t, float>;
+
+struct PerTokenGroupQuantParams {
+  const void* __restrict__ input;
+  void* __restrict__ output_q;
+  void* __restrict__ output_s;
+  int64_t num_groups;      // total groups = num_tokens * num_groups_per_row
+  int num_groups_per_row;  // hidden / group_size
+  int groups_per_block;    // groups handled by one CTA
+  int scale_stride;        // output_s.stride(1), in scale_packed_t elements
+  float eps;
+  float min_8bit;
+  float max_8bit;
+};
+
+// kGroupSize columns per group; kThreadsPerGroup threads cover one group, each
+// issuing kNumVec coalesced 128-bit loads. Lane loads are interleaved
+// (chunk v at element (v*kThreadsPerGroup + lane)*kVec) so consecutive lanes hit
+// consecutive 16B addresses -> fully coalesced even for kNumVec > 1.
+template <typename T, typename DST, int64_t kGroupSize, int kThreadsPerGroup, bool kColMajor, bool kUE8M0, bool kUsePDL>
+__global__ __launch_bounds__(256, 8) void per_token_group_quant_8bit_kernel(const PerTokenGroupQuantParams params) {
+  using namespace device;
+  namespace math = device::math;
+
+  constexpr uint32_t kVec = 16u / sizeof(T);  // 8 for bf16/fp16
+  constexpr uint32_t kElemsPerThread = kGroupSize / kThreadsPerGroup;
+  constexpr uint32_t kNumVec = kElemsPerThread / kVec;  // 128-bit loads per thread
+  static_assert(kGroupSize % (kThreadsPerGroup * kVec) == 0, "bad tiling");
+  static_assert(
+      kThreadsPerGroup >= 1 && kThreadsPerGroup <= 32 && (kThreadsPerGroup & (kThreadsPerGroup - 1)) == 0,
+      "threads-per-group must be a pow2 <= 32");
+
+  using InVec = AlignedVector<T, kVec>;
+  using scale_packed_t = scale_packed_t_t<kUE8M0>;
+  using scale_element_t = scale_element_t_t<kUE8M0>;
+
+  const int local_group = threadIdx.x / kThreadsPerGroup;
+  const int lane = threadIdx.x % kThreadsPerGroup;
+  const int64_t global_group = static_cast<int64_t>(blockIdx.x) * params.groups_per_block + local_group;
+
+  PDLWaitPrimary<kUsePDL>();
+  if (global_group >= params.num_groups) {
+    PDLTriggerSecondary<kUsePDL>();
+    return;
+  }
+
+  const T* gin = static_cast<const T*>(params.input) + global_group * kGroupSize;
+  DST* gout = static_cast<DST*>(params.output_q) + global_group * kGroupSize;
+
+  // Load kNumVec interleaved 128-bit chunks into registers.
+  float vals[kElemsPerThread];
+  float local_absmax = params.eps;
+#pragma unroll
+  for (uint32_t v = 0; v < kNumVec; ++v) {
+    InVec in_vec;
+    in_vec.load(gin + (v * kThreadsPerGroup + lane) * kVec, 0);
+#pragma unroll
+    for (uint32_t j = 0; j < kVec; ++j) {
+      const float val = static_cast<float>(in_vec[j]);
+      vals[v * kVec + j] = val;
+      local_absmax = math::max(local_absmax, math::abs(val));
+    }
+  }
+  if constexpr (kThreadsPerGroup > 1) {
+    local_absmax = warp::reduce_max<kThreadsPerGroup>(local_absmax);
+  }
+
+  // Scale (byte-identical to sgl_per_token_group_quant_8bit_v2).
+  const float kMaxInv = 1.0f / params.max_8bit;
+  float inv_scale;  // multiply input by this to quantize
+  scale_element_t scale_store;
+  if constexpr (kUE8M0) {
+    const int32_t exp = cast_to_ue8m0(local_absmax * kMaxInv);
+    inv_scale = inv_scale_ue8m0(exp);
+    scale_store = static_cast<uint8_t>(exp);
+  } else {
+    const float scale_inv = local_absmax * kMaxInv;  // stored scale
+    inv_scale = params.max_8bit / local_absmax;      // quant multiplier
+    scale_store = scale_inv;
+  }
+
+  // Quantize from registers and store kNumVec interleaved chunks.
+#pragma unroll
+  for (uint32_t v = 0; v < kNumVec; ++v) {
+    DST* o = gout + (v * kThreadsPerGroup + lane) * kVec;
+    if constexpr (std::is_same_v<DST, fp8_e4m3_t>) {
+      AlignedVector<fp8x2_e4m3_t, kVec / 2> out_vec;
+#pragma unroll
+      for (uint32_t j = 0; j < kVec / 2; ++j) {
+        out_vec[j] = pack_fp8(vals[v * kVec + 2 * j] * inv_scale, vals[v * kVec + 2 * j + 1] * inv_scale);
+      }
+      out_vec.store(o, 0);
+    } else {
+      AlignedVector<DST, kVec> out_vec;
+#pragma unroll
+      for (uint32_t j = 0; j < kVec; ++j) {
+        const float q = math::min(math::max(vals[v * kVec + j] * inv_scale, params.min_8bit), params.max_8bit);
+        out_vec[j] = static_cast<DST>(q);
+      }
+      out_vec.store(o, 0);
+    }
+  }
+
+  // One scale write per group (the leading lane).
+  if (lane == 0) {
+    scale_element_t* scale_out;
+    if constexpr (kColMajor) {
+      constexpr int kPack = static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
+      const int row = static_cast<int>(global_group / params.num_groups_per_row);    // token
+      const int col_u = static_cast<int>(global_group % params.num_groups_per_row);  // group in row
+      const int col = col_u / kPack;
+      const int pack = col_u % kPack;
+      scale_out = reinterpret_cast<scale_element_t*>(params.output_s) +
+                  (static_cast<int64_t>(col) * params.scale_stride * kPack + static_cast<int64_t>(row) * kPack + pack);
+    } else {
+      static_assert(!kUE8M0, "non-column-major UE8M0 is unsupported");
+      scale_out = static_cast<scale_element_t*>(params.output_s) + global_group;
+    }
+    *scale_out = scale_store;
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// Threads cooperating on one group. Heuristic: ~8 elems/thread for small groups
+// (one 128-bit load) and 16 elems/thread for group=128 (two loads), capped at 8
+// threads/group to keep the sub-warp reduction (and register pressure) small.
+template <int64_t kGroupSize, typename T>
+constexpr int threads_per_group() {
+  constexpr int kVec = 16 / sizeof(T);                  // 8 for bf16/fp16
+  int tpg = static_cast<int>(kGroupSize) / (2 * kVec);  // ~16 elems/thread (2 vecs)
+  if (tpg > 8) tpg = 8;
+  if (tpg < 1) tpg = 1;
+  return tpg;
+}
+
+constexpr int kBlockThreads = 256;
+
+template <int64_t kGroupSize, typename T>
+inline int pick_groups_per_block() {
+  int gpb = kBlockThreads / threads_per_group<kGroupSize, T>();
+  if (gpb < 1) gpb = 1;
+  return gpb;
+}
+
+template <typename T, typename DST, int64_t kGroupSize, bool kColMajor, bool kUE8M0, bool kUsePDL>
+void launch_quant(const PerTokenGroupQuantParams& base, int64_t num_groups, int groups_per_block, DLDevice device) {
+  using namespace host;
+  constexpr int kThreadsPerGroup = threads_per_group<kGroupSize, T>();
+  const int num_threads = groups_per_block * kThreadsPerGroup;
+  const int64_t num_blocks = (num_groups + groups_per_block - 1) / groups_per_block;
+  PerTokenGroupQuantParams params = base;
+  params.groups_per_block = groups_per_block;
+  constexpr auto kernel =
+      per_token_group_quant_8bit_kernel<T, DST, kGroupSize, kThreadsPerGroup, kColMajor, kUE8M0, kUsePDL>;
+  LaunchKernel(static_cast<uint32_t>(num_blocks), static_cast<uint32_t>(num_threads), device)
+      .enable_pdl(kUsePDL)(kernel, params);
+}
+
+template <typename DType, typename OutType, int64_t kGroupSize, bool kUsePDL>
+void per_token_group_quant_8bit(
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView output_q,
+    tvm::ffi::TensorView output_s,
+    int64_t group_size,
+    double eps,
+    double min_8bit,
+    double max_8bit,
+    bool scale_ue8m0) {
+  using namespace host;
+  static_assert(
+      kGroupSize == 16 || kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128,
+      "group_size template arg must be 16/32/64/128");
+
+  auto device = SymbolicDevice{};
+  auto M = SymbolicSize{"num_tokens"};
+  auto K = SymbolicSize{"hidden_dim"};
+  device.set_options<kDLCUDA>();
+
+  TensorMatcher({M, K}).with_dtype<DType>().with_device(device).verify(input);
+  TensorMatcher({M, K}).with_dtype<OutType>().with_device(device).verify(output_q);
+
+  RuntimeCheck(group_size == kGroupSize, "group_size does not match compiled template");
+
+  const int64_t num_tokens = M.unwrap();
+  const int64_t hidden_dim = K.unwrap();
+  const int64_t num_groups_per_row = hidden_dim / kGroupSize;
+  const int64_t num_groups = num_tokens * num_groups_per_row;
+  if (num_groups == 0) return;
+
+  const bool is_column_major = output_s.stride(0) < output_s.stride(1);
+  const int scale_stride = static_cast<int>(output_s.stride(1));
+
+  PerTokenGroupQuantParams base{};
+  base.input = input.data_ptr();
+  base.output_q = output_q.data_ptr();
+  base.output_s = output_s.data_ptr();
+  base.num_groups = num_groups;
+  base.num_groups_per_row = static_cast<int>(num_groups_per_row);
+  base.scale_stride = scale_stride;
+  base.eps = static_cast<float>(eps);
+  base.min_8bit = static_cast<float>(min_8bit);
+  base.max_8bit = static_cast<float>(max_8bit);
+
+  const auto dev = input.device();
+  const int gpb = pick_groups_per_block<kGroupSize, DType>();
+
+  // Runtime selection between compile-time-instantiated scale-layout variants
+  // (the group size itself is a template arg, supplied from Python).
+  if (is_column_major) {
+    if (scale_ue8m0) {
+      launch_quant<DType, OutType, kGroupSize, true, true, kUsePDL>(base, num_groups, gpb, dev);
+    } else {
+      launch_quant<DType, OutType, kGroupSize, true, false, kUsePDL>(base, num_groups, gpb, dev);
+    }
+  } else {
+    RuntimeCheck(!scale_ue8m0, "row-major UE8M0 unsupported");
+    launch_quant<DType, OutType, kGroupSize, false, false, kUsePDL>(base, num_groups, gpb, dev);
+  }
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh

@@ -0,0 +1,539 @@
+// JIT port of the AOT sgl_per_token_group_quant_8bit_v2 (sgl-kernel).
+//
+// Same math as the AOT v2 kernel (256-bit vectorized loads, 8 threads/128-group,
+// PDL, NaiveScheduler + MaskedLayoutScheduler, ue8m0/float scales, fp8/int8
+// output, fused silu+mul) so it is a drop-in replacement; the only changes vs the
+// AOT source are the launcher (tvm::ffi::TensorView + TensorMatcher + the JIT
+// LaunchKernel/PDL helpers) and the FP8 type alias.
+#include <sgl_kernel/tensor.h>  // TensorMatcher, SymbolicSize/Device
+#include <sgl_kernel/utils.h>   // RuntimeCheck, Panic
+
+#include <sgl_kernel/utils.cuh>  // LaunchKernel, fp8_e4m3_t, SGL_DEVICE, device::PDLWaitPrimary/TriggerSecondary
+#include <sgl_kernel/warp.cuh>   // device::warp::reduce_max
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_fp8.h>
+#include <type_traits>
+
+namespace {
+
+constexpr float LOCAL_ABSMAX_ABS = 1e-10f;
+constexpr uint32_t INPUT_PRIMARY_VEC_NUM_BYTES = 32;
+
+template <int THREADS_PER_SUBWARP>
+SGL_DEVICE float GroupReduceMax(float val) {
+  static_assert(
+      (THREADS_PER_SUBWARP & (THREADS_PER_SUBWARP - 1)) == 0 && THREADS_PER_SUBWARP <= 16 && THREADS_PER_SUBWARP >= 1,
+      "THREADS_PER_SUBWARP must be 1, 2, 4, 8, or 16");
+  // Reduce within this thread's contiguous THREADS_PER_SUBWARP-lane subgroup via
+  // the shared warp primitive, but pass an explicit subgroup mask instead of its
+  // default 0xffffffff: the block can be < 32 lanes (subwarps_per_block *
+  // THREADS_PER_SUBWARP, e.g. 1..16), where a full-warp mask names non-existent
+  // lanes and is UB / can hang.
+  constexpr device::warp::mask_t kSub = (device::warp::mask_t{1} << THREADS_PER_SUBWARP) - 1;
+  const device::warp::mask_t mask = kSub << (THREADS_PER_SUBWARP * ((threadIdx.x % 32) / THREADS_PER_SUBWARP));
+  return device::warp::reduce_max<THREADS_PER_SUBWARP>(val, mask);
+}
+
+SGL_DEVICE float silu(const float& val) {
+  // Match the AOT v2 kernel: tanh-based silu on SM100+ (Blackwell), exp-based
+  // elsewhere, so the fused silu+mul output stays bit-identical to the AOT op.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  float half = 0.5f * val;
+  float t = __tanhf(half);
+  return half * (1.0f + t);
+#else
+  return val / (1.0f + __expf(-val));
+#endif
+}
+
+SGL_DEVICE float2 fmul2_rn(float2 a, float2 b) {
+  // Match the AOT v2 kernel: use the __fmul2_rn intrinsic on SM100+.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  return __fmul2_rn(a, b);
+#else
+  float2 result;
+  result.x = a.x * b.x;
+  result.y = a.y * b.y;
+  return result;
+#endif
+}
+
+// Copied from DeepEP.
+SGL_DEVICE float fast_pow2(int x) {
+  uint32_t bits_x = (x + 127) << 23;
+  return __uint_as_float(bits_x);  // type-safe bit cast (no strict-aliasing UB)
+}
+SGL_DEVICE int fast_log2_ceil(float x) {
+  auto bits_x = __float_as_uint(x);  // type-safe bit cast (no strict-aliasing UB)
+  auto exp_x = (bits_x >> 23) & 0xff;
+  auto man_bits = bits_x & ((1 << 23) - 1);
+  return exp_x - 127 + (man_bits != 0);
+}
+
+template <typename T>
+struct DtypeInfo;
+template <>
+struct DtypeInfo<int8_t> {
+  static constexpr float MIN = -128;
+  static constexpr float MAX = 127;
+};
+template <>
+struct DtypeInfo<fp8_e4m3_t> {
+  static constexpr float MIN = -448;
+  static constexpr float MAX = 448;
+};
+
+template <bool ROUND_SCALE, typename dtype_info>
+SGL_DEVICE void calculate_fp8_scales(float amax, float& scale, float& scale_inv) {
+  constexpr float MAX_8BIT_INV = 1.0f / dtype_info::MAX;
+  if constexpr (ROUND_SCALE) {
+    auto exp_scale_inv = fast_log2_ceil(amax * MAX_8BIT_INV);
+    scale = fast_pow2(-exp_scale_inv);
+    scale_inv = fast_pow2(exp_scale_inv);
+  } else {
+    scale_inv = amax * MAX_8BIT_INV;
+    scale = dtype_info::MAX / amax;
+  }
+}
+
+template <bool SCALE_UE8M0, typename OUT_DTYPE_T = std::conditional_t<SCALE_UE8M0, uint8_t, float>>
+SGL_DEVICE OUT_DTYPE_T extract_required_scale_format(float value) {
+  if constexpr (SCALE_UE8M0) {
+    return static_cast<uint8_t>(__float_as_uint(value) >> 23);
+  } else {
+    return value;
+  }
+}
+
+template <bool FUSE_SILU_AND_MUL>
+SGL_DEVICE int compute_input_group_start_offset(
+    int expert_idx,
+    int token_idx,
+    int hidden_dim_group_idx,
+    int hidden_size,
+    int num_tokens_per_expert,
+    int group_size) {
+  return expert_idx * num_tokens_per_expert * hidden_size * (FUSE_SILU_AND_MUL ? 2 : 1) +
+         token_idx * hidden_size * (FUSE_SILU_AND_MUL ? 2 : 1) + hidden_dim_group_idx * group_size;
+}
+
+struct NaiveScheduler {
+  static void compute_exec_config(
+      int threads_per_subwarp,
+      int num_local_experts,
+      int hidden_dim_num_groups,
+      int num_groups,
+      int& subwarps_per_block,
+      dim3& grid,
+      dim3& block) {
+    subwarps_per_block = (num_groups % 16 == 0)  ? 16
+                         : (num_groups % 8 == 0) ? 8
+                         : (num_groups % 4 == 0) ? 4
+                         : (num_groups % 2 == 0) ? 2
+                                                 : 1;
+    grid = dim3(num_groups / subwarps_per_block);
+    block = dim3(subwarps_per_block * threads_per_subwarp);
+  }
+
+  template <bool FUSE_SILU_AND_MUL, int GROUP_SIZE, int THREADS_PER_SUBWARP, typename FUNC>
+  SGL_DEVICE static void execute(
+      const int subwarps_per_block,
+      const int hidden_dim_num_groups,
+      const int32_t* masked_m,
+      const int num_tokens_per_expert,
+      FUNC fn) {
+    constexpr int expert_idx = 0;
+    const int64_t subwarp_id = threadIdx.x / THREADS_PER_SUBWARP;
+    const int lane_id = threadIdx.x % THREADS_PER_SUBWARP;
+    const int64_t group_id = static_cast<int64_t>(blockIdx.x) * subwarps_per_block + subwarp_id;
+
+    int64_t input_group_start_offset;
+    if constexpr (!FUSE_SILU_AND_MUL) input_group_start_offset = group_id * GROUP_SIZE;
+    const int token_idx = group_id / hidden_dim_num_groups;
+    const int hidden_dim_group_idx = group_id % hidden_dim_num_groups;
+    if constexpr (FUSE_SILU_AND_MUL) {
+      const int hidden_size = hidden_dim_num_groups * GROUP_SIZE;
+      input_group_start_offset = compute_input_group_start_offset<FUSE_SILU_AND_MUL>(
+          expert_idx, token_idx, hidden_dim_group_idx, hidden_size, num_tokens_per_expert, GROUP_SIZE);
+    }
+    fn(expert_idx, token_idx, hidden_dim_group_idx, lane_id, input_group_start_offset);
+  }
+};
+
+struct MaskedLayoutScheduler {
+  static constexpr int TOKEN_DIM_BLOCK_NUM_PER_EXPERT = 1024;
+  static constexpr int SUBWARPS_PER_BLOCK = 16;
+
+  static void compute_exec_config(
+      int threads_per_subwarp,
+      int num_local_experts,
+      int hidden_dim_num_groups,
+      int num_groups,
+      int& subwarps_per_block,
+      dim3& grid,
+      dim3& block) {
+    subwarps_per_block = SUBWARPS_PER_BLOCK;
+    host::RuntimeCheck(hidden_dim_num_groups % subwarps_per_block == 0, "hidden_dim_num_groups not divisible by 16");
+    grid = dim3(hidden_dim_num_groups / subwarps_per_block, TOKEN_DIM_BLOCK_NUM_PER_EXPERT, num_local_experts);
+    block = dim3(subwarps_per_block * threads_per_subwarp);
+  }
+
+  template <bool FUSE_SILU_AND_MUL, int GROUP_SIZE, int THREADS_PER_SUBWARP, typename FUNC>
+  SGL_DEVICE static void execute(
+      const int subwarps_per_block,
+      const int hidden_dim_num_groups,
+      const int32_t* masked_m,
+      const int num_tokens_per_expert,
+      FUNC fn) {
+    const int64_t subwarp_id = threadIdx.x / THREADS_PER_SUBWARP;
+    const int lane_id = threadIdx.x % THREADS_PER_SUBWARP;
+    const int expert_idx = blockIdx.z;
+    const int token_idx_start = blockIdx.y;
+    const int64_t hidden_dim_group_idx = static_cast<int64_t>(blockIdx.x) * SUBWARPS_PER_BLOCK + subwarp_id;
+    const int curr_expert_token_num = masked_m[expert_idx];
+    for (int token_idx = token_idx_start; token_idx < curr_expert_token_num;
+         token_idx += TOKEN_DIM_BLOCK_NUM_PER_EXPERT) {
+      const int hidden_size = hidden_dim_num_groups * GROUP_SIZE;
+      const int64_t input_group_start_offset = compute_input_group_start_offset<FUSE_SILU_AND_MUL>(
+          expert_idx, token_idx, hidden_dim_group_idx, hidden_size, num_tokens_per_expert, GROUP_SIZE);
+      fn(expert_idx, token_idx, hidden_dim_group_idx, lane_id, input_group_start_offset);
+    }
+  }
+};
+
+template <
+    typename SCHEDULER,
+    int GROUP_SIZE,
+    int THREADS_PER_SUBWARP,
+    typename T,
+    typename DST_DTYPE,
+    bool IS_COLUMN_MAJOR,
+    bool SCALE_UE8M0,
+    bool FUSE_SILU_AND_MUL,
+    bool kUsePDL,
+    typename scale_packed_t = std::conditional_t<SCALE_UE8M0 && IS_COLUMN_MAJOR, uint32_t, float>>
+__global__ void per_token_group_quant_8bit_v2_kernel(
+    const T* __restrict__ input,
+    DST_DTYPE* __restrict__ output_q,
+    scale_packed_t* __restrict__ output_s,
+    const int32_t* __restrict__ masked_m,
+    const int subwarps_per_block,
+    const int hidden_dim_num_groups,
+    const int scale_expert_stride,
+    const int scale_hidden_stride,
+    const int num_tokens_per_expert) {
+  using dst_dtype_info = DtypeInfo<DST_DTYPE>;
+  using scale_element_t = std::conditional_t<SCALE_UE8M0 && IS_COLUMN_MAJOR, uint8_t, float>;
+  static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
+
+  device::PDLWaitPrimary<kUsePDL>();
+
+  SCHEDULER::template execute<FUSE_SILU_AND_MUL, GROUP_SIZE, THREADS_PER_SUBWARP>(
+      subwarps_per_block,
+      hidden_dim_num_groups,
+      masked_m,
+      num_tokens_per_expert,
+      [&](const int expert_idx,
+          const int token_idx,
+          const int hidden_dim_group_idx,
+          const int lane_id,
+          const int input_group_start_offset) {
+        constexpr uint32_t INPUT_PRIMARY_VEC_SIZE = INPUT_PRIMARY_VEC_NUM_BYTES / sizeof(T);
+        constexpr uint32_t INPUT_PRIMARY_INT4_SIZE = INPUT_PRIMARY_VEC_NUM_BYTES / sizeof(int4);
+
+        const int offset_num_groups = expert_idx * num_tokens_per_expert * hidden_dim_num_groups +
+                                      token_idx * hidden_dim_num_groups + hidden_dim_group_idx;
+
+        int4 input_primary_int4[INPUT_PRIMARY_INT4_SIZE];
+        T* input_primary_vec = reinterpret_cast<T*>(input_primary_int4);
+        int4 input_secondary_int4[INPUT_PRIMARY_INT4_SIZE];
+        T* input_secondary_vec = reinterpret_cast<T*>(input_secondary_int4);
+
+#pragma unroll
+        for (uint32_t j = 0; j < INPUT_PRIMARY_INT4_SIZE; ++j) {
+          // Ordinary 128-bit vectorized load (LDG.128); .nc gave no measurable
+          // gain on this streaming read-once kernel, so no inline asm.
+          input_primary_int4[j] =
+              reinterpret_cast<const int4*>(input + input_group_start_offset + lane_id * INPUT_PRIMARY_VEC_SIZE)[j];
+        }
+        if constexpr (FUSE_SILU_AND_MUL) {
+          const int secondary_offset = hidden_dim_num_groups * GROUP_SIZE;
+#pragma unroll
+          for (uint32_t j = 0; j < INPUT_PRIMARY_INT4_SIZE; ++j) {
+            input_secondary_int4[j] = reinterpret_cast<const int4*>(
+                input + input_group_start_offset + lane_id * INPUT_PRIMARY_VEC_SIZE + secondary_offset)[j];
+          }
+        }
+
+        constexpr int num_elems_per_pack = static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
+        scale_element_t* scale_output;
+        if constexpr (IS_COLUMN_MAJOR) {
+          constexpr int scale_token_stride = 1;
+          const int hidden_idx_packed = hidden_dim_group_idx / num_elems_per_pack;
+          const int pack_idx = hidden_dim_group_idx % num_elems_per_pack;
+          scale_output = reinterpret_cast<scale_element_t*>(output_s) +
+                         (expert_idx * scale_expert_stride * num_elems_per_pack +
+                          hidden_idx_packed * scale_hidden_stride * num_elems_per_pack +
+                          token_idx * scale_token_stride * num_elems_per_pack + pack_idx);
+        } else {
+          static_assert(!SCALE_UE8M0 || std::is_same_v<scale_packed_t, float>);
+          scale_output = reinterpret_cast<scale_element_t*>(output_s) + offset_num_groups;
+        }
+
+        if constexpr (IS_COLUMN_MAJOR and SCALE_UE8M0) {
+          const int remainder_num_groups = hidden_dim_num_groups % num_elems_per_pack;
+          if ((remainder_num_groups != 0) and (hidden_dim_group_idx == hidden_dim_num_groups - 1) and
+              (lane_id < num_elems_per_pack - remainder_num_groups)) {
+            const int shift = 1 + lane_id;
+            *(scale_output + shift) = 0;
+          }
+        }
+
+        float local_absmax = LOCAL_ABSMAX_ABS;
+#pragma unroll
+        for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) {
+          float val;
+          if constexpr (FUSE_SILU_AND_MUL) {
+            T val_lowprec = static_cast<T>(silu(static_cast<float>(input_primary_vec[j]))) * input_secondary_vec[j];
+            val = static_cast<float>(val_lowprec);
+            input_primary_vec[j] = val_lowprec;
+          } else {
+            val = static_cast<float>(input_primary_vec[j]);
+          }
+          local_absmax = fmaxf(local_absmax, fabsf(val));
+        }
+
+        local_absmax = GroupReduceMax<THREADS_PER_SUBWARP>(local_absmax);
+
+        float y_scale, y_scale_inv;
+        // When SCALE_UE8M0, always quantize with the rounded (power-of-2) scale
+        // — not with the exact scale followed by post-hoc rounding.
+        // This matches the official DeepSeek-V4 kernel.py act_quant(scale_fmt="ue8m0")
+        // and avoids a scale mismatch between quantization and downstream GEMM dequant,
+        // which otherwise amplifies error ~14x and degrades EAGLE accept rate on Blackwell.
+        calculate_fp8_scales<SCALE_UE8M0, dst_dtype_info>(local_absmax, y_scale, y_scale_inv);
+        if (lane_id == 0) {
+          *scale_output = extract_required_scale_format < SCALE_UE8M0 && IS_COLUMN_MAJOR > (y_scale_inv);
+        }
+        float2 y_scale_repeated = {y_scale, y_scale};
+
+        int4 output_buf;
+        if constexpr (std::is_same_v<DST_DTYPE, fp8_e4m3_t>) {
+          const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf);
+#pragma unroll
+          for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 2) {
+            float2 inputx2 = {static_cast<float>(input_primary_vec[j]), static_cast<float>(input_primary_vec[j + 1])};
+            float2 outputx2 = fmul2_rn(inputx2, y_scale_repeated);
+            outputx2.x = fminf(fmaxf(outputx2.x, dst_dtype_info::MIN), dst_dtype_info::MAX);
+            outputx2.y = fminf(fmaxf(outputx2.y, dst_dtype_info::MIN), dst_dtype_info::MAX);
+            output_buf_ptr[j / 2] = __nv_cvt_float2_to_fp8x2(outputx2, __NV_SATFINITE, __NV_E4M3);
+          }
+        } else {
+          const auto output_buf_ptr = reinterpret_cast<DST_DTYPE*>(&output_buf);
+#pragma unroll
+          for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) {
+            float val = static_cast<float>(input_primary_vec[j]);
+            float q_val = fminf(fmaxf(val * y_scale, dst_dtype_info::MIN), dst_dtype_info::MAX);
+            output_buf_ptr[j] = DST_DTYPE(q_val);
+          }
+        }
+
+        // Ordinary 128-bit vectorized store (STG.128); no inline asm.
+        *reinterpret_cast<int4*>(output_q + offset_num_groups * GROUP_SIZE + lane_id * INPUT_PRIMARY_VEC_SIZE) =
+            output_buf;
+      });
+
+  device::PDLTriggerSecondary<kUsePDL>();
+}
+
+// ----------------------------------------------------------------------------
+// Launcher (JIT). All shape-derived scalars are computed in the Python wrapper
+// and passed in, so the C++ side only needs TensorView::data_ptr()/device().
+// Runtime combos (column_major / ue8m0 / silu / masked / group_size) are
+// dispatched to the templated kernel; launch is PDL-aware via LaunchKernel.
+// ----------------------------------------------------------------------------
+template <typename S>
+struct TypeTag {
+  using type = S;
+};
+
+template <typename T, typename DST_DTYPE, bool kUsePDL>
+struct PerTokenGroupQuant8bitV2Kernel {
+  template <
+      typename SCHEDULER,
+      int GROUP_SIZE,
+      int THREADS_PER_SUBWARP,
+      bool IS_COLUMN_MAJOR,
+      bool SCALE_UE8M0,
+      bool FUSE_SILU_AND_MUL>
+  static void launch(
+      const DLDevice& device,
+      dim3 grid,
+      dim3 block,
+      int subwarps_per_block,
+      int hidden_dim_num_groups,
+      int scale_expert_stride,
+      int scale_hidden_stride,
+      int num_tokens_per_expert,
+      const void* input,
+      void* output_q,
+      void* output_s,
+      const int32_t* masked_m) {
+    using scale_packed_t = std::conditional_t<SCALE_UE8M0 && IS_COLUMN_MAJOR, uint32_t, float>;
+    auto kernel = per_token_group_quant_8bit_v2_kernel<
+        SCHEDULER,
+        GROUP_SIZE,
+        THREADS_PER_SUBWARP,
+        T,
+        DST_DTYPE,
+        IS_COLUMN_MAJOR,
+        SCALE_UE8M0,
+        FUSE_SILU_AND_MUL,
+        kUsePDL>;
+    host::LaunchKernel(grid, block, device)
+        .enable_pdl(kUsePDL)(
+            kernel,
+            static_cast<const T*>(input),
+            static_cast<DST_DTYPE*>(output_q),
+            static_cast<scale_packed_t*>(output_s),
+            masked_m,
+            subwarps_per_block,
+            hidden_dim_num_groups,
+            scale_expert_stride,
+            scale_hidden_stride,
+            num_tokens_per_expert);
+  }
+
+  template <int GROUP_SIZE>
+  static void dispatch_bools(
+      const DLDevice& device,
+      bool is_column_major,
+      bool scale_ue8m0,
+      bool fuse_silu_and_mul,
+      bool masked_layout,
+      int num_local_experts,
+      int hidden_dim_num_groups,
+      int num_groups,
+      int scale_expert_stride,
+      int scale_hidden_stride,
+      int num_tokens_per_expert,
+      const void* input,
+      void* output_q,
+      void* output_s,
+      const int32_t* masked_m) {
+    constexpr int THREADS_PER_SUBWARP = GROUP_SIZE / 16;
+
+    auto launch_with_config = [&](auto sched_tag, auto colmajor_tag, auto ue8m0_tag, auto silu_tag) {
+      using SCHEDULER = typename decltype(sched_tag)::type;
+      int subwarps_per_block;
+      dim3 grid, block;
+      SCHEDULER::compute_exec_config(
+          THREADS_PER_SUBWARP, num_local_experts, hidden_dim_num_groups, num_groups, subwarps_per_block, grid, block);
+      launch<
+          SCHEDULER,
+          GROUP_SIZE,
+          THREADS_PER_SUBWARP,
+          decltype(colmajor_tag)::value,
+          decltype(ue8m0_tag)::value,
+          decltype(silu_tag)::value>(
+          device,
+          grid,
+          block,
+          subwarps_per_block,
+          hidden_dim_num_groups,
+          scale_expert_stride,
+          scale_hidden_stride,
+          num_tokens_per_expert,
+          input,
+          output_q,
+          output_s,
+          masked_m);
+    };
+
+    if (is_column_major) {
+      if (scale_ue8m0) {
+        if (fuse_silu_and_mul) {
+          if (masked_layout)
+            launch_with_config(TypeTag<MaskedLayoutScheduler>{}, std::true_type{}, std::true_type{}, std::true_type{});
+          else
+            launch_with_config(TypeTag<NaiveScheduler>{}, std::true_type{}, std::true_type{}, std::true_type{});
+        } else {
+          launch_with_config(TypeTag<NaiveScheduler>{}, std::true_type{}, std::true_type{}, std::false_type{});
+        }
+      } else {
+        launch_with_config(TypeTag<NaiveScheduler>{}, std::true_type{}, std::false_type{}, std::false_type{});
+      }
+    } else {
+      if (scale_ue8m0) {
+        launch_with_config(TypeTag<NaiveScheduler>{}, std::false_type{}, std::true_type{}, std::false_type{});
+      } else {
+        launch_with_config(TypeTag<NaiveScheduler>{}, std::false_type{}, std::false_type{}, std::false_type{});
+      }
+    }
+  }
+
+  static void
+  run(tvm::ffi::TensorView input,
+      tvm::ffi::TensorView output_q,
+      tvm::ffi::TensorView output_s,
+      tvm::ffi::TensorView masked_m,
+      int64_t group_size,
+      bool scale_ue8m0,
+      bool fuse_silu_and_mul,
+      bool masked_layout,
+      int64_t num_groups,
+      int64_t num_local_experts,
+      bool is_column_major,
+      int64_t hidden_dim_num_groups,
+      int64_t num_tokens_per_expert,
+      int64_t scale_expert_stride,
+      int64_t scale_hidden_stride) {
+    const DLDevice dev = input.device();
+    const void* in = input.data_ptr();
+    void* oq = output_q.data_ptr();
+    void* os = output_s.data_ptr();
+    const int32_t* masked_ptr = masked_layout ? static_cast<const int32_t*>(masked_m.data_ptr()) : nullptr;
+
+    auto dispatch_gs = [&](auto gs_tag) {
+      constexpr int GS = decltype(gs_tag)::value;
+      static_assert((GS / 16) * INPUT_PRIMARY_VEC_NUM_BYTES == GS * static_cast<int>(sizeof(T)));
+      dispatch_bools<GS>(
+          dev,
+          is_column_major,
+          scale_ue8m0,
+          fuse_silu_and_mul,
+          masked_layout,
+          static_cast<int>(num_local_experts),
+          static_cast<int>(hidden_dim_num_groups),
+          static_cast<int>(num_groups),
+          static_cast<int>(scale_expert_stride),
+          static_cast<int>(scale_hidden_stride),
+          static_cast<int>(num_tokens_per_expert),
+          in,
+          oq,
+          os,
+          masked_ptr);
+    };
+    switch (group_size) {
+      case 16:
+        dispatch_gs(std::integral_constant<int, 16>{});
+        break;
+      case 32:
+        dispatch_gs(std::integral_constant<int, 32>{});
+        break;
+      case 64:
+        dispatch_gs(std::integral_constant<int, 64>{});
+        break;
+      case 128:
+        dispatch_gs(std::integral_constant<int, 128>{});
+        break;
+      default:
+        host::Panic("Unsupported group_size ", group_size);
+    }
+  }
+};
+
+}  // namespace