ref(upstream): add Deep-Spark xllm + vllm MoE/GDN reference code

Sources (Apache 2.0, cloned 2026-08-09):
- Deep-Spark/xllm: Iluvatar's official C++ inference engine
- Deep-Spark/vllm: Iluvatar's vllm fork

Key files for our EX Engine development:

MoE topk_softmax (fixes 2304 calls/token PyTorch fallback):
- xllm/kernels/cuda/moe/moe_topk_softmax_kernels.cuh
  CUB-based fused softmax+topk, power-of-2 expert count optimized
  For 64 experts: topk_gating_softmax<T,VPT=2,64,WARPS=4,BYTES=4>
- xllm/kernels/ilu/ixformer.h
  Official ixformer C++ API: topk_softmax(), paged_attention(), etc.
- xllm/kernels/ilu/fused_moe.cpp
  How xllm calls ixformer::infer::topk_softmax()
- ds_vllm/csrc/moe/topk_softmax_kernels.cu
  vllm-native topk_softmax (TensorRT-LLM derived, 874 lines)

GatedDeltaNet (fixes NaN in 4 GDN layers):
- xllm/layers/npu_torch/qwen3_gated_delta_net_base.cpp
  fp32 state accumulation, proper recurrent update

Complete FusedMoE pipeline reference:
- xllm/layers/ilu/fused_moe.cpp
  gate -> topk -> expand -> gemm1 -> act -> gemm2 -> combine

upstream_ref/README.md

@@ -0,0 +1,84 @@
+# Upstream Reference: Deep-Spark xllm + vllm
+
+Source repos (cloned 2026-08-09):
+- `Deep-Spark/xllm` — Iluvatar's C++ inference engine (Apache 2.0)
+- `Deep-Spark/vllm` — Iluvatar's vllm fork (Apache 2.0)
+
+## Call Chain: MoE topk_softmax on BI-V100
+
+```
+Our code                           xllm reference                    Iluvatar SDK
+─────────────────────────────────  ──────────────────────────────    ──────────────
+qwen3_5.py                                                          
+  Qwen3_5MoeSparseBlock.forward()                                  
+    _custom_ops.py:topk_softmax()                                   
+      ixf_F.vllm_moe_topk_softmax  ← MISSING in base image         
+      │                                                              
+      ├── xllm path (C++ native):                                  
+      │   kernels/ilu/fused_moe.cpp                                 
+      │     → ixformer::infer::topk_softmax()                      ← ixformer.h
+      │       → CUDA kernel (moe_topk_softmax_kernels.cuh)          
+      │         topk_gating_softmax<T,VPT,64,4,BYTES_PER_LDG>()    
+      │                                                              
+      ├── ds_vllm path (Python torch extension):                    
+      │   csrc/moe/topk_softmax_kernels.cu                         
+      │     → torch.ops._moe_C.topk_softmax()                      
+      │       → topk_gating_softmax_kernel_launcher<T>()            
+      │                                                              
+      └── Our EX Engine path (dlopen .so):                          
+          ex_engine/csrc/factor_moe_topk_softmax.cu                 
+            → ex_factor_0.so via ctypes                              
+            → moe_topk_softmax_kernel()                             
+```
+
+## Call Chain: GatedDeltaNet (GDN) on BI-V100
+
+```
+Our code                           xllm reference
+─────────────────────────────────  ──────────────────────────────
+qwen3_5.py                                                       
+  GatedDeltaNet.forward()                                        
+    prefill path:                                                
+      _torch_chunk_gated_delta_rule  ← produces NaN (fp16 overflow)
+      │                                                           
+      ├── xllm path:                                             
+      │   layers/npu_torch/qwen3_gated_delta_net_base.cpp        
+      │     → process_mixed_qkv() + recurrent state update       
+      │     → full fp32 accumulation                              
+      │                                                           
+      └── Our EX Engine path:                                    
+          ex_engine/csrc/factor_gdn_chunk_fwd.cu                 
+            → fp32 state accumulation, tile-based                
+```
+
+## File Index
+
+### xllm/kernels/cuda/moe/ — CUDA kernels (the actual GPU code)
+- `moe_topk_softmax_kernels.cuh` — **KEY**: fused softmax+topk, CUB-based, power-of-2 expert count optimized
+- `moe_fused_topk.cu` — sigmoid/softmax topk dispatcher  
+- `moe_topk.cuh` — topk helper functions
+- `moe_topk_sigmoid_kernels.cuh` — sigmoid variant for DeepSeek-style routing
+
+### xllm/kernels/ilu/ — Iluvatar ixformer API wrappers
+- `ixformer.h` — **KEY**: official ixformer C++ API declarations (topk_softmax, paged_attention, etc.)
+- `fused_moe.cpp` — how xllm calls ixformer::infer::topk_softmax()
+- `activation.cpp` — silu_and_mul, gelu_and_mul wrappers
+- `attention.cpp` — paged_attention wrappers
+- `norm.cpp` — rms_norm, fused_add_rms_norm wrappers
+- `rope.cpp` — rotary embedding wrappers
+
+### xllm/layers/ilu/ — Complete FusedMoE layer for Iluvatar
+- `fused_moe.cpp` — **KEY**: full MoE pipeline: gate → topk → expand → gemm1 → act → gemm2 → combine
+- `fused_moe.h` — layer interface
+
+### xllm/layers/npu_torch/ — GatedDeltaNet implementation
+- `qwen3_gated_delta_net_base.cpp` — base GDN with fp32 state management
+- `qwen3_5_gated_delta_net.cpp` — Qwen3.5 specific GDN
+
+### ds_vllm/csrc/moe/ — vllm-native MoE CUDA kernels
+- `topk_softmax_kernels.cu` — vllm's topk_softmax (TensorRT-LLM derived)
+- `moeTopKFuncs.cuh` — shared topk reduction primitives
+- `moe_align_sum_kernels.cu` — block alignment for scatter
+
+### ds_vllm/vllm/ — Python layer
+- `_custom_ops.py` — how vllm calls torch.ops._moe_C.topk_softmax

upstream_ref/ds_vllm/csrc/moe/moeTopKFuncs.cuh

@@ -0,0 +1,257 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
+ * Copyright (c) 2026, The vLLM team.
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights
+ * reserved. SPDX-License-Identifier: Apache-2.0
+ *
+ * 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 <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <cub/cub.cuh>
+
+namespace vllm {
+namespace moe {
+namespace reduce_topk {
+namespace cg = cooperative_groups;
+static constexpr int kWARP_SIZE = 32;
+
+template <typename T_>
+struct TopKRedType {
+  using T = T_;
+  static_assert(
+      std::is_same_v<T, float> || std::is_same_v<T, half> ||
+          std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
+      "Top K reduction only implemented for int, float, float16 and bfloat16");
+
+  using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
+  using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
+
+  static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
+  static constexpr int kMaxIdx = 65535;
+  TypeCmp compValIdx;
+
+  static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
+    auto valueBits = cub::Traits<T>::TwiddleIn(
+        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
+    TypeCmp compactTmp = valueBits;
+    compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
+    // Use 65535 minus idx to give higher priority to elements with smaller
+    // indices.
+    return compactTmp;
+  }
+
+  static __host__ __device__ void unpack(T& value, int32_t& index,
+                                         TypeCmp cmp) {
+    // Since “65535-idx” is always smaller than 65536 and positive, we can
+    // directly use it as the lower 16 bits
+    index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
+
+    auto compactTmp = cmp >> kMoveBits;
+    auto valueBits = cub::Traits<T>::TwiddleOut(
+        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
+    value = reinterpret_cast<T&>(valueBits);
+  }
+
+  __host__ __device__ TopKRedType() = default;
+
+  __host__ __device__ TopKRedType(T val, int32_t idx)
+      : compValIdx(makeCmpVal(val, idx)) {}
+
+  __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
+
+  __device__ inline TypeCmp reduce(
+      cg::thread_block_tile<kWARP_SIZE> const& warp) {
+    return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int K_, bool Enable_>
+struct TopKIdx {
+  // by default, empty
+};
+
+template <int K_>
+struct TopKIdx<K_, true> {
+  static constexpr int K = K_;
+  int32_t val[K];
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define TOPK_SWAP(I, J)                                         \
+  {                                                             \
+    auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
+    auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
+    topK[I].compValIdx = pairMax;                               \
+    topK[J].compValIdx = pairMin;                               \
+  }
+
+template <int N, typename RedType>
+struct Sort;
+
+template <typename RedType>
+struct Sort<1, RedType> {
+  static __device__ void run(RedType* topK) {}
+};
+
+template <typename RedType>
+struct Sort<2, RedType> {
+  static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
+};
+
+template <typename RedType>
+struct Sort<3, RedType> {
+  static __device__ void run(RedType* topK) {
+    TOPK_SWAP(0, 1);
+    TOPK_SWAP(1, 2);
+    TOPK_SWAP(0, 1);
+  }
+};
+
+template <typename RedType>
+struct Sort<4, RedType> {
+  static __device__ void run(RedType* topK) {
+    TOPK_SWAP(0, 2);
+    TOPK_SWAP(1, 3);
+    TOPK_SWAP(0, 1);
+    TOPK_SWAP(2, 3);
+    TOPK_SWAP(1, 2);
+  }
+};
+
+template <int K, typename Type>
+__forceinline__ __device__ void reduceTopK(
+    cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
+    int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue,
+    int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  using RedType = TopKRedType<Type>;
+  RedType topK{value, idx};
+  typename RedType::TypeCmp packedMax{};
+#pragma unroll
+  for (int kk = 0; kk < actualK; ++kk) {
+    topK =
+        kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
+    // get the next largest value
+    packedMax = topK.reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
+  }
+};
+
+template <int K, typename Type, int N, bool IsSorted = false>
+__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
+                               Type (&out)[K], int32_t (&outIdx)[K],
+                               Type (&value)[N], int32_t (&idx)[N],
+                               Type minValue, int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(N > 0, "Top K must have N > 0");
+  static_assert(N < 5,
+                "Only support candidates number less than or equal to 128");
+  using RedType = TopKRedType<Type>;
+  RedType topK[N];
+#pragma unroll
+  for (int nn = 0; nn < N; ++nn) {
+    topK[nn] = RedType{value[nn], idx[nn]};
+  }
+
+  if constexpr (!IsSorted) {
+    Sort<N, RedType>::run(topK);
+  }
+  typename RedType::TypeCmp packedMax{};
+#pragma unroll
+  for (int kk = 0; kk < actualK; ++kk) {
+    bool update = kk > 0 && packedMax == topK[0].compValIdx;
+#pragma unroll
+    for (int nn = 0; nn < N; ++nn) {
+      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
+                 : update              ? topK[nn + 1]
+                                       : topK[nn];
+    }
+    // get the next largest value
+    packedMax = topK[0].reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
+  }
+};
+
+template <int K, typename Type, int N>
+__forceinline__ __device__ void reduceTopK(
+    cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
+    int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
+    Type const minValue, int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(N > 0, "Top K must have N > 0");
+  static_assert(
+      N <= 16,
+      "Only support candidates number less than or equal to 16*32=512");
+  static_assert(N <= 4 || N % 4 == 0,
+                "Only support candidates number is a multiple of 4*32=128 or "
+                "less than or equal to 4");
+  using RedType = TopKRedType<Type>;
+
+  if constexpr (N <= 4) {
+    reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
+                               actualK);
+  } else {
+    constexpr int numLoops = N / 4;
+    constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
+
+    Type topKBufferValue[numResults];
+    int32_t topKBufferIdx[numResults];
+    int32_t laneIdx = threadIdx.x % kWARP_SIZE;
+
+    for (int ii = 0; ii < numResults; ++ii) {
+      topKBufferValue[ii] = minValue;
+      topKBufferIdx[ii] = ii * kWARP_SIZE - 1;
+    }
+    for (int loop = 0; loop < numLoops; ++loop) {
+      int start = loop * 4;
+      Type topKValue[K];
+      int32_t topKIdx[K];
+      Type inValue[4];
+      int32_t inIdx[4];
+      for (int i = 0; i < 4; ++i) {
+        inValue[i] = value[start + i];
+        inIdx[i] = idx[start + i];
+      }
+      reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx,
+                                 minValue, actualK);
+      int inOffset = laneIdx % K;
+      if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
+        topKBufferValue[0] = topKValue[inOffset];
+        topKBufferIdx[0] = topKIdx[inOffset];
+      }
+      if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
+        topKBufferValue[1] = topKValue[inOffset];
+        topKBufferIdx[1] = topKIdx[inOffset];
+      }
+    }
+
+    reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
+                                        topKBufferIdx, minValue, actualK);
+  }
+};
+
+#undef TOPK_SWAP
+
+}  // namespace reduce_topk
+}  // namespace moe
+}  // namespace vllm

upstream_ref/ds_vllm/csrc/moe/moe_align_sum_kernels.cu

@@ -0,0 +1,833 @@
+#include <array>
+#include <cub/cub.cuh>
+
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/macros.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "../../cuda_compat.h"
+#include "core/math.hpp"
+#include "libtorch_stable/dispatch_utils.h"
+#include "libtorch_stable/torch_utils.h"
+
+#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
+
+namespace vllm {
+namespace moe {
+namespace batched_moe_align_block_size {
+
+// Note num_threads needs to be 1024 for BlockScan Reduction in the kernel.
+static constexpr int32_t num_threads = 1024;
+static constexpr int32_t num_blocks = 1;
+__global__ void batched_moe_align_block_size_kernel(
+    int32_t const num_batches, int32_t const max_tokens_per_batch,
+    int32_t const block_size, int32_t const* __restrict__ batch_num_tokens,
+    int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids,
+    int32_t* __restrict__ num_tokens_post_pad) {
+  // TODO(varun): This is a naive implementation. Could be optimized.
+
+  size_t const batch_id = threadIdx.x;
+  size_t const stride = blockDim.x * gridDim.x;
+  int32_t const num_blocks_per_batch =
+      CEILDIV(max_tokens_per_batch, block_size);
+  int32_t const sorted_ids_size =
+      num_blocks_per_batch * num_batches * block_size;
+  int32_t const block_ids_size = sorted_ids_size / block_size;
+  int32_t const SENTINEL =
+      num_batches * max_tokens_per_batch;  // To denote invalid entries.
+  // Initialize sorted_ids
+  for (size_t i = threadIdx.x; i < sorted_ids_size; i += stride) {
+    sorted_ids[i] = SENTINEL;
+  }
+  // Initialize expert_ids with -1
+  for (size_t i = threadIdx.x; i < block_ids_size; i += stride) {
+    block_ids[i] = -1;
+  }
+
+  int32_t b_num_tokens = 0;
+  if (batch_id < num_batches) {
+    b_num_tokens = batch_num_tokens[batch_id];
+  }
+  int32_t const ceil_b_num_tokens =
+      CEILDIV(b_num_tokens, block_size) * block_size;
+
+  // Compute prefix sum over token counts per expert
+  using BlockScan = cub::BlockScan<int32_t, 1024>;
+  __shared__ typename BlockScan::TempStorage temp_storage;
+  int cumsum_val;
+  BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val);
+  __syncthreads();
+
+  bool const is_last_batch = batch_id == (num_batches - 1);
+  if (is_last_batch) {
+    *num_tokens_post_pad = cumsum_val + ceil_b_num_tokens;
+  }
+
+  if (batch_id < num_batches) {
+    int32_t const batch_offset = batch_id * max_tokens_per_batch;
+    for (size_t i = 0; i < b_num_tokens; ++i) {
+      sorted_ids[cumsum_val + i] = batch_offset + i;
+    }
+
+    int32_t const block_start = cumsum_val / block_size;
+    int32_t const num_blocks = ceil_b_num_tokens / block_size;
+    for (size_t i = 0; i < num_blocks; ++i) {
+      block_ids[block_start + i] = batch_id;
+    }
+  }
+}
+}  // namespace batched_moe_align_block_size
+
+template <typename scalar_t>
+__device__ void _moe_align_block_size(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t* __restrict__ expert_map, int32_t num_experts,
+    int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size,
+    size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded,
+    int32_t max_num_m_blocks, int32_t model_offset, int32_t inactive_expert_id,
+    int32_t topk_num, int32_t* token_mask, bool has_expert_map) {
+  extern __shared__ int32_t shared_counts[];
+
+  // Compute input buffer offsets. Typically these will all be 0, except when
+  // using Multi LoRA.
+  int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
+  int expert_ids_offset = max_num_m_blocks * model_offset;
+  int cumsum_offset = (num_experts + 1) * model_offset;
+
+  // Use separate threadblocks to fill sorted_token_ids.
+  // This is safe since the current kernel does not use sorted_token_ids.
+  if (blockIdx.x % 2) {
+    // Initialize sorted_token_ids with numel
+    for (size_t it = threadIdx.x; it < max_num_tokens_padded;
+         it += blockDim.x) {
+      sorted_token_ids[sorted_token_ids_offset + it] = numel;
+    }
+    return;
+  }
+
+  const int warp_id = threadIdx.x / WARP_SIZE;
+  const int my_expert_start = warp_id * experts_per_warp;
+
+  for (int i = 0; i < experts_per_warp; ++i) {
+    if (my_expert_start + i < padded_num_experts) {
+      shared_counts[warp_id * experts_per_warp + i] = 0;
+    }
+  }
+
+  __syncthreads();
+
+  const size_t tid = threadIdx.x;
+  const size_t stride = blockDim.x;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int expert_id = topk_ids[i];
+    if (expert_id >= num_experts) {
+      continue;
+    }
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      // filter invalid experts
+      if (expert_id == -1) continue;
+    }
+    int warp_idx = expert_id / experts_per_warp;
+    int expert_offset = expert_id % experts_per_warp;
+    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+    atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
+              mask);
+  }
+
+  __syncthreads();
+
+  // Compute prefix sum over token counts per expert
+  using BlockScan = cub::BlockScan<int32_t, 1024>;
+  __shared__ typename BlockScan::TempStorage temp_storage;
+
+  int expert_count = 0;
+  int expert_id = threadIdx.x;
+  if (expert_id < num_experts) {
+    int warp_idx = expert_id / experts_per_warp;
+    int expert_offset = expert_id % experts_per_warp;
+    expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset];
+    expert_count = CEILDIV(expert_count, block_size) * block_size;
+  }
+
+  int cumsum_val;
+  BlockScan(temp_storage).ExclusiveSum(expert_count, cumsum_val);
+  if (expert_id <= num_experts) {
+    cumsum[cumsum_offset + expert_id] = cumsum_val;
+  }
+
+  if (expert_id == num_experts) {
+    total_tokens_post_pad[model_offset] = cumsum_val;
+  }
+
+  __syncthreads();
+
+  if (threadIdx.x < num_experts) {
+    for (int i = cumsum[cumsum_offset + threadIdx.x];
+         i < cumsum[cumsum_offset + threadIdx.x + 1]; i += block_size) {
+      expert_ids[expert_ids_offset + i / block_size] = threadIdx.x;
+    }
+  }
+
+  // Fill remaining expert_ids with -1
+  const size_t fill_start_idx =
+      cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x;
+  for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) {
+    expert_ids[expert_ids_offset + i] = inactive_expert_id;
+  }
+}
+
+template <typename scalar_t, int32_t fill_threads>
+__device__ void _moe_align_block_size_small_batch_expert(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
+    size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks,
+    int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num,
+    int32_t* token_mask, bool has_expert_map) {
+  // Compute input buffer offsets. Typically these will all be 0, except when
+  // using Multi LoRA.
+  int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
+  int expert_ids_offset = max_num_m_blocks * model_offset;
+
+  // Use an additional group of threads to fill sorted_token_ids.
+  // Since the current kernel will use sorted_token_ids afterward,
+  // we fill sorted_token_ids within the same threadblock to make
+  // synchronization easier.
+  if (threadIdx.x < fill_threads) {
+    // Initialize sorted_token_ids with numel
+    for (size_t it = threadIdx.x; it < max_num_tokens_padded;
+         it += fill_threads) {
+      sorted_token_ids[sorted_token_ids_offset + it] = numel;
+    }
+    // Three __syncthreads() corresponding to the other threads
+    __syncthreads();
+    __syncthreads();
+    __syncthreads();
+    return;
+  }
+
+  const size_t tid = threadIdx.x - fill_threads;
+  const size_t stride = blockDim.x - fill_threads;
+
+  extern __shared__ int32_t shared_mem[];
+  int32_t* cumsum = shared_mem;
+  int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1);
+
+  for (int i = 0; i < num_experts; ++i) {
+    tokens_cnts[(tid + 1) * num_experts + i] = 0;
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      // filter invalid expert
+      if (expert_id == -1) continue;
+    }
+    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+    tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
+  }
+
+  __syncthreads();
+
+  if (tid < num_experts) {
+    tokens_cnts[tid] = 0;
+    for (int i = 1; i <= stride; ++i) {
+      tokens_cnts[i * num_experts + tid] +=
+          tokens_cnts[(i - 1) * num_experts + tid];
+    }
+  }
+
+  __syncthreads();
+
+  if (tid == 0) {
+    cumsum[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      cumsum[i] =
+          cumsum[i - 1] +
+          CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) *
+              block_size;
+    }
+    total_tokens_post_pad[model_offset] =
+        static_cast<int32_t>(cumsum[num_experts]);
+  }
+
+  __syncthreads();
+
+  if (tid < num_experts) {
+    for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) {
+      expert_ids[expert_ids_offset + i / block_size] = tid;
+    }
+  }
+
+  // Fill remaining expert_ids with -1
+  const size_t fill_start_idx = cumsum[num_experts] / block_size + tid;
+  for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) {
+    expert_ids[expert_ids_offset + i] = inactive_expert_id;
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      // filter invalid expert
+      if (expert_id == -1) continue;
+    }
+    int32_t rank_post_pad =
+        tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
+
+    if (token_mask == nullptr || token_mask[i / topk_num]) {
+      sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
+      ++tokens_cnts[tid * num_experts + expert_id];
+    }
+  }
+}
+
+template <typename scalar_t>
+__device__ void _count_and_sort_expert_tokens(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
+    int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
+    int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask,
+    int32_t model_offset, int32_t topk_num, bool has_expert_map) {
+  const size_t tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.y;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (expert_id >= num_experts) {
+      continue;
+    }
+
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      // filter invalid experts
+      if (expert_id == -1) continue;
+    }
+
+    if (token_mask == nullptr || token_mask[i / topk_num]) {
+      int32_t rank_post_pad = atomicAdd(
+          &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
+      sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
+          i;
+    }
+  }
+}
+
+template <typename scalar_t>
+__global__ void moe_align_block_size_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t* __restrict__ expert_map, int32_t num_experts,
+    int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size,
+    size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded,
+    int32_t topk_num, bool has_expert_map) {
+  _moe_align_block_size(
+      topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
+      num_experts, padded_num_experts, experts_per_warp, block_size, numel,
+      cumsum, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size),
+      0, -1, topk_num, nullptr, has_expert_map);
+}
+
+template <typename scalar_t>
+__global__ void count_and_sort_expert_tokens_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
+    int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
+    int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map) {
+  _count_and_sort_expert_tokens(
+      topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
+      max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
+}
+
+template <typename scalar_t, int TOPK>
+__global__ void moe_sum_kernel(
+    scalar_t* __restrict__ out,          // [..., d]
+    const scalar_t* __restrict__ input,  // [..., topk, d]
+    const int d) {
+  const int64_t token_idx = blockIdx.x;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    scalar_t x = 0.0;
+#pragma unroll
+    for (int k = 0; k < TOPK; ++k) {
+      x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
+    }
+    out[token_idx * d + idx] = x;
+  }
+}
+
+template <typename scalar_t, int32_t fill_threads>
+__global__ void moe_align_block_size_small_batch_expert_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
+    size_t numel, int32_t max_num_tokens_padded, int32_t topk_num,
+    bool has_expert_map) {
+  _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
+      topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
+      num_experts, block_size, numel, max_num_tokens_padded,
+      CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
+      has_expert_map);
+}
+
+template <typename scalar_t>
+__global__ void moe_lora_align_block_size_kernel(
+    scalar_t* __restrict__ topk_ids, int32_t* __restrict__ token_lora_mapping,
+    int64_t block_size, int32_t* __restrict__ expert_map, int num_experts,
+    int max_loras, size_t numel, int max_num_tokens_padded,
+    int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids, int32_t topk_num,
+    int32_t* total_tokens_post_pad, int32_t* adapter_enabled,
+    int32_t* __restrict__ cumsum, int32_t experts_per_warp,
+    int32_t padded_num_experts, int32_t* lora_ids,
+    int32_t* __restrict__ token_mask, bool has_expert_map) {
+  int lora_idx = blockIdx.x / 2;
+  int lora_id = lora_ids[lora_idx];
+  // Output buffers are indexed by lora_id (in [0, max_loras)). The grid
+  // iterates one extra slot to accommodate the "-1" entry that
+  // active_lora_ids may hold in position 0 for mixed base + LoRA batches;
+  // guard against any other unexpected lora_id >= max_loras to avoid
+  // out-of-bounds writes. This mirrors the `lora_id >= max_loras` guard in
+  // the Triton _fused_moe_lora_kernel.
+  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
+    return;
+  }
+
+  // Populate the token_mask based on the token-LoRA mapping
+  int num_tokens = numel / topk_num;
+  if (threadIdx.x == 0) {
+    total_tokens_post_pad[lora_id] = 0;
+
+    for (int i = 0; i < num_tokens; i++) {
+      token_mask[(lora_id * num_tokens) + i] =
+          (int)token_lora_mapping[i] == lora_id;
+    }
+  }
+
+  __syncthreads();
+
+  _moe_align_block_size(
+      topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
+      num_experts, padded_num_experts, experts_per_warp, block_size, numel,
+      cumsum, max_num_tokens_padded, max_num_m_blocks, lora_id, -1, topk_num,
+      &token_mask[(lora_id * num_tokens)], has_expert_map);
+}
+
+template <typename scalar_t>
+__global__ void lora_count_and_sort_expert_tokens_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
+    int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
+    int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask,
+    int32_t max_loras, int32_t* lora_ids, int32_t* adapter_enabled,
+    bool has_expert_map) {
+  int lora_idx = blockIdx.x;
+  int lora_id = lora_ids[lora_idx];
+  // Same guard rationale as moe_lora_align_block_size_kernel. Additionally
+  // skip disabled adapter slots: moe_lora_align_block_size_kernel early-returns
+  // for them and leaves token_mask[lora_id, :] uninitialized (token_mask is
+  // allocated with torch::empty), so running the sort loop here would traverse
+  // garbage mask bits and pollute this slot's rows of sorted_token_ids and
+  // cumsum_buffer. Downstream consumers already skip disabled slots, so the
+  // pollution is dormant today, but the check keeps behavior symmetric with
+  // the other two align kernels and avoids O(numel) wasted work per disabled
+  // slot. Short-circuit evaluation ensures adapter_enabled is only indexed
+  // after lora_id is confirmed to be in [0, max_loras).
+  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
+    return;
+  }
+
+  int num_tokens = numel / topk_num;
+
+  _count_and_sort_expert_tokens(
+      topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
+      max_num_tokens_padded, &token_mask[(lora_id * num_tokens)], lora_id,
+      topk_num, has_expert_map);
+}
+
+template <typename scalar_t, int32_t fill_threads>
+__global__ void moe_lora_align_block_size_small_batch_expert_kernel(
+    scalar_t* __restrict__ topk_ids, int32_t* token_lora_mapping,
+    int64_t block_size, int32_t* __restrict__ expert_map, int num_experts,
+    int max_loras, size_t numel, int max_num_tokens_padded,
+    int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids, int topk_num,
+    int32_t* total_tokens_post_pad, int32_t* adapter_enabled, int32_t* lora_ids,
+    int32_t* token_mask, bool has_expert_map) {
+  int lora_idx = blockIdx.x;
+  int lora_id = lora_ids[lora_idx];
+  // Same guard rationale as moe_lora_align_block_size_kernel.
+  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
+    return;
+  }
+
+  int num_tokens = numel / topk_num;
+  if (threadIdx.x == 0) {
+    total_tokens_post_pad[lora_id] = 0;
+
+    for (int i = 0; i < num_tokens; i++) {
+      token_mask[(lora_id * num_tokens) + i] =
+          (int)token_lora_mapping[i] == lora_id;
+    }
+  }
+
+  __syncthreads();
+
+  _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
+      topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
+      num_experts, block_size, numel, max_num_tokens_padded, max_num_m_blocks,
+      -1, lora_id, topk_num, &token_mask[(lora_id * num_tokens)],
+      has_expert_map);
+}
+
+}  // namespace moe
+}  // namespace vllm
+
+// taken from
+// https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
+void moe_align_block_size(
+    torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    std::optional<torch::stable::Tensor> maybe_expert_map) {
+  const cudaStream_t stream =
+      get_current_cuda_stream(topk_ids.get_device_index());
+
+  int64_t padded_num_experts =
+      ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+  int experts_per_warp = WARP_SIZE;
+  int threads = 1024;
+  threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+
+  // BlockScan uses 1024 threads and assigns one thread per expert.
+  STD_TORCH_CHECK(padded_num_experts < 1024,
+                  "padded_num_experts must be less than 1024");
+  bool has_expert_map = maybe_expert_map.has_value();
+  torch::stable::Tensor expert_map;
+  if (has_expert_map) {
+    expert_map = maybe_expert_map.value();
+  } else {
+    expert_map = torch::stable::new_empty(topk_ids, {0},
+                                          torch::headeronly::ScalarType::Int);
+  }
+
+  VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
+      topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
+        // calc needed amount of shared mem for `cumsum` tensors
+        bool small_batch_expert_mode =
+            (topk_ids.numel() < 1024) && (num_experts <= 64);
+
+        if (small_batch_expert_mode) {
+          const int32_t threads = max((int32_t)num_experts, WARP_SIZE);
+          const int32_t shared_mem_size =
+              ((threads + 1) * num_experts + (num_experts + 1)) *
+              sizeof(int32_t);
+
+          // threadIdx.x >= fill_threads: counting experts and aligning
+          // threadIdx.x < fill_threads: filling sorted_token_ids
+          constexpr int32_t fill_threads = 256;
+          auto small_batch_expert_kernel =
+              vllm::moe::moe_align_block_size_small_batch_expert_kernel<
+                  scalar_t, fill_threads>;
+          small_batch_expert_kernel<<<1, fill_threads + threads,
+                                      shared_mem_size, stream>>>(
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, block_size, topk_ids.numel(),
+              sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
+        } else {
+          torch::stable::Tensor cumsum_buffer = torch::stable::new_empty(
+              topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
+          auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>;
+
+          size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp);
+          size_t shared_mem_size =
+              num_warps * experts_per_warp * sizeof(int32_t);
+
+          // launch two threadblocks
+          // blockIdx.x == 0: counting experts and aligning
+          // blockIdx.x == 1: filling sorted_token_ids
+          align_kernel<<<2, threads, shared_mem_size, stream>>>(
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, padded_num_experts, experts_per_warp, block_size,
+              topk_ids.numel(),
+              reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
+              sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
+
+          const int block_threads = std::min(256, (int)threads);
+          const int num_blocks =
+              (topk_ids.numel() + block_threads - 1) / block_threads;
+          const int max_blocks = 65535;
+          const int actual_blocks = std::min(num_blocks, max_blocks);
+          dim3 gridDims(1, actual_blocks);
+
+          auto sort_kernel =
+              vllm::moe::count_and_sort_expert_tokens_kernel<scalar_t>;
+          sort_kernel<<<gridDims, block_threads, 0, stream>>>(
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              topk_ids.numel(), num_experts, sorted_token_ids.size(0),
+              topk_ids.size(1), has_expert_map);
+        }
+      });
+}
+
+void batched_moe_align_block_size(int64_t max_tokens_per_batch,
+                                  int64_t block_size,
+                                  const torch::stable::Tensor& batch_num_tokens,
+                                  torch::stable::Tensor sorted_ids,
+                                  torch::stable::Tensor batch_ids,
+                                  torch::stable::Tensor num_tokens_post_pad) {
+  namespace batched_kernel = vllm::moe::batched_moe_align_block_size;
+
+  const cudaStream_t stream =
+      get_current_cuda_stream(batch_num_tokens.get_device_index());
+  int32_t const B = batch_num_tokens.size(0);
+  int32_t const num_blocks_per_batch =
+      round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size;
+  int32_t const num_blocks = num_blocks_per_batch * B;
+  int64_t const sorted_ids_size = num_blocks * block_size;
+
+  STD_TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size);
+  STD_TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size);
+  STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
+  STD_TORCH_CHECK(B <= batched_kernel::num_threads);
+
+  batched_kernel::batched_moe_align_block_size_kernel<<<
+      batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>(
+      B, max_tokens_per_batch, block_size,
+      reinterpret_cast<const int32_t*>(batch_num_tokens.const_data_ptr()),
+      reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()),
+      reinterpret_cast<int32_t*>(batch_ids.mutable_data_ptr()),
+      reinterpret_cast<int32_t*>(num_tokens_post_pad.mutable_data_ptr()));
+}
+
+void moe_sum(torch::stable::Tensor& input,   // [num_tokens, topk, hidden_size]
+             torch::stable::Tensor& output)  // [num_tokens, hidden_size]
+{
+  const int hidden_size = input.size(-1);
+  const auto num_tokens = output.numel() / hidden_size;
+  const int topk = input.size(1);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(hidden_size, 1024));
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      output.get_device_index());
+  const cudaStream_t stream =
+      get_current_cuda_stream(output.get_device_index());
+
+  switch (topk) {
+    case 2:
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+          input.scalar_type(), "moe_sum_kernel", [&] {
+            vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
+                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
+                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
+                hidden_size);
+          });
+      break;
+
+    case 3:
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+          input.scalar_type(), "moe_sum_kernel", [&] {
+            vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
+                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
+                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
+                hidden_size);
+          });
+      break;
+
+    case 4:
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+          input.scalar_type(), "moe_sum_kernel", [&] {
+            vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
+                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
+                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
+                hidden_size);
+          });
+      break;
+
+    default:
+      torch::stable::sum_out(output, input, std::array<int64_t, 1>{1});
+      break;
+  }
+}
+
+void moe_lora_align_block_size(
+    torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
+    int64_t num_experts, int64_t block_size, int64_t max_loras,
+    int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
+    std::optional<torch::stable::Tensor> maybe_expert_map) {
+  const int topk_num = topk_ids.size(1);
+
+  STD_TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
+
+  int device_max_shared_mem;
+  int dev = topk_ids.get_device_index();
+  cudaDeviceGetAttribute(&device_max_shared_mem,
+                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
+  const cudaStream_t stream = get_current_cuda_stream(dev);
+
+  int64_t padded_num_experts =
+      ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+
+  // BlockScan uses 1024 threads and assigns one thread per expert.
+  STD_TORCH_CHECK(padded_num_experts < 1024,
+                  "padded_num_experts must be less than 1024");
+
+  torch::stable::Tensor token_mask =
+      torch::stable::new_empty(topk_ids, {max_loras * topk_ids.size(0)},
+                               torch::headeronly::ScalarType::Int);
+  bool has_expert_map = maybe_expert_map.has_value();
+  torch::stable::Tensor expert_map;
+  if (has_expert_map) {
+    expert_map = maybe_expert_map.value();
+  } else {
+    expert_map = torch::stable::new_empty(topk_ids, {0},
+                                          torch::headeronly::ScalarType::Int);
+  }
+
+  VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(
+      topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] {
+        bool small_batch_expert_mode =
+            (topk_ids.numel() < 1024) && (num_experts <= 64);
+
+        if (small_batch_expert_mode) {
+          const int32_t num_thread = max((int32_t)num_experts, 128);
+          const int32_t shared_mem =
+              (num_thread + 1) * num_experts * sizeof(int32_t) +
+              (num_experts + 1) * sizeof(int32_t);
+          if (shared_mem > device_max_shared_mem) {
+            STD_TORCH_CHECK(false, "Shared memory usage exceeds device limit.");
+          }
+
+          // threadIdx.x >= fill_threads: counting experts and aligning
+          // threadIdx.x < fill_threads: filling sorted_token_ids
+          constexpr int32_t fill_threads = 256;
+
+          dim3 blockDim(num_thread + fill_threads);
+          auto kernel =
+              vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel<
+                  scalar_t, fill_threads>;
+          STD_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
+              (void*)kernel, shared_mem));
+          // Grid size is (max_loras + 1) because active_lora_ids has length
+          // max_loras + 1: sorted-unique values of token_lora_mapping, which
+          // can include -1 (base-model tokens) in addition to up to max_loras
+          // real LoRA slots. Using max_loras would drop the real LoRA slot
+          // when -1 is present at position 0 and leave output buffers
+          // uninitialized, causing illegal memory accesses in downstream
+          // MoE-LoRA kernels. This mirrors the fix made for the Triton
+          // _fused_moe_lora_kernel grid in vllm-project/vllm#32277.
+          kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>(
+              reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
+              block_size,
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
+              max_num_m_blocks,
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
+              topk_num,
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
+              has_expert_map);
+        } else {
+          int num_thread = 1024;
+          dim3 blockDim(num_thread);
+          size_t num_warps = CEILDIV(padded_num_experts, WARP_SIZE);
+
+          size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t);
+
+          // cumsum buffer
+          torch::stable::Tensor cumsum = torch::stable::new_zeros(
+              topk_ids, {max_loras * (num_experts + 1)},
+              torch::headeronly::ScalarType::Int);
+
+          auto align_kernel =
+              vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;
+
+          // Launch two threadblocks per LoRA slot, across max_loras + 1 slots
+          // to cover the extra "-1" (base-model tokens) entry that
+          // active_lora_ids may contain in addition to up to max_loras real
+          // LoRA slots. Using max_loras would drop the real LoRA slot when -1
+          // occupies position 0 and leave the output buffers uninitialized,
+          // causing illegal memory accesses downstream. Mirrors the grid fix
+          // applied to _fused_moe_lora_kernel in vllm-project/vllm#32277.
+          // blockIdx.x % 2 == 0: counting experts and aligning
+          // blockIdx.x % 2 == 1: filling sorted_token_ids
+          align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size,
+                         stream>>>(
+              reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
+              block_size,
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
+              max_num_m_blocks,
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
+              topk_num,
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()), WARP_SIZE,
+              padded_num_experts,
+              reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
+              has_expert_map);
+
+          const int block_threads = std::min(256, (int)num_thread);
+          const int num_blocks =
+              (topk_ids.numel() + block_threads - 1) / block_threads;
+
+          const int max_blocks = 65535;
+          const int actual_blocks = std::min(num_blocks, max_blocks);
+
+          // Same rationale as align_kernel above: iterate over max_loras + 1
+          // slots so the sort kernel processes the real LoRA slot even when
+          // active_lora_ids has -1 at position 0.
+          dim3 gridDims(max_loras + 1, actual_blocks);
+          auto sort_kernel =
+              vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;
+
+          sort_kernel<<<gridDims, block_threads, 0, stream>>>(
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num,
+              reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
+              max_loras,
+              reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
+              has_expert_map);
+        }
+      });
+}

