feat: add DFlash speculative decoding

[0xClandestine]

Summary

• Adds DFlash speculative decoding runtime (Sources/DFlash/) with FullAttentionEngine and HybridGDNEngine
• Implements DFlashTargetModel conformance for Qwen3NextModel (27B) and Qwen35Model (35B)
• Adds MambaSnapshotCache as a lightweight O(1) rollback alternative to RecurrentRollbackCache, selectable via dflashUseTapeRollback
• Refactors Metal kernels: branchless mask semantics via metal::select, 2D kernel cache array
• Adds DFlashKernelBench micro-benchmark target
• Reorganizes DFlash test suite into tests/DFlash/
• Adds JSON result export to bench_35b.sh; adds bench_coder_next.sh

Submodule dependency

Requires a 3-line change in SharpAI/mlx-swift-lm (commit a707519): adds public to Qwen3NextModelInner.embedTokens, Qwen3NextModel.lmHead, and Qwen3NextModelInner.callCapturing. This allows DFlashTargetModel conformance to live in Sources/SwiftLM/Qwen3Next+DFlash.swift (same pattern as Qwen35+DFlash.swift) rather than inside the submodule.

The submodule pointer in this PR references that commit on SharpAI/mlx-swift-lm.

Test plan

• swift build -c release succeeds
• Baseline inference unchanged: ./SwiftLM --model <path> --port 5414
• DFlash inference: ./SwiftLM --model <27B> --draft-model <draft> --dflash --port 5414
• SSD+DFlash: ./SwiftLM --model <35B> --stream-experts --dflash --draft-model <draft> --port 5414
• Run bench_35b.sh and verify tok/s matches benchmarks above

[Copilot]

Pull request overview

Adds a new DFlash speculative decoding runtime and integrates it into the SwiftLM server, along with kernel benchmarking and cross-language intermediate-dump tooling to validate numerical parity.

Changes:

• Introduces Sources/DFlash/ runtime (engines, draft model/backend, rollback caches, kernels, registry, dumper) and a DFlash library product.
• Integrates DFlash into SwiftLM via --dflash and draft auto-resolution/loading; adds timing fields to responses and new benchmark scripts.
• Adds kernel micro-benchmark target (DFlashKernelBench) and a reorganized DFlash test/tooling suite under tests/DFlash/.

Reviewed changes

Copilot reviewed 25 out of 26 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
tests/DFlash/dump_python_intermediates.py	Python-side reference dump of intermediates for Swift↔Python comparison.
tests/DFlash/compare_swift_python.py	Compares Swift dumps vs Python dumps using cosine similarity.
tests/DFlash/compare_cosine.py	Self-consistency and “Swift-equivalent” Python path comparison tooling.
tests/DFlash/README.md	Documentation for DFlash benchmarking/comparison tools.
tests/DFlash/DFlashProfiler.swift	Swift profiler for kernel performance and basic correctness checks.
tests/DFlash/DFlashCosSimComparison.swift	Swift tool to load .npy and compare/inspect intermediates (partial WIP).
tests/DFlash/DFlashBenchmark.swift	End-to-end benchmark harness for baseline vs DFlash performance (tooling).
bench_coder_next.sh	Benchmark script for Qwen3-Coder-Next across baseline/SSD/DFlash configs.
bench_35b.sh	Benchmark script for 35B model; adds rich JSON export for downstream tooling.
Sources/SwiftLM/Server.swift	Adds --dflash path, draft auto-resolution/loading, and timing fields in responses.
Sources/SwiftLM/Qwen3Next+DFlash.swift	Adds DFlashTargetModel conformance for Qwen3NextModel.
Sources/SwiftLM/Qwen35+DFlash.swift	Adds DFlashTargetModel conformance for Qwen35 models.
Sources/DFlashKernelBench/main.swift	New micro-benchmark executable for DFlash Metal kernels (trace-friendly).
Sources/DFlash/RecurrentRollbackCache.swift	Adds recurrent tape rollback cache + snapshot rollback alternative.
Sources/DFlash/DFlashRuntime.swift	Core DFlash generation loop + cache management + token utilities.
Sources/DFlash/DFlashKernelsOptimized.swift	Alternative optimized kernels implementation (currently excluded from build).
Sources/DFlash/DFlashKernels.swift	Main kernel implementations (tape replay, gated-delta+tape, SDPA 2-pass).
Sources/DFlash/DFlashIntermediateDumper.swift	Writes Swift intermediates to .npy for Python tooling.
Sources/DFlash/DFlashEngine.swift	Defines FullAttentionEngine and HybridGDNEngine rollback behavior.
Sources/DFlash/DFlashDraftRegistry.swift	Maps target model refs to draft model refs for auto-resolution.
Sources/DFlash/DFlashDraftModel.swift	Implements the DFlash draft model architecture and context feature extraction.
Sources/DFlash/DFlashDraftBackend.swift	Implements greedy drafting logic using target embed/lm_head and draft model.
Package.swift	Adds DFlash library product and DFlashKernelBench executable target.
.gitignore	Ignores generated intermediate dump directory.

