Add opt-in --paged-attention export for dense MLA (LATENT PagedAttention) [Slice 3B]

src/mobius/__main__.py

@@ -44,6 +44,7 @@
     "prune-prefill-prefix": "prune_prefill_prefix",
     "text-only": "text_only",
     "glm-full-attention": "glm_full_attention",
+    "paged-attention": "export_paged_attention",
 }
 
 
@@ -269,6 +270,23 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
         }
     else:
         static_cache_params = None
+
+    # PagedAttention (LATENT dense-MLA) export uses the paged-cache task with
+    # 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)
+    if export_paged_attention:
+        if static_cache_params is not None:
+            raise SystemExit(
+                "Error: --features paged-attention cannot be combined with "
+                "--features static-cache."
+            )
+        if task is not None:
+            raise SystemExit(
+                "Error: --features paged-attention cannot be combined with --task. "
+                "Remove --task to use --features paged-attention."
+            )
+        task = CausalLMTask(paged_cache=True)
     trust_remote_code = args.trust_remote_code
     revision = args.revision
     output_dir = args.output_dir
@@ -358,6 +376,16 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
                     f"model_type 'glm_moe_dsa' (got '{model_type}')."
                 )
             config = dataclasses.replace(config, use_dsa=False)
+        if export_paged_attention:
+            from mobius.components._paged_mla import paged_attention_rejection
+
+            config = dataclasses.replace(config, export_paged_attention=True)
+            reason = paged_attention_rejection(config)
+            if reason is not None:
+                raise SystemExit(
+                    f"Error: --features paged-attention is not supported for this "
+                    f"model: {reason}"
+                )
         if static_cache_params is not None:
             task = _resolve_static_cache_task(model_type)
         elif task is None:
@@ -406,6 +434,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             kv_cache_scales=kv_cache_scales,
             prune_prefill_prefix=prune_prefill_prefix,
             glm_full_attention=args.glm_full_attention,
+            export_paged_attention=export_paged_attention,
         )
 
     _save_package(pkg, output_dir, args, optimize, component_filter)

src/mobius/_configs/_base.py

@@ -503,6 +503,12 @@ class ArchitectureConfig(BaseModelConfig):
     # (reusing ``DeepSeekV3TextModel`` unchanged) on runtimes that cannot yet
     # execute ``pkg.nxrt::IndexShare``.
     use_dsa: bool = True
+    # Opt-in export of ``com.microsoft::PagedAttention`` (LATENT / absorbed-MLA
+    # mode) for property-compatible dense MLA. Default off. When off, exports are
+    # byte-identical to the current dense-MLA graph. Eligibility is decided from
+    # semantic geometry (see ``mobius.components._paged_mla``), never model names;
+    # an incompatible geometry raises rather than silently falling back.
+    export_paged_attention: bool = False
     # Per-layer indexer schedule ("full" runs the indexer; "shared" reuses the
     # top-k selection from the closest preceding "full" layer). When the
     # checkpoint config omits this list, it is derived from

src/mobius/components/__init__.py

@@ -207,6 +207,24 @@
     MLPMultiModalProjector,
 )
 from mobius.components._muse_glimmer_vision import MuseGlimmerVisionModel
+from mobius.components._paged_mla import (
+    PagedCacheState as PagedCacheState,
+)
+from mobius.components._paged_mla import (
+    PagedLatentMLA as PagedLatentMLA,
+)
+from mobius.components._paged_mla import (
+    absorb_mla_weights as absorb_mla_weights,
+)
+from mobius.components._paged_mla import (
+    mla_paged_geometry as mla_paged_geometry,
+)
+from mobius.components._paged_mla import (
+    paged_attention_eligible as paged_attention_eligible,
+)
+from mobius.components._paged_mla import (
+    paged_attention_rejection as paged_attention_rejection,
+)
 from mobius.components._parakeet_audio import ParakeetFastConformerEncoder
 from mobius.components._pixtral_vision import (
     Mistral3MultiModalProjector as Mistral3MultiModalProjector,