upstream_ref/ds_vllm/csrc/moe/moe_ops.h

@@ -0,0 +1,87 @@
+#pragma once
+
+#include <torch/csrc/stable/tensor.h>
+
+#include <optional>
+#include <tuple>
+
+void topk_softmax(torch::stable::Tensor& topk_weights,
+                  torch::stable::Tensor& topk_indices,
+                  torch::stable::Tensor& token_expert_indices,
+                  torch::stable::Tensor& gating_output, bool renormalize,
+                  std::optional<torch::stable::Tensor> bias);
+
+void topk_sigmoid(torch::stable::Tensor& topk_weights,
+                  torch::stable::Tensor& topk_indices,
+                  torch::stable::Tensor& token_expert_indices,
+                  torch::stable::Tensor& gating_output, bool renormalize,
+                  std::optional<torch::stable::Tensor> bias);
+
+void topk_softplus_sqrt(
+    torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
+    torch::stable::Tensor& token_expert_indices,
+    torch::stable::Tensor& gating_output, bool renormalize,
+    double routed_scaling_factor,
+    const std::optional<torch::stable::Tensor>& correction_bias,
+    const std::optional<torch::stable::Tensor>& input_ids,
+    const std::optional<torch::stable::Tensor>& tid2eid);
+
+void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output);
+
+void moe_align_block_size(
+    torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    std::optional<torch::stable::Tensor> maybe_expert_map);
+
+void batched_moe_align_block_size(
+    int64_t max_tokens_per_batch, int64_t block_size,
+    const torch::stable::Tensor& expert_num_tokens,
+    torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad);
+
+void moe_lora_align_block_size(
+    torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
+    int64_t num_experts, int64_t block_size, int64_t max_loras,
+    int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
+    std::optional<torch::stable::Tensor> maybe_expert_map);
+#ifndef USE_ROCM
+torch::stable::Tensor moe_wna16_gemm(
+    torch::stable::Tensor input, torch::stable::Tensor output,
+    torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales,
+    std::optional<torch::stable::Tensor> b_qzeros,
+    std::optional<torch::stable::Tensor> topk_weights,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad, int64_t top_k,
+    int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K,
+    int64_t bit);
+
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
+    const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group,
+    int64_t topk, bool renormalize, double routed_scaling_factor,
+    const torch::stable::Tensor& bias, int64_t scoring_func);
+#endif
+
+bool moe_permute_unpermute_supported();
+
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t num_expert);
+
+void shuffle_rows(const torch::stable::Tensor& input_tensor,
+                  const torch::stable::Tensor& dst2src_map,
+                  torch::stable::Tensor& output_tensor);
+
+#ifndef USE_ROCM
+// DeepSeek V3 optimized router GEMM kernel for SM90+
+// Computes output = mat_a @ mat_b.T where:
+//   mat_a: [num_tokens, hidden_dim] in bf16
+//   mat_b: [num_experts, hidden_dim] in bf16
+//   output: [num_tokens, num_experts] in bf16 or fp32
+// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
+void dsv3_router_gemm(torch::stable::Tensor& output,
+                      const torch::stable::Tensor& mat_a,
+                      const torch::stable::Tensor& mat_b);
+#endif

upstream_ref/ds_vllm/csrc/moe/topk_softmax_kernels.cu

