[Bugfix] Add TRTLLM attention custom op and execution contract (#7214)

Signed-off-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>
Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Co-authored-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>

docs/design/feature/attention_execution_contract_poc.md

@@ -131,3 +131,24 @@ and real-kernel schema/fake checks pass. The full suite must be rerun with the
 vLLM version required by upstream. NPU/ROCm device numerics and compilation remain
 unvalidated. Pre-commit passes with the CI hook skips. Upstream PyTorch/CUTLASS
 warnings remain.
+
+## TRTLLM dense execution
+
+`TRTLLM_ATTN` declares `SUPPORTED` / `CUSTOM_OP` for noncausal dense BF16 on
+B200/GB200 (SM100) and B300/GB300 (SM103), head dimension 128, and equal Q/K/V
+head counts, without parallel, paged-KV, piecewise, or HSDP boundaries. Pre-construction, SAGE, skip-softmax, packed,
+and other unverified paths remain `UNMIGRATED`. Resolution and dispatch share
+metadata validation; workspace mutation is explicit in the custom-op schema.
+
+CPU contract tests cover fullgraph replay and schema/fake consistency using a
+substituted dispatcher. Real-kernel validation requires Blackwell with FlashInfer:
+
+```bash
+python -m pytest tests/diffusion/attention/test_trtllm_attn.py \
+  -k dense_contract_fullgraph_matches_sdpa -q -rs
+```
+
+These hardware tests compare eager/compiled output with FP32 SDPA and check
+schema/fake agreement. The TRTLLM and contract suite passed on Blackwell before the final architecture
+restriction was added: 61 passed, no skips (19 dependency deprecation warnings).
+The SM100/SM103 restriction is additionally covered by CPU contract tests.

tests/diffusion/attention/test_trtllm_attn.py

@@ -617,3 +617,68 @@ def test_skip_end_to_end_config_path(monkeypatch):
     ctx = fc.ForwardContext(denoise_timestep=0.3)
     monkeypatch.setattr(fc, "_forward_context", ctx)
     assert impl._resolve_skip_factor(4096) == pytest.approx(expected)
+
+
+@requires_trtllm_attn
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)),
+    reason="verified TRTLLM contract requires SM100/SM103",
+)
+@pytest.mark.parametrize("batch", [1, 2])
+def test_dense_contract_fullgraph_matches_sdpa(batch):
+    from vllm_omni.diffusion.attention.capabilities import (
+        CompilationMode,
+        ExecutionContext,
+        SupportStatus,
+    )
+
+    torch.compiler.reset()
+    try:
+        impl = _impl()
+        context = ExecutionContext(platform="cuda", require_fullgraph=True)
+        compiled = torch.compile(impl.forward_cuda, fullgraph=True, dynamic=True)
+        for q_len, kv_len in ((64, 64), (128, 192), (256, 128)):
+            # Noncontiguous Q exercises the reshape before the opaque boundary.
+            q = torch.randn(batch, 8, q_len, 128, device="cuda", dtype=torch.bfloat16).transpose(1, 2)
+            k, v = (torch.randn(batch, kv_len, 8, 128, device="cuda", dtype=torch.bfloat16) for _ in range(2))
+            result = impl.resolve_execution_path(context, q, k, v, None)
+            assert result.requested_support(context).status is SupportStatus.SUPPORTED
+            assert result.compilation_mode is CompilationMode.CUSTOM_OP
+            before = [t.clone() for t in (q, k, v)]
+            eager = impl.forward_cuda(q, k, v)
+            for _ in range(2):
+                out = compiled(q, k, v)
+                assert out.dtype == q.dtype and out.device == q.device and out.is_contiguous()
+                torch.testing.assert_close(out, eager, atol=1e-2, rtol=1e-2)
+                torch.testing.assert_close(out.float(), _sdpa_ref(q, k, v, 128**-0.5), atol=1e-2, rtol=1e-2)
+            for actual, original in zip((q, k, v), before):
+                torch.testing.assert_close(actual, original, atol=0, rtol=0)
+        # Check the real kernel's mutation and output metadata against its schema/fake.
+        qf, kf, vf = (t.reshape(-1, 8, 128).contiguous() for t in (q, k, v))
+        args = (
+            qf,
+            kf,
+            vf,
+            impl._get_workspace(q.device),
+            torch.full((batch,), kv_len, device=q.device, dtype=torch.int32),
+            torch.arange(batch + 1, device=q.device, dtype=torch.int32) * q_len,
+            torch.arange(batch + 1, device=q.device, dtype=torch.int32) * kv_len,
+            None,
+            None,
+            None,
+            q_len,
+            kv_len,
+            batch,
+            128**-0.5,
+            1.0,
+            -1.0,
+            0,
+            0,
+            False,
+        )
+        checks = torch.library.opcheck(
+            tg._trtllm_ragged_attention_op, args, test_utils=("test_schema", "test_faketensor")
+        )
+        assert all(value == "SUCCESS" for value in checks.values())
+    finally:
+        torch.compiler.reset()

