From 6e9fb06cf96311398a5e755294565012b3ad064c Mon Sep 17 00:00:00 2001
From: Vladimir Cherepanov <vcherepanov@nvidia.com>
Date: Thu, 25 Jun 2026 19:35:08 +0000
Subject: [PATCH] Add L2 score mod distributed attention shape

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---
 tests/jax/test_distributed_fused_attn.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
index 78fd0ed487..2abd9824b6 100644
--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -283,6 +283,7 @@ def test_cross_attn(
 DISTRIBUTED_SCORE_MOD_DATA_SHAPES = {
     "L0": [],
     "L1": [(4, 16, 4, 64)],
+    "L2": [],
 }