metrale_model_layers/layers/qwen3_attention/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The full-attention transformer layer, `Qwen3AttentionLayer`:
4//! its construction, kernel requirements, decode and prefill attention, and
5//! the call into the layer's FFN.
6//!
7//! Owner: model-layers (attention).
8//! Invariants: none beyond the types.
9
10// 2026-09-25: The N-column-blocked W8A16 decode tier for the attention
11// projections. `init` reads its lever once and caches it on the layer
12// (`attn_ncol`); both multi-seq routes, strided QKV and contiguous o_proj, read
13// that field.
14mod attn_ncol_gemv;
15mod circuit;
16// 2026-09-25: The log-once route lines of the `METRALE_ATTN_M16_TC` tier, one
17// for the strided QKV call site and one for the o_proj call site.
18pub(crate) mod attn_m16_tc_route;
19mod decode;
20// 2026-09-25: `pub` because the DeepSeek-V4 weight loader in
21// metrale-model-arch calls `helpers::yarn_rope_mscale`.
22pub mod helpers;
23mod init;
24mod init_arch_gates;
25mod init_decode_kernels;
26mod init_kernel_dispatch;
27mod init_prefill_kernels;
28mod init_proj_kernels;
29mod kernel_requirements;
30mod op_dump;
31// 2026-09-25: `innerq_driver` uses `metrale_gpu_runtime::registry`, which
32// exists only with the `cuda` feature, so the module carries the same gate.
33#[cfg(feature = "cuda")]
34pub mod innerq_driver;
35mod prefill;
36mod prefill_qkv_w8a8;
37mod prefill_w8a8;
38mod prefill_weights;
39#[cfg(test)]
40mod rope_site_fixture;
41mod trait_impl;
42pub mod types;
43mod types_weights;
44mod w8a8_decode_arm;
45
46#[cfg(feature = "cuda")]
47pub use innerq_driver::InnerQDriver;
48pub(crate) use types::HeadGateActivation;
49pub use types::Qwen3AttentionLayer;
50pub use types_weights::{
51    CompressorWeights, Fp8TwinSet, HcHeadWeights, HcLowRank, HcSiteWeights, HcWeights, MlaWeights,
52    W8A8_PREFILL_KERNELS, w8a8_prefill_kernels_loaded,
53};
54
55/// 2026-09-25: Startup check for one `--kv-cache-dtype`: resolves every kernel
56/// the dtype's dispatch needs (the reshape/decode pair, the chunked-prefill
57/// kernel and, for a WHT-rotated side, the WHT bookends) and errors with the
58/// full missing list. The server calls it before `build_model`. See
59/// `kernel_requirements.rs`.
60pub fn validate_required_kv_kernels(
61    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
62    kv_dtype: metrale_cache::kv_cache::KvCacheDtype,
63    head_dim: usize,
64) -> anyhow::Result<()> {
65    kernel_requirements::validate_required_kernels(gpu, kv_dtype, head_dim)
66}
67
68/// 2026-09-25: The reference sequence count that the `Legacy` split-K rule
69/// (`metrale_kernels::attn_splitk::legacy_splits`) divides the SM count by:
70/// the configured max decode batch (`ModelLevers::max_decode_seqs`), raised to
71/// `num_seqs` when the live batch is wider.
72///
73/// Using the max batch rather than the live count keeps a sequence's split
74/// count, and so its split-merge reduction tree, the same whether it is decoded
75/// alone or co-batched; the merge is not associative, so a different tree can
76/// change the temp-0 argmax. The raise to `num_seqs` keeps
77/// `num_seqs * num_q_heads * num_splits <= sm_count` whenever the rule picks
78/// more than one split, which is the bound the `Legacy` workspace is sized to
79/// (`attn_splitk::workspace_slots`).
80pub(crate) fn split_ref_seqs(num_seqs: u32, max_decode_seqs: u32) -> u32 {
81    max_decode_seqs.max(num_seqs)
82}
83
84/// 2026-09-25: Host time in µs spent in the FFN of prefill attention layers,
85/// added when `METRALE_PREFILL_HOST_TIMING=1`. Summed across layers; the layer
86/// loop takes it once per prefill and reports the attention remainder as the
87/// in-layer time minus this.
88pub static FFN_HOST_US: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
89
90pub fn add_ffn_host_us(us: u64) {
91    FFN_HOST_US.fetch_add(us, std::sync::atomic::Ordering::Relaxed);
92}
93
94pub fn take_ffn_host_us() -> u64 {
95    FFN_HOST_US.swap(0, std::sync::atomic::Ordering::Relaxed)
96}
97
98/// 2026-09-25: Per-phase host time in µs for the prefill attention path, written
99/// only by `prefill/cache_skip.rs`: 0 = the Q/K/V projections, 1 = from the
100/// projections to the BR=64 attention call, 2 = that attention call. No code
101/// writes slot 3. Summed across layers; the layer loop takes them once per
102/// prefill when `METRALE_PREFILL_HOST_TIMING=1`.
103pub static ATTN_PHASE_US: [std::sync::atomic::AtomicU64; 4] = [
104    std::sync::atomic::AtomicU64::new(0),
105    std::sync::atomic::AtomicU64::new(0),
106    std::sync::atomic::AtomicU64::new(0),
107    std::sync::atomic::AtomicU64::new(0),
108];
109
110pub fn add_attn_phase_us(i: usize, us: u64) {
111    ATTN_PHASE_US[i].fetch_add(us, std::sync::atomic::Ordering::Relaxed);
112}
113
114pub fn take_attn_phase_us() -> [u64; 4] {
115    let mut o = [0u64; 4];
116    for (i, a) in ATTN_PHASE_US.iter().enumerate() {
117        o[i] = a.swap(0, std::sync::atomic::Ordering::Relaxed);
118    }
119    o
120}
121
122#[cfg(test)]
123mod split_ref_seqs_tests {
124    use super::split_ref_seqs;
125
126    #[test]
127    fn the_split_count_does_not_move_with_co_batch_size() {
128        // 2026-09-25: A sequence decoded alone and the same sequence co-batched
129        // get the same reference count, so the `Legacy` rule picks the same
130        // split count for both.
131        let pin = 16;
132        assert_eq!(split_ref_seqs(1, pin), split_ref_seqs(8, pin));
133        assert_eq!(split_ref_seqs(1, pin), pin);
134    }
135
136    #[test]
137    fn a_batch_larger_than_the_pin_clamps_up() {
138        assert_eq!(split_ref_seqs(32, 16), 32);
139    }
140
141    #[test]
142    fn two_models_can_pin_to_different_batches() {
143        assert_ne!(split_ref_seqs(1, 4), split_ref_seqs(1, 16));
144    }
145}
metrale_model_layers/layers/qwen3_attention/mod.rs

metrale_model_layers/layers/qwen3_attention/
mod.rs
