feat: CUTLASS Cu10 grouped GEMM — real device verified
BI-V100 real device results: moe_group_gemm: err=0.000015 PASS moe_decode_cutlass: NaN=False PASS cutlass grouped: 4.77ms vs torch.mm loop: 9.38ms → 1.97x speedup Fix: gemm_grouped.cu ldb=K (not N) for ColumnMajor B view Link: -lcuinfer from /usr/local/corex-3.2.3/lib64/libcuinfer.so.7
This commit is contained in:
root
203
bench_gemm.py
Normal file
203
bench_gemm.py
Normal file
@@ -0,0 +1,203 @@
|
||||
"""bench_gemm.py — Benchmark all GEMM backends on real device.
|
||||
|
||||
Tests with Qwen3.5-27B MoE shapes:
|
||||
- Decode: M=1, K=3584, N=18944*2 (gate_up) / N=3584 (down)
|
||||
- Prefill: M=variable, same K/N
|
||||
|
||||
Usage:
|
||||
python3 bench_gemm.py
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
|
||||
# Qwen3.5-27B params (per TP=4 partition)
|
||||
H = 3584 # hidden_size
|
||||
I = 18944 // 4 # intermediate per partition (4736)
|
||||
TWO_I = I * 2 # gate + up
|
||||
NUM_EXPERTS = 128
|
||||
TOPK = 8
|
||||
|
||||
WARMUP = 5
|
||||
REPEATS = 20
|
||||
|
||||
|
||||
def bench_fn(fn, *args, name=""):
|
||||
"""Benchmark a function, return ms per call."""
|
||||
for _ in range(WARMUP):
|
||||
fn(*args)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(REPEATS):
|
||||
fn(*args)
|
||||
torch.cuda.synchronize()
|
||||
elapsed = (time.perf_counter() - t0) / REPEATS * 1000
|
||||
print(f" {name}: {elapsed:.3f} ms")
|
||||
return elapsed
|
||||
|
||||
|
||||
def bench_single_gemm(device):
|
||||
"""Benchmark single GEMM: (M,K) × (K,N) for various M."""
|
||||
print("\n=== Single GEMM (M,K)×(K,N) ===")
|
||||
for M in [1, 4, 8, 32]:
|
||||
A = torch.randn(M, H, device=device, dtype=torch.float16)
|
||||
B = torch.randn(H, TWO_I, device=device, dtype=torch.float16)
|
||||
|
||||
bench_fn(torch.mm, A, B, name=f"torch.mm M={M} K={H} N={TWO_I}")
|
||||
|
||||
# Try hgemm
|
||||
try:
|
||||
import hgemm
|
||||
bench_fn(hgemm.hgemm, A, B, name=f"hgemm M={M}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Try ixformer linear
|
||||
try:
|
||||
import ix_moe_bridge as bridge
|
||||
bench_fn(bridge.linear, A, B.t().contiguous(), name=f"ixformer_linear M={M}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def bench_group_gemm(device):
|
||||
"""Benchmark group GEMM with MoE shapes."""
|
||||
print("\n=== Group GEMM (MoE w13 projection) ===")
|
||||
|
||||
# Simulate decode: 1 token → topk=8 experts, each gets ~1 token
|
||||
total_tokens = TOPK
|
||||
expert_counts = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32)
|
||||
# Distribute tokens to first TOPK experts
|
||||
for i in range(TOPK):
|
||||
expert_counts[i] = 1
|
||||
|
||||
input_t = torch.randn(total_tokens, H, device=device, dtype=torch.float16)
|
||||
w13 = torch.randn(NUM_EXPERTS, TWO_I, H, device=device, dtype=torch.float16) * 0.01
|
||||
|
||||
# PyTorch baseline
|
||||
def torch_group_gemm():
|
||||
offset = 0
|
||||
out = torch.zeros(total_tokens, TWO_I, device=device, dtype=torch.float16)
|
||||
for e in range(NUM_EXPERTS):
|
||||
c = expert_counts[e].item()
|
||||
if c <= 0: continue
|
||||
out[offset:offset+c] = torch.mm(input_t[offset:offset+c], w13[e].t())
|
||||
offset += c
|
||||
return out
|
||||
|
||||
bench_fn(torch_group_gemm, name=f"torch.mm loop (decode, {TOPK} experts)")
|
||||
|
||||
# Try gemm_grouped
|
||||
try:
|
||||
import gemm_grouped
|
||||
bench_fn(gemm_grouped.moe_group_gemm, input_t, w13, expert_counts,
|
||||
name=f"cutlass_grouped (decode, {TOPK} experts)")
|
||||
except Exception as e:
|
||||
print(f" cutlass_grouped: {e}")
|
||||
|
||||
# Try ix_moe_bridge
|
||||
try:
|
||||
import ix_moe_bridge as bridge
|
||||
bench_fn(bridge.group_gemm, input_t, w13, expert_counts, TWO_I,
|
||||
name=f"cuinfer_group_gemm (decode, {TOPK} experts)")
|
||||
except Exception as e:
|
||||
print(f" cuinfer_group_gemm: {e}")
|
||||
|
||||
# Try hgemm
|
||||
try:
|
||||
import hgemm
|
||||
bench_fn(hgemm.moe_expert_gemm, input_t, w13, expert_counts,
|
||||
name=f"hgemm_expert (decode, {TOPK} experts)")
|
||||
except Exception as e:
|
||||
print(f" hgemm_expert: {e}")
|
||||
|
||||
# Prefill shape: 32 tokens
|
||||
print("\n=== Group GEMM (MoE w13, prefill M=32) ===")
|
||||
total_pf = 32 * TOPK # 256
|
||||
expert_counts_pf = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32)
|
||||
for i in range(total_pf):
|
||||
expert_counts_pf[i % NUM_EXPERTS] += 1
|
||||
input_pf = torch.randn(total_pf, H, device=device, dtype=torch.float16)
|
||||
|
||||
def torch_group_gemm_pf():
|
||||
offset = 0
|
||||
out = torch.zeros(total_pf, TWO_I, device=device, dtype=torch.float16)
|
||||
for e in range(NUM_EXPERTS):
|
||||
c = expert_counts_pf[e].item()
|
||||
if c <= 0: continue
|
||||
out[offset:offset+c] = torch.mm(input_pf[offset:offset+c], w13[e].t())
|
||||
offset += c
|
||||
return out
|
||||
|
||||
bench_fn(torch_group_gemm_pf, name=f"torch.mm loop (prefill, 256 tokens)")
|
||||
|
||||
try:
|
||||
import gemm_grouped
|
||||
bench_fn(gemm_grouped.moe_group_gemm, input_pf, w13, expert_counts_pf,
|
||||
name=f"cutlass_grouped (prefill, 256 tokens)")
|
||||
except Exception as e:
|
||||
print(f" cutlass_grouped: {e}")
|
||||
|
||||
|
||||
def bench_decode_fused(device):
|
||||
"""Benchmark full MoE decode pipeline."""
|
||||
print("\n=== Full MoE Decode (1 token, topk=8) ===")
|
||||
hidden = torch.randn(1, H, device=device, dtype=torch.float16)
|
||||
w13_sel = torch.randn(TOPK, TWO_I, H, device=device, dtype=torch.float16) * 0.01
|
||||
w2_sel = torch.randn(TOPK, H, I, device=device, dtype=torch.float16) * 0.01
|
||||
topk_w = torch.softmax(torch.randn(TOPK), dim=0).to(device)
|
||||
|
||||
# PyTorch baseline
|
||||
def torch_decode():
|
||||
results = []
|
||||
for k in range(TOPK):
|
||||
gu = torch.mm(hidden, w13_sel[k].t())
|
||||
act = torch.silu(gu[:, :I]) * gu[:, I:]
|
||||
down = torch.mm(act, w2_sel[k].t())
|
||||
results.append(down * topk_w[k])
|
||||
return sum(results)
|
||||
|
||||
bench_fn(torch_decode, name="torch.mm loop")
|
||||
|
||||
try:
|
||||
import gemm_grouped
|
||||
bench_fn(gemm_grouped.moe_decode_cutlass,
|
||||
hidden, w13_sel, w2_sel, topk_w,
|
||||
name="cutlass_batched")
|
||||
except Exception as e:
|
||||
print(f" cutlass_batched: {e}")
|
||||
|
||||
try:
|
||||
import corex_batched_gemm
|
||||
bench_fn(corex_batched_gemm.moe_decode_fused,
|
||||
hidden, w13_sel, w2_sel, topk_w,
|
||||
name="corex_batched")
|
||||
except Exception as e:
|
||||
print(f" corex_batched: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
if not torch.cuda.is_available():
|
||||
print("No CUDA, skipping")
|
||||
sys.exit(0)
|
||||
|
||||
device = torch.device("cuda:0")
|
||||
print(f"Device: {torch.cuda.get_device_name(0)}")
|
||||
print(f"Shapes: H={H}, I={I}, 2I={TWO_I}, experts={NUM_EXPERTS}, topk={TOPK}")
|
||||
|
||||
bench_single_gemm(device)
|
||||
bench_group_gemm(device)
|
||||
bench_decode_fused(device)
|
||||
|
||||
print("\n=== Active backend ===")
|
||||
try:
|
||||
from gemm_dispatch import get_backend
|
||||
print(f" gemm_dispatch: {get_backend()}")
|
||||
except Exception:
|
||||
print(" gemm_dispatch not loaded")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
56
ex_engine/build_cuinfer_gemm.sh
Normal file
56
ex_engine/build_cuinfer_gemm.sh
Normal file
@@ -0,0 +1,56 @@
|
||||
#!/usr/bin/env bash
|
||||
# build_cuinfer_gemm.sh — Compile cuinfer GEMM wrapper
|
||||
#
|
||||
# Links: libcuinfer.so (from /usr/local/corex/lib64/)
|
||||
# Output: cuinfer_gemm_wrapper.so
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
SRC="${SCRIPT_DIR}/cuinfer_gemm_wrapper.cu"
|
||||
HDR="${SCRIPT_DIR}/cuinfer_handle.h"
|
||||
|
||||
echo "[cuinfer_gemm] Building cuinfer_gemm_wrapper.so"
|
||||
|
||||
COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
|
||||
CUINFER_LIB=""
|
||||
for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib"; do
|
||||
if [[ -f "${d}/libcuinfer.so" ]]; then
|
||||
CUINFER_LIB="${d}"
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
python3 << PYEOF
|
||||
import os, sys, shutil
|
||||
|
||||
src = "${SRC}"
|
||||
hdr_dir = "${SCRIPT_DIR}"
|
||||
cuinfer_lib = "${CUINFER_LIB}"
|
||||
|
||||
ldflags = []
|
||||
if cuinfer_lib:
|
||||
ldflags = [f"-L{cuinfer_lib}", "-lcuinfer", f"-Wl,-rpath,{cuinfer_lib}"]
|
||||
|
||||
try:
|
||||
from torch.utils.cpp_extension import load
|
||||
mod = load(
|
||||
name="cuinfer_gemm_wrapper",
|
||||
sources=[src],
|
||||
extra_include_paths=[hdr_dir],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
extra_cuda_cflags=["-O2"],
|
||||
extra_ldflags=ldflags,
|
||||
verbose=True,
|
||||
)
|
||||
print("[cuinfer_gemm] ✓ OK")
|
||||
|
||||
import importlib
|
||||
spec = importlib.util.find_spec("cuinfer_gemm_wrapper")
|
||||
if spec and spec.origin:
|
||||
shutil.copy2(spec.origin, os.path.join(hdr_dir, "cuinfer_gemm_wrapper.so"))
|
||||
print(f"[cuinfer_gemm] ✓ Saved")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[cuinfer_gemm] ERROR: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
PYEOF
|
||||
80
ex_engine/build_gemm_grouped.sh
Normal file
80
ex_engine/build_gemm_grouped.sh
Normal file
@@ -0,0 +1,80 @@
|
||||
#!/usr/bin/env bash
|
||||
# build_gemm_grouped.sh — Compile grouped GEMM kernel + bindings
|
||||
#
|
||||
# Requires: corex clang/16 + cutlass headers (on BI-V100 device)
|
||||
# Output: gemm_grouped.so (importable from Python)
|
||||
#
|
||||
# Reference: ex_engine/xllm_kernels/build_test_cutlass_batched.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
# Source files
|
||||
GEMM_CU="${SCRIPT_DIR}/csrc/gemm_grouped.cu"
|
||||
BIND_CPP="${SCRIPT_DIR}/csrc/gemm_grouped_bind.cpp"
|
||||
BATCHED_CU="${SCRIPT_DIR}/../xllm_kernels/cuda/corex_batched_gemm_kernel.cu"
|
||||
|
||||
echo "[gemm] Building gemm_grouped.so"
|
||||
|
||||
# Find cutlass include path
|
||||
SAMPLES="/usr/local/corex-samples-3.2.3_x86_64/samples/cutlass"
|
||||
CUTLASS_INCLUDE=""
|
||||
for d in "${SAMPLES}/include" "/usr/local/corex/include/cutlass" "/usr/include/cutlass"; do
|
||||
if [[ -d "$d" ]]; then
|
||||
CUTLASS_INCLUDE="$d"
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
if [[ -z "$CUTLASS_INCLUDE" ]]; then
|
||||
echo "[gemm] ERROR: cutlass include not found"
|
||||
exit 1
|
||||
fi
|
||||
echo "[gemm] cutlass: ${CUTLASS_INCLUDE}"
|
||||
|
||||
python3 << PYEOF
|
||||
import os, sys, shutil
|
||||
|
||||
script_dir = "${SCRIPT_DIR}"
|
||||
cutlass_inc = "${CUTLASS_INCLUDE}"
|
||||
|
||||
sources = [
|
||||
"${GEMM_CU}",
|
||||
"${BIND_CPP}",
|
||||
"${BATCHED_CU}",
|
||||
]
|
||||
sources = [s for s in sources if os.path.isfile(s)]
|
||||
|
||||
print(f"[gemm] Compiling {len(sources)} source files")
|
||||
for s in sources:
|
||||
print(f" {os.path.basename(s)}")
|
||||
|
||||
try:
|
||||
from torch.utils.cpp_extension import load
|
||||
mod = load(
|
||||
name="gemm_grouped",
|
||||
sources=sources,
|
||||
extra_include_paths=[cutlass_inc, script_dir],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
extra_ldflags=["/usr/local/corex/lib64/libcuinfer.so", "-Wl,-rpath,/usr/local/corex/lib64"],
|
||||
extra_cuda_cflags=["-O2", "",
|
||||
f"-I{cutlass_inc}"],
|
||||
verbose=True,
|
||||
)
|
||||
print("[gemm] ✓ Compilation successful")
|
||||
|
||||
import importlib
|
||||
spec = importlib.util.find_spec("gemm_grouped")
|
||||
if spec and spec.origin:
|
||||
dst = os.path.join(script_dir, "gemm_grouped.so")
|
||||
shutil.copy2(spec.origin, dst)
|
||||
print(f"[gemm] ✓ Saved to {dst}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[gemm] ERROR: {e}", file=sys.stderr)
|
||||
import traceback; traceback.print_exc()
|
||||
sys.exit(1)
|
||||
PYEOF
|
||||
|
||||
echo "[gemm] Done"
|
||||
161
ex_engine/csrc/cuinfer_gemm_wrapper.cu
Normal file
161
ex_engine/csrc/cuinfer_gemm_wrapper.cu
Normal file
@@ -0,0 +1,161 @@
|
||||
// cuinfer_gemm_wrapper.cu — Wrapper around cuinferCustomGemm
|
||||
//
|
||||
// ixformer::functions::cuinfer_gemm exists in libixformer.so but
|
||||
// takes ixformer::Tensor (not torch::Tensor). We need a torch-compatible
|
||||
// wrapper that calls the C API directly.
|
||||
//
|
||||
// Symbol dump shows cuinferCustomGemm in libcuinfer.so with signature:
|
||||
// cuinferCustomGemm(handle, stream, ptrMode, transa, transb,
|
||||
// m, n, k, alpha, A, Atype, lda, strideA,
|
||||
// B, Btype, ldb, strideB, beta,
|
||||
// C, Ctype, ldc, strideC, batchCount,
|
||||
// computeType, scaleType, customHostPtr, customDevicePtr, customOption)
|
||||
//
|
||||
// Reference:
|
||||
// cat_files/ixinfer.h — cuinferCustomGemm signature
|
||||
// libixformer.so — ixformer::functions::cuinfer_gemm (confirmed in symbol dump)
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include "cuinfer_handle.h"
|
||||
|
||||
// cuinferCustomGemm is already declared in cuinfer_handle.h extern "C" block
|
||||
// We add the full signature here
|
||||
extern "C" {
|
||||
int cuinferCustomGemm(
|
||||
cuinferHandle_t handle, cudaStream_t stream,
|
||||
int ptrMode, int transa, int transb,
|
||||
int m, int n, int k,
|
||||
const void* alpha,
|
||||
const void* A, int Atype, int lda, long long int strideA,
|
||||
const void* B, int Btype, int ldb, long long int strideB,
|
||||
const void* beta,
|
||||
void* C, int Ctype, int ldc, long long int strideC,
|
||||
int batchCount, int computeType, int scaleType,
|
||||
const void* customHostPtr, const void* customDevicePtr, int customOption);
|
||||
}
|
||||
|
||||
// CUDA_R_16F = 2, CUDA_R_32F = 0 (from cudaDataType_t)
|
||||
static constexpr int kFP16 = 2;
|
||||
static constexpr int kFP32 = 0;
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// cuinfer_gemm: C = alpha * A @ B + beta * C
|
||||
//
|
||||
// A: (M, K) row-major fp16
|
||||
// B: (K, N) row-major fp16 (or (N, K) if transb)
|
||||
// C: (M, N) row-major fp16
|
||||
// ============================================================================
|
||||
torch::Tensor cuinfer_gemm(
|
||||
torch::Tensor A, // (M, K)
|
||||
torch::Tensor B, // (K, N) or (N, K) if trans_b
|
||||
bool trans_b)
|
||||
{
|
||||
TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA");
|
||||
TORCH_CHECK(A.scalar_type() == torch::kHalf, "A must be fp16");
|
||||
TORCH_CHECK(B.scalar_type() == torch::kHalf, "B must be fp16");
|
||||
|
||||
int M = A.size(0);
|
||||
int K = A.size(1);
|
||||
int N = trans_b ? B.size(0) : B.size(1);
|
||||
|
||||
if (!trans_b) {
|
||||
TORCH_CHECK(B.size(0) == K, "B rows must equal K");
|
||||
} else {
|
||||
TORCH_CHECK(B.size(1) == K, "B cols must equal K when transposed");
|
||||
}
|
||||
|
||||
auto C = torch::zeros({M, N}, A.options());
|
||||
auto stream = c10::cuda::getCurrentCUDAStream().stream();
|
||||
auto handle = CuinferHandle::get(stream);
|
||||
|
||||
if (!handle) {
|
||||
// Fallback to torch::mm
|
||||
if (trans_b) {
|
||||
return torch::mm(A.to(torch::kFloat32), B.t().to(torch::kFloat32)).to(torch::kHalf);
|
||||
}
|
||||
return torch::mm(A.to(torch::kFloat32), B.to(torch::kFloat32)).to(torch::kHalf);
|
||||
}
|
||||
|
||||
float alpha = 1.0f, beta = 0.0f;
|
||||
int transa = 0; // N = no transpose
|
||||
int transb_flag = trans_b ? 1 : 0;
|
||||
|
||||
int lda = K;
|
||||
int ldb = trans_b ? K : N;
|
||||
int ldc = N;
|
||||
|
||||
int status = cuinferCustomGemm(
|
||||
handle, stream,
|
||||
0, // CUINFER_POINTER_MODE_HOST
|
||||
transa, transb_flag,
|
||||
M, N, K,
|
||||
&alpha,
|
||||
A.data_ptr(), kFP16, lda, 0,
|
||||
B.data_ptr(), kFP16, ldb, 0,
|
||||
&beta,
|
||||
C.data_ptr(), kFP16, ldc, 0,
|
||||
1, // batchCount
|
||||
kFP32, kFP32, // computeType, scaleType
|
||||
nullptr, nullptr, 0);
|
||||
|
||||
TORCH_CHECK(status == 0, "cuinferCustomGemm failed with status ", status);
|
||||
return C;
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// cuinfer_gemm_batched: batched version
|
||||
// A: (batch, M, K), B: (batch, K, N) or (batch, N, K)
|
||||
// ============================================================================
|
||||
torch::Tensor cuinfer_gemm_batched(
|
||||
torch::Tensor A,
|
||||
torch::Tensor B,
|
||||
bool trans_b)
|
||||
{
|
||||
TORCH_CHECK(A.dim() == 3 && B.dim() == 3, "inputs must be 3D");
|
||||
|
||||
int batch = A.size(0);
|
||||
int M = A.size(1);
|
||||
int K = A.size(2);
|
||||
int N = trans_b ? B.size(1) : B.size(2);
|
||||
|
||||
auto C = torch::zeros({batch, M, N}, A.options());
|
||||
auto stream = c10::cuda::getCurrentCUDAStream().stream();
|
||||
auto handle = CuinferHandle::get(stream);
|
||||
|
||||
float alpha = 1.0f, beta = 0.0f;
|
||||
int lda = K, ldb = trans_b ? K : N, ldc = N;
|
||||
long long strideA = (long long)M * K;
|
||||
long long strideB = trans_b ? (long long)N * K : (long long)K * N;
|
||||
long long strideC = (long long)M * N;
|
||||
|
||||
int status = cuinferCustomGemm(
|
||||
handle, stream,
|
||||
0,
|
||||
0, trans_b ? 1 : 0,
|
||||
M, N, K,
|
||||
&alpha,
|
||||
A.data_ptr(), kFP16, lda, strideA,
|
||||
B.data_ptr(), kFP16, ldb, strideB,
|
||||
&beta,
|
||||
C.data_ptr(), kFP16, ldc, strideC,
|
||||
batch,
|
||||
kFP32, kFP32,
|
||||
nullptr, nullptr, 0);
|
||||
|
||||
TORCH_CHECK(status == 0, "cuinferCustomGemm batched failed: ", status);
|
||||
return C;
|
||||
}
|
||||
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("cuinfer_gemm", &cuinfer_gemm,
|
||||
"GEMM via cuinferCustomGemm (fp16, Cu10)",
|
||||
py::arg("A"), py::arg("B"), py::arg("trans_b") = false);
|
||||
m.def("cuinfer_gemm_batched", &cuinfer_gemm_batched,
|
||||
"Batched GEMM via cuinferCustomGemm",
|
||||
py::arg("A"), py::arg("B"), py::arg("trans_b") = false);
|
||||
}
|
||||
65
ex_engine/csrc/cuinfer_handle.h
Normal file
65
ex_engine/csrc/cuinfer_handle.h
Normal file
@@ -0,0 +1,65 @@
|
||||
// cuinfer_handle.h — Singleton handle manager for libcuinfer.so
|
||||
//
|
||||
// cuinferCreate/Destroy is expensive. This provides a thread-safe
|
||||
// singleton that creates once and reuses.
|
||||
//
|
||||
// Usage:
|
||||
// #include "cuinfer_handle.h"
|
||||
// cuinferHandle_t h = CuinferHandle::get(stream);
|
||||
//
|
||||
// Reference: ixformer::Context::default_cuinfer_handle (in libixformer.so)
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <mutex>
|
||||
#include <cstdio>
|
||||
|
||||
// Forward-declare cuinfer C API
|
||||
extern "C" {
|
||||
|
||||
typedef struct cuinferContext* cuinferHandle_t;
|
||||
|
||||
typedef enum {
|
||||
CUINFER_STATUS_SUCCESS_H = 0,
|
||||
} cuinferStatus_h_t;
|
||||
|
||||
int cuinferCreate(cuinferHandle_t* handle);
|
||||
int cuinferDestroy(cuinferHandle_t handle);
|
||||
int cuinferSetStream(cuinferHandle_t handle, cudaStream_t stream);
|
||||
|
||||
} // extern "C"
|
||||
|
||||
|
||||
class CuinferHandle {
|
||||
public:
|
||||
static cuinferHandle_t get(cudaStream_t stream = nullptr) {
|
||||
static CuinferHandle instance;
|
||||
if (stream && stream != instance.last_stream_) {
|
||||
cuinferSetStream(instance.handle_, stream);
|
||||
instance.last_stream_ = stream;
|
||||
}
|
||||
return instance.handle_;
|
||||
}
|
||||
|
||||
private:
|
||||
cuinferHandle_t handle_ = nullptr;
|
||||
cudaStream_t last_stream_ = nullptr;
|
||||
|
||||
CuinferHandle() {
|
||||
int status = cuinferCreate(&handle_);
|
||||
if (status != 0) {
|
||||
fprintf(stderr, "[cuinfer_handle] WARNING: cuinferCreate failed (%d)\n", status);
|
||||
handle_ = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
~CuinferHandle() {
|
||||
if (handle_) {
|
||||
cuinferDestroy(handle_);
|
||||
}
|
||||
}
|
||||
|
||||
CuinferHandle(const CuinferHandle&) = delete;
|
||||
CuinferHandle& operator=(const CuinferHandle&) = delete;
|
||||
};
|
||||
175
ex_engine/csrc/cuinfer_types.h
Normal file
175
ex_engine/csrc/cuinfer_types.h
Normal file
@@ -0,0 +1,175 @@
|
||||
// cuinfer_types.h — C API types from libcuinfer.so
|
||||
//
|
||||
// Extracted from: cat_files/ixinfer.h (165952 bytes, from real device)
|
||||
// Only the types/enums needed by our GEMM and MoE code.
|
||||
//
|
||||
// This header replaces the scattered extern "C" blocks across
|
||||
// moe_ops_impl.cu, cuinfer_gemm_wrapper.cu, gemm_grouped.cu.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
// --- Handle ---
|
||||
struct cuinferContext;
|
||||
typedef struct cuinferContext* cuinferHandle_t;
|
||||
|
||||
// --- Status ---
|
||||
typedef enum {
|
||||
CUINFER_STATUS_SUCCESS = 0,
|
||||
CUINFER_STATUS_NOT_INITIALIZED = 1,
|
||||
CUINFER_STATUS_ALLOC_FAILED = 2,
|
||||
CUINFER_STATUS_BAD_PARAM = 3,
|
||||
CUINFER_STATUS_INTERNAL_ERROR = 4,
|
||||
CUINFER_STATUS_INVALID_VALUE = 5,
|
||||
CUINFER_STATUS_ARCH_MISMATCH = 6,
|
||||
CUINFER_STATUS_EXECUTION_FAILED = 8,
|
||||
CUINFER_STATUS_NOT_SUPPORTED = 9,
|
||||
} cuinferStatus_t;
|
||||
|
||||
// --- Data types ---
|
||||
typedef enum {
|
||||
CUINFER_DATA_FLOAT = 0,
|
||||
CUINFER_DATA_DOUBLE = 1,
|
||||
CUINFER_DATA_HALF = 2,
|
||||
CUINFER_DATA_INT8 = 3,
|
||||
CUINFER_DATA_INT32 = 4,
|
||||
CUINFER_DATA_INT8x4 = 5,
|
||||
CUINFER_DATA_UINT8 = 6,
|
||||
CUINFER_DATA_UINT8x4 = 7,
|
||||
CUINFER_DATA_INT16 = 8,
|
||||
CUINFER_DATA_BFLOAT16 = 9,
|
||||
} cuinferDataType_t;
|
||||
|
||||
// --- Operations ---
|
||||
typedef enum {
|
||||
CUINFER_OP_N = 0, // no transpose
|
||||
CUINFER_OP_T = 1, // transpose
|
||||
CUINFER_OP_C = 2, // conjugate transpose
|
||||
} cuinferOperation_t;
|
||||
|
||||
// --- Pointer mode ---
|
||||
typedef enum {
|
||||
CUINFER_POINTER_MODE_HOST = 0,
|
||||
CUINFER_POINTER_MODE_DEVICE = 1,
|
||||
} cuinferPointerMode_t;
|
||||
|
||||
// --- GEMM custom option ---
|
||||
typedef enum {
|
||||
CUINFER_GEMM_DEFAULT = 0,
|
||||
} cuinferGEMMCustomOption_t;
|
||||
|
||||
// --- Reduce ops ---
|
||||
typedef enum {
|
||||
CUINFER_REDUCE_TENSOR_ADD = 0,
|
||||
CUINFER_REDUCE_TENSOR_MUL = 1,
|
||||
CUINFER_REDUCE_TENSOR_MIN = 2,
|
||||
CUINFER_REDUCE_TENSOR_MAX = 3,
|
||||
} cuinferReduceTensorOp_t;
|
||||
|
||||
// --- Softmax ---
|
||||
typedef enum {
|
||||
CUINFER_SOFTMAX_FAST = 0,
|
||||
CUINFER_SOFTMAX_ACCURATE = 1,
|
||||
CUINFER_SOFTMAX_LOG = 2,
|
||||
} cuinferSoftmaxAlgorithm_t;
|
||||
|
||||
typedef enum {
|
||||
CUINFER_SOFTMAX_MODE_INSTANCE = 0,
|
||||
CUINFER_SOFTMAX_MODE_CHANNEL = 1,
|
||||
} cuinferSoftmaxMode_t;
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// Function declarations (confirmed in libcuinfer.so symbol dump)
|
||||
// ============================================================================
|
||||
|
||||
cuinferStatus_t cuinferCreate(cuinferHandle_t* handle);
|
||||
cuinferStatus_t cuinferDestroy(cuinferHandle_t handle);
|
||||
cuinferStatus_t cuinferSetStream(cuinferHandle_t handle, cudaStream_t stream);
|
||||
cuinferStatus_t cuinferGetStream(cuinferHandle_t handle, cudaStream_t* stream);
|
||||
size_t cuinferGetVersion(void);
|
||||
const char* cuinferGetErrorString(cuinferStatus_t status);
|
||||
|
||||
// GEMM
|
||||
cuinferStatus_t cuinferCustomGemm(
|
||||
cuinferHandle_t handle, cudaStream_t stream,
|
||||
cuinferPointerMode_t ptrMode,
|
||||
cuinferOperation_t transa, cuinferOperation_t transb,
|
||||
int m, int n, int k,
|
||||
const void* alpha,
|
||||
const void* A, cudaDataType_t Atype, int lda, long long int strideA,
|
||||
const void* B, cudaDataType_t Btype, int ldb, long long int strideB,
|
||||
const void* beta,
|
||||
void* C, cudaDataType_t Ctype, int ldc, long long int strideC,
|
||||
int batchCount,
|
||||
cudaDataType_t computeType, cudaDataType_t scaleType,
|
||||
const void* customHostPtr, const void* customDevicePtr,
|
||||
cuinferGEMMCustomOption_t customOption);
|
||||
|
||||
cuinferStatus_t cuinferCustomGemmEx(
|
||||
cuinferHandle_t handle, cudaStream_t stream,
|
||||
cuinferPointerMode_t ptrMode,
|
||||
cuinferOperation_t transa, cuinferOperation_t transb,
|
||||
int m, int n, int k,
|
||||
const void* alpha,
|
||||
const void* A, cudaDataType_t Atype, int lda, long long int strideA,
|
||||
const void* B, cudaDataType_t Btype, int ldb, long long int strideB,
|
||||
const void* beta,
|
||||
void* C, cudaDataType_t Ctype, int ldc, long long int strideC,
|
||||
int batchCount,
|
||||
cudaDataType_t computeType, cudaDataType_t scaleType,
|
||||
const void* customHostPtr, const void* customDevicePtr,
|
||||
cuinferGEMMCustomOption_t customOption,
|
||||
const void* workspace);
|
||||
|
||||
// TopK
|
||||
cuinferStatus_t cuinferTopK(
|
||||
cuinferHandle_t handle,
|
||||
const void* input, int n, int m, int top_k,
|
||||
int sort_dim, bool largest, bool sorted,
|
||||
void* out_value, int* out_indice,
|
||||
cuinferDataType_t datatype, void* workspace);
|
||||
|
||||
cuinferStatus_t cuinferGetTopKWorkspace(
|
||||
cuinferHandle_t handle,
|
||||
int n, int m, int top_k,
|
||||
cuinferDataType_t datatype, size_t* workspace_size);
|
||||
|
||||
cuinferStatus_t cuinferTopKBatch(
|
||||
cuinferHandle_t handle,
|
||||
const void* input, int top_k, int batch, int n, int m, int k,
|
||||
bool largest, bool sorted, int sort_dim,
|
||||
void* output, int* indice,
|
||||
cuinferDataType_t datatype, void* workspace);
|
||||
|
||||
// Softmax
|
||||
cuinferStatus_t cuinferSoftmaxForward(
|
||||
cuinferHandle_t handle,
|
||||
cuinferSoftmaxAlgorithm_t algo,
|
||||
cuinferSoftmaxMode_t mode,
|
||||
const void* alpha,
|
||||
const void* xDesc, const void* x,
|
||||
const void* beta,
|
||||
const void* yDesc, void* y);
|
||||
|
||||
// Reduce
|
||||
cuinferStatus_t cuinferReduce(
|
||||
cuinferHandle_t handle,
|
||||
const void* in, void* out,
|
||||
cuinferDataType_t in_type,
|
||||
cuinferDataType_t acc_type,
|
||||
cuinferDataType_t out_type,
|
||||
cuinferReduceTensorOp_t reduce_op,
|
||||
int n_dims, const int* dims,
|
||||
int n_reduce_dims, const int* reduce_dim_index,
|
||||
void* workspace);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // extern "C"
|
||||
#endif
|
||||
188
ex_engine/csrc/gemm_grouped.cu
Normal file
188
ex_engine/csrc/gemm_grouped.cu
Normal file
@@ -0,0 +1,188 @@
|
||||
// gemm_grouped.cu — Per-expert GEMM using CUTLASS Cu10 TensorOp
|
||||
//
|
||||
// Source lineage:
|
||||
// cat_files/batched_gemm.cu — cutlass sample from real device
|
||||
// cat_files/default_gemm_configuration.h — Cu10 half/half/float config
|
||||
// ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu — existing impl
|
||||
// ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp — moe_expert_gemm pattern
|
||||
//
|
||||
// This file provides:
|
||||
// 1. cutlass_expert_gemm() — one cutlass GEMM per expert (Cu10 TensorOp)
|
||||
// 2. cuinfer_expert_gemm() — one cuinferCustomGemm per expert (fallback)
|
||||
// 3. moe_group_gemm() — unified entry: try cutlass, fall back to cuinfer
|
||||
//
|
||||
// All use RowMajor, FP16 data, FP32 accumulation.
|
||||
// Weight layout: [num_experts, N, K] (TN format = transB in GEMM sense)
|
||||
|
||||
#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"
|
||||
|
||||
// ============================================================================
|
||||
// Cu10 TensorOp GEMM type — from default_gemm_configuration.h
|
||||
// ThreadblockShape<128,128,32>, WarpShape<32,32,32>, Instruction<16,16,16>
|
||||
// ============================================================================
|
||||
using GemmCu10 = 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, // use TCU
|
||||
cutlass::arch::Cu10 // BI-V100
|
||||
>;
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// cutlass_expert_gemm: per-expert GEMM using CUTLASS
|
||||
//
|
||||
// For each expert e with M_e tokens:
|
||||
// C[offset:offset+M_e, :N] = A[offset:offset+M_e, :K] @ B[e, :N, :K]^T
|
||||
//
|
||||
// B is stored as [num_experts, N, K] (RowMajor), we need A×B^T.
|
||||
// Cutlass RowMajor × RowMajor computes C = A × B, so we transpose:
|
||||
// C(M,N) = A(M,K) × B^T(K,N) = A(M,K) × B_orig(N,K)^T
|
||||
//
|
||||
// In row-major: A lda=K, B lda=K (it's NxK stored row-major), C ldc=N
|
||||
// We use Cutlass's NN mode on (A, B^T) which is implemented as:
|
||||
// Cutlass RowMajor NN: C[i,j] = sum_k A[i,k] * B[k,j]
|
||||
// But B is (N,K) not (K,N), so we pass B as ColumnMajor or handle via stride.
|
||||
//
|
||||
// Simpler: A is (M,K) RowMajor, we want output (M,N).
|
||||
// B_expert is (N,K) RowMajor = same as (K,N) ColumnMajor.
|
||||
// So: A(M,K) RowMajor × B(K,N) ColumnMajor → C(M,N) RowMajor
|
||||
// This is exactly GEMM with transB.
|
||||
// ============================================================================
|
||||
|
||||
using GemmCu10_TN = cutlass::gemm::device::GemmBatched<
|
||||
cutlass::half_t, // ElementA
|
||||
cutlass::layout::RowMajor, // LayoutA — A is (M,K) row-major
|
||||
cutlass::half_t, // ElementB
|
||||
cutlass::layout::ColumnMajor, // LayoutB — B is (N,K) stored row = (K,N) col
|
||||
cutlass::half_t, // ElementC
|
||||
cutlass::layout::RowMajor, // LayoutC
|
||||
float, // ElementAccumulator
|
||||
cutlass::arch::OpClassTensorOp, // TCU
|
||||
cutlass::arch::Cu10 // BI-V100
|
||||
>;
|
||||
|
||||
|
||||
int cutlass_expert_gemm(
|
||||
int num_experts,
|
||||
const int* expert_counts, // host array [num_experts]
|
||||
const int* expert_offsets, // host array [num_experts], exclusive prefix sum
|
||||
int N, int K,
|
||||
const __half* input, // (total_tokens, K) row-major
|
||||
const __half* weights, // (num_experts, N, K) row-major — TN format
|
||||
__half* output, // (total_tokens, N) row-major
|
||||
cudaStream_t stream)
|
||||
{
|
||||
GemmCu10_TN gemm_op;
|
||||
float alpha = 1.0f, beta = 0.0f;
|
||||
int failures = 0;
|
||||
|
||||
for (int e = 0; e < num_experts; e++) {
|
||||
int M_e = expert_counts[e];
|
||||
if (M_e <= 0) continue;
|
||||
|
||||
int off = expert_offsets[e];
|
||||
auto A = reinterpret_cast<cutlass::half_t const*>(input + (long long)off * K);
|
||||
auto B = reinterpret_cast<cutlass::half_t const*>(weights + (long long)e * N * K);
|
||||
auto C = reinterpret_cast<cutlass::half_t*>(output + (long long)off * N);
|
||||
|
||||
// A: (M_e, K) RowMajor, lda = K
|
||||
// B: (N, K) RowMajor → (K, N) ColumnMajor, ldb = N (col-major stride)
|
||||
// C: (M_e, N) RowMajor, ldc = N
|
||||
cutlass::Status status = gemm_op({
|
||||
{M_e, N, K},
|
||||
{A, K}, // A, lda
|
||||
0, // strideA (not batched)
|
||||
{B, K}, // B in col-major view: (N,K) row = (K,N) col, ldb = K
|
||||
0, // strideB
|
||||
{C, N}, // C, ldc
|
||||
0, // strideC
|
||||
{C, N}, // D = C
|
||||
0,
|
||||
{alpha, beta},
|
||||
1 // batch_count = 1 (we loop over experts)
|
||||
});
|
||||
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
failures++;
|
||||
}
|
||||
}
|
||||
return failures;
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// cuinfer fallback — forward-declare cuinferCustomGemm
|
||||
// ============================================================================
|
||||
extern "C" {
|
||||
typedef struct cuinferContext* cuinferHandle_t;
|
||||
typedef enum { CUINFER_STATUS_SUCCESS_GG = 0 } cuinferStatus_gg_t;
|
||||
cuinferHandle_t cuinferCreate_handle();
|
||||
|
||||
int cuinferCustomGemm(
|
||||
cuinferHandle_t handle, cudaStream_t stream,
|
||||
int ptrMode, int transa, int transb,
|
||||
int m, int n, int k,
|
||||
const void* alpha,
|
||||
const void* A, int Atype, int lda, long long int strideA,
|
||||
const void* B, int Btype, int ldb, long long int strideB,
|
||||
const void* beta,
|
||||
void* C, int Ctype, int ldc, long long int strideC,
|
||||
int batchCount, int computeType, int scaleType,
|
||||
const void* customHostPtr, const void* customDevicePtr, int customOption);
|
||||
}
|
||||
|
||||
|
||||
int cuinfer_expert_gemm(
|
||||
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,
|
||||
cuinferHandle_t handle)
|
||||
{
|
||||
float alpha = 1.0f, beta = 0.0f;
|
||||
int failures = 0;
|
||||
|
||||
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 void* A = input + (long long)off * K;
|
||||
const void* B = weights + (long long)e * N * K;
|
||||
void* C = output + (long long)off * N;
|
||||
|
||||
// cuinferCustomGemm: transa=0 (N), transb=1 (T)
|
||||
// CUDA_R_16F = 2
|
||||
int status = cuinferCustomGemm(
|
||||
handle, stream,
|
||||
0, // CUINFER_POINTER_MODE_HOST
|
||||
0, 1, // transa=N, transb=T
|
||||
M_e, N, K,
|
||||
&alpha,
|
||||
A, 2, K, 0, // A: fp16, lda=K
|
||||
B, 2, K, 0, // B: fp16, ldb=K (row-major N×K, transposed)
|
||||
&beta,
|
||||
C, 2, N, 0, // C: fp16, ldc=N
|
||||
1, // batchCount=1
|
||||
0, 0, // computeType=fp32, scaleType=fp32
|
||||
nullptr, nullptr, 0);
|
||||
|
||||
if (status != 0) failures++;
|
||||
}
|
||||
return failures;
|
||||
}
|
||||
182
ex_engine/csrc/gemm_grouped_bind.cpp
Normal file
182
ex_engine/csrc/gemm_grouped_bind.cpp
Normal file
@@ -0,0 +1,182 @@
|
||||
// gemm_grouped_bind.cpp — Python bindings for grouped GEMM
|
||||
//
|
||||
// Source lineage:
|
||||
// ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp — moe_expert_gemm pattern
|
||||
// ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp — batched pattern
|
||||
//
|
||||
// Exports:
|
||||
// moe_group_gemm(input, weights, expert_counts) → output
|
||||
// moe_group_gemm_cutlass(input, weights, expert_counts) → output
|
||||
// moe_decode_cutlass(hidden, w13, w2, topk_weights) → output
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <vector>
|
||||
|
||||
// From gemm_grouped.cu
|
||||
int cutlass_expert_gemm(
|
||||
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);
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// moe_group_gemm: per-expert GEMM using CUTLASS Cu10 TensorOp
|
||||
//
|
||||
// input: (total_tokens, K) fp16
|
||||
// weights: (num_experts, N, K) fp16, TN layout
|
||||
// expert_counts: (num_experts,) int32
|
||||
// Returns: (total_tokens, N) fp16
|
||||
// ============================================================================
|
||||
torch::Tensor moe_group_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");
|
||||
|
||||
int total_tokens = input.size(0);
|
||||
int K = input.size(1);
|
||||
int num_experts = weights.size(0);
|
||||
int N = weights.size(1);
|
||||
TORCH_CHECK(weights.size(2) == K, "weights K dim must match input K");
|
||||
|
||||
auto output = torch::zeros({total_tokens, N}, input.options());
|
||||
|
||||
// Build host arrays
|
||||
auto counts_cpu = expert_counts.to(torch::kCPU).to(torch::kInt32).contiguous();
|
||||
int32_t* c = counts_cpu.data_ptr<int32_t>();
|
||||
std::vector<int> counts(num_experts), offsets(num_experts);
|
||||
int cumsum = 0;
|
||||
for (int i = 0; i < num_experts; i++) {
|
||||
counts[i] = c[i];
|
||||
offsets[i] = cumsum;
|
||||
cumsum += c[i];
|
||||
}
|
||||
|
||||
cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
|
||||
|
||||
int fails = cutlass_expert_gemm(
|
||||
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);
|
||||
|
||||
if (fails > 0) {
|
||||
// Fallback to PyTorch F.linear per expert
|
||||
auto input_a = input.to(torch::kFloat32);
|
||||
auto output_f = torch::zeros({total_tokens, N},
|
||||
input.options().dtype(torch::kFloat32));
|
||||
for (int e = 0; e < num_experts; e++) {
|
||||
if (counts[e] <= 0) continue;
|
||||
int off = offsets[e];
|
||||
auto x = input_a.narrow(0, off, counts[e]);
|
||||
auto w = weights[e].to(torch::kFloat32); // (N, K)
|
||||
output_f.narrow(0, off, counts[e]) = torch::mm(x, w.t());
|
||||
}
|
||||
output = output_f.to(torch::kHalf);
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// moe_decode_cutlass: fused MoE decode for single-token (batch=1)
|
||||
//
|
||||
// Uses CUTLASS batched GEMM for the topk experts simultaneously.
|
||||
//
|
||||
// hidden: (1, H) fp16
|
||||
// w13_sel: (topk, 2*I, H) fp16 — already-gathered expert weights
|
||||
// w2_sel: (topk, H, I) fp16
|
||||
// topk_weights: (topk,) float32
|
||||
// Returns: (1, H) fp16
|
||||
// ============================================================================
|
||||
|
||||
// From corex_batched_gemm_kernel.cu
|
||||
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);
|
||||
|
||||
|
||||
torch::Tensor moe_decode_cutlass(
|
||||
torch::Tensor hidden, // (1, H)
|
||||
torch::Tensor w13_sel, // (topk, 2*I, H)
|
||||
torch::Tensor w2_sel, // (topk, H, I)
|
||||
torch::Tensor topk_weights) // (topk,)
|
||||
{
|
||||
int topk = 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 (topk, 1, H)
|
||||
auto x = hidden.expand({topk, 1, H}).contiguous();
|
||||
|
||||
// w13^T: (topk, 2I, H) → transpose → (topk, H, 2I)
|
||||
auto w13_t = w13_sel.transpose(1, 2).contiguous();
|
||||
|
||||
// Step 1: gate_up = x @ w13^T → (topk, 1, 2I)
|
||||
auto gate_up_3d = torch::empty({topk, 1, two_I}, x.options());
|
||||
auto status1 = cutlass_batched_hgemm(
|
||||
1, two_I, H,
|
||||
reinterpret_cast<const __half*>(x.data_ptr<at::Half>()),
|
||||
H, H,
|
||||
reinterpret_cast<const __half*>(w13_t.data_ptr<at::Half>()),
|
||||
two_I, H * two_I,
|
||||
reinterpret_cast<__half*>(gate_up_3d.data_ptr<at::Half>()),
|
||||
two_I, two_I,
|
||||
topk);
|
||||
TORCH_CHECK(status1 == cudaSuccess, "batched GEMM 1 failed");
|
||||
|
||||
auto gate_up = gate_up_3d.squeeze(1); // (topk, 2I)
|
||||
|
||||
// Step 2: SiLU activation
|
||||
auto chunks = gate_up.chunk(2, 1);
|
||||
auto act = torch::silu(chunks[0]) * chunks[1]; // (topk, I)
|
||||
act = act.unsqueeze(1).contiguous(); // (topk, 1, I)
|
||||
|
||||
// w2^T: (topk, H, I) → transpose → (topk, I, H)
|
||||
auto w2_t = w2_sel.transpose(1, 2).contiguous();
|
||||
|
||||
// Step 3: down = act @ w2^T → (topk, 1, H)
|
||||
auto down_3d = torch::empty({topk, 1, H}, x.options());
|
||||
auto status2 = cutlass_batched_hgemm(
|
||||
1, H, I,
|
||||
reinterpret_cast<const __half*>(act.data_ptr<at::Half>()),
|
||||
I, I,
|
||||
reinterpret_cast<const __half*>(w2_t.data_ptr<at::Half>()),
|
||||
H, I * H,
|
||||
reinterpret_cast<__half*>(down_3d.data_ptr<at::Half>()),
|
||||
H, H,
|
||||
topk);
|
||||
TORCH_CHECK(status2 == cudaSuccess, "batched GEMM 2 failed");
|
||||
|
||||
auto down = down_3d.squeeze(1); // (topk, H)
|
||||
|
||||
// Step 4: weighted sum
|
||||
auto out = (down * topk_weights.unsqueeze(1).to(down.dtype())).sum(0, true);
|
||||
return out;
|
||||
}
|
||||
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("moe_group_gemm", &moe_group_gemm,
|
||||
"Per-expert GEMM via CUTLASS Cu10 TensorOp",
|
||||
py::arg("input"), py::arg("weights"), py::arg("expert_counts"));
|
||||
m.def("moe_decode_cutlass", &moe_decode_cutlass,
|
||||
"Fused MoE decode via CUTLASS batched GEMM",
|
||||
py::arg("hidden"), py::arg("w13_sel"),
|
||||
py::arg("w2_sel"), py::arg("topk_weights"));
|
||||
}
|
||||
180
ex_engine/python/gemm_dispatch.py
Normal file
180
ex_engine/python/gemm_dispatch.py
Normal file
@@ -0,0 +1,180 @@
|
||||
"""gemm_dispatch.py — Unified GEMM dispatch for MoE group matmul.
|
||||
|
||||
AST Layer 2: selects best available GEMM backend on real device.
|
||||
|
||||
Backend priority:
|
||||
1. gemm_grouped.so (cutlass Cu10 TensorOp, per-expert GEMM)
|
||||
2. ix_moe_bridge.so (cuinferCustomGemm, per-expert loop)
|
||||
3. corex_batched_gemm.so (cutlass batched, decode-only)
|
||||
4. hgemm.so (blocktiling kernel from siboehm)
|
||||
5. torch.mm loop (PyTorch fallback)
|
||||
|
||||
Reference: ex_engine/python/ix_ops_dispatch.py (407L)
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
logger = logging.getLogger("gemm_dispatch")
|
||||
|
||||
# --- Backend loading ---
|
||||
_cutlass_grouped = None
|
||||
_moe_bridge = None
|
||||
_batched_gemm = None
|
||||
_hgemm = None
|
||||
_backend = "torch"
|
||||
|
||||
|
||||
def _try_load(name):
|
||||
"""Try to load a .so module by name."""
|
||||
# Search paths
|
||||
search = [
|
||||
os.path.join(os.path.dirname(__file__), f"{name}.so"),
|
||||
os.path.join(os.path.dirname(__file__), "..", "prebuilt", f"{name}.so"),
|
||||
os.path.join(os.path.dirname(__file__), "..", f"{name}.so"),
|
||||
]
|
||||
for p in search:
|
||||
if os.path.isfile(p):
|
||||
try:
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location(name, p)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
except Exception as e:
|
||||
logger.debug(f"[gemm] Failed to load {p}: {e}")
|
||||
# Try direct import
|
||||
try:
|
||||
import importlib
|
||||
return importlib.import_module(name)
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def _init_backends():
|
||||
global _cutlass_grouped, _moe_bridge, _batched_gemm, _hgemm, _backend
|
||||
|
||||
_cutlass_grouped = _try_load("gemm_grouped")
|
||||
if _cutlass_grouped and hasattr(_cutlass_grouped, "moe_group_gemm"):
|
||||
_backend = "cutlass_grouped"
|
||||
logger.info("[gemm] Backend: cutlass_grouped (Cu10 TensorOp)")
|
||||
return
|
||||
|
||||
_moe_bridge = _try_load("ix_moe_bridge")
|
||||
if _moe_bridge and hasattr(_moe_bridge, "group_gemm"):
|
||||
_backend = "cuinfer"
|
||||
logger.info("[gemm] Backend: cuinfer (via ix_moe_bridge)")
|
||||
return
|
||||
|
||||
_batched_gemm = _try_load("corex_batched_gemm")
|
||||
if _batched_gemm and hasattr(_batched_gemm, "batched_gemm_fp16"):
|
||||
_backend = "cutlass_batched"
|
||||
logger.info("[gemm] Backend: cutlass_batched")
|
||||
return
|
||||
|
||||
_hgemm = _try_load("hgemm")
|
||||
if _hgemm and hasattr(_hgemm, "moe_expert_gemm"):
|
||||
_backend = "hgemm"
|
||||
logger.info("[gemm] Backend: hgemm (blocktiling)")
|
||||
return
|
||||
|
||||
_backend = "torch"
|
||||
logger.info("[gemm] Backend: torch (F.linear fallback)")
|
||||
|
||||
|
||||
_init_backends()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Public API
|
||||
# ============================================================================
|
||||
|
||||
def group_gemm(input_tokens, weights, expert_counts, output_dim):
|
||||
"""Per-expert GEMM: output[offset:offset+count] = input[offset:offset+count] @ W[e]^T
|
||||
|
||||
Args:
|
||||
input_tokens: (total_tokens, K) fp16
|
||||
weights: (num_experts, N, K) fp16, TN layout
|
||||
expert_counts: (num_experts,) int32
|
||||
output_dim: N (output dimension)
|
||||
|
||||
Returns:
|
||||
(total_tokens, N) fp16
|
||||
"""
|
||||
if _backend == "cutlass_grouped":
|
||||
return _cutlass_grouped.moe_group_gemm(input_tokens, weights, expert_counts)
|
||||
|
||||
if _backend == "cuinfer":
|
||||
return _moe_bridge.group_gemm(input_tokens, weights, expert_counts, output_dim)
|
||||
|
||||
if _backend == "hgemm":
|
||||
return _hgemm.moe_expert_gemm(input_tokens, weights, expert_counts)
|
||||
|
||||
# torch fallback
|
||||
return _torch_group_gemm(input_tokens, weights, expert_counts)
|
||||
|
||||
|
||||
def moe_decode_gemm(hidden, w13_sel, w2_sel, topk_weights):
|
||||
"""Single-token MoE decode: batched GEMM over topk experts.
|
||||
|
||||
Args:
|
||||
hidden: (1, H) fp16
|
||||
w13_sel: (topk, 2*I, H) fp16
|
||||
w2_sel: (topk, H, I) fp16
|
||||
topk_weights: (topk,) float32
|
||||
|
||||
Returns:
|
||||
(1, H) fp16
|
||||
"""
|
||||
if _backend == "cutlass_grouped" and hasattr(_cutlass_grouped, "moe_decode_cutlass"):
|
||||
return _cutlass_grouped.moe_decode_cutlass(hidden, w13_sel, w2_sel, topk_weights)
|
||||
|
||||
if _backend == "cutlass_batched" and _batched_gemm is not None:
|
||||
return _batched_gemm.moe_decode_fused(hidden, w13_sel, w2_sel, topk_weights)
|
||||
|
||||
# torch fallback
|
||||
return _torch_moe_decode(hidden, w13_sel, w2_sel, topk_weights)
|
||||
|
||||
|
||||
def get_backend():
|
||||
return _backend
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Fallbacks
|
||||
# ============================================================================
|
||||
|
||||
def _torch_group_gemm(input_tokens, weights, expert_counts):
|
||||
"""PyTorch fallback: per-expert F.linear loop."""
|
||||
num_experts = weights.size(0)
|
||||
N = weights.size(1)
|
||||
output = torch.zeros(input_tokens.size(0), N,
|
||||
device=input_tokens.device, dtype=input_tokens.dtype)
|
||||
|
||||
counts_cpu = expert_counts.cpu().to(torch.int32)
|
||||
offset = 0
|
||||
for e in range(num_experts):
|
||||
cnt = counts_cpu[e].item()
|
||||
if cnt <= 0:
|
||||
offset += cnt
|
||||
continue
|
||||
x = input_tokens[offset:offset+cnt]
|
||||
w = weights[e] # (N, K)
|
||||
output[offset:offset+cnt] = F.linear(x, w)
|
||||
offset += cnt
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _torch_moe_decode(hidden, w13_sel, w2_sel, topk_weights):
|
||||
"""PyTorch fallback for single-token MoE decode."""
|
||||
topk = w13_sel.size(0)
|
||||
results = []
|
||||
for k in range(topk):
|
||||
gate_up = F.linear(hidden, w13_sel[k])
|
||||
inter = gate_up.shape[-1] // 2
|
||||
act = torch.silu(gate_up[:, :inter]) * gate_up[:, inter:]
|
||||
down = F.linear(act, w2_sel[k])
|
||||
results.append(down * topk_weights[k].to(down.dtype))
|
||||
return sum(results)
|
||||
165
test_gemm.py
Normal file
165
test_gemm.py
Normal file
@@ -0,0 +1,165 @@
|
||||
"""test_gemm.py — Correctness test for all GEMM backends.
|
||||
|
||||
Compares each backend against torch.mm with Qwen3.5 MoE shapes.
|
||||
Reports max absolute error and whether it passes FP16 tolerance.
|
||||
|
||||
Usage: python3 test_gemm.py
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import torch
|
||||
|
||||
H = 3584
|
||||
I = 18944 // 4 # 4736, per TP=4
|
||||
TWO_I = I * 2
|
||||
NUM_EXPERTS = 128
|
||||
TOPK = 8
|
||||
FP16_ATOL = 5e-2 # FP16 has ~1e-3 precision, allow some accumulation error
|
||||
|
||||
|
||||
def test_single_gemm(device):
|
||||
"""Test single GEMM: A(M,K) × B(K,N)."""
|
||||
print("\n=== Single GEMM ===")
|
||||
A = torch.randn(4, H, device=device, dtype=torch.float16)
|
||||
B = torch.randn(H, TWO_I, device=device, dtype=torch.float16) * 0.01
|
||||
|
||||
ref = torch.mm(A.float(), B.float()).half()
|
||||
|
||||
backends = {}
|
||||
try:
|
||||
import cuinfer_gemm_wrapper
|
||||
backends["cuinfer"] = cuinfer_gemm_wrapper.cuinfer_gemm(A, B, False)
|
||||
except Exception as e:
|
||||
print(f" cuinfer: skip ({e})")
|
||||
|
||||
try:
|
||||
import hgemm
|
||||
backends["hgemm"] = hgemm.hgemm(A, B)
|
||||
except Exception as e:
|
||||
print(f" hgemm: skip ({e})")
|
||||
|
||||
for name, out in backends.items():
|
||||
err = (out.float() - ref.float()).abs().max().item()
|
||||
ok = "✓" if err < FP16_ATOL else "✗"
|
||||
print(f" {ok} {name}: max_err={err:.6f} (tol={FP16_ATOL})")
|
||||
|
||||
|
||||
def test_group_gemm(device):
|
||||
"""Test group GEMM with per-expert variable counts."""
|
||||
print("\n=== Group GEMM (MoE w13) ===")
|
||||
total_tokens = TOPK # decode: 1 token × 8 experts
|
||||
expert_counts = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32)
|
||||
for i in range(TOPK):
|
||||
expert_counts[i] = 1
|
||||
|
||||
input_t = torch.randn(total_tokens, H, device=device, dtype=torch.float16) * 0.1
|
||||
w13 = torch.randn(NUM_EXPERTS, TWO_I, H, device=device, dtype=torch.float16) * 0.01
|
||||
|
||||
# Reference: torch.mm per expert
|
||||
ref = torch.zeros(total_tokens, TWO_I, device=device, dtype=torch.float16)
|
||||
offset = 0
|
||||
for e in range(NUM_EXPERTS):
|
||||
c = expert_counts[e].item()
|
||||
if c <= 0: continue
|
||||
ref[offset:offset+c] = torch.mm(
|
||||
input_t[offset:offset+c].float(), w13[e].t().float()
|
||||
).half()
|
||||
offset += c
|
||||
|
||||
backends = {}
|
||||
try:
|
||||
import gemm_grouped
|
||||
backends["cutlass_grouped"] = gemm_grouped.moe_group_gemm(input_t, w13, expert_counts)
|
||||
except Exception as e:
|
||||
print(f" cutlass_grouped: skip ({e})")
|
||||
|
||||
try:
|
||||
import ix_moe_bridge
|
||||
backends["cuinfer_bridge"] = ix_moe_bridge.group_gemm(input_t, w13, expert_counts, TWO_I)
|
||||
except Exception as e:
|
||||
print(f" cuinfer_bridge: skip ({e})")
|
||||
|
||||
try:
|
||||
import hgemm
|
||||
backends["hgemm"] = hgemm.moe_expert_gemm(input_t, w13, expert_counts)
|
||||
except Exception as e:
|
||||
print(f" hgemm: skip ({e})")
|
||||
|
||||
for name, out in backends.items():
|
||||
err = (out.float() - ref.float()).abs().max().item()
|
||||
ok = "✓" if err < FP16_ATOL else "✗"
|
||||
print(f" {ok} {name}: max_err={err:.6f}")
|
||||
|
||||
|
||||
def test_batched_gemm(device):
|
||||
"""Test batched GEMM for decode path."""
|
||||
print("\n=== Batched GEMM (decode, topk=8) ===")
|
||||
A = torch.randn(TOPK, 1, H, device=device, dtype=torch.float16)
|
||||
B = torch.randn(TOPK, H, TWO_I, device=device, dtype=torch.float16) * 0.01
|
||||
|
||||
# Reference
|
||||
ref = torch.bmm(A.float(), B.float()).half()
|
||||
|
||||
backends = {}
|
||||
try:
|
||||
import cuinfer_gemm_wrapper
|
||||
backends["cuinfer_batched"] = cuinfer_gemm_wrapper.cuinfer_gemm_batched(A, B, False)
|
||||
except Exception as e:
|
||||
print(f" cuinfer_batched: skip ({e})")
|
||||
|
||||
try:
|
||||
import corex_batched_gemm
|
||||
backends["corex_batched"] = corex_batched_gemm.batched_gemm_fp16(A, B)
|
||||
except Exception as e:
|
||||
print(f" corex_batched: skip ({e})")
|
||||
|
||||
for name, out in backends.items():
|
||||
err = (out.float() - ref.float()).abs().max().item()
|
||||
ok = "✓" if err < FP16_ATOL else "✗"
|
||||
print(f" {ok} {name}: max_err={err:.6f}")
|
||||
|
||||
|
||||
def test_gemm_dispatch(device):
|
||||
"""Test the unified gemm_dispatch layer."""
|
||||
print("\n=== gemm_dispatch ===")
|
||||
try:
|
||||
from gemm_dispatch import group_gemm, get_backend
|
||||
print(f" Backend: {get_backend()}")
|
||||
|
||||
total_tokens = TOPK
|
||||
expert_counts = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32)
|
||||
for i in range(TOPK):
|
||||
expert_counts[i] = 1
|
||||
input_t = torch.randn(total_tokens, H, device=device, dtype=torch.float16) * 0.1
|
||||
w13 = torch.randn(NUM_EXPERTS, TWO_I, H, device=device, dtype=torch.float16) * 0.01
|
||||
|
||||
out = group_gemm(input_t, w13, expert_counts, TWO_I)
|
||||
assert out.shape == (total_tokens, TWO_I), f"shape: {out.shape}"
|
||||
assert not torch.isnan(out).any(), "NaN in output"
|
||||
print(f" ✓ shape={out.shape}, no NaN")
|
||||
except Exception as e:
|
||||
print(f" ✗ {e}")
|
||||
|
||||
|
||||
def main():
|
||||
if not torch.cuda.is_available():
|
||||
print("No CUDA")
|
||||
sys.exit(0)
|
||||
|
||||
device = torch.device("cuda:0")
|
||||
print(f"Device: {torch.cuda.get_device_name(0)}")
|
||||
|
||||
passed, failed = 0, 0
|
||||
for test in [test_single_gemm, test_group_gemm, test_batched_gemm, test_gemm_dispatch]:
|
||||
try:
|
||||
test(device)
|
||||
passed += 1
|
||||
except Exception as e:
|
||||
print(f" ✗ FAILED: {e}")
|
||||
failed += 1
|
||||
|
||||
print(f"\nResults: {passed} passed, {failed} failed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user