dylanyunlong/project_8
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[init] baseline7 from project_6

Browse Source
This commit is contained in:
root
2026-08-25 07:05:49 +00:00
commit 0bb337bc00
1392 changed files with 321187 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
ex_engine/xllm_kernels/build_test_cutlass_batched.sh Executable file
View File

@@ -0,0 +1,25 @@
#!/bin/bash
# build_test_cutlass_batched.sh — Compile and test Cu10 TensorOp batched GEMM
set -eo pipefail
SAMPLES="/usr/local/corex-samples-3.2.3_x86_64/samples/cutlass"
SRC="ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu"
echo "=== Compile Cu10 TensorOp batched HGEMM ==="
/usr/local/corex/bin/clang++ \
--cuda-gpu-arch=ivcore10 --cuda-path=/usr/local/corex \
-I"${SAMPLES}/include" \
-I/usr/local/corex/include \
-L/usr/local/corex/lib64 -lcudart -lcutlass \
-DBUILD_STANDALONE_TEST \
-O2 -std=c++17 \
"$SRC" -o /tmp/test_cutlass_batched 2>&1
if [ -f /tmp/test_cutlass_batched ]; then
echo "Compile: SUCCESS"
echo ""
echo "=== Run ==="
/tmp/test_cutlass_batched
else
echo "Compile: FAILED"
fi

156
ex_engine/xllm_kernels/build_test_hgemm.sh Executable file
View File

@@ -0,0 +1,156 @@
#!/bin/bash
# build_test_hgemm.sh — Compile and test hgemm_blocktiling on BI-V100
#
# Usage: bash ex_engine/xllm_kernels/build_test_hgemm.sh
set -eo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CUDA_DIR="${SCRIPT_DIR}/cuda"
echo "=== 1. Compile hgemm_blocktiling ==="
python3 -c "
import torch.utils.cpp_extension as ext
import os, shutil, glob
name = 'hgemm_blocktiling'
build_dir = '${SCRIPT_DIR}/build/tmp_' + name
os.makedirs(build_dir, exist_ok=True)
try:
mod = ext.load(
name=name,
sources=[
'${CUDA_DIR}/hgemm_blocktiling.cu',
'${CUDA_DIR}/bindings/hgemm_bind.cpp',
],
extra_include_paths=['${CUDA_DIR}/headers'],
extra_cflags=['-O2', '-std=c++17'],
extra_cuda_cflags=['-O2'],
build_directory=build_dir,
verbose=True,
)
built = glob.glob(build_dir + '/' + name + '*.so')
if built:
dst = '${SCRIPT_DIR}/build/' + name + '.so'
shutil.copy2(built[0], dst)
print(f'[build] SUCCESS: {dst} ({os.path.getsize(dst)} bytes)')
else:
print('[build] WARNING: .so not found')
except Exception as e:
print(f'[build] FAILED: {e}')
import traceback
traceback.print_exc()
"
echo ""
echo "=== 2. Functional test ==="
python3 << 'PYTEST'
import torch
import sys, os, glob
# Find and load the .so
build_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)) if '__file__' in dir() else '.',
'ex_engine/xllm_kernels/build')
sys.path.insert(0, build_dir)
try:
import hgemm_blocktiling as hg
print("Module loaded successfully")
except ImportError:
# Try loading from tmp build dir
import importlib.util
so_files = glob.glob('ex_engine/xllm_kernels/build/tmp_hgemm_blocktiling/hgemm_blocktiling*.so')
if not so_files:
print("SKIP: .so not found (need GPU machine)")
sys.exit(0)
spec = importlib.util.spec_from_file_location("hgemm_blocktiling", so_files[0])
hg = importlib.util.module_from_spec(spec)
spec.loader.exec_module(hg)
print(f"Module loaded from {so_files[0]}")
# Test 1: Small GEMM correctness
print("\n--- Test 1: Small GEMM (64x64 @ 64x64) ---")
M, N, K = 64, 64, 64
A = torch.randn(M, K, dtype=torch.float16, device='cuda')
B = torch.randn(K, N, dtype=torch.float16, device='cuda')
C_ref = torch.matmul(A.float(), B.float()).half()
C_our = hg.hgemm(A, B)
diff = (C_ref.float() - C_our.float()).abs().max().item()
print(f" Max abs diff: {diff:.6f}")
assert diff < 1.0, f"FAILED: diff={diff} too large"
print(f" PASS (diff < 1.0)")
# Test 2: Larger GEMM (typical MoE dimensions)
print("\n--- Test 2: MoE-sized GEMM (256x4096 @ 4096x11008) ---")
M, N, K = 256, 11008, 4096
A = torch.randn(M, K, dtype=torch.float16, device='cuda') * 0.01
B = torch.randn(K, N, dtype=torch.float16, device='cuda') * 0.01
C_ref = torch.matmul(A.float(), B.float()).half()
C_our = hg.hgemm(A, B)
diff = (C_ref.float() - C_our.float()).abs().max().item()
rel_diff = diff / (C_ref.float().abs().max().item() + 1e-8)
print(f" Max abs diff: {diff:.6f}, rel: {rel_diff:.6f}")
assert rel_diff < 0.05, f"FAILED: rel_diff={rel_diff} too large"
print(f" PASS")
# Test 3: MoE expert GEMM with variable counts
print("\n--- Test 3: MoE expert GEMM (8 experts, variable tokens) ---")
num_experts = 8
K_dim = 128
N_dim = 256
expert_counts = torch.tensor([32, 16, 0, 48, 8, 24, 4, 12], dtype=torch.int32)
total_tokens = expert_counts.sum().item()
input_tensor = torch.randn(total_tokens, K_dim, dtype=torch.float16, device='cuda') * 0.1
weights = torch.randn(num_experts, N_dim, K_dim, dtype=torch.float16, device='cuda') * 0.1
output = hg.moe_expert_gemm(input_tensor, weights, expert_counts.cuda())
# Verify against torch reference
offset = 0
for e in range(num_experts):
cnt = expert_counts[e].item()
if cnt == 0:
continue
inp_e = input_tensor[offset:offset+cnt]
w_e = weights[e] # (N, K)
ref_e = torch.matmul(inp_e.float(), w_e.float().t()).half()
out_e = output[offset:offset+cnt]
diff_e = (ref_e.float() - out_e.float()).abs().max().item()
print(f" Expert {e} (tokens={cnt}): max_diff={diff_e:.6f}")
offset += cnt
print(f" PASS")
# Test 4: Performance benchmark
print("\n--- Test 4: Performance (256x4096 @ 4096x11008, 100 iters) ---")
M, N, K = 256, 11008, 4096
A = torch.randn(M, K, dtype=torch.float16, device='cuda')
B = torch.randn(K, N, dtype=torch.float16, device='cuda')
# Warmup
for _ in range(10):
hg.hgemm(A, B)
torch.cuda.synchronize()
import time
start = time.time()
for _ in range(100):
hg.hgemm(A, B)
torch.cuda.synchronize()
elapsed = time.time() - start
print(f" Custom kernel: {elapsed*10:.2f} ms/iter")
start = time.time()
for _ in range(100):
torch.matmul(A, B)
torch.cuda.synchronize()
elapsed2 = time.time() - start
print(f" torch.matmul: {elapsed2*10:.2f} ms/iter")
print(f" Ratio: {elapsed/elapsed2:.2f}x")
print("\n=== ALL TESTS PASSED ===")
PYTEST

131
ex_engine/xllm_kernels/build_test_hgemm_warp.sh Executable file
View File

@@ -0,0 +1,131 @@
#!/bin/bash
# build_test_hgemm_warp.sh — Compile and benchmark kernel 10 (warp tiling)
set -eo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CUDA_DIR="${SCRIPT_DIR}/cuda"
echo "=== Compile hgemm_warptiling (kernel 10, WARPSIZE=64) ==="
python3 -c "
import torch.utils.cpp_extension as ext
import os, shutil, glob
name = 'hgemm_warptiling'
build_dir = '${SCRIPT_DIR}/build/tmp_' + name
os.makedirs(build_dir, exist_ok=True)
try:
mod = ext.load(
name=name,
sources=[
'${CUDA_DIR}/hgemm_warptiling.cu',
'${CUDA_DIR}/bindings/hgemm_warp_bind.cpp',
],
extra_include_paths=['${CUDA_DIR}/headers'],
extra_cflags=['-O2', '-std=c++17'],
extra_cuda_cflags=['-O2'],
build_directory=build_dir,
verbose=True,
)
built = glob.glob(build_dir + '/' + name + '*.so')
if built:
dst = '${SCRIPT_DIR}/build/' + name + '.so'
shutil.copy2(built[0], dst)
print(f'[build] SUCCESS: {dst} ({os.path.getsize(dst)} bytes)')
except Exception as e:
print(f'[build] FAILED: {e}')
import traceback; traceback.print_exc()
"
echo ""
echo "=== Test ==="
python3 << 'PYTEST'
import torch, sys, os, glob, time
build_dir = 'ex_engine/xllm_kernels/build'
sys.path.insert(0, build_dir)
# Load kernel 10
try:
so = glob.glob(f'{build_dir}/tmp_hgemm_warptiling/hgemm_warptiling*.so')
if so:
import importlib.util
spec = importlib.util.spec_from_file_location("hgemm_warptiling", so[0])
hw = importlib.util.module_from_spec(spec)
spec.loader.exec_module(hw)
print("kernel 10 (warp tiling) loaded")
else:
print("SKIP: kernel 10 .so not found")
sys.exit(0)
except Exception as e:
print(f"SKIP: {e}")
sys.exit(0)
# Load kernel 6 for comparison
try:
so6 = glob.glob(f'{build_dir}/tmp_hgemm_blocktiling/hgemm_blocktiling*.so')
if so6:
spec6 = importlib.util.spec_from_file_location("hgemm_blocktiling", so6[0])
hb = importlib.util.module_from_spec(spec6)
spec6.loader.exec_module(hb)
has_k6 = True
print("kernel 6 (block tiling) loaded")
else:
has_k6 = False
except:
has_k6 = False
# Correctness
print("\n--- Correctness (128x128 @ 128x128) ---")
M, N, K = 128, 128, 128
A = torch.randn(M, K, dtype=torch.float16, device='cuda')
B = torch.randn(K, N, dtype=torch.float16, device='cuda')
C_ref = torch.matmul(A.float(), B.float()).half()
C_k10 = hw.hgemm_warp(A, B)
diff = (C_ref.float() - C_k10.float()).abs().max().item()
print(f" Max abs diff: {diff:.6f}")
assert diff < 2.0, f"FAIL diff={diff}"
print(" PASS")
# Correctness on MoE size
print("\n--- Correctness (256x4096 @ 4096x11008) ---")
M, N, K = 256, 11008, 4096
A = torch.randn(M, K, dtype=torch.float16, device='cuda') * 0.01
B = torch.randn(K, N, dtype=torch.float16, device='cuda') * 0.01
C_ref = torch.matmul(A.float(), B.float()).half()
C_k10 = hw.hgemm_warp(A, B)
diff = (C_ref.float() - C_k10.float()).abs().max().item()
rel = diff / (C_ref.float().abs().max().item() + 1e-8)
print(f" Max abs diff: {diff:.6f}, rel: {rel:.6f}")
print(" PASS" if rel < 0.1 else " WARN: large relative diff")
# Performance benchmark
print("\n--- Performance (256x4096 @ 4096x11008, 100 iters) ---")
M, N, K = 256, 11008, 4096
A = torch.randn(M, K, dtype=torch.float16, device='cuda')
B = torch.randn(K, N, dtype=torch.float16, device='cuda')
def bench(fn, name, iters=100, warmup=10):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
t0 = time.time()
for _ in range(iters):
fn()
torch.cuda.synchronize()
ms = (time.time() - t0) / iters * 1000
print(f" {name}: {ms:.2f} ms/iter")
return ms
t_torch = bench(lambda: torch.matmul(A, B), "torch.matmul")
t_k10 = bench(lambda: hw.hgemm_warp(A, B), "kernel 10 (warp)")
if has_k6:
t_k6 = bench(lambda: hb.hgemm(A, B), "kernel 6 (block)")
print(f"\n K10/torch = {t_k10/t_torch:.2f}x")
print(f" K6/torch = {t_k6/t_torch:.2f}x")
print(f" K10/K6 = {t_k10/t_k6:.2f}x (K10 should be faster)")
else:
print(f"\n K10/torch = {t_k10/t_torch:.2f}x")
print("\n=== DONE ===")
PYTEST

189
ex_engine/xllm_kernels/cuda/activation.cu Normal file
View File

@@ -0,0 +1,189 @@
/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <c10/cuda/CUDAGuard.h>
#include <torch/cuda.h>
#include <torch/extension.h>
#include <cstdint>
#include "device_utils.cuh"
// ref to:
// https://github.com/vllm-project/vllm/blob/main/csrc/activation_kernels.cu
namespace {
using ::xllm::kernel::cuda::xllm_ldg;
template <typename scalar_t,
scalar_t (*ACT_FN)(const scalar_t&),
bool act_first>
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
const scalar_t& y) {
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
}
// Check if pointer is 16-byte aligned for int4 vectorized access
__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
}
// Activation and gating kernel template with 128-bit vectorized access
// optimization.
template <typename scalar_t,
scalar_t (*ACT_FN)(const scalar_t&),
bool act_first>
__global__ void XLLM_KERNEL_ATTR(1024)
act_and_mul_kernel(scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., 2, d]
const int d) {
constexpr int kVecSize = 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 >= kVecSize) {
// Fast path: 128-bit vectorized loop
const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
int4* out_vec = reinterpret_cast<int4*>(out_ptr);
const int num_vecs = d / kVecSize;
const int vec_end = num_vecs * kVecSize;
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
int4 x = xllm_ldg(&x_vec[i]), y = xllm_ldg(&y_vec[i]), r;
auto* xp = reinterpret_cast<scalar_t*>(&x);
auto* yp = reinterpret_cast<scalar_t*>(&y);
auto* rp = reinterpret_cast<scalar_t*>(&r);
#pragma unroll
for (int j = 0; j < kVecSize; j++) {
rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
}
out_vec[i] = r;
}
// Scalar cleanup for remaining elements
for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(xllm_ldg(&x_ptr[i]),
xllm_ldg(&y_ptr[i]));
}
} else {
// Scalar fallback for unaligned data or small d
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
const scalar_t x = xllm_ldg(&x_ptr[idx]);
const scalar_t y = xllm_ldg(&y_ptr[idx]);
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
}
}
}
template <typename T>
__device__ __forceinline__ T silu_kernel(const T& x) {
// x * sigmoid(x)
const float f = static_cast<float>(x);
return static_cast<T>(f / (1.0f + expf(-f)));
}
template <typename T>
__device__ __forceinline__ T gelu_kernel(const T& x) {
// Equivalent to PyTorch GELU with 'none' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
const float f = static_cast<float>(x);
constexpr float kAlpha = M_SQRT1_2;
return static_cast<T>(f * 0.5f * (1.0f + ::erf(f * kAlpha)));
}
template <typename T>
__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
// Equivalent to PyTorch GELU with 'tanh' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
const float f = static_cast<float>(x);
constexpr float kBeta = M_SQRT2 * M_2_SQRTPI * 0.5f;
constexpr float kKappa = 0.044715;
float x_cube = f * f * f;
float inner = kBeta * (f + kKappa * x_cube);
return static_cast<T>(0.5f * f * (1.0f + ::tanhf(inner)));
}
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens); \
dim3 block(std::min(d, 1024)); \
if (num_tokens == 0) { \
return; \
} \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
DISPATCH_FLOATING_TYPES(input.scalar_type(), "act_and_mul_kernel", [&] { \
act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST> \
<<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
});
void silu_and_mul(torch::Tensor out, // [..., d]
torch::Tensor input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(silu_kernel, true);
}
void gelu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(gelu_kernel, true);
}
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(gelu_tanh_kernel, true);
}
} // namespace
namespace xllm::kernel::cuda {
void act_and_mul(torch::Tensor out,
torch::Tensor input,
const std::string& act_mode) {
if (act_mode != "silu" && act_mode != "gelu" && act_mode != "gelu_tanh" &&
act_mode != "gelu_pytorch_tanh") {
TORCH_CHECK(false, "Unsupported act mode: ", act_mode,
", only support silu, gelu, gelu_tanh, gelu_pytorch_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" || act_mode == "gelu_pytorch_tanh") {
// gelu_tanh or gelu_pytorch_tanh (mathematically equivalent)
gelu_tanh_and_mul(out, input);
}
}
} // namespace xllm::kernel::cuda

129
ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp Normal file
View File

@@ -0,0 +1,129 @@
/*
* corex_batched_gemm_bind.cpp — pybind11 wrapper for CUTLASS batched GEMM
*
* Kernel uses RowMajor + OpClassTensorOp + Cu10 (verified 2.462ms).
* Source: ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu
*/
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
// Implemented in corex_batched_gemm_kernel.cu
// RowMajor, FP16 data, FP32 accumulation, TCU, Cu10
cudaError_t cutlass_batched_hgemm(
int m, int n, int k,
__half const *A, int lda, long long int batch_stride_A,
__half const *B, int ldb, long long int batch_stride_B,
__half *C, int ldc, long long int batch_stride_C,
int batch_count);
/*
* batched_gemm_fp16: C[i] = A[i] @ B[i]
* A: (batch, M, K) row-major
* B: (batch, K, N) row-major
* C: (batch, M, N) row-major
*
* Both A and B must be contiguous fp16 CUDA tensors.
*/
torch::Tensor batched_gemm_fp16(
torch::Tensor A, // (batch, M, K)
torch::Tensor B) // (batch, K, N)
{
TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA tensors");
TORCH_CHECK(A.scalar_type() == torch::kFloat16 &&
B.scalar_type() == torch::kFloat16,
"inputs must be float16");
TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
"inputs must be contiguous");
TORCH_CHECK(A.dim() == 3 && B.dim() == 3,
"inputs must be 3D (batch, rows, cols)");
int batch = A.size(0);
int M = A.size(1);
int K = A.size(2);
int N = B.size(2);
TORCH_CHECK(B.size(0) == batch, "batch size mismatch");
TORCH_CHECK(B.size(1) == K, "K dimension mismatch");
auto C = torch::zeros({batch, M, N}, A.options());
// RowMajor: A is (M,K) with lda=K, B is (K,N) with ldb=N, C is (M,N) with ldc=N
auto status = cutlass_batched_hgemm(
M, N, K,
reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
K, (long long)M * K, // lda, strideA
reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
N, (long long)K * N, // ldb, strideB
reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
N, (long long)M * N, // ldc, strideC
batch);
TORCH_CHECK(status == cudaSuccess,
"CUTLASS batched HGEMM failed: ", cudaGetErrorString(status));
return C;
}
/*
* moe_decode_fused: Full MoE decode using TCU batched GEMM.
*
* hidden_states: (1, H)
* w13_sel: (K, 2*I, H) — already gathered expert weights
* w2_sel: (K, H, I) — already gathered expert weights
* topk_weights: (K,)
*
* Pipeline:
* 1. gate_up = x @ w13^T via batched GEMM (K, 1, 2I)
* 2. act = silu(gate) * up
* 3. down = act @ w2^T via batched GEMM (K, 1, H)
* 4. out = weighted sum
*/
torch::Tensor moe_decode_fused(
torch::Tensor hidden_states, // (1, H)
torch::Tensor w13_sel, // (K, 2*I, H)
torch::Tensor w2_sel, // (K, H, I)
torch::Tensor topk_weights) // (K,)
{
int K_experts = w13_sel.size(0);
int two_I = w13_sel.size(1);
int H = w13_sel.size(2);
int I = two_I / 2;
// x: (1, H) → expand to (K, 1, H)
auto x = hidden_states.expand({K_experts, 1, H}).contiguous();
// w13^T: (K, 2I, H) → transpose last two dims → (K, H, 2I)
auto w13_t = w13_sel.transpose(1, 2).contiguous(); // (K, H, 2I)
// Step 1: gate_up = x @ w13^T → (K, 1, 2I)
auto gate_up = batched_gemm_fp16(x, w13_t);
gate_up = gate_up.squeeze(1); // (K, 2I)
// Step 2: silu activation
auto chunks = gate_up.chunk(2, /*dim=*/1);
auto act = torch::sigmoid(chunks[0]) * chunks[0] * chunks[1]; // silu(gate) * up
act = act.unsqueeze(1); // (K, 1, I)
// w2^T: (K, H, I) → transpose → (K, I, H)
auto w2_t = w2_sel.transpose(1, 2).contiguous(); // (K, I, H)
// Step 3: down = act @ w2^T → (K, 1, H)
auto down = batched_gemm_fp16(act, w2_t);
down = down.squeeze(1); // (K, H)
// Step 4: weighted sum
auto out = (down * topk_weights.unsqueeze(1)).sum(0, true);
return out.to(hidden_states.dtype());
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.doc() = "CUTLASS batched GEMM for MoE decode (BI-V100 TCU, Cu10 TensorOp)";
m.def("batched_gemm_fp16", &batched_gemm_fp16,
"Batched GEMM: (B,M,K) x (B,K,N) -> (B,M,N) in fp16 via TCU",
py::arg("A"), py::arg("B"));
m.def("moe_decode_fused", &moe_decode_fused,
"Full MoE decode via TCU batched GEMM",
py::arg("hidden_states"), py::arg("w13_sel"),
py::arg("w2_sel"), py::arg("topk_weights"));
}

135
ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp Normal file
View File

@@ -0,0 +1,135 @@
// hgemm_bind.cpp — pybind11 bindings for hgemm_blocktiling.cu
//
// Exports:
// hgemm(A, B, M, N, K) → C
// moe_expert_gemm(input, weights, expert_counts) → output
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAStream.h>
#include <vector>
// Forward declarations from hgemm_blocktiling.cu
void launch_hgemm_blocktiling(
int M, int N, int K,
const __half* alpha, const __half* A, int lda,
const __half* B, int ldb,
const __half* beta, __half* C, int ldc,
cudaStream_t stream);
void launch_moe_expert_hgemm(
int num_experts,
const int* expert_counts,
const int* expert_offsets,
int N, int K,
const __half* input,
const __half* weights,
__half* output,
cudaStream_t stream);
// ============================================================================
// Python-facing wrappers
// ============================================================================
// Simple GEMM: C = A @ B
// A: (M, K) fp16, B: (K, N) fp16 → C: (M, N) fp16
torch::Tensor hgemm(torch::Tensor A, torch::Tensor B) {
TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
TORCH_CHECK(A.scalar_type() == torch::kHalf, "A must be fp16");
TORCH_CHECK(B.scalar_type() == torch::kHalf, "B must be fp16");
TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D");
TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions must match");
int M = A.size(0);
int K = A.size(1);
int N = B.size(1);
auto C = torch::zeros({M, N}, A.options());
__half alpha = __float2half(1.0f);
__half beta = __float2half(0.0f);
cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
launch_hgemm_blocktiling(
M, N, K, &alpha,
reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
A.size(1),
reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
B.size(1),
&beta,
reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
C.size(1),
stream);
return C;
}
// MoE expert GEMM: for each expert e, compute
// output[offset_e : offset_e + count_e] = input[offset_e : offset_e + count_e] @ weights[e].T
//
// input: (total_tokens, K) fp16
// weights: (num_experts, N, K) fp16 — weight layout matches vllm w13/w2 convention
// expert_counts: (num_experts,) int32 — number of tokens per expert
//
// Returns: output (total_tokens, N) fp16
torch::Tensor moe_expert_gemm(
torch::Tensor input,
torch::Tensor weights,
torch::Tensor expert_counts
) {
TORCH_CHECK(input.is_cuda() && weights.is_cuda(), "Inputs must be CUDA");
TORCH_CHECK(input.scalar_type() == torch::kHalf, "input must be fp16");
TORCH_CHECK(weights.scalar_type() == torch::kHalf, "weights must be fp16");
TORCH_CHECK(expert_counts.scalar_type() == torch::kInt32 ||
expert_counts.scalar_type() == torch::kInt64,
"expert_counts must be int32 or int64");
int total_tokens = input.size(0);
int K = input.size(1);
int num_experts = weights.size(0);
int N = weights.size(1); // output dim
TORCH_CHECK(weights.size(2) == K, "weights K dim must match input");
auto output = torch::zeros({total_tokens, N}, input.options());
// Convert expert_counts to host int array
auto counts_cpu = expert_counts.to(torch::kCPU).to(torch::kInt32).contiguous();
std::vector<int> counts(num_experts);
std::vector<int> offsets(num_experts);
int cumsum = 0;
for (int i = 0; i < num_experts; i++) {
counts[i] = counts_cpu.data_ptr<int32_t>()[i];
offsets[i] = cumsum;
cumsum += counts[i];
}
cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
launch_moe_expert_hgemm(
num_experts,
counts.data(),
offsets.data(),
N, K,
reinterpret_cast<const __half*>(input.data_ptr<at::Half>()),
reinterpret_cast<const __half*>(weights.data_ptr<at::Half>()),
reinterpret_cast<__half*>(output.data_ptr<at::Half>()),
stream);
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("hgemm", &hgemm,
"FP16 GEMM: C = A @ B (adapted from siboehm kernel 6 for BI-V100)",
py::arg("A"), py::arg("B"));
m.def("moe_expert_gemm", &moe_expert_gemm,
"MoE expert GEMM: per-expert matmul with variable token counts",
py::arg("input"), py::arg("weights"), py::arg("expert_counts"));
}

36
ex_engine/xllm_kernels/cuda/bindings/hgemm_warp_bind.cpp Normal file
View File

@@ -0,0 +1,36 @@
// hgemm_warp_bind.cpp — pybind11 for hgemm_warptiling (kernel 10, warp64)
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAStream.h>
void launch_hgemm_warptiling(
int M, int N, int K, float alpha,
const __half* A, const __half* B,
float beta, __half* C, cudaStream_t stream);
torch::Tensor hgemm_warp(torch::Tensor A, torch::Tensor B) {
TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
TORCH_CHECK(A.scalar_type() == torch::kHalf, "A must be fp16");
TORCH_CHECK(B.scalar_type() == torch::kHalf, "B must be fp16");
TORCH_CHECK(A.size(1) == B.size(0), "Inner dims must match");
int M = A.size(0), K = A.size(1), N = B.size(1);
auto C = torch::zeros({M, N}, A.options());
cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
launch_hgemm_warptiling(M, N, K, 1.0f,
reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
0.0f,
reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
stream);
return C;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("hgemm_warp", &hgemm_warp,
"FP16 GEMM warp-tiling (siboehm K10, WARPSIZE=64 for BI-V100)");
}

18
ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp Normal file
View File

@@ -0,0 +1,18 @@
// xllm_activation_bind.cpp
#include <torch/extension.h>
namespace xllm::kernel::cuda {
void act_and_mul(torch::Tensor out, torch::Tensor input,
const std::string& act_mode);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("silu_and_mul", [](torch::Tensor out, torch::Tensor input) {
xllm::kernel::cuda::act_and_mul(out, input, "silu");
}, "SiLU and Mul", py::arg("out"), py::arg("input"));
m.def("gelu_and_mul", [](torch::Tensor out, torch::Tensor input) {
xllm::kernel::cuda::act_and_mul(out, input, "gelu");
}, "GELU and Mul", py::arg("out"), py::arg("input"));
m.def("act_and_mul", &xllm::kernel::cuda::act_and_mul,
"Activation and Mul", py::arg("out"), py::arg("input"), py::arg("act_mode"));
}

19
ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp Normal file
View File

@@ -0,0 +1,19 @@
// xllm_cache_bind.cpp
#include <torch/extension.h>
namespace xllm::kernel::cuda {
void reshape_paged_cache(torch::Tensor slot_ids, torch::Tensor keys,
torch::Tensor values, torch::Tensor key_cache,
torch::Tensor value_cache);
void block_copy(torch::Tensor key_cache_ptrs, torch::Tensor value_cache_ptrs,
torch::Tensor src_block_indices, torch::Tensor dst_block_indices,
torch::Tensor cum_sum, int64_t numel_per_block,
torch::ScalarType cache_dtype);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("reshape_paged_cache", &xllm::kernel::cuda::reshape_paged_cache,
"Reshape Paged KV Cache");
m.def("block_copy", &xllm::kernel::cuda::block_copy,
"Block Copy for KV Cache");
}

38
ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp Normal file
View File

@@ -0,0 +1,38 @@
// xllm_fused_qknorm_rope_bind.cpp — pybind11 for fused QK-Norm + RoPE kernel
// Source: upstream_ref/xllm/xllm/core/kernels/cuda/fused_qknorm_rope.cu
// Saves 4 kernel launches per layer (separate q_norm, k_norm, q_rope, k_rope)
// Qwen3.5 has 32 full-attention layers → saves 128 kernel launches per forward
#include <torch/extension.h>
namespace xllm::kernel::cuda {
void fused_qk_norm_rope(
torch::Tensor& qkv,
int64_t num_heads_q,
int64_t num_heads_k,
int64_t num_heads_v,
int64_t head_dim,
double eps,
const torch::Tensor& q_weight,
const torch::Tensor& k_weight,
const torch::Tensor& cos_sin_cache,
bool interleaved,
const torch::Tensor& position_ids);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fused_qk_norm_rope",
&xllm::kernel::cuda::fused_qk_norm_rope,
"Fused QK-Norm + RoPE (xllm CUDA kernel)",
py::arg("qkv"),
py::arg("num_heads_q"),
py::arg("num_heads_k"),
py::arg("num_heads_v"),
py::arg("head_dim"),
py::arg("eps") = 1e-6,
py::arg("q_weight"),
py::arg("k_weight"),
py::arg("cos_sin_cache"),
py::arg("interleaved") = false,
py::arg("position_ids"));
}

