sgl-project--sglang/python/sglang/srt/models/utils.py

# Copyright 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from __future__ import annotations

import itertools
from collections.abc import Iterable, Mapping
from dataclasses import dataclass, field
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Optional, Tuple

import numpy as np
import torch
import triton
import triton.language as tl

from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
from sglang.srt.environ import envs
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils import get_current_device_stream_fast, is_cuda, is_hip
from sglang.srt.utils.custom_op import register_custom_op

if TYPE_CHECKING:
    from sglang.srt.layers.layernorm import RMSNorm

_is_cuda = is_cuda()
_is_hip = is_hip()

WeightsMapping = Mapping[str, Optional[str]]
"""If a key maps to a value of `None`, the corresponding weight is ignored."""


@dataclass
class WeightsMapper:
    """Maps the name of each weight if they match the following patterns."""

    orig_to_new_substr: WeightsMapping = field(default_factory=dict)
    orig_to_new_prefix: WeightsMapping = field(default_factory=dict)
    orig_to_new_suffix: WeightsMapping = field(default_factory=dict)

    def __or__(self, other: WeightsMapper) -> WeightsMapper:
        return WeightsMapper(
            orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr},
            orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix},
            orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},
        )

    def _map_name(self, key: str) -> Optional[str]:
        for substr, new_key in sorted(
            self.orig_to_new_substr.items(), key=lambda i: len(i[0]), reverse=True
        ):
            if substr in key:
                if new_key is None:
                    return None

                key = key.replace(substr, new_key, 1)
                break

        for prefix, new_key in sorted(
            self.orig_to_new_prefix.items(), key=lambda i: len(i[0]), reverse=True
        ):
            if key.startswith(prefix):
                if new_key is None:
                    return None

                key = key.replace(prefix, new_key, 1)
                break

        for suffix, new_key in sorted(
            self.orig_to_new_suffix.items(), key=lambda i: len(i[0]), reverse=True
        ):
            if key.endswith(suffix):
                if new_key is None:
                    return None

                key = new_key.join(key.rsplit(suffix, 1))
                break

        return key

    def apply(
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[tuple[str, torch.Tensor]]:
        return (
            (out_name, data)
            for name, data in weights
            if (out_name := self._map_name(name)) is not None
        )

    def apply_list(self, values: list[str]) -> list[str]:
        return [
            out_name
            for name in values
            if (out_name := self._map_name(name)) is not None
        ]

    def apply_dict(self, values: dict[str, Any]) -> dict[str, Any]:
        return {
            out_name: value
            for name, value in values.items()
            if (out_name := self._map_name(name)) is not None
        }


class AutoWeightsLoader:
    ROTARY_EMBEDS_UNUSED_WEIGHTS = [
        "rotary_pos_emb.inv_freq",
        "rotary_emb.inv_freq",
        "rotary_emb.cos_cached",
        "rotary_emb.sin_cached",
    ]

    def __init__(
        self,
        module: torch.nn.Module,
        *,
        skip_prefixes: list[str] | None = None,
        skip_substrs: list[str] | None = None,
        ignore_unexpected_prefixes: list[str] | None = None,
        ignore_unexpected_suffixes: list[str] | None = None,
    ) -> None:
        self.module = module
        self.skip_prefixes = list(skip_prefixes or [])
        self.skip_substrs = [
            *(skip_substrs or []),
            *self.ROTARY_EMBEDS_UNUSED_WEIGHTS,
        ]
        self.ignore_unexpected_prefixes = list(ignore_unexpected_prefixes or [])
        self.ignore_unexpected_suffixes = list(ignore_unexpected_suffixes or [])

    def _groupby_prefix(
        self,
        weights: Iterable[tuple[str, torch.Tensor]],
    ) -> Iterable[tuple[str, Iterable[tuple[str, torch.Tensor]]]]:
        weights_by_parts = (
            (weight_name.split(".", 1), weight_data)
            for weight_name, weight_data in weights
        )
        for prefix, group in itertools.groupby(weights_by_parts, key=lambda x: x[0][0]):
            yield (
                prefix,
                (
                    ("" if len(parts) == 1 else parts[1], weight_data)
                    for parts, weight_data in group
                ),
            )

    @staticmethod
    def _get_qualname(prefix: str, rest: str) -> str:
        if prefix == "":
            return rest
        if rest == "":
            return prefix
        return f"{prefix}.{rest}"

    def _can_skip(self, qualname: str) -> bool:
        return any(qualname.startswith(p) for p in self.skip_prefixes) or any(
            sub in qualname for sub in self.skip_substrs
        )

    def _can_ignore_unexpected(self, qualname: str) -> bool:
        return any(
            qualname.startswith(p) for p in self.ignore_unexpected_prefixes
        ) or any(qualname.endswith(s) for s in self.ignore_unexpected_suffixes)

    def _load_param(
        self,
        base_prefix: str,
        param: torch.nn.Parameter,
        weights: Iterable[tuple[str, torch.Tensor]],
    ) -> Iterable[str]:
        for weight_name, weight_data in weights:
            weight_qualname = self._get_qualname(base_prefix, weight_name)
            if self._can_skip(weight_qualname):
                continue
            if weight_name != "":
                if self._can_ignore_unexpected(weight_qualname):
                    continue
                raise ValueError(
                    f"Attempted to load nested weight {weight_qualname!r} "
                    f"into parameter {base_prefix!r}"
                )

            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, weight_data)
            yield weight_qualname

    def _load_module(
        self,
        base_prefix: str,
        module: torch.nn.Module,
        weights: Iterable[tuple[str, torch.Tensor]],
    ) -> Iterable[str]:
        if module.__class__.__name__ == "PPMissingLayer":
            return

        if module is not self.module:
            module_load_weights = getattr(module, "load_weights", None)
            if callable(module_load_weights):
                loaded = module_load_weights(weights)
                if loaded is not None:
                    yield from (
                        self._get_qualname(base_prefix, loaded_name)
                        for loaded_name in loaded
                    )
                return

        child_modules = dict(module.named_children())
        child_params = dict(module.named_parameters(recurse=False))
        child_buffers = dict(module.named_buffers(recurse=False))
        for child_prefix, child_weights in self._groupby_prefix(weights):
            prefix = self._get_qualname(base_prefix, child_prefix)
            if child_prefix in child_modules:
                if self._can_skip(prefix + "."):
                    continue
                yield from self._load_module(
                    prefix,
                    child_modules[child_prefix],
                    child_weights,
                )
                continue

            if child_prefix in child_params:
                if self._can_skip(prefix):
                    continue
                yield from self._load_param(
                    prefix, child_params[child_prefix], child_weights
                )
                continue

            if child_prefix in child_buffers:
                if self._can_skip(prefix):
                    continue
                yield from self._load_param(
                    prefix, child_buffers[child_prefix], child_weights
                )
                continue

            if self._can_skip(prefix) or self._can_skip(prefix + "."):
                continue
            if self._can_ignore_unexpected(prefix) or self._can_ignore_unexpected(
                prefix + "."
            ):
                continue
            raise ValueError(
                f"No module or parameter named {prefix!r} in {self.module._get_name()}."
            )

    def load_weights(
        self,
        weights: Iterable[tuple[str, torch.Tensor]],
        *,
        mapper: WeightsMapper | None = None,
    ) -> set[str]:
        if mapper is not None:
            weights = mapper.apply(weights)
        weights = (
            (name, weight) for name, weight in weights if not self._can_skip(name)
        )
        return set(self._load_module("", self.module, weights))


