class LlamaMLP(nn.Module):
    def __init__(self, prefix, config, weights):
        super().__init__()
        act = config.hidden_act
        # Llama Функция активации FeedForward 2 — SwishGeLU.
        # TGI реализуется путем поточечного умножения «silu(W0*x)» и «W1*x».
        # Вот, сам.действуй = torch.nn.functional.silu
        self.act = (
            ACT2FN[act]
            if "gelu" not in act
            else lambda x: torch.nn.functional.gelu(
                x,
                approximate="tanh"
                if act in ["gelu_fast", "gelu_pytorch_tanh"]
                else "none",
            )
        )
        # Fuse gate and up proj
        # «up_proj» относится к первому FFN FeedForward. Этот FFN используется для увеличения размера, поэтому он называется «up».
        # «gate_proj» относится к линейному уровню SwishGeLU, который реализует саморегулирование, поэтому его называют «воротом».
        # Их входные данные одинаковы, поэтому их веса можно объединять и объединять для умножения матриц.
        self.gate_up_proj = TensorParallelColumnLinear.load_multi(
            config,
            prefixes=[f"{prefix}.gate_proj", f"{prefix}.up_proj"],
            weights=weights,
            dim=0, # Соедините несколько весов вместе по 0-му измерению.
            bias=False,
        )
        #  «down_proj» относится ко второму FFN FeedForward. Этот FFN используется для уменьшения размерности, поэтому он называется «down».
        self.down_proj = TensorParallelRowLinear.load(
            config,
            prefix=f"{prefix}.down_proj",
            weights=weights,
            bias=False,
        )
        # Длина вектора, полученного после увеличения размерности первого FFN
        # Здесь учитывается фактическая длина Модели после раскроя.
        self.intermediate_size = (
            config.intermediate_size // weights.process_group.size()
        )  

# роды server/text_generation_server/utils/layers.pyclass TensorParallelRowLinear(SuperLayer):
    def __init__(self, linear, process_group):
        super().__init__(linear)
        # Получитеprocess_group после инициализации связи с коллекцией Pytorch.
        # Выполните коллективную связь с соответствующим графическим процессором в соответствии с Process_group.
        self.process_group = process_group

    @classmethod
    def load(cls, config, prefix: str, weights, bias: bool):
        # Веса создаются во время процесса инициализации FlashLlama (подробнее см. общую логику).
        # префикс — это название веса
        # quantizeуказать Количественная оценкаметод,некоторый Количественная Также необходимо загрузить дополнительные веса, такие как Scale.
        # вес — это Pytorch, загруженный в этот графический процессор Tensor
        # Более подробная информация ниже Реализация get_multi_weights_row()
        weight = weights.get_multi_weights_row(prefix, quantize=config.quantize)

        # Если и есть предвзятость, то только в РАНГЕ 0 выполняет операцию добавления смещения
        # Если его добавить к каждому графическому процессору, он не будет математически эквивалентен.
        if bias and weights.process_group.rank() == 0:
            # Rank is only on the first rank process
            bias = weights.get_tensor(f"{prefix}.bias")
        else:
            bias = None
        # get_linear — метод фабричного шаблона (подробнее ниже).
        # Передайте вес, смещение и квантование, чтобы создать экземпляр линейного слоя.
        return cls(
            get_linear(weight, bias, config.quantize),
            process_group=weights.process_group,
        )

