feat(qwen3): validate DSpine draft configs and refuse them until the forward lands by sysuTobo · Pull Request #1172 · pegainfer-project/pegainfer

sysuTobo · 2026-10-19T15:00:09Z

Summary

First slice of #1169: parse and validate a DSpine drafter's config, and refuse
to serve it until the bridge forward lands.

Detect DSpine draft checkpoints from dflash_config.selected_prefix_code_bridge.
Validate the bridge and serial-refinement invariants against the Tencent
reference (tencent/Qwen3-8B-DSpine-b16).
Reject unsupported or partially-specified bridge configs instead of silently
serving them as plain DFlash.
Refuse the checkpoint in from_runtime_with_decode_environment before the
target loads, so a DSpine path never profiles or allocates KV and never
reaches the drafter loader.

No DSpine weights are loaded, no memory is billed, and no draft computation
path changes in this PR. The bridge reservation lands with the allocation that
defines its layout, where it can be checked against measured memory.

What changed

`pegainfer-qwen3/src/config.rs`

DFlashConfig.dspine: Option<DSpineConfig> (marker) and uses_dspine()
DFlashConfig::from_json (testable entry point; from_file delegates)
DFlashConfig::ensure_servable() refuses DSpine checkpoints
DSpine validation:
- causal draft attention (draft_block_causal = true, is_causal != false)
- full-attention layers only, no sliding window
- anchor-drop layout (no num_anchors, no markov head)
- terminal-disabled bridge, independent writes
- strictly increasing 1-based feature layers, last layer active
- no reliability probes or next-token rows
- refinement width from nested selected_prefix_last_write_refine or
  top-level code_bridge_second_pass; both must agree when present, with
  mode = serial_table, first_exact = L-1, power-of-two head_topk in [2, 64]
- bridge rank a non-zero multiple of 128; positive code_size and eps
- no DFlash2 convolution or selector fields next to a bridge
- head_dim == 128, hidden_size % 512 == 0
DSpineConfig deliberately carries no dimensions yet; they land with the
loader and forward that read them

`pegainfer-qwen3/src/executor.rs`

Parse the draft config and call ensure_servable() before the target model
loads and before KV profiling

Tests

CI run 40206564772:
Rust formatting, CPU Clippy, CPU unit tests, Simulated frontend E2E, and every
model-line CUDA Clippy (sm_80) pass
config::tests:
- the real tencent/Qwen3-8B-DSpine-b16 config parses and validates against Qwen3-8B
- plain DFlash stays non-DSpine and servable
- invalid bridge, schema and attention-shape cases are rejected
- the DSpine refusal is pinned (ensure_servable errors, plain DFlash does not)
- head_dim / hidden_size target constraints are rejected

Notes

The next PR lands the kernel/FFI skeleton, flat-2D 3D loader, loading, forward,
and bridge reservation as one slice. The probabilistic sampled chain remains a
separate follow-up as tracked in #1169.

github-actions · 2026-10-23T15:36:27Z

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

Detect DSpine bridge checkpoints from dflash_config.selected_prefix_code_bridge, validate the bridge and serial-refinement invariants, and reserve the bridge weights, TG/TSV tables, and lane scratch before KV pool sizing. Reject unsupported bridge configs loudly, including the flat schema, so a DSpine checkpoint cannot silently load as plain DFlash. Refs pegainfer-project#1169 Signed-off-by: bohaotu <bohaotu@pony.ai>

FeathBow

The bridge weights billed here match the release's safetensors header, and the config rules follow the Tencent reference. But nothing consumes the parsed config yet: a valid DSpine checkpoint still loads and drafts as plain DFlash, now with about 926 MiB less KV. Until the bridge forward lands, could the draft load refuse uses_dspine()?

FeathBow · 2026-10-24T16:28:17Z

        // per-step bias is the dominant term). Zero for plain DFlash drafters.
        let markov = MarkovHead::reservation_bytes(config, max_decode_batch_size)?;
        let selector = SelectorHead::reservation_bytes(config, max_decode_batch_size)?;
+        let dspine = dspine_bytes(config, max_decode_batch_size)?;


Could this move into the PR that allocates the bridge? The scratch layout is not written yet and the pinned number restates the formula. Next to the allocation it can be checked against measured memory.

FeathBow · 2026-10-24T16:28:17Z

+    /// DSpine checkpoints nest this under `dflash_config`; accept the flat form too.
+    #[serde(default)]
+    draft_block_causal: Option<bool>,
+    /// Parsed only to reject the unsupported flat DSpine schema loudly.


The reference reads the bridge and draft_block_causal only from dflash_config and has no flat schema, so both top level fields can go.

FeathBow · 2026-10-24T16:28:17Z

