From f2be56a50a95e6401ffcb574d87daabc6fb26b61 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:55:46 +0800 Subject: [PATCH 01/32] [PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP (#3545) * [PyTorch] Quantize distributed BF16 weights in the fused grouped MLP A distributed (e.g. GTP) weight whose shards are not natively quantized is re-gathered as BF16. The fused grouped-MLP op assumed gathered weights were already quantized: forward quantized the shards before the gather, and backward read _columnwise_data / the NVFP4 amax straight off the gathered tensor, which fails for a module whose per-module recipe (e.g. MXFP8 experts under a global NVFP4 recipe) does not create native-quantized params. Gather first, then quantize whatever is still unquantized with the op's per-expert weight quantizers, mirroring module/linear.py's handling of a re-gathered BF16 distributed weight. Native-quantized shards gather as quantized tensors and pass through unchanged. The FC2 backward re-gather moves ahead of the NVFP4 amax / global-scale computation so that path no longer reads un-gathered shards. The per-expert weight quantizer lists are built once in forward and saved on the FC1/FC2 contexts; backward quantizes the re-gathered weights with those rather than reading the live op. A later forward on the same module under a different recipe, or with quantization disabled, resets the ops' quantizers before the earlier backward runs (pipelined microbatches, or a module shared across recipe scopes), which would otherwise hit None or quantize in the wrong format. Regression tests: a fake routed-expert DistributedWeight group whose materialize hooks return BF16 tensors (or MXFP8 tensors for the native-quantized case) drives the CuTeDSL fused grouped MLP under per-module MXFP8, optionally inside a global NVFP4 autocast. Distinct power-of-two scales in the forward/backward hooks make the gathered weights observable: output, dgrad and the main_grad wgrads must match a plain reference bit-exactly up to those scales. A unit check confirms already-quantized gathered weights pass through the helper unchanged. A further case runs a second forward under no recipe or NVFP4 between an MXFP8 forward and its backward and requires bit-exact agreement with a reference. Fails on the unpatched op with AttributeError on _columnwise_data in backward; passes with the fix. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren * [PyTorch] Drop NVFP4 rowwise input copies from the fused grouped MLP saved state Every wgrad path reads only the columnwise tiles of the saved FC1/FC2 inputs: the cuDNN grouped wgrad kernel, the single-group dense GEMM (general_gemm with the NT layout unwraps both operands columnwise) and the generic grouped-GEMM fallback (NVFP4 forces a TN layout by taking the columnwise data for a non-transposed A and a transposed B, which is the wgrad case). The save path already pruned the rowwise data and scales for MXFP8 but kept them for NVFP4, citing the generic fallback, so each quantized NVFP4 input was held (and, under activation offloading, copied to host and back) twice. Prune them for NVFP4 as well, unconditionally. The amax tensors are untouched. Per decoder expert layer of a Nemotron hybrid proxy (83k tokens, hidden 3072) this removes a 121 MiB data and a 15 MiB scale tensor from the ~1.9 GiB saved set, with unchanged numerics on both the cuDNN kernel and the forced generic fallback (NVTE_DISABLE_CUTEDSL_WGRAD_FUSED_GROUPED_MLP=1). Adds a saved_tensors_hooks test asserting exactly one FP4 data and one scale tensor is saved per quantized input on the fused NVFP4 path, for group sizes 1 and 4, with token counts chosen so the inputs cannot be confused with weights in the count, and registers it in the L0 unit-test runner. Grouped storage is saved as raw buffers; the single-group dense path (cuDNN frontend >= 1.27) saves a GroupedTensor wrapper whole, so the test inspects the wrapper's buffers directly and requires its rowwise data and scales to be None. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren * Update transformer_engine/pytorch/ops/fused/grouped_mlp.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com> * [PyTorch] Move the NVFP4 saved-state test into test_grouped_mlp.py The check that the fused grouped MLP saves only the columnwise NVFP4 input tiles lives naturally with the other fused grouped MLP tests. Move it into TestGroupedMLPFusedOp as test_nvfp4_saves_columnwise_inputs_only, drop the standalone file, and remove its separate line from the L0 launcher, which already runs test_grouped_mlp.py. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren * [PyTorch] Move the fused grouped MLP distributed-weight tests into test_grouped_mlp.py The DistributedWeight dispatch tests for the fused grouped MLP (BF16 and native-quantized shard gather, the quantizer pass-through unit check, and the backward-after-recipe-change regression) move into test_grouped_mlp.py as TestGroupedMLPDistributedWeight, with their fake DistributedWeight helpers. Drop the standalone file and its separate L0 launcher line, which already runs test_grouped_mlp.py. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren * [PyTorch] Build fused grouped MLP weight quantizer lists lazily The per-expert weight quantizer lists were built unconditionally at the top of the fused grouped MLP forward, costing 2 * num_groups quantizer lookups on every call, including the single grouped weight path that never uses them. Build them only where they are consumed: the discrete-weight FC1/FC2 quantization and the distributed-weight gather-then-quantize sites. Save the full list on the backward context only for distributed weights, which is the one path that re-quantizes in backward; every other path keeps saving just the first expert's quantizer as before. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren * Apply suggestion from @ksivaman Co-authored-by: Kirthi Shankar Sivamani Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com> * [PyTorch] Skip columnwise quantization of gathered distributed FC1 weights in forward The forward GEMMs read only the rowwise weight data; the columnwise copy exists for the dgrad GEMM in backward, which reads it from the saved weight. A distributed weight's gathered copy is not saved (the shards are, and backward re-gathers and quantizes columnwise itself), so the columnwise tiles produced in forward were built and immediately discarded. Match the FC2 distributed site and skip them. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren --------- Signed-off-by: Xiaowei Ren Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_grouped_mlp.py | 474 ++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 136 +++-- 2 files changed, 575 insertions(+), 35 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index fd661ea1e00..f5ce1f8ef83 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -5,6 +5,7 @@ from __future__ import annotations from collections.abc import Iterable +from contextlib import nullcontext import functools import os import math @@ -21,7 +22,9 @@ from transformer_engine.pytorch.constants import TE_DType import transformer_engine.pytorch.ops.fused.grouped_mlp as grouped_mlp_module from transformer_engine.pytorch.ops.fuser import OperationFuser +from transformer_engine.pytorch.ops._common import is_quantized_tensor from transformer_engine.pytorch.ops.fused.grouped_mlp import ( + GroupedMLP_CuTeGEMMGLU, _cudnn_frontend_supports_grouped_gemm_situglu, _cudnn_frontend_supports_grouped_gemm_srelu, _cudnn_frontend_version_at_least, @@ -38,10 +41,12 @@ Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, + MXFP8Tensor, NVFP4Quantizer, QuantizerRole, is_bf16_available, ) +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.pytorch.utils import get_device_compute_capability import transformer_engine_torch as tex @@ -3289,6 +3294,475 @@ def train_step( for graph_grad, param in zip(graph_param_grads, reference_module.parameters()): assert_close(graph_grad, param.grad, **tols) + # (group_size, tokens_per_group): token counts chosen so the quantized inputs' byte sizes + # never coincide with a weight's, which the saved-tensor counting below relies on. + @pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) + @pytest.mark.parametrize("group_size,tokens_per_group", ((1, 768), (4, 256))) + def test_nvfp4_saves_columnwise_inputs_only( + self, + group_size: int, + tokens_per_group: int, + *, + hidden_size: int = 256, + glu_interleave_size: int = 32, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + ) -> None: + """The fused grouped MLP must not save the NVFP4 rowwise input copies for backward. + + Every wgrad path (the cuDNN grouped kernel, the single-group dense GEMM and the generic + grouped GEMM fallback) reads only the columnwise tiles (data and scales) of the saved FC1 + and FC2 inputs, so the rowwise copies are dead after forward. Saving them doubled the + activation memory held (and offloaded) for the quantized inputs. + + Checked through ``saved_tensors_hooks``: the saved set must hold exactly one FP4 data + tensor and one scale tensor per quantized input. Grouped storage is saved as its raw + buffers; the single-group dense path (cuDNN frontend >= 1.27) instead saves a + ``GroupedTensor`` wrapper whose buffers are inspected directly. + + Weights require grad here: with frozen weights the input quantizers skip the columnwise + copy and the fused op's ``mark_grouped_tensor`` call rejects the input before the pruning + is reached. + """ + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + + total_tokens = group_size * tokens_per_group + split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device) + # The fused grouped MLP requires a split alignment divisible by 256 (the default is 128). + fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=256, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=256, + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x.requires_grad_() + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + saved = [] + + def pack(t): + if isinstance(t, GroupedTensor): + # Dense single-group path: the wrapper is saved whole, so count its buffers. + assert t.rowwise_data is None, "rowwise FP4 input data was saved" + assert t.scale_inv is None, "rowwise FP4 input scales were saved" + for buf in (t.columnwise_data, t.columnwise_scale_inv): + if buf is not None: + saved.append((buf.dtype, buf.numel())) + elif isinstance(t, torch.Tensor): + saved.append((t.dtype, t.numel())) + return t + + with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): + with te.autocast(enabled=True, recipe=make_recipe("nvfp4_rht")): + y = module(x, split_sizes, probs, split_sizes) + forward_ops = module._module_groups[0]._forward_ops + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Byte counts that identify the quantized FC1/FC2 inputs (both are [tokens, hidden_size] + # in NVFP4: 2 elements per byte, one E4M3 scale per 16 elements). Weights have other + # sizes. + fp4_data_numel = total_tokens * hidden_size // 2 + fp4_scale_numel = total_tokens * hidden_size // 16 + n_data = sum(1 for d, n in saved if d == torch.uint8 and n == fp4_data_numel) + n_scale = sum(1 for d, n in saved if d == torch.uint8 and n == fp4_scale_numel) + # One columnwise data + one columnwise scale tensor per quantized input (FC1 x, FC2 x). + assert n_data == 2, f"expected 2 FP4 input data tensors saved, got {n_data}" + assert n_scale == 2, f"expected 2 FP4 input scale tensors saved, got {n_scale}" + + # Backward still runs on what was saved. + y.backward(dy) + assert x.grad is not None and torch.isfinite(x.grad).all() + for p in module.parameters(): + assert p.grad is not None and torch.isfinite(p.grad).all() + + +# --------------------------------------------------------------------------------------------- +# DistributedWeight dispatch in the fused grouped MLP. +# +# ``test_ops_grouped_linear_distributed_weight.py`` covers the unfused ``ops.GroupedLinear`` +# path. The fused op has its own weight plumbing, so it is covered here with the same kind of +# in-repo fake: +# +# * A DistributedWeight whose shards are plain BF16 (e.g. a Megatron GTP expert under a +# per-module MXFP8 recipe, which creates no native-quantized params) gathers as BF16. The +# fused op must gather *before* quantizing in forward, and in backward must quantize the +# re-gathered BF16 weight before the dgrad GEMMs read quantized storage. +# * A DistributedWeight whose shards are native-quantized gathers as quantized tensors, which +# must pass through unchanged (no re-quantization). +# +# The per-module recipe is MXFP8 entered inside a global NVFP4 autocast, mirroring how Megatron +# applies a ``--te-precision-config-file`` override around a fused expert MLP. +# +# As in the unfused test, the fakes apply distinct power-of-two scales in the materialize hooks +# so the gathered weights are observably the ones the GEMMs consumed. A power-of-two scale +# commutes exactly through MXFP8 quantization (E8M0 block scales shift, E4M3 data is unchanged) +# and through BF16 rounding, so the scale relations below hold bit-exactly. FC1 gets no forward +# scale because the activation between the two GEMMs is nonlinear. +# --------------------------------------------------------------------------------------------- + +# Distinct powers of two: each commutes exactly through quantization and BF16 rounding. +_DIST_FC1_BWD_SCALE = 2.0 +_DIST_FC2_FWD_SCALE = 4.0 +_DIST_FC2_BWD_SCALE = 8.0 + + +def _mxfp8_quantize(tensor: torch.Tensor) -> MXFP8Tensor: + """Native-quantized stand-in: both usages, like a shard the module's own quantizer created.""" + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=True) + return quantizer(tensor) + + +class _FakeDistWeight(torch.nn.Parameter): + """Single-process fake DistributedWeight shard of a routed-expert group. + + Every shard of the group is one of these (as with a real implementer such as GTP); only the + leader (``weight0``) carries ``weight_list``/``_group`` and is called back by the op. + ``materialize_group_for_forward`` / ``materialize_group_for_backward`` return the whole + group's "gathered" weights scaled by ``fwd_scale`` / ``bwd_scale``: plain BF16 tensors (a + BF16 shard all-gather) or MXFP8 tensors when ``gather_quantized`` (a native-quantized shard + all-gather). ``finalize_group_grads`` models the main-grad contract: accumulate the wgrads + into each shard's ``main_grad`` in-place and return dummies that the op discards. + ``grad_buffer`` hands the wgrad GEMM a fresh full-shape buffer in the compute dtype. + """ + + is_distributed_weight = True + # The fused op only takes the distributed path for a routed-expert group leader. + is_routed_expert = True + + def _gathered(self, scale: float): + gathered = [w.detach() * scale for w in self._group] + if self.gather_quantized: + gathered = [_mxfp8_quantize(w) for w in gathered] + return gathered + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return self._gathered(self.fwd_scale) + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return self._gathered(self.bwd_scale) + + 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)) + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] + + def grad_buffer(self): + return torch.empty_like(self.data) + + +def _make_fake_dist_leader(op, num_gemms, *, fwd_scale, bwd_scale, gather_quantized): + """Turn ``op.weight{i}`` into fake distributed shards; return the leader (``weight0``).""" + shards = [] + for i in range(num_gemms): + shard = _FakeDistWeight(getattr(op, f"weight{i}").data) + shard.main_grad = torch.zeros_like(shard, dtype=torch.float32) + shard.grad_added_to_main_grad = False # DDP initializes this on every param + setattr(op, f"weight{i}", shard) + shards.append(shard) + leader = shards[0] + leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + leader.fwd_scale = fwd_scale + leader.bwd_scale = bwd_scale + leader.gather_quantized = gather_quantized + leader._group = shards + leader.weight_list = shards # routed-expert group membership, asserted by the op + return leader + + +def _skip_unless_fused_mxfp8_supported(device: str, dtype: torch.dtype) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + if dtype != torch.bfloat16: + pytest.skip("Fused grouped MLP requires BF16 data") + + +def _make_dist_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device +): + fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=split_alignment, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=split_alignment, + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + return module, fc1, fc2 + + +def _assert_fused(module) -> None: + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], GroupedMLP_CuTeGEMMGLU) + + +class TestGroupedMLPDistributedWeight: + """DistributedWeight (e.g. GTP) dispatch in the fused grouped MLP; see the notes above.""" + + @pytest.mark.parametrize("outer_recipe", (None, "nvfp4")) + @pytest.mark.parametrize("gather_quantized", (False, True)) + def test_fused_grouped_mlp_distributed_weight( + self, + outer_recipe, + gather_quantized, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + ) -> None: + """Fused grouped MLP with distributed weights: gather-then-quantize in fwd and bwd. + + ``gather_quantized=False`` is the BF16-shard case (per-module MXFP8 without + native-quantized params); ``gather_quantized=True`` is the native-quantized-shard case, + whose gathered tensors must pass through. ``outer_recipe="nvfp4"`` wraps the per-module + MXFP8 autocast in a global NVFP4 one, as Megatron does for an overridden expert MLP. + """ + _skip_unless_fused_mxfp8_supported(device, dtype) + if outer_recipe == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + torch.manual_seed(0) + + split_sizes = torch.tensor( + [split_alignment, 2 * split_alignment, split_alignment, 2 * split_alignment], + dtype=torch.int, + device=device, + ) + assert split_sizes.numel() == group_size + total_tokens = int(split_sizes.sum().item()) + + module, fc1, fc2 = _make_dist_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference, ref_fc1, ref_fc2 = _make_dist_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference.load_state_dict(module.state_dict()) + + fc1_leader = _make_fake_dist_leader( + fc1, + group_size, + fwd_scale=1.0, + bwd_scale=_DIST_FC1_BWD_SCALE, + gather_quantized=gather_quantized, + ) + fc2_leader = _make_fake_dist_leader( + fc2, + group_size, + fwd_scale=_DIST_FC2_FWD_SCALE, + bwd_scale=_DIST_FC2_BWD_SCALE, + gather_quantized=gather_quantized, + ) + + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x.requires_grad_(True) + ref_x = x.detach().clone().requires_grad_(True) + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + mxfp8_recipe = make_recipe("mxfp8") + outer = make_recipe(outer_recipe) if outer_recipe is not None else None + + def _run(mod, inp): + outer_ctx = ( + te.autocast(enabled=True, recipe=outer) if outer is not None else nullcontext() + ) + # Per-module MXFP8 override entered inside the global autocast (Megatron's pattern). + with outer_ctx, te.autocast(enabled=True, recipe=mxfp8_recipe): + y = mod(inp, split_sizes, probs, split_sizes) + y.backward(dy) + return y.detach() + + out = _run(module, x) + ref_out = _run(reference, ref_x) + + # The CuTeDSL fused op was formed on both (not the unfused fallback). + _assert_fused(module) + _assert_fused(reference) + + # Every dispatch hook fired on both groups. + for leader in (fc1_leader, fc2_leader): + assert ( + leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 + ) + + # fwd used the gathered FC2 weights (scale commutes exactly through quantization). + torch.testing.assert_close(out, _DIST_FC2_FWD_SCALE * ref_out, rtol=0, atol=0) + # dgrad used the re-gathered FC2 and FC1 weights. + torch.testing.assert_close( + x.grad, _DIST_FC1_BWD_SCALE * _DIST_FC2_BWD_SCALE * ref_x.grad, rtol=0, atol=0 + ) + # wgrads were reduce-scattered into main_grad: FC2's does not depend on the weight + # scales, FC1's sees the FC2-backward scale through the activation gradient. + tols = dtype_tols(dtype) + for i in range(group_size): + fc2_w, ref_fc2_w = getattr(fc2, f"weight{i}"), getattr(ref_fc2, f"weight{i}") + fc1_w, ref_fc1_w = getattr(fc1, f"weight{i}"), getattr(ref_fc1, f"weight{i}") + torch.testing.assert_close(fc2_w.main_grad.to(dtype), ref_fc2_w.grad, **tols) + torch.testing.assert_close( + fc1_w.main_grad.to(dtype), _DIST_FC2_BWD_SCALE * ref_fc1_w.grad, **tols + ) + for w in (fc1_w, fc2_w): + assert w.grad_added_to_main_grad is True + assert w.grad is not None, "distributed weights receive a dummy .grad" + + def test_quantize_unquantized_weights_passes_through_quantized( + self, + *, + group_size: int = 2, + hidden_size: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + ) -> None: + """Already-quantized gathered weights are returned as-is; plain ones are quantized.""" + from transformer_engine.pytorch.ops.fused.grouped_mlp import _quantize_unquantized_weights + + _skip_unless_fused_mxfp8_supported(device, dtype) + torch.manual_seed(0) + module, fc1, fc2 = _make_dist_grouped_mlp(group_size, hidden_size, 256, 32, dtype, device) + split_sizes = torch.tensor([256] * group_size, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + + # A forward initializes the op's per-expert weight quantizers for the recipe. + with te.autocast(enabled=True, recipe=make_recipe("mxfp8")): + module(x, split_sizes, probs, split_sizes) + pre_quantized = _mxfp8_quantize(fc2.weight0.detach()) + plain = fc2.weight1.detach() + quantizers = [fc2.get_quantizer("forward", 2 * idx + 1) for idx in range(group_size)] + out = _quantize_unquantized_weights( + [pre_quantized, plain], quantizers, rowwise=True, columnwise=True + ) + + assert out[0] is pre_quantized, "quantized gathered weight must pass through unchanged" + assert is_quantized_tensor(out[1]), "BF16 gathered weight must be quantized" + torch.testing.assert_close( + out[1]._rowwise_data, _mxfp8_quantize(plain)._rowwise_data, rtol=0, atol=0 + ) + + @pytest.mark.parametrize("later_recipe", (None, "nvfp4"), ids=["no-recipe", "nvfp4"]) + def test_backward_after_recipe_change( + self, + later_recipe, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + ) -> None: + """Backward must quantize re-gathered BF16 weights with the quantizers saved in forward. + + A later forward under another recipe (or none) resets the ops' quantizers before the + earlier backward runs, as with pipelined microbatches or a module shared across recipe + scopes. Read from the live op, backward would hit ``None`` (``later_recipe=None``) or + quantize the MXFP8 graph's weights in the wrong format (``later_recipe="nvfp4"``). + """ + _skip_unless_fused_mxfp8_supported(device, dtype) + if later_recipe == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + torch.manual_seed(0) + + split_sizes = torch.tensor( + [split_alignment, 2 * split_alignment, split_alignment, 2 * split_alignment], + dtype=torch.int, + device=device, + ) + assert split_sizes.numel() == group_size + total_tokens = int(split_sizes.sum().item()) + + module, fc1, fc2 = _make_dist_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference, ref_fc1, ref_fc2 = _make_dist_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference.load_state_dict(module.state_dict()) + for fc in (fc1, fc2, ref_fc1, ref_fc2): + # BF16 shards: the gathered weights are what backward has to quantize. + _make_fake_dist_leader( + fc, group_size, fwd_scale=1.0, bwd_scale=1.0, gather_quantized=False + ) + + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x.requires_grad_(True) + ref_x = x.detach().clone().requires_grad_(True) + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + mxfp8_recipe = make_recipe("mxfp8") + + with te.autocast(enabled=True, recipe=mxfp8_recipe): + out = module(x, split_sizes, probs, split_sizes) + ref_out = reference(ref_x, split_sizes, probs, split_sizes) + _assert_fused(module) + ref_out.backward(dy) + + # Another forward on the same module before its backward, under a different recipe. + later_ctx = ( + te.autocast(enabled=True, recipe=make_recipe(later_recipe)) + if later_recipe is not None + else nullcontext() + ) + with torch.no_grad(), later_ctx: + module(x, split_sizes, probs, split_sizes) + if later_recipe is None: + assert fc1.get_quantizer("forward", 1) is None, "later forward did not reset quantizers" + else: + assert not isinstance(fc1.get_quantizer("forward", 1), MXFP8Quantizer) + + out.backward(dy) + + torch.testing.assert_close(out, ref_out, rtol=0, atol=0) + torch.testing.assert_close(x.grad, ref_x.grad, rtol=0, atol=0) + for i in range(group_size): + for fc, ref_fc in ((fc1, ref_fc1), (fc2, ref_fc2)): + w, ref_w = getattr(fc, f"weight{i}"), getattr(ref_fc, f"weight{i}") + torch.testing.assert_close(w.main_grad, ref_w.main_grad, rtol=0, atol=0) + class TestGroupedMLPDeterminism: """Determinism coverage for the CuTe DSL fused grouped MLP. diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 46763a6629f..d678d92749c 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -383,6 +383,43 @@ def _group_quantize_with_amax_for_grouped_mlp( ) +def _weight_quantizers(op: GroupedLinear, num_groups: int) -> list[Quantizer]: + """Per-expert weight quantizers of ``op``. + + Only built where they are consumed (discrete or distributed weights), so the single grouped + weight path does not pay for ``num_groups`` lookups it never uses. + """ + return [op.get_quantizer("forward", 2 * idx + 1) for idx in range(num_groups)] + + +def _quantize_unquantized_weights( + weights: Iterable[Any], + quantizers: Sequence[Quantizer], + *, + rowwise: bool, + columnwise: bool, +) -> list: + """Quantize plain weights with the given per-expert weight quantizers. + + Quantized entries (native-quantized params, incl. native GTP shards gathered as quantized + tensors) pass through. A distributed (e.g. GTP) BF16 weight gathers as BF16 and must be + quantized with the layer's own quantizer, mirroring ``module/linear.py``'s handling of a + re-gathered BF16 distributed weight. + + ``quantizers`` must be the ones saved on the forward context, not read from the live op: + a later forward under a different (or no) recipe replaces or clears the op's quantizers + before an earlier backward runs. + """ + out = [] + for weight, quantizer in zip(weights, quantizers): + if is_quantized_tensor(weight): + out.append(weight) + continue + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + out.append(quantizer(weight)) + return out + + def _nvfp4_amax( tensors: GroupedTensor | Iterable[NVFP4TensorStorage], *, @@ -1363,6 +1400,10 @@ def fuser_forward( fc2_input_quantizer = fc2_op.get_quantizer("forward", 0) fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1) fc2_grad_output_quantizer = fc2_op.get_quantizer("backward", 0) + # Per-expert weight quantizer lists are built lazily, only on the discrete-weight and + # distributed-weight paths that consume them. + fc1_weight_quantizers = None + fc2_weight_quantizers = None # Extract split sizes from extra input fc1_split_sizes = basic_op_extra_inputs[0][0] @@ -1483,17 +1524,17 @@ def fuser_forward( ) else: fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] - quantized_fc1_weights = [] - for idx, weight in enumerate(fc1_weights): - quantizer = fc1_op.get_quantizer("forward", 2 * idx + 1) - if not is_quantized_tensor(weight): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc1_weights.append(quantizer(weight)) - 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) + fc1_weight_quantizers = _weight_quantizers(fc1_op, num_groups) + # Distributed weights gather first (native-quantized shards gather quantized, BF16 + # shards gather as BF16); anything still unquantized is quantized here. The + # gathered copy is not saved (backward re-gathers and quantizes columnwise + # itself), so its columnwise data would be thrown away: skip producing it. + grouped_fc1_weight = _quantize_unquantized_weights( + materialize_weight_for_forward(fc1_weights) if fc1_is_dist else fc1_weights, + fc1_weight_quantizers, + rowwise=True, + columnwise=input_requires_grad and not fc1_is_dist, + ) # Prepare FC2 grouped weight tensor for fused kernels. if fc2_op.single_grouped_weight: @@ -1521,15 +1562,15 @@ def fuser_forward( ) else: fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] - quantized_fc2_weights = [] - for idx, weight in enumerate(fc2_weights): - quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) - if not is_quantized_tensor(weight): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc2_weights.append(quantizer(weight)) - else: - quantized_fc2_weights.append(weight) - grouped_fc2_weight = quantized_fc2_weights + # Distributed weights are gathered (and quantized if needed) after the FC1 kernel so + # the gather overlaps FC1 compute. + if fc2_is_dist: + grouped_fc2_weight = fc2_weights + else: + fc2_weight_quantizers = _weight_quantizers(fc2_op, num_groups) + grouped_fc2_weight = _quantize_unquantized_weights( + fc2_weights, fc2_weight_quantizers, rowwise=True, columnwise=input_requires_grad + ) # Some wrapper-copy paths may drop grouped storage metadata; enforce defaults. if isinstance(grouped_fc1_weight, GroupedTensor) and not hasattr( @@ -1858,7 +1899,13 @@ def fuser_forward( activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) if fc2_is_dist: - grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) + fc2_weight_quantizers = _weight_quantizers(fc2_op, num_groups) + grouped_fc2_weight = _quantize_unquantized_weights( + materialize_weight_for_forward(grouped_fc2_weight), + fc2_weight_quantizers, + rowwise=True, + columnwise=False, + ) if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( grouped_fc2_weight, "_with_gemm_swizzled_scales" ): @@ -2201,13 +2248,11 @@ def fuser_forward( ) saved_grouped_fc2_x = None if recompute_srelu_fc2_x else grouped_fc2_x - # MXFP8 wgrad only needs columnwise tiles. NVFP4 generic GEMM fallbacks - # need the full grouped tensor state, including rowwise data and amax. - if not use_nvfp4: - for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): - if isinstance(grouped_fc_x, (GroupedTensor, GroupedTensorStorage)): - grouped_fc_x.rowwise_data = None - grouped_fc_x.scale_inv = None + # wgrad only needs column-wise data, so clear row-wise data + for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): + if isinstance(grouped_fc_x, (GroupedTensor, GroupedTensorStorage)): + grouped_fc_x.rowwise_data = None + grouped_fc_x.scale_inv = None if cpu_offloading: activation_tensors = [ @@ -2246,7 +2291,11 @@ def fuser_forward( ) fc1_ctx.input_quantizers = [fc1_input_quantizer] - fc1_ctx.weight_quantizers = [fc1_weight_quantizer] + # Backward re-quantizes re-gathered distributed weights with the quantizers saved + # here; other paths only read the first expert's quantizer. + fc1_ctx.weight_quantizers = ( + fc1_weight_quantizers if fc1_is_dist else [fc1_weight_quantizer] + ) fc1_ctx.grad_output_quantizers = [fc1_grad_output_quantizer] fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad @@ -2255,7 +2304,9 @@ def fuser_forward( fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] - fc2_ctx.weight_quantizers = [fc2_weight_quantizer] + fc2_ctx.weight_quantizers = ( + fc2_weight_quantizers if fc2_is_dist else [fc2_weight_quantizer] + ) fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad @@ -2506,6 +2557,18 @@ def fuser_backward( scales_tensor = scales_f32.reshape(-1, 1, 1) dscales_tensor = torch.zeros_like(scales_tensor) + # Distributed weight (e.g. GTP): re-gather before any use (the NVFP4 global scale below + # reads the weights). BF16 shards gather as BF16 and are quantized with the weight + # quantizers saved in forward (the live op's may have been replaced or cleared since). + fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + if is_distributed_weight(fc2_leader): + grouped_fc2_weight = _quantize_unquantized_weights( + materialize_weight_for_backward(fc2_leader), + fc2_ctx.weight_quantizers, + rowwise=False, + columnwise=True, + ) + fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: fc2_dy_global_scale = _nvfp4_global_scale( @@ -2589,10 +2652,6 @@ def fuser_backward( if self._pass_srelu_tanh_params: fc2_dactivation_kwargs.update(tanh_clamp_scale=self._cudnn_tanh_clamp_scale) - 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() @@ -2877,7 +2936,14 @@ def fuser_backward( 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) + # BF16 GTP shards gather as BF16: quantize with the weight quantizers saved in + # forward (the live op's may have been replaced or cleared since). + grouped_fc1_weight = _quantize_unquantized_weights( + materialize_weight_for_backward(fc1_leader), + fc1_ctx.weight_quantizers, + rowwise=False, + columnwise=True, + ) fc1_dgrad_sf_override = _nvfp4_sf_dtype_override(fc1_grad_output_quantizer) From 5d7bbd851d8e6a64e5ea2460c37ff8be823a29c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:44:49 +0100 Subject: [PATCH 02/32] [PyTorch] Keep None sentinels on the custom op device (#3739) * [PyTorch] Keep None sentinels on the custom op device Signed-off-by: Pawel Gadzinski * [PyTorch] Remove the sentinel device test Signed-off-by: Pawel Gadzinski * [PyTorch] Restore the original test file ending Signed-off-by: Pawel Gadzinski * [PyTorch] Scope the device to custom op result packing Signed-off-by: Pawel Gadzinski * [PyTorch] Make result device selection easier to read Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 39d59dc2907..1ed84017ecd 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -985,14 +985,24 @@ def _register_base_op( ``pack_result``. """ + def _pack_on_input_device(result: Any, flat: Sequence[Any]) -> List[torch.Tensor]: + device = torch.device("cpu") + for tensor in tree_flatten(flat)[0]: + if isinstance(tensor, torch.Tensor) and tensor.device.type != "cpu": + device = tensor.device + break + # Keep None sentinels on the op device for CUDA Graph partitioning. + with device: + return pack_result(result) + def _impl(*flat: Any) -> List[torch.Tensor]: obj = plan.unpack(dict(zip(plan.slot_names, flat))) - return pack_result(impl(obj)) + return _pack_on_input_device(impl(obj), flat) def _fake(*flat: Any) -> List[torch.Tensor]: obj = plan.unpack(dict(zip(plan.slot_names, flat))) spec_obj = _spec_view(obj, plan.tensor_field_names()) - return pack_result(fake_impl(spec_obj)) + return _pack_on_input_device(fake_impl(spec_obj), flat) op = torch.library.custom_op( f"{_TE_OP_NAMESPACE}::{op_name}", _impl, mutates_args=(), schema=schema_str, tags=tags From fac8d4a90601e37cc2d6da12bf5ba13d52537d18 Mon Sep 17 00:00:00 2001 From: Apoorva Garg Date: Mon, 26 Oct 2026 15:34:54 -0400 Subject: [PATCH 03/32] Run minimum-versions build on pushes to main to warm sccache (#3721) * Run minimum-versions build on pushes to main to warm sccache The workflow only runs on pull_request, so sccache's GHA backend stores every object in a PR-scoped cache that other PRs can't read. Running it on pushes to main populates the base-branch scope, which PR runs can restore, so a new PR's first build no longer compiles from scratch. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Rt4UYn6zqR9uS7DWNGEBYx Signed-off-by: Apoorva Garg <271577+agarg21@users.noreply.github.com> * Restrict minimum-version workflow token permissions Signed-off-by: Przemek Tredak --------- Signed-off-by: Apoorva Garg <271577+agarg21@users.noreply.github.com> Signed-off-by: Przemek Tredak Co-authored-by: Apoorva Garg <271577+agarg21@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Przemek Tredak --- .github/workflows/minimum_versions.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/minimum_versions.yml b/.github/workflows/minimum_versions.yml index cc1802b1fb8..95a17a77c8c 100644 --- a/.github/workflows/minimum_versions.yml +++ b/.github/workflows/minimum_versions.yml @@ -6,7 +6,12 @@ name: 'Minimum supported versions' on: pull_request: + push: + branches: + - main workflow_dispatch: +permissions: + contents: read concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true From 38509867458c686b3ebc3af9a713c9fec7b1ad7b Mon Sep 17 00:00:00 2001 From: Zhihao Wang <101526713+xiuhu17@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:51:51 -0700 Subject: [PATCH 04/32] Support row-only MXFP8 distributed master-weight casts (#3488) * Support row-only MXFP8 distributed master-weight casts Signed-off-by: xiuhu17 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Share the NCCL process group across row-only cast tests Signed-off-by: xiuhu17 * Address row-only MXFP8 review and preserve bidirectional performance Signed-off-by: Zhihao Wang * Select MXFP8 partial-cast directions from output storage Use columnwise output presence to select partial-cast computation in the common API and PyTorch binding. Require scales for requested columnwise output and ignore them when that output is omitted. Document the contract and cover unused scales, empty shards, and missing required scales in regression tests. Signed-off-by: Codex --------- Signed-off-by: xiuhu17 Signed-off-by: Zhihao Wang Signed-off-by: Codex Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Codex --- .../test_cast_master_weights_to_fp8.py | 84 ++++++++++++++ tests/pytorch/test_partial_cast.py | 37 ++++++ .../include/transformer_engine/recipe.h | 5 +- .../common/recipe/mxfp8_scaling.cu | 109 +++++++++++------- .../csrc/extensions/fp8_partial_cast.cpp | 16 ++- transformer_engine/pytorch/tensor/utils.py | 39 +++++-- 6 files changed, 236 insertions(+), 54 deletions(-) diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index ab047dc32ff..7bf9b3699dd 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -9,6 +9,7 @@ import subprocess import sys import pathlib +from unittest.mock import patch sys.path.append(str(pathlib.Path(__file__).resolve().parent.parent)) from utils import run_distributed @@ -958,6 +959,86 @@ def _test_mxfp8_empty_master_shard(dp_group): ) +def _test_rowwise_master_cast(group, mixed, fragments): + rank = torch.distributed.get_rank(group) + world = torch.distributed.get_world_size(group) + shape = (160, 96) # Both scale padding and non-tile-aligned logical shapes. + full = torch.linspace(-4, 7, 160 * 96, device="cuda").reshape(shape) + quantizers = [te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=False)] + if mixed: + quantizers.append(te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=True)) + weights = [q(full.to(torch.bfloat16)) for q in quantizers] + pointers = [(w._rowwise_data.data_ptr(), w._rowwise_scale_inv.data_ptr()) for w in weights] + if world == 1: + lo, hi = 0, full.numel() + elif rank == 0: + lo, hi = 0, 17 # Split inside a 32-value block. + elif rank == 1: + lo, hi = 17, full.numel() + else: + lo, hi = 0, 0 # No master shard on tail ranks. + real_reduce = torch.distributed.all_reduce + reduced_sizes = [] + + def record_reduce(tensor, *args, **kwargs): + reduced_sizes.append(tensor.numel()) + return real_reduce(tensor, *args, **kwargs) + + for step in range(2): + # Also exercise empty ranks in the existing two-rank CI launch. + if step == 1 and world > 1: + lo, hi = (0, full.numel()) if rank == 0 else (0, 0) + master = full + step * 0.25 + shard = master.flatten()[lo:hi] if lo < hi else None + outputs = [ + ( + torch.empty(hi - lo, dtype=torch.uint8, device="cuda"), + torch.empty(hi - lo, dtype=torch.uint8, device="cuda") if i else None, + ) + for i in range(len(weights)) + ] + reduced_sizes.clear() + cast = quantize_master_weights if step == 0 else cast_master_weights_to_fp8 + with patch("torch.distributed.all_reduce", side_effect=record_reduce): + cast( + weights, + [shard] * len(weights), + [lo if shard is not None else None] * len(weights), + group, + fsdp_shard_model_weights=outputs if fragments else None, + ) + expected_amax = sum(w._rowwise_scale_inv.numel() for w in weights) + if mixed: + expected_amax += weights[1]._columnwise_scale_inv.numel() + assert reduced_sizes == [expected_amax] + for idx, (weight, quantizer) in enumerate(zip(weights, quantizers)): + expected = quantizer(master.to(torch.bfloat16)) + for direction in ["rowwise", "columnwise"] if idx else ["rowwise"]: + data = getattr(weight, f"_{direction}_data") + actual = torch.zeros_like(data).flatten() + if lo < hi: + src = ( + outputs[idx][direction == "columnwise"] + if fragments + else data.flatten()[lo:hi] + ) + actual[lo:hi].copy_(src) + real_reduce(actual, op=torch.distributed.ReduceOp.MAX, group=group) + torch.testing.assert_close( + actual.view(shape), getattr(expected, f"_{direction}_data"), rtol=0, atol=0 + ) + scale = getattr(weight, f"_{direction}_scale_inv") + ref = getattr(expected, f"_{direction}_scale_inv") + rows, cols = (160, 3) if direction == "rowwise" else (5, 96) + torch.testing.assert_close(scale[:rows, :cols], ref[:rows, :cols], rtol=0, atol=0) + assert pointers[idx] == ( + weight._rowwise_data.data_ptr(), + weight._rowwise_scale_inv.data_ptr(), + ) + assert weights[0]._columnwise_data is None + assert weights[0]._columnwise_scale_inv is None + + def run_parallel_tests() -> None: """Run parallel tests""" @@ -1000,6 +1081,9 @@ def run_parallel_tests() -> None: if is_mxfp8_available(): print("starting mxfp8 empty master shard test") _test_mxfp8_empty_master_shard(dp_group) + for mixed in (False, True): + for fragments in (False, True): + _test_rowwise_master_cast(dp_group, mixed, fragments) nvfp4_available, _ = is_nvfp4_available(return_reason=True) if nvfp4_available: print("starting cast master weights to nvfp4 test") diff --git a/tests/pytorch/test_partial_cast.py b/tests/pytorch/test_partial_cast.py index bbb18503b10..b7f683b222c 100644 --- a/tests/pytorch/test_partial_cast.py +++ b/tests/pytorch/test_partial_cast.py @@ -128,6 +128,7 @@ def _pad(x, padding): def test_mxfp8_scaling_partial_cast(): torch.cuda.manual_seed(1234) + run_one_case(0, 64, 64, 0) run_one_case(3, 32, 64, 31) run_one_case(64 * 64 - 2, 64, 64, 1) run_one_case(16384 * 6144, 16384, 6144, 0) @@ -135,3 +136,39 @@ def test_mxfp8_scaling_partial_cast(): run_one_case(131072, 768, 256, 0) run_one_case(65536, 768, 256, 131072) run_one_case(98304, 128, 768, 0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("offset,length", [(0, 0), (31, 3), (1, 4094)]) +@pytest.mark.parametrize("with_columnwise_scales", [False, True]) +def test_rowwise_partial_kernels(offset, length, with_columnwise_scales): + inp = torch.randn(length, device="cuda", dtype=torch.bfloat16) + row = torch.zeros((128, 4), device="cuda", dtype=inp.dtype) + ref = torch.zeros_like(row) + col = torch.zeros((4, 128), device="cuda", dtype=inp.dtype) + # The PyTorch extension must normalize empty views with backing storage. + omitted = col.flatten()[:0].view(0, 0) + tex.mxfp8_scaling_compute_partial_amax(inp, row, omitted, 64, 64, offset) + tex.mxfp8_scaling_compute_partial_amax(inp, ref, col, 64, 64, offset) + torch.testing.assert_close(row, ref, rtol=0, atol=0) + scales = torch.full((128, 4), 127, dtype=torch.uint8, device="cuda") + columnwise_scales = ( + torch.full((4, 128), 126, dtype=torch.uint8, device="cuda") + if with_columnwise_scales + else scales.flatten()[:0].view(0, 0) + ) + out = torch.empty(length, dtype=torch.uint8, device="cuda") + tex.mxfp8_scaling_partial_cast(inp, out, out[:0], scales, columnwise_scales, 64, 64, offset) + torch.testing.assert_close(out, inp.to(torch.float8_e4m3fn).view(torch.uint8), rtol=0, atol=0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_mxfp8_partial_cast_requires_columnwise_scales(): + inp = torch.randn(3, device="cuda", dtype=torch.bfloat16) + row = torch.empty(3, device="cuda", dtype=torch.uint8) + col = torch.empty_like(row) + scales = torch.full((128, 4), 127, device="cuda", dtype=torch.uint8) + with pytest.raises(RuntimeError, match="Columnwise output requires columnwise scales"): + tex.mxfp8_scaling_partial_cast( + inp, row, col, scales, scales.flatten()[:0].view(0, 0), 64, 64, 31 + ) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index b3feeba6b75..98c5eed0985 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -290,6 +290,7 @@ void nvte_fp8_block_scaling_partial_cast(const NVTETensor inp, NVTETensor out, * \param[in] input Input tensor (continuous segment of flattened original tensor). * \param[in,out] amax_rowwise Output tensor for row-wise maximum absolute values. * \param[in,out] amax_colwise Output tensor for column-wise maximum absolute values. + * An empty NVTETensor skips column-wise computation. * \param[in] rows Number of rows in the logical tensor. * \param[in] cols Number of columns in the logical tensor. * \param[in] start_offset Starting offset in the flattened tensor. @@ -303,7 +304,9 @@ void nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor * * This function casts the input tensor to MXFP8 format, producing both row-wise and * column-wise scaled outputs. input contains a continuous segment from the flattened - * original tensor. + * original tensor. To skip column-wise computation, pass an empty NVTETensor for + * output_colwise; scale_inv_colwise is ignored in this case. Column-wise output + * requires column-wise scales. Row-wise outputs remain required. * * \param[in] input Input (continuous segment of flattened original tensor). * \param[out] output_rowwise Output tensor with row-wise scaling (MXFP8 format). diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index be692d45633..b25e3d4b00b 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -23,7 +23,7 @@ constexpr int kColsPerTile = 128; // Columns each block processes constexpr int kThreadsPerBlock = 128; -template +template __global__ void __launch_bounds__(kThreadsPerBlock) mxfp8_scaling_compute_partial_amax_kernel(const IType *input, IType *amax_rowwise, IType *amax_colwise, int amax_rowwise_stride, @@ -47,7 +47,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { float abs_input = fabs(static_cast(input_minus_offset[idx])); row_amax = fmaxf(row_amax, abs_input); - col_amax = fmaxf(col_amax, abs_input); + if constexpr (kWithColumnwise) col_amax = fmaxf(col_amax, abs_input); } #pragma unroll @@ -63,7 +63,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock) r++; } - amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + if constexpr (kWithColumnwise) { + amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + } __syncthreads(); @@ -74,7 +76,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) amax_rowwise[r * amax_rowwise_stride + c] = static_cast(smem_amax_rowwise[r_][c_]); } -template +template __global__ void __launch_bounds__(kThreadsPerBlock) mxfp8_scaling_partial_cast_kernel(const IType *input, OType *output_rowwise, OType *output_colwise, const e8m0_t *scale_inv_rowwise, @@ -95,7 +97,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) } // Load scales_colwise - { + if constexpr (kWithColumnwise) { int c_ = threadIdx.x; int r = blockIdx.y * kRowsPerTile / 32; int c = blockIdx.x * kColsPerTile + c_; @@ -108,7 +110,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) size_t end_offset = start_offset + len; const IType *input_minus_offset = input - start_offset; OType *output_rowwise_minus_offset = output_rowwise - start_offset; - OType *output_colwise_minus_offset = output_colwise - start_offset; + OType *output_colwise_minus_offset = kWithColumnwise ? output_colwise - start_offset : nullptr; int warp_idx = threadIdx.x / 32; // int lane_idx = threadIdx.x % 32; int c = blockIdx.x * kColsPerTile + threadIdx.x; @@ -121,9 +123,11 @@ __global__ void __launch_bounds__(kThreadsPerBlock) if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { float inp = static_cast(input_minus_offset[idx]); OType out_rowwise = static_cast(inp * smem_scales_rowwise[i][warp_idx]); - OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); output_rowwise_minus_offset[idx] = out_rowwise; - output_colwise_minus_offset[idx] = out_colwise; + if constexpr (kWithColumnwise) { + OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); + output_colwise_minus_offset[idx] = out_colwise; + } } r++; @@ -149,14 +153,17 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, NVTE_CHECK(amax_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); NVTE_CHECK(amax_rowwise.dtype() == input.dtype(), "Wrong dtype of amax_rowwise"); - NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); - NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, - "Wrong padding of amax_colwise's rows"); - NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, "Invalid rows"); - NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0, - "Wrong padding of amax_colwise's cols"); - NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols"); - NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise"); + const bool with_columnwise = amax_colwise.has_data(); + if (with_columnwise) { + NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); + NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of amax_colwise's rows"); + NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of amax_colwise's cols"); + NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise"); + } int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; @@ -164,11 +171,15 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, - mxfp8_scaling_compute_partial_amax_kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(amax_rowwise.data.dptr), - reinterpret_cast(amax_colwise.data.dptr), amax_rowwise.data.shape[1], - amax_colwise.data.shape[1], rows, cols, start_offset, input.data.shape[0]);) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_columnwise, kWithColumnwise, + mxfp8_scaling_compute_partial_amax_kernel + <<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(amax_rowwise.data.dptr), + with_columnwise ? reinterpret_cast(amax_colwise.data.dptr) : nullptr, + amax_rowwise.data.shape[1], with_columnwise ? amax_colwise.data.shape[1] : 0, rows, + cols, start_offset, input.data.shape[0]);)) } void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tensor output_colwise, @@ -182,16 +193,23 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso "Invalid start_offset"); NVTE_CHECK(output_rowwise.data.shape.size() == 1, "output_rowwise must be a 1D tensor"); - NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); NVTE_CHECK(output_rowwise.data.shape[0] == input.data.shape[0], "Size of input and output_rowwise mismatch"); - NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], - "Size of input and output_colwise mismatch"); NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte, "output_rowwise should be e4m3 or uint8"); - NVTE_CHECK(output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, - "output_colwise should be e4m3 or uint8"); + const bool with_columnwise = output_colwise.has_data(); + if (with_columnwise) { + NVTE_CHECK(scale_inv_colwise.has_data(), "Columnwise output requires columnwise scales"); + NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); + NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], + "Size of input and output_colwise mismatch"); + NVTE_CHECK(input.data.shape[0] == 0 || output_colwise.data.dptr != nullptr, + "Columnwise output must have a data pointer"); + NVTE_CHECK( + output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, + "output_colwise should be e4m3 or uint8"); + } NVTE_CHECK(scale_inv_rowwise.data.shape.size() == 2, "scale_inv_rowwise must be a 2D tensor"); NVTE_CHECK(scale_inv_rowwise.data.shape[0] % rowwise_row_padding == 0, @@ -202,14 +220,16 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso NVTE_CHECK(scale_inv_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); NVTE_CHECK(scale_inv_rowwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_rowwise"); - NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, "scale_inv_colwise must be a 2D tensor"); - NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0, - "Wrong padding of scale_inv_colwise's rows"); - NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, "Invalid rows"); - NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0, - "Wrong padding of scale_inv_colwise's cols"); - NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, "Invalid cols"); - NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_colwise"); + if (with_columnwise) { + NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, "scale_inv_colwise must be a 2D tensor"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of scale_inv_colwise's rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of scale_inv_colwise's cols"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_colwise"); + } int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; @@ -217,14 +237,19 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, - mxfp8_scaling_partial_cast_kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(output_rowwise.data.dptr), - reinterpret_cast(output_colwise.data.dptr), - reinterpret_cast(scale_inv_rowwise.data.dptr), - reinterpret_cast(scale_inv_colwise.data.dptr), - scale_inv_rowwise.data.shape[1], scale_inv_colwise.data.shape[1], rows, cols, - start_offset, input.data.shape[0]);) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_columnwise, kWithColumnwise, + mxfp8_scaling_partial_cast_kernel + <<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output_rowwise.data.dptr), + with_columnwise ? reinterpret_cast(output_colwise.data.dptr) : nullptr, + reinterpret_cast(scale_inv_rowwise.data.dptr), + with_columnwise ? reinterpret_cast(scale_inv_colwise.data.dptr) + : nullptr, + scale_inv_rowwise.data.shape[1], + with_columnwise ? scale_inv_colwise.data.shape[1] : 0, rows, cols, start_offset, + input.data.shape[0]);)) } } // namespace mxfp8_scaling_recipe diff --git a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp index d6693a485ee..dc02e67e5fb 100644 --- a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp @@ -57,7 +57,9 @@ void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax const TensorWrapper input_cu = makeTransformerEngineTensor(input); TensorWrapper amax_rowwise_cu = makeTransformerEngineTensor(amax_rowwise); - TensorWrapper amax_colwise_cu = makeTransformerEngineTensor(amax_colwise); + // Empty PyTorch views may have a backing pointer. Normalize omission at this boundary. + TensorWrapper amax_colwise_cu = + amax_colwise.numel() == 0 ? TensorWrapper() : makeTransformerEngineTensor(amax_colwise); nvte_mxfp8_scaling_compute_partial_amax(input_cu.data(), amax_rowwise_cu.data(), amax_colwise_cu.data(), rows, cols, start_offset, @@ -72,13 +74,19 @@ void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwi TORCH_CHECK(output_rowwise.is_contiguous(), "output_rowwise must be contiguous"); TORCH_CHECK(output_colwise.is_contiguous(), "output_colwise must be contiguous"); TORCH_CHECK(scale_inv_rowwise.is_contiguous(), "scale_inv_rowwise must be contiguous"); - TORCH_CHECK(scale_inv_colwise.is_contiguous(), "scale_inv_colwise must be contiguous"); const TensorWrapper input_cu = makeTransformerEngineTensor(input); TensorWrapper output_rowwise_cu = makeTransformerEngineTensor(output_rowwise); - TensorWrapper output_colwise_cu = makeTransformerEngineTensor(output_colwise); + const bool with_columnwise = output_colwise.numel() != 0; + if (with_columnwise) { + TORCH_CHECK(scale_inv_colwise.numel() != 0, "Columnwise output requires columnwise scales"); + TORCH_CHECK(scale_inv_colwise.is_contiguous(), "scale_inv_colwise must be contiguous"); + } + TensorWrapper output_colwise_cu = + with_columnwise ? makeTransformerEngineTensor(output_colwise) : TensorWrapper(); const TensorWrapper scale_inv_rowwise_cu = makeTransformerEngineTensor(scale_inv_rowwise); - const TensorWrapper scale_inv_colwise_cu = makeTransformerEngineTensor(scale_inv_colwise); + const TensorWrapper scale_inv_colwise_cu = + with_columnwise ? makeTransformerEngineTensor(scale_inv_colwise) : TensorWrapper(); nvte_mxfp8_scaling_partial_cast(input_cu.data(), output_rowwise_cu.data(), output_colwise_cu.data(), scale_inv_rowwise_cu.data(), diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index ccad043a656..caaf9cf5aad 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1029,6 +1029,8 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( ): # pylint: disable=unused-argument r"""Helper function to cast master weights to FP8 primary weights for mxfp8 scaling. + Row-only primaries omit columnwise amax, scale updates, and output writes. + Parameters ---------- params : List of tuple, each tuple contains a model weight, a master weight, and an offset @@ -1058,7 +1060,13 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( raise ValueError( f"rowwise_scale_inv must be 2D, got {len(rowwise_shape)}D shape {rowwise_shape}" ) - colwise_shape = model_weight._columnwise_scale_inv.shape + if (model_weight._columnwise_data is None) != (model_weight._columnwise_scale_inv is None): + raise ValueError("MXFP8 columnwise data and scales must both be present or absent") + colwise_shape = ( + model_weight._columnwise_scale_inv.shape + if model_weight._columnwise_scale_inv is not None + else (0, 0) + ) if len(colwise_shape) != 2: raise ValueError( f"columnwise_scale_inv must be 2D, got {len(colwise_shape)}D shape {colwise_shape}" @@ -1087,7 +1095,11 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( amaxes_colwise, scale_invs_colwise = [], [] for i, (model_weight, master_weight, start_offset, _) in enumerate(params): rowwise_shape = model_weight._rowwise_scale_inv.shape - colwise_shape = model_weight._columnwise_scale_inv.shape + colwise_shape = ( + model_weight._columnwise_scale_inv.shape + if model_weight._columnwise_scale_inv is not None + else (0, 0) + ) rowwise_start = cu_rowwise_amax_sizes[i] rowwise_end = cu_rowwise_amax_sizes[i + 1] colwise_start = cu_rowwise_amax_sizes[-1] + cu_colwise_amax_sizes[i] @@ -1097,7 +1109,12 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( amaxes_rowwise.append(amax_rowwise) amaxes_colwise.append(amax_colwise) scale_invs_rowwise.append(model_weight._rowwise_scale_inv) - scale_invs_colwise.append(model_weight._columnwise_scale_inv) + # Empty tensors mark an omitted direction without allocating a GPU payload. + scale_invs_colwise.append( + model_weight._columnwise_scale_inv + if model_weight._columnwise_scale_inv is not None + else torch.empty((0, 0), dtype=torch.uint8, device=device) + ) # Compute amax of the master weight and store it in packed_amaxes. if master_weight is not None: @@ -1123,8 +1140,8 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( multi_tensor_compute_scale_inv_e8m0, None, # dummy_overflow_buf [ - amaxes_rowwise + amaxes_colwise, - scale_invs_rowwise + scale_invs_colwise, + amaxes_rowwise + [amax for amax in amaxes_colwise if amax.numel()], + scale_invs_rowwise + [scale for scale in scale_invs_colwise if scale.numel()], ], ) @@ -1145,10 +1162,18 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( end_offset = start_offset + master_weight.numel() if use_fsdp_shard_model_weights: rowwise_fragment = model_weight_fragment[0] - colwise_fragment = model_weight_fragment[1] + colwise_fragment = ( + model_weight_fragment[1] + if scale_inv_colwise.numel() + else torch.empty(0, dtype=torch.uint8, device=device) + ) else: rowwise_fragment = model_weight._rowwise_data.reshape(-1)[start_offset:end_offset] - colwise_fragment = model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] + colwise_fragment = ( + model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] + if scale_inv_colwise.numel() + else torch.empty(0, dtype=torch.uint8, device=device) + ) if len(model_weight.shape) != 2: raise ValueError( "model_weight must be 2D for MXFP8 scaling partial cast, " From 0f8d92f5a7d7459799ec860fa1eba1ec767ce6c8 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:08:01 -0700 Subject: [PATCH 05/32] [PyTorch] Set device during NCCL init in symmetric memory test (#3744) Set device during NCCL init in symmetric memory test Signed-off-by: Tim Moon --- tests/pytorch/distributed/test_symmetric_memory.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/distributed/test_symmetric_memory.py b/tests/pytorch/distributed/test_symmetric_memory.py index 10497a4dd36..69c3ea2e424 100644 --- a/tests/pytorch/distributed/test_symmetric_memory.py +++ b/tests/pytorch/distributed/test_symmetric_memory.py @@ -15,7 +15,13 @@ def _check_cache(rank, init_file): torch.cuda.set_device(rank) - dist.init_process_group("nccl", init_method=f"file://{init_file}", rank=rank, world_size=2) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=2, + device_id=torch.device("cuda", rank), + ) try: group = dist.group.WORLD stream = torch.cuda.Stream() From 89668d9dee82aba9b9107b37ed85227542a578d2 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Tue, 27 Oct 2026 14:17:17 +0800 Subject: [PATCH 06/32] [Common][PyTorch] Add NVFP4 row-scaled grouped cast with graph-safe path (#3649) * [Common][PyTorch] Add NVFP4 row-scaled grouped cast with graph-safe path Single-launch grouped row-scaled NVFP4 cast + fused amax, plus graph-safe amax/cast variants driven by device first_dims/tensor_offsets for capture-once/replay MoE. Wired into PyTorch (ep.nvfp4_group_row_scaled_cast) with C++ and PyTorch capture/replay tests. Signed-off-by: Cael Ling * Address review: guard grouped row-scaled cast dispatch and inputs Restrict the FUSED_NVFP4_ROW_SCALED selection to shapes the grouped kernels accept (128-aligned last dim, <= 64 experts) so unsupported requests fall back to the per-expert path. Return early for empty routing, make the graph-safe binding's input contiguous, and reject with_rht/with_2d_quantization/with_amax_reduction on both row-scaled entries. Add regression tests for the four guarded paths. Signed-off-by: Cael Ling * Guard row-scaled grouped cast against non-2D, partial, and strided routing The grouped row-scaled cast treats its input as 2D [rows, last_dim] indexed by split counts, so route higher-rank inputs and partial splits (split sum below the row count) to the per-expert path instead of the grouped kernel. Make first_dims and tensor_offsets contiguous in the graph-safe cast so a strided routing view is not misread. Add fallback and strided-routing tests. Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fall back to per-expert cast for mixed NVFP4 row/tensor-scaled splits The bulk NVFP4 split_quantize allocation sizes its amax buffer from the first quantizer's row-scale setting, so a group mixing row-scaled and tensor-scaled quantizers writes per-row amax into a per-tensor buffer and corrupts neighbors. Detect the mixed case and fall back to the unfused per-split path. Signed-off-by: Cael Ling * Route row-scaled NVFP4 grouped quantize through nvte_group_quantize Add the kNVTEGroupedRowScaledNVFP4 grouped tensor param and dispatch row-scaled NVFP4 outputs in nvte_group_quantize to the graph-safe amax and cast kernels. Handle row-scaled quantizers in tex.group_quantize. Remove the row-scaled-specific C APIs, the host-split row-scaled cast kernel, tex.nvfp4_group_row_scaled_cast_graph_safe, the ep.py helper, and the fused row-scaled split_quantize path. Add a top-level static_assert for unregistered types in TypeInfo. Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Require 128-row experts for uniform row-scaled NVFP4 group_quantize Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fold row-scaled NVFP4 amax into nvte_group_amax Dispatch nvte_group_amax and nvte_group_amax_graph_safe on the output tensor format so that row-scaled NVFP4 outputs get per-row amaxes, and deprecate nvte_group_nvfp4_compute_amax in favor of nvte_group_amax. Rename the row-scaled amax helpers to group_row_amax(_graph_safe). In split_quantize, fall back to unfused allocation and quantization whenever any NVFP4 quantizer is row-scaled. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix missing include Review suggestion from @greptile-apps Signed-off-by: Tim Moon * Add missing header include Signed-off-by: Tim Moon --------- Signed-off-by: Cael Ling Signed-off-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon Co-authored-by: Claude Opus 5.5 --- tests/cpp/operator/CMakeLists.txt | 1 + .../test_row_scaled_nvfp4_grouped_amax.cu | 12 +- .../test_row_scaled_nvfp4_grouped_cast.cu | 419 +++++++++++++ tests/cpp/test_common.h | 11 +- .../test_nvfp4_group_row_scaled_quantize.py | 280 +++++++++ .../common/cast/cast_grouped.cu | 14 +- .../common/cast/dispatch/quantize.cuh | 23 + .../nvfp4/group_quantize_transpose_nvfp4.cuh | 570 +++++++++++++++++- transformer_engine/common/common.h | 10 +- .../graph_safe_group_hadamard_transform.cu | 11 +- .../group_hadamard_transform.cu | 28 +- .../common/include/transformer_engine/cast.h | 18 +- .../include/transformer_engine/multi_tensor.h | 12 +- .../transformer_engine/transformer_engine.h | 13 + .../common/transformer_engine.cpp | 6 + .../pytorch/csrc/extensions/cast.cpp | 69 ++- transformer_engine/pytorch/csrc/quantizer.cpp | 1 + .../pytorch/csrc/type_converters.cpp | 4 + 18 files changed, 1430 insertions(+), 72 deletions(-) create mode 100644 tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu create mode 100644 tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 4727a2be295..6377fad36c2 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -16,6 +16,7 @@ add_executable(test_operator test_cast_mxfp8_grouped_scaled_swiglu.cu test_cast_nvfp4_transpose.cu test_row_scaled_nvfp4_grouped_amax.cu + test_row_scaled_nvfp4_grouped_cast.cu test_cast_float8blockwise.cu test_cast_float8blockwise_grouped.cu test_dequantize_mxfp8.cu diff --git a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu index 8b883f25e25..226e5c4c2ed 100644 --- a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu +++ b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu @@ -17,7 +17,7 @@ #include #include -#include +#include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -80,12 +80,12 @@ void performGroupedFusedAmaxTest(const GroupedAmaxConfig& cfg) { auto outs = make_row_scaled_outputs("gout_", Ms, K); std::vector out_handles; for (auto& t : outs) out_handles.push_back(t->data()); - nvte_group_nvfp4_compute_amax(packed.data(), out_handles.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), out_handles.data(), splits.data(), num_tensors, 0); auto outs2 = make_row_scaled_outputs("gout2_", Ms, K); std::vector out_handles2; for (auto& t : outs2) out_handles2.push_back(t->data()); - nvte_group_nvfp4_compute_amax(packed.data(), out_handles2.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), out_handles2.data(), splits.data(), num_tensors, 0); cudaDeviceSynchronize(); ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); @@ -187,7 +187,7 @@ TEST(NVFP4GroupedRowScaledAmaxTestSuite, HeterogeneousColumnwise) { std::vector handles; for (auto& t : outs) handles.push_back(t->data()); const std::vector splits(Ms); - nvte_group_nvfp4_compute_amax(packed.data(), handles.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), handles.data(), splits.data(), num_tensors, 0); cudaDeviceSynchronize(); ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); @@ -265,7 +265,7 @@ void runLoopVsGroupedBench(size_t num_tensors, size_t M, size_t K, int iters) { double grouped = median_ms( [&]() { - nvte_group_nvfp4_compute_amax(packed.data(), h_g.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), h_g.data(), splits.data(), num_tensors, 0); }, iters); @@ -274,7 +274,7 @@ void runLoopVsGroupedBench(size_t num_tensors, size_t M, size_t K, int iters) { for (size_t i = 0; i < num_tensors; ++i) { NVTETensor one = outs_l[i]->data(); size_t s1 = M; - nvte_group_nvfp4_compute_amax(ins[i]->data(), &one, &s1, 1, 0); + nvte_group_amax(ins[i]->data(), &one, &s1, 1, 0); } }, iters); diff --git a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu new file mode 100644 index 00000000000..50b57304b0f --- /dev/null +++ b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu @@ -0,0 +1,419 @@ +/************************************************************************* + * 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 + +#include +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +// All row counts and K must be multiples of 128. +struct GroupedCastConfig { + std::vector Ms; + size_t K; + bool transpose; +}; + +std::vector> make_outputs(const std::string& tag, + const std::vector& Ms, size_t K, + bool transpose) { + std::vector> outs; + for (size_t i = 0; i < Ms.size(); ++i) { + const std::vector shape = {Ms[i], K}; + auto t = std::make_unique(tag + std::to_string(i), shape, DType::kFloat4E2M1, + /*rowwise=*/true, /*columnwise=*/transpose, + NVTE_NVFP4_1D_SCALING); + t->set_row_scaled_nvfp4(true); + outs.push_back(std::move(t)); + } + return outs; +} + +std::vector> make_random_inputs(const std::vector& Ms, size_t K) { + std::vector> ins; + for (size_t i = 0; i < Ms.size(); ++i) { + ins.push_back( + std::make_unique("in_" + std::to_string(i), std::vector{Ms[i], K}, + DType::kBFloat16)); + fillCase(ins[i].get(), InputsFillCase::uniform); + ins[i]->to_cpu(); + } + return ins; +} + +std::vector pack_inputs(const std::vector>& ins, + const std::vector& Ms, size_t K) { + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + std::vector packed(sum_M * K); + size_t row_off = 0; + for (size_t i = 0; i < ins.size(); ++i) { + std::copy_n(ins[i]->rowwise_cpu_dptr(), Ms[i] * K, packed.data() + row_off * K); + row_off += Ms[i]; + } + return packed; +} + +struct GroupedReference { + size_t sum_M = 0; + size_t K = 0; + size_t row_scale_cols = 0; // rowwise E4M3 scale stride (shared across experts) + std::vector row_data; + std::vector row_scale; + std::vector row_amax; + std::vector col_data; + std::vector col_scale; + std::vector col_amax; + bool transpose = false; +}; + +GroupedReference build_reference(const std::vector>& ins, + const std::vector& Ms, size_t K, bool transpose) { + const size_t num_tensors = Ms.size(); + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + + auto ref = make_outputs("ref_", Ms, K, transpose); + QuantizationConfigWrapper quant_config; + quant_config.set_stochastic_rounding(false); + for (size_t i = 0; i < num_tensors; ++i) { + nvte_quantize_v2(ins[i]->data(), ref[i]->data(), quant_config, 0); + } + cudaDeviceSynchronize(); + + GroupedReference r; + r.sum_M = sum_M; + r.K = K; + r.transpose = transpose; + r.row_scale_cols = K / 16; // contiguous (unpadded) rowwise E4M3 scale stride + + r.row_data.resize(sum_M * K / 2); + r.row_scale.resize(sum_M * r.row_scale_cols); + r.row_amax.resize(sum_M); + if (transpose) { + r.col_data.resize((K / 2) * sum_M); + r.col_scale.resize((K / 16) * sum_M); + r.col_amax.resize(num_tensors * K); + } + + auto gather_scale = [](const fp8e4m3* src, fp8e4m3* dst, size_t rows, size_t cols, + size_t src_stride) { + for (size_t row = 0; row < rows; ++row) { + std::copy_n(src + row * src_stride, cols, dst + row * cols); + } + }; + + size_t row_off = 0; + size_t col_data_off = 0; + size_t col_scale_off = 0; + for (size_t i = 0; i < num_tensors; ++i) { + ref[i]->to_cpu(); + const size_t M = Ms[i]; + std::copy_n(reinterpret_cast(ref[i]->rowwise_cpu_dptr()), M * K / 2, + r.row_data.data() + row_off * (K / 2)); + const std::array sd = get_scale_tensor_dims(M, K, 1, 16); + gather_scale(ref[i]->rowwise_cpu_scale_inv_ptr(), + reinterpret_cast(r.row_scale.data()) + row_off * r.row_scale_cols, sd[0], + sd[1], sd[3]); + std::copy_n(ref[i]->cpu_rowwise_amax_ptr(), M, r.row_amax.data() + row_off); + if (transpose) { + std::copy_n(reinterpret_cast(ref[i]->columnwise_cpu_dptr()), + K * M / 2, r.col_data.data() + col_data_off); + const std::array sd_t = get_scale_tensor_dims(K, M, 1, 16); + gather_scale(ref[i]->columnwise_cpu_scale_inv_ptr(), + reinterpret_cast(r.col_scale.data()) + col_scale_off, sd_t[0], sd_t[1], + sd_t[3]); + std::copy_n(ref[i]->cpu_columnwise_amax_ptr(), K, r.col_amax.data() + i * K); + col_data_off += K * M / 2; + col_scale_off += K * (M / 16); + } + row_off += M; + } + return r; +} + +struct GroupedBuffersNVFP4 { + size_t capacity_rows = 0; + size_t K = 0; + size_t num_tensors = 0; + size_t row_scale_cols = 0; + bool transpose = false; + + CudaPtr in_data; + CudaPtr out_data; + CudaPtr out_scale; + CudaPtr out_amax; + CudaPtr out_col_data; + CudaPtr out_col_scale; + CudaPtr out_col_amax; + CudaPtr offsets; + CudaPtr first_dims; + + std::unique_ptr input; + std::unique_ptr output; + + void allocate(size_t capacity_rows_, size_t K_, size_t num_tensors_, size_t row_scale_cols_, + bool transpose_) { + capacity_rows = capacity_rows_; + K = K_; + num_tensors = num_tensors_; + row_scale_cols = row_scale_cols_; + transpose = transpose_; + + in_data = cuda_alloc(capacity_rows * K * sizeof(bf16)); + out_data = cuda_alloc(capacity_rows * (K / 2)); + out_scale = cuda_alloc(capacity_rows * row_scale_cols); + out_amax = cuda_alloc(capacity_rows * sizeof(float)); + offsets = cuda_alloc((num_tensors + 1) * sizeof(int64_t)); + first_dims = cuda_alloc(num_tensors * sizeof(int64_t)); + if (transpose) { + out_col_data = cuda_alloc((K / 2) * capacity_rows); + out_col_scale = cuda_alloc((K / 16) * capacity_rows); + out_col_amax = cuda_alloc(num_tensors * K * sizeof(float)); + } + + input = std::make_unique(num_tensors, + std::vector{capacity_rows, K}); + output = std::make_unique( + num_tensors, std::vector{capacity_rows, K}, NVTE_NVFP4_1D_SCALING); + + input->set_rowwise_data(in_data.get(), DType::kBFloat16, + std::vector{capacity_rows, K}); + input->set_tensor_offsets(offsets.get(), DType::kInt64, + std::vector{num_tensors + 1}); + input->set_first_dims(first_dims.get(), DType::kInt64, std::vector{num_tensors}); + + output->set_row_scaled_nvfp4(true); + output->set_rowwise_data(out_data.get(), DType::kFloat4E2M1, + std::vector{capacity_rows, K}); + output->set_rowwise_scale_inv(out_scale.get(), DType::kFloat8E4M3, + std::vector{capacity_rows, row_scale_cols}); + output->set_amax(out_amax.get(), DType::kFloat32, std::vector{capacity_rows}); + output->set_tensor_offsets(offsets.get(), DType::kInt64, + std::vector{num_tensors + 1}); + output->set_first_dims(first_dims.get(), DType::kInt64, std::vector{num_tensors}); + if (transpose) { + output->set_columnwise_data(out_col_data.get(), DType::kFloat4E2M1, + std::vector{K, capacity_rows}); + output->set_columnwise_scale_inv(out_col_scale.get(), DType::kFloat8E4M3, + std::vector{K, capacity_rows / 16}); + output->set_columnwise_amax(out_col_amax.get(), DType::kFloat32, + std::vector{num_tensors * K}); + } + } + + void set_routing(const std::vector& Ms, const GroupedReference& ref, + const std::vector& packed_in) { + std::vector off(num_tensors + 1, 0); + std::vector fd(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + fd[i] = static_cast(Ms[i]); + off[i + 1] = off[i] + static_cast(Ms[i] * K); + } + NVTE_CHECK_CUDA(cudaMemcpy(offsets.get(), off.data(), (num_tensors + 1) * sizeof(int64_t), + cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(first_dims.get(), fd.data(), num_tensors * sizeof(int64_t), + cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(in_data.get(), packed_in.data(), ref.sum_M * K * sizeof(bf16), + cudaMemcpyHostToDevice)); + std::vector poison_row(capacity_rows, -1.0f); + NVTE_CHECK_CUDA(cudaMemcpy(out_amax.get(), poison_row.data(), capacity_rows * sizeof(float), + cudaMemcpyHostToDevice)); + if (transpose) { + std::vector poison_col(num_tensors * K, -1.0f); + NVTE_CHECK_CUDA(cudaMemcpy(out_col_amax.get(), poison_col.data(), + num_tensors * K * sizeof(float), cudaMemcpyHostToDevice)); + } + } + + void quantize(cudaStream_t stream) { + QuantizationConfigWrapper quant_config; + nvte_group_quantize(input->data(), output->data(), quant_config, stream); + } +}; + +void expect_matches_reference(const GroupedBuffersNVFP4& b, const GroupedReference& ref) { + const size_t sum_M = ref.sum_M; + const size_t K = ref.K; + + std::vector h_data(sum_M * K / 2); + std::vector h_scale(sum_M * ref.row_scale_cols); + std::vector h_amax(sum_M); + NVTE_CHECK_CUDA(cudaMemcpy(h_data.data(), b.out_data.get(), h_data.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA( + cudaMemcpy(h_scale.data(), b.out_scale.get(), h_scale.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA( + cudaMemcpy(h_amax.data(), b.out_amax.get(), sum_M * sizeof(float), cudaMemcpyDeviceToHost)); + + for (size_t i = 0; i < sum_M; ++i) + ASSERT_EQ(h_amax[i], ref.row_amax[i]) << "rowwise amax mismatch at row " << i; + size_t data_mismatch = 0; + for (size_t i = 0; i < h_data.size(); ++i) + if (h_data[i] != ref.row_data[i]) ++data_mismatch; + ASSERT_EQ(data_mismatch, 0u) << "rowwise FP4 data mismatch (" << data_mismatch << "/" + << h_data.size() << ")"; + size_t scale_mismatch = 0; + for (size_t i = 0; i < h_scale.size(); ++i) + if (h_scale[i] != ref.row_scale[i]) ++scale_mismatch; + ASSERT_EQ(scale_mismatch, 0u) << "rowwise scale mismatch (" << scale_mismatch << "/" + << h_scale.size() << ")"; + + if (!ref.transpose) return; + + std::vector h_col_amax(b.num_tensors * K); + std::vector h_col_data((K / 2) * sum_M); + std::vector h_col_scale((K / 16) * sum_M); + NVTE_CHECK_CUDA(cudaMemcpy(h_col_amax.data(), b.out_col_amax.get(), + h_col_amax.size() * sizeof(float), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(h_col_data.data(), b.out_col_data.get(), h_col_data.size(), + cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(h_col_scale.data(), b.out_col_scale.get(), h_col_scale.size(), + cudaMemcpyDeviceToHost)); + for (size_t i = 0; i < h_col_amax.size(); ++i) + ASSERT_EQ(h_col_amax[i], ref.col_amax[i]) << "columnwise amax mismatch at index " << i; + size_t col_data_mismatch = 0; + for (size_t i = 0; i < h_col_data.size(); ++i) + if (h_col_data[i] != ref.col_data[i]) ++col_data_mismatch; + ASSERT_EQ(col_data_mismatch, 0u) << "columnwise FP4 data mismatch (" << col_data_mismatch << "/" + << h_col_data.size() << ")"; + size_t col_scale_mismatch = 0; + for (size_t i = 0; i < h_col_scale.size(); ++i) + if (h_col_scale[i] != ref.col_scale[i]) ++col_scale_mismatch; + ASSERT_EQ(col_scale_mismatch, 0u) << "columnwise scale mismatch (" << col_scale_mismatch << "/" + << h_col_scale.size() << ")"; +} + +void performGroupedRowScaledQuantizeTest(const GroupedCastConfig& cfg) { + const std::vector& Ms = cfg.Ms; + const size_t K = cfg.K; + const bool transpose = cfg.transpose; + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + + auto ins = make_random_inputs(Ms, K); + const GroupedReference ref = build_reference(ins, Ms, K, transpose); + const std::vector packed_in = pack_inputs(ins, Ms, K); + + GroupedBuffersNVFP4 b; + b.allocate(sum_M, K, Ms.size(), ref.row_scale_cols, transpose); + b.set_routing(Ms, ref, packed_in); + + b.quantize(0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + + expect_matches_reference(b, ref); +} + +void performGroupedRowScaledQuantizeReplayTest(size_t K, bool transpose) { + const std::vector Ms_a = {256, 128, 128}; + const std::vector Ms_b = {128, 256, 128}; + const size_t num_tensors = Ms_a.size(); + const size_t capacity_rows = 768; // > sum_M (512): exercises the capacity noop + + auto ins_a = make_random_inputs(Ms_a, K); + auto ins_b = make_random_inputs(Ms_b, K); + const GroupedReference ref_a = build_reference(ins_a, Ms_a, K, transpose); + const GroupedReference ref_b = build_reference(ins_b, Ms_b, K, transpose); + const std::vector packed_a = pack_inputs(ins_a, Ms_a, K); + const std::vector packed_b = pack_inputs(ins_b, Ms_b, K); + + GroupedBuffersNVFP4 b; + b.allocate(capacity_rows, K, num_tensors, ref_a.row_scale_cols, transpose); + b.set_routing(Ms_a, ref_a, packed_a); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); + + // Warm up outside capture so one-time setup (e.g. cudaFuncSetAttribute) does + // not land in the captured graph. + b.quantize(stream); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + b.set_routing(Ms_a, ref_a, packed_a); + + cudaGraph_t graph; + cudaGraphExec_t exec; + NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed)); + b.quantize(stream); + NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph)); + NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + + NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + expect_matches_reference(b, ref_a); + + b.set_routing(Ms_b, ref_b, packed_b); + NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + expect_matches_reference(b, ref_b); + + NVTE_CHECK_CUDA(cudaGraphExecDestroy(exec)); + NVTE_CHECK_CUDA(cudaGraphDestroy(graph)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + +} // namespace + +class NVFP4GroupedRowScaledQuantizeTestSuite + : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GroupedRowScaledQuantizeTestSuite, MatchesPerExpertQuantize) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGroupedRowScaledQuantizeTest(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + NVFP4GroupedRowScaledQuantizeRowwise, NVFP4GroupedRowScaledQuantizeTestSuite, + ::testing::Values(GroupedCastConfig{{128}, 128, false}, + GroupedCastConfig{{128, 128}, 128, false}, + GroupedCastConfig{{128, 256, 128}, 256, false}, + GroupedCastConfig{{256, 128}, 512, false}, + GroupedCastConfig{{384, 128, 256}, 128, false}, + GroupedCastConfig{{512}, 1024, false})); + +INSTANTIATE_TEST_SUITE_P( + NVFP4GroupedRowScaledQuantizeTranspose, NVFP4GroupedRowScaledQuantizeTestSuite, + ::testing::Values(GroupedCastConfig{{128}, 128, true}, + GroupedCastConfig{{128, 128}, 128, true}, + GroupedCastConfig{{128, 256, 128}, 256, true}, + GroupedCastConfig{{256, 128}, 512, true}, + GroupedCastConfig{{384, 128, 256}, 128, true}, + GroupedCastConfig{{512}, 1024, true})); + +class NVFP4GroupedRowScaledQuantizeReplaySuite + : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GroupedRowScaledQuantizeReplaySuite, ReplayWithChangedRouting) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGroupedRowScaledQuantizeReplayTest(GetParam().K, GetParam().transpose); +} + +INSTANTIATE_TEST_SUITE_P(NVFP4GroupedRowScaledQuantizeReplay, + NVFP4GroupedRowScaledQuantizeReplaySuite, + ::testing::Values(GroupedCastConfig{{}, 256, false}, + GroupedCastConfig{{}, 256, true})); diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 1b5b95d7314..58ea430d8a6 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -108,11 +108,12 @@ struct TypeInfo { struct Helper { constexpr static DType getType() { constexpr int i = static_cast(current); - if (std::is_same::type>::value) { - return current; - } else { - return Helper(i + 1)>::getType(); + if constexpr (i < static_cast(std::tuple_size::value)) { + if (std::is_same::type>::value) { + return current; + } } + return Helper(i + 1)>::getType(); } }; @@ -125,6 +126,8 @@ struct TypeInfo { template constexpr static DType getType() { + static_assert(Helper::getType() != DType::kNumTypes, + "Type is not registered"); return Helper::getType(); } diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py new file mode 100644 index 00000000000..a740e0ad2cc --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py @@ -0,0 +1,280 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped row-scaled NVFP4 quantization through tex.group_quantize.""" + +from typing import Optional + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch import NVFP4Quantizer + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def _row_scaled_quantizer(**kwargs) -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + **kwargs, + ) + + +def _tensor_scaled_quantizer() -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=False, + ) + + +def _group_quantize(x: torch.Tensor, num_tensors: int, first_dims: Optional[torch.Tensor]): + return tex.group_quantize(x, _row_scaled_quantizer(), num_tensors, first_dims) + + +def _per_expert_reference(x_valid: torch.Tensor, split_sections: list[int]): + N = x_valid.shape[1] + data, scale, amax = [], [], [] + offset = 0 + for m in split_sections: + if m == 0: + continue + ref = _row_scaled_quantizer()(x_valid[offset : offset + m]) + data.append(ref._rowwise_data.view(torch.uint8).reshape(m, N // 2)) + scale.append(ref._rowwise_scale_inv.view(torch.uint8).reshape(m, -1)[:, : N // 16]) + amax.append(ref._amax_rowwise.reshape(m)) + offset += m + return torch.cat(data, dim=0), torch.cat(scale, dim=0), torch.cat(amax, dim=0) + + +def _assert_matches_reference(grouped, x_valid: torch.Tensor, split_sections: list[int]) -> None: + N = x_valid.shape[1] + sum_m = sum(split_sections) + ref_data, ref_scale, ref_amax = _per_expert_reference(x_valid, split_sections) + + capacity = int(grouped.logical_shape[0]) + scale_stride = grouped.scale_inv.numel() // capacity + g_data = grouped.rowwise_data.view(torch.uint8).reshape(capacity, N // 2)[:sum_m] + g_scale = grouped.scale_inv.view(torch.uint8).reshape(capacity, scale_stride)[:sum_m] + g_amax = grouped.amax.reshape(-1)[:sum_m] + + torch.testing.assert_close(g_amax, ref_amax, atol=0, rtol=0) + torch.testing.assert_close(g_data, ref_data, atol=0, rtol=0) + torch.testing.assert_close(g_scale[:, : N // 16], ref_scale, atol=0, rtol=0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "split_sections, N", + [ + ([128], 128), + ([128, 256, 128], 256), + ([384, 128, 256], 128), + ([256, 128], 512), + ], +) +def test_matches_per_expert_quantize(split_sections: list[int], N: int) -> None: + torch.manual_seed(0) + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = _group_quantize(x, len(split_sections), first_dims) + + assert grouped.row_scaled_nvfp4 + _assert_matches_reference(grouped, x, split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("num_tensors, m, N", [(1, 128, 128), (4, 256, 256)]) +def test_uniform_shapes(num_tensors: int, m: int, N: int) -> None: + torch.manual_seed(1) + x = torch.randn((num_tensors * m, N), dtype=torch.bfloat16, device="cuda") + + grouped = _group_quantize(x, num_tensors, None) + + _assert_matches_reference(grouped, x, [m] * num_tensors) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "split_sections, N, capacity", + [ + ([128, 128], 256, 512), + ([256, 128, 128], 128, 768), + ], +) +def test_paged_stashing_capacity_noop(split_sections: list[int], N: int, capacity: int) -> None: + sum_m = sum(split_sections) + assert capacity > sum_m + torch.manual_seed(2) + x = torch.randn((capacity, N), dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = _group_quantize(x, len(split_sections), first_dims) + + _assert_matches_reference(grouped, x[:sum_m].clone(), split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_capture_once_replay_changed_routing() -> None: + N = 256 + num_experts = 3 + capacity = 768 + splits_a = [256, 128, 128] + splits_b = [128, 256, 128] + + first_dims = torch.empty(num_experts, dtype=torch.int64, device="cuda") + x = torch.empty((capacity, N), dtype=torch.bfloat16, device="cuda") + + torch.manual_seed(3) + x_a = torch.randn((sum(splits_a), N), dtype=torch.bfloat16, device="cuda") + x_b = torch.randn((sum(splits_b), N), dtype=torch.bfloat16, device="cuda") + + def program_routing(splits: list[int], x_src: torch.Tensor) -> None: + first_dims.copy_(torch.tensor(splits, dtype=torch.int64, device="cuda")) + x.zero_() + x[: x_src.shape[0], :].copy_(x_src) + + # Warm up outside capture so one-time setup does not land in the captured graph. + program_routing(splits_a, x_a) + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + _ = _group_quantize(x, num_experts, first_dims) + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + grouped = _group_quantize(x, num_experts, first_dims) + + program_routing(splits_a, x_a) + graph.replay() + torch.cuda.synchronize() + _assert_matches_reference(grouped, x_a, splits_a) + + program_routing(splits_b, x_b) + graph.replay() + torch.cuda.synchronize() + _assert_matches_reference(grouped, x_b, splits_b) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_non_contiguous_input_matches_contiguous() -> None: + split_sections = [128, 128] + N = 256 + torch.manual_seed(4) + x_t = torch.randn((N, sum(split_sections)), dtype=torch.bfloat16, device="cuda").t() + assert not x_t.is_contiguous() + first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = _group_quantize(x_t, len(split_sections), first_dims) + + _assert_matches_reference(grouped, x_t.contiguous(), split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rejects_unaligned_uniform_shapes() -> None: + x = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda") + + with pytest.raises(RuntimeError): + _group_quantize(x, 2, None) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rejects_strided_first_dims() -> None: + x = torch.randn((384, 128), dtype=torch.bfloat16, device="cuda") + strided = torch.tensor([256, 0, 128, 0], dtype=torch.int64, device="cuda")[::2] + assert not strided.is_contiguous() + + with pytest.raises(RuntimeError): + _group_quantize(x, 2, strided) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rejects_unsupported_quantizer_settings() -> None: + x = torch.randn((256, 128), dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + + with pytest.raises(RuntimeError): + tex.group_quantize(x, _row_scaled_quantizer(with_2d_quantization=True), 2, first_dims) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "input_shape, split_sections", + [ + ((384, 128), [256, 128]), + ((256, 64), [128, 128]), + ((256, 2, 128), [128, 128]), + ((512, 128), [128, 128]), + ], +) +def test_split_quantize_matches_per_expert(input_shape, split_sections: list[int]) -> None: + torch.manual_seed(5) + x = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda") + + outputs = tex.split_quantize( + x, split_sections, [_row_scaled_quantizer() for _ in split_sections] + ) + + assert len(outputs) == len(split_sections) + offset = 0 + for out, m in zip(outputs, split_sections): + ref = _row_scaled_quantizer()(x[offset : offset + m]) + torch.testing.assert_close( + out._rowwise_data.reshape(-1).view(torch.uint8), + ref._rowwise_data.reshape(-1).view(torch.uint8), + atol=0, + rtol=0, + ) + offset += m + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_split_quantize_mixed_row_scaled_falls_back() -> None: + split_sections = [128, 128] + N = 128 + torch.manual_seed(6) + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + + quantizers = [_tensor_scaled_quantizer(), _row_scaled_quantizer()] + outputs = tex.split_quantize(x, split_sections, quantizers) + + assert len(outputs) == len(split_sections) + offset = 0 + for out, m, make_q in zip( + outputs, split_sections, [_tensor_scaled_quantizer, _row_scaled_quantizer] + ): + ref = make_q()(x[offset : offset + m]) + torch.testing.assert_close( + out._rowwise_data.reshape(-1).view(torch.uint8), + ref._rowwise_data.reshape(-1).view(torch.uint8), + atol=0, + rtol=0, + ) + offset += m + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_split_quantize_empty_routing() -> None: + split_sections = [0, 0] + x = torch.empty((0, 128), dtype=torch.bfloat16, device="cuda") + + outputs = tex.split_quantize( + x, split_sections, [_row_scaled_quantizer() for _ in split_sections] + ) + + assert len(outputs) == len(split_sections) + for out in outputs: + assert tuple(out.shape) == (0, 128) diff --git a/transformer_engine/common/cast/cast_grouped.cu b/transformer_engine/common/cast/cast_grouped.cu index 49ab39f7e79..0a33f62d5fb 100644 --- a/transformer_engine/common/cast/cast_grouped.cu +++ b/transformer_engine/common/cast/cast_grouped.cu @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -101,16 +102,5 @@ void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *out void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, const size_t num_tensors, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_nvfp4_compute_amax); - using namespace transformer_engine; - - const Tensor *input_tensor = convertNVTETensorCheck(input); - std::vector output_tensors; - output_tensors.reserve(num_tensors); - for (size_t i = 0; i < num_tensors; ++i) { - output_tensors.push_back(convertNVTETensorCheck(outputs[i])); - } - - dispatch::nvfp4::group_compute_fused_amax(*input_tensor, /*noop=*/nullptr, output_tensors, - split_sections, num_tensors, stream); + nvte_group_amax(input, outputs, split_sections, num_tensors, stream); } diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index aa1d16fd894..13c3694d040 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -560,6 +560,29 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor quant_config_cpp.force_pow_2_scales, stream); break; } + case NVTE_NVFP4_1D_SCALING: { + NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_NVFP4_1D_SCALING."); + NVTE_CHECK(output_tensor->row_scaled_nvfp4, + "Grouped NVFP4 quantization is only implemented for row-scaled NVFP4."); + NVTE_CHECK(!quant_config_cpp.stochastic_rounding, + "Grouped row-scaled NVFP4 quantization does not support stochastic rounding."); + NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, + "Grouped row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(quant_config_cpp.nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "Grouped row-scaled NVFP4 quantization does not support 4over6."); + if (input_tensor->num_tensors == 0) { + break; + } + nvfp4::group_row_amax_graph_safe(input_tensor, noop_tensor, output_tensor, stream); + const DType scale_dtype = output_tensor->scale_inv.has_data() + ? output_tensor->scale_inv.dtype + : output_tensor->columnwise_scale_inv.dtype; + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + nvfp4::group_row_scaled_cast_graph_safe( + input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream);); + break; + } default: NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + "."); } diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index 2a481074d27..cfc5dd42c81 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -23,7 +23,9 @@ #include "../../util/math.h" #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" +#include "../core/grouped_layout.cuh" #include "core_nvfp4.cuh" +#include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh" namespace transformer_engine { namespace dispatch { @@ -99,16 +101,13 @@ static __global__ void group_fused_amax_zero_kernel(GroupedFusedAmaxArgs args, c } } +// Shared amax body; host-split and graph-safe kernels differ only in the prologue. template -__global__ void __launch_bounds__(FA_THREADS_NUM) - group_compute_fused_amax_kernel(const __grid_constant__ CUtensorMap tensor_map_input, - GroupedFusedAmaxArgs args, const float *noop, const size_t rows, - const size_t cols) { +__device__ __forceinline__ void group_compute_fused_amax_body(const CUtensorMap &tensor_map_input, + float *const row_amax_out, + float *const col_amax_out, + const int split_start) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - if (noop != nullptr && noop[0] == 1.0f) { - return; - } - using IType = __nv_bfloat16; using IType2 = ptx::FPx2; using IType3D = IType[FA_BUFFS_NUM][FA_TILE_DIM_Y][FA_TILE_DIM_X]; @@ -117,15 +116,6 @@ __global__ void __launch_bounds__(FA_THREADS_NUM) const int tid = threadIdx.x; const int block_offset_Y = blockIdx.y * FA_CHUNK_DIM_Y; const int block_offset_X = blockIdx.x * FA_CHUNK_DIM_X; - - // The 128-row chunk is inside a single expert. Select its per-expert outputs - // and remap the global row to a tensor-local row. - const int tensor_id = GetTensorId(args, block_offset_Y); - const int split_start = args.split_sections_range[tensor_id]; - float *const row_amax_out = - DO_ROW ? reinterpret_cast(args.row_amax_list[tensor_id]) : nullptr; - float *const col_amax_out = - DO_COL ? reinterpret_cast(args.col_amax_list[tensor_id]) : nullptr; const int local_row = block_offset_Y - split_start + tid; const int global_col = block_offset_X + tid; @@ -262,14 +252,91 @@ __global__ void __launch_bounds__(FA_THREADS_NUM) #endif // __CUDA_ARCH__ >= 1000 } +template +__global__ void __launch_bounds__(FA_THREADS_NUM) + group_compute_fused_amax_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + GroupedFusedAmaxArgs args, const float *noop, const size_t rows, + const size_t cols) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const int block_offset_Y = blockIdx.y * FA_CHUNK_DIM_Y; + // The 128-row chunk is inside a single expert; select its per-expert outputs. + const int tensor_id = GetTensorId(args, block_offset_Y); + const int split_start = args.split_sections_range[tensor_id]; + float *const row_amax_out = + DO_ROW ? reinterpret_cast(args.row_amax_list[tensor_id]) : nullptr; + float *const col_amax_out = + DO_COL ? reinterpret_cast(args.col_amax_list[tensor_id]) : nullptr; + + group_compute_fused_amax_body(tensor_map_input, row_amax_out, col_amax_out, + split_start); +#else + NVTE_DEVICE_ERROR("Grouped fused amax kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + +// Zero the capacity-sized grouped amax buffers before atomicMax. +static __global__ void group_fused_amax_zero_graph_safe_kernel(float *row_amax_all, + float *col_amax_all, + const size_t row_count, + const size_t col_count, + const float *noop) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + const size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t stride = static_cast(gridDim.x) * blockDim.x; + if (row_amax_all != nullptr) { + for (size_t i = idx; i < row_count; i += stride) row_amax_all[i] = 0.0f; + } + if (col_amax_all != nullptr) { + for (size_t i = idx; i < col_count; i += stride) col_amax_all[i] = 0.0f; + } +} + +// Graph-safe: per-expert outputs and row start derived on device from offsets; +// CTAs past the live row sum noop. +template +__global__ void __launch_bounds__(FA_THREADS_NUM) + group_compute_fused_amax_graph_safe_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + float *const row_amax_all, float *const col_amax_all, + const int64_t *const __restrict__ offsets_ptr, + const size_t num_tensors, const float *noop, + const size_t cols) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const size_t block_offset_Y = static_cast(blockIdx.y) * FA_CHUNK_DIM_Y; + const size_t global_offset = block_offset_Y * cols; + if (global_offset >= static_cast(offsets_ptr[num_tensors])) { + return; + } + + const size_t tensor_id = + dispatch::common::find_tensor_from_offsets(offsets_ptr, num_tensors, global_offset); + const int split_start = static_cast(static_cast(offsets_ptr[tensor_id]) / cols); + float *const row_amax_out = DO_ROW ? row_amax_all + split_start : nullptr; + float *const col_amax_out = DO_COL ? col_amax_all + tensor_id * cols : nullptr; + + group_compute_fused_amax_body(tensor_map_input, row_amax_out, col_amax_out, + split_start); +#else + NVTE_DEVICE_ERROR("Grouped fused amax kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + #endif // FP4_TYPE_SUPPORTED } // namespace group_row_scaled_amax_kernel -inline void group_compute_fused_amax(const Tensor &input, const Tensor *noop, - std::vector &output_list, - const size_t *split_sections, size_t num_tensors, - cudaStream_t stream) { +inline void group_row_amax(const Tensor &input, const Tensor *noop, + std::vector &output_list, const size_t *split_sections, + size_t num_tensors, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace group_row_scaled_amax_kernel; @@ -349,6 +416,81 @@ inline void group_compute_fused_amax(const Tensor &input, const Tensor *noop, #endif // FP4_TYPE_SUPPORTED } +// Graph-safe calculation of amaxes for row-scaled NVFP4: routing comes from device +// output->tensor_offsets, amax goes to the contiguous rowwise [sum_M] / columnwise +// [num_tensors * K] buffers. BF16-only, VARYING_FIRST_DIM. +inline void group_row_amax_graph_safe(const GroupedTensor *input, const Tensor *noop, + GroupedTensor *output, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace group_row_scaled_amax_kernel; + + NVTE_CHECK(input->dtype() == DType::kBFloat16, "Grouped fused amax requires BF16 input."); + NVTE_CHECK(output->all_same_last_dim() || output->all_same_shape(), + "Graph-safe grouped fused amax only supports a common last dimension " + "(VARYING_FIRST_DIM)."); + + const size_t num_tensors = input->num_tensors; + const size_t capacity_rows = input->logical_shape.data[0]; + const size_t cols = input->logical_shape.data[1]; + NVTE_CHECK(cols % FA_CHUNK_DIM_X == 0, "Grouped fused amax requires cols to be a multiple of ", + FA_CHUNK_DIM_X, " (got ", cols, ")."); + NVTE_CHECK(capacity_rows % FA_CHUNK_DIM_Y == 0, + "Graph-safe grouped fused amax requires the capacity row dimension to be a multiple " + "of ", + FA_CHUNK_DIM_Y, " (got ", capacity_rows, ")."); + + float *const row_amax_all = reinterpret_cast(output->amax.dptr); + float *const col_amax_all = reinterpret_cast(output->columnwise_amax.dptr); + const bool do_row = row_amax_all != nullptr; + const bool do_col = col_amax_all != nullptr; + if (!do_row && !do_col) return; + + const int64_t *offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + NVTE_CHECK(offsets_ptr != nullptr, + "Graph-safe grouped fused amax requires device tensor_offsets."); + + checkCuDriverContext(stream); + const float *noop_ptr = (noop != nullptr && noop->data.dptr != nullptr) + ? reinterpret_cast(noop->data.dptr) + : nullptr; + + // Zero the capacity-sized amax buffers before atomicMax. + { + const size_t row_count = do_row ? capacity_rows : 0; + const size_t col_count = do_col ? num_tensors * cols : 0; + const size_t total = std::max(row_count, col_count); + const unsigned zero_blocks = static_cast(DIVUP(total, static_cast(256))) + 1u; + group_fused_amax_zero_graph_safe_kernel<<>>( + row_amax_all, col_amax_all, row_count, col_count, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + + alignas(64) CUtensorMap tensor_map_input{}; + create_2D_tensor_map(tensor_map_input, input->data, capacity_rows, cols, FA_TILE_DIM_Y, + FA_TILE_DIM_X, cols, 0, sizeof(__nv_bfloat16) * 8); + + constexpr size_t buff_elems_total = FA_BUFFS_NUM * FA_BUFF_IN_SIZE; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(__nv_bfloat16), TMA_SHMEM_ALIGNMENT); + constexpr size_t dshmem_size = buff_size_aligned_in + TMA_SHMEM_ALIGNMENT; + + const dim3 grid(static_cast(cols / FA_CHUNK_DIM_X), + static_cast(capacity_rows / FA_CHUNK_DIM_Y), 1); + const dim3 block(FA_THREADS_NUM, 1, 1); + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + do_row, DO_ROW, TRANSFORMER_ENGINE_SWITCH_CONDITION(do_col, DO_COL, { + auto kernel = group_compute_fused_amax_graph_safe_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); + kernel<<>>(tensor_map_input, row_amax_all, col_amax_all, + offsets_ptr, num_tensors, noop_ptr, cols); + })); + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + namespace group_quantize_transpose_kernel { using namespace core; @@ -1237,6 +1379,392 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, #endif // FP4_TYPE_SUPPORTED } +// Single-launch grouped (MoE) row-scaled NVFP4 cast. Reuses the single-tensor +// tuned_1D encode helpers verbatim (so the output is byte-identical to the +// per-expert loop) and replaces tuned_1D's TMA output stores with a per-expert +// global scatter, since each expert's transpose output has its own M_i stride. +namespace group_row_scaled_cast_kernel { + +using namespace quantize_transpose_tuned_kernel; +using namespace ptx; + +#if FP4_TYPE_SUPPORTED + +template +__device__ __forceinline__ void group_row_scaled_cast_body( + const CUtensorMap &tensor_map_input, uint8_t *const q_row_base, ScaleType *const s_dec_row_base, + const float *const row_amax_base, uint8_t *const q_col_base, ScaleType *const s_dec_col_base, + const float *const col_amax_base, const size_t scale_stride_col, const int M_t, + const int local_row_base, const int ctaid_X, const int ctaid_Y, const size_t cols, + const size_t scale_stride_row) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr bool USE_STOCHASTIC_ROUNDING = false; // row-scaled rejects SR + constexpr bool ROW_SCALED_NVFP4 = true; + + // SR off: the RNG stays inert (random_uint4 is never consumed); it exists only + // to satisfy the shared encode-helper signatures. + RNG_t rng; + rng.init(0, 0, 0); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + int rnd_idx = 0; + + const bool leading_thread = (threadIdx.x == 0); + + const int block_offset_Y = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + const int block_offset_X = ctaid_X * TunableConfig::CHUNK_DIM_X; + + const size_t data_stride_row = cols / 2; + const size_t data_stride_col = static_cast(M_t) / 2; + + // SMEM layout: identical to the single-tensor tuned_1D kernel. + constexpr int buff_elems = BUFF_DIM_Y * BUFF_IN_DIM_X; + constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; + constexpr int buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out_t = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int in_mem = buff_size_aligned_in; + constexpr int out_mem_rowwise_data = buff_size_aligned_out; + constexpr int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; + constexpr int out_mem_rowwise_scales = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); + + IType *sIn_ptr = reinterpret_cast(dshmem); + fp4e2m1x2 *sOut_ptr = reinterpret_cast(dshmem + in_mem); + fp4e2m1x2 *sOut_tr_ptr = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + ScaleType *sSFrowwise_ptr = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + ScaleType *sSFcolwise_ptr = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + + auto &sIn = *reinterpret_cast(sIn_ptr); + + constexpr int shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + __shared__ uint64_t IN_buff_readable_mbar[BUFFS_NUM]; + + if (leading_thread) { +#pragma unroll + for (int buff = 0; buff < BUFFS_NUM; ++buff) { + ptx::mbarrier_init(&IN_buff_readable_mbar[buff], 1); + } + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + // Prefetch the first input tile(s). +#pragma unroll + for (int stage = 0; stage < TunableConfig::PREFETCH_STAGES; ++stage) { + const int buff_in_p = stage; + const int stage_Y = stage / STAGES_X; + const int stage_X = stage % STAGES_X; + const int global_offset_Y = block_offset_Y + stage_Y * TILE_DIM_Y; + const int global_offset_X = block_offset_X + stage_X * TILE_DIM_X; + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(&IN_buff_readable_mbar[buff_in_p], shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sIn[buff_in_p]), + reinterpret_cast(&tensor_map_input), global_offset_X, global_offset_Y, + &IN_buff_readable_mbar[buff_in_p]); + } + } + + int buff_in = 0; + int buff_out = 0; + int buff_out_tr = 0; + int IN_buff_readable_parity[BUFFS_NUM] = {0, 0}; + +#pragma unroll + for (int stage = 0; stage < STAGES; ++stage) { + const int stage_Y = stage / STAGES_X; + const int stage_X = stage % STAGES_X; + + // Prefetch the next input tile. + if (stage < STAGES - TunableConfig::PREFETCH_STAGES) { + const int next_prefetch_buff = (buff_in + TunableConfig::PREFETCH_STAGES) % BUFFS_NUM; + const int next_prefetch_stage = stage + TunableConfig::PREFETCH_STAGES; + const int next_stage_Y = next_prefetch_stage / STAGES_X; + const int next_stage_X = next_prefetch_stage % STAGES_X; + const int next_global_offset_Y = block_offset_Y + next_stage_Y * TILE_DIM_Y; + const int next_global_offset_X = block_offset_X + next_stage_X * TILE_DIM_X; + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(&IN_buff_readable_mbar[next_prefetch_buff], shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sIn[next_prefetch_buff]), + reinterpret_cast(&tensor_map_input), next_global_offset_X, + next_global_offset_Y, &IN_buff_readable_mbar[next_prefetch_buff]); + } + ptx::fence_proxy_async_shared_cta(); + } + + // Wait for the current input tile to land. + ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&IN_buff_readable_mbar[buff_in], + IN_buff_readable_parity[buff_in]); + IN_buff_readable_parity[buff_in] ^= 1; + + // Per-row encode (byte-identical to the single-tensor path): pass the + // per-expert amax base and a local row offset so amax[row_idx] indexes the + // expert-local row. + rowwise_scaling( + sIn_ptr, sOut_ptr, sSFrowwise_ptr, /*S_enc_rowwise=*/1.0f, stage_Y, stage_X, buff_in, + buff_out, row_amax_base, static_cast(local_row_base), static_cast(M_t), rng, + random_uint4, rnd_idx); + + if constexpr (RETURN_TRANSPOSE) { + colwise_scaling( + sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, /*S_enc_colwise=*/1.0f, stage_Y, stage_X, buff_in, + buff_out_tr, col_amax_base, static_cast(block_offset_X), cols, rng, random_uint4, + rnd_idx); + } + + // Make the encode-helper SMEM writes visible before the scatter. + __syncthreads(); + + // Scatter the rowwise FP4 tile to the per-expert buffer: a verbatim copy of + // the 64-row x 32-byte tile (2 threads x 16 B per row). + { + auto &sOut = *reinterpret_cast(sOut_ptr); + const int row_in_tile = static_cast(threadIdx.x) >> 1; + const int half = static_cast(threadIdx.x) & 1; + const int local_row = local_row_base + stage_Y * TILE_DIM_Y + row_in_tile; + const size_t byte_off_X = + static_cast(block_offset_X / 2) + stage_X * (TILE_DIM_X / 2) + half * 16; + const uint4 *src = reinterpret_cast(&sOut[buff_out][row_in_tile][half * 16]); + uint4 *dst = reinterpret_cast( + q_row_base + static_cast(local_row) * data_stride_row + byte_off_X); + *dst = *src; + } + if constexpr (RETURN_TRANSPOSE) { + auto &sOut_tr = *reinterpret_cast(sOut_tr_ptr); + const int col_in_tile = static_cast(threadIdx.x) >> 1; + const int half = static_cast(threadIdx.x) & 1; + const int global_col = block_offset_X + stage_X * TILE_DIM_X + col_in_tile; + const size_t byte_off_M = + static_cast(local_row_base / 2) + stage_Y * (TILE_DIM_Y / 2) + half * 16; + const uint4 *src = + reinterpret_cast(&sOut_tr[buff_out_tr][col_in_tile][half * 16]); + uint4 *dst = reinterpret_cast( + q_col_base + static_cast(global_col) * data_stride_col + byte_off_M); + *dst = *src; + } + + // Sync so the scatter completes before the next stage overwrites the buffer. + __syncthreads(); + + buff_in = (buff_in + 1) % BUFFS_NUM_IN; + buff_out = (buff_out + 1) % BUFFS_NUM_OUT; + buff_out_tr = (buff_out_tr + 1) % BUFFS_NUM_OUT_TR; + } + + // SF epilogue: store the per-chunk SMEM scale tiles to the per-expert buffers. + { + auto &sSFrowwise = *reinterpret_cast *>(sSFrowwise_ptr); + using ScalesVec = Vec; + const size_t scales_block_offset_X_rowwise = static_cast(ctaid_X) * SCALES_PER_CHUNK_X; + for (int row = static_cast(threadIdx.x); row < TunableConfig::CHUNK_DIM_Y; + row += THREADS_NUM) { + ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); + const size_t local_row = static_cast(local_row_base) + row; + const size_t scale_idx = local_row * scale_stride_row + scales_block_offset_X_rowwise; + scales_vec.store_to_elts(&s_dec_row_base[scale_idx], 0, SCALES_PER_CHUNK_X); + } + } + if constexpr (RETURN_TRANSPOSE) { + auto &sSFcolwise = *reinterpret_cast *>(sSFcolwise_ptr); + using ScalesVec = Vec; + const size_t local_block_offset_M = static_cast(local_row_base) / SCALE_DIM; + for (int row_tr = static_cast(threadIdx.x); row_tr < TunableConfig::CHUNK_DIM_X; + row_tr += THREADS_NUM) { + ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); + const size_t global_col = static_cast(block_offset_X) + row_tr; + const size_t scale_idx = global_col * scale_stride_col + local_block_offset_M; + scales_vec.store_to_elts(&s_dec_col_base[scale_idx], 0, SCALES_PER_CHUNK_Y); + } + } + + if (leading_thread) { +#pragma unroll + for (int buff = 0; buff < BUFFS_NUM; ++buff) { + ptx::mbarrier_invalid(&IN_buff_readable_mbar[buff]); + } + } +#else + NVTE_DEVICE_ERROR("Grouped row-scaled NVFP4 cast kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + +// Grid covers the input capacity; CTAs past the live row sum noop. +template +__global__ void __launch_bounds__(THREADS_NUM) group_row_scaled_cast_nvfp4_graph_safe_kernel( + const __grid_constant__ CUtensorMap tensor_map_input, uint8_t *const q_row_all, + ScaleType *const s_dec_row_all, const float *const row_amax_all, uint8_t *const q_col_all, + ScaleType *const s_dec_col_all, const float *const col_amax_all, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, + const size_t num_tensors, const float *noop, const size_t cols, const size_t scale_stride_row) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const int ctaid_X = blockIdx.x; + const int ctaid_Y = blockIdx.y; + const size_t block_offset_Y = static_cast(ctaid_Y) * TunableConfig::CHUNK_DIM_Y; + + const size_t global_offset = block_offset_Y * cols; + if (global_offset >= static_cast(offsets_ptr[num_tensors])) { + return; + } + + const size_t tensor_id = + dispatch::common::find_tensor_from_offsets(offsets_ptr, num_tensors, global_offset); + const size_t row_start = static_cast(offsets_ptr[tensor_id]) / cols; + const int M_t = static_cast(first_dims_ptr[tensor_id]); + const int local_row_base = static_cast(block_offset_Y - row_start); + + // Rowwise buffers are contiguous by global row; transpose buffers are packed + // per expert with column stride M_t. + const size_t scale_stride_col = RETURN_TRANSPOSE ? static_cast(M_t) / SCALE_DIM : 0; + + uint8_t *const q_row_base = q_row_all + row_start * (cols / 2); + ScaleType *const s_dec_row_base = s_dec_row_all + row_start * scale_stride_row; + const float *const row_amax_base = row_amax_all + row_start; + + uint8_t *const q_col_base = RETURN_TRANSPOSE ? q_col_all + (cols / 2) * row_start : nullptr; + ScaleType *const s_dec_col_base = + RETURN_TRANSPOSE ? s_dec_col_all + (cols / SCALE_DIM) * row_start : nullptr; + const float *const col_amax_base = RETURN_TRANSPOSE ? col_amax_all + tensor_id * cols : nullptr; + + group_row_scaled_cast_body( + tensor_map_input, q_row_base, s_dec_row_base, row_amax_base, q_col_base, s_dec_col_base, + col_amax_base, scale_stride_col, M_t, local_row_base, ctaid_X, ctaid_Y, cols, + scale_stride_row); +#else + NVTE_DEVICE_ERROR("Grouped row-scaled NVFP4 cast kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + +#endif // FP4_TYPE_SUPPORTED + +} // namespace group_row_scaled_cast_kernel + +template +inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const Tensor *noop, + GroupedTensor *output, + const QuantizationConfig *quant_config, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace group_row_scaled_cast_kernel; + using namespace quantize_transpose_tuned_kernel; + + NVTE_CHECK(input->dtype() == DType::kBFloat16, + "Grouped row-scaled NVFP4 cast requires BF16 input."); + NVTE_CHECK(output->all_same_last_dim() || output->all_same_shape(), + "Graph-safe grouped row-scaled NVFP4 cast only supports a common last dimension " + "(VARYING_FIRST_DIM)."); + + const bool use_fast_math = quant_config ? quant_config->use_fast_math : false; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + + const size_t num_tensors = input->num_tensors; + const size_t capacity_rows = input->logical_shape.data[0]; + const size_t cols = input->logical_shape.data[1]; + NVTE_CHECK(cols % TunableConfig::CHUNK_DIM_X == 0, + "Grouped row-scaled NVFP4 cast requires cols to be a multiple of ", + TunableConfig::CHUNK_DIM_X, " (got ", cols, ")."); + NVTE_CHECK(capacity_rows % TunableConfig::CHUNK_DIM_Y == 0, + "Graph-safe grouped row-scaled NVFP4 cast requires the capacity row dimension to be a " + "multiple of ", + TunableConfig::CHUNK_DIM_Y, " (got ", capacity_rows, ")."); + + const bool return_transpose = output->has_columnwise_data(); + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Graph-safe cast requires rowwise scales."); + NVTE_CHECK(output->amax.dptr != nullptr, "Graph-safe cast requires rowwise amax."); + // Per-row scale stride, from a 2D [rows, stride] scale or a flat 1D buffer. + size_t scale_stride_row = 0; + { + const auto &scale_shape = output->scale_inv.shape; + const size_t scale_ndim = scale_shape.size(); + if (scale_ndim >= 2) { + scale_stride_row = scale_shape[scale_ndim - 1]; + } else { + size_t scale_numel = 1; + for (size_t i = 0; i < scale_ndim; ++i) scale_numel *= scale_shape[i]; + scale_stride_row = scale_numel / capacity_rows; + } + } + + const int64_t *offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t *first_dims_ptr = reinterpret_cast(output->first_dims.dptr); + NVTE_CHECK(offsets_ptr != nullptr && first_dims_ptr != nullptr, + "Graph-safe cast requires device tensor_offsets and first_dims."); + + uint8_t *q_row_all = reinterpret_cast(output->data.dptr); + ScaleType *s_dec_row_all = reinterpret_cast(output->scale_inv.dptr); + const float *row_amax_all = reinterpret_cast(output->amax.dptr); + uint8_t *q_col_all = + return_transpose ? reinterpret_cast(output->columnwise_data.dptr) : nullptr; + ScaleType *s_dec_col_all = + return_transpose ? reinterpret_cast(output->columnwise_scale_inv.dptr) : nullptr; + const float *col_amax_all = + return_transpose ? reinterpret_cast(output->columnwise_amax.dptr) : nullptr; + if (return_transpose) { + NVTE_CHECK(s_dec_col_all != nullptr && col_amax_all != nullptr, + "Graph-safe transpose cast requires columnwise scale and amax buffers."); + } + + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + + alignas(64) CUtensorMap tensor_map_input{}; + create_2D_tensor_map(tensor_map_input, input->data, capacity_rows, cols, BUFF_DIM_Y, BUFF_DIM_X, + cols, 0, sizeof(IType) * 8); + + constexpr int buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; + constexpr int buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out_t = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_scales = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_scales_transpose = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_X * SCALES_PER_CHUNK_Y * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + + const dim3 grid(static_cast(cols / TunableConfig::CHUNK_DIM_X), + static_cast(capacity_rows / TunableConfig::CHUNK_DIM_Y), 1); + const dim3 block(THREADS_NUM, 1, 1); + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_fast_math, USE_FAST_MATH, + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + const int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; + const int out_scales_transpose_mem = RETURN_TRANSPOSE ? buff_size_scales_transpose : 0; + const int dshmem_size = buff_size_aligned_in + buff_size_aligned_out + + out_mem_colwise_data + buff_size_scales + out_scales_transpose_mem + + TMA_SHMEM_ALIGNMENT; + auto kernel = group_row_scaled_cast_nvfp4_graph_safe_kernel; + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>(tensor_map_input, q_row_all, s_dec_row_all, + row_amax_all, q_col_all, s_dec_col_all, + col_amax_all, offsets_ptr, first_dims_ptr, + num_tensors, noop_ptr, cols, scale_stride_row); + })); + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + } // namespace nvfp4 } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 4b64d91b932..f95b6740b09 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -545,6 +545,9 @@ struct GroupedTensor { */ bool with_gemm_swizzled_scales = false; + /*! \brief Whether NVFP4 data uses row scaling instead of tensor scaling */ + bool row_scaled_nvfp4 = false; + /*! Map from NVTEGroupedTensorParam to parameter sizes */ static constexpr size_t attr_sizes[] = { sizeof(NVTEBasicTensor), // kNVTEGroupedRowwiseData @@ -557,7 +560,8 @@ struct GroupedTensor { sizeof(NVTEBasicTensor), // kNVTEGroupedFirstDims sizeof(NVTEBasicTensor), // kNVTEGroupedLastDims sizeof(NVTEBasicTensor), // kNVTEGroupedTensorOffsets - sizeof(uint8_t) // kNVTEGroupedWithGEMMSwizzledScales + sizeof(uint8_t), // kNVTEGroupedWithGEMMSwizzledScales + sizeof(uint8_t) // kNVTEGroupedRowScaledNVFP4 }; GroupedTensor(NVTEScalingMode scaling_mode, size_t num_tensors) @@ -575,7 +579,8 @@ struct GroupedTensor { tensor_offsets(nullptr, std::vector{0}, DType::kInt64), logical_shape(nvte_make_shape(nullptr, 1)), nvte_tensor(0), - with_gemm_swizzled_scales(false) {} + with_gemm_swizzled_scales(false), + row_scaled_nvfp4(false) {} explicit operator NVTEGroupedTensor() const noexcept { return nvte_tensor; } @@ -628,6 +633,7 @@ struct GroupedTensor { scaling_mode = NVTE_DELAYED_TENSOR_SCALING; nvte_tensor = 0; with_gemm_swizzled_scales = false; + row_scaled_nvfp4 = false; } }; diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu index 2316d9697ae..0528ac0927a 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -14,6 +14,7 @@ #include +#include "common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh" #include "common/common.h" #include "common/util/ptx.cuh" #include "common/utils.cuh" @@ -563,7 +564,7 @@ void nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input static_cast(random_sign_mask_t), false, stream); } -// Grouped-tensor amax without doing hadamard transform +// Grouped-tensor amax void nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_amax_graph_safe); @@ -576,5 +577,13 @@ void nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor return; } + // Per-row amax + if (output_tensor->scaling_mode == NVTE_NVFP4_1D_SCALING && output_tensor->row_scaled_nvfp4) { + dispatch::nvfp4::group_row_amax_graph_safe(input_tensor, /*noop=*/nullptr, output_tensor, + stream); + return; + } + + // Per-tensor amax group_hadamard_transform_amax_graph_safe(input_tensor, output_tensor, 0, 0, true, stream); } diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu index 24d06e5d25e..f8c04bea0cb 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -12,8 +12,10 @@ #include #include +#include #include +#include "common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh" #include "common/common.h" #include "common/util/ptx.cuh" #include "common/utils.cuh" @@ -582,7 +584,7 @@ void nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outp static_cast(random_sign_mask_t), false, stream); } -// Grouped-tensor amax without doing hadamard transform +// Grouped-tensor amax void nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* split_sections, size_t num_tensors, cudaStream_t stream) { NVTE_API_CALL(nvte_group_amax); @@ -597,6 +599,30 @@ void nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* output_list[i] = convertNVTETensorCheck(outputs[i]); } + // Check that tensors have uniform scaling modes + const auto scaling_mode = output_list[0]->scaling_mode; + for (const auto* t : output_list) { + NVTE_CHECK(t->scaling_mode == scaling_mode, "Tensors have inconsistent scaling modes (found ", + to_string(scaling_mode), " and ", to_string(t->scaling_mode), ")."); + } + + // Compute per-row amax if outputs are row-scaled NVFP4 tensors + if (scaling_mode == NVTE_NVFP4_1D_SCALING) { + auto is_row_scaled_nvfp4 = [](const Tensor* t) -> bool { return t->row_scaled_nvfp4; }; + const bool all_row_scaled_nvfp4 = + std::all_of(output_list.begin(), output_list.end(), is_row_scaled_nvfp4); + const bool any_row_scaled_nvfp4 = + std::any_of(output_list.begin(), output_list.end(), is_row_scaled_nvfp4); + NVTE_CHECK(all_row_scaled_nvfp4 == any_row_scaled_nvfp4, + "Output tensors are not consistent between row-scaling and tensor-scaling."); + if (all_row_scaled_nvfp4) { + dispatch::nvfp4::group_row_amax(*input_tensor, /*noop=*/nullptr, output_list, split_sections, + num_tensors, stream); + return; + } + } + + // Per-tensor amax group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors, 0, 0, true, stream); } diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index ca3eb13ed24..0c0c0d89d4d 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -93,6 +93,7 @@ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t strea * The type of quantized tensor in the output depends on the scaling mode of the output * tensor. See file level comments. * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. + * NVFP4 output requires row scaling, BF16 input, and dimensions that are multiples of 128. * * \param[in] input Input grouped tensor to be cast. * \param[in,out] output Output grouped tensor. @@ -494,22 +495,7 @@ void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *out const NVTEQuantizationConfig quant_config, cudaStream_t stream); -/*! \brief Computes per-expert rowwise and/or columnwise amax for a grouped - * row-scaled NVFP4 input in a single kernel launch. - * - * The input is one packed (sum_M, K) BF16 buffer; `split_sections` gives the - * per-expert row counts (each a multiple of 128) and K must also be a multiple - * of 128. For expert i the rowwise amax [M_i] is written to `outputs[i]->amax` - * and the columnwise amax [K] to `outputs[i]->columnwise_amax`, whichever are - * allocated; experts may opt into either direction independently (a null buffer - * is skipped). - * - * \param[in] input Packed grouped input tensor (BF16). - * \param[in,out] outputs Per-expert output tensors receiving the amax vectors. - * \param[in] split_sections Per-expert row counts. - * \param[in] num_tensors Number of output tensors. - * \param[in] stream CUDA stream used for the operation. - */ +/*! \deprecated Use nvte_group_amax instead. */ void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, cudaStream_t stream); diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index 09ab260f15b..30f7669392e 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -302,10 +302,12 @@ void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor ** * This function is experimental and the API is not stable. * * For each tensor split, compute the maximum absolute value (amax) - * and populate the amax of the corresponding output tensor. + * and populate the amax buffers of the corresponding output tensor. + * All output tensors must have the same tensor format, and the + * tensor format determines which amaxes are computed. * * \param[in] input Input tensor. - * \param[in,out] outputs Array of output tensors. Only the amax is updated. + * \param[in,out] outputs Array of output tensors. Only amax buffers are updated. * \param[in] split_sections Size of each tensor split along dimension 0. * \param[in] num_tensors Number of tensor splits. * \param[in] stream CUDA stream used for the operation. @@ -313,10 +315,14 @@ void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor ** void nvte_group_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, cudaStream_t stream); -/*! \brief Grouped-tensor amax without doing hadamard transform. +/*! \brief Graph-safe variant of nvte_group_amax for grouped tensors. * * This function is experimental and the API is not stable. * + * Compute the amax for each tensor in the group and populate the amax + * buffers of the output tensor. The output tensor format determines + * which amaxes are computed. + * * \param[in] input NVTEGroupedTensor Input tensor. * \param[in,out] output NVTEGroupedTensor Output tensor. * \param[in] stream CUDA stream used for the operation. diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 4c0e844c8c4..1adb47773ef 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -570,6 +570,8 @@ enum NVTEGroupedTensorParam { 9, /*!< Tensor offsets for contiguous layout (device pointer to int64_t array) */ kNVTEGroupedWithGEMMSwizzledScales = 10, /*!< Whether scaling factors are in format expected by GEMM */ + kNVTEGroupedRowScaledNVFP4 = + 11, /*!< Whether NVFP4 data uses row scaling instead of tensor scaling */ kNVTENumGroupedTensorParams }; @@ -1302,6 +1304,11 @@ class GroupedTensorWrapper { nvte_set_grouped_tensor_param(tensor_, kNVTEGroupedWithGEMMSwizzledScales, &val, sizeof(val)); } + void set_row_scaled_nvfp4(bool row_scaled_nvfp4) { + const auto val = static_cast(row_scaled_nvfp4); + nvte_set_grouped_tensor_param(tensor_, kNVTEGroupedRowScaledNVFP4, &val, sizeof(val)); + } + // Parameter getters NVTEBasicTensor get_parameter(const NVTEGroupedTensorParam param) const noexcept { NVTEBasicTensor ret; @@ -1348,6 +1355,12 @@ class GroupedTensorWrapper { return static_cast(val); } + bool get_row_scaled_nvfp4() const { + uint8_t val = 0; + nvte_get_grouped_tensor_param(tensor_, kNVTEGroupedRowScaledNVFP4, &val, sizeof(val), nullptr); + return static_cast(val); + } + /*! \brief Get an underlying NVTEGroupedTensor. * * \return NVTEGroupedTensor held by this GroupedTensorWrapper. diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index ad0d47ac99e..66693dca3a1 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -1330,6 +1330,9 @@ void nvte_set_grouped_tensor_param(NVTEGroupedTensor tensor, NVTEGroupedTensorPa case kNVTEGroupedWithGEMMSwizzledScales: t.with_gemm_swizzled_scales = static_cast(*reinterpret_cast(buf)); break; + case kNVTEGroupedRowScaledNVFP4: + t.row_scaled_nvfp4 = static_cast(*reinterpret_cast(buf)); + break; default: NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast(param), ")"); } @@ -1425,6 +1428,9 @@ void nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, NVTEGroupedTe case kNVTEGroupedWithGEMMSwizzledScales: *reinterpret_cast(buf) = static_cast(t->with_gemm_swizzled_scales); break; + case kNVTEGroupedRowScaledNVFP4: + *reinterpret_cast(buf) = static_cast(t->row_scaled_nvfp4); + break; default: NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast(param), ")"); } diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 06d03cd6e1c..3d61dc525f8 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -211,6 +211,35 @@ void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, }); } +void group_quantize_nvfp4_row_scaled_impl(const GroupedTensorWrapper &grouped_input_tensor, + GroupedTensorWrapper &grouped_output_tensor, + const NVFP4Quantizer *nvfp4_quantizer_cpp, + const std::optional &noop_flag, + cudaStream_t stream) { + NVTE_CHECK(!nvfp4_quantizer_cpp->with_rht, + "Grouped row-scaled NVFP4 quantization does not support RHT."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization, + "Grouped row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(!nvfp4_quantizer_cpp->stochastic_rounding, + "Grouped row-scaled NVFP4 quantization does not support stochastic rounding."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_amax_reduction, + "Grouped row-scaled NVFP4 quantization does not support amax reduction."); + NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "Grouped row-scaled NVFP4 quantization does not support 4over6."); + + QuantizationConfigWrapper quant_config_cpp; + if (noop_flag.has_value()) { + quant_config_cpp.set_noop_tensor(noop_flag->data()); + } + if (transformer_engine::getenv("NVTE_USE_FAST_MATH")) { + quant_config_cpp.set_use_fast_math(true); + } + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor.data(), quant_config_cpp, + stream); + }); +} + float fp8_max_for_dtype(const DType dtype) { if (!is_fp8_dtype(dtype)) { NVTE_ERROR("Expected FP8 dtype for grouped current-scaling quantization, got ", @@ -276,14 +305,15 @@ void compute_grouped_fp8_current_scaling_amax_and_scale( } // namespace -py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, +py::object group_quantize(const at::Tensor &input, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, std::optional noop_flag, const py::object &output) { using namespace transformer_engine::pytorch::detail; init_extension(); - NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(input.dim() == 2, "Tensor must be 2D"); + const at::Tensor tensor = input.contiguous(); // No-op flag for CUDA graph weight caching (skip_fp8_weight_update). std::optional noop_flag_cpp; @@ -302,6 +332,26 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const auto quantizer_cpp = convert_quantizer(quantizer); + const bool row_scaled_nvfp4 = + detail::IsNVFP4Quantizers(quantizer.ptr()) && + static_cast(quantizer_cpp.get())->row_scaled_nvfp4; + if (row_scaled_nvfp4) { + NVTE_CHECK(output.is_none(), + "group_quantize: output reuse is not supported with row-scaled NVFP4."); + if (first_dims.has_value()) { + NVTE_CHECK(first_dims->is_contiguous() && first_dims->scalar_type() == at::kLong, + "group_quantize: first_dims must be a contiguous int64 tensor."); + } else { + NVTE_CHECK(num_tensors > 0 && logical_first_dim % (num_tensors * 128) == 0, + "group_quantize: row-scaled NVFP4 with uniform shapes requires the first " + "dimension (", + logical_first_dim, ") to be a multiple of 128 * num_tensors (", num_tensors, ")."); + first_dims = at::full({static_cast(num_tensors)}, + static_cast(logical_first_dim / num_tensors), + tensor.options().dtype(at::kLong)); + } + } + // Create input GroupedTensor. auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), @@ -370,9 +420,15 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const NVTE_CHECK(!last_dims.has_value(), "group_quantize: varying last dim is not supported with NVFP4."); NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); - group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, - nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), - !nvfp4_quantizer_cpp->disable_second_level_scale); + if (row_scaled_nvfp4) { + group_quantize_nvfp4_row_scaled_impl(grouped_input_tensor, grouped_output_tensor_cpp, + nvfp4_quantizer_cpp, noop_flag_cpp, + at::cuda::getCurrentCUDAStream()); + } else { + group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, + nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), + !nvfp4_quantizer_cpp->disable_second_level_scale); + } break; } case GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE: { @@ -2115,15 +2171,16 @@ std::vector split_quantize(const at::Tensor &tensor, [](const py::handle &quantizer) -> bool { return detail::IsNVFP4Quantizers(quantizer.ptr()); })) { - allocation_method = AllocationMethod::BULK_NVFP4; const bool has_row_scaled_nvfp4 = std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), [](const std::unique_ptr &quantizer) { return static_cast(quantizer.get())->row_scaled_nvfp4; }); if (has_row_scaled_nvfp4) { + allocation_method = AllocationMethod::UNFUSED; quantization_method = QuantizationMethod::UNFUSED; } else { + allocation_method = AllocationMethod::BULK_NVFP4; quantization_method = QuantizationMethod::FUSED_NVFP4; } } diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index ba0410bcacb..960e0fe46f4 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2264,6 +2264,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso } out_cpp.set_with_gemm_swizzled_scales(this->optimize_for_gemm); + out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 7f377583d18..f9e4f487549 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -319,6 +319,10 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { } ret.set_with_gemm_swizzled_scales(with_gemm_swizzled); + if (py::hasattr(tensor, "_row_scaled_nvfp4")) { + ret.set_row_scaled_nvfp4(tensor.attr("_row_scaled_nvfp4").cast()); + } + return ret; } From 83d8c4e3370bc0117189adbbd88a50fd57440acd Mon Sep 17 00:00:00 2001 From: Deep Patel <68748791+dinodeep@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:10:34 -0700 Subject: [PATCH 07/32] TE MoE Return Expert Counts and Custom Routing for DSv4 Support (#3596) Support user-provided routing and return expert counts Add RouterComputationInfo and RoutingMapInfo to the JAX MoE API. moe() takes one router_info argument: RouterComputationInfo has TE compute the gate and top-k, while RoutingMapInfo supplies caller-computed expert indices and differentiable routing weights (e.g. DeepSeek-V4 hash routing) and skips TE's router. moe() and _MoEBlock also return global per-expert assignment counts as a fourth output. collect_expert_counts=False skips the associated all-reduce. Signed-off-by: Deep Patel --- tests/jax/moe/test_moe_fsdp_layout.py | 3 +- tests/jax/moe/test_moe_resource_api.py | 15 +- tests/jax/test_custom_call_compute.py | 15 +- tests/jax/test_te_ep_moe.py | 322 +++++++++++- transformer_engine/jax/flax/moe.py | 35 +- transformer_engine/jax/moe.py | 651 +++++++++++++++---------- 6 files changed, 745 insertions(+), 296 deletions(-) diff --git a/tests/jax/moe/test_moe_fsdp_layout.py b/tests/jax/moe/test_moe_fsdp_layout.py index 2296dfe41ac..bee1cc98f4c 100644 --- a/tests/jax/moe/test_moe_fsdp_layout.py +++ b/tests/jax/moe/test_moe_fsdp_layout.py @@ -365,4 +365,5 @@ def test_layout_argument_is_optional_and_static(): assert inspect.signature(moe.moe).parameters["cudnn_native_weight_layout"].default is None assert inspect.signature(moe._moe).parameters["cudnn_native_weight_layout"].default is None - assert 33 in moe._moe.nondiff_argnums + parameters = list(inspect.signature(moe._moe).parameters) + assert parameters.index("cudnn_native_weight_layout") in moe._moe.nondiff_argnums diff --git a/tests/jax/moe/test_moe_resource_api.py b/tests/jax/moe/test_moe_resource_api.py index f5b058e314c..ab6c02932d6 100644 --- a/tests/jax/moe/test_moe_resource_api.py +++ b/tests/jax/moe/test_moe_resource_api.py @@ -14,7 +14,11 @@ from jax.sharding import Mesh from transformer_engine.jax.flax import _MoEBlock -from transformer_engine.jax.moe import _moe_mesh_axes, _resolve_moe_mesh_resource +from transformer_engine.jax.moe import ( + RouterComputationInfo, + _moe_mesh_axes, + _resolve_moe_mesh_resource, +) from transformer_engine.jax.sharding import MeshResource, global_mesh_resource, global_shard_guard @@ -27,7 +31,7 @@ def no_global_resource(): def test_resource_required_for_function_and_block(): module = importlib.import_module("transformer_engine.jax.moe") with pytest.raises(ValueError, match="active global_shard_guard"): - module.moe(None, None, None, None, num_experts=2, num_experts_per_tok=1) + module.moe(None, None, None, None, num_experts=2) with pytest.raises(ValueError, match="active global_shard_guard"): _MoEBlock().init(jax.random.PRNGKey(0), jnp.ones((1, 1, 4))) @@ -127,20 +131,19 @@ def test_public_api_uses_selected_resource(monkeypatch, quantize): def fake_vjp(*args): captured.update(signature.bind(*args).arguments) assert global_mesh_resource().ep_resource == "ep" - return args[0], None, jnp.zeros((1,), jnp.int32) + return args[0], None, jnp.zeros((1,), jnp.int32), jnp.zeros((2,), jnp.int32) monkeypatch.setattr(module, "_moe", fake_vjp) mesh = Mesh(np.asarray(jax.devices()[:1]).reshape(1, 1, 1), ("dp", "fsdp", "ep")) kwargs = dict( num_experts=2, - num_experts_per_tok=1, mesh_resource=MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="ep"), quant_before_fsdp_ag=quantize, ) with jax.set_mesh(mesh): module.moe( jnp.ones((1, 1, 4)), - jnp.ones((4, 2)), + RouterComputationInfo(gate_kernel=jnp.ones((4, 2)), num_experts_per_tok=1), jnp.ones((2, 4, 8)), jnp.ones((2, 4, 4)), **kwargs, @@ -160,7 +163,7 @@ def test_flax_block_forwards_boolean(monkeypatch, quantize): def fake_moe(inputs, *args, **kwargs): captured.update(kwargs) - return inputs, None, jnp.zeros((1,), jnp.int32) + return inputs, None, jnp.zeros((1,), jnp.int32), jnp.zeros((2,), jnp.int32) monkeypatch.setattr(module, "moe", fake_moe) mesh = Mesh(np.asarray(jax.devices()[:1]).reshape(1, 1), ("fsdp", "ep")) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 7ee4cb8d8fb..61759d74746 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -2396,7 +2396,7 @@ def execute(*args): assert bound["use_cudnn_fusion"] is (request_fusion is not False) assert bound["cudnn_native_weight_layout"] is native_layout received.append(bound["use_cudnn_jax_fusion"]) - return args[0], None, jnp.asarray(0) + return args[0], None, jnp.asarray(0), jnp.zeros((2,), jnp.int32) monkeypatch.setattr(_moe_module, "_moe", execute) x = jnp.ones((1, 1, 128), jnp.bfloat16) @@ -2407,13 +2407,14 @@ def execute(*args): kwargs["cudnn_native_weight_layout"] = native_layout with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") - output, _, _ = _moe_module.moe( + output, _, _, _ = _moe_module.moe( x, - jnp.ones((128, 2)), + _moe_module.RouterComputationInfo( + gate_kernel=jnp.ones((128, 2)), num_experts_per_tok=1 + ), wi, jnp.ones((2, 128, 128)), num_experts=2, - num_experts_per_tok=1, mesh_resource=MeshResource(ep_resource="ep"), **kwargs, ) @@ -2489,7 +2490,7 @@ def test_flax_forwards_fusion_bool(self, monkeypatch, request_fusion): def execute(inputs, *args, **kwargs): received.append(kwargs["use_cudnn_fusion"]) - return inputs, None, jnp.asarray(0) + return inputs, None, jnp.asarray(0), jnp.zeros((2,), jnp.int32) monkeypatch.setattr(flax_moe, "moe", execute) kwargs = {} if request_fusion is None else {"use_cudnn_fusion": request_fusion} @@ -2523,7 +2524,6 @@ def test_fusion_argument_requires_bool(self, invalid): None, None, num_experts=2, - num_experts_per_tok=1, use_cudnn_fusion=invalid, ) with pytest.raises(TypeError, match="use_cudnn_fusion must be a bool"): @@ -2538,4 +2538,5 @@ def test_fusion_argument_requires_bool(self, invalid): def test_vjp_bool_is_static_and_defaults_to_true(self): assert inspect.signature(_moe_module._moe).parameters["use_cudnn_fusion"].default is True assert inspect.signature(_moe_module.moe).parameters["use_cudnn_fusion"].default is True - assert 32 in _moe_module._moe.nondiff_argnums + parameters = list(inspect.signature(_moe_module._moe).parameters) + assert parameters.index("use_cudnn_fusion") in _moe_module._moe.nondiff_argnums diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 752fa680885..ad5c897b8d5 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -122,6 +122,8 @@ def _read_mp_options(): from transformer_engine.jax.flax import _MoEBlock as MoEBlock from transformer_engine.jax.moe import ( + RouterComputationInfo, + RoutingMapInfo, get_moe_recv_capacity_per_rank, moe, record_ep_bootstrap_signature_for_moe, @@ -288,11 +290,14 @@ def _pure_jax_moe_reference( wi, wo, expert_bias=None, + hash_input_ids=None, + hash_tid2eid=None, *, num_experts, num_experts_per_tok, aux_loss_coeff: float = 0.0, score_function: str = "softmax", + scaling_factor: float = 1.0, ): B, S, H = x.shape T = B * S @@ -302,7 +307,14 @@ def _pure_jax_moe_reference( gate_kernel_cast = gate_kernel.astype(x.dtype) logits = (x_2d @ gate_kernel_cast).astype(jnp.float32) # [T, E] - if score_function == "softmax": + if hash_input_ids is not None: + if score_function != "sqrtsoftplus" or hash_tid2eid is None: + raise ValueError("Hash reference requires sqrtsoftplus and hash_tid2eid.") + scores = jnp.sqrt(jax.nn.softplus(logits)) + top_indices = hash_tid2eid.astype(jnp.int32)[hash_input_ids.reshape(-1)] + weights = jnp.take_along_axis(scores, top_indices, axis=-1) + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + elif 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 @@ -319,6 +331,7 @@ def _pure_jax_moe_reference( weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) else: raise ValueError(f"Unsupported score_function={score_function!r}") + weights = weights * scaling_factor 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) @@ -338,6 +351,8 @@ def _pure_jax_moe_reference( output = output_2d.reshape(B, S, H).astype(x.dtype) if aux_loss_coeff > 0.0: + if hash_input_ids is not None: + raise ValueError("Hash routing does not use auxiliary loss.") # 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 @@ -373,6 +388,7 @@ def _make_block( aux_loss_coeff=0.0, use_expert_routing_bias=False, score_function="softmax", + scaling_factor=1.0, expert_bias_init=None, input_axes=("batch", None, None), quantization_recipe=None, @@ -390,6 +406,7 @@ def _make_block( aux_loss_coeff=aux_loss_coeff, use_expert_routing_bias=use_expert_routing_bias, score_function=score_function, + scaling_factor=scaling_factor, dtype=DTYPE, input_axes=input_axes, quantization_recipe=quantization_recipe, @@ -455,9 +472,9 @@ def _init_apply(block, mesh, x, key): 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, _trt = jax.jit(block.apply)(variables, x_sh) + output, aux, _trt, expert_counts = jax.jit(block.apply)(variables, x_sh) jax.block_until_ready(output) - return variables, output, aux + return variables, output, aux, expert_counts def _grad_step( @@ -478,7 +495,7 @@ def _grad_step( x_sh = _shard_inputs(x, mesh) def loss_fn(variables, x): - output, aux, _trt = block.apply(variables, x) + output, aux, _trt, _expert_counts = 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) @@ -497,7 +514,7 @@ def _grad_aux_only(block, variables, mesh, x): x_sh = _shard_inputs(x, mesh) def aux_only(variables, x): - _, aux, _trt = block.apply(variables, x) + _, aux, _trt, _expert_counts = block.apply(variables, x) return aux.astype(jnp.float32) grads = jax.jit(jax.grad(aux_only))(variables, x_sh) @@ -635,10 +652,10 @@ def layout_moe(*args, **kwargs): x = _make_inputs(jax.random.PRNGKey(41)) baseline = make_block(quantization_recipe=MXFP8BlockScaling()) quantized_ag = make_block(quantization_recipe=MXFP8BlockScaling(), quant_before_fsdp_ag=True) - variables, baseline_out, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(42)) + variables, baseline_out, _, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(42)) with _ctx(mesh): x_sh = _shard_inputs(x, mesh) - quantized_out, _, _ = jax.jit(quantized_ag.apply)(variables, x_sh) + quantized_out, _, _, _ = jax.jit(quantized_ag.apply)(variables, x_sh) quantized_out.block_until_ready() np.testing.assert_allclose( _to_global_numpy(quantized_out, mesh).astype(np.float32), @@ -670,15 +687,15 @@ def test_mesh_resource_api_forward_and_backward(mesh, quant_before_fsdp_ag, make mesh_resource=MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) ) x = _make_inputs(jax.random.PRNGKey(51)) - variables, baseline_out, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(52)) + variables, baseline_out, _, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(52)) baseline_grads, baseline_dx = _grad_step(baseline, variables, mesh, x) with jax.set_mesh(mesh), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): with global_shard_guard(None): x_sh = _shard_inputs(x, mesh) - candidate_out, _, _ = jax.jit(candidate.apply)(variables, x_sh) + candidate_out, _, _, _ = jax.jit(candidate.apply)(variables, x_sh) def loss_fn(variables, inputs): - output, _, _ = candidate.apply(variables, inputs) + output, _, _, _ = candidate.apply(variables, inputs) return jnp.mean(output.astype(jnp.float32) ** 2) candidate_grads, candidate_dx = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))( @@ -728,7 +745,7 @@ class TestTeEpMoeForward: def test_forward(self, mesh, config, quantization, make_block): block = make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(0)) - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) + variables, output, aux, expert_counts = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) # Shape / dtype / finiteness (cheap; on the local shard). assert output.shape == x.shape @@ -736,6 +753,10 @@ def test_forward(self, mesh, config, quantization, make_block): 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" + counts_np = np.asarray(jax.device_get(expert_counts)) + assert expert_counts.shape == (NUM_EXPERTS,) + assert expert_counts.dtype == jnp.int32 + assert int(counts_np.sum()) == BATCH * SEQ * TOPK # Numerical parity (replicated global view -> single rank's numpy). params_np = _params_global_numpy(variables, mesh) @@ -759,6 +780,263 @@ def test_forward(self, mesh, config, quantization, make_block): ) +class TestTeEpMoeHashRouting: + """DeepSeek-V4 hash selection with differentiable gate weights.""" + + def test_hash_forward_and_backward(self, mesh, make_block, use_cudnn_fusion): + vocab_size = 64 + scaling_factor = 1.5 + token_ids = jnp.arange(BATCH * SEQ, dtype=jnp.int32).reshape(BATCH, SEQ) % vocab_size + vocab_ids = jnp.arange(vocab_size, dtype=jnp.int32) + tid2eid = jnp.stack( + [vocab_ids % NUM_EXPERTS, (vocab_ids + 3) % NUM_EXPERTS], + axis=-1, + ) + block = make_block( + score_function="sqrtsoftplus", + scaling_factor=scaling_factor, + ) + x = _make_inputs(jax.random.PRNGKey(30)) + + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + ids_sh = jax.device_put( + token_ids, + NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None)), + ) + table_sh = jax.device_put(tid2eid, NamedSharding(mesh, P())) + variables = jax.jit(block.init)(jax.random.PRNGKey(31), x_sh) + params = { + name: _unwrap(variables["params"][name]) for name in ("gate_kernel", "wi", "wo") + } + + def externally_routed_moe(params, activations, token_ids_arg, tid2eid_arg): + logits = jnp.einsum( + "bsh,he->bse", + activations, + params["gate_kernel"].astype(activations.dtype), + ).astype(jnp.float32) + scores = jnp.sqrt(jax.nn.softplus(logits)) + routing_indices = tid2eid_arg[token_ids_arg] + routing_weights = jnp.take_along_axis(scores, routing_indices, axis=-1) + routing_weights = routing_weights / ( + routing_weights.sum(axis=-1, keepdims=True) + 1e-20 + ) + routing_weights = routing_weights * scaling_factor + return moe( + activations, + RoutingMapInfo( + routing_indices=routing_indices, + routing_weights=routing_weights, + ), + params["wi"], + params["wo"], + num_experts=NUM_EXPERTS, + input_axes=("batch", None, None), + wi_kernel_axes=("exp", "embed", "mlp"), + wo_kernel_axes=("exp", "mlp", "embed"), + dtype=DTYPE, + collect_expert_counts=True, + use_cudnn_fusion=use_cudnn_fusion, + ) + + output, aux, _trt, expert_counts = jax.jit(externally_routed_moe)( + params, x_sh, ids_sh, table_sh + ) + jax.block_until_ready(output) + assert aux is None, "aux_loss must be None for caller-supplied routing" + + def te_loss(params, activations, token_ids_arg, tid2eid_arg): + out, _aux, _total_recv, _counts = externally_routed_moe( + params, activations, token_ids_arg, tid2eid_arg + ) + return jnp.mean(out.astype(jnp.float32) ** 2) + + grads_te, grad_x_te = jax.jit(jax.grad(te_loss, argnums=(0, 1)))( + params, x_sh, ids_sh, table_sh + ) + + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + ids_np = np.asarray(jax.device_get(token_ids)) + table_np = np.asarray(jax.device_get(tid2eid)) + + def ref_loss(params, activations): + out, _ = _pure_jax_moe_reference( + activations, + params["gate_kernel"], + params["wi"], + params["wo"], + hash_input_ids=jnp.asarray(ids_np), + hash_tid2eid=jnp.asarray(table_np), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + score_function="sqrtsoftplus", + scaling_factor=scaling_factor, + ) + return jnp.mean(out.astype(jnp.float32) ** 2) + + ref_params = {name: jnp.asarray(params_np[name]) for name in ("gate_kernel", "wi", "wo")} + grads_ref, grad_x_ref = jax.jit(jax.grad(ref_loss, argnums=(0, 1)))( + ref_params, jnp.asarray(x_np) + ) + + np.testing.assert_allclose( + _to_global_numpy(output, mesh).astype(np.float32), + np.asarray( + _pure_jax_moe_reference( + jnp.asarray(x_np), + ref_params["gate_kernel"], + ref_params["wi"], + ref_params["wo"], + hash_input_ids=jnp.asarray(ids_np), + hash_tid2eid=jnp.asarray(table_np), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + score_function="sqrtsoftplus", + scaling_factor=scaling_factor, + )[0] + ).astype(np.float32), + **FWD_TOLERANCE["bf16"], + ) + expected_counts = np.bincount(table_np[ids_np].reshape(-1), minlength=NUM_EXPERTS) + np.testing.assert_array_equal(np.asarray(jax.device_get(expert_counts)), expected_counts) + + # Use a slightly larger tolerance as the numerical difference is just over 1e-7. + grad_ffn_tolerance_bf16 = {"atol": 2e-7, "rtol": 1e-7} + + for name in ("gate_kernel", "wi", "wo"): + np.testing.assert_allclose( + _to_global_numpy(grads_te[name], mesh).astype(np.float32), + np.asarray(jax.device_get(grads_ref[name])).astype(np.float32), + **( + GRAD_GATE_TOLERANCE["bf16"] + if name == "gate_kernel" + else grad_ffn_tolerance_bf16 + ), + ) + np.testing.assert_allclose( + _to_global_numpy(grad_x_te, mesh).astype(np.float32), + np.asarray(jax.device_get(grad_x_ref)).astype(np.float32), + **grad_ffn_tolerance_bf16, + ) + + +class TestTeEpMoeRouterInfo: + """Functional ``moe(router_info=...)`` API surface.""" + + def test_router_computation_matches_block(self, mesh, make_block, use_cudnn_fusion): + """Calling ``moe`` with a RouterComputationInfo reproduces _MoEBlock.""" + scaling_factor = 1.5 + aux_loss_coeff = 1e-2 + block = make_block( + score_function="sigmoid", + use_expert_routing_bias=True, + expert_bias_init=_strong_expert_bias_init, + scaling_factor=scaling_factor, + aux_loss_coeff=aux_loss_coeff, + ) + x = _make_inputs(jax.random.PRNGKey(40)) + variables, out_block, aux_block, counts_block = _init_apply( + block, mesh, x, jax.random.PRNGKey(41) + ) + params = {name: _unwrap(p) for name, p in variables["params"].items()} + + def functional_moe(params, activations): + router_info = RouterComputationInfo( + gate_kernel=params["gate_kernel"], + num_experts_per_tok=TOPK, + expert_bias=params["expert_bias"], + score_function="sigmoid", + scaling_factor=scaling_factor, + aux_loss_coeff=aux_loss_coeff, + ) + return moe( + activations, + router_info, + params["wi"], + params["wo"], + num_experts=NUM_EXPERTS, + input_axes=("batch", None, None), + dtype=DTYPE, + use_cudnn_fusion=use_cudnn_fusion, + ) + + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + out_fn, aux_fn, _trt, counts_fn = jax.jit(functional_moe)(params, x_sh) + jax.block_until_ready(out_fn) + + np.testing.assert_allclose( + _to_global_numpy(out_fn, mesh).astype(np.float32), + _to_global_numpy(out_block, mesh).astype(np.float32), + **FWD_TOLERANCE["bf16"], + ) + assert aux_fn is not None and aux_block is not None + np.testing.assert_allclose( + float(_to_global_numpy(aux_fn, mesh)), + float(_to_global_numpy(aux_block, mesh)), + **AUX_TOLERANCE, + ) + np.testing.assert_array_equal( + np.asarray(jax.device_get(counts_fn)), np.asarray(jax.device_get(counts_block)) + ) + + def test_router_info_pytree_leaves(self): + """Only arrays are pytree leaves; router settings stay static.""" + gate_kernel = jnp.zeros((HIDDEN, NUM_EXPERTS), DTYPE) + expert_bias = jnp.zeros((NUM_EXPERTS,), jnp.float32) + router = RouterComputationInfo( + gate_kernel=gate_kernel, + num_experts_per_tok=TOPK, + expert_bias=expert_bias, + score_function="sigmoid", + ) + leaves = jax.tree_util.tree_leaves(router) + assert len(leaves) == 2 + assert leaves[0] is gate_kernel and leaves[1] is expert_bias + + indices = jnp.zeros((BATCH, SEQ, TOPK), jnp.int32) + weights = jnp.ones((BATCH, SEQ, TOPK), jnp.float32) + leaves = jax.tree_util.tree_leaves( + RoutingMapInfo(routing_indices=indices, routing_weights=weights) + ) + assert len(leaves) == 2 + assert leaves[0] is indices and leaves[1] is weights + + def test_rejects_unknown_router_info(self): + x = jnp.zeros((BATCH, SEQ, HIDDEN), DTYPE) + wi = jnp.zeros((NUM_EXPERTS, HIDDEN, 2 * INTER), DTYPE) + wo = jnp.zeros((NUM_EXPERTS, INTER, HIDDEN), DTYPE) + with pytest.raises(TypeError, match="router_info"): + moe( + x, + object(), + wi, + wo, + num_experts=NUM_EXPERTS, + mesh_resource=MeshResource(ep_resource=EP_AXIS), + ) + + def test_routing_map_shape_validation(self): + x = jnp.zeros((BATCH, SEQ, HIDDEN), DTYPE) + wi = jnp.zeros((NUM_EXPERTS, HIDDEN, 2 * INTER), DTYPE) + wo = jnp.zeros((NUM_EXPERTS, INTER, HIDDEN), DTYPE) + indices = jnp.zeros((BATCH, SEQ, TOPK), jnp.int32) + with pytest.raises(ValueError, match="routing_weights"): + moe( + x, + RoutingMapInfo( + routing_indices=indices, + routing_weights=jnp.ones((BATCH, SEQ, TOPK + 1), jnp.float32), + ), + wi, + wo, + num_experts=NUM_EXPERTS, + mesh_resource=MeshResource(ep_resource=EP_AXIS), + ) + + class TestTeEpMoeBackward: """Per-config backward correctness in a single run: per-tensor grads finite, non-zero AND parity vs the pure-JAX reference.""" @@ -768,7 +1046,7 @@ class TestTeEpMoeBackward: def test_backward(self, mesh, config, quantization, make_block): block = make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(2)) - variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) + 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 @@ -850,7 +1128,7 @@ def recorded_checkpoint_name(value, name): dispatch_checkpoint_name="saved_dispatch", ) inputs = _make_inputs(jax.random.PRNGKey(34)) - variables, output, _ = _init_apply(block, mesh, inputs, jax.random.PRNGKey(35)) + variables, output, _, _ = _init_apply(block, mesh, inputs, jax.random.PRNGKey(35)) grads, grad_inputs = _grad_step(block, variables, mesh, inputs) assert len(named_values["saved_dispatch"]) >= 2 @@ -877,7 +1155,7 @@ def unexpected_selection(*args, **kwargs): monkeypatch.setattr(moe_module, "_select_cudnn_jax_fusion", unexpected_selection) block = make_block(quantization_recipe=MXFP8BlockScaling(), use_cudnn_fusion=False) x = _make_inputs(jax.random.PRNGKey(53)) - variables, baseline_output, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(54)) + variables, baseline_output, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(54)) baseline_grads, baseline_dx = _grad_step(block, variables, mesh, x) flax_moe = importlib.import_module("transformer_engine.jax.flax.moe") original_moe = flax_moe.moe @@ -892,7 +1170,7 @@ def native_moe(*args, **kwargs): monkeypatch.setattr(flax_moe, "moe", native_moe) with _ctx(mesh): - output, _, _ = jax.jit(block.apply)(variables, _shard_inputs(x, mesh)) + output, _, _, _ = jax.jit(block.apply)(variables, _shard_inputs(x, mesh)) output.block_until_ready() grads, dx = _grad_step(block, variables, mesh, x) np.testing.assert_array_equal( @@ -942,7 +1220,9 @@ def unexpected_fused_kernel(*args, **kwargs): moe_module = importlib.import_module("transformer_engine.jax.moe") with monkeypatch.context() as baseline_patch: baseline_patch.setattr(moe_module, "_select_cudnn_jax_fusion", lambda _: False) - variables, baseline_output, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(52)) + variables, baseline_output, _, _ = _init_apply( + baseline, mesh, x, jax.random.PRNGKey(52) + ) baseline_grads, baseline_dx = _grad_step(baseline, variables, mesh, x) if native_layout: @@ -959,7 +1239,7 @@ def native_moe(*args, **kwargs): monkeypatch.setattr(flax_moe, "moe", native_moe) with _ctx(mesh): - output, _, _ = jax.jit(block.apply)(variables, _shard_inputs(x, mesh)) + output, _, _, _ = jax.jit(block.apply)(variables, _shard_inputs(x, mesh)) output.block_until_ready() grads, dx = _grad_step(block, variables, mesh, x) np.testing.assert_array_equal( @@ -1026,7 +1306,7 @@ def checked_dglu(*args, **kwargs): quantization_recipe=MXFP8BlockScaling(), ) x = _make_inputs(jax.random.PRNGKey(30)) - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(31)) + variables, output, aux, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(31)) assert "glu" in shared_calls, "MoE did not select the shared cuDNN GLU JAX API" grads, grad_x = _grad_step(block, variables, mesh, x) assert "dglu" in shared_calls, "MoE VJP did not select the shared cuDNN dGLU JAX API" @@ -1139,7 +1419,7 @@ def checkpointed_moe(*args, **kwargs): dispatch_checkpoint_name="saved_dispatch", ) inputs = _make_inputs(jax.random.PRNGKey(32)) - variables, output, _ = _init_apply(block, mesh, inputs, jax.random.PRNGKey(33)) + variables, output, _, _ = _init_apply(block, mesh, inputs, jax.random.PRNGKey(33)) grads, grad_inputs = _grad_step(block, variables, mesh, inputs) assert selected_calls, f"{fused_op_name} was not selected" @@ -1168,7 +1448,7 @@ def test_aux_loss(self, mesh, make_block): 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)) + 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" @@ -1211,7 +1491,7 @@ def test_combined_loss_grads(self, mesh, make_block): 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)) + 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", "wo"): g_local = np.asarray(jax.device_get(_unwrap(grads["params"][name]).addressable_data(0))) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 6c4d5a2c8cf..19b1a41778f 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -34,7 +34,7 @@ from flax import linen as nn from transformer_engine.common.recipe import Recipe -from ..moe import _moe_mesh_axes, _resolve_moe_mesh_resource, moe +from ..moe import RouterComputationInfo, _moe_mesh_axes, _resolve_moe_mesh_resource, moe from ..quantize import QuantizerSet from ..router import ScoreFunction from ..sharding import MeshResource, _get_mesh, global_shard_guard @@ -115,6 +115,10 @@ class _MoEBlock(TransformerEngineBase): dispatch_checkpoint_name : Optional[str] JAX rematerialization checkpoint name for the EP dispatch outputs. ``None`` leaves them unnamed. + collect_expert_counts : bool + Compute global per-expert assignment counts. Set to ``False`` when + callers do not consume the counts to avoid the associated all-reduce; + the returned ``expert_counts`` is then a zero-filled array. Per-expert dispatch-slot alignment is 256 tokens when fusion is requested, including fallbacks, and 128 tokens with ``use_cudnn_fusion=False``. @@ -162,6 +166,7 @@ class _MoEBlock(TransformerEngineBase): apply_topk_weights_early: bool = False recv_capacity_per_rank: Optional[int] = None dispatch_checkpoint_name: Optional[str] = None + collect_expert_counts: bool = True # Dtypes / init / misc dtype: DType = jnp.float32 @@ -183,7 +188,7 @@ def __post_init__(self): super().__post_init__() @nn.compact - def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: + def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array, Array]: """Run the MoE forward pass. Parameters @@ -201,6 +206,8 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: total_recv_tokens : jnp.ndarray Non-differentiable per-rank pre-drop recv-slot total; flags overflow when ``drop_on_overflow`` is set at ep_bootstrap. + expert_counts : jnp.ndarray + Global assignment counts for each expert. """ mesh_resource, quant_before_fsdp_ag = _resolve_moe_mesh_resource( self.mesh_resource, @@ -300,33 +307,37 @@ def make_grouped_quantizer_set(postfix): make_grouped_quantizer_set("_fc2"), ) + router_info = RouterComputationInfo( + gate_kernel=gate_kernel, + num_experts_per_tok=self.num_experts_per_tok, + expert_bias=expert_bias, + score_function=self.score_function, + use_pre_softmax=self.use_pre_softmax, + num_groups=self.num_groups, + group_topk=self.group_topk, + scaling_factor=self.scaling_factor, + aux_loss_coeff=self.aux_loss_coeff, + gate_kernel_axes=self.gate_kernel_axes, + ) return moe( inputs, - gate_kernel, + router_info, wi, wo, wi_0_bias, wi_1_bias, wo_bias, - expert_bias, num_experts=self.num_experts, - num_experts_per_tok=self.num_experts_per_tok, activation_type=self.activation_type, - score_function=self.score_function, - use_pre_softmax=self.use_pre_softmax, - num_groups=self.num_groups, - group_topk=self.group_topk, - scaling_factor=self.scaling_factor, - aux_loss_coeff=self.aux_loss_coeff, use_cudnn_fusion=self.use_cudnn_fusion, cudnn_native_weight_layout=False, apply_topk_weights_early=self.apply_topk_weights_early, quantizer_sets=quantizer_sets, recv_capacity_per_rank=self.recv_capacity_per_rank, + collect_expert_counts=self.collect_expert_counts, quant_before_fsdp_ag=quant_before_fsdp_ag, mesh_resource=mesh_resource, 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, dtype=self.dtype, diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index d762719f9ae..44214063ac4 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -3,14 +3,23 @@ # See LICENSE for license information. """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:: +This module exposes :func:`moe`, a fused MoE forward + backward path built on +top of TE's NCCL-backed Expert Parallelism primitives +(``tex.ep_dispatch`` / ``tex.ep_combine``). The ``router_info`` argument +selects how tokens are routed: - gate -> topk -> ep_dispatch -> per-expert FFN (grouped GEMMs) - -> ep_combine -> output +* :class:`RouterComputationInfo` -- TE computes the router:: -under a single ``jax.custom_vjp`` so the routing, dispatch, FFN and + gate -> topk -> ep_dispatch -> per-expert FFN (grouped GEMMs) + -> ep_combine -> output + +* :class:`RoutingMapInfo` -- the caller supplies expert indices and routing + weights (e.g. hash routing), and the ``gate -> topk`` stage is skipped:: + + external routing -> ep_dispatch -> per-expert FFN (grouped GEMMs) + -> ep_combine -> output + +Both paths run 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. @@ -62,7 +71,90 @@ from .router import ScoreFunction, _validate_score_function from .sharding import MeshResource, _get_mesh, global_mesh_resource, global_shard_guard -__all__ = ["get_moe_recv_capacity_per_rank", "moe"] +__all__ = [ + "RouterComputationInfo", + "RoutingMapInfo", + "get_moe_recv_capacity_per_rank", + "moe", +] + + +@flax.struct.dataclass +class RouterComputationInfo: + """Router parameters and settings for :func:`moe` to compute routing itself. + + Array fields are pytree leaves; the remaining fields are static. + + Parameters + ---------- + gate_kernel : jnp.ndarray + ``[hidden, num_experts]`` router projection. + num_experts_per_tok : int + Top-k value for routing. + expert_bias : Optional[jnp.ndarray] + ``[num_experts]`` learnable router bias added before the top-k + when ``score_function='sigmoid'`` or ``'sqrtsoftplus'``. ``None`` + disables it. The bias has no gradient through the top-k primitive + itself (it only steers expert selection); a zero cotangent is + returned for it. + score_function : Union[str, ScoreFunction] + ``"softmax"`` (default), ``"sigmoid"``, or ``"sqrtsoftplus"``. + use_pre_softmax : bool + Apply softmax before top-k (vs. after). + num_groups, group_topk : Optional[int] + Grouped top-k knobs (DeepSeek-style). ``None`` disables grouping. + scaling_factor : float + Multiplier on the routing weights. + 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. + gate_kernel_axes : Tuple[Optional[str], ...] + Logical sharding axes of ``gate_kernel``. + """ + + gate_kernel: jnp.ndarray + num_experts_per_tok: int = flax.struct.field(pytree_node=False) + expert_bias: Optional[jnp.ndarray] = None + score_function: Union[str, ScoreFunction] = flax.struct.field( + pytree_node=False, default="softmax" + ) + use_pre_softmax: bool = flax.struct.field(pytree_node=False, default=False) + num_groups: Optional[int] = flax.struct.field(pytree_node=False, default=None) + group_topk: Optional[int] = flax.struct.field(pytree_node=False, default=None) + scaling_factor: float = flax.struct.field(pytree_node=False, default=1.0) + aux_loss_coeff: float = flax.struct.field(pytree_node=False, default=0.0) + gate_kernel_axes: Tuple[Optional[str], ...] = flax.struct.field(pytree_node=False, default=()) + + +@flax.struct.dataclass +class RoutingMapInfo: + """Caller-computed routing for :func:`moe`, bypassing TE's router. + + Note: the shape of ``routing_indices`` provides the number of experts per token. + + Parameters + ---------- + routing_indices : jnp.ndarray + Integer expert indices with shape ``[batch, sequence, top_k]``. + Non-differentiable. Each token's ``top_k`` indices must be distinct, + otherwise the default receive-capacity bound is invalid. + routing_weights : jnp.ndarray + Routing weights with the same shape as ``routing_indices``. + Gradients are returned through the custom VJP, so caller-side + routers remain fully differentiable. + """ + + routing_indices: jnp.ndarray + routing_weights: jnp.ndarray + + +def _num_experts_per_tok(router_info: Union[RouterComputationInfo, RoutingMapInfo]) -> int: + if isinstance(router_info, RouterComputationInfo): + return router_info.num_experts_per_tok + return router_info.routing_indices.shape[-1] def _moe_mesh_axes(resource: MeshResource): @@ -231,6 +323,7 @@ def get_moe_recv_capacity_per_rank( supplied, ``use_cudnn_fusion=True`` (default) reserves 256-token alignment, including fallbacks; ``False`` reserves 128 for regular TE grouped GEMM. This keeps eager bootstrap callers in sync with the later compiled ``moe()`` call. + The bound assumes each token's top-k expert indices are distinct. """ if not isinstance(use_cudnn_fusion, bool): raise TypeError("use_cudnn_fusion must be a bool") @@ -313,6 +406,60 @@ def _constraint_bwd(dtype_ref, grad): return _constraint(x) +def _prepare_moe_input( + x: jnp.ndarray, + ep_axis: str, + data_parallelism_axes: Tuple[str, ...], + *, + caller: str, +) -> jnp.ndarray: + """Validate and constrain an MoE input to the EP communication layout. + + NCCL EP forms communication groups from consecutive global ranks, with + ``ep_axis`` as the innermost component of the compound leading mesh axis. + The expected activation layout is therefore + ``P((*data_parallelism_axes, ep_axis), None, None)``. If ``x`` has an + incompatible known sharding, this function warns before inserting the + reshard. The custom constraint also keeps the backward cotangent in the + input activation dtype. + + Parameters + ---------- + x : jnp.ndarray + Rank-3 MoE input activation. + ep_axis : str + Physical mesh axis used for expert parallelism. + data_parallelism_axes : Tuple[str, ...] + Physical outer data-parallel mesh axes. + caller : str + Public API name used in diagnostics. + + Returns + ------- + jnp.ndarray + ``x`` constrained to the EP-compatible input sharding. + """ + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError(f"{caller}(...) requires an active jax.sharding.Mesh.") + for axis in (ep_axis, *data_parallelism_axes): + if axis not in mesh.shape: + raise ValueError(f"TE MoE resource axis {axis!r} is not in the active 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"{caller}(...): 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=3, + ) + return _with_sharding_constraint_cast_bwd(x, NamedSharding(mesh, expected_spec)) + + # ============================================================================= # Process-level NCCL EP bootstrap (must run eagerly, outside jax.jit) # ============================================================================= @@ -394,17 +541,19 @@ def _te_ep_assert_compatible_bootstrap( class _Ctx: """Residuals carried from the fwd rule into the bwd rule. - 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). + Flattened automatically by jax.custom_vjp; ``cfg`` and + ``routing_weights_dtype`` are static fields (the rest are jnp.ndarray, + GroupedNoScaleTensor, or None when aux_loss_coeff == 0). ``router`` and the router-side + residuals (``logits_2d``, ``saved_scores``, ``routing_map``) are None + when routing is caller-supplied via :class:`RoutingMapInfo`. """ x: jnp.ndarray - gate_kernel: jnp.ndarray - expert_bias: jnp.ndarray - logits_2d: jnp.ndarray - saved_scores: jnp.ndarray - routing_map: jnp.ndarray + router: Optional[RouterComputationInfo] + logits_2d: Optional[jnp.ndarray] + saved_scores: Optional[jnp.ndarray] + routing_map: Optional[jnp.ndarray] + selected_experts: jnp.ndarray cfg: Any = flax.struct.field(pytree_node=False) handle_mem: jnp.ndarray recv_topk_weights: jnp.ndarray @@ -420,6 +569,8 @@ class _Ctx: aux_const_buf: Any = None aux_tokens_per_expert: Any = None aux_saved_scores: Any = None + # Caller-supplied routing only: dtype of the routing_weights cotangent. + routing_weights_dtype: Any = flax.struct.field(pytree_node=False, default=None) # ============================================================================= @@ -1135,31 +1286,23 @@ def _ffn_bwd_per_shard( def _moe_fwd_rule( x, - gate_kernel, + router_info, wi, wo, wi_0_bias, wi_1_bias, wo_bias, - expert_bias, quantizer_sets, num_experts, - num_experts_per_tok, activation_type, - score_function, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - aux_loss_coeff, mesh_resource, input_axes, - gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, dtype, apply_topk_weights_early, recv_capacity_per_rank, + collect_expert_counts, use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, @@ -1169,16 +1312,18 @@ def _moe_fwd_rule( use_cudnn_fusion: bool = True, cudnn_native_weight_layout: Optional[bool] = None, ): - """Forward: gate -> topk -> ep_dispatch -> FFN -> ep_combine. + """Forward: [gate -> route ->] ep_dispatch -> FFN -> ep_combine. - Returns ``(output, aux_loss)``. ``aux_loss`` is a zero scalar when - ``aux_loss_coeff == 0``. + Returns ``(output, aux_loss, total_recv_tokens, expert_counts)``. + ``aux_loss`` is a zero scalar unless the router computes it and ``aux_loss_coeff`` > 0. """ with global_shard_guard(mesh_resource): ep_axis, data_parallelism_axes = _moe_mesh_axes(mesh_resource) - del gate_kernel_axes # used in bwd only from jax.experimental.shard_map import shard_map + router = router_info if isinstance(router_info, RouterComputationInfo) else None + compute_aux_loss = router is not None and router.aux_loss_coeff > 0.0 + x = with_sharding_constraint_by_logical_axes(x, input_axes) mesh = _get_mesh() @@ -1201,7 +1346,7 @@ def _moe_fwd_rule( num_expert_groups=num_experts, ) B, S, H = x.shape - K = num_experts_per_tok + K = _num_experts_per_tok(router_info) cudnn_native_weight_layout = _resolve_cudnn_native_weight_layout( wi, H, cudnn_native_weight_layout ) @@ -1283,35 +1428,75 @@ def _moe_fwd_rule( 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) + # ------------------------ Routing ------------------------ + + if router is None: + # ---------------- Caller-supplied routing ---------------- + routing_indices = jax.lax.with_sharding_constraint( + router_info.routing_indices, NamedSharding(mesh, ep3_spec) + ) + external_weights = jax.lax.with_sharding_constraint( + router_info.routing_weights, NamedSharding(mesh, ep3_spec) + ) + selected_experts = routing_indices.reshape(-1, K).astype(jnp.int32) + routing_weights = external_weights.reshape(-1, K).astype(dtype) + logits_2d = routing_map = saved_scores = None + else: + # ---------------- 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 = router.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) ---------------- + # The fused top-k primitive takes an empty shape-(0,) tensor as "no bias". + if router.expert_bias is None: + eb_arg = jnp.zeros((0,), dtype=jnp.float32) + else: + eb_arg = router.expert_bias.astype(jnp.float32) + topk_kwargs = { + "topk": K, + "use_pre_softmax": router.use_pre_softmax, + "num_groups": -1 if router.num_groups is None else router.num_groups, + "group_topk": -1 if router.group_topk is None else router.group_topk, + "scaling_factor": router.scaling_factor, + "score_function": router.score_function, + "expert_bias": eb_arg, + "compute_aux_scores": False, + } + sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( + logits_2d, **topk_kwargs + ) + sparse_probs = sparse_probs.astype(dtype) + # 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) + + if collect_expert_counts: + if router is None: + expert_counts = jnp.sum( + jax.nn.one_hot(selected_experts, num_experts, dtype=jnp.int32), + axis=(0, 1), + ) + else: + expert_counts = jnp.sum(routing_map.astype(jnp.int32), axis=0) + # The token dimension is sharded over the compound data/EP axis. + # Replicating the reduced result emits the required all-reduce. + expert_counts = jax.lax.with_sharding_constraint( + expert_counts, NamedSharding(mesh, P()) + ) + else: + # Keep the custom-VJP output structure stable without paying for an + # otherwise-unused per-layer all-reduce. + expert_counts = jnp.zeros((num_experts,), dtype=jnp.int32) # ---------------- Aux loss (global view, replicated) ---------------- # ``fused_moe_aux_loss_fwd`` sums probs and tokens_per_expert across @@ -1320,18 +1505,10 @@ def _moe_fwd_rule( # 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: + if compute_aux_loss: 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, + global_logits_2d, **topk_kwargs ) aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) # compute_aux_scores=True takes a separate kernel path: clean @@ -1343,7 +1520,7 @@ def _moe_fwd_rule( num_groups=-1, group_topk=-1, scaling_factor=1.0, - score_function=score_function, + score_function=router.score_function, expert_bias=jnp.zeros((0,), dtype=jnp.float32), compute_aux_scores=True, ) @@ -1351,7 +1528,7 @@ def _moe_fwd_rule( aux_probs.astype(jnp.float32), aux_tokens_per_expert.astype(jnp.int32), topk=K, - coeff=aux_loss_coeff, + coeff=router.aux_loss_coeff, ) aux_loss = aux_loss.astype(dtype) else: @@ -1361,10 +1538,6 @@ def _moe_fwd_rule( aux_saved_scores = None # ---------------- 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 @@ -1504,11 +1677,11 @@ def _ffn_fwd_body(*args): ctx = _Ctx( x=x, - gate_kernel=gate_kernel, - expert_bias=expert_bias, + router=router, logits_2d=logits_2d, saved_scores=saved_scores, routing_map=routing_map, + selected_experts=selected_experts, cfg=cfg, handle_mem=handle_mem, recv_topk_weights=recv_topk_weights, @@ -1524,6 +1697,9 @@ def _ffn_fwd_body(*args): aux_const_buf=aux_const_buf, aux_tokens_per_expert=aux_tokens_per_expert, aux_saved_scores=aux_saved_scores, + routing_weights_dtype=( + None if router is not None else router_info.routing_weights.dtype + ), ) static = { "has_bias": has_bias, @@ -1531,28 +1707,21 @@ def _ffn_fwd_body(*args): "recv_pr": recv_pr, "cudnn_native_weight_layout": cudnn_native_weight_layout, } - # total_recv_tokens is a non-differentiable overflow signal (see moe()). - return (output, aux_loss, total_recv_tokens), (ctx, static) + # The integer metadata outputs are non-differentiable (see moe()). + return (output, aux_loss, total_recv_tokens, expert_counts), (ctx, static) def _moe_bwd_rule( num_experts, - num_experts_per_tok, activation_type, - score_function, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - aux_loss_coeff, mesh_resource, input_axes, - gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, dtype, apply_topk_weights_early, recv_capacity_per_rank, + collect_expert_counts, use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, @@ -1568,10 +1737,9 @@ def _moe_bwd_rule( with global_shard_guard(mesh_resource): ep_axis, data_parallelism_axes = _moe_mesh_axes(mesh_resource) del ( - num_groups, - group_topk, dtype, recv_capacity_per_rank, + collect_expert_counts, wi_0_checkpoint_name, wi_1_checkpoint_name, wo_checkpoint_name, @@ -1582,10 +1750,11 @@ def _moe_bwd_rule( ) # captured / unused in bwd from jax.experimental.shard_map import shard_map - # total_recv_tokens is a non-differentiable output; its cotangent is unused. - d_output, d_aux_loss, _d_total_recv_tokens = cotangents + # The integer metadata outputs are non-differentiable; their cotangents are unused. + d_output, d_aux_loss, _d_total_recv_tokens, _d_expert_counts = cotangents ctx, static = residuals + router = ctx.router has_bias = static["has_bias"] x_shape = static["x_shape"] recv_pr = static["recv_pr"] @@ -1594,7 +1763,7 @@ def _moe_bwd_rule( if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") B, S, _ = x_shape - K = num_experts_per_tok + K = ctx.selected_experts.shape[-1] if not data_parallelism_axes: batch_pspec_axis: Any = ep_axis else: @@ -1740,69 +1909,63 @@ def _ffn_bwd_body(*args): out_partition_spec=out_partition_spec, ) - # ---------------- 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), + if router is None: + # Caller-supplied routing: d_topk_w is the routing-weights cotangent, and + # the integer routing indices get a zero (None) cotangent. + d_router_info = RoutingMapInfo( + routing_indices=None, + routing_weights=d_topk_w.astype(ctx.routing_weights_dtype), ) - # 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), + d_x = d_x_from_dispatch + else: + # ---------------- Routing bwd (global view) ---------------- + # The cotangent on routing_weights is a sparse scatter into sparse_probs + # at the selected_experts indices. + 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], ctx.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=False, - scaling_factor=1.0, - score_function=score_function, - compute_aux_scores=True, + use_pre_softmax=router.use_pre_softmax, + scaling_factor=router.scaling_factor, + score_function=router.score_function, + compute_aux_scores=False, + ) + # ---------------- Aux loss bwd (global view, replicated) ---------------- + if router.aux_loss_coeff > 0.0: + d_logits_2d = d_logits_2d + _moe_aux_loss_bwd(ctx, d_aux_loss, K).astype( + d_logits_2d.dtype + ) + + # ---------------- Gate bwd (global view) ---------------- + d_gate_logits = d_logits_2d.reshape(B, S, num_experts) + gate_kernel_cast = router.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( + router.gate_kernel.dtype + ) + d_gate_kernel = with_sharding_constraint_by_logical_axes( + d_gate_kernel, router.gate_kernel_axes + ) + d_x = d_x_from_gate + d_x_from_dispatch + # expert_bias has no learnable bwd path through fused_topk: the + # primitive's bwd returns None for the bias slot, so its cotangent is zero. + d_router_info = router.replace( + gate_kernel=d_gate_kernel, + expert_bias=( + None if router.expert_bias is None else jnp.zeros_like(router.expert_bias) + ), ) - 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 = with_sharding_constraint_by_logical_axes(d_wi, wi_kernel_axes) d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_axes) if has_bias: @@ -1812,58 +1975,72 @@ def _ffn_bwd_body(*args): d_wi_1_bias = with_sharding_constraint_by_logical_axes(d_wi_1_bias, wi_bias_axes) d_wo_bias = with_sharding_constraint_by_logical_axes(d_wo_bias, wo_bias_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) - return ( d_x, - d_gate_kernel, + d_router_info, d_wi, 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, ctx.quantizer_sets, ) +def _moe_aux_loss_bwd(ctx, d_aux_loss, K): + """Aux-loss contribution to the router logits cotangent.""" + # 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. + 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) + return 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=ctx.router.score_function, + compute_aux_scores=True, + ) + + # ============================================================================= # custom_vjp + public entry # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 34))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(8, 26))) def _moe( x, - gate_kernel, + router_info, wi, wo, wi_0_bias, wi_1_bias, wo_bias, - expert_bias, quantizer_sets, num_experts, - num_experts_per_tok, activation_type, - score_function, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - aux_loss_coeff, mesh_resource, input_axes, - gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, dtype, apply_topk_weights_early, recv_capacity_per_rank, + collect_expert_counts, use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, @@ -1875,31 +2052,23 @@ def _moe( ): primal, _ = _moe_fwd_rule( x, - gate_kernel, + router_info, wi, wo, wi_0_bias, wi_1_bias, wo_bias, - expert_bias, quantizer_sets, num_experts, - num_experts_per_tok, activation_type, - score_function, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - aux_loss_coeff, mesh_resource, input_axes, - gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, dtype, apply_topk_weights_early, recv_capacity_per_rank, + collect_expert_counts, use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, @@ -1917,23 +2086,15 @@ def _moe( def moe( x: jnp.ndarray, - gate_kernel: jnp.ndarray, + router_info: Union[RouterComputationInfo, RoutingMapInfo], wi: jnp.ndarray, wo: jnp.ndarray, wi_0_bias: Optional[jnp.ndarray] = None, wi_1_bias: Optional[jnp.ndarray] = None, wo_bias: Optional[jnp.ndarray] = None, - expert_bias: Optional[jnp.ndarray] = None, *, num_experts: int, - num_experts_per_tok: int, activation_type: str = "silu", - 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, apply_topk_weights_early: bool = False, quantizer_sets: Tuple[QuantizerSet, QuantizerSet] = ( noop_quantizer_set, @@ -1942,39 +2103,35 @@ def moe( mesh_resource: Optional[MeshResource] = None, quant_before_fsdp_ag: bool = False, 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"), dtype: jnp.dtype = jnp.float32, recv_capacity_per_rank: Optional[int] = None, + collect_expert_counts: bool = True, wi_0_checkpoint_name: Optional[str] = None, wi_1_checkpoint_name: Optional[str] = None, wo_checkpoint_name: Optional[str] = None, dispatch_checkpoint_name: Optional[str] = None, use_cudnn_fusion: bool = True, cudnn_native_weight_layout: Optional[bool] = None, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray, jnp.ndarray]: """Run a full MoE block under a single fused custom_vjp on the TE EP path. - Returns ``(output, aux_loss, total_recv_tokens)``. ``aux_loss`` is ``None`` - when ``aux_loss_coeff == 0``, else a 0-d scalar. ``total_recv_tokens`` is a - non-differentiable pre-drop recv-slot total (grad ``None``); see - ``ep_dispatch`` for using it to detect overflow. + Returns ``(output, aux_loss, total_recv_tokens, expert_counts)``. ``aux_loss`` + is ``None`` when routing is caller-supplied or ``aux_loss_coeff == 0``, else + a 0-d scalar. ``total_recv_tokens`` is a non-differentiable pre-drop + recv-slot total; see ``ep_dispatch`` for using it to detect overflow. + ``expert_counts`` is a non-differentiable int32 array of shape + ``[num_experts]`` containing global top-k assignment counts (or zeros when + ``collect_expert_counts=False``). Parameters ---------- - expert_bias : Optional[jnp.ndarray] - ``[num_experts]`` learnable router bias added before the top-k - when ``score_function='sigmoid'`` or ``'sqrtsoftplus'``. 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. + router_info : Union[RouterComputationInfo, RoutingMapInfo] + :class:`RouterComputationInfo` makes TE compute the gate and top-k + from the given router parameters. :class:`RoutingMapInfo` supplies + precomputed expert indices and routing weights (e.g. hash routing) + and skips TE's router. quantizer_sets : Tuple[QuantizerSet, QuantizerSet] Independent FC1 and FC2 quantizer sets describing the global logical operation. Token quantizers have ``dp_size * num_experts`` groups and @@ -1987,6 +2144,9 @@ def moe( (default) reserves the dropless aligned worst case. The value must match the capacity used by ``ep_bootstrap``. Overflow is reported through ``total_recv_tokens`` when bootstrap used ``drop_on_overflow=True``. + collect_expert_counts : bool + Whether to compute global per-expert assignment counts. Disable this + when the returned counts are unused to avoid an extra all-reduce. wi_0_checkpoint_name : Optional[str] JAX rematerialization checkpoint name for the gate projection output. With cuDNN fusion, also names the shared fused outputs, including its quantized output. @@ -2038,9 +2198,9 @@ def moe( describes parameter storage independently of the selected execution path. MeshResource fields name physical mesh axes, not Flax logical axes. - ``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes`` and - ``wo_kernel_axes`` remain logical-axis tuples resolved through the active - Flax rules. The selected MeshResource is also used by internal TE sharding + ``input_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` and + ``RouterComputationInfo.gate_kernel_axes`` remain logical-axis tuples + resolved through the active Flax rules. The selected MeshResource is also used by internal TE sharding in both forward and backward, independently of an enclosing global context. See module docstring for the rest of the parameter semantics and the @@ -2052,46 +2212,47 @@ def moe( mesh_resource, quant_before_fsdp_ag ) ep_axis, data_parallelism_axes = _moe_mesh_axes(mesh_resource) - score_function = _validate_score_function(score_function) + + if isinstance(router_info, RouterComputationInfo): + router_info = router_info.replace( + score_function=_validate_score_function(router_info.score_function), + aux_loss_coeff=float(router_info.aux_loss_coeff), + ) + has_aux_loss = router_info.aux_loss_coeff > 0.0 + elif isinstance(router_info, RoutingMapInfo): + routing_indices = router_info.routing_indices + routing_weights = router_info.routing_weights + if routing_indices.ndim != 3 or routing_indices.shape[:2] != x.shape[:2]: + raise ValueError( + "routing_indices must have shape [batch, sequence, top_k] matching x, " + f"got x={x.shape}, routing_indices={routing_indices.shape}." + ) + if routing_weights.shape != routing_indices.shape: + raise ValueError( + "routing_weights must have the same shape as routing_indices, got " + f"{routing_weights.shape} and {routing_indices.shape}." + ) + if routing_indices.shape[-1] <= 0: + raise ValueError("routing_indices must select at least one expert per token.") + router_info = router_info.replace(routing_indices=routing_indices.astype(jnp.int32)) + has_aux_loss = False + else: + raise TypeError( + "router_info must be a RouterComputationInfo or RoutingMapInfo, got " + f"{type(router_info).__name__}." + ) with global_shard_guard(mesh_resource): # 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.") - for axis in (ep_axis, *data_parallelism_axes): - if axis not in mesh.shape: - raise ValueError(f"TE MoE resource axis {axis!r} is not in the active 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)) + x = _prepare_moe_input(x, ep_axis, data_parallelism_axes, caller="moe") cudnn_native_weight_layout = _resolve_cudnn_native_weight_layout( wi, x.shape[-1], cudnn_native_weight_layout ) - # 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) - use_cudnn_jax_fusion = False if use_cudnn_fusion: rejection_reasons = _cudnn_jax_fusion_rejection_reasons( @@ -2106,33 +2267,25 @@ def moe( ) use_cudnn_jax_fusion = _select_cudnn_jax_fusion(rejection_reasons) - output, aux_loss, total_recv_tokens = _moe( + output, aux_loss, total_recv_tokens, expert_counts = _moe( x, - gate_kernel, + router_info, wi, wo, wi_0_bias, wi_1_bias, wo_bias, - expert_bias_arg, quantizer_sets, num_experts, - num_experts_per_tok, activation_type, - score_function, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - float(aux_loss_coeff), mesh_resource, input_axes, - gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, dtype, apply_topk_weights_early, recv_capacity_per_rank, + collect_expert_counts, use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, @@ -2142,9 +2295,9 @@ def moe( use_cudnn_fusion, cudnn_native_weight_layout, ) - if aux_loss_coeff <= 0.0: + if not has_aux_loss: aux_loss = None assert ( output.dtype == x.dtype ), f"moe() output dtype {output.dtype} != input dtype {x.dtype}" - return output, aux_loss, total_recv_tokens + return output, aux_loss, total_recv_tokens, expert_counts From dd628d3a11b7b52908ea5b7bba7d38d739df6f55 Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Tue, 27 Oct 2026 10:16:05 -0700 Subject: [PATCH 08/32] [Common] Don't call CUDA from static destructors at exit (#3733) * [Common] Keep cuda_driver::call's symbol cache alive through exit rtc::KernelManager's static destructor unloads its CUDA modules through cuda_driver::call. Each instantiation of call caches driver entry points in function-local statics, and those caches are usually first constructed after the KernelManager, so they are destroyed before it. ~Kernel then reads a destroyed unordered_map: with GCC 16's libstdc++ the process segfaults at exit; with GCC 13's it writes into freed memory instead. Allocate the cache and its mutex with new and never destroy them, and add a test that the process exits cleanly after running NVRTC kernels. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Connor Baker * [Common] Never destroy rtc::KernelManager or the fused-attention graph caches Both were function-local statics whose destructors call into CUDA after main returns: ~Kernel unloads modules through the driver, and destroying a cached cuDNN plan that a CUDA graph retains releases a CUDA user object. The CUDA Programming Guide makes any use of CUDA during program termination undefined behavior, and both crashed at exit depending on static destruction order. Allocate both with new and never delete them, as TVMFFIConfigCache::create() and the fused-attention debug registry already do. Both caches only ever insert, so nothing changes while the process runs. Replace tests/pytorch/test_process_exit.py, whose LayerNorm only runs NVRTC kernels while the NVTE_NORM_FWD_USE_CUDNN and NVTE_NORM_BWD_USE_CUDNN switches are off, with MALLOC_PERTURB_=165 for the C++ util tests. UtilTest.NVRTC drives rtc::KernelManager directly, and ctest checks each test's exit status. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Connor Baker --------- Signed-off-by: Connor Baker Co-authored-by: Claude Opus 5.5 --- tests/cpp/util/CMakeLists.txt | 4 +++- transformer_engine/common/fused_attn/graph_cache.h | 4 +++- transformer_engine/common/util/cuda_driver.h | 6 ++++-- transformer_engine/common/util/rtc.cpp | 4 +++- 4 files changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 1dfd2fed4e4..b80a588fe35 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -21,4 +21,6 @@ target_link_libraries(test_util PUBLIC target_compile_options(test_util PRIVATE -O2 -fopenmp) include(GoogleTest) -gtest_discover_tests(test_util DISCOVERY_TIMEOUT 600) +# MALLOC_PERTURB_ makes glibc overwrite freed memory, so a destructor that reads a destroyed object +# at exit crashes the test process, which ctest reports as a failure, instead of passing silently. +gtest_discover_tests(test_util DISCOVERY_TIMEOUT 600 PROPERTIES ENVIRONMENT "MALLOC_PERTURB_=165") diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index b3020f25acf..d52a0daf70b 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -130,7 +130,9 @@ std::shared_ptr> cache_graph(GraphCache auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - static GraphCache cache; + // Never destroyed: destroying cached cuDNN plans calls into CUDA, which must not be used during + // program termination. + static auto &cache = *new GraphCache; cfg.check_derived(); return detail::cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, [&] { return kCreateGraphFn(cfg); }); diff --git a/transformer_engine/common/util/cuda_driver.h b/transformer_engine/common/util/cuda_driver.h index 16242347f12..3b85bd7cce4 100644 --- a/transformer_engine/common/util/cuda_driver.h +++ b/transformer_engine/common/util/cuda_driver.h @@ -40,8 +40,10 @@ template inline CUresult call(const char *symbol, ArgTs... args) { using FuncT = CUresult(ArgTs...); - static std::unordered_map symbol_cache; - static std::mutex cache_mutex; + // Never destroyed, so a lookup never reads a destroyed map, whatever order statics are + // constructed and destroyed in. + static auto &symbol_cache = *new std::unordered_map(); + static auto &cache_mutex = *new std::mutex(); FuncT *func; { diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 2ffbbf72334..f591620124c 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -153,7 +153,9 @@ int Kernel::occupancy_max_active_blocks_per_sm(int device_id, int block_size, KernelManager& KernelManager::instance() { NVTE_CHECK(is_enabled(), "NVRTC support is not enabled"); - static KernelManager instance_; + // Never destroyed: ~Kernel unloads modules through the CUDA driver, and CUDA must not be used + // during program termination. + static KernelManager& instance_ = *new KernelManager; return instance_; } From fa4b30f09b95031931f3f9378e7d30c2071cb2bb Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 27 Oct 2026 11:37:32 -0700 Subject: [PATCH 09/32] Fix DistributedWeight handling under saved-tensor hooks (#3730) * Fix DistributedWeight handling under saved-tensor hooks Signed-off-by: Kirthi Shankar Sivamani * Reduce CPU overhead of the DistributedWeight check in prepare_for_saving prepare_for_saving and restore_from_saved run for every tensor every TE module saves, so the DistributedWeight check sits on a hot path for all users, not just GTP. Test None and torch.Tensor first, read the marker with an inline getattr (DistributedWeight implementers must be Tensor subclasses), and only query torch.compiler.is_compiling() for a flagged weight. On restore, test the common None/Tensor case first. On a CPU microbenchmark this brings the non-distributed cost back to roughly main's, e.g. 4.5 -> 5.2 us (was 10.0 us) per prepare_for_saving for a 32-expert bf16 GroupedLinear. Signed-off-by: Kirthi Shankar Sivamani Co-Authored-By: Claude Opus 5.5 * Add saved-tensor hook tests for DistributedWeight modules Port the Linear and GroupedLinear saved-tensor hook tests from #3545 to the prepare_for_saving fix, and extend them: - Run each module case both under identity saved_tensors_hooks and under te.checkpoint(use_reentrant=False), whose pack hook discards saved tensors and returns recomputed ones in backward. - Cover LayerNormLinear, which also gates its backward on the saved weight being a DistributedWeight. - Add CPU-only checks that prepare_for_saving keeps a DistributedWeight by reference and that backward gets the same object back with no hooks, identity hooks, and torch.utils.checkpoint. Co-authored-by: Xiaowei Ren Signed-off-by: Kirthi Shankar Sivamani Co-Authored-By: Claude Opus 5.5 * Still use the autograd checker for the dist weight Signed-off-by: Kirthi Shankar Sivamani * Error for distweight + torch.compile Signed-off-by: Kirthi Shankar Sivamani * Rename _SavedDistributedWeight to _SavedTensorRef and drop premise test Remove test_saved_tensor_hooks_unpack_subclass_as_plain_tensor, which only checked PyTorch's saved-tensor-hook behavior rather than TE's. Signed-off-by: Kirthi Shankar Sivamani Co-Authored-By: Claude Opus 5.5 --------- Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Claude Opus 5.5 Co-authored-by: Xiaowei Ren --- tests/pytorch/test_distributed_weight.py | 337 +++++++++++++++++- .../pytorch/quantized_tensor.py | 44 ++- 2 files changed, 374 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py index 5c13d01de91..4f029570c65 100644 --- a/tests/pytorch/test_distributed_weight.py +++ b/tests/pytorch/test_distributed_weight.py @@ -2,17 +2,24 @@ # # See LICENSE for license information. -"""Unit tests for the GTP-agnostic DistributedWeight protocol and dispatchers. +"""Tests for the GTP-agnostic DistributedWeight protocol, dispatchers, and module dispatch. -These tests exercise TE's weight-parallelism extension point in isolation, with +The first half exercises 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. + +The second half checks that ``Linear`` / ``LayerNormLinear`` / ``GroupedLinear`` keep the +distributed-weight path in backward when saved-tensor hooks are installed (directly, or via +non-reentrant activation checkpointing); see the notes above those tests. """ import pytest import torch +from torch.autograd.graph import saved_tensors_hooks +from torch.utils.checkpoint import checkpoint as torch_checkpoint +import transformer_engine.pytorch as te from transformer_engine.pytorch.distributed_weight import ( DistributedWeight, is_distributed_weight, @@ -21,6 +28,14 @@ finalize_weight_grads, weight_grad_dtype, ) +from transformer_engine.pytorch.module.grouped_linear import ( + is_module_grouped_tensor_path_supported, +) +from transformer_engine.pytorch.quantized_tensor import ( + prepare_for_saving, + restore_from_func_ctx, + restore_from_saved, +) class FakeDistributedWeight(torch.Tensor): @@ -197,3 +212,321 @@ def test_weight_grad_dtype_noop_on_plain_tensor(): plain_w = torch.nn.Parameter(torch.zeros(2)) plain_w.main_grad = torch.zeros(2, dtype=torch.float32) assert weight_grad_dtype(plain_w, torch.bfloat16) == torch.bfloat16 + + +# --------------------------------------------------------------------------------------------- +# DistributedWeight dispatch under saved-tensor hooks. +# +# The modules pass the DistributedWeight parameter to ``prepare_for_saving`` and gate their +# backward on ``is_distributed_weight(saved_weight)``. Saving it through ``save_for_backward`` +# only round-trips the subclass while no ``torch.autograd.graph.saved_tensors_hooks`` are +# installed: with hooks active, autograd unpacks a saved leaf as a *fresh plain tensor* +# (whatever the unpack hook returns), dropping the Python subclass and its +# ``is_distributed_weight`` marker. Backward then took the plain-parameter branch and failed +# (e.g. "weight was removed while fuse_wgrad_accumulation=True") or used the wrong weights. +# +# Activation offloading (TE's or Megatron-LM's fine-grained offloading) and non-reentrant +# activation checkpointing both install such hooks. ``prepare_for_saving`` therefore also keeps +# a DistributedWeight by reference and restores that object in place of the unpacked saved +# tensor (which is still saved, so autograd's in-place check is kept). These tests pin that: +# with hooks installed, the distributed path must still be taken and produce the same results +# as without hooks. +# +# TE ships no DistributedWeight implementer (they live in the caller, e.g. Megatron-LM's GTP), so +# an in-repo fake is used. It applies distinct scales in the forward / backward materialize hooks, +# so taking the wrong branch fails a specific assertion rather than drifting numerically. +# --------------------------------------------------------------------------------------------- + +FWD_SCALE = 2.0 +BWD_SCALE = 4.0 +IN_F, OUT_F, TOKENS = 256, 512, 128 +DTYPE, DEVICE = torch.bfloat16, "cuda" + +FUSE_WGRAD = pytest.mark.parametrize("fuse_wgrad", [False, True], ids=["unfused", "fused-wgrad"]) +HOOK_MODES = pytest.mark.parametrize("mode", ["saved-tensor-hooks", "checkpoint"]) + + +class _FakeDistWeight(torch.nn.Parameter): + """Fake distributed weight; the group leader holds every member in ``_group``.""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return [(w.detach() * FWD_SCALE).requires_grad_(True) for w in self._group] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return [w.detach() * 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] + if not self.fuse_wgrad_accumulation: + return [g.clone() for g in wl] + for w, g in zip(self._group, wl): + w.main_grad.add_(g.to(w.main_grad.dtype)) # stands in for the reduce-scatter + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] # real grad is in main_grad + + def grad_buffer(self): + return self.wgrad_scratch + + +def _identity_hooks(): + """Hooks that change nothing; their mere presence makes autograd re-wrap saved leaves.""" + return saved_tensors_hooks(lambda t: t, lambda t: t) + + +def _run_hooked(mode, module, *args): + """Run forward and backward with saved-tensor hooks active; return the output.""" + if mode == "saved-tensor-hooks": + with _identity_hooks(): + out = module(*args) + else: + # Non-reentrant TE checkpointing discards saved tensors in forward and hands backward + # the recomputed ones through its unpack hook. + out = te.checkpoint(module, *args, use_reentrant=False) + out.sum().backward() + return out + + +def _install_fakes(module, weight_names, fuse_wgrad): + """Replace the module's weights with fakes sharing one group; return the leader.""" + fakes = [] + for name in weight_names: + fake = _FakeDistWeight(getattr(module, name).data) + fake.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + fake.fuse_wgrad_accumulation = fuse_wgrad + fake.main_grad = torch.zeros((OUT_F, IN_F), dtype=torch.float32, device=DEVICE) + fake.wgrad_scratch = torch.zeros_like(fake.main_grad) + fake.grad_added_to_main_grad = False + setattr(module, name, fake) + fakes.append(fake) + for fake in fakes: + fake._group = fakes + return fakes[0] + + +def _check_dist_path_kept( + module, weight_names, leader, reference, ref_leader, out, ref_out, x, ref_x, fuse_wgrad, mode +): + """Hooked run must have used every DistributedWeight hook and match the unhooked run.""" + assert leader.calls["fwd"] > 0, "materialize_group_for_forward never called" + assert leader.calls["bwd"] > 0, ( + "materialize_group_for_backward never called: backward lost the DistributedWeight " + "after saved_tensors_hooks unpacked the saved weight as a plain tensor" + ) + assert leader.calls["finalize"] > 0, "finalize_group_grads never called" + # Checkpointing runs the forward a second time to recompute the saved tensors. + fwd_calls = ref_leader.calls["fwd"] * (2 if mode == "checkpoint" else 1) + assert leader.calls == {**ref_leader.calls, "fwd": fwd_calls} + # Same kernels on the same data: results are bitwise equal to the unhooked run. + torch.testing.assert_close(out, ref_out, rtol=0, atol=0) + torch.testing.assert_close(x.grad, ref_x.grad, rtol=0, atol=0) + for name in weight_names: + w, ref_w = getattr(module, name), getattr(reference, name) + if fuse_wgrad: + assert w.grad_added_to_main_grad is True + torch.testing.assert_close(w.main_grad, ref_w.main_grad, rtol=0, atol=0) + assert torch.count_nonzero(w.main_grad) > 0 + else: + assert w.grad is not None and ref_w.grad is not None + torch.testing.assert_close(w.grad, ref_w.grad, rtol=0, atol=0) + + +def _skip_without_cuda(): + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + +def test_prepare_for_saving_keeps_distributed_weight_by_reference(): + """A DistributedWeight is still saved, but restored as the original object.""" + w = FakeDistributedWeight() + plain = torch.zeros(2) + tensors, tensor_objects = prepare_for_saving(plain, w, None) + assert tensors[0] is plain and tensors[1] is w and tensors[2] is None + # Under saved-tensor hooks autograd hands back a plain tensor in the weight's slot. + unpacked = [tensors[0], torch.Tensor(w), None] + restored = restore_from_saved(tensor_objects, unpacked) + assert restored[0] is plain and restored[1] is w and restored[2] is None + + +@pytest.mark.parametrize("mode", ["no-hooks", "saved-tensor-hooks", "checkpoint"]) +def test_prepare_for_saving_round_trips_distributed_weight(mode): + """Through autograd, backward gets the DistributedWeight object back under any hooks.""" + seen = {} + + class _SaveWeight(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w): + tensors_to_save, tensor_objects = prepare_for_saving(x, w) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + return x * 2 + + @staticmethod + def backward(ctx, g): + _, seen["w"] = restore_from_func_ctx(ctx) + return g * 2, None + + w = _FakeDistWeight(torch.ones(4)) + x = torch.ones(4, requires_grad=True) + if mode == "no-hooks": + out = _SaveWeight.apply(x, w) + elif mode == "saved-tensor-hooks": + with _identity_hooks(): + out = _SaveWeight.apply(x, w) + else: + out = torch_checkpoint(_SaveWeight.apply, x, w, use_reentrant=False) + out.sum().backward() + assert seen["w"] is w + + +def test_prepare_for_saving_keeps_inplace_check_for_distributed_weight(): + """Restoring the original object must not bypass autograd's in-place modification check.""" + + class _SaveWeight(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w): + tensors_to_save, tensor_objects = prepare_for_saving(x, w) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + return x * 2 + + @staticmethod + def backward(ctx, g): + restore_from_func_ctx(ctx) + return g * 2, None + + w = _FakeDistWeight(torch.ones(4)) + x = torch.ones(4, requires_grad=True) + out = _SaveWeight.apply(x, w) + with torch.no_grad(): + w.add_(1) # e.g. an optimizer step between forward and backward + with pytest.raises(RuntimeError, match="modified by an inplace operation"): + out.sum().backward() + + +@HOOK_MODES +@FUSE_WGRAD +def test_linear_distributed_weight_under_saved_tensor_hooks(fuse_wgrad, mode): + """``Linear`` must keep the distributed path in backward when saved-tensor hooks are on.""" + _skip_without_cuda() + torch.manual_seed(0) + module, reference = ( + te.Linear( + IN_F, + OUT_F, + bias=False, + device=DEVICE, + params_dtype=DTYPE, + fuse_wgrad_accumulation=fuse_wgrad, + ) + for _ in range(2) + ) + reference.load_state_dict(module.state_dict()) + leader = _install_fakes(module, ["weight"], fuse_wgrad) + ref_leader = _install_fakes(reference, ["weight"], fuse_wgrad) + + x = torch.randn(TOKENS, IN_F, dtype=DTYPE, device=DEVICE, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + ref_out = reference(ref_x) + ref_out.sum().backward() + out = _run_hooked(mode, module, x) + + _check_dist_path_kept( + module, ["weight"], leader, reference, ref_leader, out, ref_out, x, ref_x, fuse_wgrad, mode + ) + # The scales prove which weights were used: FWD_SCALE in forward, BWD_SCALE in dgrad. + plain = te.Linear(IN_F, OUT_F, bias=False, device=DEVICE, params_dtype=DTYPE) + plain.load_state_dict({"weight": leader.data}, strict=False) + px = x.detach().clone().requires_grad_(True) + pout = plain(px) + pout.sum().backward() + torch.testing.assert_close(out.float(), FWD_SCALE * pout.float(), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(x.grad.float(), BWD_SCALE * px.grad.float(), rtol=1e-5, atol=1e-5) + + +@HOOK_MODES +@FUSE_WGRAD +def test_layernorm_linear_distributed_weight_under_saved_tensor_hooks(fuse_wgrad, mode): + """``LayerNormLinear`` must keep the distributed path in backward when hooks are on.""" + _skip_without_cuda() + torch.manual_seed(0) + module, reference = ( + te.LayerNormLinear( + IN_F, + OUT_F, + bias=False, + device=DEVICE, + params_dtype=DTYPE, + fuse_wgrad_accumulation=fuse_wgrad, + ) + for _ in range(2) + ) + reference.load_state_dict(module.state_dict()) + leader = _install_fakes(module, ["weight"], fuse_wgrad) + ref_leader = _install_fakes(reference, ["weight"], fuse_wgrad) + + x = torch.randn(TOKENS, IN_F, dtype=DTYPE, device=DEVICE, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + ref_out = reference(ref_x) + ref_out.sum().backward() + out = _run_hooked(mode, module, x) + + _check_dist_path_kept( + module, ["weight"], leader, reference, ref_leader, out, ref_out, x, ref_x, fuse_wgrad, mode + ) + torch.testing.assert_close( + module.layer_norm_weight.grad, reference.layer_norm_weight.grad, rtol=0, atol=0 + ) + + +@HOOK_MODES +@FUSE_WGRAD +@pytest.mark.parametrize("num_gemms", [2, 4]) +@pytest.mark.parametrize( + "use_grouped_tensor", [False, True], ids=["split-quantize", "grouped-tensor"] +) +def test_grouped_linear_distributed_weight_under_saved_tensor_hooks( + num_gemms, fuse_wgrad, use_grouped_tensor, mode +): + """``GroupedLinear`` (both GEMM paths) must keep the distributed path under hooks.""" + _skip_without_cuda() + if use_grouped_tensor and not is_module_grouped_tensor_path_supported(None, DTYPE): + # GroupedLinear would silently fall back to split-quantize, which the other case + # already covers; skip rather than report coverage of the native backward. + pytest.skip("native grouped-tensor path unsupported on this device / cuBLASLt") + torch.manual_seed(0) + names = [f"weight{i}" for i in range(num_gemms)] + module, reference = ( + te.GroupedLinear( + num_gemms, + IN_F, + OUT_F, + bias=False, + device=DEVICE, + params_dtype=DTYPE, + fuse_wgrad_accumulation=fuse_wgrad, + use_grouped_tensor=use_grouped_tensor, + ) + for _ in range(2) + ) + reference.load_state_dict(module.state_dict()) + leader = _install_fakes(module, names, fuse_wgrad) + ref_leader = _install_fakes(reference, names, fuse_wgrad) + + m_splits = torch.full((num_gemms,), TOKENS, dtype=torch.int64, device=DEVICE) + x = torch.randn(num_gemms * TOKENS, IN_F, dtype=DTYPE, device=DEVICE, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + ref_out = reference(ref_x, m_splits) + ref_out.sum().backward() + out = _run_hooked(mode, module, x, m_splits) + + _check_dist_path_kept( + module, names, leader, reference, ref_leader, out, ref_out, x, ref_x, fuse_wgrad, mode + ) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index f6f2794bdf9..fc6b59a32f1 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -288,21 +288,46 @@ class _SavedQuantizedTensor(NamedTuple): metadata: Dict[str, Any] +class _SavedTensorRef(NamedTuple): + """A DistributedWeight's original object, restored in place of its saved tensor.""" + + weight: torch.Tensor + + def prepare_for_saving( *tensors: Union[torch.Tensor, QuantizedTensorStorage], ) -> Tuple[ list[Optional[Union[torch.Tensor, torch.nn.Parameter]]], - list[Optional[Union[QuantizedTensorStorage, _SavedQuantizedTensor]]], + list[Optional[Union[QuantizedTensorStorage, _SavedQuantizedTensor, _SavedTensorRef]]], ]: """Prepare tensors for saving. Needed because save_for_backward accepts only torch.Tensor/torch.nn.Parameter types, while we want to be able to save - the internal TensorStorage types too.""" + the internal TensorStorage types too. + + A DistributedWeight is still saved through autograd (which keeps autograd's + in-place modification check), but its original object is also kept by + reference and restored in place of the saved tensor: saved-tensor hooks + (e.g. activation offloading) unpack saved tensors as plain tensors, which + drops the subclass and its distributed-weight hooks.""" tensor_list, tensor_objects_list = [], [] for tensor in tensors: - if tensor is None or isinstance(tensor, torch.Tensor): - tensor_list.append(tensor) + if tensor is None: + tensor_list.append(None) tensor_objects_list.append(None) + elif isinstance(tensor, torch.Tensor): + # Inline is_distributed_weight(): this runs for every saved tensor. + if getattr(tensor, "is_distributed_weight", False): + if torch.compiler.is_compiling(): + raise RuntimeError( + "Saving a DistributedWeight for backward is not supported under " + "torch.compile." + ) + tensor_list.append(tensor) + tensor_objects_list.append(_SavedTensorRef(tensor)) + else: + tensor_list.append(tensor) + tensor_objects_list.append(None) elif torch.compiler.is_compiling(): inner_names, metadata = tensor.__tensor_flatten__() tensor_list.extend(getattr(tensor, name) for name in inner_names) @@ -316,7 +341,11 @@ def prepare_for_saving( def restore_from_saved( - tensors: list[Optional[Union[torch.Tensor, QuantizedTensorStorage, _SavedQuantizedTensor]]], + tensors: list[ + Optional[ + Union[torch.Tensor, QuantizedTensorStorage, _SavedQuantizedTensor, _SavedTensorRef] + ] + ], saved_tensors: list[Optional[Union[torch.Tensor, torch.nn.Parameter]]], return_saved_tensors: bool = False, ) -> ( @@ -334,6 +363,11 @@ def restore_from_saved( if tensor is None or isinstance(tensor, torch.Tensor): tensor_objects.append(saved_tensors[0]) saved_tensors = saved_tensors[1:] + elif isinstance(tensor, _SavedTensorRef): + # The saved tensor was unpacked (running autograd's version check), but + # may have lost its subclass under hooks; use the original object. + tensor_objects.append(tensor.weight) + saved_tensors = saved_tensors[1:] elif isinstance(tensor, _SavedQuantizedTensor): count = len(tensor.inner_names) inner = dict(zip(tensor.inner_names, saved_tensors[:count])) From 96c0fb0e24915683acc24ca2cf31cac16c798830 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 27 Oct 2026 11:37:52 -0700 Subject: [PATCH 10/32] Remove device side validation of QB bounds (#3734) * Remove QB bounds device side validation Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_fused_router.py | 216 +++--------------- .../fused_topk_with_score_function.cu | 71 ++---- .../include/transformer_engine/fused_router.h | 34 +-- transformer_engine/pytorch/csrc/extensions.h | 3 +- .../pytorch/csrc/extensions/pybind.cpp | 2 +- .../pytorch/csrc/extensions/router.cpp | 43 ++-- transformer_engine/pytorch/router.py | 84 ++----- 7 files changed, 104 insertions(+), 349 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index c6c82fa3d78..ac2f46826df 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -661,109 +661,55 @@ def test_qb_topk_uses_logits_device(histogram_mode): @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -@pytest.mark.parametrize("invalid_bounds", ["equal", "reversed", "nonfinite"]) -def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds): - logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) - expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - bounds = { - "equal": [1.0, 1.0], - "reversed": [1.0, -1.0], - "nonfinite": [float("nan"), 1.0], - }[invalid_bounds] - bin_bounds = torch.tensor(bounds, device="cuda", dtype=torch.float32) - with pytest.raises(ValueError, match="finite with lower < upper"): - fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - expert_bias, - qb_histogram=histogram, - qb_bin_bounds=bin_bounds, - qb_histogram_mode=histogram_mode, - ) - +@pytest.mark.parametrize("use_dense_indices", [False, True]) +@pytest.mark.parametrize( + "invalid_bounds", + [[1.0, 1.0], [1.0, -1.0], [float("nan"), 1.0], [float("-inf"), float("inf")]], + ids=["equal", "reversed", "nan", "infinite"], +) +def test_qb_topk_invalid_bin_bounds_stay_in_histogram( + histogram_mode, use_dense_indices, invalid_bounds +): + num_tokens, num_experts, num_bins, topk = 8, 16, 32, 4 + logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32) + expert_bias = torch.zeros(num_experts, device="cuda", dtype=torch.float32) + histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) + bin_bounds = torch.tensor(invalid_bounds, device="cuda", dtype=torch.float32) + topk_indices = ( + torch.empty(num_tokens, topk, device="cuda", dtype=torch.int32) + if use_dense_indices + else None + ) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_revalidates_updated_bin_bounds(histogram_mode): - logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) - expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) - fused_topk_with_score_function( + # bin_bounds values are not validated, so invalid bounds must neither raise nor write outside + # the histogram. + probs, routing_output = fused_topk_with_score_function( logits, - 4, + topk, False, None, None, None, "sigmoid", expert_bias, + topk_indices=topk_indices, qb_histogram=histogram, qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, ) - bin_bounds.fill_(0.0) - with pytest.raises(ValueError, match="finite with lower < upper"): - fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - expert_bias, - qb_histogram=histogram, - qb_bin_bounds=bin_bounds, - qb_histogram_mode=histogram_mode, - ) - - -@pytest.mark.parametrize( - "histogram_mode", - [QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC], -) -@pytest.mark.parametrize("use_dense_indices", [False, True]) -def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(histogram_mode, use_dense_indices): - logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) - expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - topk_indices = ( - torch.empty(8, 4, device="cuda", dtype=torch.int32) if use_dense_indices else None - ) - invalid_bounds = torch.tensor([1.0, 1.0], device="cuda", dtype=torch.float32) - with pytest.raises(RuntimeError, match="finite with lower < upper"): - tex.fused_topk_with_score_function_qb_fwd( - logits, - 4, - None, - expert_bias, - int(RoutingMapFormat.BYTEMAP), - topk_indices, - histogram, - invalid_bounds, - histogram_mode, - ) + torch.cuda.synchronize() - # The validation error must not poison the CUDA context. + # Routing does not read the bounds, and every (token, expert) pair is counted in some bin. valid_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) - tex.fused_topk_with_score_function_qb_fwd( - logits, - 4, - None, - expert_bias, - int(RoutingMapFormat.BYTEMAP), - topk_indices, - histogram, - valid_bounds, - histogram_mode, - ) - torch.cuda.synchronize() + reference = qb_topk_score_function_pytorch(logits, topk, expert_bias, valid_bounds, num_bins) + torch.testing.assert_close(probs, reference["probs"]) + if use_dense_indices: + torch.testing.assert_close( + topk_indices_to_routing_map(routing_output, num_experts), reference["routing_map"] + ) + else: + torch.testing.assert_close(routing_output, reference["routing_map"]) + assert histogram.sum().item() == num_tokens * num_experts @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) @@ -848,7 +794,6 @@ def run_iteration(): qb_histogram_mode=histogram_mode, ) bin_bounds.copy_(next_bounds) - mark_qb_bin_bounds_validated(bin_bounds) return probs, routing_map # Match full-iteration capture: an eager warmup ends with a bounds update, and the same @@ -871,98 +816,11 @@ def run_iteration(): torch.testing.assert_close(bin_bounds, next_bounds) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_eager_call_revalidates_after_captured_bounds_update(histogram_mode): - logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) - expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) +def test_mark_qb_bin_bounds_validated_is_deprecated(): bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) - next_bounds = torch.tensor([-0.75, 0.25], device="cuda", dtype=torch.float32) - - def run_router(): - return fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - expert_bias, - qb_histogram=histogram, - qb_bin_bounds=bin_bounds, - qb_histogram_mode=histogram_mode, - ) - - run_router() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - bin_bounds.copy_(next_bounds) + with pytest.warns(DeprecationWarning, match="mark_qb_bin_bounds_validated is deprecated"): mark_qb_bin_bounds_validated(bin_bounds) - # Replayed device writes do not advance bin_bounds._version. The capture-time marker must not - # let a subsequent eager router call inherit trust in the replay-produced values. - next_bounds.zero_() - graph.replay() - torch.cuda.synchronize() - with pytest.raises(ValueError, match="finite with lower < upper"): - run_router() - - -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_cuda_graph_rejects_unvalidated_bounds(histogram_mode): - logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) - expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) - - with pytest.raises(RuntimeError, match="current version must be validated"): - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - expert_bias, - qb_histogram=histogram, - qb_bin_bounds=bin_bounds, - qb_histogram_mode=histogram_mode, - ) - - -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_cuda_graph_rejects_stale_bounds_validation(histogram_mode): - logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) - expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) - - def run_router(): - return fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - expert_bias, - qb_histogram=histogram, - qb_bin_bounds=bin_bounds, - qb_histogram_mode=histogram_mode, - ) - - run_router() - bin_bounds.zero_() - with pytest.raises(RuntimeError, match="current version must be validated"): - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run_router() - @pytest.mark.parametrize( "histogram_mode", diff --git a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu index 32d525bb027..e9d6626a2b0 100644 --- a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu +++ b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu @@ -10,7 +10,6 @@ #include #include -#include #include #include @@ -33,22 +32,16 @@ struct QBBinParams { CompType scale; }; +// bin_bounds values are not validated, since that would need a host sync. Users of these params +// compute bin indices with __float2int_rd, which saturates instead of invoking undefined behavior, +// and clamp them to [0, num_bins), so invalid bounds give a wrong histogram but never an +// out-of-bounds access. __device__ inline QBBinParams load_qb_bin_params(const CompType *bin_bounds, int num_bins) { const CompType lower = bin_bounds[0]; const CompType upper = bin_bounds[1]; return {lower, static_cast(num_bins) / (upper - lower)}; } -void validate_qb_bin_bounds(const Tensor bin_bounds, cudaStream_t stream) { - float bounds[2]; - NVTE_CHECK_CUDA(cudaMemcpyAsync(bounds, bin_bounds.data.dptr, sizeof(bounds), - cudaMemcpyDeviceToHost, stream)); - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - NVTE_CHECK(std::isfinite(bounds[0]) && std::isfinite(bounds[1]) && bounds[0] < bounds[1], - "QB bin_bounds values must be finite with lower < upper, got [", bounds[0], ", ", - bounds[1], "]"); -} - template __device__ inline CompType extract_qb_cutoff_and_compact(int *topk_indices, CompType *topk_scores, int topk, int lane_id) { @@ -87,8 +80,7 @@ __device__ inline void accumulate_qb_histogram_epilogue(const CompType *raw_scor if constexpr (Mode == QBMode::FusedAtomic) { const QBBinParams bin_params = load_qb_bin_params(bin_bounds, num_bins); for (int expert = lane_id; expert < num_experts; expert += kThreadsPerWarp) { - int bin = static_cast( - floorf((cutoff - raw_scores[expert] - bin_params.lower) * bin_params.scale)); + int bin = __float2int_rd((cutoff - raw_scores[expert] - bin_params.lower) * bin_params.scale); bin = max(0, min(bin, num_bins - 1)); atomicAdd(histogram + static_cast(expert) * num_bins + bin, 1); } @@ -913,7 +905,7 @@ __global__ void qb_histogram_accumulate_kernel(const CompType *raw_scores, const if (token < num_tokens && expert < num_experts) { const CompType required_bias = cutoff[token] - raw_scores[static_cast(token) * num_experts + expert]; - int bin = static_cast(floorf((required_bias - bin_params.lower) * bin_params.scale)); + int bin = __float2int_rd((required_bias - bin_params.lower) * bin_params.scale); bin = max(0, min(bin, num_bins - 1)); atomicAdd(local_histogram + local_expert * num_bins + bin, 1); } @@ -931,7 +923,7 @@ __global__ void qb_histogram_accumulate_kernel(const CompType *raw_scores, const } void qb_histogram_accumulate(const Tensor raw_scores, const Tensor cutoff, const Tensor bin_bounds, - Tensor histogram, bool validate_bin_bounds, cudaStream_t stream) { + Tensor histogram, cudaStream_t stream) { NVTE_CHECK(raw_scores.data.dtype == DType::kFloat32, "QB raw_scores must have FP32 dtype"); NVTE_CHECK(cutoff.data.dtype == DType::kFloat32, "QB cutoff must have FP32 dtype"); NVTE_CHECK(bin_bounds.data.dtype == DType::kFloat32, "QB bin_bounds must have FP32 dtype"); @@ -948,9 +940,6 @@ void qb_histogram_accumulate(const Tensor raw_scores, const Tensor cutoff, const "QB histogram must have shape [num_experts, num_bins]"); const int num_bins = static_cast(histogram.data.shape[1]); NVTE_CHECK(num_bins > 0, "QB num_bins must be positive"); - if (validate_bin_bounds) { - validate_qb_bin_bounds(bin_bounds, stream); - } const int token_partitions = std::min(4, std::max(1, (num_tokens + 1023) / 1024)); const dim3 grid((num_experts + kQBExpertsPerBlock - 1) / kQBExpertsPerBlock, token_partitions); @@ -1144,8 +1133,7 @@ void fused_topk_with_score_function_forward_qb( const Tensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, const Tensor expert_bias, Tensor probs, Tensor routing_map, NVTERoutingMapFormat routing_map_format, Tensor intermediate_output, Tensor cutoff, - Tensor histogram, Tensor bin_bounds, NVTEQBHistogramMode histogram_mode, - bool validate_bin_bounds, cudaStream_t stream) { + Tensor histogram, Tensor bin_bounds, NVTEQBHistogramMode histogram_mode, cudaStream_t stream) { check_routing_map_format(routing_map_format); const int num_bins = check_qb_forward_tensors(logits, num_tokens, num_experts, topk, expert_bias, probs, @@ -1154,9 +1142,6 @@ void fused_topk_with_score_function_forward_qb( expected_routing_map_shape(num_tokens, num_experts, routing_map_format); NVTE_CHECK(routing_map.data.shape == routing_map_shape, "QB routing_map shape does not match routing_map_format"); - if (validate_bin_bounds) { - validate_qb_bin_bounds(bin_bounds, stream); - } #define QB_ROUTER_FORWARD_LAUNCH(RoutingMapFormatVal, QbModeVal) \ TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ @@ -1195,7 +1180,7 @@ void fused_topk_with_score_function_forward_qb_with_indices( const Tensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, const Tensor expert_bias, Tensor probs, Tensor topk_indices, Tensor intermediate_output, Tensor cutoff, Tensor histogram, Tensor bin_bounds, NVTEQBHistogramMode histogram_mode, - bool validate_bin_bounds, cudaStream_t stream) { + cudaStream_t stream) { const int num_bins = check_qb_forward_tensors(logits, num_tokens, num_experts, topk, expert_bias, probs, intermediate_output, cutoff, histogram, bin_bounds); @@ -1203,9 +1188,6 @@ void fused_topk_with_score_function_forward_qb_with_indices( static_cast(topk)}; NVTE_CHECK(topk_indices.data.shape == indices_shape, "QB topk_indices must have shape [num_tokens, topk]"); - if (validate_bin_bounds) { - validate_qb_bin_bounds(bin_bounds, stream); - } #define QB_ROUTER_FORWARD_WITH_INDICES_LAUNCH(DataType, IndexType, QbModeVal) \ fused_topk_with_score_function_forward_with_indices_kernel_launcher scaling_factor, at::Tensor expert_bias, int routing_map_format, std::optional topk_indices, at::Tensor histogram, - at::Tensor bin_bounds, int histogram_mode, - bool bin_bounds_validated = false); + at::Tensor bin_bounds, int histogram_mode); void fused_topk_with_score_function_bwd( at::Tensor routing_map, at::Tensor intermediate_output, at::Tensor grad_probs, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 89e7d92aca9..ebcdd17b07c 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -149,7 +149,7 @@ void init_router_bindings(pybind11::module &m) { m.def("fused_topk_with_score_function_qb_fwd", &fused_topk_with_score_function_qb_fwd, py::arg("logits"), py::arg("topk"), py::arg("scaling_factor"), py::arg("expert_bias"), py::arg("routing_map_format"), py::arg("topk_indices"), py::arg("histogram"), - py::arg("bin_bounds"), py::arg("histogram_mode"), py::arg("bin_bounds_validated") = false, + py::arg("bin_bounds"), py::arg("histogram_mode"), "Kimi K3 QB fused topk with histogram accumulation"); m.def("fused_topk_with_score_function_bwd", &fused_topk_with_score_function_bwd, py::arg("routing_map"), py::arg("intermediate_output"), py::arg("grad_probs"), diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 77ee80c0995..79391b5bf2f 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -159,8 +159,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, std::optional scaling_factor, at::Tensor expert_bias, int routing_map_format, std::optional topk_indices, at::Tensor histogram, - at::Tensor bin_bounds, int histogram_mode, - bool bin_bounds_validated) { + at::Tensor bin_bounds, int histogram_mode) { check_routing_map_format(routing_map_format); TORCH_CHECK(logits.dim() >= 1, "logits must have at least 1 dim"); TORCH_CHECK(logits.is_cuda() && logits.is_contiguous(), @@ -239,37 +238,21 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, auto stream = at::cuda::getCurrentCUDAStream(); if (topk_indices.has_value()) { - if (bin_bounds_validated) { - nvte_fused_topk_with_score_function_forward_qb_with_indices_unchecked( - logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, - scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), - intermediate_output_cu.data(), cutoff_cu.data(), histogram_cu.data(), - bin_bounds_cu.data(), mode, stream); - } else { - nvte_fused_topk_with_score_function_forward_qb_with_indices( - logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, - scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), - intermediate_output_cu.data(), cutoff_cu.data(), histogram_cu.data(), - bin_bounds_cu.data(), mode, stream); - } + nvte_fused_topk_with_score_function_forward_qb_with_indices( + logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, + scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), + intermediate_output_cu.data(), cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), + mode, stream); } else { - if (bin_bounds_validated) { - nvte_fused_topk_with_score_function_forward_qb_v2_unchecked( - logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, - scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), - static_cast(routing_map_format), intermediate_output_cu.data(), - cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), mode, stream); - } else { - nvte_fused_topk_with_score_function_forward_qb_v2( - logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, - scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), - static_cast(routing_map_format), intermediate_output_cu.data(), - cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), mode, stream); - } + nvte_fused_topk_with_score_function_forward_qb_v2( + logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, + scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), + static_cast(routing_map_format), intermediate_output_cu.data(), + cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), mode, stream); } if (mode == NVTE_QB_HISTOGRAM_TWO_KERNEL) { - nvte_qb_histogram_accumulate_unchecked(intermediate_output_cu.data(), cutoff_cu.data(), - bin_bounds_cu.data(), histogram_cu.data(), stream); + nvte_qb_histogram_accumulate(intermediate_output_cu.data(), cutoff_cu.data(), + bin_bounds_cu.data(), histogram_cu.data(), stream); } return std::make_tuple(probs, routing_output, intermediate_output, cutoff, histogram); diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 898310b0d61..508c2cba769 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -14,7 +14,7 @@ global memory. For example, the gradient is required to have the same dtype as the input. """ -import math +import warnings from typing import Optional, Union import torch @@ -39,72 +39,24 @@ "two_kernel": int(QBHistogramMode.TWO_KERNEL), "fused_atomic": int(QBHistogramMode.FUSED_ATOMIC), } -_QB_BOUNDS_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_validated_version" -_QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_capture_validated_version" -def mark_qb_bin_bounds_validated(bin_bounds: torch.Tensor) -> None: - """Mark the current QB bounds version as valid after a trusted device-side update. +def mark_qb_bin_bounds_validated( + bin_bounds: torch.Tensor, # pylint: disable=unused-argument +) -> None: + """Deprecated no-op. - This function validates metadata only; it does not inspect or schedule validation of tensor - values. The caller must guarantee finite FP32 bounds with ``lower < upper``. When called during - CUDA graph capture, call it immediately after dispatching the trusted in-place update on the - same stream. The Python marker itself is not replayed, so every value produced on replay must - satisfy the bounds contract for every possible input and control-flow path. + QB ``bin_bounds`` values are no longer validated, so there is nothing to mark. This function + will be removed in a future release. """ - if not ( - isinstance(bin_bounds, torch.Tensor) - and bin_bounds.is_cuda - and bin_bounds.is_contiguous() - and bin_bounds.dtype == torch.float32 - and bin_bounds.shape == (2,) - ): - raise ValueError("QB bin_bounds must be a contiguous FP32 CUDA tensor with shape [2]") - version = bin_bounds._version - setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, version) - with torch.cuda.device(bin_bounds.device): - capture_validated_version = version if torch.cuda.is_current_stream_capturing() else None - setattr( - bin_bounds, - _QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR, - capture_validated_version, + warnings.warn( + "mark_qb_bin_bounds_validated is deprecated and has no effect because QB bin_bounds " + "values are no longer validated. It will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, ) -def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: - """Validate QB bounds once per PyTorch tensor version.""" - if not ( - isinstance(bin_bounds, torch.Tensor) - and bin_bounds.is_cuda - and bin_bounds.is_contiguous() - and bin_bounds.dtype == torch.float32 - and bin_bounds.shape == (2,) - ): - # The C++ binding owns metadata validation and its detailed error messages. - return False - - version = bin_bounds._version - validated_version = getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) - capture_validated_version = getattr(bin_bounds, _QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR, None) - if validated_version == version and capture_validated_version != version: - return True - with torch.cuda.device(bin_bounds.device): - is_capturing = torch.cuda.is_current_stream_capturing() - if validated_version == version and is_capturing: - return True - if is_capturing: - raise RuntimeError( - "QB bin_bounds current version must be validated before CUDA graph capture" - ) - lower, upper = bin_bounds.detach().cpu().tolist() - if not (math.isfinite(lower) and math.isfinite(upper) and lower < upper): - raise ValueError( - f"QB bin_bounds values must be finite with lower < upper, got [{lower}, {upper}]" - ) - mark_qb_bin_bounds_validated(bin_bounds) - return True - - def _validate_routing_map_format( routing_map_format: Union[str, RoutingMapFormat, int], ) -> int: @@ -223,7 +175,6 @@ def forward( histogram_mode: int, ): # pylint: disable=missing-function-docstring - bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) ( probs, routing_output, @@ -240,7 +191,6 @@ def forward( histogram, bin_bounds, histogram_mode, - bin_bounds_validated, ) if topk_indices is not None: routing_output = topk_indices @@ -321,13 +271,9 @@ def fused_topk_with_score_function( Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place. qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be - finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates. - Before CUDA graph capture, validate the current version eagerly or use - :func:`mark_qb_bin_bounds_validated` after a trusted device-side update. CUDA graph replay - does not execute Python validation or advance PyTorch tensor-version counters. Any bounds - update visible to replay, whether issued outside the graph or captured inside it, is - therefore trusted and must produce finite values with ``lower < upper`` on every replay. - Invalid replay-time bounds may silently produce an incorrect histogram. + finite with ``lower < upper``. The values are only read on the device and are not + validated, so the tensor may be updated in place at any time, including inside a CUDA + graph. Invalid bounds silently produce an incorrect histogram. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. From d7cd3024865f9eb68d425369603c90540d38e067 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Tue, 27 Oct 2026 16:02:13 -0500 Subject: [PATCH 11/32] [PyTorch] FP8BS grouped GEMM test reference fix (#3722) * [PyTorch] Test grouped FP8 block-scaling GEMM against the split-tensor path MXFP8 GEMM on Blackwell accumulates at higher precision than FP8 block-scaling GEMM on Hopper. The grouped FP8 block-scaling GEMM tests compared against an FP32 reference, so they passed with the MXFP8 emulation on Blackwell but failed on Hopper. Compare against the legacy split-tensor grouped GEMM instead, consistent with the existing MXFP8 grouped GEMM tests. Also remove the FP8 block-scaling zero-size-expert and mixed-operand tests, which no other recipe has. Signed-off-by: Alp Dener * [PyTorch] Test grouped FP8 block-scaling GEMM with uneven and empty experts Dropping the varying splits and the zero-size-expert test left gaps: the remaining tests no longer run grouped FP8 block-scaling GEMM with uneven or empty experts when the weights or outputs are per-expert lists, when accumulating, or when K is not a multiple of 128. Parametrize the grouped FP8 block-scaling GEMM test over uniform, varying and zero-expert splits, checked against the split-tensor reference. Signed-off-by: Alp Dener --------- Signed-off-by: Alp Dener --- tests/pytorch/test_grouped_linear.py | 199 ++++++++++++--------------- 1 file changed, 86 insertions(+), 113 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 292188b00f4..0a0c42b03eb 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1614,17 +1614,15 @@ def _fp8bs_grouped_gemm_available() -> bool: return tex.get_cublasLt_version() >= (130600 if cc < (10, 0) else 130300) -def _fp8bs_exact_operand(shape, dtype) -> torch.Tensor: - """Random operand exactly representable in FP8 block scaling with power-of-2 scales.""" +def _tile_scaled_randn(shape, dtype) -> torch.Tensor: + """Random tensor whose 128x128 tiles have different power-of-2 magnitudes.""" + # Block scales then differ across tiles, so a GEMM that reads the wrong block's scale fails. rows, cols = shape - x = torch.randn(shape, device="cuda", dtype=torch.float32).to(torch.float8_e4m3fn).float() - # Scale each 128x128 tile by its own power of 2 so block scales differ across tiles. Every - # quantization block lies inside one tile, so its power-of-2 scale absorbs the factor exactly. - pow2 = torch.tensor([2.0**e for e in range(-3, 4)], dtype=torch.float32, device="cuda") + pow2 = torch.tensor([2.0**e for e in range(-3, 4)], device="cuda") tile_grid = ((rows + 127) // 128, (cols + 127) // 128) factors = pow2[torch.randint(0, pow2.numel(), tile_grid, device="cuda")] factors = factors.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1) - return (x * factors[:rows, :cols]).to(dtype) + return (torch.randn(shape, device="cuda") * factors[:rows, :cols]).to(dtype) def _quantize_fp8bs(tensors, *, rowwise, columnwise, block_scaling_dim, grouped, is_weight): @@ -1646,61 +1644,53 @@ def _quantize_fp8bs(tensors, *, rowwise, columnwise, block_scaling_dim, grouped, return tex.group_quantize(torch.cat(tensors, dim=0), quantizer, len(tensors), first_dims) -def _fp8bs_reference(A, B, out_init, layout, accumulate, dtype): - """Per-expert FP32 reference for TN (y = x W^T), NN (dx = dy W) and NT (dW = dy^T x).""" - refs = [] - for a, b, o in zip(A, B, out_init): - a32, b32 = a.float(), b.float() - if layout == "TN": - ref = b32 @ a32.T - elif layout == "NN": - ref = b32 @ a32 - else: - ref = b32.T @ a32 - if accumulate: - ref = ref + o.float() - refs.append(ref.to(dtype)) - return refs - - def _run_fp8bs_grouped_gemm( z, m_sizes, k, n, layout, case, weight_block_dim, accumulate, dtype ) -> None: if layout == "TN": - A = [_fp8bs_exact_operand((n, k), dtype) for _ in range(z)] # weight - B = [_fp8bs_exact_operand((ms, k), dtype) for ms in m_sizes] # input - out = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] + A = [_tile_scaled_randn((n, k), dtype) for _ in range(z)] # weight + B = [_tile_scaled_randn((ms, k), dtype) for ms in m_sizes] # input + out = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # output elif layout == "NN": - A = [_fp8bs_exact_operand((n, k), dtype) for _ in range(z)] # weight - B = [_fp8bs_exact_operand((ms, n), dtype) for ms in m_sizes] # grad_output - out = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] - else: # NT - A = [_fp8bs_exact_operand((ms, k), dtype) for ms in m_sizes] # input - B = [_fp8bs_exact_operand((ms, n), dtype) for ms in m_sizes] # grad_output - out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] - expected = _fp8bs_reference(A, B, out, layout, accumulate, dtype) + A = [_tile_scaled_randn((n, k), dtype) for _ in range(z)] # weight + B = [_tile_scaled_randn((ms, n), dtype) for ms in m_sizes] # grad_output + out = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # dgrad + else: # layout == "NT" + A = [_tile_scaled_randn((ms, k), dtype) for ms in m_sizes] # input + B = [_tile_scaled_randn((ms, n), dtype) for ms in m_sizes] # grad_output + out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad + + out_ref = [o.clone() for o in out] transa = layout[0] == "T" transb = layout[1] == "T" a_is_weight = layout != "NT" - a_block_dim = weight_block_dim if a_is_weight else 1 - grouped_A = _quantize_fp8bs( - A, + a_kwargs = dict( # Weights always keep rowwise data, as in GroupedLinear (non-grouped 2D quantize also # needs rowwise output allocated). rowwise=transa or a_is_weight, columnwise=not transa, - block_scaling_dim=a_block_dim, - grouped=case != "discrete_in", + block_scaling_dim=weight_block_dim if a_is_weight else 1, is_weight=a_is_weight, ) - grouped_B = _quantize_fp8bs( - B, - rowwise=not transb, - columnwise=transb, - block_scaling_dim=1, - grouped=True, - is_weight=False, + b_kwargs = dict(rowwise=not transb, columnwise=transb, block_scaling_dim=1, is_weight=False) + A_fp8 = _quantize_fp8bs(A, grouped=False, **a_kwargs) + B_fp8 = _quantize_fp8bs(B, grouped=False, **b_kwargs) + grouped_A = A_fp8 if case == "discrete_in" else _quantize_fp8bs(A, grouped=True, **a_kwargs) + grouped_B = _quantize_fp8bs(B, grouped=True, **b_kwargs) + + general_grouped_gemm( + A_fp8, + B_fp8, + out_ref, + [None] * z, + dtype, + m_splits=m_sizes, + grad=layout != "TN", + accumulate=accumulate, + layout=layout, + use_split_accumulator=True, + single_output=False, ) device = A[0].device @@ -1720,85 +1710,55 @@ def _run_fp8bs_grouped_gemm( out if case == "discrete_out" else grouped_out, layout=layout, accumulate=accumulate, - use_split_accumulator=True, ) results = out if case == "discrete_out" else grouped_out.split_into_quantized_tensors() - # Operands are exactly representable, so the only error is FP32 accumulation order plus - # rounding to the output dtype: rtol covers two output ulps at worst (relative ulp is at most - # 2^-7 for bf16 and 2^-10 for fp16); atol covers accumulation-order differences on near-zero - # outputs (K <= 2944). - tols = dict(rtol=1.6e-2, atol=1e-2) if dtype == torch.bfloat16 else dict(rtol=2e-3, atol=1e-2) - for got, ref in zip(results, expected): - torch.testing.assert_close(got, ref, **tols) + tols = dict(rtol=0.125, atol=0.0675) # same as the MXFP8 test + for o, o_ref in zip(results, out_ref): + torch.testing.assert_close(o, o_ref, **tols) @pytest.mark.parametrize( "shape", - [(1, 128, 128, 512), (4, 1024, 256, 384), (2, 512, 2880, 2944)], - ids=lambda s: "x".join(str(v) for v in s), + [ + (1, 128, 128, 512), + (8, 1024, 128, 512), + (16, 4096, 128, 512), + # Grouped FP8 block-scaling quantize of the weights requires n to be a multiple of 128. + (2, 256, 2880, 2944), + ], ) -@pytest.mark.parametrize("split", ["uniform", "varying"]) +@pytest.mark.parametrize("splits", ["uniform", "varying", "zero_experts"]) @pytest.mark.parametrize("weight_block_dim", [1, 2], ids=["w1D", "w2D"]) @pytest.mark.parametrize("accumulate", [False, True]) @pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) @pytest.mark.parametrize("case", ["no_discrete", "discrete_in", "discrete_out"]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_grouped_gemm_grouped_tensor_fp8_block_scaling( - shape, split, weight_block_dim, accumulate, layout, case, dtype + shape, splits, weight_block_dim, accumulate, layout, case, dtype ) -> None: - """Grouped-tensor GEMM with FP8 block scaling against an exact high-precision reference.""" if not _fp8bs_grouped_gemm_available(): pytest.skip("Grouped GEMM with FP8 block scaling is unavailable on this system.") if dtype == torch.bfloat16 and not is_bf16_available(): pytest.skip("bfloat16 is required for grouped GEMM test.") if layout == "NT" and weight_block_dim == 2: pytest.skip("NT (wgrad) has no weight operand.") - torch.manual_seed(0) z, m, k, n = shape - if split == "uniform" or z == 1: + if splits != "uniform" and z == 1: + pytest.skip("Non-uniform splits need more than one expert.") + torch.manual_seed(0) + if splits == "uniform": m_sizes = [m // z] * z - else: - # Multiples of 128 that differ per expert and sum to m. - m_sizes = [128 * (i % 2 + 1) for i in range(z)] - m_sizes[-1] += m - sum(m_sizes) + elif splits == "varying": + m_sizes = [128 * (i % 3 + 1) for i in range(z)] + else: # Empty experts first, in the middle and, for z=16, last. + m_sizes = [128 * (i % 3) for i in range(z)] # Grouped FP8 block-scaling quantize requires every expert's token count to be a multiple of # 128 (its per-expert check is a device assert that release builds strip). - assert all(ms > 0 and ms % 128 == 0 for ms in m_sizes) + assert all(ms % 128 == 0 for ms in m_sizes) _run_fp8bs_grouped_gemm(z, m_sizes, k, n, layout, case, weight_block_dim, accumulate, dtype) -@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) -@pytest.mark.parametrize("case", ["no_discrete", "discrete_out"]) -def test_grouped_gemm_grouped_tensor_fp8_block_scaling_zero_size_expert(layout, case) -> None: - """A zero-token expert must not shift the other experts' data or scale offsets.""" - if not _fp8bs_grouped_gemm_available(): - pytest.skip("Grouped GEMM with FP8 block scaling is unavailable on this system.") - torch.manual_seed(0) - _run_fp8bs_grouped_gemm(3, [256, 0, 128], 256, 384, layout, case, 2, False, torch.bfloat16) - - -def test_grouped_gemm_grouped_tensor_fp8_block_scaling_mixed_operands_rejected() -> None: - """On Blackwell+, mixing FP8 block-scaling and MXFP8 operands must raise, not miscompute.""" - if not _fp8bs_grouped_gemm_available() or torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Requires the Blackwell+ grouped GEMM emulation path.") - z, ms, k, n = 2, 128, 256, 256 - A = [_fp8bs_exact_operand((n, k), torch.bfloat16) for _ in range(z)] - B = [_fp8bs_exact_operand((ms, k), torch.bfloat16) for _ in range(z)] - grouped_A = _quantize_fp8bs( - A, rowwise=True, columnwise=False, block_scaling_dim=2, grouped=True, is_weight=True - ) - mxfp8_quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False - ) - mxfp8_quantizer.optimize_for_gemm = True - first_dims = torch.full((z,), ms, dtype=torch.int64, device="cuda") - grouped_B = tex.group_quantize(torch.cat(B, dim=0), mxfp8_quantizer, z, first_dims) - out = _make_grouped_tensor_from_splits([ms] * z, n, A[0].device, torch.bfloat16) - with pytest.raises(RuntimeError, match="Either all tensors or no tensor"): - general_grouped_gemm_for_grouped_tensor(grouped_A, grouped_B, out, layout="TN") - - @pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) @pytest.mark.parametrize("quant_type", ["mxfp8", "fp8_block_scaling"]) def test_grouped_gemm_grouped_tensor_cuda_graph_changed_splits(quant_type, layout) -> None: @@ -1836,7 +1796,7 @@ def test_grouped_gemm_grouped_tensor_cuda_graph_changed_splits(quant_type, layou transa = layout[0] == "T" transb = layout[1] == "T" # A is the weight (uniform [n, k] per expert) for TN and NN, and the input for NT. B is the - # input for TN and the grad output for NN and NT (see _fp8bs_reference). + # input for TN and the grad output for NN and NT. a_is_weight = layout != "NT" b_last_dim = k if layout == "TN" else n out_last_dim = n if layout == "TN" else k @@ -1849,7 +1809,7 @@ def test_grouped_gemm_grouped_tensor_cuda_graph_changed_splits(quant_type, layou else: static_out = torch.empty(capacity, out_last_dim, dtype=dtype, device=device) - def _make_quantizer(*, rowwise, columnwise, is_weight): + def _make_quantizer(*, rowwise, columnwise, is_weight, optimize_for_gemm=True): if quant_type == "mxfp8": quantizer = MXFP8Quantizer( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=columnwise @@ -1864,13 +1824,16 @@ def _make_quantizer(*, rowwise, columnwise, is_weight): block_scaling_dim=2 if is_weight else 1, ) # GEMM-ready scales, as GroupedLinear requests (FP8 block scaling keeps compact scales). - quantizer.optimize_for_gemm = True + quantizer.optimize_for_gemm = optimize_for_gemm return quantizer - quantizer_a = _make_quantizer( - rowwise=transa or a_is_weight, columnwise=not transa, is_weight=a_is_weight - ) - quantizer_b = _make_quantizer(rowwise=not transb, columnwise=transb, is_weight=False) + a_usage = dict(rowwise=transa or a_is_weight, columnwise=not transa, is_weight=a_is_weight) + b_usage = dict(rowwise=not transb, columnwise=transb, is_weight=False) + quantizer_a = _make_quantizer(**a_usage) + quantizer_b = _make_quantizer(**b_usage) + # Per-expert quantizers for the split-tensor reference GEMM. + ref_quantizer_a = _make_quantizer(**a_usage, optimize_for_gemm=False) + ref_quantizer_b = _make_quantizer(**b_usage, optimize_for_gemm=False) def _run(): """Quantize both operands and run the grouped GEMM without reading splits on the host.""" @@ -1910,8 +1873,8 @@ def _run(): def _prepare(splits): """Update the device-side splits and refill every static buffer in place.""" static_first_dims.copy_(torch.tensor(splits, dtype=torch.int64)) - static_a.copy_(_fp8bs_exact_operand(tuple(static_a.shape), dtype)) - static_b.copy_(_fp8bs_exact_operand(tuple(static_b.shape), dtype)) + static_a.copy_(_tile_scaled_randn(tuple(static_a.shape), dtype)) + static_b.copy_(_tile_scaled_randn(tuple(static_b.shape), dtype)) # Values from an earlier run must not survive in the output. static_out.copy_(torch.randn_like(static_out)) @@ -1927,15 +1890,25 @@ def _check(splits, run_name): else: a_parts = [static_a[bounds[e] : bounds[e + 1]] for e in experts] results = [static_out[e * n : (e + 1) * n] for e in experts] - expected = _fp8bs_reference(a_parts, b_parts, [None] * z, layout, False, dtype) - # Same tolerances and rationale as _run_fp8bs_grouped_gemm. The operands are also exactly - # representable in MXFP8: 1x32 and 32x1 blocks also lie inside a single 128x128 tile. + expected = [torch.empty_like(r) for r in results] + general_grouped_gemm( + [ref_quantizer_a(a) for a in a_parts], + [ref_quantizer_b(b) for b in b_parts], + expected, + [None] * z, + dtype, + m_splits=splits, + grad=layout != "TN", + layout=layout, + use_split_accumulator=True, + single_output=False, + ) for e in experts: torch.testing.assert_close( results[e], expected[e], - rtol=1.6e-2, - atol=1e-2, + rtol=0.125, + atol=0.0675, msg=lambda m, e=e: f"{run_name}, splits {splits}, expert {e}: {m}", ) From f8f34c66db9ad449326975675eeb6b12be5b9318 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Wed, 28 Oct 2026 05:25:35 +0800 Subject: [PATCH 12/32] Abstract CUDA hardcodes into configurable te_device_type / te_platform (#3513) * fix: patch cuda hardcodes for plugin compatibility (complete debug directory coverage) This commit completes the CUDA plugin compatibility patches across the entire codebase: 1. Debug directory (transformer_engine/debug/): - fake_quant.py: Replace is_cuda checks with te_device_type() - log_fp8_tensor_stats.py: Replace torch.device('cuda') and torch.cuda.get_device_capability() - per_tensor_scaling.py: Replace is_cuda checks with te_device_type() - utils/stats_buffer.py: Replace .cuda() calls with .to(te_device_type()) 2. Additional missing files from cuda_patch_rc2.14: - permutation.py: Replace all is_cuda checks with device.type comparisons - tensor/float8_tensor.py: Replace is_cuda checks with device.type comparisons - tensor/float8_blockwise_tensor.py: Replace is_cuda checks with device.type comparisons - tensor/mxfp8_tensor.py: Replace is_cuda checks with device.type comparisons 3. Core infrastructure: - transformer_engine/__init__.py: Export te_platform and te_device_type from pytorch module - transformer_engine/pytorch/__init__.py: Define te_platform() and te_device_type() functions All runtime CUDA hardcodes have been replaced with plugin-compatible abstractions. Total: 64 files modified (+405/-289 lines) Signed-off-by: Xianduo Li * fix: remove explicit import to avoid black parsing error The te_platform and te_device_type functions are dynamically assigned to the transformer_engine module namespace via: transformer_engine.te_platform = te_platform transformer_engine.te_platform = te_platform in transformer_engine/pytorch/__init__.py, so explicit imports in transformer_engine/__init__.py cause black to fail parsing due to circular import detection. Signed-off-by: Xianduo Li * fix: correct malformed imports that prevent module loading Move te_platform and te_device_type imports outside of parenthesized import lists to fix SyntaxError. The imports were incorrectly placed inside multi-line import statements, preventing the modules from loading. Fixed files: - transformer_engine/debug/features/log_fp8_tensor_stats.py - transformer_engine/pytorch/attention/dot_product_attention/backends.py - transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py - transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py - transformer_engine/pytorch/attention/inference.py - transformer_engine/pytorch/graph.py - transformer_engine/pytorch/quantization.py - transformer_engine/pytorch/quantized_tensor.py - transformer_engine/pytorch/transformer.py Signed-off-by: Xianduo Li * fix: add missing te_device_type import to files using it Several files imported only te_platform but also called te_device_type(), causing NameError at runtime. Add te_device_type to imports. Fixed files: - transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py - transformer_engine/pytorch/attention/rope.py - transformer_engine/pytorch/cpp_extensions/gemm.py - transformer_engine/pytorch/optimizers/fused_adam.py - transformer_engine/pytorch/utils.py Signed-off-by: Xianduo Li * fix: correct malformed imports in additional files Move te_platform and te_device_type imports outside of parenthesized import lists to fix SyntaxError across 20 more files. Fixed files: - transformer_engine/pytorch/attention/dot_product_attention/utils.py - transformer_engine/pytorch/attention/multi_head_attention.py - transformer_engine/pytorch/cpu_offload.py - transformer_engine/pytorch/module/base.py - transformer_engine/pytorch/module/grouped_linear.py - transformer_engine/pytorch/module/layernorm_linear.py - transformer_engine/pytorch/module/layernorm_mlp.py - transformer_engine/pytorch/module/linear.py - transformer_engine/pytorch/ops/basic/basic_linear.py - transformer_engine/pytorch/ops/basic/grouped_linear.py - transformer_engine/pytorch/ops/basic/l2normalization.py - transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py - transformer_engine/pytorch/ops/fused/grouped_mlp.py - transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py - transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py - transformer_engine/pytorch/ops/op.py - transformer_engine/pytorch/tensor/nvfp4_tensor.py - transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py - transformer_engine/pytorch/triton/mhc.py - transformer_engine/pytorch/triton/permutation.py Signed-off-by: Xianduo Li * fix: use routing_map device instead of hardcoded device type In make_row_id_map, use routing_map.device for tensor allocation instead of te_device_type() to respect the input tensor's device. Signed-off-by: Xianduo Li * revert: remove te_platform() replacements, keep te_device_type() Revert all te_platform() method calls back to torch.cuda.* calls. Keep te_device_type() string replacements for device type parameters. Preserve te_platform definition in transformer_engine/pytorch/__init__.py. Changes: - Reverted 146 te_platform() calls back to torch.cuda.* - Removed 'from transformer_engine import te_platform' imports - Kept all te_device_type() replacements intact - Kept te_platform definition in __init__.py for future use Signed-off-by: Xianduo Li * fix: remove remaining te_platform() calls in utils, fused_adam, and distributed Replace remaining te_platform() calls with torch.cuda equivalents: - torch.cuda.nvtx.range() / range_push() / range_pop() for NVTX profiling - torch.cuda.amp.autocast for autocast context manager - Update docstring reference in fused_adam.py All te_platform() usage has been removed except for the definition in transformer_engine/pytorch/__init__.py. Signed-off-by: Xianduo Li * keep sync with torch.cuda api usage Signed-off-by: Xianduo Li * format Signed-off-by: Xianduo Li * Fix missing te_device_type imports and import ordering Add the missing `te_device_type` import to eight modules that called it without importing it, which would raise NameError at runtime: jit, attention/inference, attention/dot_product_attention/{backends, context_parallel,utils}, ops/fused/{grouped_mlp, userbuffers_backward_linear}, and triton/mhc. Move the `te_device_type` import ahead of relative imports where pylint classifies it as a first-party import that must precede local ones. In jit and distributed, place it so it does not split the torch import group. Order `transformer_engine_torch` before `transformer_engine` in ops/_common. Add docstrings to `te_device_type` and `te_platform`. Signed-off-by: Xianduo Li * Query autocast state and dtype on the same device The device-type patch had changed the autocast dtype lookups to the configured device while leaving the enabled checks on CUDA. With a non-CUDA device type, vendor autocast enabled, and CUDA autocast disabled, these sites took the non-autocast branch and kept the input dtype instead of the requested autocast dtype. Route the operation sites through `maybe_autocast_dtype`, which already queries one device for both calls and carries the torch version guard. This covers the activation, SwiGLU, grouped linear, basic linear, and fused linear forward operations. Add `torch_is_autocast_gpu_enabled` next to `torch_get_autocast_gpu_dtype` for the two sites outside the operation layer that had the same mismatch: `set_activation_dtype`, which sets the activation dtype for every module, and the TransformerLayer forward cast. Signed-off-by: Xianduo Li * Gate per-document THD partitioning on CUDA `tex.thd_get_partitioned_indices` allocates a CUDA output tensor and launches a CUDA kernel on the current CUDA stream, so it runs only on CUDA. The device-type patch had changed its dispatch condition to the configured device type, which sent packed sequences into that kernel on a non-CUDA device instead of the device-capable fallback, failing the `is_cuda` assertion before the fallback could run. Restore the CUDA condition in `build_rank_indices` and in `_get_thd_partition_cu_seqlens`, which had the same problem when normalizing the boundaries that feed the kernel, and record why both are gated on CUDA rather than on the configured device. Keep the configured device in the NO_LOAD_BALANCE branch, which builds indices with plain tensor operations and runs no CUDA kernel. Restore the assertion message, which named the device type that triggers it. Signed-off-by: Xianduo Li * Drop incidental changes from the device-type patch Narrow the diff against main to the device-type substitutions themselves: - context_parallel.py: remove the comments on the THD partitioning device gates. - cpp_extensions/gemm.py: restore the single-line Float8BlockwiseQTensorStorage import, which did not need wrapping. - graph.py: remove an added blank line, leaving the file unchanged from main. Signed-off-by: Xianduo Li * Cover device-type checks missed in recently merged code Two paths reached from main used CUDA device checks where the work below them is not CUDA-specific: - tensor/nvfp4_tensor.py: `NVFP4Tensor` kept CUDA checks in its `record_stream` dispatch and in `copy_`, while the sibling `Float8Tensor`, `MXFP8Tensor`, and `Float8BlockwiseQTensor` classes use the configured device type for the same logic. On another device type `copy_` would move the operand to the destination device and `record_stream` would skip stream tracking. - cross_entropy.py: the input validation rejected non-CUDA tensors even though the Triton backend derives every allocation from the input device and contains nothing CUDA-specific. Leave the checks that guard compiled CUDA kernels, the `device_types` dispatch keys on custom operations, and the `is_cuda` properties that implement the PyTorch tensor contract. Signed-off-by: Xianduo Li * Apply plugin device patches after the _nv alias is registered The top-level hook in transformer_engine/__init__.py imported .patches before load_framework_extension() registered transformer_engine_torch_nv. A plugin that imports that alias at package level - a case the loader explicitly supports - failed there, and the enclosing except only warned, so apply_patches() never ran. The loader later imported the plugin successfully but never retried the patches, leaving device patches unapplied while vendor calls still went through CUDA. Move the call into load_framework_extension(), after the alias is registered and before the PyTorch submodules capture references to torch.cuda entry points. Probe the submodule with find_spec so a plugin that needs no patches may omit it, without masking real import failures inside an existing patches module. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Accept the configured device type in the shared forward input check The modules whose `device` argument now defaults to `te_device_type()` allocate their parameters on the configured device, but the `prepare_forward()` input check still rejected anything whose `is_cuda` was false. That property is a tensor attribute, so patching the `torch.cuda` namespace does not change it, and forward raised before a plugin's native operation could run. Compare `device.type` against `te_device_type()` instead, and report the configured device type in the error message. One check covers Linear, LayerNormLinear, GroupedLinear, and LayerNormMLP, which share this base method. Behavior on CUDA is unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Restore the original spacing in the top-level package Removing the plugin hook left an extra blank line behind. Drop it so the file matches main and no longer appears in the diff. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Skip patch discovery for single-file plugins find_spec() raises ModuleNotFoundError instead of returning None when the parent module is not a package, so a single-file plugin that provides load_plugins() failed the optional-patches probe. The exception handler then disabled the plugin and restored the original extension. Check for __path__ before searching for the submodule. Genuine import failures inside an existing patches module still propagate. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Use the configured device type in GroupedTensor record_stream The merged upstream GroupedTensor stream-tracking handler checked `tensor.is_cuda`, which pins the path to CUDA. Match the device abstraction used by the other quantized-tensor `record_stream` handlers so plugin device types propagate allocator stream tracking. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Accept the configured device type in the attention input checks The attention entry points asserted `is_cuda` on queries, keys, and values, so a non-CUDA configured device failed validation before any attention implementation could run, even though the surrounding state was already created on the configured device. Check the configured device type instead, in DotProductAttention and in the matching FlashAttention and FusedAttention assertions. Apply the same change to the per-document THD partitioning check in context parallelism, and drop the redundant CUDA half of the mixed-THD sequence_ids check, where the comparison against the Q/K/V device already enforces the documented requirement. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Keep the paged KV cache and EP buffer on CUDA The paged KV cache page table and the EpBuffer default device fed tensors into torch.library custom ops registered with device_types="cuda". Allocating them on a configured non-CUDA device left no matching registered implementation, so the dispatcher failed the call before the wrapped tex function could run. Restore the CUDA allocations so these tensors match their op registrations, and drop the now-unused imports. Supporting another device on these paths requires changing the op registrations, which depends on how the device is registered with PyTorch and on when the device type is configured relative to module import. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Apply plugin device patches after the plugin finishes loading apply_patches() rewrites process-global state in place - torch.cuda entry points, TE_DEVICE_TYPE and TE_PLATFORM - with no way to undo it. Running it before load_plugins() meant a load_plugins() failure left those patches active while the except block restored only sys.modules, so import continued with the CUDA extension under a vendor device type and failed later, far from the cause. Run load_plugins() first so everything that can fail is rolled back completely. The patches still precede the PyTorch submodules that capture torch.cuda entry points, and still follow the _nv alias registration that a plugin package may need to import itself. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li * Fix regex in test_cross_entropy Signed-off-by: Przemyslaw Tredak --------- Signed-off-by: Xianduo Li Signed-off-by: Przemyslaw Tredak Co-authored-by: Claude Opus 5.5 Co-authored-by: Przemyslaw Tredak --- tests/pytorch/test_cross_entropy.py | 2 +- transformer_engine/common/__init__.py | 10 ++++++ .../debug/features/fake_quant.py | 5 ++- .../debug/features/log_fp8_tensor_stats.py | 7 +++- .../debug/features/per_tensor_scaling.py | 5 ++- .../debug/features/utils/stats_buffer.py | 5 +-- transformer_engine/pytorch/__init__.py | 21 +++++++++++ .../dot_product_attention/backends.py | 18 +++++----- .../dot_product_attention/context_parallel.py | 11 +++--- .../dot_product_attention.py | 15 ++++---- .../dot_product_attention/flex_attention.py | 11 +++--- .../dot_product_attention/softmax.py | 3 +- .../attention/dot_product_attention/utils.py | 25 +++++++------ .../pytorch/attention/multi_head_attention.py | 5 +-- transformer_engine/pytorch/attention/rope.py | 3 +- .../pytorch/cpp_extensions/gemm.py | 4 ++- transformer_engine/pytorch/cross_entropy.py | 7 ++-- .../custom_recipes/quantizer_factories.py | 3 +- .../custom_recipes/quantizer_factory_zoo.py | 3 +- .../reference_current_scaling.py | 3 +- transformer_engine/pytorch/distributed.py | 19 +++++----- .../pytorch/dynamo/tensor_spec.py | 5 +-- transformer_engine/pytorch/jit.py | 21 +++++++---- transformer_engine/pytorch/module/base.py | 14 ++++---- .../pytorch/module/grouped_linear.py | 7 ++-- .../pytorch/module/layernorm_linear.py | 5 +-- .../pytorch/module/layernorm_mlp.py | 9 ++--- transformer_engine/pytorch/module/linear.py | 7 ++-- transformer_engine/pytorch/ops/_common.py | 3 +- .../pytorch/ops/basic/activation.py | 16 +++------ .../pytorch/ops/basic/basic_linear.py | 7 ++-- transformer_engine/pytorch/ops/basic/bias.py | 3 +- .../pytorch/ops/basic/grouped_linear.py | 11 ++---- .../pytorch/ops/basic/swiglu.py | 22 +++--------- .../ops/fused/forward_linear_bias_add.py | 6 ++-- .../ops/fused/forward_linear_scale_add.py | 6 ++-- .../pytorch/ops/fused/grouped_mlp.py | 6 ++-- .../ops/fused/userbuffers_backward_linear.py | 3 +- .../ops/fused/userbuffers_forward_linear.py | 10 +++--- .../pytorch/optimizers/fused_adam.py | 3 +- transformer_engine/pytorch/permutation.py | 35 ++++++++++++------- transformer_engine/pytorch/quantization.py | 16 +++++---- .../pytorch/quantized_tensor.py | 6 ++-- .../pytorch/tensor/float8_blockwise_tensor.py | 6 ++-- .../pytorch/tensor/float8_tensor.py | 3 +- .../pytorch/tensor/grouped_tensor.py | 5 +-- .../pytorch/tensor/identity_tensor.py | 3 +- .../pytorch/tensor/mxfp8_tensor.py | 3 +- .../pytorch/tensor/nvfp4_tensor.py | 5 +-- .../tensor/storage/mxfp8_tensor_storage.py | 5 +-- .../tensor/storage/nvfp4_tensor_storage.py | 5 +-- transformer_engine/pytorch/transformer.py | 9 +++-- transformer_engine/pytorch/triton/mhc.py | 3 +- .../pytorch/triton/permutation.py | 31 +++++++++------- transformer_engine/pytorch/utils.py | 22 ++++++++---- 55 files changed, 300 insertions(+), 206 deletions(-) diff --git a/tests/pytorch/test_cross_entropy.py b/tests/pytorch/test_cross_entropy.py index f34e6666d4b..0cd77f28824 100644 --- a/tests/pytorch/test_cross_entropy.py +++ b/tests/pytorch/test_cross_entropy.py @@ -417,7 +417,7 @@ def test_parallel_cross_entropy_validation_errors(): target.transpose(0, 1), overwrite_input=True, ) - with pytest.raises(ValueError, match="CUDA"): + with pytest.raises(ValueError, match="cuda"): parallel_cross_entropy(logits.cpu(), target.cpu()) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index e4a5f9256b8..6237c6732b9 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -232,6 +232,16 @@ def load_framework_extension(framework: str) -> None: sys.modules[module_name + "_nv"] = solib _plugin = importlib.import_module(_nvte_plugin) _plugin.load_plugins() + # Apply device patches last: they rewrite process-global state in + # place and cannot be rolled back, so load_plugins() runs first and + # must not depend on the patched state. Still ahead of the PyTorch + # submodules that capture torch.cuda entry points. The submodule is + # optional: a CUDA-like plugin may omit it, and a single-file plugin + # has no submodules to search. + if hasattr(_plugin, "__path__") and ( + importlib.util.find_spec(f"{_nvte_plugin}.patches") is not None + ): + importlib.import_module(f"{_nvte_plugin}.patches").apply_patches() except Exception as e: # pylint: disable=broad-exception-caught # Rollback to pre-plugin state if plugin failed to fully initialize sys.modules.pop(module_name + "_nv", None) diff --git a/transformer_engine/debug/features/fake_quant.py b/transformer_engine/debug/features/fake_quant.py index 48112abe4fe..570c0aadf73 100644 --- a/transformer_engine/debug/features/fake_quant.py +++ b/transformer_engine/debug/features/fake_quant.py @@ -13,6 +13,7 @@ from nvdlfw_inspect.utils import append_parent_docstring +from transformer_engine import te_device_type from transformer_engine.debug.features.api import TEConfigAPIMapper from transformer_engine.common.recipe import Format from transformer_engine.pytorch import DType @@ -30,7 +31,9 @@ def fake_quantize(tensor: torch.Tensor, fp8_format: str, out=None): torch.float16, torch.bfloat16, ), "[NVTORCH INSPECT ERROR] Unsupported tensor type." - assert tensor.is_cuda, "[NVTORCH INSPECT ERROR] Must be a GPU tensor." + assert ( + tensor.device.type == te_device_type() + ), f"[NVTORCH INSPECT ERROR] Must be a {te_device_type()} tensor." assert fp8_format in { "FP8E4M3", "FP8E5M2", diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index 96f1b644cfb..efc0d8793b3 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -17,6 +17,8 @@ from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor + +from transformer_engine import te_device_type from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Quantizer, Float8CurrentScalingQuantizer, @@ -57,7 +59,10 @@ def _get_new_quantizer(recipe_name, fp8_dtype): return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) if recipe_name == "fp8_current_scaling": return Float8CurrentScalingQuantizer( - fp8_dtype=fp8_dtype, device=torch.device("cuda"), rowwise=True, columnwise=True + fp8_dtype=fp8_dtype, + device=torch.device(te_device_type()), + rowwise=True, + columnwise=True, ) if recipe_name == "mxfp8": return MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) diff --git a/transformer_engine/debug/features/per_tensor_scaling.py b/transformer_engine/debug/features/per_tensor_scaling.py index e0fc44bd08f..f83b2795042 100644 --- a/transformer_engine/debug/features/per_tensor_scaling.py +++ b/transformer_engine/debug/features/per_tensor_scaling.py @@ -11,6 +11,7 @@ import nvdlfw_inspect.api as debug_api from nvdlfw_inspect.registry import Registry, api_method +from transformer_engine import te_device_type from transformer_engine.pytorch import DType from transformer_engine.pytorch.tensor import Quantizer from transformer_engine.pytorch.tensor.float8_tensor import ( @@ -33,7 +34,9 @@ def per_tensor_cast( torch.float16, torch.bfloat16, ), "[NVTORCH INSPECT ERROR] Unsupported tensor type for per tensor current scaling" - assert tensor.is_cuda, "[NVTORCH INSPECT ERROR] Must be a GPU tensor." + assert ( + tensor.device.type == te_device_type() + ), f"[NVTORCH INSPECT ERROR] Must be a {te_device_type()} tensor." assert fp8_dtype in { DType.kFloat8E4M3, DType.kFloat8E5M2, diff --git a/transformer_engine/debug/features/utils/stats_buffer.py b/transformer_engine/debug/features/utils/stats_buffer.py index ca7f22e2de5..51cc5a0c1f2 100644 --- a/transformer_engine/debug/features/utils/stats_buffer.py +++ b/transformer_engine/debug/features/utils/stats_buffer.py @@ -16,6 +16,7 @@ from nvdlfw_inspect.utils import gather_along_first_dim from nvdlfw_inspect.logging import MetricLogger +from transformer_engine import te_device_type from transformer_engine.debug.features.utils.stats_computation import ( STATS, DEPENDENCIES, @@ -41,14 +42,14 @@ def __init__(self, layer_name, tensor_name, stats, reduction_group, reduce_withi for stat in stats: self.stats_to_compute = self.stats_to_compute | DEPENDENCIES[stat] - self._buffer = torch.zeros(len(STATS), dtype=torch.float32).cuda() + self._buffer = torch.zeros(len(STATS), dtype=torch.float32).to(te_device_type()) self._new_buffer = self._buffer.clone() self._tmp_buffer = self._buffer.clone() # in case of data parallelism it is possible that layer will not be run on one node # modified is set to True if node is run # we do not take not run nodes into account - self.modified = torch.tensor([False], dtype=torch.bool).cuda() + self.modified = torch.tensor([False], dtype=torch.bool).to(te_device_type()) self.iteration = None self.skip_reduction = False diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index aea10b50087..52f95340c44 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -9,6 +9,7 @@ import functools import torch +import transformer_engine from transformer_engine.common import load_framework_extension from transformer_engine.pytorch.torch_version import torch_version @@ -16,6 +17,26 @@ assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}." load_framework_extension("torch") + +# Device type and platform globals — live here because torch cannot be imported +# at the top-level transformer_engine package (shared with JAX). +# Use plugin-set values if available, otherwise fall back to CUDA defaults. +transformer_engine.TE_DEVICE_TYPE = getattr(transformer_engine, "TE_DEVICE_TYPE", "cuda") +transformer_engine.TE_PLATFORM = getattr(transformer_engine, "TE_PLATFORM", torch.cuda) + + +def te_device_type() -> str: + """Torch device type used by Transformer Engine, e.g. `cuda`.""" + return transformer_engine.TE_DEVICE_TYPE + + +def te_platform(): + """Torch device module used by Transformer Engine, e.g. `torch.cuda`.""" + return transformer_engine.TE_PLATFORM + + +transformer_engine.te_device_type = te_device_type +transformer_engine.te_platform = te_platform from transformer_engine.pytorch import constants from transformer_engine.pytorch.constants import CPLoadBalancingStrategy, DType from transformer_engine.pytorch.module import LayerNormLinear diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index bccb4f1e26f..717bcd6a461 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -18,6 +18,8 @@ import torch import torch.nn.functional as F + +from transformer_engine import te_device_type from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, @@ -663,10 +665,10 @@ def _forward( # S/dP are forced to use DS quantizers in DPA.init_fp8_metadata; revert them here for true CS emulation if fp8_recipe.float8_current_scaling(): S_quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=S_quantizer.dtype, device="cuda" + fp8_dtype=S_quantizer.dtype, device=te_device_type() ) dP_quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=dP_quantizer.dtype, device="cuda" + fp8_dtype=dP_quantizer.dtype, device=te_device_type() ) # disable swizzle for MXFP8Quantizer for quantizer in [ @@ -1038,9 +1040,9 @@ def forward( x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer] ), "FlashAttention only supports FP16 and BF16 data types, or Float8Tensors." - assert ( - query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda - ), "FlashAttention currently only supports CUDA tensors." + assert all( + x.device.type == te_device_type() for x in (query_layer, key_layer, value_layer) + ), f"FlashAttention currently only supports {te_device_type()} tensors." assert ( qkv_layout in QKVLayouts ), f"FlashAttention does not support qkv_layout = {qkv_layout}!" @@ -2718,9 +2720,9 @@ def forward( x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, QuantizedTensorStorage) for x in [query_layer, key_layer, value_layer] ), "FusedAttention only supports FP16 and BF16 data types, or QuantizedTensors." - assert ( - query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda - ), "FusedAttention only supports CUDA tensors." + assert all( + x.device.type == te_device_type() for x in (query_layer, key_layer, value_layer) + ), f"FusedAttention only supports {te_device_type()} tensors." assert ( qkv_layout in QKVLayouts ), f"FusedAttention does not support qkv_layout = {qkv_layout}!" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 3a153a3f994..0cdd590817a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -8,6 +8,7 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.pytorch.utils import ( get_cudnn_version, nvtx_range_pop, @@ -299,7 +300,9 @@ def get_thd_partitioned_indices( cp_size, ) target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype + target_dtype = ( + torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype + ) chunk_size = total_tokens // cp_size return torch.arange( cp_rank * chunk_size, @@ -308,9 +311,9 @@ def get_thd_partitioned_indices( device=target_device, ) cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, device) - assert cu_seqlens_padded.is_cuda, ( - "Per-document THD partitioning requires CUDA cu_seqlens; pass device='cuda' " - "when the source metadata is on CPU." + assert cu_seqlens_padded.device.type == te_device_type(), ( + f"Per-document THD partitioning requires {te_device_type()} cu_seqlens; pass" + f" device='{te_device_type()}' when the source metadata is on CPU." ) if cu_seqlens_padded.dtype != torch.int32: cu_seqlens_padded = cu_seqlens_padded.to(torch.int32) 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 e8e1396652d..f3a44708df7 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 @@ -15,6 +15,7 @@ from torch.fx.experimental.symbolic_shapes import guard_scalar from torch.nn.parameter import Parameter +from transformer_engine import te_device_type from transformer_engine.common.recipe import ( Format, Recipe, @@ -962,12 +963,14 @@ def __init__( self.softmax_offset = None if self.softmax_type == "off-by-one": self.softmax_offset = torch.zeros( - self.num_attention_heads // self.tp_size, device="cuda" + self.num_attention_heads // self.tp_size, device=te_device_type() ) if self.softmax_type == "learnable": self.register_parameter( "softmax_offset", - Parameter(torch.zeros(self.num_attention_heads // self.tp_size, device="cuda")), + Parameter( + torch.zeros(self.num_attention_heads // self.tp_size, device=te_device_type()) + ), get_rng_state_tracker=get_rng_state_tracker, ) @@ -2353,9 +2356,9 @@ def forward( fp8_output = False # checks for q/k/v shapes - assert ( - query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda - ), "DotProductAttention only supports CUDA tensors." + assert all( + x.device.type == te_device_type() for x in (query_layer, key_layer, value_layer) + ), f"DotProductAttention only supports {te_device_type()} tensors." assert ( query_layer.dtype == key_layer.dtype and query_layer.dtype == value_layer.dtype ), "Queries, keys and values must have the same data type!" @@ -2516,7 +2519,7 @@ def forward( "Mixed THD policy sequence_ids must be one-dimensional int32 or " "int64 tensors." ) - if not sequence_ids.is_cuda or sequence_ids.device != query_layer.device: + if sequence_ids.device != query_layer.device: raise ValueError( "Mixed THD policy sequence_ids must be on the Q/K/V device." ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index b9593b42d9b..2c8198881de 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -10,6 +10,7 @@ from typing import Any, Callable, Dict, Optional, Tuple import torch +from transformer_engine import te_device_type _cudnn_score_mod_handles: Dict[torch.device, Any] = {} _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} @@ -140,7 +141,7 @@ def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]: """Normalize a tensor device for graph cache keys.""" - if device.type == "cuda": + if device.type == te_device_type(): index = device.index if index is None: index = torch.cuda.current_device() @@ -194,10 +195,10 @@ def _wrapped_score_mod(sdpa_graph, score_tensor): def _get_cudnn_current_stream_handle(cudnn, device: torch.device): """Return a cuDNN handle for device, bound to PyTorch's current stream.""" - if device.type != "cuda": + if device.type != te_device_type(): raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.") if device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) handle = _cudnn_score_mod_handles.get(device) with torch.cuda.device(device): @@ -287,8 +288,8 @@ def _execute_cudnn_graph( """Execute a built cuDNN frontend Python graph.""" cudnn = _import_cudnn_frontend() - if device.type == "cuda" and device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) + if device.type == te_device_type() and device.index is None: + device = torch.device(te_device_type(), torch.cuda.current_device()) workspace = torch.empty( workspace_size, device=device, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index 6e8a14402f0..b55a960481d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -10,6 +10,7 @@ from torch import nn import transformer_engine_torch as tex from transformer_engine.pytorch.export import is_in_onnx_export_mode +from transformer_engine import te_device_type THREADS_PER_WARP = 32 THREADS_PER_BLOCK = 128 @@ -236,7 +237,7 @@ def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor: def _get_mask(): diagonal_offset = sk - sq + 1 if "bottom_right" in mask_type else 1 return torch.triu( - torch.ones(sq, sk, dtype=torch.bool, device="cuda"), diagonal=diagonal_offset + torch.ones(sq, sk, dtype=torch.bool, device=te_device_type()), diagonal=diagonal_offset ) if is_in_onnx_export_mode(): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 710ed8f7fbf..67a768d70c4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -22,6 +22,7 @@ import torch.nn.functional as F import transformer_engine_torch as tex import transformer_engine as te +from transformer_engine import te_device_type from transformer_engine.pytorch.cpp_extensions.fused_attn import ( QKVLayout, QKVFormat, @@ -2038,7 +2039,7 @@ def _mask(cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor: return ( (positions.unsqueeze(0) >= seqlens.unsqueeze(1)) .view(batch_size, 1, 1, max_seqlen) - .to(device="cuda") + .to(device=te_device_type()) ) attention_mask_q = _mask(cu_seqlens_q, max_seqlen_q) @@ -2164,9 +2165,11 @@ def get_full_mask( actual_seqlens_kv = m[:, 0, 0, :].sum(dim=1) # apply SWA mask - mask = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view( + mask = torch.arange(max_seqlen_q, dtype=torch.int32, device=te_device_type()).view( 1, 1, max_seqlen_q, 1 - ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view(1, 1, 1, max_seqlen_kv) + ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view( + 1, 1, 1, max_seqlen_kv + ) swa_left = None swa_right = None if attn_mask_type == "causal_bottom_right" or ( @@ -2262,7 +2265,7 @@ def get_alibi( m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (num_heads - n), 2)) m = torch.cat([m, m_hat]) - _alibi_cache["_alibi_slopes"] = m.to(dtype=torch.float32, device="cuda") + _alibi_cache["_alibi_slopes"] = m.to(dtype=torch.float32, device=te_device_type()) _alibi_cache["_num_heads"] = num_heads _alibi_cache["_alibi_slopes_require_update"] = False @@ -2275,9 +2278,9 @@ def get_alibi( else: raise ValueError("ALiBi slopes cannot exceed 2 dimensions.") - bias = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view( + bias = torch.arange(max_seqlen_q, dtype=torch.int32, device=te_device_type()).view( 1, 1, max_seqlen_q, 1 - ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view( + ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view( 1, 1, 1, max_seqlen_kv ) if actual_seqlens_q is None and actual_seqlens_kv is None: @@ -2297,7 +2300,9 @@ def get_alibi( _alibi_cache["_max_seqlen_q"], _alibi_cache["_max_seqlen_kv"] = max_seqlen_q, max_seqlen_kv _alibi_cache["_bottom_right_alignment"] = bottom_right_alignment bias_dtype = torch.float32 if bias_dtype is None else bias_dtype - _alibi_cache["_alibi_bias"] = bias.contiguous().to(dtype=bias_dtype, device="cuda") + _alibi_cache["_alibi_bias"] = bias.contiguous().to( + dtype=bias_dtype, device=te_device_type() + ) _alibi_cache["_alibi_bias_require_update"] = False return _alibi_cache["_alibi_slopes"], _alibi_cache["_alibi_bias"] @@ -2312,7 +2317,7 @@ def get_cu_seqlens(mask: torch.Tensor) -> torch.Tensor: mask = mask.squeeze(1).squeeze(1) reduced_mask = mask.logical_not().sum(dim=1) cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32) - zero = torch.zeros(1, dtype=torch.int32, device="cuda") + zero = torch.zeros(1, dtype=torch.int32, device=te_device_type()) cu_seqlens = torch.cat((zero, cu_seqlens)) return cu_seqlens @@ -2330,7 +2335,7 @@ def get_cu_seqlens_and_indices(mask: torch.Tensor) -> Tuple[torch.Tensor, torch. reduced_mask = mask.logical_not().sum(dim=1) cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32) - zero = torch.zeros(1, dtype=torch.int32, device="cuda") + zero = torch.zeros(1, dtype=torch.int32, device=te_device_type()) cu_seqlens = torch.cat((zero, cu_seqlens)) mask = mask.reshape(-1) @@ -2366,7 +2371,7 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor: ordered, bs * max_seqlen, ) - return indices.to(dtype=torch.int64, device="cuda").unsqueeze(1).unsqueeze(1) + return indices.to(dtype=torch.int64, device=te_device_type()).unsqueeze(1).unsqueeze(1) def get_full_cu_seqlens( diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index e07e203ca7e..93ed14ad413 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -14,6 +14,7 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm from transformer_engine.pytorch.ops.basic.l2normalization import L2Normalization +from transformer_engine import te_device_type from transformer_engine.pytorch.utils import ( SplitAlongDim, divide, @@ -141,7 +142,7 @@ class MultiheadAttention(torch.nn.Module): whether to use interleaved rotary position embeddings. bias : bool, default = True if set to ``False``, the transformer layer will not learn any additive biases. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -285,7 +286,7 @@ def __init__( ub_bulk_wgrad: bool = False, bias: bool = True, normalization: str = "LayerNorm", - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), qkv_format: str = "sbhd", name: str = None, qk_norm_type: Optional[str] = None, diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index 77ad57ed8f7..18397f954b3 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -10,6 +10,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat +from transformer_engine import te_device_type __all__ = ["RotaryPositionEmbedding", "apply_rotary_pos_emb", "apply_fused_qkv_rotary_pos_emb"] @@ -76,7 +77,7 @@ def forward(self, max_seq_len: int, offset: int = 0): offset: int, default = 0 Fixed offset for frequencies. """ - with torch.autocast(enabled=False, device_type="cuda"): + with torch.autocast(enabled=False, device_type=te_device_type()): seq = ( torch.arange(max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + offset diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f2f8537cc23..bc670b7fec9 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -11,6 +11,7 @@ import functools import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from ..constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType from ..utils import ceil_div, get_cached_ones_tensor, get_sm_count, _empty_tensor @@ -983,7 +984,8 @@ def general_grouped_gemm( if grad and use_bias: grad_bias = [ - torch.empty(B[i].size(1), dtype=out[0].dtype, device="cuda") for i in range(num_gemms) + torch.empty(B[i].size(1), dtype=out[0].dtype, device=te_device_type()) + for i in range(num_gemms) ] else: grad_bias = empty_tensors diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 8ae3b853801..2318969963b 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -10,6 +10,7 @@ import torch import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy +from transformer_engine import te_device_type from transformer_engine.pytorch.jit import no_torch_dynamo __all__ = ["parallel_cross_entropy"] @@ -114,10 +115,10 @@ def _validate_inputs( raise ValueError(f"inp dimensions must be non-zero, but got shape {tuple(inp.shape)}") if inp.dtype not in (torch.bfloat16, torch.float32): raise TypeError(f"inp must have BF16 or FP32 dtype, but got {inp.dtype}") - if not inp.is_cuda: - raise ValueError("inp must be a CUDA tensor") + if inp.device.type != te_device_type(): + raise ValueError(f"inp must be a {te_device_type()} tensor, but got {inp.device}") if target.device != inp.device: - raise ValueError("target must be on the same CUDA device as inp") + raise ValueError(f"target must be on the same {te_device_type()} device as inp") if target.dtype != torch.int64: raise TypeError(f"target must have torch.int64 dtype, but got {target.dtype}") if target.numel() != inp.shape[0] * inp.shape[1]: diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py index 250ce8c9393..b6d22901dda 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py @@ -29,6 +29,7 @@ import torch from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine import te_device_type from ..constants import DType @@ -86,7 +87,7 @@ def current_scaling_factory( return Float8CurrentScalingQuantizer( fp8_dtype=fp8_dtype, - device=torch.device("cuda"), + device=torch.device(te_device_type()), force_pow_2_scales=False, amax_epsilon=0.0, ) diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 0d5fc0c4419..5cf573d61e4 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -72,6 +72,7 @@ from typing import Optional from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine import te_device_type from ..constants import DType from .quantizer_factories import mxfp8_factory, nvfp4_factory @@ -426,6 +427,6 @@ def nvfp4_linear_fp8_dpa_factory( is_dpa_boundary and "dpa_grad_input" in role.name ) fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 - return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device=te_device_type()) return nvfp4_factory(role) diff --git a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index dcb77bbbce1..6d016500c6a 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -13,6 +13,7 @@ from transformer_engine.pytorch.custom_recipes import gemm from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer +from transformer_engine import te_device_type def current_scaling_ref_factory(role): @@ -502,7 +503,7 @@ def make_empty( # Canonicalize tensor attributes if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) # Allocate quantized data qx = torch.empty(shape, dtype=self.dtype, device=device) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index e90b17e4d9f..9b59970ba88 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -31,6 +31,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.triton.pad import pad_columnwise_scale_inv +from transformer_engine import te_device_type from .torch_version import torch_version from .utils import ( is_non_tn_fp8_gemm_supported, @@ -96,7 +97,7 @@ def is_graph_safe_rng_state(state: Union[torch.Tensor, torch.Generator]) -> bool def _get_cuda_rng_state( - device: Union[int, str, torch.device] = "cuda", + device: Union[int, str, torch.device] = te_device_type(), clone: bool = False, graph_safe: bool = True, ) -> torch.Tensor: @@ -106,7 +107,7 @@ def _get_cuda_rng_state( if isinstance(device, str): device = torch.device(device) elif isinstance(device, int): - device = torch.device("cuda", device) + device = torch.device(te_device_type(), device) idx = device.index if idx is None: idx = torch.cuda.current_device() @@ -128,11 +129,11 @@ def _set_cuda_rng_state( """Sets the random number generator state of the current GPU.""" if device == -1: - device = torch.device("cuda") + device = torch.device(te_device_type()) elif isinstance(device, str): device = torch.device(device) elif isinstance(device, int): - device = torch.device("cuda", device) + device = torch.device(te_device_type(), device) def cb() -> None: idx = device.index @@ -294,10 +295,10 @@ def _get_active_autocast_contexts(): autocast_cached = torch.is_autocast_cache_enabled() if torch_version() >= (2, 4, 0): - gpu_autocast_enabled = torch.is_autocast_enabled("cuda") - gpu_autocast_dtype = torch.get_autocast_dtype("cuda") + gpu_autocast_enabled = torch.is_autocast_enabled(te_device_type()) + gpu_autocast_dtype = torch.get_autocast_dtype(te_device_type()) gpu_autocast_ctx = torch.amp.autocast( - "cuda", + te_device_type(), enabled=gpu_autocast_enabled, dtype=gpu_autocast_dtype, cache_enabled=autocast_cached, @@ -1047,7 +1048,7 @@ def _all_gather_fp8( out: Float8TensorStorage if quantizer is not None: dtype = torch.float32 - device = "cuda" + device = te_device_type() if isinstance(inp, Float8Tensor): dtype = inp.dtype device = inp.device @@ -1987,7 +1988,7 @@ def symm_mem_alloc( 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()) + device = torch.device(te_device_type(), torch.cuda.current_device()) if not HAS_TORCH_SYMMETRIC: raise RuntimeError( "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc " diff --git a/transformer_engine/pytorch/dynamo/tensor_spec.py b/transformer_engine/pytorch/dynamo/tensor_spec.py index d96c49ca7a5..e5afa0fdd0b 100644 --- a/transformer_engine/pytorch/dynamo/tensor_spec.py +++ b/transformer_engine/pytorch/dynamo/tensor_spec.py @@ -11,6 +11,7 @@ import torch from torch._prims_common import make_contiguous_strides_for +from transformer_engine import te_device_type @dataclass @@ -108,7 +109,7 @@ def create_inner_tensors(self) -> List[torch.Tensor]: Under ``register_fake`` the ``torch.empty`` calls produce ``FakeTensor``s; ``requires_grad`` is left default (managed by ``register_autograd``). """ - device = self.device if self.device is not None else torch.device("cuda") + device = self.device if self.device is not None else torch.device(te_device_type()) if self.quantizer is None: return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)] inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device) @@ -139,7 +140,7 @@ def create_tensor(self) -> torch.Tensor: inner tensors via :meth:`assemble`. """ if self.quantizer is None: - device = self.device if self.device is not None else torch.device("cuda") + device = self.device if self.device is not None else torch.device(te_device_type()) return torch.empty( tuple(self.shape), dtype=self.dtype, diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 8bd2f68a070..8052f65bc2b 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -10,6 +10,7 @@ from typing import Callable, Optional, Tuple import torch +from .. import te_device_type from .torch_version import torch_version from .export import is_in_onnx_export_mode from .utils import gpu_autocast_ctx @@ -337,9 +338,17 @@ def warmup_jit_bias_dropout_add( # Save cuda RNG state to ensure warmup does not affect reproducibility. rng_state = torch.cuda.get_rng_state() - inp = torch.rand((seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda") - residual = torch.rand((seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda") - bias = torch.rand((hidden_size), dtype=dtype, device="cuda") + inp = torch.rand( + (seq_length, micro_batch_size, hidden_size), + dtype=dtype, + device=te_device_type(), + ) + residual = torch.rand( + (seq_length, micro_batch_size, hidden_size), + dtype=dtype, + device=te_device_type(), + ) + bias = torch.rand((hidden_size), dtype=dtype, device=te_device_type()) dropout_rate = 0.1 # Warmup JIT fusions with the input grad_enable state of both forward # prop and recomputation @@ -374,11 +383,11 @@ def warmup_jit_bias_gelu( # Save cuda RNG state to ensure warmup does not affect reproducibility. rng_state = torch.cuda.get_rng_state() - bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device="cuda") + bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device=te_device_type()) inp = torch.rand( (seq_length * micro_batch_size, ffn_hidden_size_per_partition), dtype=dtype, - device="cuda", + device=te_device_type(), ) # Warmup JIT fusions with the input grad_enable state of both forward # prop and recomputation @@ -412,7 +421,7 @@ def warmup_jit_l2normalization( inp = torch.rand( (seq_length * micro_batch_size, hidden_size), dtype=dtype, - device="cuda", + device=te_device_type(), ) eps = 1e-6 # Warmup JIT fusions with the input grad_enable state of both forward diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5aea8feea1c..2e998761154 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -21,6 +21,7 @@ import transformer_engine_torch as tex +from transformer_engine import te_device_type from ._common import _ParameterInitMeta, noop_cat, sum_bias_grad from .._extra_state import ( extra_state_pickle_advisory, @@ -63,6 +64,7 @@ from ..utils import ( is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, + torch_is_autocast_gpu_enabled, get_device_compute_capability, get_nvtx_range_context, nvtx_range_push, @@ -146,7 +148,7 @@ def get_dummy_wgrad( # The dummy goes on the current device (the one the backward pass runs on), and the cache # is keyed by it so that a dummy is never returned for a weight on another device. - device = torch.device("cuda", torch.cuda.current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) if zero_tensor and not zero: return torch._efficientzerotensor(shape, dtype=dtype, device=device) key = (*shape, dtype, device) @@ -1554,7 +1556,7 @@ def set_extra_state(self, state: torch.Tensor) -> None: if not unsafe_pickle_extra_state_enabled(): raise RuntimeError(extra_state_pickle_advisory(context)) state.seek(0) - state = torch.load(state, map_location="cuda", weights_only=False) + state = torch.load(state, map_location=te_device_type(), weights_only=False) else: raise RuntimeError("Unsupported checkpoint format.") @@ -1598,7 +1600,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: def set_activation_dtype(self, inp: torch.Tensor) -> None: """Get activation data type for AMP.""" # Native AMP (`torch.autocast`) gets highest priority - if torch.is_autocast_enabled(): + if torch_is_autocast_gpu_enabled(): self.fast_setattr("activation_dtype", torch_get_autocast_gpu_dtype()) return dtype = inp.dtype @@ -1718,9 +1720,9 @@ def prepare_forward( delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta) FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) else: - if not inp.is_cuda: + if inp.device.type != te_device_type(): raise RuntimeError( - f"TransformerEngine needs CUDA. Got input on device: {inp.device}" + f"TransformerEngine needs {te_device_type()}. Got input on device: {inp.device}" ) if self.tp_size > 1: @@ -1934,7 +1936,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: param = param._local_tensor if is_dtensor else param # Ensure parameter is on a real device if param.device == torch.device("meta"): - param = torch.empty_like(param, device="cuda") + param = torch.empty_like(param, device=te_device_type()) # Initialize the parameter values on device init_fn = self.param_init_meta[name].init_fn get_rng_state_tracker = self.param_init_meta[name].get_rng_state_tracker diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index cfa59034de2..fc595357b13 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -17,6 +17,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe, Format as RecipeFormat +from transformer_engine import te_device_type from transformer_engine.pytorch.tensor.grouped_tensor import ( GroupedTensor, GroupedTensorStorage, @@ -1631,7 +1632,7 @@ class GroupedLinear(TransformerEngineBaseModule): used to get the random number generator state tracker for initializing weights. rng_tracker_name : str, default = None the param passed to get_rng_state_tracker to get the specific rng tracker. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -1709,7 +1710,7 @@ def __init__( return_bias: bool = False, params_dtype: Optional[torch.dtype] = None, parallel_mode: Optional[str] = None, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_rs: bool = False, ub_overlap_ag: bool = False, ub_name: Optional[str] = None, @@ -2418,7 +2419,7 @@ def forward( " single_grouped_bias to allow the split-quantize fallback." ) if use_grouped_tensor_path: - if m_splits.device.type != "cuda": + if m_splits.device.type != te_device_type(): raise ValueError( "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " "tensor, or set use_grouped_tensor=False." diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 444b66f2eb1..03e6cf50de8 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -17,6 +17,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom +from transformer_engine import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_cublasmp_all_gather_output, @@ -1600,7 +1601,7 @@ class LayerNormLinear(TransformerEngineBaseModule): .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -1674,7 +1675,7 @@ def __init__( return_layernorm_output_gathered: bool = False, parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None, zero_centered_gamma: bool = False, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_ag: bool = False, ub_overlap_rs: bool = False, ub_overlap_rs_dgrad: bool = False, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 10c404bb5d5..92ae0c9ac23 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -18,6 +18,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom +from transformer_engine import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_cublasmp_all_gather_output, @@ -1798,7 +1799,7 @@ def fc2_wgrad_gemm( reduce_scatter_out = None if args.ub_overlap_rs_dgrad: reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda" + fc1_dgrad_shape, dtype=args.activation_dtype, device=te_device_type() ) if ub_bulk_wgrad: gemm_out = ub_obj_fc1_wgrad.get_buffer( @@ -1896,7 +1897,7 @@ def fc2_wgrad_gemm( reduce_scatter_out = None if ub_bulk_wgrad and ub_obj_fc1_wgrad.is_fp8_ubuf(): reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda" + fc1_dgrad_shape, dtype=args.activation_dtype, device=te_device_type() ) # Arguments to include in wgrad GEMM closure @@ -2222,7 +2223,7 @@ class LayerNormMLP(TransformerEngineBaseModule): .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -2311,7 +2312,7 @@ def __init__( micro_batch_size: Optional[int] = None, set_parallel_mode: bool = False, zero_centered_gamma: bool = False, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_ag: bool = False, name: Optional[str] = None, ub_overlap_rs: bool = False, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 808de29c134..79768cd4f40 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -16,6 +16,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version +from transformer_engine import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_cublasmp_all_gather_output, @@ -1955,7 +1956,7 @@ class Linear(TransformerEngineBaseModule): values as split sizes along dim 0. The resulting parameters will have names that end in ``_weight`` or ``_bias``, so trailing underscores are stripped from any provided names. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -2031,7 +2032,7 @@ def __init__( params_dtype: Optional[torch.dtype] = None, parallel_mode: Optional[str] = None, parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_ag: bool = False, ub_overlap_rs: bool = False, ub_overlap_rs_dgrad: bool = False, @@ -2342,7 +2343,7 @@ def _prealloc_cublas_workspace(self) -> None: from torch._subclasses.fake_tensor import is_fake weight = getattr(self, self.weight_names[0], None) - if weight is None or weight.device.type != "cuda": + if weight is None or weight.device.type != te_device_type(): return if is_fake(weight) or detect_fake_mode(): return diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index cf4e66023a6..e86b0b5a778 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -11,6 +11,7 @@ import torch from transformer_engine_torch import FP8TensorMeta +from transformer_engine import te_device_type from ..dynamo.tensor_spec import TensorSpec from ..ep import ( EpBuffer, @@ -179,7 +180,7 @@ def maybe_dequantize( def maybe_autocast_dtype( *, - device_type: str = "cuda", + device_type: str = te_device_type(), default_dtype: Optional[torch.dtype] = None, ) -> torch.dtype: """Get autocast dtype if enabled""" diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index ffcb72a9554..21a376ca3d1 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -18,7 +18,7 @@ from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import _compile_safe_warn, clear_tensor_data from ..op import BasicOperation, OperationContext -from .._common import maybe_dequantize +from .._common import maybe_autocast_dtype, maybe_dequantize __all__ = [ "GELU", @@ -94,11 +94,7 @@ def op_forward( ) -> torch.Tensor: # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = input_.dtype + dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype) if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise RuntimeError(f"Unsupported dtype ({dtype})") @@ -412,12 +408,8 @@ def fuser_forward( "in the fused grouped MLP path." ) - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - extra_input_dtype = dtype - else: - dtype = input_.dtype - extra_input_dtype = extra_input.dtype + dtype = maybe_autocast_dtype(default_dtype=input_.dtype) + extra_input_dtype = maybe_autocast_dtype(default_dtype=extra_input.dtype) x = maybe_dequantize(input_.contiguous(), dtype) scales = maybe_dequantize(extra_input.contiguous(), extra_input_dtype) diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index df8e13604a9..e60662226d3 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -43,6 +43,7 @@ get_dummy_wgrads_for_params, get_main_grad_from_param, is_quantized_tensor, + maybe_autocast_dtype, maybe_dequantize, ) @@ -1108,11 +1109,7 @@ def pack_forward_args( return BasicLinearFwdArgs( input_=input_, weight=self.weight, - dtype=( - torch.get_autocast_dtype("cuda") - if torch.is_autocast_enabled() - else self.weight.dtype - ), + dtype=maybe_autocast_dtype(default_dtype=self.weight.dtype), input_requires_grad=basic_op_ctxs[0].requires_grad, weight_requires_grad=basic_op_ctxs[0].requires_grad and self.weight.requires_grad, with_quantized_compute=with_quantized_compute, diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 0dab1e9ed0b..0a26e0460b9 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -11,6 +11,7 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...dynamo import TensorOrQuantized, TensorSpec from ...dynamo.fake import fake_bgrad from ...quantization import FP8GlobalStateManager @@ -115,7 +116,7 @@ def reset_parameters(self) -> None: # Make sure parameter is initialized bias = self.bias - if bias.device.type != "cuda": + if bias.device.type != te_device_type(): bias = torch.empty_like(bias, device=self.device) else: bias = bias.to(device=self.device) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 171500d7798..0dfc8e2366a 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -49,6 +49,7 @@ get_dummy_wgrads_for_params, get_main_grad_from_param, is_quantized_tensor, + maybe_autocast_dtype, maybe_dequantize, validate_or_alloc_output, view_main_grad_as_grouped_buffer, @@ -1055,10 +1056,7 @@ def fuser_forward( weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = weight_param.dtype + dtype = maybe_autocast_dtype(default_dtype=weight_param.dtype) # Extract split sizes from extra input. Keep on GPU for graph safety. split_sizes = basic_op_extra_inputs[0][0] @@ -1209,10 +1207,7 @@ def fuser_forward_save_ctx( # ``split_sizes``, offset metadata, and related tensors are routed # through ``save_for_backward`` (see ``_fuser_forward_split_quantize`` # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). - if torch.is_autocast_enabled(): - ctx.dtype = torch.get_autocast_dtype("cuda") - else: - ctx.dtype = weight_param.dtype + ctx.dtype = maybe_autocast_dtype(default_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. diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 8009de9b4ae..71ff010ae53 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -17,7 +17,7 @@ from ...tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data from ..op import BasicOperation, OperationContext -from .._common import maybe_dequantize +from .._common import maybe_autocast_dtype, maybe_dequantize __all__ = [ "SwiGLU", @@ -114,11 +114,7 @@ def op_forward( ) -> torch.Tensor: # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = input_.dtype + dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype) if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise RuntimeError(f"Unsupported dtype ({dtype})") @@ -361,11 +357,7 @@ def op_forward( ) -> torch.Tensor: # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = input_.dtype + dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype) if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise RuntimeError(f"Unsupported dtype ({dtype})") @@ -528,12 +520,8 @@ def fuser_forward( extra_input = basic_op_extra_inputs[0][0] # Determine compute dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - extra_input_dtype = dtype - else: - dtype = input_.dtype - extra_input_dtype = extra_input.dtype + dtype = maybe_autocast_dtype(default_dtype=input_.dtype) + extra_input_dtype = maybe_autocast_dtype(default_dtype=extra_input.dtype) # Prepare the activation input in the compute dtype. input_ = maybe_dequantize(input_, dtype) diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 28586360f59..7f7a6bd97e3 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -13,6 +13,7 @@ from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer +from .._common import maybe_autocast_dtype from ..basic import AddExtraInput, BasicLinear, Bias from ..op import FusedOperation, FusibleOperation, OperationContext @@ -92,10 +93,7 @@ def fuser_forward( backward_override = None # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = linear_op.weight.dtype + dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype) # Linear forward output = basic_op_extra_inputs[self._op_idxs["add"]][0] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index 277263e0ecb..da78bec23b7 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -13,6 +13,7 @@ from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer +from .._common import maybe_autocast_dtype from ..basic import AddExtraInput, BasicLinear, ConstantScale from ..op import ( FusedOperation, @@ -74,10 +75,7 @@ def fuser_forward( extra_input = basic_op_extra_inputs[2][0] # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = linear_op.weight.dtype + dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype) # Linear forward output, x_local, w = BasicLinear._functional_forward( diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d678d92749c..bc64d8ed9ae 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -60,6 +60,7 @@ get_dummy_wgrads_for_params, get_main_grad_from_param, is_quantized_tensor, + maybe_autocast_dtype, maybe_dequantize, validate_or_alloc_output, view_main_grad_as_grouped_buffer, @@ -1381,10 +1382,7 @@ def fuser_forward( 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") - else: - dtype = fc1_weight_param.dtype + dtype = maybe_autocast_dtype(default_dtype=fc1_weight_param.dtype) # Check which grads are required requires_grad = any(ctx.requires_grad for ctx in basic_op_ctxs) diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 37c3dee9a16..39b2b6fe136 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -11,6 +11,7 @@ import torch from transformer_engine_torch import CommOverlapType, bulk_overlap_ag_with_external_gemm +from transformer_engine import te_device_type from ...cpp_extensions import general_gemm from ...distributed import gather_along_first_dim, get_distributed_world_size from ...module.base import ( @@ -182,7 +183,7 @@ def _functional_backward( else: device = grad_output.device device = canonicalize_device(device) - if device.type != "cuda": + if device.type != te_device_type(): raise ValueError(f"Only CUDA devices are supported (got {device})") # Check datatype diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index cb5f8a16e61..9bc2cfc1e75 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -11,6 +11,7 @@ import torch from transformer_engine_torch import CommOverlapType +from transformer_engine import te_device_type from ...cpp_extensions import general_gemm from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import get_distributed_world_size @@ -23,7 +24,7 @@ from ...quantized_tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ...tensor.storage.float8_tensor_storage import Float8TensorStorage -from .._common import maybe_dequantize, is_quantized_tensor +from .._common import maybe_autocast_dtype, maybe_dequantize, is_quantized_tensor from ..basic import BasicLinear, Bias, ReduceScatter from ..op import ( FusedOperation, @@ -155,7 +156,7 @@ def _functional_forward( """ # Check device - if device.type != "cuda": + if device.type != te_device_type(): raise ValueError(f"Only CUDA devices are supported (got {device})") # Check datatype @@ -321,10 +322,7 @@ def fuser_forward( ) # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = linear_op.weight.dtype + dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype) # Userbuffers options if linear_op._userbuffers_options is None: diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index b4b7f707a83..2cde7780af0 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -15,6 +15,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +from transformer_engine import te_device_type from ..constants import DType from .multi_tensor_apply import multi_tensor_applier @@ -190,7 +191,7 @@ def __init__( self._step_supports_amp_scaling = True # Skip buffer - self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda") + self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=te_device_type()) self.multi_tensor_adam = tex.multi_tensor_adam self.multi_tensor_adam_param_remainder = tex.multi_tensor_adam_param_remainder self.multi_tensor_adam_fp8 = tex.multi_tensor_adam_fp8 diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index eeb26dca632..2cb49d4a9b4 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -7,6 +7,7 @@ from typing import Optional, Tuple import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type import transformer_engine.pytorch.triton.permutation as triton_permutation from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch.quantized_tensor import ( @@ -44,10 +45,14 @@ def moe_permute_index_map_forward( if not inp.numel(): return inp.clone(), torch.empty(0, dtype=torch.int32, device=inp.device) - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not index.is_cuda: - raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.") + if inp.device.type != te_device_type(): + raise ValueError( + f"inp must be a {te_device_type()} tensor, but got tensor on {inp.device}." + ) + if index.device.type != te_device_type(): + raise ValueError( + f"index must be a {te_device_type()} tensor, but got tensor on {index.device}." + ) if inp.size(0) != index.size(0): raise ValueError( f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " @@ -299,19 +304,25 @@ def moe_permute_mask_map_forward( ) return inp.clone(), row_id_map, permuted_probs - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not routing_map.is_cuda: + if inp.device.type != te_device_type(): raise ValueError( - f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}." + f"inp must be a {te_device_type()} tensor, but got tensor on {inp.device}." + ) + if routing_map.device.type != te_device_type(): + raise ValueError( + f"routing_map must be a {te_device_type()} tensor, but got tensor on" + f" {routing_map.device}." ) if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") + if probs.device.type != te_device_type(): + raise ValueError( + f"probs must be a {te_device_type()} tensor, but got tensor on {probs.device}." + ) if pad_offsets is not None: - if not pad_offsets.is_cuda: + if pad_offsets.device.type != te_device_type(): raise ValueError( - f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + f"pad_offsets must be a {te_device_type()} tensor, but got tensor on" + f" {pad_offsets.device}." ) if inp.size(0) != routing_map.size(0): raise ValueError( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index c167e5bb243..06d02ece44e 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -16,6 +16,8 @@ import torch import transformer_engine_torch as tex + +from transformer_engine import te_device_type from transformer_engine.common.recipe import ( Recipe, DelayedScaling, @@ -480,7 +482,7 @@ def set_skip_fp8_weight_update_tensor(cls, skip: bool) -> None: """Set the skip fp8 weight update tensor""" if cls.quantization_state.skip_fp8_weight_update_tensor is None: cls.quantization_state.skip_fp8_weight_update_tensor = torch.empty( - 1, dtype=torch.float32, device="cuda" + 1, dtype=torch.float32, device=te_device_type() ) cls.quantization_state.skip_fp8_weight_update_tensor.fill_(skip) @@ -1488,7 +1490,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.scale = torch.ones(num_quantizers, dtype=torch.float32, device=device) self.amax_history = torch.zeros( recipe.amax_history_len, @@ -1537,7 +1539,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.device = device def make_quantizers(self) -> list: @@ -1580,7 +1582,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) def make_quantizers(self) -> list: # TODO(ksivamani); Find better design for this, adding here to avoid circular import. @@ -1638,7 +1640,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.device = device def make_quantizers(self) -> list: @@ -1721,7 +1723,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) def make_quantizers(self) -> list: """Build one ``NVFP4Quantizer`` per slot, dispatched by tensor type. @@ -1965,7 +1967,7 @@ def __init__( self.num_quantizers = num_quantizers self.roles = roles if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.device = device # -- Stateful sub-state slots (initialized empty) -- diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index fc6b59a32f1..794d383ff26 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -18,6 +18,8 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.constants import dist_group_type + +from transformer_engine import te_device_type from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, @@ -535,7 +537,7 @@ def make_empty( ) if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) # Handle the device passed as string device = torch.device(device) result = tex.create_empty_quantized_tensor( @@ -591,7 +593,7 @@ def alloc_tensors( Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors`` argument of the storage's ``__tensor_unflatten__``. """ - device = torch.device(device if device is not None else "cuda") + device = torch.device(device if device is not None else te_device_type()) return { attr: torch.empty(buf_shape, dtype=buf_dtype, device=device) for attr, (buf_shape, buf_dtype) in self.inner_tensor_specs(tuple(shape)).items() diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 6105b18a736..c85577fe2f8 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -11,6 +11,8 @@ import torch import transformer_engine_torch as tex + +from transformer_engine import te_device_type from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import ( Float8BlockwiseQTensorStorage, @@ -441,7 +443,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): qt._rowwise_scale_inv, qt._columnwise_scale_inv, ): - if t is not None and t.is_cuda: + if t is not None and t.device.type == te_device_type(): t.record_stream(stream) return None @@ -538,7 +540,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device def _set_from_tensor(dst: Float8BlockwiseQTensor, src: Float8BlockwiseQTensor): dst._rowwise_data = src._rowwise_data diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index cf37c36c598..d37b8e2800e 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -10,6 +10,7 @@ from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.common.recipe import ( DelayedScaling, Float8CurrentScaling, @@ -1139,7 +1140,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device if not devices_match(new_device, tensor.device): tensor = tensor.to(device=new_device) diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index f50061f3b5a..83c3816af84 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -10,6 +10,7 @@ import torch from torch.utils._pytree import tree_map +from transformer_engine import te_device_type from ..quantized_tensor import QuantizedTensorStorage, Quantizer from .storage.grouped_tensor_storage import GroupedTensorStorage from ..constants import DType @@ -133,7 +134,7 @@ def __new__( device = maybe_tensor.device break if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) # Match QuantizedTensor __new__: accept externally-computed stride to # avoid Python-side stride computation overhead for C++ construction. @@ -240,7 +241,7 @@ def make_wrapper_like(src: GroupedTensor) -> GroupedTensor: if not isinstance(src, GroupedTensor): raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") for tensor in src.get_data_tensors(): - if tensor is not None and tensor.is_cuda: + if tensor is not None and tensor.device.type == te_device_type(): tensor.record_stream(stream) return None diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index ec171564fe2..8f8bda44cbc 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -19,6 +19,7 @@ import torch from torch.ops import aten +from transformer_engine import te_device_type from .storage.identity_tensor_storage import IdentityTensorStorage from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer @@ -114,7 +115,7 @@ def make_empty( pin_memory: bool = False, ) -> Union["IdentityTensor", IdentityTensorStorage]: if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) device = torch.device(device) data_dtype = self.dtype if self.dtype is not None else dtype data = torch.empty(tuple(shape), dtype=data_dtype, device=device, pin_memory=pin_memory) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 267806e43ee..c4081b75c64 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -13,6 +13,7 @@ from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe from ..constants import MXFP8_BLOCK_SCALING_SIZE, DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -823,7 +824,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device if not devices_match(new_device, tensor.device): tensor = tensor.to(device=new_device) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 938d9a77d4c..29be015b6ec 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -13,6 +13,7 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType from ..utils import ( @@ -796,7 +797,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): qt._amax_rowwise, qt._amax_columnwise, ): - if t is not None and t.is_cuda: + if t is not None and t.device.type == te_device_type(): t.record_stream(stream) return None @@ -954,7 +955,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device if not devices_match(new_device, tensor.device): tensor = tensor.to(device=new_device) diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index f4a95607d48..cccd048cb83 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -12,6 +12,7 @@ import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr @@ -44,8 +45,8 @@ def forward( te_dtype = torch_to_transformer_engine_dtype[dtype] # ``tex.dequantize`` requires CUDA-resident buffers. src_device = tensor.device - if src_device.type != "cuda": - cuda_tensor = tensor.to(device=torch.device("cuda")) + if src_device.type != te_device_type(): + cuda_tensor = tensor.to(device=torch.device(te_device_type())) result = tex.dequantize(cuda_tensor, te_dtype) return result.to(device=src_device) return tex.dequantize(tensor, te_dtype) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 97b66a8341f..9931325eec3 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -15,6 +15,7 @@ import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr @@ -54,8 +55,8 @@ def forward( # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor has src_device = tensor.device - if src_device.type != "cuda": - cuda_tensor = tensor.to(device=torch.device("cuda")) + if src_device.type != te_device_type(): + cuda_tensor = tensor.to(device=torch.device(te_device_type())) result = tex.dequantize(cuda_tensor, torch_to_transformer_engine_dtype[dtype]) return result.to(device=src_device) return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype]) diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index db9a352881c..cc361454a4c 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -14,6 +14,8 @@ from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention from transformer_engine.pytorch.attention.inference import InferenceParams + +from transformer_engine import te_device_type from transformer_engine.pytorch.jit import ( set_jit_fusion_options, warmup_jit_bias_dropout_add_all_dtypes, @@ -25,6 +27,7 @@ cast_if_needed, get_default_init_method, torch_get_autocast_gpu_dtype, + torch_is_autocast_gpu_enabled, ) from transformer_engine.pytorch.constants import ( AttnMaskTypes, @@ -199,7 +202,7 @@ class TransformerLayer(torch.nn.Module): supports ``'limit'``, ``'alpha'``, and ``'glu_linear_offset'`` parameters. You can set these as ``activation_params={'limit': 7.0, 'alpha': 1.702, 'glu_linear_offset': 1.0}``. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -352,7 +355,7 @@ def __init__( activation: str = "gelu", activation_params: Optional[dict] = None, normalization: str = "LayerNorm", - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), attn_input_format: str = "sbhd", name: str = None, qk_norm_type: Optional[str] = None, @@ -907,7 +910,7 @@ def forward( ) # For AMP - if torch.is_autocast_enabled(): + if torch_is_autocast_gpu_enabled(): hidden_states = cast_if_needed(hidden_states, torch_get_autocast_gpu_dtype()) # Self attention. diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py index 41f71966a23..ce9a507b448 100644 --- a/transformer_engine/pytorch/triton/mhc.py +++ b/transformer_engine/pytorch/triton/mhc.py @@ -25,6 +25,7 @@ _mhc_sinkhorn_bwd_fused, ) from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm +from transformer_engine import te_device_type ENFORCE_DETERMINISTIC = os.environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") == "0" @@ -50,7 +51,7 @@ def _init_tma_allocator(): def alloc_fn( size: int, alignment: int, stream: Optional[int] ): # pylint: disable=unused-argument - return torch.empty(size, device="cuda", dtype=torch.int8) + return torch.empty(size, device=te_device_type(), dtype=torch.int8) triton.set_allocator(alloc_fn) _tma_allocator_initialized = True diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 4f8a2ad7a42..64d3c89128e 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -9,6 +9,7 @@ import torch import triton +from transformer_engine import te_device_type from transformer_engine.common.triton.permutation import ( _row_id_map_pass_1_kernel, _row_id_map_pass_2_kernel, @@ -53,10 +54,12 @@ def make_row_id_map( The [num_experts, num_experts + n_routed) items are the indices of the experts corresponding to the first n_routed row indices above. """ - row_id_map = torch.empty((num_tokens, num_experts * 2 + 1), dtype=torch.int32, device="cuda") + row_id_map = torch.empty( + (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=routing_map.device + ) block_size = 1024 grid = (num_experts, triton.cdiv(num_tokens, block_size)) - workspace_tensor = torch.empty(grid, dtype=torch.int32, device="cuda") + workspace_tensor = torch.empty(grid, dtype=torch.int32, device=routing_map.device) # supposing num_tokens == 5, num_experts == 3, block_size == 3 # and we have a routing_map like this: @@ -165,12 +168,14 @@ def permute_with_mask_map( # Use torch.zeros when pad_offsets is provided to ensure padding regions are zeroed. # The kernel writes only to valid positions, leaving padding positions at zero. alloc = torch.zeros if pad_offsets is not None else torch.empty - output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda") + output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device=te_device_type()) permuted_probs = ( - alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None + alloc((num_out_tokens,), dtype=probs.dtype, device=te_device_type()) + if probs is not None + else None ) permuted_scale = ( - alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda") + alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device=te_device_type()) if scale is not None else None ) @@ -248,10 +253,10 @@ def unpermute_with_mask_map( hidden_size : int Hidden size of the permuted tensor. """ - output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda") + output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=te_device_type()) if permuted_probs is not None: unpermuted_probs = torch.empty( - (num_tokens, num_experts), dtype=permuted_probs.dtype, device="cuda" + (num_tokens, num_experts), dtype=permuted_probs.dtype, device=te_device_type() ) else: unpermuted_probs = None @@ -330,9 +335,11 @@ def unpermute_with_mask_map_bwd_with_merging_probs( # by the kernel. This matches the behavior of Fp8Unpadding.backward which zeros # out the padding slots. alloc = torch.zeros if pad_offsets is not None else torch.empty - act_grad = alloc((num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device="cuda") + act_grad = alloc( + (num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device=te_device_type() + ) merging_probs_grad = torch.empty( - (num_tokens, num_experts), dtype=merging_probs.dtype, device="cuda" + (num_tokens, num_experts), dtype=merging_probs.dtype, device=te_device_type() ) grid = (num_tokens,) _unpermute_bwd_with_merging_probs_kernel[grid]( @@ -383,7 +390,7 @@ def make_chunk_sort_map( num_splits : int Number of splits of split_sizes and sorted_indices. """ - row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="cuda") + row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device=te_device_type()) grid = (num_tokens,) _make_chunk_sort_map_kernel[grid]( split_sizes, @@ -421,9 +428,9 @@ def sort_chunks_by_map( is_forward : bool Whether the sort is for forward or backward. """ - output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda") + output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=te_device_type()) if probs is not None: - permuted_probs = torch.empty((num_tokens,), dtype=probs.dtype, device="cuda") + permuted_probs = torch.empty((num_tokens,), dtype=probs.dtype, device=te_device_type()) else: permuted_probs = None # pylint: disable=unnecessary-lambda-assignment diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 62afd10f153..a3208db5026 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -14,6 +14,7 @@ import numpy as np import torch +from transformer_engine import te_device_type from .torch_version import torch_version from ..common import decode_cudnn_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor @@ -789,12 +790,12 @@ def canonicalize_device(device: Optional[torch.device | str]) -> torch.device: if device is None: # Use default CUDA device device = torch.get_default_device() - if device.type != "cuda": - device = torch.device("cuda", torch.cuda.current_device()) + if device.type != te_device_type(): + device = torch.device(te_device_type(), torch.cuda.current_device()) elif not isinstance(device, torch.device): device = torch.device(device) - if device.type == "cuda" and device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) + if device.type == te_device_type() and device.index is None: + device = torch.device(te_device_type(), torch.cuda.current_device()) return device @@ -816,7 +817,7 @@ def devices_match(device1: torch.device, device2: torch.device) -> bool: device2 = torch.device(device2) if device1.type != device2.type: return False - if device1.type == "cuda": + if device1.type == te_device_type(): index1 = device1.index index2 = device2.index if index1 == index2: @@ -952,15 +953,22 @@ def canonicalize_process_group( return group +def torch_is_autocast_gpu_enabled() -> bool: + """Check whether PyTorch autocast is enabled for the GPU device.""" + if torch_version() >= (2, 4, 0): + return torch.is_autocast_enabled(te_device_type()) + return torch.is_autocast_enabled() + + def torch_get_autocast_gpu_dtype() -> torch.dtype: """Get PyTorch autocast GPU dtype.""" if torch_version() >= (2, 4, 0): - return torch.get_autocast_dtype("cuda") + return torch.get_autocast_dtype(te_device_type()) return torch.get_autocast_gpu_dtype() if torch_version() >= (2, 4, 0): - gpu_autocast_ctx = functools.partial(torch.amp.autocast, device_type="cuda") + gpu_autocast_ctx = functools.partial(torch.amp.autocast, device_type=te_device_type()) else: gpu_autocast_ctx = torch.cuda.amp.autocast From 268fdd54d39ba401d69a2fc3fe1ccd37565aa2fb Mon Sep 17 00:00:00 2001 From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:59:59 +0530 Subject: [PATCH 13/32] Guard fused MoE auxiliary loss integer offsets (#3569) * Guard fused MoE auxiliary loss integer offsets Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> * Update tests/pytorch/test_fused_router.py Co-authored-by: Alp Dener Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> * Use 64-bit indexing for fused MoE auxiliary loss Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> * Unroll forward row loop in fused MoE aux loss Signed-off-by: arcusbuilds * Revert 64-bit indexing, restore INT_MAX guard in fused MoE aux loss Match the other fused router kernels, which reject num_tokens * num_experts > INT_MAX. Keep #pragma unroll 4 on the forward row loops. Signed-off-by: arcusbuilds * Use 64-bit loop counters in fused MoE aux loss kernels Shapes at the INT_MAX guard (e.g. INT_MAX x 1) overflowed the int loop counters on their last increment and the backward grid-size math. Flattened offsets stay 32-bit, which the guard bounds. Signed-off-by: arcusbuilds --------- Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> Signed-off-by: arcusbuilds Co-authored-by: Alp Dener --- tests/pytorch/test_fused_router.py | 89 +++++++++++++++++++ .../common/fused_router/fused_moe_aux_loss.cu | 37 +++++--- 2 files changed, 114 insertions(+), 12 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index ac2f46826df..35dc510388b 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1041,6 +1041,95 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl torch.testing.assert_close(probs.grad, probs_clone.grad, atol=atol, rtol=rtol) +@pytest.mark.parametrize("path", ["forward", "graph_safe_forward", "backward"]) +def test_fused_moe_aux_loss_rejects_int_offset_overflow(path): + num_rows, num_cols = 8_388_609, 256 # Above INT_MAX elements. + bytes_needed = num_rows * num_cols * 2 + torch.cuda.empty_cache() + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < bytes_needed + (1 << 30): + pytest.skip("Needs 4.0 GiB for BF16 input/output plus 1 GiB headroom") + + tokens_per_expert = torch.ones(num_cols, device="cuda", dtype=torch.int32) + try: + if path == "backward": + tex.fused_moe_aux_loss_bwd( + Const_buf=torch.ones(2, device="cuda", dtype=torch.float32), + tokens_per_expert=tokens_per_expert, + num_rows=num_rows, + num_cols=num_cols, + grad_aux_loss=torch.ones((), device="cuda", dtype=torch.bfloat16), + ) + else: + probs = torch.empty((num_rows, num_cols), device="cuda", dtype=torch.bfloat16) + arguments = dict( + probs=probs, + tokens_per_expert=tokens_per_expert, + num_experts=num_cols, + num_rows=num_rows, + num_cols=num_cols, + topk=1, + coeff=0.01, + ) + if path == "forward": + tex.fused_moe_aux_loss_fwd(total_num_tokens=num_rows, **arguments) + else: + total = torch.tensor(num_rows, device="cuda", dtype=torch.int64) + tex.fused_moe_aux_loss_fwd_graph_safe(total_num_tokens=total, **arguments) + except torch.cuda.OutOfMemoryError: + pytest.skip("Could not allocate the BF16 input/output tensor") + except RuntimeError as exc: + assert "num_rows * num_cols exceeds INT_MAX" in str(exc) + else: + pytest.fail("Expected the int-offset guard to reject the shape") + + +@pytest.mark.parametrize("path", ["forward", "graph_safe_forward", "backward"]) +def test_fused_moe_aux_loss_largest_accepted_shape(path): + # Exactly INT_MAX elements with one column: loop counters and grid math must not overflow. + num_rows, num_cols = 2**31 - 1, 1 + bytes_needed = num_rows * num_cols * 2 + torch.cuda.empty_cache() + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < bytes_needed + (1 << 30): + pytest.skip("Needs 4.0 GiB for BF16 input/output plus 1 GiB headroom") + + tokens_per_expert = torch.ones(num_cols, device="cuda", dtype=torch.int32) + try: + if path == "backward": + grad_probs = tex.fused_moe_aux_loss_bwd( + Const_buf=torch.tensor([0.25, 0.0], device="cuda", dtype=torch.float32), + tokens_per_expert=tokens_per_expert, + num_rows=num_rows, + num_cols=num_cols, + grad_aux_loss=torch.tensor(2.0, device="cuda", dtype=torch.bfloat16), + ) + for start in range(0, num_rows, 1 << 24): + assert torch.all(grad_probs[start : start + (1 << 24)] == 0.5) + else: + probs = torch.zeros((num_rows, num_cols), device="cuda", dtype=torch.bfloat16) + probs[-2:] = 1 # Nonzero only where the row counter is closest to INT_MAX. + arguments = dict( + probs=probs, + tokens_per_expert=tokens_per_expert, + num_experts=num_cols, + num_rows=num_rows, + num_cols=num_cols, + topk=1, + coeff=1.0, + ) + if path == "forward": + aux_loss, _ = tex.fused_moe_aux_loss_fwd(total_num_tokens=1, **arguments) + else: + total = torch.tensor(1, device="cuda", dtype=torch.int64) + aux_loss, _ = tex.fused_moe_aux_loss_fwd_graph_safe( + total_num_tokens=total, **arguments + ) + torch.testing.assert_close(aux_loss, probs.new_tensor(2.0), atol=0, rtol=0) + except torch.cuda.OutOfMemoryError: + pytest.skip("Could not allocate the BF16 input/output tensor") + + def test_fused_moe_aux_loss_cuda_graph_capture(): """CUDA-graph-safe path: total_num_tokens is a device tensor whose value changes between replays. Forward and backward must both observe the new diff --git a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu index 9fb5adcc1f4..238f834f5b9 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -37,12 +37,13 @@ __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, // Grid-stride over rows so that every row is processed exactly once. // Each thread processes a subset of columns. - for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { + for (int64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { CompType col_sum = CompType(0); // Accumulate probs over the rows assigned to this CTA (grid-stride). - for (int row = blockIdx.x; row < num_rows; row += gridDim.x) { - col_sum += CompType(probs[row * num_cols + col]); +#pragma unroll 4 + for (int64_t row = blockIdx.x; row < num_rows; row += gridDim.x) { + col_sum += CompType(probs[static_cast(row) * num_cols + static_cast(col)]); } // Multiply by the token count for this expert. @@ -89,6 +90,9 @@ void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, cudaStream_t stream) { NVTE_CHECK(num_cols > 0, "num_cols must be positive, got ", num_cols); NVTE_CHECK(num_experts > 0, "num_experts must be positive, got ", num_experts); + NVTE_CHECK(static_cast(num_rows) * num_cols <= INT_MAX, + "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", + static_cast(num_rows) * num_cols); // Sequence aux loss batches independent sequences along the expert dimension. NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols, ") must be a multiple of number of experts (", num_experts, ")."); @@ -144,10 +148,11 @@ __global__ void fused_moe_aux_loss_forward_kernel_graph_safe( int num_experts, int num_rows, int num_cols, int topk, float coeff, float* Coeff_buf) { // Reduction body matches the scalar-input kernel above. CompType thread_sum = CompType(0); - for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { + for (int64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { CompType col_sum = CompType(0); - for (int row = blockIdx.x; row < num_rows; row += gridDim.x) { - col_sum += CompType(probs[row * num_cols + col]); +#pragma unroll 4 + for (int64_t row = blockIdx.x; row < num_rows; row += gridDim.x) { + col_sum += CompType(probs[static_cast(row) * num_cols + static_cast(col)]); } col_sum *= CompType(tokens_per_expert[col]); thread_sum += col_sum; @@ -182,6 +187,9 @@ void fused_moe_aux_loss_forward_kernel_launcher_graph_safe( float* Coeff_buf, cudaStream_t stream) { NVTE_CHECK(num_cols > 0, "num_cols must be positive, got ", num_cols); NVTE_CHECK(num_experts > 0, "num_experts must be positive, got ", num_experts); + NVTE_CHECK(static_cast(num_rows) * num_cols <= INT_MAX, + "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", + static_cast(num_rows) * num_cols); NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols, ") must be a multiple of number of experts (", num_experts, ")."); @@ -231,18 +239,20 @@ __global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf, const IndexType* tokens_per_expert, int num_rows, int num_cols, DataType* grad_aux_loss, DataType* grad_probs) { - int global_warp_num = gridDim.x * blockDim.x / kThreadsPerWarp; - int global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / kThreadsPerWarp; + int64_t global_warp_num = static_cast(gridDim.x) * blockDim.x / kThreadsPerWarp; + int64_t global_warp_id = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / kThreadsPerWarp; int lane_id = threadIdx.x % kThreadsPerWarp; // Loop: for all positions in each row - for (int i = lane_id; i < num_cols; i += kThreadsPerWarp) { + for (int64_t i = lane_id; i < num_cols; i += kThreadsPerWarp) { float C_coeff = Const_buf[0]; CompType tokens_per_expert_i = static_cast(tokens_per_expert[i]); CompType grad_aux_loss_value = static_cast(grad_aux_loss[0]); // Loop: for all rows - for (int j = global_warp_id; j < num_rows; j += global_warp_num) { - grad_probs[j * num_cols + i] = C_coeff * tokens_per_expert_i * grad_aux_loss_value; + for (int64_t j = global_warp_id; j < num_rows; j += global_warp_num) { + grad_probs[static_cast(j) * num_cols + static_cast(i)] = + C_coeff * tokens_per_expert_i * grad_aux_loss_value; } } } @@ -252,9 +262,12 @@ void fused_moe_aux_loss_backward_kernel_launcher(const float* Const_buf, const IndexType* tokens_per_expert, int num_rows, int num_cols, DataType* grad_aux_loss, DataType* grad_probs, cudaStream_t stream) { + NVTE_CHECK(static_cast(num_rows) * num_cols <= INT_MAX, + "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", + static_cast(num_rows) * num_cols); // Meta data for the kernel int block_size = 256; - int grid_size = (num_rows + block_size - 1) / block_size; + int grid_size = (static_cast(num_rows) + block_size - 1) / block_size; fused_moe_aux_loss_backward_kernel<<>>( Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss, grad_probs); NVTE_CHECK_CUDA(cudaGetLastError()); From 9b7e959fe47137e4bb1abf2294eb145c763f40c2 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 27 Oct 2026 15:35:13 -0700 Subject: [PATCH 14/32] Rely on Recipe object for fusion (#3717) * rely on recipe object for fusion, reduce cpu overhedads for deciding the grouped tensor path Signed-off-by: Varun Thumbe * Keep fusion decisions local to op fuser Signed-off-by: Varun Thumbe * Use recipe representation for fusion cache key Signed-off-by: Varun Thumbe * Include CUDA device in grouped tensor support cache Signed-off-by: Varun Thumbe * Declare derived quantizer settings as recipe fields Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/ops/fuser.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/module/grouped_linear.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/ops/basic/grouped_linear.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- tests/pytorch/test_fusible_ops.py | 57 +++++++++++++++++-- transformer_engine/common/recipe/__init__.py | 39 ++++++++++++- .../pytorch/module/grouped_linear.py | 53 +++++++++++++---- .../pytorch/ops/basic/grouped_linear.py | 51 +++++++++++++---- transformer_engine/pytorch/ops/fuser.py | 10 ++-- 5 files changed, 177 insertions(+), 33 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index d4455c09d68..17efba6f515 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -272,6 +272,13 @@ def track_first_forward(): assert first_forward_calls == 1 assert fusion_calls == 1 + # A recipe representation change selects a new plan without resetting the + # op's quantizers when the changed setting does not affect their state. + equivalent_current_scaling.fp8_dpa = True + fuser.maybe_fuse_ops(False, equivalent_current_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling] + assert fusion_calls == 2 + # Backward override affects both recipe state and fusion topology, so it # triggers one reset and constructs a distinct cached plan. overridden_current_scaling = transformer_engine.common.recipe.Float8CurrentScaling( @@ -280,7 +287,7 @@ def track_first_forward(): fuser.maybe_fuse_ops(False, overridden_current_scaling, x, extra_inputs) assert reset_recipes == [current_scaling, overridden_current_scaling] assert first_forward_calls == 1 - assert fusion_calls == 2 + assert fusion_calls == 3 delayed_scaling = transformer_engine.common.recipe.DelayedScaling( amax_history_len=8, @@ -289,10 +296,10 @@ def track_first_forward(): fuser.maybe_fuse_ops(False, delayed_scaling, x, extra_inputs) assert reset_recipes == [current_scaling, overridden_current_scaling, delayed_scaling] assert first_forward_calls == 1 - assert fusion_calls == 3 + assert fusion_calls == 4 - # Amax history length only affects delayed-scaling recipe state. Reset that - # state, but restore the existing DelayedScaling fusion plan from the cache. + # Amax history length changes the recipe representation, so reset the + # delayed-scaling state and build a new fusion plan. resized_delayed_scaling = transformer_engine.common.recipe.DelayedScaling( amax_history_len=16, backward_override=None, @@ -305,13 +312,51 @@ def track_first_forward(): resized_delayed_scaling, ] assert first_forward_calls == 1 - assert fusion_calls == 3 + assert fusion_calls == 5 # Repeating the exact recipe parameters performs neither operation again. fuser.maybe_fuse_ops(False, resized_delayed_scaling, x, extra_inputs) assert len(reset_recipes) == 4 assert first_forward_calls == 1 - assert fusion_calls == 3 + assert fusion_calls == 5 + + +def test_operation_fuser_rebuilds_for_extra_recipe_attribute(monkeypatch) -> None: + """A monkey-patched recipe flag must select a new fusion plan.""" + + seen_flags = [] + + def track_fusion(ops, *, recipe): + seen_flags.append(getattr(recipe, "enable_cutedsl_fused_grouped_mlp", False)) + return ops + + monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion]) + monkeypatch.setattr(OperationFuser, "forward_fusion_functions", []) + monkeypatch.setattr(OperationFuser, "backward_fusion_functions", []) + + op = te_ops.Identity() + monkeypatch.setattr(op, "reset_recipe_state", lambda **_: None) + monkeypatch.setattr(op, "pre_first_fuser_forward", lambda: None) + fuser = OperationFuser([op]) + recipe = transformer_engine.common.recipe.CustomRecipe(qfactory=lambda role: None) + x = torch.ones(1) + extra_inputs = [()] + + fuser.maybe_fuse_ops(False, recipe, x, extra_inputs) + assert seen_flags == [False] + + recipe.enable_cutedsl_fused_grouped_mlp = True + fuser.maybe_fuse_ops(False, recipe, x, extra_inputs) + fuser.maybe_fuse_ops(False, recipe, x, extra_inputs) + assert seen_flags == [False, True] + + recipe.enable_cutedsl_fused_grouped_mlp = False + fuser.maybe_fuse_ops(False, recipe, x, extra_inputs) + assert seen_flags == [False, True, False] + + del recipe.enable_cutedsl_fused_grouped_mlp + fuser.maybe_fuse_ops(False, recipe, x, extra_inputs) + assert seen_flags == [False, True, False] @torch.no_grad() diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 456500523ee..258504f9443 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -118,11 +118,28 @@ class Recipe: _cached_repr: Optional[str] = None def __setattr__(self, name: str, value: Any) -> None: - # Invalidate the cached repr on any attribute mutation. + # Invalidate the cached repr on any attribute mutation. Track public + # attributes added outside the recipe's declared dataclass fields. if name != "_cached_repr": object.__setattr__(self, "_cached_repr", None) + if not name.startswith("_") and name not in getattr( + type(self), "__dataclass_fields__", {} + ): + extra_attrs = self.__dict__.get("_extra_attrs") + if extra_attrs is None: + extra_attrs = set() + object.__setattr__(self, "_extra_attrs", extra_attrs) + extra_attrs.add(name) object.__setattr__(self, name, value) + def __delattr__(self, name: str) -> None: + object.__delattr__(self, name) + if name != "_cached_repr": + object.__setattr__(self, "_cached_repr", None) + extra_attrs = self.__dict__.get("_extra_attrs") + if extra_attrs is not None: + extra_attrs.discard(name) + @abc.abstractmethod def _make_repr(self) -> str: """Build the string representation for this recipe. @@ -133,7 +150,15 @@ def _make_repr(self) -> str: def __repr__(self) -> str: if self._cached_repr is None: - self._cached_repr = self._make_repr() + recipe_repr = self._make_repr() + extra_attrs = self.__dict__.get("_extra_attrs") + if extra_attrs: + extras = ", ".join( + f"{name}={getattr(self, name)!r}" for name in sorted(extra_attrs) + ) + separator = "" if recipe_repr.endswith(", ") else ", " + recipe_repr += f"{separator}{extras}" + self._cached_repr = recipe_repr return self._cached_repr @classmethod @@ -310,6 +335,9 @@ class Float8CurrentScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + fp8_quant_fwd_inp: QParams = field(init=False, repr=False, compare=False) + fp8_quant_fwd_weight: QParams = field(init=False, repr=False, compare=False) + fp8_quant_bwd_grad: QParams = field(init=False, repr=False, compare=False) def __post_init__(self) -> None: assert self.fp8_format in (Format.E4M3, Format.HYBRID), "Unsupported FP8 format." @@ -447,6 +475,9 @@ class Float8BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + fp8_quant_fwd_inp: QParams = field(init=False, repr=False, compare=False) + fp8_quant_fwd_weight: QParams = field(init=False, repr=False, compare=False) + fp8_quant_bwd_grad: QParams = field(init=False, repr=False, compare=False) def __post_init__(self) -> None: assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" @@ -586,6 +617,9 @@ class NVFP4BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + fp4_quant_fwd_inp: QParams = field(init=False, repr=False, compare=False) + fp4_quant_fwd_weight: QParams = field(init=False, repr=False, compare=False) + fp4_quant_bwd_grad: QParams = field(init=False, repr=False, compare=False) def __post_init__(self) -> None: assert self.fp4_format == Format.E2M1, "Only E2M1 is supported for NVFP4 scaling" @@ -633,6 +667,7 @@ def _make_repr(self) -> str: f"fp8_dpa={self.fp8_dpa}, " f"fp8_mha={self.fp8_mha}, " f"backward_override={self.backward_override}, " + f"disable_rht={self.disable_rht}, " f"row_scaled_activation={self.row_scaled_activation}, " f"nvfp4_4over6={self.nvfp4_4over6}, " f"nvfp4_4over6_e4m3_use_256={self.nvfp4_4over6_e4m3_use_256}, " diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index fc595357b13..76eded6fa84 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -92,6 +92,7 @@ def is_module_grouped_tensor_path_supported( recipe: Optional[Recipe], dtype: torch.dtype, + device: Optional[torch.device] = None, ) -> bool: """Whether the module grouped-tensor path supports this recipe and dtype. @@ -121,14 +122,43 @@ def is_module_grouped_tensor_path_supported( Runtime-only restrictions such as debug mode, CPU offloading, calibration, output quantization, and backend selection are checked separately by ``GroupedLinear``. + + If device is omitted, the current CUDA device is used. """ + # Recipe objects are not hashable. Resolve the device before the cached call so + # support results from different GPUs cannot share a cache entry. + device_index = device.index if device is not None else None + if device_index is None: + device_index = torch.cuda.current_device() + return _is_module_grouped_tensor_path_supported( + type(recipe) if recipe is not None else None, + getattr(recipe, "backward_override", None), + getattr(recipe, "disable_rht", None), + getattr(recipe, "row_scaled_activation", None), + getattr(recipe, "fp8_format", None), + dtype, + device_index, + ) + + +@functools.lru_cache +def _is_module_grouped_tensor_path_supported( + recipe_type: Optional[type[Recipe]], + backward_override: Optional[str], + disable_rht: Optional[bool], + row_scaled_activation: Optional[bool], + fp8_format: Optional[RecipeFormat], + dtype: torch.dtype, + device_index: int, +) -> bool: + """Cache routing for an immutable recipe snapshot, dtype, and CUDA device.""" if tex.get_cublasLt_build_version() < 130300: return False if dtype not in (torch.bfloat16, torch.float16): return False - device_capability = get_device_compute_capability() + device_capability = get_device_compute_capability(torch.device("cuda", device_index)) if not (9, 0) <= device_capability <= (11, 0): return False cublaslt_version = tex.get_cublasLt_version() @@ -137,27 +167,27 @@ def is_module_grouped_tensor_path_supported( if device_capability < (10, 0) and cublaslt_version < 130400: return False - if recipe is None: + if recipe_type is None: return True - if recipe.custom(): + if recipe_type.custom(): return False - if recipe.backward_override is not None: + if backward_override is not None: return False - if recipe.float8_current_scaling(): + if recipe_type.float8_current_scaling(): return device_capability >= (10, 0) or cublaslt_version >= 130500 - if recipe.float8_block_scaling(): + if recipe_type.float8_block_scaling(): # Blackwell and newer emulate FP8 block scaling with MXFP8 grouped GEMM, so it is # supported wherever MXFP8 is. On Hopper, cuBLASLt 13.6 fixes grouped GEMM algo # selection for block-scaled FP8. return device_capability >= (10, 0) or cublaslt_version >= 130600 - if recipe.mxfp8(): + if recipe_type.mxfp8(): return device_capability >= (10, 0) - if recipe.nvfp4(): + if recipe_type.nvfp4(): return ( device_capability >= (10, 0) - and not recipe.disable_rht - and not recipe.row_scaled_activation - and recipe.fp8_format == RecipeFormat.E4M3 + and not disable_rht + and not row_scaled_activation + and fp8_format == RecipeFormat.E4M3 ) return False @@ -2406,6 +2436,7 @@ def forward( use_grouped_tensor_path = is_module_grouped_tensor_path_supported( _recipe, self.activation_dtype, + device=inp.device, ) if ( self.use_grouped_tensor diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 0dfc8e2366a..fe5e77fc580 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -79,6 +79,7 @@ def is_op_fuser_grouped_tensor_path_supported( recipe: Optional[Recipe], dtype: torch.dtype, split_alignment: int = 128, + device: Optional[torch.device] = None, ) -> bool: """Whether the op-fuser grouped-tensor path supports this recipe and dtype. @@ -105,11 +106,40 @@ def is_op_fuser_grouped_tensor_path_supported( is excluded because cuBLASLt grouped GEMM does not support it. * Single grouped parameters have no split-quantize fallback, so callers must reject them when this function returns ``False``. + + If device is omitted, the current CUDA device is used. """ + # Recipe objects are not hashable. Resolve the device before the cached call so + # support results from different GPUs cannot share a cache entry. + device_index = device.index if device is not None else None + if device_index is None: + device_index = torch.cuda.current_device() + return _is_op_fuser_grouped_tensor_path_supported( + type(recipe) if recipe is not None else None, + getattr(recipe, "disable_rht", None), + getattr(recipe, "row_scaled_activation", None), + getattr(recipe, "fp8_format", None), + dtype, + split_alignment, + device_index, + ) + + +@functools.lru_cache +def _is_op_fuser_grouped_tensor_path_supported( + recipe_type: Optional[type[Recipe]], + disable_rht: Optional[bool], + row_scaled_activation: Optional[bool], + fp8_format: Optional[RecipeFormat], + dtype: torch.dtype, + split_alignment: int, + device_index: int, +) -> bool: + """Cache routing for a recipe snapshot, dtype, alignment, and CUDA device.""" if dtype not in (torch.bfloat16, torch.float16): return False - device_capability = get_device_compute_capability() + device_capability = get_device_compute_capability(torch.device("cuda", device_index)) if not (9, 0) <= device_capability <= (11, 0): return False cublaslt_version = tex.get_cublasLt_version() @@ -118,27 +148,27 @@ def is_op_fuser_grouped_tensor_path_supported( if device_capability < (10, 0) and cublaslt_version < 130400: return False - if recipe is None: + if recipe_type is None: return True - if recipe.custom(): + if recipe_type.custom(): return False - if recipe.float8_current_scaling(): + if recipe_type.float8_current_scaling(): return device_capability >= (10, 0) or cublaslt_version >= 130500 - if recipe.float8_block_scaling(): + if recipe_type.float8_block_scaling(): # Blackwell and newer emulate FP8 block scaling with MXFP8 grouped GEMM, so it is # supported wherever MXFP8 is. On Hopper, cuBLASLt 13.6 fixes grouped GEMM algo # selection for block-scaled FP8. return ( device_capability >= (10, 0) or cublaslt_version >= 130600 ) and split_alignment % 128 == 0 - if recipe.mxfp8(): + if recipe_type.mxfp8(): return device_capability >= (10, 0) and split_alignment % 128 == 0 - if recipe.nvfp4(): + if recipe_type.nvfp4(): return ( device_capability >= (10, 0) - and not recipe.disable_rht - and not recipe.row_scaled_activation - and recipe.fp8_format == RecipeFormat.E4M3 + and not disable_rht + and not row_scaled_activation + and fp8_format == RecipeFormat.E4M3 and split_alignment % 128 == 0 ) return False @@ -1085,6 +1115,7 @@ def fuser_forward( recipe, dtype, self.split_alignment, + device=device, ) if (self.single_grouped_weight or self.single_grouped_bias) and not use_grouped_tensor_path: raise RuntimeError( diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 5c7371ce644..6d13ba43942 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -51,7 +51,7 @@ def _is_graph_capturing() -> bool: "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]" ) _FusedOpList: TypeAlias = list[tuple[FusibleOperation, list[int]]] -_FusionParams: TypeAlias = tuple[type, int, Optional[str]] +_FusionParams: TypeAlias = tuple[str, int] class _OperationFuserAutogradFunction(torch.autograd.Function): @@ -669,8 +669,7 @@ def maybe_fuse_ops( # paths that reuse a cached fused operation configuration. self.first_op_requiring_backward = first_op_requiring_backward - # Check if recipe parameters don't match cached values. In this case, - # the recipe state in the basic ops might be invalid, so reset it. + # Check for changes in recipe parameters that might have invalidated recipe state in basic ops recipe_type = type(recipe) need_to_reset_recipe_state = self.recipe_type != recipe_type @@ -702,7 +701,10 @@ def maybe_fuse_ops( # backward boundary in the key, but pay construction cost only once for # each configuration. Full recompute therefore builds at most one # no-grad plan and one grad-enabled plan for a stable recipe. - fusion_params = (recipe_type, first_op_requiring_backward, backward_override) + fusion_params: _FusionParams = ( + str(recipe) if recipe is not None else "", + first_op_requiring_backward, + ) cached_ops = self._fused_ops_cache.get(fusion_params) if cached_ops is not None: self._forward_ops, self._backward_ops = cached_ops From e614b7a9d3c3082ab3ae38367ef1a39fb0b34377 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Wed, 28 Oct 2026 11:50:20 -0700 Subject: [PATCH 15/32] [Common] Optimized kernels for swizzling and remove obsolete narrow scaled kernels (#3682) Optimize scale swizzle kernels Use optimized rowwise full-width and columnwise chunk kernels for MXFP8 and NVFP4 scales, including multi-tensor swizzling. Keep the 16-byte scale-buffer alignment contract and cover its validation in C++ tests. Choose the rowwise full-width kernel after the fallback grid exceeds one SM wave, based on paired GB200 and VR200 measurements. Signed-off-by: Varun Thumbe --- docs/envvars.rst | 6 - tests/cpp/operator/test_multi_swizzle.cu | 224 +++- tests/cpp/operator/test_swizzle.cu | 196 +++- .../include/transformer_engine/swizzle.h | 3 + transformer_engine/common/swizzle/swizzle.cu | 1044 +++++------------ 5 files changed, 676 insertions(+), 797 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 09e2b462780..5ff430eabea 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -351,12 +351,6 @@ Kernel Configuration :Default: ``0`` :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. NVFP4 4over6 quantization, fused softmax, and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_NVFP4``, ``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_MXFP8_PER_TILE_SWIZZLE - - :Type: ``int`` (0 or 1) - :Default: ``0`` - :Description: Swizzle MXFP8 scaling factors for GEMM with kernels that give each 128-row tile of scaling factors (row-wise) or each 128-column by 32-row block (column-wise) its own thread block, instead of batching 32 tiles in one block. This lowers the latency of swizzling the scaling factors of tensors with few tiles, such as activations. - .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM :Type: ``int`` (0 or 1) diff --git a/tests/cpp/operator/test_multi_swizzle.cu b/tests/cpp/operator/test_multi_swizzle.cu index 4bd1800cea2..c0ac1501391 100644 --- a/tests/cpp/operator/test_multi_swizzle.cu +++ b/tests/cpp/operator/test_multi_swizzle.cu @@ -51,18 +51,25 @@ void compute_ref_swizzle(const uint8_t *h_input, uint8_t *h_output, } } -static void zero_scale_inv_padding(uint8_t *buf, - size_t padded_rows, size_t padded_cols, - size_t orig_rows, size_t orig_cols) { +static void set_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols, + uint8_t value) { for (size_t r = 0; r < padded_rows; ++r) { for (size_t c = 0; c < padded_cols; ++c) { if (r >= orig_rows || c >= orig_cols) { - buf[r * padded_cols + c] = 0; + buf[r * padded_cols + c] = value; } } } } +static void zero_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols) { + set_scale_inv_padding(buf, padded_rows, padded_cols, orig_rows, orig_cols, 0); +} + // =================================================================== // Multi-tensor swizzle test // =================================================================== @@ -93,12 +100,12 @@ void performTestMultiTensorSwizzle(const int num_tensors, const size_t M, const const NVTEShape rs = input->rowwise_scale_inv_shape(); zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); } else { const NVTEShape cs = input->columnwise_scale_inv_shape(); zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); } input->from_cpu(); @@ -141,6 +148,159 @@ void performTestMultiTensorSwizzle(const int num_tensors, const size_t M, const } } +void performTestMultiTensorSwizzleNVFP4(const int num_tensors, const size_t M, const size_t K, + bool rowwise) { + using namespace test; + constexpr size_t BLOCK_SIZE = 16; + std::vector> inputs; + std::vector> outputs; + std::vector input_handles; + std::vector output_handles; + + for (int i = 0; i < num_tensors; ++i) { + auto input = std::make_unique("input_" + std::to_string(i), + std::vector{M, K}, DType::kFloat4E2M1, rowwise, + !rowwise, NVTE_NVFP4_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), + std::vector{M, K}, DType::kFloat4E2M1, rowwise, + !rowwise, NVTE_NVFP4_1D_SCALING); + output->set_with_gemm_swizzled_scales(true); + if (rowwise) { + input->fill_uniform_rowwise_scale_inv(); + } else { + input->fill_uniform_columnwise_scale_inv(); + } + const NVTEShape scale_shape = + rowwise ? input->rowwise_scale_inv_shape() : input->columnwise_scale_inv_shape(); + uint8_t *input_scales = rowwise ? input->rowwise_cpu_scale_inv_ptr() + : input->columnwise_cpu_scale_inv_ptr(); + const size_t original_rows = rowwise ? M : K; + const size_t original_cols = + rowwise ? divide_round_up(K, BLOCK_SIZE) : divide_round_up(M, BLOCK_SIZE); + set_scale_inv_padding(input_scales, scale_shape.data[0], scale_shape.data[1], original_rows, + original_cols, 0xFF); + input->from_cpu(); + input_handles.push_back(input->data()); + output_handles.push_back(output->data()); + inputs.emplace_back(std::move(input)); + outputs.emplace_back(std::move(output)); + } + + nvte_multi_tensor_swizzle_scaling_factors(input_handles.data(), output_handles.data(), + num_tensors, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + for (int i = 0; i < num_tensors; ++i) { + outputs[i]->to_cpu(); + const NVTEShape scale_shape = + rowwise ? inputs[i]->rowwise_scale_inv_shape() : inputs[i]->columnwise_scale_inv_shape(); + const uint8_t *input_scales = rowwise ? inputs[i]->rowwise_cpu_scale_inv_ptr() + : inputs[i]->columnwise_cpu_scale_inv_ptr(); + const uint8_t *actual = rowwise ? outputs[i]->rowwise_cpu_scale_inv_ptr() + : outputs[i]->columnwise_cpu_scale_inv_ptr(); + std::vector reference_input( + input_scales, input_scales + scale_shape.data[0] * scale_shape.data[1]); + const size_t original_rows = rowwise ? M : K; + const size_t original_cols = + rowwise ? divide_round_up(K, BLOCK_SIZE) : divide_round_up(M, BLOCK_SIZE); + zero_scale_inv_padding(reference_input.data(), scale_shape.data[0], scale_shape.data[1], + original_rows, original_cols); + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, true>(reference_input.data(), reference.data(), scale_shape.data[0], + scale_shape.data[1]); + compareResults("multi_nvfp4_swizzle", actual, reference.data(), reference.size()); + } +} + +TEST(MultiTensorSwizzleTest, MXFP8ColumnwiseBatchIgnoresUnusedRowwiseBuffer) { + using namespace test; + constexpr int num_tensors = 2; + const std::vector shape{256, 128}; + std::vector> inputs; + std::vector> outputs; + std::vector input_handles; + std::vector output_handles; + + for (int i = 0; i < num_tensors; ++i) { + const bool has_rowwise_scales = i == 0; + auto input = std::make_unique("input_" + std::to_string(i), shape, + DType::kFloat8E4M3, has_rowwise_scales, true, + NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), shape, + DType::kFloat8E4M3, has_rowwise_scales, true, + NVTE_MXFP8_1D_SCALING); + if (has_rowwise_scales) input->fill_uniform_rowwise_scale_inv(); + input->fill_uniform_columnwise_scale_inv(); + output->set_with_gemm_swizzled_scales(true); + input->from_cpu(); + input_handles.push_back(input->data()); + output_handles.push_back(output->data()); + inputs.emplace_back(std::move(input)); + outputs.emplace_back(std::move(output)); + } + + nvte_multi_tensor_swizzle_scaling_factors(input_handles.data(), output_handles.data(), + num_tensors, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + for (int i = 0; i < num_tensors; ++i) { + outputs[i]->to_cpu(); + const NVTEShape scale_shape = inputs[i]->columnwise_scale_inv_shape(); + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, false>( + inputs[i]->columnwise_cpu_scale_inv_ptr(), reference.data(), scale_shape.data[1], + scale_shape.data[0]); + compareResults("multi_mixed_buffers_columnwise_" + std::to_string(i), + outputs[i]->columnwise_cpu_scale_inv_ptr(), reference.data(), + reference.size()); + } +} + +TEST(MultiTensorSwizzleTest, NVFP4FullWidthChunkLaunchSplitting) { + for (const bool rowwise : {false, true}) { + performTestMultiTensorSwizzleNVFP4(4, 160, 96, rowwise); + performTestMultiTensorSwizzleNVFP4(4, 128, 128, rowwise); + performTestMultiTensorSwizzleNVFP4(65, 128, 128, rowwise); + } +} + +TEST(MultiTensorSwizzleTest, RowwiseMixedSharedMemoryFallback) { + using namespace test; + // The first tensor fits the full-width shared-memory limit, but the second does not. + const std::vector> shapes{{128, 7168}, {128, 8064}}; + std::vector> inputs; + std::vector> outputs; + std::vector input_handles; + std::vector output_handles; + for (size_t i = 0; i < shapes.size(); ++i) { + const auto [M, K] = shapes[i]; + auto input = std::make_unique("input_" + std::to_string(i), + std::vector{M, K}, DType::kFloat8E4M3, true, + false, NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), + std::vector{M, K}, DType::kFloat8E4M3, true, + false, NVTE_MXFP8_1D_SCALING); + output->set_with_gemm_swizzled_scales(true); + input->fill_uniform_rowwise_scale_inv(); + input->from_cpu(); + input_handles.push_back(input->data()); + output_handles.push_back(output->data()); + inputs.emplace_back(std::move(input)); + outputs.emplace_back(std::move(output)); + } + + nvte_multi_tensor_swizzle_scaling_factors( + input_handles.data(), output_handles.data(), static_cast(inputs.size()), 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + for (size_t i = 0; i < inputs.size(); ++i) { + outputs[i]->to_cpu(); + const NVTEShape scale_shape = inputs[i]->rowwise_scale_inv_shape(); + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, true>(inputs[i]->rowwise_cpu_scale_inv_ptr(), + reference.data(), scale_shape.data[0], scale_shape.data[1]); + compareResults("multi_rowwise_fallback", outputs[i]->rowwise_cpu_scale_inv_ptr(), + reference.data(), reference.size()); + } +} + // =================================================================== // Multi-tensor unswizzle test (uses single-tensor swizzle to prepare) // =================================================================== @@ -172,12 +332,12 @@ void performTestMultiTensorUnswizzle(const int num_tensors, const size_t M, cons const NVTEShape rs = orig->rowwise_scale_inv_shape(); zero_scale_inv_padding(orig->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); } else { const NVTEShape cs = orig->columnwise_scale_inv_shape(); zero_scale_inv_padding(orig->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); } orig->from_cpu(); @@ -249,12 +409,12 @@ void performTestMultiTensorRoundtrip(const int num_tensors, const size_t M, cons const NVTEShape rs = orig->rowwise_scale_inv_shape(); zero_scale_inv_padding(orig->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); } else { const NVTEShape cs = orig->columnwise_scale_inv_shape(); zero_scale_inv_padding(orig->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); } orig->from_cpu(); @@ -300,22 +460,11 @@ void performTestMultiTensorRoundtrip(const int num_tensors, const size_t M, cons // Test suites // =================================================================== -// Sets NVTE_MXFP8_PER_TILE_SWIZZLE while it lives, so that MXFP8 swizzles use the per-tile -// kernels (enable) or the default ones. -class PerTileSwizzle { - public: - explicit PerTileSwizzle(bool enable) { - setenv("NVTE_MXFP8_PER_TILE_SWIZZLE", enable ? "1" : "0", 1); - } - ~PerTileSwizzle() { unsetenv("NVTE_MXFP8_PER_TILE_SWIZZLE"); } -}; - class MultiTensorSwizzleTestSuite - : public ::testing::TestWithParam, bool>> {}; + : public ::testing::TestWithParam> {}; TEST_P(MultiTensorSwizzleTestSuite, TestMultiTensorSwizzle) { - const auto& [num_tensors, M, K, rowwise] = std::get<0>(GetParam()); - const PerTileSwizzle per_tile_swizzle(std::get<1>(GetParam())); + const auto& [num_tensors, M, K, rowwise] = GetParam(); performTestMultiTensorSwizzle(num_tensors, M, K, rowwise); } @@ -343,17 +492,7 @@ TEST_P(MultiTensorRoundtripTestSuite, TestMultiTensorRoundtrip) { namespace { -// Shapes that exercise the narrow_k kernel (rowwise) / narrow_m kernel (colwise): -// Narrow-K fires when ALL tensors have scale num_tiles_k < TB_DIM (32), -// i.e. padded ceil(K/32) < 128. -// Narrow-M fires analogously for colwise when padded K < 4096 -// (since colwise m = K padded to 128, num_tiles_m = m / 128 < 32). -// -// Shapes that bypass narrow and use the regular multi_tensor kernel: -// K >= 4096 makes num_tiles_k >= 32 (rowwise) and num_tiles_m >= 32 (colwise). - std::vector> multi_tensor_test_cases = { - // --- Narrow path cases (K small → narrow_k for row, narrow_m for col) --- // M and K both aligned to 128 {3, 256, 256, true}, {3, 256, 256, false}, @@ -365,7 +504,7 @@ std::vector> multi_tensor_test_cases = { {3, 192, 256, false}, {2, 64, 256, true}, {2, 64, 256, false}, - // Larger narrow K (num_tiles_k = 8, shared mem = 128 KB) + // Larger K {2, 128, 1024, true}, {2, 128, 1024, false}, // K not divisible by 128 @@ -380,13 +519,16 @@ std::vector> multi_tensor_test_cases = { {4, 32, 64, true}, {4, 32, 64, false}, - // --- Non-narrow path cases (K >= 4096 → regular multi_tensor kernel) --- {3, 256, 4096, true}, {3, 256, 4096, false}, {2, 128, 8192, true}, {2, 128, 8192, false}, // Rowwise num_tiles_k = 34 selects vec_load_size = 2 (int2 kernel) {2, 128, 4352, true}, + // Wide row-wise scales fall back for an underfilled grid and use the full-width kernel once + // the aggregate grid is large enough. + {2, 2048, 7168, true}, // 32 aggregate row tiles: regular fallback. + {16, 2048, 7168, true}, // 256 aggregate row tiles: rowwise full-width kernel. // Colwise num_tiles_k = 512 / 32 / 4 = 4 selects vec_load_size = 4 (int4 kernel) {2, 512, 4096, false}, // MXFP8 swizzle: more tensors than one launch takes, several row-wise tiles per CTA, @@ -404,14 +546,12 @@ std::vector> multi_tensor_test_cases = { INSTANTIATE_TEST_SUITE_P( OperatorTest, MultiTensorSwizzleTestSuite, - ::testing::Combine(::testing::ValuesIn(multi_tensor_test_cases), ::testing::Values(false, true)), + ::testing::ValuesIn(multi_tensor_test_cases), [](const testing::TestParamInfo& info) { - const auto& test_case = std::get<0>(info.param); - return "n" + std::to_string(std::get<0>(test_case)) + - "_M" + std::to_string(std::get<1>(test_case)) + - "_K" + std::to_string(std::get<2>(test_case)) + - (std::get<3>(test_case) ? "_row" : "_col") + - (std::get<1>(info.param) ? "_per_tile" : ""); + return "n" + std::to_string(std::get<0>(info.param)) + + "_M" + std::to_string(std::get<1>(info.param)) + + "_K" + std::to_string(std::get<2>(info.param)) + + (std::get<3>(info.param) ? "_row" : "_col"); }); INSTANTIATE_TEST_SUITE_P( diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index a442215ae15..4ebb23fdc51 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -204,18 +205,167 @@ void performTestUnswizzle1D(const size_t M, const size_t K, bool rowwise, bool c // The buffer is stored in leading-dim-major order (row-major for rowwise, // column-major for colwise). `padded_rows x padded_cols` is the full // (padded) shape; `orig_rows` / `orig_cols` are the unpadded extents. -static void zero_scale_inv_padding(uint8_t *buf, - size_t padded_rows, size_t padded_cols, - size_t orig_rows, size_t orig_cols) { +static void set_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols, + uint8_t value) { for (size_t r = 0; r < padded_rows; ++r) { for (size_t c = 0; c < padded_cols; ++c) { if (r >= orig_rows || c >= orig_cols) { - buf[r * padded_cols + c] = 0; + buf[r * padded_cols + c] = value; } } } } +static void zero_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols) { + set_scale_inv_padding(buf, padded_rows, padded_cols, orig_rows, orig_cols, 0); +} + +void performTestSwizzleNVFP4(const size_t M, const size_t K, bool rowwise) { + using namespace test; + constexpr size_t BLOCK_SIZE = 16; + + const std::vector shape{M, K}; + Tensor input("input", shape, DType::kFloat4E2M1, rowwise, !rowwise, NVTE_NVFP4_1D_SCALING); + Tensor output("output", shape, DType::kFloat4E2M1, rowwise, !rowwise, NVTE_NVFP4_1D_SCALING); + output.set_with_gemm_swizzled_scales(true); + if (rowwise) { + input.fill_uniform_rowwise_scale_inv(); + } else { + input.fill_uniform_columnwise_scale_inv(); + } + + const NVTEShape scale_shape = + rowwise ? input.rowwise_scale_inv_shape() : input.columnwise_scale_inv_shape(); + uint8_t *input_scales = rowwise ? input.rowwise_cpu_scale_inv_ptr() + : input.columnwise_cpu_scale_inv_ptr(); + const size_t original_rows = rowwise ? M : K; + const size_t original_cols = + rowwise ? divide_round_up(K, BLOCK_SIZE) : divide_round_up(M, BLOCK_SIZE); + set_scale_inv_padding(input_scales, scale_shape.data[0], scale_shape.data[1], original_rows, + original_cols, 0xFF); + std::vector reference_input( + input_scales, input_scales + scale_shape.data[0] * scale_shape.data[1]); + zero_scale_inv_padding(reference_input.data(), scale_shape.data[0], scale_shape.data[1], + original_rows, original_cols); + input.from_cpu(); + + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, true>(reference_input.data(), reference.data(), scale_shape.data[0], + scale_shape.data[1]); + nvte_swizzle_scaling_factors(input.data(), output.data(), 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + output.to_cpu(); + + const uint8_t *actual = rowwise ? output.rowwise_cpu_scale_inv_ptr() + : output.columnwise_cpu_scale_inv_ptr(); + compareResults("nvfp4_swizzle", actual, reference.data(), reference.size()); +} + +TEST(SwizzleTest, NVFP4FullWidthChunkAndFallback) { + performTestSwizzleNVFP4(128, 128, true); + performTestSwizzleNVFP4(128, 128, false); + performTestSwizzleNVFP4(160, 96, true); // Padding in both scale dimensions. + performTestSwizzleNVFP4(160, 96, false); // Transposed padding in both scale dimensions. + performTestSwizzleNVFP4(128, 3584, true); // 28.5 KiB, 1 tile: rowwise full-width kernel. + performTestSwizzleNVFP4(2048, 3584, true); // 28.5 KiB, 16 tiles: rowwise full-width kernel. + performTestSwizzleNVFP4(2176, 3584, true); // 28.5 KiB, 17 tiles: regular fallback. + performTestSwizzleNVFP4(128, 4032, true); // 31.5 KiB: regular fallback. + performTestSwizzleNVFP4(128, 4096, true); // 32.5 KiB: regular fallback. + performTestSwizzleNVFP4(3584, 128, false); // 28.5 KiB, 1 tile: transposed full-width kernel. + performTestSwizzleNVFP4(3584, 2048, false); // 28.5 KiB, 16 tiles: transposed full-width kernel. + performTestSwizzleNVFP4(3584, 2176, false); // 28.5 KiB, 17 tiles: regular fallback. + performTestSwizzleNVFP4(4032, 128, false); // 31.5 KiB: regular fallback. + performTestSwizzleNVFP4(4096, 128, false); // 32.5 KiB: regular fallback. +} + +constexpr size_t kAlignmentTestM = 128; +constexpr size_t kAlignmentTestK = 128; +constexpr size_t kAlignmentTestBlockSize = 32; +constexpr size_t kAlignmentOffsetBytes = 4; + +void performTestRejectUnalignedMXFP8Scales(bool rowwise, bool offset_input) { + using namespace test; + constexpr size_t kDataSize = kAlignmentTestM * kAlignmentTestK; + constexpr size_t kScaleSize = kDataSize / kAlignmentTestBlockSize; + auto input_data = cuda_alloc(kDataSize); + auto output_data = cuda_alloc(kDataSize); + auto input_storage = cuda_alloc(kScaleSize + kAlignmentOffsetBytes); + auto output_storage = cuda_alloc(kScaleSize + kAlignmentOffsetBytes); + uint8_t* input_scales = input_storage.get() + (offset_input ? kAlignmentOffsetBytes : 0); + uint8_t* output_scales = output_storage.get() + (offset_input ? 0 : kAlignmentOffsetBytes); + TensorWrapper input(NVTE_MXFP8_1D_SCALING); + TensorWrapper output(NVTE_MXFP8_1D_SCALING); + const std::vector data_shape{kAlignmentTestM, kAlignmentTestK}; + const std::vector scale_shape = + rowwise ? std::vector{kAlignmentTestM, kAlignmentTestK / kAlignmentTestBlockSize} + : std::vector{kAlignmentTestM / kAlignmentTestBlockSize, kAlignmentTestK}; + if (rowwise) { + input.set_rowwise_data(input_data.get(), DType::kFloat8E4M3, data_shape); + input.set_rowwise_scale_inv(input_scales, DType::kFloat8E8M0, scale_shape); + output.set_rowwise_data(output_data.get(), DType::kFloat8E4M3, data_shape); + output.set_rowwise_scale_inv(output_scales, DType::kFloat8E8M0, scale_shape); + } else { + input.set_columnwise_data(input_data.get(), DType::kFloat8E4M3, data_shape); + input.set_columnwise_scale_inv(input_scales, DType::kFloat8E8M0, scale_shape); + output.set_columnwise_data(output_data.get(), DType::kFloat8E4M3, data_shape); + output.set_columnwise_scale_inv(output_scales, DType::kFloat8E8M0, scale_shape); + } + output.set_with_gemm_swizzled_scales(true); + EXPECT_THROW(nvte_swizzle_scaling_factors(input.data(), output.data(), 0), std::runtime_error); +} + +TEST(SwizzleTest, RejectUnalignedMXFP8ScaleBuffers) { + for (bool rowwise : {true, false}) { + performTestRejectUnalignedMXFP8Scales(rowwise, true); + performTestRejectUnalignedMXFP8Scales(rowwise, false); + } +} + +TEST(SwizzleTest, RejectUnalignedMultiMXFP8ScaleBuffers) { + using namespace test; + constexpr size_t kNumTensors = 2; + constexpr size_t kDataSize = kAlignmentTestM * kAlignmentTestK; + constexpr size_t kScaleSize = kDataSize / kAlignmentTestBlockSize; + const std::vector data_shape{kAlignmentTestM, kAlignmentTestK}; + const std::vector scale_shape{kAlignmentTestM, + kAlignmentTestK / kAlignmentTestBlockSize}; + for (bool offset_input : {true, false}) { + std::vector> input_data, output_data, input_storage, output_storage; + std::vector input_tensors, output_tensors; + std::vector inputs, outputs; + input_tensors.reserve(kNumTensors); + output_tensors.reserve(kNumTensors); + for (size_t i = 0; i < kNumTensors; i++) { + input_data.push_back(cuda_alloc(kDataSize)); + output_data.push_back(cuda_alloc(kDataSize)); + input_storage.push_back(cuda_alloc(kScaleSize + kAlignmentOffsetBytes)); + output_storage.push_back(cuda_alloc(kScaleSize + kAlignmentOffsetBytes)); + uint8_t* input_scales = + input_storage.back().get() + (i == 0 && offset_input ? kAlignmentOffsetBytes : 0); + uint8_t* output_scales = + output_storage.back().get() + (i == 0 && !offset_input ? kAlignmentOffsetBytes : 0); + input_tensors.emplace_back(NVTE_MXFP8_1D_SCALING); + output_tensors.emplace_back(NVTE_MXFP8_1D_SCALING); + input_tensors.back().set_rowwise_data(input_data.back().get(), DType::kFloat8E4M3, + data_shape); + input_tensors.back().set_rowwise_scale_inv(input_scales, DType::kFloat8E8M0, scale_shape); + output_tensors.back().set_rowwise_data(output_data.back().get(), DType::kFloat8E4M3, + data_shape); + output_tensors.back().set_rowwise_scale_inv(output_scales, DType::kFloat8E8M0, scale_shape); + output_tensors.back().set_with_gemm_swizzled_scales(true); + inputs.push_back(input_tensors.back().data()); + outputs.push_back(output_tensors.back().data()); + } + EXPECT_THROW( + nvte_multi_tensor_swizzle_scaling_factors(inputs.data(), outputs.data(), inputs.size(), 0), + std::runtime_error); + } +} + void performTestGroupedSwizzleMXFP8(const int num_tensors, const size_t M, const size_t K) { using namespace transformer_engine; using namespace test; @@ -308,17 +458,9 @@ void performTestGroupedSwizzleMXFP8(const int num_tensors, const size_t M, const num_tensors * col_numel); } -// Sets NVTE_MXFP8_PER_TILE_SWIZZLE while it lives, so that MXFP8 swizzles use the per-tile -// kernels (enable) or the default ones. -class PerTileSwizzle { - public: - explicit PerTileSwizzle(bool enable) { - setenv("NVTE_MXFP8_PER_TILE_SWIZZLE", enable ? "1" : "0", 1); - } - ~PerTileSwizzle() { unsetenv("NVTE_MXFP8_PER_TILE_SWIZZLE"); } -}; - -class SwizzleTestSuite : public ::testing::TestWithParam, std::pair, bool, bool>> {}; +class SwizzleTestSuite + : public ::testing::TestWithParam< + std::tuple, std::pair, bool>> {}; TEST_P(SwizzleTestSuite, TestSwizzle) { @@ -328,7 +470,6 @@ TEST_P(SwizzleTestSuite, TestSwizzle) { const auto num_tiles = std::get<0>(GetParam()); const auto scaling_mode = std::get<1>(GetParam()); const auto transa = std::get<2>(GetParam()); - const PerTileSwizzle per_tile_swizzle(std::get<3>(GetParam())); performTestSwizzle1D(num_tiles.first, num_tiles.second, scaling_mode.first, scaling_mode.second, @@ -1007,14 +1148,15 @@ std::vector> num_tiles = { {65, 257}, {65, 258}, {65, 259}, - // Additional narrow-path coverage: narrow_k (row) when num_tiles_K < 32, - // narrow_m (col) when num_tiles_M < 32. - {1, 4}, // narrow_k with 4 K-tiles - {1, 8}, // narrow_k with 8 K-tiles - {4, 1}, // narrow_m with 4 M-tiles - {8, 1}, // narrow_m with 8 M-tiles - {31, 1}, // narrow_m at boundary (31 < TB_DIM=32) - {1, 31}, // narrow_k at boundary (31 < TB_DIM=32) + {1, 4}, + {1, 8}, + {4, 1}, + {8, 1}, + {31, 1}, + {1, 31}, + {2, 63}, // Above 30 KiB: regular fallback. + {2, 64}, // Above 32 KiB: regular fallback. + {16, 56}, // 28.5 KiB: rowwise full-width kernel. {1600, 1}, // MXFP8: more row-wise tiles than one CTA per tile allows }; @@ -1051,8 +1193,7 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Combine( ::testing::ValuesIn(num_tiles), ::testing::ValuesIn(scaling_mode), - ::testing::ValuesIn(transa), - ::testing::Values(false, true) + ::testing::ValuesIn(transa) ), [](const testing::TestParamInfo& info) { std::string name = "ntiles" + @@ -1060,8 +1201,7 @@ INSTANTIATE_TEST_SUITE_P( std::to_string(std::get<0>(info.param).second) + "smode" + std::to_string(std::get<1>(info.param).first) + "X"+ std::to_string(std::get<1>(info.param).second) + "trans" + - std::to_string(std::get<2>(info.param)) + - (std::get<3>(info.param) ? "PerTile" : ""); + std::to_string(std::get<2>(info.param)); return name; }); diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index f85d88fdce8..1c54e5f8981 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -26,6 +26,7 @@ extern "C" { * Requirements: * - scale_inv is stored in row-major. * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - Input and output scale_inv buffer pointers used by this operation must be 16-byte aligned. * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. */ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream); @@ -40,6 +41,7 @@ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cud * Requirements: * - scale_inv is stored in row-major. * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - Input and output scale_inv buffer pointers used by this operation must be 16-byte aligned. * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. */ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs, @@ -50,6 +52,7 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen * * Use this variant when the data and scale_inv tensors intentionally have * different shapes, e.g. when scale_invs have been transposed for attention. + * The 16-byte alignment requirement for scale_inv buffer pointers used by this operation still applies. */ void nvte_multi_tensor_swizzle_scaling_factors_unchecked(const NVTETensor* inputs, NVTETensor* outputs, diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 19fb8e848ad..7089a981dc2 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -33,6 +33,38 @@ constexpr __device__ __host__ int N_SF_PER_TD_PER_TILE = 4; constexpr __device__ __host__ int NEW_SF_TILE_DIM_K_I32 = NEW_SF_TILE_DIM_K / 4; constexpr __device__ __host__ int NEW_SF_TILE_DIM_M_I32 = 32; +// Shared layout for the optimized rowwise full-width and columnwise chunk kernels. Their local K +// dimension indexes blocks along the quantized axis; local M is the other axis. An M tile spans +// 128 rows in rowwise scales or 128 columns in columnwise scales. Each scale is one byte; a +// uint32_t word packs four scales. One output block holds 128 M positions by four K scale +// positions, or 512 bytes. +constexpr int kSwizzleTileDimM = 128; +constexpr int kScalesPerWord = 4; +constexpr int kBytesPerWord = sizeof(uint32_t); +// One byte per M position makes a 128-position span eight 16-byte vectors. +constexpr int kVectorsPerMSpan = kSwizzleTileDimM / sizeof(int4); +constexpr int kMPositionsPerSwizzleGroup = kSwizzleTileDimM / kScalesPerWord; +constexpr int kSwizzleOutputBlockBytes = kSwizzleTileDimM * kScalesPerWord; + +// Both optimized multi-tensor variants store at most this many descriptors per launch. +constexpr int kMaxTensorsPerLaunch = 64; + +// Rowwise full-width kernels and their dispatch thresholds. Narrow rows always use this kernel; +// wider rows use it for tiny grids or grids large enough to occupy the GPU. +constexpr int kRowwiseThreadsPerBlock = 128; +constexpr int kRowwiseMaxSharedMemoryBytes = 32 * 1024; +constexpr int kRowwiseNarrowSharedMemoryBytes = 17 * 1024; +constexpr int kRowwiseWideSharedMemoryBytes = 30 * 1024; +constexpr int kRowwiseMaxTinyTiles = 16; +// Largest rowwise grid before several row tiles share a block. +constexpr int kRowwiseMaxBlocks = 1536; + +// Columnwise chunk kernels and the dispatch threshold for switching from 32 to 128 scale rows. +constexpr int kColumnwiseSmallChunkRows = 32; +constexpr int kColumnwiseLargeChunkRows = 128; +constexpr int kColumnwiseThreadsPerBlock = 256; +constexpr int kColumnwiseMaxBlocksWith32RowChunks = 1024; + template __device__ inline void regs_shuffle_with_bit_shifts(LType* regs_vec) { // inp, 4-byte chunks [0,1,2,3, 4,5,6,7, 8,9,10,11, 12,13,14,15] @@ -223,17 +255,6 @@ __device__ __forceinline__ void dispatch_swizzle_col_scaling_kernel_impl( } } -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - swizzle_col_scaling_kernel(const void* input, void* output, const int M, const int K, - const int original_M, const int original_K) { - const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M); - const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K); - dispatch_swizzle_col_scaling_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y, - padding_k, padding_m); -} - template __device__ inline void regs_shuffle(LType* regs_vec) { constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); @@ -394,163 +415,6 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) padding_k, padding_m); } -// Narrow-K specialization for row scaling swizzle. -// When K is small (num_tiles_k < TB_DIM), the standard kernel wastes threadIdx.x -// because there aren't enough K-tiles to distribute across threads. -// This kernel repurposes the thread dimensions: threadIdx.x iterates rows within -// an M-tile, threadIdx.y indexes M-tiles within the block, processing TB_DIM -// M-tiles per block with full thread utilization. -template -__device__ void swizzle_row_scaling_narrow_k_kernel_impl(const void* input, void* output, - const int M, const int K, - const int original_M, const int original_K, - const int bid, const int grid_dim) { - constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4; - const int K_i32 = K / 4; - const int num_tiles_m = M / SF_TILE_DIM_M; - - const int m_tile = bid * blockDim.y + threadIdx.y; - const bool active = (m_tile < num_tiles_m); - - extern __shared__ int4 slm_v4i[]; - const int slm_tile_v4i = K_i32 * (SF_TILE_SIZE_I32 / 4); - - if (active) { - const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M); - const bool padding_k = (original_K < K); - - int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i; - - for (int k = 0; k < K_i32; k++) { - const int input_base = m_tile * SF_TILE_DIM_M * K_i32 + k; - const int* input_i32 = reinterpret_cast(input) + input_base; - - int regs[N_SF_PER_TD_PER_TILE]; -#pragma unroll - for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { - const int row = i * TB_DIM + threadIdx.x; - regs[i] = __ldg(input_i32 + row * K_i32); - if (padding_m || padding_k) { - for (int j = 0; j < 4; j++) { - const int byte_row = m_tile * SF_TILE_DIM_M + row; - const int byte_col = k * 4 + j; - if (byte_row >= original_M || byte_col >= original_K) { - reinterpret_cast(®s[i])[j] = 0; - } - } - } - } - - my_slm[k * (SF_TILE_SIZE_I32 / 4) + threadIdx.x] = *reinterpret_cast(regs); - } - } - - __syncthreads(); - - if (active) { - int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i; - int4* out_v4i = - reinterpret_cast(reinterpret_cast(output) + m_tile * SF_TILE_DIM_M * K_i32); - - for (int i = threadIdx.x; i < slm_tile_v4i; i += blockDim.x) { - out_v4i[i] = my_slm[i]; - } - } -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - swizzle_row_scaling_narrow_k_kernel(const void* input, void* output, const int M, const int K, - const int original_M, const int original_K) { - swizzle_row_scaling_narrow_k_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x); -} - -// Narrow-M variant of the column scaling swizzle kernel, for when num_tiles_m < TB_DIM. -// Analogous to the narrow-K row kernel: when the M dimension is small, the normal -// col kernel underutilizes threads in the load phase because threadIdx.x covers M -// positions with vectorized loads, leaving many threads idle. This kernel repurposes -// thread dimensions: threadIdx.y indexes K-tiles within the block, threadIdx.x covers -// one int32 column of an M-tile, and M-tiles are iterated serially. -template -__device__ void swizzle_col_scaling_narrow_m_kernel_impl(const void* input, void* output, - const int M, const int K, - const int original_M, const int original_K, - const int bid, const int grid_dim) { - constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4; - constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4; - constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K; - - const int M_i32 = M / 4; - const int K_i32 = K; - const int num_tiles_m = M / SF_TILE_DIM_M; - const int num_tiles_k = K / SF_TILE_DIM_K; - - const int k_tile = bid * blockDim.y + threadIdx.y; - const bool active = (k_tile < num_tiles_k); - const int remaining = num_tiles_k - bid * static_cast(blockDim.y); - const int k_tiles_in_block = remaining <= 0 ? 0 : (remaining < TB_DIM ? remaining : TB_DIM); - - extern __shared__ int slm_narrow_m[]; - - if (active) { - const bool padding_k = (k_tile == num_tiles_k - 1) && (original_K < K); - const int32_t* input_i32 = reinterpret_cast(input); - - for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) { - const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M); - - int regs[N_SF_PER_TD_PER_TILE]; -#pragma unroll - for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { - const int k_row = k_tile * SF_TILE_DIM_K_I32 + i; - const int m_col = m_tile * SF_TILE_DIM_M_I32 + threadIdx.x; - regs[i] = __ldg(input_i32 + k_row * M_i32 + m_col); - if (padding_m || padding_k) { - for (int j = 0; j < 4; j++) { - if (m_col * 4 + j >= original_M || k_row >= original_K) { - reinterpret_cast(®s[i])[j] = 0; - } - } - } - } - - regs_shuffle_with_bit_shifts(regs); - - int tM = threadIdx.x * N_SF_PER_TD_PER_TILE; - int* slm_tile = - slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32 + threadIdx.y * SF_TILE_SIZE_I32; -#pragma unroll - for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { - slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 + - ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32] = regs[i]; - } - } - } - - __syncthreads(); - - const int linear_id = threadIdx.y * blockDim.x + threadIdx.x; - for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) { - int4* out_v4i = reinterpret_cast(reinterpret_cast(output) + - m_tile * SF_TILE_DIM_M_I32 * K_i32 + - bid * TB_DIM * SF_TILE_SIZE_I32); - int4* slm_v4i = reinterpret_cast(slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32); - const int n_v4i = k_tiles_in_block * SF_TILE_SIZE_I32 / 4; - for (int j = linear_id; j < n_v4i; j += blockDim.x * blockDim.y) { - out_v4i[j] = slm_v4i[j]; - } - } -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - swizzle_col_scaling_narrow_m_kernel(const void* input, void* output, const int M, const int K, - const int original_M, const int original_K) { - swizzle_col_scaling_narrow_m_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x); -} - template __device__ void unswizzle_row_scaling_kernel_impl(const void* input, void* output, const int M, const int K, const int bid_x, const int bid_y, @@ -712,10 +576,6 @@ struct MultiSwizzleArgs { int num_tensors; }; -constexpr size_t round_up_to_multiple(size_t value, size_t multiple) { - return DIVUP(value, multiple) * multiple; -} - template __global__ void __launch_bounds__(TB_DIM* TB_DIM) grouped_swizzle_row_scaling_uniform_shape_kernel(const void* input, void* output, const int M, @@ -867,108 +727,16 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) padding_m); } -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_args) { - // Find tensor corresponding to block - const int bid = blockIdx.x; - int tensor_id = 0; - while (kernel_args.block_range[tensor_id + 1] <= bid) { - ++tensor_id; - } - // Get args corresponding to block - const void* input = kernel_args.input_list[tensor_id]; - void* output = kernel_args.output_list[tensor_id]; - const int M = kernel_args.m_list[tensor_id]; - const int K = kernel_args.k_list[tensor_id]; - const int original_M = kernel_args.original_m_list[tensor_id]; - const int original_K = kernel_args.original_k_list[tensor_id]; - - constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); - - // Get block index in grid. Emulate 2D grid. - const int num_tiles_k = K / SF_TILE_DIM_K; - const int num_tiles_m = M / SF_TILE_DIM_M; - const int grid_dim_x = DIVUP(num_tiles_k, TB_DIM); - const int grid_dim_y = DIVUP(num_tiles_m, N_TILE_PER_TD); - const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y; - const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y; - - const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M); - const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K); - dispatch_swizzle_col_scaling_kernel_impl( - input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k, - padding_m); -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - multi_tensor_swizzle_row_scaling_narrow_k_kernel(MultiSwizzleArgs kernel_args) { - const int bid = blockIdx.x; - int tensor_id = 0; - while (kernel_args.block_range[tensor_id + 1] <= bid) { - ++tensor_id; - } - const void* input = kernel_args.input_list[tensor_id]; - void* output = kernel_args.output_list[tensor_id]; - const int M = kernel_args.m_list[tensor_id]; - const int K = kernel_args.k_list[tensor_id]; - const int original_M = kernel_args.original_m_list[tensor_id]; - const int original_K = kernel_args.original_k_list[tensor_id]; - const int flat_bid = bid - kernel_args.block_range[tensor_id]; - const int num_tiles_m = M / SF_TILE_DIM_M; - const int grid_dim = DIVUP(num_tiles_m, TB_DIM); - - swizzle_row_scaling_narrow_k_kernel_impl( - input, output, M, K, original_M, original_K, flat_bid, grid_dim); -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - multi_tensor_swizzle_col_scaling_narrow_m_kernel(MultiSwizzleArgs kernel_args) { - const int bid = blockIdx.x; - int tensor_id = 0; - while (kernel_args.block_range[tensor_id + 1] <= bid) { - ++tensor_id; - } - const void* input = kernel_args.input_list[tensor_id]; - void* output = kernel_args.output_list[tensor_id]; - const int M = kernel_args.m_list[tensor_id]; - const int K = kernel_args.k_list[tensor_id]; - const int original_M = kernel_args.original_m_list[tensor_id]; - const int original_K = kernel_args.original_k_list[tensor_id]; - const int flat_bid = bid - kernel_args.block_range[tensor_id]; - const int num_tiles_k = K / SF_TILE_DIM_K; - const int grid_dim = DIVUP(num_tiles_k, TB_DIM); - - swizzle_col_scaling_narrow_m_kernel_impl( - input, output, M, K, original_M, original_K, flat_bid, grid_dim); -} - -// MXFP8 scale swizzle with one thread block per scale tile, used when -// NVTE_MXFP8_PER_TILE_SWIZZLE=1. The narrow-K / narrow-M kernels above put 32 tiles in one -// 1024-thread block, so e.g. the scales of a 4096-row activation are swizzled by a single block. -// These kernels give every 128-row tile (row-wise) or 128-column x 32-row block (column-wise) its -// own block, which moves a contiguous chunk with coalesced 16-byte loads and stores through shared -// memory. Multi-tensor launches put several tiles in one block only when the grid would otherwise -// be very large. -constexpr int kRowwiseThreadsPerBlock = 128; -constexpr int kColumnwiseThreadsPerBlock = 256; -constexpr int kMaxTensorsPerLaunch = 64; -// Row-wise: the largest grid before several tiles share a block. -constexpr int kRowwiseMaxBlocks = 1536; -// Column-wise: the largest grid with 32-row chunks before switching to 128-row chunks. -constexpr int kColumnwiseMaxBlocksWith32RowChunks = 1024; - -// Whether MXFP8 swizzles use the per-tile kernels (NVTE_MXFP8_PER_TILE_SWIZZLE=1). -bool use_per_tile_swizzle() { return getenv("NVTE_MXFP8_PER_TILE_SWIZZLE"); } - +// Optimized scale swizzle for narrow scale widths. "Rowwise full width" means that a block stages +// the entire padded K dimension for one or more 128-row tiles. "Columnwise chunk" means that a +// block stages only 32 (or 128 for large grids) K rows for one 128-column M tile. Both paths use +// coalesced 16-byte loads and stores through shared memory. // The 4 scales packed in `word` (the first one in column first_column), with those in columns // >= num_columns set to zero. __device__ __forceinline__ uint32_t zero_padding_columns(uint32_t word, int first_column, int num_columns) { const int valid_columns = num_columns - first_column; - if (valid_columns >= 4) return word; + if (valid_columns >= kScalesPerWord) return word; if (valid_columns <= 0) return 0u; return word & ((1u << (8 * valid_columns)) - 1u); } @@ -980,31 +748,33 @@ __device__ __forceinline__ uint32_t divide_by_magic(uint32_t numerator, uint32_t } inline uint32_t division_magic(int divisor) { - return divisor <= 1 ? 0u : static_cast(((1ull << 32) + divisor - 1) / divisor); + return divisor <= 1 + ? 0u + : static_cast(DIVUP(uint64_t{1} << 32, static_cast(divisor))); } -// Row-wise: tiles [first_tile, first_tile + num_tiles) of one tensor, a contiguous chunk of the +// Row-wise: tiles [first_tile, first_tile + num_tiles) of one tensor, a consecutive chunk of the // input and of the output. Shared-memory rows are padded to an odd number of words, so the column // gathers are free of bank conflicts. Requires num_tiles * 128 * padded_k / 4 < 2^16 (guaranteed -// by the 48 KB shared-memory limit). +// by the row-wise shared-memory limit = 30 KB). __device__ __forceinline__ void swizzle_rowwise_tiles(const uint8_t* __restrict__ input, uint8_t* __restrict__ output, int padded_k, uint32_t words_per_row_magic, int original_m, int original_k, int first_tile, int num_tiles, uint32_t* smem) { - const int words_per_row = padded_k >> 2; + const int words_per_row = padded_k / kScalesPerWord; const int smem_row_words = words_per_row | 1; - const int num_vectors = - num_tiles * 8 * padded_k; // 16-byte vectors: num_tiles * 128 * padded_k / 16 - const size_t chunk_offset = static_cast(first_tile) * 128 * padded_k; + const int num_vectors = num_tiles * kVectorsPerMSpan * padded_k; + const size_t chunk_offset = static_cast(first_tile) * kSwizzleTileDimM * padded_k; const int4* input_vectors = reinterpret_cast(input + chunk_offset); for (int i = threadIdx.x; i < num_vectors; i += kRowwiseThreadsPerBlock) { const int4 vector = __ldg(input_vectors + i); - const uint32_t words[4] = {static_cast(vector.x), static_cast(vector.y), - static_cast(vector.z), static_cast(vector.w)}; + const uint32_t words[kScalesPerWord] = { + static_cast(vector.x), static_cast(vector.y), + static_cast(vector.z), static_cast(vector.w)}; #pragma unroll - for (int w = 0; w < 4; w++) { - const uint32_t word_index = 4 * i + w; + for (int w = 0; w < kScalesPerWord; w++) { + const uint32_t word_index = kScalesPerWord * i + w; const uint32_t row = divide_by_magic(word_index, words_per_row_magic); smem[row * smem_row_words + (word_index - row * words_per_row)] = words[w]; } @@ -1013,22 +783,22 @@ __device__ __forceinline__ void swizzle_rowwise_tiles(const uint8_t* __restrict_ // An output tile holds padded_k / 4 blocks of 512 bytes, i.e. 32 vectors each; vector r of a // block packs rows r, r + 32, r + 64 and r + 96 of its 4 scale columns. - const int vectors_per_tile = 8 * padded_k; + const int vectors_per_tile = kVectorsPerMSpan * padded_k; for (int tile = 0; tile < num_tiles; tile++) { - int4* output_vectors = - reinterpret_cast(output + chunk_offset + static_cast(tile) * 128 * padded_k); - const int tile_first_row = tile * 128; - const int global_first_row = first_tile * 128 + tile_first_row; + int4* output_vectors = reinterpret_cast( + output + chunk_offset + static_cast(tile) * kSwizzleTileDimM * padded_k); + const int tile_first_row = tile * kSwizzleTileDimM; + const int global_first_row = first_tile * kSwizzleTileDimM + tile_first_row; for (int v = threadIdx.x; v < vectors_per_tile; v += kRowwiseThreadsPerBlock) { - const int k_block = v >> 5; - const int row_in_group = v & 31; - uint32_t words[4]; + const int k_block = v / kMPositionsPerSwizzleGroup; + const int row_in_group = v % kMPositionsPerSwizzleGroup; + uint32_t words[kScalesPerWord]; #pragma unroll - for (int group = 0; group < 4; group++) { - const int row = row_in_group + 32 * group; + for (int group = 0; group < kScalesPerWord; group++) { + const int row = row_in_group + kMPositionsPerSwizzleGroup * group; uint32_t word = smem[(tile_first_row + row) * smem_row_words + k_block]; if (global_first_row + row >= original_m) word = 0u; - words[group] = zero_padding_columns(word, k_block * 4, original_k); + words[group] = zero_padding_columns(word, k_block * kScalesPerWord, original_k); } output_vectors[v] = make_int4(static_cast(words[0]), static_cast(words[1]), static_cast(words[2]), static_cast(words[3])); @@ -1045,92 +815,97 @@ __device__ __forceinline__ void swizzle_columnwise_chunk(const uint8_t* __restri int padded_k, int original_m, int original_k, int m_tile, int first_k, int4* smem) { - const int num_rows = min(kChunkRows, padded_k - first_k); // multiple of 4 - const int num_vectors = num_rows * 8; + const int num_rows = min(kChunkRows, padded_k - first_k); // multiple of kScalesPerWord + const int num_vectors = num_rows * kVectorsPerMSpan; for (int i = threadIdx.x; i < num_vectors; i += kColumnwiseThreadsPerBlock) { - const int row = i >> 3; - const int vector_in_row = i & 7; - smem[i] = __ldg(reinterpret_cast( - input + static_cast(first_k + row) * padded_m + m_tile * 128) + - vector_in_row); + const int row = i / kVectorsPerMSpan; + const int vector_in_row = i % kVectorsPerMSpan; + smem[i] = + __ldg(reinterpret_cast(input + static_cast(first_k + row) * padded_m + + m_tile * kSwizzleTileDimM) + + vector_in_row); } __syncthreads(); const uint8_t* smem_bytes = reinterpret_cast(smem); - const size_t first_output_block = static_cast(m_tile) * (padded_k >> 2) + (first_k >> 2); + const size_t first_output_block = + static_cast(m_tile) * (padded_k / kScalesPerWord) + (first_k / kScalesPerWord); for (int v = threadIdx.x; v < num_vectors; v += kColumnwiseThreadsPerBlock) { - const int k_block = v >> 5; - const int column_in_group = v & 31; - uint32_t words[4]; + const int k_block = v / kMPositionsPerSwizzleGroup; + const int column_in_group = v % kMPositionsPerSwizzleGroup; + uint32_t words[kScalesPerWord]; #pragma unroll - for (int group = 0; group < 4; group++) { - const int column = column_in_group + 32 * group; + for (int group = 0; group < kScalesPerWord; group++) { + const int column = column_in_group + kMPositionsPerSwizzleGroup * group; uint32_t word = 0u; - if (m_tile * 128 + column < original_m) { + if (m_tile * kSwizzleTileDimM + column < original_m) { #pragma unroll - for (int j = 0; j < 4; j++) { - const int row = k_block * 4 + j; + for (int j = 0; j < kScalesPerWord; j++) { + const int row = k_block * kScalesPerWord + j; if (first_k + row < original_k) { - word |= static_cast(smem_bytes[row * 128 + column]) << (8 * j); + word |= static_cast(smem_bytes[row * kSwizzleTileDimM + column]) << (8 * j); } } } words[group] = word; } - reinterpret_cast(output + (first_output_block + k_block) * 512)[column_in_group] = + reinterpret_cast(output + (first_output_block + k_block) * + kSwizzleOutputBlockBytes)[column_in_group] = make_int4(static_cast(words[0]), static_cast(words[1]), static_cast(words[2]), static_cast(words[3])); } } __global__ void __launch_bounds__(kRowwiseThreadsPerBlock) - swizzle_rowwise_per_tile_kernel(const uint8_t* __restrict__ input, uint8_t* __restrict__ output, - int padded_k, uint32_t words_per_row_magic, int original_m, - int original_k) { + swizzle_rowwise_full_width_kernel(const uint8_t* __restrict__ input, + uint8_t* __restrict__ output, int padded_k, + uint32_t words_per_row_magic, int original_m, + int original_k) { extern __shared__ uint32_t rowwise_smem[]; swizzle_rowwise_tiles(input, output, padded_k, words_per_row_magic, original_m, original_k, blockIdx.x, 1, rowwise_smem); } __global__ void __launch_bounds__(kColumnwiseThreadsPerBlock) - swizzle_columnwise_per_tile_kernel(const uint8_t* __restrict__ input, - uint8_t* __restrict__ output, int padded_m, int padded_k, - int original_m, int original_k) { - __shared__ int4 columnwise_smem[32 * 8]; - swizzle_columnwise_chunk<32>(input, output, padded_m, padded_k, original_m, original_k, - blockIdx.x, blockIdx.y * 32, columnwise_smem); + swizzle_columnwise_chunk_kernel(const uint8_t* __restrict__ input, uint8_t* __restrict__ output, + int padded_m, int padded_k, int original_m, int original_k) { + __shared__ int4 columnwise_smem[kColumnwiseSmallChunkRows * kVectorsPerMSpan]; + swizzle_columnwise_chunk( + input, output, padded_m, padded_k, original_m, original_k, blockIdx.x, + blockIdx.y * kColumnwiseSmallChunkRows, columnwise_smem); } // The scales of one tensor: padded_m (a multiple of 128) by padded_k (a multiple of 4) in // row-wise terms; for column-wise scales, m runs along the columns of the data. Scales at // m >= original_m or k >= original_k are padding and written as zeros. -struct PerTileSwizzleTensor { +struct SwizzleTensorDescriptor { const uint8_t* input; uint8_t* output; int padded_m, padded_k, original_m, original_k; }; -struct PerTileMultiSwizzleArgs { - PerTileSwizzleTensor tensors[kMaxTensorsPerLaunch]; +struct MultiTensorSwizzleArgs { + SwizzleTensorDescriptor tensors[kMaxTensorsPerLaunch]; int first_block[kMaxTensorsPerLaunch + 1]; uint32_t words_per_row_magic[kMaxTensorsPerLaunch]; // row-wise only - int tiles_per_block; // row-wise only + int row_tiles_per_block; // row-wise only int num_tensors; }; -__device__ __forceinline__ int tensor_of_block(const PerTileMultiSwizzleArgs& args, int block) { +__device__ __forceinline__ int tensor_of_block(const MultiTensorSwizzleArgs& args, int block) { int tensor = 0; while (tensor + 1 < args.num_tensors && args.first_block[tensor + 1] <= block) tensor++; return tensor; } __global__ void __launch_bounds__(kRowwiseThreadsPerBlock) - multi_tensor_swizzle_rowwise_per_tile_kernel(const PerTileMultiSwizzleArgs args) { + multi_tensor_swizzle_rowwise_full_width_kernel(const MultiTensorSwizzleArgs args) { extern __shared__ uint32_t rowwise_multi_smem[]; const int t = tensor_of_block(args, blockIdx.x); - const int tiles_per_block = args.tiles_per_block; - const int first_tile = (blockIdx.x - args.first_block[t]) * tiles_per_block; - const int num_tiles = min(tiles_per_block, args.tensors[t].padded_m / 128 - first_tile); + const int row_tiles_per_block = args.row_tiles_per_block; + const int first_tile = (blockIdx.x - args.first_block[t]) * row_tiles_per_block; + const int num_tiles = + min(row_tiles_per_block, args.tensors[t].padded_m / kSwizzleTileDimM - first_tile); swizzle_rowwise_tiles(args.tensors[t].input, args.tensors[t].output, args.tensors[t].padded_k, args.words_per_row_magic[t], args.tensors[t].original_m, args.tensors[t].original_k, first_tile, num_tiles, rowwise_multi_smem); @@ -1138,11 +913,11 @@ __global__ void __launch_bounds__(kRowwiseThreadsPerBlock) template __global__ void __launch_bounds__(kColumnwiseThreadsPerBlock) - multi_tensor_swizzle_columnwise_per_tile_kernel(const PerTileMultiSwizzleArgs args) { - __shared__ int4 columnwise_multi_smem[kChunkRows * 8]; + multi_tensor_swizzle_columnwise_chunk_kernel(const MultiTensorSwizzleArgs args) { + __shared__ int4 columnwise_multi_smem[kChunkRows * kVectorsPerMSpan]; const int t = tensor_of_block(args, blockIdx.x); const int block_in_tensor = blockIdx.x - args.first_block[t]; - const int num_m_tiles = args.tensors[t].padded_m / 128; + const int num_m_tiles = args.tensors[t].padded_m / kSwizzleTileDimM; const int m_tile = block_in_tensor % num_m_tiles; const int first_k = (block_in_tensor / num_m_tiles) * kChunkRows; swizzle_columnwise_chunk(args.tensors[t].input, args.tensors[t].output, @@ -1155,116 +930,119 @@ inline bool is_16_byte_aligned(const void* ptr) { return (reinterpret_cast(ptr) & 15u) == 0; } -inline bool per_tile_swizzle_supported(const PerTileSwizzleTensor& tensor, bool rowwise) { - if (tensor.padded_m <= 0 || tensor.padded_k <= 0 || tensor.padded_m % 128 != 0 || - tensor.padded_k % 4 != 0) { - return false; +inline void check_scale_inv_alignment(const Tensor* tensor, const char* name, bool rowwise) { + const auto& scales = rowwise ? tensor->scale_inv : tensor->columnwise_scale_inv; + if (scales.has_data()) { + NVTE_CHECK(is_16_byte_aligned(scales.dptr), name, rowwise ? " row-wise" : " column-wise", + " scale_inv buffer must be 16-byte aligned."); } - if (!is_16_byte_aligned(tensor.input) || !is_16_byte_aligned(tensor.output)) return false; - if (rowwise && 128 * ((tensor.padded_k >> 2) | 1) * 4 > 48 * 1024) return false; - return true; } -// One kernel for up to kMaxTensorsPerLaunch tensors, all row-wise or all column-wise. Row-wise: -// the tiles per block are the smallest power of two that keeps the grid within kRowwiseMaxBlocks -// (shared memory permitting). Column-wise: 128-row chunks once 32-row chunks would exceed -// kColumnwiseMaxBlocksWith32RowChunks blocks. Returns false, without launching anything, if any -// tensor is not supported. -inline bool launch_multi_tensor_per_tile_swizzle(const PerTileSwizzleTensor* tensors, - int num_tensors, bool rowwise, - cudaStream_t stream) { - if (num_tensors <= 0 || num_tensors > kMaxTensorsPerLaunch) return false; - PerTileMultiSwizzleArgs args; +inline bool should_use_rowwise_full_width_swizzle(int padded_k, int64_t total_row_tiles) { + const int smem_bytes = kSwizzleTileDimM * ((padded_k / kScalesPerWord) | 1) * kBytesPerWord; + if (smem_bytes <= kRowwiseNarrowSharedMemoryBytes) return true; + // The fallback launches one block per row tile at these widths. Switch after its grid + // exceeds one wave on the current GPU. + return smem_bytes <= kRowwiseWideSharedMemoryBytes && + (total_row_tiles <= kRowwiseMaxTinyTiles || total_row_tiles > cuda::sm_count()); +} + +// Launches the rowwise full-width kernel, where each block stages the entire padded K dimension, +// for up to kMaxTensorsPerLaunch tensors. The row tiles per block are the smallest power of two +// that keeps the grid within kRowwiseMaxBlocks, subject to the shared-memory limit. +inline void launch_multi_tensor_rowwise_full_width_swizzle(const SwizzleTensorDescriptor* tensors, + int num_tensors, cudaStream_t stream) { + MultiTensorSwizzleArgs args{}; args.num_tensors = num_tensors; args.first_block[0] = 0; - int64_t total_tiles = 0; + int64_t total_row_tiles = 0; int max_smem_row_words = 1; for (int i = 0; i < num_tensors; i++) { - if (!per_tile_swizzle_supported(tensors[i], rowwise)) return false; args.tensors[i] = tensors[i]; - args.words_per_row_magic[i] = division_magic(tensors[i].padded_k >> 2); - if (rowwise) { - total_tiles += tensors[i].padded_m / 128; - max_smem_row_words = max(max_smem_row_words, (tensors[i].padded_k >> 2) | 1); - } else { - total_tiles += - static_cast(tensors[i].padded_m / 128) * ((tensors[i].padded_k + 31) / 32); - } + args.words_per_row_magic[i] = division_magic(tensors[i].padded_k / kScalesPerWord); + total_row_tiles += tensors[i].padded_m / kSwizzleTileDimM; + max_smem_row_words = max(max_smem_row_words, (tensors[i].padded_k / kScalesPerWord) | 1); + } + int row_tiles_per_block = 1; + const int max_row_tiles_per_block = + kRowwiseMaxSharedMemoryBytes / (kSwizzleTileDimM * max_smem_row_words * kBytesPerWord); + while (total_row_tiles > static_cast(kRowwiseMaxBlocks) * row_tiles_per_block && + row_tiles_per_block * 2 <= max_row_tiles_per_block) { + row_tiles_per_block *= 2; + } + args.row_tiles_per_block = row_tiles_per_block; + for (int i = 0; i < num_tensors; i++) { + const SwizzleTensorDescriptor& tensor = tensors[i]; + const int blocks = DIVUP(tensor.padded_m / kSwizzleTileDimM, row_tiles_per_block); + args.first_block[i + 1] = args.first_block[i] + blocks; } - int tiles_per_block = 1, chunk_rows = 32; - if (rowwise) { - const int max_tiles_per_block = (48 * 1024) / (128 * max_smem_row_words * 4); - while (total_tiles > static_cast(kRowwiseMaxBlocks) * tiles_per_block && - tiles_per_block * 2 <= max_tiles_per_block) { - tiles_per_block *= 2; - } - } else { - if (total_tiles > kColumnwiseMaxBlocksWith32RowChunks) chunk_rows = 128; + for (int i = num_tensors + 1; i <= kMaxTensorsPerLaunch; i++) { + args.first_block[i] = args.first_block[num_tensors]; } - args.tiles_per_block = tiles_per_block; + const int grid = args.first_block[num_tensors]; + if (grid == 0) return; + multi_tensor_swizzle_rowwise_full_width_kernel<<< + grid, kRowwiseThreadsPerBlock, + row_tiles_per_block * kSwizzleTileDimM * max_smem_row_words * kBytesPerWord, stream>>>(args); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// Launches the columnwise chunk kernel, where each block stages part of K for one 128-column M +// tile, for up to kMaxTensorsPerLaunch tensors. It uses 128-row K chunks once 32-row chunks would +// exceed kColumnwiseMaxBlocksWith32RowChunks blocks. +inline void launch_multi_tensor_columnwise_chunk_swizzle(const SwizzleTensorDescriptor* tensors, + int num_tensors, cudaStream_t stream) { + MultiTensorSwizzleArgs args{}; + args.num_tensors = num_tensors; + args.first_block[0] = 0; + int64_t total_chunks = 0; for (int i = 0; i < num_tensors; i++) { - const PerTileSwizzleTensor& tensor = tensors[i]; - const int blocks = - rowwise ? (tensor.padded_m / 128 + tiles_per_block - 1) / tiles_per_block - : (tensor.padded_m / 128) * ((tensor.padded_k + chunk_rows - 1) / chunk_rows); + args.tensors[i] = tensors[i]; + total_chunks += static_cast(tensors[i].padded_m / kSwizzleTileDimM) * + DIVUP(tensors[i].padded_k, kColumnwiseSmallChunkRows); + } + const int chunk_rows = total_chunks > kColumnwiseMaxBlocksWith32RowChunks + ? kColumnwiseLargeChunkRows + : kColumnwiseSmallChunkRows; + for (int i = 0; i < num_tensors; i++) { + const SwizzleTensorDescriptor& tensor = tensors[i]; + const int blocks = (tensor.padded_m / kSwizzleTileDimM) * DIVUP(tensor.padded_k, chunk_rows); args.first_block[i + 1] = args.first_block[i] + blocks; } for (int i = num_tensors + 1; i <= kMaxTensorsPerLaunch; i++) { args.first_block[i] = args.first_block[num_tensors]; } const int grid = args.first_block[num_tensors]; - if (grid == 0) return true; - if (rowwise) { - multi_tensor_swizzle_rowwise_per_tile_kernel<<< - grid, kRowwiseThreadsPerBlock, tiles_per_block * 128 * max_smem_row_words * 4, stream>>>( - args); - } else if (chunk_rows == 128) { - multi_tensor_swizzle_columnwise_per_tile_kernel<128> + if (grid == 0) return; + if (chunk_rows == kColumnwiseLargeChunkRows) { + multi_tensor_swizzle_columnwise_chunk_kernel <<>>(args); } else { - multi_tensor_swizzle_columnwise_per_tile_kernel<32> + multi_tensor_swizzle_columnwise_chunk_kernel <<>>(args); } NVTE_CHECK_CUDA(cudaGetLastError()); - return true; } -inline bool launch_rowwise_per_tile_swizzle(const PerTileSwizzleTensor& tensor, - cudaStream_t stream) { - if (!per_tile_swizzle_supported(tensor, true)) return false; - swizzle_rowwise_per_tile_kernel<<> 2) | 1) * 4, stream>>>( - tensor.input, tensor.output, tensor.padded_k, division_magic(tensor.padded_k >> 2), - tensor.original_m, tensor.original_k); +// Each block stages the full padded K dimension for one 128-row tile. +inline void launch_rowwise_full_width_swizzle(const SwizzleTensorDescriptor& tensor, + cudaStream_t stream) { + const int grid = tensor.padded_m / kSwizzleTileDimM; + const int smem = kSwizzleTileDimM * ((tensor.padded_k / kScalesPerWord) | 1) * kBytesPerWord; + swizzle_rowwise_full_width_kernel<<>>( + tensor.input, tensor.output, tensor.padded_k, + division_magic(tensor.padded_k / kScalesPerWord), tensor.original_m, tensor.original_k); NVTE_CHECK_CUDA(cudaGetLastError()); - return true; } -inline bool launch_columnwise_per_tile_swizzle(const PerTileSwizzleTensor& tensor, - cudaStream_t stream) { - if (!per_tile_swizzle_supported(tensor, false)) return false; - dim3 grid(tensor.padded_m / 128, (tensor.padded_k + 31) / 32); - swizzle_columnwise_per_tile_kernel<<>>( +// Each block stages a 32-row K chunk for one 128-column M tile. +inline void launch_columnwise_chunk_swizzle(const SwizzleTensorDescriptor& tensor, + cudaStream_t stream) { + dim3 grid(tensor.padded_m / kSwizzleTileDimM, DIVUP(tensor.padded_k, kColumnwiseSmallChunkRows)); + swizzle_columnwise_chunk_kernel<<>>( tensor.input, tensor.output, tensor.padded_m, tensor.padded_k, tensor.original_m, tensor.original_k); NVTE_CHECK_CUDA(cudaGetLastError()); - return true; -} - -// A single tensor uses the one-tile-per-block kernels unless its grid would be large; several -// tensors go through one multi-tensor launch. -inline bool launch_per_tile_swizzle(const PerTileSwizzleTensor* tensors, int num_tensors, - bool rowwise, cudaStream_t stream) { - const bool grid_too_large = - num_tensors == 1 && (rowwise ? tensors[0].padded_m / 128 > kRowwiseMaxBlocks - : static_cast(tensors[0].padded_m / 128) * - ((tensors[0].padded_k + 31) / 32) > - kColumnwiseMaxBlocksWith32RowChunks); - if (num_tensors == 1 && !grid_too_large) { - return rowwise ? launch_rowwise_per_tile_swizzle(tensors[0], stream) - : launch_columnwise_per_tile_swizzle(tensors[0], stream); - } - return launch_multi_tensor_per_tile_swizzle(tensors, num_tensors, rowwise, stream); } } // namespace @@ -1275,7 +1053,6 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING, "Input tensor has invalid scaling mode (", to_string(input->scaling_mode), ")."); - // Check tensors CheckInputTensor(*input, "scaling_factor_input"); CheckInputTensor(*output, "scaling_factor_output"); NVTE_CHECK(!input->with_gemm_swizzled_scales, @@ -1341,9 +1118,10 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_ERROR("Invalid scaling mode"); } - // Check dims constexpr int SF_TILE_DIM_M = 128; constexpr int SF_TILE_DIM_K = 4; + + // Check dims NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Input should be padded in M/N dimension!"); NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Input should be padded in K dimension!"); @@ -1367,6 +1145,9 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s return; } + check_scale_inv_alignment(input, "Input", has_rowwise_scale_inv); + check_scale_inv_alignment(output, "Output", has_rowwise_scale_inv); + // Choose swizzle implementation bool rowwise_swizzle{false}, columnwise_swizzle{false}; switch (scaling_mode) { @@ -1420,31 +1201,24 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_ERROR("Invalid scaling mode"); } - const int narrow_k_slm_size = - TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - const PerTileSwizzleTensor per_tile_tensor{static_cast(input_scale_inv_ptr), - static_cast(output_scale_inv_ptr), - m, - k, - original_M, - original_K}; - if (scaling_mode == NVTE_MXFP8_1D_SCALING && use_per_tile_swizzle() && - per_tile_swizzle_supported(per_tile_tensor, true)) { - launch_per_tile_swizzle(&per_tile_tensor, 1, true, stream); - } else if (num_tiles_k < TB_DIM && static_cast(narrow_k_slm_size) <= - cuda::max_shared_memory_per_block_optin()) { - // Narrow-K: batch TB_DIM M-tiles per block, fully utilizing all threads. - dim3 num_blocks_narrow(DIVUP(num_tiles_m, TB_DIM)); - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_row_scaling_narrow_k_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_k_slm_size)); - swizzle_row_scaling_narrow_k_kernel - <<>>( - input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); + int vec_load_size = (num_tiles_k - 1) % 4 + 1; + /* there is no int3 and misaligned if using int4/int2 */ + if (vec_load_size == 3) vec_load_size = 1; + // Shared by the row-wise dispatch heuristic and optimized kernels. + const SwizzleTensorDescriptor descriptor{static_cast(input_scale_inv_ptr), + static_cast(output_scale_inv_ptr), + m, + k, + original_M, + original_K}; + if (should_use_rowwise_full_width_swizzle(descriptor.padded_k, + descriptor.padded_m / SF_TILE_DIM_M)) { + if (m / SF_TILE_DIM_M > kRowwiseMaxBlocks) { + launch_multi_tensor_rowwise_full_width_swizzle(&descriptor, 1, stream); + } else { + launch_rowwise_full_width_swizzle(descriptor, stream); + } } else { - int vec_load_size = (num_tiles_k - 1) % 4 + 1; - /* there is no int3 and misaligned if using int4/int2 */ - if (vec_load_size == 3) vec_load_size = 1; int n_tiles_in_tb = TB_DIM * vec_load_size; dim3 num_blocks(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m); int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); @@ -1487,77 +1261,27 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s const int original_M = input->flat_last_dim(); const int original_K = input->flat_first_dim() / MXFP8_BLOCK_SIZE; - const int narrow_m_slm_size = - TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - const PerTileSwizzleTensor per_tile_tensor{ + const SwizzleTensorDescriptor descriptor{ static_cast(input->columnwise_scale_inv.dptr), static_cast(output->columnwise_scale_inv.dptr), m, k, original_M, original_K}; - if (use_per_tile_swizzle() && per_tile_swizzle_supported(per_tile_tensor, false)) { - launch_per_tile_swizzle(&per_tile_tensor, 1, false, stream); - } else if (num_tiles_m < TB_DIM && static_cast(narrow_m_slm_size) <= - cuda::max_shared_memory_per_block_optin()) { - // Narrow-M: batch TB_DIM K-tiles per block, fully utilizing all threads. - dim3 num_blocks_narrow(DIVUP(num_tiles_k, TB_DIM)); - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_narrow_m_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_m_slm_size)); - swizzle_col_scaling_narrow_m_kernel - <<>>( - input->columnwise_scale_inv.dptr, output->columnwise_scale_inv.dptr, m, k, original_M, - original_K); + const int64_t grid_size = + static_cast(m / SF_TILE_DIM_M) * DIVUP(k, kColumnwiseSmallChunkRows); + if (grid_size > kColumnwiseMaxBlocksWith32RowChunks) { + launch_multi_tensor_columnwise_chunk_swizzle(&descriptor, 1, stream); } else { - int vec_load_size = (num_tiles_m - 1) % 4 + 1; - if (vec_load_size == 3) vec_load_size = 1; /* no int3 and misaligned if using int4/int2 */ - int n_tiles_in_tb = TB_DIM * vec_load_size; - dim3 num_blocks(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size)); - int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, - k, original_M, original_K); - break; - case 2: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, - k, original_M, original_K); - break; - case 1: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, - k, original_M, original_K); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; - } + launch_columnwise_chunk_swizzle(descriptor, stream); } NVTE_CHECK_CUDA(cudaGetLastError()); } } template -void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, - const int vec_load_size, const bool is_rowwise, - const bool use_narrow_k, const bool use_narrow_m, - cudaStream_t stream) { +void launch_multi_tensor_row_scaling_swizzle(MultiSwizzleArgs& kernel_args, const int vec_load_size, + cudaStream_t stream) { // cudaFuncSetAttribute is a host-synchronous driver call; cache the max shared memory // setting per kernel variant so we only pay the cost when slm_size actually increases. auto set_smem_if_needed = [](auto kernel_fn, int slm, int& cached) { @@ -1570,125 +1294,40 @@ void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, dim3 block_size(TB_DIM, TB_DIM); - if (is_rowwise && use_narrow_k) { - // Narrow-K path: each block handles TB_DIM M-tiles with full thread utilization. - // slm_size depends on num_tiles_k, which can vary per tensor — use the max. - int max_num_tiles_k = 0; - for (size_t j = 0; j < kernel_args.num_tensors; j++) { - const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; - const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; - max_num_tiles_k = std::max(max_num_tiles_k, num_tiles_k); - kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_m, TB_DIM); - } - int slm_size = TB_DIM * max_num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; - - static int cached_narrow_k = -1; - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_narrow_k_kernel, slm_size, - cached_narrow_k); - multi_tensor_swizzle_row_scaling_narrow_k_kernel - <<>>(kernel_args); - } else if (!is_rowwise && use_narrow_m) { - // Narrow-M path: each block handles TB_DIM K-tiles with full thread utilization. - // slm_size depends on num_tiles_m, which can vary per tensor — use the max. - int max_num_tiles_m = 0; - for (size_t j = 0; j < kernel_args.num_tensors; j++) { - const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; - const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; - max_num_tiles_m = std::max(max_num_tiles_m, num_tiles_m); - kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_k, TB_DIM); - } - int slm_size = TB_DIM * max_num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; - - static int cached_narrow_m = -1; - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_narrow_m_kernel, slm_size, - cached_narrow_m); - multi_tensor_swizzle_col_scaling_narrow_m_kernel - <<>>(kernel_args); - } else { - int n_tiles_in_tb = TB_DIM * vec_load_size; - int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - /* Calculate number of CUDA blocks needed for each tensor. - * We have to do it here because we have to iterate over all tensors in this batch to - * get the minimum vec_load_size. - */ - for (size_t j = 0; j < kernel_args.num_tensors; j++) { - const int m = kernel_args.m_list[j]; - const int k = kernel_args.k_list[j]; - int num_tiles_m = m / SF_TILE_DIM_M; - int num_tiles_k = k / SF_TILE_DIM_K; - if (is_rowwise) { - kernel_args.block_range[j + 1] = - kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m; - } else { - kernel_args.block_range[j + 1] = - kernel_args.block_range[j] + - DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size); - } - } - const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; - - static int cached_row_int4 = -1, cached_row_int2 = -1, cached_row_int1 = -1; - static int cached_col_int4 = -1, cached_col_int2 = -1, cached_col_int1 = -1; - - if (is_rowwise) { - switch (vec_load_size) { - case 4: - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_kernel, slm_size, - cached_row_int4); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - case 2: - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_kernel, slm_size, - cached_row_int2); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - case 1: - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_kernel, slm_size, - cached_row_int1); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; - } - } else { - switch (vec_load_size) { - case 4: - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_kernel, slm_size, - cached_col_int4); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - case 2: - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_kernel, slm_size, - cached_col_int2); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - case 1: - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_kernel, slm_size, - cached_col_int1); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; - } - } + int n_tiles_in_tb = TB_DIM * vec_load_size; + int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + for (size_t j = 0; j < kernel_args.num_tensors; j++) { + const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; + const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; + kernel_args.block_range[j + 1] = + kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m; + } + const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; + + static int cached_row_int4 = -1, cached_row_int2 = -1, cached_row_int1 = -1; + switch (vec_load_size) { + case 4: + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_kernel, slm_size, + cached_row_int4); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 2: + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_kernel, slm_size, + cached_row_int2); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 1: + set_smem_if_needed(multi_tensor_swizzle_row_scaling_kernel, + slm_size, cached_row_int1); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -1778,36 +1417,50 @@ void launch_multi_tensor_unswizzle_scaling_factors(MultiSwizzleArgs& kernel_args namespace { -// MXFP8 multi-tensor swizzle with the per-tile kernels. Returns false without launching anything -// if any tensor is not supported (the caller then uses the generic kernels, which also report -// shape errors). -bool multi_tensor_per_tile_swizzle(const std::vector& input, std::vector& output, - bool rowwise, cudaStream_t stream) { +// Multi-tensor swizzle using rowwise full-width or columnwise chunk kernels with aligned int4 +// (16-byte) loads and stores. +void int4_vectorized_multi_tensor_swizzle(const std::vector& input, + std::vector& output, bool rowwise, + cudaStream_t stream) { const size_t num_tensors = input.size(); - std::vector tensors(num_tensors); + std::vector tensors(num_tensors); for (size_t i = 0; i < num_tensors; i++) { - if (!is_mxfp8_scaling(input[i]->scaling_mode)) return false; - const auto& in_scales = rowwise ? input[i]->scale_inv : input[i]->columnwise_scale_inv; - const auto& out_scales = rowwise ? output[i]->scale_inv : output[i]->columnwise_scale_inv; - if (in_scales.shape.size() != 2) return false; - const int padded_m = static_cast(rowwise ? in_scales.shape[0] : in_scales.shape[1]); - const int padded_k = static_cast(rowwise ? in_scales.shape[1] : in_scales.shape[0]); - if (padded_m <= 0 || padded_k <= 0 || padded_m % 128 != 0 || padded_k % 4 != 0) return false; - if (out_scales.numel() != static_cast(padded_m) * padded_k) return false; + const bool nvfp4 = is_nvfp4_scaling(input[i]->scaling_mode); + const bool use_rowwise_buffer = input[i]->scale_inv.has_data(); + const bool selected_rowwise = rowwise && use_rowwise_buffer; + const auto& in_scales = selected_rowwise ? input[i]->scale_inv : input[i]->columnwise_scale_inv; + const auto& out_scales = + selected_rowwise ? output[i]->scale_inv : output[i]->columnwise_scale_inv; + NVTE_CHECK(in_scales.shape.size() == 2, "Expected two-dimensional scaling factors."); + const int padded_m = + static_cast((rowwise || nvfp4) ? in_scales.shape[0] : in_scales.shape[1]); + const int padded_k = + static_cast((rowwise || nvfp4) ? in_scales.shape[1] : in_scales.shape[0]); + NVTE_CHECK(padded_m > 0 && padded_k > 0 && padded_m % kSwizzleTileDimM == 0 && + padded_k % kScalesPerWord == 0, + "Invalid scaling-factor shape (", padded_m, ", ", padded_k, ")."); + NVTE_CHECK(in_scales.numel() == static_cast(padded_m) * padded_k && + out_scales.numel() == static_cast(padded_m) * padded_k, + "Input and output scaling-factor shapes must match."); const auto [first_dim, last_dim] = input[i]->flat_2d_dims(); + const int block_size = nvfp4 ? NVFP4_BLOCK_SIZE : MXFP8_BLOCK_SIZE; + const int original_m = static_cast(selected_rowwise ? first_dim : last_dim); + const int original_k = static_cast((selected_rowwise ? last_dim : first_dim) / block_size); tensors[i] = {static_cast(in_scales.dptr), static_cast(out_scales.dptr), padded_m, padded_k, - static_cast(rowwise ? first_dim : last_dim), - static_cast((rowwise ? last_dim : first_dim) / MXFP8_BLOCK_SIZE)}; - if (!per_tile_swizzle_supported(tensors[i], rowwise)) return false; + original_m, + original_k}; } for (size_t first = 0; first < num_tensors; first += kMaxTensorsPerLaunch) { const int count = static_cast(std::min(kMaxTensorsPerLaunch, num_tensors - first)); - launch_per_tile_swizzle(tensors.data() + first, count, rowwise, stream); + if (rowwise) { + launch_multi_tensor_rowwise_full_width_swizzle(tensors.data() + first, count, stream); + } else { + launch_multi_tensor_columnwise_chunk_swizzle(tensors.data() + first, count, stream); + } } - return true; } } // namespace @@ -1847,11 +1500,30 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, const bool rowwise_swizzle = all_has_data || all_nvfp4; const bool columnwise_swizzle = all_has_columnwise_data && !all_nvfp4; - if (!all_nvfp4 && use_per_tile_swizzle() && - multi_tensor_per_tile_swizzle(input, output, rowwise_swizzle, stream)) { + int max_rowwise_padded_k = 0; + int64_t total_rowwise_tiles = 0; + for (size_t i = 0; i < num_tensors; i++) { + const bool use_rowwise_buffer = input[i]->scale_inv.has_data(); + const bool selected_rowwise = rowwise_swizzle && use_rowwise_buffer; + const auto& in_scales = selected_rowwise ? input[i]->scale_inv : input[i]->columnwise_scale_inv; + check_scale_inv_alignment(input[i], "Input", selected_rowwise); + check_scale_inv_alignment(output[i], "Output", selected_rowwise); + if (rowwise_swizzle) { + NVTE_CHECK(in_scales.shape.size() == 2, "Expected two-dimensional scaling factors."); + max_rowwise_padded_k = std::max(max_rowwise_padded_k, static_cast(in_scales.shape[1])); + total_rowwise_tiles += static_cast(in_scales.shape[0]) / kSwizzleTileDimM; + } + } + if (rowwise_swizzle) { + if (should_use_rowwise_full_width_swizzle(max_rowwise_padded_k, total_rowwise_tiles)) { + int4_vectorized_multi_tensor_swizzle(input, output, true, stream); + return; + } + } else if (columnwise_swizzle) { + int4_vectorized_multi_tensor_swizzle(input, output, false, stream); return; } - + // Fallback to the generic rowwise swizzle kernel. constexpr int SF_TILE_DIM_M = 128; constexpr int SF_TILE_DIM_K = 4; if (rowwise_swizzle) { @@ -1859,18 +1531,16 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args.num_tensors = 0; kernel_args.block_range[0] = 0; int vec_load_size = 4; - bool all_narrow_k = true; for (size_t i = 0; i < num_tensors; i++) { //Launch kernel if argument struct is full if (kernel_args.num_tensors == kMaxTensorsPerKernel) { // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, true, all_narrow_k, false, stream); + launch_multi_tensor_row_scaling_swizzle( + kernel_args, vec_load_size, stream); // Reset the argument struct and vec_load_size kernel_args.num_tensors = 0; vec_load_size = 4; - all_narrow_k = true; } int m, k; @@ -1904,11 +1574,6 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, } int num_tiles_k = k / SF_TILE_DIM_K; - const int narrow_k_slm = - TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - all_narrow_k = - all_narrow_k && (num_tiles_k < TB_DIM) && - (static_cast(narrow_k_slm) <= cuda::max_shared_memory_per_block_optin()); int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; // We use the minimum vec_load_size across all tensors. // TODO(zhongbo): fix vec_load_size for NVFP4 @@ -1938,69 +1603,8 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, // Launch the remaining tensors // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, true, all_narrow_k, false, stream); - } - - if (columnwise_swizzle) { - // NVFP4 shouldn't end up here because it only needs rowwise swizzle - NVTE_CHECK(!all_nvfp4, "NVFP4 shouldn't end up here because it only needs rowwise swizzle"); - - MultiSwizzleArgs kernel_args; - kernel_args.num_tensors = 0; - kernel_args.block_range[0] = 0; - int vec_load_size = 4; - bool all_narrow_m = true; - for (size_t i = 0; i < num_tensors; i++) { - //Launch kernel if argument struct is full - if (kernel_args.num_tensors == kMaxTensorsPerKernel) { - // There is no int3 and misaligned if using int4/int2. - if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, false, false, all_narrow_m, stream); - // Reset the argument struct and vec_load_size - kernel_args.num_tensors = 0; - vec_load_size = 4; - all_narrow_m = true; - } - const int m = input[i]->columnwise_scale_inv.shape[1]; - const int k = input[i]->columnwise_scale_inv.shape[0]; - - NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Input should be padded in M/N dimension!"); - NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Input should be padded in K dimension!"); - NVTE_CHECK(k > 0, "Input scale inverse should be 2D!"); - NVTE_CHECK(m * k == std::accumulate(output[i]->columnwise_scale_inv.shape.begin(), - output[i]->columnwise_scale_inv.shape.end(), 1, - std::multiplies()), - "Input.columnwise_scale_inv size is not equal to " - "Output.columnwise_scale_inv size!"); - - int num_tiles_m = m / SF_TILE_DIM_M; - int num_tiles_k = k / SF_TILE_DIM_K; - const int narrow_m_slm = - TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - all_narrow_m = - all_narrow_m && (num_tiles_m < TB_DIM) && - (static_cast(narrow_m_slm) <= cuda::max_shared_memory_per_block_optin()); - int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; - // We use the minimum vec_load_size across all tensors. - vec_load_size = std::min(vec_load_size, vec_load_size_i); - - const int pos = kernel_args.num_tensors; - kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr); - kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr; - kernel_args.m_list[pos] = m; - kernel_args.k_list[pos] = k; - const auto [first_dim, last_dim] = input[i]->flat_2d_dims(); - kernel_args.original_m_list[pos] = last_dim; - kernel_args.original_k_list[pos] = first_dim / MXFP8_BLOCK_SIZE; - kernel_args.num_tensors++; - } - // Launch the remaining tensors - // There is no int3 and misaligned if using int4/int2. - if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, false, false, all_narrow_m, stream); + launch_multi_tensor_row_scaling_swizzle(kernel_args, + vec_load_size, stream); } } @@ -2406,8 +2010,8 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) size_t m = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); size_t k = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); - size_t padded_m = round_up_to_multiple(m, 128); - size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + size_t padded_m = DIVUP_TO_MULTIPLE(m, 128); + size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); int num_tiles_m = padded_m / SF_TILE_DIM_M; int num_tiles_k = padded_k / SF_TILE_DIM_K; @@ -2446,8 +2050,8 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) M = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); K = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); - size_t padded_m = round_up_to_multiple(M, 128); - size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4); + size_t padded_m = DIVUP_TO_MULTIPLE(M, 128); + size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4); int num_tiles_m = padded_m / SF_TILE_DIM_M; int num_tiles_k = padded_k / SF_TILE_DIM_K; @@ -2475,9 +2079,9 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) const uint8_t* input_base = reinterpret_cast(input) + current_scale_base; uint8_t* output_base = reinterpret_cast(output) + current_scale_base; - const int padded_m = static_cast(round_up_to_multiple(M, 128)); + const int padded_m = static_cast(DIVUP_TO_MULTIPLE(M, 128)); const int padded_k = - static_cast(round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4)); + static_cast(DIVUP_TO_MULTIPLE(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4)); const int original_M = static_cast(M); const int original_K = static_cast(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE))); const bool padding_m = (block_y == grid_dim_y - 1) && (original_M < padded_m); @@ -2577,9 +2181,8 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* auto launch_grouped_swizzle = [&](bool rowwise) { const size_t m = rowwise ? first_dim : last_dim; const size_t k = rowwise ? last_dim : first_dim; - const size_t padded_m = round_up_to_multiple(m, 128); - const size_t padded_k = - round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + const size_t padded_m = DIVUP_TO_MULTIPLE(m, 128); + const size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); // Per-tensor scale-element counts: // - "padded" layout: each tensor occupies padded_m * padded_k elements // (total buffer = num_tensors * padded_m * padded_k). @@ -2599,8 +2202,8 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* const size_t compact_scale_elems = rowwise ? m * padded_k : DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)) * padded_m; const size_t compact_total_scale_elems = - rowwise ? round_up_to_multiple(input->num_tensors * m, 128) * padded_k - : round_up_to_multiple( + rowwise ? DIVUP_TO_MULTIPLE(input->num_tensors * m, 128) * padded_k + : DIVUP_TO_MULTIPLE( input->num_tensors * DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4) * padded_m; @@ -2765,9 +2368,8 @@ void unswizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor auto launch_grouped_unswizzle = [&](bool rowwise) { const size_t m = rowwise ? first_dim : last_dim; const size_t k = rowwise ? last_dim : first_dim; - const size_t padded_m = round_up_to_multiple(m, 128); - const size_t padded_k = - round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + const size_t padded_m = DIVUP_TO_MULTIPLE(m, 128); + const size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); const size_t scale_elems = padded_m * padded_k; const size_t scale_elem_size = rowwise ? typeToSize(output->scale_inv.dtype) From 390cd037f0a29dc39e9d9ab6f38c96cc06c59715 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 28 Oct 2026 13:35:41 -0700 Subject: [PATCH 16/32] Fix GH CI documenation (#3770) Fix GH CI docs Signed-off-by: Kirthi Shankar Sivamani --- .../debug/features/fake_quant.py | 2 +- .../debug/features/log_fp8_tensor_stats.py | 2 +- .../debug/features/per_tensor_scaling.py | 2 +- .../debug/features/utils/stats_buffer.py | 2 +- transformer_engine/pytorch/__init__.py | 19 +++------------ .../dot_product_attention/backends.py | 2 +- .../dot_product_attention/context_parallel.py | 2 +- .../dot_product_attention.py | 2 +- .../dot_product_attention/flex_attention.py | 2 +- .../dot_product_attention/softmax.py | 2 +- .../attention/dot_product_attention/utils.py | 2 +- .../pytorch/attention/multi_head_attention.py | 2 +- transformer_engine/pytorch/attention/rope.py | 2 +- .../pytorch/cpp_extensions/gemm.py | 2 +- transformer_engine/pytorch/cross_entropy.py | 2 +- .../custom_recipes/quantizer_factories.py | 2 +- .../custom_recipes/quantizer_factory_zoo.py | 2 +- .../reference_current_scaling.py | 2 +- transformer_engine/pytorch/device.py | 23 +++++++++++++++++++ transformer_engine/pytorch/distributed.py | 2 +- .../pytorch/dynamo/tensor_spec.py | 2 +- transformer_engine/pytorch/jit.py | 2 +- transformer_engine/pytorch/module/base.py | 2 +- .../pytorch/module/grouped_linear.py | 2 +- .../pytorch/module/layernorm_linear.py | 2 +- .../pytorch/module/layernorm_mlp.py | 2 +- transformer_engine/pytorch/module/linear.py | 2 +- transformer_engine/pytorch/ops/_common.py | 2 +- transformer_engine/pytorch/ops/basic/bias.py | 2 +- .../ops/fused/userbuffers_backward_linear.py | 2 +- .../ops/fused/userbuffers_forward_linear.py | 2 +- .../pytorch/optimizers/fused_adam.py | 2 +- transformer_engine/pytorch/permutation.py | 2 +- transformer_engine/pytorch/quantization.py | 2 +- .../pytorch/quantized_tensor.py | 2 +- .../pytorch/tensor/float8_blockwise_tensor.py | 2 +- .../pytorch/tensor/float8_tensor.py | 2 +- .../pytorch/tensor/grouped_tensor.py | 2 +- .../pytorch/tensor/identity_tensor.py | 2 +- .../pytorch/tensor/mxfp8_tensor.py | 2 +- .../pytorch/tensor/nvfp4_tensor.py | 2 +- .../tensor/storage/mxfp8_tensor_storage.py | 2 +- .../tensor/storage/nvfp4_tensor_storage.py | 2 +- transformer_engine/pytorch/transformer.py | 2 +- transformer_engine/pytorch/triton/mhc.py | 2 +- .../pytorch/triton/permutation.py | 2 +- transformer_engine/pytorch/utils.py | 2 +- 47 files changed, 71 insertions(+), 61 deletions(-) create mode 100644 transformer_engine/pytorch/device.py diff --git a/transformer_engine/debug/features/fake_quant.py b/transformer_engine/debug/features/fake_quant.py index 570c0aadf73..720725c8456 100644 --- a/transformer_engine/debug/features/fake_quant.py +++ b/transformer_engine/debug/features/fake_quant.py @@ -13,7 +13,7 @@ from nvdlfw_inspect.utils import append_parent_docstring -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.debug.features.api import TEConfigAPIMapper from transformer_engine.common.recipe import Format from transformer_engine.pytorch import DType diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index efc0d8793b3..cd6ca7657c1 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -18,7 +18,7 @@ from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Quantizer, Float8CurrentScalingQuantizer, diff --git a/transformer_engine/debug/features/per_tensor_scaling.py b/transformer_engine/debug/features/per_tensor_scaling.py index f83b2795042..b9af84ee777 100644 --- a/transformer_engine/debug/features/per_tensor_scaling.py +++ b/transformer_engine/debug/features/per_tensor_scaling.py @@ -11,7 +11,7 @@ import nvdlfw_inspect.api as debug_api from nvdlfw_inspect.registry import Registry, api_method -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch import DType from transformer_engine.pytorch.tensor import Quantizer from transformer_engine.pytorch.tensor.float8_tensor import ( diff --git a/transformer_engine/debug/features/utils/stats_buffer.py b/transformer_engine/debug/features/utils/stats_buffer.py index 51cc5a0c1f2..666e2e2e932 100644 --- a/transformer_engine/debug/features/utils/stats_buffer.py +++ b/transformer_engine/debug/features/utils/stats_buffer.py @@ -16,7 +16,7 @@ from nvdlfw_inspect.utils import gather_along_first_dim from nvdlfw_inspect.logging import MetricLogger -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.debug.features.utils.stats_computation import ( STATS, DEPENDENCIES, diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 52f95340c44..02ce5cba9a8 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -18,22 +18,9 @@ load_framework_extension("torch") -# Device type and platform globals — live here because torch cannot be imported -# at the top-level transformer_engine package (shared with JAX). -# Use plugin-set values if available, otherwise fall back to CUDA defaults. -transformer_engine.TE_DEVICE_TYPE = getattr(transformer_engine, "TE_DEVICE_TYPE", "cuda") -transformer_engine.TE_PLATFORM = getattr(transformer_engine, "TE_PLATFORM", torch.cuda) - - -def te_device_type() -> str: - """Torch device type used by Transformer Engine, e.g. `cuda`.""" - return transformer_engine.TE_DEVICE_TYPE - - -def te_platform(): - """Torch device module used by Transformer Engine, e.g. `torch.cuda`.""" - return transformer_engine.TE_PLATFORM - +# Defined in a standalone module so that internal modules can import them without +# going through this package's __init__. +from transformer_engine.pytorch.device import te_device_type, te_platform transformer_engine.te_device_type = te_device_type transformer_engine.te_platform = te_platform diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 717bcd6a461..b4716f70407 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -19,7 +19,7 @@ import torch import torch.nn.functional as F -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 0cdd590817a..42fea51cecf 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -8,7 +8,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.utils import ( get_cudnn_version, nvtx_range_pop, 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 f3a44708df7..c07884bf415 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 @@ -15,7 +15,7 @@ from torch.fx.experimental.symbolic_shapes import guard_scalar from torch.nn.parameter import Parameter -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.common.recipe import ( Format, Recipe, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index 2c8198881de..8ba48682b42 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -10,7 +10,7 @@ from typing import Any, Callable, Dict, Optional, Tuple import torch -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type _cudnn_score_mod_handles: Dict[torch.device, Any] = {} _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index b55a960481d..613d7a6b9a9 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -10,7 +10,7 @@ from torch import nn import transformer_engine_torch as tex from transformer_engine.pytorch.export import is_in_onnx_export_mode -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type THREADS_PER_WARP = 32 THREADS_PER_BLOCK = 128 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 67a768d70c4..c215d0729b3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -22,7 +22,7 @@ import torch.nn.functional as F import transformer_engine_torch as tex import transformer_engine as te -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.cpp_extensions.fused_attn import ( QKVLayout, QKVFormat, diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 93ed14ad413..0e5af900141 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -14,7 +14,7 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm from transformer_engine.pytorch.ops.basic.l2normalization import L2Normalization -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.utils import ( SplitAlongDim, divide, diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index 18397f954b3..83369547f1f 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -10,7 +10,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type __all__ = ["RotaryPositionEmbedding", "apply_rotary_pos_emb", "apply_fused_qkv_rotary_pos_emb"] diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index bc670b7fec9..8acc547c821 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -11,7 +11,7 @@ import functools import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ..constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType from ..utils import ceil_div, get_cached_ones_tensor, get_sm_count, _empty_tensor diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 2318969963b..e6ab2b9a10f 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -10,7 +10,7 @@ import torch import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.jit import no_torch_dynamo __all__ = ["parallel_cross_entropy"] diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py index b6d22901dda..9cdec5b8922 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py @@ -29,7 +29,7 @@ import torch from transformer_engine.pytorch.quantization import QuantizerRole -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ..constants import DType diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 5cf573d61e4..3f074df177e 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -72,7 +72,7 @@ from typing import Optional from transformer_engine.pytorch.quantization import QuantizerRole -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ..constants import DType from .quantizer_factories import mxfp8_factory, nvfp4_factory diff --git a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index 6d016500c6a..52f192e4340 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -13,7 +13,7 @@ from transformer_engine.pytorch.custom_recipes import gemm from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type def current_scaling_ref_factory(role): diff --git a/transformer_engine/pytorch/device.py b/transformer_engine/pytorch/device.py new file mode 100644 index 00000000000..551cc1cb0d3 --- /dev/null +++ b/transformer_engine/pytorch/device.py @@ -0,0 +1,23 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Device type and platform used by Transformer Engine""" +import torch +import transformer_engine + +# Device type and platform globals — live here because torch cannot be imported +# at the top-level transformer_engine package (shared with JAX). +# Use plugin-set values if available, otherwise fall back to CUDA defaults. +transformer_engine.TE_DEVICE_TYPE = getattr(transformer_engine, "TE_DEVICE_TYPE", "cuda") +transformer_engine.TE_PLATFORM = getattr(transformer_engine, "TE_PLATFORM", torch.cuda) + + +def te_device_type() -> str: + """Torch device type used by Transformer Engine, e.g. `cuda`.""" + return transformer_engine.TE_DEVICE_TYPE + + +def te_platform(): + """Torch device module used by Transformer Engine, e.g. `torch.cuda`.""" + return transformer_engine.TE_PLATFORM diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 9b59970ba88..916ef2bfb15 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -31,7 +31,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.triton.pad import pad_columnwise_scale_inv -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from .torch_version import torch_version from .utils import ( is_non_tn_fp8_gemm_supported, diff --git a/transformer_engine/pytorch/dynamo/tensor_spec.py b/transformer_engine/pytorch/dynamo/tensor_spec.py index e5afa0fdd0b..f023dd081db 100644 --- a/transformer_engine/pytorch/dynamo/tensor_spec.py +++ b/transformer_engine/pytorch/dynamo/tensor_spec.py @@ -11,7 +11,7 @@ import torch from torch._prims_common import make_contiguous_strides_for -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type @dataclass diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 8052f65bc2b..0e28ae15f20 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -10,7 +10,7 @@ from typing import Callable, Optional, Tuple import torch -from .. import te_device_type +from .device import te_device_type from .torch_version import torch_version from .export import is_in_onnx_export_mode from .utils import gpu_autocast_ctx diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 2e998761154..9e90bcd12c9 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -21,7 +21,7 @@ import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ._common import _ParameterInitMeta, noop_cat, sum_bias_grad from .._extra_state import ( extra_state_pickle_advisory, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 76eded6fa84..a2e1c4210b7 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -17,7 +17,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe, Format as RecipeFormat -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.tensor.grouped_tensor import ( GroupedTensor, GroupedTensorStorage, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 03e6cf50de8..b4a91b2a3ea 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -17,7 +17,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_cublasmp_all_gather_output, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 92ae0c9ac23..ba15f1c7bd1 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -18,7 +18,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_cublasmp_all_gather_output, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 79768cd4f40..082e88a2c76 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -16,7 +16,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_cublasmp_all_gather_output, diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index e86b0b5a778..8d9e343499f 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -11,7 +11,7 @@ import torch from transformer_engine_torch import FP8TensorMeta -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ..dynamo.tensor_spec import TensorSpec from ..ep import ( EpBuffer, diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 0a26e0460b9..c5a9dda4f6c 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -11,7 +11,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ...dynamo import TensorOrQuantized, TensorSpec from ...dynamo.fake import fake_bgrad from ...quantization import FP8GlobalStateManager diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 39b2b6fe136..c924f9d7b16 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -11,7 +11,7 @@ import torch from transformer_engine_torch import CommOverlapType, bulk_overlap_ag_with_external_gemm -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ...cpp_extensions import general_gemm from ...distributed import gather_along_first_dim, get_distributed_world_size from ...module.base import ( diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 9bc2cfc1e75..f2f9f91d105 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -11,7 +11,7 @@ import torch from transformer_engine_torch import CommOverlapType -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ...cpp_extensions import general_gemm from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import get_distributed_world_size diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 2cde7780af0..3cd5fcd387b 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -15,7 +15,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ..constants import DType from .multi_tensor_apply import multi_tensor_applier diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index 2cb49d4a9b4..5e9da706d87 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -7,7 +7,7 @@ from typing import Optional, Tuple import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type import transformer_engine.pytorch.triton.permutation as triton_permutation from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch.quantized_tensor import ( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 06d02ece44e..355e13f0989 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -17,7 +17,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.common.recipe import ( Recipe, DelayedScaling, diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 794d383ff26..9c44dae9bf3 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -19,7 +19,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.constants import dist_group_type -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index c85577fe2f8..ef0eb2a672b 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -12,7 +12,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import ( Float8BlockwiseQTensorStorage, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index d37b8e2800e..7eede479fb9 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -10,7 +10,7 @@ from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.common.recipe import ( DelayedScaling, Float8CurrentScaling, diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 83c3816af84..3aa1a21737e 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -10,7 +10,7 @@ import torch from torch.utils._pytree import tree_map -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ..quantized_tensor import QuantizedTensorStorage, Quantizer from .storage.grouped_tensor_storage import GroupedTensorStorage from ..constants import DType diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 8f8bda44cbc..31d3de1c391 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -19,7 +19,7 @@ import torch from torch.ops import aten -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from .storage.identity_tensor_storage import IdentityTensorStorage from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index c4081b75c64..d97c05fe37d 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -13,7 +13,7 @@ from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe from ..constants import MXFP8_BLOCK_SCALING_SIZE, DType from ..utils import devices_match, round_up_to_nearest_multiple diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 29be015b6ec..43d5280fdb5 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -13,7 +13,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType from ..utils import ( diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index cccd048cb83..4d18560dd26 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -12,7 +12,7 @@ import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 9931325eec3..84ceb2a076e 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -15,7 +15,7 @@ import transformer_engine_torch as tex -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index cc361454a4c..77a11ff401f 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -15,7 +15,7 @@ from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention from transformer_engine.pytorch.attention.inference import InferenceParams -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.jit import ( set_jit_fusion_options, warmup_jit_bias_dropout_add_all_dtypes, diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py index ce9a507b448..180c9be78e0 100644 --- a/transformer_engine/pytorch/triton/mhc.py +++ b/transformer_engine/pytorch/triton/mhc.py @@ -25,7 +25,7 @@ _mhc_sinkhorn_bwd_fused, ) from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type ENFORCE_DETERMINISTIC = os.environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") == "0" diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 64d3c89128e..33a827c280e 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -9,7 +9,7 @@ import torch import triton -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from transformer_engine.common.triton.permutation import ( _row_id_map_pass_1_kernel, _row_id_map_pass_2_kernel, diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index a3208db5026..0a027468d29 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -14,7 +14,7 @@ import numpy as np import torch -from transformer_engine import te_device_type +from transformer_engine.pytorch.device import te_device_type from .torch_version import torch_version from ..common import decode_cudnn_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor From e619cea17e6fe48d05b900e3a7a01531623773b3 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Wed, 28 Oct 2026 13:40:43 -0700 Subject: [PATCH 17/32] [PyTorch][CP] Make THD CP helper kernels to be shared memory size aware (#3698) * [CUDA] Handle THD boundary capacity across CP utility kernels THD metadata can exceed per-block shared memory even when the token tensors are valid. Use device-derived default and opt-in capacities before reading the existing global boundary arrays, without allocating scratch buffers. Centralize staging and normalization so indexing helpers consume consistent boundaries while retaining pre-scaled shared reads. Preserve the current old_lse output-accumulation behavior and exercise its rescaling in boundary regression tests. Focused native validation passed 348 boundary/dtype cases, 60 output-rescaling cases including 24 exact upstream comparisons, 33 graph captures with three replays each, and 324 partition stress cases up to one million documents. SM86 and SM100 compilation passed. Signed-off-by: Sudhakar Singh * [PyTorch] Guard THD boundary tests on available device limits Older PyTorch builds do not expose the shared-memory properties used to choose exact boundary-test sizes. Check both capabilities inside the two affected tests and skip only those tests when either field is missing, without guessing a version cutoff or querying CUDA during import. Include the running PyTorch version in the skip reason to make compatibility skips identifiable. Production kernels and the remaining functional tests are unchanged. Missing/present-property checks and both targeted native boundary tests passed. Signed-off-by: Sudhakar Singh * Use runtime storage selection for THD boundary tables Avoid duplicating each dtype, layout, and functor specialization solely for boundary storage. Pass a block-uniform storage flag and plain pointers with an explicit read divisor, preserving shared pre-scaling and immutable global metadata. Use the device-reported default shared-memory budget and the existing global table as fallback, removing kernel-specific opt-in state and the boundary-view type. Retain per-boundary integer division and CP validation, and test the two remaining storage regimes. Signed-off-by: Sudhakar Singh * Restore compile-time storage selection for THD boundaries Runtime storage selection obscures the shared address space and unit read divisor, adding work to boundary searches. Restore shared/global kernel specializations so the compiler can eliminate global-only normalization from the shared path. Keep plain pointer and explicit divisor arguments instead of the boundary view. Retain the default shared-memory budget, immutable global fallback, and the existing default-boundary regression tests without restoring shared-memory opt-in. Signed-off-by: Sudhakar Singh --------- Signed-off-by: Sudhakar Singh --- tests/pytorch/attention/test_cp_utils.py | 110 ++++++ .../common/fused_attn/context_parallel.cu | 345 ++++++++++-------- 2 files changed, 312 insertions(+), 143 deletions(-) diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py index 25e1fcbfe74..52813026a0b 100644 --- a/tests/pytorch/attention/test_cp_utils.py +++ b/tests/pytorch/attention/test_cp_utils.py @@ -1017,6 +1017,116 @@ def test_thd_get_partitioned_indices_matches_dual_chunk_expected_indices(self): self.assertTrue(torch.equal(rank0, expected_rank0)) self.assertTrue(torch.equal(rank1, expected_rank1)) + def test_thd_get_partitioned_indices_shared_memory_boundary(self): + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + if not hasattr(props, "shared_memory_per_block"): + self.skipTest(f"PyTorch {torch.__version__} does not expose shared-memory limits") + max_shared_docs = props.shared_memory_per_block // 4 - 1 + for documents in (max_shared_docs, max_shared_docs + 1): + for cp_size in (2, 4): + cu_seqlens = torch.arange(documents + 1, device="cuda", dtype=torch.int32) + cu_seqlens *= 2 * cp_size + tokens = torch.arange( + documents * 2 * cp_size, device="cuda", dtype=torch.int32 + ).view(documents, 2 * cp_size) + for rank in range(cp_size): + with self.subTest(documents=documents, cp_size=cp_size, rank=rank): + indices = tex.thd_get_partitioned_indices( + cu_seqlens, tokens.numel(), cp_size, rank + ) + expected = tokens[:, [rank, 2 * cp_size - 1 - rank]].reshape(-1) + torch.cuda.synchronize() + self.assertTrue(torch.equal(indices, expected)) + + def test_thd_kernels_shared_memory_boundary(self): + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + if not hasattr(props, "shared_memory_per_block"): + self.skipTest(f"PyTorch {torch.__version__} does not expose shared-memory limits") + for tables in (1, 2): + bytes_per_document = 4 * tables + for documents in ( + props.shared_memory_per_block // bytes_per_document - 1, + props.shared_memory_per_block // bytes_per_document, + ): + with self.subTest(tables=tables, documents=documents): + cu = torch.arange(documents + 1, dtype=torch.int32, device="cuda") * 4 + total = documents * 4 + x = torch.arange(total * 4, dtype=torch.float32, device="cuda").view( + total, 1, 4 + ) + docs = x.view(documents, 4, 1, 4) + if tables == 2: + out = torch.full_like(x, -1) + expected = out.clone().view_as(docs) + expected[:, :2] = docs[:, :2] + tex.thd_copy_valid_tokens_from_per_split_to_rank_local(out, x, cu, cu // 2) + self.assertTrue(torch.equal(out, expected.reshape_as(out))) + continue + + reordered = tex.thd_sequence_order_to_cp_rank_order(x, cu, 2, total) + expected = torch.cat( + [docs[:, [rank, 3 - rank]].reshape(-1, 1, 4) for rank in range(2)] + ) + self.assertTrue(torch.equal(reordered, expected)) + restored = tex.thd_cp_rank_order_to_sequence_order(reordered, cu, 2, total) + self.assertTrue(torch.equal(restored, x)) + for half in (0, 1): + actual = tex.thd_read_half_tensor(x, cu, half) + expected = docs[:, 2 * half : 2 * half + 2].reshape_as(actual) + self.assertTrue(torch.equal(actual, expected)) + + for packed in (False, True): + shape = (1, total) if packed else (documents, 1, 4) + half_shape = (1, total // 2) if packed else (documents, 1, 2) + lse = torch.arange(total, dtype=torch.float32, device="cuda").view(shape) + lse /= total + half_lse = tex.thd_read_second_half_lse( + lse, cu, packed, total // 2 if packed else 2 + ) + expected_half = lse.reshape(documents, 4)[:, 2:].reshape(half_shape) + self.assertTrue(torch.equal(half_lse, expected_half)) + expected = lse.clone() + expected.reshape(documents, 4)[:, 2:] = torch.logaddexp( + expected_half.reshape(documents, 2), + torch.zeros_like(expected_half).reshape(documents, 2), + ) + tex.thd_second_half_lse_correction( + lse, torch.zeros_like(half_lse), cu, packed + ) + torch.testing.assert_close(lse, expected) + + for second_half in (False, True): + out = torch.full_like(x, 3) + step = torch.ones_like(x[: total // (2 if second_half else 1)]) + old_lse = torch.full(shape, 0.5, device="cuda").log() + expected = out.clone().view_as(docs) + expected[:, 2 if second_half else 0 :] *= 0.5 + expected[:, 2 if second_half else 0 :] += 1 + tex.thd_out_correction( + out, + step, + old_lse, + torch.zeros(shape, device="cuda"), + torch.zeros(half_shape if second_half else shape, device="cuda"), + cu, + second_half, + packed, + ) + torch.testing.assert_close(out, expected.reshape_as(out)) + + grad = x.clone() + tex.thd_grad_correction(grad, torch.ones_like(x), cu, "add", "copy") + expected = docs.clone() + expected[:, :2] += 1 + expected[:, 2:] = 1 + self.assertTrue(torch.equal(grad, expected.reshape_as(grad))) + grad.fill_(-1) + tex.thd_grad_correction(grad, x[: total // 2], cu, "copy", "zero") + expected.zero_() + expected[:, :2] = x[: total // 2].view(documents, 2, 1, 4) + torch.cuda.synchronize() + self.assertTrue(torch.equal(grad, expected.reshape_as(grad))) + def test_thd_copy_valid_tokens_from_per_split_matches_legacy_slice_copy_loop(self): cu_seqlens_padded = torch.tensor([2, 6, 12], dtype=torch.int32, device="cuda") cu_seqlens = torch.tensor([0, 3, 7], dtype=torch.int32, device="cuda") diff --git a/transformer_engine/common/fused_attn/context_parallel.cu b/transformer_engine/common/fused_attn/context_parallel.cu index 69a15dddd7b..9da561abfa4 100644 --- a/transformer_engine/common/fused_attn/context_parallel.cu +++ b/transformer_engine/common/fused_attn/context_parallel.cu @@ -14,6 +14,24 @@ namespace transformer_engine { namespace context_parallel { +// Return the actual launch allocation, or zero when boundaries must stay in global memory. +static size_t thd_boundary_shared_bytes(size_t requested) { + // These kernels have no static shared memory; the default limit is the full boundary budget. + static const std::vector limits = [] { + int count; + NVTE_CHECK_CUDA(cudaGetDeviceCount(&count)); + std::vector result(count); + for (int device = 0; device < count; ++device) { + NVTE_CHECK_CUDA( + cudaDeviceGetAttribute(&result[device], cudaDevAttrMaxSharedMemoryPerBlock, device)); + } + return result; + }(); + int device; + NVTE_CHECK_CUDA(cudaGetDevice(&device)); + return requested <= static_cast(limits[device]) ? requested : 0; +} + struct LseCorrectionFunctor { __forceinline__ __device__ static void run(float *lse, float *half_lse, size_t idx, size_t half_idx) { @@ -70,11 +88,12 @@ struct AddFunctor { * Support THD format for Context Parallel: Binary search an array for a target value **************************************************************************************************/ -__forceinline__ __device__ int binary_search(int target, int *array, int len) { +__forceinline__ __device__ int binary_search(int target, const int *array, int len, + int read_divisor) { int left = 1, right = len - 1; while (left < right) { int mid = (left + right) / 2; - if (array[mid] <= target) { + if (array[mid] / read_divisor <= target) { left = mid + 1; } else { right = mid; @@ -83,38 +102,45 @@ __forceinline__ __device__ int binary_search(int target, int *array, int len) { return left - 1; } -// Dual-chunk source index for THD CP partitioning. cu_seqlens_s must already be divided by -// world_size. Single source of truth shared by thd_partition_indices_kernel and -// thd_reorder_between_sequence_and_cp_rank_order_kernel so the two never diverge. -__forceinline__ __device__ int thd_partition_src_index(int token_id, int *cu_seqlens_s, int batch, - int world_size, int rank) { - int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1); - int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id]; - int index = token_id - cu_seqlens_s[seq_id]; +// Shared boundaries are pre-scaled (read_divisor=1); global boundaries remain unmodified. +__forceinline__ __device__ int thd_partition_src_index(int token_id, const int *boundaries, + int batch, int world_size, int rank, + int read_divisor) { + int seq_id = binary_search(token_id, boundaries, batch + 1, read_divisor); + int seq_start = boundaries[seq_id] / read_divisor; + int seq_len = boundaries[seq_id + 1] / read_divisor - seq_start; + int index = token_id - seq_start; int offset = index < seq_len / 2 ? rank : (world_size - 1) * 2 - rank; - return index + cu_seqlens_s[seq_id] * world_size + seq_len / 2 * offset; + return index + seq_start * world_size + seq_len / 2 * offset; } /*************************************************************************************************** * Support THD format for Context Parallel: Generate partitioned indices for input tokens **************************************************************************************************/ +template __global__ void thd_partition_indices_kernel(int *output, int *cu_seqlens, int batch, int total_tokens, int world_size, int rank) { - extern __shared__ int cu_seqlens_s[]; - for (int i = threadIdx.x; i <= batch; i += blockDim.x) { - int seqlen = cu_seqlens[i]; - // Currently we assume that each sequence length is divisible by (world_size*2) since we have - // to distribute each sequence evenly to different GPUs. - assert(seqlen % (world_size * 2) == 0); - cu_seqlens_s[i] = seqlen / world_size; + extern __shared__ int boundary_storage[]; + if constexpr (use_shared) { + for (int i = threadIdx.x; i <= batch; i += blockDim.x) { + int seqlen = cu_seqlens[i]; + assert(seqlen % (world_size * 2) == 0); + boundary_storage[i] = seqlen / world_size; + } + __syncthreads(); + } else { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i <= batch; i += blockDim.x * gridDim.x) { + assert(cu_seqlens[i] % (world_size * 2) == 0); + } } - __syncthreads(); + const int *boundaries = use_shared ? boundary_storage : cu_seqlens; + const int read_divisor = use_shared ? 1 : world_size; int tid = blockIdx.x * blockDim.x + threadIdx.x; int num_threads = blockDim.x * gridDim.x; - for (int token_id = tid; token_id < total_tokens / world_size; token_id += num_threads) { - output[token_id] = thd_partition_src_index(token_id, cu_seqlens_s, batch, world_size, rank); + output[token_id] = + thd_partition_src_index(token_id, boundaries, batch, world_size, rank, read_divisor); } } @@ -124,14 +150,19 @@ __global__ void thd_partition_indices_kernel(int *output, int *cu_seqlens, int b * cp_rank_to_sequence_order=false: out[gi]=inp[src(gi)]. * cp_rank_to_sequence_order=true: out[src(gi)]=inp[gi]. **************************************************************************************************/ +template __global__ void thd_reorder_between_sequence_and_cp_rank_order_kernel( void *out, void *inp, int *cu_seqlens, int batch, int total_tokens, int world_size, int hidden_size_in_bytes, bool cp_rank_to_sequence_order) { - extern __shared__ int cu_seqlens_s[]; - for (int i = threadIdx.x; i <= batch; i += blockDim.x) { - cu_seqlens_s[i] = cu_seqlens[i] / world_size; + extern __shared__ int boundary_storage[]; + if constexpr (use_shared) { + for (int i = threadIdx.x; i <= batch; i += blockDim.x) { + boundary_storage[i] = cu_seqlens[i] / world_size; + } + __syncthreads(); } - __syncthreads(); + const int *cu_seqlens_s = use_shared ? boundary_storage : cu_seqlens; + const int read_divisor = use_shared ? 1 : world_size; int warpid = (blockIdx.x * blockDim.x + threadIdx.x) / 32; int laneid = threadIdx.x % 32; @@ -142,7 +173,8 @@ __global__ void thd_reorder_between_sequence_and_cp_rank_order_kernel( for (int gi = warpid; gi < total_tokens; gi += num_warps) { int rank = gi / tpr; int token_id = gi % tpr; - int src = thd_partition_src_index(token_id, cu_seqlens_s, batch, world_size, rank); + int src = + thd_partition_src_index(token_id, cu_seqlens_s, batch, world_size, rank, read_divisor); int rd = cp_rank_to_sequence_order ? gi : src; int wr = cp_rank_to_sequence_order ? src : gi; float4 *src_tok = reinterpret_cast(reinterpret_cast(inp) + @@ -157,16 +189,21 @@ __global__ void thd_reorder_between_sequence_and_cp_rank_order_kernel( * Copy valid token entries from a per-split THD tensor into a rank-local accumulator. * cu_seqlens_padded gives padded THD token offsets; cu_seqlens gives valid lengths. **************************************************************************************************/ +template __global__ void thd_copy_valid_tokens_from_per_split_to_rank_local_kernel( void *out, void *inp, int *cu_seqlens_padded, int *cu_seqlens, int batch, int total_tokens, int hidden_size_in_bytes) { - extern __shared__ int padded_s[]; // [0..batch] padded boundaries - int *valid_s = padded_s + (batch + 1); // [0..batch] valid boundaries - for (int i = threadIdx.x; i <= batch; i += blockDim.x) { - padded_s[i] = cu_seqlens_padded[i]; - valid_s[i] = cu_seqlens[i]; + extern __shared__ int boundary_storage[]; + int *valid_storage = use_shared ? boundary_storage + (batch + 1) : cu_seqlens; + if constexpr (use_shared) { + for (int i = threadIdx.x; i <= batch; i += blockDim.x) { + boundary_storage[i] = cu_seqlens_padded[i]; + valid_storage[i] = cu_seqlens[i]; + } + __syncthreads(); } - __syncthreads(); + const int *padded_s = use_shared ? boundary_storage : cu_seqlens_padded; + const int *valid_s = valid_storage; int warpid = (blockIdx.x * blockDim.x + threadIdx.x) / 32; int laneid = threadIdx.x % 32; @@ -174,7 +211,7 @@ __global__ void thd_copy_valid_tokens_from_per_split_to_rank_local_kernel( int num_float4s_per_token = hidden_size_in_bytes / sizeof(float4); for (int token_id = warpid; token_id < total_tokens; token_id += num_warps) { - int seq_id = binary_search(token_id, padded_s, batch + 1); + int seq_id = binary_search(token_id, padded_s, batch + 1, 1); if (seq_id < 0 || seq_id >= batch) continue; int local = token_id - padded_s[seq_id]; int valid_len = valid_s[seq_id + 1] - valid_s[seq_id]; @@ -193,19 +230,24 @@ __global__ void thd_copy_valid_tokens_from_per_split_to_rank_local_kernel( * Support THD format for Context Parallel: Read the half of a THD tensor **************************************************************************************************/ +template __global__ void thd_read_half_tensor_kernel(void *half, void *tensor, int *cu_seqlens, int batch, int hidden_size_in_bytes, int half_idx, int dim_size_of_token) { - extern __shared__ int cu_seqlens_s[]; - for (int i = threadIdx.x; i <= batch; i += blockDim.x) { - cu_seqlens_s[i] = cu_seqlens[i] / 2; + extern __shared__ int boundary_storage[]; + if constexpr (use_shared) { + for (int i = threadIdx.x; i <= batch; i += blockDim.x) { + boundary_storage[i] = cu_seqlens[i] / 2; + } + __syncthreads(); } - __syncthreads(); + const int *cu_seqlens_s = use_shared ? boundary_storage : cu_seqlens; + const int read_divisor = use_shared ? 1 : 2; int warpid = (blockIdx.x * blockDim.x + threadIdx.x) / 32; int laneid = threadIdx.x % 32; int num_warps = (blockDim.x * gridDim.x) / 32; - int num_total_tokens = cu_seqlens_s[batch]; + int num_total_tokens = cu_seqlens_s[batch] / read_divisor; int num_float4s_per_token = hidden_size_in_bytes / sizeof(float4); size_t offset = static_cast(dim_size_of_token) * hidden_size_in_bytes; @@ -213,14 +255,15 @@ __global__ void thd_read_half_tensor_kernel(void *half, void *tensor, int *cu_se tensor = reinterpret_cast(reinterpret_cast(tensor) + offset * blockIdx.y); for (int token_id = warpid; token_id < num_total_tokens; token_id += num_warps) { - int seqid = binary_search(token_id, cu_seqlens_s, batch + 1); + int seqid = binary_search(token_id, cu_seqlens_s, batch + 1, read_divisor); size_t offset_in_bytes = static_cast(token_id) * hidden_size_in_bytes; float4 *cur_half_token = reinterpret_cast(reinterpret_cast(half) + offset_in_bytes); offset_in_bytes = - (static_cast(token_id) + cu_seqlens_s[seqid + half_idx]) * hidden_size_in_bytes; + (static_cast(token_id) + cu_seqlens_s[seqid + half_idx] / read_divisor) * + hidden_size_in_bytes; float4 *cur_token = reinterpret_cast(reinterpret_cast(tensor) + offset_in_bytes); @@ -234,30 +277,34 @@ __global__ void thd_read_half_tensor_kernel(void *half, void *tensor, int *cu_se * Support THD format for Context Parallel: softmax_lse related operations **************************************************************************************************/ -template +template __global__ void thd_lse_kernel(float *lse, float *half_lse, int *cu_seqlens, int batch, int num_heads, int lse_seqlen, int second_half_lse_seqlen) { - extern __shared__ int cu_seqlens_s[]; - for (int i = threadIdx.x; i <= batch; i += blockDim.x) { - cu_seqlens_s[i] = cu_seqlens[i] / 2; + extern __shared__ int boundary_storage[]; + if constexpr (use_shared) { + for (int i = threadIdx.x; i <= batch; i += blockDim.x) { + boundary_storage[i] = cu_seqlens[i] / 2; + } + __syncthreads(); } - __syncthreads(); + const int *cu_seqlens_s = use_shared ? boundary_storage : cu_seqlens; + const int read_divisor = use_shared ? 1 : 2; int tid = blockIdx.x * blockDim.x + threadIdx.x; int num_threads = blockDim.x * gridDim.x; - int num_total_tokens = cu_seqlens_s[batch]; + int num_total_tokens = cu_seqlens_s[batch] / read_divisor; for (int token_id = tid; token_id < num_total_tokens; token_id += num_threads) { - int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1); + int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1, read_divisor); for (int head_id = blockIdx.y; head_id < num_heads; head_id += gridDim.y) { size_t idx, half_idx; if constexpr (lse_packed) { - idx = head_id * lse_seqlen + token_id + cu_seqlens_s[seq_id + 1]; + idx = head_id * lse_seqlen + token_id + cu_seqlens_s[seq_id + 1] / read_divisor; half_idx = head_id * second_half_lse_seqlen + token_id; } else { size_t row = static_cast(seq_id) * num_heads + head_id; - int col = token_id - cu_seqlens_s[seq_id]; - int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id]; + int col = token_id - cu_seqlens_s[seq_id] / read_divisor; + int seq_len = cu_seqlens_s[seq_id + 1] / read_divisor - cu_seqlens_s[seq_id] / read_divisor; idx = row * lse_seqlen + col + seq_len; half_idx = row * second_half_lse_seqlen + col; @@ -272,42 +319,47 @@ __global__ void thd_lse_kernel(float *lse, float *half_lse, int *cu_seqlens, int * Support THD format for Context Parallel: Out correction in forward **************************************************************************************************/ -template +template __global__ void thd_out_correction_kernel(dtype *out, dtype *out_per_step, float *old_lse, float *lse, float *lse_per_step, int *cu_seqlens, int batch, int num_heads, int dim_per_head, int lse_seqlen, int lse_per_step_seqlen) { - extern __shared__ int cu_seqlens_s[]; - for (int i = threadIdx.x; i <= batch; i += blockDim.x) { - cu_seqlens_s[i] = cu_seqlens[i] / (only_second_half + 1); + extern __shared__ int boundary_storage[]; + if constexpr (use_shared) { + for (int i = threadIdx.x; i <= batch; i += blockDim.x) { + boundary_storage[i] = cu_seqlens[i] / (only_second_half + 1); + } + __syncthreads(); } - __syncthreads(); + const int *cu_seqlens_s = use_shared ? boundary_storage : cu_seqlens; + const int read_divisor = use_shared ? 1 : only_second_half + 1; int tile_id = (blockIdx.x * blockDim.x + threadIdx.x) / tile_size; int lane_id = threadIdx.x % tile_size; int num_tiles = (blockDim.x * gridDim.x) / tile_size; - int num_total_tokens = cu_seqlens_s[batch]; + int num_total_tokens = cu_seqlens_s[batch] / read_divisor; int num_loops_per_head = dim_per_head * sizeof(dtype) / sizeof(float4); for (int token_id = tile_id; token_id < num_total_tokens; token_id += num_tiles) { - int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1); + int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1, read_divisor); for (int head_id = blockIdx.y; head_id < num_heads; head_id += gridDim.y) { size_t idx, idx_per_step; if constexpr (lse_packed) { - idx = head_id * lse_seqlen + token_id + cu_seqlens_s[seq_id + 1] * only_second_half; + idx = head_id * lse_seqlen + token_id + + (cu_seqlens_s[seq_id + 1] / read_divisor) * only_second_half; idx_per_step = head_id * lse_per_step_seqlen + token_id; } else { size_t row = static_cast(seq_id) * num_heads + head_id; - int col = token_id - cu_seqlens_s[seq_id]; - int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id]; + int col = token_id - cu_seqlens_s[seq_id] / read_divisor; + int seq_len = cu_seqlens_s[seq_id + 1] / read_divisor - cu_seqlens_s[seq_id] / read_divisor; idx = row * lse_seqlen + col + seq_len * only_second_half; idx_per_step = row * lse_per_step_seqlen + col; } float old_lse_corrected_exp = expf(old_lse[idx] - lse[idx]); float lse_corrected_exp = expf(lse_per_step[idx_per_step] - lse[idx]); - idx = token_id + cu_seqlens_s[seq_id + 1] * only_second_half; + idx = token_id + (cu_seqlens_s[seq_id + 1] / read_divisor) * only_second_half; idx = (idx * num_heads + head_id) * dim_per_head; idx_per_step = (static_cast(token_id) * num_heads + head_id) * dim_per_head; dtype *cur_out = out + idx; @@ -337,23 +389,24 @@ __global__ void thd_out_correction_kernel(dtype *out, dtype *out_per_step, float * Support THD format for Context Parallel: Gradients correction in backward **************************************************************************************************/ -template +template __global__ void thd_grad_correction_kernel(dtype *grad, dtype *grad_per_step, int *cu_seqlens, int batch, int hidden_size, int dim_size_of_token) { - extern __shared__ int cu_seqlens_s[]; - for (int i = threadIdx.x; i <= batch; i += blockDim.x) { - if constexpr (functor_idx < 2) { - cu_seqlens_s[i] = cu_seqlens[i] / 2; - } else { - cu_seqlens_s[i] = cu_seqlens[i]; + extern __shared__ int boundary_storage[]; + if constexpr (use_shared) { + for (int i = threadIdx.x; i <= batch; i += blockDim.x) { + boundary_storage[i] = cu_seqlens[i] / (functor_idx < 2 ? 2 : 1); } + __syncthreads(); } - __syncthreads(); + const int *cu_seqlens_s = use_shared ? boundary_storage : cu_seqlens; + const int read_divisor = use_shared ? 1 : (functor_idx < 2 ? 2 : 1); int group_id = (blockIdx.x * blockDim.x + threadIdx.x) / group_size; int lane_id = threadIdx.x % group_size; int num_groups = (blockDim.x * gridDim.x) / group_size; - int num_total_tokens = cu_seqlens_s[batch]; + int num_total_tokens = cu_seqlens_s[batch] / read_divisor; int num_inner_loops = hidden_size * sizeof(dtype) / sizeof(float4); size_t offset = static_cast(dim_size_of_token) * hidden_size; @@ -365,21 +418,21 @@ __global__ void thd_grad_correction_kernel(dtype *grad, dtype *grad_per_step, in grad = grad + offset * blockIdx.y; for (int token_id = group_id; token_id < num_total_tokens; token_id += num_groups) { - int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1); + int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1, read_divisor); if constexpr (functor_idx < 2) { - dtype *first_half_token = - &grad[(token_id + cu_seqlens_s[seq_id]) * static_cast(hidden_size)]; - dtype *second_half_token = - &grad[(token_id + cu_seqlens_s[seq_id + 1]) * static_cast(hidden_size)]; + dtype *first_half_token = &grad[(token_id + cu_seqlens_s[seq_id] / read_divisor) * + static_cast(hidden_size)]; + dtype *second_half_token = &grad[(token_id + cu_seqlens_s[seq_id + 1] / read_divisor) * + static_cast(hidden_size)]; dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { Functor_0::run(first_half_token, token_per_step, idx); Functor_1::run(second_half_token, token_per_step, idx); } } else { - int len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id]; - bool is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2); + int len = cu_seqlens_s[seq_id + 1] / read_divisor - cu_seqlens_s[seq_id] / read_divisor; + bool is_first_half = (token_id - cu_seqlens_s[seq_id] / read_divisor) < (len / 2); dtype *token = &grad[token_id * static_cast(hidden_size)]; dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { @@ -429,9 +482,12 @@ void thd_read_half_tensor(const Tensor &tensor, const Tensor &cu_seqlens, Tensor grid_y *= tensor_shape[i]; } dim3 grid = {grid_x, grid_y}; - thd_read_half_tensor_kernel<<>>( - half.data.dptr, tensor.data.dptr, reinterpret_cast(cu_seqlens.data.dptr), batch, - hidden_size_in_bytes, half_idx, tensor_shape[seq_dim]); + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + shared_bytes != 0, use_shared, + thd_read_half_tensor_kernel<<>>( + half.data.dptr, tensor.data.dptr, reinterpret_cast(cu_seqlens.data.dptr), batch, + hidden_size_in_bytes, half_idx, tensor_shape[seq_dim]);); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -484,19 +540,17 @@ void thd_second_half_lse_correction(Tensor lse, const Tensor &lse_per_step, unsigned int grid_y = num_heads; dim3 grid = {grid_x, grid_y}; - if (lse_packed) { - thd_lse_kernel<<>>( - reinterpret_cast(lse.data.dptr), reinterpret_cast(lse_per_step.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen, - second_half_lse_seqlen); - NVTE_CHECK_CUDA(cudaGetLastError()); - } else { - thd_lse_kernel<<>>( - reinterpret_cast(lse.data.dptr), reinterpret_cast(lse_per_step.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen, - second_half_lse_seqlen); - NVTE_CHECK_CUDA(cudaGetLastError()); - } + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + lse_packed, packed, + TRANSFORMER_ENGINE_SWITCH_CONDITION(shared_bytes != 0, use_shared, + thd_lse_kernel + <<>>( + reinterpret_cast(lse.data.dptr), + reinterpret_cast(lse_per_step.data.dptr), + reinterpret_cast(cu_seqlens.data.dptr), batch, + num_heads, lse_seqlen, second_half_lse_seqlen););); + NVTE_CHECK_CUDA(cudaGetLastError()); } void thd_read_second_half_lse(const Tensor &lse, const Tensor &cu_seqlens, Tensor &half_lse, @@ -535,19 +589,17 @@ void thd_read_second_half_lse(const Tensor &lse, const Tensor &cu_seqlens, Tenso unsigned int grid_y = num_heads; dim3 grid = {grid_x, grid_y}; - if (lse_packed) { - thd_lse_kernel<<>>( - reinterpret_cast(lse.data.dptr), reinterpret_cast(half_lse.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen, - second_half_lse_seqlen); - NVTE_CHECK_CUDA(cudaGetLastError()); - } else { - thd_lse_kernel<<>>( - reinterpret_cast(lse.data.dptr), reinterpret_cast(half_lse.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen, - second_half_lse_seqlen); - NVTE_CHECK_CUDA(cudaGetLastError()); - } + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + lse_packed, packed, + TRANSFORMER_ENGINE_SWITCH_CONDITION(shared_bytes != 0, use_shared, + thd_lse_kernel + <<>>( + reinterpret_cast(lse.data.dptr), + reinterpret_cast(half_lse.data.dptr), + reinterpret_cast(cu_seqlens.data.dptr), batch, + num_heads, lse_seqlen, second_half_lse_seqlen););); + NVTE_CHECK_CUDA(cudaGetLastError()); } /*************************************************************************************************** @@ -611,27 +663,21 @@ static void thd_out_correction_helper(Tensor out, const Tensor &out_per_step, co (static_cast(total_tokens) / (only_second_half + 1) * tile + block - 1) / block; dim3 grid = {grid_x, (unsigned int)num_heads}; - if (lse_packed) { - thd_out_correction_kernel - <<>>( - reinterpret_cast(out.data.dptr), - reinterpret_cast(out_per_step.data.dptr), - reinterpret_cast(old_lse.data.dptr), reinterpret_cast(lse.data.dptr), - reinterpret_cast(lse_per_step.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, dim_per_head, - lse_seqlen, lse_per_step_seqlen); - NVTE_CHECK_CUDA(cudaGetLastError()); - } else { - thd_out_correction_kernel - <<>>( - reinterpret_cast(out.data.dptr), - reinterpret_cast(out_per_step.data.dptr), - reinterpret_cast(old_lse.data.dptr), reinterpret_cast(lse.data.dptr), - reinterpret_cast(lse_per_step.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, dim_per_head, - lse_seqlen, lse_per_step_seqlen); - NVTE_CHECK_CUDA(cudaGetLastError()); - } + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + lse_packed, packed, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + shared_bytes != 0, use_shared, + thd_out_correction_kernel + <<>>(reinterpret_cast(out.data.dptr), + reinterpret_cast(out_per_step.data.dptr), + reinterpret_cast(old_lse.data.dptr), + reinterpret_cast(lse.data.dptr), + reinterpret_cast(lse_per_step.data.dptr), + reinterpret_cast(cu_seqlens.data.dptr), + batch, num_heads, dim_per_head, lse_seqlen, + lse_per_step_seqlen););); + NVTE_CHECK_CUDA(cudaGetLastError()); } void thd_out_correction(Tensor out, const Tensor &out_per_step, const Tensor &old_lse, @@ -700,11 +746,14 @@ static void thd_grad_correction_helper(Tensor grad, const Tensor &grad_per_step, } dim3 grid = {grid_x, grid_y}; - thd_grad_correction_kernel - <<>>( - reinterpret_cast(grad.data.dptr), - reinterpret_cast(grad_per_step.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, hidden_size, total_tokens); + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + shared_bytes != 0, use_shared, + thd_grad_correction_kernel + <<>>(reinterpret_cast(grad.data.dptr), + reinterpret_cast(grad_per_step.data.dptr), + reinterpret_cast(cu_seqlens.data.dptr), batch, + hidden_size, total_tokens);); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -787,9 +836,12 @@ void thd_get_partitioned_indices(const Tensor &cu_seqlens, Tensor output, int to constexpr unsigned int block = 256; unsigned int grid = (output_shape[0] + block - 1) / block; - thd_partition_indices_kernel<<>>( - reinterpret_cast(output.data.dptr), reinterpret_cast(cu_seqlens.data.dptr), - batch, total_tokens, world_size, rank); + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + shared_bytes != 0, use_shared, + thd_partition_indices_kernel<<>>( + reinterpret_cast(output.data.dptr), reinterpret_cast(cu_seqlens.data.dptr), + batch, total_tokens, world_size, rank);); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -814,10 +866,13 @@ void thd_reorder_between_sequence_and_cp_rank_order(const Tensor &inp, const Ten int batch = cu_seqlens_shape[0] - 1; constexpr unsigned int block = 256; unsigned int grid = (static_cast(total_tokens) * 32 + block - 1) / block; - thd_reorder_between_sequence_and_cp_rank_order_kernel<<>>( - out.data.dptr, inp.data.dptr, reinterpret_cast(cu_seqlens.data.dptr), batch, - total_tokens, world_size, hidden_size_in_bytes, cp_rank_to_sequence_order); + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + shared_bytes != 0, use_shared, + thd_reorder_between_sequence_and_cp_rank_order_kernel + <<>>( + out.data.dptr, inp.data.dptr, reinterpret_cast(cu_seqlens.data.dptr), batch, + total_tokens, world_size, hidden_size_in_bytes, cp_rank_to_sequence_order);); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -843,10 +898,14 @@ void thd_copy_valid_tokens_from_per_split_to_rank_local(const Tensor &inp, int batch = cu_seqlens_shape[0] - 1; constexpr unsigned int block = 256; unsigned int grid = (static_cast(total_tokens) * 32 + block - 1) / block; - thd_copy_valid_tokens_from_per_split_to_rank_local_kernel<<< - grid, block, sizeof(int) * 2 * (batch + 1), stream>>>( - out.data.dptr, inp.data.dptr, reinterpret_cast(cu_seqlens_padded.data.dptr), - reinterpret_cast(cu_seqlens.data.dptr), batch, total_tokens, hidden_size_in_bytes); + const size_t shared_bytes = thd_boundary_shared_bytes(sizeof(int) * 2 * (batch + 1)); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + shared_bytes != 0, use_shared, + thd_copy_valid_tokens_from_per_split_to_rank_local_kernel + <<>>(out.data.dptr, inp.data.dptr, + reinterpret_cast(cu_seqlens_padded.data.dptr), + reinterpret_cast(cu_seqlens.data.dptr), batch, + total_tokens, hidden_size_in_bytes);); NVTE_CHECK_CUDA(cudaGetLastError()); } From 54ecaabc97fc0ec8f6a004db580bac164082216a Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:11:40 -0700 Subject: [PATCH 18/32] [PyT] Guard for buggy quack version (#3762) * Fail fast on Hopper FA4 backward with affected Quack versions Signed-off-by: Kshitij Lakhani * Remove redundant Quack guard documentation Signed-off-by: Kshitij Lakhani --------- Signed-off-by: Kshitij Lakhani Co-authored-by: Kirthi Shankar Sivamani --- .../dot_product_attention/backends.py | 5 ++++ .../dot_product_attention.py | 24 +++++++++++++++++++ .../attention/dot_product_attention/utils.py | 2 ++ 3 files changed, 31 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index b4716f70407..2345a1808a6 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -295,6 +295,11 @@ def wrapper(*args, **kwargs): _flash_attn_bwd_v4 = no_torch_dynamo()(_fa4_with_none_window_sentinel(_flash_attn_bwd_v4)) fa_utils.v4_validate_head_dims = _fa4_validate_head_dims + fa_utils.cutlass_dsl_version = cutlass_dsl_version + try: + fa_utils.quack_version = PkgVersion(get_pkg_version("quack-kernels")) + except PackageNotFoundError: + fa_utils.quack_version = None fa_utils.set_flash_attention_4_params() # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 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 c07884bf415..1c7f57431ae 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 @@ -10,6 +10,7 @@ import warnings import logging +from packaging.version import Version as PkgVersion import torch import torch.nn.functional as F from torch.fx.experimental.symbolic_shapes import guard_scalar @@ -2916,6 +2917,29 @@ def forward( ) if use_flash_attention: + # DSL 4.6.2 reverted bulk-copy lane election; Quack PR #181 restores it. + quack_version = FlashAttentionUtils.quack_version + cutlass_version = FlashAttentionUtils.cutlass_dsl_version + affected_dependencies = ( + quack_version is not None + and PkgVersion("0.6.0") <= quack_version < PkgVersion("0.6.4") + and cutlass_version is not None + and cutlass_version >= PkgVersion("4.6.2") + ) + if ( + flash_attention_backend is not None + and flash_attention_backend.major == 4 + and attention_params.requires_backward + and affected_dependencies + and dpa_utils.get_device_compute_capability(query_layer.device) == (9, 0) + ): + raise RuntimeError( + f"FA4 {flash_attention_backend} backward on Hopper (SM90) can hang with " + f"quack-kernels {quack_version} and CUTLASS DSL " + f"{cutlass_version}. Upgrade to a compatible " + "Quack release containing the fix in 0.6.4: " + "https://github.com/Dao-AILab/quack/pull/181" + ) orig_qk_dim = None orig_v_dim = None if ( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index c215d0729b3..f33335be950 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -151,6 +151,8 @@ class FlashAttentionUtils: fa4_version = PkgVersion("0") v4_0_0_beta31 = PkgVersion("4.0.0b31") v4_0_0_beta33 = PkgVersion("4.0.0b33") + quack_version: Optional[PkgVersion] = None + cutlass_dsl_version: Optional[PkgVersion] = None use_v4 = False # Set by a signature probe in backends.py; fail-closed default. fa3_supports_softcap = False From dfc4158d1597f72972fe2863adac436fa3fe1d4d Mon Sep 17 00:00:00 2001 From: fheinecke <23390735+fheinecke@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:32:11 -0500 Subject: [PATCH 19/32] Link tests executables against libpython when CuTeDSL is enabled (#3769) Link tests executables against cpython when CuTeDSL is enabled Signed-off-by: Fred Heinecke --- tests/cpp/CMakeLists.txt | 19 +++++++++++-------- tests/cpp_distributed/CMakeLists.txt | 18 +++++++++++------- 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index af074be1c78..477ed522d27 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -43,17 +43,20 @@ include_directories(${CMAKE_SOURCE_DIR}) find_package(CUDAToolkit REQUIRED) include(${CMAKE_SOURCE_DIR}/../../transformer_engine/common/cmake/cuDNN.cmake) -# Must match how libtransformer_engine.so was built, hence the env var default. -option(NVTE_WITH_CUTEDSL - "Link the tests against libpython, required when libtransformer_engine.so was built with NVTE_WITH_CUTEDSL=ON" - "$ENV{NVTE_WITH_CUTEDSL}") - add_subdirectory(operator) add_subdirectory(util) -if(NVTE_WITH_CUTEDSL) - # Executables that link libtransformerengine.so must also link to the Python::Python to provide - # the definitions of these CPython symbols at runtime via dynamic linking. +# libtransformer_engine.so built with NVTE_WITH_CUTEDSL=ON leaves CPython symbols undefined, so +# executables linking it must provide them. Detect this from the library itself rather than from +# build flags, which are not available when the tests are built against an installed library. +execute_process(COMMAND ${CMAKE_NM} --dynamic --undefined-only "${TE_LIB}" + OUTPUT_VARIABLE TE_LIB_UNDEFINED_SYMBOLS + RESULT_VARIABLE TE_LIB_NM_RESULT) +if(NOT TE_LIB_NM_RESULT EQUAL 0) + message(FATAL_ERROR "Failed to list the undefined symbols of ${TE_LIB} with '${CMAKE_NM}'") +endif() +if(TE_LIB_UNDEFINED_SYMBOLS MATCHES "[ \t]U _?Py[A-Za-z_]") + message(STATUS "${TE_LIB} requires CPython symbols, linking the tests against libpython") # Find the interpreter too so older CMake versions use its version to locate libpython. find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED) foreach(test_target test_operator test_util) diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 2dd7b301c14..3d9d0b8e23e 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -56,10 +56,6 @@ find_package(CUDAToolkit REQUIRED) find_package(OpenMP REQUIRED) find_package(MPI REQUIRED) -option(NVTE_WITH_CUTEDSL - "Link the tests against libpython, required when libtransformer_engine.so was built with NVTE_WITH_CUTEDSL=ON" - "$ENV{NVTE_WITH_CUTEDSL}") - # -- NCCL core ---------------------------------------------------------------- # Anchor on libnccl and derive nccl.h from the same install prefix so the # header and library can't drift across installs. @@ -131,9 +127,17 @@ add_executable(test_ep test_ep.cu ../cpp/test_common.cu) target_include_directories(test_ep PRIVATE ${EP_TEST_COMMON_INCLUDES}) target_link_libraries(test_ep PUBLIC ${EP_TEST_COMMON_LIBS}) -if(NVTE_WITH_CUTEDSL) - # Executables that link libtransformerengine.so must also link to the Python::Python to provide - # the definitions of these CPython symbols at runtime via dynamic linking. +# libtransformer_engine.so built with NVTE_WITH_CUTEDSL=ON leaves CPython symbols undefined, so +# executables linking it must provide them. Detect this from the library itself rather than from +# build flags, which are not available when the tests are built against an installed library. +execute_process(COMMAND ${CMAKE_NM} --dynamic --undefined-only "${TE_LIB}" + OUTPUT_VARIABLE TE_LIB_UNDEFINED_SYMBOLS + RESULT_VARIABLE TE_LIB_NM_RESULT) +if(NOT TE_LIB_NM_RESULT EQUAL 0) + message(FATAL_ERROR "Failed to list the undefined symbols of ${TE_LIB} with '${CMAKE_NM}'") +endif() +if(TE_LIB_UNDEFINED_SYMBOLS MATCHES "[ \t]U _?Py[A-Za-z_]") + message(STATUS "${TE_LIB} requires CPython symbols, linking the tests against libpython") # Find the interpreter too so older CMake versions use its version to locate libpython. find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED) foreach(test_target test_comm_gemm test_ep) From 83641e8ede8633153f1fb1fe2c105b67bd9a4115 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 29 Oct 2026 13:07:03 +0100 Subject: [PATCH 20/32] [PyTorch] Add torch.compile support for LayerNormLinear and LayerNormMLP Register LayerNorm module forward and backward custom ops with matching fake tensor metadata, shared eligibility checks, and compiled weight caches. Extend the existing compile tests to cover the LayerNorm modules. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 912 ++++++++++++++---- .../dot_product_attention/backends.py | 6 +- transformer_engine/pytorch/csrc/extensions.h | 3 + .../pytorch/csrc/extensions/normalization.cpp | 23 +- .../pytorch/csrc/extensions/pybind.cpp | 2 + .../pytorch/dynamo/custom_op.py | 42 +- transformer_engine/pytorch/dynamo/fake.py | 56 ++ transformer_engine/pytorch/module/_common.py | 83 +- .../pytorch/module/layernorm_linear.py | 499 +++++++++- .../pytorch/module/layernorm_mlp.py | 620 +++++++++++- transformer_engine/pytorch/module/linear.py | 166 +--- 11 files changed, 1978 insertions(+), 434 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 6868de7ef18..66c925a25a2 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -22,6 +22,7 @@ except ImportError: # pragma: no cover counters = None from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode +from torch._dynamo.testing import CompileCounterWithBackend try: from torch._opaque_base import OpaqueBaseMeta @@ -45,6 +46,11 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx from transformer_engine.pytorch.module.base import TransformerEngineBaseModule +from transformer_engine.pytorch.module import ( + linear as linear_module, + layernorm_linear as layernorm_linear_module, + layernorm_mlp as layernorm_mlp_module, +) from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.ops.basic.bias import BiasBwdArgs @@ -60,7 +66,7 @@ Quantizer, ) from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec -from transformer_engine.pytorch.dynamo.fake import fake_bgrad +from transformer_engine.pytorch.dynamo.fake import fake_bgrad, fake_norm from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -129,7 +135,8 @@ def nvfp4_4over6(): if mxfp8_available: _all_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: - _all_recipes.append(recipe.NVFP4BlockScaling()) + # Disable stochastic rounding for exact eager/compiled comparisons. + _all_recipes.append(recipe.NVFP4BlockScaling(disable_stochastic_rounding=True)) _all_recipes.append(nvfp4_4over6()) _all_recipes.append(nvfp4_row_scaled()) @@ -175,30 +182,75 @@ def _assert_no_cudagraph_skips(enabled: bool): _EAGER_ATOL, _EAGER_RTOL = 0.0, 0.0 -def _assert_close_eager_compiled(fn, compiled, model, base): - """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert the - forward output and the input / weight / bias gradients match.""" +def _flatten_outputs(out): + return list(out) if isinstance(out, (tuple, list)) else [out] + + +def _assert_close_eager_compiled(fn, compiled, model, base, *, dynamic_batch=False): + """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert every + forward output and the input / parameter gradients match.""" inp_eager = base.detach().clone().requires_grad_(True) model.zero_grad(set_to_none=True) - out_eager = fn(inp_eager) - out_eager.sum().backward() - ref_out = out_eager.detach().clone() - ref_wgrad = model.weight.grad.detach().clone() + outs_eager = _flatten_outputs(fn(inp_eager)) + sum(o.sum() for o in outs_eager).backward() + ref_outs = [o.detach().clone() for o in outs_eager] ref_igrad = inp_eager.grad.detach().clone() - # bias=False keeps a 0-element ``bias`` around; grad is None then. - ref_bgrad = model.bias.grad.detach().clone() if model.bias.grad is not None else None + ref_pgrads = { + name: p.grad.detach().clone() for name, p in model.named_parameters() if p.grad is not None + } inp_compiled = base.detach().clone().requires_grad_(True) + if dynamic_batch: + torch._dynamo.mark_dynamic(inp_compiled, 0) model.zero_grad(set_to_none=True) # Clone before a later cuda-graph replay overwrites the static output buffer. - out_compiled = compiled(inp_compiled).clone() - out_compiled.sum().backward() + outs_compiled = [o.clone() for o in _flatten_outputs(compiled(inp_compiled))] + sum(o.sum() for o in outs_compiled).backward() - torch.testing.assert_close(out_compiled, ref_out, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + assert len(outs_compiled) == len(ref_outs) + for out, ref in zip(outs_compiled, ref_outs): + torch.testing.assert_close(out, ref, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) torch.testing.assert_close(inp_compiled.grad, ref_igrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) - torch.testing.assert_close(model.weight.grad, ref_wgrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) - if ref_bgrad is not None: - torch.testing.assert_close(model.bias.grad, ref_bgrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + pgrads = {name: p.grad for name, p in model.named_parameters() if p.grad is not None} + assert pgrads.keys() == ref_pgrads.keys() + for name, ref in ref_pgrads.items(): + torch.testing.assert_close( + pgrads[name], ref, atol=_EAGER_ATOL, rtol=_EAGER_RTOL, msg=f"grad mismatch: {name}" + ) + + +def _run_module_compile_test(model, fp8_recipe, compile_mode, in_features, dtype=torch.bfloat16): + device = "cuda" + + def fn(inp): + if fp8_recipe is None: + return model(inp) + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + warm = torch.randn(32, in_features, dtype=dtype, device=device, requires_grad=True) + outs = _flatten_outputs(fn(warm)) + sum(o.sum() for o in outs).backward() + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + base = torch.randn(32, in_features, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +_LINEAR_MODULES = ["Linear", "LayerNormLinear", "LayerNormMLP"] + + +def _make_module(module, *, dtype=torch.bfloat16, device="cuda", **kwargs): + """Small module with dimensions suitable for FP8 GEMMs.""" + return getattr(te, module)( + 64, 128 if module == "LayerNormMLP" else 32, params_dtype=dtype, device=device, **kwargs + ) # --------------------------------------------------------------------------- @@ -1898,6 +1950,75 @@ def fn(x): assert out.shape == () +@pytest.mark.parametrize("with_optional_shape", [False, True]) +def test_custom_op_symbolic_shape(with_optional_shape): + """Shape arguments stay dynamic across the custom-op boundary.""" + from transformer_engine.pytorch.dynamo.custom_op import register_custom_op_with_autograd + + @dataclasses.dataclass + class ShapeArgs: + input_: torch.Tensor + shape: torch.Size + optional_shape: torch.Size | None + + @dataclasses.dataclass + class ShapeBwdArgs: + grad_output: torch.Tensor | None = None + shape: torch.Size | None = None + + def setup_grad_outputs(self, grads): + (self.grad_output,) = grads + + def impl(args): + assert isinstance(args.shape, torch.Size) + shape = args.shape if args.optional_shape is None else args.optional_shape + assert isinstance(shape, torch.Size) + return args.input_.reshape(shape).clone(), (), None + + def fake_impl(args): + shape = args.shape if args.optional_shape is None else args.optional_shape + return TensorSpec(shape=shape, dtype=args.input_.dtype, device=args.input_.device), (), None + + def setup_context(bwd_args, fwd_args, *_): + bwd_args.shape = fwd_args.input_.shape + return () + + def backward_impl(args): + return (args.grad_output.reshape(args.shape).clone(),) + + def backward_fake(args): + grad = args.grad_output + return (TensorSpec(shape=args.shape, dtype=grad.dtype, device=grad.device),) + + op = register_custom_op_with_autograd( + op_name=f"test_symbolic_shape_{with_optional_shape}", + input_tensors_for_grad=["input_"], + fwd_arg_type=ShapeArgs, + fwd_impl=impl, + fwd_fake_impl=fake_impl, + setup_context=setup_context, + bwd_arg_type=ShapeBwdArgs, + bwd_impl=backward_impl, + bwd_fake_impl=backward_fake, + ) + assert op is not None + + def fn(inp): + shape = torch.Size((inp.shape[0], 2, 4)) + return op(ShapeArgs(inp, shape, shape if with_optional_shape else None)) + + counter = CompileCounterWithBackend("aot_eager") + compiled = torch.compile(fn, fullgraph=True, backend=counter) + for batch in (2, 3, 5): + inp = torch.randn(batch, 8, requires_grad=True) + torch._dynamo.mark_dynamic(inp, 0) + out = compiled(inp) + torch.testing.assert_close(out, inp.reshape(batch, 2, 4)) + out.sum().backward() + torch.testing.assert_close(inp.grad, torch.ones_like(inp)) + assert counter.frame_count == 1 + + def test_to_tensor_spec_plain(): """``to_tensor_spec`` describes a plain tensor.""" t = torch.empty(2, 3, dtype=torch.float32) @@ -1909,11 +2030,19 @@ def test_to_tensor_spec_plain(): @pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) -def test_to_tensor_spec_quantized(factory, shape): +@pytest.mark.parametrize("swizzled", [False, True]) +def test_to_tensor_spec_quantized(factory, shape, swizzled): """``to_tensor_spec`` round-trips a quantized tensor back into a spec.""" q = factory() + has_scale_layout = isinstance(q, (MXFP8Quantizer, NVFP4Quantizer)) + if has_scale_layout: + q.optimize_for_gemm = not swizzled tensor = TensorSpec( - shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu") + shape=shape, + dtype=torch.bfloat16, + quantizer=q, + device=torch.device("cpu"), + with_gemm_swizzled_scales=swizzled if has_scale_layout else None, ).create_tensor() spec = to_tensor_spec(tensor) @@ -1923,13 +2052,45 @@ def test_to_tensor_spec_quantized(factory, shape): # Same buffer layout as the original tensor. assert spec.inner_names() == tuple(q.inner_tensor_specs(shape)) # Rebuilding from the derived spec matches the original tensor's structure. - assert _signature(spec.create_tensor(), spec.inner_names()) == _signature( - tensor, spec.inner_names() - ) + rebuilt = spec.create_tensor() + assert _signature(rebuilt, spec.inner_names()) == _signature(tensor, spec.inner_names()) + if has_scale_layout: + assert tensor._with_gemm_swizzled_scales == swizzled + assert spec.with_gemm_swizzled_scales == swizzled + assert rebuilt._with_gemm_swizzled_scales == swizzled + assert spec.quantizer.optimize_for_gemm == q.optimize_for_gemm == (not swizzled) + + +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize("with_2d_quantization", [False, True]) +@pytest.mark.parametrize("optimize_for_gemm", [False, True]) +@pytest.mark.parametrize( + "use_cudnn,shape,compact", + [ + (True, (128, 128), True), + (True, (2, 64, 128), True), + (True, (64, 128), False), + (False, (128, 128), False), + ], +) +def test_fake_norm_mxfp8_scale_layout( + monkeypatch, normalization, with_2d_quantization, optimize_for_gemm, use_cudnn, shape, compact +): + """RMSNorm's 2D MXFP8 fallback preserves the quantizer's scale layout.""" + monkeypatch.setenv("NVTE_NORM_FWD_USE_CUDNN", str(int(use_cudnn))) + quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3, with_2d_quantization=with_2d_quantization) + quantizer.optimize_for_gemm = optimize_for_gemm + inp = TensorSpec(shape=shape, dtype=torch.bfloat16, device=torch.device("cpu")) + out, _, _ = fake_norm(inp, quantizer, normalization=normalization) + compact = compact and (normalization == "LayerNorm" or not with_2d_quantization) + expected = optimize_for_gemm and not compact + assert out.create_metadata()["nontensor_kwargs"]["with_gemm_swizzled_scales"] == expected + assert out.create_tensor()._with_gemm_swizzled_scales == expected + assert quantizer.optimize_for_gemm == optimize_for_gemm # --------------------------------------------------------------------------- -# te.Linear +# te.Linear / te.LayerNormLinear / te.LayerNormMLP # --------------------------------------------------------------------------- @@ -1940,17 +2101,48 @@ def test_to_tensor_spec_quantized(factory, shape): [None, *_all_recipes], ids=lambda r: "bf16" if r is None else type(r).__name__, ) -def test_te_linear_compiles(fp8_recipe, compile_mode): - """ - torch.compile(fullgraph=True) of ``te.Linear`` under every built-in - recipe (plus the bf16-only baseline with no autocast), for both the default - backend and ``mode="reduce-overhead"`` (CUDA-graph trees). - """ - dtype = torch.bfloat16 - device = "cuda" +@pytest.mark.parametrize("module", _LINEAR_MODULES) +def test_te_modules_compile(module, fp8_recipe, compile_mode): + """Fullgraph forward/backward agrees with eager across modules, recipes and modes.""" + if module == "LayerNormMLP" and getattr(fp8_recipe, "backward_override", None) is not None: + pytest.skip("LayerNormMLP does not support backward_override, including in eager mode") + model = _make_module(module) + _run_module_compile_test(model, fp8_recipe, compile_mode, 64) + + +_non_2d_cases = [ + pytest.param((64,), None, id="rank1-bf16"), + pytest.param((2, 4, 4, 64), None, id="rank4-bf16"), + *( + pytest.param( + (2, 4, 16, 128) if isinstance(r, recipe.Float8BlockScaling) else (2, 4, 4, 64), + r, + id=f"rank4-{recipe_id(r)}", + ) + for r in _all_recipes + ), + *( + pytest.param((2, 4, 8, 64), r, id=f"rank4-aligned-{recipe_id(r)}") + for r in _all_recipes + if isinstance(r, recipe.NVFP4BlockScaling) and not r.disable_rht + ), +] - # FP8 GEMMs require leading dimensions divisible by 16. - model = te.Linear(64, 32, params_dtype=dtype, device=device) + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("compile_mode", _compile_modes) +@pytest.mark.parametrize("module", _LINEAR_MODULES) +@pytest.mark.parametrize("input_shape, fp8_recipe", _non_2d_cases) +def test_te_modules_compile_non_2d_input(input_shape, fp8_recipe, module, compile_mode): + """Compiled modules preserve non-matrix activation shapes and gradients.""" + if module == "LayerNormMLP" and getattr(fp8_recipe, "backward_override", None) is not None: + pytest.skip("LayerNormMLP does not support backward_override, including in eager mode") + dtype = torch.bfloat16 + if isinstance(fp8_recipe, recipe.Float8BlockScaling): + out_features = 256 if module == "LayerNormMLP" else 128 + model = getattr(te, module)(128, out_features, params_dtype=dtype, device="cuda") + else: + model = _make_module(module, dtype=dtype) def fn(inp): if fp8_recipe is None: @@ -1962,17 +2154,15 @@ def fn(inp): if compile_mode == "reduce-overhead": _cudagraph_warmup( fn, - torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + torch.randn(input_shape, dtype=dtype, device="cuda", requires_grad=True), backward=True, ) model.zero_grad(set_to_none=True) compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) - # Iterate a few times so reduce-overhead actually replays a captured graph. - n_iters = 3 if compile_mode == "reduce-overhead" else 1 with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): - for _ in range(n_iters): - base = torch.randn(32, 64, dtype=dtype, device=device) + for _ in range(3 if compile_mode == "reduce-overhead" else 1): + base = torch.randn(input_shape, dtype=dtype, device="cuda") _assert_close_eager_compiled(fn, compiled, model, base) @@ -1987,19 +2177,22 @@ def fn(inp): + [ pytest.param(fp8_recipe, shape, id=f"{name}-{shape}") for name, fp8_recipe in ( + ("bf16", None), + ("mxfp8", recipe.MXFP8BlockScaling()), ("nvfp4", recipe.NVFP4BlockScaling()), ("nvfp4-no-rht", recipe.NVFP4BlockScaling(disable_rht=True)), ) - for shape in ((1, 128, 128), (2, 4, 16, 128)) + for shape in ((1, 128, 128), (128, 1, 128), (2, 4, 16, 128)) ], ) def test_te_linear_compiles_non_2d_input(fp8_recipe, input_shape, compile_mode): """Compiled Linear preserves non-matrix activation shapes in forward and backward.""" if isinstance(fp8_recipe, recipe.NVFP4BlockScaling) and not nvfp4_available: pytest.skip(reason_for_no_nvfp4) + if isinstance(fp8_recipe, recipe.MXFP8BlockScaling) and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) dtype = torch.bfloat16 - out_features = 32 if fp8_recipe is None else 128 - model = te.Linear(input_shape[-1], out_features, params_dtype=dtype, device="cuda") + model = te.Linear(input_shape[-1], 128, params_dtype=dtype, device="cuda") def fn(inp): if fp8_recipe is None: @@ -2026,37 +2219,14 @@ def fn(inp): @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("compile_mode", _compile_modes) -def test_te_linear_compile_with_quantized_fp8_weight(compile_mode): - """torch.compile of Linear with the weight initialized as an FP8 tensor - (exercises the wrapper op's ``register_torch_dispatch`` input flattening).""" - dtype = torch.bfloat16 - device = "cuda" +@pytest.mark.parametrize("module", ["Linear", "LayerNormLinear"]) +def test_te_modules_compile_with_quantized_fp8_weight(module, compile_mode): + """FP8 primary weights cross the wrapper-op boundary in both compile modes.""" fp8_recipe = recipe.Float8CurrentScaling() - with te.quantized_model_init(enabled=True, recipe=fp8_recipe): - model = te.Linear(64, 32, params_dtype=dtype, device=device) - + model = _make_module(module) assert isinstance(model.weight, te.Float8Tensor) - - def fn(inp): - with te.autocast(recipe=fp8_recipe): - return model(inp) - - torch._dynamo.reset() - if compile_mode == "reduce-overhead": - _cudagraph_warmup( - fn, - torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), - backward=True, - ) - model.zero_grad(set_to_none=True) - compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) - - n_iters = 3 if compile_mode == "reduce-overhead" else 1 - with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): - for _ in range(n_iters): - base = torch.randn(32, 64, dtype=dtype, device=device) - _assert_close_eager_compiled(fn, compiled, model, base) + _run_module_compile_test(model, fp8_recipe, compile_mode, 64) @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") @@ -2105,7 +2275,7 @@ def fn(inp): ) -# Configs rejected by LinearFwdArgs.compile_unsupported_reason() that a +# Configs rejected by the modules' early fallback checks that a # single-GPU unit test can construct. Distributed-only reasons (fsdp_group, # DistributedWeight) and CPU offloading need machinery this file doesn't have; # delayed scaling is a hard error (check_recipe_support), tested separately. @@ -2114,58 +2284,67 @@ def fn(inp): # output crossing the graph-break boundary gets a plain-tensor tangent); # "no_grad" = forward under no_grad. _FALLBACK_CASES = [ - "fp8_output_differentiable", - "fuse_wgrad_accumulation", - "delayed_wgrad", - "quantized_input", + (module, case) + for module in _LINEAR_MODULES + for case in ("fuse_wgrad_accumulation", "delayed_wgrad") +] + [ + ("Linear", "fp8_output_differentiable"), + ("Linear", "quantized_input"), + ("LayerNormMLP", "checkpoint"), ] -def _fallback_case(case, dtype, device): - """Build ``(model, fn, mode, post_backward, reason)`` for one case.""" +def _fallback_case(module, case, dtype, device): + """Build ``(model, fn, mode, post_backward, reason)`` for one fallback case.""" model_kwargs = {} if case == "fuse_wgrad_accumulation": model_kwargs["fuse_wgrad_accumulation"] = True elif case == "delayed_wgrad": model_kwargs["delay_wgrad_compute"] = True - model = te.Linear(64, 32, params_dtype=dtype, device=device, **model_kwargs) - - if case == "fp8_output_differentiable": - fp8_recipe = recipe.Float8CurrentScaling() + elif case == "checkpoint": + model_kwargs["checkpoint"] = True + model = _make_module(module, dtype=dtype, device=device, **model_kwargs) + fp8_recipe = recipe.Float8CurrentScaling() - def fn(inp): - with te.autocast(recipe=fp8_recipe): + def fn(inp): + # Preserve the original Linear BF16 accumulation/delayed-wgrad cases; + # the LayerNorm cases exercise those paths under FP8. + use_fp8 = module != "Linear" or case in ("fp8_output_differentiable", "quantized_input") + with te.autocast(enabled=use_fp8, recipe=fp8_recipe): + if case == "fp8_output_differentiable": return model(inp, fp8_output=True).dequantize() + return model(inp) + if case == "fp8_output_differentiable": return model, fn, "fwd_grad", None, "differentiable fp8_output=True" if case == "fuse_wgrad_accumulation": - model.weight.main_grad = torch.zeros_like(model.weight, dtype=torch.float32) - return model, model, "bwd", None, "fuse_wgrad_accumulation" + weights = ( + (model.fc1_weight, model.fc2_weight) if module == "LayerNormMLP" else (model.weight,) + ) + for weight in weights: + weight.main_grad = torch.zeros_like(weight, dtype=torch.float32) + return model, fn, "bwd", None, "fuse_wgrad_accumulation" if case == "delayed_wgrad": - return model, model, "bwd", model.backward_dw, "delayed wgrad compute" + return model, fn, "bwd", model.backward_dw, "delayed wgrad compute" if case == "quantized_input": - fp8_recipe = recipe.Float8CurrentScaling() - - def fn(inp): - with te.autocast(recipe=fp8_recipe): - return model(inp) - return model, fn, "no_grad", None, "a quantized input tensor" + if case == "checkpoint": + return model, fn, "bwd", None, "activation checkpointing" raise ValueError(case) @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize("case", _FALLBACK_CASES) -def test_te_linear_compile_eager_fallback(case): +@pytest.mark.parametrize("module,case", _FALLBACK_CASES) +def test_te_modules_compile_eager_fallback(module, case): """Configs unsupported on the compiled custom-op path must fall back to eager under ``torch.compile`` -- warning + numerics identical to eager -- and graph-break with the explicit reason under ``fullgraph=True``.""" dtype, device = torch.bfloat16, "cuda" torch.manual_seed(0) - model_ref, fn_ref, mode, post_bwd_ref, _ = _fallback_case(case, dtype, device) + model_ref, fn_ref, mode, post_bwd_ref, _ = _fallback_case(module, case, dtype, device) torch.manual_seed(0) - model, fn, _, post_bwd, reason = _fallback_case(case, dtype, device) + model, fn, _, post_bwd, reason = _fallback_case(module, case, dtype, device) def make_inp(): torch.manual_seed(1) @@ -2194,20 +2373,27 @@ def make_inp(): if post_bwd is not None: post_bwd() torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) - # The fallback must also preserve the stateful side effects that define - # these cases (main_grad accumulation, delayed wgrad), not just inp.grad. - if getattr(model_ref.weight, "main_grad", None) is not None: - assert model.weight.grad is None - torch.testing.assert_close( - model.weight.main_grad, - model_ref.weight.main_grad, - atol=_EAGER_ATOL, - rtol=_EAGER_RTOL, - ) - else: - torch.testing.assert_close( - model.weight.grad, model_ref.weight.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL - ) + # Check both gradients and stateful accumulation for every parameter. + ref_params = dict(model_ref.named_parameters()) + for name, param in model.named_parameters(): + ref_param = ref_params[name] + if getattr(ref_param, "main_grad", None) is not None: + assert param.grad is None + torch.testing.assert_close( + param.main_grad, + ref_param.main_grad, + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + msg=f"main_grad mismatch: {name}", + ) + else: + torch.testing.assert_close( + param.grad, + ref_param.grad, + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + msg=f"grad mismatch: {name}", + ) torch.testing.assert_close(out.detach(), out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL) torch._dynamo.reset() @@ -2264,14 +2450,17 @@ def fn(inp): ], ids=recipe_id, ) -def test_te_linear_compile_is_first_microbatch( - fp8_recipe, compile_mode, deferred_backward, monkeypatch +@pytest.mark.parametrize("module", _LINEAR_MODULES) +def test_te_modules_compile_is_first_microbatch( + module, fp8_recipe, compile_mode, deferred_backward, monkeypatch ): """Reuse cached weights across microbatches and refresh after optimizer updates.""" dtype, device = torch.bfloat16, "cuda" - model = te.Linear(128, 128, params_dtype=dtype, device=device) - ref_model = te.Linear(128, 128, params_dtype=dtype, device=device) + model = getattr(te, module)(128, 128, params_dtype=dtype, device=device) + ref_model = getattr(te, module)(128, 128, params_dtype=dtype, device=device) ref_model.load_state_dict(model.state_dict()) + cache_names = ("fc1_weight", "fc2_weight") if module == "LayerNormMLP" else ("weight",) + cached_workspaces = {} def fn(inp, is_first): with te.autocast(recipe=fp8_recipe): @@ -2315,17 +2504,20 @@ def replay(graph): out = compiled(inp, is_first).clone() torch.testing.assert_close(out, out_ref, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) - workspace = model._fp8_workspaces["weight"] - if is_first: - cached_workspace = workspace - else: - assert workspace is cached_workspace - torch.testing.assert_close( - workspace.dequantize(), - ref_model._fp8_workspaces["weight"].dequantize(), - atol=_EAGER_ATOL, - rtol=_EAGER_RTOL, - ) + for name in cache_names: + workspace = model._fp8_workspaces[name] + if is_first: + if name in cached_workspaces: + assert workspace is not cached_workspaces[name] + cached_workspaces[name] = workspace + else: + assert workspace is cached_workspaces[name] + torch.testing.assert_close( + workspace.dequantize(), + ref_model._fp8_workspaces[name].dequantize(), + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + ) pending.append((out, out_ref, inp, inp_ref)) if not deferred_backward or microbatch == 2: @@ -2346,7 +2538,7 @@ def replay(graph): for param, ref_param in zip(model.parameters(), ref_model.parameters()): param.add_(param.grad, alpha=-0.001) ref_param.add_(ref_param.grad, alpha=-0.001) - del cached_workspace, workspace, out, out_ref, result, ref_result + del workspace, out, out_ref, result, ref_result torch._dynamo.reset() @@ -2403,74 +2595,363 @@ def train_step(): @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") -def test_te_linear_dynamic_shapes(): - """torch.compile of ``te.Linear`` with a ``mark_dynamic`` batch dimension: - one graph must serve all batch sizes -- no recompiles -- and match eager - numerically. - - Only the leading (batch/sequence) dims may be dynamic; the last dim is - fixed by the weight's ``in_features``. - """ - dtype = torch.bfloat16 - device = "cuda" - in_features, out_features = 64, 32 - model = te.Linear(in_features, out_features, params_dtype=dtype, device=device) +@pytest.mark.parametrize( + "module,bias_gelu_fusion", + [ + ("Linear", False), + ("LayerNormLinear", False), + ("LayerNormMLP", False), + ("LayerNormMLP", True), + ], +) +def test_te_modules_dynamic_shapes(module, bias_gelu_fusion): + """Batch sizes share the warmed-up outer graph and match eager outputs/gradients.""" + dtype, device = torch.bfloat16, "cuda" + model = _make_module(module, dtype=dtype, device=device) + if module == "LayerNormMLP": + model.bias_gelu_nvfusion = bias_gelu_fusion def fn(inp): return model(inp) torch._dynamo.reset() - compiled = torch.compile(fn, fullgraph=True) - - batch_sizes = [16, 32, 48] - - # Two warmup calls: the second absorbs the one-time recompile from module - # attributes lazily created during call one (e.g. the cached ``is_fsdp2``). + # Count the outer graph independently of lazy compilation of fused helpers. + compile_counter = CompileCounterWithBackend("inductor") + compiled = torch.compile(fn, fullgraph=True, backend=compile_counter) for _ in range(2): - warm = torch.randn(batch_sizes[0], in_features, dtype=dtype, device=device) + warm = torch.randn(16, 64, dtype=dtype, device=device) torch._dynamo.mark_dynamic(warm, 0) compiled(warm.requires_grad_(True)).sum().backward() model.zero_grad(set_to_none=True) - unique_graphs_baseline = _dynamo_counter("stats", "unique_graphs") - if not unique_graphs_baseline: - warnings.warn("unique_graphs counter is stale; skipping the recompile check") + frames_after_warmup = compile_counter.frame_count + assert frames_after_warmup > 0 - for batch in batch_sizes: - inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True) - # Mark batch dim as dynamic so Dynamo traces once and reuses across batch sizes. - torch._dynamo.mark_dynamic(inp, 0) - out = compiled(inp) - assert out.shape == (batch, out_features), f"wrong output shape for batch={batch}" - out.sum().backward() - assert inp.grad is not None, f"no input gradient for batch={batch}" - assert inp.grad.shape == inp.shape, f"wrong grad shape for batch={batch}" + for batch in (16, 32, 48): + base = torch.randn(batch, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base, dynamic_batch=True) - # Verify numerics against eager on each distinct batch size. - inp_eager = inp.detach().clone().requires_grad_(True) - model.zero_grad(set_to_none=True) - out_eager = model(inp_eager) - out_eager.sum().backward() - torch.testing.assert_close( - out.detach(), - out_eager.detach(), - atol=_EAGER_ATOL, - rtol=_EAGER_RTOL, - msg=f"forward mismatch at batch={batch}", - ) - torch.testing.assert_close( - inp.grad, - inp_eager.grad, - atol=_EAGER_ATOL, - rtol=_EAGER_RTOL, - msg=f"dgrad mismatch at batch={batch}", + assert compile_counter.frame_count == frames_after_warmup + + +# --------------------------------------------------------------------------- +# Module fake/real metadata and LayerNorm variants +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("module", ["LayerNormLinear", "LayerNormMLP"]) +@pytest.mark.parametrize("loss_source", ["output", "normalization", "both"]) +@pytest.mark.parametrize("fp8_recipe", [None, recipe.Float8CurrentScaling()]) +def test_te_layernorm_output_gradients(module, loss_source, fp8_recipe): + """Route distinct gradients from each public output to the correct backward field.""" + if fp8_recipe is not None and not fp8_available: + pytest.skip(reason_for_no_fp8) + model = _make_module(module, return_layernorm_output=True) + + def fn(inp): + if fp8_recipe is None: + output, norm_output = model(inp) + else: + with te.autocast(recipe=fp8_recipe): + output, norm_output = model(inp) + if loss_source == "output": + return output + if loss_source == "normalization": + return norm_output + return output, 0.25 * norm_output + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize( + "case", ["linear", "layernorm_linear", "mlp_bias_gelu", "mlp_gemm_gelu", "mlp_swiglu"] +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("bias", [False, True]) +@pytest.mark.parametrize("fp8_recipe", [None, recipe.Float8CurrentScaling()]) +def test_te_modules_fake_matches_real(monkeypatch, case, dtype, bias, fp8_recipe): + """Compare fake outputs, saved buffers and gradients against the eager kernels.""" + from transformer_engine.pytorch.dynamo.custom_op import _parse_arg_type, _spec_view + + if fp8_recipe is not None and not fp8_available: + pytest.skip(reason_for_no_fp8) + if fp8_recipe is not None and case == "mlp_gemm_gelu": + if dtype == torch.float32 or (dtype == torch.float16 and not bias): + pytest.skip( + "cuBLAS FP8 GELU requires matching output and auxiliary dtypes (BF16 without bias)" + ) + module_name = case if case in ("linear", "layernorm_linear") else "layernorm_mlp" + module = { + "linear": linear_module, + "layernorm_linear": layernorm_linear_module, + "layernorm_mlp": layernorm_mlp_module, + }[module_name] + checked = [] + + def assert_metadata(spec, real): + if isinstance(spec, TensorSpec): + assert real is not None + assert tuple(spec.shape) == tuple(real.shape) + assert spec.dtype == to_tensor_spec(real).dtype + assert spec.device == real.device + names = ( + tuple(real.__tensor_flatten__()[0]) + if hasattr(real, "__tensor_flatten__") + else ("data",) + ) + assert spec.inner_names() == names + for name, fake_inner in zip(names, spec.create_inner_tensors()): + real_inner = real if name == "data" else getattr(real, name) + assert fake_inner.shape == real_inner.shape + assert fake_inner.dtype == real_inner.dtype + assert fake_inner.stride() == real_inner.stride() + elif isinstance(spec, tuple): + assert isinstance(real, tuple) and len(spec) == len(real) + for fake_value, real_value in zip(spec, real): + assert_metadata(fake_value, real_value) + elif isinstance(spec, dict): + # No manual FSDP or checkpointing here; [] and None both denote no shards. + for key in spec.keys() | real.keys(): + if key in ("fsdp_shapes", "checkpoint"): + assert not spec.get(key) and not real.get(key) + else: + assert spec[key] == real[key] + else: + assert spec == real + + def checked_impl(direction): + real_impl = getattr(module, f"_{module_name}_{direction}_impl") + fake_impl = getattr(module, f"_{module_name}_{direction}_fake") + + def wrapper(args): + spec_args = _spec_view(args, _parse_arg_type(type(args)).tensor_field_names()) + # Fake execution may change quantizer usage; leave the real invocation untouched. + spec_args = dataclasses.replace( + spec_args, + **{ + field.name: value.copy() + for field in dataclasses.fields(spec_args) + if isinstance(value := getattr(spec_args, field.name), Quantizer) + }, + ) + expected = fake_impl(spec_args) + actual = real_impl(args) + assert_metadata(expected, actual) + checked.append(direction) + return actual + + return wrapper + + for direction in ("forward", "backward"): + monkeypatch.setattr(module, f"_{module_name}_{direction}_impl", checked_impl(direction)) + kwargs = dict(params_dtype=dtype, device="cuda", bias=bias) + if case != "linear": + kwargs["return_layernorm_output"] = True + if case == "linear": + model = te.Linear(64, 32, **kwargs) + elif case == "layernorm_linear": + model = te.LayerNormLinear(64, 32, **kwargs) + else: + model = te.LayerNormMLP( + 64, 128, activation="swiglu" if case == "mlp_swiglu" else "gelu", **kwargs ) + model.bias_gelu_nvfusion = case == "mlp_bias_gelu" + model.gemm_gelu_fusion = case == "mlp_gemm_gelu" + inp = torch.randn(2, 16, 64, dtype=dtype, device="cuda", requires_grad=True) + with te.autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + outs = _flatten_outputs(model(inp)) + sum(out.sum() for out in outs).backward() + assert checked == ["forward", "backward"] - if unique_graphs_baseline: - unique_graphs_after = _dynamo_counter("stats", "unique_graphs") - assert unique_graphs_after == unique_graphs_baseline, ( - "Unexpected recompilation(s) across different batch sizes: " - f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize( + "dtype,bias", [(torch.bfloat16, False), (torch.bfloat16, True), (torch.float16, True)] +) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_layernorm_mlp_compile_fp8_gemm_gelu(dtype, bias, compile_mode): + """Fused FP8 GEMM-GELU must preserve the auxiliary tensor dtype through backward.""" + model = te.LayerNormMLP(64, 128, params_dtype=dtype, device="cuda", bias=bias) + model.gemm_gelu_fusion = True + _run_module_compile_test(model, recipe.Float8CurrentScaling(), compile_mode, 64, dtype=dtype) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize( + "quantization,disable_rht_fusion", + [("mxfp8", False), ("mxfp8_2d", False), ("nvfp4", False), ("nvfp4", True)], +) +@pytest.mark.parametrize("module", ["LayerNormLinear", "LayerNormMLP"]) +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize( + "use_cudnn,input_shape,return_layernorm_output", + [ + (False, (128, 128), False), + (True, (128, 128), False), + (True, (64, 128), False), + (True, (2, 64, 128), False), + (True, (128, 128), True), + ], +) +def test_te_layernorm_compile_scale_layout( + monkeypatch, + module, + normalization, + use_cudnn, + input_shape, + return_layernorm_output, + quantization, + disable_rht_fusion, +): + """Saved normalization scales retain their layout through compiled backward.""" + if quantization.startswith("mxfp8") and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if quantization == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + monkeypatch.setenv("NVTE_NORM_FWD_USE_CUDNN", str(int(use_cudnn))) + monkeypatch.setenv("NVTE_NVFP4_DISABLE_RHT_CAST_FUSION", str(int(disable_rht_fusion))) + model = getattr(te, module)( + 128, + 128, + params_dtype=torch.bfloat16, + device="cuda", + normalization=normalization, + return_layernorm_output=return_layernorm_output, + ) + if quantization == "mxfp8_2d": + + def qfactory(role): + return MXFP8Quantizer( + tex.DType.kFloat8E4M3, + with_2d_quantization=role is not None and role.tensor_type == "input", + ) + + fp8_recipe = recipe.CustomRecipe(qfactory=qfactory) + elif quantization == "mxfp8": + fp8_recipe = recipe.MXFP8BlockScaling() + else: + fp8_recipe = recipe.NVFP4BlockScaling(disable_stochastic_rounding=True) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + base = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda") + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("sequence_parallel", [False, True]) +@pytest.mark.parametrize("gathered", [False, True]) +@pytest.mark.parametrize("input_shape", [(128, 128), (2, 64, 128)]) +def test_te_layernorm_mlp_fake_gathered_output( + monkeypatch, sequence_parallel, gathered, input_shape +): + """The normalization gather is independent of set_parallel_mode.""" + from transformer_engine.pytorch.dynamo.custom_op import _parse_arg_type, _spec_view + + module = layernorm_mlp_module + real_impl = module._layernorm_mlp_forward_impl + checked = [] + + def checked_impl(args): + args = dataclasses.replace( + args, + sequence_parallel=sequence_parallel, + tp_size=2, + return_layernorm_output_gathered=gathered, ) + spec_args = _spec_view(args, _parse_arg_type(type(args)).tensor_field_names()) + expected = module._layernorm_mlp_forward_fake(spec_args) + actual = real_impl(args) + for spec, tensor in zip(expected[:2], actual[:2]): + assert tuple(spec.shape) == tuple(tensor.shape) + checked.append(True) + return actual + + # Emulate a two-rank gather while exercising the real normalization and GEMMs. + monkeypatch.setattr(module, "get_distributed_world_size", lambda group: 2) + monkeypatch.setattr( + module, + "gather_along_first_dim", + lambda tensor, group, quantizer=None: (torch.cat([tensor, tensor]), None), + ) + monkeypatch.setattr(module, "_layernorm_mlp_forward_impl", checked_impl) + model = te.LayerNormMLP( + 128, + 128, + params_dtype=torch.bfloat16, + device="cuda", + set_parallel_mode=False, + return_layernorm_output=True, + ) + with torch.no_grad(): + model(torch.randn(input_shape, dtype=torch.bfloat16, device="cuda")) + assert checked == [True] + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize("return_layernorm_output", [False, True]) +@pytest.mark.parametrize("zero_centered_gamma", [False, True]) +def test_te_layernorm_linear_compile_variants( + normalization, return_layernorm_output, zero_centered_gamma +): + """Norm-type / returned-norm-output / zero-centered-gamma variants of + ``te.LayerNormLinear`` under torch.compile, bf16 and FP8 current scaling.""" + model = te.LayerNormLinear( + 64, + 32, + params_dtype=torch.bfloat16, + device="cuda", + normalization=normalization, + return_layernorm_output=return_layernorm_output, + zero_centered_gamma=zero_centered_gamma, + ) + _run_module_compile_test(model, None, "default", 64) + if fp8_available: + _run_module_compile_test(model, recipe.Float8CurrentScaling(), "default", 64) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("activation", ["gelu", "swiglu", "relu", "qgeglu"]) +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize("return_layernorm_output", [False, True]) +def test_te_layernorm_mlp_compile_variants(activation, normalization, return_layernorm_output): + """Activation / norm-type / returned-norm-output variants of + ``te.LayerNormMLP`` under torch.compile, bf16 and FP8 current scaling.""" + model = te.LayerNormMLP( + 64, + 128, + params_dtype=torch.bfloat16, + device="cuda", + activation=activation, + normalization=normalization, + return_layernorm_output=return_layernorm_output, + ) + _run_module_compile_test(model, None, "default", 64) + if fp8_available: + _run_module_compile_test(model, recipe.Float8CurrentScaling(), "default", 64) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("bias", [True, False]) +def test_te_layernorm_mlp_compile_no_bias_or_frozen(bias): + """``te.LayerNormMLP`` without biases, and with frozen FC1 weight (exercises + the ``None`` saved-tensor / grad slots).""" + model = te.LayerNormMLP(64, 128, params_dtype=torch.bfloat16, device="cuda", bias=bias) + _run_module_compile_test(model, None, "default", 64) + model.fc1_weight.requires_grad_(False) + _run_module_compile_test(model, None, "default", 64) + if fp8_available: + _run_module_compile_test(model, recipe.Float8CurrentScaling(), "default", 64) # --------------------------------------------------------------------------- # @@ -2912,21 +3393,20 @@ def test_te_ops_bias_compile_with_fp8_input(backward): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( - "dtype,quantization,shape,disable_rht_fusion,case,grads", + "dtype,quantization,shape,disable_rht_fusion", [ - (dtype, quantization, (2, 16, 32), False, case, grads) - for dtype in (torch.float16, torch.bfloat16) - for quantization in (None, "fp8") - for case in ("linear", "bias", "pair", "unfused") - for grads in ("all", "input", "weight", "bias", "none") - ] - + [ - (torch.bfloat16, "nvfp4", (2, 16, 128), False, "pair", "all"), - (torch.bfloat16, "nvfp4", (64, 128), False, "pair", "all"), - (torch.bfloat16, "nvfp4", (64, 128), True, "pair", "all"), - (torch.bfloat16, "nvfp4-no-rht", (2, 16, 128), False, "pair", "all"), + (torch.float16, None, (2, 16, 32), False), + (torch.bfloat16, None, (2, 16, 32), False), + (torch.float16, "fp8", (2, 16, 32), False), + (torch.bfloat16, "fp8", (2, 16, 32), False), + (torch.bfloat16, "nvfp4", (2, 16, 128), False), + (torch.bfloat16, "nvfp4", (64, 128), False), + (torch.bfloat16, "nvfp4", (64, 128), True), + (torch.bfloat16, "nvfp4-no-rht", (2, 16, 128), False), ], ) +@pytest.mark.parametrize("case", ["linear", "bias", "pair", "unfused"]) +@pytest.mark.parametrize("grads", ["all", "input", "weight", "bias", "none"]) def test_te_ops_linear_bias_compile( dtype, quantization, shape, disable_rht_fusion, case, grads, monkeypatch ): @@ -3058,13 +3538,29 @@ def unexpected_allocation(*args, **kwargs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("training", [False, True]) +@pytest.mark.parametrize("initial_device", ["cuda", "cpu", "meta"]) @pytest.mark.parametrize( - "initial_device,quantization", - [("cuda", None), ("cpu", None), ("meta", None), ("cuda", "nvfp4"), ("cuda", "nvfp4-no-rht")], + "quantization", + [ + None, + "fp8", + "fp8-block", + "mxfp8", + "nvfp4", + "nvfp4-no-rht", + "nvfp4-4over6", + "nvfp4-row-scaled", + ], ) def test_te_ops_linear_bias_cold_cudagraphs(training, initial_device, quantization): """A fresh process prevents earlier GEMMs from hiding lazy workspace allocation.""" - if quantization is not None and not nvfp4_available: + if quantization == "fp8" and not fp8_available: + pytest.skip(reason_for_no_fp8) + if quantization == "fp8-block" and not fp8_block_scaling_available: + pytest.skip(reason_for_no_fp8_block_scaling) + if quantization == "mxfp8" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if quantization is not None and quantization.startswith("nvfp4") and not nvfp4_available: pytest.skip(reason_for_no_nvfp4) script = textwrap.dedent( """ @@ -3079,17 +3575,27 @@ def test_te_ops_linear_bias_cold_cudagraphs(training, initial_device, quantizati training, initial_device, quantization = sys.argv[1:] training = training == "True" - is_nvfp4 = quantization != "None" - quant_recipe = ( - recipe.NVFP4BlockScaling( - disable_rht=quantization == "nvfp4-no-rht", + is_nvfp4 = quantization.startswith("nvfp4") + if quantization == "None": + quant_recipe = None + elif is_nvfp4: + row_scaled = quantization == "nvfp4-row-scaled" + quant_recipe = recipe.NVFP4BlockScaling( + disable_rht=quantization != "nvfp4", disable_stochastic_rounding=True, + disable_2d_quantization=row_scaled, + row_scaled_activation=row_scaled, + nvfp4_4over6="all" if quantization == "nvfp4-4over6" else "none", + backward_override="dequantized" if row_scaled else None, ) - if is_nvfp4 - else None - ) - in_features, out_features = (128, 128) if is_nvfp4 else (32, 64) - batch_size = 32 + else: + quant_recipe = { + "fp8": recipe.Float8CurrentScaling, + "fp8-block": recipe.Float8BlockScaling, + "mxfp8": recipe.MXFP8BlockScaling, + }[quantization]() + in_features, out_features = (128, 128) if quant_recipe is not None else (32, 64) + batch_size = 128 if isinstance(quant_recipe, recipe.Float8BlockScaling) else 32 eager_tols = {"rtol": 0, "atol": 0} if is_nvfp4 else {} assert get_cublas_workspace.cache_info().currsize == 0 linear = te.ops.BasicLinear( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 2345a1808a6..4ceade26060 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -11,7 +11,7 @@ import functools import inspect import os -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union import warnings import logging from packaging.version import Version as PkgVersion @@ -1662,6 +1662,10 @@ class FusedAttnBwdArgs: dO_quantizer: Optional[Any] = None dP_quantizer: Optional[Any] = None + def setup_grad_outputs(self, grads: Sequence[Any]) -> None: + """Assign the output gradient, ignoring auxiliary outputs.""" + self.grad_output = grads[0] + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: """Pull saved tensors from ``ctx`` into the fields backward consumes.""" ( diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 88dfb808420..324231fc5f2 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -340,6 +340,9 @@ py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, * LayerNorm **************************************************************************************************/ +std::optional mxfp8_norm_fused_alignment(bool rmsnorm, bool with_2d_quantization, + bool use_cudnn); + std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, const at::Tensor &mu, const at::Tensor &rsigma, const at::Tensor &gamma, const int sm_margin, diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index e93c01026a9..cc990576081 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -10,6 +10,14 @@ namespace transformer_engine::pytorch { +std::optional mxfp8_norm_fused_alignment(bool rmsnorm, bool with_2d_quantization, + bool use_cudnn) { + if (rmsnorm ? with_2d_quantization : !use_cudnn) { + return std::nullopt; + } + return 128; +} + std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, const at::Tensor &mu, const at::Tensor &rsigma, const at::Tensor &gamma, const int sm_margin, @@ -106,9 +114,11 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe if (quantizer.is_none() || IsFloat8Quantizers(quantizer.ptr())) { impl = Impl::FULLY_FUSED; } else if (IsMXFP8Quantizers(quantizer.ptr())) { - if (transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN") && outer_size % 128 == 0 && - inner_size % 128 == 0) { - // cuDNN MXFP8 kernel requires full 128x128 tiles + const auto mxfp8_quantizer_cpp = static_cast(quantizer_cpp.get()); + const auto alignment = + mxfp8_norm_fused_alignment(false, mxfp8_quantizer_cpp->with_2d_quantization, + transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); + if (alignment && outer_size % *alignment == 0 && inner_size % *alignment == 0) { impl = Impl::FULLY_FUSED; } } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr()) && @@ -350,9 +360,10 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w } else if (IsMXFP8Quantizers(quantizer.ptr())) { auto mxfp8_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(mxfp8_quantizer_cpp != nullptr, "Could not cast to MXFP8 quantizer"); - // The fused MXFP8 kernels quantize 1D blocks and require full 128x128 tiles - if (!mxfp8_quantizer_cpp->with_2d_quantization && outer_size % 128 == 0 && - inner_size % 128 == 0) { + const auto alignment = + mxfp8_norm_fused_alignment(true, mxfp8_quantizer_cpp->with_2d_quantization, + transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); + if (alignment && outer_size % *alignment == 0 && inner_size % *alignment == 0) { if (transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { impl = Impl::FUSED_NORM_QUANT_UNSWIZZLED; } else { diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index ebcdd17b07c..a5a5a025c75 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -402,6 +402,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); // Other granular functions + m.def("mxfp8_norm_fused_alignment", &transformer_engine::pytorch::mxfp8_norm_fused_alignment, + py::arg("rmsnorm"), py::arg("with_2d_quantization"), py::arg("use_cudnn")); m.def("layernorm_fwd", &transformer_engine::pytorch::layernorm_fwd, "LayerNorm", py::arg("input"), py::arg("weight"), py::arg("bias"), py::arg("eps"), py::arg("ln_out"), py::arg("quantizer"), py::arg("otype"), py::arg("sm_margin"), py::arg("zero_centered_gamma")); diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 1ed84017ecd..5a784abf332 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -18,7 +18,7 @@ keeps its saved-tensor and context-metadata contract. A ``torch.library`` custom op is narrower: it only accepts flat schema slots -(tensors plus opaque objects) and returns a flat ``Tensor[]``. +(tensors, symbolic shapes and opaque objects) and returns a flat ``Tensor[]``. Bridging the two takes three parts (below): a parsed per-op *arg plan* maps the args dataclass onto the op's input slots; a per-trace *output plan*, parsed from @@ -37,7 +37,8 @@ * ``TENSOR_OR_QUANTIZED`` -- a field that may be a plain tensor, a bare quantized storage, or ``None``: three slots (the tensor, its flat inner buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. - * ``SIMPLE`` -- every remaining simple value (scalars, enums, sizes, + * ``SHAPE`` -- a ``torch.Size`` / ``Optional[torch.Size]``: one ``SymInt[]`` slot. + * ``SIMPLE`` -- every remaining simple value (scalars, enums, quantizers -- value-opaque constants baked into the graph -- and nested collections of them), gathered into one shared ``OpaqueValueBundle`` slot. * ``PROCESS_GROUP`` -- rides in the shared bundle too, as its c10d registry @@ -66,8 +67,8 @@ forward state + ``ctx_attrs`` (e.g. saved-tensor aliases) and return the tensors to persist; the plan's output ranges are stashed on ``ctx``; * on ``backward()`` the incoming flat grads are sliced per user output from the - stashed plan (a ``grad_outputs`` field on the backward args receives the - whole tuple; otherwise ``grad_output`` receives the first output's grad), + stashed plan and passed in forward-output order to the backward + container's ``setup_grad_outputs`` method, the container's optional ``setup_saved_tensors`` hook restores the saved tensors, then the *backward op* runs the real ``bwd_impl`` and returns the flat grads (``bwd_fake_impl`` is its data-free fake). @@ -397,6 +398,7 @@ class _FieldKind(Enum): TENSOR = "tensor" # one ``Tensor`` / ``Tensor?`` slot TENSOR_OR_QUANTIZED = "tensor_or_quantized" # 3 slots: tensor / inner / meta + SHAPE = "shape" # one ``SymInt[]`` / ``SymInt[]?`` slot PROCESS_GROUP = "process_group" # c10d group name inside the shared bundle SIMPLE = "simple" # value carried verbatim inside the shared bundle UNSUPPORTED = "unsupported" # no slots; only a trivial value may cross @@ -475,6 +477,9 @@ def _is_simple_annot(annot: Any) -> bool: if get_origin(annot) in (tuple, list): inner = [a for a in get_args(annot) if a is not Ellipsis] return bool(inner) and all(_is_simple_annot(a) for a in inner) + if get_origin(annot) is dict: + args = get_args(annot) + return len(args) == 2 and args[0] is str and (args[1] is Any or _is_simple_annot(args[1])) return False @@ -491,6 +496,9 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan: if stripped is torch.Tensor: slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor") return _FieldPlan(name, _FieldKind.TENSOR, (slot,)) + if stripped is torch.Size: + slot = _SlotSpec(name, "SymInt[]?" if is_optional else "SymInt[]") + return _FieldPlan(name, _FieldKind.SHAPE, (slot,)) # A union mixing tensor types with anything else is a malformed signature # (e.g. a bare quantized-storage Optional, or Tensor | int): reject it at # registration instead of silently degrading to an unsupported field. @@ -631,7 +639,7 @@ def pack(self, obj: Any) -> Dict[str, Any]: for field in self.fields: value = getattr(obj, field.name, None) match field.kind: - case _FieldKind.TENSOR: + case _FieldKind.TENSOR | _FieldKind.SHAPE: slots[field.slots[0].name] = value case _FieldKind.TENSOR_OR_QUANTIZED: _pack_tensor_or_quantized(field, value, slots) @@ -667,6 +675,9 @@ def unpack(self, slots: Dict[str, Any]) -> Any: match field.kind: case _FieldKind.TENSOR: kwargs[field.name] = slots[field.slots[0].name] + case _FieldKind.SHAPE: + value = slots[field.slots[0].name] + kwargs[field.name] = None if value is None else torch.Size(value) case _FieldKind.TENSOR_OR_QUANTIZED: kwargs[field.name] = _unpack_tensor_or_quantized(field, slots) case _FieldKind.PROCESS_GROUP: @@ -1029,12 +1040,11 @@ def _register_autograd_for_op( the saved tuple + ``ctx_attrs`` to the module's ``setup_context`` and stashes the plan on ``ctx`` so backward can slice its grads per user output. """ - bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields) + + tensor_list_positions = tuple(offset + 1 for offset in fwd_plan.tensor_or_quantized_offsets()) def _setup_context(ctx, inputs, output): - ctx.fwd_tensor_list_lengths = { - i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) - } + ctx.fwd_tensor_list_lengths = {i: len(inputs[i]) for i in tensor_list_positions} fwd_obj = fwd_plan.unpack(dict(zip(fwd_plan.slot_names, inputs))) spec_obj = _spec_view(fwd_obj, fwd_plan.tensor_field_names()) @@ -1069,10 +1079,7 @@ def _autograd_backward(ctx, *grad_outputs): ctx.tensor_objects = None user_grads = _slice_user_grads(ctx.output_ranges, grad_outputs[0]) ctx.output_ranges = None - if bwd_takes_grad_tuple: - bwd_obj.grad_outputs = tuple(user_grads) - else: - bwd_obj.grad_output = user_grads[0] + bwd_obj.setup_grad_outputs(user_grads) kwargs = bwd_plan.pack(bwd_obj) bwd_args_flat = [kwargs[name] for name in bwd_plan.slot_names] grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] @@ -1359,6 +1366,9 @@ def register_custom_op_with_autograd( non-differentiable input). * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``bwd_impl`` returning :class:`TensorSpec` grads. + * ``bwd_arg_type.setup_grad_outputs(grads)`` -- assign incoming gradients + to the container's fields. Receives one entry per forward user output, + in return order, including auxiliary outputs such as weight workspaces. * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook; skipped if absent. @@ -1367,10 +1377,8 @@ def register_custom_op_with_autograd( forward state and returns the tensors to persist; the framework saves them via ``ctx.save_for_backward``. Before ``bwd_impl`` runs, the framework restores them into the container's tensor fields through the - ``setup_saved_tensors`` hook and sets the incoming gradient directly -- - into a ``grad_outputs`` field (tuple, one grad per user output) if - ``bwd_arg_type`` declares one, else into ``grad_output`` (the first user - output's grad) -- so ``bwd_impl`` receives a fully-populated + ``setup_saved_tensors`` hook and passes the incoming gradients to + ``setup_grad_outputs``, so ``bwd_impl`` receives a fully-populated ``bwd_arg_type``. Registration touches experimental ``torch.library`` / opaque-object APIs diff --git a/transformer_engine/pytorch/dynamo/fake.py b/transformer_engine/pytorch/dynamo/fake.py index 5f856d62214..9b70526fadc 100644 --- a/transformer_engine/pytorch/dynamo/fake.py +++ b/transformer_engine/pytorch/dynamo/fake.py @@ -14,6 +14,7 @@ from .tensor_spec import TensorSpec from ..quantized_tensor import Quantizer +from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..utils import get_device_compute_capability @@ -92,3 +93,58 @@ def fake_bgrad( ) ) return db, dx + + +def fake_activation( + input_: TensorSpec, + quantizer: Optional[Quantizer] = None, + *, + gated: bool = False, + dtype: Optional[torch.dtype] = None, + fused_quantization: bool = True, +) -> TensorSpec: + """Describe an activation output, including optional fused quantization.""" + shape = (*input_.shape[:-1], input_.shape[-1] // 2 if gated else input_.shape[-1]) + dtype = dtype or input_.dtype + if fused_quantization: + return _make_fused_quantization_output_spec(shape, dtype, input_.device, quantizer) + return TensorSpec(shape=shape, dtype=dtype, device=input_.device, quantizer=quantizer) + + +@torch.compiler.assume_constant_result +def _mxfp8_norm_fused_alignment( + normalization: str, with_2d_quantization: bool, use_cudnn: bool +) -> Optional[int]: + return tex.mxfp8_norm_fused_alignment( + normalization == "RMSNorm", with_2d_quantization, use_cudnn + ) + + +def fake_norm( + input_: TensorSpec, + quantizer: Optional[Quantizer] = None, + *, + normalization: str = "LayerNorm", + dtype: Optional[torch.dtype] = None, +) -> tuple[TensorSpec, Optional[TensorSpec], TensorSpec]: + """Describe newly allocated norm output, mean and reciprocal standard deviation.""" + if normalization not in ("LayerNorm", "RMSNorm"): + raise ValueError(f"Unsupported normalization: {normalization}") + out = _make_fused_quantization_output_spec( + input_.shape, dtype or input_.dtype, input_.device, quantizer + ) + rows, cols = math.prod(input_.shape[:-1]), input_.shape[-1] + if isinstance(quantizer, MXFP8Quantizer): + use_cudnn = bool(int(os.getenv("NVTE_NORM_FWD_USE_CUDNN", "0") or "0")) + alignment = _mxfp8_norm_fused_alignment( + normalization, quantizer.with_2d_quantization, use_cudnn + ) + if use_cudnn and alignment is not None and rows % alignment == 0 and cols % alignment == 0: + out.with_gemm_swizzled_scales = False + mean = ( + TensorSpec(shape=(rows,), dtype=torch.float32, device=input_.device) + if normalization == "LayerNorm" + else None + ) + rsigma = TensorSpec(shape=(rows,), dtype=torch.float32, device=input_.device) + return out, mean, rsigma diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 3d04882698a..fd1e9b067b9 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -6,19 +6,33 @@ import dataclasses import queue -from typing import Any, Callable, List, Optional, Tuple, Union +from typing import Any, Callable, List, Optional, Sequence, Tuple, Union import torch from .. import cpp_extensions as tex from ..constants import TE_DType from ..distributed import in_fp8_activation_recompute_phase +from ..dynamo import TensorSpec, is_value_opaque_quantizer from ..export import is_in_onnx_export_mode from ..quantization import FP8GlobalStateManager +from ..quantized_tensor import Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method +def compile_unsupported_quantizer_reason( + quantizers: Sequence[Optional[Quantizer]], +) -> Optional[str]: + """Return a fallback reason if a quantizer cannot cross the custom-op boundary.""" + for quantizer in quantizers: + # Delayed-scaling and unregistered custom-recipe quantizers are not value-opaque. + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + return "a quantizer not registered as a torch.compile value-opaque type" + return None + + def sum_bias_grad(tensor: torch.Tensor) -> torch.Tensor: """Reduce all token dimensions, preserving the final feature dimension.""" if tensor.ndim == 1: @@ -343,3 +357,70 @@ def check_fp8_reduce_and_update(restore_first_module: bool = False) -> bool: if restore_first_module or in_fp8_activation_recompute_phase(): qstate.is_first_fp8_module = first_fp8_module return result + + +def get_output_first_dim_size(input_first_dim_size: int, args: Any) -> int: + """Compute the output's first dimension size from the input's. + + Sequence parallelism gathers this dimension in column-parallel mode and + scatters it in row-parallel mode. ``args`` provides ``sequence_parallel``, + ``parallel_mode``, and ``tp_size``. + """ + if not args.sequence_parallel: + return input_first_dim_size + if args.parallel_mode == "column": + return input_first_dim_size * args.tp_size + if args.parallel_mode == "row": + return input_first_dim_size // args.tp_size + return input_first_dim_size + + +def fake_workspace_valid(workspace: TensorSpec, quantizer: Optional[Quantizer]) -> bool: + """Spec-level mirror of ``_is_weight_workspace_valid``.""" + if quantizer is None: + return True + if ( + isinstance(workspace.quantizer, Float8BlockQuantizer) + and workspace.quantizer.block_scaling_dim == 2 + and workspace.quantizer.rowwise_usage + ): + # 2D block scaling can reconstruct columnwise buffers from rowwise data. + return True + required = TensorSpec( + shape=workspace.shape, dtype=workspace.dtype, quantizer=quantizer, device=workspace.device + ).inner_names() + return set(required) <= set(workspace.inner_names()) + + +def fake_quantized_weight( + weight: TensorSpec, + quantizer: Optional[Quantizer], + workspace: Optional[TensorSpec], + *, + activation_dtype: torch.dtype, + cache_weight: bool, +) -> Tuple[TensorSpec, Optional[TensorSpec], Optional[str]]: + """Spec-level mirror of ``quantize_weight``: ``(weightmat, new_workspace, + save_alias)`` where ``save_alias`` names the forward input / output the + saved weight aliases (``None`` when it is a fresh tensor).""" + if weight.is_quantized: + return weight, None, "weight" + if workspace is not None and not fake_workspace_valid(workspace, quantizer): + workspace = None + if workspace is not None: + weightmat = dataclasses.replace(workspace) + weightmat.update_usage(rowwise_usage=True) + return weightmat, None, "weight_workspace" + weightmat = TensorSpec( + shape=tuple(weight.shape), + dtype=activation_dtype, + quantizer=quantizer, + device=weight.device, + ) + new_workspace = None + if cache_weight: + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_workspace = weightmat + weightmat.update_usage(rowwise_usage=True) + return weightmat, new_workspace, "new_weight_workspace" if cache_weight else None diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index b4a91b2a3ea..a61afe9f154 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -6,8 +6,8 @@ import os import warnings import weakref -from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, Tuple, Union, List +from dataclasses import dataclass, replace as dataclass_replace +from typing import Sequence, Any, Callable, Dict, Optional, Tuple, Union, List import torch from torch.nn import init @@ -44,6 +44,9 @@ nvtx_range_push, needs_quantized_gemm, get_nvtx_range_context, + warn_compile_eager_fallback, + warn_if_compile_disabled, + check_gemm_dims, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -65,9 +68,12 @@ from ..jit import no_torch_dynamo from ..graph import is_graph_capturing from ._common import ( + compile_unsupported_quantizer_reason, apply_normalization, check_fp8_reduce_and_update, + fake_workspace_valid, noop_cat, + get_output_first_dim_size, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -79,9 +85,15 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorOrQuantized +from ..dynamo.fake import fake_norm +from ..dynamo import ( + TensorSpec, + TensorOrQuantized, + register_custom_op_with_autograd, +) from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.hybrid_tensor import HybridQuantizer from ..tensor.identity_tensor import IdentityQuantizer from ..cpu_offload import ( @@ -279,6 +291,10 @@ class LayerNormLinearBwdArgs: # --- Per-backward scratch state (populated inside the backward impl) --- ub_obj_gradout: Optional[Any] = None + def setup_grad_outputs(self, grads: Sequence[Any]) -> None: + """Unpack gradients in forward-output order, ignoring weight workspaces.""" + self.grad_output, self.grad_ln_out, _ = grads + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: """Pull saved tensors from ``ctx`` into the fields backward consumes.""" ( @@ -1466,6 +1482,353 @@ def wgrad_gemm( ) +def _layernorm_linear_forward_fake( + args: LayerNormLinearFwdArgs, +) -> Tuple[ + TensorSpec, + Optional[TensorSpec], + Optional[TensorSpec], + Optional[Tuple[Any, ...]], + Optional[Dict], +]: + """Shape/metadata-only twin of :func:`_layernorm_linear_forward_impl` for + torch.compile, returning ``TensorSpec`` descriptors for the outputs and + saved tensors instead of allocating real data.""" + if args.fsdp_group is not None and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time LayerNormLinear forward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight_quantizer = args.weight_quantizer + fp8_or_debug = args.fp8 or args.debug + with_input_all_gather = args.parallel_mode == "column" and args.sequence_parallel + backward_needs_input = args.is_grad_enabled and args.weight_requires_grad + device = args.inp.device + + out_features, in_features = args.weight.shape + inp_leading = args.inp.shape[0] if len(args.inp.shape) > 1 else 1 + inputmat_aliases_inp = args.inp.dtype == args.activation_dtype + ln_weight_aliases = args.ln_weight.dtype == args.activation_dtype + + # Norm output quantizer usage -- mirrors the impl; the ln_out spec must be + # taken now, before the later rowwise-only usage for the all-gather. + if args.fp8: + if args.input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + args.input_quantizer.set_usage( + rowwise=True, + columnwise=backward_needs_input and args.backward_override is None, + ) + if with_input_all_gather and args.input_quantizer.supports_only_rowwise_all_gather(): + args.input_quantizer.set_usage(columnwise=False) + custom = is_custom(args.input_quantizer) + hybrid = isinstance(args.input_quantizer, HybridQuantizer) + identity = isinstance(args.input_quantizer, IdentityQuantizer) + with_quantized_norm = ( + args.fp8 + and not args.debug + and not args.return_layernorm_output + and not args.return_layernorm_output_gathered + and args.backward_override is None + and not custom + and not hybrid + and not identity + ) + # A custom quantizer's ln_out stays in high precision on the plain + # all-gather path (only the gathered copy is quantized). + ln_out_quantized = fp8_or_debug and not ( + with_input_all_gather and not args.return_layernorm_output_gathered and custom + ) + ln_out_spec, mu, rsigma = fake_norm( + args.inp, + args.input_quantizer if with_quantized_norm else None, + normalization=args.normalization, + dtype=args.activation_dtype, + ) + if ln_out_quantized and not with_quantized_norm: + ln_out_spec = TensorSpec( + shape=ln_out_spec.shape, + dtype=args.activation_dtype, + quantizer=args.input_quantizer, + device=device, + ) + if with_input_all_gather and fp8_or_debug: + args.input_quantizer.set_usage(rowwise=True, columnwise=False) + # ``ln_out_return`` is the high-precision norm output (or its all-gathered copy). + ln_out_return_is_total = with_input_all_gather and args.return_layernorm_output_gathered + # Whether the quantized ln_out is a different object from the returned one. + ln_out_rebound = ln_out_quantized and not with_quantized_norm + + # ------------------------------------------------------ + # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed``. + # ------------------------------------------------------ + new_weight_workspace = None + workspace = None # args.weight_workspace after validation + weightmat = None + weightmat_is_storage = False + weightmat_aliases_weight = False + is_weight_param_quantized = False + if fp8_or_debug: + is_weight_param_quantized = args.weight.is_quantized + if is_weight_param_quantized and not args.debug: + weight_quantizer = args.weight.quantizer + elif weight_quantizer is not None: + weight_quantizer.set_usage( + rowwise=True, + columnwise=args.is_grad_enabled + and not args.is_fsdp2 + and args.backward_override is None, + ) + + if args.weight.is_quantized: + # Primary-quantized args.weight: the impl reuses it as ``weightmat``. + weightmat = args.weight + weightmat_is_storage = True + weightmat_aliases_weight = True + else: + weightmat_is_storage = True + workspace = args.weight_workspace + if workspace is not None and not fake_workspace_valid(workspace, weight_quantizer): + # quantize_weight drops a stale workspace and builds a new one. + workspace = None + if workspace is not None: + # Copy, so the ``update_usage`` below stays off the input spec. + weightmat = dataclass_replace(workspace) + else: + weightmat = TensorSpec( + shape=tuple(args.weight.shape), + dtype=args.activation_dtype, + quantizer=weight_quantizer, + device=args.weight.device, + ) + if args.cache_weight: + # Persistent cache entries are wrappers, not bare storages. + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_weight_workspace = weightmat + weightmat.update_usage(rowwise_usage=True) + else: + weightmat_aliases_weight = args.weight.dtype == args.activation_dtype + weightmat = TensorSpec( + shape=tuple(args.weight.shape), dtype=args.activation_dtype, device=args.weight.device + ) + + # Bias cast: cuBLAS has no FP8 GEMM with FP32 args.bias. + bias_dtype = args.activation_dtype + if fp8_or_debug and args.activation_dtype == torch.float32: + bias_dtype = torch.bfloat16 + bias_aliases = args.bias is not None and args.bias.dtype == bias_dtype + + if args.output_quantizer is not None: + args.output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Outputs: y = norm(x) @ w^T and the optional norm output. + # ------------------------------------------------------ + requires_grad = args.is_grad_enabled and args.any_requires_grad() + out = TensorSpec( + shape=( + (out_features,) + if len(args.inp.shape) == 1 + else ( + get_output_first_dim_size(inp_leading, args), + *tuple(args.inp.shape[1:-1]), + out_features, + ) + ), + dtype=args.activation_dtype, + quantizer=args.output_quantizer, + requires_grad=requires_grad, + device=device, + ) + ln_out_for_return = None + if args.return_layernorm_output: + ln_leading = inp_leading + if args.return_layernorm_output_gathered and with_input_all_gather: + ln_leading = inp_leading * args.tp_size + ln_out_for_return = TensorSpec( + shape=( + (ln_leading, *tuple(args.inp.shape[1:])) + if len(args.inp.shape) > 1 + else (in_features,) + ), + dtype=args.activation_dtype, + requires_grad=requires_grad, + device=device, + ) + + # ------------------------------------------------------ + # Backward state -- saved-tensor layout + # (inputmat, wt_save, saved_weight, args.bias, args.ln_weight, ln_out, mu, rsigma). + # ------------------------------------------------------ + tensors_to_save_from_forward = None + ctx_attrs = None + if args.is_grad_enabled: + ln_out_needs_gather = args.weight_requires_grad and with_input_all_gather + + # Slot 5 -- ``ln_out_to_save``. + ln_out_alias = None + ln_out_to_save = None + if args.backward_override == "high_precision": + # ``ln_out_hp`` is taken before ln_out may be dropped, so it is always saved. + ln_out_to_save = TensorSpec( + shape=tuple(args.inp.shape), dtype=args.activation_dtype, device=device + ) + if args.return_layernorm_output and not ln_out_return_is_total: + ln_out_alias = "ln_out" + elif args.weight_requires_grad or args.return_layernorm_output: + ln_out_to_save = ln_out_spec + if ( + backward_needs_input + and args.backward_override is None + and ln_out_quantized + and ( + isinstance(args.input_quantizer, (MXFP8Quantizer, Float8BlockQuantizer)) + or not ln_out_needs_gather + ) + ): + ln_out_to_save.update_usage(rowwise_usage=False) + if args.return_layernorm_output and not ln_out_return_is_total and not ln_out_rebound: + ln_out_alias = "ln_out" + + # Slot 1 -- ``wt_save``, with the impl's alias dedup. + wt_alias = None + wt_save = None + if weightmat_aliases_weight: + wt_alias = "weight" + elif args.is_fsdp2: + pass # FSDP2 re-quantizes from the gathered args.weight in backward. + elif weightmat_is_storage and new_weight_workspace is not None: + wt_alias = "new_weight_workspace" + elif weightmat_is_storage and workspace is not None: + wt_alias = "weight_workspace" + elif weightmat_is_storage: + wt_save = weightmat + else: + wt_save = TensorSpec( + shape=tuple(args.weight.shape), + dtype=args.activation_dtype, + device=args.weight.device, + ) + + saved_tensor_aliases = ( + "inp" if inputmat_aliases_inp else None, + wt_alias, + "weight", + "bias" if bias_aliases else None, + "ln_weight" if ln_weight_aliases else None, + ln_out_alias, + None, + None, + ) + tensors_to_save_from_forward = ( + ( + None + if inputmat_aliases_inp + else TensorSpec( + shape=tuple(args.inp.shape), dtype=args.activation_dtype, device=device + ) + ), + wt_save, + None, + ( + None + if (args.bias is None or bias_aliases) + else TensorSpec(shape=tuple(args.bias.shape), dtype=bias_dtype, device=device) + ), + ( + None + if ln_weight_aliases + else TensorSpec( + shape=tuple(args.ln_weight.shape), dtype=args.activation_dtype, device=device + ) + ), + None if ln_out_alias is not None else ln_out_to_save, + mu, + rsigma, + ) + ctx_attrs = { + "fsdp_shapes": [], + "saved_tensor_aliases": saved_tensor_aliases, + "is_weight_param_quantized": is_weight_param_quantized, + "ln_out_needs_gather": ln_out_needs_gather, + } + + return out, ln_out_for_return, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs + + +def _layernorm_linear_backward_fake( + args: LayerNormLinearBwdArgs, +) -> Tuple[Optional[TensorSpec], ...]: + """Allocation-free fake of :func:`_layernorm_linear_backward_impl` on + ``TensorSpec``. Returns ``(dgrad, dgamma, dbeta, wgrad, grad_bias)`` specs; + TP/SP gather/scatter happens inside the eager op, so the specs carry + rank-local shapes.""" + if args.fsdp_group is not None: + raise NotImplementedError( + "Fake LayerNormLinear backward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight = args.saved_weight + out_dtype = args.activation_dtype + out_features, in_features = weight.shape + device = args.grad_output.device + + # Mirrors the impl; affects dgrad's buffer layout. + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) + + dgrad = None + if args.requires_dgrad: + dgrad = TensorSpec( + shape=args.inp_shape, + dtype=out_dtype, + device=device, + ) + + # The norm backward always runs; its grads take the saved (cast) ln_weight dtype. + dgamma = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=device) + dbeta = None + if args.normalization == "LayerNorm": + dbeta = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=device) + + wgrad = None + # Under fuse_wgrad_accumulation the grad goes into main_grad in place. + if args.requires_wgrad and not args.fuse_wgrad_accumulation: + wgrad = TensorSpec( + shape=(out_features, in_features), + dtype=out_dtype, + quantizer=args.grad_weight_quantizer, + device=weight.device, + ) + + grad_bias = None + # FP8 backward computes bgrad in grad_output_preprocess whenever bias is + # used; in high precision it is fused into the wgrad GEMM, so it only + # exists when wgrad runs. + fp8_bwd = args.fp8 and args.backward_override is None + if args.use_bias and (args.requires_wgrad or fp8_bwd): + grad_bias = TensorSpec(shape=(out_features,), dtype=out_dtype, device=device) + + return dgrad, dgamma, dbeta, wgrad, grad_bias + + +# Custom op used under ``torch.compile``. +_layernorm_linear_op = register_custom_op_with_autograd( + op_name="layernorm_linear", + input_tensors_for_grad=["inp", "ln_weight", "ln_bias", "weight", "bias"], + fwd_arg_type=LayerNormLinearFwdArgs, + fwd_impl=_layernorm_linear_forward_impl, + fwd_fake_impl=_layernorm_linear_forward_fake, + setup_context=_layernorm_linear_setup_ctx, + bwd_arg_type=LayerNormLinearBwdArgs, + bwd_impl=_layernorm_linear_backward_impl, + bwd_fake_impl=_layernorm_linear_backward_fake, +) + + class _LayerNormLinear(torch.autograd.Function): """LayerNormLinear semi-top level module Calls custom cuda extensions. @@ -1530,8 +1893,7 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: """Backward pass: compute gradients and reduce FP8 scaling factors.""" bwd_args: LayerNormLinearBwdArgs = ctx.backward_objects - bwd_args.grad_output = grad_output - bwd_args.grad_ln_out = grad_ln_out + bwd_args.setup_grad_outputs((grad_output, grad_ln_out, _grad_weight_workspace)) bwd_args.setup_saved_tensors(ctx) nvtx_label = "transformer_engine._LayerNormLinear.backward" if bwd_args.ub_name is not None: @@ -1556,6 +1918,22 @@ def backward( ) +@no_torch_dynamo() +def _layernorm_linear_eager( + inp: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: Optional[torch.Tensor], + weight: torch.Tensor, + bias: Optional[torch.Tensor], + fwd_args: LayerNormLinearFwdArgs, + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Run ``_LayerNormLinear`` eagerly, bypassing Dynamo.""" + if is_grad_enabled: + return _LayerNormLinear.apply(inp, ln_weight, ln_bias, weight, bias, fwd_args) + return _LayerNormLinear.forward(None, inp, ln_weight, ln_bias, weight, bias, fwd_args) + + class LayerNormLinear(TransformerEngineBaseModule): r""" Applies layer normalization followed by linear transformation to the incoming data. @@ -2023,7 +2401,6 @@ def reset_parameters(self, defer_init=False): elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1) - @no_torch_dynamo() def forward( self, inp: torch.Tensor, @@ -2051,6 +2428,13 @@ def forward( * it also allows skipping gradient accumulation during the first microbatch (since it is the first gradient being produced) + + Under ``torch.compile``, the FP8 weight cache is replaced on + the first microbatch and reused on subsequent microbatches. + ``fuse_wgrad_accumulation=True`` remains unsupported. + With ``mode="reduce-overhead"``, call + ``torch.compiler.cudagraph_mark_step_begin()`` once before + each minibatch to keep the cache live across its microbatches. """ is_grad_enabled = torch.is_grad_enabled() @@ -2106,10 +2490,19 @@ def forward( weight_quantizer, weight_tensor ) + use_compiled_op = torch.compiler.is_compiling() and _layernorm_linear_op is not None + if _layernorm_linear_op is None and torch.compiler.is_compiling(): + warn_if_compile_disabled() + if use_compiled_op: + # Process groups cross the op boundary separately from quantizers. + for quantizer in (input_quantizer, grad_output_quantizer): + if getattr(quantizer, "amax_reduction_group", None) is not None: + set_quantizer_amax_reduction_group(quantizer, None) + cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" - weight_workspace = ( - self._fp8_workspaces.get(cache_name) if cache_name is not None else None - ) + weight_workspace = None + if cache_name is not None and not (use_compiled_op and is_first_microbatch): + weight_workspace = self._fp8_workspaces.get(cache_name) dgrad_use_split_accumulator = _2X_ACC_DGRAD wgrad_use_split_accumulator = _2X_ACC_WGRAD @@ -2141,6 +2534,25 @@ def forward( linear_bias_tensor = ( bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None ) + if use_compiled_op: + fallback_reason = self._compile_eager_fallback_reason( + inp, + skip_fp8_weight_update, + fp8_output, + fp8_grad, + is_grad_enabled, + debug, + quantizers, + ) + if fallback_reason is not None: + warn_compile_eager_fallback(fallback_reason) + torch._dynamo.graph_break( + msg=f"te.LayerNormLinear falling back to eager: {fallback_reason}" + ) + use_compiled_op = False + + wgrad_store = None if use_compiled_op else self.wgrad_store + fwd_args = LayerNormLinearFwdArgs( # tensors inp=inp, @@ -2210,30 +2622,24 @@ def forward( is_fsdp2=self.is_fsdp2, # weight-grad scheduling fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, - wgrad_store=self.wgrad_store, + wgrad_store=wgrad_store, # misc cpu_offloading=is_cpu_offload_enabled(), is_grad_enabled=is_grad_enabled, ) - if is_grad_enabled: - out, ln_out, new_weight_workspace = _LayerNormLinear.apply( - inp, - self.layer_norm_weight, - self.layer_norm_bias, - weight_tensor, - linear_bias_tensor, - fwd_args, - ) + if use_compiled_op: + check_gemm_dims(inp, weight_tensor, self.fp8) + out, ln_out, new_weight_workspace = _layernorm_linear_op(fwd_args) else: - out, ln_out, new_weight_workspace = _LayerNormLinear.forward( - None, + out, ln_out, new_weight_workspace = _layernorm_linear_eager( inp, self.layer_norm_weight, self.layer_norm_bias, weight_tensor, linear_bias_tensor, fwd_args, + is_grad_enabled, ) if new_weight_workspace is not None and cache_name is not None: @@ -2300,6 +2706,57 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): for name, q in zip(names, original_quantizers) ) + def _compile_eager_fallback_reason( + self, + inp: torch.Tensor, + skip_fp8_weight_update: Optional[torch.Tensor], + fp8_output: bool, + fp8_grad: bool, + is_grad_enabled: bool, + debug: bool, + quantizers: Sequence[Optional[Quantizer]], + ) -> Optional[str]: + """Return a fallback reason after preparing the forward quantizers.""" + if debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + if is_distributed_weight(weight_tensor): + return "a DistributedWeight (custom weight parallelism, e.g. GTP)" + if isinstance(inp, (QuantizedTensor, QuantizedTensorStorage)): + return "a quantized input tensor" + if self.fsdp_group is not None: + return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" + any_requires_grad = ( + inp.requires_grad + or self.layer_norm_weight.requires_grad + or (self.layer_norm_bias is not None and self.layer_norm_bias.requires_grad) + or weight_tensor.requires_grad + or (bias_tensor is not None and bias_tensor.requires_grad) + ) + if fp8_output and is_grad_enabled and any_requires_grad: + return "differentiable fp8_output=True" + if is_cpu_offload_enabled(): + return "CPU activation offloading" + # A queued store can be enabled between forward and backward. + if self.wgrad_store is not None and ( + self.wgrad_store.context is not None or self.wgrad_store.delay_wgrad_compute() + ): + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + fp8 = FP8GlobalStateManager.is_fp8_enabled() + needs_dgrad = is_grad_enabled and inp.requires_grad + if ( + fp8 + and fp8_grad + and needs_dgrad + and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad) + ): + return "a quantized input grad (fp8_grad=True)" + if fp8 and skip_fp8_weight_update is not None: + return "GPU-controlled FP8 weight cache updates" + return compile_unsupported_quantizer_reason(quantizers) + def _get_weight_and_bias_tensors(self): # Get concatenated weight and bias tensors unfused_weights = self._get_weight_tensors() diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ba15f1c7bd1..d22fe30ad2a 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -50,6 +50,8 @@ clear_tensor_data, needs_quantized_gemm, get_nvtx_range_context, + warn_compile_eager_fallback, + warn_if_compile_disabled, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -73,9 +75,11 @@ from ..tensor.hybrid_tensor import HybridQuantizer from ..tensor.identity_tensor import IdentityQuantizer from ._common import ( + compile_unsupported_quantizer_reason, sum_bias_grad, apply_normalization, check_fp8_reduce_and_update, + fake_quantized_weight, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -87,12 +91,18 @@ mark_activation_offload, ) from ..quantized_tensor import ( + QuantizedTensor, QuantizedTensorStorage, Quantizer, prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorOrQuantized +from ..dynamo.fake import fake_activation, fake_norm +from ..dynamo import ( + TensorSpec, + TensorOrQuantized, + register_custom_op_with_autograd, +) from ..cpp_extensions import ( general_gemm, ) @@ -102,8 +112,10 @@ __all__ = ["LayerNormMLP"] -def _get_act_func_supported_list(recipe: Optional[Recipe] = None): - if recipe is None: +def _get_act_func_supported_list( + recipe: Optional[Recipe] = None, dbias_fusion: Optional[bool] = None +): + if dbias_fusion is False or (dbias_fusion is None and recipe is None): # bf16 (recipe is None): return { "gelu": (tex.gelu, tex.dgelu, None), @@ -119,7 +131,7 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): "swiglu": (tex.swiglu, tex.dswiglu, None), "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu, None), } - if recipe.delayed() or recipe.mxfp8(): + if dbias_fusion is True or recipe.delayed() or recipe.mxfp8(): # Delayed scaling, fusion supported list: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] # MXFP8: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] return { @@ -162,18 +174,26 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): raise NotImplementedError(f"Unhandled recipe type {recipe}") -def _act_func(activation: str, recipe: Optional[Recipe] = None): +def _act_func( + activation: str, recipe: Optional[Recipe] = None, dbias_fusion: Optional[bool] = None +): # based on each quantization mode, we have different kernel fusion supported: # bf16 (recipe is None): [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] # Delayed scaling, fusion supported list: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] # MXFP8: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] # Per-tensor current scaling or fp8 blockwise scaling: [] - funcs = _get_act_func_supported_list(recipe) + funcs = _get_act_func_supported_list(recipe, dbias_fusion) if activation not in funcs: raise NotImplementedError("Activation type " + activation + " is not supported!") return funcs[activation] +# Activations whose output halves the last dim (gated linear units). +_GATED_ACTIVATIONS = frozenset( + {"geglu", "glu", "qgeglu", "reglu", "sreglu", "swiglu", "clamped_swiglu"} +) + + @dataclass(slots=True) class LayerNormMLPFwdArgs: """Single-argument bag for the forward path of :class:`_LayerNormMLP`.""" @@ -236,6 +256,10 @@ class LayerNormMLPFwdArgs: backward_override: Optional[str] dgrad_use_split_accumulator: bool wgrad_use_split_accumulator: bool + # Recipe properties the backward needs (the recipe itself can't cross the op boundary). + recipe_float8_block_scaling: bool + recipe_custom: bool + recipe_dbias_dact_fusion: bool debug: bool # --- Weight-workspace caching --- @@ -356,7 +380,9 @@ class LayerNormMLPBwdArgs: # --- Numerical / dtype config --- activation_dtype: Optional[torch.dtype] = None fp8: bool = False - fp8_recipe: Optional[Any] = None + recipe_float8_block_scaling: bool = False + recipe_custom: bool = False + recipe_dbias_dact_fusion: bool = False dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD backward_override: Optional[str] = None @@ -400,6 +426,10 @@ class LayerNormMLPBwdArgs: # --- Per-backward scratch state (populated inside the backward impl) --- ub_obj_gradout: Optional[Any] = None + def setup_grad_outputs(self, grads: Sequence[Any]) -> None: + """Unpack gradients in forward-output order, ignoring weight workspaces.""" + self.grad_output, self.grad_ln_out, _, _ = grads + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: """Pull saved tensors from ``ctx`` into the fields backward consumes.""" self.set_saved_tensors(restore_from_func_ctx(ctx)) @@ -1168,7 +1198,9 @@ def _layernorm_mlp_setup_ctx( # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype bwd_args.fp8 = fp8 - bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + bwd_args.recipe_float8_block_scaling = fwd_args.recipe_float8_block_scaling + bwd_args.recipe_custom = fwd_args.recipe_custom + bwd_args.recipe_dbias_dact_fusion = fwd_args.recipe_dbias_dact_fusion bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator bwd_args.backward_override = fwd_args.backward_override @@ -1630,7 +1662,7 @@ def _layernorm_mlp_backward_impl( # Whether to set grad arg in general_gemm grad_arg = True - if args.fp8 and args.fp8_recipe.float8_block_scaling(): + if args.fp8 and args.recipe_float8_block_scaling: grad_arg = False # Arguments to include in wgrad GEMM closure @@ -1676,7 +1708,7 @@ def fc2_wgrad_gemm( # Update grad bias if needed if fc2_bias_grad is None: - if args.fp8 and args.fp8_recipe.float8_block_scaling() and fc2_bias is not None: + if args.fp8 and args.recipe_float8_block_scaling and fc2_bias is not None: # BGRAD not fused with GEMM for float8 blockwise gemm. fc2_bias_grad_ = sum_bias_grad(act_out) fc2_bias_grad = fc2_bias_grad_ @@ -1709,12 +1741,13 @@ def fc2_wgrad_gemm( fc1_bias_grad = sum_bias_grad(dact) dact = args.fc1_grad_output_quantizer(dact) elif ( - _act_func(args.activation, args.fp8_recipe if args.fp8 else None)[2] is not None + _act_func(args.activation, dbias_fusion=args.fp8 and args.recipe_dbias_dact_fusion)[2] + is not None and args.fp8 ): # Fusion: gemm, bias + gelu + quantize dbias_dact_quantize_func = _act_func( - args.activation, args.fp8_recipe if args.fp8 else None + args.activation, dbias_fusion=args.fp8 and args.recipe_dbias_dact_fusion )[2] fc1_bias_grad, dact = dbias_dact_quantize_func( fc2_dgrad, @@ -1726,7 +1759,7 @@ def fc2_wgrad_gemm( # Fusion: gemm + gelu, if not fc2_dgrad_gemm_gelu_fusion: activation_func_bwd = _act_func( - args.activation, args.fp8_recipe if args.fp8 else None + args.activation, dbias_fusion=args.fp8 and args.recipe_dbias_dact_fusion )[1] dact = activation_func_bwd( fc2_dgrad, fc1_out.to(args.activation_dtype), None, **act_params @@ -1739,7 +1772,7 @@ def fc2_wgrad_gemm( args.fc1_grad_output_quantizer, (Float8BlockQuantizer, IdentityQuantizer), ) - or args.fp8_recipe.custom() + or args.recipe_custom ): fc1_bias_grad = sum_bias_grad(dact) dact = args.fc1_grad_output_quantizer(dact) @@ -2071,6 +2104,420 @@ def fc1_wgrad_gemm( ) +def _layernorm_mlp_forward_fake( + args: LayerNormMLPFwdArgs, +) -> Tuple[ + TensorSpec, + Optional[TensorSpec], + Optional[TensorSpec], + Optional[TensorSpec], + Optional[Tuple[Any, ...]], + Optional[Dict], +]: + """Shape/metadata-only twin of :func:`_layernorm_mlp_forward_impl` for + torch.compile (no args.activation checkpointing: that path falls back to eager).""" + if args.fsdp_group is not None and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time LayerNormMLP forward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + if args.checkpoint and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time LayerNormMLP forward does not support args.activation checkpointing" + ) + + fc1_weight_quantizer = args.fc1_weight_quantizer + fc2_weight_quantizer = args.fc2_weight_quantizer + fp8_or_debug = args.fp8 or args.debug + bias_gelu_fusion = args.bias_gelu_fusion + gemm_gelu_fusion = args.gemm_gelu_fusion + device = args.inp.device + + in_features = args.ln_weight.shape[0] + fc1_out_features = args.fc1_weight.shape[0] + fc2_out_features = args.fc2_weight.shape[0] + inp_leading = args.inp.shape[0] if len(args.inp.shape) > 1 else 1 + gemm_leading_shape = ( + (inp_leading * args.tp_size, *args.inp.shape[1:-1]) + if args.sequence_parallel + else tuple(args.inp.shape[:-1]) + ) + inputmat_aliases_inp = args.inp.dtype == args.activation_dtype + ln_weight_aliases = args.ln_weight.dtype == args.activation_dtype + + backwards_needs_fc1_input = args.fc1_weight_requires_grad and args.is_grad_enabled + + # Norm output quantizer usage -- mirrors the impl; the ln_out spec must be + # taken now, before the later rowwise-only usage for the all-gather. + if args.fp8: + if args.fc1_input_quantizer is None: + raise ValueError("Missing quantizer for FC1 input tensor") + args.fc1_input_quantizer.set_usage(rowwise=True, columnwise=backwards_needs_fc1_input) + if args.sequence_parallel and args.fc1_input_quantizer.supports_only_rowwise_all_gather(): + args.fc1_input_quantizer.set_usage(columnwise=False) + custom = is_custom(args.fc1_input_quantizer) + hybrid = isinstance(args.fc1_input_quantizer, HybridQuantizer) + identity = isinstance(args.fc1_input_quantizer, IdentityQuantizer) + with_quantized_norm = ( + args.fp8 + and not args.debug + and not args.return_layernorm_output + and not args.return_layernorm_output_gathered + and not custom + and not hybrid + and not identity + ) + # A custom quantizer's ln_out stays in high precision on the plain + # all-gather path (only the gathered copy is quantized). + ln_out_quantized = fp8_or_debug and not ( + args.sequence_parallel and not args.return_layernorm_output_gathered and custom + ) + ln_out, mu, rsigma = fake_norm( + args.inp, + args.fc1_input_quantizer if with_quantized_norm else None, + normalization=args.normalization, + dtype=args.activation_dtype, + ) + if ln_out_quantized and not with_quantized_norm: + ln_out = TensorSpec( + shape=ln_out.shape, + dtype=args.activation_dtype, + quantizer=args.fc1_input_quantizer, + device=device, + ) + if args.sequence_parallel and fp8_or_debug: + args.fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_return_is_total = args.sequence_parallel and args.return_layernorm_output_gathered + ln_out_rebound = ln_out_quantized and not with_quantized_norm + + # ------------------------------------------------------ + # Weights -- mirror ``quantize_weight`` / ``cast_if_needed``. + # ------------------------------------------------------ + new_fc1_weight_workspace = None + new_fc2_weight_workspace = None + if fp8_or_debug: + if args.fc1_weight.is_quantized and not args.debug: + fc1_weight_quantizer = args.fc1_weight.quantizer + elif fc1_weight_quantizer is not None: + fc1_weight_quantizer.set_usage( + rowwise=True, columnwise=args.is_grad_enabled and not args.is_fsdp2 + ) + if args.fc2_weight.is_quantized and not args.debug: + fc2_weight_quantizer = args.fc2_weight.quantizer + elif fc2_weight_quantizer is not None: + fc2_weight_quantizer.set_usage( + rowwise=True, columnwise=args.is_grad_enabled and not args.is_fsdp2 + ) + fc1_weight_final, new_fc1_weight_workspace, fc1_wt_alias = fake_quantized_weight( + args.fc1_weight, + fc1_weight_quantizer, + args.fc1_weight_workspace, + activation_dtype=args.activation_dtype, + cache_weight=args.cache_weight, + ) + fc2_weight_final, new_fc2_weight_workspace, fc2_wt_alias = fake_quantized_weight( + args.fc2_weight, + fc2_weight_quantizer, + args.fc2_weight_workspace, + activation_dtype=args.activation_dtype, + cache_weight=args.cache_weight, + ) + fc1_wt_alias = ( + None if fc1_wt_alias is None else fc1_wt_alias.replace("weight", "fc1_weight") + ) + fc2_wt_alias = ( + None if fc2_wt_alias is None else fc2_wt_alias.replace("weight", "fc2_weight") + ) + else: + fc1_weight_final = TensorSpec( + shape=tuple(args.fc1_weight.shape), + dtype=args.activation_dtype, + device=args.fc1_weight.device, + ) + fc2_weight_final = TensorSpec( + shape=tuple(args.fc2_weight.shape), + dtype=args.activation_dtype, + device=args.fc2_weight.device, + ) + fc1_wt_alias = "fc1_weight" if args.fc1_weight.dtype == args.activation_dtype else None + fc2_wt_alias = "fc2_weight" if args.fc2_weight.dtype == args.activation_dtype else None + + # Bias cast: cuBLAS has no FP8 GEMM with FP32 bias. + bias_dtype = args.activation_dtype + if fp8_or_debug and args.activation_dtype == torch.float32: + bias_dtype = torch.bfloat16 + fc1_bias_aliases = args.fc1_bias is not None and args.fc1_bias.dtype == bias_dtype + fc2_bias_aliases = args.fc2_bias is not None and args.fc2_bias.dtype == bias_dtype + + # ------------------------------------------------------ + # FC1 GEMM + args.activation (see the impl for the fusion rules). + # ------------------------------------------------------ + if args.activation != "gelu": + gemm_gelu_fusion = bias_gelu_fusion = False + else: + if not args.fp8: + gemm_gelu_fusion = True + if gemm_gelu_fusion and bias_gelu_fusion: + gemm_gelu_fusion = False + if args.debug: + gemm_gelu_fusion = False + fc1_out = None + fc1_out_without_bias = None + if bias_gelu_fusion: + fc1_out_without_bias = TensorSpec( + shape=(*gemm_leading_shape, fc1_out_features), + dtype=args.activation_dtype, + device=device, + ) + else: + # generic_gemm stores the pre-GELU auxiliary output in the bias dtype + # for low-precision inputs, defaulting to BF16 when there is no bias. + fc1_out_dtype = args.activation_dtype + if gemm_gelu_fusion and args.fp8: + fc1_out_dtype = bias_dtype if args.fc1_bias is not None else torch.bfloat16 + fc1_out = TensorSpec( + shape=(*gemm_leading_shape, fc1_out_features), + dtype=fc1_out_dtype, + device=device, + ) + act_out = fake_activation( + fc1_out_without_bias if bias_gelu_fusion else fc1_out, + args.fc2_input_quantizer if fp8_or_debug else None, + gated=args.activation in _GATED_ACTIVATIONS, + dtype=args.activation_dtype, + fused_quantization=( + args.fp8 + and not args.debug + and not gemm_gelu_fusion + and not args.recipe_float8_block_scaling + and not args.recipe_custom + ), + ) + if args.fc2_output_quantizer is not None: + args.fc2_output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Outputs. + # ------------------------------------------------------ + requires_grad = args.is_grad_enabled and args.any_requires_grad() + out_leading = inp_leading + if args.sequence_parallel and not args.set_parallel_mode: + out_leading = inp_leading * args.tp_size + fc2_out = TensorSpec( + shape=( + (fc2_out_features,) + if len(args.inp.shape) == 1 + else (out_leading, *tuple(args.inp.shape[1:-1]), fc2_out_features) + ), + dtype=args.activation_dtype, + quantizer=args.fc2_output_quantizer, + requires_grad=requires_grad, + device=device, + ) + ln_out_for_return = None + if args.return_layernorm_output: + ln_leading = inp_leading + if ln_out_return_is_total: + ln_leading = inp_leading * args.tp_size + ln_out_for_return = TensorSpec( + shape=( + (ln_leading, *tuple(args.inp.shape[1:])) + if len(args.inp.shape) > 1 + else (in_features,) + ), + dtype=args.activation_dtype, + requires_grad=requires_grad, + device=device, + ) + + # ------------------------------------------------------ + # Backward state -- saved-tensor layout (see the impl). + # ------------------------------------------------------ + tensors_to_save_from_forward = None + ctx_attrs = None + if args.is_grad_enabled: + if not args.fc1_weight_requires_grad: + ln_out = None + if not args.fc2_weight_requires_grad: + act_out = None + fc1_wt_save = fc1_weight_final + fc2_wt_save = fc2_weight_final + if args.is_fsdp2: + if fc1_wt_alias != "fc1_weight": + fc1_wt_save, fc1_wt_alias = None, None + if fc2_wt_alias != "fc2_weight": + fc2_wt_save, fc2_wt_alias = None, None + ln_out_alias = ( + "ln_out" + if args.return_layernorm_output + and ln_out is not None + and not ln_out_return_is_total + and not ln_out_rebound + else None + ) + saved_tensor_aliases = ( + "inp" if inputmat_aliases_inp else None, + "ln_weight" if ln_weight_aliases else None, + ln_out_alias, + fc1_wt_alias, + "fc1_weight", + "fc1_bias" if fc1_bias_aliases else None, + None, + None, + None, + fc2_wt_alias, + "fc2_weight", + "fc2_bias" if fc2_bias_aliases else None, + None, + None, + ) + saved = ( + TensorSpec(shape=tuple(args.inp.shape), dtype=args.activation_dtype, device=device), + TensorSpec( + shape=tuple(args.ln_weight.shape), dtype=args.activation_dtype, device=device + ), + ln_out, + fc1_wt_save, + None, + ( + None + if args.fc1_bias is None + else TensorSpec(shape=tuple(args.fc1_bias.shape), dtype=bias_dtype, device=device) + ), + fc1_out, + fc1_out_without_bias, + act_out, + fc2_wt_save, + None, + ( + None + if args.fc2_bias is None + else TensorSpec(shape=tuple(args.fc2_bias.shape), dtype=bias_dtype, device=device) + ), + mu, + rsigma, + ) + tensors_to_save_from_forward = tuple( + None if alias is not None else spec for alias, spec in zip(saved_tensor_aliases, saved) + ) + ctx_attrs = { + "saved_tensor_aliases": saved_tensor_aliases, + "fsdp_shapes": None, + "is_recomputation": False, + } + + return ( + fc2_out, + ln_out_for_return, + new_fc1_weight_workspace, + new_fc2_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) + + +def _layernorm_mlp_backward_fake( + args: LayerNormMLPBwdArgs, +) -> Tuple[Optional[TensorSpec], ...]: + """Allocation-free fake of :func:`_layernorm_mlp_backward_impl` on + ``TensorSpec``. Returns ``(dgrad, dgamma, dbeta, fc1_wgrad, fc1_bias_grad, + fc2_wgrad, fc2_bias_grad)`` specs with rank-local shapes.""" + if args.fsdp_group is not None: + raise NotImplementedError( + "Fake LayerNormMLP backward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + out_dtype = args.activation_dtype + device = args.grad_output.device + in_features = args.ln_weight.shape[-1] + fc1_out_features = args.fc1_weight.shape[0] + fc2_out_features, act_features = args.fc2_weight.shape + + if args.fc1_grad_output_quantizer is not None: + args.fc1_grad_output_quantizer.set_usage(rowwise=True, columnwise=True) + + dgrad = None + if args.requires_dgrad: + inp_leading = args.grad_output.shape[0] + if args.sequence_parallel and not args.set_parallel_mode: + inp_leading = inp_leading // args.tp_size + dgrad = TensorSpec( + shape=( + (in_features,) + if len(args.grad_output.shape) == 1 + else (inp_leading, *args.grad_output.shape[1:-1], in_features) + ), + dtype=out_dtype, + device=device, + ) + + # The norm backward always runs; its grads take the saved (cast) ln_weight dtype. + dgamma = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=device) + dbeta = None + if args.normalization == "LayerNorm": + dbeta = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=device) + + fc1_wgrad = None + if args.fc1_weight_requires_grad and not args.fuse_wgrad_accumulation: + fc1_wgrad = TensorSpec( + shape=(fc1_out_features, in_features), + dtype=out_dtype, + quantizer=args.fc1_grad_weight_quantizer, + device=device, + ) + fc1_bias_grad = None + # FP8, debug and the fused bias-gelu path always produce it; in high + # precision it is fused into the FC1 wgrad GEMM or summed when only the + # bias needs a grad. + if args.fc1_bias is not None and any( + ( + args.fp8, + args.debug, + args.bias_gelu_fusion, + args.fc1_weight_requires_grad, + args.fc1_bias_requires_grad, + ) + ): + fc1_bias_grad = TensorSpec(shape=(fc1_out_features,), dtype=out_dtype, device=device) + + fc2_wgrad = None + if args.fc2_weight_requires_grad and not args.fuse_wgrad_accumulation: + fc2_wgrad = TensorSpec( + shape=(fc2_out_features, act_features), + dtype=out_dtype, + quantizer=args.fc2_grad_weight_quantizer, + device=device, + ) + fc2_bias_grad = None + fp8_bwd = args.fp8 and args.backward_override is None + if args.use_bias and (args.fc2_weight_requires_grad or fp8_bwd): + fc2_bias_grad = TensorSpec(shape=(fc2_out_features,), dtype=out_dtype, device=device) + + return dgrad, dgamma, dbeta, fc1_wgrad, fc1_bias_grad, fc2_wgrad, fc2_bias_grad + + +# Custom op used under ``torch.compile``. +_layernorm_mlp_op = register_custom_op_with_autograd( + op_name="layernorm_mlp", + input_tensors_for_grad=[ + "inp", + "ln_weight", + "ln_bias", + "fc1_weight", + "fc1_bias", + "fc2_weight", + "fc2_bias", + ], + fwd_arg_type=LayerNormMLPFwdArgs, + fwd_impl=_layernorm_mlp_forward_impl, + fwd_fake_impl=_layernorm_mlp_forward_fake, + setup_context=_layernorm_mlp_setup_ctx, + bwd_arg_type=LayerNormMLPBwdArgs, + bwd_impl=_layernorm_mlp_backward_impl, + bwd_fake_impl=_layernorm_mlp_backward_fake, +) + + class _LayerNormMLP(torch.autograd.Function): """LayerNormMLP semi-top level module Calls custom cuda extensions. @@ -2141,8 +2588,9 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: """Backward pass: compute gradients and reduce FP8 scaling factors.""" bwd_args: LayerNormMLPBwdArgs = ctx.backward_objects - bwd_args.grad_output = grad_output - bwd_args.grad_ln_out = grad_ln_out + bwd_args.setup_grad_outputs( + (grad_output, grad_ln_out, _grad_fc1_weight_workspace, _grad_fc2_weight_workspace) + ) with get_nvtx_range_context("_LayerNormMLP_backward"): _layernorm_mlp_recompute(bwd_args, ctx) ( @@ -2173,6 +2621,18 @@ def backward( ) +@no_torch_dynamo() +def _layernorm_mlp_eager( + tensors: Tuple[Optional[torch.Tensor], ...], + fwd_args: LayerNormMLPFwdArgs, + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """Run ``_LayerNormMLP`` eagerly, bypassing Dynamo.""" + if is_grad_enabled: + return _LayerNormMLP.apply(*tensors, fwd_args) + return _LayerNormMLP.forward(None, *tensors, fwd_args) + + class LayerNormMLP(TransformerEngineBaseModule): r""" Applies layer normalization on the input followed by the MLP module, consisting of @@ -2603,7 +3063,6 @@ def reset_parameters(self, defer_init=False): if self.set_parallel_mode: setattr(self.fc2_bias, "sequence_parallel", self.sequence_parallel) - @no_torch_dynamo() def forward( self, inp: torch.Tensor, @@ -2629,6 +3088,13 @@ def forward( * it also allows skipping gradient accumulation during the first microbatch (since it is the first gradient being produced) + + Under ``torch.compile``, the FP8 weight cache is replaced on + the first microbatch and reused on subsequent microbatches. + ``fuse_wgrad_accumulation=True`` remains unsupported. + With ``mode="reduce-overhead"``, call + ``torch.compiler.cudagraph_mark_step_begin()`` once before + each minibatch to keep the cache live across its microbatches. """ is_grad_enabled = torch.is_grad_enabled() @@ -2703,21 +3169,32 @@ def forward( if self.bias_gelu_nvfusion and not use_reentrant_activation_recompute(): self.fast_setattr("bias_gelu_nvfusion", False) + use_compiled_op = torch.compiler.is_compiling() and _layernorm_mlp_op is not None + if _layernorm_mlp_op is None and torch.compiler.is_compiling(): + warn_if_compile_disabled() + if use_compiled_op: + # Process groups cross the op boundary separately from quantizers. + for quantizer in (fc1_input_quantizer, fc2_grad_output_quantizer): + if getattr(quantizer, "amax_reduction_group", None) is not None: + set_quantizer_amax_reduction_group(quantizer, None) + cache_name_fc1 = ( None if (is_first_microbatch is None or self.is_fsdp2) else "fc1_weight" ) cache_name_fc2 = ( None if (is_first_microbatch is None or self.is_fsdp2) else "fc2_weight" ) - fc1_weight_workspace = ( - self._fp8_workspaces.get(cache_name_fc1) if cache_name_fc1 is not None else None - ) - fc2_weight_workspace = ( - self._fp8_workspaces.get(cache_name_fc2) if cache_name_fc2 is not None else None - ) + fc1_weight_workspace = None + fc2_weight_workspace = None + if cache_name_fc1 is not None and not (use_compiled_op and is_first_microbatch): + fc1_weight_workspace = self._fp8_workspaces.get(cache_name_fc1) + fc2_weight_workspace = self._fp8_workspaces.get(cache_name_fc2) dgrad_use_split_accumulator = _2X_ACC_DGRAD wgrad_use_split_accumulator = _2X_ACC_WGRAD + recipe_float8_block_scaling = False + recipe_custom = False + recipe_dbias_dact_fusion = False if self.fp8: _recipe = FP8GlobalStateManager.get_fp8_recipe() backward_override = _recipe.backward_override @@ -2725,6 +3202,9 @@ def forward( dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator if hasattr(_recipe, "fp8_gemm_wgrad"): wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator + recipe_float8_block_scaling = _recipe.float8_block_scaling() + recipe_custom = _recipe.custom() + recipe_dbias_dact_fusion = _recipe.delayed() or _recipe.mxfp8() else: backward_override = None @@ -2744,6 +3224,24 @@ def forward( fc2_bias_tensor = ( fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None ) + if use_compiled_op: + fallback_reason = self._compile_eager_fallback_reason( + inp, + skip_fp8_weight_update, + fp8_output, + is_grad_enabled, + debug, + quantizers, + ) + if fallback_reason is not None: + warn_compile_eager_fallback(fallback_reason) + torch._dynamo.graph_break( + msg=f"te.LayerNormMLP falling back to eager: {fallback_reason}" + ) + use_compiled_op = False + + wgrad_store = None if use_compiled_op else self.wgrad_store + fwd_args = LayerNormMLPFwdArgs( # tensors inp=inp, @@ -2804,6 +3302,9 @@ def forward( backward_override=backward_override, dgrad_use_split_accumulator=dgrad_use_split_accumulator, wgrad_use_split_accumulator=wgrad_use_split_accumulator, + recipe_float8_block_scaling=recipe_float8_block_scaling, + recipe_custom=recipe_custom, + recipe_dbias_dact_fusion=recipe_dbias_dact_fusion, debug=debug, # weight-workspace caching is_first_microbatch=is_first_microbatch, @@ -2827,7 +3328,7 @@ def forward( is_fsdp2=self.is_fsdp2, # weight-grad scheduling fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, - wgrad_store=self.wgrad_store, + wgrad_store=wgrad_store, # activation checkpointing checkpoint=self.checkpoint, fp8_meta=self.fp8_meta if self.checkpoint else None, @@ -2837,28 +3338,21 @@ def forward( is_grad_enabled=is_grad_enabled, ) - if is_grad_enabled: - out, ln_out, new_fc1_ws, new_fc2_ws = _LayerNormMLP.apply( - inp, - self.layer_norm_weight, - self.layer_norm_bias, - fc1_weight, - fc1_bias, - fc2_weight, - fc2_bias_tensor, - fwd_args, - ) + if use_compiled_op: + out, ln_out, new_fc1_ws, new_fc2_ws = _layernorm_mlp_op(fwd_args) else: - out, ln_out, new_fc1_ws, new_fc2_ws = _LayerNormMLP.forward( - None, - inp, - self.layer_norm_weight, - self.layer_norm_bias, - fc1_weight, - fc1_bias, - fc2_weight, - fc2_bias_tensor, + out, ln_out, new_fc1_ws, new_fc2_ws = _layernorm_mlp_eager( + ( + inp, + self.layer_norm_weight, + self.layer_norm_bias, + fc1_weight, + fc1_bias, + fc2_weight, + fc2_bias_tensor, + ), fwd_args, + is_grad_enabled, ) if new_fc1_ws is not None and cache_name_fc1 is not None: @@ -2884,6 +3378,44 @@ def forward( return out, ln_out return out + def _compile_eager_fallback_reason( + self, + inp: torch.Tensor, + skip_fp8_weight_update: Optional[torch.Tensor], + fp8_output: bool, + is_grad_enabled: bool, + debug: bool, + quantizers: Sequence[Optional[Quantizer]], + ) -> Optional[str]: + """Return a fallback reason after preparing the forward quantizers.""" + if debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + if self.checkpoint and is_grad_enabled: + return "activation checkpointing (checkpoint=True)" + if isinstance(inp, (QuantizedTensor, QuantizedTensorStorage)): + return "a quantized input tensor" + if self.fsdp_group is not None: + return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" + if ( + fp8_output + and is_grad_enabled + and (inp.requires_grad or any(p.requires_grad for p in self.parameters())) + ): + return "differentiable fp8_output=True" + if is_cpu_offload_enabled(): + return "CPU activation offloading" + # A queued store can be enabled between forward and backward. + if self.wgrad_store is not None and ( + self.wgrad_store.context is not None or self.wgrad_store.delay_wgrad_compute() + ): + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + fp8 = FP8GlobalStateManager.is_fp8_enabled() + if fp8 and skip_fp8_weight_update is not None: + return "GPU-controlled FP8 weight cache updates" + return compile_unsupported_quantizer_reason(quantizers) + def _get_quantizers(self, fp8_output, is_grad_enabled): if self.fp8: self._warn_missing_output_quantizer_role(fp8_output, False) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 082e88a2c76..8cd4cf0d1e3 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -5,7 +5,7 @@ """Linear API""" from dataclasses import dataclass, replace as dataclass_replace -from typing import Any, Callable, Dict, Optional, Tuple, Union, List +from typing import Sequence, Any, Callable, Dict, Optional, Tuple, Union, List import warnings import weakref @@ -32,8 +32,11 @@ _2X_ACC_WGRAD, ) from ._common import ( + compile_unsupported_quantizer_reason, can_reconstruct_wgrad_input_from_original, check_fp8_reduce_and_update, + fake_workspace_valid, + get_output_first_dim_size, noop_cat, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, @@ -90,7 +93,6 @@ TensorSpec, TensorOrQuantized, register_custom_op_with_autograd, - is_value_opaque_quantizer, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -182,54 +184,6 @@ class LinearFwdArgs: cpu_offloading: bool is_grad_enabled: bool - def compile_unsupported_reason(self) -> Optional[str]: - """Reason this config can't use the torch.compile custom-op path (else None).""" - if self.debug: - return "debug instrumentation (nvidia-dlfw-inspect)" - if is_distributed_weight(self.weight): - return "a DistributedWeight (custom weight parallelism, e.g. GTP)" - if isinstance(self.inp, (QuantizedTensor, QuantizedTensorStorage)): - return "a quantized input tensor" - if self.fsdp_group is not None: - return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" - if ( - self.fp8_output - and self.is_grad_enabled - and (self.input_requires_grad or self.weight_requires_grad or self.bias_requires_grad) - ): - return "differentiable fp8_output=True" - if self.cpu_offloading: - return "CPU activation offloading" - if self.wgrad_store is not None: - # Non-None only when delayed wgrad compute is on (see Linear.forward). - return "delayed wgrad compute (wgrad_store)" - if ( - self.grad_input_quantizer is not None - and self.is_grad_enabled - and self.input_requires_grad - and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad) - ): - # A quantized dgrad can't cross the op boundary: grads are packed - # one plain Tensor[] slot each (_pack_bwd_result). - return "a quantized input grad (fp8_grad=True)" - if self.fp8 and self.skip_fp8_weight_update is not None: - return "GPU-controlled FP8 weight cache updates" - if self.fuse_wgrad_accumulation: - return "fuse_wgrad_accumulation (main_grad)" - for quantizer in ( - self.input_quantizer, - self.weight_quantizer, - self.output_quantizer, - self.grad_input_quantizer, - self.grad_weight_quantizer, - self.grad_output_quantizer, - ): - # e.g. delayed-scaling Float8Quantizer and unregistered custom-recipe - # quantizers are not value-opaque and can't cross the custom-op boundary. - if quantizer is not None and not is_value_opaque_quantizer(quantizer): - return "a quantizer not registered as a torch.compile value-opaque type" - return None - @dataclass(slots=True) class LinearBwdArgs: @@ -304,6 +258,10 @@ class LinearBwdArgs: # --- Per-backward scratch state (populated inside _linear_backward_impl) --- ub_obj_gradout: Optional[Any] = None + def setup_grad_outputs(self, grads: Sequence[Any]) -> None: + """Unpack gradients in forward-output order, ignoring weight workspaces.""" + self.grad_output, _ = grads + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: """Pull saved tensors from ``ctx`` into the fields backward consumes.""" ( @@ -316,40 +274,6 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: ) # pylint: disable=unbalanced-tuple-unpacking -def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: - """Output's leading (sequence) dim from the input's: sequence parallelism - gathers it (column-parallel) or scatters it (row-parallel).""" - if not args.sequence_parallel: - return leading - if args.parallel_mode == "column": - return leading * args.tp_size - if args.parallel_mode == "row": - return leading // args.tp_size - return leading - - -def _inp_leading_from_out(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: - """Inverse of :func:`_out_leading_from_inp`.""" - if not args.sequence_parallel: - return leading - if args.parallel_mode == "column": - return leading // args.tp_size - if args.parallel_mode == "row": - return leading * args.tp_size - return leading - - -def _fake_workspace_valid(workspace: TensorSpec, quantizer: Optional[Quantizer]) -> bool: - """Spec-level mirror of ``_is_weight_workspace_valid``: the cached workspace - must already hold every inner buffer the quantizer's current usage needs.""" - if quantizer is None: - return True - required = TensorSpec( - shape=workspace.shape, dtype=workspace.dtype, quantizer=quantizer, device=workspace.device - ).inner_names() - return set(required) <= set(workspace.inner_names()) - - def _linear_forward_impl( args: LinearFwdArgs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]: @@ -862,7 +786,7 @@ def _linear_forward_fake( else: weightmat_is_storage = True workspace = args.weight_workspace - if workspace is not None and not _fake_workspace_valid(workspace, weight_quantizer): + if workspace is not None and not fake_workspace_valid(workspace, weight_quantizer): # quantize_weight drops a stale workspace and builds a new one. workspace = None if workspace is not None: @@ -894,7 +818,7 @@ def _linear_forward_fake( # Output tensor: y = x @ w^T (quantized iff an output quantizer is set). # ------------------------------------------------------ inp_leading = inp.shape[0] if len(inp.shape) > 1 else 1 - out_leading = _out_leading_from_inp(inp_leading, args) + out_leading = get_output_first_dim_size(inp_leading, args) out = TensorSpec( shape=( (out_features,) @@ -1021,10 +945,7 @@ def _linear_setup_ctx( bwd_args.use_bias = bias is not None bwd_args.requires_dgrad = fwd_args.input_requires_grad bwd_args.requires_wgrad = fwd_args.weight_requires_grad - # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True) - # inp.shape contains SymInt dims which are not hashable in OpaqueValueBundle. - # The backward reconstructs inp_shape from grad_output + weight + SP config. - bwd_args.inp_shape = None + bwd_args.inp_shape = inp.shape # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype @@ -1176,16 +1097,6 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). - if bwd_args.inp_shape is None: - in_features = saved_weight.shape[-1] - inp_leading = _inp_leading_from_out(grad_output.shape[0], bwd_args) - bwd_args.inp_shape = ( - torch.Size([in_features]) - if grad_output.ndim == 1 - else torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) - ) - # Configure Userbuffers communication (comm+GEMM overlap) bwd_args.ub_obj_gradout = None ub_obj_dgrad = None @@ -1769,20 +1680,13 @@ def _linear_backward_fake( dgrad = None if args.requires_dgrad: - # Input shape rederived from grad_output + SP config (inp_shape is not - # stored: torch.Size with SymInt cannot cross in OpaqueValueBundle). - dgrad_leading = _inp_leading_from_out(args.grad_output.shape[0], args) # Under UB reduce-scatter or bulk-wgrad overlap the returned dgrad is a # plain tensor; the quantizer only feeds the comm buffer. dgrad_quantizer = ( None if (args.ub_overlap_rs_dgrad or args.ub_bulk_wgrad) else args.grad_input_quantizer ) dgrad = TensorSpec( - shape=( - (in_features,) - if len(args.grad_output.shape) == 1 - else (dgrad_leading, *args.grad_output.shape[1:-1], in_features) - ), + shape=args.inp_shape, dtype=out_dtype, quantizer=dgrad_quantizer, device=args.grad_output.device, @@ -1890,7 +1794,7 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: """Backward pass: compute gradients and reduce FP8 scaling factors.""" bwd_args: LinearBwdArgs = ctx.backward_objects - bwd_args.grad_output = grad_output + bwd_args.setup_grad_outputs((grad_output, _grad_weight_workspace)) bwd_args.setup_saved_tensors(ctx) nvtx_label = "transformer_engine._Linear.backward" if bwd_args.ub_name is not None: @@ -2418,16 +2322,6 @@ def forward( if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()): fp8_grad = True - if torch.compiler.is_compiling() and _linear_op is not None: - reason = self._compile_eager_fallback_reason( - inp, skip_fp8_weight_update, fp8_output, fp8_grad, is_grad_enabled, debug - ) - if reason is not None: - # A break inside the try/finally below would skip the whole frame. - warn_compile_eager_fallback(reason) - torch._dynamo.graph_break(msg=f"te.Linear falling back to eager: {reason}") - return self._forward_eager_fallback(inp, is_first_microbatch, fp8_output, fp8_grad) - inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor)) try: weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() @@ -2570,8 +2464,15 @@ def forward( ) if use_compiled_op: - # Safety net for quantizer-dependent conditions only. - fallback_reason = fwd_args.compile_unsupported_reason() + fallback_reason = self._compile_eager_fallback_reason( + inp, + skip_fp8_weight_update, + fp8_output, + fp8_grad, + is_grad_enabled, + debug, + quantizers, + ) if fallback_reason is not None: warn_compile_eager_fallback(fallback_reason) torch._dynamo.graph_break( @@ -2670,9 +2571,9 @@ def _compile_eager_fallback_reason( fp8_grad: bool, is_grad_enabled: bool, debug: bool, + quantizers: Sequence[Optional[Quantizer]], ) -> Optional[str]: - """Why this call can't use the compiled op (else None), decided before - prepare_forward. Quantizer checks stay in compile_unsupported_reason.""" + """Return a fallback reason after preparing the forward quantizers.""" if debug: return "debug instrumentation (nvidia-dlfw-inspect)" weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() @@ -2706,24 +2607,7 @@ def _compile_eager_fallback_reason( return "a quantized input grad (fp8_grad=True)" if fp8 and skip_fp8_weight_update is not None: return "GPU-controlled FP8 weight cache updates" - return None - - @torch._dynamo.disable - def _forward_eager_fallback( - self, - inp: torch.Tensor, - is_first_microbatch: Optional[bool], - fp8_output: bool, - fp8_grad: bool, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: - """Re-run forward outside Dynamo (unsupported-config fallback).""" - return Linear.forward( - self, - inp, - is_first_microbatch=is_first_microbatch, - fp8_output=fp8_output, - fp8_grad=fp8_grad, - ) + return compile_unsupported_quantizer_reason(quantizers) def _get_weight_and_bias_tensors(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: # Get concatenated weight and bias tensors From a0a7e920ccf71266946cbfd8bf823ee2c78c2550 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Thu, 29 Oct 2026 19:08:29 +0100 Subject: [PATCH 21/32] [Docs] Add Mixture of Experts guide (#3494) * [Docs] Add MoE feature overview Signed-off-by: Pawel Gadzinski * [Docs] Refine MoE overview, add supporting snippets/figures and JAX API entries - Add code snippets and SVG figures referenced by mixture_of_experts.rst (moe_permute / moe_unpermute / grouped_linear tabbed examples for both PyTorch and JAX) - Add JAX API reference entries for token_dispatch, token_combine and grouped_dense so the cross-references from the MoE page resolve - Make wording framework-neutral where it was PyTorch-only (Grouped GEMM instead of GroupedLinear/grouped linear in shared sections, both m_splits and group_sizes mentioned, figure labels generalized) - Tighten routing-kernel intro: consolidate the redundant "multiple variants exist / see API ref" notes into one paragraph next to the example, and explicitly state that the kernels are differentiable - Sharpen merging_probs explanation (top-1 vs top-k) and explicitly describe what token_dispatch / token_combine return - Snippet cleanups: define previously undefined symbols, drop the JAX probs= argument from the basic example and explain its purpose in a comment, document the ignored permuted_probs / pad_offsets outputs - Reorder MoE entry in the docs/index.rst toctree Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [Docs] Extend MoE overview: router, end-to-end layer, padding, low-precision Add Router (score function + top-k + load-balancing loss) and Putting-it-together sections, plus token-probabilities / padding-and-alignment / chunk-sort subsections and a fused-expert-MLP note. New SVG figures and PyTorch/JAX snippets. Add the router and moe_permute_and_pad_with_probs API reference entries (PyTorch and JAX) and sort_chunks_by_index (JAX). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski * [Docs] MoE: fused grouped MLP and expert parallelism sections Signed-off-by: Pawel Gadzinski * [Docs] MoE: split into per-topic pages, match current docs style, document EP APIs Signed-off-by: Pawel Gadzinski * [Docs] MoE: single article with sections, redraw expert parallelism diagram Signed-off-by: Pawel Gadzinski * [Docs] MoE: warn about uneven low-precision support Signed-off-by: Pawel Gadzinski * [Docs] MoE: shorten low-precision warning Signed-off-by: Pawel Gadzinski * [Docs] MoE: reword low-precision note Signed-off-by: Pawel Gadzinski * [Docs] MoE: drop API-reference pointer from low-precision note Signed-off-by: Pawel Gadzinski * [Docs] MoE: reorganize into introduction, routing kernels, grouped GEMM, grouped MLP, EP Signed-off-by: Pawel Gadzinski * [Docs] MoE: move end-to-end example to the end of the article Signed-off-by: Pawel Gadzinski * [Docs] MoE: add expert-parallel layer diagram to the end-to-end example Signed-off-by: Pawel Gadzinski * [Docs] MoE: use expert-parallel layer diagram as Figure 1, mention EP in introduction Signed-off-by: Pawel Gadzinski * [Docs] MoE: fold EP into the stage list, surface load balancing in the introduction Signed-off-by: Pawel Gadzinski * [Docs] MoE: explain grouped routing, simplify router wording Signed-off-by: Pawel Gadzinski * [Docs] MoE: deduplicate expert_bias description Signed-off-by: Pawel Gadzinski * [Docs] MoE: move routing-kernels lead paragraphs above the router subsection Signed-off-by: Pawel Gadzinski * [Docs] MoE: describe what the blocks do and how to use them; move framework specifics to snippets Signed-off-by: Pawel Gadzinski * [Docs] MoE: describe grouped GEMM execution paths; shorten grouped MLP conditions Signed-off-by: Pawel Gadzinski * [Docs] MoE: move grouped GEMM paths above snippet; simplify grouped MLP figure Signed-off-by: Pawel Gadzinski * [Docs] MoE: EP dispatch/combine replace token dispatch/combine; redraw layer figure Signed-off-by: Pawel Gadzinski * [Docs] MoE: split routing kernels into Router and Token permutation sections Signed-off-by: Pawel Gadzinski * [Docs] MoE: fix EP figure dispatch arrows; animate dispatch and combine Signed-off-by: Pawel Gadzinski * [Docs] MoE: drop EP figure animation Signed-off-by: Pawel Gadzinski * [Docs] MoE: restructure expert parallelism (overview, common API, per-framework API) Signed-off-by: Pawel Gadzinski * [Docs] MoE: EP back in tabs without C API; describe sort_chunks_by_index Signed-off-by: Pawel Gadzinski * [Docs] MoE: EP placement figure matches flow figure; clarify capacity, MXFP8 dispatch, shared experts Signed-off-by: Pawel Gadzinski * [Docs] MoE: break long paragraphs into lists Signed-off-by: Pawel Gadzinski * [Docs] MoE: explain NCCL EP and zero-copy mode; link examples Signed-off-by: Pawel Gadzinski * [Docs] MoE: EP wording (optimized ops, worst-case capacity, MXFP8 status) Signed-off-by: Pawel Gadzinski * [Docs] MoE: describe EP buffer allocation and data flow (PyTorch) Signed-off-by: Pawel Gadzinski * [Docs] MoE: clarify what EpBuffer holds and when to reuse it Signed-off-by: Pawel Gadzinski * [Docs] MoE: spell out the JAX EP primitives Signed-off-by: Pawel Gadzinski * [Docs] MoE: precise wording for MXFP8 dispatch output Signed-off-by: Pawel Gadzinski * [Docs] MoE: restructure EP overview along one step (buffer, payload, transfer path) Signed-off-by: Pawel Gadzinski * [Docs] MoE: EP overview headings and terser bullets Signed-off-by: Pawel Gadzinski * [Docs] MoE figures: larger fonts, fix label anchors and clipping, drop placement footer Signed-off-by: Pawel Gadzinski * [Docs] MoE: EP overview - one-step permutation, buffer sizing modes, zero-copy vs CUDA graphs Signed-off-by: Pawel Gadzinski * [Docs] MoE: drop zero-copy block from EP overview Signed-off-by: Pawel Gadzinski * [Docs] MoE: shorten EpBuffer bullet Signed-off-by: Pawel Gadzinski * [Docs] MoE: merge PyTorch EP API and data-flow lists Signed-off-by: Pawel Gadzinski * [Docs] MoE: describe EP backward (PyTorch) Signed-off-by: Pawel Gadzinski * [Docs] MoE: one-step dispatch figure; capacity sizing guidance Signed-off-by: Pawel Gadzinski * [Docs] MoE: single-device example before expert parallelism Signed-off-by: Pawel Gadzinski * [Docs] MoE: receive buffer block as nested bullets with overflow behaviour Signed-off-by: Pawel Gadzinski * [Docs] MoE: drop EP backward bullet Signed-off-by: Pawel Gadzinski * [Docs] MoE: larger row labels in layer figure Signed-off-by: Pawel Gadzinski * [Docs] Fix MoE routing guidance and examples Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [Docs] Fix RST in router and moe docstrings so the API pages build with -W Signed-off-by: Pawel Gadzinski * [Docs] MoE: mask padded receive slots before combine in the PyTorch EP snippet Signed-off-by: Pawel Gadzinski * [Docs] Fix MoE EP buffer handling and grouped MLP guidance Signed-off-by: Pawel Gadzinski * [Docs] Update MoE EP operations and fix guide examples Document MoeDispatch and MoeCombine composition with internal channels. Preserve actual routing decisions for auxiliary loss, use the fusion-compatible GLU layout, and distinguish PyTorch and JAX EP behavior. Signed-off-by: Pawel Gadzinski * [Docs] Link grouped MLP fusion requirements to TE v2.19 Signed-off-by: Pawel Gadzinski * Document compact top-k routing indices Explain the PyTorch topk_indices output buffer, the unchanged probability tensor shape, and frontend support in the MoE router guide. Signed-off-by: Pawel Gadzinski * Clarify top-k index buffer requirements Document the CUDA device, contiguous layout and supported integer dtypes, show an allocation example, and clarify the routing map format restriction. Signed-off-by: Pawel Gadzinski * Stop documenting standalone JAX token dispatch and combine Signed-off-by: Pawel Gadzinski * Clarify that integrated JAX MoE requires multiple GPUs Signed-off-by: Pawel Gadzinski * Fix lint and documentation tooling compatibility Signed-off-by: Pawel Gadzinski * Install wheel for Python lint setup analysis Signed-off-by: Pawel Gadzinski * docs: update grouped MLP alignment and weighting examples Signed-off-by: Pawel Gadzinski * docs: use eager receive sizing in PyTorch EP example Signed-off-by: Pawel Gadzinski * docs: normalize JAX grouped linear split sizes to int32 Signed-off-by: Pawel Gadzinski * docs: use nonsquare expert weights in JAX MoE example Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor Co-authored-by: Claude Opus 4.8 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/_static/css/diagram-colors.css | 32 + docs/api/jax.rst | 24 + docs/api/pytorch.rst | 39 + .../mixture_of_experts/grouped_linear_jax.py | 31 + .../grouped_linear_pytorch.py | 43 + .../mixture_of_experts/grouped_mlp_pytorch.py | 25 + .../mixture_of_experts/img/grouped_linear.svg | 148 ++++ .../img/moe_ep_one_step.svg | 44 ++ .../img/moe_expert_parallel.svg | 76 ++ .../img/moe_expert_placement.svg | 43 + .../img/moe_grouped_mlp.svg | 40 + .../mixture_of_experts/img/moe_layer_ep.svg | 45 ++ .../mixture_of_experts/img/moe_padding.svg | 56 ++ .../mixture_of_experts/img/moe_permute.svg | 86 ++ .../mixture_of_experts/img/moe_router.svg | 103 +++ .../mixture_of_experts/img/moe_unpermute.svg | 81 ++ .../mixture_of_experts/mixture_of_experts.rst | 733 ++++++++++++++++++ .../moe_expert_parallel_jax.py | 71 ++ .../moe_expert_parallel_pytorch.py | 77 ++ .../mixture_of_experts/moe_layer_pytorch.py | 67 ++ .../moe_permute_pad_pytorch.py | 43 + .../mixture_of_experts/moe_permute_pytorch.py | 25 + .../moe_unpermute_pytorch.py | 26 + .../features/mixture_of_experts/router_jax.py | 44 ++ .../mixture_of_experts/router_pytorch.py | 55 ++ docs/index.rst | 1 + transformer_engine/jax/moe.py | 12 +- transformer_engine/jax/router.py | 15 +- transformer_engine/pytorch/router.py | 9 +- 29 files changed, 2079 insertions(+), 15 deletions(-) create mode 100644 docs/features/mixture_of_experts/grouped_linear_jax.py create mode 100644 docs/features/mixture_of_experts/grouped_linear_pytorch.py create mode 100644 docs/features/mixture_of_experts/grouped_mlp_pytorch.py create mode 100644 docs/features/mixture_of_experts/img/grouped_linear.svg create mode 100644 docs/features/mixture_of_experts/img/moe_ep_one_step.svg create mode 100644 docs/features/mixture_of_experts/img/moe_expert_parallel.svg create mode 100644 docs/features/mixture_of_experts/img/moe_expert_placement.svg create mode 100644 docs/features/mixture_of_experts/img/moe_grouped_mlp.svg create mode 100644 docs/features/mixture_of_experts/img/moe_layer_ep.svg create mode 100644 docs/features/mixture_of_experts/img/moe_padding.svg create mode 100644 docs/features/mixture_of_experts/img/moe_permute.svg create mode 100644 docs/features/mixture_of_experts/img/moe_router.svg create mode 100644 docs/features/mixture_of_experts/img/moe_unpermute.svg create mode 100644 docs/features/mixture_of_experts/mixture_of_experts.rst create mode 100644 docs/features/mixture_of_experts/moe_expert_parallel_jax.py create mode 100644 docs/features/mixture_of_experts/moe_expert_parallel_pytorch.py create mode 100644 docs/features/mixture_of_experts/moe_layer_pytorch.py create mode 100644 docs/features/mixture_of_experts/moe_permute_pad_pytorch.py create mode 100644 docs/features/mixture_of_experts/moe_permute_pytorch.py create mode 100644 docs/features/mixture_of_experts/moe_unpermute_pytorch.py create mode 100644 docs/features/mixture_of_experts/router_jax.py create mode 100644 docs/features/mixture_of_experts/router_pytorch.py diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index 9ee5827bd1b..f1e90a09b87 100644 --- a/docs/_static/css/diagram-colors.css +++ b/docs/_static/css/diagram-colors.css @@ -279,6 +279,38 @@ html[data-theme="dark"] .subtitle, html[data-theme="dark"] .memory-label { fill: #e0e0e0; } html[data-theme="dark"] .connector { stroke: #bdbdbd; } +/* mixture_of_experts diagrams */ +html[data-theme="dark"] .small-label, +html[data-theme="dark"] .tiny-label, +html[data-theme="dark"] .row-label, +html[data-theme="dark"] .rank-label { fill: #bdbdbd; } +html[data-theme="dark"] .grid-title, +html[data-theme="dark"] .brace-label, +html[data-theme="dark"] .row-label-l, +html[data-theme="dark"] .row-label-r, +html[data-theme="dark"] .mono { fill: #e0e0e0; } +html[data-theme="dark"] .expert-e0 { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .expert-e1 { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .expert-e2 { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .expert-e3 { fill: #5c1f38; stroke: #f48fb1; } +html[data-theme="dark"] .stage, +html[data-theme="dark"] .grouped, +html[data-theme="dark"] .a2a, +html[data-theme="dark"] .fused-mlp { fill: #3a2f5c; stroke: #b39ddb; } +html[data-theme="dark"] .moe-op { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .act, +html[data-theme="dark"] .sel { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .linear, +html[data-theme="dark"] .param, +html[data-theme="dark"] .hbm, +html[data-theme="dark"] .cell, +html[data-theme="dark"] .box { fill: #2b2b2b; stroke: #9e9e9e; } +html[data-theme="dark"] .unsel, +html[data-theme="dark"] .pad { fill: #1f1f1f; stroke: #616161; } +html[data-theme="dark"] .soft-arrow, +html[data-theme="dark"] .brace { stroke: #bdbdbd; } +html[data-theme="dark"] .skip, +html[data-theme="dark"] .farrow { stroke: #b39ddb; } /* fine_grained_quantization diagrams */ html[data-theme="dark"] .fmt-mxfp8 { fill: #10375c; stroke: #64b5f6; } diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 7a31c9d3797..99282be422c 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -59,3 +59,27 @@ Modules :members: __call__ .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules + + +Mixture of Experts +------------------ +Router, routing kernels, and grouped dense for MoE layers. See +:doc:`Mixture of Experts <../features/mixture_of_experts/mixture_of_experts>` for an overview. + +.. autoapifunction:: transformer_engine.jax.router.fused_topk_with_score_function + +.. autoapifunction:: transformer_engine.jax.router.fused_moe_aux_loss + +.. autoapifunction:: transformer_engine.jax.permutation.sort_chunks_by_index + +.. autoapifunction:: transformer_engine.jax.dense.grouped_dense + +.. autoapifunction:: transformer_engine.jax.moe.moe + +.. autoapifunction:: transformer_engine.jax.ep.ep_bootstrap + +.. autoapifunction:: transformer_engine.jax.ep.ep_finalize + +.. autoapifunction:: transformer_engine.jax.ep.ep_dispatch + +.. autoapifunction:: transformer_engine.jax.ep.ep_combine diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 31d469cb77f..6a705e57902 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -99,16 +99,55 @@ Recipe availability Mixture of Experts (MoE) functions ---------------------------------- +See :doc:`Mixture of Experts <../features/mixture_of_experts/mixture_of_experts>` for an overview of +how these functions fit together. ``GroupedLinear`` (the grouped GEMM used for the +expert layers) is documented above with the other modules. + .. autoapifunction:: transformer_engine.pytorch.moe_permute .. autoapifunction:: transformer_engine.pytorch.moe_permute_with_probs +.. autoapifunction:: transformer_engine.pytorch.moe_permute_and_pad_with_probs + .. autoapifunction:: transformer_engine.pytorch.moe_unpermute .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index_with_probs +.. autoapifunction:: transformer_engine.pytorch.router.fused_topk_with_score_function + +.. autoapifunction:: transformer_engine.pytorch.router.fused_compute_score_for_moe_aux_loss + +.. autoapifunction:: transformer_engine.pytorch.router.fused_moe_aux_loss + +Expert parallelism +~~~~~~~~~~~~~~~~~~ + +NCCL-based dispatch and combine for experts sharded across ranks. See +:doc:`Expert parallelism <../features/mixture_of_experts/mixture_of_experts>`. + +``MoeDispatch`` and ``MoeCombine`` can wrap the local expert operations in a +``transformer_engine.pytorch.ops.Sequential``. They share an ``EpConfig`` and +an ``EpBuffer``; the guide shows how to connect the per-expert counts and +routing weights through the sequence's extra-tensor channels. + +.. autoapiclass:: transformer_engine.pytorch.ops.MoeDispatch + +.. autoapiclass:: transformer_engine.pytorch.ops.MoeCombine + +.. autoapiclass:: transformer_engine.pytorch.ep.EpConfig + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_bootstrap + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_finalize + +.. autoapiclass:: transformer_engine.pytorch.ep.EpBuffer + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_dispatch + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_combine + Communication-computation overlap --------------------------------- diff --git a/docs/features/mixture_of_experts/grouped_linear_jax.py b/docs/features/mixture_of_experts/grouped_linear_jax.py new file mode 100644 index 00000000000..6ed773d1649 --- /dev/null +++ b/docs/features/mixture_of_experts/grouped_linear_jax.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_GROUPED_LINEAR_JAX +import jax.numpy as jnp +from transformer_engine.jax import dense as te_dense + +# x: [sum(group_sizes), hidden_size], expert-contiguous tokens +# kernel: [num_experts, hidden_size, ffn_hidden_size], stacked per-expert weights +# bias: [num_experts, ffn_hidden_size], stacked per-expert biases +# group_sizes: [num_experts] int32 array; group_sizes[i] is the number of routed +# tokens for expert i +group_sizes = jnp.asarray(group_sizes, dtype=jnp.int32) +split_indices = jnp.cumsum(group_sizes)[:-1] +x_by_expert = jnp.split(x, split_indices, axis=0) + +# Baseline: one matmul per expert. +loop_out = jnp.concatenate( + [x_i @ kernel_i + bias_i for x_i, kernel_i, bias_i in zip(x_by_expert, kernel, bias)], + axis=0, +) + +# Transformer Engine: one grouped dense call. group_sizes is a device array. +grouped_out = te_dense.grouped_dense( + x, + kernel, + group_sizes=group_sizes, + bias=bias, +) +# END_GROUPED_LINEAR_JAX diff --git a/docs/features/mixture_of_experts/grouped_linear_pytorch.py b/docs/features/mixture_of_experts/grouped_linear_pytorch.py new file mode 100644 index 00000000000..313098ac9fd --- /dev/null +++ b/docs/features/mixture_of_experts/grouped_linear_pytorch.py @@ -0,0 +1,43 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_GROUPED_LINEAR_PYTORCH +import torch +import transformer_engine.pytorch as te + +# x: [sum(m_splits), hidden_size], expert-contiguous tokens +# m_splits: list[int] of length num_experts; m_splits[i] is the number +# of routed tokens for expert i +# torch_experts: list[torch.nn.Linear] of length num_experts, one per expert +# (used only by the baseline loop below) +x_by_expert = torch.split(x, m_splits, dim=0) + +# Baseline: one Linear call per expert. +loop_out = torch.cat( + [expert(x_i) for expert, x_i in zip(torch_experts, x_by_expert)], + dim=0, +) + +# Transformer Engine: one grouped linear call. +grouped_linear = te.GroupedLinear( + num_experts, + hidden_size, + ffn_hidden_size, + bias=True, + params_dtype=torch.bfloat16, +).cuda() +grouped_out = grouped_linear(x, m_splits) + +# Keep token counts on the device with use_grouped_tensor=True. +grouped_linear = te.GroupedLinear( + num_experts, + hidden_size, + ffn_hidden_size, + bias=True, + params_dtype=torch.bfloat16, + use_grouped_tensor=True, +).cuda() +m_splits_dev = torch.tensor(m_splits, dtype=torch.int64, device="cuda") +grouped_out = grouped_linear(x, m_splits_dev) +# END_GROUPED_LINEAR_PYTORCH diff --git a/docs/features/mixture_of_experts/grouped_mlp_pytorch.py b/docs/features/mixture_of_experts/grouped_mlp_pytorch.py new file mode 100644 index 00000000000..614631e9c0a --- /dev/null +++ b/docs/features/mixture_of_experts/grouped_mlp_pytorch.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_GROUPED_MLP_PYTORCH +import transformer_engine.pytorch as te + +# FC1 produces gate and value features interleaved in blocks of 32. +expert_mlp = te.ops.Sequential( + te.ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size, split_alignment=256), + te.ops.ScaledSwiGLU(glu_interleave_size=32), + te.ops.GroupedLinear( + num_experts, ffn_hidden_size, hidden_size, scale_bias=True, split_alignment=256 + ), +) + +# Inputs from padded dispatch with align_size=256. +expert_out = expert_mlp( + padded, + padded_tokens_per_expert, + permuted_probs, + padded_tokens_per_expert, + permuted_probs, +) +# END_GROUPED_MLP_PYTORCH diff --git a/docs/features/mixture_of_experts/img/grouped_linear.svg b/docs/features/mixture_of_experts/img/grouped_linear.svg new file mode 100644 index 00000000000..686593a9b40 --- /dev/null +++ b/docs/features/mixture_of_experts/img/grouped_linear.svg @@ -0,0 +1,148 @@ + + + + + + + + + + + + + + Loop over experts + + + + X0 + + W0 + + b0 + + + + Expert 0 (Linear) + + + Y0 + + + + + X1 + + W1 + + b1 + + + + Expert 1 (Linear) + + + Y1 + + + + + X2 + + W2 + + b2 + + + + Expert 2 (Linear) + + + Y2 + + + + + + + Grouped GEMM + + + + + W0 + + b0 + + W1 + + b1 + + W2 + + b2 + + + + + + + + + + + X0 + X1 + X2 + + + tokens for expert 0 + + tokens for expert 1 + + tokens for expert 2 + + + + + + + + Grouped GEMM + multiple experts + + + + tokens per expert + + + + + + + + + + + Y0 + Y1 + Y2 + + + diff --git a/docs/features/mixture_of_experts/img/moe_ep_one_step.svg b/docs/features/mixture_of_experts/img/moe_ep_one_step.svg new file mode 100644 index 00000000000..6f5c6f8f84a --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_ep_one_step.svg @@ -0,0 +1,44 @@ + + + + + + + + + generic all-to-all + + tokens + + + permute + sort by rank + + + all-to-all + + + permute + group by expert + + + experts + + NCCL EP dispatch + + tokens + + + dispatch + token → expert slot + + + experts + diff --git a/docs/features/mixture_of_experts/img/moe_expert_parallel.svg b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg new file mode 100644 index 00000000000..d82d4e30714 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg @@ -0,0 +1,76 @@ + + + + + + + Expert parallelism: experts E0, E1 on rank 0 and E2, E3 on rank 1 + + tokens + dispatch all-to-all + receive buffer + local experts + combine all-to-all + output + + + Rank 0 + + Rank 1 + + + t0 → E0 + t1 → E2 + t2 → E0 + t3 → E1 + t4 → E3 + t5 → E2 + + + + + + + + + E0: t0 + E0: t2 + E1: t3 + E2: t1 + E2: t5 + E3: t4 + + + + Grouped MLPE0, E1 + + Grouped MLPE2, E3 + + + + + + + + + y0 + y1 + y2 + y3 + y4 + y5 + diff --git a/docs/features/mixture_of_experts/img/moe_expert_placement.svg b/docs/features/mixture_of_experts/img/moe_expert_placement.svg new file mode 100644 index 00000000000..6080f120859 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_expert_placement.svg @@ -0,0 +1,43 @@ + + + + + + Rank 0 + tokens + + t0 → E0 + + t1 → E2 + + t2 → E0 + local experts + + E0 + + E1 + + Rank 1 + tokens + + t3 → E1 + + t4 → E3 + + t5 → E2 + local experts + + E2 + + E3 + diff --git a/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg b/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg new file mode 100644 index 00000000000..0f3840d6543 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg @@ -0,0 +1,40 @@ + + + + + + + + + FC1 grouped GEMM + + + + activation + + + + FC2 grouped GEMM + + + operation fuser + + + fused grouped-MLP operation + + FC1 grouped GEMM + + activation + + intermediate + device memory + + FC2 grouped GEMM + diff --git a/docs/features/mixture_of_experts/img/moe_layer_ep.svg b/docs/features/mixture_of_experts/img/moe_layer_ep.svg new file mode 100644 index 00000000000..ca4cda5546b --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_layer_ep.svg @@ -0,0 +1,45 @@ + + + + + + + + + + single device + + Router + + + Token dispatch + + + Grouped MLP + + + Token combine + + + expert parallelism (experts sharded across ranks) + + Router + + + All-to-all dispatch + grouped by local expert + + + Grouped MLP + local experts + + + All-to-all combine + back in token order + diff --git a/docs/features/mixture_of_experts/img/moe_padding.svg b/docs/features/mixture_of_experts/img/moe_padding.svg new file mode 100644 index 00000000000..d4ddddebb31 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_padding.svg @@ -0,0 +1,56 @@ + + + + + + + + + Fused padding to expert-aligned blocks + + + permuted + expert 0 + + + + expert 1 + + expert 2 + + + tokens_per_expert = [3, 1, 2] + + + + token dispatch + align_size = 2 + + + padded + expert 0 + + + + pad + expert 1 + + pad + expert 2 + + + padded = [4, 2, 2] + diff --git a/docs/features/mixture_of_experts/img/moe_permute.svg b/docs/features/mixture_of_experts/img/moe_permute.svg new file mode 100644 index 00000000000..99b60b884fc --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_permute.svg @@ -0,0 +1,86 @@ + + + + + + + + + Token Dispatch + + + + tokens + + t0 + + t1 + + t2 + + t3 + + t4 + + t5 + + + + + + + + Token Dispatch + group tokens by + destination expert + + + + + + + routing_map + + + + + row_id_map + + + + + permuted tokens + + t2 + + t3 + + t0 + + t4 + + t1 + + t5 + + + expert 0 + + expert 1 + + expert 2 + + diff --git a/docs/features/mixture_of_experts/img/moe_router.svg b/docs/features/mixture_of_experts/img/moe_router.svg new file mode 100644 index 00000000000..952082db105 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_router.svg @@ -0,0 +1,103 @@ + + + + + + + + + Router: scoring and top-k selection + + + logits + e0 + e1 + e2 + e3 + + t0 + 1.2 + 3.1 + 0.4 + 2.7 + + t1 + 2.9 + 0.8 + 2.2 + 1.0 + + t2 + 0.5 + 1.1 + 3.3 + 2.5 + + + + + + + score function + + top-k + + + + + + + routing_map + e0 + e1 + e2 + e3 + + 0 + 1 + 0 + 1 + + 1 + 0 + 1 + 0 + + 0 + 0 + 1 + 1 + + + probs + e0 + e1 + e2 + e3 + + 0 + .6 + 0 + .4 + + .7 + 0 + .3 + 0 + + 0 + 0 + .6 + .4 + diff --git a/docs/features/mixture_of_experts/img/moe_unpermute.svg b/docs/features/mixture_of_experts/img/moe_unpermute.svg new file mode 100644 index 00000000000..aa27a7642e1 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_unpermute.svg @@ -0,0 +1,81 @@ + + + + + + + + + Token Combine + + + + expert outputs + + y2 + + y3 + + y0 + + y4 + + y1 + + y5 + + + expert 0 + + expert 1 + + expert 2 + + + + + + + + Token Combine + restore original + token order + + + + + + + row_id_map + + + + + tokens (original order) + + y0 + + y1 + + y2 + + y3 + + y4 + + y5 + + diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst new file mode 100644 index 00000000000..f9a9110a47a --- /dev/null +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -0,0 +1,733 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _moe-overview: + +Mixture of Experts +=================================== + +.. warning:: + + MoE feature availability depends on the GPU, framework, Transformer Engine + and dependency versions, and precision recipe. Not every operation or + optimization works in every configuration. This guide explains the building + blocks; consult the :doc:`PyTorch API reference ` and + :doc:`JAX API reference ` for compatibility and input requirements. + + The integrated JAX MoE layer currently requires multiple GPUs (at least + two). A single-GPU JAX MoE example is not available through the public API. + +Introduction +------------ + +A Mixture of Experts (MoE) layer replaces a dense feed-forward network with a set +of expert networks and a router that sends each token to one or more experts. +A token passes through the layer in the following stages: + +#. The **router** scores the experts for each token and selects the top-k of + them. +#. **Token dispatch** gathers the tokens into expert-contiguous order. +#. The **grouped MLP** (the experts) runs a single batched computation over all + expert blocks. +#. **Token combine** scatters the expert outputs back into the original token + order, merging the contributions when a token was sent to more than one + expert. + +With expert parallelism the experts are sharded across ranks, and an +**all-to-all dispatch** and **all-to-all combine** take the place of token +dispatch and token combine: the dispatch takes the router output directly and +delivers each rank's tokens already grouped by local expert, and the combine +returns the outputs to the source rank in the original token order. + +.. raw:: html + :file: img/moe_layer_ep.svg + +*Figure 1. The stages of an MoE layer on a single device and with expert +parallelism.* + +Transformer Engine provides a building block for each stage. They are exposed as +standalone functions, so they can be assembled into a complete MoE layer or +dropped into an existing implementation one piece at a time: + +* :ref:`Router `: fused score function and top-k selection, and a + fused :ref:`load-balancing loss `. +* :ref:`Token permutation `: token dispatch and combine + kernels that move tokens between their original order and the + expert-contiguous layout. +* :ref:`Grouped GEMM `: the expert linear layers as one call + over expert-contiguous blocks; the :ref:`grouped MLP ` fuses + activation and supported quantization steps into the expert GEMMs. +* :ref:`Expert parallelism `: all-to-all dispatch and + combine for experts sharded across devices. + +The :ref:`single-device example ` wires the blocks +into a complete MoE layer; :ref:`Expert parallelism ` +covers the sharded case. + +.. _moe-router: + +Router +------ + +The router decides which experts each token is sent to. It applies a score +function to the gating logits, selects the top-k experts per token, and produces +the two tensors that drive the rest of the layer: + +* ``routing_map`` - by default, a ``[num_tokens, num_experts]`` mask marking + the selected experts. Token dispatch uses it to lay the tokens out by expert. +* ``probs`` - the routing weight of each selected expert. Token combine uses + these as merging weights when a token was routed to more than one expert. + +In PyTorch, passing a preallocated ``topk_indices`` buffer of shape +``[num_tokens, topk]`` to ``fused_topk_with_score_function`` selects a compact +list of expert indices instead. The buffer must be a contiguous CUDA tensor on +the same device as ``logits``, with dtype ``torch.int16``, ``torch.int32`` or +``torch.int64``. For example: + +.. code-block:: python + + topk_indices = torch.empty( + (num_tokens, topk), dtype=torch.int32, device=logits.device + ) + +Pass this buffer as ``topk_indices=topk_indices`` and leave +``routing_map_format`` at its default value. The function fills the buffer and +returns it as its second output without materializing the routing mask. The +``probs`` output still has shape ``[num_tokens, num_experts]``. The JAX router +does not currently expose this index-output option. Both frontends also support +a bit-packed mask through ``routing_map_format="bitmap_u8"``. The examples and +figure below use the default mask format. + +``fused_topk_with_score_function`` runs the score function and the top-k +selection in a single differentiable kernel. All internal math runs in FP32, +regardless of the logits dtype. + +.. raw:: html + :file: img/moe_router.svg + +*Figure 2. The router scores the experts for each token and keeps the top-k. +The selected entries populate* ``routing_map`` *(a 0/1 mask) and* ``probs`` *(the +routing weights); all other entries are zero.* + +Options: + +* **Score function:** softmax, sigmoid or sqrtsoftplus. With softmax, ``use_pre_softmax`` + selects whether the softmax is applied before or after the top-k. +* **Grouped routing:** the experts are split into ``num_groups`` equal groups. + Each group is scored by the sum of its best expert scores, the top + ``group_topk`` groups are kept, and the top-k experts are chosen only from + those groups. +* **Expert bias:** ``expert_bias`` is added to the scores before the top-k + selection (see :ref:`Load balancing `). +* **Scaling:** ``scaling_factor`` rescales the returned probabilities. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: router_pytorch.py + :language: python + :start-after: # START_ROUTER_PYTORCH + :end-before: # END_ROUTER_PYTORCH + + .. tab:: JAX + + .. literalinclude:: router_jax.py + :language: python + :start-after: # START_ROUTER_JAX + :end-before: # END_ROUTER_JAX + +.. _moe-load-balancing: + +Load balancing +~~~~~~~~~~~~~~ + +``fused_moe_aux_loss`` computes the auxiliary load-balancing loss that penalizes +uneven token counts across experts. It takes the per-expert token counts and the +*dense* routing scores (one value per expert, not only the selected top-k), so +the loss has a gradient with respect to every expert's logit. The dense scores +are returned by the router functions shown below; add the scaled loss to the +training loss. + +Count tokens from the original router's ``routing_map``. The auxiliary-score +helpers ignore expert bias and grouped routing, so their returned maps can +select different experts. Use the same score function for routing and the +auxiliary scores. + +``expert_bias`` balances the load without an extra loss term: + +* with sigmoid or sqrtsoftplus it is added to the scores only for the top-k + selection, so it changes which experts are picked but not the returned + routing weights; +* update it between steps: lower it for overloaded experts, raise it for + under-used ones. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: router_pytorch.py + :language: python + :start-after: # START_ROUTER_AUX_PYTORCH + :end-before: # END_ROUTER_AUX_PYTORCH + + .. tab:: JAX + + .. literalinclude:: router_jax.py + :language: python + :start-after: # START_ROUTER_AUX_JAX + :end-before: # END_ROUTER_AUX_JAX + +.. _moe-token-permutation: + +Token permutation +----------------- + +Token dispatch moves the tokens into the expert-contiguous layout expected by +the grouped GEMM, and token combine moves the expert outputs back. + +* The PyTorch kernels shown below are differentiable. +* The PyTorch snippets below use the mask-map routing variant. Other variants + (for example index-map routing) follow the same pattern, see the + :doc:`PyTorch API reference `. + +Token dispatch +~~~~~~~~~~~~~~ + +Token dispatch takes the token tensor and a routing map describing each +token's destination experts, and returns a permuted token buffer in which all +rows assigned to the same expert are stored contiguously. This is the layout the +grouped GEMM consumes, together with the per-expert token counts. + +.. raw:: html + :file: img/moe_permute.svg + +*Figure 3. Token dispatch consumes the input token tensor together with the +routing map and produces an expert-contiguous token tensor; rows assigned to the +same expert are stored back-to-back.* + +A typical call looks like: + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_permute_pytorch.py + :language: python + :start-after: # START_MOE_PERMUTE_PYTORCH + :end-before: # END_MOE_PERMUTE_PYTORCH + + .. tab:: JAX + + Standalone token permutation is currently unavailable through the public + Transformer Engine JAX API. + +The PyTorch call returns the permuted token buffer of shape +``[num_out_tokens, hidden_size]`` together with a ``row_id_map`` that token +combine uses to restore the original token order after the experts have run. + +Token combine +~~~~~~~~~~~~~ + +Token combine is the inverse operation: it takes the expert-contiguous output +of the grouped GEMM and the ``row_id_map`` returned by token dispatch, and +returns a tensor of shape ``[num_tokens, hidden_size]`` with the rows written +back into the original token order. + +Pass the routing weights as ``merging_probs``; the kernel then computes the +weighted sum of the per-expert contributions in the same fused pass. Without +them the contributions are summed unweighted. They may be omitted for top-1 +only when the router guarantees that every selected weight is one, for example +post-top-k softmax with ``scaling_factor=1``. Pre-softmax, sigmoid, +sqrtsoftplus, or an additional scaling factor can produce non-unit top-1 +weights that must still be applied. + +.. raw:: html + :file: img/moe_unpermute.svg + +*Figure 4. Token combine reads the expert-contiguous output tensor and the* +``row_id_map``\ *, and writes each row back to its original token slot. With* +``merging_probs``\ *, contributions from multiple experts to the same token are +combined in the same fused kernel.* + +A typical call looks like: + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_unpermute_pytorch.py + :language: python + :start-after: # START_MOE_UNPERMUTE_PYTORCH + :end-before: # END_MOE_UNPERMUTE_PYTORCH + + .. tab:: JAX + + Standalone token unpermutation is currently unavailable through the public + Transformer Engine JAX API. + +Token probabilities +~~~~~~~~~~~~~~~~~~~ + +The routing weights must be applied exactly once: + +* **In token combine.** Pass the original routing weights as ``merging_probs``, + as in the examples above. +* **Before token combine.** Pass the weights to token dispatch as ``probs`` to + receive them in expert-contiguous order, multiply the completed expert + outputs by the permuted weights, and call token combine without + ``merging_probs``. An implementation may instead fold the weights into a + mathematically equivalent point, such as immediately before the expert's + final bias-free linear layer. + +Token dispatch only permutes ``probs``; it does not apply them. In particular, +scaling the input to a nonlinear expert is not equivalent to weighting that +expert's output. + +Padding and alignment +~~~~~~~~~~~~~~~~~~~~~ + +Grouped GEMM backends require or prefer each expert's token block to start at +an aligned offset (for example, a multiple of 128 rows). Token dispatch can pad +each block up to a multiple of ``align_size`` in the same kernel. + +.. raw:: html + :file: img/moe_padding.svg + +*Figure 5. Each expert's block is rounded up to a multiple of* ``align_size``\ *. +The per-expert padding offsets are returned so that token combine can drop the +padding again.* + +The padded dispatch returns the padded token buffer, the aligned per-expert +token counts to pass to the grouped GEMM, and the per-expert ``pad_offsets`` +that token combine needs to remove the padding. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_permute_pad_pytorch.py + :language: python + :start-after: # START_MOE_PERMUTE_PAD_PYTORCH + :end-before: # END_MOE_PERMUTE_PAD_PYTORCH + + .. tab:: JAX + + Standalone token permutation with padding is currently unavailable through + the public Transformer Engine JAX API. + +Reordering expert chunks +~~~~~~~~~~~~~~~~~~~~~~~~ + +``sort_chunks_by_index`` reorders whole blocks of rows: + +* the input ``[num_tokens, hidden_size]`` is split along the first dimension + into chunks of the given ``split_sizes``; +* the chunks are concatenated again in the order given by ``sorted_indices``: + output chunk ``i`` is input chunk ``sorted_indices[i]``, rows inside a chunk + keep their order; +* the operation is differentiable, and a ``_with_probs`` variant moves a + per-row probability tensor along with the rows. + +The typical use is expert parallelism over a generic all-to-all, where the +received buffer is ordered by source rank and then by expert, while the grouped +GEMM needs all rows of one expert together. With two source ranks and two local +experts: + +* received chunks: ``(rank 0, E4)``, ``(rank 0, E5)``, ``(rank 1, E4)``, + ``(rank 1, E5)``; +* ``sorted_indices = [0, 2, 1, 3]`` regroups them into ``E4, E4, E5, E5`` for + the grouped GEMM; +* after the experts have run, the inverse permutation restores the rank-major + order for the combine all-to-all. + +.. _moe-grouped-gemm: + +Grouped GEMM +------------ + +The grouped GEMM applies the per-expert linear layers in one call, replacing a +loop of one ``Linear`` call per expert and producing the same outputs. + +Let ``G`` be the number of experts. For expert ``i``, ``X_i`` is the routed +token block, ``W_i`` is the expert weight, and ``b_i`` is the optional bias: + +.. math:: + + Y_i = X_i W_i^T + b_i,\quad i = 0, \ldots, G - 1 + +The full layer output is the concatenation of all expert outputs: + +.. math:: + + Y = \mathrm{concat}(Y_0, Y_1, \ldots, Y_{G-1}) + +The number of token rows belonging to each expert is passed as a per-expert +token-count argument. + +.. raw:: html + :file: img/grouped_linear.svg + +*Figure 6. Both paths produce the same outputs from the same inputs. The +baseline launches one* ``Linear`` *per expert; the grouped GEMM replaces the +loop with one call.* + +Grouped GEMMs can use :doc:`low-precision training recipes +`: the inputs are quantized per expert +and the expert GEMMs run in the recipe's precision. + +The computation can be implemented in two ways: + +* **Per-expert GEMMs.** The input is split by expert and each expert's matrix + multiplication is launched separately. +* **Single grouped GEMM.** The experts run in one grouped operation. Keeping + token counts on the device avoids a host synchronization and helps with + CUDA graph capture. + +The snippets assume the tokens have already been permuted into +expert-contiguous order. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: grouped_linear_pytorch.py + :language: python + :start-after: # START_GROUPED_LINEAR_PYTORCH + :end-before: # END_GROUPED_LINEAR_PYTORCH + + .. tab:: JAX + + .. literalinclude:: grouped_linear_jax.py + :language: python + :start-after: # START_GROUPED_LINEAR_JAX + :end-before: # END_GROUPED_LINEAR_JAX + + With MXFP8, pad each expert's token block to a multiple of 128 rows and + pass the padded counts in ``group_sizes``. Aligning only the total token + buffer is insufficient. The high-level JAX ``moe`` API handles this + per-expert padding during dispatch. + +.. _moe-grouped-mlp: + +Grouped MLP +----------- + +An expert MLP is two grouped GEMMs with an activation between them. Transformer +Engine can fuse the activation and quantization into the expert GEMMs. In the +fused path illustrated below, FC1 includes the activation and FC2 runs as a +separate grouped GEMM, reading the intermediate tensor from device memory. + +.. raw:: html + :file: img/moe_grouped_mlp.svg + +*Figure 7. The operation fuser replaces the two grouped GEMMs and the activation +between them with one grouped-MLP operation that launches separate FC1 and FC2 +kernels.* + +The fusion is applied by the :doc:`operation fuser `: +a grouped linear, a scaled GLU (or SReLU) activation and another grouped linear +in sequence are replaced with one fused grouped-MLP operation. +Fusion requires compatible layer dimensions, activation layout, and precision +recipe. + +Fusion eligibility may change between releases; this guide summarizes the +requirements, while the implementation contains the current checks. For example, +see the `fusion eligibility checks in Transformer Engine v2.19 +`_. + +For GLU fusion, use ``glu_interleave_size=32``: FC1 must produce alternating +blocks of 32 gate features and 32 value features. When loading a checkpoint +whose FC1 outputs store all gate features before all value features, convert +each expert's FC1 weight and bias with ``te.interleave_glu_tensor(tensor, 32)``. +``ScaledClampedQGeGLU`` uses the same layout. For ``ScaledSReLU``, FC1 instead +produces ``ffn_hidden_size`` features, without the GLU's doubled width. + +``ScaledSwiGLU`` multiplies the activation output by the routing weights before +FC2. To weight the entire expert output, FC2 must also scale its bias: + +.. math:: + + p\,(h W_2 + b_2) = (p\,h) W_2 + p\,b_2, + +where :math:`h` is the activation output and :math:`p` is the routing weight. +Set ``scale_bias=True`` on FC2 and pass the same routing weights to both the +activation and FC2. With ``bias=False`` on FC2, bias scaling is unnecessary. + +For grouped MLP fusion, set ``split_alignment=256`` on both ``GroupedLinear`` +operations. Pad each expert's input token block to a multiple of 256 and pass +the corresponding padded counts as ``m_splits`` to both operations. For example, +``moe_permute_and_pad_with_probs`` with ``align_size=256`` produces the padded +input and counts; setting ``split_alignment`` alone does not pad them. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: grouped_mlp_pytorch.py + :language: python + :start-after: # START_GROUPED_MLP_PYTORCH + :end-before: # END_GROUPED_MLP_PYTORCH + +The arguments after ``padded`` are consumed in operation order: FC1 takes the +padded per-expert token counts, the activation takes the routing weights, and +FC2 takes the counts and the same weights for its bias. ``expert_out`` is already +weighted; the subsequent combine must not apply the routing weights again. + +.. _moe-putting-it-together: + +Example: MoE layer on a single device +------------------------------------- + +The runnable PyTorch example below wires the blocks together for top-k routing +on a single supported NVIDIA GPU: route, dispatch, run the experts, combine. +To keep the example short, each expert is represented by one grouped linear +layer; a full expert MLP uses the same routing around two grouped linear layers +and an activation. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_layer_pytorch.py + :language: python + :start-after: # START_MOE_LAYER_PYTORCH + :end-before: # END_MOE_LAYER_PYTORCH + + .. tab:: JAX + + The integrated JAX MoE layer currently requires at least two GPUs. + Standalone token permutation and unpermutation are not exposed through + the public JAX API, so a single-GPU example is currently unavailable. See + :ref:`Expert parallelism `. + +The example uses dropless routing (``num_out_tokens = num_tokens * top_k``), so +the dispatch buffer is sized statically rather than from a device-to-host sync. +Every stage is differentiable, so the assembled layer trains end to end. + +.. _moe-expert-parallelism: + +Expert parallelism +------------------ + +With expert parallelism (EP) the experts are sharded across ranks: every rank +keeps its own shard of the tokens and holds only a slice of the experts. + +.. raw:: html + :file: img/moe_expert_placement.svg + +*Figure 8. Expert placement: each rank holds its token shard and a subset of the +experts. Tokens t1, t3 and t5 are routed to experts on the other rank.* + +A token routed to an expert on another rank has to travel there and back. Two +all-to-all collectives wrap the local expert computation: + +* **Dispatch** sends each token to the rank that owns its expert. It takes the + router output (expert indices and weights) directly and delivers a receive + buffer grouped by local expert. +* The local grouped GEMM runs on the receive buffer. +* **Combine** returns the results to the source rank and writes them back in the + original token order. + +No separate token dispatch or token combine is needed. Shared experts, which +every token passes through, are not part of the dispatch: they run as a regular +dense MLP on the local tokens on every rank. + +.. raw:: html + :file: img/moe_expert_parallel.svg + +*Figure 9. Dispatch routes each token to the rank owning its expert, the local +experts run on the receive buffer, and combine returns the outputs to the source +rank.* + +Transformer Engine provides dispatch and combine as ready, differentiable +operations built on the NCCL EP library (``libnccl_ep``, loaded at runtime), so +an MoE layer with expert parallelism is just router, dispatch, local experts and +combine. + +**Communication and permutation in one step** + +Over a generic all-to-all the tokens have to be permuted twice: sorted by +destination rank before the collective and regrouped by expert after it. The +NCCL EP kernels fold both permutations into the transfer: + +* dispatch computes the destination slot of every token from the top-k indices + and writes it there directly, so no permutation kernel runs before or after + the communication; +* combine does the inverse in one pass: it returns each expert output to its + source rank and sums it into the original token position. + +.. raw:: html + :file: img/moe_ep_one_step.svg + +*Figure 10. Dispatch over a generic all-to-all needs a permutation on each side; +the NCCL EP dispatch writes every token straight into its expert slot.* + +**Receive buffer** + +The number of tokens a rank receives depends on the routing. Local expert +blocks are packed consecutively, with +each block optionally padded to the configured alignment. Their offsets are +the cumulative per-expert counts; any remaining capacity forms an unused tail. + +* **Fixed capacity**, ``recv_capacity_per_rank`` (an integer): the total buffer + size is fixed, while the expert block sizes and offsets depend on the routing. + + * Sizing needs no host synchronization, allowing CUDA graph capture when + the expert operations also support it. The low-level zero-copy EP path + additionally requires persistent caller-owned receive and gradient buffers. + * Without per-expert padding, ``ep_size * max_tokens_per_rank * top_k`` is a + dropless upper bound. With alignment enabled, capacity must also include + padding for every local expert. For example, counts of 129 and 127 need + 384 rows at alignment 128, even though only 256 tokens arrive. JAX provides + ``get_moe_recv_capacity_per_rank`` to compute an aligned bound. + * With balanced routing, a rank receives about + ``max_tokens_per_rank * top_k`` assignments. A smaller capacity based on + this estimate saves memory but can overflow; include alignment padding + when sizing it. + * On overflow the dispatch fails with a device-side error. With + ``drop_on_overflow=True`` it instead drops the tokens that do not fit; they + are not sent, so their experts contribute nothing to the output, like + token dropping in capacity-limited MoE. + * ``total_recv_tokens`` reports the required receive rows before dropping, + including alignment padding; compare it with capacity after the step to + detect an overflow. + +* **Eager (PyTorch)**, no capacity given: the buffer is sized from the actual + receive count each step. + + * This costs a host synchronization per step and is not CUDA-graph + capturable; ``drop_on_overflow`` does not apply. + +JAX requires a static receive capacity. In the high-level ``moe`` API, +``recv_capacity_per_rank=None`` selects an aligned dropless worst-case bound; +it does not size the buffer from each step's actual receive count. Pass the +same static capacity to ``ep_bootstrap``. + +**Framework APIs** + +.. tabs:: + + .. tab:: PyTorch + + ``te.ops.MoeDispatch`` and ``te.ops.MoeCombine`` compose with the local + expert MLP in one ``te.ops.Sequential``. The sequence runs dispatch, + grouped linear, scaled activation, grouped linear, and combine, with + autograd through both the tokens and routing weights. + + Initialize NCCL EP once per process with ``ep_bootstrap(ep_group, ...)``. + Both operations take the same immutable ``EpConfig`` and ``EpBuffer``. + Their settings must match the bootstrap configuration, including the + process group, token limit, receive capacity, and overflow policy. + ``EpBuffer`` stores the routing handle and per-expert counts until backward + completes; use a separate buffer for each concurrently in-flight layer + call, including overlapping pipeline microbatches. + + Dispatch produces two extra tensors: the per-expert row counts (including + alignment padding) and the received routing weights. Named channels pass + the counts to both grouped linear operations and the weights to the + activation. Bind these channels before constructing or calling the + sequence. ``output_to_caller=False`` keeps the metadata internal, so the + sequence returns only the combined output. + + This assembly example assumes an initialized NCCL process group, the + current CUDA device set for each rank, and router outputs for the local + BF16 tokens. ``num_experts`` must be divisible by the EP group size, and + ``max_tokens_per_rank`` must bound each rank's input token count. The + expert MLP uses BF16 without bias or per-expert padding. + + The example uses eager sizing (``recv_capacity_per_rank=None``) so the + receive buffer matches the actual per-expert counts. This synchronizes + with the host each step and prevents CUDA graph capture. + + .. literalinclude:: moe_expert_parallel_pytorch.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH + :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH + + The public extra arguments follow operation order: dispatch consumes + ``topk_idx`` and ``topk_w``, and combine consumes ``topk_idx`` again. The + grouped linear and activation inputs come from the internal channels. + NCCL EP uses the shared buffer's routing state for combine, but the + operation still requires the routing-index argument. + + ``ScaledSwiGLU`` applies the routing weights before the bias-free FC2, so + combine only sums the already weighted expert outputs. For an FC2 with + bias, set ``scale_bias=True`` and bind its second extra input to the same + ``routing_weights`` channel, as described in :ref:`moe-grouped-mlp`. + + PyTorch dispatch can quantize tokens before sending them, reducing + communication volume; expert GEMMs can consume the quantized groups. + Communication precision is configured separately from the expert GEMMs: + set ``dispatch_fwd_quant_recipe`` and ``combine_bwd_quant_recipe`` in both + ``EpConfig`` and ``EpBuffer`` for the corresponding communication paths. + Use ``te.autocast`` to select the expert-compute recipe, and choose an + alignment and receive capacity compatible with both the communication and + expert kernels. See the configurations in + `tests/pytorch/distributed/run_ep.py `_. + + The lower-level ``ep_dispatch`` and ``ep_combine`` functions remain + available for custom expert code or explicit receive-buffer management. + ``ep_dispatch`` returns ``(recv_tokens, recv_topk_weights, + tokens_per_expert)``; unlike the operation's extra outputs, the weights + precede the counts. Weight the expert outputs before ``ep_combine``, which + returns the summed results in the original token order. When operating on + a full fixed-capacity buffer, clear unused tail rows and weights before + multiplying so uninitialized values cannot introduce NaNs. NCCL EP zeroes + alignment padding, and combine ignores slots without a routed token. + + .. tab:: JAX + + The integrated JAX MoE layer currently supports multi-GPU execution only + and requires at least two GPUs. + + JAX EP sends BF16 tokens. Low-precision expert computation quantizes the + received tokens after dispatch; it does not change communication precision. + + JAX offers two levels of API, both experimental: + + * ``transformer_engine.jax.moe.moe`` runs the whole MoE block (router, + dispatch, expert MLPs, combine) as a single differentiable call. It is + executed inside a ``Mesh``; ``MeshResource.ep_resource`` names the mesh + axis the experts are sharded over and dispatch and combine become all-to-all + collectives over that axis. It also returns the load-balancing loss when + ``aux_loss_coeff`` is non-zero. Before the first call, initialize EP + eagerly with ``ep_bootstrap`` and record the same configuration with + ``record_ep_bootstrap_signature_for_moe``. The snippet below shows the + required TE EP bootstrap sequence. + * ``transformer_engine.jax.ep`` exposes the primitives separately. Unlike + the PyTorch ``EpBuffer``, the routing state is not kept in an object: + dispatch returns it as arrays and the caller passes them on to combine. + + * ``ep_bootstrap(world_size, rank, num_experts, max_tokens_per_rank, + recv_capacity_per_rank, hidden_dim, ...)`` initializes the EP group + once per process. It runs inside the active ``Mesh`` and reads the EP + axis (and the data-parallel axes) from ``MeshResource``. Multiple local + devices per process require TE's support for borrowing the XLA NCCL + communicator; the example below uses one process per device. + * ``EpLayerConfig(top_k, ...)`` is a small per-layer configuration that + every per-step call takes as its first argument. + * ``ep_dispatch(cfg, topk_idx, tokens, topk_weights, + recv_capacity_per_rank)`` scatters the tokens to the expert ranks and + returns ``(recv_tokens, recv_topk_weights, handle_mem, token_counts, + total_recv_tokens)``: the receive buffer grouped by local expert, the + weights of the received tokens, the routing handle and per-expert + token counts needed by combine, and the pre-drop receive total that can + be used to detect overflow. + * ``ep_combine(cfg, handle_mem, token_counts, expert_out, + num_local_tokens)`` sums the expert outputs back on the source ranks in + the original token order. It is unweighted: multiply ``expert_out`` by + ``recv_topk_weights`` before calling it. NCCL EP zeroes alignment + padding; combine ignores slots without a routed token. + ``num_local_tokens`` must be static because it fixes the output shape. + + .. literalinclude:: moe_expert_parallel_jax.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_JAX + :end-before: # END_MOE_EXPERT_PARALLEL_JAX + +Complete runnable examples: +`examples/pytorch/ep `_ +and `examples/jax/ep `_. diff --git a/docs/features/mixture_of_experts/moe_expert_parallel_jax.py b/docs/features/mixture_of_experts/moe_expert_parallel_jax.py new file mode 100644 index 00000000000..8b66b0a6206 --- /dev/null +++ b/docs/features/mixture_of_experts/moe_expert_parallel_jax.py @@ -0,0 +1,71 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_EXPERT_PARALLEL_JAX +import jax + +from transformer_engine.jax.ep import ep_bootstrap +from transformer_engine.jax.moe import ( + get_moe_recv_capacity_per_rank, + moe, + record_ep_bootstrap_signature_for_moe, +) +from transformer_engine.jax.sharding import MeshResource, global_shard_guard + +num_experts = 8 +top_k = 2 +ep_axis = "ep" +ep_size = mesh.shape[ep_axis] +max_tokens_per_rank = x.shape[0] * x.shape[1] // jax.process_count() +recv_capacity_per_rank = get_moe_recv_capacity_per_rank( + num_experts=num_experts, + num_experts_per_tok=top_k, + max_tokens_per_rank=max_tokens_per_rank, + ep_size=ep_size, +) + +# Initialize EP eagerly once per process. The mesh has one device per process. +mesh_resource = MeshResource(ep_resource=ep_axis) +with mesh, global_shard_guard(mesh_resource): + ep_bootstrap( + world_size=jax.process_count(), + 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=x.shape[-1], + max_token_dtype=x.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=x.shape[-1], + ep_size=ep_size, +) + +# mesh: jax.sharding.Mesh with an "ep" axis and one device per process +# x: [batch, seq, 256], BF16 and sharded over the mesh +# gate_kernel: [256, num_experts] router projection +# wi: [num_experts, 256, 1024] gated FC1, ffn = 512 +# wo: [num_experts, 512, 256] FC2 +with mesh, global_shard_guard(mesh_resource): + output, aux_loss, total_recv_tokens = moe( + x, + gate_kernel, + wi, + wo, + num_experts=num_experts, + num_experts_per_tok=top_k, + activation_type="silu", + score_function="softmax", + aux_loss_coeff=1e-2, # load-balancing loss; 0 disables it + mesh_resource=mesh_resource, + dtype=x.dtype, + recv_capacity_per_rank=recv_capacity_per_rank, + ) +# output: [batch, seq, hidden_size] +# aux_loss: scalar load-balancing loss (None when aux_loss_coeff == 0) +# total_recv_tokens: receive count before any capacity-based token dropping +# END_MOE_EXPERT_PARALLEL_JAX diff --git a/docs/features/mixture_of_experts/moe_expert_parallel_pytorch.py b/docs/features/mixture_of_experts/moe_expert_parallel_pytorch.py new file mode 100644 index 00000000000..744f3689143 --- /dev/null +++ b/docs/features/mixture_of_experts/moe_expert_parallel_pytorch.py @@ -0,0 +1,77 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_EXPERT_PARALLEL_PYTORCH +import torch +import torch.distributed as dist +import transformer_engine.pytorch as te +from transformer_engine.pytorch.ep import EpBuffer, EpConfig, ep_bootstrap + +# ep_group: process group the experts are sharded over +# tokens: [num_tokens, hidden_size] bf16 tokens local to this rank +# topk_idx: [num_tokens, top_k] global expert index per selected expert +# topk_w: [num_tokens, top_k] fp32 routing weights from the router +ep_size = dist.get_world_size(ep_group) +num_local_experts = num_experts // ep_size +config = EpConfig( + top_k=top_k, + hidden_dim=hidden_size, + num_local_experts=num_local_experts, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=None, + ep_group=ep_group, + alignment=0, +) + +# Once per process: sets up NCCL EP on ep_group's communicator. +ep_bootstrap( + ep_group, + num_experts=num_experts, + max_tokens_per_rank=config.max_tokens_per_rank, + hidden_dim=hidden_size, + num_topk=top_k, + recv_capacity_per_rank=config.recv_capacity_per_rank, +) +# One buffer per in-flight layer call (e.g. per pipeline microbatch). +buffer = EpBuffer( + top_k=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + recv_capacity_per_rank=config.recv_capacity_per_rank, + hidden_dim=config.hidden_dim, + num_local_experts=config.num_local_experts, + alignment=config.alignment, + device=tokens.device, +) + +dispatch = te.ops.MoeDispatch(config, buffer) +fc1 = te.ops.GroupedLinear( + num_local_experts, + hidden_size, + 2 * ffn_hidden_size, + bias=False, + device=tokens.device, + dtype=torch.bfloat16, +) +activation = te.ops.ScaledSwiGLU() +fc2 = te.ops.GroupedLinear( + num_local_experts, + ffn_hidden_size, + hidden_size, + bias=False, + device=tokens.device, + dtype=torch.bfloat16, +) +combine = te.ops.MoeCombine(config, buffer) + +# Bind dispatch metadata before constructing or calling the sequence. +dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) +dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) +fc1.set_extra_input_channel(0, "tokens_per_expert") +activation.set_extra_input_channel(0, "routing_weights") +fc2.set_extra_input_channel(0, "tokens_per_expert") +moe = te.ops.Sequential(dispatch, fc1, activation, fc2, combine) + +# Dispatch takes indices and weights; combine takes the same indices. +output = moe(tokens, topk_idx, topk_w, topk_idx) # [num_tokens, hidden_size] +# END_MOE_EXPERT_PARALLEL_PYTORCH diff --git a/docs/features/mixture_of_experts/moe_layer_pytorch.py b/docs/features/mixture_of_experts/moe_layer_pytorch.py new file mode 100644 index 00000000000..211bd417f51 --- /dev/null +++ b/docs/features/mixture_of_experts/moe_layer_pytorch.py @@ -0,0 +1,67 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_LAYER_PYTORCH +import torch +import transformer_engine.pytorch as te +from transformer_engine.pytorch.router import fused_topk_with_score_function + +num_tokens, hidden_size, num_experts, top_k = 16, 64, 4, 2 +hidden_states = torch.randn( + num_tokens, + hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, +) +gate = torch.nn.Linear( + hidden_size, + num_experts, + bias=False, + device="cuda", + dtype=torch.bfloat16, +) +experts = te.GroupedLinear( + num_experts, + hidden_size, + hidden_size, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", +) + +# 1. Router: score the experts and pick the top-k for each token. +logits = gate(hidden_states) +probs, routing_map = fused_topk_with_score_function( + logits, + topk=top_k, + use_pre_softmax=False, + num_groups=None, + group_topk=None, + scaling_factor=None, + score_function="softmax", + expert_bias=None, +) + +# 2. Dispatch: gather tokens into expert-contiguous order. +routing_map = routing_map.to(torch.int32) +permuted, row_id_map = te.moe_permute( + hidden_states, + routing_map, + num_out_tokens=num_tokens * top_k, +) + +# 3. Experts: one grouped call over all expert token blocks. +m_splits = routing_map.sum(dim=0).tolist() # tokens routed to each expert +expert_out = experts(permuted, m_splits) + +# 4. Combine: scatter the outputs back and merge the top-k contributions. +output = te.moe_unpermute( + expert_out, + row_id_map, + merging_probs=probs, + restore_shape=(num_tokens, hidden_size), +) +output.square().mean().backward() +# END_MOE_LAYER_PYTORCH diff --git a/docs/features/mixture_of_experts/moe_permute_pad_pytorch.py b/docs/features/mixture_of_experts/moe_permute_pad_pytorch.py new file mode 100644 index 00000000000..dfacf52594f --- /dev/null +++ b/docs/features/mixture_of_experts/moe_permute_pad_pytorch.py @@ -0,0 +1,43 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_PERMUTE_PAD_PYTORCH +from transformer_engine.pytorch import moe_permute_and_pad_with_probs, moe_unpermute + +# tokens: [num_tokens, hidden_size] +# probs: [num_tokens, num_experts] routing probabilities +# routing_map: [num_tokens, num_experts] int32 mask +# +# Pad each expert's token block up to a multiple of align_size (here 256) so the +# grouped GEMM sees aligned blocks. Permutation and padding happen in one kernel. +tokens_per_expert = routing_map.sum(dim=0) # [num_experts] +padded, permuted_probs, row_id_map, pad_offsets, padded_tokens_per_expert = ( + moe_permute_and_pad_with_probs( + tokens, + probs, + routing_map, + tokens_per_expert, + align_size=256, + ) +) + +# padded: [sum(padded_tokens_per_expert), hidden_size] +# pad_offsets: per-expert cumulative padding (None if already aligned) +# padded_tokens_per_expert: aligned per-expert counts -> m_splits for GroupedLinear + +# ... run an unweighted expert block on `padded`, producing expert_out ... + +# Apply weights once. Skip this multiplication for the scaled grouped MLP, +# which already weights its outputs through ScaledSwiGLU and scale_bias=True. +expert_out = expert_out * permuted_probs[:, None] + +# Combine already weighted outputs without merging_probs. Pass pad_offsets to +# remove padding and restore_shape for the original [num_tokens, hidden_size]. +output = moe_unpermute( + expert_out, + row_id_map, + restore_shape=tokens.shape, + pad_offsets=pad_offsets, +) +# END_MOE_PERMUTE_PAD_PYTORCH diff --git a/docs/features/mixture_of_experts/moe_permute_pytorch.py b/docs/features/mixture_of_experts/moe_permute_pytorch.py new file mode 100644 index 00000000000..cb5a4e263df --- /dev/null +++ b/docs/features/mixture_of_experts/moe_permute_pytorch.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_PERMUTE_PYTORCH +import torch +from transformer_engine.pytorch import moe_permute + +# tokens: [num_tokens, hidden_size] +# routing_map: [num_tokens, num_experts] mask, 1 if token routed to expert +# +# num_out_tokens is the number of rows in the permuted buffer. Reading it from +# the routing map (``int(routing_map.sum())``) triggers a device-to-host sync; +# when the value is known statically (e.g. ``num_tokens * top_k`` for dropless +# routing), prefer passing that constant directly. +permuted, row_id_map = moe_permute( + tokens, + routing_map, + num_out_tokens=int(routing_map.sum()), +) + +# permuted: [num_out_tokens, hidden_size], expert-contiguous +# row_id_map: opaque tensor used by ``moe_unpermute`` to reverse the +# permutation after the experts have run. +# END_MOE_PERMUTE_PYTORCH diff --git a/docs/features/mixture_of_experts/moe_unpermute_pytorch.py b/docs/features/mixture_of_experts/moe_unpermute_pytorch.py new file mode 100644 index 00000000000..4c7cc5c07c6 --- /dev/null +++ b/docs/features/mixture_of_experts/moe_unpermute_pytorch.py @@ -0,0 +1,26 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_UNPERMUTE_PYTORCH +from transformer_engine.pytorch import moe_unpermute + +# expert_out: [num_out_tokens, hidden_size], expert-contiguous, +# produced by an unweighted GroupedLinear or expert MLP. +# row_id_map: returned by moe_permute. +# merging_probs: [num_tokens, num_experts]; routing probabilities used to +# weight the per-expert contributions to each token. They may +# be omitted when every selected routing weight is one or the +# expert outputs are already weighted, as in the scaled grouped MLP. +# restore_shape: the original [num_tokens, hidden_size]. Required when the +# permuted buffer has more rows than the input (top-k routing); +# for top-1 it can be omitted and is inferred from expert_out. +tokens_out = moe_unpermute( + expert_out, + row_id_map, + merging_probs=merging_probs, + restore_shape=(num_tokens, hidden_size), +) + +# tokens_out: [num_tokens, hidden_size], in the original token order +# END_MOE_UNPERMUTE_PYTORCH diff --git a/docs/features/mixture_of_experts/router_jax.py b/docs/features/mixture_of_experts/router_jax.py new file mode 100644 index 00000000000..0b7b51cd123 --- /dev/null +++ b/docs/features/mixture_of_experts/router_jax.py @@ -0,0 +1,44 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_ROUTER_JAX +from transformer_engine.jax.router import fused_topk_with_score_function + +# logits: [num_tokens, num_experts], produced by the gating (router) projection. +# +# Select the top-k experts for each token. The score function and the top-k +# selection run in a single fused kernel. Most arguments have defaults, so a +# basic call only needs the logits, topk and score_function. +score_function = "softmax" # "softmax", "sigmoid" or "sqrtsoftplus" +probs, routing_map = fused_topk_with_score_function( + logits, + topk=2, + score_function=score_function, +) + +# probs: [num_tokens, num_experts], non-zero only at the selected experts. +# routing_map: [num_tokens, num_experts] bool mask of the selected experts. +# END_ROUTER_JAX + + +# START_ROUTER_AUX_JAX +from transformer_engine.jax.router import fused_moe_aux_loss + +# Use dense scores, but keep the actual router's map for counts: auxiliary mode +# ignores expert bias, grouping and scaling. +scores, _ = fused_topk_with_score_function( + logits, + topk=2, + score_function=score_function, + compute_aux_scores=True, +) +tokens_per_expert = routing_map.sum(axis=0) # [num_experts] + +aux_loss = fused_moe_aux_loss( + scores, + tokens_per_expert, + topk=2, + coeff=1e-2, # loss weight; add aux_loss to the training loss +) +# END_ROUTER_AUX_JAX diff --git a/docs/features/mixture_of_experts/router_pytorch.py b/docs/features/mixture_of_experts/router_pytorch.py new file mode 100644 index 00000000000..c2188ecabc9 --- /dev/null +++ b/docs/features/mixture_of_experts/router_pytorch.py @@ -0,0 +1,55 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_ROUTER_PYTORCH +from transformer_engine.pytorch.router import fused_topk_with_score_function + +# logits: [num_tokens, num_experts], produced by the gating (router) projection. +# +# Select the top-k experts for each token and return their routing weights. The +# score function and the top-k selection run in a single fused kernel (all math +# is done in fp32 internally for numerical stability). +score_function = "softmax" # "softmax", "sigmoid" or "sqrtsoftplus" +probs, routing_map = fused_topk_with_score_function( + logits, + topk=2, + use_pre_softmax=False, # softmax after top-k; True selects softmax-then-top-k + num_groups=None, # set with group_topk to enable grouped (device-limited) routing + group_topk=None, + scaling_factor=None, # optional scalar multiplied into the returned probs + score_function=score_function, + expert_bias=None, # [num_experts] selection bias for sigmoid or sqrtsoftplus +) + +# probs: [num_tokens, num_experts], non-zero only at the selected experts. +# Pass directly to moe_unpermute as merging_probs. +# routing_map: [num_tokens, num_experts] bool mask, True at the selected experts. +# Cast to int32 and pass to moe_permute. +# END_ROUTER_PYTORCH + + +# START_ROUTER_AUX_PYTORCH +from transformer_engine.pytorch.router import ( + fused_compute_score_for_moe_aux_loss, + fused_moe_aux_loss, +) + +# Use dense scores, but keep the actual router's map for counts: this helper +# does not apply expert bias or grouped routing. +_, scores = fused_compute_score_for_moe_aux_loss( + logits, + topk=2, + score_function=score_function, +) +tokens_per_expert = routing_map.sum(dim=0) # [num_experts] + +aux_loss = fused_moe_aux_loss( + scores, + tokens_per_expert, + total_num_tokens=logits.shape[0], + num_experts=logits.shape[1], + topk=2, + coeff=1e-2, # loss weight; add aux_loss to the training loss +) +# END_ROUTER_AUX_PYTORCH diff --git a/docs/index.rst b/docs/index.rst index 68060974f45..0d3746fcc2b 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -55,6 +55,7 @@ Transformer Engine documentation :caption: Features features/low_precision_training/index.rst + features/mixture_of_experts/mixture_of_experts.rst features/other_optimizations/index.rst diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 44214063ac4..4e9bfe4aa8d 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -2177,10 +2177,6 @@ def moe( FSDP shards the gated dimension, layout conversion permutes gathered FP8 data and inverse scales, preserving global gate/up pairing without gathering or requantizing full-precision weights. - Per-expert dispatch-slot alignment defaults to 128 tokens (``_ALIGN_SIZE``). - Requesting cuDNN fusion reserves 256 tokens, also when falling back, to - preserve compatibility with EP bootstrap buffer sizing. - use_cudnn_fusion : bool Defaults to ``True``: try cuDNN's shared JAX grouped MXFP8 GLU/dGLU APIs on SM100+ GPUs; cuDNN selects the architecture-specific kernels. ``False`` uses unfused TE grouped GEMM. @@ -2189,7 +2185,6 @@ def moe( weights fall back to unfused TE grouped GEMM. Ineligible calls warn and fall back to TE's regular grouped-GEMM implementation. API signatures and GPU capability determine support. - cudnn_native_weight_layout : Optional[bool] ``False`` selects standard contiguous gate/up [E,K,2N] storage; ``True`` selects native [E,2N,K] storage with alternating 32-column @@ -2197,6 +2192,13 @@ def moe( shape. Ambiguous square shapes require an explicit flag. This option describes parameter storage independently of the selected execution path. + Notes + ----- + Per-expert dispatch-slot alignment is 128 tokens (``_ALIGN_SIZE``) with + ``use_cudnn_fusion=False``. Requesting cuDNN fusion (the default) reserves + 256 tokens, also when falling back, to preserve compatibility with EP + bootstrap buffer sizing. + MeshResource fields name physical mesh axes, not Flax logical axes. ``input_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` and ``RouterComputationInfo.gate_kernel_axes`` remain logical-axis tuples diff --git a/transformer_engine/jax/router.py b/transformer_engine/jax/router.py index cd5c1dec6dd..0778155e143 100644 --- a/transformer_engine/jax/router.py +++ b/transformer_engine/jax/router.py @@ -145,14 +145,15 @@ def fused_topk_with_score_function( Returns ------- probs_or_scores : jnp.ndarray - When compute_aux_scores=False: Sparse probability tensor, shape [num_tokens, num_experts]. - Non-zero only at selected expert positions. - When compute_aux_scores=True: Dense score tensor, shape [num_tokens, num_experts]. - All expert positions contain scores. + When ``compute_aux_scores=False``: sparse probability tensor, shape + ``[num_tokens, num_experts]``, non-zero only at selected expert positions. + When ``compute_aux_scores=True``: dense score tensor, shape + ``[num_tokens, num_experts]``, all expert positions contain scores. routing_map : jnp.ndarray - Shape/dtype depend on routing_map_format: - - BYTEMAP: bool[num_tokens, num_experts]; True at selected expert positions. - - BITMAP_U8: uint8[num_tokens, ceil(num_experts/8)]; LSB-first bit-packed + Shape/dtype depend on ``routing_map_format``: + + - BYTEMAP: ``bool[num_tokens, num_experts]``; True at selected expert positions. + - BITMAP_U8: ``uint8[num_tokens, ceil(num_experts/8)]``; LSB-first bit-packed along the expert axis. """ if not isinstance(scaling_factor, (int, float)): diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 508c2cba769..3e09286733d 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -243,6 +243,7 @@ def fused_topk_with_score_function( ): """ Fused topk with score function router. + Parameters ---------- logits : torch.Tensor in fp32/bf16/fp16 @@ -283,9 +284,10 @@ def fused_topk_with_score_function( Same shape as ``logits``. routing_map : torch.Tensor Same leading dims as ``logits``; trailing dim and dtype depend on - routing_map_format, or dense top-k indices when topk_indices is provided: - - BYTEMAP: bool[*logits.shape[:-1], num_experts] - - BITMAP_U8: uint8[*logits.shape[:-1], ceil(num_experts/8)] + ``routing_map_format``, or dense top-k indices when ``topk_indices`` is provided: + + - BYTEMAP: ``bool[*logits.shape[:-1], num_experts]`` + - BITMAP_U8: ``uint8[*logits.shape[:-1], ceil(num_experts/8)]``, LSB-first bit-packed. """ if logits.dtype == torch.float64: @@ -483,6 +485,7 @@ def fused_moe_aux_loss( ) -> torch.Tensor: """ Fused MoE aux loss. + Parameters ---------- probs : torch.Tensor in fp32/bf16/fp16 From 6f6b8f7680bcfebb7a5cc9135ebc4ae30ea99be6 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Thu, 29 Oct 2026 12:38:12 -0700 Subject: [PATCH 22/32] [PyTorch]Validate grouped split alignment on device (#3745) * Validate grouped split alignment on device Signed-off-by: Varun Thumbe * Use versioned offset scan API for split validation Signed-off-by: Varun Thumbe * Consolidate grouped split alignment replay test Signed-off-by: Varun Thumbe * Simplify device alignment test and isolate L0 invocation Signed-off-by: Varun Thumbe * Move split alignment test into grouped tensor suite Signed-off-by: Varun Thumbe * Reduce grouped split alignment scan overhead Signed-off-by: Varun Thumbe * Test grouped split alignment during graph replay Signed-off-by: Varun Thumbe * Cover both grouped split scan sizes Signed-off-by: Varun Thumbe * Remove CUDA context poisoning alignment test Signed-off-by: Varun Thumbe * Apply suggestion from @timmoon10 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * Apply suggestion from @timmoon10 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- .../transformer_engine/transformer_engine.h | 13 ++++ .../common/util/splits_to_offsets.cu | 63 +++++++++++++++---- transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/misc.cpp | 9 +-- .../pytorch/csrc/extensions/pybind.cpp | 3 +- .../pytorch/ops/basic/grouped_linear.py | 7 ++- .../pytorch/ops/fused/grouped_mlp.py | 1 + 7 files changed, 76 insertions(+), 22 deletions(-) diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 1adb47773ef..9a08b15a8a8 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -546,6 +546,19 @@ void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, const int64_t *strides, const int *include_leading_zero, size_t num_outputs, cudaStream_t stream); +/*! \brief Compute multiple scaled offsets and assert row split alignment on the device. + * + * This uses the same prefix-sum kernel as nvte_splits_to_offsets_multi and does not synchronize + * the host, so the validation remains active during CUDA graph replay. + * + * \param[in] split_alignment Positive required divisor of every split size. + * \see nvte_splits_to_offsets_multi for the remaining parameters. + */ +void nvte_splits_to_offsets_multi_v2(NVTETensor split_sizes, size_t split_alignment, + NVTETensor *outputs, const int64_t *strides, + const int *include_leading_zero, size_t num_outputs, + cudaStream_t stream); + /*! \brief TE Grouped Tensor type * * NVTEGroupedTensor is a collection of tensors with potentially different shapes diff --git a/transformer_engine/common/util/splits_to_offsets.cu b/transformer_engine/common/util/splits_to_offsets.cu index 2a06603dafc..169052fa9de 100644 --- a/transformer_engine/common/util/splits_to_offsets.cu +++ b/transformer_engine/common/util/splits_to_offsets.cu @@ -18,6 +18,7 @@ namespace transformer_engine::splits_to_offsets { namespace { constexpr size_t kThreadsPerBlock = 256; +constexpr size_t kSmallThreadsPerBlock = 128; struct KernelArgs { static constexpr size_t kMaxNumOutputs = 8; @@ -29,6 +30,7 @@ struct KernelArgs { int64_t strides[kMaxNumOutputs] = {}; bool include_leading_zero[kMaxNumOutputs] = {}; size_t num_outputs = 0; + size_t split_alignment = 1; }; __device__ __forceinline__ int64_t load_split_size(const void *__restrict__ split_sizes, @@ -58,7 +60,8 @@ __device__ __forceinline__ void store_output(void *__restrict__ output, DType dt } } -__global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) { +template +__global__ void __launch_bounds__(kScanThreads) kernel(KernelArgs args) { const size_t tid = threadIdx.x; // Fill leading zeros if needed @@ -71,7 +74,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) { } // Workspace for prefix sum chunk - __shared__ int64_t block_scan[kThreadsPerBlock]; + __shared__ int64_t block_scan[kScanThreads]; // Sum from previous chunks __shared__ int64_t chunk_prefix; @@ -81,19 +84,35 @@ __global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) { __syncthreads(); // Perform prefix sum in chunks - for (size_t chunk_start = 0; chunk_start < args.num_splits; chunk_start += kThreadsPerBlock) { + for (size_t chunk_start = 0; chunk_start < args.num_splits; chunk_start += kScanThreads) { const size_t idx = chunk_start + tid; - // Load input from global memory into shared memory - if (idx < args.num_splits) { - block_scan[tid] = load_split_size(args.split_sizes, args.split_sizes_dtype, idx); - } else { - block_scan[tid] = 0; + // Validate before the scan. All threads leave together so no thread waits at a barrier + // after another thread has trapped on invalid split metadata. + const int64_t split = + idx < args.num_splits ? load_split_size(args.split_sizes, args.split_sizes_dtype, idx) : 0; + if (args.split_alignment > 1) { + const size_t alignment_mask = args.split_alignment - 1; + bool misaligned; + if ((args.split_alignment & alignment_mask) == 0) { + misaligned = (static_cast(split) & alignment_mask) != 0; + } else { + misaligned = static_cast(split) % args.split_alignment != 0; + } + const bool invalid = idx < args.num_splits && (split < 0 || misaligned); + if (__syncthreads_or(invalid)) { + if (tid == 0) { + NVTE_DEVICE_ERROR("Grouped split size is not divisible by split_alignment."); + __trap(); // NVTE_DEVICE_ERROR's assert is compiled out of release builds. + } + return; + } } + block_scan[tid] = split; __syncthreads(); // Prefix sum in shared memory - for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) { + for (size_t offset = 1; offset < kScanThreads; offset <<= 1) { const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0; __syncthreads(); block_scan[tid] += addend; @@ -112,7 +131,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) { // Update sum for later chunks __syncthreads(); - if (tid == kThreadsPerBlock - 1) { + if (tid == kScanThreads - 1) { chunk_prefix += block_scan[tid]; } __syncthreads(); @@ -187,7 +206,11 @@ void nvte_splits_to_offsets(const int64_t *split_sizes, int64_t *output, size_t args.strides[0] = stride; args.include_leading_zero[0] = true; args.num_outputs = 1; - s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args); + if (num_splits <= s2o::kSmallThreadsPerBlock) { + s2o::kernel<<<1, s2o::kSmallThreadsPerBlock, 0, stream>>>(args); + } else { + s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args); + } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -230,10 +253,19 @@ void nvte_splits_to_offsets_2d(const NVTETensor first_dims, const NVTETensor las void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, const int64_t *strides, const int *include_leading_zero, size_t num_outputs, cudaStream_t stream) { - NVTE_API_CALL(nvte_splits_to_offsets_multi); + nvte_splits_to_offsets_multi_v2(split_sizes, 1, outputs, strides, include_leading_zero, + num_outputs, stream); +} + +void nvte_splits_to_offsets_multi_v2(NVTETensor split_sizes, size_t split_alignment, + NVTETensor *outputs, const int64_t *strides, + const int *include_leading_zero, size_t num_outputs, + cudaStream_t stream) { + NVTE_API_CALL(nvte_splits_to_offsets_multi_v2); using namespace transformer_engine; namespace s2o = transformer_engine::splits_to_offsets; + NVTE_CHECK(split_alignment > 0, "split_alignment must be positive."); if (num_outputs == 0) { return; } @@ -279,6 +311,7 @@ void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, args.split_sizes_dtype = split_sizes_dtype; args.num_splits = num_splits; args.num_outputs = chunk_size; + args.split_alignment = split_alignment; for (size_t i = 0; i < chunk_size; ++i) { const size_t out_idx = chunk_start + i; args.outputs[i] = output_tensors[out_idx]->data.dptr; @@ -286,7 +319,11 @@ void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, args.strides[i] = strides[out_idx]; args.include_leading_zero[i] = include_leading_zero[out_idx] != 0; } - s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args); + if (num_splits <= s2o::kSmallThreadsPerBlock) { + s2o::kernel<<<1, s2o::kSmallThreadsPerBlock, 0, stream>>>(args); + } else { + s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args); + } NVTE_CHECK_CUDA(cudaGetLastError()); } } diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 88dfb808420..c459f963fba 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -585,7 +585,7 @@ at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_ std::tuple> splits_to_offsets_multi( const at::Tensor &split_sizes, const c10::Device &device, const std::vector &strides, const std::vector &include_leading_zero, const std::vector &dtypes, - bool bulk_allocate_outputs); + bool bulk_allocate_outputs, size_t split_alignment); at::Tensor copy_data_ptrs_to_device(const std::vector &tensors, const c10::Device &device); diff --git a/transformer_engine/pytorch/csrc/extensions/misc.cpp b/transformer_engine/pytorch/csrc/extensions/misc.cpp index 01bcd7eb05a..e17e362438b 100644 --- a/transformer_engine/pytorch/csrc/extensions/misc.cpp +++ b/transformer_engine/pytorch/csrc/extensions/misc.cpp @@ -42,7 +42,7 @@ at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_ std::tuple> splits_to_offsets_multi( const at::Tensor &split_sizes, const c10::Device &device, const std::vector &strides, const std::vector &include_leading_zero, const std::vector &dtypes, - bool bulk_allocate_outputs) { + bool bulk_allocate_outputs, size_t split_alignment) { const size_t num_outputs = strides.size(); const size_t num_splits = static_cast(split_sizes.numel()); @@ -51,6 +51,7 @@ std::tuple> splits_to_offsets_multi( "strides, include_leading_zero, and dtypes must have matching lengths, but got ", strides.size(), ", ", include_leading_zero.size(), ", and ", dtypes.size(), "."); NVTE_CHECK(device.is_cuda(), "device must be CUDA, but got ", device.str(), "."); + NVTE_CHECK(split_alignment > 0, "split_alignment must be positive."); // Convert split sizes to int64 GPU tensor. const at::Tensor split_sizes_i64 = @@ -94,9 +95,9 @@ std::tuple> splits_to_offsets_multi( auto split_sizes_nvte = makeTransformerEngineTensor(split_sizes_out); NVTE_SCOPED_GIL_RELEASE({ - nvte_splits_to_offsets_multi(split_sizes_nvte.data(), outputs_nvte.data(), strides.data(), - include_leading_zero_int.data(), num_outputs, - at::cuda::getCurrentCUDAStream()); + nvte_splits_to_offsets_multi_v2(split_sizes_nvte.data(), split_alignment, outputs_nvte.data(), + strides.data(), include_leading_zero_int.data(), num_outputs, + at::cuda::getCurrentCUDAStream()); }); return {split_sizes_out, std::move(outputs)}; diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index ebcdd17b07c..bd500274ee8 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -625,7 +625,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("splits_to_offsets_multi", &transformer_engine::pytorch::splits_to_offsets_multi, "Compute multiple scaled inclusive-scan offsets from a split-sizes vector", py::arg("split_sizes"), py::arg("device"), py::kw_only(), py::arg("strides"), - py::arg("include_leading_zero"), py::arg("dtypes"), py::arg("bulk_allocate") = false); + py::arg("include_leading_zero"), py::arg("dtypes"), py::arg("bulk_allocate") = false, + py::arg("split_alignment") = 1); m.def("get_num_cublas_streams", &nvte_get_num_compute_streams, "Get number of compute streams", py::call_guard()); diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index fe5e77fc580..86fa740c5b6 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -186,8 +186,7 @@ class GroupedLinear(BasicOperation): The grouped-tensor path is selected automatically where it is supported (see ``is_op_fuser_grouped_tensor_path_supported``). With MXFP8 or FP8 block scaling, it relies on every split size being a - multiple of 128, as ``split_alignment`` declares. This is checked only - on the device, so other split sizes can produce incorrect results. + multiple of 128, as ``split_alignment`` declares. Parameters ---------- @@ -232,7 +231,8 @@ class GroupedLinear(BasicOperation): in the forward pass. The scale tensor has shape ``(total_tokens,)`` and is split according to the split sizes. split_alignment : int, default = ``128`` - Guaranteed alignment of each input split size. + Expected alignment for input split sizes. Split sizes that do not + match alignment may produce errors or incorrect results. """ # Operation expects input split sizes (and optionally scales tensor) @@ -1389,6 +1389,7 @@ def _fuser_forward_grouped_tensor( include_leading_zero=[False, True, True, True], dtypes=[torch.int32, torch.int64, torch.int64, torch.int64], bulk_allocate=True, + split_alignment=self.split_alignment, ) split_points = grouped_tensor_offsets[0] base_split_offsets = grouped_tensor_offsets[1] diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index bc64d8ed9ae..bba320aad17 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1489,6 +1489,7 @@ def fuser_forward( torch.int64, ], bulk_allocate=True, + split_alignment=256, ) # Prepare FC1 grouped weight tensor for fused kernels. From 5f3abe14458cbcb08ef9d4f9bc0e5beee524fcc5 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:42:50 -0700 Subject: [PATCH 23/32] [PyTorch] Skip FA 2.8.4 for unsupported arches in L3 (#3772) * [PyTorch] Skip FA2 in L3 tests on archs without native FA2 support Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [PyTorch] Build FA2 for the family base arch instead of skipping Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- qa/L3_pytorch_FA_versions_test/test.sh | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 5a159fa6e95..4fb41314b14 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -37,14 +37,31 @@ export FLASH_ATTN_CUDA_ARCHS=$sm_arch # FA4: routine tests are here; sm100+ only; non-CP and CP FA2_SOURCE_VERSION="2.8.4" FA2_SOURCE_COMMIT="1fc0221183e799363001394a12dd9e78e15f5dda" +FA2_SOURCE_GENCODE_ARCHS="80 90 100 110 120 121" +# FA2 names only family base archs, and a base build covers the rest of its major +# family, so sm103/sm107 run the compute_100f cubin that "100" produces. +fa2_arch=$(( sm_arch / 10 * 10 )) +case " $FA2_SOURCE_GENCODE_ARCHS " in + *" $sm_arch "*) fa2_arch=$sm_arch ;; +esac +FA2_versions=() +case " $FA2_SOURCE_GENCODE_ARCHS " in + *" $fa2_arch "*) + FA2_versions=("$FA2_SOURCE_VERSION") + export FLASH_ATTN_CUDA_ARCHS=$fa2_arch + [ "$fa2_arch" = "$sm_arch" ] \ + || echo "Building FA2 for sm${fa2_arch} on sm${sm_arch} (FA2 has no --gencode for sm${sm_arch})" + ;; + *) echo "Skipping FA2 on sm${sm_arch} as FA2 has no --gencode for this arch" ;; +esac CP_FA_VERSION="" if [ $sm_arch -gt 90 ] then - FA_versions=(2.8.4 4.0.0b33) + FA_versions=("${FA2_versions[@]}" 4.0.0b33) CP_FA_VERSION="4.0.0b33" elif [ $sm_arch -eq 90 ] then - FA_versions=(2.8.4 3.0.0b1) + FA_versions=("${FA2_versions[@]}" 3.0.0b1) CP_FA_VERSION="3.0.0b1" else error_exit "No L3 FlashAttention generation is defined for sm${sm_arch}" From c83e856caebff2be63561026b873bfc4771044da Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 16:49:36 -0700 Subject: [PATCH 24/32] [JAX] Disable TE C API CuTeDSL backend at import time (#3784) [JAX] Disable CuTeDSL at import time Signed-off-by: Jeremy Berchtold --- transformer_engine/jax/__init__.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/transformer_engine/jax/__init__.py b/transformer_engine/jax/__init__.py index d0afc1ff258..45379dbf0d9 100644 --- a/transformer_engine/jax/__init__.py +++ b/transformer_engine/jax/__init__.py @@ -22,6 +22,9 @@ # pylint: disable=wrong-import-position +import os +import warnings + # This unused import is needed because the top level `transformer_engine/__init__.py` # file catches an `ImportError` as a guard for cases where the given framework's # extensions are not available. @@ -29,6 +32,13 @@ from transformer_engine.common import load_framework_extension +if os.getenv("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0": + warnings.warn( + "TE/JAX is not yet compatible with CuTeDSL due to XLA, TE/JAX, and Python GC conflicts. " + "We are working on a fix; in the meantime NVTE_ENABLE_CUTEDSL_BACKEND is forced off." + ) +os.environ["NVTE_ENABLE_CUTEDSL_BACKEND"] = "0" + load_framework_extension("jax") from . import flax From 43f61f79eb915d0d7483ca799185edd82cc12ae5 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:25:39 -0700 Subject: [PATCH 25/32] Add code quality guidelines to AGENTS.md (#3773) Signed-off-by: Tim Moon Co-authored-by: Claude Opus 5.5 --- AGENTS.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index 7dceca65f63..cb33db2f1f9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -374,6 +374,33 @@ requirements, or performance-relevant behavior. - Standalone examples live under `examples/pytorch/` and `examples/jax/`. - Tutorials and notebooks live under `docs/examples/`. +## Code Quality + +It is not sufficient for code to be correct and performant: it must be +human-readable as well. Human review is the bottleneck in TE +development, and PRs that are unpleasant to read are more likely to be +neglected by reviewers. These guidelines are defaults, not substitutes +for judgment. + +- Public APIs, including the C API, must be backward compatible. Design + new ones to be clean, general, and extensible. Within reason, generalize + existing APIs rather than adding new ones. +- Limit changes to the task's scope; do not refactor unrelated code. +- Within that scope, restructure abstractions when needed to keep the code + coherent rather than adding workarounds that minimize the diff. Avoid + overengineering and unnecessary indirection. +- Write code that is understandable without its comments, and add + comments where they aid understanding. Keep them concise, since + humans skim long comments. Describe the code as it is for a reader + without context, not its edit history. Treat an incorrect comment + like incorrect code. +- Use descriptive names that read naturally. Follow the conventions of the + surrounding code and domain, and avoid inventing idiosyncratic + terminology. +- Extend existing infrastructure rather than building parallel systems. + For example, add cases to an existing test suite instead of creating + a new test file for each change. + ## Contribution Requirements Follow `CONTRIBUTING.rst` and `.github/PULL_REQUEST_TEMPLATE.md`. From 45a6d7cf5f98693a39483d60b9ecb562d63b4b35 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 29 Oct 2026 19:03:07 -0700 Subject: [PATCH 26/32] [PyT] Disable TE torch compilation in KV-cache CI tests (#3781) Disable TE torch compilation in KV-cache CI tests to avoid numerical mismatches Signed-off-by: Kshitij Lakhani --- qa/L0_pytorch_unittest/test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d10..9051e85f686 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -72,7 +72,7 @@ NVTE_KDA_TEST_REQUIRED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pyt NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla.xml $TE_PATH/tests/pytorch/attention/test_fused_mla.py || test_fail "test_fused_mla.py" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" +NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mixed_thd_attention.xml $TE_PATH/tests/pytorch/attention/test_mixed_thd_attention.py || test_fail "test_mixed_thd_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" From aac5cb4abb470bf7c73028a7e2545bb6abb76b23 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:16:44 -0700 Subject: [PATCH 27/32] Grouped Linear CPU Overhead Reductions By Removing Unnecessary Python Logic (#3613) * Reduce module GroupedLinear CPU preparation overhead Signed-off-by: zhongboz * reduce cpu overhead for op fuser Signed-off-by: zhongboz * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * some fixes Signed-off-by: zhongboz * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix GroupedLinear bias gradients with delayed WGRAD Signed-off-by: zhongboz * Fix GroupedLinear bias gradients with frozen weights Signed-off-by: zhongboz * resolve comments Signed-off-by: zhongboz * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * resolve comments Signed-off-by: zhongboz * grep tile comments Signed-off-by: Zhongbo Zhu * Remove quantizer caching in grouped linear op Signed-off-by: Tim Moon * Test grouped linear across training and inference Signed-off-by: Tim Moon * Refactor grouped linear module implementation Disentangle backend selection from quantizer extraction. Remove overly-specific helper function for backend selection. Remove test that calls internal helper functions. Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make quantized single grouped weights own their quantizer Make sure to set column-wise usage in quantized weight param quantizer during forward pass. Signed-off-by: Tim Moon --------- Signed-off-by: zhongboz Signed-off-by: Zhongbo Zhu Signed-off-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon --- tests/pytorch/test_fusible_ops.py | 101 +++++ tests/pytorch/test_grouped_linear.py | 389 +++++++++++++++++- tests/pytorch/test_grouped_mlp.py | 82 ++++ ...odule_grouped_linear_distributed_weight.py | 93 +++++ tests/pytorch/test_quantized_tensor.py | 46 +++ .../pytorch/cpp_extensions/gemm.py | 10 +- .../pytorch/module/grouped_linear.py | 262 +++++++----- .../pytorch/ops/basic/grouped_linear.py | 86 ++-- .../pytorch/quantized_tensor.py | 18 +- 9 files changed, 942 insertions(+), 145 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 17efba6f515..b3c3d3d62dc 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -4,6 +4,7 @@ from __future__ import annotations + from collections.abc import Iterable, Sequence import io import os @@ -5706,3 +5707,103 @@ def test_linear_inference_loop( quantization=quantization, recipe=recipe, ) + + @pytest.mark.parametrize("quantization", (None, "mxfp8", "fp8_current_scaling")) + def test_grouped_linear_training_loop( + self, + *, + stages: Sequence[str] = ("train", "infer", "train", "infer"), + num_groups: int = 3, + weight_shape: tuple[int, int] = (128, 128), + dtype: Optional[torch.dtype] = None, + device: torch.device = "cuda", + quantization: Optional[str], + ) -> None: + """Training loops with grouped linear op.""" + if dtype is None: + dtype = torch.bfloat16 if is_bf16_available() else torch.float32 + + # Tensor shapes + split_sizes = [128 * (1 + group_idx % 2) for group_idx in range(num_groups)] + splits = torch.tensor(split_sizes, dtype=torch.int64, device=device) + out_features, in_features = weight_shape + in_shape = [sum(split_sizes), in_features] + out_shape = in_shape[:-1] + [out_features] + + # Quantization recipe + quantized_compute = quantization is not None + recipe = make_recipe(quantization) + maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) + maybe_skip_quantization(quantization, dims=out_shape) + + # Initialize weights + with te.quantized_model_init(enabled=quantized_compute, recipe=recipe): + module = te_ops.GroupedLinear( + num_groups, + in_features, + out_features, + bias=False, + dtype=dtype, + device=device, + ) + weights = list(module.parameters()) + assert len(weights) == len(split_sizes) + + # Expected numerical error + tols = dtype_tols(dtype) + if dtype == torch.float32: + tols = dtype_tols(torch.float16) # TF32 GEMM + if quantization is not None: + tols = quantization_tols(quantization) + + # Training loop stages + for stage in stages: + training = stage == "train" + if training: + with torch.no_grad(): + # Update parameters with random values to simulate + # optimizer step or FSDP param all-gather + for weight in weights: + weight.copy_(torch.empty_like(weight).uniform_()) + weight.grad = None + + # Forward reference implementation + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + quantization=quantization, + test_dtype=dtype, + test_device=device, + requires_grad=training, + ) + weights_ref = [to_cpu(weight) for weight in weights] + outputs_ref = [ + torch.nn.functional.linear(x_part, weight_ref) + for x_part, weight_ref in zip(torch.split(x_ref, split_sizes), weights_ref) + ] + y_ref = torch.cat(outputs_ref) + + # Grouped linear forward + with ( + torch.inference_mode(not training), + te.autocast(enabled=quantized_compute, recipe=recipe), + ): + y_test = module(x_test, splits) + + # Backward pass + if training: + dy_ref, dy_test = make_reference_and_test_tensors( + out_shape, + quantization=quantization, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + y_ref.backward(dy_ref) + y_test.backward(dy_test) + + # Check results + assert_close(y_test, y_ref, **tols) + if training: + assert_close_grads(x_test, x_ref, **tols) + for weight, weight_ref in zip(weights, weights_ref): + assert_close_grads(weight, weight_ref, **tols) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 0a0c42b03eb..19b27798c5d 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -4,6 +4,7 @@ import os import random +from unittest.mock import Mock from typing import Dict, List, Optional, Sequence import pytest @@ -60,6 +61,35 @@ seed = 1234 reset_rng_states() + +@pytest.mark.parametrize("has_skip", [False, True]) +@pytest.mark.parametrize("capturing", [False, True]) +def test_grouped_linear_skip_update_capture_short_circuit(monkeypatch, has_skip, capturing): + """Only an active capture skip tensor may override first-microbatch accumulation.""" + FP8GlobalStateManager.reset() + module = GroupedLinear(2, 128, 128, bias=False, params_dtype=torch.bfloat16) + x = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16) + skip = torch.ones(1, device="cuda") if has_skip else None + monkeypatch.setattr( + FP8GlobalStateManager.quantization_state, "skip_fp8_weight_update_tensor", skip + ) + capture = Mock(return_value=capturing) + monkeypatch.setattr(FP8GlobalStateManager, "fp8_graph_capturing", capture) + original = _GroupedLinear.forward + seen = [] + + def forward(ctx, inp, fwd_args, *args): + seen.append((fwd_args.is_first_microbatch, fwd_args.skip_fp8_weight_update)) + return original(ctx, inp, fwd_args, *args) + + monkeypatch.setattr(_GroupedLinear, "forward", forward) + with torch.no_grad(): + module(x, [128, 128], is_first_microbatch=True) + assert capture.call_count == int(has_skip) + assert seen[0][0] == (not (has_skip and capturing)) + assert seen[0][1] is (skip if capturing else None) + + NVTE_TEST_NVINSPECT_ENABLED = int(os.environ.get("NVTE_TEST_NVINSPECT_ENABLED", "0")) if NVTE_TEST_NVINSPECT_ENABLED: @@ -2310,6 +2340,161 @@ def test_single_grouped_primary_mxfp8_bypasses_weight_workspace(): assert "weight" not in grouped_linear._fp8_workspaces +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +@pytest.mark.parametrize("primary_fp8", [False, True]) +@pytest.mark.parametrize("save_original_input", [False, True]) +def test_grouped_runtime_usage_transitions(primary_fp8, save_original_input): + """Preserve primary usages and rebuild missing columnwise workspaces after inference.""" + from transformer_engine.pytorch.distributed import activation_recompute_forward + + fp8_recipe = recipe.MXFP8BlockScaling() + if not is_module_grouped_tensor_path_supported(fp8_recipe, torch.bfloat16): + pytest.skip("Native MXFP8 grouped GEMM is unavailable.") + FP8GlobalStateManager.reset() + with quantized_model_init(enabled=primary_fp8, recipe=fp8_recipe): + module = GroupedLinear( + 2, + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + single_grouped_weight=True, + use_grouped_tensor=True, + save_original_input=save_original_input, + ) + splits = torch.tensor([128, 128], device="cuda") + x = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def run(*, grad_enabled=True, first=False): + with torch.set_grad_enabled(grad_enabled), autocast(enabled=True, recipe=fp8_recipe): + return module(x, splits, is_first_microbatch=first) + + run(first=True).sum().backward() + reference_grad = x.grad.clone() + reference_wgrad = module.weight.grad.clone() + weight_q = module.weight.quantizer if primary_fp8 else module._get_weight_quantizers()[0] + assert weight_q.columnwise_usage + module.eval() + run(grad_enabled=False, first=True) + # Persistent parameters retain columnwise usage; temporary workspaces follow this call. + assert weight_q.columnwise_usage == primary_fp8 + # Force a rowwise-only workspace, as on a first inference microbatch. + if not primary_fp8: + module._fp8_workspaces.clear() + run(grad_enabled=False, first=True) + assert module._fp8_workspaces["weight"].columnwise_data is None + module.train() + module.weight.grad = None + x.grad = None + # Keep two forwards alive before either backward; quantizers are shared references. + y1, y2 = run(), run() + assert weight_q.columnwise_usage + y1.sum().backward() + y2.sum().backward() + torch.testing.assert_close(x.grad, 2 * reference_grad, rtol=1e-2, atol=5e-3) + torch.testing.assert_close(module.weight.grad, 2 * reference_wgrad, rtol=1e-2, atol=5e-3) + with activation_recompute_forward(activation_recompute=True, recompute_phase=False): + run(grad_enabled=False, first=True) + assert weight_q.columnwise_usage + with activation_recompute_forward(activation_recompute=True, recompute_phase=True): + run().sum().backward() + assert weight_q.columnwise_usage + module.weight.requires_grad_(False) + run().sum().backward() + input_q = module.quantizers["scaling_fwd"][module._offsets["input"]] + assert not input_q.columnwise_usage + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +def test_grouped_runtime_recipe_rebuild(): + """Recipe rebuilds replace cached quantizers while preserving saved forwards.""" + recipes = [ + recipe.MXFP8BlockScaling(), + recipe.Float8CurrentScaling(), + recipe.MXFP8BlockScaling(), + ] + if not all(is_module_grouped_tensor_path_supported(r, torch.bfloat16) for r in recipes): + pytest.skip("Required native grouped GEMM recipes are unavailable.") + FP8GlobalStateManager.reset() + module = GroupedLinear( + 2, + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + single_grouped_weight=True, + use_grouped_tensor=True, + ) + reference = GroupedLinear( + 2, + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + use_grouped_tensor=True, + ) + reference.load_state_dict(module.state_dict()) + x = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + splits = torch.tensor([128, 128], device="cuda") + pending = [] + workspaces = [] + for fp8_recipe in recipes: + with autocast(enabled=True, recipe=fp8_recipe): + y = module(x, splits, is_first_microbatch=True) + ref_y = reference(ref_x, splits, is_first_microbatch=True) + pending.append((y, ref_y)) + workspace = module._fp8_workspaces["weight"] + assert workspace.quantizer is module._get_weight_quantizers()[0] + assert all(workspace.quantizer is not old.quantizer for old in workspaces) + workspaces.append(workspace) + for y, ref_y in reversed(pending): + torch.testing.assert_close(y, ref_y, rtol=1e-2, atol=5e-3) + y.sum().backward() + ref_y.sum().backward() + torch.testing.assert_close(x.grad, ref_x.grad, rtol=1e-2, atol=5e-3) + ref_wgrad = torch.stack([reference.weight0.grad, reference.weight1.grad]) + torch.testing.assert_close(module.weight.grad, ref_wgrad, rtol=1e-2, atol=5e-3) + + +@pytest.mark.parametrize( + "use_grouped_tensor,single_bias", [(False, False), (True, False), (True, True)] +) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_grouped_linear_return_bias_with_frozen_weights(use_grouped_tensor, single_bias, use_bias): + """A trainable returned bias keeps the output differentiable when input/weight are frozen.""" + if use_grouped_tensor and not is_module_grouped_tensor_path_supported(None, torch.bfloat16): + pytest.skip("Native grouped GEMM is unavailable.") + module = GroupedLinear( + 2, + 128, + 128, + bias=use_bias, + return_bias=True, + params_dtype=torch.bfloat16, + single_grouped_weight=single_bias, + single_grouped_bias=single_bias, + use_grouped_tensor=use_grouped_tensor, + ) + for weight in _grouped_linear_weight_params(module): + weight.requires_grad_(False) + x = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16) + splits = torch.tensor([128, 128], device="cuda" if use_grouped_tensor else "cpu") + y, biases = module(x, splits) + assert y.requires_grad == use_bias + if not use_bias: + assert isinstance(biases, list) and len(biases) == 2 + assert all(b.numel() == 0 for b in biases) + return + params = _grouped_linear_bias_params(module) + returned = [biases] if single_bias else biases + assert len(returned) == len(params) + assert all(value is param for value, param in zip(returned, params)) + (y.sum() + sum(value.sum() for value in returned)).backward() + for param in params: + torch.testing.assert_close(param.grad, torch.ones_like(param.grad), rtol=0, atol=0) + + def _clone_outputs(outputs): return [None if out is None else out.detach().clone() for out in outputs] @@ -3145,8 +3330,200 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(): ) -def test_grouped_linear_delay_wgrad_rejects_implicit_fallback(monkeypatch): +@pytest.mark.parametrize( + "use_grouped_tensor,single_weight,single_bias", + [ + (False, False, False), + (True, False, False), + (True, True, False), + (True, False, True), + (True, True, True), + ], + ids=["split", "grouped-discrete", "single-weight", "single-bias", "single-both"], +) +@pytest.mark.parametrize( + "recipe_name", ["bf16", "delayed", "current", "mxfp8", "high_precision", "dequantized"] +) +@pytest.mark.parametrize("freeze_weight", [False, True]) +@pytest.mark.parametrize("input_requires_grad", [False, True]) +@pytest.mark.parametrize("delay_wgrad", [False, True]) +@pytest.mark.parametrize("grad_storage", ["grad", "preallocated_grad", "main_grad"]) +def test_grouped_linear_bias_gradient_accumulation( + use_grouped_tensor, + single_weight, + single_bias, + recipe_name, + freeze_weight, + input_requires_grad, + delay_wgrad, + grad_storage, +): + """Bias gradients accumulate during main backward, independently of deferred/frozen weights.""" + if recipe_name == "bf16": + fp8_recipe = None + elif recipe_name == "delayed": + fp8_recipe = recipe.DelayedScaling() + elif recipe_name == "current": + fp8_recipe = recipe.Float8CurrentScaling() + # Use exact power-of-two scales for the analytical FP32 main_grad reference. + fp8_recipe.fp8_quant_fwd_inp = recipe.QParams(power_2_scale=True) + fp8_recipe.fp8_quant_fwd_weight = recipe.QParams(power_2_scale=True) + fp8_recipe.fp8_quant_bwd_grad = recipe.QParams(power_2_scale=True) + else: + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + fp8_recipe = recipe.MXFP8BlockScaling( + backward_override=None if recipe_name == "mxfp8" else recipe_name + ) + if fp8_recipe is not None and not fp8_available: + pytest.skip(reason_for_no_fp8) + if use_grouped_tensor and not is_module_grouped_tensor_path_supported( + fp8_recipe, torch.bfloat16 + ): + pytest.skip("Recipe is unsupported by the native grouped-tensor path") + + FP8GlobalStateManager.reset() + module = GroupedLinear( + 2, + 128, + 128, + bias=True, + params_dtype=torch.bfloat16, + use_grouped_tensor=use_grouped_tensor, + single_grouped_weight=single_weight, + single_grouped_bias=single_bias, + delay_wgrad_compute=delay_wgrad, + fuse_wgrad_accumulation=grad_storage == "main_grad", + ) + weights = _grouped_linear_weight_params(module) + biases = _grouped_linear_bias_params(module) + with torch.no_grad(): + for param in weights: + data = param.rowwise_data if isinstance(param, GroupedTensor) else param + data.fill_(0.125) + param.requires_grad_(not freeze_weight) + for param in biases: + data = param.rowwise_data if isinstance(param, GroupedTensor) else param + data.fill_(0.5) + + hook_calls = {id(param): 0 for param in module.parameters()} + flat_main_grad = torch.zeros(2 * 128 * 128, device="cuda", dtype=torch.float32) + weight_views = flat_main_grad.view(2, 128, 128) + weight_views = [weight_views] if single_weight else list(weight_views.unbind()) + main_grad_ptrs = [] + if grad_storage == "main_grad": + # Model MCore's ownership: weight GEMM accumulates into an external flat buffer, + # while parameter hooks move unfused grads to main_grad and clear the handoff. + for param, view in zip(weights, weight_views): + param.main_grad = view + param.grad_added_to_main_grad = False + param.overwrite_main_grad = False + main_grad_ptrs.append(view.data_ptr()) + for param in biases: + param.main_grad = torch.zeros(param.shape, device="cuda", dtype=torch.float32) + param.grad_added_to_main_grad = False + + def accumulate(param): + hook_calls[id(param)] += 1 + if param.grad is not None and not param.grad_added_to_main_grad: + param.main_grad.add_(param.grad) + param.grad = None + + for param in module.parameters(): + if not param.requires_grad: + continue + if getattr(param, "skip_backward_post_hook", False): + module.register_wgrad_accumulation_and_reduce_hooks( + lambda param=param: accumulate(param) + ) + else: + param.register_post_accumulate_grad_hook(accumulate) + elif grad_storage == "preallocated_grad": + for param in biases: + param.grad = torch.zeros(param.shape, device="cuda", dtype=param.dtype) + + splits = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + total_scale = 0 + try: + # Decreasing dY avoids delayed-scaling overflow from a larger second amax. + # Powers of two make output and gradients exactly representable in all tested recipes. + for step, scale in enumerate((2, 1)): + total_scale += scale + for param in weights: + param.grad = None + x = torch.full( + (256, 128), + 0.25, + device="cuda", + dtype=torch.bfloat16, + requires_grad=input_requires_grad, + ) + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + y = module(x, splits, is_first_microbatch=step == 0) + torch.testing.assert_close(y, torch.full_like(y, 4.5), rtol=0, atol=0) + y.backward(torch.full_like(y, scale)) + if input_requires_grad: + torch.testing.assert_close( + x.grad, torch.full_like(x.grad, 16 * scale), rtol=0, atol=0 + ) + else: + assert x.grad is None + + def check_bias_grads(): + for param in biases: + grad = param.main_grad if grad_storage == "main_grad" else param.grad + assert grad is not None, "bias gradient was not returned by main backward" + torch.testing.assert_close( + grad, torch.full_like(grad, 128 * total_scale), rtol=0, atol=0 + ) + if grad_storage == "main_grad": + assert param.grad is None + assert hook_calls[id(param)] == step + 1 + + check_bias_grads() + if delay_wgrad: + module.backward_dw() + check_bias_grads() # Deferred weight hooks must not add dbias a second time. + for param in weights: + if freeze_weight: + assert param.grad is None + continue + grad = param.main_grad if grad_storage == "main_grad" else param.grad + assert grad is not None + expected = 32 * (total_scale if grad_storage == "main_grad" else scale) + torch.testing.assert_close(grad, torch.full_like(grad, expected), rtol=0, atol=0) + if grad_storage == "main_grad": + assert [param.main_grad.data_ptr() for param in weights] == main_grad_ptrs + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.parametrize("device", ["cuda", "meta"]) +@pytest.mark.parametrize("single_weight", [False, True]) +@pytest.mark.parametrize("single_bias", [False, True]) +def test_grouped_linear_delayed_wgrad_hook_ownership(device, single_weight, single_bias): + """Only weight hooks are deferred, including before meta parameters are materialized.""" + module = GroupedLinear( + 2, + 128, + 128, + bias=True, + device=device, + params_dtype=torch.bfloat16, + use_grouped_tensor=True, + single_grouped_weight=single_weight, + single_grouped_bias=single_bias, + delay_wgrad_compute=True, + ) + for name, param in module.named_parameters(): + assert getattr(param, "skip_backward_post_hook", False) == name.startswith("weight") + + +@pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) +def test_grouped_linear_delay_wgrad_rejects_implicit_fallback(monkeypatch, backward_override): """Delayed wgrad reports when a grouped-tensor request used the legacy path.""" + if backward_override is not None and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) monkeypatch.setattr( "transformer_engine.pytorch.module.grouped_linear.is_module_grouped_tensor_path_supported", lambda *_args, **_kwargs: False, @@ -3161,10 +3538,14 @@ def test_grouped_linear_delay_wgrad_rejects_implicit_fallback(monkeypatch): delay_wgrad_compute=True, use_grouped_tensor=True, ) - x = torch.randn(16, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) - m_splits = torch.tensor([8, 8], dtype=torch.int64, device="cuda") + x = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + m_splits = torch.tensor([32, 32], dtype=torch.int64, device="cuda") - grouped_linear(x, m_splits).sum().backward() + with autocast( + enabled=backward_override is not None, + recipe=recipe.MXFP8BlockScaling(backward_override=backward_override), + ): + grouped_linear(x, m_splits).sum().backward() with pytest.raises( RuntimeError, match="implicit fallback is unsupported with delay_wgrad_compute=True", diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index f5ce1f8ef83..ec6012f8456 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -13,6 +13,7 @@ import sys import types from typing import Optional +from unittest.mock import Mock import pytest @@ -489,6 +490,87 @@ def test_meta_single_grouped_weight_with_delayed_wgrad(self) -> None: assert grouped_weight.skip_backward_post_hook + @pytest.mark.parametrize("num_groups", [1, 8]) + @pytest.mark.parametrize("accumulate_into_main_grad", [False, True]) + @pytest.mark.parametrize("delay_wgrad_compute", [False, True]) + def test_wgrad_storage_and_hooks( + self, + monkeypatch: pytest.MonkeyPatch, + num_groups: int, + accumulate_into_main_grad: bool, + delay_wgrad_compute: bool, + ) -> None: + """Check independent wgrad storage, external main_grad views, and delayed hooks.""" + from transformer_engine.pytorch.ops.basic import grouped_linear as gl + + if not is_op_fuser_grouped_tensor_path_supported(None, torch.bfloat16): + pytest.skip("Native grouped GEMM is unavailable.") + op = te.ops.GroupedLinear( + num_groups, + 128, + 128, + bias=False, + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + model = te.ops.Sequential(op) + flat = torch.full((num_groups, 128, 128), 0.25, device="cuda", dtype=torch.float32) + weights = [getattr(op, f"weight{i}") for i in range(num_groups)] + for i, weight in enumerate(weights): + with torch.no_grad(): + weight.fill_(1) + if accumulate_into_main_grad: + # Use views of an external flat buffer to check that GEMM preserves aliases. + weight.main_grad = flat[i] + weight.overwrite_main_grad = False + weight.grad_added_to_main_grad = False + x = torch.full( + (num_groups * 128, 128), 0.125, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + splits = torch.full((num_groups,), 128, device="cuda", dtype=torch.int64) + seen = [] + original = gl.general_grouped_gemm_for_grouped_tensor + + def gemm(a, b, out, **kwargs): + if kwargs.get("layout") == "NT": + assert len(out) == num_groups + num_storages = len({g.untyped_storage().data_ptr() for g in out}) + assert num_storages == (1 if accumulate_into_main_grad else num_groups) + assert len({g.data_ptr() for g in out}) == num_groups + if accumulate_into_main_grad: + assert out[0].data_ptr() == flat.data_ptr() + else: + assert all(g._base is None for g in out) + seen.append(out) + return original(a, b, out, **kwargs) + + monkeypatch.setattr(gl, "general_grouped_gemm_for_grouped_tensor", gemm) + hook = Mock() + op.register_wgrad_accumulation_and_reduce_hooks(hook) + y = model(x, splits) + torch.testing.assert_close(y, torch.full_like(y, 16), rtol=0, atol=0) + y.backward(torch.full_like(y, 0.5)) + if delay_wgrad_compute: + assert not seen + hook.assert_not_called() + op.backward_dw() + hook.assert_called_once() + assert len(seen) == 1 + if accumulate_into_main_grad: + MegatronTrainingHelper.verify_main_grad_accumulation( + weights, + expected_main_grads=[torch.full_like(flat[i], 8.25) for i in range(num_groups)], + ) + for i, weight in enumerate(weights): + assert weight.main_grad.data_ptr() == flat[i].data_ptr() + else: + for weight in weights: + torch.testing.assert_close( + weight.grad, torch.full_like(weight.grad, 8), rtol=0, atol=0 + ) + def test_single_grouped_bias_uses_registered_packed_storage(self) -> None: """The grouped bias compute view must alias the registered trainable parent.""" op = te.ops.GroupedLinear( diff --git a/tests/pytorch/test_module_grouped_linear_distributed_weight.py b/tests/pytorch/test_module_grouped_linear_distributed_weight.py index 05fba219a99..2a664084c29 100644 --- a/tests/pytorch/test_module_grouped_linear_distributed_weight.py +++ b/tests/pytorch/test_module_grouped_linear_distributed_weight.py @@ -20,6 +20,9 @@ a specific assertion rather than drifting numerically. """ +import gc +import weakref + import pytest import torch @@ -283,3 +286,93 @@ def _finalize(wgrads, **kwargs): assert calls == ["fwd", "bwd", "finalize"], f"unexpected hook dispatch: {calls}" assert torch.count_nonzero(weight.wgrad_scratch) > 0, "wgrad never reached grad_buffer" + + +@pytest.mark.parametrize("single_weight", [False, True]) +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("fused", [False, True]) +@pytest.mark.parametrize("requires_dgrad", [False, True]) +def test_gathered_weights_are_transient( + monkeypatch, single_weight, quantized, fused, requires_dgrad +): + """No gathered copy survives forward; backward must gather even when dgrad is disabled.""" + from transformer_engine.pytorch.module.grouped_linear import ( + _GroupedLinear, + is_module_grouped_tensor_path_supported, + ) + from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor + + fp8_recipe = MXFP8BlockScaling() if quantized else None + if not is_module_grouped_tensor_path_supported(fp8_recipe, DTYPE): + pytest.skip("Requested native grouped GEMM is unavailable.") + te.quantization.FP8GlobalStateManager.reset() + module = te.GroupedLinear( + 2, + IN_F, + OUT_F, + bias=False, + params_dtype=DTYPE, + device=DEVICE, + single_grouped_weight=single_weight, + use_grouped_tensor=True, + fuse_wgrad_accumulation=fused, + ) + weights = [module.weight] if single_weight else [module.weight0, module.weight1] + for weight in weights: + weight.main_grad = torch.zeros(weight.shape, device=DEVICE, dtype=torch.float32) + weight.wgrad_scratch = torch.zeros_like(weight.main_grad) + weight.grad_buffer = lambda w=weight: w.wgrad_scratch + gathered_refs, calls = [], [] + + def gather(phase): + calls.append(phase) + if single_weight: + copy = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=2, + tensor_shape=(OUT_F, IN_F), + rowwise_data=weights[0].rowwise_data.detach().clone(), + dtype=DTYPE, + ) + copies = [copy] + else: + copies = [w.detach().clone() for w in weights] + assert all(not w.requires_grad for w in copies) + gathered_refs.extend(weakref.ref(w) for w in copies) + return copies + + def finalize(grads, **kwargs): + calls.append("finalize") + grads = [grads] if single_weight else grads + if fused: + for weight, grad in zip(weights, grads): + assert grad.data_ptr() == weight.wgrad_scratch.data_ptr() + weight.main_grad.add_(grad) + return [None] * len(weights) + return list(grads) + + leader = weights[0] + leader.is_distributed_weight = True + leader.materialize_group_for_forward = lambda: gather("fwd") + leader.materialize_group_for_backward = lambda **kw: gather("bwd") + leader.finalize_group_grads = finalize + prep = [] + original = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm + + def prepare(*args, **kwargs): + prep.append(kwargs["cache_weight"]) + return original(*args, **kwargs) + + monkeypatch.setattr(_GroupedLinear, "_prepare_weights_for_grouped_tensor_gemm", prepare) + splits, x = _inputs(2) + x.requires_grad_(requires_dgrad) + with te.autocast(enabled=quantized, recipe=fp8_recipe): + y = module(x, splits, is_first_microbatch=True) + gc.collect() + assert not module._fp8_workspaces + assert all(ref() is None for ref in gathered_refs) + y.sum().backward() + assert calls == ["fwd", "bwd", "finalize"] + assert prep == [False] * (1 + int(requires_dgrad)) + for weight in weights: + grad = weight.main_grad if fused else weight.grad + assert grad is not None and torch.count_nonzero(grad) > 0 diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 4e64215cdf5..6cd1bd403b1 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -25,6 +25,7 @@ QuantizedTensor, ) +from transformer_engine.pytorch.quantized_tensor import prepare_for_saving, restore_from_saved from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported from references.ref_per_tensor_cs import ref_per_tensor_cs_cast @@ -1187,3 +1188,48 @@ def clone_quantized_output(): assert torch.equal( buf_without_noop, buf_with_noop ), f"{quantization}/{usage}: noop flag fails to take effect because {attr} changed." + + +def test_saved_storage_protocol(monkeypatch): + """Preserve custom metadata, pass the full suffix, and use the returned remainder.""" + value, tail, replacement = (torch.tensor(i) for i in range(3)) + remainder = [replacement, tail] + + class Storage: + __slots__ = () + + def prepare_for_saving(self): + return [value], self + + def restore_from_saved(self, tensors): + assert isinstance(tensors, list) + assert len(tensors) == 2 and tensors[0] is value and tensors[1] is tail + return remainder + + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False) + storage = Storage() + saved, objects = prepare_for_saving(value, storage) + restored, remaining = restore_from_saved( + objects + [None], saved + [tail], return_saved_tensors=True + ) + assert restored[0] is value and restored[1] is storage and restored[2] is replacement + assert isinstance(remaining, list) and len(remaining) == 1 and remaining[0] is tail + + +@pytest.mark.parametrize("compiled_metadata", [False, True]) +def test_restore_real_quantized_storage(monkeypatch, compiled_metadata): + """Restore actual MXFP8 buffers and None/plain tensor slots with either metadata form.""" + if not te.is_mxfp8_available(): + pytest.skip("MXFP8 is unavailable.") + q = MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=True) + q.internal = True + x = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) + quantized = q(x) + expected = quantized.dequantize().clone() + with monkeypatch.context() as context: + context.setattr(torch.compiler, "is_compiling", lambda: compiled_metadata) + saved, objects = prepare_for_saving(x, quantized, None) + restored, tail = restore_from_saved(objects, tuple(saved) + (x,), return_saved_tensors=True) + assert restored[0] is x and restored[2] is None + assert isinstance(tail, tuple) and tail[0] is x + torch.testing.assert_close(restored[1].dequantize(), expected, rtol=0, atol=0) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 8acc547c821..10c27399c01 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -13,7 +13,13 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.device import te_device_type from ..constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType -from ..utils import ceil_div, get_cached_ones_tensor, get_sm_count, _empty_tensor +from ..utils import ( + ceil_div, + get_cached_ones_tensor, + get_device_compute_capability, + get_sm_count, + _empty_tensor, +) from ..quantized_tensor import QuantizedTensorStorage, Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer @@ -1189,7 +1195,7 @@ def _is_fp8_blockwise(operand) -> bool: # Hopper (SM90) uses a single shared alpha/beta scalar; # Blackwell+ (SM100) supports per-group alpha/beta arrays. - per_group = torch.cuda.get_device_capability() >= (10, 0) + per_group = get_device_compute_capability() >= (10, 0) num_alphabeta = num_tensors if per_group else 1 if alpha is None: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index a2e1c4210b7..c22a820ed22 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -260,6 +260,7 @@ class GroupedLinearBwdArgs: m_splits: Optional[List[int]] = None num_gemms: int = 0 + num_bias_args: int = 0 weights_shape_1: int = 0 use_bias: bool = False @@ -291,11 +292,15 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: self.inputmats = saved[:n] self.weights_fp8 = saved[n : 2 * n] self.saved_weights = saved[2 * n : 3 * n] - self.biases = saved[3 * n : 4 * n] + self.biases = saved[3 * n : 3 * n + self.num_bias_args] -def _grouped_linear_configure_quantizers(args: GroupedLinearFwdArgs, weights) -> None: - """Configure operand usage after materializing weights.""" +def _grouped_linear_configure_split_quantizers(args: GroupedLinearFwdArgs, weights) -> None: + """Configure per-expert quantizers for the split-quantize path after materialization. + + The grouped-tensor path configures its input quantizer before group_quantize + and its weight quantizers in _prepare_weights_for_grouped_tensor_gemm. + """ if args.input_quantizers[0] is not None: for input_quantizer in args.input_quantizers: input_quantizer.set_usage( @@ -321,6 +326,11 @@ def _grouped_linear_configure_quantizers(args: GroupedLinearFwdArgs, weights) -> weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) elif isinstance(weights[0], QuantizedTensorStorage): args.weight_quantizers = [weight._quantizer for weight in weights] + for weight_quantizer in args.weight_quantizers: + weight_quantizer.set_usage( + rowwise=True, + columnwise=True if columnwise_usage else None, + ) if args.output_quantizers[0] is not None: for output_quantizer in args.output_quantizers: output_quantizer.set_usage(rowwise=True, columnwise=False) @@ -352,7 +362,7 @@ def _grouped_linear_forward_impl( if is_dist_weight: weights = materialize_weight_for_forward(weights) - _grouped_linear_configure_quantizers(args, weights) + _grouped_linear_configure_split_quantizers(args, weights) weight_quantizers = args.weight_quantizers # Initialize input tensors @@ -486,6 +496,7 @@ def _grouped_linear_setup_ctx( bwd_args.m_splits = fwd_args.m_splits bwd_args.num_gemms = num_gemms + bwd_args.num_bias_args = len(fwd_args.biases) bwd_args.weights_shape_1 = fwd_args.inp.shape[-1] bwd_args.use_bias = fwd_args.use_bias @@ -588,6 +599,12 @@ def _grouped_linear_backward_impl( # Preprocess grad output out_features = grad_output.shape[-1] grad_output_view = grad_output.contiguous().view(-1, out_features) + delay_wgrad = args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute() + # Bias gradients belong to the main backward even when weights are frozen + # or their GEMM is deferred. Keep fused BF16 wgrad+dbias for immediate wgrad. + compute_dbias = args.use_bias and ( + args.fp8 or args.debug or not args.weights_requires_grad or delay_wgrad + ) grad_output_reference = args.grad_output_quantizers[0] if args.fp8 and isinstance(grad_output_reference, HybridQuantizer): # Usage is a runtime decision, not part of generation validation. @@ -604,7 +621,7 @@ def _grouped_linear_backward_impl( args.grad_output_quantizers, args.activation_dtype, with_quantized_output=args.fp8 or args.debug, - compute_dbias=(args.fp8 or args.debug) and args.use_bias, + compute_dbias=compute_dbias, disable_bulk_allocation=( args.cpu_offloading and isinstance(grad_output_reference, HybridQuantizer) @@ -675,11 +692,7 @@ def _grouped_linear_backward_impl( ) if args.weights_requires_grad: - if ( - is_dist_weight - and args.wgrad_store is not None - and args.wgrad_store.delay_wgrad_compute() - ): + if is_dist_weight and delay_wgrad: raise RuntimeError( "distributed-weight GroupedLinear requires delay_wgrad_compute=False." ) @@ -733,7 +746,9 @@ def _grouped_linear_backward_impl( layout="NT", grad=True, m_splits=m_splits, - use_bias=args.use_bias if grad_biases[0] is None else None, + # Preserve the high-precision GEMM epilogue and algorithm selection, + # including when delayed wgrad no longer owns bias accumulation. + use_bias=args.use_bias and not (args.fp8 or args.debug), bias=biases, use_split_accumulator=wgrad_use_split_accumulator, accumulate=( @@ -743,14 +758,12 @@ def _grouped_linear_backward_impl( ), ) # WGRAD - if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + if delay_wgrad: args.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) else: _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) - - for i in range(num_gemms): - if grad_biases[i] is None: - grad_biases[i] = grad_biases_[i] + if not compute_dbias: + grad_biases = grad_biases_ del grad_biases_ clear_tensor_data(*inputmats) @@ -765,10 +778,8 @@ def _grouped_linear_backward_impl( else: wgrad_list = [None] * num_gemms - if not args.use_bias or ( - args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute() and not args.fp8 - ): - grad_biases = [None] * num_gemms + if not args.use_bias: + grad_biases = [None] * args.num_bias_args dgrad_out = None if args.requires_dgrad: @@ -794,6 +805,7 @@ class GroupedLinearGraphSafeBwdArgs: grad_output_quantizers: List[Quantizer] = None is_dist_weight: bool = False num_gemms: int = 0 + num_bias_args: int = 0 in_features: int = 0 out_features: int = 0 activation_dtype: Optional[torch.dtype] = None @@ -846,7 +858,6 @@ def _grouped_linear_graph_safe_forward(args: GroupedLinearFwdArgs) -> Tuple[Any, is_dist_weight = is_distributed_weight(weights[0]) if is_dist_weight: weights = materialize_weight_for_forward(weights) - _grouped_linear_configure_quantizers(args, weights) weight_quantizers = args.weight_quantizers biases = args.biases out = args.out @@ -903,7 +914,9 @@ def _grouped_linear_graph_safe_forward(args: GroupedLinearFwdArgs) -> Tuple[Any, dtype=activation_dtype, ) - columnwise_usage = is_grad_enabled and args.input_requires_grad + columnwise_usage = (is_grad_enabled and args.input_requires_grad) or ( + is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase() + ) weights_for_gemm, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm( weights, weight_quantizers, @@ -1005,6 +1018,7 @@ def _grouped_linear_graph_safe_setup( bwd_args.grad_output_quantizers = fwd_args.grad_output_quantizers bwd_args.m_splits_tensor = fwd_args.m_splits_tensor bwd_args.num_gemms = fwd_args.num_gemms + bwd_args.num_bias_args = len(fwd_args.biases) bwd_args.activation_dtype = fwd_args.activation_dtype bwd_args.fp8 = fwd_args.fp8 bwd_args.use_bias = fwd_args.use_bias @@ -1110,7 +1124,9 @@ def _grouped_linear_graph_safe_backward( grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) dy_2d = cast_if_needed(grad_output_view, args.activation_dtype) dbias_packed = None - if args.fp8: + grouped_dy = None + needs_grouped_dy = args.requires_dgrad or args.weights_requires_grad + if args.fp8 and needs_grouped_dy: grad_output_quantizer = args.grad_output_quantizers[0] grad_output_quantizer.set_usage( rowwise=args.requires_dgrad, @@ -1138,7 +1154,7 @@ def _grouped_linear_graph_safe_backward( split_sizes, tensor_offsets=output_tensor_offsets, ) - else: + elif needs_grouped_dy: grouped_dy = _GroupedLinear._make_grouped_tensor( dy_2d, num_gemms=N, @@ -1156,8 +1172,7 @@ def _grouped_linear_graph_safe_backward( else: grad_bias_args = [dbias_packed[i].to(dtype=args.activation_dtype) for i in range(N)] else: - num_bias_args = 1 if args.single_grouped_bias else N - grad_bias_args = [None] * num_bias_args + grad_bias_args = [None] * args.num_bias_args dgrad = None if args.requires_dgrad: @@ -1265,7 +1280,7 @@ def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights): use_split_accumulator=wgrad_use_split_accumulator, accumulate=accumulate, ) - return None, [None] * N, None + return None, (), None if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): args.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], grouped_gemm_wgrad) @@ -1393,6 +1408,11 @@ def _prepare_weights_for_grouped_tensor_gemm( raise RuntimeError( "Quantized single grouped weights require quantized grouped GEMM compute." ) + # Preserve usages already enabled on the parameter quantizer. + weight.quantizer.set_usage( + rowwise=True, + columnwise=True if columnwise_usage else None, + ) return weight, new_workspaces if not with_quantized_compute: @@ -1472,8 +1492,15 @@ def _prepare_weights_for_grouped_tensor_gemm( update_ws = is_first_microbatch is None or is_first_microbatch for idx, weight in enumerate(weights): - weight_quantizer = weight_quantizers[idx] - weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + if isinstance(weight, QuantizedTensorStorage): + weight_quantizer = weight._quantizer + weight_quantizer.set_usage( + rowwise=True, + columnwise=True if columnwise_usage else None, + ) + else: + weight_quantizer = weight_quantizers[idx] + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) weight_fp8, new_workspaces[idx] = quantize_weight( tensor=weight, quantizer=weight_quantizer, @@ -1569,7 +1596,7 @@ def _forward_grouped_tensor(ctx, fwd_args: GroupedLinearFwdArgs) -> Tuple[torch. ctx.inp_shape = fwd_args.inp.shape ctx.backward_objects = bwd_args if fwd_args.fp8 and requires_grad( - fwd_args.inp, fwd_args.weights[0], fwd_args.biases[0] + fwd_args.inp, fwd_args.weights[0], fwd_args.biases[0] if fwd_args.biases else None ): bwd_args.reduce_and_update_bwd_fp8_tensors = ( FP8GlobalStateManager.is_first_fp8_module() @@ -1593,6 +1620,7 @@ def forward( fwd_args.inp = inp num_weight_args = 1 if fwd_args.single_grouped_weight else num_gemms fwd_args.weights = list(weights_and_biases[:num_weight_args]) + # Only bias=False omits bias inputs; returned biases retain their autograd edges. fwd_args.biases = list(weights_and_biases[num_weight_args:]) if fwd_args.use_grouped_tensor_path: @@ -1607,7 +1635,9 @@ def forward( ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects ctx.backward_objects = bwd_args - if fwd_args.fp8 and requires_grad(inp, fwd_args.weights[0], fwd_args.biases[0]): + if fwd_args.fp8 and requires_grad( + inp, fwd_args.weights[0], fwd_args.biases[0] if fwd_args.biases else None + ): bwd_args.reduce_and_update_bwd_fp8_tensors = ( FP8GlobalStateManager.is_first_fp8_module() ) @@ -1689,7 +1719,8 @@ class GroupedLinear(TransformerEngineBaseModule): the model is trained with lower precision and the original FP32 parameters would not fit in GPU memory. delay_wgrad_compute : bool, default = False - Whether to delay weight gradient computation + Whether to delay weight gradient computation until ``backward_dw()``. + Bias gradients are computed and accumulated during the main backward. save_original_input : bool, default = False If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, @@ -1886,13 +1917,13 @@ def __init__( self.reset_parameters(defer_init=is_meta) if self.wgrad_store.delay_wgrad_compute(): - for name, param in self.named_parameters(): - if name in ("weight", "bias"): - param.skip_backward_post_hook = True - continue + # Only weights are deferred. Biases use normal autograd accumulation hooks. + grouped_weight = getattr(self, "weight", None) + if grouped_weight is not None: + grouped_weight.skip_backward_post_hook = True + else: for i in range(self.num_gemms): - if name in (f"weight{i}", f"bias{i}"): - param.skip_backward_post_hook = True + getattr(self, f"weight{i}").skip_backward_post_hook = True def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" @@ -2028,7 +2059,7 @@ def make_grouped_weights(self, defer_init=False) -> None: grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=self.num_gemms, shapes=[(self.out_features, self.in_features)] * self.num_gemms, - quantizer=weight_quantizers[0], + quantizer=weight_quantizers[0].copy() if weight_quantizers[0] is not None else None, dtype=self.params_dtype, device=weights[0].device, ) @@ -2302,9 +2333,9 @@ def forward( 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() num_gemms = self.num_gemms + cpu_offloading = is_cpu_offload_enabled() # Make sure splits are in expected format if (self.single_grouped_weight or self.single_grouped_bias) and not self.use_grouped_tensor: @@ -2324,11 +2355,10 @@ def forward( f"does not match number of GEMMs ({num_gemms})." ) - if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = ( - FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor - ) - else: + skip_fp8_weight_update = ( + FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor + ) + if skip_fp8_weight_update is not None and not FP8GlobalStateManager.fp8_graph_capturing(): skip_fp8_weight_update = None if skip_fp8_weight_update is not None: is_first_microbatch = False @@ -2339,16 +2369,22 @@ def forward( inp = self.prepare_forward(inp, num_gemms=self.num_gemms) try: + _recipe = FP8GlobalStateManager.get_fp8_recipe() if self.fp8 else None + weight_tensors = self._get_weight_tensors() - bias_tensors = self._get_bias_tensors() + # Keep real bias inputs even when returned for the caller to apply: they + # preserve the output's autograd dependency when input/weight are frozen. + bias_tensors = self._get_bias_tensors() if self.use_bias else [] use_grouped_bias = self.use_bias and self.single_grouped_bias - quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() - + # Check whether to enable debug features + debug = self.is_debug_iter() + debug_quantizers = None if debug: - if self.no_debug_features_active(list(chain(*quantizers))): + debug_quantizers = self._get_debug_quantizers() + if self.no_debug_features_active(list(chain(*debug_quantizers))): debug = False - quantizers = self._get_quantizers() + debug_quantizers = None if debug and (self.single_grouped_weight or self.single_grouped_bias): raise RuntimeError( "TE debug features do not support single grouped parameters. DebugQuantizer " @@ -2357,6 +2393,35 @@ def forward( "features for this GroupedLinear." ) + # Choose between grouped tensor path or split-quantize path + use_grouped_tensor_path = False + if self.use_grouped_tensor and not (self.fp8_calibration or debug or cpu_offloading): + use_grouped_tensor_path = is_module_grouped_tensor_path_supported( + _recipe, + self.activation_dtype, + device=inp.device, + ) + if ( + self.use_grouped_tensor + and not use_grouped_tensor_path + and (self.single_grouped_weight or use_grouped_bias) + ): + raise RuntimeError( + "Single grouped parameters require the native grouped-tensor path, but the" + " active device, cuBLASLt version, quantization recipe, or GroupedLinear" + " feature configuration does not support it. Disable single_grouped_weight and" + " single_grouped_bias to allow the split-quantize fallback." + ) + + # Get quantizers + if use_grouped_tensor_path: + quantizers = self._get_quantizers_for_grouped_tensor_path() + elif debug: + if debug_quantizers is None: + debug_quantizers = self._get_debug_quantizers() + quantizers = debug_quantizers + else: + quantizers = self._get_quantizers() ( input_quantizers, weight_quantizers, @@ -2365,6 +2430,7 @@ def forward( grad_weight_quantizers, grad_output_quantizers, ) = quantizers + if not debug and weight_quantizers[0] is not None: # Experts share shape and recipe settings: compute once and broadcast. optimize_for_gemm = self._enable_weight_preswizzle( @@ -2385,10 +2451,8 @@ def forward( weight_requires_grad = weight_tensors[0].requires_grad if self.fp8: - _recipe = FP8GlobalStateManager.get_fp8_recipe() backward_override = _recipe.backward_override else: - _recipe = None backward_override = None save_original_input = self.save_original_input @@ -2425,30 +2489,6 @@ def forward( ) save_original_input = False - cpu_offloading = is_cpu_offload_enabled() - use_grouped_tensor_path = False - if self.use_grouped_tensor and not ( - self.fp8_calibration - or debug - or cpu_offloading - or any(q is not None for q in output_quantizers) - ): - use_grouped_tensor_path = is_module_grouped_tensor_path_supported( - _recipe, - self.activation_dtype, - device=inp.device, - ) - if ( - self.use_grouped_tensor - and not use_grouped_tensor_path - and (self.single_grouped_weight or use_grouped_bias) - ): - raise RuntimeError( - "Single grouped parameters require the native grouped-tensor path, but the" - " active device, cuBLASLt version, quantization recipe, or GroupedLinear" - " feature configuration does not support it. Disable single_grouped_weight and" - " single_grouped_bias to allow the split-quantize fallback." - ) if use_grouped_tensor_path: if m_splits.device.type != te_device_type(): raise ValueError( @@ -2521,6 +2561,9 @@ def forward( self.end_forward() if self.return_bias: + # bias=False still returns the legacy list of empty tensors when requested. + if not self.use_bias: + bias_tensors = self._get_bias_tensors() if use_grouped_bias: return out, bias_tensors[0] return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors] @@ -2536,7 +2579,7 @@ def backward_dw(self): if self.wgrad_store.context is None or self.wgrad_store.context.empty(): return with get_nvtx_range_context("_GroupedLinear_wgrad"): - (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() + (_, grad_biases, _), tensor_list = self.wgrad_store.pop() wgrad_output = tensor_list[2] weight_params = self._get_weight_tensors() if not self.fuse_wgrad_accumulation: @@ -2547,23 +2590,21 @@ def backward_dw(self): else: for i in range(self.num_gemms): weight_params[i].grad = wgrad_output[i].to(weight_params[i].dtype) - has_grad_biases = [ - grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_ - ] - if self.use_bias and any(has_grad_biases): - if self.use_grouped_tensor: - raise RuntimeError( - "GroupedLinear(use_grouped_tensor=True) fell back to the split-quantize " - "path, which produced per-expert bias gradients during delayed wgrad. " - "This implicit fallback is unsupported with delay_wgrad_compute=True. " - "Use a configuration supported by the grouped-tensor path, or set " - "use_grouped_tensor=False to select the legacy path explicitly." - ) - bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] - for i in range(self.num_gemms): - if has_grad_biases[i] and bias_params[i].grad is None: - bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) - del grad_biases_ + if ( + self.use_grouped_tensor + and self.apply_bias + and any( + grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases + ) + ): + raise RuntimeError( + "GroupedLinear(use_grouped_tensor=True) fell back to the split-quantize " + "path, which produced per-expert bias gradients during delayed wgrad. " + "This implicit fallback is unsupported with delay_wgrad_compute=True. " + "Use a configuration supported by the grouped-tensor path, or set " + "use_grouped_tensor=False to select the legacy path explicitly." + ) + del grad_biases del wgrad_output del tensor_list self._trigger_wgrad_accumulation_and_reduce_hooks() @@ -2715,3 +2756,36 @@ def _get_debug_quantizers(self): ] for name, qs in zip(names, original_quantizers) ) + + def _get_quantizers_for_grouped_tensor_path(self): + """Get quantizers consumed by grouped tensor path. + + Grouped tensor only uses the first input quantizer and grad + output quantizer. Single grouped weight uses first weight + quantizer, while discrete weights use num_gemm weight + quantizers. The remaining quantizers are unused. + + """ + input_quantizers, grad_output_quantizers = [None], [None] + num_weights = 1 if self.single_grouped_weight else self.num_gemms + weight_quantizers = [None] * num_weights + if self.fp8 or self.primary_weights_in_fp8: + weight_quantizers = [ + self.quantizers["scaling_fwd"][ + self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"] + ] + for i in range(num_weights) + ] + for quantizer in weight_quantizers: + quantizer.internal = not self.primary_weights_in_fp8 + if self.fp8: + input_quantizer = self.quantizers["scaling_fwd"][self._offsets["input"]] + input_quantizer.internal = True + input_quantizer.optimize_for_gemm = True + input_quantizers = [input_quantizer] + if torch.is_grad_enabled(): + grad_output_quantizer = self.quantizers["scaling_bwd"][self._offsets["input"]] + grad_output_quantizer.internal = True + grad_output_quantizer.optimize_for_gemm = True + grad_output_quantizers = [grad_output_quantizer] + return input_quantizers, weight_quantizers, (), (), (), grad_output_quantizers diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 86fa740c5b6..8f3c99b38e0 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -596,7 +596,7 @@ def make_grouped_weights(self) -> None: grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=self.num_groups, shapes=[(self.out_features, self.in_features)] * self.num_groups, - quantizer=quantizer, + quantizer=quantizer.copy() if quantizer is not None else None, dtype=self.weight0.dtype, device=self.weight0.device, ) @@ -823,14 +823,13 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: ) weight_requires_grad = requires_grad and weight_requires_grad - # Configure quantizer usages - for group_idx in range(self.num_groups): - input_quantizer = self.get_quantizer("forward", 2 * group_idx) - weight_quantizer = self.get_quantizer("forward", 2 * group_idx + 1) - grad_output_quantizer = self.get_quantizer("backward", group_idx) - input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) - weight_quantizer.set_usage(rowwise=True, columnwise=requires_grad) - grad_output_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + # Configure the first group's input and grad output + # quantizers since they are exposed externally and may be + # used by other ops for early quantization + input_quantizer = self.get_quantizer("forward", 0) + grad_output_quantizer = self.get_quantizer("backward", 0) + input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + grad_output_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: super().reset_recipe_state(recipe=recipe) @@ -922,7 +921,7 @@ def op_backward(self, *args, **kwargs): def _get_grouped_weight_for_gemm( self, weight_param: GroupedTensor, - weight_quantizers: list[Optional[Quantizer]], + weight_quantizers: Sequence[Optional[Quantizer]], columnwise_usage: bool, with_quantized_compute: bool, dtype: torch.dtype, @@ -933,7 +932,9 @@ def _get_grouped_weight_for_gemm( num_groups = self.num_groups is_weight_quantized = weight_param.quantizer is not None if is_weight_quantized and with_quantized_compute: - # GGEMM can use it as it is + weight_param.quantizer.set_usage( + rowwise=True, columnwise=True if columnwise_usage else None + ) return weight_param if is_weight_quantized and not with_quantized_compute: # This use-case isnt optimized yet. Involves a per-group @@ -978,7 +979,7 @@ def _get_grouped_weight_for_gemm( def _get_discrete_weights_for_gemm( self, weight_params: Optional[GroupedTensor] | list[torch.Tensor], - weight_quantizers: list[Optional[Quantizer]], + weight_quantizers: Sequence[Optional[Quantizer]], columnwise_usage: bool, with_quantized_compute: bool, dtype: torch.dtype, @@ -990,9 +991,16 @@ def _get_discrete_weights_for_gemm( for w, quantizer in zip(weight_params, weight_quantizers): if not with_quantized_compute: w = maybe_dequantize(w, dtype) - elif with_quantized_compute and not is_quantized_tensor(w): + else: quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - w = quantizer(w) + if is_quantized_tensor(w): + # Honor the op's requirements without disabling existing parameter usages. + required_columnwise = True if columnwise_usage else None + w.update_usage(rowwise_usage=True, columnwise_usage=required_columnwise) + if w._quantizer is not None: + w._quantizer.set_usage(rowwise=True, columnwise=required_columnwise) + else: + w = quantizer(w) out.append(w) return out @@ -1076,14 +1084,7 @@ def fuser_forward( input_requires_grad = ctx.requires_grad weight_requires_grad = ctx.requires_grad and weight_param.requires_grad - # Quantizers - input_quantizers = [None] * num_groups - weight_quantizers = [None] * num_groups with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() - if with_quantized_compute: - for group_idx in range(num_groups): - input_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx) - weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) # Get autocast dtype if needed dtype = maybe_autocast_dtype(default_dtype=weight_param.dtype) @@ -1126,12 +1127,19 @@ def fuser_forward( ) if use_grouped_tensor_path: + if self.single_grouped_weight: + weight_quantizers = [self.get_quantizer("forward", 1)] + else: + weight_quantizers = [ + self.get_quantizer("forward", 2 * group_idx + 1) + for group_idx in range(num_groups) + ] out, tensors_to_save = self._fuser_forward_grouped_tensor( input_=input_, split_sizes=split_sizes, scales=scales, with_quantized_compute=with_quantized_compute, - input_quantizers=input_quantizers, + input_quantizer=self.get_quantizer("forward", 0), weight_quantizers=weight_quantizers, dtype=dtype, input_requires_grad=input_requires_grad, @@ -1140,6 +1148,12 @@ def fuser_forward( out_buffer=out_buffer, ) else: + input_quantizers = [ + self.get_quantizer("forward", 2 * group_idx) for group_idx in range(num_groups) + ] + weight_quantizers = [ + self.get_quantizer("forward", 2 * group_idx + 1) for group_idx in range(num_groups) + ] out, tensors_to_save = self._fuser_forward_split_quantize( ctx=ctx, input_=input_, @@ -1219,20 +1233,16 @@ def fuser_forward_save_ctx( num_groups = self.num_groups weight_param = self.weight if self.single_grouped_weight else self.weight0 - with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() - input_quantizers = [None] * num_groups - weight_quantizers = [None] * num_groups - grad_output_quantizers = [None] * num_groups - if with_quantized_compute: - for group_idx in range(num_groups): - input_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx) - weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) - grad_output_quantizers[group_idx] = self.get_quantizer("backward", group_idx) + if use_grouped_tensor_path: + grad_output_quantizers = [self.get_quantizer("backward", 0)] + else: + grad_output_quantizers = [ + self.get_quantizer("backward", group_idx) for group_idx in range(num_groups) + ] + with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() ctx.use_grouped_tensor_path = use_grouped_tensor_path ctx.with_quantized_compute = with_quantized_compute - ctx.input_quantizers = input_quantizers - ctx.weight_quantizers = weight_quantizers ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_input_quantizers = None # ``split_sizes``, offset metadata, and related tensors are routed @@ -1256,8 +1266,8 @@ def _fuser_forward_split_quantize( split_sizes: torch.Tensor, scales: Optional[torch.Tensor], with_quantized_compute: bool, - input_quantizers: list[Optional[Quantizer]], - weight_quantizers: list[Optional[Quantizer]], + input_quantizers: Sequence[Optional[Quantizer]], + weight_quantizers: Sequence[Optional[Quantizer]], dtype: torch.dtype, input_requires_grad: bool, weight_requires_grad: bool, @@ -1372,8 +1382,8 @@ def _fuser_forward_grouped_tensor( split_sizes: torch.Tensor, scales: Optional[torch.Tensor], with_quantized_compute: bool, - input_quantizers: list[Optional[Quantizer]], - weight_quantizers: list[Optional[Quantizer]], + input_quantizer: Optional[Quantizer], + weight_quantizers: Sequence[Optional[Quantizer]], dtype: torch.dtype, input_requires_grad: bool, weight_requires_grad: bool, @@ -1398,8 +1408,6 @@ def _fuser_forward_grouped_tensor( original_shape = list(input_.size()) prequantized_input = with_quantized_compute and isinstance(input_, GroupedTensor) if with_quantized_compute: - input_quantizer = input_quantizers[0] - input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) input_quantizer.optimize_for_gemm = True if prequantized_input: # GroupedTensor forbids reshape and is already in the canonical diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 9c44dae9bf3..204e703e1d9 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -361,28 +361,34 @@ def restore_from_saved( Note: please use `restore_from_func_ctx` instead if you are restoring tensors from a function context to make sure tensor_objects is detached and its memory can be freed """ tensor_objects = [] + offset = 0 for tensor in tensors: if tensor is None or isinstance(tensor, torch.Tensor): - tensor_objects.append(saved_tensors[0]) - saved_tensors = saved_tensors[1:] + tensor_objects.append(saved_tensors[offset]) + offset += 1 elif isinstance(tensor, _SavedTensorRef): # The saved tensor was unpacked (running autograd's version check), but # may have lost its subclass under hooks; use the original object. tensor_objects.append(tensor.weight) - saved_tensors = saved_tensors[1:] + offset += 1 elif isinstance(tensor, _SavedQuantizedTensor): count = len(tensor.inner_names) - inner = dict(zip(tensor.inner_names, saved_tensors[:count])) + inner = dict(zip(tensor.inner_names, saved_tensors[offset : offset + count])) tensor_objects.append( QuantizedTensorStorage.__tensor_unflatten__(inner, tensor.metadata, None, None) ) - saved_tensors = saved_tensors[count:] + offset += count else: + # Storage restorers receive the full remaining list/tuple and own the + # returned remainder. Preserve that protocol for custom implementations. + if offset: + saved_tensors = saved_tensors[offset:] saved_tensors = tensor.restore_from_saved(saved_tensors) + offset = 0 tensor_objects.append(tensor) if return_saved_tensors: - return tensor_objects, saved_tensors + return tensor_objects, saved_tensors[offset:] if offset else saved_tensors return tensor_objects From 3fda6d9863881d3018b01d3ce8e92dcb505a7a6d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 30 Oct 2026 13:18:34 +0100 Subject: [PATCH 28/32] Fix captured FP8 state and share normalization output layout policy Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 14 ++++-- transformer_engine/pytorch/csrc/extensions.h | 4 +- .../pytorch/csrc/extensions/normalization.cpp | 43 ++++++++++++------- .../pytorch/csrc/extensions/pybind.cpp | 5 ++- transformer_engine/pytorch/dynamo/fake.py | 26 ++++++----- .../pytorch/module/layernorm_linear.py | 13 +++--- .../pytorch/module/layernorm_mlp.py | 22 +++++----- transformer_engine/pytorch/module/linear.py | 34 ++++++++------- 8 files changed, 93 insertions(+), 68 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 66c925a25a2..d7a3890b66b 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -2076,13 +2076,13 @@ def test_to_tensor_spec_quantized(factory, shape, swizzled): def test_fake_norm_mxfp8_scale_layout( monkeypatch, normalization, with_2d_quantization, optimize_for_gemm, use_cudnn, shape, compact ): - """RMSNorm's 2D MXFP8 fallback preserves the quantizer's scale layout.""" + """2D MXFP8 fallback preserves the quantizer's scale layout.""" monkeypatch.setenv("NVTE_NORM_FWD_USE_CUDNN", str(int(use_cudnn))) quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3, with_2d_quantization=with_2d_quantization) quantizer.optimize_for_gemm = optimize_for_gemm inp = TensorSpec(shape=shape, dtype=torch.bfloat16, device=torch.device("cpu")) out, _, _ = fake_norm(inp, quantizer, normalization=normalization) - compact = compact and (normalization == "LayerNorm" or not with_2d_quantization) + compact = compact and not with_2d_quantization expected = optimize_for_gemm and not compact assert out.create_metadata()["nontensor_kwargs"]["with_gemm_swizzled_scales"] == expected assert out.create_tensor()._with_gemm_swizzled_scales == expected @@ -2098,7 +2098,15 @@ def test_fake_norm_mxfp8_scale_layout( @pytest.mark.parametrize("compile_mode", _compile_modes) @pytest.mark.parametrize( "fp8_recipe", - [None, *_all_recipes], + [ + None, + *_all_recipes, + recipe.CustomRecipe( + qfactory=lambda _role: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ) + ), + ], ids=lambda r: "bf16" if r is None else type(r).__name__, ) @pytest.mark.parametrize("module", _LINEAR_MODULES) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 324231fc5f2..dbcd851593c 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -340,8 +340,8 @@ py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, * LayerNorm **************************************************************************************************/ -std::optional mxfp8_norm_fused_alignment(bool rmsnorm, bool with_2d_quantization, - bool use_cudnn); +std::optional> normalization_fwd_output_layout( + const std::string &normalization, py::handle quantizer, bool use_cudnn); std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, const at::Tensor &mu, const at::Tensor &rsigma, diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index cc990576081..e1083df3048 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -10,12 +10,30 @@ namespace transformer_engine::pytorch { -std::optional mxfp8_norm_fused_alignment(bool rmsnorm, bool with_2d_quantization, - bool use_cudnn) { - if (rmsnorm ? with_2d_quantization : !use_cudnn) { +namespace { + +std::optional> normalization_output_layout(const std::string &normalization, + const Quantizer &quantizer, + bool use_cudnn) { + NVTE_CHECK(normalization == "LayerNorm" || normalization == "RMSNorm", + "Unsupported normalization: ", normalization); + const auto *mxfp8 = dynamic_cast(&quantizer); + if (mxfp8 == nullptr || mxfp8->with_2d_quantization || + (normalization == "LayerNorm" && !use_cudnn)) { return std::nullopt; } - return 128; + return std::pair{128, !use_cudnn && quantizer.optimize_for_gemm}; +} + +} // namespace + +std::optional> normalization_fwd_output_layout( + const std::string &normalization, py::handle quantizer, bool use_cudnn) { + init_extension(); + if (!detail::IsMXFP8Quantizers(quantizer.ptr())) { + return std::nullopt; + } + return normalization_output_layout(normalization, *convert_quantizer(quantizer), use_cudnn); } std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, @@ -114,11 +132,9 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe if (quantizer.is_none() || IsFloat8Quantizers(quantizer.ptr())) { impl = Impl::FULLY_FUSED; } else if (IsMXFP8Quantizers(quantizer.ptr())) { - const auto mxfp8_quantizer_cpp = static_cast(quantizer_cpp.get()); - const auto alignment = - mxfp8_norm_fused_alignment(false, mxfp8_quantizer_cpp->with_2d_quantization, - transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); - if (alignment && outer_size % *alignment == 0 && inner_size % *alignment == 0) { + const auto layout = normalization_output_layout( + "LayerNorm", *quantizer_cpp, transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); + if (layout && outer_size % layout->first == 0 && inner_size % layout->first == 0) { impl = Impl::FULLY_FUSED; } } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr()) && @@ -358,12 +374,9 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w if (quantizer.is_none() || IsFloat8Quantizers(quantizer.ptr())) { impl = Impl::FULLY_FUSED; } else if (IsMXFP8Quantizers(quantizer.ptr())) { - auto mxfp8_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); - NVTE_CHECK(mxfp8_quantizer_cpp != nullptr, "Could not cast to MXFP8 quantizer"); - const auto alignment = - mxfp8_norm_fused_alignment(true, mxfp8_quantizer_cpp->with_2d_quantization, - transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); - if (alignment && outer_size % *alignment == 0 && inner_size % *alignment == 0) { + const auto layout = normalization_output_layout( + "RMSNorm", *quantizer_cpp, transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); + if (layout && outer_size % layout->first == 0 && inner_size % layout->first == 0) { if (transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { impl = Impl::FUSED_NORM_QUANT_UNSWIZZLED; } else { diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index a5a5a025c75..6dc36b4c78a 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -402,8 +402,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); // Other granular functions - m.def("mxfp8_norm_fused_alignment", &transformer_engine::pytorch::mxfp8_norm_fused_alignment, - py::arg("rmsnorm"), py::arg("with_2d_quantization"), py::arg("use_cudnn")); + m.def("normalization_fwd_output_layout", + &transformer_engine::pytorch::normalization_fwd_output_layout, py::arg("normalization"), + py::arg("quantizer"), py::arg("use_cudnn")); m.def("layernorm_fwd", &transformer_engine::pytorch::layernorm_fwd, "LayerNorm", py::arg("input"), py::arg("weight"), py::arg("bias"), py::arg("eps"), py::arg("ln_out"), py::arg("quantizer"), py::arg("otype"), py::arg("sm_margin"), py::arg("zero_centered_gamma")); diff --git a/transformer_engine/pytorch/dynamo/fake.py b/transformer_engine/pytorch/dynamo/fake.py index 9b70526fadc..29c177f33d0 100644 --- a/transformer_engine/pytorch/dynamo/fake.py +++ b/transformer_engine/pytorch/dynamo/fake.py @@ -14,7 +14,6 @@ from .tensor_spec import TensorSpec from ..quantized_tensor import Quantizer -from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..utils import get_device_compute_capability @@ -112,12 +111,11 @@ def fake_activation( @torch.compiler.assume_constant_result -def _mxfp8_norm_fused_alignment( - normalization: str, with_2d_quantization: bool, use_cudnn: bool -) -> Optional[int]: - return tex.mxfp8_norm_fused_alignment( - normalization == "RMSNorm", with_2d_quantization, use_cudnn - ) +def _normalization_fwd_output_layout( + normalization: str, quantizer: Quantizer, use_cudnn: bool +) -> Optional[tuple[int, bool]]: + """Return the alignment and scale-layout override for a normalization output.""" + return tex.normalization_fwd_output_layout(normalization, quantizer, use_cudnn) def fake_norm( @@ -134,13 +132,13 @@ def fake_norm( input_.shape, dtype or input_.dtype, input_.device, quantizer ) rows, cols = math.prod(input_.shape[:-1]), input_.shape[-1] - if isinstance(quantizer, MXFP8Quantizer): - use_cudnn = bool(int(os.getenv("NVTE_NORM_FWD_USE_CUDNN", "0") or "0")) - alignment = _mxfp8_norm_fused_alignment( - normalization, quantizer.with_2d_quantization, use_cudnn - ) - if use_cudnn and alignment is not None and rows % alignment == 0 and cols % alignment == 0: - out.with_gemm_swizzled_scales = False + if quantizer is not None: + use_cudnn = bool(int(os.environ.get("NVTE_NORM_FWD_USE_CUDNN", "0") or "0")) + layout = _normalization_fwd_output_layout(normalization, quantizer, use_cudnn) + if layout is not None: + alignment, with_gemm_swizzled_scales = layout + if rows % alignment == 0 and cols % alignment == 0: + out.with_gemm_swizzled_scales = with_gemm_swizzled_scales mean = ( TensorSpec(shape=(rows,), dtype=torch.float32, device=input_.device) if normalization == "LayerNorm" diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a61afe9f154..d9e3925a2b7 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -156,6 +156,7 @@ class LayerNormLinearFwdArgs: fp8: bool fp8_calibration: bool backward_override: Optional[str] + fprop_use_split_accumulator: bool dgrad_use_split_accumulator: bool wgrad_use_split_accumulator: bool debug: bool @@ -379,7 +380,7 @@ def _layernorm_linear_forward_impl( ln_bias = cast_if_needed(ln_bias, activation_dtype) nvtx_range_pop(f"{nvtx_label}.norm_input_cast") - if is_cpu_offload_enabled(): + if cpu_offloading: start_offload(inputmat) tp_world_size = get_distributed_world_size(tp_group) @@ -561,11 +562,7 @@ def _layernorm_linear_forward_impl( weight_quantizer.calibrate(weight) # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + use_split_accumulator = args.fprop_use_split_accumulator # Configure output quantizer if output_quantizer is not None: @@ -2504,11 +2501,14 @@ def forward( if cache_name is not None and not (use_compiled_op and is_first_microbatch): weight_workspace = self._fp8_workspaces.get(cache_name) + fprop_use_split_accumulator = _2X_ACC_FPROP dgrad_use_split_accumulator = _2X_ACC_DGRAD wgrad_use_split_accumulator = _2X_ACC_WGRAD if self.fp8: _recipe = FP8GlobalStateManager.get_fp8_recipe() backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_fprop"): + fprop_use_split_accumulator = _recipe.fp8_gemm_fprop.use_split_accumulator if hasattr(_recipe, "fp8_gemm_dgrad"): dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator if hasattr(_recipe, "fp8_gemm_wgrad"): @@ -2595,6 +2595,7 @@ def forward( fp8=self.fp8, fp8_calibration=self.fp8_calibration, backward_override=backward_override, + fprop_use_split_accumulator=fprop_use_split_accumulator, dgrad_use_split_accumulator=dgrad_use_split_accumulator, wgrad_use_split_accumulator=wgrad_use_split_accumulator, debug=debug, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index d22fe30ad2a..622b2818b7e 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -254,9 +254,10 @@ class LayerNormMLPFwdArgs: fp8: bool fp8_calibration: bool backward_override: Optional[str] + fprop_use_split_accumulator: bool dgrad_use_split_accumulator: bool wgrad_use_split_accumulator: bool - # Recipe properties the backward needs (the recipe itself can't cross the op boundary). + # The recipe itself can't cross the op boundary. recipe_float8_block_scaling: bool recipe_custom: bool recipe_dbias_dact_fusion: bool @@ -570,7 +571,7 @@ def _layernorm_mlp_forward_impl( assert_dim_for_fp8_exec(inputmat, fc1_weight, fc2_weight) activation_func = _act_func( - activation, FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + activation, dbias_fusion=fp8 and args.recipe_dbias_dact_fusion )[0] # Cast for native AMP @@ -578,7 +579,7 @@ def _layernorm_mlp_forward_impl( ln_weight_cast = cast_if_needed(ln_weight, activation_dtype) if ln_bias is not None: ln_bias = cast_if_needed(ln_bias, activation_dtype) - if is_cpu_offload_enabled(): + if cpu_offloading: start_offload(inputmat) tp_world_size = get_distributed_world_size(tp_group) @@ -597,11 +598,7 @@ def _layernorm_mlp_forward_impl( ub_overlap_rs = ub_overlap_rs and is_grad_enabled # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + use_split_accumulator = args.fprop_use_split_accumulator # Configure quantizer for norm output if fp8: @@ -847,12 +844,11 @@ def _layernorm_mlp_forward_impl( else: fc1_out, *_ = fc1_outputs if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if recipe.float8_block_scaling(): + if args.recipe_float8_block_scaling: # tex.quantize does not support GELU fusion for blockwise act_out = activation_func(fc1_out, None, **act_params) act_out = tex.quantize(act_out, fc2_input_quantizer) - elif recipe.custom(): + elif args.recipe_custom: # tex.quantize does not support custom quantizers act_out = activation_func(fc1_out, None, **act_params) act_out = fc2_input_quantizer(act_out) @@ -3190,6 +3186,7 @@ def forward( fc1_weight_workspace = self._fp8_workspaces.get(cache_name_fc1) fc2_weight_workspace = self._fp8_workspaces.get(cache_name_fc2) + fprop_use_split_accumulator = _2X_ACC_FPROP dgrad_use_split_accumulator = _2X_ACC_DGRAD wgrad_use_split_accumulator = _2X_ACC_WGRAD recipe_float8_block_scaling = False @@ -3198,6 +3195,8 @@ def forward( if self.fp8: _recipe = FP8GlobalStateManager.get_fp8_recipe() backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_fprop"): + fprop_use_split_accumulator = _recipe.fp8_gemm_fprop.use_split_accumulator if hasattr(_recipe, "fp8_gemm_dgrad"): dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator if hasattr(_recipe, "fp8_gemm_wgrad"): @@ -3300,6 +3299,7 @@ def forward( fp8=self.fp8, fp8_calibration=self.fp8_calibration, backward_override=backward_override, + fprop_use_split_accumulator=fprop_use_split_accumulator, dgrad_use_split_accumulator=dgrad_use_split_accumulator, wgrad_use_split_accumulator=wgrad_use_split_accumulator, recipe_float8_block_scaling=recipe_float8_block_scaling, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 8cd4cf0d1e3..0b73a4de17d 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -143,8 +143,11 @@ class LinearFwdArgs: fp8_output: bool save_original_input: bool backward_override: Optional[str] + fprop_use_split_accumulator: bool dgrad_use_split_accumulator: bool wgrad_use_split_accumulator: bool + recipe_custom: bool + fp8_recompute_forward: bool custom: bool debug: bool @@ -324,7 +327,7 @@ def _linear_forward_impl( # to reuse the high-precision DPA output saved by attention backward. Validate # the resolved quantizer since this is not necessarily an informed user opt-in. if isinstance(input_quantizer, Float8Quantizer): - if FP8GlobalStateManager.get_fp8_recipe().custom(): + if args.recipe_custom: warnings.warn( "save_original_input is incompatible with delayed-scaling quantizers " "(Float8Quantizer). Disabling save_original_input for this module.", @@ -448,7 +451,7 @@ def _linear_forward_impl( inputmat = cast_if_needed(inp, activation_dtype) # Cast for AMP inputmat_total = inputmat - if is_cpu_offload_enabled(): + if cpu_offloading: start_offload(inputmat) nvtx_range_pop(f"{nvtx_label}.input_cast_comm") # ------------------------------------------------------ @@ -477,10 +480,7 @@ def _linear_forward_impl( if backward_override is not None: columnwise_usage = False if not columnwise_usage: - columnwise_usage = ( - is_fp8_activation_recompute_enabled() - and not in_fp8_activation_recompute_phase() - ) + columnwise_usage = args.fp8_recompute_forward weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) elif isinstance(weight, QuantizedTensor): weight_quantizer = weight._quantizer @@ -519,11 +519,7 @@ def _linear_forward_impl( weight_quantizer.calibrate(weight) # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + use_split_accumulator = args.fprop_use_split_accumulator # Configure output quantizer if output_quantizer is not None: @@ -770,10 +766,7 @@ def _linear_forward_fake( if args.backward_override is not None: columnwise_usage = False if not columnwise_usage: - columnwise_usage = ( - is_fp8_activation_recompute_enabled() - and not in_fp8_activation_recompute_phase() - ) + columnwise_usage = args.fp8_recompute_forward weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) elif weight.is_quantized: weight_quantizer = weight.quantizer @@ -2364,15 +2357,20 @@ def forward( if cache_name is not None and not (use_compiled_op and is_first_microbatch): weight_workspace = self._fp8_workspaces.get(cache_name) + fprop_use_split_accumulator = _2X_ACC_FPROP dgrad_use_split_accumulator = _2X_ACC_DGRAD wgrad_use_split_accumulator = _2X_ACC_WGRAD + recipe_custom = False if self.fp8: _recipe = FP8GlobalStateManager.get_fp8_recipe() backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_fprop"): + fprop_use_split_accumulator = _recipe.fp8_gemm_fprop.use_split_accumulator if hasattr(_recipe, "fp8_gemm_dgrad"): dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator if hasattr(_recipe, "fp8_gemm_wgrad"): wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator + recipe_custom = _recipe.custom() else: backward_override = None custom = is_custom(input_quantizer) or is_custom(weight_quantizer) @@ -2428,8 +2426,14 @@ def forward( fp8_output=fp8_output, save_original_input=self.save_original_input, backward_override=backward_override, + fprop_use_split_accumulator=fprop_use_split_accumulator, dgrad_use_split_accumulator=dgrad_use_split_accumulator, wgrad_use_split_accumulator=wgrad_use_split_accumulator, + recipe_custom=recipe_custom, + fp8_recompute_forward=( + is_fp8_activation_recompute_enabled() + and not in_fp8_activation_recompute_phase() + ), custom=custom, debug=debug, # weight-workspace caching From 0a7f8bf4a305aea48459f6def6ac76751369636d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 30 Oct 2026 13:42:50 +0100 Subject: [PATCH 29/32] Gate custom compile recipe on FP8 support Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index d7a3890b66b..b4c82ee4b48 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -2098,15 +2098,18 @@ def test_fake_norm_mxfp8_scale_layout( @pytest.mark.parametrize("compile_mode", _compile_modes) @pytest.mark.parametrize( "fp8_recipe", - [ - None, - *_all_recipes, - recipe.CustomRecipe( - qfactory=lambda _role: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" + [None, *_all_recipes] + + ( + [ + recipe.CustomRecipe( + qfactory=lambda _role: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ) ) - ), - ], + ] + if fp8_available + else [] + ), ids=lambda r: "bf16" if r is None else type(r).__name__, ) @pytest.mark.parametrize("module", _LINEAR_MODULES) From 3add45caa70b633b0192270410f51cb10cfac891 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:44:12 +0000 Subject: [PATCH 30/32] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/module/layernorm_mlp.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 622b2818b7e..3cad214627d 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -570,9 +570,7 @@ def _layernorm_mlp_forward_impl( if fp8: assert_dim_for_fp8_exec(inputmat, fc1_weight, fc2_weight) - activation_func = _act_func( - activation, dbias_fusion=fp8 and args.recipe_dbias_dact_fusion - )[0] + activation_func = _act_func(activation, dbias_fusion=fp8 and args.recipe_dbias_dact_fusion)[0] # Cast for native AMP inputmat = cast_if_needed(inputmat, activation_dtype) From 9108ab557b828741f1209124a8a2f76582ad6fa1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 30 Oct 2026 13:59:08 +0100 Subject: [PATCH 31/32] Clarify normalization scale layout override naming Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/normalization.cpp | 15 ++++++++------- .../pytorch/csrc/extensions/pybind.cpp | 4 ++-- transformer_engine/pytorch/dynamo/fake.py | 8 ++++---- 4 files changed, 15 insertions(+), 14 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index dbcd851593c..1924c9c081f 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -340,7 +340,7 @@ py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, * LayerNorm **************************************************************************************************/ -std::optional> normalization_fwd_output_layout( +std::optional> normalization_scale_layout_override( const std::string &normalization, py::handle quantizer, bool use_cudnn); std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index e1083df3048..de2ed37420c 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -12,9 +12,9 @@ namespace transformer_engine::pytorch { namespace { -std::optional> normalization_output_layout(const std::string &normalization, - const Quantizer &quantizer, - bool use_cudnn) { +// (row/column alignment, swizzled scales); nullopt keeps the quantizer's default layout. +std::optional> normalization_scale_layout_override_impl( + const std::string &normalization, const Quantizer &quantizer, bool use_cudnn) { NVTE_CHECK(normalization == "LayerNorm" || normalization == "RMSNorm", "Unsupported normalization: ", normalization); const auto *mxfp8 = dynamic_cast(&quantizer); @@ -27,13 +27,14 @@ std::optional> normalization_output_layout(const std::st } // namespace -std::optional> normalization_fwd_output_layout( +std::optional> normalization_scale_layout_override( const std::string &normalization, py::handle quantizer, bool use_cudnn) { init_extension(); if (!detail::IsMXFP8Quantizers(quantizer.ptr())) { return std::nullopt; } - return normalization_output_layout(normalization, *convert_quantizer(quantizer), use_cudnn); + return normalization_scale_layout_override_impl(normalization, *convert_quantizer(quantizer), + use_cudnn); } std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, @@ -132,7 +133,7 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe if (quantizer.is_none() || IsFloat8Quantizers(quantizer.ptr())) { impl = Impl::FULLY_FUSED; } else if (IsMXFP8Quantizers(quantizer.ptr())) { - const auto layout = normalization_output_layout( + const auto layout = normalization_scale_layout_override_impl( "LayerNorm", *quantizer_cpp, transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); if (layout && outer_size % layout->first == 0 && inner_size % layout->first == 0) { impl = Impl::FULLY_FUSED; @@ -374,7 +375,7 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w if (quantizer.is_none() || IsFloat8Quantizers(quantizer.ptr())) { impl = Impl::FULLY_FUSED; } else if (IsMXFP8Quantizers(quantizer.ptr())) { - const auto layout = normalization_output_layout( + const auto layout = normalization_scale_layout_override_impl( "RMSNorm", *quantizer_cpp, transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); if (layout && outer_size % layout->first == 0 && inner_size % layout->first == 0) { if (transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 6dc36b4c78a..3aa55833dc5 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -402,8 +402,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); // Other granular functions - m.def("normalization_fwd_output_layout", - &transformer_engine::pytorch::normalization_fwd_output_layout, py::arg("normalization"), + m.def("normalization_scale_layout_override", + &transformer_engine::pytorch::normalization_scale_layout_override, py::arg("normalization"), py::arg("quantizer"), py::arg("use_cudnn")); m.def("layernorm_fwd", &transformer_engine::pytorch::layernorm_fwd, "LayerNorm", py::arg("input"), py::arg("weight"), py::arg("bias"), py::arg("eps"), py::arg("ln_out"), py::arg("quantizer"), diff --git a/transformer_engine/pytorch/dynamo/fake.py b/transformer_engine/pytorch/dynamo/fake.py index 29c177f33d0..d9b22f57501 100644 --- a/transformer_engine/pytorch/dynamo/fake.py +++ b/transformer_engine/pytorch/dynamo/fake.py @@ -111,11 +111,11 @@ def fake_activation( @torch.compiler.assume_constant_result -def _normalization_fwd_output_layout( +def _normalization_scale_layout_override( normalization: str, quantizer: Quantizer, use_cudnn: bool ) -> Optional[tuple[int, bool]]: - """Return the alignment and scale-layout override for a normalization output.""" - return tex.normalization_fwd_output_layout(normalization, quantizer, use_cudnn) + """Return (alignment, swizzled scales), or None to keep the quantizer's layout.""" + return tex.normalization_scale_layout_override(normalization, quantizer, use_cudnn) def fake_norm( @@ -134,7 +134,7 @@ def fake_norm( rows, cols = math.prod(input_.shape[:-1]), input_.shape[-1] if quantizer is not None: use_cudnn = bool(int(os.environ.get("NVTE_NORM_FWD_USE_CUDNN", "0") or "0")) - layout = _normalization_fwd_output_layout(normalization, quantizer, use_cudnn) + layout = _normalization_scale_layout_override(normalization, quantizer, use_cudnn) if layout is not None: alignment, with_gemm_swizzled_scales = layout if rows % alignment == 0 and cols % alignment == 0: From 267b0d4f68c4e7f31fc92669031c89b8347ffcb1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 30 Oct 2026 14:14:07 +0100 Subject: [PATCH 32/32] Use normalization_scale_layout and explain returned metadata Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/csrc/extensions.h | 5 +++-- .../pytorch/csrc/extensions/normalization.cpp | 18 ++++++++++-------- .../pytorch/csrc/extensions/pybind.cpp | 5 ++--- transformer_engine/pytorch/dynamo/fake.py | 6 +++--- 4 files changed, 18 insertions(+), 16 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1924c9c081f..2221f014bae 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -340,8 +340,9 @@ py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, * LayerNorm **************************************************************************************************/ -std::optional> normalization_scale_layout_override( - const std::string &normalization, py::handle quantizer, bool use_cudnn); +std::optional> normalization_scale_layout(const std::string &normalization, + py::handle quantizer, + bool use_cudnn); std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, const at::Tensor &mu, const at::Tensor &rsigma, diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index de2ed37420c..a32899fd47d 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -12,8 +12,10 @@ namespace transformer_engine::pytorch { namespace { -// (row/column alignment, swizzled scales); nullopt keeps the quantizer's default layout. -std::optional> normalization_scale_layout_override_impl( +// Return (row/column alignment, GEMM-swizzled scales); apply only to aligned shapes. +// cuDNN emits unswizzled MXFP8 scales even with optimize_for_gemm=true. +// nullopt keeps the quantizer's default layout. +std::optional> normalization_scale_layout_impl( const std::string &normalization, const Quantizer &quantizer, bool use_cudnn) { NVTE_CHECK(normalization == "LayerNorm" || normalization == "RMSNorm", "Unsupported normalization: ", normalization); @@ -27,14 +29,14 @@ std::optional> normalization_scale_layout_override_impl( } // namespace -std::optional> normalization_scale_layout_override( - const std::string &normalization, py::handle quantizer, bool use_cudnn) { +std::optional> normalization_scale_layout(const std::string &normalization, + py::handle quantizer, + bool use_cudnn) { init_extension(); if (!detail::IsMXFP8Quantizers(quantizer.ptr())) { return std::nullopt; } - return normalization_scale_layout_override_impl(normalization, *convert_quantizer(quantizer), - use_cudnn); + return normalization_scale_layout_impl(normalization, *convert_quantizer(quantizer), use_cudnn); } std::vector layernorm_bwd(const at::Tensor &dz, const at::Tensor &x, @@ -133,7 +135,7 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe if (quantizer.is_none() || IsFloat8Quantizers(quantizer.ptr())) { impl = Impl::FULLY_FUSED; } else if (IsMXFP8Quantizers(quantizer.ptr())) { - const auto layout = normalization_scale_layout_override_impl( + const auto layout = normalization_scale_layout_impl( "LayerNorm", *quantizer_cpp, transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); if (layout && outer_size % layout->first == 0 && inner_size % layout->first == 0) { impl = Impl::FULLY_FUSED; @@ -375,7 +377,7 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w if (quantizer.is_none() || IsFloat8Quantizers(quantizer.ptr())) { impl = Impl::FULLY_FUSED; } else if (IsMXFP8Quantizers(quantizer.ptr())) { - const auto layout = normalization_scale_layout_override_impl( + const auto layout = normalization_scale_layout_impl( "RMSNorm", *quantizer_cpp, transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")); if (layout && outer_size % layout->first == 0 && inner_size % layout->first == 0) { if (transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 3aa55833dc5..9d72aa8b5ed 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -402,9 +402,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); // Other granular functions - m.def("normalization_scale_layout_override", - &transformer_engine::pytorch::normalization_scale_layout_override, py::arg("normalization"), - py::arg("quantizer"), py::arg("use_cudnn")); + m.def("normalization_scale_layout", &transformer_engine::pytorch::normalization_scale_layout, + py::arg("normalization"), py::arg("quantizer"), py::arg("use_cudnn")); m.def("layernorm_fwd", &transformer_engine::pytorch::layernorm_fwd, "LayerNorm", py::arg("input"), py::arg("weight"), py::arg("bias"), py::arg("eps"), py::arg("ln_out"), py::arg("quantizer"), py::arg("otype"), py::arg("sm_margin"), py::arg("zero_centered_gamma")); diff --git a/transformer_engine/pytorch/dynamo/fake.py b/transformer_engine/pytorch/dynamo/fake.py index d9b22f57501..0ced4d17c88 100644 --- a/transformer_engine/pytorch/dynamo/fake.py +++ b/transformer_engine/pytorch/dynamo/fake.py @@ -111,11 +111,11 @@ def fake_activation( @torch.compiler.assume_constant_result -def _normalization_scale_layout_override( +def _normalization_scale_layout( normalization: str, quantizer: Quantizer, use_cudnn: bool ) -> Optional[tuple[int, bool]]: """Return (alignment, swizzled scales), or None to keep the quantizer's layout.""" - return tex.normalization_scale_layout_override(normalization, quantizer, use_cudnn) + return tex.normalization_scale_layout(normalization, quantizer, use_cudnn) def fake_norm( @@ -134,7 +134,7 @@ def fake_norm( rows, cols = math.prod(input_.shape[:-1]), input_.shape[-1] if quantizer is not None: use_cudnn = bool(int(os.environ.get("NVTE_NORM_FWD_USE_CUDNN", "0") or "0")) - layout = _normalization_scale_layout_override(normalization, quantizer, use_cudnn) + layout = _normalization_scale_layout(normalization, quantizer, use_cudnn) if layout is not None: alignment, with_gemm_swizzled_scales = layout if rows % alignment == 0 and cols % alignment == 0: