# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. """Triton fused rotary embedding kernels.""" from __future__ import annotations from typing import Any, Optional import torch from tokenspeed_kernel._triton import tl, triton from tokenspeed_kernel.platform import CapabilityRequirement from tokenspeed_kernel.registry import Priority, register_kernel from tokenspeed_kernel.signature import format_signatures def _next_power_of_2(n: int) -> int: p = 1 while p < n: p <<= 1 return p @triton.jit def _rope_apply_kernel( q_ptr, k_ptr, q_out_ptr, k_out_ptr, cos_sin_cache_ptr, positions_ptr, offsets_ptr, value_ptr, k_buffer_ptr, v_buffer_ptr, cache_loc_ptr, q_stride_t, q_stride_h, k_stride_t, k_stride_h, q_out_stride_t, q_out_stride_h, k_out_stride_t, k_out_stride_h, value_stride_t, value_stride_h, k_buffer_stride_t, k_buffer_stride_h, v_buffer_stride_t, v_buffer_stride_h, cache_stride_p, num_q_heads, num_k_heads, head_size, rotary_dim, HALF_DIM_PADDED: tl.constexpr, HEAD_DIM_PADDED: tl.constexpr, HAS_OFFSETS: tl.constexpr, HAS_Q_OUT: tl.constexpr, HAS_K_OUT: tl.constexpr, HAS_FUSED_KV: tl.constexpr, IS_NEOX: tl.constexpr, POSITION_INT64: tl.constexpr, CACHE_LOC_INT64: tl.constexpr, ): """Apply rotary embedding to one (token, head) pair in-place. Grid: (num_tokens, num_q_heads + num_k_heads). Heads in [0, num_q_heads) belong to Q; heads in [num_q_heads, num_q_heads + num_k_heads) belong to K. Each program loads cos/sin for `rotary_dim // 2` channels, applies the NEOX or GPT-J style rotation to the first `rotary_dim` lanes of the head, and leaves the trailing `head_size - rotary_dim` lanes untouched. """ token_idx = tl.program_id(0) head_idx = tl.program_id(1) is_query = head_idx < num_q_heads kv_head_idx = head_idx - num_q_heads if is_query: base_ptr = q_ptr + token_idx * q_stride_t + head_idx * q_stride_h out_ptr = ( q_out_ptr + token_idx * q_out_stride_t + head_idx * q_out_stride_h if HAS_Q_OUT else base_ptr ) else: base_ptr = k_ptr + token_idx * k_stride_t + kv_head_idx * k_stride_h out_ptr = ( k_out_ptr + token_idx * k_out_stride_t + kv_head_idx * k_out_stride_h if HAS_K_OUT else base_ptr ) if POSITION_INT64: pos = tl.load(positions_ptr + token_idx).to(tl.int64) else: pos = tl.load(positions_ptr + token_idx).to(tl.int32) if HAS_OFFSETS: if POSITION_INT64: pos = pos + tl.load(offsets_ptr + token_idx).to(tl.int64) else: pos = pos + tl.load(offsets_ptr + token_idx).to(tl.int32) half = rotary_dim // 2 half_offs = tl.arange(0, HALF_DIM_PADDED) half_mask = half_offs < half cos = tl.load( cos_sin_cache_ptr + pos * cache_stride_p + half_offs, mask=half_mask, other=0.0, ).to(tl.float32) sin = tl.load( cos_sin_cache_ptr + pos * cache_stride_p + half + half_offs, mask=half_mask, other=0.0, ).to(tl.float32) if IS_NEOX: # NEOX layout: x is split into [first_half | second_half]. # Output: [x1 * cos - x2 * sin, x2 * cos + x1 * sin]. x1 = tl.load(base_ptr + half_offs, mask=half_mask, other=0.0) x2 = tl.load(base_ptr + half + half_offs, mask=half_mask, other=0.0) x1_f = x1.to(tl.float32) x2_f = x2.to(tl.float32) o1 = x1_f * cos - x2_f * sin o2 = x2_f * cos + x1_f * sin tl.store(out_ptr + half_offs, o1.to(x1.dtype), mask=half_mask) tl.store(out_ptr + half + half_offs, o2.to(x2.dtype), mask=half_mask) else: # GPT-J layout: x is interleaved [x0, x1, x0, x1, ...]. # Pairs are (x[2i], x[2i+1]); output: # y[2i] = x[2i] * cos - x[2i+1] * sin # y[2i+1] = x[2i+1] * cos + x[2i] * sin x1 = tl.load(base_ptr + 2 * half_offs, mask=half_mask, other=0.0) x2 = tl.load(base_ptr + 2 * half_offs + 1, mask=half_mask, other=0.0) x1_f = x1.to(tl.float32) x2_f = x2.to(tl.float32) o1 = x1_f * cos - x2_f * sin o2 = x2_f * cos + x1_f * sin tl.store(out_ptr + 2 * half_offs, o1.to(x1.dtype), mask=half_mask) tl.store(out_ptr + 2 * half_offs + 1, o2.to(x2.dtype), mask=half_mask) head_offs = tl.arange(0, HEAD_DIM_PADDED) tail_mask = (head_offs >= rotary_dim) & (head_offs < head_size) if HAS_Q_OUT or HAS_K_OUT: tail = tl.load(base_ptr + head_offs, mask=tail_mask, other=0.0) tl.store(out_ptr + head_offs, tail, mask=tail_mask) if HAS_FUSED_KV and not is_query: if CACHE_LOC_INT64: cache_loc = tl.load(cache_loc_ptr + token_idx).to(tl.int64) else: cache_loc = tl.load(cache_loc_ptr + token_idx).to(tl.int32) head_mask = head_offs < head_size k_value = tl.load(out_ptr + head_offs, mask=head_mask, other=0.0) v_value = tl.load( value_ptr + token_idx * value_stride_t + kv_head_idx * value_stride_h + head_offs, mask=head_mask, other=0.0, ) tl.store( k_buffer_ptr + cache_loc * k_buffer_stride_t + kv_head_idx * k_buffer_stride_h + head_offs, k_value, mask=head_mask, ) tl.store( v_buffer_ptr + cache_loc * v_buffer_stride_t + kv_head_idx * v_buffer_stride_h + head_offs, v_value, mask=head_mask, ) def apply_rope_triton( positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, head_size: int, cos_sin_cache: torch.Tensor, is_neox: bool = True, offsets: Optional[torch.Tensor] = None, rotary_dim: Optional[int] = None, fused_set_kv_buffer_arg=None, output_q_rope: Optional[torch.Tensor] = None, output_k_rope: Optional[torch.Tensor] = None, ) -> None: """Apply rotary positional embedding to query and key in-place. Args: positions: Token positions, 1D [num_tokens]. int32 or int64. query: [num_tokens, num_q_heads * head_size] (will be viewed as [num_tokens, num_q_heads, head_size]). key: [num_tokens, num_k_heads * head_size] (will be viewed as [num_tokens, num_k_heads, head_size]). head_size: Per-head dimension. cos_sin_cache: [max_position, rotary_dim] packed as concat(cos, sin) along the last dimension. Float32 is strongly recommended for numerical stability; other dtypes are accepted. is_neox: If True, use NEOX-style rotation (x split in halves). If False, use GPT-J-style rotation (interleaved pairs). offsets: Optional [num_tokens] int tensor added to positions. rotary_dim: Rotary dimension. Defaults to cos_sin_cache.shape[-1]. Must be even and <= head_size. """ assert ( positions.dim() == 1 ), f"triton rope expects 1D positions, got shape {tuple(positions.shape)}" assert positions.dtype in ( torch.int32, torch.int64, ), f"positions dtype must be int32 or int64, got {positions.dtype}" assert ( query.dtype == key.dtype ), f"query/key dtype mismatch: {query.dtype} vs {key.dtype}" if rotary_dim is None: rotary_dim = cos_sin_cache.shape[-1] assert rotary_dim % 2 == 0, f"rotary_dim must be even, got {rotary_dim}" assert ( rotary_dim <= head_size ), f"rotary_dim ({rotary_dim}) must be <= head_size ({head_size})" assert cos_sin_cache.shape[-1] == rotary_dim, ( f"cos_sin_cache last dim ({cos_sin_cache.shape[-1]}) must equal " f"rotary_dim ({rotary_dim})" ) num_tokens = positions.shape[0] if num_tokens == 0: return q_view = query.view(num_tokens, -1, head_size) k_view = key.view(num_tokens, -1, head_size) num_q_heads = q_view.shape[1] num_k_heads = k_view.shape[1] if offsets is not None: assert ( offsets.dim() == 1 and offsets.shape[0] == num_tokens ), f"offsets must have shape [{num_tokens}], got {tuple(offsets.shape)}" if fused_set_kv_buffer_arg is not None: if ( fused_set_kv_buffer_arg.k_scale is not None or fused_set_kv_buffer_arg.v_scale is not None ): raise ValueError("k_scale/v_scale are not supported yet") if fused_set_kv_buffer_arg.cache_loc is None: raise ValueError("fused_set_kv_buffer_arg.cache_loc is required") if fused_set_kv_buffer_arg.cache_loc.dtype not in (torch.int32, torch.int64): raise ValueError( f"cache_loc must be int32 or int64, got {fused_set_kv_buffer_arg.cache_loc.dtype}" ) half = rotary_dim // 2 half_padded = max(_next_power_of_2(half), 16) head_padded = max(_next_power_of_2(head_size), 16) q_out_view = ( output_q_rope.view(num_tokens, num_q_heads, head_size) if output_q_rope is not None else q_view ) k_out_view = ( output_k_rope.view(num_tokens, num_k_heads, head_size) if output_k_rope is not None else k_view ) if fused_set_kv_buffer_arg is not None: value = fused_set_kv_buffer_arg.value value_view = value.view(num_tokens, num_k_heads, -1) assert ( value_view.shape[-1] == head_size ), f"fused value head size {value_view.shape[-1]} must match head_size {head_size}" k_buffer_view = fused_set_kv_buffer_arg.k_buffer.view( fused_set_kv_buffer_arg.k_buffer.shape[0], num_k_heads, head_size ) v_buffer_view = fused_set_kv_buffer_arg.v_buffer.view( fused_set_kv_buffer_arg.v_buffer.shape[0], num_k_heads, head_size ) cache_loc = fused_set_kv_buffer_arg.cache_loc else: value_view = k_view k_buffer_view = k_view v_buffer_view = k_view cache_loc = positions grid = (num_tokens, num_q_heads + num_k_heads) _rope_apply_kernel[grid]( q_view, k_view, q_out_view, k_out_view, cos_sin_cache, positions, offsets if offsets is not None else positions, value_view, k_buffer_view, v_buffer_view, cache_loc, q_view.stride(0), q_view.stride(1), k_view.stride(0), k_view.stride(1), q_out_view.stride(0), q_out_view.stride(1), k_out_view.stride(0), k_out_view.stride(1), value_view.stride(0), value_view.stride(1), k_buffer_view.stride(0), k_buffer_view.stride(1), v_buffer_view.stride(0), v_buffer_view.stride(1), cos_sin_cache.stride(0), num_q_heads, num_k_heads, head_size, rotary_dim, HALF_DIM_PADDED=half_padded, HEAD_DIM_PADDED=head_padded, HAS_OFFSETS=offsets is not None, HAS_Q_OUT=output_q_rope is not None, HAS_K_OUT=output_k_rope is not None, HAS_FUSED_KV=fused_set_kv_buffer_arg is not None, IS_NEOX=bool(is_neox), POSITION_INT64=positions.dtype == torch.int64, CACHE_LOC_INT64=cache_loc.dtype == torch.int64, ) @triton.jit def _fp8_quantize_kernel( x, out, scale, x_stride_t: tl.constexpr, x_stride_h: tl.constexpr, out_stride_t: tl.constexpr, out_stride_h: tl.constexpr, num_heads: tl.constexpr, n_cols: tl.constexpr, BLOCK_N: tl.constexpr, HAS_SCALE_TENSOR: tl.constexpr, ): token = tl.program_id(0) head = tl.program_id(1) offsets = tl.arange(0, BLOCK_N) mask = offsets < n_cols values = tl.load( x + token * x_stride_t + head * x_stride_h + offsets, mask=mask, other=0.0, ).to(tl.float32) if HAS_SCALE_TENSOR: scale = tl.load(scale) values = values * scale values_fp8 = values.to(tl.float8e4nv) tl.store( out + token * out_stride_t + head * out_stride_h + offsets, values_fp8, mask=(head < num_heads) & mask, ) def _fp8_quantize( x: torch.Tensor, out: torch.Tensor, scale: float | torch.Tensor, *, enable_pdl: bool, ) -> None: if x.dim() != 3 or out.dim() != 3: raise ValueError( f"MLA FP8 quantize expects rank-3 tensors, got {x.shape} and {out.shape}" ) if x.shape != out.shape: raise ValueError(f"MLA FP8 quantize shape mismatch: {x.shape} vs {out.shape}") if out.dtype != torch.float8_e4m3fn: raise TypeError(f"MLA FP8 quantize output must be e4m3fn, got {out.dtype}") if isinstance(scale, torch.Tensor): scale = scale.contiguous() block_n = max(16, _next_power_of_2(x.shape[-1])) extra_kwargs = {"launch_pdl": True} if enable_pdl else {} _fp8_quantize_kernel[(x.shape[0], x.shape[1])]( x, out, scale, x.stride(0), x.stride(1), out.stride(0), out.stride(1), num_heads=x.shape[1], n_cols=x.shape[2], BLOCK_N=block_n, HAS_SCALE_TENSOR=isinstance(scale, torch.Tensor), num_warps=4, num_stages=1, **extra_kwargs, ) def mla_rope_quantize_fp8_triton( *, positions: torch.Tensor, q_rope: torch.Tensor, k_rope: torch.Tensor, q_nope: torch.Tensor, k_nope: torch.Tensor, cos_sin_cache: torch.Tensor, q_rope_out: torch.Tensor, k_rope_out: torch.Tensor, q_nope_out: torch.Tensor, k_nope_out: torch.Tensor, is_neox: bool = True, quant_scale_q: float | torch.Tensor = 1.0, quant_scale_kv: float | torch.Tensor = 1.0, enable_pdl: bool = False, ) -> None: if q_rope.shape[-1] != k_rope.shape[-1]: raise ValueError( "q_rope and k_rope must have the same rope dim, got " f"{q_rope.shape[-1]} and {k_rope.shape[-1]}" ) if q_rope.shape[0] != k_rope.shape[0] or q_rope.shape[0] != positions.numel(): raise ValueError( "MLA RoPE token count mismatch: " f"q={q_rope.shape[0]}, k={k_rope.shape[0]}, pos={positions.numel()}" ) q_rope_tmp = torch.empty(q_rope.shape, dtype=q_rope.dtype, device=q_rope.device) k_rope_tmp = torch.empty(k_rope.shape, dtype=k_rope.dtype, device=k_rope.device) apply_rope_triton( positions=positions, query=q_rope, key=k_rope, head_size=q_rope.shape[-1], cos_sin_cache=cos_sin_cache, is_neox=is_neox, rotary_dim=q_rope.shape[-1], output_q_rope=q_rope_tmp, output_k_rope=k_rope_tmp, ) _fp8_quantize(q_rope_tmp, q_rope_out, quant_scale_q, enable_pdl=enable_pdl) _fp8_quantize(k_rope_tmp, k_rope_out, quant_scale_kv, enable_pdl=enable_pdl) _fp8_quantize(q_nope, q_nope_out, quant_scale_q, enable_pdl=enable_pdl) _fp8_quantize(k_nope, k_nope_out, quant_scale_kv, enable_pdl=enable_pdl) @register_kernel( "embedding", "rope", name="triton_embedding_rope", solution="triton", capability=CapabilityRequirement(vendors=frozenset({"amd", "nvidia"})), signatures=format_signatures(("q", "k"), "dense", {torch.float16, torch.bfloat16}), priority=Priority.PORTABLE, traits={ "partial_rotary": frozenset({True, False}), "is_neox": frozenset({True, False}), "has_fused_kv": frozenset({True, False}), "has_q_out": frozenset({True, False}), "has_k_out": frozenset({True, False}), }, tags={"portability"}, ) def triton_embedding_rope( *, positions: torch.Tensor, q: torch.Tensor, k: torch.Tensor, head_size: int, cos_sin_cache: torch.Tensor, is_neox: bool = True, fused_set_kv_buffer_arg: Any = None, q_rope_out: torch.Tensor | None = None, k_rope_out: torch.Tensor | None = None, enable_pdl: bool = False, ) -> None: apply_rope_triton( positions=positions, query=q, key=k, head_size=head_size, cos_sin_cache=cos_sin_cache, is_neox=is_neox, fused_set_kv_buffer_arg=fused_set_kv_buffer_arg, output_q_rope=q_rope_out, output_k_rope=k_rope_out, ) @register_kernel( "embedding", "rope_mla", name="triton_embedding_rope_mla", solution="triton", capability=CapabilityRequirement(vendors=frozenset({"amd", "nvidia"})), signatures=format_signatures( ("q_rope", "k_rope", "q_nope", "k_nope"), "dense", {torch.float16, torch.bfloat16}, ), priority=Priority.PORTABLE, traits={ "is_neox": frozenset({True, False}), "quantize_dtype": frozenset({torch.float8_e4m3fn}), "has_scale_q_tensor": frozenset({True, False}), "has_scale_kv_tensor": frozenset({True, False}), }, tags={"portability"}, ) def triton_embedding_rope_mla( *, positions: torch.Tensor, q_rope: torch.Tensor, k_rope: torch.Tensor, q_nope: torch.Tensor, k_nope: torch.Tensor, cos_sin_cache: torch.Tensor, q_rope_out: torch.Tensor, k_rope_out: torch.Tensor, q_nope_out: torch.Tensor, k_nope_out: torch.Tensor, is_neox: bool = True, quant_scale_q: float | torch.Tensor = 1.0, quant_scale_kv: float | torch.Tensor = 1.0, enable_pdl: bool = False, ) -> None: mla_rope_quantize_fp8_triton( positions=positions, q_rope=q_rope, k_rope=k_rope, q_nope=q_nope, k_nope=k_nope, cos_sin_cache=cos_sin_cache, q_rope_out=q_rope_out, k_rope_out=k_rope_out, q_nope_out=q_nope_out, k_nope_out=k_nope_out, is_neox=is_neox, quant_scale_q=quant_scale_q, quant_scale_kv=quant_scale_kv, enable_pdl=enable_pdl, )