[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>
This commit is contained in:
Rahul SteigerandRahul Steiger authored and GitHub committed
1 parent
b65f97bc48
commit
4934741df9
4 files changed
+609
-47
No files matched your search
@@ -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()
|
||||
@@ -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
|
||||
Reference in new issue
Block a user