Comments suppressed due to low confidence (1)

Sources/SwiftLM/Server.swift:1013

• This extension ModelContainer block is empty (only doc comments), so it adds no functionality and can confuse readers about missing API. Either implement the intended helper (e.g. extractDFlashTargetModel()) or remove the empty extension.

        }
        // Use the most recently modified snapshot
        let sorted = snapshots
            .filter { (try? $0.resourceValues(forKeys: [.isDirectoryKey]).isDirectory) == true }

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[0xClandestine]

@solderzzc thanks for the review, will post benchmarks from m3 ultra soon

[solderzzc]

@solderzzc thanks for the review, will post benchmarks from m3 ultra soon

@0xClandestine Thanks for your PR, I'm working on setting up Github Action for test automation. And collecting benchmarks on my M5 Pro 64GB.

[hankbobtheresearchoor]

🐛 Bugs & Required Changes — feat/add-dflash branch

Pulled this branch down onto an M3 Ultra (275 GB RAM) to benchmark Kimi K2.5 (384-expert MoE, 612 GB 4-bit) with --dflash --stream-experts --ssd-prefetch. Hit a chain of issues before I could get a single clean inference. Documenting everything here so it's visible.

---

1. language_model. weight prefix crash

Model: Kimi K2.5 safetensors ship with a language_model. prefix on every key (e.g. language_model.model.norm.weight). The DeepseekV3 configuration doesn't strip this in sanitize(), so quantize() / update(parameters:) fails with Key model.norm.weight not found.

Fix: Patched DeepseekV3.sanitize() to strip the language_model. prefix:

weights = Dictionary(uniqueKeysWithValues: weights.map { k, v in
    (k.hasPrefix("language_model.") ? String(k.dropFirst("language_model.".count)) : k, v)
})

This is standard for models exported from HuggingFace with a language_model wrapper.

---

2. quantize() crash — mixed dense/MoE layers

