Skip to content

[Engine] support the official Flash-Next DFlash2 checkpoint - #3322

Closed
j-luwierski wants to merge 84 commits into
Niko1221:mainfrom
j-luwierski:codex/dflash2-main
Closed

j-luwierski wants to merge 84 commits into
Niko1221:mainfrom
j-luwierski:codex/dflash2-main

Conversation

@j-luwierski

@j-luwierski j-luwierski commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Problem and result

The official incoai/Qwen3.8-Flash-Next-DFlash2 checkpoint requires five attention layers with 128-wide heads, raw HC-mean target features, dynamic convolution and a predecessor-conditioned selector. This change implements that checkpoint contract and supports original BF16 weights plus Q8_0, Q5_0 and Q4_0 backbone quantizations, using the full shared target output head. This PR resolves the Qwen3.8-Flash-Next speculative-decoding gap tracked in issue #3304 by adding opt-in DFlash2 support and setup integration. The measured results do not establish the suggested 1.5–2× gain across workloads; performance depends on target, prompts, hardware and configuration, as detailed below.

Setup offers MTP (default), experimental DFlash2, or no model drafter. DFlash2 selection enables the local CUDA build, pinned download/conversion and artifact-specific calibration. Weights, private KV pools and attention scratch are accounted for before auto-sizing the expert cache, including the configured context length and elastic growth. MTP and DFlash2 cannot be selected together.

Related discussion

Explicit maintainer confirmation of the expanded official-checkpoint/setup scope remains pending. This remains a draft while that scope and missing backend validation are resolved.

Approach

DFlash2 is opt-in behind STRATA_ENABLE_DFLASH2=ON. The converter validates the checkpoint geometry and preserves all 45 norm, convolution and selector tensors in BF16; the 36 backbone matrices can be quantized. Setup offers BF16, Q8_0 and Q4_0; Q5_0 is also supported by the converter/runtime and included in the measurements below.

Prompt and conversation snapshots retain the draft context. Cache controls enforce matching actual expert placement before checking output. The shared runtime includes the infrastructure required for proposal verification, a fix for state publication on the opt-in QFUSE path and a fix for the CPU-only loader-test build dependency. The existing default drafter stays MTP.

Current implementation and speed comparison — measured 2026-10-23

The official DFlash2 path and opt-in runtime optimizations are implemented. The suggested +32.23% long-prompt decode advantage has not been reached. The new 256K code run measures +14.84% decode speed over MTP; full request time is 0.72% slower because prompt processing dominates.

The short-context rows are the three-repeat small-page-on audit from 2026-10-22. The 256K row is a paired run from 2026-10-23. Hardware: RTX 4070 Ti SUPER 16 GiB, Ryzen 9 5900X, 125.7 GiB RAM, Linux, CUDA 13.4 / SM89. Target: local Flash-Next NVFP4, BF16 embeddings, FP8 PLE, full shared Q8_0 head and INT8 KV. Native decode tok/s excludes startup and prompt processing.

Task Input tokens Context / resident KV MTP, tok/s DFlash2 Q8_0, tok/s DFlash2 profile DFlash2 vs MTP
Code, short 27 8192 / 8192 38.030 [37.788–38.083], n=3 49.077 [49.045–49.183], n=3 Small-page on +29.05%
Code, repeated-note padding 2427 8192 / 8192 36.813 [36.401–36.909], n=3 40.915 [40.865–41.041], n=3 Small-page on +11.15%
Cache, repeated notes 2425 8192 / 8192 51.026 [51.020–51.064], n=3 51.135 [51.110–51.186], n=3 Small-page on +0.21%
Code, 256K context 261747 262144 / 32768 34.900, n=1 40.080, n=1 Recorded Q8_0 profile +14.84%

For the first three rows, ranges are observed min–max, not confidence intervals. Small-page off → on changed DFlash2 speed from 48.952 → 49.077 tok/s (+0.26%) for short code, 40.759 → 40.915 (+0.38%) for padded code and 50.622 → 51.135 (+1.01%) for cache. DFlash2 accepted/offered counts were 70/85, 58/87 and 66/78. These results are fixture-specific measurements from one machine.

256K code run: decode and full-request cost

The paired run used the same 261,747-token prompt, target model, executable, expert profile, 96-token greedy output and seed in both arms. All 96 output token IDs matched.

Metric MTP, max T=4 DFlash2 Q8_0, block 7 DFlash2 vs MTP
Decode speed 34.900 tok/s 40.080 tok/s +14.84%
Decode time for 96 tokens 2.751 s 2.395 s 0.356 s saved
Full request time, prompt + decode 1012.19 s 1019.47 s 0.72% slower
Accepted / offered draft tokens 50 / 88 (56.82%) 53 / 62 (85.48%) —

Each method ran once because processing the prompt took about 17 minutes per arm; the 256K tok/s values have no statistical interval. Full request time is not a prompt-only measurement. DFlash2 saved 0.356 seconds during decode, while its whole request took 7.280 seconds longer. The +32.23% target would require 46.149 tok/s at this MTP rate. Shared settings: 2,700 actual expert slots, 11 workers, verifier capacity 8, prefill 512, short-read 64 and 480 MiB reserve; prompt cache, adaptive swaps and PCIe sharing were disabled. Method settings were held fixed: DFlash2 block 7 / spec_min_p=0.7; MTP max T=4 / spec_min_p=0.5. No launch-parameter sweep was performed.

