Update test_flash_triton_torch.py

mosaicml · Jan 13, 2024 · 3ad5374 · 3ad5374
1 parent 52a760c
commit 3ad5374
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/models/layers/test_flash_triton_torch.py b/tests/models/layers/test_flash_triton_torch.py
@@ -28,6 +28,11 @@ def allclose_helper(t0: torch.Tensor,
 ])
 @pytest.mark.parametrize('clip_qkv', [True, False])
 @pytest.mark.parametrize('qk_ln', [True, False])
+@pytest.mark.parametrize('qk_ln, qk_gn', [
+    (True, False),
+    (False, True),
+    (False, False),
+])
 @pytest.mark.parametrize('pos_emb_config', [{
     'alibi': False,
     'rope': False