Investigation: DeepSeek-V4-Flash onboarding blocked on missing HF reference

[Copilot]

Onboarding deepseek-ai/DeepSeek-V4-Flash (model_type: deepseek_v4) as requested in the issue. After investigation, no code changes are proposed in this PR — the task cannot be completed as a minimal, verifiable change given the current state of the upstream model. Details below for reviewer decision on scope.

Architectural delta vs. existing deepseek.py (V2/V3)

From config.json on the HF repo, V4 introduces components with no precedent in mobius:

• Hybrid attention — Compressed Sparse Attention + Heavily Compressed Attention, driven by a per-layer compress_ratios list, plus index_head_dim, index_n_heads, index_topk, compress_rope_theta. Not expressible via the existing DeepSeekMLA component.
• Manifold-Constrained Hyper-Connections (mHC) replacing residual connections, parameterized by hc_eps, hc_mult, hc_sinkhorn_iters, num_hash_layers (Sinkhorn-normalized routing between layers).
• Output LoRA + grouping — o_lora_rank=1024, o_groups=8.
• New MoE scoring — scoring_func: "sqrtsoftplus" (current gate supports softmax/sigmoid only).
• Clipped SwiGLU — swiglu_limit: 10.0.

Blocker: no HuggingFace transformers reference

The HF repo ships weights (46 shards), a technical-report PDF, and a standalone inference/model.py + custom kernel.py. It does not ship modeling_deepseek_v4.py / configuration_deepseek_v4.py, and deepseek_v4 is not registered in transformers. AutoConfig.from_pretrained(..., trust_remote_code=True) fails.

Consequences for the mobius onboarding pattern:

• _registry.py keys on HF model_type → requires an HF config class.
• tests/integration_test.py requires parity with AutoModelForCausalLM at atol=1e-4 → no reference to compare against.
• Weight-name alignment (preprocess_weights) has no canonical HF state_dict to target; only the custom inference engine's naming exists.

A scaffold (config + registry + graph builder) is technically possible but would require inventing implementations of CSA/HCA/mHC/sqrtsoftplus/Sinkhorn hyper-connections with no way to verify correctness against a reference, violating the repo's root-cause-before-remedy rule and risking silent incorrect output from build("deepseek-ai/DeepSeek-V4-Flash").

Paths forward (need reviewer input)

1. Wait for transformers support to land, then port against it (the V3 pattern).
2. Research-grade port from inference/model.py, building a custom golden-output harness that runs the reference engine on tiny inputs. Multi-session effort; new components required.
3. Scaffold-only (config class + registry entry marked not-yet-functional), with CSA/HCA/mHC deferred.

Please confirm which path before I proceed.

[rui-ren]

@copilot you can find model architecture here:

https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/blob/main/inference/model.py

[Copilot]

@copilot you can find model architecture here:

https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/blob/main/inference/model.py

Thanks, I've now read inference/model.py (and kernel.py) from that link. It's a standalone inference engine, not a transformers-compatible reference. Reviewing it reinforces rather than reduces the scope concerns from the PR description:

• The forward path depends on custom CUDA kernels with no ONNX equivalent: sparse_attn (top-k indexed sparse attention), hc_split_sinkhorn (mHC routing), fp4_gemm / fp8_gemm with per-block ue8m0 scales, and fast_hadamard_transform. Porting each requires either a new mobius rewrite-rule + ORT custom op, or a dense fallback that changes numerics.
• New architectural modules with no mobius precedent: Compressor (overlapping gated-pool KV compression with per-position APE and decode-time state buffers), Indexer (FP4-simulated scoring → top-k over compressed KV), grouped output LoRA (o_groups=8, o_lora_rank=1024), sqrtsoftplus gate, clipped SwiGLU, and MTP (n_mtp_layers) heads.
• With no HF transformers class and the reference's kernels being custom, there is no atol=1e-4 parity target — the repo's standard verification path for new models doesn't apply.

A responsible implementation is multi-session and needs new reusable components (CSA, HCA, mHC, Compressor, Indexer), plus a custom golden-tensor harness driven by the reference engine on tiny configs.

Before I start that, could you confirm which of the 3 paths in the PR description you want? Specifically:

1. Wait for transformers support, then port (V3 pattern), or
2. Multi-session research port from inference/model.py with a custom golden-tensor harness, or
3. Scaffold-only (config + registry stub marked not-yet-functional).