Full 256K report, exact prompt IDs and provenance, commands, raw measurements and logs · Machine-readable summary · 2026-10-22 raw commands and output controls

The short and padded-code target continuations diverge at output index 11, and the cache task uses different content. These rows do not isolate context length from continuation difficulty. The older +32.23% short-code result used a different profile and is historical; do not combine it with these measurements.

How to enable the pushed optimization options

Use a build containing this PR (implementation dbaa12db, evidence follow-up 3aef66c3) and a compatible DFlash2 artifact. These are engine environment variables, not new command-line flags. Set them before starting the engine, then restart the engine after changes; several switches are read only at startup or first use. Selecting MTP does not activate DFlash2-specific kernels.

For a source CUDA build, enable the extension in the normal CUDA configuration:

cmake -S . -B build-dflash2 -DSTRATA_ENABLE_CUDA=ON -DSTRATA_ENABLE_DFLASH2=ON
cmake --build build-dflash2 --target strata -j4

Keep the existing model paths and DFlash2 launch arguments. Set the desired options in the same shell that starts the engine/server; an already-running process does not inherit later shell changes.

Environment variable What setting it to 1 enables Default / important limits
STRATA_DFLASH2_SMALL_PAGE_ATTENTION Specialized attention indexing and aligned paired FP16 value loads for the runtime's four-cell pages Off. Official 32Q/8KV geometry; other geometries retain the general path. This is the new kernel option measured in the table.
STRATA_DFLASH2_PROMPT_TAIL Skips whole context-fusion batches that future local attention cannot read Off. Fresh native requests with --prompt-cache 0; reference/parity/debug paths keep full processing. Improves prompt processing only, not token acceptance.
STRATA_DFLASH2_PARALLEL_PROJECTIONS Runs proposal Q/K/V projections on separate streams, and overlaps gate/up projections Off. Reference, tensor-dump, debug and section-timing paths retain serial execution. This does not enable the rejected context K/V prototype.
STRATA_DFLASH2_REGISTER_TOPK Register-based local candidate top-k within the tiled path Off. Requires the tiled path and workspace; tiled top-k is already enabled by default. Do not force STRATA_DFLASH2_SERIAL_TOPK=1 when testing it.
STRATA_DFLASH2_FUSED_CONV Fuses convolution epilogues with following conversion/residual work Off. Reference/parity/debug paths retain their diagnostic behavior.
STRATA_DFLASH2_LATTICE_SELECTOR Computes candidate transition scores in parallel before walking the conditioned path Off. Requires normal selector workspace; serial/reference selector paths bypass it.
STRATA_DFLASH_HEAD_IL + STRATA_MMVQ_IL_Q8_HEAD Enables DFlash2 interleaved head dispatch and the compatible Q8_0 shared-head kernel Both off. Set both to 1. The Q8_0 specialization is gated to measured SM89, vocabulary 248320 and supported row counts; unsupported cases retain the existing path. No head weights are replaced.

For the small-page-on arm of the speed table, use the two shared-head options in both MTP and DFlash2 comparisons, and enable only the new small-page option on the DFlash2 arm:

export STRATA_DFLASH_HEAD_IL=1
export STRATA_MMVQ_IL_Q8_HEAD=1
export STRATA_DFLASH2_SMALL_PAGE_ATTENTION=1
# Start the existing DFlash2 engine command from this shell.

Optional, independently testable additions for the existing DFlash2 command:

export STRATA_DFLASH2_PARALLEL_PROJECTIONS=1
export STRATA_DFLASH2_PROMPT_TAIL=1
# PROMPT_TAIL requires --prompt-cache 0 and the fresh native request path.
# Restart the engine after changing these variables.

PowerShell equivalent (syntax reference; Windows execution has not been validated):

$env:STRATA_DFLASH_HEAD_IL = "1"
$env:STRATA_MMVQ_IL_Q8_HEAD = "1"
$env:STRATA_DFLASH2_SMALL_PAGE_ATTENTION = "1"
$env:STRATA_DFLASH2_PARALLEL_PROJECTIONS = "1"
$env:STRATA_DFLASH2_PROMPT_TAIL = "1"
# Start/restart the existing DFlash2 engine command from this PowerShell session.

The table is not a measurement of all options enabled together. Parallel projections, prompt-tail, register top-k, fused convolution and lattice selector are separate experiments; do not infer an additive speedup. To restore these opt-in paths to their defaults, unset the corresponding variables (or set them to 0) and restart. In Bash use unset STRATA_DFLASH2_PARALLEL_PROJECTIONS; in PowerShell use Remove-Item Env:STRATA_DFLASH2_PARALLEL_PROJECTIONS, and likewise for each enabled option.

The rejected STRATA_DFLASH2_PARALLEL_CONTEXT, STRATA_DFLASH2_WARP_ATTENTION_MAX and STRATA_DFLASH2_BF16_STREAM_MEAN prototypes are not present in production and must not be advertised as usable optimization options.

Prompt-processing and acceptance findings

