Add opt-in --paged-attention export for dense MLA (LATENT PagedAttention) [Slice 3B] - #599
Conversation
…ion) Slice 3B: emit `com.microsoft::PagedAttention` v1 in LATENT/absorbed-MLA mode for property-compatible *dense* MLA (DeepSeek-V2/V3, GLM-5.2 `--glm-full-attention`). Feature is default off; feature-off exports are byte-identical to the current dense-MLA graph. Eligibility is decided purely from semantic geometry (latent width, v_head_size, partial-RoPE suffix, cache dtype, page/block constraints), never from model names. Active query-dependent sparse selection (GLM DSA, DeepSeek-V4 CSA/HCA), MTP, quantized cache, head_sink, qk-norm and sliding window are typed-rejected at construction/build; feature-on with an incompatible geometry errors rather than silently falling back to dense. GLM's vestigial indexer config (dropped under --glm-full-attention) does not reject. The op consumes/mutates caller-owned page buffers in place; Mobius never allocates or manages pages (onnx-genai-kv stays the sole cache authority). The graph binds caller-owned block_table / slot_mapping / cumulative + past lengths / per-layer LATENT key_cache, aliases the cache in place, and derives token positions from the length tensors (no position_ids input). kv_b_proj is absorbed into the query/output projections at weight-apply time (numeric contract mirrors the onnx-genai equivalence oracle). Wiring: - config flag `export_paged_attention` (default False) - `components/_paged_mla.py`: geometry, typed eligibility, weight absorption, `PagedLatentMLA` component, `PagedCacheState` - `CausalLMTask(paged_cache=True)`: caller-owned LATENT cache IO - DeepSeek/GLM model integration + weight absorption - CLI `--features paged-attention` (+ task resolution, mutual exclusion with static-cache / --task, typed eligibility error) Tests (same commit): - component: eligibility rejects, absorption, numpy decomposed-vs-absorbed LATENT parity (rel<1e-6), structural emission - model-level: feature-off byte-identical, feature-on structural + exact op attrs/inputs/outputs/model IO, typed rejects (DSA/CSA/HCA/MTP/window), torch-tensor weight absorption - CLI: feature plumbing + mutual-exclusion errors No full-size performance claim; native runtime full-size verification is separate. Requires ORT >=1.29 CUDA to execute (contrib op unknown to the base onnx checker, so model-level assertions are structural). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
…ad_size %16 Address the code-review BLOCKER: use_dsa defaults to True on every ArchitectureConfig (and via the HF resolver), and is only set False by --glm-full-attention (gated to glm_moe_dsa). Plain DeepSeek-V2/V3 never reads use_dsa (its text model is always dense), so gating the DSA typed rejection on use_dsa alone wrongly rejected real DeepSeek-V3 paged export. DSA is now treated as active only when use_dsa is set AND an indexer is actually configured (index_n_heads / index_head_dim / index_topk / indexer_types). This keeps: - DeepSeek-V3 (use_dsa=True default, no indexer) -> eligible - GLM --glm-full-attention (use_dsa=False, vestigial indexer) -> eligible - GLM DSA active (use_dsa=True + indexer) -> typed reject Also tighten the LATENT geometry gate to mirror the native validator's check_rotary_caches (validate.rs:510): dense MLA always emits do_rotary=1 with cos/sin caches, so head_size must be a multiple of 16, which (given rotary_dim %16) forces kv_lora_rank %16 == 0. Closes the gap where an 8-aligned-but-not-16-aligned latent width would pass export yet be rejected at native load. Tests: DSA-reject tests now configure an indexer; new regression that DeepSeek-V3 with the default use_dsa=True still emits paged nodes; geometry reject test now uses kv_lora_rank=8 (8- but not 16-aligned). All paged component + model + CLI + GLM/DeepSeek suites pass; ruff 0.16.2 clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Independent review round 1 — verdict: 1 BLOCKER (fixed), 1 should-fix (fixed), 1 nitAn independent reviewer (not Leon/Sapper) ran the full suite (84 paged/GLM/DeepSeek/CLI tests pass), ruff clean, and cross-checked the emitted BLOCKER (fixed): SHOULD-FIX (fixed): Tightened the LATENT geometry gate to mirror the native validator's NIT (acknowledged): byte-identical feature-off test is close to tautological by construction (feature default-off); kept as a guard against accidental default flips. No behavioral change when the feature is off (still byte-identical). No full-size perf claims. Native runtime full-size verification remains a separate gate. Do not merge — awaiting Gaff/Roy final approval. |
…ess, real qk-norm gate No-BC cleanup of the eligibility validator (no compat shims). Every field the rejection probes is a declared ArchitectureConfig field, so replace the defensive getattr(...) with direct attribute access (use_dsa, indexer_types, compress_ratios, o_lora_rank, o_groups, hc_mult, num_nextn_predict_layers, sliding_window) and the export_paged_attention reads in deepseek.py/glm_moe_dsa.py. Behavior is identical: each getattr default matched the declared default. Correctness fix (mandatory typed refusal): per-head QK-norm is consumed into the canonical attn_qk_norm / attn_qk_norm_full fields at config-extraction time (the HF use_qk_norm / qk_layernorm passthrough is not an ArchitectureConfig field, so the old check was always False and unreachable). Reject on the real fields instead — a qk-norm MLA would need the operator's q_norm/k_norm inputs, which this slice does not claim. DeepSeek-V2/V3 and GLM full-attention leave both False, so the targets are unaffected. The refusal is now reachable and tested at both component and model level. Remove the provably-dead head_sink / attention_sink branch: neither is an ArchitectureConfig field, so the guard was always False. Coverage is subsumed — the only sink-bearing MLA family (DeepSeek-V4) is already rejected by its CSA/HCA fields, a non-MLA sink model (gpt-oss) is rejected by the MLA-geometry check, and PagedLatentMLA structurally never emits a head_sink input. Documented in place; no silent-miscompute path is introduced. Tests: add component test_qk_norm_rejected (attn_qk_norm / attn_qk_norm_full) and test_window_rejected; add a model-level qk-norm typed-reject case. Full paged/GLM/DeepSeek/V4/tasks/CLI suites pass; ruff 0.16.2 clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
3a3eb2d to
8a17cfb
Compare
APPROVE 3B FOR MERGEIndependently reviewed at exact HEAD Op-contract verification (against onnx-genai
|
Follow-up: no-backward-compat honesty cleanup (
|
Roy — Final re-review of PR #599 at HEAD
|
There was a problem hiding this comment.
Pull request overview
Adds opt-in LATENT PagedAttention export for dense MLA models, with paged-cache wiring, geometry validation, weight absorption, and CLI support.
Changes:
- Adds LATENT
com.microsoft::PagedAttentiongraph emission. - Integrates paged cache handling across tasks, builders, and supported models.
- Adds structural, parity, rejection, and CLI tests.
Final review findings:
src/mobius/__main__.py: Reject or supportort-genaipaged-cache ABI (critical, 1 vote).src/mobius/components/_paged_mla.py: Reconcile%16gates with the advertised runtime contract (moderate, 2 votes).src/mobius/integrations/transformers/_builder.py: Select the paged task when enabled (critical, 4 votes).src/mobius/models/deepseek.py: Handle or reject quantized weight absorption (critical, 4 votes).src/mobius/tasks/_causal_lm.py: Use packed token/logit shapes for varlen paged attention (critical, 1 vote).src/mobius/tasks/_causal_lm.py: Preserve or reject paged mode with prefix pruning (critical, 2 votes).
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Summary |
|---|---|
tests/cli_test.py |
Tests paged-attention CLI plumbing and conflicts. |
src/mobius/tasks/_causal_lm.py |
Adds paged cache graph inputs and outputs. |
src/mobius/models/paged_mla_export_test.py |
Tests structural export and rejection cases. |
src/mobius/models/glm_moe_dsa.py |
Handles GLM dense-attention eligibility. |
src/mobius/models/deepseek.py |
Integrates paged MLA and weight absorption. |
src/mobius/integrations/transformers/_builder.py |
Adds builder support and validation. |
src/mobius/components/_paged_mla.py |
Implements LATENT attention, validation, and absorption. |
src/mobius/components/_paged_mla_test.py |
Tests geometry and numerical parity. |
src/mobius/components/__init__.py |
Exposes paged MLA components. |
src/mobius/_configs/_base.py |
Adds paged-attention configuration. |
src/mobius/__main__.py |
Adds CLI feature handling and validation. |
Suppressed comments (10)
src/mobius/main.py:289
- The CLI allows
--features paged-attention,prune-prefill-prefix, butbuild_from_module()then reconstructs this task through_enable_prefill_prefix_pruning_task()without preservingpaged_cache=True. The paged model consequently receives ordinary tuple caches and fails in_forward_paged()whendataclasses.replace()is called on them. Reject this feature combination or preserve the paged task and implement packed per-request last-token selection.
task = CausalLMTask(paged_cache=True)
src/mobius/main.py:277
export_paged_attentionis resolved here, but the diffusers and NeMo source-dispatch branches below call their own builders and return without passing or validating this flag. Thus--features paged-attentionsilently succeeds while exporting a non-paged diffusion/NeMo package, instead of producing the required typed incompatibility error. Add a fail-fast check for these source types before dispatch.
# 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)
src/mobius/main.py:283
- Paged and FP8-cache features are not mutually excluded here. With CUDA, the existing FP8 pass sees no
GroupQueryAttentionin this graph and raises after the paged graph is built; on unsupported EPs it only warns and leaves the cache unquantized. Either way--features paged-attention,fp8-kv-cacheis not a valid typed-rejected request as required for quantized cache modes. Reject the combination before build.
if export_paged_attention:
if static_cache_params is not None:
raise SystemExit(
"Error: --features paged-attention cannot be combined with "
"--features static-cache."
)
src/mobius/main.py:437
- With
--runtime onnx-genai, this newly exposed package reaches the generic decoder workflow builder, which does not classifyblock_table,slot_mapping,cumulative_sequence_length, orpast_seqlensas paged state. It treatsblock_tableas a fallback position input and raises on the other unsupported inputs, so this export cannot writeinference_metadata.yaml. Add a paged control/state mapping or reject this runtime explicitly.
export_paged_attention=export_paged_attention,
src/mobius/components/init.py:212
- These new symbols are re-exported from
mobius.componentsbelow but are missing from the module's__all__, unlike the existing public cache/component symbols. That makes wildcard imports and API introspection omitPagedLatentMLA/PagedCacheStateand the new helpers; add the intended public names to__all__, or keep them private consistently.
from mobius.components._paged_mla import (
PagedCacheState as PagedCacheState,
)
src/mobius/integrations/transformers/_builder.py:259
fp8_kv_cacheis accepted alongsideexport_paged_attention, but the FP8 pass only convertsGroupQueryAttention; this graph containsPagedAttention, so optimization warns/ignores the request and leaves a non-FP8 cache instead of typed-rejecting the unsupported combination. Reject the combination before building (for both CLI and direct builder APIs).
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 ValueError(
"export_paged_attention=True (--features paged-attention) is not "
f"supported for model_type '{model_type}': {reason}"
)
src/mobius/integrations/transformers/_builder.py:259
- The eligibility predicate admits any config with matching MLA geometry, but not every such model implements the paged forward/cache contract (for example, LongCat Flash has MLA-shaped config fields but its attention layers still consume normal two-part KV state). After task selection is corrected, this produces a paged task feeding
PagedCacheStateinto an unsupported implementation. Add an implementation capability check/protocol so unsupported architectures receive a typed rejection instead of a runtime type/shape failure.
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 ValueError(
"export_paged_attention=True (--features paged-attention) is not "
f"supported for model_type '{model_type}': {reason}"
)
src/mobius/tasks/_causal_lm.py:562
- These new paged control inputs are not handled by the existing ONNX GenAI metadata builder. It only admits a rank-2 token/attention/position input;
block_tablebecomes the fallback position input, while rank-1slot_mapping,cumulative_sequence_length, andpast_seqlensremain unsupported andwrite_onnx_genai_config()raises. Add a dedicated paged control/state-service contract (including paged cache layout), or reject--runtime onnx-genaifor this export instead of emitting an unusable package.
block_table = builder.input(
"block_table", dtype=ir.DataType.INT32, shape=[batch, max_blocks]
)
slot_mapping = builder.input("slot_mapping", dtype=ir.DataType.INT32, shape=[num_tokens])
cumulative_sequence_length = builder.input(
"cumulative_sequence_length", dtype=ir.DataType.INT32, shape=["batch + 1"]
)
past_seqlens = builder.input("past_seqlens", dtype=ir.DataType.INT32, shape=[batch])
src/mobius/tasks/_causal_lm.py:297
- Unlike the normal path above,
_build_pagedinvokes the module outsideprefill_prefix_pruning(...)and never calls_validate_pruned_logits(...). ACausalLMTask(paged_cache=True, prune_prefill_prefix=True)therefore silently emits full-sequence logits, breaking the documented prefix-pruning option even if task reconstruction is fixed. Apply the same context and validation around the paged module call.
result = module(
op,
input_ids=input_ids,
attention_mask=None,
position_ids=None,
past_key_values=paged_states,
)
src/mobius/tasks/_causal_lm.py:301
- The ordinary
CausalLMTaskpath preserves a third model result and registershidden_states.*foroutput_layer_indices, but this path binds that result to_intermediateand drops it. Selecting paged attention therefore changes the model's declared outputs for any compatible module/config that requests intermediate hidden states. Mirror the normal unpacking and_register_intermediate_hidden_stateshandling.
if len(result) == 3:
logits, present_key_values, _intermediate = result
else:
logits, present_key_values = result
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| "prune-prefill-prefix": "prune_prefill_prefix", | ||
| "text-only": "text_only", | ||
| "glm-full-attention": "glm_full_attention", | ||
| "paged-attention": "export_paged_attention", |
| if head_size % 16 != 0: | ||
| return f"PagedAttention LATENT requires head_size % 16 == 0; got {head_size}." | ||
| if l % 16 != 0: | ||
| return f"PagedAttention LATENT requires latent_dim (kv_lora_rank) % 16 == 0; got {l}." |
| 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 ValueError( | ||
| "export_paged_attention=True (--features paged-attention) is not " | ||
| f"supported for model_type '{model_type}': {reason}" | ||
| ) |
| if self.config.export_paged_attention: | ||
| renamed = self._absorb_paged_mla_weights(renamed) |
| paged_states = _make_paged_cache_inputs( | ||
| builder, | ||
| config.num_hidden_layers, | ||
| geom.head_size, | ||
| config.dtype, | ||
| batch, | ||
| ) |
| static_cache: bool = False, | ||
| paged_cache: bool = False, | ||
| max_seq_len: int | None = None, | ||
| prune_prefill_prefix: bool = False, | ||
| ): | ||
| if static_cache and paged_cache: | ||
| raise ValueError("static_cache and paged_cache are mutually exclusive.") | ||
| self._static_cache = static_cache | ||
| self._paged_cache = paged_cache |
Slice 3B — Mobius opt-in
--paged-attentionexport (dense MLA, LATENT)Emits
com.microsoft::PagedAttentionv1 in LATENT / absorbed-MLA mode for property-compatible dense MLA (DeepSeek-V2/V3, GLM-5.2--glm-full-attention). Builds on the merged onnx-genai native runtime contract: audit/typed-validator/oracle (#1940), KV index emission (#1955), CUDA LATENT kernel (#1978).Guarantees
v_head_size, partial-RoPE suffix (rotary_offset = kv_lora_rank), cache dtype (fp16/bf16), and the onnx-genai-kv page/block constraints (head_size %8, kv_lora %8, rotary_dim %16, rotary_offset %8, block pow2 ≥16).--glm-full-attention) does not reject.onnx-genai-kvstays the sole authority. The graph binds caller-ownedblock_table/slot_mapping/ cumulative + past lengths / per-layer LATENTkey_cache, aliases the cache in place, and derives token positions from the length tensors (noposition_idsinput).Emitted contract
PagedAttentionnode per layer,_domain="com.microsoft", LATENT attrs (kv_cache_layout=LATENT,kv_num_heads=1, explicitscale,v_head_size=kv_lora_rank,rotary_offset=kv_lora_rank,do_rotary=1);rotary_dimis derived fromcos_cache, never emitted.input_ids,block_table(i32),slot_mapping(i32),cumulative_sequence_length(i32),past_seqlens(i32), per-layerkey_cache.{i}[num_blocks, block_size, 1, head_size]. Outputs:logits,updated_key_cache.{i}(aliases input).kv_b_projis absorbed into the query/output projections at weight-apply time; numeric contract mirrors the onnx-genai equivalence oracle.Tests (same commit)
components/_paged_mla_test.py): eligibility rejects, absorption, numpy decomposed-vs-absorbed LATENT parity (rel < 1e-6), structural emission.models/paged_mla_export_test.py): feature-off byte-identical; feature-on structural + exact op attrs/inputs/outputs/model IO; typed rejects (DSA/CSA/HCA/MTP/window); torch-tensor weight absorption; DeepSeek-V3 path proves eligibility is not name-gated.tests/cli_test.py):--features paged-attentionplumbing + task resolution + mutual-exclusion errors.Scope / gates
Draft — do not merge. Independent review required (reviewer excluding Leon/Sapper; Gaff or Roy final approval).
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com