Maxtimer97
/

GLM2NSA

Text Generation

Model card Files Files and versions

Maxtimer97 commited on Oct 3

Commit

bb26ab9

·

1 Parent(s): c8b397a

Removed wrong edge case assertion

Files changed (1) hide show

compressed_attention.py +1 -1

compressed_attention.py CHANGED Viewed

@@ -541,7 +541,7 @@ def _compressed_attention_fwd(
     k_len, num_k_heads, head_dim = k.shape
     v_len, num_v_heads, head_dim = v.shape
     batch_size = cu_seqlens_q.shape[0] - 1
-    assert k_len == v_len and q_len > k_len
     # gqa
     assert num_k_heads == num_v_heads
     assert num_q_heads % num_k_heads == 0

     k_len, num_k_heads, head_dim = k.shape
     v_len, num_v_heads, head_dim = v.shape
     batch_size = cu_seqlens_q.shape[0] - 1
+    assert k_len == v_len and q_len >= k_len
     # gqa
     assert num_k_heads == num_v_heads
     assert num_q_heads % num_k_heads == 0