Entry point called by bench.py. Must match reference.rotary_embedding_ref signature. Args: x: [..., head_dim] tensor to apply rotary embeddings to cos: [..., head_dim // 2] precomputed cosines sin: [..., head_dim // 2] precomputed sines Returns: Tensor
(
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
)
| 73 | |
| 74 | |
| 75 | def kernel_fn( |
| 76 | x: torch.Tensor, |
| 77 | cos: torch.Tensor, |
| 78 | sin: torch.Tensor, |
| 79 | ) -> torch.Tensor: |
| 80 | """ |
| 81 | Entry point called by bench.py. Must match reference.rotary_embedding_ref signature. |
| 82 | |
| 83 | Args: |
| 84 | x: [..., head_dim] tensor to apply rotary embeddings to |
| 85 | cos: [..., head_dim // 2] precomputed cosines |
| 86 | sin: [..., head_dim // 2] precomputed sines |
| 87 | |
| 88 | Returns: |
| 89 | Tensor of same shape as x with rotary embeddings applied. |
| 90 | """ |
| 91 | assert x.is_cuda |
| 92 | |
| 93 | orig_shape = x.shape |
| 94 | head_dim = x.shape[-1] |
| 95 | half_dim = head_dim // 2 |
| 96 | |
| 97 | assert head_dim % 2 == 0, "head_dim must be even for RoPE" |
| 98 | assert cos.shape[-1] == half_dim |
| 99 | assert sin.shape[-1] == half_dim |
| 100 | |
| 101 | # Flatten to 2D: [n_rows, head_dim] |
| 102 | x_flat = x.contiguous().view(-1, head_dim) |
| 103 | n_rows = x_flat.shape[0] |
| 104 | |
| 105 | # Flatten cos/sin and broadcast to match x rows |
| 106 | cos_flat = cos.contiguous().view(-1, half_dim) |
| 107 | sin_flat = sin.contiguous().view(-1, half_dim) |
| 108 | |
| 109 | # Handle broadcasting: if cos/sin have fewer rows, expand to match |
| 110 | if cos_flat.shape[0] < n_rows: |
| 111 | repeat_factor = (n_rows + cos_flat.shape[0] - 1) // cos_flat.shape[0] |
| 112 | cos_flat = cos_flat.repeat(repeat_factor, 1)[:n_rows] |
| 113 | sin_flat = sin_flat.repeat(repeat_factor, 1)[:n_rows] |
| 114 | |
| 115 | out = torch.empty_like(x_flat) |
| 116 | |
| 117 | BLOCK_SIZE = triton.next_power_of_2(half_dim) |
| 118 | |
| 119 | grid = (n_rows,) |
| 120 | rotary_embedding_kernel[grid]( |
| 121 | x_flat, |
| 122 | cos_flat, |
| 123 | sin_flat, |
| 124 | out, |
| 125 | n_rows, |
| 126 | head_dim, |
| 127 | x_flat.stride(0), |
| 128 | cos_flat.stride(0), |
| 129 | sin_flat.stride(0), |
| 130 | out.stride(0), |
| 131 | half_dim, |
| 132 | BLOCK_SIZE=BLOCK_SIZE, |
nothing calls this directly
no outgoing calls
no test coverage detected