qwen

package
v1.3.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 21, 2026 License: MIT Imports: 18 Imported by: 0

Documentation

Index

Constants

View Source 
const KVLogitDriftMax = float32(0.05)

KVLogitDriftMax is max allowed |full-inc| logit diff on the same device.

Variables

View Source 
var (
	Model05B = ModelSpec{
		Name: "Qwen2.5-0.5B-Instruct",
		URL:  "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct/resolve/main",
	}
	Model15B = ModelSpec{
		Name: "Qwen2.5-1.5B-Instruct",
		URL:  "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct/resolve/main",
	}
	Model3B = ModelSpec{
		Name: "Qwen2.5-3B-Instruct",
		URL:  "https://huggingface.co/Qwen/Qwen2.5-3B-Instruct/resolve/main",
		Shards: []string{
			"model-00001-of-00002.safetensors",
			"model-00002-of-00002.safetensors",
		},
	}
	Model7B = ModelSpec{
		Name: "Qwen2.5-7B-Instruct",
		URL:  "https://huggingface.co/Qwen/Qwen2.5-7B-Instruct/resolve/main",
		Shards: []string{
			"model-00001-of-00004.safetensors",
			"model-00002-of-00004.safetensors",
			"model-00003-of-00004.safetensors",
			"model-00004-of-00004.safetensors",
		},
	}
)
View Source 
var DisableKVCache = false

DisableKVCache forces full-sequence forward each step.

View Source 
var ProfileEnabled = false

ProfileEnabled prints prefill/decode timing after each reply (--profile).

Functions

func ActiveModelName 

func ActiveModelName() string

func BuildKVTestIDs added in v1.3.0 

func BuildKVTestIDs(tok *Tokenizer) []int

BuildKVTestIDs builds a short multi-turn chat prefix for KV cache checks.

func CompareKVPrefill added in v1.3.0 

func CompareKVPrefill(model *Model, cfg *Config, tok *Tokenizer, cacheSize int) float32

CompareKVPrefill checks batched ForwardPrefill vs token-by-token ForwardNext.

func EnsureModelFiles 

func EnsureModelFiles(u ui.UI)

EnsureModelFiles downloads config, weights, and tokenizer if missing.

func FormatChat 

func FormatChat(user string) string

FormatChat is an alias for a single-turn instruct prompt.

func FormatSystemPrefix added in v1.3.0 

func FormatSystemPrefix() string

FormatSystemPrefix is the default system block at conversation start.

func FormatUserTurn 

func FormatUserTurn(user string) string

FormatUserTurn builds one user message plus assistant header for generation.

func GraphDisplayName added in v1.3.0 

func GraphDisplayName() string

GraphDisplayName is the short label shown in the graph UI.

func Qwen 

func Qwen(modelArg string, u ui.UI)

Qwen runs interactive chat. Optional model arg: 0.5b, 1.5b, 3b, or 7b (default 3b).

func RunChat 

func RunChat(u ui.UI)

RunChat starts an interactive terminal chat loop.

func RunKVCacheTest added in v1.3.0 

func RunKVCacheTest(u ui.UI) error

RunKVCacheTest loads the active model and compares KV incremental vs full forward.

func SetModel 

func SetModel(name string) error

SetModel selects which Qwen2.5 checkpoint to load (0.5b, 1.5b, 3b, 7b, or full name).

Types

type Attention 

type Attention struct {
	HiddenSize int
	NumHeads   int
	NumKVHeads int
	HeadDim    int
	KVRepeat   int
	Scale      float32
	RoPE       *RoPECache

	QOut int
	KOut int
	VOut int

	QKVProj *neural.Linear // D -> q+k+v (fused q/k/v)
	OProj   *neural.Linear // H*hd -> D
}

Attention implements multi-head self-attention with GQA and RoPE.

func NewAttention 

func NewAttention(cfg *Config, rope *RoPECache, rep ...tensor.AllocReporter) *Attention

func (*Attention) Forward 

func (a *Attention) Forward(x *tensor.Tensor, startPos int, tr neural.Trace) (*tensor.Tensor, neural.Trace)

func (*Attention) ForwardNext 

func (a *Attention) ForwardNext(cache *llm.KVCacheBlock, x *tensor.Tensor, pos int, tr neural.Trace) (*tensor.Tensor, neural.Trace)

