fix(deepseek-v4): per-layer CSA/HCA compressed-record symbolic dim

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 []

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,

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])