Functions — github.com/skyzh/tiny-llm

hub / github.com/skyzh/tiny-llm / functions

Functions344 in github.com/skyzh/tiny-llm

⨍ Functions 344◇ Types & classes 51

↓ 46 callersFunctionassert_allclose

( a: mx.array, b: mx.array, precision: mx.Dtype, rtol: float | None = None, atol: float |

tests_refsol/utils.py:11

↓ 22 callersMethodfrom_mlx_layer

(mlx_layer: Any)

src/tiny_llm_ref/quantize.py:22

↓ 18 callersFunctionquantized_linear

( x: mx.array, w: QuantizedWeights, bias: mx.array | None = None, )

src/tiny_llm_ref/quantize.py:32

↓ 14 callersFunctionattention_helper

( stream: mx.Stream, H_q, H, L, E, S, BATCH, use_flash_attention: bool = False )

tests_refsol/test_week_2_day_6.py:54

↓ 13 callersFunctionlinear

( x: mx.array, w: mx.array, bias: mx.array | None = None, )

src/tiny_llm_ref/basics.py:10

↓ 9 callersFunctiondequantize_linear

(mx_layer: Any)

src/tiny_llm_ref/quantize.py:50

↓ 8 callersFunction_quantized_layer

( out_dim: int, in_dim: int, group_size: int = 128 )

tests_refsol/test_week_3_day_1.py:23

↓ 8 callersFunction_quantized_layer

( out_dim: int, in_dim: int, group_size: int = 128 )

tests_refsol/test_week_3_day_2.py:26

↓ 8 callersFunctionquantized_matmul_helper

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

tests_refsol/test_week_2_day_2.py:6

↓ 8 callersFunctionsafe_div

(num: float, den: float)

bench.py:346

↓ 7 callersFunctionattention_helper

(stream: mx.Stream, H_q, H, L, E, S, BATCH, mask_mode: str)

tests_refsol/test_week_2_day_4.py:8

↓ 7 callersFunctionflash_attention

( query: mx.array, key: mx.array, value: mx.array, scale: float | None = None, mask: mx.ar

src/tiny_llm/attention.py:51

↓ 7 callersMethodfrom_mlx_layer

(mlx_layer: Any)

src/tiny_llm/quantize.py:32

↓ 7 callersMethodupdate_and_fetch

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

src/tiny_llm/paged_kv_cache.py:61

↓ 6 callersFunction_random_chunk

( length: int, num_heads: int = 2, head_dim: int = 4 )

tests_refsol/test_week_3_day_2.py:18

↓ 6 callersMethodadd_request

(self, prefilled: TinyKvCache, id: int)

src/tiny_llm/kv_cache.py:47

↓ 6 callersMethodcreate_kv_cache

(self)

src/tiny_llm_ref/qwen3_week2.py:269

↓ 5 callersFunction_step

(model, y)

src/tiny_llm_ref/generate.py:20

↓ 5 callersFunctionassert_allclose

( a: mx.array, b: mx.array, precision: mx.Dtype, rtol: float | None = None, atol: float |

benches/utils.py:11

↓ 5 callersMethodcreate_kv_cache

(self)

src/tiny_llm/qwen3_week3.py:20

↓ 5 callersFunctionscaled_dot_product_attention_grouped

( query: mx.array, key: mx.array, value: mx.array, scale: float | None = None, mask: mx.ar

src/tiny_llm/attention.py:41

↓ 4 callersFunction_random_chunk

( length: int, num_heads: int = 2, head_dim: int = 4 )

tests_refsol/test_week_3_day_1.py:15

↓ 4 callersMethodgather_dense

(self)

src/tiny_llm/paged_kv_cache.py:58

↓ 4 callersFunctionsilu

(x: mx.array)

src/tiny_llm_ref/basics.py:21

↓ 4 callersFunctiontime_flash_attention

( stream: mx.Stream, query: mx.array, key: mx.array, value: mx.array, scale: float, ma

tests_refsol/test_week_2_day_4.py:46

↓ 4 callersMethodupdate_and_fetch

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

src/tiny_llm/kv_cache.py:59

↓ 3 callersFunction_fake_qwen3_mlx_model

()

tests_refsol/test_week_3_day_1.py:39

↓ 3 callersFunction_release_kv_cache

(kv_cache)

src/tiny_llm_ref/generate.py:9

↓ 3 callersMethodas_linear

(self, x: mx.array)

src/tiny_llm_ref/embedding.py:20

↓ 3 callersFunctioncausal_mask

(L: int, S: int, dtype: mx.Dtype)

src/tiny_llm_ref/attention.py:24

↓ 3 callersFunctionget_test_attention_data

()

benches/test_attention.py:8

↓ 3 callersFunctiongrouped_expert_linear

( x: mx.array, w_experts: QuantizedWeights, expert_ids: mx.array, )

src/tiny_llm_ref/moe.py:7

↓ 3 callersFunctionhelper_test_task_3

Tests for continuous batching of decode requests.

tests_refsol/test_week_2_day_6.py:256

↓ 3 callersFunctionhelper_test_task_3

(model_name: str, iters: int = 10)

tests_refsol/test_week_1_day_5.py:110

↓ 3 callersFunctionhelper_test_task_3

(model_name: str, iters: int = 10)

tests_refsol/test_week_2_day_1.py:33

↓ 3 callersFunctionhelper_test_task_4

( model_name: str, seq_len: int, iters: int = 1, )

tests_refsol/test_week_2_day_1.py:69

↓ 3 callersFunctionpaged_attention

Paged attention backed by the C++/Metal extension. The Python wrapper keeps the model-facing shape as [B, H_q, L, D], while the extensio

src/tiny_llm_ref/attention.py:69

↓ 3 callersMethodread_page

(self, page_id: int)

src/tiny_llm/paged_kv_cache.py:34

↓ 3 callersMethodrelease

Release all resources owned by this cache. Request-scoped caches use this when generation finishes or a batch slot is remove

src/tiny_llm_ref/kv_cache.py:30

↓ 3 callersFunctionsample_next_week2

(model, y: mx.array, offset: int, kv_cache: list)

bench.py:151

↓ 2 callersMethod_append_chunk

(self, key: mx.array, value: mx.array)

src/tiny_llm_ref/paged_kv_cache.py:154

↓ 2 callersFunction_prefill

(model, tokenizer, prompt, kv_cache)

src/tiny_llm_ref/generate.py:104

↓ 2 callersFunction_print_progress

( requests: list[Request | None], pending_prefill_request: Request | None, queue_size: int, pr

src/tiny_llm/batch.py:62

↓ 2 callersFunction_print_progress

( requests: list[Request | None], pending_prefill_request: Request | None, queue_size: int, pr

src/tiny_llm_ref/batch.py:84

↓ 2 callersFunction_print_text

(text, progress)

src/tiny_llm_ref/generate.py:144

↓ 2 callersFunction_rewind_cache

(kv_cache, revert_len)

src/tiny_llm_ref/generate.py:140

↓ 2 callersFunction_step

(model, y, offsets, kv_cache)

src/tiny_llm_ref/batch.py:9

↓ 2 callersMethodas_linear

(self, x: mx.array)

src/tiny_llm/embedding.py:12

↓ 2 callersFunctionassert_causal_mask_faster_than_all_zero_mask

( stream: mx.Stream, batch: int, h_q: int, h: int, l: int, s: int, e: int, sca

tests_refsol/test_week_2_day_4.py:68

↓ 2 callersFunctioncausal_mask

(L: int, S: int, dtype: mx.Dtype)

src/tiny_llm/attention.py:37

↓ 2 callersFunctioncopy_test

(args, skip_if_exists=False, force=False)

scripts/dev-tools.py:8

↓ 2 callersMethoddecode_done

(self, token, update_offset=True)

src/tiny_llm_ref/batch.py:69

↓ 2 callersFunctiondequantize_linear

(mx_layer: Any)

src/tiny_llm/quantize.py:5

↓ 2 callersFunctiondraft_generate

(model, last_token, offset, kv_cache, num_drafts)

src/tiny_llm_ref/generate.py:127

↓ 2 callersFunctionexport_public_members

(module)

scripts/diff_api.py:9

↓ 2 callersMethodfree_page

(self, page_id: int)

src/tiny_llm_ref/paged_kv_cache.py:120

↓ 2 callersMethodgather_dense

(self)

src/tiny_llm_ref/paged_kv_cache.py:191

↓ 2 callersFunctionget_test_matmul_data

()

benches/test_quantized_matmul.py:8

↓ 2 callersFunctiongrouped_attention_helper

( stream: mx.Stream, precision: mx.Dtype, batch_dimension: int, scale: float | None, is_ca

tests_refsol/test_week_1_day_3.py:7

↓ 2 callersFunctionmedian

(values: list[float])

tests_refsol/test_week_2_day_4.py:63

↓ 2 callersFunctionquantized_matmul

( scales: mx.array, biases: mx.array, group_size: int, bits: int, a: mx.array, b: mx.a

src/tiny_llm/quantize.py:42

↓ 2 callersFunctionquantized_matmul

( scales: mx.array, biases: mx.array, group_size: int, bits: int, a: mx.array, b: mx.a

src/tiny_llm_ref/quantize.py:61

↓ 2 callersMethodremove_request

(self, id: int)

src/tiny_llm/kv_cache.py:50

↓ 2 callersFunctionrope_helper

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

tests_refsol/test_week_1_day_2.py:8

↓ 2 callersFunctionroute_topk

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

src/tiny_llm/moe.py:14

↓ 2 callersFunctionrun_benchmark

( bench_requests: list[BenchRequest], )

bench.py:384

↓ 2 callersFunctionrun_one_request

(request: BenchRequest)

bench.py:369

↓ 2 callersFunctionsample_next_week1

(model, y: mx.array)

bench.py:145

↓ 2 callersFunctionscaled_dot_product_attention_grouped

Potential input of the mask: - mx.array that can broadcast to B * H_q * L * S, which needs to be reshaped to match multi-head dimensions

src/tiny_llm_ref/attention.py:30

↓ 2 callersFunctionscaled_dot_product_attention_simple

( query: mx.array, key: mx.array, value: mx.array, scale: float | None = None, mask: mx.ar

src/tiny_llm/attention.py:5

↓ 2 callersFunctionsoftmax

(x: mx.array, axis: int)

src/tiny_llm_ref/basics.py:5

↓ 2 callersFunctionstringify_member

(members)

scripts/diff_api.py:44

↓ 2 callersMethodtext

(self)

src/tiny_llm_ref/batch.py:80

↓ 2 callersMethodupdate_and_fetch

Update the key-value cache and fetch the updated key-value cache. Args: key: The key to update the cache with.

src/tiny_llm_ref/kv_cache.py:10

↓ 2 callersMethodupdate_and_fetch_paged

Update this cache and return paged attention metadata. Week 3 caches override this. Dense caches intentionally do not provide

src/tiny_llm_ref/kv_cache.py:40

↓ 2 callersMethodwrite_page_slice

( self, page_id: int, start: int, key: mx.array, value: mx.array,

src/tiny_llm_ref/paged_kv_cache.py:94

↓ 1 callersMethod_check_page_chunk

(self, x: mx.array)

src/tiny_llm_ref/paged_kv_cache.py:44

↓ 1 callersFunction_decode_one

(token, tokenizer)

src/tiny_llm_ref/generate.py:120

↓ 1 callersMethod_ensure_page_storage

(self, key: mx.array, value: mx.array)

src/tiny_llm_ref/paged_kv_cache.py:69

↓ 1 callersFunction_fake_qwen3_mlx_model

()

tests_refsol/test_week_3_day_2.py:42

↓ 1 callersFunction_step

(model, y, offsets, kv_cache)

src/tiny_llm/batch.py:9

↓ 1 callersMethodadd_request

(self, prefilled: TinyKvCache, id: int)

src/tiny_llm_ref/kv_cache.py:192

↓ 1 callersMethodallocate_page

(self)

src/tiny_llm_ref/paged_kv_cache.py:48

↓ 1 callersFunctionbatch_generate

( model: any, tokenizer: TokenizerWrapper, prompts: list[str], max_seq_len=512, batch_size

src/tiny_llm_ref/batch.py:121

↓ 1 callersMethodblock_table

(self, max_pages: int | None = None)

src/tiny_llm_ref/paged_kv_cache.py:221

↓ 1 callersFunctionbuild_requests

( *, rng: Random, num_seqs: int, vocab_size: int, eos_token_id: int, min_input_len: in

bench.py:118

↓ 1 callersMethodcontext_lens

(self)

src/tiny_llm_ref/paged_kv_cache.py:228

↓ 1 callersFunctionflash_attention

( query: mx.array, key: mx.array, value: mx.array, scale: float | None = None, mask: mx.ar

src/tiny_llm_ref/attention.py:116

↓ 1 callersFunctiongrouped_expert_linear

( x: mx.array, w_experts: QuantizedWeights, expert_ids: mx.array, )

src/tiny_llm/moe.py:6

↓ 1 callersFunctionis_qwen3_moe_sparse_layer

(args: Any, layer_idx: int)

src/tiny_llm_ref/qwen3_week3.py:176

↓ 1 callersFunctionlinear

( x: mx.array, w: mx.array, bias: mx.array | None = None, )

src/tiny_llm/basics.py:10

↓ 1 callersFunctionload_solution_modules

(solution: str)

bench.py:93

↓ 1 callersFunctionmain

()

bench.py:350

↓ 1 callersFunctionmain

()

scripts/dev-tools.py:51

↓ 1 callersMethodpaged_metadata

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

src/tiny_llm_ref/paged_kv_cache.py:231

↓ 1 callersFunctionparse_args

()

bench.py:29

↓ 1 callersFunctionrandom_token_id

(rng: Random, low: int, high: int, eos_token_id: int)

bench.py:107

↓ 1 callersMethodread_page

(self, page_id: int)

src/tiny_llm_ref/paged_kv_cache.py:59

↓ 1 callersMethodrelease

(self)

src/tiny_llm/paged_kv_cache.py:95

↓ 1 callersMethodrelease

(self)

src/tiny_llm_ref/paged_kv_cache.py:280

next →1–100 of 344, ranked by callers