From 3ad5374865d14b45e34b107db379bf479a61f32b Mon Sep 17 00:00:00 2001
From: Vitaliy Chiley <6439018+vchiley@users.noreply.github.com>
Date: Fri, 12 Jan 2024 19:16:19 -0800
Subject: [PATCH] Update test_flash_triton_torch.py

---
 tests/models/layers/test_flash_triton_torch.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/models/layers/test_flash_triton_torch.py b/tests/models/layers/test_flash_triton_torch.py
index 4ca5c7b668..2d0e7ad05b 100644
--- 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