vllm-omni/vllm_omni/diffusion/attention/backends/trtllm_attn.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import math
from dataclasses import dataclass, replace
from typing import NamedTuple, cast

import torch
from packaging.version import InvalidVersion, Version
from vllm.logger import init_logger

from vllm_omni.diffusion.attention.backends.abstract import (
    AttentionBackend,
    AttentionImpl,
    AttentionMetadata,
    PackedPaddingMetadata,
)
from vllm_omni.diffusion.attention.capabilities import (
    CapabilityResult,
    CompilationMode,
    ExecutionContext,
    ExecutionPathResult,
    MaskMode,
    PackingMode,
    ParallelStrategy,
)
from vllm_omni.platforms import current_omni_platform

logger = init_logger(__name__)


def _validate_control(value, name: str, lo: float, hi: float | None) -> float | None:
    if value is None:
        return None
    v = float(value)
    if not math.isfinite(v) or v < lo or (hi is not None and v > hi):
        rng = f"in [{lo}, {hi}]" if hi is not None else f">= {lo}"
        raise ValueError(f"{name} must be finite and {rng}; got {value!r}.")
    return v


@dataclass(frozen=True)
class SkipSoftmaxConfig:
    threshold: float | None = None
    target_sparsity: float | None = None
    disabled_until_timestep: float = 0.0
    a: float | None = None
    b: float | None = None

    @classmethod
    def from_backend_kwargs(cls, backend_kwargs: dict | None) -> "SkipSoftmaxConfig":
        bk = backend_kwargs or {}
        return cls(
            threshold=_validate_control(bk.get("skip_softmax_threshold"), "skip_softmax_threshold", 0.0, None),
            target_sparsity=_validate_control(bk.get("target_sparsity"), "target_sparsity", 0.0, 1.0),
            disabled_until_timestep=_validate_control(
                bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0
            )
            or 0.0,
        )

    @property
    def enabled(self) -> bool:
        return self.threshold is not None or (
            self.target_sparsity is not None and self.a is not None and self.b is not None
        )

    @property
    def configured(self) -> bool:
        return self.threshold is not None or self.target_sparsity is not None

    @property
    def gated(self) -> bool:
        return self.disabled_until_timestep > 0.0

    def resolve_factor(self, seqlen: int, timestep: float | None) -> float | None:
        if self.threshold is not None:
            factor = self.threshold * seqlen
        elif self.target_sparsity is not None and self.a is not None and self.b is not None:
            factor = self.a * math.exp(self.b * self.target_sparsity)
        else:
            return None
        if self.gated and timestep is not None and timestep > self.disabled_until_timestep:
            return None
        return factor


try:
    import flashinfer
    from flashinfer.prefill import trtllm_ragged_attention_deepseek

    HAS_FLASHINFER = True
except Exception as e:  # pragma: no cover - import guard
    HAS_FLASHINFER = False
    logger.warning(
        "FlashInfer is unavailable; TRTLLM_ATTN backend will not work. Reason: %s",
        e,
    )


