Export DeepSeek V4 ratio-4 compressed sparse attention

src/mobius/__main__.py

@@ -46,6 +46,7 @@
     "text-only": "text_only",
     "glm-full-attention": "glm_full_attention",
     "paged-attention": "export_paged_attention",
+    "native-csa": "native_csa",
 }
 
 
@@ -298,6 +299,22 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
     # caller-owned page buffers. It is a distinct cache authority, so it cannot
     # be combined with the static-cache task or an explicit --task.
     export_paged_attention = getattr(args, "export_paged_attention", False)
+    native_csa = getattr(args, "native_csa", False)
+    if native_csa:
+        if static_cache_params is not None:
+            raise SystemExit(
+                "Error: --features native-csa cannot be combined with --features static-cache."
+            )
+        if task is not None:
+            raise SystemExit(
+                "Error: --features native-csa cannot be combined with --task. "
+                "The DeepSeek-V4 task owns the compressed-state ABI."
+            )
+        if not keep_quantized:
+            raise SystemExit(
+                "Error: --features native-csa cannot be combined with "
+                "--dequantize; dense reconstruction is not a CSA capability path."
+            )
     if export_paged_attention:
         if static_cache_params is not None:
             raise SystemExit(
@@ -442,7 +459,24 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
         )
 
         compressed_tensors_config = CompressedTensorsConfig.from_hf_config(parent_config)
-        config = _config_from_hf(hf_config, parent_config=parent_config)
+        if native_csa and model_type != "deepseek_v4":
+            raise SystemExit(
+                "Error: --features native-csa is only supported for model_type "
+                f"'deepseek_v4' (got {model_type!r})."
+            )
+        if native_csa:
+            config = _config_from_hf(
+                hf_config,
+                parent_config=parent_config,
+                allow_block_fp8_dense_fallback=True,
+            )
+        else:
+            config = _config_from_hf(
+                hf_config,
+                parent_config=parent_config,
+            )
+        if native_csa:
+            config = dataclasses.replace(config, native_csa=True)
         if dtype_override is not None:
             config = dataclasses.replace(config, dtype=dtype_override)
         elif compressed_tensors_config is not None and keep_quantized:
@@ -492,7 +526,34 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
         for name, model in pkg.items():
             model.graph.name = f"{config_path}/{name}"
         if load_weights:
-            if compressed_tensors_config is not None:
+            if config.block_quant_scheme is not None and hasattr(
+                model_module, "build_block_quant_streaming_plan"
+            ):
+                from mobius.integrations._weight_loading import (
+                    stream_preprocessed_safetensors_to_model,
+                )
+
+                checkpoint_dir = (
+                    os.path.dirname(config_path)
+                    if os.path.basename(config_path) == "config.json"
+                    else config_path
+                )
+                reports = {}
+                for component_name, model in pkg.items():
+                    reports[component_name] = stream_preprocessed_safetensors_to_model(
+                        model,
+                        checkpoint_dir,
+                        lambda key_index, initializers, name=component_name: (
+                            model_module.build_block_quant_streaming_plan(
+                                name, key_index, initializers
+                            )
+                        ),
+                    )
+                pkg.weight_loading_report = {
+                    "format": "mobius.weight-loading-report.v1",
+                    "components": reports,
+                }
+            elif compressed_tensors_config is not None:
                 # Packed FP4 weights cannot pass through ordinary apply_weights.
                 # The same loader owns both faithful native storage and the
                 # explicit keep_quantized=False dense reconstruction policy.
@@ -545,6 +606,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             prune_prefill_prefix=prune_prefill_prefix,
             glm_full_attention=args.glm_full_attention,
             export_paged_attention=export_paged_attention,
+            native_csa=native_csa,
             keep_quantized=keep_quantized,
             input_sampling_rate=input_sampling_rate,
             bwe_sampling_rate=bwe_sampling_rate,

src/mobius/_configs/_base.py

@@ -707,12 +707,9 @@ class ArchitectureConfig(BaseModelConfig):
     # Deferred block-scaled FP8 / packed-FP4 scheme (DeepSeek-V4-Flash native
     # CSA). Recorded by ``from_transformers`` when ``native_csa`` opts into
     # deferring #602's config-resolution block-quant reject so that graph
-    # construction can progress past the former generic weight-shape mismatch.
-    # The runtime-capability gate
-    # (``mobius.models._deepseek_v4_csa.assert_native_runtime_supports_block_quant``)
-    # then fails closed on the *runnable* full export until nxrt advertises real
-    # block-FP8 / planar-FP4 format strings. ``None`` for every ordinary,
-    # per-tensor-fp8, or non-native path.
+    # construction can select the canonical block-FP8 / planar-FP4 nxrt v1
+    # producer instead of the ordinary dense/INT4 factories. ``None`` for every
+    # ordinary, per-tensor-fp8, or non-native path.
     block_quant_scheme: BlockQuantScheme | None = None
     # HuggingFace model_type and special token IDs — populated by from_transformers()
     # so that genai_config.json can be written without re-fetching the HF config.