def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
    """Enable fused set_kv_buffer on CUDA with bfloat16 KV cache and HIP with bf16/fp16/fp8 KV cache.

    SHUFFLE 5D pools on HIP also work — the underlying triton kernel
    (`fused_qk_rope_reshape_and_cache`) natively supports the 5D
    SHUFFLE layout (key_cache.ndim==5, value_cache.ndim==5). We just need
    the per-layer arg builder to pass the raw 5D buffers without the
    `.view(-> 4D NHD)` reshape, and let the rotary forward pass
    `flash_layout=False`. See `create_fused_set_kv_buffer_arg` below.
    """
    pool = get_token_to_kv_pool()
    return (
        _is_cuda
        and pool.dtype == torch.bfloat16
        and not isinstance(pool, SWAKVPool)
        and not is_prefill_context_parallel_enabled()
        and getattr(forward_batch, "dcp_kv_mask", None) is None
    ) or (
        _is_hip
        and not is_prefill_context_parallel_enabled()
        and getattr(forward_batch, "dcp_kv_mask", None) is None
    )


def create_fused_set_kv_buffer_arg(
    value: torch.Tensor,
    layer: RadixAttention,
    forward_batch: ForwardBatch,
):
    from sglang.jit_kernel.rope import FusedSetKVBufferArg

    layer_id = layer.layer_id
    token_to_kv_pool = get_token_to_kv_pool()

    k_buffer = token_to_kv_pool.get_key_buffer(layer_id)
    v_buffer = token_to_kv_pool.get_value_buffer(layer_id)

    if not _is_hip:
        assert layer.k_scale is None and layer.v_scale is None, "scale not supported"
        return FusedSetKVBufferArg(
            value=value,
            k_buffer=k_buffer.view(k_buffer.shape[0], -1),
            v_buffer=v_buffer.view(v_buffer.shape[0], -1),
            cache_loc=forward_batch.out_cache_loc,
        )
    else:
        page_size = token_to_kv_pool.page_size
        slot_mapping_swa = (
            token_to_kv_pool.full_to_swa_index_mapping.long()
            if layer.sliding_window_size > 0
            else None
        )
        # SHUFFLE 5D pools (k_buffer.ndim == 5) consumed natively by
        # fused_qk_rope_reshape_and_cache via flash_layout=False. For the
        # legacy NHD 3D pool we reshape to the (num_blocks, page_size, H, D)
        # paged view the kernel expects under flash_layout=True.
        if k_buffer.ndim == 5:
            key_cache = k_buffer
            value_cache = v_buffer
        else:
            key_cache = k_buffer.view(
                -1, page_size, layer.tp_k_head_num, layer.qk_head_dim
            )
            value_cache = v_buffer.view(
                -1, page_size, layer.tp_v_head_num, layer.v_head_dim
            )
        return {
            "v": value.view(-1, layer.tp_v_head_num, layer.v_head_dim),
            "k_scale": layer.k_scale,
            "v_scale": layer.v_scale,
            "key_cache": key_cache,
            "value_cache": value_cache,
            "slot_mapping": forward_batch.out_cache_loc,
            "swa_slot_mapping": slot_mapping_swa,
        }


