style: cargo fmt --workspace (whitespace/wrapping only, no semantic change)

Whole-workspace rustfmt pass picked up while iterating on Mamba GPU
backward work. Verified formatting-only via diff sampling; no logic
changed.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
This commit is contained in:
osobh
2026-08-10 07:09:36 -07:00
co-authored by Claude Sonnet 5
parent ad6405663f
commit 4aaa36a57a
305 changed files with 25537 additions and 18337 deletions
+48 -22
View File
@@ -101,7 +101,14 @@ impl MedusaHead {
.collect();
let b2 = vec![0.0_f32; vocab_size];
Self { w1, b1, w2, b2, hidden_dim, vocab_size }
Self {
w1,
b1,
w2,
b2,
hidden_dim,
vocab_size,
}
}
/// Run the head forward pass.
@@ -123,8 +130,7 @@ impl MedusaHead {
// Layer 1: act = SiLU(W1 @ hidden + b1)
let mut act = vec![0.0_f32; h];
for i in 0..h {
let dot: f32 = (0..h).map(|j| self.w1[i * h + j] * hidden[j]).sum::<f32>()
+ self.b1[i];
let dot: f32 = (0..h).map(|j| self.w1[i * h + j] * hidden[j]).sum::<f32>() + self.b1[i];
// SiLU(x) = x * sigmoid(x) = x / (1 + e^{-x})
act[i] = dot * (1.0 / (1.0 + (-dot).exp()));
}
@@ -133,8 +139,7 @@ impl MedusaHead {
let v = self.vocab_size;
let mut logits = vec![0.0_f32; v];
for i in 0..v {
logits[i] =
(0..h).map(|j| self.w2[i * h + j] * act[j]).sum::<f32>() + self.b2[i];
logits[i] = (0..h).map(|j| self.w2[i * h + j] * act[j]).sum::<f32>() + self.b2[i];
}
logits
}
@@ -255,7 +260,9 @@ impl MedusaHeads {
config.hidden_dim,
config.vocab_size,
// Distinct seed per head so weights differ.
(i as u64).wrapping_mul(0xdead_beef_cafe_babe).wrapping_add(42),
(i as u64)
.wrapping_mul(0xdead_beef_cafe_babe)
.wrapping_add(42),
)
})
.collect();
@@ -392,16 +399,17 @@ impl MedusaHeads {
let logits = head.forward(hidden);
// Numerically stable CE: loss = log_sum_exp(logits) - logits[target]
let max_l = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let log_sum =
max_l + logits.iter().map(|l| (l - max_l).exp()).sum::<f32>().ln();
let log_sum = max_l + logits.iter().map(|l| (l - max_l).exp()).sum::<f32>().ln();
-(logits[target as usize] - log_sum)
})
.collect();
let total_loss =
per_head_loss.iter().sum::<f32>() / per_head_loss.len() as f32;
let total_loss = per_head_loss.iter().sum::<f32>() / per_head_loss.len() as f32;
MedusaLossResult { per_head_loss, total_loss }
MedusaLossResult {
per_head_loss,
total_loss,
}
}
}
@@ -453,7 +461,10 @@ mod tests {
let head = MedusaHead::new_random(cfg.hidden_dim, cfg.vocab_size, 2);
let hidden = ones_hidden(cfg.hidden_dim);
let logits = head.forward(&hidden);
assert!(logits.iter().all(|v| v.is_finite()), "all logits must be finite");
assert!(
logits.iter().all(|v| v.is_finite()),
"all logits must be finite"
);
}
// ── MedusaHead::top_k ────────────────────────────────────────────────────
@@ -527,10 +538,7 @@ mod tests {
let logits = head.forward(&vec![0.0; h]);
// All activations are SiLU(0) = 0, so all logits should be 0.
for l in &logits {
assert!(
l.abs() < 1e-6,
"SiLU(0) path: expected logit ≈ 0, got {l}"
);
assert!(l.abs() < 1e-6, "SiLU(0) path: expected logit ≈ 0, got {l}");
}
}
@@ -541,14 +549,24 @@ mod tests {
let v = 2;
// Identity W1 with positive bias forces positive pre-activation.
let w1 = vec![1.0_f32, 0.0, 0.0, 1.0]; // 2×2 identity
let b1 = vec![1.0_f32, 1.0]; // shift pre-activation up
let b1 = vec![1.0_f32, 1.0]; // shift pre-activation up
let w2 = vec![1.0_f32; v * h];
let b2 = vec![0.0_f32; v];
let head = MedusaHead { w1, b1, w2, b2, hidden_dim: h, vocab_size: v };
let head = MedusaHead {
w1,
b1,
w2,
b2,
hidden_dim: h,
vocab_size: v,
};
let logits = head.forward(&vec![0.0; h]);
// Pre-activation = 1.0; SiLU(1.0) = 1 / (1 + e^{-1}) ≈ 0.731 > 0.
for l in &logits {
assert!(*l > 0.0, "SiLU of positive input should produce positive output, got {l}");
assert!(
*l > 0.0,
"SiLU of positive input should produce positive output, got {l}"
);
}
}
@@ -692,7 +710,10 @@ mod tests {
let base = tree.paths[0][0];
let result = heads.verify(&tree, base, &|_| base);
// num_accepted == accepted_tokens.len() - 1
assert_eq!(result.num_accepted, result.accepted_tokens.len().saturating_sub(1));
assert_eq!(
result.num_accepted,
result.accepted_tokens.len().saturating_sub(1)
);
}
// ── training_loss ─────────────────────────────────────────────────────────
@@ -715,7 +736,10 @@ mod tests {
let targets = vec![0_u32; cfg.num_heads];
let result = heads.training_loss(&hidden, &targets);
for (i, &loss) in result.per_head_loss.iter().enumerate() {
assert!(loss >= 0.0, "head {i}: CE loss must be non-negative, got {loss}");
assert!(
loss >= 0.0,
"head {i}: CE loss must be non-negative, got {loss}"
);
}
}
@@ -751,7 +775,9 @@ mod tests {
let head = MedusaHead {
w1: {
let mut m = vec![0.0_f32; h * h];
for i in 0..h { m[i * h + i] = 1.0; } // identity
for i in 0..h {
m[i * h + i] = 1.0;
} // identity
m
},
b1: vec![0.0; h],