34
ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp Normal file
View File

@@ -0,0 +1,34 @@
// xllm_moe_bind.cpp — pybind11 for MoE CUDA kernels
#include <torch/extension.h>
#include <optional>
#include <tuple>
namespace xllm::kernel::cuda {
std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
torch::Tensor& gating_output, int64_t topk, bool renormalize,
const std::optional<torch::Tensor>& correction_bias,
const std::string& scoring_func);
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> moe_compute_index(
const torch::Tensor& expert_id, int64_t num_experts);
torch::Tensor moe_combine_result(
const torch::Tensor& gemm2, const torch::Tensor& reduce_weight,
int64_t N, int32_t topk);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("moe_fused_topk", &xllm::kernel::cuda::moe_fused_topk,
"MoE fused topk (softmax or sigmoid routing)",
py::arg("gating_output"), py::arg("topk"),
py::arg("renormalize") = true,
py::arg("correction_bias") = py::none(),
py::arg("scoring_func") = "softmax");
m.def("moe_compute_index", &xllm::kernel::cuda::moe_compute_index,
"MoE compute permutation index (histogram + prefix_sum + place)",
py::arg("expert_id"), py::arg("num_experts"));
m.def("moe_combine_result", &xllm::kernel::cuda::moe_combine_result,
"MoE combine (reorder + weighted sum)",
py::arg("gemm2"), py::arg("reduce_weight"),
py::arg("N"), py::arg("topk"));
}

24
ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp Normal file
View File

@@ -0,0 +1,24 @@
// xllm_norm_bind.cpp — pybind11 entry point for xllm norm kernels
// Compiled together with norm.cu to produce xllm_norm.so
//
// Exports: rms_norm, fused_add_rms_norm
#include <torch/extension.h>
namespace xllm::kernel::cuda {
void rms_norm(torch::Tensor output, torch::Tensor input,
torch::Tensor weight, double eps);
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
torch::Tensor& weight, double epsilon);
} // namespace xllm::kernel::cuda
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rms_norm", &xllm::kernel::cuda::rms_norm,
"RMS Norm (xllm CUDA kernel)",
py::arg("output"), py::arg("input"),
py::arg("weight"), py::arg("eps") = 1e-6);
m.def("fused_add_rms_norm", &xllm::kernel::cuda::fused_add_rms_norm,
"Fused Add + RMS Norm (xllm CUDA kernel)",
py::arg("input"), py::arg("residual"),
py::arg("weight"), py::arg("epsilon") = 1e-6);
}

17
ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp Normal file
View File

@@ -0,0 +1,17 @@
// xllm_rope_bind.cpp
#include <torch/extension.h>
#include <optional>
namespace xllm::kernel::cuda {
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
std::optional<torch::Tensor> key,
torch::Tensor& cos_sin_cache, bool is_neox);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rotary_embedding", &xllm::kernel::cuda::rotary_embedding,
"Rotary Position Embedding (xllm CUDA kernel)",
py::arg("positions"), py::arg("query"),
py::arg("key"), py::arg("cos_sin_cache"),
py::arg("is_neox") = true);
}

210
ex_engine/xllm_kernels/cuda/block_copy.cu Normal file
View File

@@ -0,0 +1,210 @@
/* Copyright 2025-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 <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <torch/extension.h>
#include <cstdint>
#include <type_traits>
#include "device_utils.cuh"
namespace xllm::kernel::cuda {
namespace {
template <typename scalar_t>
struct VecType;
template <>
struct VecType<c10::Half> {
using type = uint4;
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<c10::BFloat16> {
using type = uint4;
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<float> {
using type = float4;
static constexpr int32_t vec_width = 4;
};
DEVICE_INLINE int32_t find_group_idx(const int32_t* __restrict__ cum_sum,
const int32_t num_groups,
const int32_t dst_idx) {
int32_t left = 0;
int32_t right = num_groups - 1;
while (left < right) {
const int32_t mid = left + ((right - left) >> 1);
const bool move_left = dst_idx < cum_sum[mid];
right = move_left ? mid : right;
left = move_left ? left : mid + 1;
}
return left;
}
template <typename scalar_t, bool kVectorized>
__global__ void block_copy_kernel(const int64_t* __restrict__ key_cache_ptrs,
const int64_t* __restrict__ value_cache_ptrs,
const int32_t* __restrict__ src_block_indices,
const int32_t* __restrict__ dst_block_indices,
const int32_t* __restrict__ cum_sum,
const int32_t num_groups,
const int64_t numel_per_block) {
const int64_t layer_idx = static_cast<int64_t>(blockIdx.x);
const int32_t dst_linear_idx = static_cast<int32_t>(blockIdx.y);
const int64_t tile_idx = static_cast<int64_t>(blockIdx.z);
scalar_t* __restrict__ key_cache = reinterpret_cast<scalar_t*>(
static_cast<uintptr_t>(key_cache_ptrs[layer_idx]));
scalar_t* __restrict__ value_cache = reinterpret_cast<scalar_t*>(
static_cast<uintptr_t>(value_cache_ptrs[layer_idx]));
const int32_t group_idx = find_group_idx(cum_sum, num_groups, dst_linear_idx);
const int32_t src_block = src_block_indices[group_idx];
const int32_t dst_block = dst_block_indices[dst_linear_idx];
const int64_t src_offset = static_cast<int64_t>(src_block) * numel_per_block;
const int64_t dst_offset = static_cast<int64_t>(dst_block) * numel_per_block;
if constexpr (kVectorized) {
using VecTypeT = typename VecType<scalar_t>::type;
constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
const int64_t num_vecs_per_block = numel_per_block / kVecWidth;
const int64_t vec_idx = tile_idx * static_cast<int64_t>(blockDim.x) +
static_cast<int64_t>(threadIdx.x);
if (vec_idx >= num_vecs_per_block) {
return;
}
const int64_t elem_offset = vec_idx * kVecWidth;
const auto* key_src_vec =
reinterpret_cast<const VecTypeT*>(key_cache + src_offset + elem_offset);
const auto* value_src_vec = reinterpret_cast<const VecTypeT*>(
value_cache + src_offset + elem_offset);
auto* key_dst_vec =
reinterpret_cast<VecTypeT*>(key_cache + dst_offset + elem_offset);
auto* value_dst_vec =
reinterpret_cast<VecTypeT*>(value_cache + dst_offset + elem_offset);
*key_dst_vec = *key_src_vec;
*value_dst_vec = *value_src_vec;
} else {
const int64_t elem_idx = tile_idx * static_cast<int64_t>(blockDim.x) +
static_cast<int64_t>(threadIdx.x);
if (elem_idx >= numel_per_block) {
return;
}
key_cache[dst_offset + elem_idx] = key_cache[src_offset + elem_idx];
value_cache[dst_offset + elem_idx] = value_cache[src_offset + elem_idx];
}
}
} // namespace
void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor value_cache_ptrs,
torch::Tensor src_block_indices,
torch::Tensor dst_block_indices,
torch::Tensor cum_sum,
int64_t numel_per_block,
torch::ScalarType cache_dtype) {
if (src_block_indices.numel() == 0) {
return;
}
TORCH_CHECK(key_cache_ptrs.is_cuda());
TORCH_CHECK(value_cache_ptrs.is_cuda());
TORCH_CHECK(src_block_indices.is_cuda());
TORCH_CHECK(dst_block_indices.is_cuda());
TORCH_CHECK(cum_sum.is_cuda());
TORCH_CHECK(key_cache_ptrs.scalar_type() == torch::kInt64);
TORCH_CHECK(value_cache_ptrs.scalar_type() == torch::kInt64);
TORCH_CHECK(src_block_indices.scalar_type() == torch::kInt32);
TORCH_CHECK(dst_block_indices.scalar_type() == torch::kInt32);
TORCH_CHECK(cum_sum.scalar_type() == torch::kInt32);
TORCH_CHECK(key_cache_ptrs.dim() == 1);
TORCH_CHECK(value_cache_ptrs.dim() == 1);
TORCH_CHECK(src_block_indices.dim() == 1);
TORCH_CHECK(dst_block_indices.dim() == 1);
TORCH_CHECK(cum_sum.dim() == 1);
TORCH_CHECK(key_cache_ptrs.is_contiguous());
TORCH_CHECK(value_cache_ptrs.is_contiguous());
TORCH_CHECK(src_block_indices.is_contiguous());
TORCH_CHECK(dst_block_indices.is_contiguous());
TORCH_CHECK(cum_sum.is_contiguous());
TORCH_CHECK(key_cache_ptrs.size(0) == value_cache_ptrs.size(0));
TORCH_CHECK(src_block_indices.size(0) == cum_sum.size(0));
TORCH_CHECK(numel_per_block > 0);
const at::cuda::OptionalCUDAGuard device_guard(key_cache_ptrs.device());
constexpr int32_t kThreadsPerBlock = 256;
const int32_t num_layers = static_cast<int32_t>(key_cache_ptrs.size(0));
const int32_t num_groups = static_cast<int32_t>(src_block_indices.size(0));
const int32_t num_dst_blocks =
static_cast<int32_t>(dst_block_indices.size(0));
const cudaStream_t stream =
c10::cuda::getCurrentCUDAStream(key_cache_ptrs.get_device());
DISPATCH_FLOATING_TYPES(cache_dtype, "block_copy_kernel", [&] {
constexpr bool kHasVecType = std::is_same_v<scalar_t, float> ||
std::is_same_v<scalar_t, c10::Half> ||
std::is_same_v<scalar_t, c10::BFloat16>;
if constexpr (kHasVecType) {
constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
if (numel_per_block % kVecWidth == 0) {
const int64_t tiles_per_block =
ceil_div<int64_t>(numel_per_block / kVecWidth, kThreadsPerBlock);
const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
block_copy_kernel<scalar_t, true>
<<<grid, kThreadsPerBlock, 0, stream>>>(
key_cache_ptrs.data_ptr<int64_t>(),
value_cache_ptrs.data_ptr<int64_t>(),
src_block_indices.data_ptr<int32_t>(),
dst_block_indices.data_ptr<int32_t>(),
cum_sum.data_ptr<int32_t>(),
num_groups,
numel_per_block);
C10_CUDA_KERNEL_LAUNCH_CHECK();
return;
}
}
const int64_t tiles_per_block =
ceil_div<int64_t>(numel_per_block, kThreadsPerBlock);
const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
block_copy_kernel<scalar_t, false><<<grid, kThreadsPerBlock, 0, stream>>>(
key_cache_ptrs.data_ptr<int64_t>(),
value_cache_ptrs.data_ptr<int64_t>(),
src_block_indices.data_ptr<int32_t>(),
dst_block_indices.data_ptr<int32_t>(),
cum_sum.data_ptr<int32_t>(),
num_groups,
numel_per_block);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
}
} // namespace xllm::kernel::cuda

67
ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu Normal file
View File

@@ -0,0 +1,67 @@
/*
* corex_batched_gemm_kernel.cu — FP16 Cu10 TensorOp batched GEMM
*
* Uses cutlass::gemm::device::GemmBatched with:
* - OpClassTensorOp (TCU, not SIMT)
* - arch::Cu10 (BI-V100)
* - float accumulation (FP32, not FP16)
*
* Source: ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu (verified 2.462ms)
*/
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/gemm/device/gemm_batched.h"
cudaError_t cutlass_batched_hgemm(
int m, int n, int k,
__half const *A, int lda, long long int batch_stride_A,
__half const *B, int ldb, long long int batch_stride_B,
__half *C, int ldc, long long int batch_stride_C,
int batch_count)
{
using Gemm = cutlass::gemm::device::GemmBatched<
cutlass::half_t, // ElementA
cutlass::layout::RowMajor, // LayoutA
cutlass::half_t, // ElementB
cutlass::layout::RowMajor, // LayoutB
cutlass::half_t, // ElementC
cutlass::layout::RowMajor, // LayoutC
float, // ElementAccumulator — FP32!
cutlass::arch::OpClassTensorOp, // OperatorClass — TCU!
cutlass::arch::Cu10 // ArchTag — BI-V100!
// Defaults from DefaultGemmConfiguration<OpClassTensorOp, Cu10, half, half, half, float>:
// ThreadblockShape = <128, 128, 32>
// WarpShape = <32, 32, 32>
// InstructionShape = <16, 16, 16>
// Stages = 2
>;
float alpha = 1.0f;
float beta = 0.0f;
Gemm gemm_op;
cutlass::Status status = gemm_op({
{m, n, k},
{reinterpret_cast<cutlass::half_t const *>(A), lda},
batch_stride_A,
{reinterpret_cast<cutlass::half_t const *>(B), ldb},
batch_stride_B,
{reinterpret_cast<cutlass::half_t const *>(C), ldc},
batch_stride_C,
{reinterpret_cast<cutlass::half_t *>(C), ldc},
batch_stride_C,
{alpha, beta},
batch_count
});
if (status != cutlass::Status::kSuccess) {
return cudaErrorUnknown;
}
return cudaSuccess;
}

463
ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu Normal file
View File

