Skip to content

vllm.models.step5.nvidia.ops.mtp_norm

Normalize MTP embedding and hidden states directly into the projection input.

Functions:

  • norm_concat –

    Preserve the BF16 input rounding used with FP32 residual connections.

norm_concat(embeds, hidden, embed_weight, hidden_weight, eps, zero_centered, fp32_residual_connection)

Preserve the BF16 input rounding used with FP32 residual connections.

Source code in vllm/models/step5/nvidia/ops/mtp_norm.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def norm_concat(
    embeds: torch.Tensor,
    hidden: torch.Tensor,
    embed_weight: torch.Tensor,
    hidden_weight: torch.Tensor,
    eps: float,
    zero_centered: bool,
    fp32_residual_connection: bool,
) -> torch.Tensor:
    """Preserve the BF16 input rounding used with FP32 residual connections."""
    dtype = (
        torch.bfloat16
        if fp32_residual_connection
        else torch.promote_types(embeds.dtype, hidden.dtype)
    )
    output = torch.empty(
        (embeds.shape[0], 2 * embeds.shape[1]), device=embeds.device, dtype=dtype
    )
    if embeds.shape[0]:
        _norm_concat[(embeds.shape[0],)](
            embeds,
            hidden,
            embed_weight,
            hidden_weight,
            output,
            embeds.stride(0),
            hidden.stride(0),
            embeds.shape[1],
            eps,
            zero_centered,
            fp32_residual_connection,
            triton.next_power_of_2(embeds.shape[1]),
            enable_fp_fusion=False,
        )
    return output