diff --git a/openinfer-glm52/src/bookend.rs b/openinfer-glm52/src/bookend.rs index 4377184fe..52a48238a 100644 --- a/openinfer-glm52/src/bookend.rs +++ b/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, lm_head: &DeviceMatrix, - out: &mut Rows, -) -> Result<()> { + out: &mut Rows, +) -> Result { 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) } diff --git a/openinfer-glm52/src/config.rs b/openinfer-glm52/src/config.rs index ef2865351..56968724e 100644 --- a/openinfer-glm52/src/config.rs +++ b/openinfer-glm52/src/config.rs @@ -1,5 +1,8 @@ //! GLM5.2 constants and config probing. +use std::collections::HashSet; +use std::path::Path; + use anyhow::Context; use anyhow::Result; use anyhow::bail; @@ -7,7 +10,13 @@ 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 { + 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::(&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 { + 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::().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": "" }] + }); + let config = serde_json::json!({ + "added_tokens_decoder": { + "4": { "content": "" }, + "5": { "content": "" }, + "bad": { "content": "" } + } + }); + 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()); + } +} diff --git a/openinfer-glm52/src/dspark.rs b/openinfer-glm52/src/dspark.rs index d63348139..a2edbd9cc 100644 --- a/openinfer-glm52/src/dspark.rs +++ b/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)?, diff --git a/openinfer-glm52/src/dspark_test_support.rs b/openinfer-glm52/src/dspark_test_support.rs index faa62e0eb..862c2ecd8 100644 --- a/openinfer-glm52/src/dspark_test_support.rs +++ b/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, diff --git a/openinfer-glm52/src/lib.rs b/openinfer-glm52/src/lib.rs index fb86f8ab9..d6d648102 100644 --- a/openinfer-glm52/src/lib.rs +++ b/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(); diff --git a/openinfer-glm52/src/model/launch_ahead.rs b/openinfer-glm52/src/model/launch_ahead.rs index 9ab70950a..973baa4f2 100644 --- a/openinfer-glm52/src/model/launch_ahead.rs +++ b/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" ); } diff --git a/openinfer-glm52/src/model/mod.rs b/openinfer-glm52/src/model/mod.rs index 9b1437f4c..9a48cbc6e 100644 --- a/openinfer-glm52/src/model/mod.rs +++ b/openinfer-glm52/src/model/mod.rs @@ -55,6 +55,7 @@ use crate::config::GLM52_INDEX_HEAD_DIM; use crate::config::GLM52_INDEX_TOPK; use crate::config::GLM52_LAYERS; use crate::config::GLM52_ROPE_HALF; +use crate::config::GLM52_SELECTION_VOCAB; use crate::config::GLM52_SM_SCALE; use crate::config::GLM52_VOCAB; use crate::config::glm52_layer_has_full_indexer; @@ -581,11 +582,11 @@ impl Glm52RankModel { let (decode_lm_head, decode_vocab_start) = if let Some(rank) = attn_shard { let ranks = moe_topo.device_count(); ensure!( - rank < ranks && GLM52_VOCAB.is_multiple_of(ranks), + rank < ranks && GLM52_SELECTION_VOCAB.is_multiple_of(ranks), "GLM5.2 vocab TP shard {rank}/{ranks} cannot partition {} rows", - GLM52_VOCAB + GLM52_SELECTION_VOCAB ); - let rows = GLM52_VOCAB / ranks; + let rows = GLM52_SELECTION_VOCAB / ranks; let start = rank * rows; let mut data = ctx.stream.alloc_zeros::(rows * GLM52_HIDDEN)?; ctx.stream.memcpy_dtod( @@ -757,7 +758,7 @@ impl Glm52RankModel { cos: ctx.stream.alloc_zeros::(batch * GLM52_ROPE_HALF)?, sin: ctx.stream.alloc_zeros::(batch * GLM52_ROPE_HALF)?, token_ids: ctx.stream.alloc_zeros::(batch)?, - sampling_scratch: BatchSamplingScratch::new(ctx, batch, GLM52_VOCAB)?, + sampling_scratch: BatchSamplingScratch::new(ctx, batch, GLM52_SELECTION_VOCAB)?, speculated: None, device_positions: [0; GLM52_MAX_BATCH_PER_RANK], }) @@ -881,7 +882,7 @@ impl Glm52RankModel { shape.active_rows ); ensure!( - !effectively_greedy(&s.params, GLM52_VOCAB), + !effectively_greedy(&s.params, GLM52_SELECTION_VOCAB), "GLM5.2 effectively-greedy row {} routed to the sampler (coordinator bug)", s.row ); @@ -904,7 +905,7 @@ impl Glm52RankModel { } let logits = HiddenStatesRef { data: bucket.scratch.logits.data(), - hidden_dim: GLM52_VOCAB, + hidden_dim: GLM52_SELECTION_VOCAB, seq_len: shape.bucket, }; let as_row = |s: &crate::runner::Glm52RowSample| BatchSamplingRow { diff --git a/openinfer-glm52/src/model/step_body.rs b/openinfer-glm52/src/model/step_body.rs index b60d54514..9e340c72c 100644 --- a/openinfer-glm52/src/model/step_body.rs +++ b/openinfer-glm52/src/model/step_body.rs @@ -14,13 +14,12 @@ use openinfer_kernels::tensor::DeviceContext; use openinfer_kernels::tensor::DeviceMatrix; use openinfer_kernels::tensor::DeviceVec; -use super::GLM52_MAX_BATCH_PER_RANK; -use super::VOCAB_AR_SLOT; use crate::bookend::glm52_embed_into; use crate::bookend::glm52_final_norm_into; use crate::bookend::glm52_lm_head_into; use crate::config::GLM52_HIDDEN; use crate::config::GLM52_RMS_EPS; +use crate::config::GLM52_SELECTION_VOCAB; use crate::config::GLM52_VOCAB; use crate::dense::glm52_dense_mlp_forward_into; use crate::layer::Glm52DecodeStep; @@ -200,7 +199,7 @@ pub(super) fn run_step_body( } glm52_final_norm_into(ctx, &s.hidden, final_norm, &mut s.final_normed)?; - glm52_lm_head_into(ctx, &s.final_normed, lm_head, &mut s.logits)?; + let logit_rows = glm52_lm_head_into(ctx, &s.final_normed, lm_head, &mut s.logits)?; // Device greedy argmax per row (same semantics as a host scan: lowest // index wins ties, NaN never wins) — the step's egress shrinks from the // full vocab rows to 6 bytes per row, and the kernel chain ends on-device @@ -212,7 +211,7 @@ pub(super) fn run_step_body( ctx, s.logits.data(), batch, - lm_head.rows, + logit_rows, &mut s.argmax_partial_values, &mut s.argmax_partial_indices, &mut s.argmax_values, @@ -221,10 +220,10 @@ pub(super) fn run_step_body( if let Some(rank) = tp { ensure!( - lm_head.rows * rank.state.ranks() == GLM52_VOCAB - && vocab_start == rank.state.rank() * lm_head.rows, + logit_rows * rank.state.ranks() == GLM52_SELECTION_VOCAB + && vocab_start == rank.state.rank() * logit_rows, "GLM5.2 vocab shard [{vocab_start}..{}) does not match TP rank {}/{}", - vocab_start + lm_head.rows, + vocab_start + logit_rows, rank.state.rank(), rank.state.ranks() ); @@ -254,7 +253,7 @@ pub(super) fn run_step_body( )?; } else { ensure!( - lm_head.rows == GLM52_VOCAB && vocab_start == 0, + lm_head.rows == GLM52_VOCAB && logit_rows == GLM52_SELECTION_VOCAB && vocab_start == 0, "GLM5.2 non-TP decode received a sharded vocabulary head" ); } diff --git a/openinfer-glm52/src/oracle/bookend.rs b/openinfer-glm52/src/oracle/bookend.rs index fcb4b491b..fb34f00d8 100644 --- a/openinfer-glm52/src/oracle/bookend.rs +++ b/openinfer-glm52/src/oracle/bookend.rs @@ -30,6 +30,7 @@ use crate::bookend::glm52_embed_into; use crate::bookend::glm52_final_norm_into; use crate::bookend::glm52_lm_head_into; use crate::config::GLM52_HIDDEN; +use crate::config::GLM52_SELECTION_VOCAB; use crate::config::GLM52_VOCAB; use crate::rows::Rows; @@ -59,9 +60,13 @@ fn glm52_lm_head( ctx: &DeviceContext, normed: &Rows, lm_head: &DeviceMatrix, -) -> Result> { +) -> Result> { let mut out = Rows::zeros(ctx, normed.tokens())?; - glm52_lm_head_into(ctx, normed, lm_head, &mut out)?; + let rows = glm52_lm_head_into(ctx, normed, lm_head, &mut out)?; + ensure!( + rows == GLM52_SELECTION_VOCAB, + "bookend oracle expected the full selectable logits prefix" + ); Ok(out) } @@ -248,7 +253,7 @@ fn bookend_oracle_gate() -> Result<()> { // ---- final norm + lm_head: probes + exact argmax ---- ensure!(ORACLE_ARGMAX.len() == ORACLE_CTX, "argmax length mismatch"); - let mut logits_all: Vec = Vec::with_capacity(ORACLE_CTX * GLM52_VOCAB); + let mut logits_all: Vec = Vec::with_capacity(ORACLE_CTX * GLM52_SELECTION_VOCAB); for position in 0..ORACLE_CTX { let mut hidden = Rows::::zeros(&ctx, 1)?; ctx.stream.memcpy_htod( @@ -281,17 +286,25 @@ fn bookend_oracle_gate() -> Result<()> { let tol = ORACLE_LOGITS_REL_TOL * ORACLE_LOGITS_RMS; let failures: Vec<_> = ORACLE_LOGITS_PROBES .iter() - .filter(|&&(idx, expected)| (logits_all[idx] - expected).abs() > tol) + .filter(|&&(checkpoint_idx, expected)| { + let row = checkpoint_idx / GLM52_VOCAB; + let token = checkpoint_idx % GLM52_VOCAB; + token >= GLM52_SELECTION_VOCAB + || (logits_all[row * GLM52_SELECTION_VOCAB + token] - expected).abs() > tol + }) .collect(); println!( "bookend logits: {}/{} probes within tol={tol:.6e}", ORACLE_LOGITS_PROBES.len() - failures.len(), ORACLE_LOGITS_PROBES.len() ); - for &&(idx, expected) in failures.iter().take(10) { + for &&(checkpoint_idx, expected) in failures.iter().take(10) { + let row = checkpoint_idx / GLM52_VOCAB; + let token = checkpoint_idx % GLM52_VOCAB; + let actual = (token < GLM52_SELECTION_VOCAB) + .then(|| logits_all[row * GLM52_SELECTION_VOCAB + token]); println!( - " probe[{idx}]: oracle {expected:.6} vs engine {:.6}", - logits_all[idx] + " checkpoint probe[{checkpoint_idx}] token {token}: oracle {expected:.6} vs engine {actual:?}" ); } ensure!( diff --git a/openinfer-glm52/src/scheduler/plan.rs b/openinfer-glm52/src/scheduler/plan.rs index 3a05c1e4f..a074a0fd2 100644 --- a/openinfer-glm52/src/scheduler/plan.rs +++ b/openinfer-glm52/src/scheduler/plan.rs @@ -6,9 +6,9 @@ use openinfer_sample::SamplingParams; -use super::PAGE; use super::RankSlots; use super::slot::Glm52SlotState; +use crate::config::GLM52_SELECTION_VOCAB; use crate::config::GLM52_VOCAB; use crate::model::GLM52_DECODE_BUCKETS; use crate::model::GLM52_MAX_BATCH_PER_RANK; @@ -182,7 +182,7 @@ pub(super) fn launch_ahead_flags( /// The SAME predicate gates lease-granting and sampling-row collection, which /// is what keeps "sampled row never rides a launch-ahead step" structural. fn takes_argmax(params: &SamplingParams) -> bool { - openinfer_sample::effectively_greedy(params, GLM52_VOCAB) + openinfer_sample::effectively_greedy(params, GLM52_SELECTION_VOCAB) } /// Whether one active request's KV position permits leasing the next step: a @@ -345,7 +345,7 @@ mod tests { // the argmax token) takes the argmax path, so it may ride the lease. let tiny_top_p = decoding_fleet(openinfer_sample::SamplingParams { temperature: 0.7, - top_p: 0.5 / GLM52_VOCAB as f32, + top_p: 0.5 / GLM52_SELECTION_VOCAB as f32, ..Default::default() }); assert!( @@ -437,7 +437,7 @@ mod tests { req: request( vec![10], openinfer_sample::SamplingParams { - top_p: 0.5 / GLM52_VOCAB as f32, + top_p: 0.5 / GLM52_SELECTION_VOCAB as f32, ..sampled(0.8) }, 8, diff --git a/openinfer-glm52/src/scratch.rs b/openinfer-glm52/src/scratch.rs index b4261abc3..7f06ba3d8 100644 --- a/openinfer-glm52/src/scratch.rs +++ b/openinfer-glm52/src/scratch.rs @@ -32,7 +32,7 @@ use openinfer_kernels::tensor::DeviceContext; use crate::config::GLM52_DENSE_INTERMEDIATE; use crate::config::GLM52_HIDDEN; -use crate::config::GLM52_VOCAB; +use crate::config::GLM52_SELECTION_VOCAB; use crate::dspark::GLM52_DSPARK_CONTEXT_DIM; use crate::fp8::Glm52MlpScratch; use crate::indexer::Glm52IndexerScratch; @@ -69,7 +69,7 @@ pub(crate) struct Glm52DecodeScratch { /// copy nodes exist. pub(crate) captured: Option>, pub(crate) final_normed: Rows, - pub(crate) logits: Rows, + pub(crate) logits: Rows, /// Device greedy argmax outputs: each row's top logit bf16 value (for the /// crash-early non-finite guard) and its index — the step's per-row /// 6-byte D2H egress. The two-stage argmax stages per-4096-tile partials @@ -138,12 +138,12 @@ impl Glm52DecodeScratch { .transpose()?, final_normed: Rows::zeros(ctx, tokens)?, logits: Rows::zeros(ctx, tokens)?, - argmax_partial_values: ctx - .stream - .alloc_zeros::(argmax_batch_bf16_split_partials_len(tokens, GLM52_VOCAB))?, - argmax_partial_indices: ctx - .stream - .alloc_zeros::(argmax_batch_bf16_split_partials_len(tokens, GLM52_VOCAB))?, + argmax_partial_values: ctx.stream.alloc_zeros::( + argmax_batch_bf16_split_partials_len(tokens, GLM52_SELECTION_VOCAB), + )?, + argmax_partial_indices: ctx.stream.alloc_zeros::( + argmax_batch_bf16_split_partials_len(tokens, GLM52_SELECTION_VOCAB), + )?, argmax_values: ctx.stream.alloc_zeros::(tokens)?, argmax_indices: ctx.stream.alloc_zeros::(tokens)?, }) diff --git a/openinfer-qwen3/src/config.rs b/openinfer-qwen3/src/config.rs index 25ae2df2d..e5d99fe33 100644 --- a/openinfer-qwen3/src/config.rs +++ b/openinfer-qwen3/src/config.rs @@ -1,7 +1,9 @@ +use std::collections::HashSet; use std::fs; use anyhow::Context; use anyhow::Result; +use log::warn; use serde::Deserialize; pub(crate) const PREFILL_ATTENTION_CTA_TILE_Q: i32 = 64; @@ -30,6 +32,8 @@ pub(crate) struct Config { pub(crate) num_key_value_heads: usize, pub(crate) head_dim: usize, pub(crate) vocab_size: usize, + #[serde(skip)] + pub(crate) selection_vocab: usize, pub(crate) rms_norm_eps: f32, pub(crate) rope_theta: f32, #[serde(default = "default_max_position_embeddings")] @@ -57,6 +61,7 @@ pub(crate) struct DFlashConfig { pub(crate) num_target_layers: usize, pub(crate) head_dim: usize, pub(crate) vocab_size: usize, + pub(crate) selection_vocab: usize, pub(crate) rms_norm_eps: f32, pub(crate) rope_theta: f32, pub(crate) max_position_embeddings: usize, @@ -190,6 +195,7 @@ impl Config { let config_path = format!("{}/config.json", model_path); let content = fs::read_to_string(&config_path)?; let mut config: Config = serde_json::from_str(&content)?; + config.selection_vocab = config.vocab_size; anyhow::ensure!( config.num_key_value_heads > 0 && config @@ -294,6 +300,7 @@ impl DFlashConfig { num_target_layers: raw.num_target_layers, head_dim: raw.head_dim, vocab_size: raw.vocab_size, + selection_vocab: raw.vocab_size, rms_norm_eps: raw.rms_norm_eps, rope_theta, max_position_embeddings: raw.max_position_embeddings, @@ -323,7 +330,7 @@ impl DFlashConfig { self.anchor_first } - pub(crate) fn validate_for_target(&self, target: &Config) -> Result<()> { + pub(crate) fn validate_for_target(&mut self, target: &Config) -> Result<()> { anyhow::ensure!( self.hidden_size == target.hidden_size, "DFlash hidden_size {} does not match target {}", @@ -402,6 +409,7 @@ impl DFlashConfig { self.markov_head_type ); } + self.selection_vocab = target.selection_vocab; Ok(()) } } @@ -459,6 +467,84 @@ impl Eagle3Config { } } +#[derive(Deserialize)] +#[allow(dead_code, clippy::struct_excessive_bools)] +struct AddedTokenConfig { + #[serde(default)] + id: Option, + content: String, + #[serde(default)] + single_word: bool, + #[serde(default)] + lstrip: bool, + #[serde(default)] + rstrip: bool, + #[serde(default)] + normalized: bool, + #[serde(default)] + special: bool, +} + +#[derive(Deserialize)] +struct TokenizerJsonIds { + model: TokenizerModelIds, + #[serde(default)] + added_tokens: Vec, +} + +#[derive(Deserialize)] +struct TokenizerModelIds { + vocab: std::collections::HashMap, +} + +#[derive(Deserialize)] +struct TokenizerConfigIds { + #[serde(default)] + added_tokens_decoder: std::collections::HashMap, +} + +/// Width of the frontend-decodable id space, mirroring the pinned frontend's +/// three-source merge: `tokenizer.json` model vocab + added tokens, then +/// `tokenizer_config.json`'s added-token decoder. A row-range output bound can +/// only represent a dense prefix, so sparse id spaces fail model loading. +pub(crate) fn tokenizer_effective_vocab(model_path: &str) -> Result { + let path = format!("{model_path}/tokenizer.json"); + let content = + fs::read_to_string(&path).map_err(|e| anyhow::anyhow!("cannot read {path}: {e}"))?; + let tokenizer: TokenizerJsonIds = + serde_json::from_str(&content).map_err(|e| anyhow::anyhow!("cannot parse {path}: {e}"))?; + anyhow::ensure!( + !tokenizer.model.vocab.is_empty(), + "{path} model.vocab is empty" + ); + let mut ids: HashSet = tokenizer.model.vocab.into_values().collect(); + ids.extend(tokenizer.added_tokens.iter().filter_map(|token| token.id)); + + let config_path = format!("{model_path}/tokenizer_config.json"); + if let Ok(text) = fs::read_to_string(&config_path) { + match serde_json::from_str::(&text) { + Ok(config) => ids.extend( + config + .added_tokens_decoder + .keys() + .filter_map(|key| key.parse::().ok()), + ), + Err(error) => warn!( + "cannot parse {config_path}: {error}; skipping its added tokens like the frontend does" + ), + } + } + + let width = ids.len(); + let max_id = *ids.iter().max().expect("vocab checked non-empty") as usize; + anyhow::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) +} + impl TensorParallelConfig { pub(crate) fn validate_for(self, config: &Config) -> Result<()> { if self.world_size == 0 { @@ -504,3 +590,48 @@ impl TensorParallelConfig { self.world_size > 1 } } + +#[cfg(test)] +mod tests { + #[test] + fn effective_vocab_merges_all_frontend_id_sources() { + let dir = tempfile::tempdir().unwrap(); + let tokenizer = r#"{ + "model": { "vocab": { "a": 0, "b": 1, "c": 2 } }, + "added_tokens": [ { "id": 3, "content": "" } ] +}"#; + std::fs::write(dir.path().join("tokenizer.json"), tokenizer).unwrap(); + let config = r#"{ "added_tokens_decoder": { "4": { "content": "" }, "5": { "content": "" }, "x": { "content": "" } } }"#; + std::fs::write(dir.path().join("tokenizer_config.json"), config).unwrap(); + + assert_eq!( + super::tokenizer_effective_vocab(dir.path().to_str().unwrap()).unwrap(), + 6 + ); + } + + #[test] + fn effective_vocab_rejects_sparse_id_spaces() { + let dir = tempfile::tempdir().unwrap(); + let tokenizer = r#"{ "model": { "vocab": { "a": 0, "b": 1 } } }"#; + std::fs::write(dir.path().join("tokenizer.json"), tokenizer).unwrap(); + let config = r#"{ "added_tokens_decoder": { "5": { "content": "" } } }"#; + std::fs::write(dir.path().join("tokenizer_config.json"), config).unwrap(); + + assert!(super::tokenizer_effective_vocab(dir.path().to_str().unwrap()).is_err()); + } + + #[test] + fn invalid_decoder_entry_drops_all_decoder_tokens() { + let dir = tempfile::tempdir().unwrap(); + let tokenizer = r#"{ "model": { "vocab": { "a": 0, "b": 1 } } }"#; + std::fs::write(dir.path().join("tokenizer.json"), tokenizer).unwrap(); + let config = r#"{ "added_tokens_decoder": { "2": { "content": "" }, "3": { "content": "", "special": "not-a-bool" } } }"#; + std::fs::write(dir.path().join("tokenizer_config.json"), config).unwrap(); + + assert_eq!( + super::tokenizer_effective_vocab(dir.path().to_str().unwrap()).unwrap(), + 2 + ); + } +} diff --git a/openinfer-qwen3/src/dflash.rs b/openinfer-qwen3/src/dflash.rs index 9a9000f68..64718f536 100644 --- a/openinfer-qwen3/src/dflash.rs +++ b/openinfer-qwen3/src/dflash.rs @@ -252,7 +252,7 @@ impl DFlashBatchScratch { up_out: HiddenStates::zeros(ctx, inter_dim, batch_rows)?, act_out: HiddenStates::zeros(ctx, inter_dim, batch_rows)?, logits_normed: HiddenStates::zeros(ctx, hidden_size, batch_rows)?, - logits: HiddenStates::zeros(ctx, config.vocab_size, batch_rows)?, + logits: HiddenStates::zeros(ctx, config.selection_vocab, batch_rows)?, tail_input: HiddenStates::zeros(ctx, hidden_size, tail_capacity)?, k_tail: HiddenStates::zeros(ctx, kv_dim, tail_capacity)?, v_tail: HiddenStates::zeros(ctx, kv_dim, tail_capacity)?, @@ -757,7 +757,7 @@ impl DFlashDraftModel { for (i, state) in states.iter_mut().enumerate() { state.committed_len += context_lens[i]; } - self.compute_logits_with_target_head_into(target, scratch); + self.compute_logits_with_target_head_into(target, scratch)?; Ok(&scratch.logits) } @@ -823,7 +823,7 @@ impl DFlashDraftModel { &self, target: &Qwen3Model, scratch: &mut DFlashBatchScratch, - ) { + ) -> Result<()> { let ctx = target.device_ctx(); ops::rms_norm_batch_into( ctx, @@ -832,12 +832,14 @@ impl DFlashDraftModel { self.config.rms_norm_eps, &mut scratch.logits_normed, ); - ops::gemm_into( + ops::gemm_rows_into_checked( ctx, target.output_projection(), + 0, + self.config.selection_vocab, &scratch.logits_normed, &mut scratch.logits, - ); + ) } } @@ -846,7 +848,7 @@ pub(crate) fn validate_dflash_config_for_target( dflash_path: &str, target_config: &crate::config::Config, ) -> Result { - let config = DFlashConfig::from_file(dflash_path)?; + let mut config = DFlashConfig::from_file(dflash_path)?; config.validate_for_target(target_config)?; Ok(config) } diff --git a/openinfer-qwen3/src/dflash/loading.rs b/openinfer-qwen3/src/dflash/loading.rs index d53f8caff..c2fad9e77 100644 --- a/openinfer-qwen3/src/dflash/loading.rs +++ b/openinfer-qwen3/src/dflash/loading.rs @@ -7,6 +7,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; @@ -26,7 +27,7 @@ impl DFlashDraftModel { model_path: &str, target: &Qwen3Model, ) -> Result { - let config = DFlashConfig::from_file(model_path) + let mut config = DFlashConfig::from_file(model_path) .with_context(|| format!("load DFlash config from {model_path}"))?; config.validate_for_target(target.config())?; @@ -139,8 +140,22 @@ impl DFlashDraftModel { // embed_tokens/lm_head are intentionally skipped: the head is byte-identical // to the target's, which we reuse for the verify-equivalent logits. let markov = if config.uses_markov_head() { - let w1 = load_tensor_2d(ctx, &shards, &weight_map, MARKOV_W1_TENSOR)?; - let w2 = load_tensor_2d(ctx, &shards, &weight_map, MARKOV_W2_TENSOR)?; + let w1 = load_tensor_2d_row_shard( + ctx, + &shards, + &weight_map, + MARKOV_W1_TENSOR, + 0, + config.selection_vocab, + )?; + let w2 = load_tensor_2d_row_shard( + ctx, + &shards, + &weight_map, + MARKOV_W2_TENSOR, + 0, + config.selection_vocab, + )?; if config.enable_confidence_head { log::info!( "DSpark confidence head present in {model_path} but unused in Phase 1 \ diff --git a/openinfer-qwen3/src/dflash/reservation.rs b/openinfer-qwen3/src/dflash/reservation.rs index c71a9932c..a60279025 100644 --- a/openinfer-qwen3/src/dflash/reservation.rs +++ b/openinfer-qwen3/src/dflash/reservation.rs @@ -27,8 +27,13 @@ pub(crate) struct DFlashMemoryReservation { } impl DFlashMemoryReservation { - pub(crate) fn from_path(draft_path: &str, max_decode_batch_size: usize) -> Result { - let config = DFlashConfig::from_file(draft_path)?; + pub(crate) fn from_path( + draft_path: &str, + selection_vocab: usize, + max_decode_batch_size: usize, + ) -> Result { + let mut config = DFlashConfig::from_file(draft_path)?; + config.selection_vocab = selection_vocab; Ok(Self::from_config(&config, max_decode_batch_size)) } @@ -55,7 +60,7 @@ impl DFlashMemoryReservation { // Same total magnitude as the old per-request scratch summed over the // batch, but now one contiguous allocation. let dense_scratch_per_block_row = - BF16 * (config.vocab_size + 5 * hidden + 2 * q_dim + 3 * inter); + BF16 * (config.selection_vocab + 5 * hidden + 2 * q_dim + 3 * inter); let scratch_total = dense_scratch_per_block_row * config.block_size * max_decode_batch_size; // Draft weights (5 transformer layers + the context projection), +10% slack diff --git a/openinfer-qwen3/src/dspark.rs b/openinfer-qwen3/src/dspark.rs index d8720b233..e23cc11cd 100644 --- a/openinfer-qwen3/src/dspark.rs +++ b/openinfer-qwen3/src/dspark.rs @@ -146,7 +146,7 @@ impl MarkovHead { if !config.uses_markov_head() { return 0; } - let vocab = config.vocab_size; + let vocab = config.selection_vocab; let rank = config.markov_rank; let weights = 2 * vocab * rank * BF16; let scratch = MarkovScratch::bytes(vocab, rank, config.block_size, max_decode_batch_size); @@ -179,7 +179,7 @@ impl MarkovScratch { max_decode_batch_size > 0, "DSpark markov scratch needs a non-zero batch size" ); - let vocab = config.vocab_size; + let vocab = config.selection_vocab; let rank = config.markov_rank; let partials = markov_step_argmax_partials_len(max_decode_batch_size, vocab); let sampled = max_decode_batch_size * config.block_size; diff --git a/openinfer-qwen3/src/executor.rs b/openinfer-qwen3/src/executor.rs index 3ffbfb620..46b21e656 100644 --- a/openinfer-qwen3/src/executor.rs +++ b/openinfer-qwen3/src/executor.rs @@ -1357,6 +1357,7 @@ impl Qwen3Executor { Some(path) => { let reservation = crate::dflash::DFlashMemoryReservation::from_path( path, + model.config().selection_vocab, *BATCH_BUCKETS.last().unwrap(), )?; memory_options.kv_cache_memory_margin_bytes += reservation.fixed_bytes; diff --git a/openinfer-qwen3/src/lora.rs b/openinfer-qwen3/src/lora.rs index b385d95f7..2e2edf687 100644 --- a/openinfer-qwen3/src/lora.rs +++ b/openinfer-qwen3/src/lora.rs @@ -792,6 +792,7 @@ mod tests { num_key_value_heads: 2, head_dim: 2, vocab_size: 16, + selection_vocab: 16, rms_norm_eps: 1e-6, rope_theta: 1_000_000.0, max_position_embeddings: 40960, diff --git a/openinfer-qwen3/src/weights.rs b/openinfer-qwen3/src/weights.rs index 5ece61280..1abcffaef 100644 --- a/openinfer-qwen3/src/weights.rs +++ b/openinfer-qwen3/src/weights.rs @@ -27,6 +27,7 @@ use openinfer_kv_cache::KvBuffer; use super::config::Config; use super::config::TensorParallelConfig; +use super::config::tokenizer_effective_vocab; use crate::batch_decode_buffers::BatchDecodeBuffers; use crate::lora::DeviceLoraAdapter; use crate::lora::DeviceLoraLayer; @@ -391,7 +392,21 @@ impl Qwen3Model { debug!("Initializing GPU device {}", runtime.device_ordinal); let ctx = DeviceContext::new_with_device(runtime.device_ordinal)?; - let config = Config::from_file(model_path)?; + let mut config = Config::from_file(model_path)?; + let effective_vocab = tokenizer_effective_vocab(model_path)?; + anyhow::ensure!( + effective_vocab <= config.vocab_size, + "tokenizer defines ids up to {} but checkpoint vocab_size is {}", + effective_vocab - 1, + config.vocab_size, + ); + config.selection_vocab = effective_vocab; + if effective_vocab < config.vocab_size { + info!( + "output projection: selection bounded to decodable vocab {} (checkpoint pads to {})", + effective_vocab, config.vocab_size + ); + } let tensor_parallel = runtime.tensor_parallel.unwrap_or_default(); tensor_parallel.validate_for(&config)?;