@@ -0,0 +1,463 @@
/* 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 <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <torch/cuda.h>
#include <cmath>
#include <type_traits>
#include "cuda_ops_api.h"
#include "type_convert.cuh"
#include "utils.h"
using at::device_of;
// Borrowed from:
// https://github.com/vllm-project/vllm/blob/022f3cea5327cc720a325c50931e1edcfdf2d32b/csrc/fused_qknorm_rope_kernel.cu
constexpr uint32_t kFinalMask = 0xffffffffu;
namespace {
using namespace xllm::kernel::cuda;
template <typename T, int num>
struct packed_as;
// Specialization for packed_as used in this kernel.
template <>
struct packed_as<uint, 1> {
using type = uint;
};
template <>
struct packed_as<uint, 2> {
using type = uint2;
};
template <>
struct packed_as<uint, 4> {
using type = uint4;
};
template <typename T>
__inline__ __device__ T warp_reduce_sum(T val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1)
val += __shfl_xor_sync(kFinalMask, val, mask, 32);
return val;
}
template <typename T>
inline __device__ __host__ T div_up(T m, T n) {
return (m + n - 1) / n;
}
// Perform per-head QK Norm and RoPE in a single kernel.
// scalar_t_in: data type of QKV and RMSNorm weights
// scalar_t_cache: data type of cos/sin cache
// head_dim: the dimension of each head
// interleave: interleave=!is_neox.
template <typename scalar_t_in,
typename scalar_t_cache,
int head_dim,
bool interleave>
__global__ void fused_qknorm_rope_kernel(
void* qkv_void, // Combined QKV tensor
int const num_heads_q, // Number of query heads
int const num_heads_k, // Number of key heads
int const num_heads_v, // Number of value heads
float const eps, // Epsilon for RMS normalization
void const* q_weight_void, // RMSNorm weights for query
void const* k_weight_void, // RMSNorm weights for key
void const* cos_sin_cache_void, // Pre-computed cos/sin cache
int64_t const* position_ids, // Position IDs for RoPE
int const num_tokens, // Number of tokens
int const rotary_dim // Dimension for RoPE
) {
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
std::is_same_v<scalar_t_cache, c10::BFloat16>) {
return;
} else {
#endif
using Converter = _typeConvert<scalar_t_in>;
static_assert(Converter::exists,
"Input QKV data type is not supported for this CUDA "
"architecture or toolkit version.");
using T_in = typename Converter::hip_type;
using T2_in = typename Converter::packed_hip_type;
using CacheConverter = _typeConvert<scalar_t_cache>;
static_assert(CacheConverter::exists,
"Cache data type is not supported for this CUDA architecture "
"or toolkit version.");
using T_cache = typename CacheConverter::hip_type;
T_in* qkv = reinterpret_cast<T_in*>(qkv_void);
T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void);
T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void);
T_cache const* cos_sin_cache =
reinterpret_cast<T_cache const*>(cos_sin_cache_void);
int const warpsPerBlock = blockDim.x / 32;
int const warpId = threadIdx.x / 32;
int const laneId = threadIdx.x % 32;
// Calculate global warp index to determine which head/token this warp
// processes
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
// Total number of attention heads (Q and K)
int const total_qk_heads = num_heads_q + num_heads_k;
// Determine which token and head type (Q or K) this warp processes
int const tokenIdx = globalWarpIdx / total_qk_heads;
int const localHeadIdx = globalWarpIdx % total_qk_heads;
// Skip if this warp is assigned beyond the number of tokens
if (tokenIdx >= num_tokens) return;
bool const isQ = localHeadIdx < num_heads_q;
int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
int const num_heads = num_heads_q + num_heads_k + num_heads_v;
static_assert(head_dim % (32 * 2) == 0,
"head_dim must be divisible by 64 (each warp processes one "
"head, and each thread gets even number of "
"elements)");
constexpr int numElemsPerThread = head_dim / 32;
float elements[numElemsPerThread];
constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
static_assert(elemSizeBytes % 4 == 0,
"numSizeBytes must be a multiple of 4");
constexpr int vecSize =
elemSizeBytes /
4; // Use packed_as<uint, vecSize> to perform loading/saving.
using vec_T = typename packed_as<uint, vecSize>::type;
int offsetWarp; // Offset for the warp
if (isQ) {
// Q segment: token offset + head offset within Q segment
offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
} else {
// K segment: token offset + entire Q segment + head offset within K
// segment
offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
headIdx * head_dim;
}
int offsetThread = offsetWarp + laneId * numElemsPerThread;
// Sum of squares for RMSNorm
float sumOfSquares = 0.0f;
// Load.
{
vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]);
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
#pragma unroll
for (int i = 0; i < num_packed_elems; i++) {
// Interpret the generic vector chunk as the specific packed type
T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
// Convert to float2 for computation
float2 vals = Converter::convert(packed_val);
sumOfSquares += vals.x * vals.x;
sumOfSquares += vals.y * vals.y;
elements[2 * i] = vals.x;
elements[2 * i + 1] = vals.y;
}
}
// Reduce sum across warp using the utility function
sumOfSquares = warp_reduce_sum(sumOfSquares);
// Compute RMS normalization factor
float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
// Normalize elements
#pragma unroll
for (int i = 0; i < numElemsPerThread; i++) {
int dim = laneId * numElemsPerThread + i;
float weight = isQ ? Converter::convert(q_weight[dim])
: Converter::convert(k_weight[dim]);
elements[i] *= rms_rcp * weight;
}
// Apply RoPE to normalized elements
float elements2[numElemsPerThread]; // Additional buffer required for RoPE.
int64_t pos_id = position_ids[tokenIdx];
// Calculate cache pointer for this position - similar to
// pos_encoding_kernels.cu
T_cache const* cache_ptr = cos_sin_cache + pos_id * rotary_dim;
int const embed_dim = rotary_dim / 2;
T_cache const* cos_ptr = cache_ptr;
T_cache const* sin_ptr = cache_ptr + embed_dim;
int const rotary_lanes = rotary_dim / numElemsPerThread; // rotary range
if (laneId < rotary_lanes) {
if constexpr (interleave) {
// Perform interleaving. Use pre-computed cos/sin values.
#pragma unroll
for (int i = 0; i < numElemsPerThread / 2; ++i) {
int const idx0 = 2 * i;
int const idx1 = 2 * i + 1;
// Global dimension index in the head
int const dim_idx = laneId * numElemsPerThread + idx0;
float const val0 = elements[idx0];
float const val1 = elements[idx1];
int const half_dim = dim_idx / 2;
float const cos_val =
CacheConverter::convert(__ldg(cos_ptr + half_dim));
float const sin_val =
CacheConverter::convert(__ldg(sin_ptr + half_dim));
elements[idx0] = val0 * cos_val - val1 * sin_val;
elements[idx1] = val0 * sin_val + val1 * cos_val;
}
} else {
// Before data exchange with in warp, we need to sync.
__syncwarp();
int pairOffset = (rotary_dim / 2) / numElemsPerThread;
// Get the data from the other half of the warp. Use pre-computed
// cos/sin values.
#pragma unroll
for (int i = 0; i < numElemsPerThread; i++) {
elements2[i] = __shfl_xor_sync(kFinalMask, elements[i], pairOffset);
if (laneId < pairOffset) {
elements2[i] = -elements2[i];
}
int dim_idx = laneId * numElemsPerThread + i;
dim_idx = (dim_idx * 2) % rotary_dim;
int half_dim = dim_idx / 2;
float cos_val = CacheConverter::convert(__ldg(cos_ptr + half_dim));
float sin_val = CacheConverter::convert(__ldg(sin_ptr + half_dim));
elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
}
// __shfl_xor_sync does not provide memfence. Need to sync again.
__syncwarp();
}
}
// Store.
{
vec_T vec;
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
#pragma unroll
for (int i = 0; i < num_packed_elems; i++) {
// Convert from float2 back to the specific packed type
float2 vals = {elements[2 * i], elements[2 * i + 1]};
T2_in packed_val = Converter::convert(vals);
// Place it into the generic vector
*(reinterpret_cast<T2_in*>(&vec) + i) = packed_val;
}
*reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
}
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
}
#endif
}
// Borrowed from
// https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
if (interleave) { \
const bool INTERLEAVE = true; \
__VA_ARGS__ \
} else { \
const bool INTERLEAVE = false; \
__VA_ARGS__ \
}
template <typename scalar_t_in, typename scalar_t_cache>
void launch_fused_qknorm_rope(void* qkv,
int const num_tokens,
int const num_heads_q,
int const num_heads_k,
int const num_heads_v,
int const head_dim,
int const rotary_dim,
float const eps,
void const* q_weight,
void const* k_weight,
void const* cos_sin_cache,
bool const interleave,
int64_t const* position_ids,
cudaStream_t stream) {
constexpr int blockSize = 256;
int const warpsPerBlock = blockSize / 32;
int const totalQKHeads = num_heads_q + num_heads_k;
int const totalWarps = num_tokens * totalQKHeads;
int const gridSize = div_up(totalWarps, warpsPerBlock);
dim3 gridDim(gridSize);
dim3 blockDim(blockSize);
switch (head_dim) {
case 64:
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>
<<<gridDim, blockDim, 0, stream>>>(qkv,
num_heads_q,
num_heads_k,
num_heads_v,
eps,
q_weight,
k_weight,
cos_sin_cache,
position_ids,
num_tokens,
rotary_dim);
});
break;
case 128:
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>
<<<gridDim, blockDim, 0, stream>>>(qkv,
num_heads_q,
num_heads_k,
num_heads_v,
eps,
q_weight,
k_weight,
cos_sin_cache,
position_ids,
num_tokens,
rotary_dim);
});
break;
case 256:
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>
<<<gridDim, blockDim, 0, stream>>>(qkv,
num_heads_q,
num_heads_k,
num_heads_v,
eps,
q_weight,
k_weight,
cos_sin_cache,
position_ids,
num_tokens,
rotary_dim);
});
break;
default:
CHECK(false) << "Unsupported head dimension for fusedQKNormRope: "
<< head_dim;
}
}
} // namespace
namespace xllm::kernel::cuda {
void fused_qk_norm_rope(
torch::Tensor& qkv, // Combined QKV tensor [num_tokens,
// (num_heads_q+num_heads_k+num_heads_v)*head_dim]
int64_t num_heads_q, // Number of query heads
int64_t num_heads_k, // Number of key heads
int64_t num_heads_v, // Number of value heads
int64_t head_dim, // Dimension per head
double eps, // Epsilon for RMS normalization
const torch::Tensor& q_weight, // RMSNorm weights for query [head_dim]
const torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
const torch::Tensor&
cos_sin_cache, // Cos/sin cache [max_position, rotary_dim]
bool interleaved, // Whether RoPE is applied in interleaved style
const torch::Tensor& position_ids // Position IDs for RoPE [num_tokens]
) {
// Input validation
CHECK(qkv.is_cuda()) << "qkv must be a CUDA tensor";
CHECK(qkv.is_contiguous()) << "qkv must be contiguous";
CHECK(position_ids.is_cuda()) << "position_ids must be a CUDA tensor";
CHECK(position_ids.is_contiguous()) << "position_ids must be contiguous";
CHECK(q_weight.is_cuda()) << "q_weight must be a CUDA tensor";
CHECK(q_weight.is_contiguous()) << "q_weight must be contiguous";
CHECK(k_weight.is_cuda()) << "k_weight must be a CUDA tensor";
CHECK(k_weight.is_contiguous()) << "k_weight must be contiguous";
CHECK(cos_sin_cache.is_cuda()) << "cos_sin_cache must be a CUDA tensor";
CHECK(cos_sin_cache.is_contiguous()) << "cos_sin_cache must be contiguous";
CHECK(position_ids.scalar_type() == torch::kInt64)
<< "position_ids dtype is " << position_ids.scalar_type()
<< ", while Int64 is expected";
CHECK(qkv.dim() == 2) << "QKV tensor must be 2D: [num_tokens, "
<< "(num_heads_q+num_heads_k+num_heads_v)*head_dim]";
CHECK(position_ids.dim() == 1) << "Position IDs must be 1D: [num_tokens]";
CHECK(q_weight.dim() == 1) << "Query weights must be 1D: [head_dim]";
CHECK(k_weight.dim() == 1) << "Key weights must be 1D: [head_dim]";
CHECK(cos_sin_cache.dim() == 2)
<< "Cos/sin cache must be 2D: [max_position, rotary_dim]";
CHECK(q_weight.size(0) == head_dim)
<< "Query weights size must match head dimension";
CHECK(k_weight.size(0) == head_dim)
<< "Key weights size must match head dimension";
CHECK(cos_sin_cache.size(1) % 2 == 0) << "rotary_dim must be even";
CHECK(cos_sin_cache.size(1) <= head_dim)
<< "rotary_dim must be less than or equal to head_dim";
CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
qkv.scalar_type() == k_weight.scalar_type())
<< "qkv, q_weight and k_weight must have the same dtype";
int64_t num_tokens = qkv.size(0);
CHECK(position_ids.size(0) == num_tokens)
<< "Number of tokens in position_ids must match QKV";
int64_t total_heads = num_heads_q + num_heads_k + num_heads_v;
CHECK(qkv.size(1) == total_heads * head_dim)
<< "QKV tensor size must match total number of heads and head dimension";
const at::cuda::OptionalCUDAGuard device_guard(device_of(qkv));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
using qkv_scalar_t = scalar_t;
DISPATCH_FLOATING_TYPES(
cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
using cache_scalar_t = scalar_t;
launch_fused_qknorm_rope<qkv_scalar_t, cache_scalar_t>(
qkv.data_ptr(),
static_cast<int>(num_tokens),
static_cast<int>(num_heads_q),
static_cast<int>(num_heads_k),
static_cast<int>(num_heads_v),
static_cast<int>(head_dim),
static_cast<int>(cos_sin_cache.size(1)),
static_cast<float>(eps),
q_weight.data_ptr(),
k_weight.data_ptr(),
cos_sin_cache.data_ptr(),
interleaved,
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
stream);
});
});
}
} // namespace xllm::kernel::cuda

108
ex_engine/xllm_kernels/cuda/headers/arch_condition.h Normal file
View File

@@ -0,0 +1,108 @@
/*
* Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// refers to
// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/kernels/archCondition.h
#pragma once
namespace xllm::kernel::cuda {
namespace detail {
#ifdef __CUDA_ARCH__
// __CUDA_ARCH_SPECIFIC__ is only available starting from CUDA 12.9
#if (__CUDACC_VER_MAJOR__ > 12 || \
(__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
#define HAS_CUDA_SPECIFIC_MACRO 1
#if __CUDA_ARCH__ >= 900
#if !defined(__CUDA_ARCH_SPECIFIC__) && !defined(__CUDA_ARCH_FAMILY_SPECIFIC__)
#error \
"Compiling for SM90 or newer architectures must use Arch specific or Arch Family specific target"
#endif
#endif
#else
#define HAS_CUDA_SPECIFIC_MACRO 0
#endif
// For CUDA < 12.9, we assume that sm90 or newer architectures are always built
// with arch specific.
#if defined(__CUDA_ARCH_SPECIFIC__) || \
(!HAS_CUDA_SPECIFIC_MACRO && __CUDA_ARCH__ >= 900)
static constexpr bool isArchSpecific = true;
#else
static constexpr bool isArchSpecific = false;
#endif
struct arch_info {
static constexpr bool mIsDevice = true;
static constexpr bool mArchSpecific = isArchSpecific;
static constexpr int mMajor = __CUDA_ARCH__ / 100;
static constexpr int mMinor = __CUDA_ARCH__ / 10 % 10;
static constexpr int mArch = __CUDA_ARCH__ / 10;
};
#else
struct arch_info {
static constexpr bool mIsDevice = false;
static constexpr bool mArchSpecific = false;
static constexpr int mMajor = 0;
static constexpr int mMinor = 0;
static constexpr int mArch = 0;
};
#endif
} // namespace detail
namespace arch {
struct is_device : std::bool_constant<detail::arch_info::mIsDevice> {};
struct is_arch_specific : std::bool_constant<detail::arch_info::mArchSpecific> {
};
template <int Arch>
struct is_match
: std::bool_constant<is_device::value && detail::arch_info::mArch == Arch> {
};
template <int Major>
struct is_major : std::bool_constant<is_device::value &&
detail::arch_info::mMajor == Major> {};
template <int Arch>
struct is_compatible : std::bool_constant<is_major<Arch>::value &&
detail::arch_info::mArch >= Arch> {};
inline constexpr bool is_device_v = is_device::value;
inline constexpr bool is_arch_specific_v = is_arch_specific::value;
template <int Arch>
inline constexpr bool is_match_v = is_match<Arch>::value;
template <int Major>
inline constexpr bool is_major_v = is_major<Major>::value;
template <int Arch>
inline constexpr bool is_compatible_v = is_compatible<Arch>::value;
} // namespace arch
} // namespace xllm::kernel::cuda

37
ex_engine/xllm_kernels/cuda/headers/corex_compat_utils.h Normal file
View File

@@ -0,0 +1,37 @@
// corex_compat_utils.h — Lightweight replacement for xllm's utils.h
// Removes glog/tvm dependencies for BI-V100 corex compilation
// Provides CHECK macro via TORCH_CHECK and DISPATCH macros from device_utils.cuh
#pragma once
#include <torch/torch.h>
#include <c10/cuda/CUDAGuard.h>
// Replace glog CHECK with TORCH_CHECK
#ifndef CHECK
#define CHECK(cond) TORCH_CHECK(cond)
#endif
#ifndef CHECK_EQ
#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b))
#endif
#ifndef CHECK_GE
#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b))
#endif
// Include device_utils for DISPATCH_HALF_TYPES etc
#include "device_utils.cuh"
// ffi namespace stub (some headers reference it)
namespace ffi {
template <typename T>
using Array = std::vector<T>;
}
// HOST_DEVICE_INLINE
#if defined(__CUDACC__) || defined(_NVHPC_CUDA)
#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
#else
#define HOST_DEVICE_INLINE inline
#endif

306
ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h Normal file
View File

@@ -0,0 +1,306 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <ATen/DynamicLibrary.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <glog/logging.h>
#include <optional>
#include <tuple>
#include <vector>
#include "utils.h"
namespace xllm::kernel::cuda {
// TODO: add head_size parameter
void rotary_embedding(torch::Tensor& positions,
torch::Tensor& query,
std::optional<torch::Tensor> key,
torch::Tensor& cos_sin_cache,
// int64_t head_size,
bool is_neox);
// 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 slot_ids, // [n_tokens]
torch::Tensor keys, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor values, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor key_cache, // [n_blocks, block_size, n_heads, head_dim]
torch::Tensor value_cache);
void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor value_cache_ptrs,
torch::Tensor src_block_indices,
torch::Tensor dst_block_indices,
torch::Tensor cum_sum,
int64_t numel_per_block,
torch::ScalarType cache_dtype);
#if !defined(USE_DCU)
void batch_prefill(const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor key,
torch::Tensor value,
torch::Tensor q_cu_seq_lens,
torch::Tensor kv_cu_seq_lens,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& mask = std::nullopt);
// Wrapper function for batch_prefill that conditionally uses AttentionRunner
// for piecewise CUDA Graph capture
void batch_prefill_with_optional_piecewise_capture(
const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor key,
torch::Tensor value,
torch::Tensor q_cu_seq_lens,
torch::Tensor kv_cu_seq_lens,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse);
void batch_prefill_non_causal(
const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor key,
torch::Tensor value,
torch::Tensor q_cu_seq_lens,
torch::Tensor kv_cu_seq_lens,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& mask = std::nullopt);
void batch_chunked_prefill(
const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor k_cache,
torch::Tensor v_cache,
torch::Tensor paged_kv_indptr,
torch::Tensor paged_kv_indices,
torch::Tensor paged_kv_last_page_len,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
std::optional<torch::Tensor> qo_indptr = std::nullopt,
bool causal = true);
void batch_decode(const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor k_cache,
torch::Tensor v_cache,
torch::Tensor paged_kv_indptr,
torch::Tensor paged_kv_indices,
torch::Tensor paged_kv_last_page_len,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
bool use_tensor_core,
std::optional<torch::Tensor> qo_indptr = std::nullopt);
#endif // !defined(USE_DCU)
void rms_norm(torch::Tensor output,
torch::Tensor input,
torch::Tensor weight,
double eps);
void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
torch::Tensor& residual, // [..., hidden_size]
torch::Tensor& weight, // [hidden_size]
double epsilon);
torch::Tensor matmul(torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias);
void cutlass_scaled_mm(torch::Tensor& c,
torch::Tensor const& a,
torch::Tensor const& b,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
std::optional<torch::Tensor> const& bias);
// Static scaled FP8 quantization
// Quantizes input tensor to FP8 using a pre-computed scale factor
void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
torch::Tensor const& input, // [..., d]
torch::Tensor const& scale); // [1]
// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format
// Returns: (quantized_output, scale)
std::tuple<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
const torch::Tensor& input,
const std::optional<torch::Tensor>& output = std::nullopt,
const std::optional<torch::Tensor>& scale = std::nullopt);
// ============================================================================
// Fused RMSNorm + Static FP8 Quantization
// ============================================================================
// These functions combine RMSNorm and FP8 quantization to reduce memory
// bandwidth by avoiding the intermediate write-back to global memory.
// Fused RMSNorm + Static FP8 Quantization (without residual)
// Combines RMSNorm normalization and FP8 quantization in a single kernel.
// This is optimal for the first layer where no residual connection exists.
void rms_norm_static_fp8_quant(
torch::Tensor& out, // [..., hidden_size], FP8 output
torch::Tensor& input, // [..., hidden_size], input tensor
torch::Tensor& weight, // [hidden_size], RMSNorm weight
torch::Tensor& scale, // [1], FP8 quantization scale
double epsilon); // RMSNorm epsilon
// Fused Add + RMSNorm + Static FP8 Quantization (with residual)
// Combines residual addition, RMSNorm, and FP8 quantization in a single kernel.
// The residual tensor is updated in-place with the sum of input and residual.
void fused_add_rms_norm_static_fp8_quant(
torch::Tensor& out, // [..., hidden_size], FP8 output
torch::Tensor& input, // [..., hidden_size], input tensor
torch::Tensor& residual, // [..., hidden_size], residual (updated in-place)
torch::Tensor& weight, // [hidden_size], RMSNorm weight
torch::Tensor& scale, // [1], FP8 quantization scale
double epsilon); // RMSNorm epsilon
// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels
// Performs: c = (a @ b.T) with scales applied
torch::Tensor fp8_scaled_matmul(
const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& a_scale,
const torch::Tensor& b_scale,
torch::ScalarType output_dtype,
const std::optional<torch::Tensor>& bias = std::nullopt,
const std::optional<torch::Tensor>& output = std::nullopt);
std::pair<torch::Tensor, torch::Tensor> compute_topk_for_beam_search(
torch::Tensor combined_probs,
uint32_t batch_size,
uint32_t beam_size,
uint32_t top_k,
torch::Device device);
std::pair<torch::Tensor, torch::Tensor> compute_topk_general(
torch::Tensor input,
uint32_t batch_size,
uint32_t input_length,
uint32_t k,
torch::Device device);
torch::Tensor air_log_softmax_last_dim(const torch::Tensor& input,
const torch::Tensor& temperatures);
void fused_qk_norm_rope(
torch::Tensor& qkv, // Combined QKV tensor [num_tokens,
// (num_heads_q+num_heads_k+num_heads_v)*head_dim]
int64_t num_heads_q, // Number of query heads
int64_t num_heads_k, // Number of key heads
int64_t num_heads_v, // Number of value heads
int64_t head_dim, // Dimension per head
double eps, // Epsilon for RMS normalization
const torch::Tensor& q_weight, // RMSNorm weights for query [head_dim]
const torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
const torch::Tensor&
cos_sin_cache, // Cos/sin cache [max_position, rotary_dim]
bool interleaved, // Whether RoPE is applied in interleaved style
const torch::Tensor& position_ids // Position IDs for RoPE [num_tokens]
);
std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
torch::Tensor& gating_output,
int64_t topk,
bool renormalize,
const std::optional<torch::Tensor>& correction_bias,
const std::string& scoring_func);
torch::Tensor random_sample(const torch::Tensor& probs);
torch::Tensor cutlass_fused_moe(
const torch::Tensor& input, // [num_tokens, hidden]
const torch::Tensor& token_selected_experts, // [num_tokens, top_k]
const torch::Tensor& token_final_scales, // [num_tokens, top_k]
const torch::Tensor&
fc1_expert_weights, // [num_experts, inter_dim, hidden]
const torch::Tensor&
fc2_expert_weights, // [num_experts, hidden, inter_dim]
torch::ScalarType output_dtype,
const std::vector<torch::Tensor>& quant_scales,
int32_t tp_size,
int32_t tp_rank,
int32_t ep_size,
int32_t ep_rank,
int32_t cluster_size,
int32_t cluster_rank,
const std::optional<torch::Tensor>& fc1_expert_biases = std::nullopt,
const std::optional<torch::Tensor>& fc2_expert_biases = std::nullopt,
const std::optional<torch::Tensor>& input_sf = std::nullopt,
const std::optional<torch::Tensor>& swiglu_alpha = std::nullopt,
const std::optional<torch::Tensor>& swiglu_beta = std::nullopt,
const std::optional<torch::Tensor>& swiglu_limit = std::nullopt,
const std::optional<torch::Tensor>& output = std::nullopt,
bool enable_alltoall = false,
bool use_deepseek_fp8_block_scale = false,
bool use_w4_group_scaling = false,
bool use_mxfp8_act_scaling = false,
bool min_latency_mode = false,
bool use_packed_weights = false,
int32_t tune_max_num_tokens = 8192,
ActivationType activation_type = ActivationType::SWIGLU);
// ---- moe_compute_index (moe_compute_index.cu) ----
// Fused routing index: bincount + argsort replacement.
// Returns {src_dst, dst_src, expert_sizes}.
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> moe_compute_index(
const torch::Tensor& expert_id,
int64_t num_experts);
// ---- moe_combine_result (moe_combine.cu) ----
// Fused combine: reorder + weighted sum in one pass.
torch::Tensor moe_combine_result(const torch::Tensor& gemm2,
const torch::Tensor& reduce_weight,
int64_t N,
int32_t topk);
} // namespace xllm::kernel::cuda

150
ex_engine/xllm_kernels/cuda/headers/device_utils.cuh Normal file
View File

@@ -0,0 +1,150 @@
/* Copyright 2025-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
#if defined(USE_DCU)
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hipcub/hipcub.hpp>
namespace cub = hipcub;
#else
#include <cub/cub.cuh>
#if CUB_VERSION >= 200800
#include <cuda/functional>
#endif
#endif
namespace xllm::kernel::cuda {
#if !defined(USE_DCU)
using BFloat16Type = __nv_bfloat16;
#define WARP_SIZE 32
#define XLLM_KERNEL_ATTR(MAX_THREADS)
#else
using BFloat16Type = hip_bfloat16;
#define WARP_SIZE 64
#define XLLM_KERNEL_ATTR(MAX_THREADS) __launch_bounds__(MAX_THREADS, 1)
#endif
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
// Aligned array type
template <typename T,
// Number of elements in the array
int N,
// Alignment requirement in bytes
int Alignment = sizeof(T) * N>
class alignas(Alignment) AlignedArray {
T data[N];
};
#define XLLM_SHFL_XOR_SYNC(mask, var, lane_mask) \
__shfl_xor_sync((mask), (var), (lane_mask))
#define XLLM_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \
__shfl_xor_sync((mask), (var), (lane_mask), (width))
template <typename T>
__device__ __forceinline__ T xllm_ldg(const T* ptr) {
#if defined(USE_DCU)
return *ptr;
#else
return __ldg(ptr);
#endif
}
// Define reduction operators based on CUB version.
#if defined(USE_DCU)
using MaxReduceOp = hipcub::Max;
using MinReduceOp = hipcub::Min;
#elif CUB_VERSION >= 200800
using MaxReduceOp = ::cuda::maximum<>;
using MinReduceOp = ::cuda::minimum<>;
#else
using MaxReduceOp = cub::Max;
using MinReduceOp = cub::Min;
#endif
template <typename T>
__device__ float convert_to_float(T x) {
if constexpr (std::is_same_v<T, __half>) {
return __half2float(x);
#if defined(USE_DCU)
} else if constexpr (std::is_same_v<T, hip_bfloat16>) {
return __bfloat162float(reinterpret_cast<const __hip_bfloat16&>(x));
#else
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
return __bfloat162float(x);
#endif
} else if constexpr (std::is_same_v<T, float>) {
return x;
} else {
return static_cast<float>(x);
}
}
// Constructs some constants needed to partition the work across threads at
// compile time.
template <typename T, int EXPERTS, int BYTES_PER_LDG>
struct TopkConstants {
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 ||
EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0,
"");
static constexpr int VECs_PER_THREAD =
MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
};
} // namespace xllm::kernel::cuda
// ============================================================================
// Portable macros and utilities (from xllm/core/kernels/cuda/utils.h)
// ============================================================================
#ifndef DEVICE_INLINE
#define DEVICE_INLINE __device__ __forceinline__
#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
#endif
template <typename T>
HOST_DEVICE_INLINE constexpr std::enable_if_t<std::is_integral_v<T>, T>
ceil_div(T a, T b) {
return (a + b - 1) / b;
}
// ============================================================================
// Dispatch macros (from xllm/core/kernels/cuda/utils.h)
// These wrap AT_DISPATCH_SWITCH for float16/bfloat16/float32 dispatch.
// Placed here because cuda_ops_api.h → utils.h is not available on corex
// (glog/logging.h dependency).
// ============================================================================
#ifndef DISPATCH_FLOATING_TYPES
#define DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#define DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
#define DISPATCH_CASE_HALF_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#define DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
#endif

239
ex_engine/xllm_kernels/cuda/headers/fp8_quant_utils.cuh Normal file
View File

@@ -0,0 +1,239 @@
/* Copyright 2025-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
// clang-format off
#include <c10/util/Float8_e4m3fn.h>
#include <cmath>
#include <torch/types.h>
// clang-format on
namespace xllm {
namespace kernel {
namespace cuda {
// FP8 type max value definitions
template <typename T,
typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
std::is_same_v<T, int8_t>>>
struct quant_type_max {
static constexpr T val() { return std::numeric_limits<T>::max(); }
};
template <typename T>
__host__ __device__ static constexpr T quant_type_max_v =
quant_type_max<T>::val();
// Minimum scaling factor for quantization types
template <typename T,
typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
std::is_same_v<T, int8_t>>>
struct min_scaling_factor {
__device__ __host__ static inline float val() {
return 1.0f / (quant_type_max_v<T> * 512.0f);
}
};
template <>
struct min_scaling_factor<int8_t> {
__device__ __host__ static inline float val() {
return std::numeric_limits<float>::epsilon();
}
};
// Vectorization containers
template <typename scalar_t, size_t vec_size>
struct __align__(vec_size * sizeof(scalar_t)) vec_n_t {
scalar_t val[vec_size];
};
template <typename quant_type_t, size_t vec_size>
struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t {
static_assert(std::is_same_v<quant_type_t, int8_t> ||
std::is_same_v<quant_type_t, c10::Float8_e4m3fn>);
quant_type_t val[vec_size];
};
// Atomic max for float
__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
float old;
old = (value >= 0)
? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
: __uint_as_float(
atomicMin((unsigned int*)addr, __float_as_uint(value)));
return old;
}
// FP8 conversion functions
namespace fp8 {
#ifdef ENABLE_FP8
#include <cuda_fp8.h>
// float -> c10::Float8_e4m3fn conversion
template <typename Tout, typename Tin>
__inline__ __device__ Tout
vec_conversion(const Tin& x,
const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) {
return x;
}
template <>
__inline__ __device__ c10::Float8_e4m3fn
vec_conversion<c10::Float8_e4m3fn, float>(
const float& a,
const __nv_fp8_interpretation_t fp8_type) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
return static_cast<c10::Float8_e4m3fn>(a);
#else
return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(a, __NV_SATFINITE, fp8_type),
c10::Float8_e4m3fn::from_bits());
#endif
}
#endif // ENABLE_FP8
} // namespace fp8
// Scaled FP8 conversion with saturation
template <bool is_scale_inverted, typename fp8_type>
__device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val,
float const scale) {
float x = 0.0f;
if constexpr (is_scale_inverted) {
x = val * scale;
} else {
x = val / scale;
}
float r =
fmaxf(-quant_type_max_v<fp8_type>, fminf(x, quant_type_max_v<fp8_type>));
#ifdef ENABLE_FP8
// Use hardware cvt instruction for fp8 on nvidia
return fp8::vec_conversion<fp8_type, float>(r);
#else
return static_cast<fp8_type>(r);
#endif
}
// Vectorization utilities
template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
struct DefaultVecOp {
ScaOp scalar_op;
__device__ __forceinline__ void operator()(
vec_n_t<OutT, VEC_SIZE>& dst,
const vec_n_t<InT, VEC_SIZE>& src) const {
#pragma unroll
for (int i = 0; i < VEC_SIZE; ++i) {
scalar_op(dst.val[i], src.val[i]);
}
}
};
template <int VEC_SIZE,
typename InT,
typename OutT,
typename VecOp,
typename ScaOp>
__device__ inline void vectorize_with_alignment(
const InT* in,
OutT* out,
int len,
int tid,
int stride,
VecOp&& vec_op, // vec_n_t<InT,16> -> vec_n_t<OutT,16>
ScaOp&& scalar_op) { // InT -> OutT
static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
"VEC_SIZE must be a positive power-of-two");
constexpr int WIDTH = VEC_SIZE * sizeof(InT);
uintptr_t addr = reinterpret_cast<uintptr_t>(in);
// Fast path when the whole region is already aligned
bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
if (can_vec) {
int num_vec = len / VEC_SIZE;
using vin_t = vec_n_t<InT, VEC_SIZE>;
using vout_t = vec_n_t<OutT, VEC_SIZE>;
auto* v_in = reinterpret_cast<const vin_t*>(in);
auto* v_out = reinterpret_cast<vout_t*>(out);
for (int i = tid; i < num_vec; i += stride) {
vout_t tmp;
vin_t src = v_in[i];
vec_op(tmp, src);
v_out[i] = tmp;
}
return;
}
int misalignment_offset = addr & (WIDTH - 1);
int alignment_bytes = WIDTH - misalignment_offset;
int prefix_elems = alignment_bytes & (WIDTH - 1);
prefix_elems /= sizeof(InT);
prefix_elems = min(prefix_elems, len);
// Prefix handling
for (int i = tid; i < prefix_elems; i += stride) {
scalar_op(out[i], in[i]);
}
in += prefix_elems;
out += prefix_elems;
len -= prefix_elems;
int num_vec = len / VEC_SIZE;
using vin_t = vec_n_t<InT, VEC_SIZE>;
using vout_t = vec_n_t<OutT, VEC_SIZE>;
auto* v_in = reinterpret_cast<const vin_t*>(in);
auto* v_out = reinterpret_cast<vout_t*>(out);
// Vectorized main part
for (int i = tid; i < num_vec; i += stride) {
vout_t tmp;
vin_t src = v_in[i];
vec_op(tmp, src);
v_out[i] = tmp;
}
// Tail handling
int tail_start = num_vec * VEC_SIZE;
for (int i = tid + tail_start; i < len; i += stride) {
scalar_op(out[i], in[i]);
}
}
template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
__device__ __forceinline__ void vectorize_with_alignment(const InT* in,
OutT* out,
int len,
int tid,
int stride,
ScaOp&& scalar_op) {
using Vec = DefaultVecOp<VEC_SIZE, InT, OutT, std::decay_t<ScaOp>>;
vectorize_with_alignment<VEC_SIZE>(in,
out,
len,
tid,
stride,
Vec{scalar_op},
std::forward<ScaOp>(scalar_op));
}
} // namespace cuda
} // namespace kernel
} // namespace xllm

2114
ex_engine/xllm_kernels/cuda/headers/topk_last_dim.cuh Normal file
View File

File diff suppressed because it is too large Load Diff

231
ex_engine/xllm_kernels/cuda/headers/type_convert.cuh Normal file
View File

@@ -0,0 +1,231 @@
/* 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.
==============================================================================*/
#pragma once
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <torch/all.h>
// ref to:
// https://github.com/vllm-project/vllm/blob/main/csrc/type_convert.cuh
/* Converter helpers for the conversion from torch types to HIP/CUDA types,
and the associated type conversions within HIP/CUDA. These helpers need
to be implemented for now because the relevant type conversion
operators/constructors are not consistently implemented by HIP/CUDA, so
a generic conversion via type casts cannot be implemented.
Each helper should have the member static constexpr bool `exists`:
If false, the optimized kernel is not used for the corresponding torch type.
If true, the helper should be fully defined as shown in the examples below.
*/
namespace xllm::kernel::cuda {
template <typename torch_type>
class _typeConvert {
public:
static constexpr bool exists = false;
};
template <>
class _typeConvert<float> {
public:
static constexpr bool exists = true;
using hip_type = float;
using packed_hip_type = float2;
using packed_hip_type4 = float4; // For 128-bit vectorization
__device__ static __forceinline__ float convert(hip_type x) { return x; }
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return x;
}
__device__ static __forceinline__ float4 convert(packed_hip_type4 x) {
return x;
}
};
#if defined(USE_DCU) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) || \
defined(USE_MACA)
// CUDA < 12.0 runs into issues with packed type conversion
template <>
class _typeConvert<c10::Half> {
public:
static constexpr bool exists = true;
using hip_type = __half;
using packed_hip_type = __half2;
__device__ static __forceinline__ float convert(hip_type x) {
return __half2float(x);
}
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return __half22float2(x);
}
__device__ static __forceinline__ hip_type convert(float x) {
return __float2half_rn(x);
}
__device__ static __forceinline__ packed_hip_type convert(float2 x) {
return __float22half2_rn(x);
}
};
#endif // defined(USE_DCU) || CUDA_VERSION >= 12000
#if defined(USE_DCU)
template <>
class _typeConvert<c10::BFloat16> {
public:
static constexpr bool exists = true;
using hip_type = __hip_bfloat16;
using packed_hip_type = __hip_bfloat162;
__device__ static __forceinline__ float convert(hip_type x) {
return __bfloat162float(x);
}
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return __bfloat1622float2(x);
}
__device__ static __forceinline__ hip_type convert(float x) {
return __float2bfloat16(x);
}
__device__ static __forceinline__ packed_hip_type convert(float2 x) {
return __float22bfloat162_rn(x);
}
};
#elif defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) && \
defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) || \
defined(USE_MACA)
// CUDA_ARCH < 800 does not have BF16 support.
template <>
class _typeConvert<c10::BFloat16> {
public:
static constexpr bool exists = true;
using hip_type = __nv_bfloat16;
using packed_hip_type = __nv_bfloat162;
__device__ static __forceinline__ float convert(hip_type x) {
return __bfloat162float(x);
}
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return __bfloat1622float2(x);
}
__device__ static __forceinline__ hip_type convert(float x) {
return __float2bfloat16(x);
}
__device__ static __forceinline__ packed_hip_type convert(float2 x) {
return __float22bfloat162_rn(x);
}
};
#endif
/* Vector helper to generate vectorized and packed FP16/BF16 ops
for appropriate specializations of fused_add_rms_norm_kernel.
Only functions that are necessary in that kernel are implemented.
Alignment to 16 bytes is required to use 128-bit global memory ops.
*/
template <typename scalar_t, int width>
class alignas(16) _f16Vec {
public:
/* Not theoretically necessary that width is a power of 2 but should
almost always be the case for optimization purposes */
static_assert(width > 0 && (width & (width - 1)) == 0,
"Width is not a positive power of 2!");
using Converter = _typeConvert<scalar_t>;
using T1 = typename Converter::hip_type;
using T2 = typename Converter::packed_hip_type;
T1 data[width];
__device__ _f16Vec& operator+=(const _f16Vec<scalar_t, width>& other) {
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
if constexpr (std::is_same_v<T2, float2>) {
data[i] += other.data[i];
data[i + 1] += other.data[i + 1];
} else {
T2 temp{data[i], data[i + 1]};
temp += T2{other.data[i], other.data[i + 1]};
data[i] = temp.x;
data[i + 1] = temp.y;
}
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) data[i] += other.data[i];
}
return *this;
}
__device__ _f16Vec& operator*=(const _f16Vec<scalar_t, width>& other) {
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
if constexpr (std::is_same_v<T2, float2>) {
data[i] *= other.data[i];
data[i + 1] *= other.data[i + 1];
} else {
T2 temp{data[i], data[i + 1]};
temp *= T2{other.data[i], other.data[i + 1]};
data[i] = temp.x;
data[i + 1] = temp.y;
}
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) data[i] *= other.data[i];
}
return *this;
}
__device__ _f16Vec& operator*=(const float scale) {
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
float2 temp_f = Converter::convert(T2{data[i], data[i + 1]});
temp_f.x *= scale;
temp_f.y *= scale;
T2 temp = Converter::convert(temp_f);
data[i] = temp.x;
data[i + 1] = temp.y;
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) {
float temp = Converter::convert(data[i]) * scale;
data[i] = Converter::convert(temp);
}
}
return *this;
}
__device__ float sum_squares() const {
float result = 0.0f;
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
float2 z = Converter::convert(T2{data[i], data[i + 1]});
result += z.x * z.x + z.y * z.y;
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(data[i]);
result += x * x;
}
}
return result;
}
};
} // namespace xllm::kernel::cuda

163
ex_engine/xllm_kernels/cuda/headers/utils.h Normal file
View File

@@ -0,0 +1,163 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <ATen/DynamicLibrary.h>
#if defined(USE_DCU)
#include <c10/hip/HIPGuard.h>
#else
#include <c10/cuda/CUDAGuard.h>
#endif
#include <glog/logging.h>
#include <torch/torch.h>
#if !defined(USE_DCU)
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/optional.h>
#endif
#include <string>
#include <tuple>
#include <type_traits>
#include <unordered_map>
#if defined(__CUDACC__) || defined(_NVHPC_CUDA) || defined(__HIPCC__)
#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
#define DEVICE_INLINE __device__ __forceinline__
#define HOST_INLINE __host__ __forceinline__
#else
#define HOST_DEVICE_INLINE inline
#define DEVICE_INLINE inline
#define HOST_INLINE inline
#endif
#if !defined(USE_DCU)
namespace ffi = tvm::ffi;
#endif
namespace xllm::kernel::cuda {
template <typename T>
HOST_DEVICE_INLINE constexpr std::enable_if_t<std::is_integral_v<T>, T>
ceil_div(T a, T b) {
return (a + b - 1) / b;
}
enum class ActivationType : int8_t {
GELU = 0,
RELU = 1,
SILU = 2,
SWIGLU = 3,
GEGLU = 4,
SWIGLU_BIAS = 5,
RELU2 = 6,
IDENTITY = 7,
INVALID_TYPE = 8
};
// torch tensor is only on cpu
torch::Tensor get_cache_buffer(const int32_t seq_len,
const torch::Device& device);
// NOLINTBEGIN(cppcoreguidelines-macro-usage)
#define DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#define DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
#define DISPATCH_CASE_HALF_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#define DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
// NOLINTEND(cppcoreguidelines-macro-usage)
bool should_use_tensor_core(torch::ScalarType kv_cache_dtype,
int64_t num_attention_heads,
int64_t num_kv_heads);
bool support_pdl();
std::string path_to_uri_so_lib(const std::string& uri);
std::string determine_attention_backend(int64_t pos_encoding_mode,
bool use_fp16_qk_reduction,
bool use_custom_mask);
std::string get_batch_prefill_uri(const std::string& backend,
torch::ScalarType dtype_q,
torch::ScalarType dtype_kv,
torch::ScalarType dtype_o,
torch::ScalarType dtype_idx,
int64_t head_dim_qk,
int64_t head_dim_vo,
int64_t pos_encoding_mode,
bool use_sliding_window,
bool use_logits_soft_cap,
bool use_fp16_qk_reduction);
std::string get_batch_decode_uri(torch::ScalarType dtype_q,
torch::ScalarType dtype_kv,
torch::ScalarType dtype_o,
torch::ScalarType dtype_idx,
int64_t head_dim_qk,
int64_t head_dim_vo,
int64_t pos_encoding_mode,
bool use_sliding_window,
bool use_logits_soft_cap);
std::tuple<torch::Tensor, double> split_scale_param(const torch::Tensor& scale);
#if !defined(USE_DCU)
DLDataType to_dl_data_type(torch::ScalarType scalar_type);
// below are tvm-ffi related functions
ffi::Tensor to_ffi_tensor(const torch::Tensor& torch_tensor);
ffi::Optional<ffi::Tensor> to_ffi_optional_tensor(
const std::optional<torch::Tensor>& optional);
ffi::Array<ffi::Tensor> to_ffi_array_tensors(
const std::vector<torch::Tensor>& torch_tensors);
ffi::Optional<ffi::Array<ffi::Tensor>> to_ffi_optional_array_tensors(
const std::optional<std::vector<torch::Tensor>>& optional);
ffi::Module get_module(const std::string& uri);
ffi::Function get_function(const std::string& uri,
const std::string& func_name);
inline void bind_tvmffi_stream_to_current_torch_stream(
const torch::Device& device) {
const auto cur = c10::cuda::getCurrentCUDAStream(device.index());
// DLPack device type for CUDA is 2 (kDLCUDA).
void* original_stream = nullptr;
const int rc = TVMFFIEnvSetStream(
/*device_type=*/2,
/*device_id=*/device.index(),
reinterpret_cast<void*>(cur.stream()),
&original_stream);
if (rc != 0) {
LOG(WARNING) << "[tvmffi.stream] failed to set stream, rc=" << rc
<< " dev=" << device.index();
}
}
#endif // !defined(USE_DCU)
} // namespace xllm::kernel::cuda

