fix(Qwen3/GLM5.2): exclude frontend-undecodable padded vocab rows

openinfer-glm52/src/bookend.rs

@@ -9,6 +9,7 @@
 //! their row count and width by construction.
 
 use anyhow::Result;
+use anyhow::ensure;
 use cudarc::driver::CudaSlice;
 use openinfer_kernels::ops::embedding_rows_into;
 use openinfer_kernels::ops::gemm_strided_batched_bf16;
@@ -19,6 +20,7 @@ use openinfer_kernels::tensor::DeviceVec;
 
 use crate::config::GLM52_HIDDEN;
 use crate::config::GLM52_RMS_EPS;
+use crate::config::GLM52_SELECTION_VOCAB;
 use crate::config::GLM52_VOCAB;
 use crate::rows::Rows;
 
@@ -55,23 +57,34 @@ pub(crate) fn glm52_final_norm_into(
     )
 }
 
-/// lm_head projection over the buffers' `tokens()` rows: `lm_head @ normed ->
-/// [T, lm_head.rows]` compact logits. EP8 passes the full-vocabulary head;
-/// attention-TP decode passes this rank's contiguous vocabulary shard. One
+/// lm_head projection over the buffers' `tokens()` rows. EP passes the
+/// full checkpoint head but emits only the tokenizer-selectable prefix;
+/// attention-TP decode passes this rank's already-trimmed contiguous shard. One
 /// cuBLAS GEMM puts tokens on the n dimension, so the col-major
-/// `[lm_head.rows, T]` output is the compact row-major layout argmax consumes.
+/// `[logit_rows, T]` output is the compact row-major layout argmax consumes.
 pub(crate) fn glm52_lm_head_into(
     ctx: &DeviceContext,
     normed: &Rows<GLM52_HIDDEN>,
     lm_head: &DeviceMatrix,
-    out: &mut Rows<GLM52_VOCAB>,
-) -> Result<()> {
+    out: &mut Rows<GLM52_SELECTION_VOCAB>,
+) -> Result<usize> {
     let tokens = out.tokens();
+    let logit_rows = if lm_head.rows == GLM52_VOCAB {
+        GLM52_SELECTION_VOCAB
+    } else {
+        lm_head.rows
+    };
+    ensure!(
+        logit_rows <= GLM52_SELECTION_VOCAB
+            && (lm_head.rows == GLM52_VOCAB || GLM52_SELECTION_VOCAB.is_multiple_of(logit_rows)),
+        "GLM5.2 lm_head rows {} are neither the checkpoint head nor a selectable-vocab shard",
+        lm_head.rows
+    );
     gemm_strided_batched_bf16(
         ctx,
         true,
         false,
-        lm_head.rows,
+        logit_rows,
         tokens,
         GLM52_HIDDEN,
         &lm_head.data,
@@ -81,8 +94,9 @@ pub(crate) fn glm52_lm_head_into(
         GLM52_HIDDEN,
         0,
         out.data_mut(),
-        lm_head.rows,
+        logit_rows,
         0,
         1,
-    )
+    )?;
+    Ok(logit_rows)
 }

openinfer-glm52/src/config.rs

@@ -1,13 +1,22 @@
 //! GLM5.2 constants and config probing.
 
+use std::collections::HashSet;
+use std::path::Path;
+
 use anyhow::Context;
 use anyhow::Result;
 use anyhow::bail;
 use anyhow::ensure;
 use serde_json::Value;
 
 pub const GLM52_HIDDEN: usize = 6144;
+/// Physical checkpoint width. The embedding keeps every row because DSpark's
+/// internal mask token is the first row beyond the frontend token space.
 pub const GLM52_VOCAB: usize = 154_880;