tests/diffusion/attention/test_trtllm_contract.py

@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, PackedPaddingMetadata
+from vllm_omni.diffusion.attention.capabilities import (
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _impl(**kwargs):
+    return tg.TrtllmAttentionImpl(2, 128, 128**-0.5, **kwargs)
+
+
+def _inputs(q_len=8, kv_len=12, dtype=torch.bfloat16):
+    return (
+        torch.randn(1, q_len, 2, 128, dtype=dtype),
+        torch.randn(1, kv_len, 2, 128, dtype=dtype),
+        torch.randn(1, kv_len, 2, 128, dtype=dtype),
+    )
+
+
+@pytest.fixture
+def dispatcher(monkeypatch, tmp_path):
+    marker = tmp_path / "kernel"
+    marker.write_text("loaded")
+
+    @functools.cache
+    def load_kernel():
+        return marker.read_text()
+
+    def fake_attention(query, key, value, workspace_buffer, batch_size, bmm1_scale, **kwargs):
+        assert load_kernel() == "loaded"
+        assert kwargs["skip_all_rows_active_check"] is True
+        workspace_buffer.add_(1)
+        q, k, v = (t.reshape(batch_size, -1, t.shape[-2], t.shape[-1]).transpose(1, 2) for t in (query, key, value))
+        return (
+            F.scaled_dot_product_attention(q, k, v, scale=bmm1_scale)
+            .transpose(1, 2)
+            .reshape(query.shape[0], query.shape[1], value.shape[-1])
+            .contiguous()
+        )
+
+    def quantize(q, k, v, **kwargs):
+        scale = torch.ones(1)
+        return q, k, v, scale, scale, scale, None
+
+    # CPU tests substitute device selection; the real kernel tests use the GPU.
+    monkeypatch.setattr(tg, "_is_trtllm_blackwell_device", lambda query: True)
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+    monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention, raising=False)
+    monkeypatch.setattr(
+        tg, "flashinfer", SimpleNamespace(__version__="0.6.18", trtllm_sage_attention_quantize=quantize), raising=False
+    )
+    workspace = torch.zeros(16, dtype=torch.uint8)
+    monkeypatch.setattr(tg.TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: workspace))
+    torch.compiler.reset()
+    yield workspace
+    torch.compiler.reset()
+
+
+def test_dense_contract_uses_initialized_state(dispatcher):
+    context = ExecutionContext(platform="cuda", kernel_variant="caller_sage", causal=True, require_fullgraph=True)
+    result = _impl().resolve_execution_path(context, *_inputs(), None)
+    assert result.path == "trtllm_dense" and result.kernel_variant == "trtllm"
+    assert result.support.status is SupportStatus.SUPPORTED
+    assert result.compilation_mode is CompilationMode.CUSTOM_OP
+    assert result.requested_support(context).status is SupportStatus.SUPPORTED
+    pre = tg.TrtllmAttentionBackend.resolve_capabilities(context)
+    assert pre.kernel_variant is None and pre.support.status is SupportStatus.UNMIGRATED
+    assert pre.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+@pytest.mark.parametrize(
+    "change",
+    [
+        {"platform": "rocm"},
+        {"paged_kv": True},
+        {"parallel_strategy": ParallelStrategy.ULYSSES},
+        {"parallel_strategy": ParallelStrategy.RING},
+        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
+        {"parallel_strategy": ParallelStrategy.ALLGATHER_KV},
+        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
+    ],
+)
+def test_outer_paths_are_unmigrated(dispatcher, change):
+    context = replace(ExecutionContext(platform="cuda", require_fullgraph=True), **change)
+    result = _impl().resolve_execution_path(context, *_inputs(), None)
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+@pytest.mark.parametrize(
+    "kwargs,path",
+    [
+        ({"causal": True}, "trtllm_dense"),
+        ({"backend_kwargs": {"quant": {"dtype_qk": "int8"}}}, "trtllm_sage"),
+        ({"backend_kwargs": {"skip_softmax_threshold": 0.01}}, "trtllm_skip_softmax"),
+        ({"backend_kwargs": {"target_sparsity": 0.5}}, "trtllm_skip_softmax"),
+    ],
+)
+def test_unverified_backend_settings(dispatcher, kwargs, path):
+    result = _impl(**kwargs).resolve_execution_path(ExecutionContext(platform="cuda"), *_inputs(), None)
+    assert result.path == path and result.support.status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize("padding", [False, True])
+def test_packed_paths_are_not_dense(dispatcher, padding):
+    q, k, v = _inputs()
+    cuq, cuk = torch.tensor([0, 8], dtype=torch.int32), torch.tensor([0, 12], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cuq,
+            "cu_seqlens_k": cuk,
+            "max_seqlen_q": 8,
+            "max_seqlen_k": 12,
+        },
+        packed_padding=PackedPaddingMetadata(8, 12, cuq, cuk) if padding else None,
+    )
+    impl = _impl()
+    result = impl.resolve_execution_path(ExecutionContext(platform="cuda"), q, k, v, metadata)
+    assert result.path == ("trtllm_packed_padding" if padding else "trtllm_packed_varlen")
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert (
+        impl.resolve_execution_path(ExecutionContext(platform="cuda"), q, k, v, None).support.status
+        is SupportStatus.SUPPORTED
+    )
+
+
+@pytest.mark.parametrize(
+    "metadata,message",
+    [
+        (AttentionMetadata(attn_mask=torch.ones(1, 12)), "does not support attn_mask"),
+        (AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8])}), "Incomplete packed"),
+        (AttentionMetadata(packed_padding="invalid"), "must be PackedPaddingMetadata"),
+    ],
+)
+def test_resolution_and_execution_share_metadata_validation(dispatcher, metadata, message):
+    impl = _impl()
+    inputs = _inputs()
+    with pytest.raises(ValueError, match=message):
+        impl.resolve_execution_path(ExecutionContext(platform="cuda"), *inputs, metadata)
+    with pytest.raises(ValueError, match=message):
+        impl.forward_cuda(*inputs, metadata)
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    [
+        AttentionMetadata(full_attn_spans=[[(0, 4)]]),
+        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
+    ],
+)
+def test_special_metadata_is_unmigrated(dispatcher, metadata):
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *_inputs(), metadata)
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_invalid_geometry_and_missing_dependency(dispatcher, monkeypatch):
+    q, k, v = _inputs()
+    context = ExecutionContext(platform="cuda")
+    for inputs in ((q[0], k, v), (q, k.float(), v), (q, k, v[:, :-1])):
+        result = _impl().resolve_execution_path(context, *inputs, None)
+        assert result.support.status is SupportStatus.UNSUPPORTED and result.support.reason
+    for inputs in (
+        _inputs(dtype=torch.float32),
+        (q[..., :64], k[..., :64], v[..., :64]),
+        (q, k[:, :, :1], v[:, :, :1]),
+    ):
+        assert _impl().resolve_execution_path(context, *inputs, None).support.status is SupportStatus.UNMIGRATED
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", False)
+    result = _impl().resolve_execution_path(context, q, k, v, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED and "FlashInfer" in result.support.reason
+
+
+@pytest.mark.parametrize("use_sage", [False, True], ids=["dense", "sage"])
+def test_trtllm_dispatcher_is_opaque_to_torch_compile(dispatcher, use_sage):
+    impl = _impl(backend_kwargs={"quant": {"dtype_qk": "int8"}} if use_sage else None)
+    compiled = torch.compile(impl.forward_cuda, fullgraph=True, dynamic=True)
+    for q_len, kv_len in ((8, 12), (11, 17), (8, 12)):
+        q, k, v = _inputs(q_len, kv_len)
+        before = [t.clone() for t in (q, k, v)]
+        eager = impl.forward_cuda(q, k, v)
+        out = compiled(q, k, v)
+        torch.testing.assert_close(out, eager)
+        for actual, original in zip((q, k, v), before):
+            torch.testing.assert_close(actual, original, rtol=0, atol=0)
+    assert torch.all(dispatcher == 6)
+
+
+def test_custom_op_schema_fake_and_noncontiguous_output(dispatcher):
+    q = torch.randn(2, 7, 128).transpose(0, 1)
+    k, v = torch.randn(9, 2, 128), torch.randn(9, 2, 64)
+    args = (
+        q,
+        k,
+        v,
+        dispatcher,
+        torch.tensor([9], dtype=torch.int32),
+        torch.tensor([0, 7], dtype=torch.int32),
+        torch.tensor([0, 9], dtype=torch.int32),
+        None,
+        None,
+        None,
+        7,
+        9,
+        1,
+        128**-0.5,
+        1.0,
+        -1.0,
+        0,
+        0,
+        False,
+    )
+    result = torch.library.opcheck(tg._trtllm_ragged_attention_op, args, test_utils=("test_schema", "test_faketensor"))
+    assert all(v == "SUCCESS" for v in result.values())
+    out = tg._trtllm_ragged_attention_op(*args)
+    assert out.shape == (7, 2, 64) and out.is_contiguous()
+
+
+@pytest.mark.parametrize(
+    "capability,expected",
+    [
+        ((9, 0), False),
+        ((10, 0), True),
+        ((10, 3), True),
+        ((10, 7), False),
+        ((12, 0), False),
+        (None, False),
+    ],
+)
+def test_trtllm_device_capability_is_conservative(monkeypatch, capability, expected):
+    monkeypatch.setattr(
+        tg,
+        "current_omni_platform",
+        SimpleNamespace(
+            is_cuda=lambda: True,
+            get_device_capability=lambda device_id: capability,
+        ),
+    )
+    assert tg._is_trtllm_blackwell_device(SimpleNamespace(device=torch.device("cuda:0"))) is expected
+    assert not tg._is_trtllm_blackwell_device(SimpleNamespace(device=torch.device("cpu")))
+
+
+def test_non_blackwell_cannot_claim_fullgraph(dispatcher, monkeypatch):
+    monkeypatch.setattr(tg, "_is_trtllm_blackwell_device", lambda query: False)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    result = _impl().resolve_execution_path(context, *_inputs(), None)
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED

vllm_omni/diffusion/attention/backends/trtllm_attn.py

@@ -15,6 +15,16 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     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__)
 
