Skip to content

cuda: run streaming selected load for single-token ffn batch encodes#497

Open
giannisanni wants to merge 1 commit into
antirez:mainfrom
giannisanni:streaming-single-token-batch-fix
Open

cuda: run streaming selected load for single-token ffn batch encodes#497
giannisanni wants to merge 1 commit into
antirez:mainfrom
giannisanni:streaming-single-token-batch-fix

Conversation

@giannisanni

Copy link
Copy Markdown

Fixes the concrete failure behind the --mtp + --ssd-streaming guard on CUDA (context: #495).

metal_graph_cuda_stream_prefill_batch_selected_load() skips when n_tokens <= 1, assuming batch encodes only come from multi-token prefill chunks. The MTP verifier can issue single-position ffn batch encodes; with the selected load skipped, the routed MoE falls through to allocating full resident expert ranges, which by definition cannot fit under SSD streaming:

ds4: CUDA model arena alloc failed for moe_gate (1792.00 MiB chunk): out of memory
ds4: gpu layer 2 ffn batch encode failed
ds4: decode failed: MTP verifier failed

The one-line fix lets n_tokens == 1 take the selected load; the batch machinery already handles a one-token union.

Tested on RTX 4060 Ti 16GB / Linux, Flash q2 streaming from NVMe, with the engine guard bypassed locally for testing:

config before after
margin --mtp-draft 2 hard crash (above) runs, 1.32 t/s
margin --mtp-draft 4 OOM + fallback loops, 0.85 t/s clean, 0.84 t/s
strict --mtp-draft 2 1.25 t/s, token-identical to greedy 1.18 t/s, still token-identical
non-MTP baseline 2.65 t/s 2.55 t/s (noise)

This does not touch the engine guard itself; whether to relax it (e.g. allow with strict verify) seems like your call. Happy to follow up with that change plus docs if you want it.

metal_graph_cuda_stream_prefill_batch_selected_load() skipped when
n_tokens <= 1, on the assumption that batch encodes always come from
multi-token prefill chunks. The MTP verifier can issue single-position
ffn batch encodes, and with the selected load skipped the routed MoE
falls through to allocating full resident expert ranges, which by
definition cannot fit in device memory under SSD streaming:

    ds4: CUDA model arena alloc failed for moe_gate (1792.00 MiB chunk): out of memory
    ds4: gpu layer 2 ffn batch encode failed
    ds4: decode failed: MTP verifier failed

Allow n_tokens == 1 through the selected load. The batch machinery
already handles it (a one-token union of DS4_N_EXPERT_USED experts).

Tested on RTX 4060 Ti 16GB / Linux, DeepSeek V4 Flash q2 streaming
from NVMe, with the --mtp guard bypassed locally:
- margin --mtp-draft 2, previously a hard crash, now runs
- margin --mtp-draft 4, previously OOM/fallback loops, now clean
- strict --mtp-draft 2 output stays token-identical to plain greedy
- non-MTP baseline throughput unchanged
fulvius31 added a commit to fulvius31/ds4 that referenced this pull request Jul 24, 2026
…ez#524)

server: when thinking mode is on and the model never emits </think>
(typically truncated at max_tokens mid-thought), surface the whole text
as unfinished reasoning with empty content instead of letting it leak
into content — applied to both the DSML and GLM tool-call parse paths
(upstream antirez#524, adapted to the split parser). Verified live:
/v1/messages now returns a thinking block plus empty text for the
truncated case; closed-thinking responses are unchanged.

cuda: treat every model range as accessible when the HMM direct path is
active (upstream antirez#158, the is_cached hunk only). The PR's
full-model prefetch half is deliberately not ported: prefetching a
model larger than RAM would thrash the page cache on the SSD-streaming
Spark boxes. Dormant in current configs; single-token identity exact.

Reviewed against the queue: antirez#497 already fixed here (b715001), antirez#472
superseded by the newer bounded-reserve design, antirez#513 not reachable
(expert tiles ship disabled for iq2-down), antirez#460/antirez#528/antirez#504 deferred as
future work.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant