diff --git a/docs/design/feature/attention_execution_contract_poc.md b/docs/design/feature/attention_execution_contract_poc.md index 8c7486239..35af5455b 100644 --- a/docs/design/feature/attention_execution_contract_poc.md +++ b/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. diff --git a/tests/diffusion/attention/test_trtllm_attn.py b/tests/diffusion/attention/test_trtllm_attn.py index fad20e6be..771d73d7b 100644 --- a/tests/diffusion/attention/test_trtllm_attn.py +++ b/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() diff --git a/tests/diffusion/attention/test_trtllm_contract.py b/tests/diffusion/attention/test_trtllm_contract.py new file mode 100644 index 000000000..1e02af463 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/trtllm_attn.py b/vllm_omni/diffusion/attention/backends/trtllm_attn.py index c9257f988..813c5fd0d 100644 --- a/vllm_omni/diffusion/attention/backends/trtllm_attn.py +++ b/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(