if not hasattr(torch.ops.vllm_omni, "trtllm_ragged_attention"):

    @torch.library.custom_op(
        "vllm_omni::trtllm_ragged_attention",
        mutates_args=("workspace_buffer",),
    )
    def _trtllm_ragged_attention_op(
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        workspace_buffer: torch.Tensor,
        seq_lens: torch.Tensor,
        cum_seq_lens_q: torch.Tensor,
        cum_seq_lens_kv: torch.Tensor,
        sage_q_sf: torch.Tensor | None,
        sage_k_sf: torch.Tensor | None,
        sage_v_sf: torch.Tensor | None,
        max_q_len: int,
        max_kv_len: int,
        batch_size: int,
        bmm1_scale: float,
        bmm2_scale: float,
        skip_softmax_threshold_scale_factor: float,
        sage_q_block_size: int,
        sage_k_block_size: int,
        is_causal: bool,
    ) -> torch.Tensor:
        sage_kwargs = {}
        if sage_q_sf is not None:
            sage_kwargs = {
                "sage_attn_sfs": (sage_q_sf, sage_k_sf, None, sage_v_sf),
                "num_elts_per_sage_attn_blk": (sage_q_block_size, sage_k_block_size, 0, 1),
            }
        return trtllm_ragged_attention_deepseek(
            query=query,
            key=key,
            value=value,
            workspace_buffer=workspace_buffer,
            seq_lens=seq_lens,
            max_q_len=max_q_len,
            max_kv_len=max_kv_len,
            bmm1_scale=bmm1_scale,
            bmm2_scale=bmm2_scale,
            o_sf_scale=-1.0,
            batch_size=batch_size,
            window_left=-1,
            cum_seq_lens_q=cum_seq_lens_q,
            cum_seq_lens_kv=cum_seq_lens_kv,
            enable_pdl=False,
            is_causal=is_causal,
            return_lse=False,
            skip_softmax_threshold_scale_factor=(
                None if skip_softmax_threshold_scale_factor < 0.0 else skip_softmax_threshold_scale_factor
            ),
            skip_all_rows_active_check=True,
            **sage_kwargs,
        )

    @_trtllm_ragged_attention_op.register_fake
    def _trtllm_ragged_attention_fake(
        query,
        key,
        value,
        workspace_buffer,
        seq_lens,
        cum_seq_lens_q,
        cum_seq_lens_kv,
        sage_q_sf,
        sage_k_sf,
        sage_v_sf,
        max_q_len,
        max_kv_len,
        batch_size,
        bmm1_scale,
        bmm2_scale,
        skip_softmax_threshold_scale_factor,
        sage_q_block_size,
        sage_k_block_size,
        is_causal,
    ):
        # Match FlashInfer's default output dtype for quantized queries.
        out_dtype = (
            torch.bfloat16 if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2, torch.int8) else query.dtype
        )
        return torch.empty((*query.shape[:-1], value.shape[-1]), dtype=out_dtype, device=query.device)


_trtllm_ragged_attention_op = torch.ops.vllm_omni.trtllm_ragged_attention


_QK_QUANT_DTYPES = {
    "int8": torch.int8,
    "fp8_e4m3": torch.float8_e4m3fn,
}


@dataclass(frozen=True)
class QuantConfig:
    dtype_qk: str | None = None
    q_block_size: int = 1
    k_block_size: int = 16

    @classmethod
    def from_backend_kwargs(cls, backend_kwargs: dict | None) -> "QuantConfig":
        q = (backend_kwargs or {}).get("quant") or {}
        return cls(
            dtype_qk=q.get("dtype_qk"),
            q_block_size=int(q.get("q_block_size", 1) or 1),
            k_block_size=int(q.get("k_block_size", 16) or 16),
        )

    @property
    def enabled(self) -> bool:
        return self.dtype_qk is not None

    def quantize(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        quantize_fn,
        cu_seq_lens_q: torch.Tensor,
        cu_seq_lens_kv: torch.Tensor,
    ):
        qk_quant_dtype = _QK_QUANT_DTYPES[cast(str, self.dtype_qk)]
        q_q, k_q, v_q, q_sfs, k_sfs, v_sfs, _ = quantize_fn(
            q,
            k,
            v,
            q_block_size=self.q_block_size,
            k_block_size=self.k_block_size,
            qk_quant_dtype=qk_quant_dtype,
            smooth_k=True,
            cum_seq_lens_q=cu_seq_lens_q,
            cum_seq_lens_kv=cu_seq_lens_kv,
        )
        sage_attn_sfs = (q_sfs, k_sfs, None, v_sfs)
        num_elts_per_sage_attn_blk = (self.q_block_size, self.k_block_size, 0, 1)
        return q_q, k_q, v_q, sage_attn_sfs, num_elts_per_sage_attn_blk


class _PackedLayout(NamedTuple):
    """Active token ranges and kernel metadata for a flattened ragged batch."""

    q_tokens: int
    kv_tokens: int
    cu_seqlens_q: torch.Tensor
    cu_seqlens_kv: torch.Tensor
    batch_size: int
    seq_lens: torch.Tensor


def _workspace_bytes() -> int:
    import vllm.envs as envs

    return getattr(envs, "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", 394 * 1024 * 1024)


def _is_trtllm_blackwell_device(query: torch.Tensor) -> bool:
    # Capability inspection is eager; leave device queries outside forward.
    if query.device.type != "cuda" or not current_omni_platform.is_cuda():
        return False
    capability = current_omni_platform.get_device_capability(query.device.index)
    return capability is not None and tuple(capability) in ((10, 0), (10, 3))


