cuda: run streaming selected load for single-token ffn batch encodes#497
Open
giannisanni wants to merge 1 commit into
Open
cuda: run streaming selected load for single-token ffn batch encodes#497giannisanni wants to merge 1 commit into
giannisanni wants to merge 1 commit into
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes the concrete failure behind the
--mtp+--ssd-streamingguard on CUDA (context: #495).metal_graph_cuda_stream_prefill_batch_selected_load()skips whenn_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:The one-line fix lets
n_tokens == 1take 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:
--mtp-draft 2--mtp-draft 4--mtp-draft 2This 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.