# роды server/text_generation_server/utils/weights.pydef get_multi_weights_row(self, prefix: str, quantize: str):
        if quantize == "gptq":
            # 如果Количественная Метод оценки — «gptq», который загружает несколько весов из файла. Логика здесь опущена.
            weight = (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama)
        elif quantize == "awq":
            # Как и выше, опустите
            weight = (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama)
        else:
            # еще нет Количественная вес определения (float32/float16), следуйте этой логике загрузки
            # Аннотировано ниже
            weight = self.get_sharded(f"{prefix}.weight", dim=1)
        return weightdef get_sharded(self, tensor_name: str, dim: int):
    # При инициализации класса Weights устанавливается связь между именами весов и файлами весов.
    # Найдите файл и имя веса (которое, возможно, потребуется исправить) с помощью этого сопоставления.
    filename, tensor_name = self.get_filename(tensor_name)
    # Получите дескриптор файла, тип — Safetensors.
    f = self._get_handle(filename)
    # Получите соответствующий Тензор из защитных тензоров по названию веса.
    # Убедитесь, что измерение, на которое разбиты веса, делится на количество графических процессоров.
    slice_ = f.get_slice(tensor_name)
    world_size = self.process_group.size()
    size = slice_.get_shape()[dim]
    assert (
        size % world_size == 0
    ), f"The choosen size {size} is not compatible with sharding on {world_size} shards"
    return self.get_partial_sharded(tensor_name, dim) # Аннотировано нижеdef get_partial_sharded(self, tensor_name: str, dim: int):
    # Логика здесь аналогична get_sharded() и не будет повторяться.
    filename, tensor_name = self.get_filename(tensor_name)
    f = self._get_handle(filename)
    slice_ = f.get_slice(tensor_name)
    world_size = self.process_group.size()
    rank = self.process_group.rank()

    # Рассчитайте смещение сегментации в зависимости от ранга
    size = slice_.get_shape()[dim]
    block_size = size // world_size
    start = rank * block_size
    stop = (rank + 1) * block_size

    # Вырежьте часть гирь, которая должна быть нагружена.
    if dim == 0:
        tensor = slice_[start:stop]
    elif dim == 1:
        tensor = slice_[:, start:stop]
    else:
        raise NotImplementedError("Let's make that generic when needed")
    # Special case for gptq which shouldn't convert
    # u4 which are disguised as int32
    if tensor.dtype != torch.int32:
        tensor = tensor.to(dtype=self.dtype)
    tensor = tensor.to(device=self.device)
    return tensor

# роды server/text_generation_server/utils/layers.pydef get_linear(weight, bias, quantize):
    # В зависимости от того, Количественная оценкаи Количественная Различные способы создания экземпляров различных типов Linear (унаследованы от nn.Module)
    if quantize is None:
        # Реализация FastLinear размещена ниже.
        linear = FastLinear(weight, bias)
    elif quantize == "eetq":
        if HAS_EETQ:
            linear = EETQLinear(weight, bias)
        else:
            raise ImportError(
                "Please install EETQ from https://github.com/NetEase-FuXi/EETQ"
            )
    # 其他из Количественная оценкаметод实例化Линейная логика аналогичная, опущена
    # elif quantize == "bitsandbytes":
    # elif quantize == "bitsandbytes-fp4":
    # elif quantize == "bitsandbytes-nf4":
    # elif quantize == "gptq":
    # elif quantize == "awq":
    else:
        raise NotImplementedError(f"Quantization `{quantize}` is not implemented yet.")
    return linear# Просто обычный torch.nn.functional.linearclass FastLinear(nn.Module):
    def __init__(
        self,
        weight,
        bias,
    ) -> None:
        super().__init__()
        self.weight = nn.Parameter(weight)
        if bias is not None:
            self.bias = nn.Parameter(bias)
        else:
            self.bias = None

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return F.linear(input, self.weight, self.bias)

class FlashLlamaAttention(torch.nn.Module):
    def __init__(
        self,
        prefix: str,
        config,
        weights,
    ):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_size = self.hidden_size // self.num_heads

        # для RoPE (поворотный Position Вложение) расчет
        # self.rotary_emb = PositionRotaryEmbedding.load(
        #     config=config, prefix=f"{prefix}.rotary_emb", weights=weights
        # )
        self.rotary_emb = PositionRotaryEmbedding.static(
            config=config,
            dim=self.head_size,
            base=config.rope_theta,
            device=weights.device,
        )

        self.softmax_scale = self.head_size**-0.5

        if self.num_heads % weights.process_group.size() != 0:
            raise ValueError(
                f"`num_heads` must be divisible by `num_shards` (got `num_heads`: {self.num_heads} "
                f"and `num_shards`: {weights.process_group.size()}"
            )
        self.num_heads = self.num_heads // weights.process_group.size()
        self.num_key_value_heads = (
            config.num_key_value_heads // weights.process_group.size()
        )

        # Реализация load_attention() выглядит следующим образом.
        self.query_key_value = load_attention(config, prefix, weights)

        # Разделите o_proj по строкам и загрузите его.
        self.o_proj = TensorParallelRowLinear.load(
            config,
            prefix=f"{prefix}.o_proj",
            weights=weights,
            bias=False,
        )
        # Входной параметр PagedAttention, используемый для поддержки GQA/MQA.
        self.num_groups = self.num_heads // self.num_key_value_heads
        self.kv_head_mapping = torch.arange(
            0, self.num_key_value_heads, dtype=torch.int32, device=weights.device
        ).repeat_interleave(self.num_groups)

    def load_attention(config, prefix, weights):
        # Загрузка GQA/MQA
        if config.num_attention_heads != config.num_key_value_heads:
            return _load_gqa(config, prefix, weights)
        # Загрузка MHA
        else:
            # Baichuan и Llama имеет определенные различия в части «Внимание» и использует разные реализации загрузки.
            # Но они все загружены по колонкам.
            if config.model_type == "baichuan":
                return TensorParallelColumnLinear.load_qkv(
                    config,
                    prefix=f"{prefix}.W_pack",
                    weights=weights,
                    bias=False,
                )
            else:
                return TensorParallelColumnLinear.load_multi(
                    config,
                    prefixes=[f"{prefix}.q_proj", f"{prefix}.k_proj", f"{prefix}.v_proj"],
                    dim=0,
                    weights=weights,
                    bias=False,
                )

