Skip to content

Commit 2d330a7

Browse files
committed
Fixup RMSnorm test to only target hopper
Signed-off-by: Bram Wasti <bwasti@meta.com>
1 parent 6b909fc commit 2d330a7

File tree

1 file changed

+12
-12
lines changed

1 file changed

+12
-12
lines changed

tests/v1/generation/test_rms_norm_batch_invariant.py

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616

1717

1818
@pytest.mark.skipif(
19-
current_platform.has_device_capability(100),
20-
reason="Batch invariance tests not yet supported on B200 (SM100)",
19+
not current_platform.has_device_capability(90),
20+
reason="Batch invariance tests only supported on Hopper (SM90)",
2121
)
2222
@pytest.mark.skipif(
2323
not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
@@ -71,8 +71,8 @@ def test_rms_norm_batch_invariant_vs_standard(
7171

7272

7373
@pytest.mark.skipif(
74-
current_platform.has_device_capability(100),
75-
reason="Batch invariance tests not yet supported on B200 (SM100)",
74+
not current_platform.has_device_capability(90),
75+
reason="Batch invariance tests only supported on Hopper (SM90)",
7676
)
7777
@pytest.mark.skipif(
7878
not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
@@ -119,8 +119,8 @@ def test_rms_norm_3d_input(batch_size: int, seq_len: int, hidden_size: int):
119119

120120

121121
@pytest.mark.skipif(
122-
current_platform.has_device_capability(100),
123-
reason="Batch invariance tests not yet supported on B200 (SM100)",
122+
not current_platform.has_device_capability(90),
123+
reason="Batch invariance tests only supported on Hopper (SM90)",
124124
)
125125
@pytest.mark.skipif(
126126
not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
@@ -185,8 +185,8 @@ def test_rms_norm_numerical_stability():
185185

186186

187187
@pytest.mark.skipif(
188-
current_platform.has_device_capability(100),
189-
reason="Batch invariance tests not yet supported on B200 (SM100)",
188+
not current_platform.has_device_capability(90),
189+
reason="Batch invariance tests only supported on Hopper (SM90)",
190190
)
191191
@pytest.mark.skipif(
192192
not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
@@ -224,8 +224,8 @@ def test_rms_norm_formula():
224224

225225

226226
@pytest.mark.skipif(
227-
current_platform.has_device_capability(100),
228-
reason="Batch invariance tests not yet supported on B200 (SM100)",
227+
not current_platform.has_device_capability(90),
228+
reason="Batch invariance tests only supported on Hopper (SM90)",
229229
)
230230
@pytest.mark.skipif(
231231
not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
@@ -268,8 +268,8 @@ def test_rms_norm_different_hidden_sizes(hidden_size: int):
268268

269269

270270
@pytest.mark.skipif(
271-
current_platform.has_device_capability(100),
272-
reason="Batch invariance tests not yet supported on B200 (SM100)",
271+
not current_platform.has_device_capability(90),
272+
reason="Batch invariance tests only supported on Hopper (SM90)",
273273
)
274274
@pytest.mark.skipif(
275275
not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"

0 commit comments

Comments
 (0)