chore: import upstream snapshot with attribution

tokenspeed-kernel/python/tokenspeed_kernel/ops/embedding/__init__.py

@@ -0,0 +1,354 @@
+# 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 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.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from tokenspeed_kernel.profiling import ShapeCapture, kernel_scope
+from tokenspeed_kernel.selection import select_kernel
+from tokenspeed_kernel.signature import dense_tensor_format, format_signature
+
+
+@dataclass
+class FusedSetKVBufferArg:
+    value: torch.Tensor
+    k_buffer: torch.Tensor
+    v_buffer: torch.Tensor
+    k_scale: Optional[float]
+    v_scale: Optional[float]
+    cache_loc: torch.Tensor
+
+
+def apply_rope(
+    # embedding inputs
+    positions: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    # embedding options
+    is_neox: bool = True,
+    fused_set_kv_buffer_arg: FusedSetKVBufferArg | None = None,
+    q_rope_out: torch.Tensor | None = None,
+    k_rope_out: torch.Tensor | None = None,
+    # dispatch options
+    enable_pdl: bool = False,
+    solution: str | None = None,
+    override: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply rotary embedding through the registered embedding.rope kernel.
+
+    Args:
+        positions: Token positions. Flattened to [num_tokens] before dispatch.
+        q: Query tensor with shape [num_tokens, num_q_heads * head_size].
+        k: Key tensor with shape [num_tokens, num_kv_heads * head_size].
+        head_size: Per-head hidden dimension.
+        cos_sin_cache: Packed RoPE cache with shape [max_position, rotary_dim]
+            as concat(cos, sin) along the last dimension.
+        is_neox: Whether to use Neox-style half-split rotation. False uses
+            GPT-J interleaved-pair rotation.
+        fused_set_kv_buffer_arg: Optional fused KV-cache write arguments. Both
+            CUDA and Triton implementations currently require k_scale and
+            v_scale to be None.
+        q_rope_out: Optional output buffer for the rotated query. If omitted,
+            q is updated in place.
+        k_rope_out: Optional output buffer for the rotated key. If omitted,
+            k is updated in place.
+        enable_pdl: Passed through to kernels that support PDL.
+        solution: Optional registered solution to select.
+        override: Optional exact kernel-name or solution override.
+
+    Returns:
+        (rotated_q, rotated_k). These are q_rope_out / k_rope_out when provided,
+        otherwise the input q / k.
+    """
+    rotary_dim = cos_sin_cache.shape[-1]
+    assert rotary_dim % 2 == 0, "embedding.rope requires even rotary_dim"
+    assert rotary_dim <= head_size, "embedding.rope requires rotary_dim <= head_size"
+
+    positions = positions.flatten()
+    num_tokens = positions.shape[0]
+    if num_tokens == 0:
+        return (
+            q_rope_out if q_rope_out is not None else q,
+            k_rope_out if k_rope_out is not None else k,
+        )
+    num_q_heads = q.shape[-1] // head_size
+    num_kv_heads = k.shape[-1] // head_size
+
+    traits = {
+        "head_size": head_size,
+        "partial_rotary": rotary_dim != head_size,
+        "is_neox": is_neox,
+        "has_fused_kv": fused_set_kv_buffer_arg is not None,
+        "has_q_out": q_rope_out is not None,
+        "has_k_out": k_rope_out is not None,
+    }
+    signature = format_signature(
+        q=dense_tensor_format(q.dtype),
+        k=dense_tensor_format(k.dtype),
+    )
+    kernel = select_kernel(
+        "embedding",
+        "rope",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+
+    shape_params = {
+        "num_tokens": num_tokens,
+        "num_q_heads": num_q_heads,
+        "num_kv_heads": num_kv_heads,
+        "head_size": head_size,
+        "rotary_dim": rotary_dim,
+        "has_fused_kv": fused_set_kv_buffer_arg is not None,
+        "has_q_out": q_rope_out is not None,
+        "has_k_out": k_rope_out is not None,
+    }
+    ShapeCapture.get().record(
+        "embedding",
+        "rope",
+        kernel.name,
+        q.dtype,
+        shape_params,
+    )
+
+    with kernel_scope(
+        "embedding",
+        "rope",
+        q.dtype,
+        kernel_name=kernel.name,
+        **shape_params,
+    ):
+        kernel(
+            positions=positions,
+            q=q,
+            k=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,
+            q_rope_out=q_rope_out,
+            k_rope_out=k_rope_out,
+            enable_pdl=enable_pdl,
+        )
+
+    return (
+        q_rope_out if q_rope_out is not None else q,
+        k_rope_out if k_rope_out is not None else k,
+    )
+
+
+def apply_rope_mla(
+    # embedding inputs
+    positions: torch.Tensor,
+    q_rope: torch.Tensor,
+    k_rope: torch.Tensor,
+    q_nope: torch.Tensor,
+    k_nope: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    # embedding options
+    is_neox: bool = True,
+    quantize_dtype: torch.dtype = torch.float8_e4m3fn,
+    quant_scale_q: float | torch.Tensor = 1.0,
+    quant_scale_kv: float | torch.Tensor = 1.0,
+    q_rope_out: torch.Tensor | None = None,
+    k_rope_out: torch.Tensor | None = None,
+    q_nope_out: torch.Tensor | None = None,
+    k_nope_out: torch.Tensor | None = None,
+    enable_pdl: bool = False,
+    # dispatch options
+    solution: str | None = None,
+    override: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply MLA RoPE and quantize query/key parts to FP8.
+
+    Args:
+        positions: Token positions. Flattened to [tokens] before dispatch.
+        q_rope: Query RoPE slice with shape [tokens, q_heads, rope_dim].
+        k_rope: Key RoPE slice with shape [tokens, kv_heads, rope_dim].
+        q_nope: Query non-RoPE slice with shape [tokens, q_heads, nope_dim].
+        k_nope: Key non-RoPE slice with shape [tokens, kv_heads, nope_dim].
+        cos_sin_cache: Packed RoPE cache as concat(cos, sin) on the last dim.
+        is_neox: Whether to use Neox-style half-split rotation.
+        quantize_dtype: Output FP8 dtype. Currently only e4m3fn is supported.
+        quant_scale_q: Quantization scale multiplied into query tensors before
+            the FP8 cast.
+        quant_scale_kv: Quantization scale multiplied into key tensors before
+            the FP8 cast.
+        q_rope_out: Optional FP8 output buffer for rotated q_rope. If omitted
+            together with q_nope_out, a combined query output is allocated and
+            this slice is derived from it.
+        k_rope_out: Optional FP8 output buffer for rotated k_rope. If omitted
+            together with k_nope_out, a combined key output is allocated and this
+            slice is derived from it.
+        q_nope_out: Optional FP8 output buffer for q_nope. If omitted together
+            with q_rope_out, a combined query output is allocated and this slice
+            is derived from it.
+        k_nope_out: Optional FP8 output buffer for k_nope. If omitted together
+            with k_rope_out, a combined key output is allocated and this slice is
+            derived from it.
+        enable_pdl: Passed through to kernels that support PDL.
+        solution: Optional registered solution to select.
+        override: Optional exact kernel-name or solution override.
+
+    Returns:
+        (query_fp8, key_fp8), where the last dimension is concat(nope, rope).
+    """
+    positions = positions.flatten()
+    query_fp8 = None
+    key_fp8 = None
+
+    if q_rope_out is None and q_nope_out is None:
+        query_fp8 = torch.empty(
+            q_nope.shape[:-1] + (q_nope.shape[-1] + q_rope.shape[-1],),
+            dtype=quantize_dtype,
+            device=q_nope.device,
+        )
+        q_nope_out = query_fp8[..., : q_nope.shape[-1]]
+        q_rope_out = query_fp8[..., q_nope.shape[-1] :]
+    else:
+        q_rope_out = (
+            torch.empty(q_rope.shape, dtype=quantize_dtype, device=q_rope.device)
+            if q_rope_out is None
+            else q_rope_out
+        )
+        q_nope_out = (
+            torch.empty(q_nope.shape, dtype=quantize_dtype, device=q_nope.device)
+            if q_nope_out is None
+            else q_nope_out
+        )
+
+    if k_rope_out is None and k_nope_out is None:
+        key_fp8 = torch.empty(
+            k_nope.shape[:-1] + (k_nope.shape[-1] + k_rope.shape[-1],),
+            dtype=quantize_dtype,
+            device=k_nope.device,
+        )
+        k_nope_out = key_fp8[..., : k_nope.shape[-1]]
+        k_rope_out = key_fp8[..., k_nope.shape[-1] :]
+    else:
+        k_rope_out = (
+            torch.empty(k_rope.shape, dtype=quantize_dtype, device=k_rope.device)
+            if k_rope_out is None
+            else k_rope_out
+        )
+        k_nope_out = (
+            torch.empty(k_nope.shape, dtype=quantize_dtype, device=k_nope.device)
+            if k_nope_out is None
+            else k_nope_out
+        )
+
+    num_tokens = q_rope.shape[0]
+    if num_tokens == 0:
+        return (
+            (
+                query_fp8
+                if query_fp8 is not None
+                else torch.cat((q_nope_out, q_rope_out), dim=-1)
+            ),
+            (
+                key_fp8
+                if key_fp8 is not None
+                else torch.cat((k_nope_out, k_rope_out), dim=-1)
+            ),
+        )
+
+    traits = {
+        "is_neox": bool(is_neox),
+        "quantize_dtype": quantize_dtype,
+        "has_scale_q_tensor": isinstance(quant_scale_q, torch.Tensor),
+        "has_scale_kv_tensor": isinstance(quant_scale_kv, torch.Tensor),
+    }
+    signature = format_signature(
+        q_rope=dense_tensor_format(q_rope.dtype),
+        k_rope=dense_tensor_format(k_rope.dtype),
+        q_nope=dense_tensor_format(q_nope.dtype),
+        k_nope=dense_tensor_format(k_nope.dtype),
+    )
+    kernel = select_kernel(
+        "embedding",
+        "rope_mla",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+
+    shape_params = {
+        "num_tokens": num_tokens,
+        "q_heads": q_rope.shape[1],
+        "kv_heads": k_rope.shape[1],
+        "q_nope_dim": q_nope.shape[-1],
+        "k_nope_dim": k_nope.shape[-1],
+        "rope_dim": q_rope.shape[-1],
+        "is_neox": bool(is_neox),
+    }
+    ShapeCapture.get().record(
+        "embedding",
+        "rope_mla",
+        kernel.name,
+        q_rope.dtype,
+        shape_params,
+    )
+    with kernel_scope(
+        "embedding",
+        "rope_mla",
+        q_rope.dtype,
+        kernel_name=kernel.name,
+        **shape_params,
+    ):
+        kernel(
+            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,
+        )
+
+    query_fp8 = (
+        query_fp8
+        if query_fp8 is not None
+        else torch.cat((q_nope_out, q_rope_out), dim=-1)
+    )
+    key_fp8 = (
+        key_fp8 if key_fp8 is not None else torch.cat((k_nope_out, k_rope_out), dim=-1)
+    )
+    return query_fp8, key_fp8
+
+
+__all__ = ["FusedSetKVBufferArg", "apply_rope", "apply_rope_mla"]
+
+
+# Backend registration (side-effect imports).
+import tokenspeed_kernel.ops.embedding.cuda  # noqa: E402,F401
+import tokenspeed_kernel.ops.embedding.flashinfer  # noqa: E402,F401
+import tokenspeed_kernel.ops.embedding.triton  # noqa: E402,F401

tokenspeed-kernel/python/tokenspeed_kernel/ops/embedding/cuda.py

@@ -0,0 +1,78 @@
+# 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 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.
+
+"""CUDA rotary embedding kernels."""
+
+from typing import Any
+
+import torch
+from tokenspeed_kernel.platform import CapabilityRequirement, current_platform
+from tokenspeed_kernel.registry import Priority, register_kernel
+from tokenspeed_kernel.signature import format_signatures
+
+platform = current_platform()
+
+if platform.is_nvidia:
+    from tokenspeed_kernel.thirdparty.cuda.rope import (
+        apply_rope_with_cos_sin_cache_inplace,
+    )
+
+    @register_kernel(
+        "embedding",
+        "rope",
+        name="cuda_embedding_rope",
+        solution="cuda",
+        capability=CapabilityRequirement(vendors=frozenset({"nvidia"})),
+        signatures=format_signatures(
+            ("q", "k"), "dense", {torch.float16, torch.bfloat16}
+        ),
+        priority=Priority.PERFORMANT,
+        traits={
+            "head_size": frozenset({64, 128, 256, 512}),
+            "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={"latency"},
+    )
+    def cuda_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_with_cos_sin_cache_inplace(
+            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,
+            enable_pdl=enable_pdl,
+        )

tokenspeed-kernel/python/tokenspeed_kernel/ops/embedding/flashinfer.py

@@ -0,0 +1,83 @@
+# 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 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.
+
+"""FlashInfer rotary embedding kernels."""
+
+import torch
+from tokenspeed_kernel.platform import CapabilityRequirement, current_platform
+from tokenspeed_kernel.registry import Priority, register_kernel
+from tokenspeed_kernel.signature import format_signatures
+
+platform = current_platform()
+
+
+if platform.is_nvidia:
+    from flashinfer.rope import mla_rope_quantize_fp8
+
+    @register_kernel(
+        "embedding",
+        "rope_mla",
+        name="flashinfer_embedding_rope_mla",
+        solution="flashinfer",
+        capability=CapabilityRequirement(vendors=frozenset({"nvidia"})),
+        signatures=format_signatures(
+            ("q_rope", "k_rope", "q_nope", "k_nope"),
+            "dense",
+            {torch.float16, torch.bfloat16},
+        ),
+        priority=Priority.SPECIALIZED,
+        traits={
+            "is_neox": frozenset({True, False}),
+            "quantize_dtype": frozenset({torch.float8_e4m3fn}),
+            "has_scale_q_tensor": frozenset({False}),
+            "has_scale_kv_tensor": frozenset({False}),
+        },
+    )
+    def flashinfer_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 = 1.0,
+        quant_scale_kv: float = 1.0,
+        enable_pdl: bool = False,
+    ) -> None:
+        mla_rope_quantize_fp8(
+            q_rope=q_rope,
+            k_rope=k_rope,
+            q_nope=q_nope,
+            k_nope=k_nope,
+            cos_sin_cache=cos_sin_cache,
+            pos_ids=positions,
+            is_neox=is_neox,
+            quantize_dtype=torch.float8_e4m3fn,
+            q_rope_out=q_rope_out,
+            k_rope_out=k_rope_out,
+            q_nope_out=q_nope_out,
+            k_nope_out=k_nope_out,
+            quant_scale_q=quant_scale_q,
+            quant_scale_kv=quant_scale_kv,
+            enable_pdl=enable_pdl,
+        )

tokenspeed-kernel/python/tokenspeed_kernel/ops/embedding/triton.py

@@ -0,0 +1,572 @@
+# 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,
+    )