MCPcopy Create free account

hub / github.com/ScalingIntelligence/tokasaurus / functions

Functions507 in github.com/ScalingIntelligence/tokasaurus

⨍ Functions 507◇ Types & classes 95↳ Endpoints 10
↓ 30 callersMethodget
(self, sid: str)
tokasaurus/manager/types.py:167
↓ 28 callersMethodsanity_checks
(self, seq_ids: set[str] | None = None)
tokasaurus/manager/allocator.py:158
↓ 28 callersMethodto
(self, device: DeviceType, non_blocking: bool = False)
tokasaurus/model/types.py:235
↓ 18 callersMethodallocate_with_prefix_match
( self, seq_id: str, prompt_ids: list[int], num_reserved_blocks: int = 0,
tokasaurus/manager/allocator.py:324
↓ 17 callersMethodadd
(self, value: float)
tokasaurus/manager/stopping_predictor.py:59
↓ 14 callersMethodrun
( self, input_batch_state: BatchState, non_blocking: bool = False, )
tokasaurus/model/utils.py:840
↓ 13 callersMethodallocate_up_to_length
( self, seq_id: str, kv_indices: list[int], length: int, available_lea
tokasaurus/manager/allocator.py:265
↓ 13 callersMethodempty
(self)
tokasaurus/manager/stopping_predictor.py:56
↓ 13 callersMethodprompt_total
(self)
tokasaurus/manager/types.py:67
↓ 12 callersMethodfrom_pretrained
( cls, model_name_or_path: str, extra_config: ExtraModelConfig | None = None,
tokasaurus/model/llama.py:774
↓ 11 callersMethodremove
(self, sid: str)
tokasaurus/manager/types.py:207
↓ 10 callersMethodwait
(self)
tokasaurus/common_types.py:18
↓ 8 callersFunctionall_gather
(x: Tensor, extra_config: ExtraModelConfig)
tokasaurus/model/llama.py:56
↓ 8 callersMethodmean
(self)
tokasaurus/manager/stopping_predictor.py:141
↓ 8 callersMethodprompt_to_schedule
(self)
tokasaurus/manager/types.py:73
↓ 8 callersMethodtotal_scheduled
(self)
tokasaurus/manager/types.py:70
↓ 7 callersMethod__init__
Taken from LlamaRMSNorm.
tokasaurus/model/llama.py:36
↓ 7 callersMethodexpected_completion_length
(self, add_buffer: bool = False)
tokasaurus/manager/types.py:76
↓ 7 callersFunctionfind_free_port
()
tokasaurus/utils.py:312
↓ 7 callersMethodfree_and_update
(self, seq_id: str, kv_indices: list[int], token_ids: list[int])
tokasaurus/manager/allocator.py:387
↓ 7 callersMethodkv_cache_num_blocks
(self)
tokasaurus/common_types.py:130
↓ 7 callersMethodplan
(self, input_batch_state: BatchState, non_blocking: bool = False)
tokasaurus/model/utils.py:827
↓ 6 callersFunctionlprint
(*args, **kwargs)
tokasaurus/benchmarks/bench_model.py:114
↓ 6 callersFunctionmake_messages
(word: str)
tests/test_basic.py:174
↓ 6 callersMethodset_wrappers
(self, wrappers: WrapperCollection)
tokasaurus/model/llama.py:669
↓ 6 callersMethodto
(self, device: DeviceType | None = None, dtype: torch.dtype | None = None)
tokasaurus/model/llama.py:766
↓ 5 callersMethodadd_sequence
( self, kv_indices: list[int], kv_seq_len: int, num_qtokens: int, page
tokasaurus/model/types.py:122
↓ 5 callersFunctioncalc_block_usage_over_time
For each sequence, calculate how many kv blocks will be used right before the seq is deallocated and after all earlier-finishing sequences ha
tokasaurus/manager/scheduler.py:47
↓ 5 callersFunctionformat_value
(value: int)
tokasaurus/manager/monitoring.py:256
↓ 5 callersMethodhead_dim
(self)
tokasaurus/model/llama.py:158
↓ 5 callersMethodin_decoding
(self, sid: str)
tokasaurus/manager/types.py:217
↓ 5 callersMethodis_initialized
(self)
tokasaurus/manager/monitoring.py:252
↓ 5 callersMethodis_warmed_up
(self)
tokasaurus/manager/stopping_predictor.py:285
↓ 5 callersFunctionlaunch_server
(config: BaseConfig)
tokasaurus/benchmarks/utils.py:101
↓ 5 callersFunctionnowstamp
()
tokasaurus/server/types.py:15
↓ 5 callersFunctionpad_and_slice_tensor
( x: Tensor, num_padding: int, tp_rank: int, tp_size: int )
tokasaurus/model/utils.py:104
↓ 5 callersMethodrunning_seqs
(self)
tokasaurus/manager/types.py:238
↓ 5 callersFunctionsetup_logging
(config: "ServerConfig")
tokasaurus/utils.py:460
↓ 5 callersFunctiontry_onboarding_seqs
( block_usage: BlockUsageOverTime, seqs: list[Sequence], existing_prefill_seqs: list[Sequence],
tokasaurus/manager/scheduler.py:408
↓ 4 callersFunctionadd_decoding_ids_to_batch_state
( input_batch_state: BatchState, decoding_input_ids: Tensor, tp_rank: int = 0, tp_size: int =
tokasaurus/model/utils.py:205
↓ 4 callersMethodadd_sequence
(self, temperature: float, top_p: float)
tokasaurus/model/types.py:302
↓ 4 callersFunctioncalc_kv_token_indices
( kv_block_indices: list[int], page_size: int, start_idx: int, num_tokens: int )
tokasaurus/manager/input_building.py:138
↓ 4 callersMethodclient
(self)
tokasaurus/benchmarks/utils.py:39
↓ 4 callersFunctionconvert_unit
(milis: float, unit: str)
tokasaurus/utils.py:188
↓ 4 callersMethoddeallocate
(self, seq: Sequence)
tokasaurus/manager/types.py:324
↓ 4 callersMethodelapsed_time
(self)
tokasaurus/manager/monitoring.py:136
↓ 4 callersMethodexpected_num_additional_blocks
(self, page_size: int, add_buffer: bool = False)
tokasaurus/manager/types.py:94
↓ 4 callersFunctionmake_decision
(num_prefill: float)
tokasaurus/manager/scheduler.py:601
↓ 4 callersFunctionmake_input_batch_state
( inp: ModelInput, tp_rank: int = 0, tp_size: int = 1, pp_rank: int = 0, pp_size: int = 1,
tokasaurus/model/utils.py:113
↓ 4 callersFunctionmake_random_ids
(length: int)
tests/test_scheduler.py:33
↓ 4 callersFunctionmaybe_save_results
(config: BaseConfig, results)
tokasaurus/benchmarks/utils.py:214
↓ 4 callersFunctionmove_batch_state
( input_batch_state: BatchState, device: DeviceType, non_blocking: bool = False, )
tokasaurus/model/utils.py:227
↓ 4 callersMethodnum_prefill_tokens
(self)
tokasaurus/manager/types.py:280
↓ 4 callersMethodnum_running_seqs
(self)
tokasaurus/manager/types.py:249
↓ 4 callersMethodnum_unfinished_seqs
(self)
tokasaurus/manager/types.py:244
↓ 4 callersFunctionparallelize
( fn, items, num_workers: int | None = None, processes: bool = True, allow_unordered: bool
tokasaurus/benchmarks/utils.py:127
↓ 4 callersMethodremove_decoding
(self, sid: str)
tokasaurus/manager/types.py:198
↓ 4 callersMethodremove_queued
(self, sid: str)
tokasaurus/manager/types.py:204
↓ 4 callersFunctionserver_manager
(config: ServerConfig, finalize=True)
tokasaurus/entry.py:91
↓ 4 callersFunctionshuffle_and_limit
(ds, config: BaseConfig)
tokasaurus/benchmarks/utils.py:182
↓ 4 callersMethodupdate_seq_predictions
(self, seq: "Sequence")
tokasaurus/manager/stopping_predictor.py:190
↓ 3 callersMethodbatch_size
(self)
tokasaurus/manager/types.py:283
↓ 3 callersFunctionencode_array
(array: np.ndarray)
tokasaurus/server/utils.py:642
↓ 3 callersFunctiongenerate_output
( state: ServerState, request: CompletionsRequest | ChatCompletionRequest )
tokasaurus/server/utils.py:831
↓ 3 callersFunctionget_global_rank
parallelism order from outer to inner: dp -> pp -> tp
tokasaurus/model/utils.py:47
↓ 3 callersMethodhead_dim
(self)
tokasaurus/model/qwen3.py:42
↓ 3 callersMethodin_prefilling
(self, sid: str)
tokasaurus/manager/types.py:220
↓ 3 callersMethodin_queued
(self, sid: str)
tokasaurus/manager/types.py:223
↓ 3 callersFunctionmake_model
( config: ServerConfig, device: str, dtype: torch.dtype, pp_rank: int = 0, tp_rank: int =
tokasaurus/model/utils.py:330
↓ 3 callersFunctionmake_sequences
( num_shared_decoding_seqs: int = 128, num_unique_decoding_seqs: int = 128, num_prefilling_seqs: i
tests/test_scheduler.py:16
↓ 3 callersMethodnum_decoding_tokens
(self)
tokasaurus/manager/types.py:277
↓ 3 callersMethodnum_lm_head_tokens
(self)
tokasaurus/model/types.py:366
↓ 3 callersMethodnum_prefill_tokens
(self)
tokasaurus/model/types.py:360
↓ 3 callersFunctionprocess_chat_completions_output
( state: ServerState, crequest: ChatCompletionRequest, request: TokasaurusRequest, output: Req
tokasaurus/server/utils.py:694
↓ 3 callersFunctionreal_allocate
( state: ManagerState, seq: Sequence, available_leaf_heap: list[PrefixTreeBlock], )
tokasaurus/manager/manager.py:451
↓ 3 callersFunctionreduce_scatter
(x: Tensor, extra_config: ExtraModelConfig)
tokasaurus/model/llama.py:75
↓ 3 callersFunctionreorder_decoding_seqs_for_hydragen
Our Hydragen implementation requires us to reorder the decoding sequences so that seqs in the same shared-prefix group are adjacent to each o
tokasaurus/manager/hydragen.py:10
↓ 3 callersFunctionset_async_tp_enabled
(enabled: bool)
tokasaurus/model/utils.py:383
↓ 3 callersMethodtotal_tokens
(self)
tokasaurus/benchmarks/bench_model.py:69
↓ 2 callersMethodadd_prefilling
(self, seq: Sequence)
tokasaurus/manager/types.py:192
↓ 2 callersMethodassign_block_to_seq
(self, block: PrefixTreeBlock, seq_id: str)
tokasaurus/manager/allocator.py:242
↓ 2 callersMethodattn_fn
( ragged_q: Tensor, ragged_k: Tensor, ragged_v: Tensor, k_cach
tokasaurus/model/llama.py:166
↓ 2 callersMethodbuffer_len
(self)
tokasaurus/manager/stopping_predictor.py:282
↓ 2 callersMethodbuffer_mean
(self)
tokasaurus/manager/stopping_predictor.py:276
↓ 2 callersMethodbuild
Convert builders to tensors
tokasaurus/model/types.py:179
↓ 2 callersMethodcalc_conditional_mean_stds
For each v in vals, calculate the mean and std of the buffer values that are greater than or equal to v.
tokasaurus/manager/stopping_predictor.py:63
↓ 2 callersMethodcalc_padding
returns (all_tokens_padding, lm_head_tokens_padding)
tokasaurus/model/utils.py:790
↓ 2 callersFunctioncalc_pipeline_layer_start_and_end
( config: LlamaConfig, extra_config: ExtraModelConfig, )
tokasaurus/model/llama.py:585
↓ 2 callersFunctioncleanup_processes
(processes: list[mp.Process])
tokasaurus/entry.py:13
↓ 2 callersFunctioncreate_workspace_buffer
(device: DeviceType)
tokasaurus/model/attention_utils.py:16
↓ 2 callersFunctiondecode_completion
( state: ServerState, request: TokasaurusRequest, output: RequestOutput )
tokasaurus/server/utils.py:380
↓ 2 callersMethoddecoding_batch_indices
(self)
tokasaurus/model/types.py:372
↓ 2 callersFunctionextract_answer_gsm8k
(completion)
tokasaurus/benchmarks/monkeys_gsm8k.py:58
↓ 2 callersFunctionextract_answer_gsm8k
(completion)
tokasaurus/benchmarks/standalone_monkeys_gsm8k.py:276
↓ 2 callersFunctionfree_seq
(seq: Sequence)
tests/test_scheduler.py:126
↓ 2 callersFunctionget_dtype
(dtype_str: str)
tokasaurus/model/utils.py:34
↓ 2 callersFunctionget_eos_token_ids
(generation_config: GenerationConfig)
tokasaurus/utils.py:447
↓ 2 callersFunctiongroup_for_hydragen
Iterative version of depth-first search - we make a group for a given prefix if it meets the minimum group size/prefix length requirements,
tokasaurus/manager/hydragen.py:66
↓ 2 callersFunctionis_local
()
tokasaurus/utils.py:103
↓ 2 callersMethodis_root
(self)
tokasaurus/manager/allocator.py:20
next →1–100 of 507, ranked by callers