diff --git a/METAL_DSPARK.md b/METAL_DSPARK.md new file mode 100644 index 000000000..d6964b4c2 --- /dev/null +++ b/METAL_DSPARK.md @@ -0,0 +1,101 @@ +# DSpark speculative decoding on Apple Silicon (Metal) + +This branch makes the **DSpark block drafter** run end-to-end on the Metal +backend for DeepSeek‑V4‑Flash. It combines two pieces: + +1. The three previously‑stubbed DSpark Metal ops (PR #2): + `capture_mean`, `gather_concat`, `markov_step`. +2. Four correctness fixes on the continuous‑batching + drafter path that were + blocking the drafter (and, for the first, blocking non‑drafter continuous + batching too). + +With all of it applied, the continuous path serves without falling back to the +serial path, and the DSpark drafter produces real, accepted drafts. + +## The four fixes + +Each is small, self‑contained, and (where noted) byte‑identical to the prior +behavior on the paths it does not intend to change. + +- **Compressed‑KV banked emit wrote F32 through an F32‑sized view of the F16 + cache.** `ds4_gpu_compressor_update_tensor` runs pool/rms/rope in F32 and is + shared by two callers with different cache dtypes: the serial decode caller + passes the F32 stage (correct), while the banked `ms_row_emit` caller passes + the F16 `layer_attn_comp_cache[il]` directly. Writing F32 into an + `sizeof(float)`‑strided view of a half‑width buffer corrupts one element per + comp row (a NaN at a fixed dim), which then propagates to 100% NaN at the + first compressed‑KV attention layer once `n_comp` pushes the visible key set + past the boundary. Fix: a caller‑scoped `output_is_f16` flag — the banked + caller runs the emit in an F32 scratch and `copy_f32_to_f16`s into the + distinct cache row; every F32 caller is unchanged. + +- **`ds4_gpu_matmul_f32_tensor` was GEMV‑only** (`n_tok != 1` returned failure). + The main model's HC‑mixer weight is F16 and takes the batched F16 path, but + the DSpark drafter's HC‑mixer weight is F32 and the block runs it at + `n_tok = block size`, hitting the early return. Fix: for `n_tok != 1`, run + `n_tok` independent GEMVs over per‑row views. Byte‑identical for `n_tok == 1`. + +- **`ds4_gpu_tensor_copy` needs an open batch command buffer.** It blits into + the current `g_batch_cb` and returns failure when none is open (unlike + `ds4_gpu_tensor_copy_f32_to_f16`, which owns its own). The DSpark inject block + ends the batch cb, so the checkpoint‑restore copies in + `mtp_cont_rollback_restore_all` had no cb and failed. Fix: open a cb around + the rollback only when none is open (`begin_commands` returns 0 when one is + already open, so the MTP path is unchanged). + +- **The HC split/sinkhorn's "numerically stable" sigmoid overflows to NaN.** + `0.5*tanh(0.5*z)+0.5` was introduced to bound the sigmoid, but under Metal's + fast‑math `tanh` (which is `exp`‑based) it evaluates `exp(large) → inf`, then + `inf/inf → NaN`, for large‑positive `z`. A drafter block row whose HC‑mixer + logit is large enough (the real seed token) trips it. Fix: clamp the sigmoid + input to `[-30, 30]` before the transcendental — sigmoid is fully saturated + by ±30, so in‑range values are unchanged and only the overflow is removed. + +## Observed behavior + +Speculative decoding acceptance is content‑dependent (as it always is). Rough +single‑run numbers on an M3 Ultra, Q4‑imatrix target + the public Q2K DSpark +drafter, `temperature = 0`, with the per‑sequence yield guard disabled so the +drafter runs on every step (`DS4_DSPARK_QUENCH=0`): + +| workload | acceptance | tokens/step | +|---------------------|-----------:|------------:| +| repetitive text | ~68% | ~2.4x | +| stepwise arithmetic | ~22% | ~1.3x | +| code | ~20% | ~1.2x | +| free prose | ~13% | ~1.2x | + +These are illustrative, not a rigorous benchmark (single runs, one machine, one +quant). The yield guard is on by default and shuts the drafter off per sequence +when it is not paying for itself. A higher‑precision drafter (the public GGUF is +2‑bit; the reference checkpoint's experts are FP4) would raise the harder‑content +numbers, but is not required for a working, speedup‑positive drafter. + +## Running it + +``` +./ds4-server -m .gguf --dspark .gguf --ctx 100000 ... +``` + +The drafter arms on the continuous path. `DS4_DSPARK_QUENCH=0` disables the +per‑sequence yield guard (useful for measuring raw acceptance). + +## Credits + +This work stands on: + +- **Entrpi/ds4** — the fork this branch targets (CUDA/GB10 DSpark drafter, + block scorer, continuous batching). +- **antirez/ds4** — the upstream DeepSeek‑V4‑Flash inference engine and its + Metal backend. +- **DeepSeek** — the DeepSeek‑V4‑Flash model and the DSpark speculative‑decoding + method and reference drafter checkpoint (MIT). +- **Unsloth** — DeepSeek‑V4‑Flash GGUF quantization work. +- **llama.cpp** contributors — the Metal GEMV/GEMM and softmax kernel patterns + these ops and fixes build on. + +## Note on authorship + +The four fixes, the root‑cause analysis, and this document were produced with AI +assistance (Claude) and reviewed before submission. Measurements are from a +single machine and should be reproduced independently before being relied on. diff --git a/ds4.c b/ds4.c index 66861675d..3090b266d 100644 --- a/ds4.c +++ b/ds4.c @@ -11586,7 +11586,8 @@ static bool metal_graph_encode_decode_layer_impl( DS4_RMS_EPS, il, /* PC2: per-layer substrate index for emit-row */ DS4_COMPRESSOR_ROW_COMP, /* PC2: primary compressor -> ls->comp_row */ - NULL, NULL, 0u /* C2: no per-row substrate (serial width-1) */) != 0; + NULL, NULL, 0u /* C2: no per-row substrate (serial width-1) */, + /*output_is_f16=*/0) != 0; if (ok && emit) { /* Step 4c R1': read comp_row from g_layer_dev[il].comp_row in * the per-layer substrate (populated at top of token in @@ -11715,7 +11716,8 @@ static bool metal_graph_encode_decode_layer_impl( DS4_RMS_EPS, il, /* PC2: per-layer substrate index for emit-row */ DS4_COMPRESSOR_ROW_INDEX, /* PC2: indexer compressor -> ls->index_row */ - NULL, NULL, 0u /* C2: no per-row substrate (serial width-1) */) != 0; + NULL, NULL, 0u /* C2: no per-row substrate (serial width-1) */, + /*output_is_f16=*/0) != 0; /* TEMPORARY DIAGNOSTIC: state_kv AFTER compressor_update_tensor * returns (post compressor_store_batch + post update_pool + * shift). Slot 230+il. If this matches BE vs BC but the cache @@ -16913,7 +16915,8 @@ static bool metal_graph_encode_layer_attention_batch( * published table value). */ cap_step ? g->batch_positions : NULL, cap_step ? g->batch_seq_id : NULL, - t) != 0; + t, + /*output_is_f16=*/(!ms_fp8_primary && DS4_GPU_ATTN_COMP_CACHE_F16)) != 0; ds4_gpu_stage_prof_end(sp_cupd); if (g_cont_prof > 0) g_cont_lyr_cupd_s += now_sec() - cupd_t0; if (ok && emit && cap_step) { @@ -17338,7 +17341,8 @@ static bool metal_graph_encode_layer_attention_batch( DS4_RMS_EPS, UINT32_MAX, /* PC2: decode2-exact, no substrate; inline comp_row */ DS4_COMPRESSOR_ROW_COMP, /* PC2: row_field ignored when il==UINT32_MAX */ - NULL, NULL, 0u) != 0; + NULL, NULL, 0u, + /*output_is_f16=*/0) != 0; if (ok && emit) { ds4_gpu_tensor *comp_row_view = row_fp8_primary ? ds4_gpu_tensor_view(row_scratch, 0, @@ -17757,7 +17761,8 @@ static bool metal_graph_encode_layer_attention_batch( : DS4_COMPRESSOR_ROW_INDEX, icap_step ? g->batch_positions : NULL, icap_step ? g->batch_seq_id : NULL, - t) != 0; + t, + /*output_is_f16=*/0) != 0; if (ok && emit && icap_step) { /* C1: capture-safe QAT -- row-variant reads the emit * row live from the per-ROW table entry (row t; C2 @@ -18136,7 +18141,8 @@ static bool metal_graph_encode_layer_attention_batch( DS4_RMS_EPS, UINT32_MAX, /* PC2: decode2-exact indexer, no substrate */ DS4_COMPRESSOR_ROW_INDEX, /* PC2: row_field ignored when il==UINT32_MAX */ - NULL, NULL, 0u) != 0; + NULL, NULL, 0u, + /*output_is_f16=*/0) != 0; if (ok && emit) { ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( g->layer_index_comp_cache[il], @@ -35418,8 +35424,15 @@ int ds4_engine_continuous_generate(ds4_batch_ctx *ctx, * write). The deferred commit-reforward below is skipped on this path (no_defer set). */ if (can_rollback && !ctx->mtp_force_norollback) { const double rb_t0 = dspark_prof && dspark_now ? now_sec() : 0.0; + /* mtp_cont_rollback_restore_all restores checkpoints with + * ds4_gpu_tensor_copy, which blits into g_batch_cb and returns 0 + * when none is open. The DSpark inject block above ends the batch + * cb, so open one iff none is open (begin_commands returns 0 when + * already open -> MTP path unchanged); close it after. */ + const int rb_opened = ds4_gpu_begin_commands(); ok = mtp_cont_rollback_restore_all(g, MS, committed_rows, vnr, ctx->rollback_verify, &ctx->rollback_verify_fails); + if (rb_opened) ok = (ds4_gpu_end_commands() != 0) && ok; if (dspark_prof && dspark_now) dspark_t_rollback_s += now_sec() - rb_t0; if (!ok) { fprintf(stderr, "ds4: cont-mtp step %u FAIL @rollback (live=%u)\n", diff --git a/ds4_gpu.h b/ds4_gpu.h index 8ef3853d0..eebe89e44 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -1298,7 +1298,11 @@ int ds4_gpu_compressor_update_tensor( * serial / eager caller; Metal backend rejects seq_id != NULL. */ const ds4_gpu_tensor *positions, const ds4_gpu_tensor *seq_id, - uint32_t row_idx); + uint32_t row_idx, + /* 1 => comp_cache is an F16 buffer (the banked attn emit passes + * layer_attn_comp_cache[il] direct): run pool/rms/rope in an F32 + * scratch then copy_f32_to_f16 into the row. 0 => F32 cache. */ + int output_is_f16); /* M2-Inc5: emit tail of ds4_gpu_compressor_update_tensor only (pool + norm + * rope + ratio4 shift; no store) -- the follow-up call after a successful @@ -1333,7 +1337,8 @@ int ds4_gpu_compressor_update_tail_tensor( int row_field, const ds4_gpu_tensor *positions, const ds4_gpu_tensor *seq_id, - uint32_t row_idx); + uint32_t row_idx, + int output_is_f16); int ds4_gpu_compressor_store_batch_tensor( const ds4_gpu_tensor *kv, diff --git a/ds4_metal.m b/ds4_metal.m index 89d028461..6e9ac0fb9 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -4861,10 +4861,47 @@ int ds4_gpu_dspark_markov_step_tensor( const ds4_gpu_tensor *bias, uint32_t vocab, uint32_t nb) { - (void)cand_out; (void)base_logits; (void)blk_pos; (void)B; - (void)bias; (void)vocab; (void)nb; - fprintf(stderr, "ds4: dspark_markov_step is CUDA-only (DSpark draft)\n"); - return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!cand_out || !base_logits || !bias || vocab == 0 || nb == 0 || B == 0 || blk_pos >= B) + return 0; + if (ds4_gpu_tensor_bytes(cand_out) < (uint64_t)nb * sizeof(int32_t) || + ds4_gpu_tensor_bytes(base_logits) < (uint64_t)nb * B * vocab * sizeof(float) || + ds4_gpu_tensor_bytes(bias) < (uint64_t)nb * vocab * sizeof(float)) { + fprintf(stderr, "ds4: Metal dspark_markov_step received undersized buffers\n"); + return 0; + } + + @autoreleasepool { + id pipeline = + ds4_gpu_get_pipeline("kernel_dspark_markov_step"); + if (!pipeline) return 0; + + /* The shader's reduction assumes a power-of-two thread count and its + * threadgroup arrays are sized for 1024. */ + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 1024u) nth = 1024u; + NSUInteger pow2 = 1u; + while (pow2 * 2u <= nth) pow2 *= 2u; + nth = pow2; + + struct { uint32_t blk_pos, B, vocab, nb; } args = + { blk_pos, B, vocab, nb }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(cand_out) offset:ds4_gpu_tensor_offset(cand_out) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(base_logits) offset:ds4_gpu_tensor_offset(base_logits) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(bias) offset:ds4_gpu_tensor_offset(bias) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nb, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "dspark markov step")) return 0; + } + return 1; } int ds4_gpu_dspark_gather_concat_tensor( @@ -4873,9 +4910,39 @@ int ds4_gpu_dspark_gather_concat_tensor( const ds4_gpu_tensor *src_rows, uint32_t n_rows, uint32_t row_floats) { - (void)dst; (void)src; (void)src_rows; (void)n_rows; (void)row_floats; - fprintf(stderr, "ds4: dspark_gather_concat is CUDA-only (DSpark draft)\n"); - return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!dst || !src || !src_rows || n_rows == 0 || row_floats == 0) return 0; + const uint64_t dst_bytes = (uint64_t)n_rows * row_floats * sizeof(float); + if (ds4_gpu_tensor_bytes(dst) < dst_bytes || + ds4_gpu_tensor_bytes(src_rows) < (uint64_t)n_rows * sizeof(int32_t)) { + fprintf(stderr, "ds4: Metal dspark_gather_concat received undersized buffers\n"); + return 0; + } + + @autoreleasepool { + id pipeline = + ds4_gpu_get_pipeline("kernel_dspark_gather_concat"); + if (!pipeline) return 0; + + struct { uint32_t n_rows, row_floats; } args = { n_rows, row_floats }; + const uint64_t n = (uint64_t)n_rows * row_floats; + const NSUInteger groups = (NSUInteger)((n + 255u) / 256u); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(dst) offset:ds4_gpu_tensor_offset(dst) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(src) offset:ds4_gpu_tensor_offset(src) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(src_rows) offset:ds4_gpu_tensor_offset(src_rows) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "dspark gather concat")) return 0; + } + return 1; } int ds4_gpu_dspark_capture_mean_tensor( @@ -4886,10 +4953,39 @@ int ds4_gpu_dspark_capture_mean_tensor( uint32_t n_tokens, uint32_t slot, uint32_t n_slots) { - (void)concat; (void)hc; (void)n_embd; (void)n_hc; - (void)n_tokens; (void)slot; (void)n_slots; - fprintf(stderr, "ds4: dspark_capture_mean is CUDA-only (DSpark draft)\n"); - return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!concat || !hc || n_embd == 0 || n_hc == 0 || n_tokens == 0 || slot >= n_slots) + return 0; + if (ds4_gpu_tensor_bytes(concat) < (uint64_t)n_tokens * n_slots * n_embd * sizeof(float) || + ds4_gpu_tensor_bytes(hc) < (uint64_t)n_tokens * n_hc * n_embd * sizeof(float)) { + fprintf(stderr, "ds4: Metal dspark_capture_mean received undersized buffers\n"); + return 0; + } + + @autoreleasepool { + id pipeline = + ds4_gpu_get_pipeline("kernel_dspark_capture_mean"); + if (!pipeline) return 0; + + struct { uint32_t n_embd, n_hc, n_tokens, slot, n_slots; } args = + { n_embd, n_hc, n_tokens, slot, n_slots }; + const uint64_t n = (uint64_t)n_embd * n_tokens; + const NSUInteger groups = (NSUInteger)((n + 255u) / 256u); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(concat) offset:ds4_gpu_tensor_offset(concat) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(hc) offset:ds4_gpu_tensor_offset(hc) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "dspark capture mean")) return 0; + } + return 1; } void ds4_gpu_cleanup(void) { @@ -6780,7 +6876,25 @@ int ds4_gpu_matmul_f32_tensor( const ds4_gpu_tensor *x, uint64_t n_tok) { if (!g_initialized && !ds4_gpu_init()) return 0; - if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || n_tok != 1) return 0; + if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || n_tok == 0) return 0; + + /* F32 plain matmul previously supported n_tok==1 only (GEMV); batched + * callers (e.g. the DSpark drafter block's F32 HC-mixer at n_tok=R) + * hit the n_tok!=1 early-return. Run n_tok independent GEMVs over + * per-row views; byte-identical for n_tok==1. */ + if (n_tok != 1) { + int ok_all = 1; + for (uint64_t t = 0; t < n_tok && ok_all; t++) { + ds4_gpu_tensor *xr = ds4_gpu_tensor_view(x, t * in_dim * sizeof(float), in_dim * sizeof(float)); + ds4_gpu_tensor *outr = ds4_gpu_tensor_view(out, t * out_dim * sizeof(float), out_dim * sizeof(float)); + ok_all = xr && outr && + ds4_gpu_matmul_f32_tensor(outr, model_map, model_size, weight_offset, + in_dim, out_dim, xr, 1); + if (xr) ds4_gpu_tensor_free(xr); + if (outr) ds4_gpu_tensor_free(outr); + } + return ok_all; + } @autoreleasepool { id xbuf = ds4_gpu_tensor_buffer(x); @@ -9486,7 +9600,9 @@ int ds4_gpu_compressor_update_tail_tensor( int row_field, const ds4_gpu_tensor *positions, const ds4_gpu_tensor *seq_id, - uint32_t row_idx) { + uint32_t row_idx, + int output_is_f16) { + (void)output_is_f16; (void)kv_cur; (void)sc_cur; (void)state_kv; (void)state_score; (void)comp_cache; (void)model_map; (void)model_size; (void)ape_offset; (void)ape_type; (void)norm_offset; (void)norm_type; (void)head_dim; @@ -9526,7 +9642,8 @@ int ds4_gpu_compressor_update_tensor( int row_field, const ds4_gpu_tensor *positions, const ds4_gpu_tensor *seq_id, - uint32_t row_idx) { + uint32_t row_idx, + int output_is_f16) { (void)il; /* Metal kernels read inline args; no per-layer substrate. */ (void)row_field; /* PC2: substrate selector ignored on Metal. */ (void)positions; (void)row_idx; /* CUDA-only substrate. */ @@ -9545,10 +9662,16 @@ int ds4_gpu_compressor_update_tensor( const uint32_t width = coff * head_dim; const uint32_t state_rows = coff * ratio; const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u; + /* The banked attn emit passes an F16 comp cache; pool/rms/rope write + * device float*, so an f32-offset view into the F16 buffer corrupts + * the row. When output_is_f16, emit into an F32 scratch then + * copy_f32_to_f16 into the cache row. */ + const bool use_f16_fix = output_is_f16 != 0; + const uint64_t comp_elem = use_f16_fix ? sizeof(uint16_t) : sizeof(float); const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; const uint64_t kv_bytes = (uint64_t)width * sizeof(float); const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * comp_elem; const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); @@ -9604,7 +9727,11 @@ int ds4_gpu_compressor_update_tensor( } if (!emit) return 1; - ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view( + /* F16 cache -> emit into an F32 scratch (converted below); else the + * original f32 view straight into the (F32) comp cache row. */ + ds4_gpu_tensor *comp_row_view = use_f16_fix + ? ds4_gpu_tensor_alloc((uint64_t)head_dim * sizeof(float)) + : ds4_gpu_tensor_view( comp_cache, (uint64_t)comp_row * head_dim * sizeof(float), (uint64_t)head_dim * sizeof(float)); @@ -9648,6 +9775,13 @@ int ds4_gpu_compressor_update_tensor( beta_fast, beta_slow) != 0; } + if (ok && use_f16_fix) { + /* Convert the F32 emit scratch into the distinct F16 cache row. */ + ok = ds4_gpu_tensor_copy_f32_to_f16( + comp_cache, + (uint64_t)comp_row * head_dim * sizeof(uint16_t), + comp_row_view, 0, head_dim) != 0; + } if (ok && ratio == 4u) { cb = ds4_gpu_command_buffer(&owned); ok = cb && diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal index 4d721b569..47c427cce 100644 --- a/metal/dsv4_hc.metal +++ b/metal/dsv4_hc.metal @@ -89,16 +89,16 @@ struct ds4_metal_args_dsv4_hc_expand { // logits on M5 Max, while the historical inline exp form remains finite and is // the decode throughput baseline. #ifdef DS4_METAL_HC_STABLE -static inline float ds4_hc_sigmoid(float z) { return 0.5f * tanh(0.5f * z) + 0.5f; } -static inline float4 ds4_hc_sigmoid(float4 z) { return 0.5f * tanh(0.5f * z) + 0.5f; } +static inline float ds4_hc_sigmoid(float z) { z = clamp(z, -30.0f, 30.0f); return 0.5f * tanh(0.5f * z) + 0.5f; } +static inline float4 ds4_hc_sigmoid(float4 z) { z = clamp(z, float4(-30.0f), float4(30.0f)); return 0.5f * tanh(0.5f * z) + 0.5f; } // 2 * sigmoid(z) == 1 + tanh(z/2). -static inline float ds4_hc_twice_sigmoid(float z) { return 1.0f + tanh(0.5f * z); } -static inline float4 ds4_hc_twice_sigmoid(float4 z) { return 1.0f + tanh(0.5f * z); } +static inline float ds4_hc_twice_sigmoid(float z) { z = clamp(z, -30.0f, 30.0f); return 1.0f + tanh(0.5f * z); } +static inline float4 ds4_hc_twice_sigmoid(float4 z) { z = clamp(z, float4(-30.0f), float4(30.0f)); return 1.0f + tanh(0.5f * z); } #else -static inline float ds4_hc_sigmoid(float z) { return 1.0f / (1.0f + exp(-z)); } -static inline float4 ds4_hc_sigmoid(float4 z) { return 1.0f / (1.0f + exp(-z)); } -static inline float ds4_hc_twice_sigmoid(float z) { return 2.0f / (1.0f + exp(-z)); } -static inline float4 ds4_hc_twice_sigmoid(float4 z) { return 2.0f / (1.0f + exp(-z)); } +static inline float ds4_hc_sigmoid(float z) { z = clamp(z, -30.0f, 30.0f); return 1.0f / (1.0f + exp(-z)); } +static inline float4 ds4_hc_sigmoid(float4 z) { z = clamp(z, float4(-30.0f), float4(30.0f)); return 1.0f / (1.0f + exp(-z)); } +static inline float ds4_hc_twice_sigmoid(float z) { z = clamp(z, -30.0f, 30.0f); return 2.0f / (1.0f + exp(-z)); } +static inline float4 ds4_hc_twice_sigmoid(float4 z) { z = clamp(z, float4(-30.0f), float4(30.0f)); return 2.0f / (1.0f + exp(-z)); } #endif // Splits an HC mixer row into pre weights, post gates, and the HC-to-HC diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index 7e8cbd8a7..68c85f8b7 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -1325,3 +1325,99 @@ kernel void kernel_dsv4_softmax_pool( *((device float *) (dst + id*args.nb0 + ic*args.nb1)) = acc/sum; } + +/* --------------------------------------------------------------------------- + * DSpark drafter support kernels (Metal ports of the CUDA originals in + * ds4_cuda.cu: dspark_capture_mean_kernel, dspark_gather_concat_kernel, + * dspark_markov_step_kernel). + * ------------------------------------------------------------------------ */ + +typedef struct { + uint n_embd; + uint n_hc; + uint n_tokens; + uint slot; + uint n_slots; +} ds4_dspark_capture_mean_args; + +kernel void kernel_dspark_capture_mean( + constant ds4_dspark_capture_mean_args &args [[buffer(0)]], + device float *concat [[buffer(1)]], + const device float *hc [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const uint n = args.n_embd * args.n_tokens; + if (gid >= n) return; + const uint d = gid % args.n_embd; + const uint t = gid / args.n_embd; + float acc = 0.0f; + for (uint h = 0; h < args.n_hc; h++) + acc += hc[(ulong)t * args.n_hc * args.n_embd + (ulong)h * args.n_embd + d]; + concat[(ulong)t * args.n_slots * args.n_embd + (ulong)args.slot * args.n_embd + d] = + acc / (float)args.n_hc; +} + +typedef struct { + uint n_rows; + uint row_floats; +} ds4_dspark_gather_concat_args; + +kernel void kernel_dspark_gather_concat( + constant ds4_dspark_gather_concat_args &args [[buffer(0)]], + device float *dst [[buffer(1)]], + const device float *src [[buffer(2)]], + const device int *src_rows [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint n = args.n_rows * args.row_floats; + if (gid >= n) return; + const uint d = gid % args.row_floats; + const uint r = gid / args.row_floats; + const int sr = src_rows[r]; + if (sr < 0) return; + dst[(ulong)r * args.row_floats + d] = src[(ulong)(uint)sr * args.row_floats + d]; +} + +typedef struct { + uint blk_pos; + uint B; + uint vocab; + uint nb; +} ds4_dspark_markov_step_args; + +/* One threadgroup per bank; parallel argmax over vocab of + * base_logits[(b*B+blk_pos)*vocab + v] + bias[b*vocab + v]. + * Tie -> lowest index, matching the CUDA kernel and the host refine's + * strict-greater scan. nth is a power of two (enforced host-side). */ +kernel void kernel_dspark_markov_step( + constant ds4_dspark_markov_step_args &args [[buffer(0)]], + device int *cand_out [[buffer(1)]], + const device float *base_logits [[buffer(2)]], + const device float *bias [[buffer(3)]], + uint b [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint nth [[threads_per_threadgroup]]) { + threadgroup float sm_val[1024]; + threadgroup int sm_idx[1024]; + if (b >= args.nb) return; + const device float *blog = base_logits + (ulong)(b * args.B + args.blk_pos) * args.vocab; + const device float *brow = bias + (ulong)b * args.vocab; + float local_v = -INFINITY; + int local_i = 0; + for (uint v = tid; v < args.vocab; v += nth) { + const float s = blog[v] + brow[v]; + if (s > local_v) { local_v = s; local_i = (int)v; } + } + sm_val[tid] = local_v; + sm_idx[tid] = local_i; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = nth / 2u; s > 0u; s >>= 1) { + if (tid < s) { + const float vr = sm_val[tid + s]; + const int ir = sm_idx[tid + s]; + const float vl = sm_val[tid]; + const int il = sm_idx[tid]; + if ((vr > vl) || (vr == vl && ir < il)) { sm_val[tid] = vr; sm_idx[tid] = ir; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0) cand_out[b] = sm_idx[0]; +}