Skip to content

Commit d68ed8f

Browse files
committed
Add shared DFlash2 preparation for Qwen Flash Next
Share draft configuration, context rings, preparation and speculative rounds with Qwen 27B. Select Flash Next DFlash2 explicitly while retaining MTP defaults. Support MXFP8 draft projections, quantized draft heads and W4A16 routed decode with calibrated A4 prefill on the same packed expert weights. Order compiled QSA metadata before selector execution and enumerate declared draft weight roots once during export. Keep serving policy in Helios unchanged. Validation: 98 focused CPU tests, 12 compiled/eager QSA GPU cases and two draft state GPU cases pass. The SM121 TP1 package passed 36 native requests, including 32 concurrent requests and prompt-cache reuse. DFlash2 TP2/4 was not GPU tested. Design-Signed-Off-By: Codex (GPT-6); design-review=a38f
1 parent 585159c commit d68ed8f

46 files changed

Lines changed: 1500 additions & 344 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎docs/cli-design.md‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -391,8 +391,9 @@ explicit flags take precedence.
391391
| `--prefill-chunk-size N` | Most prompt tokens processed in one prefill step, a multiple of 64. A larger chunk reaches a long prompt's first token sooner; while it runs, other active requests wait for their next token. Command line only. | The recipe's chunk |
392392
| `--mixed-prefill` | Capture mixed prompt/decode graphs alongside separate graphs. Supported by Qwen with concurrency of at least two, with or without speculation. Command line only; existing packages retain their captured catalog. | Disabled |
393393
| `--expert-scale-compression [MODE]` | Losslessly compress routed-expert scales: DeepSeek V4.1 backbone and DSpark MXFP4 scales, or Step-5 NVFP4 scales including BF16 experts converted to W4A16. `csf` is a compressed scale encoding expanded during execution. Estimated resident GPU weight savings are 3% for DeepSeek V4.1 and 6% for Step-5; actual savings depend on encoded scales. A bare flag selects `csf`; `native` disables compression. Retain compressed scales and share expansion scratch across serial layers. | Recipe default: `csf` for DeepSeek V4.1, native otherwise |
394+
| `--speculative-method {mtp,dflash2}` | Select the recipe's learned proposal method. Flash Next defaults to MTP and also supports its pinned DFlash2 checkpoint; Qwen3.8 dense uses DFlash2. Overrides `[models."MODEL"] speculative_method`, then `[defaults] speculative_method`. The selection is recorded in preparation and rebuild metadata; native serving replays the exported graphs. | Recipe declaration |
394395
| `--no-mtp` | Disable multi-token speculation: load no draft model, keep no verification rows, and decode one token per step. | Speculation on |
395-
| `--draft-head-quantization {nvfp4,mxfp8}` | Prepare a separate proposal vocabulary head while retaining target-head precision. NVFP4 supports DeepSeek V4/V4.1; MXFP8 supports DeepSeek V4.1. Qwen3.8 dense supports both when the source checkpoint head is BF16, including `qwen3.8-27b-qad`; the packed NVIDIA head cannot be converted through this option. This does not quantize the draft body. | Checkpoint precision |
396+
| `--draft-head-quantization {nvfp4,mxfp8}` | Prepare a separate proposal vocabulary head while retaining target-head precision. NVFP4 supports DeepSeek V4/V4.1; MXFP8 supports DeepSeek V4.1. Qwen3.8 dense supports both when the source checkpoint head is BF16, including `qwen3.8-27b-qad`; the packed NVIDIA head cannot be converted through this option. Flash Next supports both with `--speculative-method dflash2` and a BF16 target vocabulary head. This does not quantize the draft body. | Checkpoint precision |
396397
| `--mtp-draft-tokens N` | Most draft tokens verified in one round; see `MODEL --help` for its limit. Verification capacities cover counts from 0 through `N` within the speculative token capacity. Each complete speculative program records its actual proposal depth independently of target width and aggregate rows. Prefix selection uses measured complete-program cost and completed acceptance or confidence. Zero retains draft state with width-one execution; `--no-mtp` excludes draft weights and state. | Recipe declaration |
397398
| `--speculative-token-capacity N` | Maximum physical target-token rows in wider speculative graphs; width one covers every request bucket. Applies to every speculative recipe. | 128 except Gemma 4, which derives it from concurrency and draft depth (6 at its defaults) |
398399
| `--speculative-graph-budget N` | Maximum packed round graphs to capture, independently of physical target-token rows. Excludes mandatory width-one and ordinary graphs; requires packed adaptive learned speculation at positive draft depth. Must cover both physical-row endpoints for every eligible request bucket. | 64 for packed recipes |
@@ -408,8 +409,9 @@ explicit flags take precedence.
408409
| `--compile-workers N` | CPU compiler process budget per host, divided among local preparation ranks and bounded by available RAM and device memory. Applies to b12x kernel compilation and complete-forward compilation. Overrides `[preparation] compile_workers`. | Automatic |
409410
| `--tuning-gpus LIST` | Local kernel-tuning GPUs. Start with the serving GPU list, then append complete replicas of that TP group on matching hardware. Serving TP size stays fixed. Command line only; cannot be combined with multi-host `--ranks`. | Serving GPUs |
410411
| `--no-compile` | Development only: capture graphs without compiling forwards. A recipe never sets it. | off |
412+
| `--force-moe-decode-a16` / `--no-force-moe-decode-a16` | Prepare Qwen3.8 Flash NVFP4 target routed experts as one packed weight owner: decode and speculative verification use W4A16 at every token-row count, and separate prefill graphs use calibrated A4. Requires SM120/121, checkpoint activation scales, and rank-local expert intermediate width divisible by 64 (TP1/2). Mixed prefill and other recipes/precisions reject the flag. Shared experts and stored A16 MTP experts retain their precision. Overrides `[preparation] force_moe_decode_a16`; the resolved value is recorded in the plan and package rebuild metadata. | off |
411413
| `--native-dense-activations` | Replace dense projection `auto` precision with b12x `quantized`: A8 for FP8 weights and A4 for FP4 weights. Explicit recipe precision and routed MoE precision are retained. NVFP4 requires a checkpoint activation global scale; unsupported native paths fail preparation. | Recipe precision |
412-
| `--online-mxfp8` | Store eligible BF16 projections as MXFP8 while loading. Kimi K3 converts KDA projections, MLA latent/gate, query and output projections, shared experts, MoE latent down/up projections, and DSpark context, outputs and MLPs; batched MLA key/value projections retain BF16. Qwen3.8 dense converts target attention/MLP projections and DFlash2 context, attention and MLP projections. Existing packed weights, normalization, convolutions, selectors and target vocabulary weights retain their precision. Recipes that name no such projections reject it. | Checkpoint precision |
414+
| `--online-mxfp8` | Store eligible BF16 projections as MXFP8 while loading. Kimi K3 converts KDA projections, MLA latent/gate, query and output projections, shared experts, MoE latent down/up projections, and DSpark context, outputs and MLPs; batched MLA key/value projections retain BF16. Qwen3.8 dense converts target attention/MLP projections and DFlash2 context, attention and MLP projections. Flash Next with `--speculative-method dflash2` converts only DFlash2 projections; target precision is retained. Existing packed weights, normalization, convolutions, selectors and target vocabulary weights retain their precision. Recipes that name no such projections reject it. | Checkpoint precision |
413415
| `--language-model-only` | Prepare no vision encoder and serve text only, as vLLM's flag of the same name; requests with images are rejected and the encoder's weights, scratch, and graphs become KV pages. Recipes without a separable encoder reject it. | The recipe's image support |
414416