type Block 

type Block struct {
	InputNorm *RMSNorm
	Attention *Attention
	PostNorm  *RMSNorm
	MLP       *MLP
}

Block is one Qwen2 decoder layer (pre-norm).

func NewBlock 

func NewBlock(cfg *Config, rope *RoPECache, rep ...tensor.AllocReporter) *Block

func (*Block) Forward 

func (b *Block) Forward(x *tensor.Tensor, startPos int, tr neural.Trace) (*tensor.Tensor, neural.Trace)

func (*Block) ForwardNext 

func (b *Block) ForwardNext(cache *llm.KVCacheBlock, x *tensor.Tensor, pos int, tr neural.Trace) (*tensor.Tensor, neural.Trace)

type ChatSession 

type ChatSession struct {
	Model     *Model
	Tokenizer *Tokenizer
	Graph     *graph.Collector
	Cache     *llm.KVCache
	Tokens    []int
	RNG       *rand.Rand

	ContextSize int
	MaxGen      int
	Temperature float32
	TopK        int
	// contains filtered or unexported fields
}

ChatSession holds loaded model state for multi-turn chat.

func NewChatSession 

func NewChatSession(u ui.UI) (*ChatSession, error)

NewChatSession loads model, tokenizer, and prepares KV cache.

func (*ChatSession) Generate 

func (s *ChatSession) Generate(u ui.UI, logits *tensor.Tensor, prof *TurnProfile) string

Generate streams assistant tokens to u until stop or max length.

func (*ChatSession) Prefill 

func (s *ChatSession) Prefill(ids []int) *tensor.Tensor

Prefill runs a user turn through the cache and returns logits for the first reply token.

func (*ChatSession) Reply 

func (s *ChatSession) Reply(user string, u ui.UI)

Reply encodes a user turn, prefills it, and generates the assistant response.

func (*ChatSession) Reset 

func (s *ChatSession) Reset()

Reset clears conversation history and KV cache.

type Config 

type Config struct {
	VocabSize            int     `json:"vocab_size"`
	HiddenSize           int     `json:"hidden_size"`
	NumHiddenLayers      int     `json:"num_hidden_layers"`
	NumAttentionHeads    int     `json:"num_attention_heads"`
	NumKeyValueHeads     int     `json:"num_key_value_heads"`
	IntermediateSize     int     `json:"intermediate_size"`
	MaxPositionEmbedding int     `json:"max_position_embeddings"`
	RopeTheta            float64 `json:"rope_theta"`
	RMSNormEps           float32 `json:"rms_norm_eps"`
	TieWordEmbeddings    bool    `json:"tie_word_embeddings"`
	BosTokenID           int     `json:"bos_token_id"`
	EosTokenID           int     `json:"eos_token_id"`
}

Config holds Qwen2 model hyperparameters from config.json.

func LoadConfig 

func LoadConfig(path string) (*Config, error)

func (*Config) HeadDim 

func (c *Config) HeadDim() int

func (*Config) KVRepeat 

func (c *Config) KVRepeat() int

type KVForwardCompare added in v1.3.0 

type KVForwardCompare struct {
	TokenCount int
	MaxDiff    float32
	WorstIndex int
	FullArgmax int
	IncArgmax  int
}

KVForwardCompare summarizes incremental KV decode vs one-shot full forward.

func CompareKVForward added in v1.3.0 

func CompareKVForward(model *Model, cfg *Config, ids []int, cacheSize int) KVForwardCompare

CompareKVForward checks each position: incremental ForwardNext vs full Forward logits.

type MLP 

type MLP struct {
	Gate *neural.Linear
	Up   *neural.Linear
	Down *neural.Linear
}

MLP is SwiGLU: down(SiLU(gate(x)) * up(x)).

func NewMLP 

func NewMLP(cfg *Config, rep ...tensor.AllocReporter) *MLP

func (*MLP) Forward 

func (m *MLP) Forward(x *tensor.Tensor, tr neural.Trace) (*tensor.Tensor, neural.Trace)

type Model 

type Model struct {
	Config *Config
	RoPE   *RoPECache

	Embed  *neural.Embeddings
	Blocks []*Block
	Norm   *RMSNorm
	LMHead *tensor.Tensor // [D, V] tied to embed when configured
}

