kernels: fix uninitialized shared-memory read in M=1 decode GEMV

gemv_bf16_fused_kernel returned early on out-of-range columns
(`if (col >= N) return;`) BEFORE the cooperative load of x into shared
memory and the `__syncthreads()`. When N is not a multiple of GEMV_TILE_N
(128), the last column-block's out-of-range threads exited without loading
their slice of x_shared, so the in-range threads then read uninitialized
shared memory in the dot product — and __syncthreads with exited threads is
itself UB. Result: intermittent huge/garbage outputs (~1e33) that, after
the next RMSNorm, collapsed the whole forward pass to a degenerate logit
distribution (argmax → vocab_size-1, or NaN), derailing generation.

This hit every M=1 BF16 GEMV (n>=256) with n % 128 != 0 — i.e. gpt-oss
decode o_proj and the MoE projections (n=2880). q/k/v (4096) and lm_head
(201088) are 128-aligned and were unaffected, as is Qwen3 (hidden 4096),
which is why this manifested as intermittent gpt-oss-only decode failures.

Fix: all threads participate in the shared-memory load and reach the
barrier; the col>=N check moves to AFTER __syncthreads.

Verified on dash5 (TP=2): a prompt that reliably produced garbage ~70% of
runs now yields clean logits 16/16; the multi-turn Chinese chat that
collapsed mid-conversation completes coherently with 0 NaN warnings.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xserv-model/src/bin/xserv-chat.rs