@@ -0,0 +1,874 @@
+/*
+ * Adapted from https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
+ * Copyright (c) 2024, The vLLM team.
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * 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 <type_traits>
+
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include "../../cuda_compat.h"
+#include "../../cub_helpers.h"
+#include "libtorch_stable/torch_utils.h"
+
+#ifndef USE_ROCM
+    #include <cuda_bf16.h>
+    #include <cuda_fp16.h>
+#else
+    #include <hip/hip_bf16.h>
+    #include <hip/hip_fp16.h>
+    typedef __hip_bfloat16 __nv_bfloat16;
+    typedef __hip_bfloat162 __nv_bfloat162;
+#endif
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+namespace vllm {
+namespace moe {
+
+/// Aligned array type
+template <
+    typename T,
+    /// Number of elements in the array
+    int N,
+    /// Alignment requirement in bytes
+    int Alignment = sizeof(T) * N
+>
+struct alignas(Alignment) AlignedArray {
+    T data[N];
+};
+
+template <typename T>
+__device__ __forceinline__ float toFloat(T value) {
+    if constexpr (std::is_same_v<T, float>) {
+        return value;
+    } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+        return __bfloat162float(value);
+    } else if constexpr (std::is_same_v<T, __half>) {
+        return __half2float(value);
+    }
+}
+
+// Scoring function enums
+enum ScoringFunc {
+  SCORING_SOFTMAX = 0, // apply softmax
+  SCORING_SIGMOID = 1  // apply sigmoid
+};
+
+// ====================== Softmax things ===============================
+// We have our own implementation of softmax here so we can support transposing the output
+// in the softmax kernel when we extend this module to support expert-choice routing.
+template <int TPB, typename InputType>
+__launch_bounds__(TPB) __global__
+    void moeSoftmax(const InputType* input, const bool* finished, float* output, const int num_cols)
+{
+    using BlockReduce = cub::BlockReduce<float, TPB>;
+    __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+    __shared__ float normalizing_factor;
+    __shared__ float float_max;
+
+    const int thread_row_offset = blockIdx.x * num_cols;
+
+    float threadData(-FLT_MAX);
+
+    // Don't touch finished rows.
+    if ((finished != nullptr) && finished[blockIdx.x])
+    {
+        return;
+    }
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        const float val = toFloat(input[idx]);
+        threadData = max(val, threadData);
+    }
+
+    const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, CubMaxOp());
+    if (threadIdx.x == 0)
+    {
+        float_max = maxElem;
+    }
+    __syncthreads();
+
+    threadData = 0;
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        const float val = toFloat(input[idx]);
+        threadData += expf(val - float_max);
+    }
+
+    const auto Z = BlockReduce(tmpStorage).Reduce(threadData, CubAddOp());
+
+    if (threadIdx.x == 0)
+    {
+        normalizing_factor = 1.f / Z;
+    }
+    __syncthreads();
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        const float val = toFloat(input[idx]);
+        float softmax_val = expf(val - float_max) * normalizing_factor;
+        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
+        if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f;
+        output[idx] = softmax_val;
+    }
+}
+
+template <int TPB, typename InputType>
+__launch_bounds__(TPB) __global__
+    void moeSigmoid(const InputType* input, const bool* finished, float* output, const int num_cols)
+{
+    const int thread_row_offset = blockIdx.x * num_cols;
+
+    // Don't touch finished rows.
+    if ((finished != nullptr) && finished[blockIdx.x])
+    {
+        return;
+    }
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        const float val = toFloat(input[idx]);
+        float sigmoid_val = 1.0f / (1.0f + __expf(-val));
+        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
+        if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f;
+        output[idx] = sigmoid_val;
+    }
+}
+
+template <int TPB, typename IndType>
+__launch_bounds__(TPB) __global__ void moeTopK(
+    const float* inputs_after_softmax,
+    const bool* finished,
+    float* output,
+    IndType* indices,
+    int* source_rows,
+    const int num_experts,
+    const int k,
+    const int start_expert,
+    const int end_expert,
+    const bool renormalize,
+    const float* bias)
+{
+
+    using cub_kvp = cub::KeyValuePair<int, float>;
+    using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+    __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+    cub_kvp thread_kvp;
+    cub::ArgMax arg_max;
+
+    const int num_rows = gridDim.x;
+    const int block_row = blockIdx.x;
+
+    const bool row_is_active = finished ? !finished[block_row] : true;
+    const int thread_read_offset = blockIdx.x * num_experts;
+    float selected_sum = 0.f;
+    for (int k_idx = 0; k_idx < k; ++k_idx)
+    {
+        thread_kvp.key = 0;
+        thread_kvp.value = -1.f; // This is OK because inputs are probabilities
+
+        cub_kvp inp_kvp;
+        for (int expert = threadIdx.x; expert < num_experts; expert += TPB)
+        {
+            const int idx = thread_read_offset + expert;
+            inp_kvp.key = expert;
+
+            // Apply correction bias if provided
+            if (bias != nullptr) {
+              inp_kvp.value = inputs_after_softmax[idx] + bias[expert];
+            } else {
+              inp_kvp.value = inputs_after_softmax[idx];
+            }
+
+            for (int prior_k = 0; prior_k < k_idx; ++prior_k)
+            {
+                const int prior_winning_expert = indices[k * block_row + prior_k];
+
+                if (prior_winning_expert == expert)
+                {
+                    inp_kvp = thread_kvp;
+                }
+            }
+
+            thread_kvp = arg_max(inp_kvp, thread_kvp);
+        }
+
+        const cub_kvp result_kvp = BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+        if (threadIdx.x == 0)
+        {
+            // Ignore experts the node isn't responsible for with expert parallelism
+            const int expert = result_kvp.key;
+            const bool node_uses_expert = expert >= start_expert && expert < end_expert;
+            const bool should_process_row = row_is_active && node_uses_expert;
+
+            const int idx = k * block_row + k_idx;
+            // Return the unbiased scores for output weights
+            output[idx] = inputs_after_softmax[thread_read_offset + expert];
+            indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
+            assert(indices[idx] >= 0);
+            source_rows[idx] = k_idx * num_rows + block_row;
+            if (renormalize) {
+                selected_sum += inputs_after_softmax[thread_read_offset + expert];
+            }
+        }
+        __syncthreads();
+    }
+
+    // Renormalize the k weights for this row to sum to 1, if requested.
+    if (renormalize) {
+        if (threadIdx.x == 0) {
+            const float denom = selected_sum > 0.f ? selected_sum : 1.f;
+            for (int k_idx = 0; k_idx < k; ++k_idx) {
+                const int idx = k * block_row + k_idx;
+                output[idx] = output[idx] / denom;
+            }
+        }
+    }
+}
+
+// ====================== TopK softmax things ===============================
+
+/*
+  A Top-K gating softmax written to exploit when the number of experts in the MoE layers
+  are a small power of 2. This allows us to cleanly share the rows among the threads in
+  a single warp and eliminate communication between warps (so no need to use shared mem).
+
+  It fuses the softmax, max and argmax into a single kernel.
+
+  Limitations:
+  1) This implementation is optimized for when the number of experts is a small power of 2.
+     Additionally it also supports when number of experts is multiple of 64 which is still
+     faster than the computing softmax and topK separately (only tested on CUDA yet).
+  2) This implementation assumes k is small, but will work for any k.
+*/
+
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType,
+          typename InputType = float, ScoringFunc SF>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
+    void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
+        int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
+        const float* bias)
+{
+    static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
+                      std::is_same_v<InputType, __half>,
+                  "InputType must be float, __nv_bfloat16, or __half");
+
+    // We begin by enforcing compile time assertions and setting up compile time constants.
+    static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2");
+    static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+    // Number of bytes each thread pulls in per load
+    static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
+    static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+    static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+    static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+    if constexpr (std::is_same_v<InputType, __nv_bfloat16> || std::is_same_v<InputType, __half>) {
+        static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
+            "ELTS_PER_LDG must be 1 or even for 16-bit conversion");
+    }
+
+    // Restrictions based on previous section.
+    static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
+    static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
+    static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
+    static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM, "THREADS_PER_ROW can be at most warp size");
+
+    // We have NUM_EXPERTS elements per row. We specialize for small #experts
+    static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
+    static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+    static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+    // Restrictions for previous section.
+    static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp");
+
+    // ===================== From this point, we finally start computing run-time variables. ========================
+
+    // Compute CTA and warp rows. We pack multiple rows into a single warp, and a block contains WARPS_PER_CTA warps.
+    // This, each block processes a chunk of rows. We start by computing the start row for each block.
+    const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+
+    // Now, using the base row per thread block, we compute the base row per warp.
+    const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+
+    // The threads in a warp are split into sub-groups that will work on a row.
+    // We compute row offset for each thread sub-group
+    const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+    const int thread_row = warp_base_row + thread_row_in_warp;
+
+    // Threads with indices out of bounds should early exit here.
+    if (thread_row >= num_rows)
+    {
+        return;
+    }
+    const bool row_is_active = finished ? !finished[thread_row] : true;
+
+    // We finally start setting up the read pointers for each thread. First, each thread jumps to the start of the
+    // row it will read.
+    const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+
+    // Now, we compute the group each thread belong to in order to determine the first column to start loads.
+    const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+    const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+    const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+    // Finally, we pull in the data from global mem
+    float row_chunk[VPT];
+
+    // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert to float
+    if constexpr (std::is_same_v<InputType, float>) {
+        using VecType = AlignedArray<float, ELTS_PER_LDG>;
+        VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
+        const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
+#pragma unroll
+        for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+            row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+        }
+    } else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
+        if constexpr (ELTS_PER_LDG >= 2) {
+            using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
+            float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
+            const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
+#pragma unroll
+            for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+                VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+                int base_idx_f2 = ii * ELTS_PER_LDG / 2;
+#pragma unroll
+                for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
+                    row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
+                        *reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2)
+                    );
+                }
+            }
+        } else { // ELTS_PER_LDG == 1
+#pragma unroll
+            for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+                const __nv_bfloat16* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
+                row_chunk[ii] = __bfloat162float(*scalar_ptr);
+            }
+        }
+    } else if constexpr (std::is_same_v<InputType, __half>) {
+        if constexpr (ELTS_PER_LDG >= 2) {
+            using VecType = AlignedArray<__half, ELTS_PER_LDG>;
+            float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
+            const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
+#pragma unroll
+            for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+                VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+                int base_idx_f2 = ii * ELTS_PER_LDG / 2;
+#pragma unroll
+                for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
+                    row_chunk_f2[base_idx_f2 + jj] = __half22float2(
+                        *reinterpret_cast<const __half2*>(vec.data + jj * 2)
+                    );
+                }
+            }
+        } else { // ELTS_PER_LDG == 1
+#pragma unroll
+            for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+                const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
+                row_chunk[ii] = __half2float(*scalar_ptr);
+            }
+        }
+    }
+
+    if constexpr (SF == SCORING_SOFTMAX) {
+      // First, we perform a max reduce within the thread.
+      float thread_max = row_chunk[0];
+#pragma unroll
+      for (int ii = 1; ii < VPT; ++ii) {
+        thread_max = max(thread_max, row_chunk[ii]);
+      }
+
+// Now, we find the max within the thread group and distribute among the threads. We use a butterfly reduce.
+#pragma unroll
+      for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+      {
+        thread_max = max(thread_max, VLLM_SHFL_XOR_SYNC_WIDTH(thread_max, mask, THREADS_PER_ROW));
+      }
+
+      // From this point, thread max in all the threads have the max within the row.
+      // Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum.
+      float row_sum = 0;
+#pragma unroll
+      for (int ii = 0; ii < VPT; ++ii)
+      {
+        row_chunk[ii] = expf(row_chunk[ii] - thread_max);
+        row_sum += row_chunk[ii];
+      }
+
+// Now, we perform the sum reduce within each thread group. Similar to the max reduce, we use a bufferfly pattern.
+#pragma unroll
+      for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+      {
+        row_sum += VLLM_SHFL_XOR_SYNC_WIDTH(row_sum, mask, THREADS_PER_ROW);
+      }
+
+      // From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables
+      // respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to
+      // compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row.
+      // However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the
+      // argmax after computing the softmax.
+      const float reciprocal_row_sum = 1.f / row_sum;
+
+#pragma unroll
+      for (int ii = 0; ii < VPT; ++ii)
+      {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+      }
+    } else if constexpr (SF == SCORING_SIGMOID) {
+#pragma unroll
+      for (int ii = 0; ii < VPT; ++ii)
+      {
+        row_chunk[ii] = 1.0f / (1.0f + __expf(-row_chunk[ii]));
+      }
+    }
+
+    // Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs.
+    // NaN gating (from degenerate hidden states in CUDA graph padding) causes
+    // softmax to produce all-NaN, which makes the argmax loop always pick
+    // expert 0 for every top-k slot, producing duplicate expert IDs that
+    // crash FlashInfer's three-step MoE sort.
+    // With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1].
+#pragma unroll
+    for (int ii = 0; ii < VPT; ++ii) {
+      if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
+        row_chunk[ii] = 0.f;
+      }
+    }
+
+    static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+    // If bias is not null, use biased value for selection
+    float row_chunk_for_choice[VPT];
+    // Apply correction bias
+    if (bias != nullptr) {
+#pragma unroll
+      for (int ldg = 0; ldg < LDG_PER_THREAD; ++ldg) {
+#pragma unroll
+        for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+          const int expert = first_elt_read_by_thread + ldg * COLS_PER_GROUP_LDG + ii;
+          float bias_val = expert < NUM_EXPERTS ? bias[expert] : 0.0f;
+          row_chunk_for_choice[ldg * ELTS_PER_LDG + ii] = row_chunk[ldg * ELTS_PER_LDG + ii] + bias_val;
+        }
+      }
+    } else {
+#pragma unroll
+      for (int ii = 0; ii < VPT; ++ii) {
+        row_chunk_for_choice[ii] = row_chunk[ii];
+      }
+    }
+
+    // Now, row_chunk contains the softmax / sigmoid of the row chunk. Now, I want to find the topk elements in each row, along
+    // with the max index.
+    int start_col = first_elt_read_by_thread;
+
+    float selected_sum = 0.f;
+    for (int k_idx = 0; k_idx < k; ++k_idx)
+    {
+        // First, each thread does the local argmax
+        float max_val_for_choice = row_chunk_for_choice[0];
+        float max_val = row_chunk[0];
+        int expert = start_col;
+#pragma unroll
+        for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG)
+        {
+#pragma unroll
+            for (int ii = 0; ii < ELTS_PER_LDG; ++ii)
+            {
+                float val_for_choice = row_chunk_for_choice[ldg * ELTS_PER_LDG + ii];
+                float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+
+                // No check on the experts here since columns with the smallest index are processed first and only
+                // updated if > (not >=)
+                if (val_for_choice > max_val_for_choice)
+                {
+                    max_val_for_choice = val_for_choice;
+                    max_val = val;
+                    expert = col + ii;
+                }
+            }
+        }
+
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads reach consensus about the max.
+// This will be useful for K > 1 so that the threads can agree on "who" had the max value. That thread can
+// then blank out their max with -inf and the warp can run more iterations...
+#pragma unroll
+        for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+        {
+            float other_max_for_choice = VLLM_SHFL_XOR_SYNC_WIDTH(max_val_for_choice, mask, THREADS_PER_ROW);
+            float other_max = VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
+            int other_expert = VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
+
+            // We want lower indices to "win" in every thread so we break ties this way
+            if (other_max_for_choice > max_val_for_choice || (other_max_for_choice == max_val_for_choice && other_expert < expert))
+            {
+                max_val_for_choice = other_max_for_choice;
+                max_val = other_max;
+                expert = other_expert;
+            }
+        }
+
+        // Write the max for this k iteration to global memory.
+        if (thread_group_idx == 0)
+        {
+            // Add a guard to ignore experts not included by this node
+            const bool node_uses_expert = expert >= start_expert && expert < end_expert;
+            const bool should_process_row = row_is_active && node_uses_expert;
+
+            // The lead thread from each sub-group will write out the final results to global memory. (This will be a
+            // single) thread per row of the input/output matrices.
+            const int idx = k * thread_row + k_idx;
+            output[idx] = max_val;
+            indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+            source_rows[idx] = k_idx * num_rows + thread_row;
+            if (renormalize) {
+                selected_sum += max_val;
+            }
+        }
+
+        // Finally, we clear the value in the thread with the current max if there is another iteration to run.
+        if (k_idx + 1 < k)
+        {
+            const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+            const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+            // Only the thread in the group which produced the max will reset the "winning" value to -inf.
+            if (thread_group_idx == thread_to_clear_in_group)
+            {
+                const int offset_for_expert = expert % ELTS_PER_LDG;
+                // Safe to set to any negative value since row_chunk values must be between 0 and 1.
+                row_chunk_for_choice[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
+            }
+        }
+    }
+
+    // Renormalize the k weights for this row to sum to 1, if requested.
+    if (renormalize) {
+        if (thread_group_idx == 0)
+        {
+            const float denom = selected_sum > 0.f ? selected_sum : 1.f;
+            for (int k_idx = 0; k_idx < k; ++k_idx)
+            {
+                const int idx = k * thread_row + k_idx;
+                output[idx] = output[idx] / denom;
+            }
+        }
+    }
+}
+
+namespace detail
+{
+// Constructs some constants needed to partition the work across threads at compile time.
+template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename InputType>
+struct TopkConstants
+{
+    static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
+    static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0, "");
+    static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
+    static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+    static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+    static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
+};
+} // namespace detail
+
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
+void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
+    int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
+    const float* bias, cudaStream_t stream)
+{
+    static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
+    using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
+    static constexpr int VPT = Constants::VPT;
+    static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+    const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+    const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+    dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
+    topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
+        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias);
+}
+
+#ifndef USE_ROCM
+  #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                   \
+    static_assert(WARP_SIZE == 32,                                            \
+                  "Unsupported warp size. Only 32 is supported for CUDA");    \
+    topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, MAX_BYTES, \
+                             IndType, InputType, SF>(                         \
+        gating_output, nullptr, topk_weights, topk_indices,                   \
+        token_expert_indices, num_tokens, topk, 0, num_experts, renormalize,  \
+        bias, stream);
+#else
+  #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                    \
+    if (WARP_SIZE == 64) {                                                     \
+      topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, MAX_BYTES,       \
+                               IndType, InputType, SF>(                        \
+          gating_output, nullptr, topk_weights, topk_indices,                  \
+          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
+          bias, stream);                                                       \
+    } else if (WARP_SIZE == 32) {                                              \
+      topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES,       \
+                               IndType, InputType, SF>(                        \
+          gating_output, nullptr, topk_weights, topk_indices,                  \
+          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
+          bias, stream);                                                       \
+    } else {                                                                   \
+      assert(false &&                                                          \
+             "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
+    }
+#endif
+
+template <typename IndType, typename InputType, ScoringFunc SF>
+void topkGatingKernelLauncher(
+    const InputType* gating_output,
+    float* topk_weights,
+    IndType* topk_indices,
+    int* token_expert_indices,
+    float* workspace,
+    const int num_tokens,
+    const int num_experts,
+    const int topk,
+    const bool renormalize,
+    const float* bias,
+    cudaStream_t stream) {
+    static constexpr int WARPS_PER_TB = 4;
+    static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
+#ifndef USE_ROCM
+    // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
+    // elements can be loaded by a warp
+    static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
+    (std::is_same_v<InputType, __nv_bfloat16> || std::is_same_v<InputType, __half>) ? 4 : 8;
+#endif
+    switch (num_experts) {
+        case 1:
+            LAUNCH_TOPK(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 2:
+            LAUNCH_TOPK(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 4:
+            LAUNCH_TOPK(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 8:
+            LAUNCH_TOPK(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 16:
+            LAUNCH_TOPK(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 32:
+            LAUNCH_TOPK(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 64:
+            LAUNCH_TOPK(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 128:
+            LAUNCH_TOPK(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 256:
+            LAUNCH_TOPK(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        case 512:
+            LAUNCH_TOPK(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            break;
+        // (CUDA only) support multiples of 64 when num_experts is not power of 2.
+        // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of num_experts,
+        // alternatively we can test 4 bytes loading and enable it in future.
+#ifndef USE_ROCM
+        case 192:
+            LAUNCH_TOPK(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            break;
+        case 320:
+            LAUNCH_TOPK(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            break;
+        case 384:
+            LAUNCH_TOPK(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            break;
+        case 448:
+            LAUNCH_TOPK(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            break;
+        case 576:
+            LAUNCH_TOPK(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            break;
+#endif
+        default: {
+            STD_TORCH_CHECK(workspace != nullptr,
+                "workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
+            static constexpr int TPB = 256;
+            if constexpr (SF == SCORING_SOFTMAX) {
+              moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
+                gating_output, nullptr, workspace, num_experts);
+            } else if constexpr (SF == SCORING_SIGMOID) {
+              moeSigmoid<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
+                gating_output, nullptr, workspace, num_experts);
+            } else {
+                STD_TORCH_CHECK(false, "Unsupported scoring func");
+            }
+            moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
+                workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
+                num_experts, topk, 0, num_experts, renormalize, bias);
+        }
+    }
+}
+
+} // namespace moe
+} // namespace vllm
+
+
+template<typename ComputeType, vllm::moe::ScoringFunc SF>
+void dispatch_topk_launch(
+    torch::stable::Tensor& gating_output,
+    torch::stable::Tensor& topk_weights,
+    torch::stable::Tensor& topk_indices,
+    torch::stable::Tensor& token_expert_indices,
+    torch::stable::Tensor& softmax_workspace,
+    int num_tokens, int num_experts, int topk, bool renormalize,
+    std::optional<torch::stable::Tensor> bias,
+    cudaStream_t stream)
+ {
+    const float* bias_ptr = nullptr;
+    if (bias.has_value()) {
+      const torch::stable::Tensor& bias_tensor = bias.value();
+      STD_TORCH_CHECK(bias_tensor.scalar_type() == torch::headeronly::ScalarType::Float,
+                      "bias tensor must be float32");
+      STD_TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
+      STD_TORCH_CHECK(bias_tensor.size(0) == num_experts,
+                      "bias size mismatch, expected: ", num_experts);
+      STD_TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
+      bias_ptr = bias_tensor.const_data_ptr<float>();
+    }
+
+    if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
+        vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
+            reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
+            topk_weights.mutable_data_ptr<float>(),
+            topk_indices.mutable_data_ptr<int>(),
+            token_expert_indices.mutable_data_ptr<int>(),
+            softmax_workspace.mutable_data_ptr<float>(),
+            num_tokens, num_experts, topk, renormalize,
+            bias_ptr, stream);
+    } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
+        vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
+            reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
+            topk_weights.mutable_data_ptr<float>(),
+            topk_indices.mutable_data_ptr<uint32_t>(),
+            token_expert_indices.mutable_data_ptr<int>(),
+            softmax_workspace.mutable_data_ptr<float>(),
+            num_tokens, num_experts, topk, renormalize,
+            bias_ptr, stream);
+    } else {
+        STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
+        vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
+            reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
+            topk_weights.mutable_data_ptr<float>(),
+            topk_indices.mutable_data_ptr<int64_t>(),
+            token_expert_indices.mutable_data_ptr<int>(),
+            softmax_workspace.mutable_data_ptr<float>(),
+            num_tokens, num_experts, topk, renormalize,
+            bias_ptr, stream);
+    }
+}
+
+void topk_softmax(
+    torch::stable::Tensor& topk_weights,                // [num_tokens, topk]
+    torch::stable::Tensor& topk_indices,                // [num_tokens, topk]
+    torch::stable::Tensor& token_expert_indices,        // [num_tokens, topk]
+    torch::stable::Tensor& gating_output,               // [num_tokens, num_experts]
+    bool renormalize,
+    std::optional<torch::stable::Tensor> bias)
+{
+    const int num_experts = gating_output.size(-1);
+    const auto num_tokens = gating_output.numel() / num_experts;
+    const int topk = topk_weights.size(-1);
+
+    const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    const bool needs_workspace = !is_pow_2 || num_experts > 256;
+    const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
+
+    torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
+    const cudaStream_t stream =
+        get_current_cuda_stream(gating_output.get_device_index());
+    auto softmax_workspace = torch::stable::new_empty(
+        gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
+
+    if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
+        dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
+        dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+        dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else {
+        STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
+    }
+}
+
+void topk_sigmoid(
+    torch::stable::Tensor& topk_weights,                // [num_tokens, topk]
+    torch::stable::Tensor& topk_indices,                // [num_tokens, topk]
+    torch::stable::Tensor& token_expert_indices,        // [num_tokens, topk]
+    torch::stable::Tensor& gating_output,               // [num_tokens, num_experts]
+    bool renormalize,
+    std::optional<torch::stable::Tensor> bias)
+{
+    const int num_experts = gating_output.size(-1);
+    const auto num_tokens = gating_output.numel() / num_experts;
+    const int topk = topk_weights.size(-1);
+
+    const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    const bool needs_workspace = !is_pow_2 || num_experts > 256;
+    const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
+
+    torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
+    const cudaStream_t stream =
+        get_current_cuda_stream(gating_output.get_device_index());
+    auto workspace = torch::stable::new_empty(
+        gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
+
+    if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
+        dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
+        dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+        dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else {
+        STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
+    }
+}

upstream_ref/ds_vllm/vllm/qwen3_5.py

@@ -0,0 +1,819 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 The vLLM team.
+# Copyright 2025 The Qwen Team.
+# Copyright 2025 The HuggingFace Inc. team.
+# All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+"""Inference-only Qwen3.5 Series compatible with HuggingFace weights."""
+
+import typing
+from collections.abc import Callable, Iterable
+
+import torch
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import (
+    get_pp_group,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import (
+    GemmaRMSNorm as Qwen3_5RMSNorm,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.qwen3_5 import (
+    Qwen3_5Config,
+    Qwen3_5TextConfig,
+)
+from vllm.transformers_utils.configs.qwen3_5_moe import (
+    Qwen3_5MoeConfig,
+    Qwen3_5MoeTextConfig,
+)
+
+from .interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    MultiModalEmbeddings,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsPP,
+    _require_is_multimodal,
+)
+from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
+from .qwen3_next import (
+    Qwen3NextAttention,
+    Qwen3NextDecoderLayer,
+    Qwen3NextModel,
+    Qwen3NextSparseMoeBlock,
+    QwenNextMixtureOfExperts,
+)
+from .qwen3_vl import (
+    Qwen3_VisionTransformer,
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMultiModalProcessor,
+    Qwen3VLProcessingInfo,
+)
+from .utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    _merge_multimodal_embeddings,
+    extract_layer_index,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+
+class Qwen3_5ProcessingInfo(Qwen3VLProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config(Qwen3_5Config)
+
+
+class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config(Qwen3_5MoeConfig)
+
+
+class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        layer_type: str,
+        prefix: str = "",
+    ) -> None:
+        super(Qwen3NextDecoderLayer, self).__init__()
+
+        config = vllm_config.model_config.hf_text_config
+        model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        self.layer_type = layer_type
+        self.layer_idx = extract_layer_index(prefix)
+
+        if self.layer_type == "linear_attention":
+            self.linear_attn = QwenGatedDeltaNetAttention(
+                config=config,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.linear_attn",
+                gqa_interleaved_layout=False,
+            )
+        elif self.layer_type == "full_attention":
+            self.self_attn = Qwen3NextAttention(
+                config,
+                model_config=model_config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+            )
+        else:
+            raise ValueError(f"Invalid layer_type {self.layer_type}")
+
+        # NOTE: Determine the MLP type based on the model type
+        # Qwen3.5 use all layers for MLP / Qwen3.5-MoE use sparse MoE blocks
+        if config.model_type == "qwen3_5_moe_text":
+            self.mlp = Qwen3NextSparseMoeBlock(
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.mlp",
+            )
+        elif config.model_type == "qwen3_5_text":
+            self.mlp = Qwen3NextMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            raise ValueError(f"Invalid model_type {config.model_type}")
+
+        self.input_layernorm = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        self.layer_scale = getattr(config, "layer_scale", False)
+        if self.layer_scale:
+            self.attn_layer_scale = torch.nn.Parameter(
+                torch.zeros(
+                    1,
+                    1,
+                    config.hidden_size,
+                ),
+            )
+            self.ffn_layer_scale = torch.nn.Parameter(
+                torch.zeros(
+                    1,
+                    1,
+                    config.hidden_size,
+                ),
+            )
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    }
+)
+class Qwen3_5Model(Qwen3NextModel):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super(Qwen3NextModel, self).__init__()
+
+        config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = (
+            vllm_config.model_config.hf_text_config
+        )
+        parallel_config = vllm_config.parallel_config
+
+        eplb_config = parallel_config.eplb_config
+        self.num_redundant_experts = eplb_config.num_redundant_experts
+
+        self.config = config
+
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+        )
+
+        def get_layer(prefix: str):
+            return Qwen3_5DecoderLayer(
+                vllm_config,
+                layer_type=config.layer_types[extract_layer_index(prefix)],
+                prefix=prefix,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
+        )
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+
+        if get_pp_group().is_last_rank:
+            self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.aux_hidden_state_layers: tuple[int, ...] = ()
+
+    def load_fused_expert_weights(
+        self,
+        name: str,
+        params_dict: dict,
+        loaded_weight: torch.Tensor,
+        shard_id: str,
+        num_experts: int,
+    ) -> bool:
+        param = params_dict[name]
+        weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
+        loaded_local_expert = False
+        for expert_id in range(num_experts):
+            curr_expert_weight = loaded_weight[expert_id]
+            success = weight_loader(
+                param,
+                curr_expert_weight,
+                name,
+                shard_id=shard_id,
+                expert_id=expert_id,
+                return_success=True,
+            )
+            if success:
+                loaded_local_expert = True
+
+        return loaded_local_expert
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            # GDN
+            ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
+            ("in_proj_qkvz", "in_proj_z", 3),
+            # self attention
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            # mlp
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+            ("in_proj_ba", "in_proj_b", 0),
+            ("in_proj_ba", "in_proj_a", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        expert_params_mapping = self.get_expert_mapping()
+        is_fused_expert = False
+        fused_expert_params_mapping: list[tuple[str, str, int, str]] = []
+        for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_up_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="gate_up_proj",
+            num_experts=1,
+        ):
+            if shard_id == "w3":
+                continue
+            parts = ckpt_name.split(".")
+            fused_expert_params_mapping.append(
+                (f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id)
+            )
+        num_experts = (
+            self.config.num_experts if hasattr(self.config, "num_experts") else 0
+        )
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            if name.startswith("mtp."):
+                continue
+
+            # Remapping the name of FP8 kv-scale.
+            if name.endswith("scale"):
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if "experts.gate_up_proj" in name or "experts.down_proj" in name:
+                    is_fused_expert = True
+                    expert_params_mapping = fused_expert_params_mapping
+
+                if weight_name not in name:
+                    continue
+
+                if "mlp.experts" in name:
+                    continue
+
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                # Skip layers on other devices.
+                if is_pp_missing_parameter(name, self):
+                    continue
+                # name = apply_attn_prefix(name, params_dict)
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                is_expert_weight = False
+                for mapping in expert_params_mapping:
+                    param_name, weight_name, expert_id, shard_id = mapping
+                    if weight_name not in name:
+                        continue
+                    is_expert_weight = True
+                    name_mapped = name.replace(weight_name, param_name)
+                    # Skip layers on other devices.
+                    if is_pp_missing_parameter(name_mapped, self):
+                        continue
+                    if is_fused_expert:
+                        # qwen3.5 no need to transpose
+                        # loaded_weight = loaded_weight.transpose(-1, -2)
+                        if "experts.gate_up_proj" in name:
+                            loaded_weight = loaded_weight.chunk(2, dim=-2)
+                            success_w1 = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight[0],
+                                "w1",
+                                num_experts,
+                            )
+                            success_w3 = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight[1],
+                                "w3",
+                                num_experts,
+                            )
+                            success = success_w1 and success_w3
+                        else:
+                            # down_proj
+                            success = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight,
+                                shard_id,
+                                num_experts,
+                            )
+                        if success:
+                            name = name_mapped
+                            break
+                    else:
+                        # Skip loading extra bias for GPTQ models.
+                        if (
+                            name_mapped.endswith(".bias")
+                            or name_mapped.endswith("_bias")
+                        ) and name_mapped not in params_dict:
+                            continue
+                        param = params_dict[name_mapped]
+                        weight_loader = param.weight_loader
+                        success = weight_loader(
+                            param,
+                            loaded_weight,
+                            name_mapped,
+                            shard_id=shard_id,
+                            expert_id=expert_id,
+                            return_success=True,
+                        )
+                    if success:
+                        name = name_mapped
+                        break
+                else:
+                    if is_expert_weight:
+                        # We've checked that this is an expert weight
+                        # However it's not mapped locally to this rank
+                        # So we simply skip it
+                        continue
+                    # Skip loading extra bias for GPTQ models.
+                    if name.endswith(".bias") and name not in params_dict:
+                        continue
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    if name not in params_dict:
+                        logger.warning_once(
+                            f"Parameter {name} not found in params_dict, skip loading"
+                        )
+                        continue
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Qwen3_5ForCausalLMBase(
+    nn.Module,
+    HasInnerState,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsPP,
+):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        # GDN fused projections.
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        config = vllm_config.model_config.hf_text_config
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
+
+        scheduler_config = vllm_config.scheduler_config
+        if cache_config.mamba_cache_mode == "all":
+            raise NotImplementedError(
+                "Qwen3.5 currently does not support 'all' prefix caching, "
+                "please use '--mamba-cache-mode=align' instead"
+            )
+        self.quant_config = vllm_config.quant_config
+
+        super().__init__()
+        self.config = config
+        self.scheduler_config = scheduler_config
+        self.model = Qwen3_5Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+        if get_pp_group().is_last_rank:
+            if config.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
+        self.model.aux_hidden_state_layers = layers
+
+    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
+        num_layers = len(self.model.layers)
+        return (2, num_layers // 2, num_layers - 3)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ):
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
+
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["mtp."],
+        )
+        return loader.load_weights(weights)
+
+
+class Qwen3_5ForCausalLM(Qwen3_5ForCausalLMBase):
+    pass
+
+
+class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+
+        # set MoE hyperparameters
+        self.set_moe_parameters()
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return self.model.get_expert_mapping()
+
+
+########################################################
+# Qwen3_5-Dense
+########################################################
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen3_5ProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid):
+    # Qwen3.5 does not support multimodal pruning (EVS).
+    supports_multimodal_pruning = False
+
+    packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | {
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
+        # protocols have not __init__ method, so we need to use nn.Module.__init__
+        nn.Module.__init__(self)
+        config: Qwen3_5Config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        # Qwen3.5 does not support multimodal pruning (EVS).
+        self.is_multimodal_pruning_enabled = False
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3_VisionTransformer(
+                config.vision_config,
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen3_5ForCausalLM(
+                vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids,
+            self.language_model.embed_input_ids,
+            is_multimodal=is_multimodal,
+        )
+
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        is_multimodal = _require_is_multimodal(is_multimodal)
+
+        inputs_embeds = _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+        return inputs_embeds
+
+    def recompute_mrope_positions(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Qwen3.5 does not support multimodal pruning (EVS). "
+            "recompute_mrope_positions should never be called."
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        """Run forward pass for Qwen3.5.
+
+        Args:
+            input_ids: Flattened (concatenated) input_ids corresponding to a
+                batch.
+            positions: Flattened (concatenated) position ids corresponding to a
+                batch.
+                **NOTE**: If mrope is enabled (default setting for Qwen3VL
+                opensource models), the shape will be `(3, seq_len)`,
+                otherwise it will be `(seq_len,).
+            intermediate_tensors: Intermediate tensors from previous pipeline
+                stages.
+            inputs_embeds: Pre-computed input embeddings.
+            **kwargs: Additional keyword arguments including:
+                - pixel_values: Pixel values to be fed to a model.
+                    `None` if no images are passed.
+                - image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in
+                    LLM. `None` if no images are passed.
+                - pixel_values_videos: Pixel values of videos to be fed to a
+                    model. `None` if no videos are passed.
+                - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in
+                    LLM. `None` if no videos are passed.
+        """
+
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        hidden_states = self.language_model.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["mtp."],
+        )
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: "VllmConfig",
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: "VllmConfig"
+    ) -> tuple[tuple[int, int], tuple[int, int]]:
+        parallel_config = vllm_config.parallel_config
+        hf_config = vllm_config.model_config.hf_text_config
+        tp_size = parallel_config.tensor_parallel_size
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+            tp_size,
+            hf_config.linear_num_key_heads,
+            hf_config.linear_num_value_heads,
+            hf_config.linear_key_head_dim,
+            hf_config.linear_value_head_dim,
+            hf_config.linear_conv_kernel_dim,
+            num_spec,
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
+        return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
+
+
+########################################################
+# Qwen3_5-MoE
+########################################################
+
+
+class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ) -> None:
+        assert self.num_local_physical_experts == num_local_physical_experts
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+        for layer in self.language_model.model.layers:
+            if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
+                moe = layer.mlp
+                moe.n_local_physical_experts = num_local_physical_experts
+                moe.n_physical_experts = num_physical_experts
+                moe.n_redundant_experts = self.num_redundant_experts
+                moe.experts.update_expert_map()
+
+    def set_moe_parameters(self):
+        self.expert_weights = []
+
+        self.moe_layers = []
+        example_moe = None
+        for layer in self.language_model.model.layers:
+            if isinstance(layer, Qwen3_5DecoderLayer) and isinstance(
+                layer.mlp, Qwen3NextSparseMoeBlock
+            ):
+                example_moe = layer.mlp
+                self.moe_layers.append(layer.mlp.experts)
+
+        if example_moe is None:
+            raise RuntimeError(
+                "No Qwen3_5 layer found in the language_model.model.layers."
+            )
+
+        # Set MoE hyperparameters
+        self.num_moe_layers = len(self.moe_layers)
+        self.num_expert_groups = 1
+        self.num_shared_experts = 0
+        self.num_logical_experts = example_moe.n_logical_experts
+        self.num_physical_experts = example_moe.n_physical_experts
+        self.num_local_physical_experts = example_moe.n_local_physical_experts
+        self.num_routed_experts = example_moe.n_routed_experts
+        self.num_redundant_experts = example_moe.n_redundant_experts
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen3_5MoeProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class Qwen3_5MoeForConditionalGeneration(
+    Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts
+):
+    # For MoE LoRA weights loading
+    is_3d_moe_weight: bool = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
+        # protocols have not __init__ method, so we need to use nn.Module.__init__
+        nn.Module.__init__(self)
+        config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        # Qwen3.5 does not support multimodal pruning (EVS).
+        self.is_multimodal_pruning_enabled = False
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3_VisionTransformer(
+                config.vision_config,
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen3_5MoeForCausalLM(
+                vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+        # set MoE hyperparameters
+        self.set_moe_parameters()

upstream_ref/xllm/kernels/cuda/activation.cu

@@ -0,0 +1,183 @@
+/* Copyright 2025 The vLLM Authors and The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <c10/cuda/CUDAGuard.h>
+#include <torch/cuda.h>
+
+#include "cuda_ops_api.h"
+
+// ref to:
+// https://github.com/vllm-project/vllm/blob/main/csrc/activation_kernels.cu
+
+namespace {
+
+// Use read-only cache load for CUDA kernels.
+#define XLLM_LDG(arg) __ldg(arg)
+
+template <typename scalar_t,
+          scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
+__device__ __forceinline__ scalar_t compute(const scalar_t& x,
+                                            const scalar_t& y) {
+  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
+}
+
+// Check if pointer is 16-byte aligned for int4 vectorized access
+__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
+  return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
+}
+
+// Activation and gating kernel template with 128-bit vectorized access
+// optimization.
+template <typename scalar_t,
+          scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
+__global__ void act_and_mul_kernel(
+    scalar_t* __restrict__ out,          // [..., d]
+    const scalar_t* __restrict__ input,  // [..., 2, d]
+    const int d) {
+  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* x_ptr = input + token_idx * 2 * d;
+  const scalar_t* y_ptr = x_ptr + d;
+  scalar_t* out_ptr = out + token_idx * d;
+
+  // Check alignment for 128-bit vectorized access.
+  // All three pointers must be 16-byte aligned for safe int4 operations.
+  const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
+                       is_16byte_aligned(out_ptr);
+
+  if (aligned && d >= VEC_SIZE) {
+    // Fast path: 128-bit vectorized loop
+    const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
+    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
+    const int num_vecs = d / VEC_SIZE;
+    const int vec_end = num_vecs * VEC_SIZE;
+
+    for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
+      int4 x = XLLM_LDG(&x_vec[i]), y = XLLM_LDG(&y_vec[i]), r;
+      auto* xp = reinterpret_cast<scalar_t*>(&x);
+      auto* yp = reinterpret_cast<scalar_t*>(&y);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
+#pragma unroll
+      for (int j = 0; j < VEC_SIZE; j++) {
+        rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
+      }
+      out_vec[i] = r;
+    }
+    // Scalar cleanup for remaining elements
+    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
+      out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(XLLM_LDG(&x_ptr[i]),
+                                                        XLLM_LDG(&y_ptr[i]));
+    }
+  } else {
+    // Scalar fallback for unaligned data or small d
+    for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+      const scalar_t x = XLLM_LDG(&x_ptr[idx]);
+      const scalar_t y = XLLM_LDG(&y_ptr[idx]);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
+    }
+  }
+}
+
+template <typename T>
+__device__ __forceinline__ T silu_kernel(const T& x) {
+  // x * sigmoid(x)
+  return (T)(((float)x) / (1.0f + expf((float)-x)));
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_kernel(const T& x) {
+  // Equivalent to PyTorch GELU with 'none' approximation.
+  // Refer to:
+  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
+  const float f = (float)x;
+  constexpr float ALPHA = M_SQRT1_2;
+  return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
+  // Equivalent to PyTorch GELU with 'tanh' approximation.
+  // Refer to:
+  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
+  const float f = (float)x;
+  constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
+  constexpr float KAPPA = 0.044715;
+  float x_cube = f * f * f;
+  float inner = BETA * (f + KAPPA * x_cube);
+  return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
+}
+
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                   \
+  int d = input.size(-1) / 2;                                              \
+  int64_t num_tokens = input.numel() / input.size(-1);                     \
+  dim3 grid(num_tokens);                                                   \
+  dim3 block(std::min(d, 1024));                                           \
+  if (num_tokens == 0) {                                                   \
+    return;                                                                \
+  }                                                                        \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));        \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();            \
+  DISPATCH_FLOATING_TYPES(input.scalar_type(), "act_and_mul_kernel", [&] { \
+    act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>              \
+        <<<grid, block, 0, stream>>>(                                      \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);      \
+  });
+
+void silu_and_mul(torch::Tensor out,    // [..., d]
+                  torch::Tensor input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(silu_kernel, true);
+}
+
+void gelu_and_mul(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(gelu_kernel, true);
+}
+
+void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
+                       torch::Tensor& input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(gelu_tanh_kernel, true);
+}
+}  // namespace
+
+namespace xllm::kernel::cuda {
+
+void act_and_mul(torch::Tensor out,
+                 torch::Tensor input,
+                 const std::string& act_mode) {
+  if (act_mode != "silu" && act_mode != "gelu" && act_mode != "gelu_tanh") {
+    LOG(FATAL) << "Unsupported act mode: " << act_mode
+               << ", only support silu, gelu, gelu_tanh";
+  }
+
+  // flashinfer act_and_mul ops
+  // std::string uri = act_mode + "_and_mul";
+  // FunctionFactory::get_instance().act_and_mul(uri).call(
+  //     out, input, support_pdl());
+
+  if (act_mode == "silu") {
+    silu_and_mul(out, input);
+  } else if (act_mode == "gelu") {
+    gelu_and_mul(out, input);
+  } else if (act_mode == "gelu_tanh") {
+    gelu_tanh_and_mul(out, input);
+  }
+}
+
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/device_utils.cuh

@@ -0,0 +1,80 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <cub/cub.cuh>
+
+namespace xllm::kernel::cuda {
+
+#define WARP_SIZE 32
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+// Aligned array type
+template <typename T,
+          // Number of elements in the array
+          int N,
+          // Alignment requirement in bytes
+          int Alignment = sizeof(T) * N>
+class alignas(Alignment) AlignedArray {
+  T data[N];
+};
+
+#define XLLM_SHFL_XOR_SYNC(mask, var, lane_mask) \
+  __shfl_xor_sync((mask), (var), (lane_mask))
+#define XLLM_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \
+  __shfl_xor_sync((mask), (var), (lane_mask), (width))
+
+// Define reduction operators based on CUDA version
+// CUDA 13 (12.9+) deprecated cub::Max/Min in favor of cuda::maximum/minimum
+#if CUDA_VERSION >= 12090
+using MaxReduceOp = ::cuda::maximum<>;
+using MinReduceOp = ::cuda::minimum<>;
+#else
+using MaxReduceOp = cub::Max;
+using MinReduceOp = cub::Min;
+#endif
+
+template <typename T>
+__device__ float convert_to_float(T x) {
+  if constexpr (std::is_same_v<T, __half>) {
+    return __half2float(x);
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+    return __bfloat162float(x);
+  } else if constexpr (std::is_same_v<T, float>) {
+    return x;
+  } else {
+    return static_cast<float>(x);
+  }
+}
+
+// Constructs some constants needed to partition the work across threads at
+// compile time.
+template <typename T, int EXPERTS, int BYTES_PER_LDG>
+struct TopkConstants {
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 ||
+                    EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0,
+                "");
+  static constexpr int VECs_PER_THREAD =
+      MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+  static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+};
+
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/moe/fused_moe.cpp

@@ -0,0 +1,123 @@
+/* Copyright 2026 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "kernels/cuda/cuda_ops_api.h"
+#include "kernels/cuda/utils.h"
+#include "platform/device.h"
+
+namespace xllm::kernel::cuda {
+
+torch::Tensor cutlass_fused_moe(
+    const torch::Tensor& input,                   // [num_tokens, hidden]
+    const torch::Tensor& token_selected_experts,  // [num_tokens, top_k]
+    const torch::Tensor& token_final_scales,      // [num_tokens, top_k]
+    const torch::Tensor&
+        fc1_expert_weights,  // [num_experts, inter_dim, hidden]
+    const torch::Tensor&
+        fc2_expert_weights,  // [num_experts, hidden, inter_dim]
+    torch::ScalarType output_dtype,
+    const std::vector<torch::Tensor>& quant_scales,
+    int32_t tp_size,
+    int32_t tp_rank,
+    int32_t ep_size,
+    int32_t ep_rank,
+    int32_t cluster_size,
+    int32_t cluster_rank,
+    const std::optional<torch::Tensor>& fc1_expert_biases,
+    const std::optional<torch::Tensor>& fc2_expert_biases,
+    const std::optional<torch::Tensor>& input_sf,
+    const std::optional<torch::Tensor>& swiglu_alpha,
+    const std::optional<torch::Tensor>& swiglu_beta,
+    const std::optional<torch::Tensor>& swiglu_limit,
+    const std::optional<torch::Tensor>& output,
+    bool enable_alltoall,
+    bool use_deepseek_fp8_block_scale,
+    bool use_w4_group_scaling,
+    bool use_mxfp8_act_scaling,
+    bool min_latency_mode,
+    bool use_packed_weights,
+    int32_t tune_max_num_tokens,
+    ActivationType activation_type) {
+  int64_t num_rows = input.size(0);
+  int64_t hidden_size = fc2_expert_weights.size(1);
+
+  if (min_latency_mode) {
+    num_rows *= fc2_expert_weights.size(0);
+  }
+
+  std::vector<int64_t> output_shape = {num_rows, hidden_size};
+  torch::Tensor result_output;
+  if (output.has_value() && output.value().defined()) {
+    result_output = output.value();
+  } else {
+    torch::TensorOptions options = input.options().dtype(output_dtype);
+    result_output = torch::empty(output_shape, options);
+  }
+
+  std::string fused_moe_uri = "fused_moe";
+  if (Device::is_support_sm90a()) {
+    fused_moe_uri += "_90";
+  } else if (Device::is_support_sm100a() || Device::is_support_sm100f()) {
+    fused_moe_uri += "_100";
+  } else if (Device::is_support_sm120a()) {
+    fused_moe_uri += "_120";
+  } else {
+    LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120.";
+  }
+
+  bind_tvmffi_stream_to_current_torch_stream(input.device());
+
+  ffi::Module fused_moe_runner =
+      get_function(fused_moe_uri, "init")(
+          to_dl_data_type(input.scalar_type()),
+          to_dl_data_type(fc1_expert_weights.scalar_type()),
+          to_dl_data_type(output_dtype),
+          use_deepseek_fp8_block_scale,
+          use_w4_group_scaling,
+          use_mxfp8_act_scaling,
+          use_packed_weights)
+          .cast<ffi::Module>();
+
+  fused_moe_runner->GetFunction("run_moe").value()(
+      to_ffi_tensor(result_output),
+      to_ffi_tensor(input),
+      to_ffi_tensor(token_selected_experts),
+      to_ffi_optional_tensor(token_final_scales),
+      to_ffi_tensor(fc1_expert_weights),
+      to_ffi_optional_tensor(fc1_expert_biases),
+      to_ffi_tensor(fc2_expert_weights),
+      to_ffi_optional_tensor(fc2_expert_biases),
+      to_ffi_optional_array_tensors(quant_scales),
+      to_ffi_optional_tensor(input_sf),
+      to_ffi_optional_tensor(swiglu_alpha),
+      to_ffi_optional_tensor(swiglu_beta),
+      to_ffi_optional_tensor(swiglu_limit),
+      tp_size,
+      tp_rank,
+      ep_size,
+      ep_rank,
+      cluster_size,
+      cluster_rank,
+      enable_alltoall,
+      min_latency_mode,
+      /*profile_ids=*/ffi::Optional<ffi::Array<int64_t>>(),  // TODO: support
+                                                             // auto tuning
+                                                             // profile ids
+      support_pdl(),
+      activation_type);
+
+  return result_output;
+}
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/moe/moe_fused_topk.cu

@@ -0,0 +1,56 @@
+/* Copyright 2026 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "kernels/cuda/cuda_ops_api.h"
+#include "moe_topk_sigmoid_kernels.cuh"
+#include "moe_topk_softmax_kernels.cuh"
+
+namespace xllm::kernel::cuda {
+
+std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
+    torch::Tensor& gating_output,
+    int64_t topk,
+    bool renormalize,
+    const std::optional<torch::Tensor>& correction_bias,
+    const std::string& scoring_func) {
+  int64_t num_tokens = gating_output.size(0);
+
+  torch::Tensor topk_weights = torch::empty(
+      {num_tokens, topk},
+      torch::dtype(torch::kFloat32).device(gating_output.device()));
+  torch::Tensor topk_ids =
+      torch::empty({num_tokens, topk},
+                   torch::dtype(torch::kInt32).device(gating_output.device()));
+
+  if (scoring_func == "softmax") {
+    std::optional<torch::Tensor> none_correction_bias = std::nullopt;
+    topk_softmax(topk_weights,
+                 topk_ids,
+                 gating_output,
+                 renormalize,
+                 /*moe_softcapping=*/0.0,
+                 none_correction_bias);
+  } else if (scoring_func == "sigmoid") {
+    topk_sigmoid(
+        topk_weights, topk_ids, gating_output, renormalize, correction_bias);
+  } else {
+    LOG(FATAL) << "Unsupported scoring function for moe topk: " << scoring_func
+               << "only softmax and sigmoid are supported";
+  }
+
+  return std::make_tuple(topk_weights, topk_ids);
+}
+
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/moe/moe_topk.cuh

