From 0cd078dcce75848e37167844e10db1b4fc55f99c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 29 Aug 2026 10:11:04 +0000 Subject: [PATCH] feat: export DeepSeek V4 native compressed sparse attention Add the ratio-4 CSA producer, canonical planar block-quant operators, strict native checkpoint streaming, and focused fail-closed coverage while preserving ratio-128 and non-CSA paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/__main__.py | 66 ++- src/mobius/_configs/_base.py | 16 +- src/mobius/_configs/_quantization.py | 6 +- src/mobius/components/__init__.py | 2 + src/mobius/components/_quantized_linear.py | 83 ++- .../components/_quantized_linear_test.py | 67 ++- src/mobius/integrations/_block_quant.py | 49 +- src/mobius/integrations/_block_quant_test.py | 53 +- src/mobius/integrations/_weight_loading.py | 76 ++- .../integrations/_weight_loading_test.py | 76 ++- .../gguf/_block_quantized_moe_builder_test.py | 131 ++--- src/mobius/integrations/gguf/_builder.py | 11 +- .../integrations/transformers/_builder.py | 50 +- .../transformers/_builder_test.py | 78 +++ .../transformers/_config_resolver.py | 18 +- src/mobius/models/_deepseek_v4_csa.py | 33 +- src/mobius/models/deepseek_v4.py | 517 +++++++++++++++--- src/mobius/models/deepseek_v4_flash_test.py | 189 +++++-- .../_block_quantized_moe_fusion.py | 75 +-- .../_block_quantized_moe_fusion_test.py | 62 +-- src/mobius/tasks/_deepseek_v4.py | 12 +- tests/cli_test.py | 40 ++ 22 files changed, 1285 insertions(+), 425 deletions(-) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 581e48253..ce962235f 100644 --- a/src/mobius/__main__.py +++ b/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, diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 683922515..fa32fdac8 100644 --- a/src/mobius/_configs/_base.py +++ b/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, diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index 419b80a9b..701a72675 100644 --- a/src/mobius/_configs/_quantization.py +++ b/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 diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 65389842f..00b0ff3e5 100644 --- a/src/mobius/components/__init__.py +++ b/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, diff --git a/src/mobius/components/_quantized_linear.py b/src/mobius/components/_quantized_linear.py index 4bcf2444c..1535cd90d 100644 --- a/src/mobius/components/_quantized_linear.py +++ b/src/mobius/components/_quantized_linear.py @@ -295,7 +295,7 @@ 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 @@ -303,6 +303,8 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: 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. diff --git a/src/mobius/components/_quantized_linear_test.py b/src/mobius/components/_quantized_linear_test.py index e1d4feb20..6717db536 100644 --- a/src/mobius/components/_quantized_linear_test.py +++ b/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.""" diff --git a/src/mobius/integrations/_block_quant.py b/src/mobius/integrations/_block_quant.py index db9bbd813..b2bc616eb 100644 --- a/src/mobius/integrations/_block_quant.py +++ b/src/mobius/integrations/_block_quant.py @@ -696,13 +696,13 @@ def stack_expert_bank( # Runtime (nxrt) emission gate — prove representability or typed-reject # --------------------------------------------------------------------------- -#: Block formats the onnx-genai ``nxrt`` CPU kernel's ``BlockFormat::parse`` -#: accepts (``crates/onnx-runtime-ep-cpu/src/kernels/block_quantized_matmul.rs``). -#: MXFP4 here means the *interleaved* llama.cpp ``block_mxfp4`` layout -#: (``QK=32``, 17 bytes/block: 1 E8M0 byte + 16 nibble bytes) packed into a -#: single tensor — NOT a planar (separate nibble + separate scale) layout. +#: Block formats accepted by the canonical onnx-genai ``pkg.nxrt`` v1 ABI. +#: The planar formats use the dedicated auxiliary-scale input; the remaining +#: formats are self-describing interleaved blocks. NXRT_BLOCK_FORMATS: frozenset[str] = frozenset( { + "block_fp8", + "fp4_planar", "mxfp4", "iq4_nl", "iq4_xs", @@ -731,24 +731,8 @@ def runtime_representation_gap( if desc.kind is QuantKind.ORDINARY: return None - if desc.kind is QuantKind.BLOCK_FP8: - return ( - "nxrt has no block-FP8 BlockFormat: its BlockFormat::parse accepts only " - f"{sorted(NXRT_BLOCK_FORMATS)} and there is no E4M3-weight x 2D-UE8M0-block-scale " - "dequant path in block_quantized_{matmul,moe}.rs. Emitting a block-fp8 " - f"projection ({desc.name}, block {desc.block_shape}) would be unrunnable." - ) - - if desc.kind is QuantKind.FP4_PACKED: - return ( - "nxrt MXFP4 requires the interleaved single-tensor llama.cpp block_mxfp4 " - "layout (QK=32, 17 bytes/block = 1 E8M0 byte + 16 nibble bytes); this " - f"checkpoint stores fp4 experts *planar* ({desc.name}: I8 packed " - f"{desc.packed_shape} nibbles + a separate F8_E8M0 {desc.scale_shape} block-32 " - "micro-scale). No planar-fp4 bank ABI exists, and a planar->interleaved " - "transcode is unproven (E2M1 nibble order + E8M0 exponent bias vs llama.cpp " - "must be verified). Emitting with the split tensors would be unrunnable." - ) + if desc.kind in (QuantKind.BLOCK_FP8, QuantKind.FP4_PACKED): + return None return f"{desc.name}: unsupported tensor ({desc.unsupported_reason})" @@ -762,16 +746,8 @@ def plan_routed_expert_bank( """Prove a routed-expert bank is runtime-representable, else typed-reject. Validates that every routed expert shares one packed shape / dtype / scale - layout, then checks :func:`runtime_representation_gap`. Because *runtime* - cannot represent either quantized family today, this raises - :class:`BlockQuantExportError` naming the exact gap rather than emitting an - unrunnable node. It never falls back to a dense per-expert graph and never - dequantizes. - - When (and only when) a future runtime *can* represent the bank, the caller - supplies the byte-exact per-expert payloads via *per_expert_bytes* and a - :class:`PackedExpertBank` is returned — the signature is stable for that - path and for Deckard's #593 integration. + layout, then checks :func:`runtime_representation_gap`. The caller supplies + byte-exact per-expert payloads; no dense fallback or dequantization occurs. """ if not expert_descriptors: raise BlockQuantValidationError("cannot plan an empty routed-expert bank") @@ -802,12 +778,11 @@ def plan_routed_expert_bank( f"Routed-expert bank ({len(expert_descriptors)} x {first.kind.value}) is not " f"representable by the {runtime!r} runtime, so no BlockQuantizedMoE node is " f"emitted (fail closed, no dense fallback, no dequantization). ABI gap: {gap} " - "Resolving this needs a runtime ABI extension (a planar-fp4 / block-fp8 " - "BlockFormat) or a proven byte-exact planar->interleaved transcode primitive." + "Resolving this needs a runtime ABI extension or a proven byte-exact " + "layout conversion primitive." ) - # Representable path (not reachable for today's nxrt ABI). Build the bank - # only from caller-supplied byte-exact payloads — never from placeholders. + # Build the bank only from caller-supplied byte-exact payloads. if per_expert_bytes is None: raise BlockQuantValidationError( "runtime can represent the bank but no per_expert_bytes were supplied to pack it" diff --git a/src/mobius/integrations/_block_quant_test.py b/src/mobius/integrations/_block_quant_test.py index 5a51416b0..9bc7556c6 100644 --- a/src/mobius/integrations/_block_quant_test.py +++ b/src/mobius/integrations/_block_quant_test.py @@ -427,7 +427,7 @@ def test_gap_none_for_ordinary(self): d = classify_tensor("gate.weight", "BF16", (8, 8)) assert runtime_representation_gap(d) is None - def test_gap_block_fp8_names_missing_format(self): + def test_block_fp8_is_representable_by_planar_v1(self): scheme = _real_scheme() d = classify_tensor( "attn.wq_a.weight", @@ -437,21 +437,24 @@ def test_gap_block_fp8_names_missing_format(self): scale_shape=(8, 32), scheme=scheme, ) - gap = runtime_representation_gap(d) - assert gap is not None and "block-FP8" in gap + assert runtime_representation_gap(d) is None - def test_gap_fp4_names_planar_vs_interleaved(self): + def test_fp4_planar_is_representable_by_planar_v1(self): d = _routed_fp4_descs(1)[0] - gap = runtime_representation_gap(d) - assert gap is not None - assert "planar" in gap and "block_mxfp4" in gap + assert runtime_representation_gap(d) is None - def test_plan_routed_bank_typed_rejects_fp4(self): - with pytest.raises(BlockQuantExportError) as ei: - plan_routed_expert_bank(_routed_fp4_descs(3)) - msg = str(ei.value) - assert "not representable" in msg - assert "no dense fallback" in msg + def test_plan_routed_bank_preserves_fp4_bytes(self): + descriptors = _routed_fp4_descs(3) + payloads = [ + bytes([expert + 1]) * descriptor.weight_num_bytes + for expert, descriptor in enumerate(descriptors) + ] + bank = plan_routed_expert_bank( + descriptors, + per_expert_bytes=payloads, + ) + assert bank.num_experts == 3 + assert [bank.expert_bytes(i) for i in range(3)] == payloads def test_plan_rejects_mixed_bank(self): scheme = _real_scheme() @@ -529,7 +532,7 @@ class TestRealCheckpointHeaders: def _index(self) -> dict: return json.loads((REAL_CHECKPOINT / "model.safetensors.index.json").read_text()) - def test_real_layer0_classifies_and_gate_rejects(self): + def test_real_layer0_classifies_and_gate_accepts(self): scheme = _real_scheme() wm = self._index()["weight_map"] # Build a header index for layer 0 by reading only shard headers. @@ -552,8 +555,9 @@ def test_real_layer0_classifies_and_gate_rejects(self): d for d in descs.values() if d.is_routed_expert and d.name.endswith("w1.weight") ] assert routed and all(d.kind is QuantKind.FP4_PACKED for d in routed) - with pytest.raises(BlockQuantExportError): - plan_routed_expert_bank(routed) + payloads = [bytes(descriptor.weight_num_bytes) for descriptor in routed] + bank = plan_routed_expert_bank(routed, per_expert_bytes=payloads) + assert bank.num_experts == len(routed) def test_real_expert_bytes_preserved(self): wm = self._index()["weight_map"] @@ -564,6 +568,23 @@ def test_real_expert_bytes_preserved(self): assert len(raw) == end - start assert shape == (2048, 2048) + def test_real_ratio4_index_query_has_canonical_fp8_scale_pair(self): + wm = self._index()["weight_map"] + weight = "layers.2.attn.indexer.wq_b.weight" + scale = "layers.2.attn.indexer.wq_b.scale" + header_index = {} + for key in (weight, scale): + shard = wm[key] + entry = read_safetensors_header(REAL_CHECKPOINT / shard)[key] + header_index[key] = (entry["dtype"], tuple(entry["shape"])) + descriptor = build_descriptors(header_index, _real_scheme())[weight] + assert descriptor.kind is QuantKind.BLOCK_FP8 + assert descriptor.packed_shape == (8192, 1024) + assert descriptor.scale_name == scale + assert descriptor.scale_dtype == "F8_E8M0" + assert descriptor.scale_shape == (64, 8) + assert descriptor.block_shape == (128, 128) + def _span(path, key): with open(path, "rb") as f: diff --git a/src/mobius/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index 95cd9bdb9..a48d83b87 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -59,6 +59,9 @@ "F64": 8, "F8_E4M3": 1, "F8_E5M2": 1, + "F8_E8M0": 1, + "I8": 1, + "I32": 4, "I64": 8, } _SAFETENSORS_TO_IR_DTYPE = { @@ -66,6 +69,8 @@ "F16": ir.DataType.FLOAT16, "F32": ir.DataType.FLOAT, "F8_E4M3": ir.DataType.FLOAT8E4M3FN, + "F8_E8M0": ir.DataType.FLOAT8E8M0, + "I8": ir.DataType.INT8, } @@ -74,7 +79,7 @@ class StreamingWeightSource: """One checkpoint tensor bound to one dense ONNX initializer.""" source_name: str - mode: Literal["direct", "fp8_scalar", "fp8_block_128"] = "direct" + mode: Literal["direct", "native", "fp8_scalar", "fp8_block_128"] = "direct" scale_name: str | None = None expected_scale: float | None = None @@ -614,6 +619,14 @@ def _materialize_preprocessed_source( "classified as scaled FP8" ) return tensor if tensor.dtype == target_dtype else tensor.to(target_dtype) + if source.mode == "native": + if tensor.dtype != target_dtype: + raise ValueError( + f"Native source '{source.source_name}' has torch dtype " + f"{tensor.dtype}, expected {target_dtype}; native storage is " + "never cast or reconstructed" + ) + return tensor raise AssertionError(f"Unknown streaming weight mode: {source.mode}") @@ -804,18 +817,29 @@ def validate_source( target_bytes = (math.prod(expected_shape) * initializer.dtype.bitwidth + 7) // 8 if isinstance(source, StreamingWeightSource): source_shape, source_bytes, scale_bytes = validate_source(source) + if source.mode == "native": + source_dtype = key_index[source.source_name][2] + if _SAFETENSORS_TO_IR_DTYPE.get(source_dtype) != initializer.dtype: + raise ValueError( + f"Native source '{source.source_name}' has dtype " + f"{source_dtype}, but initializer '{target_name}' expects " + f"{initializer.dtype}" + ) if expected_shape != source_shape: raise ValueError( f"Weight shape mismatch for '{target_name}': model expects " f"{expected_shape}, checkpoint source '{source.source_name}' has " f"{source_shape}" ) - bf16_bytes = math.prod(source_shape) * 2 - cast_bytes = target_bytes if initializer.dtype != ir.DataType.BFLOAT16 else 0 largest_source_tensor_bytes = max(largest_source_tensor_bytes, source_bytes) + if source.mode == "native": + working_set_bytes = source_bytes + else: + bf16_bytes = math.prod(source_shape) * 2 + cast_bytes = target_bytes if initializer.dtype != ir.DataType.BFLOAT16 else 0 + working_set_bytes = source_bytes + bf16_bytes + cast_bytes + scale_bytes largest_reconstruction_working_set_bytes = max( - largest_reconstruction_working_set_bytes, - source_bytes + bf16_bytes + cast_bytes + scale_bytes, + largest_reconstruction_working_set_bytes, working_set_bytes ) else: if len(expected_shape) != 3 or len(source.experts) != expected_shape[0]: @@ -828,31 +852,45 @@ def validate_source( rows = 0 for projection in projections: source_shape, source_bytes, scale_bytes = validate_source(projection) + if projection.mode == "native": + source_dtype = key_index[projection.source_name][2] + if _SAFETENSORS_TO_IR_DTYPE.get(source_dtype) != initializer.dtype: + raise ValueError( + f"Native expert source '{projection.source_name}' has " + f"dtype {source_dtype}, but initializer " + f"'{target_name}' expects {initializer.dtype}" + ) if len(source_shape) != 2 or source_shape[1] != expected_shape[2]: raise ValueError( f"Expert source '{projection.source_name}' has shape " f"{source_shape}; expected [rows, {expected_shape[2]}]" ) rows += source_shape[0] - dense_projection_bytes = ( - math.prod(source_shape) * initializer.dtype.bitwidth + 7 - ) // 8 - bf16_projection_bytes = math.prod(source_shape) * 2 - cast_projection_bytes = ( - dense_projection_bytes - if initializer.dtype != ir.DataType.BFLOAT16 - else 0 - ) largest_source_tensor_bytes = max( largest_source_tensor_bytes, source_bytes, ) + if projection.mode == "native": + projection_working_set = source_bytes + else: + dense_projection_bytes = ( + math.prod(source_shape) * initializer.dtype.bitwidth + 7 + ) // 8 + bf16_projection_bytes = math.prod(source_shape) * 2 + cast_projection_bytes = ( + dense_projection_bytes + if initializer.dtype != ir.DataType.BFLOAT16 + else 0 + ) + projection_working_set = ( + source_bytes + + bf16_projection_bytes + + cast_projection_bytes + + scale_bytes + ) max_transient_bytes = max( max_transient_bytes, - source_bytes - + bf16_projection_bytes - + cast_projection_bytes - + scale_bytes, + projection_working_set, ) if rows != expected_shape[1]: raise ValueError( @@ -1819,7 +1857,7 @@ def stream_safetensors_to_model( def external_data_checksums( - output_dir: str | pathlib.PathLike, + output_dir: str | pathlib.Path, *, pattern: str = "*.onnx.data", chunk_size: int = 1 << 20, diff --git a/src/mobius/integrations/_weight_loading_test.py b/src/mobius/integrations/_weight_loading_test.py index 35b0f2493..0b20b4478 100644 --- a/src/mobius/integrations/_weight_loading_test.py +++ b/src/mobius/integrations/_weight_loading_test.py @@ -33,7 +33,14 @@ from mobius._builder import build_from_module from mobius._model_package import ModelPackage from mobius._testing import make_config -from mobius.integrations._weight_loading import _download_weights, apply_weights +from mobius.integrations._weight_loading import ( + StreamingExpertBankSource, + StreamingWeightPlan, + StreamingWeightSource, + _download_weights, + apply_weights, + stream_preprocessed_safetensors_to_model, +) from mobius.models.base import CausalLMModel from mobius.tasks import CausalLMTask, ModelTask @@ -62,6 +69,73 @@ def get_state_dict(self) -> dict[str, torch.Tensor]: return self.state_dict +def test_native_streaming_preserves_fp8_scale_and_fp4_bank(tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + weight = torch.arange(8, dtype=torch.uint8).reshape(2, 4).view(torch.float8_e4m3fn) + scale = torch.tensor([127], dtype=torch.uint8).reshape(1, 1).view(torch.float8_e8m0fnu) + expert0 = torch.arange(8, dtype=torch.int8).reshape(2, 4) + expert1 = expert0 + 8 + safetensors.torch.save_file( + { + "w.weight": weight, + "w.scale": scale, + "experts.0.weight": expert0, + "experts.1.weight": expert1, + }, + checkpoint / "model.safetensors", + ) + + graph = ir.Graph([], [], nodes=[], name="native_stream") + graph.initializers["weight"] = ir.Value( + name="weight", + type=ir.TensorType(ir.DataType.FLOAT8E4M3FN), + shape=ir.Shape([2, 4]), + ) + graph.initializers["scale"] = ir.Value( + name="scale", + type=ir.TensorType(ir.DataType.FLOAT8E8M0), + shape=ir.Shape([1, 1]), + ) + graph.initializers["bank"] = ir.Value( + name="bank", + type=ir.TensorType(ir.DataType.INT8), + shape=ir.Shape([2, 2, 4]), + ) + model = ir.Model(graph, ir_version=11) + + def planner(_index, _initializers): + return StreamingWeightPlan( + targets={ + "weight": StreamingWeightSource("w.weight", mode="native"), + "scale": StreamingWeightSource("w.scale", mode="native"), + "bank": StreamingExpertBankSource( + experts=( + (StreamingWeightSource("experts.0.weight", mode="native"),), + (StreamingWeightSource("experts.1.weight", mode="native"),), + ) + ), + }, + report={"output_weight_format": "native_planar_block_quant"}, + ) + + report = stream_preprocessed_safetensors_to_model(model, str(checkpoint), planner) + assert report["output_weight_format"] == "native_planar_block_quant" + assert report["largest_reconstruction_working_set_bytes"] == 24 + assert ( + model.graph.initializers["weight"].const_value.numpy().view("uint8").tobytes() + == weight.view(torch.uint8).numpy().tobytes() + ) + assert ( + model.graph.initializers["scale"].const_value.numpy().view("uint8").tobytes() + == scale.view(torch.uint8).numpy().tobytes() + ) + assert torch.equal( + torch.from_numpy(model.graph.initializers["bank"].const_value.numpy()), + torch.stack([expert0, expert1]), + ) + + def _build_model_with_weights() -> tuple[ir.Model, list[str]]: """Build a small test model and return it with its initializer names.""" config = make_config() diff --git a/src/mobius/integrations/gguf/_block_quantized_moe_builder_test.py b/src/mobius/integrations/gguf/_block_quantized_moe_builder_test.py index 2971cfeab..759caab5c 100644 --- a/src/mobius/integrations/gguf/_block_quantized_moe_builder_test.py +++ b/src/mobius/integrations/gguf/_block_quantized_moe_builder_test.py @@ -7,18 +7,15 @@ (:func:`_fuse_native_block_moe`, :func:`_assert_sparse_moe_graph`) on synthetic exported packages -- i.e. the graph state the builder hands to the fusion after ``pkg.apply_weights``. They are distinct from the rewrite-rule unit tests: the -focus here is the *builder's* honesty policy (default fail-closed for -mixed-format per-projection v2 with no environment opt-in, opt-in dense -retention, and the final-graph-state backstop), byte-preserving expert banks, -and deterministic external data. +focus here is the *builder's* honesty policy, canonical per-projection v1 +wiring, byte-preserving expert banks, and deterministic external data. Native IQ/GGUF blocks are codebook-based and cannot be dequantized in NumPy, so correctness is proven structurally (the routed expert storm collapses; the shared expert is untouched) and by byte-preservation (the native blocks are stacked verbatim). GLM-5.2 UD-IQ1 layers mix native formats across their fc1/fc2/fc3 -banks, so they require the ``block_layout_version=2`` per-projection ABI; until -that ships in the runtime the builder must typed-reject them rather than emit an -unrunnable node. +banks; the canonical v1 ABI represents them with per-projection format +attributes. """ from __future__ import annotations @@ -60,26 +57,17 @@ def _moe(pkg: ModelPackage) -> ir.Node: # --------------------------------------------------------------------------- # -# The retired env flag cannot force an unrunnable v2 node through the builder # +# The retired env flag cannot change the canonical v1 node # # --------------------------------------------------------------------------- # @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on"]) -def test_builder_env_flag_cannot_force_v2(monkeypatch, value) -> None: - """The retired ``MOBIUS_ENABLE_BQMOE_PERPROJ_V2`` env flag cannot bypass the blocker. - - A mixed-format (v2-only) layer must typed-reject through the builder no matter - how the old opt-in env var is set: production never emits an unrunnable v2 - node from an environment variable. The reject is atomic -- the dense graph is - left untouched. - """ +def test_builder_env_flag_cannot_change_v1(monkeypatch, value) -> None: monkeypatch.setenv(_V2_ENV, value) pkg, _ = _pkg(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") - equal_before = _count(pkg, "Equal") - - with pytest.raises(SparseMoEExportError, match=r"block_layout_version=2"): - _fuse_native_block_moe(pkg, allow_dense=False) - - assert _count(pkg, "BlockQuantizedMoE") == 0 - assert _count(pkg, "Equal") == equal_before # untouched + assert _fuse_native_block_moe(pkg, allow_dense=False) == 1 + attrs = _moe(pkg).attributes + assert attrs["block_layout_version"].value == 1 + assert attrs["fc1_format"].value == "iq1_s" + assert attrs["fc2_format"].value == "iq4_xs" # --------------------------------------------------------------------------- # @@ -102,58 +90,36 @@ def test_builder_fuses_uniform_native_moe(monkeypatch) -> None: # projections survive. assert _count(pkg, "BlockQuantizedMatMul") == 3 attrs = _moe(pkg).attributes - assert "block_layout_version" not in attrs # v1 (uniform) - assert attrs["format"].value == "iq4_xs" + assert attrs["block_layout_version"].value == 1 + assert attrs["fc1_format"].value == "iq4_xs" + assert attrs["fc2_format"].value == "iq4_xs" -def test_builder_mixed_format_typed_rejects_without_v2_runtime(monkeypatch) -> None: - """GLM-5.2-style mixed formats need v2; with no v2 runtime, fail closed. - - The reject must be atomic: the dense graph is left untouched so the caller - sees the original storm, never a half-rewritten graph. - """ +def test_builder_mixed_format_fuses_with_canonical_v1(monkeypatch) -> None: monkeypatch.delenv(_V2_ENV, raising=False) pkg, _ = _pkg(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") - equal_before = _count(pkg, "Equal") - - with pytest.raises(SparseMoEExportError, match=r"block_layout_version=2"): - _fuse_native_block_moe(pkg, allow_dense=False) - - assert _count(pkg, "BlockQuantizedMoE") == 0 - assert _count(pkg, "Equal") == equal_before # untouched - + assert _fuse_native_block_moe(pkg, allow_dense=False) == 1 + attrs = _moe(pkg).attributes + assert attrs["block_layout_version"].value == 1 + assert attrs["fc1_format"].value == "iq1_s" + assert attrs["fc2_format"].value == "iq4_xs" -def test_builder_env_flag_with_allow_dense_keeps_runnable_fallback(monkeypatch) -> None: - """Env flag set + allow_dense keeps the runnable dense storm, never a v2 node. - The only non-reject escape for a mixed-format layer is the explicit dense - fallback (every expert per token, no throughput claim). Even then the builder - must not emit a ``block_layout_version=2`` node from the retired env flag. - """ +def test_builder_env_flag_with_allow_dense_still_fuses(monkeypatch) -> None: monkeypatch.setenv(_V2_ENV, "1") pkg, _ = _pkg(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") - equal_before = _count(pkg, "Equal") - fused = _fuse_native_block_moe(pkg, allow_dense=True) - - assert fused == 0 # nothing fused - assert _count(pkg, "BlockQuantizedMoE") == 0 # no v2 node emitted from the env flag - assert _count(pkg, "Equal") == equal_before # runnable dense fallback preserved + assert fused == 1 + assert _count(pkg, "BlockQuantizedMoE") == 1 -def test_builder_allow_dense_retains_dense_fallback(monkeypatch) -> None: - """Opting into the dense fallback keeps the runnable per-expert storm.""" +def test_builder_allow_dense_does_not_disable_supported_fusion(monkeypatch) -> None: monkeypatch.delenv(_V2_ENV, raising=False) pkg, _ = _pkg(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") - equal_before = _count(pkg, "Equal") - fused = _fuse_native_block_moe(pkg, allow_dense=True) - # The gate is a no-op under the opt-in (the fusion already warned). _assert_sparse_moe_graph(pkg, source="mixed.gguf", allow_dense=True) - - assert fused == 0 - assert _count(pkg, "BlockQuantizedMoE") == 0 - assert _count(pkg, "Equal") == equal_before # dense fallback preserved + assert fused == 1 + assert _count(pkg, "BlockQuantizedMoE") == 1 def test_builder_bank_bytes_are_byte_preserved(monkeypatch) -> None: @@ -358,50 +324,39 @@ def test_e2e_uniform_native_moe_fuses_through_builder(monkeypatch, tmp_path) -> assert moe == 1 assert matmul == 4 # attn q/k/v/o only; the 12 routed expert matmuls collapsed attrs = _moe(pkg).attributes - assert "block_layout_version" not in attrs # uniform -> v1 - assert attrs["format"].value == "iq1_s" + assert attrs["block_layout_version"].value == 1 + assert attrs["fc1_format"].value == "iq1_s" + assert attrs["fc2_format"].value == "iq1_s" -def test_e2e_mixed_format_typed_rejects_through_builder(monkeypatch, tmp_path) -> None: - """A GLM-5.2-style per-projection mixed GGUF fails closed end-to-end. - - Without the v2 runtime the builder must raise ``SparseMoEExportError`` rather - than emit an unrunnable v2 node -- this is exactly the honest path GLM-5.2 - UD-IQ1 takes until the per-projection runtime ships. - """ +def test_e2e_mixed_format_fuses_through_builder(monkeypatch, tmp_path) -> None: from mobius.integrations.gguf import build_from_gguf monkeypatch.delenv(_V2_ENV, raising=False) path = tmp_path / "mixed-moe.gguf" _write_native_moe_gguf(path, down_fmt="iq4_xs") - with pytest.raises(SparseMoEExportError, match=r"block_layout_version=2"): - build_from_gguf(path, keep_quantized=True) - + pkg = build_from_gguf(path, keep_quantized=True) + moe, matmul = _e2e_ops(pkg) + assert (moe, matmul) == (1, 4) + attrs = _moe(pkg).attributes + assert attrs["block_layout_version"].value == 1 + assert attrs["fc1_format"].value == "iq1_s" + assert attrs["fc2_format"].value == "iq4_xs" -def test_e2e_env_flag_cannot_force_v2_through_builder(monkeypatch, tmp_path) -> None: - """The retired env flag cannot make ``build_from_gguf`` emit a v2 node. - End-to-end, a GLM-5.2-style per-projection mixed GGUF must still raise - ``SparseMoEExportError`` with the old opt-in env var set: there is no - production, CLI, or environment path to an unrunnable v2 node. - """ +def test_e2e_env_flag_cannot_change_v1_through_builder(monkeypatch, tmp_path) -> None: from mobius.integrations.gguf import build_from_gguf monkeypatch.setenv(_V2_ENV, "1") path = tmp_path / "mixed-moe-env.gguf" _write_native_moe_gguf(path, down_fmt="iq4_xs") - with pytest.raises(SparseMoEExportError, match=r"block_layout_version=2"): - build_from_gguf(path, keep_quantized=True) - + pkg = build_from_gguf(path, keep_quantized=True) + assert _moe(pkg).attributes["block_layout_version"].value == 1 -def test_e2e_allow_dense_retains_storm_through_builder(monkeypatch, tmp_path) -> None: - """``allow_dense_moe=True`` keeps the runnable per-expert dense fallback. - It must warn+retain, never count as a fused success: the 12 routed expert - matmuls stay, so no BlockQuantizedMoE is produced. - """ +def test_e2e_allow_dense_still_fuses_supported_layer(monkeypatch, tmp_path) -> None: from mobius.integrations.gguf import build_from_gguf monkeypatch.delenv(_V2_ENV, raising=False) @@ -410,9 +365,7 @@ def test_e2e_allow_dense_retains_storm_through_builder(monkeypatch, tmp_path) -> pkg = build_from_gguf(path, keep_quantized=True, allow_dense_moe=True) moe, matmul = _e2e_ops(pkg) - assert moe == 0 - # 4 attention + 4 experts * 3 projections all remain as dense BlockQuantizedMatMul. - assert matmul == 16 + assert (moe, matmul) == (1, 4) def test_e2e_shared_expert_survives_fusion_through_builder(monkeypatch, tmp_path) -> None: diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index ab099328d..295cf7b44 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -360,13 +360,8 @@ def _fuse_native_block_moe(pkg, *, allow_dense: bool) -> int: :class:`SparseMoEExportError` atomically with the graph untouched, unless ``allow_dense`` downgrades it to a warning + dense keep. - A layer that mixes native formats across its fc1/fc2/fc3 banks (GLM-5.2 - UD-IQ1) can only be expressed with the ``block_layout_version=2`` - per-projection ABI, which no shipped onnx-genai runtime executes yet. The - production builder therefore never enables v2: such layers always - typed-reject here rather than emit an unrunnable node. There is no - environment or CLI opt-in -- v2 stays a schema-construction test path until a - real typed runtime-capability handshake exists. + The canonical v1 ABI carries each projection's block format independently, + so mixed-format layers such as GLM-5.2 UD-IQ1 remain sparse and runnable. """ # Imported lazily: the generic rewrite lives in the rewrite_rules package and # must not be pulled into the GGUF import graph at module load time. @@ -374,8 +369,6 @@ def _fuse_native_block_moe(pkg, *, allow_dense: bool) -> int: fused = 0 for model in pkg.values(): - # No ``_allow_perproj_v2_schema`` argument: the production authority path - # always fails closed for mixed-format v2 (fail-safe default). fused += fuse_block_quantized_moe(model, allow_dense_moe=allow_dense) return fused diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 72099a822..2ed9253f4 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -185,6 +185,7 @@ def build_transformers_model( export_paged_attention: bool = False, input_sampling_rate: int | None = None, bwe_sampling_rate: int | None = None, + native_csa: bool = False, ) -> ModelPackage: """Build a model package from a Transformers checkpoint. @@ -311,11 +312,31 @@ def build_transformers_model( task, allow_parent_architecture_override=not text_only, ) - config = _config_from_hf( - hf_config, - parent_config=parent_config, - module_class=module_class, - ) + if native_csa and model_type != "deepseek_v4": + raise ValueError( + "native_csa=True (--features native-csa) is only supported for " + f"model_type 'deepseek_v4' (got {model_type!r})" + ) + if native_csa: + config = _config_from_hf( + hf_config, + parent_config=parent_config, + module_class=module_class, + allow_block_fp8_dense_fallback=True, + ) + else: + config = _config_from_hf( + hf_config, + parent_config=parent_config, + module_class=module_class, + ) + if native_csa: + if not keep_quantized: + raise ValueError( + "native_csa=True requires keep_quantized=True; dense " + "reconstruction is not a CSA capability path" + ) + config = dataclasses.replace(config, native_csa=True) if ( compressed_tensors_config is not None and fp8_kv_cache @@ -403,6 +424,25 @@ def build_transformers_model( if load_weights: if config.block_quant_scheme is not None and hasattr( + model_module, "build_block_quant_streaming_plan" + ): + reports = {} + for component_name, model in package.items(): + reports[component_name] = stream_preprocessed_safetensors_to_model( + model, + model_id, + lambda key_index, initializers, name=component_name: ( + model_module.build_block_quant_streaming_plan( + name, key_index, initializers + ) + ), + revision=revision, + ) + package.weight_loading_report = { + "format": "mobius.weight-loading-report.v1", + "components": reports, + } + elif config.block_quant_scheme is not None and hasattr( model_module, "build_fp8_streaming_plan" ): if len(package) != 1: diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 02d5672c3..f0f275e31 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -31,6 +31,84 @@ def build_fp8_streaming_plan(self, *_args): raise AssertionError("mock streaming function should own planning") +class _DummyNativeCsaModule(_DummyModule): + def build_block_quant_streaming_plan(self, *_args): + return "native-plan" + + +def test_native_csa_activates_planar_streaming(monkeypatch) -> None: + hf_config = SimpleNamespace(model_type="deepseek_v4") + scheme = BlockQuantScheme.from_quantization_config( + {"quant_method": "fp8", "weight_block_size": [128, 128]}, + expert_dtype="fp4", + ) + assert scheme is not None + model = ir.Model(ir.Graph([], [], nodes=[], name="model"), ir_version=11) + package = ModelPackage({"model": model}) + built_configs = [] + loader_calls = [] + + monkeypatch.setattr( + transformers_builder, + "_load_transformers_config", + lambda *_args, **_kwargs: (hf_config, False), + ) + monkeypatch.setattr( + transformers_builder, + "_resolve_module_class", + lambda *_args, **_kwargs: ( + _DummyNativeCsaModule, + "deepseek-v4", + "deepseek_v4", + ), + ) + + def resolve_config( + primary, + *, + parent_config, + module_class, + allow_block_fp8_dense_fallback, + ): + assert primary is hf_config + assert parent_config is hf_config + assert module_class is _DummyNativeCsaModule + assert allow_block_fp8_dense_fallback is True + return make_config( + model_type="deepseek_v4", + block_quant_scheme=scheme, + ) + + monkeypatch.setattr(_config_resolver, "_config_from_hf", resolve_config) + + def build_module(_module, config, *_args, **_kwargs): + built_configs.append(config) + package.config = config + return package + + monkeypatch.setattr(transformers_builder, "build_from_module", build_module) + + def stream(_model, _model_id, planner, **kwargs): + loader_calls.append((planner({}, {}), kwargs)) + return {"output_weight_format": "native_planar_block_quant"} + + monkeypatch.setattr( + transformers_builder, + "stream_preprocessed_safetensors_to_model", + stream, + ) + + result = transformers_builder.build_transformers_model( + "deepseek-ai/DeepSeek-V4-Flash", + native_csa=True, + ) + assert result is package + assert built_configs[0].native_csa is True + assert loader_calls == [ + ("native-plan", {"revision": None}), + ] + + @pytest.mark.parametrize( ("keep_quantized", "expected_loader"), [(True, "qdq"), (False, "dense")], diff --git a/src/mobius/integrations/transformers/_config_resolver.py b/src/mobius/integrations/transformers/_config_resolver.py index c16ea0351..3332a3330 100644 --- a/src/mobius/integrations/transformers/_config_resolver.py +++ b/src/mobius/integrations/transformers/_config_resolver.py @@ -16,6 +16,7 @@ "_try_load_config_json", ] +import inspect import logging from mobius._configs import ( @@ -28,7 +29,13 @@ logger = logging.getLogger(__name__) -def _config_from_hf(hf_config, parent_config=None, module_class=None) -> BaseModelConfig: +def _config_from_hf( + hf_config, + parent_config=None, + module_class=None, + *, + allow_block_fp8_dense_fallback: bool = False, +) -> BaseModelConfig: """Select the right config class for a HuggingFace config object. Resolution order: @@ -66,6 +73,15 @@ def _config_from_hf(hf_config, parent_config=None, module_class=None) -> BaseMod # Call from_transformers — pass parent_config for ArchitectureConfig tree if issubclass(config_cls, ArchitectureConfig): + if ( + "allow_block_fp8_dense_fallback" + in inspect.signature(config_cls.from_transformers).parameters + ): + return config_cls.from_transformers( + hf_config, + parent_config=parent_config, + allow_block_fp8_dense_fallback=allow_block_fp8_dense_fallback, + ) return config_cls.from_transformers(hf_config, parent_config=parent_config) return config_cls.from_transformers(hf_config) diff --git a/src/mobius/models/_deepseek_v4_csa.py b/src/mobius/models/_deepseek_v4_csa.py index 8c4fa3208..47d63d49e 100644 --- a/src/mobius/models/_deepseek_v4_csa.py +++ b/src/mobius/models/_deepseek_v4_csa.py @@ -284,6 +284,18 @@ def present_index_carry_name(self) -> str: def selected_indices_name(self) -> str: return f"selected_indices.{self.layer_id}" + @property + def past_records_axis_name(self) -> str: + return f"past_compressed_records.{self.layer_id}" + + @property + def present_records_axis_name(self) -> str: + return f"present_compressed_records.{self.layer_id}" + + @property + def selected_records_axis_name(self) -> str: + return f"selected_records.{self.layer_id}" + def _layer_compress_ratio(config: ArchitectureConfig, layer_id: int) -> int: ratios = config.compress_ratios or [] @@ -428,13 +440,9 @@ def plan_native_csa( # packed-FP4 DeepSeek-V4-Flash checkpoint is allowed to *progress* (the # ``ArchitectureConfig`` records the deferred ``block_quant_scheme`` rather than # rejecting at config resolution), so the CSA nodes and their compressed state -# IO are built and inspectable. The *runnable* full export, however, must fail -# closed until the native ``nxrt`` runtime can actually execute the block-quant -# weights. That capability is owned by ``mobius.integrations._block_quant`` -# (#602): ``runtime_representation_gap`` returns a precise ABI-gap string while -# ``nxrt`` lacks a block-FP8 / planar-FP4 ``BlockFormat`` and ``None`` once the -# real format strings land -- at which point this gate opens automatically with -# no change here. +# IO can compose with the canonical planar block-quant producer. Runtime +# representability remains a property gate owned by +# ``mobius.integrations._block_quant``. def _representative_block_fp8_descriptor( @@ -520,11 +528,9 @@ def assert_native_runtime_supports_block_quant( ) -> None: """Full-export runtime-capability gate for a deferred block-quant scheme. - No-op unless ``config.block_quant_scheme`` is set (only ``from_transformers`` - records it, and only when ``native_csa`` deferred #602's config-resolution - reject). Raises the typed :class:`BlockQuantExportError` -- fail-closed, - never a silent dense fallback -- while *runtime* cannot execute the - checkpoint's block-FP8 / planar-FP4 weights. + No-op when the canonical runtime represents the recorded scheme. Raises the + typed :class:`BlockQuantExportError` for an unknown/unrepresentable runtime; + never selects a silent dense fallback. """ scheme = getattr(config, "block_quant_scheme", None) gap = native_runtime_block_quant_gap(scheme, runtime=runtime) @@ -534,8 +540,7 @@ def assert_native_runtime_supports_block_quant( "native CSA full export requires a runtime that can execute the " f"checkpoint's block-quant weights, but {runtime!r} cannot yet. Graph " "construction progressed (CSA nodes + compressed state IO are built), " - "but the runnable export is blocked until the native block-FP8 / " - "planar-FP4 format strings land (mobius.integrations._block_quant." + "and no dense fallback is permitted (mobius.integrations._block_quant." "runtime_representation_gap / plan_routed_expert_bank). ABI gap:\n" f"{gap}" ) diff --git a/src/mobius/models/deepseek_v4.py b/src/mobius/models/deepseek_v4.py index eb4f272d4..91d4fe85b 100644 --- a/src/mobius/models/deepseek_v4.py +++ b/src/mobius/models/deepseek_v4.py @@ -1,16 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""DeepSeek-V4 export with a sink-aware dense CSA fallback and MTP sidecar. +"""DeepSeek-V4 export with native compressed sparse attention and MTP sidecar. The released V4 architecture replaces V3 MLA with compressed sparse attention and adds Hyper-Connections. This module implements the V4 projections, -Hyper-Connections, hash/sqrt-softplus MoE routing, and a dense causal-attention -fallback with the learned attention sinks. The official per-layer compression -schedule is represented by exported compressor/indexer tensors, and the -checkpoint's MTP block is exported as a standalone sidecar. Executing learned -KV compression and sparse selection still requires runtime cache/sparse ops; -until those land, the target and MTP graphs attend densely for correctness. +Hyper-Connections, hash/sqrt-softplus MoE routing, and optional native +``CompressedSparseAttention``. The official per-layer compression schedule is +represented by live compressor/indexer projections, while non-native exports +retain the sink-aware dense fallback. The checkpoint's MTP block is exported as +a standalone sidecar. """ from __future__ import annotations @@ -34,6 +33,7 @@ Embedding, Linear, MoELayer, + PlanarBlockQuantizedLinear, QuantizedEmbedding, RMSNorm, create_attention_bias, @@ -109,6 +109,27 @@ def _gqa_kv_lengths(op: OpBuilder, attention_mask: ir.Value) -> tuple[ir.Value, def _projection_class(config: ArchitectureConfig): + scheme = config.block_quant_scheme + if scheme is not None: + if not scheme.is_block_scaled_fp8 or len(scheme.weight_block_size) != 2: + raise NativeCsaExportError( + "native DeepSeek-V4 block-quant projections require block-scaled " + f"FP8 with a 2-D block geometry; got {scheme}" + ) + block_out, block_in = scheme.weight_block_size + + def planar_projection(in_features: int, out_features: int, bias: bool = False): + return PlanarBlockQuantizedLinear( + in_features, + out_features, + format="block_fp8", + block_size_out=block_out, + block_size_in=block_in, + model_dtype=config.dtype, + bias=bias, + ) + + return planar_projection quantization = config.quantization if quantization is None or quantization.quant_method == "none": return Linear @@ -254,6 +275,114 @@ def _validate_hash_routing_tables(state_dict: dict[str, torch.Tensor]) -> None: ) +def _pack_planar_expert_weights( + state_dict: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Stack per-expert FP4 weight/scale planes into canonical nxrt banks.""" + grouped: dict[str, dict[int, dict[str, torch.Tensor]]] = {} + passthrough: dict[str, torch.Tensor] = {} + marker = ".mlp.experts." + for key, value in state_dict.items(): + if marker not in key: + passthrough[key] = value + continue + prefix, suffix = key.split(marker, 1) + parts = suffix.split(".") + if ( + len(parts) != 3 + or parts[1] + not in { + "gate_proj", + "up_proj", + "down_proj", + } + or parts[2] not in {"weight", "scale"} + ): + raise NativeCsaExportError(f"malformed planar routed-expert tensor name {key!r}") + try: + expert = int(parts[0]) + except ValueError as exc: + raise NativeCsaExportError(f"invalid routed-expert index in {key!r}") from exc + grouped.setdefault(prefix, {}).setdefault(expert, {})[f"{parts[1]}.{parts[2]}"] = value + + required = { + "gate_proj.weight", + "gate_proj.scale", + "up_proj.weight", + "up_proj.scale", + "down_proj.weight", + "down_proj.scale", + } + for prefix, experts in grouped.items(): + ids = sorted(experts) + if ids != list(range(len(ids))): + raise NativeCsaExportError( + f"{prefix} planar expert ids must be contiguous 0..E-1, got {ids}" + ) + for expert, tensors in experts.items(): + missing = sorted(required - set(tensors)) + extra = sorted(set(tensors) - required) + if missing or extra: + raise NativeCsaExportError( + f"{prefix} expert {expert} malformed planar tensor set: " + f"missing={missing}, extra={extra}" + ) + passthrough[f"{prefix}.mlp.moe.fc1_experts_weights"] = torch.stack( + [experts[i]["gate_proj.weight"] for i in ids] + ) + passthrough[f"{prefix}.mlp.moe.fc2_experts_weights"] = torch.stack( + [experts[i]["down_proj.weight"] for i in ids] + ) + passthrough[f"{prefix}.mlp.moe.fc3_experts_weights"] = torch.stack( + [experts[i]["up_proj.weight"] for i in ids] + ) + passthrough[f"{prefix}.mlp.moe.fc1_experts_aux_scale"] = torch.stack( + [experts[i]["gate_proj.scale"] for i in ids] + ) + passthrough[f"{prefix}.mlp.moe.fc2_experts_aux_scale"] = torch.stack( + [experts[i]["down_proj.scale"] for i in ids] + ) + passthrough[f"{prefix}.mlp.moe.fc3_experts_aux_scale"] = torch.stack( + [experts[i]["up_proj.scale"] for i in ids] + ) + return passthrough + + +def _map_checkpoint_weight_name(key: str) -> str: + """Map one official DeepSeek-V4 tensor name to the exported module name.""" + new_key = key + if new_key == "embed.weight": + return "model.embed_tokens.weight" + if new_key == "head.weight": + return "lm_head.weight" + if new_key == "norm.weight": + return "model.norm.weight" + if new_key.startswith("hc_head_"): + return f"model.{new_key}.weight" if new_key == "hc_head_fn" else f"model.{new_key}" + if new_key.startswith("layers."): + new_key = f"model.{new_key}" + if new_key.startswith(("model.layers.", "mtp.")): + new_key = new_key.replace(".attn.wq_a.", ".self_attn.q_a_proj.") + new_key = new_key.replace(".attn.q_norm.", ".self_attn.q_a_layernorm.") + new_key = new_key.replace(".attn.wq_b.", ".self_attn.q_b_proj.") + new_key = new_key.replace(".attn.wkv.", ".self_attn.kv_proj.") + new_key = new_key.replace(".attn.kv_norm.", ".self_attn.kv_layernorm.") + new_key = new_key.replace(".attn.wo_a.", ".self_attn.o_a_proj.") + new_key = new_key.replace(".attn.wo_b.", ".self_attn.o_b_proj.") + new_key = new_key.replace(".attn.", ".self_attn.") + new_key = new_key.replace(".attn_norm.", ".input_layernorm.") + new_key = new_key.replace(".ffn_norm.", ".post_attention_layernorm.") + new_key = new_key.replace(".ffn.gate.", ".mlp.moe.gate.") + new_key = new_key.replace(".ffn.experts.", ".mlp.experts.") + new_key = new_key.replace(".ffn.shared_experts.", ".mlp.shared_experts.") + new_key = new_key.replace(".w1.", ".gate_proj.") + new_key = new_key.replace(".w2.", ".down_proj.") + new_key = new_key.replace(".w3.", ".up_proj.") + if ".hc_" in new_key and new_key.endswith("_fn"): + new_key = f"{new_key}.weight" + return new_key + + class DeepSeekV4Gate(nn.Module): """V4 sqrt-softplus router with hash routing for the first layers.""" @@ -369,6 +498,126 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return self.down_proj(op, op.Mul(op.Swish(gate), up)) +class _DeepSeekV4PlanarMoE(nn.Module): + """Sparse routed experts using the canonical 12-input nxrt v1 ABI.""" + + def __init__(self, config: ArchitectureConfig, gate: DeepSeekV4Gate): + super().__init__() + assert config.num_local_experts is not None + assert config.num_experts_per_tok is not None + assert config.moe_intermediate_size is not None + scheme = config.block_quant_scheme + if scheme is None or not scheme.has_packed_fp4_experts: + raise NativeCsaExportError( + "planar DeepSeek-V4 MoE requires packed FP4 routed experts" + ) + self.gate = gate + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + self.hidden_size = config.hidden_size + self.intermediate_size = config.moe_intermediate_size + self.model_dtype = config.dtype + self.swiglu_limit = config.swiglu_limit + + packed_hidden = config.hidden_size // 2 + packed_intermediate = config.moe_intermediate_size // 2 + scale_hidden = config.hidden_size // 32 + scale_intermediate = config.moe_intermediate_size // 32 + if config.hidden_size % 32 or config.moe_intermediate_size % 32: + raise NativeCsaExportError( + "fp4_planar routed experts require hidden/intermediate widths " + "divisible by 32; got " + f"{config.hidden_size}/{config.moe_intermediate_size}" + ) + self.fc1_experts_weights = nn.Parameter( + [self.num_experts, self.intermediate_size, packed_hidden], + dtype=ir.DataType.INT8, + ) + self.fc2_experts_weights = nn.Parameter( + [self.num_experts, self.hidden_size, packed_intermediate], + dtype=ir.DataType.INT8, + ) + self.fc3_experts_weights = nn.Parameter( + [self.num_experts, self.intermediate_size, packed_hidden], + dtype=ir.DataType.INT8, + ) + self.fc1_experts_aux_scale = nn.Parameter( + [self.num_experts, self.intermediate_size, scale_hidden], + dtype=ir.DataType.FLOAT8E8M0, + ) + self.fc2_experts_aux_scale = nn.Parameter( + [self.num_experts, self.hidden_size, scale_intermediate], + dtype=ir.DataType.FLOAT8E8M0, + ) + self.fc3_experts_aux_scale = nn.Parameter( + [self.num_experts, self.intermediate_size, scale_hidden], + dtype=ir.DataType.FLOAT8E8M0, + ) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + input_ids: ir.Value, + ) -> ir.Value: + routing_weights, selected_experts = self.gate(op, hidden_states, input_ids) + routing_weights = op.CastLike(routing_weights, hidden_states) + router_logits, router_weights = _scatter_selected_to_full( + op, routing_weights, selected_experts, self.num_experts + ) + router_logits = op.Cast( + op.Reshape(router_logits, [-1, self.num_experts]), + to=ir.DataType.FLOAT, + ) + router_weights = op.Cast( + op.Reshape(router_weights, [-1, self.num_experts]), + to=ir.DataType.FLOAT, + ) + activation = ( + hidden_states + if hidden_states.dtype == ir.DataType.FLOAT + else op.Cast(hidden_states, to=ir.DataType.FLOAT) + ) + op.builder.graph.opset_imports["pkg.nxrt"] = 1 + result = op.BlockQuantizedMoE( + activation, + router_logits, + self.fc1_experts_weights, + None, + self.fc2_experts_weights, + None, + self.fc3_experts_weights, + None, + router_weights, + self.fc1_experts_aux_scale, + self.fc2_experts_aux_scale, + self.fc3_experts_aux_scale, + k=self.top_k, + activation_type="swiglu", + normalize_routing_weights=0, + swiglu_fusion=0, + swiglu_limit=(self.swiglu_limit if self.swiglu_limit > 0 else float("inf")), + fc1_format="fp4_planar", + fc2_format="fp4_planar", + fc3_format="fp4_planar", + fc1_block_size_out=1, + fc1_block_size_in=32, + fc2_block_size_out=1, + fc2_block_size_in=32, + fc3_block_size_out=1, + fc3_block_size_in=32, + block_layout_version=1, + _domain="pkg.nxrt", + ) + result.dtype = ir.DataType.FLOAT + result.shape = hidden_states.shape + if self.model_dtype != ir.DataType.FLOAT: + result = op.Cast(result, to=self.model_dtype) + result.dtype = self.model_dtype + result.shape = hidden_states.shape + return result + + class DeepSeekV4MoE(nn.Module): def __init__(self, config: ArchitectureConfig, layer_id: int): super().__init__() @@ -381,16 +630,19 @@ def __init__(self, config: ArchitectureConfig, layer_id: int): # swiglu_limit=0.0 as "clip to zero" -- math.inf is required to # disable clipping at the op level. swiglu_limit = config.swiglu_limit if config.swiglu_limit > 0 else math.inf - self.moe = MoELayer( - config, - gate=gate, - expert_factory=lambda expert_config, _linear_class: _DeepSeekV4Expert( - expert_config, expert_config.intermediate_size - ), - activation_alpha=1.0, - activation_beta=0.0, - swiglu_limit=swiglu_limit, - ) + if config.block_quant_scheme is not None: + self.moe = _DeepSeekV4PlanarMoE(config, gate) + else: + self.moe = MoELayer( + config, + gate=gate, + expert_factory=lambda expert_config, _linear_class: _DeepSeekV4Expert( + expert_config, expert_config.intermediate_size + ), + activation_alpha=1.0, + activation_beta=0.0, + swiglu_limit=swiglu_limit, + ) shared_size = config.moe_intermediate_size * (config.n_shared_experts or 1) self.shared_experts = _DeepSeekV4Expert(config, shared_size) @@ -458,11 +710,18 @@ def __init__(self, config: ArchitectureConfig): assert config.index_head_dim is not None self.index_n_heads = config.index_n_heads self.index_head_dim = config.index_head_dim - self.wq_b = DeepSeekV4DeferredProjection( - config, - config.q_lora_rank, - config.index_n_heads * config.index_head_dim, - ) + if config.block_quant_scheme is not None: + self.wq_b = _projection_class(config)( + config.q_lora_rank, + config.index_n_heads * config.index_head_dim, + bias=False, + ) + else: + self.wq_b = DeepSeekV4DeferredProjection( + config, + config.q_lora_rank, + config.index_n_heads * config.index_head_dim, + ) self.weights_proj = DeepSeekV4DeferredProjection( config, config.hidden_size, config.index_n_heads ) @@ -748,8 +1007,10 @@ def _forward_native_csa( # parameters. Routed through the compressor's ``__call__`` so every # compressor parameter is realized (named + registered as an # initializer) rather than left dangling. + assert self.compressor is not None compressor_kv, compressor_gate, ape, norm_weight = self.compressor(op, hidden_states) + present_compressed: tuple common = dict( query=_cast_to_f32(op, query_4d, self._dtype), current_kv=_cast_to_f32(op, current_kv, self._dtype), @@ -1361,7 +1622,7 @@ class DeepSeekV4CausalLMModel(CausalLMModel): def __init__(self, config: ArchitectureConfig): nn.Module.__init__(self) - if any(config.compress_ratios or ()): + if any(config.compress_ratios or ()) and not config.native_csa: logger.warning( "DeepSeek-V4 sparse cache execution requires runtime support; " "exporting sink-aware dense attention with CSA/HCA tensors retained." @@ -1400,78 +1661,185 @@ def forward( ) return self.lm_head(op, hidden_states), presents + def build_block_quant_streaming_plan( + self, + component_name: str, + key_index: dict[str, tuple[str, list[int], str]], + initializers: dict[str, ir.Value], + ): + """Build a complete header-validated native planar weight plan.""" + from mobius.integrations._block_quant import ( + QuantKind, + build_descriptors, + ) + from mobius.integrations._weight_loading import ( + StreamingExpertBankSource, + StreamingWeightPlan, + StreamingWeightSource, + ) + + assert_native_runtime_supports_block_quant(self.config) + if self.config.block_quant_scheme is None: + raise NativeCsaExportError( + "native block-quant streaming requested without a block-quant scheme" + ) + if component_name not in {"model", "mtp"}: + raise NativeCsaExportError( + f"unknown DeepSeek-V4 package component {component_name!r}" + ) + descriptors = build_descriptors( + {name: (dtype, tuple(shape)) for name, (_path, shape, dtype) in key_index.items()}, + self.config.block_quant_scheme, + ) + malformed = [ + descriptor + for descriptor in descriptors.values() + if descriptor.kind is QuantKind.UNSUPPORTED + ] + if malformed: + sample = malformed[0] + raise NativeCsaExportError( + f"malformed block-quant checkpoint tensor {sample.name!r}: " + f"{sample.unsupported_reason}" + ) + + mtp_component = component_name == "mtp" + targets: dict[str, StreamingWeightSource | StreamingExpertBankSource] = {} + ignored: dict[str, str] = {} + + def owned(source_name: str) -> bool: + return source_name.startswith("mtp.") == mtp_component + + def native(source_name: str) -> StreamingWeightSource: + if source_name not in key_index: + raise NativeCsaExportError( + f"missing required planar checkpoint tensor {source_name!r}" + ) + return StreamingWeightSource(source_name, mode="native") + + expert_prefixes: set[str] = set() + for source_name in key_index: + if not owned(source_name): + ignored[source_name] = ( + "target decoder tensor" if mtp_component else "MTP sidecar tensor" + ) + continue + if ".ffn.experts." in source_name: + expert_prefixes.add(source_name.split(".ffn.experts.", 1)[0]) + continue + + target_name = _map_checkpoint_weight_name(source_name) + if target_name not in initializers: + # Hash-routed layers consume the gate matrix to compute routing + # weights, but their checkpoint bias is intentionally unused. + if ( + source_name.endswith(".ffn.gate.bias") + and source_name.startswith("layers.") + and int(source_name.split(".", 2)[1]) < self.config.num_hash_layers + ): + ignored[source_name] = "unused hash-router bias" + continue + raise NativeCsaExportError( + f"checkpoint tensor {source_name!r} maps to missing " + f"{component_name} initializer {target_name!r}" + ) + source_dtype = key_index[source_name][2] + if source_dtype in {"F8_E4M3", "F8_E8M0", "I8"}: + targets[target_name] = StreamingWeightSource(source_name, mode="native") + else: + targets[target_name] = StreamingWeightSource(source_name, mode="direct") + + experts = self.config.num_local_experts + if experts is None or experts <= 0: + raise NativeCsaExportError( + f"invalid routed expert count {experts} for native planar export" + ) + for source_prefix in sorted(expert_prefixes): + target_prefix = _map_checkpoint_weight_name( + f"{source_prefix}.ffn.experts.0.w1.weight" + ).split(".mlp.experts.", 1)[0] + bank_sources = { + f"{target_prefix}.mlp.moe.fc1_experts_weights": (("w1", "weight"),), + f"{target_prefix}.mlp.moe.fc2_experts_weights": (("w2", "weight"),), + f"{target_prefix}.mlp.moe.fc3_experts_weights": (("w3", "weight"),), + f"{target_prefix}.mlp.moe.fc1_experts_aux_scale": (("w1", "scale"),), + f"{target_prefix}.mlp.moe.fc2_experts_aux_scale": (("w2", "scale"),), + f"{target_prefix}.mlp.moe.fc3_experts_aux_scale": (("w3", "scale"),), + } + for target_name, projections in bank_sources.items(): + if target_name not in initializers: + raise NativeCsaExportError( + f"planar expert bank maps to missing {component_name} " + f"initializer {target_name!r}" + ) + targets[target_name] = StreamingExpertBankSource( + experts=tuple( + tuple( + native(f"{source_prefix}.ffn.experts.{expert}.{projection}.{kind}") + for projection, kind in projections + ) + for expert in range(experts) + ) + ) + + missing_targets = sorted( + name + for name, initializer in initializers.items() + if initializer.const_value is None and name not in targets + ) + if missing_targets: + raise NativeCsaExportError( + f"{len(missing_targets)} {component_name} initializer(s) have no " + f"checkpoint mapping (e.g. {missing_targets[:5]})" + ) + + return StreamingWeightPlan( + targets=targets, + ignored=ignored, + report={ + "output_weight_format": "native_planar_block_quant", + "native_fp8": True, + "native_planar_fp4": True, + "runtime_execution_proven": False, + "runtime_dependency": "justinchuby/onnx-genai#2321", + "component": component_name, + }, + ) + def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: """Map the official DeepSeek checkpoint names to mobius modules.""" # Full-export runtime-capability gate. When a native-CSA export deferred # #602's block-quant reject (so graph construction could progress), the - # runnable export still fails closed here -- before any weight is - # mapped/assigned, replacing the former generic "Weight shape mismatch" - # with a typed, actionable blocker -- until nxrt advertises the real - # block-FP8 / planar-FP4 format strings. No-op on every ordinary path. + # canonical planar producer is representable before any weight is + # mapped/assigned. No-op on every ordinary path. assert_native_runtime_supports_block_quant(self.config) # Same predicate as MoELayer/_supported_qmoe_quantization so the # repacked weights and the emitted graph never disagree. use_qmoe = supported_qmoe_quantization(self.config.quantization) is not None + use_planar = self.config.block_quant_scheme is not None renamed: dict[str, torch.Tensor] = {} skipped = 0 for key, value in state_dict.items(): if key.startswith("mtp.") and len(self.mtp) == 0: skipped += 1 continue - new_key = key - if new_key == "embed.weight": - new_key = "model.embed_tokens.weight" - elif new_key == "head.weight": - new_key = "lm_head.weight" - elif new_key == "norm.weight": - new_key = "model.norm.weight" - elif new_key.startswith("hc_head_"): - new_key = ( - f"model.{new_key}.weight" - if new_key == "hc_head_fn" - else f"model.{new_key}" - ) - elif new_key.startswith("layers."): - new_key = f"model.{new_key}" - elif new_key.startswith("mtp."): + if key.startswith("mtp."): try: - mtp_index = int(new_key.split(".", 2)[1]) + mtp_index = int(key.split(".", 2)[1]) except (IndexError, ValueError): skipped += 1 continue if mtp_index >= len(self.mtp): skipped += 1 continue - if new_key.startswith(("model.layers.", "mtp.")): - new_key = new_key.replace(".attn.wq_a.", ".self_attn.q_a_proj.") - new_key = new_key.replace(".attn.q_norm.", ".self_attn.q_a_layernorm.") - new_key = new_key.replace(".attn.wq_b.", ".self_attn.q_b_proj.") - new_key = new_key.replace(".attn.wkv.", ".self_attn.kv_proj.") - new_key = new_key.replace(".attn.kv_norm.", ".self_attn.kv_layernorm.") - new_key = new_key.replace(".attn.wo_a.", ".self_attn.o_a_proj.") - new_key = new_key.replace(".attn.wo_b.", ".self_attn.o_b_proj.") - new_key = new_key.replace(".attn.", ".self_attn.") - new_key = new_key.replace(".attn_norm.", ".input_layernorm.") - new_key = new_key.replace(".ffn_norm.", ".post_attention_layernorm.") - # DeepSeekV4MoE composes the shared MoELayer, so the gate - # lives at mlp.moe.gate.* (see DeepSeekV4MoE.__init__). - new_key = new_key.replace(".ffn.gate.", ".mlp.moe.gate.") - new_key = new_key.replace(".ffn.experts.", ".mlp.experts.") - new_key = new_key.replace(".ffn.shared_experts.", ".mlp.shared_experts.") - new_key = new_key.replace(".w1.", ".gate_proj.") - new_key = new_key.replace(".w2.", ".down_proj.") - new_key = new_key.replace(".w3.", ".up_proj.") - if ".hc_" in new_key and new_key.endswith("_fn"): - new_key = f"{new_key}.weight" - # Dense fallback (unquantized or non-QMoE-eligible): experts - # are a plain ModuleList under moe.experts.{i}.*. Skipped for - # the QMoE path -- stack_per_expert_moe_weights below expects - # this same per-index ".mlp.experts.{i}.*" layout as input, - # fusing it into the tensors pack_qmoe_expert_weights expects. - if not use_qmoe: - new_key = new_key.replace(".mlp.experts.", ".mlp.moe.experts.") + new_key = _map_checkpoint_weight_name(key) + # Dense fallback (unquantized or non-QMoE-eligible): experts are a + # ModuleList under ``mlp.moe.experts``. Fused QMoE/planar paths + # consume the intermediate ``mlp.experts`` names as banks. + if not use_qmoe and not use_planar: + new_key = new_key.replace(".mlp.experts.", ".mlp.moe.experts.") renamed[new_key] = value if skipped: logger.warning( @@ -1480,7 +1848,10 @@ def preprocess_weights( skipped, ) processed = super().preprocess_weights(renamed) - if use_qmoe: + if use_planar: + _validate_hash_routing_tables(processed) + processed = _pack_planar_expert_weights(processed) + elif use_qmoe: _validate_hash_routing_tables(processed) # DeepSeek-V4 checkpoints store routed experts per-index # (".mlp.experts.{i}.gate_proj/up_proj/down_proj.*"), unlike diff --git a/src/mobius/models/deepseek_v4_flash_test.py b/src/mobius/models/deepseek_v4_flash_test.py index 7f726cc61..03fd440fe 100644 --- a/src/mobius/models/deepseek_v4_flash_test.py +++ b/src/mobius/models/deepseek_v4_flash_test.py @@ -26,7 +26,6 @@ from mobius.models._deepseek_v4_csa import ( CSA_COMPRESSION_RATIO, CSA_DOMAIN, - CSA_OP_TYPE, HCA_COMPRESSION_RATIO, NativeCsaExportError, assert_native_runtime_supports_block_quant, @@ -880,8 +879,16 @@ def test_native_csa_threads_compressed_state_io(): pres_kv = outputs["present_compressed_kv.1"] assert past_kv.dtype == ir.DataType.FLOAT assert pres_kv.dtype == ir.DataType.FLOAT - assert [str(d) for d in past_kv.shape] == ["batch", "past_compressed_records", "16"] - assert [str(d) for d in pres_kv.shape] == ["batch", "present_compressed_records", "16"] + assert [str(d) for d in past_kv.shape] == [ + "batch", + "past_compressed_records.1", + "16", + ] + assert [str(d) for d in pres_kv.shape] == [ + "batch", + "present_compressed_records.1", + "16", + ] past_carry = inputs["past_compression_carry.1"] pres_carry = outputs["present_compression_carry.1"] @@ -1057,6 +1064,16 @@ def test_native_csa_emits_both_ratios_for_interleaved_schedule(): assert ratio4.attributes["cache_format"].value == "fp8_e4m3_block64" assert ratio128.attributes["index_topk"].value == 0 assert ratio4.attributes["index_topk"].value == 4 + inputs = _named(graph.inputs) + outputs = _named(graph.outputs) + assert ( + str(inputs["past_compressed_kv.1"].shape[1]) + == (str(inputs["past_index_key.1"].shape[1])) + == "past_compressed_records.1" + ) + assert str(inputs["past_compressed_kv.2"].shape[1]) == "past_compressed_records.2" + assert str(outputs["present_compressed_kv.1"].shape[1]) == ("present_compressed_records.1") + assert str(outputs["present_compressed_kv.2"].shape[1]) == ("present_compressed_records.2") def test_native_csa_ratio4_threads_index_state_io(): @@ -1273,18 +1290,12 @@ def test_plan_native_csa_ratio4_layer_matches_contract(): # (E4M3 weight + 2D UE8M0 [128, 128] block scale) with *packed-fp4* routed # experts (top-level ``expert_dtype=fp4``). # -# The two slices compose in TWO stages, mandated by the directive -# "graph construction should progress past the former generic shape mismatch, -# full export must typed-reject at the runtime capability gate": +# The two slices compose in two stages: # # * A non-native export keeps #602's early, loud config-resolution reject. -# * A native-CSA export DEFERS that reject (records ``block_quant_scheme``) -# so graph construction PROGRESSES -- CSA nodes + compressed state IO are -# built and inspectable -- and the runnable *full export* then fails closed -# at the runtime-capability gate (``assert_native_runtime_supports_block_quant`` -# / ``preprocess_weights``) until nxrt advertises the real block-FP8 / -# planar-FP4 format strings. Never a silent dense fallback, never a partial -# native + silent dense graph. +# * A native-CSA export records ``block_quant_scheme`` and emits canonical +# block-FP8 MatMul plus planar-FP4 MoE nodes alongside CSA state IO. Never a +# silent dense fallback, never a partial native + silent dense graph. # --------------------------------------------------------------------------- # Frozen verbatim from the official config.json @ @@ -1421,73 +1432,151 @@ def test_native_csa_defers_block_quant_so_graph_construction_can_progress(): assert config.block_quant_scheme.has_packed_fp4_experts -def test_native_csa_full_export_typed_rejects_at_runtime_capability_gate(): - # The deferred config still fails closed at the runtime-capability gate: - # the runnable full export (assert_native_runtime_supports_block_quant, and - # therefore preprocess_weights) raises the typed BlockQuantExportError while - # nxrt cannot execute block-FP8 / planar-FP4 weights. Never silent dense. +def test_native_csa_runtime_capability_gate_accepts_planar_v1(): hf = _v4_hf_config( quantization_config=_REAL_BLOCK_FP8_QUANT_CONFIG, expert_dtype="fp4", native_csa=True, ) config = ArchitectureConfig.from_transformers(hf) - with pytest.raises(BlockQuantExportError) as exc: - assert_native_runtime_supports_block_quant(config) - msg = str(exc.value) - assert "block-FP8" in msg or "block-fp8" in msg.lower() - assert "planar" in msg.lower() - # The gap string is sourced from the #602 runtime contract, so the gate - # tracks the real format strings and opens with no change here. + assert_native_runtime_supports_block_quant(config) gap = native_runtime_block_quant_gap(config.block_quant_scheme) - assert gap is not None - assert "nxrt" in gap + assert gap is None -def test_native_csa_graph_construction_progresses_past_block_quant(): - # With the deferred block-quant scheme present, graph construction PROGRESSES - # past the former generic weight-shape mismatch: build_from_module emits the - # frozen CSA node and its compressed state IO instead of raising. The gate - # lives on the *full export* (weight load), not on construction, so the same - # config that builds a graph still rejects the runnable export. +def test_native_csa_graph_emits_planar_weights_and_sparse_experts(): scheme = BlockQuantScheme.from_quantization_config( _REAL_BLOCK_FP8_QUANT_CONFIG, expert_dtype="fp4" ) assert scheme is not None and scheme.is_owned - config = _tiny_config( - num_hidden_layers=2, - compress_ratios=[0, 128], + config = _ratio4_config( + num_hidden_layers=1, + compress_ratios=[4], native_csa=True, block_quant_scheme=scheme, + moe_intermediate_size=32, + num_nextn_predict_layers=0, ) graph = build_from_module(DeepSeekV4CausalLMModel(config), config, task="deepseek-v4")[ "model" ].graph csa = [n for n in _csa_nodes(graph) if n.domain == CSA_DOMAIN] assert len(csa) == 1 - assert csa[0].op_type == CSA_OP_TYPE - # ... but the runnable full export of that same config still fails closed. - with pytest.raises(BlockQuantExportError): - assert_native_runtime_supports_block_quant(config) + assert csa[0].attributes["compression_ratio"].value == 4 + + index_mm = next( + node + for node in graph + if node.domain == "pkg.nxrt" + and node.op_type == "BlockQuantizedMatMul" + and node.inputs[1].name.endswith("self_attn.indexer.wq_b.weight") + ) + assert len(index_mm.inputs) == 4 + assert [value.name for value in index_mm.inputs[:3]] == [ + index_mm.inputs[0].name, + "model.layers.0.self_attn.indexer.wq_b.weight", + "model.layers.0.self_attn.indexer.wq_b.scale", + ] + assert index_mm.attributes["format"].value == "block_fp8" + assert index_mm.attributes["block_size_out"].value == 128 + assert index_mm.attributes["block_size_in"].value == 128 + assert graph.initializers[index_mm.inputs[1].name].dtype == ir.DataType.FLOAT8E4M3FN + assert graph.initializers[index_mm.inputs[2].name].dtype == ir.DataType.FLOAT8E8M0 + index_query_reshape = csa[0].inputs[11].producer() + assert index_query_reshape is not None + assert index_query_reshape.inputs[0].producer() is index_mm + assert "q_a_layernorm" in index_mm.inputs[0].name + + moe = next( + node + for node in graph + if node.domain == "pkg.nxrt" and node.op_type == "BlockQuantizedMoE" + ) + assert len(moe.inputs) == 12 + assert [moe.inputs[index].name for index in (2, 4, 6, 9, 10, 11)] == [ + "model.layers.0.mlp.moe.fc1_experts_weights", + "model.layers.0.mlp.moe.fc2_experts_weights", + "model.layers.0.mlp.moe.fc3_experts_weights", + "model.layers.0.mlp.moe.fc1_experts_aux_scale", + "model.layers.0.mlp.moe.fc2_experts_aux_scale", + "model.layers.0.mlp.moe.fc3_experts_aux_scale", + ] + assert moe.attributes["fc1_format"].value == "fp4_planar" + assert moe.attributes["fc2_format"].value == "fp4_planar" + assert moe.attributes["fc3_format"].value == "fp4_planar" + assert moe.attributes["block_layout_version"].value == 1 + assert count_op_type(graph, "QMoE") == 0 + assert count_op_type(graph, "GroupQueryAttention") == 0 + assert not any(".mlp.moe.experts." in name for name in graph.initializers) + assert_native_runtime_supports_block_quant(config) -def test_preprocess_weights_enforces_runtime_capability_gate(): - # The weight-load hook (full export) enforces the same gate: a native-CSA - # module whose config carries a deferred block-quant scheme rejects in - # preprocess_weights before mapping a single tensor -- the typed blocker - # replaces the former generic "Weight shape mismatch". +def test_preprocess_weights_maps_real_planar_names_and_banks(): scheme = BlockQuantScheme.from_quantization_config( _REAL_BLOCK_FP8_QUANT_CONFIG, expert_dtype="fp4" ) - config = _tiny_config( - num_hidden_layers=2, - compress_ratios=[0, 128], + config = _ratio4_config( + num_hidden_layers=1, + compress_ratios=[4], native_csa=True, block_quant_scheme=scheme, + moe_intermediate_size=32, + num_nextn_predict_layers=0, ) module = DeepSeekV4CausalLMModel(config) - with pytest.raises(BlockQuantExportError): - module.preprocess_weights({}) + + def e8m0(*shape): + return torch.zeros(shape, dtype=torch.uint8).view(torch.float8_e8m0fnu) + + state = { + "layers.0.attn.indexer.wq_b.weight": torch.zeros(64, 8, dtype=torch.float8_e4m3fn), + "layers.0.attn.indexer.wq_b.scale": e8m0(1, 1), + } + for expert in range(2): + for projection in ("w1", "w2", "w3"): + state[f"layers.0.ffn.experts.{expert}.{projection}.weight"] = torch.zeros( + 32, 16, dtype=torch.int8 + ) + state[f"layers.0.ffn.experts.{expert}.{projection}.scale"] = e8m0(32, 1) + processed = module.preprocess_weights(state) + assert ( + processed["model.layers.0.self_attn.indexer.wq_b.scale"].dtype == torch.float8_e8m0fnu + ) + assert processed["model.layers.0.mlp.moe.fc1_experts_weights"].shape == ( + 2, + 32, + 16, + ) + assert processed["model.layers.0.mlp.moe.fc3_experts_aux_scale"].shape == ( + 2, + 32, + 1, + ) + + +def test_streaming_plan_fails_closed_for_missing_planar_scale(): + scheme = BlockQuantScheme.from_quantization_config( + _REAL_BLOCK_FP8_QUANT_CONFIG, expert_dtype="fp4" + ) + config = _ratio4_config( + num_hidden_layers=1, + compress_ratios=[4], + native_csa=True, + block_quant_scheme=scheme, + moe_intermediate_size=32, + num_nextn_predict_layers=0, + ) + module = DeepSeekV4CausalLMModel(config) + graph = build_from_module(module, config, task="deepseek-v4")["model"].graph + malformed_header = { + "layers.0.attn.indexer.wq_b.weight": ( + "shard.safetensors", + [64, 8], + "F8_E4M3", + ) + } + with pytest.raises(NativeCsaExportError, match="malformed block-quant checkpoint tensor"): + module.build_block_quant_streaming_plan("model", malformed_header, graph.initializers) def test_native_csa_is_property_gated_not_a_blanket_v4_refusal(): diff --git a/src/mobius/rewrite_rules/_block_quantized_moe_fusion.py b/src/mobius/rewrite_rules/_block_quantized_moe_fusion.py index 10b13e96c..04b40299a 100644 --- a/src/mobius/rewrite_rules/_block_quantized_moe_fusion.py +++ b/src/mobius/rewrite_rules/_block_quantized_moe_fusion.py @@ -24,10 +24,8 @@ * Experts are ``BlockQuantizedMatMul`` (native blocks) not ``MatMulNBits`` (int4 affine); there are no scale/zero-point inputs to carry. * GLM-5.2 UD-IQ1 quantises the ``gate``/``up`` projections and the ``down`` - projection to *different* native formats (e.g. ``iq1_s`` gate/up, ``iq4_xs`` - down). ``BlockQuantizedMoE`` expresses this with ``block_layout_version=2`` - and per-projection ``fc1_format`` / ``fc2_format`` / ``fc3_format`` attributes - (a uniform-format layer stays on ``block_layout_version=1``). + projection to *different* native formats. The canonical v1 ABI carries + per-projection ``fc1_format`` / ``fc2_format`` / ``fc3_format`` attributes. * Routing is reconstructed **gate-agnostically** from the already-computed ``selected_experts`` / ``routing_weights`` tensors, so a sigmoid+bias+scaling GLM gate fuses identically to a plain softmax top-k gate. The routing decision @@ -42,8 +40,8 @@ The rewrite is **property-gated and fails closed**: a layer that is not a native dense-expert storm is skipped, and a *routed native-block* storm whose experts -cannot be expressed as one expert-major bank (mixed native formats, per-expert -bias, an incomplete/untraceable expert group, ...) raises +cannot be expressed as one expert-major bank (per-expert bias, an +incomplete/untraceable expert group, ...) raises :class:`~mobius.integrations.gguf.SparseMoEExportError` -- the same typed capability error the GGUF builder's sparse-MoE honesty gate raises -- so export fails closed rather than silently shipping a dense-all-expert graph. Pass @@ -477,6 +475,7 @@ def _build_routing( weights so that (with ``normalize_routing_weights=0``) the kernel gathers exactly those weights. The result is independent of the originating gate. """ + assert layer.routing_weights is not None neg_one = _make_initializer( graph, f"{prefix}.neg_one", np.array([-1], dtype=np.int64), ir.DataType.INT64 ) @@ -609,14 +608,8 @@ def _plan_layer( ) -> _FusionPlan: """Validate and byte-stack one native MoE layer without touching the graph. - ``allow_perproj_v2_schema`` is a private, test-only switch. A layer that - mixes native formats across its fc1/fc2/fc3 banks can only be expressed as a - ``block_layout_version=2`` per-projection node, and no shipped onnx-genai - runtime executes that ABI yet. It therefore defaults off on every production, - CLI, and environment path, and planning fails closed for such a layer (before - any graph mutation) rather than emitting an unrunnable node. It is set - ``True`` only by schema-construction tests that assert the v2 node's shape; - the node they build is not runnable and must never be shipped. + ``allow_perproj_v2_schema`` is retained as an ignored compatibility + parameter for callers predating the canonical per-projection v1 ABI. """ ids = sorted(layer.experts) gate_nodes = [layer.experts[i].gate for i in ids] @@ -668,29 +661,12 @@ def _plan_layer( "normalize_routing_weights": 0, "swiglu_fusion": swiglu_fusion, } - distinct = set(projection_formats.values()) - if len(distinct) == 1: - attributes["format"] = next(iter(distinct)) - else: - if not allow_perproj_v2_schema: - raise _UnfusableError( - "mixed per-projection native formats " - f"{sorted(distinct)} can only be expressed with the " - "block_layout_version=2 per-projection BlockQuantizedMoE ABI, " - "which no shipped onnx-genai runtime implements -- emitting a v2 " - "node would build an unrunnable graph (overclaim), so export " - "fails closed. There is no production, CLI, or environment opt-in; " - "v2 stays a schema-construction test path until a typed runtime " - "capability handshake ships" - ) - attributes["block_layout_version"] = 2 - # ``format`` is the required base/fallback; per-projection attributes - # override it where a projection uses a different native format. - attributes["format"] = projection_formats["fc1"] - attributes["fc1_format"] = projection_formats["fc1"] - attributes["fc2_format"] = projection_formats["fc2"] - if "fc3" in projection_formats: - attributes["fc3_format"] = projection_formats["fc3"] + del allow_perproj_v2_schema + attributes["block_layout_version"] = 1 + attributes["fc1_format"] = projection_formats["fc1"] + attributes["fc2_format"] = projection_formats["fc2"] + if "fc3" in projection_formats: + attributes["fc3_format"] = projection_formats["fc3"] return _FusionPlan( layer=layer, @@ -708,6 +684,9 @@ def _emit_layer(graph: ir.Graph, plan: _FusionPlan) -> None: """Materialise a validated :class:`_FusionPlan` into the graph.""" prefix = plan.prefix layer = plan.layer + assert layer.routed_out is not None + routed_producer = layer.routed_out.producer() + assert routed_producer is not None fc3_w_v: ir.Value | None = None if plan.fc3_w is not None: @@ -739,6 +718,9 @@ def _emit_layer(graph: ir.Graph, plan: _FusionPlan) -> None: fc3_w_v, None, router_weights, + None, + None, + None, ], attributes=plan.attributes, domain=_NXRT_DOMAIN, @@ -768,7 +750,7 @@ def _emit_layer(graph: ir.Graph, plan: _FusionPlan) -> None: else: final_out = moe_out - graph.insert_after(layer.routed_out.producer(), new_nodes) + graph.insert_after(routed_producer, new_nodes) layer.routed_out.replace_all_uses_with(final_out) graph.opset_imports[_NXRT_DOMAIN] = 1 @@ -843,8 +825,8 @@ def fuse_block_quantized_moe( (pure expert-major stack/concat, no requantization). Routing is reconstructed gate-agnostically from the graph's ``selected_experts`` / ``routing_weights`` tensors, so softmax, sigmoid+bias+scaling, and other top-k gates all fuse. - Mixed per-projection native formats are expressed with - ``block_layout_version=2``. + Mixed per-projection native formats are expressed by the canonical + per-projection ``block_layout_version=1`` attributes. Every candidate layer is fully validated and byte-stacked *before* any node is emitted, so a fail-closed layer never leaves the graph half-rewritten. A @@ -864,17 +846,7 @@ def fuse_block_quantized_moe( ``allow_dense_moe_experts`` flag / ``MOBIUS_ALLOW_DENSE_MOE_EXPERTS`` environment variable is used. This is a research/correctness path only and makes no throughput claim. - _allow_perproj_v2_schema: Private, test-only switch. A layer that mixes - native formats across its fc1/fc2/fc3 banks (e.g. GLM-5.2 UD-IQ1 with - an iq1 gate/up and a higher-bit down) can only be expressed with the - ``block_layout_version=2`` per-projection ``BlockQuantizedMoE`` ABI. - No shipped onnx-genai runtime implements that ABI, so this defaults to - ``False`` (fail closed): the production ``build_from_gguf`` path never - sets it, and a mixed-format layer typed-rejects rather than emitting an - unrunnable v2 node. It is ``True`` only in schema-construction tests - that assert the v2 node's shape; that node is not runnable and must - never be shipped. Until a real typed runtime-capability handshake - exists there is no production, CLI, or environment path to v2. + _allow_perproj_v2_schema: Deprecated compatibility argument; ignored. Returns: The number of MoE layers fused. @@ -882,6 +854,7 @@ def fuse_block_quantized_moe( Raises: SparseMoEExportError: A routed native-block MoE layer cannot be sparse fused and ``allow_dense_moe`` is not set. + """ if allow_dense_moe is None: from mobius._flags import flags diff --git a/src/mobius/rewrite_rules/_block_quantized_moe_fusion_test.py b/src/mobius/rewrite_rules/_block_quantized_moe_fusion_test.py index 5c1ebcd0d..1ff9576b5 100644 --- a/src/mobius/rewrite_rules/_block_quantized_moe_fusion_test.py +++ b/src/mobius/rewrite_rules/_block_quantized_moe_fusion_test.py @@ -423,8 +423,8 @@ def test_input_wiring_matches_runtime_abi() -> None: model, _ = _build_dense_graph(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") fuse_block_quantized_moe(model, _allow_perproj_v2_schema=True) moe = _moe(model.graph) - # 9-input BlockQuantizedMoE ABI order (fused SwiGLU: no fc3). - assert len(moe.inputs) == 9 + # Canonical 12-input BlockQuantizedMoE ABI (fused SwiGLU: no fc3/scales). + assert len(moe.inputs) == 12 assert moe.inputs[0].name == "hidden" # input assert moe.inputs[1].name.endswith("router_logits") # router_logits assert moe.inputs[2].name.endswith("fc1_experts_weights") @@ -434,6 +434,9 @@ def test_input_wiring_matches_runtime_abi() -> None: assert moe.inputs[6] is None # fc3 (fused -> absent) assert moe.inputs[7] is None # fc3 bias assert moe.inputs[8].name.endswith("router_weights") + assert moe.inputs[9] is None # interleaved fc1 has no auxiliary scale + assert moe.inputs[10] is None # interleaved fc2 has no auxiliary scale + assert moe.inputs[11] is None # absent fc3 has no auxiliary scale # --------------------------------------------------------------------------- # @@ -477,27 +480,27 @@ def test_expert_major_bank_dtype_is_uint8() -> None: # --------------------------------------------------------------------------- # -# Per-projection format attributes (v1 vs v2) # +# Per-projection format attributes (canonical v1) # # --------------------------------------------------------------------------- # def test_uniform_format_stays_layout_version_1() -> None: model, _ = _build_dense_graph(gate_fmt="iq4_xs", up_fmt="iq4_xs", down_fmt="iq4_xs") fuse_block_quantized_moe(model) moe = _moe(model.graph) attrs = moe.attributes - assert "block_layout_version" not in attrs # defaults to v1 - assert attrs["format"].value == "iq4_xs" - assert "fc1_format" not in attrs - assert "fc2_format" not in attrs + assert attrs["block_layout_version"].value == 1 + assert "format" not in attrs + assert attrs["fc1_format"].value == "iq4_xs" + assert attrs["fc2_format"].value == "iq4_xs" assert "fc3_format" not in attrs -def test_fused_mixed_format_emits_layout_version_2() -> None: +def test_fused_mixed_format_emits_canonical_layout_version_1() -> None: model, _ = _build_dense_graph(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") fuse_block_quantized_moe(model, _allow_perproj_v2_schema=True) moe = _moe(model.graph) attrs = moe.attributes - assert attrs["block_layout_version"].value == 2 - assert attrs["format"].value == "iq1_s" # base/fallback + assert attrs["block_layout_version"].value == 1 + assert "format" not in attrs assert attrs["fc1_format"].value == "iq1_s" assert attrs["fc2_format"].value == "iq4_xs" assert "fc3_format" not in attrs # fused -> no fc3 @@ -509,7 +512,7 @@ def test_unfused_mixed_format_emits_all_projection_formats() -> None: fuse_block_quantized_moe(model, _allow_perproj_v2_schema=True) moe = _moe(model.graph) attrs = moe.attributes - assert attrs["block_layout_version"].value == 2 + assert attrs["block_layout_version"].value == 1 assert attrs["fc1_format"].value == "iq1_s" assert attrs["fc3_format"].value == "iq3_xxs" assert attrs["fc2_format"].value == "iq4_xs" @@ -526,22 +529,15 @@ def test_core_attributes_match_bqmoe_abi() -> None: # --------------------------------------------------------------------------- # -# Per-projection (block_layout_version=2): fail-closed default + schema hook # +# Per-projection canonical-v1 production path # # --------------------------------------------------------------------------- # -def test_mixed_format_rejects_v2_by_default() -> None: - """A direct public call rejects a mixed-format (v2-only) layer by default. - - No shipped onnx-genai runtime executes the ``block_layout_version=2`` - per-projection ABI, so ``fuse_block_quantized_moe`` fails closed for a - mixed-format layer with no opt-in of any kind. The reject is atomic: the - dense graph is left untouched. - """ +def test_mixed_format_fuses_to_v1_by_default() -> None: model, _ = _build_dense_graph(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") - graph = model.graph - with pytest.raises(SparseMoEExportError, match=r"block_layout_version=2"): - fuse_block_quantized_moe(model) - assert _count(graph, "BlockQuantizedMoE") == 0 - assert _count(graph, "Equal") == E # dense graph left untouched + assert fuse_block_quantized_moe(model) == 1 + attrs = _moe(model.graph).attributes + assert attrs["block_layout_version"].value == 1 + assert attrs["fc1_format"].value == "iq1_s" + assert attrs["fc2_format"].value == "iq4_xs" def test_uniform_format_still_fuses_to_v1() -> None: @@ -549,22 +545,16 @@ def test_uniform_format_still_fuses_to_v1() -> None: model, _ = _build_dense_graph(gate_fmt="iq4_xs", up_fmt="iq4_xs", down_fmt="iq4_xs") assert fuse_block_quantized_moe(model) == 1 attrs = _moe(model.graph).attributes - assert "block_layout_version" not in attrs # v1 - assert attrs["format"].value == "iq4_xs" - + assert attrs["block_layout_version"].value == 1 + assert attrs["fc1_format"].value == "iq4_xs" + assert attrs["fc2_format"].value == "iq4_xs" -def test_perproj_v2_schema_hook_is_test_only() -> None: - """The private ``_allow_perproj_v2_schema`` hook builds the v2 node shape. - This is the schema-construction test path only: the emitted v2 node is not - runnable on any shipped runtime and is unreachable from the production - builder, the CLI, or an environment variable. It exists so the per-projection - v2 schema stays covered until a real typed runtime-capability handshake ships. - """ +def test_legacy_schema_hook_still_emits_canonical_v1() -> None: model, _ = _build_dense_graph(gate_fmt="iq1_s", up_fmt="iq1_s", down_fmt="iq4_xs") assert fuse_block_quantized_moe(model, _allow_perproj_v2_schema=True) == 1 attrs = _moe(model.graph).attributes - assert attrs["block_layout_version"].value == 2 + assert attrs["block_layout_version"].value == 1 assert attrs["fc1_format"].value == "iq1_s" assert attrs["fc2_format"].value == "iq4_xs" diff --git a/src/mobius/tasks/_deepseek_v4.py b/src/mobius/tasks/_deepseek_v4.py index 78bc62839..a6df63430 100644 --- a/src/mobius/tasks/_deepseek_v4.py +++ b/src/mobius/tasks/_deepseek_v4.py @@ -84,17 +84,17 @@ def _compressed_inputs(builder, module, batch): When ``config.native_csa`` is off every plan is ``None``, so no inputs are created and the returned list is all-``None`` (byte-identical to - the pre-CSA graph). The compressed-record axis is a shared dynamic - symbolic dim because a layer's attention cache and index cache advance - in lockstep, and every CSA layer advances its cache together. + the pre-CSA graph). Each layer has a distinct dynamic record axis; + within a ratio-4 layer its compressed and index caches still advance in + lockstep. """ - records = ir.SymbolicDim("past_compressed_records") past_compressed_states: list = [] for layer in module.model.layers: plan = layer.self_attn.csa_plan if plan is None: past_compressed_states.append(None) continue + records = ir.SymbolicDim(plan.past_records_axis_name) past_compressed_kv = builder.input( plan.past_compressed_kv_name, dtype=plan.cache_dtype, @@ -152,12 +152,12 @@ def _compressed_outputs( top-k result ``[batch, index_num_heads, sequence, min(records, topk)]`` (inspection-only; not threaded back as state). """ - present_records = ir.SymbolicDim("present_compressed_records") - selected_records = ir.SymbolicDim("selected_records") for layer, present in zip(module.model.layers, present_compressed_states): plan = layer.self_attn.csa_plan if plan is None: continue + present_records = ir.SymbolicDim(plan.present_records_axis_name) + selected_records = ir.SymbolicDim(plan.selected_records_axis_name) present_compressed_kv = present[0] present_compression_carry = present[1] present_compressed_kv.shape = ir.Shape([batch, present_records, plan.stored_width]) diff --git a/tests/cli_test.py b/tests/cli_test.py index 4e8aee556..15657db96 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -596,6 +596,46 @@ def test_features_prune_prefill_prefix_passed_through(self): ) assert mock_build.call_args.kwargs.get("prune_prefill_prefix") is True + def test_features_native_csa_passed_through(self): + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius.integrations.diffusers._builder._load_diffusers_pipeline_index", + return_value=None, + ), + mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as mock_build, + mock.patch("mobius.__main__._save_package"), + ): + main( + [ + "build", + "--model", + "some/model", + tmpdir, + "--no-weights", + "--features", + "native-csa", + ] + ) + assert mock_build.call_args.kwargs.get("native_csa") is True + + def test_features_native_csa_rejects_dequantize(self): + with ( + tempfile.TemporaryDirectory() as tmpdir, + pytest.raises(SystemExit, match=r"native-csa.*--dequantize"), + ): + main( + [ + "build", + "--model", + "some/model", + tmpdir, + "--features", + "native-csa", + "--dequantize", + ] + ) + def test_features_comma_separated_multiple(self): """A single --features accepts a comma-separated list.""" with (