def permute_inv(perm: torch.Tensor) -> torch.Tensor:
    inv_perm = torch.empty_like(perm)
    inv_perm[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype)
    return inv_perm


def compute_cu_seqlens_from_grid_numpy(grid_thw: torch.Tensor) -> torch.Tensor:
    """
    Compute cu_seqlens from grid_thw using NumPy.

    grid_thw: [T, 3] int tensor on CPU.
              columns: [repeat_count, H, W]
    Returns:
        cu_seqlens: 1D int32 tensor on CPU, shape [N + 1]
    """
    assert (
        grid_thw.device.type == "cpu"
    ), "compute_cu_seqlens_from_grid_numpy expects a CPU tensor"
    arr = grid_thw.numpy()

    cu_seqlens = np.repeat(arr[:, 1] * arr[:, 2], arr[:, 0]).cumsum(
        axis=0, dtype=np.int32
    )
    cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
    cu_seqlens = torch.from_numpy(cu_seqlens)
    return cu_seqlens


class RotaryPosMixin:
    @staticmethod
    @lru_cache(maxsize=1024)
    def rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor:
        if isinstance(h, torch.Tensor):
            h = int(h.item())
        if isinstance(w, torch.Tensor):
            w = int(w.item())
        if isinstance(spatial_merge_size, torch.Tensor):
            spatial_merge_size = int(spatial_merge_size.item())
        hpos_ids = np.broadcast_to(np.arange(h).reshape(h, 1), (h, w))
        h_div = h // spatial_merge_size
        w_div = w // spatial_merge_size
        hpos_ids = hpos_ids.reshape(
            h_div,
            spatial_merge_size,
            w_div,
            spatial_merge_size,
        )
        hpos_ids = hpos_ids.transpose(0, 2, 1, 3)
        hpos_ids = hpos_ids.flatten()

        wpos_ids = np.broadcast_to(np.arange(w).reshape(1, w), (h, w))
        wpos_ids = wpos_ids.reshape(
            h_div,
            spatial_merge_size,
            w_div,
            spatial_merge_size,
        )
        wpos_ids = wpos_ids.transpose(0, 2, 1, 3)
        wpos_ids = wpos_ids.flatten()

        return torch.from_numpy(np.stack([hpos_ids, wpos_ids], axis=-1))


