Skip to content

Commit 2403b59

Browse files
authored
Merge pull request #72 from raindaywhu/br_wjh_eplb
add get_tok_ids
2 parents c600494 + 2125fe0 commit 2403b59

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

vllm_ascend/models/deepseek_v2.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -810,6 +810,9 @@ def get_all_expert_map(self,num_moe_layers):
810810

811811
return torch.stack(all_loads, dim=0)
812812

813+
def get_topk_ids(self,layer_id):
814+
return self.model.layers[layer_id+3].mlp.experts.topk_ids
815+
813816
def get_all_moe_loads(
814817
self,
815818
num_moe_layers: int,

0 commit comments

Comments
 (0)