167
ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu Normal file
View File

@@ -0,0 +1,167 @@
// hgemm_blocktiling.cu — FP16 GEMM for BI-V100
//
// 1:1 from siboehm/SGEMM_CUDA kernel 6 (sgemmVectorize).
// Changes: float→__half, float4→load 4 halfs, FP32 accumulator.
// No WARPSIZE usage. No cooperative_groups. CUDA 10.2 safe.
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
template <const int BM, const int BN, const int BK, const int TM, const int TN>
__global__ void hgemmVectorize(int M, int N, int K, float alpha,
const __half *A, const __half *B,
float beta, __half *C) {
const uint cRow = blockIdx.y;
const uint cCol = blockIdx.x;
// BN/TN are the number of threads to span a column
const int threadCol = threadIdx.x % (BN / TN);
const int threadRow = threadIdx.x / (BN / TN);
// allocate space for the current blocktile in smem
// A stored transposed: As[BK][BM], B normal: Bs[BK][BN]
__shared__ __half As[BM * BK];
__shared__ __half Bs[BK * BN];
// Move blocktile to beginning of A's row and B's column
A += cRow * BM * K;
B += cCol * BN;
C += cRow * BM * N + cCol * BN;
// calculating the indices that this thread will load into SMEM
// FP16: load 4 halfs (8 bytes) per step. 4 halfs per thread.
// siboehm: float4 = 4 floats = 128bit. We do 4 halfs = 64bit.
const uint innerRowA = threadIdx.x / (BK / 4);
const uint innerColA = threadIdx.x % (BK / 4);
const uint innerRowB = threadIdx.x / (BN / 4);
const uint innerColB = threadIdx.x % (BN / 4);
// allocate thread-local cache for results in registerfile
// FP32 accumulation to avoid FP16 precision loss
float threadResults[TM * TN] = {0.0f};
__half regM[TM];
__half regN[TN];
// outer-most loop over block tiles
for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
// populate the SMEM caches
// transpose A while loading it (same as siboehm)
// Load 4 halfs from A
__half a0 = A[innerRowA * K + innerColA * 4 + 0];
__half a1 = A[innerRowA * K + innerColA * 4 + 1];
__half a2 = A[innerRowA * K + innerColA * 4 + 2];
__half a3 = A[innerRowA * K + innerColA * 4 + 3];
As[(innerColA * 4 + 0) * BM + innerRowA] = a0;
As[(innerColA * 4 + 1) * BM + innerRowA] = a1;
As[(innerColA * 4 + 2) * BM + innerRowA] = a2;
As[(innerColA * 4 + 3) * BM + innerRowA] = a3;
// Load 4 halfs from B (no transpose)
Bs[innerRowB * BN + innerColB * 4 + 0] = B[innerRowB * N + innerColB * 4 + 0];
Bs[innerRowB * BN + innerColB * 4 + 1] = B[innerRowB * N + innerColB * 4 + 1];
Bs[innerRowB * BN + innerColB * 4 + 2] = B[innerRowB * N + innerColB * 4 + 2];
Bs[innerRowB * BN + innerColB * 4 + 3] = B[innerRowB * N + innerColB * 4 + 3];
__syncthreads();
// advance blocktile
A += BK; // move BK columns to right
B += BK * N; // move BK rows down
// calculate per-thread results
for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
// block into registers
for (uint i = 0; i < TM; ++i) {
regM[i] = As[dotIdx * BM + threadRow * TM + i];
}
for (uint i = 0; i < TN; ++i) {
regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
}
// FP32 accumulation
for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
float aVal = __half2float(regM[resIdxM]);
for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
threadResults[resIdxM * TN + resIdxN] +=
aVal * __half2float(regN[resIdxN]);
}
}
}
__syncthreads();
}
// write out the results
for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
for (uint resIdxN = 0; resIdxN < TN; resIdxN += 1) {
uint row = cRow * BM + threadRow * TM + resIdxM;
uint col = cCol * BN + threadCol * TN + resIdxN;
if (row < M && col < N) {
float c_old = __half2float(C[(threadRow * TM + resIdxM) * N +
threadCol * TN + resIdxN]);
C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN] =
__float2half(alpha * threadResults[resIdxM * TN + resIdxN] +
beta * c_old);
}
}
}
}
// ============================================================================
// Launch wrapper — matches siboehm runSgemmVectorize
// ============================================================================
void launch_hgemm_blocktiling(
int M, int N, int K,
const __half* alpha_ptr,
const __half* A, int lda,
const __half* B, int ldb,
const __half* beta_ptr,
__half* C, int ldc,
cudaStream_t stream)
{
constexpr int BM = 128;
constexpr int BN = 128;
constexpr int BK = 8;
constexpr int TM = 8;
constexpr int TN = 8;
// 256 threads — same as siboehm
constexpr int NUM_THREADS = (BM * BN) / (TM * TN);
dim3 grid(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
dim3 block(NUM_THREADS);
float alpha = 1.0f, beta = 0.0f;
if (alpha_ptr) alpha = __half2float(*alpha_ptr);
if (beta_ptr) beta = __half2float(*beta_ptr);
hgemmVectorize<BM, BN, BK, TM, TN>
<<<grid, block, 0, stream>>>(M, N, K, alpha, A, B, beta, C);
}
// ============================================================================
// MoE expert GEMM — C++ loop over experts (replaces Python for-loop)
// ============================================================================
void launch_moe_expert_hgemm(
int num_experts,
const int* expert_counts, // host, [num_experts]
const int* expert_offsets, // host, [num_experts]
int N, int K,
const __half* input, // (total_tokens, K)
const __half* weights, // (num_experts, N, K)
__half* output, // (total_tokens, N)
cudaStream_t stream)
{
for (int e = 0; e < num_experts; e++) {
int M_e = expert_counts[e];
if (M_e == 0) continue;
int off = expert_offsets[e];
const __half* A = input + off * K;
const __half* B = weights + (long long)e * N * K;
__half* C_e = output + off * N;
launch_hgemm_blocktiling(M_e, N, K,
nullptr, A, K, B, N, nullptr, C_e, N, stream);
}
}

199
ex_engine/xllm_kernels/cuda/hgemm_warptiling.cu Normal file
View File

@@ -0,0 +1,199 @@
// hgemm_warptiling.cu — FP16 warp-tiling GEMM for BI-V100 (warp_size=64)
//
// 1:1 from siboehm/SGEMM_CUDA kernel 10 (sgemmWarptiling).
// Changes from original:
// 1. WARPSIZE = 32 → 64 (BI-V100 confirmed)
// 2. float → __half for A/B/C data and shared memory
// 3. float4 vectorized load → 4 scalar __half loads
// 4. threadResults accumulator stays float (FP32 accumulation)
// 5. C writeback: scalar instead of float4
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
const int WARPSIZE = 64; // BI-V100 confirmed
namespace wt {
template <const int BM, const int BN, const int BK, const int rowStrideA,
const int rowStrideB>
__device__ void loadFromGmem(int N, int K, const __half *A, const __half *B,
__half *As, __half *Bs, int innerRowA, int innerColA,
int innerRowB, int innerColB) {
for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
// Load 4 halfs from A, transpose while storing
__half a0 = A[(innerRowA + offset) * K + innerColA * 4 + 0];
__half a1 = A[(innerRowA + offset) * K + innerColA * 4 + 1];
__half a2 = A[(innerRowA + offset) * K + innerColA * 4 + 2];
__half a3 = A[(innerRowA + offset) * K + innerColA * 4 + 3];
As[(innerColA * 4 + 0) * BM + innerRowA + offset] = a0;
As[(innerColA * 4 + 1) * BM + innerRowA + offset] = a1;
As[(innerColA * 4 + 2) * BM + innerRowA + offset] = a2;
As[(innerColA * 4 + 3) * BM + innerRowA + offset] = a3;
}
for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
// Load 4 halfs from B, no transpose
Bs[(innerRowB + offset) * BN + innerColB * 4 + 0] =
B[(innerRowB + offset) * N + innerColB * 4 + 0];
Bs[(innerRowB + offset) * BN + innerColB * 4 + 1] =
B[(innerRowB + offset) * N + innerColB * 4 + 1];
Bs[(innerRowB + offset) * BN + innerColB * 4 + 2] =
B[(innerRowB + offset) * N + innerColB * 4 + 2];
Bs[(innerRowB + offset) * BN + innerColB * 4 + 3] =
B[(innerRowB + offset) * N + innerColB * 4 + 3];
}
}
template <const int BM, const int BN, const int BK, const int WM, const int WN,
const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
const int TM, const int TN>
__device__ void
processFromSmem(float *regM, float *regN, float *threadResults, const __half *As,
const __half *Bs, const uint warpRow, const uint warpCol,
const uint threadRowInWarp, const uint threadColInWarp) {
for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
// populate registers for whole warptile
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
for (uint i = 0; i < TM; ++i) {
regM[wSubRowIdx * TM + i] = __half2float(
As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
threadRowInWarp * TM + i]);
}
}
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
for (uint i = 0; i < TN; ++i) {
regN[wSubColIdx * TN + i] = __half2float(
Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
threadColInWarp * TN + i]);
}
}
// execute warptile matmul — FP32 accumulation
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
(wSubColIdx * TN) + resIdxN] +=
regM[wSubRowIdx * TM + resIdxM] *
regN[wSubColIdx * TN + resIdxN];
}
}
}
}
}
}
} // namespace wt
template <const int BM, const int BN, const int BK, const int WM, const int WN,
const int WNITER, const int TM, const int TN, const int NUM_THREADS>
__global__ void __launch_bounds__(NUM_THREADS)
hgemmWarptiling(int M, int N, int K, float alpha, const __half *A,
const __half *B, float beta, __half *C) {
const uint cRow = blockIdx.y;
const uint cCol = blockIdx.x;
// Placement of the warp in the threadblock tile
const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
const uint warpCol = warpIdx % (BN / WN);
const uint warpRow = warpIdx / (BN / WN);
// size of the warp subtile
constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
constexpr uint WSUBM = WM / WMITER;
constexpr uint WSUBN = WN / WNITER;
// Placement of the thread in the warp subtile
const uint threadIdxInWarp = threadIdx.x % WARPSIZE; // [0, 63]
const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN);
const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN);
// allocate space for the current blocktile in SMEM
__shared__ __half As[BM * BK];
__shared__ __half Bs[BK * BN];
// Move blocktile to beginning of A's row and B's column
A += cRow * BM * K;
B += cCol * BN;
// Move C_ptr to warp's output tile
C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
// calculating the indices that this thread will load into SMEM
// FP16: 4 halfs per thread per step
const uint innerRowA = threadIdx.x / (BK / 4);
const uint innerColA = threadIdx.x % (BK / 4);
constexpr uint rowStrideA = (NUM_THREADS * 4) / BK;
const uint innerRowB = threadIdx.x / (BN / 4);
const uint innerColB = threadIdx.x % (BN / 4);
constexpr uint rowStrideB = NUM_THREADS / (BN / 4);
// allocate thread-local cache for results in registerfile
float threadResults[WMITER * TM * WNITER * TN] = {0.0f};
// we cache into registers on the warptile level
float regM[WMITER * TM] = {0.0f};
float regN[WNITER * TN] = {0.0f};
// outer-most loop over block tiles
for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
wt::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB);
__syncthreads();
wt::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
TN>(regM, regN, threadResults, As, Bs, warpRow, warpCol,
threadRowInWarp, threadColInWarp);
A += BK; // move BK columns to right
B += BK * N; // move BK rows down
__syncthreads();
}
// write out the results — scalar writeback (no float4 for __half)
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
__half *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
for (uint resIdxN = 0; resIdxN < TN; resIdxN += 1) {
uint idx = (threadRowInWarp * TM + resIdxM) * N +
threadColInWarp * TN + resIdxN;
float c_old = __half2float(C_interim[idx]);
const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
wSubColIdx * TN + resIdxN;
C_interim[idx] = __float2half(alpha * threadResults[i] + beta * c_old);
}
}
}
}
}
// ============================================================================
// Launch wrapper
// ============================================================================
void launch_hgemm_warptiling(
int M, int N, int K,
float alpha,
const __half* A,
const __half* B,
float beta,
__half* C,
cudaStream_t stream)
{
// Config B — best on BI-V100 (beats cublas 0.7x on 256x4096@4096x11008):
// probe_k10_configs.sh confirmed: 7.6ms vs cublas 10.5ms
// 128 threads = 2 warps of 64
// WMITER = (64*64)/(64*8*4*2) = 4096/4096 = 1
// WSUBM = 64/1 = 64, WSUBN = 64/2 = 32
// threads_per_warp = (64/8)*(32/4) = 8*8 = 64 ✓
constexpr int NUM_THREADS = 128;
constexpr int BM = 128, BN = 128, BK = 16;
constexpr int WM = 64, WN = 64;
constexpr int WNITER = 2;
constexpr int TM = 8, TN = 4;
dim3 grid(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
dim3 block(NUM_THREADS);
hgemmWarptiling<BM, BN, BK, WM, WN, WNITER, TM, TN, NUM_THREADS>
<<<grid, block, 0, stream>>>(M, N, K, alpha, A, B, beta, C);
}

124
ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp Normal file
View File

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

105
ex_engine/xllm_kernels/cuda/moe/moe_combine.cu Executable file
View File

@@ -0,0 +1,105 @@
/* Copyright 2025-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.
==============================================================================*/
// Fused MoE combine kernel — reorder + weighted sum in one pass.
// Replaces: torch::zeros + index_copy_ + view + multiply + sum
//
// Algorithm per token (each block handles one token):
// 1. For each of its topk experts, read gemm2 at flat_idx directly
// (gemm2 is flat-index-ordered after scatter via index_copy_ with dst_src)
// 2. Multiply by router weight
// 3. Accumulate into output[token]
//
// Grid: num_tokens (N) blocks
// Block: HIDDEN_DIM / HIDDEN_TILE threads
#include <c10/cuda/CUDAGuard.h>
#include "device_utils.cuh"
#include <torch/extension.h>
namespace xllm::kernel::cuda {
constexpr int32_t kCombineBlockSize = 256;
template <typename scalar_t>
__global__ void XLLM_KERNEL_ATTR(kCombineBlockSize) moe_combine_kernel(
const scalar_t* __restrict__ gemm2, // [N*topk, H] flat-index-ordered
const float* __restrict__ reduce_weight, // [N, topk]
scalar_t* __restrict__ output, // [N, H]
int64_t N,
int32_t topk,
int64_t H) {
int64_t token_id = blockIdx.x; // 0 .. N-1
if (token_id >= N) return;
int32_t tid = threadIdx.x;
int32_t stride = kCombineBlockSize;
// Accumulate over topk experts for this token
for (int64_t h = tid; h < H; h += stride) {
float acc = 0.0f;
for (int32_t k = 0; k < topk; ++k) {
int64_t flat_idx = token_id * topk + k;
float w = reduce_weight[flat_idx];
acc += w * static_cast<float>(gemm2[flat_idx * H + h]);
}
output[token_id * H + h] = static_cast<scalar_t>(acc);
}
}
// ---- Host-side orchestrator ----
torch::Tensor moe_combine_result(
const torch::Tensor& gemm2, // [N*topk, H] flat-index-ordered
const torch::Tensor& reduce_weight, // [N, topk] float or same as gemm2
int64_t N,
int32_t topk) {
auto stream = at::cuda::getCurrentCUDAStream();
int64_t H = gemm2.size(1);
auto dtype = gemm2.scalar_type();
auto output = torch::empty({N, H}, gemm2.options());
auto rw = reduce_weight.to(gemm2.device(), torch::kFloat32).contiguous();
if (dtype == torch::kFloat16) {
moe_combine_kernel<c10::Half>
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<c10::Half>(),
rw.data_ptr<float>(),
output.data_ptr<c10::Half>(),
N,
topk,
H);
} else if (dtype == torch::kBFloat16) {
moe_combine_kernel<c10::BFloat16>
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<c10::BFloat16>(),
rw.data_ptr<float>(),
output.data_ptr<c10::BFloat16>(),
N,
topk,
H);
} else {
moe_combine_kernel<float>
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<float>(),
rw.data_ptr<float>(),
output.data_ptr<float>(),
N,
topk,
H);
}
return output;
}
} // namespace xllm::kernel::cuda

156
ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu Normal file
View File

@@ -0,0 +1,156 @@
/* Copyright 2025-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.
==============================================================================*/
// Fused MoE token index computation — 3 kernels replacing:
// torch::bincount + 2 × torch::argsort + torch::cumsum + CPU sync
//
// Phase 1 histogram: atomicAdd per-expert token counts
// Phase 2 prefix_sum: 1 block, exclusive scan → expert_offsets
// Phase 3 place_indices: atomicAdd on offsets, write dst_src + src_dst
//
// expert_sizes = per-expert token count [num_experts] (preserved)
// expert_offsets = exclusive prefix sum of counts (scratch, reused)
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <cub/block/block_scan.cuh>
#include "device_utils.cuh"
namespace xllm::kernel::cuda {
constexpr int32_t kMoeIndexBlock = 256;
// ---- Phase 1: histogram ----
__global__ void
#ifdef USE_DCU
__launch_bounds__(kMoeIndexBlock, 1)
#endif
moe_histogram_kernel(const int32_t* __restrict__ expert_id,
int32_t* __restrict__ expert_sizes,
int64_t num_elements,
int32_t num_experts) {
int64_t tid = int64_t(blockIdx.x) * kMoeIndexBlock + threadIdx.x;
if (tid < num_elements) {
int32_t eid = expert_id[tid];
if (eid >= 0 && eid < num_experts) {
atomicAdd(&expert_sizes[eid], 1);
}
}
}
// ---- Phase 2: exclusive prefix sum (1 block) ----
// input: expert_sizes (per-expert counts)
// output: expert_offsets (exclusive scan of counts)
// total_out (total number of tokens, scalar)
__global__ void
#ifdef USE_DCU
__launch_bounds__(kMoeIndexBlock, 1)
#endif
moe_prefix_sum_kernel(const int32_t* __restrict__ expert_sizes,
int32_t* __restrict__ expert_offsets,
int32_t num_experts,
int64_t* __restrict__ total_out) {
using BlockScan = cub::BlockScan<int32_t, kMoeIndexBlock>;
__shared__ typename BlockScan::TempStorage s_scan;
int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0;
int32_t offset;
BlockScan(s_scan).ExclusiveSum(val, offset);
__syncthreads();
// total = all elements sum = last thread's exclusive output + its input
int32_t total = offset + val;
if (threadIdx.x < num_experts) {
expert_offsets[threadIdx.x] = offset;
}
if (threadIdx.x == 0 && total_out != nullptr) {
*total_out = total;
}
}
// ---- Phase 3: place indices ----
// atomicAdd on expert_offsets to assign a unique position within
// [start(e), start(e)+count(e)), then write both direction mappings.
__global__ void
#ifdef USE_DCU
__launch_bounds__(kMoeIndexBlock, 1)
#endif
moe_place_indices_kernel(const int32_t* __restrict__ expert_id,
int32_t* __restrict__ expert_offsets,
int32_t* __restrict__ dst_src,
int32_t* __restrict__ src_dst,
int64_t num_elements,
int32_t num_experts) {
int64_t flat_idx = int64_t(blockIdx.x) * kMoeIndexBlock + threadIdx.x;
if (flat_idx >= num_elements) return;
int32_t eid = expert_id[flat_idx];
if (eid < 0 || eid >= num_experts) return;
int32_t pos = atomicAdd(&expert_offsets[eid], 1);
dst_src[pos] = static_cast<int32_t>(flat_idx);
src_dst[flat_idx] = pos;
}
// ---- Host-side orchestrator ----
// Returns {src_dst, dst_src, expert_sizes}
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> moe_compute_index(
const torch::Tensor& expert_id,
int64_t num_experts) {
auto device = expert_id.device();
auto stream = at::cuda::getCurrentCUDAStream();
int64_t N = expert_id.numel();
int32_t E = static_cast<int32_t>(num_experts);
TORCH_CHECK(E <= kMoeIndexBlock, "num_experts cannot exceed ", kMoeIndexBlock);
auto expert_id_i32 = expert_id.to(torch::kInt32).contiguous();
auto opt_i32 = expert_id_i32.options();
auto expert_sizes = torch::zeros({num_experts}, opt_i32);
auto expert_offsets = torch::empty({num_experts}, opt_i32);
auto dst_src = torch::empty({N}, opt_i32);
auto src_dst = torch::empty({N}, opt_i32);
int64_t grid = (N + kMoeIndexBlock - 1) / kMoeIndexBlock;
// Phase 1: histogram
moe_histogram_kernel<<<grid, kMoeIndexBlock, 0, stream>>>(
expert_id_i32.data_ptr<int32_t>(),
expert_sizes.data_ptr<int32_t>(),
N,
E);
// Phase 2: prefix sum (1 block)
moe_prefix_sum_kernel<<<1, kMoeIndexBlock, 0, stream>>>(
expert_sizes.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
E,
nullptr);
// Phase 3: place indices
moe_place_indices_kernel<<<grid, kMoeIndexBlock, 0, stream>>>(
expert_id_i32.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
dst_src.data_ptr<int32_t>(),
src_dst.data_ptr<int32_t>(),
N,
E);
return std::make_tuple(src_dst, dst_src, expert_sizes);
}
} // namespace xllm::kernel::cuda

59
ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu Normal file
View File

@@ -0,0 +1,59 @@
/* Copyright 2025-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.
==============================================================================*/
#if defined(USE_DCU)
#include "kernels/dcu/dcu_ops_api.h"
#else
#include "device_utils.cuh"
#include <torch/extension.h>
#endif
#include "moe_topk_sigmoid_kernels.cuh"
#include "moe_topk_softmax_kernels.cuh"
namespace xllm::kernel::cuda {
std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
torch::Tensor& gating_output,
int64_t topk,
bool renormalize,
const std::optional<torch::Tensor>& correction_bias,
const std::string& scoring_func) {
int64_t num_tokens = gating_output.size(0);
torch::Tensor topk_weights = torch::empty(
{num_tokens, topk},
torch::dtype(torch::kFloat32).device(gating_output.device()));
torch::Tensor topk_ids =
torch::empty({num_tokens, topk},
torch::dtype(torch::kInt32).device(gating_output.device()));
if (scoring_func == "softmax") {
std::optional<torch::Tensor> none_correction_bias = std::nullopt;
topk_softmax(topk_weights,
topk_ids,
gating_output,
renormalize,
/*moe_softcapping=*/0.0,
none_correction_bias);
} else if (scoring_func == "sigmoid") {
topk_sigmoid(
topk_weights, topk_ids, gating_output, renormalize, correction_bias);
} else {
TORCH_CHECK(false, "Unsupported scoring function: ", scoring_func);
}
return std::make_tuple(topk_weights, topk_ids);
}
} // namespace xllm::kernel::cuda

345
ex_engine/xllm_kernels/cuda/moe/moe_topk.cuh Normal file
View File

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

608
ex_engine/xllm_kernels/cuda/moe/moe_topk_sigmoid_kernels.cuh Normal file
View File

@@ -0,0 +1,608 @@
// Adapt from
// https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu
// which is originally adapted from
// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <cub/util_type.cuh>
#if !defined(USE_DCU) && !defined(USE_MACA)
#endif
#include "device_utils.cuh"
namespace {
using namespace xllm::kernel::cuda;
#if defined(USE_DCU)
static constexpr unsigned long long kSigmoidFullMask = 0xffffffffffffffffULL;
#else
static constexpr unsigned int kSigmoidFullMask = 0xffffffffU;
#endif
// ====================== Sigmoid things ===============================
// We have our own implementation of sigmoid here so we can support transposing
// the output in the sigmoid kernel when we extend this module to support
// expert-choice routing.
template <typename T, int TPB>
__launch_bounds__(TPB) __global__
void moe_sigmoid(const T* input,
const bool* finished,
float* output,
const int num_cols,
const float* correction_bias) {
const int thread_row_offset = blockIdx.x * num_cols;
// Don't touch finished rows.
if ((finished != nullptr) && finished[blockIdx.x]) {
return;
}
// First pass: Apply transformation, find max, and write transformed values to
// output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
float val = convert_to_float<T>(input[idx]);
val = 1.0f / (1.0f + expf(-val));
// Apply correction bias if provided
if (correction_bias != nullptr) {
val = val + correction_bias[ii];
}
output[idx] = val; // Store transformed value
}
}
template <int TPB>
__launch_bounds__(TPB) __global__
void moe_topK(const float* inputs_after_sigmoid,
const bool* finished,
float* output,
int* indices,
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float* correction_bias) {
using cub_kvp = cub::KeyValuePair<int, float>;
using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
cub_kvp thread_kvp;
cub::ArgMax arg_max;
const int block_row = blockIdx.x;
const bool row_is_active = finished ? !finished[block_row] : true;
const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0;
for (int k_idx = 0; k_idx < k; ++k_idx) {
thread_kvp.key = 0;
thread_kvp.value = -1.f; // This is OK because inputs are probabilities
cub_kvp inp_kvp;
for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
const int idx = thread_read_offset + expert;
inp_kvp.key = expert;
inp_kvp.value = inputs_after_sigmoid[idx];
for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
const int prior_winning_expert = indices[k * block_row + prior_k];
if (prior_winning_expert == expert) {
inp_kvp = thread_kvp;
}
}
thread_kvp = arg_max(inp_kvp, thread_kvp);
}
const cub_kvp result_kvp =
BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
if (threadIdx.x == 0) {
// Ignore experts the node isn't responsible for with expert parallelism
const int expert = result_kvp.key;
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const int idx = k * block_row + k_idx;
float val = result_kvp.value;
if (correction_bias != nullptr) {
val -= correction_bias[expert];
}
output[idx] = val;
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
row_sum_for_renormalize += val;
}
__syncthreads();
}
if (renormalize && threadIdx.x == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * block_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
// ====================== TopK sigmoid things ===============================
/*
A Top-K gating sigmoid written to exploit when the number of experts in the
MoE layers are a small power of 2. This allows us to cleanly share the rows
among the threads in a single warp and eliminate communication between warps
(so no need to use shared mem).
It fuses the sigmoid, max and argmax into a single kernel.
Limitations:
1) This implementation is intended for when the number of experts is a small
power of 2. 2) This implementation assumes k is small, but will work for any
k.
*/
template <typename T,
int VPT,
int NUM_EXPERTS,
int WARPS_PER_CTA,
int BYTES_PER_LDG>
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
void topk_gating_sigmoid(const T* input,
const bool* finished,
float* output,
const int num_rows,
int* indices,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float* correction_bias) {
// We begin by enforcing compile time assertions and setting up compile time
// constants.
static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
"NUM_EXPERTS must be power of 2");
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
"BYTES_PER_LDG must be power of 2");
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load
static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section.
static_assert(
VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE % kThreadsPerRow == 0,
"The threads per row must cleanly divide the threads per warp");
static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
"THREADS_PER_ROW must be power of 2");
static_assert(kThreadsPerRow <= 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 kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section.
static_assert(kEltsPerWarp % kEltsPerRow == 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 * kRowsPerCta;
// 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 * kRowsPerWarp;
// 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 / kThreadsPerRow;
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 * kEltsPerRow;
// 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 % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
// Determine the pointer type to use to read in the data depending on the
// BYTES_PER_LDG template param. In theory, this can support all powers of 2
// up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
// array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS.
using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem
T row_chunk_temp[VPT];
AccessType* row_chunk_vec_ptr =
reinterpret_cast<AccessType*>(&row_chunk_temp);
const AccessType* vec_thread_read_ptr =
reinterpret_cast<const AccessType*>(thread_read_ptr);
#pragma unroll
// Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ...
for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
}
float row_chunk[VPT];
#pragma unroll
// Note(Byron): upcast logits to float32
for (int ii = 0; ii < VPT; ++ii) {
float val = convert_to_float<T>(row_chunk_temp[ii]);
val = 1.0f / (1.0f + expf(-val));
// Apply correction bias if provided
if (correction_bias != nullptr) {
/*
LDG is interleaved
|thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG|
|--------- group0 --------| |----------group1 --------|
^ local2
*/
const int group_id = ii / kEltsPerLdg;
const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread +
group_id * kThreadsPerRow * kEltsPerLdg + 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 kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
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 < kLdgPerThread;
++ldg, col += kColsPerGroupLdg) {
#pragma unroll
for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * kEltsPerLdg + 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 = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
kSigmoidFullMask, max_val, mask, kThreadsPerRow);
int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
kSigmoidFullMask, expert, mask, kThreadsPerRow);
// 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 / kColsPerGroupLdg;
const int thread_to_clear_in_group =
(expert / kEltsPerLdg) % kThreadsPerRow;
// 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 % kEltsPerLdg;
// Safe to set to any negative value since row_chunk values must be
// between 0 and 1.
row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
-10000.f;
}
}
}
// Fuse renormalization of topk_weights into this kernel
if (renormalize && thread_group_idx == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * thread_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
template <typename T, int EXPERTS, int WARPS_PER_TB>
void topk_gating_sigmoid_launcher_helper(const T* input,
const bool* finished,
float* output,
int* indices,
const int num_rows,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float* correction_bias,
cudaStream_t stream) {
static constexpr std::size_t kMaxBytesPerLdg = 16;
static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
static constexpr int kVpt = Constants::VPT;
static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
topk_gating_sigmoid<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
<<<num_blocks, block_dim, 0, stream>>>(input,
finished,
output,
num_rows,
indices,
k,
start_expert,
end_expert,
renormalize,
correction_bias);
}
#define LAUNCH_SIGMOID(TYPE, NUM_EXPERTS, WARPS_PER_TB) \
topk_gating_sigmoid_launcher_helper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
gating_output, \
nullptr, \
topk_weights, \
topk_indices, \
num_tokens, \
topk, \
0, \
num_experts, \
renormalize, \
correction_bias, \
stream);
template <typename T>
void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
float* topk_weights,
int* topk_indices,
float* sigmoid_workspace,
const int num_tokens,
const int num_experts,
const int topk,
const bool renormalize,
const float* correction_bias,
cudaStream_t stream) {
static constexpr int kWarpsPerTb = 4;
switch (num_experts) {
case 1:
LAUNCH_SIGMOID(T, 1, kWarpsPerTb);
break;
case 2:
LAUNCH_SIGMOID(T, 2, kWarpsPerTb);
break;
case 4:
LAUNCH_SIGMOID(T, 4, kWarpsPerTb);
break;
case 8:
LAUNCH_SIGMOID(T, 8, kWarpsPerTb);
break;
case 16:
LAUNCH_SIGMOID(T, 16, kWarpsPerTb);
break;
case 32:
LAUNCH_SIGMOID(T, 32, kWarpsPerTb);
break;
case 64:
LAUNCH_SIGMOID(T, 64, kWarpsPerTb);
break;
case 128:
LAUNCH_SIGMOID(T, 128, kWarpsPerTb);
break;
case 256:
LAUNCH_SIGMOID(T, 256, kWarpsPerTb);
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 kTpb = 256;
moe_sigmoid<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr,
sigmoid_workspace,
num_experts,
correction_bias);
moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(sigmoid_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize,
correction_bias);
}
}
}
} // namespace
namespace xllm::kernel::cuda {
void topk_sigmoid(torch::Tensor& topk_weights, // [num_tokens, topk]
torch::Tensor& topk_indices, // [num_tokens, topk]
torch::Tensor& gating_output, // [num_tokens, num_experts]
const bool renormalize,
const std::optional<torch::Tensor>& correction_bias) {
// Check data type
CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
gating_output.scalar_type() == at::ScalarType::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16)
<< "gating_output must be float32, float16, or bfloat16";
// Check dimensions
CHECK(gating_output.dim() == 2)
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
CHECK(topk_weights.dim() == 2)
<< "topk_weights must be 2D tensor [num_tokens, topk]";
CHECK(topk_indices.dim() == 2)
<< "topk_indices must be 2D tensor [num_tokens, topk]";
// Check shapes
CHECK(gating_output.size(0) == topk_weights.size(0))
<< "First dimension of topk_weights must match num_tokens in "
"gating_output";
CHECK(gating_output.size(0) == topk_indices.size(0))
<< "First dimension of topk_indices must match num_tokens in "
"gating_output";
CHECK(topk_weights.size(-1) == topk_indices.size(-1))
<< "Second dimension of topk_indices must match topk in topk_weights";
CHECK(topk_weights.size(-1) <= gating_output.size(-1))
<< "topk must be less than or equal to num_experts";
const int num_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0));
const int topk = static_cast<int>(topk_weights.size(-1));
const bool is_pow_2 =
(num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
const bool needs_workspace = !is_pow_2 || num_experts > 256;
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
torch::Tensor sigmoid_workspace = torch::empty(
{workspace_size}, gating_output.options().dtype(at::ScalarType::Float));
const at::ScalarType dtype = gating_output.scalar_type();
// Validate correction_bias if provided - must always be float32
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
CHECK(bias_tensor.dim() == 1)
<< "correction_bias must be 1D tensor [num_experts]";
CHECK(bias_tensor.size(0) == num_experts)
<< "correction_bias size must match num_experts";
CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
bias_ptr = bias_tensor.data_ptr<float>();
}
if (dtype == at::ScalarType::Float) {
topk_gating_sigmoid_kernel_launcher<float>(
gating_output.data_ptr<float>(),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
sigmoid_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::Half) {
topk_gating_sigmoid_kernel_launcher<__half>(
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
sigmoid_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::BFloat16) {
topk_gating_sigmoid_kernel_launcher<BFloat16Type>(
reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
sigmoid_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
bias_ptr,
stream);
} else {
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
}
}
} // namespace xllm::kernel::cuda

866
ex_engine/xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh Normal file
View File

@@ -0,0 +1,866 @@
// Adapt from
// https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu
// which is originally adapted from
// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <cub/util_type.cuh>
#if !defined(USE_DCU) && !defined(USE_MACA)
#endif
#include "device_utils.cuh"
using cub_kvp = cub::KeyValuePair<int, float>;
namespace {
using namespace xllm::kernel::cuda;
#if defined(USE_DCU)
static constexpr unsigned long long kSoftmaxFullMask = 0xffffffffffffffffULL;
#else
static constexpr unsigned int kSoftmaxFullMask = 0xffffffffU;
#endif
// ====================== Softmax things ===============================
// We have our own implementation of softmax here so we can support transposing
// the output in the softmax kernel when we extend this module to support
// expert-choice routing.
template <typename T, int TPB>
__launch_bounds__(TPB) __global__
void moe_softmax(const T* input,
const bool* finished,
float* output,
const int num_cols,
const float moe_softcapping,
const float* correction_bias) {
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
__shared__ float normalizing_factor;
__shared__ float float_max;
const int thread_row_offset = blockIdx.x * num_cols;
float threadData(-FLT_MAX);
// Don't touch finished rows.
if ((finished != nullptr) && finished[blockIdx.x]) {
return;
}
// First pass: Apply transformation, find max, and write transformed values to
// output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
float val = convert_to_float<T>(input[idx]);
// Apply tanh softcapping if enabled
if (moe_softcapping != 0.0f) {
val = tanhf(val / moe_softcapping) * moe_softcapping;
}
// Apply correction bias if provided
if (correction_bias != nullptr) {
val = val + correction_bias[ii];
}
output[idx] = val; // Store transformed value
threadData = max(val, threadData);
}
const float maxElem =
BlockReduce(tmpStorage).Reduce(threadData, MaxReduceOp());
if (threadIdx.x == 0) {
float_max = maxElem;
}
__syncthreads();
// Second pass: Compute sum using transformed values from output
threadData = 0;
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
threadData += exp((output[idx] - float_max));
}
const auto Z = BlockReduce(tmpStorage).Sum(threadData);
if (threadIdx.x == 0) {
normalizing_factor = 1.f / Z;
}
__syncthreads();
// Third pass: Compute final softmax using transformed values from output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
const float softmax_val =
exp((output[idx] - float_max)) * normalizing_factor;
output[idx] = softmax_val;
}
}
namespace moe {
class TopKPair {
public:
static constexpr int kPair = 2;
static constexpr int kMaxIndex = 0;
cub_kvp max;
cub_kvp secondMax;
__device__ TopKPair() {}
__device__ TopKPair(cub_kvp max, cub_kvp secondMax)
: max(max), secondMax(secondMax) {}
};
class TopKPairArgMax {
public:
__device__ TopKPairArgMax() {}
__device__ __forceinline__ TopKPair
operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
cub_kvp globalMax, globalSecondMax;
// Determine the global maximum
if (candidate1.max.value > candidate2.max.value) {
globalMax = candidate1.max;
} else {
globalMax = candidate2.max;
}
// Determine the global second maximum
if (globalMax.key == candidate1.max.key) {
// If candidate1 contributed the max, compare its secondMax with
// candidate2's max
globalSecondMax = (candidate1.secondMax.value > candidate2.max.value)
? candidate1.secondMax
: candidate2.max;
} else {
// If candidate2 contributed the max, compare its secondMax with
// candidate1's max
globalSecondMax = (candidate2.secondMax.value > candidate1.max.value)
? candidate2.secondMax
: candidate1.max;
}
return TopKPair(globalMax, globalSecondMax);
}
};
} // namespace moe
template <int TPB>
__launch_bounds__(TPB) __global__
void moe_topk_fast(float* inputs_after_softmax,
const bool* finished,
float* output,
int* indices,
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize) {
using namespace moe;
using BlockReduce = cub::BlockReduce<TopKPair, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
TopKPair thread_pair;
const int block_row = blockIdx.x;
const bool row_is_active = finished ? !finished[block_row] : true;
const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0;
// Each loop finds the top 2 elements,
// thus requiring only ceil(k / 2) loops (calculated as (k + 1) / 2).
for (int k_idx = 0; k_idx < (k + TopKPair::kPair - 1) / TopKPair::kPair;
++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::kPair; i++) {
if (k_idx * 2 + i >= k) {
break;
}
cub_kvp result = (i == TopKPair::kMaxIndex) ? result_pair.max
: result_pair.secondMax;
int expert = result.key;
bool node_uses_expert = expert >= start_expert && expert < end_expert;
bool should_process_row = row_is_active && node_uses_expert;
// The inputs_after_softmax is modified in-place to avoid unnecessary
// loops for finding the top k-1 value. 1.f represents the minimum
// value.
inputs_after_softmax[thread_read_offset + expert] = -1.f;
int idx = k * block_row + k_idx * 2 + i;
output[idx] = result.value;
indices[idx] =
should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
row_sum_for_renormalize += result.value;
}
}
__syncthreads();
}
if (renormalize && threadIdx.x == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * block_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
template <int TPB>
__launch_bounds__(TPB) __global__ void moe_topK(float* inputs_after_softmax,
const bool* finished,
float* output,
int* indices,
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize) {
using cub_kvp = cub::KeyValuePair<int, float>;
using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
cub_kvp thread_kvp;
cub::ArgMax arg_max;
const int block_row = blockIdx.x;
const bool row_is_active = finished ? !finished[block_row] : true;
const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0;
for (int k_idx = 0; k_idx < k; ++k_idx) {
thread_kvp.key = 0;
thread_kvp.value = -1.f; // This is OK because inputs are probabilities
cub_kvp inp_kvp;
for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
const int idx = thread_read_offset + expert;
inp_kvp.key = expert;
inp_kvp.value = inputs_after_softmax[idx];
thread_kvp = arg_max(inp_kvp, thread_kvp);
}
const cub_kvp result_kvp =
BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
if (threadIdx.x == 0) {
// Ignore experts the node isn't responsible for with expert parallelism
const int expert = result_kvp.key;
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const int idx = k * block_row + k_idx;
output[idx] = result_kvp.value;
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
row_sum_for_renormalize += result_kvp.value;
// The inputs_after_softmax is modified in-place to avoid unnecessary
// loops for finding the top k-1 value. 1.f represents the minimum value.
inputs_after_softmax[thread_read_offset + expert] = -1.f;
}
__syncthreads();
}
if (renormalize && threadIdx.x == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * block_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
// ====================== TopK softmax things ===============================
/*
A Top-K gating softmax written to exploit when the number of experts in the
MoE layers are a small power of 2. This allows us to cleanly share the rows
among the threads in a single warp and eliminate communication between warps
(so no need to use shared mem).
It fuses the softmax, max and argmax into a single kernel.
Limitations:
1) This implementation is intended for when the number of experts is a small
power of 2. 2) This implementation assumes k is small, but will work for any
k.
*/
template <typename T,
int VPT,
int NUM_EXPERTS,
int WARPS_PER_CTA,
int BYTES_PER_LDG>
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
void topk_gating_softmax(const T* input,
const bool* finished,
float* output,
const int num_rows,
int* indices,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias) {
// We begin by enforcing compile time assertions and setting up compile time
// constants.
static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
"NUM_EXPERTS must be power of 2");
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
"BYTES_PER_LDG must be power of 2");
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load
static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section.
static_assert(
VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE % kThreadsPerRow == 0,
"The threads per row must cleanly divide the threads per warp");
static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
"THREADS_PER_ROW must be power of 2");
static_assert(kThreadsPerRow <= 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 kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section.
static_assert(kEltsPerWarp % kEltsPerRow == 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 * kRowsPerCta;
// 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 * kRowsPerWarp;
// 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 / kThreadsPerRow;
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 * kEltsPerRow;
// 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 % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
// Determine the pointer type to use to read in the data depending on the
// BYTES_PER_LDG template param. In theory, this can support all powers of 2
// up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
// array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS.
using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem
T row_chunk_temp[VPT];
AccessType* row_chunk_vec_ptr =
reinterpret_cast<AccessType*>(&row_chunk_temp);
const AccessType* vec_thread_read_ptr =
reinterpret_cast<const AccessType*>(thread_read_ptr);
#pragma unroll
// Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ...
for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
}
float row_chunk[VPT];
#pragma unroll
// Note(Byron): upcast logits to float32
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = convert_to_float<T>(row_chunk_temp[ii]);
}
// Apply tanh softcapping and correction bias
if (moe_softcapping != 0.0f || correction_bias != nullptr) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
// Apply tanh softcapping if enabled
if (moe_softcapping != 0.0f) {
val = tanhf(val / moe_softcapping) * moe_softcapping;
}
// Apply correction bias if provided
if (correction_bias != nullptr) {
/*
LDG is interleaved
|thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG|
|--------- group0 --------| |----------group1 --------|
^ local2
*/
const int group_id = ii / kEltsPerLdg;
const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread +
group_id * kThreadsPerRow * kEltsPerLdg +
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 = kThreadsPerRow / 2; mask > 0; mask /= 2) {
// butterfly reduce with (lane id ^ mask)
thread_max = max(thread_max,
XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, thread_max, mask, kThreadsPerRow));
}
// 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 = kThreadsPerRow / 2; mask > 0; mask /= 2) {
row_sum += XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, row_sum, mask, kThreadsPerRow);
}
// 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 kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
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 < kLdgPerThread;
++ldg, col += kColsPerGroupLdg) {
#pragma unroll
for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * kEltsPerLdg + 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 = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, max_val, mask, kThreadsPerRow);
int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, expert, mask, kThreadsPerRow);
// 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 / kColsPerGroupLdg;
const int thread_to_clear_in_group =
(expert / kEltsPerLdg) % kThreadsPerRow;
// 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 % kEltsPerLdg;
// Safe to set to any negative value since row_chunk values must be
// between 0 and 1.
row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
-10000.f;
}
}
}
// Fuse renormalization of topk_weights into this kernel
if (renormalize && thread_group_idx == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * thread_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
template <typename T, int EXPERTS, int WARPS_PER_TB>
void topk_gating_softmax_launcher_helper(const T* input,
const bool* finished,
float* output,
int* indices,
const int num_rows,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr std::size_t kMaxBytesPerLdg = 16;
static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
static constexpr int kVpt = Constants::VPT;
static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
topk_gating_softmax<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
<<<num_blocks, block_dim, 0, stream>>>(input,
finished,
output,
num_rows,
indices,
k,
start_expert,
end_expert,
renormalize,
moe_softcapping,
correction_bias);
}
#define LAUNCH_SOFTMAX(TYPE, NUM_EXPERTS, WARPS_PER_TB) \
topk_gating_softmax_launcher_helper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
gating_output, \
nullptr, \
topk_weights, \
topk_indices, \
num_tokens, \
topk, \
0, \
num_experts, \
renormalize, \
moe_softcapping, \
correction_bias, \
stream);
template <typename T>
void topk_gating_softmax_kernel_launcher(const T* gating_output,
float* topk_weights,
int* topk_indices,
float* softmax_workspace,
const int num_tokens,
const int num_experts,
const int topk,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr int kWarpsPerTb = 4;
switch (num_experts) {
case 1:
LAUNCH_SOFTMAX(T, 1, kWarpsPerTb);
break;
case 2:
LAUNCH_SOFTMAX(T, 2, kWarpsPerTb);
break;
case 4:
LAUNCH_SOFTMAX(T, 4, kWarpsPerTb);
break;
case 8:
LAUNCH_SOFTMAX(T, 8, kWarpsPerTb);
break;
case 16:
LAUNCH_SOFTMAX(T, 16, kWarpsPerTb);
break;
case 32:
LAUNCH_SOFTMAX(T, 32, kWarpsPerTb);
break;
case 64:
LAUNCH_SOFTMAX(T, 64, kWarpsPerTb);
break;
case 128:
LAUNCH_SOFTMAX(T, 128, kWarpsPerTb);
break;
case 256:
LAUNCH_SOFTMAX(T, 256, kWarpsPerTb);
break;
default: {
CHECK(softmax_workspace != nullptr)
<< "softmax_workspace must be provided for num_experts that are "
"not a power of 2.";
static constexpr int kTpb = 256;
moe_softmax<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr,
softmax_workspace,
num_experts,
moe_softcapping,
correction_bias);
if (topk == 1) {
// Note: As an optimization for better performance,
// the softmax_workspace is overwritten in-place by both moeTopK and
// moe_topk_fast.
moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
} else {
moe_topk_fast<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
}
}
}
}
} // namespace
namespace xllm::kernel::cuda {
void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
torch::Tensor& topk_indices, // [num_tokens, topk]
torch::Tensor& gating_output, // [num_tokens, num_experts]
const bool renormalize,
const double moe_softcapping,
const std::optional<torch::Tensor>& correction_bias) {
// Check data type
CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
gating_output.scalar_type() == at::ScalarType::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16)
<< "gating_output must be float32, float16, or bfloat16";
// Check dimensions
CHECK(gating_output.dim() == 2)
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
CHECK(topk_weights.dim() == 2)
<< "topk_weights must be 2D tensor [num_tokens, topk]";
CHECK(topk_indices.dim() == 2)
<< "topk_indices must be 2D tensor [num_tokens, topk]";
// Check shapes
CHECK(gating_output.size(0) == topk_weights.size(0))
<< "First dimension of topk_weights must match num_tokens in "
"gating_output"
<< "First dimension of topk_indices must match num_tokens in "
"gating_output";
CHECK(topk_weights.size(-1) == topk_indices.size(-1))
<< "Second dimension of topk_indices must match topk in topk_weights"
<< "topk must be less than or equal to num_experts";
const int num_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0));
const int topk = static_cast<int>(topk_weights.size(-1));
const bool is_pow_2 =
(num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
const bool needs_workspace = !is_pow_2 || num_experts > 256;
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
torch::Tensor softmax_workspace = torch::empty(
{workspace_size}, gating_output.options().dtype(at::ScalarType::Float));
const at::ScalarType dtype = gating_output.scalar_type();
// Validate correction_bias if provided - must always be float32
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
CHECK(bias_tensor.dim() == 1)
<< "correction_bias must be 1D tensor [num_experts]";
CHECK(bias_tensor.size(0) == num_experts)
<< "correction_bias size must match num_experts";
CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
bias_ptr = bias_tensor.data_ptr<float>();
}
// Cast moe_softcapping from double to float for CUDA kernels
const float moe_softcapping_f = static_cast<float>(moe_softcapping);
if (dtype == at::ScalarType::Float) {
topk_gating_softmax_kernel_launcher<float>(
gating_output.data_ptr<float>(),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
softmax_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::Half) {
topk_gating_softmax_kernel_launcher<__half>(
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
softmax_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::BFloat16) {
topk_gating_softmax_kernel_launcher<BFloat16Type>(
reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
softmax_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else {
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
}
}
} // namespace xllm::kernel::cuda

233
ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu Normal file
View File

@@ -0,0 +1,233 @@
// moe_cutlass_batched.cu — FP16 Cu10 TensorOp batched GEMM for MoE on BI-V100
//
// Adapted from corex-samples cutlass/examples/05_batched_gemm/batched_gemm.cu
// Changes from original:
// 1. float → cutlass::half_t (FP16 data)
// 2. arch::OpClassSimt → arch::OpClassTensorOp (use TCU)
// 3. arch::Sm61 → arch::Cu10 (BI-V100 arch)
// 4. ElementAccumulator = float (FP32 accumulation)
// 5. Row-major layout (PyTorch convention) instead of column-major
//
// Default Cu10 FP16 TensorOp config from default_gemm_configuration.h:
// ThreadblockShape = GemmShape<128, 128, 32>
// WarpShape = GemmShape<32, 32, 32>
// InstructionShape = GemmShape<16, 16, 16>
// kStages = 2
//
// This uses __ivcorex_matrix_mad_f32x4_f16x4 under the hood (via mma_cu10.h).
#include <iostream>
#include <vector>
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/gemm/device/gemm_batched.h"
// FP16 batched GEMM using Cu10 TensorOp
// C[i] = alpha * A[i] @ B[i] + beta * C[i]
// All matrices row-major, FP16 in/out, FP32 accumulation.
cudaError_t cutlass_batched_hgemm_tensorop(
int m, int n, int k,
float alpha,
cutlass::half_t const *A, int lda, long long int batch_stride_A,
cutlass::half_t const *B, int ldb, long long int batch_stride_B,
cutlass::half_t *C, int ldc, long long int batch_stride_C,
float beta,
int batch_count)
{
using Gemm = cutlass::gemm::device::GemmBatched<
cutlass::half_t, // ElementA
cutlass::layout::RowMajor, // LayoutA
cutlass::half_t, // ElementB
cutlass::layout::RowMajor, // LayoutB
cutlass::half_t, // ElementC
cutlass::layout::RowMajor, // LayoutC
float, // ElementAccumulator
cutlass::arch::OpClassTensorOp, // OperatorClass — use TCU
cutlass::arch::Cu10 // ArchTag — BI-V100
// Remaining params use defaults from DefaultGemmConfiguration:
// ThreadblockShape = <128, 128, 32>
// WarpShape = <32, 32, 32>
// InstructionShape = <16, 16, 16>
// Stages = 2
>;
Gemm gemm_op;
cutlass::Status status = gemm_op({
{m, n, k},
{A, lda},
batch_stride_A,
{B, ldb},
batch_stride_B,
{C, ldc},
batch_stride_C,
{C, ldc},
batch_stride_C,
{alpha, beta},
batch_count
});
if (status != cutlass::Status::kSuccess) {
return cudaErrorUnknown;
}
return cudaSuccess;
}
// ============================================================================
// Standalone test
// ============================================================================
#ifdef BUILD_STANDALONE_TEST
#include <cuda_fp16.h>
#include <cstdio>
#include <cstdlib>
#include <cmath>
int main() {
// Test: 8 batches of (1, 256) @ (256, 128) — simulates decode MoE
int m = 1, n = 128, k = 256;
int batch_count = 8;
float alpha = 1.0f, beta = 0.0f;
int lda = k; // row-major: (m, k), stride = k
int ldb = n; // row-major: (k, n), stride = n
int ldc = n; // row-major: (m, n), stride = n
long long int stride_A = (long long)m * k;
long long int stride_B = (long long)k * n;
long long int stride_C = (long long)m * n;
size_t size_A = batch_count * stride_A * sizeof(cutlass::half_t);
size_t size_B = batch_count * stride_B * sizeof(cutlass::half_t);
size_t size_C = batch_count * stride_C * sizeof(cutlass::half_t);
// Allocate host
std::vector<cutlass::half_t> h_A(batch_count * stride_A);
std::vector<cutlass::half_t> h_B(batch_count * stride_B);
std::vector<cutlass::half_t> h_C(batch_count * stride_C, cutlass::half_t(0.0f));
// Fill with small values
for (auto &v : h_A) v = cutlass::half_t(0.01f * (rand() % 100 - 50));
for (auto &v : h_B) v = cutlass::half_t(0.01f * (rand() % 100 - 50));
// Allocate device
cutlass::half_t *d_A, *d_B, *d_C;
cudaMalloc(&d_A, size_A);
cudaMalloc(&d_B, size_B);
cudaMalloc(&d_C, size_C);
cudaMemcpy(d_A, h_A.data(), size_A, cudaMemcpyHostToDevice);
cudaMemcpy(d_B, h_B.data(), size_B, cudaMemcpyHostToDevice);
cudaMemcpy(d_C, h_C.data(), size_C, cudaMemcpyHostToDevice);
// Run CUTLASS batched GEMM
cudaError_t result = cutlass_batched_hgemm_tensorop(
m, n, k, alpha,
d_A, lda, stride_A,
d_B, ldb, stride_B,
d_C, ldc, stride_C,
beta, batch_count);
cudaDeviceSynchronize();
if (result != cudaSuccess) {
printf("CUTLASS batched GEMM FAILED: %s\n", cudaGetErrorString(result));
cudaError_t last = cudaGetLastError();
if (last != cudaSuccess)
printf("Last CUDA error: %s\n", cudaGetErrorString(last));
cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
return -1;
}
// Copy back
cudaMemcpy(h_C.data(), d_C, size_C, cudaMemcpyDeviceToHost);
// Verify against CPU reference
bool pass = true;
for (int b = 0; b < batch_count; b++) {
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
float ref = 0.0f;
for (int p = 0; p < k; p++) {
float a_val = float(h_A[b * stride_A + i * k + p]);
float b_val = float(h_B[b * stride_B + p * n + j]);
ref += a_val * b_val;
}
float got = float(h_C[b * stride_C + i * n + j]);
if (fabs(ref - got) > 1.0f) {
printf("MISMATCH batch=%d [%d,%d]: ref=%.4f got=%.4f\n",
b, i, j, ref, got);
pass = false;
}
}
}
}
if (pass) {
printf("CUTLASS Cu10 TensorOp batched HGEMM: PASSED (%d batches of %dx%d@%dx%d)\n",
batch_count, m, k, k, n);
}
// Benchmark
cudaEvent_t t0, t1;
cudaEventCreate(&t0);
cudaEventCreate(&t1);
// Warmup
for (int i = 0; i < 5; i++)
cutlass_batched_hgemm_tensorop(m, n, k, alpha,
d_A, lda, stride_A, d_B, ldb, stride_B,
d_C, ldc, stride_C, beta, batch_count);
cudaDeviceSynchronize();
cudaEventRecord(t0);
for (int i = 0; i < 100; i++)
cutlass_batched_hgemm_tensorop(m, n, k, alpha,
d_A, lda, stride_A, d_B, ldb, stride_B,
d_C, ldc, stride_C, beta, batch_count);
cudaEventRecord(t1);
cudaEventSynchronize(t1);
float ms;
cudaEventElapsedTime(&ms, t0, t1);
printf("Perf: %.3f ms/iter (8 batches of 1x256 @ 256x128)\n", ms / 100.0f);
// Also test MoE-sized: 8 batches of (1, 4096) @ (4096, 11008)
int m2 = 1, n2 = 11008, k2 = 4096;
long long stride_A2 = (long long)m2 * k2;
long long stride_B2 = (long long)k2 * n2;
long long stride_C2 = (long long)m2 * n2;
cutlass::half_t *d_A2, *d_B2, *d_C2;
cudaMalloc(&d_A2, batch_count * stride_A2 * sizeof(cutlass::half_t));
cudaMalloc(&d_B2, batch_count * stride_B2 * sizeof(cutlass::half_t));
cudaMalloc(&d_C2, batch_count * stride_C2 * sizeof(cutlass::half_t));
for (int i = 0; i < 5; i++)
cutlass_batched_hgemm_tensorop(m2, n2, k2, alpha,
d_A2, k2, stride_A2, d_B2, n2, stride_B2,
d_C2, n2, stride_C2, beta, batch_count);
cudaDeviceSynchronize();
cudaEventRecord(t0);
for (int i = 0; i < 20; i++)
cutlass_batched_hgemm_tensorop(m2, n2, k2, alpha,
d_A2, k2, stride_A2, d_B2, n2, stride_B2,
d_C2, n2, stride_C2, beta, batch_count);
cudaEventRecord(t1);
cudaEventSynchronize(t1);
cudaEventElapsedTime(&ms, t0, t1);
printf("Perf: %.3f ms/iter (8 batches of 1x4096 @ 4096x11008 — MoE decode)\n", ms / 20.0f);
cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
cudaFree(d_A2); cudaFree(d_B2); cudaFree(d_C2);
cudaEventDestroy(t0);
cudaEventDestroy(t1);
return pass ? 0 : -1;
}
#endif // BUILD_STANDALONE_TEST

595
ex_engine/xllm_kernels/cuda/norm.cu Normal file
View File

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

600
ex_engine/xllm_kernels/cuda/norm.cu.orig Normal file
View File

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

102
ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu Normal file
View File