def _reshape_for_qk_norm(x: torch.Tensor, head_dim: int) -> torch.Tensor:
    """Reshape a (..., H*D) tensor into (..., H, D) ahead of QK RMSNorm.

    On CUDA with the inductor piecewise-cuda-graph compiler, return a
    stride-preserving view so inductor can fuse this reshape with the
    subsequent RMSNorm (and any upstream/downstream FP8 quant) into a
    single triton kernel -- the original motivation of #21734.

    Everywhere else (ROCm, or CUDA with the eager PCG fallback), use the
    flat 2D reshape that forces a copy when the input is a non-contiguous
    QKV-split stride-trick view. ROCm's RMSNorm kernels assume contiguous
    inputs and fault on strided tensors (root cause of the #21734 revert
    in #23159).
    """

    if (
        _is_cuda
        and get_server_args().cuda_graph_config.prefill.tc_compiler == "inductor"
    ):
        return x.view(*x.shape[:-1], -1, head_dim)
    return x.reshape(-1, head_dim)


def apply_qk_norm(
    q: torch.Tensor,
    k: torch.Tensor,
    q_norm: RMSNorm,
    k_norm: RMSNorm,
    head_dim: int,
    alt_stream: Optional[torch.cuda.Stream] = None,
    allow_inplace: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Apply QK normalization for query and key tensors.
    If eligible, we will use JIT fused inplace QK normalization for better performance.

    Args:
        q: Query tensor of shape [batch_size, ...]
        k: Key tensor of shape [batch_size, ...]
        q_norm: RMSNorm layer for query normalization
        k_norm: RMSNorm layer for key normalization
        head_dim: Dimension of each attention head
        alt_stream: Optional alternative CUDA stream for overlapping computation
        allow_inplace: Whether to allow inplace normalization. (True for better performance)

    Returns:
        Tuple of normalized query and key tensors
    """

    batch_size = q.size(0)
    q_eps = q_norm.variance_epsilon
    k_eps = k_norm.variance_epsilon

    if (
        _is_cuda  # TODO(dark): have not tested on ROCm or other backends
        and allow_inplace  # TODO(dark): this can be relaxed if needed
        and (q_eps == k_eps)  # TODO(dark): this can also be relaxed
        and not envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
        and get_server_args().cuda_graph_config.prefill.tc_compiler
        != "inductor"  # let inductor fuse QK norm
        and can_use_fused_inplace_qknorm(head_dim, q.dtype)
    ):
        fused_inplace_qknorm(
            q=q.view(batch_size, -1, head_dim),
            k=k.view(batch_size, -1, head_dim),
            q_weight=q_norm.weight,
            k_weight=k_norm.weight,
            head_dim=head_dim,
            eps=q_eps,
        )
        return q, k

    if alt_stream is not None and get_is_capture_mode():
        current_stream = get_current_device_stream_fast()
        alt_stream.wait_stream(current_stream)
        q_by_head = _reshape_for_qk_norm(q, head_dim)
        q_by_head = q_norm(q_by_head)
        with torch.cuda.stream(alt_stream):
            k_by_head = _reshape_for_qk_norm(k, head_dim)
            k_by_head = k_norm(k_by_head)
        current_stream.wait_stream(alt_stream)
    else:
        q_by_head = _reshape_for_qk_norm(q, head_dim)
        q_by_head = q_norm(q_by_head)
        k_by_head = _reshape_for_qk_norm(k, head_dim)
        k_by_head = k_norm(k_by_head)
    q = q_by_head.view(q.shape)
    k = k_by_head.view(k.shape)
    return q, k


# ---------------------------------------------------------------------------
# Fused QK GemmaRMSNorm Triton kernel
# grid = q_rows (the larger dimension in GQA).  Every block computes Q norm
# for its row; the first k_rows blocks also compute K norm.  No torch.cat,
# no tl.where for weight selection, no output slice.
# ---------------------------------------------------------------------------
@triton.jit
def _fused_qk_gemma_rmsnorm_kernel(
    Q_ptr,
    K_ptr,
    Q_out_ptr,
    K_out_ptr,
    QW_ptr,
    KW_ptr,
    q_stride,
    k_stride,
    k_rows,
    HEAD_DIM: tl.constexpr,
    BLOCK_HD: tl.constexpr,
    EPS: tl.constexpr,
    FP16: tl.constexpr,
):
    pid = tl.program_id(0)
    cols = tl.arange(0, BLOCK_HD)
    mask = cols < HEAD_DIM
    out_dtype = tl.float16 if FP16 else tl.bfloat16

    # Q norm (every block) — use q_stride to handle non-contiguous input
    q_off = pid * q_stride + cols
    q = tl.load(Q_ptr + q_off, mask=mask, other=0.0).to(tl.float32)
    w_q = tl.load(QW_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    q_var = tl.sum(q * q, axis=0) / HEAD_DIM
    q_normed = (q * tl.rsqrt(q_var + EPS) * (w_q + 1.0)).to(out_dtype)
    # output is always contiguous
    q_out_off = pid * HEAD_DIM + cols
    tl.store(Q_out_ptr + q_out_off, q_normed, mask=mask)

    # K norm (first k_rows blocks only) — use k_stride for input
    if pid < k_rows:
        k_off = pid * k_stride + cols
        k = tl.load(K_ptr + k_off, mask=mask, other=0.0).to(tl.float32)
        w_k = tl.load(KW_ptr + cols, mask=mask, other=0.0).to(tl.float32)
        k_var = tl.sum(k * k, axis=0) / HEAD_DIM
        k_normed = (k * tl.rsqrt(k_var + EPS) * (w_k + 1.0)).to(out_dtype)
        k_out_off = pid * HEAD_DIM + cols
        tl.store(K_out_ptr + k_out_off, k_normed, mask=mask)


def fused_qk_gemma_rmsnorm(
    q: torch.Tensor,
    k: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    eps: float,
    head_dim: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Fused QK GemmaRMSNorm — single Triton kernel for both q_norm and k_norm.

    grid = q_rows; every block processes its Q row, and the first k_rows
    blocks also process K.  No torch.cat, no slice, no tl.where.
    Passes input strides to the kernel so non-contiguous tensors (e.g. from
    qkv.split()) are read correctly without an extra .contiguous() copy.
    """
    q_flat = q.reshape(-1, head_dim)
    k_flat = k.reshape(-1, head_dim)

    q_rows = q_flat.shape[0]
    k_rows = k_flat.shape[0]

    q_out = torch.empty(q_rows, head_dim, dtype=q.dtype, device=q.device)
    k_out = torch.empty(k_rows, head_dim, dtype=k.dtype, device=k.device)

    BLOCK_HD = triton.next_power_of_2(head_dim)

    _fused_qk_gemma_rmsnorm_kernel[(q_rows,)](
        q_flat,
        k_flat,
        q_out,
        k_out,
        q_weight,
        k_weight,
        q_flat.stride(0),
        k_flat.stride(0),
        k_rows,
        HEAD_DIM=head_dim,
        BLOCK_HD=BLOCK_HD,
        EPS=eps,
        FP16=(q.dtype == torch.float16),
    )

    return q_out, k_out


# ---------------------------------------------------------------------------
# Fused QK GemmaRMSNorm + gate extraction kernel
# For models with attn_output_gate (e.g. Qwen3.5) where q and gate are
# interleaved per head: [q_h0, gate_h0, q_h1, gate_h1, ...].
# Reads q from the interleaved buffer, normalizes it, and copies gate to a
# contiguous output — all in a single kernel launch.  Eliminates two
# elementwise copy kernels that would otherwise be needed to deinterleave.
# ---------------------------------------------------------------------------
@triton.jit
def _fused_qk_gemma_rmsnorm_gate_kernel(
    QG_ptr,
    K_ptr,
    Q_out_ptr,
    K_out_ptr,
    Gate_out_ptr,
    QW_ptr,
    KW_ptr,
    qg_token_stride,
    qg_head_stride,
    k_token_stride,
    k_head_stride,
    num_heads,
    num_kv_heads,
    k_rows,
    HEAD_DIM: tl.constexpr,
    BLOCK_HD: tl.constexpr,
    EPS: tl.constexpr,
    FP16: tl.constexpr,
):
    pid = tl.program_id(0)
    cols = tl.arange(0, BLOCK_HD)
    mask = cols < HEAD_DIM
    out_dtype = tl.float16 if FP16 else tl.bfloat16

    token_idx = pid // num_heads
    head_idx = pid % num_heads

    base = token_idx * qg_token_stride + head_idx * qg_head_stride

    # Q norm
    q = tl.load(QG_ptr + base + cols, mask=mask, other=0.0).to(tl.float32)
    w_q = tl.load(QW_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    q_var = tl.sum(q * q, axis=0) / HEAD_DIM
    q_normed = (q * tl.rsqrt(q_var + EPS) * (w_q + 1.0)).to(out_dtype)
    out_off = pid * HEAD_DIM + cols
    tl.store(Q_out_ptr + out_off, q_normed, mask=mask)

    # Gate copy
    gate = tl.load(QG_ptr + base + HEAD_DIM + cols, mask=mask, other=0.0)
    tl.store(Gate_out_ptr + out_off, gate, mask=mask)

    # K norm (first k_rows blocks only)
    if pid < k_rows:
        token_idx_k = pid // num_kv_heads
        head_idx_k = pid % num_kv_heads
        k_off = token_idx_k * k_token_stride + head_idx_k * k_head_stride + cols
        k = tl.load(K_ptr + k_off, mask=mask, other=0.0).to(tl.float32)
        w_k = tl.load(KW_ptr + cols, mask=mask, other=0.0).to(tl.float32)
        k_var = tl.sum(k * k, axis=0) / HEAD_DIM
        k_normed = (k * tl.rsqrt(k_var + EPS) * (w_k + 1.0)).to(out_dtype)
        k_out_off = pid * HEAD_DIM + cols
        tl.store(K_out_ptr + k_out_off, k_normed, mask=mask)


def fused_qk_gemma_rmsnorm_with_gate(
    q_gate: torch.Tensor,
    k: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    eps: float,
    head_dim: int,
    num_heads: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Fused QK GemmaRMSNorm + gate extraction from interleaved q_gate buffer.

    q_gate: (seq, q_size*2) where q and gate are interleaved per head,
            i.e. [q_h0, gate_h0, q_h1, gate_h1, ...] with q_size = num_heads * head_dim.
            Can be a non-contiguous view from qkv.split().
    k: (seq, kv_size) — same as fused_qk_gemma_rmsnorm.

    Returns (q_out, k_out, gate_out) all contiguous with shape
    (seq*num_heads, head_dim), (seq*num_kv_heads, head_dim), (seq*num_heads, head_dim).
    """
    seq_len = q_gate.shape[0]
    qg_3d = q_gate.view(seq_len, num_heads, 2 * head_dim)
    num_kv_heads = k.shape[-1] // head_dim
    k_3d = k.view(seq_len, num_kv_heads, head_dim)

    q_rows = seq_len * num_heads
    k_rows = seq_len * num_kv_heads

    q_out = torch.empty(q_rows, head_dim, dtype=q_gate.dtype, device=q_gate.device)
    k_out = torch.empty(k_rows, head_dim, dtype=k.dtype, device=k.device)
    gate_out = torch.empty(q_rows, head_dim, dtype=q_gate.dtype, device=q_gate.device)

    BLOCK_HD = triton.next_power_of_2(head_dim)

    _fused_qk_gemma_rmsnorm_gate_kernel[(q_rows,)](
        qg_3d,
        k_3d,
        q_out,
        k_out,
        gate_out,
        q_weight,
        k_weight,
        qg_3d.stride(0),
        qg_3d.stride(1),
        k_3d.stride(0),
        k_3d.stride(1),
        num_heads,
        num_kv_heads,
        k_rows,
        HEAD_DIM=head_dim,
        BLOCK_HD=BLOCK_HD,
        EPS=eps,
        FP16=(q_gate.dtype == torch.float16),
    )

    return q_out, k_out, gate_out


# Register the inplace op
fused_inplace_qknorm = register_custom_op(fused_inplace_qknorm, mutates_args=["q", "k"])