A separate 7227-token code fixture measures prompt-tail processing at 16897.4 → 16759.0 ms (−0.82%), with identical output IDs and 20/28 acceptance. Decode medians are 713.4 and 713.9 ms: no steady-state win. Prompt-time changes on the 2427/2425-token fixtures are negligible. This evidence is separate from the decode-speed table.

The joint diagnostic capture records native BF16 proposals and their exact target features at 95 anchor positions per task. The published BF16 PyTorch drafter, consuming those same features, chooses the same first proposal at all 285 positions: 90/95 correct for short code, 81/95 for padded code and 83/95 for cache. Later proposals can differ; this is not full bitwise PyTorch parity or validation of native target features against the upstream teacher. Reference/native correct-prefix sums are 256/258 for padded code and 348/348 for cache. Simply reproducing reference BF16 arithmetic does not recover the missing acceptance on these fixtures.

Context deduplication, two-pass self-conditioning and causal per-request selector feedback did not consistently improve acceptance. Extra context K/V parallelism and warp-maximum reduction did not establish a speed benefit and were removed. Their patches, results and limitations remain as negative evidence. No weights were fine-tuned in these experiments.

Latest validation

Passed: CUDA builds with DFlash2 enabled and disabled; dflash_load_test, dflash_policy_test; default/enabled dflash2_parity --bench; 30 before/default/enabled attention outputs bit-identical, including four-cell pages, unaligned fallback and physical rings; independent FP64 attention oracle; long NumPy stage oracle for all five layers/live KV/final normalization/selector within documented tolerances; fresh native output controls, complete joint feature/proposal capture and Python syntax compilation of the modified stage oracle.

Diagnostic prototypes were removed and production rebuilt. SHA256 42a56c9a217ddb47c58fa6f42ad331c77ae2c296a532d84c53b3610fedb4a3f2 matches the executable used for the latest table.

Failed / rejected: the follow-up algorithm and scheduling experiments did not meet their acceptance/performance goals. Initial incomplete captures were retained and excluded from final findings. Earlier test failures remain documented below and in their original reports.

Not run on the latest head: HIP/SYCL compilation/execution (compilers unavailable), Metal, Windows, other target families, multi-GPU, vision, concurrent requests, a held-out long-prompt suite or full 128K prompt sweep. Earlier validation at another revision is not claimed as a latest-head rerun. This PR remains a draft pending scope agreement and missing backend validation.

Historical performance and validation reports

These links preserve previous tables and their original settings. They are not the current speed table:

Historical validation — original revisions

Earlier checkpoint validation used engine source at af85f5ce3b4411833b5529361ae18ffdadbdce3f, based on main e88ee39d1d6be6e3364183af3b39b62a7ed6f211. Performance evidence commit: 86af7286254d1efb148cf0aa8c9a96868d4e9f09; it adds scripts/results without changing engine code. Linux, RTX 4070 Ti SUPER 16 GiB, Ryzen 9 5900X, 125.7 GiB RAM, CUDA 13.4, GCC 14.2, SM89.

Passed:

  • Follow-up fix at c7bebe215a6feb754e08dfbc02a3ac12d38fa55f: CUDA extension ON/OFF builds; three ON and two OFF CTests; native fixed-block default/explicit-zero checks; 192/192 target-only fixture IDs match all threshold/block repeats at matching 3279 expert slots. Same-day MTP fixture IDs also match.
  • CUDA extension ON/OFF builds; 12 selected CUDA CTests, plus extension-OFF loader rejection.
  • 203 setup/conversion/cache/tool tests and 342 server tests.
  • CPU-only full build after fixing the loader test's CUDA-header dependency.
  • Main/OFF/ON regression comparison: 768 candidate IDs match main across AR/MTP, greedy/sampled requests and three prompts. Final rebuild executable hashes are unchanged.
  • Official Q8 DFlash2 cache test: eight cases, 256/256 IDs equal both fresh DFlash2 and target-only controls at matching 2000 expert slots. Checks live continuation, parked conversation restoration, prompt snapshots, rewind and sampled-to-greedy transition.
  • Full 128K draft-memory allocation with BF16/Q8/Q4 and short generation; one actual 8192→16384-cell growth with an 8448-token prompt.

Commands, source/executable hashes, raw IDs and logs are in validation evidence.

Failed:

  • One follow-up output-control startup provided 3275 rather than the requested 3279 slots; the placement assertion stopped before any request. The successful unmeasured controls use a 480 MiB reserve to retain exactly 3279 actual slots. Performance runs keep 512 MiB. The failed startup is retained in the follow-up evidence.

  • Four CPU CTests fail with the available dependency/tooling: three require missing Trellis types and one reports an AVX512 allowlist violation. All four reproduce on unmodified main. The full candidate CPU run has 44 passed, one skipped and four failed tests.

  • Initial CPU-only build and setup golden-fixture failures were corrected and rerun successfully; original failures are retained.

  • The initial default-path harness rejected a valid early EOS from a truncated prompt. The corrected harness uses the full prompt and compares actual IDs; its failed attempt is retained.

  • An earlier 256 MiB VRAM-reserve startup failed to fit verifier allocations. That failed attempt and an interrupted fixed-cache experiment are retained in the earlier quantization evidence and excluded from measured medians.

  • A preliminary performance restart overlapped teardown of the preceding engine and failed the host memory guard. The previous process tree was fully stopped before the final complete series. That failure and the interrupted prefill-512 attempt are retained and excluded from medians.

