Upload HfMoondream

text.py

@@ -28,18 +28,24 @@ def attn(
     qkv_out = w.qkv(x)  # shape: (bsz, q_len, (n_heads + 2*n_kv_heads)*head_dim)
     q_dim = n_heads * head_dim
     kv_dim = n_kv_heads * head_dim
-
-    q = qkv_out[..., :q_dim].view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
-    k = (
-        qkv_out[..., q_dim : q_dim + kv_dim]
-        .view(bsz, q_len, n_kv_heads, head_dim)
-        .transpose(1, 2)
-    )
-    v = (
-        qkv_out[..., q_dim + kv_dim :]
-        .view(bsz, q_len, n_kv_heads, head_dim)
-        .transpose(1, 2)
-    )
+    q, k, v = qkv_out.split([q_dim, kv_dim, kv_dim], dim=-1)
+    del qkv_out
+
+    q = q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
+    k = k.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
+    v = v.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
+
+    # q = qkv_out[..., :q_dim].view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
+    # k = (
+    #     qkv_out[..., q_dim : q_dim + kv_dim]
+    #     .view(bsz, q_len, n_kv_heads, head_dim)
+    #     .transpose(1, 2)
+    # )
+    # v = (
+    #     qkv_out[..., q_dim + kv_dim :]
+    #     .view(bsz, q_len, n_kv_heads, head_dim)
+    #     .transpose(1, 2)
+    # )
 
     q = apply_rotary_emb(q, freqs_cis, position_ids, n_heads)
     k = apply_rotary_emb(k, freqs_cis, position_ids, n_kv_heads)