Implement exact Qwen Sparse Attention and long-context state #97

New Issue

hugo · 2026-09-01T18:24:30Z

hugo commented

2026-09-01 18:24:30 +00:00

Area

Implement exact Qwen Sparse Attention and all long-context cache, rollback, checkpoint, and admission behavior through the native 262,144-token window.

Sequence

Step 5 of 7. Previous: #96. Next: #98. Start only after short-context generation including PLE is correct; #98 relies on this issue's rollback contract.

Implementation

Maintain K/V, raw indexer keys, and ratio-4 pooled blocks for all 12 full-attention layers.
Implement the official four-head query/one-head key indexer, ReLU score aggregation, deterministic tie handling, top-512 block selection, each query's causal partial-block tail, and the 2,048-token attention budget.
Reuse DS4Server's GLM sparse-indexer/top-k, gathered-attention, flash-attention, and checkpoint infrastructure only where semantics match Qwen's selection and 24Q/2KV geometry.
Provide an exact portable Metal path for the macOS 15+ floor. Any Metal 4/NAX large-prefill path must have a real capability/geometry probe and exact fallback.
Reserve QSA capacity for later batched MTP verification without moving live buffers or invalidating saved state.
Snapshot every indexer/QSA component needed for prefix reuse, rollback, session save/restore, and capture-commit.
Extend admission and diagnostics for QSA cache growth through the native context window.

Verification

Prove dense/sparse equality while all history remains within budget and match golden selected block IDs beyond it, including ties and partial blocks.
Test chunked prefill, cache-hit continuation, multi-turn prompts, rollback/replay, and checkpoint/resume around 2,048, 16,384, 64K, 131,072, and 262,144 tokens.
Include regression cases for multi-segment state corruption, recurrent rollback, context-depth throughput collapse, and the native-window boundary failures seen in early upstream implementations.
Record logits, token hashes, prefill/AR throughput, cache bytes, and peak memory at every depth.

Completion requirements

Sparse block selections, logits, and generated tokens match pinned references.
Throughput does not collapse merely because history crosses the sparse-selection threshold.
Rollback and checkpoint restore preserve recurrent, PLE, KV, raw-indexer, and pooled-indexer state together.
The admitted native window either runs correctly or is refused before allocation with an exact memory shortfall.

Non-goals

MTP, context beyond 262,144, KV quantization, multi-host inference, vision, or final speed tuning.

## Area Implement exact Qwen Sparse Attention and all long-context cache, rollback, checkpoint, and admission behavior through the native 262,144-token window. ## Sequence Step 5 of 7. Previous: #96. Next: #98. Start only after short-context generation including PLE is correct; #98 relies on this issue's rollback contract. ## Implementation - Maintain K/V, raw indexer keys, and ratio-4 pooled blocks for all 12 full-attention layers. - Implement the official four-head query/one-head key indexer, ReLU score aggregation, deterministic tie handling, top-512 block selection, each query's causal partial-block tail, and the 2,048-token attention budget. - Reuse DS4Server's GLM sparse-indexer/top-k, gathered-attention, flash-attention, and checkpoint infrastructure only where semantics match Qwen's selection and 24Q/2KV geometry. - Provide an exact portable Metal path for the macOS 15+ floor. Any Metal 4/NAX large-prefill path must have a real capability/geometry probe and exact fallback. - Reserve QSA capacity for later batched MTP verification without moving live buffers or invalidating saved state. - Snapshot every indexer/QSA component needed for prefix reuse, rollback, session save/restore, and capture-commit. - Extend admission and diagnostics for QSA cache growth through the native context window. ## Verification - Prove dense/sparse equality while all history remains within budget and match golden selected block IDs beyond it, including ties and partial blocks. - Test chunked prefill, cache-hit continuation, multi-turn prompts, rollback/replay, and checkpoint/resume around 2,048, 16,384, 64K, 131,072, and 262,144 tokens. - Include regression cases for multi-segment state corruption, recurrent rollback, context-depth throughput collapse, and the native-window boundary failures seen in early upstream implementations. - Record logits, token hashes, prefill/AR throughput, cache bytes, and peak memory at every depth. ## Completion requirements - Sparse block selections, logits, and generated tokens match pinned references. - Throughput does not collapse merely because history crosses the sparse-selection threshold. - Rollback and checkpoint restore preserve recurrent, PLE, KV, raw-indexer, and pooled-indexer state together. - The admitted native window either runs correctly or is refused before allocation with an exact memory shortfall. ## Non-goals - MTP, context beyond 262,144, KV quantization, multi-host inference, vision, or final speed tuning.