@@ -88,6 +98,96 @@ except Exception as e:  # pragma: no cover - import guard
     )
 
 
+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,
@@ -156,9 +256,53 @@ def _workspace_bytes() -> int:
     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
@@ -231,6 +375,96 @@ class TrtllmAttentionImpl(AttentionImpl):
                 )
             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)
 
@@ -361,26 +595,7 @@ class TrtllmAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
-        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."
-            )
+        extra, has_packed_metadata, packed_padding = _normalize_trtllm_metadata(attn_metadata)
 
         if not HAS_FLASHINFER:
             raise ImportError(
@@ -449,10 +664,8 @@ class TrtllmAttentionImpl(AttentionImpl):
 
         _skip_factor = self._resolve_skip_factor(max_kv_len)
 
-        # SAGE kwargs are only understood by newer FlashInfer builds; pass them exclusively when
-        # SAGE quant is active (which already requires the kernel, checked at init) so the dense
-        # path stays compatible with older builds that lack these parameters.
-        sage_kwargs: dict = {}
+        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,
@@ -462,30 +675,29 @@ class TrtllmAttentionImpl(AttentionImpl):
                 cu_seq_lens_q,
                 cu_seq_lens_kv,
             )
-            sage_kwargs["sage_attn_sfs"] = sage_attn_sfs
-            sage_kwargs["num_elts_per_sage_attn_blk"] = sage_block_sizes
+            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_deepseek(
-            query=q,
-            key=k,
-            value=v,
-            workspace_buffer=workspace,
-            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,
-            window_left=-1,
-            cum_seq_lens_q=cu_seq_lens_q,
-            cum_seq_lens_kv=cu_seq_lens_kv,
-            enable_pdl=False,
-            is_causal=self.causal,
-            return_lse=False,
-            skip_softmax_threshold_scale_factor=_skip_factor,
-            skip_all_rows_active_check=True,
-            **sage_kwargs,
+        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(