hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Support DeepSeek V4 Flash 0731

Browse Source
This commit is contained in:
Georg Bauer
2026-08-29 20:28:50 +02:00
parent ad855b321e
commit f1c177b754
23 changed files with 10510 additions and 2324 deletions
Download Patch File Download Diff File Expand all files Collapse all files

285
metal/flash_attn.metal
View File

@@ -1398,24 +1398,26 @@ constant int32_t FC_flash_attn_ext_vec_reduce_NWG [[function_constant(FC_FLASH_A
// Reduces split-K decode FlashAttention partials. It combines each workgroup's
// output vector and softmax (sum,max) pair into the final attention result.
kernel void kernel_flash_attn_ext_vec_reduce(
/* Shared and deliberately noinline so the split-K reduction is compiled once and
* every caller gets identical codegen. The RoPE-fused sibling in dsv4_rope.metal
* calls this same body, which is what keeps the fusion bit-exact. */
static __attribute__((noinline)) void ds4_flash_attn_vec_reduce_row(
constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
device const char * htmp,
device char * dst,
uint tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
#define NWG (FC_flash_attn_ext_vec_reduce_NWG)
#define DV (FC_flash_attn_ext_vec_reduce_DV)
uint tgpig,
ushort tiisg,
ushort sgitg,
short NWG_,
short DV_) {
const uint64_t rid = tgpig;
const short iwg = tiisg;
device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV*NWG;
device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV_*NWG_;
float S = ss[rid*(2*NWG) + 2*iwg + 0];
float M = ss[rid*(2*NWG) + 2*iwg + 1];
float S = ss[rid*(2*NWG_) + 2*iwg + 0];
float M = ss[rid*(2*NWG_) + 2*iwg + 1];
const float m = simd_max(M);
const float ms = exp(M - m);
@@ -1423,19 +1425,268 @@ kernel void kernel_flash_attn_ext_vec_reduce(
S = simd_sum(S*ms);
S = S == 0.0f ? 0.0f : 1.0f/S;
const short DV4 = DV/4;
const short DV4 = DV_/4;
device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG;
device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG_;
device float4 * dst4 = (device float4 *) dst + rid*DV4;
for (short i = sgitg; i < DV4; i += NWG) {
const float4 v = simd_sum(htmp4[i*NWG + iwg]*ms);
for (short i = sgitg; i < DV4; i += NWG_) {
const float4 v = simd_sum(htmp4[i*NWG_ + iwg]*ms);
if (iwg == 0) {
dst4[i] = v*S;
}
}
#undef NWG
#undef DV
}
kernel void kernel_flash_attn_ext_vec_reduce(
constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
device const char * htmp,
device char * dst,
uint tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
ds4_flash_attn_vec_reduce_row(args, htmp, dst, tgpig, tiisg, sgitg,
(short)FC_flash_attn_ext_vec_reduce_NWG,
(short)FC_flash_attn_ext_vec_reduce_DV);
}
// M5 decode specialization: time-slice all 32 split-K workgroups through eight
// physical simdgroups, then reduce through the same 32-lane topology without a
// device partial buffer. The host gate fixes the exact F16 512-wide geometry.
static inline void ds4_flash_attn_vec_packed8_reduce_f16_512(
constant ds4_metal_args_flash_attn_ext_vec & args,
device const char * q,
device const char * k,
device const char * v,
device const char * mask,
device const char * sinks,
device const char * pad,
device char * dst,
threadgroup char * shmem,
uint head,
ushort tiisg,
ushort sgitg) {
constexpr short NW = 32;
constexpr short C = 32;
constexpr short NSG = 8;
constexpr short NWG = 32;
constexpr short DK4 = 128;
constexpr short DV4 = 128;
constexpr short SH = 128;
/* 24,448 dynamic bytes: shared Q, eight score/mask banks, all 32
* split-local weights and stats, sink scales, and a padded 32x33 F32
* float4 partial plane. */
threadgroup half4 *q_shared = (threadgroup half4 *)shmem;
threadgroup half *score_banks =
(threadgroup half *)(q_shared + DK4);
threadgroup volatile float *weights =
(threadgroup volatile float *)(score_banks + NSG * SH);
threadgroup volatile float *stats = weights + NWG * C;
threadgroup volatile float *sink_scale = stats + 2 * NWG;
threadgroup volatile float4 *partial_plane =
(threadgroup volatile float4 *)(sink_scale + NWG);
const short lane = (short)tiisg;
threadgroup half *bank = score_banks + (short)sgitg * SH;
threadgroup float *ss = (threadgroup float *)bank;
threadgroup half *sm = bank + 2 * C;
device const float4 *q4 =
(device const float4 *)(q + (uint64_t)head * args.nb02);
if (sgitg == 0) {
for (short i = lane; i < DK4; i += NW) {
q_shared[i] = (half4)q4[i];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
/* Eight physical simdgroups time-slice the exact 32 legacy split-K
* workgroups. The official <=1024-key gate gives each virtual split at
* most one 32-row block, so its value partial can be formed later from
* these materialized weights without changing online-softmax order. */
for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) {
float S = 0.0f;
float M = -FLT_MAX / 2;
float out_scale = 1.0f;
const int ic_original = (int)iwg * C;
weights[(uint)iwg * C + (uint)lane] = 0.0f;
ss[lane] = 0.0f;
sm[lane] = (half)0.0h;
simdgroup_barrier(mem_flags::mem_threadgroup);
if (ic_original < args.ne11) {
device const char *k_block = k;
device const char *v_block = v;
device const half *pm = (device const half *)mask;
int ic = ic_original;
if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
k_block = pad;
const uint64_t k_pad_bytes =
args.nb11 * (uint64_t)C *
(uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
const uint64_t v_pad_bytes =
args.nb21 * (uint64_t)C *
(uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
if (FC_flash_attn_ext_vec_shared_kvpad) {
v_block = k_block;
pm = (device const half *)(k_block +
k_pad_bytes + v_pad_bytes);
} else {
v_block = k_block + k_pad_bytes;
pm = (device const half *)(v_block + v_pad_bytes);
}
ic = 0;
}
sm[lane] = pm[ic + lane];
if (simd_max(sm[lane]) > -MAXHALF) {
device const half4 *pk4 =
(device const half4 *)(k_block +
(uint64_t)ic * args.nb11);
threadgroup const half4 *pq4 = q_shared;
pk4 += lane;
pq4 += lane;
float lane_mqk = 0.0f;
FOR_UNROLL (short cc = 0; cc < C; ++cc) {
float mqk = 0.0f;
FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) {
mqk += dot((float4)pk4[cc * DK4 + ii * NW],
(float4)pq4[ii * NW]);
}
mqk = simd_sum(mqk);
if (lane == cc) {
lane_mqk = mqk;
}
}
ss[lane] = fma(lane_mqk, args.scale,
(float)sm[lane]);
simdgroup_barrier(mem_flags::mem_threadgroup);
const float old_m = M;
const float score = ss[lane];
M = simd_max(max(M, score));
const float ms = exp(old_m - M);
const float vs = exp(score - M);
S = S * ms + simd_sum(vs);
ss[lane] = vs;
simdgroup_barrier(mem_flags::mem_threadgroup);
weights[(uint)iwg * C + (uint)lane] = ss[lane];
}
if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) {
const float old_m = M;
const float sink = lane == 0
? ((device const float *)sinks)[head]
: -FLT_MAX / 2;
M = simd_max(max(M, sink));
const float ms = exp(old_m - M);
const float vs = exp(sink - M);
S = S * ms + simd_sum(vs);
out_scale = ms;
}
} else if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) {
const float old_m = M;
const float sink = lane == 0
? ((device const float *)sinks)[head]
: -FLT_MAX / 2;
M = simd_max(max(M, sink));
const float ms = exp(old_m - M);
const float vs = exp(sink - M);
S = S * ms + simd_sum(vs);
out_scale = ms;
}
if (lane == 0) {
stats[2 * (uint)iwg + 0] = S;
stats[2 * (uint)iwg + 1] = M;
sink_scale[(uint)iwg] = out_scale;
}
simdgroup_barrier(mem_flags::mem_threadgroup);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
/* Recreate the legacy 32-lane reducer exactly: lane is the virtual split
* index, including neutral idle splits in their original tree positions. */
const short split = lane;
float reduce_S = stats[2 * (uint)split + 0];
float reduce_M = stats[2 * (uint)split + 1];
const float reduce_max = simd_max(reduce_M);
const float reduce_ms = exp(reduce_M - reduce_max);
reduce_S = simd_sum(reduce_S * reduce_ms);
const float reduce_inv =
reduce_S == 0.0f ? 0.0f : 1.0f / reduce_S;
device float4 *dst4 =
(device float4 *)(dst +
(uint64_t)head * 512u * sizeof(float));
/* Form one 32-float4 output quadrant at a time. During production each
* physical simdgroup time-slices four virtual splits while SIMD lanes are
* contiguous output columns, exactly matching the legacy V loads and
* cc-major accumulation. A padded 33-column plane avoids a 32-way TG-bank
* conflict when the reducer transposes lanes back to virtual splits. */
for (short quadrant = 0; quadrant < 4; ++quadrant) {
for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) {
float4 lo = float4(0.0f);
const int ic_original = (int)iwg * C;
if (ic_original < args.ne11) {
device const char *v_block = v;
int ic = ic_original;
if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
device const char *k_block = pad;
const uint64_t k_pad_bytes =
args.nb11 * (uint64_t)C *
(uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
if (FC_flash_attn_ext_vec_shared_kvpad) {
v_block = k_block;
} else {
v_block = k_block + k_pad_bytes;
}
ic = 0;
}
device const half4 *pv4 =
(device const half4 *)(v_block +
(uint64_t)ic * args.nb21);
threadgroup volatile float *split_weights =
weights + (uint)iwg * C;
const short oc = quadrant * NW + lane;
FOR_UNROLL (short cc = 0; cc < C; ++cc) {
lo += float4(pv4[cc * DV4 + oc]) *
float4(split_weights[cc]);
}
float4 acc = float4(0.0f);
acc += lo;
if (iwg == 0) {
acc *= sink_scale[0];
}
lo = acc;
}
partial_plane[(uint)iwg * 33u + (uint)lane] = lo;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
/* lane is now the legacy split index and each physical simdgroup
* reduces four output columns through the identical simd_sum tree. */
for (short out_lane = (short)sgitg; out_lane < NW; out_lane += NSG) {
const float4 materialized =
(float4)partial_plane[(uint)lane * 33u + (uint)out_lane];
const float4 reduced = simd_sum(materialized * reduce_ms);
if (lane == 0) {
dst4[quadrant * NW + out_lane] = reduced * reduce_inv;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
}