metrale_model_layers/layer/transformer_layer/
capabilities.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `LayerCapabilities`, the side-effect-free questions the model asks a layer when it
4//! picks a route: graph eligibility, batched decode and verify support, rollback, SSM
5//! pool use and the MLA prefill mode.
6//!
7//! Owner: model-layers.
8//! Invariants: none beyond the types.
9
10/// 2026-09-26: A supertrait of `TransformerLayer`; see the module header.
11pub trait LayerCapabilities {
12    /// 2026-09-25: True when this layer's prefill attends only over the tokens it is
13    /// handed, so skipping a cached prefix would hide it from attention. The Qwen3
14    /// attention layer returns true when it runs MLA; the prefix lookup then matches
15    /// nothing unless `METRALE_MLA_PREFIX_SKIP=1` (`mla_prefill_needs_full_recompute`).
16    fn uses_local_mla_prefill(&self) -> bool {
17        false
18    }
19
20    /// 2026-09-25: Whether this layer's online FP8-KV calibration has frozen its scale;
21    /// `None` when the layer runs no online calibration (the default). The model lifts its
22    /// CUDA-graph suppression once no layer reports `Some(false)`
23    /// (`graphs_ready_after_fp8_kv_cal`).
24    fn fp8_calibration_frozen(&self) -> Option<bool> {
25        None
26    }
27
28    /// 2026-09-25: True when this layer's decode cannot run inside a CUDA graph, such as
29    /// the QSA indexer's host top-k. The model ORs it across layers into
30    /// `decode_graph_veto` when it is built.
31    fn decode_graph_unsupported(&self) -> bool {
32        false
33    }
34
35    /// 2026-09-25: True when this layer cannot serve a batched multi-sequence decode step.
36    /// The batched decode ORs it across layers into `hc_perseq` (`decode_a2.rs`) and then
37    /// runs each sequence through `decode`; the single-GPU fused decode+prefill ORs it into
38    /// `hc_qsa_perseq` (`decode_b.rs`) and then runs the batched decode and the prefill
39    /// separately.
40    fn decode_multi_seq_unsupported(&self) -> bool {
41        false
42    }
43
44    /// 2026-09-25: True when this layer keeps per-sequence state that lowering the
45    /// sequence's KV cursor does not rewind, such as a monotonic cache count or an n-gram
46    /// history. The model ORs it across layers, and the scheduler's `rollback_to_boundary`
47    /// then declines the rollback (`LayerStateNotRewindable`).
48    fn decode_rollback_unsupported(&self) -> bool {
49        false
50    }
51
52    /// 2026-09-25: True when this layer cannot serve a batched multi-sequence verify
53    /// (`decode_verify_multi`); `can_batch_verify_dispatch` then refuses the batch. It is
54    /// separate from [`Self::decode_multi_seq_unsupported`] because the answers can
55    /// differ.
56    fn decode_verify_multi_unsupported(&self) -> bool {
57        false
58    }
59
60    #[allow(clippy::too_many_arguments)]
61    /// 2026-09-25: True when a captured decode graph goes stale once a new sequence takes
62    /// this slot. Decode graphs are keyed by `slot_idx`, which is safe only while every
63    /// per-sequence address a capture bakes lives in the slot-addressed SSM pool. A layer
64    /// that allocates its own per-sequence buffers (GLM-5.3's DSA indexer cache) returns
65    /// true, and `free_sequence_dispatch` then drops that slot's graphs, except for a slot
66    /// index past the SSM pool's `max_slots`.
67    fn graph_stale_on_new_sequence(&self) -> bool {
68        false
69    }
70
71    /// 2026-09-25: True for an SSM layer. The model's split prefill then runs
72    /// `prefill_phase1`, `prefill_gdn_full` and `prefill_phase3` instead of `prefill`.
73    fn is_ssm_layer(&self) -> bool {
74        false
75    }
76
77    /// 2026-09-25: Whether this layer's recurrent state lives in the shared SSM pool. When
78    /// true (the default), sequence setup hands a linear-attention layer an `SsmLayerState`
79    /// with pool addresses and never calls its `alloc_state`. A linear-attention layer with
80    /// its own state type returns false; the Kimi K3 layer (`kimi_k3/bound.rs`) does, and
81    /// downcasts its state to `K3CpuFallbackState`.
82    fn uses_ssm_pool(&self) -> bool {
83        true
84    }
85
86    /// 2026-10-01: True when this layer's `prefill` honours a replay-tail capture
87    /// (`ForwardContext::midchunk_capture` with `replay_tail`) on every build: the rows
88    /// from `cap_local` on get the bits a prefill pass starting at `cap_local` gives them
89    /// (an SSM layer splits its conv and recurrence there and copies its state into the
90    /// plan's slots; an attention layer runs those rows as their own paged call). The
91    /// model plans a replay-tail capture only when every layer says so
92    /// (`TransformerModel::inpass_cut_capture_supported`); otherwise it keeps the
93    /// two-pass tail split.
94    fn supports_replay_tail_split(&self) -> bool {
95        false
96    }
97}
metrale_model_layers/layer/transformer_layer/capabilities.rs

metrale_model_layers/layer/transformer_layer/
capabilities.rs
