diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index f58609e5f6e7..daa1e0d8b70c 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -2,15 +2,15 @@ schema_version: 1 vendors: flashinfer-prims-ts: url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev-20261012 - commit: ae312f7cc1723b06fda5168105ade0f328ee47d8 + branch: bug-6791733-prims-ts-cuda-graph--source-q-tma-release13 + commit: ae3851ab5f958ba22f154059df5b38321f6369d2 source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch patch_digest: sha256:65cc08e20a9cc7bdc190eb6f3e79124986dced74394c1f3250ae2d100a8ee340 - digest: sha256-tree-v1:e411b6c6679474d362b166e586df4276f52322d75a35c59c782d6266143e560e + digest: sha256-tree-v1:c9bcfbbac9db243110b7f9249ec0edd8ba4447ea03a7c89b1c7f6649d6d4bb6e sana-sol-attn: url: https://github.com/NVlabs/Sana.git branch: sol-engine diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index 1163324ec632..9c20e4356d79 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -328,6 +328,9 @@ def _make_context_kernel( enable_skip_correction=True, uses_ldtm_stat=uses_ldtm_stat, causal_single_kv_tile=(causal_single_kv_tile and not use_paged_kv), + use_blackwell_tma_oob_workaround=( + torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3)) + ), **paged_kwargs, ) return fmha diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py index c74fbcc1287c..9b23b5dd7d38 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py @@ -200,6 +200,22 @@ from cutlass.experimental import primitives as prims +def _compact_q_prefix_min_rows(q: cute.Tensor) -> int | None: + """Rows needed for a proven 128 KiB contiguous prefix, or None if unproven.""" + if len(q.shape) != 3: + return None + h, d = q.shape[1:] + element_bits = q.element_type.width + if not all(isinstance(x, int) for x in (h, d, element_bits, *q.stride)): + return None + if h <= 0 or d <= 0 or element_bits not in (8, 16): + return None + if tuple(q.stride) != (h * d, d, 1): + return None + row_bytes = h * d * (element_bits // 8) + return (128 * 1024 + row_bytes - 1) // row_bytes + + def _as_i32(x: int | Int32) -> Int32: return Int32(x) if isinstance(x, int) else x @@ -2743,6 +2759,9 @@ class FmhaTs: Use the fixed causal one-K/V-tile task domains. The context runner enables this only for query-paired, fixed-length inputs whose K/V extent fits one 128-token tile (default: False). + use_blackwell_tma_oob_workaround : bool, optional + Use conservative OOB address generation for ragged Q tensor maps. + Enable only for a Blackwell execution device (default: False). """ def __init__( @@ -2772,6 +2791,7 @@ def __init__( causal_single_kv_tile: bool = False, exhaustive_deadlock_race_check: bool = True, d_v: int | None = None, + use_blackwell_tma_oob_workaround: bool = False, ) -> None: """Initialize mode-specific tiling, dtype, and schedule configuration.""" head_paired = resolve_head_paired_mode( @@ -2815,6 +2835,7 @@ def __init__( self.is_causal = is_causal self.is_clc_dynamic = is_clc_dynamic self.exhaustive_deadlock_race_check = exhaustive_deadlock_race_check + self.use_blackwell_tma_oob_workaround = use_blackwell_tma_oob_workaround q_dtype = in_qk_dtype or cutlass.Float16 k_dtype = in_qk_dtype or cutlass.Float16 @@ -3162,6 +3183,24 @@ def __call__( # Zero-fill partial Q rows and head-dimension fragments. oob_fill=cuda.TensorMapFloatOOBFill.none, ) + if cutlass.const_expr(self.use_blackwell_tma_oob_workaround): + # Only the real Q layout can establish a mapped prefix; + # synthetic ragged dimensions and padded head sizes cannot. + minimum_rows = _compact_q_prefix_min_rows(q_cute) + correct_oob_addressing = True + if cutlass.const_expr(minimum_rows is not None): + correct_oob_addressing = q_cute.shape[0] < minimum_rows + if correct_oob_addressing: + # Work around the Blackwell driver TMA descriptor issue documented at + # https://docs.nvidia.com/datacenter/tesla/tesla-release-notes-615-71-09/index.html#known-issues + # Clear Blackwell descriptor bit 85 on the host when a + # contiguous 128 KiB prefix has not been established. + descriptor_word1 = ( + cutlass.Pointer(tma_q_desc.get_ptr(), dtype=cutlass.Uint64) + 1 + ) + descriptor_word1.store( + descriptor_word1.load() & ~cutlass.Uint64(1 << 21) + ) else: tma_q_desc = cuda.create_tensor_map_tiled_from_view( q_cute, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index ed023f9e176f..14fbe8a6ebb9 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -100,8 +100,11 @@ l0_b200: # by absolute index, and they gate on >= 2 GPUs plus Blackwell. They also # rebind mpi4py's serializer at import time, which applies process-wide on # every stage that collects them, including where the tests themselves skip. - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + # test_prims_ts_attention_backend.py keeps its own entry so its kernel JIT time does not + # count against this entry's timeout. + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/test_prims_ts_attention_backend.py - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120) + - unittest/_torch/attention/test_prims_ts_attention_backend.py - unittest/_torch/compilation - unittest/_torch/debugger - unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 3a0e50c09bae..924a25b3fc48 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -24,11 +24,14 @@ l0_b300: # rebind mpi4py's serializer at import time, which applies process-wide on # every stage that collects them, including where the tests themselves skip. # The sparse/dsa/test_cute_dsl_* files ignored below keep their own entries (pre-existing). - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py + # test_prims_ts_attention_backend.py keeps its own entry so its kernel JIT time does not + # count against this entry's timeout. + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py --ignore=unittest/_torch/attention/test_prims_ts_attention_backend.py - unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py - unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120) - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py + - unittest/_torch/attention/test_prims_ts_attention_backend.py - unittest/_torch/modeling/test_qsa_runtime_wiring.py - unittest/_torch/modeling/test_qwen4_exp_support.py - unittest/_torch/thop/parallel TIMEOUT (90) diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index 2adfac26c11c..878512866e7c 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -16,6 +16,7 @@ import functools import inspect import math +from collections.abc import Callable from types import SimpleNamespace from unittest.mock import Mock @@ -30,6 +31,15 @@ ) +@pytest.fixture(autouse=True) +def _trtllm_backend_only(monkeypatch: pytest.MonkeyPatch) -> None: + # PrimTS runs inside the TRTLLM backend. Other backends have their own + # tests and would only add JIT time here. + import backend_case + + monkeypatch.setattr(backend_case, "BACKENDS_UNDER_TEST", ("TRTLLM",)) + + _QWEN2_7B = { "num_heads": 28, "num_kv_heads": 4, @@ -82,9 +92,10 @@ }, id="generation", ), + # The context request matches the context case so both reuse one kernel. pytest.param( { - "seq_lens": [41, 1], + "seq_lens": [65, 1], "num_cached_tokens": [0, 63], "num_contexts": 1, }, @@ -192,29 +203,6 @@ def run_with_zeroed_v_tail_check( assert zeroed_tails == [(0, 1), (1, 5)] -def test_prims_ts_fp16_dense_context_with_alternate_shape( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") - case = BackendCase( - num_heads=8, - num_kv_heads=2, - head_dim=256, - seq_lens=[67, 33], - num_cached_tokens=[0, 0], - num_contexts=2, - dtype="float16", - causal=False, - kv_layout="HND", - page_size=64, - use_kv_cache_manager_v2=True, - ) - - results = run_case(case) - - assert "TRTLLM" in results - - def test_prims_ts_uses_compact_preprocessing_and_separate_decode_workspace( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -371,8 +359,11 @@ def test_prims_ts_deepseek_v3_lite_mla_generation( run_case(case) -@pytest.mark.parametrize("num_heads", [6, 12, 96]) -@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +# 6 and 96 heads select the M8 and repeated-M64 Q tiles; the scales test below +# covers the M16 tile with 12 heads. +@pytest.mark.parametrize( + "use_kv_cache_manager_v2,num_heads", [(False, 6), (True, 96)], ids=["v1-6", "v2-96"] +) def test_prims_ts_fp8_mla_preprocessing( monkeypatch: pytest.MonkeyPatch, num_heads: int, use_kv_cache_manager_v2: bool ) -> None: @@ -436,7 +427,7 @@ def run_and_capture(fmha: PrimsTSFmha, params: FmhaParams) -> None: assert calls == [0, 1, 0, 1] -@pytest.mark.parametrize("num_heads", [6, 12, 96]) +@pytest.mark.parametrize("num_heads", [12, 96]) @pytest.mark.parametrize("batch_size", [2, 65]) @pytest.mark.parametrize("max_seq_len", [64, 128], ids=["short-kv", "regular-kv"]) def test_prims_ts_fp8_mla_scales_and_graph_replay( @@ -1310,3 +1301,271 @@ def counted_prims_context(self, *args, **kwargs): assert calls["fallback"] > 0 assert calls["prims_context"] == 0 + + +@pytest.mark.parametrize( + "window_left,paged", + [(-1, False), (-1, True), (31, False)], + ids=["causal-contiguous", "causal-paged", "sliding-window-contiguous"], +) +def test_prims_ts_context_packed_query_tails(window_left: int, paged: bool) -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BatchPrefillPagedTSWrapper, + BatchPrefillTSWrapper, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + num_qo_heads, num_kv_heads, head_dim = 8, 2, 128 + page_size = 32 + # Cover a tile crossing into another request, a partial second Q tile, + # and a final tile that extends beyond the packed allocation. + q_lengths = [3, 130, 1] + kv_lengths = [33, 160, 17] + q_offsets = [0, 3, 133, 134] + kv_offsets = [0, 33, 193, 210] + batch_size = len(q_lengths) + max_pages = math.ceil(max(kv_lengths) / page_size) + query = torch.randn(sum(q_lengths), num_qo_heads, head_dim, device=device, dtype=dtype) + k_cache = torch.randn( + batch_size * max_pages, num_kv_heads, page_size, head_dim, device=device, dtype=dtype + ) + v_cache = torch.randn_like(k_cache) + block_tables = torch.arange(batch_size * max_pages, device=device, dtype=torch.int32).view( + batch_size, max_pages + ) + qo_indptr = torch.tensor(q_offsets, device=device, dtype=torch.int32) + kv_indptr = torch.tensor(kv_offsets, device=device, dtype=torch.int32) + seq_lens = torch.tensor(kv_lengths, device=device, dtype=torch.int32) + keys, values = [], [] + for batch_idx, kv_length in enumerate(kv_lengths): + pages = block_tables[batch_idx].long() + keys.append(k_cache[pages].transpose(1, 2).reshape(-1, num_kv_heads, head_dim)[:kv_length]) + values.append( + v_cache[pages].transpose(1, 2).reshape(-1, num_kv_heads, head_dim)[:kv_length] + ) + output = torch.empty_like(query) + plan_args = { + "device": device, + "batch_size": batch_size, + "max_seq_len_q": max(q_lengths), + "max_kv_len": max(kv_lengths), + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + "q_dtype": dtype, + "k_dtype": dtype, + "v_dtype": dtype, + "out_dtype": dtype, + "mask_type": "causal", + "window_left": window_left, + } + if paged: + wrapper = BatchPrefillPagedTSWrapper(kv_layout="HND") + wrapper.plan(**plan_args, page_size=page_size) + run = functools.partial( + wrapper.run, query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output + ) + else: + wrapper = BatchPrefillTSWrapper() + wrapper.plan(**plan_args, packed=True) + run = functools.partial( + wrapper.run, query, torch.cat(keys), torch.cat(values), qo_indptr, kv_indptr, out=output + ) + + def reference() -> torch.Tensor: + outputs = [] + for batch_idx, (q_length, kv_length) in enumerate(zip(q_lengths, kv_lengths)): + q = query[q_offsets[batch_idx] : q_offsets[batch_idx + 1]].float() + k = keys[batch_idx].float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1) + v = values[batch_idx].float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1) + scores = torch.einsum("qhd,khd->hqk", q, k) / math.sqrt(head_dim) + q_positions = torch.arange(q_length, device=device) + kv_length - q_length + k_positions = torch.arange(kv_length, device=device) + allowed = k_positions[None, :] <= q_positions[:, None] + if window_left > 0: + allowed &= k_positions[None, :] >= q_positions[:, None] - window_left + scores.masked_fill_(~allowed, float("-inf")) + outputs.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), v)) + return torch.cat(outputs).to(dtype) + + run() + torch.testing.assert_close(output, reference(), atol=3e-2, rtol=3e-3) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(validate=False) + query.neg_() + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, reference(), atol=3e-2, rtol=3e-3) + + +def _prims_ts_causal_context_reference( + queries: list[torch.Tensor], keys: list[torch.Tensor], values: list[torch.Tensor] +) -> torch.Tensor: + outputs = [] + for query, key, value in zip(queries, keys, values): + q_length, num_qo_heads, head_dim = query.shape + kv_length, num_kv_heads, _ = key.shape + key = key.float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1) + value = value.float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1) + scores = torch.einsum("qhd,khd->hqk", query.float(), key) / math.sqrt(head_dim) + q_positions = torch.arange(q_length, device=query.device) + kv_length - q_length + k_positions = torch.arange(kv_length, device=query.device) + scores.masked_fill_(k_positions[None, :] > q_positions[:, None], float("-inf")) + outputs.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), value)) + return torch.cat(outputs) + + +def _check_prims_ts_context_replays( + query: torch.Tensor, + output: torch.Tensor, + run: Callable[..., torch.Tensor], + reference: Callable[[], torch.Tensor], +) -> None: + is_fp8 = query.dtype == torch.float8_e4m3fn + atol, rtol, relative_l2_limit = (0.13, 0.05, 0.1) if is_fp8 else (0.03, 0.003, 0.02) + + def check() -> None: + expected = reference() + actual = output.reshape(expected.shape).float() + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + relative_l2 = torch.linalg.vector_norm(actual - expected) / torch.linalg.vector_norm( + expected + ) + assert relative_l2.item() < relative_l2_limit + + run() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(validate=False) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + check() + # copy_ also supports FP8, whose in-place arithmetic is not implemented. + query.copy_(-query.float()) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + check() + + +def _check_prims_ts_compact_query( + paged: bool, + dtype: torch.dtype, + head_dim: int, + total_q: int, + storage_offset_rows: int = 0, +) -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BatchPrefillPagedTSWrapper, + BatchPrefillTSWrapper, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + num_qo_heads, num_kv_heads, page_size = 8, 2, 32 + # Plan with the packed-tail test's bounds. Every Q size then reuses one + # compiled kernel per dtype, head size, and layout (BF16 D128 reuses the + # packed-tail kernels), and the kernel chooses the descriptor path at run time. + max_seq_len_q, max_kv_len = 130, 160 + assert total_q - 4 <= max_seq_len_q and total_q + 28 <= max_kv_len + q_lengths = [3, total_q - 4, 1] + kv_lengths = [33, total_q + 28, 17] + q_offsets = [0, 3, total_q - 1, total_q] + kv_offsets = [0, 33, total_q + 61, total_q + 78] + query_storage = torch.randn( + total_q + storage_offset_rows, + num_qo_heads, + head_dim, + device=device, + dtype=torch.bfloat16, + ).to(dtype) + query = query_storage[storage_offset_rows:] + assert query.is_contiguous() + assert query.storage_offset() == storage_offset_rows * num_qo_heads * head_dim + keys = [ + torch.randn(length, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16).to(dtype) + for length in kv_lengths + ] + values = [ + torch.randn(length, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16).to(dtype) + for length in kv_lengths + ] + output = torch.empty(query.shape, device=device, dtype=torch.bfloat16) + qo_indptr = torch.tensor(q_offsets, device=device, dtype=torch.int32) + plan_args = { + "device": device, + "batch_size": len(q_lengths), + "max_seq_len_q": max_seq_len_q, + "max_kv_len": max_kv_len, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + "q_dtype": dtype, + "k_dtype": dtype, + "v_dtype": dtype, + "out_dtype": torch.bfloat16, + "mask_type": "causal", + } + if paged: + max_pages = math.ceil(max_kv_len / page_size) + cache_shape = (len(q_lengths) * max_pages, num_kv_heads, page_size, head_dim) + k_cache = torch.zeros(cache_shape, device=device, dtype=dtype) + v_cache = torch.zeros_like(k_cache) + for batch_idx, (key, value) in enumerate(zip(keys, values)): + for start in range(0, key.shape[0], page_size): + page = batch_idx * max_pages + start // page_size + count = min(page_size, key.shape[0] - start) + k_cache[page, :, :count].copy_(key[start : start + count].transpose(0, 1)) + v_cache[page, :, :count].copy_(value[start : start + count].transpose(0, 1)) + block_tables = torch.arange( + len(q_lengths) * max_pages, device=device, dtype=torch.int32 + ).view(len(q_lengths), max_pages) + seq_lens = torch.tensor(kv_lengths, device=device, dtype=torch.int32) + wrapper = BatchPrefillPagedTSWrapper(kv_layout="HND") + wrapper.plan(**plan_args, page_size=page_size) + run = functools.partial( + wrapper.run, query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output + ) + else: + kv_indptr = torch.tensor(kv_offsets, device=device, dtype=torch.int32) + wrapper = BatchPrefillTSWrapper() + wrapper.plan(**plan_args, packed=True) + run = functools.partial( + wrapper.run, query, torch.cat(keys), torch.cat(values), qo_indptr, kv_indptr, out=output + ) + + def reference() -> torch.Tensor: + queries = [query[start:end] for start, end in zip(q_offsets[:-1], q_offsets[1:])] + return _prims_ts_causal_context_reference(queries, keys, values) + + _check_prims_ts_context_replays(query, output, run, reference) + + +@pytest.mark.parametrize("delta", [-1, 0, 1], ids=["below", "at", "above"]) +@pytest.mark.parametrize( + "dtype,head_dim,boundary_q,paged", + [ + pytest.param(torch.bfloat16, 128, 64, False, id="bf16-d128-contiguous"), + pytest.param(torch.bfloat16, 128, 64, True, id="bf16-d128-paged"), + pytest.param(torch.bfloat16, 256, 32, True, id="bf16-d256-paged"), + pytest.param(torch.float8_e4m3fn, 128, 128, False, id="fp8-d128-contiguous"), + ], +) +def test_prims_ts_context_compact_query_boundary( + delta: int, dtype: torch.dtype, head_dim: int, boundary_q: int, paged: bool +) -> None: + # These compact Q tensors span just below, at, and above 128 KiB. The groups + # cover both dtypes, head sizes, and KV layouts with three row thresholds. + _check_prims_ts_compact_query(paged, dtype, head_dim, boundary_q + delta) + + +def test_prims_ts_context_compact_query_storage_offset() -> None: + _check_prims_ts_compact_query( + paged=False, dtype=torch.bfloat16, head_dim=128, total_q=64, storage_offset_rows=1 + )