llmcompressor.modeling.deepseekv32.model

def forward(
    self,
    x: torch.Tensor,
    residual: torch.Tensor,
    start_pos: int,
    freqs_cis: torch.Tensor,
    mask: Optional[torch.Tensor],
) -> torch.Tensor:
    """
    Forward pass for the Transformer block.

    Args:
        x (torch.Tensor): Input tensor.
        start_pos (int): Starting position in the sequence.
        freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
        mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.

    Returns:
        torch.Tensor: Output tensor after block computation.
    """
    if residual is None:
        x, residual = self.input_layernorm(x), x
    else:
        x, residual = self.input_layernorm(x, residual)
    x = self.self_attn(x, start_pos, freqs_cis, mask)
    x, residual = self.post_attention_layernorm(x, residual)
    x = self.mlp(x)
    return x, residual

def forward(
    self,
    input_ids: torch.LongTensor | None = None,
    labels: torch.LongTensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    **kwargs: Unpack[TransformersKwargs],
) -> CausalLMOutputWithPast:
    r"""
    Example:

    ```python
    >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM

    >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf")
    >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf")

    >>> prompt = "Hey, are you conscious? Can you talk to me?"
    >>> inputs = tokenizer(prompt, return_tensors="pt")

    >>> # Generate
    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
    ```"""
    hidden_states = self.model(
        input_ids=input_ids,
    )

    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
    slice_indices = (
        slice(-logits_to_keep, None)
        if isinstance(logits_to_keep, int)
        else logits_to_keep
    )
    logits = self.lm_head(hidden_states[:, slice_indices, :])

    loss = None
    if labels is not None:
        loss = self.loss_function(
            logits=logits,
            labels=labels,
            vocab_size=self.config.vocab_size,
            **kwargs,
        )

    return CausalLMOutputWithPast(
        loss=loss,
        logits=logits,
    )

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass for the Expert layer.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Output tensor after expert computation.
    """
    return self.down_proj(
        (F.silu(self.gate_proj(x).float()) * self.up_proj(x).float()).type_as(x)
    )

def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Forward pass for the gating mechanism.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Routing weights and selected expert indices.
    """
    scores = F.linear(x.float(), self.weight.float())
    if self.score_func == "softmax":
        scores = scores.softmax(dim=-1)
    else:
        scores = scores.sigmoid()
    original_scores = scores
    if self.e_score_correction_bias is not None:
        scores = scores + self.e_score_correction_bias
    if self.n_groups > 1:
        scores = scores.view(x.size(0), self.n_groups, -1)
        if self.bias is None:
            group_scores = scores.amax(dim=-1)
        else:
            group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)
        indices = group_scores.topk(self.topk_groups, dim=-1)[1]
        mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(
            1, indices, False
        )
        scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1)
    indices = scores.topk(self.topk, dim=-1)[1]
    weights = original_scores.gather(1, indices)
    if self.score_func == "sigmoid":
        weights /= weights.sum(dim=-1, keepdim=True)
    weights *= self.route_scale
    return weights, indices

