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
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