fixing kdiff

arch/hourglass/axial_rope.py

@@ -21,7 +21,6 @@ def rotate_half(x):
     return x.view(*shape, d * r)
 
 
-@flags.compile_wrap
 def apply_rotary_emb(freqs, t, start_index=0, scale=1.0):
     freqs = freqs.to(t)
     rot_dim = freqs.shape[-1]

arch/hourglass/image_transformer_v2.py

@@ -87,7 +87,6 @@ def filter_params(function, module):
 
 # Kernels
 
-@flags.compile_wrap
 def linear_geglu(x, weight, bias=None):
     x = x @ weight.mT
     if bias is not None:
@@ -96,7 +95,6 @@ def linear_geglu(x, weight, bias=None):
     return x * F.gelu(gate)
 
 
-@flags.compile_wrap
 def rms_norm(x, scale, eps):
     dtype = reduce(torch.promote_types, (x.dtype, scale.dtype, torch.float32))
     mean_sq = torch.mean(x.to(dtype)**2, dim=-1, keepdim=True)
@@ -104,7 +102,6 @@ def rms_norm(x, scale, eps):
     return x * scale.to(x.dtype)
 
 
-@flags.compile_wrap
 def scale_for_cosine_sim(q, k, scale, eps):
     dtype = reduce(torch.promote_types, (q.dtype, k.dtype, scale.dtype, torch.float32))
     sum_sq_q = torch.sum(q.to(dtype)**2, dim=-1, keepdim=True)
@@ -115,7 +112,6 @@ def scale_for_cosine_sim(q, k, scale, eps):
     return q * scale_q.to(q.dtype), k * scale_k.to(k.dtype)
 
 
-@flags.compile_wrap
 def scale_for_cosine_sim_qkv(qkv, scale, eps):
     q, k, v = qkv.unbind(2)
     q, k = scale_for_cosine_sim(q, k, scale[:, None], eps)
@@ -179,7 +175,6 @@ class AdaRMSNorm(nn.Module):
 
 # Rotary position embeddings
 
-@flags.compile_wrap
 def apply_rotary_emb(x, theta, conj=False):
     out_dtype = x.dtype
     dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32))
@@ -195,7 +190,6 @@ def apply_rotary_emb(x, theta, conj=False):
     return torch.cat((y1, y2, x3), dim=-1)
 
 
-@flags.compile_wrap
 def _apply_rotary_emb_inplace(x, theta, conj):
     dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32))
     d = theta.shape[-1]