From f778260c6aa03aa8ce062e3a4d025e01085af7ac Mon Sep 17 00:00:00 2001 From: "ryan.u" Date: Fri, 11 Sep 2026 21:12:03 +0900 Subject: [PATCH 1/3] [PyTorch] Add head-parallel FA4 backward for cuDNN CP attention Signed-off-by: ryan.u --- benchmarks/benchmark_mixed_cp.py | 143 ++++++++++ docs/envvars.rst | 6 + qa/L3_pytorch_FA_versions_test/test.sh | 4 + tests/pytorch/attention/run_mixed_cp.py | 264 ++++++++++++++++++ tests/pytorch/attention/test_mixed_cp.py | 51 ++++ .../common/triton/cp_packing.py | 155 ++++++++++ .../dot_product_attention/context_parallel.py | 57 ++++ .../dot_product_attention/mixed_cp.py | 173 ++++++++++++ 8 files changed, 853 insertions(+) create mode 100644 benchmarks/benchmark_mixed_cp.py create mode 100644 tests/pytorch/attention/run_mixed_cp.py create mode 100644 tests/pytorch/attention/test_mixed_cp.py create mode 100644 transformer_engine/common/triton/cp_packing.py create mode 100644 transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py diff --git a/benchmarks/benchmark_mixed_cp.py b/benchmarks/benchmark_mixed_cp.py new file mode 100644 index 0000000000..d2ff9a65c1 --- /dev/null +++ b/benchmarks/benchmark_mixed_cp.py @@ -0,0 +1,143 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Compare native and mixed CP attention graph replay on one GPU group.""" + +import argparse +import json +import logging +import os +from pathlib import Path +import statistics +import time + +import torch +import torch.distributed as dist + +from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention.mixed_cp import _is_supported + + +def main(): + """Measure both variants in ABBA order with the same tensors and process group.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--sequence", type=int, default=32768) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl", device_id=torch.device("cuda", torch.cuda.current_device())) + try: + cp_size = dist.get_world_size() + if args.sequence % (2 * cp_size): + raise ValueError("Sequence length must be divisible by twice the CP size") + sequence = args.sequence // cp_size + torch.manual_seed(2026 + dist.get_rank()) + inputs = [ + torch.randn( + sequence, args.batch, 32, width, device="cuda", dtype=torch.bfloat16 + ).requires_grad_() + for width in (192, 192, 128) + ] + gradient = torch.randn(sequence, args.batch, 32 * 128, device="cuda", dtype=torch.bfloat16) + if not _is_supported(*inputs, cp_size): + raise ValueError("This configuration cannot exercise mixed CP backward") + module = ( + DotProductAttention( + num_attention_heads=32, + kv_channels=(192, 128), + attention_dropout=0, + qkv_format="sbhd", + attn_mask_type="causal", + softmax_scale=0.083, + cp_group=dist.group.WORLD, + cp_global_ranks=list(range(cp_size)), + cp_stream=torch.cuda.Stream(), + cp_comm_type="p2p", + ) + .cuda() + .train() + ) + records = [] + reference = None + for variant in ("native", "mixed", "mixed", "native"): + os.environ["NVTE_FUSED_ATTN_CP_USE_FAv4_BWD"] = "1" if variant == "mixed" else "0" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + for tensor in inputs: + tensor.grad = None + output = module(*inputs) + output.backward(gradient) + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + values = (output.detach().clone(), [t.grad.detach().clone() for t in inputs]) + if reference is None: + reference = values + else: + torch.testing.assert_close(values[0], reference[0], rtol=0, atol=0) + for value, expected in zip(values[1], reference[1], strict=True): + torch.testing.assert_close(value, expected, rtol=0.04, atol=0.025) + assert ( + value.float() - expected.float() + ).norm() / expected.float().norm() < 0.008 + del values, output + for tensor in inputs: + tensor.grad = None + dist.barrier() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = module(*inputs) + output.backward(gradient) + for _ in range(10): + graph.replay() + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - before + dist.barrier() + samples = [] + for _ in range(50): + start = time.perf_counter() + graph.replay() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1000) + ranks = [None] * cp_size + dist.all_gather_object( + ranks, dict(samples_ms=samples, capture_incremental_peak_bytes=peak) + ) + maxima = [max(rank["samples_ms"][i] for rank in ranks) for i in range(50)] + records.append(dict(variant=variant, median_ms=statistics.median(maxima), ranks=ranks)) + graph.reset() + del graph, output + for tensor in inputs: + tensor.grad = None + torch.cuda.synchronize() + if dist.get_rank() == 0: + report = dict( + sequence=args.sequence, + batch=args.batch, + cp_size=cp_size, + scope=( + "Attention forward/backward CUDA graph replay only; excludes projections, MoE," + " optimizer and checkpoint I/O" + ), + timing=( + "ABBA; 50 synchronized wall-clock samples per run; maximum across ranks per" + " sample" + ), + memory="Incremental allocated peak during capture and replay; not full GPU memory", + runs=records, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + logging.getLogger(__name__).info("Saved %s", args.output) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + main() diff --git a/docs/envvars.rst b/docs/envvars.rst index 0fee105fd0..3173dbc6e1 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -196,6 +196,12 @@ backend-selection overview. :Default: ``0`` :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention. +.. envvar:: NVTE_FUSED_ATTN_CP_USE_FAv4_BWD + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Use head-parallel FlashAttention-4 backward with cuDNN P2P context-parallel forward. Input packing and gradient restoration each use one all-to-all exchange. Requires an installed FA4 backend, SM100, FP16 or BF16, causal ``sbhd`` self-attention, equal query/key/value head counts divisible by CP size, CP size 2, 4, or 8, and head dimensions (128, 128) or (192, 128). Dropout, bias, softcapping, explicit sequence lengths, FP8, max-logit output, hierarchical CP, and ``torch.compile`` use the existing backward implementation. Unsupported configurations also retain that implementation. Performance and temporary memory depend on the sequence length and CP topology; benchmark before enabling. + .. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG :Type: ``int`` (0, 1 or 2), optionally followed by ``:`` diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index e7fd3189f3..9cfadf62d3 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -135,6 +135,10 @@ do fi NVTE_TORCH_COMPILE=0 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" fi + if [ "$sm_arch" -eq 100 ] && [[ "$fa_version" == 4.* ]]; then + python3 -m pytest -v -s --junitxml="$XML_LOG_DIR/pytest_test_mixed_cp.xml" \ + "$TE_PATH/tests/pytorch/attention/test_mixed_cp.py" || test_fail "test_mixed_cp.py" + fi done if [ "$RET" -ne 0 ]; then diff --git a/tests/pytorch/attention/run_mixed_cp.py b/tests/pytorch/attention/run_mixed_cp.py new file mode 100644 index 0000000000..356ede700b --- /dev/null +++ b/tests/pytorch/attention/run_mixed_cp.py @@ -0,0 +1,264 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Distributed worker for mixed context-parallel attention tests.""" + +import os +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import mixed_cp + +GROUP = None + + +@pytest.fixture(scope="module", autouse=True) +def distributed(): + """Use separate CP groups to exercise attention heads partitioned by TP.""" + global GROUP + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl", device_id=torch.device("cuda", torch.cuda.current_device())) + tp_size = int(os.environ.get("TEST_TP_SIZE", "1")) + cp_size = dist.get_world_size() // tp_size + for tp in range(tp_size): + ranks = list(range(tp * cp_size, (tp + 1) * cp_size)) + group = dist.new_group(ranks) + if dist.get_rank() in ranks: + GROUP = group + yield + dist.destroy_process_group() + + +def local_slice(tensor): + """Select the two balanced sequence chunks owned by this CP rank.""" + chunks = tensor.chunk(2 * dist.get_world_size(GROUP), dim=0) + rank = dist.get_rank(GROUP) + return torch.cat((chunks[rank], chunks[-1 - rank]), dim=0).contiguous() + + +def make_case( + batch=1, head_dim=192, dtype=torch.bfloat16, *, kv_heads=None, mask="causal", dropout=0 +): + """Return local attention inputs and an independent full-sequence reference.""" + torch.manual_seed(2026) + heads = 16 // int(os.environ.get("TEST_TP_SIZE", "1")) + full = [ + torch.randn(1024, batch, h, d, device="cuda", dtype=dtype) + for h, d in ((heads, head_dim), (kv_heads or heads, head_dim), (kv_heads or heads, 128)) + ] + gradient = torch.randn(1024, batch, heads, 128, device="cuda", dtype=dtype) + inputs = [local_slice(t).requires_grad_() for t in full] + cp_size = dist.get_world_size(GROUP) + first_rank = dist.get_rank() // cp_size * cp_size + module = ( + DotProductAttention( + num_attention_heads=heads, + num_gqa_groups=kv_heads or heads, + kv_channels=(head_dim, 128), + attention_dropout=dropout, + qkv_format="sbhd", + attn_mask_type=mask, + softmax_scale=0.083, + cp_group=GROUP, + cp_global_ranks=list(range(first_rank, first_rank + cp_size)), + cp_stream=torch.cuda.Stream(), + cp_comm_type="p2p", + ) + .cuda() + .train() + ) + return module, inputs, local_slice(gradient).flatten(-2), full, gradient + + +def run(module, inputs, gradient): + for tensor in inputs: + tensor.grad = None + output = module(*inputs) + output.backward(gradient) + torch.cuda.synchronize() + return output.detach().clone(), [t.grad.detach().clone() for t in inputs] + + +def check_gradients(actual, expected): + for value, reference in zip(actual, expected, strict=True): + torch.testing.assert_close(value.float(), reference.float(), rtol=0.04, atol=0.025) + relative_l2 = (value.float() - reference.float()).norm() / reference.float().norm() + assert relative_l2 < 0.008, relative_l2 + + +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("head_dim", [128, 192]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_output_and_gradients(monkeypatch, batch, head_dim, dtype): + module, inputs, gradient, full, full_gradient = make_case(batch, head_dim, dtype) + monkeypatch.setenv("NVTE_FUSED_ATTN_CP_USE_FAv4_BWD", "0") + baseline = run(module, inputs, gradient) + calls = [] + backward = mixed_cp._backward + + def observed(*args): + calls.append(True) + return backward(*args) + + monkeypatch.setattr(mixed_cp, "_backward", observed) + monkeypatch.setenv("NVTE_FUSED_ATTN_CP_USE_FAv4_BWD", "1") + actual = run(module, inputs, gradient) + assert len(calls) == 1 + torch.testing.assert_close(actual[0], baseline[0], rtol=0, atol=0) + check_gradients(actual[1], baseline[1]) + + # Compare all elements against full-sequence FP32 math attention as well. + reference_inputs = [t.float().detach().requires_grad_() for t in full] + with sdpa_kernel(SDPBackend.MATH): + output = F.scaled_dot_product_attention( + *(t.permute(1, 2, 0, 3) for t in reference_inputs), + is_causal=True, + scale=0.083, + ).permute(2, 0, 1, 3) + output.backward(full_gradient.float()) + torch.testing.assert_close( + actual[0].float(), local_slice(output).flatten(-2), rtol=0.04, atol=0.025 + ) + check_gradients(actual[1], [local_slice(t.grad) for t in reference_inputs]) + + +def gather_bits(tensor): + """Gather raw storage, including dtypes not supported by NCCL collectives.""" + wire = tensor.contiguous().view(torch.uint8) + pieces = [torch.empty_like(wire) for _ in range(dist.get_world_size(GROUP))] + dist.all_gather(pieces, wire, group=GROUP) + return [piece.view(tensor.dtype).reshape(tensor.shape) for piece in pieces] + + +def gather_reference(tensor, *, owner_shift=0): + """Assemble global sequence order independently of the packing implementation.""" + size, rank = dist.get_world_size(GROUP), dist.get_rank(GROUP) + pieces = gather_bits(tensor) + chunks = [None] * (2 * size) + for peer, piece in enumerate(pieces): + owner = (peer + owner_shift) % size + chunks[owner], chunks[-1 - owner] = piece.chunk(2, dim=0) + return torch.cat(chunks).chunk(size, dim=2)[rank].contiguous() + + +@pytest.mark.parametrize("batch,head_dim", [(1, 128), (2, 192)]) +def test_exchange_bits(batch, head_dim): + torch.manual_seed(78 + dist.get_rank()) + sequence, heads = 128, 16 + tensors = [ + torch.randint( + -32768, 32767, (sequence * 2, batch, heads, width), device="cuda", dtype=torch.int16 + )[::2].view(torch.bfloat16) + for width in (head_dim, head_dim, 128, 128, 128) + ] + bits = torch.randint( + -(2**31), 2**31 - 1, (batch, heads, sequence * 2), device="cuda", dtype=torch.int32 + )[..., ::2] + bits[0, 0, :4] = torch.tensor( + [0, -(2**31), 0x7F800000, 0x7FC12345], device="cuda", dtype=torch.int32 + ) + actual = mixed_cp._to_heads(*tensors, bits.view(torch.float32), GROUP) + for value, tensor, shift in zip(actual[:5], tensors, (0, 1, 1, 0, 0), strict=True): + expected = gather_reference(tensor.view(torch.int16), owner_shift=shift) + torch.testing.assert_close(value.view(torch.int16), expected, rtol=0, atol=0) + expected_lse = gather_reference(bits.permute(2, 0, 1).unsqueeze(-1)) + torch.testing.assert_close( + actual[5].view(torch.int32), expected_lse.squeeze(-1).permute(1, 2, 0), rtol=0, atol=0 + ) + returned = mixed_cp._to_sequence(*actual[:3], GROUP) + for value, tensor in zip(returned, actual[:3], strict=True): + parts = gather_bits(tensor) + expected = local_slice(torch.cat(parts, dim=2)) + torch.testing.assert_close( + value.view(torch.int16), expected.view(torch.int16), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("batch", [1, 2]) +def test_changed_input_graph_replay(monkeypatch, batch): + module, inputs, gradient, _, _ = make_case(batch=batch) + monkeypatch.setenv("NVTE_FUSED_ATTN_CP_USE_FAv4_BWD", "1") + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + run(module, inputs, gradient) + torch.cuda.current_stream().wait_stream(warmup) + for tensor in inputs: + tensor.grad = None + dist.barrier(group=GROUP) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = module(*inputs) + output.backward(gradient) + for factor in (0.9, 1.1): + with torch.no_grad(): + for tensor in inputs: + tensor.mul_(factor) + gradient.mul_(factor) + reference_inputs = [t.detach().clone().requires_grad_() for t in inputs] + monkeypatch.setenv("NVTE_FUSED_ATTN_CP_USE_FAv4_BWD", "0") + reference = run(module, reference_inputs, gradient) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, reference[0], rtol=0, atol=0) + check_gradients([t.grad for t in inputs], reference[1]) + graph.reset() + + +@pytest.mark.parametrize( + "mode", ["disabled", "missing_fa4", "gqa", "noncausal", "dropout", "custom_lengths"] +) +def test_native_fallback(monkeypatch, mode): + kwargs = {} + if mode == "gqa": + kwargs["kv_heads"] = 4 + if mode == "noncausal": + kwargs["mask"] = "no_mask" + if mode == "dropout": + kwargs["dropout"] = 0.1 + module, inputs, gradient, _, _ = make_case(**kwargs) + monkeypatch.setenv("NVTE_FUSED_ATTN_CP_USE_FAv4_BWD", "0" if mode == "disabled" else "1") + if mode == "missing_fa4": + from transformer_engine.pytorch.attention.dot_product_attention import backends + + monkeypatch.setattr(backends, "_flash_attn_bwd_v4", None) + + def unexpected(*args): + raise AssertionError("Unsupported configuration selected mixed backward") + + monkeypatch.setattr(mixed_cp, "_backward", unexpected) + if mode == "custom_lengths": + lengths = torch.tensor([0, 1024], dtype=torch.int32, device="cuda") + module(*inputs, cu_seqlens_q=lengths, cu_seqlens_kv=lengths).backward(gradient) + else: + run(module, inputs, gradient) + assert all(torch.isfinite(t.grad).all() for t in inputs) + + +def test_deterministic_mode(monkeypatch): + """Exercise the optional backward with deterministic attention enabled.""" + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + test_output_and_gradients(monkeypatch, 2, 192, torch.bfloat16) + + +if __name__ == "__main__": + args = [__file__, "-q", "-x", "-c", "/dev/null"] + if os.environ.get("XML_LOG_DIR"): + path = Path(os.environ["XML_LOG_DIR"]) + path.mkdir(parents=True, exist_ok=True) + args.extend( + [ + f"--junitxml={path / ('rank' + os.environ['RANK'] + '.xml')}", + "-o", + f"cache_dir={path / '.pytest-cache'}", + ] + ) + raise SystemExit(pytest.main(args)) diff --git a/tests/pytorch/attention/test_mixed_cp.py b/tests/pytorch/attention/test_mixed_cp.py new file mode 100644 index 0000000000..2bfd765b38 --- /dev/null +++ b/tests/pytorch/attention/test_mixed_cp.py @@ -0,0 +1,51 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Launch the mixed CP backward checks on supported GPU configurations.""" + +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + + +@pytest.mark.parametrize("cp_size,tp_size", [(2, 1), (4, 1), (8, 1), (2, 2), (4, 2)]) +def test_mixed_context_parallel(cp_size, tp_size, tmp_path): + if torch.cuda.device_count() < cp_size * tp_size: + pytest.skip("Not enough GPUs for this CP/TP configuration") + if torch.cuda.get_device_capability() != (10, 0): + pytest.skip("Mixed CP backward currently targets SM100") + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_bwd_v4, + ) + + if _flash_attn_bwd_v4 is None: + pytest.skip("FlashAttention-4 backward is unavailable") + env = dict( + os.environ, + TEST_TP_SIZE=str(tp_size), + XML_LOG_DIR=str(tmp_path), + NVTE_FLASH_ATTN="0", + NVTE_FUSED_ATTN="1", + NVTE_FUSED_ATTN_BACKEND="1", + ) + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc-per-node={cp_size * tp_size}", + str(Path(__file__).with_name("run_mixed_cp.py")), + ], + env=env, + capture_output=True, + text=True, + timeout=900, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/transformer_engine/common/triton/cp_packing.py b/transformer_engine/common/triton/cp_packing.py new file mode 100644 index 0000000000..9763c2577c --- /dev/null +++ b/transformer_engine/common/triton/cp_packing.py @@ -0,0 +1,155 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Bit-preserving packing between balanced sequence and head partitions.""" + +import triton +import triton.language as tl + + +@triton.jit +def _pack_cp_tensors( + Q, + K, + V, + O, + DO, + LSE, + WIRE, + S: tl.constexpr, + B: tl.constexpr, + DQ: tl.constexpr, + DV: tl.constexpr, + H: tl.constexpr, + CP: tl.constexpr, + Q0: tl.constexpr, + Q1: tl.constexpr, + Q2: tl.constexpr, + K0: tl.constexpr, + K1: tl.constexpr, + K2: tl.constexpr, + V0: tl.constexpr, + V1: tl.constexpr, + V2: tl.constexpr, + O0: tl.constexpr, + O1: tl.constexpr, + O2: tl.constexpr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + L0: tl.constexpr, + L1: tl.constexpr, + L2: tl.constexpr, + FORWARD: tl.constexpr, + W: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = i < CP * S * B * H * W + f = i % W + h = i // W % H + b = i // (W * H) % B + s = i // (W * H * B) % S + peer = i // (W * H * B * S) + if FORWARD: + source_s = s + source_h = peer * H + h + else: + source_s = tl.where( + s < S // 2, peer * (S // 2) + s, (2 * CP - 1 - peer) * (S // 2) + s - S // 2 + ) + source_h = h + bits = tl.load(Q + source_s * Q0 + b * Q1 + source_h * Q2 + f, valid & (f < DQ), other=0) + bits += tl.load( + K + source_s * K0 + b * K1 + source_h * K2 + f - DQ, + valid & (f >= DQ) & (f < (2 * DQ)), + other=0, + ) + bits += tl.load( + V + source_s * V0 + b * V1 + source_h * V2 + f - (2 * DQ), + valid & (f >= (2 * DQ)) & (f < (2 * DQ + DV)), + other=0, + ) + if FORWARD: + bits += tl.load( + O + source_s * O0 + b * O1 + source_h * O2 + f - (2 * DQ + DV), + valid & (f >= (2 * DQ + DV)) & (f < (2 * DQ + 2 * DV)), + other=0, + ) + bits += tl.load( + DO + source_s * D0 + b * D1 + source_h * D2 + f - (2 * DQ + 2 * DV), + valid & (f >= (2 * DQ + 2 * DV)) & (f < (2 * DQ + 3 * DV)), + other=0, + ) + lse = tl.load( + LSE + b * L0 + source_h * L1 + source_s * L2, + valid & (f >= (2 * DQ + 3 * DV)) & (f < (2 * DQ + 3 * DV + 2)), + other=0, + ).to(tl.uint32) + half = tl.where(f == (2 * DQ + 3 * DV), lse & 65535, lse >> 16).to(tl.int16) + bits += tl.where((f >= (2 * DQ + 3 * DV)) & (f < (2 * DQ + 3 * DV + 2)), half, 0) + tl.store(WIRE + i, bits, valid) + + +@triton.jit +def _unpack_cp_tensors( + WIRE, + Q, + K, + V, + O, + DO, + LSE, + S: tl.constexpr, + B: tl.constexpr, + DQ: tl.constexpr, + DV: tl.constexpr, + H: tl.constexpr, + CP: tl.constexpr, + FORWARD: tl.constexpr, + W: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = i < CP * S * B * H * W + f = i % W + h = i // W % H + b = i // (W * H) % B + s = i // (W * H * B) % S + peer = i // (W * H * B * S) + bits = tl.load(WIRE + i, valid, other=0) + if FORWARD: + target_s = tl.where( + s < S // 2, peer * (S // 2) + s, (2 * CP - 1 - peer) * (S // 2) + s - S // 2 + ) + kv_owner = (peer + 1) % CP + kv_s = tl.where( + s < S // 2, + kv_owner * (S // 2) + s, + (2 * CP - 1 - kv_owner) * (S // 2) + s - S // 2, + ) + row = (target_s * B + b) * H + h + kv_row = (kv_s * B + b) * H + h + else: + row = (s * B + b) * (CP * H) + peer * H + h + kv_row = row + tl.store(Q + row * DQ + f, bits, valid & (f < DQ)) + tl.store(K + kv_row * DQ + f - DQ, bits, valid & (f >= DQ) & (f < (2 * DQ))) + tl.store(V + kv_row * DV + f - (2 * DQ), bits, valid & (f >= (2 * DQ)) & (f < (2 * DQ + DV))) + if FORWARD: + tl.store( + O + row * DV + f - (2 * DQ + DV), + bits, + valid & (f >= (2 * DQ + DV)) & (f < (2 * DQ + 2 * DV)), + ) + tl.store( + DO + row * DV + f - (2 * DQ + 2 * DV), + bits, + valid & (f >= (2 * DQ + 2 * DV)) & (f < (2 * DQ + 3 * DV)), + ) + tl.store( + LSE + 2 * ((b * H + h) * S * CP + target_s) + f - (2 * DQ + 3 * DV), + bits, + valid & (f >= (2 * DQ + 3 * DV)) & (f < (2 * DQ + 3 * DV + 2)), + ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 510d14ac63..2898b73f8a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1654,6 +1654,44 @@ def forward( device_compute_capability < (10, 0) and cp_size == 2 ) + ctx.use_fa4_cp_bwd = False + if ( + os.getenv("NVTE_FUSED_ATTN_CP_USE_FAv4_BWD", "0") == "1" + and not torch.compiler.is_compiling() + and all( + ( + is_training, + use_fused_attention, + not fp8, + dropout_p == 0, + softcap == 0, + not return_max_logit, + attn_mask_type == "causal", + attn_bias_type == "no_bias", + attn_bias is None, + qkv_format == "sbhd", + cp_group_a2a is None, + cu_seqlens_q_padded is None, + cu_seqlens_kv_padded is None, + ) + ) + ): + from .mixed_cp import _is_supported + + if _is_supported(q, k, v, cp_size): + # FusedAttention slices the cached lengths without copying. Match + # that storage without reading device values during graph capture. + full_lengths = dpa_utils.get_full_cu_seqlens( + q.shape[1], q.shape[0] * cp_size, q.device + ) + ctx.use_fa4_cp_bwd = all( + lengths is not None + and lengths.shape == full_lengths.shape + and lengths.stride() == full_lengths.stride() + and lengths.data_ptr() == full_lengths.data_ptr() + for lengths in (cu_seqlens_q, cu_seqlens_kv) + ) + # set up attention args if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) @@ -2492,6 +2530,25 @@ def backward(ctx, dout, *_args): cu_seqlens_kv_padded, *other_tensors, ) = restore_from_func_ctx(ctx) + if ctx.use_fa4_cp_bwd: + from .mixed_cp import _backward + + # The final ring buffer belongs to the next CP rank. Packing restores + # its token order while distributing heads over the same group. + dq, dk, dv = _backward( + q.reshape(ctx.orig_q_shape), + kv[: ctx.k_numel].reshape(ctx.orig_k_shape), + kv[ctx.k_numel :].reshape(ctx.orig_v_shape), + out.reshape(ctx.orig_o_shape), + dout.reshape(ctx.orig_o_shape), + softmax_lse, + ctx.cp_group, + ctx.softmax_scale, + ctx.deterministic, + ) + nvtx_range_pop(f"{nvtx_label}") + return (None, dq, dk, dv) + (None,) * 27 + cu_seqlens_q_per_step = other_tensors[:cp_size] cu_seqlens_kv_per_step = other_tensors[cp_size : cp_size * 2] rng_states = other_tensors[cp_size * 2 : cp_size * 3] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py b/transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py new file mode 100644 index 0000000000..cf5c8ebe89 --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py @@ -0,0 +1,173 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused ring forward with head-parallel FlashAttention backward.""" + +import torch +import torch.distributed as dist +import triton + +from transformer_engine.common.triton.cp_packing import _pack_cp_tensors, _unpack_cp_tensors +from transformer_engine.pytorch.utils import get_device_compute_capability + + +def _is_supported(q, k, v, cp_size): + """Check static geometry and optional backend availability before dispatch.""" + from .backends import _flash_attn_bwd_v4 + + return ( + _flash_attn_bwd_v4 is not None + and get_device_compute_capability() == (10, 0) + and cp_size in (2, 4, 8) + and all(t.ndim == 4 and t.is_cuda for t in (q, k, v)) + and q.shape == k.shape + and q.shape[:3] == v.shape[:3] + and q.shape[0] > 0 + and q.shape[0] % 2 == 0 + and q.shape[1] > 0 + and q.shape[2] > 0 + and q.shape[2] % cp_size == 0 + and (q.shape[-1], v.shape[-1]) in ((128, 128), (192, 128)) + and q.dtype in (torch.float16, torch.bfloat16) + and all(t.dtype == q.dtype and t.device == q.device and t.stride(-1) == 1 for t in (k, v)) + and q.stride(-1) == 1 + ) + + +def _backward(q, k, v, output, gradient, lse, group, softmax_scale, deterministic): + """Exchange the saved ring tensors once in each direction around FA4 backward.""" + from .backends import _flash_attn_bwd_v4 + + q, k, v, output, gradient, lse = _to_heads(q, k, v, output, gradient, lse, group) + dq, dk, dv = _flash_attn_bwd_v4( + q.transpose(0, 1), + k.transpose(0, 1), + v.transpose(0, 1), + output.transpose(0, 1), + gradient.transpose(0, 1), + lse, + softmax_scale=softmax_scale, + causal=True, + deterministic=deterministic, + ) + return _to_sequence(dq.transpose(0, 1), dk.transpose(0, 1), dv.transpose(0, 1), group) + + +def _exchange(tensors: tuple, group: object, *, forward: bool, lse=None) -> tuple: + size = dist.get_world_size(group) + q, k, v = tensors[:3] + sequence, batch, heads, _ = q.shape + if ( + batch < 1 + or size not in (2, 4, 8) + or any( + t.dtype != q.dtype + or t.dtype not in (torch.float16, torch.bfloat16) + or t.stride(-1) != 1 + for t in tensors + ) + ): + raise ValueError("Fused packing requires FP16/BF16, CP2/4/8 and contiguous head widths") + dq, dv = q.shape[-1], v.shape[-1] + if (dq, dv) not in ((128, 128), (192, 128)) or k.shape != q.shape or v.shape[:3] != q.shape[:3]: + raise ValueError("Unexpected MLA head geometry") + local_s, local_h = (sequence, heads // size) if forward else (sequence // size, heads) + if local_s % 2 or (forward and heads % size): + raise ValueError("Invalid balanced CP partition") + width = triton.cdiv(2 * dq + 3 * dv + 2, 64) * 64 if forward else 2 * dq + dv + # NCCL byte transport and integer kernel loads preserve every FP32 LSE bit. + wire = torch.empty( + size * local_s * batch * local_h * width * 2, dtype=torch.uint8, device=q.device + ) + received = torch.empty_like(wire) + o, do = tensors[3:] if forward else (q, q) + if forward and (lse.dtype != torch.float32 or lse.shape != (batch, heads, sequence)): + raise ValueError("Unexpected native LSE geometry") + _pack_cp_tensors[(triton.cdiv(wire.numel() // 2, 1024),)]( + q.view(torch.int16), + k.view(torch.int16), + v.view(torch.int16), + o.view(torch.int16), + do.view(torch.int16), + lse.view(torch.int32) if forward else None, + wire.view(torch.int16), + local_s, + batch, + dq, + dv, + local_h, + size, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + o.stride(0), + o.stride(1), + o.stride(2), + do.stride(0), + do.stride(1), + do.stride(2), + lse.stride(0) if forward else 0, + lse.stride(1) if forward else 0, + lse.stride(2) if forward else 0, + forward, + width, + 1024, + ) + dist.all_to_all_single(received, wire, group=group) + # The sent buffer is dead after the collective and becomes the output slab. + out_s, out_h = (sequence * size, local_h) if forward else (local_s, heads * size) + count = out_s * batch * out_h + widths = (dq, dq, dv, dv, dv) if forward else (dq, dq, dv) + slab = wire.view(q.dtype) + outputs = [] + offset = 0 + for head_width in widths: + outputs.append( + slab[offset : offset + count * head_width].view(out_s, batch, out_h, head_width) + ) + offset += count * head_width + global_lse = ( + slab[offset : offset + count * 2].view(torch.float32).view(batch, out_h, out_s) + if forward + else None + ) + qo, ko, vo = outputs[0], outputs[1], outputs[2] + oo, d_o = outputs[3:] if forward else (qo, qo) + _unpack_cp_tensors[(triton.cdiv(wire.numel() // 2, 1024),)]( + received.view(torch.int16), + qo.view(torch.int16), + ko.view(torch.int16), + vo.view(torch.int16), + oo.view(torch.int16), + d_o.view(torch.int16), + global_lse.view(torch.int16) if forward else None, + local_s, + batch, + dq, + dv, + local_h, + size, + forward, + width, + 1024, + ) + return (*outputs, global_lse) if forward else tuple(outputs) + + +def _to_heads(query, key, value, output, gradient, lse, group) -> tuple: + """Exchange native ring tensors, accounting for the final KV buffer owner.""" + result = _exchange((query, key, value, output, gradient), group, forward=True, lse=lse) + return result[0], result[1], result[2], result[3], result[4], result[5] + + +def _to_sequence(dq, dk, dv, group) -> tuple: + """Return all head partitions to their original balanced token owner.""" + result = _exchange((dq, dk, dv), group, forward=False) + return result[0], result[1], result[2] From 37813c67fec821fe7c1bda4df1d14ab7e9919cc1 Mon Sep 17 00:00:00 2001 From: "ryan.u" Date: Fri, 11 Sep 2026 21:48:08 +0900 Subject: [PATCH 2/3] Support wide offsets in CP tensor packing Signed-off-by: ryan.u --- tests/pytorch/attention/run_mixed_cp.py | 16 ++++++++++++++++ transformer_engine/common/triton/cp_packing.py | 12 ++++++++++-- .../attention/dot_product_attention/mixed_cp.py | 13 ++++++++++--- 3 files changed, 36 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/attention/run_mixed_cp.py b/tests/pytorch/attention/run_mixed_cp.py index 356ede700b..d2df4ee096 100644 --- a/tests/pytorch/attention/run_mixed_cp.py +++ b/tests/pytorch/attention/run_mixed_cp.py @@ -181,6 +181,22 @@ def test_exchange_bits(batch, head_dim): ) +def test_exchange_large_stride(): + """Exercise offsets beyond INT_MAX without exchanging a large payload.""" + torch.manual_seed(78 + dist.get_rank()) + query = torch.empty_strided( + (4, 1, 16, 128), (2**30, 2048, 128, 1), device="cuda", dtype=torch.bfloat16 + ) + query.view(torch.int16).copy_( + torch.randint(-32768, 32767, query.shape, device="cuda", dtype=torch.int16) + ) + tensors = [query] + [torch.randn_like(query.contiguous()) for _ in range(4)] + lse = torch.randn(1, 16, 4, device="cuda", dtype=torch.float32) + actual = mixed_cp._to_heads(*tensors, lse, GROUP) + expected = gather_reference(query.view(torch.int16)) + torch.testing.assert_close(actual[0].view(torch.int16), expected, rtol=0, atol=0) + + @pytest.mark.parametrize("batch", [1, 2]) def test_changed_input_graph_replay(monkeypatch, batch): module, inputs, gradient, _, _ = make_case(batch=batch) diff --git a/transformer_engine/common/triton/cp_packing.py b/transformer_engine/common/triton/cp_packing.py index 9763c2577c..08af2d136d 100644 --- a/transformer_engine/common/triton/cp_packing.py +++ b/transformer_engine/common/triton/cp_packing.py @@ -44,8 +44,12 @@ def _pack_cp_tensors( FORWARD: tl.constexpr, W: tl.constexpr, BLOCK: tl.constexpr, + INDEX64: tl.constexpr, ): - i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + block = tl.program_id(0) + if INDEX64: + block = block.to(tl.int64) + i = block * BLOCK + tl.arange(0, BLOCK) valid = i < CP * S * B * H * W f = i % W h = i // W % H @@ -110,8 +114,12 @@ def _unpack_cp_tensors( FORWARD: tl.constexpr, W: tl.constexpr, BLOCK: tl.constexpr, + INDEX64: tl.constexpr, ): - i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + block = tl.program_id(0) + if INDEX64: + block = block.to(tl.int64) + i = block * BLOCK + tl.arange(0, BLOCK) valid = i < CP * S * B * H * W f = i % W h = i // W % H diff --git a/transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py b/transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py index cf5c8ebe89..92d6585ee1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/mixed_cp.py @@ -76,10 +76,15 @@ def _exchange(tensors: tuple, group: object, *, forward: bool, lse=None) -> tupl if local_s % 2 or (forward and heads % size): raise ValueError("Invalid balanced CP partition") width = triton.cdiv(2 * dq + 3 * dv + 2, 64) * 64 if forward else 2 * dq + dv - # NCCL byte transport and integer kernel loads preserve every FP32 LSE bit. - wire = torch.empty( - size * local_s * batch * local_h * width * 2, dtype=torch.uint8, device=q.device + elements = size * local_s * batch * local_h * width + # Strided inputs can require wide offsets even when the payload is small. + index64 = elements >= 2**31 or any( + sum((dim - 1) * stride for dim, stride in zip(t.shape, t.stride())) >= 2**31 + for t in (*tensors, lse) + if t is not None ) + # NCCL byte transport and integer kernel loads preserve every FP32 LSE bit. + wire = torch.empty(elements * 2, dtype=torch.uint8, device=q.device) received = torch.empty_like(wire) o, do = tensors[3:] if forward else (q, q) if forward and (lse.dtype != torch.float32 or lse.shape != (batch, heads, sequence)): @@ -119,6 +124,7 @@ def _exchange(tensors: tuple, group: object, *, forward: bool, lse=None) -> tupl forward, width, 1024, + index64, ) dist.all_to_all_single(received, wire, group=group) # The sent buffer is dead after the collective and becomes the output slab. @@ -157,6 +163,7 @@ def _exchange(tensors: tuple, group: object, *, forward: bool, lse=None) -> tupl forward, width, 1024, + index64, ) return (*outputs, global_lse) if forward else tuple(outputs) From b517bbfd657836a9fe69eb6c36338bbd5fdd4cf7 Mon Sep 17 00:00:00 2001 From: "ryan.u" Date: Sat, 12 Sep 2026 18:24:14 +0900 Subject: [PATCH 3/3] Keep wide-index CP tests small by default Signed-off-by: ryan.u --- tests/pytorch/attention/run_mixed_cp.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/run_mixed_cp.py b/tests/pytorch/attention/run_mixed_cp.py index d2df4ee096..81b10aba43 100644 --- a/tests/pytorch/attention/run_mixed_cp.py +++ b/tests/pytorch/attention/run_mixed_cp.py @@ -5,6 +5,7 @@ """Distributed worker for mixed context-parallel attention tests.""" import os +from functools import partial from pathlib import Path import pytest @@ -148,8 +149,17 @@ def gather_reference(tensor, *, owner_shift=0): return torch.cat(chunks).chunk(size, dim=2)[rank].contiguous() +@pytest.mark.parametrize("index64", [False, True]) @pytest.mark.parametrize("batch,head_dim", [(1, 128), (2, 192)]) -def test_exchange_bits(batch, head_dim): +def test_exchange_bits(monkeypatch, batch, head_dim, index64): + if index64: + # Exercise both wide-index kernels with the same small reference tensors. + def run_wide(run, *args, **kwargs): + return run(*args[:-1], True, **kwargs) + + for kernel in (mixed_cp._pack_cp_tensors, mixed_cp._unpack_cp_tensors): + monkeypatch.setattr(kernel, "run", partial(run_wide, kernel.run)) + torch.manual_seed(78 + dist.get_rank()) sequence, heads = 128, 16 tensors = [ @@ -181,6 +191,10 @@ def test_exchange_bits(batch, head_dim): ) +@pytest.mark.skipif( + os.environ.get("NVTE_TEST_CP_LARGE_STRIDE") != "1", + reason="Set NVTE_TEST_CP_LARGE_STRIDE=1 for the 6 GiB-per-rank offset regression", +) def test_exchange_large_stride(): """Exercise offsets beyond INT_MAX without exchanging a large payload.""" torch.manual_seed(78 + dist.get_rank())