gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Qwen3.6 MoE inference support

Browse Source
This commit is contained in:
Gahow Wang
2026-07-13 20:24:41 +08:00
parent 588bfd9df3
commit a2de146fb6
27 changed files with 3153 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
csrc/attention/flash_attention.cu
View File

@@ -13,13 +13,13 @@
// O [batch, num_q_heads, q_len, head_dim]
//
// Shared memory (BF16):
// smem_q[BR][head_dim] — 64 * 128 * 2 = 16 KB (loaded once per Q tile)
// smem_kv[BC][head_dim] — 64 * 128 * 2 = 16 KB (alternates K and V)
// Total: 32 KB (fits in default 48 KB shared memory)
// Use 32-row tiles so head_dim=256 still fits in the default 48 KB shared
// memory limit: 2 * 32 * 256 * 2 = 32 KB.
#define BR 64
#define BC 64
#define BR 32
#define BC 32
#define THREADS_PER_BLOCK 128
#define FLASH_HEAD_DIM_MAX 256
__global__ void flash_attention_bf16_kernel(
const __nv_bfloat16* __restrict__ Q,
@@ -73,8 +73,8 @@ __global__ void flash_attention_bf16_kernel(
// Thread t (0 <= t < q_tile_rows) owns Q row t
bool owns_row = (tid < q_tile_rows);
// Per-thread FP32 accumulators (head_dim up to 128)
float O_acc[128];
// Per-thread FP32 accumulators (head_dim up to 256)
float O_acc[FLASH_HEAD_DIM_MAX];
float m_val = -INFINITY;
float l_val = 0.0f;
if (owns_row) {
@@ -250,7 +250,7 @@ __global__ void flash_attention_sinks_bf16_kernel(
bool owns_row = (tid < q_tile_rows);
float O_acc[128];
float O_acc[FLASH_HEAD_DIM_MAX];
float m_val = -INFINITY;
float l_val = 0.0f;
if (owns_row) {
@@ -384,7 +384,7 @@ __global__ void flash_attention_sinks_bf16_kernel(
// ============================================================
#define DECODE_THREADS 256
#define HEAD_DIM_MAX 128
#define HEAD_DIM_MAX 256
__global__ void decode_attention_bf16_kernel(
const __nv_bfloat16* __restrict__ Q,
@@ -413,7 +413,7 @@ __global__ void decode_attention_bf16_kernel(
const __nv_bfloat16* V_base = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
__nv_bfloat16* O_ptr = O + ((long long)batch_idx * num_q_heads + q_head) * head_dim;
// Load Q vector into registers (head_dim <= 128)
// Load Q vector into registers (head_dim <= 256)
float q_reg[HEAD_DIM_MAX];
for (int d = 0; d < head_dim; d++) {
q_reg[d] = __bfloat162float(Q_ptr[d]);