Feature/qwen35 tp fused prefill ggml backend - #103
Merged
Conversation
…ckend
TP >= 2 prefill ran a per-layer loop that round-tripped the [N, hidden]
activation through host memory several times per layer, measuring 2.4x
SLOWER than TP=1 (9B Q8: 612 vs 1459 tok/s at N=512). This adds the
whole-model fused TP prefill: tp_mode in the Qwen3.5 verify kernel builds
one N-token graph per rank over that rank's shards (packed GDN in-proj,
in-graph KV append, MRoPE, expert-parallel MoE, column-parallel LM head)
and returns a segmented plan the existing tp_execute_plans driver runs
with ~2 AllReduce cut points per layer — the same architecture as the
fused TP decode and Gemma4's TP verify.
MoE routing under expert parallelism runs a per-rank top-k over
mask-zeroed router probabilities instead of remapping the global top-k
through an id LUT: ggml's batched mul_mat_id paths (CUDA mm_ids_helper
and the generic fallback) require the ids within one token to be
DISTINCT, and mapping every foreign route to a filler id silently
scrambled expert outputs at N > 1 (fluent but wrong text on the 35B).
The rank-local top-k always yields distinct owned experts; gating
weights keep exactly the global top-k members (prob >= the k-th largest
global prob, read through a DESC argsort since CUDA's TOP_K output is
unsorted), so the ranks' partials sum to the single-GPU MoE.
The per-layer GDN states ride back to the host through the plan's new
extra_out downloads; the driver then invalidates the host-keyed device
copies so the fused decode re-uploads the post-prefill state.
Measured on 2x RTX 2000 Ada, ggml_cuda, prefill/decode tok/s:
Qwen3.5-9B Q8 TP=1: 1469/22.7 TP=2: 612/41 -> 2051/41 (N=512)
TP=2: 2365/40 (N=2048)
Qwen3.5-35B-A3B IQ4 TP=2: ~100/42 -> 1702/42 (N=512), 2122/42 (N=2048)
Output verified against TP=1 / pre-change baselines: short prompts,
2K-token document recall, multi-turn KV reuse, and image+MRoPE
multimodal prefill all match.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Nemotron's TP MoE fed the flattened [seqLen, numExperts] router logits to SelectNemotronTopKExperts as if they were ONE token's row, indexing the router bias out of bounds on any multi-token prefill (TP >= 2 crashed on the first prompt) and applying a single expert set to the whole chunk. Route each token independently and bucket the assignments by expert so each expert runs one batched matmul over its tokens — the same dispatch shape as the gpt-oss TP MoE. The router now also computes once on rank 0 instead of per rank. Qwen3.5 TP prefill: bias the mask-zeroed router probabilities so an owned expert always outranks a foreign one when softmax underflows to exactly zero — a zero-for-zero tie in top_k could select a foreign index and reintroduce the duplicate-ids hazard. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
ggml-vulkan's GET_ROWS asserts zero buffer-offset remainders, so an element-offset view of the argsort output cannot be its id input (the k-th-column gather aborted ggml_vk_build_graph). Gather all sorted probabilities through the whole (aligned) argsort tensor once and take w_min / the top-k normalizer as views of the RESULT — elementwise ops and cont carry offsets fine on every backend. Verified on ggml_vulkan --tp 2 (35B fused prefill active, output matches CUDA) and re-verified on ggml_cuda. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Engine comparison — TensorSharp vs llama.cpp (PR smoke)No report artifact was produced — the benchmark failed before generating results (see the workflow logs). |
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
1 participant
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.
No description provided.