# роды server/text_generation_server/models/custom_modeling/flash_llama_modeling.pyclass LlamaMLP(nn.Module):
    # Логика __init__() была прокомментирована выше и не будет здесь повторяться.
    def __init__(self, prefix, config, weights):
        super().__init__()
        act = config.hidden_act
        self.act = () # Параметр опущен
        # Fuse gate and up proj
        self.gate_up_proj = TensorParallelColumnLinear.load_multi() # Параметр опущен
        self.down_proj = TensorParallelRowLinear.load() # Параметр опущен
        self.intermediate_size = (
            config.intermediate_size // weights.process_group.size()
        )

    def forward(self, hidden_states):
        # через Gate_up_proj, Вычислениеgate_states иup_states одновременно за одно матричное вычисление
        gate_up_states = self.gate_up_proj(hidden_states)
        # Разделить Gate_up_states на Gate_states и Up_states с помощью операции Reshape и Slice.
        gate_up_states = gate_up_states.view(-1, 2, self.intermediate_size)
        # здесь self.act для torch.nn.functional.silu()
        # Через вычисления в скобках self.down_proj объединитеgate_states иup_states, чтобы получить результат SwishGeLU
        # Наконец, выходные данные второго Linear в FeedForward рассчитываются с помощью self.down_proj().
        return self.down_proj(self.act(gate_up_states[:, 0]) * gate_up_states[:, 1])