+/// Dense tokenizer prefix that may be emitted by target or draft sampling.
+/// The official tokenizer defines ids 0..=154855; checkpoint rows 154856..154879
+/// are not frontend-decodable.
+pub(crate) const GLM52_SELECTION_VOCAB: usize = 154_856;
 pub const GLM52_LAYERS: usize = 78;
 pub const GLM52_DENSE_LAYERS: usize = 3;
 pub const GLM52_MOE_LAYERS: usize = GLM52_LAYERS - GLM52_DENSE_LAYERS;
@@ -66,6 +75,76 @@ pub(crate) fn glm52_layer_has_full_indexer(layer: usize) -> bool {
         .is_multiple_of(GLM52_INDEX_TOPK_FREQ)
 }
 
+/// Width of the frontend-decodable id space. This mirrors the Qwen tokenizer
+/// merge: `tokenizer.json` model vocab + added tokens, then
+/// `tokenizer_config.json`'s added-token decoder. The model uses a compile-time
+/// output width for captured graphs, so startup validates the supplied
+/// tokenizer against [`GLM52_SELECTION_VOCAB`] instead of silently accepting
+/// a different token space.
+pub(crate) fn tokenizer_effective_vocab(model_path: &Path) -> Result<usize> {
+    let path = model_path.join("tokenizer.json");
+    let content = std::fs::read_to_string(&path)
+        .map_err(|err| anyhow::anyhow!("read {}: {err}", path.display()))?;
+    let tokenizer: Value = serde_json::from_str(&content)
+        .map_err(|err| anyhow::anyhow!("parse {}: {err}", path.display()))?;
+
+    let config_path = model_path.join("tokenizer_config.json");
+    let tokenizer_config = std::fs::read_to_string(&config_path).ok().and_then(|text| {
+        match serde_json::from_str::<Value>(&text) {
+            Ok(config) => Some(config),
+            Err(err) => {
+                log::warn!(
+                    "cannot parse {}: {err}; skipping its added tokens like the frontend does",
+                    config_path.display()
+                );
+                None
+            }
+        }
+    });
+    tokenizer_effective_vocab_from_json(&tokenizer, tokenizer_config.as_ref())
+}
+
+fn tokenizer_effective_vocab_from_json(
+    tokenizer: &Value,
+    tokenizer_config: Option<&Value>,
+) -> Result<usize> {
+    let vocab = tokenizer
+        .pointer("/model/vocab")
+        .and_then(Value::as_object)
+        .ok_or_else(|| anyhow::anyhow!("tokenizer.json model.vocab must be an object"))?;
+    ensure!(!vocab.is_empty(), "tokenizer.json model.vocab is empty");
+    let mut ids = HashSet::with_capacity(vocab.len());
+    for (token, id) in vocab {
+        let id = id
+            .as_u64()
+            .with_context(|| format!("tokenizer model.vocab id for `{token}` is not unsigned"))?;
+        ids.insert(u32::try_from(id).context("tokenizer model.vocab id exceeds u32")?);
+    }
+    if let Some(added) = tokenizer.get("added_tokens").and_then(Value::as_array) {
+        for token in added {
+            let Some(id) = token.get("id").and_then(Value::as_u64) else {
+                continue;
+            };
+            ids.insert(u32::try_from(id).context("tokenizer added-token id exceeds u32")?);
+        }
+    }
+    if let Some(decoder) = tokenizer_config
+        .and_then(|config| config.get("added_tokens_decoder"))
+        .and_then(Value::as_object)
+    {
+        ids.extend(decoder.keys().filter_map(|key| key.parse::<u32>().ok()));
+    }
+
+    let width = ids.len();
+    let max_id = *ids.iter().max().expect("vocab checked non-empty") as usize;
+    ensure!(
+        max_id + 1 == width,
+        "tokenizer id space is not dense (max id {max_id}, {width} distinct ids); a row-range \
+         selection bound cannot mask holes"
+    );
+    Ok(width)
+}
+
 pub fn probe_config_json(json: &Value) -> Result<()> {
     let model_type = string_field(json, "model_type")?;
     if model_type != "glm_moe_dsa" {
@@ -302,3 +381,35 @@ fn ensure_float_close(actual: f64, expected: f64, tolerance: f64, label: &str) -
     );
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::tokenizer_effective_vocab_from_json;
+
+    #[test]
+    fn effective_vocab_merges_all_frontend_id_sources() {
+        let tokenizer = serde_json::json!({
+            "model": { "vocab": { "a": 0, "b": 1, "c": 2 } },
+            "added_tokens": [{ "id": 3, "content": "<x>" }]
+        });
+        let config = serde_json::json!({
+            "added_tokens_decoder": {
+                "4": { "content": "<y>" },
+                "5": { "content": "<z>" },
+                "bad": { "content": "<ignored>" }
+            }
+        });
+        assert_eq!(
+            tokenizer_effective_vocab_from_json(&tokenizer, Some(&config)).unwrap(),
+            6
+        );
+    }
+
+    #[test]
+    fn effective_vocab_rejects_sparse_id_spaces() {
+        let tokenizer = serde_json::json!({
+            "model": { "vocab": { "a": 0, "b": 2 } }
+        });
+        assert!(tokenizer_effective_vocab_from_json(&tokenizer, None).is_err());
+    }
+}

