Functions — github.com/skyzh/tiny-llm

↓ 1 callersMethodremove_request

(self, id: int)

src/tiny_llm_ref/kv_cache.py:205

↓ 1 callersMethodrewind

(self, n: int)

src/tiny_llm/paged_kv_cache.py:92

↓ 1 callersMethodrewind

Remove the newest n logical tokens from this cache. This is needed by speculative decoding when some draft tokens are reject

src/tiny_llm_ref/kv_cache.py:55

↓ 1 callersMethodrewind

(self, n: int)

src/tiny_llm_ref/kv_cache.py:242

↓ 1 callersFunctionrope_helper

(stream: mx.Stream, traditional: bool, precision: mx.Dtype)

tests_refsol/test_week_2_day_6.py:10

↓ 1 callersFunctionroute_topk

( x: mx.array, w_router: QuantizedWeights, top_k: int, norm_topk_prob: bool = False, )

src/tiny_llm_ref/moe.py:36

↓ 1 callersFunctionrun_batch_requests_week2

Benchmark Week 2 continuous batching without tokenizer/detokenizer overhead.

bench.py:216

↓ 1 callersFunctionrun_one_request_week1

( model, request: BenchRequest, )

bench.py:165

↓ 1 callersFunctionrun_one_request_week2

( model, request: BenchRequest, kv_cache_cls, )

bench.py:188

↓ 1 callersFunctionsample_next_week2_batched

( model, y: mx.array, offsets: list[int], kv_cache: list )

bench.py:157

↓ 1 callersFunctionscaled_dot_product_attention_simple

A simple implementation of scaled dot product attention. Assuming Q, K, V are of the same shape. Assuming mask is always a float array that y

src/tiny_llm_ref/attention.py:6

↓ 1 callersFunctionshortcut_name_to_full_name

(shortcut_name: str)

src/tiny_llm/models.py:6

↓ 1 callersFunctionshortcut_name_to_full_name

(shortcut_name: str)

src/tiny_llm_ref/models.py:6

↓ 1 callersFunctionsilu

(x: mx.array)

src/tiny_llm/basics.py:18

↓ 1 callersFunctionsimple_generate

