support qwen 3.8 flash next (#37500)
Co-authored-by: ch-wan <54331508+ch-wan@users.noreply.github.com> Co-authored-by: ispobock <26454835+ispobock@users.noreply.github.com> Co-authored-by: JustinTong0323 <115166877+JustinTong0323@users.noreply.github.com> Co-authored-by: samuellees <26428561+samuellees@users.noreply.github.com> Co-authored-by: YAMY1234 <74099316+YAMY1234@users.noreply.github.com> Co-authored-by: yhyang201 <47235274+yhyang201@users.noreply.github.com> Co-authored-by: yizhang2077 <25844240+yizhang2077@users.noreply.github.com> Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com> Co-authored-by: Shinto C V <cshintov@gmail.com> Co-authored-by: Julian Huang <huangzhilin.hzl@antgroup.com> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: yhyang201 <yhyang201@gmail.com>
This commit is contained in:
Qiaolin Yu
co-authored by
ch-wan
ispobock
JustinTong0323
samuellees
YAMY1234
yhyang201
yizhang2077
zijiexia
Shinto C V
Julian Huang
Xiaoyu Zhang
yhyang201
ch-wan
ispobock
JustinTong0323
samuellees
YAMY1234
yhyang201
yizhang2077
zijiexia
Shinto C V
Julian Huang
Xiaoyu Zhang
yhyang201
parent
afe90a8bc9
commit
52fecfdf09
@@ -0,0 +1,486 @@
|
||||
// Fused QSA (Qwen4-Exp sparse attention) indexer-prep kernels.
|
||||
// Outputs are bit-identical to the eager bf16/fp16 aten chain, mirrored step by step.
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/math.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
/// \brief Round a float to the storage dtype and back (one eager aten step).
|
||||
/// The inline cvt keeps nvcc from folding the round-trip away: every eager
|
||||
/// bf16/fp16 aten op rounds its result, and this kernel must round likewise.
|
||||
template <typename T>
|
||||
SGL_DEVICE float eager_round(float x) {
|
||||
return static_cast<float>(DTypeTrait<T>::from(x));
|
||||
}
|
||||
|
||||
template <>
|
||||
SGL_DEVICE float eager_round<bf16_t>(float x) {
|
||||
#ifndef USE_ROCM
|
||||
uint16_t u;
|
||||
asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(u) : "f"(x));
|
||||
return __bfloat162float(__ushort_as_bfloat16(u));
|
||||
#else
|
||||
return static_cast<float>(DTypeTrait<bf16_t>::from(x));
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
SGL_DEVICE float eager_round<fp16_t>(float x) {
|
||||
#ifndef USE_ROCM
|
||||
uint16_t u;
|
||||
asm("cvt.rn.f16.f32 %0, %1;" : "=h"(u) : "f"(x));
|
||||
return __ushort_as_half(u);
|
||||
#else
|
||||
return static_cast<float>(DTypeTrait<fp16_t>::from(x));
|
||||
#endif
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Apply (M)RoPE to a normed row staged in shared memory.
|
||||
*
|
||||
* Reproduces the eager `get_cos_sin_with_position` + `apply_rotary_emb` chain:
|
||||
* `axis_map[i]` selects which position axis feeds pair index i (plain RoPE
|
||||
* uses all-zero maps; Qwen interleaved/sectioned MRoPE maps are built on the
|
||||
* host). The cos/sin cache row is [cos(half), sin(half)] of width rotary_dim.
|
||||
*
|
||||
* \tparam T Storage element type: bf16_t | fp16_t.
|
||||
* \tparam kHeadDim Compile-time head dimension (multiple of 32).
|
||||
* \tparam kIsNeox true -> NeoX pairing (d, d+half); false -> GPT-J (2i, 2i+1).
|
||||
* \param smem_row Normed row [kHeadDim], one warp cooperates.
|
||||
* \param out_row Destination row [kHeadDim].
|
||||
* \param cos_sin_cache [num_positions, rotary_dim] fp32 cache.
|
||||
* \param axis_map [rotary_dim/2] int32 position-axis selector per pair.
|
||||
* \param pos Resolved per-axis positions for this token (>= 3 entries).
|
||||
* \param rotary_dim Rotated prefix length; tail dims pass through.
|
||||
*/
|
||||
template <typename T, int kHeadDim, bool kIsNeox>
|
||||
SGL_DEVICE void qsa_mrope_apply(
|
||||
const T* __restrict__ smem_row,
|
||||
T* __restrict__ out_row,
|
||||
const float* __restrict__ cos_sin_cache,
|
||||
const int32_t* __restrict__ axis_map,
|
||||
const int64_t* pos,
|
||||
const int32_t rotary_dim) {
|
||||
using namespace device;
|
||||
constexpr int kPerLane = kHeadDim / kWarpThreads;
|
||||
using vec_t = AlignedVector<T, kPerLane>;
|
||||
const uint32_t lane = threadIdx.x % kWarpThreads;
|
||||
const int32_t half = rotary_dim / 2;
|
||||
|
||||
vec_t ov;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kPerLane; ++i) {
|
||||
const int32_t d = static_cast<int32_t>(lane * kPerLane) + i;
|
||||
T o;
|
||||
if constexpr (kIsNeox) {
|
||||
if (d < half) {
|
||||
const int32_t p = d + half;
|
||||
const float* row = cos_sin_cache + pos[axis_map[d]] * rotary_dim;
|
||||
const float c = eager_round<T>(row[d]);
|
||||
const float s = eager_round<T>(row[half + d]);
|
||||
const float nd = static_cast<float>(smem_row[d]);
|
||||
const float np = static_cast<float>(smem_row[p]);
|
||||
o = DTypeTrait<T>::from(eager_round<T>(nd * c) - eager_round<T>(np * s));
|
||||
} else if (d < rotary_dim) {
|
||||
const int32_t p = d - half;
|
||||
const float* row = cos_sin_cache + pos[axis_map[p]] * rotary_dim;
|
||||
const float c = eager_round<T>(row[p]);
|
||||
const float s = eager_round<T>(row[half + p]);
|
||||
const float nd = static_cast<float>(smem_row[d]);
|
||||
const float np = static_cast<float>(smem_row[p]);
|
||||
o = DTypeTrait<T>::from(eager_round<T>(nd * c) + eager_round<T>(np * s));
|
||||
} else {
|
||||
o = smem_row[d];
|
||||
}
|
||||
} else {
|
||||
if (d < rotary_dim) {
|
||||
const int32_t p = d / 2;
|
||||
const float* row = cos_sin_cache + pos[axis_map[p]] * rotary_dim;
|
||||
const float c = eager_round<T>(row[p]);
|
||||
const float s = eager_round<T>(row[half + p]);
|
||||
const int32_t q = (d % 2 == 0) ? d + 1 : d - 1;
|
||||
const float nd = static_cast<float>(smem_row[d]);
|
||||
const float nq = static_cast<float>(smem_row[q]);
|
||||
const float t1 = eager_round<T>(nd * c);
|
||||
const float t2 = eager_round<T>(nq * s);
|
||||
o = DTypeTrait<T>::from((d % 2 == 0) ? t1 - t2 : t1 + t2);
|
||||
} else {
|
||||
o = smem_row[d];
|
||||
}
|
||||
}
|
||||
ov[i] = o;
|
||||
}
|
||||
ov.store(out_row, lane); // offset is in vector units
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Gemma RMSNorm of one row into shared memory (warp-cooperative).
|
||||
*
|
||||
* out = x * rsqrt(mean(x^2) + eps) * (1 + w), fp32 math with one rounding.
|
||||
* The sum-of-squares uses flashinfer RMSNormKernel's exact partial layout
|
||||
* (vec_size = 8 for 2-byte dtypes, thread t sums elements [8t, 8t+8), inactive
|
||||
* lanes contribute 0, xor butterfly over the warp) so results stay bit-equal
|
||||
* to the eager sgl_kernel gemma_rmsnorm this replaces.
|
||||
*/
|
||||
template <typename T, int kHeadDim>
|
||||
SGL_DEVICE void qsa_gemma_norm_row(
|
||||
const T* __restrict__ x_row, const T* __restrict__ weight, const float eps, T* __restrict__ smem_row) {
|
||||
using namespace device;
|
||||
constexpr int kPerLane = kHeadDim / kWarpThreads;
|
||||
using vec_t = AlignedVector<T, kPerLane>;
|
||||
const uint32_t lane = threadIdx.x % kWarpThreads;
|
||||
|
||||
vec_t xv, wv;
|
||||
xv.load(x_row, lane); // offset is in vector units
|
||||
wv.load(weight, lane);
|
||||
|
||||
float xf[kPerLane];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kPerLane; ++i) {
|
||||
xf[i] = static_cast<float>(xv[i]);
|
||||
}
|
||||
|
||||
static_assert(kHeadDim % 8 == 0);
|
||||
constexpr uint32_t kNormThreads = kHeadDim / 8;
|
||||
float ss = 0.0f;
|
||||
if (lane < kNormThreads) {
|
||||
AlignedVector<T, 4> va, vb;
|
||||
va.load(x_row, lane * 2); // elements [8*lane, 8*lane+4)
|
||||
vb.load(x_row, lane * 2 + 1); // elements [8*lane+4, 8*lane+8)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const float f = static_cast<float>(va[i]);
|
||||
ss += f * f;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const float f = static_cast<float>(vb[i]);
|
||||
ss += f * f;
|
||||
}
|
||||
}
|
||||
ss = warp::reduce_sum(ss);
|
||||
const float nf = math::rsqrt(ss / kHeadDim + eps);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kPerLane; ++i) {
|
||||
const float wf = static_cast<float>(wv[i]);
|
||||
smem_row[lane * kPerLane + i] = DTypeTrait<T>::from(xf[i] * nf * (1.0f + wf));
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
struct QsaIndexQPrepParams {
|
||||
const void* qk; // [tokens, (num_q_heads + 1) * kHeadDim]
|
||||
void* q_out; // [tokens, q_heads_padded, kHeadDim]
|
||||
const void* weight; // [kHeadDim]
|
||||
const float* cos_sin_cache; // [positions_capacity, rotary_dim]
|
||||
const int32_t* axis_map; // [rotary_dim / 2]
|
||||
const int64_t* positions; // [num_axes, tokens] (row stride may exceed tokens)
|
||||
const int64_t* cache_loc; // [tokens]
|
||||
void* key_state_buffer; // [slots, kHeadDim]
|
||||
int64_t* rope_position_buffer; // [slots, 3]
|
||||
int64_t positions_stride;
|
||||
int32_t num_axes;
|
||||
int32_t num_q_heads;
|
||||
int32_t q_heads_padded;
|
||||
int32_t rotary_dim;
|
||||
float eps;
|
||||
};
|
||||
|
||||
/**
|
||||
* \brief Per-token fused index-Q prep: gemma norm + MRoPE for every query
|
||||
* head, zero-fill of padded heads, raw token-K store and RoPE-position store.
|
||||
* One CTA (4 warps) per token; one warp per query head.
|
||||
*/
|
||||
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
|
||||
__global__ __launch_bounds__(128) void qsa_index_q_prep_kernel(const QsaIndexQPrepParams __grid_constant__ params) {
|
||||
using namespace device;
|
||||
constexpr int kPerLane = kHeadDim / kWarpThreads;
|
||||
using vec_t = AlignedVector<T, kPerLane>;
|
||||
const uint32_t token = blockIdx.x;
|
||||
const uint32_t warp = threadIdx.x / kWarpThreads;
|
||||
const uint32_t lane = threadIdx.x % kWarpThreads;
|
||||
__shared__ T smem_rows[4][kHeadDim];
|
||||
|
||||
device::PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
const int64_t qk_row = static_cast<int64_t>(token) * (params.num_q_heads + 1) * kHeadDim;
|
||||
const int64_t loc = params.cache_loc[token];
|
||||
int64_t pos[3];
|
||||
#pragma unroll
|
||||
for (int a = 0; a < 3; ++a) {
|
||||
const int64_t ax = a < params.num_axes ? a : 0;
|
||||
pos[a] = params.positions[ax * params.positions_stride + token];
|
||||
}
|
||||
|
||||
for (int32_t h = static_cast<int32_t>(warp); h < params.q_heads_padded; h += 4) {
|
||||
T* out_row = static_cast<T*>(params.q_out) + (static_cast<int64_t>(token) * params.q_heads_padded + h) * kHeadDim;
|
||||
if (h < params.num_q_heads) {
|
||||
const T* x_row = static_cast<const T*>(params.qk) + qk_row + h * kHeadDim;
|
||||
qsa_gemma_norm_row<T, kHeadDim>(x_row, static_cast<const T*>(params.weight), params.eps, smem_rows[warp]);
|
||||
qsa_mrope_apply<T, kHeadDim, kIsNeox>(
|
||||
smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim);
|
||||
} else {
|
||||
vec_t zv;
|
||||
zv.fill(DTypeTrait<T>::from(0.0f));
|
||||
zv.store(out_row, lane); // offset is in vector units
|
||||
}
|
||||
}
|
||||
|
||||
// Raw token-K and RoPE coordinates are stored for every token, whether or
|
||||
// not the token completes a compression group.
|
||||
if (warp == 0) {
|
||||
vec_t kv;
|
||||
kv.load(
|
||||
static_cast<const T*>(params.qk) + qk_row + params.num_q_heads * kHeadDim,
|
||||
lane); // offset is in vector units
|
||||
kv.store(static_cast<T*>(params.key_state_buffer) + loc * kHeadDim, lane);
|
||||
}
|
||||
if (warp == 1 && lane < 3) {
|
||||
params.rope_position_buffer[loc * 3 + lane] = pos[lane];
|
||||
}
|
||||
|
||||
device::PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
struct QsaIndexKCompressParams {
|
||||
const void* key_state_buffer; // [slots, kHeadDim]
|
||||
const int32_t* group_locs; // [groups, compress_ratio]
|
||||
const int64_t* rope_position_buffer; // [slots, 3]
|
||||
const float* cos_sin_cache; // [positions_capacity, rotary_dim]
|
||||
const int32_t* axis_map; // [rotary_dim / 2]
|
||||
const void* weight; // [kHeadDim]
|
||||
const int32_t* write_locs; // [groups]
|
||||
void* compressed_k_buffer; // [compressed_slots, kHeadDim]
|
||||
int32_t compress_ratio;
|
||||
int32_t rotary_dim;
|
||||
int32_t num_groups;
|
||||
float eps;
|
||||
};
|
||||
|
||||
/**
|
||||
* \brief Per-group compressed-K prep: fp32 mean over the group, gemma norm,
|
||||
* MRoPE at the group-start position, store into the compressed cache.
|
||||
* One warp per group.
|
||||
*/
|
||||
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
|
||||
__global__
|
||||
__launch_bounds__(128) void qsa_index_k_compress_kernel(const QsaIndexKCompressParams __grid_constant__ params) {
|
||||
using namespace device;
|
||||
constexpr int kPerLane = kHeadDim / kWarpThreads;
|
||||
using vec_t = AlignedVector<T, kPerLane>;
|
||||
const uint32_t warp = threadIdx.x / kWarpThreads;
|
||||
const uint32_t lane = threadIdx.x % kWarpThreads;
|
||||
const uint32_t group = blockIdx.x * 4 + warp;
|
||||
if (group >= static_cast<uint32_t>(params.num_groups)) {
|
||||
return;
|
||||
}
|
||||
__shared__ T smem_rows[4][kHeadDim];
|
||||
|
||||
device::PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
const int32_t* locs = params.group_locs + group * params.compress_ratio;
|
||||
const int32_t loc0 = locs[0];
|
||||
|
||||
// fp32 mean over the group, rounded to the storage dtype exactly like
|
||||
// average_pool_qsa_keys (float().mean(dim=1).to(dtype)).
|
||||
float mf[kPerLane];
|
||||
{
|
||||
float acc[kPerLane];
|
||||
for (int32_t r = 0; r < params.compress_ratio; ++r) {
|
||||
vec_t v;
|
||||
v.load(
|
||||
static_cast<const T*>(params.key_state_buffer) + static_cast<int64_t>(locs[r]) * kHeadDim,
|
||||
lane); // offset is in vector units
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kPerLane; ++i) {
|
||||
const float f = static_cast<float>(v[i]);
|
||||
acc[i] = (r == 0) ? f : acc[i] + f;
|
||||
}
|
||||
}
|
||||
const float inv_ratio = 1.0f / static_cast<float>(params.compress_ratio);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kPerLane; ++i) {
|
||||
const T m = DTypeTrait<T>::from(acc[i] * inv_ratio);
|
||||
mf[i] = static_cast<float>(m);
|
||||
smem_rows[warp][lane * kPerLane + i] = m;
|
||||
}
|
||||
__syncwarp();
|
||||
// Sum of squares in flashinfer RMSNormKernel's exact partial layout
|
||||
// (thread t sums elements [8t, 8t+8), inactive lanes contribute 0), so
|
||||
// the result stays bit-equal to the eager k_layernorm this replaces.
|
||||
static_assert(kHeadDim % 8 == 0);
|
||||
float ss = 0.0f;
|
||||
if (lane < kHeadDim / 8) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
const float f = static_cast<float>(smem_rows[warp][lane * 8 + j]);
|
||||
ss += f * f;
|
||||
}
|
||||
}
|
||||
ss = warp::reduce_sum(ss);
|
||||
const float nf = math::rsqrt(ss / kHeadDim + params.eps);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kPerLane; ++i) {
|
||||
const float wf = static_cast<float>(static_cast<const T*>(params.weight)[lane * kPerLane + i]);
|
||||
smem_rows[warp][lane * kPerLane + i] = DTypeTrait<T>::from(mf[i] * nf * (1.0f + wf));
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
int64_t pos[3];
|
||||
#pragma unroll
|
||||
for (int a = 0; a < 3; ++a) {
|
||||
pos[a] = params.rope_position_buffer[static_cast<int64_t>(loc0) * 3 + a];
|
||||
}
|
||||
|
||||
T* out_row = static_cast<T*>(params.compressed_k_buffer) + static_cast<int64_t>(params.write_locs[group]) * kHeadDim;
|
||||
qsa_mrope_apply<T, kHeadDim, kIsNeox>(
|
||||
smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim);
|
||||
|
||||
device::PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Validate inputs and launch `qsa_index_q_prep_kernel` (one CTA per token).
|
||||
*
|
||||
* \tparam T Element type: bf16_t | fp16_t.
|
||||
* \tparam kHeadDim Index head dimension: 64 | 128 | 256.
|
||||
* \tparam kIsNeox RoPE pairing style.
|
||||
* \tparam kUsePDL Whether to launch with PDL enabled.
|
||||
*/
|
||||
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
|
||||
void qsa_index_q_prep(
|
||||
tvm::ffi::TensorView qk,
|
||||
tvm::ffi::TensorView q_out,
|
||||
tvm::ffi::TensorView weight,
|
||||
tvm::ffi::TensorView cos_sin_cache,
|
||||
tvm::ffi::TensorView axis_map,
|
||||
tvm::ffi::TensorView positions,
|
||||
int64_t num_axes,
|
||||
tvm::ffi::TensorView cache_loc,
|
||||
tvm::ffi::TensorView key_state_buffer,
|
||||
tvm::ffi::TensorView rope_position_buffer,
|
||||
int64_t num_q_heads,
|
||||
int64_t rotary_dim,
|
||||
float eps) {
|
||||
using namespace host;
|
||||
auto tokens = SymbolicSize{"tokens"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
constexpr int64_t D = kHeadDim;
|
||||
|
||||
TensorMatcher({tokens, (num_q_heads + 1) * D}).with_dtype<T>().with_device(device).verify(qk);
|
||||
auto heads_padded = SymbolicSize{"heads_padded"};
|
||||
TensorMatcher({tokens, heads_padded, D}).with_dtype<T>().with_device(device).verify(q_out);
|
||||
TensorMatcher({D}).with_dtype<T>().with_device(device).verify(weight);
|
||||
auto cache_rows = SymbolicSize{"cos_sin_cache_rows"};
|
||||
TensorMatcher({cache_rows, rotary_dim}).with_dtype<fp32_t>().with_device(device).verify(cos_sin_cache);
|
||||
TensorMatcher({rotary_dim / 2}).with_dtype<int32_t>().with_device(device).verify(axis_map);
|
||||
TensorMatcher({num_axes, tokens}).with_dtype<int64_t>().with_device(device).with_strides({-1, 1}).verify(positions);
|
||||
TensorMatcher({tokens}).with_dtype<int64_t>().with_device(device).verify(cache_loc);
|
||||
auto slots = SymbolicSize{"state_slots"};
|
||||
TensorMatcher({slots, D}).with_dtype<T>().with_device(device).verify(key_state_buffer);
|
||||
TensorMatcher({slots, 3}).with_dtype<int64_t>().with_device(device).verify(rope_position_buffer);
|
||||
|
||||
const int64_t num_tokens = tokens.unwrap();
|
||||
const int64_t q_heads_padded = heads_padded.unwrap();
|
||||
CHECK_HOST(num_tokens > 0) << "qsa_index_q_prep: no tokens";
|
||||
CHECK_HOST(num_axes == 1 || num_axes == 3) << "qsa_index_q_prep: positions must have 1 or 3 axes, got " << num_axes;
|
||||
CHECK_HOST(q_heads_padded >= num_q_heads)
|
||||
<< "qsa_index_q_prep: padded heads " << q_heads_padded << " < num_q_heads " << num_q_heads;
|
||||
CHECK_HOST(rotary_dim > 0 && rotary_dim % 2 == 0 && rotary_dim <= D)
|
||||
<< "qsa_index_q_prep: invalid rotary_dim " << rotary_dim;
|
||||
|
||||
const auto params = QsaIndexQPrepParams{
|
||||
.qk = qk.data_ptr(),
|
||||
.q_out = q_out.data_ptr(),
|
||||
.weight = weight.data_ptr(),
|
||||
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
|
||||
.axis_map = static_cast<const int32_t*>(axis_map.data_ptr()),
|
||||
.positions = static_cast<const int64_t*>(positions.data_ptr()),
|
||||
.cache_loc = static_cast<const int64_t*>(cache_loc.data_ptr()),
|
||||
.key_state_buffer = key_state_buffer.data_ptr(),
|
||||
.rope_position_buffer = static_cast<int64_t*>(rope_position_buffer.data_ptr()),
|
||||
.positions_stride = positions.stride(0),
|
||||
.num_axes = static_cast<int32_t>(num_axes),
|
||||
.num_q_heads = static_cast<int32_t>(num_q_heads),
|
||||
.q_heads_padded = static_cast<int32_t>(q_heads_padded),
|
||||
.rotary_dim = static_cast<int32_t>(rotary_dim),
|
||||
.eps = eps,
|
||||
};
|
||||
LaunchKernel(static_cast<uint32_t>(num_tokens), 128, device.unwrap())
|
||||
.enable_pdl(kUsePDL)(qsa_index_q_prep_kernel<T, kHeadDim, kIsNeox, kUsePDL>, params);
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Validate inputs and launch `qsa_index_k_compress_kernel` (one warp per group).
|
||||
*/
|
||||
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
|
||||
void qsa_index_k_compress(
|
||||
tvm::ffi::TensorView key_state_buffer,
|
||||
tvm::ffi::TensorView group_locs,
|
||||
tvm::ffi::TensorView rope_position_buffer,
|
||||
tvm::ffi::TensorView cos_sin_cache,
|
||||
tvm::ffi::TensorView axis_map,
|
||||
tvm::ffi::TensorView weight,
|
||||
tvm::ffi::TensorView write_locs,
|
||||
tvm::ffi::TensorView compressed_k_buffer,
|
||||
int64_t compress_ratio,
|
||||
int64_t rotary_dim,
|
||||
float eps) {
|
||||
using namespace host;
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
constexpr int64_t D = kHeadDim;
|
||||
|
||||
auto slots = SymbolicSize{"state_slots"};
|
||||
TensorMatcher({slots, D}).with_dtype<T>().with_device(device).verify(key_state_buffer);
|
||||
auto groups = SymbolicSize{"groups"};
|
||||
TensorMatcher({groups, compress_ratio}).with_dtype<int32_t>().with_device(device).verify(group_locs);
|
||||
TensorMatcher({slots, 3}).with_dtype<int64_t>().with_device(device).verify(rope_position_buffer);
|
||||
auto cache_rows = SymbolicSize{"cos_sin_cache_rows"};
|
||||
TensorMatcher({cache_rows, rotary_dim}).with_dtype<fp32_t>().with_device(device).verify(cos_sin_cache);
|
||||
TensorMatcher({rotary_dim / 2}).with_dtype<int32_t>().with_device(device).verify(axis_map);
|
||||
TensorMatcher({D}).with_dtype<T>().with_device(device).verify(weight);
|
||||
TensorMatcher({groups}).with_dtype<int32_t>().with_device(device).verify(write_locs);
|
||||
auto compressed_slots = SymbolicSize{"compressed_slots"};
|
||||
TensorMatcher({compressed_slots, D}).with_dtype<T>().with_device(device).verify(compressed_k_buffer);
|
||||
|
||||
const int64_t num_groups = groups.unwrap();
|
||||
CHECK_HOST(num_groups > 0) << "qsa_index_k_compress: no groups";
|
||||
CHECK_HOST(compress_ratio > 0 && compress_ratio <= 16)
|
||||
<< "qsa_index_k_compress: invalid compress_ratio " << compress_ratio;
|
||||
CHECK_HOST(rotary_dim > 0 && rotary_dim % 2 == 0 && rotary_dim <= D)
|
||||
<< "qsa_index_k_compress: invalid rotary_dim " << rotary_dim;
|
||||
|
||||
const auto params = QsaIndexKCompressParams{
|
||||
.key_state_buffer = key_state_buffer.data_ptr(),
|
||||
.group_locs = static_cast<const int32_t*>(group_locs.data_ptr()),
|
||||
.rope_position_buffer = static_cast<const int64_t*>(rope_position_buffer.data_ptr()),
|
||||
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
|
||||
.axis_map = static_cast<const int32_t*>(axis_map.data_ptr()),
|
||||
.weight = weight.data_ptr(),
|
||||
.write_locs = static_cast<const int32_t*>(write_locs.data_ptr()),
|
||||
.compressed_k_buffer = compressed_k_buffer.data_ptr(),
|
||||
.compress_ratio = static_cast<int32_t>(compress_ratio),
|
||||
.rotary_dim = static_cast<int32_t>(rotary_dim),
|
||||
.num_groups = static_cast<int32_t>(num_groups),
|
||||
.eps = eps,
|
||||
};
|
||||
LaunchKernel(static_cast<uint32_t>(div_ceil(num_groups, 4)), 128, device.unwrap())
|
||||
.enable_pdl(kUsePDL)(qsa_index_k_compress_kernel<T, kHeadDim, kIsNeox, kUsePDL>, params);
|
||||
}
|
||||
|
||||
} // namespace sglang
|
||||
@@ -0,0 +1,291 @@
|
||||
// Radix-select fast top-k with the AOT fast_topk_v2 semantics: for each row b,
|
||||
// select the kTopK largest scores in [row_starts[b], row_starts[b] + lengths[b])
|
||||
// and write their indices relative to row_starts[b]; order within a row is unspecified.
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
namespace fast_topk_detail {
|
||||
|
||||
constexpr uint32_t kThreadsPerBlock = 1024;
|
||||
// Each radix pass needs at most ~kTopK candidates in the threshold bin, so
|
||||
// 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
|
||||
constexpr size_t kSmemBytes = 8 * 1024 * sizeof(uint32_t); // 32KB
|
||||
|
||||
struct FastTopKParams {
|
||||
const float* __restrict__ input; // [B, input_stride]
|
||||
const int32_t* __restrict__ row_starts; // [B]
|
||||
int32_t* __restrict__ indices; // [B, kTopK]
|
||||
const int32_t* __restrict__ lengths; // [B]
|
||||
int64_t input_stride;
|
||||
};
|
||||
|
||||
SGL_DEVICE auto convert_to_uint8(float x) -> uint8_t {
|
||||
const __half h = __float2half_rn(x);
|
||||
const uint16_t bits = __half_as_ushort(h);
|
||||
const uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
|
||||
return static_cast<uint8_t>(key >> 8);
|
||||
}
|
||||
|
||||
SGL_DEVICE auto convert_to_uint32(float x) -> uint32_t {
|
||||
const uint32_t bits = __float_as_uint(x);
|
||||
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
|
||||
}
|
||||
|
||||
// When length <= kTopK, write the indices directly.
|
||||
template <int kTopK>
|
||||
SGL_DEVICE void naive_topk(const float* __restrict__ score, int32_t* __restrict__ indice, int32_t length) {
|
||||
const auto tid = threadIdx.x;
|
||||
for (int i = tid; i < kTopK; i += kThreadsPerBlock) {
|
||||
indice[i] = (i < length) ? i : -1;
|
||||
}
|
||||
}
|
||||
|
||||
// Radix-select top-k. Assumes length > kTopK (checked by the caller).
|
||||
template <int kTopK>
|
||||
SGL_DEVICE void radix_select_topk(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
|
||||
int topk = kTopK;
|
||||
constexpr auto BLOCK_SIZE = kThreadsPerBlock;
|
||||
constexpr auto RADIX = 256;
|
||||
constexpr auto SMEM_INPUT_SIZE = kSmemBytes / (2 * sizeof(int));
|
||||
|
||||
alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
|
||||
alignas(128) __shared__ int s_counter;
|
||||
alignas(128) __shared__ int s_threshold_bin_id;
|
||||
alignas(128) __shared__ int s_num_input[2];
|
||||
|
||||
auto& s_histogram = s_histogram_buf[0];
|
||||
// allocate for two rounds
|
||||
extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
|
||||
|
||||
const int tx = threadIdx.x;
|
||||
|
||||
// stage 1: 8bit coarse histogram
|
||||
if (tx < RADIX + 1) s_histogram[tx] = 0;
|
||||
__syncthreads();
|
||||
|
||||
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const auto bin = convert_to_uint8(input[idx + row_start]);
|
||||
::atomicAdd(&s_histogram[bin], 1);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto run_cumsum = [&] {
|
||||
#pragma unroll 8
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
static_assert(1 << 8 == RADIX);
|
||||
if (tx < RADIX) {
|
||||
const auto j = 1 << i;
|
||||
const auto k = i & 1;
|
||||
auto value = s_histogram_buf[k][tx];
|
||||
if (tx < RADIX - j) {
|
||||
value += s_histogram_buf[k][tx + j];
|
||||
}
|
||||
s_histogram_buf[k ^ 1][tx] = value;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
};
|
||||
|
||||
run_cumsum();
|
||||
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
|
||||
s_threshold_bin_id = tx;
|
||||
s_num_input[0] = 0;
|
||||
s_counter = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto threshold_bin = s_threshold_bin_id;
|
||||
topk -= s_histogram[threshold_bin + 1];
|
||||
|
||||
if (topk == 0) {
|
||||
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return;
|
||||
} else {
|
||||
__syncthreads();
|
||||
if (tx < RADIX + 1) {
|
||||
s_histogram[tx] = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const auto raw_input = input[idx + row_start];
|
||||
const auto bin = static_cast<int>(convert_to_uint8(raw_input));
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
} else if (bin == threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_num_input[0], 1);
|
||||
// fuse the histogram computation here
|
||||
if (pos < int(SMEM_INPUT_SIZE)) {
|
||||
s_input_idx[0][pos] = idx;
|
||||
const auto bin = convert_to_uint32(raw_input);
|
||||
const auto sub_bin = (bin >> 24) & 0xFF;
|
||||
::atomicAdd(&s_histogram[sub_bin], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// stage 2: refine with 8bit radix passes
|
||||
#pragma unroll 4
|
||||
for (int round = 0; round < 4; ++round) {
|
||||
__shared__ int s_last_remain;
|
||||
const auto r_idx = round % 2;
|
||||
|
||||
// clip here to prevent overflow
|
||||
const auto _raw_num_input = s_num_input[r_idx];
|
||||
const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
|
||||
|
||||
run_cumsum();
|
||||
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
|
||||
s_threshold_bin_id = tx;
|
||||
s_num_input[r_idx ^ 1] = 0;
|
||||
s_last_remain = topk - s_histogram[tx + 1];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto threshold_bin = s_threshold_bin_id;
|
||||
topk -= s_histogram[threshold_bin + 1];
|
||||
|
||||
if (topk == 0) {
|
||||
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
|
||||
const auto idx = s_input_idx[r_idx][i];
|
||||
const auto offset = 24 - round * 8;
|
||||
const auto bin = (convert_to_uint32(input[idx + row_start]) >> offset) & 0xFF;
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
break;
|
||||
} else {
|
||||
__syncthreads();
|
||||
if (tx < RADIX + 1) {
|
||||
s_histogram[tx] = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
|
||||
const auto idx = s_input_idx[r_idx][i];
|
||||
const auto raw_input = input[idx + row_start];
|
||||
const auto offset = 24 - round * 8;
|
||||
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
} else if (bin == threshold_bin) {
|
||||
if (round == 3) {
|
||||
const auto pos = ::atomicAdd(&s_last_remain, -1);
|
||||
if (pos > 0) {
|
||||
index[kTopK - pos] = idx;
|
||||
}
|
||||
} else {
|
||||
const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
|
||||
if (pos < int(SMEM_INPUT_SIZE)) {
|
||||
// fuse the histogram computation here
|
||||
s_input_idx[r_idx ^ 1][pos] = idx;
|
||||
const auto bin = convert_to_uint32(raw_input);
|
||||
const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
|
||||
::atomicAdd(&s_histogram[sub_bin], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int kTopK, bool kUsePDL>
|
||||
__global__ __launch_bounds__(fast_topk_detail::kThreadsPerBlock) void fast_topk_kernel(
|
||||
const fast_topk_detail::FastTopKParams __grid_constant__ params) {
|
||||
using namespace fast_topk_detail;
|
||||
device::PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
const auto bid = static_cast<uint64_t>(blockIdx.x);
|
||||
const auto row_start = params.row_starts == nullptr ? 0 : params.row_starts[bid];
|
||||
const auto length = params.lengths[bid];
|
||||
const auto indice = params.indices + bid * kTopK;
|
||||
const auto score = params.input + bid * params.input_stride;
|
||||
if (length <= kTopK) {
|
||||
naive_topk<kTopK>(score, indice, length);
|
||||
} else {
|
||||
radix_select_topk<kTopK>(score, indice, row_start, length);
|
||||
}
|
||||
|
||||
device::PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
} // namespace fast_topk_detail
|
||||
|
||||
/**
|
||||
* \brief Per-row top-k selection over ragged rows of a fp32 score matrix.
|
||||
*
|
||||
* Row b selects the kTopK largest values in
|
||||
* score[b, row_starts[b] : row_starts[b] + lengths[b]) and writes their
|
||||
* indices (relative to row_starts[b]) into indices[b]. Unfilled slots are
|
||||
* -1 when lengths[b] < kTopK.
|
||||
*/
|
||||
template <int kTopK, bool kUsePDL>
|
||||
struct FastTopKKernel {
|
||||
static constexpr auto kernel = fast_topk_detail::fast_topk_kernel<kTopK, kUsePDL>;
|
||||
|
||||
static void
|
||||
run(const tvm::ffi::TensorView score,
|
||||
const tvm::ffi::TensorView row_starts,
|
||||
const tvm::ffi::TensorView indices,
|
||||
const tvm::ffi::TensorView lengths) {
|
||||
using namespace host;
|
||||
auto B = SymbolicSize{"batch"};
|
||||
auto L = SymbolicSize{"length"};
|
||||
auto S = SymbolicSize{"input_stride"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B, L}) // score
|
||||
.with_strides({S, 1})
|
||||
.with_dtype<fp32_t>()
|
||||
.with_device(device)
|
||||
.verify(score);
|
||||
TensorMatcher({B}) // row_starts
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device)
|
||||
.verify(row_starts);
|
||||
TensorMatcher({B, kTopK}) // indices
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device)
|
||||
.verify(indices);
|
||||
TensorMatcher({B}) // lengths
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device)
|
||||
.verify(lengths);
|
||||
|
||||
const auto params = fast_topk_detail::FastTopKParams{
|
||||
.input = static_cast<const float*>(score.data_ptr()),
|
||||
.row_starts = static_cast<const int32_t*>(row_starts.data_ptr()),
|
||||
.indices = static_cast<int32_t*>(indices.data_ptr()),
|
||||
.lengths = static_cast<const int32_t*>(lengths.data_ptr()),
|
||||
.input_stride = S.unwrap(),
|
||||
};
|
||||
|
||||
const auto num_rows = static_cast<uint32_t>(B.unwrap());
|
||||
LaunchKernel(num_rows, fast_topk_detail::kThreadsPerBlock, device.unwrap(), fast_topk_detail::kSmemBytes)
|
||||
.enable_pdl(kUsePDL)(kernel, params);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang
|
||||
@@ -0,0 +1,178 @@
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/math.cuh>
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
struct GroupedGemmaRMSNormParams {
|
||||
const void* input;
|
||||
const void* __restrict__ weight;
|
||||
void* output;
|
||||
uint32_t num_groups;
|
||||
float eps;
|
||||
};
|
||||
|
||||
/**
|
||||
* \brief Grouped Gemma-style RMSNorm: out = x * rsqrt(mean(x^2) + eps) * (1 + w).
|
||||
*
|
||||
* The last dim of the input is split into `num_groups` chunks of `kGroupSize`
|
||||
* elements. Variance is computed per (token, group) chunk, so a [M, H] input
|
||||
* with H = num_groups * kGroupSize behaves like M * num_groups independent
|
||||
* RMSNorm rows whose weight rows are the matching kGroupSize slice of `weight`.
|
||||
*
|
||||
* One CTA handles one (token, group) chunk. Since chunks are contiguous in
|
||||
* memory, block `bid` reads/writes elements [bid * kGroupSize, (bid + 1) *
|
||||
* kGroupSize) and uses weight slice (bid % num_groups) * kGroupSize.
|
||||
*
|
||||
* \tparam kGroupSize Elements per group. Must be a multiple of 512.
|
||||
* \tparam kUsePDL Whether to emit the PDL wait/trigger pair.
|
||||
* \tparam Float Element type: bf16_t | fp16_t.
|
||||
*/
|
||||
template <int64_t kGroupSize, bool kUsePDL, typename Float>
|
||||
__global__ __launch_bounds__(kGroupSize / 16) void grouped_gemma_rmsnorm_kernel(
|
||||
const GroupedGemmaRMSNormParams __grid_constant__ params) {
|
||||
using namespace device;
|
||||
using Float2 = packed_t<Float>;
|
||||
#if SGL_ARCH_BLACKWELL_OR_GREATER
|
||||
// Blackwell: 32B vector, each thread loads/stores once
|
||||
using Storage = AlignedVector<Float2, 8>;
|
||||
constexpr uint32_t kNumLoads = 1;
|
||||
#else
|
||||
// Pre-Blackwell: 16B vector, each thread loads/stores twice
|
||||
using Storage = AlignedVector<Float2, 4>;
|
||||
constexpr uint32_t kNumLoads = 2;
|
||||
#endif
|
||||
constexpr uint32_t kVecLen = kNumLoads == 1 ? 8 : 4;
|
||||
constexpr auto kNumThreads = kGroupSize / 16;
|
||||
constexpr auto kNumWarps = kNumThreads / kWarpThreads;
|
||||
|
||||
const uint32_t bid = blockIdx.x;
|
||||
const uint32_t group = bid % params.num_groups;
|
||||
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
|
||||
// Warp 0 writes smem[tx] for all 32 lanes in the cross-warp reduce below,
|
||||
// so this must hold kWarpThreads entries, not kNumWarps.
|
||||
__shared__ float smem[kWarpThreads];
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
const auto input_ptr = pointer::offset<Float>(params.input, static_cast<int64_t>(bid) * kGroupSize);
|
||||
const auto output_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(bid) * kGroupSize);
|
||||
const auto weight_ptr = pointer::offset<Float>(params.weight, static_cast<int64_t>(group) * kGroupSize);
|
||||
|
||||
Storage input_vec[kNumLoads];
|
||||
Storage weight_vec[kNumLoads];
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kNumLoads; ++j) {
|
||||
input_vec[j] = gmem.load(input_ptr, j);
|
||||
weight_vec[j] = gmem.load(weight_ptr, j);
|
||||
}
|
||||
|
||||
float sum_of_squares = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kNumLoads; ++j) {
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecLen; ++i) {
|
||||
const auto [x, y] = cast<fp32x2_t>(input_vec[j][i]);
|
||||
sum_of_squares += x * x + y * y;
|
||||
}
|
||||
}
|
||||
|
||||
sum_of_squares = warp::reduce_sum(sum_of_squares);
|
||||
float norm_factor;
|
||||
if constexpr (kNumWarps == 1) {
|
||||
norm_factor = math::rsqrt(sum_of_squares / kGroupSize + params.eps);
|
||||
} else {
|
||||
const auto warp_id = threadIdx.x / kWarpThreads;
|
||||
smem[warp_id] = sum_of_squares;
|
||||
__syncthreads();
|
||||
if (warp_id == 0) {
|
||||
const auto tx = threadIdx.x;
|
||||
const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
|
||||
sum_of_squares = warp::reduce_sum(local_sum);
|
||||
smem[tx] = math::rsqrt(sum_of_squares / kGroupSize + params.eps);
|
||||
}
|
||||
__syncthreads();
|
||||
norm_factor = smem[warp_id];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kNumLoads; ++j) {
|
||||
Storage output_vec;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecLen; ++i) {
|
||||
const auto [ix, iy] = cast<fp32x2_t>(input_vec[j][i]);
|
||||
const auto [wx, wy] = cast<fp32x2_t>(weight_vec[j][i]);
|
||||
output_vec[i] = cast<Float2>(fp32x2_t{ix * norm_factor * (1.0f + wx), iy * norm_factor * (1.0f + wy)});
|
||||
}
|
||||
gmem.store(output_ptr, output_vec, j);
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
template <int64_t kGroupSize, bool kUsePDL, typename DType>
|
||||
struct GroupedGemmaRMSNormKernel {
|
||||
static_assert(sizeof(DType) == 2, "GroupedGemmaRMSNorm only supports 2-byte dtypes");
|
||||
static_assert(kGroupSize % 512 == 0, "kGroupSize must be a multiple of 512");
|
||||
static constexpr auto kernel = grouped_gemma_rmsnorm_kernel<kGroupSize, kUsePDL, DType>;
|
||||
static constexpr auto kBlockSize = static_cast<uint32_t>(kGroupSize / 16);
|
||||
|
||||
/**
|
||||
* \brief Validate tensors and launch one CTA per (token, group) chunk.
|
||||
* \param input [M, H] contiguous, H % kGroupSize == 0
|
||||
* \param weight [H]
|
||||
* \param output [M, H] contiguous, same shape/dtype/device as input
|
||||
* \param eps RMSNorm epsilon
|
||||
*/
|
||||
static void
|
||||
run(const tvm::ffi::TensorView input,
|
||||
const tvm::ffi::TensorView weight,
|
||||
const tvm::ffi::TensorView output,
|
||||
float eps) {
|
||||
using namespace host;
|
||||
auto M = SymbolicSize{"num_tokens"};
|
||||
auto H = SymbolicSize{"hidden_size"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({M, H}) // input
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(input);
|
||||
TensorMatcher({H}) // weight
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(weight);
|
||||
TensorMatcher({M, H}) // output
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(output);
|
||||
|
||||
const int64_t hidden_size = H.unwrap();
|
||||
CHECK_HOST(hidden_size % kGroupSize == 0) << "grouped_gemma_rmsnorm: hidden_size (" << hidden_size
|
||||
<< ") must be divisible by group_size (" << kGroupSize << ")";
|
||||
|
||||
const auto params = GroupedGemmaRMSNormParams{
|
||||
.input = input.data_ptr(),
|
||||
.weight = weight.data_ptr(),
|
||||
.output = output.data_ptr(),
|
||||
.num_groups = static_cast<uint32_t>(hidden_size / kGroupSize),
|
||||
.eps = eps,
|
||||
};
|
||||
|
||||
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
|
||||
const uint32_t num_blocks = num_tokens * params.num_groups;
|
||||
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(kernel, params);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang
|
||||
@@ -0,0 +1,381 @@
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/math.cuh>
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
struct HcCombineParams {
|
||||
const void* block_output; // [M, H]
|
||||
const void* residual; // [M, HC * H]
|
||||
const void* normed_residual; // [M, HC * H]
|
||||
const void* inject_weight; // [HC, HC * H]
|
||||
void* output; // [M, HC * H]
|
||||
};
|
||||
|
||||
/**
|
||||
* \brief Fused HyperConnection (gated residual) combine:
|
||||
*
|
||||
* a[m, c] = 2 * sigmoid(dot(normed_residual[m, :], inject_weight[c, :]) / kHcCount)
|
||||
* output[m, c*H + i] = residual[m, c*H + i] + a[m, c] * block_output[m, i]
|
||||
*
|
||||
* One CTA handles one token row. Phase 1 computes the kHcCount gate values with a
|
||||
* block reduction over the full HC*H row (fp32 accumulation). Phase 2 streams the
|
||||
* HC*H output elements with vectorized 16B accesses; the block_output row is only
|
||||
* H elements, so its re-read per branch stays in L2.
|
||||
*
|
||||
* \tparam kHcCount Number of hyper-connection branches (4 in production).
|
||||
* \tparam kHiddenSize Per-branch hidden size H. HC*H must be a multiple of
|
||||
* kNumThreads * kVecLen so the row maps exactly onto the CTA.
|
||||
* \tparam kUsePDL Whether to emit the PDL wait/trigger pair.
|
||||
* \tparam Float Element type: bf16_t | fp16_t.
|
||||
*/
|
||||
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
|
||||
__global__ __launch_bounds__(256) void hc_combine_kernel(const HcCombineParams __grid_constant__ params) {
|
||||
using namespace device;
|
||||
using Float2 = packed_t<Float>;
|
||||
using Storage = AlignedVector<Float2, 4>; // 8 elements, 16 bytes
|
||||
constexpr uint32_t kVecLen = 8;
|
||||
constexpr uint32_t kNumThreads = 256;
|
||||
constexpr int64_t kRowSize = kHcCount * kHiddenSize;
|
||||
constexpr uint32_t kVecsPerRow = kRowSize / kVecLen; // 1280 for 4x2560
|
||||
constexpr uint32_t kVecsPerThread = kVecsPerRow / kNumThreads; // 5 for 4x2560
|
||||
constexpr uint32_t kVecsPerBranch = kHiddenSize / kVecLen; // 320 for 2560
|
||||
constexpr uint32_t kNumWarps = kNumThreads / kWarpThreads;
|
||||
|
||||
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
|
||||
const uint32_t m = blockIdx.x;
|
||||
|
||||
const auto y_ptr = pointer::offset<Float>(params.block_output, static_cast<int64_t>(m) * kHiddenSize);
|
||||
const auto r_ptr = pointer::offset<Float>(params.residual, static_cast<int64_t>(m) * kRowSize);
|
||||
const auto n_ptr = pointer::offset<Float>(params.normed_residual, static_cast<int64_t>(m) * kRowSize);
|
||||
const auto w_ptr = static_cast<const Float*>(params.inject_weight);
|
||||
const auto out_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(m) * kRowSize);
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
// Phase 1: gate values a_c, accumulated in fp32 and reduced across the CTA.
|
||||
Storage n_vec[kVecsPerThread];
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
|
||||
n_vec[j] = gmem.load(n_ptr, j);
|
||||
}
|
||||
|
||||
float acc[kHcCount];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kHcCount; ++c) {
|
||||
const auto wc_ptr = w_ptr + static_cast<int64_t>(c) * kRowSize;
|
||||
float sum = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
|
||||
const Storage w_vec = gmem.load(wc_ptr, j);
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
|
||||
const auto [nx, ny] = cast<fp32x2_t>(n_vec[j][i]);
|
||||
const auto [wx, wy] = cast<fp32x2_t>(w_vec[i]);
|
||||
sum += nx * wx + ny * wy;
|
||||
}
|
||||
}
|
||||
acc[c] = warp::reduce_sum(sum);
|
||||
}
|
||||
|
||||
__shared__ float smem[kHcCount][kNumWarps];
|
||||
const uint32_t warp_id = threadIdx.x / kWarpThreads;
|
||||
const uint32_t lane = threadIdx.x % kWarpThreads;
|
||||
if (lane == 0) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kHcCount; ++c) {
|
||||
smem[c][warp_id] = acc[c];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
__shared__ float a_shared[kHcCount];
|
||||
if (threadIdx.x < kHcCount) {
|
||||
float total = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t w = 0; w < kNumWarps; ++w) {
|
||||
total += smem[threadIdx.x][w];
|
||||
}
|
||||
a_shared[threadIdx.x] = 2.0f / (1.0f + math::exp(-total / kHcCount));
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Phase 2: stream the output row. Vector `vec_idx` lies entirely inside
|
||||
// branch `vec_idx / kVecsPerBranch` (H is a multiple of kVecLen).
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
|
||||
const uint32_t vec_idx = threadIdx.x + j * kNumThreads;
|
||||
const uint32_t branch = vec_idx / kVecsPerBranch;
|
||||
const uint32_t col_in_branch = (vec_idx % kVecsPerBranch) * kVecLen;
|
||||
const float a = a_shared[branch];
|
||||
|
||||
const Storage r_vec = gmem.load(r_ptr, j);
|
||||
Storage y_vec;
|
||||
y_vec.load(y_ptr, col_in_branch / kVecLen);
|
||||
Storage out_vec;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
|
||||
const auto [rx, ry] = cast<fp32x2_t>(r_vec[i]);
|
||||
const auto [yx, yy] = cast<fp32x2_t>(y_vec[i]);
|
||||
out_vec[i] = cast<Float2>(fp32x2_t{rx + a * yx, ry + a * yy});
|
||||
}
|
||||
gmem.store(out_ptr, out_vec, j);
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename DType>
|
||||
struct HcCombineKernel {
|
||||
static_assert(sizeof(DType) == 2, "HcCombine only supports 2-byte dtypes");
|
||||
static_assert(kHcCount > 0, "kHcCount must be positive");
|
||||
static_assert(kHiddenSize > 0 && kHiddenSize % 8 == 0, "kHiddenSize must be a multiple of 8");
|
||||
static_assert((kHcCount * kHiddenSize) % (256 * 8) == 0, "kHcCount * kHiddenSize must be a multiple of 2048");
|
||||
static constexpr auto kernel = hc_combine_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
|
||||
static constexpr uint32_t kBlockSize = 256;
|
||||
|
||||
/**
|
||||
* \brief Validate tensors and launch one CTA per token row.
|
||||
* \param block_output [M, H] contiguous
|
||||
* \param residual [M, HC * H] contiguous
|
||||
* \param normed_residual [M, HC * H] contiguous, same dtype/device as residual
|
||||
* \param inject_weight [HC, HC * H] contiguous, same dtype/device as residual
|
||||
* \param output [M, HC * H] contiguous, same shape/dtype/device as residual
|
||||
*/
|
||||
static void
|
||||
run(const tvm::ffi::TensorView block_output,
|
||||
const tvm::ffi::TensorView residual,
|
||||
const tvm::ffi::TensorView normed_residual,
|
||||
const tvm::ffi::TensorView inject_weight,
|
||||
const tvm::ffi::TensorView output) {
|
||||
using namespace host;
|
||||
auto M = SymbolicSize{"num_tokens"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({M, kHiddenSize}) // block_output
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(block_output);
|
||||
TensorMatcher({M, kHcCount * kHiddenSize}) // residual, normed_residual, output
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(residual)
|
||||
.verify(normed_residual)
|
||||
.verify(output);
|
||||
TensorMatcher({kHcCount, kHcCount * kHiddenSize}) // inject_weight
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(inject_weight);
|
||||
|
||||
const auto params = HcCombineParams{
|
||||
.block_output = block_output.data_ptr(),
|
||||
.residual = residual.data_ptr(),
|
||||
.normed_residual = normed_residual.data_ptr(),
|
||||
.inject_weight = inject_weight.data_ptr(),
|
||||
.output = output.data_ptr(),
|
||||
};
|
||||
|
||||
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
|
||||
LaunchKernel(num_tokens, kBlockSize, device.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(kernel, params);
|
||||
}
|
||||
};
|
||||
|
||||
struct HcCombineSplitParams {
|
||||
const void* block_output;
|
||||
const void* residual;
|
||||
const void* normed_residual;
|
||||
const void* inject_weight;
|
||||
void* output;
|
||||
float* partials;
|
||||
};
|
||||
|
||||
namespace hc_combine_split_detail {
|
||||
|
||||
// One CTA per warp of the original kernel: CTA s replays exactly the work of
|
||||
// threads [32s, 32s+32) of the 256-thread reference, so every float is
|
||||
// accumulated in the reference order and results stay bit-identical.
|
||||
constexpr uint32_t kSplit = 8;
|
||||
constexpr uint32_t kRefThreads = 256;
|
||||
constexpr uint32_t kGateThreads = 32;
|
||||
constexpr uint32_t kApplyThreads = 160;
|
||||
constexpr uint32_t kVecLen = 8;
|
||||
|
||||
} // namespace hc_combine_split_detail
|
||||
|
||||
/**
|
||||
* \brief Stage 1 of the split combine: partial gate dots over a K slice.
|
||||
*
|
||||
* Grid is (rows, kSplit) so the [HC, HC*H] inject weight is read once across
|
||||
* the whole grid instead of once per row, and the row's traffic is spread over
|
||||
* kSplit CTAs. Each CTA writes its own partials slot, so no atomics and no
|
||||
* buffer clearing are needed.
|
||||
*/
|
||||
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
|
||||
__global__ __launch_bounds__(hc_combine_split_detail::kGateThreads) void hc_combine_gate_kernel(
|
||||
const HcCombineSplitParams __grid_constant__ params) {
|
||||
using namespace device;
|
||||
using namespace hc_combine_split_detail;
|
||||
using Float2 = packed_t<Float>;
|
||||
using Storage = AlignedVector<Float2, 4>;
|
||||
constexpr uint32_t kVecLen = 8;
|
||||
constexpr int64_t kRowSize = kHcCount * kHiddenSize;
|
||||
constexpr uint32_t kVecsPerRow = kRowSize / kVecLen;
|
||||
constexpr uint32_t kVecsPerThread = kVecsPerRow / kRefThreads;
|
||||
static_assert(kVecsPerRow % kRefThreads == 0);
|
||||
static_assert(kRefThreads / kGateThreads == kSplit);
|
||||
|
||||
const uint32_t m = blockIdx.x;
|
||||
const uint32_t split = blockIdx.y / kHcCount;
|
||||
const uint32_t c = blockIdx.y % kHcCount;
|
||||
const uint32_t ref_tid = split * kGateThreads + threadIdx.x;
|
||||
|
||||
const auto n_ptr = pointer::offset<Float>(params.normed_residual, static_cast<int64_t>(m) * kRowSize);
|
||||
const auto w_ptr = static_cast<const Float*>(params.inject_weight);
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
Storage n_vec[kVecsPerThread];
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
|
||||
n_vec[j].load(n_ptr, ref_tid + j * kRefThreads);
|
||||
}
|
||||
|
||||
{
|
||||
const auto wc_ptr = w_ptr + static_cast<int64_t>(c) * kRowSize;
|
||||
float sum = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
|
||||
Storage w_vec;
|
||||
w_vec.load(wc_ptr, ref_tid + j * kRefThreads);
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
|
||||
const auto [nx, ny] = cast<fp32x2_t>(n_vec[j][i]);
|
||||
const auto [wx, wy] = cast<fp32x2_t>(w_vec[i]);
|
||||
sum += nx * wx + ny * wy;
|
||||
}
|
||||
}
|
||||
sum = warp::reduce_sum(sum);
|
||||
if (threadIdx.x == 0) {
|
||||
params.partials[(static_cast<int64_t>(m) * kSplit + split) * kHcCount + c] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Stage 2: reduce the partial dots and stream the combined row.
|
||||
*
|
||||
* Each CTA owns one contiguous vector slice of the row. kSplit divides
|
||||
* kHiddenSize, so a slice never straddles two branches and the gate is a
|
||||
* per-CTA scalar.
|
||||
*/
|
||||
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
|
||||
__global__ __launch_bounds__(hc_combine_split_detail::kApplyThreads) void hc_combine_apply_kernel(
|
||||
const HcCombineSplitParams __grid_constant__ params) {
|
||||
using namespace device;
|
||||
using namespace hc_combine_split_detail;
|
||||
using Float2 = packed_t<Float>;
|
||||
using Storage = AlignedVector<Float2, 4>;
|
||||
constexpr int64_t kRowSize = kHcCount * kHiddenSize;
|
||||
constexpr uint32_t kVecsPerRow = kRowSize / kVecLen;
|
||||
constexpr uint32_t kVecsPerSplit = kVecsPerRow / kSplit;
|
||||
constexpr uint32_t kVecsPerThread = kVecsPerSplit / kApplyThreads;
|
||||
constexpr uint32_t kVecsPerBranch = kHiddenSize / kVecLen;
|
||||
static_assert(kVecsPerBranch % kVecsPerSplit == 0);
|
||||
|
||||
const uint32_t m = blockIdx.x;
|
||||
const uint32_t split = blockIdx.y;
|
||||
const uint32_t vec_base = split * kVecsPerSplit;
|
||||
const uint32_t branch = vec_base / kVecsPerBranch;
|
||||
|
||||
const auto y_ptr = pointer::offset<Float>(params.block_output, static_cast<int64_t>(m) * kHiddenSize);
|
||||
const auto r_ptr = pointer::offset<Float>(params.residual, static_cast<int64_t>(m) * kRowSize);
|
||||
const auto out_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(m) * kRowSize);
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
float total = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t s = 0; s < kSplit; ++s) {
|
||||
total += params.partials[(static_cast<int64_t>(m) * kSplit + s) * kHcCount + branch];
|
||||
}
|
||||
const float a = 2.0f / (1.0f + math::exp(-total / kHcCount));
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
|
||||
const uint32_t vec_idx = vec_base + threadIdx.x + j * kApplyThreads;
|
||||
const uint32_t col_in_branch = (vec_idx % kVecsPerBranch);
|
||||
Storage r_vec;
|
||||
r_vec.load(r_ptr, vec_idx);
|
||||
Storage y_vec;
|
||||
y_vec.load(y_ptr, col_in_branch);
|
||||
Storage out_vec;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
|
||||
const auto [rx, ry] = cast<fp32x2_t>(r_vec[i]);
|
||||
const auto [yx, yy] = cast<fp32x2_t>(y_vec[i]);
|
||||
out_vec[i] = cast<Float2>(fp32x2_t{rx + a * yx, ry + a * yy});
|
||||
}
|
||||
out_vec.store(out_ptr, vec_idx);
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename DType>
|
||||
struct HcCombineSplitKernel {
|
||||
static_assert(sizeof(DType) == 2, "HcCombine only supports 2-byte dtypes");
|
||||
static constexpr auto gate_kernel = hc_combine_gate_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
|
||||
static constexpr auto apply_kernel = hc_combine_apply_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
|
||||
|
||||
static void
|
||||
run(const tvm::ffi::TensorView block_output,
|
||||
const tvm::ffi::TensorView residual,
|
||||
const tvm::ffi::TensorView normed_residual,
|
||||
const tvm::ffi::TensorView inject_weight,
|
||||
const tvm::ffi::TensorView output,
|
||||
const tvm::ffi::TensorView partials) {
|
||||
using namespace host;
|
||||
using namespace hc_combine_split_detail;
|
||||
auto M = SymbolicSize{"num_tokens"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({M, kHiddenSize}).with_dtype<DType>().with_device(device).verify(block_output);
|
||||
TensorMatcher({M, kHcCount * kHiddenSize})
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(residual)
|
||||
.verify(normed_residual)
|
||||
.verify(output);
|
||||
TensorMatcher({kHcCount, kHcCount * kHiddenSize}).with_dtype<DType>().with_device(device).verify(inject_weight);
|
||||
auto part_rows = SymbolicSize{"partial_rows"};
|
||||
TensorMatcher({part_rows, kSplit, kHcCount}).with_dtype<fp32_t>().with_device(device).verify(partials);
|
||||
|
||||
const auto params = HcCombineSplitParams{
|
||||
.block_output = block_output.data_ptr(),
|
||||
.residual = residual.data_ptr(),
|
||||
.normed_residual = normed_residual.data_ptr(),
|
||||
.inject_weight = inject_weight.data_ptr(),
|
||||
.output = output.data_ptr(),
|
||||
.partials = static_cast<float*>(partials.data_ptr()),
|
||||
};
|
||||
|
||||
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
|
||||
LaunchKernel(dim3(num_tokens, kSplit * kHcCount, 1), kGateThreads, device.unwrap())
|
||||
.enable_pdl(kUsePDL)(gate_kernel, params);
|
||||
LaunchKernel(dim3(num_tokens, kSplit, 1), kApplyThreads, device.unwrap()).enable_pdl(kUsePDL)(apply_kernel, params);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang
|
||||
@@ -20,6 +20,7 @@ load Triton, CUTLASS, or compile a JIT extension.
|
||||
| Diffusion residual-gate add | `residual_gate_add_jit.py` | [sgl-project/sglang#29361](https://github.com/sgl-project/sglang/pull/29361), merge commit `495f13fa12` |
|
||||
| LTX2 QK-norm split-RoPE | `ltx2_qknorm_split_rope_jit.py` | [sgl-project/sglang#29708](https://github.com/sgl-project/sglang/pull/29708), merge commit `fcb9f229b3` |
|
||||
| FLUX.2 FP8 producer and QKV packing fusions | `layernorm_modulate_triton.py`, `flux2_qkv_epilogue_jit.py`, `flux2_token_cat_fp8_triton.py` | [sgl-project/sglang#37162](https://github.com/sgl-project/sglang/pull/37162), merge commit `1c3ad92438` |
|
||||
| Qwen3.8 QSA packed-varlen decode on SM121 | `qwen38_qsa_sm121/` | [radixark/KDA-1.5#4](https://github.com/radixark/KDA-1.5/pull/4) at `414ce456e14a`; see the package README |
|
||||
|
||||
For JIT kernels, the Python entry module and the corresponding source under
|
||||
`csrc/` move together. The shared `sglang.kernels.jit` loader remains build
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
# Qwen3.8 QSA packed-varlen decode for SM121
|
||||
|
||||
This implementation was optimized by Codex and Kimi K3 agents through
|
||||
[KDA-1.5](https://github.com/radixark/KDA-1.5). The task and immutable real
|
||||
tensor replay were registered in [radixark/KDA-1.5 PR #4](https://github.com/radixark/KDA-1.5/pull/4)
|
||||
at commit `414ce456e14ae8546f77d9356d2c4d955c5bb7f1`. This package integrates
|
||||
winning submission `b4181149c8884ddb`; its byte-exact submitted source has SHA256
|
||||
`4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909`.
|
||||
|
||||
The kernel is specialized for the packed QSA decode tensors captured from
|
||||
`RadixArk/Qwen3.8-Flash-Next-NVFP4` on NVIDIA GB10 (SM121):
|
||||
|
||||
- BF16 query, key, value, and output with head dimension 256
|
||||
- one packed query row per sequence and device-side `cu_seqlens`
|
||||
- 12 query heads per KV head: TP1 uses 24Q/2KV and TP2 uses 12Q/1KV
|
||||
- all query-row counts in the validated `1 <= bs <= 128` envelope
|
||||
- `max_seqlen_k` capacity up to 2055 and captured logical selected-KV lengths
|
||||
up to 2051 rows per sequence
|
||||
|
||||
The implementation groups the 12 query heads that share one KV head into one
|
||||
CTA, uses BF16 tensor-core QK/PV products with FP32 online-softmax state, and
|
||||
splits long KV rows across multiple CTAs. The last arriving split performs a
|
||||
stable FP32 merge and resets its device counter in the same launch. A
|
||||
host-visible shape/topology policy selects the two measured schedules, while
|
||||
the live device `cu_seqlens_k` selects one, two, four, or eight active splits
|
||||
without a host synchronization.
|
||||
|
||||
SM121 dispatch checks the exact Qwen3.8 contract and routes directly to this
|
||||
kernel; it is the only packed-QSA attention implementation added by this PR.
|
||||
The KDA replay passed all 15 TP1/TP2 production tensors on two independent GB10
|
||||
GPUs, and the final source passed 150,000 consecutive launches with all
|
||||
counters returning to zero.
|
||||
|
||||
After adaptation into SGLang, the packaged kernel passed the same 15/15 replay
|
||||
with exactly one CUDA activity per row and a 2.0702x all-shape geomean over the
|
||||
generic Triton fallback (1.6951x large, 2.3653x small). On one DGX Spark running
|
||||
the full TP1 NVFP4 model with NEXTN, three-round low-concurrency serving A/B
|
||||
improved total token throughput by 4.45% at concurrency 1 and 4.00% at
|
||||
concurrency 4. A 50-example, five-shot GSM8K A/B with a 2048-token output limit
|
||||
scored 49/50 for both Triton and KDA, with the same single failed example.
|
||||
|
||||
An additional synthetic GB10 sweep covers both TP topologies, every batch size
|
||||
from 1 through 16, and short plus saturated KV rows. All 64 cases passed; the
|
||||
maximum relative L2 against the original correct Triton implementation was
|
||||
0.002422, and speedup ranged from 1.41x to 5.09x.
|
||||
|
||||
A follow-up extended-batch sweep covers both TP topologies, batch sizes
|
||||
17/24/32/48/64/96/128, and short, medium, plus saturated KV rows. All 42 cases
|
||||
passed with maximum relative L2 0.002410. Geomean speedup was 4.48x, the slowest
|
||||
case still improved by 1.58x, and no case regressed. The packaged scratch space
|
||||
is therefore sized for the largest tested batch, 128.
|
||||
@@ -0,0 +1,83 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SUPPORTED_HEAD_TOPOLOGIES = frozenset({(12, 1), (24, 2)})
|
||||
# Largest batch qualified by the extended GB10 baseline sweep.
|
||||
_MAX_BATCH = 128
|
||||
_MAX_SELECTED_KV = 2055
|
||||
_logged_fast_path = False
|
||||
|
||||
|
||||
def can_use_qwen38_qsa_sm121(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
max_seqlen_k: int,
|
||||
) -> bool:
|
||||
"""Return whether this call matches the captured Qwen3.8 SM121 contract."""
|
||||
if not q.is_cuda or q.ndim != 3 or q.dtype != torch.bfloat16:
|
||||
return False
|
||||
batch, num_q_heads, head_dim = q.shape
|
||||
if not (0 < batch <= _MAX_BATCH) or head_dim != 256:
|
||||
return False
|
||||
if k.ndim != 3 or v.shape != k.shape or k.dtype != q.dtype or v.dtype != q.dtype:
|
||||
return False
|
||||
num_kv_heads = k.shape[1]
|
||||
if (num_q_heads, num_kv_heads) not in _SUPPORTED_HEAD_TOPOLOGIES:
|
||||
return False
|
||||
if k.shape[2] != head_dim or not (0 < max_seqlen_k <= _MAX_SELECTED_KV):
|
||||
return False
|
||||
if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous():
|
||||
return False
|
||||
if q.device != k.device or q.device != v.device:
|
||||
return False
|
||||
if cu_seqlens_q.device != q.device or cu_seqlens_k.device != q.device:
|
||||
return False
|
||||
if cu_seqlens_q.dtype != torch.int32 or cu_seqlens_k.dtype != torch.int32:
|
||||
return False
|
||||
if cu_seqlens_q.ndim != 1 or cu_seqlens_k.ndim != 1:
|
||||
return False
|
||||
if not cu_seqlens_q.is_contiguous() or not cu_seqlens_k.is_contiguous():
|
||||
return False
|
||||
if cu_seqlens_q.numel() != batch + 1 or cu_seqlens_k.numel() != batch + 1:
|
||||
return False
|
||||
properties = torch.cuda.get_device_properties(q.device)
|
||||
return (properties.major, properties.minor) == (12, 1)
|
||||
|
||||
|
||||
def qwen38_qsa_sm121(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
max_seqlen_k: int,
|
||||
softmax_scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Run the KDA-generated Qwen3.8 packed QSA decode kernel."""
|
||||
global _logged_fast_path
|
||||
if not can_use_qwen38_qsa_sm121(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k):
|
||||
raise ValueError("unsupported call for the KDA Qwen3.8 SM121 QSA kernel")
|
||||
|
||||
from .kernel import qwen38_qsa_sm121 as run_kernel
|
||||
|
||||
if not _logged_fast_path:
|
||||
logger.info(
|
||||
"Using the Codex/Kimi K3 KDA Qwen3.8 QSA kernel on SM121 "
|
||||
"(radixark/KDA-1.5#4, submission b4181149c8884ddb)"
|
||||
)
|
||||
_logged_fast_path = True
|
||||
return run_kernel(q, k, v, cu_seqlens_q, cu_seqlens_k, softmax_scale)
|
||||
|
||||
|
||||
__all__ = ["can_use_qwen38_qsa_sm121", "qwen38_qsa_sm121"]
|
||||
@@ -0,0 +1,269 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# KDA-1.5 submission b4181149c8884ddb (https://github.com/radixark/KDA-1.5/pull/4);
|
||||
# source SHA256 4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909.
|
||||
"""Shape-specialized Qwen3.8 packed QSA decode kernel for SM121."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _qsa_split_kernel(
|
||||
q_ptr,
|
||||
k_ptr,
|
||||
v_ptr,
|
||||
out_ptr,
|
||||
cu_q_ptr,
|
||||
cu_k_ptr,
|
||||
partial_max_ptr,
|
||||
partial_sum_ptr,
|
||||
partial_acc_ptr,
|
||||
counter_ptr,
|
||||
softmax_scale,
|
||||
NUM_Q_HEADS: tl.constexpr,
|
||||
NUM_KV_HEADS: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_KV: tl.constexpr,
|
||||
MAX_SPLITS: tl.constexpr,
|
||||
q_stride_t: tl.constexpr,
|
||||
q_stride_h: tl.constexpr,
|
||||
k_stride_t: tl.constexpr,
|
||||
k_stride_h: tl.constexpr,
|
||||
v_stride_t: tl.constexpr,
|
||||
v_stride_h: tl.constexpr,
|
||||
out_stride_t: tl.constexpr,
|
||||
out_stride_h: tl.constexpr,
|
||||
):
|
||||
sequence_idx = tl.program_id(0)
|
||||
split_program = tl.program_id(1)
|
||||
kv_head_idx = split_program // MAX_SPLITS
|
||||
split_idx = split_program - kv_head_idx * MAX_SPLITS
|
||||
slot = sequence_idx * NUM_KV_HEADS + kv_head_idx
|
||||
queries_per_kv = NUM_Q_HEADS // NUM_KV_HEADS
|
||||
|
||||
query_idx = tl.load(cu_q_ptr + sequence_idx)
|
||||
kv_begin = tl.load(cu_k_ptr + sequence_idx)
|
||||
kv_end = tl.load(cu_k_ptr + sequence_idx + 1)
|
||||
kv_count = kv_end - kv_begin
|
||||
tile_count = tl.cdiv(kv_count, BLOCK_KV)
|
||||
|
||||
# The split count depends only on live device metadata and the static
|
||||
# launch geometry, so CUDA Graph replay needs no host readback.
|
||||
batch = tl.num_programs(0)
|
||||
n_splits = 1
|
||||
if tile_count >= 1536 // BLOCK_KV:
|
||||
n_splits = 2
|
||||
elif tile_count >= 512 // BLOCK_KV and batch * NUM_KV_HEADS <= 4:
|
||||
n_splits = 4
|
||||
if batch == 1:
|
||||
if tile_count >= 1024 // BLOCK_KV:
|
||||
n_splits = 8
|
||||
elif tile_count >= 512 // BLOCK_KV:
|
||||
n_splits = 4
|
||||
elif tile_count >= 256 // BLOCK_KV:
|
||||
n_splits = 2
|
||||
if split_idx >= n_splits:
|
||||
return
|
||||
|
||||
tile_lo = (tile_count * split_idx) // n_splits
|
||||
tile_hi = (tile_count * (split_idx + 1)) // n_splits
|
||||
kv_start = kv_begin + tile_lo * BLOCK_KV
|
||||
kv_stop = tl.minimum(kv_begin + tile_hi * BLOCK_KV, kv_end)
|
||||
|
||||
m = tl.arange(0, BLOCK_M)
|
||||
d = tl.arange(0, HEAD_DIM)
|
||||
q_head = kv_head_idx * queries_per_kv + m
|
||||
q_mask = m < queries_per_kv
|
||||
query = tl.load(
|
||||
q_ptr + query_idx * q_stride_t + q_head[:, None] * q_stride_h + d[None, :],
|
||||
mask=q_mask[:, None],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
n0 = tl.arange(0, BLOCK_KV)
|
||||
k_rows = (
|
||||
k_ptr
|
||||
+ kv_head_idx * k_stride_h
|
||||
+ kv_start.to(tl.int64) * k_stride_t
|
||||
+ n0[:, None] * k_stride_t
|
||||
)
|
||||
v_rows = (
|
||||
v_ptr
|
||||
+ kv_head_idx * v_stride_h
|
||||
+ kv_start.to(tl.int64) * v_stride_t
|
||||
+ n0[:, None] * v_stride_t
|
||||
)
|
||||
|
||||
running_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
|
||||
running_sum = tl.zeros([BLOCK_M], tl.float32)
|
||||
accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
|
||||
|
||||
kv_len = kv_stop - kv_start
|
||||
full_end = kv_start + (kv_len // BLOCK_KV) * BLOCK_KV
|
||||
for block_start in range(kv_start, full_end, BLOCK_KV):
|
||||
keys = tl.load(k_rows + d[None, :])
|
||||
scores = tl.dot(query, tl.trans(keys)) * softmax_scale
|
||||
new_max = tl.maximum(running_max, tl.max(scores, axis=1))
|
||||
old_scale = tl.exp(running_max - new_max)
|
||||
probabilities = tl.exp(scores - new_max[:, None])
|
||||
running_sum = running_sum * old_scale + tl.sum(probabilities, axis=1)
|
||||
values = tl.load(v_rows + d[None, :])
|
||||
accumulator = accumulator * old_scale[:, None] + tl.dot(
|
||||
probabilities.to(tl.bfloat16), values
|
||||
)
|
||||
running_max = new_max
|
||||
k_rows += BLOCK_KV * k_stride_t
|
||||
v_rows += BLOCK_KV * v_stride_t
|
||||
|
||||
if full_end < kv_stop:
|
||||
n = full_end + n0
|
||||
n_mask = n < kv_stop
|
||||
keys = tl.load(k_rows + d[None, :], mask=n_mask[:, None], other=0.0)
|
||||
scores = tl.dot(query, tl.trans(keys)) * softmax_scale
|
||||
scores = tl.where(n_mask[None, :], scores, -float("inf"))
|
||||
new_max = tl.maximum(running_max, tl.max(scores, axis=1))
|
||||
old_scale = tl.exp(running_max - new_max)
|
||||
probabilities = tl.exp(scores - new_max[:, None])
|
||||
running_sum = running_sum * old_scale + tl.sum(probabilities, axis=1)
|
||||
values = tl.load(v_rows + d[None, :], mask=n_mask[:, None], other=0.0)
|
||||
accumulator = accumulator * old_scale[:, None] + tl.dot(
|
||||
probabilities.to(tl.bfloat16), values
|
||||
)
|
||||
running_max = new_max
|
||||
|
||||
if n_splits == 1:
|
||||
output = accumulator / tl.where(running_sum > 0.0, running_sum, 1.0)[:, None]
|
||||
tl.store(
|
||||
out_ptr
|
||||
+ query_idx * out_stride_t
|
||||
+ q_head[:, None] * out_stride_h
|
||||
+ d[None, :],
|
||||
output.to(out_ptr.dtype.element_ty),
|
||||
mask=q_mask[:, None],
|
||||
)
|
||||
return
|
||||
|
||||
partial_row = (slot * MAX_SPLITS + split_idx) * BLOCK_M + m
|
||||
tl.store(partial_max_ptr + partial_row, running_max)
|
||||
tl.store(partial_sum_ptr + partial_row, running_sum)
|
||||
tl.store(
|
||||
partial_acc_ptr + partial_row[:, None] * HEAD_DIM + d[None, :],
|
||||
accumulator,
|
||||
)
|
||||
tl.debug_barrier()
|
||||
arrival = tl.atomic_add(counter_ptr + slot, 1, sem="acq_rel", scope="gpu")
|
||||
if arrival == n_splits - 1:
|
||||
merged_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
|
||||
for j in tl.static_range(MAX_SPLITS):
|
||||
j_ok = (j < n_splits) & (m < BLOCK_M)
|
||||
row = (slot * MAX_SPLITS + j) * BLOCK_M + m
|
||||
mj = tl.load(partial_max_ptr + row, mask=j_ok, other=-float("inf"))
|
||||
merged_max = tl.maximum(merged_max, mj)
|
||||
merged_sum = tl.zeros([BLOCK_M], tl.float32)
|
||||
merged_acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
|
||||
for j in tl.static_range(MAX_SPLITS):
|
||||
j_ok = (j < n_splits) & (m < BLOCK_M)
|
||||
row = (slot * MAX_SPLITS + j) * BLOCK_M + m
|
||||
mj = tl.load(partial_max_ptr + row, mask=j_ok, other=-float("inf"))
|
||||
lj = tl.load(partial_sum_ptr + row, mask=j_ok, other=0.0)
|
||||
weight = tl.exp(mj - merged_max)
|
||||
merged_sum += weight * lj
|
||||
partial = tl.load(
|
||||
partial_acc_ptr + row[:, None] * HEAD_DIM + d[None, :],
|
||||
mask=j_ok[:, None],
|
||||
other=0.0,
|
||||
)
|
||||
merged_acc += weight[:, None] * partial
|
||||
output = merged_acc / tl.where(merged_sum > 0.0, merged_sum, 1.0)[:, None]
|
||||
tl.store(
|
||||
out_ptr
|
||||
+ query_idx * out_stride_t
|
||||
+ q_head[:, None] * out_stride_h
|
||||
+ d[None, :],
|
||||
output.to(out_ptr.dtype.element_ty),
|
||||
mask=q_mask[:, None],
|
||||
)
|
||||
tl.atomic_xchg(counter_ptr + slot, 0, sem="release", scope="gpu")
|
||||
|
||||
|
||||
# The worst-case TP1 scratch allocation at the qualified limit is 32.3 MiB.
|
||||
_MAX_BATCH = 128
|
||||
_MAX_KV_HEADS = 2
|
||||
_BLOCK_M = 16
|
||||
_MAX_SPLITS = 8
|
||||
_MAX_SLOTS = _MAX_BATCH * _MAX_KV_HEADS
|
||||
_HEAD_DIM = 256
|
||||
_scratch: dict[int, tuple[torch.Tensor, ...]] = {}
|
||||
|
||||
|
||||
def _get_scratch(device: torch.device) -> tuple[torch.Tensor, ...]:
|
||||
device_index = (
|
||||
device.index if device.index is not None else torch.cuda.current_device()
|
||||
)
|
||||
scratch = _scratch.get(device_index)
|
||||
if scratch is None:
|
||||
rows = _MAX_SLOTS * _MAX_SPLITS * _BLOCK_M
|
||||
scratch = (
|
||||
torch.empty(rows, dtype=torch.float32, device=device),
|
||||
torch.empty(rows, dtype=torch.float32, device=device),
|
||||
torch.empty(rows * _HEAD_DIM, dtype=torch.float32, device=device),
|
||||
torch.zeros(_MAX_SLOTS, dtype=torch.int32, device=device),
|
||||
)
|
||||
_scratch[device_index] = scratch
|
||||
return scratch
|
||||
|
||||
|
||||
def qwen38_qsa_sm121(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
softmax_scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Run the KDA-generated SM121 Qwen3.8 QSA kernel."""
|
||||
output = torch.empty_like(q)
|
||||
partial_max, partial_sum, partial_acc, counters = _get_scratch(q.device)
|
||||
batch, num_q_heads, head_dim = q.shape
|
||||
num_kv_heads = k.shape[1]
|
||||
|
||||
# Measured schedule. The TP1 q_rows=4 shape stays on BK64 on purpose:
|
||||
# its short and saturated rows cannot be told apart from the host.
|
||||
use_bk32 = (num_kv_heads == 1 and batch < 12) or (num_kv_heads == 2 and batch < 4)
|
||||
block_kv = 32 if use_bk32 else 64
|
||||
stages = 3 if use_bk32 else 2
|
||||
_qsa_split_kernel[(batch, num_kv_heads * _MAX_SPLITS)](
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
partial_max,
|
||||
partial_sum,
|
||||
partial_acc,
|
||||
counters,
|
||||
softmax_scale,
|
||||
NUM_Q_HEADS=num_q_heads,
|
||||
NUM_KV_HEADS=num_kv_heads,
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_M=_BLOCK_M,
|
||||
BLOCK_KV=block_kv,
|
||||
MAX_SPLITS=_MAX_SPLITS,
|
||||
q_stride_t=q.stride(0),
|
||||
q_stride_h=q.stride(1),
|
||||
k_stride_t=k.stride(0),
|
||||
k_stride_h=k.stride(1),
|
||||
v_stride_t=v.stride(0),
|
||||
v_stride_h=v.stride(1),
|
||||
out_stride_t=output.stride(0),
|
||||
out_stride_h=output.stride(1),
|
||||
num_warps=4,
|
||||
num_stages=stages,
|
||||
)
|
||||
return output
|
||||
@@ -6,8 +6,21 @@ The Triton kernels migrated here live in this package
|
||||
KV-cache index/write kernels went to the ``kvcache`` group instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sglang.kernels.registry import register_kernel
|
||||
from sglang.kernels.spec import KernelBackend, KernelSpec
|
||||
from sglang.kernels.selector import get_kernel
|
||||
from sglang.kernels.spec import (
|
||||
CapabilityRequirement,
|
||||
FormatSignature,
|
||||
KernelBackend,
|
||||
KernelSpec,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
# (module, public_fn) migrated from layers/attention/triton_ops + model_executor.
|
||||
_TRITON_KERNELS = [
|
||||
@@ -45,7 +58,78 @@ for _mod, _fn in _TRITON_KERNELS:
|
||||
)
|
||||
del _mod, _fn
|
||||
|
||||
__all__ = []
|
||||
register_kernel(
|
||||
KernelSpec(
|
||||
op="attention.kda_qwen38_qsa_sm121",
|
||||
backend=KernelBackend.TRITON,
|
||||
target=("sglang.kernels.kda_kernels.qwen38_qsa_sm121:qwen38_qsa_sm121"),
|
||||
capabilities=frozenset(
|
||||
{CapabilityRequirement.cuda(min_sm=(12, 1), max_sm=(12, 1))}
|
||||
),
|
||||
format_signature=FormatSignature(
|
||||
supported_dtypes=("bfloat16",),
|
||||
description=(
|
||||
"Qwen3.8 packed QSA decode: D=256, 12:1 GQA, 1 <= q_rows <= 128"
|
||||
),
|
||||
),
|
||||
description=(
|
||||
"SM121 Qwen3.8 QSA decode optimized by Codex/Kimi K3 through KDA-1.5."
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def can_use_kda_qwen38_qsa_sm121(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
max_seqlen_k: int,
|
||||
) -> bool:
|
||||
"""Check the exact E2E-qualified Qwen3.8/SM121 QSA contract."""
|
||||
from sglang.kernels.kda_kernels.qwen38_qsa_sm121 import (
|
||||
can_use_qwen38_qsa_sm121,
|
||||
)
|
||||
|
||||
return can_use_qwen38_qsa_sm121(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k)
|
||||
|
||||
|
||||
def qwen38_qsa_sm121_varlen(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
max_seqlen_q: int = 1,
|
||||
max_seqlen_k: int = 0,
|
||||
softmax_scale: float = 1.0,
|
||||
causal: bool = True,
|
||||
**_: object,
|
||||
) -> torch.Tensor:
|
||||
"""Run the only SM121 packed-QSA kernel for its qualified contract."""
|
||||
del causal
|
||||
if max_seqlen_q != 1:
|
||||
raise ValueError(f"QSA requires max_seqlen_q=1, got {max_seqlen_q}")
|
||||
if not can_use_kda_qwen38_qsa_sm121(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k
|
||||
):
|
||||
raise ValueError(
|
||||
"unsupported SM121 QSA call: expected BF16 D=256, 12:1 GQA, "
|
||||
"TP1 24Q/2KV or TP2 12Q/1KV, bs<=128, and selected KV<=2055"
|
||||
)
|
||||
return get_kernel("attention.kda_qwen38_qsa_sm121", KernelBackend.TRITON)(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seqlen_k,
|
||||
softmax_scale,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["can_use_kda_qwen38_qsa_sm121", "qwen38_qsa_sm121_varlen"]
|
||||
|
||||
|
||||
# Vendored linear-attention (flash-linear-attention port) kernels relocated
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
"""Fused QSA (Qwen4-Exp sparse attention) indexer-prep kernels.
|
||||
|
||||
``qsa_index_q_norm_rope_store`` fuses, per token, the eager chain
|
||||
split -> GemmaRMSNorm(index q) -> MRoPE(index q) -> raw-K store ->
|
||||
RoPE-position store into one kernel launch.
|
||||
|
||||
``qsa_index_k_compress_store`` fuses, per completed compress group, the eager
|
||||
chain gather -> fp32 mean -> GemmaRMSNorm -> MRoPE(group-start position) ->
|
||||
compressed-cache store into one kernel launch.
|
||||
|
||||
Both kernels reproduce the eager numerics step by step (fp32 norm reduction,
|
||||
per-op rounding to the storage dtype during RoPE, fp32 group mean rounded to
|
||||
the storage dtype before the norm). Outputs are bit-comparable to the eager
|
||||
indexer path: the eager RMSNorm (flashinfer's CuTe DSL kernel) reduces sums of
|
||||
squares in an order that cannot be reproduced exactly, so a small fraction of
|
||||
rows (~1 in 30k) may flip by one bf16 ulp on a rounding boundary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_qsa_indexer_module(
|
||||
dtype: torch.dtype, head_dim: int, is_neox_style: bool
|
||||
) -> Module:
|
||||
"""Compile and cache the JIT QSA indexer module for one specialisation."""
|
||||
if dtype not in (torch.bfloat16, torch.float16):
|
||||
raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
|
||||
if head_dim not in (64, 128, 256):
|
||||
raise RuntimeError(
|
||||
f"Unsupported index head_dim {head_dim}. Supported: 64, 128, 256"
|
||||
)
|
||||
args = make_cpp_args(dtype, head_dim, is_neox_style, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
"qsa_indexer",
|
||||
*args,
|
||||
cuda_files=["attention/qsa_indexer.cuh"],
|
||||
cuda_wrappers=[
|
||||
("q_prep", f"qsa_index_q_prep<{args}>"),
|
||||
("k_compress", f"qsa_index_k_compress<{args}>"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def qsa_index_q_norm_rope_store(
|
||||
qk: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
axis_map: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
cache_loc: torch.Tensor,
|
||||
key_state_buffer: torch.Tensor,
|
||||
rope_position_buffer: torch.Tensor,
|
||||
num_q_heads: int,
|
||||
rotary_dim: int,
|
||||
eps: float,
|
||||
is_neox_style: bool,
|
||||
q_heads_padded: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Per-token fused index-Q prep.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
qk : CUDA bf16/fp16 [tokens, (num_q_heads + 1) * head_dim],
|
||||
contiguous fused index Q/K projection output
|
||||
positions : CUDA int64 [tokens] or [3, tokens] RoPE positions
|
||||
(a strided trailing slice is accepted)
|
||||
cos_sin_cache : CUDA fp32 [capacity, rotary_dim] RoPE cache
|
||||
axis_map : CUDA int32 [rotary_dim // 2] position-axis per pair index
|
||||
weight : [head_dim] gemma norm weight (kernel applies 1 + w)
|
||||
cache_loc : CUDA int64 [tokens] state slots of each token
|
||||
key_state_buffer : CUDA [slots, head_dim] raw-K state buffer (written)
|
||||
rope_position_buffer : CUDA int64 [slots, 3] position buffer (written)
|
||||
num_q_heads : number of index query heads
|
||||
rotary_dim : rotated prefix of each head row
|
||||
eps : RMSNorm epsilon
|
||||
is_neox_style : NeoX (True) or GPT-J (False) RoPE pairing
|
||||
q_heads_padded : output head count; heads >= num_q_heads are zero-filled
|
||||
(defaults to num_q_heads)
|
||||
|
||||
Returns
|
||||
-------
|
||||
CUDA tensor [tokens, q_heads_padded, head_dim]: normed + rotated index Q.
|
||||
"""
|
||||
num_tokens = qk.shape[0]
|
||||
head_dim = weight.shape[0]
|
||||
if q_heads_padded is None:
|
||||
q_heads_padded = num_q_heads
|
||||
if positions.ndim == 1:
|
||||
positions = positions.unsqueeze(0)
|
||||
q_out = torch.empty(
|
||||
(num_tokens, q_heads_padded, head_dim), dtype=qk.dtype, device=qk.device
|
||||
)
|
||||
module = _jit_qsa_indexer_module(qk.dtype, head_dim, is_neox_style)
|
||||
module.q_prep(
|
||||
qk,
|
||||
q_out,
|
||||
weight,
|
||||
cos_sin_cache,
|
||||
axis_map,
|
||||
positions,
|
||||
positions.shape[0],
|
||||
cache_loc,
|
||||
key_state_buffer,
|
||||
rope_position_buffer,
|
||||
num_q_heads,
|
||||
rotary_dim,
|
||||
eps,
|
||||
)
|
||||
return q_out
|
||||
|
||||
|
||||
def qsa_index_k_compress_store(
|
||||
key_state_buffer: torch.Tensor,
|
||||
group_locs: torch.Tensor,
|
||||
rope_position_buffer: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
axis_map: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
write_locs: torch.Tensor,
|
||||
compressed_k_buffer: torch.Tensor,
|
||||
compress_ratio: int,
|
||||
rotary_dim: int,
|
||||
eps: float,
|
||||
is_neox_style: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Per-group compressed-K prep (in-place store into ``compressed_k_buffer``).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
key_state_buffer : CUDA bf16/fp16 [slots, head_dim] raw-K state buffer
|
||||
group_locs : CUDA int32 [groups, compress_ratio] state slots of each
|
||||
completed group, group-start slot in column 0 (its RoPE
|
||||
position rotates the group; other columns any order)
|
||||
rope_position_buffer : CUDA int64 [slots, 3] per-slot RoPE coordinates
|
||||
cos_sin_cache : CUDA fp32 [capacity, rotary_dim] RoPE cache
|
||||
axis_map : CUDA int32 [rotary_dim // 2] position-axis per pair index
|
||||
weight : [head_dim] gemma norm weight (kernel applies 1 + w)
|
||||
write_locs : CUDA int32 [groups] compressed-cache slots to write
|
||||
compressed_k_buffer : CUDA [compressed_slots, head_dim] (written)
|
||||
compress_ratio : raw keys per compressed key
|
||||
rotary_dim : rotated prefix of each head row
|
||||
eps : RMSNorm epsilon
|
||||
is_neox_style : NeoX (True) or GPT-J (False) RoPE pairing
|
||||
"""
|
||||
head_dim = weight.shape[0]
|
||||
module = _jit_qsa_indexer_module(key_state_buffer.dtype, head_dim, is_neox_style)
|
||||
module.k_compress(
|
||||
key_state_buffer,
|
||||
group_locs,
|
||||
rope_position_buffer,
|
||||
cos_sin_cache,
|
||||
axis_map,
|
||||
weight,
|
||||
write_locs,
|
||||
compressed_k_buffer,
|
||||
compress_ratio,
|
||||
rotary_dim,
|
||||
eps,
|
||||
)
|
||||
@@ -171,6 +171,8 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
|
||||
a,
|
||||
mixed_qkvz,
|
||||
mixed_ba,
|
||||
qkvz_row_stride,
|
||||
ba_row_stride,
|
||||
NUM_HEADS_QK: tl.constexpr,
|
||||
NUM_HEADS_V: tl.constexpr,
|
||||
HEAD_QK: tl.constexpr,
|
||||
@@ -185,19 +187,19 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
|
||||
TOTAL_Q: tl.constexpr = NUM_HEADS_QK * HEAD_QK
|
||||
TOTAL_K: tl.constexpr = NUM_HEADS_QK * HEAD_QK
|
||||
TOTAL_V: tl.constexpr = NUM_HEADS_V * HEAD_V
|
||||
TOTAL_QKVZ: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V + TOTAL_V
|
||||
TOTAL_BA: tl.constexpr = NUM_HEADS_V * 2
|
||||
|
||||
# ── Output dimensions ──
|
||||
QKV_DIM_T: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V
|
||||
|
||||
# ── Read from contiguous input ──
|
||||
# q for head group i_qk: in the all_q region, offset i_qk * HEAD_QK
|
||||
blk_q_ptr = mixed_qkvz + i_bs * TOTAL_QKVZ + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
|
||||
blk_q_ptr = (
|
||||
mixed_qkvz + i_bs * qkvz_row_stride + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
|
||||
)
|
||||
# k for head group i_qk: in the all_k region
|
||||
blk_k_ptr = (
|
||||
mixed_qkvz
|
||||
+ i_bs * TOTAL_QKVZ
|
||||
+ i_bs * qkvz_row_stride
|
||||
+ TOTAL_Q
|
||||
+ i_qk * HEAD_QK
|
||||
+ tl.arange(0, HEAD_QK)
|
||||
@@ -209,7 +211,11 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
|
||||
# vector access. V_POW2 arrives as a wrapper-computed constexpr so the
|
||||
# dead branch is pruned before tl.arange validation.
|
||||
v_ld_base = (
|
||||
mixed_qkvz + i_bs * TOTAL_QKVZ + TOTAL_Q + TOTAL_K + i_qk * V_PER_GROUP * HEAD_V
|
||||
mixed_qkvz
|
||||
+ i_bs * qkvz_row_stride
|
||||
+ TOTAL_Q
|
||||
+ TOTAL_K
|
||||
+ i_qk * V_PER_GROUP * HEAD_V
|
||||
)
|
||||
z_ld_base = v_ld_base + TOTAL_V
|
||||
|
||||
@@ -250,12 +256,14 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
|
||||
|
||||
# ── b and a from contiguous [all_b | all_a] ──
|
||||
for i in tl.static_range(V_PER_GROUP):
|
||||
blk_b_ptr = mixed_ba + i_bs * TOTAL_BA + i_qk * V_PER_GROUP + i
|
||||
blk_b_ptr = mixed_ba + i_bs * ba_row_stride + i_qk * V_PER_GROUP + i
|
||||
blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
|
||||
tl.store(blk_b_st_ptr, tl.load(blk_b_ptr))
|
||||
|
||||
for i in tl.static_range(V_PER_GROUP):
|
||||
blk_a_ptr = mixed_ba + i_bs * TOTAL_BA + NUM_HEADS_V + i_qk * V_PER_GROUP + i
|
||||
blk_a_ptr = (
|
||||
mixed_ba + i_bs * ba_row_stride + NUM_HEADS_V + i_qk * V_PER_GROUP + i
|
||||
)
|
||||
blk_a_st_ptr = a + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
|
||||
tl.store(blk_a_st_ptr, tl.load(blk_a_ptr))
|
||||
|
||||
@@ -319,6 +327,8 @@ def fused_qkvzba_split_reshape_cat_contiguous(
|
||||
a,
|
||||
mixed_qkvz,
|
||||
mixed_ba,
|
||||
mixed_qkvz.stride(0),
|
||||
mixed_ba.stride(0),
|
||||
num_heads_qk,
|
||||
num_heads_v,
|
||||
head_qk,
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
_FAST_TOPK_SUPPORTED_K = (512, 2048)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_fast_topk_module(topk: int) -> Module:
|
||||
"""Compile and cache the JIT fast top-k module for a given k."""
|
||||
# Checks on the compile key live here, not in `fast_topk`: `cache_once`
|
||||
# keys on `topk`, so this runs once per specialisation.
|
||||
if topk not in _FAST_TOPK_SUPPORTED_K:
|
||||
raise RuntimeError(
|
||||
f"Unsupported topk {topk}. Supported: {_FAST_TOPK_SUPPORTED_K}"
|
||||
)
|
||||
args = make_cpp_args(topk, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
"fast_topk",
|
||||
*args,
|
||||
cuda_files=["elementwise/fast_topk.cuh"],
|
||||
cuda_wrappers=[("fast_topk", f"FastTopKKernel<{args}>::run")],
|
||||
)
|
||||
|
||||
|
||||
def fast_topk(
|
||||
score: torch.Tensor,
|
||||
lengths: torch.Tensor,
|
||||
topk: int,
|
||||
row_starts: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Per-row top-k selection over a fp32 score matrix.
|
||||
|
||||
Row b selects the `topk` largest values in
|
||||
``score[b, row_starts[b] : row_starts[b] + lengths[b]]`` and returns their
|
||||
indices relative to ``row_starts[b]``. Slots beyond ``lengths[b]`` are -1.
|
||||
Output order within a row is unspecified (atomic collection order).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
score : CUDA fp32 tensor [B, L]
|
||||
lengths : CUDA int32 tensor [B]
|
||||
topk : number of indices per row; 512 or 2048
|
||||
row_starts : optional CUDA int32 tensor [B]; defaults to zeros
|
||||
|
||||
Returns
|
||||
-------
|
||||
CUDA int32 tensor [B, topk]
|
||||
"""
|
||||
batch = score.shape[0]
|
||||
if row_starts is None:
|
||||
row_starts = torch.zeros(batch, dtype=torch.int32, device=score.device)
|
||||
indices = score.new_empty((batch, topk), dtype=torch.int32)
|
||||
|
||||
module = _jit_fast_topk_module(topk)
|
||||
module.fast_topk(score, row_starts, indices, lengths)
|
||||
return indices
|
||||
@@ -0,0 +1,135 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_hc_combine_module(
|
||||
hc_count: int, hidden_size: int, dtype: torch.dtype
|
||||
) -> Module:
|
||||
"""Compile and cache the JIT HC combine module for a given shape/dtype."""
|
||||
# Validation lives here rather than in `hc_combine`,
|
||||
# so `cache_once` runs it once per (hc_count, hidden_size, dtype), not per call.
|
||||
if dtype not in (torch.bfloat16, torch.float16):
|
||||
raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
|
||||
if hidden_size <= 0 or hidden_size % 8 != 0:
|
||||
raise RuntimeError(
|
||||
f"Unsupported hidden_size {hidden_size}. Must be a multiple of 8."
|
||||
)
|
||||
if hc_count <= 0 or (hc_count * hidden_size) % 2048 != 0:
|
||||
raise RuntimeError(
|
||||
f"Unsupported hc_count * hidden_size {hc_count * hidden_size}. "
|
||||
"Must be a multiple of 2048."
|
||||
)
|
||||
args = make_cpp_args(hc_count, hidden_size, is_arch_support_pdl(), dtype)
|
||||
return load_jit(
|
||||
"hc_combine",
|
||||
*args,
|
||||
cuda_files=["elementwise/hc_combine.cuh"],
|
||||
cuda_wrappers=[
|
||||
("hc_combine", f"HcCombineKernel<{args}>::run"),
|
||||
("hc_combine_split", f"HcCombineSplitKernel<{args}>::run"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def hc_combine(
|
||||
block_output: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
normed_residual: torch.Tensor,
|
||||
inject_weight: torch.Tensor,
|
||||
hc_count: int,
|
||||
hidden_size: int,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Fused HyperConnection (gated residual) combine.
|
||||
|
||||
a[m, c] = 2 * sigmoid(dot(normed_residual[m], inject_weight[c]) / hc_count)
|
||||
out[m, c*H + i] = residual[m, c*H + i] + a[m, c] * block_output[m, i]
|
||||
|
||||
Mirrors ``GatedResidual._combine_compute`` in
|
||||
``sglang.srt.layers.hyperconnection``. All math is accumulated in fp32.
|
||||
|
||||
Supported dtypes: torch.bfloat16, torch.float16.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
block_output : CUDA tensor [..., hidden_size]
|
||||
residual : CUDA tensor [..., hc_count * hidden_size]
|
||||
normed_residual : CUDA tensor, same shape/dtype as residual
|
||||
inject_weight : CUDA tensor [hc_count, hc_count * hidden_size]
|
||||
hc_count : number of hyper-connection branches
|
||||
hidden_size : per-branch hidden size
|
||||
out : optional pre-allocated output tensor (same shape/dtype as residual)
|
||||
|
||||
Returns
|
||||
-------
|
||||
Combined tensor, same shape/dtype as residual.
|
||||
"""
|
||||
y = block_output.reshape(-1, hidden_size)
|
||||
r = residual.reshape(-1, hc_count * hidden_size)
|
||||
n = normed_residual.reshape(-1, hc_count * hidden_size)
|
||||
if out is None:
|
||||
out = torch.empty_like(r)
|
||||
else:
|
||||
out = out.reshape(-1, hc_count * hidden_size)
|
||||
|
||||
module = _jit_hc_combine_module(hc_count, hidden_size, residual.dtype)
|
||||
module.hc_combine(y, r, n, inject_weight, out)
|
||||
return out.reshape(residual.shape)
|
||||
|
||||
|
||||
_SPLIT = 8
|
||||
_MAX_ROWS = 32
|
||||
_partials_cache = {}
|
||||
|
||||
|
||||
def _get_partials(hc_count: int, device: torch.device, rows: int) -> torch.Tensor:
|
||||
key = (device, hc_count)
|
||||
buf = _partials_cache.get(key)
|
||||
if buf is None or buf.shape[0] < rows:
|
||||
# The gate kernel writes one slot per (row, split, hc); a buffer shorter
|
||||
# than `rows` is written past its end rather than truncated.
|
||||
buf = torch.empty(
|
||||
(max(rows, _MAX_ROWS), _SPLIT, hc_count),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
_partials_cache[key] = buf
|
||||
return buf
|
||||
|
||||
|
||||
def hc_combine_split(
|
||||
block_output: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
normed_residual: torch.Tensor,
|
||||
inject_weight: torch.Tensor,
|
||||
hc_count: int,
|
||||
hidden_size: int,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
y = block_output.reshape(-1, hidden_size)
|
||||
r = residual.reshape(-1, hc_count * hidden_size)
|
||||
n = normed_residual.reshape(-1, hc_count * hidden_size)
|
||||
if out is None:
|
||||
out = torch.empty_like(r)
|
||||
else:
|
||||
out = out.reshape(-1, hc_count * hidden_size)
|
||||
rows = r.shape[0]
|
||||
partials = _get_partials(hc_count, r.device, rows)[:rows]
|
||||
module = _jit_hc_combine_module(hc_count, hidden_size, residual.dtype)
|
||||
module.hc_combine_split(y, r, n, inject_weight, out, partials)
|
||||
return out.reshape(residual.shape)
|
||||
@@ -0,0 +1,118 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
_MAX_ROWS = 32
|
||||
_KPAD_ALIGN = 128
|
||||
|
||||
|
||||
def pad_lowrank(lowrank: int) -> int:
|
||||
return (lowrank + _KPAD_ALIGN - 1) // _KPAD_ALIGN * _KPAD_ALIGN
|
||||
|
||||
|
||||
def hc_mix_shape_ok(hc_count: int, hidden_size: int, lowrank: int) -> bool:
|
||||
k = hc_count * hidden_size
|
||||
return (
|
||||
hidden_size % 8 == 0
|
||||
and lowrank > 0
|
||||
and lowrank % 8 == 0
|
||||
and k % _KPAD_ALIGN == 0
|
||||
and k % 128 == 0
|
||||
)
|
||||
|
||||
|
||||
def permute_pad_up_weight(w_up: torch.Tensor, hc_count: int) -> torch.Tensor:
|
||||
k, lowrank = w_up.shape
|
||||
hidden_size = k // hc_count
|
||||
padded = torch.zeros(
|
||||
(k, pad_lowrank(lowrank)), dtype=w_up.dtype, device=w_up.device
|
||||
)
|
||||
padded[:, :lowrank] = w_up.detach()
|
||||
src = torch.arange(k, device=w_up.device)
|
||||
perm = torch.empty(k, dtype=torch.long, device=w_up.device)
|
||||
perm[(src % hidden_size) * hc_count + (src // hidden_size)] = src
|
||||
return padded[perm].contiguous()
|
||||
|
||||
|
||||
_scratch_cache = {}
|
||||
_tactic_cache = {}
|
||||
|
||||
|
||||
def _get_scratch(lowrank: int, dtype: torch.dtype, device: torch.device):
|
||||
key = (device, lowrank, dtype)
|
||||
buf = _scratch_cache.get(key)
|
||||
if buf is None:
|
||||
buf = torch.zeros((_MAX_ROWS, pad_lowrank(lowrank)), dtype=dtype, device=device)
|
||||
_scratch_cache[key] = buf
|
||||
return buf
|
||||
|
||||
|
||||
def _get_tactic(m: int, n: int, k: int):
|
||||
key = (m, n, k)
|
||||
tactic = _tactic_cache.get(key)
|
||||
if tactic is None:
|
||||
from sglang.kernels.ops.gemm.dense_bf16_gemm_sm100_splitk_epilogue import (
|
||||
SplitKTactic,
|
||||
default_tactic,
|
||||
validate_tactic,
|
||||
)
|
||||
|
||||
if n <= 512 and k % (128 * 16) == 0:
|
||||
tactic = SplitKTactic(mma_m=64, mma_n=8, split_k=16, ab_stages=6)
|
||||
try:
|
||||
validate_tactic(tactic, m, n, k)
|
||||
except ValueError:
|
||||
tactic = default_tactic(m, n, k)
|
||||
else:
|
||||
tactic = default_tactic(m, n, k)
|
||||
_tactic_cache[key] = tactic
|
||||
return tactic
|
||||
|
||||
|
||||
def hc_mix(
|
||||
hyper_input_normed: torch.Tensor,
|
||||
w_down: torch.Tensor,
|
||||
w_up_permuted_padded: torch.Tensor,
|
||||
hc_count: int,
|
||||
hidden_size: int,
|
||||
) -> torch.Tensor:
|
||||
from sglang.kernels.ops.gemm.dense_bf16_gemm_sm100_splitk_epilogue import (
|
||||
run_splitk_dense_gate,
|
||||
run_splitk_dense_silu,
|
||||
)
|
||||
|
||||
rows, k = hyper_input_normed.shape
|
||||
lowrank = w_down.shape[0]
|
||||
kpad = w_up_permuted_padded.shape[1]
|
||||
out = torch.empty(
|
||||
(rows, hidden_size),
|
||||
dtype=hyper_input_normed.dtype,
|
||||
device=hyper_input_normed.device,
|
||||
)
|
||||
if rows == 0:
|
||||
return out
|
||||
inv_hc = 1.0 / hc_count
|
||||
t_pad = _get_scratch(lowrank, hyper_input_normed.dtype, hyper_input_normed.device)
|
||||
for row_start in range(0, rows, _MAX_ROWS):
|
||||
row_end = min(row_start + _MAX_ROWS, rows)
|
||||
m = row_end - row_start
|
||||
chunk = hyper_input_normed[row_start:row_end]
|
||||
run_splitk_dense_silu(
|
||||
chunk,
|
||||
w_down.T,
|
||||
t_pad[:m, :lowrank],
|
||||
True,
|
||||
_get_tactic(m, lowrank, k),
|
||||
inv_hc,
|
||||
)
|
||||
run_splitk_dense_gate(
|
||||
t_pad[:m],
|
||||
w_up_permuted_padded.T,
|
||||
chunk,
|
||||
out[row_start:row_end],
|
||||
True,
|
||||
_get_tactic(m, k, kpad),
|
||||
inv_hc,
|
||||
hc_count,
|
||||
)
|
||||
return out
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,77 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_grouped_gemma_rmsnorm_module(group_size: int, dtype: torch.dtype) -> Module:
|
||||
"""Compile and cache the JIT grouped Gemma RMSNorm module."""
|
||||
# Validation sits under `cache_once`: once per (group_size, dtype) specialisation,
|
||||
# not once per call.
|
||||
if dtype not in (torch.bfloat16, torch.float16):
|
||||
raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
|
||||
if group_size <= 0 or group_size % 512 != 0:
|
||||
raise RuntimeError(
|
||||
f"Unsupported group_size {group_size}. Must be a multiple of 512."
|
||||
)
|
||||
args = make_cpp_args(group_size, is_arch_support_pdl(), dtype)
|
||||
return load_jit(
|
||||
"grouped_gemma_rmsnorm",
|
||||
*args,
|
||||
cuda_files=["elementwise/grouped_gemma_rmsnorm.cuh"],
|
||||
cuda_wrappers=[
|
||||
("grouped_gemma_rmsnorm", f"GroupedGemmaRMSNormKernel<{args}>::run")
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def grouped_gemma_rmsnorm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
group_size: int,
|
||||
eps: float = 1e-6,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Grouped Gemma-style RMSNorm: out = x * rsqrt(mean(x^2) + eps) * (1 + weight).
|
||||
|
||||
The last dimension is split into groups of `group_size` elements; variance
|
||||
is computed per group. With group_size == input.size(-1) this reduces to a
|
||||
plain Gemma RMSNorm.
|
||||
|
||||
Supported dtypes: torch.bfloat16, torch.float16.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input : CUDA tensor [..., hidden_size], hidden_size % group_size == 0
|
||||
weight : CUDA tensor [hidden_size]
|
||||
group_size : elements per variance group (multiple of 512)
|
||||
eps : RMSNorm epsilon
|
||||
out : optional pre-allocated output tensor (same shape/dtype as input)
|
||||
|
||||
Returns
|
||||
-------
|
||||
Normalized tensor, same shape/dtype as input.
|
||||
"""
|
||||
hidden_size = input.size(-1)
|
||||
x = input.reshape(-1, hidden_size)
|
||||
if out is None:
|
||||
out = torch.empty_like(x)
|
||||
else:
|
||||
out = out.reshape(-1, hidden_size)
|
||||
|
||||
module = _jit_grouped_gemma_rmsnorm_module(group_size, input.dtype)
|
||||
module.grouped_gemma_rmsnorm(x, weight, out, eps)
|
||||
return out.reshape(input.shape)
|
||||
@@ -740,13 +740,15 @@ def _fused_commit_track_indices_kernel(
|
||||
last_correct_out_ptr,
|
||||
track_steps_out_ptr,
|
||||
dtn,
|
||||
accept_stride,
|
||||
interval,
|
||||
HAS_TRACK: tl.constexpr,
|
||||
):
|
||||
b = tl.program_id(0).to(tl.int64)
|
||||
al = tl.load(accept_lens_ptr + b).to(tl.int64)
|
||||
row = b * accept_stride
|
||||
base = b * dtn
|
||||
last = tl.load(accept_index_ptr + base + al - 1).to(tl.int64) - base
|
||||
last = tl.load(accept_index_ptr + row + al - 1).to(tl.int64) - base
|
||||
tl.store(last_correct_out_ptr + b, last)
|
||||
if HAS_TRACK:
|
||||
pre = tl.load(seq_lens_ptr + b).to(tl.int64)
|
||||
@@ -755,7 +757,7 @@ def _fused_commit_track_indices_kernel(
|
||||
tp = (post // interval) * interval
|
||||
ti = tp - pre - 1
|
||||
ti = tl.where(ti < 0, 0, ti)
|
||||
cand = tl.load(accept_index_ptr + base + ti).to(tl.int64) - base
|
||||
cand = tl.load(accept_index_ptr + row + ti).to(tl.int64) - base
|
||||
tl.store(track_steps_out_ptr + b, tl.where(cross, cand, -1))
|
||||
|
||||
|
||||
@@ -766,8 +768,9 @@ def fused_commit_track_indices(
|
||||
draft_token_num: int,
|
||||
mamba_track_interval: int,
|
||||
):
|
||||
"""Single-launch replacement for the eager index math in
|
||||
``commit_mamba_states_after_verify`` (index ranges, gathers, floordiv chain)."""
|
||||
"""Single-launch replacement for the eager verify-commit index math;
|
||||
accept_index is [bs, tree_depth] but its values index bs * draft_token_num rows."""
|
||||
accept_index = accept_index.contiguous()
|
||||
bs = accept_lens.shape[0]
|
||||
last_correct_step_indices = torch.empty(
|
||||
bs, dtype=torch.int64, device=accept_lens.device
|
||||
@@ -785,6 +788,7 @@ def fused_commit_track_indices(
|
||||
last_correct_step_indices,
|
||||
mamba_steps_to_track,
|
||||
draft_token_num,
|
||||
accept_index.shape[1],
|
||||
mamba_track_interval,
|
||||
HAS_TRACK=has_track,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,309 @@
|
||||
"""Bitwise-exact fused kernels for decode-sized Qwen4 PLE paths."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
_QWEN4_NGRAM_SIZE = 3
|
||||
_QWEN4_HEADS_PER_NGRAM = 8
|
||||
_QWEN4_NGRAM_HEADS = 16
|
||||
_QWEN4_HC_COUNT = 4
|
||||
_QWEN4_HIDDEN_SIZE = 2560
|
||||
_QWEN4_MAX_SHORT_CONV_STATE_LEN = 16
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _round_bf16_to_fp32(value):
|
||||
"""RNE-round fp32 to BF16 precision while retaining an fp32 register."""
|
||||
|
||||
bits = value.to(tl.int32, bitcast=True)
|
||||
rounding_bias = 0x7FFF + ((bits >> 16) & 1)
|
||||
rounded_bits = (bits + rounding_bias) & -65536
|
||||
return rounded_bits.to(tl.float32, bitcast=True)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _qwen4_ngram_hash_kernel(
|
||||
contexts_ptr,
|
||||
multipliers_ptr,
|
||||
vocab_sizes_ptr,
|
||||
offsets_ptr,
|
||||
output_ptr,
|
||||
num_outputs,
|
||||
eos_token_id,
|
||||
NGRAM_SIZE: tl.constexpr,
|
||||
HEADS_PER_NGRAM: tl.constexpr,
|
||||
NGRAM_HEADS: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
output_idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = output_idx < num_outputs
|
||||
token_idx = output_idx // NGRAM_HEADS
|
||||
head_idx = output_idx % NGRAM_HEADS
|
||||
context_base = token_idx * NGRAM_SIZE
|
||||
|
||||
token_0 = tl.load(contexts_ptr + context_base, mask=mask, other=0)
|
||||
token_1 = tl.load(contexts_ptr + context_base + 1, mask=mask, other=0)
|
||||
token_2 = tl.load(contexts_ptr + context_base + 2, mask=mask, other=0)
|
||||
multiplier_0 = tl.load(multipliers_ptr)
|
||||
multiplier_1 = tl.load(multipliers_ptr + 1)
|
||||
multiplier_2 = tl.load(multipliers_ptr + 2)
|
||||
|
||||
# Only the final position of each three-token context is materialized.
|
||||
previous_1 = tl.where(token_1 == eos_token_id, eos_token_id, token_1)
|
||||
previous_2 = tl.where(
|
||||
(token_0 == eos_token_id) | (token_1 == eos_token_id),
|
||||
eos_token_id,
|
||||
token_0,
|
||||
)
|
||||
mixed = (token_2 * multiplier_0) ^ (previous_1 * multiplier_1)
|
||||
mixed_3 = mixed ^ (previous_2 * multiplier_2)
|
||||
mixed = tl.where(head_idx < HEADS_PER_NGRAM, mixed, mixed_3)
|
||||
|
||||
vocab_size = tl.load(vocab_sizes_ptr + head_idx, mask=mask, other=1)
|
||||
offset = tl.load(offsets_ptr + head_idx, mask=mask, other=0)
|
||||
tl.store(output_ptr + output_idx, mixed % vocab_size + offset, mask=mask)
|
||||
|
||||
|
||||
def can_fuse_qwen4_ngram_hash(
|
||||
contexts: torch.Tensor,
|
||||
multipliers: torch.Tensor,
|
||||
vocab_sizes: torch.Tensor,
|
||||
offsets: torch.Tensor,
|
||||
) -> bool:
|
||||
"""Return whether inputs match the fixed Qwen4 PLE hash contract."""
|
||||
|
||||
return (
|
||||
contexts.is_cuda
|
||||
and contexts.dtype == torch.long
|
||||
and contexts.dim() == 2
|
||||
and contexts.shape[1] == _QWEN4_NGRAM_SIZE
|
||||
and contexts.is_contiguous()
|
||||
and multipliers.is_cuda
|
||||
and multipliers.dtype == torch.long
|
||||
and multipliers.numel() == _QWEN4_NGRAM_SIZE
|
||||
and vocab_sizes.is_cuda
|
||||
and vocab_sizes.dtype == torch.long
|
||||
and vocab_sizes.numel() == _QWEN4_NGRAM_HEADS
|
||||
and offsets.is_cuda
|
||||
and offsets.dtype == torch.long
|
||||
and offsets.numel() == _QWEN4_NGRAM_HEADS
|
||||
)
|
||||
|
||||
|
||||
def fused_qwen4_ngram_hash(
|
||||
contexts: torch.Tensor,
|
||||
multipliers: torch.Tensor,
|
||||
vocab_sizes: torch.Tensor,
|
||||
offsets: torch.Tensor,
|
||||
eos_token_id: int,
|
||||
) -> torch.Tensor:
|
||||
"""Return the 16 Qwen4 PLE N-gram IDs in one kernel launch."""
|
||||
|
||||
if not can_fuse_qwen4_ngram_hash(contexts, multipliers, vocab_sizes, offsets):
|
||||
raise ValueError("unsupported input for fused Qwen4 PLE N-gram hash")
|
||||
output = torch.empty(
|
||||
(contexts.shape[0], _QWEN4_NGRAM_HEADS),
|
||||
dtype=torch.long,
|
||||
device=contexts.device,
|
||||
)
|
||||
num_outputs = output.numel()
|
||||
if num_outputs:
|
||||
block_size = 256
|
||||
_qwen4_ngram_hash_kernel[(triton.cdiv(num_outputs, block_size),)](
|
||||
contexts,
|
||||
multipliers,
|
||||
vocab_sizes,
|
||||
offsets,
|
||||
output,
|
||||
num_outputs,
|
||||
eos_token_id,
|
||||
NGRAM_SIZE=_QWEN4_NGRAM_SIZE,
|
||||
HEADS_PER_NGRAM=_QWEN4_HEADS_PER_NGRAM,
|
||||
NGRAM_HEADS=_QWEN4_NGRAM_HEADS,
|
||||
BLOCK_SIZE=block_size,
|
||||
num_warps=4,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _qwen4_gate_value_kernel(
|
||||
gate_ptr,
|
||||
value_ptr,
|
||||
output_ptr,
|
||||
num_tokens,
|
||||
HC_COUNT: tl.constexpr,
|
||||
HIDDEN_SIZE: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
token_group = tl.program_id(0)
|
||||
token = token_group // HC_COUNT
|
||||
hidden = tl.arange(0, BLOCK_SIZE)
|
||||
mask = (token < num_tokens) & (hidden < HIDDEN_SIZE)
|
||||
|
||||
# `gate` arrives already rounded to bf16 by the eager multiply/reduce/divide;
|
||||
# every remaining eager bf16 rounding boundary is reproduced below.
|
||||
gate = tl.load(gate_ptr + token_group).to(tl.float32)
|
||||
magnitude = tl.maximum(tl.abs(gate), 1.0e-6)
|
||||
root = _round_bf16_to_fp32(tl.sqrt(magnitude))
|
||||
sign = tl.where(gate > 0.0, 1.0, tl.where(gate < 0.0, -1.0, 0.0))
|
||||
transformed = _round_bf16_to_fp32(root * sign)
|
||||
activated = _round_bf16_to_fp32(tl.sigmoid(transformed))
|
||||
|
||||
value = tl.load(value_ptr + token * HIDDEN_SIZE + hidden, mask=mask, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
output_offset = token_group * HIDDEN_SIZE + hidden
|
||||
tl.store(output_ptr + output_offset, activated * value, mask=mask)
|
||||
|
||||
|
||||
def can_fuse_qwen4_gate_value(gate: torch.Tensor, value: torch.Tensor) -> bool:
|
||||
"""Return whether inputs match Qwen4's fixed BF16 gate/value contract."""
|
||||
|
||||
return (
|
||||
gate.is_cuda
|
||||
and gate.dtype == torch.bfloat16
|
||||
and gate.dim() == 3
|
||||
and gate.shape[1:] == (_QWEN4_HC_COUNT, 1)
|
||||
and gate.is_contiguous()
|
||||
and value.is_cuda
|
||||
and value.dtype == gate.dtype
|
||||
and value.shape == (gate.shape[0], _QWEN4_HIDDEN_SIZE)
|
||||
and value.is_contiguous()
|
||||
)
|
||||
|
||||
|
||||
def fused_qwen4_gate_value(gate: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply Qwen4's post-reduction gate and value broadcast in one kernel."""
|
||||
|
||||
if not can_fuse_qwen4_gate_value(gate, value):
|
||||
raise ValueError("unsupported input for fused Qwen4 PLE gate/value")
|
||||
output = torch.empty(
|
||||
(gate.shape[0], _QWEN4_HC_COUNT, _QWEN4_HIDDEN_SIZE),
|
||||
dtype=value.dtype,
|
||||
device=value.device,
|
||||
)
|
||||
if gate.shape[0]:
|
||||
_qwen4_gate_value_kernel[(gate.shape[0] * _QWEN4_HC_COUNT,)](
|
||||
gate,
|
||||
value,
|
||||
output,
|
||||
gate.shape[0],
|
||||
HC_COUNT=_QWEN4_HC_COUNT,
|
||||
HIDDEN_SIZE=_QWEN4_HIDDEN_SIZE,
|
||||
BLOCK_SIZE=4096,
|
||||
num_warps=8,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _qwen4_short_conv_state_kernel(
|
||||
state_ptr,
|
||||
state_indices_ptr,
|
||||
x_ptr,
|
||||
conv_input_ptr,
|
||||
num_tokens,
|
||||
CHANNELS: tl.constexpr,
|
||||
STATE_LEN: tl.constexpr,
|
||||
BLOCK_CHANNELS: tl.constexpr,
|
||||
BLOCK_STATE_LEN: tl.constexpr,
|
||||
):
|
||||
token = tl.program_id(0)
|
||||
channel = tl.program_id(1) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS)[:, None]
|
||||
state_col = tl.arange(0, BLOCK_STATE_LEN)[None, :]
|
||||
channel_mask = (token < num_tokens) & (channel < CHANNELS)
|
||||
state_mask = channel_mask & (state_col < STATE_LEN)
|
||||
state_index = tl.load(state_indices_ptr + token, mask=token < num_tokens, other=0)
|
||||
state_base = state_index * CHANNELS * STATE_LEN
|
||||
state_offset = state_base + channel * STATE_LEN + state_col
|
||||
output_base = token * CHANNELS * (STATE_LEN + 1)
|
||||
output_offset = output_base + channel * (STATE_LEN + 1) + state_col
|
||||
|
||||
# Materialize every old state value before advancing the in-place cache,
|
||||
# so the convolution input equals the native index_select + cat result.
|
||||
old_state = tl.load(state_ptr + state_offset, mask=state_mask, other=0.0)
|
||||
tl.store(conv_input_ptr + output_offset, old_state, mask=state_mask)
|
||||
x = tl.load(x_ptr + token * CHANNELS + channel, mask=channel_mask, other=0.0)
|
||||
tl.store(
|
||||
conv_input_ptr + output_base + channel * (STATE_LEN + 1) + STATE_LEN,
|
||||
x,
|
||||
mask=channel_mask,
|
||||
)
|
||||
tl.debug_barrier()
|
||||
|
||||
# Slot 0 is the CUDA-graph padding slot and may appear in several rows;
|
||||
# its post-step value is unobservable, so skip it to avoid duplicate writers.
|
||||
update_mask = state_mask & (state_col < STATE_LEN - 1) & (state_index != 0)
|
||||
next_value = tl.load(
|
||||
conv_input_ptr + output_offset + 1, mask=update_mask, other=0.0
|
||||
)
|
||||
tl.store(state_ptr + state_offset, next_value, mask=update_mask)
|
||||
tl.store(
|
||||
state_ptr + state_base + channel * STATE_LEN + STATE_LEN - 1,
|
||||
x,
|
||||
mask=channel_mask & (state_index != 0),
|
||||
)
|
||||
|
||||
|
||||
def can_fuse_qwen4_short_conv_state(
|
||||
state: torch.Tensor,
|
||||
state_indices: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
) -> bool:
|
||||
"""Return whether decode state movement can use the exact fused kernel."""
|
||||
|
||||
return (
|
||||
state.is_cuda
|
||||
and state.dtype in (torch.bfloat16, torch.float16)
|
||||
and state.dim() == 3
|
||||
and state.is_contiguous()
|
||||
and 0 < state.shape[2] <= _QWEN4_MAX_SHORT_CONV_STATE_LEN
|
||||
and state_indices.is_cuda
|
||||
and state_indices.dtype == torch.long
|
||||
and state_indices.dim() == 1
|
||||
and state_indices.is_contiguous()
|
||||
and x.is_cuda
|
||||
and x.dtype == state.dtype
|
||||
and x.dim() == 2
|
||||
and x.is_contiguous()
|
||||
and x.shape == (state_indices.shape[0], state.shape[1])
|
||||
)
|
||||
|
||||
|
||||
def fused_qwen4_short_conv_state(
|
||||
state: torch.Tensor,
|
||||
state_indices: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Build ``[selected state, x]`` and advance real decode slots in one launch."""
|
||||
|
||||
if not can_fuse_qwen4_short_conv_state(state, state_indices, x):
|
||||
raise ValueError("unsupported input for fused Qwen4 short-conv state")
|
||||
state_len = state.shape[2]
|
||||
conv_input = torch.empty(
|
||||
(x.shape[0], x.shape[1], state_len + 1),
|
||||
dtype=x.dtype,
|
||||
device=x.device,
|
||||
)
|
||||
if x.shape[0]:
|
||||
block_channels = 128
|
||||
block_state_len = triton.next_power_of_2(state_len)
|
||||
_qwen4_short_conv_state_kernel[
|
||||
(x.shape[0], triton.cdiv(x.shape[1], block_channels))
|
||||
](
|
||||
state,
|
||||
state_indices,
|
||||
x,
|
||||
conv_input,
|
||||
x.shape[0],
|
||||
CHANNELS=state.shape[1],
|
||||
STATE_LEN=state_len,
|
||||
BLOCK_CHANNELS=block_channels,
|
||||
BLOCK_STATE_LEN=block_state_len,
|
||||
num_warps=8,
|
||||
)
|
||||
return conv_input
|
||||
Reference in New Issue
Block a user