@@ -0,0 +1,102 @@
/* Copyright 2025-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 <c10/cuda/CUDAStream.h>
#include <torch/extension.h>
#include "device_utils.cuh"
namespace xllm::kernel::cuda {
template <typename T>
__global__ void XLLM_KERNEL_ATTR(1024) reshape_paged_cache_kernel(
const int* __restrict__ slot_ids, // [n_tokens]
const T* __restrict__ keys, // [n_tokens, n_heads, head_dim]
const T* __restrict__ values, // [n_tokens, n_heads, head_dim]
T* __restrict__ key_cache,
T* __restrict__ value_cache,
int64_t k_stride,
int64_t v_stride,
int64_t n_kv_heads,
int64_t head_dim,
int64_t block_size) {
// block/token index
const int64_t bid = blockIdx.x;
// which slot to write to
const int64_t slot_id = slot_ids[bid];
if (slot_id < 0) {
return;
}
// block index
const int64_t block_idx = slot_id / block_size;
// offset within block
const int64_t block_offset = slot_id % block_size;
// base index for the block in cache
const int64_t block_base_idx = block_idx * block_size * n_kv_heads * head_dim;
// copy value one by one for the token
for (int64_t i = threadIdx.x; i < n_kv_heads * head_dim; i += blockDim.x) {
const int64_t k_src_idx = bid * k_stride + i;
const int64_t v_src_idx = bid * v_stride + i;
// cache: [n_blocks, block_size, n_heads, head_dim]
const int64_t head_base_idx =
block_base_idx + block_offset * n_kv_heads * head_dim;
// which head to write to
const int head_idx = i / head_dim;
// which dim within head to write to
const int head_offset = i % head_dim;
const int64_t dst_idx = head_base_idx + head_idx * head_dim + head_offset;
key_cache[dst_idx] = keys[k_src_idx];
value_cache[dst_idx] = values[v_src_idx];
}
}
void reshape_paged_cache(
torch::Tensor slot_ids, // [n_tokens]
torch::Tensor keys, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor values, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor key_cache, // [n_blocks, block_size, n_heads, head_dim]
torch::Tensor value_cache) {
// keys and values should be continuous at n_kv_heads and head_dim dims
CHECK(keys.stride(-1) == 1 && keys.stride(-2) == keys.size(-1));
CHECK(values.stride(-1) == 1 && values.stride(-2) == values.size(-1));
const int64_t n_tokens = keys.size(-3);
const int64_t n_kv_heads = keys.size(-2);
const int64_t head_dim = keys.size(-1);
const int64_t block_size = key_cache.size(-3);
// it is possible that keys and values have different strides
const int64_t k_stride = keys.stride(-3);
const int64_t v_stride = values.stride(-3);
const int64_t n = n_kv_heads * head_dim;
dim3 grid(n_tokens);
dim3 block(std::min<int>(n, 1024));
DISPATCH_FLOATING_TYPES(
keys.scalar_type(), "reshape_paged_cache_kernel", [&] {
reshape_paged_cache_kernel<scalar_t>
<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
slot_ids.data_ptr<int>(),
keys.data_ptr<scalar_t>(),
values.data_ptr<scalar_t>(),
key_cache.data_ptr<scalar_t>(),
value_cache.data_ptr<scalar_t>(),
k_stride,
v_stride,
n_kv_heads,
head_dim,
block_size);
});
}
} // namespace xllm::kernel::cuda

258
ex_engine/xllm_kernels/cuda/rope.cu Normal file
View File

@@ -0,0 +1,258 @@
/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include "device_utils.cuh"
// ref to:
// https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu
namespace {
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_token_rotary_embedding(
scalar_t* __restrict__ arr,
const scalar_t* __restrict__ cos_ptr,
const scalar_t* __restrict__ sin_ptr,
int rot_offset,
int embed_dim) {
int x_index, y_index;
scalar_t cos, sin;
if (IS_NEOX) {
// GPT-NeoX style rotary embedding.
x_index = rot_offset;
y_index = embed_dim + rot_offset;
cos = *(cos_ptr + x_index);
sin = *(sin_ptr + x_index);
} else {
// GPT-J style rotary embedding.
x_index = 2 * rot_offset;
y_index = 2 * rot_offset + 1;
cos = *(cos_ptr + x_index / 2);
sin = *(sin_ptr + x_index / 2);
}
const scalar_t x = arr[x_index];
const scalar_t y = arr[y_index];
arr[x_index] = x * cos - y * sin;
arr[y_index] = y * cos + x * sin;
}
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_rotary_embedding(
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* cache_ptr,
const int head_size,
const int num_heads,
const int num_kv_heads,
const int rot_dim,
const int token_idx,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride) {
const int embed_dim = rot_dim / 2;
const scalar_t* cos_ptr = cache_ptr;
const scalar_t* sin_ptr = cache_ptr + embed_dim;
const int nq = num_heads * embed_dim;
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head =
token_idx * query_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
if (key != nullptr) {
const int nk = num_kv_heads * embed_dim;
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head =
token_idx * key_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
}
}
template <typename scalar_t, bool IS_NEOX>
__global__ void XLLM_KERNEL_ATTR(512) rotary_embedding_kernel(
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
// [num_tokens]
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2,
// rot_dim // 2]
const int rot_dim,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride,
const int num_heads,
const int num_kv_heads,
const int head_size) {
// Each thread block is responsible for one token.
const int token_idx = blockIdx.x;
int64_t pos = positions[token_idx];
const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
apply_rotary_embedding<scalar_t, IS_NEOX>(query,
key,
cache_ptr,
head_size,
num_heads,
num_kv_heads,
rot_dim,
token_idx,
query_stride,
key_stride,
head_stride);
}
} // namespace
namespace xllm::kernel::cuda {
// flashinfer rope ops
// void apply_rope_pos_ids_cos_sin_cache(torch::Tensor q,
// torch::Tensor k,
// torch::Tensor cos_sin_cache,
// torch::Tensor pos_ids,
// bool interleave) {
// const int64_t head_dim = cos_sin_cache.size(-1) / 2;
// q = q.view({q.size(0), -1, head_dim});
// k = k.view({k.size(0), -1, head_dim});
// FunctionFactory::get_instance().rope_func("rope").call(
// q, k, q, k, cos_sin_cache, pos_ids, interleave);
// }
void rotary_embedding(
torch::Tensor& positions, // [batch_size, seq_len] or [num_tokens]
torch::Tensor& query, // [batch_size, seq_len, num_heads * head_size] or
// [num_tokens, num_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
std::optional<torch::Tensor> key,
// null or
// [batch_size, seq_len, num_kv_heads * head_size] or
// [num_tokens, num_kv_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
// int64_t head_size,
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
bool is_neox) {
// num_tokens = batch_size * seq_len
const int positions_ndim = positions.dim();
const int query_ndim = query.dim();
// For partial rotary models, e.g. MiniMax-M2 with head_dim=128 and
// rotary_dim=64, the cache width is the rotary dimension rather than the
// physical per-head stride. When query is already shaped as
// [*, num_heads, head_size], infer the real head_size from query itself.
int64_t head_size = (query_ndim == positions_ndim + 2)
? query.size(-1)
: cos_sin_cache.size(-1);
int64_t num_tokens = positions.numel();
// 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
int64_t head_stride =
(query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
dim3 grid(num_tokens);
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_FLOATING_TYPES(
query.scalar_type(), "apply_rope_pos_ids_cos_sin_cache", [&] {
if (is_neox) {
rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cos_sin_cache.data_ptr<scalar_t>(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size);
} else {
rotary_embedding_kernel<scalar_t, false><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cos_sin_cache.data_ptr<scalar_t>(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size);
}
});
}
} // namespace xllm::kernel::cuda

129
ex_engine/xllm_kernels/cuda/xattention/beam_search.cpp Normal file
View File

@@ -0,0 +1,129 @@
/* 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 <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <torch/script.h>
#include <torch/torch.h>
#include "cuda.h"
namespace xllm::kernel::cuda {
void beam_search(torch::Tensor acc_logprob,
torch::Tensor in_sequence_group,
torch::Tensor top_tokens,
torch::Tensor top_logprobs,
torch::Tensor out_acc_logprob,
torch::Tensor out_token_ids,
torch::Tensor out_token_index,
torch::Tensor out_beam_count_prefix_sums,
torch::Tensor out_sequence_group,
uint32_t batch_size,
uint32_t current_step) {
torch::Device device = acc_logprob.device();
uint32_t beam_size = in_sequence_group.size(1);
uint32_t top_k = top_tokens.size(1);
uint32_t total_rounds = in_sequence_group.size(2);
CHECK_EQ(beam_size, top_k) << "beam_size must be equal with top_k.";
if (current_step == 0) {
auto tokens_view =
top_tokens.view({batch_size, top_k}).slice(1, 0, beam_size);
auto init_probs_view =
top_logprobs.view({batch_size, top_k}).slice(1, 0, beam_size);
out_token_ids.view({batch_size, beam_size}).copy_(tokens_view);
out_acc_logprob.view({batch_size, beam_size}).copy_(init_probs_view);
auto indices =
torch::arange(
beam_size,
torch::TensorOptions().dtype(torch::kInt32).device(device))
.unsqueeze(0)
.expand({batch_size, -1})
.reshape({-1, 1});
out_token_index.copy_(indices);
auto sequence_view =
out_sequence_group.view({batch_size, beam_size, total_rounds});
sequence_view.slice(2, 0, 1).squeeze(2).copy_(tokens_view);
} else {
auto combined_probs =
(acc_logprob + top_logprobs).view({batch_size, beam_size * top_k});
auto topk_result = torch::topk(combined_probs, beam_size, -1);
auto new_probs = std::get<0>(topk_result); // [batch_size, beam_size]
auto new_indices = std::get<1>(topk_result); // [batch_size, beam_size]
auto ordered_indices = new_indices.argsort(static_cast<int64_t>(1), false);
// Reorder new_probs (and corresponding new_indices) by ordered_indices to
// keep alignment.
if (current_step < total_rounds - 1) {
new_probs = new_probs.gather(1, ordered_indices);
new_indices = new_indices.gather(1, ordered_indices);
}
auto parent_beam = (new_indices / top_k).to(torch::kLong);
auto token_in_beam = (new_indices % top_k).to(torch::kLong);
auto top_tokens_reshaped = top_tokens.view({batch_size, beam_size, top_k});
auto batch_idx =
torch::arange(batch_size,
torch::TensorOptions().dtype(torch::kLong).device(device))
.unsqueeze(1)
.expand_as(parent_beam);
using torch::indexing::TensorIndex;
auto new_tokens = top_tokens_reshaped.index({TensorIndex(batch_idx),
TensorIndex(parent_beam),
TensorIndex(token_in_beam)});
out_acc_logprob.view({batch_size, beam_size}).copy_(new_probs);
out_token_index.view({batch_size, beam_size})
.copy_(new_indices.to(torch::kInt32));
out_token_ids.view({batch_size, beam_size}).copy_(new_tokens);
auto batch_range =
torch::arange(
batch_size,
torch::TensorOptions().dtype(torch::kInt32).device(device))
.unsqueeze(1)
.expand({-1, beam_size});
auto beam_range =
torch::arange(
beam_size,
torch::TensorOptions().dtype(torch::kInt32).device(device))
.unsqueeze(0)
.expand({batch_size, -1});
using torch::indexing::Slice;
using torch::indexing::TensorIndex;
out_sequence_group.slice(2, 0, current_step) =
in_sequence_group.index({TensorIndex(batch_range),
TensorIndex(parent_beam.to(torch::kInt32)),
Slice(0, current_step)});
out_sequence_group.slice(2, current_step, current_step + 1) =
new_tokens.unsqueeze(2);
}
}
} // namespace xllm::kernel::cuda

312
ex_engine/xllm_kernels/cuda/xattention/cache_select.cu Normal file
View File

@@ -0,0 +1,312 @@
/* 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 <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <glog/logging.h>
#include <torch/extension.h>
#include <cstdint>
#include <vector>
#include "xattention_ops_api.h"
namespace {
// In-place cache selection kernel for Xattention.
// Reorders KV cache entries based on beam search results. After beam search,
// the beam indices may have changed, and this kernel copies KV cache data from
// old beam positions to new beam positions to maintain consistency.
// Inputs:
// k_ptrs_i64 : [Layer] - pointers to K cache tensors for each layer
// v_ptrs_i64 : [Layer] - pointers to V cache tensors for each layer
// beam_index : [B*Beam] - mapping from new beam index to old beam index
// block_table : [B] - request ID per batch item (extracted from [B*Beam,
// 1]) B : batch size (actual batch size, not batch_size *
// beam_size) Beam : beam width Kv : number of KV
// heads MaxStep : maximum decode steps D : head
// dimension MaxReq : maximum number of requests Layer :
// number of transformer layers decode_step : current decode step
// (0-indexed)
// Cache layout: [MaxReq, Beam, MaxStep, Kv, D]
// The kernel performs two passes to avoid overwriting data:
// pass-1: copy from old_beam > new_beam (increasing new_beam)
// pass-2: copy from old_beam < new_beam (decreasing new_beam)
template <typename scalar_t>
__global__ void cache_select_inplace_ptrs_kernel(
const int64_t* __restrict__ k_ptrs_i64, // [Layer]
const int64_t* __restrict__ v_ptrs_i64, // [Layer]
const int32_t* __restrict__ beam_index, // [B*Beam]
const int32_t* __restrict__ block_table, // [B]
int32_t B,
int32_t Beam,
int32_t Kv,
int32_t MaxStep,
int32_t D,
int32_t MaxReq,
int32_t Layer,
int32_t decode_step) {
const int32_t b = static_cast<int32_t>(blockIdx.x);
const int32_t kv = static_cast<int32_t>(blockIdx.y);
const int32_t layer = static_cast<int32_t>(blockIdx.z);
if (b >= B || kv >= Kv || layer >= Layer) {
return;
}
const int32_t step_end =
decode_step < (MaxStep - 1) ? decode_step : (MaxStep - 1);
const int32_t req = block_table[b];
if (req < 0 || req >= MaxReq) {
return;
}
scalar_t* __restrict__ k_cache =
reinterpret_cast<scalar_t*>(static_cast<uintptr_t>(k_ptrs_i64[layer]));
scalar_t* __restrict__ v_cache =
reinterpret_cast<scalar_t*>(static_cast<uintptr_t>(v_ptrs_i64[layer]));
// base(req, beam, s, kv, d) = ((((req*Beam + beam)*MaxStep + s)*Kv + kv) * D
// + d)
const int64_t req_base = static_cast<int64_t>(req) * Beam;
const int64_t step_kv_stride = static_cast<int64_t>(Kv) * D;
const int64_t kv_d_base = static_cast<int64_t>(kv) * D;
// grid_step is typically small; loop over s in-kernel to reduce launch
// blocks.
for (int32_t s = 0; s <= step_end; ++s) {
// pass-1: new_beam increasing, copy if old_beam > new_beam
for (int32_t new_beam = 0; new_beam < Beam; ++new_beam) {
const int32_t old_beam = beam_index[b * Beam + new_beam] / Beam;
if (old_beam >= 0 && old_beam < Beam && old_beam > new_beam) {
const int64_t dst_base =
((req_base + new_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
const int64_t src_base =
((req_base + old_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
for (int32_t d = static_cast<int32_t>(threadIdx.x); d < D;
d += static_cast<int32_t>(blockDim.x)) {
k_cache[dst_base + d] = k_cache[src_base + d];
v_cache[dst_base + d] = v_cache[src_base + d];
}
}
}
// pass-2: new_beam decreasing, copy if old_beam < new_beam
for (int32_t new_beam = Beam - 1; new_beam >= 0; --new_beam) {
const int32_t old_beam = beam_index[b * Beam + new_beam] / Beam;
if (old_beam >= 0 && old_beam < Beam && old_beam < new_beam) {
const int64_t dst_base =
((req_base + new_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
const int64_t src_base =
((req_base + old_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
for (int32_t d = static_cast<int32_t>(threadIdx.x); d < D;
d += static_cast<int32_t>(blockDim.x)) {
k_cache[dst_base + d] = k_cache[src_base + d];
v_cache[dst_base + d] = v_cache[src_base + d];
}
}
}
}
}
void cache_select_cuda_launch_ptrs(
torch::Tensor k0,
torch::Tensor v0,
torch::Tensor k_ptrs_i64, // [Layer] int64 (CUDA)
torch::Tensor v_ptrs_i64, // [Layer] int64 (CUDA)
torch::Tensor beam_index_i32, // [B*Beam, 1] int32
torch::Tensor block_table_i32, // [B] int32
int64_t decode_step,
int64_t layer_num) {
CHECK(k_ptrs_i64.is_cuda() && v_ptrs_i64.is_cuda())
<< "k_ptrs_i64/v_ptrs_i64 must be CUDA";
CHECK_EQ(k_ptrs_i64.scalar_type(), torch::kInt64)
<< "k_ptrs_i64/v_ptrs_i64 must be int64";
CHECK_EQ(v_ptrs_i64.scalar_type(), torch::kInt64)
<< "k_ptrs_i64/v_ptrs_i64 must be int64";
CHECK(k_ptrs_i64.is_contiguous() && v_ptrs_i64.is_contiguous())
<< "k_ptrs_i64/v_ptrs_i64 must be contiguous";
const int64_t B64 = block_table_i32.size(0);
const int64_t Beam64 = k0.size(1);
const int64_t MaxStep64 = k0.size(2);
const int64_t Kv64 = k0.size(3);
const int64_t D64 = k0.size(4);
const int64_t MaxReq64 = k0.size(0);
const int64_t Layer64 = layer_num;
const int32_t B = static_cast<int32_t>(B64);
const int32_t Beam = static_cast<int32_t>(Beam64);
const int32_t Kv = static_cast<int32_t>(Kv64);
const int32_t MaxStep = static_cast<int32_t>(MaxStep64);
const int32_t D = static_cast<int32_t>(D64);
const int32_t MaxReq = static_cast<int32_t>(MaxReq64);
const int32_t Layer = static_cast<int32_t>(Layer64);
const int32_t decode_step_i32 = static_cast<int32_t>(decode_step);
// Warp-aligned threads, capped to keep occupancy reasonable.
int threads_per_block = ((D + 31) / 32) * 32;
if (threads_per_block < 32) {
threads_per_block = 32;
}
if (threads_per_block > 256) {
threads_per_block = 256;
}
dim3 block_dim(static_cast<unsigned int>(threads_per_block), 1, 1);
CHECK_LE(Kv64, static_cast<int64_t>(UINT32_MAX)) << "Kv too large for grid.y";
CHECK_LE(Layer64, 65535) << "layer_num too large for grid.z";
dim3 grid_dim(static_cast<unsigned int>(B),
static_cast<unsigned int>(Kv),
static_cast<unsigned int>(Layer));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_FLOATING_TYPES_AND2(torch::ScalarType::Half,
torch::ScalarType::BFloat16,
k0.scalar_type(),
"cache_select_inplace_ptrs_kernel",
[&] {
cache_select_inplace_ptrs_kernel<scalar_t>
<<<grid_dim, block_dim, 0, stream>>>(
k_ptrs_i64.data_ptr<int64_t>(),
v_ptrs_i64.data_ptr<int64_t>(),
beam_index_i32.data_ptr<int32_t>(),
block_table_i32.data_ptr<int32_t>(),
B,
Beam,
Kv,
MaxStep,
D,
MaxReq,
Layer,
decode_step_i32);
});
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
} // namespace
namespace xllm::kernel::cuda {
void cache_select(const torch::Tensor& beam_index, // [B*Beam, 1]
std::vector<torch::Tensor>& unshared_k_cache,
std::vector<torch::Tensor>& unshared_v_cache,
const torch::Tensor& block_table, // [B*Beam, 1]
int64_t decode_step,
int64_t beam_size,
int64_t layer_num) {
CHECK_GE(layer_num, 0) << "layer_num must be >= 0";
if (layer_num == 0) {
return;
}
CHECK_EQ(static_cast<int64_t>(unshared_k_cache.size()), layer_num)
<< "unshared_k_cache length mismatch";
CHECK_EQ(static_cast<int64_t>(unshared_v_cache.size()), layer_num)
<< "unshared_v_cache length mismatch";
CHECK(beam_index.is_cuda()) << "beam_index must be CUDA";
CHECK(block_table.is_cuda()) << "block_table must be CUDA";
CHECK_EQ(block_table.dim(), 2) << "block_table must be [B*Beam, 1]";
CHECK_EQ(block_table.size(1), 1) << "block_table must be [B*Beam, 1]";
CHECK_EQ(beam_index.dim(), 2) << "beam_index must be [B*Beam, 1]";
CHECK_EQ(beam_index.size(1), 1) << "beam_index must be [B*Beam, 1]";
CHECK_GE(decode_step, 0) << "decode_step must be >= 0";
CHECK_GT(beam_size, 0) << "beam_size must be > 0";
// block_table is [B*Beam, 1] with sequential values [0,1,2,3,...]
// Infer actual batch_size
CHECK_EQ(block_table.size(0) % beam_size, 0)
<< "block_table.size(0) must be divisible by beam_size";
const int64_t B = block_table.size(0) / beam_size;
CHECK_EQ(beam_index.size(0), B * beam_size)
<< "beam_index size mismatch with B*beam_size";
// Prepare indices (int32, contiguous).
auto beam_index_i32 = beam_index.to(torch::kInt32).contiguous();
auto block_table_i32 = torch::arange(
0,
B,
torch::TensorOptions().dtype(torch::kInt32).device(block_table.device()));
// Validate shapes/dtypes against layer 0.
const auto& k0 = unshared_k_cache[0];
const auto& v0 = unshared_v_cache[0];
CHECK(k0.is_cuda() && v0.is_cuda()) << "cache must be CUDA";
CHECK(k0.is_contiguous() && v0.is_contiguous()) << "cache must be contiguous";
CHECK_EQ(k0.dim(), 5) << "cache must be 5D [MaxReq, Beam, MaxStep, Kv, D]";
CHECK_EQ(v0.sizes(), k0.sizes()) << "k/v cache shapes must match";
CHECK_EQ(k0.size(1), beam_size) << "beam_size mismatch with cache";
CHECK_LT(decode_step, k0.size(2)) << "decode_step must be < max_decode_step";
// Pack layer pointers into CUDA int64 tensors so we can launch once.
// Note: pointer values are produced on host (data_ptr()), then copied to GPU.
c10::cuda::CUDAGuard device_guard(k0.device());
auto ptr_cuda_opts =
torch::TensorOptions().dtype(torch::kInt64).device(k0.device());
auto k_ptrs_i64 = torch::empty({layer_num}, ptr_cuda_opts);
auto v_ptrs_i64 = torch::empty({layer_num}, ptr_cuda_opts);
std::vector<int64_t> k_ptrs_host(static_cast<size_t>(layer_num));
std::vector<int64_t> v_ptrs_host(static_cast<size_t>(layer_num));
for (int64_t layer = 0; layer < layer_num; ++layer) {
auto k = unshared_k_cache[static_cast<size_t>(layer)];
auto v = unshared_v_cache[static_cast<size_t>(layer)];
CHECK(k.is_cuda() && v.is_cuda()) << "cache must be CUDA";
CHECK(k.is_contiguous() && v.is_contiguous()) << "cache must be contiguous";
CHECK_EQ(k.sizes(), k0.sizes()) << "all layers must have same cache shape";
CHECK_EQ(v.sizes(), k0.sizes()) << "all layers must have same cache shape";
CHECK_EQ(k.scalar_type(), k0.scalar_type())
<< "all layers must have same dtype";
CHECK_EQ(v.scalar_type(), k0.scalar_type())
<< "all layers must have same dtype";
CHECK_EQ(k.get_device(), k0.get_device())
<< "all layers must be on the same CUDA device";
CHECK_EQ(v.get_device(), k0.get_device())
<< "all layers must be on the same CUDA device";
k_ptrs_host[static_cast<size_t>(layer)] =
static_cast<int64_t>(reinterpret_cast<uintptr_t>(k.data_ptr()));
v_ptrs_host[static_cast<size_t>(layer)] =
static_cast<int64_t>(reinterpret_cast<uintptr_t>(v.data_ptr()));
}
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
C10_CUDA_CHECK(
cudaMemcpyAsync(k_ptrs_i64.data_ptr<int64_t>(),
k_ptrs_host.data(),
static_cast<size_t>(layer_num) * sizeof(int64_t),
cudaMemcpyHostToDevice,
stream));
C10_CUDA_CHECK(
cudaMemcpyAsync(v_ptrs_i64.data_ptr<int64_t>(),
v_ptrs_host.data(),
static_cast<size_t>(layer_num) * sizeof(int64_t),
cudaMemcpyHostToDevice,
stream));
cache_select_cuda_launch_ptrs(k0,
v0,
k_ptrs_i64,
v_ptrs_i64,
beam_index_i32,
block_table_i32,
decode_step,
layer_num);
}
} // namespace xllm::kernel::cuda

298
ex_engine/xllm_kernels/cuda/xattention/decoder_reshape_and_cache.cu Normal file
View File

@@ -0,0 +1,298 @@
/* 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 <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <torch/cuda.h>
#include <cstdint>
#include <type_traits>
#include "kernels/cuda/utils.h"
#include "xattention_ops_api.h"
namespace {
template <typename scalar_t>
struct VecType;
template <>
struct VecType<c10::Half> {
using type = uint4; // 8 elements * 2 bytes = 16 bytes
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<c10::BFloat16> {
using type = uint4; // 8 elements * 2 bytes = 16 bytes
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<float> {
using type = float4; // 4 elements * 4 bytes = 16 bytes
static constexpr int32_t vec_width = 4;
};
// decoder reshape and cache kernel.
// Copies proj_k and proj_v into unshared_k_cache / unshared_v_cache.
// Inputs:
// proj_k : [batch_size, beam_size, kv_heads, head_dim]
// proj_v : [batch_size, beam_size, kv_heads, head_dim]
// step : [1] - current decode step
// batch_size : batch size
// beam_size : beam size
// kv_heads : number of kv heads
// head_dim : head dimension
// k_stride0 : proj_k.stride(0)
// k_stride1 : proj_k.stride(1)
// v_stride0 : proj_v.stride(0)
// v_stride1 : proj_v.stride(1)
// cache_stride0 : unshared_k_cache.stride(0)
// cache_stride1 : unshared_k_cache.stride(1)
// cache_stride2 : unshared_k_cache.stride(2)
// cache_stride3 : unshared_k_cache.stride(3)
// Outputs:
// unshared_k_cache : [max_batch_size, beam_size, max_step, kv_heads,
// head_dim]
// unshared_v_cache : [max_batch_size, beam_size, max_step, kv_heads,
// head_dim]
template <typename scalar_t>
__global__ void decoder_reshape_and_cache_kernel(
const scalar_t* __restrict__ proj_k,
const scalar_t* __restrict__ proj_v,
scalar_t* __restrict__ unshared_k_cache,
scalar_t* __restrict__ unshared_v_cache,
const int32_t* __restrict__ step,
const int64_t batch_size,
const int64_t beam_size,
const int64_t kv_heads,
const int64_t head_dim,
const int64_t k_stride0,
const int64_t k_stride1,
const int64_t v_stride0,
const int64_t v_stride1,
const int64_t cache_stride0,
const int64_t cache_stride1,
const int64_t cache_stride2,
const int64_t cache_stride3) {
using VecTypeT = typename VecType<scalar_t>::type;
constexpr int32_t VEC_WIDTH = VecType<scalar_t>::vec_width;
const int64_t token_idx = static_cast<int64_t>(blockIdx.y);
const int64_t total_tokens = batch_size * beam_size;
if (token_idx >= total_tokens) {
return;
}
const int64_t batch_idx = token_idx / beam_size;
const int64_t beam_idx = token_idx - batch_idx * beam_size;
__shared__ int32_t current_step_s;
if (threadIdx.x == 0) {
current_step_s = __ldg(step);
}
__syncthreads();
const int64_t current_step = static_cast<int64_t>(current_step_s);
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
const int64_t total_vecs = kv_heads * vecs_per_head;
const int64_t k_token_base = batch_idx * k_stride0 + beam_idx * k_stride1;
const int64_t v_token_base = batch_idx * v_stride0 + beam_idx * v_stride1;
const int64_t dst_token_base = batch_idx * cache_stride0 +
beam_idx * cache_stride1 +
current_step * cache_stride2;
for (int64_t linear_idx = static_cast<int64_t>(threadIdx.x);
linear_idx < total_vecs;
linear_idx += static_cast<int64_t>(blockDim.x)) {
const int64_t head_idx = linear_idx / vecs_per_head;
const int64_t vec_idx = linear_idx - head_idx * vecs_per_head;
const int64_t vec_offset = vec_idx * VEC_WIDTH;
const auto* k_src_vec = reinterpret_cast<const VecTypeT*>(
proj_k + k_token_base + head_idx * head_dim + vec_offset);
const auto* v_src_vec = reinterpret_cast<const VecTypeT*>(
proj_v + v_token_base + head_idx * head_dim + vec_offset);
auto* k_dst_vec =
reinterpret_cast<VecTypeT*>(unshared_k_cache + dst_token_base +
head_idx * cache_stride3 + vec_offset);
auto* v_dst_vec =
reinterpret_cast<VecTypeT*>(unshared_v_cache + dst_token_base +
head_idx * cache_stride3 + vec_offset);
*k_dst_vec = *k_src_vec;
*v_dst_vec = *v_src_vec;
}
}
} // namespace
namespace xllm::kernel::cuda {
void decoder_reshape_and_cache(torch::Tensor proj_k,
torch::Tensor proj_v,
torch::Tensor unshared_k_cache,
torch::Tensor unshared_v_cache,
torch::Tensor step) {
CHECK_EQ(proj_k.dim(), 4) << "proj_k must be 4-dimensional";
CHECK_EQ(proj_v.dim(), 4) << "proj_v must be 4-dimensional";
CHECK_EQ(unshared_k_cache.dim(), 5)
<< "unshared_k_cache must be 5-dimensional";
CHECK_EQ(unshared_v_cache.dim(), 5)
<< "unshared_v_cache must be 5-dimensional";
CHECK(proj_k.is_cuda() && proj_v.is_cuda() && unshared_k_cache.is_cuda() &&
unshared_v_cache.is_cuda() && step.is_cuda())
<< "all tensors must be CUDA tensors";
CHECK_EQ(step.dim(), 1) << "step must be 1-dimensional";
CHECK_EQ(step.size(0), 1) << "step must have shape [1]";
CHECK_EQ(step.scalar_type(), at::ScalarType::Int)
<< "step must be int32 (torch::kInt32)";
const int64_t batch_size = proj_k.size(0);
const int64_t beam_size = proj_k.size(1);
const int64_t kv_heads = proj_k.size(2);
const int64_t head_dim = proj_k.size(3);
CHECK_EQ(proj_v.sizes(), proj_k.sizes())
<< "proj_v and proj_k must have same shape";
CHECK_EQ(unshared_k_cache.size(3), kv_heads)
<< "unshared_k_cache kv_heads mismatch";
CHECK_EQ(unshared_k_cache.size(4), head_dim)
<< "unshared_k_cache head_dim mismatch";
CHECK(unshared_v_cache.sizes() == unshared_k_cache.sizes())
<< "unshared_v_cache and unshared_k_cache must have same shape";
// This kernel is specialized for qkv-slice layouts:
// last dim contiguous and kv head stride tightly packed by head_dim.
CHECK_EQ(proj_k.stride(3), 1) << "proj_k must satisfy stride(3)=1";
CHECK_EQ(proj_v.stride(3), 1) << "proj_v must satisfy stride(3)=1";
CHECK_EQ(proj_k.stride(2), head_dim)
<< "proj_k must satisfy stride(2)=head_dim";
CHECK_EQ(proj_v.stride(2), head_dim)
<< "proj_v must satisfy stride(2)=head_dim";
CHECK_EQ(unshared_k_cache.stride(4), 1)
<< "unshared_k_cache must satisfy stride(4)=1";
CHECK_EQ(unshared_v_cache.stride(4), 1)
<< "unshared_v_cache must satisfy stride(4)=1";
CHECK_EQ(unshared_k_cache.stride(3), head_dim)
<< "unshared_k_cache must satisfy stride(3)=head_dim";
CHECK_EQ(unshared_v_cache.stride(3), head_dim)
<< "unshared_v_cache must satisfy stride(3)=head_dim";
const at::cuda::OptionalCUDAGuard device_guard(device_of(proj_k));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const int64_t k_stride0 = proj_k.stride(0);
const int64_t k_stride1 = proj_k.stride(1);
const int64_t v_stride0 = proj_v.stride(0);
const int64_t v_stride1 = proj_v.stride(1);
const int64_t cache_stride0 = unshared_k_cache.stride(0);
const int64_t cache_stride1 = unshared_k_cache.stride(1);
const int64_t cache_stride2 = unshared_k_cache.stride(2);
const int64_t cache_stride3 = unshared_k_cache.stride(3);
// Launch kernel: one block per (batch, beam), threads cover
// kv_heads*head_dim.
const int64_t total_tokens = batch_size * beam_size;
dim3 grid_dim(1, static_cast<unsigned int>(total_tokens), 1);
DISPATCH_FLOATING_TYPES(
proj_k.scalar_type(), "decoder_reshape_and_cache_kernel", [&] {
constexpr int32_t VEC_WIDTH = (std::is_same_v<scalar_t, c10::Half> ||
std::is_same_v<scalar_t, c10::BFloat16>)
? 8
: 4; // FP16/BF16: 8, Float: 4
constexpr int32_t kWarpSize = 32;
constexpr int32_t kMaxThreadsPerBlock = 256;
constexpr int32_t kAlignmentBytes = 16; // 128-bit alignment
CHECK(head_dim % VEC_WIDTH == 0)
<< "head_dim must be divisible by vector width: " << VEC_WIDTH;
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
const int64_t total_vecs = kv_heads * vecs_per_head;
CHECK(total_vecs > 0) << "total_vecs must be > 0";
int32_t threads_per_block = static_cast<int32_t>(
total_vecs > kMaxThreadsPerBlock ? kMaxThreadsPerBlock
: total_vecs);
threads_per_block =
((threads_per_block + kWarpSize - 1) / kWarpSize) * kWarpSize;
if (threads_per_block < kWarpSize) {
threads_per_block = kWarpSize;
}
dim3 block_dim(threads_per_block, 1, 1);
const auto proj_k_ptr =
reinterpret_cast<std::uintptr_t>(proj_k.data_ptr<scalar_t>());
const auto proj_v_ptr =
reinterpret_cast<std::uintptr_t>(proj_v.data_ptr<scalar_t>());
const auto k_cache_ptr = reinterpret_cast<std::uintptr_t>(
unshared_k_cache.data_ptr<scalar_t>());
const auto v_cache_ptr = reinterpret_cast<std::uintptr_t>(
unshared_v_cache.data_ptr<scalar_t>());
CHECK(proj_k_ptr % kAlignmentBytes == 0)
<< "proj_k data_ptr must be 16-byte aligned";
CHECK(proj_v_ptr % kAlignmentBytes == 0)
<< "proj_v data_ptr must be 16-byte aligned";
CHECK(k_cache_ptr % kAlignmentBytes == 0)
<< "unshared_k_cache data_ptr must be 16-byte aligned";
CHECK(v_cache_ptr % kAlignmentBytes == 0)
<< "unshared_v_cache data_ptr must be 16-byte aligned";
const int64_t scalar_bytes = static_cast<int64_t>(sizeof(scalar_t));
CHECK((k_stride0 * scalar_bytes) % kAlignmentBytes == 0)
<< "proj_k stride(0) bytes must be 16-byte aligned";
CHECK((k_stride1 * scalar_bytes) % kAlignmentBytes == 0)
<< "proj_k stride(1) bytes must be 16-byte aligned";
CHECK((v_stride0 * scalar_bytes) % kAlignmentBytes == 0)
<< "proj_v stride(0) bytes must be 16-byte aligned";
CHECK((v_stride1 * scalar_bytes) % kAlignmentBytes == 0)
<< "proj_v stride(1) bytes must be 16-byte aligned";
CHECK((cache_stride0 * scalar_bytes) % kAlignmentBytes == 0)
<< "cache stride(0) bytes must be 16-byte aligned";
CHECK((cache_stride1 * scalar_bytes) % kAlignmentBytes == 0)
<< "cache stride(1) bytes must be 16-byte aligned";
CHECK((cache_stride2 * scalar_bytes) % kAlignmentBytes == 0)
<< "cache stride(2) bytes must be 16-byte aligned";
CHECK((cache_stride3 * scalar_bytes) % kAlignmentBytes == 0)
<< "cache stride(3) bytes must be 16-byte aligned";
decoder_reshape_and_cache_kernel<scalar_t>
<<<grid_dim, block_dim, 0, stream>>>(
proj_k.data_ptr<scalar_t>(),
proj_v.data_ptr<scalar_t>(),
unshared_k_cache.data_ptr<scalar_t>(),
unshared_v_cache.data_ptr<scalar_t>(),
step.data_ptr<int32_t>(),
batch_size,
beam_size,
kv_heads,
head_dim,
k_stride0,
k_stride1,
v_stride0,
v_stride1,
cache_stride0,
cache_stride1,
cache_stride2,
cache_stride3);
});
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
} // namespace xllm::kernel::cuda

168
ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu Normal file
View File

@@ -0,0 +1,168 @@
/* 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 <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <torch/cuda.h>
#include <cmath>
#include "kernels/cuda/utils.h"
#include "xattention_ops_api.h"
namespace {
// Fused log-sum-exp combine kernel.
//
// Layout and strategy (aligned with the TileLang version):
// - Each block is responsible for one (batch_idx, head_idx) pair, i.e. one
// row in the flattened [B * H, D] layout.
// - Threads within a block parallelize along the head_dim (D) dimension to
// ensure coalesced global memory access.
//
// Tensors:
// shared_o : [B, H, D] - shared attention output
// shared_lse : [B, H, 1] - shared log-sum-exp (FP32)
// unshared_o : [B, H, D] - unshared attention output
// unshared_lse: [B, H, 1] - unshared log-sum-exp (FP32)
// output : [B, H, D] - combined output
template <typename scalar_t, typename out_scalar_t>
__global__ void lse_combine_kernel(
out_scalar_t* __restrict__ output, // [B, H, D]
const scalar_t* __restrict__ shared_o, // [B, H, D]
const float* __restrict__ shared_lse, // [B, H, 1], always FP32
const scalar_t* __restrict__ unshared_o, // [B, H, D]
const float* __restrict__ unshared_lse, // [B, H, 1], always FP32
const int64_t B, // batch_size * beam_size
const int64_t H, // num_heads
const int64_t D) { // head_dim
const int64_t total_elements = B * H;
const int64_t idx = static_cast<int64_t>(blockIdx.y);
if (idx >= total_elements) {
return;
}
// Load LSE scalars for this (batch, head) pair.
const float shared_lse_val = shared_lse[idx];
const float unshared_lse_val = unshared_lse[idx];
// 1. Compute element-wise max LSE.
const float lse_max = fmaxf(shared_lse_val, unshared_lse_val);
// 2. Compute base-2 exponentials relative to max.
const float exp_shared = exp2f(shared_lse_val - lse_max);
const float exp_unshared = exp2f(unshared_lse_val - lse_max);
// 3. Compute merged LSE.
const float lse_new = lse_max + log2f(exp_shared + exp_unshared);
// 4. Compute normalized weights.
const float w_shared = exp2f(shared_lse_val - lse_new);
const float w_unshared = exp2f(unshared_lse_val - lse_new);
// 5. Weighted combine along the head_dim.
const int64_t base_idx = idx * D;
// Threads in the block parallelize along D with stride blockDim.x for
// coalesced global memory access.
for (int64_t d = threadIdx.x; d < D; d += blockDim.x) {
const float shared_val = static_cast<float>(shared_o[base_idx + d]);
const float unshared_val = static_cast<float>(unshared_o[base_idx + d]);
const float combined = w_shared * shared_val + w_unshared * unshared_val;
output[base_idx + d] = static_cast<out_scalar_t>(combined);
}
}
} // namespace
namespace xllm::kernel::cuda {
// Host wrapper for the fused LSE combine kernel.
//
// All inputs are expected to be on the same CUDA device:
// shared_o : [B, H, D], floating type (including Half/BFloat16)
// shared_lse : [B, H, 1], float32
// unshared_o : [B, H, D], same type/shape as shared_o
// unshared_lse: [B, H, 1], float32
// output : [B, H, D], will be resized/allocated as needed.
void lse_combine(torch::Tensor output,
torch::Tensor shared_o,
torch::Tensor shared_lse,
torch::Tensor unshared_o,
torch::Tensor unshared_lse) {
CHECK_EQ(shared_o.dim(), 3) << "shared_o must be 3D [B, H, D]";
CHECK_EQ(unshared_o.dim(), 3) << "unshared_o must be 3D [B, H, D]";
CHECK_EQ(shared_lse.dim(), 3) << "shared_lse must be 3D [B, H, 1]";
CHECK_EQ(unshared_lse.dim(), 3) << "unshared_lse must be 3D [B, H, 1]";
const int64_t B = shared_o.size(0);
const int64_t H = shared_o.size(1);
const int64_t D = shared_o.size(2);
CHECK_EQ(shared_o.sizes(), unshared_o.sizes())
<< "shared_o and unshared_o must have same shape";
CHECK_EQ(shared_lse.scalar_type(), torch::kFloat32)
<< "shared_lse must be float32";
CHECK_EQ(unshared_lse.scalar_type(), torch::kFloat32)
<< "unshared_lse must be float32";
CHECK_EQ(shared_lse.size(0), B)
<< "shared_lse shape mismatch, expected [B, H, 1]";
CHECK_EQ(shared_lse.size(1), H)
<< "shared_lse shape mismatch, expected [B, H, 1]";
CHECK_EQ(shared_lse.size(2), 1)
<< "shared_lse shape mismatch, expected [B, H, 1]";
CHECK_EQ(unshared_lse.size(0), B)
<< "unshared_lse shape mismatch, expected [B, H, 1]";
CHECK_EQ(unshared_lse.size(1), H)
<< "unshared_lse shape mismatch, expected [B, H, 1]";
CHECK_EQ(unshared_lse.size(2), 1)
<< "unshared_lse shape mismatch, expected [B, H, 1]";
// Ensure output has the correct shape and dtype.
if (!output.defined() || output.sizes() != shared_o.sizes()) {
output = torch::empty_like(shared_o);
}
const at::cuda::OptionalCUDAGuard device_guard(device_of(shared_o));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
// Launch kernel: one block per (batch, head) pair, threads along D.
const int64_t total_elements = B * H;
const int threads_per_block = 128;
dim3 block_dim(threads_per_block, 1, 1);
dim3 grid_dim(1, static_cast<unsigned int>(total_elements), 1);
DISPATCH_FLOATING_TYPES(
shared_o.scalar_type(), "lse_combine_kernel_input", [&] {
using in_t = scalar_t;
DISPATCH_FLOATING_TYPES(
output.scalar_type(), "lse_combine_kernel_output", [&] {
using out_t = scalar_t;
lse_combine_kernel<in_t, out_t>
<<<grid_dim, block_dim, 0, stream>>>(
output.data_ptr<out_t>(),
shared_o.data_ptr<in_t>(),
shared_lse.data_ptr<float>(),
unshared_o.data_ptr<in_t>(),
unshared_lse.data_ptr<float>(),
B,
H,
D);
});
});
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
} // namespace xllm::kernel::cuda

220
ex_engine/xllm_kernels/cuda/xattention/prefill_reshape_and_cache.cu Normal file
View File

@@ -0,0 +1,220 @@
/* 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 <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <glog/logging.h>
#include <torch/cuda.h>
#include <cstdint>
#include <type_traits>
#include "kernels/cuda/cuda_ops_api.h"
#include "kernels/cuda/utils.h"
using at::device_of;
namespace {
template <typename scalar_t>
struct VecType;
template <>
struct VecType<c10::Half> {
using type = uint4; // 8 elements * 2 bytes = 16 bytes
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<c10::BFloat16> {
using type = uint4; // 8 elements * 2 bytes = 16 bytes
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<float> {
using type = float4; // 4 elements * 4 bytes = 16 bytes
static constexpr int32_t vec_width = 4;
};
template <typename scalar_t>
__global__ void prefill_reshape_and_cache_kernel(
const scalar_t* __restrict__ proj_k, // [shared_len, kv_heads, head_dim]
const scalar_t* __restrict__ proj_v, // [shared_len, kv_heads, head_dim]
scalar_t* __restrict__ shared_k_cache, // [shared_len, kv_heads, head_dim]
scalar_t* __restrict__ shared_v_cache, // [shared_len, kv_heads, head_dim]
const int64_t shared_len,
const int64_t kv_heads,
const int64_t head_dim,
const int64_t k_stride0, // proj_k.stride(0)
const int64_t v_stride0, // proj_v.stride(0)
const int64_t v_stride1) { // proj_v.stride(1), same as head_dim
using VecTypeT = typename VecType<scalar_t>::type;
constexpr int32_t VEC_WIDTH = VecType<scalar_t>::vec_width;
const int64_t token_idx = static_cast<int64_t>(blockIdx.y);
if (token_idx >= shared_len) {
return;
}
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
const int64_t total_vecs = kv_heads * vecs_per_head;
const int64_t k_token_base = token_idx * k_stride0;
const int64_t v_token_base = token_idx * v_stride0;
const int64_t dst_token_base = token_idx * kv_heads * head_dim;
for (int64_t linear_idx = threadIdx.x; linear_idx < total_vecs;
linear_idx += blockDim.x) {
const int64_t head_idx = linear_idx / vecs_per_head;
const int64_t vec_idx = linear_idx - head_idx * vecs_per_head;
const int64_t head_offset = head_idx * head_dim;
const int64_t vec_offset = vec_idx * VEC_WIDTH;
const auto* k_src_vec = reinterpret_cast<const VecTypeT*>(
proj_k + k_token_base + head_offset + vec_offset);
const auto* v_src_vec = reinterpret_cast<const VecTypeT*>(
proj_v + v_token_base + head_idx * v_stride1 + vec_offset);
auto* k_dst_vec = reinterpret_cast<VecTypeT*>(
shared_k_cache + dst_token_base + head_offset + vec_offset);
auto* v_dst_vec = reinterpret_cast<VecTypeT*>(
shared_v_cache + dst_token_base + head_offset + vec_offset);
*k_dst_vec = *k_src_vec;
*v_dst_vec = *v_src_vec;
}
}
} // namespace
namespace xllm::kernel::cuda {
void prefill_reshape_and_cache(
torch::Tensor proj_k, // [shared_len, kv_heads, head_dim]
torch::Tensor proj_v, // [shared_len, kv_heads, head_dim]
torch::Tensor
shared_k_cache, // [num_shared_kv_seq_len, kv_heads, head_dim]
torch::Tensor shared_v_cache) {
CHECK(proj_k.dim() == 3) << "proj_k must be 3-dimensional";
CHECK(proj_v.dim() == 3) << "proj_v must be 3-dimensional";
CHECK(shared_k_cache.dim() == 3) << "shared_k_cache must be 3-dimensional";
CHECK(shared_v_cache.dim() == 3) << "shared_v_cache must be 3-dimensional";
CHECK(proj_k.is_cuda() && proj_v.is_cuda() && shared_k_cache.is_cuda() &&
shared_v_cache.is_cuda())
<< "all tensors must be CUDA tensors";
const int64_t shared_len = proj_k.size(0);
const int64_t kv_heads = proj_k.size(1);
const int64_t head_dim = proj_k.size(2);
CHECK(proj_v.sizes() == proj_k.sizes())
<< "proj_v and proj_k must have same shape";
CHECK(shared_k_cache.size(0) >= shared_len &&
shared_k_cache.size(1) == kv_heads &&
shared_k_cache.size(2) == head_dim)
<< "shared_k_cache shape mismatch";
CHECK(shared_v_cache.size(0) >= shared_len &&
shared_v_cache.size(1) == kv_heads &&
shared_v_cache.size(2) == head_dim)
<< "shared_v_cache shape mismatch";
shared_k_cache = shared_k_cache.slice(0, 0, shared_len);
shared_v_cache = shared_v_cache.slice(0, 0, shared_len);
// This kernel is specialized for qkv-slice layouts:
// last dim contiguous and head stride tightly packed by head_dim.
CHECK(proj_k.stride(2) == 1 && proj_v.stride(2) == 1)
<< "proj_k/proj_v must be contiguous on head_dim (stride(2)=1)";
CHECK(proj_k.stride(1) == head_dim && proj_v.stride(1) == head_dim)
<< "proj_k/proj_v must satisfy stride(1)=head_dim for qkv-slice layout";
CHECK(shared_k_cache.stride(2) == 1 && shared_v_cache.stride(2) == 1)
<< "shared caches must be contiguous on head_dim (stride(2)=1)";
CHECK(shared_k_cache.stride(1) == head_dim &&
shared_v_cache.stride(1) == head_dim)
<< "shared caches must satisfy stride(1)=head_dim";
CHECK(shared_k_cache.stride(0) == kv_heads * head_dim &&
shared_v_cache.stride(0) == kv_heads * head_dim)
<< "shared caches must be contiguous on token stride";
const at::cuda::OptionalCUDAGuard device_guard(device_of(proj_k));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const int64_t k_stride0 = proj_k.stride(0);
const int64_t v_stride0 = proj_v.stride(0);
const int64_t v_stride1 = proj_v.stride(1);
dim3 grid_dim(1, static_cast<unsigned int>(shared_len), 1);
DISPATCH_FLOATING_TYPES(
proj_k.scalar_type(), "prefill_reshape_and_cache_kernel", [&] {
constexpr int32_t VEC_WIDTH = (std::is_same_v<scalar_t, c10::Half> ||
std::is_same_v<scalar_t, c10::BFloat16>)
? 8
: 4; // FP16/BF16: 8, Float: 4
constexpr int32_t kWarpSize = 32;
constexpr int32_t kMaxThreadsPerBlock = 256;
CHECK(head_dim % VEC_WIDTH == 0)
<< "head_dim must be divisible by vector width: " << VEC_WIDTH;
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
const int64_t total_vecs = kv_heads * vecs_per_head;
CHECK(total_vecs > 0) << "total_vecs must be > 0";
int32_t threads_per_block = static_cast<int32_t>(
total_vecs > kMaxThreadsPerBlock ? kMaxThreadsPerBlock
: total_vecs);
threads_per_block =
((threads_per_block + kWarpSize - 1) / kWarpSize) * kWarpSize;
if (threads_per_block < kWarpSize) {
threads_per_block = kWarpSize;
}
dim3 block_dim(threads_per_block, 1, 1);
const auto proj_k_ptr =
reinterpret_cast<std::uintptr_t>(proj_k.data_ptr<scalar_t>());
const auto proj_v_ptr =
reinterpret_cast<std::uintptr_t>(proj_v.data_ptr<scalar_t>());
const auto k_cache_ptr = reinterpret_cast<std::uintptr_t>(
shared_k_cache.data_ptr<scalar_t>());
const auto v_cache_ptr = reinterpret_cast<std::uintptr_t>(
shared_v_cache.data_ptr<scalar_t>());
constexpr int32_t alignment_bytes = 16; // 128-bit alignment
CHECK(proj_k_ptr % alignment_bytes == 0)
<< "proj_k data_ptr must be 16-byte aligned";
CHECK(proj_v_ptr % alignment_bytes == 0)
<< "proj_v data_ptr must be 16-byte aligned";
CHECK(k_cache_ptr % alignment_bytes == 0)
<< "shared_k_cache data_ptr must be 16-byte aligned";
CHECK(v_cache_ptr % alignment_bytes == 0)
<< "shared_v_cache data_ptr must be 16-byte aligned";
const int64_t scalar_bytes = static_cast<int64_t>(sizeof(scalar_t));
CHECK((k_stride0 * scalar_bytes) % alignment_bytes == 0)
<< "proj_k stride(0) bytes must be 16-byte aligned";
CHECK((v_stride0 * scalar_bytes) % alignment_bytes == 0)
<< "proj_v stride(0) bytes must be 16-byte aligned";
CHECK((v_stride1 * scalar_bytes) % alignment_bytes == 0)
<< "proj_v stride(1) bytes must be 16-byte aligned";
prefill_reshape_and_cache_kernel<scalar_t>
<<<grid_dim, block_dim, 0, stream>>>(
proj_k.data_ptr<scalar_t>(),
proj_v.data_ptr<scalar_t>(),
shared_k_cache.data_ptr<scalar_t>(),
shared_v_cache.data_ptr<scalar_t>(),
shared_len,
kv_heads,
head_dim,
k_stride0,
v_stride0,
v_stride1);
});
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
} // namespace xllm::kernel::cuda

63
ex_engine/xllm_kernels/cuda/xattention/xattention_ops_api.h Normal file
View File

@@ -0,0 +1,63 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <torch/torch.h>
#include <vector>
namespace xllm::kernel::cuda {
void decoder_reshape_and_cache(torch::Tensor proj_k,
torch::Tensor proj_v,
torch::Tensor unshared_k_cache,
torch::Tensor unshared_v_cache,
torch::Tensor step);
void cache_select(const torch::Tensor& beam_index,
std::vector<torch::Tensor>& unshared_k_cache,
std::vector<torch::Tensor>& unshared_v_cache,
const torch::Tensor& block_table,
int64_t decode_step,
int64_t beam_size,
int64_t layer_num);
void lse_combine(torch::Tensor output,
torch::Tensor shared_o,
torch::Tensor shared_lse,
torch::Tensor unshared_o,
torch::Tensor unshared_lse);
void prefill_reshape_and_cache(
torch::Tensor proj_k, // [shared_len, kv_heads, head_dim]
torch::Tensor proj_v, // [shared_len, kv_heads, head_dim]
torch::Tensor
shared_k_cache, // [num_shared_kv_seq_len, kv_heads, head_dim]
torch::Tensor shared_v_cache);
void beam_search(torch::Tensor acc_logprob,
torch::Tensor in_sequence_group,
torch::Tensor top_tokens,
torch::Tensor top_logprobs,
torch::Tensor out_acc_logprob,
torch::Tensor out_token_ids,
torch::Tensor out_token_index,
torch::Tensor out_beam_count_prefix_sums,
torch::Tensor out_sequence_group,
uint32_t batch_size,
uint32_t current_step);
} // namespace xllm::kernel::cuda

28
ex_engine/xllm_kernels/ilu/CMakeLists.txt Normal file
View File

@@ -0,0 +1,28 @@
include(cc_library)
set(CMAKE_CUDA_ARCHITECTURES ivcore11)
file(GLOB_RECURSE ILU_HEADER_FILES
"${CMAKE_CURRENT_LIST_DIR}/*.h"
)
file(GLOB_RECURSE ILU_SOURCE_FILES
"${CMAKE_CURRENT_LIST_DIR}/*.cpp"
"${CMAKE_CURRENT_LIST_DIR}/*.cu"
)
find_package(Python3 REQUIRED COMPONENTS Interpreter Development)
cc_library(
NAME
ilu_kernels
HDRS
${ILU_HEADER_FILES}
SRCS
${ILU_SOURCE_FILES}
DEPS
torch
:util
ixformer_kernels
ixformer
${Python3_LIBRARIES}
cuinfer
)

32
ex_engine/xllm_kernels/ilu/activation.cpp Normal file
View File

@@ -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

163
ex_engine/xllm_kernels/ilu/attention.cpp Normal file
View File

@@ -0,0 +1,163 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
#include "ixinfer.h"
#include "utils.h"
using namespace ixformer;
namespace xllm::kernel::ilu {
void reshape_paged_cache(torch::Tensor& key,
std::optional<torch::Tensor>& value,
torch::Tensor& key_cache,
std::optional<torch::Tensor>& value_cache,
torch::Tensor& slot_mapping) {
auto value_ = value.value_or(torch::Tensor());
auto value_cache_ = value_cache.value_or(torch::Tensor());
int64_t key_token_stride = key.stride(0);
int64_t value_token_stride = 0;
if (value_.defined()) {
value_token_stride = value_.stride(0);
}
slot_mapping = slot_mapping.to(at::kLong);
infer::xllm_reshape_and_cache(key,
value_,
key_cache,
value_cache_,
slot_mapping,
key_token_stride,
value_token_stride);
}
void batch_prefill(torch::Tensor& query,
const torch::Tensor& key,
const std::optional<torch::Tensor>& value,
torch::Tensor& output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_cu_seq_lens,
const std::optional<torch::Tensor>& kv_cu_seq_lens,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& attn_bias,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_quant_scale,
const std::optional<torch::Tensor>& v_quant_scale,
const torch::Tensor& block_tables,
int64_t max_query_len,
int64_t max_seq_len,
float scale,
bool is_causal,
int64_t window_size_left,
int64_t window_size_right,
const std::string& compute_dtype,
bool return_lse) {
double softcap = 0.0;
bool sqrt_alibi = false;
auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor());
auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor());
auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor());
auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor());
auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor());
auto block_tables_ = block_tables;
auto key_ = key;
auto value_ = value.value();
infer::ixinfer_flash_attn_unpad_with_block_tables(query,
key_,
value_,
output,
block_tables_,
q_cu_seq_lens_,
kv_cu_seq_lens_,
max_query_len,
max_seq_len,
is_causal,
window_size_left,
window_size_right,
static_cast<double>(scale),
softcap,
sqrt_alibi,
alibi_slope,
c10::nullopt,
output_lse);
}
void batch_decode(torch::Tensor& query,
const torch::Tensor& k_cache,
torch::Tensor& output,
const torch::Tensor& block_table,
const torch::Tensor& seq_lens,
const std::optional<torch::Tensor>& v_cache,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_cache_quant_scale,
const std::optional<torch::Tensor>& v_cache_quant_scale,
const std::optional<torch::Tensor>& out_quant_scale,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& mask,
const std::string& compute_dtype,
int64_t max_seq_len,
int64_t window_size_left,
int64_t window_size_right,
float scale,
bool return_lse,
bool is_causal,
int64_t kv_cache_quant_bit_size) {
if (query.dim() == 4) {
query =
query
.view({query.size(0) * query.size(1), query.size(2), query.size(3)})
.contiguous();
}
if (output.dim() == 4) {
output = output
.view({output.size(0) * output.size(1),
output.size(2),
output.size(3)})
.contiguous();
;
}
auto v_cache_ = v_cache.value_or(torch::Tensor());
int64_t num_kv_heads = k_cache.size(1);
int64_t page_block_size = k_cache.size(2);
double softcap = 0.0;
bool enable_cuda_graph = false;
bool use_sqrt_alibi = false;
auto block_table_ = block_table;
auto k_cache_ = k_cache;
auto seq_lens_ = seq_lens;
infer::xllm_paged_attention(output,
query,
k_cache_,
v_cache_,
num_kv_heads,
scale,
block_table_,
seq_lens_,
page_block_size,
max_seq_len,
alibi_slope,
is_causal,
(int32_t)window_size_left,
(int32_t)window_size_right,
softcap,
enable_cuda_graph,
use_sqrt_alibi,
c10::nullopt);
}
} // namespace xllm::kernel::ilu

99
ex_engine/xllm_kernels/ilu/fused_moe.cpp Normal file
View File

@@ -0,0 +1,99 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <glog/logging.h>
#include "ilu_ops_api.h"
namespace xllm::kernel::ilu {
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
const torch::Tensor& input,
int64_t topk,
int64_t num_expert_group,
int64_t topk_group,
bool normalize,
const std::optional<torch::Tensor>& mask,
const std::string& normed_by,
const std::string& scoring_func,
double route_scale,
const std::optional<torch::Tensor>& e_score_correction_bias) {
torch::Tensor input_ = input.to(torch::kFloat32);
auto reduce_weight =
torch::empty({input.size(0), topk},
torch::dtype(torch::kFloat).device(input.device()));
auto topk_indices =
torch::empty({input.size(0), topk},
torch::dtype(torch::kInt32).device(input.device()));
auto token_expert_indices =
torch::empty({input.size(0), topk},
torch::dtype(torch::kInt32).device(input.device()));
infer::topk_softmax(
reduce_weight, topk_indices, token_expert_indices, input_, false);
auto tt = reduce_weight.sum(-1);
if (normalize) {
reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1);
}
return std::make_tuple(reduce_weight, topk_indices);
}
std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
int64_t expert_num) {
auto src_dst = expert_id.new_empty({expert_id.numel()});
auto dst_src = torch::empty_like(src_dst);
auto expert_sizes_gpu = expert_id.new_empty({expert_num});
auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
infer::moe_compute_token_index_api(expert_id,
src_dst,
dst_src,
expert_sizes_gpu,
/*expert_mask=*/std::nullopt,
/*expert_sizes_cpu*/ std::nullopt,
/*expert_sizes_gpu*/ std::nullopt,
0,
expert_num,
expert_num);
expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
}
torch::Tensor moe_expand_input(const torch::Tensor& input,
const torch::Tensor& gather_index,
const torch::Tensor& combine_idx,
int64_t topk) {
int64_t dst_tokens = input.size(0) * topk;
auto output = input.new_empty({dst_tokens, input.size(1)});
infer::moe_expand_input(
output, input, combine_idx, gather_index, dst_tokens, topk);
return output;
}
torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) {
input = input.view({-1, weight.size(1), input.size(1)});
auto output = input.new_empty({input.size(0), input.size(2)});
infer::moe_output_reduce_sum(output,
input,
weight,
/*mask=*/std::nullopt,
/*extra_residual*/ std::nullopt,
/*scaling_factor=*/1.0);
return output;
}
} // namespace xllm::kernel::ilu