openinfer-glm52/src/dspark.rs

@@ -32,6 +32,7 @@ use openinfer_core::weight_loader::deserialize_shards;
 use openinfer_core::weight_loader::load_shard_info;
 use openinfer_core::weight_loader::load_tensor_1d;
 use openinfer_core::weight_loader::load_tensor_2d;
+use openinfer_core::weight_loader::load_tensor_2d_row_shard;
 use openinfer_core::weight_loader::mmap_shards;
 use openinfer_core::weight_loader::precompute_rope;
 use openinfer_kernels::ops::add_batch_into;
@@ -52,6 +53,7 @@ use openinfer_kernels::tensor::DeviceVec;
 use openinfer_kernels::tensor::HiddenStates;
 
 use crate::config::GLM52_HIDDEN;
+use crate::config::GLM52_SELECTION_VOCAB;
 use crate::config::GLM52_VOCAB;
 use crate::model::GLM52_MAX_BATCH_PER_RANK;
 
@@ -110,7 +112,7 @@ const DSPARK_HEAD_DIM: usize = 64;
 const DSPARK_QKV_DIM: usize = DSPARK_HEADS * DSPARK_HEAD_DIM;
 const DSPARK_INTER: usize = 12_288;
 const DSPARK_MARKOV_RANK: usize = 256;
-const DSPARK_MASK_TOKEN: u32 = 154_856;
+const DSPARK_MASK_TOKEN: u32 = GLM52_SELECTION_VOCAB as u32;
 const DSPARK_ROPE_THETA: f32 = 8_000_000.0;
 const DSPARK_RMS_EPS: f32 = 1.0e-5;
 
