# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Private descriptor-first BF16 API for SM100 grouped GEMM dGLU."""

from __future__ import annotations

import os
import weakref
from typing import Optional, Tuple

import cutlass
import cutlass.cute as cute
from cuda.bindings import driver as cuda
from cutlass.cute.nvgpu import OperandMajorMode
from cutlass.cute.runtime import from_dlpack, make_fake_stream

from cudnn.api_base import APIBase, TensorDesc
from cudnn.datatypes import _convert_to_cutlass_data_type
from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _pointer_values, _validate_pointer_tensor
from cudnn.tensor_adapter import (
    allocate_byte_workspace,
    canonicalize_unit_dim_strides,
    cuda_is_available,
    default_stream,
    detect_framework,
    get_compute_capability,
    get_data_ptr,
    get_device,
    get_shape,
    get_strides,
    get_version,
    is_torch_tensor,
    to_host_list,
)

from ..moe_utils import MoEWeightMode
from .moe_grouped_gemm_dglu_dbias import MoEGroupedGemmDgluDbiasBf16Kernel


def _output_dtypes():
    return [cutlass.BFloat16, cutlass.Float16, cutlass.Float32]


class GroupedGemmDgluBf16API(APIBase):
    """Descriptor-first lifecycle API for the source BF16 dGLU kernel."""

    def __init__(
        self,
        sample_a: torch.Tensor,
        sample_c: torch.Tensor,
        sample_d_row: torch.Tensor,
        sample_padded_offsets: torch.Tensor,
        sample_alpha: torch.Tensor,
        sample_beta: torch.Tensor,
        sample_prob: torch.Tensor,
        sample_dprob: torch.Tensor,
        sample_b: Optional[torch.Tensor] = None,
        sample_dbias: Optional[torch.Tensor] = None,
        num_experts: Optional[int] = None,
        b_shape: Optional[Tuple[int, ...]] = None,
        b_dtype: Optional[torch.dtype] = None,
        acc_dtype: Optional[torch.dtype] = None,
        mma_tiler_mn: Tuple[int, int] = (256, 256),
        cluster_shape_mn: Optional[Tuple[int, int]] = None,
        vector_f32: bool = False,
        m_aligned: int = 256,
        act_func: str = "dswiglu",
        b_major: str = "k",
        use_dynamic_sched: bool = False,
    ) -> None:
        if acc_dtype is None:
            acc_dtype = cutlass.Float32
        super().__init__()
        self._warn_experimental_api()
        self._framework = detect_framework(sample_a)

        if sample_b is not None and num_experts is None:
            self.weight_mode = MoEWeightMode.DENSE
        elif sample_b is None and num_experts is not None:
            self.weight_mode = MoEWeightMode.DISCRETE
            if b_shape is None or b_dtype is None:
                raise ValueError("b_shape and b_dtype are required in discrete mode")
        else:
            raise ValueError("Provide sample_b for dense mode or (num_experts, b_shape, b_dtype) " "for discrete mode, but not both")

        # A canonical TensorDesc is rebuilt and re-checked for every operand on every
        # launch, which is the largest single cost in execute(). See _validate_live_tensor.
        self._live_desc_cache: dict = {}
        self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True)
        self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True)
        self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True)
        self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row", canonical=True)
        self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True)
        self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True)
        self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta", canonical=True)
        self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True)
        self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob", canonical=True)
        self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias", canonical=True)

        self._sample_offset_values = self._copy_values_to_host(sample_padded_offsets)
        self._sample_offsets_ref = weakref.ref(sample_padded_offsets)
        self._sample_offsets_version = get_version(sample_padded_offsets)
        self._sample_data_ptrs = {
            name: get_data_ptr(tensor)
            for name, tensor in (
                ("sample_a", sample_a),
                ("sample_b", sample_b),
                ("sample_c", sample_c),
                ("sample_d_row", sample_d_row),
                ("sample_padded_offsets", sample_padded_offsets),
                ("sample_alpha", sample_alpha),
                ("sample_beta", sample_beta),
                ("sample_prob", sample_prob),
                ("sample_dprob", sample_dprob),
                ("sample_dbias", sample_dbias),
            )
            if tensor is not None
        }

        self.expert_cnt = self.b_desc.shape[2] if self.weight_mode == MoEWeightMode.DENSE and self.b_desc.ndim == 3 else int(num_experts or 0)
        self.b_shape = tuple(b_shape) if b_shape is not None else None
        self.b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else self.b_desc.dtype
        self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype)
        self.mma_tiler_mn = tuple(mma_tiler_mn)
        self.use_2cta_instrs = self.mma_tiler_mn[0] == 256
        self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1)))
        self.vector_f32 = vector_f32
        self.m_aligned = m_aligned
        self.act_func = act_func
        self.b_major = b_major
        self.use_dynamic_sched = use_dynamic_sched
        self._has_dbias = self.dbias_desc is not None
        self._kernel = MoEGroupedGemmDgluDbiasBf16Kernel
        self._workspace: Optional[torch.Tensor] = None
        self._live_b_ptrs = None
        self._compile_b_ptrs = None
        self._validated_offsets: dict[int, tuple] = {}
        self._validated_pointer_values: dict[int, tuple] = {}
        self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0"))

    @staticmethod
    def _expect_shape(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None:
        if desc.shape != expected:
            raise ValueError(f"{name} shape mismatch: expected {expected}, got {desc.shape}")

    @staticmethod
    def _expect_stride(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None:
        if desc.stride != expected:
            raise ValueError(f"{name} must use the source-compatible layout with stride " f"{expected}, got {desc.stride}")

    @staticmethod
    def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None:
        if desc.device != device:
            raise ValueError(f"{name} must be on {device}, got {desc.device}")

    @staticmethod
    def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]:
        return tuple(int(value) for value in to_host_list(tensor))

    @staticmethod
    def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool:
        cached = cache.get(id(tensor))
        return bool(cached and cached[0]() is tensor and cached[1] == get_version(tensor) and cached[2] == extra)

    @staticmethod
    def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None:
        key = id(tensor)

        def discard(_reference, *, cache=cache, key=key):
            cache.pop(key, None)

        cache[key] = (weakref.ref(tensor, discard), get_version(tensor), extra)

    @staticmethod
    def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tensor_m: int) -> None:
        if len(values) != expert_cnt:
            raise ValueError(f"padded_offsets length mismatch: expected {expert_cnt}, got {len(values)}")
        previous = 0
        for index, value in enumerate(values):
            if value < previous:
                raise ValueError("padded_offsets must be a non-decreasing cumulative sum; " f"index {index} is {value} after {previous}")
            if value % MoEGroupedGemmDgluDbiasBf16Kernel.FIX_PAD_SIZE != 0:
                raise ValueError(f"padded_offsets[{index}] must be 256-aligned, got {value}")
            previous = value
        if not values or values[-1] <= 0 or values[-1] > tensor_m:
            raise ValueError(f"padded_offsets last value must be in [1, {tensor_m}], got " f"{values[-1] if values else None}")

    def _validate_offsets_once(self, offsets: torch.Tensor, *, tensor_m: int) -> None:
        extra = (self.expert_cnt, tensor_m)
        if self._is_validation_cached(self._validated_offsets, offsets, extra):
            return
        values = self._copy_values_to_host(offsets)
        self._validate_offset_sequence(values, expert_cnt=self.expert_cnt, tensor_m=tensor_m)
        self._remember_validation(self._validated_offsets, offsets, extra)

    def _validate_pointer_values_once(self, b_ptrs: torch.Tensor) -> None:
        if self._is_validation_cached(self._validated_pointer_values, b_ptrs, self.expert_cnt):
            return
        values = _pointer_values(b_ptrs)
        if any(value == 0 or value % 16 != 0 for value in values):
            raise ValueError("b_ptrs entries must be non-null and 16-byte aligned")
        self._remember_validation(self._validated_pointer_values, b_ptrs, self.expert_cnt)

    @staticmethod
    def _validate_data_alignment(tensor: torch.Tensor, name: str) -> None:
        if get_data_ptr(tensor) % 16 != 0:
            raise ValueError(f"{name} data pointer must be 16-byte aligned")

    @staticmethod
    def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None:
        if get_data_ptr(tensor) % 8 != 0:
            raise ValueError("b_ptrs data pointer must be 8-byte aligned")

    def _record_pointer_stream(self, b_ptrs: torch.Tensor, stream: cuda.CUstream) -> None:
        if not is_torch_tensor(b_ptrs):
            # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the
            # array referenced until the next execute so its buffer outlives the launch.
            self._live_b_ptrs = b_ptrs
            return
        import torch

        handle = int(stream)
        torch_current = torch.cuda.current_stream(b_ptrs.device)
        torch_default = torch.cuda.default_stream(b_ptrs.device)
        if handle == torch_current.cuda_stream:
            launch_stream = torch_current
        elif handle == torch_default.cuda_stream:
            launch_stream = torch_default
        else:
            launch_stream = torch.cuda.ExternalStream(handle, device=b_ptrs.device)
        b_ptrs.record_stream(launch_stream)

    def check_support(self) -> bool:
        if self.a_desc.ndim != 3:
            raise ValueError(f"sample_a must be rank-3, got {self.a_desc.shape}")
        tensor_m, k, one = self.a_desc.shape
        if one != 1:
            raise ValueError(f"sample_a trailing dimension must be 1, got {one}")

        if self.weight_mode == MoEWeightMode.DENSE:
            if self.b_desc.ndim != 3:
                raise ValueError(f"sample_b must be rank-3, got {self.b_desc.shape}")
            n, b_k, experts = self.b_desc.shape
            if b_k != k:
                raise ValueError(f"sample_b K dimension ({b_k}) must match sample_a ({k})")
            if experts != self.expert_cnt:
                raise ValueError("sample_b expert dimension is inconsistent")
            self._expect_stride(self.b_desc, (k, 1, n * k), "sample_b")
        else:
            if len(self.b_shape) not in (2, 3):
                raise ValueError(f"b_shape must be rank-2 or rank-3, got {self.b_shape}")
            n, b_k = self.b_shape[:2]
            if len(self.b_shape) == 3 and self.b_shape[2] != 1:
                raise ValueError(f"b_shape trailing dimension must be 1, got {self.b_shape}")
            if b_k != k:
                raise ValueError(f"b_shape K dimension ({b_k}) must match sample_a ({k})")
        if n <= 0 or n % 32 != 0:
            raise ValueError(f"N must be positive and divisible by 32, got {n}")

        two_n = 2 * n
        self._expect_shape(self.c_desc, (tensor_m, two_n, 1), "sample_c")
        self._expect_shape(self.d_row_desc, (tensor_m, two_n, 1), "sample_d_row")
        self._expect_shape(self.padded_offsets_desc, (self.expert_cnt,), "sample_padded_offsets")
        self._expect_shape(self.alpha_desc, (self.expert_cnt,), "sample_alpha")
        self._expect_shape(self.beta_desc, (self.expert_cnt,), "sample_beta")
        self._expect_shape(self.prob_desc, (tensor_m, 1, 1), "sample_prob")
        self._expect_shape(self.dprob_desc, (tensor_m, 1, 1), "sample_dprob")

        self._expect_stride(self.a_desc, (k, 1, tensor_m * k), "A tensor")
        self._expect_stride(self.c_desc, (two_n, 1, tensor_m * two_n), "sample_c")
        self._expect_stride(self.d_row_desc, (two_n, 1, tensor_m * two_n), "sample_d_row")
        self._expect_stride(self.padded_offsets_desc, (1,), "sample_padded_offsets")
        self._expect_stride(self.alpha_desc, (1,), "sample_alpha")
        self._expect_stride(self.beta_desc, (1,), "sample_beta")
        self._expect_stride(self.prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "sample_prob")
        self._expect_stride(self.dprob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "sample_dprob")

        self._check_dtype(self.a_desc, cutlass.BFloat16, "sample_a")
        if self.b_desc is not None:
            self._check_dtype(self.b_desc, cutlass.BFloat16, "sample_b")
        self._check_dtype(self.b_dtype, cutlass.BFloat16, "b_dtype")
        self._check_dtype(self.c_desc, _output_dtypes(), "sample_c")
        self._check_dtype(self.d_row_desc, _output_dtypes(), "sample_d_row")
        self._check_dtype(self.padded_offsets_desc, cutlass.Int32, "sample_padded_offsets")
        self._check_dtype(self.alpha_desc, cutlass.Float32, "sample_alpha")
        self._check_dtype(self.beta_desc, cutlass.Float32, "sample_beta")
        self._check_dtype(self.prob_desc, cutlass.Float32, "sample_prob")
        self._check_dtype(self.dprob_desc, cutlass.Float32, "sample_dprob")

        device = self.a_desc.device
        for desc, name in (
            (self.c_desc, "sample_c"),
            (self.d_row_desc, "sample_d_row"),
            (self.padded_offsets_desc, "sample_padded_offsets"),
            (self.alpha_desc, "sample_alpha"),
            (self.beta_desc, "sample_beta"),
            (self.prob_desc, "sample_prob"),
            (self.dprob_desc, "sample_dprob"),
        ):
            self._expect_device(desc, device, name)
        if self.b_desc is not None:
            self._expect_device(self.b_desc, device, "sample_b")

        if self.dbias_desc is not None:
            self._expect_shape(self.dbias_desc, (self.expert_cnt, two_n, 1), "sample_dbias")
            self._expect_stride(self.dbias_desc, canonicalize_unit_dim_strides((self.expert_cnt, two_n, 1), (two_n, 1, 1)), "sample_dbias")
            self._check_dtype(self.dbias_desc, cutlass.BFloat16, "sample_dbias")
            self._expect_device(self.dbias_desc, device, "sample_dbias")

        for name, pointer in self._sample_data_ptrs.items():
            if pointer % 16 != 0:
                raise ValueError(f"{name} data pointer must be 16-byte aligned")

        if self.acc_dtype is not cutlass.Float32:
            raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}")
        if self.m_aligned != 256:
            raise ValueError(f"m_aligned must be 256, got {self.m_aligned}")
        if tensor_m % 256 != 0:
            raise ValueError(f"sample_a M dimension must be 256-aligned, got {tensor_m}")
        if self.act_func not in ("dswiglu", "dgeglu"):
            raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {self.act_func}")
        if self.b_major not in ("k", "n"):
            raise ValueError(f"b_major must be 'k' or 'n', got {self.b_major}")
        if self.expert_cnt <= 0 or self.expert_cnt > 1024:
            raise ValueError(f"expert count must be in [1, 1024], got {self.expert_cnt}")

        self._validate_offset_sequence(self._sample_offset_values, expert_cnt=self.expert_cnt, tensor_m=tensor_m)
        sample_offsets = self._sample_offsets_ref()
        if sample_offsets is not None and get_version(sample_offsets) == self._sample_offsets_version:
            self._remember_validation(self._validated_offsets, sample_offsets, (self.expert_cnt, tensor_m))
        elif sample_offsets is not None:
            self._validate_offsets_once(sample_offsets, tensor_m=tensor_m)

        if not self._kernel.can_implement(
            cutlass.BFloat16,
            _convert_to_cutlass_data_type(self.c_desc.dtype),
            _convert_to_cutlass_data_type(self.d_row_desc.dtype),
            _convert_to_cutlass_data_type(self.acc_dtype),
            self.use_2cta_instrs,
            self.mma_tiler_mn,
            self.cluster_shape_mn,
            tensor_m,
            n,
            k,
            self.expert_cnt,
            "k",
            self.b_major,
            "n",
            self.m_aligned,
            self.act_func,
        ):
            raise ValueError("Unsupported BF16 grouped GEMM dGLU configuration")

        if not cuda_is_available():
            raise RuntimeError("CUDA is not available")
        major, minor = get_compute_capability()
        capability = major * 10 + minor
        if capability < 100:
            raise RuntimeError(f"GroupedGemmDgluSm100 requires SM100+, found SM{capability}")
        self._is_supported = True
        return True

    def compile(self) -> None:
        self._ensure_support_checked()
        if self._compiled_kernel is not None:
            return

        kernel = self._kernel(
            acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype),
            use_2cta_instrs=self.use_2cta_instrs,
            mma_tiler_mn=self.mma_tiler_mn,
            cluster_shape_mn=self.cluster_shape_mn,
            vectorized_f32=self.vector_f32,
            expert_cnt=self.expert_cnt,
            weight_mode=self.weight_mode,
            use_dynamic_sched=self.use_dynamic_sched,
            act_func=self.act_func,
        )
        hardware_info = cutlass.utils.HardwareInfo()
        max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self.num_cluster_overlap_margin
        if max_active_clusters <= 0:
            raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN")

        workspace_bytes = kernel.get_workspace_bytes()
        # Internal scratch in the caller's framework allocator; kernels write through its
        # raw pointer and it is never surfaced as a framework array.
        self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device)
        if get_data_ptr(self._workspace) % 128 != 0:
            raise RuntimeError("workspace allocation must be 128-byte aligned")
        workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator
        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False)

        a_fake = self._make_fake_cute_compact_tensor(
            self.a_desc.dtype,
            self.a_desc.shape,
            self.a_desc.stride_order,
            dynamic_mode=0,
            divisibility=256,
        )
        c_fake = self._make_fake_cute_compact_tensor(
            self.c_desc.dtype,
            self.c_desc.shape,
            self.c_desc.stride_order,
            dynamic_mode=0,
            divisibility=256,
        )
        d_fake = self._make_fake_cute_compact_tensor(
            self.d_row_desc.dtype,
            self.d_row_desc.shape,
            self.d_row_desc.stride_order,
            dynamic_mode=0,
            divisibility=256,
        )
        valid_m = cute.sym_int(divisibility=256)
        prob_fake = self._make_fake_cute_tensor(self.prob_desc.dtype, (valid_m, 1, 1), self.prob_desc.stride)
        dprob_fake = self._make_fake_cute_tensor(self.dprob_desc.dtype, (valid_m, 1, 1), self.dprob_desc.stride)

        if self.weight_mode == MoEWeightMode.DENSE:
            b_fake = self._make_fake_cute_tensor_from_desc(self.b_desc)
            n_value = cutlass.Int32(0)
            k_value = cutlass.Int32(0)
            b_stride = cutlass.Int64(0)
            b_major_mode = OperandMajorMode.K
        else:
            # Compile-time placeholder for the pointer-array argument: real device bytes
            # (fake tensors have dummy iterators) allocated in the caller's framework,
            # retyped to Int64 via the element_type override.
            self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device)
            self._validate_pointer_array_alignment(self._compile_b_ptrs)
            placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8)
            placeholder.element_type = cutlass.Int64
            b_fake = placeholder.iterator
            n, k = self.b_shape[:2]
            n_value = cutlass.Int32(n)
            k_value = cutlass.Int32(k)
            b_stride = cutlass.Int64(k if self.b_major == "k" else n)
            b_major_mode = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN

        raw_compiled = cute.compile(
            kernel,
            a=a_fake,
            b=b_fake,
            n=n_value,
            k=k_value,
            b_stride_size=b_stride,
            b_major_mode=b_major_mode,
            workspace_ptr=workspace_ptr,
            c=c_fake,
            d=d_fake,
            padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc),
            alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc),
            beta=self._make_fake_cute_tensor_from_desc(self.beta_desc),
            prob=prob_fake,
            dprob=dprob_fake,
            linear_offset=cutlass.Float32(0.0),
            dbias_tensor=self._make_fake_cute_tensor_from_desc(self.dbias_desc),
            max_active_clusters=max_active_clusters,
            stream=fake_stream,
            options="--enable-tvm-ffi",
        )
        cached_n, cached_k, cached_b_stride = n_value, k_value, b_stride

        def tensor_api(
            a_tensor,
            c_tensor,
            d_row_tensor,
            padded_offsets,
            alpha_tensor,
            beta_tensor,
            prob_tensor,
            dprob_tensor,
            b_tensor,
            b_ptrs,
            dbias_tensor,
            stream,
            linear_offset,
        ) -> None:
            b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(get_data_ptr(b_ptrs))
            raw_compiled(
                a_tensor,
                b_arg,
                cached_n,
                cached_k,
                cached_b_stride,
                workspace_ptr,
                c_tensor,
                d_row_tensor,
                padded_offsets,
                alpha_tensor,
                beta_tensor,
                prob_tensor,
                dprob_tensor,
                cutlass.Float32(linear_offset),
                dbias_tensor,
                stream,
            )

        self._compiled_kernel = tensor_api

    def _validate_live_tensor(self, tensor: torch.Tensor, sample: TensorDesc, name: str, *, dynamic_m: bool = False) -> TensorDesc:
        # The outcome of this check is decided entirely by (shape, stride, dtype, device),
        # so memoize on exactly that. An operand differing in any of them takes a different
        # key, misses, and is checked in full; nothing is skipped on the strength of a
        # tensor's identity. Data-pointer alignment is checked by the caller, not here.
        key = (name, dynamic_m, get_shape(tensor), get_strides(tensor), tensor.dtype, get_device(tensor))
        hit = self._live_desc_cache.get(key)
        if hit is not None:
            return hit
        desc = self._make_tensor_desc(tensor, name=name, canonical=True)
        if desc.dtype != sample.dtype:
            raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}")
        if desc.device != sample.device:
            raise ValueError(f"{name} device mismatch: expected {sample.device}, got {desc.device}")
        if dynamic_m:
            if desc.shape[1:] != sample.shape[1:]:
                raise ValueError(f"{name} shape suffix mismatch: expected {sample.shape[1:]}, got {desc.shape[1:]}")
            if desc.stride_order != sample.stride_order:
                raise ValueError(f"{name} layout mismatch: expected stride order {sample.stride_order}, got {desc.stride_order}")
        elif desc.shape != sample.shape or desc.stride != sample.stride:
            raise ValueError(f"{name} descriptor mismatch: expected shape/stride {sample.shape}/{sample.stride}, " f"got {desc.shape}/{desc.stride}")
        self._live_desc_cache[key] = desc
        return desc

    def execute(
        self,
        a_tensor: torch.Tensor,
        c_tensor: torch.Tensor,
        d_row_tensor: torch.Tensor,
        padded_offsets: torch.Tensor,
        alpha_tensor: torch.Tensor,
        beta_tensor: torch.Tensor,
        prob_tensor: torch.Tensor,
        dprob_tensor: torch.Tensor,
        b_tensor: Optional[torch.Tensor] = None,
        b_ptrs: Optional[torch.Tensor] = None,
        dbias_tensor: Optional[torch.Tensor] = None,
        linear_offset: float = 0.0,
        current_stream: Optional[cuda.CUstream] = None,
    ) -> None:
        if current_stream is None:
            # torch inputs stay ordered with the caller's current torch stream;
            # other frameworks (e.g. JAX) default to the CUDA legacy default stream.
            current_stream = default_stream(detect_framework(a_tensor))
        if self._compiled_kernel is None:
            raise RuntimeError("Kernel not compiled; call compile() first")

        a_desc = self._validate_live_tensor(a_tensor, self.a_desc, "a_tensor", dynamic_m=True)
        tensor_m, k, _ = a_desc.shape
        if tensor_m % 256 != 0:
            raise ValueError(f"a_tensor M dimension must be 256-aligned, got {tensor_m}")
        c_desc = self._validate_live_tensor(c_tensor, self.c_desc, "c_tensor", dynamic_m=True)
        d_desc = self._validate_live_tensor(d_row_tensor, self.d_row_desc, "d_row_tensor", dynamic_m=True)
        prob_desc = self._validate_live_tensor(prob_tensor, self.prob_desc, "prob_tensor", dynamic_m=True)
        dprob_desc = self._validate_live_tensor(dprob_tensor, self.dprob_desc, "dprob_tensor", dynamic_m=True)
        self._validate_live_tensor(padded_offsets, self.padded_offsets_desc, "padded_offsets")
        self._validate_live_tensor(alpha_tensor, self.alpha_desc, "alpha_tensor")
        self._validate_live_tensor(beta_tensor, self.beta_desc, "beta_tensor")

        two_n = c_desc.shape[1]
        self._expect_shape(c_desc, (tensor_m, two_n, 1), "c_tensor")
        self._expect_shape(d_desc, (tensor_m, two_n, 1), "d_row_tensor")
        self._expect_shape(prob_desc, (tensor_m, 1, 1), "prob_tensor")
        self._expect_shape(dprob_desc, (tensor_m, 1, 1), "dprob_tensor")
        self._expect_stride(a_desc, (k, 1, tensor_m * k), "a_tensor")
        self._expect_stride(c_desc, (two_n, 1, tensor_m * two_n), "c_tensor")
        self._expect_stride(d_desc, (two_n, 1, tensor_m * two_n), "d_row_tensor")
        self._expect_stride(prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "prob_tensor")
        self._expect_stride(dprob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "dprob_tensor")
        self._validate_offsets_once(padded_offsets, tensor_m=tensor_m)

        for tensor, name in (
            (a_tensor, "a_tensor"),
            (c_tensor, "c_tensor"),
            (d_row_tensor, "d_row_tensor"),
            (padded_offsets, "padded_offsets"),
            (alpha_tensor, "alpha_tensor"),
            (beta_tensor, "beta_tensor"),
            (prob_tensor, "prob_tensor"),
            (dprob_tensor, "dprob_tensor"),
        ):
            self._validate_data_alignment(tensor, name)

        if self._has_dbias:
            if dbias_tensor is None:
                raise ValueError("dbias_tensor is required because the API was compiled with sample_dbias")
            self._validate_live_tensor(dbias_tensor, self.dbias_desc, "dbias_tensor")
            self._validate_data_alignment(dbias_tensor, "dbias_tensor")
        elif dbias_tensor is not None:
            raise ValueError("dbias_tensor must be omitted because the API was compiled without sample_dbias")

        if self.weight_mode == MoEWeightMode.DENSE:
            if b_tensor is None or b_ptrs is not None:
                raise ValueError("Dense execution requires b_tensor and forbids b_ptrs")
            self._validate_live_tensor(b_tensor, self.b_desc, "b_tensor")
            self._validate_data_alignment(b_tensor, "b_tensor")
        else:
            if b_tensor is not None or b_ptrs is None:
                raise ValueError("Discrete execution requires b_ptrs and forbids b_tensor")
            _validate_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt)
            if get_device(b_ptrs) != self.a_desc.device:
                raise ValueError(f"b_ptrs must be on the same device as a_tensor ({self.a_desc.device}), got {get_device(b_ptrs)}")
            if get_data_ptr(b_ptrs) % 8 != 0:
                raise ValueError("b_ptrs data pointer must be 8-byte aligned")
            self._validate_pointer_values_once(b_ptrs)
            self._record_pointer_stream(b_ptrs, current_stream)

        self._compiled_kernel(
            a_tensor,
            c_tensor,
            d_row_tensor,
            padded_offsets,
            alpha_tensor,
            beta_tensor,
            prob_tensor,
            dprob_tensor,
            b_tensor,
            b_ptrs,
            dbias_tensor,
            current_stream,
            linear_offset,
        )
