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:
@@ -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],
|
||||
|
||||
Reference in New Issue
Block a user