def _normalize_trtllm_metadata(
    attn_metadata: AttentionMetadata | None,
) -> tuple[dict, bool, PackedPaddingMetadata | None]:
    """Inspect host metadata only; never read device tensor values."""
    extra = getattr(attn_metadata, "extra", {}) if attn_metadata is not None else {}
    packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
    present_packed_keys = [key for key in packed_keys if key in extra]
    if present_packed_keys and len(present_packed_keys) != len(packed_keys):
        missing = sorted(set(packed_keys) - set(present_packed_keys))
        raise ValueError(f"Incomplete packed TRTLLM attention metadata; missing {missing}")
    has_packed_metadata = len(present_packed_keys) == len(packed_keys)

    attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None
    packed_padding = getattr(attn_metadata, "packed_padding", None) if attn_metadata is not None else None
    if packed_padding is not None and not isinstance(packed_padding, PackedPaddingMetadata):
        raise ValueError("packed_padding must be PackedPaddingMetadata")
    if packed_padding is not None and not has_packed_metadata:
        raise ValueError("PackedPaddingMetadata requires complete packed TRTLLM attention metadata")
    if attn_mask is not None:
        raise ValueError(
            "TRTLLM_ATTN does not support attn_mask. Represent structural suffix padding "
            "with packed-padding metadata, or select a mask-capable backend such as "
            "CUDNN_ATTN or TORCH_SDPA."
        )

    return extra, has_packed_metadata, packed_padding


class TrtllmAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @classmethod
    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
        # Construction decides SAGE/skip settings; caller-provided kernel identity
        # cannot establish the initialized path or a fullgraph guarantee.
        return ExecutionPathResult.unmigrated(
            cls.get_name(), replace(context, kernel_variant=None), path="uninitialized"
        )

    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        return True

    @classmethod
    def supports_multi_doc_packed_varlen(cls) -> bool:
        return True

    @staticmethod
    def get_supported_head_sizes() -> list[int]:
        return [128]

    @staticmethod
    def get_name() -> str:
        return "TRTLLM_ATTN"

    @staticmethod
    def get_impl_cls() -> type["TrtllmAttentionImpl"]:
        return TrtllmAttentionImpl