Not run:

HIP/SYCL compilation and execution (ROCm/icpx unavailable), Metal, Windows, multi-GPU, concurrent requests and a full 128K-token prompt. The CUDA extension is explicitly restricted to one GPU and serial requests. Shared backend changes still require the missing builds; keep this draft pending targeted validation.

Limits

Sampled DFlash2 requests use target-only fallback. Cache equality checks cover tested arithmetic paths and matching placement, not every configuration. Disk session restoration is unsupported. The CUDA extension is restricted to one GPU and serial requests. The latest table uses equal actual expert placement; older auto-cache reports allowed different placement, so their results must not be mixed. Performance conclusions apply to the measured hardware and prompts. No universal speedup or +32% long-prompt result is claimed.

Additional setup checkpoint (2026-10-20)

Setup now offers the pinned Turbo-Brilliance DFlash2 fine-tune with --dflash2-head turbo-brilliance; the official checkpoint remains the default. This fine-tune was trained for DavidAU's Turbo-Brilliance GSQ-RCO IQ3_S target. Setup downloads its BF16 model.safetensors, converts it to GGUF and keeps it in a separate cache. The repository's t3000-q8-dv.gguf is a separate plain-DFlash artifact and is not used by this CUDA DFlash2 path.

Validation: pinned revision and SHA-256 checked; config and all 81 BF16 tensor shapes matched the converter using ranged reads of the checkpoint header (the full 1.2 GB weights were not downloaded). Setup choice, cache/conversion, quantizer and calibration-key tests pass. The CUDA DFlash2 conversion of this fine-tune has not been performance-benchmarked; the upstream Apple M4 timings are for the separate plain-DFlash GGUF.

@gopinath87607

Copy link
Copy Markdown
Contributor

hi this is a great feature i was expecting dflash2 for a long time btw if you have time can you check this prs

mw00/project-maya#283 (dflash+mtp space allocation)
&
#3307 (mtp space allocation) maybe once your pr got merged maybe i can work on this loan allocation this will save vrams.

@j-luwierski

Copy link
Copy Markdown
Contributor Author

Thanks a lot for the kind words! I really appreciate it! :)

Sure, I'd be happy to take a look at your PRs. And thanks for offering to work on the loan allocation improvements — saving VRAM would definitely be valuable!

I'm still working on improving DFlash2 performance for long prompts, so there's more optimization work ahead.

@bsorensen110

Copy link
Copy Markdown
Contributor

DFlash2 on the AMD HIP backend (RDNA4 / gfx1201) — buildable, parity-clean, measured

The STRATA_ENABLE_DFLASH2 block refused every non-CUDA backend, so I ported it to HIP to find out whether the restriction was a code constraint or a validation gap. It is the latter: the kernels are already portable.

Branch: https://github.com/bsorensen110/Strata/tree/feat/dflash2-assistant-prefix — tip d46a21b9, based on dea6920b (v0.2.1.2) with #3356 and this PR's 56 commits cherry-picked, plus the HIP port and a merge of my not-yet-upstream RDNA4 wave.

