iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CI] Use explicit devices in IR tests (#49513)

Browse Source 
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
This commit is contained in:
Andreas Karatzas
2026-07-23 20:08:14 -07:00
committed by GitHub
parent 7bdf8cc37c
commit 2ac125123a
2 changed files with 66 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/ir/test_layernorm.py
+51 -24
View File
@@ -50,13 +50,13 @@ def test_rms_norm_registration():
reason="Currently only kernels on CUDA, ROCm and XPU",
)
class TestRMSNorm:
@classmethod
def setup_class(cls, **kwargs):
torch.set_default_device(current_platform.device_type)
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
num_tokens=4,
hidden_size=8,
dtype=dtype,
epsilon=epsilon,
device=current_platform.device_type,
)
out = rms_norm_native(x, weight, epsilon=epsilon)
@@ -87,7 +87,11 @@ class TestRMSNorm:
def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
impl = ir.ops.rms_norm.impls[provider]
x, weight, eps = ir.ops.rms_norm.generate_inputs(
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
num_tokens=n_tokens,
hidden_size=hidden_size,
dtype=dtype,
epsilon=epsilon,
device=current_platform.device_type,
)
args = (x, weight, eps)
@@ -119,7 +123,11 @@ class TestRMSNorm:
pytest.skip(f"{provider} impl not supported on this platform")
args = ir.ops.rms_norm.generate_inputs(
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
num_tokens=n_tokens,
hidden_size=hidden_size,
dtype=dtype,
epsilon=epsilon,
device=current_platform.device_type,
)
# When checking the torch op, we have to set priority and use dispatch
@@ -132,11 +140,14 @@ class TestRMSNorm:
reason="aiter is only supported on ROCm",
)
def test_aiter_rejects_unsupported_dtypes():
torch.set_default_device(current_platform.device_type)
impl = ir.ops.rms_norm.impls["aiter"]
for dtype in [torch.float32, torch.float64]:
args = ir.ops.rms_norm.generate_inputs(
num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
num_tokens=8,
hidden_size=4096,
dtype=dtype,
epsilon=1e-5,
device=current_platform.device_type,
)
assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"
@@ -146,15 +157,16 @@ def test_aiter_rejects_unsupported_dtypes():
reason="ROCm vllm_c RMSNorm needs explicit ND input handling",
)
def test_vllm_c_rms_norm_accepts_nd_input():
torch.set_default_device(current_platform.device_type)
impl = ir.ops.rms_norm.impls["vllm_c"]
if not impl.supported:
pytest.skip("vllm_c impl not supported on this platform")
base = torch.randn(3, 8, 192, dtype=torch.float16)
base = torch.randn(
3, 8, 192, dtype=torch.float16, device=current_platform.device_type
)
x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
assert not x.is_contiguous()
weight = torch.randn(16, dtype=torch.float16)
weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type)
epsilon = 1e-5
output = impl.impl_fn(x, weight, epsilon)
@@ -195,18 +207,21 @@ def test_fused_add_rms_norm_registration():
reason="ROCm vllm_c fused_add_rms_norm needs explicit ND input handling",
)
def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
torch.set_default_device(current_platform.device_type)
impl = ir.ops.fused_add_rms_norm.impls["vllm_c"]
if not impl.supported:
pytest.skip("vllm_c impl not supported on this platform")
base = torch.randn(3, 8, 192, dtype=torch.float16)
residual_base = torch.randn(3, 8, 192, dtype=torch.float16)
base = torch.randn(
3, 8, 192, dtype=torch.float16, device=current_platform.device_type
)
residual_base = torch.randn(
3, 8, 192, dtype=torch.float16, device=current_platform.device_type
)
x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
x_residual = residual_base.split(64, dim=-1)[0].view(3, 8, 4, 16)
assert not x.is_contiguous()
assert not x_residual.is_contiguous()
weight = torch.randn(16, dtype=torch.float16)
weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type)
epsilon = 1e-5
output, residual = impl.impl_fn(x.clone(), x_residual.clone(), weight, epsilon)
@@ -227,13 +242,13 @@ def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
reason="Currently only kernels on CUDA, ROCm and XPU",
)
class TestFusedAddRMSNorm:
@classmethod
def setup_class(cls, **kwargs):
torch.set_default_device(current_platform.device_type)
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
num_tokens=4,
hidden_size=8,
dtype=dtype,
epsilon=epsilon,
device=current_platform.device_type,
)
out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps)
@@ -278,7 +293,11 @@ class TestFusedAddRMSNorm:
def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
impl = ir.ops.fused_add_rms_norm.impls[provider]
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
num_tokens=n_tokens,
hidden_size=hidden_size,
dtype=dtype,
epsilon=epsilon,
device=current_platform.device_type,
)
args = (x, x_residual, weight, eps, None)
@@ -324,7 +343,11 @@ class TestFusedAddRMSNorm:
pytest.skip(f"{provider} impl not supported on this platform")
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
num_tokens=n_tokens,
hidden_size=hidden_size,
dtype=dtype,
epsilon=epsilon,
device=current_platform.device_type,
)
# Test default overload - should NOT modify inputs even with inplace impl
@@ -368,7 +391,11 @@ class TestFusedAddRMSNorm:
@pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
args = ir.ops.fused_add_rms_norm.generate_inputs(
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
num_tokens=n_tokens,
hidden_size=hidden_size,
dtype=dtype,
epsilon=epsilon,
device=current_platform.device_type,
)
args = args + (None,) # Add variance_size parameter
vllm/ir/ops/layernorm.py
+15 -7
View File
@@ -23,10 +23,14 @@ def rms_norm(
@rms_norm.register_input_generator
def _rms_norm_input_generator(
num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5
num_tokens: int,
hidden_size: int,
dtype: torch.dtype,
epsilon: float = 1e-5,
device: torch.device | str | None = None,
) -> tuple:
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
weight = torch.randn(hidden_size, dtype=dtype)
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
weight = torch.randn(hidden_size, dtype=dtype, device=device)
return x, weight, epsilon
@@ -64,9 +68,13 @@ fused_add_rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
@fused_add_rms_norm.register_input_generator
def _fused_add_rms_norm_input_generator(
num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5
num_tokens: int,
hidden_size: int,
dtype: torch.dtype,
epsilon: float = 1e-5,
device: torch.device | str | None = None,
) -> tuple:
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype)
weight = torch.randn(hidden_size, dtype=dtype)
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
weight = torch.randn(hidden_size, dtype=dtype, device=device)
return x, x_residual, weight, epsilon