recipes/docker-images/glm53-kindling/experimental/fixes/_fa_common.py

"""
Copyright (c) 2026 by FlashInfer team.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
"""

import functools
import math
from typing import Callable, ClassVar, Optional, Protocol, Tuple, TypeVar, Union, cast

import torch

from ....jit import gen_batch_mla_module
from ....utils import MaskMode, check_shape_dtype_device, get_compute_capability
from ._capabilities import MLAPlanCapabilities, plan_capability_rejection_reason
from .._planning import _MLAPlanArguments


class _GeneratedBatchMLAModule(Protocol):
    def plan(self, *args: object) -> object: ...

    def plan_with_staged_workspace_bytes(self, *args: object) -> tuple[object, int]: ...

    def run(self, *args: object) -> object: ...


@functools.cache
def get_batch_mla_module(
    backend: str,
    dtype_q: torch.dtype,
    dtype_kv: torch.dtype,
    dtype_o: torch.dtype,
    dtype_idx: torch.dtype,
    head_dim_ckv: int,
    head_dim_kpe: int,
    use_profiler: bool,
) -> _GeneratedBatchMLAModule:
    return gen_batch_mla_module(
        backend,
        dtype_q,
        dtype_kv,
        dtype_o,
        dtype_idx,
        head_dim_ckv,
        head_dim_kpe,
        use_profiler,
    ).build_and_load()


def _validate_generated_fa_plan(
    *,
    backend: str,
    device: torch.device,
    qo_indptr: torch.Tensor,
    kv_indptr: torch.Tensor,
    kv_indices: torch.Tensor,
    kv_len_arr: torch.Tensor,
    head_dim_ckv: int,
    head_dim_kpe: int,
    q_data_type: torch.dtype,
    kv_data_type: torch.dtype,
    output_dtype: torch.dtype,
    scale_mode: str,
) -> None:
    for name, tensor in (
        ("qo_indptr", qo_indptr),
        ("kv_indptr", kv_indptr),
        ("kv_indices", kv_indices),
        ("kv_len_arr", kv_len_arr),
    ):
        if tensor.dtype != torch.int32:
            raise ValueError(f"{name} must have dtype torch.int32, got {tensor.dtype}.")
    if q_data_type not in (torch.float16, torch.bfloat16):
        raise ValueError(
            f"MLA q_data_type {q_data_type} is not supported by the {backend} backend."
        )
    supported_kv_dtypes = (torch.float16, torch.bfloat16, torch.float8_e4m3fn)
    if kv_data_type not in supported_kv_dtypes:
        raise ValueError(
            f"MLA kv_data_type {kv_data_type} is not supported by the {backend} "
            f"backend. Supported dtypes: {list(supported_kv_dtypes)}."
        )
    if output_dtype != q_data_type:
        raise ValueError(
            f"{backend} backend output_dtype must match q_data_type, got "
            f"{output_dtype} and {q_data_type}."
        )
    if head_dim_kpe < 0:
        raise ValueError(f"head_dim_kpe must be >= 0, got {head_dim_kpe}.")
    if kv_data_type == torch.float8_e4m3fn:
        major, minor = get_compute_capability(device)
        if major not in (9, 12):
            raise ValueError(
                "FP8 kv_data_type for MLA requires an SM90 (Hopper) device, "
                f"got SM{major}{minor}."
            )
        if q_data_type != torch.bfloat16:
            raise ValueError(
                "FP8 kv_data_type for MLA currently only supports "
                f"q_data_type=torch.bfloat16, got {q_data_type}."
            )
        if head_dim_ckv != 512 or head_dim_kpe not in (0, 64):
            raise ValueError(
                "FP8 kv_data_type for MLA currently only supports "
                "head_dim_ckv=512 and head_dim_kpe in (0, 64), got "
                f"head_dim_ckv={head_dim_ckv}, head_dim_kpe={head_dim_kpe}."
            )
        if scale_mode != "kv-per-tensor":
            raise ValueError("FP8 MLA plans require scale_mode='kv-per-tensor'.")
    elif scale_mode != "default":
        raise ValueError("non-FP8 MLA plans require scale_mode='default'.")