class TrtllmAttentionImpl(AttentionImpl):
    _workspace: torch.Tensor | None = None

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
        num_kv_heads: int | None = None,
        prefix: str = "",
        qkv_layout: str | None = None,
        backend_kwargs: dict | None = None,
        role: str = "self",
        **extra_impl_args,
    ) -> None:
        self.num_heads = num_heads
        self.head_size = head_size
        self.softmax_scale = softmax_scale
        self.causal = causal
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
        self.role = role

        self.skip = SkipSoftmaxConfig.from_backend_kwargs(backend_kwargs)
        self._warned_missing_timestep = False

        # Override: SageAttention does not support causal attention
        if causal:
            self.quant = QuantConfig(dtype_qk=None, q_block_size=0, k_block_size=0)
        else:
            self.quant = QuantConfig.from_backend_kwargs(backend_kwargs)
        # Resolve the SAGE quantize fn once at init so the compiled forward path
        # does not import it every step.
        self._sage_quantize_fn = None
        if self.quant.enabled:
            if self.quant.dtype_qk not in _QK_QUANT_DTYPES:
                raise RuntimeError(
                    f"TRTLLM_ATTN quant (SAGE) supports dtype_qk in {sorted(_QK_QUANT_DTYPES)}, got "
                    f"{self.quant.dtype_qk!r}. FLASHINFER_ATTN dtypes (float16/bfloat16) are not SAGE."
                )
            try:
                flashinfer_version = Version(flashinfer.__version__)
            except (AttributeError, InvalidVersion):
                flashinfer_version = None
            if flashinfer_version is not None and flashinfer_version < Version("0.6.18rc10"):
                raise RuntimeError(
                    f"FlashInfer {flashinfer_version} is too old for TRTLLM_ATTN quant (SAGE); "
                    "install flashinfer >= 0.6.18rc10."
                )
            self._sage_quantize_fn = flashinfer.trtllm_sage_attention_quantize

    def resolve_execution_path(
        self,
        context: ExecutionContext,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None,
    ) -> ExecutionPathResult:
        extra, packed, packed_padding = _normalize_trtllm_metadata(attn_metadata)
        packing_mode = PackingMode.NONE
        if packed_padding is not None:
            packing_mode = PackingMode.PACKED_PADDING
        elif packed:
            packing_mode = PackingMode.MULTI_DOCUMENT
        context = replace(
            context,
            kernel_variant="trtllm_sage" if self.quant.enabled else "trtllm",
            dtype=str(query.dtype).removeprefix("torch."),
            causal=self.causal,
            mask_mode=MaskMode.NONE,
            packing_mode=packing_mode,
            piecewise=attn_metadata is not None and attn_metadata.full_attn_spans is not None,
            kv_cache_dtype=extra.get("kv_cache_dtype"),
        )
        if self.quant.enabled:
            path = "sage"
        elif self.skip.configured:
            path = "skip_softmax"
        elif packed_padding is not None:
            path = "packed_padding"
        elif packed:
            path = "packed_varlen"
        else:
            path = "dense"
        result = ExecutionPathResult.unmigrated("TRTLLM_ATTN", context, path="trtllm_" + path)
        if not HAS_FLASHINFER:
            return replace(
                result,
                support=CapabilityResult.unsupported(
                    "TRTLLM_ATTN requires FlashInfer; install it or select another backend."
                ),
            )
        if (
            context.platform != "cuda"
            or not _is_trtllm_blackwell_device(query)
            or path != "dense"
            or self.causal
            or context.dtype != "bfloat16"
            or context.piecewise
            or context.paged_kv
            or context.kv_cache_dtype not in (None, "auto", "float")
            or context.parallel_strategy is not ParallelStrategy.NONE
            or context.outer_boundaries
        ):
            return replace(
                result,
                support=CapabilityResult.unmigrated(
                    "Only dense noncausal BF16 TRTLLM on Blackwell SM100/SM103 "
                    "without parallel or outer boundaries is migrated."
                ),
            )
        if any(t.ndim != 4 for t in (query, key, value)):
            reason = "Q, K, and V must have rank 4 (batch, sequence, heads, head dimension)."
        elif query.dtype != key.dtype or query.dtype != value.dtype:
            reason = "Q, K, and V dtypes must match."
        elif query.device != key.device or query.device != value.device:
            reason = "Q, K, and V devices must match."
        elif key.shape != value.shape or query.shape[0] != key.shape[0] or query.shape[-1] != key.shape[-1]:
            reason = "K/V shapes, Q/K batch sizes, and Q/K head dimensions must match."
        elif any(size == 0 for t in (query, key, value) for size in t.shape):
            reason = "Q, K, and V dimensions must be nonzero."
        else:
            reason = None
        if reason:
            return replace(
                result,
                support=CapabilityResult.unsupported(
                    "TRTLLM_ATTN: " + reason + " Select compatible inputs or another backend."
                ),
            )
        # This is a verified geometry, not an exhaustive list of kernel limits.
        if query.shape[-1] != 128 or query.shape[2] != key.shape[2]:
            return replace(
                result,
                support=CapabilityResult.unmigrated(
                    "TRTLLM fullgraph validation covers head dimension 128 with equal Q/K/V head counts."
                ),
            )
        return replace(result, support=CapabilityResult.supported(), compilation_mode=CompilationMode.CUSTOM_OP)

    def set_layer_calibration(self, a: float, b: float) -> None:
        self.skip = replace(self.skip, a=a, b=b)

    def _resolve_skip_factor(self, seqlen: int) -> float | None:
        if not self.skip.enabled:
            return None

        timestep = None
        if self.skip.gated:
            from vllm_omni.diffusion.forward_context import get_forward_context

            timestep = getattr(get_forward_context(), "denoise_timestep", None)
            if timestep is None:
                if not self._warned_missing_timestep:
                    logger.warning(
                        "TRTLLM skip: disabled_until_timestep=%s set but this pipeline does not "
                        "publish denoise_timestep; staying dense. Have the pipeline call "
                        "DenoiseProgressMixin.record_denoise_step to enable timestep gating.",
                        self.skip.disabled_until_timestep,
                    )
                    self._warned_missing_timestep = True
                return None
        return self.skip.resolve_factor(seqlen, timestep)

    @classmethod
    def _get_workspace(cls, device: torch.device) -> torch.Tensor:
        nbytes = _workspace_bytes()
        ws = cls._workspace
        if ws is None or ws.device != device or ws.numel() < nbytes:
            ws = torch.zeros(nbytes, dtype=torch.uint8, device=device)
            cls._workspace = ws
        return ws

    @staticmethod
    def _prepare_packed_padding_layout(
        physical_batch: int,
        q: torch.Tensor,
        k: torch.Tensor,
        packed_padding: PackedPaddingMetadata,
        extra: dict,
    ) -> _PackedLayout:
        """Prepare a [real, pad] layout from an explicit host-side boundary."""
        valid_q_tokens = packed_padding.q_length
        valid_kv_tokens = packed_padding.kv_length
        max_q_len = extra["max_seqlen_q"]
        max_kv_len = extra["max_seqlen_k"]
        lengths = (valid_q_tokens, valid_kv_tokens, max_q_len, max_kv_len)
        if any(isinstance(value, bool) or not isinstance(value, int) for value in lengths):
            raise ValueError("Mask-free packed TRTLLM attention lengths must be Python integers")
        if physical_batch != 1:
            raise ValueError("Mask-free packed TRTLLM attention requires a single packed batch")

        total_q_tokens = q.shape[0]
        total_kv_tokens = k.shape[0]
        if not 0 < valid_q_tokens <= total_q_tokens:
            raise ValueError(
                "PackedPaddingMetadata.q_length must be within the packed Q sequence, "
                f"got {valid_q_tokens} for length {total_q_tokens}"
            )
        if not 0 < valid_kv_tokens <= total_kv_tokens:
            raise ValueError(
                "PackedPaddingMetadata.kv_length must be within the packed K/V sequence, "
                f"got {valid_kv_tokens} for length {total_kv_tokens}"
            )
        if max_q_len != valid_q_tokens or max_kv_len != valid_kv_tokens:
            raise ValueError("Mask-free packed TRTLLM attention lengths must match max_seqlen_q/k")
        published_kv_length = extra.get("valid_kv_length")
        if published_kv_length is not None:
            if isinstance(published_kv_length, bool) or not isinstance(published_kv_length, int):
                raise ValueError("valid_kv_length must be a Python integer")
            if published_kv_length != valid_kv_tokens:
                raise ValueError("PackedPaddingMetadata.kv_length must match valid_kv_length")

        cu_seq_lens_q = packed_padding.cu_seqlens_q
        cu_seq_lens_kv = packed_padding.cu_seqlens_k
        if cu_seq_lens_q.dtype != torch.int32 or cu_seq_lens_kv.dtype != torch.int32:
            raise ValueError("Mask-free packed TRTLLM attention requires int32 cu_seqlens")
        if cu_seq_lens_q.device != q.device or cu_seq_lens_kv.device != k.device:
            raise ValueError("Mask-free packed TRTLLM attention metadata must be on the Q/K device")
        if cu_seq_lens_q.shape != (2,) or cu_seq_lens_kv.shape != (2,):
            raise ValueError("Mask-free packed TRTLLM attention requires canonical two-element cu_seqlens")
        return _PackedLayout(
            q_tokens=valid_q_tokens,
            kv_tokens=valid_kv_tokens,
            cu_seqlens_q=cu_seq_lens_q,
            cu_seqlens_kv=cu_seq_lens_kv,
            batch_size=1,
            # Canonical packed-padding metadata is [0, valid_kv_tokens], so this slice
            # is the one-element sequence-length view expected by the kernel.
            seq_lens=cu_seq_lens_kv[1:],
        )

    @staticmethod
    def _prepare_generic_packed_layout(
        q: torch.Tensor,
        k: torch.Tensor,
        cu_seq_lens_q: torch.Tensor,
        cu_seq_lens_kv: torch.Tensor,
    ) -> _PackedLayout:
        """Validate and normalize a general ragged batch."""
        if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]:
            raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens")

        # Packed producers may retain a zero-length trailing segment when the
        # input is already aligned. It is not a real sequence and is removed
        # before deriving the batch size and per-sequence lengths.
        while (
            cu_seq_lens_q.numel() > 2
            and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item())
            and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item())
        ):
            cu_seq_lens_q = cu_seq_lens_q[:-1]
            cu_seq_lens_kv = cu_seq_lens_kv[:-1]

        return _PackedLayout(
            q_tokens=q.shape[0],
            kv_tokens=k.shape[0],
            cu_seqlens_q=cu_seq_lens_q,
            cu_seqlens_kv=cu_seq_lens_kv,
            batch_size=cu_seq_lens_q.numel() - 1,
            seq_lens=(cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous(),
        )

    def forward_cuda(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        extra, has_packed_metadata, packed_padding = _normalize_trtllm_metadata(attn_metadata)

        if not HAS_FLASHINFER:
            raise ImportError(
                "TRTLLM_ATTN backend requires flashinfer. Install it or select "
                "another backend via --diffusion-attention-backend."
            )

        physical_batch, q_len, num_q_heads, head_dim = query.shape
        kv_len, num_kv_heads = key.shape[1], key.shape[2]
        device = query.device

        q = query.reshape(physical_batch * q_len, num_q_heads, head_dim).contiguous()
        k = key.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous()
        v = value.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous()
        output_tokens = q.shape[0]

        if has_packed_metadata:
            cu_seq_lens_q = extra["cu_seqlens_q"]
            cu_seq_lens_kv = extra["cu_seqlens_k"]
            if cu_seq_lens_q.ndim != 1 or cu_seq_lens_kv.ndim != 1:
                raise ValueError("Packed TRTLLM attention cu_seqlens tensors must be one-dimensional")
            if cu_seq_lens_q.numel() != cu_seq_lens_kv.numel() or cu_seq_lens_q.numel() < 2:
                raise ValueError("Packed TRTLLM attention requires matching non-empty Q and KV sequence batches")

            # PackedPaddingMetadata is an optional producer-published shortcut
            # for structural suffix padding. Inputs without it retain the
            # general cu_seqlens path for arbitrary ragged Q/KV batches.
            if packed_padding is not None:
                packed_layout = self._prepare_packed_padding_layout(
                    physical_batch,
                    q,
                    k,
                    packed_padding,
                    extra,
                )
            else:
                # Generic metadata can describe any number of real sequences,
                # not only a single [real, pad] pair.
                packed_layout = self._prepare_generic_packed_layout(
                    q,
                    k,
                    cu_seq_lens_q,
                    cu_seq_lens_kv,
                )

            q = q[: packed_layout.q_tokens]
            k = k[: packed_layout.kv_tokens]
            v = v[: packed_layout.kv_tokens]
            cu_seq_lens_q = packed_layout.cu_seqlens_q
            cu_seq_lens_kv = packed_layout.cu_seqlens_kv
            batch = packed_layout.batch_size
            seq_lens = packed_layout.seq_lens
            max_q_len = int(extra["max_seqlen_q"])
            max_kv_len = int(extra["max_seqlen_k"])
        else:
            batch = physical_batch
            seq_lens = torch.full((batch,), kv_len, dtype=torch.int32, device=device)
            cu_seq_lens_q = torch.arange(0, (batch + 1) * q_len, step=q_len, dtype=torch.int32, device=device)
            cu_seq_lens_kv = torch.arange(0, (batch + 1) * kv_len, step=kv_len, dtype=torch.int32, device=device)
            max_q_len = q_len
            max_kv_len = kv_len
        workspace = self._get_workspace(device)

        bmm1_scale = self.softmax_scale
        bmm2_scale = 1.0

        _skip_factor = self._resolve_skip_factor(max_kv_len)

        sage_q_sf = sage_k_sf = sage_v_sf = None
        sage_q_block_size = sage_k_block_size = 0
        if self.quant.enabled:
            q, k, v, sage_attn_sfs, sage_block_sizes = self.quant.quantize(
                q,
                k,
                v,
                self._sage_quantize_fn,
                cu_seq_lens_q,
                cu_seq_lens_kv,
            )
            sage_q_sf, sage_k_sf, _, sage_v_sf = sage_attn_sfs
            sage_q_block_size, sage_k_block_size, _, _ = sage_block_sizes

        out = _trtllm_ragged_attention_op(
            q,
            k,
            v,
            workspace,
            seq_lens,
            cu_seq_lens_q,
            cu_seq_lens_kv,
            sage_q_sf,
            sage_k_sf,
            sage_v_sf,
            max_q_len,
            max_kv_len,
            batch,
            bmm1_scale,
            bmm2_scale,
            -1.0 if _skip_factor is None else _skip_factor,
            sage_q_block_size,
            sage_k_block_size,
            self.causal,
        )
        if out.shape[0] != output_tokens:
            padded_out = torch.zeros(
                (output_tokens, num_q_heads, head_dim),
                dtype=out.dtype,
                device=out.device,
            )
            padded_out[: out.shape[0]] = out
            out = padded_out
        return out.reshape(physical_batch, q_len, num_q_heads, head_dim)