From 3e7368010189344b96fcfafd2e533d650f2e671c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 25 Aug 2026 07:53:47 +0000 Subject: [PATCH] fix(deepseek-v4): per-layer CSA/HCA compressed-record symbolic dim The native CompressedSparseAttention exporter stamped a single shared symbolic dim (`past_compressed_records` / `present_compressed_records`) on every CSA layer's compressed-record axis, on the assumption that "every CSA layer advances its cache together". That is false for the official interleaved schedule: a ratio-4 CSA layer pools one record per ~4 tokens while a ratio-128 HCA layer pools one per ~128, so their record counts diverge as the sequence grows. A shared symbol then forces ORT to bind the same dim to two different sizes (e.g. 2 and 0 at prefill), which it correctly rejects: symbol component.model.past_compressed_records bound to conflicting sizes 2 and 0 across bound inputs Give each layer its own record-axis symbolic dim (`past_compressed_records.{layer_id}` / `present_compressed_records.{layer_id}` / `selected_records.{layer_id}`) via new `CsaLayerPlan` properties. Within a layer the attention cache and the learned-index cache still share the one per-layer symbol because they advance in lockstep at the same ratio -- a real constraint worth expressing -- but layers of different ratios no longer alias. Extends `test_native_csa_emits_both_ratios_for_interleaved_schedule` to assert the ratio-4 and ratio-128 layers carry distinct record axes and that the ratio-4 attention/index caches share their layer's axis; updates the ratio-128 IO test to the per-layer symbol name. No BC shim (dev-time API change). Surfaced while bringing up the onnx-genai native-decode CSA/HCA E2E proof against a tiny alternating ratio-4/ratio-128 fixture built by this exporter. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/_deepseek_v4_csa.py | 20 ++++++++++++++++++ src/mobius/models/deepseek_v4_flash_test.py | 23 +++++++++++++++++++-- src/mobius/tasks/_deepseek_v4.py | 21 +++++++++++-------- 3 files changed, 53 insertions(+), 11 deletions(-) diff --git a/src/mobius/models/_deepseek_v4_csa.py b/src/mobius/models/_deepseek_v4_csa.py index 8c4fa3208..bcf62931b 100644 --- a/src/mobius/models/_deepseek_v4_csa.py +++ b/src/mobius/models/_deepseek_v4_csa.py @@ -284,6 +284,26 @@ def present_index_carry_name(self) -> str: def selected_indices_name(self) -> str: return f"selected_indices.{self.layer_id}" + # -- dynamic record-axis symbolic dim names --------------------------- + # Each layer's compressed-record axis is a *per-layer* symbolic dim. A + # mixed schedule pools at different ratios per layer (ratio-4 keeps one + # record per 4 tokens, ratio-128 one per 128), so their record counts + # diverge and MUST NOT share one symbol -- ORT rejects binding the same + # symbolic dim to two sizes. Within a layer the attention cache and the + # learned-index cache advance in lockstep (same ratio), so they share this + # one per-layer symbol; that lockstep constraint is real and worth stating. + @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 [] diff --git a/src/mobius/models/deepseek_v4_flash_test.py b/src/mobius/models/deepseek_v4_flash_test.py index 7f726cc61..35d031cfb 100644 --- a/src/mobius/models/deepseek_v4_flash_test.py +++ b/src/mobius/models/deepseek_v4_flash_test.py @@ -880,8 +880,8 @@ 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"] @@ -1058,6 +1058,25 @@ def test_native_csa_emits_both_ratios_for_interleaved_schedule(): assert ratio128.attributes["index_topk"].value == 0 assert ratio4.attributes["index_topk"].value == 4 + # Regression: a mixed schedule pools ratio-4 and ratio-128 layers at + # different rates, so their compressed-record axes MUST be distinct + # symbolic dims (a shared symbol makes ORT reject binding the same dim to + # two record counts). Within the ratio-4 layer the attention and index + # caches advance in lockstep, so they SHARE that layer's one record symbol. + inputs = _named(graph.inputs) + outputs = _named(graph.outputs) + r4_kv_axis = str(inputs["past_compressed_kv.1"].shape[1]) + r4_ik_axis = str(inputs["past_index_key.1"].shape[1]) + r128_kv_axis = str(inputs["past_compressed_kv.2"].shape[1]) + assert r4_kv_axis == r4_ik_axis == "past_compressed_records.1" + assert r128_kv_axis == "past_compressed_records.2" + assert r4_kv_axis != r128_kv_axis + r4_present = str(outputs["present_compressed_kv.1"].shape[1]) + r128_present = str(outputs["present_compressed_kv.2"].shape[1]) + assert r4_present == "present_compressed_records.1" + assert r128_present == "present_compressed_records.2" + assert r4_present != r128_present + def test_native_csa_ratio4_threads_index_state_io(): # ratio-4 threads the packed uint8 attention/index caches, the f32 carries, diff --git a/src/mobius/tasks/_deepseek_v4.py b/src/mobius/tasks/_deepseek_v4.py index 78bc62839..ba6acdbdf 100644 --- a/src/mobius/tasks/_deepseek_v4.py +++ b/src/mobius/tasks/_deepseek_v4.py @@ -84,17 +84,19 @@ 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). The compressed-record axis is a *per-layer* dynamic + symbolic dim: within a layer the attention cache and index cache advance + in lockstep (same ratio), but a mixed ratio-4/ratio-128 schedule pools + at different rates per layer, so their record counts diverge and cannot + share one symbol. """ - 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, @@ -143,21 +145,22 @@ def _compressed_outputs( symbolic-shape-inference function (like ``pkg.nxrt::IndexShare`` in the GLM DSA task), so each present output is stamped with an explicit type or it would export untyped. The compressed-record axis is a distinct - dynamic symbolic dim (present record count = past + newly pooled - blocks, not a simple ``past + sequence`` sum); the carry tensors are - records-independent ``[batch, slots, planes, width]``. + *per-layer* dynamic symbolic dim (present record count = past + newly + pooled blocks, not a simple ``past + sequence`` sum, and each layer + pools at its own ratio); the carry tensors are records-independent + ``[batch, slots, planes, width]``. Ratio-4 additionally emits the packed uint8 ``present_index_key``, the f32 ``present_index_carry``, and the transient int32 ``selected_indices`` 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])