hugo added the idea label 2026-09-01 18:24:30 +00:00

~~hugo referenced this issue 2026-09-01 18:27:01 +00:00~~

Prepare exact Qwen3.8 Flash Next Bare Speed artifacts #93

hugo referenced this issue

2026-09-01 18:36:51 +00:00

Implement the Qwen3.8 Flash Next Rust and Metal text core #95

hugo referenced this issue

2026-09-01 18:36:51 +00:00

Stream Qwen3.8 Flash Next PLE from mapped storage #96

hugo referenced this issue

2026-09-01 18:36:52 +00:00

Add Qwen3.8 Flash Next native MTP speculation #98

hugo added enhancement and removed idea labels 2026-09-01 18:42:26 +00:00

hugo referenced this issue

2026-09-03 18:32:28 +00:00

Implement the Qwen3.8 Flash Next Rust and Metal text core #95

hugo referenced this issue

2026-09-03 19:13:12 +00:00

Stream Qwen3.8 Flash Next PLE from mapped storage #96

hugo commented

2026-09-03 19:37:49 +00:00

Implemented and verified in final commit 714b39f (Implement exact Qwen sparse attention). The earlier hash was amended before closure so issue #97 remains exactly one commit.

Implementation:

Added the official Qwen indexer graph in Rust: one affine 640-wide projection split into four 128-dimensional query heads and one raw 128-dimensional key, zero-centered RMS normalization, partial 64-dimensional RoPE, ReLU head aggregation, deterministic top-512 block selection, chronological expansion, and the causal partial-block tail.
Added BF16 raw indexer-key state and ratio-4 pooled BF16 block state to all 12 attention layers. Pooling performs the official float mean, BF16 cast, key normalization, and RoPE at the block first position. Main attention stays dense through 2,048 tokens and uses exactly 2,048 selected tokens plus the 0-3 token causal tail after that.
Added bounded threadgroup attention kernels so attention work remains capped after QSA selection. All orchestration, state ownership, selection policy, checkpoint logic, tests, and admission math are Rust; only permitted data-parallel kernels were added to metal/qwen38.metal.
Extended checkpoint v3, reset, resident session swap, prefix rebuild, rollback/replay, and disk restore to preserve KV, recurrent, PLE, raw-indexer, and pooled-indexer state together. Checkpoints write only live raw rows and completed pooled blocks.
Preallocated KV/raw/pooled capacity for later three-token MTP verification without moving live buffers. Exact 262,144-token engine state is 7,564,812,288 bytes. Admission is 79,335,550,955 bytes without MTP weights and 81,008,126,487 bytes with them, including QSA score and worst-case top-k scratch.

Behavioral review:

DS4 has no Qwen/QSA implementation, so this is a new DS4Server feature. The applicable DS4 sparse-state invariant is preserved: every selectable compressed prefix row is checkpointed, selection is consumed in original causal order, and incomplete causal state is explicit. Qwen math follows the pinned official Qwen/Hugging Face behavior; no foreign C or Python implementation was copied.
Dense and sparse attention are bit-identical when selection covers the visible prefix. Equal index scores deterministically retain lower block IDs. The causal partial tail is appended exactly once.
The existing dense oracle is unchanged: logits SHA-256 89f485fdc9d6c490f9821c3f7c4b20281094947b05325aa5652cdfa43e36a456, samples [6.406557, 2.0828183, -3.0580456, -0.1210103], greedy token 89648.

