Functions — github.com/EverMind-AI/MSA

hub / github.com/EverMind-AI/MSA / functions

Functions265 in github.com/EverMind-AI/MSA

⨍ Functions 265◇ Types & classes 54

↓ 15 callersMethodupdate

( self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int,

src/utils/cache.py:152

↓ 8 callersFunction_rag

(name: str)

src/benchmarks.py:60

↓ 7 callersMethodexpect

(q: mp.Queue, constant)

src/types.py:14

↓ 7 callersMethodsend

(q: mp.Queue, constant, data=None, block=True)

src/types.py:20

↓ 7 callersMethodstart

启动显存监控 Returns: bool: 是否成功启动

src/utils/gpu_monitor.py:100

↓ 6 callersFunctionconvert_tensor

转换结构体的 tensor device，如果cuda_device非 None，则将cpu 转换到 cuda，否则将 cuda 转换到 cpu 支持dict, list, tuple, set

src/utils/cache.py:391

↓ 6 callersMethodrecord_kwargs

记录层的元数据信息 Args: layer_idx: 层索引 kwargs: 包含路由层信息的字典

src/utils/cache.py:30

↓ 3 callersMethod_convert_units

转换字节为指定单位

src/utils/gpu_monitor.py:60

↓ 3 callersFunction_run

(command)

src/utils/misc.py:54

↓ 3 callersFunctioneval_recall

(predict, label, at=10)

src/app/benchmark.py:106

↓ 3 callersMethodstop

停止显存监控并返回峰值显存使用量 Args: verbose: 是否打印详细信息 Returns: float: 峰值显存使用量（单位：GB或MB）

src/utils/gpu_monitor.py:140

↓ 2 callersMethod__init__

该 worker 被MemoryWorker创建并仅执行prefill stage 1获取block 的kv cache

src/msa_service.py:272

↓ 2 callersMethod__init__

(self, config: Qwen3Config)

src/msa/model.py:138

↓ 2 callersMethod__init__

(self, _distributed_cache_data=None)

src/utils/cache.py:19

↓ 2 callersMethod_get_current_usage

获取当前显存使用量（按指定单位）

src/utils/gpu_monitor.py:67

↓ 2 callersMethod_get_memory_info

获取当前GPU显存信息

src/utils/gpu_monitor.py:51

↓ 2 callersFunction_get_package_version

(name: str)

src/utils/data_utils.py:92

↓ 2 callersMethod_post_process

(self)

src/msa_service.py:584

↓ 2 callersFunction_rag_0108

(name: str)

src/benchmarks.py:64

↓ 2 callersMethod_resolve

Return local path if cached, otherwise download from HF into data/.

src/benchmarks.py:31

↓ 2 callersFunctionbase_it

(predict, label, at, score_func)

src/app/benchmark.py:88

↓ 2 callersMethodchunks

(self)

src/msa_service.py:226

↓ 2 callersFunctioncompose_input

组建reference的 input

src/utils/tools.py:95

↓ 2 callersMethodcopy

创建缓存的深拷贝 Returns: CustomDynamicCache: 缓存的新副本

src/utils/cache.py:101

↓ 2 callersMethodcreate_global_doc_ids

(self, local_doc_ids)

src/msa_service.py:218

↓ 2 callersFunctionformat_num

(num)

src/utils/common.py:155

↓ 2 callersMethodgenerate

(self, prompts: Union[str, List[str]], userdata=None, requi

src/msa_service.py:1798

↓ 2 callersMethodget_idx_to_doc

(self)

src/msa_service.py:1587

↓ 2 callersMethodget_model_envs

(self)

src/config/memory_config.py:40

↓ 2 callersFunctionget_sequence_parallel_preprocess

(stage, tokenizer, cutoff_len=None, sequence_parallel_size=1, sequence_parallel_mode="ulysses")

src/utils/data_utils.py:237

↓ 2 callersFunctionis_dist_avail_and_initialized

()

src/utils/misc.py:29

↓ 2 callersMethodnum_model_layers

(self)

src/msa_service.py:347

↓ 2 callersMethodstop_all

停止所有GPU监控并返回结果

src/utils/gpu_monitor.py:284

↓ 2 callersMethodupdate_router_kcache

( self, key_states: torch.Tensor, layer_idx: int, )

src/utils/cache.py:117

↓ 1 callersMethod_apply_template

(self, prompt)

src/msa_service.py:1722

↓ 1 callersMethod_apply_template_regenerate

(self, prompt)

src/msa_service.py:1734

↓ 1 callersMethod_calculate_routing_scores_adaptive

( self, query_states: torch.Tensor, # [B, H, Q_len, D] pooled_k_bched: torch.Tenso

src/msa/memory_sparse_attention.py:407

↓ 1 callersFunction_convert_recursive

(obj)

src/utils/cache.py:398

↓ 1 callersMethod_create_generate_args

(self)

src/msa_service.py:1042

↓ 1 callersMethod_forward

( self, hidden_states: torch.Tensor, doc_ids: torch.LongTensor, attention_mask

src/msa/memory_sparse_attention.py:531

↓ 1 callersMethod_gather_querys

处理变长 seqlen：填充到统一长度并 all_gather

src/msa_service.py:1088

↓ 1 callersMethod_generate_slice

(self, chunks)

src/msa_service.py:644

↓ 1 callersFunction_get_sequence_parallel_dataset

(dataset, num_works, tokenizer=None, cutoff_len=10000, sequence_parallel_s

src/utils/data_utils.py:248

↓ 1 callersMethod_inference

重新加载单个memory block 完整复制eval_anything_v2_batch.py中reload_memory的逻辑 Args: block: memory block数据 [(doc_id, doc_str)

src/prefill.py:248

↓ 1 callersMethod_init_block_data

(self, shape: torch.Size, dtype: torch.dtype, has_rk=False)

src/msa_service.py:481

↓ 1 callersMethod_load_memory_file

加载memory文件

src/msa_service.py:1502

↓ 1 callersMethod_load_model

加载模型和tokenizer

src/prefill.py:58

↓ 1 callersMethod_prepare_block_inputs

重新加载单个memory block 完整复制eval_anything_v2_batch.py中reload_memory的逻辑 Args: block: memory block数据 [(doc_id, doc_str)

src/prefill.py:188

↓ 1 callersMethod_prepare_template

(self)

src/msa_service.py:1702

↓ 1 callersMethod_prepare_template

重新加载单个memory block 完整复制eval_anything_v2_batch.py中reload_memory的逻辑 Args: block: memory block数据 [(doc_id, doc_str)

src/prefill.py:159

↓ 1 callersMethod_process_buckets

发送 blocks 到各 worker

src/msa_service.py:1554

↓ 1 callersMethod_setup_device

设置CUDA设备

src/utils/gpu_worker.py:16

↓ 1 callersMethod_setup_environment

设置环境变量

src/utils/gpu_worker.py:11

↓ 1 callersMethod_sort_reference

对文档进行重排序并且分 bucket和block，分配到一个 GPU 上的文档被称为 bucket， bucket 和 bucket 之间的 chunk 数量尽可能均衡 Args: docs: 文档数据列表

src/msa_service.py:1478

↓ 1 callersMethod_start_worker

(self)

src/msa_service.py:323

↓ 1 callersMethod_stop_worker

(self)

src/msa_service.py:339

↓ 1 callersMethod_validate_inputs

(self, input_ids: List[List[int]])

src/msa_service.py:1688

↓ 1 callersMethod_wait_ready_signal

(self)

src/msa_service.py:1444

↓ 1 callersMethod_worker_all_gather

(self, cmdname: str, timeout=600)

src/msa_service.py:1427

↓ 1 callersMethodacquire

获取执行权限，如果超过最大并发数则阻塞 Returns: bool: 是否成功获取权限

src/utils/tools.py:71

↓ 1 callersMethodbalanced_bucket_partition

全局均衡划分算法：先全局分blocks，再分配到buckets, 这里并不需要每个 block 必须是最多max_chunk_per_block个 chunk，而是根据max_chunk_per_block这个数大致决定分成多少 block（GPU

src/msa_service.py:1452

↓ 1 callersFunctionbuild_match_prompt

(gold_answer, model_answer)

src/evaluation/llm_judge.py:16

↓ 1 callersFunctionbuild_score_prompt

(gold_answer, model_answer, query)

src/evaluation/llm_judge.py:38

↓ 1 callersMethodcaculate_aux_loss

(self, aux_loss, outputs, batch_aux_labels, device, dtype)

src/msa/model.py:679

↓ 1 callersMethodcaculate_bce_loss

(self, logits, label)

src/msa/model.py:571

↓ 1 callersMethodcaculate_infonce_loss

(self, logits, label, num_pos)

src/msa/model.py:557

↓ 1 callersMethodcalculate_decoupled_infonce_loss

改进版：Decoupled InfoNCE 解决了多正样本之间的互斥问题，检索任务推荐使用。

src/msa/model.py:524

↓ 1 callersMethodcalculate_focal_infonce_loss

Args: logits: 模型输出的 logits label: 正样本的 mask (通常是 multi-hot) num_pos: 正样本的数量 gamma: Focal Loss

src/msa/model.py:638

↓ 1 callersFunctioncalculate_ir_metrics

计算信息检索中的 Precision, Recall, F1, 和 IoU。

src/app/benchmark.py:124

↓ 1 callersMethodcalculate_multi_pos_focal_infonce

Args: logits: (Batch_Size, Num_Candidates) 或者是 (N, 1) 的形式 label: (Batch_Size, Num_Candidates) Multi-hot 标签，1为正，0为负

src/msa/model.py:576

↓ 1 callersMethodclear_kvcache

(self)

src/utils/cache.py:26

↓ 1 callersMethodclose_worker

(q: mp.Queue)

src/prefill.py:98

↓ 1 callersMethodcount_chunks_per_batch

(self, doc_ids, attention_mask, kernel_size)

src/msa/memory_sparse_attention.py:511

↓ 1 callersFunctioncreate_cache

根据参数创建合适的缓存实例 Args: args: 包含量化相关参数的命名空间对象 Returns: CustomDynamicCache or CustomQuantizeDynamicCache: 缓存实例

src/utils/cache.py:337

↓ 1 callersFunctioncumulative_concat

(tensors)

src/utils/tools.py:34

↓ 1 callersMethoddeserialize

流式反序列化：每次只加载一个文件进内存，处理完后立即释放或转移到 GPU，将系统内存（RAM）占用控制在最低。

src/msa_service.py:410

↓ 1 callersMethoddoc_query

集成版：检索 + 跨卡提取 + 查表加速 + 结果归集返回: final_k, final_v: [H, Total_Selected_C, D] 格式，直接用于 Flash Attn 散射 final_score

src/msa_service.py:1133

↓ 1 callersFunctioneval_mrr

(predict, label, at=10)

src/app/benchmark.py:111

↓ 1 callersFunctionformat_bytes

将字节数转换为人类可读的格式单位为 M 及以下不要小数点，单位为 G 以上保留一位小数，数值不能小于 1

src/utils/tools.py:5

↓ 1 callersMethodforward_with_kvcache_for_batch_parrallel

( self, hidden_states: torch.Tensor, doc_ids: torch.LongTensor, attention_mask

src/msa/memory_sparse_attention.py:115

↓ 1 callersMethodfrom_dict

从字典创建配置对象时，确保 msa_config 被转换为 DotDict。这是关键方法，AutoConfig.from_pretrained() 最终会调用这个方法。

src/msa/configuration_msa.py:41

↓ 1 callersMethodgenerate

(self, req: GenerateRequest)

src/msa_service.py:1316

↓ 1 callersMethodgenerate_blocks

(self, docs: List[Document])

src/msa_service.py:497

↓ 1 callersMethodget_bench_files

(self)

src/benchmarks.py:116

↓ 1 callersMethodget_input_embeddings

(self)

src/msa/model.py:314

↓ 1 callersMethodget_kvcache

获取指定层的KV缓存 Args: layer_idx: 层索引 Returns: Tuple[torch.Tensor, torch.Tensor]: (key_cache, value_cache

src/utils/cache.py:46

↓ 1 callersMethodget_max_local_pool_doc_id

(self)

src/msa_service.py:697

↓ 1 callersMethodget_output_embeddings

(self)

src/msa/model.py:320

↓ 1 callersFunctionget_rank

()

src/utils/misc.py:37

↓ 1 callersMethodget_router_k

获取用于检索的k

src/msa_service.py:240

↓ 1 callersMethodget_router_kcache

(self, layer_idx)

src/utils/cache.py:60

↓ 1 callersMethodget_seq_length

返回缓存状态的序列长度 Args: layer_idx: 可选的层索引 Returns: int: 序列长度

src/utils/cache.py:83

↓ 1 callersMethodget_template_prefix_kvcaches

(self, layer_idx: int)

src/msa_service.py:700

↓ 1 callersFunctiongreedy_knapsack

r""" An efficient greedy algorithm with binary search for the knapsack problem.

src/utils/data_utils.py:113

↓ 1 callersMethodinference

处理memory block Args: memory_block: 分配给此GPU的memory block

src/prefill.py:142

↓ 1 callersMethodinit_docs

(self, docs: List[Document], device)

src/msa_service.py:207

↓ 1 callersMethodinitialize

初始化Memory服务

src/msa_service.py:1401

↓ 1 callersMethodload_model

加载模型和tokenizer

src/msa_service.py:1070

↓ 1 callersFunctionmain

(origin_model_path, save_model_path)

src/utils/resave_model.py:13

↓ 1 callersMethodmap_tensor_to_group_ids

根据相邻元素是否相同，将输入 Tensor a 的值映射到递增的组 ID Tensor b。 a[i] == a[i-1] => b[i] = b[i-1] a[i] != a[i-1] => b[i] = b[i-1] + 1 b

src/msa_service.py:598

next →1–100 of 265, ranked by callers