44 lines
1.7 KiB
Python
44 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""No vLLM/model/MATE dependency: python -m pytest -q this_file.py."""
|
|
|
|
import importlib
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
if hasattr(torch.version, "musa") and torch.version.musa is not None:
|
|
importlib.import_module("torchada")
|
|
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.musa]
|
|
|
|
|
|
@triton.jit
|
|
def _atomic_sum(values, output, count: tl.constexpr, block: tl.constexpr):
|
|
offsets = tl.program_id(0) * block + tl.arange(0, block)
|
|
increments = tl.load(values + offsets, mask=offsets < count, other=0)
|
|
tl.atomic_add(output + offsets % 32, increments, mask=offsets < count)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
|
def test_atomic_add(dtype):
|
|
cuda = getattr(torch, "cuda")
|
|
available = cuda.is_available() or (hasattr(torch, "musa") and torch.musa.is_available())
|
|
if not available:
|
|
pytest.skip("No CUDA/MUSA device")
|
|
if dtype == torch.bfloat16 and getattr(torch.version, "musa", None) is not None:
|
|
if tuple(int(part) for part in triton.__version__.split(".")[:2]) < (3, 6):
|
|
pytest.skip("Triton-MUSA < 3.6 does not support BF16 atomic_add")
|
|
# Each destination receives eight exactly representable increments from
|
|
# two independent blocks. Reset between runs to exercise real atomics.
|
|
values = torch.ones(256, device="cuda", dtype=dtype)
|
|
output = torch.zeros(32, device="cuda", dtype=dtype)
|
|
for _ in range(5):
|
|
output.zero_()
|
|
_atomic_sum[(2,)](values, output, 256, 128)
|
|
cuda.synchronize()
|
|
torch.testing.assert_close(output, torch.full_like(output, 8), rtol=0, atol=0)
|