Pinned logical-depth oracle records (depth: logits SHA-256, greedy token, prefill tok/s, AR tok/s, live attention bytes, peak RSS bytes, checkpoint bytes):

2,048: 4cd0daf267d046d71f23291e756f81b2ebcb549712b441b2d1e8074b41e1da9f, 44,496, 0.601, 0.611, 58,195,968, 735,133,696, 174,839,908.
16,384: f2513cb75775f0704a52b66053e5e47060234db57b3339e9de5a7d1ba6898bae, 44,496, 0.610, 0.616, 465,567,744, 2,018,492,416, 582,269,028.
65,536: f104bc1aa6c0e44fbfff2746300374d7c8edf59a6ee5eb43a5f5469a0219e637, 197,597, 0.612, 0.618, 1,862,270,976, 6,098,698,240, 1,979,168,868.
131,072: 660c5f431a6271748c2e4cb326eed6af21a4e67f3937f71b78e635a25fce2669, 180,094, 0.609, 0.626, 3,724,541,952, 13,487,177,728, 3,841,701,988.
262,144: ead42c990931a88a4196bdc2593013fd471e4f1559e80ccb1779d662e2d08725, 19,559, 0.584, 0.605, 7,449,083,904, 20,488,683,520, 7,566,768,228.

Each depth ran the real Rust-owned model prefill and decode graph, asserted its pinned logits/token oracle, enforced at least 90% of dense-boundary throughput, round-tripped the live KV/raw/pooled frontier through resident session swap, wrote and restored the exact disk checkpoint, and checked its byte size. The full-depth run remained within 4.5% of boundary AR throughput.

Additional isolated verification:

Production tensor bindings executed the first real sparse layer at position 2,048.
Release selection passed at every required depth; full-depth selection took 2.455 ms. Sparse attention at depth 262,143 stayed capped at 2,051 visible tokens and took 51.257 ms.
Full 262,144-token preallocation succeeded on Apple M5 Max with zero swaps.
Raw and pooled QSA state, recurrent state, PLE state, and logits round-tripped checkpoint/reset/replay. Tie, partial-tail, dense/sparse equality, corruption boundary, context boundary, and exact memory-layout regressions pass.
cargo fmt --all -- --check passed.
cargo clippy --all-targets --all-features -- -D warnings passed.
make bundle passed.
cargo test --all-features passed: 236 main tests plus 4 artifact-tool tests; only declared hardware/artifact tests were ignored.
The isolated Metal primitives, full-window benchmark, full-context preallocation, first-sparse-layer, per-depth oracle/checkpoint suite, and pinned model tests were explicitly run and passed.

