diff --git a/tests/pytorch/attention/test_fused_attn_real_strides.py b/tests/pytorch/attention/test_fused_attn_real_strides.py new file mode 100644 index 0000000000..920463a3e7 --- /dev/null +++ b/tests/pytorch/attention/test_fused_attn_real_strides.py @@ -0,0 +1,321 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for real-stride plumbing (nvte_fused_attn_fwd/bwd_v2). + +For dense (non-THD, non-paged) f16 layouts, the PyTorch extension passes +the real torch strides of Q/K/V (and dO) to the cuDNN fused-attention +graph instead of strides reconstructed from the NVTE_QKV_Layout enum. +Strided views into a packed QKV buffer then compute correctly even when +declared with the plain *separate* layout enum -- the enum no longer needs +to encode memory geometry -- and DotProductAttention no longer needs the +pointer-based (data_ptr/storage_offset) layout detection that graph-breaks +under torch.compile. +""" + +import pytest +import torch + +import transformer_engine.pytorch # noqa: F401 (loads libtransformer_engine.so) +import transformer_engine_torch as tex +from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, +) +import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils +from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + fused_attn_bwd, + fused_attn_fwd, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") + +_BACKEND = tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen +_B, _S, _H, _D = 2, 128, 8, 64 +_DTYPE = torch.bfloat16 + + +def _cu_seqlens(): + return torch.arange(0, (_B + 1) * _S, _S, dtype=torch.int32, device="cuda") + + +def _run(q, k, v, d_o, qkv_layout, fmt): + """One fwd+bwd through the F16 arbitrary-seqlen backend; returns (out, dq, dk, dv).""" + cu = _cu_seqlens() + out, aux = fused_attn_fwd( + True, _S, _S, cu, cu, q, k, v, _DTYPE, _BACKEND, + dropout=0.0, qkv_layout=qkv_layout, o_format=fmt, + attn_bias_type="no_bias", attn_mask_type="no_mask", + ) + dq, dk, dv, _, _ = fused_attn_bwd( + _S, _S, cu, cu, q, k, v, out, d_o, _DTYPE, aux, _BACKEND, + dropout=0.0, qkv_layout=qkv_layout, o_format=fmt, do_format=fmt, + dqkv_layout=qkv_layout, attn_bias_type="no_bias", attn_mask_type="no_mask", + deterministic=True, + ) + return out, dq, dk, dv + + +def _assert_bit_exact(result, reference): + for name, x, y in zip(("out", "dq", "dk", "dv"), result, reference): + assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" + + +def _backend_supported(): + try: + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + _run(q, q.clone(), q.clone(), q.clone(), "bshd_bshd_bshd", "bshd") + return True + except Exception: + return False + + +requires_backend = pytest.mark.skipif( + not (torch.cuda.is_available() and _backend_supported()), + reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", +) + + +@requires_backend +def test_packed_bs3hd_views_declared_separate(): + """Headline: strided views into a packed [b,s,3,h,d] buffer, declared with the + plain separate layout enum, are bit-exact vs the contiguous baseline because the + real strides are passed to the cuDNN graph.""" + torch.manual_seed(0) + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + d_o = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + + reference = _run( + qkv[:, :, 0].contiguous(), qkv[:, :, 1].contiguous(), qkv[:, :, 2].contiguous(), + d_o, "bshd_bshd_bshd", "bshd", + ) + + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + assert not q.is_contiguous() + result = _run(q, k, v, d_o, "bshd_bshd_bshd", "bshd") + _assert_bit_exact(result, reference) + + +@requires_backend +def test_packed_sbh3d_views_declared_separate(): + """Same as above for the sbh3d interleave (packing at dim -2, sbhd format).""" + torch.manual_seed(0) + qkv = torch.randn(_S, _B, _H, 3, _D, dtype=_DTYPE, device="cuda") + d_o = torch.randn(_S, _B, _H, _D, dtype=_DTYPE, device="cuda") + + reference = _run( + qkv[:, :, :, 0].contiguous(), qkv[:, :, :, 1].contiguous(), + qkv[:, :, :, 2].contiguous(), d_o, "sbhd_sbhd_sbhd", "sbhd", + ) + + q, k, v = qkv[:, :, :, 0], qkv[:, :, :, 1], qkv[:, :, :, 2] + assert not q.is_contiguous() + result = _run(q, k, v, d_o, "sbhd_sbhd_sbhd", "sbhd") + _assert_bit_exact(result, reference) + + +@requires_backend +def test_packed_views_with_packed_enum_unchanged(): + """Regression: the historical path (packed views with the matching packed enum) + still computes the same values as the separate-contiguous baseline. For packed + views the real strides coincide with the enum-derived ones, so passing them is + a no-op numerically.""" + torch.manual_seed(0) + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + d_o = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + + reference = _run( + qkv[:, :, 0].contiguous(), qkv[:, :, 1].contiguous(), qkv[:, :, 2].contiguous(), + d_o, "bshd_bshd_bshd", "bshd", + ) + result = _run(q, k, v, d_o, "bs3hd", "bshd") + _assert_bit_exact(result, reference) + + +# --------------------------------------------------------------------------- +# DotProductAttention end-to-end: DPA skips pointer-based qkv layout detection +# (data_ptr/storage_offset games) for dense f16 layouts and declares the +# format-derived separate layout instead. +# --------------------------------------------------------------------------- + + +def _force_backend(monkeypatch, backend): + """Force a single attention backend via env and invalidate the selection cache.""" + flash, fused = {"flash": ("1", "0"), "fused": ("0", "1")}[backend] + monkeypatch.setenv("NVTE_FLASH_ATTN", flash) + monkeypatch.setenv("NVTE_FUSED_ATTN", fused) + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + if backend == "flash": + # flash-attn bwd uses atomics unless deterministic; the fused (cuDNN) + # backend gets disabled outright by the deterministic flag on some + # devices, and its bwd is run-to-run deterministic anyway. + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + dpa_module._attention_backends["backend_selection_requires_update"] = True + + +def _make_dpa(qkv_format): + return DotProductAttention( + _H, _D, attention_dropout=0.0, qkv_format=qkv_format, attn_mask_type="no_mask" + ) + + +def _dpa_separate(qkv_format): + """Fwd+bwd on three contiguous leaves; returns (out, dq, dk, dv).""" + torch.manual_seed(0) + shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + q, k, v = [ + torch.randn(*shape, dtype=_DTYPE, device="cuda", requires_grad=True) for _ in range(3) + ] + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(q, k, v) + out.backward(torch.ones_like(out)) + return out, q.grad, k.grad, v.grad + + +def _dpa_packed_views(qkv_format): + """Fwd+bwd on strided views into one packed leaf; returns (out, dq, dk, dv).""" + torch.manual_seed(0) + shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + parts = [torch.randn(*shape, dtype=_DTYPE, device="cuda") for _ in range(3)] + qkv = torch.stack(parts, dim=2).requires_grad_() # bs3hd / sb3hd packing + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + assert not q.is_contiguous() + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(q, k, v) + out.backward(torch.ones_like(out)) + return out, qkv.grad[:, :, 0], qkv.grad[:, :, 1], qkv.grad[:, :, 2] + + +@requires_backend +@pytest.mark.parametrize("fmt", ["bshd", "sbhd"]) +def test_dpa_fused_packed_views(monkeypatch, fmt): + """Fused backend: packed stack views (declared separate, real strides) are + bit-exact vs contiguous separate q/k/v, fwd and grads.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate(fmt) + result = _dpa_packed_views(fmt) + _assert_bit_exact(result, reference) + + +def test_dpa_flash_packed_views(monkeypatch): + """Flash backend smoke: flash consumes real strides natively, so packed views + declared separate are bit-exact vs separate.""" + _force_backend(monkeypatch, "flash") + try: + reference = _dpa_separate("bshd") + except Exception as exc: + pytest.skip(f"flash attention backend not available: {exc}") + result = _dpa_packed_views("bshd") + _assert_bit_exact(result, reference) + + +@requires_backend +def test_dpa_layout_detection_skipped_dense_kept_thd(monkeypatch): + """get_qkv_layout is not called for dense bshd, but still runs for thd + (ragged layouts ignore strides in C++, so detection stays).""" + _force_backend(monkeypatch, "fused") + + calls = [] + orig = dpa_utils.get_qkv_layout + + def counting(*args, **kwargs): + calls.append(kwargs.get("qkv_format")) + return orig(*args, **kwargs) + + monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) + + _dpa_separate("bshd") + assert not calls, "get_qkv_layout should be skipped for dense bshd" + + # thd: full sequences, padding mask + torch.manual_seed(0) + t = _B * _S + q, k, v = [torch.randn(t, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] + cu = _cu_seqlens() + dpa_module._attention_backends["backend_selection_requires_update"] = True + dpa = DotProductAttention( + _H, _D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" + ) + try: + dpa(q, k, v, cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=_S, max_seqlen_kv=_S) + except ValueError: + # No thd-capable backend on this device; the layout step (the subject + # of this test) runs before backend dispatch, so the assertion below + # still holds. + pass + assert calls == ["thd"], "get_qkv_layout must still run for thd" + + +@requires_backend +def test_dpa_noncontiguous_head_dim_normalized(monkeypatch): + """Inputs with stride(-1) != 1 are normalized with .contiguous() in the bypass + path (mirrors the old detection's contiguous-retry) and stay bit-exact.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate("bshd") + + torch.manual_seed(0) + q, k, v = [ + torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) + for _ in range(3) + ] + # transpose(-1, -2) of a transposed copy: same values, stride(-1) != 1 + q_t = q.detach().transpose(2, 3).contiguous().transpose(2, 3).requires_grad_() + k_t = k.detach().transpose(2, 3).contiguous().transpose(2, 3).requires_grad_() + v_t = v.detach().transpose(2, 3).contiguous().transpose(2, 3).requires_grad_() + assert q_t.stride(-1) != 1 + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd")(q_t, k_t, v_t) + out.backward(torch.ones_like(out)) + _assert_bit_exact((out, q_t.grad, k_t.grad, v_t.grad), reference) + + +def _compiled_graph_breaks(): + """Compile DPA (after an eager warm-up that caches backend selection) and + return dynamo's graph_break counters for one fwd+bwd.""" + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + torch.manual_seed(0) + q, k, v = [ + torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) + for _ in range(3) + ] + dpa = _make_dpa("bshd") + dpa_module._attention_backends["backend_selection_requires_update"] = True + dpa(q, k, v) # eager warm-up: backend selection happens outside dynamo + out = torch.compile(dpa)(q, k, v) + out.backward(torch.ones_like(out)) + breaks = dict(torch._dynamo.utils.counters["graph_break"]) + torch._dynamo.reset() + return breaks + + +def _pointer_breaks(breaks): + return { + reason: count + for reason, count in breaks.items() + if "data_ptr" in reason or "UntypedStorage" in reason + } + + +@requires_backend +def test_dpa_torch_compile_no_data_ptr_graph_breaks(monkeypatch): + """Compiling DPA no longer graph-breaks on data_ptr/UntypedStorage: the + pointer-based layout detection is bypassed for dense layouts. Negative + control: forcing the detection back on reintroduces those breaks.""" + _force_backend(monkeypatch, "fused") + + # negative control: force pointer-based detection, expect data_ptr breaks + monkeypatch.setattr(dpa_module, "_skip_pointer_layout_detection", lambda *a: False) + baseline = _compiled_graph_breaks() + assert _pointer_breaks(baseline), ( + f"expected data_ptr/UntypedStorage breaks with detection forced on: {baseline}" + ) + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + breaks = _compiled_graph_breaks() + assert not _pointer_breaks(breaks), ( + f"data_ptr/UntypedStorage graph breaks should be gone: {_pointer_breaks(breaks)}" + ) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fc21771297..b1ad326d62 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -529,22 +529,20 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return backend; } -// NVTE fused attention FWD with separate Q, K and V -void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - const NVTETensor page_table_v, const NVTETensor rng_state, - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, - bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_flash_attn_fwd); +// NVTE fused attention FWD with separate Q, K and V and explicit strides +void nvte_fused_attn_fwd_v2( + const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, + const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, + const NVTETensor page_table_k, const NVTETensor page_table_v, const NVTETensor rng_state, + size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, bool return_max_logit, + bool cuda_graph, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, + NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_attn_fwd_v2); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); @@ -622,10 +620,11 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, is_training, return_max_logit, attn_scale, dropout, qkv_layout, o_format, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, input_Q, input_K, - input_V, input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, - input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle); + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_strides, + input_Q, input_K, input_V, input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, + input_cu_seqlens_q, input_cu_seqlens_kv, input_cu_seqlens_q_padded, + input_cu_seqlens_kv_padded, input_page_table_k, input_page_table_v, input_rng_state, + wkspace, stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, is_training, attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, @@ -637,23 +636,47 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } } -// NVTE fused attention BWD with separate Q, K and V -void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, - const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, - NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, + +// NVTE fused attention FWD with separate Q, K and V +void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, + const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, + NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, - size_t max_seqlen_kv, float attn_scale, float dropout, + const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, + const NVTETensor page_table_v, const NVTETensor rng_state, + size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, + bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, - bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_flash_attn_bwd); + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, + NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_flash_attn_fwd); + nvte_fused_attn_fwd_v2(Q, K, V, Bias, SoftmaxOffset, S, O, Aux_CTX_Tensors, cu_seqlens_q, + cu_seqlens_kv, cu_seqlens_q_padded, cu_seqlens_kv_padded, page_table_k, + page_table_v, rng_state, max_seqlen_q, max_seqlen_kv, is_training, + return_max_logit, cuda_graph, attn_scale, dropout, qkv_layout, o_format, + qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, bottom_right_diagonal, + NVTEQKVStrides{nullptr, nullptr, nullptr, nullptr}, workspace, stream); +} + +// NVTE fused attention BWD with separate Q, K and V and explicit strides +void nvte_fused_attn_bwd_v2( + const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor O, + const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, + NVTETensor dQ, NVTETensor dK, NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, + size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, + NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, bool cuda_graph, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_attn_bwd_v2); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); @@ -712,11 +735,11 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso fused_attn_arbitrary_seqlen_bwd( b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale, dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, input_Q, input_K, - input_V, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, output_dQ, output_dK, - output_dV, output_dBias, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, - input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, wkspace, stream, - handle); + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_strides, + input_Q, input_K, input_V, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, + output_dQ, output_dK, output_dV, output_dBias, output_dSoftmaxOffset, input_cu_seqlens_q, + input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, + wkspace, stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { size_t i = 0; const Tensor *input_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); @@ -742,6 +765,32 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso } } +// NVTE fused attention BWD with separate Q, K and V +void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, + const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, + const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, + NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, + const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, + size_t max_seqlen_kv, float attn_scale, float dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, + bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_flash_attn_bwd); + nvte_fused_attn_bwd_v2(Q, K, V, O, dO, S, dP, Aux_CTX_Tensors, dQ, dK, dV, dBias, dSoftmaxOffset, + cu_seqlens_q, cu_seqlens_kv, cu_seqlens_q_padded, cu_seqlens_kv_padded, + max_seqlen_q, max_seqlen_kv, attn_scale, dropout, qkv_layout, o_format, + do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, + bias_type, attn_mask_type, softmax_type, window_size_left, + window_size_right, bottom_right_diagonal, deterministic, cuda_graph, + NVTEQKVStrides{nullptr, nullptr, nullptr, nullptr}, workspace, stream); +} + uint32_t nvte_get_runtime_num_segments(NVTETensor cu_seqlen, NVTETensor workspace, size_t len, cudaStream_t stream) { NVTE_API_CALL(nvte_get_runtime_num_segments); diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index 6df7ad35c8..9d542bc022 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -55,12 +55,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( bool is_training, bool return_max_logit, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, void *devPtrQ, void *devPtrK, - void *devPtrV, void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, - void *devPtrO, void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, - void *devPtrCuSeqlensKV, void *devPtrPageTableK, void *devPtrPageTableV, - void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, - void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + int64_t window_size_right, bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, + void *devPtrQ, void *devPtrK, void *devPtrV, void *devPtrBias, void *devPtrSoftmaxOffset, + void *devPtrS1, void *devPtrS2, void *devPtrO, void *devPtrDropoutSeed, + void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, + void *devPtrPageTableK, void *devPtrPageTableV, void *devPtrSeqOffsetsQ, + void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, void *workspace, + size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); @@ -94,6 +95,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } + // Real strides provided by the caller replace the enum-derived Q/K/V strides for + // dense (non-THD, non-paged) layouts. + const bool has_real_qkv_strides = + qkv_strides.q != nullptr && qkv_strides.k != nullptr && qkv_strides.v != nullptr; + const bool use_real_strides = + has_real_qkv_strides && !is_paged_kv && !is_ragged_q && !is_ragged_kv; + // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { @@ -153,6 +161,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cudnn_frontend::DataType_t::NOT_SET, return_max_logit, }; + if (use_real_strides) { + for (int i = 0; i < 4; ++i) { + descriptor.real_strides[i] = qkv_strides.q[i]; + descriptor.real_strides[4 + i] = qkv_strides.k[i]; + descriptor.real_strides[8 + i] = qkv_strides.v[i]; + } + } namespace fe = cudnn_frontend; using graph_and_tensors = @@ -219,6 +234,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); } + if (use_real_strides) { + // use the real tensor strides instead of the enum-derived ones + q_stride.assign(qkv_strides.q, qkv_strides.q + 4); + k_stride.assign(qkv_strides.k, qkv_strides.k + 4); + v_stride.assign(qkv_strides.v, qkv_strides.v + 4); + } Q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Q") @@ -558,10 +579,10 @@ void fused_attn_arbitrary_seqlen_bwd_impl( NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, void *devPtrQ, void *devPtrKTranspose, - void *devPtrVTranspose, void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, - void *devPtrSoftmaxOffset, void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, - void *devPtrdBias, void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, + bool bottom_right_diagonal, bool deterministic, NVTEQKVStrides qkv_strides, void *devPtrQ, + void *devPtrKTranspose, void *devPtrVTranspose, void *devPtrO, void *devPtrSoftmaxStats, + void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrdQ, void *devPtrdK, void *devPtrdV, + void *devPtrdO, void *devPtrdBias, void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { @@ -598,6 +619,14 @@ void fused_attn_arbitrary_seqlen_bwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } + // Real strides provided by the caller replace the enum-derived strides of the Q/K/V and + // dO inputs for dense (non-THD, non-paged) layouts. Outputs dQ/dK/dV keep the + // enum-derived strides (TE allocates them). + const bool is_dense = !is_paged_kv && !is_ragged_q && !is_ragged_kv; + const bool use_real_strides = + qkv_strides.q != nullptr && qkv_strides.k != nullptr && qkv_strides.v != nullptr && is_dense; + const bool use_real_do_strides = qkv_strides.d_o != nullptr && is_dense; + // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { @@ -656,6 +685,18 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cudnn_frontend::DataType_t::NOT_SET, false, }; + if (use_real_strides) { + for (int i = 0; i < 4; ++i) { + descriptor.real_strides[i] = qkv_strides.q[i]; + descriptor.real_strides[4 + i] = qkv_strides.k[i]; + descriptor.real_strides[8 + i] = qkv_strides.v[i]; + } + } + if (use_real_do_strides) { + for (int i = 0; i < 4; ++i) { + descriptor.real_strides[12 + i] = qkv_strides.d_o[i]; + } + } namespace fe = cudnn_frontend; using graph_and_tensors = @@ -722,6 +763,21 @@ void fused_attn_arbitrary_seqlen_bwd_impl( generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_O_Matrix); + // dQ/dK/dV outputs and O keep the enum-derived strides; the real-stride override + // only applies to the Q/K/V and dO inputs. + std::vector dq_stride(q_stride); + std::vector dk_stride(k_stride); + std::vector dv_stride(v_stride); + std::vector do_stride(o_stride); + if (use_real_strides) { + q_stride.assign(qkv_strides.q, qkv_strides.q + 4); + k_stride.assign(qkv_strides.k, qkv_strides.k + 4); + v_stride.assign(qkv_strides.v, qkv_strides.v + 4); + } + if (use_real_do_strides) { + do_stride.assign(qkv_strides.d_o, qkv_strides.d_o + 4); + } + q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Q") .set_dim({b, h, s_q, d_qk}) @@ -741,7 +797,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( dO = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("dO") .set_dim({b, h, s_q, d_v}) - .set_stride(o_stride)); + .set_stride(do_stride)); if (is_ragged_q) { offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("offset_q") @@ -893,9 +949,9 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto [dQ, dK, dV] = mha_graph->sdpa_backward(q, k, v, o, dO, stats, sdpa_backward_options); - dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(q_stride); - dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(k_stride); - dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(v_stride); + dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(dq_stride); + dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(dk_stride); + dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(dv_stride); if (is_ragged_q) { dQ->set_ragged_offset(offset_q); } @@ -1079,12 +1135,12 @@ void fused_attn_arbitrary_seqlen_fwd( bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, - const Tensor *page_table_k, const Tensor *page_table_v, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { + bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, const Tensor *input_Bias, + const Tensor *input_SoftmaxOffset, Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, const Tensor *page_table_v, + const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto QKV_type = input_Q->data.dtype; @@ -1214,11 +1270,11 @@ void fused_attn_arbitrary_seqlen_fwd( max_batch_size, max_tokens_q, max_tokens_kv, num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h, bias_sq, bias_skv, is_training, return_max_logit, attn_scale, p_dropout, qkv_layout, o_format, bias_type, - mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, - devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, - devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, - devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), - workspace->data.dptr, &workspace_size, stream, handle); + mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + qkv_strides, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, + devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, + devPtrPageTableK, devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, + get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1242,8 +1298,8 @@ void fused_attn_arbitrary_seqlen_bwd( NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, + bool deterministic, NVTEQKVStrides qkv_strides, const Tensor *input_Q, const Tensor *input_K, + const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, @@ -1313,11 +1369,12 @@ void fused_attn_arbitrary_seqlen_bwd( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, max_batch_size, max_tokens_q, max_tokens_kv, bias_b, bias_h, bias_sq, bias_skv, attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, - devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, - devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, devPtrDropoutOffset, - devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_strides, + devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, + devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, + devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, + devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, + stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h index 8f79b5bb4a..2a9414d5c2 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h @@ -25,12 +25,12 @@ void fused_attn_arbitrary_seqlen_fwd( bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, - const Tensor *page_table_k, const Tensor *page_table_v, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, const Tensor *input_Bias, + const Tensor *input_SoftmaxOffset, Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, const Tensor *page_table_v, + const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); void fused_attn_arbitrary_seqlen_bwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, @@ -39,8 +39,8 @@ void fused_attn_arbitrary_seqlen_bwd( NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, + bool deterministic, NVTEQKVStrides qkv_strides, const Tensor *input_Q, const Tensor *input_K, + const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 41656062a4..3719616f41 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -312,6 +313,10 @@ struct FADescriptor_v1 { cudnn_frontend::DataType_t do_tensor_type; cudnn_frontend::DataType_t dqkv_tensor_type; bool return_max_logit; + // Real (torch) strides for Q, K, V, dO in cuDNN dim order [b, h, s, d], concatenated + // ([0:4]=Q, [4:8]=K, [8:12]=V, [12:16]=dO). All zeros when the enum-derived strides + // are used. Part of the cache key so plans built for different strides are not reused. + std::array real_strides = {}; bool operator<(const FADescriptor_v1 &rhs) const { return std::tie(b, h, hg, s_q, s_kv, d_qk, d_v, num_pages_k, num_pages_v, page_size_k, @@ -320,7 +325,7 @@ struct FADescriptor_v1 { do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, bias_type, qkv_tensor_type, o_tensor_type, do_tensor_type, - dqkv_tensor_type, return_max_logit) < + dqkv_tensor_type, return_max_logit, real_strides) < std::tie(rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d_qk, rhs.d_v, rhs.num_pages_k, rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k, rhs.max_pages_per_seq_v, rhs.bias_b, rhs.bias_h, rhs.bias_sq, rhs.bias_skv, @@ -329,7 +334,7 @@ struct FADescriptor_v1 { rhs.do_scale_inv_format, rhs.mask_type, rhs.softmax_type, rhs.window_size_left, rhs.window_size_right, rhs.bottom_right_diagonal, rhs.deterministic, rhs.bias_type, rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, - rhs.dqkv_tensor_type, rhs.return_max_logit); + rhs.dqkv_tensor_type, rhs.return_max_logit, rhs.real_strides); } }; diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 41e4b136bd..6f3acb4984 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -388,6 +388,85 @@ void nvte_populate_rng_state_async(NVTETensor rng_state_dst, const NVTETensor se size_t q_max_seqlen, size_t kv_max_seqlen, NVTE_Fused_Attn_Backend backend, cudaStream_t stream); +/*! \struct NVTEQKVStrides + * \brief Real memory strides of the Q, K, V (and dO in backward) tensors. + * + * \warning This API is **experimental** and subject to change. + * + * Each member points to 4 int64 strides in cuDNN dimension order [b, h, s, d], in units + * of elements, or is NULL, in which case the strides of the corresponding tensor are + * derived from the NVTE_QKV_Layout enum (the historical behavior). `q`, `k` and `v` must + * either all be provided or all be NULL. `d_o` is only consulted by the backward pass. + * + * Only the F16 arbitrary-seqlen backend consumes these strides, and only for dense + * (non-THD, non-paged) layouts; the FP8 and max-512 backends, as well as THD/paged + * layouts, ignore them and always use the enum-derived strides. Output tensors + * (O, dQ/dK/dV) always keep the enum-derived strides since Transformer Engine + * allocates them. + */ +typedef struct NVTEQKVStrides { + const int64_t *q; /*!< Q strides [b, h, s, d] or NULL */ + const int64_t *k; /*!< K strides [b, h, s, d] or NULL */ + const int64_t *v; /*!< V strides [b, h, s, d] or NULL */ + const int64_t *d_o; /*!< dO strides [b, h, s, d] or NULL (backward only) */ +} NVTEQKVStrides; + +/*! \brief Compute dot product attention with separate Q, K and V, with explicit strides. + * + * \warning This API is **experimental** and subject to change. + * + * Identical to nvte_fused_attn_fwd(), with an additional `qkv_strides` parameter that + * provides the real memory strides of Q/K/V to the cuDNN graph instead of strides + * reconstructed from `qkv_layout` (see NVTEQKVStrides for the exact semantics). + * nvte_fused_attn_fwd() is equivalent to calling this function with all-NULL strides. + * + * \param[in] qkv_strides Real strides of Q/K/V; NULL members fall back + * to the enum-derived strides. + * + * All other parameters are as in nvte_fused_attn_fwd(). + */ +void nvte_fused_attn_fwd_v2( + const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, + const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, + const NVTETensor page_table_k, const NVTETensor page_table_v, const NVTETensor rng_state, + size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, bool return_max_logit, + bool cuda_graph, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, + NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream); + +/*! \brief Compute the backward of the dot product attention with separate Q, K and V, + * with explicit strides. + * + * \warning This API is **experimental** and subject to change. + * + * Identical to nvte_fused_attn_bwd(), with an additional `qkv_strides` parameter that + * provides the real memory strides of Q/K/V and dO to the cuDNN graph instead of strides + * reconstructed from `qkv_layout` (see NVTEQKVStrides for the exact semantics). + * nvte_fused_attn_bwd() is equivalent to calling this function with all-NULL strides. + * + * \param[in] qkv_strides Real strides of Q/K/V and dO; NULL members fall + * back to the enum-derived strides. + * + * All other parameters are as in nvte_fused_attn_bwd(). + */ +void nvte_fused_attn_bwd_v2( + const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor O, + const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, + NVTETensor dQ, NVTETensor dK, NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, + size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, + NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, bool cuda_graph, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream); + /*! \brief Get KV format for a given QKV layout. * * \warning This API is **experimental** and subject to change. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 03008bb2d7..7ae2610947 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -190,6 +190,33 @@ def _pad_qkv_head_dim(query_layer, key_layer, value_layer): return query_layer, key_layer, value_layer, orig_head_dim_qk, orig_head_dim_v +def _skip_pointer_layout_detection(qkv_format, fp8_dpa, inference_params) -> bool: + """Decide whether pointer-based qkv layout detection can be skipped. + + ``get_qkv_layout`` inspects ``untyped_storage().data_ptr()`` and + ``storage_offset()`` of q/k/v on every forward to detect packed layouts + (e.g. ``bs3hd``). This is torch.compile-hostile (data_ptr access causes + graph breaks) and adds per-step CPU overhead. For dense f16 layouts the + detection is redundant for memory correctness: the cuDNN fused-attention + backend receives the live tensor strides via the v2 C API, + flash-attention consumes real strides natively, and the unfused backend + uses plain torch ops. The format-derived separate layout string is then + always safe to declare. + + Excluded cases (detection must still run): + - ``thd``: the C++ side ignores strides for ragged layouts, so a packed + ``t3hd``/``th3d`` input declared as ``thd_thd_thd`` would silently + compute on wrong memory. + - FP8 DPA: the FP8 fused backend does not consume the real-strides + parameter, and ``combine_and_quantize`` relies on detected packedness. + - KV caching (``inference_params``): layouts may need the ``paged_kv_`` + prefix and mixed q/kv formats that only detection derives. + """ + if inference_params is not None or fp8_dpa: + return False + return qkv_format in ("bshd", "sbhd") + + def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim_v): """Trim FlashAttention output after padding V to a larger head dimension.""" out_shape = attn_out.shape[:-1] @@ -1449,6 +1476,24 @@ def forward( qkv_format=qkv_format, inference_params=inference_params, ) + elif _skip_pointer_layout_detection( + qkv_format, + self.fp8 and self.fp8_meta["recipe"].fp8_dpa, + inference_params, + ): + # The backends consume the live tensor strides for dense + # layouts, so the layout enum no longer needs to encode the + # memory geometry: declare the format-derived separate layout + # and pass q/k/v through as-is. Only the last (head) dimension + # must be packed; normalize it if needed (stride(-1) is plain + # tensor metadata, so this stays torch.compile-traceable). + query_layer, key_layer, value_layer = [ + x if x.stride(-1) == 1 else x.contiguous() + for x in [query_layer, key_layer, value_layer] + ] + qkv_layout = f"{qkv_format}_{qkv_format}_{qkv_format}" + q_format = qkv_format + kv_format = qkv_format else: ( qkv_layout, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index eb8813d4a0..310eb8a8b6 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -12,6 +12,46 @@ namespace { constexpr int block_size = 512; +// Cast a py::handle to at::Tensor if it wraps a plain torch tensor. +bool get_plain_torch_tensor(pybind11::handle handle, at::Tensor &out) { + try { + out = handle.cast(); + return true; + } catch (const pybind11::cast_error &) { + return false; + } +} + +// Map a 4-D torch tensor's strides to cuDNN dim order [b, h, s, d] given its format. +// Returns false for unsupported formats (e.g. THD) or non-4D tensors. +bool to_bhsd_strides(const at::Tensor &t, NVTE_QKV_Format format, int64_t *out) { + if (t.dim() != 4) { + return false; + } + switch (format) { + case NVTE_QKV_Format::NVTE_BSHD: // torch [b, s, h, d] + out[0] = t.stride(0); + out[1] = t.stride(2); + out[2] = t.stride(1); + out[3] = t.stride(3); + return true; + case NVTE_QKV_Format::NVTE_SBHD: // torch [s, b, h, d] + out[0] = t.stride(1); + out[1] = t.stride(2); + out[2] = t.stride(0); + out[3] = t.stride(3); + return true; + case NVTE_QKV_Format::NVTE_BHSD: // torch [b, h, s, d] + out[0] = t.stride(0); + out[1] = t.stride(1); + out[2] = t.stride(2); + out[3] = t.stride(3); + return true; + default: + return false; + } +} + // fast zero-fills of tensors void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &start_index) { std::vector shape = transformer_engine::pytorch::convertShape(self.shape()); @@ -243,16 +283,34 @@ std::vector fused_attn_fwd( // create workspace TensorWrapper workspace; + // Real torch strides of Q/K/V for the cuDNN graph; NULL members fall back to + // the enum-derived strides (quantized tensors, THD, non-4D tensors). + int64_t q_strides[4], k_strides[4], v_strides[4]; + NVTEQKVStrides qkv_strides{nullptr, nullptr, nullptr, nullptr}; + { + at::Tensor q_torch, k_torch, v_torch; + if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && + get_plain_torch_tensor(V, v_torch)) { + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + if (to_bhsd_strides(q_torch, q_format, q_strides) && + to_bhsd_strides(k_torch, kv_format, k_strides) && + to_bhsd_strides(v_torch, kv_format, v_strides)) { + qkv_strides = NVTEQKVStrides{q_strides, k_strides, v_strides, nullptr}; + } + } + } + // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_fwd( + nvte_fused_attn_fwd_v2( te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_SoftmaxOffset.data(), te_S.data(), te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); + window_size[1], bottom_right_diagonal, qkv_strides, workspace.data(), + at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace and auxiliary output tensors @@ -302,14 +360,15 @@ std::vector fused_attn_fwd( // execute the kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_fwd( + nvte_fused_attn_fwd_v2( te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_SoftmaxOffset.data(), te_S.data(), te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); + window_size[1], bottom_right_diagonal, qkv_strides, workspace.data(), + at::cuda::getCurrentCUDAStream()); }); // destroy tensor wrappers, but not allocated memory @@ -571,17 +630,36 @@ std::vector fused_attn_bwd( // create workspace TensorWrapper workspace; + // Real torch strides of Q/K/V and dO for the cuDNN graph; NULL members fall + // back to the enum-derived strides (quantized tensors, THD, non-4D tensors). + int64_t q_strides[4], k_strides[4], v_strides[4], do_strides[4]; + NVTEQKVStrides qkv_strides{nullptr, nullptr, nullptr, nullptr}; + { + at::Tensor q_torch, k_torch, v_torch, do_torch; + if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && + get_plain_torch_tensor(V, v_torch)) { + if (to_bhsd_strides(q_torch, q_format, q_strides) && + to_bhsd_strides(k_torch, kv_format, k_strides) && + to_bhsd_strides(v_torch, kv_format, v_strides)) { + const bool have_do_strides = get_plain_torch_tensor(dO, do_torch) && + to_bhsd_strides(do_torch, do_format, do_strides); + qkv_strides = + NVTEQKVStrides{q_strides, k_strides, v_strides, have_do_strides ? do_strides : nullptr}; + } + } + } + // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( + nvte_fused_attn_bwd_v2( te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), - at::cuda::getCurrentCUDAStream()); + window_size[1], bottom_right_diagonal, deterministic, cuda_graph, qkv_strides, + workspace.data(), at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace @@ -591,15 +669,15 @@ std::vector fused_attn_bwd( // execute kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( + nvte_fused_attn_bwd_v2( te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), - at::cuda::getCurrentCUDAStream()); + window_size[1], bottom_right_diagonal, deterministic, cuda_graph, qkv_strides, + workspace.data(), at::cuda::getCurrentCUDAStream()); }); // destroy tensor wrappers