Conversation
The per-batch output offset is computed in int32. For a large f32 destination it wraps once a batch base reaches 2^31 elements, storing that batch ~8.59 GiB below the tensor and leaving its own region unwritten. Upstream: ggml-org#28210.
- add kernel_flash_attn_ext_vec_idx: compacts finite mask entries into a per-row index list (Hillis-Steele scan, one threadgroup per row) - extend vec FA kernel with optional sparse index gathering (FC slot 5) - add host-side gate: sparse path when n_kv_max > 0, mask present, supported head sizes / KV types, n_kv_max <= 4096 - new buffer region extra_idx for the index list - pipeline getter extended with has_sparse param - add test cases: head sizes, quant types, nb>1, nr23 variants, sinks, ALiBi, softcap, permute, v_view_of_k, no-mask fallback Note: multi-row (nb*nr23[1] > 1) cases still failing - rid mapping in the store phase needs revisiting for the sparse path.
- kernel_flash_attn_ext_vec_idx: mask param is half* but nb31 is a byte stride, so the per-row mask offset was scaled by 2x; cast to char* before applying the byte strides - kernel_flash_attn_ext_vec: sparse pidx param is char* so the per-row element offset was under-scaled by sizeof(int); scale it by sizeof(int) to get the correct byte offset - fixes the multi-row (nb*nr23[1] > 1) sparse flash attention failures
The idx kernel previously read the mask row twice: once to count the finite entries (for the prefix scan) and again to recover their positions. Since the kernel is memory-bound, this doubled the mask traffic. Keep the finite positions in a per-thread register array during the count pass and write them out directly, avoiding the second mask read. A dense mask with more than NLOCAL finite entries in a slice falls back to re-reading the mask to write the remaining positions.
Measure the sparse vec FA kernel across KV sizes, n_kv_max hints and batch
sizes. Run with:
./build/bin/test-backend-ops -b MTL0 -o FLASH_ATTN_EXT -p "n_kv_max=[1-9]" perf
Add an optional fast path for producers that can fill a host-visible tensor in place. Metal shared buffers expose their CPU-writable address while private buffers and other backends retain the safe staging fallback.
Explicit-value signal/wait plus a CPU-side signal and a notification callback. Together these let the GPU hand work to a CPU servicer and resume without the graph splitting: the GPU signals, a listener callback runs on the CPU, and the CPU signals back. Round-trip latency is what decides whether that beats letting ggml_backend_sched split the graph, so measure it before relying on it. Notifications run on one private serial queue; a servicer must not block it.
Measurement only. With every switch unset the encode path is byte-for-byte the usual one - the env parse is a pthread_once, and the profiler pointer is null. - GGML_METAL_KPROF=<stride>: split the encoder every <stride> non-noop nodes and sample the GPU timestamp counter at each segment boundary, emitting one JSONL record per segment. The node map is keyed by graph shape, so equal graphs deduplicate and a stride sweep calibrates the split's own overhead. - GGML_METAL_GPU_PROFILE: total GPU busy time per context from Metal's own GPUStartTime/GPUEndTime. One llama_context gets one Metal context, so under speculative decoding the target and draft report separately. - GGML_METAL_DEBUG_GROUPS: name every dispatch after its ggml op so an external profiler can attribute GPU time per op without starting a capture, which would conflict with a profiler already recording. - GGML_METAL_SMEM_PAD: pad a threadgroup allocation without the kernel using the bytes, to probe occupancy. Arithmetic is untouched, so output is unchanged. - examples/moe-routing: capture the routed experts per layer per token.
The FA guard was `#pragma unroll (MIN(DK8/2, 4*NSG))`, but NSG is a template
parameter and the host picks nsg=8 for DK=512, so the MIN folds to 32 - a 32x unroll
of the DK loop, measured 3.97x slower than a 4x unroll on M1 Max.
Index the MXFP4 dequant table by the whole BYTE rather than by each nibble: one load
yields both the low and the high nibble, so 16 scalar lookups per iteration become 8
half2 lookups. half is exact for E2M1 - every value in {0,.5,1,1.5,2,3,4,6} is
representable - so the result stays bit-identical to the f32 table. The table grows
from 128 B to 1 KiB, and the threadgroup allocation with it.
Read each four-byte quant vector as two naturally aligned halfwords before reconstructing the signed lanes. This avoids conservative byte-aligned loads for alternating 34-byte q8_0 blocks.
Keeps the routed expert weights on disk and pages them into a fixed-size per-layer cache on demand, so a model whose experts do not fit in RAM still runs at a useful rate. Based on PR ggml-org#25294 with Metal/Apple adaptations. --moe-stream enable --moe-stream-cache budget in GiB, or exact slots per layer with an 's' --moe-stream-io-threads reader threads --moe-stream-direct O_DIRECT reads, falling back to buffered Rebase note: upstream took TENSOR_READ_LAZY on bit 5 for its own on-demand row reads, so TENSOR_STREAMED moves to bit 6. The two are not equivalent - READ_LAZY reads rows through mmap, STREAMED means moe-stream owns the tensor's I/O entirely and the loader must not allocate or read it.
Squashes the original partitioning change with its two follow-up fixes: flooring the pair chunk at n_tokens, and splitting a hot expert across waves instead of aborting the request. Measured: +37% prefill. Wave count is nearly free once partitioned, so the pair imbalance rather than the wave count is what bounds the chunk.
One work item per weight SLAB rather than per expert. A miss read its 2-3 slabs sequentially on one thread, so the device saw queue depth 1 even though the reads are independent: measured 1.00 ms/slab read against 0.065 ms/slab upload, i.e. 94% of a miss is the read, serialised at the drive's QD1 rate (~2.9 GB/s against 7.3 GB/s at QD8). Issuing them together is what lifts the depth. Halve route hotness every 1024 tokens, not 64. At 64 an expert accumulates only ~1.5 uses between halvings (256 experts, 6 per token), so counters sit at 0-3, cannot rank experts, and eviction degenerates into plain LRU. Measured on decode, three runs each: 64 -> 7.54/7.61/7.82 t/s at ~6.8% miss, 1024 -> 9.06/8.85/9.11 t/s at ~5.4%.
Two prefetch paths, both output-identical - they change when an expert is read, never which expert runs. Verified byte-identical over 60 greedy tokens each. hash layers: deepseek4 sets hash_layer_count=3, so layers 0-2 pick experts by a token-id lookup rather than a router matmul. Their loads can start before layer 0. Worth its own path: those 3 of 43 layers are 36% of steady-state decode misses, because a token-id hash has no locality to gain as the cache warms. lookahead (LLAMA_MOE_STREAM_LOOKAHEAD=K): at layer L predict layer L+1's routing from L's router input and start those loads a layer early. Rides the existing remap op as a second src, so it adds no graph split. The prediction skips the attention and FFN terms between the layers - the exact version would need layer L+1's attention output, i.e. attention twice per layer. misses/token 17.2 -> 8.2, stall 23.2% -> 16.1%, decode +9.5% (3/3 pairs), and +9.3% with speculative decoding enabled. Also here: read-latency histogram for the streaming worker (the mean cannot tell a page-cache hit from an SSD read), reads that land straight in the expert cache when the backend offers a host pointer, and LLAMA_MOE_STREAM_TAIL_HOT (default 0, measured worse - kept as a knob with its numbers in the comment).
llama-bench parses its own arguments rather than going through common/arg.cpp, so it had no way to enable expert streaming. Without it the tool cannot load a model whose experts do not fit in RAM - which is exactly the configuration worth benchmarking on a memory-constrained machine. Adds --moe-stream, --moe-stream-cache <GiB> (implies --moe-stream) and --moe-stream-io-threads, plumbed into llama_model_params.
Move residency off the graph's critical path. A Metal kernel resolves hits from a device-side slot table, evicts by least-recent use for misses, and hands the slot list to a CPU servicer over a shared event - so decode has no CPU custom op and no graph split for the remap (~34 us against ~930 us). Staged behind LLAMA_MOE_STREAM_GPU_SLOT: mode 1 verifies the kernel against the CPU remap while the CPU answer is still what reaches the GEMM, mode 2 publishes the table without inserting a node (isolating the buffer's cost from the kernel's), and mode 3 gives the GPU ownership. Only for small ubatches - the kernel resolves the pair list serially in one thread, which is right for a handful of pairs and wrong for a 4096-token prefill. Also hardens the asynchronous load path: a slot is published only when its LAST slab lands, so no consumer can observe a half-filled expert, and in-flight loads for an evicted occupant are recognised as stale by a per-slot generation counter.
Qwen3.8's PLE table is far larger than RAM but each ubatch touches only a compact set of rows. Unlike an expert tensor it needs no persistent slot cache: the raw quantised rows are read directly into one compact graph input and the normal get_rows op dequantises them. Reads are deduplicated (a repeated row is read once and memcpy'd), sorted by file offset so the reader and the drive see a less hostile pattern during large prefills, and issued in parallel across the existing expert readers with their own queue and completion count. Where the backend exposes a host pointer the rows land straight in the graph input with no staging copy.
Partitioning gives each (token, expert) pair to one wave instead of running every wave over every pair and masking the rest. It measured about 37 percent prompt processing on the streamed models, and it already self-disables where it cannot pay: the graph skips it for single-wave ubatches, so decode is unchanged. The old check tested only for the variable's presence, so LLAMA_MOE_STREAM_PARTITION=0 turned partitioning ON. Parse the value, matching LLAMA_MOE_STREAM_LRU in the same file. Belongs in the MoE layer next to the partitioning commit; kept here to avoid a conflicting insert.
Three fixes to the pair-partition planner. The split pass could hard-abort. A split moves a hot expert's surplus pairs into another wave, which needs pair room AND a free expert slot there. Pair room always exists (n_waves*chunk >= n_pairs); slot room does not - at cap 52 with 256 experts, ceil(256/52) = 5 waves leaves 5*52 - 256 = 4 spare slots for the whole layer, and exhausting them killed the server mid-request. Size the wave count from cap-1 when partitioning, which takes that to 6 waves and 56 spare slots and is near-free: the chunk tracks the mean, so n_waves*chunk stays ~1.5*n_pairs whatever the count. The receiver search had the test BACKWARDS - it skipped waves that already stage the expert, which are exactly the waves needing no extra slot, because the pairs land on a slot that wave has already reserved. Prefer those; require a free slot only when the wave does not already stage it. LLAMA_MOE_STREAM_WAVE_CAP was silently overridden: the partition wave-budget branch recomputed stream_wave_cap unconditionally, so a sweep could run at a cap other than the one asked for - and only for the smaller ubatches, which corrupts PART of a sweep and is harder to spot than corrupting all of it. cap_forced now makes the override win outright, with a WARN naming the pairs/wave floor it violates. The abort, if it is ever reached, now reports which constraint bound. LLAMA_MOE_WAVE_SLACK=0 disables the cap-1 sizing for A/B. Cost measured at +1.09% (0.59 sd, n=4 interleaved) - indistinguishable from zero.
Tie KPROF batches to their owning command buffers and graph maps, wait for profile callbacks during teardown, validate profiling knobs, and propagate routing-capture I/O failures.
GGML_METAL_GPU_PROFILE and GGML_METAL_DEBUG_GROUPS were presence tests, so =0 ENABLED them - the same footgun already fixed in LLAMA_MOE_STREAM_PARTITION, where it silently inverted any A/B that tried to turn the feature off. Both now go through ggml_metal_positive_env(), which existed but was file-static; it is declared in ggml-metal-device.h with a comment saying why a presence test is not acceptable, so the next switch has one obvious thing to call.
Share the selected CSA rows across blocks of eight queries while preserving each query's exact membership, so the attention kernel walks one deduplicated row list per block instead of a dense masked row space. ggml_union_build dedups a block's top-k selections into an ASCENDING list, packing the row id in the low 24 bits and an 8-bit membership mask in the high byte. A threadgroup bitmap yields that order with no sort, and a row's union index is its rank in the bitmap. ggml_flash_attn_union then attends the dense prefix plus exactly the entries whose membership bit is set - exact, not an approximation. On by default where the shape qualifies. Includes a CPU reference for both ops and their tests; a device-memory fence ordering the output zeroing against the atomic ORs that follow (a threadgroup-only fence lets a zero land after another lane's OR and silently drop a selection, changing attention output non-deterministically per run); the GQA head mapping fix in the reference; and tensor/tail hardening.
kernel_union_build held the whole CSA row space in a fixed 2048-word threadgroup bitmap, so n_csa above 65536 fell back to dense attention - silently, with no log line, at exactly the contexts union-8 was built for. Walk the row space in chunks of 65536 instead, carrying the union offset across chunks so ids stay globally ascending; empty chunks skip their prefix sum and second pass. Threadgroup memory is unchanged at 16 KB, so occupancy is unchanged. The only bound left is the 24 bits an id is packed into. supports_op carried the same 65536 cap, and saying no THERE does not disable the path - it moves union_build to the CPU backend and adds a graph split. Both raised. A one-time WARN now names the reason when the union path declines, suppressed for the two expected cases (decode, short context). Tests: UNION_BUILD at n_csa 65535/65536/65537/131072 straddling the chunk boundary, and FLASH_ATTN_UNION at kv 65536 and 131072. 10/10 and 13/13 against the CPU reference.
Use ggml_top_k for expert routing instead of argsort_top_k - a native contiguous top-k rather than a full sort of every expert score. Metal has pipelines for it, so there is no CPU fallback and no extra graph split; caller-provided non-contiguous selections stay valid for MoE streaming. Keep the hyper-connection mix copy-free. The grouped RMSNorm scaled after a reshape, which put a reshape between the norm and its consumer and blocked the Metal RMS_NORM -> MUL fusion; scale the norm output directly and reshape the gamma instead. Collapse the streams by starting the ADD chain from a stream view rather than a cont of it - a stream view already has contiguous rows, which ADD accepts - which also leaves the chain eligible for fusion. Select the sparse indexer blocks directly: scores and bias are already block granular, so selecting there avoids expanding an [n_blocks, n_tps] surface to [n_kv, n_tps] before the top-k.
Keep the raw indexer K cache in f16 so quantization cannot change discrete top-k block selection. Preserve the requested type for the unused V side to reduce the memory increase, and provide LLAMA_QWEN4EXP_INDEXER_F16=0 for A/B.
The model ships a NextN block: its own fc_embd/fc_hidden projections plus a hyper-connection norm, reusing the trunk's LM head. Supporting it needs the trunk to export res->t_h_nextn - the WIDE pre-mixer hyper-connection stream - with an explicit ggml_build_forward_expand. The combiner must run per hyper-connection stream on the wide hidden state; mean-pooling first makes the head read a hidden state that is never computed and acceptance collapses to ~1% SILENTLY. An mtp- sidecar carries only the NextN block plus embeddings, so the trunk tensors and the PLE table are optional when one is loaded. Recurrent-state rollback comes with it: qwen4exp shares the delta-net base and build_rs path with qwen35, and its extra caches already mirror rollback - the indexer cache is addressed by the attention cells, and the conv state writes every ring bank.
The block path returned cache-cell ids as blk*r + k, assuming block b owns cells [b*r, b*r+r). set_input_qsa does not produce that: it hands out COMPACTED block ids and records members as blk_cells[blk*r + k] = cell, so blk*r + k indexes that table, not the cell array - which is what the header of llama-memory-hybrid-idx.h warns about (blocks cut the position line, not the cell array). The two coincide only when every live cell sits at its own position: one sequence, no holes, never shifted. They diverge after a context shift (seq_add reaches this cache - hparams_idx.rope_type is NONE precisely because K-shift runs on it), after a partial seq_rm leaves holes, with more than one sequence in a stream, and always on the ranked mrope ordering. Nothing asserted it, and the failure is silent: wrong cells get unmasked and the model still writes fluent text. Gather through the table at BLOCK granularity - rows of r ids - which also drops the f32 detour the arithmetic needed (an [r, n_blk_sel, n_tps, n_stream] repeat of an arange, two casts and an add). set_input_qsa now also fills the SPARE block's member list with the unpooled tail cells, at their own position slot. The arithmetic reached those only by the same identity coincidence, so without this the fix would drop the newest tokens from the selection entirely. LLAMA_QSA_GATHER=0 restores the arithmetic for A/B. Cost measured at +1.09% (0.59 sd, n=4 interleaved) - indistinguishable from zero.
Enable native MTP rollback for qwen4exp, and bound what a draft context costs. - the draft inherits the target's expert-cache budget, which for a small MoE drafter can accidentally cover every expert, disable streaming, and consume memory the target needs. Pin the draft's slots instead. - DFlash and DSpark chunk the target's prompt batches into their own ubatches, so the draft context needs only n_seq*(n_max+1) rows rather than the target's batch buffers. Not applied to draft-simple, whose process() forwards the target batch and must keep the target's acceptance bound. - MTP must accept the target's complete logical batch, but llama_decode may split it into smaller physical batches, so cap the draft ubatch independently of the target. - --spec-max-prompt disables speculation entirely above a prompt length. A draft head is an extra layer over the WHOLE prompt, so its cost scales with context while its per-token saving does not; the threshold is really a bet on answer length. Context checkpoints are server state and are not part of llama_state_seq_save_file(), so a saved slot keeps them in a best-effort sidecar and can resume near the end of its prompt after a restart instead of reprocessing from position zero.
Pass an explicit argument to the variadic server log macro so the saved-slot checkpoint warning builds under the configured C++ standard.
slot_ckpt_read_blob checked a declared blob size only against SLOT_CKPT_MAX_BLOB (4 GiB) before blob.resize(size), so a sidecar truncated by a power loss, a full disk or a killed server could commit a 4 GiB allocation before the read discovered the bytes were not there. The throw is caught, so not a crash - but a 4 GiB transient on a box sitting at ~50 GiB wired starts paging, and a truncated sidecar is exactly what those situations leave behind. Pass the remaining byte count down and reject any size past it, recomputed per blob since each read tightens the bound. A malformed sidecar now fails at the header instead of at the allocator.
The perf suite covered 4096x14336 only, which is neither model's expert shape, so the format choice for a checkpoint could not be measured from it. Add MUL_MAT cases at Qwen3.8-Flash-Next's 2560 -> 640 and DeepSeek-V4-Flash's 4096 -> 2048 expert projections, across every quant type the checkpoints use plus reference formats, at n=512 (a wave's rows for one expert) and n=1. Ranked by time per EFFECTIVE bit-per-weight these invert the intuitive order: the i-quant kernels are occupancy-bound on this GPU, so IQ3_XXS is last in all four tables despite reading the fewest bytes, and at n=1 it is slower in absolute terms than F16. This is what the checkpoint tables in the research logs are generated from.
Rewrite the README and both research logs around what is reproducible. Each log now carries the checkpoint's full tensor table read from the GGUFs, side by side with the reference quant it is derived from, so the custom mixes are shown rather than described: Qwen's splice differs from its base in exactly two rows and is 0.9 GiB smaller, and DeepSeek's is 10.3 GiB larger than the previous default almost entirely in ffn_gate/up_exps. Each also carries a kernel survey - every quant format measured at that model's own expert-GEMM shape, PP-like and TG-like, ranked by time per effective bit-per-weight. That ranking is the argument for both custom mixes and it does not follow bit width. Model throughput tables are removed. The env-var table in the README replaces them and records that pair partitioning and lookahead are now on by default.
github-actions
Bot
added
the
documentation
Owner
|
Can you please add the decode number as well? |
mihailescu2m
force-pushed
the
master
branch
2 times, most recently
from
September 15, 2026 01:04
28b8edb to
d1762fc
Compare
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
5 participants
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.
Thanks for this fork and the write-up — it's what made the model runnable here at all.
Adding one measurement from a different operating point, in case it's useful to others: on an M1 Max 64 GB with the experts fully resident (no
--moe-stream), prefill peaks at-ub 1024rather than 4096. Without a streaming budget bounding the experts, the larger ubatch takes its compute buffer out of the page cache — the same one the 90-byte PLE rows are read through, so it seemed to belong next to that section.Docs only. Single run per configuration on one machine, so indicative rather than rigorous. Happy to reword or drop it if it doesn't fit the log's shape.
🤖 Generated with Claude Code