What the port needed — two files

  1. CMakeLists.txt: the FATAL_ERROR "DFlash2 currently requires the NVIDIA CUDA backend" becomes a warning when a backend is present. src/kernels/cuda/dflash2.cu was already in the shared strata_kernels source list and compiles unchanged through include/strata/platform/hip_compat — only <<<>>>, __syncthreads, __shfl_xor_sync, __half/__half2, no tensor cores, no cuBLAS.
  2. src/kernels/dflash2_parity.cpp: __half_as_ushort / __ushort_as_half are device-only in HIP, so the host test did not compile. Replaced with strata::kernels::f16_from_f32 / f32_from_f16 from include/strata/kernels/f16_bits.hpp (round-to-nearest-even, exact inverse — the tree's existing fp16-bits source of truth, worth using here for the CUDA build too given the note in that header).

src/core/dflash*.cpp needed no change: I found no CUDA-only calls in it.

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DSTRATA_ENABLE_HIP=ON \
      -DSTRATA_ENABLE_CUDA=OFF -DCMAKE_HIP_ARCHITECTURES=gfx1201 -DSTRATA_ENABLE_DFLASH2=ON
cmake --build build --target strata dflash2_parity -j16

Validation on gfx1201 (Radeon AI PRO R9700, ROCm nightly 10.2)

dflash2_parity: PASS — convolution, block isolation, top-k, conditioned path, confidence and nonfinite refusal. Neighbouring suites on the same build: iq_multi_parity 0 failures, shared_expert_parity 0 failures, gr_parity --selftest OK.

Artifact: official checkpoint at pinned a22a1744, model.safetensors sha256 verified against d6723f2a…, converted to Q8_0 (776,057,536 B, 81 tensors = 45 BF16 + 36 Q8_0, selector_top_k 16, taps [5,15,25,35,45]).

Measurements

Target: Qwen3.8-Flash-Next IQ3_S pack, --kv int8, greedy, 4-token prompt, --dflash-block 3 --dflash-window 0 --spec 8.

Determinism first, since it changes the numbers: with STRATA_PREFILL_CPU_SHARE=0 and STRATA_ROUTE_TAIL_SKIP=0 set, repeated runs of the same configuration are byte-identical on this machine. Note that --spec 2 is prompt-lookup speculation, not a target-only control, so it is not a valid "expected tokens" reference — my first cross-check used it and produced a false alarm.

Engine (--generate), three runs each, auto expert cache:

Drafter 16 new tokens 48 new tokens
MTP --spec 4 52.53, 53.00 tok/s 64.27 tok/s
DFlash2 Q8_0 58.96, 59.62 tok/s 60.41 tok/s

Service (/v1/chat/completions, 60-token prompt, 36 generated, same binary and config apart from the drafter):

Drafter decode drafts accepted text
DFlash2 Q8_0 110.7, 115.0, 92.9, 98.4 tok/s 24/26, 24/26, 24/29, 23/25 "The ocean is vast and blue."
MTP 84.0, 86.9 tok/s 23/29, 24/32 "The ocean is vast and blue."

So on this card DFlash2 is ahead on short generations (+11% engine, +~22% service) and roughly tied on 48 tokens; acceptance ran 88–96% here. The draft path also reserves less VRAM than MTP, which of course changes expert residency: 11,808 experts / 22.50 GiB with DFlash2 vs 11,616 / 22.14 GiB with MTP at 262144 context. Those samples are single observations, not confidence intervals.

Notes

  • Cross-drafter token divergence I saw initially tracks expert residency, not the drafter: with a cache small enough that the CPU pool serves most experts, MTP and DFlash2 emitted identical ids for 13 tokens. If you want A/B token equality in your own runs, pin --expert-cache and turn the CPU share off.
  • --dflash-block 3 is the configuration setup writes and the one I measured. I did not validate --dflash-block 0.
  • Not tested: CUDA-vs-HIP numerics on AMD vs NVIDIA, long-prompt latency, other AMD architectures, Windows.

Thanks — the checkpoint contract and the artifact validation tools made this port small; the only real work was the backend gate and the host-side fp16 helpers.

@bsorensen110

Copy link
Copy Markdown
Contributor

Very good 👍

@fukc-gihtub

fukc-gihtub commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

I think the DFlash2 will outperform MTP on coding tasks.

@j-luwierski

Copy link
Copy Markdown
Contributor Author

That’s exactly what I’m seeing in my local tests as well. DFlash2 seems to have an advantage over MTP on coding tasks so far.

@Gamal-ElDeen

Copy link
Copy Markdown

Registering a fork branch so this checkpoint is not started again as a separate pull request.

https://github.com/Gamal-ElDeen/Strata/tree/dflash2 is one commit on v0.2.1.4 (9de04cd). It adds an opt-in --dflash PACK for a local DFL2 file produced by tools/dflash_rt.py from a local incoai/Qwen3.8-Flash-Next-DFlash2 snapshot. The flag is exclusive with --mtp. After the pack loads, the process exits before serve, because the verify window still drafts only with MTP.

This is not a request to merge that branch, and it is not a second implementation of the drafter in this pull request. #3457 covers the Metal --dflash path. #3183 is the open request for the checkpoint.

@Gamal-ElDeen

Copy link
Copy Markdown

Measured this branch on an RTX 5090 (SM_120, CUDA 13.4), target Qwen3.8-Flash-Next FP4-Local, drafter the BF16 GGUF from the official safetensors. SHA tested: 7841145. Built as-is with STRATA_ENABLE_DFLASH2=ON. The seven dflash ctests passed. No source patch.

Same binary for both arms, spec 4, spec-min-p 0.5, pool-workers 19, learned expert profile, w4a8, tail skip off. One run per cell.

Fair width is --dflash-block 3. With the default suffix drafter left on, --spec 4 raises the verify cap and this tree drafts K=5 while MTP stays at 4.

cell acceptance decode tok/s
MTP short greedy, 141 prompt, 200 tokens 129/175 166.2
DFlash K=3 short greedy 129/172 145.5
MTP long greedy, 5269 prompt, 200 tokens 122/162 163.0
DFlash K=3 long greedy 122/157 143.1
MTP short sampled, temperature 0.7 45/47 197.4
DFlash K=3 sampled 13/24 114.4

Sampled requests log that DFlash uses target-only decoding. K=5 was not faster on the short prompt (144.7 tok/s) and was tied on the long one (163.7 vs 163.0). Peak VRAM was 32178 MiB for MTP and for K=3. These are single runs, so a few percent is noise. Not merged here.

@j-luwierski

j-luwierski commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor Author

@Gamal-ElDeen Thank you so much for taking the time to test this and share such detailed results! I really appreciate it, especially since I don’t have an RTX 5090 myself.

Could you also try testing DFlash2 with a smaller quantization of the drafter head, for example Q8? My suspicion is that using a smaller quantization for the DFlash2 head could reduce its VRAM footprint and leave more room for caching the target model’s experts, potentially improving performance.

Thanks again for doing these tests and sharing the results!

@j-luwierski

Copy link
Copy Markdown
Contributor Author

Only CUDA path.

@j-luwierski

Copy link
Copy Markdown
Contributor Author

Hi @Niko1221. DFlash2 is now outperforming MTP! 🚀

Could you please test it on your more powerful hardware? I'd be really interested to see how the performance scales on a higher-end GPU and whether the improvements hold up.

@j-luwierski
j-luwierski marked this pull request as ready for review October 23, 2026 03:17
@j-luwierski
j-luwierski force-pushed the codex/dflash2-main branch 2 times, most recently from 3b56edb to 4335151 Compare October 23, 2026 12:57
@j-luwierski

Copy link
Copy Markdown
Contributor Author

Hey @Niko1221! I noticed you've added your own DFlash2 implementation. Nice work! :)

I wanted to ask whether you think my implementation could still be useful for the project, and if it makes sense for me to keep working on improving it. I'd be happy to continue, especially on the performance side, but I'd also like to avoid duplicating effort.

What do you think?

@3chet

3chet commented Oct 24, 2026

Copy link
Copy Markdown

Build failure on gfx906 — the root cause is a macro gap, and the fix is two lines

We hit this on v0.2.3.1 (ROCm 7.14, gfx906), and a repo-wide search for dflash2 gfx906 returns nothing, so we may be the first:

src/core/dflash2.cpp:205: error: identifier "cudaStreamCreateWithPriority" is undefined
src/core/dflash2.cpp:206: error: identifier "cudaDeviceGetStreamPriorityRange" is undefined

Root cause. STRATA_USE_HIP is defined in exactly one place: cmake/hip_backend.cmake:98. That file is only included on the STRATA_ENABLE_HIP (wave32 / RDNA) branch. gfx906 takes the other branch — CMakeLists.txt:76-80 sets STRATA_ENABLE_CUDA=ON — so it never includes hip_backend.cmake and never defines STRATA_USE_HIP anywhere. The guards at dflash2.cpp:198 and :446 (v0.2.3.1 line numbers) therefore evaluate false, the #else CUDA branch is compiled, and it references two names the gfx906 compat header doesn't map.

Upstream already knows the disconnect: CMakeLists.txt:157-159 comments that hip_backend.cmake is "never included here" — the rocBLAS gating was copied without the macro.

Why nobody reported it. @bsorensen110's HIP port above is the best evidence: he builds with -DSTRATA_ENABLE_HIP=ON -DSTRATA_ENABLE_CUDA=OFF on gfx1201, so STRATA_USE_HIP is defined for him, the native branch always compiles, and his conclusion "src/core/dflash*.cpp needed no change: I found no CUDA-only calls in it" is correct — for his build. On gfx906 the guards evaluate false, so the CUDA branch is the one compiled and those calls do exist. CI doesn't run the opt-in gfx906 build either, so nothing catches it.

Suggested fix. Widen the two guards:

#if defined(STRATA_USE_HIP) || defined(STRATA_HIP_GFX906)

This matches the idiom already used in generate.cpp (7 occurrences) and leaves the other 36 STRATA_USE_HIP sites untouched. We just rebuilt 0.1.48.1 with exactly this change: it compiles clean on gfx906, and the HIP branch is what gets built — which also restores the intended defaults (high-priority streams off, single stream) that the CUDA branch would have silently flipped on. Defining STRATA_USE_HIP=1 for gfx906 would also work but flips 45 sites, so we avoided it.

Two caveats for whoever picks this up: on gfx906 we are still gated out at runtime by the single-card/CLI restriction at generate.cpp:5477-5479 (--serve, --batch, or layer split all fall back to MTP), so for us this is a build-correctness fix, not a feature unlock. And the two cuda* → hip* mappings we added to strata_hip.h as a stopgap do compile, but they enable code paths upstream deliberately disabled on HIP — hence the guard fix instead.

Happy to test a patch.

…try, CLI)