( model: Qwen3ModelWeek1, tokenizer: TokenizerWrapper, prompt: str, sampler: Callable[[mx.arra

src/tiny_llm_ref/generate.py:14

↓ 1 callersFunctionsimple_generate_with_kv_cache

( model: Qwen3ModelWeek2, tokenizer: TokenizerWrapper, prompt: str )

src/tiny_llm_ref/generate.py:47

↓ 1 callersFunctionsoftmax

(x: mx.array, axis: int)

src/tiny_llm/basics.py:5

↓ 1 callersFunctionspeculative_generate

( draft_model: Qwen3ModelWeek2, model: Qwen3ModelWeek2, draft_tokenizer: TokenizerWrapper, tok

src/tiny_llm_ref/generate.py:82

↓ 1 callersFunctiontest

(args)

scripts/dev-tools.py:25

↓ 1 callersFunctiontest_refsol

(args)

scripts/dev-tools.py:38

↓ 1 callersMethodtext

(self)

src/tiny_llm/batch.py:58

↓ 1 callersMethodtry_prefill

Prefill this request up to max_step size, returns None if prefill is not done

src/tiny_llm/batch.py:42

↓ 1 callersMethodtry_prefill

Prefill this request up to max_step size, returns None if prefill is not done

src/tiny_llm_ref/batch.py:42

↓ 1 callersMethodupdate_and_fetch

( self, keys: mx.array, values: mx.array, mask_length: int | None = None,

src/tiny_llm/kv_cache.py:38

↓ 1 callersMethodupdate_and_fetch_paged

( self, key: mx.array, value: mx.array, mask_length: int | None = None,

src/tiny_llm/paged_kv_cache.py:83

↓ 1 callersMethodupdate_and_fetch_paged

( self, keys: mx.array, values: mx.array, mask_length: int | None = None,

src/tiny_llm_ref/kv_cache.py:132

↓ 1 callersFunctionvalidate_args

(args: argparse.Namespace)

bench.py:68

Method__call__

( self, x: mx.array, offsets: list[int], cache: TinyKvCache, mask: mx.

src/tiny_llm/qwen3_week2.py:32

Method__call__

(self, x: mx.array)

src/tiny_llm/qwen3_week2.py:53

Method__call__

( self, x: mx.array, offset: int, cache: TinyKvCache, mask: mx.array |

src/tiny_llm/qwen3_week2.py:83

Method__call__

( self, inputs: mx.array, offset: int, cache: list[TinyKvCache], )

src/tiny_llm/qwen3_week2.py:102

Method__call__

(self, x: mx.array)

src/tiny_llm/moe.py:35

Method__call__

( self, x: mx.array, offset: list[slice] | slice | None = None )

src/tiny_llm/positional_encoding.py:14

Method__call__

( self, query: mx.array, key: mx.array, value: mx.array, mask: mx.arra

src/tiny_llm/attention.py:27

Method__call__

( self, x: mx.array, mask: mx.array | str | None = None, )

src/tiny_llm/qwen3_week1.py:30

Method__call__

(self, x: mx.array)

src/tiny_llm/qwen3_week1.py:49

Method__call__

( self, x: mx.array, mask: mx.array | str | None = None, )

src/tiny_llm/qwen3_week1.py:78

Method__call__

( self, inputs: mx.array, )

src/tiny_llm/qwen3_week1.py:90

Method__call__

(self, x: mx.array)

src/tiny_llm/layer_norm.py:8

Method__call__

( self, inputs: mx.array, offset: int | list[int] | mx.array, cache: list[Tiny

src/tiny_llm/qwen3_week3.py:23

Method__call__

(self, x: mx.array)

src/tiny_llm/embedding.py:9

Method__call__

(self, x: mx.array)

src/tiny_llm/embedding.py:20

Method__call__

( self, x: mx.array, offsets: list[int], cache: TinyKvCache, mask: mx.

src/tiny_llm_ref/qwen3_week2.py:55

Method__call__

(self, x: mx.array)

src/tiny_llm_ref/qwen3_week2.py:122

Method__call__

( self, x: mx.array, offset: int, cache: TinyKvCache, mask: mx.array |

src/tiny_llm_ref/qwen3_week2.py:177

Method__call__

( self, inputs: mx.array, offset: int, cache: list[TinyKvCache], )

src/tiny_llm_ref/qwen3_week2.py:274

Method__call__

(self, x: mx.array)

src/tiny_llm_ref/moe.py:68

Method__call__

( self, x: mx.array, offset: list[slice] | slice | None = None )

src/tiny_llm_ref/positional_encoding.py:26

Method__call__

( self, query: mx.array, key: mx.array, value: mx.array, mask: mx.arra

src/tiny_llm_ref/attention.py:181

Method__call__

( self, x: mx.array, mask: mx.array | str | None = None, )

src/tiny_llm_ref/qwen3_week1.py:45

Method__call__

(self, x: mx.array)

src/tiny_llm_ref/qwen3_week1.py:95

Method__call__

( self, x: mx.array, mask: mx.array | str | None = None, )

src/tiny_llm_ref/qwen3_week1.py:145

Method__call__

( self, inputs: mx.array, )

src/tiny_llm_ref/qwen3_week1.py:211

Method__call__

(self, x: mx.array)

src/tiny_llm_ref/layer_norm.py:10

Method__call__

( self, x: mx.array, offsets: int | list[int] | mx.array, cache: TinyKvCache,

src/tiny_llm_ref/qwen3_week3.py:48

Method__call__

(self, x: mx.array)

src/tiny_llm_ref/qwen3_week3.py:112

Method__call__

( self, x: mx.array, offset: int | list[int] | mx.array, cache: TinyKvCache,

src/tiny_llm_ref/qwen3_week3.py:162

Method__call__

( self, inputs: mx.array, offset: int | list[int] | mx.array, cache: list[Tiny

src/tiny_llm_ref/qwen3_week3.py:291

Method__call__

(self, x: mx.array)

src/tiny_llm_ref/embedding.py:17

Method__call__

(self, x: mx.array)

src/tiny_llm_ref/embedding.py:35

Method__init__

( self, hidden_size: int, num_heads: int, num_kv_heads: int, head_dim:

src/tiny_llm/qwen3_week2.py:13

Method__init__

( self, dim: int, hidden_dim: int, w_gate: QuantizedWeights, w_up: Qua

src/tiny_llm/qwen3_week2.py:43

Method__init__

( self, num_attention_heads: int, num_kv_heads: int, hidden_size: int,

src/tiny_llm/qwen3_week2.py:58

Method__init__

( self, mlx_model: Any, enable_flash_attn: bool = False, )

src/tiny_llm/qwen3_week2.py:94

Method__init__

(self, page_size: int = 128)

src/tiny_llm/paged_kv_cache.py:20

Method__init__

(self, pool: TinyKvPagedPool)

src/tiny_llm/paged_kv_cache.py:51

Method__init__

( self, model: any, tokenizer: TokenizerWrapper, prompt: str, prefill_

src/tiny_llm/batch.py:19

Method__init__

( self, w_router: QuantizedWeights, w_gate: QuantizedWeights, w_up: QuantizedW

src/tiny_llm/moe.py:24

Method__init__

( self, dims: int, seq_len: int, base: int = 10000, traditional: bool

src/tiny_llm/positional_encoding.py:5

Method__init__

( self, hidden_size: int, num_heads: int, wq: mx.array, wk: mx.array,

src/tiny_llm/attention.py:16

Method__init__

( self, hidden_size: int, num_heads: int, num_kv_heads: int, head_dim:

src/tiny_llm/qwen3_week1.py:12

Method__init__

( self, dim: int, hidden_dim: int, w_gate: mx.array, w_up: mx.array,

src/tiny_llm/qwen3_week1.py:39

Method__init__

( self, num_attention_heads: int, num_kv_heads: int, hidden_size: int,

src/tiny_llm/qwen3_week1.py:54

Method__init__

(self, mlx_model: Any)

src/tiny_llm/qwen3_week1.py:87

Method__init__

( self, scales: mx.array, biases: mx.array, group_size: int, bits: int

src/tiny_llm/quantize.py:17

Method__init__

(self, max_active_requests: int, max_seq_len: int)

src/tiny_llm/kv_cache.py:34

Method__init__

(self)

src/tiny_llm/kv_cache.py:55

Method__init__

(self, dim: int, weight: mx.array, eps: float = 1e-5)

src/tiny_llm/layer_norm.py:5

Method__init__

( self, mlx_model: Any, page_size: int = 128, )

src/tiny_llm/qwen3_week3.py:12

Method__init__

(self, vocab_size: int, embedding_dim: int, weight: mx.array)

src/tiny_llm/embedding.py:6

Method__init__

(self, vocab_size: int, embedding_dim: int, weight: QuantizedWeights)

src/tiny_llm/embedding.py:17

Method__init__

( self, hidden_size: int, num_heads: int, num_kv_heads: int, head_dim:

src/tiny_llm_ref/qwen3_week2.py:16

Method__init__

( self, dim: int, hidden_dim: int, w_gate: QuantizedWeights, w_up: Qua

src/tiny_llm_ref/qwen3_week2.py:108

Method__init__

( self, num_attention_heads: int, num_kv_heads: int, hidden_size: int,

src/tiny_llm_ref/qwen3_week2.py:130

Method__init__

( self, mlx_model: Any, enable_flash_attn: bool = False, )

src/tiny_llm_ref/qwen3_week2.py:192

Method__init__

(self, page_size: int = 128)

src/tiny_llm_ref/paged_kv_cache.py:27

Method__init__

(self, pool: TinyKvPagedPool)

src/tiny_llm_ref/paged_kv_cache.py:137

Method__init__

( self, model: any, tokenizer: TokenizerWrapper, prompt: str, prefill_

src/tiny_llm_ref/batch.py:19

Method__init__

( self, w_router: QuantizedWeights, w_gate: QuantizedWeights, w_up: QuantizedW

src/tiny_llm_ref/moe.py:52

Method__init__

( self, dims: int, seq_len: int, base: int = 10000, traditional: bool

src/tiny_llm_ref/positional_encoding.py:5

Method__init__

( self, hidden_size: int, num_heads: int, wq: mx.array, wk: mx.array,

src/tiny_llm_ref/attention.py:158

Method__init__

( self, hidden_size: int, num_heads: int, num_kv_heads: int, head_dim:

src/tiny_llm_ref/qwen3_week1.py:12

Method__init__

( self, dim: int, hidden_dim: int, w_gate: mx.array, w_up: mx.array,

src/tiny_llm_ref/qwen3_week1.py:81

Method__init__

( self, num_attention_heads: int, num_kv_heads: int, hidden_size: int,

src/tiny_llm_ref/qwen3_week1.py:100

Method__init__

( self, mlx_model: Any, )

src/tiny_llm_ref/qwen3_week1.py:158

Method__init__

( self, scales: mx.array, biases: mx.array, group_size: int, bits: int

src/tiny_llm_ref/quantize.py:7

Method__init__

(self, max_active_requests: int, max_seq_len: int)

src/tiny_llm_ref/kv_cache.py:67

Method__init__

(self)

src/tiny_llm_ref/kv_cache.py:213

Method__init__

(self, dim: int, weight: mx.array, eps: float = 1e-5)

src/tiny_llm_ref/layer_norm.py:5

Method__init__

( self, hidden_size: int, num_heads: int, num_kv_heads: int, head_dim:

src/tiny_llm_ref/qwen3_week3.py:16

Functions344 in github.com/skyzh/tiny-llm