olmo2: stop generation at the end-of-sequence token

A 200-token run walked straight through <|endoftext|> and started a fresh
document mid-response, which is what a base model does when nothing stops it.
The runner now reads eos_token_id from config.json and stops before emitting
the sentinel rather than after, because a caller wants the text and not the
marker. The serving lane does the same.

Also verified while looking: 200 tokens stay finite with no NaN and no
collapse, so the f16 residual holds across sixteen post-norm layers, and the
same prompt produces bit-identical output across runs.

Error paths all explain themselves rather than failing obscurely. A Qwen config
is refused by name, a missing file by path, and a SmolLM-shaped head_dim 64 with
the reason it cannot work here.

Author: maincode-dave

crates/mainarch-cli/src/olmo_lane.rs

@@ -157,12 +157,17 @@ impl OlmoLane {
                     continue;
                 }
 
+                let eos = runner.eos_token();
                 let mut ttft_ms = 0.0f64;
                 let mut generated = 0usize;
                 for i in 0..req.max_new {
                     if i == 0 {
                         ttft_ms = started.elapsed().as_secs_f64() * 1e3;
                     }
+                    // Stop before emitting the sentinel, not after.
+                    if Some(next) == eos {
+                        break;
+                    }
                     let piece = tokenizer
                         .decode(&[next], false)
                         .unwrap_or_else(|_| String::new());

crates/mainarch-core/src/olmo2.rs

@@ -93,6 +93,8 @@ pub struct Olmo2Config {
     pub vocab_size: u64,
     pub rope_theta: u64,
     pub tie_word_embeddings: bool,
+    /// End-of-sequence token, when config.json declares one.
+    pub eos_token_id: Option<u64>,
 }
 
 impl Olmo2Config {
@@ -132,6 +134,10 @@ impl Olmo2Config {
                 .map(|v| v.as_bool("tie_word_embeddings"))
                 .transpose()?
                 .unwrap_or(false),
+            eos_token_id: obj
+                .get("eos_token_id")
+                .map(|v| v.as_u64("eos_token_id"))
+                .transpose()?,
         })
     }
 
@@ -378,6 +384,7 @@ pub fn synthetic_olmo2_config(num_hidden_layers: u64) -> Olmo2Config {
         vocab_size: 100_352,
         rope_theta: 500_000,
         tie_word_embeddings: false,
+        eos_token_id: Some(100_257),
     }
 }
 
@@ -1552,6 +1559,11 @@ impl Olmo2Runner {
         Ok(best)
     }
 
+    /// The configured end-of-sequence token, if the checkpoint declares one.
+    pub fn eos_token(&self) -> Option<u32> {
+        self.weights.config.eos_token_id.map(|e| e as u32)
+    }
+
     /// Greedy generation. The prompt is consumed one token at a time through the
     /// decode path, which is prefill done the slow honest way: no prefill GEMM
     /// kernel is needed because the decode loop already grows the KV cache.
@@ -1570,8 +1582,15 @@ impl Olmo2Runner {
             next = self.step(dev, t, pos)?;
             pos += 1;
         }
+        let eos = self.weights.config.eos_token_id.map(|e| e as u32);
         let mut out = Vec::with_capacity(max_new);
         for _ in 0..max_new {
+            // Stop *before* emitting EOS. A caller wants the text, not the
+            // sentinel, and a base model that runs past it starts a new
+            // document mid-response.
+            if Some(next) == eos {
+                break;
+            }
             out.push(next);
             if pos as usize >= self.max_seq as usize {
                 break;