MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / init_experts

Method init_experts

tensorrt_llm/layers/moe.py:871–891  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

869 self.max_low_rank = None
870
871 def init_experts(self):
872 # Note we use horizontal fusion for gated activation to do the operation in one GEMM invocation
873 # The left matrix is a linear projection (no activation applied)
874 # The right matrix is the gating value (activation applied)
875 # The naming convention is the inverse of GatedMLP, but the same as `tensorrt_llm/functional.py`
876 fc_out_size = self.expert_inter_size * 2 if is_gated_activation(
877 self.hidden_act) else self.expert_inter_size
878 groupwise_quant_algo = self.zero * GroupwiseQuantAlgo.ZERO + self.pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + self.use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA
879
880 self.fc = MOEWeightWrapper(self.hidden_size, fc_out_size,
881 self.experts_per_node, self.quant_mode,
882 groupwise_quant_algo, self.group_size,
883 self.dtype, self.weight_dtype, self.bias,
884 self.wrapper_tllm_to_externel_key_dict,
885 self.mapping.moe_tp_size, 0)
886 self.proj = MOEWeightWrapper(self.expert_inter_size, self.hidden_size,
887 self.experts_per_node, self.quant_mode,
888 groupwise_quant_algo, self.group_size,
889 self.dtype, self.weight_dtype, self.bias,
890 self.wrapper_tllm_to_externel_key_dict,
891 self.mapping.moe_tp_size, 1)
892
893 def default_routing(self, logits):
894 topk_values, topk_indices = topk(softmax(cast(logits, trt.float32),

Callers 1

__init__Method · 0.95

Calls 2

MOEWeightWrapperClass · 0.85
is_gated_activationFunction · 0.50

Tested by

no test coverage detected