I'd like to avoid landing either an unverified graph builder or a partial stub without your go-ahead on scope.

[justinchuby]

@copilot as long as the model is representable in onnx, you can implement the new components in mobius. If there are constructs not representable (maybe sparse attention? Is that correct?), please comment here. A reference is xadupre/mbext#330 which seems to do the job. Feel free to use it as reference or spot any issues if there are.

[Copilot]

@copilot as long as the model is representable in onnx, you can implement the new components in mobius. If there are constructs not representable (maybe sparse attention? Is that corre...

Thanks for the green light. Quick assessment on ONNX representability and a flag on the mbext reference:

ONNX representability — everything maps, with one perf caveat:

• sparse_attn (top-k indexed sparse attention): representable as Gather over K/V along the seq axis using topk_idxs → standard scaled-dot-product attention with the per-row attention sink concatenated as an extra logit. No native sparse-attention ONNX op, so on CPU/CUDA EP this becomes a dense gather+SDPA (correct, but no kernel-level sparsity speedup vs. the reference's custom kernel).
• hc_split_sinkhorn (Sinkhorn-normalized hyper-connection routing): pure tensor ops — alternating row/col normalization in a fixed hc_sinkhorn_iters loop. Unrollable to a static graph, fully representable.
• sqrtsoftplus gate: trivial Softplus → Sqrt, just a new MoEGate scoring branch alongside softmax/sigmoid.
• Clipped SwiGLU (swiglu_limit): Clip(up, -L, L) * Silu(Clip(gate, max=L)). Trivial.
• Compressor (overlapping gated-pool KV compression with APE + RoPE): pure tensor ops; the decode-time kv_state/score_state buffers are KV-cache-style state and map to the existing past_key/past_value pattern.
• Indexer (FP4-simulated scoring → top-k over compressed KV): Linear → einsum → Relu → weighted-sum → TopK. ONNX TopK covers it.
• Output LoRA + groups (o_groups=8, o_lora_rank=1024): grouped einsum, representable.
• FP8/FP4 weights with ue8m0 per-block scales (block_size=128, fp4_block_size=32): not representable as runtime-quantized matmul without ORT custom ops, but dequantizing weights at load time to bf16/fp16 (the standard mobius pattern) is fine and matches the reference's own bf16 simulation paths.
• Hadamard rotation: MatMul against a constant Hadamard matrix.
• MTP heads: representable but expand the ModelPackage shape (extra logit outputs / extra block).
• Hyper-Connections (mHC): the residual stream becomes [b, s, hc_mult, d] everywhere — invasive but pure ops.

Nothing in the V4 forward path is ONNX-non-representable. The only thing the reference does that we cannot match 1:1 is the fused FP4/FP8 gemm + fused sparse-attn kernels, and that's a perf gap, not a correctness gap.

Concern about the mbext reference (xadupre/mbext#330): that PR routes DeepseekV4ForCausalLM into the existing V3 builder unchanged ("Both … share the same MLA+MoE architecture and identical HuggingFace config fields"). Looking at inference/model.py, V4-Flash is not architecturally V3:

• Residual stream shape is [b, s, hc_mult, d] (Hyper-Connections), not [b, s, d].
• Attention is sliding-window CSA + Compressor + Indexer with per-layer compress_ratios, not classic MLA.
• Gate is sqrtsoftplus, not softmax/sigmoid.
• Output projection is grouped LoRA (o_groups, o_lora_rank), not standard wo.
• MTP blocks exist as part of the model.
• Weights have FP8/FP4 + ue8m0 block scales.

A V3 graph loaded with V4 weights will produce wrong logits. So I'd treat that PR as a config-aliasing stub for genai builder, not a parity reference for mobius. Worth flagging back to them too.

Proposed plan for this PR (multi-session, will land in pieces):

1. New components: sqrtsoftplus gate variant, clipped SwiGLU, grouped output LoRA, HyperConnection (pre/post + Sinkhorn), KVCompressor, SparseTopKIndexer, SparseAttention (gather+SDPA fallback), MTP head wrapper.
2. DeepSeekV4Config in _configs.py, registry entry keyed on model_type="deepseek_v4" with a fallback for architectures=["DeepseekV4ForCausalLM"] until transformers registers it.
3. models/deepseek_v4.py with preprocess_weights aligned to the inference/model.py state_dict naming.
4. Verification: since there is no HF transformers reference, I'll add a torch-reference harness that imports inference/model.py (vendored or fetched into a tests fixture) on a tiny config, generates golden tensors, and compares mobius ONNX output at atol=1e-4 (bf16-relaxed for FP8 paths). This lives behind the existing @pytest.mark.integration marker.

Before I start cutting code: do you want all of (1)–(4) in this one PR, or should I split it (e.g. components first in a separate PR, then the model + verification harness)? Given the size of the diff a split would be much easier to review, but happy either way.

[justinchuby]

cc @xadupre for the analysis.

Will hand this back to you @rui-ren , thanks.

[github-actions]

Performance Comparison

Comparing f943b9b → 41565c8

Model	Metric	Baseline	Current	Delta
bert (feature-extraction)	model_size_bytes	359 KB	359 KB	+0.0%	⚪
bert (feature-extraction)	num_nodes	60	60	+0.0%	⚪
falcon	model_size_bytes	364 KB	364 KB	+0.0%	⚪
falcon	num_nodes	66	66	+0.0%	⚪
gemma2	model_size_bytes	428 KB	428 KB	+0.0%	⚪
gemma2	num_nodes	107	107	+0.0%	⚪
gpt2	model_size_bytes	388 KB	388 KB	+0.0%	⚪
gpt2	num_nodes	53	53	+0.0%	⚪
llama	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama	num_nodes	61	61	+0.0%	⚪
llama (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama (static-cache)	num_nodes	58	58	+0.0%	⚪
mamba (ssm-text-generation)	model_size_bytes	296 KB	296 KB	+0.0%	⚪
mamba (ssm-text-generation)	num_nodes	98	98	+0.0%	⚪
phi3	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3	num_nodes	59	59	+0.0%	⚪
phi3 (static-cache)	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3 (static-cache)	num_nodes	56	56	+0.0%	⚪
qwen2	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2	num_nodes	61	61	+0.0%	⚪
qwen2 (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2 (static-cache)	num_nodes	58	58	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	model_size_bytes	506 KB	506 KB	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	num_nodes	275	275	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	model_size_bytes	458 KB	458 KB	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	num_nodes	129	129	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	model_size_bytes	977 KB	977 KB	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	num_nodes	408	408	+0.0%	⚪
t5 (seq2seq)	model_size_bytes	836 KB	836 KB	+0.0%	⚪
t5 (seq2seq)	num_nodes	166	166	+0.0%	⚪
whisper (speech-to-text)	model_size_bytes	1008 KB	1008 KB	+0.0%	⚪
whisper (speech-to-text)	num_nodes	128	128	+0.0%	⚪

No performance regressions.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[justinchuby]

@copilot check if this is done is main or if anything is missing

[Copilot]

@copilot check if this is done is main or if anything is missing

Done in main. DeepSeek-V4-Flash export landed in #405 (5544f2c) with a follow-up cleanup in #429 (58c7472):

• deepseek_v4 registered in _registry.py; config in _configs/_base.py; model in src/mobius/models/deepseek_v4.py; task in src/mobius/tasks/_deepseek_v4.py.
• Exports Q/KV/grouped-output projections, Hyper-Connections, hash / sqrtsoftplus MoE routing, shared experts, and clamped SwiGLU.
• GGUF deepseek4 architecture + tensor/config mapping wired up; ORT-GenAI runtime aliases added.
• Tests: deepseek_v4_flash_test.py, _test_configs.py, model_coverage_test.py, plus GGUF reader/mapping tests.

Deferred follow-ups (explicitly documented in the PR description and DSV4_FLASH_EXPORT.md):

• CSA/HCA learned KV compression, sparse indexer, attention sinks, and per-layer compressed RoPE — the shipped path is a dense causal-attention fallback and is not numerically equivalent to the reference.
• Optional MTP component export.
• Direct packed dynamic int2/int1 and MXFP4 expert import (needs an ORT custom op / MatMulNBits extension).
• Optimized split-GGUF packed-expert repacking for the Unsloth shards.

So the onboarding itself is complete; the remaining items are the sparse-attention / quantization follow-ups called out at merge time. This PR can be closed. Nothing to do here.