415417
Speculative token capacity and packed graph budget also accept configuration-file
@@ -1145,8 +1147,10 @@ api_key = "" # Empty: no authentication.
11451147
[defaults]
11461148
model = "qwen3.8-flash"
11471149
max_concurrency = 4
1150+
# speculative_method = "mtp" # Flash Next also supports dflash2.
11481151

11491152
[preparation]
1153+
# force_moe_decode_a16 = false # Qwen Flash NVFP4: A16 decode/verify, calibrated A4 prefill.
11501154
# compile_workers = 24 # Optional CPU process budget per host.
11511155
# autotune_effort = "default" # off, default, audit.
11521156
# collective_tuning = "static" # static, runtime; omitted uses the recipe default.

‎docs/model-preparation.md‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,31 @@ rectangular: the default catalog covers five drafts per request through batch 16
440440
and three drafts at batch 32. Already exported packages retain their captured
441441
capacities.
442442

443+
`ignition export qwen3.8-flash --speculative-method dflash2` selects
444+
[`incoai/Qwen3.8-Flash-Next-DFlash2`](https://huggingface.co/incoai/Qwen3.8-Flash-Next-DFlash2)
445+
at revision `a22a1744e85c17364dcc54a6a420af7135d4c7f6` instead of the checkpoint's
446+
MTP module. Draft proposals use an eight-token block and up to seven verified
447+
drafts; the default remains five drafts and 128 target-token rows. The target
448+
captures completed zero-based layers 5, 15, 25, 35, and 45 after averaging its four
449+
hyperconnection streams. Five independent draft attention layers retain the last
450+
2,048 accepted target positions. These rings are included in continuation
451+
snapshots and restored with the target's GDN, QSA, PLE, and position state.
452+
Mixed prefill is rejected; separate prefill and verification graphs retain the
453+
same target arithmetic. `--online-mxfp8` converts the draft's eligible BF16
454+
projections, and `--draft-head-quantization nvfp4` or `mxfp8` prepares a separate
455+
proposal vocabulary head. Target weights retain checkpoint precision.
456+
457+
`ignition export qwen3.8-flash --force-moe-decode-a16` prepares the NVFP4 target
458+
routed experts as one packed W4A16 weight owner. Decode and speculative
459+
verification use A16 at every row count; separate prefill graphs use calibrated
460+
A4 from the checkpoint activation scales. The option requires SM120/121 and
461+
rank-local intermediate width divisible by 64 (TP1/2); mixed prefill is rejected
462+
because its prompt and decode rows share one expert call. Shared experts and MTP
463+
retain their declared precision. `[preparation] force_moe_decode_a16 = true`
464+
sets the default, and `--no-force-moe-decode-a16` overrides it. Preparation records
465+
the resolved option in the plan and rebuild metadata; compatible weight reuse
466+
requires the packed layouts and activation scales prepared for that option.
467+
443468
`ignition export qwen3.8-flash --gpus 0,1 --shard-hyperconnections` partitions
444469
the 320-wide hyperconnection bottleneck across ranks, including target and MTP
445470
final mixers. Each rank loads its down-projection rows and up-projection columns;

‎helios/proto/runtime_package.proto‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,8 @@ message RebuildManifest {
112112
string collective_tuning = 25;
113113
// Preparation-only offline geometry and runtime profile recommendations.
114114
string collective_ladder_json = 26;
115+
// Preparation-only routed W4A16 decode/verification with calibrated A4 prefill.
116+
bool force_moe_decode_a16 = 27;
115117
}
116118
message RebuildSpeculation {
117119
uint32 draft_tokens = 1;
@@ -126,6 +128,8 @@ message RebuildSpeculation {
126128
optional uint64 graph_budget = 10;
127129
// Preparation arithmetic and communication choice; serving replays captured graphs.
128130
bool vocabulary_parallel_sampling = 11;
131+
// Preparation-only learned proposal method; serving replays captured graphs.
132+
string method = 12;
129133
}
130134
// Views of retained prepared weight storage, with no additional payload copies.
131135
// Names are paths in the recipe's weight objects, not original checkpoint keys.

‎ignition/cli/__init__.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,8 @@ def _preparation_arguments(running, *, model_default="defaults.model"):
161161
speculative = running.add_argument_group("speculative decoding",
162162
"MTP flags configure speculative decoding, including DSpark and DFlash.\n"
163163
"Recipe defaults and limits: ignition export MODEL --help.")
164+
speculative.add_argument("--speculative-method", choices=("mtp", "dflash2"),
165+
help="learned proposal method; supported choices and the default depend on MODEL")
164166
speculative.add_argument("--no-mtp", action="store_true",
165167
help="disable multi-token speculation: load no draft model and decode one token"
166168
" per step; cannot accompany other speculation flags")
@@ -179,13 +181,16 @@ def _preparation_arguments(running, *, model_default="defaults.model"):
179181
" excludes mandatory width-one and ordinary graphs. Requires packed adaptive"
180182
" learned speculation (default: see MODEL --help)")
181183
advanced = running.add_argument_group("advanced options")
184+
advanced.add_argument("--force-moe-decode-a16", action=argparse.BooleanOptionalAction, default=None,
185+
help="use W4A16 target routed experts for decode and verification, with calibrated A4 prefill;"
186+
" requires Qwen3.8 Flash NVFP4 on SM120/121 and separate prefill graphs")
182187
advanced.add_argument("--native-dense-activations", action="store_true",
183188
help="use quantized activations instead of auto precision for dense projections:"
184189
" A8 for FP8 weights and A4 for FP4 weights; explicit recipe precision is retained")
185190
advanced.add_argument("--draft-head-quantization", choices=("nvfp4", "mxfp8"),
186191
help="prepare a separate quantized vocabulary head for learned draft proposals;"
187192
" target verification retains its precision (NVFP4: DeepSeek V4/V4.1;"
188-
" MXFP8: DeepSeek V4.1; both: Qwen3.8 dense with a BF16 checkpoint head;"
193+
" MXFP8: DeepSeek V4.1; both: Qwen3.8 dense and Flash Next DFlash2 with a BF16 checkpoint head;"
189194
" default: checkpoint precision)")
190195
advanced.add_argument("--markov-head-quantization", choices=("nvfp4", "fp8"),
191196
help="prepare the draft Markov transition projection as NVFP4"
@@ -235,7 +240,7 @@ def _preparation_arguments(running, *, model_default="defaults.model"):
235240
help="store the recipe's eligible BF16 projections as MXFP8 while loading"
236241
" (Kimi K3 KDA, MLA except batched key/value, shared"
237242
" experts, MoE latent projections, and DSpark draft context, outputs and MLPs;"
238-
" Qwen3.8 dense target attention/MLPs and DFlash2 context, attention and MLPs;"
243+
" Qwen3.8 dense target attention/MLPs and DFlash2 context, attention and MLPs; Flash Next DFlash2 projections;"
239244
" default: checkpoint precision)")
240245
advanced.add_argument("--language-model-only", action="store_true",
241246
help="prepare no vision encoder and reject images; its weights and scratch become KV"

‎ignition/cli/bench.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ def run(args):
150150
compile_workers=args.compile_workers, tuning_gpus=args.tuning_gpus,
151151
autotune_effort=args.autotune_effort, collective_tuning=args.collective_tuning, collective_ladder=args.collective_ladder,
152152
lock_gpu_clocks=args.lock_gpu_clocks, native_dense_activations=args.native_dense_activations,
153+
force_moe_decode_a16=args.force_moe_decode_a16,
153154
shard_hyperconnections=args.shard_hyperconnections, joint_tuning=args.joint_tuning,
154155
image_concurrency=args.image_concurrency, mixed_prefill=args.mixed_prefill,
155156
prompt_logprobs=args.prompt_logprobs, expert_scale_compression=args.expert_scale_compression,
@@ -221,7 +222,7 @@ def _measure(model, plan, configuration, settings, progress, *, compile_forwards
221222
"requested_max_concurrency": plan.concurrency, "prepared_request_capacity": max(settings.batches),
222223
"preparation": {key: preparation[key] for key in (
223224
"tensor_parallel", "prefill_chunk", "stream_overlap", "context_parallel",
224-
"native_dense_activations", "autotune_effort", "collective_tuning", "collective_ladder", "compile_workers",
225+
"native_dense_activations", "force_moe_decode_a16", "autotune_effort", "collective_tuning", "collective_ladder", "compile_workers",
225226
"speculative_token_capacity", "speculation", "expert_scale_compression", "mla_compute_mode",
226227
"online_mxfp8", "language_model_only")}}
227228

‎ignition/cli/config.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,20 +21,21 @@
2121
"publishing": {"signing_key": str, "namespace": str, "official_source": str},
2222
"server": {"host": str, "port": int, "api_key": str, "thinking_token_budget": int},
2323
"defaults": {"model": str, "max_concurrency": int, "image_concurrency": int,
24-
"speculative_token_capacity": int, "speculative_graph_budget": int},
24+
"speculative_token_capacity": int, "speculative_graph_budget": int, "speculative_method": str},
2525
"memory": {"reserve": str},
26-
"preparation": {"compile_workers": int, "autotune_effort": str, "collective_tuning": str, "checkpoint_token_file": str},
26+
"preparation": {"compile_workers": int, "autotune_effort": str, "collective_tuning": str, "checkpoint_token_file": str, "force_moe_decode_a16": bool},
2727
"prefix_cache": {"memory": str, "disk_limit": str, "drain_on_shutdown": str, "max_age": str},
2828
"paths": {"hub_cache": str, "kernel_cache": str, "package_cache": str, "export_staging": str,
2929
"prefix_cache": str, "cuda": str},
3030
}
3131
_MODEL_KEYS = {"max_concurrency": int, "image_concurrency": int,
32-
"speculative_token_capacity": int, "speculative_graph_budget": int}
32+
"speculative_token_capacity": int, "speculative_graph_budget": int, "speculative_method": str}
3333
# Empty prefix-cache values leave the decision to Flare, which derives them from the host.
3434
_BUILT_IN = {
3535
"server.host": "127.0.0.1", "server.port": 8000, "server.api_key": "",
3636
"server.thinking_token_budget": None,
3737
"memory.reserve": "0",
38+
"preparation.force_moe_decode_a16": False,
3839
"preparation.compile_workers": 0,
3940
"preparation.autotune_effort": "default",
4041
"preparation.collective_tuning": None,
@@ -176,10 +177,12 @@ def write_initial(self, model: str, max_concurrency: int) -> None:
176177
f'# model = "{model}"',
177178
f"# max_concurrency = {max_concurrency}",
178179
'# image_concurrency = 4 # active image requests; bounded by max_concurrency',
180+
'# speculative_method = "mtp" # proposal method; Qwen Flash also supports dflash2',
179181
'# speculative_token_capacity = 8 # wider target graphs require batch * width within this limit',
180182
'# speculative_graph_budget = 64 # packed round graphs; excludes width-one and ordinary graphs',
181183
"", "[preparation]",
182184
'# compile_workers = 24 # CPU processes per host; default: automatic, bounded by memory',
185+
'# force_moe_decode_a16 = false # Qwen Flash NVFP4 decode/verify A16, calibrated A4 prefill',
183186
'# autotune_effort = "default" # off, default, audit',
184187
'# collective_tuning = "static" # static, runtime; omitted uses recipe default',
185188
"", "[memory]",

‎ignition/cli/display.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,5 @@ def speculation(settings, *, brief: bool = False) -> str:
3939
precision = " · " + ", ".join(heads) if heads else ""
4040
if brief:
4141
return depth + selection + precision + parallel
42-
return f"{depth} per round · MTP module" + selection + precision + parallel
42+
source = "DFlash2" if settings.method == "dflash2" else "MTP module"
43+
return f"{depth} per round · {source}" + selection + precision + parallel

‎ignition/cli/help.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ def recipe_help(model=None):
4242
speculation = recipe.speculation
4343
if speculation is not None:
4444
lines.extend((
45+
*([] if speculation.method is None else [
46+
f" --speculative-method: {', '.join(dict.fromkeys((speculation.method, *recipe.speculative_methods)))}; default {speculation.method}"]),
4547
f" --mtp-draft-tokens: 0..{recipe.max_draft_tokens}; default {speculation.draft_tokens}",
4648
f" Verification: {'adaptive' if speculation.adaptive else 'fixed width'} by default;"
4749
" --no-mtp excludes draft weights and state.",

0 commit comments

Comments
 (0)