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

538
metal/dsv4_rope.metal
View File

@@ -41,6 +41,23 @@ struct ds4_metal_args_dsv4_rope_affine_pair {
float beta_slow;
};
struct ds4_metal_args_dsv4_head_norm_rope {
int32_t n_head;
int32_t head_dim;
int32_t head_dim4;
int32_t n_dims;
int32_t n_ctx_orig;
int32_t pos0;
int32_t inverse;
float eps;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
};
static float rope_yarn_ramp(const float low, const float high, const int i0) {
const float y = (i0 / 2 - low) / max(0.001f, high - low);
return 1.0f - min(1.0f, max(0.0f, y));
@@ -327,36 +344,109 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4(
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
}
// DS4 positions are always affine within one RoPE dispatch. This variant
// reconstructs the same wrapped int32 position in-kernel, avoiding the host
// position array and its buffer binding while preserving the pair lane mapping
// and all floating-point operations of the specialization above.
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
device const char * src0 [[buffer(1)]],
device char * dst [[buffer(4)]],
uint tid [[thread_index_in_threadgroup]],
ushort3 ntg [[threads_per_threadgroup]],
uint3 tgpig [[threadgroup_position_in_grid]]) {
const int i1 = tgpig[0];
const int i2 = tgpig[1];
// Fuses the per-head RMSNorm and partial Q RoPE while retaining the standalone
// norm reduction tree and the mode-0 RoPE lane mapping.
kernel void kernel_dsv4_head_rms_norm_rope_tail_f32(
constant ds4_metal_args_dsv4_head_norm_rope & args,
device char * xraw,
threadgroup float * shmem_f32 [[threadgroup(0)]],
uint3 tgpig [[threadgroup_position_in_grid]],
ushort3 tpitg [[thread_position_in_threadgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort3 ntg [[threads_per_threadgroup]]) {
if (sgitg == 0) {
shmem_f32[tiisg] = 0.0f;
}
const uint head = tgpig.x;
const uint tok = tgpig.y;
device float4 * x4 = (device float4 *)xraw +
((uint64_t)tok * (uint64_t)args.n_head + head) *
(uint64_t)args.head_dim4;
float sumf = 0.0f;
for (int i00 = tpitg.x; i00 < args.head_dim4; i00 += ntg.x) {
sumf += dot(x4[i00], x4[i00]);
}
sumf = simd_sum(sumf);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tiisg == 0) {
shmem_f32[sgitg] = sumf;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
sumf = simd_sum(shmem_f32[tiisg]);
const float scale = 1.0f / sqrt(sumf / args.head_dim + args.eps);
const int n_nope = args.head_dim - args.n_dims;
if (n_nope < 0) {
return;
}
float corr_dims[2];
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base,
args.beta_fast, args.beta_slow, corr_dims);
const float theta_base = (float)(args.pos0 + (int)tok);
const float inv_ndims = -1.0f / args.n_dims;
device float * xs = (device float *)x4;
for (int i0 = tpitg.x; i0 < args.head_dim; i0 += ntg.x) {
if (i0 < n_nope) {
xs[i0] = xs[i0] * scale;
continue;
}
const int r = i0 - n_nope;
if ((r & 1) != 0) {
continue;
}
#ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta =
theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
#else
const float theta =
theta_base * pow(args.freq_base, inv_ndims * r);
#endif
float cos_theta;
float sin_theta;
rope_yarn(theta, args.freq_scale, corr_dims, r,
args.ext_factor, args.attn_factor,
&cos_theta, &sin_theta);
if (args.inverse) {
sin_theta = -sin_theta;
}
const float x0 = xs[i0] * scale;
const float x1 = xs[i0 + 1] * scale;
xs[i0] = x0 * cos_theta - x1 * sin_theta;
xs[i0 + 1] = x0 * sin_theta + x1 * cos_theta;
}
}
// DS4 positions are always affine within one RoPE dispatch. This variant
// reconstructs the same wrapped int32 position in-kernel, avoiding the host
// position array and its buffer binding while preserving the pair lane mapping
// and all floating-point operations of the specialization above.
/* Shared, deliberately noinline so that every caller gets bit-identical
* trigonometric codegen. The header note about tiny trig codegen changes
* flipping sampled tokens is exactly why this body must be compiled once and
* shared rather than inlined separately into each kernel. */
static __attribute__((noinline)) void ds4_rope_tail_pair_affine_row(
constant ds4_metal_args_dsv4_rope_affine_pair & args,
device const char * src_base,
device char * dst_base,
int n_nope,
uint raw_pos,
uint tid,
uint nthreads) {
float corr_dims[2];
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
const float theta_base = (float)as_type<int>(raw_pos);
const float theta_base = (float)as_type<int>(raw_pos);
const float inv_ndims = -1.f/args.n_dims;
device const char * src_base =
src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
device char * dst_base =
dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
for (int r = tid; r < args.n_dims; r += ntg.x) {
for (int r = tid; r < args.n_dims; r += nthreads) {
if ((r & 1) != 0) {
continue;
}
@@ -381,5 +471,413 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
*((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta;
*((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta;
}}
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
device const char * src0 [[buffer(1)]],
device char * dst [[buffer(4)]],
uint tid [[thread_index_in_threadgroup]],
ushort3 ntg [[threads_per_threadgroup]],
uint3 tgpig [[threadgroup_position_in_grid]]) {
const int i1 = tgpig[0];
const int i2 = tgpig[1];
const int n_nope = args.head_dim - args.n_dims;
if (n_nope < 0) {
return;
}
const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
device const char * src_base =
src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
device char * dst_base =
dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
ds4_rope_tail_pair_affine_row(args, src_base, dst_base, n_nope, raw_pos, tid, ntg.x);
}
// Decode-only fusion of the KV RoPE tail with the FP8/raw finalizer. Both were
// already single 64-thread threadgroups on the same row, back to back, so the
// pair cost two dispatches (~12.4 us) to touch 2 KB. The RoPE body below is a
// verbatim copy of kernel_dsv4_rope_tail_f32_inplace_pair_affine specialised to
// the decode grid (one head, one token, so i1 = i2 = 0) and the finalizer body
// is a verbatim copy of kernel_dsv4_kv_fp8_store_f32. The barrier between them
// is required because RoPE writes element pairs across lanes while the raw copy
// reads them per lane. Arithmetic, order and rounding are unchanged; the header
// warning above about trigonometric codegen still applies, so this kernel is
// gated and verified against full-vocabulary logits before promotion.
kernel void kernel_dsv4_kv_rope_fp8_store_f32(
constant ds4_metal_args_dsv4_kv_fp8_store & args,
constant ds4_metal_args_dsv4_rope_affine_pair & rope,
device float * kv,
device float * raw_cache,
threadgroup float * scratch [[threadgroup(0)]],
uint tid [[thread_index_in_threadgroup]]) {
{
const int rope_n_nope = rope.head_dim - rope.n_dims;
if (rope_n_nope < 0) {
return;
}
ds4_rope_tail_pair_affine_row(rope,
(device const char *)kv,
(device char *)kv,
rope_n_nope,
rope.pos0,
tid,
64u);
}
/* The RoPE helper writes device-memory pairs that different lanes read
* below. A threadgroup-only fence does not make those cross-lane device
* writes visible. */
threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
{
const int head_dim = args.head_dim;
const int n_rot = args.n_rot;
const int n_nope = head_dim - n_rot;
if (head_dim <= 0 || n_rot < 0 || n_nope < 0 || tid >= 64) {
return;
}
device float * raw = raw_cache + (int64_t)args.raw_row * head_dim;
for (int off = 0; off < n_nope; off += 64) {
float v = 0.0f;
if (off + (int)tid < n_nope) {
v = kv[off + tid];
scratch[tid] = abs(v);
} else {
scratch[tid] = 0.0f;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = 32; stride > 0; stride >>= 1) {
if (tid < stride) {
scratch[tid] = max(scratch[tid], scratch[tid + stride]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const float amax = max(scratch[0], 1.0e-4f);
const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
if (off + (int)tid < n_nope) {
const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
kv[off + tid] = q;
// Diagnostic only: skip the FP16 round-trip that normally matches the
// half-typed FlashAttention KV buffer's precision. With this enabled the
// indexer will see higher-precision raw values than FlashAttention does,
// which is informative but not a production-ready setting.
#ifdef DS4_METAL_KV_RAW_F32
raw[off + tid] = q;
#else
raw[off + tid] = (float)((half)q);
#endif
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
for (int i = n_nope + tid; i < head_dim; i += 64) {
#ifdef DS4_METAL_KV_RAW_F32
raw[i] = kv[i];
#else
raw[i] = (float)((half)kv[i]);
#endif
}
}
}
/* Decode-only sibling of kernel_flash_attn_ext_vec_reduce that also applies the
* inverse RoPE tail to the row it just produced, removing a whole dispatch per
* layer. Each threadgroup owns one head's entire 512-float row, so the RoPE is
* an intra-threadgroup dependency: reduce, barrier, rotate. Both halves call the
* same shared noinline helpers the standalone kernels use, so the arithmetic and
* its codegen are identical to running the two dispatches back to back. */
kernel void kernel_flash_attn_ext_vec_reduce_rope(
constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
device const char * htmp,
device char * dst,
constant ds4_metal_args_dsv4_rope_affine_pair & rope,
uint tgpig[[threadgroup_position_in_grid]],
ushort tiitg[[thread_index_in_threadgroup]],
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);
threadgroup_barrier(mem_flags::mem_device);
const int n_nope = rope.head_dim - rope.n_dims;
if (n_nope < 0) {
return;
}
device char * row = dst + (uint64_t)tgpig * rope.row_bytes;
ds4_rope_tail_pair_affine_row(rope,
(device const char *)row,
row,
n_nope,
rope.pos0,
tiitg,
(uint)(32 * FC_flash_attn_ext_vec_reduce_NWG));
}
struct ds4_metal_args_dsv4_comp_finalize {
ds4_metal_args_dsv4_rope_affine_pair rope;
float rms_eps;
uint32_t pad0;
};
/* Decode-only emit-path fusion. Every ratio-th token, each layer finalizes
* one freshly pooled compressor row per compressor: RMS norm, RoPE tail, and
* then the FP8 round-trip + F16 commit copy (attention, 512 floats) or the
* Hadamard+FP4 QAT (indexer, 128 floats). Those were seven single-row
* dispatches; this kernel is one dispatch with two threadgroups.
*
* Each phase reproduces its standalone kernel bit-exactly:
* - norm: kernel_rms_norm_mul_f32_4's tree (float4 lanes, simd_sum, zero-
* padded 32-slot cross-simdgroup reduce); 512 uses 128 virtual threads on
* simdgroups 0-3, 128 uses 32 virtual threads on simdgroup 0.
* - rope: ds4_rope_tail_pair_affine_row verbatim (lanes 0-63, nthreads=64).
* - fp8: kernel_dsv4_fp8_kv_quantize_f32's 64-lane shmem max tree and
* round-trip, src==dst so the verbatim tail copy is a no-op and dropped.
* - commit: per-element f32->f16 conversion (value-wise exact).
* - qat: kernel_dsv4_indexer_hadamard_fp4_f32's butterfly and per-32 amax
* tree on lanes 0-127.
* Threads outside a phase's virtual width still execute every barrier, so
* threadgroup barriers stay uniform across the 256-thread threadgroup. */
kernel void kernel_dsv4_comp_row_finalize_f32(
constant ds4_metal_args_dsv4_comp_finalize & args [[buffer(0)]],
device float * attn_row [[buffer(1)]],
device const float * attn_norm_w [[buffer(2)]],
device char * attn_cache [[buffer(3)]],
device float * index_row [[buffer(4)]],
device const float * index_norm_w [[buffer(5)]],
device float * attn_state_kv [[buffer(6)]],
device float * attn_state_score [[buffer(7)]],
device float * index_state_kv [[buffer(8)]],
device float * index_state_score [[buffer(9)]],
threadgroup float * shmem [[threadgroup(0)]],
uint tgpig [[threadgroup_position_in_grid]],
ushort tiitg [[thread_index_in_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
constant ds4_metal_args_dsv4_rope_affine_pair & rope_args = args.rope;
if (tgpig == 0) {
/* -------- attention compressor row (512 floats) -------- */
{
device float4 * y4 = (device float4 *)attn_row;
device const float4 * x4 = (device const float4 *)attn_row;
device const float4 * w4 = (device const float4 *)attn_norm_w;
if (sgitg == 0) {
shmem[tiisg] = 0.0f;
}
float sumf = 0.0f;
if (tiitg < 128) {
sumf = dot(x4[tiitg], x4[tiitg]);
}
sumf = simd_sum(sumf);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tiitg < 128 && tiisg == 0) {
shmem[sgitg] = sumf;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float total = 0.0f;
if (tiitg < 128) {
total = simd_sum(shmem[tiisg]);
}
const float mean = total / 512.0f;
const float scale = 1.0f/sqrt(mean + args.rms_eps);
if (tiitg < 128) {
y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg];
}
}
threadgroup_barrier(mem_flags::mem_device);
ds4_rope_tail_pair_affine_row(rope_args,
(device const char *)attn_row,
(device char *)attn_row,
512 - rope_args.n_dims,
rope_args.pos0,
tiitg,
64u);
threadgroup_barrier(mem_flags::mem_device);
for (int off = 0; off < 512 - rope_args.n_dims; off += 64) {
float v = 0.0f;
if (tiitg < 64) {
v = attn_row[off + tiitg];
shmem[tiitg] = abs(v);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = 32; stride > 0; stride >>= 1) {
if (tiitg < stride) {
shmem[tiitg] = max(shmem[tiitg], shmem[tiitg + stride]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const float amax = max(shmem[0], 1.0e-4f);
const float scale = exp2(ceil(log2(amax / 448.0f)));
if (tiitg < 64) {
const float q = dsv4_e4m3fn_dequant(clamp(v / scale, -448.0f, 448.0f)) * scale;
attn_row[off + tiitg] = q;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
threadgroup_barrier(mem_flags::mem_device);
if (tiitg < 128) {
device const float4 * x4 = (device const float4 *)attn_row;
device half4 * o4 = (device half4 *)attn_cache;
const float4 v = x4[tiitg];
o4[tiitg] = half4(v);
}
return;
}
if (tgpig >= 2u) {
/* Ratio-4 state shifts for both compressors (elementwise row move,
* so the flat gid mapping is bit-exact): 4*1024 attention elements
* then 4*256 indexer elements. */
const uint gid = (tgpig - 2u) * 256u + tiitg;
const uint n0 = 4u * 1024u;
if (gid < n0) {
attn_state_kv[gid] = attn_state_kv[n0 + gid];
attn_state_score[gid] = attn_state_score[n0 + gid];
return;
}
const uint gid1 = gid - n0;
const uint n1 = 4u * 256u;
if (gid1 >= n1) return;
index_state_kv[gid1] = index_state_kv[n1 + gid1];
index_state_score[gid1] = index_state_score[n1 + gid1];
return;
}
/* -------- indexer compressor row (128 floats) -------- */
{
device float4 * y4 = (device float4 *)index_row;
device const float4 * x4 = (device const float4 *)index_row;
device const float4 * w4 = (device const float4 *)index_norm_w;
if (sgitg == 0) {
shmem[tiisg] = 0.0f;
}
float sumf = 0.0f;
if (tiitg < 32) {
sumf = dot(x4[tiitg], x4[tiitg]);
}
sumf = simd_sum(sumf);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tiisg == 0) {
shmem[sgitg] = sumf;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float total = 0.0f;
if (tiitg < 32) {
total = simd_sum(shmem[tiisg]);
}
const float mean = total / 128.0f;
const float scale = 1.0f/sqrt(mean + args.rms_eps);
if (tiitg < 32) {
y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg];
}
}
threadgroup_barrier(mem_flags::mem_device);
ds4_rope_tail_pair_affine_row(rope_args,
(device const char *)index_row,
(device char *)index_row,
128 - rope_args.n_dims,
rope_args.pos0,
tiitg,
64u);
threadgroup_barrier(mem_flags::mem_device);
{
threadgroup float *vals = shmem;
threadgroup float *absbuf = shmem + 128;
if (tiitg < 128) {
vals[tiitg] = index_row[tiitg];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = 1u; stride < 128u; stride <<= 1u) {
if (tiitg < 128 && (tiitg & stride) == 0u) {
const uint base = (tiitg & ~(2u * stride - 1u)) + (tiitg & (stride - 1u));
const float a = vals[base];
const float b = vals[base + stride];
vals[base] = a + b;
vals[base + stride] = a - b;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float v = 0.0f;
if (tiitg < 128) {
v = vals[tiitg] * 0.08838834764831845f;
absbuf[tiitg] = abs(v);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint block = tiitg >> 5u;
const uint lane = tiitg & 31u;
const uint block_base = block * 32u;
for (uint stride = 16u; stride > 0u; stride >>= 1u) {
if (tiitg < 128 && lane < stride) {
absbuf[block_base + lane] = max(absbuf[block_base + lane],
absbuf[block_base + lane + stride]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tiitg < 128) {
const float amax = max(absbuf[block_base], 7.052966104933725e-38f);
const float scale = exp2(ceil(log2(amax / 6.0f)));
index_row[tiitg] = dsv4_e2m1fn_dequant(clamp(v / scale, -6.0f, 6.0f)) * scale;
}
}
}
// Host-visible packed FlashAttention + exact inverse-RoPE decode kernel.
kernel void kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512(
constant ds4_metal_args_flash_attn_ext_vec & args [[buffer(0)]],
device const char * q [[buffer(1)]],
device const char * k [[buffer(2)]],
device const char * v [[buffer(3)]],
device const char * mask [[buffer(4)]],
device const char * sinks [[buffer(5)]],
device const char * pad [[buffer(6)]],
device char * dst [[buffer(7)]],
constant ds4_metal_args_dsv4_rope_affine_pair & rope
[[buffer(8)]],
threadgroup char * shmem [[threadgroup(0)]],
uint head [[threadgroup_position_in_grid]],
ushort tiitg [[thread_index_in_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
/* Uniform specialization guard; host applies the same eligibility gate. */
if (!FC_flash_attn_ext_vec_has_mask ||
!FC_flash_attn_ext_vec_has_sinks ||
FC_flash_attn_ext_vec_has_bias ||
FC_flash_attn_ext_vec_has_scap ||
FC_flash_attn_ext_vec_nsg != 1 ||
FC_flash_attn_ext_vec_nwg != 32 ||
FC_flash_attn_ext_vec_ns10 != 512 ||
FC_flash_attn_ext_vec_ns20 != 512 ||
args.ne01 != 1 || args.ne02 != 64 || args.ne03 != 1 ||
args.ne_12_2 != 1 || args.ne_12_3 != 1 ||
args.ne31 != 1 || args.ne32 != 1 || args.ne33 != 1 ||
args.ne11 <= 0 || args.ne11 > 1024 || head >= (uint)args.ne02 ||
args.nb02 != 2048 || args.nb11 != 1024 || args.nb21 != 1024 ||
rope.head_dim != 512 || rope.n_dims != 64 ||
rope.row_bytes != 2048 || rope.inverse == 0) {
return;
}
ds4_flash_attn_vec_packed8_reduce_f16_512(
args, q, k, v, mask, sinks, pad, dst, shmem,
head, tiisg, sgitg);
/* Same producer/consumer boundary as the current reduce+RoPE kernel. */
threadgroup_barrier(mem_flags::mem_device);
const int n_nope = rope.head_dim - rope.n_dims;
device char * row = dst + (uint64_t)head * rope.row_bytes;
ds4_rope_tail_pair_affine_row(rope,
(device const char *)row,
row,
n_nope,
rope.pos0,
tiitg,
32u * 32u);
}