Root cause: Layer 0 uses a dense DeepseekV3MLP (bf16, no .scales), while layers 1–60 use MoE SwitchMLP (4-bit quantized, has .scales). When quantize() runs, ModuleChildren.unflattened() creates .none entries for array indices that have no quantize updates (layer 0's MLP). The default case in update(modules:) then throws unexpectedStructure(key: "layers").

Fix: Patched Module.swift update(modules:) to skip .none entries in the values array instead of throwing:

for (keyPart, v) in zip(keys, values) {
    if case .none = v { continue }  // skip layers with no updates
    // ... existing update logic
}

---

3. Weight verification fails on mixed-precision models

Load.swift calls verify: [.all] after quantize, which fails when some layers are quantized and others aren't (layer 0 dense vs rest MoE).

Fix: Changed verify: [.all] → verify: [] in Load.swift. Minimal change, skips verification that doesn't account for hybrid precision.

---

4. Tiktoken tokenizer not supported

Kimi K2.5 uses TikTokenTokenizer (tiktoken BPE), but swift-transformers only supports PreTrainedTokenizer. The server crashes with "The tokenizer type 'TikTokenTokenizer' is not supported".

Workaround: Converted tiktoken.model → tokenizer.json using the tokenizers Python library, then changed tokenizer_config.json tokenizer_class to PreTrainedTokenizer. However, the converted tokenizer produces garbled output — the ByteFallback decoder doesn't properly handle tiktoken's BPE encoding. This needs a proper tiktoken decoder implementation in swift-transformers.

---

5. File descriptor exhaustion (ulimit -n 256)

With 182 safetensor files mmap'd for the main model + draft model files + metallib + dylibs, the process exceeds the default macOS FD limit of 256. When pread_into tries to open() another file for SSD streaming, it gets Cannot open: ...model-00033-of-00182.safetensors — not because the file doesn't exist, but because the process hit the FD ceiling.

Fix: Start with ulimit -n 1024. The server should probably set this programmatically or at least document the requirement.

---

6. DFlash draft model fails to load

The DFlash draft model at the snapshot path has a model.safetensors.index.json pointing to sharded files (e.g. model-00001-of-00002.safetensors), but the loader tries to open a single model.safetensors which doesn't exist. Error:

Failed to open file .../Kimi-K2.5-DFlash/.../model.safetensors

The safetensors loader needs to handle the sharded case for draft models.

---

7. DeepseekV3 doesn't conform to DFlashTargetModel

Server logs: ⚠️ DFlash enabled but target model does NOT conform to DFlashTargetModel. The PR adds DFlash conformances for Llama, Qwen3, Qwen3MoE, Qwen3Next — but not DeepseekV3. Since Kimi K2.5 is deepseek_v3 architecture, DFlash speculative decoding is effectively a no-op.

---

8. Segfault on 3rd inference request

Consistent crash pattern: warmup request succeeds, request 1 succeeds (~0.07 tok/s), request 2 succeeds, request 3 segfaults. This matches the SSD streaming segfault seen in previous sessions. Likely a use-after-free or buffer overrun in the pread_into / expert streaming path.

---

9. Metal toolchain missing (build-time)

First swift build failed with xcrun: error: unable to find utility "metal", not a developer tool or in PATH. Required:

xcodebuild -downloadComponent MetalToolchain

The build.sh script prints this suggestion but doesn't check for it upfront. A pre-flight check would save time.

---

Summary of patches needed to run Kimi K2.5 + SSD streaming:

File	Change	Severity
DeepseekV3.swift	Strip language_model. prefix in sanitize()	Blocker
Module.swift	Handle .none entries in update(modules:)	Blocker
Load.swift	verify: [] for mixed-precision models	Blocker
swift-transformers	Tiktoken support / proper BPE decoder	Blocker for tiktoken models
Server launch	FD limit ulimit -n 1024+	Runtime crash
Draft model loader	Handle sharded safetensors for draft models	DFlash broken
DeepseekV3.swift	Add DFlashTargetModel conformance	DFlash no-op
pread_into C++	3rd-request segfault (use-after-free?)	Runtime crash

Happy to open separate issues for any of these or submit patches. The branch is very close to working — most of these are edge cases around MoE models with mixed precision and tiktoken tokenizers rather than fundamental architecture issues.

[0xClandestine]

@hankbobtheresearchoor good bot 🤖

[0xClandestine]

Thanks for the detailed writeup bot, very useful. Here's where things stand after today's commits:

Fixed in this branch:

1. language_model. prefix — DeepseekV3DFlash.sanitize() now strips the wrapper prefix at load time, so kimi_k25 HuggingFace exports resolve correctly to model.* paths.

2. Stale per-expert keys — after stacking experts.N.* weights into switch_mlp, the original per-expert keys are now removed from the dict. This prevents verify: .noUnusedKeys from throwing on the leftover keys.

3. FD limit — RLIMIT_NOFILE is raised to 4096 at server startup, covering the 182-shard case.

4. DFlash conformance for deepseek_v3 / kimi_k25 — DeepseekV3DFlashModel is registered via LLMTypeRegistry.shared.registerModelType() at startup (no submodule changes). Issue 7 should be resolved.

Needs an upstream fix (mlx-swift):

Issue 2 — the sparse update(modules:) crash. When layer 0's dense MLP is bf16 and layers 1–60 are 4-bit, quantize() produces an array update with .none at index 0. Module.update(modules:) currently tries to set that element to nil instead of skipping it, which crashes on a non-optional array slot. The fix is your suggested 2-line skip in mlx-swift/Source/MLXNN/Module.swift. That change needs to go to SharpAI/mlx-swift — will open a separate issue/PR there.

Still out of scope here: tiktoken (Issue 4), 3rd-request SSD segfault (Issue 8 — tracked separately), Metal toolchain preflight (Issue 9).