39
ex_engine/xllm_kernels/ilu/group_gemm.cpp Normal file
View File

@@ -0,0 +1,39 @@
/* 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 "ilu_ops_api.h"
namespace xllm::kernel::ilu {
torch::Tensor group_gemm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& tokens_per_experts,
const std::optional<torch::Tensor>& dst_to_src,
torch::Tensor& output) {
infer::moe_w16a16_group_gemm(
output,
input,
weight,
tokens_per_experts,
dst_to_src,
/*bias=*/std::nullopt,
/*format=*/"TN",
/*persistent=*/0,
/*output_n=*/tokens_per_experts.sum().item<int64_t>());
return output;
}
} // namespace xllm::kernel::ilu

153
ex_engine/xllm_kernels/ilu/ilu_ops_api.h Normal file
View File

@@ -0,0 +1,153 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <ATen/DynamicLibrary.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <cuda_runtime.h>
#include <glog/logging.h>
#include <torch/all.h>
#include <optional>
#include "ATen/Tensor.h"
#include "ATen/cuda/CUDAEvent.h"
#include "c10/core/Device.h"
#include "c10/core/DeviceGuard.h"
#include "c10/core/GradMode.h"
#include "c10/core/InferenceMode.h"
#include "c10/core/MemoryFormat.h"
#include "c10/core/ScalarType.h"
#include "c10/core/TensorOptions.h"
#include "c10/cuda/CUDAFunctions.h"
#include "c10/cuda/CUDAGuard.h"
#include "c10/cuda/CUDAStream.h"
#include "ixformer.h"
#include "kernels/kernels.h"
// #include "utils.h"
using namespace ixformer;
namespace xllm::kernel::ilu {
void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
torch::Tensor& key,
torch::Tensor& cos_sin_cache,
torch::Tensor& positions,
bool interleave);
// act_mode only support silu, gelu, gelu_tanh
void act_and_mul(torch::Tensor out,
torch::Tensor input,
const std::string& act_mode);
void reshape_paged_cache(
torch::Tensor& key, // (num_tokens, num_heads, head_size)
std::optional<torch::Tensor>& value, // (num_tokens, num_heads, head_size)
torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size)
std::optional<torch::Tensor>&
value_cache, // (num_blocks, num_heads, block_size, head_size)
torch::Tensor& slot_mapping); //(num_tokens)
void batch_prefill(torch::Tensor& query,
const torch::Tensor& key,
const std::optional<torch::Tensor>& value,
torch::Tensor& output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_cu_seq_lens,
const std::optional<torch::Tensor>& kv_cu_seq_lens,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& attn_bias,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_quant_scale,
const std::optional<torch::Tensor>& v_quant_scale,
const torch::Tensor& block_tables,
int64_t max_query_len,
int64_t max_seq_len,
float scale,
bool is_causal,
int64_t window_size_left,
int64_t window_size_right,
const std::string& compute_dtype,
bool return_lse);
void batch_decode(torch::Tensor& query,
const torch::Tensor& k_cache,
torch::Tensor& output,
const torch::Tensor& block_table,
const torch::Tensor& seq_lens,
const std::optional<torch::Tensor>& v_cache,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_cache_quant_scale,
const std::optional<torch::Tensor>& v_cache_quant_scale,
const std::optional<torch::Tensor>& out_quant_scale,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& mask,
const std::string& compute_dtype,
int64_t max_seq_len,
int64_t window_size_left,
int64_t window_size_right,
float scale,
bool return_lse,
bool is_causal,
int64_t kv_cache_quant_bit_size);
void residual_layer_norm(torch::Tensor& input,
torch::Tensor& output,
std::optional<torch::Tensor>& residual,
torch::Tensor& weight,
std::optional<torch::Tensor>& bias,
std::optional<torch::Tensor>& residual_out,
double eps);
void rms_norm(torch::Tensor& output,
torch::Tensor& input,
torch::Tensor& weight,
double eps);
torch::Tensor matmul(torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias);
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
const torch::Tensor& input,
int64_t topk,
int64_t num_expert_group,
int64_t topk_group,
bool normalize,
const std::optional<torch::Tensor>& mask,
const std::string& normed_by,
const std::string& scoring_func,
double route_scale,
const std::optional<torch::Tensor>& e_score_correction_bias);
std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
int64_t expert_num);
torch::Tensor moe_expand_input(const torch::Tensor& input,
const torch::Tensor& gather_index,
const torch::Tensor& combine_idx,
int64_t topk);
torch::Tensor group_gemm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& tokens_per_experts,
const std::optional<torch::Tensor>& dst_to_src,
torch::Tensor& output);
torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
} // namespace xllm::kernel::ilu

