gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Qwen3.6 MoE inference support

Browse Source
This commit is contained in:
Gahow Wang
2026-07-13 20:24:41 +08:00
parent 588bfd9df3
commit a2de146fb6
27 changed files with 3153 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
crates/xserv-model/src/bin/xserv-chat.rs
View File

@@ -5,8 +5,8 @@ use std::sync::{Arc, mpsc};
use std::thread;
use xserv_model::{
BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, SamplingParams,
loader, sample, sample_greedy_penalized,
BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, ModelFamily, PagedKVCache, Qwen3,
SamplingParams, loader, sample, sample_greedy_penalized,
};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -93,24 +93,26 @@ fn tp_worker_loop(
rank as u32,
));
let weights = loader::load_model_dir(&model_dir, Device::Cpu);
let model = if config.is_moe() {
ChatModel::GptOss(GptOss::from_weights_tp(
let model = match config.family() {
ModelFamily::GptOss => ChatModel::GptOss(GptOss::from_weights_tp(
config.clone(),
weights,
rank,
world,
rank as u32,
Some(tp),
))
} else {
ChatModel::Qwen3(Qwen3::from_weights_tp(
)),
ModelFamily::Qwen35Moe => {
panic!("Qwen3.5/3.6 MoE chat inference is not implemented yet")
}
_ => ChatModel::Qwen3(Qwen3::from_weights_tp(
config.clone(),
weights,
rank,
world,
rank as u32,
Some(tp),
))
)),
};
let local_kv = config.num_kv_heads() / world;
let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
@@ -323,20 +325,27 @@ fn main() {
);
let config = ModelConfig::from_file(&opts.model_dir.join("config.json"));
let model_type = config.model_type.as_deref().unwrap_or("unknown");
let is_moe = config.is_moe();
let model_type = config.model_type_str();
let model_family = config.family();
let is_gpt_oss = model_family == ModelFamily::GptOss;
let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1);
eprintln!(
"Model: {model_type}{}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
if is_moe { " (MoE)" } else { "" },
if config.is_gpt_oss() { " (MoE)" } else { "" },
config.num_layers(),
config.hidden(),
config.num_heads(),
config.num_kv_heads(),
config.vocab_size,
config.vocab_size(),
max_seq_len
);
if model_family == ModelFamily::Qwen35Moe {
eprintln!(
"Qwen3.5/3.6 MoE is recognized but chat inference is not implemented yet; it requires DeltaNet/linear-attention and Qwen MoE support."
);
std::process::exit(1);
}
let world = opts.tp;
if world > 1 {
@@ -374,7 +383,7 @@ fn main() {
let tp = Arc::new(xserv_distributed::TpContext::init(0, world, id, 0));
let weights = loader::load_model_dir(&opts.model_dir, Device::Cpu);
eprintln!("Loaded {} tensors", weights.len());
let m = if is_moe {
let m = if is_gpt_oss {
ChatModel::GptOss(GptOss::from_weights_tp(
config.clone(),
weights,
@@ -412,7 +421,7 @@ fn main() {
eprintln!("Loading weights...");
let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
eprintln!("Loaded {} tensors", weights.len());
let m = if is_moe {
let m = if is_gpt_oss {
ChatModel::GptOss(GptOss::from_weights(config.clone(), weights))
} else {
ChatModel::Qwen3(Qwen3::from_weights(config.clone(), weights))
@@ -465,7 +474,7 @@ fn main() {
_ => {}
}
if is_moe {
if is_gpt_oss {
// Harmony multi-turn: re-render the whole conversation (prior
// analysis dropped) and re-prefill into a freshly cleared slot.
let prompt =
@@ -495,7 +504,7 @@ fn main() {
max_new_tokens,
use_color,
&tp_handle,
is_moe,
is_gpt_oss,
opts.enable_thinking,
);
moe_history.push((input.to_string(), answer));
@@ -540,7 +549,7 @@ fn main() {
max_new_tokens,
use_color,
&tp_handle,
is_moe,
is_gpt_oss,
opts.enable_thinking,
);
match finish {
@@ -672,6 +681,7 @@ fn parse_args() -> CliOptions {
temperature,
top_k,
top_p,
..SamplingParams::default()
},
system_prompt,
enable_thinking,
@@ -846,25 +856,25 @@ fn generate_with_paged_cache(
max_tokens: usize,
use_color: bool,
tp: &Option<TpHandle>,
is_moe: bool,
is_gpt_oss: bool,
enable_thinking: bool,
) -> (Finish, String) {
let harmony_end_id = if is_moe {
let harmony_end_id = if is_gpt_oss {
tokenizer.special_token_id("<|end|>")
} else {
None
};
let harmony_channel_id = if is_moe {
let harmony_channel_id = if is_gpt_oss {
tokenizer.special_token_id("<|channel|>")
} else {
None
};
let harmony_message_id = if is_moe {
let harmony_message_id = if is_gpt_oss {
tokenizer.special_token_id("<|message|>")
} else {
None
};
let harmony_special: Vec<u32> = if is_moe {
let harmony_special: Vec<u32> = if is_gpt_oss {
[
"<|channel|>",
"<|start|>",
@@ -888,7 +898,7 @@ fn generate_with_paged_cache(
InAnalysis,
InFinal,
}
let mut hstate = if is_moe {
let mut hstate = if is_gpt_oss {
HarmonyState::InFinal
} else {
HarmonyState::Normal
@@ -931,7 +941,7 @@ fn generate_with_paged_cache(
let mut next = pick(&logits, sampling, &history);
let mut decode_buffer = Vec::new();
let mut in_thinking = false;
let show_thinking = is_moe && enable_thinking;
let show_thinking = is_gpt_oss && enable_thinking;
// 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
@@ -1046,7 +1056,7 @@ fn generate_with_paged_cache(
next = pick(&logits, sampling, &history);
continue;
}
if is_moe && hstate != HarmonyState::InFinal {
if is_gpt_oss && 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

18
crates/xserv-model/src/bin/xserv-cli.rs
View File

@@ -1,7 +1,7 @@
use std::io::{self, Write};
use std::path::PathBuf;
use xserv_model::{
BLOCK_SIZE, KVCache, ModelConfig, PagedKVCache, SamplingParams, loader, sample,
BLOCK_SIZE, KVCache, ModelConfig, ModelFamily, PagedKVCache, SamplingParams, loader, sample,
sample_greedy_penalized,
};
use xserv_tensor::{DType, Device};
@@ -43,6 +43,7 @@ fn main() {
temperature: flag(&args, "--temperature", 0.0f32),
top_k: flag(&args, "--top-k", 0usize),
top_p: flag(&args, "--top-p", 1.0f32),
..SamplingParams::default()
};
let rep_penalty = flag(&args, "--rep-penalty", 1.0f32);
let rep_window = flag(&args, "--rep-window", 512usize);
@@ -56,22 +57,29 @@ fn main() {
);
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let model_type = config.model_type.as_deref().unwrap_or("unknown");
let model_type = config.model_type_str();
let model_family = config.family();
eprintln!(
"Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
config.num_layers(),
config.hidden(),
config.num_heads(),
config.num_kv_heads(),
config.vocab_size
config.vocab_size()
);
if model_family == ModelFamily::Qwen35Moe {
eprintln!(
"Qwen3.5/3.6 MoE is recognized but CLI inference is not implemented yet; it requires DeltaNet/linear-attention and Qwen MoE support."
);
std::process::exit(1);
}
eprintln!("Loading weights...");
let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
eprintln!("Loaded {} tensors", weights.len());
let is_qwen3 = model_type.contains("qwen");
let is_gpt_oss = model_type.contains("gpt_oss");
let is_qwen3 = model_family == ModelFamily::Qwen3;
let is_gpt_oss = model_family == ModelFamily::GptOss;
let dtype = if is_qwen3 || is_gpt_oss {
DType::BF16
} else {

180
crates/xserv-model/src/config.rs
View File

@@ -1,6 +1,27 @@
use serde::Deserialize;
use std::path::Path;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelFamily {
Gpt2,
Qwen3,
GptOss,
Qwen35Moe,
Unknown,
}
impl ModelFamily {
pub fn as_str(self) -> &'static str {
match self {
ModelFamily::Gpt2 => "gpt2",
ModelFamily::Qwen3 => "qwen3",
ModelFamily::GptOss => "gpt_oss",
ModelFamily::Qwen35Moe => "qwen3_5_moe",
ModelFamily::Unknown => "unknown",
}
}
}
#[derive(Debug, Clone, Deserialize)]
pub struct RopeScaling {
pub rope_type: Option<String>,
@@ -10,10 +31,21 @@ pub struct RopeScaling {
pub beta_slow: Option<f64>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct RopeParameters {
pub rope_type: Option<String>,
pub rope_theta: Option<f64>,
pub partial_rotary_factor: Option<f64>,
pub mrope_interleaved: Option<bool>,
pub mrope_section: Option<Vec<usize>>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ModelConfig {
pub architectures: Option<Vec<String>>,
pub model_type: Option<String>,
#[serde(default)]
pub text_config: Option<Box<ModelConfig>>,
// Modern HF naming
#[serde(default)]
@@ -26,6 +58,7 @@ pub struct ModelConfig {
pub num_key_value_heads: Option<usize>,
#[serde(default)]
pub num_hidden_layers: Option<usize>,
#[serde(default)]
pub vocab_size: usize,
#[serde(default)]
pub max_position_embeddings: Option<usize>,
@@ -56,14 +89,22 @@ pub struct ModelConfig {
#[serde(default)]
pub tie_word_embeddings: Option<bool>,
// MoE (gpt-oss)
// MoE (gpt-oss / Qwen MoE)
#[serde(default)]
pub num_local_experts: Option<usize>,
#[serde(default)]
pub num_experts: Option<usize>,
#[serde(default)]
pub num_experts_per_tok: Option<usize>,
#[serde(default)]
pub moe_intermediate_size: Option<usize>,
#[serde(default)]
pub shared_expert_intermediate_size: Option<usize>,
#[serde(default)]
pub layer_types: Option<Vec<String>>,
#[serde(default)]
pub full_attention_interval: Option<usize>,
#[serde(default)]
pub sliding_window: Option<usize>,
#[serde(default)]
pub attention_bias: Option<bool>,
@@ -72,11 +113,29 @@ pub struct ModelConfig {
#[serde(default)]
pub rope_scaling: Option<RopeScaling>,
#[serde(default)]
pub rope_parameters: Option<RopeParameters>,
#[serde(default)]
pub swiglu_limit: Option<f64>,
#[serde(default)]
pub geglu_alpha: Option<f64>,
#[serde(default)]
pub hidden_act: Option<String>,
// Qwen3.5/3.6 linear attention / DeltaNet metadata
#[serde(default)]
pub linear_conv_kernel_dim: Option<usize>,
#[serde(default)]
pub linear_key_head_dim: Option<usize>,
#[serde(default)]
pub linear_num_key_heads: Option<usize>,
#[serde(default)]
pub linear_num_value_heads: Option<usize>,
#[serde(default)]
pub linear_value_head_dim: Option<usize>,
#[serde(default)]
pub partial_rotary_factor: Option<f64>,
#[serde(default)]
pub mtp_num_hidden_layers: Option<usize>,
}
impl ModelConfig {
@@ -87,80 +146,153 @@ impl ModelConfig {
.unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display()))
}
pub fn text(&self) -> &Self {
self.text_config.as_deref().unwrap_or(self)
}
pub fn model_type_str(&self) -> &str {
self.text()
.model_type
.as_deref()
.or(self.model_type.as_deref())
.unwrap_or("unknown")
}
pub fn family(&self) -> ModelFamily {
let top_type = self.model_type.as_deref().unwrap_or("");
let text_type = self.text().model_type.as_deref().unwrap_or(top_type);
let has_arch = |needle: &str| {
self.architectures
.as_ref()
.map(|arch| arch.iter().any(|a| a.contains(needle)))
.unwrap_or(false)
};
if top_type == "qwen3_5_moe" || text_type == "qwen3_5_moe_text" || has_arch("Qwen3_5Moe") {
ModelFamily::Qwen35Moe
} else if text_type == "gpt_oss" || top_type == "gpt_oss" {
ModelFamily::GptOss
} else if text_type.contains("qwen") || top_type.contains("qwen") {
ModelFamily::Qwen3
} else if text_type.contains("gpt2") || top_type.contains("gpt2") {
ModelFamily::Gpt2
} else {
ModelFamily::Unknown
}
}
pub fn hidden(&self) -> usize {
self.hidden_size
.or(self.n_embd)
let c = self.text();
c.hidden_size
.or(c.n_embd)
.expect("hidden_size or n_embd required")
}
pub fn num_heads(&self) -> usize {
self.num_attention_heads
.or(self.n_head)
let c = self.text();
c.num_attention_heads
.or(c.n_head)
.expect("num_attention_heads or n_head required")
}
pub fn num_layers(&self) -> usize {
self.num_hidden_layers
.or(self.n_layer)
let c = self.text();
c.num_hidden_layers
.or(c.n_layer)
.expect("num_hidden_layers or n_layer required")
}
pub fn max_seq_len(&self) -> usize {
self.max_position_embeddings
.or(self.n_positions)
.unwrap_or(2048)
let c = self.text();
c.max_position_embeddings.or(c.n_positions).unwrap_or(2048)
}
pub fn ffn_hidden(&self) -> usize {
self.intermediate_size
.or(self.n_inner)
.unwrap_or(self.hidden() * 4)
let c = self.text();
c.intermediate_size
.or(c.moe_intermediate_size)
.or(c.n_inner)
.unwrap_or_else(|| self.hidden() * 4)
}
pub fn vocab_size(&self) -> usize {
self.text().vocab_size
}
pub fn num_kv_heads(&self) -> usize {
self.num_key_value_heads.unwrap_or(self.num_heads())
self.text().num_key_value_heads.unwrap_or(self.num_heads())
}
pub fn head_dim(&self) -> usize {
self.explicit_head_dim
self.text()
.explicit_head_dim
.unwrap_or_else(|| self.hidden() / self.num_heads())
}
pub fn ln_eps(&self) -> f32 {
self.layer_norm_eps
.or(self.layer_norm_epsilon)
let c = self.text();
c.layer_norm_eps
.or(c.layer_norm_epsilon)
.or(c.rms_norm_eps)
.unwrap_or(1e-5) as f32
}
pub fn rope_theta_value(&self) -> Option<f64> {
let c = self.text();
c.rope_theta
.or_else(|| c.rope_parameters.as_ref().and_then(|rp| rp.rope_theta))
}
pub fn tied_embeddings(&self) -> bool {
self.tie_word_embeddings.unwrap_or(true)
self.text().tie_word_embeddings.unwrap_or(true)
}
pub fn num_experts(&self) -> usize {
self.num_local_experts.unwrap_or(0)
self.text()
.num_local_experts
.or(self.text().num_experts)
.unwrap_or(0)
}
pub fn experts_per_token(&self) -> usize {
self.num_experts_per_tok.unwrap_or(1)
self.text().num_experts_per_tok.unwrap_or(1)
}
pub fn is_moe(&self) -> bool {
self.num_local_experts.unwrap_or(0) > 1
self.num_experts() > 1
}
pub fn is_gpt_oss(&self) -> bool {
self.family() == ModelFamily::GptOss
}
pub fn is_qwen35_moe(&self) -> bool {
self.family() == ModelFamily::Qwen35Moe
}
pub fn is_sliding_layer(&self, layer_idx: usize) -> bool {
self.layer_types
self.text()
.layer_types
.as_ref()
.and_then(|lt| lt.get(layer_idx))
.map(|t| t == "sliding_attention")
.unwrap_or(false)
}
pub fn is_linear_attention_layer(&self, layer_idx: usize) -> bool {
self.text()
.layer_types
.as_ref()
.and_then(|lt| lt.get(layer_idx))
.map(|t| t == "linear_attention")
.unwrap_or(false)
}
pub fn window_size(&self) -> usize {
self.sliding_window.unwrap_or(0)
self.text().sliding_window.unwrap_or(0)
}
pub fn geglu_alpha(&self) -> f32 {
self.geglu_alpha.unwrap_or(1.702) as f32
self.text().geglu_alpha.unwrap_or(1.702) as f32
}
}

4
crates/xserv-model/src/decode_graph.rs
View File

@@ -98,9 +98,9 @@ impl DecodeGraphState {
let num_kv_heads = config.num_kv_heads();
let head_dim = config.head_dim();
let intermediate = config.ffn_hidden();
let vocab_size = config.vocab_size;
let vocab_size = config.vocab_size();
let num_layers = config.num_layers();
let eps = config.rms_norm_eps.unwrap_or(1e-6) as f32;
let eps = config.ln_eps();
let es = 2usize; // BF16 = 2 bytes
let stream = CudaStream::new().expect("create CUDA stream for graph");

10
crates/xserv-model/src/lib.rs
View File

@@ -8,10 +8,11 @@ pub mod kv_cache;
pub mod loader;
pub mod paged_kv_cache;
pub mod qwen3;
pub mod qwen35_moe;
pub mod qwen3_graph;
pub mod sampling;
pub use config::ModelConfig;
pub use config::{ModelConfig, ModelFamily};
pub use decode_graph::{DecodeGraphState, LayerWeightPtrs};
pub use gpt_oss::GptOss;
pub use gpt_oss_graph::{GptOssDecodeGraph, GraphedGptOssDecoder};
@@ -19,7 +20,12 @@ pub use gpt2::{GPT2, KVCache};
pub use kv_cache::GpuKVCache;
pub use paged_kv_cache::{BLOCK_SIZE, BlockAllocator, Location, PagedKVCache};
pub use qwen3::Qwen3;
pub use sampling::{SamplingParams, sample, sample_greedy_penalized};
pub use qwen35_moe::{
Qwen35AttentionMeta, Qwen35FullAttentionOutput, Qwen35FullAttentionWeights,
Qwen35LinearAttentionWeights, Qwen35LinearProjectionOutput, Qwen35Moe, Qwen35MoeSpec,
Qwen35RecurrentCache, Qwen35TensorMeta, Qwen35TensorSpec,
};
pub use sampling::{SamplingParams, sample, sample_greedy_penalized, sample_with_history};
/// Initialize GPU kernel hooks. Called automatically by model constructors,
/// but safe to call multiple times (idempotent via OnceLock).

60
crates/xserv-model/src/loader.rs
View File

@@ -1,10 +1,20 @@
use half::{bf16, f16};
use safetensors::SafeTensors;
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::path::Path;
use xserv_tensor::{DType, Device, Tensor};
pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
load_safetensors_filtered(path, device, |_| true)
}
/// Load only tensors accepted by `keep`. The safetensors file is still read as
/// one byte buffer, but unneeded tensor payloads are not copied into Tensors.
pub fn load_safetensors_filtered(
path: &Path,
device: Device,
keep: impl Fn(&str) -> bool,
) -> HashMap<String, Tensor> {
let data =
std::fs::read(path).unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
let st = SafeTensors::deserialize(&data)
@@ -13,6 +23,9 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor>
let mut tensors = HashMap::new();
for (name, view) in st.tensors() {
if !keep(&name) {
continue;
}
let shape: Vec<usize> = view.shape().to_vec();
let raw_bytes = view.data();
let dtype = match view.dtype() {
@@ -36,27 +49,64 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor>
/// Load from a directory containing model.safetensors (or sharded files) + config.json.
pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
load_model_dir_filtered(dir, device, |_| true)
}
/// Load a filtered subset of a model directory. For indexed sharded models,
/// consult `model.safetensors.index.json` first so shards containing no wanted
/// tensors are never read. This is critical for pipeline parallelism on large
/// models: each stage should read only its layer range, not the full checkpoint.
pub fn load_model_dir_filtered(
dir: &Path,
device: Device,
keep: impl Fn(&str) -> bool,
) -> HashMap<String, Tensor> {
let single = dir.join("model.safetensors");
if single.exists() {
return load_safetensors(&single, device);
return load_safetensors_filtered(&single, device, keep);
}
let index_path = dir.join("model.safetensors.index.json");
let wanted_shards: Option<HashSet<String>> = if index_path.exists() {
let text = std::fs::read_to_string(&index_path)
.unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
let index: serde_json::Value = serde_json::from_str(&text)
.unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
let map = index
.get("weight_map")
.and_then(|v| v.as_object())
.unwrap_or_else(|| panic!("{} has no weight_map", index_path.display()));
Some(
map.iter()
.filter(|(name, _)| keep(name))
.filter_map(|(_, shard)| shard.as_str().map(str::to_string))
.collect(),
)
} else {
None
};
// Try sharded: model-00001-of-NNNNN.safetensors
let mut all_tensors = HashMap::new();
let mut entries: Vec<_> = std::fs::read_dir(dir)
.unwrap()
.filter_map(|e| e.ok())
.filter(|e| {
e.path()
let is_safetensors = e
.path()
.file_name()
.map(|f| f.to_string_lossy().ends_with(".safetensors"))
.unwrap_or(false)
.unwrap_or(false);
let is_wanted = wanted_shards.as_ref().is_none_or(|wanted| {
wanted.contains(&e.file_name().to_string_lossy().to_string())
});
is_safetensors && is_wanted
})
.collect();
entries.sort_by_key(|e| e.file_name());
for entry in entries {
let tensors = load_safetensors(&entry.path(), device);
let tensors = load_safetensors_filtered(&entry.path(), device, &keep);
all_tensors.extend(tensors);
}

1447
crates/xserv-model/src/qwen35_moe.rs Normal file
View File

File diff suppressed because it is too large Load Diff

64
crates/xserv-model/src/sampling.rs
View File

@@ -1,5 +1,6 @@
use half::bf16;
use rand::Rng;
use std::collections::HashSet;
use xserv_tensor::{DType, Device, Tensor};
#[derive(Clone)]
@@ -7,6 +8,8 @@ pub struct SamplingParams {
pub temperature: f32,
pub top_k: usize,
pub top_p: f32,
pub presence_penalty: f32,
pub repetition_penalty: f32,
}
impl Default for SamplingParams {
@@ -15,6 +18,8 @@ impl Default for SamplingParams {
temperature: 0.0,
top_k: 0,
top_p: 1.0,
presence_penalty: 0.0,
repetition_penalty: 1.0,
}
}
}
@@ -22,12 +27,21 @@ impl Default for SamplingParams {
/// Sample a token from logits with shape [seq_len, vocab_size].
/// Uses the last position's logits. Handles both F32 and BF16 dtypes.
pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
sample_with_history(logits, params, &[])
}
/// Sample while applying penalties to tokens already present in the request.
/// Presence penalty follows the OpenAI convention (subtract once per distinct
/// token); repetition penalty follows the HF convention.
pub fn sample_with_history(logits: &Tensor, params: &SamplingParams, history: &[u32]) -> u32 {
assert_eq!(logits.ndim(), 2);
// Greedy fast path: GPU argmax + 4-byte D2H instead of copying the whole
// [seq, vocab] logits to the host and scanning it (~201k bf16/token).
// NaN logits lose every `>` comparison in the kernel, matching the
// NaN-safe host argmax below.
if params.temperature == 0.0
&& params.presence_penalty == 0.0
&& params.repetition_penalty == 1.0
&& logits.dtype() == DType::BF16
&& matches!(logits.device(), Device::Cuda(_))
&& logits.is_contiguous()
@@ -55,11 +69,6 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
_ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
};
// Greedy
if params.temperature == 0.0 {
return argmax(&last_row);
}
// NaN-safe: sampling path uses partial_cmp().unwrap() in top-k/top-p
// sorts and softmax; a single NaN logit would panic the engine thread.
// Replace NaN with -inf (equivalent to masking) instead.
@@ -74,6 +83,29 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
eprintln!("[sampling] WARNING: NaN logits encountered in sample()");
}
if !history.is_empty() && (params.presence_penalty != 0.0 || params.repetition_penalty != 1.0) {
let seen: HashSet<u32> = history.iter().copied().collect();
for id in seen {
let i = id as usize;
if i >= last_row.len() {
continue;
}
if params.repetition_penalty != 1.0 {
let value = last_row[i];
last_row[i] = if value > 0.0 {
value / params.repetition_penalty
} else {
value * params.repetition_penalty
};
}
last_row[i] -= params.presence_penalty;
}
}
if params.temperature == 0.0 {
return argmax(&last_row);
}
// Apply temperature
let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();
@@ -195,3 +227,25 @@ fn argmax(data: &[f32]) -> u32 {
}
best_i as u32
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn greedy_sampling_applies_history_penalties() {
let logits = Tensor::from_slice(&[10.0f32, 9.0, 0.0], &[1, 3]);
let presence = SamplingParams {
presence_penalty: 2.0,
..SamplingParams::default()
};
assert_eq!(sample_with_history(&logits, &presence, &[0]), 1);
let repetition = SamplingParams {
repetition_penalty: 2.0,
..SamplingParams::default()
};
assert_eq!(sample_with_history(&logits, &repetition, &[0]), 1);
}
}