@@ -328,6 +328,13 @@ fn main() {
     eprintln!("Ready (paged KV cache, tp={world}).");
     eprintln!("Commands: /exit, /quit, /clear\n");
 
+    // gpt-oss multi-turn history of (user, assistant-final) text. Harmony
+    // requires re-rendering the conversation each turn with prior analysis
+    // dropped, so the moe path re-prefills from this rather than reusing an
+    // incremental KV cache (which would accumulate CoT + <|return|> and collapse
+    // at longer context). Qwen3 ignores this and keeps the incremental cache.
+    let mut moe_history: Vec<(String, String)> = Vec::new();
+
     loop {
         let line = match read_line_edited("user> ") {
             Line::Eof => break,
@@ -341,10 +348,8 @@ fn main() {
         match input {
             "/exit" | "/quit" | "exit" | "quit" => break,
             "/clear" => {
-                if let Some(h) = &tp_handle { h.send(TpCommand::Free(SLOT)); h.wait(); }
+                reset_slot(&mut cache, &tp_handle);
-                cache.free_sequence(SLOT);
+                moe_history.clear();
-                if let Some(h) = &tp_handle { h.send(TpCommand::Register(SLOT)); h.wait(); }
-                cache.register_sequence(SLOT).expect("register chat slot");
                 eprintln!("history and KV cache cleared");
                 continue;
             }
@@ -355,21 +360,47 @@ fn main() {
             _ => {}
         }
 
-        let include_system = cache.seq_len(SLOT) == 0;
+        if is_moe {
-        let prompt = if is_moe {
+            // Harmony multi-turn: re-render the whole conversation (prior
-            build_turn_prompt_gpt_oss(
+            // analysis dropped) and re-prefill into a freshly cleared slot.
+            let prompt = build_conversation_gpt_oss(
                 opts.system_prompt.as_deref(),
-                include_system,
+                &moe_history,
                 input,
-            )
+            );
-        } else {
+            let prompt_tokens = tokenizer.encode(&prompt);
-            build_turn_prompt(
+            if prompt_tokens.is_empty() {
+                continue;
+            }
+            if prompt_tokens.len() >= max_seq_len {
+                eprintln!(
+                    "context full: conversation needs {} tokens >= max_seq_len {max_seq_len}; use /clear",
+                    prompt_tokens.len()
+                );
+                continue;
+            }
+            let max_new_tokens = opts.max_tokens.min(max_seq_len - prompt_tokens.len());
+            reset_slot(&mut cache, &tp_handle);
+            print!("assistant> ");
+            io::stdout().flush().unwrap();
+            let (_finish, answer) = generate_with_paged_cache(
+                &model, &mut cache, &tokenizer, &prompt_tokens, &opts.sampling,
+                max_new_tokens, use_color, &tp_handle, is_moe,
+            );
+            moe_history.push((input.to_string(), answer));
+            println!();
+            continue;
+        }
+
+        // Qwen3: incremental KV cache — only the new turn is prefilled and the
+        // assistant's tokens stay cached for the next turn.
+        let include_system = cache.seq_len(SLOT) == 0;
+        let prompt = build_turn_prompt(
             opts.system_prompt.as_deref(),
             include_system,
             input,
             opts.enable_thinking,
-            )
+        );
-        };
         let prompt_tokens = tokenizer.encode(&prompt);
         if prompt_tokens.is_empty() {
             continue;
@@ -388,7 +419,7 @@ fn main() {
 
         print!("assistant> ");
         io::stdout().flush().unwrap();
-        let finish = generate_with_paged_cache(
+        let (finish, _answer) = generate_with_paged_cache(
             &model,
             &mut cache,
             &tokenizer,
@@ -404,14 +435,21 @@ fn main() {
                 append_after_stop(&model, &mut cache, &tokenizer, max_seq_len, token_id, &tp_handle);
             }
             Finish::Length => {
-                let end_text = if is_moe { "<|end|>\n" } else { "<|im_end|>\n" };
+                append_text_to_cache(&model, &mut cache, &tokenizer, max_seq_len, "<|im_end|>\n", &tp_handle);
-                append_text_to_cache(&model, &mut cache, &tokenizer, max_seq_len, end_text, &tp_handle);
             }
         }
         println!();
     }
 }
 
+/// Free and re-register the single chat KV slot (clears all cached context).
+fn reset_slot(cache: &mut PagedKVCache, tp: &Option<TpHandle>) {
+    if let Some(h) = tp { h.send(TpCommand::Free(SLOT)); h.wait(); }
+    cache.free_sequence(SLOT);
+    if let Some(h) = tp { h.send(TpCommand::Register(SLOT)); h.wait(); }
+    cache.register_sequence(SLOT).expect("register chat slot");
+}
+
 fn parse_args() -> CliOptions {
     let args: Vec<String> = std::env::args().skip(1).collect();
     if args.is_empty() || args.iter().any(|a| a == "--help" || a == "-h") {
@@ -575,15 +613,30 @@ fn build_turn_prompt(
     prompt
 }
 
-fn build_turn_prompt_gpt_oss(
+/// Render the full gpt-oss harmony conversation for re-prefill. gpt-oss was
+/// trained on this exact system-message structure (identity / knowledge cutoff
+/// / current date / Reasoning level / channels — see the model's
+/// chat_template.jinja `build_system_message`). A hand-rolled substitute puts
+/// the model out of distribution and destabilizes channel selection.
+///
+/// Harmony multi-turn drops prior chain-of-thought: past assistant messages are
+/// rendered as completed `final` channels ending in `<|end|>` (not the
+/// `<|return|>` stop token). Keeping the analysis + `<|return|>` of every turn
+/// in context — as an incremental KV cache does — is out of distribution and
+/// makes the model collapse at longer context. "Reasoning: low" keeps the
+/// analysis channel short for an interactive chat.
+fn build_conversation_gpt_oss(
     system: Option<&str>,
-    include_system: bool,
+    history: &[(String, String)],
-    user_input: &str,
+    current_user: &str,
 ) -> String {
     let mut prompt = String::new();
-    if include_system {
     prompt.push_str("<|start|>system<|message|>");
-        prompt.push_str("You are a helpful assistant.\n\n# Valid channels: analysis, commentary, final. Channel must be included for every message.");
+    prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
+    prompt.push_str("Knowledge cutoff: 2024-06\n");
+    prompt.push_str(&format!("Current date: {}\n\n", today_ymd()));
+    prompt.push_str("Reasoning: low\n\n");
+    prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
     prompt.push_str("<|end|>");
     if let Some(sys) = system {
         if !sys.trim().is_empty() {
@@ -592,14 +645,39 @@ fn build_turn_prompt_gpt_oss(
             prompt.push_str("<|end|>");
         }
     }
+    for (user, assistant) in history {
+        prompt.push_str("<|start|>user<|message|>");
+        prompt.push_str(user);
+        prompt.push_str("<|end|>");
+        prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
+        prompt.push_str(assistant.trim());
+        prompt.push_str("<|end|>");
     }
     prompt.push_str("<|start|>user<|message|>");
-    prompt.push_str(user_input);
+    prompt.push_str(current_user);
     prompt.push_str("<|end|>");
     prompt.push_str("<|start|>assistant");
     prompt
 }
 
+/// Current UTC date as "YYYY-MM-DD" for the harmony system message. Rata Die
+/// civil-calendar conversion (same algorithm the server uses for strftime_now).
+fn today_ymd() -> String {
+    use std::time::{SystemTime, UNIX_EPOCH};
+    let secs = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
+    let z = (secs / 86400) as i64 + 719468;
+    let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
+    let doe = z - era * 146097;
+    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
+    let y = yoe + era * 400;
+    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
+    let mp = (5 * doy + 2) / 153;
+    let d = doy - (153 * mp + 2) / 5 + 1;
+    let m = if mp < 10 { mp + 3 } else { mp - 9 };
+    let y = if m <= 2 { y + 1 } else { y };
+    format!("{y:04}-{m:02}-{d:02}")
+}
+
 fn generate_with_paged_cache(
     model: &ChatModel,
     cache: &mut PagedKVCache,
@@ -610,7 +688,7 @@ fn generate_with_paged_cache(
     use_color: bool,
     tp: &Option<TpHandle>,
     is_moe: bool,
-) -> Finish {
+) -> (Finish, String) {
     let harmony_end_id = if is_moe { tokenizer.special_token_id("<|end|>") } else { None };
     let harmony_channel_id = if is_moe { tokenizer.special_token_id("<|channel|>") } else { None };
     let harmony_message_id = if is_moe { tokenizer.special_token_id("<|message|>") } else { None };
@@ -627,9 +705,14 @@ fn generate_with_paged_cache(
     enum HarmonyState { Normal, ReadingChannel, InAnalysis, InFinal }
     let mut hstate = if is_moe { HarmonyState::InFinal } else { HarmonyState::Normal };
 
+    // Off by default. A repetition penalty over a harmony stream penalizes the
+    // control tokens (<|channel|>, <|message|>, <|start|>) that MUST repeat to
+    // open the final channel — so a non-1.0 default makes gpt-oss stop right
+    // after the analysis block, before emitting any answer. Opt in via the env
+    // var if you want it for plain (non-harmony) generation.
     let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY").ok()
         .and_then(|s| s.parse().ok())
-        .unwrap_or(if is_moe { 1.3 } else { 1.0 });
+        .unwrap_or(1.0);
     let rep_window: usize = std::env::var("XSERV_REP_WINDOW").ok()
         .and_then(|s| s.parse().ok())
         .unwrap_or(512);
@@ -650,6 +733,11 @@ fn generate_with_paged_cache(
     let mut next = pick(&logits, sampling, &history);
     let mut decode_buffer = Vec::new();
     let mut in_thinking = false;
+    // Visible answer tokens, returned for multi-turn history. For moe this is
+    // the final-channel content only (analysis is suppressed/gray); for Qwen3
+    // it is everything printed. The caller decodes these into the assistant
+    // message it re-renders into the next prompt.
+    let mut answer_ids: Vec<u32> = Vec::new();
 
     for _ in 0..max_tokens {
         let position = cache.seq_len(SLOT);
@@ -663,7 +751,7 @@ fn generate_with_paged_cache(
                 use_color,
             );
             io::stdout().flush().unwrap();
-            return Finish::Stop { token_id: next };
+            return (Finish::Stop { token_id: next }, tokenizer.decode(&answer_ids));
         }
         if harmony_end_id == Some(next) {
             // <|end|> closes current segment; if in final channel, we're done
@@ -674,7 +762,7 @@ fn generate_with_paged_cache(
             );
             if hstate == HarmonyState::InFinal {
                 io::stdout().flush().unwrap();
-                return Finish::Stop { token_id: next };
+                return (Finish::Stop { token_id: next }, tokenizer.decode(&answer_ids));
             }
             hstate = HarmonyState::Normal;
             next = pick(&logits, sampling, &history);
@@ -724,7 +812,16 @@ fn generate_with_paged_cache(
             next = pick(&logits, sampling, &history);
             continue;
         }
+        if is_moe && hstate != HarmonyState::InFinal {
+            // Between harmony messages (after a channel's <|end|>, before the
+            // next <|channel|>): the model emits a role header like "assistant".
+            // That's structural, not user-visible content — suppress it. Only
+            // for moe/harmony; non-moe (Qwen3) stays in Normal and prints here.
+            next = pick(&logits, sampling, &history);
+            continue;
+        }
 
+        answer_ids.push(next);
         print_generated_token(
             tokenizer,
             next,
@@ -742,7 +839,7 @@ fn generate_with_paged_cache(
         use_color,
     );
     io::stdout().flush().unwrap();
-    Finish::Length
+    (Finish::Length, tokenizer.decode(&answer_ids))
 }
 
 fn append_after_stop(

crates/xserv-server/src/api.rs

@@ -198,8 +198,17 @@ fn build_prompt_hardcoded(messages: &[Message], model_type: &str) -> String {
 
 fn build_prompt_gpt_oss(messages: &[Message]) -> String {
     let mut prompt = String::new();
+    // Canonical harmony system message (mirrors the model's chat_template.jinja
+    // build_system_message macro). A hand-rolled substitute puts gpt-oss out of
+    // distribution and destabilizes channel selection. This hardcoded builder is
+    // only a fallback for gpt-oss models that ship no Jinja template; the
+    // gpt-oss-20b release does ship one, so the template path is normally used.
     prompt.push_str("<|start|>system<|message|>");
-    prompt.push_str("You are a helpful assistant.\n\n# Valid channels: analysis, commentary, final. Channel must be included for every message.");
+    prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
+    prompt.push_str("Knowledge cutoff: 2024-06\n");
+    prompt.push_str(&format!("Current date: {}\n\n", strftime_now("%Y-%m-%d".to_string())));
+    prompt.push_str("Reasoning: low\n\n");
+    prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
     prompt.push_str("<|end|>");
     let dev_instructions: String = messages
         .iter()

csrc/attention/flash_attention.cu

@@ -306,6 +306,14 @@ __global__ void flash_attention_sinks_bf16_kernel(
                 row_max = fmaxf(row_max, s);
             }
 
+            // A fully-masked KV tile (every key causal- or window-masked) has
+            // row_max == -INFINITY. Folding it in computes expf(-inf - (-inf))
+            // = NaN, and a later valid tile's 0*NaN correction then poisons the
+            // whole row. This happens for sliding-window layers whenever a
+            // query's window starts past an early tile (the causal `continue`
+            // above only skips fully-future tiles, not out-of-window ones).
+            // A masked tile contributes nothing to the softmax — skip it.
+            if (row_max != -INFINITY) {
                 float m_new = fmaxf(m_val, row_max);
                 float psum = 0.0f;
                 for (int c = 0; c < kv_tile_cols; c++) {
@@ -316,6 +324,9 @@ __global__ void flash_attention_sinks_bf16_kernel(
                 l_val = correction * l_val + psum;
                 for (int d = 0; d < head_dim; d++) O_acc[d] *= correction;
                 m_val = m_new;
+            } else {
+                for (int c = 0; c < kv_tile_cols; c++) P[c] = 0.0f;
+            }
         }
 
         __syncthreads();

csrc/gemm/gemv.cu

@@ -28,18 +28,25 @@ __global__ void gemv_bf16_fused_kernel(
     const int t = threadIdx.x;
     const int col = block_n * GEMV_TILE_N + t;
 
-    if (col >= N) return;
-
     const int k_start = block_k * GEMV_TILE_K;
     const int k_end = min(k_start + GEMV_TILE_K, K);
     const int k_len = k_end - k_start;
 
+    // Cooperative load of x into shared memory uses ALL threads in the block
+    // (indexed by t, independent of col). Threads whose column is out of range
+    // must still help load and reach the barrier — returning early here would
+    // leave part of x_shared uninitialized AND make __syncthreads divergent
+    // (UB). So the col>=N check happens only AFTER the load + barrier. This bug
+    // produced intermittent huge/garbage outputs whenever N % GEMV_TILE_N != 0
+    // (e.g. gpt-oss decode o_proj with N=2880), collapsing the forward pass.
     __shared__ float x_shared[GEMV_TILE_K];
     for (int i = t; i < k_len; i += GEMV_BLOCK) {
         x_shared[i] = __bfloat162float(x[k_start + i]);
     }
     __syncthreads();
 
+    if (col >= N) return;
+
     float sum = 0.0f;
     for (int ki = 0; ki < k_len; ki++) {
         sum += x_shared[ki] * __bfloat162float(W[(long long)(k_start + ki) * N + col]);