@@ -1359,11 +1356,8 @@ def _per_layer_value(attribute: str) -> int | None:
         # ``BlockQuantExportError`` (the INT4/per-tensor path cannot load them).
         # For a native-CSA export we *defer* that reject so graph construction
         # can progress past the former generic weight-shape mismatch; the
-        # runtime-capability gate
-        # (``mobius.models._deepseek_v4_csa.assert_native_runtime_supports_block_quant``,
-        # invoked at weight-load / full-export) then fails closed until nxrt
-        # advertises real block-FP8 / planar-FP4 format strings. Every non-native
-        # path keeps #602's early, loud reject.
+        # canonical planar block-quant producer consumes the recorded scheme.
+        # Every non-native path keeps #602's early, loud reject.
         from mobius.integrations._block_quant import (
             BlockQuantExportError,
             BlockQuantScheme,

src/mobius/_configs/_quantization.py

@@ -117,9 +117,9 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None:
                 "experts (I8-packed E2M1 nibbles + UE8M0 micro-scale). Parse and "
                 "validate these by property with mobius.integrations._block_quant "
                 "(BlockQuantScheme / classify_tensor / QuantizedTensorDescriptor); "
-                "the routed-expert emission gate (plan_routed_expert_bank) reports "
-                "the exact onnx-genai nxrt ABI gap. Native export is blocked until "
-                "the runtime gains a block-FP8 / planar-FP4 BlockFormat."
+                "the routed-expert emission gate (plan_routed_expert_bank) validates "
+                "the canonical onnx-genai nxrt planar ABI. Use the explicit native "
+                "CSA export path; ordinary dense/INT4 export remains fail-closed."
             )
         if method == "none":
             return None

src/mobius/components/__init__.py

@@ -9,6 +9,7 @@
     "BatchNorm2d",
     "BertEmbeddings",
     "BlockQuantizedLinear",
+    "PlanarBlockQuantizedLinear",
     "CausalConv1d",
     "CausalConvNd",
     "CausalDepthwiseConv1d",