Implemented and verified in final commit `714b39f` (`Implement exact Qwen sparse attention`). The earlier hash was amended before closure so issue #97 remains exactly one commit. Implementation: - Added the official Qwen indexer graph in Rust: one affine 640-wide projection split into four 128-dimensional query heads and one raw 128-dimensional key, zero-centered RMS normalization, partial 64-dimensional RoPE, ReLU head aggregation, deterministic top-512 block selection, chronological expansion, and the causal partial-block tail. - Added BF16 raw indexer-key state and ratio-4 pooled BF16 block state to all 12 attention layers. Pooling performs the official float mean, BF16 cast, key normalization, and RoPE at the block first position. Main attention stays dense through 2,048 tokens and uses exactly 2,048 selected tokens plus the 0-3 token causal tail after that. - Added bounded threadgroup attention kernels so attention work remains capped after QSA selection. All orchestration, state ownership, selection policy, checkpoint logic, tests, and admission math are Rust; only permitted data-parallel kernels were added to `metal/qwen38.metal`. - Extended checkpoint v3, reset, resident session swap, prefix rebuild, rollback/replay, and disk restore to preserve KV, recurrent, PLE, raw-indexer, and pooled-indexer state together. Checkpoints write only live raw rows and completed pooled blocks. - Preallocated KV/raw/pooled capacity for later three-token MTP verification without moving live buffers. Exact 262,144-token engine state is 7,564,812,288 bytes. Admission is 79,335,550,955 bytes without MTP weights and 81,008,126,487 bytes with them, including QSA score and worst-case top-k scratch. Behavioral review: - DS4 has no Qwen/QSA implementation, so this is a new DS4Server feature. The applicable DS4 sparse-state invariant is preserved: every selectable compressed prefix row is checkpointed, selection is consumed in original causal order, and incomplete causal state is explicit. Qwen math follows the pinned official Qwen/Hugging Face behavior; no foreign C or Python implementation was copied. - Dense and sparse attention are bit-identical when selection covers the visible prefix. Equal index scores deterministically retain lower block IDs. The causal partial tail is appended exactly once. - The existing dense oracle is unchanged: logits SHA-256 `89f485fdc9d6c490f9821c3f7c4b20281094947b05325aa5652cdfa43e36a456`, samples `[6.406557, 2.0828183, -3.0580456, -0.1210103]`, greedy token `89648`. Pinned logical-depth oracle records (depth: logits SHA-256, greedy token, prefill tok/s, AR tok/s, live attention bytes, peak RSS bytes, checkpoint bytes): - 2,048: `4cd0daf267d046d71f23291e756f81b2ebcb549712b441b2d1e8074b41e1da9f`, 44,496, 0.601, 0.611, 58,195,968, 735,133,696, 174,839,908. - 16,384: `f2513cb75775f0704a52b66053e5e47060234db57b3339e9de5a7d1ba6898bae`, 44,496, 0.610, 0.616, 465,567,744, 2,018,492,416, 582,269,028. - 65,536: `f104bc1aa6c0e44fbfff2746300374d7c8edf59a6ee5eb43a5f5469a0219e637`, 197,597, 0.612, 0.618, 1,862,270,976, 6,098,698,240, 1,979,168,868. - 131,072: `660c5f431a6271748c2e4cb326eed6af21a4e67f3937f71b78e635a25fce2669`, 180,094, 0.609, 0.626, 3,724,541,952, 13,487,177,728, 3,841,701,988. - 262,144: `ead42c990931a88a4196bdc2593013fd471e4f1559e80ccb1779d662e2d08725`, 19,559, 0.584, 0.605, 7,449,083,904, 20,488,683,520, 7,566,768,228. Each depth ran the real Rust-owned model prefill and decode graph, asserted its pinned logits/token oracle, enforced at least 90% of dense-boundary throughput, round-tripped the live KV/raw/pooled frontier through resident session swap, wrote and restored the exact disk checkpoint, and checked its byte size. The full-depth run remained within 4.5% of boundary AR throughput. Additional isolated verification: - Production tensor bindings executed the first real sparse layer at position 2,048. - Release selection passed at every required depth; full-depth selection took 2.455 ms. Sparse attention at depth 262,143 stayed capped at 2,051 visible tokens and took 51.257 ms. - Full 262,144-token preallocation succeeded on Apple M5 Max with zero swaps. - Raw and pooled QSA state, recurrent state, PLE state, and logits round-tripped checkpoint/reset/replay. Tie, partial-tail, dense/sparse equality, corruption boundary, context boundary, and exact memory-layout regressions pass. - `cargo fmt --all -- --check` passed. - `cargo clippy --all-targets --all-features -- -D warnings` passed. - `make bundle` passed. - `cargo test --all-features` passed: 236 main tests plus 4 artifact-tool tests; only declared hardware/artifact tests were ignored. - The isolated Metal primitives, full-window benchmark, full-context preallocation, first-sparse-layer, per-depth oracle/checkpoint suite, and pinned model tests were explicitly run and passed.

hugo closed this issue

2026-09-03 19:37:53 +00:00

hugo reopened this issue

2026-09-03 19:38:07 +00:00

hugo closed this issue

2026-09-03 19:46:53 +00:00

Sign in to join this conversation.

1 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: hugo/DS4Server#97