@@ -0,0 +1,285 @@
+
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// refers to
+// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
+
+#pragma once
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+#include <cub/cub.cuh>
+
+#include "core/kernels/cuda/arch_condition.h"
+
+namespace xllm::kernel::cuda {
+namespace reduce_topk {
+namespace cg = cooperative_groups;
+static constexpr int kWARP_SIZE = 32;
+static constexpr bool kTLLM_GEN_HAS_FAST_REDUX = arch::is_major_v<10>;
+
+template <typename T_>
+struct TopKRedType {
+  using T = T_;
+  static_assert(
+      std::is_same_v<T, float> || std::is_same_v<T, half> ||
+          std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
+      "Top K reduction only implemented for int, float, float16 and bfloat16");
+
+  using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
+  using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
+
+  static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
+  static constexpr int kMaxIdx = 65535;
+  TypeCmp compValIdx;
+
+  static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
+    auto valueBits = cub::Traits<T>::TwiddleIn(
+        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
+    TypeCmp compactTmp = valueBits;
+    compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
+    // Use 65535 minus idx to give higher priority to elements with smaller
+    // indices.
+    return compactTmp;
+  }
+
+  static __host__ __device__ void unpack(T& value,
+                                         int32_t& index,
+                                         TypeCmp cmp) {
+    // Since “65535-idx” is always smaller than 65536 and positive, we can
+    // directly use it as the lower 16 bits
+    index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
+
+    auto compactTmp = cmp >> kMoveBits;
+    auto valueBits = cub::Traits<T>::TwiddleOut(
+        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
+    value = reinterpret_cast<T&>(valueBits);
+  }
+
+  __host__ __device__ TopKRedType() = default;
+
+  __host__ __device__ TopKRedType(T val, int32_t idx)
+      : compValIdx(makeCmpVal(val, idx)) {}
+
+  __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
+
+  __device__ inline TypeCmp reduce(
+      cg::thread_block_tile<kWARP_SIZE> const& warp) {
+    if constexpr (!kTLLM_GEN_HAS_FAST_REDUX || sizeof(TypeCmp) == 8) {
+      return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
+    } else {
+      TypeCmp result;
+      asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
+          : "=r"(result)
+          : "r"(compValIdx));
+      return result;
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int K_, bool Enable_>
+struct TopKIdx {
+  // by default, empty
+};
+
+template <int K_>
+struct TopKIdx<K_, true> {
+  static constexpr int K = K_;
+  int32_t val[K];
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define TOPK_SWAP(I, J)                                         \
+  {                                                             \
+    auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
+    auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
+    topK[I].compValIdx = pairMax;                               \
+    topK[J].compValIdx = pairMin;                               \
+  }
+
+template <int N, typename RedType>
+struct Sort;
+
+template <typename RedType>
+struct Sort<1, RedType> {
+  static __device__ void run(RedType* topK) {}
+};
+
+template <typename RedType>
+struct Sort<2, RedType> {
+  static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
+};
+
+template <typename RedType>
+struct Sort<3, RedType> {
+  static __device__ void run(RedType* topK) {
+    TOPK_SWAP(0, 1);
+    TOPK_SWAP(1, 2);
+    TOPK_SWAP(0, 1);
+  }
+};
+
+template <typename RedType>
+struct Sort<4, RedType> {
+  static __device__ void run(RedType* topK) {
+    TOPK_SWAP(0, 2);
+    TOPK_SWAP(1, 3);
+    TOPK_SWAP(0, 1);
+    TOPK_SWAP(2, 3);
+    TOPK_SWAP(1, 2);
+  }
+};
+
+template <int K, typename Type>
+__forceinline__ __device__ void reduceTopK(
+    cg::thread_block_tile<kWARP_SIZE> const& warp,
+    Type (&out)[K],
+    int32_t (&outIdx)[K],
+    Type value,
+    int32_t idx,
+    Type const minValue,
+    int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  using RedType = TopKRedType<Type>;
+  RedType topK{value, idx};
+  typename RedType::TypeCmp packedMax{};
+#pragma unroll
+  for (int kk = 0; kk < actualK; ++kk)  //@todo: check if actualK is correct
+  {
+    topK =
+        kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
+    // get the next largest value
+    packedMax = topK.reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
+  }
+};
+
+template <int K, typename Type, int N, bool IsSorted = false>
+__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
+                               Type (&out)[K],
+                               int32_t (&outIdx)[K],
+                               Type (&value)[N],
+                               int32_t (&idx)[N],
+                               Type minValue,
+                               int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(N > 0, "Top K must have N > 0");
+  static_assert(N < 5,
+                "Only support candidates number less than or equal to 128");
+  using RedType = TopKRedType<Type>;
+  RedType topK[N];
+#pragma unroll
+  for (int nn = 0; nn < N; ++nn) {
+    topK[nn] = RedType{value[nn], idx[nn]};
+  }
+
+  if constexpr (!IsSorted) {
+    Sort<N, RedType>::run(topK);
+  }
+  typename RedType::TypeCmp packedMax{};
+#pragma unroll
+  for (int kk = 0; kk < actualK; ++kk) {
+    bool update = kk > 0 && packedMax == topK[0].compValIdx;
+#pragma unroll
+    for (int nn = 0; nn < N; ++nn) {
+      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
+                 : update              ? topK[nn + 1]
+                                       : topK[nn];
+    }
+    // get the next largest value
+    packedMax = topK[0].reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
+  }
+};
+
+template <int K, typename Type, int N>
+__forceinline__ __device__ void reduceTopK(
+    cg::thread_block_tile<kWARP_SIZE> const& warp,
+    Type (&out)[K],
+    int32_t (&outIdx)[K],
+    Type (&value)[N],
+    int32_t (&idx)[N],
+    Type const minValue,
+    int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(N > 0, "Top K must have N > 0");
+  static_assert(
+      N <= 16,
+      "Only support candidates number less than or equal to 16*32=512");
+  static_assert(N <= 4 || N % 4 == 0,
+                "Only support candidates number is a multiple of 4*32=128 or "
+                "less than or equal to 4");
+  using RedType = TopKRedType<Type>;
+
+  if constexpr (N <= 4) {
+    reduceTopKFunc<K, Type, N>(
+        warp, out, outIdx, value, idx, minValue, actualK);
+  } else {
+    constexpr int numLoops = N / 4;
+    constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
+
+    Type topKBufferValue[numResults];
+    int32_t topKBufferIdx[numResults];
+    int32_t laneIdx = threadIdx.x % kWARP_SIZE;
+
+    // Sentinel index must be in [0, kMaxIdx] to survive makeCmpVal pack/unpack
+    // (kMaxIdx - idx is stored in 16 bits; -1 would become 0 and unpack to
+    // 65535). Use kMaxIdx so sentinel slots have smallest compValIdx for
+    // minValue and lose to any real candidate.
+    for (int ii = 0; ii < numResults; ++ii) {
+      topKBufferValue[ii] = minValue;
+      topKBufferIdx[ii] = RedType::kMaxIdx;
+    }
+    for (int loop = 0; loop < numLoops; ++loop) {
+      int start = loop * 4;
+      Type topKValue[K];
+      int32_t topKIdx[K];
+      Type inValue[4];
+      int32_t inIdx[4];
+      for (int i = 0; i < 4; ++i) {
+        inValue[i] = value[start + i];
+        inIdx[i] = idx[start + i];
+      }
+      reduceTopKFunc<K, Type, 4>(
+          warp, topKValue, topKIdx, inValue, inIdx, minValue, actualK);
+      int inOffset = laneIdx % K;
+      if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
+        topKBufferValue[0] = topKValue[inOffset];
+        topKBufferIdx[0] = topKIdx[inOffset];
+      }
+      if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
+        topKBufferValue[1] = topKValue[inOffset];
+        topKBufferIdx[1] = topKIdx[inOffset];
+      }
+    }
+
+    reduceTopKFunc<K, Type, numResults>(
+        warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK);
+  }
+};
+
+#undef TOPK_SWAP
+
+}  // namespace reduce_topk
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/moe/moe_topk_sigmoid_kernels.cuh

@@ -0,0 +1,602 @@
+// Adapt from
+// https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu
+// which is originally adapted from
+// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
+/* 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 <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include <cub/util_type.cuh>
+#include <cuda/functional>
+
+#include "kernels/cuda/device_utils.cuh"
+
+namespace {
+
+using namespace xllm::kernel::cuda;
+
+// ====================== Sigmoid things ===============================
+// We have our own implementation of sigmoid here so we can support transposing
+// the output in the sigmoid kernel when we extend this module to support
+// expert-choice routing.
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__
+    void moe_sigmoid(const T* input,
+                     const bool* finished,
+                     float* output,
+                     const int num_cols,
+                     const float* correction_bias) {
+  const int thread_row_offset = blockIdx.x * num_cols;
+
+  // Don't touch finished rows.
+  if ((finished != nullptr) && finished[blockIdx.x]) {
+    return;
+  }
+
+  // First pass: Apply transformation, find max, and write transformed values to
+  // output
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    float val = convert_to_float<T>(input[idx]);
+
+    val = 1.0f / (1.0f + expf(-val));
+
+    // Apply correction bias if provided
+    if (correction_bias != nullptr) {
+      val = val + correction_bias[ii];
+    }
+
+    output[idx] = val;  // Store transformed value
+  }
+}
+
+template <int TPB>
+__launch_bounds__(TPB) __global__
+    void moe_topK(const float* inputs_after_sigmoid,
+                  const bool* finished,
+                  float* output,
+                  int* indices,
+                  const int num_experts,
+                  const int k,
+                  const int start_expert,
+                  const int end_expert,
+                  const bool renormalize,
+                  const float* correction_bias) {
+  using cub_kvp = cub::KeyValuePair<int, float>;
+  using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  cub_kvp thread_kvp;
+  cub::ArgMax arg_max;
+
+  const int block_row = blockIdx.x;
+
+  const bool row_is_active = finished ? !finished[block_row] : true;
+  const int thread_read_offset = blockIdx.x * num_experts;
+  float row_sum_for_renormalize = 0;
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    thread_kvp.key = 0;
+    thread_kvp.value = -1.f;  // This is OK because inputs are probabilities
+
+    cub_kvp inp_kvp;
+    for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
+      const int idx = thread_read_offset + expert;
+      inp_kvp.key = expert;
+      inp_kvp.value = inputs_after_sigmoid[idx];
+
+      for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
+        const int prior_winning_expert = indices[k * block_row + prior_k];
+
+        if (prior_winning_expert == expert) {
+          inp_kvp = thread_kvp;
+        }
+      }
+
+      thread_kvp = arg_max(inp_kvp, thread_kvp);
+    }
+
+    const cub_kvp result_kvp =
+        BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+    if (threadIdx.x == 0) {
+      // Ignore experts the node isn't responsible for with expert parallelism
+      const int expert = result_kvp.key;
+      const bool node_uses_expert =
+          expert >= start_expert && expert < end_expert;
+      const bool should_process_row = row_is_active && node_uses_expert;
+
+      const int idx = k * block_row + k_idx;
+      float val = result_kvp.value;
+      if (correction_bias != nullptr) {
+        val -= correction_bias[expert];
+      }
+      output[idx] = val;
+      indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
+      assert(indices[idx] >= 0);
+      row_sum_for_renormalize += val;
+    }
+    __syncthreads();
+  }
+
+  if (renormalize && threadIdx.x == 0) {
+    float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * block_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+// ====================== TopK sigmoid things ===============================
+
+/*
+  A Top-K gating sigmoid written to exploit when the number of experts in the
+  MoE layers are a small power of 2. This allows us to cleanly share the rows
+  among the threads in a single warp and eliminate communication between warps
+  (so no need to use shared mem).
+
+  It fuses the sigmoid, max and argmax into a single kernel.
+
+  Limitations:
+  1) This implementation is intended for when the number of experts is a small
+  power of 2. 2) This implementation assumes k is small, but will work for any
+  k.
+*/
+
+template <typename T,
+          int VPT,
+          int NUM_EXPERTS,
+          int WARPS_PER_CTA,
+          int BYTES_PER_LDG>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
+    void topk_gating_sigmoid(const T* input,
+                             const bool* finished,
+                             float* output,
+                             const int num_rows,
+                             int* indices,
+                             const int k,
+                             const int start_expert,
+                             const int end_expert,
+                             const bool renormalize,
+                             const float* correction_bias) {
+  // We begin by enforcing compile time assertions and setting up compile time
+  // constants.
+  static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
+  static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
+                "NUM_EXPERTS must be power of 2");
+  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
+                "BYTES_PER_LDG must be power of 2");
+  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+  // Number of bytes each thread pulls in per load
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+  // Restrictions based on previous section.
+  static_assert(
+      VPT % ELTS_PER_LDG == 0,
+      "The elements per thread must be a multiple of the elements per ldg");
+  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+                "The threads per row must cleanly divide the threads per warp");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+                "THREADS_PER_ROW must be power of 2");
+  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+                "THREADS_PER_ROW can be at most warp size");
+
+  // We have NUM_EXPERTS elements per row. We specialize for small #experts
+  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+  // Restrictions for previous section.
+  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+                "The elts per row must cleanly divide the total elt per warp");
+
+  // ===================== From this point, we finally start computing run-time
+  // variables. ========================
+
+  // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
+  // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
+  // rows. We start by computing the start row for each block.
+  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+
+  // Now, using the base row per thread block, we compute the base row per warp.
+  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+
+  // The threads in a warp are split into sub-groups that will work on a row.
+  // We compute row offset for each thread sub-group
+  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row = warp_base_row + thread_row_in_warp;
+
+  // Threads with indices out of bounds should early exit here.
+  if (thread_row >= num_rows) {
+    return;
+  }
+  const bool row_is_active = finished ? !finished[thread_row] : true;
+
+  // We finally start setting up the read pointers for each thread. First, each
+  // thread jumps to the start of the row it will read.
+  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+
+  // Now, we compute the group each thread belong to in order to determine the
+  // first column to start loads.
+  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+  // Determine the pointer type to use to read in the data depending on the
+  // BYTES_PER_LDG template param. In theory, this can support all powers of 2
+  // up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
+  // array here. We defined our own aligned array and use it here to avoid the
+  // dependency on CUTLASS.
+  using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+
+  // Finally, we pull in the data from global mem
+  T row_chunk_temp[VPT];
+  AccessType* row_chunk_vec_ptr =
+      reinterpret_cast<AccessType*>(&row_chunk_temp);
+  const AccessType* vec_thread_read_ptr =
+      reinterpret_cast<const AccessType*>(thread_read_ptr);
+#pragma unroll
+  // Note(Byron): interleaved loads to achieve better memory coalescing
+  // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
+  // thread[2] | thread[3] | ...
+  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+  }
+
+  float row_chunk[VPT];
+#pragma unroll
+  // Note(Byron): upcast logits to float32
+  for (int ii = 0; ii < VPT; ++ii) {
+    float val = convert_to_float<T>(row_chunk_temp[ii]);
+    val = 1.0f / (1.0f + expf(-val));
+    // Apply correction bias if provided
+    if (correction_bias != nullptr) {
+      /*
+      LDG is interleaved
+      |thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG|
+      |--------- group0 --------| |----------group1 --------|
+                                    ^ local2
+      */
+      const int group_id = ii / ELTS_PER_LDG;
+      const int local_id = ii % ELTS_PER_LDG;
+      const int expert_idx = first_elt_read_by_thread +
+                             group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                             local_id;
+      val = val + correction_bias[expert_idx];
+    }
+
+    row_chunk[ii] = val;
+  }
+
+  // Now, row_chunk contains the sigmoid of the row chunk. Now, I want to find
+  // the topk elements in each row, along with the max index.
+  int start_col = first_elt_read_by_thread;
+  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+  float row_sum_for_renormalize = 0;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    // First, each thread does the local argmax
+    float max_val = row_chunk[0];
+    int expert = start_col;
+#pragma unroll
+    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+         ++ldg, col += COLS_PER_GROUP_LDG) {
+#pragma unroll
+      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+
+        // No check on the experts here since columns with the smallest index
+        // are processed first and only updated if > (not >=)
+        if (val > max_val) {
+          max_val = val;
+          expert = col + ii;
+        }
+      }
+    }
+
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads
+// reach consensus about the max. This will be useful for K > 1 so that the
+// threads can agree on "who" had the max value. That thread can then blank out
+// their max with -inf and the warp can run more iterations...
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+      float other_max =
+          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
+      int other_expert =
+          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
+
+      // We want lower indices to "win" in every thread so we break ties this
+      // way
+      if (other_max > max_val ||
+          (other_max == max_val && other_expert < expert)) {
+        max_val = other_max;
+        expert = other_expert;
+      }
+    }
+
+    // Write the max for this k iteration to global memory.
+    if (thread_group_idx == 0) {
+      // Add a guard to ignore experts not included by this node
+      const bool node_uses_expert =
+          expert >= start_expert && expert < end_expert;
+      const bool should_process_row = row_is_active && node_uses_expert;
+
+      // The lead thread from each sub-group will write out the final results to
+      // global memory. (This will be a single) thread per row of the
+      // input/output matrices.
+      const int idx = k * thread_row + k_idx;
+      if (correction_bias != nullptr) {
+        max_val -= correction_bias[expert];
+      }
+      output[idx] = max_val;
+      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+      row_sum_for_renormalize += max_val;
+    }
+
+    // Finally, we clear the value in the thread with the current max if there
+    // is another iteration to run.
+    if (k_idx + 1 < k) {
+      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int thread_to_clear_in_group =
+          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+      // Only the thread in the group which produced the max will reset the
+      // "winning" value to -inf.
+      if (thread_group_idx == thread_to_clear_in_group) {
+        const int offset_for_expert = expert % ELTS_PER_LDG;
+        // Safe to set to any negative value since row_chunk values must be
+        // between 0 and 1.
+        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+            -10000.f;
+      }
+    }
+  }
+
+  // Fuse renormalization of topk_weights into this kernel
+  if (renormalize && thread_group_idx == 0) {
+    float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+#pragma unroll
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * thread_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+template <typename T, int EXPERTS, int WARPS_PER_TB>
+void topk_gating_sigmoid_launcher_helper(const T* input,
+                                         const bool* finished,
+                                         float* output,
+                                         int* indices,
+                                         const int num_rows,
+                                         const int k,
+                                         const int start_expert,
+                                         const int end_expert,
+                                         const bool renormalize,
+                                         const float* correction_bias,
+                                         cudaStream_t stream) {
+  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+
+  static constexpr int BYTES_PER_LDG =
+      MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int VPT = Constants::VPT;
+  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+  dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+  topk_gating_sigmoid<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+      <<<num_blocks, block_dim, 0, stream>>>(input,
+                                             finished,
+                                             output,
+                                             num_rows,
+                                             indices,
+                                             k,
+                                             start_expert,
+                                             end_expert,
+                                             renormalize,
+                                             correction_bias);
+}
+
+#define LAUNCH_SIGMOID(TYPE, NUM_EXPERTS, WARPS_PER_TB)                 \
+  topk_gating_sigmoid_launcher_helper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
+      gating_output,                                                    \
+      nullptr,                                                          \
+      topk_weights,                                                     \
+      topk_indices,                                                     \
+      num_tokens,                                                       \
+      topk,                                                             \
+      0,                                                                \
+      num_experts,                                                      \
+      renormalize,                                                      \
+      correction_bias,                                                  \
+      stream);
+
+template <typename T>
+void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
+                                         float* topk_weights,
+                                         int* topk_indices,
+                                         float* sigmoid_workspace,
+                                         const int num_tokens,
+                                         const int num_experts,
+                                         const int topk,
+                                         const bool renormalize,
+                                         const float* correction_bias,
+                                         cudaStream_t stream) {
+  static constexpr int WARPS_PER_TB = 4;
+  switch (num_experts) {
+    case 1:
+      LAUNCH_SIGMOID(T, 1, WARPS_PER_TB);
+      break;
+    case 2:
+      LAUNCH_SIGMOID(T, 2, WARPS_PER_TB);
+      break;
+    case 4:
+      LAUNCH_SIGMOID(T, 4, WARPS_PER_TB);
+      break;
+    case 8:
+      LAUNCH_SIGMOID(T, 8, WARPS_PER_TB);
+      break;
+    case 16:
+      LAUNCH_SIGMOID(T, 16, WARPS_PER_TB);
+      break;
+    case 32:
+      LAUNCH_SIGMOID(T, 32, WARPS_PER_TB);
+      break;
+    case 64:
+      LAUNCH_SIGMOID(T, 64, WARPS_PER_TB);
+      break;
+    case 128:
+      LAUNCH_SIGMOID(T, 128, WARPS_PER_TB);
+      break;
+    case 256:
+      LAUNCH_SIGMOID(T, 256, WARPS_PER_TB);
+      break;
+    default: {
+      TORCH_CHECK(sigmoid_workspace != nullptr,
+                  "sigmoid_workspace must be provided for num_experts that are "
+                  "not a power of 2.");
+      static constexpr int TPB = 256;
+      moe_sigmoid<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
+                                                          nullptr,
+                                                          sigmoid_workspace,
+                                                          num_experts,
+                                                          correction_bias);
+      moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(sigmoid_workspace,
+                                                    nullptr,
+                                                    topk_weights,
+                                                    topk_indices,
+                                                    num_experts,
+                                                    topk,
+                                                    0,
+                                                    num_experts,
+                                                    renormalize,
+                                                    correction_bias);
+    }
+  }
+}
+}  // namespace
+
+namespace xllm::kernel::cuda {
+void topk_sigmoid(torch::Tensor& topk_weights,   // [num_tokens, topk]
+                  torch::Tensor& topk_indices,   // [num_tokens, topk]
+                  torch::Tensor& gating_output,  // [num_tokens, num_experts]
+                  const bool renormalize,
+                  const std::optional<torch::Tensor>& correction_bias) {
+  // Check data type
+  CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
+        gating_output.scalar_type() == at::ScalarType::Half ||
+        gating_output.scalar_type() == at::ScalarType::BFloat16)
+      << "gating_output must be float32, float16, or bfloat16";
+
+  // Check dimensions
+  CHECK(gating_output.dim() == 2)
+      << "gating_output must be 2D tensor [num_tokens, num_experts]";
+  CHECK(topk_weights.dim() == 2)
+      << "topk_weights must be 2D tensor [num_tokens, topk]";
+  CHECK(topk_indices.dim() == 2)
+      << "topk_indices must be 2D tensor [num_tokens, topk]";
+
+  // Check shapes
+  CHECK(gating_output.size(0) == topk_weights.size(0))
+      << "First dimension of topk_weights must match num_tokens in "
+         "gating_output";
+  CHECK(gating_output.size(0) == topk_indices.size(0))
+      << "First dimension of topk_indices must match num_tokens in "
+         "gating_output";
+  CHECK(topk_weights.size(-1) == topk_indices.size(-1))
+      << "Second dimension of topk_indices must match topk in topk_weights";
+  CHECK(topk_weights.size(-1) <= gating_output.size(-1))
+      << "topk must be less than or equal to num_experts";
+
+  const int num_experts = static_cast<int>(gating_output.size(-1));
+  const int num_tokens = static_cast<int>(gating_output.size(0));
+  const int topk = static_cast<int>(topk_weights.size(-1));
+
+  const bool is_pow_2 =
+      (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+  const bool needs_workspace = !is_pow_2 || num_experts > 256;
+  const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  torch::Tensor sigmoid_workspace = torch::empty(
+      {workspace_size}, gating_output.options().dtype(at::ScalarType::Float));
+
+  const at::ScalarType dtype = gating_output.scalar_type();
+
+  // Validate correction_bias if provided - must always be float32
+  const float* bias_ptr = nullptr;
+  if (correction_bias.has_value()) {
+    const torch::Tensor& bias_tensor = correction_bias.value();
+    CHECK(bias_tensor.dim() == 1)
+        << "correction_bias must be 1D tensor [num_experts]";
+    CHECK(bias_tensor.size(0) == num_experts)
+        << "correction_bias size must match num_experts";
+    CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
+        << "correction_bias must be float32, got " << bias_tensor.scalar_type();
+    bias_ptr = bias_tensor.data_ptr<float>();
+  }
+
+  if (dtype == at::ScalarType::Float) {
+    topk_gating_sigmoid_kernel_launcher<float>(
+        gating_output.data_ptr<float>(),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(),
+        sigmoid_workspace.data_ptr<float>(),
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        bias_ptr,
+        stream);
+  } else if (dtype == at::ScalarType::Half) {
+    topk_gating_sigmoid_kernel_launcher<__half>(
+        reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(),
+        sigmoid_workspace.data_ptr<float>(),
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        bias_ptr,
+        stream);
+  } else if (dtype == at::ScalarType::BFloat16) {
+    topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>(
+        reinterpret_cast<const __nv_bfloat16*>(
+            gating_output.data_ptr<at::BFloat16>()),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(),
+        sigmoid_workspace.data_ptr<float>(),
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        bias_ptr,
+        stream);
+  } else {
+    LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
+  }
+}
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/moe/moe_topk_softmax_kernels.cuh

@@ -0,0 +1,855 @@
+// Adapt from
+// https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu
+// which is originally adapted from
+// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
+/* 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 <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include <cub/util_type.cuh>
+#include <cuda/functional>
+
+#include "kernels/cuda/device_utils.cuh"
+
+using cub_kvp = cub::KeyValuePair<int, float>;
+
+namespace {
+
+using namespace xllm::kernel::cuda;
+
+// ====================== Softmax things ===============================
+// We have our own implementation of softmax here so we can support transposing
+// the output in the softmax kernel when we extend this module to support
+// expert-choice routing.
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__
+    void moe_softmax(const T* input,
+                     const bool* finished,
+                     float* output,
+                     const int num_cols,
+                     const float moe_softcapping,
+                     const float* correction_bias) {
+  using BlockReduce = cub::BlockReduce<float, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  __shared__ float normalizing_factor;
+  __shared__ float float_max;
+
+  const int thread_row_offset = blockIdx.x * num_cols;
+
+  float threadData(-FLT_MAX);
+
+  // Don't touch finished rows.
+  if ((finished != nullptr) && finished[blockIdx.x]) {
+    return;
+  }
+
+  // First pass: Apply transformation, find max, and write transformed values to
+  // output
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    float val = convert_to_float<T>(input[idx]);
+
+    // Apply tanh softcapping if enabled
+    if (moe_softcapping != 0.0f) {
+      val = tanhf(val / moe_softcapping) * moe_softcapping;
+    }
+
+    // Apply correction bias if provided
+    if (correction_bias != nullptr) {
+      val = val + correction_bias[ii];
+    }
+
+    output[idx] = val;  // Store transformed value
+    threadData = max(val, threadData);
+  }
+
+  const float maxElem =
+      BlockReduce(tmpStorage).Reduce(threadData, MaxReduceOp());
+
+  if (threadIdx.x == 0) {
+    float_max = maxElem;
+  }
+  __syncthreads();
+
+  // Second pass: Compute sum using transformed values from output
+  threadData = 0;
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData += exp((output[idx] - float_max));
+  }
+
+  const auto Z = BlockReduce(tmpStorage).Sum(threadData);
+
+  if (threadIdx.x == 0) {
+    normalizing_factor = 1.f / Z;
+  }
+  __syncthreads();
+
+  // Third pass: Compute final softmax using transformed values from output
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    const float softmax_val =
+        exp((output[idx] - float_max)) * normalizing_factor;
+    output[idx] = softmax_val;
+  }
+}
+
+namespace moe {
+struct TopKPair {
+  static const int PAIR = 2;
+  static const int MAX_INDEX = 0;
+  cub_kvp max;
+  cub_kvp secondMax;
+
+  __device__ TopKPair() {}
+  __device__ TopKPair(cub_kvp max, cub_kvp secondMax)
+      : max(max), secondMax(secondMax) {}
+};
+
+struct TopKPairArgMax {
+  __device__ TopKPairArgMax() {}
+  __device__ __forceinline__ TopKPair
+  operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
+    cub_kvp globalMax, globalSecondMax;
+
+    // Determine the global maximum
+    if (candidate1.max.value > candidate2.max.value) {
+      globalMax = candidate1.max;
+    } else {
+      globalMax = candidate2.max;
+    }
+
+    // Determine the global second maximum
+    if (globalMax.key == candidate1.max.key) {
+      // If candidate1 contributed the max, compare its secondMax with
+      // candidate2's max
+      globalSecondMax = (candidate1.secondMax.value > candidate2.max.value)
+                            ? candidate1.secondMax
+                            : candidate2.max;
+    } else {
+      // If candidate2 contributed the max, compare its secondMax with
+      // candidate1's max
+      globalSecondMax = (candidate2.secondMax.value > candidate1.max.value)
+                            ? candidate2.secondMax
+                            : candidate1.max;
+    }
+    return TopKPair(globalMax, globalSecondMax);
+  }
+};
+}  // namespace moe
+
+template <int TPB>
+__launch_bounds__(TPB) __global__
+    void moe_topk_fast(float* inputs_after_softmax,
+                       const bool* finished,
+                       float* output,
+                       int* indices,
+                       const int num_experts,
+                       const int k,
+                       const int start_expert,
+                       const int end_expert,
+                       const bool renormalize) {
+  using namespace moe;
+  using BlockReduce = cub::BlockReduce<TopKPair, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+  TopKPair thread_pair;
+
+  const int block_row = blockIdx.x;
+
+  const bool row_is_active = finished ? !finished[block_row] : true;
+  const int thread_read_offset = blockIdx.x * num_experts;
+  float row_sum_for_renormalize = 0;
+  // Each loop finds the top 2 elements,
+  // thus requiring only ⌈k/2⌉ loops (calculated as (k + 1) / 2).
+  for (int k_idx = 0; k_idx < (k + TopKPair::PAIR - 1) / TopKPair::PAIR;
+       ++k_idx) {
+    // Initializing the top 2 elements by the minimum value.
+    thread_pair.max.key = 0;
+    thread_pair.max.value = -1.f;
+    thread_pair.secondMax.key = 0;
+    thread_pair.secondMax.value = -1.f;
+
+    cub_kvp inp_kvp;
+    for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
+      const int idx = thread_read_offset + expert;
+      inp_kvp.key = expert;
+      inp_kvp.value = inputs_after_softmax[idx];
+      // updating the thread_pair according to inp_kvp's value
+      if (inp_kvp.value > thread_pair.max.value) {
+        thread_pair.secondMax = thread_pair.max;
+        thread_pair.max = inp_kvp;
+      } else if (inp_kvp.value > thread_pair.secondMax.value) {
+        thread_pair.secondMax = inp_kvp;
+      }
+    }
+
+    TopKPairArgMax reducer;
+    const TopKPair result_pair =
+        BlockReduce(tmpStorage).Reduce(thread_pair, reducer);
+    if (threadIdx.x == 0) {
+#pragma unroll
+      // updating 2 elements to the result.
+      for (int i = 0; i < TopKPair::PAIR; i++) {
+        if (k_idx * 2 + i >= k) break;
+        cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max
+                                                    : result_pair.secondMax;
+        int expert = result.key;
+        bool node_uses_expert = expert >= start_expert && expert < end_expert;
+        bool should_process_row = row_is_active && node_uses_expert;
+        // The inputs_after_softmax is modified in-place to avoid unnecessary
+        // loops for finding the top k-1 value. 1.f represents the minimum
+        // value.
+        inputs_after_softmax[thread_read_offset + expert] = -1.f;
+        int idx = k * block_row + k_idx * 2 + i;
+        output[idx] = result.value;
+        indices[idx] =
+            should_process_row ? (expert - start_expert) : num_experts;
+        assert(indices[idx] >= 0);
+        row_sum_for_renormalize += result.value;
+      }
+    }
+    __syncthreads();
+  }
+
+  if (renormalize && threadIdx.x == 0) {
+    float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * block_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+template <int TPB>
+__launch_bounds__(TPB) __global__ void moe_topK(float* inputs_after_softmax,
+                                                const bool* finished,
+                                                float* output,
+                                                int* indices,
+                                                const int num_experts,
+                                                const int k,
+                                                const int start_expert,
+                                                const int end_expert,
+                                                const bool renormalize) {
+  using cub_kvp = cub::KeyValuePair<int, float>;
+  using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  cub_kvp thread_kvp;
+  cub::ArgMax arg_max;
+
+  const int block_row = blockIdx.x;
+
+  const bool row_is_active = finished ? !finished[block_row] : true;
+  const int thread_read_offset = blockIdx.x * num_experts;
+  float row_sum_for_renormalize = 0;
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    thread_kvp.key = 0;
+    thread_kvp.value = -1.f;  // This is OK because inputs are probabilities
+
+    cub_kvp inp_kvp;
+    for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
+      const int idx = thread_read_offset + expert;
+      inp_kvp.key = expert;
+      inp_kvp.value = inputs_after_softmax[idx];
+      thread_kvp = arg_max(inp_kvp, thread_kvp);
+    }
+
+    const cub_kvp result_kvp =
+        BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+    if (threadIdx.x == 0) {
+      // Ignore experts the node isn't responsible for with expert parallelism
+      const int expert = result_kvp.key;
+      const bool node_uses_expert =
+          expert >= start_expert && expert < end_expert;
+      const bool should_process_row = row_is_active && node_uses_expert;
+
+      const int idx = k * block_row + k_idx;
+      output[idx] = result_kvp.value;
+      indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
+      assert(indices[idx] >= 0);
+      row_sum_for_renormalize += result_kvp.value;
+      // The inputs_after_softmax is modified in-place to avoid unnecessary
+      // loops for finding the top k-1 value. 1.f represents the minimum value.
+      inputs_after_softmax[thread_read_offset + expert] = -1.f;
+    }
+    __syncthreads();
+  }
+
+  if (renormalize && threadIdx.x == 0) {
+    float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * block_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+// ====================== TopK softmax things ===============================
+
+/*
+  A Top-K gating softmax written to exploit when the number of experts in the
+  MoE layers are a small power of 2. This allows us to cleanly share the rows
+  among the threads in a single warp and eliminate communication between warps
+  (so no need to use shared mem).
+
+  It fuses the softmax, max and argmax into a single kernel.
+
+  Limitations:
+  1) This implementation is intended for when the number of experts is a small
+  power of 2. 2) This implementation assumes k is small, but will work for any
+  k.
+*/
+
+template <typename T,
+          int VPT,
+          int NUM_EXPERTS,
+          int WARPS_PER_CTA,
+          int BYTES_PER_LDG>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
+    void topk_gating_softmax(const T* input,
+                             const bool* finished,
+                             float* output,
+                             const int num_rows,
+                             int* indices,
+                             const int k,
+                             const int start_expert,
+                             const int end_expert,
+                             const bool renormalize,
+                             const float moe_softcapping,
+                             const float* correction_bias) {
+  // We begin by enforcing compile time assertions and setting up compile time
+  // constants.
+  static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
+  static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
+                "NUM_EXPERTS must be power of 2");
+  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
+                "BYTES_PER_LDG must be power of 2");
+  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+  // Number of bytes each thread pulls in per load
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+  // Restrictions based on previous section.
+  static_assert(
+      VPT % ELTS_PER_LDG == 0,
+      "The elements per thread must be a multiple of the elements per ldg");
+  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+                "The threads per row must cleanly divide the threads per warp");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+                "THREADS_PER_ROW must be power of 2");
+  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+                "THREADS_PER_ROW can be at most warp size");
+
+  // We have NUM_EXPERTS elements per row. We specialize for small #experts
+  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+  // Restrictions for previous section.
+  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+                "The elts per row must cleanly divide the total elt per warp");
+
+  // ===================== From this point, we finally start computing run-time
+  // variables. ========================
+
+  // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
+  // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
+  // rows. We start by computing the start row for each block.
+  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+
+  // Now, using the base row per thread block, we compute the base row per warp.
+  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+
+  // The threads in a warp are split into sub-groups that will work on a row.
+  // We compute row offset for each thread sub-group
+  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row = warp_base_row + thread_row_in_warp;
+
+  // Threads with indices out of bounds should early exit here.
+  if (thread_row >= num_rows) {
+    return;
+  }
+  const bool row_is_active = finished ? !finished[thread_row] : true;
+
+  // We finally start setting up the read pointers for each thread. First, each
+  // thread jumps to the start of the row it will read.
+  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+
+  // Now, we compute the group each thread belong to in order to determine the
+  // first column to start loads.
+  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+  // Determine the pointer type to use to read in the data depending on the
+  // BYTES_PER_LDG template param. In theory, this can support all powers of 2
+  // up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
+  // array here. We defined our own aligned array and use it here to avoid the
+  // dependency on CUTLASS.
+  using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+
+  // Finally, we pull in the data from global mem
+  T row_chunk_temp[VPT];
+  AccessType* row_chunk_vec_ptr =
+      reinterpret_cast<AccessType*>(&row_chunk_temp);
+  const AccessType* vec_thread_read_ptr =
+      reinterpret_cast<const AccessType*>(thread_read_ptr);
+#pragma unroll
+  // Note(Byron): interleaved loads to achieve better memory coalescing
+  // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
+  // thread[2] | thread[3] | ...
+  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+  }
+
+  float row_chunk[VPT];
+#pragma unroll
+  // Note(Byron): upcast logits to float32
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = convert_to_float<T>(row_chunk_temp[ii]);
+  }
+
+  // Apply tanh softcapping and correction bias
+  if (moe_softcapping != 0.0f || correction_bias != nullptr) {
+#pragma unroll
+    for (int ii = 0; ii < VPT; ++ii) {
+      float val = row_chunk[ii];
+
+      // Apply tanh softcapping if enabled
+      if (moe_softcapping != 0.0f) {
+        val = tanhf(val / moe_softcapping) * moe_softcapping;
+      }
+
+      // Apply correction bias if provided
+      if (correction_bias != nullptr) {
+        /*
+        LDG is interleaved
+        |thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG|
+        |--------- group0 --------| |----------group1 --------|
+                                      ^ local2
+        */
+        const int group_id = ii / ELTS_PER_LDG;
+        const int local_id = ii % ELTS_PER_LDG;
+        const int expert_idx = first_elt_read_by_thread +
+                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                               local_id;
+        val = val + correction_bias[expert_idx];
+      }
+
+      row_chunk[ii] = val;
+    }
+  }
+
+  // First, we perform a max reduce within the thread. We can do the max in fp16
+  // safely (I think) and just convert to float afterwards for the exp + sum
+  // reduction.
+  float thread_max = row_chunk[0];
+#pragma unroll
+  for (int ii = 1; ii < VPT; ++ii) {
+    thread_max = max(thread_max, row_chunk[ii]);
+  }
+
+  /*********************************/
+  /********* Softmax Begin *********/
+  /*********************************/
+
+// Now, we find the max within the thread group and distribute among the
+// threads. We use a butterfly reduce. lane id: 0-31 within a warp
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    // butterfly reduce with (lane id ^ mask)
+    thread_max = max(thread_max,
+                     XLLM_SHFL_XOR_SYNC_WIDTH(
+                         0xffffffff, thread_max, mask, THREADS_PER_ROW));
+  }
+
+  // From this point, thread max in all the threads have the max within the row.
+  // Now, we subtract the max from each element in the thread and take the exp.
+  // We also compute the thread local sum.
+  float row_sum = 0;
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = expf(row_chunk[ii] - thread_max);
+    row_sum += row_chunk[ii];
+  }
+
+// Now, we perform the sum reduce within each thread group. Similar to the max
+// reduce, we use a bufferfly pattern.
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    row_sum +=
+        XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW);
+  }
+
+  // From this point, all threads have the max and the sum for their rows in the
+  // thread_max and thread_sum variables respectively. Finally, we can scale the
+  // rows for the softmax. Technically, for top-k gating we don't need to
+  // compute the entire softmax row. We can likely look at the maxes and only
+  // compute for the top-k values in the row. However, this kernel will likely
+  // not be a bottle neck and it seems better to closer match torch and find the
+  // argmax after computing the softmax.
+  const float reciprocal_row_sum = 1.f / row_sum;
+
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+  }
+  /*******************************/
+  /********* Softmax End *********/
+  /*******************************/
+
+  // Now, softmax_res contains the softmax of the row chunk. Now, I want to find
+  // the topk elements in each row, along with the max index.
+  int start_col = first_elt_read_by_thread;
+  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+  float row_sum_for_renormalize = 0;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    // First, each thread does the local argmax
+    float max_val = row_chunk[0];
+    int expert = start_col;
+#pragma unroll
+    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+         ++ldg, col += COLS_PER_GROUP_LDG) {
+#pragma unroll
+      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+
+        // No check on the experts here since columns with the smallest index
+        // are processed first and only updated if > (not >=)
+        if (val > max_val) {
+          max_val = val;
+          expert = col + ii;
+        }
+      }
+    }
+
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads
+// reach consensus about the max. This will be useful for K > 1 so that the
+// threads can agree on "who" had the max value. That thread can then blank out
+// their max with -inf and the warp can run more iterations...
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+      float other_max =
+          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
+      int other_expert =
+          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
+
+      // We want lower indices to "win" in every thread so we break ties this
+      // way
+      if (other_max > max_val ||
+          (other_max == max_val && other_expert < expert)) {
+        max_val = other_max;
+        expert = other_expert;
+      }
+    }
+
+    // Write the max for this k iteration to global memory.
+    if (thread_group_idx == 0) {
+      // Add a guard to ignore experts not included by this node
+      const bool node_uses_expert =
+          expert >= start_expert && expert < end_expert;
+      const bool should_process_row = row_is_active && node_uses_expert;
+
+      // The lead thread from each sub-group will write out the final results to
+      // global memory. (This will be a single) thread per row of the
+      // input/output matrices.
+      const int idx = k * thread_row + k_idx;
+      output[idx] = max_val;
+      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+      row_sum_for_renormalize += max_val;
+    }
+
+    // Finally, we clear the value in the thread with the current max if there
+    // is another iteration to run.
+    if (k_idx + 1 < k) {
+      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int thread_to_clear_in_group =
+          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+      // Only the thread in the group which produced the max will reset the
+      // "winning" value to -inf.
+      if (thread_group_idx == thread_to_clear_in_group) {
+        const int offset_for_expert = expert % ELTS_PER_LDG;
+        // Safe to set to any negative value since row_chunk values must be
+        // between 0 and 1.
+        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+            -10000.f;
+      }
+    }
+  }
+
+  // Fuse renormalization of topk_weights into this kernel
+  if (renormalize && thread_group_idx == 0) {
+    float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+#pragma unroll
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * thread_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+template <typename T, int EXPERTS, int WARPS_PER_TB>
+void topk_gating_softmax_launcher_helper(const T* input,
+                                         const bool* finished,
+                                         float* output,
+                                         int* indices,
+                                         const int num_rows,
+                                         const int k,
+                                         const int start_expert,
+                                         const int end_expert,
+                                         const bool renormalize,
+                                         const float moe_softcapping,
+                                         const float* correction_bias,
+                                         cudaStream_t stream) {
+  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+
+  static constexpr int BYTES_PER_LDG =
+      MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int VPT = Constants::VPT;
+  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+  dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+  topk_gating_softmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+      <<<num_blocks, block_dim, 0, stream>>>(input,
+                                             finished,
+                                             output,
+                                             num_rows,
+                                             indices,
+                                             k,
+                                             start_expert,
+                                             end_expert,
+                                             renormalize,
+                                             moe_softcapping,
+                                             correction_bias);
+}
+
+#define LAUNCH_SOFTMAX(TYPE, NUM_EXPERTS, WARPS_PER_TB)                 \
+  topk_gating_softmax_launcher_helper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
+      gating_output,                                                    \
+      nullptr,                                                          \
+      topk_weights,                                                     \
+      topk_indices,                                                     \
+      num_tokens,                                                       \
+      topk,                                                             \
+      0,                                                                \
+      num_experts,                                                      \
+      renormalize,                                                      \
+      moe_softcapping,                                                  \
+      correction_bias,                                                  \
+      stream);
+
+template <typename T>
+void topk_gating_softmax_kernel_launcher(const T* gating_output,
+                                         float* topk_weights,
+                                         int* topk_indices,
+                                         float* softmax_workspace,
+                                         const int num_tokens,
+                                         const int num_experts,
+                                         const int topk,
+                                         const bool renormalize,
+                                         const float moe_softcapping,
+                                         const float* correction_bias,
+                                         cudaStream_t stream) {
+  static constexpr int WARPS_PER_TB = 4;
+  switch (num_experts) {
+    case 1:
+      LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB);
+      break;
+    case 2:
+      LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB);
+      break;
+    case 4:
+      LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB);
+      break;
+    case 8:
+      LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB);
+      break;
+    case 16:
+      LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB);
+      break;
+    case 32:
+      LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB);
+      break;
+    case 64:
+      LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB);
+      break;
+    case 128:
+      LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB);
+      break;
+    case 256:
+      LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
+      break;
+    default: {
+      CHECK(softmax_workspace != nullptr)
+          << "softmax_workspace must be provided for num_experts that are "
+             "not a power of 2.";
+      static constexpr int TPB = 256;
+      moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
+                                                          nullptr,
+                                                          softmax_workspace,
+                                                          num_experts,
+                                                          moe_softcapping,
+                                                          correction_bias);
+      if (topk == 1) {
+        // Note: As an optimization for better performance,
+        // the softmax_workspace is overwritten in-place by both moeTopK and
+        // moe_topk_fast.
+        moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
+                                                      nullptr,
+                                                      topk_weights,
+                                                      topk_indices,
+                                                      num_experts,
+                                                      topk,
+                                                      0,
+                                                      num_experts,
+                                                      renormalize);
+      } else {
+        moe_topk_fast<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
+                                                           nullptr,
+                                                           topk_weights,
+                                                           topk_indices,
+                                                           num_experts,
+                                                           topk,
+                                                           0,
+                                                           num_experts,
+                                                           renormalize);
+      }
+    }
+  }
+}
+}  // namespace
+
+namespace xllm::kernel::cuda {
+void topk_softmax(torch::Tensor& topk_weights,   // [num_tokens, topk]
+                  torch::Tensor& topk_indices,   // [num_tokens, topk]
+                  torch::Tensor& gating_output,  // [num_tokens, num_experts]
+                  const bool renormalize,
+                  const double moe_softcapping,
+                  const std::optional<torch::Tensor>& correction_bias) {
+  // Check data type
+  CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
+        gating_output.scalar_type() == at::ScalarType::Half ||
+        gating_output.scalar_type() == at::ScalarType::BFloat16)
+      << "gating_output must be float32, float16, or bfloat16";
+
+  // Check dimensions
+  CHECK(gating_output.dim() == 2)
+      << "gating_output must be 2D tensor [num_tokens, num_experts]";
+  CHECK(topk_weights.dim() == 2)
+      << "topk_weights must be 2D tensor [num_tokens, topk]";
+  CHECK(topk_indices.dim() == 2)
+      << "topk_indices must be 2D tensor [num_tokens, topk]";
+
+  // Check shapes
+  CHECK(gating_output.size(0) == topk_weights.size(0))
+      << "First dimension of topk_weights must match num_tokens in "
+         "gating_output"
+      << "First dimension of topk_indices must match num_tokens in "
+         "gating_output";
+
+  CHECK(topk_weights.size(-1) == topk_indices.size(-1))
+      << "Second dimension of topk_indices must match topk in topk_weights"
+      << "topk must be less than or equal to num_experts";
+
+  const int num_experts = static_cast<int>(gating_output.size(-1));
+  const int num_tokens = static_cast<int>(gating_output.size(0));
+  const int topk = static_cast<int>(topk_weights.size(-1));
+
+  const bool is_pow_2 =
+      (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+  const bool needs_workspace = !is_pow_2 || num_experts > 256;
+  const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  torch::Tensor softmax_workspace = torch::empty(
+      {workspace_size}, gating_output.options().dtype(at::ScalarType::Float));
+
+  const at::ScalarType dtype = gating_output.scalar_type();
+
+  // Validate correction_bias if provided - must always be float32
+  const float* bias_ptr = nullptr;
+  if (correction_bias.has_value()) {
+    const torch::Tensor& bias_tensor = correction_bias.value();
+    CHECK(bias_tensor.dim() == 1)
+        << "correction_bias must be 1D tensor [num_experts]";
+    CHECK(bias_tensor.size(0) == num_experts)
+        << "correction_bias size must match num_experts";
+    CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
+        << "correction_bias must be float32, got " << bias_tensor.scalar_type();
+    bias_ptr = bias_tensor.data_ptr<float>();
+  }
+
+  // Cast moe_softcapping from double to float for CUDA kernels
+  const float moe_softcapping_f = static_cast<float>(moe_softcapping);
+
+  if (dtype == at::ScalarType::Float) {
+    topk_gating_softmax_kernel_launcher<float>(
+        gating_output.data_ptr<float>(),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(),
+        softmax_workspace.data_ptr<float>(),
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        moe_softcapping_f,
+        bias_ptr,
+        stream);
+  } else if (dtype == at::ScalarType::Half) {
+    topk_gating_softmax_kernel_launcher<__half>(
+        reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(),
+        softmax_workspace.data_ptr<float>(),
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        moe_softcapping_f,
+        bias_ptr,
+        stream);
+  } else if (dtype == at::ScalarType::BFloat16) {
+    topk_gating_softmax_kernel_launcher<__nv_bfloat16>(
+        reinterpret_cast<const __nv_bfloat16*>(
+            gating_output.data_ptr<at::BFloat16>()),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(),
+        softmax_workspace.data_ptr<float>(),
+        num_tokens,
+        num_experts,
+        topk,
+        renormalize,
+        moe_softcapping_f,
+        bias_ptr,
+        stream);
+  } else {
+    LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
+  }
+}
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/norm.cu

@@ -0,0 +1,590 @@
+/* Copyright 2025 The vLLM Authors and The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <c10/cuda/CUDAGuard.h>
+#include <torch/cuda.h>
+
+#include <cub/cub.cuh>
+
+#include "cuda_ops_api.h"
+#include "fp8_quant_utils.cuh"
+#include "type_convert.cuh"
+
+// ref to:
+// https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_kernels.cu
+
+#if CUB_VERSION >= 200800
+#include <cuda/std/functional>
+using CubAddOp = ::cuda::std::plus<>;
+using CubMaxOp = ::cuda::maximum<>;
+#else   // if CUB_VERSION < 200800
+using CubAddOp = cub::Sum;
+using CubMaxOp = cub::Max;
+#endif  // CUB_VERSION
+
+namespace {
+
+using namespace xllm::kernel::cuda;
+
+template <typename scalar_t>
+__global__ void rms_norm_kernel(
+    scalar_t* __restrict__ out,          // [..., hidden_size]
+    const scalar_t* __restrict__ input,  // [..., hidden_size]
+    const int64_t input_stride,
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float epsilon,
+    const int num_tokens,
+    const int hidden_size) {
+  __shared__ float s_variance;
+  float variance = 0.0f;
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    const float x = (float)input[blockIdx.x * input_stride + idx];
+    variance += x * x;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    float x = (float)input[blockIdx.x * input_stride + idx];
+    out[blockIdx.x * hidden_size + idx] =
+        ((scalar_t)(x * s_variance)) * weight[idx];
+  }
+}
+
+/* Function specialization in the case of FP16/BF16 tensors.
+   Additional optimizations we can make in this case are
+   packed and vectorized operations, which help with the
+   memory latency bottleneck. */
+template <typename scalar_t, int width>
+__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
+fused_add_rms_norm_kernel(
+    scalar_t* __restrict__ input,  // [..., hidden_size]
+    const int64_t input_stride,
+    scalar_t* __restrict__ residual,      // [..., hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float epsilon,
+    const int num_tokens,
+    const int hidden_size) {
+  // Sanity checks on our vector struct and type-punned pointer arithmetic
+  static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
+  static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
+
+  const int vec_hidden_size = hidden_size / width;
+  const int64_t vec_input_stride = input_stride / width;
+  __shared__ float s_variance;
+  float variance = 0.0f;
+  /* These and the argument pointers are all declared `restrict` as they are
+     not aliased in practice. Argument pointers should not be dereferenced
+     in this kernel as that would be undefined behavior */
+  auto* __restrict__ input_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(input);
+  auto* __restrict__ residual_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
+  auto* __restrict__ weight_v =
+      reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);
+
+  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
+    int id = blockIdx.x * vec_hidden_size + idx;
+    int64_t strided_id = blockIdx.x * vec_input_stride + idx;
+    _f16Vec<scalar_t, width> temp = input_v[strided_id];
+    temp += residual_v[id];
+    variance += temp.sum_squares();
+    residual_v[id] = temp;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
+    int id = blockIdx.x * vec_hidden_size + idx;
+    int64_t strided_id = blockIdx.x * vec_input_stride + idx;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
+    input_v[strided_id] = temp;
+  }
+}
+
+/* Generic fused_add_rms_norm_kernel
+   The width field is not used here but necessary for other specializations.
+ */
+template <typename scalar_t, int width>
+__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
+fused_add_rms_norm_kernel(
+    scalar_t* __restrict__ input,  // [..., hidden_size]
+    const int64_t input_stride,
+    scalar_t* __restrict__ residual,      // [..., hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float epsilon,
+    const int num_tokens,
+    const int hidden_size) {
+  __shared__ float s_variance;
+  float variance = 0.0f;
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    scalar_t z = input[blockIdx.x * input_stride + idx];
+    z += residual[blockIdx.x * hidden_size + idx];
+    float x = (float)z;
+    variance += x * x;
+    residual[blockIdx.x * hidden_size + idx] = z;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    float x = (float)residual[blockIdx.x * hidden_size + idx];
+    input[blockIdx.x * input_stride + idx] =
+        ((scalar_t)(x * s_variance)) * weight[idx];
+  }
+}
+
+#define LAUNCH_FUSED_ADD_RMS_NORM(width)                                \
+  DISPATCH_FLOATING_TYPES(                                              \
+      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {           \
+        fused_add_rms_norm_kernel<scalar_t, width>                      \
+            <<<grid, block, 0, stream>>>(input.data_ptr<scalar_t>(),    \
+                                         input_stride,                  \
+                                         residual.data_ptr<scalar_t>(), \
+                                         weight.data_ptr<scalar_t>(),   \
+                                         epsilon,                       \
+                                         num_tokens,                    \
+                                         hidden_size);                  \
+      });
+
+// ============================================================================
+// Fused RMSNorm + Static FP8 Quantization Kernels
+// ============================================================================
+// These kernels combine RMSNorm and FP8 quantization to reduce memory
+// bandwidth by avoiding the intermediate write-back to global memory.
+
+// Dispatch macro for FP8 types
+#define DISPATCH_FP8_TYPES(TYPE, NAME, ...)         \
+  [&] {                                             \
+    const auto& the_type = TYPE;                    \
+    switch (the_type) {                             \
+      case at::ScalarType::Float8_e4m3fn: {         \
+        using fp8_t = c10::Float8_e4m3fn;           \
+        return __VA_ARGS__();                       \
+      }                                             \
+      default:                                      \
+        AT_ERROR(#NAME,                             \
+                 " not implemented for FP8 type '", \
+                 toString(the_type),                \
+                 "'");                              \
+    }                                               \
+  }()
+
+/**
+ * Fused RMSNorm + Static FP8 Quantization kernel (without residual)
+ * Combines RMSNorm and FP8 quantization in a single kernel to reduce
+ * memory bandwidth by avoiding intermediate write-back.
+ *
+ * @tparam scalar_t Input data type (float, half, bfloat16)
+ * @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn)
+ * @param out Output FP8 tensor [num_tokens, hidden_size]
+ * @param input Input tensor [num_tokens, hidden_size]
+ * @param input_stride Stride of input tensor in the token dimension
+ * @param weight RMSNorm weight tensor [hidden_size]
+ * @param scale FP8 quantization scale (scalar)
+ * @param epsilon RMSNorm epsilon
+ * @param num_tokens Number of tokens
+ * @param hidden_size Hidden dimension size
+ */
+template <typename scalar_t, typename fp8_type>
+__global__ void rms_norm_static_fp8_quant_kernel(
+    fp8_type* __restrict__ out,          // [num_tokens, hidden_size]
+    const scalar_t* __restrict__ input,  // [num_tokens, hidden_size]
+    const int64_t input_stride,
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float* __restrict__ scale,      // [1]
+    const float epsilon,
+    const int num_tokens,
+    const int hidden_size) {
+  __shared__ float s_variance;
+  float variance = 0.0f;
+
+  const scalar_t* input_row = input + blockIdx.x * input_stride;
+
+  // Step 1: Compute variance for RMSNorm
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    const float x = static_cast<float>(input_row[idx]);
+    variance += x * x;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  // Step 2: Precompute scale inverse to avoid division
+  const float scale_inv = 1.0f / (*scale);
+
+  // Step 3: Fused RMSNorm + FP8 quantization
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    float x = static_cast<float>(input_row[idx]);
+    float out_norm = (static_cast<scalar_t>(x * s_variance)) *
+                     static_cast<float>(weight[idx]);
+    out[blockIdx.x * hidden_size + idx] =
+        xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(out_norm,
+                                                                  scale_inv);
+  }
+}
+
+/**
+ * Fused Add + RMSNorm + Static FP8 Quantization kernel (with residual)
+ * Optimized version with packed + vectorized operations for FP16/BF16.
+ *
+ * @tparam scalar_t Input data type (float, half, bfloat16)
+ * @tparam width Vector width for optimization (0, 8)
+ * @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn)
+ */
+template <typename scalar_t, int width, typename fp8_type>
+__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
+fused_add_rms_norm_static_fp8_quant_kernel(
+    fp8_type* __restrict__ out,    // [num_tokens, hidden_size]
+    scalar_t* __restrict__ input,  // [num_tokens, hidden_size]
+    const int64_t input_stride,
+    scalar_t* __restrict__ residual,      // [num_tokens, hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float* __restrict__ scale,      // [1]
+    const float epsilon,
+    const int num_tokens,
+    const int hidden_size) {
+  static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
+  static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
+
+  const int vec_hidden_size = hidden_size / width;
+  const int64_t vec_input_stride = input_stride / width;
+  __shared__ float s_variance;
+  float variance = 0.0f;
+
+  auto* __restrict__ input_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(input);
+  auto* __restrict__ residual_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
+  auto* __restrict__ weight_v =
+      reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);
+
+  // Step 1: Fused add and compute variance
+  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
+    int id = blockIdx.x * vec_hidden_size + idx;
+    int64_t strided_id = blockIdx.x * vec_input_stride + idx;
+    _f16Vec<scalar_t, width> temp = input_v[strided_id];
+    temp += residual_v[id];
+    variance += temp.sum_squares();
+    residual_v[id] = temp;  // Store updated residual
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  // Step 2: Precompute scale inverse
+  const float scale_inv = 1.0f / (*scale);
+
+  // Step 3: Fused RMSNorm + FP8 quantization
+  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
+    int id = blockIdx.x * vec_hidden_size + idx;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
+
+    // Convert each element to FP8
+#pragma unroll
+    for (int i = 0; i < width; ++i) {
+      float val = _typeConvert<scalar_t>::convert(temp.data[i]);
+      out[id * width + i] =
+          xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(val,
+                                                                    scale_inv);
+    }
+  }
+}
+
+/**
+ * Generic fused add + RMSNorm + FP8 quant kernel (fallback for unaligned data)
+ */
+template <typename scalar_t, int width, typename fp8_type>
+__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
+fused_add_rms_norm_static_fp8_quant_kernel(
+    fp8_type* __restrict__ out,    // [num_tokens, hidden_size]
+    scalar_t* __restrict__ input,  // [num_tokens, hidden_size]
+    const int64_t input_stride,
+    scalar_t* __restrict__ residual,      // [num_tokens, hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float* __restrict__ scale,      // [1]
+    const float epsilon,
+    const int num_tokens,
+    const int hidden_size) {
+  __shared__ float s_variance;
+  float variance = 0.0f;
+
+  // Step 1: Fused add and compute variance
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    scalar_t z = input[blockIdx.x * input_stride + idx];
+    z += residual[blockIdx.x * hidden_size + idx];
+    float x = static_cast<float>(z);
+    variance += x * x;
+    residual[blockIdx.x * hidden_size + idx] = z;  // Store updated residual
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  // Step 2: Precompute scale inverse
+  const float scale_inv = 1.0f / (*scale);
+
+  // Step 3: Fused RMSNorm + FP8 quantization
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    float x = static_cast<float>(residual[blockIdx.x * hidden_size + idx]);
+    float out_norm = (static_cast<scalar_t>(x * s_variance)) *
+                     static_cast<float>(weight[idx]);
+    out[blockIdx.x * hidden_size + idx] =
+        xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(out_norm,
+                                                                  scale_inv);
+  }
+}
+
+#define LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(width)                     \
+  DISPATCH_FLOATING_TYPES(                                                    \
+      input.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] {       \
+        DISPATCH_FP8_TYPES(                                                   \
+            out.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] {   \
+              fused_add_rms_norm_static_fp8_quant_kernel<scalar_t,            \
+                                                         width,               \
+                                                         fp8_t>               \
+                  <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),         \
+                                               input.data_ptr<scalar_t>(),    \
+                                               input_stride,                  \
+                                               residual.data_ptr<scalar_t>(), \
+                                               weight.data_ptr<scalar_t>(),   \
+                                               scale.data_ptr<float>(),       \
+                                               epsilon,                       \
+                                               num_tokens,                    \
+                                               hidden_size);                  \
+            });                                                               \
+      });
+
+}  // namespace
+
+namespace xllm::kernel::cuda {
+
+// flashinfer rmsnorm ops
+// void rmsnorm(torch::Tensor output,
+//              torch::Tensor input,
+//              torch::Tensor weight,
+//              double eps) {
+//   FunctionFactory::get_instance().rmsnorm_func("norm").call(
+//       output, input, weight, eps, support_pdl());
+// }
+
+void rms_norm(torch::Tensor output,  // [..., hidden_size]
+              torch::Tensor input,   // [..., hidden_size]
+              torch::Tensor weight,  // [hidden_size]
+              double eps) {
+  CHECK(output.is_contiguous());
+  CHECK(input.stride(-1) == 1);
+  CHECK(weight.is_contiguous());
+
+  int hidden_size = input.size(-1);
+  int num_tokens = input.numel() / hidden_size;
+  int64_t input_stride = input.stride(-2);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(hidden_size, 1024));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
+    rms_norm_kernel<scalar_t>
+        <<<grid, block, 0, stream>>>(output.data_ptr<scalar_t>(),
+                                     input.data_ptr<scalar_t>(),
+                                     input_stride,
+                                     weight.data_ptr<scalar_t>(),
+                                     eps,
+                                     num_tokens,
+                                     hidden_size);
+  });
+}
+
+void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
+                        torch::Tensor& residual,  // [..., hidden_size]
+                        torch::Tensor& weight,    // [hidden_size]
+                        double epsilon) {
+  CHECK(weight.scalar_type() == input.scalar_type());
+  CHECK(input.scalar_type() == residual.scalar_type());
+  CHECK(residual.is_contiguous());
+  CHECK(weight.is_contiguous());
+  int hidden_size = input.size(-1);
+  int64_t input_stride = input.stride(-2);
+  int num_tokens = input.numel() / hidden_size;
+
+  dim3 grid(num_tokens);
+  /* This kernel is memory-latency bound in many scenarios.
+     When num_tokens is large, a smaller block size allows
+     for increased block occupancy on CUs and better latency
+     hiding on global mem ops. */
+  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  dim3 block(std::min(hidden_size, max_block_size));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  /*If the tensor types are FP16/BF16, try to use the optimized kernel
+    with packed + vectorized ops.
+    Max optimization is achieved with a width-8 vector of FP16/BF16s
+    since we can load at most 128 bits at once in a global memory op.
+    However, this requires each tensor's data to be aligned to 16
+    bytes.
+   */
+  auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
+  auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
+  auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
+  constexpr int vector_width = 8;
+  constexpr int req_alignment_bytes =
+      vector_width * 2;  // vector_width * sizeof(bfloat16 or float16) (float32
+                         // falls back to non-vectorized version anyway)
+  bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
+                          res_ptr % req_alignment_bytes == 0 &&
+                          wt_ptr % req_alignment_bytes == 0;
+  bool offsets_are_multiple_of_vector_width =
+      hidden_size % vector_width == 0 && input_stride % vector_width == 0;
+  if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) {
+    LAUNCH_FUSED_ADD_RMS_NORM(8);
+  } else {
+    LAUNCH_FUSED_ADD_RMS_NORM(0);
+  }
+}
+
+// ============================================================================
+// Fused RMSNorm + Static FP8 Quantization Host Functions
+// ============================================================================
+
+void rms_norm_static_fp8_quant(torch::Tensor& out,    // [..., hidden_size], FP8
+                               torch::Tensor& input,  // [..., hidden_size]
+                               torch::Tensor& weight,  // [hidden_size]
+                               torch::Tensor& scale,   // [1]
+                               double epsilon) {
+  CHECK(out.is_contiguous());
+  CHECK(input.stride(-1) == 1);
+  CHECK(weight.is_contiguous());
+  CHECK(scale.is_contiguous());
+
+  int hidden_size = input.size(-1);
+  int64_t input_stride = input.stride(-2);
+  int num_tokens = input.numel() / hidden_size;
+
+  // For large num_tokens, use smaller blocks to increase SM concurrency
+  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  dim3 grid(num_tokens);
+  dim3 block(std::min(hidden_size, max_block_size));
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "rms_norm_static_fp8_quant", [&] {
+        DISPATCH_FP8_TYPES(out.scalar_type(), "rms_norm_static_fp8_quant", [&] {
+          rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t>
+              <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),
+                                           input.data_ptr<scalar_t>(),
+                                           input_stride,
+                                           weight.data_ptr<scalar_t>(),
+                                           scale.data_ptr<float>(),
+                                           epsilon,
+                                           num_tokens,
+                                           hidden_size);
+        });
+      });
+}
+
+void fused_add_rms_norm_static_fp8_quant(
+    torch::Tensor& out,       // [..., hidden_size], FP8
+    torch::Tensor& input,     // [..., hidden_size]
+    torch::Tensor& residual,  // [..., hidden_size]
+    torch::Tensor& weight,    // [hidden_size]
+    torch::Tensor& scale,     // [1]
+    double epsilon) {
+  CHECK(out.is_contiguous());
+  CHECK(residual.is_contiguous());
+  CHECK(weight.is_contiguous());
+  CHECK(scale.is_contiguous());
+  CHECK(residual.scalar_type() == input.scalar_type());
+  CHECK(weight.scalar_type() == input.scalar_type());
+
+  int hidden_size = input.size(-1);
+  int64_t input_stride = input.stride(-2);
+  int num_tokens = input.numel() / hidden_size;
+
+  dim3 grid(num_tokens);
+  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  dim3 block(std::min(hidden_size, max_block_size));
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  // Check alignment for vectorized kernel
+  auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
+  auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
+  auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
+  constexpr int vector_width = 8;
+  constexpr int req_alignment_bytes = vector_width * 2;
+
+  bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
+                          res_ptr % req_alignment_bytes == 0 &&
+                          wt_ptr % req_alignment_bytes == 0;
+  bool offsets_are_multiple_of_vector_width =
+      hidden_size % vector_width == 0 && input_stride % vector_width == 0;
+
+  if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) {
+    LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(8);
+  } else {
+    LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(0);
+  }
+}
+
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/cuda/rope.cu

