Skip to content

vllm.models.step5.nvidia.ops.add

Add low-precision tensors without materializing their FP32 conversions.

Functions:

  • add_fp32 –

    Return x.float() + y.float(), preserving the FP32 addition.

add_fp32(x, y)

Return x.float() + y.float(), preserving the FP32 addition.

Source code in vllm/models/step5/nvidia/ops/add.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def add_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Return x.float() + y.float(), preserving the FP32 addition."""
    if (
        not x.is_cuda
        or x.device != y.device
        or x.shape != y.shape
        or not x.is_contiguous()
        or not y.is_contiguous()
        or x.dtype not in (torch.bfloat16, torch.float16, torch.float32)
        or y.dtype not in (torch.bfloat16, torch.float16, torch.float32)
        or (x.dtype == torch.float32 and y.dtype == torch.float32)
    ):
        return x.float() + y.float()
    output = torch.empty_like(x, dtype=torch.float32)
    if x.numel():
        _add_fp32[(triton.cdiv(x.numel(), 1024),)](
            x, y, output, x.numel(), 1024, enable_fp_fusion=False
        )
    return output