From a8df9ef1730b9cf30f9341fd1298214a3b406dce Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 12 Aug 2026 14:53:25 +0700 Subject: [PATCH] Keep direct RoPE path self-contained --- src/h3_blackwell_runtime/attention.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/h3_blackwell_runtime/attention.py b/src/h3_blackwell_runtime/attention.py index f8362ee..0878ef5 100644 --- a/src/h3_blackwell_runtime/attention.py +++ b/src/h3_blackwell_runtime/attention.py @@ -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":