Model is Qwen2 causal LM.

func NewModel 

func NewModel(cfg *Config, maxContext int, u ui.UI) *Model

func (*Model) CaptureGraph added in v1.3.0 

func (m *Model) CaptureGraph(tokens []int, g *graph.Collector)

CaptureGraph records one full forward pass for tokens, then freezes g.

func (*Model) Forward 

func (m *Model) Forward(tokens [][]int) *tensor.Tensor

Forward runs full sequence [B][T] -> logits [B,T,V].

func (*Model) ForwardNext 

func (m *Model) ForwardNext(cache *llm.KVCache, token int, predict bool, tr neural.Trace) (*tensor.Tensor, neural.Trace)

ForwardNext runs one token with KV cache; returns logits [B,1,V] when predict=true.

func (*Model) ForwardPrefill added in v1.3.0 

func (m *Model) ForwardPrefill(cache *llm.KVCache, ids []int, predict bool, tr neural.Trace) (*tensor.Tensor, neural.Trace)

ForwardPrefill runs a token chunk into KV cache; returns logits [B,T,V] when predict=true.

func (*Model) Load 

func (m *Model) Load(weightsDir string, u ui.UI) error

Load reads weights from a model directory (single or sharded safetensors).

type ModelSpec 

type ModelSpec struct {
	Name   string
	URL    string
	Shards []string // empty = single model.safetensors
}

ModelSpec describes a HuggingFace Qwen2.5 instruct checkpoint.

type RMSNorm 

type RMSNorm struct {
	Weight *tensor.Tensor // [D]
	Eps    float32
}

RMSNorm scales by root-mean-square without mean centering.

func NewRMSNorm 

func NewRMSNorm(d int, eps float32, rep ...tensor.AllocReporter) *RMSNorm

func (*RMSNorm) Forward 

func (ln *RMSNorm) Forward(x *tensor.Tensor, tr neural.Trace) (*tensor.Tensor, neural.Trace)

Forward on [B,T,D] or [B,D].

type RoPECache 

type RoPECache struct {
	HeadDim int
	Cos     *tensor.Tensor // [maxPos, headDim]
	Sin     *tensor.Tensor // [maxPos, headDim]
}

RoPECache holds precomputed cos/sin tables for rotary embeddings.

func NewRoPECache 

func NewRoPECache(headDim int, maxPos int, theta float64) *RoPECache

type Tokenizer 

type Tokenizer struct {
	Vocab    map[string]int
	InvVocab []string
	Merges   map[string]int // "a b" -> rank
	ByteEnc  map[byte]rune
	ByteDec  map[rune]byte
	Special  map[string]int
	Regex    *regexp.Regexp
}

Tokenizer encodes/decodes Qwen2 BPE text.

func NewTokenizer 

func NewTokenizer(path string) (*Tokenizer, error)

func (*Tokenizer) ChatEosIDs added in v1.3.0 

func (t *Tokenizer) ChatEosIDs(eosFromConfig int) map[int]bool

ChatEosIDs returns end-of-sequence ids for soft stop heuristics (not im_start).

func (*Tokenizer) ChatStopIDs added in v1.3.0 

func (t *Tokenizer) ChatStopIDs(eosFromConfig int) map[int]bool

ChatStopIDs returns token ids that end assistant generation.

func (*Tokenizer) Decode 

func (t *Tokenizer) Decode(ids []int) string

Decode converts token ids to text.

func (*Tokenizer) DecodeID 

func (t *Tokenizer) DecodeID(id int) string

func (*Tokenizer) Encode 

func (t *Tokenizer) Encode(text string) []int

Encode converts text to token ids.

type TurnProfile added in v1.3.0 

type TurnProfile struct {
	Prefill       time.Duration
	PrefillTokens int
	Decode        time.Duration
	DecodeTokens  int
	Forward       time.Duration
	ForwardSteps  int
	CUDA          cuda.TimingStats
	GPUMemEnd     cuda.MemStats
}

TurnProfile holds wall-clock timing for one user turn.

func (TurnProfile) Print added in v1.3.0 

func (p TurnProfile) Print(u ui.UI, cacheLen int)

Print writes a one-line prefill vs decode summary.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.