147
ex_engine/xllm_kernels/ilu/ixformer.h Normal file
View File

@@ -0,0 +1,147 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <torch/all.h>
#include "ATen/Tensor.h"
#include "utils.h"
namespace ixformer::infer {
torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& out,
torch::Tensor& block_tables,
torch::Tensor& cu_seq_q,
torch::Tensor& cu_seq_k,
int64_t max_seq_q,
int64_t max_seq_k,
bool is_causal,
int64_t window_left,
int64_t window_right,
double scale,
double softcap,
bool sqrt_alibi,
const std::optional<torch::Tensor>& alibi_slopes,
const std::optional<torch::Tensor>& sinks,
std::optional<torch::Tensor>& lse);
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
torch::Tensor xllm_paged_attention(
torch::Tensor& out,
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
int64_t num_kv_heads,
double scale,
torch::Tensor& block_tables,
torch::Tensor& context_lens,
int64_t block_size,
int64_t max_context_len,
const std::optional<torch::Tensor>& alibi_slopes,
bool causal,
int32_t window_left,
int32_t window_right,
double softcap,
bool enable_cuda_graph,
bool use_sqrt_alibi,
const std::optional<torch::Tensor>& sinks);
torch::Tensor ixformer_linear(torch::Tensor& input,
torch::Tensor& weight,
int64_t act_type,
const std::optional<torch::Tensor>& bias,
const std::optional<torch::Tensor>& out,
const std::optional<bool> persistent);
torch::Tensor ixformer_linear_ex(torch::Tensor& input,
torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& out);
void xllm_reshape_and_cache(torch::Tensor& key,
torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
int64_t key_token_stride,
int64_t value_token_stride);
void xllm_rotary_embedding(torch::Tensor& positions,
torch::Tensor& query,
torch::Tensor& key,
int64_t head_size,
torch::Tensor& cos_sin_cache,
bool is_neox);
void residual_rms_norm(torch::Tensor& input,
torch::Tensor& residual,
torch::Tensor& weight,
torch::Tensor& output,
torch::Tensor& residual_output,
const std::optional<torch::Tensor>& fused_bias,
double alpha,
double eps,
bool is_post);
void rms_norm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& output,
const std::optional<torch::Tensor>& fused_bias,
double eps);
void topk_softmax(torch::Tensor& topk_weights,
torch::Tensor& topk_indices,
torch::Tensor& token_expert_indices,
torch::Tensor& gating_output,
bool renormalize);
void moe_compute_token_index_api(
torch::Tensor& topk_ids,
torch::Tensor& src_dst,
torch::Tensor& dst_src,
torch::Tensor& expert_sizes_gpu,
const c10::optional<torch::Tensor>& expert_mask,
const c10::optional<torch::Tensor>& expert_sizes_cpu,
const c10::optional<torch::Tensor>& expand_tokens_gpu,
int64_t start_expert_id,
int64_t end_expert_id,
int64_t num_experts);
void moe_expand_input(torch::Tensor outputs,
torch::Tensor inputs,
torch::Tensor dst_to_src,
const c10::optional<torch::Tensor>& src_to_dst,
int64_t dst_tokens,
int64_t expand_factor);
void moe_w16a16_group_gemm(torch::Tensor output,
torch::Tensor inputs,
torch::Tensor weights,
torch::Tensor tokens_per_experts,
const c10::optional<torch::Tensor>& dst_to_src,
const c10::optional<torch::Tensor>& bias,
std::string format,
int64_t persistent,
int64_t output_n);
void moe_output_reduce_sum(torch::Tensor outputs,
torch::Tensor inputs,
const c10::optional<torch::Tensor>& mul_weight,
const c10::optional<torch::Tensor>& mask,
const c10::optional<torch::Tensor>& extra_residual,
double scaling_factor);
} // namespace ixformer::infer

73
ex_engine/xllm_kernels/ilu/matmul.cpp Normal file
View File

@@ -0,0 +1,73 @@
/* 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 "util/env_var.h"
namespace xllm::kernel::ilu {
bool gemv_conditions(const torch::Tensor& input,
const torch::Tensor& weight,
const torch::Tensor& bias,
int64_t gemv_max_batch) {
// gemv input:[m,k] weight:[n,k]
// 1. m <= gemv_max_batch
// 2. k % 32 == 0 && n % 2 == 0
// 3. bias is None
torch::Tensor input_view = input.view({-1, input.size(-1)});
torch::Tensor weight_view = weight.view({-1, weight.size(-1)});
int64_t m = input_view.size(0);
int64_t k = input_view.size(1);
int64_t n = weight_view.size(0);
if (bias.defined() == false && m <= gemv_max_batch && k % 32 == 0 &&
n % 2 == 0) {
return true;
}
return false;
}
torch::Tensor matmul(torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias) {
int64_t act_type = -1;
bool persistent = false;
std::vector<int64_t> output_shape = a.sizes().vec();
if (!output_shape.empty()) {
output_shape[output_shape.size() - 1] = b.size(0);
}
torch::Tensor output = a.new_empty(output_shape);
bool use_gemv = true;
const int64_t gemv_max_batch = 1;
const bool disable_infer_gemm_ex =
xllm::util::get_bool_env("DISABLE_INFER_GEMM_EX", false);
use_gemv =
use_gemv &&
gemv_conditions(a, b, bias.value_or(at::Tensor()), gemv_max_batch) &&
!disable_infer_gemm_ex && (act_type == -1);
if (use_gemv) {
output = infer::ixformer_linear_ex(a, b, bias, output);
} else {
output = infer::ixformer_linear(a, b, act_type, bias, output, persistent);
}
return output;
}
} // namespace xllm::kernel::ilu

51
ex_engine/xllm_kernels/ilu/norm.cpp Normal file
View File

@@ -0,0 +1,51 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
#include "utils.h"
using namespace ixformer;
namespace xllm::kernel::ilu {
void residual_layer_norm(torch::Tensor& input,
torch::Tensor& output,
std::optional<torch::Tensor>& residual,
torch::Tensor& weight,
std::optional<torch::Tensor>& bias,
std::optional<torch::Tensor>& residual_out,
double eps) {
auto residual_ = residual.value_or(torch::zeros_like(input));
torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input));
infer::residual_rms_norm(input,
residual_,
weight,
output,
residual_out_,
bias,
/*alpha=*/1.0,
eps,
false);
}
void rms_norm(torch::Tensor& output,
torch::Tensor& input,
torch::Tensor& weight,
double eps) {
std::optional<torch::Tensor> fused_bias = std::nullopt;
infer::rms_norm(input, weight, output, fused_bias, eps);
}
} // namespace xllm::kernel::ilu

31
ex_engine/xllm_kernels/ilu/rope.cpp Normal file
View File

@@ -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

63
ex_engine/xllm_kernels/ilu/utils.h Normal file
View File

@@ -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

11
ex_engine/xllm_kernels/kernels.h Normal file
View File

@@ -0,0 +1,11 @@
/* Auto-generated aggregation header for xllm::kernel namespace.
* Equivalent to CMake cc_library(NAME kernels HDRS param.h ops_api.h).
*
* AST Layer 3: kernel dispatch interface
* Called by: xllm_layers/ (Layer 2)
* Calls: xllm_kernels/ilu/ (Layer 4)
*/
#pragma once
#include "param.h"
#include "ops_api.h"

59
ex_engine/xllm_kernels/npu/npu_causal_conv1d.cpp Normal file
View File

@@ -0,0 +1,59 @@
/* 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 "core/kernels/npu/aclnn/pytorch_npu_helper.hpp"
#include "core/kernels/npu/utils.h"
#include "core/kernels/npu/xllm_ops/xllm_ops_api.h"
namespace xllm::kernel::npu {
torch::Tensor causal_conv1d(const torch::Tensor& x,
const torch::Tensor& weight,
const torch::Tensor& conv_state,
const std::optional<torch::Tensor>& bias_opt,
const torch::IntArrayRef query_start_loc_opt,
const torch::IntArrayRef cache_indices_opt,
const torch::IntArrayRef initial_state_mode_opt,
const torch::IntArrayRef num_accepted_tokens_opt,
int64_t activation_mode,
int64_t pad_slot_id,
int64_t run_mode) {
check_tensor(x, "x", "causal_conv1d");
check_tensor(weight, "weight", "causal_conv1d");
check_tensor(conv_state, "conv_state", "causal_conv1d");
c10::optional<torch::Tensor> bias_tensor = c10::nullopt;
if (bias_opt.has_value() && bias_opt.value().defined()) {
bias_tensor = bias_opt.value();
}
torch::Tensor output = torch::empty(x.sizes(), x.options());
EXEC_NPU_CMD(aclnnCausalConv1d,
x,
weight,
bias_tensor,
conv_state,
query_start_loc_opt,
cache_indices_opt,
initial_state_mode_opt,
num_accepted_tokens_opt,
activation_mode,
pad_slot_id,
run_mode,
output);
return output;
}
} // namespace xllm::kernel::npu

83
ex_engine/xllm_kernels/npu/npu_recurrent_gated_delta_rule.cpp Normal file
View File

@@ -0,0 +1,83 @@
/* 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 <glog/logging.h>
#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp"
#include "core/kernels/npu/npu_ops_api.h"
#include "core/kernels/npu/utils.h"
namespace {
c10::optional<torch::Tensor> to_c10_optional_tensor(
const std::optional<torch::Tensor>& tensor_opt) {
if (tensor_opt.has_value() && tensor_opt.value().defined()) {
return tensor_opt.value();
}
return c10::nullopt;
}
} // namespace
namespace xllm::kernel::npu {
torch::Tensor npu_recurrent_gated_delta_rule(
const torch::Tensor& query,
const torch::Tensor& key,
const torch::Tensor& value,
torch::Tensor& state,
const std::optional<torch::Tensor>& beta,
const std::optional<double> scale,
const std::optional<torch::Tensor>& actual_seq_lengths,
const std::optional<torch::Tensor>& ssm_state_indices,
const std::optional<torch::Tensor>& num_accepted_tokens,
const std::optional<torch::Tensor>& g,
const std::optional<torch::Tensor>& gk) {
check_tensor(query, "query", "recurrent_gated_delta_rule");
check_tensor(key, "key", "recurrent_gated_delta_rule");
check_tensor(value, "value", "recurrent_gated_delta_rule");
check_tensor(state, "state", "recurrent_gated_delta_rule");
CHECK(scale.has_value())
<< "recurrent_gated_delta_rule requires a valid scale value";
c10::optional<torch::Tensor> beta_tensor = to_c10_optional_tensor(beta);
c10::optional<torch::Tensor> actual_seq_lengths_tensor =
to_c10_optional_tensor(actual_seq_lengths);
c10::optional<torch::Tensor> ssm_state_indices_tensor =
to_c10_optional_tensor(ssm_state_indices);
c10::optional<torch::Tensor> num_accepted_tokens_tensor =
to_c10_optional_tensor(num_accepted_tokens);
c10::optional<torch::Tensor> g_tensor = to_c10_optional_tensor(g);
c10::optional<torch::Tensor> gk_tensor = to_c10_optional_tensor(gk);
float scale_value = static_cast<float>(scale.value());
torch::Tensor output = torch::empty_like(value);
EXEC_NPU_CMD(aclnnRecurrentGatedDeltaRule,
query,
key,
value,
beta_tensor,
state,
actual_seq_lengths_tensor,
ssm_state_indices_tensor,
g_tensor,
gk_tensor,
num_accepted_tokens_tensor,
scale_value,
output);
return output;
}
} // namespace xllm::kernel::npu

1101
ex_engine/xllm_kernels/ops_api.cpp Normal file
View File

File diff suppressed because it is too large Load Diff

177
ex_engine/xllm_kernels/ops_api.h Normal file
View File

@@ -0,0 +1,177 @@
/* 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 "param.h"
namespace xllm::kernel {
static const std::string kActModeSilu = "silu";
static const std::string kActModeGelu = "gelu";
static const std::string kActModeQuickGelu = "quick_gelu";
static const std::string kActModeSwish = "swish";
void apply_rotary(RotaryParams& params);
void active(ActivationParams& params);
void reshape_paged_cache(ReshapePagedCacheParams& params);
void reshape_from_cache(ReshapeFromCacheParams& params);
// Quantize and store KV cache to paged cache (INT8 quantization)
// Only supported on MLU backend
void quant_to_paged_cache(ReshapePagedCacheParams& params);
// Dequantize KV cache from paged cache (INT8 to FP16/BF16)
// Only supported on MLU backend
void dequant_from_paged_cache(ReshapeFromCacheParams& params);
void fused_layernorm(FusedLayerNormParams& params);
torch::Tensor matmul(MatmulParams& params);
torch::Tensor group_gemm(GroupGemmParams& params);
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
MoeFusedTopkParams& params);
std::vector<torch::Tensor> moe_gen_idx(MoeGenIdxParams& params);
torch::Tensor moe_expand_input(MoeExpandInputParams& params);
torch::Tensor moe_combine_result(MoeCombineResultParams& params);
torch::Tensor moe_all2all_gen_send_layout(
MoeAll2AllGenSendLayoutParams& params);
std::vector<torch::Tensor> moe_all2all_gen_gather_index(
MoeAll2AllGenGatherIndexParams& params);
std::vector<torch::Tensor> moe_all2all_create(MoeAll2AllCreateParams& params);
void moe_all2all_init(MoeAll2AllInitParams& params);
void moe_all2all_dispatch(MoeAll2AllDispatchParams& params);
void moe_all2all_combine(MoeAll2AllCombineParams& params);
void moe_all2all_destroy(MoeAll2AllDestroyParams& params);
std::tuple<torch::Tensor, torch::Tensor> scaled_quantize(
ScaledQuantizeParams& params);
torch::Tensor scaled_matmul(ScaledMatmulParams& params);
torch::Tensor apply_top_k_top_p(TopKPParams& params);
torch::Tensor random_sample(RandomSampleParams& params);
torch::Tensor rejection_sample(RejectionSampleParams& params);
void masked_indexer_select_paged_kv(MaskedIndexerSelectPagedKVParams& params);
void gather_split(GatherSplitParams& params);
void fused_mla_q(FusedMlaQParams& params);
void fused_mla_kv(FusedMlaKVParams& params);
void fused_indexer_q(FusedIndexerQParams& params);
void fused_indexer_k(FusedIndexerKParams& params);
// L2 normalization along the last dimension
torch::Tensor l2_norm(torch::Tensor& x, double eps = 1e-6);
// TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx + moe_expand_input
// (and token_count/cusum outputs) on other backends.
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
moe_init_routing_v2(MoeInitRoutingV2Params& params);
// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format
// Returns: (quantized_output, scale)
std::tuple<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
Fp8ScaledQuantizeParams& params);
// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels
// Performs: c = (a @ b.T) with scales applied
torch::Tensor fp8_scaled_matmul(Fp8ScaledMatmulParams& params);
// Static scaled FP8 quantization helper
// Quantizes input tensor to FP8 using a pre-computed scale factor
void static_scaled_fp8_quant(StaticScaledFp8QuantParams& params);
// Fused RMSNorm + Static FP8 Quantization
// These fused operations combine RMSNorm and FP8 quantization to reduce memory
// bandwidth by avoiding the intermediate write-back to global memory.
// Fused RMSNorm + Static FP8 Quantization
// Returns: FP8 quantized output tensor
torch::Tensor rms_norm_static_fp8_quant(RmsNormStaticFp8QuantParams& params);
// Fused Add + RMSNorm + Static FP8 Quantization (with residual)
// Returns: tuple of (FP8 quantized output, updated residual)
std::tuple<torch::Tensor, torch::Tensor> fused_add_rms_norm_static_fp8_quant(
FusedAddRmsNormStaticFp8QuantParams& params);
std::pair<torch::Tensor, torch::Tensor> fused_gdn_gating(
FusedGdnGatingParams& params);
std::pair<torch::Tensor, torch::Tensor> fused_recurrent_gated_delta_rule(
FusedRecurrentGatedDeltaRuleParams& params);
torch::Tensor causal_conv1d_update(CausalConv1dUpdateParams& params);
torch::Tensor gated_layer_norm(GatedLayerNormParams& params);
std::pair<torch::Tensor, torch::Tensor> partial_rotary_embedding(
PartialRotaryEmbeddingParams& params);
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
fused_qkvzba_split_reshape_cat(FusedQkvzbaSplitReshapeParams& params);
void gemma_rms_norm(GemmaRMSNormParams& params);
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
split_qkv_rmsnorm_mrope(SplitQkvRmsnormMropeParams& params);
bool has_split_qkv_rmsnorm_mrope_specialization(int64_t num_q_heads,
int64_t num_kv_heads,
int64_t head_size);
torch::Tensor build_split_qkv_rmsnorm_mrope_gather_pattern(
int64_t rope_dim,
const std::vector<int64_t>& mrope_section,
bool is_interleaved,
const torch::Device& device);
std::pair<torch::Tensor, torch::Tensor> chunk_gated_delta_rule(
ChunkGatedDeltaRuleParams& params);
torch::Tensor recurrent_gated_delta_rule(
const torch::Tensor& query,
const torch::Tensor& key,
const torch::Tensor& value,
torch::Tensor& state,
const std::optional<torch::Tensor>& beta,
const std::optional<double> scale,
const std::optional<torch::Tensor>& actual_seq_lengths,
const std::optional<torch::Tensor>& ssm_state_indices,
const std::optional<torch::Tensor>& num_accepted_tokens,
const std::optional<torch::Tensor>& g,
const std::optional<torch::Tensor>& gk);
} // namespace xllm::kernel

1441
ex_engine/xllm_kernels/param.h Normal file
View File

File diff suppressed because it is too large Load Diff

122
ex_engine/xllm_kernels/rebuild_test_k10.sh Executable file
View File

@@ -0,0 +1,122 @@
#!/bin/bash
# rebuild_test_k10.sh — Clean rebuild and test kernel 10 Config B
set -eo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CUDA_DIR="${SCRIPT_DIR}/cuda"
echo "=== Clean old builds ==="
rm -rf "${SCRIPT_DIR}/build/tmp_hgemm_warptiling"
rm -f "${SCRIPT_DIR}/build/hgemm_warptiling.so"
echo "=== Compile ==="
python3 -c "
import torch.utils.cpp_extension as ext
import os, shutil, glob
name = 'hgemm_warptiling'
build_dir = '${SCRIPT_DIR}/build/tmp_' + name
os.makedirs(build_dir, exist_ok=True)
mod = ext.load(
name=name,
sources=[
'${CUDA_DIR}/hgemm_warptiling.cu',
'${CUDA_DIR}/bindings/hgemm_warp_bind.cpp',
],
extra_include_paths=['${CUDA_DIR}/headers'],
extra_cflags=['-O2', '-std=c++17'],
extra_cuda_cflags=['-O2'],
build_directory=build_dir,
verbose=True,
)
built = glob.glob(build_dir + '/' + name + '*.so')
if built:
dst = '${SCRIPT_DIR}/build/' + name + '.so'
shutil.copy2(built[0], dst)
print(f'[build] SUCCESS: {dst}')
"
echo ""
echo "=== Test ==="
python3 << 'PYTEST'
import torch, sys, os, glob, time, importlib.util
build_dir = 'ex_engine/xllm_kernels/build'
so = glob.glob(f'{build_dir}/tmp_hgemm_warptiling/hgemm_warptiling*.so')
if not so:
print("SKIP: .so not found")
sys.exit(0)
spec = importlib.util.spec_from_file_location("hgemm_warptiling", so[0])
hw = importlib.util.module_from_spec(spec)
spec.loader.exec_module(hw)
print(f"Loaded: {so[0]}")
# Test 1: tiny
print("\n--- 16x16 @ 16x16 ---")
A = torch.eye(16, dtype=torch.float16, device='cuda')
B = torch.ones(16, 16, dtype=torch.float16, device='cuda')
C = hw.hgemm_warp(A, B)
diff = (C.float() - B.float()).abs().max().item()
print(f" I @ ones = ones? diff={diff:.6f}")
# Test 2: 128x128
print("\n--- 128x128 @ 128x128 ---")
A = torch.randn(128, 128, dtype=torch.float16, device='cuda') * 0.1
B = torch.randn(128, 128, dtype=torch.float16, device='cuda') * 0.1
C_ref = torch.matmul(A.float(), B.float()).half()
C_k10 = hw.hgemm_warp(A, B)
diff = (C_ref.float() - C_k10.float()).abs().max().item()
print(f" max_diff={diff:.6f}")
if diff > 2.0:
# Debug: print a few values
print(f" C_ref[0,:5] = {C_ref[0,:5].tolist()}")
print(f" C_k10[0,:5] = {C_k10[0,:5].tolist()}")
print(f" C_ref[-1,-5:] = {C_ref[-1,-5:].tolist()}")
print(f" C_k10[-1,-5:] = {C_k10[-1,-5:].tolist()}")
print(" FAIL")
else:
print(" PASS")
# Test 3: MoE size
print("\n--- 256x4096 @ 4096x11008 ---")
A = torch.randn(256, 4096, dtype=torch.float16, device='cuda') * 0.01
B = torch.randn(4096, 11008, dtype=torch.float16, device='cuda') * 0.01
C_ref = torch.matmul(A.float(), B.float()).half()
C_k10 = hw.hgemm_warp(A, B)
diff = (C_ref.float() - C_k10.float()).abs().max().item()
rel = diff / (C_ref.float().abs().max().item() + 1e-8)
print(f" max_diff={diff:.6f}, rel={rel:.6f}")
if diff > 2.0:
print(f" C_ref[0,:5] = {C_ref[0,:5].tolist()}")
print(f" C_k10[0,:5] = {C_k10[0,:5].tolist()}")
print(" FAIL")
else:
print(" PASS")
# Test 4: Performance
print("\n--- Performance 256x4096 @ 4096x11008 ---")
for _ in range(10):
hw.hgemm_warp(A, B)
torch.cuda.synchronize()
t0 = time.time()
for _ in range(100):
hw.hgemm_warp(A, B)
torch.cuda.synchronize()
ms_k10 = (time.time() - t0) / 100 * 1000
for _ in range(10):
torch.matmul(A, B)
torch.cuda.synchronize()
t0 = time.time()
for _ in range(100):
torch.matmul(A, B)
torch.cuda.synchronize()
ms_torch = (time.time() - t0) / 100 * 1000
print(f" kernel 10: {ms_k10:.2f} ms")
print(f" torch.matmul: {ms_torch:.2f} ms")
print(f" ratio: {ms_k10/ms_torch:.2f}x")
PYTEST