daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Keep direct RoPE path self-contained

Browse source
This commit is contained in:
Daniel Maddern 2026-08-12 14:53:25 +07:00
parent 79bdce15e0
commit a8df9ef173
1 changed files with 4 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/h3_blackwell_runtime/attention.py
View file

@ -86,15 +86,12 @@ class H3SageAttention(nn.Module):
sequence = x.shape[0]
inner = self.heads * self.head_dim
q, k, v = self.qkv_proj(x).split(inner, dim=-1)
q = q.view(1, sequence, self.heads, self.head_dim)
k = k.view(1, sequence, self.heads, self.head_dim)
q = rms_norm(q.view(1, sequence, self.heads, self.head_dim), self.q_norm_weight, self.eps)
k = rms_norm(k.view(1, sequence, self.heads, self.head_dim), self.k_norm_weight, self.eps)
v = v.view(1, sequence, self.heads, self.head_dim)
from comfy_kitchen import rms_rope_split_half_
rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, epsilon=self.eps, rot_dim=rope_rotation.shape[-3] * 2)
q = q.transpose(1, 2).contiguous()
k = k.transpose(1, 2).contiguous()
q = apply_split_half_rope(q, rope_rotation).transpose(1, 2).contiguous()
k = apply_split_half_rope(k, rope_rotation).transpose(1, 2).contiguous()
v = v.transpose(1, 2).contiguous()
if self.backend == "sage2":