Repository navigation
Conversation
added 30 commits
October 4, 2026 14:49
…9 to 32-row windows (MTPLX_NAX_FLASH_WIDE, off by default) MLX falls back to its unfused SDPA for more than 8 query rows at head dim 256, so context-copy verify blocks of 9 to 32 rows re-read the dense KV cache and materialise scores: 112 ms at 16 rows and 131 ms at 24 rows per verify at an 80K prefix. The head-dim-split TensorOps flash kernel now walks groups of 32 query rows on a grid axis inside one launch (q_len 9..32), shipping contract for q_len <= 5 bit-identical, kernel self-check lane and tests included. M5 Pro, stock MLX 0.32.2, Qwen3.8-27B: attention at 80K 2.70/3.69/5.67/6.63 ms per layer at 9/16/24/32 rows (stock 9.83/6.80/7.88/8.45); verify forward at 80K 286 to 164 ms (9 rows), 241 to 190 (16), 336 to 294 (24).
…X_NAX_FLASH_DSPLIT4, off by default) The head-dim-split flash kernel read the KV cache at 180 to 190 GB/s of a 288 GB/s roof at 20K to 80K tokens; with the matrix multiply removed the same loads ran at 248 GB/s, so it was limited by live state per thread. NDH=4 (NSGM x 4 simdgroups, 32 accumulators each, 4-way score exchange in the same 32 KB buffer) raises occupancy: attention at 4 rows 27 to 32% faster (80K: 1.73 to 1.31 ms per layer), whole verify forward -2.3/-4.1/-6.1% at 20K/50K/80K, switch off bit-identical. NDH is a template parameter next to TG_M; the wide route (9 to 32 rows) keeps NDH=2. Kernel self-check lane nax_flash_dsplit4_sdpa, tests, changelog. (cherry picked from commit 18359ef)
… of runtime.make_cache() A restore from RAM builds the cache with the request's cache_factory; the SSD fallback (_restore_cold) ignored it and built the runtime's default layout, the paged KV cache for a long context: about 5x slower prefill and slower decode after every restore from disk. restore() now hands both factories to _restore_cold, which uses them exactly as the RAM restore does (meta state not imposed on the request's cache). Without a factory nothing changes.
…kernel, bank and admission
…he merge; real-kernel test
…ts instead of a gather; prefill route only above 32 rows; route counters per window size
…s storage (MLX copied the strided view on every kernel launch)
…rmat of the same rows (no history gather), guards behind MTPLX_SEGMENTED_KV_SSD
…o on the segment), /health block, SSD counters
…ts hold (nothing merged before: every turn's snapshot holds its segments)
…ith segments too (the fix itself is the base commit)
…attention not hooked, no flash route) keep the stock cache, log one line and report the reason on /health
…defensive lse fallback, tests (kernel tests pending a power-connected run)
…lx-lm Qwen3/Llama attention has no head_dim attribute)
…e hook body on a segmented cache (it would gather the history on every call otherwise); support check requires the layers to be reachable by that body
…ken, no context fence, no MIN_TOKENS); single-segment calls take the stock route on a zero-copy view instead of counting as gathers; several segments take the kernels at any length
…also holds on stock MLX
…ill over segments then gathers the history)
…(no NOMASK parameter); the segment kernel switches the mask off in its own derived source The NOMASK template parameter changed the stock kernel's source and instantiation name without a function: _fused_source replaced the expression anyway. It now replaces the stock `gp <= row_limit` directly by `(nomask_seg || gp <= row_limit)`, so the file the switch-off path runs is byte-identical to upstream. The kernel limits are named constants (MAX_M_ROWS = 32, MAX_Q_LEN = 10) shared with the route gate, and the docstrings no longer point at a gather_attention that does not exist.
…1 and 2 gathered windows of 11 to 32 rows) The sub-window size was 32 // gqa: 32 rows at GQA 1 and 16 at GQA 2, both above the kernel's q_len <= 10, so segments_supported refused and the window gathered the history. It is now min(MAX_Q_LEN, MAX_M_ROWS // gqa). The measured models (GQA 4, 6, 8) were not affected. Tests: the sub-window sizes at GQA 1 to 8 (CPU) and GQA 1 and 2 windows of 12 and 32 rows against an fp32 reference on the kernel (GPU). Docstring of decode_segments_attention describes the current route (1 to 32 rows, sub-windows, when the stock ladder runs).
…runtime load starts without one segmented_kv_enabled() was True while no model had been checked. The verdict is made in configure_split_full_attention, which the Laguna load route skips: such a model got segment caches without the split-attention hook, and its attention gathered every segment on every call. Without a verdict the switch now behaves as off (stock caches, bank and SSD path), /health reports model_not_checked, and runtime.load resets the verdict so a model never inherits the previous load's. The test that encoded the old behaviour now asserts the new one; the other segmented test modules give their caches a supported verdict, as a real load does. The SSD-restore test checks that a request's factory builds segment caches with the switch on.
…, one helper each for unique bytes and env switches, flag-off paths skip the segment walks - Reference counting (KVSegment holders, hold/release, refcount) steered nothing since the merge also merges held segments; removed with its test. Segments are shared by plain references; the module docstring says so. - Unused in production and removed: SegmentedKVState.prefix/segments/__len__, _prefix_refs, SegmentedKVCache.tail_rows/stats() and its per-cache counters (route_counts has them), compact(max_segments=), is_segmented, the _mtplx_segmentable tag. - "Unique bytes over segment id" lives in one place (unique_state_bytes over _unique_segment_bytes), used by the bank's overcount and /health; SegmentedKVCache.gather goes through _gather_refs. - adapt_layer_for_restore(entry, state) -> state; restore_cache reads the switch once per call instead of importing and reading it per layer. - One _env_on() for the module's switches; attention_split asks segmented_kv_requested(); the route counter is public (count_route) and the codec uses its own _count_ssd everywhere; a known segment block and a reuse hit share one block entry in the codec. - Constants instead of literals (VERIFY_WINDOW_MAX in the counter label, MAX_FUSED_SEGMENTS as the default segment limit), and one sentence on why small segmented tensors are concatenated in the codec. - With the switch off, SessionBank.total_nbytes and the admission no longer walk the snapshot trees for segments. - The state getter documents its side effect (it seals the tail). - Tests: the test of the never-shipped MIN_TOKENS setting is gone, the merge test no longer says "unshared", counters are read from route_counts; ruff I001 (attention_split imports, a test) and RUF059 fixed.
…behaviour only (no 32K fence history, windows 1 to 32 rows)
…t-attention hook mlx-lm's llama Attention has no q/k norm, so the hook did not take it and the support check reported attention_not_hooked (Mistral 7B v0.3 kept the stock cache). The hook body now serves an ungated attention with or without q/k norm, on a segmented cache only; with a stock cache the forward is mlx-lm's own, switch on or off. Tests: detection of the three forms, flag-off and stock-cache identity (with and without bias), support verdict, and a GPU test at Mistral shapes (32/8 heads, head_dim 128) through turns and rollback with no gather.
…hes only, the self-check probes the segment kernel Segments only engaged with MTPLX_GQA_PACKED_SDPA and MTPLX_NAX_FLASH_ROUTE on, which only the turbo profile sets; under sustained (the recommended profile of Qwen3.6-35B-A3B and Apodex-1.1-mini) the support check switched them off. Two holes besides: the verdict was made before the kernel self-check, so a self-check that disabled the kernel left segmented caches that gathered the history on every multi-segment call; and under a profile without an attention route the hook sent a segmented cache to the stock forward (a gather per call). - Verdict: neither switch is required; either one set to 0 still switches segments off (gqa_packed_sdpa_off, nax_flash_route_off). nax_unavailable stays (M1 to M4: off). - Hook: a segmented cache always takes the segment route, whatever the profile. - Self-check: runs whenever MTPLX_SEGMENTED_KV is on and probes the segment kernel (lane segmented_kv_sdpa, two segments against stock SDPA) at the model's head_dim and GQA factor. A failure, or a self-check that cannot run, turns the verdict off before the first cache is built: stock caches, reason segment_kernel_selfcheck_failed on /health, one log line. With the switch off nothing changes: every path sits behind segmented_kv_requested() or a segmented cache. Tests: verdict with the switches unset, on and 0; self-check probe shapes, failure (mismatch, raise, bail, no run) to stock caches; hook routing without the switches for gated, q/k-norm and Llama-style attention; flag off leaves the self-check and verdict alone; GPU: the segment kernel serves without the switches, no gather, and the probe passes on the real kernel at head_dim 128 and 256.
… kill switches, self-check on the segment kernel, Qwen3.6-35B-A3B sustained numbers)
…egments instead of repaging Past the dense-decode ceiling (auto: 15 % of RAM over 64 KiB a token, 157,286 tokens on a 64 GB Mac) the auto layout repaged the prefill cache into VllmMetalPagedKVCache, which is built from each layer's ``keys``: on a segment cache that gathers every segment into one paged copy of the history. The admission priced that copy correctly and refused the turn. 2026-10-10, Qwen3.8-27B, 64 GB: a follow-up turn of 161,835 tokens (161,751 restored, 84 new) was priced 11.6 GiB of growth for the repage (repage_copy_bytes = (prompt + 16,387 reserved rows) x 68 KiB) and refused with 10.1 GiB free, eleven retries in a row; the turn before it had decoded from 157,089 to 161,751 tokens on segments at 19 tok/s. With segments on (and KV unquantized) the auto layout is contiguous_dense_decode at any length, so the restore keeps its segments, nothing repages and the admission prices the new rows only (about 1.1 GiB for that turn). The stock cache, quantized KV and an explicit MTPLX_SUSTAINED_PREFILL_LAYOUT keep their repage.
…present, keep return_lse as fallback All fused logsumexp calls go through sdpa_with_lse; the probe records which API answered and /health reports it as lse_api. The per-segment merge and the reference route no longer produce NaN for rows without a visible key (lse -inf): such rows get weight 0 and merge to zeros.
…ce after the merge The reference route rounded every segment's output to the query dtype before the logsumexp merge, and the merge rounded again. On Qwen3.8-27B tensors at 20K (one verify window, 16 layers) that put 26 to 29% of the attention elements one ulp away from the exact result rounded once, against 15 to 22% for the stock kernel and 3.5 to 7% with a single rounding. The per-segment outputs now stay float32 and merge_segment_outputs rounds once to the query dtype (MTPLX_SEGMENTED_KV_PREFILL=lse, reference route; the fused sdpa_with_lse kernel still returns the query dtype per segment). The existing exactness test compared bfloat16 output with a loose tolerance and could not see the second rounding; the new tests require the route to equal the float32 attention rounded once and to be no further from stock SDPA than that exact result is.
…l when the build offers it The fused lse prefill route still rounded twice: the kernel returned each segment's output in the query dtype and merge_segment_outputs rounded again. Our MLX build (fork branch prod-lse-f32out) can return the output in float32 next to the logsumexp (scaled_dot_product_attention(..., return_lse=True, float32_output=True)). sdpa_with_lse takes a float32_output flag, sdpa_lse_float32_output_available() probes the option once per process (a bfloat16 chunk of 16 rows must come back float32), and the kernel route uses it, so the merge rounds once to the query dtype like the reference route. Without the option the behaviour is unchanged. /health reports lse_float32_output.
…_dot_product_attention_with_lse scaled_dot_product_attention_with_lse returns the per-segment output in float32 for bfloat16/float16 queries (and the lse in float32), so the merge rounds once with no probing. sdpa_with_lse prefers that function and falls back to scaled_dot_product_attention(return_lse=True), whose output keeps the query dtype (the earlier behaviour, rounding a second time in the merge). Removes the float32_output keyword and its availability probe. Tests: the route rounds once only with the function (mocked kernels), and a real-kernel test compares the fused route with the same route rounded per segment. The mocked plumbing test hides the function so its keyword fake is used on builds that have it.
…_product_attention_with_lse
…PLIT4, row groups without empty blocks A wide verify window (9 to 32 rows) costs per 16-row block of query rows, not per pass over the KV cache: one group of 64 rows per threadgroup was no faster than two groups of 32, and reusing each loaded K/V fragment for two blocks was no faster either. The quarter split runs each block faster, so the wide route now follows the MTPLX_NAX_FLASH_DSPLIT4 switch (and its self-check lane). Row groups are 32 rows when the 16-row blocks pair up and 16 rows otherwise: an empty block in the last group cost as much as a full one. MTPLX_NAX_FLASH_WIDE_TGM still overrides, capped at 2 with quarters (the score exchange holds 2 x 2 KB per simdgroup). Attention for 16 layers at Qwen3.8-27B shapes, 100K tokens, M5 Pro: 9 rows 56.9 to 47.5 ms, 13 rows 72.8 to 59.4, 17 rows 95.7 to 77.5, 25 rows 118.3 to 105.2. With the switch off the wide route is bit-identical to before. Tests: reference and stock envelope for both splits, the shape policy, quarters within the halves route's distance to stock SDPA.
…) into feat/segmented-kv (youssofal#610) # Conflicts: # CHANGELOG.md
…de verify route The segment kernel is derived from the head-dim-split kernel's source. With the quarter split (MTPLX_NAX_FLASH_DSPLIT4) and the wide route (MTPLX_NAX_FLASH_WIDE) that source takes TG_M (16-row blocks per threadgroup) and NDH (head-dim parts) as template parameters, which the segment launch did not pass, and the segment route split every window past 32 live rows into sub-windows of 5 rows at GQA 6, each a full walk over the keys. The segment launch now takes the shape the contiguous dispatcher picks (_route_shape): halves, or quarters with DSPLIT4=1, for up to 32 live rows per KV head; with WIDE=1 the wide route's split and row groups (wide_shape) for every wider window of up to 32 rows, in one launch, with the matching key-block table. One segment calls the contiguous entry with the same shape. Sub-windows remain for WIDE off or its lane disabled. head_dim 128 keeps halves. Tests: one launch, one segment bit-identical to the contiguous wide kernel, several within bf16 ulps, no further from stock SDPA than the halves route. The two head_dim 128 sub-window tests clear MTPLX_NAX_FLASH_WIDE: they test the sub-window route.
jvmenen
pushed a commit
to jvmenen/MTPLX
that referenced
this pull request
Oct 17, 2026
The segment-kernel coupling with youssofal#593/youssofal#594 is on prod since cb017d5; youssofal#639 now holds it upstream. This brings the docstrings and the two hd128 tests that pin MTPLX_NAX_FLASH_WIDE off to the youssofal#639 + youssofal#630 text, so prod carries exactly what the PRs hold. No code change (AST-identical).
…for the rest Three changes to how a spill of segmented rows finds blocks already on disk: - A merge of two sealed segments keeps both segments' block digests at the shifted rows (the merged rows are a copy). Merges dropped them, so every block of a merged segment was evaluated, copied and hashed again on the next spill, only to be found on disk. - A block a segment names is referenced before any fingerprint is taken, and the fingerprints of segmented rows are taken per run of blocks that need one, the first time one does. With an empty memo a block inside one sealed segment needs none: its segment names it once captured. A session's first spill would otherwise read its whole history on the GPU for fingerprints nothing can match. - A block across two segments is never named by a segment; it is referenced through the incremental memo when one is given (confirmed by known_blob like a segment hit), and a segment remembers a memo hit inside it for the next spill.
This branch has not been deployed
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.
Why
The segment kernel of #610 is not a kernel of its own:
sdpa_segmented._fused_sourcederives it at import from the source of the head-dim-split kernel insdpa_nax_flash_dsplit. #593 and #594 change that source. Row blocks per threadgroup (TG_M) and head-dim parts (NDH) become template parameters, so one body serves halves, quarters and the wide row groups. Each PR is correct on its own, but on a tree with all three:use of undeclared identifier 'TG_M','NDH'). Every launch over two or more segments fails, with or without the switches. The kernel self-check probe of the segment kernel fails, so a model load keeps the stock cache. A call that reaches the kernel raises when it is evaluated. 43 tests intest_sdpa_segmented.pyandtest_segmented_kv_attention.pyfail. One segment still works, because it calls the contiguous entry;gqa * q_len <= 32,q_len <= 10) and runs every wider window as sub-windows, while the contiguous route serves that window in one launch with Wide verify attention: flash kernel for 9 to 32-row windows (MTPLX_NAX_FLASH_WIDE, opt-in) #593. At GQA 6 (Qwen3.8-27B) that is every window past 5 rows: 6 rows run as 5 + 1, each sub-window a full walk over the keys. On the segment route of Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610 at 100K tokens, attention for 16 layers takes 27.7 ms at 5 rows and 52.0 ms at 6 rows. The number of segments costs little next to that: 27.2 ms at 4 rows over one segment, 27.7 over nine, 28.7 over sixteen.How it works, in plain terms
The segment kernel launches with the shape the contiguous dispatcher would pick for the same window, so the kernel the segments use and the kernel a contiguous cache uses stay the same kernel.
_route_shape(q_len, gqa_factor, head_dim)returns(wide, ndh, tg_m), orNonewhen no segment route takes the window. Up to 32 live rows per KV head (and at most 10 rows) take the shipping contract: halves, or quarters withMTPLX_NAX_FLASH_DSPLIT4=1unless the quarter self-check lane failed. WithMTPLX_NAX_FLASH_WIDE=1a wider window of up to 32 rows takes the wide contract:wide_shape()gives the split and the row groups, the same as for a contiguous cache. Its lane being disabled meansNone.TG_MandNDH, sizes the threadgroup as32 * TG_M * NDHthreads and tiles the rows overgrid.y(row groups), as the contiguous wide entry does.segment_block_countsreads the key-block table that belongs to the shape: halves, quarters (_default_dsplit4_blocks) or wide (_default_wide_blocksfor that many live rows), so a list of segments is cut into the blocks the contiguous dispatcher uses for the same total.sdpa_nax_flash_dsplit_wideorsdpa_nax_flash_dsplit), so it stays bit-identical to a contiguous cache.segments_supporteduses_route_shape, sodecode_segments_attentionsends a window to the kernel whole whenever the wide route takes it. Sub-windows remain forMTPLX_NAX_FLASH_WIDEoff or its lane disabled.With both switches off the segment route is the one in #610: the same shapes, the same blocks, the same sub-windows (bit-identical in the comparison below).
What changes
mtplx/kernels/sdpa_segmented.py:_route_shape, a shape-awaresegment_block_counts, theTG_M/NDHtemplate arguments and row groups in the launch, the wide entry for one segment,segments_supportedthrough_route_shape. A refused shape bails asroute_shape(it wasm_rows_gt_32/q_len).mtplx/segmented_kv.py: docstring and comment ofdecode_segments_attentiononly.tests/test_sdpa_segmented.py: the unsupported-shape test now uses a window past the wide route's 32 rows; a new test for 6 to 32-row windows over 1, 3 and 9 segments, with halves and quarters.tests/test_sdpa_segmented_hd128.py: the two sub-window tests clearMTPLX_NAX_FLASH_WIDE, since sub-windows are the route with it off.CHANGELOG.md: one entry.No new switch. The route follows
MTPLX_NAX_FLASH_WIDEandMTPLX_NAX_FLASH_DSPLIT4as the contiguous route does.Results
Attention for 16 full-attention layers at Qwen3.8-27B shapes (24 query heads, 4 KV heads, head_dim 256, bf16) through
decode_segments_attention, on an M5 Pro 64 GB withMTPLX_SEGMENTED_KV=1 MTPLX_NAX_FLASH_WIDE=1 MTPLX_NAX_FLASH_DSPLIT4=1. One process per variant, 25 timed runs after 3 warm-up runs, two passes in reverse order (the first pass shown; in the second, after is within 6 % of it and before is 1 to 13 % slower). "Before" is the merge of #610 and #594 without this commit. "Nine segments" is a realistic layout: one large segment, seven of 8K down to 256 rows, a 640-row tail.Before, a window of 6 to 10 rows took 2 kernel calls per layer, 12 or 13 rows took 3, 16 to 20 rows took 4 and 24 or 25 rows took 5. After, every window takes one. Rows 1 to 5 keep their shape and cost.
Compared with a segment route that already gave 9 to 32-row windows to the wide kernel with halves only, and still split 6 to 8 rows into sub-windows (same bench, two runs per variant, median):
At 32K the pattern is the same: 6 to 8 rows 27 to 29 % less, 9 to 25 rows 11 to 22 % less.
Tests
tests/test_sdpa_segmented.py::test_wide_windows_take_one_launch_within_the_halves_envelope_of_stock(48 cases: 6 to 32 rows, 1, 3 and 9 segments, halves and quarters). The window is one launch. One segment is bit-identical to the contiguous wide kernel and several segments are within 3 bf16 ulps of it. The result is no further from stock SDPA than the contiguous halves route plus one ulp.test_unsupported_shapes_bail_without_a_kernel_launch: 33 rows bail asroute_shape;segments_supportedfollowsMTPLX_NAX_FLASH_WIDE.decode_segments_attentionover 9 layouts (GQA 6 at head_dim 256 with 1 to 15 segments; GQA 4 and 8 at head_dim 128), 1 to 32 rows each, with every combination ofMTPLX_NAX_FLASH_WIDEandMTPLX_NAX_FLASH_DSPLIT4. Compared with the build in which this change was first written (the three PRs combined into one tree, with the same commit), all 288 windows per combination are bit-identical, also with Segmented KV on Macs without tensor units (M1-M4): per-segment attention with a logsumexp merge #630 merged on top. With both switches off that includes the sub-window route of Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610.test_sdpa_nax_flash_dsplit_wide.py,test_sdpa_nax_flash.py,test_sdpa_segmented*.py,test_segmented_kv*.py, the dsplit and wide tests): 828 passed, 1 skipped, on the branch as filed (Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610 including its lse docs and tests commits, merged with Verify attention: quarter head-dim split for 2 to 5 rows (MTPLX_NAX_FLASH_DSPLIT4, opt-in, depends on #593) #594).MTPLX_SEGMENTED_KV=1 MTPLX_NAX_FLASH_WIDE=1 MTPLX_NAX_FLASH_DSPLIT4=1set for the whole run, the same suites pass apart from those two andtest_selfcheck_enabled_gating, which expects no self-check switch in the environment.test_release_pins,test_hc_verify_read,test_qwen4_m4_routed_glu, two intest_qwen4_route_kernel. The two lse test failures seen then are fixed by those Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610 commits.