Repository navigation
MoE drafts: DeepSeek-V4's own DSpark draft architecture (stacked on #1056) - #1057
Draft
yushengsu-thu wants to merge 41 commits into
Draft
yushengsu-thu wants to merge 41 commits into
yushengsu-thu wants to merge 41 commits into
Conversation
The MoE layer, experts, top-k router, the qwen3_5_moe family and the build_ffn switch, with the Qwen3-8B MoE DSpark config and recipe. No load balancing: the router is stateless and the trainer is untouched.
Router fixes what every family's router provides, forward(x) -> scores, indices, weights, and is the point where other families (DeepSeek) add their own routing; TopKRouter keeps the algorithm. Load balancing builds on this contract in a later PR.
DSpark accepts use_liger_kernel since sgl-project#963. Liger replaces the draft's RMSNorms; the MoE FFN keeps its grouped_mm experts, which Liger does not cover yet.
SharedExpert.reset_parameters, and MoELayer.reset_parameters covers the shared expert and its gate too, so resetting the layer means the whole layer. HF's init already reached these Linears; the values are the same.
Within an MoE draft (moe_preset), ffn_layer_types marks the layers that keep the dense MLP (intermediate_size); the others take the family's MoE layer (moe_intermediate_size per expert). No list means every layer is MoE, as before. The DSpark MoE config trains ["dense", "moe"].
…he score Router implements the shared steps, top-k selection and the combine weights (renormalized, scaled), and forward always returns scores, indices, weights. A family overrides the steps it needs: TopKRouter adds the linear gate and its softmax score; a DeepSeek-style router would override score (fp32 sigmoid) and select (group-limited). MoELayer takes any Router. TopKRouter's numerics are unchanged.
The step that gathers, renormalizes and scales the chosen experts' scores returns the weights that combine their outputs; routing_weights says so (combine did not).
Follows the draft-layer layout of sgl-project#1025: the MoE FFN is a common FFN component next to the dense one, in specforge/modeling/draft/layers/ffn/moe. A pure move plus import paths; layers/__init__.py and layers/ffn/__init__.py stay empty so sgl-project#1025's versions merge in cleanly.
Sigmoid top-k sums can underflow to 0 and turn the weights into NaN; HF DeepSeek-V3/V4 and torchtitan's route_norm_epsilon guard it with 1e-20. For softmax the sum is at least K / E, so the addition is a bitwise no-op in fp32: Qwen parity and the e2e losses are unchanged.
The shipped draft config now matches the 8 x 1536 experiment: 5 layers, 32 query / 8 KV heads, every FFN an MoE layer (8 experts, top-1), and the recipe trains for 6 epochs with Liger on. ffn_layer_types stays explicit so the per-layer choice is visible; the dense + MoE mix is still covered by the tests.
The DFlash-family model passes moe_token_mask, the kept-block mask expanded to draft tokens, on every draft call; the draft and its decoder layers pass it on, and only an MoELayer takes it. Nothing reads it yet: the load balancing that follows counts with it. The stand-in drafts in the loss tests accept the argument.
TrainerCore.register_step_hook takes a StepHook: before_backward adjusts each micro-batch's StepOutput before the backend sees it, on_step runs after every optimizer step and its metrics join that step's. Nothing registers one yet; MoE load balancing will, once, at assembly.
A fp32 expert_bias buffer (None unless a family asks for one) shifts the top-k picks, never the weights; bias_update_rate 0 keeps a loaded bias fixed; update_expert_bias is one DeepSeek-V3 sign step with a zero-mean delta, called with the step's pooled kept-slot counts.
A non-persistent int64 tokens_per_expert buffer; in training each forward adds its kept slots (moe_token_mask) with scatter_add_. Integer buffers survive dtype casts and stay exact; reset_parameters zeroes the counts since from_pretrained builds on the meta device.
MoEBalancing finds a model's MoELayers by type and, at every optimizer step, all-reduces their counts in one collective, hands each router its pooled load through update_expert_bias, resets the counts and reports moe/step_balance (1 when uniform, E / K when collapsed). trainer.py registers it when the model has MoE layers. MoE drafts refuse activation checkpointing: the layers count in forward, a recompute would count twice.
…ecipe turns them on dflash_config.moe_router_enable_expert_bias and moe_router_bias_update_rate (0 keeps a loaded bias fixed), both off by default.
Router contract and bias step, the layer's masked counts, the step hook with a mocked all-reduce, the mask through DFlash and DFlash2 layers, TrainerCore running the registered hook, a 2-rank pooled update, the launch gate's bias assertions, the host-sync MoE case with the bias, and a from_pretrained round trip of the bias with zeroed counts.
select now receives the shifted scores (choice) and only picks the top-k, so a family that overrides select, such as DeepSeek-V3's group-limited routing, needs no knowledge of the bias. The weights still come from the unbiased scores. No behavior change: the bitwise probe against sgl-project#1003 stays identical.
MoELayer takes aux_loss_coef and aux_loss_type (microbatch_wise, or the default batch_wise over the step's counts pooled across ranks) and, in training with a non-zero coefficient, leaves this forward's loss in aux_loss for the balancing hook to fold before backward. At coefficient 0 nothing is computed: no graph and no per-micro-batch collective.
… step hook MoEBalancing.before_backward collects each MoE layer's loss times its coefficient, adds the sum to the loss and, when the strategy reports loss_terms, to the numerator scaled by the denominator, and logs the raw mean as moe/aux_loss. Folded rather than injected: SpecForge's trainer normalizes gradients by the global denominator, so an injected gradient would need a per-micro-batch scale handed to the layer before backward.
…pe trains with batch_wise 0.005 dflash_config.moe_router_aux_loss_coef and moe_router_aux_loss_type; the default coefficient 0 keeps drafts without the keys loss-free.
The loss against transformers' load_balancing_loss_func with and without the mask, the layer's two load types and its skip at coefficient 0 (no collective), the hook's fold with and without loss_terms, the gradient parity of the hook's fold against an objective-side fold through TrainerCore, a 2-rank batch-wise load, the launch gate's aux metric, the host-sync MoE case with the loss on, and the family's aux knobs.
build_moe assembles DeepSeek-V4's MoE layer from its config keys on the stack's parts: TopKRouter with sqrt(softplus) scores, routed_scaling_factor as route_scale and the expert bias, clamped SwiGLU (swiglu_limit) for the routed experts and the one ungated shared expert. As in transformers, the top-k weights are always renormalized. No new router class.
moe_preset: "deepseek_v4" now reaches the family through build_ffn.
The family's block against transformers' DeepseekV4SparseMoeBlock on the grouped_mm and loop paths (same weights, a non-zero expert bias; only two state-dict names differ), the CUDA bf16 case for both families, and the config reading: sqrtsoftplus scores, route_scale, the shared expert, the always-on renormalization, the clamp, the refusals, the preset.
Pure relocation, no behavior change. Split from sgl-project#1023. - specforge/layers/{embedding,linear,lm_head,ring} -> specforge/modeling/draft/layers/ - modeling/draft/{flex_attention,flex_attention_backend}.py -> layers/attention/backends/ - generate_eagle3_mask -> modeling/draft/eagle3_mask.py - modeling/draft/llama3_eagle.py -> modeling/draft/eagle3.py - LlamaMLP -> layers/ffn/dense.py:GatedMLP (re-exported from eagle3 as LlamaMLP) - tests/test_layers -> tests/test_modeling/test_layers (test_decoder.py -> test_ring_attention.py) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
specforge.layers, modeling/draft/{llama3_eagle,flex_attention,flex_attention_backend}
now resolve to the relocated modules (same module objects, so registries and
class identity are preserved) and emit a DeprecationWarning. This commit can
be dropped if a clean break is preferred.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Split from sgl-project#1023. DFlash-family attention moves into specforge/modeling/draft/layers: - attention/base.py: Attention/AttentionContext/AttentionOutput contract and DenseAttention (projections + cache update + backend dispatch + o_proj) - attention/{gqa,mla,rotary}.py: GQA/MHA and expanded-MLA projections, RoPE - attention/backends/dense.py: eager/SDPA/Flash/Flex dispatch, unchanged from Qwen3DFlashAttentionBase.forward - config.py (LayerKernels, AttentionConfig), factory.py (build_attention, build_ffn), decoder.py (pre-norm DecoderLayer) Qwen3DFlashAttention/Qwen3DFlashMLAAttention remain as compatibility constructors; DFlashKernels aliases LayerKernels. Mask handling is unchanged. The parity test embeds the pre-refactor classes verbatim and checks outputs, gradients and two AdamW steps bit for bit across DFlash/DFlash2/DSpark/Domino, GQA/MLA, eager/SDPA, sliding and causal layouts, plus cached block decoding. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Split from sgl-project#1023. DecoderLayer (and DFlash2's convolution decoder) read and write the residual stream through Residual modules: - residual/standard.py: the existing pre-norm add; parameter-free, so default checkpoints, parameter order and numerics are unchanged (parity test still bit-exact) - residual/gr.py: Qwen3.8-Flash-Next gated residual (per-branch Gemma RMSNorm, low-rank gated read, gated write, learned final readout) - residual/mhc.py: DeepSeek-V4.1 single-pass mHC (lagged pre-mix, FP32 coefficients and Sinkhorn, stream-zero initial read) Opt-in via draft_residual in the DFlash-family config. Non-standard residuals force non-reentrant gradient checkpointing and record their non-parameter semantics in the DFlash/Domino/DSpark resume contracts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 2, 2026
A new attention kind, ``deepseek_v4``, following transformers' DeepseekV4Attention: q_a_proj -> q_a_norm -> q_b_proj -> a parameter-free per-head RMS norm (q_b_norm); one 512-wide KV head shared by every query head, serving as both K and V (kv_proj -> kv_norm); interleaved RoPE on the trailing qk_rope_head_dim channels with the plain rope_theta (the draft layers have no YaRN); per-head sinks; and a grouped low-rank output, o_a_proj (block-diagonal, o_groups x o_lora_rank) then o_b_proj. Because K is V, the rope slice of the attention output carries the keys' rotation; the conjugate rotation at the query positions turns it into the reference's relative rotation before the output projection. The class owns its forward for that step, as transformers, torchtitan and Megatron-Core do; DenseAttention and the dense backends are unchanged. The kind is registered in AttentionConfig/build_attention and accepted as dflash_config.attention_mode; the model-level rotary carries the partial-RoPE width as it does for MLA. Tests compare the forward output and every gradient against a plain-PyTorch port of DeepSeek's inference/model.py attention on eager (CPU) and flex (CUDA), check the parameter names and the V4-Flash shapes against the checkpoint, and build DFlash and DSpark drafts with the kind. Second commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
A new attention kind, ``deepseek_v4``, following transformers' DeepseekV4Attention: q_a_proj -> q_a_norm -> q_b_proj -> a parameter-free per-head RMS norm (q_b_norm); one 512-wide KV head shared by every query head, serving as both K and V (kv_proj -> kv_norm); interleaved RoPE on the trailing qk_rope_head_dim channels with the plain rope_theta (the draft layers have no YaRN); per-head sinks; and a grouped low-rank output, o_a_proj (block-diagonal, o_groups x o_lora_rank) then o_b_proj. Because K is V, the rope slice of the attention output carries the keys' rotation; the conjugate rotation at the query positions turns it into the reference's relative rotation before the output projection. The class owns its forward for that step; DenseAttention and the dense backends are unchanged. Why its own kind rather than a variant of MLA: V4 has no kv_b_proj (the latent itself is K and V for every head, with the rope part inside V), needs the inverse rotation on the output and a grouped low-rank output projection, so the forward graph differs from MLA, not only the sizes. transformers (DeepseekV4Attention written separately from DeepseekV3Attention), torchtitan (DeepSeekV4Attention on the common Attention base, apart from DeepSeekMLAAttention), Megatron-Core (DSv4HybridAttention on Attention, not on MLASelfAttention) and SGLang (MqaAttentionBase, not DeepseekV2AttentionMLA) all keep it as a separate class for the same reason, and all let it own its forward. The kind is registered in AttentionConfig/build_attention and accepted as dflash_config.attention_mode; the model-level rotary carries the partial-RoPE width as it does for MLA. Tests compare the forward output and every gradient against a plain-PyTorch port of DeepSeek's inference/model.py attention on eager (CPU) and flex (CUDA), check the parameter names and the V4-Flash shapes against the checkpoint, and build DFlash and DSpark drafts with the kind. Second commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
yushengsu-thu
force-pushed
the
dsk-v4-official-dspark
branch
from
November 2, 2026 00:49
c39ac14 to
09c4030
Compare
``dflash_config.context_window`` limits every draft block to the last that many context keys before its anchor, the same keys for each query of the block, as DeepSeek-V4's DSpark does (get_dspark_topk_idxs: 128 cached context positions plus the whole block, for every block query). In-block visibility is unchanged and still follows ``is_causal``; the existing ``sliding_window`` keeps its per-query, FA2-style meaning, and the two intersect when both are set. The window reaches the two mask builders, the training masks of every DFlash-family layer type, the masks spec_generate builds itself (a full layer with a window now gets a mask instead of None) and the attention resume contract, which records the window only when it is set so that contracts of existing runs keep their shape. Existing configs are unaffected: the option defaults to null. Tests check the visible keys per query offset for both in-block directions, clipping at the sequence start, agreement of the dense and flex builders with the element-level reference across batches, blocks and a combined sliding window, the training-path plumbing, the config validation and the resume contract. Third commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
DeepSeek-V4 (and its DSpark draft) reads each sublayer with the pre-mix computed from that sublayer's own input and collapses the streams at the end through hc_head, a learned sigmoid read with its own fn/scale/base. The mHC shipped so far is DeepSeek-V4.1's single-pass variant, which reads with the previous sublayer's pre-mix and collapses with the last one, without head parameters. Following SGLang, which expresses the difference as one flag (HcConfig.pre_from_prev) on one mHC implementation, ResidualConfig gains ``pre_from_prev`` (default true, the existing behaviour). With false, ManifoldHyperConnection reads with its own pre-mix and carries nothing forward, and the model's residual_readout is the new HyperConnectionHead (mix_weight [streams, streams x hidden], mix_scale [1], mix_bias [streams]) before the final norm. The coefficient path, Sinkhorn and the write are shared and unchanged. build_residual_readout in the factory now owns the GR readout, the V4 head and the identity; the resume contract records the flag. Tests compare the V4 read and write and all gradients with the same reference math as the single-pass tests, the head with DeepSeek's hc_head formula (also under bf16 autocast), and run every residual model test (forward and gradients across DFlash, DFlash2 and DSpark with GQA and MLA, the DSpark objective, checkpoint recomputation, reentrant rejection, bf16 CUDA, save/load, resume contract) for the V4 variant as well. Fourth commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
configs/deepseek-v4-flash-moe-dspark.json is DeepSeek's own DSpark draft architecture for DeepSeek-V4-Flash, field for field from the official DSpark config and checkpoint: three stages of V4 attention (64 heads on one 512-wide KV head, q rank 1024, grouped output 8 x 1024, rope 64 with the plain theta, sinks), the block-anchored 128-token context window with a bidirectional block of 5, V4 mHC (pre_from_prev false, 4 streams, 20 Sinkhorn iterations), the 256-expert MoE (2048 wide, top-6, one shared expert, sqrtsoftplus routing scaled by 1.5, clamped SwiGLU, expert bias), target layers 40-42, mask token 128799, a rank-256 Markov head and the confidence head on hidden plus Markov embedding. The confidence projection gains dflash_config.confidence_head_bias (default true); V4's has none. The recipe examples/configs/online/disaggregated/external/deepseek-v4-flash-moe-dspark.yaml is the existing V4-Flash external-capture recipe with this draft. The test builds the full-size model on the meta device and checks every state-dict name and shape against the official mtp.* layout, and that the element count equals the official tensors' (three stages plus main_proj, main_norm and the head stage); a small-scale run trains every parameter through the DSpark objective. The DSpark config guard admits attention_mode deepseek_v4 with one KV head. Fifth commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
DeepSeek-V4's DSpark scores confidence on the readout after hc_head and
before the final norm (forward_head: ``logits = head(norm(x))``,
``confidence = confidence_head(x, markov_embed)``), and SGLang serves it
that way. SpecForge fed the normed hidden, so a V4 draft trained here
would have served a miscalibrated confidence.
``dflash_config.confidence_head_input`` ("post_norm", the default, or
"pre_norm") selects the input. With "pre_norm", the DSpark objective
asks the draft for the readout before the norm; the family model carries
it as one more aligned tensor through the chunked, checkpointed
objective and hands it to predict_confidence. Drafts that do not ask
keep the three-value draft-block output and never see the new keyword.
The confidence projection now computes in fp32 with autocast disabled,
as DeepSeek's does; the features may be an unnormalized readout. The V4
draft config sets confidence_head_input to pre_norm.
Tests check that the requested readout norms to the regular output, that
the family model passes the pre-norm chunk matching each logits chunk
(and nothing extra in the default mode), the validation errors, and the
fp32 projection under bf16 autocast.
Sixth commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
specforge/export/deepseek_v4_dspark.py converts DeepSeek's DSpark draft
between its official checkpoint and the SpecForge draft of
configs/deepseek-v4-flash-moe-dspark.json. The official layout keeps the
stages under mtp.{stage}.* with the inference code's names: FP8 linears
(e4m3 with one e8m0 scale per 128x128 block), FP4 routed experts (two e2m1
codes per byte along the input, low nibble first, one e8m0 scale per 32
inputs) and fp32 hyper-connection, sink and gate-bias tensors.
Import dequantizes to bf16, renames to transformers' names, fuses the
experts' w1/w3 into gate_up_proj and w2 into down_proj, checks every
tensor name and shape against the draft built on the meta device and
writes Hub-style shards; only the shards holding mtp.* are downloaded.
Export reverses the mapping, unquantized, for SGLang's
DeepseekV4ForCausalLMDSpark, with the official config minus its
quantization fields. scripts/convert_deepseek_v4_dspark.py is the CLI.
Verified on deepseek-ai/DeepSeek-V4-Flash-DSpark: 4705 tensors from three
shards become the draft's 78 (19.846B parameters) that AutoDraftModel
loads without missing or unexpected keys and runs; the FP8 and FP4
results equal SGLang's dequantization bit for bit on real tensors. Tests
cover the block and nibble conventions, a round trip, exact import of
synthetic quantized tensors, both directory conversions and the config.
Seventh commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
Unrelated to DeepSeek-V4: purely a speed change for any flex_attention draft with head_dim > 256 on Hopper and Blackwell. torch 2.13 picks 16x16 backward tiles past head_dim 256; 32x32 tiles with 8 warps and 3 stages are 3x faster on H200 and B200. On SM100 the default forward tiles fault (misaligned address) and an aliased key/value tensor makes the compiled backward 4.7x slower, so flex_kernel_options also sets 32x32 forward tiles and clones the value there. Caller kernel_options win.
Unrelated to DeepSeek-V4: purely a memory saving for FULL_SHARD with gradient accumulation. FSDP1 no_sync keeps the full unsharded gradient on every rank until the boundary micro-step (37 GiB for a 19.7B bf16 draft against 9 GiB sharded), which defeats FULL_SHARD and OOMs a 4-rank B200 run. Reduce-scattering every micro-step accumulates into the sharded gradient instead; SHARD_GRAD_OP and DDP are unchanged.
Contributor
Author
yushengsu-thu
marked this pull request as ready for review
November 3, 2026 04:27
yushengsu-thu
requested review from
FlamingoPg,
FrankLeeeee,
shuaills,
sleepcoo and
zyksir
as code owners
November 3, 2026 04:27
yushengsu-thu
force-pushed
the
dsk-v4-official-dspark
branch
from
November 3, 2026 04:27
5258e28 to
9e12b96
Compare
yushengsu-thu
marked this pull request as draft
November 3, 2026 22:58
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 4, 2026
A module that defines a per-head ``sinks`` parameter (HF's name, as in gpt-oss and DeepSeek-V4) now has it applied by the dense backend: - eager: transformers' DeepSeek-V4 eager_attention_forward, which puts the sink column into the softmax. The CPU reference path. - flex_attention on CUDA: request the logsumexp and scale the output by sigmoid(lse - sink), the same softmax over [keys, sink]; a query with no visible key keeps its zeros. The fast GPU path. - sdpa, flash_attention_* and CPU flex refuse sinks with a clear error. Tests compare the forward output and every gradient against a softmax reference with the extra column on both backends, and train the sinks alone back to hidden target values. First commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 4, 2026
``dflash_config.context_window`` limits every draft block to the last that many context keys before its anchor, the same keys for each query of the block, as DeepSeek-V4's DSpark does (get_dspark_topk_idxs: 128 cached context positions plus the whole block, for every block query). In-block visibility is unchanged and still follows ``is_causal``; the existing ``sliding_window`` keeps its per-query, FA2-style meaning, and the two intersect when both are set. The window reaches the two mask builders, the training masks of every DFlash-family layer type, the masks spec_generate builds itself (a full layer with a window now gets a mask instead of None) and the attention resume contract, which records the window only when it is set so that contracts of existing runs keep their shape. Existing configs are unaffected: the option defaults to null. Tests check the visible keys per query offset for both in-block directions, clipping at the sequence start, agreement of the dense and flex builders with the element-level reference across batches, blocks and a combined sliding window, the training-path plumbing, the config validation and the resume contract. Third commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 4, 2026
A module that defines a per-head ``sinks`` parameter (HF's name, as in gpt-oss and DeepSeek-V4) now has it applied by the dense backend: - eager: transformers' DeepSeek-V4 eager_attention_forward, which puts the sink column into the softmax. The CPU reference path. - flex_attention on CUDA: request the logsumexp and scale the output by sigmoid(lse - sink), the same softmax over [keys, sink]; a query with no visible key keeps its zeros. The fast GPU path. - sdpa, flash_attention_* and CPU flex refuse sinks with a clear error. Tests compare the forward output and every gradient against a softmax reference with the extra column on both backends, and train the sinks alone back to hidden target values. First commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 4, 2026
``dflash_config.context_window`` limits every draft block to the last that many context keys before its anchor, the same keys for each query of the block, as DeepSeek-V4's DSpark does (get_dspark_topk_idxs: 128 cached context positions plus the whole block, for every block query). In-block visibility is unchanged and still follows ``is_causal``; the existing ``sliding_window`` keeps its per-query, FA2-style meaning, and the two intersect when both are set. The window reaches the two mask builders, the training masks of every DFlash-family layer type, the masks spec_generate builds itself (a full layer with a window now gets a mask instead of None) and the attention resume contract, which records the window only when it is set so that contracts of existing runs keep their shape. Existing configs are unaffected: the option defaults to null. Tests check the visible keys per query offset for both in-block directions, clipping at the sequence start, agreement of the dense and flex builders with the element-level reference across batches, blocks and a combined sliding window, the training-path plumbing, the config validation and the resume contract. Third commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 4, 2026
``dflash_config.confidence_head_bias`` (default true, the existing behaviour) builds AcceptRatePredictor's projection without a bias when false. DeepSeek-V4.1's DSpark ships ``confidence_head.proj.weight`` and no bias, so its draft config will set the flag to false. Taken from the DeepSeek-V4 draft config commit of sgl-project#1057 without that config and its test; the fp32 confidence test already exercised the keyword.
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 4, 2026
A new attention kind, ``deepseek_v41``, following DeepSeek-V4.1's inference DSparkAttention: q_a_proj -> q_a_norm -> q_b_proj; one 512-wide KV head shared by every query head, serving as both K and V (kv_proj -> kv_norm); interleaved RoPE on the trailing qk_rope_head_dim channels with the plain rope_theta (the draft layers have no YaRN); per-head sinks; and a grouped low-rank output, o_a_proj (block-diagonal, o_groups x o_lora_rank) then o_b_proj. Unlike DeepSeek-V4, V4.1 applies no parameter-free per-head RMS norm to q after q_b_proj. Because K is V, the rope slice of the attention output carries the keys' rotation; the conjugate rotation at the query positions turns it into the reference's relative rotation before the output projection. The class owns its forward for that step; DenseAttention and the dense backends are unchanged. The kind is registered in AttentionConfig/build_attention and accepted as dflash_config.attention_mode; the model-level rotary carries the partial-RoPE width as it does for MLA. Tests compare the forward output and every gradient against a plain-PyTorch port of V4.1's inference/model.py DSparkAttention on eager (CPU) and flex (CUDA), check the parameter names and the V4.1-Flash shapes against the checkpoint (hidden 5120, q rank 1280), and build DFlash and DSpark drafts with the kind. Adapted from the DeepSeek-V4 attention kind of sgl-project#1057.
yushengsu-thu
added a commit
to yushengsu-thu/SpecForge
that referenced
this pull request
Nov 4, 2026
specforge/export/deepseek_v41_dspark.py converts DeepSeek-V4.1's DSpark
draft between its official checkpoint and the SpecForge draft of
configs/deepseek-v4.1-flash-moe-dspark.json. The official layout keeps
the stages under mtp.{stage}.* with the inference code's names: FP8
linears (e4m3 with one e8m0 scale per 32x32 block; V4 used 128x128),
FP4 routed experts (two e2m1 codes per byte along the input, low nibble
first, one e8m0 scale per 32 inputs) and fp32 hyper-connection, sink and
gate-bias tensors.
Import dequantizes to bf16, renames to transformers' names (the Markov
head's embed/head become markov_w1/markov_w2), fuses the experts' w1/w3
into gate_up_proj and w2 into down_proj, drops the image-token routing
bias gate.bias_vl that a text draft does not carry, checks every tensor
name and shape against the draft built on the meta device and writes
Hub-style shards; only the shards holding mtp.* are downloaded. Export
reverses the mapping, unquantized, writing gate.bias_vl as a copy of
gate.bias so the official layout stays loadable, with the official
config minus its quantization fields.
scripts/convert_deepseek_v41_dspark.py is the CLI.
Verified on deepseek-ai/DeepSeek-V4.1-Flash: 2401 tensors from three
shards become the draft's 75 (14.225B parameters) that AutoDraftModel
loads without missing or unexpected keys and runs; export then import
reproduces them bit for bit. The official draft's routing biases are all
zero. Tests cover the block and nibble conventions, a round trip, exact
import of synthetic quantized tensors, both directory conversions and
the config.
Adapted from the DeepSeek-V4 converter of sgl-project#1057.
This branch has not been deployed
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
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.
Summary [cherry-pick some PRs - after those PRs merges, rebase this PR first]
DeepSeek's own DeepSeek-V4 DSpark draft architecture in SpecForge, trained against
deepseek-ai/DeepSeek-V4-Flash; stacked on #1056; increment: yushengsu-thu/SpecForge@dsk-v4-moe-dspark...dsk-v4-official-dsparkWork in progress; scope under discussion. The target is the module DeepSeek ships under
mtp.*inDeepSeek-V4-Flash-DSpark: three full V4 blocks (V4 attention with the 128-token window over the projected target states, mHC hyper-connections, target-shaped MoE from #1056'sdeepseek_v4family,main_projover the last three target layers, markov and confidence heads, block size 5), sharing the target's embedding and head.Stacked: #1033 -> #1034 -> #1035 -> #1056 -> this PR.