From a9cd697fa2fb1da2df187f49e706e3e698ae6dc5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 7 Jul 2026 17:35:28 +0200 Subject: [PATCH 01/35] [PyTorch] DotProductAttention: declarative packed qkv/kv inputs Fused QKV projections naturally produce one packed buffer, but DotProductAttention forces callers to slice it into q/k/v views that TE then reverse-engineers with pointer-based layout detection (get_qkv_layout inspects data_ptr/storage_offset on every forward, which graph-breaks under torch.compile and adds CPU overhead). Let callers declare the packing instead (JAX-style): * DotProductAttention.forward gains optional qkv_layer (fully packed QKV: [b,s,3,h,d]/[s,b,3,h,d]/[b,s,h,3,d]/[s,b,h,3,d] dense, [t,3,h,d]/ [t,h,3,d] thd), kv_layer (packed KV used with query_layer), and qkv_interleave_dim (-3 or -2; explicit knob rather than shape inference since h==3 or hg==2 would be ambiguous). * Q/K/V are derived as zero-copy select() views and the exact layout enum (bs3hd, bsh3d, sb3hd, bshd_bs2hd, t3hd, ...) is constructed declaratively -- it is truthful by construction, so get_qkv_layout is never called on this path, including for thd and FP8 DPA. * combine_and_quantize no longer re-combines what is already combined: a new optional combined= argument carries the caller's original packed buffer, which is quantized directly instead of rebuilding the packed buffer from q/k/v views via combine_tensors (a raw set_ with a silent adjacency/interleave assumption). The packed original is threaded from DPA.forward through FusedAttention to FusedAttnFunc.forward; all legacy call sites are untouched (combined=None preserves exact behavior), and backward combine calls are unchanged (gradients have no pre-packed original). Tests: dense fwd+grad bit-exactness vs separate contiguous q/k/v for bs3hd/bsh3d/sb3hd/kv-packed/GQA (fused + flash), validation errors, torch.compile (no data_ptr/UntypedStorage graph breaks), FP8 combined-vs-views bit equivalence, and detection-free declared t3hd. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 355 ++++++++++++++++++ .../dot_product_attention/backends.py | 12 +- .../dot_product_attention.py | 126 ++++++- .../attention/dot_product_attention/utils.py | 30 +- 4 files changed, 515 insertions(+), 8 deletions(-) create mode 100644 tests/pytorch/attention/test_dpa_packed_inputs.py diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py new file mode 100644 index 0000000000..7e63630014 --- /dev/null +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -0,0 +1,355 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for declarative packed QKV/KV inputs to DotProductAttention. + +Instead of slicing a fused-projection buffer into q/k/v views (which TE then +reverse-engineers via pointer-based layout detection), callers can pass the +packed tensor directly (``qkv_layer``/``kv_layer`` + ``qkv_interleave_dim``). +Q/K/V are derived as zero-copy views and the exact layout string (e.g. +``bs3hd``) is declared, not detected -- including for thd and FP8 DPA. +""" + +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.attention.dot_product_attention.utils import ( + combine_and_quantize, +) +from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + fused_attn_fwd, +) +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") + +_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 _fused_backend_supported(): + try: + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + fused_attn_fwd( + True, _S, _S, _cu_seqlens(), _cu_seqlens(), q, q.clone(), q.clone(), _DTYPE, + tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + dropout=0.0, qkv_layout="bshd_bshd_bshd", o_format="bshd", + attn_bias_type="no_bias", attn_mask_type="no_mask", + ) + return True + except Exception: + return False + + +requires_fused = pytest.mark.skipif( + not (torch.cuda.is_available() and _fused_backend_supported()), + reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", +) + + +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 + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + dpa_module._attention_backends["backend_selection_requires_update"] = True + + +def _make_dpa(qkv_format, num_gqa_groups=None): + return DotProductAttention( + _H, + _D, + num_gqa_groups=num_gqa_groups, + attention_dropout=0.0, + qkv_format=qkv_format, + attn_mask_type="no_mask", + ) + + +def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): + for name, x, y in zip(names, result, reference): + assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" + + +def _fresh_parts(qkv_format, num_heads, seed=0): + torch.manual_seed(seed) + shape = (_B, _S, num_heads, _D) if qkv_format == "bshd" else (_S, _B, num_heads, _D) + return [torch.randn(*shape, dtype=_DTYPE, device="cuda") for _ in range(3)] + + +def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): + """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" + hg = num_gqa_groups or _H + torch.manual_seed(0) + q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + kv_shape = (_B, _S, hg, _D) if qkv_format == "bshd" else (_S, _B, hg, _D) + q = torch.randn(*q_shape, dtype=_DTYPE, device="cuda") + k = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") + v = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") + q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) + out.backward(torch.ones_like(out)) + return out, q.grad, k.grad, v.grad + + +# --------------------------------------------------------------------------- +# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact +# --------------------------------------------------------------------------- + + +@requires_fused +@pytest.mark.parametrize( + "qkv_format, interleave_dim", + [ + pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd + pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) + pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd + ], +) +def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): + """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads + flow back into the packed tensor itself.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline(qkv_format) + + torch.manual_seed(0) + q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + parts = [torch.randn(*q_shape, dtype=_DTYPE, device="cuda") for _ in range(3)] + stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d + qkv = torch.stack(parts, dim=stack_dim).requires_grad_() + + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) + out.backward(torch.ones_like(out)) + + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + grads = [qkv.grad.select(stack_dim, i) for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +@requires_fused +@pytest.mark.parametrize( + "num_gqa_groups", + [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) +) +def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): + """kv_layer packed input (with separate query) is bit-exact vs separate + contiguous q/k/v; grads flow into the packed kv tensor.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline("bshd", num_gqa_groups) + + hg = num_gqa_groups or _H + torch.manual_seed(0) + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + k = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") + v = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") + q = q.clone().requires_grad_() + kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] + + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) + out.backward(torch.ones_like(out)) + + assert kv.grad is not None and kv.grad.shape == kv.shape + _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) + + +# --------------------------------------------------------------------------- +# 2. Flash backend smoke +# --------------------------------------------------------------------------- + + +def test_dpa_flash_qkv_layer(monkeypatch): + """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" + _force_backend(monkeypatch, "flash") + try: + reference = _dpa_separate_baseline("bshd") + except Exception as exc: + pytest.skip(f"flash attention backend not available: {exc}") + + torch.manual_seed(0) + parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] + qkv = torch.stack(parts, dim=2).requires_grad_() + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd")(qkv_layer=qkv) + out.backward(torch.ones_like(out)) + grads = [qkv.grad[:, :, i] for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +# --------------------------------------------------------------------------- +# 3. Validation errors +# --------------------------------------------------------------------------- + + +def test_dpa_packed_input_validation(): + dpa = _make_dpa("bshd") + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") + k = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + + with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): + dpa(qkv_layer=qkv, key_layer=k) + with pytest.raises(ValueError, match="query_layer is required when kv_layer"): + dpa(kv_layer=kv) + with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): + dpa(qkv_layer=qkv, qkv_interleave_dim=-1) + with pytest.raises(ValueError, match="mutually exclusive"): + dpa(qkv_layer=qkv, kv_layer=kv) + with pytest.raises(ValueError, match="must have size 3 at dim"): + dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 + with pytest.raises(ValueError, match="required unless packed"): + dpa() + + +# --------------------------------------------------------------------------- +# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer +# --------------------------------------------------------------------------- + + +@requires_fused +def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): + _force_backend(monkeypatch, "fused") + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + + torch.manual_seed(0) + parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] + qkv = torch.stack(parts, dim=2).requires_grad_() + dpa = _make_dpa("bshd") + + def fn(x): + return dpa(qkv_layer=x) + + dpa_module._attention_backends["backend_selection_requires_update"] = True + eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo + compiled_out = torch.compile(fn)(qkv) + compiled_out.backward(torch.ones_like(compiled_out)) + + breaks = dict(torch._dynamo.utils.counters["graph_break"]) + torch._dynamo.reset() + pointer_breaks = { + reason: count + for reason, count in breaks.items() + if "data_ptr" in reason or "UntypedStorage" in reason + } + assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" + assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" + + +# --------------------------------------------------------------------------- +# 5. FP8 combine refactor: combined= path is bit-identical to combine_tensors path +# --------------------------------------------------------------------------- + + +def _fp8_quantizer(): + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + ) + + +def test_combine_and_quantize_combined_matches_views(): + """Quantizing the caller's packed buffer directly (combined=) produces the same + _data bits and scale_inv as rebuilding the packed buffer from q/k/v views via + combine_tensors (the old set_-based path).""" + torch.manual_seed(0) + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + + old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined=qkv) + + assert old[3] == new[3] == "bs3hd" + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +def test_combine_and_quantize_combined_kv_matches_views(): + """Same for the kv-packed (group 2) layout.""" + torch.manual_seed(0) + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") + k, v = kv[:, :, 0], kv[:, :, 1] + + old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined=kv) + + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +# --------------------------------------------------------------------------- +# 6. thd declarative: t3hd is declared, get_qkv_layout is never called +# --------------------------------------------------------------------------- + + +def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): + """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling + get_qkv_layout; full forward+backward runs if a thd backend is available.""" + calls = [] + orig_get_qkv_layout = dpa_utils.get_qkv_layout + + def counting(*args, **kwargs): + calls.append(kwargs.get("qkv_format")) + return orig_get_qkv_layout(*args, **kwargs) + + monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) + + seen_layouts = [] + orig_get_backend = dpa_utils.get_attention_backend + + def recording(params): + seen_layouts.append(params.qkv_layout) + return orig_get_backend(params) + + monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) + + torch.manual_seed(0) + t = _B * _S + qkv = torch.randn(t, 3, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) + cu = _cu_seqlens() + dpa = DotProductAttention( + _H, _D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" + ) + dpa_module._attention_backends["backend_selection_requires_update"] = True + ran_full = True + try: + out = dpa( + qkv_layer=qkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=_S, + max_seqlen_kv=_S, + ) + out.backward(torch.ones_like(out)) + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + except ValueError: + # No thd-capable backend on this device; the layout step (the subject of + # this test) runs before backend dispatch, so the assertions still hold. + ran_full = False + + assert not calls, "get_qkv_layout must not be called for declared packed thd input" + assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" + if not ran_full: + pytest.skip("layout declaration verified; no thd backend available for full run") diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8f42983553..50565bebf3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1328,6 +1328,7 @@ def forward( fp8_output, layer_number, return_max_logit, + packed_qkv=None, ): # pylint: disable=missing-function-docstring @@ -1390,7 +1391,13 @@ def forward( q_fp8, k_fp8, v_fp8 = q, k, v else: q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( - qkv_layout, q, k, v, QKV_quantizer, used_in_backward=is_training + qkv_layout, + q, + k, + v, + QKV_quantizer, + used_in_backward=is_training, + combined=packed_qkv, ) # print quantizers @@ -1897,6 +1904,7 @@ def backward(ctx, d_out, *_args): None, None, None, + None, # packed_qkv ) @@ -1991,6 +1999,7 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + packed_qkv: Optional[torch.Tensor] = None, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2210,6 +2219,7 @@ def forward( fp8_output, self.layer_number, self.return_max_logit, + packed_qkv, ) if self.return_max_logit: 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..e968817b97 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 @@ -1008,9 +1008,9 @@ def get_quantizer_roles( @no_torch_dynamo(recursive=False) def forward( self, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, + query_layer: Optional[torch.Tensor] = None, + key_layer: Optional[torch.Tensor] = None, + value_layer: Optional[torch.Tensor] = None, attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None, qkv_format: str = None, cu_seqlens_q: torch.Tensor = None, @@ -1035,6 +1035,9 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + qkv_layer: Optional[torch.Tensor] = None, + kv_layer: Optional[torch.Tensor] = None, + qkv_interleave_dim: int = -3, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1245,8 +1248,113 @@ def forward( Runtime tensors exposed to score_mod_bprop as cuDNN graph tensors. Keys are user-defined string names consumed by the callback through ``tensors[name]``; there is no predefined set of accepted keys. + qkv_layer: Optional[torch.Tensor], default = None + Fully packed QKV tensor. When the QKV projection produces one packed buffer + (e.g. a fused QKV GEMM), it can be passed here directly instead of slicing + it into :attr:`query_layer`/:attr:`key_layer`/:attr:`value_layer` views. + For :attr:`qkv_format` = {"bshd", "sbhd"}, it must be a 5D tensor of shape + ``[b, s, 3, h, d]``/``[s, b, 3, h, d]`` (:attr:`qkv_interleave_dim` = -3) or + ``[b, s, h, 3, d]``/``[s, b, h, 3, d]`` (:attr:`qkv_interleave_dim` = -2); + for :attr:`qkv_format` = "thd", a 4D tensor of shape ``[t, 3, h, d]`` or + ``[t, h, 3, d]``. Q/K/V are derived as zero-copy views and the memory layout + (e.g. ``bs3hd``) is declared from the packing itself, so no pointer-based + layout detection runs on this path -- including for "thd" and FP8 attention. + Mutually exclusive with :attr:`query_layer`, :attr:`key_layer`, + :attr:`value_layer` and :attr:`kv_layer`. + kv_layer: Optional[torch.Tensor], default = None + Packed KV tensor, used together with :attr:`query_layer` + (e.g. ``[b, s, 2, hg, d]`` for :attr:`qkv_interleave_dim` = -3, or + ``[b, s, hg, 2, d]`` for :attr:`qkv_interleave_dim` = -2). K/V are derived + as zero-copy views and the layout (e.g. ``bshd_bs2hd``) is declared, not + detected. Mutually exclusive with :attr:`key_layer`, :attr:`value_layer` + and :attr:`qkv_layer`. + qkv_interleave_dim: int, default = -3 + Dimension of :attr:`qkv_layer`/:attr:`kv_layer` where the 3 (QKV) or 2 (KV) + interleave sits; must be -3 (e.g. ``bs3hd``) or -2 (e.g. ``bsh3d``, + Megatron-style). This is an explicit knob rather than shape inference, + since e.g. ``h == 3`` would make the shapes ambiguous. """ + # Declarative packed inputs: derive q/k/v as zero-copy views of the packed + # buffer and construct the exact layout string from the declaration. The + # layout enum is truthful by construction, so no pointer-based detection + # is needed downstream (this also covers thd and FP8 DPA). + packed_tensor = None + declared_qkv_layout = None + if qkv_layer is not None or kv_layer is not None: + if qkv_layer is not None and kv_layer is not None: + raise ValueError("qkv_layer and kv_layer are mutually exclusive.") + if inference_params is not None: + raise ValueError( + "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" + " (inference_params); pass separate query/key/value tensors instead." + ) + if qkv_interleave_dim not in (-3, -2): + raise ValueError( + "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" + f" {qkv_interleave_dim}." + ) + packed_format = qkv_format if qkv_format is not None else self.qkv_format + + def _packed_layout(fmt: str, num: int) -> str: + # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d + pos = len(fmt) + qkv_interleave_dim + 1 + return fmt[:pos] + str(num) + fmt[pos:] + + if qkv_layer is not None: + if any(x is not None for x in (query_layer, key_layer, value_layer)): + raise ValueError( + "qkv_layer already packs Q, K and V: query_layer, key_layer and" + " value_layer must be None when qkv_layer is provided." + ) + expected_rank = 4 if packed_format == "thd" else 5 + if qkv_layer.dim() != expected_rank: + raise ValueError( + f"qkv_layer must be a {expected_rank}D tensor for" + f" qkv_format={packed_format!r}, got {qkv_layer.dim()}D." + ) + if qkv_layer.shape[qkv_interleave_dim] != 3: + raise ValueError( + f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." + ) + query_layer, key_layer, value_layer = ( + qkv_layer.select(qkv_interleave_dim, i) for i in range(3) + ) + packed_tensor = qkv_layer + declared_qkv_layout = _packed_layout(packed_format, 3) + else: + if query_layer is None: + raise ValueError( + "kv_layer packs only K and V: query_layer is required when kv_layer" + " is provided." + ) + if key_layer is not None or value_layer is not None: + raise ValueError( + "kv_layer already packs K and V: key_layer and value_layer must be" + " None when kv_layer is provided." + ) + if kv_layer.dim() != query_layer.dim() + 1: + raise ValueError( + f"kv_layer must have one more dimension than query_layer, got" + f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." + ) + if kv_layer.shape[qkv_interleave_dim] != 2: + raise ValueError( + f"kv_layer must have size 2 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." + ) + key_layer, value_layer = ( + kv_layer.select(qkv_interleave_dim, i) for i in range(2) + ) + packed_tensor = kv_layer + declared_qkv_layout = f"{packed_format}_{_packed_layout(packed_format, 2)}" + elif query_layer is None or key_layer is None or value_layer is None: + raise ValueError( + "query_layer, key_layer and value_layer are required unless packed" + " inputs (qkv_layer or query_layer + kv_layer) are provided." + ) + with self.prepare_forward_ctx( query_layer, num_gemms=3, @@ -1432,7 +1540,15 @@ def forward( cu_seqlens_kv_padded = None # get qkv's memory layout - if all( + if declared_qkv_layout is not None: + # Packed inputs (qkv_layer/kv_layer) declare the layout: the enum is + # truthful by construction, so the pointer-based detection in + # get_qkv_layout is skipped entirely -- for dense, thd (t3hd/th3d) + # and FP8 DPA alike. + qkv_layout = declared_qkv_layout + q_format = qkv_format + kv_format = qkv_format + elif all( isinstance(x, Float8TensorStorage) for x in [query_layer, key_layer, value_layer] ): ( @@ -1812,6 +1928,7 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, + packed_qkv=packed_tensor, ) return self.fused_attention( query_layer, @@ -1847,6 +1964,7 @@ def forward( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, + packed_qkv=packed_tensor, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 9913b78dfc..3176200ffa 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2783,8 +2783,18 @@ def combine_and_quantize( used_in_forward=True, used_in_backward=False, keep_same_data_and_scale_inv_format=False, + combined: Optional[torch.Tensor] = None, ): - """Combine Q, K, V tensors based on qkv_layout and quantize them together.""" + """Combine Q, K, V tensors based on qkv_layout and quantize them together. + + When ``combined`` is provided, it must be the caller's original packed buffer + matching the packed group in ``qkv_layout`` (the full QKV tensor for ``3`` + layouts such as ``bs3hd``, or the KV tensor for ``2`` layouts such as + ``bshd_bs2hd``). It is then quantized directly instead of re-deriving the + packed buffer from the q/k/v views via ``combine_tensors`` (which rebuilds it + with a raw ``set_`` under a silent adjacency/interleave assumption). Ignored + for MXFP8 quantization. + """ if isinstance(qkv_quantizer, MXFP8Quantizer): qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) assert qkv_format in ("bshd", "sbhd"), ( @@ -2884,12 +2894,26 @@ def combine_and_quantize( match qkv_group: case 1: dim = qkv_layout.find("3") - qkv = combine_tensors([q, k, v], dim) + if combined is not None: + assert combined.shape[dim] == 3, ( + f"combined QKV tensor does not match qkv_layout {qkv_layout}: expected" + f" size 3 at dim {dim}, got shape {tuple(combined.shape)}." + ) + qkv = combined + else: + qkv = combine_tensors([q, k, v], dim) qkv_fp8 = qkv_quantizer(qkv) q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True) case 2: dim = qkv_layout.split("_")[1].find("2") - kv = combine_tensors([k, v], dim) + if combined is not None: + assert combined.shape[dim] == 2, ( + f"combined KV tensor does not match qkv_layout {qkv_layout}: expected" + f" size 2 at dim {dim}, got shape {tuple(combined.shape)}." + ) + kv = combined + else: + kv = combine_tensors([k, v], dim) tensors = [q, kv] num_tensors = len(tensors) shapes = [x.shape for x in tensors] From 5f70d121989747878eb6baa5015b647aee98cd93 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 12:55:58 +0200 Subject: [PATCH 02/35] [PyTorch] MultiheadAttention: pass packed projection output to DPA declaratively Adopt the new DotProductAttention packed API inside MultiheadAttention: * self-attention (np == ng): the fused QKV projection output, already viewed as [.., h, 3, d] (qkv_weight_interleaved) or [.., 3, h, d], is handed to DPA directly as qkv_layer with the matching qkv_interleave_dim (-2 / -3) -- no SplitAlongDim slicing in MHA and no pointer-based layout detection in DPA. * cross-attention: the packed KV projection output is exposed as [.., hg, 2, d] / [.., 2, hg, d] and passed as kv_layer. * The pass-through only engages when no per-tensor operation needs the individual q/k/v slices: it is skipped for RoPE, QK normalization, KV caching (inference_params), CPU offloading, GQA (np != ng, not a uniform 3-interleave) and quantized (FP8) projection outputs; those keep the legacy sliced-views path unchanged. Tests: MHA self (interleaved + non-interleaved) and cross (both interleaves) are bit-exact vs the same MHA with packed inputs converted back to separate contiguous q/k/v (output, input grad, weight grads); spy asserts the packed argument and interleave dim actually reach DPA; GQA and RoPE fall back to the views path. TransformerLayer regression suite unchanged; test_kv_cache failures on this device are pre-existing on origin/main (verified). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 201 +++++++++++++++++- .../pytorch/attention/multi_head_attention.py | 143 +++++++++---- 2 files changed, 300 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py index 7e63630014..b87e20a344 100644 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -16,7 +16,7 @@ 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 import DotProductAttention, MultiheadAttention from transformer_engine.pytorch.attention.dot_product_attention import ( dot_product_attention as dpa_module, ) @@ -353,3 +353,202 @@ def recording(params): assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" if not ran_full: pytest.skip("layout declaration verified; no thd backend available for full run") + + +# --------------------------------------------------------------------------- +# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed +# packed (qkv_layer/kv_layer) to DotProductAttention +# --------------------------------------------------------------------------- + + +def _spy_packed_dpa(monkeypatch, record): + """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" + orig = DotProductAttention.forward + + def spy(self, *args, **kwargs): + record.append( + ( + kwargs.get("qkv_layer") is not None, + kwargs.get("kv_layer") is not None, + kwargs.get("qkv_interleave_dim", None), + ) + ) + return orig(self, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", spy) + + +def _strip_packed_dpa(monkeypatch): + """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" + orig = DotProductAttention.forward + + def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): + qkv = kwargs.pop("qkv_layer", None) + kv = kwargs.pop("kv_layer", None) + dim = kwargs.pop("qkv_interleave_dim", -3) + if qkv is not None: + query_layer, key_layer, value_layer = ( + qkv.select(dim, i).contiguous() for i in range(3) + ) + elif kv is not None: + key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) + return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", stripped) + + +def _run_mha(mha, x, encoder_output=None): + dpa_module._attention_backends["backend_selection_requires_update"] = True + if encoder_output is not None: + out = mha(x, encoder_output=encoder_output) + else: + out = mha(x) + out.backward(torch.ones_like(out)) + wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] + xgrad = x.grad.clone() + x.grad = None + mha.zero_grad(set_to_none=True) + return out, xgrad, wgrads + + +def _assert_mha_equal(result, reference): + out, xgrad, wgrads = result + out_ref, xgrad_ref, wgrads_ref = reference + assert torch.equal(out, out_ref), "output differs" + assert torch.equal(xgrad, xgrad_ref), "input grad differs" + assert len(wgrads) == len(wgrads_ref) + for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): + assert torch.equal(w, w_ref), f"weight grad {i} differs" + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): + """MHA self-attention passes its packed projection output straight to DPA as + qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with + packed inputs converted back to separate contiguous q/k/v.""" + _force_backend(monkeypatch, "fused") + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x) + assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x) + _assert_mha_equal(result, reference) + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): + """MHA cross-attention passes its packed KV projection output to DPA as + kv_layer, bit-exact vs the separate contiguous reference.""" + _force_backend(monkeypatch, "fused") + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + attention_type="cross", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + enc = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda") + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x, encoder_output=enc) + assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x, encoder_output=enc) + _assert_mha_equal(result, reference) + + +@requires_fused +def test_mha_gqa_falls_back_to_views(monkeypatch): + """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy + sliced-views path and still work.""" + _force_backend(monkeypatch, "fused") + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + num_gqa_groups=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + + record = [] + _spy_packed_dpa(monkeypatch, record) + out, _, _ = _run_mha(mha, x) + assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" + assert out.shape == (_S, _B, hidden) + + +@requires_fused +def test_mha_rope_falls_back_to_views(monkeypatch): + """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" + _force_backend(monkeypatch, "fused") + from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding + + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_DTYPE, + device="cuda", + ) + rope = RotaryPositionEmbedding(_D)(max_seq_len=_S).to("cuda") + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + + record = [] + _spy_packed_dpa(monkeypatch, record) + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = mha(x, rotary_pos_emb=rope) + assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" + assert out.shape == (_S, _B, hidden) diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 70ae9dfc21..c3b32e7c03 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -9,6 +9,7 @@ import torch from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm @@ -903,6 +904,22 @@ def forward( self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + # Packed pass-through to DotProductAttention: the fused QKV/KV projection + # already produces one packed buffer, which DPA accepts directly via its + # declarative qkv_layer/kv_layer arguments (deriving q/k/v as zero-copy + # views and skipping pointer-based layout detection). Only possible when + # no per-tensor operation (RoPE, QK normalization, KV caching, CPU + # offloading) needs the individual q/k/v slices. + packed_dpa_eligible = ( + rotary_pos_emb is None + and self.q_norm is None + and inference_params is None + and not is_cpu_offload_enabled() + ) + packed_qkv_layer = None + packed_kv_layer = None + packed_interleave_dim = -3 + layernorm_output = None if self.attention_type == "self": # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn] @@ -947,28 +964,43 @@ def forward( mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - # qkv_weight_interleaved: - # [sq, b, ng, (np/ng + 2), hn] - # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] - # not qkv_weight_interleaved: - # [sq, b, (np/ng + 2), ng, hn] - # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] - query_layer, key_layer, value_layer = SplitAlongDim.apply( - mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) - ) - - if self.qkv_format == "thd": - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) - ) + if ( + num_queries_per_key_value == 1 + and packed_dpa_eligible + and not isinstance(mixed_x_layer, QuantizedTensorStorage) + ): + # np == ng: the projection output is a uniform 3-interleave + # ([.., h, 3, d] interleaved / [.., 3, h, d] otherwise), which + # DotProductAttention accepts directly as a declared packed + # qkv_layer -- no slicing here, no layout detection there. + packed_qkv_layer = mixed_x_layer + packed_interleave_dim = split_dim + query_layer = None + key_layer = None + value_layer = None else: - # query: -> [sq, b, np, hn] - # key, value: -> [sq, b, ng, hn] - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) + # qkv_weight_interleaved: + # [sq, b, ng, (np/ng + 2), hn] + # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] + # not qkv_weight_interleaved: + # [sq, b, (np/ng + 2), ng, hn] + # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] + query_layer, key_layer, value_layer = SplitAlongDim.apply( + mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) ) + + if self.qkv_format == "thd": + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) + else: + # query: -> [sq, b, np, hn] + # key, value: -> [sq, b, ng, hn] + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) elif self.attention_type == "cross": # Attention heads [sk, b, h] --> [sk, b, (ng * 2 * hn)] mixed_kv_layer = self.key_value( @@ -996,34 +1028,56 @@ def forward( mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) - # mixed_kv_layer --> 2 [sk, b, ng, hn] - key_layer, value_layer = SplitAlongDim.apply( - mixed_kv_layer, - split_dim, - mixed_kv_layer.shape[split_dim] // 2, - ) - key_layer, value_layer = ( - x.reshape( - x.size(0), - x.size(1), - -1, - self.hidden_size_per_attention_head, - ) - for x in (key_layer, value_layer) - ) - - if self.qkv_format == "thd": - key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (key_layer, value_layer) - ) + if packed_dpa_eligible and not isinstance(mixed_kv_layer, QuantizedTensorStorage): + # Declare the packed KV to DotProductAttention instead of + # slicing it: expose the 2-interleave as its own dimension. + if self.qkv_weight_interleaved: + # [.., ng, 2 * hn] --> [.., ng, 2, hn] + packed_kv_shape = mixed_kv_layer.size()[:-1] + ( + 2, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -2 + else: + # [.., 2 * ng, hn] --> [.., 2, ng, hn] + packed_kv_shape = mixed_kv_layer.size()[:-2] + ( + 2, + self.num_gqa_groups_per_partition, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -3 + packed_kv_layer = mixed_kv_layer.view(*packed_kv_shape) + key_layer = None + value_layer = None else: - # key, value: -> [sq, b, ng, hn] + # mixed_kv_layer --> 2 [sk, b, ng, hn] + key_layer, value_layer = SplitAlongDim.apply( + mixed_kv_layer, + split_dim, + mixed_kv_layer.shape[split_dim] // 2, + ) key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + x.reshape( + x.size(0), + x.size(1), + -1, + self.hidden_size_per_attention_head, + ) for x in (key_layer, value_layer) ) + if self.qkv_format == "thd": + key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + else: + # key, value: -> [sq, b, ng, hn] + key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + # Attention head [sq, b, h] --> [sq, b, hp] if self.input_layernorm: layernorm_query_outputs = self.layernorm_query( @@ -1143,6 +1197,9 @@ def forward( inference_params=inference_params, pad_between_seqs=pad_between_seqs, fp8_output=dpa_fp8_output, + qkv_layer=packed_qkv_layer, + kv_layer=packed_kv_layer, + qkv_interleave_dim=packed_interleave_dim, ) # =================== From 5713048326289dceaae168d940ec267844f330c9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 14:52:12 +0200 Subject: [PATCH 03/35] [PyTorch] DotProductAttention: factor packed-input handling into _unpack_packed_qkv Signed-off-by: Pawel Gadzinski --- .../dot_product_attention.py | 187 ++++++++++-------- 1 file changed, 109 insertions(+), 78 deletions(-) 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 e968817b97..8feca1b630 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 @@ -197,6 +197,104 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1) +def _unpack_packed_qkv( + qkv_layer: Optional[torch.Tensor], + kv_layer: Optional[torch.Tensor], + query_layer: Optional[torch.Tensor], + key_layer: Optional[torch.Tensor], + value_layer: Optional[torch.Tensor], + qkv_format: str, + qkv_interleave_dim: int, + inference_params: Optional[InferenceParams], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[str]]: + """Resolve declarative packed inputs into q/k/v. + + Derives q/k/v as zero-copy views of the packed buffer (``qkv_layer`` or + ``kv_layer``) and constructs the exact layout string from the declaration. + The layout enum is truthful by construction, so no pointer-based detection + is needed downstream (this also covers thd and FP8 DPA). + + Returns ``(query_layer, key_layer, value_layer, packed_tensor, + declared_qkv_layout)``; the last two are ``None`` when no packed input is + given. + """ + if qkv_layer is None and kv_layer is None: + if query_layer is None or key_layer is None or value_layer is None: + raise ValueError( + "query_layer, key_layer and value_layer are required unless packed" + " inputs (qkv_layer or query_layer + kv_layer) are provided." + ) + return query_layer, key_layer, value_layer, None, None + + if qkv_layer is not None and kv_layer is not None: + raise ValueError("qkv_layer and kv_layer are mutually exclusive.") + if inference_params is not None: + raise ValueError( + "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" + " (inference_params); pass separate query/key/value tensors instead." + ) + if qkv_interleave_dim not in (-3, -2): + raise ValueError( + "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" + f" {qkv_interleave_dim}." + ) + + def _packed_layout(fmt: str, num: int) -> str: + # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d + pos = len(fmt) + qkv_interleave_dim + 1 + return fmt[:pos] + str(num) + fmt[pos:] + + if qkv_layer is not None: + if any(x is not None for x in (query_layer, key_layer, value_layer)): + raise ValueError( + "qkv_layer already packs Q, K and V: query_layer, key_layer and" + " value_layer must be None when qkv_layer is provided." + ) + expected_rank = 4 if qkv_format == "thd" else 5 + if qkv_layer.dim() != expected_rank: + raise ValueError( + f"qkv_layer must be a {expected_rank}D tensor for" + f" qkv_format={qkv_format!r}, got {qkv_layer.dim()}D." + ) + if qkv_layer.shape[qkv_interleave_dim] != 3: + raise ValueError( + f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." + ) + query_layer, key_layer, value_layer = ( + qkv_layer.select(qkv_interleave_dim, i) for i in range(3) + ) + return query_layer, key_layer, value_layer, qkv_layer, _packed_layout(qkv_format, 3) + + if query_layer is None: + raise ValueError( + "kv_layer packs only K and V: query_layer is required when kv_layer" " is provided." + ) + if key_layer is not None or value_layer is not None: + raise ValueError( + "kv_layer already packs K and V: key_layer and value_layer must be" + " None when kv_layer is provided." + ) + if kv_layer.dim() != query_layer.dim() + 1: + raise ValueError( + f"kv_layer must have one more dimension than query_layer, got" + f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." + ) + if kv_layer.shape[qkv_interleave_dim] != 2: + raise ValueError( + f"kv_layer must have size 2 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." + ) + key_layer, value_layer = (kv_layer.select(qkv_interleave_dim, i) for i in range(2)) + return ( + query_layer, + key_layer, + value_layer, + kv_layer, + f"{qkv_format}_{_packed_layout(qkv_format, 2)}", + ) + + class DotProductAttention(TransformerEngineBaseModule): r"""Allows the model to jointly attend to information from different representation subspaces as described in the paper: @@ -1275,85 +1373,18 @@ def forward( since e.g. ``h == 3`` would make the shapes ambiguous. """ - # Declarative packed inputs: derive q/k/v as zero-copy views of the packed - # buffer and construct the exact layout string from the declaration. The - # layout enum is truthful by construction, so no pointer-based detection - # is needed downstream (this also covers thd and FP8 DPA). - packed_tensor = None - declared_qkv_layout = None - if qkv_layer is not None or kv_layer is not None: - if qkv_layer is not None and kv_layer is not None: - raise ValueError("qkv_layer and kv_layer are mutually exclusive.") - if inference_params is not None: - raise ValueError( - "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" - " (inference_params); pass separate query/key/value tensors instead." - ) - if qkv_interleave_dim not in (-3, -2): - raise ValueError( - "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" - f" {qkv_interleave_dim}." - ) - packed_format = qkv_format if qkv_format is not None else self.qkv_format - - def _packed_layout(fmt: str, num: int) -> str: - # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d - pos = len(fmt) + qkv_interleave_dim + 1 - return fmt[:pos] + str(num) + fmt[pos:] - - if qkv_layer is not None: - if any(x is not None for x in (query_layer, key_layer, value_layer)): - raise ValueError( - "qkv_layer already packs Q, K and V: query_layer, key_layer and" - " value_layer must be None when qkv_layer is provided." - ) - expected_rank = 4 if packed_format == "thd" else 5 - if qkv_layer.dim() != expected_rank: - raise ValueError( - f"qkv_layer must be a {expected_rank}D tensor for" - f" qkv_format={packed_format!r}, got {qkv_layer.dim()}D." - ) - if qkv_layer.shape[qkv_interleave_dim] != 3: - raise ValueError( - f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" - f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." - ) - query_layer, key_layer, value_layer = ( - qkv_layer.select(qkv_interleave_dim, i) for i in range(3) - ) - packed_tensor = qkv_layer - declared_qkv_layout = _packed_layout(packed_format, 3) - else: - if query_layer is None: - raise ValueError( - "kv_layer packs only K and V: query_layer is required when kv_layer" - " is provided." - ) - if key_layer is not None or value_layer is not None: - raise ValueError( - "kv_layer already packs K and V: key_layer and value_layer must be" - " None when kv_layer is provided." - ) - if kv_layer.dim() != query_layer.dim() + 1: - raise ValueError( - f"kv_layer must have one more dimension than query_layer, got" - f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." - ) - if kv_layer.shape[qkv_interleave_dim] != 2: - raise ValueError( - f"kv_layer must have size 2 at dim {qkv_interleave_dim}" - f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." - ) - key_layer, value_layer = ( - kv_layer.select(qkv_interleave_dim, i) for i in range(2) - ) - packed_tensor = kv_layer - declared_qkv_layout = f"{packed_format}_{_packed_layout(packed_format, 2)}" - elif query_layer is None or key_layer is None or value_layer is None: - raise ValueError( - "query_layer, key_layer and value_layer are required unless packed" - " inputs (qkv_layer or query_layer + kv_layer) are provided." + query_layer, key_layer, value_layer, packed_tensor, declared_qkv_layout = ( + _unpack_packed_qkv( + qkv_layer, + kv_layer, + query_layer, + key_layer, + value_layer, + qkv_format if qkv_format is not None else self.qkv_format, + qkv_interleave_dim, + inference_params, ) + ) with self.prepare_forward_ctx( query_layer, From 2255bcc8aa0ea17e4a5e412c57b6ff140eec4589 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:06:23 +0200 Subject: [PATCH 04/35] [PyTorch] DotProductAttention: validate packed-input last-dim stride Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 19 ++++++++++++++++--- .../dot_product_attention.py | 9 +++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py index b87e20a344..7de1d938eb 100644 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -43,10 +43,21 @@ def _fused_backend_supported(): try: q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") fused_attn_fwd( - True, _S, _S, _cu_seqlens(), _cu_seqlens(), q, q.clone(), q.clone(), _DTYPE, + True, + _S, + _S, + _cu_seqlens(), + _cu_seqlens(), + q, + q.clone(), + q.clone(), + _DTYPE, tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - dropout=0.0, qkv_layout="bshd_bshd_bshd", o_format="bshd", - attn_bias_type="no_bias", attn_mask_type="no_mask", + dropout=0.0, + qkv_layout="bshd_bshd_bshd", + o_format="bshd", + attn_bias_type="no_bias", + attn_mask_type="no_mask", ) return True except Exception: @@ -215,6 +226,8 @@ def test_dpa_packed_input_validation(): dpa(qkv_layer=qkv, kv_layer=kv) with pytest.raises(ValueError, match="must have size 3 at dim"): dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 + with pytest.raises(ValueError, match="stride 1 in its last"): + dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory with pytest.raises(ValueError, match="required unless packed"): dpa() 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 8feca1b630..88e496fa28 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 @@ -238,6 +238,15 @@ def _unpack_packed_qkv( "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" f" {qkv_interleave_dim}." ) + packed = qkv_layer if qkv_layer is not None else kv_layer + # The declared layout describes the packed buffer's memory, so it must have + # stride 1 in its last dimension (the check get_qkv_layout would otherwise + # perform on the derived q/k/v views). + if packed.stride(-1) != 1: + raise ValueError( + "The packed tensor (qkv_layer/kv_layer) must have stride 1 in its last" + f" dimension, got strides {tuple(packed.stride())}." + ) def _packed_layout(fmt: str, num: int) -> str: # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d From dc4469bf88f30f6afc2316e512dcb1954a99fc59 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:28:06 +0200 Subject: [PATCH 05/35] [PyTorch] DotProductAttention: pass packed qkv/kv buffers as separate arguments Rename the single layout-dependent packed_qkv plumbing argument (which held the full QKV buffer for *3* layouts but the KV buffer for *_2* layouts) into explicit packed_qkv/packed_kv, mirroring the public qkv_layer/kv_layer API. combine_and_quantize's combined= is split into combined_qkv=/combined_kv= accordingly, and _unpack_packed_qkv no longer returns the packed tensor since the callers already hold qkv_layer/kv_layer. Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 13 ++++--- .../dot_product_attention/backends.py | 7 +++- .../dot_product_attention.py | 38 +++++++++---------- .../attention/dot_product_attention/utils.py | 36 +++++++++--------- 4 files changed, 49 insertions(+), 45 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py index 7de1d938eb..9fc71cee18 100644 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -268,7 +268,8 @@ def fn(x): # --------------------------------------------------------------------------- -# 5. FP8 combine refactor: combined= path is bit-identical to combine_tensors path +# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to +# the combine_tensors path # --------------------------------------------------------------------------- @@ -281,15 +282,15 @@ def _fp8_quantizer(): def test_combine_and_quantize_combined_matches_views(): - """Quantizing the caller's packed buffer directly (combined=) produces the same - _data bits and scale_inv as rebuilding the packed buffer from q/k/v views via - combine_tensors (the old set_-based path).""" + """Quantizing the caller's packed buffer directly (combined_qkv=) produces the + same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views + via combine_tensors (the old set_-based path).""" torch.manual_seed(0) qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined=qkv) + new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) assert old[3] == new[3] == "bs3hd" for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): @@ -305,7 +306,7 @@ def test_combine_and_quantize_combined_kv_matches_views(): k, v = kv[:, :, 0], kv[:, :, 1] old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined=kv) + new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 50565bebf3..79fe695db2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1329,6 +1329,7 @@ def forward( layer_number, return_max_logit, packed_qkv=None, + packed_kv=None, ): # pylint: disable=missing-function-docstring @@ -1397,7 +1398,8 @@ def forward( v, QKV_quantizer, used_in_backward=is_training, - combined=packed_qkv, + combined_qkv=packed_qkv, + combined_kv=packed_kv, ) # print quantizers @@ -1905,6 +1907,7 @@ def backward(ctx, d_out, *_args): None, None, None, # packed_qkv + None, # packed_kv ) @@ -2000,6 +2003,7 @@ def forward( score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, packed_qkv: Optional[torch.Tensor] = None, + packed_kv: Optional[torch.Tensor] = None, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2220,6 +2224,7 @@ def forward( self.layer_number, self.return_max_logit, packed_qkv, + packed_kv, ) if self.return_max_logit: 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 88e496fa28..990239e388 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 @@ -206,7 +206,7 @@ def _unpack_packed_qkv( qkv_format: str, qkv_interleave_dim: int, inference_params: Optional[InferenceParams], -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[str]]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[str]]: """Resolve declarative packed inputs into q/k/v. Derives q/k/v as zero-copy views of the packed buffer (``qkv_layer`` or @@ -214,9 +214,8 @@ def _unpack_packed_qkv( The layout enum is truthful by construction, so no pointer-based detection is needed downstream (this also covers thd and FP8 DPA). - Returns ``(query_layer, key_layer, value_layer, packed_tensor, - declared_qkv_layout)``; the last two are ``None`` when no packed input is - given. + Returns ``(query_layer, key_layer, value_layer, declared_qkv_layout)``; + the layout is ``None`` when no packed input is given. """ if qkv_layer is None and kv_layer is None: if query_layer is None or key_layer is None or value_layer is None: @@ -224,7 +223,7 @@ def _unpack_packed_qkv( "query_layer, key_layer and value_layer are required unless packed" " inputs (qkv_layer or query_layer + kv_layer) are provided." ) - return query_layer, key_layer, value_layer, None, None + return query_layer, key_layer, value_layer, None if qkv_layer is not None and kv_layer is not None: raise ValueError("qkv_layer and kv_layer are mutually exclusive.") @@ -273,7 +272,7 @@ def _packed_layout(fmt: str, num: int) -> str: query_layer, key_layer, value_layer = ( qkv_layer.select(qkv_interleave_dim, i) for i in range(3) ) - return query_layer, key_layer, value_layer, qkv_layer, _packed_layout(qkv_format, 3) + return query_layer, key_layer, value_layer, _packed_layout(qkv_format, 3) if query_layer is None: raise ValueError( @@ -299,7 +298,6 @@ def _packed_layout(fmt: str, num: int) -> str: query_layer, key_layer, value_layer, - kv_layer, f"{qkv_format}_{_packed_layout(qkv_format, 2)}", ) @@ -1382,17 +1380,15 @@ def forward( since e.g. ``h == 3`` would make the shapes ambiguous. """ - query_layer, key_layer, value_layer, packed_tensor, declared_qkv_layout = ( - _unpack_packed_qkv( - qkv_layer, - kv_layer, - query_layer, - key_layer, - value_layer, - qkv_format if qkv_format is not None else self.qkv_format, - qkv_interleave_dim, - inference_params, - ) + query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( + qkv_layer, + kv_layer, + query_layer, + key_layer, + value_layer, + qkv_format if qkv_format is not None else self.qkv_format, + qkv_interleave_dim, + inference_params, ) with self.prepare_forward_ctx( @@ -1968,7 +1964,8 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, - packed_qkv=packed_tensor, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) return self.fused_attention( query_layer, @@ -2004,7 +2001,8 @@ def forward( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, - packed_qkv=packed_tensor, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 3176200ffa..e1643283d2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2783,17 +2783,17 @@ def combine_and_quantize( used_in_forward=True, used_in_backward=False, keep_same_data_and_scale_inv_format=False, - combined: Optional[torch.Tensor] = None, + combined_qkv: Optional[torch.Tensor] = None, + combined_kv: Optional[torch.Tensor] = None, ): """Combine Q, K, V tensors based on qkv_layout and quantize them together. - When ``combined`` is provided, it must be the caller's original packed buffer - matching the packed group in ``qkv_layout`` (the full QKV tensor for ``3`` - layouts such as ``bs3hd``, or the KV tensor for ``2`` layouts such as - ``bshd_bs2hd``). It is then quantized directly instead of re-deriving the - packed buffer from the q/k/v views via ``combine_tensors`` (which rebuilds it - with a raw ``set_`` under a silent adjacency/interleave assumption). Ignored - for MXFP8 quantization. + When ``combined_qkv`` (for ``3`` layouts such as ``bs3hd``) or ``combined_kv`` + (for ``2`` layouts such as ``bshd_bs2hd``) is provided, it must be the + caller's original packed buffer that q/k/v are views of. It is then quantized + directly instead of re-deriving the packed buffer from the q/k/v views via + ``combine_tensors`` (which rebuilds it with a raw ``set_`` under a silent + adjacency/interleave assumption). Ignored for MXFP8 quantization. """ if isinstance(qkv_quantizer, MXFP8Quantizer): qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) @@ -2894,24 +2894,24 @@ def combine_and_quantize( match qkv_group: case 1: dim = qkv_layout.find("3") - if combined is not None: - assert combined.shape[dim] == 3, ( - f"combined QKV tensor does not match qkv_layout {qkv_layout}: expected" - f" size 3 at dim {dim}, got shape {tuple(combined.shape)}." + if combined_qkv is not None: + assert combined_qkv.shape[dim] == 3, ( + f"combined_qkv does not match qkv_layout {qkv_layout}: expected" + f" size 3 at dim {dim}, got shape {tuple(combined_qkv.shape)}." ) - qkv = combined + qkv = combined_qkv else: qkv = combine_tensors([q, k, v], dim) qkv_fp8 = qkv_quantizer(qkv) q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True) case 2: dim = qkv_layout.split("_")[1].find("2") - if combined is not None: - assert combined.shape[dim] == 2, ( - f"combined KV tensor does not match qkv_layout {qkv_layout}: expected" - f" size 2 at dim {dim}, got shape {tuple(combined.shape)}." + if combined_kv is not None: + assert combined_kv.shape[dim] == 2, ( + f"combined_kv does not match qkv_layout {qkv_layout}: expected" + f" size 2 at dim {dim}, got shape {tuple(combined_kv.shape)}." ) - kv = combined + kv = combined_kv else: kv = combine_tensors([k, v], dim) tensors = [q, kv] From f82ba7033712e300e22048be00f10e3c0cfff1de Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:37:59 +0200 Subject: [PATCH 06/35] [PyTorch] Move packed qkv/kv input tests into test_attention.py Fold the tests from the new test_dpa_packed_inputs.py file into the existing attention test suite as a dedicated section, reusing its imports. No new test file; test logic unchanged apart from renaming the module-level constants (_B/_S/_H/_D/_DTYPE -> _PACKED_*) to avoid collisions. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 580 ++++++++++++++++++ .../attention/test_dpa_packed_inputs.py | 568 ----------------- 2 files changed, 580 insertions(+), 568 deletions(-) delete mode 100644 tests/pytorch/attention/test_dpa_packed_inputs.py diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 2dbf94fc20..aa70f52565 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -26,10 +26,13 @@ from transformer_engine.pytorch.attention.dot_product_attention import ( _attention_backends, ) +import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils, check_set_window_size, + combine_and_quantize, ) +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( @@ -3002,3 +3005,580 @@ def forward( self.quantizers, ) return out + + +# --------------------------------------------------------------------------- +# Declarative packed QKV/KV inputs (qkv_layer/kv_layer + qkv_interleave_dim) +# +# Instead of slicing a fused-projection buffer into q/k/v views (which TE then +# reverse-engineers via pointer-based layout detection), callers can pass the +# packed tensor directly. Q/K/V are derived as zero-copy views and the exact +# layout string (e.g. bs3hd) is declared, not detected -- including for thd +# and FP8 DPA. +# --------------------------------------------------------------------------- + +_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D = 2, 128, 8, 64 +_PACKED_DTYPE = torch.bfloat16 + + +def _cu_seqlens(): + return torch.arange(0, (_PACKED_B + 1) * _PACKED_S, _PACKED_S, dtype=torch.int32, device="cuda") + + +def _fused_backend_supported(): + try: + q = torch.randn( + _PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + fused_attn_fwd( + True, + _PACKED_S, + _PACKED_S, + _cu_seqlens(), + _cu_seqlens(), + q, + q.clone(), + q.clone(), + _PACKED_DTYPE, + tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + dropout=0.0, + qkv_layout="bshd_bshd_bshd", + o_format="bshd", + attn_bias_type="no_bias", + attn_mask_type="no_mask", + ) + return True + except Exception: + return False + + +requires_fused = pytest.mark.skipif( + not (torch.cuda.is_available() and _fused_backend_supported()), + reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", +) + + +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 + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + _attention_backends["backend_selection_requires_update"] = True + + +def _make_dpa(qkv_format, num_gqa_groups=None): + return DotProductAttention( + _PACKED_H, + _PACKED_D, + num_gqa_groups=num_gqa_groups, + attention_dropout=0.0, + qkv_format=qkv_format, + attn_mask_type="no_mask", + ) + + +def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): + for name, x, y in zip(names, result, reference): + assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" + + +def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): + """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" + hg = num_gqa_groups or _PACKED_H + torch.manual_seed(0) + q_shape = ( + (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) + if qkv_format == "bshd" + else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) + ) + kv_shape = ( + (_PACKED_B, _PACKED_S, hg, _PACKED_D) + if qkv_format == "bshd" + else (_PACKED_S, _PACKED_B, hg, _PACKED_D) + ) + q = torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") + k = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") + v = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") + q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) + out.backward(torch.ones_like(out)) + return out, q.grad, k.grad, v.grad + + +# --------------------------------------------------------------------------- +# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact +# --------------------------------------------------------------------------- + + +@requires_fused +@pytest.mark.parametrize( + "qkv_format, interleave_dim", + [ + pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd + pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) + pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd + ], +) +def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): + """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads + flow back into the packed tensor itself.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline(qkv_format) + + torch.manual_seed(0) + q_shape = ( + (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) + if qkv_format == "bshd" + else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) + ) + parts = [torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") for _ in range(3)] + stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d + qkv = torch.stack(parts, dim=stack_dim).requires_grad_() + + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) + out.backward(torch.ones_like(out)) + + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + grads = [qkv.grad.select(stack_dim, i) for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +@requires_fused +@pytest.mark.parametrize( + "num_gqa_groups", + [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) +) +def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): + """kv_layer packed input (with separate query) is bit-exact vs separate + contiguous q/k/v; grads flow into the packed kv tensor.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline("bshd", num_gqa_groups) + + hg = num_gqa_groups or _PACKED_H + torch.manual_seed(0) + q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + k = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + v = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + q = q.clone().requires_grad_() + kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] + + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) + out.backward(torch.ones_like(out)) + + assert kv.grad is not None and kv.grad.shape == kv.shape + _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) + + +# --------------------------------------------------------------------------- +# 2. Flash backend smoke +# --------------------------------------------------------------------------- + + +def test_dpa_flash_qkv_layer(monkeypatch): + """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" + _force_backend(monkeypatch, "flash") + try: + reference = _dpa_separate_baseline("bshd") + except Exception as exc: + pytest.skip(f"flash attention backend not available: {exc}") + + torch.manual_seed(0) + parts = [ + torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + for _ in range(3) + ] + qkv = torch.stack(parts, dim=2).requires_grad_() + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd")(qkv_layer=qkv) + out.backward(torch.ones_like(out)) + grads = [qkv.grad[:, :, i] for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +# --------------------------------------------------------------------------- +# 3. Validation errors +# --------------------------------------------------------------------------- + + +def test_dpa_packed_input_validation(): + dpa = _make_dpa("bshd") + qkv = torch.randn( + _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + kv = torch.randn( + _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + k = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + + with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): + dpa(qkv_layer=qkv, key_layer=k) + with pytest.raises(ValueError, match="query_layer is required when kv_layer"): + dpa(kv_layer=kv) + with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): + dpa(qkv_layer=qkv, qkv_interleave_dim=-1) + with pytest.raises(ValueError, match="mutually exclusive"): + dpa(qkv_layer=qkv, kv_layer=kv) + with pytest.raises(ValueError, match="must have size 3 at dim"): + dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 + with pytest.raises(ValueError, match="stride 1 in its last"): + dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory + with pytest.raises(ValueError, match="required unless packed"): + dpa() + + +# --------------------------------------------------------------------------- +# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer +# --------------------------------------------------------------------------- + + +@requires_fused +def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): + _force_backend(monkeypatch, "fused") + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + + torch.manual_seed(0) + parts = [ + torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + for _ in range(3) + ] + qkv = torch.stack(parts, dim=2).requires_grad_() + dpa = _make_dpa("bshd") + + def fn(x): + return dpa(qkv_layer=x) + + _attention_backends["backend_selection_requires_update"] = True + eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo + compiled_out = torch.compile(fn)(qkv) + compiled_out.backward(torch.ones_like(compiled_out)) + + breaks = dict(torch._dynamo.utils.counters["graph_break"]) + torch._dynamo.reset() + pointer_breaks = { + reason: count + for reason, count in breaks.items() + if "data_ptr" in reason or "UntypedStorage" in reason + } + assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" + assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" + + +# --------------------------------------------------------------------------- +# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to +# the combine_tensors path +# --------------------------------------------------------------------------- + + +def _fp8_quantizer(): + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + ) + + +def test_combine_and_quantize_combined_matches_views(): + """Quantizing the caller's packed buffer directly (combined_qkv=) produces the + same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views + via combine_tensors (the old set_-based path).""" + torch.manual_seed(0) + qkv = torch.randn( + _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + + old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) + + assert old[3] == new[3] == "bs3hd" + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +def test_combine_and_quantize_combined_kv_matches_views(): + """Same for the kv-packed (group 2) layout.""" + torch.manual_seed(0) + q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + kv = torch.randn( + _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + k, v = kv[:, :, 0], kv[:, :, 1] + + old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) + + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +# --------------------------------------------------------------------------- +# 6. thd declarative: t3hd is declared, get_qkv_layout is never called +# --------------------------------------------------------------------------- + + +def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): + """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling + get_qkv_layout; full forward+backward runs if a thd backend is available.""" + calls = [] + orig_get_qkv_layout = dpa_utils.get_qkv_layout + + def counting(*args, **kwargs): + calls.append(kwargs.get("qkv_format")) + return orig_get_qkv_layout(*args, **kwargs) + + monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) + + seen_layouts = [] + orig_get_backend = dpa_utils.get_attention_backend + + def recording(params): + seen_layouts.append(params.qkv_layout) + return orig_get_backend(params) + + monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) + + torch.manual_seed(0) + t = _PACKED_B * _PACKED_S + qkv = torch.randn( + t, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + cu = _cu_seqlens() + dpa = DotProductAttention( + _PACKED_H, _PACKED_D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" + ) + _attention_backends["backend_selection_requires_update"] = True + ran_full = True + try: + out = dpa( + qkv_layer=qkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=_PACKED_S, + max_seqlen_kv=_PACKED_S, + ) + out.backward(torch.ones_like(out)) + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + except ValueError: + # No thd-capable backend on this device; the layout step (the subject of + # this test) runs before backend dispatch, so the assertions still hold. + ran_full = False + + assert not calls, "get_qkv_layout must not be called for declared packed thd input" + assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" + if not ran_full: + pytest.skip("layout declaration verified; no thd backend available for full run") + + +# --------------------------------------------------------------------------- +# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed +# packed (qkv_layer/kv_layer) to DotProductAttention +# --------------------------------------------------------------------------- + + +def _spy_packed_dpa(monkeypatch, record): + """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" + orig = DotProductAttention.forward + + def spy(self, *args, **kwargs): + record.append( + ( + kwargs.get("qkv_layer") is not None, + kwargs.get("kv_layer") is not None, + kwargs.get("qkv_interleave_dim", None), + ) + ) + return orig(self, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", spy) + + +def _strip_packed_dpa(monkeypatch): + """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" + orig = DotProductAttention.forward + + def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): + qkv = kwargs.pop("qkv_layer", None) + kv = kwargs.pop("kv_layer", None) + dim = kwargs.pop("qkv_interleave_dim", -3) + if qkv is not None: + query_layer, key_layer, value_layer = ( + qkv.select(dim, i).contiguous() for i in range(3) + ) + elif kv is not None: + key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) + return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", stripped) + + +def _run_mha(mha, x, encoder_output=None): + _attention_backends["backend_selection_requires_update"] = True + if encoder_output is not None: + out = mha(x, encoder_output=encoder_output) + else: + out = mha(x) + out.backward(torch.ones_like(out)) + wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] + xgrad = x.grad.clone() + x.grad = None + mha.zero_grad(set_to_none=True) + return out, xgrad, wgrads + + +def _assert_mha_equal(result, reference): + out, xgrad, wgrads = result + out_ref, xgrad_ref, wgrads_ref = reference + assert torch.equal(out, out_ref), "output differs" + assert torch.equal(xgrad, xgrad_ref), "input grad differs" + assert len(wgrads) == len(wgrads_ref) + for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): + assert torch.equal(w, w_ref), f"weight grad {i} differs" + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): + """MHA self-attention passes its packed projection output straight to DPA as + qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with + packed inputs converted back to separate contiguous q/k/v.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x) + assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x) + _assert_mha_equal(result, reference) + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): + """MHA cross-attention passes its packed KV projection output to DPA as + kv_layer, bit-exact vs the separate contiguous reference.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + attention_type="cross", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + enc = torch.randn(_PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda") + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x, encoder_output=enc) + assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x, encoder_output=enc) + _assert_mha_equal(result, reference) + + +@requires_fused +def test_mha_gqa_falls_back_to_views(monkeypatch): + """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy + sliced-views path and still work.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + num_gqa_groups=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + + record = [] + _spy_packed_dpa(monkeypatch, record) + out, _, _ = _run_mha(mha, x) + assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" + assert out.shape == (_PACKED_S, _PACKED_B, hidden) + + +@requires_fused +def test_mha_rope_falls_back_to_views(monkeypatch): + """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + rope = RotaryPositionEmbedding(_PACKED_D)(max_seq_len=_PACKED_S).to("cuda") + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + + record = [] + _spy_packed_dpa(monkeypatch, record) + _attention_backends["backend_selection_requires_update"] = True + out = mha(x, rotary_pos_emb=rope) + assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" + assert out.shape == (_PACKED_S, _PACKED_B, hidden) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py deleted file mode 100644 index 9fc71cee18..0000000000 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ /dev/null @@ -1,568 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Tests for declarative packed QKV/KV inputs to DotProductAttention. - -Instead of slicing a fused-projection buffer into q/k/v views (which TE then -reverse-engineers via pointer-based layout detection), callers can pass the -packed tensor directly (``qkv_layer``/``kv_layer`` + ``qkv_interleave_dim``). -Q/K/V are derived as zero-copy views and the exact layout string (e.g. -``bs3hd``) is declared, not detected -- including for thd and FP8 DPA. -""" - -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, MultiheadAttention -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.attention.dot_product_attention.utils import ( - combine_and_quantize, -) -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( - fused_attn_fwd, -) -from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") - -_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 _fused_backend_supported(): - try: - q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - fused_attn_fwd( - True, - _S, - _S, - _cu_seqlens(), - _cu_seqlens(), - q, - q.clone(), - q.clone(), - _DTYPE, - tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - dropout=0.0, - qkv_layout="bshd_bshd_bshd", - o_format="bshd", - attn_bias_type="no_bias", - attn_mask_type="no_mask", - ) - return True - except Exception: - return False - - -requires_fused = pytest.mark.skipif( - not (torch.cuda.is_available() and _fused_backend_supported()), - reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", -) - - -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 - monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") - dpa_module._attention_backends["backend_selection_requires_update"] = True - - -def _make_dpa(qkv_format, num_gqa_groups=None): - return DotProductAttention( - _H, - _D, - num_gqa_groups=num_gqa_groups, - attention_dropout=0.0, - qkv_format=qkv_format, - attn_mask_type="no_mask", - ) - - -def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): - for name, x, y in zip(names, result, reference): - assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" - - -def _fresh_parts(qkv_format, num_heads, seed=0): - torch.manual_seed(seed) - shape = (_B, _S, num_heads, _D) if qkv_format == "bshd" else (_S, _B, num_heads, _D) - return [torch.randn(*shape, dtype=_DTYPE, device="cuda") for _ in range(3)] - - -def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): - """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" - hg = num_gqa_groups or _H - torch.manual_seed(0) - q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) - kv_shape = (_B, _S, hg, _D) if qkv_format == "bshd" else (_S, _B, hg, _D) - q = torch.randn(*q_shape, dtype=_DTYPE, device="cuda") - k = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") - v = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") - q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) - out.backward(torch.ones_like(out)) - return out, q.grad, k.grad, v.grad - - -# --------------------------------------------------------------------------- -# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact -# --------------------------------------------------------------------------- - - -@requires_fused -@pytest.mark.parametrize( - "qkv_format, interleave_dim", - [ - pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd - pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) - pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd - ], -) -def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): - """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads - flow back into the packed tensor itself.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline(qkv_format) - - torch.manual_seed(0) - q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) - parts = [torch.randn(*q_shape, dtype=_DTYPE, device="cuda") for _ in range(3)] - stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d - qkv = torch.stack(parts, dim=stack_dim).requires_grad_() - - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) - out.backward(torch.ones_like(out)) - - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - grads = [qkv.grad.select(stack_dim, i) for i in range(3)] - _assert_bit_exact((out, *grads), reference) - - -@requires_fused -@pytest.mark.parametrize( - "num_gqa_groups", - [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) -) -def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): - """kv_layer packed input (with separate query) is bit-exact vs separate - contiguous q/k/v; grads flow into the packed kv tensor.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline("bshd", num_gqa_groups) - - hg = num_gqa_groups or _H - torch.manual_seed(0) - q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - k = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") - v = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") - q = q.clone().requires_grad_() - kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] - - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) - out.backward(torch.ones_like(out)) - - assert kv.grad is not None and kv.grad.shape == kv.shape - _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) - - -# --------------------------------------------------------------------------- -# 2. Flash backend smoke -# --------------------------------------------------------------------------- - - -def test_dpa_flash_qkv_layer(monkeypatch): - """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" - _force_backend(monkeypatch, "flash") - try: - reference = _dpa_separate_baseline("bshd") - except Exception as exc: - pytest.skip(f"flash attention backend not available: {exc}") - - torch.manual_seed(0) - parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] - qkv = torch.stack(parts, dim=2).requires_grad_() - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd")(qkv_layer=qkv) - out.backward(torch.ones_like(out)) - grads = [qkv.grad[:, :, i] for i in range(3)] - _assert_bit_exact((out, *grads), reference) - - -# --------------------------------------------------------------------------- -# 3. Validation errors -# --------------------------------------------------------------------------- - - -def test_dpa_packed_input_validation(): - dpa = _make_dpa("bshd") - qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") - kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") - k = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - - with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): - dpa(qkv_layer=qkv, key_layer=k) - with pytest.raises(ValueError, match="query_layer is required when kv_layer"): - dpa(kv_layer=kv) - with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): - dpa(qkv_layer=qkv, qkv_interleave_dim=-1) - with pytest.raises(ValueError, match="mutually exclusive"): - dpa(qkv_layer=qkv, kv_layer=kv) - with pytest.raises(ValueError, match="must have size 3 at dim"): - dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 - with pytest.raises(ValueError, match="stride 1 in its last"): - dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory - with pytest.raises(ValueError, match="required unless packed"): - dpa() - - -# --------------------------------------------------------------------------- -# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer -# --------------------------------------------------------------------------- - - -@requires_fused -def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): - _force_backend(monkeypatch, "fused") - torch._dynamo.reset() - torch._dynamo.utils.counters.clear() - - torch.manual_seed(0) - parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] - qkv = torch.stack(parts, dim=2).requires_grad_() - dpa = _make_dpa("bshd") - - def fn(x): - return dpa(qkv_layer=x) - - dpa_module._attention_backends["backend_selection_requires_update"] = True - eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo - compiled_out = torch.compile(fn)(qkv) - compiled_out.backward(torch.ones_like(compiled_out)) - - breaks = dict(torch._dynamo.utils.counters["graph_break"]) - torch._dynamo.reset() - pointer_breaks = { - reason: count - for reason, count in breaks.items() - if "data_ptr" in reason or "UntypedStorage" in reason - } - assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" - assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" - - -# --------------------------------------------------------------------------- -# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to -# the combine_tensors path -# --------------------------------------------------------------------------- - - -def _fp8_quantizer(): - return Float8Quantizer( - scale=torch.ones(1, dtype=torch.float32, device="cuda"), - amax=torch.zeros(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, - ) - - -def test_combine_and_quantize_combined_matches_views(): - """Quantizing the caller's packed buffer directly (combined_qkv=) produces the - same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views - via combine_tensors (the old set_-based path).""" - torch.manual_seed(0) - qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") - q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - - old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) - - assert old[3] == new[3] == "bs3hd" - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -def test_combine_and_quantize_combined_kv_matches_views(): - """Same for the kv-packed (group 2) layout.""" - torch.manual_seed(0) - q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") - k, v = kv[:, :, 0], kv[:, :, 1] - - old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) - - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -# --------------------------------------------------------------------------- -# 6. thd declarative: t3hd is declared, get_qkv_layout is never called -# --------------------------------------------------------------------------- - - -def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): - """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling - get_qkv_layout; full forward+backward runs if a thd backend is available.""" - calls = [] - orig_get_qkv_layout = dpa_utils.get_qkv_layout - - def counting(*args, **kwargs): - calls.append(kwargs.get("qkv_format")) - return orig_get_qkv_layout(*args, **kwargs) - - monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) - - seen_layouts = [] - orig_get_backend = dpa_utils.get_attention_backend - - def recording(params): - seen_layouts.append(params.qkv_layout) - return orig_get_backend(params) - - monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) - - torch.manual_seed(0) - t = _B * _S - qkv = torch.randn(t, 3, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) - cu = _cu_seqlens() - dpa = DotProductAttention( - _H, _D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" - ) - dpa_module._attention_backends["backend_selection_requires_update"] = True - ran_full = True - try: - out = dpa( - qkv_layer=qkv, - cu_seqlens_q=cu, - cu_seqlens_kv=cu, - max_seqlen_q=_S, - max_seqlen_kv=_S, - ) - out.backward(torch.ones_like(out)) - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - except ValueError: - # No thd-capable backend on this device; the layout step (the subject of - # this test) runs before backend dispatch, so the assertions still hold. - ran_full = False - - assert not calls, "get_qkv_layout must not be called for declared packed thd input" - assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" - if not ran_full: - pytest.skip("layout declaration verified; no thd backend available for full run") - - -# --------------------------------------------------------------------------- -# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed -# packed (qkv_layer/kv_layer) to DotProductAttention -# --------------------------------------------------------------------------- - - -def _spy_packed_dpa(monkeypatch, record): - """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" - orig = DotProductAttention.forward - - def spy(self, *args, **kwargs): - record.append( - ( - kwargs.get("qkv_layer") is not None, - kwargs.get("kv_layer") is not None, - kwargs.get("qkv_interleave_dim", None), - ) - ) - return orig(self, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", spy) - - -def _strip_packed_dpa(monkeypatch): - """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" - orig = DotProductAttention.forward - - def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): - qkv = kwargs.pop("qkv_layer", None) - kv = kwargs.pop("kv_layer", None) - dim = kwargs.pop("qkv_interleave_dim", -3) - if qkv is not None: - query_layer, key_layer, value_layer = ( - qkv.select(dim, i).contiguous() for i in range(3) - ) - elif kv is not None: - key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) - return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", stripped) - - -def _run_mha(mha, x, encoder_output=None): - dpa_module._attention_backends["backend_selection_requires_update"] = True - if encoder_output is not None: - out = mha(x, encoder_output=encoder_output) - else: - out = mha(x) - out.backward(torch.ones_like(out)) - wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] - xgrad = x.grad.clone() - x.grad = None - mha.zero_grad(set_to_none=True) - return out, xgrad, wgrads - - -def _assert_mha_equal(result, reference): - out, xgrad, wgrads = result - out_ref, xgrad_ref, wgrads_ref = reference - assert torch.equal(out, out_ref), "output differs" - assert torch.equal(xgrad, xgrad_ref), "input grad differs" - assert len(wgrads) == len(wgrads_ref) - for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): - assert torch.equal(w, w_ref), f"weight grad {i} differs" - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): - """MHA self-attention passes its packed projection output straight to DPA as - qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with - packed inputs converted back to separate contiguous q/k/v.""" - _force_backend(monkeypatch, "fused") - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x) - assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x) - _assert_mha_equal(result, reference) - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): - """MHA cross-attention passes its packed KV projection output to DPA as - kv_layer, bit-exact vs the separate contiguous reference.""" - _force_backend(monkeypatch, "fused") - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - attention_type="cross", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - enc = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda") - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x, encoder_output=enc) - assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x, encoder_output=enc) - _assert_mha_equal(result, reference) - - -@requires_fused -def test_mha_gqa_falls_back_to_views(monkeypatch): - """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy - sliced-views path and still work.""" - _force_backend(monkeypatch, "fused") - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - num_gqa_groups=2, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - - record = [] - _spy_packed_dpa(monkeypatch, record) - out, _, _ = _run_mha(mha, x) - assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" - assert out.shape == (_S, _B, hidden) - - -@requires_fused -def test_mha_rope_falls_back_to_views(monkeypatch): - """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" - _force_backend(monkeypatch, "fused") - from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding - - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_DTYPE, - device="cuda", - ) - rope = RotaryPositionEmbedding(_D)(max_seq_len=_S).to("cuda") - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - - record = [] - _spy_packed_dpa(monkeypatch, record) - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = mha(x, rotary_pos_emb=rope) - assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" - assert out.shape == (_S, _B, hidden) From 5cd0351649d811421bff6f9d748f45d58732609c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:40:50 +0200 Subject: [PATCH 07/35] [PyTorch] Deprecate pointer-based detection of packed qkv layouts get_qkv_layout now emits a DeprecationWarning when it recognizes q/k/v as views of a packed buffer purely from data pointers/strides/offsets (detected *3*/*_2* layouts), pointing callers at the declarative qkv_layer/kv_layer API. Separate q/k/v tensors (hd_hd_hd layouts) never warn since there is nothing to declare. Signed-off-by: Pawel Gadzinski --- .../pytorch/attention/dot_product_attention/utils.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index e1643283d2..bf709753f4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2428,6 +2428,17 @@ def run_iteratively(q, k, v): if qkv_layout == "not_supported": raise RuntimeError("The provided qkv memory layout is not supported!") + if len(qkv_layout.split("_")) < 3: + # q/k/v were recognized as views of a packed buffer only by inspecting + # their data pointers, strides and storage offsets. + warnings.warn( + "Relying on pointer-based detection of packed q/k/v layouts" + f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" + " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" + " DotProductAttention instead.", + DeprecationWarning, + ) + if inference_params is not None and inference_params.is_paged: qkv_layout = "paged_kv_" + qkv_layout From 361a2a0f37c32a5e24ec0f512504608a65c46e81 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:55:24 +0200 Subject: [PATCH 08/35] [PyTorch] Fix implicit string concatenation lint warning Signed-off-by: Pawel Gadzinski --- .../attention/dot_product_attention/dot_product_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 990239e388..cab47964da 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 @@ -276,7 +276,7 @@ def _packed_layout(fmt: str, num: int) -> str: if query_layer is None: raise ValueError( - "kv_layer packs only K and V: query_layer is required when kv_layer" " is provided." + "kv_layer packs only K and V: query_layer is required when kv_layer is provided." ) if key_layer is not None or value_layer is not None: raise ValueError( From 627dd3fe69c179bf6b9c56b320f738b4d3cd9c01 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 16:39:51 +0200 Subject: [PATCH 09/35] [PyTorch] Trim packed-input test section to equivalence and MHA tests Drop the API-validation, torch.compile graph-break, combine_and_quantize equivalence and thd no-detection spy tests; keep the dense/flash bit-exact equivalence tests and the MHA packed pass-through/fallback tests. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 180 ---------------------- 1 file changed, 180 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index aa70f52565..845968c137 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -26,13 +26,10 @@ from transformer_engine.pytorch.attention.dot_product_attention import ( _attention_backends, ) -import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils, check_set_window_size, - combine_and_quantize, ) -from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( @@ -3202,183 +3199,6 @@ def test_dpa_flash_qkv_layer(monkeypatch): _assert_bit_exact((out, *grads), reference) -# --------------------------------------------------------------------------- -# 3. Validation errors -# --------------------------------------------------------------------------- - - -def test_dpa_packed_input_validation(): - dpa = _make_dpa("bshd") - qkv = torch.randn( - _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - kv = torch.randn( - _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - k = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - - with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): - dpa(qkv_layer=qkv, key_layer=k) - with pytest.raises(ValueError, match="query_layer is required when kv_layer"): - dpa(kv_layer=kv) - with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): - dpa(qkv_layer=qkv, qkv_interleave_dim=-1) - with pytest.raises(ValueError, match="mutually exclusive"): - dpa(qkv_layer=qkv, kv_layer=kv) - with pytest.raises(ValueError, match="must have size 3 at dim"): - dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 - with pytest.raises(ValueError, match="stride 1 in its last"): - dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory - with pytest.raises(ValueError, match="required unless packed"): - dpa() - - -# --------------------------------------------------------------------------- -# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer -# --------------------------------------------------------------------------- - - -@requires_fused -def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): - _force_backend(monkeypatch, "fused") - torch._dynamo.reset() - torch._dynamo.utils.counters.clear() - - torch.manual_seed(0) - parts = [ - torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - for _ in range(3) - ] - qkv = torch.stack(parts, dim=2).requires_grad_() - dpa = _make_dpa("bshd") - - def fn(x): - return dpa(qkv_layer=x) - - _attention_backends["backend_selection_requires_update"] = True - eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo - compiled_out = torch.compile(fn)(qkv) - compiled_out.backward(torch.ones_like(compiled_out)) - - breaks = dict(torch._dynamo.utils.counters["graph_break"]) - torch._dynamo.reset() - pointer_breaks = { - reason: count - for reason, count in breaks.items() - if "data_ptr" in reason or "UntypedStorage" in reason - } - assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" - assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" - - -# --------------------------------------------------------------------------- -# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to -# the combine_tensors path -# --------------------------------------------------------------------------- - - -def _fp8_quantizer(): - return Float8Quantizer( - scale=torch.ones(1, dtype=torch.float32, device="cuda"), - amax=torch.zeros(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, - ) - - -def test_combine_and_quantize_combined_matches_views(): - """Quantizing the caller's packed buffer directly (combined_qkv=) produces the - same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views - via combine_tensors (the old set_-based path).""" - torch.manual_seed(0) - qkv = torch.randn( - _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - - old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) - - assert old[3] == new[3] == "bs3hd" - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -def test_combine_and_quantize_combined_kv_matches_views(): - """Same for the kv-packed (group 2) layout.""" - torch.manual_seed(0) - q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - kv = torch.randn( - _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - k, v = kv[:, :, 0], kv[:, :, 1] - - old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) - - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -# --------------------------------------------------------------------------- -# 6. thd declarative: t3hd is declared, get_qkv_layout is never called -# --------------------------------------------------------------------------- - - -def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): - """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling - get_qkv_layout; full forward+backward runs if a thd backend is available.""" - calls = [] - orig_get_qkv_layout = dpa_utils.get_qkv_layout - - def counting(*args, **kwargs): - calls.append(kwargs.get("qkv_format")) - return orig_get_qkv_layout(*args, **kwargs) - - monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) - - seen_layouts = [] - orig_get_backend = dpa_utils.get_attention_backend - - def recording(params): - seen_layouts.append(params.qkv_layout) - return orig_get_backend(params) - - monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) - - torch.manual_seed(0) - t = _PACKED_B * _PACKED_S - qkv = torch.randn( - t, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - cu = _cu_seqlens() - dpa = DotProductAttention( - _PACKED_H, _PACKED_D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" - ) - _attention_backends["backend_selection_requires_update"] = True - ran_full = True - try: - out = dpa( - qkv_layer=qkv, - cu_seqlens_q=cu, - cu_seqlens_kv=cu, - max_seqlen_q=_PACKED_S, - max_seqlen_kv=_PACKED_S, - ) - out.backward(torch.ones_like(out)) - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - except ValueError: - # No thd-capable backend on this device; the layout step (the subject of - # this test) runs before backend dispatch, so the assertions still hold. - ran_full = False - - assert not calls, "get_qkv_layout must not be called for declared packed thd input" - assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" - if not ran_full: - pytest.skip("layout declaration verified; no thd backend available for full run") - - # --------------------------------------------------------------------------- # 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed # packed (qkv_layer/kv_layer) to DotProductAttention From a6c454a43275ae44ee10a6f741b02b033a6c68c8 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:22:15 +0200 Subject: [PATCH 10/35] [PyTorch] Drop MHA packed pass-through tests Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 205 ---------------------- 1 file changed, 205 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 845968c137..64548cdd19 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -3197,208 +3197,3 @@ def test_dpa_flash_qkv_layer(monkeypatch): out.backward(torch.ones_like(out)) grads = [qkv.grad[:, :, i] for i in range(3)] _assert_bit_exact((out, *grads), reference) - - -# --------------------------------------------------------------------------- -# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed -# packed (qkv_layer/kv_layer) to DotProductAttention -# --------------------------------------------------------------------------- - - -def _spy_packed_dpa(monkeypatch, record): - """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" - orig = DotProductAttention.forward - - def spy(self, *args, **kwargs): - record.append( - ( - kwargs.get("qkv_layer") is not None, - kwargs.get("kv_layer") is not None, - kwargs.get("qkv_interleave_dim", None), - ) - ) - return orig(self, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", spy) - - -def _strip_packed_dpa(monkeypatch): - """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" - orig = DotProductAttention.forward - - def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): - qkv = kwargs.pop("qkv_layer", None) - kv = kwargs.pop("kv_layer", None) - dim = kwargs.pop("qkv_interleave_dim", -3) - if qkv is not None: - query_layer, key_layer, value_layer = ( - qkv.select(dim, i).contiguous() for i in range(3) - ) - elif kv is not None: - key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) - return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", stripped) - - -def _run_mha(mha, x, encoder_output=None): - _attention_backends["backend_selection_requires_update"] = True - if encoder_output is not None: - out = mha(x, encoder_output=encoder_output) - else: - out = mha(x) - out.backward(torch.ones_like(out)) - wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] - xgrad = x.grad.clone() - x.grad = None - mha.zero_grad(set_to_none=True) - return out, xgrad, wgrads - - -def _assert_mha_equal(result, reference): - out, xgrad, wgrads = result - out_ref, xgrad_ref, wgrads_ref = reference - assert torch.equal(out, out_ref), "output differs" - assert torch.equal(xgrad, xgrad_ref), "input grad differs" - assert len(wgrads) == len(wgrads_ref) - for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): - assert torch.equal(w, w_ref), f"weight grad {i} differs" - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): - """MHA self-attention passes its packed projection output straight to DPA as - qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with - packed inputs converted back to separate contiguous q/k/v.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x) - assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x) - _assert_mha_equal(result, reference) - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): - """MHA cross-attention passes its packed KV projection output to DPA as - kv_layer, bit-exact vs the separate contiguous reference.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - attention_type="cross", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - enc = torch.randn(_PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda") - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x, encoder_output=enc) - assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x, encoder_output=enc) - _assert_mha_equal(result, reference) - - -@requires_fused -def test_mha_gqa_falls_back_to_views(monkeypatch): - """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy - sliced-views path and still work.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - num_gqa_groups=2, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - - record = [] - _spy_packed_dpa(monkeypatch, record) - out, _, _ = _run_mha(mha, x) - assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" - assert out.shape == (_PACKED_S, _PACKED_B, hidden) - - -@requires_fused -def test_mha_rope_falls_back_to_views(monkeypatch): - """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - rope = RotaryPositionEmbedding(_PACKED_D)(max_seq_len=_PACKED_S).to("cuda") - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - - record = [] - _spy_packed_dpa(monkeypatch, record) - _attention_backends["backend_selection_requires_update"] = True - out = mha(x, rotary_pos_emb=rope) - assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" - assert out.shape == (_PACKED_S, _PACKED_B, hidden) From f45338968c84e923ee87675f49d709d7cdf44505 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:30:00 +0200 Subject: [PATCH 11/35] [PyTorch] Fold packed-input tests into test_dpa_qkv_layout via a declarative param Parametrize test_dpa_qkv_layout and test_dpa_qkv_layout_thd with declarative={views,declarative}: the declarative mode passes the packed buffer to DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read off the packed buffer) instead of slicing it into q/k/v views for pointer-based detection. This reuses the whole existing config matrix (masks, bias, SWA, cross-attention, thd, all backends) for the declarative API, replacing the dedicated packed-input test section. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 298 +++++++--------------- 1 file changed, 94 insertions(+), 204 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 64548cdd19..d08c98f619 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -134,6 +134,7 @@ def test_dot_product_attention( qkv_layout, swa, pad_between_seqs, + declarative_packed=False, ): """Test DotProductAttention module""" @@ -151,6 +152,8 @@ def test_dot_product_attention( qkv_layout = "bshd_bs2hd" if not is_mla and not is_mqa_gqa else "bshd_bshd_bshd" if "3" in qkv_layout and config.attn_type == "cross": pytest.skip("No need to test this layout for cross attention") + if declarative_packed and not any(c.isdigit() for c in qkv_layout): + pytest.skip("Declarative packed inputs only apply to packed qkv layouts.") if config.window_size == (-1, -1) and swa: config.window_size = [2, 2] @@ -225,6 +228,7 @@ def test_dot_product_attention( workspace_opt, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # FlashAttention backend @@ -238,6 +242,7 @@ def test_dot_product_attention( workspace_opt, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # Compare results @@ -929,9 +934,26 @@ def test_dpa_alibi_slopes(dtype, model_configs, model): @pytest.mark.parametrize("model_configs", [model_configs_layout]) @pytest.mark.parametrize("model", model_configs_layout.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts) -def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): - """Test DotProductAttention module with different QKV layouts""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) +@pytest.mark.parametrize( + "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] +) +def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout, declarative): + """Test DotProductAttention module with different QKV layouts. + + declarative=False passes q/k/v as views of the packed buffer (pointer-based + layout detection); declarative=True passes the packed buffer itself via + qkv_layer/kv_layer (declared layout, gradients on the packed buffer).""" + test_dot_product_attention( + dtype, + model_configs, + model, + False, + True, + qkv_layout, + False, + False, + declarative_packed=declarative, + ) qkv_layouts_thd = ["t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"] @@ -998,7 +1020,10 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) @pytest.mark.parametrize("model", model_configs_layout_thd.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts_thd) -def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): +@pytest.mark.parametrize( + "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] +) +def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative): """Test DotProductAttention module with different QKV layouts""" config = model_configs[model] if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -1006,14 +1031,30 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = True") pad_between_seqs = True test_dot_product_attention( - dtype, model_configs, model, False, True, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + True, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative, ) if get_cudnn_version() >= (9, 3, 0): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = False") # cuDNN 9.3.0+ is required to run pad_between_seqs = False/True in the same run pad_between_seqs = False test_dot_product_attention( - dtype, model_configs, model, False, True, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + True, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative, ) @@ -1026,8 +1067,14 @@ def _run_dot_product_attention( workspace_opt: bool, pad_between_seqs: bool, is_training: bool, + declarative_packed: bool = False, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Run DotProductAttention module with one forward pass and one backward pass""" + """Run DotProductAttention module with one forward pass and one backward pass. + + With declarative_packed=True (packed qkv_layout only), the packed buffer is + passed to DotProductAttention directly via qkv_layer/kv_layer instead of + slicing it into q/k/v views, and input gradients are read off the packed + buffer itself.""" # Set RNG and environment varables reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" @@ -1225,6 +1272,12 @@ def _run_dot_product_attention( tensor_count = int(l) split_dim = dim break + if declarative_packed and split_dim != 0: + # The packed buffer is the autograd leaf; q/k/v below are non-leaf + # views of it, and DPA receives the buffer via qkv_layer/kv_layer. + tensor.requires_grad_() + packed_tensor = tensor + packed_interleave_dim = split_dim - tensor.dim() tensors = torch.split(tensor, 1, dim=split_dim) if split_dim != 0 else [tensor] tensors_orig = ( torch.split(tensor_orig, 1, dim=split_dim) if split_dim != 0 else [tensor_orig] @@ -1237,8 +1290,10 @@ def _run_dot_product_attention( inp.append(tensors[j]) inp_orig.append(tensors_orig[j]) for i in range(3): - inp[i].requires_grad = True - inp_orig[i].requires_grad = True + if inp[i].is_leaf: + inp[i].requires_grad = True + if inp_orig[i].is_leaf: + inp_orig[i].requires_grad = True # Create output gradient qkv_format_kv = "_".join(qkv_format) @@ -1320,10 +1375,21 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: k = inp[1] v = inp[2] d_out = out_grad + packed_kwargs = {} + if declarative_packed: + assert backend in ["FusedAttention", "FlashAttention"] + packed_kwargs["qkv_interleave_dim"] = packed_interleave_dim + if len(qkv_layout.split("_")) == 1: + packed_kwargs["qkv_layer"] = packed_tensor + q, k, v = None, None, None + else: + packed_kwargs["kv_layer"] = packed_tensor + k, v = None, None out = block( q, k, v, + **packed_kwargs, window_size=config.window_size, attention_mask=attention_mask, qkv_format=qkv_format, @@ -1353,6 +1419,25 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: out.backward(d_out) + if is_training and declarative_packed: + # Input gradients live on the packed buffer; slice them back out (and + # wrap them in .grad holders) so the cross-backend comparisons below + # stay uniform with the separate-q/k/v path. + assert packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape + + class _PackedGrad: + def __init__(self, grad): + self.grad = grad + + packed_grads = [ + _PackedGrad(packed_tensor.grad.select(packed_interleave_dim, j)) + for j in range(packed_tensor.shape[packed_interleave_dim]) + ] + if len(qkv_layout.split("_")) == 1: + q, k, v = packed_grads + else: + k, v = packed_grads + d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad @@ -3002,198 +3087,3 @@ def forward( self.quantizers, ) return out - - -# --------------------------------------------------------------------------- -# Declarative packed QKV/KV inputs (qkv_layer/kv_layer + qkv_interleave_dim) -# -# Instead of slicing a fused-projection buffer into q/k/v views (which TE then -# reverse-engineers via pointer-based layout detection), callers can pass the -# packed tensor directly. Q/K/V are derived as zero-copy views and the exact -# layout string (e.g. bs3hd) is declared, not detected -- including for thd -# and FP8 DPA. -# --------------------------------------------------------------------------- - -_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D = 2, 128, 8, 64 -_PACKED_DTYPE = torch.bfloat16 - - -def _cu_seqlens(): - return torch.arange(0, (_PACKED_B + 1) * _PACKED_S, _PACKED_S, dtype=torch.int32, device="cuda") - - -def _fused_backend_supported(): - try: - q = torch.randn( - _PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - fused_attn_fwd( - True, - _PACKED_S, - _PACKED_S, - _cu_seqlens(), - _cu_seqlens(), - q, - q.clone(), - q.clone(), - _PACKED_DTYPE, - tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - dropout=0.0, - qkv_layout="bshd_bshd_bshd", - o_format="bshd", - attn_bias_type="no_bias", - attn_mask_type="no_mask", - ) - return True - except Exception: - return False - - -requires_fused = pytest.mark.skipif( - not (torch.cuda.is_available() and _fused_backend_supported()), - reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", -) - - -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 - monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") - _attention_backends["backend_selection_requires_update"] = True - - -def _make_dpa(qkv_format, num_gqa_groups=None): - return DotProductAttention( - _PACKED_H, - _PACKED_D, - num_gqa_groups=num_gqa_groups, - attention_dropout=0.0, - qkv_format=qkv_format, - attn_mask_type="no_mask", - ) - - -def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): - for name, x, y in zip(names, result, reference): - assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" - - -def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): - """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" - hg = num_gqa_groups or _PACKED_H - torch.manual_seed(0) - q_shape = ( - (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) - if qkv_format == "bshd" - else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) - ) - kv_shape = ( - (_PACKED_B, _PACKED_S, hg, _PACKED_D) - if qkv_format == "bshd" - else (_PACKED_S, _PACKED_B, hg, _PACKED_D) - ) - q = torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") - k = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") - v = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") - q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) - out.backward(torch.ones_like(out)) - return out, q.grad, k.grad, v.grad - - -# --------------------------------------------------------------------------- -# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact -# --------------------------------------------------------------------------- - - -@requires_fused -@pytest.mark.parametrize( - "qkv_format, interleave_dim", - [ - pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd - pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) - pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd - ], -) -def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): - """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads - flow back into the packed tensor itself.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline(qkv_format) - - torch.manual_seed(0) - q_shape = ( - (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) - if qkv_format == "bshd" - else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) - ) - parts = [torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") for _ in range(3)] - stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d - qkv = torch.stack(parts, dim=stack_dim).requires_grad_() - - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) - out.backward(torch.ones_like(out)) - - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - grads = [qkv.grad.select(stack_dim, i) for i in range(3)] - _assert_bit_exact((out, *grads), reference) - - -@requires_fused -@pytest.mark.parametrize( - "num_gqa_groups", - [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) -) -def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): - """kv_layer packed input (with separate query) is bit-exact vs separate - contiguous q/k/v; grads flow into the packed kv tensor.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline("bshd", num_gqa_groups) - - hg = num_gqa_groups or _PACKED_H - torch.manual_seed(0) - q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - k = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - v = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - q = q.clone().requires_grad_() - kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] - - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) - out.backward(torch.ones_like(out)) - - assert kv.grad is not None and kv.grad.shape == kv.shape - _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) - - -# --------------------------------------------------------------------------- -# 2. Flash backend smoke -# --------------------------------------------------------------------------- - - -def test_dpa_flash_qkv_layer(monkeypatch): - """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" - _force_backend(monkeypatch, "flash") - try: - reference = _dpa_separate_baseline("bshd") - except Exception as exc: - pytest.skip(f"flash attention backend not available: {exc}") - - torch.manual_seed(0) - parts = [ - torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - for _ in range(3) - ] - qkv = torch.stack(parts, dim=2).requires_grad_() - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd")(qkv_layer=qkv) - out.backward(torch.ones_like(out)) - grads = [qkv.grad[:, :, i] for i in range(3)] - _assert_bit_exact((out, *grads), reference) From 2e0f32de143e8ed45bb1329518b73b9947ff47da Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:41:52 +0200 Subject: [PATCH 12/35] [PyTorch] Shrink declarative packed-input tests to a dedicated small matrix Revert the declarative parametrization of test_dpa_qkv_layout(_thd) (which doubled their whole config x layout product) and instead add test_dpa_qkv_layout(_thd)_declarative covering all packed layouts on a trimmed config dimension: one self-attention and one cross-attention config (kv_layer path) for dense, one config for thd. Past the input handling the backend code is identical to the views mode, so the full config matrix added no coverage. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 62 ++++++++++++++--------- 1 file changed, 39 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index d08c98f619..3569ca3cdf 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -934,25 +934,28 @@ def test_dpa_alibi_slopes(dtype, model_configs, model): @pytest.mark.parametrize("model_configs", [model_configs_layout]) @pytest.mark.parametrize("model", model_configs_layout.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts) -@pytest.mark.parametrize( - "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] -) -def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout, declarative): - """Test DotProductAttention module with different QKV layouts. +def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): + """Test DotProductAttention module with different QKV layouts""" + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + - declarative=False passes q/k/v as views of the packed buffer (pointer-based - layout detection); declarative=True passes the packed buffer itself via - qkv_layer/kv_layer (declared layout, gradients on the packed buffer).""" +qkv_layouts_packed = [l for l in qkv_layouts if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 5), reason="cuDNN 8.9.5+ is required.") +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout]) +@pytest.mark.parametrize("model", ["layout_1_1", "layout_1_2"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_packed) +def test_dpa_qkv_layout_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed inputs: the packed buffer is passed to + DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read + off the packed buffer) instead of q/k/v views + pointer-based detection. + Layout coverage is complete; the model-config dimension is trimmed to one + self-attention and one cross-attention config, since past the input + handling the backend code is identical to test_dpa_qkv_layout.""" test_dot_product_attention( - dtype, - model_configs, - model, - False, - True, - qkv_layout, - False, - False, - declarative_packed=declarative, + dtype, model_configs, model, False, True, qkv_layout, False, False, declarative_packed=True ) @@ -1020,10 +1023,7 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout, declarative): @pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) @pytest.mark.parametrize("model", model_configs_layout_thd.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts_thd) -@pytest.mark.parametrize( - "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] -) -def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative): +def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=False): """Test DotProductAttention module with different QKV layouts""" config = model_configs[model] if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -1039,7 +1039,7 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative qkv_layout, False, pad_between_seqs, - declarative_packed=declarative, + declarative_packed=declarative_packed, ) if get_cudnn_version() >= (9, 3, 0): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = False") @@ -1054,10 +1054,26 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative qkv_layout, False, pad_between_seqs, - declarative_packed=declarative, + declarative_packed=declarative_packed, ) +qkv_layouts_thd_packed = [l for l in qkv_layouts_thd if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (9, 0, 0), reason="cuDNN 9.0.0+ is required.") +@pytest.mark.skipif( + get_device_compute_capability() < (9, 0), reason="THD is only supported on Hopper+." +) +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) +@pytest.mark.parametrize("model", ["layout_0_0"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_thd_packed) +def test_dpa_qkv_layout_thd_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed thd inputs, see test_dpa_qkv_layout_declarative.""" + test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=True) + + def _run_dot_product_attention( dtype: torch.dtype, config: ModelConfig, From 36cf09e5f6bf0a879ae1fbb98845a2d8afad16dc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:52:13 +0200 Subject: [PATCH 13/35] [PyTorch] Use explicit q/k/v grad variables in the DPA test harness Replace the .grad-holder objects substituted for q/k/v in declarative packed mode with q_grad/k_grad/v_grad variables computed right after backward, used uniformly by all return paths. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 52 +++++++++++------------ 1 file changed, 24 insertions(+), 28 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 3569ca3cdf..e0b2758440 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1435,34 +1435,33 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: out.backward(d_out) - if is_training and declarative_packed: - # Input gradients live on the packed buffer; slice them back out (and - # wrap them in .grad holders) so the cross-backend comparisons below - # stay uniform with the separate-q/k/v path. - assert packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape - - class _PackedGrad: - def __init__(self, grad): - self.grad = grad - - packed_grads = [ - _PackedGrad(packed_tensor.grad.select(packed_interleave_dim, j)) - for j in range(packed_tensor.shape[packed_interleave_dim]) - ] - if len(qkv_layout.split("_")) == 1: - q, k, v = packed_grads + q_grad, k_grad, v_grad = None, None, None + if is_training: + if declarative_packed: + # Input gradients live on the packed buffer; slice them back out so + # the cross-backend comparisons below stay uniform with the + # separate-q/k/v path. + assert ( + packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape + ) + packed_grads = [ + packed_tensor.grad.select(packed_interleave_dim, j) + for j in range(packed_tensor.shape[packed_interleave_dim]) + ] + if len(qkv_layout.split("_")) == 1: + q_grad, k_grad, v_grad = packed_grads + else: + q_grad = q.grad + k_grad, v_grad = packed_grads else: - k, v = packed_grads + q_grad, k_grad, v_grad = q.grad, k.grad, v.grad d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad if backend in ["UnfusedDotProductAttention"]: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) if backend in ["FusedAttention", "FlashAttention"]: if qkv_format == "thd" and pad_between_seqs: out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) @@ -1482,13 +1481,13 @@ def __init__(self, grad): out_orig = torch.cat([out_orig, out[valid_range_q[0] : valid_range_q[1]]], dim=0) if is_training: q_grad_orig = torch.cat( - [q_grad_orig, q.grad[valid_range_q[0] : valid_range_q[1]]], dim=0 + [q_grad_orig, q_grad[valid_range_q[0] : valid_range_q[1]]], dim=0 ) k_grad_orig = torch.cat( - [k_grad_orig, k.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [k_grad_orig, k_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) v_grad_orig = torch.cat( - [v_grad_orig, v.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [v_grad_orig, v_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) if is_training: return ( @@ -1499,10 +1498,7 @@ def __init__(self, grad): else: return out_orig, max_logit, (None, None, None, d_softmax_offset) else: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) model_configs_te_layer = { From 560e24649c7cf75931c9ba3769e30c9c4d807c02 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:09:52 +0000 Subject: [PATCH 14/35] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/dot_product_attention/dot_product_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 55ee89c701..82a23ff021 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 @@ -284,7 +284,7 @@ def _packed_layout(fmt: str, num: int) -> str: ) if kv_layer.dim() != query_layer.dim() + 1: raise ValueError( - f"kv_layer must have one more dimension than query_layer, got" + "kv_layer must have one more dimension than query_layer, got" f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." ) if kv_layer.shape[qkv_interleave_dim] != 2: From 86bb5bc354e6058b43d41dbd23fc052d2716a418 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 13:56:38 +0200 Subject: [PATCH 15/35] [PyTorch] Address review: warning stacklevel, offload suppression, explicit k_norm - get_qkv_layout deprecation warning: add stacklevel=2 and skip it while CPU offloading is enabled (offloading forces MultiheadAttention onto the sliced-views fallback, so the caller has no migration option there). - MultiheadAttention: gate the packed pass-through on k_norm explicitly instead of relying on q_norm/k_norm being created together. Signed-off-by: Pawel Gadzinski --- .../attention/dot_product_attention/utils.py | 22 ++++++++++++------- .../pytorch/attention/multi_head_attention.py | 1 + 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 54ce11950e..4949641810 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -35,6 +35,7 @@ META_DP, ) from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Tensor, @@ -2434,14 +2435,19 @@ def run_iteratively(q, k, v): if len(qkv_layout.split("_")) < 3: # q/k/v were recognized as views of a packed buffer only by inspecting - # their data pointers, strides and storage offsets. - warnings.warn( - "Relying on pointer-based detection of packed q/k/v layouts" - f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" - " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" - " DotProductAttention instead.", - DeprecationWarning, - ) + # their data pointers, strides and storage offsets. Skip the nudge while + # CPU offloading is enabled: offloading forces MultiheadAttention onto + # its sliced-views fallback, so packed views reaching detection are + # expected there and the caller has no migration option. + if not is_cpu_offload_enabled(): + warnings.warn( + "Relying on pointer-based detection of packed q/k/v layouts" + f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" + " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" + " DotProductAttention instead.", + DeprecationWarning, + stacklevel=2, + ) if inference_params is not None and inference_params.is_paged: qkv_layout = "paged_kv_" + qkv_layout diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index c3b32e7c03..0b73daf667 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -913,6 +913,7 @@ def forward( packed_dpa_eligible = ( rotary_pos_emb is None and self.q_norm is None + and self.k_norm is None and inference_params is None and not is_cpu_offload_enabled() ) From b782ea6ed498410188e4f6e3881c62697f24b223 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 10 Jul 2026 10:07:12 -0700 Subject: [PATCH 16/35] TE EP integration to MoEBlock (#3116) * [JAX] Resync onto upstream PR #3036, restore TE-EP-only MoE block Reset 33 local commits onto phuong/ep-3-jax @ c34771d4 (her latest with EpConfig + EpLayerConfig API, NCCL bumped to 808d2433) and re-applied the three deltas uniquely ours: * transformer_engine/jax/moe.py: replaces upstream's multi-backend MoE block with our TE-EP-only single-custom-vjp rewrite. Adapted to her new API surface: tex.EpLayerConfig replaces tex.ep_make_handle (no more EpHandle pool/cache); 5 EP callsites rewired (cfg passed in place of handle, ep_prepare arg order swapped, top_k= dropped from ep_dispatch_bwd since it's now in cfg. * tests/jax/test_te_ep_moe.py: TE-EP MoE test (kept), with ep_bootstrap kwargs ep_size= and allow_handle_mem_reloc= dropped (no longer supported; ep_size is derived from mesh axes and the handle_mem reloc gating is gone). * tests/jax/run_te_ep_moe.sh: multi-process launcher (kept). Pre-sync state preserved at branch teddy/te_ep_integration.backup-pre-phuong-sync. EOF ) Signed-off-by: Teddy Do * tests/jax: trim TE-EP MoE suite (drop bootstrap, flax-wrapper, bias-zero) * drop ``TestZZZTeEpMoeBootstrap``: the re-bootstrap mismatch is a one-line guard in ``ep_bootstrap`` and not the MoE block's concern; exercising it from this suite also taints the per-process NCCL bootstrap cache for the rest of the file with no real upside. * drop ``TestTeEpMoEBlockFlax::test_init_apply_parity``: every config in ``_CONFIGS`` already runs ``MoEBlock`` (the Flax wrapper) end-to-end via ``test_forward`` / ``test_backward``, so this was a duplicate of ``softmax`` parity in another wrapper -- leave wrapper refactors to devs without paying for an extra CI run each time. * drop ``sigmoid-bias-zero``: with a zero-init bias buffer the routing math collapses to the no-bias case, so ``sigmoid`` already covers that numerical path. The bias-aware codepath is still exercised by ``sigmoid-bias-strong`` (non-zero bias). * refresh the module-level docstring to list intentional non-coverage so future readers don't re-add these tests. Signed-off-by: Teddy Do * jax/router: fix two bwd custom_partitioning bugs (aux-loss rank, topk closure) Two unrelated one-line bugs in the bwd custom_partitioning machinery that only surface once the MoE block's aux-loss path is lifted out of shard_map (the custom_partitioning_sharding_rule check is skipped under shard_map, which is why these never tripped before). 1. FusedMoEAuxLossBwdPrimitive.shardy_sharding_rule: ``grad_aux_loss`` is the cotangent of a scalar loss and is rank-0; declaring it with a spurious ``grad_one`` factor gave it rank-1 and tripped JAX's custom_partitioning_sharding_rule rank check at global view. Change the rule's third operand entry to empty: "const_buf_one, num_experts, grad_one -> i num_experts" -> "const_buf_one, num_experts, -> i num_experts" 2. FusedTopkWithScoreFunctionBwdPrimitive.partition: ``del result_infos, routing_map_format`` removed ``routing_map_format`` from the enclosing scope before the nested ``sharded_impl`` closure was invoked. Python closures resolve names at call time, not definition time, so when XLA finally invoked ``sharded_impl`` for the bwd partitioned impl it raised ``NameError: cannot access free variable 'routing_map_format'``. Drop ``routing_map_format`` from the ``del`` and leave a NOTE so future cleanups don't reintroduce the bug. Sibling partition methods (fwd topk, both aux-loss directions) already only ``del result_infos`` and need no change. Signed-off-by: Teddy Do * jax/ep: skip size-1 dp/fsdp axis in _ep_outer_axis A dp_resource or fsdp_resource that exists in the active mesh resource config but is sized 1 in the actual mesh would still be returned by ``_ep_outer_axis()``, pinning EP-output PartitionSpecs to a degenerate axis. JAX collapses size-1 mesh axes during lowering, which made the EP-output specs reference an axis that no longer exists at runtime -- breaking shard_map output stitching on configs where DP or FSDP is optional. Treat a size-1 axis as absent: prefer dp -> fsdp, but only when the candidate axis is actually sized > 1 in the current mesh. Falls back to the previous behaviour when no axis is configured at all. Signed-off-by: Teddy Do * jax/flax: realign _MoEBlock with post-resync moe() signature After the upstream PR #3036 resync the moe() API surface lost PermutationBackend (TE-EP is the only backend now), gate_inside_vjp (always True), and the per-call quantizer_sets knob (quantization flows through the standard TE autocast / with_quantizer_set context). It also gained apply_topk_weights_early and renamed the wrapper's private _align_size to the public align_size the test suite already uses. The Flax _MoEBlock wrapper was still passing the old kwargs, which broke every test that touched the wrapper. Wrapper changes: * drop "from ..moe import PermutationBackend" plus the dataclass field, the isinstance(..., PermutationBackend) validation in __post_init__, and the pass-through to moe(). * drop "from ..quantize import noop_quantizer_set" and the quantizer_sets=(noop, noop, noop) pass-through. * drop gate_inside_vjp=True. * rename _align_size: int = 0 -> align_size: int = 0 (matches what tests/jax/test_te_ep_moe.py already passes). * add apply_topk_weights_early: bool = False and pass it through to moe(). * refresh class docstring: drop permutation_backend / _align_size / quantizer_sets descriptions, add apply_topk_weights_early / align_size, note that quantization currently flows only through fp8_autocast. Signed-off-by: Teddy Do * jax/moe: plumb token_counts to grouped_gemm and zero 0-token wgrad slices Two correctness fixes for the TE-EP MoE custom_vjp that together let the bwd parity tests pass on 0-token-globally experts, and drop a workaround that is no longer needed. (1) Plumb per-expert padded token_counts into grouped_gemm group_sizes. NCCL EP HT dispatch lays out recv_tokens expert-major as [expert_0_padded | expert_1_padded | ... | overalloc_tail] where each per-expert block already includes the dispatch_output_per_expert_alignment zero-padding and only the trailing overalloc tail (slack between sum(token_counts) and the worst-case recv_pr) is unused. Previously _ffn_fwd_per_shard built a static local_group_sizes = jnp.full((num_local_experts,), slots_per_expert), which over-counted by the overalloc tail and forced cuBLAS to run the GEMM for every group including 0-token-routed experts. Pipe the real per-shard token_counts (1, num_local_experts) from ep_prepare through _moe_fwd_rule (added to ffn_in_specs/ffn_in_args with ep2_spec), into _ffn_fwd_per_shard as token_counts_local, and reshape into local_group_sizes for both grouped_quantize and grouped_gemm. cuBLAS now skips both 0-token experts and the trailing overalloc tail. Mirror the residual spec change on the bwd (local_group_sizes residual moves from P() to ep2_spec). (2) Per-group jnp.where zero-fill on wgrad outputs. cuBLAS grouped_gemm skips groups with size_g == 0 without zero-filling the corresponding out[g, :, :] slice (cublaslt_grouped_gemm.cu lines 2086/2096). For a shard hosting an expert that received zero tokens globally, d_wo / d_wi_combined for that expert is left uninit, which propagates NaN straight into the user's optimizer state. Add wgrad_group_active = (local_group_sizes > 0)[:, None, None] in _ffn_bwd_per_shard and apply via jnp.where on d_wo (right after the wo wgrad) and d_wi_combined (right after the fused wi_0+wi_1 wgrad). Mask shape is (num_local_experts, 1, 1) so cost is negligible. (3) Drop the lax.cond zero-init guard on r_tok in _moe_fwd_rule._body. Previously a jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like) wrapper around recv_tokens worked around tex.ep_dispatch_fwd leaving the recv buffer uninit on fully-empty-receiver ranks. With (1) in place, cuBLAS skips experts whose group_sizes == 0 and the per-row trailing tail of dispatched recv_tokens is unread by every downstream consumer (subsequent grouped_gemms read only sum(group_sizes) rows; ep_combine and ep_dispatch_bwd are handle_mem-aware). The only per-row consumer that would propagate the tail is grouped_dbias (per-row segment_sum), which only runs when has_bias=True, and that FFN bias path is currently gated upstream (cuBLAS grouped_gemm has no fused bias on Hopper yet; PR 3083 adds the pure-JAX bias add). With (2) handling the user-visible wgrad-NaN risk on 0-token experts, the lax.cond is now redundant. Replace with a NOTE pointing at the two follow-ups that would force its reintroduction: - a future caller that reads the full recv tile non-group-aware (e.g. an inspect probe), or - the FFN bias path landing, which would resurrect grouped_dbias. Also rewrite the _ffn_fwd_per_shard and _ffn_bwd_per_shard docstrings to spell out the per-row vs per-group uninit semantics so the next person debugging a NaN here has the invariants written down. Signed-off-by: Teddy Do * jax/flax,tests: rename use_bias/use_expert_bias for symmetry (PR #3116) Address jberchtold-nvidia's PR #3116 nit "rename use_bias -> use_ffn_bias and use_expert_bias -> use_expert_routing_bias". The two flags are siblings (they enable two different bias buffers) but the old names suggested ``use_bias`` was the general fallback, which wasn't the intent. The new names make the FFN-vs-routing distinction obvious from the call site. * transformer_engine/jax/flax/moe.py use_bias -> use_ffn_bias (dataclass field + branch in __call__ + docstring entry) use_expert_bias -> use_expert_routing_bias (same) * tests/jax/test_te_ep_moe.py _make_block(use_expert_bias=...) -> use_expert_routing_bias sigmoid-bias-strong config key updated _reference_kwargs_from_config now reads use_expert_routing_bias ``_MoEBlock`` is still the experimental underscore-prefixed alias (no public ``MoEBlock`` export yet), so the rename is API-safe. The pre-resync legacy tests (``test_moe_vjp.py``, ``test_multiprocess_moe_vjp.py``) are intentionally not updated -- they already reference removed APIs like ``PermutationBackend`` and need a separate post-resync cleanup pass. Signed-off-by: Teddy Do * jax/moe: address PR #3116 review feedback (hardcode align + expand inline justifications) Responds to jberchtold-nvidia's PR #3116 review threads on ``transformer_engine/jax/moe.py``. All changes are confined to a single file because each review thread targets a localized region and splitting mid-file would risk reordering bugs. Per review thread: 1. "Why do we need _with_sharding_constraint_cast_bwd? I haven't seen something like this required for our other VJPs." -- Expand the helper's docstring to spell out exactly why MoE needs it: unlike LN+MLP, the MoE bwd composes a bf16 cotangent from ep_dispatch_bwd with an fp32 cotangent from fused_topk_with_score_function_bwd (which the fwd's logits_2d -> fp32 promotion forces). Without the cast, ``d_x`` surfaces at fp32 even when ``x`` is bf16, doubling activation grad bandwidth and breaking any downstream LN bwd that pins a bf16 layout. (Review thread "Why do we need this utility function?".) 2. "Why is this dtype casting required? I don't recall us needing it for the non-MoE LNMLP block." -- Expand the comment above the bwd activation fp32 promotion to explain the MoE-specific math: LN+MLP's silu sits behind a downstream LN that absorbs the bf16 rounding error, while MoE's silu sits on the *expert* side of routing -- the bf16 rounding rides directly into expert_outputs and is summed across topk experts by ep_combine. Bf16 silu alone drifts ~1% vs fp32 silu and compounds through wo->combine into the ~1.4% per-element parity gap we measured against the pure-JAX softmax reference. Mirroring the fwd's fp32 promotion in the bwd keeps silu' in lock-step with silu. (Review thread on "# Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply".) 3. "Do we have a use-case for user-specified alignments beyond 128 currently? ... it'd make sense to instead hardcode _ALIGN_SIZE = 128 as a constant at the top of the file for now to simplify this MoEBlock API. We can always expand the API to support a user-specified align size in the future." -- Implement the suggestion. Drop ``align_size`` from ``_moe_fwd_rule`` / ``_moe_bwd_rule`` / ``_moe`` / public ``moe()``; shift the ``custom_vjp`` ``nondiff_argnums`` from ``range(9, 27)`` -> ``range(9, 26)``; replace ``effective_align = max(int(align_size), 128)`` with the new module-level ``_ALIGN_SIZE = 128`` constant. Trim the ``moe()`` docstring accordingly. (Review thread on "natural_spe = num_ep * max_tokens_per_rank".) 4. "Which axis name inputs are physical mesh axes and why can be logical axes? ... No need to make any changes for now, I just want to assess which are which and then we can discuss if it makes sense to support logical on some/all or if some are required to be physical axes." -- Add an "Axis-name parameters" section to ``moe()``'s docstring listing which kwargs are physical mesh axes (``ep_axis``, ``data_parallelism_axes`` -- they index ``Mesh.shape`` directly to compute ``num_ep`` / ``dp_size`` and to construct the ``P((dp..., ep), None, None)`` for ``jax.lax.with_sharding_constraint``) vs logical axes (``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` -- resolved via the Flax logical-axis rules). Also document why ``ep_axis`` / ``data_parallelism_axes`` are intentionally non-logical: the EP comm-group construction (``dp_color = rank // ep_size``) and the bootstrap signature check both require concrete integer sizes. (Review thread on "batch_pspec_axis = (*data_parallelism_axes, ep_axis)".) 5. "Is this NaN filtering a debugging artifact or something we need in the final version?" -- Strengthen the inline comment above ``sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, ...)`` to explicitly call this out as a CORRECTNESS REQUIREMENT, not a debugging artifact: it covers the sigmoid+K>1 underflow path where top-K sigmoid scores all round to zero and the ``weights / (weights.sum + 1e-20)`` normalisation emits NaN. Observationally the filter is a no-op on the dense unit-test distributions, but it must stay in for sparse / production routing. (Review thread on "sparse_probs = jnp.where(jnp.isnan(sparse_probs), ...).") Not addressed in this commit (intentional): * Review thread on the ``align_size: int = 0`` placeholder in ``flax/moe.py`` ("Placeholder comment for me to fix this so align_size is inferred automatically based on the recipe and doesn't need to be specified by the user"). That's jberchtold's own follow-up. * Review thread on the explicit ``tree_flatten`` / ``tree_unflatten`` on ``_Ctx`` ("better to use the ``@flax_struct.dataclass``"). Deferred to a separate, testable commit because changing a ``custom_vjp`` residual's pytree registration touches subtle ordering / None-handling semantics that warrant their own bisect surface. * Review thread on ``use_bias`` / ``use_expert_bias`` renames -- handled in the immediately preceding commit ``jax/flax,tests: rename use_bias/use_expert_bias for symmetry``. * Review thread on the ``expert_bias`` fp32 init -- already resolved during the Phuong PR #3036 resync (the redundant ``jnp.float32`` second-dtype argument on ``self.param`` was dropped; ``expert_bias`` now lives at ``self.dtype``). Signed-off-by: Teddy Do * jax/moe: strip PR-response framing from comments; drop sparse_probs NaN sanitizer * Rewrite the inline justifications added in 078a7d80 so each one reads as standalone code documentation, not as a reply to a reviewer: drop "per PR #3116 review", "review feedback", "Renamed from ... per PR ..." and similar PR/thread references from moe.py, flax/moe.py, and tests/jax/test_te_ep_moe.py. Technical content (why the fp32 promotion is needed for the MoE silu+multiply, why _with_sharding_constraint_cast_bwd exists, physical-vs-logical axis split in moe() docstring, the 128 alignment rationale) is preserved and reframed to be useful to a reader who has no PR context. * Drop the jnp.where(jnp.isnan(sparse_probs), 0, sparse_probs) guard. Tracing fused_topk_with_score_function.cu shows the kernel divides by sum_scores + 1e-20, so finite non-negative sigmoid scores cannot produce NaN here; the filter was only defense against upstream NaNs, which would mask a real regression if anything ever did start producing them. Signed-off-by: Teddy Do * jax/moe: drop fp32 island around silu+multiply (fwd, bwd, reference) The SwiGLU intermediate (activation inputs gate_proj_out/up_proj_out, silu+multiply, and activation output) was previously promoted to fp32 in _ffn_fwd_per_shard and again in _ffn_bwd_per_shard, then cast back to the wi/wo GEMM dtype. The promotion bought nothing: the activation inputs come out of the wi grouped_gemm in bf16, the activation output is consumed by the wo GEMM (or wo's quantizer for FP8/FP4) in the same dtype, and storing higher precision than either consumer is wasted bandwidth. * _ffn_fwd_per_shard: drop the .astype(jnp.float32) on gate_proj_out and up_proj_out and the trailing .astype(sorted_x.dtype). The multiply now stays in the wi GEMM output dtype end-to-end. * _ffn_bwd_per_shard: symmetric simplification. jax.vjp(act_fn, ...) runs at bf16, both d_intermediate * silu' and d_intermediate * up stay at bf16, no casts. silu' is now consistent with silu (both bf16) so the chain rule composes cleanly without the prior fp32 detour. * tests/jax/test_te_ep_moe.py::_pure_jax_moe_reference: drop the matching fp32 silu in the parity reference so the test compares bf16-vs-bf16. Parity tolerance was not loosened; expect the comparison to tighten now that both sides round silu identically. Also fix an inaccurate inline comment at the apply_topk_weights_early fwd branch: the bf16 requirement on expert_outputs is enforced by ep_bootstrap (which rejects max_token_dtype != bf16 and sizes the NCCL EP HT mega-buffer for 2-byte slots accordingly), not by a runtime assert in the combine FFI. Signed-off-by: Teddy Do * remove useless comments Signed-off-by: Teddy Do * tests/jax: remove legacy MoE VJP tests + launcher; point CI at TE-EP successor test_moe_vjp.py and test_multiprocess_moe_vjp.py both import PermutationBackend from transformer_engine.jax.moe -- an API that was removed during the Phuong PR #3036 resync. Both files have been dead-on-import ever since; the multiprocess launcher run_multiprocess_moe_vjp.sh only points at the dead test. test_te_ep_moe.py (the TE-EP-only custom_vjp suite) already covers everything the legacy files exercised that is still meaningful: fwd, bwd parity vs the pure-JAX reference, aux loss, both score functions, multi-process. The legacy parametrize axis (PermutationBackend.PURE_JAX vs TRITON) no longer exists. * Delete tests/jax/test_moe_vjp.py * Delete tests/jax/test_multiprocess_moe_vjp.py * Delete tests/jax/run_multiprocess_moe_vjp.sh * qa/L0_jax_distributed_unittest/test.sh: switch the MoE VJP distributed suite invocation from run_multiprocess_moe_vjp.sh / test_multiprocess_moe_vjp.py to run_te_ep_moe.sh / test_te_ep_moe.py. * tests/jax/conftest.py: docstring reference updated. * tests/jax/test_te_ep_moe.py: drop stale "successor to ..." aside and the "mirroring run_multiprocess_moe_vjp.sh" parenthetical. Net: -981 / +9. Signed-off-by: Teddy Do * jax/moe: swap _Ctx to @flax.struct.dataclass, drop manual pytree boilerplate Per reviewer feedback (Jaberchtold on PR #3036): the manual tree_flatten / tree_unflatten on _Ctx duplicate exactly what @flax.struct.dataclass auto-generates, and the permutation dataclasses elsewhere in this module already use flax.struct. Switching to @flax.struct.dataclass: * Removes ~75 lines of mechanical tree_flatten / tree_unflatten that have to be kept in sync with the field list by hand. * Keeps cfg as the single static field via flax.struct.field(pytree_node=False), so the fwd -> bwd boundary behavior under jax.custom_vjp is unchanged. * Drops two now-unused imports (dataclasses.dataclass, jax.tree_util.register_pytree_node_class) and adds flax.struct. Field order and the (children, aux_data) split are byte-equivalent to the previous manual implementation, so the pytree treedef seen by jax.custom_vjp is identical. Signed-off-by: Teddy Do * jax/moe: drop bwd recv_topk_weights NaN sanitizer; trust the dispatch contract Mirrors the sparse_probs NaN-sanitizer removal in fe446974: we trust ep_dispatch_fwd's contract that recv_topk_weights does not contain NaN, and would rather see NaN propagate (catching a contract violation immediately) than silently sanitize it. The mask_bool dance itself stays: ctx.expert_outputs and grad_pre_combine still carry NaN at padded slots (ep_dispatch_fwd leaves uninit memory in recv_tokens, FFN and combine_bwd propagate it), and IEEE NaN * 0 = NaN means jnp.where is structurally needed to overwrite padded positions with literal zeros before the sum reduction. What changed: * Drop `recv_w_clean = jnp.where(jnp.isnan(...), 0, ...)` and thread ctx.recv_topk_weights directly into w / mask_bool. * Replace the NaN-defensive comment block with a shorter note that explains the structural reason the mask is still needed (NaN in expert_outputs / grad_pre_combine at padded slots), without claiming anything about recv_topk_weights. Addresses Greptile P1 by removing the asymmetry (fwd had no sanitizer, bwd did) -- chosen direction is "remove the bwd sanitizer", matching the project-wide stance of trusting kernel contracts rather than papering over violations. Signed-off-by: Teddy Do * jax/moe: assert output dtype; tests cover d_x parity (dtype + values) Two related dtype-contract changes: 1. moe.py: one-line assert at the moe() return path that output.dtype == x.dtype. Cheap structural guard against any future bug that lets the public output drift wider than the user-supplied input dtype. 2. test_te_ep_moe.py: extend test_backward to also check d_x, the gradient propagated back to the previous layer in backprop. _grad_step now uses jax.grad(loss_fn, argnums=(0, 1)) and returns (grads_variables, grad_x); the reference path does the same so we can compare. d_x is checked for: * shape == x.shape * dtype == x.dtype (protects the _with_sharding_constraint_cast_bwd wrapper that casts the fp32-promoted gate path back to the primal dtype on bwd; a regression in that wrapper would silently double activation gradient bandwidth) * finiteness + non-zero * numerical parity vs the pure-JAX reference d_x Addresses jberchtold review comment on test_te_ep_moe.py:650 ("we also need to check the final propagated gradient that will be passed onto the next layer in backprop"). test_combined_loss_grads is adjusted to ``grads, _`` unpacking; it doesn't need d_x for its main+aux finiteness check. Signed-off-by: Teddy Do * tests/jax/test_te_ep_moe: strip docstring to just "what this suite covers" Drops two paragraphs whose content was agent-flavoured PR-review notes rather than user-facing test docs: * The final "FP8 / MXFP8 deferred" paragraph that referenced an internal review artifact (``.pr3036-review/INTEGRATION_DESIGN.md``) not in the repo. * The "Intentional non-coverage" section that explained which tests deliberately do not exist (no Flax-wrapper smoke, no re-bootstrap-mismatch test) and why -- exactly the kind of defensive / forward-looking justification prose CLAUDE.md says to keep out of the codebase. The remaining docstring covers what readers actually need: how to launch the suite, what each test class exercises, and a short note on the parametrize-vs-class layout. Addresses jberchtold review comment on test_te_ep_moe.py:54. Signed-off-by: Teddy Do * jax/moe: address TE EP alignment review feedback Signed-off-by: Teddy Do * jax/moe: fix early topk weighting padded-slot masking Signed-off-by: Teddy Do * jax/moe: remove unused EP mesh size Signed-off-by: Teddy Do * jax/moe: tighten TE EP recv capacity bound Signed-off-by: Teddy Do * jax/moe: simplify late TE EP weighting Signed-off-by: Teddy Do * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * jax/moe: reduce padded-slot recv weight masking Signed-off-by: Teddy Do --------- Signed-off-by: Teddy Do Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- qa/L0_jax_distributed_unittest/test.sh | 8 +- tests/jax/conftest.py | 4 +- ...ltiprocess_moe_vjp.sh => run_te_ep_moe.sh} | 48 +- tests/jax/test_moe_vjp.py | 443 --- tests/jax/test_multiprocess_moe_vjp.py | 406 --- tests/jax/test_te_ep_moe.py | 745 +++++ transformer_engine/jax/cpp_extensions/ep.py | 9 +- .../jax/cpp_extensions/router.py | 14 +- transformer_engine/jax/flax/moe.py | 67 +- transformer_engine/jax/moe.py | 2785 ++++++----------- 10 files changed, 1769 insertions(+), 2760 deletions(-) rename tests/jax/{run_multiprocess_moe_vjp.sh => run_te_ep_moe.sh} (61%) delete mode 100644 tests/jax/test_moe_vjp.py delete mode 100644 tests/jax/test_multiprocess_moe_vjp.py create mode 100644 tests/jax/test_te_ep_moe.py diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index f86cea284e..a563e6908d 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -41,12 +41,12 @@ wait TE_PATH=$TE_PATH bash $TE_PATH/examples/jax/ep/run_test_ep.sh || test_fail "run_test_ep.sh" wait -# MoE custom_vjp distributed suite. Runs one Python process per GPU -# via tests/jax/run_multiprocess_moe_vjp.sh (mirrors the pattern in +# TE-EP MoE custom_vjp distributed suite. Runs one Python process per +# GPU via tests/jax/run_te_ep_moe.sh (mirrors the pattern in # examples/jax/encoder/run_test_multiprocessing_encoder.sh). Requires # >=4 visible GPUs. -TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_multiprocess_moe_vjp.sh \ - || test_fail "test_multiprocess_moe_vjp.py" +TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ + || test_fail "test_te_ep_moe.py" # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs; # auto-skips otherwise). CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)" diff --git a/tests/jax/conftest.py b/tests/jax/conftest.py index 74cb91202c..d729bfd1c7 100644 --- a/tests/jax/conftest.py +++ b/tests/jax/conftest.py @@ -90,8 +90,8 @@ def pytest_addoption(parser): """CLI options used by multiprocess JAX tests. ``--num-process`` and ``--process-id`` let a multiprocess launcher - (see ``tests/jax/run_multiprocess_moe_vjp.sh``) fork one pytest - process per GPU and tell each child its rank, so the test module + (see ``tests/jax/run_te_ep_moe.sh``) fork one pytest process per + GPU and tell each child its rank, so the test module can call ``jax.distributed.initialize(...)`` with the right ``local_device_ids``. Both default to 0; non-multiprocess tests ignore them. diff --git a/tests/jax/run_multiprocess_moe_vjp.sh b/tests/jax/run_te_ep_moe.sh similarity index 61% rename from tests/jax/run_multiprocess_moe_vjp.sh rename to tests/jax/run_te_ep_moe.sh index 8dc1d2eb04..32d5f21956 100755 --- a/tests/jax/run_multiprocess_moe_vjp.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -3,46 +3,43 @@ # # See LICENSE for license information. # -# Multiprocess (one-GPU-per-process) launcher for the unified MoE VJP +# Multiprocess (one-GPU-per-process) launcher for the TE-EP MoE custom_vjp # test suite. Forks one pytest invocation per visible GPU, passing each -# its own --num-process=N --process-id=i, and waits for all of them. -# Each child calls jax.distributed.initialize(..., local_device_ids= -# process_id) so each Python process only sees its one GPU as a local -# device and the participating processes form a global mesh. +# its own --num-process=N --process-id=i, and waits for all of them. Each +# child calls jax.distributed.initialize(..., local_device_ids=process_id) +# so each Python process only sees its one GPU as a local device and the +# participating processes form a global (ep, fsdp) mesh. set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" -TEST_FILE="$TE_ROOT/tests/jax/test_multiprocess_moe_vjp.py" +TEST_FILE="$TE_ROOT/tests/jax/test_te_ep_moe.py" PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" if [ "$NUM_GPUS" -lt 4 ]; then - echo "[run_multiprocess_moe_vjp.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2 + echo "[run_te_ep_moe.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2 exit 1 fi export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}" export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.5}" -export MOE_VJP_COORDINATOR_ADDRESS="${MOE_VJP_COORDINATOR_ADDRESS:-127.0.0.1:13456}" +export TE_EP_MOE_COORDINATOR_ADDRESS="${TE_EP_MOE_COORDINATOR_ADDRESS:-127.0.0.1:13457}" echo "============================================================" -echo "MoE VJP MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" +echo "TE-EP MoE MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" echo " test file : $TEST_FILE" -echo " coordinator : $MOE_VJP_COORDINATOR_ADDRESS" +echo " coordinator : $TE_EP_MOE_COORDINATOR_ADDRESS" echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" echo "============================================================" -# Per-process logs. MOE_VJP_MP_LOG_DIR can be set to a host-mounted dir -# (e.g. when running inside a container that throws away /tmp on exit) -# so logs survive for postmortem inspection. Defaults to a fresh /tmp. -if [ -n "${MOE_VJP_MP_LOG_DIR:-}" ]; then - LOG_DIR="$MOE_VJP_MP_LOG_DIR" +if [ -n "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then + LOG_DIR="$TE_EP_MOE_MP_LOG_DIR" mkdir -p "$LOG_DIR" else - LOG_DIR=$(mktemp -d -t moe_vjp_mp_XXXXXX) + LOG_DIR=$(mktemp -d -t te_ep_moe_mp_XXXXXX) fi echo "Per-process logs: $LOG_DIR" @@ -63,8 +60,6 @@ cleanup() { } trap cleanup EXIT INT TERM -# Launch one pytest per GPU. Process 0 streams to stdout; others log -# only to file so the live output isn't a mosaic. for i in $(seq 0 $((NUM_GPUS - 1))); do LOG_FILE="$LOG_DIR/proc_${i}.log" PYTEST_CMD=( @@ -84,7 +79,6 @@ for i in $(seq 0 $((NUM_GPUS - 1))); do PIDS+=("$!") done -# Wait for all and collect exit codes. EXITS=() for pid in "${PIDS[@]}"; do if wait "$pid"; then @@ -94,7 +88,6 @@ for pid in "${PIDS[@]}"; do fi done -# Summary. echo echo "============================================================" echo "Per-process exit codes:" @@ -102,12 +95,9 @@ for i in "${!EXITS[@]}"; do echo " proc $i -> ${EXITS[$i]}" done -# Final pass/fail. Any non-zero in any process fails the suite, but -# we tolerate non-zero on the non-zero processes only if proc 0 -# reports PASS (this matches the encoder launcher's logic). Simplest -# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which -# the file emits via ``pytest.skip(allow_module_level=True)`` on -# pre-Blackwell GPUs) as success. Anything else is a failure. +# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which the +# file emits via pytest.skip(allow_module_level=True) on pre-Blackwell +# GPUs) as success. FAILED=0 for e in "${EXITS[@]}"; do if [ "$e" != "0" ] && [ "$e" != "5" ]; then @@ -118,14 +108,14 @@ done echo if [ "$FAILED" -eq 0 ]; then - echo "[run_multiprocess_moe_vjp.sh] all processes PASSED" - if [ -z "${MOE_VJP_MP_LOG_DIR:-}" ]; then + echo "[run_te_ep_moe.sh] all processes PASSED" + if [ -z "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then rm -rf "$LOG_DIR" fi exit 0 fi -echo "[run_multiprocess_moe_vjp.sh] at least one process FAILED" +echo "[run_te_ep_moe.sh] at least one process FAILED" echo " retaining logs at $LOG_DIR for diagnosis" echo " process 0 tail:" tail -20 "$LOG_DIR/proc_0.log" 2>/dev/null || true diff --git a/tests/jax/test_moe_vjp.py b/tests/jax/test_moe_vjp.py deleted file mode 100644 index cc458d039e..0000000000 --- a/tests/jax/test_moe_vjp.py +++ /dev/null @@ -1,443 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Single-device tests for the unified MoE custom_vjp at -``transformer_engine.jax.moe.moe`` (and its Flax wrapper -``transformer_engine.jax.flax._MoEBlock``). - -Strategy --------- - -Rather than reproducing every internal kernel residual, we rely on a -single end-to-end pure-JAX *reference* implementation of the whole -MoE block (``_pure_jax_moe_reference`` below) and compare the TE -``moe(...)`` forward output AND parameter gradients against it. This -gives us coverage of: - -* the gate GEMM, -* the fused top-k routing primitive (and its bwd), -* the dispatch / per-expert FFN / combine pipeline (and their bwds - threaded through the absorbed primitives), -* the optional aux-loss path (and its bwd). - -The reference uses only ``jnp`` ops + ``jax.vjp``, so we get a -"definitive" pullback to compare against without needing the TE -primitive bwd kernels. - -Distributed (EP + FSDP) testing is intentionally NOT in this file -- -that needs a multi-device setup and lives in -``tests/jax/test_distributed_moe_vjp.py`` (follow-up). -""" - -from functools import partial -from typing import Optional, Tuple - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -from transformer_engine_jax import get_device_compute_capability -from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import PermutationBackend, moe - -# The MoE custom_vjp uses grouped GEMM, which is currently -# Blackwell-only (sm_100+). Skip the whole file on older arches. -if get_device_compute_capability(0) < 100: - pytest.skip( - "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", - allow_module_level=True, - ) - -# Parametrize values for the dispatch / combine backend. Only the -# ``triton`` variant is gated by the ``triton`` marker (so the -# ``pure_jax`` variant still runs on environments without Triton). -BACKEND_PARAMS = [ - pytest.param("pure_jax", id="pure_jax"), - pytest.param("triton", id="triton", marks=pytest.mark.triton), -] - - -# ----------------------------------------------------------------------------- -# Test config -# ----------------------------------------------------------------------------- - -DTYPE = jnp.float32 # use fp32 for tighter parity assertions -BATCH_SIZE = 2 -SEQUENCE_LENGTH = 16 -HIDDEN_SIZE = 32 -INTERMEDIATE_SIZE = 64 -NUM_EXPERTS = 8 -NUM_EXPERTS_PER_TOK = 2 - - -def _make_inputs(key: jax.Array, *, batch=BATCH_SIZE, seq=SEQUENCE_LENGTH) -> jax.Array: - return jax.random.normal(key, (batch, seq, HIDDEN_SIZE), dtype=DTYPE) - - -# ----------------------------------------------------------------------------- -# Pure-JAX reference MoE -# ----------------------------------------------------------------------------- -# -# Implements EXACTLY the same math as ``moe(...)`` for the no-EP, -# softmax-routing, no-bias, silu activation, no-quantization path. -# Returns ``(output, aux_loss_or_zero)``. Used as ground truth for both -# fwd and bwd parity. - - -@partial( - jax.jit, - static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff"), -) -def _pure_jax_moe_reference( - x: jnp.ndarray, - gate_kernel: jnp.ndarray, - wi_0: jnp.ndarray, - wi_1: jnp.ndarray, - wo: jnp.ndarray, - *, - num_experts: int, - num_experts_per_tok: int, - aux_loss_coeff: float = 0.0, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Reference no-EP MoE forward (pure JAX, no TE primitives). - - Mirrors :func:`transformer_engine.jax.moe._body_fwd` for the - PURE_JAX backend, no biases, softmax routing, silu activation, - no quantization. Linear ops only -- ``jax.vjp`` over this gives - the canonical bwd to compare against. - """ - B, S, H = x.shape - T = B * S - x_2d = x.reshape(T, H) - - # Gate - logits = x_2d @ gate_kernel # [T, E] - - # Softmax + topk (no expert_bias, no grouping, scale=1.0) - probs_full = jax.nn.softmax(logits, axis=-1) # [T, E] - # top-k by probability: - sorted_idx = jnp.argsort(probs_full, axis=-1) # ascending - selected = sorted_idx[:, -num_experts_per_tok:] # [T, K] - weights = jnp.take_along_axis(probs_full, selected, axis=-1) # [T, K] - # Normalize topk weights to sum to 1 (matches softmax->topk semantics - # of fused_topk_with_score_function with use_pre_softmax=False): - weights = weights / jnp.sum(weights, axis=-1, keepdims=True) - - # Build a sparse routing_map [T, E] with weights at selected positions - routing_weights_full = jnp.zeros_like(probs_full) - routing_weights_full = routing_weights_full.at[jnp.arange(T)[:, None], selected].set(weights) - - # Per-expert FFN: replicate each token K times, gather by expert, - # run through wi_0 / wi_1 / wo, gather back, weighted-sum. - # - # Vectorize the gather without sorting: for each (token, slot k), - # multiply the corresponding expert's FFN by routing_weights[t, k] - # and sum over experts. - # x_2d: [T, H], wi_0: [E, H, M], wi_1: [E, H, M], wo: [E, M, H] - # For each expert e: layer_w0_e = x_2d @ wi_0[e]; layer_w1_e = x_2d @ wi_1[e] - # intermediate_e = silu(layer_w0_e) * layer_w1_e - # expert_out_e = intermediate_e @ wo[e] - # output[t, h] = sum_e routing_weights_full[t, e] * expert_out_e[t, h] - layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) # [T, E, M] - layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) # [T, E, M] - intermediate = jax.nn.silu(layer_w0) * layer_w1 # [T, E, M] - expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] - output_2d = jnp.einsum("te,teh->th", routing_weights_full, expert_out) # [T, H] - output = output_2d.reshape(B, S, H) - - if aux_loss_coeff > 0.0: - # aux scores: clean per-expert softmax (compute_aux_scores=True - # kernel uses a clean softmax, no bias, scale=1, no grouping). - aux_probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1) - # tokens_per_expert from REAL routing_map (post-grouping); here - # there's no grouping so == count of non-zero positions per expert. - routing_map = (routing_weights_full > 0).astype(jnp.int32) - tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] - # aux_loss formula: (E * coeff / (k * T^2)) * sum_e - # (sum_t aux_probs[t, e]) * tokens_per_expert[e] - sum_probs_per_expert = jnp.sum(aux_probs, axis=0) # [E] - aux_loss = (num_experts * aux_loss_coeff / (num_experts_per_tok * (T**2))) * jnp.sum( - sum_probs_per_expert * tokens_per_expert.astype(jnp.float32) - ) - else: - aux_loss = jnp.zeros((), dtype=DTYPE) - - return output, aux_loss - - -# ----------------------------------------------------------------------------- -# Helpers -# ----------------------------------------------------------------------------- - - -def _init_params(key: jax.Array) -> dict: - k_g, k_w0, k_w1, k_wo = jax.random.split(key, 4) - init = jax.nn.initializers.variance_scaling(1.0, "fan_in", "truncated_normal") - return dict( - gate_kernel=init(k_g, (HIDDEN_SIZE, NUM_EXPERTS), DTYPE), - wi_0=init(k_w0, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), - wi_1=init(k_w1, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), - wo=init(k_wo, (NUM_EXPERTS, INTERMEDIATE_SIZE, HIDDEN_SIZE), DTYPE), - ) - - -@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) -def _run_te_moe( - x: jnp.ndarray, - params: dict, - *, - permutation_backend, - aux_loss_coeff: float = 0.0, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - return moe( - x, - params["gate_kernel"], - params["wi_0"], - params["wi_1"], - params["wo"], - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - activation_type="silu", - score_function="softmax", - use_pre_softmax=False, - scaling_factor=1.0, - aux_loss_coeff=aux_loss_coeff, - permutation_backend=permutation_backend, - align_size=0, - dtype=DTYPE, - ) - - -@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) -def _grads_te_main_loss(params, x, *, permutation_backend, aux_loss_coeff: float = 0.0): - """jit'd grad of ``mean(out**2)`` w.r.t. params (no aux contribution).""" - - def loss(params, x): - out, _ = _run_te_moe( - x, params, permutation_backend=permutation_backend, aux_loss_coeff=aux_loss_coeff - ) - return jnp.mean(out**2) - - return jax.grad(loss)(params, x) - - -@partial(jax.jit, static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff")) -def _grads_ref_main_loss(params, x, *, num_experts, num_experts_per_tok, aux_loss_coeff=0.0): - """jit'd grad of ``mean(out**2)`` w.r.t. params on the pure-JAX ref.""" - - def loss(params, x): - out, _ = _pure_jax_moe_reference( - x, - **params, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - aux_loss_coeff=aux_loss_coeff, - ) - return jnp.mean(out**2) - - return jax.grad(loss)(params, x) - - -@partial(jax.jit, static_argnames=("permutation_backend",)) -def _grad_te_aux_only(params, x, *, permutation_backend): - """jit'd grad of just the aux loss scalar (no main contribution).""" - - def aux_only(params, x): - _, aux = _run_te_moe( - x, params, permutation_backend=permutation_backend, aux_loss_coeff=1e-2 - ) - return aux.astype(jnp.float32) - - return jax.grad(aux_only)(params, x) - - -# ----------------------------------------------------------------------------- -# Tests -# ----------------------------------------------------------------------------- - - -class TestMoeVjpForward: - """Forward shape / finiteness / parity vs pure-JAX reference.""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_forward_shape_and_finite(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(0) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out, aux = _run_te_moe(x, params, permutation_backend=backend) - assert out.shape == x.shape - assert out.dtype == x.dtype - assert jnp.all(jnp.isfinite(out)) - assert aux is None - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_forward_parity_vs_pure_jax_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(1) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out_te, _ = _run_te_moe(x, params, permutation_backend=backend) - out_ref, _ = _pure_jax_moe_reference( - x, - **params, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - ) - # FP32, small shapes -> tight tolerance - np.testing.assert_allclose(np.array(out_te), np.array(out_ref), atol=2e-5, rtol=2e-5) - - def test_pure_jax_triton_equivalence(self): - key = jax.random.PRNGKey(2) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out_pj, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.PURE_JAX) - out_tr, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.TRITON) - np.testing.assert_allclose(np.array(out_pj), np.array(out_tr), atol=2e-5, rtol=2e-5) - - -class TestMoeVjpBackward: - """Backward parity vs pure-JAX reference (which uses ``jax.vjp`` over - plain JAX ops, giving us the canonical pullback).""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_grads_finite_and_nonzero(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(3) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - grads = _grads_te_main_loss(params, x, permutation_backend=backend) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g = grads[name] - assert jnp.all(jnp.isfinite(g)), f"{name} grad has NaN/Inf" - assert jnp.any(g != 0.0), f"{name} grad is identically zero" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_grads_match_pure_jax_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(4) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - grads_te = _grads_te_main_loss(params, x, permutation_backend=backend) - grads_ref = _grads_ref_main_loss( - params, - x, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - ) - # Loose-ish tol on grads: routing path has discrete topk so the - # softmax cotangent paths through the non-topk experts diverge - # slightly between TE (which uses the fused topk bwd) and the - # reference (which uses argsort-based take_along_axis). - # Tighter than the bf16 tests. - for name in ("wi_0", "wi_1", "wo"): - np.testing.assert_allclose( - np.array(grads_te[name]), - np.array(grads_ref[name]), - atol=5e-5, - rtol=5e-5, - err_msg=f"grad mismatch on {name}", - ) - # Gate grad has more error budget because it propagates through - # the topk derivative kernel (which differs in zero-pattern - # treatment from a plain take_along_axis). - np.testing.assert_allclose( - np.array(grads_te["gate_kernel"]), - np.array(grads_ref["gate_kernel"]), - atol=5e-4, - rtol=5e-4, - err_msg="grad mismatch on gate_kernel", - ) - - -class TestMoeVjpAuxLoss: - """Aux-loss path: forward + grad parity.""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_returned_and_finite(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(5) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - _, aux = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) - assert aux is not None - assert aux.shape == () - assert jnp.isfinite(aux) - assert jnp.abs(aux) < 1e2 - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_parity_vs_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(6) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - _, aux_te = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) - _, aux_ref = _pure_jax_moe_reference( - x, - **params, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - aux_loss_coeff=1e-2, - ) - np.testing.assert_allclose(float(aux_te), float(aux_ref), atol=1e-5, rtol=1e-5) - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_grads_propagate_to_logits(self, backend_name): - """The aux-loss bwd path must produce non-zero gate-kernel grads - when only the aux-loss scalar is differentiated (no main-output - contribution).""" - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(7) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - g_gate = _grad_te_aux_only(params, x, permutation_backend=backend)["gate_kernel"] - assert jnp.all(jnp.isfinite(g_gate)) - assert jnp.any( - g_gate != 0.0 - ), "aux_loss bwd should propagate to gate_kernel via fused_topk bwd" - - -# ----------------------------------------------------------------------------- -# Flax wrapper smoke test -# ----------------------------------------------------------------------------- - - -class TestMoEBlockFlaxWrapper: - """Sanity-check the thin Flax wrapper: forward + grad on init.""" - - def test_init_and_apply(self): - block = MoEBlock( - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - intermediate_size=INTERMEDIATE_SIZE, - permutation_backend=PermutationBackend.PURE_JAX, - dtype=DTYPE, - ) - key = jax.random.PRNGKey(8) - ki, kx = jax.random.split(key) - x = _make_inputs(kx) - variables = jax.jit(block.init)(ki, x) - out, aux = jax.jit(block.apply)(variables, x) - assert out.shape == x.shape - assert aux is None - - @jax.jit - def grad_fn(variables, x): - return jax.grad(lambda v, x: jnp.mean(block.apply(v, x)[0] ** 2))(variables, x) - - grads = grad_fn(variables, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g = grads["params"][name] - g = g.value if hasattr(g, "value") else g - assert jnp.all(jnp.isfinite(g)), f"{name} grad NaN/Inf" - assert jnp.any(g != 0.0), f"{name} grad zero" diff --git a/tests/jax/test_multiprocess_moe_vjp.py b/tests/jax/test_multiprocess_moe_vjp.py deleted file mode 100644 index 97044780f0..0000000000 --- a/tests/jax/test_multiprocess_moe_vjp.py +++ /dev/null @@ -1,406 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Multi-process (one-GPU-per-process) tests for the unified MoE custom_vjp. - -The launcher ``tests/jax/run_multiprocess_moe_vjp.sh`` forks one pytest -process per visible GPU (mirroring -``examples/jax/encoder/run_test_multiprocessing_encoder.sh``). Each -process binds to exactly one device via -``jax.distributed.initialize(..., local_device_ids=process_id)``; the -participating processes form a global mesh through JAX's distributed -runtime. - -How to run ----------- - -You typically do NOT invoke pytest on this file directly -- use the -launcher, which passes ``--num-process=N --process-id=i`` to each -forked process. Driving it directly with only one process will skip -every test because :func:`jax.distributed.initialize` requires -multiple participants. - - bash tests/jax/run_multiprocess_moe_vjp.sh - -CI invocation lives in ``qa/L0_jax_distributed_unittest/test.sh``. -""" - -import os - -# NCCL needs HBM headroom that JAX's default 90% preallocation does -# not leave. Set before any jax import below. -os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") -os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") - -import sys - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -from jax.experimental import mesh_utils -from jax.sharding import Mesh, NamedSharding, PartitionSpec as P -from flax.linen import partitioning as nn_partitioning - - -# Per-process distributed bootstrap. Each pytest invocation initializes -# JAX with exactly one local device (its assigned GPU). Once -# initialized, the four processes form one global mesh of 4 devices. -def _init_distributed(num_process: int, process_id: int) -> bool: - """Initialize jax.distributed for this pytest process. - - Returns True if initialization succeeded (i.e. this is a real - multi-process launch), False if num_process == 0 / 1 meaning the - file is being collected without a launcher and tests should be - skipped at module level. - """ - if num_process <= 1: - return False - coord = os.environ.get("MOE_VJP_COORDINATOR_ADDRESS", "127.0.0.1:1234") - jax.distributed.initialize( - coordinator_address=coord, - num_processes=num_process, - process_id=process_id, - local_device_ids=process_id, - ) - assert jax.local_device_count() == 1, "one GPU per process is the whole point" - assert ( - jax.device_count() == num_process - ), f"global device_count {jax.device_count()} != num_process {num_process}" - return True - - -# Read --num-process / --process-id BEFORE pytest collects any tests so -# we can fast-skip the whole module when not in a multiprocess launch. -def _read_mp_options(): - # Use pytest's option lookup via the request fixture isn't available - # at module top-level; parse argv ourselves the same way encoder - # test does. CLI form is e.g. "pytest ... --num-process=4 --process-id=0". - num = int(os.environ.get("MP_NUM_PROCESS", "0") or "0") - pid = int(os.environ.get("MP_PROCESS_ID", "0") or "0") - for i, a in enumerate(sys.argv): - if a.startswith("--num-process="): - num = int(a.split("=", 1)[1]) - elif a == "--num-process" and i + 1 < len(sys.argv): - num = int(sys.argv[i + 1]) - elif a.startswith("--process-id="): - pid = int(a.split("=", 1)[1]) - elif a == "--process-id" and i + 1 < len(sys.argv): - pid = int(sys.argv[i + 1]) - return num, pid - - -_MP_NUM_PROCESS, _MP_PROCESS_ID = _read_mp_options() -_MP_ACTIVE = _init_distributed(_MP_NUM_PROCESS, _MP_PROCESS_ID) - -if not _MP_ACTIVE: - # Skip the entire module if not launched via the multiprocess - # runner. Lets `pytest tests/jax/` collect this file harmlessly. - pytest.skip( - "test_multiprocess_moe_vjp.py requires the multiprocess launcher " - "(run_multiprocess_moe_vjp.sh). Skipping.", - allow_module_level=True, - ) - -from transformer_engine_jax import get_device_compute_capability - -# Grouped GEMM in the MoE custom_vjp currently requires Blackwell -# (sm_100+). Skip the whole file on older arches. -if get_device_compute_capability(0) < 100: - pytest.skip( - "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", - allow_module_level=True, - ) - -import transformer_engine.jax as te -from transformer_engine.common import recipe as te_recipe -from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import PermutationBackend -from transformer_engine.jax.sharding import MeshResource, global_shard_guard - -# Parametrize values for the dispatch / combine backend. Only the -# ``triton`` variant carries the ``triton`` marker, so the -# ``pure_jax`` variant still runs on environments without Triton. -BACKEND_PARAMS = [ - pytest.param("pure_jax", id="pure_jax"), - pytest.param("triton", id="triton", marks=pytest.mark.triton), -] - - -EP_AXIS = "ep" -FSDP_AXIS = "fsdp" -EP_SIZE = 2 -# FSDP_SIZE adapts to whatever the launcher gave us: dlcluster GB200 -# gives 4 GPUs (FSDP=2), CI B200 gives 8 GPUs (FSDP=4). Both stay -# 128-aligned for MXFP8 and divide num_experts/topk cleanly. -assert ( - jax.device_count() % EP_SIZE == 0 -), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" -FSDP_SIZE = jax.device_count() // EP_SIZE -NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE - -LOGICAL_AXIS_RULES = ( - ("exp", EP_AXIS), - ("embed", FSDP_AXIS), - ("mlp", None), - ("batch", (EP_AXIS, FSDP_AXIS)), -) - - -@pytest.fixture(scope="module") -def mesh(): - if jax.device_count() < NUM_DEVICES_REQUIRED: - pytest.skip( - f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" - f" have {jax.device_count()}" - ) - devices = mesh_utils.create_device_mesh((EP_SIZE, FSDP_SIZE)) - return Mesh(devices, axis_names=(EP_AXIS, FSDP_AXIS)) - - -# ``recipe`` parametrize values used across all tests below. ``None`` -# = plain bf16; the named recipes route through TE's autocast and -# exercise the FP8/MXFP8 quantization paths in _body_fwd/_body_bwd. -# Only recipes that work on TE Blackwell are included; older GPUs -# skip via the ``hardware_supports`` guard below. -RECIPE_NAMES = ("bf16", "MXFP8BlockScaling") - - -def _resolve_recipe(name): - """Return ``(use_fp8, recipe_instance)`` for the parametrize id.""" - if name == "bf16": - return False, None - if name == "MXFP8BlockScaling": - return True, te_recipe.MXFP8BlockScaling() - raise ValueError(f"unknown recipe name: {name!r}") - - -def _hardware_supports(recipe_name): - """Skip an FP8 recipe on GPUs that don't have the hw for it.""" - if recipe_name == "bf16": - return True - from transformer_engine_jax import get_device_compute_capability - - arch = get_device_compute_capability(0) - if recipe_name == "MXFP8BlockScaling": - return arch >= 100 - return False - - -def _autocast_ctx(recipe_name): - """Context manager that turns FP8 on for non-bf16 recipes.""" - use_fp8, recipe_inst = _resolve_recipe(recipe_name) - return te.autocast(enabled=use_fp8, recipe=recipe_inst) - - -def _tol_finite_grad(recipe_name): - """Per-recipe absolute tolerance for parity grad comparison.""" - if recipe_name == "bf16": - return 5e-2 - # MXFP8 grads carry block-scale quantization noise; loosen accordingly. - return 3e-1 - - -# ----------------------------------------------------------------------------- -# Helpers -# ----------------------------------------------------------------------------- - - -def _make_block( - *, - num_experts, - num_experts_per_tok, - intermediate_size, - permutation_backend, - aux_loss_coeff=0.0, - dtype=jnp.bfloat16, - align_size=0, -): - return MoEBlock( - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - intermediate_size=intermediate_size, - permutation_backend=permutation_backend, - data_parallelism_axes=(FSDP_AXIS,), - aux_loss_coeff=aux_loss_coeff, - dtype=dtype, - _align_size=align_size, - ) - - -def _shard_inputs(x, mesh): - return jax.lax.with_sharding_constraint( - x, NamedSharding(mesh, P((EP_AXIS, FSDP_AXIS), None, None)) - ) - - -def _init_apply(block, mesh, x, key): - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x = _shard_inputs(x, mesh) - variables = jax.jit(block.init)(key, x) - jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) - output, aux = jax.jit(block.apply)(variables, x) - jax.block_until_ready(output) - return variables, output, aux - - -def _grad_step(block, variables, mesh, x): - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x = _shard_inputs(x, mesh) - - def loss_fn(variables, x): - output, aux = block.apply(variables, x) - main = jnp.mean(output.astype(jnp.float32) ** 2) - return main + (aux.astype(jnp.float32) if aux is not None else 0.0) - - grads = jax.jit(jax.grad(loss_fn))(variables, x) - jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) - return grads - - -def _unwrap(x): - return x.value if hasattr(x, "value") else x - - -def _local_shard(x): - """Return the local (this-process) shard of a global JAX Array as numpy. - - Every assertion in this file is structural (finite-ness, non-zero, - parity within tolerance). For all of these, checking the local - shard on each process is sufficient and avoids any cross-process - collective in the test machinery. ``arr.addressable_data(0)`` - returns the local-device view of the sharded array -- with one - GPU per process there is exactly one addressable shard. - """ - return np.asarray(jax.device_get(x.addressable_data(0))) - - -# ----------------------------------------------------------------------------- -# Mixtral-style shapes, sized to fit on a single 4-GPU bf16 box (a -# 4-way data-parallel shard of a Mixtral-8 block). -# ----------------------------------------------------------------------------- - -BATCH = EP_SIZE * FSDP_SIZE * 4 # 16 on 4-GPU, 32 on 8-GPU -SEQ = 2048 -HIDDEN = 1024 -INTER = 4096 -NUM_EXPERTS = 8 -TOPK = 2 - - -class TestMoeVjpMultiprocess: - """Multiprocess (one-GPU-per-process) correctness checks for the - unified MoE custom_vjp. - """ - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_fwd_and_bwd(self, mesh, backend_name, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - backend = PermutationBackend(backend_name) - block = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=backend, - ) - x = jax.random.normal( - jax.random.PRNGKey(0), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - with _autocast_ctx(recipe_name): - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) - # Local-shard checks (see _local_shard docstring for why). - out_local = _local_shard(output) - assert output.dtype == x.dtype - assert np.all(np.isfinite(out_local)), "output has NaN/Inf" - assert aux is None - with _autocast_ctx(recipe_name): - grads = _grad_step(block, variables, mesh, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_local = _local_shard(_unwrap(grads["params"][name])) - assert np.all(np.isfinite(g_local)), f"{name} grad has NaN/Inf" - assert np.any(g_local != 0.0), f"{name} grad is identically zero" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_aux_loss(self, mesh, backend_name, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - backend = PermutationBackend(backend_name) - block = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=backend, - aux_loss_coeff=1e-2, - ) - x = jax.random.normal( - jax.random.PRNGKey(4), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - with _autocast_ctx(recipe_name): - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(5)) - out_local = _local_shard(output) - assert np.all(np.isfinite(out_local)), "output has NaN/Inf under aux" - assert aux is not None - assert aux.shape == () - aux_local = _local_shard(aux) - assert np.isfinite(aux_local), "aux is NaN/Inf" - with _autocast_ctx(recipe_name): - grads = _grad_step(block, variables, mesh, x) - g_gate_local = _local_shard(_unwrap(grads["params"]["gate_kernel"])) - assert np.all(np.isfinite(g_gate_local)), "gate grad NaN/Inf under aux" - - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_pure_jax_triton_parity(self, mesh, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - block_pj = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=PermutationBackend.PURE_JAX, - ) - block_tr = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=PermutationBackend.TRITON, - ) - x = jax.random.normal( - jax.random.PRNGKey(6), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - tol = _tol_finite_grad(recipe_name) - with _autocast_ctx(recipe_name): - variables, out_pj, _ = _init_apply(block_pj, mesh, x, jax.random.PRNGKey(7)) - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x_sh = _shard_inputs(x, mesh) - out_tr, _ = jax.jit(block_tr.apply)(variables, x_sh) - - out_pj_local = _local_shard(out_pj) - out_tr_local = _local_shard(out_tr) - diff = float(np.max(np.abs(out_pj_local - out_tr_local))) - assert diff < tol, f"forward parity breach: max_abs_diff={diff} (tol={tol})" - - with _autocast_ctx(recipe_name): - grads_pj = _grad_step(block_pj, variables, mesh, x) - grads_tr = _grad_step(block_tr, variables, mesh, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_pj = _local_shard(_unwrap(grads_pj["params"][name])) - g_tr = _local_shard(_unwrap(grads_tr["params"][name])) - d = float(np.max(np.abs(g_pj - g_tr))) - assert d < tol, f"grad parity breach on {name}: max_abs_diff={d} (tol={tol})" diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py new file mode 100644 index 0000000000..d08765e184 --- /dev/null +++ b/tests/jax/test_te_ep_moe.py @@ -0,0 +1,745 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-process (one-GPU-per-process) tests for the TE-EP MoE custom_vjp. + +The launcher ``tests/jax/run_te_ep_moe.sh`` forks one pytest process per +visible GPU. Each process binds to exactly one device via +``jax.distributed.initialize(..., local_device_ids=process_id)``; the +participating processes form a global ``(ep, fsdp)`` mesh through JAX's +distributed runtime. + +How to run +---------- + +You typically do NOT invoke pytest on this file directly -- use the +launcher, which passes ``--num-process=N --process-id=i`` to each +forked process. Driving it directly with only one process will skip +every test because :func:`jax.distributed.initialize` requires +multiple participants, and the TE EP NCCL primitives require at +least four ranks. + + bash tests/jax/run_te_ep_moe.sh + +What this suite covers +---------------------- + +Each test exercises one MoE-block run and bundles every check that +single run supports — shape, dtype, +finiteness AND numerical parity vs a pure-JAX reference. Variations +on the block are pytest parametrize values rather than separate test +classes: + +* ``test_forward`` covers the forward across a curated set of + configurations (softmax/sigmoid scoring, optional non-zero + expert_bias). Each config asserts shape, dtype, finiteness and + numerical parity vs the reference in one run. +* ``test_backward`` mirrors that for gradients. +* ``TestTeEpMoeAuxLoss`` covers the second return value end-to-end + (returned + parity + aux-only grad propagates to gate + combined + main+aux grads stay finite) in two consolidated tests. +""" + +import os + +os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") +os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") + +import sys +from functools import partial + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from jax.experimental import mesh_utils +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from flax.linen import partitioning as nn_partitioning + + +def _init_distributed(num_process: int, process_id: int) -> bool: + """Initialize jax.distributed for this pytest process. + + Returns True on a real multi-process launch, False otherwise so + the module can fast-skip when pytest collects it without the + launcher. + """ + if num_process <= 1: + return False + coord = os.environ.get("TE_EP_MOE_COORDINATOR_ADDRESS", "127.0.0.1:13457") + jax.distributed.initialize( + coordinator_address=coord, + num_processes=num_process, + process_id=process_id, + local_device_ids=process_id, + ) + assert jax.local_device_count() == 1, "one GPU per process is required for TE EP" + assert ( + jax.device_count() == num_process + ), f"global device_count {jax.device_count()} != num_process {num_process}" + return True + + +def _read_mp_options(): + num = int(os.environ.get("MP_NUM_PROCESS", "0") or "0") + pid = int(os.environ.get("MP_PROCESS_ID", "0") or "0") + for i, a in enumerate(sys.argv): + if a.startswith("--num-process="): + num = int(a.split("=", 1)[1]) + elif a == "--num-process" and i + 1 < len(sys.argv): + num = int(sys.argv[i + 1]) + elif a.startswith("--process-id="): + pid = int(a.split("=", 1)[1]) + elif a == "--process-id" and i + 1 < len(sys.argv): + pid = int(sys.argv[i + 1]) + return num, pid + + +_MP_NUM_PROCESS, _MP_PROCESS_ID = _read_mp_options() +_MP_ACTIVE = _init_distributed(_MP_NUM_PROCESS, _MP_PROCESS_ID) + +if not _MP_ACTIVE: + pytest.skip( + "test_te_ep_moe.py requires the multiprocess launcher (run_te_ep_moe.sh). Skipping.", + allow_module_level=True, + ) + +from transformer_engine_jax import get_device_compute_capability + +# Grouped GEMM in the MoE custom_vjp requires Blackwell (sm_100+). The +# TE EP NCCL primitives themselves need SM>=90, but the FFN body uses +# grouped_gemm, so the file as a whole gates on sm_100+. +if get_device_compute_capability(0) < 100: + pytest.skip( + "MoE TE EP tests require Blackwell (sm_100+) for grouped GEMM", + allow_module_level=True, + ) + +from transformer_engine.jax.flax import _MoEBlock as MoEBlock +from transformer_engine.jax.moe import _ALIGN_SIZE, moe, record_ep_bootstrap_signature_for_moe +from transformer_engine.jax.ep import ep_bootstrap +from transformer_engine.jax.sharding import MeshResource, global_shard_guard + + +# ----------------------------------------------------------------------------- +# Mesh / shape config +# ----------------------------------------------------------------------------- + +EP_AXIS = "ep" +FSDP_AXIS = "fsdp" +EP_SIZE = 2 +assert ( + jax.device_count() % EP_SIZE == 0 +), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" +FSDP_SIZE = jax.device_count() // EP_SIZE +NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE + +LOGICAL_AXIS_RULES = ( + ("exp", EP_AXIS), + ("embed", FSDP_AXIS), + ("mlp", None), + ("batch", (EP_AXIS, FSDP_AXIS)), +) + +# Small shapes so the parity tests stay tight on bf16. The block still +# has all four ranks participating in dispatch/combine. +DTYPE = jnp.bfloat16 +BATCH = EP_SIZE * FSDP_SIZE * 2 # 8 on 4-GPU, 16 on 8-GPU +SEQ = 32 +HIDDEN = 64 +INTER = 128 +NUM_EXPERTS = 8 +TOPK = 2 + +# bf16 grouped_gemm + softmax-topk + ep all-to-all stack drifts ~1e-1 vs a +# fp32 numpy reference. Keep these tight enough to catch real bugs but +# loose enough to absorb expected bf16 rounding. +FWD_ATOL = 5e-2 +FWD_RTOL = 5e-2 +GRAD_FFN_ATOL = 1e-1 +GRAD_FFN_RTOL = 1e-1 +GRAD_GATE_ATOL = 5e-1 +GRAD_GATE_RTOL = 5e-1 + +# Two TE EP runs that should be bitwise-equal modulo XLA fusion order +# (slot alignment rounding, etc.). +TE_TO_TE_ATOL = 5e-3 +TE_TO_TE_RTOL = 5e-3 + +# Aux loss is computed in float32 from the SAME logits as the routing +# path. Numerical drift between TE-EP and the reference is dominated by +# the bf16-rounded softmax inside the topk kernel. +AUX_ATOL = 1e-3 +AUX_RTOL = 1e-3 + + +# ----------------------------------------------------------------------------- +# Fixtures +# ----------------------------------------------------------------------------- + + +def _compute_worst_case_recv_pr(): + """Per-rank recv buffer the bootstrap must reserve. + + NCCL EP HT expert-major uses one flat recv buffer with variable + per-expert zones. Each non-empty expert zone is padded to + ``_ALIGN_SIZE`` slots, so the reserve must cover the worst-case + total assignments plus independent per-zone padding. + """ + num_procs = jax.device_count() + num_local_experts = NUM_EXPERTS // EP_SIZE + max_tokens_per_rank = (BATCH // num_procs) * SEQ + tokens_per_ep_group = EP_SIZE * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(TOPK, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts + aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + ) + return min(per_expert_bound, aligned_total_bound) + + +@pytest.fixture(scope="module") +def mesh(): + if jax.device_count() < NUM_DEVICES_REQUIRED: + pytest.skip( + f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" + f" have {jax.device_count()}" + ) + # ``ep`` must be the inner axis: ``ep_bootstrap`` forms NCCL EP groups + # from consecutive global ranks via ``dp_color = rank // ep_size``, so + # only an (outer_fsdp, inner_ep) device layout groups ranks correctly. + devices = mesh_utils.create_device_mesh((FSDP_SIZE, EP_SIZE)) + mesh_obj = Mesh(devices, axis_names=(FSDP_AXIS, EP_AXIS)) + + num_procs = jax.process_count() + max_tokens_per_rank = (BATCH // num_procs) * SEQ + recv_capacity_per_rank = _compute_worst_case_recv_pr() + + # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather + # and cannot run from inside jax.jit. Sized to the worst-case recv_pr + # across _CONFIGS so every parametrized config is bootstrap-compatible. + with mesh_obj, global_shard_guard(MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS)): + ep_bootstrap( + world_size=num_procs, + rank=jax.process_index(), + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_capacity_per_rank, + hidden_dim=HIDDEN, + max_token_dtype=DTYPE, + ) + record_ep_bootstrap_signature_for_moe( + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_capacity_per_rank, + hidden_dim=HIDDEN, + ep_size=EP_SIZE, + ) + return mesh_obj + + +# ----------------------------------------------------------------------------- +# Pure-JAX reference MoE (no EP). Mirrors the exact math of TE's fused +# router primitive (see tests/jax/test_fused_router.py for the same +# reference applied to the standalone router kernel): +# +# softmax + post-softmax (use_pre_softmax=False, the default): +# 1. top_k by raw logits +# 2. softmax over just the K selected logits (so weights sum to 1) +# +# sigmoid + optional expert_bias: +# 1. scores = sigmoid(logits) +# 2. top_k by (scores + expert_bias) [bias only steers selection] +# 3. weights = scores at top_k positions, normalized when K > 1 +# +# Then for both: +# * weights *= scaling_factor (we leave scaling_factor=1.0 in this +# suite, matching _make_block's default). +# * per-expert FFN: silu(layer_w0) * layer_w1 → wo. +# ----------------------------------------------------------------------------- + + +@partial( + jax.jit, + static_argnames=( + "num_experts", + "num_experts_per_tok", + "aux_loss_coeff", + "score_function", + ), +) +def _pure_jax_moe_reference( + x, + gate_kernel, + wi_0, + wi_1, + wo, + expert_bias=None, + *, + num_experts, + num_experts_per_tok, + aux_loss_coeff: float = 0.0, + score_function: str = "softmax", +): + B, S, H = x.shape + T = B * S + K = num_experts_per_tok + x_2d = x.reshape(T, H) + + gate_kernel_cast = gate_kernel.astype(x.dtype) + logits = (x_2d @ gate_kernel_cast).astype(jnp.float32) # [T, E] + + if score_function == "softmax": + # use_pre_softmax=False: topk on raw logits, then softmax over K. + top_logits, top_indices = jax.lax.top_k(logits, k=K) + weights = jax.nn.softmax(top_logits, axis=-1) # [T, K], sums to 1 + elif score_function == "sigmoid": + scores = jax.nn.sigmoid(logits) # [T, E] + if expert_bias is not None and expert_bias.shape != (0,): + scores_for_routing = scores + expert_bias.astype(jnp.float32)[None, :] + _, top_indices = jax.lax.top_k(scores_for_routing, k=K) + weights = jnp.take_along_axis(scores, top_indices, axis=-1) + else: + weights, top_indices = jax.lax.top_k(scores, k=K) + # Sigmoid weights are normalized when K > 1 (matches the kernel). + if K > 1: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + else: + raise ValueError(f"Unsupported score_function={score_function!r}") + + routing_weights_full = jnp.zeros((T, num_experts), dtype=jnp.float32) + routing_weights_full = routing_weights_full.at[jnp.arange(T)[:, None], top_indices].set(weights) + + # FFN. ``apply_topk_weights_early`` is a fusion knob that doesn't + # change the math (wo is linear), so the reference is identical for + # both placements. + layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) + layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) + # Activation runs in x.dtype (typically bf16) to mirror the impl -- + # the impl keeps silu+multiply in the wi GEMM output dtype because + # storing higher precision than the consumer (wo) GEMM buys nothing. + intermediate = jax.nn.silu(layer_w0) * layer_w1 + expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] + output_2d = jnp.einsum("te,teh->th", routing_weights_full.astype(x.dtype), expert_out) + output = output_2d.reshape(B, S, H).astype(x.dtype) + + if aux_loss_coeff > 0.0: + # tex.fused_moe_aux_loss formula (matches the same + # reference_aux_loss helper from test_fused_router.py). The + # "aux scores" use the same score_function but always with + # K-normalised sigmoid (when sigmoid) / plain softmax (when + # softmax) — see tex.fused_topk_with_score_function_fwd with + # compute_aux_scores=True. + if score_function == "softmax": + aux_scores = jax.nn.softmax(logits, axis=-1) + else: # sigmoid + aux_scores = jax.nn.sigmoid(logits) + if K > 1: + aux_scores = aux_scores / (aux_scores.sum(axis=-1, keepdims=True) + 1e-20) + routing_map = (routing_weights_full > 0).astype(jnp.int32) + tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] + sum_probs_per_expert = jnp.sum(aux_scores, axis=0) # [E] + aux_loss = (num_experts * aux_loss_coeff / (K * (T**2))) * jnp.sum( + sum_probs_per_expert * tokens_per_expert.astype(jnp.float32) + ) + aux_loss = aux_loss.astype(x.dtype) + else: + aux_loss = jnp.zeros((), dtype=x.dtype) + return output, aux_loss + + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + + +def _make_block( + *, + apply_topk_weights_early=False, + aux_loss_coeff=0.0, + use_expert_routing_bias=False, + score_function="softmax", + expert_bias_init=None, +): + kwargs = dict( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + data_parallelism_axes=(FSDP_AXIS,), + apply_topk_weights_early=apply_topk_weights_early, + aux_loss_coeff=aux_loss_coeff, + use_expert_routing_bias=use_expert_routing_bias, + score_function=score_function, + dtype=DTYPE, + ) + # Custom expert_bias_init lets tests inject a non-zero expert_bias without + # poking variables['params'] post-init. + if expert_bias_init is not None: + kwargs["expert_bias_init"] = expert_bias_init + return MoEBlock(**kwargs) + + +def _strong_expert_bias_init(key, shape, dtype): + """Half +5, half -5 — large enough to force topk onto the +ve half.""" + del key + n = shape[0] + return jnp.concatenate( + [ + jnp.full((n // 2,), 5.0, dtype=dtype), + jnp.full((n - n // 2,), -5.0, dtype=dtype), + ] + ) + + +def _shard_inputs(x, mesh): + # Match the layout moe.py re-pins to: outer dp axes, then ep innermost. + return jax.lax.with_sharding_constraint( + x, NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None, None)) + ) + + +def _ctx(mesh): + """Combined mesh + global_shard_guard + axis_rules context.""" + + class _Combo: + def __enter__(self_inner): + self_inner._m = mesh.__enter__() + self_inner._gs = global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ) + self_inner._gs.__enter__() + self_inner._ar = nn_partitioning.axis_rules(LOGICAL_AXIS_RULES) + self_inner._ar.__enter__() + return self_inner._m + + def __exit__(self_inner, *args): + self_inner._ar.__exit__(*args) + self_inner._gs.__exit__(*args) + mesh.__exit__(*args) + + return _Combo() + + +def _init_apply(block, mesh, x, key): + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + variables = jax.jit(block.init)(key, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) + output, aux = jax.jit(block.apply)(variables, x_sh) + jax.block_until_ready(output) + return variables, output, aux + + +def _grad_step(block, variables, mesh, x, *, include_aux=False): + """Run jax.grad of mean(out^2) [+ aux if include_aux] vs (params, x). + + Returns ``(grads_variables, grad_x)`` so callers can check both the + weight gradients and the input-activation gradient that propagates + back to the previous layer. + """ + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + + def loss_fn(variables, x): + output, aux = block.apply(variables, x) + loss = jnp.mean(output.astype(jnp.float32) ** 2) + if include_aux and aux is not None: + loss = loss + aux.astype(jnp.float32) + return loss + + grads_v, grad_x = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))(variables, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(grads_v)[0]) + jax.block_until_ready(grad_x) + return grads_v, grad_x + + +def _grad_aux_only(block, variables, mesh, x): + """Jit'd grad of just the aux loss scalar — proves it reaches the + gate even when no main-output contribution is present.""" + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + + def aux_only(variables, x): + _, aux = block.apply(variables, x) + return aux.astype(jnp.float32) + + grads = jax.jit(jax.grad(aux_only))(variables, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) + return grads + + +def _unwrap(x): + return x.value if hasattr(x, "value") else x + + +def _to_global_numpy(arr, mesh): + """Replicate a sharded JAX array onto every rank and return as numpy. + + Triggers an all-gather inside JIT. The resulting addressable_data(0) + contains the full global array on every process, so we can run the + pure-JAX reference and compare against it from any process. + """ + rep = NamedSharding(mesh, P()) + with mesh: + full = jax.jit(lambda a: jax.lax.with_sharding_constraint(a, rep))(arr) + full.block_until_ready() + return np.asarray(jax.device_get(full.addressable_data(0))) + + +def _params_global_numpy(variables, mesh): + """Pull every entry of variables['params'] to a replicated numpy array.""" + params = variables["params"] + return {name: _to_global_numpy(_unwrap(p), mesh) for name, p in params.items()} + + +def _make_inputs(key): + """Generate a globally-identical input tensor on every process.""" + return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) + + +# ----------------------------------------------------------------------------- +# Tests +# ----------------------------------------------------------------------------- + + +# ----------------------------------------------------------------------------- +# Parametrize variants exercised by both the forward and the backward +# parity tests. Each config is one MoE-block configuration the suite +# wants covered; the test body checks shape, dtype, finiteness AND +# numerical parity vs the same pure-JAX reference (which understands +# the same set of knobs). +# ----------------------------------------------------------------------------- + +_CONFIGS = [ + pytest.param( + dict(score_function="softmax"), + id="softmax", + ), + pytest.param( + dict(score_function="softmax", apply_topk_weights_early=True), + id="softmax-early-weighting", + ), + pytest.param( + dict(score_function="sigmoid"), + id="sigmoid", + ), + # NOTE: a ``sigmoid-bias-zero`` config (use_expert_routing_bias=True + # with a zero-initialised bias buffer) was previously exercised + # here. It was dropped because the routing math collapses to the + # no-bias case when the buffer is zero -- ``sigmoid`` already + # covers that numerical path. The bias-aware codepath is still + # exercised by ``sigmoid-bias-strong`` below, which uses a + # non-zero bias. + pytest.param( + dict( + score_function="sigmoid", + use_expert_routing_bias=True, + expert_bias_init=_strong_expert_bias_init, + ), + id="sigmoid-bias-strong", + ), +] + + +def _reference_kwargs_from_config(config, params_np): + """Pick out the reference-relevant pieces of a parametrize config.""" + return dict( + score_function=config.get("score_function", "softmax"), + expert_bias=( + jnp.asarray(params_np["expert_bias"]) + if config.get("use_expert_routing_bias", False) + else None + ), + ) + + +class TestTeEpMoeForward: + """Per-config forward correctness in a single run: shape, dtype, + finiteness AND numerical parity vs the pure-JAX reference.""" + + @pytest.mark.parametrize("config", _CONFIGS) + def test_forward(self, mesh, config): + block = _make_block(**config) + x = _make_inputs(jax.random.PRNGKey(0)) + variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) + + # Shape / dtype / finiteness (cheap; on the local shard). + assert output.shape == x.shape + assert output.dtype == x.dtype + out_local = np.asarray(jax.device_get(output.addressable_data(0))) + assert np.all(np.isfinite(out_local)), "output has NaN/Inf" + assert aux is None, "aux_loss should be None when aux_loss_coeff == 0" + + # Numerical parity (replicated global view -> single rank's numpy). + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + out_te_np = _to_global_numpy(output, mesh) + + out_ref, _ = _pure_jax_moe_reference( + jnp.asarray(x_np), + jnp.asarray(params_np["gate_kernel"]), + jnp.asarray(params_np["wi_0"]), + jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wo"]), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + **_reference_kwargs_from_config(config, params_np), + ) + np.testing.assert_allclose( + out_te_np.astype(np.float32), + np.asarray(jax.device_get(out_ref)).astype(np.float32), + atol=FWD_ATOL, + rtol=FWD_RTOL, + err_msg=f"forward parity breach for config={config}", + ) + + +class TestTeEpMoeBackward: + """Per-config backward correctness in a single run: per-tensor + grads finite, non-zero AND parity vs the pure-JAX reference.""" + + @pytest.mark.parametrize("config", _CONFIGS) + def test_backward(self, mesh, config): + block = _make_block(**config) + x = _make_inputs(jax.random.PRNGKey(2)) + variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) + grads_te, grad_x_te = _grad_step(block, variables, mesh, x) + + # Reference grads via jax.grad over the pure-JAX MoE with the + # same config. argnums=(0, 1) so the reference also produces a + # d_x for the propagated-gradient parity check below. + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + ref_kwargs = _reference_kwargs_from_config(config, params_np) + ref_expert_bias = ref_kwargs.pop("expert_bias") + + def loss_fn(params, x): + out, _ = _pure_jax_moe_reference( + x, + params["gate_kernel"], + params["wi_0"], + params["wi_1"], + params["wo"], + ref_expert_bias, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + **ref_kwargs, + ) + return jnp.mean(out.astype(jnp.float32) ** 2) + + grads_ref, grad_x_ref = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))( + {k: jnp.asarray(v) for k, v in params_np.items() if k != "expert_bias"}, + jnp.asarray(x_np), + ) + grads_ref_np = {k: np.asarray(jax.device_get(v)) for k, v in grads_ref.items()} + grad_x_ref_np = np.asarray(jax.device_get(grad_x_ref)) + + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + # Per-tensor: finite + non-zero + parity in one pass. + g_te = _to_global_numpy(_unwrap(grads_te["params"][name]), mesh) + assert np.all(np.isfinite(g_te)), f"{name} grad has NaN/Inf [config={config}]" + assert np.any(g_te != 0.0), f"{name} grad identically zero [config={config}]" + atol, rtol = ( + (GRAD_GATE_ATOL, GRAD_GATE_RTOL) + if name == "gate_kernel" + else (GRAD_FFN_ATOL, GRAD_FFN_RTOL) + ) + np.testing.assert_allclose( + g_te.astype(np.float32), + grads_ref_np[name].astype(np.float32), + atol=atol, + rtol=rtol, + err_msg=f"grad parity breach on {name} [config={config}]", + ) + + # d_x: the gradient propagated back to the previous layer. Checks + # shape, dtype (must match x.dtype — protects the + # _with_sharding_constraint_cast_bwd wrapper that casts the + # fp32-promoted gate path back to bf16), finiteness, non-zero + # AND numerical parity vs the pure-JAX reference d_x. + grad_x_te_np = _to_global_numpy(grad_x_te, mesh) + assert ( + grad_x_te.shape == x.shape + ), f"d_x shape {grad_x_te.shape} != x.shape {x.shape} [config={config}]" + assert ( + grad_x_te.dtype == x.dtype + ), f"d_x dtype {grad_x_te.dtype} != x.dtype {x.dtype} [config={config}]" + assert np.all(np.isfinite(grad_x_te_np)), f"d_x has NaN/Inf [config={config}]" + assert np.any(grad_x_te_np != 0.0), f"d_x identically zero [config={config}]" + np.testing.assert_allclose( + grad_x_te_np.astype(np.float32), + grad_x_ref_np.astype(np.float32), + atol=GRAD_FFN_ATOL, + rtol=GRAD_FFN_RTOL, + err_msg=f"d_x parity breach [config={config}]", + ) + + +class TestTeEpMoeAuxLoss: + """Aux-loss path. Consolidated into: + * ``test_aux_loss``: one run that checks the returned scalar's + shape / dtype / finiteness / magnitude AND numerical parity vs the + reference AND that the aux-only bwd propagates to gate_kernel. + * ``test_combined_loss_grads``: one run for joint main+aux bwd + finite + non-zero per tensor. + """ + + def test_aux_loss(self, mesh): + coeff = 1e-2 + block = _make_block(aux_loss_coeff=coeff) + x = _make_inputs(jax.random.PRNGKey(20)) + variables, _, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(21)) + + # Shape / dtype / finiteness / magnitude. + assert aux is not None, "aux_loss should be returned when coeff > 0" + assert aux.shape == (), f"aux_loss must be 0-d scalar, got {aux.shape}" + assert aux.dtype == DTYPE, f"aux_loss dtype {aux.dtype} != {DTYPE}" + aux_np = _to_global_numpy(aux, mesh) + assert np.isfinite(aux_np), "aux_loss is NaN/Inf" + assert abs(float(aux_np)) < 1e2, f"aux_loss looks unreasonable: {aux_np}" + + # Numerical parity vs the reference. + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + _, aux_ref = _pure_jax_moe_reference( + jnp.asarray(x_np), + jnp.asarray(params_np["gate_kernel"]), + jnp.asarray(params_np["wi_0"]), + jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wo"]), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + aux_loss_coeff=coeff, + ) + np.testing.assert_allclose( + float(aux_np), + float(jax.device_get(aux_ref)), + atol=AUX_ATOL, + rtol=AUX_RTOL, + ) + + # Aux-only bwd must propagate to gate_kernel — proves the + # fused_moe_aux_loss_bwd → topk(compute_aux_scores)_bwd chain is + # wired. + aux_grads = _grad_aux_only(block, variables, mesh, x) + g_gate = np.asarray( + jax.device_get(_unwrap(aux_grads["params"]["gate_kernel"]).addressable_data(0)) + ) + assert np.all(np.isfinite(g_gate)), "gate grad NaN/Inf under aux-only loss" + assert np.any(g_gate != 0.0), "aux bwd should propagate to gate_kernel" + + def test_combined_loss_grads(self, mesh): + """Joint main + aux loss bwd: per-tensor finite + non-zero in + one pass.""" + block = _make_block(aux_loss_coeff=1e-2) + x = _make_inputs(jax.random.PRNGKey(22)) + variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(23)) + grads, _ = _grad_step(block, variables, mesh, x, include_aux=True) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g_local = np.asarray(jax.device_get(_unwrap(grads["params"][name]).addressable_data(0))) + assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf under main+aux" + assert np.any(g_local != 0.0), f"{name} grad zero under main+aux" diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 34ec0a33b6..77e60afbcd 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -23,7 +23,7 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive -from ..sharding import global_mesh_resource +from ..sharding import global_mesh_resource, get_mesh_axis_size __all__ = [ "EpConfig", @@ -125,8 +125,15 @@ def _ep_outer_axis(): When set, EP-output globals carry an extra leading ``dp_size`` dim so SPMD sees each DP color's slab as distinct (rather than replicated across DP). + + A dp/fsdp axis that is sized 1 in the active mesh is treated as absent so + we don't pin EP-output specs to a degenerate axis that JAX may collapse. """ gsr = global_mesh_resource() + if gsr.dp_resource is not None and get_mesh_axis_size(gsr.dp_resource) > 1: + return gsr.dp_resource + if gsr.fsdp_resource is not None and get_mesh_axis_size(gsr.fsdp_resource) > 1: + return gsr.fsdp_resource return gsr.dp_resource or gsr.fsdp_resource diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index 8cc94fcaaf..46f51c9d33 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -412,6 +412,11 @@ def partition( arg_infos, result_infos, ): + # NOTE: do NOT include ``routing_map_format`` in this ``del``: the + # ``sharded_impl`` closure below resolves it by name at call time + # (when XLA invokes the partitioned impl), so deleting it here + # raises ``NameError: cannot access free variable 'routing_map_format'`` + # at execution time of the bwd custom_partitioning. del result_infos grad_spec = get_padded_spec(arg_infos[2]) out_sharding = NamedSharding(mesh, PartitionSpec(*grad_spec)) @@ -645,7 +650,14 @@ def shardy_sharding_rule(*args): # backward reconstructs the full [num_tokens, num_experts] grad_probs from # scalar inputs. Shardy will leave num_tokens unsharded, which matches the # replicated PartitionSpec(None, None) in partition(). - return "const_buf_one, num_experts, grad_one -> i num_experts" + # + # grad_aux_loss is the cotangent of a scalar loss and is therefore + # rank-0; the third operand entry is empty (no factor labels). Declaring + # it with the spurious "grad_one" factor gave it rank-1 and tripped + # JAX's custom_partitioning_sharding_rule check once the MoE block + # lifted its aux-loss path out of shard_map (the rule is skipped under + # shard_map, which is why this surfaces only at global view). + return "const_buf_one, num_experts, -> i num_experts" register_primitive(FusedMoEAuxLossBwdPrimitive) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 91346a7a48..3629346e33 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -37,8 +37,7 @@ # import P`` without a second jax.sharding import. from jax.sharding import PartitionSpec as P # noqa: F401 # pylint: disable=unused-import -from ..moe import PermutationBackend, moe -from ..quantize import noop_quantizer_set +from ..moe import moe from ..router import ScoreFunction from ..sharding import get_active_resource_axis from .module import TransformerEngineBase @@ -50,7 +49,7 @@ Initializer = Callable[[PRNGKey, Shape, DType], Array] -__all__ = ["PermutationBackend", "_MoEBlock"] +__all__ = ["_MoEBlock"] class _MoEBlock(TransformerEngineBase): @@ -82,10 +81,11 @@ class _MoEBlock(TransformerEngineBase): Grouped top-k knobs (DeepSeek-style). ``None`` disables grouping. scaling_factor : float Multiplier on the routing weights. - use_expert_bias : bool - If ``True``, registers a per-expert routing bias (shape ``[E]``). - Only meaningful with ``score_function="sigmoid"``; the underlying - primitive validates the pairing. + use_expert_routing_bias : bool + If ``True``, registers a per-expert routing bias (shape ``[E]``) + used by the topk selection. Only meaningful with + ``score_function="sigmoid"``; the underlying primitive validates + the pairing. aux_loss_coeff : float If ``> 0``, return the MoE auxiliary load-balancing loss scalar in addition to the main output. @@ -100,23 +100,27 @@ class _MoEBlock(TransformerEngineBase): replicated across non-EP axes within an EP group; set e.g. ``("fsdp",)`` for true FSDP-of-batch where each device owns a unique slice of the batch. - permutation_backend : PermutationBackend - ``PURE_JAX`` (default) or ``TRITON``. - _align_size : int - Per-expert group-size alignment (``0`` disables; required > 0 - for quantized grouped GEMM). Internal knob; will be inferred - from the active quantization recipe in a follow-up PR. + apply_topk_weights_early : bool + If ``True``, multiply expert outputs by their top-k weights + *inside* each shard before ``ep_combine`` (saves one global + reduction at the cost of an extra broadcast). Default ``False``. + + The per-expert dispatch-slot alignment is fixed internally at 128 + tokens (see ``moe._ALIGN_SIZE``) -- the value required by NCCL EP + HT and satisfied by every current TE grouped-GEMM recipe -- and is + therefore not exposed as a per-instance knob. dtype : jnp.dtype Compute / parameter dtype. kernel_init, bias_init, expert_bias_init : Initializers. - use_bias : bool - Register per-expert FFN biases. + use_ffn_bias : bool + Register per-expert FFN biases (``wi_0_bias``, ``wi_1_bias``, + ``wo_bias``). Quantization is currently configured via the standard TE autocast - context (``fp8_autocast``/``with_quantizer_set``); per-call - quantizer sets can also be passed through ``__call__``'s - ``quantizer_sets`` keyword once we stabilise the recipe pipeline. + context (``fp8_autocast``/``with_quantizer_set``) and threaded + through ``moe()`` internally; this wrapper does not expose a + per-call ``quantizer_sets`` knob yet. """ # Architecture @@ -131,7 +135,7 @@ class _MoEBlock(TransformerEngineBase): num_groups: Optional[int] = None group_topk: Optional[int] = None scaling_factor: float = 1.0 - use_expert_bias: bool = False + use_expert_routing_bias: bool = False aux_loss_coeff: float = 0.0 # Sharding (logical axes) @@ -143,16 +147,15 @@ class _MoEBlock(TransformerEngineBase): # Parallelism data_parallelism_axes: Tuple[str, ...] = () - # Permutation - permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX - _align_size: int = 0 + # MoE knobs forwarded to ``moe()`` + apply_topk_weights_early: bool = False # Dtypes / init / misc dtype: DType = jnp.float32 kernel_init: Optional[Initializer] = None bias_init: Initializer = nn.initializers.zeros expert_bias_init: Initializer = nn.initializers.zeros - use_bias: bool = False + use_ffn_bias: bool = False def __post_init__(self): if self.kernel_init is None: @@ -163,11 +166,6 @@ def __post_init__(self): 1.0, "fan_in", "truncated_normal", dtype=self.dtype ), ) - if not isinstance(self.permutation_backend, PermutationBackend): - raise TypeError( - "permutation_backend must be a PermutationBackend, got" - f" {self.permutation_backend!r}" - ) super().__post_init__() @nn.compact @@ -221,7 +219,7 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: self.dtype, ) wi_0_bias = wi_1_bias = wo_bias = None - if self.use_bias: + if self.use_ffn_bias: wi_0_bias = self.param( "wi_0_bias", nn.with_logical_partitioning(self.bias_init, ("exp", "mlp")), @@ -241,12 +239,14 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: self.dtype, ) expert_bias = None - if self.use_expert_bias: + if self.use_expert_routing_bias: + # The router logits are promoted to fp32 before fused top-k; keep + # the routing bias in the same dtype so it only affects selection. expert_bias = self.param( "expert_bias", nn.with_logical_partitioning(self.expert_bias_init, ("exp",)), (self.num_experts,), - self.dtype, + jnp.float32, ) ep_axis = get_active_resource_axis("ep_resource") @@ -270,15 +270,12 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: group_topk=self.group_topk, scaling_factor=self.scaling_factor, aux_loss_coeff=self.aux_loss_coeff, - permutation_backend=self.permutation_backend, - align_size=self._align_size, - gate_inside_vjp=True, + apply_topk_weights_early=self.apply_topk_weights_early, ep_axis=ep_axis, data_parallelism_axes=self.data_parallelism_axes, input_axes=self.input_axes, gate_kernel_axes=self.gate_kernel_axes, wi_kernel_axes=self.wi_kernel_axes, wo_kernel_axes=self.wo_kernel_axes, - quantizer_sets=(noop_quantizer_set, noop_quantizer_set, noop_quantizer_set), dtype=self.dtype, ) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 2a1c818cb3..887e005de6 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -1,76 +1,51 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Functional Mixture-of-Experts (MoE) entry point with a single fused VJP. - -This module exposes :func:`moe`, the framework-agnostic flat function that -implements an entire MoE block (gate -> top-k routing -> token dispatch -> -per-expert FFN -> token combine, plus optional expert parallelism via a -shard_map / ragged_all_to_all collective) under a *single* -``jax.custom_vjp``. It is the moral analog of -:func:`transformer_engine.jax.layernorm_mlp.layernorm_mlp` for MoE: one -custom_vjp boundary covers the whole block so future fusions (FP8 over the -EP wire, fused ``ragged_all_to_all + grouped_gemm``, gate+route+dispatch -fusion) can land without re-architecting the call site. - -Design rationale ----------------- - -The earlier MoE block (:class:`transformer_engine.jax.flax.moe._MoEBlock`) -composed many narrower custom_vjps -- one per :func:`grouped_dense`, one -per :func:`token_dispatch`, etc. Every nested custom_vjp is a place where -a quantized :class:`ScaledTensor` cannot survive (JAX requires custom_vjp -inputs / outputs to be plain ``jnp.ndarray`` ish pytrees). To enable -end-to-end FP8 flow -- in particular FP8 carried over the EP -ragged_all_to_all -- the dispatch's quantize, the a2a, the per-expert -FFN, the inverse a2a, and the combine all have to live inside the same -VJP. This file collapses them into one. - -Implementation conventions --------------------------- - -* No nested ``custom_vjp``. Every primitive's ``_fwd`` and ``_bwd`` is - called directly (e.g. :func:`tex.fused_topk_with_score_function_fwd` / - ``_bwd``, :func:`unpermute_with_mask_map`, - :func:`unpermute_bwd_with_merging_probs`, - :func:`sort_chunks_by_map(is_forward=False)`, - forward + reverse :func:`jax.lax.ragged_all_to_all`) so the outer - ``_moe_bwd_rule`` controls the bwd graph end-to-end without invoking - ``jax.vjp`` for re-linearization. -* The fwd/bwd context (``ctx``) is a plain ``dict`` whose keys depend on - the static configuration (permutation backend, EP active or not, - presence of biases, aux loss enabled). The ``_moe_fwd_rule`` builds a - matching ``ctx_specs`` dict in lockstep when opening the EP shard_map - so ``out_specs`` structurally matches the body's return. -* :func:`_dispatch` is the helper that wraps - ``permute -> a2a -> local_permute`` (forward); :func:`_combine` is its - inverse. Their ``_bwd`` siblings drive the inverse collectives in the - bwd rule. None of these helpers form a custom_vjp boundary. +"""Mixture-of-Experts (MoE) layer for TransformerEngine JAX. + +This module exposes :func:`moe`, a single fused MoE forward pass + bwd +built on top of TE's NCCL-backed Expert Parallelism primitives +(``tex.ep_dispatch`` / ``tex.ep_combine``). The block runs:: + + gate -> topk -> ep_dispatch -> per-expert FFN (grouped GEMMs) + -> ep_combine -> output + +under a single ``jax.custom_vjp`` so the routing, dispatch, FFN and +combine steps fuse cleanly under XLA without leaking intermediate +residuals into the user-facing autograd graph. + +Sharding model +-------------- +* Inbound activations are 3D ``[B, S, H]`` sharded + ``((*data_parallelism_axes, ep_axis), None, None)``. The public + :func:`moe` soft-repins this on entry and warns when a reshard is + inserted. +* The EP primitives operate at global view (their custom_partitioning + rules handle per-shard execution). The FFN GEMMs run per-shard inside + a small ``shard_map`` whose ``in_specs`` and ``out_specs`` mirror the + same ``((dp, ep), ...)`` layout. + +Out-of-scope (for now) +---------------------- +FP8 / MXFP8 quantizer sets are not yet wired on this path; turning +them on requires recipe-aware residual specs and ``ScaledTensor`` +leaves across the ``shard_map`` boundary. ``aux_loss_coeff`` and +``expert_bias`` are supported (the former forces a per-step +all-gather over the routing-side logits, which lives off the critical +path and overlaps with the dispatch collective). """ -import math -from dataclasses import dataclass -from enum import Enum from functools import partial -from typing import Any, NewType, Optional, Tuple, Union +from typing import Any, Optional, Tuple, Union +import warnings +import flax.struct import jax import jax.numpy as jnp -from flax import struct as flax_struct -from jax.sharding import PartitionSpec as P +from jax.sharding import NamedSharding, PartitionSpec as P from . import cpp_extensions as tex -from .permutation import ( - PureJaxPermState, - compute_ragged_all_to_all_params, - compute_reverse_ragged_all_to_all_params, - pure_jax_token_combine, - pure_jax_token_dispatch, - routing_map_to_selected_experts, -) from .quantize import ( - QuantizerSet, - ScaledTensor, TensorUsage, noop_quantizer_set, with_sharding_constraint_by_logical_axes, @@ -79,1070 +54,251 @@ from .router import ScoreFunction, _validate_score_function from .sharding import _get_mesh -# Triton-backed primitives are imported lazily: callers on the PURE_JAX -# permutation backend should not need ``triton`` installed. The TRITON -# branches in this module call ``_require_triton()`` first to raise a -# clear error if the import failed. -try: - from .triton_extensions.permutation import ( - make_chunk_sort_map, - make_row_id_map, - permute_with_mask_map, - permute_with_mask_map_and_pad, - sort_chunks_by_map, - unpermute_bwd_with_merging_probs, - unpermute_bwd_with_merging_probs_and_unpad, - unpermute_with_mask_map, - unpermute_with_mask_map_and_unpad, - ) - - _TRITON_AVAILABLE = True -except ImportError: - _TRITON_AVAILABLE = False - make_chunk_sort_map = None - make_row_id_map = None - permute_with_mask_map = None - permute_with_mask_map_and_pad = None - sort_chunks_by_map = None - unpermute_bwd_with_merging_probs = None - unpermute_bwd_with_merging_probs_and_unpad = None - unpermute_with_mask_map = None - unpermute_with_mask_map_and_unpad = None - - -def _require_triton(): - """Raise a clear error if Triton permutation kernels are unavailable.""" - if not _TRITON_AVAILABLE: - raise ImportError( - "PermutationBackend.TRITON requires" - " ``transformer_engine.jax.triton_extensions`` (and ``triton``)." - " Install Triton or pass PermutationBackend.PURE_JAX." - ) - - -PRNGKey = Any -Shape = Tuple[int, ...] -DType = NewType("DType", jnp.dtype) -Array = NewType("Array", jnp.ndarray) - +__all__ = ["moe"] -__all__ = ["moe", "PermutationBackend"] +# Per-expert dispatch-slot alignment fed to ``tex.ep_prepare`` as +# ``dispatch_output_per_expert_alignment``. NCCL EP HT requires the +# per-expert recv block to be at least 128-token aligned, and all current +# TE grouped-GEMM recipes (bf16/fp16/fp8/mxfp8) are satisfied by the +# same 128-token tile, so a single constant covers every supported path. +_ALIGN_SIZE = 128 -# ============================================================================= -# Enums -# ============================================================================= +def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: + """Sharding constraint that keeps bwd cotangents in the primal dtype. -class PermutationBackend(Enum): - """Token-dispatch / combine backend used by :func:`moe`. - - * ``TRITON``: TE's fused Triton kernels. Faster than ``PURE_JAX`` - on current hardware and the recommended default. - * ``PURE_JAX``: ``jnp.argsort`` + gather paths compiled as plain - XLA; useful as a numerical reference and on builds without - Triton available. - """ + Plain ``jax.lax.with_sharding_constraint`` is identity on the fwd + but does not constrain the dtype of the cotangent that flows back + through it. In this MoE bwd, ``d_x`` is built from two paths: - PURE_JAX = "pure_jax" - TRITON = "triton" + * ``d_x_from_dispatch`` from ``ep_dispatch_bwd`` -- primal dtype + (bf16 in mixed precision). + * ``d_x_from_gate = d_logits_2d @ gate_kernel.T`` where + ``d_logits_2d`` is produced by + ``fused_topk_with_score_function_bwd``. That primitive runs at + fp32 because the fwd promoted ``logits_2d`` to fp32 (the fused + topk/softmax/sigmoid kernels are only validated at fp32). - -# ============================================================================= -# Dispatch-state records (carried _dispatch -> _combine / *_bwd) -# ============================================================================= -# -# Two NamedTuples (one per permutation backend) so we get type -# discrimination at the consumer side via ``isinstance``. The backend- -# specific residuals are required fields; the EP-only residuals are -# Optional and are populated only when the run is EP-active. Each field -# is either an ``ndarray`` or ``None`` -- nothing static, since these -# values cross the shard_map pytree boundary and would otherwise be -# coerced into JitTracers. - - -@flax_struct.dataclass -class _PureJaxDispatchState: - """Residuals saved by :func:`_dispatch` on the PURE_JAX path. - - Registered as a JAX pytree via ``flax.struct.dataclass``: each - annotated field is a leaf, ``None`` is a non-leaf sentinel. The - matching spec built by :func:`_build_dispatch_specs` mirrors this - layout so shard_map's value and spec trees line up. + JAX's type promotion then makes ``d_x_from_gate + d_x_from_dispatch`` + fp32, so the user-visible ``d_x`` ends up wider than ``x``. That + doubles activation-grad bandwidth and breaks any downstream kernel + that pins a bf16 input layout. This wrapper inserts an explicit + cast back to the primal dtype on the bwd side and re-asserts the + same sharding there as well. """ - group_sizes: jnp.ndarray - sorted_indices: jnp.ndarray - routing_weights: jnp.ndarray - # EP-only: - all_shards_tokens_per_expert: Optional[jnp.ndarray] = None - local_perm_row_id_map: Optional[jnp.ndarray] = None - - -@flax_struct.dataclass -class _TritonDispatchState: - """Residuals saved by :func:`_dispatch` on the TRITON path.""" - - group_sizes: jnp.ndarray - row_id_map: jnp.ndarray - pad_offsets: Optional[jnp.ndarray] # populated only when align_size > 0 - merging_probs: jnp.ndarray - # EP-only: - all_shards_tokens_per_expert: Optional[jnp.ndarray] = None - local_perm_row_id_map: Optional[jnp.ndarray] = None - - -_DispatchState = Union[_PureJaxDispatchState, _TritonDispatchState] - + @jax.custom_vjp + def _constraint(y): + return jax.lax.with_sharding_constraint(y, sharding) -@flax_struct.dataclass -class _BodyCtx: - """Residuals carried fwd_rule -> bwd_rule by :func:`_body_fwd`. + def _constraint_fwd(y): + return jax.lax.with_sharding_constraint(y, sharding), jnp.zeros((), dtype=y.dtype) - Optional fields (``expert_bias``, ``aux_*``) are ``None`` when the - matching feature is disabled. :func:`_build_ctx_specs` mirrors that - layout so the shard_map spec and value trees match leaf-for-leaf. - """ + def _constraint_bwd(dtype_ref, grad): + return (jax.lax.with_sharding_constraint(grad.astype(dtype_ref.dtype), sharding),) - # Always present. - x: Any - gate_kernel: Any - logits_2d: Any - saved_scores: Any - routing_map: Any - dispatch: Any # _DispatchState - casted_sorted_x_lhs_trans: Any - casted_wi_rhs_trans: Any # combined [E, H, 2M] residual for fused wi_0|wi_1 bwd - gate_proj_out: Any - up_proj_out: Any - casted_intermediate_lhs_trans: Any - casted_wo_rhs_trans: Any - expert_outputs: Any - local_group_sizes: Any - # Feature-gated. - expert_bias: Any = None - aux_const_buf: Any = None - aux_tokens_per_expert: Any = None - aux_logits_for_score: Any = None - aux_saved_scores: Any = None + _constraint.defvjp(_constraint_fwd, _constraint_bwd) + return _constraint(x) # ============================================================================= -# ctx / dispatch-state key conventions +# Process-level NCCL EP bootstrap (must run eagerly, outside jax.jit) # ============================================================================= # -# Both ``ctx`` (carried fwd_rule -> bwd_rule) and the dispatch state -# (carried _dispatch -> _combine / _dispatch_bwd / _combine_bwd) are plain -# python dicts. Using a dict (rather than a flax_struct.dataclass) lets us -# vary the populated keys with the static config without breaking -# ``shard_map``'s ``out_specs`` structural match: the spec dict and the -# value dict are built with the SAME keys via :func:`_build_ctx_specs`. -# -# Below is the key glossary so the rest of the file reads cleanly. -# -# DispatchState (dict): values are jnp.ndarray unless noted -# Always present: -# "group_sizes" [n_groups] per-expert token counts -# (n_groups = E for no-EP, -# E_local for EP) -# "ep_active" bool (carried as a Python flag, -# not in the dict; passed -# alongside) -# PURE_JAX backend: -# "sorted_indices" [num_real + padding] argsort indices -# "routing_weights" [num_tokens, topk] per-token-per-expert weights -# TRITON backend: -# "row_id_map" [num_tokens, 2*E + 1] -# "pad_offsets" [E] or None -# "merging_probs" [num_tokens, E] -# EP-only: -# "all_shards_tokens_per_expert" [num_ep, E] -# "local_perm_row_id_map" [recv_buffer_rows] -# "local_perm_inv_row_id_map" [recv_buffer_rows] -# -# NOTE: per-shard compile-time-constant shapes (num_real_tokens, -# padding_size, pre/post_a2a_buffer_shape) are NOT stored in this -# dict; they are recomputed in _body_fwd/_body_bwd via -# _compute_static_shape_info and passed as Python ints / int tuples to -# the dispatch/combine helpers. Storing them in the dict would cause -# JAX's pytree-flatten across the shard_map boundary to coerce them -# into JitTracer 0-d arrays, which breaks Python-level control flow -# (e.g. ``if padding > 0``) and ``jnp.zeros(shape)`` in the bwd. -# -# See :class:`_BodyCtx` (NamedTuple) for the ctx layout and field -# documentation. :func:`_build_ctx_specs` returns a matching ``_BodyCtx`` -# of ``P(...)`` specs so shard_map's value/spec trees line up -# leaf-for-leaf. - - -# ============================================================================= -# Static shape helper -# ============================================================================= -# -# A set of per-shard shape/size values that the dispatch and combine -# helpers (both fwd and bwd) need. They're all derivable from existing -# static args, so we recompute them in both ``_body_fwd`` and -# ``_body_bwd`` and pass them as Python ints / int-tuples through -# explicit kwargs. We MUST NOT stash them inside the dynamic -# ``state`` / ``ctx`` dict: when the dict crosses the EP shard_map's -# out_specs/in_specs boundary, JAX's pytree-flatten coerces any Python -# int leaves into traced 0-d arrays, which then breaks dependent Python -# code in the bwd (e.g. ``if padding > 0`` and ``jnp.zeros(shape)``). - - -@dataclass(frozen=True) -class _StaticShapeInfo: - """Per-shard compile-time-constant shape info used by dispatch / - combine fwd and bwd. Fields are Python ints / int tuples (NOT jnp - arrays) so they can be passed as ordinary static keyword args. - - Attributes - ---------- - num_real_tokens : int - Per-shard count of real (non-padding) permuted tokens, - i.e. ``per_shard_num_tokens * num_experts_per_tok``. - padding_size : int - Per-shard number of alignment-padding tokens appended to the - sort buffer (``num_experts * (align_size - 1)`` when - ``align_size > 0``, else ``0``). - pre_a2a_buffer_shape : tuple[int, int] - ``(num_real_tokens + padding_size, hidden)`` -- the per-shard - shape of the sorted-inputs buffer sent over the EP - ragged_all_to_all in the fwd direction. - post_a2a_buffer_shape : Optional[tuple[int, int]] - ``(recv_buffer_rows, hidden)`` when EP is active, ``None`` - otherwise. - """ +# ``tex.ep_bootstrap`` does a NCCL UID allgather over the JAX runtime, which +# cannot run from inside a jit-traced function. The caller must bootstrap +# eagerly once per process before any jitted MoE call, then record the +# bootstrap signature via ``record_ep_bootstrap_signature_for_moe``. The +# per-call check below verifies the recorded signature is wide enough for +# the current MoE invocation (smaller per-call usage is fine since the C++ +# backend reserves worst-case buffers at bootstrap time). - num_real_tokens: int - padding_size: int - pre_a2a_buffer_shape: Tuple[int, int] - post_a2a_buffer_shape: Optional[Tuple[int, int]] +_te_ep_bootstrap_signature: Optional[Tuple[int, int, int, int, int]] = None -def _compute_static_shape_info( - *, - batch_size: int, - sequence_length: int, - hidden: int, +def record_ep_bootstrap_signature_for_moe( num_experts: int, - num_experts_per_tok: int, - align_size: int, - ep_active: bool, - num_ep: int = 1, - fsdp_sizes: Tuple[int, ...] = (), - recv_buffer_rows: int = 0, - batch_is_per_shard: bool = True, -) -> _StaticShapeInfo: - """Build a :class:`_StaticShapeInfo` for the current rank. - - ``batch_is_per_shard`` controls whether ``batch_size`` is already - sharded (True -- e.g. when this is called from inside a shard_map - body, where ``x.shape[0]`` reports the per-shard batch size) or - global (False -- e.g. when computing from x.shape outside the - shard_map body). + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + ep_size: int, +) -> None: + """Record the params passed to ``ep_bootstrap`` so the per-call check + in ``_moe_fwd_rule`` can verify compatibility. Call this once per + process immediately after ``ep_bootstrap``. """ - if ep_active and not batch_is_per_shard: - dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 - per_shard_batch = batch_size // (num_ep * dp_size) - else: - per_shard_batch = batch_size - per_shard_num_tokens = per_shard_batch * sequence_length - num_real_tokens = per_shard_num_tokens * num_experts_per_tok - padding_size = num_experts * (align_size - 1) if align_size > 0 else 0 - pre_a2a_buffer_shape = (num_real_tokens + padding_size, hidden) - post_a2a_buffer_shape = (recv_buffer_rows, hidden) if ep_active else None - return _StaticShapeInfo( - num_real_tokens=num_real_tokens, - padding_size=padding_size, - pre_a2a_buffer_shape=pre_a2a_buffer_shape, - post_a2a_buffer_shape=post_a2a_buffer_shape, + global _te_ep_bootstrap_signature + _te_ep_bootstrap_signature = ( + num_experts, + max_tokens_per_rank, + recv_capacity_per_rank, + hidden_dim, + ep_size, ) -# ============================================================================= -# Dispatch / combine helpers (no VJP boundary -- pure Python) -# ============================================================================= - - -def _dispatch( - inputs_2d: jnp.ndarray, - sparse_probs: jnp.ndarray, - routing_map: jnp.ndarray, - *, - backend: PermutationBackend, +def _te_ep_assert_compatible_bootstrap( num_experts: int, - num_experts_per_tok: int, - align_size: int, - # EP-only: - ep_active: bool, - ep_axis: Optional[str], - num_ep: int, - recv_buffer_rows: int, - shard_id: Optional[jnp.ndarray] = None, -) -> Tuple[jnp.ndarray, dict]: - """``permute -> (a2a -> local_permute) iff ep_active``. - - Returns ``(sorted_x, state)`` where ``sorted_x`` has shape - ``[buffer_rows, hidden]`` -- ``E`` groups (no-EP) or ``E_local`` groups - (EP) -- and ``state`` is a dict carrying everything :func:`_combine` - and the bwd helpers need to reverse the operation. - - Bypasses the ``custom_vjp``-wrapped public ``token_dispatch`` / - ``pure_jax_token_dispatch`` wrappers (well, mostly: PURE_JAX still - composes through ``pure_jax_token_dispatch`` because that helper has - no ``custom_vjp`` itself -- only its inner ``_sort_activations`` does, - which is fine since we never auto-diff through it from this layer). - For TRITON we call the underlying ``permute_with_mask_map`` / - ``permute_with_mask_map_and_pad`` primitives directly. - """ - num_tokens, hidden = inputs_2d.shape - topk = num_experts_per_tok - - # Backend-specific residuals collected here, then packaged into the - # appropriate _*DispatchState below. - sorted_indices = None - routing_weights_kept = None - row_id_map = None - pad_offsets = None - merging_probs = None - - # ------------------------------------------------------------------ - # Step 1: global permute (every shard routes its own tokens over the - # full expert axis). Backend-specific. - # ------------------------------------------------------------------ - if backend is PermutationBackend.PURE_JAX: - selected_experts, routing_weights = routing_map_to_selected_experts( - sparse_probs, routing_map, topk - ) - sorted_inputs, perm_state, group_sizes = pure_jax_token_dispatch( - inputs_2d, - selected_experts, - num_experts=num_experts, - num_experts_per_tok=topk, - align_size=align_size, - ) - # NOTE: ``perm_state.num_real_tokens`` and ``perm_state.padding_size`` - # are compile-time Python ints; intentionally NOT stored in the - # returned state (would be coerced to JitTracer 0-d arrays under - # the EP shard_map's pytree flatten). Recompute via - # ``_compute_static_shape_info`` in the bwd / EP-combine - # call sites that need them. - sorted_indices = perm_state.sorted_indices - routing_weights_kept = routing_weights - else: - # TRITON backend -- inline the underlying primitive sequence - # (mirrors ``_token_dispatch_fwd_rule`` but exposes the residuals - # to our ctx instead of saving them inside another custom_vjp). - num_out_tokens = num_tokens * topk - row_id_map = make_row_id_map(routing_map, num_tokens, num_experts) - tokens_per_expert = jnp.sum(routing_map, axis=0).astype(jnp.int32) - if align_size > 0: - target_tokens_per_expert = ( - jnp.ceil(tokens_per_expert / align_size) * align_size - ).astype(jnp.int32) - pad_lengths = target_tokens_per_expert - tokens_per_expert - cum_pad = jnp.cumsum(pad_lengths) - pad_offsets = jnp.concatenate([jnp.array([0], dtype=cum_pad.dtype), cum_pad[:-1]]) - worst_case_out_tokens = ( - (num_out_tokens + num_experts * (align_size - 1)) // align_size - ) * align_size - sorted_inputs, _ = permute_with_mask_map_and_pad( - inputs_2d, - row_id_map, - None, - pad_offsets, - num_tokens, - num_experts, - worst_case_out_tokens, - hidden, - align_size=align_size, - ) - group_sizes = target_tokens_per_expert - else: - sorted_inputs, _ = permute_with_mask_map( - inputs_2d, - row_id_map, - None, - num_tokens, - num_experts, - num_out_tokens, - hidden, - ) - pad_offsets = None - group_sizes = tokens_per_expert - merging_probs = sparse_probs - - def _build_state(group_sizes_val, ep_all=None, ep_local=None): - if backend is PermutationBackend.PURE_JAX: - return _PureJaxDispatchState( - group_sizes=group_sizes_val, - sorted_indices=sorted_indices, - routing_weights=routing_weights_kept, - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - return _TritonDispatchState( - group_sizes=group_sizes_val, - row_id_map=row_id_map, - pad_offsets=pad_offsets, - merging_probs=merging_probs, - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - - if not ep_active: - return sorted_inputs, _build_state(group_sizes) - - # ------------------------------------------------------------------ - # Step 2 (EP only): all_gather per-expert counts so every shard knows - # the [num_ep, num_experts] token-count matrix. - # ------------------------------------------------------------------ - all_shards_tokens_per_expert = jax.lax.all_gather( - group_sizes[None, :], - axis_name=ep_axis, - axis=0, - tiled=True, - ) - - # ------------------------------------------------------------------ - # Step 3 (EP only): forward ragged_all_to_all over the EP axis. - # ------------------------------------------------------------------ - in_off, send_sz, out_off, recv_sz = compute_ragged_all_to_all_params( - all_shards_tokens_per_expert, shard_id, num_ep - ) - post_a2a_buffer_shape = (recv_buffer_rows, hidden) - recv_buf = jnp.zeros(post_a2a_buffer_shape, dtype=sorted_inputs.dtype) - x_recv = jax.lax.ragged_all_to_all( - sorted_inputs, recv_buf, in_off, send_sz, out_off, recv_sz, axis_name=ep_axis - ) - - # ------------------------------------------------------------------ - # Step 4 (EP only): local permute -- (source_shard, expert) -> - # (expert, shard). Inlined ``local_permute_after_a2a`` so we control - # both the row_id_map and its inverse for the bwd. - # ------------------------------------------------------------------ - num_experts_local = num_experts // num_ep - local_expert_start = shard_id * num_experts_local - local_expert_columns = jax.lax.dynamic_slice( - all_shards_tokens_per_expert, - start_indices=(0, local_expert_start), - slice_sizes=(num_ep, num_experts_local), - ) - split_sizes = local_expert_columns.reshape(-1) # source-major - indices_matrix = jnp.arange(num_ep * num_experts_local, dtype=jnp.int32).reshape( - num_ep, num_experts_local - ) - sorted_chunk_indices = indices_matrix.T.reshape(-1) # source-major -> expert-major - num_chunks = num_ep * num_experts_local - # Build a SINGLE row_id_map. ``is_forward=True`` permutes - # source-major -> expert-major; ``is_forward=False`` is the exact - # inverse (this is exactly what ``_sort_chunks_by_index_bwd_rule`` - # uses on the saved residual). _MoEBlock builds two row_id_maps - # only because it calls ``sort_chunks_by_index`` twice -- once in - # ``local_permute_after_a2a`` and again in ``local_unpermute_before_a2a``; - # each of those wrappers calls ``make_chunk_sort_map`` internally. - # Here we share one map across (fwd permute, fwd inverse-permute, - # bwd permute, bwd inverse-permute). - local_perm_row_id_map = make_chunk_sort_map( - split_sizes, sorted_chunk_indices, recv_buffer_rows, num_chunks - ) - sorted_x, _ = sort_chunks_by_map( - x_recv, local_perm_row_id_map, None, recv_buffer_rows, hidden, is_forward=True - ) - local_group_sizes = jnp.sum(local_expert_columns, axis=0) - - # NOTE: pre_a2a_buffer_shape and post_a2a_buffer_shape are compile- - # time int tuples; intentionally NOT stored in the returned state - # (would be coerced to JitTracer 0-d arrays under the EP shard_map's - # pytree flatten). Recompute via ``_compute_static_shape_info`` in - # the bwd call sites that need them. For EP, ``group_sizes`` here is - # the per-local-expert count (the FFN runs over E_local groups, not - # E). The global ``group_sizes`` lives inside - # ``all_shards_tokens_per_expert`` if anyone needs it for - # diagnostics. - return sorted_x, _build_state( - local_group_sizes, - ep_all=all_shards_tokens_per_expert, - ep_local=local_perm_row_id_map, - ) - - -def _combine( - expert_outputs: jnp.ndarray, - state: _DispatchState, - *, - backend: PermutationBackend, - ep_active: bool, - batch_size: int, - sequence_length: int, - dtype: jnp.dtype, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # Computed by _compute_static_shape_info in the caller, passed here - # rather than stored in ``state`` to survive shard_map crossings. - num_real_tokens: int, - padding_size: int, - pre_a2a_buffer_shape: Tuple[int, int], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Inverse of :func:`_dispatch`. - - Returns ``(output, expert_outputs_post_ep)``. ``output`` is the - ``[B, S, H]`` combined activations. ``expert_outputs_post_ep`` is - the FFN-output tensor in the shape that Step 3 of the combine - actually consumed (i.e. after the reverse ragged_all_to_all on EP - runs, or the original input on non-EP). The caller stashes this as - the bwd residual so that ``_combine_bwd``'s Step-3 inverse sees - the same tensor the forward Step 3 used. - """ - if ep_active: - # Step 1 (EP): inverse local permute. Reuse the SAME row_id_map - # built in _dispatch by setting is_forward=False (this is the - # exact inverse, identical to what - # ``_sort_chunks_by_index_bwd_rule`` does with the saved residual). - recv_buffer_rows, hidden = expert_outputs.shape - x_send_back, _ = sort_chunks_by_map( - expert_outputs, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=False, - ) - # Step 2 (EP): reverse ragged_all_to_all. - in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - send_back_buf = jnp.zeros(pre_a2a_buffer_shape, dtype=expert_outputs.dtype) - expert_outputs = jax.lax.ragged_all_to_all( - x_send_back, - send_back_buf, - in_off_r, - send_sz_r, - out_off_r, - recv_sz_r, - axis_name=ep_axis, - ) - - # Step 3: global combine. ``expert_outputs`` here is the post-A2A - # tensor under EP, or the original input under non-EP -- whichever - # value Step 3 actually consumes. Returned as the second tuple - # element so the caller can stash it as the bwd residual. - if backend is PermutationBackend.PURE_JAX: - # Reuse the reference pure-jax implementation; it has no - # custom_vjp on its outer surface so we can call it freely. - perm_state = PureJaxPermState( - sorted_indices=state.sorted_indices, - num_real_tokens=num_real_tokens, - padding_size=padding_size, - ) - output = pure_jax_token_combine( - expert_outputs, - perm_state, - state.routing_weights, - num_experts_per_tok=num_experts_per_tok, - batch_size=batch_size, - sequence_length=sequence_length, + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + ep_size: int, +) -> None: + """Verify a prior eager ``ep_bootstrap`` is wide enough for this call.""" + if _te_ep_bootstrap_signature is None: + raise RuntimeError( + "TE EP was not bootstrapped. Call" + " transformer_engine.jax.ep.ep_bootstrap(...) eagerly (outside" + " any jax.jit) once per process, then" + " transformer_engine.jax.moe.record_ep_bootstrap_signature_for_moe(...)" + " with the same params, before invoking moe()." ) - return output, expert_outputs - # TRITON - num_tokens = state.row_id_map.shape[0] - num_experts = (state.row_id_map.shape[1] - 1) // 2 - hidden = expert_outputs.shape[-1] - if state.pad_offsets is not None: - out_2d, _ = unpermute_with_mask_map_and_unpad( - expert_outputs, - state.row_id_map, - state.merging_probs, - None, - state.pad_offsets, - num_tokens, - num_experts, - hidden, - ) - else: - out_2d, _ = unpermute_with_mask_map( - expert_outputs, - state.row_id_map, - state.merging_probs, - None, - num_tokens, - num_experts, - hidden, + b_num_experts, b_max_tpr, b_recv_pr, b_hidden, b_ep_size = _te_ep_bootstrap_signature + if ( + num_experts != b_num_experts + or hidden_dim != b_hidden + or ep_size != b_ep_size + or max_tokens_per_rank > b_max_tpr + or recv_capacity_per_rank > b_recv_pr + ): + raise ValueError( + "TE EP was already bootstrapped with signature" + f" (num_experts={b_num_experts}, max_tokens_per_rank={b_max_tpr}," + f" recv_capacity_per_rank={b_recv_pr}, hidden_dim={b_hidden}," + f" ep_size={b_ep_size}); this moe() call needs" + f" (num_experts={num_experts}, max_tokens_per_rank={max_tokens_per_rank}," + f" recv_capacity_per_rank={recv_capacity_per_rank}, hidden_dim={hidden_dim}," + f" ep_size={ep_size}). Re-bootstrap with wider params (or matching exact" + " sizes) is required." ) - return out_2d.reshape(batch_size, sequence_length, hidden).astype(dtype), expert_outputs -def _combine_bwd( # pylint: disable=unused-argument - d_output: jnp.ndarray, - state: _DispatchState, - expert_outputs: jnp.ndarray, - *, - backend: PermutationBackend, - ep_active: bool, - batch_size: int, - sequence_length: int, - dtype: jnp.dtype, - num_experts: int, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # See ``_compute_static_shape_info`` and the note in ``_dispatch`` - # for why these are kwargs rather than state-dict entries. - num_real_tokens: int, - padding_size: int, - post_a2a_buffer_shape: Optional[Tuple[int, int]], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Inverse of :func:`_combine` on the cotangent. - - Returns ``(d_expert_outputs, d_routing_weights_or_merging_probs)``. +# ============================================================================= +# Residual container threaded fwd -> bwd +# ============================================================================= - ``expert_outputs`` is the *forward* output of the FFN (same value the - fwd handed to :func:`_combine`). It's required by the TRITON - combine_bwd kernel; for PURE_JAX we don't need it but accept it for - a symmetric signature. - """ - # Step 3 inverse: global combine bwd. - d_output_2d = d_output.reshape(-1, d_output.shape[-1]) - if backend is PermutationBackend.PURE_JAX: - # The pure-jax combine is: - # unsort = _sort_activations(expert_outputs, argsort(sorted_indices)) - # if pad: unsort = unsort[:num_real] - # reshape -> einsum BKE,BK -> BE -> reshape to BSE - # Hand-derive the bwd in plain JAX (no custom_vjp involved): - unsort_indices = jnp.argsort(state.sorted_indices) - topk = num_experts_per_tok - num_real = num_real_tokens - padding = padding_size - # Recover the unsorted intermediate that the fwd produced (we - # need it for the d_routing_weights pullback). Apply the same - # gather the fwd did. - unsort_intermediate = expert_outputs[unsort_indices] - if padding > 0: - unsort_intermediate = unsort_intermediate[:num_real] - # Bwd of einsum/reshape: - # output[B, E] = sum_K intermediate[B, K, E] * weights[B, K] - # d_intermediate[B, K, E] = d_output[B, E] * weights[B, K] - # d_weights[B, K] = sum_E d_output[B, E] * intermediate[B, K, E] - rw = state.routing_weights.reshape(-1, topk) - intermediate_3d = unsort_intermediate.reshape(rw.shape[0], topk, -1) - rw_cast = rw.astype(intermediate_3d.dtype) - d_intermediate_3d = jnp.einsum("BE,BK -> BKE", d_output_2d, rw_cast) - d_routing_weights = jnp.einsum("BE,BKE -> BK", d_output_2d, intermediate_3d).astype( - state.routing_weights.dtype - ) - d_routing_weights = d_routing_weights.reshape(state.routing_weights.shape) - d_unsort_intermediate = d_intermediate_3d.reshape(num_real, -1) - # Pad back with zeros if the fwd stripped padding. - if padding > 0: - d_unsort_intermediate = jnp.concatenate( - [ - d_unsort_intermediate, - jnp.zeros( - (padding, d_unsort_intermediate.shape[-1]), - dtype=d_unsort_intermediate.dtype, - ), - ], - axis=0, - ) - # Bwd of the gather is gather-by-original-indices: - # sorted = unsort[argsort(sorted_indices)] - # d_sorted = scatter d_unsort via argsort(sorted_indices) - # = d_unsort[sorted_indices] (gather by original sorted_indices, - # which is the inverse of argsort(sorted_indices)). - d_expert_outputs_global = d_unsort_intermediate[state.sorted_indices] - else: - # TRITON combine bwd: requires fwd_input (expert_outputs). - num_tokens = state.row_id_map.shape[0] - n_experts = (state.row_id_map.shape[1] - 1) // 2 - hidden = d_output_2d.shape[-1] - num_out_tokens = expert_outputs.shape[0] - if state.pad_offsets is not None: - d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs_and_unpad( - d_output_2d, - state.row_id_map, - expert_outputs, - state.merging_probs, - state.pad_offsets, - num_tokens, - n_experts, - num_out_tokens, - hidden, - ) - # The kernel only writes positions tokens map to; padded - # positions may contain NaN. Replace with zeros (matches - # ``_token_combine_bwd_rule``). - d_expert_outputs_global = jnp.where( - jnp.isnan(d_expert_outputs_global), 0.0, d_expert_outputs_global - ) - else: - d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs( - d_output_2d, - state.row_id_map, - expert_outputs, - state.merging_probs, - num_tokens, - n_experts, - num_out_tokens, - hidden, - ) - d_routing_weights = d_merging_probs - - if not ep_active: - return d_expert_outputs_global, d_routing_weights - - # Step 2 (EP) inverse: bwd of reverse ragged_all_to_all is a forward - # ragged_all_to_all using the SAME forward parameters (sender / - # receiver roles swap from the reverse direction back to forward). - in_off_f, send_sz_f, out_off_f, recv_sz_f = compute_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - recv_buf_for_bwd = jnp.zeros(post_a2a_buffer_shape, dtype=d_expert_outputs_global.dtype) - d_x_send_back = jax.lax.ragged_all_to_all( - d_expert_outputs_global, - recv_buf_for_bwd, - in_off_f, - send_sz_f, - out_off_f, - recv_sz_f, - axis_name=ep_axis, - ) - # Step 1 (EP) inverse: combine fwd applied is_forward=False; the - # bwd is is_forward=True with the SAME row_id_map. - recv_buffer_rows, hidden = d_x_send_back.shape - d_expert_outputs, _ = sort_chunks_by_map( - d_x_send_back, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=True, - ) - return d_expert_outputs, d_routing_weights +@flax.struct.dataclass +class _Ctx: + """Residuals carried from the fwd rule into the bwd rule. -def _dispatch_bwd( - d_sorted_x: jnp.ndarray, - state: _DispatchState, - inputs_2d_shape: Tuple[int, ...], - *, - backend: PermutationBackend, - ep_active: bool, - num_experts: int, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # See ``_compute_static_shape_info`` and the note in ``_dispatch`` - # for why these are kwargs rather than state-dict entries. - num_real_tokens: int, - padding_size: int, - pre_a2a_buffer_shape: Tuple[int, int], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> jnp.ndarray: - """Inverse of :func:`_dispatch` on the cotangent. Returns ``d_inputs_2d``. - - The probs path through dispatch is always discarded (PURE_JAX never - threads probs through dispatch; TRITON technically does but the - caller drops ``permuted_probs``, so its cotangent is structurally - zero). The probs gradient instead flows back through - :func:`_combine_bwd`. + Flattened automatically by jax.custom_vjp; ``cfg`` is the only + static field (the rest are jnp.ndarray, GroupedNoScaleTensor, or + None when aux_loss_coeff == 0). """ - if ep_active: - # Step 4 inverse: dispatch fwd applied is_forward=True; bwd is - # is_forward=False with the SAME row_id_map. - recv_buffer_rows, hidden = d_sorted_x.shape - d_x_recv, _ = sort_chunks_by_map( - d_sorted_x, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=False, - ) - # Step 3 inverse: bwd of forward ragged_a2a is the reverse-direction - # ragged_a2a using the SAME params with sender/receiver swapped. - in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - recv_buf_pre = jnp.zeros(pre_a2a_buffer_shape, dtype=d_x_recv.dtype) - d_sorted_x = jax.lax.ragged_all_to_all( - d_x_recv, - recv_buf_pre, - in_off_r, - send_sz_r, - out_off_r, - recv_sz_r, - axis_name=ep_axis, - ) - # Step 1 inverse: global permute bwd. - if backend is PermutationBackend.PURE_JAX: - # Fwd was: replicated = repeat(inputs_2d, topk, axis=0) - # padded = pad(replicated, (0, padding_size)) - # sorted = padded[sorted_indices] - # Bwd: d_padded = scatter via sorted_indices - # = d_sorted[argsort(sorted_indices)] - # d_replicated = d_padded[:num_real] - # d_inputs_2d = d_replicated.reshape(T, topk, H).sum(axis=1) - sorted_indices = state.sorted_indices - num_real = num_real_tokens - padding = padding_size - topk = num_experts_per_tok - unsort_indices = jnp.argsort(sorted_indices) - d_padded = d_sorted_x[unsort_indices] - if padding > 0: - d_replicated = d_padded[:num_real] - else: - d_replicated = d_padded - num_tokens = inputs_2d_shape[0] - hidden = inputs_2d_shape[-1] - d_inputs_2d = d_replicated.reshape(num_tokens, topk, hidden).sum(axis=1) - return d_inputs_2d - - # TRITON: bwd is unpermute_with_mask_map[_and_unpad]. - num_tokens = inputs_2d_shape[0] - hidden = inputs_2d_shape[-1] - if state.pad_offsets is not None: - d_inputs_2d, _ = unpermute_with_mask_map_and_unpad( - d_sorted_x, - state.row_id_map, - None, - None, - state.pad_offsets, - num_tokens, - num_experts, - hidden, - ) - else: - d_inputs_2d, _ = unpermute_with_mask_map( - d_sorted_x, - state.row_id_map, - None, - None, - num_tokens, - num_experts, - hidden, - ) - return d_inputs_2d + x: jnp.ndarray + gate_kernel: jnp.ndarray + expert_bias: jnp.ndarray + logits_2d: jnp.ndarray + saved_scores: jnp.ndarray + routing_map: jnp.ndarray + cfg: Any = flax.struct.field(pytree_node=False) + handle_mem: jnp.ndarray + token_counts: jnp.ndarray + recv_topk_weights: jnp.ndarray + casted_sorted_x_lhs_trans: Any + casted_wi_rhs_trans: Any + gate_proj_out: jnp.ndarray + up_proj_out: jnp.ndarray + casted_intermediate_lhs_trans: Any + casted_wo_rhs_trans: Any + expert_outputs: jnp.ndarray + local_group_sizes: jnp.ndarray + aux_const_buf: Any = None + aux_tokens_per_expert: Any = None + aux_saved_scores: Any = None # ============================================================================= -# Per-shard body +# Per-shard FFN body (runs inside shard_map) # ============================================================================= -def _body_fwd( # pylint: disable=unused-argument - captured: dict, +def _ffn_fwd_per_shard( + recv_tokens_local: jnp.ndarray, + recv_topk_weights_local: jnp.ndarray, + token_counts_local: jnp.ndarray, + wi_0: jnp.ndarray, + wi_1: jnp.ndarray, + wo: jnp.ndarray, + wi_0_bias: Optional[jnp.ndarray], + wi_1_bias: Optional[jnp.ndarray], + wo_bias: Optional[jnp.ndarray], *, - # Statics - num_experts: int, - num_experts_per_tok: int, + num_local_experts: int, activation_type: str, - score_function: ScoreFunction, - use_pre_softmax: bool, - num_groups: Optional[int], - group_topk: Optional[int], - scaling_factor: float, - aux_loss_coeff: float, - permutation_backend: PermutationBackend, - align_size: int, - gate_inside_vjp: bool, - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], - dtype: jnp.dtype, - # EP-only statics - ep_active: bool, - ep_axis: Optional[str], - data_parallelism_axes: Tuple[str, ...], - fsdp_sizes: Tuple[int, ...], - num_ep: int, - num_experts_local: int, - recv_buffer_rows: int, -) -> Tuple[jnp.ndarray, jnp.ndarray, dict]: - """Per-shard forward body. Returns ``(output, aux_loss, ctx_dict)``. - - ``aux_loss`` is always materialized (zeros scalar when disabled) so - the ``shard_map``'s ``out_specs`` has a static structure. + apply_topk_weights_early: bool, +): + """Per-shard FFN forward. + + Operates on the shard-local ``[1, recv_pr, H]`` slice that + ``tex.ep_dispatch`` produces. Returns the expert outputs (shaped + ``[1, recv_pr, H_out]`` so the surrounding ``shard_map`` reassembles + them as ``[num_procs, recv_pr, H_out]``) plus the residuals consumed + by the bwd. + + ``token_counts_local`` (``[1, num_local_experts]``, from + ``tex.ep_prepare``) is passed to ``grouped_gemm`` as ``group_sizes`` + so cuBLAS skips both 0-token-routed experts and the dispatch + overalloc tail. """ - if not gate_inside_vjp: - raise NotImplementedError( - "gate_inside_vjp=False is deferred to a follow-up PR; for now" - " the gate GEMM lives inside the MoE VJP." - ) - - x = captured["inputs"] - gate_kernel = captured["gate_kernel"] - wi_0 = captured["wi_0"] - wi_1 = captured["wi_1"] - wo = captured["wo"] - wi_0_bias = captured.get("wi_0_bias") - wi_1_bias = captured.get("wi_1_bias") - wo_bias = captured.get("wo_bias") - expert_bias = captured.get("expert_bias") - - batch_size, sequence_length, hidden = x.shape - - # ---------------- Stage 1: gate ---------------- - gate_kernel_cast = gate_kernel.astype(x.dtype) - gate_logits = jnp.einsum("bsh,he->bse", x, gate_kernel_cast) - logits_2d = gate_logits.reshape(-1, num_experts) - inputs_2d = x.reshape(-1, hidden) - - # ---------------- Stage 2: routing ---------------- - # Under EP, expert_bias is sharded P(ep_axis); the router needs the - # full E-dim view, so all_gather it. - if ep_active and expert_bias is not None: - full_expert_bias = jax.lax.all_gather(expert_bias, axis_name=ep_axis, tiled=True) - else: - full_expert_bias = expert_bias - # Pass an empty array sentinel when expert_bias is unused (the - # underlying primitive expects a real ndarray, not None). - eb_arg = ( - full_expert_bias if full_expert_bias is not None else jnp.zeros((0,), dtype=jnp.float32) - ) - sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( - logits_2d, - topk=num_experts_per_tok, - use_pre_softmax=use_pre_softmax, - num_groups=-1 if num_groups is None else num_groups, - group_topk=-1 if group_topk is None else group_topk, - scaling_factor=scaling_factor, - score_function=score_function, - expert_bias=eb_arg, - compute_aux_scores=False, - ) - sparse_probs = sparse_probs.astype(dtype) - - # ---------------- Stage 2b: aux loss ---------------- - if aux_loss_coeff > 0.0: - if ep_active: - collective_axes: Any = ( - ep_axis if not data_parallelism_axes else (ep_axis, *data_parallelism_axes) - ) - global_logits_2d = jax.lax.all_gather( - logits_2d, axis_name=collective_axes, axis=0, tiled=True - ) - _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( - global_logits_2d, - topk=num_experts_per_tok, - use_pre_softmax=use_pre_softmax, - num_groups=-1 if num_groups is None else num_groups, - group_topk=-1 if group_topk is None else group_topk, - scaling_factor=scaling_factor, - score_function=score_function, - expert_bias=eb_arg, - compute_aux_scores=False, - ) - aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) - aux_logits_for_score = global_logits_2d - else: - aux_tokens_per_expert = jnp.sum(routing_map.astype(jnp.int32), axis=0) - aux_logits_for_score = logits_2d - # Aux-side scores: clean per-expert scores (no grouped routing, - # no bias). compute_aux_scores=True takes a separate path that - # ignores the grouping knobs. - aux_probs, _aux_routing_map, aux_saved_scores = tex.fused_topk_with_score_function_fwd( - aux_logits_for_score.astype(jnp.float32), - topk=num_experts_per_tok, - use_pre_softmax=False, - num_groups=-1, - group_topk=-1, - scaling_factor=1.0, - score_function=score_function, - expert_bias=jnp.zeros((0,), dtype=jnp.float32), - compute_aux_scores=True, - ) - aux_loss, aux_const_buf = tex.fused_moe_aux_loss_fwd( - aux_probs.astype(jnp.float32), - aux_tokens_per_expert.astype(jnp.int32), - topk=num_experts_per_tok, - coeff=aux_loss_coeff, - ) - else: - aux_loss = jnp.zeros((), dtype=dtype) - aux_const_buf = None - aux_tokens_per_expert = None - aux_logits_for_score = None - aux_saved_scores = None - - # ---------------- Stage 3: dispatch ---------------- - shard_id = jax.lax.axis_index(ep_axis) if ep_active else None - sorted_x, dispatch_state = _dispatch( - inputs_2d, - sparse_probs, - routing_map, - backend=permutation_backend, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - ep_axis=ep_axis, - num_ep=num_ep, - recv_buffer_rows=recv_buffer_rows, - shard_id=shard_id, - ) - local_group_sizes = dispatch_state.group_sizes - - # ---------------- Stage 4: per-expert FFN (inlined) ---------------- - q_set_w0, q_set_w1, q_set_wo = quantizer_sets - if q_set_w0 == noop_quantizer_set: - wi_0 = wi_0.astype(sorted_x.dtype) - if q_set_w1 == noop_quantizer_set: - wi_1 = wi_1.astype(sorted_x.dtype) - if q_set_wo == noop_quantizer_set: - wo = wo.astype(sorted_x.dtype) - - # GEMM 1+2 (fused): up_proj_combined = sorted_x @ wi where - # wi := concat([wi_0, wi_1], axis=-1) -> shape [E, H, 2M] - # combined_out := sorted_x @ wi -> shape [T, 2M] - # Splitting the output back into ``gate_proj_out`` / ``up_proj_out`` - # is free (it's a slicing reshape). This collapses two grouped - # GEMMs and two grouped quantizes of ``sorted_x`` (one per kernel) - # into one of each. Bias is concatenated the same way. - # - # FP8/MXFP8 caveat: per-expert amax is now computed over [H, 2M] - # rather than [H, M] for each of wi_0 / wi_1 separately, so the - # representable range for one of the two halves may shift slightly - # vs. the pre-fusion code. Numerics tests cover this. - inter_M = wi_0.shape[-1] + hidden = recv_tokens_local.shape[-1] + sorted_x = recv_tokens_local.reshape(-1, hidden) + recv_w_flat = recv_topk_weights_local.reshape(-1) + local_group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) + + wi_0 = wi_0.astype(sorted_x.dtype) + wi_1 = wi_1.astype(sorted_x.dtype) + wo = wo.astype(sorted_x.dtype) + + # Concat wi_0/wi_1 along the trailing axis (NOT stack on a new + # axis). grouped_gemm requires the 3D (G, K, N) weight layout with + # contracting_dims=((1,), (1,)); a 4D stack variant walks off the + # end of the RHS and returns NaN. wi_combined = jnp.concatenate([wi_0, wi_1], axis=-1) wi_combined_bias = ( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) - casted_sorted_x = tex.grouped_quantize(sorted_x, q_set_w0.x, local_group_sizes, flatten_axis=-1) - casted_wi = tex.grouped_quantize(wi_combined, q_set_w0.kernel, flatten_axis=-1) + + q_set = noop_quantizer_set + casted_sorted_x = tex.grouped_quantize(sorted_x, q_set.x, local_group_sizes, flatten_axis=-1) + casted_wi = tex.grouped_quantize(wi_combined, q_set.kernel, flatten_axis=-1) combined_out = tex.grouped_gemm( casted_sorted_x.get_tensor(usage=TensorUsage.LHS), casted_wi.get_tensor(usage=TensorUsage.RHS), contracting_dims=((1,), (1,)), bias=wi_combined_bias, ) - gate_proj_out = combined_out[..., :inter_M] - up_proj_out = combined_out[..., inter_M:] + gate_proj_out, up_proj_out = jnp.split(combined_out, 2, axis=-1) casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS) casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS) - if isinstance(casted_sorted_x_lhs_trans, ScaledTensor): - casted_sorted_x_lhs_trans = casted_sorted_x_lhs_trans.checkpoint(q_set_w0.x) - if isinstance(casted_wi_rhs_trans, ScaledTensor): - casted_wi_rhs_trans = casted_wi_rhs_trans.checkpoint(q_set_w0.kernel) - # Activation: intermediate = act(gate_proj_out) * up_proj_out + # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM + # output dtype; the activation output (`intermediate`) stays in the + # dtype the wo GEMM / wo's quantized input consumes. For bf16 compute + # that's all bf16; for FP8/FP4 the downstream grouped_quantize is what + # transitions to the target precision. act_fn = _convert_to_activation_function(activation_type) intermediate = act_fn(gate_proj_out) * up_proj_out - # GEMM 3: expert_outputs = intermediate @ wo + if apply_topk_weights_early: + # Fold the per-token combine weights into the FFN intermediate; + # the downstream wo GEMM is linear so this is equivalent to the + # late-weighting path. Padded recv slots can contain uninitialized + # data, so overwrite inactive rows with literal zeros instead of + # relying on multiplication by a zero mask (IEEE NaN * 0 = NaN). + # ``w_b`` is cast to ``intermediate.dtype`` so the multiply doesn't + # promote expert_outputs above the EP buffer's element width. + w_b = recv_w_flat[:, None].astype(intermediate.dtype) + active = (recv_w_flat != 0)[:, None] + intermediate = jnp.where(active, intermediate * w_b, jnp.zeros_like(intermediate)) + casted_intermediate = tex.grouped_quantize( - intermediate, q_set_wo.x, local_group_sizes, flatten_axis=-1 + intermediate, q_set.x, local_group_sizes, flatten_axis=-1 ) - casted_wo = tex.grouped_quantize(wo, q_set_wo.kernel, flatten_axis=-1) + casted_wo = tex.grouped_quantize(wo, q_set.kernel, flatten_axis=-1) expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), casted_wo.get_tensor(usage=TensorUsage.RHS), @@ -1151,524 +307,148 @@ def _body_fwd( # pylint: disable=unused-argument ) casted_intermediate_lhs_trans = casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS) casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) - if isinstance(casted_intermediate_lhs_trans, ScaledTensor): - casted_intermediate_lhs_trans = casted_intermediate_lhs_trans.checkpoint(q_set_wo.x) - if isinstance(casted_wo_rhs_trans, ScaledTensor): - casted_wo_rhs_trans = casted_wo_rhs_trans.checkpoint(q_set_wo.kernel) - - # ---------------- Stage 5: combine ---------------- - # Compute per-shard static shape info once and pass through both - # _combine and (later) the bwd helpers via kwargs -- never via the - # state dict, which gets pytree-flattened across shard_map and would - # coerce Python ints into JitTracer 0-d arrays. - _static_shape = _compute_static_shape_info( - batch_size=batch_size, - sequence_length=sequence_length, - hidden=hidden, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - num_ep=num_ep, - fsdp_sizes=fsdp_sizes, - recv_buffer_rows=recv_buffer_rows, - ) - # ``expert_outputs_residual`` is the post-A2A FFN-output tensor that - # Step 3 of the combine actually consumed. Saving this (rather than - # the pre-A2A shard-local FFN output) is what makes - # ``_combine_bwd``'s Step-3 inverse see the same value the forward - # Step 3 saw -- otherwise EP + TRITON yields wrong d_expert_outputs. - output, expert_outputs_residual = _combine( - expert_outputs, - dispatch_state, - backend=permutation_backend, - ep_active=ep_active, - batch_size=batch_size, - sequence_length=sequence_length, - dtype=dtype, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) - # ---------------- Build ctx ---------------- - aux_enabled = aux_loss_coeff > 0.0 - ctx = _BodyCtx( - x=x, - gate_kernel=gate_kernel, - logits_2d=logits_2d, - saved_scores=saved_scores, - routing_map=routing_map, - dispatch=dispatch_state, - casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, - casted_wi_rhs_trans=casted_wi_rhs_trans, - gate_proj_out=gate_proj_out, - up_proj_out=up_proj_out, - casted_intermediate_lhs_trans=casted_intermediate_lhs_trans, - casted_wo_rhs_trans=casted_wo_rhs_trans, - expert_outputs=expert_outputs_residual, - local_group_sizes=local_group_sizes, - expert_bias=expert_bias if expert_bias is not None else None, - aux_const_buf=aux_const_buf if aux_enabled else None, - aux_tokens_per_expert=aux_tokens_per_expert if aux_enabled else None, - aux_logits_for_score=aux_logits_for_score if aux_enabled else None, - aux_saved_scores=aux_saved_scores if aux_enabled else None, + expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) + # Reshape local_group_sizes to (1, num_local_experts) so the + # surrounding shard_map can stitch per-shard counts back into the + # global (num_procs, num_local_experts) layout matching token_counts. + local_group_sizes_3d = local_group_sizes.reshape(1, num_local_experts) + residuals = ( + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out, + up_proj_out, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes_3d, ) - - return output, aux_loss, ctx - - -def _body_bwd( # pylint: disable=unused-argument - ctx: _BodyCtx, - dy_pair: Tuple[jnp.ndarray, jnp.ndarray], + return expert_outputs_3d, residuals + + +def _ffn_bwd_per_shard( + d_expert_outputs_local: jnp.ndarray, + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out: jnp.ndarray, + up_proj_out: jnp.ndarray, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes: jnp.ndarray, + recv_topk_weights_local: jnp.ndarray, *, - num_experts: int, - num_experts_per_tok: int, activation_type: str, - score_function: ScoreFunction, - use_pre_softmax: bool, - num_groups: Optional[int], - group_topk: Optional[int], - scaling_factor: float, - aux_loss_coeff: float, - permutation_backend: PermutationBackend, - align_size: int, - gate_inside_vjp: bool, - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], - dtype: jnp.dtype, - ep_active: bool, - ep_axis: Optional[str], - data_parallelism_axes: Tuple[str, ...], - fsdp_sizes: Tuple[int, ...], - num_ep: int, - num_experts_local: int, - recv_buffer_rows: int, - # Static side info (kept here rather than inside ctx because they're - # python flags / shapes, not array leaves): - has_wi_bias: bool, - has_wo_bias: bool, - has_expert_bias: bool, - x_shape: Tuple[int, ...], -) -> dict: - """Per-shard backward body. Returns a dict of grads keyed identically - to the ``captured`` dict consumed by :func:`_body_fwd`.""" - if not gate_inside_vjp: - raise NotImplementedError("gate_inside_vjp=False is deferred to a follow-up PR.") - - d_output, d_aux_loss = dy_pair - # The fused FFN bwd quantizes via ``q_set_w0`` only (one quantize for - # the [E, H, 2M] fused wi tensor and one for the [T, 2M] fused dgrad), - # so ``q_set_w1`` is intentionally unused here. - q_set_w0, _q_set_w1, q_set_wo = quantizer_sets - batch_size, sequence_length, hidden = x_shape - shard_id = jax.lax.axis_index(ep_axis) if ep_active else None - - # Recompute per-shard static shape info from existing statics - # (Python ints / int tuples). Plumbed via kwargs to _combine_bwd - # and _dispatch_bwd -- NOT through the ctx dict, because the - # dict gets pytree-flattened across the bwd shard_map's in_specs - # and Python ints would be coerced into JitTracer 0-d arrays - # (breaking ``if padding > 0`` and ``jnp.zeros(shape)`` callsites). - # ``batch_size`` here is the GLOBAL batch size (captured in - # ``x_shape`` by the outer fwd rule), hence ``batch_is_per_shard=False``. - _static_shape = _compute_static_shape_info( - batch_size=batch_size, - sequence_length=sequence_length, - hidden=hidden, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - num_ep=num_ep, - fsdp_sizes=fsdp_sizes, - recv_buffer_rows=recv_buffer_rows, - batch_is_per_shard=False, - ) - - # Compute per-shard input shape: under the EP shard_map body, the - # gradient tensors live at per-shard shape, so the dispatch_bwd - # reshape target and ``d_x_from_dispatch.reshape(x_shape)`` below - # must use the per-shard shape rather than the captured global - # ``x_shape``. - if ep_active: - dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 - per_shard_batch = batch_size // (num_ep * dp_size) - per_shard_x_shape: Tuple[int, ...] = (per_shard_batch, sequence_length, hidden) - else: - per_shard_x_shape = x_shape - - # ---------------- Combine bwd ---------------- - d_expert_outputs, d_routing_weights = _combine_bwd( - d_output, - ctx.dispatch, - ctx.expert_outputs, - backend=permutation_backend, - ep_active=ep_active, - batch_size=batch_size, - sequence_length=sequence_length, - dtype=dtype, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - post_a2a_buffer_shape=_static_shape.post_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) + apply_topk_weights_early: bool, + has_bias: bool, +): + """Per-shard FFN backward. - # ---------------- FFN bwd: GEMM 3 (wo) ---------------- - casted_d_eo = tex.grouped_quantize( - d_expert_outputs, q_set_wo.dgrad, ctx.local_group_sizes, flatten_axis=-1 - ) + Mirrors :func:`_ffn_fwd_per_shard`. Returns + ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], + d_wi_0, d_wi_1, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. + """ + local_group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) + d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) + recv_w_flat = recv_topk_weights_local.reshape(-1) + q_set = noop_quantizer_set + # cuBLAS grouped_gemm skips size_g == 0 groups without zero-filling + # the output slice; mask 0-token-expert wgrads to zero so the + # optimizer never sees uninit memory. + wgrad_group_active = (local_group_sizes > 0)[:, None, None] + + # wo bwd + casted_d_eo = tex.grouped_quantize(d_eo_2d, q_set.dgrad, local_group_sizes, flatten_axis=-1) + _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) + _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) d_intermediate = tex.grouped_gemm( - casted_d_eo.get_tensor(usage=TensorUsage.LHS), - ctx.casted_wo_rhs_trans, + _casted_d_eo_lhs, + casted_wo_rhs_trans, contracting_dims=((1,), (2,)), ) d_wo = tex.grouped_gemm( - ctx.casted_intermediate_lhs_trans, - casted_d_eo.get_tensor(usage=TensorUsage.RHS), + casted_intermediate_lhs_trans, + _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) - d_wo_bias = tex.grouped_dbias(d_expert_outputs, ctx.local_group_sizes) if has_wo_bias else None + d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) + d_wo_bias = tex.grouped_dbias(d_eo_2d, local_group_sizes) if has_bias else None - # ---------------- Activation bwd ---------------- - # intermediate = act(gate_proj_out) * up_proj_out - # d(gate_proj_out) = vjp(act, gate_proj_out)(d_intermediate * up_proj_out) - # d(up_proj_out) = d_intermediate * act(gate_proj_out) act_fn = _convert_to_activation_function(activation_type) - act_gate_proj_out, dact_gate_proj_pullback = jax.vjp(act_fn, ctx.gate_proj_out) - d_up_proj_out = d_intermediate * act_gate_proj_out - (d_gate_proj_out,) = dact_gate_proj_pullback(d_intermediate * ctx.up_proj_out) - - # ---------------- FFN bwd: GEMM 1+2 fused (wi_0 | wi_1) ---------------- - # Concat the two upstream grads along the output (M) axis, do one - # grouped quantize + one dgrad GEMM + one wgrad GEMM, then split. - # ``ctx.casted_wi_rhs_trans`` has shape [E, H, 2M] from the fwd - # fused quantize, so the dgrad math is: - # d_sorted_x = [d_gate | d_up] @ wi_rhs_trans - # = d_gate @ wi_0^T + d_up @ wi_1^T - inter_M = d_gate_proj_out.shape[-1] + if apply_topk_weights_early: + # intermediate' = intermediate * w * mask. Split the cotangent + # across both factors before the activation bwd consumes it. Padded + # recv slots may still be NaN in the saved activation residuals, so + # use zero-filled residuals on inactive rows before the activation VJP. + w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) + active = (recv_w_flat != 0)[:, None] + gate_proj_for_bwd = jnp.where(active, gate_proj_out, jnp.zeros_like(gate_proj_out)) + up_proj_for_bwd = jnp.where(active, up_proj_out, jnp.zeros_like(up_proj_out)) + intermediate_unweighted = act_fn(gate_proj_for_bwd) * up_proj_for_bwd + d_recv_w_from_intermediate = jnp.sum( + d_intermediate * intermediate_unweighted, + axis=-1, + ).astype(recv_w_flat.dtype) + d_intermediate = jnp.where(active, d_intermediate * w_b, jnp.zeros_like(d_intermediate)) + else: + gate_proj_for_bwd = gate_proj_out + up_proj_for_bwd = up_proj_out + d_recv_w_from_intermediate = jnp.zeros_like(recv_w_flat) + + # Activation bwd, symmetric with the fwd: silu' and the two + # elementwise products run in the GEMM dtype (no fp32 island), so + # the chain rule composes through at the same precision the wi/wo + # GEMMs consume. + act_gp, dact_pullback = jax.vjp(act_fn, gate_proj_for_bwd) + d_up_proj_out = d_intermediate * act_gp + (d_gate_proj_out,) = dact_pullback(d_intermediate * up_proj_for_bwd) + + # wi bwd (fused gate/up via concat). Mirror the fused fwd: pack the + # gate/up cotangents along the trailing axis, run a single + # grouped_quantize + two grouped_gemm pair (one dgrad, one wgrad) + # against the fused casted_wi_rhs_trans residual, then split the + # wgrad result back into d_wi_0 / d_wi_1 halves with jnp.split. d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) casted_d_combined = tex.grouped_quantize( - d_combined, q_set_w0.dgrad, ctx.local_group_sizes, flatten_axis=-1 + d_combined, q_set.dgrad, local_group_sizes, flatten_axis=-1 ) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), - ctx.casted_wi_rhs_trans, + casted_wi_rhs_trans, contracting_dims=((1,), (2,)), ) d_wi_combined = tex.grouped_gemm( - ctx.casted_sorted_x_lhs_trans, + casted_sorted_x_lhs_trans, casted_d_combined.get_tensor(usage=TensorUsage.RHS), contracting_dims=((0,), (0,)), ) - d_wi_0 = d_wi_combined[..., :inter_M] - d_wi_1 = d_wi_combined[..., inter_M:] - if has_wi_bias: - d_wi_combined_bias = tex.grouped_dbias(d_combined, ctx.local_group_sizes) - d_wi_0_bias = d_wi_combined_bias[..., :inter_M] - d_wi_1_bias = d_wi_combined_bias[..., inter_M:] + d_wi_combined = jnp.where(wgrad_group_active, d_wi_combined, jnp.zeros_like(d_wi_combined)) + d_wi_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1) + if has_bias: + d_wi_combined_bias = tex.grouped_dbias(d_combined, local_group_sizes) + d_wi_0_bias, d_wi_1_bias = jnp.split(d_wi_combined_bias, 2, axis=-1) else: d_wi_0_bias = None d_wi_1_bias = None - # ---------------- Dispatch bwd ---------------- - inputs_2d_shape = (per_shard_x_shape[0] * per_shard_x_shape[1], hidden) - d_inputs_2d = _dispatch_bwd( - d_sorted_x, - ctx.dispatch, - inputs_2d_shape=inputs_2d_shape, - backend=permutation_backend, - ep_active=ep_active, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) - d_x_from_dispatch = d_inputs_2d.reshape(per_shard_x_shape) - - # ---------------- Routing bwd ---------------- - # The probs cotangent comes from _combine_bwd. For PURE_JAX it's the - # cotangent of routing_weights (post-routing_map_to_selected_experts); - # we need to bridge back to sparse_probs. For TRITON it's already the - # cotangent of merging_probs == sparse_probs. - if d_routing_weights is not None: - if permutation_backend is PermutationBackend.PURE_JAX: - # routing_map_to_selected_experts: - # selected_experts = argsort(routing_map)[..., -topk:] - # weights = take_along_axis(sparse_probs, selected_experts, axis=-1) - # routing_map is bool (non-diff); the gradient of weights - # w.r.t. sparse_probs is a scatter-into-zero along the - # selected_experts indices. - selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -num_experts_per_tok:] - d_sparse_probs = jnp.zeros_like(ctx.saved_scores).astype(d_routing_weights.dtype) - d_sparse_probs = jnp.take_along_axis(d_sparse_probs, selected_experts, axis=-1) - # Actually scatter: build via jnp.zeros + .at[].set - d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_routing_weights.dtype) - d_sparse_probs = d_sparse_probs.at[ - jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts - ].set(d_routing_weights) - else: - d_sparse_probs = d_routing_weights.astype(jnp.float32) - else: - d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=jnp.float32) - - # Topk bwd primitive: returns d_logits (no d_expert_bias). - d_logits_2d_main = tex.fused_topk_with_score_function_bwd( - ctx.routing_map, - ctx.saved_scores, - d_sparse_probs.astype(ctx.saved_scores.dtype), - topk=num_experts_per_tok, - use_pre_softmax=use_pre_softmax, - scaling_factor=scaling_factor, - score_function=score_function, - compute_aux_scores=False, - ) - - # ---------------- Aux loss bwd ---------------- - if aux_loss_coeff > 0.0: - # Step 1: aux_loss bwd -> d_aux_probs - aux_num_tokens = ctx.aux_logits_for_score.shape[0] - d_aux_probs = tex.fused_moe_aux_loss_bwd( - ctx.aux_const_buf, - ctx.aux_tokens_per_expert.astype(jnp.int32), - d_aux_loss.reshape(()), - num_tokens=aux_num_tokens, - ) - # Step 2: aux-side topk bwd (compute_aux_scores=True path). - # The routing_map argument is ignored in this branch (the kernel - # uses saved_scores); pass any shape-correct integer tensor. - d_aux_logits = tex.fused_topk_with_score_function_bwd( - jnp.zeros(ctx.aux_logits_for_score.shape, dtype=jnp.bool_), - ctx.aux_saved_scores, - d_aux_probs.astype(ctx.aux_saved_scores.dtype), - topk=num_experts_per_tok, - use_pre_softmax=False, - scaling_factor=1.0, - score_function=score_function, - compute_aux_scores=True, - ) - # Step 3: under EP the aux logits were all_gathered along - # ``(ep_axis, *data_parallelism_axes)`` (the latter being FSDP - # axes that shard the batch). The bwd is the inverse of that - # multi-axis tiled all_gather: ``dynamic_slice`` to pick out - # this shard's local rows from the global cotangent. - # - # JAX's convention for tiled ``all_gather(axis_name=(a, b, ...))`` - # is row-major over the tuple: the shard at mesh position - # ``(i_a, i_b, ...)`` writes to rows - # ``[(i_a * size_b * ... + i_b * ... + ...) * local_T : - # + local_T)``. We invert that by computing the same flat - # index here and slicing. - if ep_active: - local_T_aux = ctx.logits_2d.shape[0] - flat_shard = shard_id # ep is the outermost axis in the gather tuple - for ax, sz in zip(data_parallelism_axes, fsdp_sizes): - flat_shard = flat_shard * sz + jax.lax.axis_index(ax) - d_aux_logits_local = jax.lax.dynamic_slice( - d_aux_logits.astype(ctx.logits_2d.dtype), - start_indices=(flat_shard * local_T_aux, 0), - slice_sizes=(local_T_aux, num_experts), - ) - else: - d_aux_logits_local = d_aux_logits.astype(d_logits_2d_main.dtype) - d_logits_2d = d_logits_2d_main + d_aux_logits_local.astype(d_logits_2d_main.dtype) - else: - d_logits_2d = d_logits_2d_main - - # ---------------- Gate bwd ---------------- - d_gate_logits = d_logits_2d.reshape(per_shard_x_shape[0], per_shard_x_shape[1], num_experts) - gate_kernel_cast = ctx.gate_kernel.astype(ctx.x.dtype) - d_x_from_gate = jnp.einsum("bse,he->bsh", d_gate_logits, gate_kernel_cast) - d_gate_kernel = jnp.einsum("bsh,bse->he", ctx.x, d_gate_logits).astype(ctx.gate_kernel.dtype) - d_x = d_x_from_gate + d_x_from_dispatch - - # Reduce per-rank partial contributions to match the out_specs - # declared by _build_grads_specs: - # gate_kernel : P() -> psum across (ep, *fsdp) - # wi_0/wi_1/wo : P(ep_axis, ...) -> psum across (*fsdp) only - # inputs : P((ep, fsdp), ...) -> already shard-local, no reduction - if ep_active: - replicate_all = (ep_axis,) + tuple(data_parallelism_axes) - d_gate_kernel = jax.lax.psum(d_gate_kernel, axis_name=replicate_all) - if data_parallelism_axes: - replicate_fsdp = tuple(data_parallelism_axes) - d_wi_0 = jax.lax.psum(d_wi_0, axis_name=replicate_fsdp) - d_wi_1 = jax.lax.psum(d_wi_1, axis_name=replicate_fsdp) - d_wo = jax.lax.psum(d_wo, axis_name=replicate_fsdp) - if has_wi_bias: - d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=replicate_fsdp) - d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=replicate_fsdp) - if has_wo_bias: - d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=replicate_fsdp) - - grads: dict = { - "inputs": d_x, - "gate_kernel": d_gate_kernel, - "wi_0": d_wi_0, - "wi_1": d_wi_1, - "wo": d_wo, - } - if has_wi_bias: - grads["wi_0_bias"] = d_wi_0_bias - grads["wi_1_bias"] = d_wi_1_bias - if has_wo_bias: - grads["wo_bias"] = d_wo_bias - if has_expert_bias: - # expert_bias has no gradient through topk (the topk bwd returns - # None for it). Emit a structural zero so the outer rule has - # something to package. - grads["expert_bias"] = jnp.zeros_like(ctx.expert_bias) - return grads - - -# ============================================================================= -# Spec builders for shard_map (lockstep with ctx_dict / captured_dict) -# ============================================================================= - - -def _build_in_specs( - ep_axis: str, - batch_pspec_axis: Any, - *, - has_bias: bool, - has_expert_bias: bool, -) -> dict: - """Build the ``in_specs`` dict for the EP fwd shard_map.""" - specs: dict = { - "inputs": P(batch_pspec_axis, None, None), - "gate_kernel": P(), - "wi_0": P(ep_axis, None, None), - "wi_1": P(ep_axis, None, None), - "wo": P(ep_axis, None, None), - } - if has_bias: - for name in ("wi_0_bias", "wi_1_bias", "wo_bias"): - specs[name] = P(ep_axis, None) - if has_expert_bias: - specs["expert_bias"] = P(ep_axis) - return specs - - -def _build_dispatch_specs( # pylint: disable=unused-argument - ep_axis: str, - *, - backend: PermutationBackend, - ep_active: bool, - align_size: int, -) -> _DispatchState: - """Build the shard_map ``out_specs`` for the dispatch state. - - Returns a :data:`_DispatchState` (either :class:`_PureJaxDispatchState` - or :class:`_TritonDispatchState`) whose fields are - :class:`PartitionSpec` placeholders. Optional fields are set to - ``P()`` when populated by :func:`_dispatch` and to ``None`` when - intentionally omitted, so the spec's pytree structure mirrors the - value's structure leaf-for-leaf. - """ - ep_all = P() if ep_active else None - ep_local = P() if ep_active else None - if backend is PermutationBackend.PURE_JAX: - return _PureJaxDispatchState( - group_sizes=P(), - sorted_indices=P(), - routing_weights=P(), - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - return _TritonDispatchState( - group_sizes=P(), - row_id_map=P(), - pad_offsets=P() if align_size > 0 else None, - merging_probs=P(), - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - - -def _build_ctx_specs( # pylint: disable=unused-argument - ep_axis: str, - batch_pspec_axis: Any, - *, - backend: PermutationBackend, - ep_active: bool, - has_bias: bool, - has_expert_bias: bool, - aux_loss_enabled: bool, - align_size: int, -) -> _BodyCtx: - """Build the spec :class:`_BodyCtx` mirroring :func:`_body_fwd`'s ctx. - - Fields gated off by the static config (``expert_bias``, ``aux_*``) - are ``None`` here so the spec pytree matches the value pytree - leaf-for-leaf. - """ - return _BodyCtx( - # Per-shard local activations along the batch axis. - x=P(batch_pspec_axis, None, None), - gate_kernel=P(), - logits_2d=P(batch_pspec_axis, None), - saved_scores=P(batch_pspec_axis, None), - routing_map=P(batch_pspec_axis, None), - dispatch=_build_dispatch_specs( - ep_axis, backend=backend, ep_active=ep_active, align_size=align_size - ), - # FFN residuals: the LHS_TRANS / RHS_TRANS variants of - # grouped_quantize have leading "rows"/"experts" dims that are - # already shard-local (post-dispatch). Use P(ep_axis,...) on - # leading dim; that works whether the leaf is a plain ndarray - # or a ScaledTensor (shard_map applies the spec leaf-wise to - # the registered ScaledTensor pytree). - casted_sorted_x_lhs_trans=P(), - casted_wi_rhs_trans=P(ep_axis, None, None), - gate_proj_out=P(), - up_proj_out=P(), - casted_intermediate_lhs_trans=P(), - casted_wo_rhs_trans=P(ep_axis, None, None), - expert_outputs=P(), - local_group_sizes=P(), - expert_bias=P(ep_axis) if has_expert_bias else None, - aux_const_buf=P() if aux_loss_enabled else None, - aux_tokens_per_expert=P() if aux_loss_enabled else None, - aux_logits_for_score=P() if aux_loss_enabled else None, - aux_saved_scores=P() if aux_loss_enabled else None, - ) - - -def _build_grads_specs( - ep_axis: str, - batch_pspec_axis: Any, - *, - has_bias: bool, - has_expert_bias: bool, -) -> dict: - """Spec dict for the grads dict returned by :func:`_body_bwd`.""" - return _build_in_specs( - ep_axis, - batch_pspec_axis, - has_bias=has_bias, - has_expert_bias=has_expert_bias, + d_sorted_x_3d = d_sorted_x.reshape(1, d_sorted_x.shape[0], d_sorted_x.shape[1]) + d_recv_w_3d = d_recv_w_from_intermediate.reshape(1, -1) + return ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, ) # ============================================================================= -# Top-level VJP rules +# Full fwd / bwd rules (custom_vjp halves) # ============================================================================= -def _moe_fwd_rule( # pylint: disable=unused-argument - # Args MUST match the positional order of ``_moe`` (diff first, - # then nondiff). See ``_moe_bwd_rule`` for the opposite convention. +def _moe_fwd_rule( x, gate_kernel, wi_0, @@ -1687,170 +467,328 @@ def _moe_fwd_rule( # pylint: disable=unused-argument group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, ): - x = with_sharding_constraint_by_logical_axes(x, input_axes) - ep_active = ep_axis is not None - body_kwargs = { - "num_experts": num_experts, - "num_experts_per_tok": num_experts_per_tok, - "activation_type": activation_type, - "score_function": score_function, - "use_pre_softmax": use_pre_softmax, - "num_groups": num_groups, - "group_topk": group_topk, - "scaling_factor": scaling_factor, - "aux_loss_coeff": aux_loss_coeff, - "permutation_backend": permutation_backend, - "align_size": align_size, - "gate_inside_vjp": gate_inside_vjp, - "quantizer_sets": quantizer_sets, - "dtype": dtype, - "ep_axis": ep_axis, - "data_parallelism_axes": data_parallelism_axes, - } - captured: dict = { - "inputs": x, - "gate_kernel": gate_kernel, - "wi_0": wi_0, - "wi_1": wi_1, - "wo": wo, - } - has_bias = wi_0_bias is not None - has_expert_bias = expert_bias is not None - if has_bias: - captured["wi_0_bias"] = wi_0_bias - captured["wi_1_bias"] = wi_1_bias - captured["wo_bias"] = wo_bias - if has_expert_bias: - captured["expert_bias"] = expert_bias - - if not ep_active: - output, aux_loss, ctx = _body_fwd( - captured, - **body_kwargs, - ep_active=False, - fsdp_sizes=(), - num_ep=1, - num_experts_local=num_experts, - recv_buffer_rows=0, - ) - # Carry static side info to the bwd rule alongside ctx. These - # are Python ints/bools/tuples (NOT pytree leaves), so we - # bundle them as a plain dict rather than putting them on the - # ``_BodyCtx`` NamedTuple where shard_map would try to flatten - # them into JitTracers. - static = { - "has_wi_bias": has_bias, - "has_wo_bias": has_bias, - "has_expert_bias": has_expert_bias, - "x_shape": x.shape, - "num_experts_local": num_experts, - "recv_buffer_rows": 0, - } - return (output, aux_loss), (ctx, static) - - # ---------------- EP path ---------------- + """Forward: gate -> topk -> ep_dispatch -> shard_map(FFN) -> ep_combine. + + Returns ``(output, aux_loss)``. ``aux_loss`` is a zero scalar when + ``aux_loss_coeff == 0``. + """ + del gate_kernel_axes, wi_kernel_axes, wo_kernel_axes # used in bwd only from jax.experimental.shard_map import shard_map + x = with_sharding_constraint_by_logical_axes(x, input_axes) + mesh = _get_mesh() if mesh is None or mesh.empty: - raise ValueError("moe(...) requires an active jax.sharding.Mesh when ep_axis is set.") + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + if ep_axis is None: + raise ValueError("moe(...) requires ep_axis to be set (TE EP backend).") num_ep = mesh.shape[ep_axis] if num_experts % num_ep != 0: raise ValueError(f"num_experts={num_experts} must be divisible by EP size={num_ep}") - num_experts_local = num_experts // num_ep + num_local_experts = num_experts // num_ep - # Reject overlapping EP / FSDP axes. Listing ep_axis in - # data_parallelism_axes would produce a duplicate-axis PartitionSpec - # ((ep, ep, ...)) which JAX rejects, and would also double-count - # num_ep in dp_size (under-sizing recv_buffer_rows by a factor of - # num_ep). Catch it up front with a clear error. + dp_size = 1 for ax in data_parallelism_axes: - if ax not in mesh.shape: - raise ValueError( - f"data_parallelism_axes contains {ax!r} but mesh has" - f" axes {tuple(mesh.shape.keys())}" - ) - if ax == ep_axis: - raise ValueError( - f"data_parallelism_axes={data_parallelism_axes!r} contains the EP" - f" axis {ep_axis!r}; EP is implicit in the batch sharding and must" - " not also be listed as a data-parallel axis." - ) + dp_size *= mesh.shape[ax] + num_procs = num_ep * dp_size + + B, S, H = x.shape + K = num_experts_per_tok + if B % num_procs != 0: + raise ValueError(f"batch={B} not divisible by ep*dp={num_procs}") + + # Per-rank send capacity: B/num_procs rows x S tokens per rank. + max_tokens_per_rank = (B // num_procs) * S + # Per-rank receive capacity. NCCL EP HT expert-major lays out variable + # per-expert zones in one flat recv buffer, with each non-empty zone padded + # to ``dispatch_output_per_expert_alignment``. + tokens_per_ep_group = num_ep * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(K, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts + aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + ) + recv_pr = min(per_expert_bound, aligned_total_bound) + + _te_ep_assert_compatible_bootstrap( + num_experts=num_experts, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_pr, + hidden_dim=H, + ep_size=num_ep, + ) if not data_parallelism_axes: batch_pspec_axis: Any = ep_axis else: - batch_pspec_axis = (ep_axis, *data_parallelism_axes) - dp_size = 1 - for ax in data_parallelism_axes: - dp_size *= mesh.shape[ax] + # ep must be innermost: ep_bootstrap forms NCCL EP comms from + # consecutive global ranks (dp_color = rank // ep_size), so the + # comm only stays within one model replica under (outer_dp, ep). + batch_pspec_axis = (*data_parallelism_axes, ep_axis) + ep3_spec = P(batch_pspec_axis, None, None) + ep2_spec = P(batch_pspec_axis, None) + x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, ep3_spec)) + + # ---------------- Gate (global view) ---------------- + # tex.fused_topk_with_score_function is only validated against its + # pytorch reference at fp32 (see tests/pytorch/test_fused_router.py: + # parametrize gates dtype on torch.float32 only; the tolerance helper + # raises NotImplementedError for any other dtype). Keeping logits in + # the activation dtype (e.g. bf16) lets sigmoid / softmax / topk + # accumulate at low precision and silently produce NaNs on tokens + # whose normalised weights underflow. Cast to fp32 here to stay in + # the validated regime. + gate_kernel_cast = gate_kernel.astype(x.dtype) + gate_logits = jnp.einsum("bsh,he->bse", x, gate_kernel_cast) + logits_2d = gate_logits.reshape(-1, num_experts).astype(jnp.float32) + + # ---------------- Routing (global view) ---------------- + # expert_bias is an empty (shape-(0,)) sentinel when the caller did + # not enable it; the primitive treats that as "no bias". + eb_arg = expert_bias if expert_bias.shape != (0,) else jnp.zeros((0,), dtype=jnp.float32) + sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( + logits_2d, + topk=K, + use_pre_softmax=use_pre_softmax, + num_groups=-1 if num_groups is None else num_groups, + group_topk=-1 if group_topk is None else group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=eb_arg, + compute_aux_scores=False, + ) + sparse_probs = sparse_probs.astype(dtype) - global_batch_size, sequence_length, _hidden = x.shape - topk = num_experts_per_tok - if global_batch_size % (num_ep * dp_size) != 0: - raise ValueError(f"batch={global_batch_size} not divisible by ep*dp={num_ep * dp_size}") - recv_buffer_rows = (global_batch_size // dp_size) * sequence_length * topk - if align_size > 0: - recv_buffer_rows += num_experts * (align_size - 1) + # ---------------- Aux loss (global view, replicated) ---------------- + # ``fused_moe_aux_loss_fwd`` sums probs and tokens_per_expert across + # all tokens, which is wrong when T is sharded. Force-replicate the + # gate logits and recompute the routing map at global view so the + # kernel sees a complete [T_global, E] tensor. The replication is a + # single all-gather over (*dp, ep) and lives off the dispatch + # critical path. + if aux_loss_coeff > 0.0: + global_logits_2d = jax.lax.with_sharding_constraint(logits_2d, NamedSharding(mesh, P())) + _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( + global_logits_2d, + topk=K, + use_pre_softmax=use_pre_softmax, + num_groups=-1 if num_groups is None else num_groups, + group_topk=-1 if group_topk is None else group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=eb_arg, + compute_aux_scores=False, + ) + aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) + # compute_aux_scores=True takes a separate kernel path: clean + # per-expert softmax, no grouping / bias / scaling. + aux_probs, _aux_rm, aux_saved_scores = tex.fused_topk_with_score_function_fwd( + global_logits_2d.astype(jnp.float32), + topk=K, + use_pre_softmax=False, + num_groups=-1, + group_topk=-1, + scaling_factor=1.0, + score_function=score_function, + expert_bias=jnp.zeros((0,), dtype=jnp.float32), + compute_aux_scores=True, + ) + aux_loss, aux_const_buf = tex.fused_moe_aux_loss_fwd( + aux_probs.astype(jnp.float32), + aux_tokens_per_expert.astype(jnp.int32), + topk=K, + coeff=aux_loss_coeff, + ) + aux_loss = aux_loss.astype(dtype) + else: + aux_loss = jnp.zeros((), dtype=dtype) + aux_const_buf = None + aux_tokens_per_expert = None + aux_saved_scores = None - in_specs = _build_in_specs( - ep_axis, - batch_pspec_axis, - has_bias=has_bias, - has_expert_bias=has_expert_bias, + # ---------------- Routing -> (topk_idx, topk_w) at 3D ---------------- + # argsort on a bool tensor places True last (False=0 < True=1), so the + # last K indices are the selected expert IDs. + selected_experts = jnp.argsort(routing_map, axis=-1)[..., -K:] + routing_weights = jnp.take_along_axis(sparse_probs, selected_experts, axis=-1) + topk_idx_3d = selected_experts.reshape(B, S, K).astype(jnp.int32) + topk_w_3d = routing_weights.reshape(B, S, K).astype(jnp.float32) + # tex.ep_prepare/dispatch's partition only folds ep_axis into a replicated + # leading dim, not the outer dp/fsdp axes, so a replicated topk_idx makes + # each rank see B/ep rows (not B/num_procs) and overrun the bootstrap-sized + # send buffer. Pin both routing tensors to the (outer, ep) leading sharding + # so per-rank token counts match max_tokens_per_rank. + topk_idx_3d = jax.lax.with_sharding_constraint(topk_idx_3d, NamedSharding(mesh, ep3_spec)) + topk_w_3d = jax.lax.with_sharding_constraint(topk_w_3d, NamedSharding(mesh, ep3_spec)) + + # ---------------- TE EP dispatch (global view) ---------------- + cfg = tex.EpLayerConfig( + top_k=K, + dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) - output_spec = P(batch_pspec_axis, None, None) - aux_spec = P() - ctx_spec = _build_ctx_specs( - ep_axis, - batch_pspec_axis, - backend=permutation_backend, - ep_active=True, - has_bias=has_bias, - has_expert_bias=has_expert_bias, - aux_loss_enabled=(aux_loss_coeff > 0.0), - align_size=align_size, + token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( + cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr + ) + recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3_spec)) + recv_topk_weights = jax.lax.with_sharding_constraint( + recv_topk_weights, NamedSharding(mesh, ep2_spec) + ) + + # ---------------- FFN (per-shard via shard_map) ---------------- + has_bias = wi_0_bias is not None + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) if has_bias else None + # token_counts is the per-shard (1, num_local_experts) padded + # per-expert count from ep_prepare; piped into _ffn_fwd_per_shard + # as the grouped_gemm group_sizes so cuBLAS skips both 0-token + # experts and the trailing overalloc tail. + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec, kernel_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi_0, wi_1, wo] + if has_bias: + ffn_in_specs = ffn_in_specs + (bias_spec, bias_spec, bias_spec) + ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) + + # FFN residuals live entirely on the local ep rank, so the leading + # "experts" / "rows" dims map to P() (already shard-local). wi is + # fused via jnp.concatenate along the trailing (output) axis + # (see _ffn_fwd_per_shard for rationale), so the residual is a + # single 3D casted_wi_rhs_trans of shape + # (num_local_experts, hidden, 2*H_inter). local_group_sizes is + # now per-shard dynamic (= per-shard token_counts), so its + # residual spec mirrors ep2_spec (one row per ep rank). + residuals_spec = ( + P(), # casted_sorted_x_lhs_trans + P(ep_axis, None, None), # casted_wi_rhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans + P(ep_axis, None, None), # casted_wo_rhs_trans + ep2_spec, # local_group_sizes (1, num_local_experts) per shard ) + out_specs = (ep3_spec, residuals_spec) - _fsdp_sizes: Tuple[int, ...] = tuple(mesh.shape[ax] for ax in data_parallelism_axes) - - def _shardmap_body(captured_local): - return _body_fwd( - captured_local, - **body_kwargs, - ep_active=True, - fsdp_sizes=_fsdp_sizes, - num_ep=num_ep, - num_experts_local=num_experts_local, - recv_buffer_rows=recv_buffer_rows, + def _body(*args): + if has_bias: + (r_tok, r_w, tc, w0, w1, w_o, w0b, w1b, wob) = args + else: + (r_tok, r_w, tc, w0, w1, w_o) = args + w0b = w1b = wob = None + # NOTE: tex.ep_dispatch_fwd's NCCL EP HT path leaves the recv + # buffer uninitialised on fully-empty-receiver ranks (and at + # padded slots on partially-loaded ranks). We don't need a + # zero-init guard here anymore because: + # 1. ``tc`` (per-expert padded counts) is plumbed into + # grouped_gemm as group_sizes, so cuBLAS skips both + # 0-token experts and the trailing overalloc tail. + # 2. The per-group wgrad masks in _ffn_bwd_per_shard zero + # ``d_wo`` / ``d_wi_combined`` slices for 0-token-globally + # experts (cuBLAS skips size_g==0 groups without + # zero-filling, which would otherwise leak NaN into the + # user's optimizer). + # 3. All other downstream consumers (ep_combine, + # ep_dispatch_bwd) are handle_mem-aware and read only + # valid positions. + # If a future caller adds a non-group-aware reader of r_tok + # (e.g. an inspect probe over the full recv tile), re-add the + # ``jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like)`` + # guard here. + return _ffn_fwd_per_shard( + r_tok, + r_w, + tc, + w0, + w1, + w_o, + w0b, + w1b, + wob, + num_local_experts=num_local_experts, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, ) - output, aux_loss, ctx = shard_map( - _shardmap_body, + expert_outputs, ffn_residuals = shard_map( + _body, mesh=mesh, - in_specs=(in_specs,), - out_specs=(output_spec, aux_spec, ctx_spec), + in_specs=ffn_in_specs, + out_specs=out_specs, check_rep=False, - )(captured) + )(*ffn_in_args) + expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) + + # ---------------- TE EP combine (global view) ---------------- + out_partition_spec = (batch_pspec_axis, None, None) + if apply_topk_weights_early: + # expert_outputs is already weighted upstream. + output = tex.ep_combine_fwd( + cfg, + handle_mem, + expert_outputs, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + else: + # HT combine is unweighted; apply routing weights before calling it. + # Padded recv slots are ignored by combine via handle_mem metadata. + w = recv_topk_weights[..., None].astype(expert_outputs.dtype) + weighted = expert_outputs * w + output = tex.ep_combine_fwd( + cfg, + handle_mem, + weighted, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + + ( + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out, + up_proj_out, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes, + ) = ffn_residuals + + ctx = _Ctx( + x=x, + gate_kernel=gate_kernel, + expert_bias=expert_bias, + logits_2d=logits_2d, + saved_scores=saved_scores, + routing_map=routing_map, + cfg=cfg, + handle_mem=handle_mem, + token_counts=token_counts, + recv_topk_weights=recv_topk_weights, + casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, + casted_wi_rhs_trans=casted_wi_rhs_trans, + gate_proj_out=gate_proj_out, + up_proj_out=up_proj_out, + casted_intermediate_lhs_trans=casted_intermediate_lhs_trans, + casted_wo_rhs_trans=casted_wo_rhs_trans, + expert_outputs=expert_outputs, + local_group_sizes=local_group_sizes, + aux_const_buf=aux_const_buf, + aux_tokens_per_expert=aux_tokens_per_expert, + aux_saved_scores=aux_saved_scores, + ) static = { - "has_wi_bias": has_bias, - "has_wo_bias": has_bias, - "has_expert_bias": has_expert_bias, + "has_bias": has_bias, "x_shape": x.shape, - "num_experts_local": num_experts_local, - "recv_buffer_rows": recv_buffer_rows, + "recv_pr": recv_pr, } return (output, aux_loss), (ctx, static) @@ -1865,128 +803,259 @@ def _moe_bwd_rule( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, - ctx, - dy_pair, + apply_topk_weights_early, + residuals, + cotangents, ): - ctx, static = ctx # split tensor residuals from static side info - has_wi_bias = static["has_wi_bias"] - has_wo_bias = static["has_wo_bias"] - has_expert_bias = static["has_expert_bias"] - x_shape = static["x_shape"] - num_experts_local = static["num_experts_local"] - recv_buffer_rows = static["recv_buffer_rows"] + """Backward mirror of :func:`_moe_fwd_rule`.""" + del num_groups, group_topk, dtype # captured in residuals / unused in bwd + from jax.experimental.shard_map import shard_map - ep_active = ep_axis is not None - mesh = _get_mesh() if ep_active else None - fsdp_sizes: Tuple[int, ...] = ( - tuple(mesh.shape[ax] for ax in data_parallelism_axes) if ep_active else () - ) - body_kwargs = { - "num_experts": num_experts, - "num_experts_per_tok": num_experts_per_tok, - "activation_type": activation_type, - "score_function": score_function, - "use_pre_softmax": use_pre_softmax, - "num_groups": num_groups, - "group_topk": group_topk, - "scaling_factor": scaling_factor, - "aux_loss_coeff": aux_loss_coeff, - "permutation_backend": permutation_backend, - "align_size": align_size, - "gate_inside_vjp": gate_inside_vjp, - "quantizer_sets": quantizer_sets, - "dtype": dtype, - "ep_axis": ep_axis, - "data_parallelism_axes": data_parallelism_axes, - "fsdp_sizes": fsdp_sizes, - "num_ep": 1 if not ep_active else mesh.shape[ep_axis], - "num_experts_local": num_experts_local, - "recv_buffer_rows": recv_buffer_rows, - "has_wi_bias": has_wi_bias, - "has_wo_bias": has_wo_bias, - "has_expert_bias": has_expert_bias, - "x_shape": x_shape, - } + d_output, d_aux_loss = cotangents - if not ep_active: - grads = _body_bwd(ctx, dy_pair, ep_active=False, **body_kwargs) - # Apply sharding constraints on grads. - grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( - grads["gate_kernel"], gate_kernel_axes - ) - grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) - grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) - grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) - grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) - return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + ctx, static = residuals + has_bias = static["has_bias"] + x_shape = static["x_shape"] + recv_pr = static["recv_pr"] - from jax.experimental.shard_map import shard_map + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + dp_size = 1 + for ax in data_parallelism_axes: + dp_size *= mesh.shape[ax] + B, S, _ = x_shape + K = num_experts_per_tok if not data_parallelism_axes: batch_pspec_axis: Any = ep_axis else: - batch_pspec_axis = (ep_axis, *data_parallelism_axes) - ctx_spec = _build_ctx_specs( - ep_axis, - batch_pspec_axis, - backend=permutation_backend, - ep_active=True, - has_bias=has_wi_bias, - has_expert_bias=has_expert_bias, - aux_loss_enabled=(aux_loss_coeff > 0.0), - align_size=align_size, + batch_pspec_axis = (*data_parallelism_axes, ep_axis) + ep3_spec = P(batch_pspec_axis, None, None) + ep2_spec = P(batch_pspec_axis, None) + out_partition_spec = (batch_pspec_axis, None, None) + + # ---------------- Combine bwd (global view) ---------------- + d_output = jax.lax.with_sharding_constraint(d_output, NamedSharding(mesh, ep3_spec)) + grad_pre_combine = tex.ep_combine_bwd(ctx.cfg, ctx.handle_mem, d_output, recv_pr) + grad_pre_combine = jax.lax.with_sharding_constraint( + grad_pre_combine, NamedSharding(mesh, ep3_spec) ) - dy_specs = (P(batch_pspec_axis, None, None), P()) - grads_spec = _build_grads_specs( - ep_axis, batch_pspec_axis, has_bias=has_wi_bias, has_expert_bias=has_expert_bias + + if apply_topk_weights_early: + # combine_fwd consumed already-weighted expert_outputs; the recv_w + # cotangent flows through the early-weighting step inside the FFN bwd. + d_expert_outputs = grad_pre_combine + d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) + else: + # Reverse the late-weighting multiply. Padded expert-major rows are + # part of the physical grouped-GEMM ranges, so write literal zero + # cotangents for inactive rows instead of relying on NaN * 0. + w = ctx.recv_topk_weights[..., None].astype(grad_pre_combine.dtype) + mask_bool = (ctx.recv_topk_weights != 0)[..., None] + d_expert_outputs = jnp.where( + mask_bool, grad_pre_combine * w, jnp.zeros_like(grad_pre_combine) + ) + d_recv_w_from_combine = (grad_pre_combine * ctx.expert_outputs).sum(axis=-1) + d_recv_w_from_combine = d_recv_w_from_combine.astype(ctx.recv_topk_weights.dtype) + + # ---------------- FFN bwd (per-shard via shard_map) ---------------- + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) if has_bias else None + + bwd_in_specs = ( + ep3_spec, # d_expert_outputs + P(), # casted_sorted_x_lhs_trans + P(ep_axis, None, None), # casted_wi_rhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans + P(ep_axis, None, None), # casted_wo_rhs_trans + ep2_spec, # local_group_sizes (1, num_local_experts) per shard + ep2_spec, # recv_topk_weights + ) + bwd_in_args = [ + d_expert_outputs, + ctx.casted_sorted_x_lhs_trans, + ctx.casted_wi_rhs_trans, + ctx.gate_proj_out, + ctx.up_proj_out, + ctx.casted_intermediate_lhs_trans, + ctx.casted_wo_rhs_trans, + ctx.local_group_sizes, + ctx.recv_topk_weights, + ] + bwd_out_specs = ( + ep3_spec, # d_sorted_x + ep2_spec, # d_recv_w_from_intermediate + kernel_spec, # d_wi_0 + kernel_spec, # d_wi_1 + kernel_spec, # d_wo + bias_spec if has_bias else None, # d_wi_0_bias + bias_spec if has_bias else None, # d_wi_1_bias + bias_spec if has_bias else None, # d_wo_bias ) - def _bwd_body(ctx_local, dy_local): - return _body_bwd(ctx_local, dy_local, ep_active=True, **body_kwargs) + def _bwd_body(*args): + ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = _ffn_bwd_per_shard( + *args, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + has_bias=has_bias, + ) + # Weight grads accumulate per-DP-shard inside the body; psum across + # DP axes so each replica sees the full sum (matches out_specs + # P(ep_axis, ...) which is DP-replicated). + if data_parallelism_axes: + dp = tuple(data_parallelism_axes) + d_wi_0 = jax.lax.psum(d_wi_0, axis_name=dp) + d_wi_1 = jax.lax.psum(d_wi_1, axis_name=dp) + d_wo = jax.lax.psum(d_wo, axis_name=dp) + if has_bias: + d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=dp) + d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=dp) + d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=dp) + return ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) - grads = shard_map( + ( + d_sorted_x, + d_recv_w_from_intermediate, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = shard_map( _bwd_body, mesh=mesh, - in_specs=(ctx_spec, dy_specs), - out_specs=grads_spec, + in_specs=bwd_in_specs, + out_specs=bwd_out_specs, check_rep=False, - )(ctx, dy_pair) + )( + *bwd_in_args + ) - grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( - grads["gate_kernel"], gate_kernel_axes + d_recv_w_total = d_recv_w_from_combine + d_recv_w_from_intermediate + + # ---------------- Dispatch bwd (global view) ---------------- + d_sorted_x = jax.lax.with_sharding_constraint(d_sorted_x, NamedSharding(mesh, ep3_spec)) + d_recv_w_total = jax.lax.with_sharding_constraint(d_recv_w_total, NamedSharding(mesh, ep2_spec)) + d_x_from_dispatch, d_topk_w = tex.ep_dispatch_bwd( + ctx.cfg, + ctx.handle_mem, + d_sorted_x, + d_recv_w_total, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, ) - grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) - grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) - grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) - grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) - return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + # ---------------- Routing bwd (global view) ---------------- + # The cotangent on routing_weights is a sparse scatter into sparse_probs + # at the selected_experts indices. + selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] + d_topk_w_flat = d_topk_w.reshape(-1, K) + d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_topk_w_flat.dtype) + d_sparse_probs = d_sparse_probs.at[ + jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts + ].set(d_topk_w_flat) + + d_logits_2d = tex.fused_topk_with_score_function_bwd( + ctx.routing_map, + ctx.saved_scores, + d_sparse_probs.astype(ctx.saved_scores.dtype), + topk=K, + use_pre_softmax=use_pre_softmax, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=False, + ) + + # ---------------- Aux loss bwd (global view, replicated) ---------------- + # Reverse the fwd's all-gather/aux pipeline: aux_loss_bwd produces + # d_aux_probs, then topk_bwd(compute_aux_scores=True) produces the + # extra d_logits contribution. The replicated tensor adds into the + # T-sharded routing-side d_logits via JAX's normal broadcast. + if aux_loss_coeff > 0.0: + T_global = ctx.logits_2d.shape[0] + d_aux_loss_scalar = d_aux_loss.reshape(()).astype(jnp.float32) + d_aux_probs = tex.fused_moe_aux_loss_bwd( + ctx.aux_const_buf, + ctx.aux_tokens_per_expert.astype(jnp.int32), + d_aux_loss_scalar, + num_tokens=int(T_global), + ) + # routing_map is ignored by the kernel when compute_aux_scores=True, + # so pass a zero placeholder of the right shape/dtype. + zero_routing_map = jnp.zeros(ctx.aux_saved_scores.shape, dtype=ctx.routing_map.dtype) + d_logits_aux = tex.fused_topk_with_score_function_bwd( + zero_routing_map, + ctx.aux_saved_scores, + d_aux_probs.astype(ctx.aux_saved_scores.dtype), + topk=K, + use_pre_softmax=False, + scaling_factor=1.0, + score_function=score_function, + compute_aux_scores=True, + ) + d_logits_2d = d_logits_2d + d_logits_aux.astype(d_logits_2d.dtype) + + # ---------------- Gate bwd (global view) ---------------- + d_gate_logits = d_logits_2d.reshape(B, S, num_experts) + gate_kernel_cast = ctx.gate_kernel.astype(ctx.x.dtype) + d_x_from_gate = jnp.einsum("bse,he->bsh", d_gate_logits, gate_kernel_cast) + d_gate_kernel = jnp.einsum("bsh,bse->he", ctx.x, d_gate_logits).astype(ctx.gate_kernel.dtype) + d_x = d_x_from_gate + d_x_from_dispatch + + # Pin output grads to the declared logical axes so downstream + # optimizers see consistent shardings. + d_x = with_sharding_constraint_by_logical_axes(d_x, input_axes) + d_gate_kernel = with_sharding_constraint_by_logical_axes(d_gate_kernel, gate_kernel_axes) + d_wi_0 = with_sharding_constraint_by_logical_axes(d_wi_0, wi_kernel_axes) + d_wi_1 = with_sharding_constraint_by_logical_axes(d_wi_1, wi_kernel_axes) + d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_axes) + + # expert_bias has no learnable bwd path through fused_topk: the + # primitive's bwd returns None for the bias slot. Match that with a + # zero cotangent of the right shape so custom_vjp's arity check + # passes. + d_expert_bias = jnp.zeros_like(ctx.expert_bias) -def _grads_dict_to_tuple( - grads: dict, has_wi_bias: bool, has_wo_bias: bool, has_expert_bias: bool -) -> Tuple: - """Pack the body_bwd's grads dict into the positional tuple JAX expects.""" return ( - grads["inputs"], - grads["gate_kernel"], - grads["wi_0"], - grads["wi_1"], - grads["wo"], - grads.get("wi_0_bias") if has_wi_bias else None, - grads.get("wi_1_bias") if has_wi_bias else None, - grads.get("wo_bias") if has_wo_bias else None, - grads.get("expert_bias") if has_expert_bias else None, + d_x, + d_gate_kernel, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias if has_bias else None, + d_wi_1_bias if has_bias else None, + d_wo_bias if has_bias else None, + d_expert_bias, ) @@ -1995,7 +1064,7 @@ def _grads_dict_to_tuple( # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 29))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 26))) def _moe( x, gate_kernel, @@ -2015,23 +1084,16 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, ): - # Call in `_moe`'s own signature order to match what JAX will pass - # the fwd rule via ``_argnums_partial``. See the comment block at - # the top of ``_moe_fwd_rule`` for why this differs from - # ``_moe_bwd_rule``'s convention. - output_pair, _ = _moe_fwd_rule( + primal, _ = _moe_fwd_rule( x, gate_kernel, wi_0, @@ -2050,19 +1112,16 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, ) - return output_pair + return primal _moe.defvjp(_moe_fwd_rule, _moe_bwd_rule) @@ -2079,56 +1138,106 @@ def moe( wo_bias: Optional[jnp.ndarray] = None, expert_bias: Optional[jnp.ndarray] = None, *, - # Architecture num_experts: int, num_experts_per_tok: int, activation_type: str = "silu", - # Routing score_function: Union[str, ScoreFunction] = "softmax", use_pre_softmax: bool = False, num_groups: Optional[int] = None, group_topk: Optional[int] = None, scaling_factor: float = 1.0, aux_loss_coeff: float = 0.0, - # Permutation - permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX, - align_size: int = 0, - # Gate placement (Phuong: "perhaps as an option") - gate_inside_vjp: bool = True, - # Parallelism (resolved by caller from MeshResource) - ep_axis: Optional[str] = None, + apply_topk_weights_early: bool = False, + ep_axis: str, data_parallelism_axes: Tuple[str, ...] = (), - # Logical axes for sharding constraints input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed"), - # Quantization - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet] = ( - noop_quantizer_set, - noop_quantizer_set, - noop_quantizer_set, - ), dtype: jnp.dtype = jnp.float32, ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Run a full MoE block under a single fused custom_vjp. + """Run a full MoE block under a single fused custom_vjp on the TE EP path. + + Returns ``(output, aux_loss)``. ``aux_loss`` is ``None`` when + ``aux_loss_coeff == 0`` and a 0-d scalar otherwise. - Parameters and return are documented at the call site of - ``_MoEBlock.__call__``. See module docstring for design rationale. + Parameters + ---------- + expert_bias : Optional[jnp.ndarray] + ``[num_experts]`` learnable router bias added before the top-k + when ``score_function='sigmoid'``. Pass ``None`` to disable. + The bias has no gradient through the top-k primitive itself (it + only steers expert selection); a zero cotangent is returned for + it. + aux_loss_coeff : float + Per-step expert-load-balance loss coefficient. ``0.0`` (default) + disables the aux loss entirely. When non-zero, an extra + all-gather over the routing-side logits is inserted so the + ``fused_moe_aux_loss`` kernel sees a global ``[T_global, E]`` + view; this lives off the dispatch critical path. + + Note that the per-expert dispatch-slot alignment is fixed internally + at 128 tokens (``_ALIGN_SIZE``); see that constant's docstring for + rationale and how to extend if a future recipe needs >128. + + Axis-name parameters: + + * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh + axis names* -- they index ``jax.sharding.Mesh.shape`` directly + (to compute ``num_ep`` / ``dp_size`` and to construct + ``P((dp..., ep), None, None)`` for the per-shard + ``jax.lax.with_sharding_constraint`` calls that JAX requires + to refer to real mesh axes). + * ``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, + ``wo_kernel_axes`` are *logical axis names* (e.g. + ``"batch"``, ``"embed"``, ``"mlp"``, ``"exp"``) -- they get + resolved via the active Flax logical-axis rules and consumed + by ``with_sharding_constraint_by_logical_axes``. They are + ``Optional[str]`` tuples so a rule of ``None`` means + "replicated on this axis". + + Logical-axis support for ``ep_axis`` / ``data_parallelism_axes`` + is intentionally out of scope: the EP comm-group construction + (``dp_color = rank // ep_size``) and the bootstrap signature + check both require concrete integer sizes, so a logical name + would have to be resolved to a physical one anyway before any + EP primitive is called. If a downstream pipeline needs to plumb + logical names all the way to ``moe()``, do the rule lookup at + the call site. + + See module docstring for the rest of the parameter semantics and the + surrounding design rationale. """ - if not isinstance(permutation_backend, PermutationBackend): - raise TypeError( - f"permutation_backend must be a PermutationBackend, got {permutation_backend!r}" - ) - if permutation_backend is PermutationBackend.TRITON: - _require_triton() - # Normalize string score_function ("softmax" / "sigmoid") to the - # ScoreFunction enum once here. The underlying primitive - # ``tex.fused_topk_with_score_function_fwd`` expects an int-coercible - # value (the enum has integer .value), and the public router wrapper - # we bypass also normalizes here. score_function = _validate_score_function(score_function) + # Enforce ((outer_dp..., ep), None, None) on inbound activations. The + # EP comm groups consecutive global ranks (dp_color = rank // ep_size), + # so ep MUST be innermost in the partition spec. Soft re-pin: free if + # upstream already matches, single reshard otherwise. + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + expected_leading: Any = (*data_parallelism_axes, ep_axis) if data_parallelism_axes else ep_axis + expected_spec = P(expected_leading, None, None) + actual_spec = getattr(getattr(x, "sharding", None), "spec", None) + if actual_spec is not None and tuple(actual_spec) != tuple(expected_spec): + warnings.warn( + f"moe(...): inbound x sharding {actual_spec} does not match expected " + f"{expected_spec}; inserting a reshard. Apply " + "jax.lax.with_sharding_constraint upstream to avoid this overhead.", + UserWarning, + stacklevel=2, + ) + x = _with_sharding_constraint_cast_bwd(x, NamedSharding(mesh, expected_spec)) + + # custom_vjp can't trace through None args; lower expert_bias to an + # empty shape-(0,) tensor that fused_topk_with_score_function treats + # as "no bias". + if expert_bias is None: + expert_bias_arg = jnp.zeros((0,), dtype=jnp.float32) + else: + expert_bias_arg = expert_bias.astype(jnp.float32) + output, aux_loss = _moe( x, gate_kernel, @@ -2138,28 +1247,26 @@ def moe( wi_0_bias, wi_1_bias, wo_bias, - expert_bias, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - activation_type=activation_type, - score_function=score_function, - use_pre_softmax=use_pre_softmax, - num_groups=num_groups, - group_topk=group_topk, - scaling_factor=scaling_factor, - aux_loss_coeff=aux_loss_coeff, - permutation_backend=permutation_backend, - align_size=align_size, - gate_inside_vjp=gate_inside_vjp, - ep_axis=ep_axis, - data_parallelism_axes=data_parallelism_axes, - input_axes=input_axes, - gate_kernel_axes=gate_kernel_axes, - wi_kernel_axes=wi_kernel_axes, - wo_kernel_axes=wo_kernel_axes, - quantizer_sets=quantizer_sets, - dtype=dtype, + expert_bias_arg, + num_experts, + num_experts_per_tok, + activation_type, + score_function, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + float(aux_loss_coeff), + ep_axis, + data_parallelism_axes, + input_axes, + gate_kernel_axes, + wi_kernel_axes, + wo_kernel_axes, + dtype, + apply_topk_weights_early, ) if aux_loss_coeff <= 0.0: aux_loss = None + assert output.dtype == x.dtype, f"moe() output dtype {output.dtype} != input dtype {x.dtype}" return output, aux_loss From 6377ca161c0e9859083e369909ac37ad95bd94f4 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Fri, 10 Jul 2026 17:41:19 -0500 Subject: [PATCH 17/35] [PyTorch] Fix GIL/refcount abort in Comm+GEMM overlap and NCCL-EP init bindings (#3203) --- .../pytorch/csrc/extensions/ep.cpp | 11 +++++++--- .../pytorch/csrc/extensions/pybind.cpp | 21 ++++++++++++------- 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 8173df947e..118f14a01f 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -145,7 +145,13 @@ void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t nu .max_token_dtype = static_cast(GetTransformerEngineDType(torch_dtype)), .zero_copy = zero_copy ? 1 : 0, }; - nvte_ep_initialize(static_cast(ep_comm), &cfg); + // Release the GIL only around the native init. It must stay held while pybind11 casts + // the ``max_token_dtype`` object above and destroys the by-value ``pybind11::object`` + // parameter on return; releasing it across those trips pybind11's dec_ref GIL assertion. + { + pybind11::gil_scoped_release nogil; + nvte_ep_initialize(static_cast(ep_comm), &cfg); + } g_zero_copy_enabled.store(zero_copy, std::memory_order_relaxed); g_ep_initialized = true; g_ep_group_name = group_name; @@ -366,8 +372,7 @@ void register_ep_bindings(pybind11::module_& m) { "Initialize the EP backend; borrows torch's NCCL comm pointed to by ``comm_ptr``.", py::arg("comm_ptr"), py::arg("group_name"), py::arg("num_experts"), py::arg("max_tokens_per_rank"), py::arg("max_recv_tokens_per_rank"), py::arg("hidden_dim"), - py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false, - py::call_guard()); + py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false); m.def("ep_finalize", &ep_finalize, "Tear down the EP backend. Idempotent.", py::call_guard()); m.def("ep_get_zero_copy", &ep_get_zero_copy, "Return the current EP zero-copy toggle state."); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 9c9ec36138..7e9d114be8 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -716,6 +716,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { int num_max_streams, int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool atomic_gemm, bool rs_overlap_first_gemm) { + // Release the GIL only around the native construction (blocking collectives) to avoid + // tripping pybind11's inc_ref/dec_ref GIL assertions. + py::gil_scoped_release nogil; if (use_cublasmp) { return std::make_shared(helper, helper->mylocal, tp_size, comm_type, buffer_shape, buffer_dtype, num_comm_sm, @@ -726,8 +729,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm); }), - py::call_guard(), py::arg("buffer_shape"), - py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), + py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), py::arg("use_cublasmp") = false, py::arg("comm_type") = transformer_engine::CommOverlapType::RS, py::arg("num_splits") = 4, py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, @@ -751,6 +753,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool atomic_gemm, bool use_ce, bool aggregate, bool use_cublasmp) { + // Release the GIL only around the native construction (blocking collectives) to avoid + // tripping pybind11's inc_ref/dec_ref GIL assertions. + py::gil_scoped_release nogil; if (use_cublasmp) { return std::make_shared(helper, helper->mylocal, tp_size, comm_type, buffer_shape, buffer_dtype, num_comm_sm, @@ -761,12 +766,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, use_ce, aggregate); }), - py::call_guard(), py::arg("buffer_shape"), - py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), py::arg("comm_type"), - py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, py::arg("comm_cga_size") = 1, - py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 1, - py::arg("set_sm_margin") = false, py::arg("atomic_gemm") = false, - py::arg("use_ce") = true, py::arg("aggregate") = false, py::arg("use_cublasmp") = false) + py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), + py::arg("comm_type"), py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, + py::arg("comm_cga_size") = 1, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, + py::arg("num_comm_sm") = 1, py::arg("set_sm_margin") = false, + py::arg("atomic_gemm") = false, py::arg("use_ce") = true, py::arg("aggregate") = false, + py::arg("use_cublasmp") = false) .def("copy_into_buffer", static_cast( &CommOverlapP2P::copy_into_buffer), From aef96db0c0ee959f8197007e4ffad13bd4074003 Mon Sep 17 00:00:00 2001 From: Carlos Gomes Date: Tue, 14 Jul 2026 23:53:51 +0200 Subject: [PATCH 18/35] Migrate norms and softmax kernels to NVRTC (#3156) * [Common] NVRTC for fused softmax and normalization (Phase 0) Move the fused-softmax and LayerNorm/RMSNorm kernels from build-time template instantiation to runtime NVRTC compilation, with full coverage of the existing kernel set so the NVRTC path is the default. Fused softmax: - RTC compile/launch path for scaled / scaled-masked / scaled-upper-triangular / scaled-aligned-causal softmax, keyed by dtype, shape and mask/causal mode. - NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX (default OFF) restores the static template dispatch. Normalization (LayerNorm + RMSNorm, forward + backward): - Replace the static REGISTER_NORM_LAUNCHER template fanout with an NVRTC registry that compiles the selected (norm type, direction, dtypes, hidden size, CTA config) kernel on first use and caches it. - NVTE_BUILD_LEGACY_STATIC_NORM (default OFF) restores the static launchers. - NVRTC-safe kernel sources: kernel sources/headers avoid common.h under __CUDACC_RTC__; add the dtype aliases and a minimal std::is_same/conditional_t in the RTC build, and replace a zero-length padding array (a GNU extension nvcc accepts but NVRTC rejects) with a no-padding union specialization. KernelManager (util/rtc.{h,cpp}) gains occupancy / function-attribute / cooperative-launch helpers needed by the norm launchers. Validated on sm_89 (RTX 6000 Ada): full normalization operator suite 192/192, softmax + NVRTC unit tests pass; libtransformer_engine.so shrinks ~72 MB -> ~65 MB. On sm_100a the NVRTC norm forward kernel builds where the static instantiation crashed the compiler. Signed-off-by: CarlosGomes98 * Add fully qualified name to softmax kernels Signed-off-by: CarlosGomes98 * Add static fallback option, fix softmax acc_t dtype Signed-off-by: CarlosGomes98 * add missing license Signed-off-by: CarlosGomes98 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * greptile changes Signed-off-by: CarlosGomes98 * fix formatting, .clang-format Signed-off-by: CarlosGomes98 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * import cleanup Signed-off-by: CarlosGomes98 * Test more columns for softmax, mr changes Signed-off-by: CarlosGomes98 * Fix tests Signed-off-by: Carlos Gomes --------- Signed-off-by: CarlosGomes98 Signed-off-by: Carlos Gomes Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- .clang-format | 3 +- docs/envvars.rst | 2 +- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_softmax.cu | 243 ++++++ tests/cpp/util/CMakeLists.txt | 10 +- tests/cpp/util/test_nvrtc.cpp | 155 +++- transformer_engine/common/CMakeLists.txt | 47 ++ .../scaled_aligned_causal_masked_softmax.cu | 103 ++- .../fused_softmax/scaled_masked_softmax.cu | 575 ++++++++------ .../scaled_upper_triang_masked_softmax.cu | 425 ++++++---- .../customized_pipeline.cuh | 2 +- .../common/normalization/common.h | 102 +-- .../common/normalization/kernel_params.h | 88 +++ .../common/normalization/kernel_traits.h | 21 +- .../layernorm/ln_bwd_kernels.cuh | 5 + .../layernorm/ln_bwd_semi_cuda_kernel.cu | 37 +- .../layernorm/ln_fwd_cuda_kernel.cu | 38 +- .../layernorm/ln_fwd_kernels.cuh | 9 +- .../layernorm/rtc/ln_bwd_kernel.cu | 8 + .../layernorm/rtc/ln_fwd_kernel.cu | 12 + .../rmsnorm/rmsnorm_bwd_kernels.cuh | 30 +- .../rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu | 65 +- .../rmsnorm/rmsnorm_fwd_cuda_kernel.cu | 36 +- .../rmsnorm/rmsnorm_fwd_kernels.cuh | 9 +- .../rmsnorm/rtc/rmsnorm_bwd_kernel.cu | 8 + .../rmsnorm/rtc/rmsnorm_fwd_kernel.cu | 8 + .../common/normalization/rtc_dispatch.cpp | 742 ++++++++++++++++++ .../common/normalization/rtc_dispatch.h | 65 ++ transformer_engine/common/util/rtc.cpp | 73 +- transformer_engine/common/util/rtc.h | 62 +- transformer_engine/common/utils.cuh | 39 +- 31 files changed, 2468 insertions(+), 555 deletions(-) create mode 100644 tests/cpp/operator/test_softmax.cu create mode 100644 transformer_engine/common/normalization/kernel_params.h create mode 100644 transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu create mode 100644 transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu create mode 100644 transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu create mode 100644 transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu create mode 100644 transformer_engine/common/normalization/rtc_dispatch.cpp create mode 100644 transformer_engine/common/normalization/rtc_dispatch.h diff --git a/.clang-format b/.clang-format index aec13e3762..7860641ede 100644 --- a/.clang-format +++ b/.clang-format @@ -261,7 +261,7 @@ SpacesInParensOptions: InEmptyParentheses: false Other: false SpacesInSquareBrackets: false -Standard: Auto +Standard: c++17 StatementAttributeLikeMacros: - Q_EMIT StatementMacros: @@ -277,4 +277,3 @@ WhitespaceSensitiveMacros: - PP_STRINGIZE - STRINGIZE ... - diff --git a/docs/envvars.rst b/docs/envvars.rst index aa55621477..e9c3091c18 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -274,7 +274,7 @@ Kernel Configuration :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. This can be useful in environments where NVRTC is not available or not desired. + :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. Fused softmax and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX`` or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise no static fallback is available. .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 832177c637..2d5953c513 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -34,6 +34,7 @@ add_executable(test_operator test_multi_cast_transpose.cu test_multi_padding.cu test_multi_unpadding.cu + test_softmax.cu test_causal_softmax.cu test_swizzle.cu test_multi_swizzle.cu diff --git a/tests/cpp/operator/test_softmax.cu b/tests/cpp/operator/test_softmax.cu new file mode 100644 index 0000000000..b535c5c6cd --- /dev/null +++ b/tests/cpp/operator/test_softmax.cu @@ -0,0 +1,243 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include "../test_common.h" + +using namespace transformer_engine; + +namespace { + +template +void ref_softmax_row(Type *out, const Type *in, const uint8_t *mask, int cols, float scale) { + float max_value = -10000.0f; + bool has_unmasked = false; + for (int j = 0; j < cols; ++j) { + if (mask != nullptr && mask[j] == 1) continue; + max_value = std::max(max_value, static_cast(in[j]) * scale); + has_unmasked = true; + } + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + if (mask != nullptr && mask[j] == 1) { + out[j] = static_cast(0.0f); + continue; + } + const float val = has_unmasked ? std::exp(static_cast(in[j]) * scale - max_value) : 0.0f; + sum += val; + out[j] = static_cast(val); + } + for (int j = 0; j < cols; ++j) { + out[j] = static_cast(static_cast(out[j]) / sum); + } +} + +template +void ref_softmax_bwd(Type *grad_in, const Type *grad, const Type *softmax, int cols, float scale) { + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + sum += static_cast(grad[j]) * static_cast(softmax[j]); + } + for (int j = 0; j < cols; ++j) { + grad_in[j] = + static_cast(scale * (static_cast(grad[j]) - sum) * + static_cast(softmax[j])); + } +} + +template +void ref_upper_row(Type *out, const Type *in, int row, int cols, float scale) { + float max_value = -10000.0f; + for (int j = 0; j <= row; ++j) { + max_value = std::max(max_value, static_cast(in[j]) * scale); + } + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + if (j <= row) { + const float val = std::exp(static_cast(in[j]) * scale - max_value); + sum += val; + out[j] = static_cast(val); + } else { + out[j] = static_cast(0.0f); + } + } + for (int j = 0; j <= row; ++j) { + out[j] = static_cast(static_cast(out[j]) / sum); + } +} + +template +void test_scaled_softmax(DType dtype, size_t cols) { + using namespace test; + constexpr size_t batches = 2; + constexpr size_t heads = 2; + constexpr size_t rows = 8; + constexpr float scale = 0.7f; + const size_t elements_total = batches * heads * rows * cols; + Tensor input("input", std::vector{batches, heads, rows, cols}, dtype); + Tensor softmax("softmax", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad("grad", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad_out("grad_out", std::vector{batches, heads, rows, cols}, dtype); + fillUniform(&input); + fillUniform(&grad); + nvte_scaled_softmax_forward(input.data(), softmax.data(), scale, 0); + nvte_scaled_softmax_backward(grad.data(), softmax.data(), grad_out.data(), scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (size_t row = 0; row < elements_total / cols; ++row) { + ref_softmax_row(ref.get() + row * cols, input_cpu + row * cols, nullptr, cols, scale); + ref_softmax_bwd(ref_grad.get() + row * cols, grad_cpu + row * cols, ref.get() + row * cols, + cols, scale); + } + auto [atol, rtol] = getTolerances(dtype); + // Fused fp16/bf16 softmax backward differs from the fp32 reference by a few ULP. + atol = (dtype == DType::kBFloat16) ? 2e-3 : 1e-4; + compareResults("scaled_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("scaled_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +template +void test_masked_softmax(DType dtype, size_t cols) { + using namespace test; + constexpr size_t batches = 2; + constexpr size_t heads = 2; + constexpr size_t rows = 8; + constexpr float scale = -0.3f; + const size_t elements_total = batches * heads * rows * cols; + Tensor input("input", std::vector{batches, heads, rows, cols}, dtype); + Tensor mask("mask", std::vector{1, 1, rows, cols}, DType::kByte); + Tensor softmax("softmax", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad("grad", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad_out("grad_out", std::vector{batches, heads, rows, cols}, dtype); + fillUniform(&input); + fillUniform(&grad); + uint8_t *mask_cpu = mask.rowwise_cpu_dptr(); + for (size_t i = 0; i < rows * cols; ++i) { + mask_cpu[i] = (i % 7 == 0) ? 1 : 0; + } + mask.from_cpu(); + nvte_scaled_masked_softmax_forward(input.data(), mask.data(), softmax.data(), scale, 0); + nvte_scaled_masked_softmax_backward(grad.data(), softmax.data(), grad_out.data(), scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (int row = 0; row < batches * heads * rows; ++row) { + const int mask_row = row % rows; + ref_softmax_row(ref.get() + row * cols, input_cpu + row * cols, mask_cpu + mask_row * cols, + cols, scale); + ref_softmax_bwd(ref_grad.get() + row * cols, grad_cpu + row * cols, ref.get() + row * cols, + cols, scale); + } + auto [atol, rtol] = getTolerances(dtype); + // Fused fp16/bf16 softmax backward differs from the fp32 reference by a few ULP. + atol = (dtype == DType::kBFloat16) ? 2e-3 : 1e-4; + compareResults("masked_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("masked_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +template +void test_upper_softmax(DType dtype, size_t seq) { + using namespace test; + // attn_batches must be a multiple of the kernel's batches_per_block (up to 8). + constexpr size_t attn_batches = 8; + constexpr float scale = 1.2f; + const size_t elements_total = attn_batches * seq * seq; + Tensor input("input", std::vector{attn_batches, seq, seq}, dtype); + Tensor softmax("softmax", std::vector{attn_batches, seq, seq}, dtype); + Tensor grad("grad", std::vector{attn_batches, seq, seq}, dtype); + Tensor grad_out("grad_out", std::vector{attn_batches, seq, seq}, dtype); + fillUniform(&input); + fillUniform(&grad); + nvte_scaled_upper_triang_masked_softmax_forward(input.data(), softmax.data(), scale, 0); + nvte_scaled_upper_triang_masked_softmax_backward(grad.data(), softmax.data(), grad_out.data(), + scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (int batch = 0; batch < attn_batches; ++batch) { + for (int row = 0; row < seq; ++row) { + const size_t offset = (batch * seq + row) * seq; + ref_upper_row(ref.get() + offset, input_cpu + offset, row, seq, scale); + ref_softmax_bwd(ref_grad.get() + offset, grad_cpu + offset, ref.get() + offset, seq, scale); + for (int col = row + 1; col < seq; ++col) { + ref_grad[offset + col] = static_cast(0.0f); + } + } + } + auto [atol, rtol] = getTolerances(dtype); + // Upper-triangular backward diverges more near the causal diagonal; looser atol. + atol = (dtype == DType::kBFloat16) ? 1e-2 : 2e-3; + compareResults("upper_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("upper_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +} // namespace + +// Dispatch a 16-bit float dtype to a templated test body. Mirrors +// TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT but uses the test harness's own fp16/bf16 +// aliases so we don't have to include common.h here -- doing so would make the +// test's Tensor type ambiguous with transformer_engine::Tensor. +#define SOFTMAX_TEST_DISPATCH_16BIT(dtype, fn, cols) \ + switch (dtype) { \ + case DType::kFloat16: \ + fn(dtype, cols); \ + break; \ + case DType::kBFloat16: \ + fn(dtype, cols); \ + break; \ + default: \ + GTEST_FAIL() << "Unsupported 16-bit dtype for test"; \ + } + +using SoftmaxTestParams = std::tuple; + +class SoftmaxApiTestSuite : public ::testing::TestWithParam {}; + +TEST_P(SoftmaxApiTestSuite, ScaledSoftmax) { + const auto [dtype, cols] = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_scaled_softmax, cols); +} + +TEST_P(SoftmaxApiTestSuite, MaskedSoftmax) { + const auto [dtype, cols] = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_masked_softmax, cols); +} + +TEST_P(SoftmaxApiTestSuite, UpperTriangularSoftmax) { + const auto [dtype, cols] = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_upper_softmax, cols); +} + +INSTANTIATE_TEST_SUITE_P(OperatorTest, SoftmaxApiTestSuite, + ::testing::Combine( + ::testing::Values(DType::kFloat16, DType::kBFloat16), + ::testing::Values(32, 112, 1024)), + [](const testing::TestParamInfo &info) { + const DType dtype = std::get<0>(info.param); + const int cols = std::get<1>(info.param); + return test::typeName(dtype) + "_Cols" + std::to_string(cols); + }); diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 6d70b7b84f..1dfd2fed4e 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -9,7 +9,15 @@ add_executable(test_util find_package(OpenMP REQUIRED) -target_link_libraries(test_util PUBLIC CUDA::cudart GTest::gtest_main ${TE_LIB} CUDA::nvrtc CUDNN::cudnn OpenMP::OpenMP_CXX) +find_package(Threads REQUIRED) +target_link_libraries(test_util PUBLIC + CUDA::cudart + GTest::gtest_main + ${TE_LIB} + CUDA::nvrtc + CUDNN::cudnn + OpenMP::OpenMP_CXX + Threads::Threads) target_compile_options(test_util PRIVATE -O2 -fopenmp) include(GoogleTest) diff --git a/tests/cpp/util/test_nvrtc.cpp b/tests/cpp/util/test_nvrtc.cpp index d41084449e..5228c60972 100644 --- a/tests/cpp/util/test_nvrtc.cpp +++ b/tests/cpp/util/test_nvrtc.cpp @@ -4,11 +4,14 @@ * See LICENSE for license information. ************************************************************************/ +#include + +#include +#include #include +#include #include -#include - #include "util/rtc.h" using namespace transformer_engine; @@ -19,10 +22,10 @@ TEST(UtilTest, NVRTC) { } // GPU data buffer - int *device_buffer; + int* device_buffer; std::vector host_buffer(2); - cudaMalloc((void**)&device_buffer, 2*sizeof(int)); // NOLINT(*) - cudaMemset(device_buffer, 0, 2*sizeof(int)); + cudaMalloc((void**)&device_buffer, 2 * sizeof(int)); // NOLINT(*) + cudaMemset(device_buffer, 0, 2 * sizeof(int)); // CUDA kernel implementations const char code1[] = R"code( @@ -38,45 +41,151 @@ __global__ void my_kernel(uint32_t *data) { data[0] = 789; data[1] = 12; } +)code"; + const char code3[] = R"code( +#ifndef NVTE_GTEST_RTC_VALUE +#error "NVTE_GTEST_RTC_VALUE must be provided" +#endif +__global__ void my_kernel(int *data) { + data[0] = NVTE_GTEST_RTC_VALUE; + data[1] = 34; +} +)code"; + const char header4[] = R"code( +#define NVTE_GTEST_RTC_HEADER_VALUE 78 +)code"; + const char code4[] = R"code( +#include "test_nvrtc_header.h" +__global__ void my_kernel(int *data) { + data[0] = NVTE_GTEST_RTC_HEADER_VALUE; + data[1] = 90; +} )code"; // Make sure kernels are not available auto& nvrtc_manager = rtc::KernelManager::instance(); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel2")); - EXPECT_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, - device_buffer), + EXPECT_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, device_buffer), std::runtime_error); - EXPECT_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, - device_buffer), + EXPECT_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, device_buffer), std::runtime_error); // Compile and run first kernel - EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel1", - "my_kernel", - code1, - "test_nvrtc_kernel1.cu")); + EXPECT_NO_THROW( + nvrtc_manager.compile("my gtest kernel1", "my_kernel", code1, "test_nvrtc_kernel1.cu")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel2")); - EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, - device_buffer)); - EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2*sizeof(int), - cudaMemcpyDeviceToHost), + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), cudaSuccess); EXPECT_EQ(host_buffer[0], 123); EXPECT_EQ(host_buffer[1], -456); // Compile and run second kernel - EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel2", - "my_kernel", - code2, - "test_nvrtc_kernel2.cu")); + EXPECT_NO_THROW( + nvrtc_manager.compile("my gtest kernel2", "my_kernel", code2, "test_nvrtc_kernel2.cu")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel2")); EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, device_buffer)); - EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2*sizeof(int), - cudaMemcpyDeviceToHost), + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), cudaSuccess); EXPECT_EQ(host_buffer[0], 789); EXPECT_EQ(host_buffer[1], 12); + + // Compile and run kernel with extra compile options + EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel3", "my_kernel", code3, + "test_nvrtc_kernel3.cu", {"-DNVTE_GTEST_RTC_VALUE=56"})); + EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel3")); + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel3", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 56); + EXPECT_EQ(host_buffer[1], 34); + + // Compile and run kernel with an extra in-memory header + EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel4", "my_kernel", code4, + "test_nvrtc_kernel4.cu", {}, + {{header4, "test_nvrtc_header.h"}})); + EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel4")); + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel4", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 78); + EXPECT_EQ(host_buffer[1], 90); + + EXPECT_EQ(cudaFree(device_buffer), cudaSuccess); +} + +TEST(UtilTest, NVRTCConcurrentCompile) { + if (!rtc::is_enabled()) { + GTEST_SKIP() << "NVRTC not enabled, skipping tests"; + } + + constexpr int num_threads = 8; + constexpr char kernel_label[] = "my concurrent gtest kernel"; + const char code[] = R"code( +__global__ void my_kernel(int *data) { + data[0] = 314; + data[1] = 159; +} +)code"; + + int device_id = 0; + ASSERT_EQ(cudaGetDevice(&device_id), cudaSuccess); + + auto& nvrtc_manager = rtc::KernelManager::instance(); + ASSERT_FALSE(nvrtc_manager.is_compiled(kernel_label)); + + std::atomic ready{0}; + std::atomic start{false}; + std::vector errors(num_threads); + std::vector threads; + threads.reserve(num_threads); + for (int thread_id = 0; thread_id < num_threads; ++thread_id) { + threads.emplace_back([&, thread_id] { + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + try { + const cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(status)); + } + (void)nvrtc_manager.is_compiled(kernel_label); + nvrtc_manager.compile(kernel_label, "my_kernel", code, "test_nvrtc_concurrent_kernel.cu"); + } catch (...) { + errors[thread_id] = std::current_exception(); + } + }); + } + + while (ready.load(std::memory_order_acquire) != num_threads) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) { + thread.join(); + } + for (const auto& error : errors) { + if (error != nullptr) { + try { + std::rethrow_exception(error); + } catch (const std::exception& e) { + ADD_FAILURE() << e.what(); + } + } + } + + ASSERT_TRUE(nvrtc_manager.is_compiled(kernel_label)); + int* device_buffer = nullptr; + ASSERT_EQ(cudaMalloc(reinterpret_cast(&device_buffer), 2 * sizeof(int)), cudaSuccess); + ASSERT_NO_THROW(nvrtc_manager.launch(kernel_label, 1, 1, 0, 0, device_buffer)); + std::vector host_buffer(2); + ASSERT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 314); + EXPECT_EQ(host_buffer[1], 159); + EXPECT_EQ(cudaFree(device_buffer), cudaSuccess); } diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index be64fcb2be..d7f20137a2 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -183,6 +183,7 @@ list(APPEND transformer_engine_cpp_sources fused_attn/fused_attn.cpp gemm/config.cpp normalization/common.cpp + normalization/rtc_dispatch.cpp normalization/layernorm/ln_api.cpp normalization/rmsnorm/rmsnorm_api.cpp util/cuda_driver.cpp @@ -551,13 +552,59 @@ make_string_header_from_file(transpose/rtc/transpose.cu string_code_transpose_rtc_transpose_cu) make_string_header_from_file(transpose/rtc/swap_first_dims.cu string_code_transpose_rtc_swap_first_dims_cu) +make_string_header_from_file(fused_softmax/scaled_masked_softmax.cu + string_code_fused_softmax_scaled_masked_softmax_cu) +make_string_header_from_file(fused_softmax/scaled_upper_triang_masked_softmax.cu + string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu) +make_string_header_from_file(fused_softmax/scaled_aligned_causal_masked_softmax.cu + string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu) make_string_header_from_file(utils.cuh string_code_utils_cuh) make_string_header_from_file(util/math.h string_code_util_math_h) + +# Norm NVRTC bundled headers + RTC source files +make_string_header_from_file(normalization/kernel_params.h + string_code_normalization_kernel_params_h) +make_string_header_from_file(normalization/kernel_traits.h + string_code_normalization_kernel_traits_h) +make_string_header_from_file(normalization/layernorm/ln_fwd_kernels.cuh + string_code_normalization_layernorm_ln_fwd_kernels_cuh) +make_string_header_from_file(normalization/layernorm/ln_bwd_kernels.cuh + string_code_normalization_layernorm_ln_bwd_kernels_cuh) +make_string_header_from_file(normalization/rmsnorm/rmsnorm_fwd_kernels.cuh + string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh) +make_string_header_from_file(normalization/rmsnorm/rmsnorm_bwd_kernels.cuh + string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh) +make_string_header_from_file(normalization/layernorm/rtc/ln_fwd_kernel.cu + string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu) +make_string_header_from_file(normalization/layernorm/rtc/ln_bwd_kernel.cu + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu) +make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu + string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu) +make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu) target_include_directories(transformer_engine PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/string_headers") +option(NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + "Also compile legacy static fused softmax kernels for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=$) + +# Default OFF: LayerNorm/RMSNorm require NVRTC. Set ON to additionally compile +# the legacy static template instantiations, which are selected only when +# NVTE_DISABLE_NVRTC=1 at runtime. This compiles all 556 +# REGISTER_NORM_LAUNCHER variants up front. +option(NVTE_BUILD_LEGACY_STATIC_NORM + "Also compile static norm kernels for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_NORM=$) + # Compiler options set(nvte_sources_with_fast_math) list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu diff --git a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu index 6ea8017e07..1f8e748479 100644 --- a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -19,10 +21,23 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using fp16 = half; +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -92,7 +107,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -345,6 +360,42 @@ __global__ void scaled_aligned_causal_masked_softmax_warp_backward( } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=aligned_causal,direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *direction) { + NVTE_ERROR("Fused aligned causal softmax RTC path is disabled for ", direction, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void call_kernel_scaled_aligned_causal_masked_softmax_forward( dim3 grid_size, dim3 block_size, const int shmem_size, cudaStream_t stream, output_t *dst, @@ -454,6 +505,24 @@ void dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const dim3 block_size(warp_width, warps_per_block); dim3 grid_size(blocks); + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_aligned_causal_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, dst, src, rtc_scale, + microbatches, query_seq_len, key_seq_len); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX // create an array of pointers to functions using ForwardFuncType = typename FunctionWrapper::ForwardType; static std::array forwardFunctionArray; @@ -466,6 +535,9 @@ void dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const // Call the corresponding kernel forwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, dst, src, scale, microbatches, query_seq_len, key_seq_len); +#else + throw_nvrtc_required("forward"); +#endif } template @@ -499,6 +571,23 @@ void dispatch_scaled_aligned_causal_masked_softmax_backward( dim3 block_size(warp_width, warps_per_block); dim3 grid_size(blocks); + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_aligned_causal_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, grad_input, grad, output, + scale, microbatches, query_seq_len, key_seq_len); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX // create an array of pointers to functions using BackwardFuncType = typename FunctionWrapper::BackwardType; static std::array backwardFunctionArray; @@ -511,6 +600,9 @@ void dispatch_scaled_aligned_causal_masked_softmax_backward( // Call the corresponding kernel backwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, grad_input, grad, output, scale, microbatches, query_seq_len, key_seq_len); +#else + throw_nvrtc_required("backward"); +#endif } void scaled_aligned_causal_masked_softmax_forward(const Tensor &input, Tensor *softmax_results, @@ -546,8 +638,13 @@ void scaled_aligned_causal_masked_softmax_backward(Tensor output_grads, const Te reinterpret_cast(softmax_results.data.dptr), scale_factor, query_seq_len, key_seq_len, batches, attn_heads, stream);); } + +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_aligned_causal_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { @@ -568,3 +665,5 @@ void nvte_scaled_aligned_causal_masked_softmax_backward(const NVTETensor incomin *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 27f86673c5..6f295ad00a 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -17,10 +19,27 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + +template +__device__ __forceinline__ T neg_infinity() { + return -static_cast(__int_as_float(0x7f800000)); +} + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -67,7 +86,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -144,7 +163,7 @@ __global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, c } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -167,7 +186,7 @@ __global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, c for (int i = 0; i < WARP_BATCH; ++i) { #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -269,7 +288,7 @@ __global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -299,7 +318,7 @@ __global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t for (int i = 0; i < WARP_BATCH; ++i) { #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -420,6 +439,42 @@ __global__ void scaled_masked_softmax_warp_backward(output_t *gradInput, const i } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *variant, const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=", variant, ",direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *variant) { + NVTE_ERROR("Fused softmax RTC path is disabled for ", variant, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const input_t scale, int query_seq_len, int key_seq_len, int batches, @@ -448,70 +503,89 @@ void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const in NVTE_CHECK(query_seq_len % batches_per_block == 0, "Unsupported shape."); dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 1: // 2 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 2: // 4 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 3: // 8 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 4: // 16 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 5: // 32 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 6: // 64 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 7: // 128 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 8: // 256 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 9: // 512 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 10: // 1024 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 11: // 2048 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 12: // 4096 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 13: // 8192 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 14: // 16384 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("scaled", "forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile( + kernel_label, + make_softmax_rtc_kernel_name("scaled_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, rtc_scale, batch_count, + key_seq_len); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 1: // 2 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 2: // 4 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 3: // 8 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 4: // 16 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 5: // 32 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 6: // 64 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 7: // 128 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 8: // 256 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 9: // 512 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 10: // 1024 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 11: // 2048 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 12: // 4096 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 13: // 8192 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 14: // 16384 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled softmax forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -546,85 +620,105 @@ void dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, c NVTE_CHECK(query_seq_len % batches_per_block == 0, "Unsupported shape."); dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 1: // 2 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 2: // 4 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 3: // 8 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 4: // 16 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 5: // 32 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 6: // 64 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 7: // 128 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 8: // 256 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 9: // 512 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 10: // 1024 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 11: // 2048 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 12: // 4096 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 13: // 8192 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 14: // 16384 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("masked", "forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, mask, rtc_scale, + batch_count, key_seq_len, pad_batches); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 1: // 2 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 2: // 4 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 3: // 8 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 4: // 16 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 5: // 32 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 6: // 64 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 7: // 128 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 8: // 256 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 9: // 512 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 10: // 1024 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 11: // 2048 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 12: // 4096 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 13: // 8192 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 14: // 16384 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled masked softmax forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -658,85 +752,104 @@ void dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t int batches_per_block = warps_per_block * batches_per_warp; int blocks = batch_count / batches_per_block; dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 1: // 2 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 2: // 4 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 3: // 8 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 4: // 16 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 5: // 32 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 6: // 64 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 7: // 128 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 8: // 256 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 9: // 512 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 10: // 1024 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 11: // 2048 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 12: // 4096 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 13: // 8192 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 14: // 16384 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("masked", "backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, grad_input, grad, output, scale, + batch_count, key_seq_len); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 1: // 2 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 2: // 4 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 3: // 8 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 4: // 16 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 5: // 32 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 6: // 64 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 7: // 128 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 8: // 256 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 9: // 512 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 10: // 1024 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 11: // 2048 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 12: // 4096 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 13: // 8192 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 14: // 16384 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled softmax backward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -812,8 +925,12 @@ void scaled_masked_softmax_backward(Tensor output_grads, const Tensor incoming_g query_seq_len, key_seq_len, batches, attn_heads, stream);); } +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { NVTE_API_CALL(nvte_scaled_softmax_forward); @@ -850,3 +967,5 @@ void nvte_scaled_masked_softmax_backward(const NVTETensor incoming_grads, *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu index 431148cd1d..c18444f646 100644 --- a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -17,10 +19,28 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using fp16 = half; +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + +template +__device__ __forceinline__ T neg_infinity() { + return -static_cast(__int_as_float(0x7f800000)); +} + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -90,7 +110,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -166,13 +186,13 @@ __global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, c if ((element_index + element) < batch_element_count) { elements[i][it + element] = (acc_t)temp_data[element] * scale; } else { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -196,7 +216,7 @@ __global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, c #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { if (it < warp_iteration_limit) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -333,6 +353,42 @@ __global__ void scaled_upper_triang_masked_softmax_warp_backward(output_t *gradI } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=upper_triang,direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *direction) { + NVTE_ERROR("Fused upper-triangular softmax RTC path is disabled for ", direction, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const input_t *src, const input_t scale, int softmax_elements, @@ -365,85 +421,104 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const in int blocks_per_seq = attn_batches / batches_per_block; dim3 blocks(seq_len, blocks_per_seq, 1); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 1: // 2 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 2: // 4 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 3: // 8 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 4: // 16 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 5: // 32 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 6: // 64 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 7: // 128 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 8: // 256 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 9: // 512 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 10: // 1024 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 11: // 2048 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 12: // 4096 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 13: // 8192 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 14: // 16384 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_upper_triang_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, rtc_scale, batch_count, + softmax_elements_stride, softmax_elements); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 1: // 2 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 2: // 4 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 3: // 8 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 4: // 16 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 5: // 32 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 6: // 64 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 7: // 128 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 8: // 256 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 9: // 512 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 10: // 1024 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 11: // 2048 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 12: // 4096 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 13: // 8192 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 14: // 16384 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + default: + break; + } +#else + throw_nvrtc_required("forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -482,85 +557,103 @@ void dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, int blocks_per_seq = attn_batches / batches_per_block; dim3 blocks(seq_len, blocks_per_seq, 1); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 1: // 2 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 2: // 4 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 3: // 8 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 4: // 16 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 5: // 32 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 6: // 64 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 7: // 128 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 8: // 256 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 9: // 512 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 10: // 1024 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 11: // 2048 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 12: // 4096 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 13: // 8192 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 14: // 16384 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_upper_triang_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, grad_input, grad, output, scale, + batch_count, softmax_elements_stride, softmax_elements); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 1: // 2 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 2: // 4 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 3: // 8 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 4: // 16 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 5: // 32 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 6: // 64 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 7: // 128 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 8: // 256 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 9: // 512 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 10: // 1024 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 11: // 2048 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 12: // 4096 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 13: // 8192 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 14: // 16384 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + default: + break; + } +#else + throw_nvrtc_required("backward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -595,8 +688,12 @@ void scaled_upper_triang_masked_softmax_backward(Tensor output_grads, const Tens seq_len, attn_batches, stream);); } +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_upper_triang_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { @@ -615,3 +712,5 @@ void nvte_scaled_upper_triang_masked_softmax_backward(const NVTETensor incoming_ *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh index bc46341e88..11754fbb50 100644 --- a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh +++ b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh @@ -18,7 +18,7 @@ namespace detail { // by producer_commit. Use case, accumulator generation as // the result of MMA instructions. template , - class AtomThrShape_MNK_ = Shape<_1, _1, _1> > + class AtomThrShape_MNK_ = Shape<_1, _1, _1>> class CustomizedPipelineTmaUmmaAsync { public: static constexpr uint32_t Stages = Stages_; diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index f5dce64193..31a547f62c 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -24,6 +24,7 @@ #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" +#include "kernel_params.h" namespace transformer_engine { @@ -51,97 +52,6 @@ struct LaunchParams { } }; -struct KernelParamsBase { - KernelParamsBase() - : ctas_per_col(0), - rows(0), - cols(0), - x(nullptr), - mu(nullptr), - rs(nullptr), - gamma(nullptr), - workspace(nullptr), - barrier(nullptr), - zero_centered_gamma(false) {} - - // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x. - int ctas_per_col; - // Size of CTA group. - int ctas_per_row; - - // Input is interpreted as matrix. We normalize across columns. - int rows; - int cols; - - // Common data pointers. - void* x; - void* mu; - void* rs; - void* gamma; - - // Multi-CTA workspace in gmem. - void* workspace; - - // Multi-CTA sync barriers in gmem. - int* barrier; - - // Whether gamma is centered around 0 - bool zero_centered_gamma; -}; - -struct ForwardKernelParams : public KernelParamsBase { - ForwardKernelParams() - : KernelParamsBase(), z(nullptr), beta(nullptr), epsilon(0.f), fp8_out(false) {} - - // Output of LN FWD. - void* z; - void* beta; - float epsilon; - - // Scaling factor - void* scale; - int scale_byte_size; - - // Inverse of scaling factor - void* scale_inv; - - // AMax output - void* amax; - int amax_byte_size; - - // Whether to compute scale and amax - bool fp8_out; -}; - -struct BackwardKernelParams : public KernelParamsBase { - BackwardKernelParams() - : KernelParamsBase(), - dz(nullptr), - dbeta_part(nullptr), - dgamma_part(nullptr), - dx(nullptr), - dbeta(nullptr), - dgamma(nullptr) {} - - // Input: gradient wrt. LN FWD output. - void* dz; - - // Input: extra tensor to add for fused backward+add - void* add; - - // Workspace for Wgrad pre-reduction. - void* dbeta_part; - void* dgamma_part; - - // Output: Dgrad. - void* dx; - // Output: Wgrad. - void* dbeta; - void* dgamma; -}; - -using BackwardAddKernelParams = BackwardKernelParams; - enum class NVTE_Norm_Backend { Te, Cudnn }; enum class NVTE_Norm_Stage { Forward, Backward, BackwardAdd }; @@ -191,6 +101,16 @@ class TeNormalizationRegistry { return 0; } + // Overload for capturing-callable dispatchers (e.g. NVRTC closures). + static int registerFunction(TupleKeyType key, Function func) { + auto [general_key, batch_size, hidden_size, is_tuned] = key; + if (is_tuned) + getInstance().tuned_function_map.emplace(key, std::move(func)); + else + getInstance().general_function_map[general_key].emplace(hidden_size, std::move(func)); + return 0; + } + static Function getKernel(TupleKeyType key) { auto& instance = getInstance(); auto [general_key, batch_size, hidden_size, is_tuned] = key; diff --git a/transformer_engine/common/normalization/kernel_params.h b/transformer_engine/common/normalization/kernel_params.h new file mode 100644 index 0000000000..42dc423dd0 --- /dev/null +++ b/transformer_engine/common/normalization/kernel_params.h @@ -0,0 +1,88 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ +#define TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ + +// POD kernel-parameter structs shared between the host norm dispatchers and +// the NVRTC kernel sources. This header is intentionally free of host-only +// includes (no , no cudnn) so it can be compiled by +// NVRTC as well. + +namespace transformer_engine { +namespace normalization { + +struct KernelParamsBase { + // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x. + int ctas_per_col = 0; + // Size of CTA group. + int ctas_per_row = 0; + + // Input is interpreted as matrix. We normalize across columns. + int rows = 0; + int cols = 0; + + // Common data pointers. + void* x = nullptr; + void* mu = nullptr; + void* rs = nullptr; + void* gamma = nullptr; + + // Multi-CTA workspace in gmem. + void* workspace = nullptr; + + // Multi-CTA sync barriers in gmem. + int* barrier = nullptr; + + // Whether gamma is centered around 0 + bool zero_centered_gamma = false; +}; + +struct ForwardKernelParams : public KernelParamsBase { + // Output of LN FWD. + void* z = nullptr; + void* beta = nullptr; + float epsilon = 0.f; + + // Scaling factor + void* scale = nullptr; + int scale_byte_size = 0; + + // Inverse of scaling factor + void* scale_inv = nullptr; + + // AMax output + void* amax = nullptr; + int amax_byte_size = 0; + + // Whether to compute scale and amax + bool fp8_out = false; +}; + +struct BackwardKernelParams : public KernelParamsBase { + // Input: gradient wrt. LN FWD output. + void* dz = nullptr; + + // Input: extra tensor to add for fused backward+add + void* add = nullptr; + + // Workspace for Wgrad pre-reduction. + void* dbeta_part = nullptr; + void* dgamma_part = nullptr; + + // Output: Dgrad. + void* dx = nullptr; + // Output: Wgrad. + void* dbeta = nullptr; + void* dgamma = nullptr; +}; + +using BackwardAddKernelParams = BackwardKernelParams; + +} // namespace normalization +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ diff --git a/transformer_engine/common/normalization/kernel_traits.h b/transformer_engine/common/normalization/kernel_traits.h index 12fc095c38..41be82e742 100644 --- a/transformer_engine/common/normalization/kernel_traits.h +++ b/transformer_engine/common/normalization/kernel_traits.h @@ -7,12 +7,29 @@ #ifndef TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_TRAITS_H_ #define TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_TRAITS_H_ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include "../common.h" #include "../utils.cuh" +#endif namespace transformer_engine { namespace normalization { +#ifdef __CUDACC_RTC__ +// Under NVRTC the kernel sources do not include common.h (it pulls in cuDNN and +// other host-only headers), so the dtype aliases that the RTC name expressions +// reference are nototherwise visible (e.g. ::transformer_engine::normalization::fp16). +// Mirror the definitions from common.h here for the RTC. +// The underlying CUDA types come from utils.cuh (cuda_fp16/bf16/fp8 headers). +using fp32 = float; +using fp16 = half; +using bf16 = nv_bfloat16; +using fp8e4m3 = __nv_fp8_e4m3; +using fp8e5m2 = __nv_fp8_e5m2; +#endif // __CUDACC_RTC__ + template struct Kernel_traits_base { @@ -31,7 +48,7 @@ template > + compute_t_, index_t_, THREADS_PER_CTA_>> struct Kernel_traits_finalize : public Base { enum { ROWS_PER_CTA = Base::THREADS_PER_CTA / Base::THREADS_PER_WARP }; static_assert(static_cast(ROWS_PER_CTA) <= static_cast(Base::THREADS_PER_WARP)); @@ -69,7 +86,7 @@ template > + WARPS_M_ * WARPS_N_ * THREADS_PER_WARP>> struct Kernel_traits : public Base { using input_t = typename Base::input_t; using weight_t = typename Base::weight_t; diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh index c4b00b87c3..0480b588b9 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh @@ -7,8 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_BWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu index 68aa0942c1..2e3d5fd622 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu @@ -7,6 +7,7 @@ #include "../../common.h" #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_bwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -132,6 +133,30 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, NVTE_CHECK_CUDA(cudaGetLastError()); } +#define REGISTER_NORM_LAUNCHER_LN_BWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_bwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_LN_BWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_bwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -141,10 +166,16 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, launch_ln_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_LN_BWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ - } // namespace + } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_LN_BWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, ... diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu index 464df8d276..5be13357b1 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_fwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -101,6 +102,33 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, } } +// Register a single RTC-first dispatcher. When the static fallback is enabled, +// its function pointer is passed to the dispatcher and selected only when +// NVTE_DISABLE_NVRTC=1. +#define REGISTER_NORM_LAUNCHER_LN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_fwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, \ + BYTES_PER_LDG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_LN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_fwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -110,10 +138,16 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, launch_ln_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_LN_FWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_LN_FWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index 5a37cf46da..0e350146ea 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -7,11 +7,16 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include #include #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -147,7 +152,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_tuned_kernel( if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } @@ -323,7 +328,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_general_kerne if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } diff --git a/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu b/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu new file mode 100644 index 0000000000..f93f9e6edb --- /dev/null +++ b/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "ln_bwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu b/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu new file mode 100644 index 0000000000..5c34f05e8d --- /dev/null +++ b/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu @@ -0,0 +1,12 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC source file for LayerNorm forward kernels. The host bundles this as a +// string header; specific template instantiations are requested via +// nvrtcAddNameExpression at runtime. + +#include "kernel_traits.h" +#include "ln_fwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh index d620ee5260..492edd65f6 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh @@ -7,10 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ -#include - +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -18,19 +21,34 @@ namespace normalization { struct maybe_not_t {}; template -using maybe_t = std::conditional_t; - -template +using maybe_t = transformer_engine::detail::conditional_t; + +// dx and add share storage; `add` is positioned at the tail of the `dx` +// storage via leading padding. NeedsPadding is false when dx_t and add_t are +// the same size (or add_t is larger), in which case the padding array would be +// zero-length -- legal as a GNU extension under nvcc but rejected by NVRTC. The +// no-padding specialization below covers that case so both compilers are happy +// while keeping an identical layout. +template sizeof(maybe_t))> union dx_add_t { using add_t = maybe_t; using dx_t = Ivec; struct { - char _padding[sizeof(dx_t) > sizeof(add_t) ? sizeof(dx_t) - sizeof(add_t) : 0]; + char _padding[sizeof(dx_t) - sizeof(add_t)]; add_t add; }; dx_t dx; }; +template +union dx_add_t { + using add_t = maybe_t; + using dx_t = Ivec; + add_t add; + dx_t dx; +}; + template __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_bwd_tuned_kernel( BackwardKernelParams params) { diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu index 60238f256d..7e8a1cd624 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_bwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -133,6 +134,60 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para NVTE_CHECK_CUDA(cudaGetLastError()); } +// Two-level dispatch: stage (Backward / BackwardAdd) and launch type +// (tuned / general). The optional static fallback is selected only when +// NVTE_DISABLE_NVRTC=1. +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_tuned_Backward(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_bwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, false, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_tuned_BackwardAdd(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK) \ + static_assert(ADD_FLAG, "RMSNorm BackwardAdd registrations require ADD_FLAG=true"); \ + [[maybe_unused]] static const int \ + _rmsn_bwd_tuned_add_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_Backward( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_bwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + false, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_BackwardAdd(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, \ + CTYPE, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK) \ + static_assert(ADD_FLAG, "RMSNorm BackwardAdd registrations require ADD_FLAG=true"); \ + [[maybe_unused]] static const int \ + _rmsn_bwd_general_add_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + ADD_FLAG, STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -142,10 +197,16 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para launch_rmsnorm_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_RMSN_BWD_##LAUNCH_TYPE##_##NORM_STAGE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_RMSN_BWD_##LAUNCH_TYPE##_##NORM_STAGE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, \ + CTYPE, __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create rmsnorm bwd tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, ... diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu index 5522fd5c6b..c3308cbce0 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_fwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -102,6 +103,31 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param } } +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_fwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, \ + BYTES_PER_LDG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_fwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -111,10 +137,16 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param launch_rmsnorm_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_RMSN_FWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_RMSN_FWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create rmsnorm tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index 900fb58be2..bb284a9ab1 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -7,11 +7,16 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include #include #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -139,7 +144,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_tuned_ke if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } @@ -298,7 +303,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_general_ if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } diff --git a/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu new file mode 100644 index 0000000000..706e79870b --- /dev/null +++ b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "rmsnorm_bwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu new file mode 100644 index 0000000000..4d1bd10a7b --- /dev/null +++ b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "rmsnorm_fwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rtc_dispatch.cpp b/transformer_engine/common/normalization/rtc_dispatch.cpp new file mode 100644 index 0000000000..5ae81f4233 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.cpp @@ -0,0 +1,742 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC-backed registry registration for LayerNorm/RMSNorm forward + backward +// launchers. Each per-config REGISTER_NORM_LAUNCHER macro expands to a call +// into one of the register_*_tuned/general functions defined here. When +// NVTE_BUILD_LEGACY_STATIC_NORM=ON, that call also supplies a static fallback. +// The registered closure prefers NVRTC and selects the fallback only when +// NVTE_DISABLE_NVRTC=1. + +#include "rtc_dispatch.h" + +#include +#include +#include + +#include "../util/cuda_driver.h" +#include "../util/rtc.h" +#include "../util/string.h" +#include "common.h" + +// NVRTC source strings for the four kernel families. These are tiny stub +// files (each is a couple of #include lines that pull in kernel_traits + the +// matching kernel header); NVRTC then instantiates a specific +// (Kernel_traits) on demand via nvrtcAddNameExpression. +#include "string_code_normalization_kernel_params_h.h" +#include "string_code_normalization_kernel_traits_h.h" +#include "string_code_normalization_layernorm_ln_bwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_ln_fwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu.h" +#include "string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu.h" + +namespace transformer_engine { +namespace normalization { +namespace rtc_norm { + +namespace { + +const std::vector& norm_headers() { + static const std::vector headers = { + {string_code_normalization_kernel_params_h, "kernel_params.h"}, + {string_code_normalization_kernel_traits_h, "kernel_traits.h"}, + {string_code_normalization_layernorm_ln_fwd_kernels_cuh, "ln_fwd_kernels.cuh"}, + {string_code_normalization_layernorm_ln_bwd_kernels_cuh, "ln_bwd_kernels.cuh"}, + {string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh, "rmsnorm_fwd_kernels.cuh"}, + {string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh, "rmsnorm_bwd_kernels.cuh"}, + }; + return headers; +} + +void compile_norm_kernel(rtc::KernelManager& manager, const std::string& label, + const std::string& kernel_expr, const char* rtc_source, + const char* filename) { + manager.compile(label, kernel_expr, rtc_source, filename, {}, norm_headers()); +} + +// Map our DType enum onto the C++ type names used inside the norm RTC sources. +// Aliases come from normalization/common.h (`using bf16 = nv_bfloat16;` etc.). +const char* cpp_name_for(DType dt) { + switch (dt) { + case DType::kFloat32: + return "::transformer_engine::normalization::fp32"; + case DType::kFloat16: + return "::transformer_engine::normalization::fp16"; + case DType::kBFloat16: + return "::transformer_engine::normalization::bf16"; + case DType::kFloat8E4M3: + return "::transformer_engine::normalization::fp8e4m3"; + case DType::kFloat8E5M2: + return "::transformer_engine::normalization::fp8e5m2"; + default: + NVTE_ERROR("Unsupported DType for norm RTC dispatch"); + } +} + +int byte_size_of(DType dt) { + switch (dt) { + case DType::kFloat32: + return 4; + case DType::kFloat16: + case DType::kBFloat16: + return 2; + case DType::kFloat8E4M3: + case DType::kFloat8E5M2: + return 1; + default: + NVTE_ERROR("Unsupported DType for norm RTC dispatch"); + } +} + +// Build a Kernel_traits template-argument list as a string. The C++ argument +// list matches: +// Kernel_traits +std::string kernel_traits_expr(DType wt, DType it, DType ot, DType ct, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg) { + return concat_strings("::transformer_engine::normalization::Kernel_traits<", cpp_name_for(wt), + ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", cpp_name_for(ct), + ", uint32_t, ", hidden_size, ", ", ctas_per_row, ", ", warps_m, ", ", + warps_n, ", ", bytes_per_ldg, ">"); +} + +// Stats::SMEM_BYTES expressed in host code. +// stats_t = TypeToVec2::Type — for compute_t == fp32 that's float2 +// (8 bytes); for fp16 it's half2 (4 bytes); for bf16 it's nv_bfloat162 (4 +// bytes). Matches the formulas in utils.cuh: +// WARPS_N == 1 → 0 +// else → WARPS_M * WARPS_N * sizeof(stats_t) * 2 +int stats_smem_bytes(DType ctype, int warps_m, int warps_n) { + if (warps_n == 1) return 0; + int sizeof_stats_t; + switch (ctype) { + case DType::kFloat32: + sizeof_stats_t = 8; + break; + case DType::kFloat16: + case DType::kBFloat16: + sizeof_stats_t = 4; + break; + default: + NVTE_ERROR("Unsupported compute dtype for norm smem calc"); + } + return warps_m * warps_n * sizeof_stats_t * 2; +} + +// Reducer::SMEM_BYTES. +// reduce_t = TypeToVec2::Type — same sizes as stats_t. +int reducer_smem_bytes(DType ctype, int warps_m, int warps_n) { + return stats_smem_bytes(ctype, warps_m, warps_n); +} + +// Kernel_traits::SMEM_BYTES (used by backward launchers): +// SMEM_BYTES_DGRAD = Reducer::SMEM_BYTES +// SMEM_BYTES_WGRAD = (CTAS_PER_ROW > 1) ? 0 : WARPS_M * HIDDEN * sizeof(compute_t) +// SMEM_BYTES = DGRAD + WGRAD +int bwd_smem_bytes(DType ctype, int hidden_size, int ctas_per_row, int warps_m, int warps_n) { + const int dgrad = reducer_smem_bytes(ctype, warps_m, warps_n); + const int wgrad = (ctas_per_row > 1) ? 0 : warps_m * hidden_size * byte_size_of(ctype); + return dgrad + wgrad; +} + +template +bool try_static_fallback(StaticFallback static_fallback, + LaunchParams& launch_params, const bool configure_params) { + if (rtc::is_enabled()) { + return false; + } + NVTE_CHECK(static_fallback != nullptr, + "NVRTC is disabled and this normalization build has no static fallback. Rebuild with " + "NVTE_BUILD_LEGACY_STATIC_NORM=ON."); + static_fallback(launch_params, configure_params); + return true; +} + +// Common per-launch configure/launch helper. `kernel_expr` is the full +// templated kernel symbol (used as the nvrtcAddNameExpression argument); the +// closure captures everything needed. +template +void register_launcher(const std::string& label, const std::string& kernel_expr, + const char* rtc_source, const char* filename, TupleKeyType key, + int threads_per_cta, int dynamic_smem_bytes, int ctas_per_row, + bool needs_cooperative, int barrier_bytes_per_col, + int workspace_bytes_per_col, int dgamma_part_bytes_per_col, + StaticFallback static_fallback) { + auto closure = [label, kernel_expr, rtc_source, filename, threads_per_cta, dynamic_smem_bytes, + ctas_per_row, needs_cooperative, barrier_bytes_per_col, workspace_bytes_per_col, + dgamma_part_bytes_per_col, static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, rtc_source, filename); + } + + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, dynamic_smem_bytes); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = + launch_params.multiprocessorCount * ctas_per_sm / ctas_per_row; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = barrier_bytes_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_bytes_per_col * launch_params.params.ctas_per_col; + } + if (dgamma_part_bytes_per_col > 0) { + launch_params.dgamma_part_bytes = + dgamma_part_bytes_per_col * launch_params.params.ctas_per_col; + } + return; + } + + // Real launch. + if (dynamic_smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + dynamic_smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_col), dim3(threads_per_cta), dynamic_smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + dynamic_smem_bytes, stream, launch_params.params); + } + (void)needs_cooperative; + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +} // namespace + +// ============================================================================ +// LayerNorm Forward +// ============================================================================ + +void register_ln_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl, StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Forward, wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("ln_fwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::ln_fwd_tuned_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = stats_smem_bytes(ct, wm, wn); + // tuned path multi-CTA workspace formula: + // barrier_bytes = 2 * ctas_per_col * sizeof(index_t == uint32_t == 4 bytes) + // workspace_bytes = ctas_per_col * WARPS_M * CTAS_PER_ROW * sizeof(stats_t) * 2 + const int sizeof_stats_t = (ct == DType::kFloat32) ? 8 : 4; // float2 vs half2/bf162 + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_stats_t * 2; + register_launcher( + label, kernel_expr, string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, + "ln_fwd_kernel.cu", key, threads_per_cta, smem_bytes, cr, /*needs_cooperative=*/cr > 1, + barrier_per_col, workspace_per_col, /*dgamma_part_bytes_per_col=*/0, static_fallback); +} + +void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, int wn, + int bl, StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Forward, wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("ln_fwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl); + // "general" path always uses CTAS_PER_ROW=1 in the Kernel_traits. + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::ln_fwd_general_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const auto closure = [label, kernel_expr, threads_per_cta, hidden, wm, ct, static_fallback]( + LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, + string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, + "ln_fwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, /*smem=*/0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + // compute_t bytes + const int ctype_bytes = byte_size_of(ct); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * ctype_bytes * 2; + } + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, + stream, launch_params.params); + } + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// RMSNorm Forward (same shape as LayerNorm Forward) +// ============================================================================ + +void register_rmsnorm_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, NVTE_Norm_Stage::Forward, + wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("rmsnorm_fwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::rmsnorm_fwd_tuned_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = stats_smem_bytes(ct, wm, wn); + const int sizeof_stats_t = (ct == DType::kFloat32) ? 8 : 4; + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_stats_t * 2; + register_launcher( + label, kernel_expr, string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, + "rmsnorm_fwd_kernel.cu", key, threads_per_cta, smem_bytes, cr, /*needs_cooperative=*/cr > 1, + barrier_per_col, workspace_per_col, 0, static_fallback); +} + +void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, + int wn, int bl, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, NVTE_Norm_Stage::Forward, + wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("rmsnorm_fwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::rmsnorm_fwd_general_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const auto closure = [label, kernel_expr, threads_per_cta, hidden, wm, ct, static_fallback]( + LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, + string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, + "rmsnorm_fwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + const int ctype_bytes = byte_size_of(ct); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * ctype_bytes * 2; + } + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, + stream, launch_params.params); + } + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// LayerNorm Backward (main kernel + finalize kernel) +// ============================================================================ + +void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl_main, int bl_final, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Backward, wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("ln_bwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string main_traits = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl_main); + const std::string main_kexpr = concat_strings( + "&::transformer_engine::normalization::ln_bwd_tuned_kernel<", main_traits, ">"); + // Kernel_traits_finalize + const std::string finalize_traits = + concat_strings("::transformer_engine::normalization::Kernel_traits_finalize<", hidden, ", ", + cpp_name_for(wt), ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", + cpp_name_for(ct), ", uint32_t, 1024, ", bl_final, ">"); + const std::string finalize_kexpr = concat_strings( + "&::transformer_engine::normalization::ln_bwd_finalize_tuned_kernel<", finalize_traits, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = bwd_smem_bytes(ct, hidden, cr, wm, wn); + const int sizeof_reduce_t = (ct == DType::kFloat32) ? 8 : 4; + // tuned backward: workspace = ctas_per_col * WARPS_M * CTAS_PER_ROW * sizeof(reduce_t) * 2 + // dgamma_part_bytes = ctas_per_col * cols * sizeof(compute_t) + // barrier_bytes = 2 * ctas_per_col * sizeof(index_t) + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_reduce_t * 2; + const int dgamma_part_per_col = hidden * byte_size_of(ct); + + // Finalize kernel dims: Kernel_traits_finalize::THREADS_PER_CTA == 1024 + // Kernel_traits_finalize::CTAS = HIDDEN_SIZE / 32 (since COLS%32==0) + // COLS = HIDDEN_SIZE * sizeof(compute_t) / BYTES_PER_LDG_FINAL + // CTAS = COLS / 32 + const int colspass = hidden * byte_size_of(ct) / bl_final; + const int finalize_ctas = colspass / 32; + const int finalize_threads_per_cta = 1024; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, + smem_bytes, cr, barrier_per_col, workspace_per_col, dgamma_part_per_col, + finalize_ctas, finalize_threads_per_cta, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, smem_bytes); + launch_params.params.ctas_per_row = cr; + launch_params.params.ctas_per_col = launch_params.multiprocessorCount * ctas_per_sm / cr; + if (cr > 1) { + launch_params.barrier_bytes = barrier_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_per_col * launch_params.params.ctas_per_col; + } + launch_params.dgamma_part_bytes = dgamma_part_per_col * launch_params.params.ctas_per_col; + return; + } + if (smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (cr == 1) { + mgr.launch(main_label, dim3(ctas_per_col), dim3(threads_per_cta), smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(cr * ctas_per_col), dim3(threads_per_cta), smem_bytes, + stream, launch_params.params); + } + mgr.launch(finalize_label, dim3(finalize_ctas), dim3(finalize_threads_per_cta), 0, stream, + launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, int wn, + int bl_main, int bl_final, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Backward, wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("ln_bwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl_main); + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::ln_bwd_general_kernel<", traits, ">"); + // ln_bwd_finalize_general_kernel + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::ln_bwd_finalize_general_kernel<", + cpp_name_for(wt), ", ", cpp_name_for(ct), ", 4, 1, ", bl_final, ", 32>"); + const int threads_per_cta = wm * wn * 32; + // general bwd uses ctas_per_row = ceil_div(cols, HIDDEN_SIZE); smem=0 for main kernel call. + const int ctype_bytes = byte_size_of(ct); + const int finalize_threads_per_warp = 32; + const int finalize_warps_n = 1; + const int finalize_warps_m = 4; + const int finalize_elts_n_per_cta = + finalize_threads_per_warp * finalize_warps_n * bl_final / ctype_bytes; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, hidden, + wm, ctype_bytes, finalize_elts_n_per_cta, finalize_warps_n, finalize_warps_m, + finalize_threads_per_warp, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * (ctype_bytes * 2) * 2; + } + launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + 0, stream, launch_params.params); + } + const dim3 fin_block(finalize_threads_per_warp * finalize_warps_n, finalize_warps_m); + const dim3 fin_grid(ceil_div(cols, finalize_elts_n_per_cta), 1); + mgr.launch(finalize_label, fin_grid, fin_block, 0, stream, launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// RMSNorm Backward (main kernel + finalize kernel; same shape as LayerNorm bwd) +// ============================================================================ + +void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl_main, int bl_final, bool with_add, + StaticFallback static_fallback) { + const auto stage = with_add ? NVTE_Norm_Stage::BackwardAdd : NVTE_Norm_Stage::Backward; + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, stage, wt, it, ot, ct, 0, + hidden, false, true); + const std::string add_tag = with_add ? "_add" : ""; + const std::string label = concat_strings( + "rmsnorm_bwd_tuned", add_tag, ",w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, ",cr=", cr, + ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl_main); + const char* add_flag = with_add ? "true" : "false"; + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_tuned_kernel<", traits, + ", ", add_flag, ">"); + const std::string finalize_traits = + concat_strings("::transformer_engine::normalization::Kernel_traits_finalize<", hidden, ", ", + cpp_name_for(wt), ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", + cpp_name_for(ct), ", uint32_t, 1024, ", bl_final, ">"); + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_finalize_tuned_kernel<", + finalize_traits, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = bwd_smem_bytes(ct, hidden, cr, wm, wn); + const int sizeof_reduce_t = (ct == DType::kFloat32) ? 8 : 4; + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_reduce_t * 2; + const int dgamma_part_per_col = hidden * byte_size_of(ct); + const int colspass = hidden * byte_size_of(ct) / bl_final; + const int finalize_ctas = colspass / 32; + const int finalize_threads_per_cta = 1024; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, + smem_bytes, cr, barrier_per_col, workspace_per_col, dgamma_part_per_col, + finalize_ctas, finalize_threads_per_cta, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, smem_bytes); + launch_params.params.ctas_per_row = cr; + launch_params.params.ctas_per_col = launch_params.multiprocessorCount * ctas_per_sm / cr; + if (cr > 1) { + launch_params.barrier_bytes = barrier_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_per_col * launch_params.params.ctas_per_col; + } + launch_params.dgamma_part_bytes = dgamma_part_per_col * launch_params.params.ctas_per_col; + return; + } + if (smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (cr == 1) { + mgr.launch(main_label, dim3(ctas_per_col), dim3(threads_per_cta), smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(cr * ctas_per_col), dim3(threads_per_cta), smem_bytes, + stream, launch_params.params); + } + mgr.launch(finalize_label, dim3(finalize_ctas), dim3(finalize_threads_per_cta), 0, stream, + launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +void register_rmsnorm_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, + int wn, int bl_main, int bl_final, bool with_add, + StaticFallback static_fallback) { + const auto stage = with_add ? NVTE_Norm_Stage::BackwardAdd : NVTE_Norm_Stage::Backward; + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, stage, wt, it, ot, ct, 0, + hidden, false, false); + const std::string add_tag = with_add ? "_add" : ""; + const std::string label = concat_strings( + "rmsnorm_bwd_general", add_tag, ",w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, ",wm=", wm, + ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl_main); + const char* add_flag = with_add ? "true" : "false"; + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_general_kernel<", traits, + ", ", add_flag, ">"); + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_finalize_general_kernel<", + cpp_name_for(wt), ", ", cpp_name_for(ct), ", 4, 1, ", bl_final, ", 32>"); + const int threads_per_cta = wm * wn * 32; + const int ctype_bytes = byte_size_of(ct); + const int finalize_warps_m = 4; + const int finalize_warps_n = 1; + const int finalize_threads_per_warp = 32; + const int finalize_elts_n_per_cta = + finalize_threads_per_warp * finalize_warps_n * bl_final / ctype_bytes; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, hidden, + wm, ctype_bytes, finalize_elts_n_per_cta, finalize_warps_n, finalize_warps_m, + finalize_threads_per_warp, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * (ctype_bytes * 2) * 2; + } + launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + 0, stream, launch_params.params); + } + const dim3 fin_block(finalize_threads_per_warp * finalize_warps_n, finalize_warps_m); + const dim3 fin_grid(ceil_div(cols, finalize_elts_n_per_cta), 1); + mgr.launch(finalize_label, fin_grid, fin_block, 0, stream, launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +} // namespace rtc_norm +} // namespace normalization +} // namespace transformer_engine diff --git a/transformer_engine/common/normalization/rtc_dispatch.h b/transformer_engine/common/normalization/rtc_dispatch.h new file mode 100644 index 0000000000..314f150f70 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.h @@ -0,0 +1,65 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ +#define TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ + +#include + +#include "common.h" + +namespace transformer_engine { +namespace normalization { +namespace rtc_norm { + +template +using StaticFallback = void (*)(LaunchParams&, const bool); + +// Register an RTC-backed launcher for a single LayerNorm Forward "tuned" +// (multi-CTA-capable) config. Compiles and launches via NVRTC on first use. +void register_ln_fwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Forward "general" +// (no multi-CTA) config. +void register_ln_fwd_general(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Backward "tuned" config. +void register_ln_bwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Backward "general" config. +void register_ln_bwd_general(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, + StaticFallback static_fallback = nullptr); + +// Same set for RMSNorm. +void register_rmsnorm_fwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_fwd_general(DType wtype, DType itype, DType otype, DType ctype, + int hidden_size, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_bwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, bool with_add, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_bwd_general(DType wtype, DType itype, DType otype, DType ctype, + int hidden_size, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, bool with_add, + StaticFallback static_fallback = nullptr); + +} // namespace rtc_norm +} // namespace normalization +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 70024a202c..20616f8cb6 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -47,12 +47,7 @@ inline int max_supported_sm_arch() { } // namespace bool is_enabled() { - static bool is_enabled_ = false; - static bool need_to_check_env = true; - if (need_to_check_env) { - is_enabled_ = !getenv("NVTE_DISABLE_NVRTC"); - need_to_check_env = false; - } + static const bool is_enabled_ = !getenv("NVTE_DISABLE_NVRTC"); return is_enabled_; } @@ -132,6 +127,18 @@ void Kernel::set_function_cache_config(int device_id, CUfunc_cache cache_config) NVTE_CALL_CHECK_CUDA_DRIVER(cuFuncSetCacheConfig, get_function(device_id), cache_config); } +void Kernel::set_function_attribute(int device_id, CUfunction_attribute attr, int value) { + NVTE_CALL_CHECK_CUDA_DRIVER(cuFuncSetAttribute, get_function(device_id), attr, value); +} + +int Kernel::occupancy_max_active_blocks_per_sm(int device_id, int block_size, + std::size_t dynamic_smem_bytes) { + int num_blocks = 0; + NVTE_CALL_CHECK_CUDA_DRIVER(cuOccupancyMaxActiveBlocksPerMultiprocessor, &num_blocks, + get_function(device_id), block_size, dynamic_smem_bytes); + return num_blocks; +} + KernelManager& KernelManager::instance() { NVTE_CHECK(is_enabled(), "NVRTC support is not enabled"); static KernelManager instance_; @@ -139,11 +146,17 @@ KernelManager& KernelManager::instance() { } void KernelManager::compile(const std::string& kernel_label, const std::string& kernel_name, - const std::string& code, const std::string& filename) { - std::lock_guard lock_guard_(lock_); + const std::string& code, const std::string& filename, + const std::vector& extra_options, + const std::vector
& extra_headers) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::unique_lock lock_guard_(lock_); + if (kernel_cache_.count(key) > 0) { + return; + } // Choose whether to compile to PTX or cubin - const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); const int compile_sm_arch = std::min(sm_arch_, max_supported_sm_arch()); const bool compile_ptx = sm_arch_ != compile_sm_arch; @@ -160,6 +173,7 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& opts.push_back(concat_strings("--gpu-architecture=sm_", compile_sm_arch)); } opts.push_back(concat_strings("-I", cuda::include_directory(true))); + opts.insert(opts.end(), extra_options.begin(), extra_options.end()); std::vector opts_ptrs; for (const auto& opt : opts) { opts_ptrs.push_back(opt.c_str()); @@ -167,11 +181,19 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& // Compile source nvrtcProgram program; - constexpr int num_headers = 2; - constexpr const char* headers[num_headers] = {string_code_utils_cuh, string_code_util_math_h}; - constexpr const char* include_names[num_headers] = {"utils.cuh", "util/math.h"}; - NVTE_CHECK_NVRTC(nvrtcCreateProgram(&program, code.c_str(), filename.c_str(), num_headers, - headers, include_names)); + std::vector headers = {string_code_utils_cuh, string_code_util_math_h}; + std::vector include_names = {"utils.cuh", "util/math.h"}; + headers.reserve(headers.size() + extra_headers.size()); + include_names.reserve(include_names.size() + extra_headers.size()); + for (const auto& header : extra_headers) { + NVTE_CHECK(header.content != nullptr && header.include_name != nullptr, + "NVRTC header content and include name must not be null"); + headers.push_back(header.content); + include_names.push_back(header.include_name); + } + NVTE_CHECK_NVRTC(nvrtcCreateProgram(&program, code.c_str(), filename.c_str(), + static_cast(headers.size()), headers.data(), + include_names.data())); NVTE_CHECK_NVRTC(nvrtcAddNameExpression(program, kernel_name.c_str())); const nvrtcResult compile_result = nvrtcCompileProgram(program, opts_ptrs.size(), opts_ptrs.data()); @@ -239,7 +261,6 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& } // Cache compiled code - const auto key = get_kernel_cache_key(kernel_label, device_id); kernel_cache_.insert({key, Kernel(mangled_name, std::move(compiled_code))}); kernel_cache_.at(key).get_function(device_id); // Make sure kernel is available on device @@ -250,12 +271,34 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& void KernelManager::set_cache_config(const std::string& kernel_label, CUfunc_cache cache_config) { const int device_id = cuda::current_device(); const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to configure RTC kernel before compilation"); kernel_cache_.at(key).set_function_cache_config(device_id, cache_config); } +void KernelManager::set_function_attribute(const std::string& kernel_label, + CUfunction_attribute attr, int value) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to configure RTC kernel before compilation"); + kernel_cache_.at(key).set_function_attribute(device_id, attr, value); +} + +int KernelManager::occupancy_max_active_blocks_per_sm(const std::string& kernel_label, + int block_size, + std::size_t dynamic_smem_bytes) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to query occupancy before compilation"); + return kernel_cache_.at(key).occupancy_max_active_blocks_per_sm(device_id, block_size, + dynamic_smem_bytes); +} + bool KernelManager::is_compiled(const std::string& kernel_label, int device_id) const { const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); return kernel_cache_.count(key) > 0; } diff --git a/transformer_engine/common/util/rtc.h b/transformer_engine/common/util/rtc.h index 65faf7bcc2..745566ee60 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -33,6 +34,12 @@ namespace rtc { */ bool is_enabled(); +/*! \brief Header made available to an NVRTC program */ +struct Header { + const char *content; + const char *include_name; +}; + /*! \brief Wrapper class for a runtime-compiled CUDA kernel */ class Kernel { public: @@ -79,6 +86,29 @@ class Kernel { */ void set_function_cache_config(int device_id, CUfunc_cache cache_config); + /*! \brief Set a kernel function attribute (driver-API wrapper of + * cuFuncSetAttribute, e.g. for dynamic shared memory size). + */ + void set_function_attribute(int device_id, CUfunction_attribute attr, int value); + + /*! \brief Wrapper of cuOccupancyMaxActiveBlocksPerMultiprocessor for a + * runtime-compiled function. + */ + int occupancy_max_active_blocks_per_sm(int device_id, int block_size, + std::size_t dynamic_smem_bytes); + + /*! \brief Cooperative launch of an RTC kernel via cuLaunchCooperativeKernel. + */ + template + void launch_cooperative(int device_id, const dim3 grid_dim, const dim3 block_dim, + unsigned int shared_mem_bytes, cudaStream_t stream, ArgTs &&...args) { + cuda_driver::ensure_context_exists(); + void *arg_ptrs[] = {const_cast(static_cast(&args))...}; + NVTE_CALL_CHECK_CUDA_DRIVER(cuLaunchCooperativeKernel, get_function(device_id), grid_dim.x, + grid_dim.y, grid_dim.z, block_dim.x, block_dim.y, block_dim.z, + shared_mem_bytes, static_cast(stream), arg_ptrs); + } + private: /*! \brief Mangled function name */ std::string mangled_name_; @@ -113,9 +143,13 @@ class KernelManager { * \param[in] code Kernel source code * \param[in] filename Path to associate with source code, * primarily for debugging + * \param[in] extra_options Additional NVRTC compiler options + * \param[in] extra_headers Additional in-memory headers available to the program */ void compile(const std::string &kernel_label, const std::string &kernel_name, - const std::string &code, const std::string &filename); + const std::string &code, const std::string &filename, + const std::vector &extra_options = {}, + const std::vector
&extra_headers = {}); /*! \brief Whether CUDA kernel has been compiled for CUDA device * @@ -143,6 +177,7 @@ class KernelManager { unsigned int shared_mem_bytes, cudaStream_t stream, ArgTs &&...args) { const int device_id = cuda::current_device(); const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to launch RTC kernel before compilation"); kernel_cache_.at(key).launch(device_id, grid_dim, block_dim, shared_mem_bytes, stream, std::forward(args)...); @@ -157,11 +192,32 @@ class KernelManager { */ void set_cache_config(const std::string &kernel_label, CUfunc_cache cache_config); + /*! \brief Set a function attribute (e.g. cuFuncAttributeMaxDynamicSharedMemorySize). */ + void set_function_attribute(const std::string &kernel_label, CUfunction_attribute attr, + int value); + + /*! \brief Query cuOccupancyMaxActiveBlocksPerMultiprocessor for a compiled kernel. */ + int occupancy_max_active_blocks_per_sm(const std::string &kernel_label, int block_size, + std::size_t dynamic_smem_bytes); + + /*! \brief Cooperative launch wrapper (cuLaunchCooperativeKernel). */ + template + void launch_cooperative(const std::string &kernel_label, const dim3 grid_dim, + const dim3 block_dim, unsigned int shared_mem_bytes, cudaStream_t stream, + ArgTs &&...args) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to launch RTC kernel before compilation"); + kernel_cache_.at(key).launch_cooperative(device_id, grid_dim, block_dim, shared_mem_bytes, + stream, std::forward(args)...); + } + private: /*! \brief Compiled kernels */ std::unordered_map kernel_cache_; - /*! \brief Mutex for thread-safe compilation */ - std::mutex lock_; + /*! \brief Mutex for thread-safe cache access */ + mutable std::shared_mutex lock_; KernelManager() = default; ~KernelManager() = default; diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 635c7a36d2..0b75da622c 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -30,6 +30,37 @@ static_assert(sizeof(uint32_t) == 4); static_assert(sizeof(uint64_t) == 8); #endif +// Minimal subset of used by RTC kernel headers. Keep these in a +// project-owned namespace because adding primary templates to std is undefined. +namespace transformer_engine { +namespace detail { + +template +struct is_same { + static constexpr bool value = false; +}; +template +struct is_same { + static constexpr bool value = true; +}; + +template +inline constexpr bool is_same_v = is_same::value; + +template +struct conditional { + using type = T; +}; +template +struct conditional { + using type = F; +}; +template +using conditional_t = typename conditional::type; + +} // namespace detail +} // namespace transformer_engine + //////////////////////////////////////////////////////////////////////////////////////////////////// constexpr uint32_t THREADS_PER_WARP = 32; @@ -940,13 +971,13 @@ __device__ __forceinline__ float ordered_uint_to_float(unsigned int u) { template __device__ __forceinline__ T abs_val(T val) { - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { #if __CUDA_ARCH__ >= 800 return __habs(val); #else return static_cast<__nv_bfloat16>(fabsf(static_cast(val))); #endif - } else if constexpr (std::is_same_v) { + } else if constexpr (detail::is_same_v) { return __habs(val); } else { return fabsf(val); @@ -955,13 +986,13 @@ __device__ __forceinline__ T abs_val(T val) { template __device__ __forceinline__ T max_val(T a, T b) { - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { #if __CUDA_ARCH__ >= 800 return __hmax(a, b); #else return static_cast<__nv_bfloat16>(fmaxf(static_cast(a), static_cast(b))); #endif - } else if constexpr (std::is_same_v) { + } else if constexpr (detail::is_same_v) { return __hmax(a, b); } else { return fmaxf(a, b); From 215e15e0ecd8ac2c5b83c4ac7e3e47a8d2da9fd7 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:20:44 -0700 Subject: [PATCH 19/35] [Common] Support scaled & clamped swiglu, srelu for BF16 (#3132) * support scaled swiglu, scaled srelu and scaled clamp swiglu Signed-off-by: zhongboz * vectorized loading improvement Signed-off-by: Zhongbo Zhu * fix bug for backward kernel Signed-off-by: Zhongbo Zhu * optimize Signed-off-by: Zhongbo Zhu * fix unit test failure Signed-off-by: Zhongbo Zhu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update tests/cpp/operator/test_scaled_activation.cu Signed-off-by: vthumbe1503 * resolve comments Signed-off-by: Zhongbo Zhu * refactor, resolve comments Signed-off-by: Zhongbo Zhu * address review comment Signed-off-by: Varun Thumbe * adaptive cta to fix slow block reduce for scale grads Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor to have gated and unary activation in activation infra Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reuse scale grad kernel for non scale grad since it is faster anyway Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: zhongboz Signed-off-by: Zhongbo Zhu Signed-off-by: vthumbe1503 Signed-off-by: Varun Thumbe Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 --- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_scaled_activation.cu | 291 ++++++++ transformer_engine/common/CMakeLists.txt | 6 + .../common/activation/scaled_activation.cu | 642 ++++++++++++++++++ .../common/activation/scaled_activation.h | 90 +++ .../common/activation/scaled_srelu.cu | 29 + .../common/activation/scaled_swiglu.cu | 56 ++ .../include/transformer_engine/activation.h | 105 +++ 8 files changed, 1220 insertions(+) create mode 100644 tests/cpp/operator/test_scaled_activation.cu create mode 100644 transformer_engine/common/activation/scaled_activation.cu create mode 100644 transformer_engine/common/activation/scaled_activation.h create mode 100644 transformer_engine/common/activation/scaled_srelu.cu create mode 100644 transformer_engine/common/activation/scaled_swiglu.cu diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 2d5953c513..06ed56cc5e 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -27,6 +27,7 @@ add_executable(test_operator test_cast_transpose_dbias_dgelu.cu test_cast_transpose_dgeglu.cu test_act.cu + test_scaled_activation.cu test_normalization.cu test_normalization_mxfp8.cu test_memset.cu diff --git a/tests/cpp/operator/test_scaled_activation.cu b/tests/cpp/operator/test_scaled_activation.cu new file mode 100644 index 0000000000..cf12ca8633 --- /dev/null +++ b/tests/cpp/operator/test_scaled_activation.cu @@ -0,0 +1,291 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include "../test_common.h" + +using namespace transformer_engine; + +namespace { + +enum class ScaledActivationCase { + kSwiGLU, + kClampedSwiGLU, + kSReLU, +}; + +constexpr float kClampedLimit = 0.5f; +constexpr float kClampedAlpha = 1.702f; +constexpr float kClampedLinearOffset = 0.5f; + +const char *activation_name(ScaledActivationCase activation) { + switch (activation) { + case ScaledActivationCase::kSwiGLU: + return "scaled_swiglu"; + case ScaledActivationCase::kClampedSwiGLU: + return "scaled_clamped_swiglu"; + case ScaledActivationCase::kSReLU: + return "scaled_srelu"; + } + return "unknown"; +} + +inline void glu_indices(const size_t row, const size_t col, const size_t hidden, + const int64_t interleave, size_t *act_idx, size_t *linear_idx) { + if (interleave > 0) { + const size_t block = col / static_cast(interleave); + const size_t lane = col % static_cast(interleave); + const size_t base = row * hidden * 2 + block * static_cast(interleave) * 2 + lane; + *act_idx = base; + *linear_idx = base + static_cast(interleave); + } else { + const size_t base = row * hidden * 2; + *act_idx = base + col; + *linear_idx = base + hidden + col; + } +} + +inline void gated_grads(const ScaledActivationCase activation, const float act_in, + const float linear_in, float *dact, float *dlinear, float *unscaled) { + switch (activation) { + case ScaledActivationCase::kSwiGLU: { + const float act = test::silu(act_in); + *unscaled = act * linear_in; + *dact = test::dsilu(act_in) * linear_in; + *dlinear = act; + return; + } + case ScaledActivationCase::kClampedSwiGLU: { + const bool dlinear_mask = linear_in <= kClampedLimit && linear_in >= -kClampedLimit; + const float act = test::qgelu(fminf(kClampedLimit, act_in)); + const float dact_base = + act_in <= kClampedLimit ? test::dqgelu(fminf(kClampedLimit, act_in)) : 0.0f; + const float linear = + fminf(fmaxf(-kClampedLimit, linear_in), kClampedLimit) + kClampedLinearOffset; + *unscaled = act * linear; + *dact = dact_base * linear; + *dlinear = dlinear_mask ? act : 0.0f; + return; + } + case ScaledActivationCase::kSReLU: + *unscaled = test::srelu(act_in); + *dact = test::dsrelu(act_in); + *dlinear = 0.0f; + return; + } +} + +template +void compute_reference(ScaledActivationCase activation, const DataT *input, const ScaleT *scales, + const DataT *grad_output, DataT *output, DataT *grad_input, + DataT *grad_scales, const size_t rows, const size_t hidden, + const int64_t interleave, const bool compute_grad_scales) { + const bool is_gated = activation != ScaledActivationCase::kSReLU; + const size_t input_cols = is_gated ? hidden * 2 : hidden; + std::fill(grad_input, grad_input + rows * input_cols, static_cast(0.0f)); + + for (size_t row = 0; row < rows; ++row) { + const float scale = static_cast(scales[row]); + float scale_grad = 0.0f; + for (size_t col = 0; col < hidden; ++col) { + const size_t out_idx = row * hidden + col; + float unscaled = 0.0f; + float dact = 0.0f; + float dlinear = 0.0f; + if (is_gated) { + size_t act_idx = 0; + size_t linear_idx = 0; + glu_indices(row, col, hidden, interleave, &act_idx, &linear_idx); + const float act_in = static_cast(input[act_idx]); + const float linear_in = static_cast(input[linear_idx]); + gated_grads(activation, act_in, linear_in, &dact, &dlinear, &unscaled); + + const float scaled_grad = static_cast(grad_output[out_idx]) * scale; + grad_input[act_idx] = static_cast(scaled_grad * dact); + grad_input[linear_idx] = static_cast(scaled_grad * dlinear); + } else { + const float x = static_cast(input[out_idx]); + unscaled = test::srelu(x); + const float scaled_grad = static_cast(grad_output[out_idx]) * scale; + grad_input[out_idx] = static_cast(scaled_grad * test::dsrelu(x)); + } + + output[out_idx] = static_cast(unscaled * scale); + scale_grad += static_cast(grad_output[out_idx]) * unscaled; + } + if (compute_grad_scales) { + grad_scales[row] = static_cast(scale_grad); + } + } +} + +template +void run_scaled_activation_test(ScaledActivationCase activation, const size_t rows, + const size_t hidden, const int64_t interleave, + const bool compute_grad_scales) { + using namespace test; + const DType data_type = TypeInfo::dtype; + const DType scale_type = TypeInfo::dtype; + const bool is_gated = activation != ScaledActivationCase::kSReLU; + const size_t input_cols = is_gated ? hidden * 2 : hidden; + + Tensor input("input", std::vector{rows, input_cols}, data_type); + Tensor scales("act_scales", std::vector{rows}, scale_type); + Tensor output("output", std::vector{rows, hidden}, data_type); + Tensor grad_output("grad_output", std::vector{rows, hidden}, data_type); + Tensor grad_input("grad_input", std::vector{rows, input_cols}, data_type); + Tensor grad_scales("grad_scales", std::vector{rows}, data_type); + + fillUniform(&input); + fillUniform(&scales); + fillUniform(&grad_output); + + std::unique_ptr ref_output = std::make_unique(rows * hidden); + std::unique_ptr ref_grad_input = std::make_unique(rows * input_cols); + std::unique_ptr ref_grad_scales = std::make_unique(rows); + + compute_reference(activation, input.rowwise_cpu_dptr(), scales.rowwise_cpu_dptr(), + grad_output.rowwise_cpu_dptr(), ref_output.get(), + ref_grad_input.get(), ref_grad_scales.get(), rows, hidden, interleave, + compute_grad_scales); + + switch (activation) { + case ScaledActivationCase::kSwiGLU: + nvte_scaled_swiglu(input.data(), scales.data(), output.data(), interleave, 0); + nvte_scaled_dswiglu(grad_output.data(), input.data(), scales.data(), grad_input.data(), + compute_grad_scales ? grad_scales.data() : nullptr, interleave, 0); + break; + case ScaledActivationCase::kClampedSwiGLU: + nvte_scaled_clamped_swiglu(input.data(), scales.data(), output.data(), kClampedLimit, + kClampedAlpha, kClampedLinearOffset, interleave, 0); + nvte_scaled_clamped_dswiglu( + grad_output.data(), input.data(), scales.data(), grad_input.data(), + compute_grad_scales ? grad_scales.data() : nullptr, kClampedLimit, kClampedAlpha, + kClampedLinearOffset, interleave, 0); + break; + case ScaledActivationCase::kSReLU: + nvte_scaled_srelu(input.data(), scales.data(), output.data(), 0); + nvte_scaled_dsrelu(grad_output.data(), input.data(), scales.data(), grad_input.data(), + compute_grad_scales ? grad_scales.data() : nullptr, 0); + break; + } + + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + auto [atol, rtol] = getTolerances(data_type); + if (data_type == DType::kFloat32) { + atol = 5e-5; + rtol = 5e-5; + } + compareResults("scaled_activation_output", output, ref_output.get(), true, atol, rtol); + compareResults("scaled_activation_grad_input", grad_input, ref_grad_input.get(), true, atol, + rtol); + if (compute_grad_scales) { + compareResults("scaled_activation_grad_scales", grad_scales, ref_grad_scales.get(), true, atol, + rtol); + } +} + +class ScaledActivationTest + : public ::testing::TestWithParam< + std::tuple, int64_t, + bool>> { +}; + +std::string test_name_generator( + const testing::TestParamInfo &info) { + const auto activation = std::get<0>(info.param); + const auto data_type = std::get<1>(info.param); + const auto scale_type = std::get<2>(info.param); + const auto shape = std::get<3>(info.param); + const auto interleave = std::get<4>(info.param); + const auto compute_grad_scales = std::get<5>(info.param); + return std::string(activation_name(activation)) + "_data_" + test::typeName(data_type) + + "_scale_" + test::typeName(scale_type) + "_m_" + std::to_string(shape.first) + "_h_" + + std::to_string(shape.second) + "_interleave_" + std::to_string(interleave) + + (compute_grad_scales ? "_with_scale_grad" : "_no_scale_grad"); +} + +} // namespace + +TEST_P(ScaledActivationTest, ForwardBackward) { + const auto activation = std::get<0>(GetParam()); + const auto data_type = std::get<1>(GetParam()); + const auto scale_type = std::get<2>(GetParam()); + const auto shape = std::get<3>(GetParam()); + const auto interleave = std::get<4>(GetParam()); + const auto compute_grad_scales = std::get<5>(GetParam()); + + if (activation == ScaledActivationCase::kSReLU && interleave != 0) { + GTEST_SKIP() << "Interleave has no meaning for SReLU."; + } + if (activation != ScaledActivationCase::kSReLU && interleave > 0 && + shape.second % static_cast(interleave) != 0) { + GTEST_SKIP() << "Hidden size must be divisible by GLU interleave."; + } + + using namespace test; + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(data_type, DataT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(scale_type, ScaleT, { + run_scaled_activation_test(activation, shape.first, shape.second, interleave, + compute_grad_scales); + }); + }); +} + +// Test axes (the six tuple elements consumed by ScaledActivationTest): +// 1. Activation : SwiGLU and ClampedSwiGLU are gated (input is [M, 2H]); +// SReLU is unary (input is [M, H], no gate split). +// 2. Data dtype : dtype of the activation input/output tensors. +// 3. Scale dtype : dtype of act_scales / grad_act_scales. +// 4. Shape {rows, hidden}: rows = M (tokens), hidden = H (output width; gated input is 2H). +// 5. GLU interleave : 0 = contiguous [a | b]; 32 = interleaved a/b blocks. Only valid +// for gated activations with hidden % 32 == 0; SReLU skips != 0. +// 6. compute_grad_scales : whether the backward also reduces grad_act_scales. + +// Interleave is swept over {0, 32}; invalid combinations -- SReLU with any nonzero interleave, or +// a gated activation whose hidden is not divisible by the interleave -- are skipped at runtime by +// the GTEST_SKIP guards in the test body. +INSTANTIATE_TEST_SUITE_P( + OperatorTest_ScaledActivation, ScaledActivationTest, + ::testing::Combine( + ::testing::Values(ScaledActivationCase::kSwiGLU, ScaledActivationCase::kClampedSwiGLU, + ScaledActivationCase::kSReLU), + ::testing::Values(DType::kFloat32, DType::kBFloat16), // data dtype + ::testing::Values(DType::kFloat32, DType::kBFloat16), // scale dtype + ::testing::Values(std::pair{17, 64}, // aligned + interleaved + std::pair{13, 100}, // scalar fallback + std::pair{1024, 2048}), // large FFN-ish width + ::testing::Values(0, 32), // contiguous + interleaved + ::testing::Values(false, true)), // grad_act_scales off / on + test_name_generator); + +// Keep FP16 coverage focused on representative aligned and scalar-fallback shapes instead of +// multiplying it across the full shape matrix above. +INSTANTIATE_TEST_SUITE_P( + OperatorTest_ScaledActivation_FP16, ScaledActivationTest, + ::testing::Combine( + ::testing::Values(ScaledActivationCase::kSwiGLU, ScaledActivationCase::kClampedSwiGLU, + ScaledActivationCase::kSReLU), + ::testing::Values(DType::kFloat16), // data dtype + ::testing::Values(DType::kFloat32, DType::kFloat16), // scale dtype + ::testing::Values(std::pair{17, 64}, // aligned/interleaved + std::pair{13, 100}), // scalar fallback + ::testing::Values(0, 32), // contiguous + interleaved + ::testing::Values(false, true)), // grad_act_scales off / on + test_name_generator); diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index d7f20137a2..d09876c990 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -253,6 +253,9 @@ list(APPEND transformer_engine_cuda_arch_specific_sources activation/relu_dbias.cu activation/relu_grouped.cu activation/relu_grouped_dbias.cu + activation/scaled_activation.cu + activation/scaled_srelu.cu + activation/scaled_swiglu.cu activation/swiglu.cu activation/swiglu_dbias.cu activation/swiglu_grouped.cu @@ -630,6 +633,9 @@ if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) activation/relu_dbias.cu activation/relu_grouped.cu activation/relu_grouped_dbias.cu + activation/scaled_activation.cu + activation/scaled_srelu.cu + activation/scaled_swiglu.cu activation/swiglu.cu activation/swiglu_dbias.cu activation/swiglu_grouped.cu diff --git a/transformer_engine/common/activation/scaled_activation.cu b/transformer_engine/common/activation/scaled_activation.cu new file mode 100644 index 0000000000..4dce3185e0 --- /dev/null +++ b/transformer_engine/common/activation/scaled_activation.cu @@ -0,0 +1,642 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* Vectorization model (gated): + * + * With no GLU interleave, the row is laid out as: + * [ act[0:H] | gate[0:H] ] + * With GLU interleave (e.g. 32): + * [ act[0:32] | gate[0:32] | act[32:64] | gate[32:64] | ... ] + * + * Backward uses fused SiLU / ClampedSiLU specializations (one sigmoid) matching + * gated_fp8.cuh; other ActOP/DActOP pairs call the ops directly. + */ + +#include +#include + +#include "../common.h" +#include "../util/math.h" +#include "./scaled_activation.h" + +namespace transformer_engine { +namespace { + +using namespace detail::scaled_activation; + +using WarpReducer = Reducer; + +// blockDim.x must be a multiple of warp size and <= kReductionThreads. +__device__ __forceinline__ float block_reduce_sum(float value, float *smem) { + const int lane = threadIdx.x % THREADS_PER_WARP; + const int warp = threadIdx.x / THREADS_PER_WARP; + const int num_warps = blockDim.x / THREADS_PER_WARP; + Empty params = {}; + WarpReducer reducer(params, /*bidm=*/0, /*bidn=*/0, /*warp_m=*/0, /*warp_n=*/0, lane, + /*smem=*/nullptr); + Sum sum; + + value = reducer.reduce(value, sum); + if (lane == 0) { + smem[warp] = value; + } + __syncthreads(); + + value = threadIdx.x < num_warps ? smem[lane] : 0.0f; + return warp == 0 ? reducer.reduce(value, sum) : value; +} + +// --------------------------------------------------------------------------- +// Device helpers: fused Act / DAct (IEEE sigmoid / expf) +// --------------------------------------------------------------------------- + +template +__device__ __forceinline__ float gated_forward_value(const float act_in, const float gate_in, + const ParamOP ¶m) { + if constexpr (std::is_same::value) { + const float gate = fminf(fmaxf(-param.limit, gate_in), param.limit) + param.glu_linear_offset; + return ActOP(act_in, param) * gate; + } else { + return ActOP(act_in, param) * gate_in; + } +} + +template +__device__ __forceinline__ void gated_backward_values(const float act_in, const float gate_in, + const ParamOP ¶m, float *dact, + float *dgate, float *unscaled) { + Empty empty = {}; + float act_x = 0.0f; + float dact_x = 0.0f; + float gate = gate_in; + bool dgate_mask = true; + + if constexpr (std::is_same::value) { + dgate_mask = gate_in <= param.limit && gate_in >= -param.limit; + gate = fminf(fmaxf(-param.limit, gate_in), param.limit) + param.glu_linear_offset; + const bool dact_mask = act_in <= param.limit; + const float clamped_act_in = fminf(act_in, param.limit); + const float s = sigmoid(param.alpha * clamped_act_in, empty); + act_x = clamped_act_in * s; + dact_x = dact_mask ? s + param.alpha * clamped_act_in * s * (1.0f - s) : 0.0f; + } else if constexpr ((ActOP == &silu) && (DActOP == &dsilu)) { + const float s = sigmoid(act_in, empty); + act_x = act_in * s; + dact_x = s + act_in * s * (1.0f - s); + } else { + act_x = ActOP(act_in, param); + dact_x = DActOP(act_in, param); + } + + *unscaled = act_x * gate; + *dact = dact_x * gate; + *dgate = dgate_mask ? act_x : 0.0f; +} + +// --------------------------------------------------------------------------- +// Gated kernels +// --------------------------------------------------------------------------- + +template +__global__ void __launch_bounds__(kThreads, 4) + scaled_gated_forward_kernel(const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, OutputT *__restrict__ output, + const size_t rows, const size_t hidden, const size_t segment_size, + const size_t num_segments, const size_t num_vectors_per_segment, + const ParamOP param) { + const size_t total_vectors = rows * num_segments * num_vectors_per_segment; + for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < total_vectors; + tid += gridDim.x * blockDim.x) { + const size_t vector_idx = tid % num_vectors_per_segment; + const size_t segment = (tid / num_vectors_per_segment) % num_segments; + const size_t row = tid / (num_vectors_per_segment * num_segments); + const size_t input_segment_offset = row * hidden * 2 + segment * segment_size * 2; + const size_t output_segment_offset = row * hidden + segment * segment_size; + + VectorizedLoader act_loader(input + input_segment_offset, segment_size); + VectorizedLoader gate_loader(input + input_segment_offset + segment_size, + segment_size); + VectorizedStorer output_storer(output + output_segment_offset, + segment_size); + act_loader.load(vector_idx, segment_size); + gate_loader.load(vector_idx, segment_size); + const float scale = static_cast(act_scales[row]); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + const float unscaled = gated_forward_value( + static_cast(act_loader.separate()[lane]), + static_cast(gate_loader.separate()[lane]), param); + output_storer.separate()[lane] = static_cast(unscaled * scale); + } + output_storer.store(vector_idx, segment_size); + } +} + +template +__global__ void __launch_bounds__(kReductionThreads, 4) + scaled_gated_backward_kernel(const GradT *__restrict__ grad_output, + const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, + OutputT *__restrict__ grad_input, + GradScaleT *__restrict__ grad_act_scales, const size_t rows, + const size_t hidden, const size_t segment_size, + const size_t num_segments, const size_t num_vectors_per_segment, + const ParamOP param) { + __shared__ float smem[kReductionWarps]; + const size_t row = blockIdx.x; + (void)rows; + float scale_grad = 0.0f; + const float scale = static_cast(act_scales[row]); + + const size_t row_vectors = num_segments * num_vectors_per_segment; + for (size_t row_vector_idx = threadIdx.x; row_vector_idx < row_vectors; + row_vector_idx += blockDim.x) { + const size_t segment = row_vector_idx / num_vectors_per_segment; + const size_t vector_idx = row_vector_idx % num_vectors_per_segment; + const size_t input_segment_offset = row * hidden * 2 + segment * segment_size * 2; + const size_t output_segment_offset = row * hidden + segment * segment_size; + VectorizedLoader grad_loader(grad_output + output_segment_offset, + segment_size); + VectorizedLoader act_loader(input + input_segment_offset, segment_size); + VectorizedLoader gate_loader(input + input_segment_offset + segment_size, + segment_size); + VectorizedStorer act_storer(grad_input + input_segment_offset, + segment_size); + VectorizedStorer gate_storer( + grad_input + input_segment_offset + segment_size, segment_size); + + grad_loader.load(vector_idx, segment_size); + act_loader.load(vector_idx, segment_size); + gate_loader.load(vector_idx, segment_size); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + float dact = 0.0f; + float dgate = 0.0f; + float unscaled = 0.0f; + gated_backward_values( + static_cast(act_loader.separate()[lane]), + static_cast(gate_loader.separate()[lane]), param, &dact, &dgate, &unscaled); + const float grad = static_cast(grad_loader.separate()[lane]); + if constexpr (ComputeScaleGrad) { + scale_grad += grad * unscaled; + } + + const float scaled_grad = grad * scale; + act_storer.separate()[lane] = static_cast(scaled_grad * dact); + gate_storer.separate()[lane] = static_cast(scaled_grad * dgate); + } + act_storer.store(vector_idx, segment_size); + gate_storer.store(vector_idx, segment_size); + } + + if constexpr (ComputeScaleGrad) { + scale_grad = block_reduce_sum(scale_grad, smem); + if (threadIdx.x == 0) { + grad_act_scales[row] = static_cast(scale_grad); + } + } +} + +// --------------------------------------------------------------------------- +// Unary kernels +// --------------------------------------------------------------------------- + +template +__global__ void __launch_bounds__(kThreads, 4) + scaled_unary_forward_kernel(const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, OutputT *__restrict__ output, + const size_t rows, const size_t hidden, + const size_t num_vectors_per_row, const ParamOP param) { + const size_t total_vectors = rows * num_vectors_per_row; + for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < total_vectors; + tid += gridDim.x * blockDim.x) { + const size_t vector_idx = tid % num_vectors_per_row; + const size_t row = tid / num_vectors_per_row; + VectorizedLoader input_loader(input + row * hidden, hidden); + VectorizedStorer output_storer(output + row * hidden, hidden); + input_loader.load(vector_idx, hidden); + const float scale = static_cast(act_scales[row]); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + const float unscaled = ActOP(static_cast(input_loader.separate()[lane]), param); + output_storer.separate()[lane] = static_cast(unscaled * scale); + } + output_storer.store(vector_idx, hidden); + } +} + +template +__global__ void __launch_bounds__(kReductionThreads, 4) + scaled_unary_backward_kernel(const GradT *__restrict__ grad_output, + const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, + OutputT *__restrict__ grad_input, + GradScaleT *__restrict__ grad_act_scales, const size_t rows, + const size_t hidden, const size_t num_vectors_per_row, + const ParamOP param) { + __shared__ float smem[kReductionWarps]; + const size_t row = blockIdx.x; + (void)rows; + float scale_grad = 0.0f; + const float scale = static_cast(act_scales[row]); + + VectorizedLoader grad_loader(grad_output + row * hidden, hidden); + VectorizedLoader input_loader(input + row * hidden, hidden); + VectorizedStorer grad_input_storer(grad_input + row * hidden, hidden); + for (size_t vector_idx = threadIdx.x; vector_idx < num_vectors_per_row; + vector_idx += blockDim.x) { + grad_loader.load(vector_idx, hidden); + input_loader.load(vector_idx, hidden); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + const float x = static_cast(input_loader.separate()[lane]); + const float grad = static_cast(grad_loader.separate()[lane]); + if constexpr (ComputeScaleGrad) { + scale_grad += grad * ActOP(x, param); + } + grad_input_storer.separate()[lane] = static_cast(grad * scale * DActOP(x, param)); + } + grad_input_storer.store(vector_idx, hidden); + } + + if constexpr (ComputeScaleGrad) { + scale_grad = block_reduce_sum(scale_grad, smem); + if (threadIdx.x == 0) { + grad_act_scales[row] = static_cast(scale_grad); + } + } +} + +// --------------------------------------------------------------------------- +// Tensor checks +// --------------------------------------------------------------------------- + +void check_gated_forward_tensors(const Tensor *input, const Tensor *act_scales, + const Tensor *output, const int64_t glu_interleave_size, + const char *api_name, size_t *rows, size_t *hidden) { + const auto input_dims = input->flat_2d_dims(); + const auto output_dims = output->flat_2d_dims(); + NVTE_CHECK(input_dims[0] == output_dims[0], api_name, ": input/output row mismatch."); + NVTE_CHECK(input_dims[1] == output_dims[1] * 2, api_name, + ": gated input last dimension must be twice output last dimension."); + NVTE_CHECK(glu_interleave_size >= 0, api_name, ": glu_interleave_size must be non-negative."); + if (glu_interleave_size > 0) { + NVTE_CHECK(glu_interleave_size % 32 == 0, api_name, + ": nonzero glu_interleave_size must be a multiple of 32."); + NVTE_CHECK(output_dims[1] % static_cast(glu_interleave_size) == 0, api_name, + ": output last dimension must be divisible by glu_interleave_size."); + } + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + *rows = input_dims[0]; + *hidden = output_dims[1]; +} + +void check_gated_backward_tensors(const Tensor *grad_output, const Tensor *input, + const Tensor *act_scales, const Tensor *grad_input, + const Tensor *grad_act_scales, const int64_t glu_interleave_size, + const char *api_name, size_t *rows, size_t *hidden) { + const auto grad_dims = grad_output->flat_2d_dims(); + const auto input_dims = input->flat_2d_dims(); + const auto grad_input_dims = grad_input->flat_2d_dims(); + NVTE_CHECK(grad_dims[0] == input_dims[0] && input_dims[0] == grad_input_dims[0], api_name, + ": input/grad row mismatch."); + NVTE_CHECK(input_dims[1] == grad_dims[1] * 2 && grad_input_dims[1] == input_dims[1], api_name, + ": gated backward dimensions are inconsistent."); + NVTE_CHECK(glu_interleave_size >= 0, api_name, ": glu_interleave_size must be non-negative."); + if (glu_interleave_size > 0) { + NVTE_CHECK(glu_interleave_size % 32 == 0, api_name, + ": nonzero glu_interleave_size must be a multiple of 32."); + NVTE_CHECK(grad_dims[1] % static_cast(glu_interleave_size) == 0, api_name, + ": grad last dimension must be divisible by glu_interleave_size."); + } + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + if (grad_act_scales != nullptr) { + NVTE_CHECK(grad_act_scales->numel() == input_dims[0], api_name, + ": grad_act_scales must have one value per row."); + } + *rows = input_dims[0]; + *hidden = grad_dims[1]; +} + +void check_unary_forward_tensors(const Tensor *input, const Tensor *act_scales, + const Tensor *output, const char *api_name, size_t *rows, + size_t *hidden) { + const auto input_dims = input->flat_2d_dims(); + const auto output_dims = output->flat_2d_dims(); + NVTE_CHECK(input_dims[0] == output_dims[0] && input_dims[1] == output_dims[1], api_name, + ": input/output shapes must match."); + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + *rows = input_dims[0]; + *hidden = output_dims[1]; +} + +void check_unary_backward_tensors(const Tensor *grad_output, const Tensor *input, + const Tensor *act_scales, const Tensor *grad_input, + const Tensor *grad_act_scales, const char *api_name, size_t *rows, + size_t *hidden) { + const auto grad_dims = grad_output->flat_2d_dims(); + const auto input_dims = input->flat_2d_dims(); + const auto grad_input_dims = grad_input->flat_2d_dims(); + NVTE_CHECK(grad_dims[0] == input_dims[0] && input_dims[0] == grad_input_dims[0], api_name, + ": input/grad row mismatch."); + NVTE_CHECK(grad_dims[1] == input_dims[1] && input_dims[1] == grad_input_dims[1], api_name, + ": unary backward dimensions are inconsistent."); + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + if (grad_act_scales != nullptr) { + NVTE_CHECK(grad_act_scales->numel() == input_dims[0], api_name, + ": grad_act_scales must have one value per row."); + } + *rows = input_dims[0]; + *hidden = grad_dims[1]; +} + +} // namespace + +// --------------------------------------------------------------------------- +// Launch implementations +// --------------------------------------------------------------------------- + +using namespace detail::scaled_activation; + +template +void launch_scaled_gated_forward(const NVTETensor nvte_input, const NVTETensor nvte_act_scales, + NVTETensor nvte_output, ParamOP param, int64_t glu_interleave_size, + cudaStream_t stream, const char *api_name) { + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *output = convertNVTETensorCheck(nvte_output); + size_t rows = 0; + size_t hidden = 0; + check_gated_forward_tensors(input, act_scales, output, glu_interleave_size, api_name, &rows, + &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(output->data.dtype, OutputT, { + constexpr int nvec = 32 / static_cast(std::max(sizeof(InputT), sizeof(OutputT))); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto output_ptr = reinterpret_cast(output->data.dptr); + const size_t segment_size = + glu_interleave_size > 0 ? static_cast(glu_interleave_size) : hidden; + const size_t num_segments = glu_interleave_size > 0 ? hidden / segment_size : 1; + const auto align = row_vector_alignment(segment_size, nvec, input_ptr, + input_ptr + segment_size, output_ptr); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, segment_size, nvec, sizeof(InputT)) + : segment_size; + const int blocks = static_cast(std::min( + DIVUP(rows * num_segments * num_vectors, static_cast(kThreads)), 65535)); + if (use_vector) { + scaled_gated_forward_kernel + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + segment_size, num_segments, num_vectors, param); + } else { + scaled_gated_forward_kernel<1, InputT, ScaleT, OutputT, ParamOP, ActOP> + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + segment_size, num_segments, segment_size, param); + } + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void launch_scaled_gated_backward(const NVTETensor nvte_grad_output, const NVTETensor nvte_input, + const NVTETensor nvte_act_scales, NVTETensor nvte_grad_input, + NVTETensor nvte_grad_act_scales, ParamOP param, + int64_t glu_interleave_size, cudaStream_t stream, + const char *api_name) { + const Tensor *grad_output = convertNVTETensorCheck(nvte_grad_output); + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *grad_input = convertNVTETensorCheck(nvte_grad_input); + Tensor *grad_act_scales = + nvte_grad_act_scales == nullptr ? nullptr : convertNVTETensorCheck(nvte_grad_act_scales); + size_t rows = 0; + size_t hidden = 0; + check_gated_backward_tensors(grad_output, input, act_scales, grad_input, grad_act_scales, + glu_interleave_size, api_name, &rows, &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_output->data.dtype, GradT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_input->data.dtype, OutputT, { + constexpr int nvec = + 32 / static_cast(std::max({sizeof(GradT), sizeof(InputT), sizeof(OutputT)})); + const auto grad_ptr = reinterpret_cast(grad_output->data.dptr); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto grad_input_ptr = reinterpret_cast(grad_input->data.dptr); + const size_t segment_size = + glu_interleave_size > 0 ? static_cast(glu_interleave_size) : hidden; + const size_t num_segments = glu_interleave_size > 0 ? hidden / segment_size : 1; + const auto align = row_vector_alignment(segment_size, nvec, grad_ptr, input_ptr, + input_ptr + segment_size, grad_input_ptr, + grad_input_ptr + segment_size); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, segment_size, nvec, sizeof(InputT)) + : segment_size; + const int reduction_threads = choose_reduction_threads(num_segments * num_vectors); + if (grad_act_scales == nullptr) { + if (use_vector) { + scaled_gated_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, + segment_size, num_segments, num_vectors, param); + } else { + scaled_gated_backward_kernel<1, false, GradT, InputT, ScaleT, OutputT, OutputT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, + segment_size, num_segments, segment_size, param); + } + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_act_scales->data.dtype, GradScaleT, { + auto grad_act_scales_ptr = reinterpret_cast(grad_act_scales->data.dptr); + if (use_vector) { + scaled_gated_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, segment_size, num_segments, num_vectors, param); + } else { + scaled_gated_backward_kernel<1, true, GradT, InputT, ScaleT, OutputT, GradScaleT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, segment_size, num_segments, segment_size, param); + } + }); + } + }); + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void launch_scaled_unary_forward(const NVTETensor nvte_input, const NVTETensor nvte_act_scales, + NVTETensor nvte_output, ParamOP param, cudaStream_t stream, + const char *api_name) { + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *output = convertNVTETensorCheck(nvte_output); + size_t rows = 0; + size_t hidden = 0; + check_unary_forward_tensors(input, act_scales, output, api_name, &rows, &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(output->data.dtype, OutputT, { + constexpr int nvec = 32 / static_cast(std::max(sizeof(InputT), sizeof(OutputT))); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto output_ptr = reinterpret_cast(output->data.dptr); + const auto align = row_vector_alignment(hidden, nvec, input_ptr, output_ptr); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, hidden, nvec, sizeof(InputT)) : hidden; + const int blocks = static_cast( + std::min(DIVUP(rows * num_vectors, static_cast(kThreads)), 65535)); + if (use_vector) { + scaled_unary_forward_kernel + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + num_vectors, param); + } else { + scaled_unary_forward_kernel<1, InputT, ScaleT, OutputT, ParamOP, ActOP> + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + hidden, param); + } + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void launch_scaled_unary_backward(const NVTETensor nvte_grad_output, const NVTETensor nvte_input, + const NVTETensor nvte_act_scales, NVTETensor nvte_grad_input, + NVTETensor nvte_grad_act_scales, ParamOP param, + cudaStream_t stream, const char *api_name) { + const Tensor *grad_output = convertNVTETensorCheck(nvte_grad_output); + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *grad_input = convertNVTETensorCheck(nvte_grad_input); + Tensor *grad_act_scales = + nvte_grad_act_scales == nullptr ? nullptr : convertNVTETensorCheck(nvte_grad_act_scales); + size_t rows = 0; + size_t hidden = 0; + check_unary_backward_tensors(grad_output, input, act_scales, grad_input, grad_act_scales, + api_name, &rows, &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_output->data.dtype, GradT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_input->data.dtype, OutputT, { + constexpr int nvec = + 32 / static_cast(std::max({sizeof(GradT), sizeof(InputT), sizeof(OutputT)})); + const auto grad_ptr = reinterpret_cast(grad_output->data.dptr); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto grad_input_ptr = reinterpret_cast(grad_input->data.dptr); + const auto align = + row_vector_alignment(hidden, nvec, grad_ptr, input_ptr, grad_input_ptr); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, hidden, nvec, sizeof(InputT)) + : hidden; + const int reduction_threads = choose_reduction_threads(num_vectors); + if (grad_act_scales == nullptr) { + if (use_vector) { + scaled_unary_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, + num_vectors, param); + } else { + scaled_unary_backward_kernel<1, false, GradT, InputT, ScaleT, OutputT, OutputT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, hidden, + param); + } + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_act_scales->data.dtype, GradScaleT, { + auto grad_act_scales_ptr = reinterpret_cast(grad_act_scales->data.dptr); + if (use_vector) { + scaled_unary_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, num_vectors, param); + } else { + scaled_unary_backward_kernel<1, true, GradT, InputT, ScaleT, OutputT, GradScaleT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, hidden, param); + } + }); + } + }); + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// --------------------------------------------------------------------------- +// Explicit instantiations +// --------------------------------------------------------------------------- + +template void launch_scaled_gated_forward>(const NVTETensor, + const NVTETensor, NVTETensor, + Empty, int64_t, cudaStream_t, + const char *); +template void launch_scaled_gated_backward, dsilu>( + const NVTETensor, const NVTETensor, const NVTETensor, NVTETensor, NVTETensor, Empty, int64_t, + cudaStream_t, const char *); + +template void launch_scaled_gated_forward>( + const NVTETensor, const NVTETensor, NVTETensor, ClampedSwiGLUParam, int64_t, cudaStream_t, + const char *); +template void launch_scaled_gated_backward, + clamped_dsilu>( + const NVTETensor, const NVTETensor, const NVTETensor, NVTETensor, NVTETensor, + ClampedSwiGLUParam, int64_t, cudaStream_t, const char *); + +template void launch_scaled_unary_forward>(const NVTETensor, + const NVTETensor, NVTETensor, + Empty, cudaStream_t, + const char *); +template void launch_scaled_unary_backward, dsrelu>( + const NVTETensor, const NVTETensor, const NVTETensor, NVTETensor, NVTETensor, Empty, + cudaStream_t, const char *); + +} // namespace transformer_engine diff --git a/transformer_engine/common/activation/scaled_activation.h b/transformer_engine/common/activation/scaled_activation.h new file mode 100644 index 0000000000..391bff6841 --- /dev/null +++ b/transformer_engine/common/activation/scaled_activation.h @@ -0,0 +1,90 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* Scaled activations apply an activation, multiply by a per-row scale + * (act_scales[row]), do all math in fp32, and cast once at the store. The + * backward path optionally also reduces the gradient of the per-row scale. + * + * Public launch APIs are templated on ParamOP / ActOP / DActOP (same shape as + * gated_act_fn). Kernel definitions and explicit instantiations live in + * scaled_activation.cu. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_ACTIVATION_SCALED_ACTIVATION_H_ +#define TRANSFORMER_ENGINE_COMMON_ACTIVATION_SCALED_ACTIVATION_H_ + +#include + +#include "../common.h" +#include "../util/math.h" +#include "../util/vectorized_pointwise.h" +#include "../utils.cuh" + +namespace transformer_engine { +namespace detail { +namespace scaled_activation { + +constexpr int kThreads = unary_kernel_threads; +constexpr int kReductionThreads = 256; +constexpr int kReductionWarps = kReductionThreads / THREADS_PER_WARP; + +// Pick a CTA size for one-block-per-row scale-grad: enough threads to cover +// row_vectors, rounded up to a warp multiple, capped at kReductionThreads. +inline int choose_reduction_threads(const size_t row_vectors) { + if (row_vectors >= static_cast(kReductionThreads)) { + return kReductionThreads; + } + const int needed = static_cast(row_vectors); + int rounded = (needed + THREADS_PER_WARP - 1) / THREADS_PER_WARP * THREADS_PER_WARP; + if (rounded < THREADS_PER_WARP) { + rounded = THREADS_PER_WARP; + } + return rounded; +} + +template +Alignment row_vector_alignment(const size_t lead_dim, const int nvec, const Ptrs... ptrs) { + if (nvec == 1) { + return Alignment::SAME_ALIGNED; + } + if (lead_dim % static_cast(nvec) != 0) { + return Alignment::DIFFERENT; + } + const auto align = CheckAlignment(lead_dim, nvec, ptrs...); + return align == Alignment::SAME_ALIGNED ? Alignment::SAME_ALIGNED : Alignment::DIFFERENT; +} + +} // namespace scaled_activation +} // namespace detail + +template +void launch_scaled_gated_forward(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, ParamOP param, int64_t glu_interleave_size, + cudaStream_t stream, const char *api_name); + +template +void launch_scaled_gated_backward(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, ParamOP param, + int64_t glu_interleave_size, cudaStream_t stream, + const char *api_name); + +template +void launch_scaled_unary_forward(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, ParamOP param, cudaStream_t stream, + const char *api_name); + +template +void launch_scaled_unary_backward(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, ParamOP param, cudaStream_t stream, + const char *api_name); + +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_ACTIVATION_SCALED_ACTIVATION_H_ diff --git a/transformer_engine/common/activation/scaled_srelu.cu b/transformer_engine/common/activation/scaled_srelu.cu new file mode 100644 index 0000000000..2e81bab669 --- /dev/null +++ b/transformer_engine/common/activation/scaled_srelu.cu @@ -0,0 +1,29 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../common.h" +#include "../util/math.h" +#include "./scaled_activation.h" + +void nvte_scaled_srelu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_srelu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_unary_forward>(input, act_scales, output, param, stream, + "nvte_scaled_srelu"); +} + +void nvte_scaled_dsrelu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_dsrelu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_unary_backward, dsrelu>( + grad, input, act_scales, grad_input, grad_act_scales, param, stream, "nvte_scaled_dsrelu"); +} diff --git a/transformer_engine/common/activation/scaled_swiglu.cu b/transformer_engine/common/activation/scaled_swiglu.cu new file mode 100644 index 0000000000..ee0692b1de --- /dev/null +++ b/transformer_engine/common/activation/scaled_swiglu.cu @@ -0,0 +1,56 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../common.h" +#include "../util/math.h" +#include "./scaled_activation.h" + +void nvte_scaled_swiglu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + int64_t glu_interleave_size, cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_swiglu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_gated_forward>( + input, act_scales, output, param, glu_interleave_size, stream, "nvte_scaled_swiglu"); +} + +void nvte_scaled_dswiglu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, + int64_t glu_interleave_size, cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_dswiglu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_gated_backward, dsilu>( + grad, input, act_scales, grad_input, grad_act_scales, param, glu_interleave_size, stream, + "nvte_scaled_dswiglu"); +} + +void nvte_scaled_clamped_swiglu(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_clamped_swiglu); + using namespace transformer_engine; + ClampedSwiGLUParam param = {limit, alpha, glu_linear_offset}; + launch_scaled_gated_forward>( + input, act_scales, output, param, glu_interleave_size, stream, "nvte_scaled_clamped_swiglu"); +} + +void nvte_scaled_clamped_dswiglu(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_clamped_dswiglu); + using namespace transformer_engine; + ClampedSwiGLUParam param = {limit, alpha, glu_linear_offset}; + launch_scaled_gated_backward, + clamped_dsilu>( + grad, input, act_scales, grad_input, grad_act_scales, param, glu_interleave_size, stream, + "nvte_scaled_clamped_dswiglu"); +} diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 4ed083740d..3df84954bf 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -368,6 +368,43 @@ void nvte_clamped_swiglu(const NVTETensor input, NVTETensor output, float limit, void nvte_clamped_swiglu_v2(const NVTETensor input, NVTETensor output, float limit, float alpha, float glu_linear_offset, cudaStream_t stream); +/*! \brief Computes ScaledSwiGLU without materializing GLU deinterleave. + * + * Computes output = SwiGLU(input) * act_scales[:, None]. + * If glu_interleave_size > 0, input is interpreted as interleaved + * [activation_block, linear_block] chunks of that size. + * + * \param[in] input Input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] output Output tensor of shape [N, H]. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_swiglu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + int64_t glu_interleave_size, cudaStream_t stream); + +/*! \brief Computes ScaledClampedSwiGLU without materializing GLU deinterleave. + * + * Computes output = ClampedSwiGLU(input) * act_scales[:, None]. + * This uses the same clamping, alpha, and linear-offset semantics as + * nvte_clamped_swiglu_v2. + * + * \param[in] input Input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] output Output tensor of shape [N, H]. + * \param[in] limit Clipping limit. + * \param[in] alpha Activation sigmoid alpha. + * \param[in] glu_linear_offset Offset added to linear component after clamping. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_clamped_swiglu(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream); + /*! \brief Computes the gated ReLU activation of the input. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -473,6 +510,47 @@ void nvte_clamped_dswiglu_v2(const NVTETensor grad, const NVTETensor input, NVTE float limit, float alpha, float glu_linear_offset, cudaStream_t stream); +/*! \brief Computes ScaledSwiGLU backward without materializing GLU deinterleave. + * + * The optional grad_act_scales tensor may be null. When present, it receives + * sum(dY * SwiGLU(input), dim=-1). + * + * \param[in] grad Incoming gradient of shape [N, H]. + * \param[in] input Forward input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] grad_input Outgoing gradient of shape [N, H * 2]. + * \param[in,out] grad_act_scales Optional row-wise scale gradient of shape [N], or null. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_dswiglu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, + int64_t glu_interleave_size, cudaStream_t stream); + +/*! \brief Computes ScaledClampedSwiGLU backward without materializing GLU deinterleave. + * + * The optional grad_act_scales tensor may be null. When present, it receives + * sum(dY * ClampedSwiGLU(input), dim=-1). + * + * \param[in] grad Incoming gradient of shape [N, H]. + * \param[in] input Forward input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] grad_input Outgoing gradient of shape [N, H * 2]. + * \param[in,out] grad_act_scales Optional row-wise scale gradient of shape [N], or null. + * \param[in] limit Clipping limit. + * \param[in] alpha Activation sigmoid alpha. + * \param[in] glu_linear_offset Offset added to linear component after clamping. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_clamped_dswiglu(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream); + /*! \brief Computes the gated ReLU activation gradient. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -509,6 +587,33 @@ void nvte_dqgeglu(const NVTETensor grad, const NVTETensor input, NVTETensor outp void nvte_dsreglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes ScaledSReLU. + * + * Computes output = SReLU(input) * act_scales[:, None]. + * + * \param[in] input Input tensor for activation. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] output Output tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_srelu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + cudaStream_t stream); + +/*! \brief Computes ScaledSReLU backward. + * + * The optional grad_act_scales tensor may be null. When present, it receives + * sum(dY * SReLU(input), dim=-1). + * + * \param[in] grad Incoming gradient. + * \param[in] input Forward input tensor. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] grad_input Outgoing input gradient. + * \param[in,out] grad_act_scales Optional row-wise scale gradient of shape [N], or null. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_dsrelu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif From 70957ad5aa8337a12c254e86a6e0b7ffdbbd4b59 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Tue, 14 Jul 2026 17:41:17 -0600 Subject: [PATCH 20/35] [Common] Pass cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA fprop (#3186) * [Common] Pass cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA fprop cuDNN >= 9.24 SDPA (unified engine) accepts cumulative sequence lengths directly (cu_seq_len_q/kv) and can scale ragged offsets stored in coarser units back to elements via a per-tensor ragged offset multiplier. Use both in the f16/bf16 forward to skip the two conversion kernels (cu_seqlens_to_actual_seqlens and cu_seqlens_padded_to_offsets) that previously ran before every varlen fprop: - Bind the user's int32 cu_seqlens buffers as CU_SEQ_LEN_Q/KV for the padding mask, and the token-unit cu_seqlens_padded buffers as ragged offsets for Q/K/V/O/Stats with elements-per-token multipliers. - Gate on cudnn >= 9.24 and !dropout (the FE rejects dropout together with generated stats on the unified engine; TE always generates stats). CU_SEQ_LEN inputs pin implementation selection to the unified engine. - Keep the true batch size on the direct path: cuDNN reads the user's [actual_b+1] buffers, so the quantized max_b graph batch would read out of bounds. Token-dim bucketing (max_t) is unaffected. - No conversion workspace is needed on the direct path. - Factor the layout-group -> multiplier mapping into RaggedOffsetMultipliers (utils.h), shared by the graph builder and the legacy conversion kernel so the two cannot drift. The kernel rewrite also removes a cross-thread read (offsets_v[tid] = offsets_k[cu_seqlens_id]) that raced for quantized-batch tail entries with interleaved layouts. - Backward is unchanged (no backend support yet). NVTE_FUSED_ATTN_DIRECT_SEQLENS=0 disables the new path (testing aid, to be removed before merging). Validated on H100 and Blackwell against cuDNN 9.25: test_dpa_softmax_thd 15/15 in both modes, and direct-vs-legacy fused outputs/grads match for all THD layouts (thd_thd_thd, t3hd, th3d, thd_t2hd, thd_th2d) x MHA/GQA x padding/padding_causal x pad_between_seqs {false,true}. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam * [Common] Pass cu_seqlens directly to cuDNN SDPA FP8/MXFP8 fprop Extend the direct-seqlens path to the FP8/MXFP8 forward: bind the user's int32 cu_seqlens buffers as CU_SEQ_LEN_Q/KV for the padding mask instead of converting them to per-batch lengths with the cu_seqlens_to_actual_seqlens kernel before every call. (Unlike the F16 path, the FP8 path has no THD/ragged support, so this is the only conversion kernel there.) FP8/MXFP8 on the unified engine requires cuDNN >= 9.25 and cuDNN frontend >= 1.26. The frontend is header-only, so its version is a compile-time property; the gate uses a constant-folded CUDNN_FRONTEND_VERSION check (all referenced symbols exist in 1.25, so no preprocessor guards are needed). Dropout with generated stats stays on the legacy path, same as F16. Backward is unchanged (no backend support yet). Validated against cuDNN 9.25 + frontend 1.26 (test_dpa_fp8_vs_f16, padding configs, direct path on with no fallback): 56 passed on H100 (delayed + current scaling), 168 passed on Blackwell (adds MXFP8); zero failures. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [Common] Harden direct-seqlens version gates Address review feedback and a version-mix bug found in testing: - Remove the NVTE_FUSED_ATTN_DIRECT_SEQLENS env override (unnecessary; the version gates fully determine the path). - Check the compile-time CUDNN_VERSION in addition to the runtime version. The cuDNN frontend gates cu_seq_len support on min(compile-time, runtime) version, so e.g. a binary built against 9.24 headers running on a 9.25 library must take the legacy path; a runtime-only check let it attempt the direct fp8 graph, which the frontend rejects ("No suitable implementation") with no fallback. - Add a (currently redundant) CUDNN_FRONTEND_VERSION >= 1.25 check to the f16 gate for symmetry with the fp8 gate. - Raise the fp8 frontend floor from 1.26 to 1.27: 1.26 suffices for this C++ API use, but 1.27 is the floor for the python FE API's fp8 cu_seq_len support (exposed post-1.26-cut), and a single version story per feature avoids a silent gap when TE moves to the python FE API. Smoke-tested on H100: f16 THD 15/15 (direct path, cuDNN 9.24), fp8 padding subset 56 passed via legacy on 9.24, and 56 passed via legacy on the 9.24-compile/9.25-runtime mix that previously failed 56/56. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam * [Common] Lower fp8 direct-seqlens frontend floor to 1.26 Per TE team discussion: 1.26 is all the C++ FE API needs for fp8 + cu_seqlens (the support surface made the 1.26 cut; SDPA_fp8_attributes has had the setters since 1.25). Keep a comment noting that the python FE API requires 1.27 (its sdpa_fp8 binding gained cu_seq_len_q/kv post-1.26-cut), so a future migration to the python FE API knows to raise the floor. Smoke-tested on H100: f16 THD 15/15 (direct, cuDNN 9.24), fp8 padding subset 56 passed via legacy on 9.24 and on the 9.24-compile/9.25-runtime mix. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam * [Common] Fix sm120 THD softmax-stats layout and allocation use_ragged_stats excludes sm120, but the forward Stats declaration used the weaker condition (is_ragged_q && cudnn >= 9.6). On sm120 with THD, fwd therefore declared the ragged-style [b][s][h] stats stride with a null ragged offset (i.e. dense token-major), while bwd read the stats tensor as dense [b][h][s] -- a fwd/bwd layout mismatch. It would also have let the direct-seqlens path set a ragged-offset multiplier on a null ragged offset, a frontend validation error. Use use_ragged_stats for the fwd declaration so fwd and bwd agree, and give the stats allocation the same sm120 exception Max already has: without it the buffer is [num_tokens_q, h, 1], undersized for the dense [b, h, s_q, 1] graph whenever num_tokens_q < b * s_q. Pre-existing issue, independent of the direct-seqlens work. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam * [Common] Rename use_direct_seqlens to use_cu_seqlens_directly Clearer name for the flag controlling whether cu_seqlens buffers are passed straight to cuDNN SDPA; comment wording updated to match. No functional change. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam * [Common] Suppress fn_size lint on fused_attn_arbitrary_seqlen_fwd_impl The direct-seqlens additions push the function to 508 non-comment lines, over cpplint's 500 limit. Per TE team, refactoring this long-standing function is beyond the scope of this PR, so suppress with NOLINT for now. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam * [Common] Pin the UNIFIED implementation on the direct cu_seqlens path cu_seq_len (and the ragged offset multiplier) are unified-engine-only, so with those inputs attached AUTO can only ever resolve to UNIFIED anyway. Pinning changes only the failure mode: an unsupported config fails with the unified engine's specific error instead of auto-selection's generic "no suitable implementation". Ordinary graphs (no cu_seq_len attached) keep AUTO. Matches the cudnn-frontend cu_seq_len sample, which pins for the same reason. Smoke-tested on H100: f16 THD 15/15 via the pinned direct path (cuDNN 9.24); fp8 padding subsets 56 passed via legacy on 9.24 and on the 9.24-compile/9.25-runtime mix. Co-Authored-By: Claude Fable 5 Signed-off-by: Emil Gilliam --------- Signed-off-by: Emil Gilliam Co-authored-by: Claude Fable 5 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sudhakar Singh --- .../fused_attn_f16_arbitrary_seqlen.cu | 176 +++++++++++++----- .../common/fused_attn/fused_attn_fp8.cu | 94 +++++++--- transformer_engine/common/fused_attn/utils.cu | 56 ++---- transformer_engine/common/fused_attn/utils.h | 37 +++- 4 files changed, 252 insertions(+), 111 deletions(-) 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..bf34758a35 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 @@ -94,6 +94,20 @@ void fused_attn_arbitrary_seqlen_fwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } + // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative + // tensor, and can accept ragged offsets in arbitrary units (such as tokens) instead + // of elements. Take advantage of this if possible to avoid 2 extra kernel calls. + const bool use_cu_seqlens_directly = + CUDNN_FRONTEND_VERSION >= 12500 && + // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN + // version, so we'll do the same. + (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && + // This extra restriction is needed because cuDNN frontend doesn't yet allow + // the combination of dropout and stats generation for the fprop unified engine, + // so any such request would always get routed to the old composite SDPA engine + // (which doesn't support cu_seqlens). Remove this restriction when possible. + !is_dropout; + // 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) { @@ -103,14 +117,27 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // so the check passes; ragged offset still provides variable-length boundaries. if (sm_arch_ != 120) { // replace batch size and maximum sequence lengths with maximum token counts - // for query and key/value so the graph is static within each quantization bucket - b = max_b; + // for query and key/value so the graph is static within each quantization bucket. + // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size: + // cuDNN reads the user's [actual_b+1] cu_seqlens buffers, so a quantized batch + // would read out of bounds. + if (!use_cu_seqlens_directly) { + b = max_b; + } s_q = is_ragged_q ? max_t_q : s_q; s_kv = is_ragged_kv ? max_t_kv : s_kv; } } - const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; + const DType ragged_offset_type = + use_cu_seqlens_directly + ? DType::kInt32 // cu_seqlens* are given to us as int32; keep it that way. + : (cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32); + + // Ragged offset multipliers (elements per token); shared with the legacy conversion + // kernel (cu_seqlens_padded_to_offsets) so the two paths cannot drift apart. + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + bool generate_stats = true; // Always return stats try { FADescriptor_v1 descriptor{ @@ -166,8 +193,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr, // S2 std::shared_ptr, // bias std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv + std::shared_ptr, // seq_q / cu_seq_len_q + std::shared_ptr, // seq_kv / cu_seq_len_kv std::shared_ptr, // page_table_k std::shared_ptr, // page_table_v std::shared_ptr, // offset_q @@ -231,6 +258,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_stride({1, 1, 1, 1}) .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); Q->set_ragged_offset(offset_q); + if (use_cu_seqlens_directly) { + Q->set_ragged_offset_multiplier(offset_mults.q); + } } K = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("K").set_stride(k_stride)); V = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("V").set_stride(v_stride)); @@ -250,6 +280,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); K->set_dim({b, hg, s_kv, d_qk}).set_ragged_offset(offset_k); V->set_dim({b, hg, s_kv, d_v}).set_ragged_offset(offset_v); + if (use_cu_seqlens_directly) { + K->set_ragged_offset_multiplier(offset_mults.k); + V->set_ragged_offset_multiplier(offset_mults.v); + } } else { K->set_dim({b, hg, s_kv, d_qk}); V->set_dim({b, hg, s_kv, d_v}); @@ -293,17 +327,38 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding) + .set_cu_seq_len_q(seq_q) + .set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } } if (is_paged_kv) { @@ -363,6 +418,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_data_type(fe::DataType_t::FLOAT)); if (use_ragged_stats) { Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Max->set_ragged_offset_multiplier(offset_mults.stats); + } } else { Max->set_stride({h * s_q, s_q, 1, 1}); } @@ -382,11 +440,17 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_stride({1, 1, 1, 1}) .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); O->set_ragged_offset(offset_o); + if (use_cu_seqlens_directly) { + O->set_ragged_offset_multiplier(offset_mults.o); + } } Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Stats->set_ragged_offset_multiplier(offset_mults.stats); + } } else { Stats->set_stride({h * s_q, s_q, 1, 1}); } @@ -437,18 +501,23 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // Exit to request upper level API to allocate memory if needed // n.b. Care should be taken to align each of the added worksapce tensors to their type. // We do this by adding padding at the end of each separate allocation. + // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is + // needed: cuDNN consumes the user's cu_seqlens buffers as-is. auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; const size_t num_bytes_per_ragged_offset = alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + size_t actual_seqlen_workspace_size = 0; size_t seqlen_offsets_workspace_size = 0; - if (is_ragged_q || is_ragged_kv) { - size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - if (use_ragged_stats) { - seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; - } else { - seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; + if (!use_cu_seqlens_directly) { + if (is_padding) { + actual_seqlen_workspace_size = 2 * num_bytes_per_seqlen; + } + if (is_ragged_q || is_ragged_kv) { + const size_t count = + 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + seqlen_offsets_workspace_size = + (use_ragged_stats ? count + 1 : count) * num_bytes_per_ragged_offset; } } if (workspace == nullptr) { @@ -475,17 +544,22 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } if (is_padding) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void *devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; - cu_seqlens_to_actual_seqlens<<>>( - actual_b, b, static_cast(devPtrCuSeqlensQ), - static_cast(devPtrCuSeqlensKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; + if (use_cu_seqlens_directly) { + variant_pack[seq_q] = devPtrCuSeqlensQ; + variant_pack[seq_kv] = devPtrCuSeqlensKV; + } else { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; + void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; + void *devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; + cu_seqlens_to_actual_seqlens<<>>( + actual_b, b, static_cast(devPtrCuSeqlensQ), + static_cast(devPtrCuSeqlensKV), + static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); + NVTE_CHECK_CUDA(cudaGetLastError()); + variant_pack[seq_q] = devActualSeqlenQ; + variant_pack[seq_kv] = devActualSeqlenKV; + } } if (is_paged_kv) { @@ -493,7 +567,22 @@ void fused_attn_arbitrary_seqlen_fwd_impl( variant_pack[page_table_v] = devPtrPageTableV; } - if (is_ragged_q || is_ragged_kv) { + if (use_cu_seqlens_directly) { + // The token-unit cu_seqlens_padded buffers serve as the ragged offsets; the engine + // applies the per-tensor multipliers set at graph build time. + if (is_ragged_q) { + variant_pack[offset_q] = devPtrSeqOffsetsQ; + variant_pack[offset_o] = devPtrSeqOffsetsQ; + } + if (is_ragged_kv) { + void *devOffsetsKV = offset_mults.kv_from_q ? devPtrSeqOffsetsQ : devPtrSeqOffsetsKV; + variant_pack[offset_k] = devOffsetsKV; + variant_pack[offset_v] = devOffsetsKV; + } + if (use_ragged_stats) { + variant_pack[offset_stats] = devPtrSeqOffsetsQ; + } + } else if (is_ragged_q || is_ragged_kv) { constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block) / nthreads_per_block; void *devOffsets = @@ -517,9 +606,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -549,7 +637,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } catch (cudnn_frontend::cudnnException &e) { NVTE_ERROR(e.what()); } -} +} // NOLINT(readability/fn_size) void fused_attn_arbitrary_seqlen_bwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, @@ -1034,9 +1122,10 @@ void fused_attn_arbitrary_seqlen_bwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + const RaggedOffsetMultipliers offset_mults(nvte_get_qkv_layout_group(qkv_layout), h, hg, d_qk, + d_v); cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1142,7 +1231,10 @@ void fused_attn_arbitrary_seqlen_fwd( Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + // sm120 does not use ragged stats: the graph declares a dense + // [b, h, s_q, 1] stats tensor, so allocate to match (same as Max below). + if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && + (sm_arch_ != 120)) { output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index eab1ae02e6..000af41aee 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -60,6 +60,24 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative + // tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike + // the F16 path, the FP8 path has no THD/ragged-offset support, so only the + // cu_seqlens_to_actual_seqlens conversion applies here. Also note that the + // needed versions of cuDNN backend and frontend are higher than for F16.) + const bool use_cu_seqlens_directly = + // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). + // Note: For the Python API, 1.27 is required. + CUDNN_FRONTEND_VERSION >= 12600 && + // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN + // version, so we'll do the same. + (CUDNN_VERSION >= 92500 && cudnn_runtime_version >= 92500) && + // This extra restriction is needed because cuDNN frontend doesn't yet allow + // the combination of dropout and stats generation for the fprop unified engine, + // so any such request would always get routed to the old composite SDPA engine + // (which doesn't support cu_seqlens). Remove this restriction when possible. + !is_dropout; + try { FADescriptor_v1 descriptor{b, h, @@ -262,17 +280,38 @@ void fused_attn_fp8_fwd_impl( // } if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding) + .set_cu_seq_len_q(seq_q) + .set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } } if (is_dropout) { @@ -379,8 +418,10 @@ void fused_attn_fp8_fwd_impl( auto plan_workspace_size = mha_graph->get_workspace_size(); - // Exit to request upper level API to allocate memory if needed - size_t actual_seqlen_workspace_size = 2 * b * sizeof(int32_t); + // Exit to request upper level API to allocate memory if needed. + // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is + // needed: cuDNN consumes the user's cu_seqlens buffers as-is. + size_t actual_seqlen_workspace_size = use_cu_seqlens_directly ? 0 : 2 * b * sizeof(int32_t); if (workspace == nullptr) { *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; @@ -417,17 +458,22 @@ void fused_attn_fp8_fwd_impl( } */ if (is_padding) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); - cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) - static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; + if (use_cu_seqlens_directly) { + variant_pack[seq_q] = devPtrcuSeqlensQ; + variant_pack[seq_kv] = devPtrcuSeqlensKV; + } else { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; + void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; + void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + cu_seqlens_to_actual_seqlens<<>>( + b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), + static_cast(devActualSeqlenKV)); + NVTE_CHECK_CUDA(cudaGetLastError()); + variant_pack[seq_q] = devActualSeqlenQ; + variant_pack[seq_kv] = devActualSeqlenKV; + } } if (is_dropout) { diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index 3e628b6581..9b54a64cbe 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -429,71 +429,43 @@ __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, // convert cu_seqlens_padded to offsets template __device__ void cu_seqlens_padded_to_offsets_impl( - NVTE_QKV_Layout_Group layout_group, int64_t actual_b, int64_t max_b, int64_t h, int64_t hg, - int64_t d_qk, int64_t d_v, const int32_t *cu_seqlens_q_padded, - const int32_t *cu_seqlens_kv_padded, OFFSETS_T *offsets_q, OFFSETS_T *offsets_k, - OFFSETS_T *offsets_v, OFFSETS_T *offsets_o, OFFSETS_T *offsets_s) { + const RaggedOffsetMultipliers &mults, int64_t actual_b, int64_t max_b, + const int32_t *cu_seqlens_q_padded, const int32_t *cu_seqlens_kv_padded, OFFSETS_T *offsets_q, + OFFSETS_T *offsets_k, OFFSETS_T *offsets_v, OFFSETS_T *offsets_o, OFFSETS_T *offsets_s) { size_t tid = blockIdx.x * blockDim.x + threadIdx.x; auto cu_seqlens_id = min(tid, actual_b); if (tid <= max_b) { if (offsets_s != nullptr) { - offsets_s[tid] = h * cu_seqlens_q_padded[cu_seqlens_id]; + offsets_s[tid] = mults.stats * cu_seqlens_q_padded[cu_seqlens_id]; } if (offsets_q != nullptr && offsets_o != nullptr) { - offsets_o[tid] = h * d_v * cu_seqlens_q_padded[cu_seqlens_id]; - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD: - offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_3HD: - case NVTE_QKV_Layout_Group::NVTE_H3D: - offsets_q[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - break; - } + offsets_q[tid] = mults.q * cu_seqlens_q_padded[cu_seqlens_id]; + offsets_o[tid] = mults.o * cu_seqlens_q_padded[cu_seqlens_id]; } if (offsets_k != nullptr && offsets_v != nullptr) { - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD: - offsets_k[tid] = hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id]; - offsets_v[tid] = hg * d_v * cu_seqlens_kv_padded[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_3HD: - case NVTE_QKV_Layout_Group::NVTE_H3D: - offsets_k[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - offsets_v[tid] = offsets_k[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - offsets_k[tid] = 2 * hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id]; - offsets_v[tid] = offsets_k[cu_seqlens_id]; - break; - } + const int32_t *cu_seqlens_kv_src = + mults.kv_from_q ? cu_seqlens_q_padded : cu_seqlens_kv_padded; + offsets_k[tid] = mults.k * cu_seqlens_kv_src[cu_seqlens_id]; + offsets_v[tid] = mults.v * cu_seqlens_kv_src[cu_seqlens_id]; } } } -__global__ void cu_seqlens_padded_to_offsets(NVTE_QKV_Layout_Group layout_group, int64_t actual_b, - int64_t max_b, int64_t h, int64_t hg, int64_t d_qk, - int64_t d_v, const int32_t *cu_seqlens_q_padded, +__global__ void cu_seqlens_padded_to_offsets(RaggedOffsetMultipliers mults, int64_t actual_b, + int64_t max_b, const int32_t *cu_seqlens_q_padded, const int32_t *cu_seqlens_kv_padded, DType offset_dtype, void *offsets_q, void *offsets_k, void *offsets_v, void *offsets_o, void *offsets_s) { if (offset_dtype == DType::kInt32) { cu_seqlens_padded_to_offsets_impl( - layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded, + mults, actual_b, max_b, cu_seqlens_q_padded, cu_seqlens_kv_padded, reinterpret_cast(offsets_q), reinterpret_cast(offsets_k), reinterpret_cast(offsets_v), reinterpret_cast(offsets_o), reinterpret_cast(offsets_s)); } else { assert(offset_dtype == DType::kInt64 && "expect int64"); cu_seqlens_padded_to_offsets_impl( - layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded, + mults, actual_b, max_b, cu_seqlens_q_padded, cu_seqlens_kv_padded, reinterpret_cast(offsets_q), reinterpret_cast(offsets_k), reinterpret_cast(offsets_v), reinterpret_cast(offsets_o), reinterpret_cast(offsets_s)); diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index b2bbd31f68..1864f9417d 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -309,14 +309,45 @@ struct FADescriptor_v1 { } }; +// Per-tensor scale factors relating cu_seqlens_padded (token units) to tensor-element +// ragged offsets, as a function of the QKV layout group. Single source of truth shared +// by the cu_seqlens_padded_to_offsets conversion kernel and the direct-seqlens path +// (which passes them to cuDNN as ragged offset multipliers). +struct RaggedOffsetMultipliers { + RaggedOffsetMultipliers(NVTE_QKV_Layout_Group layout_group, int64_t h, int64_t hg, int64_t d_qk, + int64_t d_v) + : q(h * d_qk), k(hg * d_qk), v(hg * d_v), o(h * d_v), stats(h), kv_from_q(false) { + switch (layout_group) { + case NVTE_QKV_Layout_Group::NVTE_3HD: + case NVTE_QKV_Layout_Group::NVTE_H3D: + q = k = v = 3 * h * d_qk; + kv_from_q = true; + break; + case NVTE_QKV_Layout_Group::NVTE_HD_2HD: + case NVTE_QKV_Layout_Group::NVTE_HD_H2D: + k = v = 2 * hg * d_qk; + break; + default: + break; + } + } + + int64_t q; + int64_t k; + int64_t v; + int64_t o; + int64_t stats; + // K/V offsets scale the Q-side cu_seqlens_padded (interleaved QKV layouts) + bool kv_from_q; +}; + __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, int32_t const *const q_cu_seqlens, int32_t const *const kv_cu_seqlens, int32_t *q_seqlens, int32_t *kv_seqlens); -__global__ void cu_seqlens_padded_to_offsets(NVTE_QKV_Layout_Group layout_group, int64_t actual_b, - int64_t max_b, int64_t h, int64_t hg, int64_t d_qk, - int64_t d_v, const int32_t *cu_seqlens_q_padded, +__global__ void cu_seqlens_padded_to_offsets(RaggedOffsetMultipliers mults, int64_t actual_b, + int64_t max_b, const int32_t *cu_seqlens_q_padded, const int32_t *cu_seqlens_kv_padded, DType offset_dtype, void *offsets_q, void *offsets_k, void *offsets_v, void *offsets_o, void *offsets_s); From bfdf24db59b57ba52a30c49c19da8030394e2e81 Mon Sep 17 00:00:00 2001 From: wdykas <73254672+wdykas@users.noreply.github.com> Date: Wed, 15 Jul 2026 11:41:28 -0400 Subject: [PATCH 21/35] =?UTF-8?q?[PyTorch]=20Add=20per-version=20FlashAtte?= =?UTF-8?q?ntion=20env=20vars=20(NVTE=5FFLASH=5FATTN=5FV2=E2=80=A6=20(#320?= =?UTF-8?q?4)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [PyTorch] Add per-version FlashAttention env vars (NVTE_FLASH_ATTN_V2/V3/V4) NVTE_FLASH_ATTN enables or disables the whole FlashAttention family, but the choice between FlashAttention 2, 3, and 4 is automatic (package presence and compute capability) with no user override. Some workloads need to pin the FlashAttention generation, e.g. RL training that must produce bitwise-identical logprobs to an inference engine running a specific FlashAttention version: different generations use different tile sizes and online-softmax accumulation orders, so mixed versions between training and inference break batch-invariant / train-inference parity guarantees. Add NVTE_FLASH_ATTN_V2, NVTE_FLASH_ATTN_V3, and NVTE_FLASH_ATTN_V4 (default 1) that disable a specific FlashAttention version even when it is installed, following the existing NVTE_FLASH_ATTN filter pattern. Behavior is unchanged when the variables are unset. Signed-off-by: wdykas --- docs/envvars.rst | 18 ++++++++++++++++++ .../attention/dot_product_attention/utils.py | 12 ++++++------ 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index e9c3091c18..b3765a06bd 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -141,6 +141,24 @@ backend-selection overview. :Default: ``1`` :Description: Enable or disable FlashAttention backend for DotProductAttention. When set to ``0``, FlashAttention will not be used. +.. envvar:: NVTE_FLASH_ATTN_V2 + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable or disable FlashAttention 2 (the ``flash-attn`` package) for DotProductAttention, without affecting FlashAttention 3 or 4. When set to ``0``, FlashAttention 2 will not be used even if it is installed. Useful for pinning the FlashAttention version, e.g. so training-side attention runs the same kernel generation as an inference engine. + +.. envvar:: NVTE_FLASH_ATTN_V3 + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable or disable FlashAttention 3 (the ``flash-attn-3`` package) for DotProductAttention, without affecting FlashAttention 2 or 4. When set to ``0``, FlashAttention 3 will not be used even if it is installed. + +.. envvar:: NVTE_FLASH_ATTN_V4 + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable or disable FlashAttention 4 (the ``flash-attn-4`` package) for DotProductAttention, without affecting FlashAttention 2 or 3. When set to ``0``, FlashAttention 4 will not be used even if it is installed. + .. envvar:: NVTE_FUSED_ATTN :Type: ``int`` (0 or 1) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7be94a6fa1..4034351dcf 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -455,18 +455,18 @@ def get_attention_backend( # Filter: Environment variables use_flash_attention = int(os.getenv("NVTE_FLASH_ATTN", "1")) - use_flash_attention_2 = use_flash_attention - use_flash_attention_3 = use_flash_attention - use_flash_attention_4 = use_flash_attention + use_flash_attention_2 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V2", "1")) + use_flash_attention_3 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V3", "1")) + use_flash_attention_4 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) flash_attention_backend = None use_fused_attention = int(os.getenv("NVTE_FUSED_ATTN", "1")) use_unfused_attention = int(os.getenv("NVTE_UNFUSED_ATTN", "1")) if not use_flash_attention_2 and FlashAttentionUtils.is_installed: - logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0") + logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V2=0") if not use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: - logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0") + logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V3=0") if not use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0") + logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V4=0") if not use_fused_attention: logger.debug("Disabling FusedAttention due to NVTE_FUSED_ATTN=0") if not use_unfused_attention: From eabdd463afd011f69815fa8cbc9a5b457f315c3f Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:57:38 -0700 Subject: [PATCH 22/35] Update list of authorized CI users (#3211) Signed-off-by: Tim Moon --- .github/workflows/trigger-ci.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml index de26531a98..68d1d7d71f 100644 --- a/.github/workflows/trigger-ci.yml +++ b/.github/workflows/trigger-ci.yml @@ -60,6 +60,8 @@ jobs: || github.actor == 'kainzhong' || github.actor == 'cspades' || github.actor == 'jomitchellnv' + || github.actor == 'fheinecke' + || github.actor == 'janekb04' ) steps: - name: Check if comment is issued by authorized person From 9d92fa058e71f1e17d4a7d3186d16ca2b0be7266 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Wed, 15 Jul 2026 15:21:10 -0700 Subject: [PATCH 23/35] Fix Cuda Graph based MOE Tests Hang in CI (#3210) * fix grouped linear hang Signed-off-by: Varun Thumbe * make the same change in grouped mlp as well Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe --- tests/pytorch/test_grouped_linear.py | 13 ++- tests/pytorch/test_grouped_mlp.py | 148 +++++++++++++++++++-------- 2 files changed, 115 insertions(+), 46 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index b561506b31..5dff390c9a 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1857,6 +1857,15 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch params_dtype=dtype, device=device, ) + reference_grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=bias, + params_dtype=dtype, + device=device, + ) + reference_grouped_linear.load_state_dict(grouped_linear.state_dict()) static_x = torch.randn(total_tokens, in_features, dtype=dtype, device=device) static_x.requires_grad_(True) @@ -1919,7 +1928,7 @@ def _train_step(x, dy, out_buf, *, use_graphed): expected_x = fresh_x.detach().clone().requires_grad_(True) expected_dy = fresh_dy.detach().clone() with autocast(enabled=use_fp8, recipe=fp8_recipe): - expected_out = grouped_linear(expected_x, static_m_splits) + expected_out = reference_grouped_linear(expected_x, static_m_splits) expected_out.backward(expected_dy) tols = dict(rtol=1e-2, atol=5e-3) @@ -1927,7 +1936,7 @@ def _train_step(x, dy, out_buf, *, use_graphed): tols = dict(rtol=0.05, atol=0.05) torch.testing.assert_close(graph_out.float(), expected_out.float(), **tols) torch.testing.assert_close(graph_dx.float(), expected_x.grad.float(), **tols) - for graph_grad, param in zip(graph_param_grads, grouped_linear.parameters()): + for graph_grad, param in zip(graph_param_grads, reference_grouped_linear.parameters()): assert param.grad is not None torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index e24fff9049..483f477f2f 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -515,7 +515,8 @@ def test_grouped_linear_cuda_graph_safe( split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) # Pad input tokens to validate the sync-free flow - in_shape = (split_sizes.sum().item() + token_padding, in_features) + num_active_tokens = split_sizes.sum().item() + in_shape = (num_active_tokens + token_padding, in_features) out_shape = (in_shape[0], out_features) recipe = make_recipe(quantization) @@ -531,33 +532,45 @@ def test_grouped_linear_cuda_graph_safe( single_grouped_weight=single_grouped_weight, single_grouped_bias=single_grouped_bias, ) + reference_op = te.ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=bias, + device=device, + dtype=dtype, + accumulate_into_main_grad=accumulate_into_main_grad, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + ) + reference_op.load_state_dict(op.state_dict()) - def _weight_params() -> list[torch.nn.Parameter]: + def _weight_params(module: torch.nn.Module) -> list[torch.nn.Parameter]: if single_grouped_weight: - return [op.weight] - return [getattr(op, f"weight{i}") for i in range(group_size)] + return [module.weight] + return [getattr(module, f"weight{i}") for i in range(group_size)] - def _bias_params() -> list[torch.nn.Parameter]: + def _bias_params(module: torch.nn.Module) -> list[torch.nn.Parameter]: if not bias: return [] if single_grouped_bias: - return [op.bias] - return [getattr(op, f"bias{i}") for i in range(group_size)] + return [module.bias] + return [getattr(module, f"bias{i}") for i in range(group_size)] - def _init_main_grads(value: float = 0.0) -> None: + def _init_main_grads(module: torch.nn.Module, value: float = 0.0) -> None: if not accumulate_into_main_grad: return with torch.no_grad(): - for w in _weight_params(): + for w in _weight_params(module): if getattr(w, "main_grad", None) is None: w.main_grad = torch.empty(w.size(), device=device, dtype=torch.float32) w.main_grad.fill_(value) - def _collect_main_grads() -> list[torch.Tensor]: - return [w.main_grad.detach().clone() for w in _weight_params()] + def _collect_main_grads(module: torch.nn.Module) -> list[torch.Tensor]: + return [w.main_grad.detach().clone() for w in _weight_params(module)] - def _zero_param_grads() -> None: - for param in op.parameters(): + def _zero_param_grads(module: torch.nn.Module) -> None: + for param in module.parameters(): if param.grad is None: param.grad = torch.zeros_like(param) else: @@ -582,7 +595,7 @@ def train_step( out_buf.copy_(out) return out_buf - _init_main_grads(0.0) + _init_main_grads(op, 0.0) static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True) static_dy = torch.randn(out_shape, device=device, dtype=dtype) @@ -605,8 +618,8 @@ def train_step( static_dy.copy_(fresh_dy) # Reset grads & main_grads so the captured iteration starts fresh. - _zero_param_grads() - _init_main_grads(0.5) + _zero_param_grads(op) + _init_main_grads(op, 0.5) if static_x.grad is not None: static_x.grad.zero_() @@ -617,35 +630,36 @@ def train_step( torch.cuda.synchronize() graph_dx = static_x.grad.detach().clone() if accumulate_into_main_grad: - graph_main_grads = _collect_main_grads() + graph_main_grads = _collect_main_grads(op) graph_param_grads: list[torch.Tensor] = [] else: graph_main_grads = [] graph_param_grads = [param.grad.detach().clone() for param in op.parameters()] - # Reference: same op invoked eagerly with the same fresh inputs and - # the same starting grad/main_grad state. - _zero_param_grads() - _init_main_grads(0.5) - static_x.grad.zero_() + # Reference: an independent op invoked eagerly with the same fresh + # inputs and starting grad/main_grad state. + _zero_param_grads(reference_op) + _init_main_grads(reference_op, 0.5) expected_x = fresh_x.detach().clone().requires_grad_(True) expected_dy = fresh_dy.detach().clone() with te.autocast(enabled=quantization is not None, recipe=recipe): - expected_out = op(expected_x, static_split_sizes) + expected_out = reference_op(expected_x, static_split_sizes) expected_out.backward(expected_dy) tols = dtype_tols(dtype) if quantization is not None: tols = quantization_tols(quantization) - assert_close(graph_out, expected_out, **tols) - assert_close(graph_dx, expected_x.grad, **tols) + # Grouped GEMM only defines rows covered by split_sizes. The padded tail + # is intentionally outside every group and remains uninitialized. + assert_close(graph_out[:num_active_tokens], expected_out[:num_active_tokens], **tols) + assert_close(graph_dx[:num_active_tokens], expected_x.grad[:num_active_tokens], **tols) if accumulate_into_main_grad: - for g, w in zip(graph_main_grads, _weight_params()): + for g, w in zip(graph_main_grads, _weight_params(reference_op)): assert_close(g, w.main_grad, **tols) else: - for g, param in zip(graph_param_grads, op.parameters()): + for g, param in zip(graph_param_grads, reference_op.parameters()): assert_close(g, param.grad, **tols) @@ -1491,7 +1505,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( random.shuffle(split_sizes) split_sizes = torch.tensor(split_sizes, dtype=torch.int64, device=device) # Pad the input tokens to validate the sync-free MOE - in_shape = (split_sizes.sum().item() + token_padding, hidden_size) + num_active_tokens = split_sizes.sum().item() + in_shape = (num_active_tokens + token_padding, hidden_size) recipe = make_recipe("mxfp8") with te.quantized_model_init(enabled=True, recipe=recipe): fc1 = te.ops.GroupedLinear( @@ -1524,8 +1539,43 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( scaled_act, fc2, ) + reference_fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + ) + reference_fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + ) + reference_scaled_act = ( + te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + if activation == "scaled_swiglu" + else te.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size) + ) + reference_module = te.ops.Sequential( + reference_fc1, + reference_scaled_act, + reference_fc2, + ) + reference_module.load_state_dict(module.state_dict()) - def _init_main_grads(value: float = 0.0) -> None: + def _init_main_grads( + fc1: torch.nn.Module, + fc2: torch.nn.Module, + value: float = 0.0, + ) -> None: if not accumulate_into_main_grad: return with torch.no_grad(): @@ -1563,7 +1613,10 @@ def _init_main_grads(value: float = 0.0) -> None: fc1_weight.main_grad.fill_(value) fc2_weight.main_grad.fill_(value) - def _collect_main_grads() -> tuple[torch.Tensor, torch.Tensor]: + def _collect_main_grads( + fc1: torch.nn.Module, + fc2: torch.nn.Module, + ) -> tuple[torch.Tensor, torch.Tensor]: if single_grouped_weight: fc1_main_grad = fc1.weight.main_grad.detach().clone() fc2_main_grad = fc2.weight.main_grad.detach().clone() @@ -1604,7 +1657,7 @@ def train_step( out_buf.copy_(out) return out_buf - _init_main_grads(0.0) + _init_main_grads(fc1, fc2, 0.0) static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True) static_probs = torch.randn((in_shape[0],), device=device, dtype=dtype, requires_grad=True) @@ -1643,7 +1696,7 @@ def train_step( for param in module.parameters(): param.grad = torch.zeros_like(param) - _init_main_grads(0.5) + _init_main_grads(fc1, fc2, 0.5) if static_x.grad is not None: static_x.grad.zero_() if static_probs.grad is not None: @@ -1658,21 +1711,19 @@ def train_step( graph_dx = static_x.grad.detach().clone() graph_dprobs = static_probs.grad.detach().clone() if accumulate_into_main_grad: - graph_fc1_main_grad, graph_fc2_main_grad = _collect_main_grads() + graph_fc1_main_grad, graph_fc2_main_grad = _collect_main_grads(fc1, fc2) else: graph_param_grads = [param.grad.detach().clone() for param in module.parameters()] - for param in module.parameters(): - param.grad.zero_() - _init_main_grads(0.5) - static_x.grad.zero_() - static_probs.grad.zero_() + for param in reference_module.parameters(): + param.grad = torch.zeros_like(param) + _init_main_grads(reference_fc1, reference_fc2, 0.5) expected_x = fresh_x.detach().clone().requires_grad_(True) expected_probs = fresh_probs.detach().clone().requires_grad_(True) expected_dy = fresh_dy.detach().clone() with te.autocast(enabled=True, recipe=recipe): - expected_out = module( + expected_out = reference_module( expected_x, static_split_sizes, expected_probs, @@ -1681,15 +1732,24 @@ def train_step( expected_out.backward(expected_dy) tols = dtype_tols(dtype) - assert_close(graph_out, expected_out, **tols) - assert_close(graph_dx, expected_x.grad, **tols) - assert_close(graph_dprobs, expected_probs.grad, **tols) + # The padded tail is outside every expert and its outputs/gradients are + # intentionally left uninitialized. + assert_close(graph_out[:num_active_tokens], expected_out[:num_active_tokens], **tols) + assert_close(graph_dx[:num_active_tokens], expected_x.grad[:num_active_tokens], **tols) + assert_close( + graph_dprobs[:num_active_tokens], + expected_probs.grad[:num_active_tokens], + **tols, + ) if accumulate_into_main_grad: - expected_fc1_main_grad, expected_fc2_main_grad = _collect_main_grads() + expected_fc1_main_grad, expected_fc2_main_grad = _collect_main_grads( + reference_fc1, + reference_fc2, + ) assert_close(graph_fc1_main_grad, expected_fc1_main_grad, **tols) assert_close(graph_fc2_main_grad, expected_fc2_main_grad, **tols) else: - for graph_grad, param in zip(graph_param_grads, module.parameters()): + for graph_grad, param in zip(graph_param_grads, reference_module.parameters()): assert_close(graph_grad, param.grad, **tols) From 68493d2d55ac37e540301467b278bdb1c2019e81 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 17 Jul 2026 04:53:19 +0200 Subject: [PATCH 24/35] [PyTorch] Add optional caller-provided output/grad-input buffers to GroupedLinear and fused grouped MLP (#3161) * Add optional caller-provided output/grad-input buffers to GroupedLinear module and fusible ops Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route per-op kwargs through Sequential via module-keyed op_kwargs mapping Signed-off-by: Phuong Nguyen * Write fused grouped MLP MXFP8 output and dgrad directly into caller buffers, eliminating the D2D copy + cleanup Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use 256-aligned splits in caller-buffer grouped MLP test Signed-off-by: Phuong Nguyen * use basic_ops to track op kwargs Signed-off-by: YangFei1990 * add doc and resolve comments Signed-off-by: YangFei1990 * move out/dgrad_out out from the non_tensor_args Signed-off-by: YangFei1990 --------- Signed-off-by: Phuong Nguyen Signed-off-by: YangFei1990 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: YangFei1990 Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com> --- tests/pytorch/test_fusible_ops.py | 84 +++++++++++++ tests/pytorch/test_grouped_linear.py | 96 +++++++++++++++ tests/pytorch/test_grouped_mlp.py | 101 ++++++++++++++++ .../pytorch/module/grouped_linear.py | 111 +++++++++++++++--- transformer_engine/pytorch/ops/_common.py | 30 +++++ .../pytorch/ops/basic/grouped_linear.py | 30 +++-- .../pytorch/ops/fused/grouped_mlp.py | 54 +++++++-- transformer_engine/pytorch/ops/sequential.py | 57 ++++++++- 8 files changed, 525 insertions(+), 38 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3ad15b6a05..5ae657c5d2 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -19,6 +19,10 @@ import transformer_engine.common.recipe import transformer_engine.pytorch as te import transformer_engine.pytorch.ops as te_ops +from transformer_engine.pytorch.ops.basic.grouped_linear import ( + OUTPUT_BUFFER_KEY, + GRAD_INPUT_BUFFER_KEY, +) from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV from transformer_engine.pytorch.ops.fused import ( @@ -2196,6 +2200,86 @@ def test_grouped_linear( if bias: assert_close_grads(getattr(op, f"bias{group_idx}"), bs_ref[group_idx], **tols) + def test_grouped_linear_caller_buffers( + self, + *, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + """Caller output/grad_input buffers routed to the last/first op of a Sequential. + + The chain has two GroupedLinears with distinct inner dims, so ``output`` + can only fit the last op's output and ``grad_input`` only the first op's + dgrad -- a mis-route would fail the buffer's shape check. + """ + group_size = 3 + in_features, hidden, out_features = 128, 256, 64 + split_sizes = torch.tensor([128, 256, 128], dtype=torch.int32, device=device) + num_tokens = int(split_sizes.sum()) + + torch.manual_seed(1234) + x = (0.1 * torch.randn(num_tokens, in_features, device=device)).to(dtype) + dy = (0.1 * torch.randn(num_tokens, out_features, device=device)).to(dtype) + + def build() -> te_ops.Sequential: + fc1 = te_ops.GroupedLinear( + group_size, in_features, hidden, bias=False, device=device, dtype=dtype + ) + fc2 = te_ops.GroupedLinear( + group_size, hidden, out_features, bias=False, device=device, dtype=dtype + ) + return te_ops.Sequential(fc1, fc2) + + # Reference: internal allocation. + model_ref = build() + x_ref = x.detach().clone().requires_grad_(True) + y_ref = model_ref(x_ref, split_sizes, split_sizes) + y_ref.backward(dy) + + # Caller-provided buffers, same weights as the reference. + model = build() + with torch.no_grad(): + for op_idx in range(2): + for i in range(group_size): + getattr(model[op_idx], f"weight{i}").copy_( + getattr(model_ref[op_idx], f"weight{i}") + ) + sentinel = 7.0 + out_buf = torch.full((num_tokens, out_features), sentinel, dtype=dtype, device=device) + dgrad_buf = torch.full((num_tokens, in_features), sentinel, dtype=dtype, device=device) + x_test = x.detach().clone().requires_grad_(True) + y = model( + x_test, + split_sizes, + split_sizes, + op_kwargs={ + model[0]: {GRAD_INPUT_BUFFER_KEY: dgrad_buf}, + model[1]: {OUTPUT_BUFFER_KEY: out_buf}, + }, + ) + + # Forward output aliases the last op's output buffer with no copy. + assert y.data_ptr() == out_buf.data_ptr() + torch.testing.assert_close(y, y_ref, rtol=0, atol=0) + + y.backward(dy) + + # grad_input written into the first op's dgrad buffer. + assert not torch.all(dgrad_buf == sentinel) + torch.testing.assert_close(dgrad_buf, x_ref.grad, rtol=0, atol=0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0, atol=0) + + # A buffer whose shape does not match the output is rejected. + bad = torch.empty(num_tokens + 1, out_features, dtype=dtype, device=device) + with pytest.raises(ValueError): + model_bad = build() + model_bad( + x.detach(), + split_sizes, + split_sizes, + op_kwargs={model_bad[1]: {OUTPUT_BUFFER_KEY: bad}}, + ) + @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("scales_requires_grad", (False, True)) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 5dff390c9a..64951a43b8 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1710,6 +1710,102 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() +@pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"]) +@pytest.mark.parametrize("supply", ["out", "dgrad_out", "both"]) +def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatch): + """Caller-provided forward out and/or backward dgrad_out buffers. + + Checks that a supplied buffer is written in place (bit-for-bit vs internal allocation) + and, on the fused path with a padded input, that only the valid rows are touched. + """ + if use_fused_path: + device_capability = torch.cuda.get_device_capability() + if not (9, 0) <= device_capability <= (11, 0): + pytest.skip( + "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell" + " (SM10x and SM110)." + ) + cublaslt_version = tex.get_cublasLt_version() + if device_capability < (10, 0) and cublaslt_version < 130400: + pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") + if cublaslt_version < 130300: + pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if use_fused_path else "0") + give_out = supply in ("out", "both") + give_dgrad = supply in ("dgrad_out", "both") + + dtype = torch.bfloat16 + num_gemms = 3 + in_features = 128 + out_features = 128 + m_splits_list = [64, 96, 80] + valid_tokens = sum(m_splits_list) # 240 + # The fused path supports a padded input; the legacy path requires tight packing. + num_rows = valid_tokens + (80 if use_fused_path else 0) + m_splits = ( + torch.tensor(m_splits_list, dtype=torch.int64, device="cuda") + if use_fused_path + else m_splits_list + ) + + torch.manual_seed(1234) + x_base = (0.1 * torch.randn(num_rows, in_features, device="cuda")).to(dtype) + dy = torch.zeros(num_rows, out_features, dtype=dtype, device="cuda") + dy[:valid_tokens] = (0.1 * torch.randn(valid_tokens, out_features, device="cuda")).to(dtype) + + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=dtype, + device="cuda", + ) + + # Reference: internal allocation. + x_ref = x_base.detach().clone().requires_grad_(True) + y_ref = grouped_linear(x_ref, m_splits) + y_ref.backward(dy) + + # Caller-provided buffers with a sentinel-filled tail (only the requested ones). + sentinel = 7.0 + out_buf = ( + torch.full((num_rows, out_features), sentinel, dtype=dtype, device="cuda") + if give_out + else None + ) + dgrad_buf = ( + torch.full((num_rows, in_features), sentinel, dtype=dtype, device="cuda") + if give_dgrad + else None + ) + x = x_base.detach().clone().requires_grad_(True) + y = grouped_linear(x, m_splits, out=out_buf, dgrad_out=dgrad_buf) + + if give_out: + # Forward output is the caller buffer itself (no copy); padded tail untouched. + assert y.data_ptr() == out_buf.data_ptr() + assert tuple(y.shape) == (num_rows, out_features) + assert torch.all(out_buf[valid_tokens:] == sentinel) + torch.testing.assert_close(y[:valid_tokens], y_ref[:valid_tokens], rtol=0, atol=0) + + y.backward(dy) + + if give_dgrad: + # dgrad written into the caller buffer; padded tail untouched. + assert torch.all(dgrad_buf[valid_tokens:] == sentinel) + torch.testing.assert_close( + dgrad_buf[:valid_tokens], x_ref.grad[:valid_tokens], rtol=0, atol=0 + ) + torch.testing.assert_close(x.grad[:valid_tokens], x_ref.grad[:valid_tokens], rtol=0, atol=0) + + # A buffer whose row count does not match the input rows is rejected. + bad_out = torch.empty(num_rows + 1, out_features, dtype=dtype, device="cuda") + with pytest.raises(ValueError): + grouped_linear(x, m_splits, out=bad_out) + + @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct. diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 483f477f2f..76550db5f8 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -19,6 +19,10 @@ _cudnn_frontend_supports_grouped_gemm_srelu, _cudnn_frontend_version_supported, ) +from transformer_engine.pytorch.ops.basic.grouped_linear import ( + OUTPUT_BUFFER_KEY, + GRAD_INPUT_BUFFER_KEY, +) from transformer_engine.pytorch import ( QuantizedTensor, Float8CurrentScalingQuantizer, @@ -1331,6 +1335,103 @@ def _run_case(single_grouped_weight: bool) -> tuple[torch.Tensor, ...]: torch.testing.assert_close(fc1_db_false, fc1_db_true, **bias_tols) torch.testing.assert_close(fc2_db_false, fc2_db_true, **bias_tols) + @pytest.mark.parametrize("quantization", ("mxfp8", "nvfp4_rht")) + def test_grouped_mlp_caller_buffers( + self, + quantization: str, + *, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + ) -> None: + """Caller-provided output/grad_input buffers on the fused MXFP8/NVFP4 grouped MLP.""" + if quantization == "mxfp8" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if quantization == "nvfp4_rht" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP is not supported on this system") + + split_sizes = torch.tensor( + [split_alignment * (i + 1) for i in range(group_size)], + dtype=torch.int64, + device=device, + ) + num_tokens = int(split_sizes.sum()) + recipe = make_recipe(quantization) + + x_base = torch.empty((num_tokens, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + probs_base = torch.empty((num_tokens,), device=device, dtype=dtype).uniform_(-0.25, 0.25) + dy_base = torch.empty((num_tokens, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + fc1_ws_base = [ + torch.empty((2 * hidden_size, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + for _ in range(group_size) + ] + fc2_ws_base = [ + torch.empty((hidden_size, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + for _ in range(group_size) + ] + + def build() -> te.ops.Sequential: + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype + ) + scaled_act = te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + fc2 = te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + module = te.ops.Sequential(fc1, scaled_act, fc2) + with torch.no_grad(): + for i in range(group_size): + getattr(fc1, f"weight{i}").copy_(fc1_ws_base[i]) + getattr(fc2, f"weight{i}").copy_(fc2_ws_base[i]) + return module + + def run(module, *, output=None, grad_input=None): + x = x_base.detach().clone().requires_grad_(True) + probs = probs_base.detach().clone().requires_grad_(True) + op_kwargs = {} + if grad_input is not None: + op_kwargs[module[0]] = {GRAD_INPUT_BUFFER_KEY: grad_input} + if output is not None: + op_kwargs[module[2]] = {OUTPUT_BUFFER_KEY: output} + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes, op_kwargs=op_kwargs or None) + y.backward(dy_base) + return y, x.grad + + # Reference: internal allocation. + y_ref, dx_ref = run(build()) + + # Caller-provided buffers. + module = build() + sentinel = 7.0 + out_buf = torch.full((num_tokens, hidden_size), sentinel, dtype=dtype, device=device) + dgrad_buf = torch.full((num_tokens, hidden_size), sentinel, dtype=dtype, device=device) + y, _ = run(module, output=out_buf, grad_input=dgrad_buf) + + # The fused op ran, the returned output aliases the buffer, and both buffers were written. + assert isinstance( + module._module_groups[0]._forward_ops[0][0], + te.ops.fused.GroupedMLP_CuTeGEMMGLU, + ) + assert y.data_ptr() == out_buf.data_ptr() + assert not torch.all(dgrad_buf == sentinel) + torch.testing.assert_close(y, y_ref, rtol=0, atol=0) + torch.testing.assert_close(dgrad_buf, dx_ref, rtol=0, atol=0) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) @pytest.mark.parametrize("zero_out_wgrad", (False, True)) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 82ee7953c1..a65ee3b5c3 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -236,6 +236,41 @@ def _prepare_weights_for_grouped_tensor_gemm( weights_for_gemm.append(weight_fp8) return weights_for_gemm, new_workspaces + @staticmethod + def _validate_or_alloc_output( + buffer: Optional[torch.Tensor], + rows: int, + cols: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Validate and return the caller's output buffer, or allocate one if it is None. + + The buffer must be a 2D, contiguous, non-grad tensor matching the required shape, + dtype, and device. Validation reads host-side metadata only, with no device sync. + """ + if buffer is None: + return torch.empty((rows, cols), dtype=dtype, device=device) + if buffer.dim() != 2: + raise ValueError(f"Output buffer must be 2D, got {buffer.dim()}D.") + if buffer.size(0) != rows: + raise ValueError(f"Output buffer rows {buffer.size(0)} must match input rows {rows}.") + if buffer.size(1) != cols: + raise ValueError( + f"Output buffer last dim {buffer.size(1)} does not match required {cols}." + ) + if buffer.dtype != dtype: + raise ValueError(f"Output buffer dtype {buffer.dtype} does not match required {dtype}.") + if buffer.device != device: + raise ValueError( + f"Output buffer device {buffer.device} does not match required {device}." + ) + if not buffer.is_contiguous(): + raise ValueError("Output buffer must be contiguous.") + if buffer.requires_grad: + raise ValueError("Output buffer must not require gradient.") + return buffer + @staticmethod def _forward_grouped_tensor( ctx, @@ -259,6 +294,8 @@ def _forward_grouped_tensor( skip_fp8_weight_update: Optional[torch.Tensor], weights: Tuple[torch.Tensor, ...], biases: Tuple[torch.Tensor, ...], + out: Optional[torch.Tensor] = None, + dgrad_out: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, list]: """Forward path backed by GroupedTensor + cuBLASLt grouped GEMM.""" num_gemms = len(m_splits) @@ -303,10 +340,12 @@ def _forward_grouped_tensor( cache_weight=cache_weight, ) - out = torch.empty( - [x.size(0), out_features], - dtype=activation_dtype, - device=device, + out = _GroupedLinear._validate_or_alloc_output( + out, + x.size(0), + out_features, + activation_dtype, + device, ) grouped_out = _GroupedLinear._make_grouped_tensor( out, @@ -381,6 +420,7 @@ def _forward_grouped_tensor( lambda j=i: weights[j].main_grad for i in range(num_gemms) ] ctx.device = device + ctx.dgrad_out = dgrad_out ctx.m_splits = None ctx.num_gemms = num_gemms ctx.activation_dtype = activation_dtype @@ -413,6 +453,8 @@ def forward( inp: torch.Tensor, m_splits: torch.Tensor, non_tensor_args: Tuple, + out: Optional[torch.Tensor], + dgrad_out: Optional[torch.Tensor], *weights_and_biases, ) -> Tuple[torch.Tensor, list]: # pylint: disable=missing-function-docstring @@ -544,6 +586,8 @@ def forward( skip_fp8_weight_update=skip_fp8_weight_update, weights=weights, biases=biases, + out=out, + dgrad_out=dgrad_out, ) # Convert splits to list of ints for compatibility with split functions @@ -598,10 +642,12 @@ def forward( bias_dtype = torch.bfloat16 # FP8 GEMM only supports BF16/FP16 bias biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases # Initialize output tensor - out = torch.empty( - [sum(m_splits), weights_fp8[0].size(0)], - dtype=activation_dtype, - device=device, + out = _GroupedLinear._validate_or_alloc_output( + out, + sum(m_splits), + weights_fp8[0].size(0), + activation_dtype, + device, ) # Choose whether to use split accumulator @@ -719,6 +765,7 @@ def forward( ctx.debug = debug ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers + ctx.dgrad_out = dgrad_out # backward overrides if backward_override is not None: @@ -813,10 +860,12 @@ def _backward_grouped_tensor( for weight in weights: if isinstance(weight, QuantizedTensorStorage): weight.update_usage(columnwise_usage=True) - dgrad = torch.empty( - (dy_2d.size(0), ctx.weights_shape_1), - dtype=ctx.activation_dtype, - device=ctx.device, + dgrad = _GroupedLinear._validate_or_alloc_output( + ctx.dgrad_out, + dy_2d.size(0), + ctx.weights_shape_1, + ctx.activation_dtype, + ctx.device, ) grouped_dgrad = _GroupedLinear._make_grouped_tensor( dgrad, @@ -918,6 +967,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits None, # non_tensor_args + None, # out + None, # dgrad_out *wgrad_list, *grad_biases, ) @@ -1019,10 +1070,12 @@ def backward( dgrad_gemm_use_split_accumulator = ( recipe.fp8_gemm_dgrad.use_split_accumulator ) - dgrad = torch.empty( - (sum(ctx.m_splits), ctx.weights_shape_1), - dtype=ctx.activation_dtype, - device=ctx.device, + dgrad = _GroupedLinear._validate_or_alloc_output( + ctx.dgrad_out, + sum(ctx.m_splits), + ctx.weights_shape_1, + ctx.activation_dtype, + ctx.device, ) weights_for_dgrad = weights if ctx.backward_override == "dequantized": @@ -1191,6 +1244,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits None, # non_tensor_args + None, # out + None, # dgrad_out *wgrad_list, *grad_biases, ) @@ -1684,6 +1739,8 @@ def forward( inp: torch.Tensor, m_splits: torch.Tensor, is_first_microbatch: Optional[bool] = None, + out: Optional[torch.Tensor] = None, + dgrad_out: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """ Apply the linear transformation to the input. @@ -1707,6 +1764,19 @@ def forward( * it also allows skipping gradient accumulation during the first microbatch (since it is the first gradient being produced) + out : torch.Tensor, default = None + Optional preallocated buffer for the forward output; the returned tensor + aliases it with no copy. Must be a 2D, contiguous, non-grad tensor of shape + [num_tokens, out_features] in the activation dtype. Only the first + sum(m_splits) rows are written; any padded trailing rows are left unchanged. + Can be given independently of dgrad_out. If the buffer is reused across + iterations, pass ``buffer.detach()`` so autograd does not set its + ``requires_grad`` (which would trip the non-grad check on the next call). + dgrad_out : torch.Tensor, default = None + Optional preallocated buffer for the backward input gradient, of shape + [num_tokens, in_features] with the same constraints as out. Receives the + final gradient only when inp has a single consumer in the autograd graph; + otherwise autograd accumulates into a new tensor. """ debug = self.is_debug_iter() is_grad_enabled = torch.is_grad_enabled() @@ -1805,7 +1875,14 @@ def forward( debug, ) out, new_workspaces = linear_fn( - *autograd_ctx, inp, m_splits, non_tensor_args, *weight_tensors, *bias_tensors + *autograd_ctx, + inp, + m_splits, + non_tensor_args, + out, + dgrad_out, + *weight_tensors, + *bias_tensors, ) if cache_weight: diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 675a102240..607346ce30 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -18,6 +18,36 @@ from ..utils import canonicalize_dtype +def validate_or_alloc_output( + buffer: Optional[torch.Tensor], + shape: tuple[int, ...] | list[int], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Return the caller's output buffer, or allocate one if it is None. + + The buffer must be a contiguous, non-grad tensor matching the required + shape, dtype, and device. Validation reads host-side metadata only. If the + buffer is reused across iterations, pass ``buffer.detach()`` so autograd does + not set its ``requires_grad`` (which would trip the non-grad check here on the + next call). + """ + shape = tuple(shape) + if buffer is None: + return torch.empty(shape, dtype=dtype, device=device) + if tuple(buffer.shape) != shape: + raise ValueError(f"Output buffer shape {tuple(buffer.shape)} does not match {shape}.") + if buffer.dtype != dtype: + raise ValueError(f"Output buffer dtype {buffer.dtype} does not match {dtype}.") + if buffer.device != device: + raise ValueError(f"Output buffer device {buffer.device} does not match {device}.") + if not buffer.is_contiguous(): + raise ValueError("Output buffer must be contiguous.") + if buffer.requires_grad: + raise ValueError("Output buffer must not require gradient.") + return buffer + + def is_quantized_tensor(tensor: torch.Tensor | QuantizedTensorStorage) -> bool: """Check if tensor is a quantized tensor""" return isinstance(tensor, QuantizedTensorStorage) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5c96d4658e..5ef0fa4339 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -48,6 +48,7 @@ get_main_grad_from_param, is_quantized_tensor, maybe_dequantize, + validate_or_alloc_output, view_main_grad_as_grouped_buffer, ) from ..op import BasicOperation, OperationContext @@ -58,6 +59,12 @@ ) +# Keys for passing caller-provided output and grad-input buffers to a grouped +# linear (or fused grouped MLP) through Sequential's ``op_kwargs``. +OUTPUT_BUFFER_KEY = "output" +GRAD_INPUT_BUFFER_KEY = "grad_input" + + class GroupedLinear(BasicOperation): r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` @@ -986,6 +993,9 @@ def fuser_forward( if self._scale_bias: scales = basic_op_extra_inputs[0][1] + # Caller-provided output buffer (backward grad-input buffer is read in save_ctx). + out_buffer = basic_op_kwargs[0].get(OUTPUT_BUFFER_KEY) + # Dispatch: graph-safe GroupedTensor flow whenever it can be used. # See ``_is_graph_safe_path_supported`` for the gating rationale -- # in short it requires Hopper (SM90+) plus a supported dtype / @@ -1010,6 +1020,7 @@ def fuser_forward( input_requires_grad=input_requires_grad, weight_requires_grad=weight_requires_grad, device=device, + out_buffer=out_buffer, ) else: out, tensors_to_save = self._fuser_forward_split_quantize( @@ -1023,6 +1034,7 @@ def fuser_forward( input_requires_grad=input_requires_grad, weight_requires_grad=weight_requires_grad, device=device, + out_buffer=out_buffer, ) # Save tensors and autograd metadata on the basic-op context. @@ -1111,6 +1123,8 @@ def fuser_forward_save_ctx( ctx.dtype = weight_param.dtype ctx.input_requires_grad = requires_grad[0] ctx.weight_requires_grad = requires_grad[0] and weight_param.requires_grad + # Caller-provided backward grad-input buffer. + ctx.dgrad_out = basic_op_kwargs[0].get(GRAD_INPUT_BUFFER_KEY) # ================================================================== # Legacy `tex.split_quantize` + `general_grouped_gemm` flow. @@ -1129,6 +1143,7 @@ def _fuser_forward_split_quantize( input_requires_grad: bool, weight_requires_grad: bool, device: torch.device, + out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow.""" num_groups = self.num_groups @@ -1173,7 +1188,7 @@ def _fuser_forward_split_quantize( # Allocate output tensor in_shape = list(input_.size()) out_shape = in_shape[:-1] + [self.out_features] - out = torch.empty(out_shape, dtype=dtype, device=device) + out = validate_or_alloc_output(out_buffer, out_shape, dtype, device) # Perform GEMMs use_gemm_bias = has_bias and not self._scale_bias @@ -1239,6 +1254,7 @@ def _fuser_forward_grouped_tensor( input_requires_grad: bool, weight_requires_grad: bool, device: torch.device, + out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: """Graph-safe GroupedTensor forward path (pure compute). Returns ``(output, tensors_to_save)``. ``split_sizes``, @@ -1299,7 +1315,7 @@ def _fuser_forward_grouped_tensor( # Allocate output buffer and wrap as a GroupedTensor view. out_shape = original_shape[:-1] + [self.out_features] - out = torch.empty(out_shape, dtype=dtype, device=device) + out = validate_or_alloc_output(out_buffer, out_shape, dtype, device) grouped_out = GroupedTensorStorage( shape=(total_tokens, self.out_features), dtype=dtype, @@ -1489,10 +1505,8 @@ def _fuser_backward_split_quantize( if ctx.input_requires_grad: out_shape = list(grad_output.size()) in_shape = out_shape[:-1] + [self.in_features] - grad_input = torch.empty( - in_shape, - dtype=ctx.dtype, - device=device, + grad_input = validate_or_alloc_output( + getattr(ctx, "dgrad_out", None), in_shape, ctx.dtype, device ) general_grouped_gemm( ws, @@ -1668,7 +1682,9 @@ def _fuser_backward_grouped_tensor( grad_input = None if ctx.input_requires_grad: grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] - grad_input = torch.empty(grad_input_shape, dtype=dtype, device=device) + grad_input = validate_or_alloc_output( + getattr(ctx, "dgrad_out", None), grad_input_shape, dtype, device + ) grouped_grad_input = GroupedTensorStorage( shape=(total_tokens, self.in_features), dtype=dtype, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index e0961044ad..31189af09c 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -48,8 +48,10 @@ get_main_grad_from_param, is_quantized_tensor, maybe_dequantize, + validate_or_alloc_output, view_main_grad_as_grouped_buffer, ) +from ..basic.grouped_linear import GRAD_INPUT_BUFFER_KEY, OUTPUT_BUFFER_KEY @functools.lru_cache(maxsize=None) @@ -768,7 +770,9 @@ def fuse_grouped_mlp_ops( class _GroupedMLP_CuTeGEMMBase(FusedOperation): """Joint fused op for block-scaled GroupedLinear + activation + GroupedLinear. - Uses experimental CuTe DSL kernels from cuDNN front-end. + MXFP8 uses CuTe DSL grouped-GEMM kernels via the cuDNN front-end. When + caller output/grad_input buffers are provided, the GEMMs write into them + directly. """ @@ -879,6 +883,21 @@ def fuser_forward( fc1_op, activation_op, fc2_op = self.basic_ops fc1_ctx, _activation_ctx, fc2_ctx = basic_op_ctxs + # Caller-provided buffers: FC2 forward output and FC1 backward grad-input. Guard against + # a buffer routed to the wrong op (output belongs on FC2, grad-input on FC1). + if OUTPUT_BUFFER_KEY in basic_op_kwargs[0]: + raise ValueError( + f"'{OUTPUT_BUFFER_KEY}' buffer can only be provided to FC2 (the last op) of " + "the fused grouped MLP, not FC1." + ) + if GRAD_INPUT_BUFFER_KEY in basic_op_kwargs[-1]: + raise ValueError( + f"'{GRAD_INPUT_BUFFER_KEY}' buffer can only be provided to FC1 (the first op) of " + "the fused grouped MLP, not FC2." + ) + output_buffer = basic_op_kwargs[-1].get(OUTPUT_BUFFER_KEY) + fc1_ctx.dgrad_out = basic_op_kwargs[0].get(GRAD_INPUT_BUFFER_KEY) + # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) @@ -1301,7 +1320,7 @@ def fuser_forward( tensor_offsets=fc2_x_tensor_offsets, ) - fc2_out_buf = torch.empty(fc2_out_shape, dtype=dtype, device=device) + fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) if ( num_groups == 1 and grouped_fc2_x.columnwise_data is not None @@ -1329,9 +1348,9 @@ def fuser_forward( fc2_bias_packed.transpose(0, 1).contiguous().expand(in_shape[0], -1) ) if fc2_scales is not None: - fc2_out_buf = fc2_out_buf + token_bias * fc2_scales.view(-1, 1) + fc2_out_buf += token_bias * fc2_scales.view(-1, 1) else: - fc2_out_buf = fc2_out_buf + token_bias + fc2_out_buf += token_bias else: fc2_out_grouped = GroupedTensorStorage( shape=(in_shape[0], fc2_weight_shape[0]), @@ -1434,8 +1453,15 @@ def fuser_forward( fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn fc2_quant_kwargs["b_major"] = "k" - fc2_kernel_out = self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) - fc2_out = fc2_kernel_out["d_tensor"].permute(2, 0, 1).view(fc2_out_shape).contiguous() + # Always allocate the output (the caller's buffer if provided, else a fresh one) and + # pass it as the kernel's d_tensor, so the kernel writes in place and the call is uniform. + output_buffer = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) + fc2_quant_kwargs["d_tensor"] = output_buffer.as_strided( + (in_shape[0], fc2_weight_shape[0], 1), + (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), + ) + self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) + fc2_out = output_buffer # Save state for backward pass if requires_grad: @@ -1962,11 +1988,12 @@ def fuser_backward( # FC1 dgrad GEMM grad_input = None + grad_input_buffer = getattr(fc1_ctx, "dgrad_out", None) if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] if use_nvfp4: - grad_input = torch.empty(in_shape, dtype=dtype, device=device) + grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) if num_groups == 1: if fc1_op.single_grouped_weight: fc1_w_single = grouped_fc1_weight.split_into_quantized_tensors()[0] @@ -2063,8 +2090,17 @@ def fuser_backward( fc1_dgrad_kwargs["b_dtype"] = torch.float8_e4m3fn fc1_dgrad_kwargs["b_major"] = "n" - fc1_dgrad_kernel_out = self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) - grad_input = fc1_dgrad_kernel_out["d_tensor"].view(in_shape) + # Always allocate the grad-input (the caller's buffer if provided, else a fresh + # one) and pass it as the kernel's d_tensor, decoupling it from the kernel call. + grad_input_buffer = validate_or_alloc_output( + grad_input_buffer, in_shape, dtype, device + ) + fc1_dgrad_kwargs["d_tensor"] = grad_input_buffer.as_strided( + (out_shape[0], fc1_weight_shape[1], 1), + (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), + ) + self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) + grad_input = grad_input_buffer # FC1 wgrad GEMM fc1_grad_params = _compute_grad_params( diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index 592ddae23a..cb5dfecb9f 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -6,11 +6,11 @@ from __future__ import annotations from collections.abc import Iterable, Iterator -from typing import Optional +from typing import Any, Optional import torch -from transformer_engine.pytorch.ops.op import FusibleOperation +from transformer_engine.pytorch.ops.op import BasicOperation, FusibleOperation from transformer_engine.pytorch.ops.fuser import OperationFuser @@ -168,23 +168,34 @@ def forward( self, input: torch.Tensor, # pylint: disable=redefined-builtin *extra_inputs: torch.Tensor, + op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]] = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: - """Forward pass""" + """Forward pass. + + op_kwargs : optional mapping from a contained op, keyed by module + or index, to extra keyword arguments forwarded to that op. + Only fusible operations can be targeted, for example to pass + preallocated output or grad input buffers to a grouped linear + or grouped MLP. + """ # Create module groups if needed if self._module_groups is None: self._module_groups = self._make_module_groups(self._modules.values()) + # Route op kwargs to each module group's basic ops + group_op_kwargs = self._resolve_op_kwargs(op_kwargs) + # Forward pass for each module group x = input extra_outputs: list[torch.Tensor] = [] - for module_group in self._module_groups: + for group_idx, module_group in enumerate(self._module_groups): if isinstance(module_group, OperationFuser): xs, extra_inputs = ( (x,) + extra_inputs[: module_group.num_extra_inputs], extra_inputs[module_group.num_extra_inputs :], ) - xs = module_group(*xs) + xs = module_group(*xs, basic_op_kwargs=group_op_kwargs[group_idx]) if isinstance(xs, tuple): x, ys = xs[0], xs[1:] extra_outputs.extend(ys) @@ -196,3 +207,39 @@ def forward( if extra_outputs: return (x,) + tuple(extra_outputs) return x + + def _resolve_op_kwargs( + self, + op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]], + ) -> list[Optional[list[dict[str, Any]]]]: + """Map per-op kwargs onto each module group's basic-op kwargs list.""" + group_kwargs: list[Optional[list[dict[str, Any]]]] = [None] * len(self._module_groups) + if not op_kwargs: + return group_kwargs + + # Construct map from basic-op id to its kwarg dict + resolved: dict[int, dict[str, Any]] = {} + for key, kwargs in op_kwargs.items(): + module = self[key] if isinstance(key, int) else key + if not isinstance(module, BasicOperation): + raise ValueError( + f"Attempted to provide forward kwargs to {type(module).__name__}, but " + "Sequential only allows providing forward kwargs to a BasicOperation." + ) + resolved[id(module)] = kwargs + + # Slot each op's kwargs into its module group by matching basic-op identity + for group_idx, module_group in enumerate(self._module_groups): + if not isinstance(module_group, OperationFuser): + continue + # pylint: disable-next=protected-access + group_kwargs[group_idx] = [{} for _ in module_group._basic_ops] + # pylint: disable-next=protected-access + for idx, op in enumerate(module_group._basic_ops): + if id(op) in resolved: + group_kwargs[group_idx][idx] = resolved.pop(id(op)) + + # Any keys left over target an op that is not in this Sequential + if resolved: + raise ValueError("op_kwargs contains keys that are not in this Sequential") + return group_kwargs From 868d8d9216da361c666519652115e23688db5211 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 20 Jul 2026 10:28:26 -0700 Subject: [PATCH 25/35] Fix FusedAdam empty tensor handling (#3212) * Fix FusedAdam empty tensor handling Signed-off-by: Jingyue Wu * Move empty tensor filtering into MultiTensorApply Signed-off-by: Jingyue Wu --------- Signed-off-by: Jingyue Wu Co-authored-by: vthumbe1503 --- tests/pytorch/test_fused_optimizer.py | 26 ++++++++++++++++++ .../multi_tensor/multi_tensor_apply.cuh | 3 +++ .../pytorch/optimizers/fused_sgd.py | 27 +++++++++---------- .../pytorch/optimizers/multi_tensor_apply.py | 10 +++++++ 4 files changed, 52 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index a2863cba98..6832ef89dd 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -166,6 +166,19 @@ def test_frozen_model(self): torch.testing.assert_close(ref_param, tst_param) + def test_empty_param_at_end_of_group(self): + tensors = [ + torch.ones(4, dtype=torch.float, device="cuda"), + torch.empty(0, dtype=torch.float, device="cuda"), + ] + ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(tensors, self.options) + + self.gen_grad(ref_param, tst_param) + ref_optim.step() + tst_optim.step() + + torch.testing.assert_close(ref_param, tst_param) + def gen_precision_aware_test( self, use_fp8_params, @@ -796,6 +809,19 @@ def test_float(self): def test_half(self): self.gen_single_type_test(param_type=torch.float16) + def test_empty_param_at_end_of_group(self): + tensors = [ + torch.ones(4, dtype=torch.float, device="cuda"), + torch.empty(0, dtype=torch.float, device="cuda"), + ] + ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(tensors, self.options) + + self.gen_grad(ref_param, tst_param) + ref_optim.step() + tst_optim.step() + + torch.testing.assert_close(ref_param, tst_param) + class Model(torch.nn.Module): def __init__(self): diff --git a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh index 6710a161b3..d6ac23d4af 100644 --- a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh +++ b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh @@ -75,6 +75,9 @@ void multi_tensor_apply(int64_t block_size, int64_t chunk_size, loc_tensor_info++; auto chunks_this_tensor = (tensor_lists[0][t]->numel() + chunk_size - 1) / chunk_size; + NVTE_CHECK(chunks_this_tensor > 0, + "multi_tensor_apply expects tensors with at least one chunk; zero-sized tensors " + "must be filtered before launch because they skip the chunk loop"); for (auto chunk = 0; chunk < chunks_this_tensor; chunk++) { tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; diff --git a/transformer_engine/pytorch/optimizers/fused_sgd.py b/transformer_engine/pytorch/optimizers/fused_sgd.py index d7ab3fe9fe..10151e1406 100644 --- a/transformer_engine/pytorch/optimizers/fused_sgd.py +++ b/transformer_engine/pytorch/optimizers/fused_sgd.py @@ -295,20 +295,19 @@ def step(self, closure=None): for _, (launch_set, first_run) in enumerate(zip(launch_sets, first_runs)): assert len(launch_set[0]) == len(launch_set[1]) assert len(launch_set[0]) == len(launch_set[2]) - if len(launch_set[0]) > 0: - multi_tensor_applier( - self.multi_tensor_sgd, - self._dummy_overflow_buf, - launch_set, - weight_decay, - momentum, - dampening, - group["lr"], - nesterov, - first_run, - self.wd_after_momentum, - 1.0 / self.most_recent_scale, - ) + multi_tensor_applier( + self.multi_tensor_sgd, + self._dummy_overflow_buf, + launch_set, + weight_decay, + momentum, + dampening, + group["lr"], + nesterov, + first_run, + self.wd_after_momentum, + 1.0 / self.most_recent_scale, + ) self.most_recent_scale = 1.0 self.scale_set_by_backward = False diff --git a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py b/transformer_engine/pytorch/optimizers/multi_tensor_apply.py index a5cbd27337..e7d4fa0db3 100644 --- a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py +++ b/transformer_engine/pytorch/optimizers/multi_tensor_apply.py @@ -18,6 +18,16 @@ def __call__(self, op, noop_flag_buffer, tensor_lists, *args): if isinstance(t, DTensor): tensor_lists[i][j] = t._local_tensor + if any(len(tensors) != len(tensor_lists[0]) for tensors in tensor_lists): + raise RuntimeError("Expected aligned multi-tensor lists.") + + keep_slot = [tensor.numel() > 0 for tensor in tensor_lists[0]] + for i, tensors in enumerate(tensor_lists): + tensor_lists[i] = [tensor for tensor, keep in zip(tensors, keep_slot) if keep] + + if not tensor_lists[0]: + return None + return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args) From 870a68ef6e6b7f888371752caef8e99cd820f680 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:19:17 -0700 Subject: [PATCH 26/35] Changed VERSION to 2.19.0.dev0 (#3228) Signed-off-by: Kshitij Lakhani --- build_tools/VERSION.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build_tools/VERSION.txt b/build_tools/VERSION.txt index 830a65a39c..e3daa3c6a7 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.18.0.dev0 +2.19.0.dev0 From 622a3eea14042a38bad2b2b64bb32a788eacfec1 Mon Sep 17 00:00:00 2001 From: Fei Wu <33940270+YangFei1990@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:09:25 -0700 Subject: [PATCH 27/35] [PyTorch] NCCL EP zero copy with symmem pool and user provided recv_topk_weight tensor (#3187) * expose user-provided weights * adding pool based symm allocation; remove the persistent buffer in EpBuffer * add zero copy tests Signed-off-by: YangFei1990 --------- Signed-off-by: YangFei1990 Co-authored-by: Phuong Nguyen --- tests/pytorch/distributed/run_ep.py | 59 ++++---- transformer_engine/pytorch/distributed.py | 56 +++++++- transformer_engine/pytorch/ep.py | 155 ++++++++++------------ 3 files changed, 159 insertions(+), 111 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index ee6a97ffea..534fd9642a 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -19,6 +19,7 @@ ep_dispatch, ep_combine, symm_mem_alloc, + is_symm_backed, _ep_combine_raw, _ep_dispatch_raw, ) @@ -155,13 +156,7 @@ def setUp(self): ): self.skipTest("not exercised in zero-copy mode") - def _make_buffer( - self, - alignment=0, - top_k=TOP_K, - dispatch_recv_tokens=None, - combine_grad_expert_out=None, - ): + def _make_buffer(self, alignment=0, top_k=TOP_K): return EpBuffer( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, @@ -169,8 +164,6 @@ def _make_buffer( hidden_dim=HIDDEN_DIM, num_local_experts=NUM_LOCAL_EXPERTS, alignment=alignment, - dispatch_recv_tokens=dispatch_recv_tokens, - combine_grad_expert_out=combine_grad_expert_out, ) def _expert_out(self, expert_out): @@ -253,10 +246,10 @@ def test_dispatch_autograd(self): ] for label, recv_tokens in cases: with self.subTest(case=label): - buf = self._make_buffer(dispatch_recv_tokens=recv_tokens) + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) - rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w) + rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w, recv_tokens=recv_tokens) if recv_tokens is not None: # caller-supplied recv_tokens must be used in place self.assertEqual(rt.data_ptr(), recv_tokens.data_ptr()) rt = self._stage_grad_symm(rt) @@ -269,20 +262,17 @@ def test_dispatch_autograd(self): @_zero_copy_test_include def test_caller_provides_dispatch_recv_tokens(self): - """Caller-supplied recv_tokens: EpBuffer adopts it (recv_topk_weights stays - owned) and ep_dispatch returns a view of the caller's buffer.""" + """Caller-supplied recv_tokens (symm-mem-backed in zero-copy): ep_dispatch + writes into it and returns a view of the caller's buffer.""" if ZERO_COPY: rc = self.cfg.recv_capacity_per_rank rt_buf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: rt_buf, _rw_buf, _ = self._make_raw_recv() - buf = self._make_buffer(dispatch_recv_tokens=rt_buf) - self.assertEqual(buf.recv_tokens_symm_buf.data_ptr(), rt_buf.data_ptr()) - if ZERO_COPY: # recv_topk_weights is always buffer-owned in zero-copy - self.assertIsNotNone(buf.recv_topk_weights_symm_buf) + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) - rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w) + rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w, recv_tokens=rt_buf) self.assertEqual(rt.data_ptr(), rt_buf.data_ptr()) rt = self._stage_grad_symm(rt) rw = self._stage_grad_symm(rw) @@ -294,22 +284,45 @@ def test_caller_provides_dispatch_recv_tokens(self): @_zero_copy_test_include def test_caller_provides_grad_expert_out(self): - """Caller-supplied grad_expert_out: EpBuffer adopts it as the combine - backward grad target (symm-mem under zero-copy).""" + """Caller-supplied grad_out (symm-mem-backed in zero-copy): ep_combine's + backward scatters the expert-out grad into it.""" rc = self.cfg.recv_capacity_per_rank if ZERO_COPY: gbuf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: gbuf = torch.empty(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) - buf = self._make_buffer(combine_grad_expert_out=gbuf) - self.assertEqual(buf.grad_expert_out_symm_buf.data_ptr(), gbuf.data_ptr()) + gbuf.zero_() + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) recv_t = self._stage_grad_symm(recv_t) recv_w = self._stage_grad_symm(recv_w) expert_out = self._expert_out(self._weighted(recv_t, recv_w)) - out = ep_combine(buf, expert_out) + out = ep_combine(buf, expert_out, grad_out=gbuf) + (0.5 * (out.float() ** 2).sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) + # the caller-owned buffer was used as the combine-bwd scatter target + self.assertGreater(gbuf.abs().sum().item(), 0.0) + + @_zero_copy_test_include + def test_zero_copy_pool_auto_alloc(self): + """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO + tensors from the symm-mem pool (is_symm_backed). This is the primary mcore + path — mcore hands no caller buffers, TE pools them on the fly.""" + if not ZERO_COPY: + self.skipTest("zero-copy pool auto-alloc only") + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) # recv_tokens=None -> pool + self.assertTrue(is_symm_backed(recv_t)) # dispatch recv came from the symm-mem pool + recv_t = self._stage_grad_symm(recv_t) + recv_w = self._stage_grad_symm(recv_w) + expert_out = self._expert_out(self._weighted(recv_t, recv_w)) + out = ep_combine(buf, expert_out) # grad_out=None -> bwd allocs the grad from the pool (0.5 * (out.float() ** 2).sum()).backward() torch.cuda.synchronize() torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index c050f26869..569335d93f 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1842,13 +1842,54 @@ def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_gr return msg +_SYMM_MEM_POOL = None +_SYMM_MEM_POOL_BACKEND = None + + +def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): + """Process-wide torch MemPool backed by the symmetric-memory allocator, created once (each rank + drives one device). The pool/allocator captures the backend at creation and there is no per-pool + backend arg, so the (process-global) backend is always set before the pool is created. The + collective rendezvous cost is amortized across allocations (paid per new segment, not per buffer). + """ + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND + if _SYMM_MEM_POOL is None: + symm_mem.set_backend(backend) + _SYMM_MEM_POOL_BACKEND = backend + if hasattr(symm_mem, "get_mem_pool"): + _SYMM_MEM_POOL = symm_mem.get_mem_pool(device) + elif hasattr(torch.cuda, "MemPool") and hasattr(symm_mem, "get_mempool_allocator"): + _SYMM_MEM_POOL = torch.cuda.MemPool(symm_mem.get_mempool_allocator(device)) + else: + raise RuntimeError( + "No symmetric-memory MemPool API available (need torch symm-mem get_mem_pool, or " + "torch.cuda.MemPool + get_mempool_allocator)." + ) + elif backend != _SYMM_MEM_POOL_BACKEND: + raise RuntimeError( + f"symm-mem pool already created with backend {_SYMM_MEM_POOL_BACKEND!r}; " + f"cannot switch to {backend!r}" + ) + return _SYMM_MEM_POOL + + def symm_mem_alloc( shape, dtype: torch.dtype, ep_group: dist_group_type, device: Optional[torch.device] = None, + use_pool: bool = False, + backend: str = "NCCL", ) -> torch.Tensor: - """Allocate and rendezvous a symm-mem buffer on ep_group. Collective on ep_group.""" + """Allocate a symm-mem buffer on ep_group. + + ``use_pool=False`` (default): freshly allocate and do one explicit collective ``rendezvous`` per + buffer, for caller-owned static buffers (fp8 zero-copy). ``use_pool=True``: allocate from a + process-wide symm-mem MemPool whose segments are auto-registered (implicit mempool), so no explicit + rendezvous is needed and torch manages the tensor lifecycle (freed back to the pool) — for + lifecycle-managed zero-copy, e.g. bf16, where the recv buffer is saved for backward and so cannot + be a shared static buffer. ``backend`` selects the symm-mem backend (default NCCL; for the pool it + is captured at pool creation).""" if device is None: device = torch.device("cuda", torch.cuda.current_device()) if not HAS_TORCH_SYMMETRIC: @@ -1856,10 +1897,15 @@ def symm_mem_alloc( "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc " "requires PyTorch built with NCCL symm-mem support." ) - if symm_mem.get_backend(device) != "NCCL": - symm_mem.set_backend("NCCL") - t = symm_mem.empty(*shape, dtype=dtype, device=device) - symm_mem.rendezvous(t, group=ep_group) + if use_pool: + pool = _get_symm_mem_pool(device, backend) + with torch.cuda.use_mem_pool(pool): + t = torch.empty(*shape, dtype=dtype, device=device) + else: + if symm_mem.get_backend(device) != backend: + symm_mem.set_backend(backend) + t = symm_mem.empty(*shape, dtype=dtype, device=device) + symm_mem.rendezvous(t, group=ep_group) return t diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index e57caa1f42..0ae4805caa 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -25,6 +25,7 @@ "ep_dispatch", "ep_combine", "symm_mem_alloc", + "is_symm_backed", ] @@ -158,18 +159,32 @@ def ep_finalize() -> None: _EP_GROUP = None +def is_symm_backed(t: torch.Tensor) -> bool: + """Whether ``t`` is symm-mem-backed on the EP group. Prefer torch's local ``is_symm_mem_tensor`` + when the build provides it (no collective, no exception); otherwise fall back to the rendezvous + probe the C++ ep kernel uses (``maybe_make_window``): cached for an already-registered tensor, + raises for a plain one.""" + from torch.distributed import _symmetric_memory as _symm + + if hasattr(_symm, "is_symm_mem_tensor"): + return bool(_symm.is_symm_mem_tensor(t)) + if _EP_GROUP is None: + raise RuntimeError( + "is_symm_backed called before ensure_nccl_ep_bootstrapped(); no EP group registered." + ) + try: + _symm.rendezvous(t, _EP_GROUP.group_name) + return True + except Exception: # pylint: disable=broad-exception-caught + return False + + # Buffer class EpBuffer: - """Per-microbatch EP layer state holding handle_mem and token_counts. + """Per-microbatch EP layer state: handle_mem, token_counts, and shape/dtype config. Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). - - In zero-copy mode the buffer owns the symm-mem buffers the one-sided path - requires: the dispatch recv outputs (recv_tokens, recv_topk_weights) and the - combine backward grad target. One set per buffer, so each layer/microbatch is - isolated. In normal mode these are None and allocated in-flight instead (recv - outputs in the dispatch forward, the combine grad in the backward). """ __slots__ = ( @@ -184,39 +199,8 @@ class EpBuffer: "device", "token_counts", "zero_copy", - "recv_tokens_symm_buf", - "recv_topk_weights_symm_buf", - "grad_expert_out_symm_buf", ) - def _alloc_symm_buffers(self) -> None: - """Fill in buffer-owned symm-mem buffers the caller did not supply. - recv_topk_weights is always owned. In normal mode caller-supplied - tensors are kept as-is and the rest stay None (allocated in-flight).""" - if not self.zero_copy: - self.recv_topk_weights_symm_buf = None - return - if _EP_GROUP is None: - raise RuntimeError( - "ep_bootstrap must be called before constructing a zero-copy EpBuffer" - ) - rc, h = self.recv_capacity_per_rank, self.hidden_dim - # Persistent across microbatches; keep resident under CPU offloading. - self.recv_topk_weights_symm_buf = symm_mem_alloc( - (rc,), torch.float32, _EP_GROUP, device=self.device - ) - mark_not_offload(self.recv_topk_weights_symm_buf) - if self.recv_tokens_symm_buf is None: - self.recv_tokens_symm_buf = symm_mem_alloc( - (rc, h), self.payload_dtype, _EP_GROUP, device=self.device - ) - mark_not_offload(self.recv_tokens_symm_buf) - if self.grad_expert_out_symm_buf is None: - self.grad_expert_out_symm_buf = symm_mem_alloc( - (rc, h), self.payload_dtype, _EP_GROUP, device=self.device - ) - mark_not_offload(self.grad_expert_out_symm_buf) - def __init__( self, top_k: int, @@ -227,14 +211,7 @@ def __init__( alignment: int = 0, payload_dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, - dispatch_recv_tokens: Optional[torch.Tensor] = None, - combine_grad_expert_out: Optional[torch.Tensor] = None, ) -> None: - """Pass ``dispatch_recv_tokens`` (dispatch recv output) and/or - ``combine_grad_expert_out`` (combine backward grad target) to use caller-owned - buffers; the buffer then skips allocating them. Both must be symm-mem-backed - under zero-copy. Whatever is left None is buffer-owned (zero-copy) or allocated - in-flight (normal mode). recv_topk_weights is always owned by the buffer.""" if device is None: device = torch.device("cuda", torch.cuda.current_device()) alignment = int(alignment) @@ -249,15 +226,12 @@ def __init__( self.payload_dtype = payload_dtype self.device = device self.zero_copy = bool(tex.ep_get_zero_copy()) - self.recv_tokens_symm_buf = dispatch_recv_tokens - self.grad_expert_out_symm_buf = combine_grad_expert_out size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) self.token_counts = torch.empty(self.num_local_experts, dtype=torch.int32, device=device) # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) - self._alloc_symm_buffers() # torch.library custom ops (so they don't graph-break under torch.compile) @@ -476,14 +450,15 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: class _EpCombine(torch.autograd.Function): """Autograd combine. - bwd scatters the expert_out grad into ``grad_symm_buf`` (EpBuffer-owned - symm-mem, one-sided) in zero-copy mode, or into a plain tensor allocated - in-flight here otherwise. The latter keeps allocation torch.compile / - CUDA-graph safe and lets autograd own the grad's lifetime. + bwd scatters the expert_out grad into ``grad_out``. When the caller supplies it + (mcore-managed mode) that buffer is used as-is; otherwise it is allocated on the + fly here — from the symm-mem pool in zero-copy mode (one-sided target), or a plain + tensor in normal mode (keeps allocation torch.compile / CUDA-graph safe and lets + autograd own the grad's lifetime). - ``grad_symm_buf`` is the backward's scatter target (an output it writes, never - reads), so it is stashed as a plain ctx attribute rather than via - save_for_backward, which would version-track a tensor we mutate. + ``grad_out`` is the backward's scatter target (an output it writes, never reads), + so it is stashed as a plain ctx attribute rather than via save_for_backward, which + would version-track a tensor we mutate. """ @staticmethod @@ -492,7 +467,7 @@ def forward( # type: ignore[override] handle_mem: torch.Tensor, num_local_tokens: int, hidden_dim: int, - grad_symm_buf: Optional[torch.Tensor], + grad_out: Optional[torch.Tensor], expert_out: torch.Tensor, ): """Combine fwd; stashes the bwd grad target or expert_out shape to size it.""" @@ -500,8 +475,8 @@ def forward( # type: ignore[override] result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) ctx.save_for_backward(handle_mem) - ctx.grad_symm_buf = grad_symm_buf - if grad_symm_buf is None: + ctx.grad_out = grad_out + if grad_out is None: ctx.expert_out_shape = expert_out.shape ctx.expert_out_dtype = expert_out.dtype ctx.device = device @@ -513,17 +488,17 @@ def backward(ctx, g_result): # type: ignore[override] if not g_result.is_contiguous(): g_result = g_result.contiguous() (handle_mem,) = ctx.saved_tensors - grad_expert_out = ctx.grad_symm_buf + grad_expert_out = ctx.grad_out if grad_expert_out is None: - grad_expert_out = torch.empty( - ctx.expert_out_shape, dtype=ctx.expert_out_dtype, device=ctx.device + grad_expert_out = _alloc_io( + ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() ) torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) return ( None, # handle_mem None, # num_local_tokens None, # hidden_dim - None, # grad_symm_buf + None, # grad_out grad_expert_out, ) @@ -539,18 +514,33 @@ def _require_bf16(name: str, t: torch.Tensor) -> None: ) +def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tensor: + """Allocate a dispatch/combine IO tensor the caller did not supply: from the symm-mem pool in + zero-copy mode (auto-registered segment, lifecycle managed by torch refcount), else plain.""" + if zero_copy: + t = symm_mem_alloc(shape, dtype, _EP_GROUP, device=device, use_pool=True) + # symm-mem storage is non-resizable; exempt it from CPU activation offloading (which + # releases via storage.resize_(0)). Matters for bf16 recv_tokens (the saved activation). + mark_not_offload(t) + return t + return torch.empty(*shape, dtype=dtype, device=device) + + def ep_dispatch( buffer: EpBuffer, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, + *, + recv_tokens: Optional[torch.Tensor] = None, + recv_topk_weights: Optional[torch.Tensor] = None, ): """Prepare + dispatch with autograd. topk_idx must be int32 or int64. - recv_tokens comes from the EpBuffer (caller-supplied or buffer-owned under - zero-copy) or is allocated in-flight (normal mode). recv_topk_weights is always - owned by the buffer. Returns (recv_tokens, recv_topk_weights, token_counts); - token_counts is non-diff. + ``recv_tokens`` / ``recv_topk_weights`` are the dispatch recv outputs: pass caller-owned buffers + (mcore-managed mode; in zero-copy they must be symm-mem-backed) or leave them None to allocate on + the fly (zero-copy: symm-mem pool; normal: plain). Returns (recv_tokens, recv_topk_weights, + token_counts); token_counts is non-diff. """ _require_bf16("tokens", tokens) if topk_weights.dtype is not torch.float32: @@ -558,19 +548,17 @@ def ep_dispatch( f"topk_weights must be float32; got dtype={topk_weights.dtype}. " "Cast with topk_weights.float() before calling." ) - recv_tokens = buffer.recv_tokens_symm_buf if recv_tokens is None: - recv_tokens = torch.empty( - buffer.recv_capacity_per_rank, - buffer.hidden_dim, - dtype=buffer.payload_dtype, - device=buffer.device, + recv_tokens = _alloc_io( + (buffer.recv_capacity_per_rank, buffer.hidden_dim), + buffer.payload_dtype, + buffer.device, + buffer.zero_copy, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (buffer.recv_capacity_per_rank,), torch.float32, buffer.device, buffer.zero_copy ) - recv_topk_weights = ( - buffer.recv_topk_weights_symm_buf - if buffer.zero_copy - else torch.empty(buffer.recv_capacity_per_rank, dtype=torch.float32, device=buffer.device) - ) return _EpDispatch.apply( buffer.handle_mem, buffer.top_k, @@ -589,22 +577,23 @@ def ep_combine( expert_out: torch.Tensor, *, num_local_tokens: Optional[int] = None, + grad_out: Optional[torch.Tensor] = None, ): """Combine with autograd; caller pre-applies topk weighting. - The backward scatters the expert_out grad into the EpBuffer grad target - (caller-supplied or buffer-owned under zero-copy), or a tensor allocated - in-flight (normal mode). Result shape is (num_local_tokens, buffer.hidden_dim); - defaults to buffer.max_tokens_per_rank rows. + ``expert_out`` is the combine input (always caller-supplied; in zero-copy it must be symm-mem- + backed). ``grad_out`` is the backward's grad target: pass a caller-owned buffer (mcore-managed + mode) or leave it None to allocate on the fly in the backward (zero-copy: symm-mem pool; normal: + plain). Result shape is (num_local_tokens, buffer.hidden_dim); defaults to + buffer.max_tokens_per_rank rows. """ _require_bf16("expert_out", expert_out) if num_local_tokens is None: num_local_tokens = buffer.max_tokens_per_rank - grad_expert_out = buffer.grad_expert_out_symm_buf return _EpCombine.apply( buffer.handle_mem, num_local_tokens, buffer.hidden_dim, - grad_expert_out, + grad_out, expert_out, ) From 73c72ff01122bb0fca6ddc79562c4485cd3f68fb Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 22 Jul 2026 09:43:29 -0700 Subject: [PATCH 28/35] [Common] Migrate NCCL EP submodule to NVIDIA/nccl-extensions (#3222) * Migrate NCCL EP submodule to NVIDIA/nccl-extensions * Drop PYTHONPATH override from EP test, example, and bench launchers * Drop cross-mode recv comparison in EP zero-copy IdentityAllSymm test Signed-off-by: Phuong Nguyen * [Common] Rename 3rdparty/nccl submodule directory to nccl-extensions Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- .gitmodules | 4 ++-- 3rdparty/nccl | 1 - 3rdparty/nccl-extensions | 1 + examples/jax/ep/run_test_ep.sh | 4 ++-- examples/pytorch/ep/bench/run_ep_bench.sh | 2 -- .../pytorch/ep/bench/run_nccl_ep_bench.sh | 4 ++-- examples/pytorch/ep/run_test_ep.sh | 1 - setup.py | 12 +++++------ tests/cpp_distributed/test_ep.cu | 21 ++++++------------- tests/jax/multi_process_launch_ep.sh | 4 ++-- tests/pytorch/distributed/run_test_ep.sh | 2 -- transformer_engine/common/CMakeLists.txt | 4 ++-- 12 files changed, 23 insertions(+), 37 deletions(-) delete mode 160000 3rdparty/nccl create mode 160000 3rdparty/nccl-extensions diff --git a/.gitmodules b/.gitmodules index 07647e915e..88d431cba5 100644 --- a/.gitmodules +++ b/.gitmodules @@ -5,5 +5,5 @@ path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git [submodule "3rdparty/nccl"] - path = 3rdparty/nccl - url = https://github.com/NVIDIA/nccl.git + path = 3rdparty/nccl-extensions + url = https://github.com/NVIDIA/nccl-extensions.git diff --git a/3rdparty/nccl b/3rdparty/nccl deleted file mode 160000 index b87848fbc5..0000000000 --- a/3rdparty/nccl +++ /dev/null @@ -1 +0,0 @@ -Subproject commit b87848fbc52da65b5a898b4ac6633fcf51cec4ed diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions new file mode 160000 index 0000000000..2c6135a721 --- /dev/null +++ b/3rdparty/nccl-extensions @@ -0,0 +1 @@ +Subproject commit 2c6135a721824ff792af7b72900b0ab758fa1f98 diff --git a/examples/jax/ep/run_test_ep.sh b/examples/jax/ep/run_test_ep.sh index 1305ca6fd1..86aa6ca087 100755 --- a/examples/jax/ep/run_test_ep.sh +++ b/examples/jax/ep/run_test_ep.sh @@ -32,8 +32,8 @@ export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" COORD="${COORD:-127.0.0.1:12345}" TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-300}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') diff --git a/examples/pytorch/ep/bench/run_ep_bench.sh b/examples/pytorch/ep/bench/run_ep_bench.sh index fefecd7fa9..3b0977e4c3 100755 --- a/examples/pytorch/ep/bench/run_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_ep_bench.sh @@ -26,10 +26,8 @@ if [ "${NSYS}" -eq 1 ] && [ "${KINETO}" -eq 1 ]; then fi SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}" diff --git a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh index 8f6da04a00..ac5fcacc25 100755 --- a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh @@ -23,8 +23,8 @@ TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -BIN="${TE_REPO_ROOT}/3rdparty/nccl/build/test/nccl_ep/ep_bench" -LIB="${TE_REPO_ROOT}/3rdparty/nccl/build/lib" +BIN="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/test/nccl_ep/ep_bench" +LIB="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/lib" [ -x "${BIN}" ] || { echo "ep_bench not built at ${BIN}" >&2; exit 2; } NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) diff --git a/examples/pytorch/ep/run_test_ep.sh b/examples/pytorch/ep/run_test_ep.sh index 13b41f4cb2..d8e6b50556 100755 --- a/examples/pytorch/ep/run_test_ep.sh +++ b/examples/pytorch/ep/run_test_ep.sh @@ -17,7 +17,6 @@ if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi : ${TEST_TIMEOUT_S:=120} SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py" -export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" # Stage JIT cubins on tmpfs for fast iteration. : ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} diff --git a/setup.py b/setup.py index 150d92969c..c0c37364a9 100644 --- a/setup.py +++ b/setup.py @@ -207,11 +207,11 @@ def _discover_nccl_home() -> str: def build_nccl_ep_submodule() -> str: - """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME.""" - nccl_root = current_file_path / "3rdparty" / "nccl" - if not (nccl_root / "Makefile").exists(): + """Build libnccl_ep.a from the 3rdparty/nccl-extensions submodule and return NCCL_HOME.""" + nccl_root = current_file_path / "3rdparty" / "nccl-extensions" + if not (nccl_root / "nccl_ep" / "Makefile").exists(): raise RuntimeError( - f"NCCL submodule not found at {nccl_root}. " + f"NCCL EP submodule not found at {nccl_root}. " "Run `git submodule update --init --recursive`." ) @@ -267,13 +267,13 @@ def build_nccl_ep_submodule() -> str: "rebuilding libnccl_ep.a" ) subprocess.check_call( - ["make", "-C", "contrib/nccl_ep", "clean"], + ["make", "-C", "nccl_ep", "clean"], cwd=str(nccl_root), env=env, ) print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')") subprocess.check_call( - ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"], + ["make", "-j", str(nproc), "-C", "nccl_ep", "lib"], cwd=str(nccl_root), env=env, ) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 7dbbcdce9d..47732196ed 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -675,7 +675,7 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface TEST_P(EPPipelineTest, FullForwardBackward) { const DType dtype = GetParam(); // NCCL EP backend currently asserts ncclBfloat16 in ncclEpDispatch - // (contrib/nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. + // (nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. if (dtype != DType::kBFloat16) { GTEST_SKIP() << test::typeName(dtype) << " not yet supported by NCCL EP backend"; } @@ -750,8 +750,9 @@ class EPZeroCopyTest : public EpOpTestBase { }; TYPED_TEST_SUITE(EPZeroCopyTest, EPBf16Only); -// Identity round-trip with symm-mem on dispatch i/o + combine input. Bit-exact -// vs HBM reference (same routing, same input). +// Identity round-trip with symm-mem on dispatch i/o + combine input. The combined +// result is bit-exact vs the HBM reference; the intermediate recv buffer is not, +// since zero-copy and HBM dispatch use different per-expert layouts. TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { using Tok = TypeParam; EP_PULL_FIXTURE(); @@ -776,10 +777,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { ref_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); std::vector ref_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), - ref_recv.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); @@ -818,24 +816,17 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { symm_window(sym_recv), sym_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); std::vector sym_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, - sym_recv_host.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), sym_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); - // Compare per filled recv slot (HBM ref vs symm) and full result. - int total_recv = this->template read_total_recv(sym_buf); - for (int i = 0; i < total_recv * hidden_dim_; ++i) - ASSERT_EQ(tok_to_float(sym_recv_host[i]), tok_to_float(ref_recv[i])) - << "recv mismatch at " << i; + // Combined result is the cross-mode invariant (see note above). for (size_t i = 0; i < sym_result.size(); ++i) ASSERT_EQ(tok_to_float(sym_result[i]), tok_to_float(ref_result[i])) << "result mismatch at " << i; if (g_process_id == 0) - printf(" IdentityAllSymm: passed (recv_slots=%d, bit-exact vs HBM)\n", total_recv); + printf(" IdentityAllSymm: passed (result bit-exact vs HBM)\n"); NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh index ff89f712eb..8547d77f2b 100755 --- a/tests/jax/multi_process_launch_ep.sh +++ b/tests/jax/multi_process_launch_ep.sh @@ -17,8 +17,8 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 68b691f787..62c9a0207f 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -9,8 +9,6 @@ set -uo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) if [ "${DETECTED_GPUS}" -lt 4 ]; then diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index d09876c990..8eba515e5e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -446,10 +446,10 @@ option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) if(NVTE_WITH_NCCL_EP) # SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. # -- NCCL EP headers -------------------------------------------------------- -# Headers + libs are produced by the in-tree 3rdparty/nccl submodule build +# Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build # (auto-built by setup.py via build_nccl_ep_submodule). set(NCCL_EP_SUBMODULE_ROOT - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl-extensions") set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include") if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") message(FATAL_ERROR From 098b49678c3b57f252dd6caef82465caac2c3994 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Wed, 22 Jul 2026 15:34:58 -0500 Subject: [PATCH 29/35] [PyTorch] Enable fused FP8 block-scaling path in GroupedLinear module (#3171) * [Common/PyTorch] Support power-of-2 scales in grouped FP8 block-scaling quantize The default Float8BlockScaling recipe constrains scales to powers of 2, so the fused grouped path must honor the flag to stay numerically consistent with the unfused path. Thread a runtime pow_2_scales argument through the grouped quantize kernels (the shared scale helper already implements the rounding) and drop the force_pow_2_scales rejections. Also add a quantization-config parameter to nvte_group_quantize_dbias, which previously had no way to receive force_pow_2_scales or amax_epsilon on the bgrad path. Signed-off-by: Alp Dener * [PyTorch] Enable fused grouped FP8 block-scaling path in GroupedLinear module Admit Float8BlockQuantizer in the fused GroupedTensor path on Hopper. The existing usage flags already match the Hopper TN-only mapping and the grouped GEMM selects transposed columnwise storage for NN/NT layouts, so only the path predicate changes. The fused path is an explicit opt-in via NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM, so raise on Blackwell (SM100/SM110) instead of silently falling back; the fused path has no MXFP8-broadcast emulation. Extend the fused dbias path (tex.bgrad_group_quantize) to FP8 block scaling when dgrad is required (dbias is computed in the rowwise pass). Add fp8_block_scaling to the fused-path tests with a Hopper-only gate, assert the fused path engages via a group_quantize spy, and add a Blackwell error-path test. Signed-off-by: Alp Dener * [PyTorch] Enable FP8 block-scaling in GroupedLinear fusible op Replace the blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state with a per-op supports_float8_block_scaling flag and opt in the GroupedLinear op. Mirror the module-path predicate and fused-bgrad changes; since the graph-safe flow is default-on here (no env-var opt-in), other architectures fall back to the split-quantize flow instead of raising. Force use_split_accumulator=True for FP8 block-scaling operands in general_grouped_gemm_for_grouped_tensor, matching non-grouped general_gemm: cuBLAS has no fast-accum FP8 block-scaling algorithm, so the ops-layer forward failed algo selection without it. Add fp8_block_scaling coverage to the ops GroupedLinear tests. The CUDA-graph-safe test skips it for now: the replayed wgrad for the last expert diverges between replays depending on process allocation history; under investigation. Graph capture remains covered by the module-path test. Signed-off-by: Alp Dener * [PyTorch] Use persistent workspaces in grouped-tensor GEMM general_grouped_gemm_for_grouped_tensor allocated its setup workspace (the cuBLAS per-group pointer/dimension arrays) and its cuBLAS workspace with per-call torch.empty. Under make_graphed_callables the forward and backward graphs share one capture memory pool, and a per-call allocation's block returns to that pool as soon as the Python reference dies, so blocks alias across the two graphs and captured kernels from one graph overwrite the GEMM metadata the other graph reads at replay. Observed as allocation-history-dependent failures in the ops-layer GroupedLinear cuda-graph test: capture-time cublasLtMatmulAlgoGetHeuristic NOT_SUPPORTED errors and corrupted wgrad outputs. This is also the likely mechanism behind the FP8 block-scaling wgrad corruption under CUDA graphs previously observed on Hopper and attributed to cuBLAS. Cache the setup workspace per (device, group size) and reuse the cached per-device cuBLAS workspace from the non-grouped path; consecutive GEMMs reusing one workspace are ordered by the stream. Signed-off-by: Alp Dener * [PyTorch] Fix grouped FP8 block-scaling CUDA-graph deadlock via per-role cuBLAS workspaces The grouped-tensor GEMM path shared one persistent cuBLAS workspace across all grouped matmuls. cuBLAS's grouped GEMM keeps a grid-synchronization flag in the first bytes of that workspace and zeros it (via a captured memset) before each matmul. When the dgrad and wgrad grouped matmuls of a GroupedLinear backward share one workspace inside a replayed CUDA graph, that flag is aliased between the two matmuls; on the second graph replay the second matmul's cooperative kernel deadlocks with cuBLAS 13.6 (and corrupts the last expert's wgrad on cuBLAS < 13.6). The two matmuls are strictly stream-ordered (single stream, all-DEFAULT graph edges, no programmatic dependent launch), so this is shared-workspace reuse, not concurrent co-scheduling. Give dgrad/forward (slot 0) and wgrad (slot 1) distinct persistent cuBLAS workspaces, dedicated to the grouped path. Each slot remains a single persistent allocation, so CUDA-graph capture safety is preserved. Also drop the cuBLAS-version gate that skipped the FP8 block-scaling GroupedLinear CUDA-graph test, so it now exercises the fix on all supported cuBLAS versions. Signed-off-by: Alp Dener * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: document split-accumulator override, fix stale dbias comment - general_grouped_gemm_for_grouped_tensor: expand the comment to state that the fused grouped FP8 block-scaling GEMM forces use_split_accumulator=True and intentionally overrides the caller-supplied value, consistent with the Float8BlockScaling recipe (which fixes it True for fprop/dgrad/wgrad). - Float8BlockScaling recipe docstring: document that FP8 block scaling always uses split accumulation and that the fused grouped GEMM path ignores any caller- or recipe-supplied use_split_accumulator value. - GroupedLinear ops backward: correct the stale "BF16/FP16 path" comment; that branch also handles quantized paths where bgrad fusion did not apply (e.g. FP8 block scaling without a dgrad pass). Signed-off-by: Alp Dener * [PyTorch] Revert fusible-ops FP8 block-scaling; scope PR to GroupedLinear module Restrict this PR to the GroupedLinear module fused-quantize path. Revert the fusible-ops FP8 block-scaling enablement -- the BasicOperation opt-in gate, the GroupedLinear op support, and the fusible-ops test coverage -- back to main. Enabling fusible-ops FP8 block-scaling for both grouped and non-grouped paths is deferred to a separate PR. The blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state is restored. The split-accumulator guard in general_grouped_gemm_for_grouped_tensor is retained: it is correct for the module's FP8 block-scaling grouped GEMM. Signed-off-by: Alp Dener * [PyTorch] Isolate grouped wgrad cuBLAS workspace by NT layout, not out-discreteness _get_grouped_cublas_workspace slots were keyed on is_discrete_out as a proxy for "this is the wgrad GEMM", which only holds when wgrad writes a list of per-expert grads. With single_grouped_weight=True, wgrad writes a single grouped weight-grad (GroupedTensor out, not a list), so is_discrete_out is False and it collided with dgrad on slot 0 -- reintroducing the FP8 block-scaling grid-sync-flag aliasing deadlock/corruption under CUDA-graph replay. Key the slot on the wgrad layout (NT / transb) instead: fprop (TN) and dgrad (NN) share slot 0, wgrad (NT) is always isolated on slot 1. Signed-off-by: Alp Dener * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: isolate grouped cuBLAS workspace per layout; drop redundant test spy - _get_grouped_cublas_workspace now keys the persistent workspace on the grouped GEMM layout, so fprop (TN), dgrad (NN), and wgrad (NT) each get a distinct workspace. The previous NT-vs-rest scheme left fprop and dgrad sharing one workspace; those have also been reported to conflict under CUDA-graph replay. Documents that the deadlock is deterministic and present through cuBLAS 13.7. - Drop the group_quantize call-counting spy in test_grouped_linear_grouped_tensor_path_matches_legacy; fused-path engagement is covered by the graph-safe test. Signed-off-by: Alp Dener * updated grouped GEMM workspace comment on stale TMA descriptor related deadlocks Signed-off-by: Alp Dener * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alp Dener Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../test_cast_float8blockwise_grouped.cu | 16 ++++-- tests/cpp/operator/test_cast_mxfp8_grouped.cu | 4 +- tests/pytorch/test_grouped_linear.py | 55 ++++++++++++++++--- tests/pytorch/test_grouped_tensor.py | 14 ++++- .../common/cast/cast_grouped_dbias.cu | 5 +- .../common/cast/dispatch/quantize.cuh | 30 ++++------ .../group_quantize_fp8_blockwise.cuh | 33 ++++++----- .../common/include/transformer_engine/cast.h | 4 +- transformer_engine/common/recipe/__init__.py | 6 ++ .../pytorch/cpp_extensions/gemm.py | 55 +++++++++++++++---- .../pytorch/csrc/extensions/cast.cpp | 15 ++++- transformer_engine/pytorch/csrc/quantizer.cpp | 8 --- .../pytorch/module/grouped_linear.py | 40 ++++++++++++-- 13 files changed, 201 insertions(+), 84 deletions(-) diff --git a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu index bc9f104e17..90418d6287 100644 --- a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu +++ b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu @@ -368,6 +368,7 @@ struct TestConfig { ScalingDir dir; std::vector first_dims; size_t K; + bool force_pow_2_scales; }; class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam {}; @@ -375,7 +376,7 @@ class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam TEST_P(GroupedFP8BlockwiseTestSuite, Test) { const TestConfig& cfg = GetParam(); perform_test(cfg.shape_rep, cfg.block_dim, cfg.dir, cfg.first_dims, cfg.K, - /*force_pow_2_scales=*/false, /*epsilon=*/0.0f); + cfg.force_pow_2_scales, /*epsilon=*/0.0f); } std::vector make_configs() { @@ -387,11 +388,13 @@ std::vector make_configs() { for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) { for (auto dir : {ScalingDir::ROWWISE, ScalingDir::COLWISE, ScalingDir::BOTH}) { for (size_t K : Ks) { - for (const auto& v : uniform) { - configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K}); - } - for (const auto& v : jagged) { - configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K}); + for (bool pow2 : {false, true}) { + for (const auto& v : uniform) { + configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K, pow2}); + } + for (const auto& v : jagged) { + configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K, pow2}); + } } } } @@ -408,6 +411,7 @@ std::string make_name(const ::testing::TestParamInfo& info) { s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size()); s += "_M"; for (size_t m : c.first_dims) s += "_" + std::to_string(m); + s += (c.force_pow_2_scales ? "_POW2" : "_FP32SC"); return s; } diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index de72299be1..80b80da0d8 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -509,9 +509,9 @@ void performTest(const ProcessingMethod processing_method, break; } case ProcessingMethod::CAST_DBIAS: { - nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0); + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0); workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); - nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0); + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0); break; } case ProcessingMethod::CAST_DBIAS_DACT: { diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 64951a43b8..10d0dbd227 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1499,6 +1499,9 @@ def test_fp8_grouped_gemm(shape, accumulate): _fp8_available, _reason_for_no_fp8 = fp8_available, reason_for_no_fp8 _mxfp8_available, _reason_for_no_mxfp8 = mxfp8_available, reason_for_no_mxfp8 _nvfp4_available, _reason_for_no_nvfp4 = nvfp4_available, reason_for_no_nvfp4 +_fp8_block_scaling_available, _reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) @pytest.fixture(autouse=True) @@ -1590,8 +1593,14 @@ def _run_grouped_linear_path( recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling + ), + ), ], - ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], + ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) @pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN) @@ -1605,10 +1614,13 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( pytest.skip( "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor - # current scaling also runs on the Hopper grouped GEMM path. + # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor + # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - if use_fp8 and not is_current_scaling and device_capability < (10, 0): + is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() + if is_block_scaling and not (9, 0) <= device_capability < (10, 0): + pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") + if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): pytest.skip( "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." ) @@ -1906,8 +1918,14 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling + ), + ), ], - ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], + ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch): @@ -1918,10 +1936,13 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch pytest.skip( "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor - # current scaling also runs on the Hopper grouped GEMM path. + # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor + # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - if use_fp8 and not is_current_scaling and device_capability < (10, 0): + is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() + if is_block_scaling and not (9, 0) <= device_capability < (10, 0): + pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") + if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): pytest.skip( "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." ) @@ -2037,6 +2058,24 @@ def _train_step(x, dy, out_buf, *, use_graphed): torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols) +@pytest.mark.skipif(not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling) +@pytest.mark.skipif( + not (10, 0) <= torch.cuda.get_device_capability() <= (11, 0), + reason="Error path only triggers on Blackwell (SM100/SM110).", +) +def test_grouped_linear_fused_path_fp8_block_scaling_blackwell_error(monkeypatch): + """FP8BS + fused env var on Blackwell must raise, not silently fall back.""" + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") + FP8GlobalStateManager.reset() + dtype = torch.bfloat16 + grouped_linear = GroupedLinear(2, 128, 128, bias=False, params_dtype=dtype, device="cuda") + x = torch.randn(256, 128, device="cuda", dtype=dtype, requires_grad=True) + m_splits = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + with pytest.raises(RuntimeError, match="Hopper-only"): + with autocast(enabled=True, recipe=recipe.Float8BlockScaling()): + grouped_linear(x, m_splits) + + @pytest.mark.parametrize("swizzle_type", ["mxfp8_rowwise", "mxfp8_columnwise", "nvfp4"]) def test_swizzle_scales_and_pack_ptrs_for_discrete_weights( swizzle_type: str, diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index eeb6e7a394..4dd52ee2bd 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -1008,11 +1008,19 @@ def _assert_fp8_cs_group_quantize_matches_reference( @pytest.mark.parametrize("shape_case", ["uniform", "varying_first"]) @pytest.mark.parametrize("direction", ["rowwise", "columnwise", "both"]) @pytest.mark.parametrize("output_dbias", [False, True]) + @pytest.mark.parametrize( + "force_pow_2_scales", [False, True], ids=["fp32_scales", "pow2_scales"] + ) @pytest.mark.skipif( not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped ) def test_quantize_grouped_fp8_blockwise( - self, block_scaling_dim: int, shape_case: str, direction: str, output_dbias: bool + self, + block_scaling_dim: int, + shape_case: str, + direction: str, + output_dbias: bool, + force_pow_2_scales: bool, ) -> None: """Test grouped FP8 block-scaling quantization against per-tensor quantization. @@ -1061,7 +1069,7 @@ def test_quantize_grouped_fp8_blockwise( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=columnwise, - force_pow_2_scales=False, + force_pow_2_scales=force_pow_2_scales, amax_epsilon=0.0, block_scaling_dim=block_scaling_dim, ) @@ -1081,7 +1089,7 @@ def test_quantize_grouped_fp8_blockwise( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=False, + force_pow_2_scales=force_pow_2_scales, amax_epsilon=0.0, block_scaling_dim=block_scaling_dim, ) diff --git a/transformer_engine/common/cast/cast_grouped_dbias.cu b/transformer_engine/common/cast/cast_grouped_dbias.cu index 5290255a00..b7ced30b11 100644 --- a/transformer_engine/common/cast/cast_grouped_dbias.cu +++ b/transformer_engine/common/cast/cast_grouped_dbias.cu @@ -11,7 +11,8 @@ #include "dispatch/quantize.cuh" void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { + NVTEGroupedTensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias); using namespace transformer_engine; @@ -20,5 +21,5 @@ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor constexpr const NVTEGroupedTensor activation_input = nullptr; dispatch::group_quantize_bwd_helper( - input, activation_input, output, dbias, workspace, nullptr, stream); + input, activation_input, output, dbias, workspace, quant_config, stream); } diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 033d464bcf..54f374b551 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -475,22 +475,16 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor } case NVTE_BLOCK_SCALING_1D: { NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_1D."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); fp8_blockwise::group_quantize_blockwise_1d(input_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream); + quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream); break; } case NVTE_BLOCK_SCALING_2D: { NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_2D."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); fp8_blockwise::group_quantize_blockwise_2d(input_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream); + quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream); break; } default: @@ -537,22 +531,18 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe case NVTE_BLOCK_SCALING_1D: case NVTE_BLOCK_SCALING_2D: { NVTE_CHECK(!IS_DACT, "IS_DACT is not implemented for grouped FP8 block scaling."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); // dbias is computed in-kernel and reduced per-expert inside group_quantize_blockwise_{1d,2d} // (mirrors MXFP8); those also handle the two-call workspace sizing protocol. GroupedTensor *dbias_arg = IS_DBIAS ? dbias_tensor : nullptr; Tensor *workspace_arg = IS_DBIAS ? workspace_tensor : nullptr; if (scaling_mode == NVTE_BLOCK_SCALING_1D) { - fp8_blockwise::group_quantize_blockwise_1d(grad_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream, dbias_arg, - workspace_arg); + fp8_blockwise::group_quantize_blockwise_1d( + grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg); } else { - fp8_blockwise::group_quantize_blockwise_2d(grad_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream, dbias_arg, - workspace_arg); + fp8_blockwise::group_quantize_blockwise_2d( + grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg); } break; } diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index 1fd1738f93..33aadb5bfa 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -284,7 +284,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_stride_y, - const float epsilon, const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { + const float epsilon, const bool pow_2_scales, const float* __restrict__ noop_ptr, + float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -378,8 +379,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma for (int w = 1; w < kNumWarps; ++w) { block_amax = fmaxf(block_amax, warp_amaxes[w]); } - const CType scale = - compute_scale_from_types(block_amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(block_amax, epsilon, pow_2_scales); // The 2D colwise per-expert scale offset requires a CTA-cooperative prefix // sum in the VARYING_FIRST_DIM case, so compute it across all threads before @@ -475,7 +475,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t R_total, const float epsilon, - const float* __restrict__ noop_ptr) { + const bool pow_2_scales, const float* __restrict__ noop_ptr) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -530,8 +530,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) CType amax = compute_row_amax(in_vec[it]); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; if (thr_col == 0 && r_global < R_total) { // Per-expert layout: (blocks_X, roundup(M_t, 4)). Compute expert base @@ -572,8 +571,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_t_stride_aligned_K, - const size_t R_total, const float epsilon, const float* __restrict__ noop_ptr, - float* __restrict__ dbias_workspace) { + const size_t R_total, const float epsilon, const bool pow_2_scales, + const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -656,8 +655,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType amax = compute_row_amax(in_vec); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t r_global = global_row_base + row_local; @@ -733,8 +731,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke } amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t c_global = global_col_base + col_local; @@ -814,7 +811,8 @@ inline GroupedBlockwiseLaunchInfo prepare_grouped_blockwise_launch(const Grouped // reports the [total_row_blocks, K] fp32 shape and returns without launching. inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -883,7 +881,7 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_stride_y, epsilon, - noop_ptr, dbias_workspace); + pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS @@ -907,7 +905,8 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso // per-tile column partial can be computed. inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -967,7 +966,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso reinterpret_cast(output->scale_inv.dptr), info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, - info.R_total, epsilon, noop_ptr); + info.R_total, epsilon, pow_2_scales, noop_ptr); } } else if constexpr (kRowwise || kColwise) { // CW-only, BOTH, or RW-only WITH dbias: smem-cached TMA kernel. @@ -998,7 +997,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_t_stride_aligned_K, - info.R_total, epsilon, noop_ptr, dbias_workspace); + info.R_total, epsilon, pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 554d8c1ac9..4d6d24ba65 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -161,10 +161,12 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d * \param[in,out] output Output grouped FP8/MXFP8 tensor. * \param[out] dbias Result of the reduction of the input along columns. * \param[out] workspace Workspace tensor. + * \param[in] quant_config Quantization configuration. * \param[in] stream CUDA stream used for the operation. */ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); /*! \brief Computes backward of GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8a03f2f51a..8c209ace15 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -404,6 +404,12 @@ class Float8BlockScaling(Recipe): NOTE: To relax the default constraint that scales be powers of 2, set env variable NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 to override it for the recipe defaults. + NOTE: FP8 block scaling requires split accumulation for numerical accuracy, so + ``fp8_gemm_fprop``/``fp8_gemm_dgrad``/``fp8_gemm_wgrad`` all fix + ``use_split_accumulator=True`` (enforced in ``__post_init__``). The fused grouped + GEMM path (GroupedLinear) always uses split accumulation for FP8 block scaling and + ignores any caller- or recipe-supplied ``use_split_accumulator`` value. + Parameters ---------- fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.E4M3 diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3b066d50b..d5446e7b61 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -13,6 +13,7 @@ from ..utils import get_sm_count, _empty_tensor from ..quantized_tensor import Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -435,6 +436,30 @@ def _get_fp32_zeros_tensor(num_tensors: int, device: torch.device) -> torch.Tens return torch.zeros(num_tensors, dtype=torch.float32, device=device) +@functools.lru_cache(maxsize=None) +def _get_grouped_gemm_setup_workspace(device: int, num_tensors: int) -> torch.Tensor: + """Persistent setup workspace (per-group pointer/dim arrays) for grouped-tensor GEMM.""" + return torch.empty( + get_grouped_gemm_setup_workspace_size(num_tensors), + dtype=torch.uint8, + device=device, + ) + + +@functools.lru_cache(maxsize=None) +def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: + """Persistent cuBLAS workspace for the grouped-tensor GEMM path, one per GEMM layout. + + Grouped cuBlasLt GEMM kernels in cuBLAS versions <= 13.7 leave behind stale descriptors in the + workspace that cause back-to-back GEMM kernels to crash/deadlock on 2nd CUDA-graph replay. As a + workaround, we allocate a different workspace for each GEMM layout (TN, NN, NT) to avoid + contamination between subsequent GEMM calls (when there is no other graph node between GEMM + kernels). + """ + assert layout in ("TN", "NN", "NT"), f"unexpected grouped GEMM layout {layout}" + return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -474,6 +499,20 @@ def general_grouped_gemm_for_grouped_tensor( if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + def _is_fp8_blockwise(operand) -> bool: + if isinstance(operand, (list, tuple)): + return any(isinstance(t, Float8BlockwiseQTensorStorage) for t in operand) + if isinstance(operand, GroupedTensorStorage): + return isinstance(operand.quantizer, Float8BlockQuantizer) + return False + + if _is_fp8_blockwise(A) or _is_fp8_blockwise(B): + # The fused grouped FP8 block-scaling GEMM only supports split accumulation, + # so force it on and intentionally override any caller-supplied value. This + # matches the Float8BlockScaling recipe, which fixes use_split_accumulator=True + # for all of fprop/dgrad/wgrad, so no user-configurable setting is discarded. + use_split_accumulator = True + if is_discrete_out: # wgrad case. grouped_gemm_impl = tex.te_general_grouped_gemm_for_discrete_out @@ -513,16 +552,12 @@ def general_grouped_gemm_for_grouped_tensor( if not alpha.is_cuda or not beta.is_cuda: raise ValueError("alpha and beta must be CUDA tensors.") - workspace_setup = torch.empty( - get_grouped_gemm_setup_workspace_size(num_tensors), - dtype=torch.uint8, - device=device, - ) - workspace_cublas = torch.empty( - get_cublas_workspace_size_bytes(), - dtype=torch.uint8, - device=device, - ) + workspace_setup = _get_grouped_gemm_setup_workspace(device.index, num_tensors) + # Each grouped-GEMM layout gets its own persistent cuBLAS workspace: two grouped + # GEMMs sharing one workspace can deadlock under CUDA-graph replay (see + # _get_grouped_cublas_workspace). wgrad (NT) is the case seen in TE; fprop (TN) and + # dgrad (NN) have also been reported to conflict, so all three layouts are isolated. + workspace_cublas = _get_grouped_cublas_workspace(device.index, layout) sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8d77a9e349..8b1cd384aa 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -500,6 +500,15 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto quantizer_cpp = convert_quantizer(quantizer); + // MXFP8 carries no quantization knobs; FP8 block-scaling reads scale constraints + // off the quantizer, matching the forward group_quantize dispatch. + QuantizationConfigWrapper quant_config_cpp; + if (detail::IsFloat8BlockwiseQuantizers(quantizer.ptr())) { + auto *fp8_block_quantizer_cpp = static_cast(quantizer_cpp.get()); + quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); + quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + } + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), @@ -530,7 +539,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto stream = at::cuda::getCurrentCUDAStream(); NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); if (workspace_nvte.ndim() > 0 && workspace_nvte.numel() > 0) { at::Tensor workspace_torch = allocateSpace(workspace_nvte.shape(), workspace_nvte.dtype()); @@ -539,7 +549,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); return py::make_tuple(py::reinterpret_borrow(grouped_output_py), py::cast(std::move(dbias_torch))); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4308464c52..17237fa9b8 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1154,14 +1154,6 @@ std::pair Float8BlockQuantizer::create_grouped const size_t logical_last_dim) const { using namespace pybind11::literals; - // The fused grouped FP8 block-scaling path uses unconstrained FP32 scales and does not - // implement power-of-2 scaling. Reject force_pow_2_scales rather than silently ignoring it; - // the unfused per-tensor split-quantize path still honors it. - NVTE_CHECK(!force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support force_pow_2_scales=True. " - "Set force_pow_2_scales=False, or use the unfused split-quantize path " - "(NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0) which supports power-of-2 scales."); - const auto tensor_offsets = resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, logical_first_dim, logical_last_dim); diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index a65ee3b5c3..5b12e4a4f7 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -55,7 +55,13 @@ from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..triton.grouped_dbias_dscales import compute_grouped_dbias -from ..tensor import Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, NVFP4Quantizer +from ..tensor import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -103,11 +109,15 @@ def _is_grouped_tensor_path_supported( and be incompatible with CUDA Graphs. Supported Compute Capability (CC) and precisions: - * Hopper (CC 9.0): BF16/FP16 and FP8 per-tensor current scaling. + * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 + block scaling (1D/2D, including power-of-2 scales). * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8 per-tensor current scaling. - FP8 delayed scaling and FP8 block scaling are not supported because the - corresponding grouped quantization kernels are missing. + FP8 delayed scaling is not supported because the corresponding grouped + quantization kernels are missing. FP8 block scaling on Blackwell (SM100 and + SM110) raises instead of falling back: the fused path is Hopper-only and has + no MXFP8-broadcast emulation. Architectures outside the fused-path window + (e.g. SM120) fall back to the legacy path like every other recipe. Grouped GEMM requires cuBLAS 13.3+ (13.4+ on Hopper, 13.5+ for FP8 per-tensor current scaling on Hopper); otherwise the legacy path is used. Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization @@ -147,6 +157,21 @@ def _is_grouped_tensor_path_supported( if device_capability < (10, 0) and cublaslt_version < 130500: return False return True + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM + # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast + # emulation. On Blackwell (SM100/SM110, the only other arch that reaches + # this branch) fail loudly rather than silently falling back to the + # unfused path the user explicitly opted out of. + if get_device_compute_capability() >= (10, 0): + raise RuntimeError( + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the" + " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped" + " FP8 block-scaling path is Hopper-only. Unset" + " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused" + " path (emulated via MXFP8 GEMM on Blackwell)." + ) + return True # MXFP8 and NVFP4 require Blackwell+. if not (10, 0) <= device_capability <= (11, 0): return False @@ -820,7 +845,12 @@ def _backward_grouped_tensor( columnwise=ctx.weights_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True - if ctx.use_bias and isinstance(grad_output_quantizer, MXFP8Quantizer): + # The grouped FP8 block-scaling bgrad kernel computes dbias in the rowwise + # pass, so the fusion needs rowwise output (i.e. dgrad required). + fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( + isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.requires_dgrad + ) + if ctx.use_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, From 8280f022840b1c885a1d7dce96c171ed9644c1b5 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 22 Jul 2026 13:53:40 -0700 Subject: [PATCH 30/35] Fix nccl-extensions submodule name (#3239) update nccl-ext submodule name Signed-off-by: Phuong Nguyen --- .gitmodules | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitmodules b/.gitmodules index 88d431cba5..c350f75227 100644 --- a/.gitmodules +++ b/.gitmodules @@ -4,6 +4,6 @@ [submodule "3rdparty/cutlass"] path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git -[submodule "3rdparty/nccl"] +[submodule "3rdparty/nccl-extensions"] path = 3rdparty/nccl-extensions url = https://github.com/NVIDIA/nccl-extensions.git From 7b55d302bea1ead6047eac285afa0443c68fbdca Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 22 Jul 2026 14:55:55 -0700 Subject: [PATCH 31/35] [JAX] Schedule EP dispatch/combine on XLA collective stream (#3231) * [JAX] Schedule EP dispatch/combine on XLA collective stream * [JAX] Gate EP collective-stream annotation on JAX/XLA version Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- tests/jax/test_multi_process_ep.py | 50 +++++++++++++++++++++ transformer_engine/jax/cpp_extensions/ep.py | 17 +++++++ transformer_engine/jax/version_utils.py | 11 +++++ 3 files changed, 78 insertions(+) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 0b8bb25f3f..d3effcdc07 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -36,6 +36,7 @@ ep_combine_fwd, get_ep_config, ) +from transformer_engine.jax.version_utils import is_collective_stream_supported # ── Test config ───────────────────────────────────────────────────────────── @@ -660,6 +661,55 @@ def run(idx, toks, w): expected = (("dp", "ep"),) if self.dp > 1 else ("ep",) self.assertEqual(tuple(compiled.output_shardings.spec), expected) + @unittest.skipUnless( + is_collective_stream_supported(), + "JAX/XLA lacks the gpu_stream:collective annotation (openxla/xla#39604)", + ) + def test_z_dispatch_combine_on_collective_stream(self): + """Every EP FFI custom call must carry the collective-stream annotation + so XLA schedules them on the collective stream instead of overlapping + them with other collectives.""" + T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() + dp_spec = PartitionSpec(("dp", "ep"), None) + ep_spec_3d = PartitionSpec(("dp", "ep"), None, None) + ep_spec_2d = PartitionSpec(("dp", "ep"), None) + + with self.mesh, global_shard_guard(self.mr): + + @jax.jit + def run(idx, toks, w): + idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) + toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) + w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) + recv_t, recv_w, hm, tc = ep_dispatch( + self.hk, idx, toks, w, self.recv_capacity_per_rank + ) + recv_t = jax.lax.with_sharding_constraint( + recv_t, NamedSharding(self.mesh, ep_spec_3d) + ) + recv_w = jax.lax.with_sharding_constraint( + recv_w, NamedSharding(self.mesh, ep_spec_2d) + ) + weighted = self._preweight_expert_out(recv_t, recv_w) + out = ep_combine(self.hk, hm, tc, weighted, T_dp, out_sharding=(("dp", "ep"), None)) + return jax.lax.with_sharding_constraint(out, NamedSharding(self.mesh, dp_spec)) + + hlo = run.lower(topk_idx, tokens, topk_w).compile().as_text() + + # Every te_ep_* FFI custom call must carry the collective-stream + # annotation so XLA places it on the collective stream. + ep_lines = [l for l in hlo.splitlines() if 'custom_call_target="te_ep_' in l] + self.assertTrue(ep_lines, f"no te_ep_* custom calls in compiled HLO:\n{hlo}") + missing = [ + l.strip()[:200] + for l in ep_lines + if '_xla_stream_annotation="collective"' not in l.replace(" ", "") + ] + self.assertFalse( + missing, + "te_ep_* custom calls missing collective-stream annotation:\n" + "\n".join(missing), + ) + def test_z_no_unexpected_reshard_in_hlo_bwd(self): """Compiled bwd HLO must not insert XLA collectives outside the EP FFI.""" T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 77e60afbcd..806e7ae480 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -24,6 +24,18 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive from ..sharding import global_mesh_resource, get_mesh_axis_size +from ..version_utils import is_collective_stream_supported + + +def _on_collective_stream(func): + """Pin ``func``'s ops to XLA's collective stream so the scheduler serializes + them with native collectives. No-op on JAX that lacks the annotation.""" + if not is_collective_stream_supported(): + return func + from jax.experimental.compute_on import compute_on + + return compute_on("gpu_stream:collective")(func) # pylint: disable=not-callable + __all__ = [ "EpConfig", @@ -894,6 +906,7 @@ def shardy_sharding_rule(*args): # ── Public-ish helpers (used by jax/ep.py) ────────────────────────────────── +@_on_collective_stream def ep_prepare(cfg: EpLayerConfig, topk_idx): """Exchange routing metadata for ``cfg``; return ``(token_counts, handle_mem)``.""" return EpPreparePrimitive.outer_primitive.bind( @@ -904,6 +917,7 @@ def ep_prepare(cfg: EpLayerConfig, topk_idx): ) +@_on_collective_stream def ep_dispatch_fwd( cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ): @@ -920,6 +934,7 @@ def ep_dispatch_fwd( ) +@_on_collective_stream def ep_combine_fwd( cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None ): @@ -935,6 +950,7 @@ def ep_combine_fwd( ) +@_on_collective_stream def ep_dispatch_bwd( cfg: EpLayerConfig, handle_mem, @@ -956,6 +972,7 @@ def ep_dispatch_bwd( ) +@_on_collective_stream def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank): """Backward of combine; returns grad_expert_out [num_procs, recv_capacity_per_rank, H].""" return EpCombineBwdPrimitive.outer_primitive.bind( diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e4619d8670..8765a4fce0 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -64,6 +64,16 @@ def is_triton_autotuned_alias_safe() -> bool: return v >= PkgVersion(_TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR) +# XLA gained the ``gpu_stream:collective`` stream annotation in openxla/xla#39604, +# which ships in the JAX 0.10.0 release. Older XLA fatally fails on it. +_COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.0" + + +def is_collective_stream_supported() -> bool: + """Return True if the installed JAX supports the gpu_stream:collective annotation.""" + return jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION) + + def is_triton_extension_supported() -> bool: """Return True if the current JAX version supports Triton kernel dispatch. @@ -77,6 +87,7 @@ def is_triton_extension_supported() -> bool: __all__ = [ "jax_version_meet_requirement", "is_triton_autotuned_alias_safe", + "is_collective_stream_supported", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION", From c9a1f5a9307a658dce9c5bdec6f7d7d71d5afe2e Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Thu, 23 Jul 2026 06:11:12 +0800 Subject: [PATCH 32/35] Generalized Tensor Parallelism (GTP) (#3005) * Generalized Tensor Parallelism (GTP) init commit Co-authored-by: Jieming Zhang Signed-off-by: Shiqing Fan * GTP + gmm fusion Signed-off-by: Shiqing Fan * [fix] Respect per-op activation-offload markers in fused grouped MLP Signed-off-by: Shiqing Fan * Code clean: rename GTP weight-sharding axis to gtp_remat Signed-off-by: Shiqing Fan * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert "[fix] Respect per-op activation-offload markers in fused grouped MLP" This reverts commit 8bb26f047f690bb8a7e5884a1f9bf419a34fe002. Signed-off-by: Shiqing Fan * Make TE GTP-agnostic at construction Signed-off-by: Shiqing Fan * GTP+nvfp4: fix GTP backward GEMM scaling-mode mismatch for bf16-gathered weights Signed-off-by: Shiqing Fan * Make TE runtime GTP-agnostic via a DistributedWeight protocol Signed-off-by: Shiqing Fan * Code clean - Take a single leader weight in the DistributedWeight dispatchers - Gather the FC2 grouped weight late in the fused grouped MLP Signed-off-by: Shiqing Fan * Simplify the NVFP4 gather post-process; Materialize the EGTP FC1 weight before the NVFP4 dgrad dispatch Signed-off-by: Shiqing Fan * Code clean - Rename gather coalescing flag grouped -> external_coalescing; - Clean up DistributedWeight wiring in TE modules - Restructure _all_gather_nvfp4 Signed-off-by: Shiqing Fan * fix comments Signed-off-by: Shiqing Fan * Support DistributedWeight in the fusible grouped-linear ops path - Add a self-contained dispatch test with a fake DistributedWeight implementer Signed-off-by: Shiqing Fan * Unify distributed-weight wgrad finalize to return a graph-safe dummy - `finalize_weight_grads` now accepts a weight list or a bare leader, mirroring materialize_weight_for_backward; - Centralize the in-place / dummy / async-None finalize contract in DistributedWeight.finalize_group_grads and delegate the dispatcher docstring to it. Signed-off-by: Shiqing Fan --------- Signed-off-by: Shiqing Fan Co-authored-by: Jieming Zhang Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- qa/L0_pytorch_unittest/test.sh | 2 + tests/pytorch/test_distributed_weight.py | 172 ++++++++++++++++++ ...t_ops_grouped_linear_distributed_weight.py | 131 +++++++++++++ transformer_engine/pytorch/distributed.py | 94 +++++++--- .../pytorch/distributed_weight.py | 117 ++++++++++++ .../pytorch/module/grouped_linear.py | 48 ++++- .../pytorch/module/layernorm_linear.py | 33 +++- transformer_engine/pytorch/module/linear.py | 58 +++++- .../pytorch/ops/basic/grouped_linear.py | 64 +++++-- .../pytorch/ops/fused/grouped_mlp.py | 61 ++++++- 10 files changed, 720 insertions(+), 60 deletions(-) create mode 100644 tests/pytorch/test_distributed_weight.py create mode 100644 tests/pytorch/test_ops_grouped_linear_distributed_weight.py create mode 100644 transformer_engine/pytorch/distributed_weight.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 91d3be63d0..5d767ba4d1 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -45,6 +45,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml $TE_PATH/tests/pytorch/test_multi_tensor.py || test_fail "test_multi_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py" @@ -66,6 +67,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py" +PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py" if [ "$RET" -ne 0 ]; then diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py new file mode 100644 index 0000000000..69b5917d90 --- /dev/null +++ b/tests/pytorch/test_distributed_weight.py @@ -0,0 +1,172 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for the GTP-agnostic DistributedWeight protocol and dispatchers. + +These tests exercise TE's weight-parallelism extension point in isolation, with +a tiny in-repo ``FakeDistributedWeight`` stub standing in for a real implementer +(e.g. Megatron's GTPShardedParam). No GPU, process group, or Megatron import is +required: the whole contract lives in TE and is verified against the fake. +""" + +import pytest +import torch + +from transformer_engine.pytorch.distributed_weight import ( + DistributedWeight, + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) + + +class FakeDistributedWeight(torch.Tensor): + """Minimal DistributedWeight implementer for dispatcher tests. + + Records how it was called and returns marker tensors so the dispatcher's + behavior (delegation, list normalization, no-op fallback) is observable. + """ + + is_distributed_weight = True + + def __new__(cls, group_size=1): + t = torch.zeros(1).as_subclass(cls) + t.group_size = group_size + t.calls = [] + return t + + def materialize_group_for_forward(self): + self.calls.append("fwd") + out = [torch.full((2, 2), float(i)) for i in range(self.group_size)] + # Match the real GTP contract: single weight returns a bare tensor. + return out if self.group_size > 1 else out[0] + + def materialize_group_for_backward(self, **kwargs): + self.calls.append(("bwd", kwargs)) + out = [torch.full((2, 2), float(10 + i)) for i in range(self.group_size)] + return out if self.group_size > 1 else out[0] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls.append(("finalize", wgrads)) + wl = wgrads if isinstance(wgrads, (list, tuple)) else [wgrads] + out = [w + 100 for w in wl] + return out if self.group_size > 1 else out[0] + + def grad_buffer(self): + return torch.full((2, 2), -1.0) + + +class FakeNonTensorWeight: + """DistributedWeight-shaped object that is NOT a torch.Tensor (contract violation).""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + return torch.zeros(2, 2) + + def materialize_group_for_backward(self, **kwargs): + return torch.zeros(2, 2) + + def finalize_group_grads(self, wgrads, **kwargs): + return wgrads + + def grad_buffer(self): + return torch.zeros(2, 2) + + +def test_protocol_runtime_checkable(): + """A conforming object passes isinstance; a plain tensor does not.""" + assert isinstance(FakeDistributedWeight(), DistributedWeight) + assert not isinstance(torch.zeros(2), DistributedWeight) + + +def test_is_distributed_weight(): + assert is_distributed_weight(FakeDistributedWeight()) + assert not is_distributed_weight(torch.zeros(2)) + assert not is_distributed_weight(torch.nn.Parameter(torch.zeros(2))) + + +def test_non_tensor_implementer_rejected(): + """Implementers must be torch.Tensor subclasses; a non-Tensor fails loudly.""" + with pytest.raises(TypeError, match="torch.Tensor subclass"): + is_distributed_weight(FakeNonTensorWeight()) + + +def test_forward_noop_on_plain_tensor(): + """Plain weights pass through unchanged — the critical non-regression.""" + w = torch.nn.Parameter(torch.randn(4, 4)) + out = materialize_weight_for_forward(w) + assert out == [w] + assert out[0] is w + + +@pytest.mark.parametrize("group_size", [1, 3]) +def test_forward_dispatches(group_size): + """Linear (N=1) and GroupedLinear (N=k): one coalesced call, full list returned.""" + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_forward(w) + assert isinstance(out, list) and len(out) == group_size + # Leader is delegated to exactly once (coalesced), not once per weight. + assert w.calls == ["fwd"] + + +def test_forward_accepts_weight_list(): + """The dispatcher accepts the full per-expert list; the leader (index 0) coalesces it.""" + leader = FakeDistributedWeight(group_size=3) + followers = [torch.zeros(2, 2), torch.zeros(2, 2)] + out = materialize_weight_for_forward([leader, *followers]) + assert isinstance(out, list) and len(out) == 3 + # Leader delegated exactly once; the follower entries are not materialized separately. + assert leader.calls == ["fwd"] + + +def test_forward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + out = materialize_weight_for_forward(ws) + assert out == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_backward_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_backward(w) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_accepts_weight_list(): + """Backward dispatcher also accepts the full per-expert list; the leader coalesces it.""" + leader = FakeDistributedWeight(group_size=2) + out = materialize_weight_for_backward([leader, torch.zeros(2, 2)]) + assert isinstance(out, list) and len(out) == 2 + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_noop_on_plain_tensor(): + plain = torch.zeros(2) + assert materialize_weight_for_backward(plain) == [plain] + + +def test_backward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + assert materialize_weight_for_backward(ws) == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_finalize_grads_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + wgrads = [torch.zeros(2, 2) for _ in range(group_size)] + out = finalize_weight_grads(w, wgrads) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 100.0)) + + +def test_finalize_grads_noop_on_plain_tensor(): + """No-op path leaves the grads untouched.""" + plain_w = torch.nn.Parameter(torch.zeros(2)) + g = [torch.ones(2)] + assert finalize_weight_grads(plain_w, g) == g diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py new file mode 100644 index 0000000000..700be4138e --- /dev/null +++ b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py @@ -0,0 +1,131 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DistributedWeight dispatch in the fusible ``ops.GroupedLinear`` (unfused fallback path). + +TE ships no DistributedWeight implementer (GTP etc. live in the caller), so this validates the +*dispatch wiring* with an in-repo fake. The fake applies a DISTINCT, observable scale in each +materialize hook so every plumbing point is independently checked -- a wiring bug fails a specific +assertion: + + * ``materialize_group_for_forward`` scales the weights by ``FWD_SCALE`` -> the fwd GEMM must use + the materialized weights, so ``out == FWD_SCALE * plain_out``. + * ``materialize_group_for_backward`` scales by ``BWD_SCALE`` -> dgrad must use the + re-materialized weights, so ``dgrad == BWD_SCALE * plain_dgrad``. + * ``finalize_group_grads`` reduce-scatters the wgrads into ``main_grad`` in-place and returns a + dummy -> the real wgrad lands in ``main_grad`` and the ops path returns a throwaway ``.grad`` + (it discards finalize's return; see DistributedWeight.finalize_group_grads). + +Real all-gather / reduce-scatter math is exercised by the caller's distributed tests; here the fake +is single-process and only proves the ops.GroupedLinear integration routes weights/grads through +the hooks (and does not bypass them). +""" + +import pytest +import torch + +import transformer_engine.pytorch as te + +# Distinct powers of two so each scale commutes exactly through the (possibly TF32) GEMM rounding, +# making the linear scale relations bit-exact under a tight tolerance. +FWD_SCALE = 2.0 +BWD_SCALE = 4.0 + + +class _FakeDistWeight(torch.nn.Parameter): + """Single-process fake DistributedWeight leader for dispatch testing. + + Each materialize hook applies its own scale (see module docstring) so fwd/bwd routing is + observable; ``finalize_group_grads`` models the real main-grad contract -- reduce-scatter (here + an identity accumulate) the wgrads into each shard's ``main_grad`` in-place, flag + ``grad_added_to_main_grad``, and return a dummy that the ops path discards. + """ + + is_distributed_weight = True + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return [w * FWD_SCALE for w in self._group] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return [w * BWD_SCALE for w in self._group] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + for w, g in zip(self._group, wl): + w.main_grad.add_(g.to(w.main_grad.dtype)) # in-place accumulate into main_grad + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] # dummy grads (real value is now in main_grad) + + def grad_buffer(self): + return self.data + + +def _make_fake_dist_leader(op, num_gemms): + """Replace ``op.weight0`` with a fake distributed leader referencing the whole group.""" + w0 = op.weight0 + leader = _FakeDistWeight(w0.data) + leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + leader._group = [leader] + [getattr(op, f"weight{i}") for i in range(1, num_gemms)] + op.weight0 = leader + return leader + + +@pytest.mark.parametrize("num_gemms", [2, 4]) +def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms): + """Every DistributedWeight hook must be routed through the GEMM flow (and not bypassed). + + fwd/bwd use the scaled materialized weights; finalize reduce-scatters the wgrad into + ``main_grad`` in-place, and the ops path returns a throwaway dummy ``.grad``. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.manual_seed(0) + # fp32 with power-of-two scales: each scale commutes exactly through the GEMM rounding (even + # TF32), so the linear scale relations below are bit-exact under a tight tolerance. + in_f, out_f, total_tokens = 32, 64, num_gemms * 8 + dtype, device = torch.float32, "cuda" + + op = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference.load_state_dict(op.state_dict()) + + leader = _make_fake_dist_leader(op, num_gemms) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + w.main_grad = torch.zeros((out_f, in_f), dtype=torch.float32, device=device) + w.grad_added_to_main_grad = False # DDP initializes this on every param + + m_splits = [total_tokens // num_gemms] * num_gemms + m_splits[-1] += total_tokens - sum(m_splits) + split_sizes = torch.tensor(m_splits, dtype=torch.int64, device=device) + + x = torch.randn(total_tokens, in_f, dtype=dtype, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + out = op(x, split_sizes) + out.sum().backward() + ref_out = reference(ref_x, split_sizes) + ref_out.sum().backward() + + # All three dispatch hooks actually fired. + assert leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 + + tols = dict(rtol=1e-5, atol=1e-5) + # fwd used the materialized (FWD_SCALE) weights. + torch.testing.assert_close(out, FWD_SCALE * ref_out, **tols) + # dgrad used the re-materialized (BWD_SCALE) weights. + torch.testing.assert_close(x.grad, BWD_SCALE * ref_x.grad, **tols) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + ref_w = getattr(reference, f"weight{i}") + # The distributed op's grad of record is main_grad (finalize reduce-scattered the identity + # wgrad there); compare it to the plain reference module's ordinary autograd .grad. + torch.testing.assert_close(w.main_grad.to(dtype), ref_w.grad, **tols) + # main_grad is flagged, and op's own .grad is a discarded dummy (not the real grad, which + # lives in main_grad). Nobody writes the real grad into op.weight.grad by design. + assert w.grad_added_to_main_grad is True + assert w.grad is not None, f"weight{i} should receive a dummy .grad" diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 569335d93f..d1525b53f0 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Iterable -from contextlib import contextmanager, AbstractContextManager, ContextDecorator +from contextlib import contextmanager, AbstractContextManager, ContextDecorator, nullcontext from functools import lru_cache from dataclasses import dataclass import math @@ -926,7 +926,10 @@ def fork(self, name: str = "model-parallel-rng"): def reduce_scatter_along_first_dim( - inp: torch.Tensor, tp_group: dist_group_type, async_op: bool = False + inp: torch.Tensor, + tp_group: dist_group_type, + async_op: bool = False, + output: torch.Tensor = None, ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: """Reduce-scatter the input tensor across model parallel group.""" world_size = get_distributed_world_size(tp_group) @@ -944,7 +947,8 @@ def reduce_scatter_along_first_dim( dim_size[0] = dim_size[0] // world_size - output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) + if output is None: + output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) handle = torch.distributed.reduce_scatter_tensor( output, inp.contiguous(), group=tp_group, async_op=async_op ) @@ -1311,7 +1315,8 @@ def wait(self) -> None: """Wait for the async operation to complete and post-process the tensor.""" if self._synchronized: return - self.async_handle.wait() + if self.async_handle is not None: + self.async_handle.wait() _post_process_nvfp4_gather( self.output, self.columnwise_data_interleaved, @@ -1328,6 +1333,8 @@ def _all_gather_nvfp4( async_op: bool = False, quantizer: NVFP4Quantizer, out_shape: Optional[list[int]] = None, + output_tensor=None, + external_coalescing=False, ) -> tuple[NVFP4TensorStorage, Optional[torch.distributed.Work]]: """All-gather NVFP4 tensor along first dimension.""" @@ -1404,15 +1411,23 @@ def _all_gather_nvfp4( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct NVFP4 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) # Coalesce NCCL collectives for gathering data and scale inverses. - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as gather_coalescing_manager: + if not external_coalescing: + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather NVFP4 data for row-wise usage if quantizer.rowwise_usage: @@ -1493,10 +1508,10 @@ def _all_gather_nvfp4( # Transfer amax to output. out._amax_columnwise = inp._amax_columnwise - handle = gather_coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None # Fixes interleaved data for transposed tensor/scale inv and pads scale inv if needed. - if async_op and quantizer.columnwise_usage: + if (async_op or external_coalescing) and quantizer.columnwise_usage: handle = _NVFP4AllGatherAsyncHandle( out, out_columnwise_data, out_scale_inv, world_size, handle ) @@ -1513,6 +1528,8 @@ def _all_gather_mxfp8( async_op: bool = False, quantizer: MXFP8Quantizer, out_shape: Optional[list[int]] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[MXFP8TensorStorage, Optional[torch.distributed.Work]]: """All-gather MXFP8 tensor along first dimension.""" @@ -1578,15 +1595,23 @@ def _all_gather_mxfp8( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct MXFP8 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - # Coalesce NCCL collectives - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as coalescing_manager: + if not external_coalescing: + # Coalesce NCCL collectives for gathering data and scale inverses. + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather MXFP8 data for row-wise usage if quantizer.rowwise_usage: @@ -1633,7 +1658,7 @@ def _all_gather_mxfp8( group=process_group, ) - handle = coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None return out, handle @@ -1642,9 +1667,17 @@ def gather_along_first_dim( process_group: dist_group_type, async_op: bool = False, quantizer: Optional[Quantizer] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[torch.Tensor, Optional[torch.distributed.Work]]: """ All-gather tensors and concatenate along first dimension. + + ``external_coalescing``: composability flag for callers that batch several gathers into + one outer ``torch.distributed._coalescing_manager``. Coalescing managers cannot nest, so + when set this call skips opening its own manager and defers any post-gather fixup + (e.g. NVFP4 columnwise de-interleave) into the returned handle; ``handle.wait()`` completes + it once the outer manager has closed. Leave ``False`` for standalone gathers. """ # Return immediately if no communication is required @@ -1732,6 +1765,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # NVFP4 case @@ -1746,6 +1781,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # High-precision communication for quantized tensors @@ -1775,19 +1812,20 @@ def gather_along_first_dim( inp = inp.dequantize() # Communication for plain PyTorch tensors - out = torch.empty( - out_shape, - dtype=inp.dtype, - device=inp.device, - memory_format=torch.contiguous_format, - ) + if output_tensor is None: + output_tensor = torch.empty( + out_shape, + dtype=inp.dtype, + device=inp.device, + memory_format=torch.contiguous_format, + ) handle = torch.distributed.all_gather_into_tensor( - out, + output_tensor, inp.contiguous(), group=process_group, async_op=async_op, ) - return out, handle + return output_tensor, handle # Global cache to store symmetric memory tensors diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py new file mode 100644 index 0000000000..0d491430e5 --- /dev/null +++ b/transformer_engine/pytorch/distributed_weight.py @@ -0,0 +1,117 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""GTP-agnostic weight-parallelism extension point. + +TE owns this contract but ships no implementation; the caller (e.g. Megatron GTP) implements the +protocol on the weight and injects it at construction. Dispatchers are list-shaped (Linear -> 1, +GroupedLinear -> N; leader is ``weights[0]``) and no-op on plain tensors. +""" + +from typing import Any, List, Protocol, runtime_checkable + +import torch + +__all__ = [ + "DistributedWeight", + "is_distributed_weight", + "materialize_weight_for_forward", + "materialize_weight_for_backward", + "finalize_weight_grads", +] + + +@runtime_checkable +class DistributedWeight(Protocol): + """Structural interface for a custom-weight-parallel weight (AG for the GEMM, reduce/RS the + grad, re-materialize in backward). Duck-typed ``typing.Protocol``: implementers need not + subclass it, and all state (shards, group, async handles) lives outside TE on the implementer. + + Implementers MUST be ``torch.Tensor`` subclasses (needed by ``ctx.save_for_backward``, DDP + backward hooks, and ``torch.compile``); enforced at runtime by :func:`is_distributed_weight`. + """ + + # Capability marker: True on an implementer, absent on plain tensors; TE's fwd/bwd gate on it. + is_distributed_weight: bool + + def materialize_group_for_forward(self) -> Any: + """Return the tensor(s) to feed the forward GEMM (may all-gather shards).""" + + def materialize_group_for_backward(self) -> Any: + """Re-materialize the full weight(s) for the backward GEMMs.""" + + def finalize_group_grads(self, wgrads: Any) -> Any: + """Post-process freshly computed weight grad(s) (e.g. reduce-scatter). + + May consume ``wgrads`` in-place -- reduce-scatter into ``main_grad`` and set + ``grad_added_to_main_grad`` -- returning a dummy grad (or ``None`` for an async collective) + that callers use as the parameter grad(s) or discard. + """ + + def grad_buffer(self) -> torch.Tensor: + """The gradient accumulation buffer for this weight.""" + + +def is_distributed_weight(weight: Any) -> bool: + """True if ``weight`` participates in custom weight parallelism (False on plain tensors). + + Enforces the :class:`DistributedWeight` requirement that an implementer be a ``torch.Tensor`` + subclass, failing loudly here rather than silently breaking autograd downstream. + """ + flag = bool(getattr(weight, "is_distributed_weight", False)) + if flag and not isinstance(weight, torch.Tensor): + raise TypeError( + "DistributedWeight implementers must be torch.Tensor subclasses; got " + f"{type(weight).__name__}." + ) + return flag + + +def materialize_weight_for_forward(weights: Any) -> List[Any]: + """Prepare the weight(s) fed to the forward GEMM, always returned as a list. + + Args: + weights: the module's weight(s) -- a single weight (Linear) or the full per-expert list + (GroupedLinear). A bare weight is treated as a one-element list. + + Returns: + - Distributed group: the leader ``weights[0]`` all-gathers/coalesces the whole group and + returns all N materialized weights; the follower entries ``weights[1:]`` are ignored + here (the leader already holds references to its group). + - Otherwise: the input weights, unchanged. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_forward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def materialize_weight_for_backward(weights: Any) -> List[Any]: + """Backward-GEMM mirror of :func:`materialize_weight_for_forward` (same contract).""" + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_backward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def finalize_weight_grads(weights: Any, wgrads: List[Any]) -> List[Any]: + """Finalize a weight group's grad(s), mirroring :func:`materialize_weight_for_backward`. + + Delegates to the leader's :meth:`DistributedWeight.finalize_group_grads` (which defines the + in-place / dummy / async-``None`` return contract); returns ``wgrads`` unchanged when not + distributed. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(wgrads) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 5b12e4a4f7..32963c07ca 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -46,6 +46,12 @@ is_fp8_activation_recompute_enabled, in_fp8_activation_recompute_phase, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, @@ -524,6 +530,11 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad + origin_weights = weights + is_dist_weight = is_distributed_weight(weights[0]) + if is_dist_weight: + weights = materialize_weight_for_forward(weights) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -734,6 +745,10 @@ def forward( if backward_override == "high_precision" and inp.requires_grad else [None] * num_gemms ) + if is_dist_weight: + # GTP: gathered workspace is transient (re-gathered in backward), don't save it. + weights_fp8 = [None] * num_gemms + saved_weights = origin_weights tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, @@ -761,6 +776,8 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + elif is_dist_weight: + ctx.main_grad_funcs = [origin_weights[i].grad_buffer for i in range(num_gemms)] else: ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) @@ -1024,7 +1041,12 @@ def backward( # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + is_dist_weight = is_distributed_weight(saved_weights[0]) + if is_dist_weight: + origin_weights = saved_weights + if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] @@ -1085,13 +1107,18 @@ def backward( ctx.m_splits, ) - if ctx.is_first_microbatch is not None: + if is_dist_weight: + accumulate_wgrad_into_param_main_grad = False + elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch ) else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation + if is_dist_weight: + weights = materialize_weight_for_backward(origin_weights) + if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD if ctx.fp8 or ctx.debug: @@ -1162,6 +1189,9 @@ def backward( device=ctx.device, ) wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] + if is_dist_weight: + # Gathered weights are no longer needed after dgrad GEMM. + del weights if ctx.save_original_input: inp = inputmats[0] @@ -1212,7 +1242,8 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + if not is_dist_weight + and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), ) @@ -1254,10 +1285,13 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - wgrad_list = [ - handle_custom_ddp_from_mcore(weight, main_grad, wgrad) - for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) - ] + if is_dist_weight: + wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) + else: + wgrad_list = [ + handle_custom_ddp_from_mcore(weight, main_grad, wgrad) + for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) + ] else: wgrad_list = [None] * ctx.num_gemms diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 43799b003c..cffc5ba1e6 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -56,6 +56,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing @@ -310,6 +316,11 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + origin_weight = weight + is_dist_weight = is_distributed_weight(origin_weight) + if is_dist_weight: + weight = materialize_weight_for_forward(weight)[0] + out_features = weight.shape[0] new_weight_workspace = None weightmat = weight is_weight_param_quantized = False @@ -500,10 +511,15 @@ def forward( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the gathered quantized workspace; + # backward re-gathers from the saved (sharded) weight and re-quantizes. + if is_dist_weight: + wt_save = None + tensors_to_save, tensor_objects = prepare_for_saving( inputmat, wt_save, - weight, + origin_weight, bias, ln_weight, ln_out_to_save, @@ -530,6 +546,8 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad + elif is_dist_weight: + ctx.main_grad_func = origin_weight.grad_buffer else: ctx.main_grad_func = lambda: weight.main_grad ctx.grad_input_quantizer = grad_input_quantizer @@ -623,6 +641,9 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) + is_dist_weight = is_distributed_weight(saved_weight) + if is_dist_weight: + weight = materialize_weight_for_backward(saved_weight)[0] # Restore from weakref to get original weight python object # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) # Only needed when fuse_wgrad_accumulation is enabled. @@ -640,7 +661,7 @@ def backward( ), "weight was removed while fuse_wgrad_accumulation=True" # Since main_grad can be modified inplace, it should not be a part of saved_tensors main_grad = ctx.main_grad_func() if weight is not None else None - if main_grad is not None: + if main_grad is not None and not is_dist_weight: origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -984,7 +1005,10 @@ def backward( use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch ) @@ -1056,6 +1080,9 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 78a4d31852..cef6d3d533 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -55,6 +55,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_gemm, ) @@ -265,6 +271,7 @@ def _linear_forward_impl( """ weight = args.weight + is_dist_weight = is_distributed_weight(args.weight) inp = args.inp bias = args.bias input_quantizer = args.input_quantizer @@ -413,6 +420,14 @@ def _linear_forward_impl( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + # Distributed weight (e.g. GTP): rebind `weight` to the all-gathered tensor; + # `args.weight` keeps the sharded-param reference for backward re-gather / grad + # finalize. No-op for a plain weight. + if is_dist_weight: + weight = materialize_weight_for_forward(args.weight)[0] + # Refresh out_features from the gathered weight (captured sharded above, pre-gather). + out_features = weight.shape[0] + new_weight_workspace = None weightmat = weight if fp8 or debug: @@ -599,6 +614,9 @@ def _linear_forward_impl( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the workspace; backward re-gathers it. + if is_dist_weight: + wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. @@ -704,6 +722,8 @@ def _linear_setup_ctx( bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): bwd_args.main_grad_func = weight.get_main_grad + elif is_distributed_weight(weight): + bwd_args.main_grad_func = weight.grad_buffer else: bwd_args.main_grad_func = lambda: weight.main_grad @@ -749,6 +769,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. inputmat = args.inputmat weight_fp8 = args.weight_fp8 saved_weight = args.saved_weight + is_dist_weight = is_distributed_weight(saved_weight) bias = args.bias input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer @@ -793,7 +814,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - origin_weight_python_object.main_grad = main_grad + if not is_dist_weight: + origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already @@ -963,6 +985,12 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. dgrad = None dgrad_work = None + + # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when + # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd. + if is_dist_weight: + weight_fp8 = materialize_weight_for_backward(saved_weight)[0] + if bwd_args.requires_dgrad: # FSDP2: Re-create workspace from all-gathered weight when @@ -977,6 +1005,16 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. elif bwd_args.weight_quantizer is not None: bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) + elif ( + is_dist_weight + and bwd_args.fp8 + and bwd_args.weight_quantizer is not None + and not isinstance(weight_fp8, QuantizedTensorStorage) + ): + # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer + # so the dgrad operand isn't cast by the delayed recipe. + bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight_fp8 = bwd_args.weight_quantizer(weight_fp8) # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): @@ -1163,7 +1201,10 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = bwd_args.wgrad_use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if bwd_args.is_first_microbatch is not None: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif bwd_args.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( bwd_args.fuse_wgrad_accumulation and not bwd_args.is_first_microbatch ) @@ -1239,6 +1280,11 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) + # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad + # (async; overlap with the next layer's bwd via the cascade). + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ @@ -1287,15 +1333,19 @@ def wgrad_gemm( origin_weight_python_object, "grad_added_to_main_grad" ): origin_weight_python_object.grad_added_to_main_grad = True + # Use the param's local shape (sharded under GTP) so the dummy wgrad + # matches the saved weight shape; main_grad_func() under GTP returns + # an unsharded scratch and would otherwise mismatch. + wgrad_shape = list(origin_weight_python_object.shape) if getattr(origin_weight_python_object, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, ) elif bwd_args.fuse_wgrad_accumulation: diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5ef0fa4339..af3f8b5930 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -52,6 +52,12 @@ view_main_grad_as_grouped_buffer, ) from ..op import BasicOperation, OperationContext +from ...distributed_weight import ( + finalize_weight_grads, + is_distributed_weight, + materialize_weight_for_backward, + materialize_weight_for_forward, +) from ...tensor import GroupedTensor, GroupedTensorStorage from ...triton.grouped_dbias_dscales import ( compute_grouped_dbias, @@ -905,6 +911,26 @@ def _get_weight_tensors(self) -> list[torch.nn.Parameter]: return [self.weight] return [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + def _forward_weight_list(self) -> list[torch.Tensor]: + """Per-expert forward weights, materialized (all-gathered) when distributed.""" + weights = [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + if is_distributed_weight(weights[0]): + weights = materialize_weight_for_forward(weights) + return weights + + def _backward_weight_setup(self): + """Return ``(origin_weights, is_dist_weight, dgrad_weights)``; dgrad weights are the + re-materialized (all-gathered) weights when distributed, else ``None``.""" + origin_weights = self._get_weight_tensors() + is_dist_weight = is_distributed_weight(origin_weights[0]) + dgrad_weights = materialize_weight_for_backward(origin_weights) if is_dist_weight else None + return origin_weights, is_dist_weight, dgrad_weights + + def _is_distributed_weight(self) -> bool: + """Whether this op's weights are distributed (materialized per fwd/bwd, not saved).""" + leader = self.weight if self.single_grouped_weight else self.weight0 + return is_distributed_weight(leader) + def _get_grouped_bias_for_gemm( self, dtype: torch.dtype, @@ -1158,7 +1184,7 @@ def _fuser_forward_split_quantize( if weights is None: weights = self.weight.split_into_quantized_tensors() else: - weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] + weights = self._forward_weight_list() # materialized when distributed bs = None if has_bias: bs = self._get_bias_tensors(dtype) @@ -1213,7 +1239,8 @@ def _fuser_forward_split_quantize( out_splits[i].add_(bs[i].unsqueeze(0) * scales_splits[i].unsqueeze(-1)) # Prepare weight tensors for backward pass - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so we never save the gathered weight + if not input_requires_grad or self._is_distributed_weight(): ws = [None] * num_groups elif with_quantized_compute: for w, weight_param in zip(ws, weights): @@ -1306,7 +1333,7 @@ def _fuser_forward_grouped_tensor( else: # Discrete weights grouped_weights = self._get_discrete_weights_for_gemm( - [getattr(self, f"weight{idx}") for idx in range(num_groups)], + self._forward_weight_list(), weight_quantizers, columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, @@ -1349,7 +1376,8 @@ def _fuser_forward_grouped_tensor( bias_scale=bias_scale, ) - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so never save the gathered weight. + if not input_requires_grad or self._is_distributed_weight(): grouped_weights = None if self.single_grouped_weight else [None] * num_groups if not weight_requires_grad: @@ -1409,7 +1437,7 @@ def _fuser_backward_split_quantize( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device # Saved tensors from forward pass. Layout: @@ -1461,7 +1489,8 @@ def _fuser_backward_split_quantize( grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] # Initialize grad weight buffers. - accumulate_into_main_grad = self._accumulate_into_main_grad + # Distributed weights reduce their own grads (finalize); never accumulate into main_grad. + accumulate_into_main_grad = self._accumulate_into_main_grad and not is_dist_weight grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups @@ -1509,7 +1538,7 @@ def _fuser_backward_split_quantize( getattr(ctx, "dgrad_out", None), in_shape, ctx.dtype, device ) general_grouped_gemm( - ws, + dist_dgrad_weights if is_dist_weight else ws, dys, [grad_input], [None] * num_groups, # quantization_params @@ -1554,11 +1583,16 @@ def _fuser_backward_split_quantize( if not delay_wgrad: clear_tensor_data(*xs) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, grad_weights) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups @@ -1592,7 +1626,7 @@ def _fuser_backward_grouped_tensor( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device dtype = ctx.dtype @@ -1695,7 +1729,7 @@ def _fuser_backward_grouped_tensor( tensor_offsets=base_split_offsets * self.in_features, ) general_grouped_gemm_for_grouped_tensor( - ws, + dist_dgrad_weights if is_dist_weight else ws, grouped_dy, grouped_grad_input, layout="NN", @@ -1741,7 +1775,8 @@ def _fuser_backward_grouped_tensor( final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad else: - if self._accumulate_into_main_grad: + # Distributed weights finalize wgrads (below); never accumulate into main_grad. + if self._accumulate_into_main_grad and not is_dist_weight: final_weight_grads = [ get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] @@ -1771,11 +1806,16 @@ def _fuser_backward_grouped_tensor( else: wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, final_weight_grads) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 31189af09c..83954a9b3d 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -19,6 +19,12 @@ from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor +from ...distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ...module.base import _2X_ACC_WGRAD from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer @@ -541,6 +547,7 @@ def _compute_grad_params( wgrad_output = None op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" weights = fc_op._get_weight_tensors() + is_dist_weight = is_distributed_weight(weights[0]) if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: @@ -573,7 +580,9 @@ def _compute_grad_params( else: w_list = [None] * num_groups if ctx.weight_requires_grad: - if fc_op._accumulate_into_main_grad: + # Distributed weight: the GEMM produces full-sized wgrads but main_grad is sharded, + # so use a full-sized scratch buffer (the reduce-scatter below lands it in main_grad). + if fc_op._accumulate_into_main_grad and not is_dist_weight: w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: @@ -589,6 +598,10 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() + if is_dist_weight and delay_wgrad: + raise RuntimeError( + "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False." + ) if cudnn_wgrad_kernel_fn is not None: offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) gemm_fn = functools.partial( @@ -627,9 +640,14 @@ def _compute_grad_params( fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: gemm_fn(grouped_x, grouped_dy, wgrad_output) + # Distributed weight: reduce-scatter the wgrads into main_grad. + # Return discarded (see finalize_weight_grads); dummy wgrads returned below. + if is_dist_weight: + finalize_weight_grads(weights, w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added - if fc_op._accumulate_into_main_grad: + # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add. + if fc_op._accumulate_into_main_grad or is_dist_weight: w_list = get_dummy_wgrads_for_params(weights) elif delay_wgrad: w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups @@ -909,6 +927,19 @@ def fuser_forward( num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + + # Distributed weight: expert weights are sharded 1/N along out_features; the fused kernels + # read the full shape, so all-gather the full weight first. A plain weight is a no-op. + fc1_is_dist = is_distributed_weight(fc1_weight_param) + fc2_is_dist = is_distributed_weight(fc2_weight_param) + assert fc1_is_dist == fc2_is_dist, "FC1/FC2 must share one distributed-weight group." + if fc1_is_dist: + assert ( + not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight + ), "distributed-weight fused grouped-MLP requires single_grouped_weight=False." + assert fc1_op.weight0.is_routed_expert and fc1_op.weight0.weight_list is not None + assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None + device = fc1_weight_param.device if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") @@ -997,6 +1028,8 @@ def fuser_forward( else: quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights + if fc1_is_dist: + grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight) # Prepare FC2 grouped weight tensor for fused kernels. if fc2_op.single_grouped_weight: @@ -1035,10 +1068,6 @@ def fuser_forward( grouped_fc1_weight, "_with_gemm_swizzled_scales" ): grouped_fc1_weight._with_gemm_swizzled_scales = False - if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( - grouped_fc2_weight, "_with_gemm_swizzled_scales" - ): - grouped_fc2_weight._with_gemm_swizzled_scales = False # Group-quantize input tensor and convert dtypes if needed fc1_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) @@ -1271,6 +1300,13 @@ def fuser_forward( else: fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + if fc2_is_dist: + grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) + if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( + grouped_fc2_weight, "_with_gemm_swizzled_scales" + ): + grouped_fc2_weight._with_gemm_swizzled_scales = False + # Unpack kernel outputs # Note: Fused kernel outputs tensors with non-contiguous # logical dims. @@ -1504,6 +1540,11 @@ def fuser_forward( fc2_weight_tensors = ( [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight ) + # Save the joint op's internal layout; distributed weights save the small shards. + if fc1_is_dist: + fc1_weight_tensors = fc1_weights + if fc2_is_dist: + fc2_weight_tensors = fc2_weights fc1_ctx.save_for_backward( split_sizes, base_split_offsets, @@ -1767,6 +1808,10 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) + fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + if is_distributed_weight(fc2_leader): + grouped_fc2_weight = materialize_weight_for_backward(fc2_leader) + if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc2_weight_for_gemm = grouped_fc2_weight.copy() @@ -1992,6 +2037,10 @@ def fuser_backward( if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] + fc1_leader = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 + if is_distributed_weight(fc1_leader): + grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) + if use_nvfp4: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) if num_groups == 1: From 8abfa2e1c600721bcd0e5b646f49a6b7167cbdcb Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:24:05 -0700 Subject: [PATCH 33/35] Update list of authorized CI users (#3241) Signed-off-by: Tim Moon --- .github/workflows/trigger-ci.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml index 68d1d7d71f..c6c91906cd 100644 --- a/.github/workflows/trigger-ci.yml +++ b/.github/workflows/trigger-ci.yml @@ -62,6 +62,7 @@ jobs: || github.actor == 'jomitchellnv' || github.actor == 'fheinecke' || github.actor == 'janekb04' + || github.actor == 'YangFei1990' ) steps: - name: Check if comment is issued by authorized person From 4adad4c218c115cd9af235fb3d4e13ef4cec55a8 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 22 Jul 2026 15:29:14 -0700 Subject: [PATCH 34/35] [Common] Fix Build: NCCL EP build to respect `MAX_JOBS` (#3138) fix nproc Signed-off-by: Phuong Nguyen --- setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index c0c37364a9..45e079c601 100644 --- a/setup.py +++ b/setup.py @@ -25,6 +25,7 @@ remove_dups, min_python_version_str, nccl_ep_enabled, + get_max_jobs_for_parallel_build, ) frameworks = get_frameworks() @@ -250,7 +251,7 @@ def build_nccl_ep_submodule() -> str: ) gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) - nproc = os.cpu_count() or 8 + nproc = get_max_jobs_for_parallel_build() env = os.environ.copy() env["NVCC_GENCODE"] = gencode # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build From 457772bc7b3ceaa9cf809623d0efeb04242db9c3 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 23 Jul 2026 14:07:49 +0200 Subject: [PATCH 35/35] [PyTorch] Address review nits: clarify DPA packed-input docstrings - DotProductAttention.forward: mark query/key/value_layer as Optional and note they are required only when no packed input (qkv_layer/kv_layer) is given (Charlene). - combine_and_quantize: describe combined_qkv/combined_kv in terms of qkv_group=1 / qkv_group=2 layouts instead of '3'/'2' layouts (Charlene). Signed-off-by: Pawel Gadzinski --- .../dot_product_attention.py | 15 +++++++++------ .../attention/dot_product_attention/utils.py | 4 ++-- 2 files changed, 11 insertions(+), 8 deletions(-) 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 82a23ff021..4cc4cab1b8 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 @@ -1221,12 +1221,15 @@ def forward( Parameters ---------- - query_layer : torch.Tensor - Query tensor. - key_layer : torch.Tensor - Key tensor. - value_layer : torch.Tensor - Value tensor. + query_layer : Optional[torch.Tensor], default = None + Query tensor. Required unless a packed input (``qkv_layer``, or + ``kv_layer`` together with ``query_layer``) is provided instead. + key_layer : Optional[torch.Tensor], default = None + Key tensor. Required unless a packed input (``qkv_layer`` or ``kv_layer``) + is provided instead. + value_layer : Optional[torch.Tensor], default = None + Value tensor. Required unless a packed input (``qkv_layer`` or ``kv_layer``) + is provided instead. attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]], default = None. Boolean tensor(s) used to mask out attention softmax input. It should be ``None`` for causal masks and ``"no_mask"``. For padding masks, it should be diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 1205896085..1713a98129 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2809,8 +2809,8 @@ def combine_and_quantize( ): """Combine Q, K, V tensors based on qkv_layout and quantize them together. - When ``combined_qkv`` (for ``3`` layouts such as ``bs3hd``) or ``combined_kv`` - (for ``2`` layouts such as ``bshd_bs2hd``) is provided, it must be the + When ``combined_qkv`` (for ``qkv_group=1`` layouts such as ``bs3hd``) or + ``combined_kv`` (for ``qkv_group=2`` layouts such as ``bshd_bs2hd``) is provided, it must be the caller's original packed buffer that q/k/v are views of. It is then quantized directly instead of re-deriving the packed buffer from the q/k/v views via ``combine_tensors`` (which rebuilds it with a raw ``set_`` under a silent