@@ -0,0 +1,251 @@
+/* Copyright 2025 The vLLM Authors and The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include "cuda_ops_api.h"
+
+// ref to:
+// https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu
+
+namespace {
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_token_rotary_embedding(
+    scalar_t* __restrict__ arr,
+    const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr,
+    int rot_offset,
+    int embed_dim) {
+  int x_index, y_index;
+  scalar_t cos, sin;
+  if (IS_NEOX) {
+    // GPT-NeoX style rotary embedding.
+    x_index = rot_offset;
+    y_index = embed_dim + rot_offset;
+    cos = *(cos_ptr + x_index);
+    sin = *(sin_ptr + x_index);
+  } else {
+    // GPT-J style rotary embedding.
+    x_index = 2 * rot_offset;
+    y_index = 2 * rot_offset + 1;
+    cos = *(cos_ptr + x_index / 2);
+    sin = *(sin_ptr + x_index / 2);
+  }
+
+  const scalar_t x = arr[x_index];
+  const scalar_t y = arr[y_index];
+  arr[x_index] = x * cos - y * sin;
+  arr[y_index] = y * cos + x * sin;
+}
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_rotary_embedding(
+    scalar_t* __restrict__ query,  // [batch_size, seq_len, num_heads,
+                                   // head_size] or [num_tokens, num_heads,
+                                   // head_size]
+    scalar_t* __restrict__ key,    // nullptr or
+                                   // [batch_size, seq_len, num_kv_heads,
+                                   // head_size] or [num_tokens, num_kv_heads,
+                                   // head_size]
+    const scalar_t* cache_ptr,
+    const int head_size,
+    const int num_heads,
+    const int num_kv_heads,
+    const int rot_dim,
+    const int token_idx,
+    const int64_t query_stride,
+    const int64_t key_stride,
+    const int64_t head_stride) {
+  const int embed_dim = rot_dim / 2;
+  const scalar_t* cos_ptr = cache_ptr;
+  const scalar_t* sin_ptr = cache_ptr + embed_dim;
+
+  const int nq = num_heads * embed_dim;
+  for (int i = threadIdx.x; i < nq; i += blockDim.x) {
+    const int head_idx = i / embed_dim;
+    const int64_t token_head =
+        token_idx * query_stride + head_idx * head_stride;
+    const int rot_offset = i % embed_dim;
+    apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+  }
+
+  if (key != nullptr) {
+    const int nk = num_kv_heads * embed_dim;
+    for (int i = threadIdx.x; i < nk; i += blockDim.x) {
+      const int head_idx = i / embed_dim;
+      const int64_t token_head =
+          token_idx * key_stride + head_idx * head_stride;
+      const int rot_offset = i % embed_dim;
+      apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+    }
+  }
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__global__ void rotary_embedding_kernel(
+    const int64_t* __restrict__ positions,  // [batch_size, seq_len] or
+                                            // [num_tokens]
+    scalar_t* __restrict__ query,           // [batch_size, seq_len, num_heads,
+                                   // head_size] or [num_tokens, num_heads,
+                                   // head_size]
+    scalar_t* __restrict__ key,  // nullptr or
+                                 // [batch_size, seq_len, num_kv_heads,
+                                 // head_size] or [num_tokens, num_kv_heads,
+                                 // head_size]
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
+    const int rot_dim,
+    const int64_t query_stride,
+    const int64_t key_stride,
+    const int64_t head_stride,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size) {
+  // Each thread block is responsible for one token.
+  const int token_idx = blockIdx.x;
+  int64_t pos = positions[token_idx];
+  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
+
+  apply_rotary_embedding<scalar_t, IS_NEOX>(query,
+                                            key,
+                                            cache_ptr,
+                                            head_size,
+                                            num_heads,
+                                            num_kv_heads,
+                                            rot_dim,
+                                            token_idx,
+                                            query_stride,
+                                            key_stride,
+                                            head_stride);
+}
+}  // namespace
+
+namespace xllm::kernel::cuda {
+
+// flashinfer rope ops
+// void apply_rope_pos_ids_cos_sin_cache(torch::Tensor q,
+//                                       torch::Tensor k,
+//                                       torch::Tensor cos_sin_cache,
+//                                       torch::Tensor pos_ids,
+//                                       bool interleave) {
+//   const int64_t head_dim = cos_sin_cache.size(-1) / 2;
+//   q = q.view({q.size(0), -1, head_dim});
+//   k = k.view({k.size(0), -1, head_dim});
+
+//   FunctionFactory::get_instance().rope_func("rope").call(
+//       q, k, q, k, cos_sin_cache, pos_ids, interleave);
+// }
+
+void rotary_embedding(
+    torch::Tensor& positions,  // [batch_size, seq_len] or [num_tokens]
+    torch::Tensor& query,  // [batch_size, seq_len, num_heads * head_size] or
+                           // [num_tokens, num_heads * head_size] or
+                           // [batch_size, seq_len, num_heads, head_size] or
+                           // [num_tokens, num_heads, head_size]
+    std::optional<torch::Tensor> key,
+    // null or
+    // [batch_size, seq_len, num_kv_heads * head_size] or
+    // [num_tokens, num_kv_heads * head_size] or
+    // [batch_size, seq_len, num_heads, head_size] or
+    // [num_tokens, num_heads, head_size]
+    // int64_t head_size,
+    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    bool is_neox) {
+  // num_tokens = batch_size * seq_len
+  int64_t head_size = cos_sin_cache.size(-1);
+  int64_t num_tokens = positions.numel();
+  int positions_ndim = positions.dim();
+
+  // Make sure num_tokens dim is consistent across positions, query, and key
+  CHECK(positions_ndim == 1 || positions_ndim == 2)
+      << "positions must have shape [num_tokens] or [batch_size, seq_len]";
+
+  if (positions_ndim == 1) {
+    CHECK(query.size(0) == positions.size(0) &&
+          (!key.has_value() || key->size(0) == positions.size(0)))
+        << "query, key and positions must have the same number of tokens";
+  }
+  if (positions_ndim == 2) {
+    CHECK(query.size(0) == positions.size(0) &&
+          (!key.has_value() || key->size(0) == positions.size(0)) &&
+          query.size(1) == positions.size(1) &&
+          (!key.has_value() || key->size(1) == positions.size(1)))
+        << "query, key and positions must have the same batch_size and seq_len";
+  }
+
+  // Make sure head_size is valid for query and key
+  // hidden_size = num_heads * head_size
+  int query_hidden_size = query.numel() / num_tokens;
+  int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
+  CHECK(query_hidden_size % head_size == 0);
+  CHECK(key_hidden_size % head_size == 0);
+
+  // Make sure query and key have consistent number of heads
+  int num_heads = query_hidden_size / head_size;
+  int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
+  CHECK(num_heads % num_kv_heads == 0);
+
+  int rot_dim = cos_sin_cache.size(1);
+  int seq_dim_idx = positions_ndim - 1;
+  int64_t query_stride = query.stride(seq_dim_idx);
+  int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
+  // Determine head stride: for [*, heads, head_size] use stride of last dim;
+  // for flat [*, heads*head_size], heads blocks are contiguous of size
+  // head_size
+  int query_ndim = query.dim();
+  int64_t head_stride =
+      (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  DISPATCH_FLOATING_TYPES(
+      query.scalar_type(), "apply_rope_pos_ids_cos_sin_cache", [&] {
+        if (is_neox) {
+          rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(),
+              query.data_ptr<scalar_t>(),
+              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
+              cos_sin_cache.data_ptr<scalar_t>(),
+              rot_dim,
+              query_stride,
+              key_stride,
+              head_stride,
+              num_heads,
+              num_kv_heads,
+              head_size);
+        } else {
+          rotary_embedding_kernel<scalar_t, false><<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(),
+              query.data_ptr<scalar_t>(),
+              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
+              cos_sin_cache.data_ptr<scalar_t>(),
+              rot_dim,
+              query_stride,
+              key_stride,
+              head_stride,
+              num_heads,
+              num_kv_heads,
+              head_size);
+        }
+      });
+}
+
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/kernels/ilu/activation.cpp

@@ -0,0 +1,32 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "ilu_ops_api.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void act_and_mul(torch::Tensor out,
+                 torch::Tensor input,
+                 const std::string& act_mode) {
+  if (act_mode == "silu") {
+    infer::silu_and_mul(input, out);
+  } else {
+    LOG(FATAL) << "Unsupported act mode: " << act_mode
+               << ", only support silu, gelu, gelu_tanh";
+  }
+}
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm/kernels/ilu/attention.cpp

@@ -0,0 +1,163 @@
+
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "ilu_ops_api.h"
+#include "ixinfer.h"
+#include "utils.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void reshape_paged_cache(torch::Tensor& key,
+                         std::optional<torch::Tensor>& value,
+                         torch::Tensor& key_cache,
+                         std::optional<torch::Tensor>& value_cache,
+                         torch::Tensor& slot_mapping) {
+  auto value_ = value.value_or(torch::Tensor());
+  auto value_cache_ = value_cache.value_or(torch::Tensor());
+
+  int64_t key_token_stride = key.stride(0);
+  int64_t value_token_stride = 0;
+  if (value_.defined()) {
+    value_token_stride = value_.stride(0);
+  }
+  slot_mapping = slot_mapping.to(at::kLong);
+  infer::xllm_reshape_and_cache(key,
+                                value_,
+                                key_cache,
+                                value_cache_,
+                                slot_mapping,
+                                key_token_stride,
+                                value_token_stride);
+}
+
+void batch_prefill(torch::Tensor& query,
+                   const torch::Tensor& key,
+                   const std::optional<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& v_quant_scale,
+                   const torch::Tensor& block_tables,
+                   int64_t max_query_len,
+                   int64_t max_seq_len,
+                   float scale,
+                   bool is_causal,
+                   int64_t window_size_left,
+                   int64_t window_size_right,
+                   const std::string& compute_dtype,
+                   bool return_lse) {
+  double softcap = 0.0;
+  bool sqrt_alibi = false;
+  auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor());
+  auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor());
+  auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor());
+  auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor());
+  auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor());
+  auto block_tables_ = block_tables;
+  auto key_ = key;
+  auto value_ = value.value();
+  infer::ixinfer_flash_attn_unpad_with_block_tables(query,
+                                                    key_,
+                                                    value_,
+                                                    output,
+                                                    block_tables_,
+                                                    q_cu_seq_lens_,
+                                                    kv_cu_seq_lens_,
+                                                    max_query_len,
+                                                    max_seq_len,
+                                                    is_causal,
+                                                    window_size_left,
+                                                    window_size_right,
+                                                    static_cast<double>(scale),
+                                                    softcap,
+                                                    sqrt_alibi,
+                                                    alibi_slope,
+                                                    c10::nullopt,
+                                                    output_lse);
+}
+
+void batch_decode(torch::Tensor& query,
+                  const torch::Tensor& k_cache,
+                  torch::Tensor& output,
+                  const torch::Tensor& block_table,
+                  const torch::Tensor& seq_lens,
+                  const std::optional<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& mask,
+                  const std::string& compute_dtype,
+                  int64_t max_seq_len,
+                  int64_t window_size_left,
+                  int64_t window_size_right,
+                  float scale,
+                  bool return_lse,
+                  bool is_causal,
+                  int64_t kv_cache_quant_bit_size) {
+  if (query.dim() == 4) {
+    query =
+        query
+            .view({query.size(0) * query.size(1), query.size(2), query.size(3)})
+            .contiguous();
+  }
+  if (output.dim() == 4) {
+    output = output
+                 .view({output.size(0) * output.size(1),
+                        output.size(2),
+                        output.size(3)})
+                 .contiguous();
+    ;
+  }
+  auto v_cache_ = v_cache.value_or(torch::Tensor());
+  int64_t num_kv_heads = k_cache.size(1);
+  int64_t page_block_size = k_cache.size(2);
+  double softcap = 0.0;
+  bool enable_cuda_graph = false;
+  bool use_sqrt_alibi = false;
+  auto block_table_ = block_table;
+  auto k_cache_ = k_cache;
+  auto seq_lens_ = seq_lens;
+  infer::xllm_paged_attention(output,
+                              query,
+                              k_cache_,
+                              v_cache_,
+                              num_kv_heads,
+                              scale,
+                              block_table_,
+                              seq_lens_,
+                              page_block_size,
+                              max_seq_len,
+                              alibi_slope,
+                              is_causal,
+                              (int32_t)window_size_left,
+                              (int32_t)window_size_right,
+                              softcap,
+                              enable_cuda_graph,
+                              use_sqrt_alibi,
+                              c10::nullopt);
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm/kernels/ilu/fused_moe.cpp

@@ -0,0 +1,99 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <glog/logging.h>
+
+#include "ilu_ops_api.h"
+
+namespace xllm::kernel::ilu {
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& e_score_correction_bias) {
+  torch::Tensor input_ = input.to(torch::kFloat32);
+  auto reduce_weight =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kFloat).device(input.device()));
+  auto topk_indices =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kInt32).device(input.device()));
+  auto token_expert_indices =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kInt32).device(input.device()));
+
+  infer::topk_softmax(
+      reduce_weight, topk_indices, token_expert_indices, input_, false);
+
+  auto tt = reduce_weight.sum(-1);
+  if (normalize) {
+    reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1);
+  }
+  return std::make_tuple(reduce_weight, topk_indices);
+}
+
+std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
+                                       int64_t expert_num) {
+  auto src_dst = expert_id.new_empty({expert_id.numel()});
+  auto dst_src = torch::empty_like(src_dst);
+  auto expert_sizes_gpu = expert_id.new_empty({expert_num});
+  auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
+  infer::moe_compute_token_index_api(expert_id,
+                                     src_dst,
+                                     dst_src,
+                                     expert_sizes_gpu,
+                                     /*expert_mask=*/std::nullopt,
+                                     /*expert_sizes_cpu*/ std::nullopt,
+                                     /*expert_sizes_gpu*/ std::nullopt,
+                                     0,
+                                     expert_num,
+                                     expert_num);
+
+  expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
+  return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
+}
+
+torch::Tensor moe_expand_input(const torch::Tensor& input,
+                               const torch::Tensor& gather_index,
+                               const torch::Tensor& combine_idx,
+                               int64_t topk) {
+  int64_t dst_tokens = input.size(0) * topk;
+  auto output = input.new_empty({dst_tokens, input.size(1)});
+  infer::moe_expand_input(
+      output, input, combine_idx, gather_index, dst_tokens, topk);
+
+  return output;
+}
+
+torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) {
+  input = input.view({-1, weight.size(1), input.size(1)});
+  auto output = input.new_empty({input.size(0), input.size(2)});
+  infer::moe_output_reduce_sum(output,
+                               input,
+                               weight,
+                               /*mask=*/std::nullopt,
+                               /*extra_residual*/ std::nullopt,
+                               /*scaling_factor=*/1.0);
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm/kernels/ilu/ilu_ops_api.h

@@ -0,0 +1,153 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <ATen/DynamicLibrary.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <cuda_runtime.h>
+#include <glog/logging.h>
+#include <torch/all.h>
+
+#include <optional>
+
+#include "ATen/Tensor.h"
+#include "ATen/cuda/CUDAEvent.h"
+#include "c10/core/Device.h"
+#include "c10/core/DeviceGuard.h"
+#include "c10/core/GradMode.h"
+#include "c10/core/InferenceMode.h"
+#include "c10/core/MemoryFormat.h"
+#include "c10/core/ScalarType.h"
+#include "c10/core/TensorOptions.h"
+#include "c10/cuda/CUDAFunctions.h"
+#include "c10/cuda/CUDAGuard.h"
+#include "c10/cuda/CUDAStream.h"
+#include "ixformer.h"
+#include "kernels/kernels.h"
+
+// #include "utils.h"
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
+                                      torch::Tensor& key,
+                                      torch::Tensor& cos_sin_cache,
+                                      torch::Tensor& positions,
+                                      bool interleave);
+
+// act_mode only support silu, gelu, gelu_tanh
+void act_and_mul(torch::Tensor out,
+                 torch::Tensor input,
+                 const std::string& act_mode);
+
+void reshape_paged_cache(
+    torch::Tensor& key,                   //  (num_tokens, num_heads, head_size)
+    std::optional<torch::Tensor>& value,  // (num_tokens, num_heads, head_size)
+    torch::Tensor& key_cache,  // (num_blocks, num_heads, block_size, head_size)
+    std::optional<torch::Tensor>&
+        value_cache,  // (num_blocks, num_heads, block_size, head_size)
+    torch::Tensor& slot_mapping);  //(num_tokens)
+
+void batch_prefill(torch::Tensor& query,
+                   const torch::Tensor& key,
+                   const std::optional<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& v_quant_scale,
+                   const torch::Tensor& block_tables,
+                   int64_t max_query_len,
+                   int64_t max_seq_len,
+                   float scale,
+                   bool is_causal,
+                   int64_t window_size_left,
+                   int64_t window_size_right,
+                   const std::string& compute_dtype,
+                   bool return_lse);
+
+void batch_decode(torch::Tensor& query,
+                  const torch::Tensor& k_cache,
+                  torch::Tensor& output,
+                  const torch::Tensor& block_table,
+                  const torch::Tensor& seq_lens,
+                  const std::optional<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& mask,
+                  const std::string& compute_dtype,
+                  int64_t max_seq_len,
+                  int64_t window_size_left,
+                  int64_t window_size_right,
+                  float scale,
+                  bool return_lse,
+                  bool is_causal,
+                  int64_t kv_cache_quant_bit_size);
+
+void residual_layer_norm(torch::Tensor& input,
+                         torch::Tensor& output,
+                         std::optional<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& residual_out,
+                         double eps);
+
+void rms_norm(torch::Tensor& output,
+              torch::Tensor& input,
+              torch::Tensor& weight,
+              double eps);
+
+torch::Tensor matmul(torch::Tensor a,
+                     torch::Tensor b,
+                     std::optional<torch::Tensor> bias);
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& e_score_correction_bias);
+
+std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
+                                       int64_t expert_num);
+
+torch::Tensor moe_expand_input(const torch::Tensor& input,
+                               const torch::Tensor& gather_index,
+                               const torch::Tensor& combine_idx,
+                               int64_t topk);
+
+torch::Tensor group_gemm(torch::Tensor& input,
+                         torch::Tensor& weight,
+                         torch::Tensor& tokens_per_experts,
+                         const std::optional<torch::Tensor>& dst_to_src,
+                         torch::Tensor& output);
+
+torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm/kernels/ilu/ixformer.h

