Skip to content

metal: DSpark drafter working end-to-end (3 ops + 4 continuous/drafter fixes) #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
robotnursenyc wants to merge 2 commits into
base: batched-serving
Choose a base branch
from
Open

metal: DSpark drafter working end-to-end (3 ops + 4 continuous/drafter fixes) #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eab172d
metal: implement the three DSpark drafter ops (capture_mean, gather_c…
cashlabnj Jul 19, 2026
d7201ad
metal: four fixes that make the DSpark drafter run end-to-end
robotnursenyc Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 101 additions & 0 deletions METAL_DSPARK.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 <v4-flash-target>.gguf --dspark <dspark-drafter>.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.
25 changes: 19 additions & 6 deletions ds4.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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",
Expand Down
9 changes: 7 additions & 2 deletions ds4_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Loading