class _BatchMLAGeneratedFaMechanics:
    def __init__(
        self,
        *,
        backend: Optional[str] = None,
        float_workspace_buffer: torch.Tensor,
        use_cuda_graph: bool,
        qo_indptr_buf: Optional[torch.Tensor],
        kv_indptr_buf: Optional[torch.Tensor],
        kv_indices_buf: Optional[torch.Tensor],
        kv_len_arr_buf: Optional[torch.Tensor],
        int_workspace_buffer: Optional[torch.Tensor] = None,
        pin_memory_int_workspace_buffer: Optional[torch.Tensor] = None,
    ) -> None:
        self._backend = backend
        self._float_workspace_buffer = float_workspace_buffer
        self.device = float_workspace_buffer.device
        if int_workspace_buffer is None:
            int_workspace_buffer = torch.empty(
                (8 * 1024 * 1024,), dtype=torch.uint8, device=self.device
            )
        if pin_memory_int_workspace_buffer is None:
            pin_memory_int_workspace_buffer = torch.empty(
                int_workspace_buffer.shape,
                dtype=int_workspace_buffer.dtype,
                pin_memory=True,
                device="cpu",
            )
        self._int_workspace_buffer = int_workspace_buffer
        self._pin_memory_int_workspace_buffer = pin_memory_int_workspace_buffer
        self._use_cuda_graph = use_cuda_graph
        self._qo_indptr_buf = qo_indptr_buf
        self._kv_indptr_buf = kv_indptr_buf
        self._kv_indices_buf = kv_indices_buf
        self._kv_len_arr_buf = kv_len_arr_buf

    @staticmethod
    def _storage_interval(tensor: torch.Tensor) -> tuple[int, int, int]:
        start = tensor.storage_offset() * tensor.element_size()
        return (
            tensor.untyped_storage().data_ptr(),
            start,
            start + tensor.numel() * tensor.element_size(),
        )

    def _preflight_graph_metadata_buffers(
        self,
        *,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        kv_indices: torch.Tensor,
        kv_len_arr: torch.Tensor,
    ) -> None:
        named = (
            ("qo_indptr", self._qo_indptr_buf, qo_indptr, False),
            ("kv_indptr", self._kv_indptr_buf, kv_indptr, False),
            ("kv_indices", self._kv_indices_buf, kv_indices, True),
            ("kv_len_arr", self._kv_len_arr_buf, kv_len_arr, False),
        )
        for name, reserved, source, allow_larger in named:
            if reserved is None:
                raise ValueError(
                    "CUDA graph mode requires reserved qo_indptr, kv_indptr, "
                    "kv_indices, and kv_len_arr buffers."
                )
            if reserved.dtype != torch.int32:
                raise ValueError(
                    f"CUDA graph reserved {name} buffer must have dtype torch.int32."
                )
            if reserved.device != self.device:
                raise ValueError(
                    f"CUDA graph reserved {name} buffer must be on {self.device}."
                )
            if not reserved.is_contiguous():
                raise ValueError(
                    f"CUDA graph reserved {name} buffer must be contiguous."
                )
            if name == "kv_indices" and reserved.ndim != 1:
                raise ValueError(
                    "CUDA graph reserved kv_indices buffer must have rank 1."
                )
            if (allow_larger and reserved.shape[0] < source.shape[0]) or (
                not allow_larger and reserved.shape != source.shape
            ):
                raise ValueError(
                    f"CUDA graph reserved {name} buffer has insufficient or "
                    "incompatible capacity."
                )
        intervals = [
            (name, self._storage_interval(reserved))
            for name, reserved, _, _ in named
            if reserved is not None
        ]
        for index, (left_name, (left_ptr, left_start, left_end)) in enumerate(
            intervals
        ):
            for right_name, (right_ptr, right_start, right_end) in intervals[
                index + 1 :
            ]:
                if left_ptr == right_ptr and max(left_start, right_start) < min(
                    left_end, right_end
                ):
                    raise ValueError(
                        "CUDA graph reserved metadata buffers must not overlap: "
                        f"{left_name} overlaps {right_name}."
                    )
        copy_targets = {
            name: reserved[: source.shape[0]] if allow_larger else reserved
            for name, reserved, source, allow_larger in named
            if reserved is not None
        }
        for source_name, _, source, _ in named:
            source_interval = self._storage_interval(source)
            for target_name, target in copy_targets.items():
                target_interval = self._storage_interval(target)
                overlaps = source_interval[0] == target_interval[0] and max(
                    source_interval[1], target_interval[1]
                ) < min(source_interval[2], target_interval[2])
                exact_corresponding_reuse = (
                    source_name == target_name and source_interval == target_interval
                )
                if overlaps and not exact_corresponding_reuse:
                    raise ValueError(
                        f"CUDA graph {source_name} source overlaps reserved "
                        f"{target_name} target; only exact same-source/same-target "
                        "reuse is allowed."
                    )

    def _stage_metadata(
        self,
        *,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        kv_indices: torch.Tensor,
        kv_len_arr: torch.Tensor,
    ) -> None:
        if self._use_cuda_graph:
            self._preflight_graph_metadata_buffers(
                qo_indptr=qo_indptr,
                kv_indptr=kv_indptr,
                kv_indices=kv_indices,
                kv_len_arr=kv_len_arr,
            )
            assert self._qo_indptr_buf is not None
            assert self._kv_indptr_buf is not None
            assert self._kv_indices_buf is not None
            assert self._kv_len_arr_buf is not None
            targets = (
                (self._qo_indptr_buf, qo_indptr),
                (self._kv_indptr_buf, kv_indptr),
                (self._kv_indices_buf[: len(kv_indices)], kv_indices),
                (self._kv_len_arr_buf, kv_len_arr),
            )
            # GB10 fix: no rollback snapshot. It cloned kv_indices_buf on every
            # plan(), 136 MB per decode step for GLM's sparse MLA, to undo a
            # copy that the preflight above has already made safe.
            for target, source in targets:
                target.copy_(source, non_blocking=True)
        else:
            self._qo_indptr_buf = qo_indptr.to(self.device, non_blocking=True)
            self._kv_indptr_buf = kv_indptr.to(self.device, non_blocking=True)
            self._kv_indices_buf = kv_indices.to(self.device, non_blocking=True)
            self._kv_len_arr_buf = kv_len_arr.to(self.device, non_blocking=True)

    def _plan_generated_fa(
        self,
        *,
        module_loader: Callable[[], _GeneratedBatchMLAModule],
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        kv_indices: torch.Tensor,
        kv_len_arr: torch.Tensor,
        num_heads: int,
        head_dim_ckv: int,
        page_size: int,
        causal: bool,
        sm_scale: float,
        q_data_type: torch.dtype,
        kv_data_type: torch.dtype,
        use_profiler: bool,
    ) -> None:
        # ---------------------------------------------------------------------------
        # Build the generated backend plan
        # ---------------------------------------------------------------------------
        cached_module = module_loader()
        qo_indptr_host = qo_indptr.to("cpu")
        kv_indptr_host = kv_indptr.to("cpu")
        kv_len_arr_host = kv_len_arr.to("cpu")
        plan_args = (
            self._float_workspace_buffer,
            self._int_workspace_buffer,
            self._pin_memory_int_workspace_buffer,
            qo_indptr_host,
            kv_indptr_host,
            kv_len_arr_host,
            num_heads,
            head_dim_ckv,
            causal,
        )
        if hasattr(cached_module, "plan_with_staged_workspace_bytes"):
            plan_info, staged_int_workspace_bytes = (
                cached_module.plan_with_staged_workspace_bytes(*plan_args)
            )
        else:
            # Compatibility for externally cached legacy modules; generated modules
            # are always adapted at load time.
            plan_info = cached_module.plan(*plan_args)
            staged_int_workspace_bytes = 0

        # ---------------------------------------------------------------------------
        # Stage metadata and publish backend state
        # ---------------------------------------------------------------------------
        self._stage_metadata(
            qo_indptr=qo_indptr,
            kv_indptr=kv_indptr,
            kv_indices=kv_indices,
            kv_len_arr=kv_len_arr,
        )
        self._cached_module = cached_module
        self._causal = causal
        self._page_size = page_size
        self._sm_scale = sm_scale
        self._head_dim_ckv = head_dim_ckv
        self._q_data_type = q_data_type
        self._kv_data_type = kv_data_type
        self._use_profiler = use_profiler
        self._plan_info = plan_info
        self._staged_int_workspace_bytes = staged_int_workspace_bytes

    def _validate_run_input_dtypes(
        self,
        *,
        q_nope: torch.Tensor,
        q_pe: torch.Tensor,
        ckv_cache: torch.Tensor,
        kpe_cache: torch.Tensor,
    ) -> None:
        if q_nope.dtype != self._q_data_type:
            raise ValueError(
                f"q_nope.dtype={q_nope.dtype} does not match the planned "
                f"q_data_type={self._q_data_type}."
            )
        if q_pe.dtype != self._q_data_type:
            raise ValueError(
                f"q_pe.dtype={q_pe.dtype} does not match the planned "
                f"q_data_type={self._q_data_type}."
            )
        if ckv_cache.dtype != self._kv_data_type:
            raise ValueError(
                f"ckv_cache.dtype={ckv_cache.dtype} does not match the planned "
                f"kv_data_type={self._kv_data_type}."
            )
        if kpe_cache.dtype != self._kv_data_type:
            raise ValueError(
                f"kpe_cache.dtype={kpe_cache.dtype} does not match the planned "
                f"kv_data_type={self._kv_data_type}."
            )

    def _run_generated_fa(
        self,
        *,
        q_nope: torch.Tensor,
        q_pe: torch.Tensor,
        ckv_cache: torch.Tensor,
        kpe_cache: torch.Tensor,
        out: Optional[torch.Tensor],
        lse: Optional[torch.Tensor],
        return_lse: bool,
        profiler_buffer: Optional[torch.Tensor],
        return_lse_base_on_e: bool,
        ckv_scale: Optional[float],
        ckv_scale_arr: Optional[torch.Tensor],
        kpe_scale: Optional[float],
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        # ---------------------------------------------------------------------------
        # Validate inputs and scale arguments
        # ---------------------------------------------------------------------------
        self._validate_run_input_dtypes(
            q_nope=q_nope,
            q_pe=q_pe,
            ckv_cache=ckv_cache,
            kpe_cache=kpe_cache,
        )
        kv_is_fp8 = self._kv_data_type == torch.float8_e4m3fn
        if kv_is_fp8:
            if (ckv_scale is None) == (ckv_scale_arr is None):
                raise ValueError(
                    "Exactly one of ckv_scale or ckv_scale_arr is required when "
                    "kv_data_type is FP8."
                )
            if kpe_scale is None:
                raise ValueError("kpe_scale is required when kv_data_type is FP8.")
            ckv_scale_f = 1.0 if ckv_scale is None else float(ckv_scale)
            kpe_scale_f = float(kpe_scale)
            if ckv_scale is not None and (
                not math.isfinite(ckv_scale_f) or ckv_scale_f <= 0.0
            ):
                raise ValueError(
                    f"ckv_scale must be a finite positive value, got {ckv_scale}"
                )
            if not math.isfinite(kpe_scale_f) or kpe_scale_f <= 0.0:
                raise ValueError(
                    f"kpe_scale must be a finite positive value, got {kpe_scale}"
                )
        else:
            if (
                ckv_scale is not None
                or ckv_scale_arr is not None
                or kpe_scale is not None
            ):
                raise ValueError(
                    "ckv_scale / ckv_scale_arr / kpe_scale are only valid when "
                    "kv_data_type is FP8."
                )
            ckv_scale_f = 1.0
            kpe_scale_f = 1.0

        # ---------------------------------------------------------------------------
        # Prepare output and auxiliary buffers
        # ---------------------------------------------------------------------------
        if profiler_buffer is None and self._use_profiler:
            raise ValueError("Profiler is enabled, profiler_buffer must be provided")
        if out is None:
            out = torch.empty_like(q_nope)
        else:
            check_shape_dtype_device(
                out, q_nope.shape, q_nope.dtype, q_nope.device, "out"
            )
        if return_lse:
            if lse is None:
                lse = torch.empty(
                    q_nope.shape[:2], dtype=torch.float32, device=self.device
                )
            else:
                check_shape_dtype_device(
                    lse, q_nope.shape[:2], torch.float32, q_nope.device, "lse"
                )
        if ckv_scale_arr is not None:
            expected_scale_shape = (*ckv_cache.shape[:-1], self._head_dim_ckv // 128)
            check_shape_dtype_device(
                ckv_scale_arr,
                expected_scale_shape,
                torch.float32,
                ckv_cache.device,
                "ckv_scale_arr",
            )
            if not ckv_scale_arr.is_contiguous():
                raise ValueError("ckv_scale_arr must be contiguous.")

        # ---------------------------------------------------------------------------
        # Launch the generated backend
        # ---------------------------------------------------------------------------
        mask_mode = MaskMode.CAUSAL.value if self._causal else MaskMode.NON_CAUSAL.value
        profiler_args = (profiler_buffer,) if self._use_profiler else ()
        self._cached_module.run(
            self._float_workspace_buffer,
            self._int_workspace_buffer,
            self._plan_info,
            q_nope,
            q_pe,
            ckv_cache,
            kpe_cache,
            self._kv_indices_buf,
            out,
            lse,
            mask_mode,
            q_nope.shape[1],
            self._page_size,
            self._sm_scale,
            return_lse_base_on_e,
            ckv_scale_f,
            kpe_scale_f,
            ckv_scale_arr,
            *profiler_args,
        )
        return (out, lse) if return_lse else out


_FaBackendT = TypeVar("_FaBackendT", bound="_BatchMLAPagedAttentionFaBackendBase")


class _BatchMLAPagedAttentionFaBackendBase(_BatchMLAGeneratedFaMechanics):
    _plan_capabilities: ClassVar[Optional[MLAPlanCapabilities]] = None

    def __init__(
        self,
        *,
        backend: Optional[str] = None,
        float_workspace_buffer: torch.Tensor,
        use_cuda_graph: bool,
        qo_indptr_buf: Optional[torch.Tensor],
        kv_indptr_buf: Optional[torch.Tensor],
        kv_indices_buf: Optional[torch.Tensor],
        kv_len_arr_buf: Optional[torch.Tensor],
        query_split_widths: tuple[int, int],
        kv_split_widths: tuple[int, int],
        int_workspace_buffer: Optional[torch.Tensor] = None,
        pin_memory_int_workspace_buffer: Optional[torch.Tensor] = None,
    ) -> None:
        if backend is None:
            if self._plan_capabilities is None:
                raise TypeError("generated-FA backend capabilities are required.")
            backend = self._plan_capabilities.backend_name
        super().__init__(
            backend=backend,
            float_workspace_buffer=float_workspace_buffer,
            use_cuda_graph=use_cuda_graph,
            qo_indptr_buf=qo_indptr_buf,
            kv_indptr_buf=kv_indptr_buf,
            kv_indices_buf=kv_indices_buf,
            kv_len_arr_buf=kv_len_arr_buf,
            int_workspace_buffer=int_workspace_buffer,
            pin_memory_int_workspace_buffer=pin_memory_int_workspace_buffer,
        )
        self._query_split_widths = query_split_widths
        self._kv_split_widths = kv_split_widths

    def plan(
        self,
        *,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        kv_indices: torch.Tensor,
        kv_len_arr: torch.Tensor,
        num_heads: int,
        head_dim_ckv: int,
        head_dim_kpe: int,
        page_size: int,
        causal: bool,
        sm_scale: float,
        q_data_type: torch.dtype,
        kv_data_type: torch.dtype,
        output_dtype: torch.dtype,
        scale_mode: str,
        use_profiler: bool,
    ) -> None:
        raise NotImplementedError

    @classmethod
    def plan_from_wrapper(
        cls: type[_FaBackendT], args: _MLAPlanArguments
    ) -> _FaBackendT:
        assert cls._plan_capabilities is not None
        if reason := plan_capability_rejection_reason(args, cls._plan_capabilities):
            raise ValueError(reason)
        csr = args.csr()
        backend = cls(
            float_workspace_buffer=args._float_workspace_buffer,
            use_cuda_graph=args._use_cuda_graph,
            qo_indptr_buf=args._qo_indptr_buf,
            kv_indptr_buf=args._kv_indptr_buf,
            kv_indices_buf=args._kv_indices_buf,
            kv_len_arr_buf=args._kv_len_arr_buf,
            query_split_widths=(args.head_dim_ckv, args.head_dim_kpe),
            kv_split_widths=(args.head_dim_ckv, args.head_dim_kpe),
            int_workspace_buffer=args._graph_plan_int_workspace_buffer,
        )
        if args._use_cuda_graph:
            backend._preflight_graph_metadata_buffers(
                qo_indptr=csr.qo_indptr,
                kv_indptr=csr.kv_indptr,
                kv_indices=csr.kv_indices,
                kv_len_arr=csr.kv_len_arr,
            )
        backend.plan(
            qo_indptr=csr.qo_indptr,
            kv_indptr=csr.kv_indptr,
            kv_indices=csr.kv_indices,
            kv_len_arr=csr.kv_len_arr,
            num_heads=args.num_heads,
            head_dim_ckv=args.head_dim_ckv,
            head_dim_kpe=args.head_dim_kpe,
            page_size=args.page_size,
            causal=args.causal,
            sm_scale=args.sm_scale,
            q_data_type=args.q_data_type,
            kv_data_type=args.kv_data_type,
            output_dtype=args.output_dtype,
            scale_mode=args.scale_mode,
            use_profiler=args.use_profiler,
        )
        return backend

    def run_from_wrapper(
        self,
        *,
        query: object,
        kv_cache: object,
        out: Optional[torch.Tensor],
        lse: Optional[torch.Tensor],
        return_lse: bool,
        profiler_buffer: Optional[torch.Tensor],
        kv_len: Optional[torch.Tensor],
        page_table: Optional[torch.Tensor],
        return_lse_base_on_e: bool,
        o_scale: Optional[float],
        ckv_scale: Optional[float],
        ckv_scale_arr: Optional[torch.Tensor],
        kpe_scale: Optional[float],
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        if kv_len is not None:
            raise ValueError("kv_len is only supported with cutlass backend.")
        if page_table is not None:
            raise ValueError("page_table is only supported with cutlass backend.")
        if o_scale is not None:
            raise ValueError(
                "o_scale is only supported with the cutlass backend for now."
            )
        q_nope, q_pe = cast(tuple[torch.Tensor, torch.Tensor], query)
        ckv_cache, kpe_cache = cast(tuple[torch.Tensor, torch.Tensor], kv_cache)
        return self._run_generated_fa(
            q_nope=q_nope,
            q_pe=q_pe,
            ckv_cache=ckv_cache,
            kpe_cache=kpe_cache,
            out=out,
            lse=lse,
            return_lse=return_lse,
            profiler_buffer=profiler_buffer,
            return_lse_base_on_e=return_lse_base_on_e,
            ckv_scale=ckv_scale,
            ckv_scale_arr=ckv_scale_arr,
            kpe_scale=kpe_scale,
        )