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