The DeepSpec DFlash block drafter (PixelML/Qwen3.8-Flash-Next-NVFP4-DFlash) as a
standalone artifact next to the MTP layer.  docs/DFLASH.md pins the semantics
derived from the PixelML serving adapter, llama.cpp's dflash support and the
z-lab reference: the HC-contracted taps [3,15,23,35,43], the context-KV
injection (rmsnorm(fc(taps)) feeding every layer's K/V), and the DeepSpec
anchor layout (K queries, query zero fed the anchor token, every query
predicts one past its own rope position).

- include/strata/core/dflash.hpp + src/core/dflash.cpp: GGUF v3 parse,
  llama.cpp and raw-HF naming families, structural validation (all 58 tensors,
  BF16 only, shapes vs metadata, refusal of DFlash2/DSpark extras and of a
  shipped embed/lm_head), validate_supported for the geometry fast path.
- tools/dflash_gguf.py (+ BF16 in tools/gguf_writer.py): safetensors -> GGUF
  with a bit-exact read-back self-check; the published checkpoint converts
  1:1 (sha256 35a23c17c248ff2e3296e6b78882b6d955af3092498fb7b6c48be1af4bfa971a).
- CLI: --dflash FILE (exclusive with --mtp, needs --spec 2..8),
  --dflash-block (never above the trained block), --dflash-window,
  --dflash-mask-token; prompt lookup disabled out loud; the engine stops with
  an explicit milestone line until the block forward is wired.
- dflash_load_test: 20 host-only fixture cases.
…als)