@@ -0,0 +1,147 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <torch/all.h>
+
+#include "ATen/Tensor.h"
+#include "utils.h"
+
+namespace ixformer::infer {
+torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    torch::Tensor& out,
+    torch::Tensor& block_tables,
+    torch::Tensor& cu_seq_q,
+    torch::Tensor& cu_seq_k,
+    int64_t max_seq_q,
+    int64_t max_seq_k,
+    bool is_causal,
+    int64_t window_left,
+    int64_t window_right,
+    double scale,
+    double softcap,
+    bool sqrt_alibi,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::optional<torch::Tensor>& sinks,
+    std::optional<torch::Tensor>& lse);
+
+void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
+
+torch::Tensor xllm_paged_attention(
+    torch::Tensor& out,
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    int64_t num_kv_heads,
+    double scale,
+    torch::Tensor& block_tables,
+    torch::Tensor& context_lens,
+    int64_t block_size,
+    int64_t max_context_len,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    bool causal,
+    int32_t window_left,
+    int32_t window_right,
+    double softcap,
+    bool enable_cuda_graph,
+    bool use_sqrt_alibi,
+    const std::optional<torch::Tensor>& sinks);
+
+torch::Tensor ixformer_linear(torch::Tensor& input,
+                              torch::Tensor& weight,
+                              int64_t act_type,
+                              const std::optional<torch::Tensor>& bias,
+                              const std::optional<torch::Tensor>& out,
+                              const std::optional<bool> persistent);
+
+torch::Tensor ixformer_linear_ex(torch::Tensor& input,
+                                 torch::Tensor& weight,
+                                 const c10::optional<torch::Tensor>& bias,
+                                 const c10::optional<torch::Tensor>& out);
+
+void xllm_reshape_and_cache(torch::Tensor& key,
+                            torch::Tensor& value,
+                            torch::Tensor& key_cache,
+                            torch::Tensor& value_cache,
+                            torch::Tensor& slot_mapping,
+                            int64_t key_token_stride,
+                            int64_t value_token_stride);
+
+void xllm_rotary_embedding(torch::Tensor& positions,
+                           torch::Tensor& query,
+                           torch::Tensor& key,
+                           int64_t head_size,
+                           torch::Tensor& cos_sin_cache,
+                           bool is_neox);
+
+void residual_rms_norm(torch::Tensor& input,
+                       torch::Tensor& residual,
+                       torch::Tensor& weight,
+                       torch::Tensor& output,
+                       torch::Tensor& residual_output,
+                       const std::optional<torch::Tensor>& fused_bias,
+                       double alpha,
+                       double eps,
+                       bool is_post);
+
+void rms_norm(torch::Tensor& input,
+              torch::Tensor& weight,
+              torch::Tensor& output,
+              const std::optional<torch::Tensor>& fused_bias,
+              double eps);
+
+void topk_softmax(torch::Tensor& topk_weights,
+                  torch::Tensor& topk_indices,
+                  torch::Tensor& token_expert_indices,
+                  torch::Tensor& gating_output,
+                  bool renormalize);
+
+void moe_compute_token_index_api(
+    torch::Tensor& topk_ids,
+    torch::Tensor& src_dst,
+    torch::Tensor& dst_src,
+    torch::Tensor& expert_sizes_gpu,
+    const c10::optional<torch::Tensor>& expert_mask,
+    const c10::optional<torch::Tensor>& expert_sizes_cpu,
+    const c10::optional<torch::Tensor>& expand_tokens_gpu,
+    int64_t start_expert_id,
+    int64_t end_expert_id,
+    int64_t num_experts);
+
+void moe_expand_input(torch::Tensor outputs,
+                      torch::Tensor inputs,
+                      torch::Tensor dst_to_src,
+                      const c10::optional<torch::Tensor>& src_to_dst,
+                      int64_t dst_tokens,
+                      int64_t expand_factor);
+
+void moe_w16a16_group_gemm(torch::Tensor output,
+                           torch::Tensor inputs,
+                           torch::Tensor weights,
+                           torch::Tensor tokens_per_experts,
+                           const c10::optional<torch::Tensor>& dst_to_src,
+                           const c10::optional<torch::Tensor>& bias,
+                           std::string format,
+                           int64_t persistent,
+                           int64_t output_n);
+
+void moe_output_reduce_sum(torch::Tensor outputs,
+                           torch::Tensor inputs,
+                           const c10::optional<torch::Tensor>& mul_weight,
+                           const c10::optional<torch::Tensor>& mask,
+                           const c10::optional<torch::Tensor>& extra_residual,
+                           double scaling_factor);
+}  // namespace ixformer::infer

upstream_ref/xllm/kernels/ilu/norm.cpp