def forward(
    self,
    x: torch.Tensor,
    start_pos: int,
    freqs_cis: torch.Tensor,
    mask: Optional[torch.Tensor],
):
    """
    Forward pass for the Multi-Head Latent Attention (MLA) Layer.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).
        start_pos (int): Starting position in the sequence for caching.
        freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
        mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.

    Returns:
        torch.Tensor: Output tensor with the same shape as the input.
    """
    bsz, seqlen, _ = x.size()
    end_pos = start_pos + seqlen
    qr = self.q_a_layernorm(self.q_a_proj(x))
    q = self.q_b_proj(qr)
    q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
    q_nope, q_pe = torch.split(
        q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
    )
    q_pe = apply_rotary_emb(q_pe, freqs_cis)
    kv = self.kv_a_proj_with_mqa(x)
    kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
    kv = self.kv_a_layernorm(kv)
    k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)
    self.kv_cache[:bsz, start_pos:end_pos] = kv
    self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)
    if mask is not None:  # MHA prefill
        q = torch.cat([q_nope, q_pe], dim=-1)
        kv = self.kv_b_proj(kv)
        kv = kv.view(
            bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim
        )
        k_nope, v = torch.split(
            kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
        )
        k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)
        scores = torch.einsum("bshd,bthd->bsht", q, k).mul_(self.softmax_scale)

        # indexer
        topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
        index_mask = torch.full(
            (bsz, seqlen, seqlen), float("-inf"), device=x.device
        ).scatter_(-1, topk_indices, 0)
        index_mask += mask
        scores += index_mask.unsqueeze(2)

        scores = scores.softmax(dim=-1)
        x = torch.einsum("bsht,bthd->bshd", scores, v)
    else:  # MQA decode
        if self.dequant_wkv_b is None and self.kv_b_proj.scale is not None:
            self.dequant_wkv_b = weight_dequant(
                self.kv_b_proj.weight, self.kv_b_proj.scale
            )
        kv_b_proj = (
            self.kv_b_proj.weight
            if self.dequant_wkv_b is None
            else self.dequant_wkv_b
        )
        kv_b_proj = kv_b_proj.view(self.n_local_heads, -1, self.kv_lora_rank)
        q_nope = torch.einsum(
            "bshd,hdc->bshc", q_nope, kv_b_proj[:, : self.qk_nope_head_dim]
        )
        scores = (
            torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos])
            + torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos])
        ) * self.softmax_scale

        # indexer
        topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
        index_mask = torch.full(
            (bsz, 1, end_pos), float("-inf"), device=x.device
        ).scatter_(-1, topk_indices, 0)
        scores += index_mask.unsqueeze(2)

        scores = scores.softmax(dim=-1)
        x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
        x = torch.einsum("bshc,hdc->bshd", x, kv_b_proj[:, -self.v_head_dim :])
    x = self.o_proj(x.flatten(2))
    return x

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass for the MLP layer.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Output tensor after MLP computation.
    """
    return self.down_proj(
        (F.silu(self.gate_proj(x).float()) * self.up_proj(x).float()).type_as(x)
    )

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass for the MoE module.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Output tensor after expert routing and computation.
    """
    shape = x.size()
    x = x.view(-1, self.dim)
    weights, indices = self.gate(x)
    y = torch.zeros_like(x, dtype=torch.float32)
    counts = torch.bincount(
        indices.flatten(), minlength=self.n_routed_experts
    ).tolist()
    for i in range(self.experts_start_idx, self.experts_end_idx):
        if counts[i] == 0:
            continue
        expert = self.experts[i]
        idx, top = torch.where(indices == i)
        y[idx] += expert(x[idx]) * weights[idx, top, None]
    y += self.shared_experts(x)
    if world_size > 1:
        dist.all_reduce(y)
    return y.type_as(x).view(shape)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass for parallel embedding layer.

    Args:
        x (torch.Tensor): Input tensor containing token indices.

    Returns:
        torch.Tensor: Embedded representations.

    Raises:
        ValueError: If `world_size` is not defined.
    """
    if world_size > 1:
        mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
        x = x - self.vocab_start_idx
        x[mask] = 0
    y = F.embedding(x, self.weight)
    if world_size > 1:
        y[mask] = 0
        dist.all_reduce(y)
    return y

def forward(self, x: torch.Tensor, residual: Optional[torch.Tensor] = None):
    """
    Forward pass for RMSNorm.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Normalized tensor with the same shape as input.
    """
    dtype = x.dtype
    if residual is None:
        x = x.float()
        var = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return (self.weight * x).to(dtype)
    else:
        x = residual = x.float() + residual.float()
        var = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return (self.weight * x).to(dtype), residual.to(dtype)

@torch.inference_mode()
def forward(self, input_ids: torch.Tensor, start_pos: int = 0):
    """
    Forward pass for the Transformer model.

    Args:
        input_ids (torch.Tensor): Input tensor of token IDs with shape (batch_size, seq_len).
        start_pos (int, optional): Starting position in the sequence for rotary embeddings. Defaults to 0.

    Returns:
        torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
    """
    seqlen = input_ids.size(1)
    freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
    mask = (
        torch.full((seqlen, seqlen), float("-inf"), device=input_ids.device).triu_(
            1
        )
        if seqlen > 1
        else None
    )
    h, residual = self.embed_tokens(input_ids), None
    for layer in self.layers:
        h, residual = layer(h, residual, start_pos, freqs_cis, mask)
    h, _ = self.norm(h, residual)
    return h