@@ -135,7 +137,8 @@ pub(crate) struct Glm52DsparkModel {
     hidden_norm: DeviceVec,
     /// Context projection `[6144, 30720]`.
     fc: DeviceMatrix,
-    /// Markov head: `bias(prev) = w2 @ w1[prev]`, both `[154880, 256]`.
+    /// Markov head over selectable tokens:
+    /// `bias(prev) = w2 @ w1[prev]`, both `[154856, 256]`.
     markov_w1: DeviceMatrix,
     markov_w2: DeviceMatrix,
     cos_cache: DeviceVec,
@@ -315,10 +318,34 @@ impl Glm52DsparkModel {
 
         let fc = load_tensor_2d(ctx, &shards, &weight_map, "fc.weight")?;
         ensure_matrix(&fc, "fc", GLM52_HIDDEN, GLM52_DSPARK_CONTEXT_DIM)?;
-        let markov_w1 = load_tensor_2d(ctx, &shards, &weight_map, "markov_head.markov_w1.weight")?;
-        let markov_w2 = load_tensor_2d(ctx, &shards, &weight_map, "markov_head.markov_w2.weight")?;
-        ensure_matrix(&markov_w1, "markov_w1", GLM52_VOCAB, DSPARK_MARKOV_RANK)?;
-        ensure_matrix(&markov_w2, "markov_w2", GLM52_VOCAB, DSPARK_MARKOV_RANK)?;
+        let markov_w1 = load_tensor_2d_row_shard(
+            ctx,
+            &shards,
+            &weight_map,
+            "markov_head.markov_w1.weight",
+            0,
+            GLM52_SELECTION_VOCAB,
+        )?;
+        let markov_w2 = load_tensor_2d_row_shard(
+            ctx,
+            &shards,
+            &weight_map,
+            "markov_head.markov_w2.weight",
+            0,
+            GLM52_SELECTION_VOCAB,
+        )?;
+        ensure_matrix(
+            &markov_w1,
+            "markov_w1",
+            GLM52_SELECTION_VOCAB,
+            DSPARK_MARKOV_RANK,
+        )?;
+        ensure_matrix(
+            &markov_w2,
+            "markov_w2",
+            GLM52_SELECTION_VOCAB,
+            DSPARK_MARKOV_RANK,
+        )?;
 
         // embed_tokens / lm_head / confidence_head are intentionally not
         // loaded: the first two are byte-identical to the target's, the
@@ -697,7 +724,14 @@ impl Glm52DsparkModel {
                                 DSPARK_RMS_EPS,
                                 logits_normed,
                             );
-                            gemm_into_checked(ctx, lm_head, logits_normed, logits)?;
+                            gemm_rows_into_checked(
+                                ctx,
+                                lm_head,
+                                0,
+                                GLM52_SELECTION_VOCAB,
+                                logits_normed,
+                                logits,
+                            )?;
                         });
                     }
                 }
@@ -717,7 +751,14 @@ impl Glm52DsparkModel {
                     layer_tail!(l);
                 }
                 rms_norm_batch_into(ctx, hidden, &self.norm, DSPARK_RMS_EPS, logits_normed);