@@ -381,6 +382,7 @@
 from mobius.components._quantized_linear import (
     BlockQuantizedLinear,
     NVFP4QuantizedLinear,
+    PlanarBlockQuantizedLinear,
     QuantizedEmbedding,
     QuantizedLinear,
     TiedQuantizedLMHead,

src/mobius/components/_quantized_linear.py

@@ -295,14 +295,16 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
         op.builder.graph.opset_imports[_NXRT_DOMAIN] = 1
         output_dtype = x.dtype
         activation = x if x.dtype == ir.DataType.FLOAT else op.Cast(x, to=ir.DataType.FLOAT)
-        inputs: list[ir.Value | None] = [activation, self.weight]
+        inputs: list[ir.Value | None] = [activation, self.weight, None]
         if self.bias is not None:
             bias = (
                 self.bias
                 if self.bias.dtype == ir.DataType.FLOAT
                 else op.Cast(self.bias, to=ir.DataType.FLOAT)
             )
             inputs.append(bias)
+        else:
+            inputs.append(op.Constant(value=ir.tensor(np.zeros(self._n, dtype=np.float32))))
 
         result = op.BlockQuantizedMatMul(
             *inputs,
@@ -323,6 +325,85 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
         return result
 
 
+class PlanarBlockQuantizedLinear(nn.Module):
+    """Linear layer backed by the canonical nxrt planar block-quantized ABI."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        *,
+        format: str,
+        block_size_out: int,
+        block_size_in: int,
+        model_dtype: ir.DataType,
+        bias: bool = False,
+    ):
+        super().__init__()
+        if format not in {"block_fp8", "fp4_planar"}:
+            raise ValueError(f"format must be 'block_fp8' or 'fp4_planar', got {format!r}")
+        if block_size_out <= 0 or block_size_in <= 0:
+            raise ValueError(
+                f"block geometry must be positive, got [{block_size_out}, {block_size_in}]"
+            )
+        if format == "fp4_planar" and (block_size_out, block_size_in) != (1, 32):
+            raise ValueError(
+                "fp4_planar requires block geometry [1, 32], got "
+                f"[{block_size_out}, {block_size_in}]"
+            )
+        if format == "fp4_planar" and in_features % 2:
+            raise ValueError(f"fp4_planar requires an even K, got {in_features}")
+
+        self._k = in_features
+        self._n = out_features
+        self._format = format
+        self._block_size_out = block_size_out
+        self._block_size_in = block_size_in
+        self._model_dtype = model_dtype
+        packed_k = in_features if format == "block_fp8" else in_features // 2
+        weight_dtype = ir.DataType.FLOAT8E4M3FN if format == "block_fp8" else ir.DataType.INT8
+        self.weight = nn.Parameter([out_features, packed_k], dtype=weight_dtype)
+        self.scale = nn.Parameter(
+            [
+                math.ceil(out_features / block_size_out),
+                math.ceil(in_features / block_size_in),
+            ],
+            dtype=ir.DataType.FLOAT8E8M0,
+        )
+        self.bias = nn.Parameter([out_features], dtype=ir.DataType.FLOAT) if bias else None
+
+    def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
+        op.builder.graph.opset_imports[_NXRT_DOMAIN] = 1
+        output_dtype = x.dtype or self._model_dtype
+        activation = x if x.dtype == ir.DataType.FLOAT else op.Cast(x, to=ir.DataType.FLOAT)
+        result = op.BlockQuantizedMatMul(
+            activation,
+            self.weight,
+            self.scale,
+            (
+                self.bias
+                if self.bias is not None
+                else op.Constant(value=ir.tensor(np.zeros(self._n, dtype=np.float32)))
+            ),
+            K=self._k,
+            N=self._n,
+            format=self._format,
+            block_layout_version=1,
+            block_size_out=self._block_size_out,
+            block_size_in=self._block_size_in,
+            _domain=_NXRT_DOMAIN,
+        )
+        result.dtype = ir.DataType.FLOAT
+        if x.shape is not None:
+            result.shape = ir.Shape([*x.shape[:-1], self._n])
+        if output_dtype not in (None, ir.DataType.FLOAT):
+            result = op.Cast(result, to=output_dtype)
+            result.dtype = output_dtype
+            if x.shape is not None:
+                result.shape = ir.Shape([*x.shape[:-1], self._n])
+        return result
+
+
 class QuantizedEmbedding(nn.Module):
     """Embedding backed by the GatherBlockQuantized custom op.
 

src/mobius/components/_quantized_linear_test.py

@@ -20,6 +20,7 @@
 from mobius.components._quantized_linear import (
     BlockQuantizedLinear,
     NVFP4QuantizedLinear,
+    PlanarBlockQuantizedLinear,
     QuantizedLinear,
 )
 
@@ -360,7 +361,8 @@ def test_emits_native_block_contract(
         node = next(node for node in graph if node.op_type == "BlockQuantizedMatMul")
         assert node.domain == "pkg.nxrt"
         assert graph.opset_imports["pkg.nxrt"] == 1
-        assert len(node.inputs) == 3
+        assert len(node.inputs) == 4
+        assert node.inputs[2] is None
         attrs = {attribute.name: attribute.value for attribute in node.attributes.values()}
         assert attrs == {
             "K": IN_FEATURES,
@@ -374,6 +376,69 @@ def test_rejects_runtime_unsupported_iq_format(self):
             BlockQuantizedLinear(IN_FEATURES, OUT_FEATURES, format="q4_k")
 
 
+class TestPlanarBlockQuantizedLinear:
+    def test_emits_canonical_block_fp8_contract(self):
+        linear = PlanarBlockQuantizedLinear(
+            128,
+            256,
+            format="block_fp8",
+            block_size_out=128,
+            block_size_in=128,
+            model_dtype=ir.DataType.BFLOAT16,
+        )
+        assert linear.weight.shape == [256, 128]
+        assert linear.weight.dtype == ir.DataType.FLOAT8E4M3FN
+        assert linear.scale.shape == [2, 1]
+        assert linear.scale.dtype == ir.DataType.FLOAT8E8M0
+
+        builder, op, graph = create_test_builder()
+        x = create_test_input(
+            builder,
+            "x",
+            [1, 4, 128],
+            dtype=ir.DataType.BFLOAT16,
+        )
+        result = linear(op, x)
+        builder._adapt_outputs([result], "")
+
+        node = next(node for node in graph if node.op_type == "BlockQuantizedMatMul")
+        assert node.domain == "pkg.nxrt"
+        assert len(node.inputs) == 4
+        assert node.inputs[1].name == "weight"
+        assert node.inputs[2].name == "scale"
+        assert node.inputs[3].producer().op_type == "Constant"
+        assert {name: attr.value for name, attr in node.attributes.items()} == {
+            "K": 128,
+            "N": 256,
+            "format": "block_fp8",
+            "block_layout_version": 1,
+            "block_size_out": 128,
+            "block_size_in": 128,
+        }
+
+    def test_fp4_planar_packed_shape_and_geometry(self):
+        linear = PlanarBlockQuantizedLinear(
+            64,
+            32,
+            format="fp4_planar",
+            block_size_out=1,
+            block_size_in=32,
+            model_dtype=ir.DataType.FLOAT,
+        )
+        assert linear.weight.shape == [32, 32]
+        assert linear.weight.dtype == ir.DataType.INT8
+        assert linear.scale.shape == [32, 2]
+        with pytest.raises(ValueError, match=r"\[1, 32\]"):
+            PlanarBlockQuantizedLinear(
+                64,
+                32,
+                format="fp4_planar",
+                block_size_out=128,
+                block_size_in=128,
+                model_dtype=ir.DataType.FLOAT,
+            )
+
+
 class TestMakeQuantizedLinearFactory:
     """Tests for the make_quantized_linear_factory closure."""