feat(deepseek-v4): default-off native CompressedSparseAttention (HCA ratio-128) export [C1, DRAFT] - #593
Conversation
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
B1 — native CSA export extended to the full interleaved ratio-4 + ratio-128 schedule (commit
|
…on primitives (DRAFT) (#602) ## Summary DeepSeek-V4 real-checkpoint blocker **B2**: generic **block-FP8 / packed-FP4 weight-loading + expert-emission primitives**, classified by tensor *properties* (never by model name). Separate artifact from PR #591 (native BQMoE fusion) — touches none of its files, and does not touch `deepseek_v4.py` (Deckard/CSA #593). ### Root cause `QuantizationConfig.from_transformers` treated `quant_method='fp8'` as ordinary per-tensor float8 and returned `None`, so the model built bf16 `Linear` initializers. The real experts are **FP4-packed in int8** (E2M1, block-32, UE8M0 micro-scales), so a packed `[2048,2048]` expert was compared to its logical `[2048,4096]` initializer → confusing `Weight shape mismatch`. Projections are **block-FP8** (E4M3 weight + 2D UE8M0 `[128,128]` block scales). ### What this adds (`mobius.integrations._block_quant`) - **`BlockQuantScheme`** — property parser over `quantization_config` + top-level `expert_dtype`. Per-tensor fp8 (no `weight_block_size`) is *not* owned. - **`QuantizedTensorDescriptor`** — clean, breaking contract: logical **and** packed shape, weight qtype, block geometry, scale name/dtype/shape/layout, exact byte counts, routed/shared role. - **`classify_tensor` / `validate_descriptor`** — distinguish `ORDINARY`, `BLOCK_FP8`, `FP4_PACKED`, `UNSUPPORTED`. Logical-vs-packed shape, scale pairing, and wrong/missing/orphan scales fail closed. **No dequantization, no copy-to-float.** - **`read_raw_tensor_bytes` / `LazyRawTensor`** — byte-exact, header-only, bounded lazy per-shard loading (one tensor resident). - **`stack_expert_bank` / `PackedExpertBank`** — byte-exact expert-major bank packing (reusable lowering primitive; ragged banks are a hard error). - **`runtime_representation_gap` / `plan_routed_expert_bank`** — the **emission gate**. ### ABI verdict (why it typed-rejects, no fake fusion) onnx-genai `nxrt` CPU kernels (`block_quantized_{matmul,moe}.rs`) accept only the **interleaved llama.cpp `block_mxfp4`** layout (QK=32, 17 bytes/block) and the `iq*` GGUF formats: - **block-FP8** projections → no `block_fp8` BlockFormat + no E4M3×2D-UE8M0 dequant path → **typed reject**. - **FP4 experts** are numerically MXFP4 but stored **planar** (separate I8 nibbles + separate E8M0 block-32 scale), while nxrt MXFP4 needs a **single interleaved tensor** → **typed reject** (planar→interleaved transcode is unproven; no planar-FP4 bank ABI). `plan_routed_expert_bank` raises `BlockQuantExportError` naming the exact gap rather than emitting an unrunnable node. **No dense fallback, no env-flag emission.** ### `from_transformers` change (property-guarded) Block-scaled fp8 / fp4-expert checkpoints now surface a precise typed blocker instead of `None → shape-mismatch`. Ordinary per-tensor fp8 still returns `None`; GPTQ/AWQ/GGUF/ModelOpt paths unchanged. ### Preflight (B2 moved) `ArchitectureConfig.from_file(<real checkpoint>)` now raises `BlockQuantExportError` with the exact layout + ABI gap — the failure moved from a confusing shape-mismatch to an actionable typed reject. ### Tests Real `quantization_config` + a measured slice of the checkpoint index metadata + tiny synthetic packed safetensors (real `I8` / `F8_E4M3` / `F8_E8M0` dtypes): classification, logical-vs-packed validation, byte preservation, scale pairing (missing/duplicate/orphan/wrong), shared-vs-routed, bounded lazy load, byte-exact bank stacking, emission typed-reject. An opt-in suite exercises the mounted checkpoint headers directly. `ruff check` + `ruff format` clean. ### Deckard #593 handoff (stable public API) `from mobius.integrations._block_quant import (BlockQuantScheme, QuantizedTensorDescriptor, QuantKind, classify_tensor, validate_descriptor, build_descriptors, pair_weight_scales, read_raw_tensor_bytes, LazyRawTensor, stack_expert_bank, PackedExpertBank, plan_routed_expert_bank, runtime_representation_gap, BlockQuantError, BlockQuantValidationError, BlockQuantExportError)`. No edits to `deepseek_v4.py`. ### Remaining typed blockers (not this PR) 1. nxrt has **no block-FP8 BlockFormat**. 2. FP4 experts are **planar**, nxrt MXFP4 is **interleaved** — needs a proven byte-exact planar→interleaved transcode **or** a planar-FP4 bank ABI extension. 3. End-to-end runnable export additionally needs Deckard's model wiring (#593). No A100 benchmark (no runnable shape-faithful path yet). **Draft — do not merge.** 🤖 Flagged for squad review (needs review): touches the shared `QuantizationConfig.from_transformers` config seam. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
#602 consumption — local stacked-branch composition (fail-closed) ✅Per the directive to consume Mobius #602's generic B2 API What the composition proves (fail-closed, no perf claim):
Real-checkpoint operational evidence (rev Tests/lint on the stack: 80 passed (35 #593 CSA + 40 #602 + 5 new Gate for a non-fail-closed export remains the onnx-genai native block-FP8 / |
…ratio-128) export [C1]
Slice C1 of the DeepSeek-V4-Flash CSA/HCA custom-kernel path: add a
default-off, property-gated Mobius exporter feature that emits the already
-merged frozen `pkg.nxrt::CompressedSparseAttention` v1 op (ratio-128 "Heavily
Compressed Attention" subset) for property-matching layers, replacing the
zero-valued shape-anchor preservation with real op dataflow and threaded
compressed state IO. Export-only; no runner/KV changes.
What
- New `mobius.models._deepseek_v4_csa`: `plan_native_csa` property gate +
`emit_hca_attention` op emission. The gate matches native-op *properties*
(compression ratio, block size, head/latent/rope dims, dtype, carry planes),
not a model-name/shape allowlist.
- `ArchitectureConfig.native_csa: bool = False` opt-in (no HF equivalent;
set via `config_overrides={"native_csa": True}`).
- `DeepSeekV4Attention._forward_native_hca`: emits the 11-input/3-output op
with pre-rotated query/current_kv (dense sliding-window ring over the
existing dense KV IO), real (unrotated) compressor activations, f32 casts
for the f32-cache-format op inputs, and threaded past_*/present_* state.
- `DeepSeekV4Task` threads deterministic `past_compressed_kv.{i}` /
`past_compression_carry.{i}` inputs and `present_*` outputs (explicit shapes
since the custom op has no Python shape inference), with dynamic record axes
and fixed carry planes so present state chains back for >=16 decode steps.
Fail-closed, never silent dense
- Requesting native CSA for ratio-4 CSA, an MTP layer, quantized compressor
weights, or an unknown ratio raises `NativeCsaExportError`; ratio-0 dense
and feature-off legitimately return None (dense).
- Feature off is byte-identical to the existing dense correctness export
(verified) and emits zero `pkg.nxrt` ops / no compressed IO.
Does NOT claim PagedAttention compatibility. onnx-genai runner/KV wiring is a
later slice.
Tests
- Shape-faithful/structural (no weight download, no ORT run): exact op attrs
+ 11/3 IO, threaded compressed-state IO names/shapes/dtypes, decode
chainability, f32 inputs under a FLOAT16 model dtype, no dead anchor for the
HCA layer, byte-identical disabled baseline, defaults-off + opt-in, and typed
rejects (ratio-4/quantized/MTP/unknown-ratio). Existing DeepSeek-V4 tests
unchanged and passing.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ratio-4 CSA [B1/C2] Extend the default-off native pkg.nxrt::CompressedSparseAttention v1 exporter from ratio-128-only HCA to the full official interleaved DeepSeek-V4-Flash schedule (21 ratio-4 CSA + 20 ratio-128 HCA + dense). Ratio-4 layers were previously a typed reject, which made a full native_csa=True export of the official checkpoint architecturally impossible; now both ratios emit the frozen op or the export typed-rejects before serialization (never partial-native + silent dense). Contract re-derived from the frozen kernel (crates/onnx-runtime-ep-cpu/src/kernels/compressed_sparse_attention.rs), the shape-inference handler (crates/onnx-runtime-shape-inference/src/handlers/ custom_ops.rs), and the real safetensors index at rev 60d8d707...: - ratio-4 CSA: 19 inputs / 6 outputs, cache_format='fp8_e4m3_block64' (stored_width=(head_dim-rope)/64*65 + rope*2), a learned FP4 indexer (index cache 'fp4_e2m1_block32', width index_head_dim/32*17), attention carry 8 slots x 2 planes x 2*head_dim, index carry 8 x 2 x 2*index_head_dim, plus the transient int32 selected_indices top-k output. - ratio-128 HCA unchanged: 11 inputs / 3 outputs, f32 cache, carry 128 slots. Clean property-based refactor (no back-compat shims, per user directive): - HcaLayerPlan -> CsaLayerPlan carrying every native-op property for both ratios; _fp8_block64_width/_fp4_width packing helpers fail-closed on invalid divisibility. - plan_native_csa returns a ratio-4 plan; emit_csa_attention branches by ratio and raises on any missing learned-indexer input. - DeepSeekV4IndexerTensors.forward gains a live dataflow mode feeding raw index_query = reshape(wq_b(q_lora)), index_weight, and the index compressor (the op applies index RoPE/Hadamard/FP4 internally) -- the same tensors previously kept as dead anchors, now real dataflow. - _forward_native_csa wires both ratios; the task threads uint8 compressed/ index caches + f32 carries + selected_indices with deterministic layer-suffixed names and dynamic record axes. - MTP recurrence / unknown ratio / unsupported quant still fail closed. Feature stays off by default; disabled export byte-identical. Tests (shape-faithful/structural; no weight download, no ORT run) - New valid-dim _ratio4_config; ratio-4 emits 19/6 with exact attrs; an interleaved [0,4,128] schedule emits both ratios (19/6 and 11/3); ratio-4 threads uint8 index cache (fp4 width 17) + f32 index carry + int32 selected_indices; learned indexer wiring (index_query = Reshape of MatMul, index compressor MatMuls, learned ape/norm); no dead index anchor; prefill + >=16 decode chainability for both ratio-4 caches; typed rejects for invalid fp8/fp4 packing dims and missing index config; ratio-4 plan contract (stored_width=193, index_stored_width=17, carry_slots=8, ...). Existing ratio-128 and DeepSeek-V4 tests migrated atomically and passing (35 total). No performance claim: no ratio-4 layer has executed sparse and no >=16 decode has run. Real-checkpoint weight loading remains blocked on generic FP8/FP4 dequant (B2), reached only after the native gate accepts both ratios. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…able export on nxrt runtime capability Rebased #593 (native pkg.nxrt::CompressedSparseAttention export, ratio-4 CSA + ratio-128 HCA) onto Mobius main now that #578 (6992181) and #602 (d16cd4a, mobius.integrations._block_quant block-FP8 / packed-FP4 loading contract) are merged. Consumes the merged _block_quant API directly; the temporary stacked duplicate-parse branch is retired. #602 rejects the block-scaled-FP8 + packed-FP4 DeepSeek-V4-Flash checkpoint at config resolution (the INT4/per-tensor path cannot load it). That is correct for a normal export, but it also blocks a native-CSA export before any graph exists. Split the behavior into the two stages the directive requires: * Non-native (default) export keeps #602's early, loud config-resolution reject -- unchanged, its tests still pass. * native_csa opts into DEFERRING that reject: ArchitectureConfig.from_transformers catches the typed BlockQuantExportError, records the parsed block_quant_scheme on the config, and lets graph construction PROGRESS past the former generic "Weight shape mismatch" (build_from_module emits the CSA nodes + compressed state IO). The runnable FULL EXPORT then fails closed at a runtime-capability gate (assert_native_runtime_supports_block_quant, enforced at the top of DeepSeekV4 preprocess_weights, before apply_weights) while nxrt cannot execute block-FP8 / planar-FP4 weights. The gap string is sourced from _block_quant.runtime_representation_gap, so the gate tracks the real nxrt format strings and opens automatically when the native runtime slice lands -- no change here. Never a silent dense fallback, never partial-native. No BC shims: ArchitectureConfig gains a block_quant_scheme field; deepseek_v4.py adds no FP8/FP4 weight parsing (delegates to _block_quant). No performance claim -- ratio-4 sparse execution and >=16-decode remain blocked on the unmerged native block-FP8 / planar-FP4 runtime. Tests (deepseek_v4_flash_test.py, 45 pass): non-native fail-closed at config resolution; fp4-experts-alone owned; native_csa defers (scheme recorded, quant None); full export typed-rejects at the runtime-capability gate; graph construction progresses past block-quant (CSA node built); preprocess_weights enforces the gate; property-gated not a blanket V4/native_csa refusal; per-tensor fp8 not over-owned. #602 (_block_quant_test) + config (_base_test/_extractors_test) suites unchanged (69 pass). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
3e10955 to
5028076
Compare
Rebased onto main (#578 + #602 merged); two-stage block-quant compositionRebased this branch onto latest Behavior split into the two stages the directive requires — "graph construction should progress past the former generic shape mismatch, full export must typed-reject at the runtime capability gate":
The gate ( No BC shims (per the no-backward-compat directive): Validation (local; not waiting on CI):
No performance claim. Ratio-4 sparse execution and ≥16-decode stay blocked on the unmerged native block-FP8 / planar-FP4 runtime; this slice only makes graph construction progress and the runnable export fail closed at the honest capability gate. Stopping for independent review. Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com |
Roy — independent review, HEAD
|
There was a problem hiding this comment.
Pull request overview
This PR adds an exporter-only, default-off path for DeepSeek-V4(-Flash) that can emit the frozen pkg.nxrt::CompressedSparseAttention v1 custom op for property-matching layers, and threads the additional compressed state I/O through the task graph so the exported ONNX can carry/return the native compressed state alongside the existing dense KV cache.
Changes:
- Add a new DeepSeek-V4 CSA/HCA planning + emission module (
_deepseek_v4_csa.py) that property-gates native op emission and provides a typed runtime-capability gate for deferred block-quant schemes. - Extend DeepSeek-V4 model + task wiring to thread
past_*→present_*compressed state tensors, and to replace the prior “dead shape anchor” with real compressor/indexer dataflow when native CSA is enabled. - Add extensive structural tests covering node arity/attrs, compressed-state I/O wiring, enable/disable byte-identity, and fail-closed rejects.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/tasks/_deepseek_v4.py | Threads native-CSA compressed-state inputs/outputs through the DeepSeek-V4 task graph. |
| src/mobius/models/deepseek_v4.py | Adds native-CSA planning and op emission path inside attention; replaces anchors with real compressor/indexer dataflow when enabled. |
| src/mobius/models/deepseek_v4_flash_test.py | Adds structural tests for CSA/HCA op emission, state I/O threading, enable/disable behavior, and typed rejects. |
| src/mobius/models/_deepseek_v4_csa.py | New module implementing property gate, op emission helper, and runtime block-quant capability gate. |
| src/mobius/_configs/_base.py | Adds ArchitectureConfig.native_csa toggle and defers block-quant rejection only when native CSA is requested. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # DeepSeek-V4 native compressed-sparse-attention export toggle. Not an | ||
| # upstream HF field -- a Mobius export-time opt-in (default False) that | ||
| # replaces the dense CSA/HCA correctness fallback with the frozen | ||
| # ``pkg.nxrt::CompressedSparseAttention`` v1 op for property-matching | ||
| # ratio-128 (HCA) layers. Off by default so every shipped graph stays | ||
| # byte-identical and ``pkg.nxrt``-free unless explicitly requested via | ||
| # ``config_overrides={"native_csa": True}``. When requested, layers that |
Summary
Slice C1 of the DeepSeek-V4-Flash CSA/HCA custom-kernel path. Adds a default-off, property-gated Mobius exporter feature that emits the already-merged frozen
pkg.nxrt::CompressedSparseAttentionv1 op (ratio-128 "Heavily Compressed Attention" subset) for property-matching layers, replacing the zero-valued shape-anchor preservation with real op dataflow and threaded compressed state IO.This is exporter-only. No onnx-genai runner/KV changes; no PagedAttention-compatibility claim.
Design:
deckard-deepseek-v4-csa-hca-cuda-slice.md(§5 property ABI, §8 tiny-test spec, §10 C1 slice).What changed
mobius.models._deepseek_v4_csa(new):plan_native_csaproperty gate +emit_hca_attentionop emission. The gate matches native-op properties (compression ratio, block size, head/latent/rope dims, dtype, carry planes), not a model-name/shape allowlist.ArchitectureConfig.native_csa: bool = False— export-time opt-in (no HF equivalent;config_overrides={"native_csa": True}).DeepSeekV4Attention._forward_native_hca— emits the 11-input / 3-output op with pre-rotatedquery/current_kv(dense sliding-window ring over the existing dense KV IO), real unrotated compressor activations, f32 casts for the f32-cache-format op inputs, and threadedpast_*/present_*compressed state. The shared inverse RoPE + output projection run unchanged onY.DeepSeekV4Task— threads deterministicpast_compressed_kv.{i}/past_compression_carry.{i}inputs andpresent_*outputs (explicit shapes, since the custom op has no Python shape inference), with a dynamic record axis and fixed carry planes so present state chains back into past inputs for ≥16 decode steps.Fail-closed (never silent dense)
Requesting native CSA for ratio-4 CSA, an MTP layer, quantized compressor weights, or an unknown ratio raises
NativeCsaExportError. Ratio-0 dense layers and feature-off legitimately returnNone(dense) — those are not a suppressed fallback.Off = byte-identical
With the feature off, the exported graph is byte-identical to the existing dense correctness export (verified) and emits zero
pkg.nxrtops / no compressed IO.Tests (shape-faithful / structural — no weight download, no ORT run)
The frozen op has no Python shape inference and no Rust EP in the Python env, so tests assert structure: exact op attrs + 11/3 IO, threaded compressed-state IO names/shapes/dtypes, decode chainability, f32 inputs under a FLOAT16 model dtype (GLM DSA precedent), no dead anchor for the HCA layer, byte-identical disabled baseline, defaults-off + opt-in, and typed rejects (ratio-4 / quantized / MTP / unknown-ratio). Existing DeepSeek-V4 tests unchanged.
Local validation (this env):
deepseek_v4_flash_test.py+deepseek_v4_test.py: 30 passed_configs/tasks/glm_moe_dsarelated tests: 168 passedruff check+ruff format --check: cleanScope / non-goals
Status
DRAFT — stop for independent review. Do not merge. CI is async; local tests above are the validation of record.