MCPcopy Index your code

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

Functions344 in github.com/skyzh/tiny-llm

Method__init__
( self, dim: int, hidden_dim: int, w_gate: QuantizedWeights, w_up: Qua
src/tiny_llm_ref/qwen3_week3.py:98
Method__init__
( self, num_attention_heads: int, num_kv_heads: int, hidden_size: int,
src/tiny_llm_ref/qwen3_week3.py:120
Method__init__
( self, mlx_model: Any, page_size: int = 128, )
src/tiny_llm_ref/qwen3_week3.py:185
Method__init__
( self, vocab_size: int, embedding_dim: int, weight: mx.array, )
src/tiny_llm_ref/embedding.py:7
Method__init__
( self, vocab_size: int, embedding_dim: int, weight: QuantizedWeights, )
src/tiny_llm_ref/embedding.py:25
Function_step
(model, y)
src/tiny_llm/generate.py:14
Methodallocate_page
(self)
src/tiny_llm/paged_kv_cache.py:31
Methodas_linear
(self, x: mx.array)
src/tiny_llm/embedding.py:23
Methodas_linear
(self, x: mx.array)
src/tiny_llm_ref/embedding.py:45
Functionassert_allclose
( a: mx.array, b: mx.array, precision: mx.Dtype, rtol: float | None = None, atol: float |
tests/utils.py:11
Functionbatch_generate
( model: any, tokenizer: TokenizerWrapper, prompts: list[str], max_seq_len=512, batch_size
src/tiny_llm/batch.py:99
Methodblock_table
(self, max_pages: int | None = None)
src/tiny_llm/paged_kv_cache.py:70
Methodcontext_lens
(self)
src/tiny_llm/paged_kv_cache.py:73
Methodcreate_kv_cache
(self)
src/tiny_llm_ref/qwen3_week3.py:284
Methoddecode_done
(self, token, update_offset=True)
src/tiny_llm/batch.py:50
Functiondispatch_model
(model_name: str, mlx_model, week: int, **kwargs)
src/tiny_llm/models.py:22
Functiondispatch_model
(model_name: str, mlx_model, week: int, **kwargs)
src/tiny_llm_ref/models.py:22
Methodfree_page
(self, page_id: int)
src/tiny_llm/paged_kv_cache.py:46
Methodget_seq_len
(data)
src/tiny_llm_ref/kv_cache.py:102
Methodkey_values
(self)
src/tiny_llm_ref/paged_kv_cache.py:149
Functionmake_sampler
(temp: float, top_p: float, top_k: int | None)
src/tiny_llm/sampler.py:5
Functionmake_sampler
(temp: float, top_p: float, top_k: int | None)
src/tiny_llm_ref/sampler.py:5
Functionnp_type_to_mx_type
(np_type: np.dtype)
benches/utils.py:49
Functionnp_type_to_mx_type
(np_type: np.dtype)
tests_refsol/utils.py:49
Functionnp_type_to_mx_type
(np_type: np.dtype)
tests/utils.py:49
Methodnum_free_pages
(self)
src/tiny_llm/paged_kv_cache.py:28
Methodnum_free_pages
(self)
src/tiny_llm_ref/paged_kv_cache.py:41
Methodnum_pages
(self)
src/tiny_llm/paged_kv_cache.py:24
Methodnum_pages
(self)
src/tiny_llm/paged_kv_cache.py:55
Methodnum_pages
(self)
src/tiny_llm_ref/paged_kv_cache.py:37
Methodnum_pages
(self)
src/tiny_llm_ref/paged_kv_cache.py:145
Methodpaged_metadata
( self, max_pages: int | None = None, mask: mx.array | str | None = None, )
src/tiny_llm/paged_kv_cache.py:76
Functionquantized_linear
( x: mx.array, w: QuantizedWeights, bias: mx.array | None = None, )
src/tiny_llm/quantize.py:54
Functionqwen3_0_6b_model_exists
()
benches/utils.py:58
Functionqwen3_0_6b_model_exists
()
tests_refsol/utils.py:58
Functionqwen3_0_6b_model_exists
()
tests/utils.py:58
Functionqwen3_1_7b_model_exists
()
benches/utils.py:69
Functionqwen3_1_7b_model_exists
()
tests_refsol/utils.py:69
Functionqwen3_1_7b_model_exists
()
tests/utils.py:69
Functionqwen3_4b_model_exists
()
benches/utils.py:80
Functionqwen3_4b_model_exists
()
tests_refsol/utils.py:80
Functionqwen3_4b_model_exists
()
tests/utils.py:80
Methodrewind
(self, n: int)
src/tiny_llm_ref/paged_kv_cache.py:258
Functionsample
(logprobs: mx.array)
src/tiny_llm/sampler.py:6
Functionsample
(logprobs: mx.array)
src/tiny_llm_ref/sampler.py:6
Functionsimple_generate
( model: Qwen3ModelWeek1, tokenizer: TokenizerWrapper, prompt: str, sampler: Callable[[mx.arra
src/tiny_llm/generate.py:8
Functionsimple_generate_with_kv_cache
( model: Qwen3ModelWeek2, tokenizer: TokenizerWrapper, prompt: str )
src/tiny_llm/generate.py:18
Functionspeculative_generate
( draft_model: Qwen3ModelWeek2, model: Qwen3ModelWeek2, draft_tokenizer: TokenizerWrapper, tok
src/tiny_llm/generate.py:25
Functiontest_mlx_attention
(benchmark)
benches/test_attention.py:18
Functiontest_mlx_quantized_matmul
(benchmark)
benches/test_quantized_matmul.py:20
Functiontest_refsol_attention
(benchmark)
benches/test_attention.py:27
Functiontest_refsol_flash_attention
(benchmark)
benches/test_attention.py:38
Functiontest_refsol_quantized_matmul
(benchmark)
benches/test_quantized_matmul.py:31
Functiontest_rope_week2_batch_offset
( stream: mx.Stream, traditional: bool, precision: mx.Dtype )
tests_refsol/test_rope.py:10
Functiontest_task_1
()
tests_refsol/test_week_1_day_6.py:5
Functiontest_task_1
()
tests_refsol/test_week_1_day_7.py:5
Functiontest_task_1_attention_with_mask_cpu
()
tests_refsol/test_week_2_day_6.py:157
Functiontest_task_1_attention_with_mask_cpu_large
()
tests_refsol/test_week_2_day_6.py:161
Functiontest_task_1_attention_with_mask_cpu_small
()
tests_refsol/test_week_2_day_6.py:153
Functiontest_task_1_attention_with_mask_gpu
()
tests_refsol/test_week_2_day_6.py:173
Functiontest_task_1_attention_with_mask_gpu_extra_small
()
tests_refsol/test_week_2_day_6.py:165
Functiontest_task_1_attention_with_mask_gpu_large
()
tests_refsol/test_week_2_day_6.py:177
Functiontest_task_1_attention_with_mask_gpu_small
()
tests_refsol/test_week_2_day_6.py:169
Functiontest_task_1_flash_attention_with_mask_cpu
()
tests_refsol/test_week_2_day_6.py:129
Functiontest_task_1_flash_attention_with_mask_cpu_large
()
tests_refsol/test_week_2_day_6.py:133
Functiontest_task_1_flash_attention_with_mask_cpu_small
()
tests_refsol/test_week_2_day_6.py:125
Functiontest_task_1_flash_attention_with_mask_gpu
()
tests_refsol/test_week_2_day_6.py:145
Functiontest_task_1_flash_attention_with_mask_gpu_extra_small
()
tests_refsol/test_week_2_day_6.py:137
Functiontest_task_1_flash_attention_with_mask_gpu_large
()
tests_refsol/test_week_2_day_6.py:149
Functiontest_task_1_flash_attention_with_mask_gpu_small
()
tests_refsol/test_week_2_day_6.py:141
Functiontest_task_1_grouped_attention
( stream: mx.Stream, precision: mx.Dtype, batch_dimension: int, scale: float | None )
tests_refsol/test_week_1_day_3.py:66
Functiontest_task_1_grouped_expert_linear
()
tests_refsol/test_week_3_day_3.py:19
Functiontest_task_1_model_kv_caches_share_layer_pools
()
tests_refsol/test_week_3_day_1.py:183
Functiontest_task_1_model_layer_caches_keep_independent_page_metadata
()
tests_refsol/test_week_3_day_1.py:199
Functiontest_task_1_paged_attention_matches_dense_flash_attention
()
tests_refsol/test_week_3_day_2.py:105
Functiontest_task_1_paged_cache_matches_full_cache
()
tests_refsol/test_week_3_day_1.py:102
Functiontest_task_1_paged_cache_rewind
()
tests_refsol/test_week_3_day_1.py:153
Functiontest_task_1_paged_pool_reuses_freed_pages
()
tests_refsol/test_week_3_day_1.py:125
Functiontest_task_1_rms_norm
( stream: mx.Stream, precision: mx.Dtype, )
tests_refsol/test_week_1_day_4.py:11
Functiontest_task_1_rms_norm_cast_to_float32
(stream: mx.Stream)
tests_refsol/test_week_1_day_4.py:32
Functiontest_task_1_rope_mlx_traditional
( stream: mx.Stream, with_offset: bool, precision: mx.Dtype )
tests_refsol/test_week_1_day_2.py:58
Functiontest_task_1_rope_multiple_offsets
( stream: mx.Stream, traditional: bool, precision: mx.Dtype )
tests_refsol/test_week_2_day_6.py:48
Functiontest_task_1_simple_attention
Test if `scaled_dot_product_attention_simple` can process Q/K/V correctly. We assume Q/K/V are of the same dimensions and test different batc
tests_refsol/test_week_1_day_1.py:26
Functiontest_task_1_simple_attention_scale_mask
Test if `scaled_dot_product_attention_simple` can process scale and mask correctly.
tests_refsol/test_week_1_day_1.py:69
Functiontest_task_1_softmax
(stream: mx.Stream, precision: mx.Dtype)
tests_refsol/test_week_1_day_1.py:10
Functiontest_task_1_transformer_block
( stream: mx.Stream, precision: mx.Dtype, mask: str | None )
tests_refsol/test_week_1_day_5.py:10
Functiontest_task_2_batched_paged_attention_matches_dense_attention
()
tests_refsol/test_week_3_day_2.py:140
Functiontest_task_2_batching_kv_cache
()
tests_refsol/test_week_2_day_6.py:181
Functiontest_task_2_embedding_as_linear
()
tests_refsol/test_week_1_day_5.py:144
Functiontest_task_2_embedding_call
()
tests_refsol/test_week_1_day_5.py:127
Functiontest_task_2_flash_attention_cpu
(mask_mode: str)
tests_refsol/test_week_2_day_4.py:133
Functiontest_task_2_flash_attention_cpu_causal_mask_faster_than_all_zero_mask
()
tests_refsol/test_week_2_day_4.py:142
Functiontest_task_2_flash_attention_cpu_large
(mask_mode: str)
tests_refsol/test_week_2_day_4.py:138
Functiontest_task_2_flash_attention_cpu_small
(mask_mode: str)
tests_refsol/test_week_2_day_4.py:128
Functiontest_task_2_grouped_attention_causal_mask
( stream: mx.Stream, precision: mx.Dtype, batch_dimension: int, scale: float | None )
tests_refsol/test_week_1_day_3.py:128
Functiontest_task_2_linear
(stream: mx.Stream, precision: mx.Dtype)
tests_refsol/test_week_1_day_1.py:113
Functiontest_task_2_mask_only_different_dim
( stream: mx.Stream, )
tests_refsol/test_week_1_day_3.py:98
Functiontest_task_2_mask_only_same_dim
( stream: mx.Stream, )
tests_refsol/test_week_1_day_3.py:73
Functiontest_task_2_quantized_matmul_complex_bf16_cpu
()
tests_refsol/test_week_2_day_2.py:53
Functiontest_task_2_quantized_matmul_complex_f16_cpu
()
tests_refsol/test_week_2_day_2.py:61
← previousnext →201–300 of 344, ranked by callers