Summary
Introduced in #340. _register_kv_cache_outputs (src/mobius/tasks/_cache_utils.py, the ~_register_kv_cache_outputs helper, +63/-2 from a former 3-line stub) stamps explicit present.* output shapes/dtypes when all present-shape params are supplied. When only a partial set of present-shape params is provided, it currently emits logger.warning(...) and silently falls back to the known-wrong GQA shape-inference path (fail-open). A wiring bug therefore ships a structurally-wrong model with only a log line — no hard failure.
Risk
Fail-open behavior: a partial/misconfigured present-shape set produces an incorrect GQA static-cache model that may run but produce wrong results, with nothing but a warning to signal it. This is easy to miss in automated export pipelines that don't scrape logs.
Proposed fix
Add an opt-in strict: bool = False parameter (or a module-level config flag) that raises ValueError on a partial present-shape set instead of warn-and-proceed, listing the provided vs. missing params. The default stays warn for backward compatibility; production export paths opt into strict. Add a unit test asserting the raise.
Scope: small, isolated to _cache_utils.py + its test. Non-blocking.
Provenance
Note on Major-3b (separate / out of scope)
The related apply_weights() INFO-only logging of unmapped weights (src/mobius/_model_package.py:306) is pre-existing on main — _model_package.py is untouched by #340 (git diff --stat 5423421..c2bedd6 is empty for it). It is explicitly out of #340 scope and is not part of this follow-up. It is lower priority (a save-time _check_weights already raises on missing initializers); track separately if desired.
Summary
Introduced in #340.
_register_kv_cache_outputs(src/mobius/tasks/_cache_utils.py, the~_register_kv_cache_outputshelper, +63/-2 from a former 3-line stub) stamps explicitpresent.*output shapes/dtypes when all present-shape params are supplied. When only a partial set of present-shape params is provided, it currently emitslogger.warning(...)and silently falls back to the known-wrong GQA shape-inference path (fail-open). A wiring bug therefore ships a structurally-wrong model with only a log line — no hard failure.Risk
Fail-open behavior: a partial/misconfigured present-shape set produces an incorrect GQA static-cache model that may run but produce wrong results, with nothing but a warning to signal it. This is easy to miss in automated export pipelines that don't scrape logs.
Proposed fix
Add an opt-in
strict: bool = Falseparameter (or a module-level config flag) that raisesValueErroron a partial present-shape set instead of warn-and-proceed, listing the provided vs. missing params. The default stayswarnfor backward compatibility; production export paths opt intostrict. Add a unit test asserting the raise.Scope: small, isolated to
_cache_utils.py+ its test. Non-blocking.Provenance
5423421). Deferred from fp16/GQA export fixes + capture-safe static-cache attention (supersedes #328) #340's review as a tracked follow-up.Note on Major-3b (separate / out of scope)
The related
apply_weights()INFO-only logging of unmapped weights (src/mobius/_model_package.py:306) is pre-existing onmain—_model_package.pyis untouched by #340 (git diff --stat 5423421..c2bedd6is empty for it). It is explicitly out of #340 scope and is not part of this follow-up. It is lower priority (a save-time_check_weightsalready raises on missing initializers); track separately if desired.