+        );
+
+        ensure!(
+            bridge.code_size > 0 && bridge.code_size.is_multiple_of(128),


The reference has no code_size constraint. This can go, or land with the kernel that needs it.

FeathBow · 2026-10-24T16:28:18Z


+impl DSpineConfig {
+    /// `Ok(None)` for a plain DFlash/DSpark checkpoint; `Ok(Some(_))` when the
+    /// checkpoint carries a code bridge. Every DSpine-only field is validated


The reference also rejects selected_prefix_feature_reliability, selected_prefix_feature_next_token_rows, conv or selector fields next to a bridge, and a code_bridge_second_pass whose head_topk disagrees. The fields left unchecked (kappa, selected_prefix_decode_mode, the teacher fields, refine version and weight) are training only, so the doc comment should say that.

A parsed bridge is not consumed yet, so a DSpine checkpoint would still load as plain DFlash while its reservation shrank the KV pool. Refuse the load until the bridge forward lands, drop the flat-schema fields the reference never reads, keep only the kernel's rank/hidden constraints, and add the reference's reliability/next-token/conv/selector and second-pass agreement checks. The bridge scratch reservation moves to the allocation PR so it can be checked against measured memory; the checkpoint-sized weights stay. Signed-off-by: bohaotu <bohaotu@pony.ai>

FeathBow

The plain DFlash fallback is closed, but the refusal comes after the resource work. Startup loads the target, adds the DSpine reservation (about 805 MiB of weights and tables) to the KV margin, profiles and allocates the pool, and only then loads the drafter and refuses. On a card with little headroom the profile fails first with "memory profile leaves no room for KV cache", which hides the actual reason. Could the refusal move into the config check before the target loads, and the reservation land with the allocations it bills, where it can be compared with measured memory?

None of the new tests reach the refusal: with the ensure! in loading.rs removed, every path they cover stays the same. A test on the early refusal would pin it.

The reference also rejects an explicit is_causal = false, layer_types other than full attention and a sliding window; this parser does not read those yet. The description still describes the first commit: the scratch term and the 970,933,248 B total, the code_size rule, the flat bridge rejection, and no draft code path changes.

FeathBow · 2026-10-25T12:34:28Z

+    pub(crate) rank: usize,
+    // Read by the bridge forward/refinement and the weight loader in the next slice.
+    #[allow(dead_code)]
+    pub(crate) eps: f32,


eps, feature_layers and head_topk have no reader in this PR. The refusal only needs bridge detection, so these can land with the loader and forward that read them.

FeathBow · 2026-10-25T12:34:28Z

        self.markov_rank > 0
    }

+    /// DSpine code bridge is active; the draft forward injects per-layer


This says the forward injects features, while every load refuses. The refusal message in loading.rs and the doc comment on from_inner also claim near zero acceptance without a measurement; saying the checkpoint is unsupported is enough.

FeathBow · 2026-10-25T12:34:28Z

+                self.hidden_size
+            );
+            anyhow::ensure!(
+                !self.anchor_first(),


from_json already rejects num_anchors and a nonzero markov_rank, the only inputs that set anchor_first, so this check and the test that sets the field by hand can go.

FeathBow · 2026-10-25T12:34:28Z

+/// Top-level refinement declaration some DSpine checkpoints carry instead of
+/// (or alongside) `dflash_config.selected_prefix_last_write_refine`.
+#[derive(Clone, Debug, Deserialize)]
+struct RawDSpineSecondPass {


The reference raises on unknown code_bridge_second_pass keys. #[serde(deny_unknown_fields)] here would match it, so a key added later that changes the pass gets refused at load.

…rly refusal Move the memory reservation to the PR that allocates the bridge, so the resolved bridge dimensions leave DSpineConfig and a marker remains. Run the refusal before the target loads, where a test can pin it, and reject explicitly non-causal configs, non-full layer types and sliding windows as the reference does. Deny unknown second-pass keys. Signed-off-by: bohaotu <bohaotu@pony.ai>

github-actions Bot added qwen3 Qwen3 model crate (pegainfer-qwen3) enhancement New feature or request labels Oct 19, 2026

sysuTobo marked this pull request as ready for review October 20, 2026 03:26

sysuTobo force-pushed the feat/dspine-config branch from a113cbf to 4d67f53 Compare October 20, 2026 04:13

github-actions Bot added the needs-rebase Conflicts with main, needs a rebase label Oct 23, 2026

sysuTobo force-pushed the feat/dspine-config branch from 4d67f53 to 5237b4a Compare October 23, 2026 17:05

github-actions Bot removed the needs-rebase Conflicts with main, needs a rebase label Oct 23, 2026

FeathBow requested changes Oct 24, 2026

View reviewed changes

sysuTobo requested a review from FeathBow October 25, 2026 05:49

FeathBow requested changes Oct 25, 2026

View reviewed changes

sysuTobo marked this pull request as draft October 25, 2026 15:10

sysuTobo marked this pull request as ready for review October 25, 2026 15:21

sysuTobo changed the title ~~feat(qwen3): parse and reserve memory for DSpine drafters~~ feat(qwen3): validate DSpine draft configs and refuse them until the forward lands Oct 25, 2026

Navigation Menu

Repository navigation

feat(qwen3): validate DSpine draft configs and refuse them until the forward lands - #1172

feat(qwen3): validate DSpine draft configs and refuse them until the forward lands#1172
sysuTobo wants to merge 3 commits into
pegainfer-project:mainfrom
sysuTobo:feat/dspine-config

sysuTobo commented Oct 19, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 23, 2026

Uh oh!

FeathBow left a comment

Uh oh!

FeathBow Oct 24, 2026

Uh oh!

FeathBow Oct 24, 2026

Uh oh!

FeathBow Oct 24, 2026

Uh oh!

FeathBow Oct 24, 2026

Uh oh!

FeathBow left a comment

Uh oh!

FeathBow Oct 25, 2026

Uh oh!

FeathBow Oct 25, 2026

Uh oh!

FeathBow Oct 25, 2026

Uh oh!

FeathBow Oct 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

sysuTobo commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

pegainfer-qwen3/src/config.rs

pegainfer-qwen3/src/executor.rs

Tests

Notes

Uh oh!

github-actions Bot commented Oct 23, 2026

Uh oh!

FeathBow left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FeathBow left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sysuTobo commented Oct 19, 2026 •

edited

Loading

`pegainfer-qwen3/src/config.rs`

`pegainfer-qwen3/src/executor.rs`