Feat/ssd streaming 10x - #9
Merged
Merged
Conversation
Qwen3.5-122B-A10B-4bit on M1 Ultra 64GB: - Batch eval: collapse ~1400 per-expert eval calls to ~48/token - Cross-projection batching: orchestrate gate/up/down in SwitchGLU - Eval reduction: skip redundant partitionedLayerCall eval (KV cache forced transitively via lazy dependency chain) - Concurrent pread: DispatchQueue.concurrentPerform for QD=24 NVMe SwitchLayers.swift: - ExpertRange struct for cross-projection expert tracking - allocateExpertBuffers/loadExpertWeights/computeExperts helpers - SwitchGLU cross-projection SSD path with merged eval - concurrentPerform replaces sequential pread (8x3=24 parallel reads) LayerPartitioning.swift: - Skip per-layer eval when isSsdStream (next layer forces via chain) - Remove Stream.gpu.synchronize() for SSD path Benchmark: 200 tokens generation, M1 Ultra 64GB Baseline: 0.58 tok/s (1720ms/tok) This: 4.95 tok/s (202ms/tok) — 8.5x improvement Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Overlaps NVMe I/O with GPU compute using persistent expert buffers and previous-token routing prediction. ~5.04 tok/s (+4% over v1). - Persistent buffers: allocated once per SwitchGLU, reused across tokens - asyncEval(idx) starts GPU work (prev layer compute + attention/router) - Speculative pread loads predicted experts during GPU async window - Only miss experts (~30%) need on-demand pread after routing sync - Cold path (first token) falls back to baseline merged eval
Allows reducing the number of active experts per token at runtime without model recompilation. Benchmarked on Qwen3.5-122B-A10B-4bit: top-k=8 (default): 4.84 tok/s baseline top-k=4: 5.91 tok/s (+22%) top-k=2: 6.52 tok/s (+35%, still coherent) top-k=1: garbage output Recommended: SWIFTLM_TOP_K=4 for best quality/speed tradeoff.
…ybrid models Add DraftModelRef wrapper and speculative generate path to ModelContainer. Enable speculative decoding on hybrid Attention+Mamba architectures by implementing checkpoint-based rollback on MambaCache (recurrent state cannot be trimmed like attention KV caches). ModelContainer.swift: - DraftModelRef: @unchecked Sendable wrapper for LanguageModel refs - extractDraftModel(): extract model ref from loaded container - generate(speculative): routes to SpeculativeTokenIterator via context KVCache.swift: - MambaCache.checkpoint(): deep-copies recurrent state before speculation - MambaCache.trim(n): restores checkpoint on rejection (n>0), clears on full acceptance (n==0) - MambaCache.isTrimmable: now returns true Evaluate.swift: - SpeculativeTokenIterator.speculateRound(): checkpoints MambaCache layers before each draft generation round Tested with Qwen3.5-122B (hybrid Attention+Mamba+MoE). Note: speculative decoding is counterproductive for SSD-streaming MoE models where the bottleneck is I/O not compute. Works correctly for in-RAM models.
6 tasks
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
2 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.
No description provided.