use std::collections::HashSet;
use crate::error::InferenceError;
use crate::model::qwen::QwenConfig;
use crate::model::qwen35_config::Qwen35Config;
use crate::quant::quarot::hadamard::{
MAX_BLOCK_HADAMARD_BLOCKS, MAX_BLOCK_HADAMARD_LEN, RandomizedHadamard,
};
use crate::quant::quarot::rotation::{
absorb_input_rotation, absorb_input_rotation_f64, absorb_output_rotation,
absorb_output_rotation_f64,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum AbsorptionSide {
InputSide,
OutputSide,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TensorRotation {
pub side: AbsorptionSide,
pub rotation_id: RotationId,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum RotationId {
ResidualStream,
AttentionOutputR3,
MlpDownR4,
}
impl RotationId {
pub fn online_transform_site(self) -> Option<OnlineTransformSite> {
match self {
RotationId::ResidualStream => None,
RotationId::AttentionOutputR3 => Some(OnlineTransformSite::AttentionOutputPreOProj),
RotationId::MlpDownR4 => Some(OnlineTransformSite::MlpPreDownProj),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OnlineTransformSite {
AttentionOutputPreOProj,
MlpPreDownProj,
}
impl OnlineTransformSite {
pub fn weight_tensor_suffix(self) -> &'static str {
match self {
OnlineTransformSite::AttentionOutputPreOProj => "self_attn.o_proj.weight",
OnlineTransformSite::MlpPreDownProj => "mlp.down_proj.weight",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct OnlineRotationSpec {
pub id: RotationId,
pub side: AbsorptionSide,
pub seed: u64,
pub block_size: usize,
pub layer_scope: Option<Vec<usize>>,
}
fn check_block_hadamard_num_blocks_cap(
axis_name: &str,
dim: usize,
block_size: usize,
) -> Result<(), InferenceError> {
if dim == 0 {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec: {axis_name} must be non-zero — \
BlockHadamard::new refuses a zero-length rotation axis"
)));
}
if dim > MAX_BLOCK_HADAMARD_LEN {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec: {axis_name} {dim} exceeds the \
MAX_BLOCK_HADAMARD_LEN cap of {MAX_BLOCK_HADAMARD_LEN} — this \
recipe cannot be materialized by BlockHadamard::new, so it is \
refused here rather than certified as a valid artifact"
)));
}
let num_blocks = dim / block_size;
if num_blocks > MAX_BLOCK_HADAMARD_BLOCKS {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec: {axis_name} {dim} with block_size \
{block_size} needs {num_blocks} BlockHadamard blocks, exceeding \
the MAX_BLOCK_HADAMARD_BLOCKS cap of \
{MAX_BLOCK_HADAMARD_BLOCKS} — this recipe cannot be \
materialized by BlockHadamard::new, so it is refused here \
rather than certified as a valid artifact"
)));
}
Ok(())
}
impl OnlineRotationSpec {
const MAX_LAYER_SCOPE_ENTRIES: usize = 4096;
pub fn r3_full_attention(
cfg: &Qwen35Config,
seed: u64,
block_size: usize,
) -> Result<Self, InferenceError> {
if block_size == 0 || !block_size.is_power_of_two() {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::r3_full_attention requires a power-of-two \
block_size, got {block_size}"
)));
}
if !cfg.num_attention_heads.is_multiple_of(block_size) {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::r3_full_attention: block_size {block_size} \
does not divide num_attention_heads {} — R3's online factor \
is QuaRot Eq. 9's cross-head Hadamard (H_num_heads ⊗ \
I_head_dim), so block_size divides the head axis, not \
head_dim",
cfg.num_attention_heads
)));
}
check_block_hadamard_num_blocks_cap(
"num_attention_heads",
cfg.num_attention_heads,
block_size,
)?;
let layers: Vec<usize> = (0..cfg.num_hidden_layers)
.filter(|&i| cfg.is_full_attention(i))
.collect();
if layers.is_empty() {
return Err(InferenceError::Inference(
"OnlineRotationSpec::r3_full_attention: config resolved zero \
full-attention layers — refusing an empty-scope R3 artifact"
.to_string(),
));
}
Ok(Self {
id: RotationId::AttentionOutputR3,
side: AbsorptionSide::InputSide,
seed,
block_size,
layer_scope: Some(layers),
})
}
pub fn r4_dense_mlp(
cfg: &Qwen35Config,
seed: u64,
block_size: usize,
) -> Result<Self, InferenceError> {
if cfg.is_moe() {
return Err(InferenceError::Inference(
"OnlineRotationSpec::r4_dense_mlp: this config is a MoE \
configuration (loader requires mlp.experts.down_proj, not \
the dense mlp.down_proj.weight this constructor targets) — \
MoE R4 targets are not yet modeled"
.to_string(),
));
}
if block_size == 0 || !block_size.is_power_of_two() {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::r4_dense_mlp requires a power-of-two \
block_size, got {block_size}"
)));
}
if !cfg.intermediate_size.is_multiple_of(block_size) {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::r4_dense_mlp: block_size {block_size} \
does not divide intermediate_size {}",
cfg.intermediate_size
)));
}
check_block_hadamard_num_blocks_cap(
"intermediate_size",
cfg.intermediate_size,
block_size,
)?;
Ok(Self {
id: RotationId::MlpDownR4,
side: AbsorptionSide::InputSide,
seed,
block_size,
layer_scope: None,
})
}
pub fn validate(&self, cfg: Option<&Qwen35Config>) -> Result<(), InferenceError> {
if self.side != AbsorptionSide::InputSide {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: {:?} requires side=InputSide \
(the only orientation proven correct by the R3/R4 reference \
test), got {:?}",
self.id, self.side
)));
}
if self.block_size == 0 || !self.block_size.is_power_of_two() {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: {:?} requires a power-of-two \
block_size, got {}",
self.id, self.block_size
)));
}
match self.id {
RotationId::ResidualStream => {
return Err(InferenceError::Inference(
"OnlineRotationSpec::validate: RotationId::ResidualStream \
is a purely offline rotation and must never appear as an \
online OnlineRotationSpec"
.to_string(),
));
}
RotationId::AttentionOutputR3 => {
let layers = self.layer_scope.as_ref().ok_or_else(|| {
InferenceError::Inference(
"OnlineRotationSpec::validate: AttentionOutputR3 (R3) \
requires an explicit non-empty layer_scope (full-\
attention layers only) — layer_scope=None is invalid \
for R3"
.to_string(),
)
})?;
if layers.is_empty() {
return Err(InferenceError::Inference(
"OnlineRotationSpec::validate: AttentionOutputR3 (R3) \
layer_scope must not be empty"
.to_string(),
));
}
if layers.len() > Self::MAX_LAYER_SCOPE_ENTRIES {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: R3 layer_scope declares {} \
entries, exceeding the maximum of {} — a spec this large \
is rejected before the per-layer membership check to keep \
validation cost bounded regardless of input",
layers.len(),
Self::MAX_LAYER_SCOPE_ENTRIES
)));
}
for pair in layers.windows(2) {
if pair[0] >= pair[1] {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: R3 layer_scope {layers:?} \
must be strictly sorted ascending with no \
duplicates — found {} at or after {}",
pair[1], pair[0]
)));
}
}
if let Some(cfg) = cfg {
if !cfg.num_attention_heads.is_multiple_of(self.block_size) {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: R3 block_size {} \
does not divide cfg.num_attention_heads {}",
self.block_size, cfg.num_attention_heads
)));
}
check_block_hadamard_num_blocks_cap(
"num_attention_heads",
cfg.num_attention_heads,
self.block_size,
)?;
for &idx in layers {
if idx >= cfg.num_hidden_layers || !cfg.is_full_attention(idx) {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: R3 layer_scope \
includes layer {idx}, which is not a valid \
full-attention layer for this config"
)));
}
}
let scoped: HashSet<usize> = layers.iter().copied().collect();
let required_full_attention_layers: Vec<usize> = (0..cfg.num_hidden_layers)
.filter(|&idx| cfg.is_full_attention(idx))
.collect();
let missing_layers: Vec<usize> = required_full_attention_layers
.iter()
.copied()
.filter(|idx| !scoped.contains(idx))
.collect();
if !missing_layers.is_empty() {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: R3 layer_scope {layers:?} \
does not cover all of this config's full-attention \
layers {required_full_attention_layers:?} — missing \
layer(s) {missing_layers:?}"
)));
}
}
}
RotationId::MlpDownR4 => {
if self.layer_scope.is_some() {
return Err(InferenceError::Inference(
"OnlineRotationSpec::validate: MlpDownR4 (R4) \
layer_scope must be None — every layer carries a \
dense MLP, so R4 is never restricted to a subset of \
layers"
.to_string(),
));
}
if let Some(cfg) = cfg {
if cfg.is_moe() {
return Err(InferenceError::Inference(
"OnlineRotationSpec::validate: MlpDownR4 (R4) \
targets the dense mlp.down_proj.weight tensor, \
but this config is a MoE configuration (loader \
requires mlp.experts.down_proj instead) — MoE \
R4 targets are not yet modeled by this plan"
.to_string(),
));
}
if !cfg.intermediate_size.is_multiple_of(self.block_size) {
return Err(InferenceError::Inference(format!(
"OnlineRotationSpec::validate: R4 block_size {} \
does not divide cfg.intermediate_size {}",
self.block_size, cfg.intermediate_size
)));
}
check_block_hadamard_num_blocks_cap(
"intermediate_size",
cfg.intermediate_size,
self.block_size,
)?;
}
}
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RuleRequirement {
Required,
Optional,
}
#[derive(Debug, Clone)]
struct Rule {
pattern: String,
rotation: TensorRotation,
requirement: RuleRequirement,
}
#[derive(Debug, Clone)]
pub struct RotationPlan {
rules: Vec<Rule>,
}
impl RotationPlan {
pub fn qwen35_residual_stream_linear_layers() -> Self {
let r_in = TensorRotation {
side: AbsorptionSide::InputSide,
rotation_id: RotationId::ResidualStream,
};
let r_out = TensorRotation {
side: AbsorptionSide::OutputSide,
rotation_id: RotationId::ResidualStream,
};
let req = |pat: &str, rot: TensorRotation| Rule {
pattern: pat.into(),
rotation: rot,
requirement: RuleRequirement::Required,
};
let opt = |pat: &str, rot: TensorRotation| Rule {
pattern: pat.into(),
rotation: rot,
requirement: RuleRequirement::Optional,
};
Self {
rules: vec![
req("self_attn.q_proj.weight", r_in),
req("self_attn.k_proj.weight", r_in),
req("self_attn.v_proj.weight", r_in),
req("self_attn.o_proj.weight", r_out),
req("linear_attn.in_proj_qkv.weight", r_in),
req("linear_attn.in_proj_z.weight", r_in),
req("linear_attn.in_proj_b.weight", r_in),
req("linear_attn.in_proj_a.weight", r_in),
req("linear_attn.out_proj.weight", r_out),
req("mlp.gate_proj.weight", r_in),
req("mlp.up_proj.weight", r_in),
req("mlp.down_proj.weight", r_out),
req("embed_tokens.weight", r_in),
opt("lm_head.weight", r_in),
],
}
}
pub fn qwen3_residual_stream_linear_layers() -> Self {
let r_in = TensorRotation {
side: AbsorptionSide::InputSide,
rotation_id: RotationId::ResidualStream,
};
let r_out = TensorRotation {
side: AbsorptionSide::OutputSide,
rotation_id: RotationId::ResidualStream,
};
let req = |pat: &str, rot: TensorRotation| Rule {
pattern: pat.into(),
rotation: rot,
requirement: RuleRequirement::Required,
};
let opt = |pat: &str, rot: TensorRotation| Rule {
pattern: pat.into(),
rotation: rot,
requirement: RuleRequirement::Optional,
};
Self {
rules: vec![
req("self_attn.q_proj.weight", r_in),
req("self_attn.k_proj.weight", r_in),
req("self_attn.v_proj.weight", r_in),
req("self_attn.o_proj.weight", r_out),
req("mlp.gate_proj.weight", r_in),
req("mlp.up_proj.weight", r_in),
req("mlp.down_proj.weight", r_out),
req("embed_tokens.weight", r_in),
opt("lm_head.weight", r_in),
],
}
}
pub fn for_tensor(&self, name: &str) -> Option<TensorRotation> {
self.rules
.iter()
.find(|rule| name.ends_with(&rule.pattern))
.map(|rule| rule.rotation)
}
pub fn absorption_for_module(&self, module: &str) -> Option<AbsorptionSide> {
let suffix = format!("{module}.weight");
self.rules
.iter()
.find(|rule| rule.pattern.ends_with(&suffix))
.map(|rule| rule.rotation.side)
}
pub fn rule_count(&self) -> usize {
self.rules.len()
}
pub fn validate_coverage<'a, I>(&self, tensor_names: I) -> CoverageReport
where
I: IntoIterator<Item = &'a str>,
{
let names: Vec<&str> = tensor_names.into_iter().collect();
let mut matched_tensors: Vec<String> = Vec::new();
let mut unplanned_tensors: Vec<String> = Vec::new();
let mut ambiguous_tensors: Vec<String> = Vec::new();
let mut unmatched_required_rules: Vec<String> = Vec::new();
let mut unmatched_optional_rules: Vec<String> = Vec::new();
for rule in &self.rules {
let any = names.iter().any(|n| n.ends_with(&rule.pattern));
if !any {
match rule.requirement {
RuleRequirement::Required => {
unmatched_required_rules.push(rule.pattern.clone())
}
RuleRequirement::Optional => {
unmatched_optional_rules.push(rule.pattern.clone())
}
}
}
}
for name in &names {
let match_count = self
.rules
.iter()
.filter(|rule| name.ends_with(&rule.pattern))
.count();
match match_count {
0 => unplanned_tensors.push((*name).to_string()),
1 => matched_tensors.push((*name).to_string()),
_ => ambiguous_tensors.push((*name).to_string()),
}
}
CoverageReport {
matched_tensors,
unplanned_tensors,
ambiguous_tensors,
unmatched_required_rules,
unmatched_optional_rules,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CoverageReport {
pub matched_tensors: Vec<String>,
pub unplanned_tensors: Vec<String>,
pub ambiguous_tensors: Vec<String>,
pub unmatched_required_rules: Vec<String>,
pub unmatched_optional_rules: Vec<String>,
}
impl CoverageReport {
pub fn is_complete(&self) -> bool {
self.unmatched_required_rules.is_empty() && self.ambiguous_tensors.is_empty()
}
}
pub fn apply_tensor_rotation(
name: &str,
weight: &mut [f32],
rows: usize,
cols: usize,
plan: &RotationPlan,
residual_rotation: &RandomizedHadamard,
) -> Result<bool, InferenceError> {
let Some(tr) = plan.for_tensor(name) else {
return Ok(false);
};
let rotation = match tr.rotation_id {
RotationId::ResidualStream => residual_rotation,
RotationId::AttentionOutputR3 | RotationId::MlpDownR4 => {
return Err(InferenceError::Inference(
"apply_tensor_rotation: R3/R4 rotation IDs are contract-only in \
this PR (issue #703 PR1) — no plan should reference them yet, \
and no online rotation is wired into offline absorption"
.to_string(),
));
}
};
match tr.side {
AbsorptionSide::InputSide => absorb_input_rotation(weight, rows, cols, rotation)?,
AbsorptionSide::OutputSide => absorb_output_rotation(weight, rows, cols, rotation)?,
}
Ok(true)
}
pub fn apply_tensor_rotation_f64(
name: &str,
weight: &mut [f64],
rows: usize,
cols: usize,
plan: &RotationPlan,
residual_rotation: &RandomizedHadamard,
) -> Result<bool, InferenceError> {
let Some(tr) = plan.for_tensor(name) else {
return Ok(false);
};
let rotation = match tr.rotation_id {
RotationId::ResidualStream => residual_rotation,
RotationId::AttentionOutputR3 | RotationId::MlpDownR4 => {
return Err(InferenceError::Inference(
"apply_tensor_rotation_f64: R3/R4 rotation IDs are contract-only \
in this PR (issue #703 PR1) — no plan should reference them \
yet, and no online rotation is wired into offline absorption"
.to_string(),
));
}
};
match tr.side {
AbsorptionSide::InputSide => absorb_input_rotation_f64(weight, rows, cols, rotation)?,
AbsorptionSide::OutputSide => absorb_output_rotation_f64(weight, rows, cols, rotation)?,
}
Ok(true)
}
pub fn qwen3_required_tensor_names(cfg: &QwenConfig) -> Vec<String> {
let mut names: Vec<String> = vec!["embed_tokens.weight".to_string(), "norm.weight".to_string()];
for i in 0..cfg.num_hidden_layers {
let p = format!("layers.{i}");
names.extend([
format!("{p}.self_attn.q_proj.weight"),
format!("{p}.self_attn.k_proj.weight"),
format!("{p}.self_attn.v_proj.weight"),
format!("{p}.self_attn.o_proj.weight"),
format!("{p}.self_attn.q_norm.weight"),
format!("{p}.self_attn.k_norm.weight"),
format!("{p}.input_layernorm.weight"),
format!("{p}.mlp.gate_proj.weight"),
format!("{p}.mlp.up_proj.weight"),
format!("{p}.mlp.down_proj.weight"),
format!("{p}.post_attention_layernorm.weight"),
]);
}
names
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn qwen35_plan_covers_residual_stream_tensors() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
assert_eq!(plan.rule_count(), 14);
let cases = [
(
"model.layers.0.self_attn.q_proj.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.5.self_attn.k_proj.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.5.self_attn.v_proj.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.5.self_attn.o_proj.weight",
AbsorptionSide::OutputSide,
),
(
"model.layers.0.linear_attn.in_proj_qkv.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.0.linear_attn.in_proj_z.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.0.linear_attn.in_proj_b.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.0.linear_attn.in_proj_a.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.0.linear_attn.out_proj.weight",
AbsorptionSide::OutputSide,
),
(
"model.layers.23.mlp.gate_proj.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.23.mlp.up_proj.weight",
AbsorptionSide::InputSide,
),
(
"model.layers.23.mlp.down_proj.weight",
AbsorptionSide::OutputSide,
),
(
"model.language_model.embed_tokens.weight",
AbsorptionSide::InputSide,
),
("lm_head.weight", AbsorptionSide::InputSide),
];
for (name, expected_side) in cases {
let tr = plan
.for_tensor(name)
.unwrap_or_else(|| panic!("plan missed tensor {name}"));
assert_eq!(tr.side, expected_side, "wrong side for {name}");
assert_eq!(tr.rotation_id, RotationId::ResidualStream);
}
}
#[test]
fn qwen35_plan_misses_rmsnorm_and_scalar_weights() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let skipped = [
"model.layers.0.input_layernorm.weight",
"model.layers.0.post_attention_layernorm.weight",
"model.norm.weight",
"model.layers.0.linear_attn.norm.weight",
"model.layers.0.linear_attn.A_log",
"model.layers.0.linear_attn.dt_bias",
"model.layers.0.linear_attn.conv1d.weight",
];
for name in skipped {
assert!(
plan.for_tensor(name).is_none(),
"non-linear-layer tensor {name} should not match plan"
);
}
}
fn qwen35_required_residual_tensor_suffixes() -> &'static [&'static str] {
&[
"self_attn.q_proj.weight",
"self_attn.k_proj.weight",
"self_attn.v_proj.weight",
"self_attn.o_proj.weight",
"linear_attn.in_proj_qkv.weight",
"linear_attn.in_proj_z.weight",
"linear_attn.in_proj_b.weight",
"linear_attn.in_proj_a.weight",
"linear_attn.out_proj.weight",
"mlp.gate_proj.weight",
"mlp.up_proj.weight",
"mlp.down_proj.weight",
"embed_tokens.weight",
]
}
fn synthetic_qwen35_tensor_names(include_lm_head: bool) -> Vec<String> {
let mut names: Vec<String> = qwen35_required_residual_tensor_suffixes()
.iter()
.enumerate()
.map(|(i, s)| {
if s.contains("embed_tokens") {
format!("model.language_model.{s}")
} else {
format!("model.layers.{i}.{s}")
}
})
.collect();
if include_lm_head {
names.push("lm_head.weight".to_string());
}
names
}
#[test]
fn validate_coverage_tied_embeddings_is_complete_without_lm_head() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let names = synthetic_qwen35_tensor_names(false);
let report = plan.validate_coverage(names.iter().map(String::as_str));
assert!(
report.is_complete(),
"tied-embeddings tensor list should yield complete coverage, got: {report:?}"
);
assert_eq!(
report.unmatched_required_rules.len(),
0,
"no required rule should be missing: {:?}",
report.unmatched_required_rules
);
assert_eq!(
report.unmatched_optional_rules,
vec!["lm_head.weight".to_string()],
"lm_head should appear in unmatched_optional_rules"
);
assert_eq!(report.matched_tensors.len(), 13);
assert!(report.ambiguous_tensors.is_empty());
}
#[test]
fn validate_coverage_untied_embeddings_matches_lm_head() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let names = synthetic_qwen35_tensor_names(true);
let report = plan.validate_coverage(names.iter().map(String::as_str));
assert!(
report.is_complete(),
"untied case should be complete: {report:?}"
);
assert!(report.unmatched_optional_rules.is_empty());
assert_eq!(report.matched_tensors.len(), 14);
}
#[test]
fn validate_coverage_flags_unmatched_required_rule() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let names: Vec<String> = synthetic_qwen35_tensor_names(false)
.into_iter()
.filter(|n| !n.ends_with("self_attn.o_proj.weight"))
.collect();
let report = plan.validate_coverage(names.iter().map(String::as_str));
assert!(
!report.is_complete(),
"missing required rule should fail completeness"
);
assert_eq!(
report.unmatched_required_rules,
vec!["self_attn.o_proj.weight".to_string()]
);
}
#[test]
fn validate_coverage_lists_unplanned_tensors() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let mut names = synthetic_qwen35_tensor_names(false);
names.push("model.layers.0.input_layernorm.weight".to_string());
names.push("model.mtp.head.weight".to_string());
let report = plan.validate_coverage(names.iter().map(String::as_str));
assert!(
report.is_complete(),
"unplanned tensors should not break required-rule completeness"
);
assert_eq!(report.unplanned_tensors.len(), 2);
assert!(
report
.unplanned_tensors
.iter()
.any(|n| n.ends_with("input_layernorm.weight"))
);
}
#[test]
fn qwen35_plan_misses_moe_expert_weights() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let moe_names = [
"model.layers.0.mlp.experts.gate_up_proj.weight",
"model.layers.0.mlp.experts.down_proj.weight",
];
for name in moe_names {
assert!(
plan.for_tensor(name).is_none(),
"MoE expert tensor {name} should not match v0 plan"
);
}
}
#[test]
fn apply_tensor_rotation_skips_unplanned() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let hidden = 64;
let r = RandomizedHadamard::new(7, hidden).unwrap();
let mut weight = vec![1.0_f32; hidden * hidden];
let weight_copy = weight.clone();
let rotated = apply_tensor_rotation(
"model.layers.0.linear_attn.in_proj.weight",
&mut weight,
hidden,
hidden,
&plan,
&r,
)
.unwrap();
assert!(!rotated, "unplanned tensor should report not rotated");
assert_eq!(weight, weight_copy, "unplanned tensor must not be mutated");
}
#[test]
fn apply_tensor_rotation_mlp_layer_pair_rotates_output() {
let hidden = 64;
let intermediate = 128;
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let r = RandomizedHadamard::new(0xC0FFEE, hidden).unwrap();
let mut state = 1_u64;
let mut rand = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 11) as u32 as f32 / u32::MAX as f32) - 0.5
};
let gate_proj: Vec<f32> = (0..intermediate * hidden).map(|_| rand()).collect();
let down_proj: Vec<f32> = (0..hidden * intermediate).map(|_| rand()).collect();
let x: Vec<f32> = (0..hidden).map(|_| rand()).collect();
let matvec = |w: &[f32], r: usize, c: usize, v: &[f32]| -> Vec<f32> {
(0..r)
.map(|i| (0..c).map(|j| w[i * c + j] * v[j]).sum())
.collect()
};
let intermediate_out = matvec(&gate_proj, intermediate, hidden, &x);
let y_original = matvec(&down_proj, hidden, intermediate, &intermediate_out);
let mut y_expected = y_original.clone();
r.apply(&mut y_expected).unwrap();
let mut gate_proj_abs = gate_proj.clone();
let mut down_proj_abs = down_proj.clone();
assert!(
apply_tensor_rotation(
"model.layers.0.mlp.gate_proj.weight",
&mut gate_proj_abs,
intermediate,
hidden,
&plan,
&r,
)
.unwrap()
);
assert!(
apply_tensor_rotation(
"model.layers.0.mlp.down_proj.weight",
&mut down_proj_abs,
hidden,
intermediate,
&plan,
&r,
)
.unwrap()
);
let mut x_rotated = x.clone();
r.apply(&mut x_rotated).unwrap();
let intermediate_rot = matvec(&gate_proj_abs, intermediate, hidden, &x_rotated);
let y_rotated = matvec(&down_proj_abs, hidden, intermediate, &intermediate_rot);
let max_abs_diff = y_expected
.iter()
.zip(y_rotated.iter())
.map(|(a, b)| (a - b).abs())
.fold(0.0_f32, f32::max);
assert!(
max_abs_diff < 1e-3,
"MLP pair output should equal R · y_original: max_abs_diff={max_abs_diff}"
);
}
#[test]
fn absorption_for_module_input_side() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
for module in [
"q_proj",
"k_proj",
"v_proj",
"gate_proj",
"up_proj",
"in_proj_qkv",
"in_proj_z",
"in_proj_b",
"in_proj_a",
] {
assert_eq!(
plan.absorption_for_module(module),
Some(AbsorptionSide::InputSide),
"{module} should be InputSide"
);
}
}
#[test]
fn absorption_for_module_output_side() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
for module in ["o_proj", "down_proj", "out_proj"] {
assert_eq!(
plan.absorption_for_module(module),
Some(AbsorptionSide::OutputSide),
"{module} should be OutputSide"
);
}
}
#[test]
fn absorption_for_module_unknown() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
assert_eq!(plan.absorption_for_module("conv1d"), None);
assert_eq!(plan.absorption_for_module("norm"), None);
}
#[test]
fn qwen3_plan_covers_residual_stream_tensors() {
let plan = RotationPlan::qwen3_residual_stream_linear_layers();
assert_eq!(plan.rule_count(), 9);
let cases = [
(
"layers.0.self_attn.q_proj.weight",
AbsorptionSide::InputSide,
),
(
"layers.5.self_attn.k_proj.weight",
AbsorptionSide::InputSide,
),
(
"layers.5.self_attn.v_proj.weight",
AbsorptionSide::InputSide,
),
(
"layers.5.self_attn.o_proj.weight",
AbsorptionSide::OutputSide,
),
("layers.27.mlp.gate_proj.weight", AbsorptionSide::InputSide),
("layers.27.mlp.up_proj.weight", AbsorptionSide::InputSide),
("layers.27.mlp.down_proj.weight", AbsorptionSide::OutputSide),
("embed_tokens.weight", AbsorptionSide::InputSide),
("lm_head.weight", AbsorptionSide::InputSide),
];
for (name, expected_side) in cases {
let tr = plan
.for_tensor(name)
.unwrap_or_else(|| panic!("qwen3 plan missed tensor {name}"));
assert_eq!(tr.side, expected_side, "wrong side for {name}");
assert_eq!(tr.rotation_id, RotationId::ResidualStream);
}
}
#[test]
fn qwen3_plan_does_not_cover_gdn_tensors() {
let plan = RotationPlan::qwen3_residual_stream_linear_layers();
let gdn_names = [
"layers.0.linear_attn.in_proj_qkv.weight",
"layers.0.linear_attn.in_proj_z.weight",
"layers.0.linear_attn.in_proj_b.weight",
"layers.0.linear_attn.in_proj_a.weight",
"layers.0.linear_attn.out_proj.weight",
];
for name in gdn_names {
assert!(
plan.for_tensor(name).is_none(),
"GDN tensor {name} should not match qwen3 plan"
);
}
}
#[test]
fn qwen3_0_6b_required_tensor_names_count_and_spot_check() {
let cfg = QwenConfig::qwen3_embedding_0_6b();
let names = qwen3_required_tensor_names(&cfg);
assert_eq!(
names.len(),
310,
"expected 310 names for Qwen3-0.6B (28 layers × 11 + 2 global)"
);
assert!(names.contains(&"embed_tokens.weight".to_string()));
assert!(names.contains(&"norm.weight".to_string()));
assert!(names.contains(&"layers.0.self_attn.q_proj.weight".to_string()));
assert!(names.contains(&"layers.27.mlp.down_proj.weight".to_string()));
assert!(names.contains(&"layers.0.self_attn.q_norm.weight".to_string()));
assert!(names.contains(&"layers.0.self_attn.k_norm.weight".to_string()));
assert!(
!names.iter().any(|n| n.contains("linear_attn")),
"Qwen3 required names must not include GDN linear_attn tensors"
);
}
#[test]
fn qwen3_0_6b_validate_coverage_complete_without_lm_head() {
let plan = RotationPlan::qwen3_residual_stream_linear_layers();
let cfg = QwenConfig::qwen3_embedding_0_6b();
let names = qwen3_required_tensor_names(&cfg);
let report = plan.validate_coverage(names.iter().map(String::as_str));
assert!(
report.is_complete(),
"Qwen3-0.6B required names yield complete coverage: {report:?}"
);
assert_eq!(report.unmatched_required_rules.len(), 0);
assert_eq!(
report.unmatched_optional_rules,
vec!["lm_head.weight".to_string()],
"lm_head.weight should be the only unmatched-optional rule (tied embeddings)"
);
}
#[test]
fn f64_apply_matches_f32_apply_pattern() {
let plan = RotationPlan::qwen35_residual_stream_linear_layers();
let hidden = 64;
let r = RandomizedHadamard::new(11, hidden).unwrap();
let mut weight_f32: Vec<f32> = (0..hidden * hidden)
.map(|i| (i as f32 * 0.01).sin())
.collect();
let mut weight_f64: Vec<f64> = weight_f32.iter().map(|&v| f64::from(v)).collect();
let r1 = apply_tensor_rotation(
"model.layers.0.self_attn.q_proj.weight",
&mut weight_f32,
hidden,
hidden,
&plan,
&r,
)
.unwrap();
let r2 = apply_tensor_rotation_f64(
"model.layers.0.self_attn.q_proj.weight",
&mut weight_f64,
hidden,
hidden,
&plan,
&r,
)
.unwrap();
assert!(r1 && r2);
for (i, (a, b)) in weight_f32.iter().zip(weight_f64.iter()).enumerate() {
let delta = (f64::from(*a) - b).abs();
assert!(
delta < 1e-5,
"tensor[{i}]: f32={a} vs f64={b}, delta={delta}"
);
}
}
#[test]
fn r3_full_attention_scope_matches_qwen35_0_8b_layer_pattern() {
let cfg = Qwen35Config::qwen35_0_8b();
let spec = OnlineRotationSpec::r3_full_attention(&cfg, 7, 8).unwrap();
assert_eq!(spec.id, RotationId::AttentionOutputR3);
assert_eq!(spec.side, AbsorptionSide::InputSide);
assert_eq!(spec.block_size, 8);
assert_eq!(
spec.layer_scope,
Some(vec![3usize, 7, 11, 15, 19, 23]),
"R3 scope must be exactly the config's full-attention layers"
);
assert_eq!(spec.layer_scope.as_ref().unwrap().len(), 6);
for &idx in spec.layer_scope.as_ref().unwrap() {
assert!(
cfg.is_full_attention(idx),
"layer {idx} must be full-attention"
);
}
for idx in 0..cfg.num_hidden_layers {
if !spec.layer_scope.as_ref().unwrap().contains(&idx) {
assert!(
!cfg.is_full_attention(idx),
"layer {idx} is full-attention but missing from R3 scope"
);
}
}
}
#[test]
fn r3_rejects_block_size_not_dividing_num_attention_heads() {
let cfg = Qwen35Config::qwen35_0_8b();
assert!(OnlineRotationSpec::r3_full_attention(&cfg, 1, 16).is_err());
}
#[test]
fn r3_accepts_block_size_grouping_non_power_of_two_head_count() {
let cfg = crate::model::qwen35_config::Qwen35Config::qwen36_27b();
assert_eq!(cfg.num_attention_heads, 24);
let spec = OnlineRotationSpec::r3_full_attention(&cfg, 1, 8).unwrap();
assert_eq!(spec.block_size, 8);
assert!(!cfg.num_attention_heads.is_power_of_two());
assert!(OnlineRotationSpec::r3_full_attention(&cfg, 1, 24).is_err());
}
#[test]
fn r3_rejects_non_power_of_two_block_size() {
let cfg = Qwen35Config::qwen35_0_8b();
assert!(OnlineRotationSpec::r3_full_attention(&cfg, 1, 96).is_err());
}
#[test]
fn r3_rejects_zero_block_size() {
let cfg = Qwen35Config::qwen35_0_8b();
assert!(OnlineRotationSpec::r3_full_attention(&cfg, 1, 0).is_err());
}
#[test]
fn r4_dense_mlp_has_no_layer_scope_restriction() {
let cfg = Qwen35Config::qwen35_0_8b();
let spec = OnlineRotationSpec::r4_dense_mlp(&cfg, 9, 256).unwrap();
assert_eq!(spec.id, RotationId::MlpDownR4);
assert_eq!(spec.side, AbsorptionSide::InputSide);
assert_eq!(spec.layer_scope, None);
}
#[test]
fn r3_rejects_oversized_layer_scope() {
let oversized: Vec<usize> =
(0..(OnlineRotationSpec::MAX_LAYER_SCOPE_ENTRIES + 1)).collect();
let hand_built = OnlineRotationSpec {
id: RotationId::AttentionOutputR3,
side: AbsorptionSide::InputSide,
seed: 1,
block_size: 1,
layer_scope: Some(oversized),
};
let err = hand_built.validate(None).unwrap_err();
assert!(
format!("{err}").contains(&OnlineRotationSpec::MAX_LAYER_SCOPE_ENTRIES.to_string()),
"expected the layer_scope cap error naming the maximum, got: {err}"
);
}
#[test]
fn r4_dense_mlp_rejects_moe_config() {
let moe_cfg = Qwen35Config::qwen36_35b_a3b();
assert!(moe_cfg.is_moe(), "sanity: qwen36_35b_a3b is MoE");
let err = OnlineRotationSpec::r4_dense_mlp(&moe_cfg, 9, 256).unwrap_err();
assert!(format!("{err}").contains("MoE"), "got: {err}");
let hand_built = OnlineRotationSpec {
id: RotationId::MlpDownR4,
side: AbsorptionSide::InputSide,
seed: 9,
block_size: 256,
layer_scope: None,
};
let err = hand_built.validate(Some(&moe_cfg)).unwrap_err();
assert!(format!("{err}").contains("MoE"), "got: {err}");
let dense_cfg = Qwen35Config::qwen35_0_8b();
assert!(!dense_cfg.is_moe(), "sanity: qwen35_0_8b is dense");
assert!(OnlineRotationSpec::r4_dense_mlp(&dense_cfg, 9, 256).is_ok());
assert!(hand_built.validate(Some(&dense_cfg)).is_ok());
}
#[test]
fn r4_rejects_block_size_not_dividing_intermediate_size() {
let cfg = Qwen35Config::qwen35_0_8b();
assert_eq!(
cfg.intermediate_size % 1024,
512,
"sanity: 1024 does not divide 3584"
);
assert!(OnlineRotationSpec::r4_dense_mlp(&cfg, 1, 1024).is_err());
}
#[test]
fn r4_dense_mlp_rejects_num_blocks_above_block_hadamard_cap() {
let cfg = Qwen35Config::qwen36_27b();
assert!(!cfg.is_moe(), "sanity: qwen36_27b is dense");
assert_eq!(cfg.intermediate_size, 17408);
assert_eq!(
cfg.intermediate_size / 4,
4352,
"sanity: 17408 / 4 = 4352 blocks"
);
let err = OnlineRotationSpec::r4_dense_mlp(&cfg, 9, 4).unwrap_err();
assert!(
format!("{err}").contains("4352") && format!("{err}").contains("4096"),
"expected the num_blocks cap error naming both the block count \
and the cap, got: {err}"
);
let hand_built = OnlineRotationSpec {
id: RotationId::MlpDownR4,
side: AbsorptionSide::InputSide,
seed: 9,
block_size: 4,
layer_scope: None,
};
let err = hand_built.validate(Some(&cfg)).unwrap_err();
assert!(
format!("{err}").contains("4352") && format!("{err}").contains("4096"),
"expected validate() to independently reject the hand-built \
spec with the num_blocks cap error, got: {err}"
);
assert!(super::super::hadamard::BlockHadamard::new(9, cfg.intermediate_size, 4).is_err());
}
#[test]
fn r4_dense_mlp_accepts_block_size_within_num_blocks_cap() {
let cfg = Qwen35Config::qwen36_27b();
assert_eq!(
cfg.intermediate_size / 128,
136,
"sanity: 17408 / 128 = 136 blocks"
);
let spec = OnlineRotationSpec::r4_dense_mlp(&cfg, 9, 128).unwrap();
assert!(spec.validate(Some(&cfg)).is_ok());
let bh = super::super::hadamard::BlockHadamard::new(9, cfg.intermediate_size, 128)
.expect("136 blocks must be constructible — well under the 4096 cap");
assert_eq!(bh.num_blocks(), 136);
}
#[test]
fn r4_dense_mlp_rejects_dim_over_max_block_hadamard_len_even_within_block_cap() {
let mut cfg = Qwen35Config::qwen35_0_8b();
cfg.intermediate_size = 33_554_432;
assert_eq!(
cfg.intermediate_size / 8192,
4096,
"sanity: exactly at MAX_BLOCK_HADAMARD_BLOCKS, not over it"
);
let err = OnlineRotationSpec::r4_dense_mlp(&cfg, 9, 8192)
.expect_err("dim over MAX_BLOCK_HADAMARD_LEN must be rejected by the plan gate");
assert!(
format!("{err}").contains("MAX_BLOCK_HADAMARD_LEN"),
"expected the plan gate to reject on the length bound, not the \
block-count bound (which this pair sits exactly at), got: {err}"
);
assert!(
super::super::hadamard::BlockHadamard::new(9, cfg.intermediate_size, 8192).is_err()
);
let small_cfg = Qwen35Config::qwen35_0_8b();
assert!(OnlineRotationSpec::r4_dense_mlp(&small_cfg, 9, 128).is_ok());
}
#[test]
fn r4_accepts_all_named_design_doc_block_sizes() {
let cfg = Qwen35Config::qwen35_0_8b();
for &b in &[64usize, 128, 256] {
assert_eq!(
cfg.intermediate_size % b,
0,
"block size {b} must divide 3584"
);
assert!(OnlineRotationSpec::r4_dense_mlp(&cfg, 1, b).is_ok());
}
}
#[test]
fn rotation_id_online_transform_site_pins_every_variant() {
assert_eq!(RotationId::ResidualStream.online_transform_site(), None);
assert_eq!(
RotationId::AttentionOutputR3.online_transform_site(),
Some(OnlineTransformSite::AttentionOutputPreOProj)
);
assert_eq!(
RotationId::MlpDownR4.online_transform_site(),
Some(OnlineTransformSite::MlpPreDownProj)
);
assert_eq!(
OnlineTransformSite::AttentionOutputPreOProj.weight_tensor_suffix(),
"self_attn.o_proj.weight"
);
assert_eq!(
OnlineTransformSite::MlpPreDownProj.weight_tensor_suffix(),
"mlp.down_proj.weight"
);
}
#[test]
fn apply_tensor_rotation_refuses_r3_rotation_id() {
let plan = RotationPlan {
rules: vec![Rule {
pattern: "self_attn.o_proj.weight".to_string(),
rotation: TensorRotation {
side: AbsorptionSide::InputSide,
rotation_id: RotationId::AttentionOutputR3,
},
requirement: RuleRequirement::Required,
}],
};
let hidden = 64;
let r = RandomizedHadamard::new(1, hidden).unwrap();
let mut weight = vec![1.0_f32; hidden * hidden];
let err = apply_tensor_rotation(
"model.layers.0.self_attn.o_proj.weight",
&mut weight,
hidden,
hidden,
&plan,
&r,
)
.unwrap_err();
let msg = format!("{err}");
assert!(
msg.contains("contract-only"),
"expected a contract-only refusal, got: {msg}"
);
}
}