-                gemm_into_checked(ctx, lm_head, logits_normed, logits)?;
+                gemm_rows_into_checked(
+                    ctx,
+                    lm_head,
+                    0,
+                    GLM52_SELECTION_VOCAB,
+                    logits_normed,
+                    logits,
+                )?;
             }
         }
         if fwd_on {
@@ -929,7 +970,8 @@ impl Glm52DsparkScratch {
         // preallocated so a draft round never touches the allocator (and the
         // VRAM probe's ledger charged exactly this).
         let tail_capacity = cache_len;
-        let partials = markov_step_argmax_partials_len(GLM52_MAX_BATCH_PER_RANK, GLM52_VOCAB);
+        let partials =
+            markov_step_argmax_partials_len(GLM52_MAX_BATCH_PER_RANK, GLM52_SELECTION_VOCAB);
         Ok(Self {
             block_token_ids_h: vec![DSPARK_MASK_TOKEN; max_rows],
             token_ids_d: ctx.stream.alloc_zeros(max_rows)?,
@@ -943,12 +985,12 @@ impl Glm52DsparkScratch {
             up_out: HiddenStates::zeros(ctx, DSPARK_INTER, max_rows)?,
             act_out: HiddenStates::zeros(ctx, DSPARK_INTER, max_rows)?,
             logits_normed: HiddenStates::zeros(ctx, GLM52_HIDDEN, max_rows)?,
-            logits: HiddenStates::zeros(ctx, GLM52_VOCAB, max_rows)?,
+            logits: HiddenStates::zeros(ctx, GLM52_SELECTION_VOCAB, max_rows)?,
             tail_input: HiddenStates::zeros(ctx, GLM52_HIDDEN, tail_capacity)?,
             k_tail: HiddenStates::zeros(ctx, DSPARK_QKV_DIM, tail_capacity)?,
             v_tail: HiddenStates::zeros(ctx, DSPARK_QKV_DIM, tail_capacity)?,
             w1emb: HiddenStates::zeros(ctx, DSPARK_MARKOV_RANK, GLM52_MAX_BATCH_PER_RANK)?,
-            bias: HiddenStates::zeros(ctx, GLM52_VOCAB, GLM52_MAX_BATCH_PER_RANK)?,
+            bias: HiddenStates::zeros(ctx, GLM52_SELECTION_VOCAB, GLM52_MAX_BATCH_PER_RANK)?,
             partial_values: ctx.stream.alloc_zeros(partials)?,
             partial_indices: ctx.stream.alloc_zeros(partials)?,
             prev_tokens: ctx.stream.alloc_zeros(GLM52_MAX_BATCH_PER_RANK)?,

openinfer-glm52/src/dspark_test_support.rs

@@ -43,8 +43,8 @@ impl Glm52DsparkModel {
             norm: DeviceVec::zeros(ctx, GLM52_HIDDEN)?,
             hidden_norm: DeviceVec::zeros(ctx, GLM52_HIDDEN)?,
             fc: mat(GLM52_HIDDEN, GLM52_DSPARK_CONTEXT_DIM)?,
-            markov_w1: mat(GLM52_VOCAB, DSPARK_MARKOV_RANK)?,
-            markov_w2: mat(GLM52_VOCAB, DSPARK_MARKOV_RANK)?,
+            markov_w1: mat(GLM52_SELECTION_VOCAB, DSPARK_MARKOV_RANK)?,
+            markov_w2: mat(GLM52_SELECTION_VOCAB, DSPARK_MARKOV_RANK)?,
             cos_cache,
             sin_cache,
             cache_len,

openinfer-glm52/src/lib.rs

@@ -72,6 +72,8 @@ use weights::Glm52RankLoadBundle;
 use weights::Glm52WeightManifest;
 
 use crate::config::GLM52_MAX_CONTEXT;
+use crate::config::GLM52_SELECTION_VOCAB;
+use crate::config::tokenizer_effective_vocab;
 use crate::model::GLM52_MODEL_LEN_ALIGN;
 use crate::model::glm52_arena_bytes;
 use crate::model::glm52_pool_blocks;
@@ -1114,6 +1116,12 @@ fn validate_startup(
     let json: serde_json::Value = serde_json::from_str(&content)
         .map_err(|err| anyhow::anyhow!("parse {}: {err}", config_path.display()))?;
     probe_config_json(&json)?;
+    let tokenizer_vocab = tokenizer_effective_vocab(model_path)?;
+    ensure!(
+        tokenizer_vocab == GLM52_SELECTION_VOCAB,
+        "GLM5.2 tokenizer defines {tokenizer_vocab} selectable ids, this build expects \
+         {GLM52_SELECTION_VOCAB} for its captured output buffers"
+    );
 
     let expected_devices = moe_topo.device_count();
     let remote_ranks: usize = options.rank_hosts.iter().map(|host| host.ranks).sum();

openinfer-glm52/src/model/launch_ahead.rs

@@ -10,7 +10,7 @@ use openinfer_kernels::ops::glm52_decode_feed_launch;
 use openinfer_kernels::tensor::DeviceContext;
 
 use super::GLM52_MAX_BATCH_PER_RANK;
-use super::GLM52_VOCAB;
+use super::GLM52_SELECTION_VOCAB;
 use super::Glm52RankModel;
 use super::Glm52StepShape;
 
@@ -156,8 +156,8 @@ impl Glm52RankModel {
                 "GLM5.2 slot {slot} greedy argmax found no finite logit (top = {top_value})"
             );
             ensure!(
-                (0..GLM52_VOCAB as i32).contains(&top_index),
-                "GLM5.2 slot {slot} greedy argmax index {top_index} outside the vocab"
+                (0..GLM52_SELECTION_VOCAB as i32).contains(&top_index),
+                "GLM5.2 slot {slot} greedy argmax index {top_index} outside the selectable vocab"
             );
         }