@@ -0,0 +1,51 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "ilu_ops_api.h"
+#include "utils.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void residual_layer_norm(torch::Tensor& input,
+                         torch::Tensor& output,
+                         std::optional<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& residual_out,
+                         double eps) {
+  auto residual_ = residual.value_or(torch::zeros_like(input));
+  torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input));
+  infer::residual_rms_norm(input,
+                           residual_,
+                           weight,
+                           output,
+                           residual_out_,
+                           bias,
+                           /*alpha=*/1.0,
+                           eps,
+                           false);
+}
+
+void rms_norm(torch::Tensor& output,
+              torch::Tensor& input,
+              torch::Tensor& weight,
+              double eps) {
+  std::optional<torch::Tensor> fused_bias = std::nullopt;
+  infer::rms_norm(input, weight, output, fused_bias, eps);
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm/kernels/ilu/rope.cpp

@@ -0,0 +1,31 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "ilu_ops_api.h"
+#include "utils.h"
+
+namespace xllm::kernel::ilu {
+
+void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
+                                      torch::Tensor& key,
+                                      torch::Tensor& cos_sin_cache,
+                                      torch::Tensor& positions,
+                                      bool interleave) {
+  const int64_t head_size = cos_sin_cache.size(-1);
+  infer::xllm_rotary_embedding(
+      positions, query, key, head_size, cos_sin_cache, !interleave);
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm/kernels/ilu/utils.h

@@ -0,0 +1,63 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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
+namespace xllm::kernel::ilu {
+#undef check_tensor_contiguous
+#define check_tensor_contiguous(x, type) \
+  TORCH_CHECK(x.scalar_type() == type);  \
+  TORCH_CHECK(x.is_cuda());              \
+  TORCH_CHECK(x.is_contiguous());
+
+#undef check_tensor_half_bf_float
+#define check_tensor_half_bf_float(x)                       \
+  TORCH_CHECK(x.scalar_type() == at::ScalarType::Half ||    \
+              x.scalar_type() == at::ScalarType::Float ||   \
+              x.scalar_type() == at::ScalarType::BFloat16); \
+  TORCH_CHECK(x.is_cuda());
+
+// from torchCheckMsgImpl
+inline const char* ixformer_check_msg_impl(const char* msg) { return msg; }
+// // If there is just 1 user-provided C-string argument, use it.
+
+#define IXFORMER_CHECK_MSG(cond, type, ...)                 \
+  (ixformer_check_msg_impl(                                 \
+      "Expected " #cond                                     \
+      " to be true, but got false.  "                       \
+      "(Could this error message be improved?  If so, "     \
+      "please report an enhancement request to ixformer.)", \
+      ##__VA_ARGS__))
+
+#define IXFORMER_CHECK(cond, ...)                                            \
+  {                                                                          \
+    if (!(cond)) {                                                           \
+      std::cerr << __FILE__ << " (" << __LINE__ << ")"                       \
+                << "-" << __FUNCTION__ << " : "                              \
+                << IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \
+      throw std::runtime_error("IXFORMER_CHECK ERROR");                      \
+    }                                                                        \
+  }
+
+#undef CUINFER_CHECK
+#define CUINFER_CHECK(func)                                                \
+  do {                                                                     \
+    cuinferStatus_t status = (func);                                       \
+    if (status != CUINFER_STATUS_SUCCESS) {                                \
+      std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \
+                << ": " << cuinferGetErrorString(status) << std::endl;     \
+      throw std::runtime_error("CUINFER_CHECK ERROR");                     \
+    }                                                                      \
+  } while (0)
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm/layers/ilu/attention.cpp

@@ -0,0 +1,189 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "attention.h"
+
+#include "kernels/ilu/ilu_ops_api.h"
+#include "kernels/ops_api.h"
+
+namespace xllm {
+namespace layer {
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             float scale,
+                             int64_t num_kv_heads,
+                             int64_t sliding_window)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(head_size),
+      use_fused_mla_qkv_(false),
+      enable_lighting_indexer_(false),
+      enable_mla_(false),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             int64_t num_kv_heads,
+                             int64_t v_head_dim,
+                             int64_t sliding_window,
+                             float scale,
+                             bool use_fused_mla_qkv,
+                             bool enable_lighting_indexer,
+                             bool enable_mla)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(v_head_dim),
+      use_fused_mla_qkv_(use_fused_mla_qkv),
+      enable_lighting_indexer_(enable_lighting_indexer),
+      enable_mla_(enable_mla),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
+    const AttentionMetadata& attn_metadata,
+    torch::Tensor& query,
+    torch::Tensor& key,
+    torch::Tensor& value,
+    KVCache& kv_cache) {
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  torch::Tensor output;
+  if (enable_mla_) {
+    output = torch::empty({query.size(0), num_heads_ * v_head_dim_},
+                          query.options());
+  } else {
+    output = torch::empty_like(query);
+  }
+  if (attn_metadata.is_dummy) {
+    return std::make_tuple(output, output_lse);
+  }
+
+  bool only_prefill =
+      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
+  int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_;
+  torch::Tensor k_cache = kv_cache.get_k_cache();
+  std::optional<torch::Tensor> v_cache;
+  std::optional<torch::Tensor> v;
+  if (!enable_mla_) {
+    v = value.view({-1, num_kv_heads, head_size_});
+    v_cache = kv_cache.get_v_cache();
+  }
+
+  bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_);
+  if (!skip_process_cache) {
+    xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params;
+    reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_});
+    reshape_paged_cache_params.value = v;
+    reshape_paged_cache_params.k_cache = k_cache;
+    reshape_paged_cache_params.v_cache = v_cache;
+    reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping;
+    xllm::kernel::reshape_paged_cache(reshape_paged_cache_params);
+  }
+
+  if (enable_lighting_indexer_ || !only_prefill) {
+    decoder_forward(query, output, k_cache, v_cache, attn_metadata);
+  } else {
+    prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata);
+  }
+
+  int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_;
+  output = output.view({-1, num_heads_ * head_size});
+  return {output, output_lse};
+}
+
+void AttentionImpl::prefill_forward(torch::Tensor& query,
+                                    torch::Tensor& key,
+                                    torch::Tensor& value,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  query = query.view({-1, num_heads_, head_size_});
+  output = output.view({-1, num_heads_, head_size_v});
+  // torch::Tensor k_cache_ = k_cache;
+  // torch::Tensor v_cache_ = v_cache.value();
+  xllm::kernel::ilu::batch_prefill(query,
+                                   k_cache,
+                                   v_cache,
+                                   output,
+                                   output_lse,
+                                   attn_metadata.q_cu_seq_lens,
+                                   attn_metadata.kv_cu_seq_lens,
+                                   /*alibi_slope=*/std::nullopt,
+                                   /*attn_bias=*/std::nullopt,
+                                   /*q_quant_scale=*/std::nullopt,
+                                   /*k_quant_scale=*/std::nullopt,
+                                   /*v_quant_scale=*/std::nullopt,
+                                   attn_metadata.block_table,
+                                   attn_metadata.max_query_len,
+                                   attn_metadata.max_seq_len,
+                                   scale_,
+                                   attn_metadata.is_causal,
+                                   sliding_window_,
+                                   /*window_size_right=*/-1,
+                                   attn_metadata.compute_dtype,
+                                   /*return_lse=*/false);
+}
+
+void AttentionImpl::decoder_forward(torch::Tensor& query,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  query = query.view({-1, 1, num_heads_, head_size_});
+  output = output.view({-1, 1, num_heads_, head_size_v});
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+
+  int64_t block_aligned_max_seq_len =
+      attn_metadata.block_table.size(-1) * k_cache.size(2);
+
+  xllm::kernel::ilu::batch_decode(query,
+                                  k_cache,
+                                  output,
+                                  attn_metadata.block_table,
+                                  attn_metadata.kv_seq_lens,
+                                  v_cache,
+                                  output_lse,
+                                  /*q_quant_scale=*/std::nullopt,
+                                  /*k_quant_scale=*/std::nullopt,
+                                  /*v_quant_scale=*/std::nullopt,
+                                  /*out_quant_scale=*/std::nullopt,
+                                  /*alibi_slope=*/std::nullopt,
+                                  attn_metadata.attn_mask,
+                                  attn_metadata.compute_dtype,
+                                  block_aligned_max_seq_len,
+                                  sliding_window_,
+                                  /*window_size_right=*/-1,
+                                  scale_,
+                                  /*return_lse=*/false,
+                                  attn_metadata.is_causal,
+                                  /*kv_cache_quant_bit_size=*/-1);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/layers/ilu/attention.h

@@ -0,0 +1,82 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <torch/torch.h>
+
+#include <tuple>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_input_params.h"
+#include "layers/common/attention_metadata.h"
+
+namespace xllm {
+namespace layer {
+class AttentionImpl : public torch::nn::Module {
+ public:
+  AttentionImpl() = default;
+
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                float scale,
+                int64_t num_kv_heads,
+                int64_t sliding_window);
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                int64_t num_kv_heads,
+                int64_t v_head_dim,
+                int64_t sliding_window,
+                float scale,
+                bool use_fused_mla_qkv,
+                bool enable_lighting_indexer,
+                bool enable_mla);
+
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
+      const AttentionMetadata& attn_metadata,
+      torch::Tensor& query,
+      torch::Tensor& key,
+      torch::Tensor& value,
+      KVCache& kv_cache);
+
+  void prefill_forward(torch::Tensor& query,
+                       torch::Tensor& key,
+                       torch::Tensor& value,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+  void decoder_forward(torch::Tensor& query,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+ private:
+  int64_t num_heads_;
+  int64_t head_size_;
+  float scale_;
+  int64_t num_kv_heads_;
+  int64_t v_head_dim_;
+  bool use_fused_mla_qkv_;
+  bool enable_lighting_indexer_;
+  bool enable_mla_;
+  int64_t sliding_window_;
+};
+TORCH_MODULE(Attention);
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/layers/ilu/fused_moe.cpp

@@ -0,0 +1,797 @@
+/* Copyright 2026 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "fused_moe.h"
+
+#include <glog/logging.h>
+
+#include <iomanip>
+
+#include "common/global_flags.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "kernels/ops_api.h"
+#include "layers/common/dp_utils.h"
+#include "util/utils.h"
+
+namespace {
+
+int32_t get_dtype_size(torch::ScalarType dtype) {
+  return static_cast<int32_t>(torch::elementSize(dtype));
+}
+
+}  // namespace
+
+namespace xllm {
+namespace layer {
+
+FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
+                           const FusedMoEArgs& moe_args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options)
+    : num_total_experts_(static_cast<int64_t>(model_args.n_routed_experts())),
+      topk_(model_args.num_experts_per_tok()),
+      num_expert_group_(model_args.n_group()),
+      topk_group_(model_args.topk_group()),
+      route_scale_(model_args.routed_scaling_factor()),
+      hidden_size_(model_args.hidden_size()),
+      n_shared_experts_(model_args.n_shared_experts()),
+      is_gated_(moe_args.is_gated),
+      renormalize_(model_args.norm_topk_prob() ? 1 : 0),
+      hidden_act_(model_args.hidden_act()),
+      scoring_func_(model_args.scoring_func()),
+      quant_args_(quant_args),
+      parallel_args_(parallel_args),
+      options_(options),
+      device_(options.device()) {
+  const int64_t num_experts = num_total_experts_;
+  const int64_t intermediate_size =
+      static_cast<int64_t>(model_args.moe_intermediate_size());
+  const std::string& topk_method = model_args.topk_method();
+  int64_t ep_size = parallel_args.ep_size();
+  int64_t ep_rank = 0;
+  tp_pg_ = parallel_args.tp_group_;
+  if (ep_size > 1) {
+    ep_rank = parallel_args.moe_ep_group_->rank();
+    tp_pg_ = parallel_args.moe_tp_group_;
+  }
+
+  // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is
+  // supported
+  if (!quant_args.quant_method().empty()) {
+    if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 ||
+        !quant_args.activation_dynamic()) {
+      LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when "
+                    "quant_method is set. "
+                 << "Got quant_method=" << quant_args.quant_method()
+                 << ", bits=" << quant_args.bits()
+                 << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+    // If confirmed as smoothquant w8a8, set is_smoothquant_ to true
+    is_smoothquant_ = true;
+  } else {
+    is_smoothquant_ = false;
+  }
+
+  // Deep EP initialization check
+  enable_deep_ep_ = FLAGS_expert_parallel_degree == 2 && ep_size > 1;
+  if (enable_deep_ep_) {
+    // for now, we only implement the deep ep for decode stage.
+    // so we will assume the max_token_num is limited to max_batch_size * (1+K)
+    // K is the number of speculative tokens.
+    int64_t dispatch_token_size;
+    if (quant_args.quant_method() == "smoothquant") {
+      // float32 is for the scale of the quantized input
+      dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) +
+                            get_dtype_size(torch::kFloat32);
+    } else {
+      dispatch_token_size =
+          hidden_size_ * get_dtype_size(options_.dtype().toScalarType());
+    }
+    torch::ScalarType combine_dtype = options_.dtype().toScalarType();
+    int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype);
+    // Ensure calculation base is at least ep_size
+    int64_t effective_seqs =
+        std::max((int64_t)FLAGS_max_seqs_per_batch, (int64_t)ep_size);
+    // NOTE: FLAGS_max_seqs_per_batch represents the maximum total batch size,
+    // regardless of the dp size. To ensure robust scheduling and account
+    // for the worst-case scenario, we must guarantee that each rank is capable
+    // of handling the maximum possible number of tokens. Therefore, we define
+    // max_num_tokens_per_rank as the full maximum value, without dividing by
+    // either the rank count or the dp size.
+    int64_t max_num_tokens_per_rank =
+        (1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_;
+
+    // make sure that all layers share the same deep ep instance
+    //  so that the memory footprint is minimized
+    deep_ep_ = DeepEPManager::get_instance(dispatch_token_size,
+                                           combine_token_size,
+                                           max_num_tokens_per_rank,
+                                           num_experts,
+                                           parallel_args,
+                                           options_);
+
+    // obtain the buffer and parameters of deep ep
+    deep_ep_buffer_ = deep_ep_->get_buffer();
+    deep_ep_params_ = deep_ep_->get_params();
+
+    // intermediate buffer that can be initialized once
+    // we place these tensor here in order to speed up forward pass
+    int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv;
+    int64_t token_bytes = is_smoothquant_
+                              ? get_dtype_size(torch::kInt8)
+                              : get_dtype_size(options_.dtype().toScalarType());
+    token_bytes = token_bytes * hidden_size_;
+    int64_t head_size = n_tokens_recv * token_bytes;
+    dispatch_recv_token_tensor_head_ =
+        deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size)
+            .view({n_tokens_recv, token_bytes});
+    // input scale in smoothquant
+    if (is_smoothquant_) {
+      int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32);
+      dispatch_recv_token_tensor_tail_ =
+          deep_ep_buffer_.combine_send_token_tensor
+              .narrow(0, head_size, tail_size)
+              .view({n_tokens_recv, -1});
+    }
+  }
+
+  // calculate the number of experts per rank
+  num_experts_per_rank_ = num_experts / ep_size;
+  start_expert_id_ = ep_rank * num_experts_per_rank_;
+
+  if (topk_method == "noaux_tc") {
+    e_score_correction_bias_ = register_parameter(
+        "e_score_correction_bias", torch::empty({num_experts}, options), false);
+  }
+
+  gate_ = register_module(
+      "gate_proj",
+      ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options));
+  if (n_shared_experts_ > 0) {
+    ProcessGroup* shared_expert_pg;
+    if (parallel_args_.ep_size() > 1) {
+      // we use tp=1 for shared experts computation in deep ep mode
+      CHECK(parallel_args_.ep_size() == parallel_args_.world_size())
+          << "Models with shared experts only support ep_size equal to "
+             "world size for now.";
+      shared_expert_pg = parallel_args.moe_tp_group_;
+    } else {
+      shared_expert_pg = parallel_args.process_group_;
+    }
+    // The shared experts computation can proceed in parallel with the
+    // final communication step during the MoE computation, as long as it
+    // remains independent of any communication operations. For optimal
+    // performance, ensure that the shared experts layer on each rank always
+    // maintains its own unique weights.
+    shared_experts_ =
+        register_module("shared_experts",
+                        DenseMLP(hidden_size_,
+                                 intermediate_size * n_shared_experts_,
+                                 is_gated_,
+                                 false,
+                                 hidden_act_,
+                                 /*enable_result_reduction=*/true,
+                                 quant_args,
+                                 shared_expert_pg,
+                                 options));
+  }
+
+  // create weight buffer
+  const int64_t world_size = tp_pg_->world_size();
+  int64_t local_intermediate_size = intermediate_size / world_size;
+  if (is_smoothquant_) {
+    auto quant_option = options_.dtype(torch::kInt8);
+    auto fp_option = options_.dtype(torch::kFloat32);
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            quant_option),
+        false);
+    w13_scale_ = register_parameter(
+        "w13_scale",
+        torch::empty({num_experts_per_rank_, local_intermediate_size * 2},
+                     fp_option),
+        false);
+    // Note: We do not check enable_deep_ep_ here, since smooth quantization
+    // information may be needed even when deep EP mode is disabled. This allows
+    // retrieving quantization parameters for any subset of experts as required.
+    input_smooth_ = register_parameter(
+        "input_smooth",
+        torch::empty({num_total_experts_, hidden_size_}, fp_option),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            quant_option),
+        false);
+    w2_scale_ = register_parameter(
+        "w2_scale",
+        torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
+        false);
+    act_smooth_ = register_parameter(
+        "act_smooth",
+        torch::empty({num_experts_per_rank_, local_intermediate_size},
+                     fp_option),
+        false);
+
+  } else {
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            options_),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            options_),
+        false);
+  }
+}
+
+torch::Tensor FusedMoEImpl::create_group_gemm_output(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& group_list,
+    torch::ScalarType dtype,
+    torch::Tensor& workspace) {
+  // unify shape logic: define the target shape once.
+  bool is_3d_weight = (b.dim() != 2);
+  int64_t num_tokens = a.size(0);
+  int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0);
+
+  std::vector<int64_t> output_shape;
+  int64_t required_elements = num_tokens * out_dim;
+
+  if (is_3d_weight) {
+    output_shape = {num_tokens, out_dim};
+  } else {
+    output_shape = {group_list.size(0), num_tokens, out_dim};
+    required_elements *= group_list.size(0);
+  }
+
+  auto options = a.options().dtype(dtype);
+
+  // non-smoothquant: direct allocation
+  if (!is_smoothquant_) {
+    return torch::empty(output_shape, options);
+  }
+
+  // smoothquant: managed workspace logic
+  if (!workspace.defined()) {
+    // Lazy initialization: allocate max buffer for the lifecycle
+    // Note: accessing class members w13_ and w2_ directly for context
+    int64_t max_width = std::max(w13_.size(1), w2_.size(1));
+    workspace = torch::empty({num_tokens * max_width}, options);
+  }
+
+  // view construction
+  CHECK(workspace.numel() >= required_elements)
+      << "FusedMoE Workspace too small! Alloc: " << workspace.numel()
+      << ", Req: " << required_elements;
+
+  // utilize the pre-calculated output_shape
+  return workspace.slice(0, 0, required_elements).view(output_shape);
+}
+
+torch::Tensor FusedMoEImpl::select_experts(
+    const torch::Tensor& hidden_states_2d,
+    const torch::Tensor& router_logits_2d,
+    SelectedExpertInfo& selected_expert_info,
+    bool enable_all2all_communication) {
+  // prepare the parameters for select_experts
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1: apply softmax topk or sigmoid topk / routing logic
+  torch::Tensor reduce_weight;
+  torch::Tensor expert_id;
+  {
+    xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
+    moe_active_topk_params.input = router_logits_2d;
+    moe_active_topk_params.topk = topk_;
+    moe_active_topk_params.num_expert_group = num_expert_group_;
+    moe_active_topk_params.topk_group = topk_group_;
+    moe_active_topk_params.normalize = renormalize_;
+    moe_active_topk_params.normed_by = "topk_logit";
+    moe_active_topk_params.scoring_func = scoring_func_;
+    moe_active_topk_params.route_scale = route_scale_;
+    moe_active_topk_params.e_score_correction_bias = e_score_correction_bias;
+    std::tie(reduce_weight, expert_id) =
+        xllm::kernel::moe_active_topk(moe_active_topk_params);
+  }
+
+  // Step 2: generate expert ids
+  torch::Tensor gather_idx;
+  torch::Tensor combine_idx;
+  torch::Tensor token_count;
+  std::optional<torch::Tensor> cusum_token_count;
+  {
+    xllm::kernel::MoeGenIdxParams moe_gen_idx_params;
+    moe_gen_idx_params.expert_id = expert_id;
+    moe_gen_idx_params.expert_num = num_total_experts_;
+    std::vector<torch::Tensor> output_vec =
+        xllm::kernel::moe_gen_idx(moe_gen_idx_params);
+    gather_idx = output_vec[0];
+    combine_idx = output_vec[1];
+    token_count = output_vec[2];
+    // during all2all communication, we do not need cusum_token_count in the
+    // following computation
+    if (enable_all2all_communication) {
+      cusum_token_count = std::nullopt;
+    } else {
+      cusum_token_count = output_vec[3];
+    }
+  }
+
+  // Step 3: expand and quantize input if needed
+  torch::Tensor expand_hidden_states;
+  torch::Tensor hidden_states_scale;
+  torch::Tensor token_count_slice;
+  // all2all related variables
+  torch::Tensor dispatch_send_token_tensor;
+  // in all2all, the input is scattered, so there is no need to slice the token
+  // count, and we can use the dispatch buffer directly
+  if (enable_all2all_communication) {
+    token_count_slice = token_count;
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    int64_t dispatch_bytes =
+        num_token_expand * deep_ep_params_.dispatch_token_size;
+    dispatch_send_token_tensor =
+        deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes)
+            .view({num_token_expand, deep_ep_params_.dispatch_token_size});
+  } else {
+    token_count_slice =
+        token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size);
+  }
+
+  if (is_smoothquant_) {
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = hidden_states_2d;
+    // use dispatch_send_token_tensor buffer for input
+    //  to reduce memory footprint
+    if (enable_all2all_communication) {
+      scaled_quantize_params.smooth = input_smooth_;
+      scaled_quantize_params.output =
+          dispatch_send_token_tensor.slice(1, 0, hidden_size_);
+    } else {
+      scaled_quantize_params.smooth = input_smooth_.slice(
+          0, start_expert_id_, start_expert_id_ + expert_size);
+      scaled_quantize_params.gather_index_start_position =
+          cusum_token_count.value().index({start_expert_id_}).unsqueeze(0);
+    }
+    scaled_quantize_params.token_count = token_count_slice;
+    scaled_quantize_params.gather_index = gather_idx;
+    scaled_quantize_params.act_mode = "none";
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = false;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(expand_hidden_states, hidden_states_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+    if (enable_all2all_communication) {
+      // since view_as_dtype has not supported stride yet,
+      //  we need to copy the scale output to the dispatch buffer
+      torch::Tensor dispatch_scale_slice =
+          dispatch_send_token_tensor.slice(1, hidden_size_);
+      torch::Tensor hidden_states_scale_bytes =
+          view_as_dtype(hidden_states_scale, torch::kInt8)
+              .view_as(dispatch_scale_slice);
+      dispatch_scale_slice.copy_(hidden_states_scale_bytes);
+    }
+  } else {
+    xllm::kernel::MoeExpandInputParams moe_expand_input_params;
+    moe_expand_input_params.input = hidden_states_2d;
+    moe_expand_input_params.gather_index = gather_idx;
+    moe_expand_input_params.combine_idx = combine_idx;
+    moe_expand_input_params.topk = topk_;
+    expand_hidden_states =
+        xllm::kernel::moe_expand_input(moe_expand_input_params);
+    if (enable_all2all_communication) {
+      // use copy to place the output inside the dispatch buffer
+      torch::Tensor dispatch_tensor =
+          view_as_dtype(expand_hidden_states, torch::kChar);
+      dispatch_send_token_tensor.copy_(dispatch_tensor);
+    }
+  }
+
+  // collect the selected tensor
+  selected_expert_info.reduce_weight = reduce_weight;
+  selected_expert_info.combine_idx = combine_idx;
+  selected_expert_info.token_count_slice = token_count_slice;
+  selected_expert_info.cusum_token_count = cusum_token_count;
+  if (is_smoothquant_) {
+    selected_expert_info.input_scale = hidden_states_scale;
+  }
+
+  return expand_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states,
+                                            const torch::Tensor& router_logits,
+                                            bool enable_all2all_communication) {
+  if (!stream_initialized_) {
+    // update device record
+    device_ = xllm::Device(hidden_states.device());
+
+    // acquire streams from the pool again
+    routed_stream_ = device_.get_stream_from_pool();
+    shared_stream_ = device_.get_stream_from_pool();
+    stream_initialized_ = true;
+  }
+
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+
+  // prepare the parameters for MoE computation
+  torch::Tensor shared_expert_output;
+  torch::IntArrayRef hidden_states_shape = hidden_states.sizes();
+  torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType();
+  torch::Tensor hidden_states_2d =
+      hidden_states.reshape({-1, hidden_states.size(-1)});
+  torch::Tensor router_logits_2d =
+      router_logits.reshape({-1, router_logits.size(-1)});
+  int64_t group_gemm_max_dim = enable_all2all_communication
+                                   ? deep_ep_params_.max_num_tokens_recv / topk_
+                                   : hidden_states_2d.size(0);
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1-3: select experts
+  SelectedExpertInfo selected_expert_info;
+  torch::Tensor expand_hidden_states =
+      select_experts(hidden_states_2d,
+                     router_logits_2d,
+                     selected_expert_info,
+                     enable_all2all_communication);
+
+  // Communciation Step 1: Dipatch
+  // intermediate outputs that are used both in dispatch and combine
+  torch::Tensor gather_by_rank_index;
+  torch::Tensor token_sum;
+  if (enable_all2all_communication) {
+    int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_;
+
+    // 1. Dispatch Step: Generate layout and send data
+    deep_ep_->dispatch_step(dispatch_token_num,
+                            selected_expert_info.token_count_slice);
+
+    // 2. Process Result: Generate indices and unpack to computation buffer
+    // use the buffer during initialization for the output
+    expand_hidden_states = dispatch_recv_token_tensor_head_;
+    std::optional<torch::Tensor> output_tail = std::nullopt;
+    if (is_smoothquant_) {
+      output_tail = dispatch_recv_token_tensor_tail_;
+      // update selected_expert_info with the tail (input scale)
+      selected_expert_info.input_scale = output_tail;
+    }
+
+    DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result(
+        num_experts_per_rank_, expand_hidden_states, output_tail);
+
+    // Extract metadata for subsequent steps
+    gather_by_rank_index = deep_ep_meta.gather_rank_index;
+    selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice;
+    token_sum = deep_ep_meta.token_sum;
+  }
+
+  // common gemm workspace for reduce memory footprint
+  torch::Tensor gemm_workspace;
+
+  // Step 4: group gemm 1
+  torch::Tensor gemm1_out =
+      create_group_gemm_output(expand_hidden_states,
+                               w13_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    torch::ScalarType a_dtype =
+        is_smoothquant_ ? torch::kInt8 : hidden_states_dtype;
+    group_gemm_params.a =
+        view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_});
+    group_gemm_params.b = w13_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      torch::Tensor a_scale =
+          selected_expert_info.input_scale.value().flatten();
+      selected_expert_info.input_scale =
+          view_as_dtype(a_scale, torch::kFloat32);
+      group_gemm_params.a_scale = selected_expert_info.input_scale;
+      group_gemm_params.b_scale = w13_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm1_out;
+    group_gemm_params.combine_idx = std::nullopt;
+    gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Step 5: activation or scaled quantization(fused with activation)
+  torch::Tensor act_out;
+  torch::Tensor act_out_scale;
+  if (is_smoothquant_) {
+    int64_t slice_dim = gemm1_out.size(1);
+    if (is_gated_) slice_dim /= 2;
+    // slice operation is a view, does not take up extra memory, but points to
+    // the same memory
+    act_out = expand_hidden_states.slice(1, 0, slice_dim);
+    act_out_scale =
+        selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0));
+    // call scaled quantization kernel (also fused with activation)
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = gemm1_out;
+    scaled_quantize_params.smooth = act_smooth_;
+    scaled_quantize_params.token_count = selected_expert_info.token_count_slice;
+    scaled_quantize_params.output = act_out;
+    scaled_quantize_params.output_scale = act_out_scale;
+    scaled_quantize_params.act_mode = hidden_act_;
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = is_gated_;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(act_out, act_out_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+  } else {
+    act_out = is_gated_
+                  ? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous()
+                  : gemm1_out;
+    // call activation kernel
+    xllm::kernel::ActivationParams activation_params;
+    activation_params.input = gemm1_out;
+    activation_params.output = act_out;
+    activation_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    activation_params.act_mode = hidden_act_;
+    activation_params.is_gated = is_gated_;
+    activation_params.start_expert_id = start_expert_id_;
+    activation_params.expert_size = expert_size;
+    xllm::kernel::active(activation_params);
+  }
+
+  // Step 6: group gemm 2
+  torch::Tensor gemm2_out =
+      create_group_gemm_output(act_out,
+                               w2_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    group_gemm_params.a = act_out;
+    group_gemm_params.b = w2_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      group_gemm_params.a_scale = act_out_scale;
+      group_gemm_params.b_scale = w2_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm2_out;
+    group_gemm_params.combine_idx = selected_expert_info.combine_idx;
+    gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Communciation Step 2: Combine
+  if (enable_all2all_communication) {
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    // Delegate pack, layout generation and combine to DeepEP
+    torch::Tensor combine_send_layout =
+        deep_ep_->combine_step_pack(gemm2_out,
+                                    gather_by_rank_index,
+                                    token_sum,
+                                    hidden_size_,
+                                    hidden_states_dtype);
+
+    // create a wait event for the current stream to finish computation
+    auto current_stream = device_.current_stream();
+    routed_stream_->wait_stream(*current_stream);
+    // pure communciation kernel: dispatch
+    {
+      torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+      gemm2_out = deep_ep_->combine_step_comm(combine_send_layout,
+                                              num_token_expand,
+                                              hidden_size_,
+                                              hidden_states_dtype);
+    }
+
+    // pure computation kernel: shared experts
+    if (n_shared_experts_ > 0) {
+      shared_stream_->wait_stream(*current_stream);
+      torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+      shared_expert_output = shared_experts_(hidden_states);
+    }
+
+    // join for parallelization
+    current_stream->wait_stream(*routed_stream_);
+    if (n_shared_experts_ > 0) {
+      current_stream->wait_stream(*shared_stream_);
+    }
+  }
+
+  // After group gemm is finished, some tensors are no
+  // longer needed. We must explicitly release the memory.
+  expand_hidden_states = torch::Tensor();
+  selected_expert_info.input_scale = std::nullopt;
+  act_out = torch::Tensor();
+
+  // Step 7: combine the intermediate results and get the final hidden states
+  torch::Tensor final_hidden_states;
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::MoeCombineResultParams moe_combine_result_params;
+    moe_combine_result_params.input = gemm2_out;
+    moe_combine_result_params.reduce_weight =
+        selected_expert_info.reduce_weight;
+    moe_combine_result_params.gather_ids = selected_expert_info.combine_idx;
+    moe_combine_result_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    moe_combine_result_params.start_expert_id = start_expert_id_;
+    moe_combine_result_params.expert_size = expert_size;
+    moe_combine_result_params.bias = std::nullopt;
+    // if all2all communication is enabled and shared output is provided,
+    //  we will fused the add up to combine result
+    if (enable_all2all_communication && n_shared_experts_ > 0) {
+      moe_combine_result_params.residual =
+          shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+    }
+
+    final_hidden_states =
+        xllm::kernel::moe_combine_result(moe_combine_result_params);
+  }
+
+  // reshape the final hidden states to the original shape
+  final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
+
+  if (enable_all2all_communication) {
+    return final_hidden_states;
+  }
+
+  // Communciation Step 3: AllReduce for non-all2all communication
+  // shared experts can be parallelized with the final communication step
+  // during moe computation.
+  auto current_stream = device_.current_stream();
+  routed_stream_->wait_stream(*current_stream);
+  {
+    torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+    if (tp_pg_->world_size() > 1) {
+      final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
+    }
+    if (parallel_args_.ep_size() > 1) {
+      final_hidden_states = parallel_state::reduce(
+          final_hidden_states, parallel_args_.moe_ep_group_);
+    }
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_stream_->wait_stream(*current_stream);
+    torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+    // for non all2all, we compute the shared experts parallelized with the
+    // final communication step
+    shared_expert_output = shared_experts_(hidden_states);
+    shared_expert_output =
+        shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+  }
+
+  // join for parallelization
+  current_stream->wait_stream(*routed_stream_);
+  if (n_shared_experts_ > 0) {
+    current_stream->wait_stream(*shared_stream_);
+    final_hidden_states += shared_expert_output;
+  }
+
+  return final_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
+                                    const ModelInputParams& input_params) {
+  // we only support all2all communication for decode stage for now
+  bool enable_all2all_communication =
+      enable_deep_ep_ && std::all_of(input_params.dp_is_decode.begin(),
+                                     input_params.dp_is_decode.end(),
+                                     [](int32_t val) { return val == 1; });
+
+  bool is_dp_ep_parallel =
+      parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1;
+  // during all2all communication, the output has been
+  //  gathered and sliced by dispatch and combine steps,
+  //  so we do not need to gather input and slice output again
+  bool need_gather_and_slice =
+      is_dp_ep_parallel && !enable_all2all_communication;
+
+  auto input = hidden_states;
+  if (need_gather_and_slice) {
+    input = parallel_state::gather(input,
+                                   parallel_args_.dp_local_process_group_,
+                                   input_params.dp_global_token_nums);
+  }
+  // MoE Gate
+  auto router_logits = gate_(input);
+
+  // MoE Experts
+  auto output =
+      forward_experts(input, router_logits, enable_all2all_communication);
+
+  if (need_gather_and_slice) {
+    output = get_dp_local_slice(output, input_params, parallel_args_);
+  }
+
+  return output;
+}
+
+void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) {
+  if (e_score_correction_bias_.defined() &&
+      !e_score_correction_bias_is_loaded_) {
+    LOAD_WEIGHT(e_score_correction_bias);
+  }
+}
+
+void FusedMoEImpl::load_experts(const StateDict& state_dict) {
+  const int64_t rank = tp_pg_->rank();
+  const int64_t world_size = tp_pg_->world_size();
+  const int64_t start_expert_id = start_expert_id_;
+  const int64_t num_experts_per_rank = num_experts_per_rank_;
+  const int64_t num_total_experts = num_total_experts_;
+  std::vector<std::string> prefixes = {"gate_proj.", "up_proj."};
+  if (is_smoothquant_) {
+    LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13);
+    LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale);
+    // When supporting DeepEP All2All mode,
+    // we need to load the complete set of expert weights corresponding to
+    // "up_proj.smooth". Note that even if deep EP mode is not enabled, it
+    // remains possible to retrieve the smooth quantization information for a
+    // subset of experts. Therefore, we intentionally do not check whether
+    // deep_ep_ is enabled in this case.
+    LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1);
+    LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0);
+  } else {
+    LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13);
+    LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1);
+  }
+}
+
+void FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
+  if (state_dict.size() == 0) {
+    return;
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_experts_->load_state_dict(
+        state_dict.get_dict_with_prefix("shared_experts."));
+  }
+  gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
+  load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
+  load_experts(state_dict.get_dict_with_prefix("experts."));
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/layers/ilu/fused_moe.h

@@ -0,0 +1,131 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "layers/common/deep_ep.h"
+#include "layers/common/dense_mlp.h"
+#include "layers/common/fused_moe_base.h"
+#include "layers/common/linear.h"
+#include "platform/device.h"
+#include "util/tensor_helper.h"
+
+namespace xllm {
+namespace layer {
+
+class FusedMoEImpl : public torch::nn::Module {
+ public:
+  FusedMoEImpl() = default;
+  FusedMoEImpl(const ModelArgs& model_args,
+               const FusedMoEArgs& moe_args,
+               const QuantArgs& quant_args,
+               const ParallelArgs& parallel_args,
+               const torch::TensorOptions& options);
+
+  torch::Tensor forward_experts(const torch::Tensor& hidden_states,
+                                const torch::Tensor& router_logits,
+                                bool enable_all2all_communication);
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const ModelInputParams& input_params);
+  void load_state_dict(const StateDict& state_dict);
+
+ private:
+  // struct to store the selected expert info
+  struct SelectedExpertInfo {
+    torch::Tensor reduce_weight;
+    torch::Tensor combine_idx;
+    torch::Tensor token_count_slice;
+    std::optional<torch::Tensor> cusum_token_count;
+    std::optional<torch::Tensor> input_scale;
+  };
+
+  // initial steps for MoE computation, select the experts for each token
+  torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
+                               const torch::Tensor& router_logits_2d,
+                               SelectedExpertInfo& selected_expert_info,
+                               bool enable_all2all_communication);
+
+ private:
+  int64_t num_total_experts_;
+  int64_t topk_;
+  int64_t num_expert_group_;
+  int64_t topk_group_;
+  double route_scale_;
+  int64_t hidden_size_;
+  int64_t n_shared_experts_;
+  bool is_gated_;
+  int64_t renormalize_;
+  std::string hidden_act_;
+  std::string scoring_func_;
+  bool is_smoothquant_;
+
+  int64_t num_experts_per_rank_;
+  int64_t start_expert_id_;
+
+  // Deep EP related parameters
+  bool enable_deep_ep_;
+  DeepEPBuffer deep_ep_buffer_;
+  DeepEPParams deep_ep_params_;
+  torch::Tensor dispatch_recv_token_tensor_head_;
+  torch::Tensor dispatch_recv_token_tensor_tail_;
+
+  // steams for parallel shared experts
+  std::unique_ptr<Stream> shared_stream_;
+  std::unique_ptr<Stream> routed_stream_;
+  xllm::Device device_;
+  bool stream_initialized_ = false;
+
+  ReplicatedLinear gate_{nullptr};
+  DenseMLP shared_experts_{nullptr};
+  DeepEP deep_ep_{nullptr};
+
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_;
+  torch::TensorOptions options_;
+  ProcessGroup* tp_pg_;
+
+  DEFINE_WEIGHT(w13);
+  DEFINE_FUSED_WEIGHT(w1);
+  DEFINE_FUSED_WEIGHT(w3);
+  DEFINE_FUSED_WEIGHT(w2);
+  DEFINE_WEIGHT(e_score_correction_bias);
+  DEFINE_WEIGHT(w13_scale);
+  DEFINE_FUSED_WEIGHT(w1_scale);
+  DEFINE_FUSED_WEIGHT(w3_scale);
+  DEFINE_FUSED_WEIGHT(w2_scale);
+  DEFINE_FUSED_WEIGHT(input_smooth);
+  DEFINE_FUSED_WEIGHT(act_smooth);
+
+  void load_e_score_correction_bias(const StateDict& state_dict);
+  void load_experts(const StateDict& state_dict);
+  // create the group gemm output tensor with the workspace
+  torch::Tensor create_group_gemm_output(const torch::Tensor& a,
+                                         const torch::Tensor& b,
+                                         const torch::Tensor& group_list,
+                                         torch::ScalarType dtype,
+                                         torch::Tensor& workspace);
+};
+TORCH_MODULE(FusedMoE);
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/layers/npu_torch/qwen3_5_gated_delta_net.cpp

@@ -0,0 +1,185 @@
+/* Copyright 2026 The xLLM Authors. 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
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+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 "qwen3_5_gated_delta_net.h"
+
+#include <glog/logging.h>
+
+namespace xllm {
+namespace layer {
+
+Qwen3_5GatedDeltaNetImpl::Qwen3_5GatedDeltaNetImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options)
+    : Qwen3NextGatedDeltaNetImpl(args,
+                                 quant_args,
+                                 parallel_args,
+                                 options,
+                                 /*init_projections=*/false) {
+  in_proj_qkv_ = register_module("in_proj_qkv",
+                                 ColumnParallelLinear(args.hidden_size(),
+                                                      k_size_ * 2 + v_size_,
+                                                      /*bias=*/false,
+                                                      /*gather_output=*/false,
+                                                      quant_args,
+                                                      parallel_args.tp_group_,
+                                                      options));
+  in_proj_z_ = register_module("in_proj_z",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    v_size_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+  in_proj_b_ = register_module("in_proj_b",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    num_v_heads_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+  in_proj_a_ = register_module("in_proj_a",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    num_v_heads_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+}
+
+torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_qkvz_from_split_activations(
+    const torch::Tensor& qkv,
+    const torch::Tensor& z) const {
+  CHECK_EQ(qkv.dim(), 3) << "Expected qkv activation to be 3D, got "
+                         << qkv.sizes();
+  CHECK_EQ(z.dim(), 3) << "Expected z activation to be 3D, got " << z.sizes();
+  CHECK_EQ(qkv.size(0), z.size(0)) << "qkv/z batch size mismatch.";
+  CHECK_EQ(qkv.size(1), z.size(1)) << "qkv/z sequence size mismatch.";
+  CHECK_EQ(qkv.size(2), (2 * k_size_ + v_size_) / tp_size_)
+      << "Unexpected qkv hidden size for Qwen3.5.";
+  CHECK_EQ(z.size(2), v_size_ / tp_size_)
+      << "Unexpected z hidden size for Qwen3.5.";
+  CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
+  CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
+      << "linear_num_value_heads must be divisible by linear_num_key_heads.";
+
+  const int64_t bs = qkv.size(0);
+  const int64_t seqlen = qkv.size(1);
+  const int64_t local_k_heads = num_k_heads_ / tp_size_;
+  const int64_t local_v_heads = num_v_heads_ / tp_size_;
+  const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
+
+  auto qkv_split = torch::split(
+      qkv, {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}, 2);
+  auto q = qkv_split[0].view({bs, seqlen, local_k_heads, head_k_dim_});
+  auto k = qkv_split[1].view({bs, seqlen, local_k_heads, head_k_dim_});
+  auto v = qkv_split[2].view({bs, seqlen, local_v_heads, head_v_dim_});
+  auto z_view = z.view({bs, seqlen, local_v_heads, head_v_dim_});
+
+  v = v.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
+  z_view =
+      z_view.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
+
+  return torch::cat({q, k, v, z_view}, -1).view({bs, seqlen, -1}).contiguous();
+}
+
+torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
+    const torch::Tensor& b,
+    const torch::Tensor& a) const {
+  CHECK_EQ(b.dim(), 3) << "Expected b activation to be 3D, got " << b.sizes();
+  CHECK_EQ(a.dim(), 3) << "Expected a activation to be 3D, got " << a.sizes();
+  CHECK_EQ(b.size(0), a.size(0)) << "b/a batch size mismatch.";
+  CHECK_EQ(b.size(1), a.size(1)) << "b/a sequence size mismatch.";
+  CHECK_EQ(b.size(2), num_v_heads_ / tp_size_)
+      << "Unexpected b hidden size for Qwen3.5.";
+  CHECK_EQ(a.size(2), num_v_heads_ / tp_size_)
+      << "Unexpected a hidden size for Qwen3.5.";
+  CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
+  CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
+      << "linear_num_value_heads must be divisible by linear_num_key_heads.";
+
+  const int64_t bs = b.size(0);
+  const int64_t seqlen = b.size(1);
+  const int64_t local_k_heads = num_k_heads_ / tp_size_;
+  const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
+
+  auto b_view = b.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
+  auto a_view = a.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
+  return torch::cat({b_view, a_view}, -1).view({bs, seqlen, -1}).contiguous();
+}
+
+std::pair<torch::Tensor, torch::Tensor>
+Qwen3_5GatedDeltaNetImpl::project_padded_inputs(
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata) {
+  auto qkv = reshape_qkvz_with_pad(attn_metadata,
+                                   in_proj_qkv_->forward(hidden_states));
+  auto z_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states));
+  auto b_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states));
+  auto a_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states));
+  return {merge_qkvz_from_split_activations(qkv, z_proj),
+          merge_ba_from_split_activations(b_proj, a_proj)};
+}
+
+void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(
+    const StateDict& state_dict) {
+  auto in_proj_qkv_state_dict = state_dict.get_dict_with_prefix("in_proj_qkv.");
+  if (in_proj_qkv_state_dict.size() > 0 && !in_proj_qkv_->is_weight_loaded()) {
+    in_proj_qkv_->load_state_dict(
+        in_proj_qkv_state_dict,
+        /*shard_tensor_count=*/3,
+        /*shard_sizes=*/
+        {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_});
+  }
+
+  auto in_proj_z_state_dict = state_dict.get_dict_with_prefix("in_proj_z.");
+  if (in_proj_z_state_dict.size() > 0 && !in_proj_z_->is_weight_loaded()) {
+    in_proj_z_->load_state_dict(in_proj_z_state_dict);
+  }
+
+  auto in_proj_b_state_dict = state_dict.get_dict_with_prefix("in_proj_b.");
+  if (in_proj_b_state_dict.size() > 0 && !in_proj_b_->is_weight_loaded()) {
+    in_proj_b_->load_state_dict(in_proj_b_state_dict);
+  }
+
+  auto in_proj_a_state_dict = state_dict.get_dict_with_prefix("in_proj_a.");
+  if (in_proj_a_state_dict.size() > 0 && !in_proj_a_->is_weight_loaded()) {
+    in_proj_a_->load_state_dict(in_proj_a_state_dict);
+  }
+}
+
+void Qwen3_5GatedDeltaNetImpl::verify_projection_weights(
+    const std::string& prefix) const {
+  CHECK(in_proj_qkv_ && in_proj_qkv_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_qkv.weight";
+  CHECK(in_proj_z_ && in_proj_z_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_z.weight";
+  CHECK(in_proj_b_ && in_proj_b_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_b.weight";
+  CHECK(in_proj_a_ && in_proj_a_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_a.weight";
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/layers/npu_torch/qwen3_5_gated_delta_net.h

@@ -0,0 +1,58 @@
+/* Copyright 2026 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <torch/torch.h>
+
+#include <string>
+#include <utility>
+
+#include "qwen3_next_gated_delta_net.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
+ public:
+  Qwen3_5GatedDeltaNetImpl() = default;
+  Qwen3_5GatedDeltaNetImpl(const ModelArgs& args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options);
+
+ protected:
+  std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
+      const torch::Tensor& hidden_states,
+      const AttentionMetadata& attn_metadata) override;
+
+  void load_projection_state_dict(const StateDict& state_dict) override;
+  void verify_projection_weights(const std::string& prefix) const override;
+
+ private:
+  torch::Tensor merge_qkvz_from_split_activations(const torch::Tensor& qkv,
+                                                  const torch::Tensor& z) const;
+  torch::Tensor merge_ba_from_split_activations(const torch::Tensor& b,
+                                                const torch::Tensor& a) const;
+
+  ColumnParallelLinear in_proj_qkv_{nullptr};
+  ColumnParallelLinear in_proj_z_{nullptr};
+  ColumnParallelLinear in_proj_b_{nullptr};
+  ColumnParallelLinear in_proj_a_{nullptr};
+};
+TORCH_MODULE(Qwen3_5GatedDeltaNet);
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/layers/npu_torch/qwen3_gated_delta_net_base.cpp

@@ -0,0 +1,576 @@
+/* Copyright 2026 The xLLM Authors. 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
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+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 "qwen3_gated_delta_net_base.h"
+
+#include <glog/logging.h>
+#include <torch/torch.h>
+
+#include <tuple>
+
+#include "xllm/core/kernels/ops_api.h"
+
+namespace xllm {
+namespace layer {
+
+namespace {
+torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) {
+  auto norm = torch::sqrt(torch::sum(torch::square(x), dim, true) + eps);
+  return x / norm;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor g,
+    torch::Tensor beta,
+    std::optional<torch::Tensor> initial_state,
+    bool output_final_state = true,
+    bool use_qk_l2norm_in_kernel = true) {
+  auto initial_dtype = query.dtype();
+
+  if (use_qk_l2norm_in_kernel) {
+    query = l2norm(query, -1, 1e-6);
+    key = l2norm(key, -1, 1e-6);
+  }
+
+  auto to_float32_and_transpose = [](torch::Tensor x) {
+    return x.transpose(1, 2).contiguous().to(torch::kFloat32);
+  };
+  query = to_float32_and_transpose(query);
+  key = to_float32_and_transpose(key);
+  value = to_float32_and_transpose(value);
+  beta = to_float32_and_transpose(beta);
+  g = to_float32_and_transpose(g);
+
+  int64_t batch_size = key.size(0);
+  int64_t num_heads = key.size(1);
+  int64_t sequence_length = key.size(2);
+  int64_t k_head_dim = key.size(3);
+  int64_t v_head_dim = value.size(3);
+
+  float scale_val = 1.0 / std::sqrt(static_cast<float>(query.size(-1)));
+  torch::Tensor scale = torch::tensor(scale_val, query.options());
+  query = query * scale;
+  torch::Tensor core_attn_out = torch::zeros(
+      {batch_size, num_heads, sequence_length, v_head_dim},
+      torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
+  torch::Tensor last_recurrent_state;
+  if (!initial_state.has_value()) {
+    last_recurrent_state = torch::zeros(
+        {batch_size, num_heads, k_head_dim, v_head_dim},
+        torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
+  } else {
+    last_recurrent_state =
+        initial_state.value().to(value.device(), torch::kFloat32);
+  }
+
+  for (int64_t i = 0; i < sequence_length; ++i) {
+    torch::Tensor q_t = query.select(2, i);
+    torch::Tensor k_t = key.select(2, i);
+    torch::Tensor v_t = value.select(2, i);
+    torch::Tensor g_t = g.select(2, i).exp().unsqueeze(-1).unsqueeze(-1);
+    torch::Tensor beta_t = beta.select(2, i).unsqueeze(-1);
+    last_recurrent_state = last_recurrent_state * g_t;
+    torch::Tensor kv_mem =
+        torch::sum(last_recurrent_state * k_t.unsqueeze(-1), -2);
+    torch::Tensor delta = (v_t - kv_mem) * beta_t;
+    last_recurrent_state =
+        last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2);
+    core_attn_out.select(2, i) =
+        torch::sum(last_recurrent_state * q_t.unsqueeze(-1), -2);
+  }
+
+  core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
+  return std::make_tuple(core_attn_out, last_recurrent_state);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor g,
+    torch::Tensor beta,
+    int64_t chunk_size = 64,
+    c10::optional<torch::Tensor> initial_state = c10::nullopt,
+    bool output_final_state = true,
+    bool use_qk_l2norm_in_kernel = true) {
+  auto initial_dtype = query.dtype();
+  if (use_qk_l2norm_in_kernel) {
+    query = l2norm(query, -1, 1e-6);
+    key = l2norm(key, -1, 1e-6);
+  }
+  auto to_float32 = [](torch::Tensor x) {
+    return x.transpose(1, 2).contiguous().to(torch::kFloat32);
+  };
+
+  query = to_float32(query);
+  key = to_float32(key);
+  value = to_float32(value);
+  beta = to_float32(beta);
+  g = to_float32(g);
+
+  auto batch_size = query.size(0);
+  auto num_heads = query.size(1);
+  auto sequence_length = query.size(2);
+  auto k_head_dim = key.size(-1);
+  auto v_head_dim = value.size(-1);
+
+  int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
+  query = torch::nn::functional::pad(
+      query, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  key = torch::nn::functional::pad(
+      key, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  value = torch::nn::functional::pad(
+      value, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  beta = torch::nn::functional::pad(
+      beta, torch::nn::functional::PadFuncOptions({0, pad_size}));
+  g = torch::nn::functional::pad(
+      g, torch::nn::functional::PadFuncOptions({0, pad_size}));
+
+  int64_t total_sequence_length = sequence_length + pad_size;
+  float scale = 1.0 / std::sqrt(static_cast<float>(query.size(-1)));
+  query = query * scale;
+  auto v_beta = value * beta.unsqueeze(-1);
+  auto k_beta = key * beta.unsqueeze(-1);
+  auto reshape_to_chunks = [chunk_size](torch::Tensor x) {
+    auto shape = x.sizes();
+    std::vector<int64_t> new_shape = {
+        shape[0], shape[1], shape[2] / chunk_size, chunk_size, shape[3]};
+    return x.reshape(new_shape);
+  };
+
+  query = reshape_to_chunks(query);
+  key = reshape_to_chunks(key);
+  value = reshape_to_chunks(value);
+  k_beta = reshape_to_chunks(k_beta);
+  v_beta = reshape_to_chunks(v_beta);
+
+  auto g_shape = g.sizes();
+  std::vector<int64_t> g_new_shape = {
+      g_shape[0], g_shape[1], g_shape[2] / chunk_size, chunk_size};
+  g = g.reshape(g_new_shape);
+  auto mask = torch::triu(
+      torch::ones(
+          {chunk_size, chunk_size},
+          torch::TensorOptions().dtype(torch::kBool).device(query.device())),
+      0);
+
+  g = g.cumsum(-1);
+  auto g_diff = g.unsqueeze(-1) - g.unsqueeze(-2);
+  auto decay_mask = g_diff.tril().exp().to(torch::kFloat32);
+  decay_mask = decay_mask.tril();
+  auto attn = -(torch::matmul(k_beta, key.transpose(-1, -2)) * decay_mask)
+                   .masked_fill(mask, 0.0);
+  for (int64_t i = 1; i < chunk_size; ++i) {
+    if (!attn.is_contiguous()) {
+      attn = attn.contiguous();
+    }
+    auto row = attn.slice(-2, i, i + 1)
+                   .slice(-1, 0, i)
+                   .squeeze(-2)
+                   .clone()
+                   .contiguous();
+    auto sub = attn.slice(-2, 0, i).slice(-1, 0, i).clone().contiguous();
+    auto row_unsq = row.unsqueeze(-1).contiguous();
+    auto row_sub_mul = (row_unsq * sub).contiguous();
+    auto row_sub_sum = row_sub_mul.sum(-2).contiguous();
+    auto row_final = (row + row_sub_sum).contiguous();
+    attn.index_put_({torch::indexing::Ellipsis,
+                     torch::indexing::Slice(i, i + 1),
+                     torch::indexing::Slice(0, i)},
+                    row_final.unsqueeze(-2));
+  }
+
+  attn = attn +
+         torch::eye(
+             chunk_size,
+             torch::TensorOptions().dtype(attn.dtype()).device(attn.device()));
+  value = torch::matmul(attn, v_beta);
+  auto k_cumdecay = torch::matmul(attn, (k_beta * g.exp().unsqueeze(-1)));
+  torch::Tensor last_recurrent_state;
+  if (!initial_state.has_value()) {
+    last_recurrent_state = torch::zeros(
+        {batch_size, num_heads, k_head_dim, v_head_dim},
+        torch::TensorOptions().dtype(value.dtype()).device(value.device()));
+  } else {
+    last_recurrent_state = initial_state.value().to(value);
+  }
+  auto core_attn_out = torch::zeros_like(value);
+  mask = torch::triu(
+      torch::ones(
+          {chunk_size, chunk_size},
+          torch::TensorOptions().dtype(torch::kBool).device(query.device())),
+      1);
+  int64_t num_chunks = total_sequence_length / chunk_size;
+  for (int64_t i = 0; i < num_chunks; ++i) {
+    auto q_i = query.select(2, i);
+    auto k_i = key.select(2, i);
+    auto v_i = value.select(2, i);
+    auto attn_i =
+        (torch::matmul(q_i, k_i.transpose(-1, -2)) * decay_mask.select(2, i))
+            .masked_fill_(mask, 0.0);
+    auto v_prime = torch::matmul(k_cumdecay.select(2, i), last_recurrent_state);
+    auto v_new = v_i - v_prime;
+    auto attn_inter = torch::matmul(q_i * g.select(2, i).unsqueeze(-1).exp(),
+                                    last_recurrent_state);
+    core_attn_out.select(2, i) = attn_inter + torch::matmul(attn_i, v_new);
+    auto g_i_last = g.select(2, i).select(-1, -1).unsqueeze(-1);
+    auto g_exp_term = (g_i_last - g.select(2, i)).exp().unsqueeze(-1);
+    auto k_g_exp = (k_i * g_exp_term).transpose(-1, -2).contiguous();
+    last_recurrent_state = last_recurrent_state * g_i_last.unsqueeze(-1).exp() +
+                           torch::matmul(k_g_exp, v_new);
+  }
+  auto core_attn_out_shape = core_attn_out.sizes();
+  std::vector<int64_t> reshape_shape = {
+      core_attn_out_shape[0],
+      core_attn_out_shape[1],
+      core_attn_out_shape[2] * core_attn_out_shape[3],
+      core_attn_out_shape[4]};
+  core_attn_out = core_attn_out.reshape(reshape_shape);
+  core_attn_out = core_attn_out.slice(2, 0, sequence_length);
+  core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
+  return std::make_tuple(core_attn_out, last_recurrent_state);
+}
+}  // namespace
+
+Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options) {
+  tp_size_ = parallel_args.tp_group_->world_size();
+  rank_ = parallel_args.tp_group_->rank();
+  num_k_heads_ = args.linear_num_key_heads();
+  num_v_heads_ = args.linear_num_value_heads();
+  head_k_dim_ = args.linear_key_head_dim();
+  head_v_dim_ = args.linear_value_head_dim();
+  k_size_ = num_k_heads_ * head_k_dim_;
+  v_size_ = num_v_heads_ * head_v_dim_;
+  conv_kernel_size_ = args.linear_conv_kernel_dim();
+
+  // Shared causal conv projection over mixed QKV states.
+  conv1d_ = register_module("conv1d",
+                            ColumnParallelLinear(args.linear_conv_kernel_dim(),
+                                                 k_size_ * 2 + v_size_,
+                                                 /*bias=*/false,
+                                                 /*gather_output=*/false,
+                                                 quant_args,
+                                                 parallel_args.tp_group_,
+                                                 options));
+
+  auto opts = options.dtype(torch::kFloat32);
+  dt_bias_ = register_parameter("dt_bias",
+                                torch::ones({num_v_heads_ / tp_size_}, opts),
+                                /*requires_grad=*/false);
+
+  A_log_ = register_parameter("A_log",
+                              torch::empty({num_v_heads_ / tp_size_}, opts),
+                              /*requires_grad=*/false);
+
+  // Output projection and gated RMSNorm shared by hybrid variants.
+  o_proj_ = register_module("out_proj",
+                            RowParallelLinear(v_size_,
+                                              args.hidden_size(),
+                                              /*bias=*/false,
+                                              /*input_is_parallelized=*/true,
+                                              /*if_reduce_results=*/true,
+                                              quant_args,
+                                              parallel_args.tp_group_,
+                                              options));
+
+  norm_ = register_module(
+      "norm", RmsNormGated(head_v_dim_, args.rms_norm_eps(), options));
+}
+
+void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict(
+    const StateDict& state_dict) {
+  const int64_t rank = rank_;
+  const int64_t world_size = tp_size_;
+  const int32_t shard_tensor_count = 3;
+  const std::vector<int64_t> shard_sizes = {
+      k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_};
+
+  if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) {
+    conv1d_->load_state_dict(
+        StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes);
+  }
+  o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj."));
+  if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) {
+    norm_->load_state_dict(StateDict({{"weight", w}}));
+  }
+  LOAD_SHARDED_WEIGHT(dt_bias, 0);
+  LOAD_SHARDED_WEIGHT(A_log, 0);
+}
+
+void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights(
+    const std::string& prefix) const {
+  CHECK(dt_bias_is_loaded_)
+      << "Missing required weight after all shards loaded: " << prefix
+      << "dt_bias";
+  CHECK(A_log_is_loaded_) << "Missing required weight after all shards loaded: "
+                          << prefix << "A_log";
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata,
+    KVCache& kv_cache,
+    const ModelInputParams& input_params) {
+  auto [qkvz_padded, ba_padded] =
+      project_padded_inputs(hidden_states, attn_metadata);
+  int64_t batch_size = qkvz_padded.size(0);
+  int64_t seq_len = qkvz_padded.size(1);
+
+  torch::Tensor qkvz_flat =
+      qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
+  torch::Tensor ba_flat =
+      ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
+  xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
+  fused_params.mixed_qkvz = qkvz_flat;
+  fused_params.mixed_ba = ba_flat;
+  fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
+  fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
+  fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
+  fused_params.head_v = static_cast<int32_t>(head_v_dim_);
+
+  torch::Tensor mixed_qkv, z, b, a;
+  std::tie(mixed_qkv, z, b, a) =
+      xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
+
+  mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
+  z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
+  b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
+  a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
+
+  torch::Tensor conv_cache = kv_cache.get_conv_cache();
+  torch::Tensor ssm_cache = kv_cache.get_ssm_cache();
+  torch::Tensor g, beta, core_attn_out, last_recurrent_state;
+  auto device = mixed_qkv.device();
+  auto conv_weight = conv1d_->weight();
+  auto linear_state_indices = get_linear_state_indices(input_params, device);
+
+  if (attn_metadata.is_prefill) {
+    mixed_qkv = mixed_qkv.transpose(1, 2);
+    torch::Tensor conv_state =
+        (seq_len < conv_kernel_size_ - 1)
+            ? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len})
+        : (seq_len > conv_kernel_size_ - 1)
+            ? mixed_qkv.narrow(
+                  -1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1)
+            : mixed_qkv;
+    conv_state = conv_state.transpose(1, 2).contiguous();
+    conv_cache.index_put_({linear_state_indices},
+                          conv_state.to(conv_cache.dtype()));
+    torch::Tensor bias;
+    auto conv_output =
+        torch::conv1d(mixed_qkv,
+                      conv_weight.unsqueeze(1).to(device),
+                      bias,
+                      /*stride=*/std::vector<int64_t>{1},
+                      /*padding=*/std::vector<int64_t>{3},
+                      /*dilation=*/std::vector<int64_t>{1},
+                      /*groups=*/static_cast<int64_t>(mixed_qkv.size(1)));
+    mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len));
+
+  } else {
+    xllm::kernel::CausalConv1dUpdateParams conv1d_params;
+    conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)});
+    conv1d_params.conv_state = conv_cache;
+    conv1d_params.weight = conv_weight;
+    conv1d_params.conv_state_indices = linear_state_indices;
+    conv1d_params.block_idx_last_scheduled_token =
+        std::optional<torch::Tensor>();
+    conv1d_params.initial_state_idx = std::optional<torch::Tensor>();
+    conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
+    conv1d_params.max_query_len = attn_metadata.max_query_len;
+    mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
+    // Reshape back to 3D [batch_size, dim, seq_len]
+    mixed_qkv =
+        mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous();
+    mixed_qkv = mixed_qkv.transpose(1, 2);
+  }
+
+  // Compute gated delta net decay and beta terms.
+  if (attn_metadata.is_prefill) {
+    xllm::kernel::FusedGdnGatingParams gdn_params;
+    gdn_params.A_log = A_log_;
+    gdn_params.a = a.contiguous().view({-1, a.size(-1)});
+    gdn_params.b = b.contiguous().view({-1, b.size(-1)});
+    gdn_params.dt_bias = dt_bias_;
+    gdn_params.beta = 1.0f;
+    gdn_params.threshold = 20.0f;
+    std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
+    g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)});
+    beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)});
+  } else {
+    xllm::kernel::FusedGdnGatingParams gdn_params;
+    gdn_params.A_log = A_log_;
+    gdn_params.a = a.view({-1, a.size(-1)});
+    gdn_params.b = b.view({-1, b.size(-1)});
+    gdn_params.dt_bias = dt_bias_;
+    gdn_params.beta = 1.0f;
+    gdn_params.threshold = 20.0f;
+    std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
+  }
+  auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv);
+  // Apply chunked or recurrent gated-delta attention and update caches.
+  if (attn_metadata.is_prefill) {
+    xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params;
+    chunk_gated_delta_params.q = processed_q;
+    chunk_gated_delta_params.k = processed_k;
+    chunk_gated_delta_params.v = processed_v;
+    chunk_gated_delta_params.g = g;
+    chunk_gated_delta_params.beta = beta;
+    // Get initial state from ssm_cache for sequences with previous state
+    // Shape: [batch_size, num_heads, head_k_dim, head_v_dim]
+    torch::Tensor initial_state_tensor =
+        torch::index_select(ssm_cache, 0, linear_state_indices);
+    // Todo: chunked-prefill/prefix-cache use initial_state
+    initial_state_tensor.fill_(0.0);
+    chunk_gated_delta_params.initial_state = initial_state_tensor;
+    chunk_gated_delta_params.output_final_state = true;
+    chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
+    chunk_gated_delta_params.head_first = false;
+    chunk_gated_delta_params.use_qk_l2norm_in_kernel = true;
+    std::tie(core_attn_out, last_recurrent_state) =
+        xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params);
+    ssm_cache.index_put_(
+        {linear_state_indices},
+        last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
+  } else {
+    processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
+    processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
+    auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
+    torch::Tensor actual_seq_lengths =
+        torch::cat({zero, attn_metadata.q_seq_lens}, 0);
+    double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
+    core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
+                        processed_q.reshape(
+                            {-1, processed_q.size(-2), processed_q.size(-1)}),
+                        processed_k.reshape(
+                            {-1, processed_k.size(-2), processed_k.size(-1)}),
+                        processed_v.reshape(
+                            {-1, processed_v.size(-2), processed_v.size(-1)}),
+                        ssm_cache,
+                        beta.squeeze(0).contiguous(),
+                        scale,
+                        actual_seq_lengths,
+                        linear_state_indices,
+                        c10::nullopt,
+                        g.squeeze(0).contiguous(),
+                        c10::nullopt)
+                        .unsqueeze(0)
+                        .contiguous();
+  }
+
+  auto z_reshaped = z.view({-1, z.size(-1)});
+  auto core_attn_out_reshaped =
+      core_attn_out.view({-1, core_attn_out.size(-1)});
+  auto norm_out = norm_->forward(core_attn_out_reshaped, z_reshaped);
+  auto z_shape_og = z.sizes().vec();
+  norm_out = norm_out.view(z_shape_og);
+  norm_out = norm_out.view({-1, norm_out.size(2), norm_out.size(3)});
+
+  // Project the normalized attention output back to hidden size.
+  auto rearranged_norm =
+      norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)});
+  rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm);
+  auto attn_output = o_proj_->forward(rearranged_norm);
+  return attn_output;
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
+    const AttentionMetadata& attn_metadata,
+    const torch::Tensor& padded_qkvz) const {
+  if (!attn_metadata.is_prefill) {
+    return padded_qkvz;
+  }
+  std::vector<torch::Tensor> valid_batches;
+  int64_t bs = attn_metadata.q_seq_lens.size(0);
+  int64_t max_len = attn_metadata.max_query_len;
+  const auto& ori_seq_lens = attn_metadata.q_seq_lens;
+  auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1});
+  for (int64_t b = 0; b < bs; ++b) {
+    int64_t ori_len = ori_seq_lens[b].template item<int64_t>();
+    torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len);
+    valid_batches.push_back(valid_batch);
+  }
+  return torch::cat(valid_batches, 0).contiguous();
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices(
+    const ModelInputParams& input_params,
+    const torch::Device& device) const {
+  CHECK(!input_params.linear_state_ids.empty())
+      << "linear_state_ids must be populated for gated delta net";
+  if (input_params.linear_state_indices.defined()) {
+    return input_params.linear_state_indices;
+  }
+  return torch::tensor(
+      input_params.linear_state_ids,
+      torch::TensorOptions().dtype(torch::kInt).device(device));
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad(
+    const AttentionMetadata& attn_metadata,
+    const torch::Tensor& qkvz) const {
+  int64_t bs = attn_metadata.q_seq_lens.size(0);
+  int64_t max_len = attn_metadata.max_query_len;
+  const auto& start_loc = attn_metadata.q_seq_lens;
+  if (!attn_metadata.is_prefill) {
+    return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)});
+  }
+  std::vector<torch::Tensor> batches;
+  int64_t idx = 0;
+  for (int64_t b = 0; b < bs; ++b) {
+    int64_t cur_len = start_loc[b].template item<int64_t>();
+    torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous();
+    idx = idx + cur_len;
+    if (batch.size(0) != max_len) {
+      batch = batch.size(0) > max_len
+                  ? batch.slice(0, 0, max_len).contiguous()
+                  : torch::nn::functional::pad(
+                        batch,
+                        torch::nn::functional::PadFuncOptions(
+                            {0, 0, 0, max_len - batch.size(0)}))
+                        .contiguous();
+    }
+    batches.push_back(batch);
+  }
+  auto ret = torch::stack(batches, 0).contiguous();
+  return ret;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+Qwen3GatedDeltaNetBaseImpl::process_mixed_qkv(torch::Tensor& mixed_qkv) const {
+  mixed_qkv = mixed_qkv.transpose(1, 2);
+  int64_t batch_size = mixed_qkv.size(0);
+  int64_t seq_len = mixed_qkv.size(1);
+  std::vector<int64_t> split_sizes = {
+      k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_};
+  auto processed_qkv = torch::split(mixed_qkv, split_sizes, 2);
+  auto processed_q = processed_qkv[0];
+  auto processed_k = processed_qkv[1];
+  auto processed_v = processed_qkv[2];
+  processed_q = processed_q.view(
+      {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_});
+  processed_k = processed_k.view(
+      {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_});
+  processed_v = processed_v.view(
+      {batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
+  return std::make_tuple(processed_q, processed_k, processed_v);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/layers/npu_torch/qwen3_gated_delta_net_base.h

@@ -0,0 +1,90 @@
+/* Copyright 2026 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <torch/torch.h>
+
+#include <string>
+#include <tuple>
+#include <utility>
+
+#include "attention.h"
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "layers/common/linear.h"
+#include "layers/common/rms_norm_gated.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
+ public:
+  Qwen3GatedDeltaNetBaseImpl() = default;
+  Qwen3GatedDeltaNetBaseImpl(const ModelArgs& args,
+                             const QuantArgs& quant_args,
+                             const ParallelArgs& parallel_args,
+                             const torch::TensorOptions& options);
+
+  virtual void load_state_dict(const StateDict& state_dict) = 0;
+  virtual void verify_loaded_weights(const std::string& prefix) const = 0;
+
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const AttentionMetadata& attn_metadata,
+                        KVCache& kv_cache,
+                        const ModelInputParams& input_params);
+
+ protected:
+  virtual std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
+      const torch::Tensor& hidden_states,
+      const AttentionMetadata& attn_metadata) = 0;
+
+  void load_common_state_dict(const StateDict& state_dict);
+  void verify_common_loaded_weights(const std::string& prefix) const;
+
+  torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata,
+                                      const torch::Tensor& qkvz) const;
+  torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
+                                   const torch::Tensor& padded_qkvz) const;
+  torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
+                                         const torch::Device& device) const;
+
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
+      torch::Tensor& mixed_qkv) const;
+
+  int64_t num_k_heads_ = 0;
+  int64_t num_v_heads_ = 0;
+  int64_t head_k_dim_ = 0;
+  int64_t head_v_dim_ = 0;
+  int64_t k_size_ = 0;
+  int64_t v_size_ = 0;
+  int64_t tp_size_ = 1;
+  int64_t rank_ = 0;
+  int32_t conv_kernel_size_ = 0;
+
+  ColumnParallelLinear conv1d_{nullptr};
+  RowParallelLinear o_proj_{nullptr};
+  RmsNormGated norm_{nullptr};
+
+  DEFINE_WEIGHT(dt_bias);
+  DEFINE_WEIGHT(A_log);
+};
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/models/qwen3_5.h

@@ -0,0 +1,218 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 <cstdint>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "core/layers/npu_torch/qwen3_5_decoder_layer_impl.h"
+#include "models/model_registry.h"
+#include "qwen3_next.h"
+
+namespace xllm {
+
+class Qwen3_5ModelImpl : public Qwen3NextModelImpl {
+ public:
+  explicit Qwen3_5ModelImpl(const ModelContext& context)
+      : Qwen3NextModelImpl(context, /*init_decoder_layers=*/false) {
+    const int32_t n_layers = context.get_model_args().n_layers();
+    for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) {
+      add_decoder_layer(
+          std::make_shared<layer::Qwen3_5DecoderLayerImpl>(context, layer_id));
+    }
+  }
+};
+TORCH_MODULE(Qwen3_5Model);
+
+class Qwen3_5ForCausalLMImpl : public Qwen3NextForCausalLMImpl {
+ public:
+  explicit Qwen3_5ForCausalLMImpl(const ModelContext& context)
+      : Qwen3NextForCausalLMImpl(context, /*init_model=*/false) {
+    set_model_module(std::make_shared<Qwen3_5ModelImpl>(context));
+  }
+};
+TORCH_MODULE(Qwen3_5ForCausalLM);
+
+#define LOAD_ARG_TEXT_OR_ROOT(arg_name, json_key, default_value) \
+  LOAD_ARG_OR(arg_name, "text_config." json_key, default_value); \
+  LOAD_ARG_OR(arg_name, json_key, args->arg_name())
+
+#define LOAD_ARG_TEXT_OR_ROOT_CHAIN(arg_name, json_key, default_value) \
+  LOAD_ARG_TEXT_OR_ROOT(arg_name, json_key, default_value)
+
+#define LOAD_QWEN3_5_ROPE_ARG(arg_name, default_value)                       \
+  LOAD_ARG_OR(arg_name, "text_config." #arg_name, default_value);            \
+  LOAD_ARG_OR(arg_name, #arg_name, args->arg_name());                        \
+  LOAD_ARG_OR(                                                               \
+      arg_name, "text_config.rope_scaling." #arg_name, args->arg_name());    \
+  LOAD_ARG_OR(arg_name, "rope_scaling." #arg_name, args->arg_name());        \
+  LOAD_ARG_OR(                                                               \
+      arg_name, "text_config.rope_parameters." #arg_name, args->arg_name()); \
+  LOAD_ARG_OR(arg_name, "rope_parameters." #arg_name, args->arg_name())
+
+#define LOAD_QWEN3_5_NEXT_COMPAT_ARGS(default_moe_intermediate_size,           \
+                                      default_num_experts,                     \
+                                      default_num_experts_per_tok,             \
+                                      default_shared_expert_intermediate_size) \
+  LOAD_ARG_TEXT_OR_ROOT(attention_bias, "attention_bias", false);              \
+  LOAD_ARG_TEXT_OR_ROOT(attention_dropout, "attention_dropout", 0.0f);         \
+  LOAD_ARG_TEXT_OR_ROOT(bos_token_id, "bos_token_id", 151643);                 \
+  LOAD_ARG_TEXT_OR_ROOT(decoder_sparse_step, "decoder_sparse_step", 1);        \
+  LOAD_ARG_TEXT_OR_ROOT(eos_token_id, "eos_token_id", 151645);                 \
+  LOAD_ARG_TEXT_OR_ROOT(head_dim, "head_dim", 256);                            \
+  LOAD_ARG_TEXT_OR_ROOT(hidden_act, "hidden_act", "silu");                     \
+  LOAD_ARG_TEXT_OR_ROOT(hidden_size, "hidden_size", 2048);                     \
+  LOAD_ARG_TEXT_OR_ROOT(initializer_range, "initializer_range", 0.02f);        \
+  LOAD_ARG_TEXT_OR_ROOT(intermediate_size, "intermediate_size", 5120);         \
+  LOAD_ARG_TEXT_OR_ROOT(                                                       \
+      max_position_embeddings, "max_position_embeddings", 262144);             \
+  LOAD_ARG_TEXT_OR_ROOT(max_window_layers, "max_window_layers", 28);           \
+  LOAD_ARG_TEXT_OR_ROOT(moe_intermediate_size,                                 \
+                        "moe_intermediate_size",                               \
+                        default_moe_intermediate_size);                        \
+  LOAD_ARG_TEXT_OR_ROOT(norm_topk_prob, "norm_topk_prob", true);               \
+  LOAD_ARG_TEXT_OR_ROOT(n_heads, "num_attention_heads", 16);                   \
+  LOAD_ARG_TEXT_OR_ROOT(num_experts, "num_experts", default_num_experts);      \
+  LOAD_ARG_TEXT_OR_ROOT(num_experts_per_tok,                                   \
+                        "num_experts_per_tok",                                 \
+                        default_num_experts_per_tok);                          \
+  LOAD_ARG_TEXT_OR_ROOT(n_layers, "num_hidden_layers", 48);                    \
+  LOAD_ARG_OR(n_kv_heads, "text_config.num_key_value_heads", 2);               \
+  LOAD_ARG_OR(                                                                 \
+      n_kv_heads, "num_key_value_heads", args->n_kv_heads().value_or(2));      \
+  LOAD_ARG_TEXT_OR_ROOT(output_router_logits, "output_router_logits", false);  \
+  LOAD_ARG_TEXT_OR_ROOT(rms_norm_eps, "rms_norm_eps", 1e-6);                   \
+  LOAD_QWEN3_5_ROPE_ARG(rope_theta, 10000000.0f);                              \
+  LOAD_ARG_TEXT_OR_ROOT(router_aux_loss_coef, "router_aux_loss_coef", 0.001f); \
+  LOAD_ARG_TEXT_OR_ROOT(use_sliding_window, "use_sliding_window", false);      \
+  LOAD_ARG_TEXT_OR_ROOT(sliding_window, "sliding_window", 4096);               \
+  LOAD_ARG_TEXT_OR_ROOT(tie_word_embeddings, "tie_word_embeddings", false);    \
+  LOAD_ARG_TEXT_OR_ROOT(vocab_size, "vocab_size", 151936);                     \
+  LOAD_ARG_TEXT_OR_ROOT(                                                       \
+      mlp_only_layers, "mlp_only_layers", std::vector<int32_t>());             \
+  LOAD_ARG_TEXT_OR_ROOT(attn_output_gate, "attn_output_gate", true);           \
+  LOAD_ARG_TEXT_OR_ROOT(                                                       \
+      full_attention_interval, "full_attention_interval", 4);                  \
+  LOAD_ARG_TEXT_OR_ROOT(linear_conv_kernel_dim, "linear_conv_kernel_dim", 4);  \
+  LOAD_ARG_TEXT_OR_ROOT(linear_key_head_dim, "linear_key_head_dim", 128);      \
+  LOAD_ARG_TEXT_OR_ROOT(linear_num_key_heads, "linear_num_key_heads", 16);     \
+  LOAD_ARG_TEXT_OR_ROOT(linear_num_value_heads, "linear_num_value_heads", 32); \
+  LOAD_ARG_TEXT_OR_ROOT(linear_value_head_dim, "linear_value_head_dim", 128);  \
+  LOAD_QWEN3_5_ROPE_ARG(partial_rotary_factor, 0.25f);                         \
+  LOAD_ARG_OR(rope_scaling_mrope_section,                                      \
+              "text_config.rope_scaling.mrope_section",                        \
+              std::vector<int64_t>());                                         \
+  LOAD_ARG_OR(rope_scaling_mrope_section,                                      \
+              "text_config.rope_parameters.mrope_section",                     \
+              args->rope_scaling_mrope_section());                             \
+  LOAD_ARG_OR(rope_scaling_mrope_section,                                      \
+              "rope_parameters.mrope_section",                                 \
+              args->rope_scaling_mrope_section());                             \
+  LOAD_ARG_OR(rope_scaling_mrope_interleaved,                                  \
+              "text_config.rope_scaling.mrope_interleaved",                    \
+              false);                                                          \
+  LOAD_ARG_OR(rope_scaling_mrope_interleaved,                                  \
+              "text_config.rope_parameters.mrope_interleaved",                 \
+              args->rope_scaling_mrope_interleaved());                         \
+  LOAD_ARG_OR(rope_scaling_mrope_interleaved,                                  \
+              "rope_parameters.mrope_interleaved",                             \
+              args->rope_scaling_mrope_interleaved());                         \
+  LOAD_ARG_TEXT_OR_ROOT(shared_expert_intermediate_size,                       \
+                        "shared_expert_intermediate_size",                     \
+                        default_shared_expert_intermediate_size);              \
+  LOAD_ARG_OR(                                                                 \
+      num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0);       \
+  LOAD_ARG_OR(num_nextn_predict_layers,                                        \
+              "mtp_num_hidden_layers",                                         \
+              args->num_nextn_predict_layers());                               \
+  LOAD_ARG_OR(num_nextn_predict_layers,                                        \
+              "text_config.num_nextn_predict_layers",                          \
+              args->num_nextn_predict_layers());                               \
+  LOAD_ARG_OR(num_nextn_predict_layers,                                        \
+              "num_nextn_predict_layers",                                      \
+              args->num_nextn_predict_layers());                               \
+  LOAD_ARG_OR(                                                                 \
+      layer_types, "text_config.layer_types", std::vector<std::string>());     \
+  LOAD_ARG_OR(layer_types, "layer_types", args->layer_types());                \
+  LOAD_ARG_OR(                                                                 \
+      layer_types, "text_config.layers_block_type", args->layer_types());      \
+  LOAD_ARG_OR(layer_types, "layers_block_type", args->layer_types());          \
+  LOAD_ARG_OR(                                                                 \
+      n_routed_experts, "text_config.n_routed_experts", args->num_experts());  \
+  LOAD_ARG_OR(n_routed_experts, "n_routed_experts", args->num_experts());      \
+  SET_ARG(n_shared_experts,                                                    \
+          args->shared_expert_intermediate_size() > 0 ? 1 : 0);                \
+  SET_ARG(scoring_func, "softmax");                                            \
+  SET_ARG(topk_method, "");                                                    \
+  SET_ARG(n_group, -1);                                                        \
+  SET_ARG(topk_group, 0);                                                      \
+  SET_ARG(routed_scaling_factor, 1.0f);                                        \
+  SET_ARG(stop_token_ids,                                                      \
+          std::unordered_set<int32_t>({args->eos_token_id()}));                \
+  LOAD_ARG_TEXT_OR_ROOT(mamba_ssm_dtype, "mamba_ssm_dtype", "float32")
+
+#define LOAD_QWEN3_5_TYPE_AND_DTYPE(default_model_type)         \
+  LOAD_ARG_OR(model_type, "model_type", default_model_type);    \
+  LOAD_ARG_OR(dtype, "text_config.dtype", "bfloat16");          \
+  LOAD_ARG_OR(dtype, "dtype", args->dtype());                   \
+  LOAD_ARG_OR(dtype, "text_config.torch_dtype", args->dtype()); \
+  LOAD_ARG_OR(dtype, "torch_dtype", args->dtype())
+
+REGISTER_CAUSAL_MODEL(qwen3_5, Qwen3_5ForCausalLM);
+REGISTER_MODEL_ARGS(qwen3_5, [&] {
+  LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5");
+  LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0,
+                                /*num_experts=*/0,
+                                /*num_experts_per_tok=*/0,
+                                /*shared_expert_intermediate_size=*/0);
+});
+
+REGISTER_CAUSAL_MODEL(qwen3_5_text, Qwen3_5ForCausalLM);
+REGISTER_MODEL_ARGS(qwen3_5_text, [&] {
+  LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_text");
+  LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0,
+                                /*num_experts=*/0,
+                                /*num_experts_per_tok=*/0,
+                                /*shared_expert_intermediate_size=*/0);
+});
+
+REGISTER_CAUSAL_MODEL(qwen3_5_moe, Qwen3_5ForCausalLM);
+REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
+  LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe");
+  LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512,
+                                /*num_experts=*/512,
+                                /*num_experts_per_tok=*/10,
+                                /*shared_expert_intermediate_size=*/512);
+});
+
+REGISTER_CAUSAL_MODEL(qwen3_5_moe_text, Qwen3_5ForCausalLM);
+REGISTER_MODEL_ARGS(qwen3_5_moe_text, [&] {
+  LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe_text");
+  LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512,
+                                /*num_experts=*/512,
+                                /*num_experts_per_tok=*/10,
+                                /*shared_expert_intermediate_size=*/512);
+});
+
+#undef LOAD_QWEN3_5_TYPE_AND_DTYPE
+#undef LOAD_QWEN3_5_NEXT_COMPAT_ARGS
+#undef LOAD_QWEN3_5_ROPE_ARG
+#undef LOAD_ARG_TEXT_OR_ROOT_CHAIN
+#undef LOAD_ARG_TEXT_OR_ROOT
+
+}  // namespace xllm