The verify window and the prompt path both capture the ATTN-HALF HC read's
contracted residual at the DFlash boundaries (the trained tap ids + 1, i.e. the
residual after each tapped layer including its FFN write-back) - the tensor the
drafter's fusion consumes, not the raw 10240-wide HC state.

- Verifier::set_tap_layers: the copies ride in the captured window graphs
  (both the doorbell and zero-doorbell variants) into a [n_taps x max_t x
  n_embd] f32 buffer carved with the rest; tap_stride() exposes the layout.
- Prefill::set_tap_layers + on_taps: per-chunk capture in BF16 (the fusion's
  input precision; the f32 buffer did not fit next to the auto-sized expert
  cache), delivered next to on_chunk with the chunk's rows strided by
  tap_stride_rows().
- STRATA_DFLASH_TAPS=<file> + tools/dflash_taps.py: the feature-capture gate.
  Two target-only runs over a shared deterministic prompt; the prompt path and
  the verify window measure the same position.  IQ3_XXS, 400 tokens: per-tap
  cos 0.9966-0.9999 (engine-path tolerance); the gate's floor sits at cos 0.99,
  an order of magnitude above any wiring bug's signature.
…p (greedy)

DFlashDrafter: BF16 weights uploaded from the GGUF; ONE QsaState PER DRAFT
LAYER (a QsaState is a single layer's pools - sharing them made every layer
attend layer 0); the fusion (bf16 taps -> fc -> hidden_norm -> per-layer
k/v with k_norm + rope, appended at the committed positions); the block
forward (query rows [anchor, mask x (K-1)] at rope positions [P..P+K-1],
non-causal attention over [0, P+K), the target's own head, row argmaxes).
Eager, one stream, batches of <= 8 rows.  New kernels: swiglu_inplace and
bf16_gather_strided (elementwise).

Engine wiring: --dflash uploads before the expert cache is sized (its
bind_bytes joins the VRAM reserve), the head/embedding bind from the target,
the verifier's taps feed the per-cycle context update (add_context_f32), the
prompt path's taps feed the context cells (add_context), the T=1 bootstrap
window's pick is the first anchor, prompt lookup is disabled out loud,
sampled runs decode without the drafter (greedy only), --serve + --dflash is
refused for now.

STATE: target-only paths unchanged (ctest green); DFlash runs end to end
(16 GB card, IQ3_XXS target, 1.9 GiB VRAM for the drafter) but acceptance is
ZERO - the proposals are near-constant tokens, i.e. the conditioning does not
reach the queries correctly yet.  The alignment/conditioning debug continues
against a numpy reference harness (next commit).
The z-lab/DeepSpec semantics in f32 with bf16 rounding at the activation
boundaries: context = rmsnorm(fc(taps)) (one vector per committed position,
every layer's K/V from it with k_norm + NeoX rope at the position), the noise
block [anchor, mask x (K-1)] at rope positions [P..P+K-1] through five dense
layers with per-head q/k norms and NON-CAUSAL attention over the context
cells plus the block's own cells, final norm.  Stage dumps (--dump) feed the
engine-vs-reference comparison; the head stays the target's.
The drafter inherited shapes_of()'s n_rot = 64 - the Flash-Next TARGET QSA's
partial rotary - so dimensions 64..255 of every draft Q/K head were never
rotated.  The PixelML checkpoint was distilled with FULL 256-dimensional
NeoX RoPE, theta 1e7, unscaled (its config: head_dim 256, full-head, no
scaling keys).

- native_rope.cu: dflash_rope_neox_apply(), a DFlash-specific helper (the
  validated target path refuses n_rot != 64 by contract and is not stretched):
  one thread per pair, angle = pos * theta^(-2p/256), ggml NeoX reading,
  no scaling, exact in-place alias allowed.
- DFlashDrafter: shapes_.n_rot = head_dim and all four rope calls (context K,
  query Q, block K) go through the helper with the artifact's theta.
- dflash_rope_parity: GPU vs independent host reference in the same float
  math; position-scaled bound (device/host powf ulp amplified by the
  position), a check that dims 64..255 actually rotate for position > 0
  (the regression this guards against leaves them untouched), position 0
  identity, in-place aliasing.
@j-luwierski
j-luwierski force-pushed the codex/dflash2-main branch 2 times, most recently from 3b799b0 to 212a924 Compare October 24, 2026 02:23
@j-luwierski

Copy link
Copy Markdown
Contributor Author

Hi @3chet. Thanks for the suggestions and for testing! I've implemented the suggested changes. Could you please test them again and let me know if everything works as expected?

@3chet

3chet commented Oct 24, 2026

Copy link
Copy Markdown

Thanks — happy to test immediately, but I can't find the change to test against. Could you point me at it?

What I checked just now (2026-10-24 04:30 UTC):

  • src/core/dflash2.cpp on main still has #if defined(STRATA_USE_HIP) at lines 198 and 446, unchanged. Same at tag v0.2.4.1.
  • include/strata/platform/hip_compat/strata_hip.h on main has no mapping for cudaStreamCreateWithPriority or cudaDeviceGetStreamPriorityRange (its last commit is e6ad6c81d, 2026-10-16).
  • No new branch (git ls-remote shows 2 branches, none matching dflash/hip/gfx906), and your latest commit on the repo dates from 2026-09-30.
  • The 30 commits pushed since 2026-10-23 18:00 UTC are all dsv41 / launch / docs.

So either the fix lives in a fork branch you haven't pushed here yet, or it went somewhere I'm not looking. A commit sha or a branch name is enough — I'll build that exact tree on gfx906 (ROCm 7.14) and report back the same day.

For reference, the two-line change we validated locally is:

#if defined(STRATA_USE_HIP) || defined(STRATA_HIP_GFX906)

at both dflash2.cpp:198 and :446. With that, 0.1.48.1 compiles clean on gfx906 and the HIP branch is what gets built (so the deliberate HIP defaults — no high-priority stream, single stream — stay in force). If you implemented it differently (e.g. defining the two names in strata_hip.h instead), I'm still happy to test that path; I just need to know which tree to build.

@j-luwierski

j-luwierski commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor Author

@3chet Thanks for checking — both fixes are in PR #3322:

  • ddbd3b9: your exact two-line guard change at both locations in src/core/dflash2.cpp, preserving the intended HIP stream defaults.
  • 11f5d51: I’ve just added both CUDA → HIP priority API mappings in strata_hip.h as well.
    The PR head containing both fixes is 11f5d51. Thanks for offering to test on gfx906 / ROCm 7.14 — looking forward to your results.

@3chet

3chet commented Oct 26, 2026

Copy link
Copy Markdown

gfx906 build of head c90fa866: clean, with zero local patches

@j-luwierski — sorry for the delay, this needed a free window on the box. Note that the head has moved on from 11f5d512 to c90fa866 (the ~20 perf commits on top), so I tested the current head, not just the two fix commits.

I rebuilt it on the machine that originally reported the failure (2× Instinct MI50 16 GB, gfx906 / wave64). The tree is exactly c90fa866 — no local reverts, no extra shims:

cmake -S . -B build-3322 -G Ninja -DCMAKE_BUILD_TYPE=Release \
      -DSTRATA_HIP_GFX906=ON -DSTRATA_NATIVE_EXPERTS=ON \
      -DSTRATA_GGML_DIR=third_party/llamacpp-pinned -DSTRATA_BUILD_ID=c90fa866
  • 162/162 steps, 0 errors. Toolchain: the ROCm 7.14 container (AMD clang version 23.0.0) that also builds my production engines; ggml pinned at 3cf03257f219.
  • strata --version → strata 0.1.49 (the PR's base line), --families → qwen4exp.
  • Binary 36,787,080 B, md5 420dc8c05db287f7e189081770827939. Rebuilt from a clean object state a second time → byte-identical, so it is reproducible here.
  • nm -D shows U hipStreamCreateWithPriority@hip_4.2 and no cudaStreamCreateWithPriority ⇒ the HIP branch is what got compiled, which also restores the HIP stream defaults (no high-priority stream, single stream) that the CUDA branch had been silently flipping on.

Which commit is load-bearing: both work, one is redundant

Two negative controls on the same tree:

Tree state Result
head as-is (ddbd3b94 + 11f5d512) compiles
guard reverted to #if defined(STRATA_USE_HIP), shim kept still compiles
guard reverted and both #defines removed from hip_compat/strata_hip.h fails with exactly the original errors — dflash2.cpp:205 'cudaDeviceGetStreamPriorityRange' was not declared, :206 'cudaStreamCreateWithPriority'

So 11f5d512 alone is sufficient for the build, and ddbd3b94 is belt-and-braces. I'd keep both: only the guard restores the HIP-specific stream behaviour on a build where STRATA_USE_HIP is never defined, and it matches the idiom already used in generate.cpp.

Scope caveat (so nobody expects a feature unlock on gfx906)

This is a build-correctness fix here. CMakeLists.txt:90-92 still refuses STRATA_ENABLE_DFLASH2 together with STRATA_HIP_GFX906, and the runtime gate keeps DFlash2 to one card + serial CLI. src/core/dflash2.cpp is compiled unconditionally into strata_core, which is precisely why the build broke without anyone opting in.

Not tested: actually running the drafter on gfx906 (both cards are held by a production 2× MI50 layer-split server, and I did not want to take it down for a path CMake refuses anyway), multi-GPU, and the non-Qwen target families.

Happy to rebuild any follow-up head — a full build takes about 4 minutes on this box.


Two small suggestions, feel free to ignore: (a) a compile-only CI job covering gfx906 would have caught this without a GPU — the failure is pure host-side C++ in dflash2.cpp, no device needed; (b) #if defined(STRATA_USE_HIP) || defined(STRATA_HIP_GFX906) now appears twice in dflash2.cpp while generate.cpp uses it seven times — a shared STRATA_HIP_LIKE macro would stop the next file from repeating the same trap.

@Niko1221

Copy link
Copy Markdown
Owner

Thank you for the work on the official DFlash2 checkpoint. Strata already ships its own DFlash2 support, so we will not merge this second implementation. Thanks for the careful work.

@Niko1221 Niko1221 closed this Oct 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants