diff --git a/upstream_ref/README.md b/upstream_ref/README.md new file mode 100644 index 00000000..0c95d7a8 --- /dev/null +++ b/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() + │ + ├── ds_vllm path (Python torch extension): + │ csrc/moe/topk_softmax_kernels.cu + │ → torch.ops._moe_C.topk_softmax() + │ → topk_gating_softmax_kernel_launcher() + │ + └── 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 diff --git a/upstream_ref/ds_vllm/csrc/moe/moeTopKFuncs.cuh b/upstream_ref/ds_vllm/csrc/moe/moeTopKFuncs.cuh new file mode 100644 index 00000000..70e21cf8 --- /dev/null +++ b/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 +#include +#include + +namespace vllm { +namespace moe { +namespace reduce_topk { +namespace cg = cooperative_groups; +static constexpr int kWARP_SIZE = 32; + +template +struct TopKRedType { + using T = T_; + static_assert( + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v, + "Top K reduction only implemented for int, float, float16 and bfloat16"); + + using TypeCmp = std::conditional_t; + using IdxT = std::conditional_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::TwiddleIn( + reinterpret_cast::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((cmp & 0xFFFF)); + + auto compactTmp = cmp >> kMoveBits; + auto valueBits = cub::Traits::TwiddleOut( + reinterpret_cast::UnsignedBits&>(compactTmp)); + value = reinterpret_cast(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 const& warp) { + return cg::reduce(warp, compValIdx, cg::greater{}); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct TopKIdx { + // by default, empty +}; + +template +struct TopKIdx { + 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 +struct Sort; + +template +struct Sort<1, RedType> { + static __device__ void run(RedType* topK) {} +}; + +template +struct Sort<2, RedType> { + static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } +}; + +template +struct Sort<3, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 1); + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 1); + } +}; + +template +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 +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile 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; + 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 +__device__ void reduceTopKFunc(cg::thread_block_tile 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; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } + + if constexpr (!IsSorted) { + Sort::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 +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile 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; + + if constexpr (N <= 4) { + reduceTopKFunc(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(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(warp, out, outIdx, topKBufferValue, + topKBufferIdx, minValue, actualK); + } +}; + +#undef TOPK_SWAP + +} // namespace reduce_topk +} // namespace moe +} // namespace vllm diff --git a/upstream_ref/ds_vllm/csrc/moe/moe_align_sum_kernels.cu b/upstream_ref/ds_vllm/csrc/moe/moe_align_sum_kernels.cu new file mode 100644 index 00000000..d7c68ff2 --- /dev/null +++ b/upstream_ref/ds_vllm/csrc/moe/moe_align_sum_kernels.cu @@ -0,0 +1,833 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +#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; + __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 +__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; + __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 +__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(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 +__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 +__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 +__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 +__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 +__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( + 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 +__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 +__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 +__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( + 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 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(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(experts_ids.mutable_data_ptr()), + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(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; + + 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(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(experts_ids.mutable_data_ptr()), + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, padded_num_experts, experts_per_warp, block_size, + topk_ids.numel(), + reinterpret_cast(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; + sort_kernel<<>>( + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(cumsum_buffer.mutable_data_ptr()), + reinterpret_cast(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(batch_num_tokens.const_data_ptr()), + reinterpret_cast(sorted_ids.mutable_data_ptr()), + reinterpret_cast(batch_ids.mutable_data_ptr()), + reinterpret_cast(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<<>>( + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(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<<>>( + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(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<<>>( + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(input.const_data_ptr()), + hidden_size); + }); + break; + + default: + torch::stable::sum_out(output, input, std::array{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 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<<>>( + reinterpret_cast(topk_ids.mutable_data_ptr()), + reinterpret_cast(token_lora_mapping.mutable_data_ptr()), + block_size, + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, + max_num_m_blocks, + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(expert_ids.mutable_data_ptr()), + topk_num, + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(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; + + // 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(topk_ids.mutable_data_ptr()), + reinterpret_cast(token_lora_mapping.mutable_data_ptr()), + block_size, + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, + max_num_m_blocks, + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(expert_ids.mutable_data_ptr()), + topk_num, + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + reinterpret_cast(cumsum.mutable_data_ptr()), WARP_SIZE, + padded_num_experts, + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(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; + + sort_kernel<<>>( + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(cumsum.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num, + reinterpret_cast(token_mask.mutable_data_ptr()), + max_loras, + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + has_expert_map); + } + }); +} \ No newline at end of file diff --git a/upstream_ref/ds_vllm/csrc/moe/moe_ops.h b/upstream_ref/ds_vllm/csrc/moe/moe_ops.h new file mode 100644 index 00000000..43cbb7f8 --- /dev/null +++ b/upstream_ref/ds_vllm/csrc/moe/moe_ops.h @@ -0,0 +1,87 @@ +#pragma once + +#include + +#include +#include + +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 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 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& correction_bias, + const std::optional& input_ids, + const std::optional& 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 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 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 b_qzeros, + std::optional 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 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 diff --git a/upstream_ref/ds_vllm/csrc/moe/topk_softmax_kernels.cu b/upstream_ref/ds_vllm/csrc/moe/topk_softmax_kernels.cu new file mode 100644 index 00000000..e8453579 --- /dev/null +++ b/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 + +#include +#include +#include +#include +#include + +#include "../../cuda_compat.h" +#include "../../cub_helpers.h" +#include "libtorch_stable/torch_utils.h" + +#ifndef USE_ROCM + #include + #include +#else + #include + #include + 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 +__device__ __forceinline__ float toFloat(T value) { + if constexpr (std::is_same_v) { + return value; + } else if constexpr (std::is_same_v) { + return __bfloat162float(value); + } else if constexpr (std::is_same_v) { + 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 +__launch_bounds__(TPB) __global__ + void moeSoftmax(const InputType* input, const bool* finished, float* output, const int num_cols) +{ + using BlockReduce = cub::BlockReduce; + __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 +__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 +__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; + using BlockReduce = cub::BlockReduce; + __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 +__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 || std::is_same_v || + std::is_same_v, + "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 || std::is_same_v) { + 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) { + using VecType = AlignedArray; + VecType* row_chunk_vec_ptr = reinterpret_cast(&row_chunk); + const VecType* vec_thread_read_ptr = reinterpret_cast(thread_read_ptr); +#pragma unroll + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { + row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; + } + } else if constexpr (std::is_same_v) { + if constexpr (ELTS_PER_LDG >= 2) { + using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>; + float2* row_chunk_f2 = reinterpret_cast(row_chunk); + const VecType* vec_thread_read_ptr = reinterpret_cast(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(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) { + if constexpr (ELTS_PER_LDG >= 2) { + using VecType = AlignedArray<__half, ELTS_PER_LDG>; + float2* row_chunk_f2 = reinterpret_cast(row_chunk); + const VecType* vec_thread_read_ptr = reinterpret_cast(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(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 +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 +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; + 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<<>>( + 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( \ + 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( \ + 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( \ + 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 +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 || std::is_same_v) ? 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<<>>( + gating_output, nullptr, workspace, num_experts); + } else if constexpr (SF == SCORING_SIGMOID) { + moeSigmoid<<>>( + gating_output, nullptr, workspace, num_experts); + } else { + STD_TORCH_CHECK(false, "Unsupported scoring func"); + } + moeTopK<<>>( + workspace, nullptr, topk_weights, topk_indices, token_expert_indices, + num_experts, topk, 0, num_experts, renormalize, bias); + } + } +} + +} // namespace moe +} // namespace vllm + + +template +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 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(); + } + + if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { + vllm::moe::topkGatingKernelLauncher( + reinterpret_cast(gating_output.const_data_ptr()), + topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_data_ptr(), + softmax_workspace.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, + bias_ptr, stream); + } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { + vllm::moe::topkGatingKernelLauncher( + reinterpret_cast(gating_output.const_data_ptr()), + topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_data_ptr(), + softmax_workspace.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, + bias_ptr, stream); + } else { + STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); + vllm::moe::topkGatingKernelLauncher( + reinterpret_cast(gating_output.const_data_ptr()), + topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_data_ptr(), + softmax_workspace.mutable_data_ptr(), + 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 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(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 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(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()); + } +} diff --git a/upstream_ref/ds_vllm/vllm/_custom_ops.py b/upstream_ref/ds_vllm/vllm/_custom_ops.py new file mode 100644 index 00000000..3bac6972 --- /dev/null +++ b/upstream_ref/ds_vllm/vllm/_custom_ops.py @@ -0,0 +1,3998 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from enum import IntEnum +from typing import TYPE_CHECKING, Literal + +import torch + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.scalar_type import ScalarType +from vllm.utils.flashinfer import ( + flashinfer_quant_nvfp4_8x4_sf_layout, +) +from vllm.utils.math_utils import cdiv + +logger = init_logger(__name__) + +current_platform.import_kernels() + +if TYPE_CHECKING: + + def register_fake(fn): + return lambda name: fn +else: + try: + from torch.library import register_fake + except ImportError: + from torch.library import impl_abstract as register_fake + + +# scaled_fp4_quant functional + out variant for torch.compile buffer management + + +def create_fp4_scale_tensor( + m: int, + n: int, + device: torch.device, + is_sf_swizzled_layout: bool, +) -> torch.Tensor: + """ + Allocate the output scale tensor for scaled_fp4_quant. + + When is_sf_swizzled_layout=True, we use rounded values to store the + swizzled scales. Due to the requirement of the Tensor Core, the minimum + tile is 128x4 for the scales. So, we first pad the scales to multiples + of 128 (rows) and 4 (cols). Then, the scales (in float8_e4m3fn) are + packed into an int32 for every 4 values. More: + https://docs.nvidia.com/cuda/parallel-thread-execution/ + #tcgen05-mma-scale-factor-b-layout-4x + """ + from vllm.utils.math_utils import round_up + + block_size = 16 + if is_sf_swizzled_layout: + rounded_m = round_up(m, 128) + scale_n = n // block_size + rounded_n = round_up(scale_n, 4) + return torch.empty( + (rounded_m, rounded_n // 4), device=device, dtype=torch.int32 + ) + else: + return torch.empty((m, n // block_size), device=device, dtype=torch.uint8) + + +def create_fp4_output_tensors( + m: int, + n: int, + device: torch.device, + is_sf_swizzled_layout: bool, + padded_n: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Allocate both output tensors for scaled_fp4_quant: + (quantized_output, output_scale). + + Must match the C++ scaled_fp4_quant_func allocation exactly when + ``padded_n`` is ``None``. When ``padded_n`` is provided, allocate a larger + packed-FP4 output/scale buffer so the quantization kernel can write + CUTLASS-compatible K padding directly + """ + physical_n = padded_n if padded_n is not None else n + output = torch.empty((m, physical_n // 2), device=device, dtype=torch.uint8) + output_scale = create_fp4_scale_tensor(m, physical_n, device, is_sf_swizzled_layout) + return output, output_scale + + +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"): + + @register_fake("_C::scaled_fp4_quant") + def _scaled_fp4_quant_fake( + input: torch.Tensor, + input_scale: torch.Tensor, + is_sf_swizzled_layout: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + n = input.shape[-1] + m = input.numel() // n + return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout) + + @register_fake("_C::scaled_fp4_quant.out") + def _scaled_fp4_quant_out_fake( + input: torch.Tensor, + input_scale: torch.Tensor, + is_sf_swizzled_layout: bool, + *, + output: torch.Tensor, + output_scale: torch.Tensor, + ) -> None: + return None + + +# page attention ops +def paged_attention_v1( + out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: torch.Tensor | None, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + torch.ops._C.paged_attention_v1( + out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def paged_attention_v2( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: torch.Tensor | None, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + torch.ops._C.paged_attention_v2( + out, + exp_sum, + max_logits, + tmp_out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def paged_attention_rocm( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor | None, + block_size: int, + max_seq_len: int, + alibi_slopes: torch.Tensor | None, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + fp8_out_scale: torch.Tensor | None = None, + mfma_type: str = "fp8" if envs.VLLM_ROCM_FP8_MFMA_PAGE_ATTN else "f16", +) -> None: + torch.ops._rocm_C.paged_attention( + out, + exp_sum, + max_logits, + tmp_out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + query_start_loc, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + fp8_out_scale, + mfma_type, + ) + + +def mla_decode_kvcache_cpu( + out: torch.Tensor, + query: torch.Tensor, + kv_cache: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, +) -> None: + torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens) + + +# merge attn states ops +def merge_attn_states( + output: torch.Tensor, + prefix_output: torch.Tensor, + prefix_lse: torch.Tensor, + suffix_output: torch.Tensor, + suffix_lse: torch.Tensor, + output_lse: torch.Tensor | None = None, + prefill_tokens_with_context: int | None = None, + output_scale: torch.Tensor | None = None, +) -> None: + torch.ops._C.merge_attn_states( + output, + output_lse, + prefix_output, + prefix_lse, + suffix_output, + suffix_lse, + prefill_tokens_with_context, + output_scale, + ) + + +# pos encoding ops +def rotary_embedding( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox: bool, + rope_dim_offset: int = 0, + inverse: bool = False, +) -> None: + if rope_dim_offset == 0 and not inverse: + torch.ops._C.rotary_embedding( + positions, query, key, head_size, cos_sin_cache, is_neox + ) + else: + torch.ops._C.rotary_embedding( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox, + rope_dim_offset, + inverse, + ) + + +# layer norm ops +def rms_norm( + out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, epsilon: float +) -> None: + torch.ops._C.rms_norm(out, input, weight, epsilon) + + +def fused_add_rms_norm( + input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, epsilon: float +) -> None: + # Note: this func is batch invariant + torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon) + + +def fused_qk_norm_rope( + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + position_ids: torch.Tensor, + forced_token_heads_per_warp: int = -1, +) -> None: + torch.ops._C.fused_qk_norm_rope( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + eps, + q_weight, + k_weight, + cos_sin_cache, + is_neox, + position_ids, + forced_token_heads_per_warp, + ) + + +def apply_repetition_penalties_torch( + logits: torch.Tensor, + prompt_mask: torch.Tensor, + output_mask: torch.Tensor, + repetition_penalties: torch.Tensor, +) -> None: + repetition_penalties = repetition_penalties.unsqueeze(dim=1).repeat( + 1, logits.size(1) + ) + # If token appears in prompt or output, apply, otherwise use 1.0 for no-op. + penalties = torch.where(prompt_mask | output_mask, repetition_penalties, 1.0) + # If logits are positive, divide by penalty, otherwise multiply by penalty. + scaling = torch.where(logits > 0, 1.0 / penalties, penalties) + logits *= scaling + + +def apply_repetition_penalties_cuda( + logits: torch.Tensor, + prompt_mask: torch.Tensor, + output_mask: torch.Tensor, + repetition_penalties: torch.Tensor, +) -> None: + torch.ops._C.apply_repetition_penalties_( + logits, prompt_mask, output_mask, repetition_penalties + ) + + +def apply_repetition_penalties( + logits: torch.Tensor, + prompt_mask: torch.Tensor, + output_mask: torch.Tensor, + repetition_penalties: torch.Tensor, +) -> None: + """Apply repetition penalties to logits in-place. + + Args: + logits: The logits tensor of shape [num_seqs, vocab_size]. + prompt_mask: A boolean tensor indicating which tokens appear in the prompt. + output_mask: A boolean tensor indicating which tokens appear in the output. + repetition_penalties: The repetition penalties of shape (num_seqs, ). + """ + if logits.is_cuda and logits.is_contiguous(): + apply_repetition_penalties_cuda( + logits, prompt_mask, output_mask, repetition_penalties + ) + else: + apply_repetition_penalties_torch( + logits, prompt_mask, output_mask, repetition_penalties + ) + + +# fused quant layer norm ops +def rms_norm_dynamic_per_token_quant( + input: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + quant_dtype: torch.dtype, + scale_ub: torch.Tensor | None = None, + residual: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty(input.shape, dtype=quant_dtype, device=input.device) + scales = torch.empty( + (input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32 + ) + + torch.ops._C.rms_norm_dynamic_per_token_quant( + output, input, weight, scales, epsilon, scale_ub, residual + ) + return output, scales + + +# fused quant layer norm ops blocked +def rms_norm_per_block_quant( + input: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + quant_dtype: torch.dtype, + group_size: list[int], + scale_ub: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + is_scale_transposed: bool = False, + tma_alignment: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + assert len(group_size) == 2 + output = torch.empty(input.shape, dtype=quant_dtype, device=input.device) + if is_scale_transposed: + if tma_alignment == 0: + scales = torch.empty( + (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]), + device=input.device, + dtype=torch.float32, + ).transpose(0, 1) + else: + m = input.shape[-2] + sf_k = input.shape[-1] // group_size[1] + tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment + shape = input.shape[:-2] + (m, sf_k) + stride = ( + (1, tma_aligned_m) + if input.dim() == 2 + else (tma_aligned_m * sf_k, 1, tma_aligned_m) + ) + scales = torch.empty_strided( + shape, stride, device=input.device, dtype=torch.float32 + ) + else: + scales = torch.empty( + (input.numel() // input.shape[-1], input.shape[-1] // group_size[1]), + device=input.device, + dtype=torch.float32, + ) + + assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str( + tma_alignment + ) + + torch.ops._C.rms_norm_per_block_quant( + output, + input, + weight, + scales, + epsilon, + scale_ub, + residual, + group_size[1], + is_scale_transposed, + ) + return output, scales + + +# fused silu_and_mul + block quant +def silu_and_mul_per_block_quant( + input: torch.Tensor, + group_size: int, # Changed from list[int] + quant_dtype: torch.dtype, + scale_ub: torch.Tensor | None = None, + is_scale_transposed: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + assert input.ndim == 2, f"input must be 2D [batch, hidden*2], got {input.shape}" + assert input.shape[-1] % 2 == 0, ( + f"input last dim must be even (gate||up layout), got {input.shape[-1]}" + ) + + # Output is half the width of input (after silu_and_mul) + num_tokens = input.shape[0] + hidden_size = input.shape[-1] // 2 # Divide by 2 because input is [gate || up] + + # Allocate output tensor (FP8 or INT8) + output = torch.empty( + (num_tokens, hidden_size), device=input.device, dtype=quant_dtype + ) + + # Allocate scales tensor + num_groups = hidden_size // group_size # Directly use group_size + if is_scale_transposed: + scales = torch.empty( + (num_groups, num_tokens), + device=input.device, + dtype=torch.float32, + ).t() + else: + scales = torch.empty( + (num_tokens, num_groups), + device=input.device, + dtype=torch.float32, + ) + + # Call the C++ kernel + torch.ops._C.silu_and_mul_per_block_quant( + output, + input, + scales, + group_size, # Pass directly as int + scale_ub, + is_scale_transposed, + ) + + return output, scales + + +# quantization ops +# awq +def awq_dequantize( + qweight: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor, + split_k_iters: int, + thx: int, + thy: int, +) -> torch.Tensor: + if envs.VLLM_USE_TRITON_AWQ: + from vllm.model_executor.layers.quantization.awq_triton import ( + awq_dequantize_triton, + ) + + return awq_dequantize_triton(qweight, scales, zeros) + return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters, thx, thy) + + +if hasattr(torch.ops._C, "awq_dequantize"): + + @register_fake("_C::awq_dequantize") + def _awq_dequantize_fake( + qweight: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor, + split_k_iters: torch.SymInt, + thx: int, + thy: int, + ) -> torch.Tensor: + in_c = qweight.size(0) + qout_c = qweight.size(1) + out_c = qout_c * 8 + return torch.empty((in_c, out_c), dtype=scales.dtype, device=scales.device) + + +def awq_gemm( + input: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + split_k_iters: int, +) -> torch.Tensor: + if envs.VLLM_USE_TRITON_AWQ: + from vllm.model_executor.layers.quantization.awq_triton import awq_gemm_triton + + return awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters) + return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters) + + +if hasattr(torch.ops._C, "awq_gemm"): + + @register_fake("_C::awq_gemm") + def _awq_gemm_fake( + input: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + split_k_iters: torch.SymInt, + ) -> torch.Tensor: + num_in_feats = input.size(0) + return torch.empty( + (split_k_iters, num_in_feats, qweight.size(1) * 8), + dtype=input.dtype, + device=input.device, + ).sum(0) + + +# gptq +def gptq_gemm( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_gptq_qzeros: torch.Tensor, + b_gptq_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_exllama: bool, + use_v2_format: bool, + bit: int, +) -> torch.Tensor: + return torch.ops._C.gptq_gemm( + a, + b_q_weight, + b_gptq_qzeros, + b_gptq_scales, + b_g_idx, + use_exllama, + use_v2_format, + bit, + ) + + +if hasattr(torch.ops._C, "gptq_gemm"): + + @register_fake("_C::gptq_gemm") + def _gptq_gemm_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_gptq_qzeros: torch.Tensor, + b_gptq_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_exllama: bool, + use_v2_format: bool, + bit: int, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + +def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None: + torch.ops._C.gptq_shuffle(q_weight, q_perm, bit) + + +def gptq_gemm_rdna3( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, +) -> torch.Tensor: + return torch.ops._rocm_C.gptq_gemm_rdna3( + a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"): + + @register_fake("_rocm_C::gptq_gemm_rdna3") + def _gptq_gemm_rdna3_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"): + + @register_fake("_rocm_C::gptq_gemm_rdna3_wmma") + def _gptq_gemm_rdna3_wmma_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + +def moe_gptq_gemm_rdna3( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + block_size_m: int, + mul_topk_weight: bool, + output_topk: int = 0, +) -> None: + torch.ops._rocm_C.moe_gptq_gemm_rdna3( + a, + c, + b_q_weight, + b_scales, + b_qzeros, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + block_size_m, + mul_topk_weight, + output_topk, + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3"): + + @register_fake("_rocm_C::moe_gptq_gemm_rdna3") + def _moe_gptq_gemm_rdna3_fake( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + block_size_m: int, + mul_topk_weight: bool, + output_topk: int = 0, + ) -> None: + return + + +if hasattr(torch.ops._C, "allspark_w8a16_gemm"): + + @register_fake("_C::allspark_w8a16_gemm") + def _allspark_w8a16_gemm_fake( + a: torch.Tensor, + b_qweight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor | None, + n: torch.SymInt, + group_size: torch.SymInt, + sm_count: torch.SymInt, + sm_version: torch.SymInt, + CUBLAS_M_THRESHOLD: torch.SymInt, + has_zp: bool, + n32k16_reorder: bool, + ) -> torch.Tensor: + m = a.size(0) + return torch.empty((m, n), device=a.device, dtype=a.dtype) + + +if hasattr(torch.ops._C, "ggml_dequantize"): + + @register_fake("_C::ggml_dequantize") + def _ggml_dequantize_fake( + W: torch.Tensor, + quant_type: int, + m: torch.SymInt, + n: torch.SymInt, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + return torch.empty((m, n), dtype=torch.float16, device=W.device) + + @register_fake("_C::ggml_mul_mat_vec_a8") + def _ggml_mul_mat_vec_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: torch.SymInt, + ) -> torch.Tensor: + return torch.empty((X.shape[0], row), dtype=X.dtype, device=W.device) + + @register_fake("_C::ggml_mul_mat_a8") + def _ggml_mul_mat_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: torch.SymInt, + ) -> torch.Tensor: + batch = X.size(0) + return torch.empty((batch, row), dtype=X.dtype, device=W.device) + + @register_fake("_C::ggml_moe_a8") + def _ggml_moe_a8_fake( + X: torch.Tensor, + W: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + quant_type: int, + row: torch.SymInt, + top_k: torch.SymInt, + tokens: torch.SymInt, + ) -> torch.Tensor: + tokens = X.size(0) + return torch.empty((tokens * top_k, row), dtype=torch.float16, device=W.device) + + +if hasattr(torch.ops._C, "ggml_moe_a8_vec"): + + @register_fake("_C::ggml_moe_a8_vec") + def _ggml_moe_a8_vec_fake( + X: torch.Tensor, + W: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: torch.SymInt, + tokens: torch.SymInt, + ) -> torch.Tensor: + tokens = X.size(0) + return torch.empty((tokens * top_k, row), dtype=X.dtype, device=W.device) + + +# cutlass +def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool: + return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability) + + +def cutlass_scaled_fp4_mm( + a: torch.Tensor, + b: torch.Tensor, + block_scale_a: torch.Tensor, + block_scale_b: torch.Tensor, + alpha: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + assert a.ndim == 2 and b.ndim == 2 + m, n = a.shape[0], b.shape[0] + out = torch.empty((m, n), dtype=out_dtype, device=a.device) + torch.ops._C.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b, alpha) + return out + + +def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool: + return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability) + + +def cutlass_scaled_mm_supports_block_fp8(cuda_device_capability: int) -> bool: + return torch.ops._C.cutlass_scaled_mm_supports_block_fp8(cuda_device_capability) + + +def cutlass_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + """ + `cutlass_scaled_mm` implements a fused version of + `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)` + where scale_a * a and scale_b * b are implemented using numpy-style + broadcasting. + + In order to support blockwise scaling like found in DeepSeek V3 we also + support extended "group" broadcast rules. We extend the numpy-style + broadcasting rules with the following rule: + "if the extent of a dimension in the source shape is between 1 and + corresponding extent in the target shape we repeat each element along + that dimension src_shape[dim] // target_shape[dim] times consecutively" + example if we have: + a = [[1, 2], and target_shape = (2, 4) + [3, 4]] + then we would expand a to: + a = [[1, 1, 2, 2], + [3, 3, 4, 4]] + currently we only support the case: + scale_a.shape * [1, 128] == a.shape + scale_b.shape * [128, 128] == b.shape + """ + assert out_dtype is torch.bfloat16 or out_dtype is torch.float16 + assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype + + # Massage the input to be 2D + target_shape = (*a.shape[:-1], b.shape[1]) + a = a.view(-1, a.shape[-1]) + + cutlass_compatible_b = b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0 + if current_platform.is_rocm() or not cutlass_compatible_b: + from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa + triton_scaled_mm, + ) + + out = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias) + else: + out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device) + torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias) + + return out.view(*target_shape) + + +def cutlass_scaled_mm_azp( + a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + azp_adj: torch.Tensor, + azp: torch.Tensor | None = None, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Args: + azp_adj: In the per-tensor case, this should include the azp. + Always per-channel. + azp: Only set in the per-token case. Per-token if set. + """ + assert b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0 + assert out_dtype is torch.bfloat16 or out_dtype is torch.float16 + assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype + + # Massage the input to be 2D + target_shape = (*a.shape[:-1], b.shape[1]) + a = a.view(-1, a.shape[-1]) + assert azp is None or azp.numel() == a.shape[0] + + out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device) + torch.ops._C.cutlass_scaled_mm_azp(out, a, b, scale_a, scale_b, azp_adj, azp, bias) + return out.view(*target_shape) + + +def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool: + if cuda_device_capability < 90 or cuda_device_capability >= 110: + return False + try: + return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability) + except AttributeError: + # Return False on non-CUDA platforms where it is not available + return False + + +def get_cutlass_moe_mm_data( + topk_ids: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + input_permutation: torch.Tensor, + output_permutation: torch.Tensor, + num_experts: int, + n: int, + k: int, + blockscale_offsets: torch.Tensor | None = None, + is_gated: bool = True, +): + """ + Prepare data necessary to perform CUTLASS grouped matrix multiplications + used in CUTLASS-based fused MoE. + + The function takes in topk_ids (token-expert mapping) and uses it to + compute: + - expert_offsets: Indices that mark at which token index each expert begins + its computation after the input is sorted with + input_permutation. The number of tokens computed with + expert E is expert_offsets[E + 1] - expert_offsets[E] + - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's + multiplication in two grouped MMs used in + the fused MoE operation. + - input_permutation: Permutation that must be used to shuffle the input + before executing the MMs. + - output_permutation: Permutation that must be used to shuffle the output + after executing the MMs. + - blockscale_offsets: Optional argument passed for fp4 moe. Indices that + mark at which block scale index each expert begins + its computation. The number of block scale rows + computed with expert E is blockscale_offsets[E + 1] - + blockscale_offsets[E] + - is_gated: Whether the activation is gated (gate + up). When True, the + first GEMM N dimension is 2*n; when False, it is n. + """ + return torch.ops._C.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + n, + k, + blockscale_offsets, + is_gated, + ) + + +def get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + expert_first_token_offset: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + n: int, + k: int, + swap_ab: bool, +): + """Compute per-expert (M, N, K) problem sizes from expert_first_token_offset""" + return torch.ops._C.get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + expert_first_token_offset, + problem_sizes1, + problem_sizes2, + n, + k, + swap_ab, + ) + + +def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor): + """ + Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor. + This is used in MoE to permute the input tensor before performing grouped matrix multiplications. + """ + num_tokens_permuted = dst2src_map.shape[0] + output_tensor = torch.empty( + (num_tokens_permuted, input_tensor.shape[1]), + device=input_tensor.device, + dtype=input_tensor.dtype, + ) + torch.ops._moe_C.shuffle_rows(input_tensor, dst2src_map, output_tensor) + return output_tensor + + +def get_cutlass_batched_moe_mm_data( + expert_offsets: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + expert_num_tokens: torch.Tensor, + num_local_experts: int, + padded_m: int, + n: int, + k: int, +): + """ + Prepare data necessary to perform CUTLASS grouped matrix multiplications + used in CUTLASS-based fused MoE. + + The function takes in expert_num_tokens (token count per expert) and + non_zero_expert_idxs (consecutive indices of experts with non-zero token + counts) and uses them to compute: + - expert_offsets: Indices that mark at which token index each expert begins + its computation. + - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's + multiplication in two grouped MMs used in + the fused MoE operation. + """ + return torch.ops._C.get_cutlass_batched_moe_mm_data( + expert_offsets, + problem_sizes1, + problem_sizes2, + expert_num_tokens, + num_local_experts, + padded_m, + n, + k, + ) + + +def cutlass_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes: torch.Tensor, + a_strides: torch.Tensor, + b_strides: torch.Tensor, + c_strides: torch.Tensor, + per_act_token: bool, + per_out_ch: bool, +): + """ + A single grouped matrix multiplication used in CUTLASS-based fused MoE. + The function executes fp8-quantized OUT = AB matrix multiplication. + + - expert_offsets: Indices that mark at which token index each expert begins + its computation. The number of tokens computed with + expert E is expert_offsets[E + 1] - expert_offsets[E] + - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped + MMs used in the fused MoE operation. + - a/b/c_strides: The data strides passed to grouped matrix multiplication. + """ + return torch.ops._C.cutlass_moe_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + expert_offsets, + problem_sizes, + a_strides, + b_strides, + c_strides, + per_act_token, + per_out_ch, + ) + + +def cutlass_fp4_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + alphas: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + sf_offsets: torch.Tensor, +): + """ + An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs + the gemms for each combination based on the specified problem sizes. + + This is used as the MoE gemm during NVFP4 Quantized FusedMoE forward. + - a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized + input and expert weights. + - a_/b_scales: The blockscales in FP8-E4M3 precision + - expert_offsets/sf_offsets: Indices that mark at which token index + each expert begins its computation. The number of tokens + computed with expert E is expert_offsets[E + 1] - + expert_offsets[E] And the sf_size per expert is + sf_offset[E+1] - sf_offset[E] + - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped + MMs used in the fused MoE operation. + """ + return torch.ops._C.cutlass_fp4_group_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + alphas, + problem_sizes, + expert_offsets, + sf_offsets, + ) + + +def cutlass_mxfp4_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + sf_offsets: torch.Tensor, +): + """ + An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4). + + Uses mx_float4_t types with E8M0 scale factors and 32-element blocks. + - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte) + - a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout) + - Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales). + - expert_offsets/sf_offsets: expert boundary indices + - problem_sizes: (num_experts, 3) with (M, N, K) per expert + """ + return torch.ops._C.cutlass_mxfp4_group_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + problem_sizes, + expert_offsets, + sf_offsets, + ) + + +def mxfp8_experts_quant( + input_tensor: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + quant_output: torch.Tensor, + scale_factor: torch.Tensor, +) -> None: + torch.ops._C.mxfp8_experts_quant( + input_tensor, + problem_sizes, + expert_offsets, + blockscale_offsets, + quant_output, + scale_factor, + ) + + +def cutlass_mxfp8_grouped_mm( + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + out_tensors: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, +) -> None: + torch.ops._C.cutlass_mxfp8_grouped_mm( + a_tensors, + b_tensors, + a_scales, + b_scales, + out_tensors, + problem_sizes, + expert_offsets, + blockscale_offsets, + ) + + +if hasattr(torch.ops._C, "mxfp8_experts_quant"): + + @register_fake("_C::mxfp8_experts_quant") + def _mxfp8_experts_quant_fake( + input_tensor: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + quant_output: torch.Tensor, + scale_factor: torch.Tensor, + ) -> None: + return None + + +if hasattr(torch.ops._C, "cutlass_mxfp8_grouped_mm"): + + @register_fake("_C::cutlass_mxfp8_grouped_mm") + def _cutlass_mxfp8_grouped_mm_fake( + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + out_tensors: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + ) -> None: + return None + + +# gptq_marlin +def gptq_marlin_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + return torch.ops._C.gptq_marlin_repack( + b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit + ) + + +if hasattr(torch.ops._C, "gptq_marlin_repack"): + + @register_fake("_C::gptq_marlin_repack") + def _gptq_marlin_repack_fake( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: torch.SymInt, + size_n: torch.SymInt, + num_bits: int, + is_a_8bit: bool = False, + ) -> torch.Tensor: + pack_factor = 32 // num_bits + marlin_tile_size = 16 + return torch.empty( + (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor), + dtype=b_q_weight.dtype, + device=b_q_weight.device, + ) + + +# awq_marlin +def awq_marlin_repack( + b_q_weight: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + return torch.ops._C.awq_marlin_repack( + b_q_weight, size_k, size_n, num_bits, is_a_8bit + ) + + +if hasattr(torch.ops._C, "awq_marlin_repack"): + + @register_fake("_C::awq_marlin_repack") + def _awq_marlin_repack_fake( + b_q_weight: torch.Tensor, + size_k: torch.SymInt, + size_n: torch.SymInt, + num_bits: int, + is_a_8bit: bool = False, + ) -> torch.Tensor: + pack_factor = 32 // num_bits + marlin_tile_size = 16 + return torch.empty( + (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor), + dtype=b_q_weight.dtype, + device=b_q_weight.device, + ) + + +def gptq_marlin_moe_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + num_experts = b_q_weight.shape[0] + assert size_k % 16 == 0 + output = torch.empty( + (num_experts, size_k // 16, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype, + ) + for e in range(num_experts): + output[e] = torch.ops._C.gptq_marlin_repack( + b_q_weight[e], perm[e], size_k, size_n, num_bits, is_a_8bit + ) + return output + + +def awq_marlin_moe_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + num_experts = b_q_weight.shape[0] + assert size_k % 16 == 0 + output = torch.empty( + (num_experts, size_k // 16, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype, + ) + for e in range(num_experts): + output[e] = torch.ops._C.awq_marlin_repack( + b_q_weight[e], size_k, size_n, num_bits, is_a_8bit + ) + return output + + +def marlin_int4_fp8_preprocess( + qweight: torch.Tensor, + qzeros_or_none: torch.Tensor | None = None, + inplace: bool = False, +): + return torch.ops._C.marlin_int4_fp8_preprocess(qweight, qzeros_or_none, inplace) + + +def marlin_gemm( + a: torch.Tensor, + c: torch.Tensor | None, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool = True, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + return torch.ops._C.marlin_gemm( + a, + c, + b_q_weight, + b_bias, + b_scales, + a_scales, + global_scale, + b_zeros, + g_idx, + perm, + workspace, + b_q_type.id, + size_m, + size_n, + size_k, + is_k_full, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + ) + + +if hasattr(torch.ops._C, "marlin_gemm"): + + @register_fake("_C::marlin_gemm") + def _marlin_gemm_fake( + a: torch.Tensor, + c: torch.Tensor | None, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + b_q_type_id: int, + size_m: torch.SymInt, + size_n: torch.SymInt, + size_k: torch.SymInt, + is_k_full: bool = True, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, + ) -> torch.Tensor: + dtype = a.dtype + if dtype not in [torch.half, torch.bfloat16]: + dtype = b_scales.dtype + return torch.empty((size_m, size_n), device=a.device, dtype=dtype) + + +# machete +def machete_supported_schedules( + a_type: torch.dtype, + b_type: ScalarType, + group_scales_type: torch.dtype | None, + group_zeros_type: torch.dtype | None = None, + channel_scales_type: torch.dtype | None = None, + token_scales_type: torch.dtype | None = None, + out_type: torch.dtype | None = None, +) -> list[str]: + return torch.ops._C.machete_supported_schedules( + a_type, + b_type.id, + group_scales_type, + group_zeros_type, + channel_scales_type, + token_scales_type, + out_type, + ) + + +def machete_mm( + a: torch.Tensor, + # b_q Should be the tensor returned by machete_prepack_B + b_q: torch.Tensor, + b_type: ScalarType, + out_type: torch.dtype | None = None, + b_group_scales: torch.Tensor | None = None, + b_group_zeros: torch.Tensor | None = None, + b_group_size: int | None = None, + b_channel_scales: torch.Tensor | None = None, + a_token_scales: torch.Tensor | None = None, + schedule: str | None = None, +) -> torch.Tensor: + return torch.ops._C.machete_mm( + a, + b_q, + b_type.id, + out_type, + b_group_scales, + b_group_zeros, + b_group_size, + b_channel_scales, + a_token_scales, + schedule, + ) + + +if hasattr(torch.ops._C, "machete_mm"): + + @register_fake("_C::machete_mm") + def machete_mm_fake( + a: torch.Tensor, + # b_q Should be the tensor returned by machete_prepack_B + b_q: torch.Tensor, + b_type: ScalarType, + out_type: torch.dtype | None = None, + b_group_scales: torch.Tensor | None = None, + b_group_zeros: torch.Tensor | None = None, + b_group_size: int | None = None, + b_channel_scales: torch.Tensor | None = None, + a_token_scales: torch.Tensor | None = None, + schedule: str | None = None, + ) -> torch.Tensor: + m = a.size(0) + n = b_q.size(1) + return torch.empty((m, n), device=a.device, dtype=a.dtype) + + +def machete_prepack_B( + b_q_weight: torch.Tensor, + a_type: torch.dtype, + b_type: ScalarType, + group_scales_type: torch.dtype | None, +) -> torch.Tensor: + return torch.ops._C.machete_prepack_B( + b_q_weight, a_type, b_type.id, group_scales_type + ) + + +if hasattr(torch.ops._C, "machete_prepack_B"): + + @register_fake("_C::machete_prepack_B") + def machete_prepack_B_fake( + b_q_weight: torch.Tensor, + a_type: torch.dtype, + b_type: ScalarType, + group_scales_type: torch.dtype | None, + ) -> torch.Tensor: + return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format) + + +# CUTLASS W4A8 +def cutlass_w4a8_mm( + a: torch.Tensor, + # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b + b_q: torch.Tensor, + b_group_scales: torch.Tensor, + b_group_size: int, + b_channel_scales: torch.Tensor, + a_token_scales: torch.Tensor, + out_type: torch.dtype | None = None, + maybe_schedule: str | None = None, +) -> torch.Tensor: + return torch.ops._C.cutlass_w4a8_mm( + a, + b_q, + b_group_scales, + b_group_size, + b_channel_scales, + a_token_scales, + out_type, + maybe_schedule, + ) + + +if hasattr(torch.ops._C, "cutlass_w4a8_mm"): + + @register_fake("_C::cutlass_w4a8_mm") + def cutlass_w4a8_mm_fake( + a: torch.Tensor, + # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b + b_q: torch.Tensor, + b_group_scales: torch.Tensor, + b_group_size: int, + b_channel_scales: torch.Tensor, + a_token_scales: torch.Tensor, + out_type: torch.dtype | None = None, + maybe_schedule: str | None = None, + ) -> torch.Tensor: + m = a.size(0) + n = b_q.size(1) + out_dtype = out_type if out_type is not None else torch.bfloat16 + return torch.empty((m, n), device=a.device, dtype=out_dtype) + + +def cutlass_pack_scale_fp8(scales: torch.Tensor) -> torch.Tensor: + return torch.ops._C.cutlass_pack_scale_fp8(scales) + + +if hasattr(torch.ops._C, "cutlass_pack_scale_fp8"): + + @register_fake("_C::cutlass_pack_scale_fp8") + def cutlass_pack_scale_fp8_fake(scales: torch.Tensor) -> torch.Tensor: + return torch.empty_like(scales, memory_format=torch.contiguous_format) + + +def cutlass_encode_and_reorder_int4b(b: torch.Tensor) -> torch.Tensor: + return torch.ops._C.cutlass_encode_and_reorder_int4b(b) + + +if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b"): + + @register_fake("_C::cutlass_encode_and_reorder_int4b") + def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor: + return torch.empty_like(b, memory_format=torch.contiguous_format) + + +def cutlass_w4a8_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + b_group_scales: torch.Tensor, + b_group_size: int, + expert_offsets: torch.Tensor, + problem_sizes: torch.Tensor, + a_strides: torch.Tensor, + b_strides: torch.Tensor, + c_strides: torch.Tensor, + group_scale_strides: torch.Tensor, + maybe_schedule: str | None = None, +): + """ + Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the + W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8) + and both per-channel + per-token scaling in the epilogue. + + Args: + out_tensors: + Output buffer for all experts (updated in-place). + a_tensors: + FP8 (E4M3FN) activations for all experts. + b_tensors: + INT4-packed weight matrix for all experts, packed to INT32 + a_scales: + Per-token FP8 activation scales, applied in the epilogue. + b_scales: + Per-channel FP8 weight scales for each expert, applied in the epilogue. + b_group_scales: + FP8 scale values for group-wise INT4 weight blocks. + b_group_size: + Number of elements grouped under each entry of b_group_scales. + expert_offsets: + Cumulative token offsets + problem_sizes: + Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher. + a/b/c/group_scale_strides: + Strides describing the memory layout of the input tensors. + maybe_schedule: + Optional override to choose a specific kernel or epilogue schedule. + + Returns: + out_tensors updated in-place with the dequantized INT4xFP8 grouped GEMM result. + """ + return torch.ops._C.cutlass_w4a8_moe_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + b_group_scales, + b_group_size, + expert_offsets, + problem_sizes, + a_strides, + b_strides, + c_strides, + group_scale_strides, + maybe_schedule, + ) + + +def cutlass_encode_and_reorder_int4b_grouped( + b_tensors: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.cutlass_encode_and_reorder_int4b_grouped(b_tensors) + + +if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b_grouped"): + + @register_fake("_C::cutlass_encode_and_reorder_int4b_grouped") + def cutlass_encode_and_reorder_int4b_grouped_fake(b: torch.Tensor) -> torch.Tensor: + return torch.empty_like(b, memory_format=torch.contiguous_format) + + +def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: + return torch.ops._C.permute_cols(a, perm) + + +if hasattr(torch.ops._C, "permute_cols"): + + @register_fake("_C::permute_cols") + def _permute_cols_fake(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: + return torch.empty_like(a) + + +# fp4 +def scaled_fp4_quant( + input: torch.Tensor, + input_global_scale: torch.Tensor, + is_sf_swizzled_layout: bool = True, + backend: str = "none", + padded_n: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to FP4 and return quantized tensor and scale. + + This function quantizes the last dimension of the given tensor `input`. For + every 16 consecutive elements, a single dynamically computed scaling factor + is shared. This scaling factor is quantized using the `input_global_scale` + and is stored in a swizzled layout (see + https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x). + + Args: + input: The input tensor to be quantized to FP4 + input_global_scale: A scalar scaling factor for the entire tensor. + use_8x4_sf_layout: Whether to use the 8x4 or 128x4 layout for the scaling + padded_n: Optional padded K dimension. When provided, the quantized + output and scale tensors are allocated for ``padded_n`` + + Returns: + tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every + two values are packed into a uint8 and float8_e4m3 scaling factors + in the sizzled layout. + """ + assert not current_platform.is_rocm() + assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}." + other_dims = 1 if input.ndim == 1 else -1 + input = input.reshape(other_dims, input.shape[-1]) + m, n = input.shape + block_size = 16 + + assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}." + assert input.dtype in (torch.float16, torch.bfloat16), ( + f"input.dtype needs to be fp16 or bf16 but got {input.dtype}." + ) + if padded_n is not None: + assert padded_n >= n, f"padded_n must be >= n, got padded_n={padded_n}, n={n}." + assert padded_n % block_size == 0, ( + f"padded_n has to be a multiple of {block_size}, but got {padded_n}." + ) + + use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False # noqa: SIM210 + if use_8x4_sf_layout and padded_n is not None and padded_n != n: + # TODO: support this case + raise ValueError("padded_n is not supported with TRTLLM 8x4 scale layout.") + if use_8x4_sf_layout: + output, output_scale = flashinfer_quant_nvfp4_8x4_sf_layout( + input, input_global_scale + ) + else: + # Pre-allocate and call .out variant (same behavior as old in-place API) + output, output_scale = create_fp4_output_tensors( + m, + n, + input.device, + is_sf_swizzled_layout, + padded_n=padded_n, + ) + torch.ops._C.scaled_fp4_quant.out( + input, + input_global_scale, + is_sf_swizzled_layout, + output=output, + output_scale=output_scale, + ) + + output_scale = output_scale.view(torch.float8_e4m3fn) + return output, output_scale + + +def scaled_fp4_experts_quant( + input_tensor: torch.Tensor, + input_global_scale: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to NVFP4 and return quantized tensor and scale, for + packed MoE Inputs. + Args: + input_tensor: The input tensor to be quantized to NVFP4 + input_global_scale: A scalar scaling factor for the entire tensor. + expert_offsets: The expert offsets tensor + blockscale_offsets: The blockscale offsets tensor + Outputs: + output: The quantized tensor in NVFP4 + output_scales: The blockscale tensor in FP8-E4M3 + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2, ( + f"input.ndim needs to be == 2, but got {input_tensor.ndim}." + ) + + # Control the maximum number of tokens per expert supported by the + # NVFP4 MoE Expert Quantization. This is used to prevent the kernel + # from running out of memory. This value can also be increased to support + # larger models. + MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k = input_tensor.shape + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT(" + f"{MAX_TOKENS_PER_EXPERT})" + f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" + f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value." + ) + scales_k = k // 16 + padded_k = (scales_k + (4 - 1)) // 4 + + # output is uint8 and packed fp4 values + output = torch.empty( + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 + ) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.scaled_fp4_experts_quant( + output, + output_scales, + input_tensor, + input_global_scale, + expert_offsets, + blockscale_offsets, + ) + output_scales = output_scales.view(torch.float8_e4m3fn) + return output, output_scales + + +def silu_and_mul_scaled_fp4_experts_quant( + input_tensor: torch.Tensor, + input_global_scale: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations. + + Args: + input_tensor: The input tensor with gate || up layout [m_topk, k*2] + input_global_scale: A per-expert scaling factor [n_experts] + expert_offsets: The expert offsets tensor [n_experts+1] + blockscale_offsets: The blockscale offsets tensor [n_experts+1] + topk: Number of top-k experts selected + Outputs: + output: The quantized tensor in NVFP4 [m_topk, k/2] + output_scales: The blockscale tensor in FP8-E4M3 + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2, ( + f"input.ndim needs to be == 2, but got {input_tensor.ndim}." + ) + + # Control the maximum number of tokens per expert supported by the + # NVFP4 MoE Expert Quantization. This is used to prevent the kernel + # from running out of memory. This value can also be increased to support + # larger models. + MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k_times_2 = input_tensor.shape + assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)" + k = k_times_2 // 2 + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT(" + f"{MAX_TOKENS_PER_EXPERT})" + f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" + f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value." + ) + scales_k = k // 16 + padded_k = (scales_k + (4 - 1)) // 4 + + # output is uint8 and packed fp4 values + output = torch.empty( + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 + ) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.silu_and_mul_scaled_fp4_experts_quant( + output, + output_scales, + input_tensor, + input_global_scale, + expert_offsets, + blockscale_offsets, + ) + output_scales = output_scales.view(torch.float8_e4m3fn) + return output, output_scales + + +def mxfp4_experts_quant( + input_tensor: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + n_experts: int, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to MXFP4 for packed MoE inputs. + Uses 32-element blocks with E8M0 (power-of-two) scale factors. + MXFP4 has no global scale - only block-level E8M0 scale factors. + + Args: + input_tensor: [m_topk, k] BF16/FP16 activations + expert_offsets: [n_experts+1] token boundaries per expert + blockscale_offsets: [n_experts+1] SF row boundaries per expert + n_experts: number of experts + topk: number of top-k experts + Returns: + output: [m_topk, k//2] packed E2M1 values (uint8) + output_scales: E8M0 blockscales in swizzled layout (uint8 view) + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2 + + MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k = input_tensor.shape + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT(" + f"{MAX_TOKENS_PER_EXPERT})" + f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use" + f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value." + ) + scales_k = k // 32 + padded_k = (scales_k + (4 - 1)) // 4 + + output = torch.empty( + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 + ) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.mxfp4_experts_quant( + output, + output_scales, + input_tensor, + expert_offsets, + blockscale_offsets, + n_experts, + ) + # E8M0 SFs are stored as uint8 + output_scales = output_scales.view(torch.uint8) + return output, output_scales + + +def silu_and_mul_mxfp4_experts_quant( + input_tensor: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + n_experts: int, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations. + MXFP4 has no global scale - only block-level E8M0 scale factors. + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2 + + MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k_times_2 = input_tensor.shape + assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)" + k = k_times_2 // 2 + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk + scales_k = k // 32 + padded_k = (scales_k + (4 - 1)) // 4 + + output = torch.empty( + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 + ) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.silu_and_mul_mxfp4_experts_quant( + output, + output_scales, + input_tensor, + expert_offsets, + blockscale_offsets, + n_experts, + ) + output_scales = output_scales.view(torch.uint8) + return output, output_scales + + +# fp8 +def scaled_fp8_quant( + input: torch.Tensor, + scale: torch.Tensor | None = None, + num_token_padding: int | None = None, + scale_ub: torch.Tensor | None = None, + use_per_token_if_dynamic: bool = False, + output: torch.Tensor | None = None, + group_shape: tuple[int, int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to FP8 and return quantized tensor and scale. + + This function supports both static and dynamic quantization: If you + provide the scale, it will use static scaling and if you omit it, + the scale will be determined dynamically. The function also allows + optional padding of the output tensors for downstream kernels that + will benefit from padding. + + Args: + input: The input tensor to be quantized to FP8 (must be 2D: [M, N]) + scale: Optional scaling factor for the FP8 quantization. Supports: + - 0D or [1]: per-tensor scaling + - 1D: requires explicit group_shape to disambiguate per-channel + vs per-token (use (-1, 1) for per-channel, (1, -1) for per-token) + - 2D [M/group_m, N/group_n]: group scaling (e.g. [M, N/128] for + DeepSeek-style (1,128) groups, or [M/128, N/128] for (128,128)) + scale_ub: Optional upper bound for scaling factor in dynamic + per token case + num_token_padding: If specified, pad the first dimension + of the output to at least this value. + use_per_token_if_dynamic: Whether to do per_tensor or per_token + in the dynamic quantization case. + group_shape: Optional tuple (group_m, group_n) specifying the group + shape for static quantization. Use -1 for "full extent" (e.g., + (-1, -1) for per-tensor, (-1, 1) for per-channel, etc.) + Required for 1D scales; optional for 2D scales. + + Returns: + tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and + scaling factor. + """ + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + shape: tuple[int, int] | torch.Size = input.shape + # For ROCm on MI300, the output fp8 dtype is torch.float_e3m3fnuz + out_dtype: torch.dtype = current_platform.fp8_dtype() + if num_token_padding: + shape = (max(num_token_padding, input.shape[0]), shape[1]) + if output is None: + output = torch.empty(shape, device=input.device, dtype=out_dtype) + else: + assert num_token_padding is None, "padding not supported if output passed in" + assert output.dtype == out_dtype + + if scale is None: + if use_per_token_if_dynamic: + scale = torch.empty((shape[0], 1), device=input.device, dtype=torch.float32) + torch.ops._C.dynamic_per_token_scaled_fp8_quant( + output, input, scale, scale_ub + ) + else: + scale = torch.empty(1, device=input.device, dtype=torch.float32) + torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale) + else: + torch.ops._C.static_scaled_fp8_quant(output, input, scale, group_shape) + + return output, scale + + +# gptq allspark +def allspark_repack_weight( + qweight: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor | None = None, + has_zp: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format + for Ampere W8A16 Fused Gemm kernel + + Args: + qweight: uint8 weight tensor, original k x n format. + scale: fp16/bf16 weight scale tensor, 1 x n format. + zero_point: fp16/bf16 weight zero_point tensor, 1 x n format. + Must be provided for asymmetric quantization. + has_zp: if use symmetric quantization, has_zp = False. + if use asymmetric quantization, has_zp = True. + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : + rearranged weight, scale, and optionally zero_point. + """ + K = qweight.shape[0] + N = qweight.shape[1] + N_32align = (N + 32 - 1) // 32 * 32 + + qweight_reorder = torch.empty( + (N_32align, K), device=qweight.device, dtype=qweight.dtype + ) + scale_reorder = torch.empty((1, N_32align), device=scale.device, dtype=scale.dtype) + zero_point_reorder = None + if has_zp: + assert zero_point is not None, ( + "zero_point must be provided for asymmetric quantization." + ) + zero_point_reorder = torch.empty( + (1, N_32align), device=zero_point.device, dtype=zero_point.dtype + ) + + torch.ops._C.rearrange_kn_weight_as_n32k16_order( + qweight, + scale, + zero_point, + has_zp, + qweight_reorder, + scale_reorder, + zero_point_reorder, + K, + N, + N_32align, + ) + + return qweight_reorder, scale_reorder, zero_point_reorder + + +def allspark_w8a16_gemm( + a: torch.Tensor, + b_qweight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor | None, + n: int, + group_size: int, + sm_count: int, + sm_version: int, + CUBLAS_M_THRESHOLD: int, + has_zp: bool, + n32k16_reorder: bool, +) -> torch.Tensor: + return torch.ops._C.allspark_w8a16_gemm( + a, + b_qweight, + b_scales, + b_qzeros, + n, + group_size, + sm_count, + sm_version, + CUBLAS_M_THRESHOLD, + has_zp, + n32k16_reorder, + ) + + +# int8 +def scaled_int8_quant( + input: torch.Tensor, + scale: torch.Tensor | None = None, + azp: torch.Tensor | None = None, + symmetric: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """ + Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. + + Args: + input: The input tensor to be quantized to int8. + scale: Optional scaling factor for the int8 quantization. + When not provided, we invoke dynamic-per-token quantization. + azp: Optional zero-point for the int8 quantization. + Must be provided for asymmetric quantization if `scale` is provided. + symmetric: Whether to use symmetric quantization (scale only, azp ignored). + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp. + """ + output = torch.empty_like(input, dtype=torch.int8) + if scale is not None: + # static-per-tensor quantization. + assert symmetric == (azp is None), ( + "azp must only be provided for asymmetric quantization." + ) + torch.ops._C.static_scaled_int8_quant(output, input, scale, azp) + return output, scale, azp + + # dynamic-per-token quantization. + input_scales = torch.empty( + (input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32 + ) + input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32) + torch.ops._C.dynamic_scaled_int8_quant( + output, input.contiguous(), input_scales, input_azp + ) + return output, input_scales, input_azp + + +# gguf +def ggml_dequantize( + W: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None +) -> torch.Tensor: + return torch.ops._C.ggml_dequantize(W, quant_type, m, n, dtype) + + +def ggml_mul_mat_vec_a8( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, +) -> torch.Tensor: + return torch.ops._C.ggml_mul_mat_vec_a8(W, X, quant_type, row) + + +def ggml_mul_mat_a8( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, +) -> torch.Tensor: + return torch.ops._C.ggml_mul_mat_a8(W, X, quant_type, row) + + +def ggml_moe_a8( + X: torch.Tensor, + W: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + quant_type: int, + row: int, + top_k: int, + tokens: int, +) -> torch.Tensor: + return torch.ops._C.ggml_moe_a8( + X, + W, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type, + row, + top_k, + tokens, + ) + + +def ggml_moe_a8_vec( + X: torch.Tensor, + W: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: torch.SymInt, + tokens: torch.SymInt, +) -> torch.Tensor: + return torch.ops._C.ggml_moe_a8_vec(X, W, topk_ids, top_k, quant_type, row, tokens) + + +def ggml_moe_get_block_size(quant_type: int) -> int: + return torch.ops._C.ggml_moe_get_block_size(quant_type) + + +# mamba +def selective_scan_fwd( + u: torch.Tensor, + delta: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D_: torch.Tensor | None, + z_: torch.Tensor | None, + delta_bias_: torch.Tensor | None, + delta_softplus: bool, + query_start_loc: torch.Tensor | None, + cache_indices: torch.Tensor | None, + has_initial_state: torch.Tensor | None, + ssm_states: torch.Tensor, + null_block_id: int, + block_size: int = 1024, + block_idx_first_scheduled_token: torch.Tensor | None = None, + block_idx_last_scheduled_token: torch.Tensor | None = None, + initial_state_idx: torch.Tensor | None = None, + cu_chunk_seqlen: torch.Tensor | None = None, + last_chunk_indices: torch.Tensor | None = None, +): + torch.ops._C.selective_scan_fwd( + u, + delta, + A, + B, + C, + D_, + z_, + delta_bias_, + delta_softplus, + query_start_loc, + cache_indices, + has_initial_state, + ssm_states, + null_block_id, + block_size, + block_idx_first_scheduled_token, + block_idx_last_scheduled_token, + initial_state_idx, + cu_chunk_seqlen, + last_chunk_indices, + ) + + +# ROCm skinny gemms +def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor: + return torch.ops._rocm_C.LLMM1(a, b, rows_per_block) + + +def wvSplitK( + a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None +) -> torch.Tensor: + return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count) + + +def wvSplitKrc( + a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None +) -> torch.Tensor: + return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count) + + +def wvSplitKQ( + a: torch.Tensor, + b: torch.Tensor, + out_dtype: torch.dtype, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + cu_count: int, + bias: torch.Tensor = None, +) -> torch.Tensor: + out = torch.empty((b.shape[0], a.shape[0]), dtype=out_dtype, device=b.device) + torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count) + return out + + +# moe +def moe_sum(input: torch.Tensor, output: torch.Tensor): + torch.ops._moe_C.moe_sum(input, output) + + +def moe_align_block_size( + topk_ids: torch.Tensor, + num_experts: int, + block_size: int, + sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + expert_map: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.moe_align_block_size( + topk_ids, + num_experts, + block_size, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + expert_map, + ) + + +def batched_moe_align_block_size( + max_tokens_per_batch: int, + block_size: int, + expert_num_tokens: torch.Tensor, + sorted_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, +) -> None: + torch.ops._moe_C.batched_moe_align_block_size( + max_tokens_per_batch, + block_size, + expert_num_tokens, + sorted_ids, + expert_ids, + num_tokens_post_pad, + ) + + +def moe_lora_align_block_size( + topk_ids: torch.Tensor, + token_lora_mapping: torch.Tensor, + num_experts: int, + block_size: int, + max_loras: int, + max_num_tokens_padded: int, + max_num_m_blocks: int, + sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + adapter_enabled: torch.Tensor, + lora_ids: torch.Tensor, + expert_map: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.moe_lora_align_block_size( + topk_ids, + token_lora_mapping, + num_experts, + block_size, + max_loras, + max_num_tokens_padded, + max_num_m_blocks, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + adapter_enabled, + lora_ids, + expert_map, + ) + + +def moe_wna16_gemm( + input: torch.Tensor, + output: torch.Tensor, + b_qweight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor | None, + topk_weights: torch.Tensor | None, + sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + top_k: int, + BLOCK_SIZE_M: int, + BLOCK_SIZE_N: int, + BLOCK_SIZE_K: int, + bit: int, +) -> torch.Tensor: + if not current_platform.is_cuda(): + raise NotImplementedError( + "The optimized moe_wna16_gemm kernel is only available on CUDA platforms" + ) + torch.ops._moe_C.moe_wna16_gemm( + input, + output, + b_qweight, + b_scales, + b_qzeros, + topk_weights, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + top_k, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + bit, + ) + + +def dsv3_router_gemm( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + output = torch.empty( + hidden_states.shape[0], + router_weight.shape[0], + device=hidden_states.device, + dtype=output_dtype, + ) + torch.ops._moe_C.dsv3_router_gemm(output, hidden_states, router_weight) + return output + + +def fp32_router_gemm( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, +) -> torch.Tensor: + output = torch.empty( + hidden_states.shape[0], + router_weight.shape[0], + device=hidden_states.device, + dtype=torch.float32, + ) + torch.ops._C.fp32_router_gemm(output, hidden_states, router_weight) + return output + + +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"): + + @register_fake("_C::fp32_router_gemm") + def fp32_router_gemm_fake( + output: torch.Tensor, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + ) -> None: + return + + +def topk_softmax( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + e_score_correction_bias: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.topk_softmax( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize, + e_score_correction_bias, + ) + + +def topk_sigmoid( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + e_score_correction_bias: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.topk_sigmoid( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize, + e_score_correction_bias, + ) + + +def topk_hash_softplus_sqrt( + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, + input_tokens: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.topk_softplus_sqrt( + topk_weights, + topk_indices, + token_expert_indices, + gating_output, + renormalize, + routed_scaling_factor, + e_score_correction_bias, + input_tokens, + hash_indices_table, + ) + + +def grouped_topk( + scores: torch.Tensor, + num_expert_group: int, + topk_group: int, + topk: int, + renormalize: bool, + routed_scaling_factor: float, + bias: torch.Tensor, + scoring_func: int = 0, +): + """ + Perform grouped top-k routing for mixture of experts. + + Args: + scores: Raw inputs (logits if scoring_func=1, scores if scoring_func=0) + num_expert_group: Number of expert groups + topk_group: Number of groups to select + topk: Number of experts to select per token + renormalize: Whether to renormalize the output weights + routed_scaling_factor: Scaling factor for routing weights + bias: Bias tensor (e_score_correction_bias). Always fused in kernel. + scoring_func: 0=none (no activation), 1=sigmoid + """ + if not current_platform.is_cuda(): + raise NotImplementedError( + "The fused grouped_topk kernel is only available on CUDA platforms" + ) + return torch.ops._moe_C.grouped_topk( + scores, + num_expert_group, + topk_group, + topk, + renormalize, + routed_scaling_factor, + bias, + scoring_func, + ) + + +def moe_wna16_marlin_gemm( + input: torch.Tensor, + output: torch.Tensor | None, + b_qweight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_qzeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_past_padded: torch.Tensor, + topk_weights: torch.Tensor, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, + thread_k: int = -1, + thread_n: int = -1, + blocks_per_sm: int = -1, +) -> torch.Tensor: + return torch.ops._moe_C.moe_wna16_marlin_gemm( + input, + output, + b_qweight, + b_bias, + b_scales, + a_scales, + global_scale, + b_qzeros, + g_idx, + perm, + workspace, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + b_q_type.id, + size_m, + size_n, + size_k, + is_k_full, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + thread_k, + thread_n, + blocks_per_sm, + ) + + +if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_gemm"): + + @register_fake("_moe_C::moe_wna16_marlin_gemm") + def moe_wna16_marlin_gemm_fake( + input: torch.Tensor, + output: torch.Tensor | None, + b_qweight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_qzeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_past_padded: torch.Tensor, + topk_weights: torch.Tensor, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, + ): + return torch.empty( + (size_m * top_k, size_n), dtype=input.dtype, device=input.device + ) + + +def reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def concat_and_cache_mla( + kv_c: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.concat_and_cache_mla( + kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale + ) + + +def concat_and_cache_mla_rope_fused( + positions: torch.Tensor, + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + slot_mapping: torch.Tensor, + kv_cache: torch.Tensor, + kv_cache_dtype: str, + kv_cache_scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused( + positions, + q_pe, + k_pe, + kv_c, + cos_sin_cache, + is_neox, + slot_mapping, + kv_cache, + kv_cache_dtype, + kv_cache_scale, + ) + + +def swap_blocks( + src: torch.Tensor, + dst: torch.Tensor, + block_size_in_bytes: int, + block_mapping: torch.Tensor, +) -> None: + """ + Copy specific blocks from one tensor to another. + + This method assumes each of the two input tensors is composed of + consecutive contiguous blocks, of size block_size_in_bytes. + i.e. the memory layout for each tensor is: + [block0] [block1] ... [block N] + + block_mapping determines the subset of blocks to copy of the source tensor, + and their matching destination block number on the destination tensor. + block_mapping is expected to be a tensor of shape (num_blocks_to_copy, 2) + where each block_mapping[i] represents a single copy operation, copying + block #block_mapping[i][0] from the source tensor + to block #block_mapping[i][1] on the destination tensor. + block_mapping should have dtype int64. + + The source and the destination tensors can be either on cpu or gpu, + but not both on cpu. + the block mapping tensor must on cpu. + """ + torch.ops._C_cache_ops.swap_blocks(src, dst, block_size_in_bytes, block_mapping) + + +def swap_blocks_batch( + src_ptrs: torch.Tensor, + dst_ptrs: torch.Tensor, + sizes: torch.Tensor, + is_src_access_order_any: bool = False, +) -> None: + """ + Batch version of swap_blocks: submit all copies in a single driver call. + + Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i] + of sizes[i] bytes. All three tensors must be CPU tensors with the + platform-appropriate pointer dtype: int64 on CUDA/ROCm (required by + cache_kernels.cu) and uint64 on XPU (required by the XPU DMA engine). + On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission + overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync. + + is_src_access_order_any: if True, pass CU_MEMCPY_SRC_ACCESS_ORDER_ANY to + cuMemcpyBatchAsync, letting the DMA engine prefetch source bytes + out of stream order. Only safe when no GPU stream is concurrently + writing to the source. Defaults to False (STREAM ordering), which + is always safe. + """ + if current_platform.is_xpu(): + torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes) + else: + torch.ops._C_cache_ops.swap_blocks_batch( + src_ptrs, dst_ptrs, sizes, is_src_access_order_any + ) + + +def convert_fp8( + output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8" +) -> None: + torch.ops._C_cache_ops.convert_fp8(output, input, scale, kv_dtype) + + +def gather_and_maybe_dequant_cache( + src_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, + token_to_seq: torch.Tensor, + num_tokens: int, + kv_cache_dtype: str, + scale: torch.Tensor, + seq_starts: torch.Tensor | None = None, +) -> None: + torch.ops._C_cache_ops.gather_and_maybe_dequant_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + token_to_seq, + num_tokens, + kv_cache_dtype, + scale, + seq_starts, + ) + + +def cp_gather_cache( + src_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, + batch_size: int, + seq_starts: torch.Tensor | None = None, +) -> None: + torch.ops._C_cache_ops.cp_gather_cache( + src_cache, dst, block_table, cu_seq_lens, batch_size, seq_starts + ) + + +def cp_gather_and_upconvert_fp8_kv_cache( + src_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + workspace_starts: torch.Tensor, + batch_size: int, +) -> None: + """Gather and upconvert FP8 KV cache to BF16 workspace. + + Args: + src_cache: FP8 KV cache [num_blocks, block_size, 656] + dst: BF16 output workspace [total_tokens, 576] + block_table: Block indices [num_reqs, max_blocks] + seq_lens: Sequence lengths [num_reqs] + workspace_starts: Workspace start offsets [num_reqs] + batch_size: Number of requests + """ + torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache, dst, block_table, seq_lens, workspace_starts, batch_size + ) + + +def concat_mla_q( + ql_nope: torch.Tensor, + q_pe: torch.Tensor, + q_out: torch.Tensor, +) -> None: + """Concatenate query nope and rope for MLA/DSA attention. + + Args: + ql_nope: Query nope component [num_tokens, num_heads, nope_dim] + q_pe: Query rope component [num_tokens, num_heads, rope_dim] + q_out: Output tensor [num_tokens, num_heads, nope_dim + rope_dim] + """ + torch.ops._C_cache_ops.concat_mla_q(ql_nope, q_pe, q_out) + + +def indexer_k_quant_and_cache( + k: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + quant_block_size: int, + kv_cache_dtype: str, +) -> None: + torch.ops._C_cache_ops.indexer_k_quant_and_cache( + k, kv_cache, slot_mapping, quant_block_size, kv_cache_dtype + ) + + +def top_k_per_row_prefill( + logits: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + raw_topk_indices: torch.Tensor, + num_rows: int, + stride0: int, + stride1: int, + topk_tokens: int, +) -> None: + torch.ops._C.top_k_per_row_prefill( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + raw_topk_indices, + num_rows, + stride0, + stride1, + topk_tokens, + ) + + +def top_k_per_row_decode( + logits: torch.Tensor, + next_n: int, + seq_lens: torch.Tensor, + raw_topk_indices: torch.Tensor, + num_rows: int, + stride0: int, + stride1: int, + topk_tokens: int, +) -> None: + torch.ops._C.top_k_per_row_decode( + logits, + next_n, + seq_lens, + raw_topk_indices, + num_rows, + stride0, + stride1, + topk_tokens, + ) + + +def cp_gather_indexer_k_quant_cache( + kv_cache: torch.Tensor, + dst_k: torch.Tensor, + dst_scale: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.cp_gather_indexer_k_quant_cache( + kv_cache, dst_k, dst_scale, block_table, cu_seq_lens + ) + + +def get_device_attribute(attribute: int, device: int) -> int: + return torch.ops._C_cuda_utils.get_device_attribute(attribute, device) + + +def get_max_shared_memory_per_block_device_attribute(device: int) -> int: + # ruff: noqa: E501 + return torch.ops._C_cuda_utils.get_max_shared_memory_per_block_device_attribute( + device + ) + + +# custom ar +def init_custom_ar( + ipc_tensors: list[torch.Tensor], + rank_data: torch.Tensor, + rank: int, + fully_connected: bool, +) -> int: + return torch.ops._C_custom_ar.init_custom_ar( + ipc_tensors, rank_data, rank, fully_connected + ) + + +def all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + +def dispose(fa: int) -> None: + torch.ops._C_custom_ar.dispose(fa) + + +def meta_size() -> int: + return torch.ops._C_custom_ar.meta_size() + + +def register_buffer(fa: int, ipc_tensors: list[int]) -> None: + return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors) + + +def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]: + return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa) + + +def register_graph_buffers( + fa: int, handles: list[list[int]], offsets: list[list[int]] +) -> None: + torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets) + + +def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]: + return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size) + + +def open_mem_handle(mem_handle: torch.Tensor): + return torch.ops._C_custom_ar.open_mem_handle(mem_handle) + + +def free_shared_buffer(ptr: int) -> None: + torch.ops._C_custom_ar.free_shared_buffer(ptr) + + +# quick all reduce +def init_custom_qr(rank: int, world_size: int, qr_max_size: int | None = None) -> int: + return torch.ops._C_custom_ar.init_custom_qr(rank, world_size, qr_max_size) + + +def qr_destroy(fa: int) -> None: + torch.ops._C_custom_ar.qr_destroy(fa) + + +def qr_all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + quant_level: int, + cast_bf2half: bool = False, +) -> None: + torch.ops._C_custom_ar.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half) + + +def qr_get_handle(fa: int) -> torch.Tensor: + return torch.ops._C_custom_ar.qr_get_handle(fa) + + +def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None: + return torch.ops._C_custom_ar.qr_open_handles(fa, handles) + + +def qr_max_size() -> int: + return torch.ops._C_custom_ar.qr_max_size() + + +def get_flash_mla_metadata( + cache_seqlens: torch.Tensor, + num_heads_per_head_k: int, + num_heads_k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Arguments: + cache_seqlens: (batch_size), dtype torch.int32. + num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k. + num_heads_k: num_heads_k. + + Return: + tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. + num_splits: (batch_size + 1), dtype torch.int32. + """ + return torch.ops._C.get_flash_mla_metadata( + cache_seqlens, num_heads_per_head_k, num_heads_k + ) + + +def flash_mla_with_kvcache( + q: torch.Tensor, + k_cache: torch.Tensor, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + head_dim_v: int, + tile_scheduler_metadata: torch.Tensor, + num_splits: torch.Tensor, + softmax_scale: float | None = None, + causal: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Arguments: + q: (batch_size, seq_len_q, num_heads_q, head_dim). + k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). + block_table: (batch_size, max_num_blocks_per_seq), torch.int32. + cache_seqlens: (batch_size), torch.int32. + head_dim_v: Head_dim of v. + tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), torch.int32, return by get_mla_metadata. + num_splits: (batch_size + 1), torch.int32, return by get_mla_metadata. + softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim). + causal: bool. Whether to apply causal attention mask. + + Return: + out: (batch_size, seq_len_q, num_heads_q, head_dim_v). + softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. + """ + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) + out, softmax_lse = torch.ops._C.flash_mla_fwd_kvcache( + q, + k_cache, + None, + head_dim_v, + cache_seqlens, + block_table, + softmax_scale, + causal, + tile_scheduler_metadata, + num_splits, + ) + return out, softmax_lse + + +def sm100_cutlass_mla_decode( + out: torch.Tensor, + lse: torch.Tensor, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + workspace: torch.Tensor, + scale: float, + num_kv_splits: int, +) -> torch.Tensor: + torch.ops._C.sm100_cutlass_mla_decode( + out, + lse, + q_nope, + q_pe, + kv_c_and_k_pe_cache, + seq_lens, + page_table, + workspace, + scale, + num_kv_splits, + ) + return out + + +def sm100_cutlass_mla_get_workspace_size( + max_seq_len: int, num_batches: int, sm_count: int, num_kv_splits: int +) -> int: + return torch.ops._C.sm100_cutlass_mla_get_workspace_size( + max_seq_len, num_batches, sm_count, num_kv_splits + ) + + +def dsv3_fused_a_gemm( + output: torch.Tensor, + mat_a: torch.Tensor, + mat_b: torch.Tensor, +) -> None: + """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + + Computes output = mat_a @ mat_b.T where: + mat_a: [num_tokens, 7168] row-major bf16 (hidden states) + mat_b: [7168, 2112] column-major bf16 (weight transposed) + output: [num_tokens, 2112] row-major bf16 + + Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes. + Requires SM 9.0+ (Hopper). + """ + torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b) + + +if hasattr(torch.ops._C, "weight_packed_linear"): + + @register_fake("_C::weight_packed_linear") + def weight_packed_linear_fake( + mat1: torch.Tensor, + mat2: torch.Tensor, + bias: torch.Tensor | None, + is_vnni: bool, + ) -> torch.Tensor: + return torch.empty( + (mat1.size(0), mat2.size(0)), dtype=mat1.dtype, device=mat2.device + ) + + +class CPUQuantMethod(IntEnum): + UNQUANT = 0 + INT8_W8A8 = 1 + FP8_W8A16 = 2 + INT4_W4A8 = 3 + MXFP4 = 4 + + +if hasattr(torch.ops._C, "fused_experts_cpu"): + + @register_fake("_C::fused_experts_cpu") + def fused_experts_cpu_fake( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool, + moe_comp_method: CPUQuantMethod, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + w1_zero: torch.Tensor | None, + w2_zero: torch.Tensor | None, + block_size: list[int] | None, + w1_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + alpha: float | None, + limit: float | None, + is_vnni: bool, + ) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +def fused_experts_cpu( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool, + moe_comp_method: CPUQuantMethod, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + w1_zero: torch.Tensor | None, + w2_zero: torch.Tensor | None, + block_size: list[int] | None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + alpha: float | None = None, + limit: float | None = None, + is_vnni: bool = True, +) -> torch.Tensor: + return torch.ops._C.fused_experts_cpu( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + inplace, + moe_comp_method, + w1_scale, + w2_scale, + w1_zero, + w2_zero, + block_size, + w1_bias, + w2_bias, + alpha, + limit, + is_vnni, + ) + + +if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"): + + @register_fake("_C::int8_scaled_mm_with_quant") + def int8_scaled_mm_with_quant_fake( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, + ) -> torch.Tensor: + M = mat1.size(0) + N = mat2.size(0) + return torch.empty((M, N), dtype=out_dtype) + + +class CPUQuantAlgo(IntEnum): + AWQ = 0 + GPTQ = 1 + + +if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"): + + @register_fake("_C::convert_weight_packed_scale_zp") + def convert_weight_packed_scale_zp_fake( + qweight: torch.Tensor, + qzeros: torch.Tensor, + scales: torch.Tensor, + quant_method_4bit: CPUQuantAlgo, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + torch.empty_like(qweight), + torch.empty_like(qzeros), + torch.empty_like(scales), + ) + + +def convert_weight_packed_scale_zp( + qweight: torch.Tensor, + qzeros: torch.Tensor, + scales: torch.Tensor, + quant_method_4bit: CPUQuantAlgo, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops._C.convert_weight_packed_scale_zp( + qweight, + qzeros, + scales, + quant_method_4bit, + ) + + +if hasattr(torch.ops._C, "int4_scaled_mm_cpu"): + + @register_fake("_C::int4_scaled_mm_cpu") + def int4_scaled_mm_cpu_fake( + x: torch.Tensor, + w: torch.Tensor, + w_zeros: torch.Tensor, + w_scales: torch.Tensor, + bias: torch.Tensor | None, + ) -> torch.Tensor: + N = w_scales.size(0) * w_scales.size(-1) + return torch.empty((x.size(0), N), dtype=x.dtype, device=x.device) + + +def int4_scaled_mm_cpu( + x: torch.Tensor, + w: torch.Tensor, + w_zeros: torch.Tensor, + w_scales: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + x_shape = x.shape + x_2d = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x + + out = torch.ops._C.int4_scaled_mm_cpu( + x_2d, + w, + w_zeros, + w_scales, + bias, + ) + out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out + return out + + +if hasattr(torch.ops._C, "fp8_scaled_mm_cpu"): + + @register_fake("_C::fp8_scaled_mm_cpu") + def fp8_scaled_mm_cpu_fake( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + block_size: list[int], + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, + ) -> torch.Tensor: + M = mat1.size(0) + N = mat2.size(0) + return torch.empty((M, N), dtype=out_dtype, device=mat1.device) + + +_supports_cpu_fp8_w8a16 = bool(hasattr(torch.ops._C, "fp8_scaled_mm_cpu")) + + +def fp8_scaled_mm_cpu( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + block_size: list[int], + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, +) -> torch.Tensor: + return torch.ops._C.fp8_scaled_mm_cpu( + mat1, mat2, scales2, block_size, bias, out_dtype, is_vnni + ) + + +def chunk_gated_delta_rule_cpu( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor, + head_first: bool, + use_qk_l2norm_in_kernel: bool, + eps: float = 1e-5, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.chunk_gated_delta_rule_cpu( + query, + key, + value, + g, + beta, + initial_state, + output_final_state, + cu_seqlens, + head_first, + use_qk_l2norm_in_kernel, + eps, + ) + + +def fused_sigmoid_gating_delta_rule_update_cpu( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + use_qk_l2norm_in_kernel: bool, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, +) -> torch.Tensor: + return torch.ops._C.fused_sigmoid_gating_delta_rule_update_cpu( + A_log, + dt_bias, + q, + k, + v, + a, + b, + initial_state_source, + initial_state_indices, + cu_seqlens, + use_qk_l2norm_in_kernel, + softplus_beta, + softplus_threshold, + ) + + +def fused_gdn_gating_cpu( + A_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.fused_gdn_gating_cpu( + A_log, + a, + b, + dt_bias, + ) + + +def causal_conv1d_weight_pack( + weight: torch.Tensor, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_weight_pack( + weight, + ) + + +def causal_conv1d_fwd_cpu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor | None, + query_start_loc: torch.Tensor | None, + cache_indices: torch.Tensor | None, + has_initial_state: torch.Tensor | None, + silu_activation: bool, + is_vnni: bool, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_fwd_cpu( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + silu_activation, + -1, + is_vnni, + ) + + +def causal_conv1d_update_cpu( + x: torch.Tensor, + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + conv_state_indices: torch.Tensor | None, + is_vnni: bool, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_update_cpu( + x, + conv_states, + weight, + bias, + silu_activation, + None, + conv_state_indices, + -1, + is_vnni, + ) + + +class CPUDNNLGEMMHandler: + def __init__(self) -> None: + self.handler_tensor: torch.Tensor | None = None + self.n = -1 + self.k = -1 + self.dtor = torch.ops._C.release_dnnl_matmul_handler + + def __del__(self): + if self.handler_tensor is not None: + self.dtor(self.handler_tensor.item()) + + +_supports_onednn = bool(hasattr(torch.ops._C, "create_onednn_mm_handler")) + + +def is_onednn_acl_supported(): + return torch.ops._C.is_onednn_acl_supported() + + +def create_onednn_mm( + weight: torch.Tensor, # [K, N] + primitive_cache_size: int = 128, +) -> CPUDNNLGEMMHandler: + handler = CPUDNNLGEMMHandler() + handler.k, handler.n = weight.size() + # store the handler pointer in a tensor it doesn't get inlined + handler.handler_tensor = torch.tensor( + torch.ops._C.create_onednn_mm_handler(weight, primitive_cache_size), + dtype=torch.int64, + ) + return handler + + +def onednn_mm( + dnnl_handler: CPUDNNLGEMMHandler, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype) + torch.ops._C.onednn_mm( + output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler_tensor + ) + + return output + + +def create_onednn_scaled_mm( + weight: torch.Tensor, # [K, N] + weight_scales: torch.Tensor, + output_type: torch.dtype, + dynamic_quant: bool, + use_azp: bool, + primitive_cache_size: int = 128, +) -> CPUDNNLGEMMHandler: + handler = CPUDNNLGEMMHandler() + handler.k, handler.n = weight.size() + # store the handler pointer in a tensor so it doesn't get inlined + handler.handler_tensor = torch.tensor( + torch.ops._C.create_onednn_scaled_mm_handler( + weight, + weight_scales, + output_type, + dynamic_quant, + use_azp, + primitive_cache_size, + ), + dtype=torch.int64, + ) + return handler + + +def onednn_scaled_int8_quant( + input: torch.Tensor, + scale: torch.Tensor | None = None, + azp: torch.Tensor | None = None, + symmetric: bool = True, +): + """ + Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. + + Args: + input: The input tensor to be quantized to int8. + scale: Optional scaling factor for the int8 quantization. + When not provided, we invoke dynamic-per-token quantization. + azp: Optional zero-point for the int8 quantization. + Must be provided for asymmetric quantization if `scale` is provided. + symmetric: Whether to use symmetric quantization (scale only, azp ignored). + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp. + """ + output = torch.empty_like(input, dtype=torch.int8) + token_num = input.numel() // input.shape[-1] + input = input.view((token_num, input.shape[-1])) + if scale is not None: + # static-per-tensor quantization. + assert symmetric == (azp is None), ( + "azp must only be provided for asymmetric quantization." + ) + torch.ops._C.static_scaled_int8_quant(output, input, scale, azp) + return output, scale, azp + + # dynamic-per-token quantization. + input_scales = torch.empty((token_num, 1), device=input.device, dtype=torch.float32) + input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32) + torch.ops._C.dynamic_scaled_int8_quant(output, input, input_scales, input_azp) + return output, input_scales, input_azp + + +def onednn_scaled_mm( + dnnl_handler: CPUDNNLGEMMHandler, + x: torch.Tensor, + output: torch.Tensor, + input_scale: torch.Tensor | None, + input_zp: torch.Tensor | None, + input_zp_adj: torch.Tensor | None, + bias: torch.Tensor | None, +) -> torch.Tensor: + torch.ops._C.onednn_scaled_mm( + output, + x, + input_scale, + input_zp, + input_zp_adj, + bias, + dnnl_handler.handler_tensor, + ) + + return output + + +def cpu_attn_get_scheduler_metadata( + num_reqs: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + seq_lens: torch.Tensor, + dtype: torch.dtype, + query_start_loc: torch.Tensor, + causal: bool, + sliding_window_size: int, + isa: str, + enable_kv_split: bool, +) -> torch.Tensor: + scheduler_metadata = torch.ops._C.get_scheduler_metadata( + num_reqs, + num_heads, + num_kv_heads, + head_dim, + seq_lens, + dtype, + query_start_loc, + causal, + sliding_window_size, + isa, + enable_kv_split, + ) + return scheduler_metadata + + +def cpu_attn_reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + isa: str, + k_scale: float = 1.0, + v_scale: float = 1.0, + kv_cache_dtype: str = "auto", +) -> None: + torch.ops._C.cpu_attn_reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + isa, + k_scale, + v_scale, + kv_cache_dtype, + ) + + +def cpu_attention_with_kv_cache( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + output: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + scale: float, + causal: bool, + alibi_slopes: torch.Tensor | None, + sliding_window: tuple[int, int], + block_table: torch.Tensor, + softcap: float, + scheduler_metadata: torch.Tensor, + s_aux: torch.Tensor | None, + k_scale: float = 1.0, + v_scale: float = 1.0, + kv_cache_dtype: str = "auto", +) -> None: + torch.ops._C.cpu_attention_with_kv_cache( + query, + key_cache, + value_cache, + output, + query_start_loc, + seq_lens, + scale, + causal, + alibi_slopes, + sliding_window[0], + sliding_window[1], + block_table, + softcap, + scheduler_metadata, + s_aux, + k_scale, + v_scale, + kv_cache_dtype, + ) + + +def cpu_gemm_wna16( + input: torch.Tensor, + q_weight: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + bias: torch.Tensor | None, + pack_factor: int, + isa_hint: str, +) -> torch.Tensor: + output = torch.empty((input.size(0), scales.size(1)), dtype=input.dtype) + torch.ops._C.cpu_gemm_wna16( + input, + q_weight, + output, + scales, + zeros, + g_idx, + bias, + pack_factor, + isa_hint, + ) + return output + + +def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor: + out = torch.empty_like(input) + torch.ops._C.activation_lut_bf16(out, input, activation) + return out + + +def cpu_prepack_moe_weight( + weight: torch.Tensor, + isa: str, +) -> torch.Tensor: + output = torch.empty_like(weight) + torch.ops._C.prepack_moe_weight(weight, output, isa) + return output + + +def cpu_fused_moe( + input: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + act: str, + isa: str, + skip_weighted: bool = False, +) -> torch.Tensor: + output = torch.empty_like(input) + torch.ops._C.cpu_fused_moe( + output, + input, + w13, + w2, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + skip_weighted, + act, + isa, + ) + return output + + +if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"): + + @register_fake("_qutlass_C::matmul_mxf4_bf16_tn") + def _fake_matmul_mxf4_bf16_tn( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, + ): + return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16) + + +def matmul_mxf4_bf16_tn( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops._qutlass_C.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"): + + @register_fake("_qutlass_C::fusedQuantizeMxQuest") + def _fake_fused_quantize_mx_quest( + a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor + ): + return xh_e2m1, xh_e8m0 + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxAbsMax"): + + @register_fake("_qutlass_C::fusedQuantizeMxAbsMax") + def _fake_fused_quantize_mx_absmax( + a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor + ): + return xh_e2m1, xh_e8m0 + + +def fusedQuantizeMx( + a: torch.Tensor, b: torch.Tensor, *, method: Literal["quest", "abs_max"] = "quest" +) -> tuple[torch.Tensor, torch.Tensor]: + if a.dim() == 0: + raise ValueError("`a` must have at least 1 dimension.") + if a.size(-1) % 32 != 0: + raise ValueError(f"last dim of `a` must be divisible by 32, got {a.size(-1)}.") + if b.device != a.device: + raise ValueError("`a` and `b` must be on the same device.") + + xh_e2m1 = torch.empty( + *a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device + ) + + rows, cols = a.numel() // a.size(-1), a.size(-1) // 32 + n_row_blocks = cdiv(rows, 128) + n_col_blocks = cdiv(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + xh_e8m0 = torch.empty( + padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=a.device + ) + + if not hasattr(torch.ops, "_qutlass_C"): + raise RuntimeError( + "The `_qutlass_C` extension is not loaded. " + "Make sure your custom op library is imported before calling fusedQuantizeMx." + ) + + if method == "quest": + return torch.ops._qutlass_C.fusedQuantizeMxQuest(a, b, xh_e2m1, xh_e8m0) + elif method == "abs_max": + return torch.ops._qutlass_C.fusedQuantizeMxAbsMax(a, b, xh_e2m1, xh_e8m0) + else: + raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'") + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"): + + @register_fake("_qutlass_C::fusedQuantizeNv") + def _fake_fused_quantize_nv( + a: torch.Tensor, + b: torch.Tensor, + xh_e2m1: torch.Tensor, + xh_e4m3: torch.Tensor, + global_scale: torch.Tensor, + ): + return xh_e2m1, xh_e4m3 + + +def fusedQuantizeNv( + a: torch.Tensor, b: torch.Tensor, global_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + xh_e2m1 = torch.empty( + *a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device + ) + + rows, cols = a.numel() // a.size(-1), a.size(-1) // 16 + n_row_blocks = cdiv(rows, 128) + n_col_blocks = cdiv(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + xh_e4m3 = torch.empty( + padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device + ) + + return torch.ops._qutlass_C.fusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale) + + +def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor: + """ + Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832) + kernels. Note that these kernels exploit the recursive properties of + Sylvester Hadamards, and therefore do not require transform weight data + + Note that sylvester hadamard transforms are also symmetric, which means that + this function is also applies the (transpose <=> inverse) transform. + + Args: + x: value to be transformed inplace + inplace: modify value in place + + Returns: + value after transformation + """ + return torch.ops._C.hadacore_transform(x, inplace) + + +if hasattr(torch.ops._C, "hadacore_transform"): + + @register_fake("_C::hadacore_transform") + def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor: + return torch.empty_like(x) if not inplace else x + + +if hasattr(torch.ops._C, "minimax_allreduce_rms"): + + @register_fake("_C::minimax_allreduce_rms") + def _minimax_allreduce_rms_fake( + input: torch.Tensor, + norm_weight: torch.Tensor, + workspace: torch.Tensor, + rank: int, + nranks: int, + eps: float, + ) -> torch.Tensor: + return torch.empty_like(input) + + +if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): + + @register_fake("_C::minimax_allreduce_rms_qk") + def _minimax_allreduce_rms_qk_fake( + qkv: torch.Tensor, + norm_weight_q: torch.Tensor, + norm_weight_k: torch.Tensor, + workspace: torch.Tensor, + q_size: int, + kv_size: int, + rank: int, + nranks: int, + eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + token_num = qkv.shape[0] + return ( + torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device), + torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device), + ) diff --git a/upstream_ref/ds_vllm/vllm/qwen3_5.py b/upstream_ref/ds_vllm/vllm/qwen3_5.py new file mode 100644 index 00000000..43b90046 --- /dev/null +++ b/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() diff --git a/upstream_ref/xllm/kernels/cuda/activation.cu b/upstream_ref/xllm/kernels/cuda/activation.cu new file mode 100644 index 00000000..ad5a8ee0 --- /dev/null +++ b/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 +#include + +#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 +__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(ptr) & 15) == 0; +} + +// Activation and gating kernel template with 128-bit vectorized access +// optimization. +template +__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(x_ptr); + const int4* y_vec = reinterpret_cast(y_ptr); + int4* out_vec = reinterpret_cast(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(&x); + auto* yp = reinterpret_cast(&y); + auto* rp = reinterpret_cast(&r); +#pragma unroll + for (int j = 0; j < VEC_SIZE; j++) { + rp[j] = compute(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(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(x, y); + } + } +} + +template +__device__ __forceinline__ T silu_kernel(const T& x) { + // x * sigmoid(x) + return (T)(((float)x) / (1.0f + expf((float)-x))); +} + +template +__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 +__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, ACT_FIRST> \ + <<>>( \ + out.data_ptr(), input.data_ptr(), 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 diff --git a/upstream_ref/xllm/kernels/cuda/device_utils.cuh b/upstream_ref/xllm/kernels/cuda/device_utils.cuh new file mode 100644 index 00000000..e44db294 --- /dev/null +++ b/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 + +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 +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 +__device__ float convert_to_float(T x) { + if constexpr (std::is_same_v) { + return __half2float(x); + } else if constexpr (std::is_same_v) { + return __bfloat162float(x); + } else if constexpr (std::is_same_v) { + return x; + } else { + return static_cast(x); + } +} + +// Constructs some constants needed to partition the work across threads at +// compile time. +template +struct TopkConstants { + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); + static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || + EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, + ""); + static constexpr int VECs_PER_THREAD = + 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 \ No newline at end of file diff --git a/upstream_ref/xllm/kernels/cuda/moe/fused_moe.cpp b/upstream_ref/xllm/kernels/cuda/moe/fused_moe.cpp new file mode 100644 index 00000000..735e27eb --- /dev/null +++ b/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& 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& fc1_expert_biases, + const std::optional& fc2_expert_biases, + const std::optional& input_sf, + const std::optional& swiglu_alpha, + const std::optional& swiglu_beta, + const std::optional& swiglu_limit, + const std::optional& 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 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(); + + 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>(), // TODO: support + // auto tuning + // profile ids + support_pdl(), + activation_type); + + return result_output; +} +} // namespace xllm::kernel::cuda diff --git a/upstream_ref/xllm/kernels/cuda/moe/moe_fused_topk.cu b/upstream_ref/xllm/kernels/cuda/moe/moe_fused_topk.cu new file mode 100644 index 00000000..26f2a475 --- /dev/null +++ b/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 moe_fused_topk( + torch::Tensor& gating_output, + int64_t topk, + bool renormalize, + const std::optional& 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 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 diff --git a/upstream_ref/xllm/kernels/cuda/moe/moe_topk.cuh b/upstream_ref/xllm/kernels/cuda/moe/moe_topk.cuh new file mode 100644 index 00000000..8d66bb21 --- /dev/null +++ b/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 +#include + +#include + +#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 +struct TopKRedType { + using T = T_; + static_assert( + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v, + "Top K reduction only implemented for int, float, float16 and bfloat16"); + + using TypeCmp = std::conditional_t; + using IdxT = std::conditional_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::TwiddleIn( + reinterpret_cast::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((cmp & 0xFFFF)); + + auto compactTmp = cmp >> kMoveBits; + auto valueBits = cub::Traits::TwiddleOut( + reinterpret_cast::UnsignedBits&>(compactTmp)); + value = reinterpret_cast(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 const& warp) { + if constexpr (!kTLLM_GEN_HAS_FAST_REDUX || sizeof(TypeCmp) == 8) { + return cg::reduce(warp, compValIdx, cg::greater{}); + } else { + TypeCmp result; + asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(result) + : "r"(compValIdx)); + return result; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct TopKIdx { + // by default, empty +}; + +template +struct TopKIdx { + 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 +struct Sort; + +template +struct Sort<1, RedType> { + static __device__ void run(RedType* topK) {} +}; + +template +struct Sort<2, RedType> { + static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } +}; + +template +struct Sort<3, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 1); + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 1); + } +}; + +template +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 +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile 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; + 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 +__device__ void reduceTopKFunc(cg::thread_block_tile 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; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } + + if constexpr (!IsSorted) { + Sort::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 +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile 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; + + if constexpr (N <= 4) { + reduceTopKFunc( + 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( + 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( + warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK); + } +}; + +#undef TOPK_SWAP + +} // namespace reduce_topk +} // namespace xllm::kernel::cuda diff --git a/upstream_ref/xllm/kernels/cuda/moe/moe_topk_sigmoid_kernels.cuh b/upstream_ref/xllm/kernels/cuda/moe/moe_topk_sigmoid_kernels.cuh new file mode 100644 index 00000000..a8de51c2 --- /dev/null +++ b/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 +#include +#include + +#include +#include + +#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 +__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(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 +__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; + using BlockReduce = cub::BlockReduce; + __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 +__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; + + // Finally, we pull in the data from global mem + T row_chunk_temp[VPT]; + AccessType* row_chunk_vec_ptr = + reinterpret_cast(&row_chunk_temp); + const AccessType* vec_thread_read_ptr = + reinterpret_cast(thread_read_ptr); +#pragma unroll + // 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(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 +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; + 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 + <<>>(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( \ + gating_output, \ + nullptr, \ + topk_weights, \ + topk_indices, \ + num_tokens, \ + topk, \ + 0, \ + num_experts, \ + renormalize, \ + correction_bias, \ + stream); + +template +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<<>>(gating_output, + nullptr, + sigmoid_workspace, + num_experts, + correction_bias); + moe_topK<<>>(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& 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(gating_output.size(-1)); + const int num_tokens = static_cast(gating_output.size(0)); + const int topk = static_cast(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(); + } + + if (dtype == at::ScalarType::Float) { + topk_gating_sigmoid_kernel_launcher( + gating_output.data_ptr(), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else if (dtype == at::ScalarType::Half) { + topk_gating_sigmoid_kernel_launcher<__half>( + reinterpret_cast(gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else if (dtype == at::ScalarType::BFloat16) { + topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>( + reinterpret_cast( + gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else { + LOG(FATAL) << "Unsupported gating_output dtype: " << dtype; + } +} +} // namespace xllm::kernel::cuda diff --git a/upstream_ref/xllm/kernels/cuda/moe/moe_topk_softmax_kernels.cuh b/upstream_ref/xllm/kernels/cuda/moe/moe_topk_softmax_kernels.cuh new file mode 100644 index 00000000..ea74ad18 --- /dev/null +++ b/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 +#include +#include + +#include +#include + +#include "kernels/cuda/device_utils.cuh" + +using cub_kvp = cub::KeyValuePair; + +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 +__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; + __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(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 +__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; + __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 +__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; + using BlockReduce = cub::BlockReduce; + __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 +__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; + + // Finally, we pull in the data from global mem + T row_chunk_temp[VPT]; + AccessType* row_chunk_vec_ptr = + reinterpret_cast(&row_chunk_temp); + const AccessType* vec_thread_read_ptr = + reinterpret_cast(thread_read_ptr); +#pragma unroll + // 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(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 +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; + 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 + <<>>(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( \ + gating_output, \ + nullptr, \ + topk_weights, \ + topk_indices, \ + num_tokens, \ + topk, \ + 0, \ + num_experts, \ + renormalize, \ + moe_softcapping, \ + correction_bias, \ + stream); + +template +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<<>>(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<<>>(softmax_workspace, + nullptr, + topk_weights, + topk_indices, + num_experts, + topk, + 0, + num_experts, + renormalize); + } else { + moe_topk_fast<<>>(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& 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(gating_output.size(-1)); + const int num_tokens = static_cast(gating_output.size(0)); + const int topk = static_cast(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(); + } + + // Cast moe_softcapping from double to float for CUDA kernels + const float moe_softcapping_f = static_cast(moe_softcapping); + + if (dtype == at::ScalarType::Float) { + topk_gating_softmax_kernel_launcher( + gating_output.data_ptr(), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + softmax_workspace.data_ptr(), + 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(gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + softmax_workspace.data_ptr(), + 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( + gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + softmax_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + moe_softcapping_f, + bias_ptr, + stream); + } else { + LOG(FATAL) << "Unsupported gating_output dtype: " << dtype; + } +} +} // namespace xllm::kernel::cuda diff --git a/upstream_ref/xllm/kernels/cuda/norm.cu b/upstream_ref/xllm/kernels/cuda/norm.cu new file mode 100644 index 00000000..90de7e53 --- /dev/null +++ b/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 +#include + +#include + +#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 +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 +__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; + __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 +__global__ std::enable_if_t<(width > 0) && _typeConvert::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>); + static_assert(sizeof(_f16Vec) == 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*>(input); + auto* __restrict__ residual_v = + reinterpret_cast<_f16Vec*>(residual); + auto* __restrict__ weight_v = + reinterpret_cast*>(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 temp = input_v[strided_id]; + temp += residual_v[id]; + variance += temp.sum_squares(); + residual_v[id] = temp; + } + + using BlockReduce = cub::BlockReduce; + __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 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 +__global__ std::enable_if_t<(width == 0) || !_typeConvert::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; + __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 \ + <<>>(input.data_ptr(), \ + input_stride, \ + residual.data_ptr(), \ + weight.data_ptr(), \ + 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 +__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(input_row[idx]); + variance += x * x; + } + + using BlockReduce = cub::BlockReduce; + __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(input_row[idx]); + float out_norm = (static_cast(x * s_variance)) * + static_cast(weight[idx]); + out[blockIdx.x * hidden_size + idx] = + xllm::kernel::cuda::scaled_fp8_conversion(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 +__global__ std::enable_if_t<(width > 0) && _typeConvert::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>); + static_assert(sizeof(_f16Vec) == 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*>(input); + auto* __restrict__ residual_v = + reinterpret_cast<_f16Vec*>(residual); + auto* __restrict__ weight_v = + reinterpret_cast*>(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 temp = input_v[strided_id]; + temp += residual_v[id]; + variance += temp.sum_squares(); + residual_v[id] = temp; // Store updated residual + } + + using BlockReduce = cub::BlockReduce; + __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 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::convert(temp.data[i]); + out[id * width + i] = + xllm::kernel::cuda::scaled_fp8_conversion(val, + scale_inv); + } + } +} + +/** + * Generic fused add + RMSNorm + FP8 quant kernel (fallback for unaligned data) + */ +template +__global__ std::enable_if_t<(width == 0) || !_typeConvert::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(z); + variance += x * x; + residual[blockIdx.x * hidden_size + idx] = z; // Store updated residual + } + + using BlockReduce = cub::BlockReduce; + __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(residual[blockIdx.x * hidden_size + idx]); + float out_norm = (static_cast(x * s_variance)) * + static_cast(weight[idx]); + out[blockIdx.x * hidden_size + idx] = + xllm::kernel::cuda::scaled_fp8_conversion(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 \ + <<>>(out.data_ptr(), \ + input.data_ptr(), \ + input_stride, \ + residual.data_ptr(), \ + weight.data_ptr(), \ + scale.data_ptr(), \ + 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 + <<>>(output.data_ptr(), + input.data_ptr(), + input_stride, + weight.data_ptr(), + 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(input.data_ptr()); + auto res_ptr = reinterpret_cast(residual.data_ptr()); + auto wt_ptr = reinterpret_cast(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 + <<>>(out.data_ptr(), + input.data_ptr(), + input_stride, + weight.data_ptr(), + scale.data_ptr(), + 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(input.data_ptr()); + auto res_ptr = reinterpret_cast(residual.data_ptr()); + auto wt_ptr = reinterpret_cast(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 diff --git a/upstream_ref/xllm/kernels/cuda/rope.cu b/upstream_ref/xllm/kernels/cuda/rope.cu new file mode 100644 index 00000000..5860d990 --- /dev/null +++ b/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 +#include +#include + +#include "cuda_ops_api.h" + +// ref to: +// https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu + +namespace { + +template +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 +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( + 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( + key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim); + } + } +} + +template +__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(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 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(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<<>>( + positions.data_ptr(), + query.data_ptr(), + key.has_value() ? key->data_ptr() : nullptr, + cos_sin_cache.data_ptr(), + rot_dim, + query_stride, + key_stride, + head_stride, + num_heads, + num_kv_heads, + head_size); + } else { + rotary_embedding_kernel<<>>( + positions.data_ptr(), + query.data_ptr(), + key.has_value() ? key->data_ptr() : nullptr, + cos_sin_cache.data_ptr(), + rot_dim, + query_stride, + key_stride, + head_stride, + num_heads, + num_kv_heads, + head_size); + } + }); +} + +} // namespace xllm::kernel::cuda diff --git a/upstream_ref/xllm/kernels/ilu/activation.cpp b/upstream_ref/xllm/kernels/ilu/activation.cpp new file mode 100644 index 00000000..ae2a16ba --- /dev/null +++ b/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 diff --git a/upstream_ref/xllm/kernels/ilu/attention.cpp b/upstream_ref/xllm/kernels/ilu/attention.cpp new file mode 100644 index 00000000..aa257bf1 --- /dev/null +++ b/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& value, + torch::Tensor& key_cache, + std::optional& 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& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& 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(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& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& 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 \ No newline at end of file diff --git a/upstream_ref/xllm/kernels/ilu/fused_moe.cpp b/upstream_ref/xllm/kernels/ilu/fused_moe.cpp new file mode 100644 index 00000000..794f9bd9 --- /dev/null +++ b/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 + +#include "ilu_ops_api.h" + +namespace xllm::kernel::ilu { + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& 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 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 diff --git a/upstream_ref/xllm/kernels/ilu/ilu_ops_api.h b/upstream_ref/xllm/kernels/ilu/ilu_ops_api.h new file mode 100644 index 00000000..e4fd7853 --- /dev/null +++ b/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 +#include +#include +#include +#include + +#include + +#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& value, // (num_tokens, num_heads, head_size) + torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size) + std::optional& + 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& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& 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& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& 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& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& 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 bias); + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias); + +std::vector 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& dst_to_src, + torch::Tensor& output); + +torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight); +} // namespace xllm::kernel::ilu diff --git a/upstream_ref/xllm/kernels/ilu/ixformer.h b/upstream_ref/xllm/kernels/ilu/ixformer.h new file mode 100644 index 00000000..57ce66dc --- /dev/null +++ b/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 + +#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& alibi_slopes, + const std::optional& sinks, + std::optional& 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& 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& sinks); + +torch::Tensor ixformer_linear(torch::Tensor& input, + torch::Tensor& weight, + int64_t act_type, + const std::optional& bias, + const std::optional& out, + const std::optional persistent); + +torch::Tensor ixformer_linear_ex(torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& 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& fused_bias, + double alpha, + double eps, + bool is_post); + +void rms_norm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& output, + const std::optional& 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& expert_mask, + const c10::optional& expert_sizes_cpu, + const c10::optional& 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& 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& dst_to_src, + const c10::optional& 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& mul_weight, + const c10::optional& mask, + const c10::optional& extra_residual, + double scaling_factor); +} // namespace ixformer::infer diff --git a/upstream_ref/xllm/kernels/ilu/norm.cpp b/upstream_ref/xllm/kernels/ilu/norm.cpp new file mode 100644 index 00000000..c5a98595 --- /dev/null +++ b/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& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& 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 fused_bias = std::nullopt; + infer::rms_norm(input, weight, output, fused_bias, eps); +} + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/upstream_ref/xllm/kernels/ilu/rope.cpp b/upstream_ref/xllm/kernels/ilu/rope.cpp new file mode 100644 index 00000000..89370b79 --- /dev/null +++ b/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 diff --git a/upstream_ref/xllm/kernels/ilu/utils.h b/upstream_ref/xllm/kernels/ilu/utils.h new file mode 100644 index 00000000..e8af0c3c --- /dev/null +++ b/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 \ No newline at end of file diff --git a/upstream_ref/xllm/layers/ilu/attention.cpp b/upstream_ref/xllm/layers/ilu/attention.cpp new file mode 100644 index 00000000..b66f28a4 --- /dev/null +++ b/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> AttentionImpl::forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache) { + std::optional 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 v_cache; + std::optional 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& v_cache, + const AttentionMetadata& attn_metadata) { + int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_; + std::optional 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& 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 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 diff --git a/upstream_ref/xllm/layers/ilu/attention.h b/upstream_ref/xllm/layers/ilu/attention.h new file mode 100644 index 00000000..a971835f --- /dev/null +++ b/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 + +#include + +#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> 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& v_cache, + const AttentionMetadata& attn_metadata); + + void decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& 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 diff --git a/upstream_ref/xllm/layers/ilu/fused_moe.cpp b/upstream_ref/xllm/layers/ilu/fused_moe.cpp new file mode 100644 index 00000000..4238012e --- /dev/null +++ b/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 + +#include + +#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(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(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(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 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 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 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 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 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 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 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 diff --git a/upstream_ref/xllm/layers/ilu/fused_moe.h b/upstream_ref/xllm/layers/ilu/fused_moe.h new file mode 100644 index 00000000..3e477064 --- /dev/null +++ b/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 + +#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 cusum_token_count; + std::optional 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 shared_stream_; + std::unique_ptr 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 diff --git a/upstream_ref/xllm/layers/npu_torch/qwen3_5_gated_delta_net.cpp b/upstream_ref/xllm/layers/npu_torch/qwen3_5_gated_delta_net.cpp new file mode 100644 index 00000000..7d572476 --- /dev/null +++ b/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 + +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 +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 diff --git a/upstream_ref/xllm/layers/npu_torch/qwen3_5_gated_delta_net.h b/upstream_ref/xllm/layers/npu_torch/qwen3_5_gated_delta_net.h new file mode 100644 index 00000000..bec6c1c6 --- /dev/null +++ b/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 + +#include +#include + +#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 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 diff --git a/upstream_ref/xllm/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/upstream_ref/xllm/layers/npu_torch/qwen3_gated_delta_net_base.cpp new file mode 100644 index 00000000..cec9a95e --- /dev/null +++ b/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 +#include + +#include + +#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_recurrent_gated_delta_rule( + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor g, + torch::Tensor beta, + std::optional 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(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_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 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(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 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 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 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 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(num_k_heads_ / tp_size_); + fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); + fused_params.head_qk = static_cast(head_k_dim_); + fused_params.head_v = static_cast(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{1}, + /*padding=*/std::vector{3}, + /*dilation=*/std::vector{1}, + /*groups=*/static_cast(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(); + conv1d_params.initial_state_idx = std::optional(); + 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(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 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(); + 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 batches; + int64_t idx = 0; + for (int64_t b = 0; b < bs; ++b) { + int64_t cur_len = start_loc[b].template item(); + 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 +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 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 diff --git a/upstream_ref/xllm/layers/npu_torch/qwen3_gated_delta_net_base.h b/upstream_ref/xllm/layers/npu_torch/qwen3_gated_delta_net_base.h new file mode 100644 index 00000000..2994f329 --- /dev/null +++ b/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 + +#include +#include +#include + +#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 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 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 diff --git a/upstream_ref/xllm/models/qwen3_5.h b/upstream_ref/xllm/models/qwen3_5.h new file mode 100644 index 00000000..7c712d39 --- /dev/null +++ b/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 +#include +#include +#include + +#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(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(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()); \ + 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()); \ + 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()); \ + 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({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