# роды server/text_generation_server/utils/layers.py# Класс SuperLayer TensorParallelColumnLinear и TensorParallelRowLinear SuperLayer(nn.Module):
    def __init__(self, linear):
        super().__init__()
        # Содержит соответствующий тип (Количественная оценка/Нет Количественная оценка）изlinear        self.linear = linear

    def forward(self, x):
        # Просто отрегулируйте направление движения self.linear вперед.
        return self.linear.forward(x)# TensorParallelColumnLinear не отменяет переопределение вперед()# То есть класс реализации front() SuperLayer корректируется. TensorParallelColumnLinear(SuperLayer):
    @classmethod
    def load_qkv(cls, config, prefix: str, weights, bias: bool):
        # упущение
        pass

    @classmethod
    def load(cls, config, prefix: str, weights, bias: bool):
        # проанализировано выше
        return cls.load_multi(config, [prefix], weights, bias, dim=0)

    @classmethod
    def load_multi(cls, config, prefixes: List[str], weights, bias: bool, dim: int):
        # упущение
        pass# TensorParallelRowLinear перезаписывает класс вперед() TensorParallelRowLinear(SuperLayer):
    def __init__(self, linear, process_group):
        super().__init__(linear)
        self.process_group = process_group

    @classmethod
    def load(cls, config, prefix: str, weights, bias: bool):
        # проанализировано выше
        weight = weights.get_multi_weights_row(prefix, quantize=config.quantize)

        if bias and weights.process_group.rank() == 0:
            # Rank is only on the first rank process
            bias = weights.get_tensor(f"{prefix}.bias")
        else:
            bias = None
        return cls(
            get_linear(weight, bias, config.quantize),
            process_group=weights.process_group,
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        # Отрегулируйте линейную пересылку вперед() для выполнения матричных вычислений.
        out = super().forward(input)
        # если self.process_group.size() > 1 означает, что установлен тензорный параллелизм (рассуждение с несколькими картами)
        # Интегрируйте результаты вычислений через all_reduce, где это необходимо, чтобы результаты математически соответствовали рассуждениям по одной карте.
        if self.process_group.size() > 1:
            torch.distributed.all_reduce(out, group=self.process_group)
        return out

# роды server/text_generation_server/models/custom_modeling/flash_llama_modeling.pyclass FlashLlamaAttention(torch.nn.Module):
    def __init__(
        self,
        prefix: str,
        config,
        weights,
    ):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_size = self.hidden_size // self.num_heads

        self.rotary_emb = PositionRotaryEmbedding.static() # Параметр опущен

        self.softmax_scale = self.head_size**-0.5

        if self.num_heads % weights.process_group.size() != 0:
            raise ValueError
        self.num_heads = self.num_heads // weights.process_group.size()
        self.num_key_value_heads = (
            config.num_key_value_heads // weights.process_group.size()
        )

        self.query_key_value = load_attention(config, prefix, weights)

        self.o_proj = TensorParallelRowLinear.load() # Параметр опущен
        self.num_groups = self.num_heads // self.num_key_value_heads
        self.kv_head_mapping = torch.arange(
            0, self.num_key_value_heads, dtype=torch.int32, device=weights.device
        ).repeat_interleave(self.num_groups)

    def forward(
        self,
        hidden_states,
        cos,
        sin,
        cu_seqlen_prefill,
        kv_cache,
        block_tables,
        slots,
        input_lengths,
        max_s,
    ):
        # Привязать скрытые_состояния к QKV
        qkv = self.query_key_value(hidden_states)
        # ВоляQKVРасколотьдляQиKVиReshape,Рассматривается здесьGQA\MQAиз情况
        query, kv = qkv.split(
            [
                self.head_size * self.num_heads,
                2 * self.head_size * self.num_key_value_heads,
            ],
            dim=1,
        )
        query = query.view(-1, self.num_heads, self.head_size)
        kv = kv.view(-1, 2, self.num_key_value_heads, self.head_size)

        # queryиkey необходимо добавить RoPE (Rotary Position Embedding）
        # косисин был рассчитан заранее и повторно использован в каждом слое 
        self.rotary_emb(query, cos, sin)
        self.rotary_emb(torch.select(kv, dim=1, index=0), cos, sin)

        # Сохраните вновь рассчитанный тензор KeyиValue в KV, управляемый PagedAttention. В кэше
        # kv — вновь рассчитанное значение KV,
        # kv_cache — существующий kv_cache
        # slots инструктирует reshape_and_cache() скопировать kv в соответствующее место в kv_cache.
        paged_attention.reshape_and_cache(
            kv[:, 0], kv[:, 1], kv_cache[0], kv_cache[1], slots
        )

        # output tensor
        attn_output = torch.empty_like(query)

        # Чтобы узнать о различиях в вычислениях Prefill и Decode, прочтите первую часть этой серии.
        # Prefill
        if cu_seqlen_prefill is not None:
            # flash attention
            # Давайте сосредоточимся на cu_seqlen_prefillимакс_s.
            # Этот API FlashAttention поддерживает пакетную работу.
            # То есть он поддерживает Q(K\V) разных выборок. Тензоры склеены между собой
            # Вам нужно выполнить только один вывод, чтобы получить результаты внимания для каждого образца.
            # cu_seqlen_prefill используется для указания положения каждого семпла в склеенном тензоре
            # max_s используется для запроса длины самой длинной выборки для облегчения планирования FA.
            flash_attn.attention(
                query, # Query
                torch.select(kv, dim=1, index=0), # New Key
                torch.select(kv, dim=1, index=1), # New Value
                attn_output, # Предварительно выделенный выход
                cu_seqlen_prefill, 
                max_s,
                self.softmax_scale,
            )
        # Decode
        else:
            paged_attention.attention(
                attn_output, # Предварительно выделенный выход
                query, # Query
                kv_cache[0], # Key Cache
                kv_cache[1], # Value Cache
                self.kv_head_mapping, # Используется для обработки GQA/MQA, то есть для сопоставления головы KV и количества голов Q, когда d не равно.
                self.softmax_scale,
                block_tables, # Используется для указания KV, необходимого для этого расчета. Место хранения кэша
                input_lengths, # Текущая длина каждой выборки (len(KV_Cache)+1)
                max_s, # max_s используется для запроса длины самой длинной выборки для облегчения планирования FA.
            )

        # Наконец, значение, рассчитанное с помощью Attention, отправляется в self.o_proj() для аффинного
        return self.o_proj(attn_output.view(-1, self.num_heads * self.head_size))