From f3401df703909f211dc8617e8832ba2f2d3592a0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Sat, 6 Jun 2026 14:11:29 +0200 Subject: [PATCH 01/42] [PyTorch] Make tensorless quantizers opaque value objects for torch.compile Give tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling, NVFP4) value-object semantics so torch.compile can treat them as baked-in constants: - Add opt-in value identity to the base Quantizer (_value_fields / _value_key / __eq__ / __hash__). Quantizers holding live tensors (delayed-scaling Float8Quantizer) and custom quantizers keep identity semantics. - New transformer_engine/pytorch/dynamo.py houses the torch.compile glue: __fx_repr__, value-key reconstruction and register_value_opaque_quantizer (gracefully a no-op on PyTorch builds without the opaque-object API). - Register the four tensorless quantizers as value opaque types. Also fix CustomRecipe state caching in TransformerEngineBaseModule: set_meta_tensor now rebuilds quantizers when the CustomRecipe instance changes (e.g. nested te.autocast regions) instead of reusing the first recipe's state, since every CustomRecipe shares the CustomRecipeState type but carries its own qfactory. Move the quantizer value-object tests into tests/pytorch/test_torch_compile.py and add that file to the L0 pytorch unittest QA suite. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 154 ++++++++++++++++++ transformer_engine/pytorch/dynamo.py | 120 ++++++++++++++ .../pytorch/quantized_tensor.py | 67 ++++++++ .../pytorch/tensor/float8_blockwise_tensor.py | 7 + .../pytorch/tensor/float8_tensor.py | 9 + .../pytorch/tensor/mxfp8_tensor.py | 7 + .../pytorch/tensor/nvfp4_tensor.py | 25 +++ 7 files changed, 389 insertions(+) create mode 100644 transformer_engine/pytorch/dynamo.py diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 1286492a6e..8adcc88e61 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -24,6 +24,7 @@ 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.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.quantization import QuantizerRole @@ -32,6 +33,14 @@ is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, + Float8Quantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) +from transformer_engine.pytorch.dynamo import ( + register_value_opaque_quantizer, + _quantizer_from_value_key, ) from utils import recipe_id @@ -384,3 +393,148 @@ def fn(inp): out = compiled(inp) out.sum().backward() + + +# --------------------------------------------------------------------------- +# Value-opaque quantizers: eager value semantics + FX reconstruction +# +# The tensorless quantizers (current-scaling FP8, FP8 blockwise, MXFP8, NVFP4) +# are torch.compile *value* opaque types: they provide value-based +# ``__eq__`` / ``__hash__`` and an evaluable ``__fx_repr__`` (see +# ``torch._library.opaque_object`` Note [Opaque Objects]). These tests exercise +# the eager value semantics and the FX reconstruction round-trip. They are +# CPU-friendly except for NVFP4 (whose constructor touches the current CUDA +# device). +# --------------------------------------------------------------------------- + + +def _mxfp8(dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True): + return MXFP8Quantizer(fp8_dtype=dtype, rowwise=rowwise, columnwise=columnwise) + + +def _blockwise(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=True, block_scaling_dim=2): + return Float8BlockQuantizer( + fp8_dtype=dtype, + rowwise=True, + columnwise=True, + force_pow_2_scales=force_pow_2_scales, + block_scaling_dim=block_scaling_dim, + ) + + +def _current_scaling(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=False, amax_epsilon=0.0): + return Float8CurrentScalingQuantizer( + fp8_dtype=dtype, + device=torch.device("cpu"), + force_pow_2_scales=force_pow_2_scales, + amax_epsilon=amax_epsilon, + ) + + +# (factory, kwargs_for_a_different_but_valid_config) +_CPU_VALUE_QUANTIZERS = [ + pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), + pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), + pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), +] + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_value_equality(factory, other_kwargs): + """Same config -> equal & same hash; different config -> not equal.""" + a = factory() + b = factory() + assert a is not b + assert a == b + assert hash(a) == hash(b) + + c = factory(**other_kwargs) + assert a != c + # Usage flags participate in the value. + d = factory() + d.set_usage(rowwise=False, columnwise=False) + assert a != d + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_usable_in_set_and_dict(factory, other_kwargs): + a = factory() + b = factory() + c = factory(**other_kwargs) + assert len({a, b, c}) == 2 + mapping = {a: "x"} + assert mapping[b] == "x" + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_cross_type_inequality(factory, other_kwargs): + a = factory() + other = _current_scaling() if not isinstance(a, Float8CurrentScalingQuantizer) else _mxfp8() + assert a != other + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_fx_repr_roundtrip(factory, other_kwargs): + """``__fx_repr__`` returns an evaluable expression rebuilding an equal object.""" + a = factory() + repr_str, globals_ = a.__fx_repr__() + assert isinstance(repr_str, str) + assert isinstance(globals_, dict) + rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used + assert rebuilt == a + assert rebuilt is not a + assert hash(rebuilt) == hash(a) + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_value_key_reconstruction(factory, other_kwargs): + a = factory() + rebuilt = _quantizer_from_value_key(a._value_key()) + assert type(rebuilt) is type(a) + assert rebuilt == a + # The deprecated amax-reduction process group is never carried in the value. + assert getattr(rebuilt, "amax_reduction_group", None) is None + + +def test_quantizer_delayed_scaling_keeps_identity_semantics(): + """Float8Quantizer holds live tensors -> identity (not value) semantics.""" + scale = torch.ones(1) + amax = torch.zeros(1) + a = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) + b = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) + assert a._value_fields() is None + assert a == a + assert a != b # distinct instances are not equal despite identical config + assert hash(a) == object.__hash__(a) + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_registration_is_idempotent_and_tolerant(factory, other_kwargs): + """Re-registering must not raise, regardless of PyTorch opaque-object support.""" + cls = type(factory()) + register_value_opaque_quantizer(cls) + register_value_opaque_quantizer(cls) + + if not _opaque_available: + pytest.skip("PyTorch build without opaque-object API") + from torch._library.opaque_object import is_opaque_value_type + + assert is_opaque_value_type(cls) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4Quantizer requires CUDA") +def test_quantizer_nvfp4_value_semantics(): + a = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + b = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + assert a == b + assert hash(a) == hash(b) + + c = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, with_rht=not a.with_rht) + assert a != c + + rebuilt = _quantizer_from_value_key(a._value_key()) + assert rebuilt == a + assert rebuilt.amax_reduction_group is None + + repr_str, globals_ = a.__fx_repr__() + assert eval(repr_str, dict(globals_)) == a # pylint: disable=eval-used diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo.py new file mode 100644 index 0000000000..26aff6f59c --- /dev/null +++ b/transformer_engine/pytorch/dynamo.py @@ -0,0 +1,120 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile glue for Transformer Engine quantizers. + +This module isolates the torch.compile-specific plumbing that turns a +*tensorless* quantizer into a torch.compile **value** opaque type: + + * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used + by FX codegen and registers the quantizer class with + ``torch._library.opaque_object``. It is a no-op (other than populating the + local registry) on PyTorch builds without the opaque-object API, so + importing Transformer Engine never fails on older PyTorch -- only + torch.compile specialization on the quantizer is unavailable there. + * :func:`_quantizer_from_value_key` -- rebuilds a quantizer constant from its + value key inside the generated FX graph. + +The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` / +``_value_fields``) live on the quantizer itself; see +:class:`transformer_engine.pytorch.quantized_tensor.Quantizer`. + +See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a +value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` / +``__fx_repr__``). +""" + +from __future__ import annotations +from typing import Any, Dict, Tuple + +from .constants import DType + + +# Maps a quantizer class qualname to the class object. A value key stores only +# the qualname, so reconstruction looks the class up here. Populated by +# ``register_value_opaque_quantizer`` at import time of each tensor module; this +# avoids importing the tensor modules into this module (which would create an +# import cycle). +_QUANTIZER_VALUE_REGISTRY: Dict[str, type] = {} + + +def _quantizer_from_value_key(key: Tuple[Any, ...]) -> Any: + """Rebuild a tensorless quantizer from its value key. + + Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the + generated FX code calls this to materialize the quantizer constant. The + deprecated amax-reduction process group is never part of the value, so a + reconstructed quantizer always starts with no stored group. + """ + qualname, items = key[0], key[1] + cls = _QUANTIZER_VALUE_REGISTRY[qualname] + # Bypass ``__init__`` and restore the value attributes directly: the value + # key already captures every value-defining field (including derived ones), + # and the constructors have heterogeneous signatures / side effects. + obj = cls.__new__(cls) + field_names = set() + for name, value in items: + if name == "dtype": + value = DType.cast(value) + object.__setattr__(obj, name, value) + field_names.add(name) + # The deprecated amax-reduction process group is excluded from the value; + # restore it as ``None`` for quantizers that still carry the fallback so + # attribute access keeps working. + if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): + object.__setattr__(obj, "amax_reduction_group", None) + return obj + + +def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: + """``__fx_repr__`` for value-opaque quantizers (attached at registration). + + Returns an evaluable expression that rebuilds the quantizer via + :func:`_quantizer_from_value_key`, together with the globals needed to + evaluate it. + """ + return ( + f"_quantizer_from_value_key({self._value_key()!r})", + {"_quantizer_from_value_key": _quantizer_from_value_key}, + ) + + +def register_value_opaque_quantizer(cls: type) -> None: + """Register a tensorless quantizer class as a torch.compile value opaque type. + + Attaches ``__fx_repr__`` and registers the class with + ``torch._library.opaque_object``. Safe to call on any PyTorch build: on + versions without the opaque-object API it only records the class in the + local registry and attaches ``__fx_repr__`` (both harmless), so Transformer + Engine keeps importing and running in eager mode. + + The quantizer class must already provide value ``__eq__`` / ``__hash__`` and + a non-``None`` ``_value_fields`` (see + :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). + """ + _QUANTIZER_VALUE_REGISTRY[cls.__qualname__] = cls + + # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the + # class, so attach it before registering. + if "__fx_repr__" not in cls.__dict__: + cls.__fx_repr__ = _quantizer_fx_repr + + try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + register_opaque_type, + is_opaque_value_type, + ) + except (ImportError, AttributeError): + # Older PyTorch without the opaque-object API: eager value semantics + # still work; torch.compile specialization on the quantizer does not. + return + + if is_opaque_value_type(cls): + return + + try: + register_opaque_type(cls, typ="value") + except (ImportError, AttributeError, RuntimeError, TypeError): + # Tolerate partial / experimental opaque-object support. + pass diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cfe488aae5..f37b3cc63d 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -408,6 +408,73 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } + # ----- Value-object identity (torch.compile opaque value support) ----- + # A *tensorless* quantizer (one whose entire state is a handful of plain, + # reproducible scalars -- no live tensors, no process groups) behaves like a + # value: two instances with the same configuration are interchangeable. Such + # quantizers opt into value-based ``__eq__`` / ``__hash__`` by overriding + # ``_value_fields``. Quantizers that keep the default (e.g. delayed-scaling + # ``Float8Quantizer``, which holds live scale/amax tensors, and any custom + # quantizer) retain the default identity semantics. + # + # This is the eager-side half of registering the quantizer as a torch.compile + # *value* opaque type; the torch.compile glue (``__fx_repr__``, FX + # reconstruction and ``register_opaque_type``) lives in + # ``transformer_engine.pytorch.dynamo``. + + #: Attributes shared by every quantizer that take part in value identity. + _BASE_VALUE_FIELDS: Tuple[str, ...] = ( + "rowwise_usage", + "columnwise_usage", + "internal", + "optimize_for_gemm", + ) + + def _value_fields(self) -> Optional[Tuple[str, ...]]: + """Subclass-specific value-defining attribute names, or ``None``. + + Returning ``None`` (the default) means the quantizer is *not* a value + object and keeps identity-based equality/hashing. Tensorless quantizers + override this to return the tuple of attribute names that, together with + :attr:`_BASE_VALUE_FIELDS`, fully determine their value (excluding + non-value state such as a deprecated amax-reduction process group). + """ + return None + + def _value_key(self) -> Tuple[Any, ...]: + """Hashable, reproducible key identifying this quantizer's value. + + Only valid for value quantizers (``_value_fields()`` is not ``None``). + """ + fields = self._value_fields() # pylint: disable=assignment-from-none + assert fields is not None, f"{type(self).__name__} is not a value quantizer" + items = [] + for name in self._BASE_VALUE_FIELDS + tuple(fields): + value = getattr(self, name) + if name == "dtype": + # ``DType`` is an ``IntEnum``; store the int so the key stays + # plain: hashable and ``repr``-reproducible for FX codegen. + value = int(value) + items.append((name, value)) + return (type(self).__qualname__, tuple(items)) + + def __eq__(self, other: object) -> Any: + # Value quantizers compare by configuration; everything else keeps the + # default identity semantics (returning ``NotImplemented`` makes Python + # fall back to identity). + if self is other: + return True + if self._value_fields() is None or type(self) is not type(other): + return NotImplemented + if other._value_fields() is None: + return NotImplemented + return self._value_key() == other._value_key() + + def __hash__(self) -> int: + if self._value_fields() is None: + return object.__hash__(self) + return hash(self._value_key()) + class QuantizedTensor(torch.Tensor): """Abstract base class for tensor with quantized data diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ba46508d74..92c22acd0b 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -14,6 +14,7 @@ from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -69,6 +70,9 @@ def copy(self) -> Float8BlockQuantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim") + def update_quantized( self, src: torch.Tensor, @@ -211,6 +215,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return Float8BlockScaling +register_value_opaque_quantizer(Float8BlockQuantizer) + + class Float8BlockwiseQTensor(Float8BlockwiseQTensorStorage, QuantizedTensor): """Tensor class with FP8 data quantized via NxN blocks or 1xN blocks. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e26abf7df0..56b1ecfb09 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -18,6 +18,7 @@ from ..utils import canonicalize_process_group, devices_match from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import dist_group_type, DType @@ -386,6 +387,14 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group (not a value) and is restored as ``None`` on rebuild. + return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") + + +register_value_opaque_quantizer(Float8CurrentScalingQuantizer) + class Float8Tensor(Float8TensorStorage, QuantizedTensor): """Experimental tensor class with FP8 data diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index d759aaf5c4..a3746b3088 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -18,6 +18,7 @@ from ..utils import devices_match, round_up_to_nearest_multiple from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -57,6 +58,9 @@ def copy(self) -> MXFP8Quantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype",) + def update_quantized( self, src: torch.Tensor, @@ -1058,3 +1062,6 @@ def backward( ) return dgrad, None return grad.view(ctx.shape), None + + +register_value_opaque_quantizer(MXFP8Quantizer) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index aa92be004f..573c78907c 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -23,6 +23,7 @@ from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -333,6 +334,30 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type: def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group (not a value) and is restored as ``None`` on rebuild. + # ``rht_matrix_random_sign_mask_t`` is derived (from + # ``_with_random_sign_mask`` and the device) but is stored verbatim so + # reconstruction does not need to touch the device. + return ( + "dtype", + "with_rht", + "with_post_rht_amax", + "with_2d_quantization", + "stochastic_rounding", + "row_scaled_nvfp4", + "nvfp4_use_4over6", + "nvfp4_e4m3_max", + "nvfp4_4over6_err_mode", + "_with_random_sign_mask", + "rht_matrix_random_sign_mask_t", + "with_amax_reduction", + ) + + +register_value_opaque_quantizer(NVFP4Quantizer) + class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): """Quantized tensor class with FP4 data From c4ad54c1b1f28d9d3402cff92ce244fb9f1a2bc4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Sat, 6 Jun 2026 14:51:06 +0200 Subject: [PATCH 02/42] [PyTorch] Drop quantizer value registry; reconstruct via __fx_repr__ globals Follow-up to the value-opaque quantizer support: - Remove the module-level _QUANTIZER_VALUE_REGISTRY (qualname -> class) and _quantizer_from_value_key. __fx_repr__ now captures the quantizer class directly in the FX globals and reconstructs via _rebuild_quantizer(cls, items), matching how PyTorch's own value opaque types (e.g. DTensor placements) reconstruct themselves. This removes global mutable state and the qualname collision risk. - Consolidate the quantizer value-object tests in test_torch_compile.py down to two functions and exercise reconstruction through the public __fx_repr__ path instead of internal helpers. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 138 +++--------------- transformer_engine/pytorch/dynamo.py | 53 +++---- .../pytorch/quantized_tensor.py | 22 +-- 3 files changed, 51 insertions(+), 162 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8adcc88e61..4491e23881 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -36,11 +36,6 @@ Float8Quantizer, Float8BlockQuantizer, MXFP8Quantizer, - NVFP4Quantizer, -) -from transformer_engine.pytorch.dynamo import ( - register_value_opaque_quantizer, - _quantizer_from_value_key, ) from utils import recipe_id @@ -396,145 +391,60 @@ def fn(inp): # --------------------------------------------------------------------------- -# Value-opaque quantizers: eager value semantics + FX reconstruction +# Value-opaque quantizers # -# The tensorless quantizers (current-scaling FP8, FP8 blockwise, MXFP8, NVFP4) -# are torch.compile *value* opaque types: they provide value-based -# ``__eq__`` / ``__hash__`` and an evaluable ``__fx_repr__`` (see -# ``torch._library.opaque_object`` Note [Opaque Objects]). These tests exercise -# the eager value semantics and the FX reconstruction round-trip. They are -# CPU-friendly except for NVFP4 (whose constructor touches the current CUDA -# device). +# Tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling) are +# torch.compile *value* opaque types: value-based ``__eq__`` / ``__hash__`` plus +# an evaluable ``__fx_repr__`` that rebuilds an equal object (see +# ``torch._library.opaque_object`` Note [Opaque Objects]). # --------------------------------------------------------------------------- -def _mxfp8(dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True): - return MXFP8Quantizer(fp8_dtype=dtype, rowwise=rowwise, columnwise=columnwise) +def _mxfp8(dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=dtype) -def _blockwise(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=True, block_scaling_dim=2): +def _blockwise(force_pow_2_scales=True): return Float8BlockQuantizer( - fp8_dtype=dtype, + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, force_pow_2_scales=force_pow_2_scales, - block_scaling_dim=block_scaling_dim, ) -def _current_scaling(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=False, amax_epsilon=0.0): +def _current_scaling(amax_epsilon=0.0): return Float8CurrentScalingQuantizer( - fp8_dtype=dtype, + fp8_dtype=tex.DType.kFloat8E4M3, device=torch.device("cpu"), - force_pow_2_scales=force_pow_2_scales, amax_epsilon=amax_epsilon, ) -# (factory, kwargs_for_a_different_but_valid_config) -_CPU_VALUE_QUANTIZERS = [ +# (factory, kwargs producing a different-but-valid config) +_VALUE_QUANTIZERS = [ pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), ] -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_value_equality(factory, other_kwargs): - """Same config -> equal & same hash; different config -> not equal.""" - a = factory() - b = factory() +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory, other_kwargs): + """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" + a, b = factory(), factory() + # Same config -> equal, same hash, interchangeable as a dict/set key. assert a is not b assert a == b assert hash(a) == hash(b) + assert {a: "x"}[b] == "x" + # Different config -> not equal. + assert a != factory(**other_kwargs) - c = factory(**other_kwargs) - assert a != c - # Usage flags participate in the value. - d = factory() - d.set_usage(rowwise=False, columnwise=False) - assert a != d - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_usable_in_set_and_dict(factory, other_kwargs): - a = factory() - b = factory() - c = factory(**other_kwargs) - assert len({a, b, c}) == 2 - mapping = {a: "x"} - assert mapping[b] == "x" - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_cross_type_inequality(factory, other_kwargs): - a = factory() - other = _current_scaling() if not isinstance(a, Float8CurrentScalingQuantizer) else _mxfp8() - assert a != other - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_fx_repr_roundtrip(factory, other_kwargs): - """``__fx_repr__`` returns an evaluable expression rebuilding an equal object.""" - a = factory() + # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. repr_str, globals_ = a.__fx_repr__() - assert isinstance(repr_str, str) - assert isinstance(globals_, dict) rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used - assert rebuilt == a - assert rebuilt is not a + assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_value_key_reconstruction(factory, other_kwargs): - a = factory() - rebuilt = _quantizer_from_value_key(a._value_key()) - assert type(rebuilt) is type(a) - assert rebuilt == a - # The deprecated amax-reduction process group is never carried in the value. + # The deprecated amax-reduction group is never part of the value. assert getattr(rebuilt, "amax_reduction_group", None) is None - - -def test_quantizer_delayed_scaling_keeps_identity_semantics(): - """Float8Quantizer holds live tensors -> identity (not value) semantics.""" - scale = torch.ones(1) - amax = torch.zeros(1) - a = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) - b = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) - assert a._value_fields() is None - assert a == a - assert a != b # distinct instances are not equal despite identical config - assert hash(a) == object.__hash__(a) - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_registration_is_idempotent_and_tolerant(factory, other_kwargs): - """Re-registering must not raise, regardless of PyTorch opaque-object support.""" - cls = type(factory()) - register_value_opaque_quantizer(cls) - register_value_opaque_quantizer(cls) - - if not _opaque_available: - pytest.skip("PyTorch build without opaque-object API") - from torch._library.opaque_object import is_opaque_value_type - - assert is_opaque_value_type(cls) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4Quantizer requires CUDA") -def test_quantizer_nvfp4_value_semantics(): - a = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) - b = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) - assert a == b - assert hash(a) == hash(b) - - c = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, with_rht=not a.with_rht) - assert a != c - - rebuilt = _quantizer_from_value_key(a._value_key()) - assert rebuilt == a - assert rebuilt.amax_reduction_group is None - - repr_str, globals_ = a.__fx_repr__() - assert eval(repr_str, dict(globals_)) == a # pylint: disable=eval-used diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo.py index 26aff6f59c..16f73920ae 100644 --- a/transformer_engine/pytorch/dynamo.py +++ b/transformer_engine/pytorch/dynamo.py @@ -9,12 +9,14 @@ * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used by FX codegen and registers the quantizer class with - ``torch._library.opaque_object``. It is a no-op (other than populating the - local registry) on PyTorch builds without the opaque-object API, so - importing Transformer Engine never fails on older PyTorch -- only - torch.compile specialization on the quantizer is unavailable there. - * :func:`_quantizer_from_value_key` -- rebuilds a quantizer constant from its - value key inside the generated FX graph. + ``torch._library.opaque_object``. It is a no-op on PyTorch builds without + the opaque-object API, so importing Transformer Engine never fails on older + PyTorch -- only torch.compile specialization on the quantizer is + unavailable there. + * :func:`_rebuild_quantizer` -- rebuilds a quantizer constant from its value + items inside the generated FX graph. The quantizer class is captured + directly in the FX globals (see :func:`_quantizer_fx_repr`), so no global + class registry is needed. The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` / ``_value_fields``) live on the quantizer itself; see @@ -22,7 +24,10 @@ See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` / -``__fx_repr__``). +``__fx_repr__``). The ``__fx_repr__`` contract -- ``(repr_string, {name: type})`` +where ``repr_string`` references the names in the dict -- is exactly how +PyTorch's own value opaque types (e.g. DTensor placements) reconstruct +themselves, including across the on-disk compile cache. """ from __future__ import annotations @@ -31,26 +36,16 @@ from .constants import DType -# Maps a quantizer class qualname to the class object. A value key stores only -# the qualname, so reconstruction looks the class up here. Populated by -# ``register_value_opaque_quantizer`` at import time of each tensor module; this -# avoids importing the tensor modules into this module (which would create an -# import cycle). -_QUANTIZER_VALUE_REGISTRY: Dict[str, type] = {} - - -def _quantizer_from_value_key(key: Tuple[Any, ...]) -> Any: - """Rebuild a tensorless quantizer from its value key. +def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: + """Rebuild a tensorless quantizer of type *cls* from its value items. Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the generated FX code calls this to materialize the quantizer constant. The deprecated amax-reduction process group is never part of the value, so a reconstructed quantizer always starts with no stored group. """ - qualname, items = key[0], key[1] - cls = _QUANTIZER_VALUE_REGISTRY[qualname] # Bypass ``__init__`` and restore the value attributes directly: the value - # key already captures every value-defining field (including derived ones), + # items already capture every value-defining field (including derived ones), # and the constructors have heterogeneous signatures / side effects. obj = cls.__new__(cls) field_names = set() @@ -71,12 +66,15 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: """``__fx_repr__`` for value-opaque quantizers (attached at registration). Returns an evaluable expression that rebuilds the quantizer via - :func:`_quantizer_from_value_key`, together with the globals needed to - evaluate it. + :func:`_rebuild_quantizer`, capturing both the helper and the quantizer + class itself in the FX globals so codegen can resolve them with no global + registry and no qualname collisions. """ + cls = type(self) + items = self._value_key()[1] return ( - f"_quantizer_from_value_key({self._value_key()!r})", - {"_quantizer_from_value_key": _quantizer_from_value_key}, + f"_rebuild_quantizer({cls.__name__}, {items!r})", + {"_rebuild_quantizer": _rebuild_quantizer, cls.__name__: cls}, ) @@ -85,16 +83,13 @@ def register_value_opaque_quantizer(cls: type) -> None: Attaches ``__fx_repr__`` and registers the class with ``torch._library.opaque_object``. Safe to call on any PyTorch build: on - versions without the opaque-object API it only records the class in the - local registry and attaches ``__fx_repr__`` (both harmless), so Transformer - Engine keeps importing and running in eager mode. + versions without the opaque-object API it only attaches ``__fx_repr__`` + (harmless), so Transformer Engine keeps importing and running in eager mode. The quantizer class must already provide value ``__eq__`` / ``__hash__`` and a non-``None`` ``_value_fields`` (see :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). """ - _QUANTIZER_VALUE_REGISTRY[cls.__qualname__] = cls - # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index f37b3cc63d..6809893a40 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -408,20 +408,6 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } - # ----- Value-object identity (torch.compile opaque value support) ----- - # A *tensorless* quantizer (one whose entire state is a handful of plain, - # reproducible scalars -- no live tensors, no process groups) behaves like a - # value: two instances with the same configuration are interchangeable. Such - # quantizers opt into value-based ``__eq__`` / ``__hash__`` by overriding - # ``_value_fields``. Quantizers that keep the default (e.g. delayed-scaling - # ``Float8Quantizer``, which holds live scale/amax tensors, and any custom - # quantizer) retain the default identity semantics. - # - # This is the eager-side half of registering the quantizer as a torch.compile - # *value* opaque type; the torch.compile glue (``__fx_repr__``, FX - # reconstruction and ``register_opaque_type``) lives in - # ``transformer_engine.pytorch.dynamo``. - #: Attributes shared by every quantizer that take part in value identity. _BASE_VALUE_FIELDS: Tuple[str, ...] = ( "rowwise_usage", @@ -433,11 +419,9 @@ def get_usages(self) -> Dict[str, bool]: def _value_fields(self) -> Optional[Tuple[str, ...]]: """Subclass-specific value-defining attribute names, or ``None``. - Returning ``None`` (the default) means the quantizer is *not* a value - object and keeps identity-based equality/hashing. Tensorless quantizers - override this to return the tuple of attribute names that, together with - :attr:`_BASE_VALUE_FIELDS`, fully determine their value (excluding - non-value state such as a deprecated amax-reduction process group). + Returning ``None`` (the default) means the quantizer cannot be represented as + a value opaque object and keeps identity-based equality/hashing. + This also means, that torch.compile will not be able to optimize the quantizer. """ return None From a06324bd4f10354d76b5fbec133f07709bd6a1da Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Sun, 7 Jun 2026 15:42:31 +0200 Subject: [PATCH 03/42] [PyTorch] Split dynamo.py into a dynamo/ package Replace the single dynamo.py module with a dynamo/ package so the torch.compile glue can grow with a clear responsibility split across the stacked branches. This branch owns the value-opaque quantizer layer. * dynamo/quantizer_opaque.py -- register_value_opaque_quantizer and helpers * dynamo/__init__.py -- re-exports the public API so callers keep importing from transformer_engine.pytorch.dynamo unchanged Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 18 ++++++++++++++++++ .../{dynamo.py => dynamo/quantizer_opaque.py} | 7 +++---- 2 files changed, 21 insertions(+), 4 deletions(-) create mode 100644 transformer_engine/pytorch/dynamo/__init__.py rename transformer_engine/pytorch/{dynamo.py => dynamo/quantizer_opaque.py} (95%) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py new file mode 100644 index 0000000000..44ca61d470 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile glue for Transformer Engine. + +Public API is re-exported here so callers keep importing from +``transformer_engine.pytorch.dynamo`` regardless of the internal module layout: + + * :mod:`.quantizer_opaque` -- make a tensorless quantizer a torch.compile + *value* opaque type (:func:`register_value_opaque_quantizer`). +""" + +from .quantizer_opaque import register_value_opaque_quantizer + +__all__ = [ + "register_value_opaque_quantizer", +] diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py similarity index 95% rename from transformer_engine/pytorch/dynamo.py rename to transformer_engine/pytorch/dynamo/quantizer_opaque.py index 16f73920ae..c7455846ee 100644 --- a/transformer_engine/pytorch/dynamo.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -2,10 +2,9 @@ # # See LICENSE for license information. -"""torch.compile glue for Transformer Engine quantizers. +"""Value-opaque quantizers for torch.compile. -This module isolates the torch.compile-specific plumbing that turns a -*tensorless* quantizer into a torch.compile **value** opaque type: +Turns a *tensorless* quantizer into a torch.compile **value** opaque type: * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used by FX codegen and registers the quantizer class with @@ -33,7 +32,7 @@ class registry is needed. from __future__ import annotations from typing import Any, Dict, Tuple -from .constants import DType +from ..constants import DType def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: From ea5b396b9e1e3ee67905c58b50e30989f5ab095d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 8 Jun 2026 12:39:40 +0200 Subject: [PATCH 04/42] [PyTorch] Raise in quantizer __fx_repr__ when a process group is stored A value-opaque quantizer must not carry live distributed state. Scan the quantizer attributes in __fx_repr__ and raise TypeError if any holds a torch.distributed.ProcessGroup (e.g. a non-None deprecated amax_reduction_group), so it cannot be silently baked into a torch.compile FX graph. Clarify the related comments accordingly. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 7 -- transformer_engine/pytorch/dynamo/__init__.py | 9 +-- .../pytorch/dynamo/quantizer_opaque.py | 70 ++++++++++--------- .../pytorch/quantized_tensor.py | 4 +- .../pytorch/tensor/float8_tensor.py | 3 +- .../pytorch/tensor/nvfp4_tensor.py | 3 +- 6 files changed, 45 insertions(+), 51 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 4491e23881..9a7a4a356a 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -392,11 +392,6 @@ def fn(inp): # --------------------------------------------------------------------------- # Value-opaque quantizers -# -# Tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling) are -# torch.compile *value* opaque types: value-based ``__eq__`` / ``__hash__`` plus -# an evaluable ``__fx_repr__`` that rebuilds an equal object (see -# ``torch._library.opaque_object`` Note [Opaque Objects]). # --------------------------------------------------------------------------- @@ -446,5 +441,3 @@ def test_quantizer_value_object(factory, other_kwargs): rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) - # The deprecated amax-reduction group is never part of the value. - assert getattr(rebuilt, "amax_reduction_group", None) is None diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 44ca61d470..aae8b9cff6 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -2,14 +2,7 @@ # # See LICENSE for license information. -"""torch.compile glue for Transformer Engine. - -Public API is re-exported here so callers keep importing from -``transformer_engine.pytorch.dynamo`` regardless of the internal module layout: - - * :mod:`.quantizer_opaque` -- make a tensorless quantizer a torch.compile - *value* opaque type (:func:`register_value_opaque_quantizer`). -""" +"""torch.compile glue for Transformer Engine.""" from .quantizer_opaque import register_value_opaque_quantizer diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index c7455846ee..55bc8326e6 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -2,46 +2,35 @@ # # See LICENSE for license information. -"""Value-opaque quantizers for torch.compile. - -Turns a *tensorless* quantizer into a torch.compile **value** opaque type: - - * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used - by FX codegen and registers the quantizer class with - ``torch._library.opaque_object``. It is a no-op on PyTorch builds without - the opaque-object API, so importing Transformer Engine never fails on older - PyTorch -- only torch.compile specialization on the quantizer is - unavailable there. - * :func:`_rebuild_quantizer` -- rebuilds a quantizer constant from its value - items inside the generated FX graph. The quantizer class is captured - directly in the FX globals (see :func:`_quantizer_fx_repr`), so no global - class registry is needed. - -The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` / -``_value_fields``) live on the quantizer itself; see -:class:`transformer_engine.pytorch.quantized_tensor.Quantizer`. - -See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a -value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` / -``__fx_repr__``). The ``__fx_repr__`` contract -- ``(repr_string, {name: type})`` -where ``repr_string`` references the names in the dict -- is exactly how -PyTorch's own value opaque types (e.g. DTensor placements) reconstruct -themselves, including across the on-disk compile cache. -""" +"""Value-opaque quantizers for torch.compile.""" from __future__ import annotations from typing import Any, Dict, Tuple -from ..constants import DType +from ..constants import DType, dist_group_type + + +def _contains_process_group(value: Any) -> bool: + """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. + + Checks the value directly and one level of ``tuple``/``list`` nesting, which + covers the shapes a quantizer value field could plausibly take. + """ + if isinstance(value, dist_group_type): + return True + if isinstance(value, (tuple, list)): + return any(_contains_process_group(item) for item in value) + return False def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: """Rebuild a tensorless quantizer of type *cls* from its value items. Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the - generated FX code calls this to materialize the quantizer constant. The - deprecated amax-reduction process group is never part of the value, so a - reconstructed quantizer always starts with no stored group. + generated FX code calls this to materialize the quantizer constant. A + quantizer that actually stores a process group never reaches this path: + ``__fx_repr__`` raises for it. The deprecated amax-reduction group is not a + value field, so the rebuilt quantizer simply has no group attribute. """ # Bypass ``__init__`` and restore the value attributes directly: the value # items already capture every value-defining field (including derived ones), @@ -53,9 +42,9 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: value = DType.cast(value) object.__setattr__(obj, name, value) field_names.add(name) - # The deprecated amax-reduction process group is excluded from the value; - # restore it as ``None`` for quantizers that still carry the fallback so - # attribute access keeps working. + # The deprecated amax-reduction group is not a value field. Quantizers that + # actually hold a group error out in ``__fx_repr__`` before reaching here, so + # this only initializes the (groupless) attribute to keep access working. if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): object.__setattr__(obj, "amax_reduction_group", None) return obj @@ -68,8 +57,23 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: :func:`_rebuild_quantizer`, capturing both the helper and the quantizer class itself in the FX globals so codegen can resolve them with no global registry and no qualname collisions. + + Raises ``TypeError`` if the quantizer stores a process group (e.g. a + non-``None`` deprecated ``amax_reduction_group``): live distributed state + must never be baked into the graph as a constant, so such a quantizer cannot + be used with ``torch.compile``. Pass the reduction group per quantize call + instead of storing it on the quantizer. """ cls = type(self) + for name, value in vars(self).items(): + if _contains_process_group(value): + raise TypeError( + f"{cls.__name__} cannot be used with torch.compile: attribute " + f"{name!r} holds a torch.distributed.ProcessGroup, which is live " + "distributed state and must not be baked into an FX graph as a " + "constant. Pass the amax reduction group per quantize call instead " + "of storing it on the quantizer." + ) items = self._value_key()[1] return ( f"_rebuild_quantizer({cls.__name__}, {items!r})", diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 6809893a40..3612c1080d 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -421,7 +421,9 @@ def _value_fields(self) -> Optional[Tuple[str, ...]]: Returning ``None`` (the default) means the quantizer cannot be represented as a value opaque object and keeps identity-based equality/hashing. - This also means, that torch.compile will not be able to optimize the quantizer. + This also means that passing such a quantizer as an argument to a custom op + causes a graph break under torch.compile, since it cannot be baked into the + FX graph as a constant. """ return None diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 56b1ecfb09..0310c3855c 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -389,7 +389,8 @@ def supports_only_rowwise_all_gather(self) -> bool: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group (not a value) and is restored as ``None`` on rebuild. + # process group (not a value). If one is actually stored, ``__fx_repr__`` + # raises so it can never be baked into a torch.compile graph. return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 573c78907c..4bca783922 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -336,7 +336,8 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group (not a value) and is restored as ``None`` on rebuild. + # process group (not a value). If one is actually stored, ``__fx_repr__`` + # raises so it can never be baked into a torch.compile graph. # ``rht_matrix_random_sign_mask_t`` is derived (from # ``_with_random_sign_mask`` and the device) but is stored verbatim so # reconstruction does not need to touch the device. From aa65e34e2ed8ba784543caa703ef7962062d3546 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 8 Jun 2026 16:17:36 +0200 Subject: [PATCH 05/42] [PyTorch] Cover NVFP4 in quantizer value-object test NVFP4Quantizer is registered as a value-opaque quantizer but was missing from the value-semantics / __fx_repr__ round-trip test. Add it to _VALUE_QUANTIZERS (skipped without CUDA, which it needs to construct). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 9a7a4a356a..3405001e04 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -27,7 +27,7 @@ from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -416,11 +416,29 @@ def _current_scaling(amax_epsilon=0.0): ) +def _nvfp4(with_rht=False): + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=with_rht, + ) + + # (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), + pytest.param( + _nvfp4, + {"with_rht": True}, + id="nvfp4", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), + reason="NVFP4Quantizer requires CUDA to construct", + ), + ), ] From e1b1db6b1a9de58a6868f1a706659733792476cc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 17:13:15 +0200 Subject: [PATCH 06/42] Reject a value quantizer that carries an amax reduction group in __eq__/__hash__ The amax reduction group is excluded from the value key, so a value quantizer that stored one would compare/hash equal to a groupless one and let torch.compile reuse a graph that skips the reduction. __eq__/__hash__ now raise (mirroring __fx_repr__, which already rejects any process-group-bearing quantizer). The group should be passed per quantize call, not stored on the quantizer. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/quantized_tensor.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 3612c1080d..b7357b0e7b 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -444,6 +444,18 @@ def _value_key(self) -> Tuple[Any, ...]: items.append((name, value)) return (type(self).__qualname__, tuple(items)) + def _check_value_has_no_amax_reduction_group(self) -> None: + # The amax reduction group is not part of the value key, so a value + # quantizer that stores one would compare/hash equal to a groupless one + # and let torch.compile reuse a graph that skips the reduction. Reject it + # (mirrors ``__fx_repr__``); pass the group per quantize call instead. + if getattr(self, "amax_reduction_group", None) is not None: + raise TypeError( + f"{type(self).__name__} with a non-None amax_reduction_group cannot be " + "used as a value object; pass the amax reduction group per quantize call " + "instead of storing it on the quantizer." + ) + def __eq__(self, other: object) -> Any: # Value quantizers compare by configuration; everything else keeps the # default identity semantics (returning ``NotImplemented`` makes Python @@ -454,11 +466,14 @@ def __eq__(self, other: object) -> Any: return NotImplemented if other._value_fields() is None: return NotImplemented + self._check_value_has_no_amax_reduction_group() + other._check_value_has_no_amax_reduction_group() return self._value_key() == other._value_key() def __hash__(self) -> int: if self._value_fields() is None: return object.__hash__(self) + self._check_value_has_no_amax_reduction_group() return hash(self._value_key()) From 8c33d0ec4dd8ff255a75827fd643d3619e6ae9af Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 18:03:29 +0200 Subject: [PATCH 07/42] Recognize value-opaque quantizers via a class flag Add is_value_opaque_quantizer() + the _te_compile_value_opaque flag stamped at registration, so dynamo-traced code can detect registered quantizers (and fall back to eager for unregistered ones). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 3 ++- .../pytorch/dynamo/quantizer_opaque.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index aae8b9cff6..ee860c78e3 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -4,8 +4,9 @@ """torch.compile glue for Transformer Engine.""" -from .quantizer_opaque import register_value_opaque_quantizer +from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer __all__ = [ "register_value_opaque_quantizer", + "is_value_opaque_quantizer", ] diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 55bc8326e6..6cb9552b71 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,6 +10,17 @@ from ..constants import DType, dist_group_type +# Class attribute stamped on quantizers registered as torch.compile value-opaque +# types. +_VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" + + +def is_value_opaque_quantizer(quantizer: Any) -> bool: + """Whether *quantizer*'s class is registered as a torch.compile value-opaque + type.""" + return getattr(quantizer, _VALUE_OPAQUE_FLAG, False) + + def _contains_process_group(value: Any) -> bool: """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. @@ -93,6 +104,10 @@ def register_value_opaque_quantizer(cls: type) -> None: a non-``None`` ``_value_fields`` (see :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). """ + # Stamp the class so it can be recognized as value-opaque in dynamo-traced + # code (used to fall back to eager for unregistered quantizers). + setattr(cls, _VALUE_OPAQUE_FLAG, True) + # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: From 945f62dadd18d5d0bb7b68c54f4bfb9767699521 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 11:34:37 +0200 Subject: [PATCH 08/42] Address review: narrow opaque-type except, add fullgraph test, fix nvfp4 value key - Narrow register_opaque_type except to (RuntimeError, TypeError): the API is already imported above, so ImportError/AttributeError there only mask real errors. - Add test_quantizer_value_object_fullgraph exercising torch.compile(fullgraph=True) end-to-end to verify opaque-type registration took effect. - Restore missing NVFP4Quantizer._with_random_sign_mask assignment required by _value_fields()/_value_key(). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 18 ++++++++++++++++++ .../pytorch/dynamo/quantizer_opaque.py | 5 +++-- .../pytorch/tensor/nvfp4_tensor.py | 1 + 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 3405001e04..5e1f753f08 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -459,3 +459,21 @@ def test_quantizer_value_object(factory, other_kwargs): rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) + + +@pytest.mark.skipif( + not _opaque_available, + reason="torch.compile opaque-object support requires PyTorch >= 2.11", +) +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory, other_kwargs): + """Quantizer survives torch.compile(fullgraph=True) - verifies registration took effect.""" + + def fn(quantizer): + return quantizer + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + quantizer = factory() + assert compiled(quantizer) is quantizer diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 6cb9552b71..6d630d665c 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -128,6 +128,7 @@ def register_value_opaque_quantizer(cls: type) -> None: try: register_opaque_type(cls, typ="value") - except (ImportError, AttributeError, RuntimeError, TypeError): - # Tolerate partial / experimental opaque-object support. + except (RuntimeError, TypeError): + # Keep TE importable: registration must never crash the import, e.g. on + # PyTorch versions with only partial / experimental opaque-object support. pass diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 4bca783922..ffc5f97eca 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -174,6 +174,7 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") + self._with_random_sign_mask = with_random_sign_mask self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) From e3c8f430883762ac7be289616bf79f40eb522a0d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 12:05:01 +0200 Subject: [PATCH 09/42] Restore NVFP4 rht_matrix on value-key rebuild; assert quantize round-trip _rebuild_quantizer only restores value-key fields, so a reconstructed NVFP4Quantizer was missing the derived rht_matrix tensor (not hashable, so not in the value key) and failed at copy()/quantize time. Add a _rebuild_derived_state hook (called by _rebuild_quantizer) that NVFP4Quantizer uses to rebuild rht_matrix from _with_random_sign_mask (lru_cache -> cheap). Extend test_quantizer_value_object to also quantize with the original and the rebuilt quantizer and require bit-exact results (gated on HW support), so a field the kernel needs but the value key omits can no longer slip through. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 27 +++++++++++++++++-- .../pytorch/dynamo/quantizer_opaque.py | 5 ++++ .../pytorch/tensor/nvfp4_tensor.py | 10 +++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 5e1f753f08..da4f96d2f3 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -416,7 +416,10 @@ def _current_scaling(amax_epsilon=0.0): ) -def _nvfp4(with_rht=False): +def _nvfp4(with_rht=True): + # Default with_rht=True so the quantize round-trip below exercises the + # derived ``rht_matrix`` tensor (the field most likely to be dropped on + # value-key reconstruction). return NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, rowwise=True, @@ -425,6 +428,17 @@ def _nvfp4(with_rht=False): ) +def _hw_available(quantizer): + """Whether this HW can actually run the quantize kernel for *quantizer*.""" + if isinstance(quantizer, MXFP8Quantizer): + return mxfp8_available + if isinstance(quantizer, NVFP4Quantizer): + return nvfp4_available + if isinstance(quantizer, Float8BlockQuantizer): + return fp8_block_scaling_available + return fp8_available # Float8CurrentScalingQuantizer + + # (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), @@ -432,7 +446,7 @@ def _nvfp4(with_rht=False): pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), pytest.param( _nvfp4, - {"with_rht": True}, + {"with_rht": False}, id="nvfp4", marks=pytest.mark.skipif( not torch.cuda.is_available(), @@ -460,6 +474,15 @@ def test_quantizer_value_object(factory, other_kwargs): assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) + # The rebuilt quantizer must also *behave* identically, not just compare + # equal: equality only looks at the value key, so a field the kernel needs + # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would + # slip through the checks above and only blow up at quantize time. Run the + # real quantize kernel on both and require bit-exact results. + if torch.cuda.is_available() and _hw_available(a): + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) + @pytest.mark.skipif( not _opaque_available, diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 6d630d665c..97532b2790 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -58,6 +58,11 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: # this only initializes the (groupless) attribute to keep access working. if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): object.__setattr__(obj, "amax_reduction_group", None) + # Restore non-value derived state that ``__init__`` would normally build but + # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor). + finalize = getattr(obj, "_rebuild_derived_state", None) + if finalize is not None: + finalize() return obj diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ffc5f97eca..f2f30cdcc5 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -186,6 +186,16 @@ def __getstate__(self): state["amax_reduction_group"] = None return state + def _rebuild_derived_state(self) -> None: + """Restore the derived ``rht_matrix`` after value-key reconstruction. + + ``rht_matrix`` is a ``torch.Tensor`` built from ``_with_random_sign_mask`` + and the device, so it cannot be part of the (hashable) value key. + ``_rebuild_quantizer`` calls this hook to rebuild it; the ``lru_cache`` on + :func:`get_rht_matrix` makes an already-seen (flag, device) a cheap hit. + """ + self.rht_matrix = get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device()) + def update_quantized( self, src: torch.Tensor, From 3f6862137b30aa30c2980cc1cb9cc750599b02fa Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 14:46:49 +0200 Subject: [PATCH 10/42] Enforce process-group rejection in _value_key, not __fx_repr__; add test Move the ProcessGroup guard out of the (overridable) __fx_repr__ into Quantizer._value_key -- the single point every value-materialization path (__eq__/__hash__/__fx_repr__) goes through -- so a custom __fx_repr__ can no longer bypass it. Generalizes the old amax-only check to any field holding a ProcessGroup. Add a test that a value quantizer carrying a live group raises. Addresses review on NVIDIA#3152. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 21 ++++++++ .../pytorch/dynamo/quantizer_opaque.py | 34 +++---------- .../pytorch/quantized_tensor.py | 50 +++++++++++++------ 3 files changed, 61 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index da4f96d2f3..dc98e13708 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -484,6 +484,27 @@ def test_quantizer_value_object(factory, other_kwargs): torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) +def test_value_quantizer_rejects_process_group(): + """A value quantizer holding a live ProcessGroup must refuse to be turned + into a value key / FX constant (raise), not silently drop the group.""" + import torch.distributed as dist # pylint: disable=import-outside-toplevel + + created = not dist.is_initialized() + if created: + dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1) + try: + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + q.amax_reduction_group = dist.group.WORLD + # Every value-materialization path must reject it (hash, eq, __fx_repr__). + with pytest.raises(TypeError): + hash(q) + with pytest.raises(TypeError): + q.__fx_repr__() + finally: + if created: + dist.destroy_process_group() + + @pytest.mark.skipif( not _opaque_available, reason="torch.compile opaque-object support requires PyTorch >= 2.11", diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 97532b2790..37e718689d 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -7,7 +7,7 @@ from __future__ import annotations from typing import Any, Dict, Tuple -from ..constants import DType, dist_group_type +from ..constants import DType # Class attribute stamped on quantizers registered as torch.compile value-opaque @@ -21,19 +21,6 @@ def is_value_opaque_quantizer(quantizer: Any) -> bool: return getattr(quantizer, _VALUE_OPAQUE_FLAG, False) -def _contains_process_group(value: Any) -> bool: - """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. - - Checks the value directly and one level of ``tuple``/``list`` nesting, which - covers the shapes a quantizer value field could plausibly take. - """ - if isinstance(value, dist_group_type): - return True - if isinstance(value, (tuple, list)): - return any(_contains_process_group(item) for item in value) - return False - - def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: """Rebuild a tensorless quantizer of type *cls* from its value items. @@ -74,22 +61,13 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: class itself in the FX globals so codegen can resolve them with no global registry and no qualname collisions. - Raises ``TypeError`` if the quantizer stores a process group (e.g. a - non-``None`` deprecated ``amax_reduction_group``): live distributed state - must never be baked into the graph as a constant, so such a quantizer cannot - be used with ``torch.compile``. Pass the reduction group per quantize call - instead of storing it on the quantizer. + Raises ``TypeError`` (via :meth:`Quantizer._value_key`) if the quantizer + stores a process group (e.g. a non-``None`` deprecated + ``amax_reduction_group``): live distributed state must never be baked into + the graph as a constant. Pass the reduction group per quantize call instead + of storing it on the quantizer. """ cls = type(self) - for name, value in vars(self).items(): - if _contains_process_group(value): - raise TypeError( - f"{cls.__name__} cannot be used with torch.compile: attribute " - f"{name!r} holds a torch.distributed.ProcessGroup, which is live " - "distributed state and must not be baked into an FX graph as a " - "constant. Pass the amax reduction group per quantize call instead " - "of storing it on the quantizer." - ) items = self._value_key()[1] return ( f"_rebuild_quantizer({cls.__name__}, {items!r})", diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index b7357b0e7b..033a35f8e1 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -16,6 +16,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe +from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, @@ -23,6 +24,19 @@ ) +def _contains_process_group(value: Any) -> bool: + """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. + + Checks the value directly and one level of ``tuple``/``list`` nesting, which + covers the shapes a quantizer value field could plausibly take. + """ + if isinstance(value, dist_group_type): + return True + if isinstance(value, (tuple, list)): + return any(_contains_process_group(item) for item in value) + return False + + # Custom ops that should pass through __torch_dispatch__ without unwrapping # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that # handle quantized tensors internally. @@ -427,6 +441,24 @@ def _value_fields(self) -> Optional[Tuple[str, ...]]: """ return None + def _check_value_has_no_process_group(self) -> None: + # A value quantizer is baked into the FX graph as a constant via its + # value key, which cannot carry live distributed state. Enforced here -- + # the single point every value-materialization path (``__eq__`` / + # ``__hash__`` / ``__fx_repr__``) goes through -- so a custom + # ``__fx_repr__`` cannot bypass it. Reject any field holding a + # ProcessGroup (e.g. the deprecated ``amax_reduction_group``) rather than + # silently dropping it; pass the reduction group per quantize call. + for name, value in vars(self).items(): + if _contains_process_group(value): + raise TypeError( + f"{type(self).__name__} cannot be used as a torch.compile value " + f"object: attribute {name!r} holds a torch.distributed.ProcessGroup, " + "which is live distributed state and must not be baked into an FX " + "graph. Pass the amax reduction group per quantize call instead of " + "storing it on the quantizer." + ) + def _value_key(self) -> Tuple[Any, ...]: """Hashable, reproducible key identifying this quantizer's value. @@ -434,6 +466,7 @@ def _value_key(self) -> Tuple[Any, ...]: """ fields = self._value_fields() # pylint: disable=assignment-from-none assert fields is not None, f"{type(self).__name__} is not a value quantizer" + self._check_value_has_no_process_group() items = [] for name in self._BASE_VALUE_FIELDS + tuple(fields): value = getattr(self, name) @@ -444,36 +477,21 @@ def _value_key(self) -> Tuple[Any, ...]: items.append((name, value)) return (type(self).__qualname__, tuple(items)) - def _check_value_has_no_amax_reduction_group(self) -> None: - # The amax reduction group is not part of the value key, so a value - # quantizer that stores one would compare/hash equal to a groupless one - # and let torch.compile reuse a graph that skips the reduction. Reject it - # (mirrors ``__fx_repr__``); pass the group per quantize call instead. - if getattr(self, "amax_reduction_group", None) is not None: - raise TypeError( - f"{type(self).__name__} with a non-None amax_reduction_group cannot be " - "used as a value object; pass the amax reduction group per quantize call " - "instead of storing it on the quantizer." - ) - def __eq__(self, other: object) -> Any: # Value quantizers compare by configuration; everything else keeps the # default identity semantics (returning ``NotImplemented`` makes Python - # fall back to identity). + # fall back to identity). ``_value_key`` rejects a stored ProcessGroup. if self is other: return True if self._value_fields() is None or type(self) is not type(other): return NotImplemented if other._value_fields() is None: return NotImplemented - self._check_value_has_no_amax_reduction_group() - other._check_value_has_no_amax_reduction_group() return self._value_key() == other._value_key() def __hash__(self) -> int: if self._value_fields() is None: return object.__hash__(self) - self._check_value_has_no_amax_reduction_group() return hash(self._value_key()) From 32d17683f036390a6825902ebacba0e7bf592050 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:08:39 +0200 Subject: [PATCH 11/42] Strengthen fullgraph test: quantize/dequantize via a custom op, not passthrough Replace the trivial pass-through fullgraph test with one that drives each production quantizer through a minimal custom op (quantize + dequantize) under torch.compile(fullgraph=True) and compares to eager -- so the opaque-type registration is actually exercised inside the graph (a graph break would make fullgraph=True raise). Op registration sits right before the test. Also drop stale comments referencing the old __fx_repr__-side process-group guard. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 54 ++++++++++++++++--- .../pytorch/dynamo/quantizer_opaque.py | 10 ++-- .../pytorch/tensor/nvfp4_tensor.py | 3 +- 3 files changed, 51 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index dc98e13708..63cb82eca8 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -505,19 +505,59 @@ def test_value_quantizer_rejects_process_group(): dist.destroy_process_group() +if _opaque_available: + # A minimal custom op taking a tensor and a value-opaque quantizer that + # quantizes + dequantizes inside it, one per production quantizer class. + # ``test_quantizer_value_object_fullgraph`` drives this under + # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the + # graph -- proving the opaque-type registration took effect (a graph break + # would make ``fullgraph=True`` raise). + _qdq_lib = torch.library.Library("test_te_qdq", "DEF") + _QDQ_OPS = {} + for _qcls in ( + MXFP8Quantizer, + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + NVFP4Quantizer, + ): + _op = f"qdq_{_qcls.__name__}" + _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor") + + @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib) + def _qdq_impl(x, q): + return q(x).dequantize() + + @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib) + def _qdq_fake(x, q): + return torch.empty_like(x) + + _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op) + + @pytest.mark.skipif( not _opaque_available, reason="torch.compile opaque-object support requires PyTorch >= 2.11", ) @pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) def test_quantizer_value_object_fullgraph(factory, other_kwargs): - """Quantizer survives torch.compile(fullgraph=True) - verifies registration took effect.""" + """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. - def fn(quantizer): - return quantizer + A custom op quantizes+dequantizes with the (opaque value) quantizer; the + compiled result must match eager. ``fullgraph=True`` raises on any graph + break, so this proves the opaque-type registration actually took effect -- + unlike merely passing the quantizer through. + """ + q = factory() + if not (torch.cuda.is_available() and _hw_available(q)): + pytest.skip("format not supported on this HW") - torch._dynamo.reset() - compiled = torch.compile(fn, fullgraph=True) + op = _QDQ_OPS[type(q)] + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def fn(inp): + return op(inp, q) - quantizer = factory() - assert compiled(quantizer) is quantizer + ref = fn(x) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 37e718689d..409cb979d1 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -25,10 +25,7 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: """Rebuild a tensorless quantizer of type *cls* from its value items. Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the - generated FX code calls this to materialize the quantizer constant. A - quantizer that actually stores a process group never reaches this path: - ``__fx_repr__`` raises for it. The deprecated amax-reduction group is not a - value field, so the rebuilt quantizer simply has no group attribute. + generated FX code calls this to materialize the quantizer constant. """ # Bypass ``__init__`` and restore the value attributes directly: the value # items already capture every value-defining field (including derived ones), @@ -40,9 +37,8 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: value = DType.cast(value) object.__setattr__(obj, name, value) field_names.add(name) - # The deprecated amax-reduction group is not a value field. Quantizers that - # actually hold a group error out in ``__fx_repr__`` before reaching here, so - # this only initializes the (groupless) attribute to keep access working. + # The deprecated amax-reduction group is not a value field; initialize it to + # None so attribute access keeps working on the rebuilt quantizer. if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): object.__setattr__(obj, "amax_reduction_group", None) # Restore non-value derived state that ``__init__`` would normally build but diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index f2f30cdcc5..d9cd534606 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -347,8 +347,7 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group (not a value). If one is actually stored, ``__fx_repr__`` - # raises so it can never be baked into a torch.compile graph. + # process group, not a value (``_value_key`` rejects a stored group). # ``rht_matrix_random_sign_mask_t`` is derived (from # ``_with_random_sign_mask`` and the device) but is stored verbatim so # reconstruction does not need to touch the device. From 28bde9e7040a8b4c87b992cd5717b7518d19e000 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:29:11 +0200 Subject: [PATCH 12/42] Clarify comments: rht_matrix_random_sign_mask_t derivation; why the opaque flag - rht_matrix_random_sign_mask_t is a device-independent int derived from _with_random_sign_mask (the device only places a throwaway tensor); fix the misleading comment. - Explain why registration uses a class attribute, not a registry set: is_value_opaque_quantizer is traced inside the compile graph and dynamo can bake a getattr constant but cannot do 'type(q) in set' on the opaque class. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/quantizer_opaque.py | 8 ++++++-- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 6 +++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 409cb979d1..33e831ef38 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,8 +10,12 @@ from ..constants import DType -# Class attribute stamped on quantizers registered as torch.compile value-opaque -# types. +# Registration marks the class with this attribute instead of recording it in a +# module-level set. ``is_value_opaque_quantizer`` runs *inside* the torch.compile +# graph (``Linear.forward`` consults it): Dynamo can trace a ``getattr`` on the +# opaque quantizer and bake the result as a constant, but cannot evaluate +# ``type(q) in some_set`` -- it has no equality/hash rules for the opaque class +# object, so a set/dict lookup graph-breaks under ``fullgraph=True``. _VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index d9cd534606..c8c0a7b854 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -348,9 +348,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated # process group, not a value (``_value_key`` rejects a stored group). - # ``rht_matrix_random_sign_mask_t`` is derived (from - # ``_with_random_sign_mask`` and the device) but is stored verbatim so - # reconstruction does not need to touch the device. + # ``rht_matrix_random_sign_mask_t`` is a device-independent int derived + # from ``_with_random_sign_mask``; kept in the key so the rebuilt + # quantizer carries it without recomputation. return ( "dtype", "with_rht", From 2c3c5df0fb3488b8e95670066ac53f958640a6de Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:30:44 +0200 Subject: [PATCH 13/42] Reword opaque-flag comment: self-contained, no Linear reference Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 33e831ef38..98349e12ba 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,12 +10,12 @@ from ..constants import DType -# Registration marks the class with this attribute instead of recording it in a -# module-level set. ``is_value_opaque_quantizer`` runs *inside* the torch.compile -# graph (``Linear.forward`` consults it): Dynamo can trace a ``getattr`` on the -# opaque quantizer and bake the result as a constant, but cannot evaluate -# ``type(q) in some_set`` -- it has no equality/hash rules for the opaque class -# object, so a set/dict lookup graph-breaks under ``fullgraph=True``. +# Registration marks the class with this attribute rather than recording it in a +# module-level set. It looks odd but is a deliberate workaround: the check must +# stay traceable when it runs inside a torch.compile graph -- Dynamo can bake a +# ``getattr`` on the opaque quantizer into a constant, but cannot evaluate +# ``type(q) in some_set`` (no equality/hash rules for the opaque class object), +# which would graph-break under ``fullgraph=True``. _VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" From 826f271ebb28466cdd94b9ee5cc6875eb61d4a97 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:38:53 +0200 Subject: [PATCH 14/42] Cover is_opaque_value_type with the import-safety guard too is_opaque_value_type(cls) sat between the import guard and the register_opaque_type guard, so on a partial/experimental opaque-object build it could raise RuntimeError/TypeError and crash TE import. Move it inside the same except so the 'registration never crashes import' promise holds for both calls. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/quantizer_opaque.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 98349e12ba..8b8b3caa69 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -106,12 +106,11 @@ def register_value_opaque_quantizer(cls: type) -> None: # still work; torch.compile specialization on the quantizer does not. return - if is_opaque_value_type(cls): - return - try: - register_opaque_type(cls, typ="value") + if not is_opaque_value_type(cls): + register_opaque_type(cls, typ="value") except (RuntimeError, TypeError): - # Keep TE importable: registration must never crash the import, e.g. on - # PyTorch versions with only partial / experimental opaque-object support. + # Keep TE importable: neither the opaque-type query nor the registration + # must crash the import, e.g. on PyTorch versions with only partial / + # experimental opaque-object support. pass From 4cd244eee50b2ac3b12ed2f77ac02964ca99f6c5 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 29 Jun 2026 18:35:12 +0200 Subject: [PATCH 15/42] [PyTorch] Expert Parallelism: PyTorch wrapper + autograd ops with symm-mem zero-copy (#3035) * Expert Parallelism: PyTorch wrapper + autograd ops with symm-mem zero-copy Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- build_tools/pytorch.py | 10 + examples/pytorch/ep/bench/ep_bench.py | 421 ++++++++++++ examples/pytorch/ep/bench/run_ep_bench.sh | 72 +++ .../pytorch/ep/bench/run_nccl_ep_bench.sh | 62 ++ examples/pytorch/ep/ep_moe.py | 255 ++++++++ examples/pytorch/ep/run_test_ep.sh | 37 ++ qa/L1_pytorch_distributed_unittest/test.sh | 1 + setup.py | 56 +- tests/pytorch/distributed/run_ep.py | 521 +++++++++++++++ tests/pytorch/distributed/run_test_ep.sh | 74 +++ tests/pytorch/distributed/test_ep.py | 31 + transformer_engine/pytorch/csrc/extensions.h | 37 ++ .../pytorch/csrc/extensions/ep.cpp | 386 +++++++++++ .../pytorch/csrc/extensions/pybind.cpp | 4 + transformer_engine/pytorch/distributed.py | 21 + transformer_engine/pytorch/ep.py | 610 ++++++++++++++++++ 16 files changed, 2575 insertions(+), 23 deletions(-) create mode 100644 examples/pytorch/ep/bench/ep_bench.py create mode 100755 examples/pytorch/ep/bench/run_ep_bench.sh create mode 100755 examples/pytorch/ep/bench/run_nccl_ep_bench.sh create mode 100644 examples/pytorch/ep/ep_moe.py create mode 100755 examples/pytorch/ep/run_test_ep.sh create mode 100644 tests/pytorch/distributed/run_ep.py create mode 100755 tests/pytorch/distributed/run_test_ep.sh create mode 100644 tests/pytorch/distributed/test_ep.py create mode 100644 transformer_engine/pytorch/csrc/extensions/ep.cpp create mode 100644 transformer_engine/pytorch/ep.py diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index e2e6d09c29..5ed4eae9d5 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -77,6 +77,16 @@ def setup_pytorch_extension( setup_mpi_flags(include_dirs, cxx_flags) + # Mirror the NCCL EP gate from setup.py / common CMake. When disabled, the + # ep.cpp source no-ops at the #ifdef boundary; without the define it would + # produce undefined references to nvte_ep_*. + if bool(int(os.getenv("NVTE_WITH_NCCL_EP", "1"))): + cxx_flags.append("-DNVTE_WITH_NCCL_EP") + # PyTorch's symm-mem headers gate the NCCL_HAS_SYMMEM_* feature macros on + # USE_NCCL. The EP extension shares the symm-mem NCCL comm with torch, so + # it needs those macros visible. + cxx_flags.append("-DUSE_NCCL") + library_dirs = [] libraries = [] if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", 0))): diff --git a/examples/pytorch/ep/bench/ep_bench.py b/examples/pytorch/ep/bench/ep_bench.py new file mode 100644 index 0000000000..2b7a2c62e5 --- /dev/null +++ b/examples/pytorch/ep/bench/ep_bench.py @@ -0,0 +1,421 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""PyTorch EP perf bench: raw and autograd dispatch/combine on a single EP group. + +One process per GPU; launched via run_ep_bench.sh (torchrun). + +Stages (each timed in its own loop): + - dispatch_raw: _ep_dispatch_raw (no autograd, no prepare) + - ep_dispatch_fwd: ep_dispatch forward only + - ep_dispatch_fwd_bwd: ep_dispatch + backward on 0.5 * ||recv||^2 + - combine_raw: _ep_combine_raw (no autograd) + - ep_combine_fwd: ep_combine forward only + - ep_combine_fwd_bwd: ep_combine + backward + +ep_prepare runs once outside the timed loops. --kineto DIR dumps a Chrome +trace plus a per-kernel summary on rank 0. +""" + +import argparse +import gc +import os +import sys +import time +from contextlib import nullcontext + +import numpy as np +import torch +import torch.distributed as dist + +from transformer_engine.pytorch.ep import ( + EpBuffer, + ep_bootstrap, + ep_combine, + ep_dispatch, + ep_finalize, + ep_prepare, + _ep_combine_raw, + _ep_dispatch_raw, +) + + +def _parse_args(): + p = argparse.ArgumentParser(description="TE-PyTorch EP perf bench") + p.add_argument("--tokens-per-rank", type=int, default=8192) + p.add_argument("--hidden", type=int, default=7168) + p.add_argument("--top-k", type=int, default=8) + p.add_argument("--num-experts", type=int, default=256) + p.add_argument("--warmup", type=int, default=2) + p.add_argument("--iters", type=int, default=10) + p.add_argument( + "--max-num-sms", + type=int, + default=0, + help="Max SMs for dispatch/combine/preprocess kernels (0 = auto).", + ) + p.add_argument( + "--kineto", + default=None, + help="If set, dump a Kineto Chrome trace + per-kernel summary into this dir (rank 0).", + ) + p.add_argument( + "--cuda-graph", + action="store_true", + default=False, + help=( + "Capture each stage into a CUDA graph and time replay() instead of the eager call. " + "Raw + fwd-only stages use torch.cuda.graph; fwd+bwd stages use " + "torch.cuda.make_graphed_callables to capture forward and backward together." + ), + ) + p.add_argument( + "--mode-label", + default=None, + help="Optional suffix for NVTX range names (e.g. 'fused' / 'unfused').", + ) + p.add_argument( + "--caller-provides-dispatch-recv-tokens", + action="store_true", + default=False, + help="Supply recv_tokens to ep_dispatch instead of letting EpBuffer own it.", + ) + p.add_argument( + "--caller-provides-grad-expert-out", + action="store_true", + default=False, + help="Supply the combine backward grad buffer to ep_combine.", + ) + return p.parse_args() + + +def _nvtx_funcs(): + """Return push/pop helpers using torch.cuda.nvtx if available.""" + try: + push = torch.cuda.nvtx.range_push + pop = torch.cuda.nvtx.range_pop + return push, pop + except AttributeError: + return lambda _name: None, lambda: None + + +def _device_sm() -> int: + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor + + +def _make_inputs(rank, world_size, T, H, K, E, device): + """Round-robin identity routing + uniform top-k weights.""" + topk_idx = np.empty((T, K), dtype=np.int64) + for t in range(T): + for k in range(K): + topk_idx[t, k] = ((rank * T + t) * K + k) % E + rng = np.random.default_rng(seed=42 + rank) + tokens_np = (rng.standard_normal((T, H), dtype=np.float32) * 0.5).astype(np.float32) + return ( + torch.from_numpy(topk_idx).to(device), + torch.from_numpy(tokens_np).to(device=device, dtype=torch.bfloat16), + torch.full((T, K), 1.0 / K, dtype=torch.float32, device=device), + ) + + +def _time_stage_us(name, fn, iters, nvtx_suffix, push, pop): + """Time fn for iters iterations after one untimed warmup; returns mean us.""" + # Run iters+1 times; drop the first (autotune outlier) and frame NVTX from iter 1. + total_ns = 0 + counted = 0 + for i in range(iters + 1): + if i == 1: + push(f"{name}{nvtx_suffix}") + torch.cuda.synchronize() + t0 = time.perf_counter_ns() + fn() + torch.cuda.synchronize() + dt = time.perf_counter_ns() - t0 + if i == 0: + continue + total_ns += dt + counted += 1 + pop() + return total_ns / 1e3 / counted + + +def main(): + args = _parse_args() + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank))) + device = torch.device("cuda", torch.cuda.current_device()) + + if _device_sm() < 90: + if rank == 0: + print(f"[ep_bench] SKIPPED: EP requires SM>=90 (got SM{_device_sm()})") + dist.destroy_process_group() + return + if world_size < 4: + if rank == 0: + print(f"[ep_bench] SKIPPED: EP requires >=4 ranks (got {world_size})") + dist.destroy_process_group() + return + + ep_size = world_size + E = args.num_experts + assert E % ep_size == 0, f"num_experts ({E}) must be divisible by ep_size ({ep_size})" + num_local_experts = E // ep_size + T = args.tokens_per_rank + H = args.hidden + K = args.top_k + # Conservative cap: every token could land on every local expert. + recv_pr = world_size * T * K // 2 + if rank == 0: + print( + f"[ep_bench] world={world_size} ep={ep_size} T={T} H={H} K={K} " + f"E={E} (local={num_local_experts}) recv_pr={recv_pr}" + + (f" mode={args.mode_label}" if args.mode_label else ""), + flush=True, + ) + + ep_group = dist.new_group(ranks=list(range(world_size)), backend="nccl") + ep_bootstrap( + ep_group, + num_experts=E, + max_tokens_per_rank=T, + recv_capacity_per_rank=recv_pr, + hidden_dim=H, + max_num_sms=args.max_num_sms, + ) + + topk_idx, tokens_hbm, topk_w_hbm = _make_inputs(rank, world_size, T, H, K, E, device) + + # Caller-supplied buffers for the autograd ep_dispatch/ep_combine stages + # (normal mode -> plain tensors), reused across iters. None when not opted in. + caller_recv_tokens = ( + torch.empty(recv_pr, H, dtype=torch.bfloat16, device=device) + if args.caller_provides_dispatch_recv_tokens + else None + ) + caller_grad_expert_out = ( + torch.empty(recv_pr, H, dtype=torch.bfloat16, device=device) + if args.caller_provides_grad_expert_out + else None + ) + + buffer = EpBuffer( + top_k=K, + max_tokens_per_rank=T, + recv_capacity_per_rank=recv_pr, + hidden_dim=H, + num_local_experts=num_local_experts, + dispatch_recv_tokens=caller_recv_tokens, + combine_grad_expert_out=caller_grad_expert_out, + ) + + tokens = tokens_hbm + topk_w = topk_w_hbm + recv_tokens = torch.empty(recv_pr, H, dtype=torch.bfloat16, device=device) + recv_w = torch.empty(recv_pr, dtype=torch.float32, device=device) + + # -- Prepare once outside the timed loops ------------------------------ + ep_prepare(buffer, topk_idx) + torch.cuda.synchronize() + + # Pre-dispatch a steady recv_tokens / recv_w so combine stages have valid input. + _ep_dispatch_raw(buffer, topk_idx, tokens, topk_w, recv_tokens, recv_w) + torch.cuda.synchronize() + # fp-equivalent stand-in for an MLP output. + expert_out = recv_tokens.clone() + + nvtx_suffix = f"[{args.mode_label}]" if args.mode_label else "" + push, pop = _nvtx_funcs() + + # -- Stage closures ---------------------------------------------------- + # Persistent fwd+bwd inputs (make_graphed_callables needs stable storage). + tokens_p = tokens.detach().clone().requires_grad_(True) + eo_p = recv_tokens.detach().clone().requires_grad_(True) + + # Stand-in callables; the cuda-graph branch below swaps in graphed versions. + fwd_bwd_dispatch_fn = lambda x: ep_dispatch(buffer, x, topk_idx, topk_w)[0] # noqa: E731 + fwd_bwd_combine_fn = lambda expert_out: ep_combine(buffer, expert_out) # noqa: E731 + + def _dispatch_raw(): + _ep_dispatch_raw(buffer, topk_idx, tokens, topk_w, recv_tokens, recv_w) + + def _combine_raw(): + out_buf = torch.empty(T, H, dtype=torch.bfloat16, device=device) + _ep_combine_raw(buffer, expert_out, out_buf) + + def _ep_dispatch_fwd(): + ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w) + + def _ep_dispatch_fwd_bwd(): + tokens_p.grad = None + r = fwd_bwd_dispatch_fn(tokens_p) + (0.5 * (r * r).sum(dtype=torch.float32)).backward() + + def _ep_combine_fwd(): + ep_combine(buffer, recv_tokens) + + def _ep_combine_fwd_bwd(): + eo_p.grad = None + out = fwd_bwd_combine_fn(eo_p) + (0.5 * (out * out).sum(dtype=torch.float32)).backward() + + stages = [ + ("dispatch_raw", _dispatch_raw, True), + ("ep_dispatch_fwd", _ep_dispatch_fwd, True), + ("ep_dispatch_fwd_bwd", _ep_dispatch_fwd_bwd, False), + ("combine_raw", _combine_raw, True), + ("ep_combine_fwd", _ep_combine_fwd, True), + ("ep_combine_fwd_bwd", _ep_combine_fwd_bwd, False), + ] + # Third tuple element: True = direct torch.cuda.graph capture; False = use + # make_graphed_callables (autograd-aware) instead. + + # -- Warmup ----------------------------------------------------------- + for _ in range(args.warmup): + for _name, fn, _capt in stages: + fn() + torch.cuda.synchronize() + + # -- Optional CUDA-graph capture -------------------------------------- + # Capture each capturable stage on a side stream and time .replay() + # instead of the eager call. Outputs allocated inside the + # autograd.Function's forward go through the per-capture private pool + # so addresses stay stable across replays. + captured_runners = {} + if args.cuda_graph: + # Graph fwd+bwd of the autograd-wrapped ops via make_graphed_callables. + class _DispatchMod(torch.nn.Module): + def forward(self, x): + return ep_dispatch(buffer, x, topk_idx, topk_w)[0] + + class _CombineMod(torch.nn.Module): + def forward(self, expert_out): + return ep_combine(buffer, expert_out) + + disp_mod = _DispatchMod().cuda() + comb_mod = _CombineMod().cuda() + g_disp, g_comb = torch.cuda.make_graphed_callables( + (disp_mod, comb_mod), + ((tokens_p,), (eo_p,)), + ) + fwd_bwd_dispatch_fn = g_disp + fwd_bwd_combine_fn = g_comb + + # Direct torch.cuda.graph capture for raw + fwd-only stages. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for name, fn, direct_capturable in stages: + if not direct_capturable: + continue + fn() # prime the allocator for stable replay addresses + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + fn() + captured_runners[name] = g + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + # -- Optional Kineto profiling ---------------------------------------- + kineto_ctx = nullcontext() + if args.kineto and rank == 0: + os.makedirs(args.kineto, exist_ok=True) + kineto_ctx = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=False, + with_stack=False, + ) + + # -- Timed loops ------------------------------------------------------ + results = {} + with kineto_ctx as prof: + for name, fn, _ in stages: + runner = fn + if name in captured_runners: + # Time replay() instead of the eager call. + graph = captured_runners[name] + runner = graph.replay + results[name] = _time_stage_us(name, runner, args.iters, nvtx_suffix, push, pop) + + if rank == 0: + label = f" [{args.mode_label}]" if args.mode_label else "" + print("", flush=True) + print(f"| stage | mean wall (us){label} |", flush=True) + print("|----------------------|---------------:|", flush=True) + for name in ( + "dispatch_raw", + "ep_dispatch_fwd", + "ep_dispatch_fwd_bwd", + "combine_raw", + "ep_combine_fwd", + "ep_combine_fwd_bwd", + ): + print(f"| {name:20s} | {results[name]:14.1f} |", flush=True) + print( + "| (dispatch fwd-raw) |" + f" {results['ep_dispatch_fwd'] - results['dispatch_raw']:14.1f} |", + flush=True, + ) + print( + "| (dispatch bwd-fwd) |" + f" {results['ep_dispatch_fwd_bwd'] - results['ep_dispatch_fwd']:14.1f} |", + flush=True, + ) + print( + "| (combine fwd-raw) |" + f" {results['ep_combine_fwd'] - results['combine_raw']:14.1f} |", + flush=True, + ) + print( + "| (combine bwd-fwd) |" + f" {results['ep_combine_fwd_bwd'] - results['ep_combine_fwd']:14.1f} |", + flush=True, + ) + print("", flush=True) + + if args.kineto and rank == 0 and prof is not None: + trace_path = os.path.join(args.kineto, "ep_bench_trace.json") + prof.export_chrome_trace(trace_path) + print(f"[ep_bench] kineto trace: {trace_path}", flush=True) + print( + prof.key_averages().table(sort_by="cuda_time_total", row_limit=30), + flush=True, + ) + kern_csv = os.path.join(args.kineto, "ep_bench_kernels.csv") + with open(kern_csv, "w") as f: + f.write("name,cuda_time_us,cpu_time_us,count\n") + for evt in prof.key_averages(): + if evt.device_time_total == 0 and evt.cpu_time_total == 0: + continue + f.write(f"{evt.key},{evt.device_time_total},{evt.cpu_time_total},{evt.count}\n") + print(f"[ep_bench] per-kernel CSV: {kern_csv}", flush=True) + + # Captured CUDA graphs (when --cuda-graph) hold references to NCCL EP + # handles and per-pool streams; drop them and sync before ep_finalize, + # otherwise the post-finalize dist.barrier can deadlock against pending + # graph state. + torch.cuda.synchronize() + if args.cuda_graph: + fwd_bwd_dispatch_fn = None + fwd_bwd_combine_fn = None + captured_runners.clear() + del g_disp, g_comb, disp_mod, comb_mod + del tokens_p, eo_p, buffer, recv_tokens, recv_w, tokens, topk_w, expert_out + gc.collect() + torch.cuda.synchronize() + # Release NCCL EP's borrowed comm before torch destroys it. + ep_finalize() + dist.barrier() + dist.destroy_process_group() + sys.stdout.flush() + sys.stderr.flush() + + +if __name__ == "__main__": + main() diff --git a/examples/pytorch/ep/bench/run_ep_bench.sh b/examples/pytorch/ep/bench/run_ep_bench.sh new file mode 100755 index 0000000000..fefecd7fa9 --- /dev/null +++ b/examples/pytorch/ep/bench/run_ep_bench.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Launcher for examples/pytorch/ep/bench/ep_bench.py. +# Examples: +# bash run_ep_bench.sh # plain run, stdout only +# bash run_ep_bench.sh --cuda-graph # capture + replay each stage as a CUDA graph +# bash run_ep_bench.sh --kineto # Chrome trace + per-kernel CSV (rank 0) +# bash run_ep_bench.sh --nsys # nsys profile on rank 0 -> results/pyt_nsys.nsys-rep + +set -uo pipefail + +NSYS=0; KINETO=0; CGRAPH=0 +for a in "$@"; do + case "$a" in + --nsys) NSYS=1 ;; + --kineto) KINETO=1 ;; + --cuda-graph) CGRAPH=1 ;; + *) echo "unknown arg: $a" >&2; exit 2 ;; + esac +done +if [ "${NSYS}" -eq 1 ] && [ "${KINETO}" -eq 1 ]; then + echo "--nsys and --kineto both attach CUPTI; pick one." >&2; exit 2 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" +RESULTS="${SCRIPT_DIR}/results" +mkdir -p "${RESULTS}" +export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}" +if [ "${NUM_GPUS}" -lt 4 ]; then + echo "EP bench requires >=4 GPUs (found ${NUM_GPUS}); SKIPPING."; exit 0 +fi +if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi + +: "${TIMEOUT_S:=1800}" +: "${NCCL_EP_JIT_CACHE_DIR:=${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)}" +export NCCL_EP_JIT_CACHE_DIR +mkdir -p "${NCCL_EP_JIT_CACHE_DIR}" + +EXTRA_ARGS=() +TAG="pyt" +[ "${CGRAPH}" -eq 1 ] && EXTRA_ARGS+=(--cuda-graph) && TAG="${TAG}_cg" +if [ "${KINETO}" -eq 1 ]; then + EXTRA_ARGS+=(--kineto "${RESULTS}/kineto_${TAG}") +fi + +EP_BENCH_EXTRA_FLAGS="${EP_BENCH_EXTRA_FLAGS:-}" +LAUNCH=(torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_GPUS}" + "${SCRIPT_DIR}/ep_bench.py" "${EXTRA_ARGS[@]}" ${EP_BENCH_EXTRA_FLAGS}) + +if [ "${NSYS}" -eq 1 ]; then + NSYS_CMD=(nsys profile + --output "${RESULTS}/pyt_${TAG}_nsys" + --force-overwrite=true + --trace=cuda,nvtx + --gpu-metrics-devices=none + --cuda-um-cpu-page-faults=false + --cuda-um-gpu-page-faults=false) + echo "[run_ep_bench] launching with nsys (results/${TAG}_nsys.nsys-rep)" + timeout --foreground --signal=TERM "${TIMEOUT_S}" "${NSYS_CMD[@]}" "${LAUNCH[@]}" + RC=$? +else + timeout --foreground --signal=TERM "${TIMEOUT_S}" "${LAUNCH[@]}" + RC=$? +fi +exit $RC diff --git a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh new file mode 100755 index 0000000000..8f6da04a00 --- /dev/null +++ b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Launcher for the native NCCL EP ``ep_bench`` (baseline for PyTorch comparison). +# Usage: +# bash run_nccl_ep_bench.sh # plain run, stdout only +# bash run_nccl_ep_bench.sh --nsys # nsys → results/nccl_ep_nsys.nsys-rep + +set -uo pipefail + +NSYS=0 +for a in "$@"; do + case "$a" in + --nsys) NSYS=1 ;; + *) echo "unknown arg: $a" >&2; exit 2 ;; + esac +done + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" +RESULTS="${SCRIPT_DIR}/results" +mkdir -p "${RESULTS}" + +BIN="${TE_REPO_ROOT}/3rdparty/nccl/build/test/nccl_ep/ep_bench" +LIB="${TE_REPO_ROOT}/3rdparty/nccl/build/lib" +[ -x "${BIN}" ] || { echo "ep_bench not built at ${BIN}" >&2; exit 2; } + +NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +if [ "${NUM_GPUS}" -lt 4 ]; then + echo "NCCL EP bench requires >=4 GPUs (found ${NUM_GPUS}); SKIPPING."; exit 0 +fi +if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi + +if [ "${NSYS}" -eq 1 ]; then + ITERS=10 +else + ITERS=50 +fi +ARGS=(--algorithm ht --layout em --tokens 2048 --hidden 7168 --top-k 8 + --experts 256 --warmup 5 --iters "${ITERS}") +[ "${NSYS}" -eq 1 ] && ARGS+=(--profile) # enables NVTX ranges + cudaProfilerStart/Stop + +CMD=(/usr/local/mpi/bin/mpirun --allow-run-as-root --oversubscribe -np "${NUM_GPUS}" + -x LD_LIBRARY_PATH="${LIB}:${LD_LIBRARY_PATH:-}" + "${BIN}" "${ARGS[@]}") + +if [ "${NSYS}" -eq 1 ]; then + CMD=(nsys profile + --output "${RESULTS}/nccl_ep_nsys" + --force-overwrite=true + --capture-range=cudaProfilerApi + --capture-range-end=stop + --trace=cuda,nvtx,osrt + "${CMD[@]}") +fi + +[ "${NSYS}" -eq 1 ] && SUFFIX="_nsys" || SUFFIX="" +LOG="${RESULTS}/stdout_nccl_ep${SUFFIX}.txt" +"${CMD[@]}" 2>&1 | tee "${LOG}" +echo "Done. Log: ${LOG}" diff --git a/examples/pytorch/ep/ep_moe.py b/examples/pytorch/ep/ep_moe.py new file mode 100644 index 0000000000..149185f251 --- /dev/null +++ b/examples/pytorch/ep/ep_moe.py @@ -0,0 +1,255 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""End-to-end MoE example: dispatch -> batched expert linear -> combine, fwd + bwd. + +One process per GPU; launched via run_test_ep.sh (torchrun). +""" + +import argparse +import os +import sys + +import numpy as np +import torch +import torch.distributed as dist + +from transformer_engine.pytorch.ep import ( + EpBuffer, + ep_bootstrap, + ep_combine, + ep_dispatch, + ep_finalize, +) + + +def _parse_args(): + p = argparse.ArgumentParser(description="TE-PyTorch EP MoE example (fwd + bwd)") + p.add_argument("--num-tokens", type=int, default=8, help="Per-rank token count.") + p.add_argument("--top-k", type=int, default=2) + p.add_argument("--hidden", type=int, default=32) + p.add_argument("--hidden-out", type=int, default=32) + p.add_argument("--num-experts", type=int, default=None) + p.add_argument("--check", action="store_true", default=True) + p.add_argument( + "--benchmark", + action="store_true", + help="Time fwd over HBM buffers.", + ) + p.add_argument("--benchmark-iters", type=int, default=20) + p.add_argument("--benchmark-warmup", type=int, default=5) + p.add_argument( + "--caller-provides-dispatch-recv-tokens", + action="store_true", + default=False, + help="Supply recv_tokens to ep_dispatch instead of letting EpBuffer own it.", + ) + p.add_argument( + "--caller-provides-grad-expert-out", + action="store_true", + default=False, + help="Supply the combine backward grad buffer to ep_combine.", + ) + return p.parse_args() + + +def _make_routing(rank, T, K, E, num_local_experts): + """Deterministic routing: topk_idx[t, k] = (rank*NLE + t*K + k) % E.""" + topk_idx = np.empty((T, K), dtype=np.int64) + for t in range(T): + for k in range(K): + topk_idx[t, k] = (rank * num_local_experts + t * K + k) % E + return topk_idx + + +def _batched_expert_linear(recv_tokens, kernels, num_local_experts): + """Per-expert linear via bmm; ``recv_pr // num_local_experts`` slots per expert.""" + recv_pr, _H = recv_tokens.shape + H_out = kernels.shape[-1] + slots_per_expert = recv_pr // num_local_experts + grouped = recv_tokens.view(num_local_experts, slots_per_expert, recv_tokens.shape[-1]) + out = torch.bmm(grouped, kernels.to(grouped.dtype)) + return out.view(recv_pr, H_out) + + +def _reference_moe(tokens, topk_idx, topk_w, kernels): + T, K = topk_idx.shape + H_out = kernels.shape[-1] + out = np.zeros((T, H_out), dtype=np.float32) + for t in range(T): + tok = tokens[t].astype(np.float32) + for k in range(K): + e = int(topk_idx[t, k]) + out[t] += float(topk_w[t, k]) * (tok @ kernels[e].astype(np.float32)) + return out + + +def _reference_grad(tokens, topk_idx, topk_w, kernels): + T, K = topk_idx.shape + H = tokens.shape[-1] + ref_out = _reference_moe(tokens, topk_idx, topk_w, kernels) + grad = np.zeros((T, H), dtype=np.float32) + for t in range(T): + mixed = np.zeros_like(kernels[0]) + for k in range(K): + mixed = mixed + float(topk_w[t, k]) * kernels[int(topk_idx[t, k])] + grad[t] = ref_out[t] @ mixed.T + return ref_out, grad + + +def main(): + args = _parse_args() + + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank))) + device = torch.device("cuda", torch.cuda.current_device()) + + major, minor = torch.cuda.get_device_capability() + if major * 10 + minor < 90: + if rank == 0: + print(f"[ep_moe] SKIPPED: EP requires SM>=90 (got SM{major}{minor})") + dist.destroy_process_group() + return + + if world_size < 4: + if rank == 0: + print(f"[ep_moe] SKIPPED: EP requires >= 4 ranks (got {world_size})") + dist.destroy_process_group() + return + + ep_size = world_size + num_experts = args.num_experts if args.num_experts is not None else world_size + assert num_experts % ep_size == 0 + num_local_experts = num_experts // ep_size + T = args.num_tokens + recv_pr = ep_size * T * args.top_k + + ep_group = dist.new_group(ranks=list(range(world_size)), backend="nccl") + ep_bootstrap( + ep_group, + num_experts=num_experts, + max_tokens_per_rank=T, + recv_capacity_per_rank=recv_pr, + hidden_dim=args.hidden, + ) + try: + _run_layer( + args, rank, world_size, ep_size, num_experts, num_local_experts, T, recv_pr, device + ) + finally: + ep_finalize() + dist.destroy_process_group() + + +def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, T, recv_pr, device): + rng = np.random.default_rng(seed=42 + rank) + tokens_np = (rng.standard_normal((T, args.hidden), dtype=np.float32) * 0.5).astype(np.float32) + topk_idx_np = _make_routing(rank, T, args.top_k, num_experts, num_local_experts) + w_np = np.full((T, args.top_k), 1.0 / args.top_k, dtype=np.float32) + # Same seed across ranks -> identical kernel array everywhere. + kr = np.random.default_rng(seed=42) + kernels_np = ( + kr.standard_normal((num_experts, args.hidden, args.hidden_out), dtype=np.float32) + * (1.0 / np.sqrt(args.hidden)) + ).astype(np.float32) + + tokens = ( + torch.from_numpy(tokens_np).to(device=device, dtype=torch.bfloat16).requires_grad_(True) + ) + topk_idx = torch.from_numpy(topk_idx_np).to(device) + topk_w = torch.from_numpy(w_np).to(device) + kernels_local = torch.from_numpy( + kernels_np[rank * num_local_experts : (rank + 1) * num_local_experts] + ).to(device=device, dtype=torch.bfloat16) + + # Caller-supplied buffers (normal mode -> plain tensors), reused across iters. + recv_tokens = ( + torch.empty(recv_pr, args.hidden, dtype=torch.bfloat16, device=device) + if args.caller_provides_dispatch_recv_tokens + else None + ) + grad_expert_out = ( + torch.empty(recv_pr, args.hidden, dtype=torch.bfloat16, device=device) + if args.caller_provides_grad_expert_out + else None + ) + + buffer = EpBuffer( + top_k=args.top_k, + max_tokens_per_rank=T, + recv_capacity_per_rank=recv_pr, + hidden_dim=args.hidden, + num_local_experts=num_local_experts, + dispatch_recv_tokens=recv_tokens, + combine_grad_expert_out=grad_expert_out, + ) + + recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w) + expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts) + # Apply per-slot topk weighting before combine. + expert_out = expert_out * recv_w_out.unsqueeze(-1).to(expert_out.dtype) + out = ep_combine(buffer, expert_out) + + loss = 0.5 * (out.float() ** 2).sum() + loss.backward() + torch.cuda.synchronize() + + if rank == 0: + print( + f"[ep_moe] loss={loss.item():.4f} grad_tokens.shape={tuple(tokens.grad.shape)} " + f"ep={ep_size} num_experts={num_experts} recv_pr={recv_pr}" + ) + + if args.benchmark: + # Time dispatch + expert + combine over HBM buffers. + import time + + torch.cuda.synchronize() + dist.barrier() + for _ in range(args.benchmark_warmup): + rt, rw, _tc = ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w) + expert_out = _batched_expert_linear(rt, kernels_local, num_local_experts) + expert_out = expert_out * rw.unsqueeze(-1).to(expert_out.dtype) + ep_combine(buffer, expert_out) + torch.cuda.synchronize() + dist.barrier() + t0 = time.perf_counter() + for _ in range(args.benchmark_iters): + rt, rw, _tc = ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w) + expert_out = _batched_expert_linear(rt, kernels_local, num_local_experts) + expert_out = expert_out * rw.unsqueeze(-1).to(expert_out.dtype) + ep_combine(buffer, expert_out) + torch.cuda.synchronize() + dt_ms = (time.perf_counter() - t0) * 1000.0 / args.benchmark_iters + if rank == 0: + print(f"[ep_moe --benchmark] HBM: {dt_ms:.3f} ms/iter (iters={args.benchmark_iters})") + + if args.check: + # All-gather inputs/outputs/grads for a global reference comparison. + global_tokens = [torch.empty_like(tokens) for _ in range(world_size)] + global_topk_idx = [torch.empty_like(topk_idx) for _ in range(world_size)] + global_topk_w = [torch.empty_like(topk_w) for _ in range(world_size)] + global_out = [torch.empty_like(out) for _ in range(world_size)] + global_grad = [torch.empty_like(tokens.grad) for _ in range(world_size)] + dist.all_gather(global_tokens, tokens.detach()) + dist.all_gather(global_topk_idx, topk_idx) + dist.all_gather(global_topk_w, topk_w) + dist.all_gather(global_out, out.detach()) + dist.all_gather(global_grad, tokens.grad) + if rank == 0: + all_tokens = torch.cat(global_tokens).float().cpu().numpy() + all_idx = torch.cat(global_topk_idx).cpu().numpy() + all_w = torch.cat(global_topk_w).cpu().numpy() + all_out = torch.cat(global_out).float().cpu().numpy() + all_grad = torch.cat(global_grad).float().cpu().numpy() + ref_out, ref_grad = _reference_grad(all_tokens, all_idx, all_w, kernels_np) + np.testing.assert_allclose(all_out, ref_out, rtol=5e-2, atol=5e-2) + np.testing.assert_allclose(all_grad, ref_grad, rtol=5e-2, atol=5e-2) + print(f"[ep_moe] --check PASSED (ref_out.sum()={float(ref_out.sum()):.4f})") + + +if __name__ == "__main__": + main() + sys.exit(0) diff --git a/examples/pytorch/ep/run_test_ep.sh b/examples/pytorch/ep/run_test_ep.sh new file mode 100755 index 0000000000..13b41f4cb2 --- /dev/null +++ b/examples/pytorch/ep/run_test_ep.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +set -uo pipefail + +DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}" +if [ "${NUM_GPUS}" -lt 4 ]; then + echo "EP requires >= 4 GPUs (found ${NUM_GPUS}); SKIPPING." + exit 0 +fi +if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi + +: ${TE_PATH:=/opt/transformerengine} +: ${TEST_TIMEOUT_S:=120} + +SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py" +export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" + +# Stage JIT cubins on tmpfs for fast iteration. +: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} +export NCCL_EP_JIT_CACHE_DIR +mkdir -p "$NCCL_EP_JIT_CACHE_DIR" + +echo "*** Executing ep_moe.py across ${NUM_GPUS} GPUs (timeout=${TEST_TIMEOUT_S}s) ***" +timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \ + torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_GPUS}" \ + "${SCRIPT}" --check 2>&1 | tee stdout_ep_moe.txt +RC=${PIPESTATUS[0]} + +RET=0 +if [ "${RC}" -ne 0 ]; then RET=1; fi +if grep -qE "(^|]:)FAILED|(^|]:)Traceback" stdout_ep_moe.txt; then RET=1; fi +rm -f stdout_ep_moe.txt +exit $RET diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 7eb34a62e4..50a51353d1 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -50,6 +50,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_use python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cp_utils.xml $TE_PATH/tests/pytorch/attention/test_cp_utils.py || test_fail "test_cp_utils.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_ep.xml $TE_PATH/tests/pytorch/distributed/test_ep.py || test_fail "test_ep.py" # debug tests diff --git a/setup.py b/setup.py index b231a5c55d..1eead737d5 100644 --- a/setup.py +++ b/setup.py @@ -138,7 +138,12 @@ def setup_requirements() -> Tuple[List[str], List[str]]: def _discover_nccl_home() -> str: - """Resolve NCCL_HOME: honor env var, else probe well-known prefixes, else ldconfig.""" + """Resolve NCCL_HOME, preferring the NCCL the dynamic loader resolves at runtime. + + Probes in order: NCCL_HOME env var, ldconfig cache, well-known prefixes, then a + pip-installed nvidia-nccl-cu* wheel. To test a non-default NCCL (e.g. a wheel), set + NCCL_HOME and ensure the runtime loader resolves the same lib (e.g. LD_LIBRARY_PATH). + """ env_home = os.environ.get("NCCL_HOME") if env_home: if (Path(env_home) / "include" / "nccl.h").exists(): @@ -152,28 +157,11 @@ def _discover_nccl_home() -> str: # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu). lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") - # pip-installed NCCL (nvidia-nccl-cu* wheel) lives under nvidia/nccl in - # site-packages and has no top-level include/lib layout. - try: - import importlib.util - - spec = importlib.util.find_spec("nvidia.nccl") - if spec is not None and spec.submodule_search_locations: - pip_root = Path(next(iter(spec.submodule_search_locations))) - if (pip_root / "include" / "nccl.h").exists() and any( - (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names - ): - return str(pip_root) - except (ImportError, ValueError): - pass - - for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): - p = Path(cand) - if (p / "include" / "nccl.h").exists() and any( - (p / sub / name).exists() for sub in lib_subdirs for name in lib_names - ): - return str(p) - + # Prefer the NCCL the dynamic loader will actually resolve at runtime so the + # EP build links against the same libnccl that gets loaded. libtransformer_engine + # carries no NCCL RUNPATH, so the loader uses ldconfig/system paths; building + # against a different NCCL (e.g. a pip wheel) causes ABI mismatches. ldconfig is + # the ground truth for runtime resolution, so consult it before well-known prefixes. try: out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() for line in out.splitlines(): @@ -187,6 +175,28 @@ def _discover_nccl_home() -> str: except (subprocess.CalledProcessError, FileNotFoundError): pass + for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): + p = Path(cand) + if (p / "include" / "nccl.h").exists() and any( + (p / sub / name).exists() for sub in lib_subdirs for name in lib_names + ): + return str(p) + + # Fall back to a pip-installed NCCL (nvidia-nccl-cu* wheel) under nvidia/nccl + # in site-packages, used only when no system NCCL is present. + try: + import importlib.util + + spec = importlib.util.find_spec("nvidia.nccl") + if spec is not None and spec.submodule_search_locations: + pip_root = Path(next(iter(spec.submodule_search_locations))) + if (pip_root / "include" / "nccl.h").exists() and any( + (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names + ): + return str(pip_root) + except (ImportError, ValueError): + pass + raise RuntimeError( "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." ) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py new file mode 100644 index 0000000000..0acc00cd57 --- /dev/null +++ b/tests/pytorch/distributed/run_ep.py @@ -0,0 +1,521 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Multi-process PyTorch EP tests, launched via torchrun (one process per GPU).""" + +import os +import sys +import unittest + +import numpy as np +import torch +import torch.distributed as dist + +from transformer_engine.pytorch.ep import ( + EpBuffer, + ep_bootstrap, + ep_finalize, + ep_prepare, + ep_dispatch, + ep_combine, + symm_mem_alloc, + _ep_combine_raw, + _ep_dispatch_raw, +) + + +ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" + +# Must come after the transformer_engine import so libtransformer_engine.so is loaded. +import transformer_engine_torch as tex # noqa: F401 + + +NUM_LOCAL_EXPERTS = 2 +HIDDEN_DIM = 32 +TOP_K = 2 +TOKENS_PER_RANK = 4 + + +def _zero_copy_test_include(fn): + """Mark a test to also run in the zero-copy pass; others skip there.""" + fn._zero_copy_test_include = True + return fn + + +class _StageToSymm(torch.autograd.Function): + """Identity op that stages ``src`` into a symm-mem buffer; grad passes through. + Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine. + """ + + @staticmethod + def forward(ctx, src, symm_buf): # type: ignore[override] + symm_buf.copy_(src) + return symm_buf + + @staticmethod + def backward(ctx, g): # type: ignore[override] + return g, None + + +class _GradToSymm(torch.autograd.Function): + """Identity fwd; bwd stages the upstream grad into a symm-mem buffer and + returns it, so the next backward (dispatch_bwd) receives a symm-window grad + input — which zero-copy ncclEpCombine requires. + """ + + @staticmethod + def forward(ctx, x, symm_buf): # type: ignore[override] + ctx.symm_buf = symm_buf + return x + + @staticmethod + def backward(ctx, g): # type: ignore[override] + ctx.symm_buf.copy_(g) + return ctx.symm_buf, None + + +def _device_sm() -> int: + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor + + +def _build_ep_group(): + """EP group spanning all ranks of the default PG.""" + world_pg = dist.distributed_c10d._get_default_group() + ranks = list(range(world_pg.size())) + return dist.new_group(ranks=ranks, backend="nccl") + + +def _make_identity_inputs(rank, ep_size, device="cuda"): + """Per-rank identity routing + uniform weights so combine matches tokens.""" + T = TOKENS_PER_RANK + E = ep_size * NUM_LOCAL_EXPERTS + topk_idx = np.empty((T, TOP_K), dtype=np.int64) + base = rank * T + for t in range(T): + for k in range(TOP_K): + topk_idx[t, k] = ((base + t) * TOP_K + k) % E + tokens_np = np.linspace( + 0.1 + rank * 0.01, 0.9 + rank * 0.01, T * HIDDEN_DIM, dtype=np.float32 + ).reshape(T, HIDDEN_DIM) + topk_weights = np.full((T, TOP_K), 1.0 / TOP_K, dtype=np.float32) + return ( + torch.from_numpy(topk_idx).to(device), + torch.from_numpy(tokens_np).to(device=device, dtype=torch.bfloat16), + torch.from_numpy(topk_weights).to(device), + ) + + +class _Cfg: + rank: int + world_size: int + ep_size: int + num_experts: int + recv_capacity_per_rank: int + device: torch.device + + +def _make_cfg() -> _Cfg: + cfg = _Cfg() + cfg.rank = dist.get_rank() + cfg.world_size = dist.get_world_size() + cfg.ep_size = cfg.world_size + cfg.num_experts = NUM_LOCAL_EXPERTS * cfg.ep_size + T = TOKENS_PER_RANK + active = min(cfg.num_experts, T * cfg.ep_size * TOP_K) + overconc = cfg.num_experts // active + cfg.recv_capacity_per_rank = NUM_LOCAL_EXPERTS * max(T * cfg.ep_size * TOP_K, 16) * overconc * 2 + cfg.device = torch.device("cuda", torch.cuda.current_device()) + return cfg + + +class TestEP(unittest.TestCase): + cfg: _Cfg + ep_group: dist.ProcessGroup + + @classmethod + def setUpClass(cls): + if _device_sm() < 90: + raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})") + cls.cfg = _make_cfg() + cls.ep_group = _build_ep_group() + ep_bootstrap( + cls.ep_group, + num_experts=cls.cfg.num_experts, + max_tokens_per_rank=TOKENS_PER_RANK, + recv_capacity_per_rank=cls.cfg.recv_capacity_per_rank, + hidden_dim=HIDDEN_DIM, + zero_copy=ZERO_COPY, + ) + + def setUp(self): + # Only the zero-copy-capable tests run in the zero-copy pass. + if ZERO_COPY and not getattr( + getattr(self, self._testMethodName), "_zero_copy_test_include", False + ): + self.skipTest("not exercised in zero-copy mode") + + def _make_buffer( + self, + alignment=0, + top_k=TOP_K, + dispatch_recv_tokens=None, + combine_grad_expert_out=None, + ): + return EpBuffer( + top_k=top_k, + max_tokens_per_rank=TOKENS_PER_RANK, + recv_capacity_per_rank=self.cfg.recv_capacity_per_rank, + hidden_dim=HIDDEN_DIM, + num_local_experts=NUM_LOCAL_EXPERTS, + alignment=alignment, + dispatch_recv_tokens=dispatch_recv_tokens, + combine_grad_expert_out=combine_grad_expert_out, + ) + + def _expert_out(self, expert_out): + """Stage the combine input into symm-mem under zero-copy (combine requires it).""" + if not ZERO_COPY: + return expert_out + symm_buf = symm_mem_alloc(tuple(expert_out.shape), expert_out.dtype, self.ep_group) + return _StageToSymm.apply(expert_out, symm_buf) + + def _stage_grad_symm(self, x, symm_buf=None): + """Route x's upstream grad through a symm-mem buffer so dispatch_bwd gets + a symm-window grad input under zero-copy; passthrough otherwise. Pass a + pre-allocated symm_buf to avoid allocating during an interleaved schedule.""" + if not ZERO_COPY: + return x + if symm_buf is None: + symm_buf = symm_mem_alloc(tuple(x.shape), x.dtype, self.ep_group) + return _GradToSymm.apply(x, symm_buf) + + def _make_raw_recv(self, dtype=torch.bfloat16): + """Raw recv tensors + token_counts for the primitive tests.""" + rc = self.cfg.recv_capacity_per_rank + return ( + torch.empty(rc, HIDDEN_DIM, dtype=dtype, device=self.cfg.device), + torch.empty(rc, dtype=torch.float32, device=self.cfg.device), + torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int32, device=self.cfg.device), + ) + + @staticmethod + def _weighted(recv_tokens, recv_w): + """fp32 per-slot weighting + cast back; matches the upstream combine input.""" + mask = (recv_w != 0).to(torch.float32).unsqueeze(-1) + return (recv_tokens.float() * recv_w.unsqueeze(-1).float() * mask).to(recv_tokens.dtype) + + def _moe_step(self, buffer, topk_idx, tokens, w): + recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, w) + expert_out = self._weighted(recv_t, recv_w_out) + return ep_combine(buffer, expert_out) + + # Prepare + + def test_primitive_prepare(self): + buf = self._make_buffer() + topk_idx, _toks, _w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + token_counts = ep_prepare(buf, topk_idx) + torch.cuda.synchronize() + self.assertEqual(token_counts.shape, (NUM_LOCAL_EXPERTS,)) + local = int(token_counts.sum().item()) + total = torch.tensor([local], dtype=torch.int64, device=self.cfg.device) + dist.all_reduce(total, op=dist.ReduceOp.SUM, group=self.ep_group) + self.assertEqual(int(total.item()), self.cfg.world_size * TOKENS_PER_RANK * TOP_K) + + # Identity round-trip via raw primitives + + def test_primitive_dispatch_combine_identity(self): + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_tokens, recv_w, _ = self._make_raw_recv() + ep_prepare(buf, topk_idx) + _ep_dispatch_raw(buf, topk_idx, tokens, w, recv_tokens, recv_w) + result = torch.empty_like(tokens) + _ep_combine_raw(buf, self._weighted(recv_tokens, recv_w), result) + torch.cuda.synchronize() + torch.testing.assert_close(result.float(), tokens.float(), atol=5e-2, rtol=5e-2) + + # Autograd + + @_zero_copy_test_include + def test_dispatch_autograd(self): + """0.5*||recv_tokens||^2 ; grad_tokens equals TOP_K * tokens. Covers the + EpBuffer-owned recv tokens (symm-mem under zero-copy) and, in normal + mode, a caller-supplied recv_tokens buffer.""" + if ZERO_COPY: + cases = [("buffer_owned", None)] + else: + rt_buf, _rw_buf, _ = self._make_raw_recv() + cases = [ + ("default_alloc", None), + ("caller_recv", rt_buf), + ] + for label, recv_tokens in cases: + with self.subTest(case=label): + buf = self._make_buffer(dispatch_recv_tokens=recv_tokens) + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w) + if recv_tokens is not None: # caller-supplied recv_tokens must be used in place + self.assertEqual(rt.data_ptr(), recv_tokens.data_ptr()) + rt = self._stage_grad_symm(rt) + rw = self._stage_grad_symm(rw) + (0.5 * (rt.float() ** 2).sum() + 0.0 * rw.float().sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close( + tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2 + ) + + @_zero_copy_test_include + def test_caller_provides_dispatch_recv_tokens(self): + """Caller-supplied recv_tokens: EpBuffer adopts it (recv_topk_weights stays + owned) and ep_dispatch returns a view of the caller's buffer.""" + if ZERO_COPY: + rc = self.cfg.recv_capacity_per_rank + rt_buf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) + else: + rt_buf, _rw_buf, _ = self._make_raw_recv() + buf = self._make_buffer(dispatch_recv_tokens=rt_buf) + self.assertEqual(buf.recv_tokens_symm_buf.data_ptr(), rt_buf.data_ptr()) + if ZERO_COPY: # recv_topk_weights is always buffer-owned in zero-copy + self.assertIsNotNone(buf.recv_topk_weights_symm_buf) + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w) + self.assertEqual(rt.data_ptr(), rt_buf.data_ptr()) + rt = self._stage_grad_symm(rt) + rw = self._stage_grad_symm(rw) + (0.5 * (rt.float() ** 2).sum() + 0.0 * rw.float().sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close( + tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2 + ) + + @_zero_copy_test_include + def test_caller_provides_grad_expert_out(self): + """Caller-supplied grad_expert_out: EpBuffer adopts it as the combine + backward grad target (symm-mem under zero-copy).""" + rc = self.cfg.recv_capacity_per_rank + if ZERO_COPY: + gbuf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) + else: + gbuf = torch.empty(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) + buf = self._make_buffer(combine_grad_expert_out=gbuf) + self.assertEqual(buf.grad_expert_out_symm_buf.data_ptr(), gbuf.data_ptr()) + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) + recv_t = self._stage_grad_symm(recv_t) + recv_w = self._stage_grad_symm(recv_w) + expert_out = self._expert_out(self._weighted(recv_t, recv_w)) + out = ep_combine(buf, expert_out) + (0.5 * (out.float() ** 2).sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) + + # Multi-iter stability + + def test_dispatch_autograd_multiple_iterations(self): + """5 fwd+bwd iters on the same EpBuffer must be bit-stable.""" + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + + def one_step(): + tokens_p = tokens.detach().clone().requires_grad_(True) + out = self._moe_step(buf, topk_idx, tokens_p, w) + loss = 0.5 * (out.float() ** 2).sum() + loss.backward() + return out.detach().clone(), tokens_p.grad.detach().clone() + + out_ref, grad_ref = one_step() + torch.cuda.synchronize() + for _ in range(4): + out_i, grad_i = one_step() + torch.cuda.synchronize() + torch.testing.assert_close(out_i, out_ref, atol=0, rtol=0) + torch.testing.assert_close(grad_i, grad_ref, atol=0, rtol=0) + + # CUDA graph + + def test_cuda_graph_capture(self): + """Capture raw dispatch+combine into a CUDA graph; replay must be bit-stable.""" + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_tokens, recv_w, _ = self._make_raw_recv() + result = torch.empty_like(tokens) + + def step(): + ep_prepare(buf, topk_idx) + _ep_dispatch_raw(buf, topk_idx, tokens, w, recv_tokens, recv_w) + _ep_combine_raw(buf, self._weighted(recv_tokens, recv_w), result) + + for _ in range(3): + step() + torch.cuda.synchronize() + + # Routing is fixed per layer; prepare runs once before capture. + ep_prepare(buf, topk_idx) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + with torch.cuda.graph(graph): + _ep_dispatch_raw(buf, topk_idx, tokens, w, recv_tokens, recv_w) + _ep_combine_raw(buf, self._weighted(recv_tokens, recv_w), result) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + ref = result.clone() + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(result.float(), ref.float(), atol=0, rtol=0) + + # PP-1F1B handle isolation + + @_zero_copy_test_include + def test_pp_1f1b_two_handles(self): + """PP-1F1B interleave (F0 F1 B0 F2 B1 B2) over 3 per-microbatch buffers, + run eagerly and replayed from a CUDA graph capturing the full fwd+bwd + schedule (prepare included; routing is fixed so replay reproduces it).""" + for capture in (False, True): + with self.subTest(capture=capture): + self._run_1f1b(capture) + + def _run_1f1b(self, capture): + T, H = TOKENS_PER_RANK, HIDDEN_DIM + idx, _toks, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + scales = (0.13, 0.41, 0.77) + buffers, tokens, tokens_p = [], [], [] + for s in scales: + buffers.append(self._make_buffer()) + t = torch.full( + (T, H), s + self.cfg.rank * 0.01, dtype=torch.bfloat16, device=self.cfg.device + ) + tokens.append(t) + tokens_p.append(t.detach().clone().requires_grad_(True)) + + recv = [None, None, None] + # Per-microbatch grad-staging buffers, symm-mem under zero-copy and + # pre-allocated so nothing is allocated/freed mid-interleave. The recv + # outputs are owned by each EpBuffer (symm-mem under zero-copy). + recv_w = [None, None, None] + rc = self.cfg.recv_capacity_per_rank + if ZERO_COPY: + gbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales] + gbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales] + else: + gbuf_t = gbuf_w = [None, None, None] + + def fwd(k): + rt, rw, _ = ep_dispatch(buffers[k], tokens_p[k], idx, w) + recv[k] = self._stage_grad_symm(rt, gbuf_t[k]) + recv_w[k] = self._stage_grad_symm(rw, gbuf_w[k]) + + def bwd(k): + (0.5 * (recv[k].float() ** 2).sum() + 0.0 * recv_w[k].float().sum()).backward() + recv[k] = None + recv_w[k] = None + + def interleave(): + fwd(0) + fwd(1) + bwd(0) + fwd(2) + bwd(1) + bwd(2) + + def zero_grads(): + for tp in tokens_p: + if tp.grad is not None: + tp.grad.zero_() + + if not capture: + interleave() + else: + # Warmup on a side stream, then capture the full schedule and replay. + # Grads stay pre-allocated (zeroed, not None) so backward accumulates + # in place during both capture and replay. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + zero_grads() + interleave() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + zero_grads() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + interleave() + zero_grads() + graph.replay() + + torch.cuda.synchronize() + for k in range(3): + torch.testing.assert_close( + tokens_p[k].grad.float(), + tokens[k].float() * float(TOP_K), + atol=5e-2, + rtol=5e-2, + ) + + @_zero_copy_test_include + def test_combine_autograd(self): + """ep_combine fwd+bwd; bwd grad target is the EpBuffer symm buffer (zc) or in-flight.""" + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) + recv_t = self._stage_grad_symm(recv_t) + recv_w = self._stage_grad_symm(recv_w) + expert_out = self._expert_out(self._weighted(recv_t, recv_w)) + out = ep_combine(buf, expert_out) + (0.5 * (out.float() ** 2).sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) + + # Input validation + + def test_topk_int32_raises_clear_error(self): + buf = self._make_buffer() + topk_idx_int32 = torch.zeros( + TOKENS_PER_RANK, TOP_K, dtype=torch.int32, device=self.cfg.device + ) + with self.assertRaises(RuntimeError) as cm: + ep_prepare(buf, topk_idx_int32) + msg = str(cm.exception) + self.assertIn("topk_idx", msg) + self.assertIn(".long()", msg) + + +def _init_distributed(): + dist.init_process_group(backend="nccl") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + try: + from torch.distributed import _symmetric_memory as _symm_mem + + _symm_mem.set_backend("NCCL") + except (ImportError, RuntimeError): + pass + + +if __name__ == "__main__": + _init_distributed() + loader = unittest.TestLoader() + name_filter = os.environ.get("NVTE_EP_TEST_FILTER") + if name_filter: + loader.testMethodPrefix = name_filter + suite = loader.loadTestsFromTestCase(TestEP) + runner = unittest.TextTestRunner(stream=sys.stdout, verbosity=2) + result = runner.run(suite) + dist.barrier() + ep_finalize() + dist.destroy_process_group() + sys.exit(0 if result.wasSuccessful() else 1) diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh new file mode 100755 index 0000000000..ae40c8ba4b --- /dev/null +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Launcher for tests/pytorch/distributed/run_ep.py. Auto-detects GPU count. +# Short timeout by default to surface hangs early. + +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +if [ "${DETECTED_GPUS}" -lt 4 ]; then + echo "EP requires >= 4 GPUs (found ${DETECTED_GPUS}); SKIPPING." + exit 0 +fi + +# NCCL EP requires NVLink/NVSwitch between GPUs. +# On PCIe-only nodes (no NVLink) it falls back to the network +# transport and deadlocks, so skip cleanly there. +if ! nvidia-smi topo -m 2>/dev/null | grep -qE "\bNV[0-9]+\b"; then + echo "No NVLink between GPUs (PCIe-only fabric); NCCL EP is unsupported here. SKIPPING." + exit 0 +fi + +NUM_RANKS="${NVTE_TEST_EP_NUM_RANKS:-${DETECTED_GPUS}}" +if [ "${NUM_RANKS}" -gt 8 ]; then NUM_RANKS=8; fi + +# Short timeout to detect hangs early. +TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-120}" + +# Stage NCCL EP JIT cubins on tmpfs to keep iteration fast. +: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} +export NCCL_EP_JIT_CACHE_DIR +mkdir -p "$NCCL_EP_JIT_CACHE_DIR" + +SCRIPT="${SCRIPT_DIR}/run_ep.py" + +RET=0 + +# Run the suite once per IO mode. Modes can't be mixed in one process +# (ep_bootstrap is once-per-process), so zero-copy gets its own run; only the +# zero-copy-capable tests execute there (the rest self-skip). +run_pass() { + local label="$1" + local zc="$2" + local log="stdout_ep_${label}.txt" + echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" + # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun. + NVTE_EP_ZERO_COPY="${zc}" setsid timeout --foreground --kill-after=10 --signal=TERM \ + "${TEST_TIMEOUT_S}" \ + torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ + "${SCRIPT}" 2>&1 | tee "${log}" + local rc=${PIPESTATUS[0]} + pkill -9 -f "tests/pytorch/distributed/run_ep.py" 2>/dev/null || true + + if [ "${rc}" -ne 0 ]; then echo "[${label}] torchrun exited with ${rc}"; RET=1; fi + # Match unittest failure markers and unhandled Python tracebacks; torchrun + # prefixes per-rank stderr with "[rankN]:" so don't anchor at column 0. + if grep -qE "(^|]:)FAILED|(^|]:)Traceback" "${log}"; then RET=1; fi + if ! grep -qE "Ran [0-9]+ test|^OK$" "${log}"; then + echo "[${label}] ERROR: no test summary — likely hang or early crash" + RET=1 + fi + if [ -z "${KEEP_EP_LOGS:-}" ]; then rm -f "${log}"; fi +} + +run_pass "default" 0 +run_pass "zero_copy" 1 + +exit $RET diff --git a/tests/pytorch/distributed/test_ep.py b/tests/pytorch/distributed/test_ep.py new file mode 100644 index 0000000000..81eef9a3c1 --- /dev/null +++ b/tests/pytorch/distributed/test_ep.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Pytest driver — spawns run_ep.py under torchrun and asserts the suite passed.""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch + +TEST_ROOT = Path(__file__).parent.resolve() +WORKER = TEST_ROOT / "run_ep.py" +LAUNCHER = TEST_ROOT / "run_test_ep.sh" + + +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="EP requires >= 4 GPUs") +def test_multi_process_ep(): + """Launch the EP unit-test suite across all visible GPUs. + + Short timeout so a hang on any rank surfaces fast rather than burning CI time. + """ + timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) + proc = subprocess.run( + ["bash", str(LAUNCHER)], + env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(timeout_s)}, + timeout=timeout_s + 30, + check=False, + ) + assert proc.returncode == 0, f"EP test suite failed (rc={proc.returncode})" diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 9e26bb8c26..69c73fe2fa 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -9,6 +9,7 @@ #include +#include #include #include #include @@ -661,6 +662,42 @@ void inplace_multi_tensor_swizzle_scales_for_gemm_unchecked(std::vector +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "transformer_engine/comm_window.h" + +#ifdef NCCL_HAS_SYMMEM_SUPPORT +#include +#endif + +#include "../common.h" +#include "../extensions.h" +#include "transformer_engine/gemm.h" + +namespace transformer_engine::pytorch { + +namespace { + +// EP process group name, captured at ep_initialize. Used by the symm-mem +// window resolver below to look up SymmetricMemory for payload tensors. +// Empty until ep_initialize. +std::string g_ep_group_name; // NOLINT(runtime/string) + +// True while the EP backend holds a borrowed reference to torch's NCCL comm. +bool g_ep_initialized = false; + +// Zero-copy IO toggle captured at ep_initialize. Atomic so the Python-side +// toggle is safe against concurrent ep_dispatch/combine (which release the GIL). +std::atomic g_zero_copy_enabled{false}; + +// Sentinel returned by maybe_make_window when zero-copy is off or the tensor +// is not symm-mem-backed; the backend treats it as "no window, use staged copy". +constexpr NVTECommWindow kNoWindow = {nullptr, 0}; + +// Resolve ``t`` to an NCCL symm-mem window for the zero-copy one-sided path. +// Returns ``kNoWindow`` when symm-mem support isn't compiled in, zero-copy is +// disabled, no group is set, or ``t`` isn't symm-mem-backed; callers pass the +// resulting window unconditionally to the backend. +NVTECommWindow maybe_make_window(const at::Tensor& t) { +#ifdef NCCL_HAS_SYMMEM_SUPPORT + if (!g_zero_copy_enabled.load(std::memory_order_relaxed)) return kNoWindow; + if (g_ep_group_name.empty()) return kNoWindow; + c10::intrusive_ptr sm; + try { + sm = c10d::symmetric_memory::rendezvous(t, g_ep_group_name); + } catch (const std::exception&) { + return kNoWindow; // Tensor not symm-mem-backed; fall back to staged copy. + } + if (sm == nullptr) return kNoWindow; + auto* nccl_sm = dynamic_cast(sm.get()); + NVTE_CHECK(nccl_sm != nullptr, + "Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to " + "\"NCCL\" before allocating EP payload buffers."); + return NVTECommWindow{static_cast(nccl_sm->get_window()), + static_cast(nccl_sm->get_offset())}; +#else + (void)t; + return kNoWindow; +#endif +} + +// When zero-copy is enabled, the named tensor must be symm-mem-backed on the +// EP group. Throws a clear error otherwise. No-op when zero-copy is off or +// symm-mem support isn't compiled in. Mirrors maybe_make_window's resolution +// path but turns the "not symm-mem" outcome into a hard error. +void check_symm_mem_required(const at::Tensor& t, const char* name) { +#ifdef NCCL_HAS_SYMMEM_SUPPORT + if (!g_zero_copy_enabled.load(std::memory_order_relaxed)) return; + NVTE_CHECK(!g_ep_group_name.empty(), + "Zero-copy is enabled but EP group name is unset; call ep_initialize first."); + c10::intrusive_ptr sm; + try { + sm = c10d::symmetric_memory::rendezvous(t, g_ep_group_name); + } catch (const std::exception&) { + sm = nullptr; + } + NVTE_CHECK(sm != nullptr, "ep zero-copy: ", name, + " must be symm-mem-backed on the EP group (allocate via symm_mem_alloc)."); +#else + (void)t; + (void)name; +#endif +} + +// The backend only accepts int64 topk_idx. The PyTorch wrapper enforces this +// at the boundary so the per-step ops don't need an upcast workspace. +void check_topk_idx_int64(at::Tensor topk_idx) { + NVTE_CHECK(topk_idx.is_contiguous(), "topk_idx must be contiguous"); + NVTE_CHECK(topk_idx.scalar_type() == at::kLong, + "topk_idx must be int64; got dtype=", c10::toString(topk_idx.scalar_type()), + ". Cast with topk_idx.long() before calling."); +} + +using Shape = std::vector; + +} // namespace + +bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_relaxed); } + +// ── Bootstrap ──────────────────────────────────────────────────────────────── +// Borrows torch's NCCL host comm (from ``ProcessGroupNCCL._comm_ptr()``). +// ``group_name`` is captured for the symm-mem window resolver. + +void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t num_experts, + int64_t max_tokens_per_rank, int64_t max_recv_tokens_per_rank, + int64_t hidden_dim, int64_t max_num_sms, pybind11::object max_token_dtype, + bool zero_copy) { + NVTE_CHECK(!group_name.empty(), "group_name must be non-empty (used for symm-mem lookup)"); + NVTE_CHECK(comm_ptr != 0, "comm_ptr must be non-null (torch NCCL host comm pointer)"); + NVTE_CHECK(!g_ep_initialized, "ep_initialize called twice without ep_finalize"); + + auto ep_comm = reinterpret_cast(comm_ptr); + int ep_size = 0; + NVTE_CHECK(ncclCommCount(ep_comm, &ep_size) == ncclSuccess, "ncclCommCount failed"); + auto torch_dtype = max_token_dtype.cast(); + NVTEEpGroupConfig cfg{ + /*ep_size=*/ep_size, + /*num_experts=*/static_cast(num_experts), + /*max_tokens_per_rank=*/static_cast(max_tokens_per_rank), + /*max_recv_tokens_per_rank=*/static_cast(max_recv_tokens_per_rank), + /*hidden_dim=*/static_cast(hidden_dim), + /*max_num_sms=*/static_cast(max_num_sms), + /*max_token_dtype=*/static_cast(GetTransformerEngineDType(torch_dtype)), + /*zero_copy=*/zero_copy ? 1 : 0, + }; + nvte_ep_initialize(static_cast(ep_comm), cfg); + g_zero_copy_enabled.store(zero_copy, std::memory_order_relaxed); + g_ep_initialized = true; + g_ep_group_name = group_name; +} + +void ep_finalize() { + if (!g_ep_initialized) return; + // The borrowed comm is owned by torch's symm-mem layer; don't destroy it. + nvte_ep_shutdown(); + g_ep_initialized = false; + g_ep_group_name.clear(); + g_zero_copy_enabled.store(false, std::memory_order_relaxed); +} + +namespace { + +NVTEEpLayerConfig make_layer_cfg(int64_t top_k, int64_t dispatch_output_per_expert_alignment) { + return NVTEEpLayerConfig{ + /*top_k=*/static_cast(top_k), + /*dispatch_output_per_expert_alignment=*/ + static_cast(dispatch_output_per_expert_alignment), + }; +} + +} // namespace + +int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_alignment) { + return static_cast( + nvte_ep_handle_mem_size(make_layer_cfg(top_k, dispatch_output_per_expert_alignment))); +} + +// ── Per-step ops ───────────────────────────────────────────────────────────── + +void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor token_counts, int64_t top_k, + int64_t dispatch_output_per_expert_alignment) { + auto stream = at::cuda::getCurrentCUDAStream().stream(); + NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]"); + check_topk_idx_int64(topk_idx); + const size_t T_flat = topk_idx.numel() / topk_idx.size(-1); + const size_t topk_n = static_cast(topk_idx.size(-1)); + + auto topk_idx_te = + makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, DType::kInt64); + auto token_counts_te = makeTransformerEngineTensor( + token_counts.data_ptr(), Shape{static_cast(token_counts.numel())}, DType::kInt32); + auto handle_mem_te = makeTransformerEngineTensor( + handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + + nvte_ep_prepare(handle_mem_te.data(), topk_idx_te.data(), token_counts_te.data(), + make_layer_cfg(top_k, dispatch_output_per_expert_alignment), stream); +} + +void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, + at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights) { + auto stream = at::cuda::getCurrentCUDAStream().stream(); + NVTE_CHECK(tokens.dim() >= 2, "tokens must be at least 2D [..., H]"); + NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]"); + NVTE_CHECK(topk_weights.dim() >= 2, "topk_weights must be at least 2D [..., top_k]"); + NVTE_CHECK(recv_tokens.dim() >= 2, "recv_tokens must be at least 2D [..., recv_pr, H]"); + check_topk_idx_int64(topk_idx); + NVTE_CHECK(tokens.is_contiguous(), "tokens must be contiguous"); + NVTE_CHECK(topk_weights.is_contiguous(), "topk_weights must be contiguous"); + NVTE_CHECK(recv_tokens.is_contiguous(), "recv_tokens must be contiguous"); + NVTE_CHECK(recv_topk_weights.is_contiguous(), "recv_topk_weights must be contiguous"); + + const size_t H = static_cast(tokens.size(-1)); + const size_t T_flat = tokens.numel() / H; + const size_t topk_n = static_cast(topk_idx.size(-1)); + const size_t recv_pr = recv_tokens.numel() / H; + + NVTE_CHECK(static_cast(topk_weights.size(-1)) == topk_n, + "topk_weights last dim must equal topk_idx last dim"); + NVTE_CHECK(static_cast(topk_idx.numel()) == T_flat * topk_n, + "topk_idx token count must equal tokens token count"); + NVTE_CHECK(static_cast(topk_weights.numel()) == T_flat * topk_n, + "topk_weights token count must equal tokens token count"); + NVTE_CHECK(static_cast(recv_topk_weights.numel()) == recv_pr, + "recv_topk_weights total size must equal recv_tokens recv_pr"); + NVTE_CHECK(recv_tokens.scalar_type() == tokens.scalar_type(), "recv_tokens dtype (", + c10::toString(recv_tokens.scalar_type()), ") must match tokens dtype (", + c10::toString(tokens.scalar_type()), ")"); + check_symm_mem_required(recv_tokens, "recv_tokens"); + check_symm_mem_required(recv_topk_weights, "recv_topk_weights"); + + auto tok_dtype = GetTransformerEngineDType(tokens.scalar_type()); + auto handle_mem_te = makeTransformerEngineTensor( + handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + auto topk_idx_te = + makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, DType::kInt64); + auto tokens_te = makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype); + auto topk_w_te = + makeTransformerEngineTensor(topk_weights.data_ptr(), Shape{T_flat, topk_n}, DType::kFloat32); + auto recv_tokens_te = + makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype); + auto recv_topk_w_te = + makeTransformerEngineTensor(recv_topk_weights.data_ptr(), Shape{recv_pr}, DType::kFloat32); + + // top_k / alignment are carried by the cached layer_cfg seeded at ep_prepare; + // per-step ops look them up by handle_mem pointer in the backend. + NVTECommWindow tokens_win = maybe_make_window(tokens); + NVTECommWindow topk_w_win = maybe_make_window(topk_weights); + NVTECommWindow recv_tokens_win = maybe_make_window(recv_tokens); + NVTECommWindow recv_topk_w_win = maybe_make_window(recv_topk_weights); + nvte_ep_dispatch(handle_mem_te.data(), topk_idx_te.data(), tokens_te.data(), tokens_win, + topk_w_te.data(), topk_w_win, recv_tokens_te.data(), recv_tokens_win, + recv_topk_w_te.data(), recv_topk_w_win, stream); +} + +void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result) { + auto stream = at::cuda::getCurrentCUDAStream().stream(); + NVTE_CHECK(expert_out.dim() >= 2, "expert_out must be at least 2D [..., recv_pr, H]"); + NVTE_CHECK(result.dim() >= 2, "result must be at least 2D [..., H]"); + NVTE_CHECK(expert_out.is_contiguous(), "expert_out must be contiguous"); + + const size_t H = static_cast(expert_out.size(-1)); + const size_t recv_pr = expert_out.numel() / H; + const size_t T_flat = result.numel() / H; + NVTE_CHECK(static_cast(result.size(-1)) == H, + "result hidden dim must equal expert_out hidden dim"); + NVTE_CHECK(result.scalar_type() == expert_out.scalar_type(), "result dtype (", + c10::toString(result.scalar_type()), ") must match expert_out dtype (", + c10::toString(expert_out.scalar_type()), ")"); + check_symm_mem_required(expert_out, "expert_out"); + + auto eo_dtype = GetTransformerEngineDType(expert_out.scalar_type()); + auto handle_mem_te = makeTransformerEngineTensor( + handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + auto expert_out_te = + makeTransformerEngineTensor(expert_out.data_ptr(), Shape{recv_pr, H}, eo_dtype); + auto result_te = makeTransformerEngineTensor(result.data_ptr(), Shape{T_flat, H}, eo_dtype); + + NVTECommWindow expert_out_win = maybe_make_window(expert_out); + nvte_ep_combine(handle_mem_te.data(), expert_out_te.data(), expert_out_win, result_te.data(), + stream); +} + +void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_topk_weights, + at::Tensor grad_tokens, at::Tensor grad_topk_weights) { + auto stream = at::cuda::getCurrentCUDAStream().stream(); + NVTE_CHECK(grad.dim() >= 2, "grad must be at least 2D [..., recv_pr, H]"); + NVTE_CHECK(grad_tokens.dim() >= 2, "grad_tokens must be at least 2D [..., H]"); + NVTE_CHECK(grad_topk_weights.dim() >= 2, "grad_topk_weights must be at least 2D [..., top_k]"); + NVTE_CHECK(grad.is_contiguous(), "grad must be contiguous"); + NVTE_CHECK(g_recv_topk_weights.is_contiguous(), "g_recv_topk_weights must be contiguous"); + + const size_t H = static_cast(grad.size(-1)); + const size_t recv_pr = grad.numel() / H; + const size_t T_flat = grad_tokens.numel() / H; + const size_t topk_n = static_cast(grad_topk_weights.size(-1)); + NVTE_CHECK(static_cast(g_recv_topk_weights.numel()) == recv_pr, + "g_recv_topk_weights total size must equal grad recv_pr"); + NVTE_CHECK(static_cast(grad_tokens.size(-1)) == H, + "grad_tokens hidden dim must equal grad H"); + NVTE_CHECK(static_cast(grad_topk_weights.numel()) == T_flat * topk_n, + "grad_topk_weights numel (", grad_topk_weights.numel(), + ") must equal T_flat * top_k (", T_flat * topk_n, ")"); + NVTE_CHECK(grad_tokens.scalar_type() == grad.scalar_type(), "grad_tokens dtype (", + c10::toString(grad_tokens.scalar_type()), ") must match grad dtype (", + c10::toString(grad.scalar_type()), ")"); + // Upstream grads are autograd-allocated, so they take the staged-copy path. + + auto g_dtype = GetTransformerEngineDType(grad.scalar_type()); + auto handle_mem_te = makeTransformerEngineTensor( + handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + auto grad_te = makeTransformerEngineTensor(grad.data_ptr(), Shape{recv_pr, H}, g_dtype); + auto g_recv_w_te = + makeTransformerEngineTensor(g_recv_topk_weights.data_ptr(), Shape{recv_pr}, DType::kFloat32); + auto grad_tokens_te = + makeTransformerEngineTensor(grad_tokens.data_ptr(), Shape{T_flat, H}, g_dtype); + auto grad_topk_w_te = makeTransformerEngineTensor(grad_topk_weights.data_ptr(), + Shape{T_flat, topk_n}, DType::kFloat32); + + NVTECommWindow grad_win = maybe_make_window(grad); + NVTECommWindow g_recv_w_win = maybe_make_window(g_recv_topk_weights); + nvte_ep_dispatch_bwd(handle_mem_te.data(), grad_te.data(), grad_win, g_recv_w_te.data(), + g_recv_w_win, grad_tokens_te.data(), grad_topk_w_te.data(), stream); +} + +void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out) { + auto stream = at::cuda::getCurrentCUDAStream().stream(); + NVTE_CHECK(grad.dim() >= 2, "grad must be at least 2D [..., H]"); + NVTE_CHECK(grad_expert_out.dim() >= 2, "grad_expert_out must be at least 2D [..., recv_pr, H]"); + NVTE_CHECK(grad.is_contiguous(), "grad must be contiguous"); + NVTE_CHECK(grad_expert_out.is_contiguous(), "grad_expert_out must be contiguous"); + + const size_t H = static_cast(grad.size(-1)); + const size_t T_flat = grad.numel() / H; + const size_t recv_pr = grad_expert_out.numel() / H; + NVTE_CHECK(static_cast(grad_expert_out.size(-1)) == H, + "grad_expert_out hidden dim must match grad H"); + NVTE_CHECK(grad_expert_out.scalar_type() == grad.scalar_type(), "grad_expert_out dtype (", + c10::toString(grad_expert_out.scalar_type()), ") must match grad dtype (", + c10::toString(grad.scalar_type()), ")"); + // grad is autograd-allocated (staged-copy path); grad_expert_out is the + // EpBuffer-owned scatter target and must be symm-mem in zero-copy mode. + check_symm_mem_required(grad_expert_out, "grad_expert_out"); + + auto g_dtype = GetTransformerEngineDType(grad.scalar_type()); + auto handle_mem_te = makeTransformerEngineTensor( + handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + auto grad_te = makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype); + auto grad_expert_out_te = + makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype); + + // grad is autograd-allocated (staged); grad_expert_out resolves to a symm-mem + // window in zero-copy mode, else kNoWindow for the staged path. + NVTECommWindow grad_win = maybe_make_window(grad); + NVTECommWindow grad_expert_out_win = maybe_make_window(grad_expert_out); + nvte_ep_combine_bwd(handle_mem_te.data(), grad_te.data(), grad_win, grad_expert_out_te.data(), + grad_expert_out_win, stream); +} + +void register_ep_bindings(pybind11::module_& m) { + namespace py = pybind11; + m.def("ep_initialize", &ep_initialize, + "Initialize the EP backend; borrows torch's NCCL comm pointed to by ``comm_ptr``.", + py::arg("comm_ptr"), py::arg("group_name"), py::arg("num_experts"), + py::arg("max_tokens_per_rank"), py::arg("max_recv_tokens_per_rank"), py::arg("hidden_dim"), + py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false, + py::call_guard()); + m.def("ep_finalize", &ep_finalize, "Tear down the EP backend. Idempotent.", + py::call_guard()); + m.def("ep_get_zero_copy", &ep_get_zero_copy, "Return the current EP zero-copy toggle state."); + m.def("ep_handle_mem_size", &ep_handle_mem_size, + "Return the handle_mem byte size for the given layer config.", py::arg("top_k"), + py::arg("dispatch_output_per_expert_alignment") = 0); + m.def("ep_prepare", &ep_prepare, "EP prepare", py::call_guard()); + m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::call_guard()); + m.def("ep_combine", &ep_combine, "EP combine", py::call_guard()); + m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward", + py::call_guard()); + m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", + py::call_guard()); +} + +} // namespace transformer_engine::pytorch + +#endif // NVTE_WITH_NCCL_EP diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index ba92a5f143..14272c9ac4 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -293,6 +293,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "DSquaredReLU + DBias + Quantize", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer")); +#ifdef NVTE_WITH_NCCL_EP + transformer_engine::pytorch::register_ep_bindings(m); +#endif // NVTE_WITH_NCCL_EP + // Permutation functions m.def("moe_permute_fwd", transformer_engine::pytorch::moe_permute_fwd, "MOE permute FWD", py::call_guard()); diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 670eecaa5e..c050f26869 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1842,6 +1842,27 @@ def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_gr return msg +def symm_mem_alloc( + shape, + dtype: torch.dtype, + ep_group: dist_group_type, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Allocate and rendezvous a symm-mem buffer on ep_group. Collective on ep_group.""" + if device is None: + device = torch.device("cuda", torch.cuda.current_device()) + if not HAS_TORCH_SYMMETRIC: + raise RuntimeError( + "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc " + "requires PyTorch built with NCCL symm-mem support." + ) + if symm_mem.get_backend(device) != "NCCL": + symm_mem.set_backend("NCCL") + t = symm_mem.empty(*shape, dtype=dtype, device=device) + symm_mem.rendezvous(t, group=ep_group) + return t + + def symmetric_all_reduce( inp: torch.Tensor, tp_group: Optional[dist_group_type] = None, diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py new file mode 100644 index 0000000000..b41115976a --- /dev/null +++ b/transformer_engine/pytorch/ep.py @@ -0,0 +1,610 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""PyTorch Expert Parallelism (EP) API.""" + +from __future__ import annotations + +import atexit +import warnings +from typing import Optional + +import torch +import torch.distributed as dist + +import transformer_engine_torch as tex + +from .cpu_offload import mark_not_offload +from .distributed import symm_mem_alloc + + +__all__ = [ + "EpBuffer", + "ep_bootstrap", + "ep_finalize", + "ep_dispatch", + "ep_combine", + "symm_mem_alloc", +] + + +# ``symm_mem_alloc`` (imported from .distributed) allocates the symm-mem buffers +# used by the zero-copy IO path. Set ``ep_bootstrap(zero_copy=True)`` to opt in; +# the C++ backend then operates the EP group in zero-copy mode. + + +# Bootstrap + + +# NCCL EP requires NCCL >= 2.30.4 (matches the C++ backend's runtime check). +_MIN_NCCL_VERSION = (2, 30, 4) + + +def _check_nccl_runtime_version() -> None: + """Raise with a clear message if the loaded libnccl is too old for NCCL EP.""" + import ctypes + + try: + lib = ctypes.CDLL("libnccl.so.2", mode=ctypes.RTLD_GLOBAL) + v = ctypes.c_int(0) + if lib.ncclGetVersion(ctypes.byref(v)) != 0: + warnings.warn("ncclGetVersion failed; skipping NCCL EP version check.") + return + except OSError: # libnccl not findable; let the C++ side error + return + n = v.value + # NCCL packs as (major*10000 + minor*100 + patch) up to ~2.x; newer + # builds use the same scheme. Decode defensively. + major, minor, patch = n // 10000, (n // 100) % 100, n % 100 + if (major, minor, patch) < _MIN_NCCL_VERSION: + min_str = ".".join(str(x) for x in _MIN_NCCL_VERSION) + raise RuntimeError( + f"NCCL EP requires NCCL >= {min_str}, found {major}.{minor}.{patch} at runtime. " + "Set LD_LIBRARY_PATH to a newer libnccl.so before launching." + ) + + +_BOOTSTRAPPED = False +_ATEXIT_REGISTERED = False +# EP group captured at bootstrap; EpBuffer uses it to allocate the symm-mem +# combine grad buffer in zero-copy mode. +_EP_GROUP: Optional[dist.ProcessGroup] = None + + +def _atexit_finalize() -> None: + """Best-effort teardown at interpreter shutdown; swallows errors.""" + global _BOOTSTRAPPED, _EP_GROUP + if _BOOTSTRAPPED: + try: + tex.ep_finalize() + except Exception: # pylint: disable=broad-exception-caught + import traceback + + traceback.print_exc() + finally: + _BOOTSTRAPPED = False + _EP_GROUP = None + + +def ep_bootstrap( + ep_group: dist.ProcessGroup, + num_experts: int, + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + max_num_sms: int = 0, + zero_copy: bool = False, + max_token_dtype: torch.dtype = torch.bfloat16, +) -> None: + """Initialize EP by borrowing ep_group's NCCL comm. Call once per process. + + max_token_dtype sets the widest token dtype this EP group will dispatch; + it sizes NCCL EP staging buffers. + + ``zero_copy`` opts the EP group into the symm-mem zero-copy IO path; pass + ``True`` only when payload tensors are allocated via ``symm_mem_alloc``. + Defaults to ``False``. + """ + global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP + if _BOOTSTRAPPED: + raise RuntimeError("ep_bootstrap was already called in this process") + if ep_group.size() < 2: + raise ValueError(f"ep_bootstrap requires ep_group.size() >= 2 (got {ep_group.size()}).") + _check_nccl_runtime_version() + if zero_copy: + warnings.warn( + "ep_bootstrap(zero_copy=True) is experimental; the symm-mem IO path " + "and its alias contracts on EpBuffer slots are subject to change.", + stacklevel=2, + ) + + # Materialize the PG's NCCL comm before borrowing its raw handle. + dist.barrier(group=ep_group, device_ids=[torch.cuda.current_device()]) + comm_ptr = ep_group._get_backend(torch.device("cuda"))._comm_ptr() + + tex.ep_initialize( + int(comm_ptr), + str(ep_group.group_name), + int(num_experts), + int(max_tokens_per_rank), + int(recv_capacity_per_rank), + int(hidden_dim), + int(max_num_sms), + max_token_dtype, + bool(zero_copy), + ) + _BOOTSTRAPPED = True + _EP_GROUP = ep_group + if not _ATEXIT_REGISTERED: + atexit.register(_atexit_finalize) + _ATEXIT_REGISTERED = True + + +def ep_finalize() -> None: + """Optional explicit EP teardown; idempotent. + + An atexit handler covers normal interpreter shutdown, so most users do not + need to call this. Call it explicitly only before + ``dist.destroy_process_group()``, since the borrowed NCCL comm becomes + invalid once the PG is destroyed. + """ + global _BOOTSTRAPPED, _EP_GROUP + if not _BOOTSTRAPPED: + return + try: + tex.ep_finalize() + finally: + _BOOTSTRAPPED = False + _EP_GROUP = None + + +# Buffer + + +class EpBuffer: + """Per-microbatch EP layer state holding handle_mem and token_counts. + Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). + + In zero-copy mode the buffer owns the symm-mem buffers the one-sided path + requires: the dispatch recv outputs (recv_tokens, recv_topk_weights) and the + combine backward grad target. One set per buffer, so each layer/microbatch is + isolated. In normal mode these are None and allocated in-flight instead (recv + outputs in the dispatch forward, the combine grad in the backward). + """ + + __slots__ = ( + "handle_mem", + "top_k", + "alignment", + "max_tokens_per_rank", + "recv_capacity_per_rank", + "hidden_dim", + "num_local_experts", + "payload_dtype", + "device", + "token_counts", + "zero_copy", + "recv_tokens_symm_buf", + "recv_topk_weights_symm_buf", + "grad_expert_out_symm_buf", + ) + + def _alloc_symm_buffers(self) -> None: + """Fill in buffer-owned symm-mem buffers the caller did not supply. + recv_topk_weights is always owned. In normal mode caller-supplied + tensors are kept as-is and the rest stay None (allocated in-flight).""" + if not self.zero_copy: + self.recv_topk_weights_symm_buf = None + return + if _EP_GROUP is None: + raise RuntimeError( + "ep_bootstrap must be called before constructing a zero-copy EpBuffer" + ) + rc, h = self.recv_capacity_per_rank, self.hidden_dim + # Persistent across microbatches; keep resident under CPU offloading. + self.recv_topk_weights_symm_buf = symm_mem_alloc( + (rc,), torch.float32, _EP_GROUP, device=self.device + ) + mark_not_offload(self.recv_topk_weights_symm_buf) + if self.recv_tokens_symm_buf is None: + self.recv_tokens_symm_buf = symm_mem_alloc( + (rc, h), self.payload_dtype, _EP_GROUP, device=self.device + ) + mark_not_offload(self.recv_tokens_symm_buf) + if self.grad_expert_out_symm_buf is None: + self.grad_expert_out_symm_buf = symm_mem_alloc( + (rc, h), self.payload_dtype, _EP_GROUP, device=self.device + ) + mark_not_offload(self.grad_expert_out_symm_buf) + + def __init__( + self, + top_k: int, + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + num_local_experts: int, + alignment: int = 0, + payload_dtype: torch.dtype = torch.bfloat16, + device: Optional[torch.device] = None, + dispatch_recv_tokens: Optional[torch.Tensor] = None, + combine_grad_expert_out: Optional[torch.Tensor] = None, + ) -> None: + """Pass ``dispatch_recv_tokens`` (dispatch recv output) and/or + ``combine_grad_expert_out`` (combine backward grad target) to use caller-owned + buffers; the buffer then skips allocating them. Both must be symm-mem-backed + under zero-copy. Whatever is left None is buffer-owned (zero-copy) or allocated + in-flight (normal mode). recv_topk_weights is always owned by the buffer.""" + if device is None: + device = torch.device("cuda", torch.cuda.current_device()) + alignment = int(alignment) + if alignment > 1 and (alignment & (alignment - 1)) != 0: + raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") + self.top_k = int(top_k) + self.alignment = alignment + self.max_tokens_per_rank = int(max_tokens_per_rank) + self.recv_capacity_per_rank = int(recv_capacity_per_rank) + self.hidden_dim = int(hidden_dim) + self.num_local_experts = int(num_local_experts) + self.payload_dtype = payload_dtype + self.device = device + self.zero_copy = bool(tex.ep_get_zero_copy()) + self.recv_tokens_symm_buf = dispatch_recv_tokens + self.grad_expert_out_symm_buf = combine_grad_expert_out + + size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) + self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) + self.token_counts = torch.empty(self.num_local_experts, dtype=torch.int32, device=device) + # Persistent tensor; keep resident if activation CPU offloading is on. + mark_not_offload(self.handle_mem) + self._alloc_symm_buffers() + + +# torch.library custom ops (so they don't graph-break under torch.compile) + +_LIB = "transformer_engine_ep" + + +@torch.library.custom_op( + f"{_LIB}::prepare", + mutates_args=("handle_mem", "token_counts"), + device_types="cuda", +) +def _prepare_op( + handle_mem: torch.Tensor, + top_k: int, + topk_idx: torch.Tensor, + token_counts: torch.Tensor, + alignment: int, +) -> None: + tex.ep_prepare(handle_mem, topk_idx, token_counts, top_k, alignment) + + +@_prepare_op.register_fake +def _(*_args, **_kw): + return None + + +@torch.library.custom_op( + f"{_LIB}::dispatch", + mutates_args=("recv_tokens", "recv_topk_weights"), + device_types="cuda", +) +def _dispatch_op( + handle_mem: torch.Tensor, + topk_idx: torch.Tensor, + tokens: torch.Tensor, + topk_weights: torch.Tensor, + recv_tokens: torch.Tensor, + recv_topk_weights: torch.Tensor, +) -> None: + tex.ep_dispatch(handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights) + + +@_dispatch_op.register_fake +def _(*_args, **_kw): + return None + + +@torch.library.custom_op( + f"{_LIB}::combine", + mutates_args=("result",), + device_types="cuda", +) +def _combine_op( + handle_mem: torch.Tensor, + expert_out: torch.Tensor, + result: torch.Tensor, +) -> None: + tex.ep_combine(handle_mem, expert_out, result) + + +@_combine_op.register_fake +def _(*_args, **_kw): + return None + + +@torch.library.custom_op( + f"{_LIB}::dispatch_bwd", + mutates_args=("grad_tokens", "grad_topk_weights"), + device_types="cuda", +) +def _dispatch_bwd_op( + handle_mem: torch.Tensor, + grad: torch.Tensor, + g_recv_topk_weights: torch.Tensor, + grad_tokens: torch.Tensor, + grad_topk_weights: torch.Tensor, +) -> None: + tex.ep_dispatch_bwd(handle_mem, grad, g_recv_topk_weights, grad_tokens, grad_topk_weights) + + +@_dispatch_bwd_op.register_fake +def _(*_args, **_kw): + return None + + +@torch.library.custom_op( + f"{_LIB}::combine_bwd", + mutates_args=("grad_expert_out",), + device_types="cuda", +) +def _combine_bwd_op( + handle_mem: torch.Tensor, + grad: torch.Tensor, + grad_expert_out: torch.Tensor, +) -> None: + tex.ep_combine_bwd(handle_mem, grad, grad_expert_out) + + +@_combine_bwd_op.register_fake +def _(*_args, **_kw): + return None + + +# Non-autograd primitives + + +def ep_prepare(buffer: "EpBuffer", topk_idx: torch.Tensor) -> torch.Tensor: + """AllGather the routing map; fills ``buffer.handle_mem`` and returns + ``buffer.token_counts`` (int32, shape [num_local_experts]). topk_idx must + be int64. + """ + torch.ops.transformer_engine_ep.prepare( + buffer.handle_mem, buffer.top_k, topk_idx, buffer.token_counts, buffer.alignment + ) + return buffer.token_counts + + +def _ep_dispatch_raw( + buffer: "EpBuffer", + topk_idx: torch.Tensor, + tokens: torch.Tensor, + topk_weights: torch.Tensor, + recv_tokens: torch.Tensor, + recv_topk_weights: torch.Tensor, +) -> None: + """Raw dispatch; no autograd, no prepare. Caller must run ep_prepare first.""" + tex.ep_dispatch( + buffer.handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights + ) + + +def _ep_combine_raw(buffer: "EpBuffer", expert_out: torch.Tensor, result: torch.Tensor) -> None: + """Raw combine; no autograd. Caller pre-weights expert_out.""" + tex.ep_combine(buffer.handle_mem, expert_out, result) + + +# autograd.Function wrappers + + +class _EpDispatch(torch.autograd.Function): + """Autograd prepare+dispatch; bwd uses user-supplied grad inputs as-is.""" + + @staticmethod + def forward( # type: ignore[override] + ctx, + handle_mem: torch.Tensor, + top_k: int, + alignment: int, + recv_tokens: torch.Tensor, + recv_topk_weights: torch.Tensor, + token_counts: torch.Tensor, + topk_idx: torch.Tensor, + tokens: torch.Tensor, + topk_weights: torch.Tensor, + ): + """Prepare + dispatch fwd.""" + torch.ops.transformer_engine_ep.prepare( + handle_mem, top_k, topk_idx, token_counts, alignment + ) + torch.ops.transformer_engine_ep.dispatch( + handle_mem, + topk_idx, + tokens, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + ctx.save_for_backward(handle_mem) + ctx.tokens_shape = tokens.shape + ctx.tokens_dtype = tokens.dtype + ctx.topk_weights_shape = topk_weights.shape + ctx.tokens_T_flat = tokens.numel() // tokens.shape[-1] + ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1] + ctx.top_k = topk_weights.shape[-1] + ctx.recv_capacity = recv_tokens.shape[0] + ctx.hidden_dim = tokens.shape[-1] + ctx.mark_non_differentiable(token_counts) + # Detach so the long-lived buffers aren't tracked as differentiable outputs; + # autograd re-attaches grad_fn pointing back at this Function. + return recv_tokens.detach(), recv_topk_weights.detach(), token_counts + + @staticmethod + def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: ignore[override] + """Dispatch bwd; normalizes grad-input layout, otherwise passes through.""" + (handle_mem,) = ctx.saved_tensors + device = handle_mem.device + g_recv_tokens = g_recv_tokens.contiguous() + g_recv_topk_weights = g_recv_topk_weights.contiguous() + grad_tokens = torch.empty( + ctx.tokens_T_flat, ctx.hidden_dim, dtype=ctx.tokens_dtype, device=device + ) + grad_topk_weights = torch.empty( + ctx.topk_T_flat, ctx.top_k, dtype=torch.float32, device=device + ) + torch.ops.transformer_engine_ep.dispatch_bwd( + handle_mem, + g_recv_tokens, + g_recv_topk_weights, + grad_tokens, + grad_topk_weights, + ) + return ( + None, # handle_mem + None, # top_k + None, # alignment + None, # recv_tokens + None, # recv_topk_weights + None, # token_counts + None, # topk_idx + grad_tokens.view(ctx.tokens_shape), + grad_topk_weights.view(ctx.topk_weights_shape), + ) + + +class _EpCombine(torch.autograd.Function): + """Autograd combine. + + bwd scatters the expert_out grad into ``grad_symm_buf`` (EpBuffer-owned + symm-mem, one-sided) in zero-copy mode, or into a plain tensor allocated + in-flight here otherwise. The latter keeps allocation torch.compile / + CUDA-graph safe and lets autograd own the grad's lifetime. + + ``grad_symm_buf`` is the backward's scatter target (an output it writes, never + reads), so it is stashed as a plain ctx attribute rather than via + save_for_backward, which would version-track a tensor we mutate. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + handle_mem: torch.Tensor, + num_local_tokens: int, + hidden_dim: int, + grad_symm_buf: Optional[torch.Tensor], + expert_out: torch.Tensor, + ): + """Combine fwd; stashes the bwd grad target or expert_out shape to size it.""" + device = expert_out.device + result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) + torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) + ctx.save_for_backward(handle_mem) + ctx.grad_symm_buf = grad_symm_buf + if grad_symm_buf is None: + ctx.expert_out_shape = expert_out.shape + ctx.expert_out_dtype = expert_out.dtype + ctx.device = device + return result + + @staticmethod + def backward(ctx, g_result): # type: ignore[override] + """Combine bwd; scatters the result grad into the grad target.""" + if not g_result.is_contiguous(): + g_result = g_result.contiguous() + (handle_mem,) = ctx.saved_tensors + grad_expert_out = ctx.grad_symm_buf + if grad_expert_out is None: + grad_expert_out = torch.empty( + ctx.expert_out_shape, dtype=ctx.expert_out_dtype, device=ctx.device + ) + torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) + return ( + None, # handle_mem + None, # num_local_tokens + None, # hidden_dim + None, # grad_symm_buf + grad_expert_out, + ) + + +# Public high-level wrappers + + +# NCCL EP currently only supports bfloat16 payload tensors. +def _require_bf16(name: str, t: torch.Tensor) -> None: + if t.dtype is not torch.bfloat16: + raise NotImplementedError( + f"NCCL EP currently supports only bfloat16 payloads; got {name}.dtype={t.dtype}." + ) + + +def ep_dispatch( + buffer: EpBuffer, + tokens: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, +): + """Prepare + dispatch with autograd. topk_idx must be int64. + + recv_tokens comes from the EpBuffer (caller-supplied or buffer-owned under + zero-copy) or is allocated in-flight (normal mode). recv_topk_weights is always + owned by the buffer. Returns (recv_tokens, recv_topk_weights, token_counts); + token_counts is non-diff. + """ + _require_bf16("tokens", tokens) + if topk_weights.dtype is not torch.float32: + raise TypeError( + f"topk_weights must be float32; got dtype={topk_weights.dtype}. " + "Cast with topk_weights.float() before calling." + ) + recv_tokens = buffer.recv_tokens_symm_buf + if recv_tokens is None: + recv_tokens = torch.empty( + buffer.recv_capacity_per_rank, + buffer.hidden_dim, + dtype=buffer.payload_dtype, + device=buffer.device, + ) + recv_topk_weights = ( + buffer.recv_topk_weights_symm_buf + if buffer.zero_copy + else torch.empty(buffer.recv_capacity_per_rank, dtype=torch.float32, device=buffer.device) + ) + return _EpDispatch.apply( + buffer.handle_mem, + buffer.top_k, + buffer.alignment, + recv_tokens, + recv_topk_weights, + buffer.token_counts, + topk_idx, + tokens, + topk_weights, + ) + + +def ep_combine( + buffer: EpBuffer, + expert_out: torch.Tensor, + *, + num_local_tokens: Optional[int] = None, +): + """Combine with autograd; caller pre-applies topk weighting. + + The backward scatters the expert_out grad into the EpBuffer grad target + (caller-supplied or buffer-owned under zero-copy), or a tensor allocated + in-flight (normal mode). Result shape is (num_local_tokens, buffer.hidden_dim); + defaults to buffer.max_tokens_per_rank rows. + """ + _require_bf16("expert_out", expert_out) + if num_local_tokens is None: + num_local_tokens = buffer.max_tokens_per_rank + grad_expert_out = buffer.grad_expert_out_symm_buf + return _EpCombine.apply( + buffer.handle_mem, + num_local_tokens, + buffer.hidden_dim, + grad_expert_out, + expert_out, + ) From 90baf028905d23e133ecb5920a6a68c82f1ceeb7 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 30 Jun 2026 09:37:14 +0200 Subject: [PATCH 16/42] [Common] Update NCCL submodule to have the fix for MAX_SUPPORTED_TOKENS_PER_RANK (#3150) * nccl with relax num_dispatch_tokens%64!=0 Signed-off-by: Phuong Nguyen * Skip EP tests/examples on nodes without NVLink Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- 3rdparty/nccl | 2 +- examples/jax/ep/bench/run_ep_bench.sh | 7 +++++++ tests/cpp_distributed/run_test_ep.sh | 6 ++++++ tests/jax/multi_process_launch_ep.sh | 7 +++++++ tests/pytorch/distributed/run_test_ep.sh | 4 ++-- 5 files changed, 23 insertions(+), 3 deletions(-) diff --git a/3rdparty/nccl b/3rdparty/nccl index 808d2433dd..a6b5de08b6 160000 --- a/3rdparty/nccl +++ b/3rdparty/nccl @@ -1 +1 @@ -Subproject commit 808d2433dda3cccc80f8172a94a6b117359e7102 +Subproject commit a6b5de08b6af4f938cef541ae6e4d405632f89a4 diff --git a/examples/jax/ep/bench/run_ep_bench.sh b/examples/jax/ep/bench/run_ep_bench.sh index 1531dfd5cf..63133156eb 100755 --- a/examples/jax/ep/bench/run_ep_bench.sh +++ b/examples/jax/ep/bench/run_ep_bench.sh @@ -47,6 +47,13 @@ NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) if [ "${NUM_GPUS}" -lt 4 ]; then echo "EP bench requires >=4 GPUs (found ${NUM_GPUS}); SKIPPING."; exit 0 fi + +# NCCL EP requires active NVLink P2P among ranks on the node. +if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then + echo "NVLink not detected on this platform — EP bench requires NVLink; SKIPPING." + exit 0 +fi + NUM=4 COORD="${COORD:-127.0.0.1:23457}" TIMEOUT_S="${TIMEOUT_S:-1800}" diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh index d486d45f8a..da293dadfd 100755 --- a/tests/cpp_distributed/run_test_ep.sh +++ b/tests/cpp_distributed/run_test_ep.sh @@ -35,6 +35,12 @@ if (( MIN_SM > 0 && MIN_SM < 90 )); then exit 0 fi +# NCCL EP requires active NVLink P2P among ranks on the node. +if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then + echo "NVLink not detected on this platform; SKIPPING." + exit 0 +fi + TEST_BIN="${BUILD_DIR}/test_ep" if [[ ! -x "${TEST_BIN}" ]]; then echo "ERROR: binary not found: ${TEST_BIN}" diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh index d32ce5f5d3..ff89f712eb 100755 --- a/tests/jax/multi_process_launch_ep.sh +++ b/tests/jax/multi_process_launch_ep.sh @@ -32,6 +32,13 @@ if [ "${NUM_RUNS}" -lt 4 ]; then echo "NCCL EP requires at least 4 GPUs (found ${NUM_RUNS}); SKIPPING." exit 0 fi + +# NCCL EP requires active NVLink P2P among ranks on the node. +if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then + echo "NVLink not detected on this platform — EP test requires NVLink; SKIPPING." + exit 0 +fi + # Default test mesh is (2, 2); use exactly 4 ranks even on larger boxes. NUM_RUNS="${NVTE_TEST_EP_NUM_RANKS:-4}" diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index ae40c8ba4b..68b691f787 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -18,10 +18,10 @@ if [ "${DETECTED_GPUS}" -lt 4 ]; then exit 0 fi -# NCCL EP requires NVLink/NVSwitch between GPUs. +# NCCL EP requires active NVLink P2P among ranks on the node. # On PCIe-only nodes (no NVLink) it falls back to the network # transport and deadlocks, so skip cleanly there. -if ! nvidia-smi topo -m 2>/dev/null | grep -qE "\bNV[0-9]+\b"; then +if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then echo "No NVLink between GPUs (PCIe-only fabric); NCCL EP is unsupported here. SKIPPING." exit 0 fi From 46bdc85f8523dfb1b975095bec40ec1c9bd2e0df Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Tue, 30 Jun 2026 15:17:39 +0200 Subject: [PATCH 17/42] [PyTorch] Preserve fprop operands for dequantized backward override (#3141) * Preserve fprop operands for dequantized backward override Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add test_grouped_linear_backward_override_high_precision_forces_save_original_input test Signed-off-by: root --------- Signed-off-by: Evgeny Signed-off-by: root Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root --- tests/pytorch/test_backward_override.py | 212 ++++++++++++++++++ .../pytorch/module/grouped_linear.py | 2 + transformer_engine/pytorch/module/linear.py | 2 + 3 files changed, 216 insertions(+) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index 5e6f36e8b4..c0acf2e6b3 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -858,6 +858,218 @@ def test_backward_override_recipe_matches_requested_mode( assert quant_recipe.backward_override is None +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) +def test_linear_backward_override_dequantized_ignores_save_original_input( + recipe_name: str, + use_bias: bool, +) -> None: + reset_rng_states() + dtype = torch.bfloat16 + input_shape = (32, 128) + out_features = 128 + _maybe_skip_recipe_dtype(recipe_name, dtype, "linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "linear") + _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, "linear") + + mode_recipe = make_recipe(recipe_name, backward_override="dequantized") + skip_unsupported_backward_override("linear", mode_recipe, "dequantized") + + module_ref = te.Linear( + input_shape[-1], + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + save_original_input=False, + ) + module_test = te.Linear( + input_shape[-1], + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + save_original_input=True, + ) + _copy_named_parameters(module_ref, module_test) + + x = torch.randn(*input_shape, dtype=dtype, device="cuda") + dy = torch.randn(input_shape[0], out_features, dtype=dtype, device="cuda") + + y_ref, dx_ref, dw_ref, db_ref = _run_single_step(module_ref, x, dy, mode_recipe) + y_test, x_test, saved_operands = _run_single_step_with_saved_operands( + module_test, x, mode_recipe + ) + _assert_saved_quantized_operand_uses_rowwise_only(saved_operands[0], name="linear_input") + + y_test_detached = y_test.detach().clone() + y_test.backward(dy) + assert x_test.grad is not None + assert module_test.weight.grad is not None + dx_test = x_test.grad.detach().clone() + dw_test = module_test.weight.grad.detach().clone() + test_bias = getattr(module_test, "bias", None) + db_test = ( + None if test_bias is None or test_bias.grad is None else test_bias.grad.detach().clone() + ) + + assert_close(y_test_detached, y_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dx_test, dx_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dw_test, dw_ref, rtol=0, atol=0, check_dtype=True) + if use_bias: + assert db_test is not None and db_ref is not None + assert_close(db_test, db_ref, rtol=0, atol=0, check_dtype=True) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) +def test_grouped_linear_backward_override_dequantized_ignores_save_original_input( + recipe_name: str, + use_bias: bool, +) -> None: + reset_rng_states() + dtype = torch.bfloat16 + in_features = 128 + out_features = 128 + m_splits = [64, 64] + num_gemms = len(m_splits) + num_tokens = sum(m_splits) + _maybe_skip_recipe_dtype(recipe_name, dtype, "grouped_linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "grouped_linear") + _maybe_skip_unsupported_grouped_splits(recipe_name, m_splits) + + mode_recipe = make_recipe(recipe_name, backward_override="dequantized") + skip_unsupported_backward_override("grouped_linear", mode_recipe, "dequantized") + + module_ref = te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + save_original_input=False, + ) + module_test = te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + save_original_input=True, + ) + _copy_named_parameters(module_ref, module_test) + + x = torch.randn(num_tokens, in_features, dtype=dtype, device="cuda") + dy = torch.randn(num_tokens, out_features, dtype=dtype, device="cuda") + + y_ref, dx_ref, dw_ref, db_ref = _run_grouped_linear_single_step( + module_ref, x, m_splits, dy, mode_recipe + ) + y_test, x_test, saved_operands = _run_grouped_linear_step_with_saved_operands( + module_test, x, m_splits, mode_recipe + ) + saved_inputs = saved_operands[:num_gemms] + for i, saved_input in enumerate(saved_inputs): + _assert_saved_quantized_operand_uses_rowwise_only( + saved_input, name=f"grouped_linear_input{i}" + ) + + y_test_detached = y_test.detach().clone() + y_test.backward(dy) + assert x_test.grad is not None + dx_test = x_test.grad.detach().clone() + dw_test = [getattr(module_test, f"weight{i}").grad.detach().clone() for i in range(num_gemms)] + db_test: list[Optional[torch.Tensor]] = [] + for i in range(num_gemms): + if use_bias: + db_test.append(getattr(module_test, f"bias{i}").grad.detach().clone()) + else: + db_test.append(None) + + assert_close(y_test_detached, y_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dx_test, dx_ref, rtol=0, atol=0, check_dtype=True) + for test_dw, ref_dw in zip(dw_test, dw_ref): + assert_close(test_dw, ref_dw, rtol=0, atol=0, check_dtype=True) + if use_bias: + for test_db, ref_db in zip(db_test, db_ref): + assert test_db is not None and ref_db is not None + assert_close(test_db, ref_db, rtol=0, atol=0, check_dtype=True) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +def test_linear_backward_override_high_precision_forces_save_original_input( + recipe_name: str, +) -> None: + reset_rng_states() + dtype = torch.bfloat16 + input_shape = (32, 128) + _maybe_skip_recipe_dtype(recipe_name, dtype, "linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "linear") + _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, "linear") + + mode_recipe = make_recipe(recipe_name, backward_override="high_precision") + skip_unsupported_backward_override("linear", mode_recipe, "high_precision") + + module = te.Linear( + input_shape[-1], + 128, + bias=False, + params_dtype=dtype, + device="cuda", + save_original_input=False, + ) + x = torch.randn(*input_shape, dtype=dtype, device="cuda") + + _, _, saved_operands = _run_single_step_with_saved_operands(module, x, mode_recipe) + + assert isinstance(saved_operands[0], torch.Tensor) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +def test_grouped_linear_backward_override_high_precision_forces_save_original_input( + recipe_name: str, +) -> None: + reset_rng_states() + dtype = torch.bfloat16 + in_features = 128 + out_features = 128 + m_splits = [64, 64] + num_gemms = len(m_splits) + num_tokens = sum(m_splits) + _maybe_skip_recipe_dtype(recipe_name, dtype, "grouped_linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "grouped_linear") + _maybe_skip_unsupported_grouped_splits(recipe_name, m_splits) + + mode_recipe = make_recipe(recipe_name, backward_override="high_precision") + skip_unsupported_backward_override("grouped_linear", mode_recipe, "high_precision") + + module = te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=dtype, + device="cuda", + save_original_input=False, + ) + x = torch.randn(num_tokens, in_features, dtype=dtype, device="cuda") + + _, _, saved_operands = _run_grouped_linear_step_with_saved_operands( + module, x, m_splits, mode_recipe + ) + + saved_inputs = saved_operands[:num_gemms] + assert isinstance(saved_inputs[0], torch.Tensor) + assert saved_inputs[0].shape == x.shape + assert all(saved_input is None for saved_input in saved_inputs[1:]) + + saved_weights = saved_operands[2 * num_gemms : 3 * num_gemms] + for saved_weight in saved_weights: + assert isinstance(saved_weight, torch.Tensor) + + @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) @pytest.mark.parametrize("module_type", ("linear", "layernorm_linear", "ops_linear")) @pytest.mark.parametrize("input_shape,out_features", _shape_test_cases) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 8d56e423c4..f2fa8b657e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -431,6 +431,8 @@ def forward( backward_override = None if backward_override == "high_precision": save_original_input = True + elif backward_override == "dequantized": + save_original_input = False num_gemms = len(m_splits) weights = weights_and_biases[:num_gemms] diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index fed367bce5..78a4d31852 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -289,6 +289,8 @@ def _linear_forward_impl( is_fsdp2 = args.is_fsdp2 if backward_override == "high_precision": save_original_input = True + elif backward_override == "dequantized": + save_original_input = False # NVTX label for profiling nvtx_label = "transformer_engine._Linear.forward" From 353206de37e4f3e91354b82c8a199e42ea1498c2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:19:12 +0200 Subject: [PATCH 18/42] [PyTorch] Make quantized-tensor __repr__ safe (#3146) * Make quantized-tensor __repr__ fake-safe under torch.compile Under torch.compile, TE quantized-tensor __repr__ methods are invoked on FakeTensors during AOT autograd's structured logging. The repr bodies call self._scale_inv.item() and/or self.dequantize() (which dispatches to the raw C++ op tex.dequantize), both of which access a FakeTensor's data pointer and raise: RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor) ... This was the sole cause of six fp8 failures in tests/pytorch/test_torch_compile.py. Fix: add one shared helper, safe_quantized_repr, in tensor/_quantization_helpers.py (a safe leaf module importing only torch) that builds a metadata-only repr string. Each data-touching __repr__ now wraps its existing body in a try/except and falls back to the helper when the data cannot be materialized. The eager (non-fake) repr output is unchanged; only a fallback path is added. Wrapped reprs: Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, NVFP4Tensor and their *Storage counterparts. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make quantized __repr__ fallback universal, drop FakeTensor-specific logic Remove the FakeTensor-specific heuristic (_is_fake_data_access_error) and the warning path from safe_quantized_repr. The fallback is now a plain metadata-only repr triggered by any exception while materializing data, with each attribute access individually guarded so __repr__ never raises. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Co-authored-by: Claude Opus 4.8 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../pytorch/tensor/_quantization_helpers.py | 39 +++++++++++++++++++ .../pytorch/tensor/float8_blockwise_tensor.py | 20 +++++++--- .../pytorch/tensor/float8_tensor.py | 19 +++++---- .../pytorch/tensor/mxfp8_tensor.py | 7 +++- .../pytorch/tensor/nvfp4_tensor.py | 7 +++- .../float8_blockwise_tensor_storage.py | 31 +++++++++------ .../tensor/storage/float8_tensor_storage.py | 18 +++++---- .../tensor/storage/mxfp8_tensor_storage.py | 22 ++++++----- .../tensor/storage/nvfp4_tensor_storage.py | 24 +++++++----- 9 files changed, 132 insertions(+), 55 deletions(-) diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 56cf503630..1b08039dda 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -83,3 +83,42 @@ def _stride_from_shape(shape: list[int]): for d in reversed(shape[1:]): rstride.append(rstride[-1] * d) return list(reversed(rstride)) + + +def safe_quantized_repr(obj, cls_name, extras=None, error=None): + """Metadata-only repr fallback for quantized tensors whose data cannot be + materialized for any reason. + + Each attribute access is guarded so that ``__repr__`` never raises. + + Parameters + ---------- + extras : dict, optional + Additional plain-Python (non-tensor) attributes to include, e.g. + ``{"is_2D_scaled": self._is_2D_scaled}``. Values are inserted after + ``fp8_dtype`` and before ``shape``. + error : BaseException, optional + The exception that triggered the fallback. When given, its type and + message are included in the ``data=`` field so that it is visible *why* + the data could not be materialized. + """ + parts = [] + fp8_dtype = getattr(obj, "_fp8_dtype", None) + if fp8_dtype is not None: + parts.append(f"fp8_dtype={fp8_dtype}") + if extras: + for key, value in extras.items(): + parts.append(f"{key}={value}") + try: + parts.append(f"shape={tuple(obj.shape)}") + except Exception: # pylint: disable=broad-except + pass + try: + parts.append(f"dtype={obj.dtype}") + except Exception: # pylint: disable=broad-except + pass + if error is not None: + parts.append(f"data=") + else: + parts.append("data=") + return f"{cls_name}({', '.join(parts)})" diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ba46508d74..d2d28aecfb 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -14,7 +14,7 @@ from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..quantized_tensor import QuantizedTensor, Quantizer -from ._quantization_helpers import _IdentityFunc +from ._quantization_helpers import _IdentityFunc, safe_quantized_repr from ..constants import DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -267,11 +267,19 @@ def __new__( return instance def __repr__(self, *, tensor_contents=None): - return ( - f"Float8BlockwiseQTensor(fp8_dtype={self._fp8_dtype}," - f" is_2D_scaled={self._is_2D_scaled}," - f" data={self.dequantize()})" - ) + try: + return ( + f"Float8BlockwiseQTensor(fp8_dtype={self._fp8_dtype}," + f" is_2D_scaled={self._is_2D_scaled}," + f" data={self.dequantize()})" + ) + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr( + self, + "Float8BlockwiseQTensor", + extras={"is_2D_scaled": self._is_2D_scaled}, + error=exc, + ) def quantize_( self, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e26abf7df0..e90e35c40c 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -18,7 +18,7 @@ from ..utils import canonicalize_process_group, devices_match from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer -from ._quantization_helpers import _IdentityFunc +from ._quantization_helpers import _IdentityFunc, safe_quantized_repr from ..constants import dist_group_type, DType aten = torch.ops.aten @@ -423,13 +423,16 @@ class Float8Tensor(Float8TensorStorage, QuantizedTensor): amax_reduction_group: Optional[dist_group_type] = None def __repr__(self, *, tensor_contents=None): - return ( - "Float8Tensor(" - f"fp8_dtype={self._fp8_dtype}, " - f"scale_inv={self._scale_inv.item()}, " - f"data={self.dequantize()}" - ")" - ) + try: + return ( + "Float8Tensor(" + f"fp8_dtype={self._fp8_dtype}, " + f"scale_inv={self._scale_inv.item()}, " + f"data={self.dequantize()}" + ")" + ) + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr(self, "Float8Tensor", error=exc) def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index d759aaf5c4..33db63d059 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -18,7 +18,7 @@ from ..utils import devices_match, round_up_to_nearest_multiple from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func from ..quantized_tensor import QuantizedTensor, Quantizer -from ._quantization_helpers import _IdentityFunc +from ._quantization_helpers import _IdentityFunc, safe_quantized_repr aten = torch.ops.aten @@ -233,7 +233,10 @@ def __new__( ) def __repr__(self, *, tensor_contents=None): - return f"MXFP8Tensor(fp8_dtype={self._fp8_dtype}, data={self.dequantize()})" + try: + return f"MXFP8Tensor(fp8_dtype={self._fp8_dtype}, data={self.dequantize()})" + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr(self, "MXFP8Tensor", error=exc) def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index aa92be004f..f59af5637b 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -23,7 +23,7 @@ from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func from ..quantized_tensor import QuantizedTensor, Quantizer -from ._quantization_helpers import _IdentityFunc +from ._quantization_helpers import _IdentityFunc, safe_quantized_repr aten = torch.ops.aten @@ -409,7 +409,10 @@ def __new__( return instance def __repr__(self, *, tensor_contents=None): - return f"NVFP4Tensor, data={self.dequantize()})" + try: + return f"NVFP4Tensor, data={self.dequantize()})" + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr(self, "NVFP4Tensor", error=exc) def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index f7a3dae70b..993ead42ee 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -12,6 +12,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType_To_Torch, DType @@ -354,17 +355,25 @@ def _transpose_columnwise_data(self): del _old_data def __repr__(self): - if self._rowwise_data is not None: - data = self.dequantize() - descriptor = "rowwise" - else: - data = self.dequantize() - descriptor = "columnwise" - return ( - "Float8BlockwiseQTensorStorage(" - f"fp8_dtype={self._fp8_dtype}, " - f"{descriptor}_scaled_data={data})" - ) + try: + if self._rowwise_data is not None: + data = self.dequantize() + descriptor = "rowwise" + else: + data = self.dequantize() + descriptor = "columnwise" + return ( + "Float8BlockwiseQTensorStorage(" + f"fp8_dtype={self._fp8_dtype}, " + f"{descriptor}_scaled_data={data})" + ) + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr( + self, + "Float8BlockwiseQTensorStorage", + extras={"is_2D_scaled": self._is_2D_scaled}, + error=exc, + ) def update_usage( self, rowwise_usage: Optional[bool] = None, columnwise_usage: Optional[bool] = None diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index a97162f91c..374d0e1e72 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -12,6 +12,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -209,13 +210,16 @@ def view(self, shape: torch.Size): ) def __repr__(self): - return ( - "Float8TensorStorage(" - f"fp8_dtype={self._fp8_dtype}, " - f"scale_inv={self._scale_inv.item()}, " - f"data={self.dequantize()}" - ")" - ) + try: + return ( + "Float8TensorStorage(" + f"fp8_dtype={self._fp8_dtype}, " + f"scale_inv={self._scale_inv.item()}, " + f"data={self.dequantize()}" + ")" + ) + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr(self, "Float8TensorStorage", error=exc) def _create_transpose(self): """Update FP8 transpose cache""" diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index ea592cd989..606ac9e74b 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -13,6 +13,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType @@ -257,15 +258,18 @@ def view(self, shape: torch.Size): ) def __repr__(self): - data_rowwise = self.dequantize() - - return ( - "MXFP8TensorStorage(" - f"fp8_dtype={self._fp8_dtype}, " - f"rowwise_scaled_data={data_rowwise}" - f"rowwise_scale_inv={self._rowwise_scale_inv}, " - ")" - ) + try: + data_rowwise = self.dequantize() + + return ( + "MXFP8TensorStorage(" + f"fp8_dtype={self._fp8_dtype}, " + f"rowwise_scaled_data={data_rowwise}" + f"rowwise_scale_inv={self._rowwise_scale_inv}, " + ")" + ) + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr(self, "MXFP8TensorStorage", error=exc) def update_usage( self, diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 53bb5e7c11..09f040ba67 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -16,6 +16,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType from ...utils import _empty_tensor @@ -340,16 +341,19 @@ def view(self, shape: torch.Size): ) def __repr__(self): - data_rowwise = self.dequantize() - - return ( - "NVFP4TensorStorage(" - f"rowwise_scaled_data={data_rowwise}," - f"rowwise_scale_inv={self._rowwise_scale_inv}," - f"amax_rowwise={self._amax_rowwise}," - f"amax_columnwise={self._amax_columnwise}," - ")" - ) + try: + data_rowwise = self.dequantize() + + return ( + "NVFP4TensorStorage(" + f"rowwise_scaled_data={data_rowwise}," + f"rowwise_scale_inv={self._rowwise_scale_inv}," + f"amax_rowwise={self._amax_rowwise}," + f"amax_columnwise={self._amax_columnwise}," + ")" + ) + except Exception as exc: # pylint: disable=broad-except + return safe_quantized_repr(self, "NVFP4TensorStorage", error=exc) def update_usage( self, From 3df5e19b1f24325b3bfd34bca10717aaba86b960 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 30 Jun 2026 19:45:40 +0200 Subject: [PATCH 19/42] [Common] EP C API: version config structs and extend `nvte_ep_prepare` with `total_recv_tokens_per_rank` placeholder (#3154) * versioning EP C configs Signed-off-by: Phuong Nguyen * Rename EP prepare token_counts to recv_tokens_per_expert Signed-off-by: Phuong Nguyen * Add total_recv_tokens_per_rank placeholder to nvte_ep_prepare Signed-off-by: Phuong Nguyen * Adapt PyTorch EP binding to versioned nvte_ep C config API Signed-off-by: Phuong Nguyen * Rename EP group config max_num_sms to num_comm_sms Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep.cu | 44 +++++++----- tests/cpp_distributed/test_ep_common.h | 8 +-- transformer_engine/common/ep/ep_api.cpp | 69 ++++++++++++++----- transformer_engine/common/ep/ep_backend.cpp | 27 +++++--- transformer_engine/common/ep/ep_backend.h | 5 +- .../common/include/transformer_engine/ep.h | 55 ++++++++++----- transformer_engine/jax/csrc/extensions/ep.cpp | 36 ++++++---- .../pytorch/csrc/extensions/ep.cpp | 33 +++++---- 8 files changed, 177 insertions(+), 100 deletions(-) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index c7fee7720c..7dbbcdce9d 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -60,7 +60,7 @@ static std::vector generate_tokens(int rank, int num_tokens, int hidden_dim) return v; } -static std::vector expected_token_counts( +static std::vector expected_recv_tokens_per_expert( int recv_rank, int num_processes, int num_tokens, int top_k, int num_experts, int num_local_experts) { int base = recv_rank * num_local_experts; @@ -128,7 +128,7 @@ struct EPBuffers { DevBuf topk_idx; DevBuf topk_weights; DevBuf tokens; - DevBuf token_counts; + DevBuf recv_tokens_per_expert; DevBuf handle_mem; DevBuf recv_tokens; DevBuf recv_topk_weights; @@ -144,22 +144,26 @@ struct EPBuffers { size_t recv_capacity = 0; int top_k_ = 0; size_t alignment_ = 0; + NVTEEpLayerConfig layer_cfg_{}; void alloc(int num_tokens, int top_k, int hidden_dim, int num_local_experts, int ep_size, int max_tokens_per_rank, size_t alignment = 0) { top_k_ = top_k; alignment_ = alignment; + layer_cfg_ = NVTE_EP_LAYER_CONFIG_INIT; + layer_cfg_.top_k = top_k; + layer_cfg_.dispatch_output_per_expert_alignment = alignment; recv_capacity = static_cast(ep_size) * max_tokens_per_rank * 2; topk_idx.alloc(num_tokens * top_k); topk_weights.alloc(num_tokens * top_k); tokens.alloc(num_tokens * hidden_dim); - token_counts.alloc(num_local_experts); + recv_tokens_per_expert.alloc(num_local_experts); recv_tokens.alloc(recv_capacity * hidden_dim); recv_topk_weights.alloc(recv_capacity); result.alloc(num_tokens * hidden_dim); - handle_mem_size = nvte_ep_handle_mem_size(NVTEEpLayerConfig{top_k, alignment}); + handle_mem_size = nvte_ep_handle_mem_size(&layer_cfg_); handle_mem.alloc(handle_mem_size); grad_result.alloc(num_tokens * hidden_dim); @@ -174,25 +178,29 @@ struct EPBuffers { // expects. template struct EPTensors { - TensorWrapper topk_idx, topk_weights, token_counts, handle_mem, tokens; + TensorWrapper topk_idx, topk_weights, recv_tokens_per_expert, handle_mem, tokens; TensorWrapper recv_tokens, recv_topk_weights, result; TensorWrapper grad_result, grad_expert, grad_tokens; TensorWrapper g_recv_topk_weights, grad_topk_weights; int top_k_ = 0; size_t alignment_ = 0; + NVTEEpLayerConfig layer_cfg_{}; EPTensors(EPBuffers& b, int num_tokens, int top_k, int hidden_dim, int num_local_experts) { top_k_ = top_k; alignment_ = b.alignment_; + layer_cfg_ = NVTE_EP_LAYER_CONFIG_INIT; + layer_cfg_.top_k = top_k; + layer_cfg_.dispatch_output_per_expert_alignment = b.alignment_; constexpr DType kTokDType = test::TypeInfo::dtype; using Shape = std::vector; topk_idx = TensorWrapper(b.topk_idx.get(), Shape{(size_t)num_tokens, (size_t)top_k}, DType::kInt64); topk_weights = TensorWrapper(b.topk_weights.get(), Shape{(size_t)num_tokens, (size_t)top_k}, DType::kFloat32); - token_counts = TensorWrapper(b.token_counts.get(), + recv_tokens_per_expert = TensorWrapper(b.recv_tokens_per_expert.get(), Shape{(size_t)num_local_experts}, DType::kInt32); handle_mem = TensorWrapper(b.handle_mem.get(), Shape{b.handle_mem_size}, DType::kByte); @@ -259,7 +267,7 @@ class EpOpTestBase : public ::testing::Test { template int read_total_recv(const EPBuffers& buf) const { std::vector cnt(num_local_experts_); - NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(), + NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.recv_tokens_per_expert.get(), num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); int total = 0; for (int c : cnt) total += c; @@ -300,7 +308,7 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, @@ -309,9 +317,9 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) { // 1. Per-expert counts. std::vector got_counts(num_local_experts_); - NVTE_CHECK_CUDA(cudaMemcpy(got_counts.data(), buf.token_counts.get(), + NVTE_CHECK_CUDA(cudaMemcpy(got_counts.data(), buf.recv_tokens_per_expert.get(), num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); - auto exp_counts = expected_token_counts(g_process_id, g_num_processes, num_tokens_, top_k_, + auto exp_counts = expected_recv_tokens_per_expert(g_process_id, g_num_processes, num_tokens_, top_k_, num_experts_, num_local_experts_); int total_recv = 0; for (int i = 0; i < num_local_experts_; ++i) { @@ -379,7 +387,7 @@ TYPED_TEST(EPCombineTest, Combine) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, @@ -426,7 +434,7 @@ TYPED_TEST(EPCombineBwdTest, CombineBwdCheck) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, @@ -447,7 +455,7 @@ TYPED_TEST(EPCombineBwdTest, CombineBwdCheck) { int total_recv = this->template read_total_recv(buf); std::vector cnt(num_local_experts_); - NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(), + NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.recv_tokens_per_expert.get(), num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); std::vector h_ge(buf.recv_capacity * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(h_ge.data(), buf.grad_expert.get(), @@ -495,7 +503,7 @@ TYPED_TEST(EPDispatchBwdTest, DispatchBwdCheck) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, @@ -563,7 +571,7 @@ TYPED_TEST(EPDispatchBwdGradWeightsTest, RoundTrip) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.recv_topk_weights.get(), 0, buf.recv_topk_weights.bytes(), stream)); ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), @@ -634,7 +642,7 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, @@ -759,7 +767,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - ASSERT_NO_THROW(nvte_ep_prepare(ref_t.handle_mem.data(), ref_t.topk_idx.data(), ref_t.token_counts.data(), NVTEEpLayerConfig{ref_t.top_k_, ref_t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(ref_t.handle_mem.data(), ref_t.topk_idx.data(), ref_t.recv_tokens_per_expert.data(), nullptr, &ref_t.layer_cfg_, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(ref_t.handle_mem.data(), ref_t.topk_idx.data(), ref_t.tokens.data(), NVTECommWindow{}, ref_t.topk_weights.data(), NVTECommWindow{}, ref_t.recv_tokens.data(), NVTECommWindow{}, @@ -800,7 +808,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { sym_t.recv_tokens = TensorWrapper(sym_recv.ptr, std::vector{sym_buf.recv_capacity, (size_t)hidden_dim_}, kTokDType); - ASSERT_NO_THROW(nvte_ep_prepare(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.token_counts.data(), NVTEEpLayerConfig{sym_t.top_k_, sym_t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.recv_tokens_per_expert.data(), nullptr, &sym_t.layer_cfg_, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.tokens.data(), symm_window(sym_tokens), sym_t.topk_weights.data(), NVTECommWindow{}, diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index d5e006cef6..7cf6017090 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -146,7 +146,7 @@ static bool ep_bootstrap(int argc, char* argv[]) { ncclUniqueId uid{}; exchange_unique_id(&uid); - NVTEEpGroupConfig group_config{}; + NVTEEpGroupConfig group_config = NVTE_EP_GROUP_CONFIG_INIT; group_config.ep_size = g_ep_size; group_config.num_experts = g_num_experts; group_config.max_tokens_per_rank = g_max_tokens_per_rank; @@ -156,7 +156,7 @@ static bool ep_bootstrap(int argc, char* argv[]) { group_config.max_token_dtype = g_max_token_dtype; NVTE_CHECK_NCCL(ncclCommInitRank(&g_ep_comm, g_num_processes, uid, g_process_id)); - nvte_ep_initialize(static_cast(g_ep_comm), group_config); + nvte_ep_initialize(static_cast(g_ep_comm), &group_config); if (g_process_id == 0) { printf("EP initialized: ep_size=%d num_experts=%d " @@ -173,7 +173,7 @@ static bool ep_bootstrap(int argc, char* argv[]) { static void ep_reinitialize(int zero_copy) { if (!g_ep_initialized) return; nvte_ep_shutdown(); - NVTEEpGroupConfig group_config{}; + NVTEEpGroupConfig group_config = NVTE_EP_GROUP_CONFIG_INIT; group_config.ep_size = g_ep_size; group_config.num_experts = g_num_experts; group_config.max_tokens_per_rank = g_max_tokens_per_rank; @@ -181,7 +181,7 @@ static void ep_reinitialize(int zero_copy) { group_config.hidden_dim = g_hidden_dim; group_config.max_token_dtype = g_max_token_dtype; group_config.zero_copy = zero_copy; - nvte_ep_initialize(static_cast(g_ep_comm), group_config); + nvte_ep_initialize(static_cast(g_ep_comm), &group_config); } // Tear down in dependency order: backend's ep_group reads from ep_comm, diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index 66ee3dc8d9..0981289ffe 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -13,6 +13,10 @@ #include +#include +#include +#include + #include "../util/logging.h" #if defined(NVTE_WITH_NCCL_EP) @@ -24,18 +28,30 @@ using transformer_engine::ep::EPBackend; -void nvte_ep_initialize(void* ep_comm, NVTEEpGroupConfig group_config) { - NVTE_CHECK(ep_comm != nullptr, "ep_comm must not be null"); - EPBackend::initialize(static_cast(ep_comm), group_config); -} - -void nvte_ep_shutdown(void) { EPBackend::shutdown(); } - -size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig layer_cfg) { - return EPBackend::get().handle_mem_size(layer_cfg); +namespace { +// Smallest accepted struct_size: covers the base (required) fields. Frozen; +// never raise these. Later fields are read only when struct_size covers them. +constexpr size_t kGroupConfigMinSize = offsetof(NVTEEpGroupConfig, zero_copy) + sizeof(int); +constexpr size_t kLayerConfigMinSize = + offsetof(NVTEEpLayerConfig, dispatch_output_per_expert_alignment) + sizeof(size_t); + +// Copy a caller's versioned config into a full current-layout struct: fields +// the caller did not provide stay zero, extra trailing fields are dropped. +// struct_size 0 is read as the base layout. Requires a size_t struct_size +// first member. +template +Cfg normalize_ep_config(const Cfg* user, size_t min_size, const char* name) { + NVTE_CHECK(user != nullptr, name, " must not be null"); + const size_t want = (user->struct_size == 0) ? min_size : user->struct_size; + NVTE_CHECK(want >= min_size, name, ".struct_size (", user->struct_size, + ") is below the required minimum ", min_size, + "; zero-initialize the struct or set struct_size via NVTE_EP_*_CONFIG_INIT"); + Cfg cfg{}; + std::memcpy(&cfg, user, std::min(want, sizeof(Cfg))); + cfg.struct_size = sizeof(Cfg); + return cfg; } -namespace { inline void* handle_mem_ptr(NVTETensor mem) { void* p = nvte_tensor_data(mem); NVTE_CHECK(p != nullptr, "handle_mem tensor data must not be null"); @@ -43,9 +59,25 @@ inline void* handle_mem_ptr(NVTETensor mem) { } } // namespace -void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor token_counts, - NVTEEpLayerConfig layer_cfg, cudaStream_t stream) { - EPBackend::get().prepare(handle_mem_ptr(handle_mem), topk_idx, token_counts, layer_cfg, stream); +void nvte_ep_initialize(void* ep_comm, const NVTEEpGroupConfig* group_config) { + NVTE_CHECK(ep_comm != nullptr, "ep_comm must not be null"); + NVTEEpGroupConfig cfg = normalize_ep_config(group_config, kGroupConfigMinSize, "group_config"); + EPBackend::initialize(static_cast(ep_comm), cfg); +} + +void nvte_ep_shutdown(void) { EPBackend::shutdown(); } + +size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg) { + NVTEEpLayerConfig cfg = normalize_ep_config(layer_cfg, kLayerConfigMinSize, "layer_cfg"); + return EPBackend::get().handle_mem_size(cfg); +} + +void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv_tokens_per_expert, + NVTETensor total_recv_tokens_per_rank, const NVTEEpLayerConfig* layer_cfg, + cudaStream_t stream) { + NVTEEpLayerConfig cfg = normalize_ep_config(layer_cfg, kLayerConfigMinSize, "layer_cfg"); + EPBackend::get().prepare(handle_mem_ptr(handle_mem), topk_idx, recv_tokens_per_expert, + total_recv_tokens_per_rank, cfg, stream); } void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tokens, @@ -88,15 +120,18 @@ namespace { } } // namespace -void nvte_ep_initialize(void* /*ep_comm*/, NVTEEpGroupConfig /*group_config*/) { ep_not_built(); } +void nvte_ep_initialize(void* /*ep_comm*/, const NVTEEpGroupConfig* /*group_config*/) { + ep_not_built(); +} void nvte_ep_shutdown(void) {} -size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig /*layer_cfg*/) { ep_not_built(); } +size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* /*layer_cfg*/) { ep_not_built(); } void nvte_ep_prepare(NVTETensor /*handle_mem*/, NVTETensor /*topk_idx*/, - NVTETensor /*token_counts*/, NVTEEpLayerConfig /*layer_cfg*/, - cudaStream_t /*stream*/) { + NVTETensor /*recv_tokens_per_expert*/, + NVTETensor /*total_recv_tokens_per_rank*/, + const NVTEEpLayerConfig* /*layer_cfg*/, cudaStream_t /*stream*/) { ep_not_built(); } diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index f1510693bb..a82ec1c98d 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -110,8 +110,8 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { "hidden_dim * sizeof(max_token_dtype) exceeds 4 GiB; got ", row_bytes, " bytes"); NVTE_CHECK(config.num_experts % config.ep_size == 0, "num_experts (", config.num_experts, ") must be divisible by ep_size (", config.ep_size, ")"); - NVTE_CHECK(config.max_num_sms >= 0, "max_num_sms must be >= 0 (0 = auto), got ", - config.max_num_sms); + NVTE_CHECK(config.num_comm_sms >= 0, "num_comm_sms must be >= 0 (0 = auto), got ", + config.num_comm_sms); const int sm = cuda::sm_arch(); NVTE_CHECK(sm >= 90, "NCCL EP requires SM_90+ (Hopper or later), but current device is SM_", sm); @@ -207,8 +207,8 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { cfg.rdma_buffer_size = NCCL_EP_AUTO; cfg.num_qp_per_rank = NCCL_EP_AUTO; cfg.num_channels = NCCL_EP_AUTO; - cfg.max_num_sms = group_config.max_num_sms > 0 - ? static_cast(group_config.max_num_sms) + cfg.max_num_sms = group_config.num_comm_sms > 0 + ? static_cast(group_config.num_comm_sms) : NCCL_EP_AUTO; // Must be > 0; NCCL EP errors out on 0. cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); @@ -319,8 +319,11 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { return hm_size; } -void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor token_counts, - NVTEEpLayerConfig layer_cfg, cudaStream_t stream) { +void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, + NVTETensor recv_tokens_per_expert, + NVTETensor /*total_recv_tokens_per_rank*/, NVTEEpLayerConfig layer_cfg, + cudaStream_t stream) { + // total_recv_tokens_per_rank is a reserved placeholder; not yet populated. NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); NVTE_CHECK(nvte_tensor_shape(topk_idx).ndim == 2, "topk_idx must be 2D [T, top_k]"); @@ -329,13 +332,15 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx, topk_idx_shape); // ncclEpUpdateHandle writes per-expert counts via expert_counters. - NVTEShape token_counts_shape; - ncclEpTensor_t token_counts_desc; - if (token_counts != nullptr) { - token_counts_desc = make_nccl_ep_tensor(token_counts, token_counts_shape); + NVTEShape recv_tokens_per_expert_shape; + ncclEpTensor_t recv_tokens_per_expert_desc; + if (recv_tokens_per_expert != nullptr) { + recv_tokens_per_expert_desc = + make_nccl_ep_tensor(recv_tokens_per_expert, recv_tokens_per_expert_shape); } ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT; - layout_info.expert_counters = (token_counts != nullptr) ? &token_counts_desc : nullptr; + layout_info.expert_counters = + (recv_tokens_per_expert != nullptr) ? &recv_tokens_per_expert_desc : nullptr; std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index 2325baafca..80c9b9cea3 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -46,8 +46,9 @@ class EPBackend { size_t handle_mem_size(NVTEEpLayerConfig layer_cfg); // Seeds the cache for handle_mem with layer_cfg and runs the routing AllGather. - void prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor token_counts, - NVTEEpLayerConfig layer_cfg, cudaStream_t stream); + void prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor recv_tokens_per_expert, + NVTETensor total_recv_tokens_per_rank, NVTEEpLayerConfig layer_cfg, + cudaStream_t stream); // Per-step ops below require a prior prepare(). void dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTETensor tokens, diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 8928b92825..224622fd41 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -28,11 +28,16 @@ extern "C" { #endif /* -- Config structs ------------------------------------------------------- */ -/* TODO: add a struct_size/version field to these configs (and align with other - * TE public structs) once a TE-wide convention for ABI versioning lands. */ +/* Each config begins with struct_size so the API can add fields without + * breaking ABI. The backend reads only the bytes struct_size covers and + * zero-defaults the rest; struct_size 0 means the base layout. Append new + * fields at the end only; never reorder, resize, or remove existing ones. */ /*! \brief Group-level EP configuration (fixed for the EP group lifetime). */ typedef struct { + /*! Struct size in bytes, or 0 for the base layout. Set to + * sizeof(NVTEEpGroupConfig) to include fields added in newer versions. */ + size_t struct_size; /*! EP world size. */ int ep_size; /*! Total experts across all ranks. */ @@ -43,10 +48,11 @@ typedef struct { int max_recv_tokens_per_rank; /*! Token hidden dimension. */ int hidden_dim; - /*! Max SMs for EP kernels. 0 = auto. */ - int max_num_sms; + /*! Max SMs for NCCL EP dispatch/combine kernels. 0 = auto. */ + int num_comm_sms; /*! Widest token dtype the group will dispatch; sizes staging buffers. - * Per-dispatch tensors may use any dtype with element size <= this. */ + * Required (no default): must be set to a real token dtype. Per-dispatch + * tensors may use any dtype with element size <= this. */ NVTEDType max_token_dtype; /*! Zero-copy dispatch/combine. When nonzero, payload tensors must be backed * by NVTECommWindow handles and transfer in place (no staging copies); @@ -59,6 +65,9 @@ typedef struct { * overflow policy, ...). */ typedef struct { + /*! Struct size in bytes, or 0 for the base layout. Set to + * sizeof(NVTEEpLayerConfig) to include fields added in newer versions. */ + size_t struct_size; /*! Per-token expert fan-out (> 0). */ int top_k; /*! Per-expert recv-slab alignment in tokens (power of two; 0/1 disables). @@ -67,6 +76,12 @@ typedef struct { size_t dispatch_output_per_expert_alignment; } NVTEEpLayerConfig; +/* Zero-init a config with struct_size set to the current layout: + * NVTEEpGroupConfig cfg = NVTE_EP_GROUP_CONFIG_INIT; + * cfg.ep_size = ...; */ +#define NVTE_EP_GROUP_CONFIG_INIT {sizeof(NVTEEpGroupConfig)} +#define NVTE_EP_LAYER_CONFIG_INIT {sizeof(NVTEEpLayerConfig)} + /* -- Bootstrap ------------------------------------------------------------ */ /*! \brief Bootstrap the EP backend from an existing NCCL EP sub-communicator. @@ -78,9 +93,9 @@ typedef struct { * group per process, bound to the current CUDA device. * * \param[in] ep_comm Opaque ncclComm_t for the EP sub-group. - * \param[in] group_config Group-level EP configuration. + * \param[in] group_config Group-level EP configuration (struct_size set). */ -void nvte_ep_initialize(void* ep_comm, NVTEEpGroupConfig group_config); +void nvte_ep_initialize(void* ep_comm, const NVTEEpGroupConfig* group_config); /*! \brief Tear down the EP backend. Idempotent. Does not destroy ep_comm. */ void nvte_ep_shutdown(void); @@ -94,10 +109,10 @@ void nvte_ep_shutdown(void); * for that layer (the backend keys its cache on the pointer). Host-only; * size is stable for a given (group, layer) pair. * - * \param[in] layer_cfg Per-call layer configuration. + * \param[in] layer_cfg Per-call layer configuration (struct_size set). * \return size in bytes for the handle_mem buffer. */ -size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig layer_cfg); +size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg); /* -- Per-step ops (all allocation-free, CUDA graph-capturable) ------------ */ @@ -106,17 +121,19 @@ size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig layer_cfg); * AllGathers topk_idx across the EP group and stages per-expert offsets and * counts into handle_mem so the matching dispatch/combine/_bwd can run with * no further routing computation. Must precede every dispatch/combine/_bwd - * that uses this handle_mem. token_counts becomes host-valid after a stream - * sync. - * - * \param[in] handle_mem uint8 routing-state buffer. - * \param[in] topk_idx [T, top_k] int64 routing indices. - * \param[out] token_counts [num_local_experts] int32 counts. - * \param[in] layer_cfg Per-call layer configuration. - * \param[in] stream CUDA stream. + * that uses this handle_mem. recv_tokens_per_expert becomes host-valid after a + * stream sync. + * + * \param[in] handle_mem uint8 routing-state buffer. + * \param[in] topk_idx [T, top_k] int64 routing indices. + * \param[out] recv_tokens_per_expert [num_local_experts] int32 counts. + * \param[out] total_recv_tokens_per_rank Reserved placeholder; may be null. Unused for now. + * \param[in] layer_cfg Per-call layer configuration (struct_size set). + * \param[in] stream CUDA stream. */ -void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor token_counts, - NVTEEpLayerConfig layer_cfg, cudaStream_t stream); +void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv_tokens_per_expert, + NVTETensor total_recv_tokens_per_rank, const NVTEEpLayerConfig* layer_cfg, + cudaStream_t stream); /*! \brief Dispatch tokens (and routing weights) to expert ranks. * diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index ee204e7594..bfd96776c7 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -45,16 +45,17 @@ class EpResources { NVTE_CHECK_NCCL(ncclCommInitRank(&comm_, p.ep_size, uid, p.rank_within_group)); // zero_copy=0: JAX EP path always stages payloads; the zero-copy fast path // requires NVTECommWindow-backed tensors, which JAX bindings don't expose. - NVTEEpGroupConfig cfg{.ep_size = p.ep_size, + NVTEEpGroupConfig cfg{.struct_size = sizeof(NVTEEpGroupConfig), + .ep_size = p.ep_size, .num_experts = p.num_experts, .max_tokens_per_rank = p.max_tokens_per_rank, .max_recv_tokens_per_rank = p.max_recv_tokens_per_rank, .hidden_dim = p.hidden_dim, - .max_num_sms = p.max_num_sms, + .num_comm_sms = p.max_num_sms, .max_token_dtype = p.max_token_dtype, .zero_copy = 0}; try { - nvte_ep_initialize(static_cast(comm_), cfg); + nvte_ep_initialize(static_cast(comm_), &cfg); } catch (...) { ncclCommDestroy(comm_); comm_ = nullptr; @@ -162,8 +163,11 @@ void ReleaseEpResources() { } size_t EpHandleMemSize(int top_k, size_t dispatch_output_per_expert_alignment) { - NVTEEpLayerConfig layer_cfg{top_k, dispatch_output_per_expert_alignment}; - return nvte_ep_handle_mem_size(layer_cfg); + NVTEEpLayerConfig layer_cfg{ + .struct_size = sizeof(NVTEEpLayerConfig), + .top_k = top_k, + .dispatch_output_per_expert_alignment = dispatch_output_per_expert_alignment}; + return nvte_ep_handle_mem_size(&layer_cfg); } pybind11::capsule GetEpInstanceStateTypeIdCapsule() { @@ -192,8 +196,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpInstantiateHandler, EpInstantiateImpl, FFI::Bind // ── ep_prepare ──────────────────────────────────────────────────────────────── Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type topk_idx, - Result_Type token_counts, Result_Type handle_mem, Result_Type workspace, - EpConfig config) { + Result_Type recv_tokens_per_expert, Result_Type handle_mem, + Result_Type workspace, EpConfig config) { (void)ep_state; // lifetime only. auto topk_dims = topk_idx.dimensions(); NVTE_CHECK(topk_dims.size() >= 2, @@ -218,15 +222,19 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T } auto topk_idx_ = TensorWrapper(topk_idx_data, topk_shape, DType::kInt64); - std::vector tc_shape = {static_cast(token_counts->element_count())}; - auto token_counts_ = TensorWrapper(token_counts->untyped_data(), tc_shape, DType::kInt32); + std::vector tc_shape = {static_cast(recv_tokens_per_expert->element_count())}; + auto recv_tokens_per_expert_ = + TensorWrapper(recv_tokens_per_expert->untyped_data(), tc_shape, DType::kInt32); std::vector hm_shape = {static_cast(handle_mem->element_count())}; auto handle_mem_ = TensorWrapper(handle_mem->untyped_data(), hm_shape, DType::kByte); - NVTEEpLayerConfig layer_cfg{static_cast(config.top_k), - static_cast(config.dispatch_output_per_expert_alignment)}; - nvte_ep_prepare(handle_mem_.data(), topk_idx_.data(), token_counts_.data(), layer_cfg, stream); + NVTEEpLayerConfig layer_cfg{.struct_size = sizeof(NVTEEpLayerConfig), + .top_k = static_cast(config.top_k), + .dispatch_output_per_expert_alignment = + static_cast(config.dispatch_output_per_expert_alignment)}; + nvte_ep_prepare(handle_mem_.data(), topk_idx_.data(), recv_tokens_per_expert_.data(), + /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); return ffi_with_cuda_error_check(); } @@ -235,8 +243,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, .Ctx() // stream .Ctx<::xla::ffi::State>() // EP state .Arg() // topk_idx - .Ret() // token_counts - .Ret() // handle_mem + .Ret() // recv_tokens_per_expert + .Ret() // handle_mem .Ret() // workspace (FFI scratch) .Attrs(), FFI_CudaGraph_Traits); diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index d1ef76af40..ae23c705e5 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -136,16 +136,17 @@ void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t nu NVTE_CHECK(ncclCommCount(ep_comm, &ep_size) == ncclSuccess, "ncclCommCount failed"); auto torch_dtype = max_token_dtype.cast(); NVTEEpGroupConfig cfg{ - /*ep_size=*/ep_size, - /*num_experts=*/static_cast(num_experts), - /*max_tokens_per_rank=*/static_cast(max_tokens_per_rank), - /*max_recv_tokens_per_rank=*/static_cast(max_recv_tokens_per_rank), - /*hidden_dim=*/static_cast(hidden_dim), - /*max_num_sms=*/static_cast(max_num_sms), - /*max_token_dtype=*/static_cast(GetTransformerEngineDType(torch_dtype)), - /*zero_copy=*/zero_copy ? 1 : 0, + .struct_size = sizeof(NVTEEpGroupConfig), + .ep_size = ep_size, + .num_experts = static_cast(num_experts), + .max_tokens_per_rank = static_cast(max_tokens_per_rank), + .max_recv_tokens_per_rank = static_cast(max_recv_tokens_per_rank), + .hidden_dim = static_cast(hidden_dim), + .num_comm_sms = static_cast(max_num_sms), + .max_token_dtype = static_cast(GetTransformerEngineDType(torch_dtype)), + .zero_copy = zero_copy ? 1 : 0, }; - nvte_ep_initialize(static_cast(ep_comm), cfg); + nvte_ep_initialize(static_cast(ep_comm), &cfg); g_zero_copy_enabled.store(zero_copy, std::memory_order_relaxed); g_ep_initialized = true; g_ep_group_name = group_name; @@ -164,17 +165,18 @@ namespace { NVTEEpLayerConfig make_layer_cfg(int64_t top_k, int64_t dispatch_output_per_expert_alignment) { return NVTEEpLayerConfig{ - /*top_k=*/static_cast(top_k), - /*dispatch_output_per_expert_alignment=*/ - static_cast(dispatch_output_per_expert_alignment), + .struct_size = sizeof(NVTEEpLayerConfig), + .top_k = static_cast(top_k), + .dispatch_output_per_expert_alignment = + static_cast(dispatch_output_per_expert_alignment), }; } } // namespace int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_alignment) { - return static_cast( - nvte_ep_handle_mem_size(make_layer_cfg(top_k, dispatch_output_per_expert_alignment))); + auto layer_cfg = make_layer_cfg(top_k, dispatch_output_per_expert_alignment); + return static_cast(nvte_ep_handle_mem_size(&layer_cfg)); } // ── Per-step ops ───────────────────────────────────────────────────────────── @@ -194,8 +196,9 @@ void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor token_cou auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + auto layer_cfg = make_layer_cfg(top_k, dispatch_output_per_expert_alignment); nvte_ep_prepare(handle_mem_te.data(), topk_idx_te.data(), token_counts_te.data(), - make_layer_cfg(top_k, dispatch_output_per_expert_alignment), stream); + /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); } void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, From 613c545312cf7772db8e50e2408964e65ae415cc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 30 Jun 2026 23:44:26 +0200 Subject: [PATCH 20/42] Stamp value-opaque flag only after successful registration Move the _VALUE_OPAQUE_FLAG setattr to the end of register_value_opaque_quantizer, after register_opaque_type succeeds (or the type is already opaque). Previously the flag was set up front, so is_value_opaque_quantizer reported True even when the opaque-object API was missing or registration raised, since both paths are swallowed. Eager value semantics (__eq__/__hash__/__fx_repr__) are independent of the flag, so this only tightens the predicate to mean torch actually knows the type as opaque. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 8b8b3caa69..a1106d6292 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -87,12 +87,10 @@ def register_value_opaque_quantizer(cls: type) -> None: a non-``None`` ``_value_fields`` (see :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). """ - # Stamp the class so it can be recognized as value-opaque in dynamo-traced - # code (used to fall back to eager for unregistered quantizers). - setattr(cls, _VALUE_OPAQUE_FLAG, True) - # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the - # class, so attach it before registering. + # class, so attach it before registering. Eager value semantics + # (``__eq__`` / ``__hash__`` / ``__fx_repr__``) work regardless of whether + # the opaque-object registration below succeeds. if "__fx_repr__" not in cls.__dict__: cls.__fx_repr__ = _quantizer_fx_repr @@ -113,4 +111,9 @@ def register_value_opaque_quantizer(cls: type) -> None: # Keep TE importable: neither the opaque-type query nor the registration # must crash the import, e.g. on PyTorch versions with only partial / # experimental opaque-object support. - pass + return + + # Stamp the class only once torch actually knows it as an opaque value type, + # so ``is_value_opaque_quantizer`` never reports a quantizer as opaque when + # the registration was skipped or failed. + setattr(cls, _VALUE_OPAQUE_FLAG, True) From 9db604f4f321a9919134827036a59825faf4d786 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 30 Jun 2026 23:45:25 +0200 Subject: [PATCH 21/42] Drop verbose comments around value-opaque flag stamping Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/quantizer_opaque.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index a1106d6292..4ba4761421 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -88,9 +88,7 @@ def register_value_opaque_quantizer(cls: type) -> None: :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). """ # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the - # class, so attach it before registering. Eager value semantics - # (``__eq__`` / ``__hash__`` / ``__fx_repr__``) work regardless of whether - # the opaque-object registration below succeeds. + # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: cls.__fx_repr__ = _quantizer_fx_repr @@ -113,7 +111,4 @@ def register_value_opaque_quantizer(cls: type) -> None: # experimental opaque-object support. return - # Stamp the class only once torch actually knows it as an opaque value type, - # so ``is_value_opaque_quantizer`` never reports a quantizer as opaque when - # the registration was skipped or failed. setattr(cls, _VALUE_OPAQUE_FLAG, True) From 3011dfd879a7359186332417da951551de3e8d97 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 30 Jun 2026 23:50:37 +0200 Subject: [PATCH 22/42] Narrow value process-group check to amax_reduction_group _check_value_has_no_process_group ran on every guard eval (via __eq__/__hash__) and scanned all of vars(self) recursively. The only attribute that can hold a ProcessGroup is the deprecated amax_reduction_group, so check it directly (O(1)) and drop the _contains_process_group helper. Same guarantee, off the hot path. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../pytorch/quantized_tensor.py | 38 +++++++------------ 1 file changed, 13 insertions(+), 25 deletions(-) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 033a35f8e1..93af20cc9b 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -24,19 +24,6 @@ ) -def _contains_process_group(value: Any) -> bool: - """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. - - Checks the value directly and one level of ``tuple``/``list`` nesting, which - covers the shapes a quantizer value field could plausibly take. - """ - if isinstance(value, dist_group_type): - return True - if isinstance(value, (tuple, list)): - return any(_contains_process_group(item) for item in value) - return False - - # Custom ops that should pass through __torch_dispatch__ without unwrapping # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that # handle quantized tensors internally. @@ -446,18 +433,19 @@ def _check_value_has_no_process_group(self) -> None: # value key, which cannot carry live distributed state. Enforced here -- # the single point every value-materialization path (``__eq__`` / # ``__hash__`` / ``__fx_repr__``) goes through -- so a custom - # ``__fx_repr__`` cannot bypass it. Reject any field holding a - # ProcessGroup (e.g. the deprecated ``amax_reduction_group``) rather than - # silently dropping it; pass the reduction group per quantize call. - for name, value in vars(self).items(): - if _contains_process_group(value): - raise TypeError( - f"{type(self).__name__} cannot be used as a torch.compile value " - f"object: attribute {name!r} holds a torch.distributed.ProcessGroup, " - "which is live distributed state and must not be baked into an FX " - "graph. Pass the amax reduction group per quantize call instead of " - "storing it on the quantizer." - ) + # ``__fx_repr__`` cannot bypass it. The only attribute that can hold a + # ProcessGroup is the deprecated ``amax_reduction_group`` (a scalar group + # excluded from the value key); reject it rather than silently dropping + # it -- otherwise a stored group would compare/hash equal to a groupless + # quantizer. Pass the reduction group per quantize call instead. + if isinstance(getattr(self, "amax_reduction_group", None), dist_group_type): + raise TypeError( + f"{type(self).__name__} cannot be used as a torch.compile value " + "object: 'amax_reduction_group' holds a torch.distributed.ProcessGroup, " + "which is live distributed state and must not be baked into an FX " + "graph. Pass the amax reduction group per quantize call instead of " + "storing it on the quantizer." + ) def _value_key(self) -> Tuple[Any, ...]: """Hashable, reproducible key identifying this quantizer's value. From fe5e5dba70f5fe5cd11620a432ec6470d011c1ac Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 30 Jun 2026 23:52:46 +0200 Subject: [PATCH 23/42] Shorten amax_reduction_group check comment Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/quantized_tensor.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 93af20cc9b..6161539b9c 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -429,15 +429,9 @@ def _value_fields(self) -> Optional[Tuple[str, ...]]: return None def _check_value_has_no_process_group(self) -> None: - # A value quantizer is baked into the FX graph as a constant via its - # value key, which cannot carry live distributed state. Enforced here -- - # the single point every value-materialization path (``__eq__`` / - # ``__hash__`` / ``__fx_repr__``) goes through -- so a custom - # ``__fx_repr__`` cannot bypass it. The only attribute that can hold a - # ProcessGroup is the deprecated ``amax_reduction_group`` (a scalar group - # excluded from the value key); reject it rather than silently dropping - # it -- otherwise a stored group would compare/hash equal to a groupless - # quantizer. Pass the reduction group per quantize call instead. + # A value quantizer cannot carry live distributed state into the FX + # graph; reject a stored ``amax_reduction_group`` and pass it per + # quantize call instead. if isinstance(getattr(self, "amax_reduction_group", None), dist_group_type): raise TypeError( f"{type(self).__name__} cannot be used as a torch.compile value " From 6f66c3e2429ec1dcbdf97c5ecc2aaa5e8a879b53 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 1 Jul 2026 00:03:16 +0200 Subject: [PATCH 24/42] Drop trivial value-equality boilerplate from quantizer test Remove the a==b / hash / dict-key block that just exercised Python's own dict semantics; equality and hashing are still covered by the __fx_repr__ round-trip (rebuilt == a, hash match) and the bit-exact kernel check. other_kwargs is now unused, so drop it from the parametrization and both test signatures. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 24 ++++++++---------------- 1 file changed, 8 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 63cb82eca8..c7f4ceb71c 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -441,12 +441,11 @@ def _hw_available(quantizer): # (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ - pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), - pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), - pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), + pytest.param(_mxfp8, id="mxfp8"), + pytest.param(_blockwise, id="float8_blockwise"), + pytest.param(_current_scaling, id="float8_current_scaling"), pytest.param( _nvfp4, - {"with_rht": False}, id="nvfp4", marks=pytest.mark.skipif( not torch.cuda.is_available(), @@ -456,17 +455,10 @@ def _hw_available(quantizer): ] -@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) -def test_quantizer_value_object(factory, other_kwargs): +@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory): """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" - a, b = factory(), factory() - # Same config -> equal, same hash, interchangeable as a dict/set key. - assert a is not b - assert a == b - assert hash(a) == hash(b) - assert {a: "x"}[b] == "x" - # Different config -> not equal. - assert a != factory(**other_kwargs) + a = factory() # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. repr_str, globals_ = a.__fx_repr__() @@ -538,8 +530,8 @@ def _qdq_fake(x, q): not _opaque_available, reason="torch.compile opaque-object support requires PyTorch >= 2.11", ) -@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) -def test_quantizer_value_object_fullgraph(factory, other_kwargs): +@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory): """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. A custom op quantizes+dequantizes with the (opaque value) quantizer; the From 25ad5cb904c602d04291d88ee04c6f6f3ea3467d Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 30 Jun 2026 17:51:38 -0700 Subject: [PATCH 25/42] Graph Safe Current Scaling Support for GroupedLinear Module/Ops + Fix CUBLAS GGEMM heuristics (#3143) * support in grouped linear and relevant tests Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Unecessary details remove Removed details about FP8 current scaling methods. Signed-off-by: vthumbe1503 * fix grouped linear module's grouped tensor path Signed-off-by: Varun Thumbe * allow more current scaling use-cases.. block nvfp4+rht+single grouped weight being cuda graphable Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * some minor comment fixing Signed-off-by: Varun Thumbe * fix heuristics Signed-off-by: Varun Thumbe * only dealyed scaling skip in failure comment Signed-off-by: vthumbe1503 * address review comment Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix for other 2 nvte APIs Signed-off-by: Varun Thumbe * fix m and n Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_grouped_linear.py | 31 ++++++++++--- tests/pytorch/test_grouped_mlp.py | 45 ++++++++++++++---- .../common/gemm/cublaslt_grouped_gemm.cu | 27 ++++++----- .../pytorch/module/grouped_linear.py | 14 ++++-- .../pytorch/ops/basic/grouped_linear.py | 46 +++++++++++++------ 5 files changed, 120 insertions(+), 43 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index caa84ec02a..01a7cf2415 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1496,6 +1496,7 @@ def test_fp8_grouped_gemm(shape, accumulate): _FUSED_GROUPED_GEMM_ENV = "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM" _ALL_BOOLEAN = all_boolean +_fp8_available, _reason_for_no_fp8 = fp8_available, reason_for_no_fp8 _mxfp8_available, _reason_for_no_mxfp8 = mxfp8_available, reason_for_no_mxfp8 _nvfp4_available, _reason_for_no_nvfp4 = nvfp4_available, reason_for_no_nvfp4 @@ -1577,6 +1578,10 @@ def _run_grouped_linear_path( "fp8_recipe", [ None, + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + ), pytest.param( recipe.MXFP8BlockScaling(), marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), @@ -1586,7 +1591,7 @@ def _run_grouped_linear_path( marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), ], - ids=["bf16", "mxfp8", "nvfp4"], + ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) @pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN) @@ -1600,8 +1605,13 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( pytest.skip( "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." ) - if use_fp8 and device_capability < (10, 0): - pytest.skip("Quantized GroupedTensor grouped GEMM path requires Blackwell (SM100+).") + # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor + # current scaling also runs on the Hopper grouped GEMM path. + is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() + if use_fp8 and not is_current_scaling and device_capability < (10, 0): + pytest.skip( + "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." + ) cublaslt_version = tex.get_cublasLt_version() if device_capability < (10, 0) and cublaslt_version < 130400: pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") @@ -1786,6 +1796,10 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): "fp8_recipe", [ None, + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + ), pytest.param( recipe.MXFP8BlockScaling(), marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), @@ -1795,7 +1809,7 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), ], - ids=["bf16", "mxfp8", "nvfp4"], + ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch): @@ -1806,8 +1820,13 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch pytest.skip( "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." ) - if use_fp8 and device_capability < (10, 0): - pytest.skip("Quantized GroupedTensor grouped GEMM path requires Blackwell (SM100+).") + # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor + # current scaling also runs on the Hopper grouped GEMM path. + is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() + if use_fp8 and not is_current_scaling and device_capability < (10, 0): + pytest.skip( + "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." + ) cublaslt_version = tex.get_cublasLt_version() if device_capability < (10, 0) and cublaslt_version < 130400: pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index cb90ac6bd9..e24fff9049 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -288,13 +288,9 @@ def test_grouped_linear( if single_grouped_bias and not bias: pytest.skip("single_grouped_bias requires bias=True") - if ( - single_grouped_weight - and quantized_weight - and quantization in ("fp8_delayed_scaling", "fp8_current_scaling") - ): + if single_grouped_weight and quantized_weight and quantization in ("fp8_delayed_scaling"): pytest.skip( - "single_grouped_weight does not support FP8 delayed/current scaling " + "single_grouped_weight does not support FP8 delayed scaling " "with quantized_model_init" ) @@ -439,7 +435,10 @@ def test_grouped_linear( @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) @pytest.mark.parametrize( "quantization", - [None] + (["mxfp8"] if mxfp8_available else []), + [None] + + (["fp8_current_scaling"] if fp8_available else []) + + (["mxfp8"] if mxfp8_available else []) + + (["nvfp4_rht"] if nvfp4_available else []), ) @pytest.mark.parametrize("quantized_weight", (False, True)) @pytest.mark.parametrize("bias", (False, True)) @@ -475,10 +474,38 @@ def test_grouped_linear_cuda_graph_safe( "single_grouped_weight/single_grouped_bias requires" " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" ) - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") + device_capability = torch.cuda.get_device_capability() + if device_capability < (9, 0): + pytest.skip( + "Grouped GEMM CUDA-graph-safe path requires Hopper (SM90) or Blackwell (SM100+)" + ) + # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path, + # but MXFP8/NVFP4 grouped quantization kernels require Blackwell (SM100+). + requires_blackwell = quantization is not None and quantization != "fp8_current_scaling" + if requires_blackwell and device_capability < (10, 0): + pytest.skip("MXFP8/NVFP4 grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") + # Grouped GEMM on Hopper requires cuBLAS 13.4+; Blackwell requires cuBLAS 13.3+. + cublaslt_version = tex.get_cublasLt_version() + if device_capability < (10, 0) and cublaslt_version < 130400: + pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") + if cublaslt_version < 130300: + pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") if quantization is None and quantized_weight: pytest.skip("quantized_weight requires a quantization recipe") + if ( + quantization is not None + and quantization.startswith("nvfp4") + and dtype != torch.bfloat16 + ): + pytest.skip("NVFP4 grouped GEMM only supports BF16 output") + if single_grouped_weight and quantization is not None and quantization.startswith("nvfp4"): + # Currently, split_quantization is used which is not cuda graph safe. + # We should either support grouped weight quantization without rht or need to do + # inplace per tensor weight quantization to make this use-case cuda graphable if needed. + pytest.skip( + "NVFP4 grouped GEMM with single_grouped_weight is not supported yet; " + "only discrete weights (single_grouped_weight=False) are supported." + ) single_grouped_bias = bias and single_grouped_weight diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 481b3ac1ea..44502422cd 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1625,7 +1625,7 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT inputC->dtype(), outputD->dtype()); // Compute average dimensions for heuristics - // K dimension: if transa, K is A's first dim; if not, K is A's last dim + // K dimension: if transa, K is A's last dim; if not, K is A's first dim // Use original inputA and transa for heuristics (not modified A_sel.trans) GroupedGemmConfig gemm_config; gemm_config.use_split_accumulator = config_.use_split_accumulator; @@ -1633,10 +1633,12 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT gemm_config.use_per_group_alpha_beta = use_per_group_alpha_beta; gemm_config.alpha_dptr = alpha_tensor->data.dptr; gemm_config.beta_dptr = beta_tensor->data.dptr; - gemm_config.avg_m = config_.avg_m.value_or(compute_avg_first_dim(outputD)); - gemm_config.avg_n = config_.avg_n.value_or(compute_avg_last_dim(outputD)); + // avg m = avg num of rows of D in column-major = avg last dim of D + // avg n = avg num of cols of D in column-major = avg first dim of D + gemm_config.avg_m = config_.avg_m.value_or(compute_avg_last_dim(outputD)); + gemm_config.avg_n = config_.avg_n.value_or(compute_avg_first_dim(outputD)); gemm_config.avg_k = - config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); + config_.avg_k.value_or(transa ? compute_avg_last_dim(inputA) : compute_avg_first_dim(inputA)); gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, gemm_config, workspace.cublas_workspace_ptr, stream); @@ -1783,10 +1785,9 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num gemm_config.use_per_group_alpha_beta = use_per_group_alpha_beta; gemm_config.alpha_dptr = alpha_tensor->data.dptr; gemm_config.beta_dptr = beta_tensor->data.dptr; - gemm_config.avg_m = config_.avg_m.value_or(compute_avg_first_dim(outputD)); - gemm_config.avg_n = - config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB)); - gemm_config.avg_k = config_.avg_k.value_or(transa ? avg_first_dim : avg_last_dim); + gemm_config.avg_m = config_.avg_m.value_or(compute_avg_last_dim(outputD)); + gemm_config.avg_n = config_.avg_n.value_or(compute_avg_first_dim(outputD)); + gemm_config.avg_k = config_.avg_k.value_or(transa ? avg_last_dim : avg_first_dim); gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, gemm_config, workspace.cublas_workspace_ptr, stream); @@ -1869,12 +1870,16 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, gemm_config.use_per_group_alpha_beta = use_per_group_alpha_beta; gemm_config.alpha_dptr = alpha_tensor->data.dptr; gemm_config.beta_dptr = beta_tensor->data.dptr; + // D is a discrete list here, so derive the heuristic dims from the grouped inputs instead. + // avg m (rows of D in column-major) is derived from A and avg n (cols of D) from B. The + // reduction dim K is A's last dim when transa, otherwise its first dim (cuBLAS views operands + // transposed) -- matching nvte_grouped_gemm / _with_discrete_inputA. gemm_config.avg_m = - config_.avg_m.value_or(transa ? compute_avg_last_dim(inputA) : compute_avg_first_dim(inputA)); + config_.avg_m.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); gemm_config.avg_n = - config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB)); + config_.avg_n.value_or(transb ? compute_avg_last_dim(inputB) : compute_avg_first_dim(inputB)); gemm_config.avg_k = - config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); + config_.avg_k.value_or(transa ? compute_avg_last_dim(inputA) : compute_avg_first_dim(inputA)); gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, d_dtype, num_tensors, gemm_config, workspace.cublas_workspace_ptr, stream); diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f2fa8b657e..c79adbf7e4 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -103,9 +103,10 @@ def _is_grouped_tensor_path_supported( and be incompatible with CUDA Graphs. Supported Compute Capability (CC) and precisions: - * Hopper (CC 9.0): BF16/FP16. - * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT. - FP8 delayed / current scaling, and FP8 block scaling are not supported because the + * Hopper (CC 9.0): BF16/FP16 and FP8 per-tensor current scaling. + * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8 + per-tensor current scaling. + FP8 delayed scaling and FP8 block scaling are not supported because the corresponding grouped quantization kernels are missing. Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization currently requires RHT. @@ -133,6 +134,9 @@ def _is_grouped_tensor_path_supported( return False # 5. Filter by quantization recipes. if fp8: + if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): + return True + # MXFP8 and NVFP4 require Blackwell+. if not (10, 0) <= get_device_compute_capability() <= (11, 0): return False return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( @@ -328,7 +332,9 @@ def _forward_grouped_tensor( if is_grad_enabled: if weight_requires_grad: - if fp8: + # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data + # in backward pass) + if fp8 and grouped_x.columnwise_data is not None: grouped_x.rowwise_data = None grouped_x.scale_inv = None else: diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 4bbd75bc64..057c9da0e6 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -26,7 +26,13 @@ from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...quantized_tensor import QuantizedTensorStorage -from ...tensor import MXFP8Quantizer, MXFP8Tensor, NVFP4Quantizer, Quantizer +from ...tensor import ( + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + MXFP8Tensor, + NVFP4Quantizer, + Quantizer, +) from ...utils import ( canonicalize_device, canonicalize_dtype, @@ -417,11 +423,8 @@ def make_grouped_weights(self) -> None: quantizer = self.get_quantizer("forward", 1) recipe = None if quantizer is None else quantizer._get_compatible_recipe() - if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()): - raise RuntimeError( - "Delayed scaling or float8 current scaling is not supported with" - " single_grouped_weight=True" - ) + if recipe is not None and recipe.delayed(): + raise RuntimeError("Delayed scaling is not supported with single_grouped_weight=True") grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=self.num_groups, @@ -759,6 +762,7 @@ def _is_graph_safe_path_supported( with_quantized_compute: bool, input_quantizers: Sequence[Optional[Quantizer]], dtype: torch.dtype, + single_grouped_weight: bool, ) -> bool: """Whether the graph-safe grouped-tensor flow can be used. @@ -768,8 +772,13 @@ def _is_graph_safe_path_supported( requirement without duplicating its cuBLAS version checks. * Quantized compute supports MXFP8 and NVFP4 on Blackwell GPUs with Compute Capability (CC) 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped quantization currently - requires RHT; - Every other quantization recipe (fp8 delayed / current scaling, fp8 block scaling, ...) + requires RHT. NVFP4 is additionally restricted to discrete weights: with + ``single_grouped_weight=True`` the weight quantizer is non-RHT and cannot use the + graph-safe grouped quantize kernel, so we fall back to the split-quantize flow. + * FP8 per-tensor current scaling is backed by grouped current-scaling quantization + (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling, + which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0). + Every other quantization recipe (fp8 delayed scaling, fp8 block scaling, ...) falls back to the legacy flow because the corresponding grouped quantization kernels are missing. * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0) @@ -780,11 +789,20 @@ def _is_graph_safe_path_supported( if not (9, 0) <= get_device_compute_capability() <= (11, 0): return False if with_quantized_compute: + # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM + # path; the compute-capability range was already checked above. + if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): + return True + # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. if not (10, 0) <= get_device_compute_capability() <= (11, 0): return False - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( - isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers - ) + if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): + return True + # NVFP4 graph-safe grouped quantization requires RHT and only supports + # discrete weights; otherwise fall back to the split-quantize flow. + if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): + return not single_grouped_weight + return False return dtype in (torch.bfloat16, torch.float16) def _get_grouped_weight_for_gemm( @@ -971,6 +989,7 @@ def fuser_forward( with_quantized_compute=with_quantized_compute, input_quantizers=input_quantizers, dtype=dtype, + single_grouped_weight=self.single_grouped_weight, ) if use_grouped_tensor_path: @@ -1318,8 +1337,9 @@ def _fuser_forward_grouped_tensor( # [split_sizes, base_split_offsets, split_points, # (scales if _scale_bias), grouped_x, *weights] if grouped_x is not None: - if with_quantized_compute: - # only columnwise data is needed for wgrad + # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data + # in backward pass) + if with_quantized_compute and grouped_x.columnwise_data is not None: grouped_x.rowwise_data = None grouped_x.scale_inv = None saved: list[Optional[torch.Tensor]] = [split_sizes, base_split_offsets, split_points] From 4cd705b75394563c0246bdddfa5d3148106c9285 Mon Sep 17 00:00:00 2001 From: HaochenYuan <106647990+HaochenYuan@users.noreply.github.com> Date: Wed, 1 Jul 2026 23:57:17 +0800 Subject: [PATCH 26/42] [PyT] [Common] add support for enabling cuda graph under thd format in megatron. (#2898) * add support for THD CUDA graph Signed-off-by: HaochenYuan * modify comment Signed-off-by: HaochenYuan * address @timmoon10: drop FAv2-bwd alloc gate, rely on THD tail zero-fill Signed-off-by: HaochenYuan * Support graph-safe MoE aux loss token count Signed-off-by: HaochenYuan * add graph guard for one zero fill Signed-off-by: HaochenYuan * remove redundant zero-fill Signed-off-by: HaochenYuan * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename & remove prelude kernel Signed-off-by: HaochenYuan * Update warp reduction function Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: HaochenYuan Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- .../attention/run_attention_with_cp.py | 13 +++ tests/pytorch/test_fused_router.py | 70 ++++++++++++ .../common/fused_router/fused_moe_aux_loss.cu | 108 ++++++++++++++++++ .../include/transformer_engine/fused_router.h | 41 +++++-- .../dot_product_attention/backends.py | 23 ++++ .../dot_product_attention/context_parallel.py | 44 ++++++- .../dot_product_attention.py | 24 ++-- transformer_engine/pytorch/csrc/extensions.h | 4 + .../pytorch/csrc/extensions/pybind.cpp | 6 +- .../pytorch/csrc/extensions/router.cpp | 31 +++++ transformer_engine/pytorch/router.py | 48 +++++--- 11 files changed, 373 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 3d2f99b51b..82b9df262f 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -411,6 +411,13 @@ def run_dpa_with_cp( cu_seqlens_kv=cu_seqlens_kv, cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_padded, + # Test runner sets cu_seqlens_q == cu_seqlens_q_padded for the + # FlashAttention path, i.e. no inter-sequence padding. Declare this + # explicitly so the sync-free auto-detect (which conservatively + # picks True when padded cu_seqlens are present) does not disable FA. + pad_between_seqs=( + (kernel_backend != "FlashAttention") if qkv_format == "thd" else None + ), fp8_output=fp8_mha, ) if config.return_max_logit: @@ -528,6 +535,12 @@ def run_dpa_with_cp( cu_seqlens_kv=cu_seqlens_kv, cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_padded, + # See note above (non-CP branch): same explicit declaration so + # FlashAttention isn't disabled by the conservative sync-free + # auto-detect when this test path constructs no inter-seq padding. + pad_between_seqs=( + (kernel_backend != "FlashAttention") if qkv_format == "thd" else None + ), fp8_output=fp8_mha, ) if config.return_max_logit: diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index ab12216df8..68d3ed9565 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -523,6 +523,76 @@ 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) +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 + value via the device-side coefficient computation.""" + dtype = torch.float32 + num_tokens = 4096 + num_experts = 128 + topk = 4 + num_cols = num_experts + coeff = 0.01 + + offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 + probs = ( + torch.arange(-num_cols // 2, num_cols // 2, device="cuda", dtype=dtype) * 1e-2 + ).unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) + probs = probs.contiguous().requires_grad_(True) + tokens_per_expert = torch.randint(1, 1000, (num_cols,), device="cuda", dtype=torch.int32) + + total_num_tokens_dev = torch.tensor(num_tokens, dtype=torch.int64, device="cuda") + + # Warmup on a side stream to satisfy CUDA Graph capture requirements. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + warmup_out = fused_moe_aux_loss( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=total_num_tokens_dev, + num_experts=num_experts, + topk=topk, + coeff=coeff, + ) + torch.autograd.grad(warmup_out, probs) + del warmup_out + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = fused_moe_aux_loss( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=total_num_tokens_dev, + num_experts=num_experts, + topk=topk, + coeff=coeff, + ) + (grad_probs,) = torch.autograd.grad(out, probs) + + atol, rtol = _get_tolerances(dtype, num_cols) + # Replay with several distinct token counts; the captured graph must pick + # up each new value through total_num_tokens_dev. + for new_total in (num_tokens, num_tokens // 2, num_tokens * 2 - 17): + total_num_tokens_dev.fill_(new_total) + g.replay() + torch.cuda.synchronize() + ref_probs = probs.detach().clone().requires_grad_(True) + ref = aux_loss_pytorch( + probs=ref_probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=new_total, + topk=topk, + num_experts=num_experts, + moe_aux_loss_coeff=coeff, + ) + (ref_grad_probs,) = torch.autograd.grad(ref, ref_probs) + torch.testing.assert_close(out, ref, atol=atol, rtol=rtol) + torch.testing.assert_close(grad_probs, ref_grad_probs, atol=atol, rtol=rtol) + + def _bytemap_to_bitmap_u8(bytemap: torch.Tensor) -> torch.Tensor: """Reference packer: bool[T, E] -> uint8[T, ceil(E/8)] LSB-first. 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 a80e90c0fd..9fb5adcc1f 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -132,6 +132,100 @@ void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_ex reinterpret_cast(Coeff_buf.data.dptr), stream););); } +/* ------------------------------------------------------------------------- + * CUDA-graph-safe variant: total_num_tokens lives in a 0-dim int64 device + * tensor whose value can change between graph replays. Each CTA's reduction + * lane computes C_coeff directly from the device value. The first CTA also + * writes C_coeff into Coeff_buf[0] for backward. + * ------------------------------------------------------------------------- */ +template +__global__ void fused_moe_aux_loss_forward_kernel_graph_safe( + const DataType* probs, const IndexType* tokens_per_expert, const int64_t* total_num_tokens_ptr, + 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) { + CompType col_sum = CompType(0); + for (int row = blockIdx.x; row < num_rows; row += gridDim.x) { + col_sum += CompType(probs[row * num_cols + col]); + } + col_sum *= CompType(tokens_per_expert[col]); + thread_sum += col_sum; + } + + extern __shared__ float shmem[]; + CompType* shmem_block = reinterpret_cast(shmem); + shmem_block[threadIdx.x] = thread_sum; + __syncthreads(); + + const int warp_id = threadIdx.x / kThreadsPerWarp; + const int lane_id = threadIdx.x % kThreadsPerWarp; + if (warp_id == 0) { + CompType block_sum = warp_reduce_on_shmem( + shmem_block, static_cast(blockDim.x), lane_id); + if (lane_id == 0) { + const float total_num_tokens = static_cast(*total_num_tokens_ptr); + const float C_coeff = (static_cast(num_experts) * coeff) / static_cast(topk) / + total_num_tokens / total_num_tokens; + if (blockIdx.x == 0) { + Coeff_buf[0] = C_coeff; + } + atomicAdd(&Coeff_buf[1], static_cast(block_sum * C_coeff)); + } + } +} + +template +void fused_moe_aux_loss_forward_kernel_launcher_graph_safe( + const DataType* probs, const IndexType* tokens_per_expert, const int64_t* total_num_tokens_dev, + int num_experts, int num_rows, int num_cols, int topk, float coeff, DataType* aux_loss, + 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(num_cols % num_experts == 0, "Number of input columns (", num_cols, + ") must be a multiple of number of experts (", num_experts, ")."); + + const int block_size = ((std::min(1024, num_cols) + static_cast(kThreadsPerWarp) - 1) / + static_cast(kThreadsPerWarp)) * + static_cast(kThreadsPerWarp); + const int grid_size = cuda::sm_count() * 2; + const size_t smem_size = block_size * sizeof(CompType); + check_shared_memory_capacity_num_experts(smem_size, num_cols); + + // Zero the float accumulator. The main kernel writes Coeff_buf[0] for backward. + NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream)); + fused_moe_aux_loss_forward_kernel_graph_safe + <<>>(probs, tokens_per_expert, total_num_tokens_dev, + num_experts, num_rows, num_cols, topk, coeff, + Coeff_buf); + NVTE_CHECK_CUDA(cudaGetLastError()); + + convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void fused_moe_aux_loss_forward_graph_safe(const Tensor& probs, const Tensor& tokens_per_expert, + const Tensor& total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff, + Tensor& aux_loss, Tensor& Coeff_buf, + cudaStream_t stream) { + NVTE_CHECK(total_num_tokens.data.dtype == DType::kInt64, + "total_num_tokens must be a 0-dim int64 tensor; got dtype ", + static_cast(total_num_tokens.data.dtype)); + NVTE_CHECK(total_num_tokens.numel() == 1, + "total_num_tokens must contain exactly one element; got ", total_num_tokens.numel()); + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( + probs.data.dtype, DataType, + TE_ROUTER_INDEX_TYPE_SWITCH_ALL( + tokens_per_expert.data.dtype, IndexType, + fused_moe_aux_loss_forward_kernel_launcher_graph_safe( + reinterpret_cast(probs.data.dptr), + reinterpret_cast(tokens_per_expert.data.dptr), + reinterpret_cast(total_num_tokens.data.dptr), num_experts, num_rows, + num_cols, topk, coeff, reinterpret_cast(aux_loss.data.dptr), + reinterpret_cast(Coeff_buf.data.dptr), stream););); +} + template __global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf, const IndexType* tokens_per_expert, int num_rows, @@ -195,6 +289,20 @@ void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor to *convertNVTETensorCheck(Coeff_buf), stream); } +void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs, + const NVTETensor tokens_per_expert, + const NVTETensor total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff, + NVTETensor aux_loss, NVTETensor Coeff_buf, + cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_moe_aux_loss_forward_graph_safe); + using namespace transformer_engine; + fused_router::fused_moe_aux_loss_forward_graph_safe( + *convertNVTETensorCheck(probs), *convertNVTETensorCheck(tokens_per_expert), + *convertNVTETensorCheck(total_num_tokens), num_experts, num_rows, num_cols, topk, coeff, + *convertNVTETensorCheck(aux_loss), *convertNVTETensorCheck(Coeff_buf), stream); +} + void nvte_fused_moe_aux_loss_backward(const NVTETensor Const_buf, const NVTETensor tokens_per_expert, int num_rows, int num_cols, NVTETensor grad_aux_loss, NVTETensor grad_probs, diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 08f347c616..03fec2eb2c 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -210,25 +210,44 @@ void nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_ou int num_experts, int topk, int score_function, NVTETensor grad_logits, cudaStream_t stream); -/*! \brief Forward pass for auxiliary loss. +/*! \brief Forward pass for auxiliary loss. Host-int total_num_tokens path: + * the coefficient is folded on the host and passed as a kernel argument. + * Prefer this path when total_num_tokens is statically known and the call + * is not captured into a CUDA Graph. * - * \param[in] probs Probabilities from the forward pass. + * \param[in] probs Probabilities from the forward pass. * \param[in] tokens_per_expert Number of tokens per expert. - * \param[in] total_num_tokens Number of total tokens. Will be used in seq/global aux loss. - * \param[in] num_experts Number of experts. - * \param[in] num_rows Number of rows of probs. - * \param[in] num_cols Number of columns of probs. - * \param[in] topk Topk value. - * \param[in] coeff Coefficient. - * \param[out] aux_loss Output GPU scalar for auxiliary loss. - * \param[out] Const_buf Output GPU scalar for temporary constant buffer for backward pass. - * \param[in] stream CUDA stream used for the operation. + * \param[in] total_num_tokens Number of total tokens. Used in seq/global aux loss. + * \param[in] num_experts Number of experts. + * \param[in] num_rows Number of rows of probs. + * \param[in] num_cols Number of columns of probs. + * \param[in] topk Topk value. + * \param[in] coeff Coefficient. + * \param[out] aux_loss Output GPU scalar for auxiliary loss. + * \param[out] Const_buf Output GPU scalar for temporary constant buffer for backward + * pass. + * \param[in] stream CUDA stream used for the operation. */ void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor tokens_per_expert, int total_num_tokens, int num_experts, int num_rows, int num_cols, int topk, float coeff, NVTETensor aux_loss, NVTETensor Const_buf, cudaStream_t stream); +/*! \brief Forward pass for auxiliary loss. Device-tensor total_num_tokens path: + * the coefficient is computed on device from a 0-dim int64 GPU tensor so its + * value stays dynamic across CUDA Graph replays. Prefer this path when the + * caller needs CUDA-graph-safe semantics with a dynamic token count. + * + * \param[in] total_num_tokens 0-dim int64 GPU tensor with the total token count. + * Other parameters as in :c:func:`nvte_fused_moe_aux_loss_forward`. + */ +void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs, + const NVTETensor tokens_per_expert, + const NVTETensor total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff, + NVTETensor aux_loss, NVTETensor Const_buf, + cudaStream_t stream); + /*! \brief Backward pass for auxiliary loss. * * \param[in] Const_buf Constant buffer from the forward pass. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8f42983553..785e438cda 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1708,6 +1708,29 @@ def backward(ctx, d_out, *_args): dq = dq[..., : d_out.shape[-1]] dk = dk[..., : d_out.shape[-1]] dv = dv[..., : d_out.shape[-1]] + # Zero-fill positions beyond cu_seqlens_*_padded[-1] in dQ/dK/dV for THD. + # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV. + # Sync-free `arange + mask` so capture and eager paths run the same code. + _qkv_format = ctx.qkv_layout.split("_")[0].replace("3", "").replace("2", "") + if _qkv_format == "thd": + if ( + cu_seqlens_q_padded is not None + and isinstance(dq, torch.Tensor) + and dq.shape[0] > 0 + ): + q_pad_mask = ( + torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1] + ) + dq[q_pad_mask] = 0 + if cu_seqlens_kv_padded is not None: + kv_actual_t = cu_seqlens_kv_padded[-1] + for d_tensor in (dk, dv): + if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0: + kv_pad_mask = ( + torch.arange(d_tensor.shape[0], device=d_tensor.device) + >= kv_actual_t + ) + d_tensor[kv_pad_mask] = 0 else: with get_nvtx_range_context("FusedAttnFunc.backward"): # get nominal data type of dq, dk, dv 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 61a46a8652..a62ca73187 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -2194,6 +2194,7 @@ def forward( ctx.S_quantizer.scale = S_quantizer.scale.clone() nvtx_range_pop(f"{nvtx_label}") + if return_max_logit: return out_ret, max_logit return out_ret @@ -2846,10 +2847,28 @@ def backward(ctx, dout, *_args): dim = ctx.qkv_format.index("s") dq, dk, dv = [x.view(*x.shape[:dim], -1, *x.shape[dim + 2 :]) for x in [dq, dk, dv]] - if ctx.qkv_format == "thd" and not ctx.use_fused_attention: - dq[cu_seqlens_q_padded[-1] :].fill_(0) - dk[cu_seqlens_kv_padded[-1] :].fill_(0) - dv[cu_seqlens_kv_padded[-1] :].fill_(0) + # Zero-fill dQ/dK/dV at positions beyond cu_seqlens_*_padded[-1]. + if ( + ctx.qkv_format == "thd" + and not ctx.use_fused_attention + and cu_seqlens_q_padded is not None + and cu_seqlens_kv_padded is not None + ): + if is_graph_capturing(): + # arange+mask under capture: `tensor[scalar_tensor:]` slicing would + # force a GPU->CPU sync that is forbidden during CUDA graph capture. + q_pad_mask = torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1] + kv_pad_mask = ( + torch.arange(dk.shape[0], device=dk.device) >= cu_seqlens_kv_padded[-1] + ) + dq[q_pad_mask] = 0 + dk[kv_pad_mask] = 0 + dv[kv_pad_mask] = 0 + else: + # Pre-existing TE eager-mode behaviour. + dq[cu_seqlens_q_padded[-1] :].fill_(0) + dk[cu_seqlens_kv_padded[-1] :].fill_(0) + dv[cu_seqlens_kv_padded[-1] :].fill_(0) if ctx.fp8 and ctx.is_input_fp8: dq, dk, dv, _, _ = combine_and_quantize(ctx.qkv_layout, dq, dk, dv, ctx.dQKV_quantizer) @@ -2904,6 +2923,23 @@ def backward(ctx, dout, *_args): nvtx_range_pop(f"{nvtx_label}") + # Zero-fill dQ/dK/dV at positions beyond the actual sequence end (THD CUDA Graph). + # cu_seqlens_*_padded are already local to this CP rank in the THD path. + # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV. + # Skip the corresponding zero-fill when its padded cu_seqlens is absent. + if ctx.qkv_format == "thd": + if cu_seqlens_q_padded is not None and isinstance(dq, torch.Tensor) and dq.shape[0] > 0: + q_pad_mask = torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1] + dq[q_pad_mask] = 0 + if cu_seqlens_kv_padded is not None: + kv_actual_t = cu_seqlens_kv_padded[-1] + for d_tensor in [dk, dv]: + if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0: + kv_pad_mask = ( + torch.arange(d_tensor.shape[0], device=d_tensor.device) >= kv_actual_t + ) + d_tensor[kv_pad_mask] = 0 + return ( None, dq, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 03008bb2d7..aa1481a384 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 @@ -1558,16 +1558,24 @@ def forward( False ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss, 111s} shapes" - # check if there is padding between sequences when qkv_format='thd' + # Default pad_between_seqs auto-detect. For THD, infer presence of + # inter-sequence padding from whether padded cu_seqlens were supplied -- + # sync-free, and stable across eager and CUDA graph capture (the auto-detect + # must return the same value in both modes for backend selection to match). + # If padded cu_seqlens are the *same object* as the unpadded ones, no real + # inter-sequence padding exists (only THD tail padding) -- treat as False so + # FlashAttention v2/v4 remain eligible. if pad_between_seqs is None: if qkv_format == "thd": - pad_between_seqs = ( - cu_seqlens_q_padded is not None - and not torch.equal(cu_seqlens_q_padded[:-1], cu_seqlens_q[:-1]) - ) or ( - cu_seqlens_kv_padded is not None - and not torch.equal(cu_seqlens_kv_padded[:-1], cu_seqlens_kv[:-1]) - ) + if ( + cu_seqlens_q_padded is cu_seqlens_q + and cu_seqlens_kv_padded is cu_seqlens_kv + ): + pad_between_seqs = False + else: + pad_between_seqs = ( + cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None + ) else: pad_between_seqs = False diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 69c73fe2fa..6edfbdc00e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -57,6 +57,10 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, int num_rows, int num_cols, int topk, float coeff); +std::tuple fused_moe_aux_loss_fwd_graph_safe( + at::Tensor probs, at::Tensor tokens_per_expert, at::Tensor total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff); + at::Tensor fused_moe_aux_loss_bwd(at::Tensor Const_buf, at::Tensor tokens_per_expert, int num_rows, int num_cols, at::Tensor grad_aux_loss); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 14272c9ac4..9c9ec36138 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -158,7 +158,11 @@ void init_router_bindings(pybind11::module &m) { m.def("fused_moe_aux_loss_fwd", &fused_moe_aux_loss_fwd, py::arg("probs"), py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"), py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"), - "Fused aux loss fwd"); + "Fused aux loss fwd (host-int total_num_tokens, host-folded C_coeff)"); + m.def("fused_moe_aux_loss_fwd_graph_safe", &fused_moe_aux_loss_fwd_graph_safe, py::arg("probs"), + py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"), + py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"), + "Fused aux loss fwd (device-tensor total_num_tokens, CUDA-graph-safe)"); m.def("fused_moe_aux_loss_bwd", &fused_moe_aux_loss_bwd, py::arg("Const_buf"), py::arg("tokens_per_expert"), py::arg("num_rows"), py::arg("num_cols"), py::arg("grad_aux_loss"), "Fused aux loss bwd"); diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 70762e3729..d59d6bc415 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -321,6 +321,37 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, return std::make_tuple(aux_loss, Const_buf); } +std::tuple fused_moe_aux_loss_fwd_graph_safe( + at::Tensor probs, at::Tensor tokens_per_expert, at::Tensor total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff) { + TORCH_CHECK(topk > 0, "topk must be greater than 0"); + TORCH_CHECK(num_experts > 0, "num_experts must be greater than 0"); + // Device-tensor path: keep total_num_tokens dynamic across CUDA Graph replays. + // Validate shape and dtype on the host; do not read its value (avoids a sync). + TORCH_CHECK(total_num_tokens.is_cuda(), "total_num_tokens must be a CUDA tensor"); + TORCH_CHECK(total_num_tokens.numel() == 1, + "total_num_tokens must contain exactly one element; got ", total_num_tokens.numel()); + TORCH_CHECK(total_num_tokens.scalar_type() == at::kLong, "total_num_tokens must be int64; got ", + total_num_tokens.scalar_type()); + + // Create the output tensor + at::Tensor aux_loss = at::empty({}, at::dtype(probs.scalar_type()).device(at::kCUDA)); + at::Tensor Const_buf = at::empty({2}, at::dtype(at::kFloat).device(at::kCUDA)); + + auto probs_cu = makeTransformerEngineTensor(probs); + auto tokens_per_expert_cu = makeTransformerEngineTensor(tokens_per_expert); + auto total_num_tokens_cu = makeTransformerEngineTensor(total_num_tokens); + auto aux_loss_cu = makeTransformerEngineTensor(aux_loss); + auto Const_buf_cu = makeTransformerEngineTensor(Const_buf); + + nvte_fused_moe_aux_loss_forward_graph_safe(probs_cu.data(), tokens_per_expert_cu.data(), + total_num_tokens_cu.data(), num_experts, num_rows, + num_cols, topk, coeff, aux_loss_cu.data(), + Const_buf_cu.data(), at::cuda::getCurrentCUDAStream()); + + return std::make_tuple(aux_loss, Const_buf); +} + at::Tensor fused_moe_aux_loss_bwd(at::Tensor Const_buf, at::Tensor tokens_per_expert, int num_rows, int num_cols, at::Tensor grad_aux_loss) { // Create the output tensor diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 519d06ce23..51350ab2fd 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -280,7 +280,9 @@ def fused_compute_score_for_moe_aux_loss( class FusedAuxLoss(torch.autograd.Function): """ - Fused MoE aux loss. + Fused MoE aux loss. ``total_num_tokens`` may be either a Python int + (host-folded coefficient, original fast path) or a 0-dim int64 CUDA + tensor (device-folded coefficient, CUDA-graph-safe path). """ @staticmethod @@ -288,7 +290,7 @@ def forward( ctx, probs: torch.Tensor, tokens_per_expert: torch.Tensor, - total_num_tokens: int, + total_num_tokens: Union[int, torch.Tensor], num_experts: int, topk: int, coeff: float, @@ -296,16 +298,28 @@ def forward( # pylint: disable=missing-function-docstring num_rows = probs.size(0) num_cols = probs.size(1) - aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd( - probs=probs, - tokens_per_expert=tokens_per_expert, - total_num_tokens=total_num_tokens, - num_experts=num_experts, - num_rows=num_rows, - num_cols=num_cols, - topk=topk, - coeff=coeff, - ) + if isinstance(total_num_tokens, torch.Tensor): + aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd_graph_safe( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=total_num_tokens, + num_experts=num_experts, + num_rows=num_rows, + num_cols=num_cols, + topk=topk, + coeff=coeff, + ) + else: + aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=int(total_num_tokens), + num_experts=num_experts, + num_rows=num_rows, + num_cols=num_cols, + topk=topk, + coeff=coeff, + ) ctx.save_for_backward(Const_buf, tokens_per_expert) ctx.num_rows = num_rows ctx.num_cols = num_cols @@ -328,7 +342,7 @@ def backward(ctx, grad_aux_loss): def fused_moe_aux_loss( probs: torch.Tensor, tokens_per_expert: torch.Tensor, - total_num_tokens: int, + total_num_tokens: Union[int, torch.Tensor], num_experts: int, topk: int, coeff: float, @@ -340,8 +354,12 @@ def fused_moe_aux_loss( probs : torch.Tensor in fp32/bf16/fp16 tokens_per_expert : torch.Tensor in int32/int64/fp32/bf16 the number of tokens per expert. - total_num_tokens : int - the total number of tokens used in the aux loss calculation. + total_num_tokens : int or 0-dim int64 CUDA torch.Tensor + the total number of tokens used in the aux loss calculation. Pass a + Python int for the fastest path (coefficient folded on the host). + Pass a 0-dim int64 CUDA tensor when the call is captured into a + CUDA Graph and the value must stay dynamic across replays; the + coefficient is computed on device by the main reduction kernel. num_experts : int topk : int coeff : float From 9f2074e4bccaa3a8f386ad053b9b74f983238c5e Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Wed, 1 Jul 2026 12:04:13 -0500 Subject: [PATCH 27/42] [Common/PyTorch] Grouped-quantize kernels for 1D and 2D FP8 block-scaling (#3135) Implements grouped-tensor quantize for the FP8 1D (1x128) and 2D (128x128) block-scaling recipes in row-wise (RW), column-wise (CW) and BOTH quantization directions. A single CUDA kernel launch walks 128x128 tiles across every tensor in the group, with each CTA decoding its owning tensor from the device-side GroupedTensor metadata with (N, R, K) shapes. Supports SAME_BOTH_DIMS (all tensors identical) and VARYING_FIRST_DIM (constant K, varying R) shape representations. Three kernels share the dispatcher in group_quantize_blockwise_{1d,2d}: - group_block_scaled_1d_rw_kernel: RW-only dispatch; 8 threads/row, reads global memory directly into vec-16 registers; bypasses TMA since the shared-memory roundtrip and ptx::mbarrier do not buy anything without re-use in the CW path. - group_block_scaled_1d_tma_kernel: CW-only and BOTH dispatch. TMA bulk-load fills shared memory input cache. BOTH runs an RW pass (8 threads/row, vec-16 read from shared memory) then a CW pass; CW-only skips the RW pass. The CW pass uses 4 t/col with 32-row reg_data and two column passes in the BOTH instantiation (keeps the per-thread register footprint under the sm_90 3-CTAs/SM threshold) and 2 t/col with 64-row reg_data in the CW-only instantiation (avoids doubling the smem-load bank-conflict footprint that 4 t/col would introduce). - group_block_scaled_2d_tma_kernel: RW-only, CW-only and BOTH dispatch. TMA bulk-load fills shared memory input cache. Pass 1 stages 8 IVecs/thread in registers while computing the per-tile scalar amax. Pass 2 quantizes from registers, emits row-wise output, stages column-wise output to the shared memory transpose staging buffer, then drains smem_T to global memory. Per-expert scale offsets: - 1D RW: closed-form O(1) for both SAME_BOTH_DIMS and VARYING_FIRST_DIM (each M_i is a multiple of kTileDim=128, hence of kScaleColAlign=4, so DIVUP_TO_MULTIPLE collapses and the prefix sum reduces to a single tensor_offsets_ptr[tensor_id]/K load). - 2D CW: closed-form O(1) for SAME_BOTH_DIMS; CTA-cooperative warp-shuffle prefix sum for VARYING_FIRST_DIM (non-linear DIVUP_TO_MULTIPLE on blocks_y_t prevents a closed form). The cooperative reduction uses the existing warp_allreduce_sum helper from common/utils.cuh. Dequantize and bias-gradient (bgrad): - group_dequantize_fp8_blockwise.cuh: kernels for all four modes (1D/2D x rowwise/columnwise), inverting the per-expert layouts the quantize kernels write. - bgrad_group_quantize accepts Float8Block quantizers and computes dbias per-tile column-partial in-kernel (mirroring MXFP8); reduced per expert via the existing common::grouped_reduce_dbias. Scale constraints: the fused grouped FP8BS path supports only unconstrained FP32 scales (Float8BlockQuantizer::create_grouped_tensor rejects force_pow_2_scales=True). Power-of-2 scales remain available on the non-grouped/unfused split-quantize path used for Blackwell MXFP8 emulation. Tests: existing parametrized grouped quantize / dequantize / bgrad tests in test_grouped_tensor.py cover MXFP8, NVFP4, FP8 current scaling and the newly-added FP8 block scaling recipe. tests/cpp/operator/ test_cast_float8blockwise_grouped.cu adds 72 C++ unit-test cases over uniform/jagged shapes, all four (BD x direction) modes, K in {128, 256, 512}, and CUDA-graph capture coverage. Kernels are gated to Hopper (sm_90) at the host dispatcher (cuBlasLt grouped GEMM supports FP8 block-scaling only on Hopper). JAX integration is intentionally left out of scope and deferred to a follow-up PR. Resolves #2525 Signed-off-by: Alp Dener --- tests/cpp/operator/CMakeLists.txt | 2 + .../test_cast_float8blockwise_grouped.cu | 416 +++++++ ...test_dequantize_float8blockwise_grouped.cu | 317 +++++ tests/cpp/operator/test_grouped_gemm.cu | 291 +++++ tests/pytorch/test_grouped_tensor.py | 310 +++++ .../common/cast/dispatch/dequantize.cuh | 6 + .../common/cast/dispatch/quantize.cuh | 43 + .../group_dequantize_fp8_blockwise.cuh | 519 +++++++++ .../group_quantize_fp8_blockwise.cuh | 1022 +++++++++++++++++ transformer_engine/common/util/ptx.cuh | 48 +- .../pytorch/csrc/extensions/cast.cpp | 37 +- transformer_engine/pytorch/csrc/quantizer.cpp | 42 +- 12 files changed, 3024 insertions(+), 29 deletions(-) create mode 100644 tests/cpp/operator/test_cast_float8blockwise_grouped.cu create mode 100644 tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu create mode 100644 transformer_engine/common/cast/fp8_blockwise/group_dequantize_fp8_blockwise.cuh create mode 100644 transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 1c4d86a3a8..832177c637 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -15,8 +15,10 @@ add_executable(test_operator test_cast_mxfp8_grouped.cu test_cast_nvfp4_transpose.cu test_cast_float8blockwise.cu + test_cast_float8blockwise_grouped.cu test_dequantize_mxfp8.cu test_dequantize_mxfp8_grouped.cu + test_dequantize_float8blockwise_grouped.cu test_dequantize_nvfp4.cu test_transpose.cu test_cast_transpose.cu diff --git a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu new file mode 100644 index 0000000000..c7d7a73475 --- /dev/null +++ b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu @@ -0,0 +1,416 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "../test_common.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum class ShapeRep { SAME_BOTH_DIMS = 0, VARYING_FIRST_DIM = 1 }; +enum class ScalingDir { ROWWISE = 0, COLWISE = 1, BOTH = 2 }; +enum class BlockDim { ONE_D = 1, TWO_D = 2 }; + +constexpr size_t kBlock = 128; + +inline size_t align4(size_t x) { return ((x + 3) / 4) * 4; } + +// Configure split-quantize reference: call non-grouped nvte_quantize_v2 on each tensor slice. +// Returns flat host buffers for per-tensor outputs and scales (in their per-tensor natural +// layout) so the test can index them and compare element-wise against the grouped layout. +struct PerTensorRef { + std::vector> output; // per tensor, FP8 raw bytes (R_t * K) + std::vector> output_t; // per tensor, FP8 raw bytes (K * R_t) + std::vector> scale_inv; // per tensor, layout per non-grouped impl + std::vector> scale_inv_t; // per tensor, layout per non-grouped impl +}; + +// Per-expert scale layout helpers mirroring the kernel + cuBLAS grouped GEMM +// expectation. Each expert's scales occupy a contiguous sub-block of the global +// scale buffer; these compute per-expert padded sizes (in floats) so the test +// can both size the buffer and compute per-expert base offsets. +// 1D rowwise : blocks_X * roundup(M_t, 4) +// 1D colwise : blocks_y_t * roundup(K, 4) +// 2D rowwise : blocks_y_t * roundup(blocks_X, 4) +// 2D colwise : blocks_X * roundup(blocks_y_t, 4) +inline size_t per_expert_scale_floats(BlockDim block_dim, bool columnwise, size_t M_t, size_t K) { + constexpr size_t kBlk = 128; + const size_t blocks_X = (K + kBlk - 1) / kBlk; + const size_t blocks_y = (M_t + kBlk - 1) / kBlk; + if (block_dim == BlockDim::ONE_D) { + if (!columnwise) return blocks_X * align4(M_t); + return blocks_y * align4(K); + } + // 2D + if (!columnwise) return blocks_y * align4(blocks_X); + return blocks_X * align4(blocks_y); +} + +// Cumulative per-expert offset (in floats) for tensor `t`. +inline size_t per_expert_scale_offset(const std::vector& first_dims, size_t t, + BlockDim block_dim, bool columnwise, size_t K) { + size_t offset = 0; + for (size_t i = 0; i < t; ++i) { + offset += per_expert_scale_floats(block_dim, columnwise, first_dims[i], K); + } + return offset; +} + +template +void perform_test(ShapeRep shape_rep, BlockDim block_dim, ScalingDir dir, + const std::vector& first_dims_h, size_t K, + bool force_pow_2_scales, float epsilon) { + if (getDeviceComputeCapability() < hopperComputeCapability) { + GTEST_SKIP(); + } + + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + const size_t num_tensors = first_dims_h.size(); + size_t R_total = 0; + for (size_t m : first_dims_h) { + ASSERT_EQ(m % kBlock, 0u) << "Per-tensor first dim must be multiple of 128"; + R_total += m; + } + ASSERT_EQ(K % 16u, 0u); + + // Host data + std::mt19937 gen(0xC0FFEEu); + std::uniform_real_distribution dist(-2.0f, 1.0f); + std::vector input_h(R_total * K); + for (auto& v : input_h) v = static_cast(dist(gen)); + + // Tensor offsets (element offsets) + std::vector offsets_h(num_tensors + 1, 0); + for (size_t t = 0; t < num_tensors; ++t) { + offsets_h[t + 1] = offsets_h[t] + static_cast(first_dims_h[t] * K); + } + std::vector first_dims_i64(num_tensors); + for (size_t t = 0; t < num_tensors; ++t) first_dims_i64[t] = static_cast(first_dims_h[t]); + + const bool use_rowwise = (dir == ScalingDir::ROWWISE || dir == ScalingDir::BOTH); + const bool use_colwise = (dir == ScalingDir::COLWISE || dir == ScalingDir::BOTH); + + const NVTEScalingMode mode = + (block_dim == BlockDim::ONE_D) ? NVTE_BLOCK_SCALING_1D : NVTE_BLOCK_SCALING_2D; + + // Allocate grouped device buffers. + InputType* input_d = nullptr; + OutputType* output_d = nullptr; + OutputType* output_t_d = nullptr; + float* scale_inv_d = nullptr; + float* scale_inv_t_d = nullptr; + int64_t* offsets_d = nullptr; + int64_t* first_dims_d = nullptr; + + const size_t blocks_X = (K + kBlock - 1) / kBlock; + + // Grouped scale buffers are sized as the sum of per-expert padded sub-blocks, + // matching cuBLAS grouped FP8 block-scaling GEMM's per-expert layout. + size_t scale_inv_elems = 0; + size_t scale_inv_t_elems = 0; + for (size_t t = 0; t < num_tensors; ++t) { + scale_inv_elems += per_expert_scale_floats(block_dim, /*columnwise=*/false, first_dims_h[t], K); + scale_inv_t_elems += + per_expert_scale_floats(block_dim, /*columnwise=*/true, first_dims_h[t], K); + } + std::vector scale_inv_shape = {scale_inv_elems}; + std::vector scale_inv_t_shape = {scale_inv_t_elems}; + + const size_t input_bytes = R_total * K * sizeof(InputType); + const size_t output_bytes = R_total * K * sizeof(OutputType); + + cudaMalloc(&input_d, input_bytes); + cudaMemcpy(input_d, input_h.data(), input_bytes, cudaMemcpyHostToDevice); + cudaMalloc(&offsets_d, (num_tensors + 1) * sizeof(int64_t)); + cudaMemcpy(offsets_d, offsets_h.data(), (num_tensors + 1) * sizeof(int64_t), + cudaMemcpyHostToDevice); + if (shape_rep == ShapeRep::VARYING_FIRST_DIM) { + cudaMalloc(&first_dims_d, num_tensors * sizeof(int64_t)); + cudaMemcpy(first_dims_d, first_dims_i64.data(), num_tensors * sizeof(int64_t), + cudaMemcpyHostToDevice); + } + if (use_rowwise) { + cudaMalloc(&output_d, output_bytes); + cudaMemset(output_d, 0, output_bytes); + cudaMalloc(&scale_inv_d, scale_inv_elems * sizeof(float)); + cudaMemset(scale_inv_d, 0, scale_inv_elems * sizeof(float)); + } + if (use_colwise) { + cudaMalloc(&output_t_d, output_bytes); + cudaMemset(output_t_d, 0, output_bytes); + cudaMalloc(&scale_inv_t_d, scale_inv_t_elems * sizeof(float)); + cudaMemset(scale_inv_t_d, 0, scale_inv_t_elems * sizeof(float)); + } + + // Build grouped tensors. + std::vector logical_shape_vec = {R_total, K}; + NVTEShape logical_shape = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size()); + + NVTEGroupedTensor in_gt = nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, + logical_shape); + NVTEGroupedTensor out_gt = nvte_create_grouped_tensor(mode, num_tensors, logical_shape); + + NVTEBasicTensor in_data = {input_d, static_cast(itype), logical_shape}; + nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedRowwiseData, &in_data, sizeof(in_data)); + + NVTEShape offsets_shape; + offsets_shape.ndim = 1; + offsets_shape.data[0] = num_tensors + 1; + NVTEBasicTensor offsets_bt = {offsets_d, kNVTEInt64, offsets_shape}; + if (shape_rep == ShapeRep::VARYING_FIRST_DIM) { + NVTEShape first_dims_shape; + first_dims_shape.ndim = 1; + first_dims_shape.data[0] = num_tensors; + NVTEBasicTensor first_dims_bt = {first_dims_d, kNVTEInt64, first_dims_shape}; + nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedFirstDims, &first_dims_bt, + sizeof(first_dims_bt)); + nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedFirstDims, &first_dims_bt, + sizeof(first_dims_bt)); + nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedTensorOffsets, &offsets_bt, + sizeof(offsets_bt)); + nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedTensorOffsets, &offsets_bt, + sizeof(offsets_bt)); + } + + if (use_rowwise) { + NVTEBasicTensor out_data = {output_d, static_cast(otype), logical_shape}; + NVTEShape scale_inv_shape_nv = nvte_make_shape(scale_inv_shape.data(), scale_inv_shape.size()); + NVTEBasicTensor scale_bt = {scale_inv_d, kNVTEFloat32, scale_inv_shape_nv}; + nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseData, &out_data, sizeof(out_data)); + nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseScaleInv, &scale_bt, sizeof(scale_bt)); + } + if (use_colwise) { + NVTEBasicTensor out_t_data = {output_t_d, static_cast(otype), logical_shape}; + NVTEShape scale_inv_t_shape_nv = nvte_make_shape(scale_inv_t_shape.data(), + scale_inv_t_shape.size()); + NVTEBasicTensor scale_t_bt = {scale_inv_t_d, kNVTEFloat32, scale_inv_t_shape_nv}; + nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedColumnwiseData, &out_t_data, + sizeof(out_t_data)); + nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedColumnwiseScaleInv, &scale_t_bt, + sizeof(scale_t_bt)); + } + + // Run grouped quantize. + QuantizationConfigWrapper quant_config; + quant_config.set_force_pow_2_scales(force_pow_2_scales); + quant_config.set_amax_epsilon(epsilon); + nvte_group_quantize(in_gt, out_gt, quant_config, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + // Pull grouped outputs back to host. + std::vector output_h(use_rowwise ? R_total * K : 0); + std::vector output_t_h(use_colwise ? R_total * K : 0); + std::vector scale_inv_h(use_rowwise ? scale_inv_elems : 0); + std::vector scale_inv_t_h(use_colwise ? scale_inv_t_elems : 0); + if (use_rowwise) { + cudaMemcpy(output_h.data(), output_d, R_total * K, cudaMemcpyDeviceToHost); + cudaMemcpy(scale_inv_h.data(), scale_inv_d, scale_inv_elems * sizeof(float), + cudaMemcpyDeviceToHost); + } + if (use_colwise) { + cudaMemcpy(output_t_h.data(), output_t_d, R_total * K, cudaMemcpyDeviceToHost); + cudaMemcpy(scale_inv_t_h.data(), scale_inv_t_d, scale_inv_t_elems * sizeof(float), + cudaMemcpyDeviceToHost); + } + + // Run split-quantize reference per tensor and compare element-wise. + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t row_offset = static_cast(offsets_h[t]) / K; + + std::vector tshape = {M, K}; + Tensor ref_in("ref_in_" + std::to_string(t), tshape, itype); + // The non-grouped 2D kernel requires rowwise output to be allocated even when only colwise + // data is consumed. We always allocate both and compare only what the grouped kernel produced. + const bool ref_rowwise = (block_dim == BlockDim::TWO_D) ? true : use_rowwise; + const bool ref_colwise = use_colwise; + Tensor ref_out("ref_out_" + std::to_string(t), tshape, otype, ref_rowwise, ref_colwise, mode); + + // Copy this tensor's input slice into ref_in. + { + auto* dst = ref_in.rowwise_dptr(); + const InputType* src = reinterpret_cast(input_d) + row_offset * K; + cudaMemcpy(dst, src, M * K * sizeof(InputType), cudaMemcpyDeviceToDevice); + } + + QuantizationConfigWrapper qc; + qc.set_force_pow_2_scales(force_pow_2_scales); + qc.set_amax_epsilon(epsilon); + nvte_quantize_v2(ref_in.data(), ref_out.data(), qc, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + ref_out.to_cpu(); // sync output and scale_inv buffers from GPU to CPU + + // Compare data. + if (use_rowwise) { + const OutputType* ref_data = ref_out.rowwise_cpu_dptr(); + for (size_t r = 0; r < M; ++r) { + for (size_t c = 0; c < K; ++c) { + const uint8_t got = output_h[(row_offset + r) * K + c]; + const uint8_t exp = reinterpret_cast(ref_data)[r * K + c]; + ASSERT_EQ(got, exp) << "rowwise data mismatch t=" << t << " r=" << r << " c=" << c; + } + } + } + if (use_colwise) { + const OutputType* ref_data_t = ref_out.columnwise_cpu_dptr(); + // Per-expert columnwise data: contiguous (K, M_t) block at element offset + // K * row_offset, matching cuBLAS grouped GEMM's per-expert data pointer. + const size_t expert_data_off = static_cast(row_offset) * K; + for (size_t c = 0; c < K; ++c) { + for (size_t r = 0; r < M; ++r) { + const uint8_t got = output_t_h[expert_data_off + c * M + r]; + const uint8_t exp = reinterpret_cast(ref_data_t)[c * M + r]; + ASSERT_EQ(got, exp) << "colwise data mismatch t=" << t << " c=" << c << " r=" << r; + } + } + } + + // Compare scales. Per-expert layout: each expert's scales live in a + // contiguous sub-block at per_expert_scale_offset(...). + if (block_dim == BlockDim::ONE_D) { + const size_t M_pad = align4(M); + const size_t K_pad = align4(K); + const size_t blocks_y_per_tensor = M / kBlock; + if (use_rowwise) { + const float* ref_sc = ref_out.rowwise_cpu_scale_inv_ptr(); + // Per-expert RW: (blocks_X, roundup(M_t, 4)). + const size_t expert_off = + per_expert_scale_offset(first_dims_h, t, block_dim, false, K); + for (size_t bx = 0; bx < blocks_X; ++bx) { + for (size_t r = 0; r < M; ++r) { + const float got = scale_inv_h[expert_off + bx * M_pad + r]; + const float exp = ref_sc[bx * M_pad + r]; + ASSERT_EQ(got, exp) << "1D rowwise scale mismatch t=" << t << " bx=" << bx + << " r=" << r; + } + } + } + if (use_colwise) { + const float* ref_sct = ref_out.columnwise_cpu_scale_inv_ptr(); + // Per-expert CW: (blocks_y_t, roundup(K, 4)) compact. + const size_t expert_off = per_expert_scale_offset(first_dims_h, t, block_dim, true, K); + for (size_t by = 0; by < blocks_y_per_tensor; ++by) { + for (size_t c = 0; c < K; ++c) { + const float got = scale_inv_t_h[expert_off + by * K_pad + c]; + const float exp = ref_sct[by * K_pad + c]; + ASSERT_EQ(got, exp) << "1D colwise scale mismatch t=" << t << " by=" << by + << " c=" << c; + } + } + } + } else { + // 2D per-expert: rowwise (blocks_y_t, roundup(blocks_X, 4)); colwise + // (blocks_X, roundup(blocks_y_t, 4)) compact. + const size_t blocks_y_per_tensor = M / kBlock; + const size_t bx_pad = align4(blocks_X); + const size_t by_pad_t = align4(blocks_y_per_tensor); + if (use_rowwise) { + const float* ref_sc = ref_out.rowwise_cpu_scale_inv_ptr(); + const size_t expert_off = + per_expert_scale_offset(first_dims_h, t, block_dim, false, K); + for (size_t by = 0; by < blocks_y_per_tensor; ++by) { + for (size_t bx = 0; bx < blocks_X; ++bx) { + const float got = scale_inv_h[expert_off + by * bx_pad + bx]; + const float exp = ref_sc[by * bx_pad + bx]; + ASSERT_EQ(got, exp) << "2D rowwise scale mismatch t=" << t << " by=" << by + << " bx=" << bx; + } + } + } + if (use_colwise) { + const float* ref_sct = ref_out.columnwise_cpu_scale_inv_ptr(); + const size_t expert_off = per_expert_scale_offset(first_dims_h, t, block_dim, true, K); + for (size_t bx = 0; bx < blocks_X; ++bx) { + for (size_t by = 0; by < blocks_y_per_tensor; ++by) { + const float got = scale_inv_t_h[expert_off + bx * by_pad_t + by]; + const float exp = ref_sct[bx * by_pad_t + by]; + ASSERT_EQ(got, exp) << "2D colwise scale mismatch t=" << t << " bx=" << bx + << " by=" << by; + } + } + } + } + } + + nvte_destroy_grouped_tensor(in_gt); + nvte_destroy_grouped_tensor(out_gt); + cudaFree(input_d); + if (output_d) cudaFree(output_d); + if (output_t_d) cudaFree(output_t_d); + if (scale_inv_d) cudaFree(scale_inv_d); + if (scale_inv_t_d) cudaFree(scale_inv_t_d); + cudaFree(offsets_d); + if (first_dims_d) cudaFree(first_dims_d); +} + +struct TestConfig { + ShapeRep shape_rep; + BlockDim block_dim; + ScalingDir dir; + std::vector first_dims; + size_t K; +}; + +class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam {}; + +TEST_P(GroupedFP8BlockwiseTestSuite, Test) { + const TestConfig& cfg = GetParam(); + perform_test(cfg.shape_rep, cfg.block_dim, cfg.dir, cfg.first_dims, cfg.K, + /*force_pow_2_scales=*/false, /*epsilon=*/0.0f); +} + +std::vector make_configs() { + std::vector configs; + std::vector> uniform = {{128, 128}, {256, 256, 256, 256}}; + std::vector> jagged = { + {128, 256, 384, 512}, {256, 128, 512, 384, 1024}}; + std::vector Ks = {128, 256, 512}; + for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) { + for (auto dir : {ScalingDir::ROWWISE, ScalingDir::COLWISE, ScalingDir::BOTH}) { + for (size_t K : Ks) { + for (const auto& v : uniform) { + configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K}); + } + for (const auto& v : jagged) { + configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K}); + } + } + } + } + return configs; +} + +std::string make_name(const ::testing::TestParamInfo& info) { + const auto& c = info.param; + std::string s = (c.shape_rep == ShapeRep::SAME_BOTH_DIMS ? "SAME" : "VARYFIRST"); + s += "_BD" + std::to_string(static_cast(c.block_dim)); + s += (c.dir == ScalingDir::ROWWISE ? "_RW" + : c.dir == ScalingDir::COLWISE ? "_CW" : "_BOTH"); + s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size()); + s += "_M"; + for (size_t m : c.first_dims) s += "_" + std::to_string(m); + return s; +} + +INSTANTIATE_TEST_SUITE_P(GroupedFP8Blockwise, GroupedFP8BlockwiseTestSuite, + ::testing::ValuesIn(make_configs()), make_name); + +} // namespace diff --git a/tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu b/tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu new file mode 100644 index 0000000000..48b0a54c9b --- /dev/null +++ b/tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu @@ -0,0 +1,317 @@ +/************************************************************************* + * 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" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum class ShapeRep { SAME_BOTH_DIMS = 0, VARYING_FIRST_DIM = 1 }; +enum class ScalingDir { ROWWISE = 0, COLWISE = 1 }; +enum class BlockDim { ONE_D = 1, TWO_D = 2 }; + +constexpr size_t kBlock = 128; + +inline size_t align4(size_t x) { return ((x + 3) / 4) * 4; } + +// Per-expert padded scale size (in floats), matching the grouped FP8 block-scaling layout that +// the grouped quantize kernel writes and cuBLAS grouped GEMM consumes: +// 1D rowwise : blocks_X * roundup(M_t, 4) (scale shape {blocks_X, roundup(M_t, 4)}) +// 1D colwise : blocks_y_t * roundup(K, 4) (scale shape {blocks_y_t, roundup(K, 4)}) +// 2D rowwise : blocks_y_t * roundup(blocks_X,4) (scale shape {blocks_y_t, roundup(blocks_X, 4)}) +// 2D colwise : blocks_X * roundup(blocks_y_t,4)(scale shape {blocks_X, roundup(blocks_y_t,4)}) +inline void per_expert_scale_shape(BlockDim block_dim, bool columnwise, size_t M_t, size_t K, + size_t& scale_y, size_t& scale_x) { + const size_t blocks_X = (K + kBlock - 1) / kBlock; + const size_t blocks_y = (M_t + kBlock - 1) / kBlock; + if (block_dim == BlockDim::ONE_D) { + if (!columnwise) { + scale_y = blocks_X; + scale_x = align4(M_t); + } else { + scale_y = blocks_y; + scale_x = align4(K); + } + } else { + if (!columnwise) { + scale_y = blocks_y; + scale_x = align4(blocks_X); + } else { + scale_y = blocks_X; + scale_x = align4(blocks_y); + } + } +} + +inline size_t per_expert_scale_floats(BlockDim block_dim, bool columnwise, size_t M_t, size_t K) { + size_t y, x; + per_expert_scale_shape(block_dim, columnwise, M_t, K, y, x); + return y * x; +} + +// Grouped FP8 block-scaling dequantize test. +// +// Methodology mirrors test_dequantize_mxfp8_grouped.cu: the grouped dequantize kernel is +// validated against single-tensor nvte_dequantize called in a loop for each tensor; results must +// be bitwise identical. We generate random FP8 data and random FP32 scales laid out in the +// grouped per-expert format, run nvte_group_dequantize, and for each expert slice out its data + +// scale sub-block, feed it to a per-tensor (direction-only) dequantize, and compare. +template +void performTest(ShapeRep shape_rep, BlockDim block_dim, bool rowwise, + const std::vector& first_dims_h, size_t K) { + // FP8 block-scaling grouped kernels are Hopper-only (SM90-SM99). + if (getDeviceComputeCapability() < hopperComputeCapability || + getDeviceComputeCapability() >= blackwellComputeCapability) { + GTEST_SKIP(); + } + + const DType itype = TypeInfo::dtype; + const DType otype = TypeInfo::dtype; + const bool columnwise = !rowwise; + + const size_t num_tensors = first_dims_h.size(); + size_t R_total = 0; + for (size_t m : first_dims_h) { + ASSERT_EQ(m % kBlock, 0u) << "Per-tensor first dim must be a multiple of 128"; + R_total += m; + } + ASSERT_EQ(K % 16u, 0u); + + const NVTEScalingMode mode = + (block_dim == BlockDim::ONE_D) ? NVTE_BLOCK_SCALING_1D : NVTE_BLOCK_SCALING_2D; + + // Element offsets (both data and, for columnwise, the transposed (K, M_t) block are contiguous + // per expert at element offset row_offset * K). + std::vector offsets_h(num_tensors + 1, 0); + for (size_t t = 0; t < num_tensors; ++t) + offsets_h[t + 1] = offsets_h[t] + static_cast(first_dims_h[t] * K); + std::vector first_dims_i64(num_tensors); + for (size_t t = 0; t < num_tensors; ++t) + first_dims_i64[t] = static_cast(first_dims_h[t]); + + // Per-expert scale sub-block offsets (in floats). + std::vector scale_off(num_tensors + 1, 0); + for (size_t t = 0; t < num_tensors; ++t) + scale_off[t + 1] = + scale_off[t] + per_expert_scale_floats(block_dim, columnwise, first_dims_h[t], K); + const size_t total_scales = scale_off[num_tensors]; + + // ---- Random FP8 data (valid normals) + random FP32 scales ---- + std::mt19937 gen(0xD3C0DEu); + const double minAbs = Numeric_Traits::minNorm; + const double maxAbs = Numeric_Traits::maxNorm; + std::uniform_real_distribution<> dis(minAbs, maxAbs); + std::uniform_real_distribution<> dis_sign(-1.0, 1.0); + std::uniform_real_distribution scale_dis(0.25f, 4.0f); + + std::vector data_h(R_total * K); + for (auto& v : data_h) { + double val = dis(gen); + if (dis_sign(gen) < 0.0) val = -val; + v = static_cast(val); + } + std::vector scales_h(total_scales); + for (auto& s : scales_h) s = scale_dis(gen); + + // ---- Device buffers ---- + InputType* data_d = nullptr; + float* scales_d = nullptr; + OutputType* out_grouped_d = nullptr; + int64_t* offsets_d = nullptr; + int64_t* first_dims_d = nullptr; + + cudaMalloc(&data_d, R_total * K * sizeof(InputType)); + cudaMemcpy(data_d, data_h.data(), R_total * K * sizeof(InputType), cudaMemcpyHostToDevice); + cudaMalloc(&scales_d, total_scales * sizeof(float)); + cudaMemcpy(scales_d, scales_h.data(), total_scales * sizeof(float), cudaMemcpyHostToDevice); + cudaMalloc(&out_grouped_d, R_total * K * sizeof(OutputType)); + cudaMemset(out_grouped_d, 0, R_total * K * sizeof(OutputType)); + cudaMalloc(&offsets_d, (num_tensors + 1) * sizeof(int64_t)); + cudaMemcpy(offsets_d, offsets_h.data(), (num_tensors + 1) * sizeof(int64_t), + cudaMemcpyHostToDevice); + if (shape_rep == ShapeRep::VARYING_FIRST_DIM) { + cudaMalloc(&first_dims_d, num_tensors * sizeof(int64_t)); + cudaMemcpy(first_dims_d, first_dims_i64.data(), num_tensors * sizeof(int64_t), + cudaMemcpyHostToDevice); + } + + // ---- Build grouped input (quantized) + output (high precision) tensors ---- + std::vector logical_shape_vec = {R_total, K}; + NVTEShape logical_shape = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size()); + std::vector data_1d = {R_total * K}; + NVTEShape data_shape = nvte_make_shape(data_1d.data(), data_1d.size()); + std::vector scale_1d = {total_scales}; + NVTEShape scale_shape = nvte_make_shape(scale_1d.data(), scale_1d.size()); + + NVTEShape offsets_shape; + offsets_shape.ndim = 1; + offsets_shape.data[0] = num_tensors + 1; + NVTEShape first_dims_shape; + first_dims_shape.ndim = 1; + first_dims_shape.data[0] = num_tensors; + NVTEBasicTensor offsets_bt = {offsets_d, kNVTEInt64, offsets_shape}; + NVTEBasicTensor first_dims_bt = {first_dims_d, kNVTEInt64, first_dims_shape}; + auto set_shape_meta = [&](NVTEGroupedTensor gt) { + if (shape_rep == ShapeRep::VARYING_FIRST_DIM) { + nvte_set_grouped_tensor_param(gt, kNVTEGroupedFirstDims, &first_dims_bt, + sizeof(first_dims_bt)); + nvte_set_grouped_tensor_param(gt, kNVTEGroupedTensorOffsets, &offsets_bt, sizeof(offsets_bt)); + } + }; + + NVTEGroupedTensor in_gt = nvte_create_grouped_tensor(mode, num_tensors, logical_shape); + NVTEBasicTensor in_data_bt = {data_d, static_cast(itype), data_shape}; + NVTEBasicTensor in_scale_bt = {scales_d, kNVTEFloat32, scale_shape}; + if (rowwise) { + nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedRowwiseData, &in_data_bt, sizeof(in_data_bt)); + nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedRowwiseScaleInv, &in_scale_bt, + sizeof(in_scale_bt)); + } else { + nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedColumnwiseData, &in_data_bt, + sizeof(in_data_bt)); + nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedColumnwiseScaleInv, &in_scale_bt, + sizeof(in_scale_bt)); + } + set_shape_meta(in_gt); + + NVTEGroupedTensor out_gt = + nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape); + NVTEBasicTensor out_data_bt = {out_grouped_d, static_cast(otype), data_shape}; + nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseData, &out_data_bt, sizeof(out_data_bt)); + set_shape_meta(out_gt); + + // ---- Grouped dequantize ---- + nvte_group_dequantize(in_gt, out_gt, 0); + cudaDeviceSynchronize(); + { + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + } + std::vector out_grouped_h(R_total * K); + cudaMemcpy(out_grouped_h.data(), out_grouped_d, R_total * K * sizeof(OutputType), + cudaMemcpyDeviceToHost); + + // ---- Reference: host dequantize (out = float(fp8) * scale_inv) ---- + // + // There is no non-grouped FP8 block-scaling dequantize to loop over (only the grouped path + // implements it), so the reference is computed on the host. The per-expert scale sub-block + // indexing mirrors what test_cast_float8blockwise_grouped.cu validates the grouped quantize + // kernel writes; this makes the host reference an independent check of the grouped dequant. + // 1D rowwise : scale[bx * roundup(M,4) + r] (bx = c/128) + // 1D colwise : scale[by * roundup(K,4) + c] (by = r/128) + // 2D rowwise : scale[by * roundup(blocks_X,4) + bx] + // 2D colwise : scale[bx * roundup(blocks_y,4) + by] + // Data is contiguous (M,K) for rowwise and transposed (K,M) for columnwise. + auto scale_index = [&](size_t r, size_t c, size_t M) -> size_t { + const size_t blocks_X = (K + kBlock - 1) / kBlock; + const size_t blocks_y = (M + kBlock - 1) / kBlock; + const size_t bx = c / kBlock; + const size_t by = r / kBlock; + if (block_dim == BlockDim::ONE_D) { + return columnwise ? (by * align4(K) + c) : (bx * align4(M) + r); + } + return columnwise ? (bx * align4(blocks_y) + by) : (by * align4(blocks_X) + bx); + }; + + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t row_offset = static_cast(offsets_h[t]) / K; + const size_t data_off = row_offset * K; // rowwise (M,K) or colwise transposed (K,M) + const size_t s_off = scale_off[t]; + for (size_t r = 0; r < M; ++r) { + for (size_t c = 0; c < K; ++c) { + const size_t d_idx = columnwise ? (c * M + r) : (r * K + c); + const float fp8v = static_cast(data_h[data_off + d_idx]); + const float sc = scales_h[s_off + scale_index(r, c, M)]; + const float ref = fp8v * sc; + const float got = static_cast(out_grouped_h[(row_offset + r) * K + c]); + const float rel = std::fabs(got - ref) / std::max(std::fabs(ref), 1e-3f); + ASSERT_LT(rel, 1e-2f) << "dequant mismatch t=" << t << " r=" << r << " c=" << c + << " got=" << got << " ref=" << ref << " (fp8=" << fp8v + << " scale=" << sc << ")"; + } + } + } + + nvte_destroy_grouped_tensor(in_gt); + nvte_destroy_grouped_tensor(out_gt); + cudaFree(data_d); + cudaFree(scales_d); + cudaFree(out_grouped_d); + cudaFree(offsets_d); + if (first_dims_d) cudaFree(first_dims_d); +} + +struct TestConfig { + ShapeRep shape_rep; + BlockDim block_dim; + bool rowwise; + std::vector first_dims; + size_t K; +}; + +std::vector make_configs() { + std::vector configs; + std::vector> uniform = {{128, 128}, {256, 256, 256, 256}}; + std::vector> jagged = {{128, 256, 384, 512}, {256, 128, 512, 384, 1024}}; + std::vector Ks = {128, 256, 512}; + for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) { + for (bool rowwise : {true, false}) { + for (size_t K : Ks) { + for (const auto& v : uniform) + configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, rowwise, v, K}); + for (const auto& v : jagged) + configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, rowwise, v, K}); + } + } + } + return configs; +} + +} // namespace + +class GroupedDequantizeFP8BlockwiseTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(GroupedDequantizeFP8BlockwiseTestSuite, Test) { + const TestConfig cfg = std::get<0>(GetParam()); + const DType output_type = std::get<1>(GetParam()); + // FP8 block scaling is E4M3-centric (matches the grouped quantize test scope). + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( + output_type, OutputType, + performTest(cfg.shape_rep, cfg.block_dim, cfg.rowwise, cfg.first_dims, + cfg.K);); +} + +INSTANTIATE_TEST_SUITE_P( + GroupedFP8Blockwise, GroupedDequantizeFP8BlockwiseTestSuite, + ::testing::Combine(::testing::ValuesIn(make_configs()), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)), + [](const testing::TestParamInfo& info) { + const TestConfig& c = std::get<0>(info.param); + std::string s = (c.shape_rep == ShapeRep::SAME_BOTH_DIMS ? "SAME" : "VARYFIRST"); + s += "_BD" + std::to_string(static_cast(c.block_dim)); + s += (c.rowwise ? "_RW" : "_CW"); + s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size()); + s += "_" + test::typeName(std::get<1>(info.param)); + return s; + }); diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu index 12b4703469..09fcbad8df 100644 --- a/tests/cpp/operator/test_grouped_gemm.cu +++ b/tests/cpp/operator/test_grouped_gemm.cu @@ -385,6 +385,36 @@ inline AlphaBetaTensors make_alpha_beta(size_t num_gemms) { // Compare each tensor inside a grouped D buffer (with per-tensor offsets) against the // reference D_multi[i] tensors. +// Capture `body(stream)` into a CUDA graph and replay it `n_replays` times, calling +// `verify(iter)` after each launch+sync. Used to assert that the grouped GEMM kernels +// are graph-safe (capture succeeds, replay is correct, and replays are deterministic). +template +inline void capture_and_replay_grouped_gemm(cudaStream_t stream, int n_replays, + Body&& body, Verify&& verify) { + // Warmup off-graph so cuBLASLt can initialize its internal heuristic / handle state + // outside capture (the very first matmul on a handle may do host-side setup that + // isn't capturable in some cuBLAS versions). + body(stream); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed)); + body(stream); + cudaGraph_t graph; + NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph)); + + cudaGraphExec_t exec; + NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + + for (int i = 0; i < n_replays; ++i) { + NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + verify(i); + } + + NVTE_CHECK_CUDA(cudaGraphExecDestroy(exec)); + NVTE_CHECK_CUDA(cudaGraphDestroy(graph)); +} + inline void compare_grouped_d_to_multi( const GroupedBuffers& grouped_D, const std::vector>& shapes, @@ -649,6 +679,255 @@ void run_grouped_gemm_discrete_in_case(const TestParams& params) { compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, "grouped_discrete_in_vs_multi"); } +// Graph-capture variant of run_grouped_gemm_case. Captures nvte_grouped_gemm on a +// non-default stream and replays it twice, verifying outputs against the multi-tensor +// reference after each replay. Asserts capture succeeds, replay is correct, and the +// operation is deterministic across replays. +void run_grouped_gemm_graph_case(const TestParams& params) { + if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) { + GTEST_SKIP() << reason; + } + auto ref = make_grouped_gemm_ref(params); + const auto& shapes = ref.shapes; + const size_t num_gemms = ref.num_gemms; + + std::vector A_views, B_views; + A_views.reserve(num_gemms); + B_views.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) { + A_views.push_back(&ref.A_tensors[i]); + B_views.push_back(&ref.B_tensors[i]); + } + GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode()); + GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode()); + + std::vector C_tensors, D_group_tensors; + C_tensors.reserve(num_gemms); + D_group_tensors.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) { + const auto [M, N, K] = shapes[i]; + (void)K; + if (!params.use_null_c) { + C_tensors.emplace_back( + Tensor("C" + std::to_string(i), std::vector{M, N}, params.output_dtype)); + } + D_group_tensors.emplace_back( + Tensor("D_group" + std::to_string(i), std::vector{M, N}, params.output_dtype)); + } + + std::vector C_views, D_views; + for (size_t i = 0; i < num_gemms; ++i) { + if (!params.use_null_c) C_views.push_back(&C_tensors[i]); + D_views.push_back(&D_group_tensors[i]); + } + + std::optional grouped_C; + if (!params.use_null_c) { + grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING); + } + GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING); + + AlphaBetaTensors ab = make_alpha_beta(num_gemms); + const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms); + Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); + Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte); + GroupedMatmulConfigWrapper grouped_config; + if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + + auto body = [&](cudaStream_t s) { + for (auto& d : D_group_tensors) { + NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0, + bytes(d.rowwise_shape(), d.dtype()), s)); + } + nvte_grouped_gemm(grouped_A.get_handle(), params.transa, grouped_B.get_handle(), + params.transb, + params.use_null_c ? nullptr : grouped_C->get_handle(), + grouped_D.get_handle(), ab.alpha.data(), ab.beta.data(), + setup_ws.data(), cublas_ws.data(), grouped_config, s); + }; + + capture_and_replay_grouped_gemm( + stream, /*n_replays=*/2, body, + [&](int iter) { + const std::string tag = "grouped_graph_replay_" + std::to_string(iter); + compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, tag.c_str()); + }); + + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// Graph-capture variant of run_grouped_gemm_discrete_out_case. +void run_grouped_gemm_discrete_out_graph_case(const TestParams& params) { + if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) { + GTEST_SKIP() << reason; + } + auto ref = make_grouped_gemm_ref(params); + const auto& shapes = ref.shapes; + const size_t num_gemms = ref.num_gemms; + + std::vector A_views, B_views; + A_views.reserve(num_gemms); + B_views.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) { + A_views.push_back(&ref.A_tensors[i]); + B_views.push_back(&ref.B_tensors[i]); + } + GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode()); + GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode()); + + std::vector C_tensors, D_list_tensors; + C_tensors.reserve(num_gemms); + D_list_tensors.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) { + const auto [M, N, K] = shapes[i]; + (void)K; + if (!params.use_null_c) { + C_tensors.emplace_back( + Tensor("C" + std::to_string(i), std::vector{M, N}, params.output_dtype)); + } + D_list_tensors.emplace_back( + Tensor("D_list" + std::to_string(i), std::vector{M, N}, params.output_dtype)); + } + + std::vector C_list_ptrs, D_list_ptrs; + if (!params.use_null_c) C_list_ptrs.reserve(num_gemms); + D_list_ptrs.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) { + if (!params.use_null_c) C_list_ptrs.push_back(C_tensors[i].data()); + D_list_ptrs.push_back(D_list_tensors[i].data()); + } + + AlphaBetaTensors ab = make_alpha_beta(num_gemms); + const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms); + Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); + Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte); + GroupedMatmulConfigWrapper grouped_config; + if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + + auto body = [&](cudaStream_t s) { + for (auto& d : D_list_tensors) { + NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0, + bytes(d.rowwise_shape(), d.dtype()), s)); + } + nvte_grouped_gemm_with_discrete_out( + grouped_A.get_handle(), params.transa, grouped_B.get_handle(), params.transb, + params.use_null_c ? nullptr : C_list_ptrs.data(), + params.use_null_c ? 0 : num_gemms, D_list_ptrs.data(), num_gemms, ab.alpha.data(), + ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, s); + }; + + auto verify = [&](int iter) { + const std::string tag = "discrete_out_graph_replay_" + std::to_string(iter); + for (size_t i = 0; i < num_gemms; ++i) { + D_list_tensors[i].to_cpu(); + ref.D_multi[i].to_cpu(); + auto [atol, rtol] = getTolerances(ref.D_multi[i].dtype()); + switch (ref.D_multi[i].dtype()) { + case DType::kBFloat16: + compareResults(tag.c_str(), D_list_tensors[i], + ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); + break; + case DType::kFloat16: + compareResults(tag.c_str(), D_list_tensors[i], + ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); + break; + case DType::kFloat32: + compareResults(tag.c_str(), D_list_tensors[i], + ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); + break; + default: + NVTE_ERROR("Unsupported D dtype in test: " + + std::to_string(static_cast(ref.D_multi[i].dtype()))); + } + } + }; + + capture_and_replay_grouped_gemm(stream, /*n_replays=*/2, body, verify); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// Graph-capture variant of run_grouped_gemm_discrete_in_case. +void run_grouped_gemm_discrete_in_graph_case(const TestParams& params) { + if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) { + GTEST_SKIP() << reason; + } + auto ref = make_grouped_gemm_ref(params); + const auto& shapes = ref.shapes; + const size_t num_gemms = ref.num_gemms; + + std::vector B_views; + B_views.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) B_views.push_back(&ref.B_tensors[i]); + GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode()); + + std::vector C_tensors, D_group_tensors; + C_tensors.reserve(num_gemms); + D_group_tensors.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) { + const auto [M, N, K] = shapes[i]; + (void)K; + if (!params.use_null_c) { + C_tensors.emplace_back(Tensor("C" + std::to_string(i), + std::vector{M, N}, params.output_dtype)); + } + D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i), + std::vector{M, N}, params.output_dtype)); + } + + std::vector C_views, D_views; + for (size_t i = 0; i < num_gemms; ++i) { + if (!params.use_null_c) C_views.push_back(&C_tensors[i]); + D_views.push_back(&D_group_tensors[i]); + } + + std::optional grouped_C; + if (!params.use_null_c) { + grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING); + } + GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING); + + AlphaBetaTensors ab = make_alpha_beta(num_gemms); + const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms); + Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); + Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte); + + std::vector A_list_ptrs; + A_list_ptrs.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) A_list_ptrs.push_back(ref.A_tensors[i].data()); + + GroupedMatmulConfigWrapper grouped_config; + if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + + auto body = [&](cudaStream_t s) { + for (auto& d : D_group_tensors) { + NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0, + bytes(d.rowwise_shape(), d.dtype()), s)); + } + nvte_grouped_gemm_with_discrete_inputA( + A_list_ptrs.data(), num_gemms, params.transa, grouped_B.get_handle(), params.transb, + params.use_null_c ? nullptr : grouped_C->get_handle(), grouped_D.get_handle(), + ab.alpha.data(), ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, s); + }; + + capture_and_replay_grouped_gemm( + stream, /*n_replays=*/2, body, + [&](int iter) { + const std::string tag = "discrete_in_graph_replay_" + std::to_string(iter); + compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, tag.c_str()); + }); + + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + class GroupedGemmTest : public ::testing::TestWithParam {}; TEST_P(GroupedGemmTest, CompareWithMultiTensorGemm) { @@ -663,6 +942,18 @@ TEST_P(GroupedGemmTest, CompareWithMultiTensorGemmDiscreteIn) { run_grouped_gemm_discrete_in_case(GetParam()); } +TEST_P(GroupedGemmTest, CudaGraphCapture) { + run_grouped_gemm_graph_case(GetParam()); +} + +TEST_P(GroupedGemmTest, CudaGraphCaptureDiscreteOut) { + run_grouped_gemm_discrete_out_graph_case(GetParam()); +} + +TEST_P(GroupedGemmTest, CudaGraphCaptureDiscreteIn) { + run_grouped_gemm_discrete_in_graph_case(GetParam()); +} + std::string MakeGroupedGemmTestName(const testing::TestParamInfo& info) { constexpr const char* kShapeNames[] = {"AllSameMul128", "SameMMul128", "SameNMul128", "AllDiffMul128", "AllSameMul32"}; diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index e1e15e6875..eeb6e7a394 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -33,6 +33,21 @@ mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +# The fused grouped FP8 block-scaling quantize/dequantize kernels are Hopper-only: they gate on +# SM90-SM99 (NVTE_CHECK(sm >= 90 && sm < 100)). FP8 block scaling is still reported "available" on +# Blackwell (SM100+) for the emulated/non-grouped paths, so ``fp8_block_scaling_available`` alone +# does not exclude SM100 — add the Hopper arch bound for the grouped tests. +_device_cc = torch.cuda.get_device_capability() if torch.cuda.is_available() else (0, 0) +fp8_block_scaling_grouped_available = fp8_block_scaling_available and (9, 0) <= _device_cc < (10, 0) +reason_for_no_fp8_block_scaling_grouped = ( + reason_for_no_fp8_block_scaling + if not fp8_block_scaling_available + else ( + "Fused grouped FP8 block-scaling quantize/dequantize is only supported on Hopper" + " (SM90-SM99)." + ) +) + _quantization_params = [ pytest.param( "fp8_delayed_scaling", @@ -115,6 +130,33 @@ def _rowwise_offset_bytes(numel: int, quantization: str) -> int: return numel +def _fp8bs_per_expert_scale_floats( + block_scaling_dim: int, columnwise: bool, m_t: int, k: int +) -> int: + """Per-expert padded scale size (in floats) for grouped FP8 block-scaling. + + Mirrors the per-expert sub-block layout that cuBLAS grouped GEMM consumes (and that the C++ + test ``test_cast_float8blockwise_grouped.cu`` verifies against):: + + 1D rowwise : blocks_X * roundup(M_t, 4) + 1D colwise : blocks_y_t * roundup(K, 4) + 2D rowwise : blocks_y_t * roundup(blocks_X, 4) + 2D colwise : blocks_X * roundup(blocks_y_t, 4) + + The 2D columnwise roundup of each expert's block-rows to a multiple of 4 is the source of the + per-expert slack reserved in ``Float8BlockQuantizer.create_grouped_tensor``. + """ + + def align4(x: int) -> int: + return ((x + 3) // 4) * 4 + + blocks_x = (k + 127) // 128 + blocks_y = (m_t + 127) // 128 + if block_scaling_dim == 1: + return blocks_y * align4(k) if columnwise else blocks_x * align4(m_t) + return blocks_x * align4(blocks_y) if columnwise else blocks_y * align4(blocks_x) + + class TestGroupedTensor: @staticmethod def setup_class(cls) -> None: @@ -715,6 +757,54 @@ def _run_group_quantize(input_tensor): if output_dbias: assert torch.allclose(static_dbias, expected_dbias) + @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"]) + @pytest.mark.skipif( + not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped + ) + def test_group_quantize_fp8_blockwise_cudagraph_capturable( + self, block_scaling_dim: int + ) -> None: + """Ensure grouped FP8 block-scaling quantize is CUDA graph capturable (parity with MXFP8).""" + first_dims_host = [256, 128, 384] + num_tensors = len(first_dims_host) + hidden = 512 + shape = [(r, hidden) for r in first_dims_host] + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + grouped_input = torch.cat(input_tensors, dim=0) + first_dims = torch.tensor(first_dims_host, dtype=torch.int64, device="cuda") + + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=block_scaling_dim, + ) + + torch.cuda.synchronize() + static_input = grouped_input.clone() + static_first_dims = first_dims.clone() + + def _run(inp): + return tex.group_quantize(inp, quantizer, num_tensors, static_first_dims) + + _ = _run(static_input) # warmup allocator/kernels + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output = _run(static_input) + + # Replay with fresh input copied into the captured buffer. + static_input.copy_(torch.randn_like(grouped_input)) + graph.replay() + torch.cuda.synchronize() + + expected = _run(static_input) + assert torch.equal(static_output.rowwise_data, expected.rowwise_data) + assert torch.equal(static_output.scale_inv, expected.scale_inv) + @pytest.mark.parametrize("mode", ["rowwise", "columnwise", "both"]) @pytest.mark.parametrize( "shape_case", @@ -914,6 +1004,117 @@ def _assert_fp8_cs_group_quantize_matches_reference( expected = torch.cat(expected_columnwise) assert torch.equal(grouped_output.columnwise_data[: expected.numel()], expected) + @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"]) + @pytest.mark.parametrize("shape_case", ["uniform", "varying_first"]) + @pytest.mark.parametrize("direction", ["rowwise", "columnwise", "both"]) + @pytest.mark.parametrize("output_dbias", [False, True]) + @pytest.mark.skipif( + not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped + ) + def test_quantize_grouped_fp8_blockwise( + self, block_scaling_dim: int, shape_case: str, direction: str, output_dbias: bool + ) -> None: + """Test grouped FP8 block-scaling quantization against per-tensor quantization. + + Covers rowwise, columnwise and both directions. Each expert's data sub-block (and, for + rowwise, its scale sub-block placed at the cumulative padded offset from + ``_fp8bs_per_expert_scale_floats``) is compared against an independent per-tensor reference + quantizer. The columnwise scale layout carries 2D padding plus the per-expert slack + reserved in ``Float8BlockQuantizer.create_grouped_tensor``; it is validated end to end by + ``test_group_dequantize_fp8_blockwise`` rather than by a fragile byte-compare here. + + FP8 block-scaling supports only SAME_BOTH_DIMS (``uniform``) and VARYING_FIRST_DIM + (``varying_first``); ``varying_last``/``varying_both`` are rejected at the kernel level. + Per-tensor first dim must be a multiple of 128 (kernel tile size). + """ + rowwise = direction in ("rowwise", "both") + columnwise = direction in ("columnwise", "both") + + # dbias is the bias gradient (per-column input sum) emitted by the bgrad path, which + # requires rowwise output; the columnwise-only + dbias combination is not applicable. + if output_dbias and not rowwise: + pytest.skip("bgrad (dbias) requires rowwise output; columnwise-only does not apply.") + + if shape_case == "uniform": + per_tensor_shapes = [(128, 512)] * 3 + first_dims_host = None + else: # varying_first + per_tensor_shapes = [(128, 512), (256, 512), (384, 512)] + first_dims_host = [s[0] for s in per_tensor_shapes] + + num_tensors = len(per_tensor_shapes) + + input_tensors = [ + torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in per_tensor_shapes + ] + flat_buffer = torch.cat([t.reshape(-1) for t in input_tensors]) + common_last = per_tensor_shapes[0][1] + grouped_input = flat_buffer.view(-1, common_last) + + first_dims = ( + torch.tensor(first_dims_host, dtype=torch.int64, device="cuda") + if first_dims_host is not None + else None + ) + + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=block_scaling_dim, + ) + + if output_dbias: + grouped_output, dbias = tex.bgrad_group_quantize( + grouped_input, quantizer, num_tensors, first_dims + ) + else: + grouped_output = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims) + + # Compare each expert's sub-block against an independent per-tensor reference. The + # reference enables both directions: the non-grouped 2D kernel requires rowwise output to + # be allocated even when only columnwise is consumed, and the columnwise data/scale it + # emits are independent of whether rowwise is also computed. + ref_quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=block_scaling_dim, + ) + # Data sub-blocks are contiguous (rowwise (M_t, K) / columnwise transposed (K, M_t)) with + # no inter-expert padding. The rowwise scale sub-block is placed at the cumulative + # per-expert padded offset, so a wrong stride fails byte-equality here. The columnwise + # scale carries 2D ``roundup(blocks_y_t, 4)`` padding columns (and the per-expert slack), + # so its layout is validated end to end by ``test_group_dequantize_fp8_blockwise`` instead. + data_off = 0 + rw_scale_off = 0 + for tensor in input_tensors: + m_t, k = tensor.shape + numel = m_t * k + ref = ref_quantizer(tensor) + if rowwise: + ref_rw = ref._rowwise_data.reshape(-1) + assert torch.equal(grouped_output.rowwise_data[data_off : data_off + numel], ref_rw) + ref_rs = ref._rowwise_scale_inv.reshape(-1) + assert torch.equal( + grouped_output.scale_inv[rw_scale_off : rw_scale_off + ref_rs.numel()], ref_rs + ) + rw_scale_off += _fp8bs_per_expert_scale_floats(block_scaling_dim, False, m_t, k) + if columnwise: + ref_cw = ref._columnwise_data.reshape(-1) + assert torch.equal( + grouped_output.columnwise_data[data_off : data_off + numel], ref_cw + ) + data_off += numel + + if output_dbias: + expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) + assert torch.allclose(dbias, expected_dbias) + @pytest.mark.parametrize( "shape", [[(512, 1024), (512, 1024)], [(256, 512), (512, 512), (768, 512)]], @@ -995,6 +1196,115 @@ def test_group_dequantize_cudagraph_capturable(self) -> None: for exp, got in zip(expected_tensors, static_tensors): assert torch.equal(got, exp) + @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"]) + @pytest.mark.parametrize("direction", ["rowwise", "columnwise"]) + @pytest.mark.skipif( + not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped + ) + def test_group_dequantize_fp8_blockwise_cudagraph_capturable( + self, block_scaling_dim: int, direction: str + ) -> None: + """Ensure grouped FP8 block-scaling dequantize is CUDA graph capturable (parity with MXFP8).""" + rowwise = direction == "rowwise" + columnwise = direction == "columnwise" + num_tensors = 2 + shape = [(512, 1024) for _ in range(num_tensors)] + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + grouped_input = torch.cat(input_tensors, dim=0) + + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=block_scaling_dim, + ) + first_dims = torch.tensor( + [shape[0][0] for _ in range(num_tensors)], dtype=torch.int64, device="cuda" + ) + + quantized = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims) + + # Warmup dequantize. + torch.cuda.synchronize() + _ = tex.group_dequantize(quantized, te.DType.kBFloat16) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output = tex.group_dequantize(quantized, te.DType.kBFloat16) + + # Replay with fresh quantized data copied into the captured input buffers. + fresh_input = torch.cat( + [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape], dim=0 + ) + fresh_quantized = tex.group_quantize(fresh_input, quantizer, num_tensors, first_dims) + if rowwise: + quantized.rowwise_data.copy_(fresh_quantized.rowwise_data) + quantized.scale_inv.copy_(fresh_quantized.scale_inv) + else: + quantized.columnwise_data.copy_(fresh_quantized.columnwise_data) + quantized.columnwise_scale_inv.copy_(fresh_quantized.columnwise_scale_inv) + + graph.replay() + torch.cuda.synchronize() + + expected = tex.group_dequantize(quantized, te.DType.kBFloat16) + expected_tensors = expected.split_into_quantized_tensors() + static_tensors = static_output.split_into_quantized_tensors() + for exp, got in zip(expected_tensors, static_tensors): + assert torch.equal(got, exp) + + @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"]) + @pytest.mark.parametrize("direction", ["rowwise", "columnwise"]) + @pytest.mark.parametrize( + "shape", + [[(512, 1024), (512, 1024)], [(128, 512), (256, 512), (384, 512)]], + ) + @pytest.mark.skipif( + not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped + ) + def test_group_dequantize_fp8_blockwise( + self, block_scaling_dim: int, direction: str, shape: List[Tuple[int, int]] + ) -> None: + """Test grouped FP8 block-scaling dequantize round-trip for rowwise and columnwise. + + The columnwise + ``varying_first`` + 2D case exercises the per-expert columnwise + scale-buffer slack end to end: a wrong slack/stride would place scales at the wrong + offsets and corrupt the dequantized values. ``group_dequantize`` consumes exactly one of + rowwise / columnwise data, so ``both`` is not a valid round-trip and is not tested here. + """ + rowwise = direction == "rowwise" + columnwise = direction == "columnwise" + num_tensors = len(shape) + + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + grouped_input = torch.cat(input_tensors, dim=0) + + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=block_scaling_dim, + ) + first_dims = torch.tensor([s[0] for s in shape], dtype=torch.int64, device="cuda") + + quantized = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims) + dequantized = tex.group_dequantize(quantized, te.DType.kBFloat16) + + assert dequantized.num_tensors == num_tensors + assert dequantized.logical_shape == quantized.logical_shape + assert torch.equal(dequantized.first_dims, quantized.first_dims) + assert torch.equal(dequantized.tensor_offsets, quantized.tensor_offsets) + + dequantized_tensors = dequantized.split_into_quantized_tensors() + assert len(dequantized_tensors) == num_tensors + for orig, deq in zip(input_tensors, dequantized_tensors): + torch.testing.assert_close(deq, orig, atol=0.125, rtol=0.1) + def test_clear(self) -> None: """Test clear method""" num_tensors = 3 diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh index 63c1b046ff..bf4a021811 100644 --- a/transformer_engine/common/cast/dispatch/dequantize.cuh +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -15,6 +15,7 @@ #include "../../common.h" #include "../fp8/dequantize_fp8.cuh" +#include "../fp8_blockwise/group_dequantize_fp8_blockwise.cuh" #include "../mxfp8/dequantize_mxfp8.cuh" #include "../mxfp8/group_dequantize_mxfp8.cuh" #include "../nvfp4/dequantize_nvfp4.cuh" @@ -69,6 +70,11 @@ inline void group_dequantize_helper(const GroupedTensor &input, GroupedTensor *o } break; } + case NVTE_BLOCK_SCALING_1D: + case NVTE_BLOCK_SCALING_2D: { + fp8_blockwise::group_dequantize(&input, output, stream); + break; + } default: NVTE_ERROR("Grouped dequantize not implemented for scaling mode: " + to_string(input.scaling_mode) + "."); diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 97dd27aec6..031122d966 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -19,6 +19,7 @@ #include "../core/common.cuh" #include "../fp8/group_quantize_fp8.cuh" #include "../fp8/quantize_fp8.cuh" +#include "../fp8_blockwise/group_quantize_fp8_blockwise.cuh" #include "../mxfp8/group_quantize_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" #include "../nvfp4/group_quantize_transpose_nvfp4.cuh" @@ -472,6 +473,26 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor workspace_tensor, &quant_config_cpp, stream); break; } + case NVTE_BLOCK_SCALING_1D: { + NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_1D."); + NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, + "Fused grouped FP8 block-scaling quantize does not support " + "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " + "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); + fp8_blockwise::group_quantize_blockwise_1d(input_tensor, output_tensor, noop_tensor, + quant_config_cpp.amax_epsilon, stream); + break; + } + case NVTE_BLOCK_SCALING_2D: { + NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_2D."); + NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, + "Fused grouped FP8 block-scaling quantize does not support " + "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " + "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); + fp8_blockwise::group_quantize_blockwise_2d(input_tensor, output_tensor, noop_tensor, + quant_config_cpp.amax_epsilon, stream); + break; + } default: NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + "."); } @@ -513,6 +534,28 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe &quant_config_cpp, stream); break; } + case NVTE_BLOCK_SCALING_1D: + case NVTE_BLOCK_SCALING_2D: { + NVTE_CHECK(!IS_DACT, "IS_DACT is not implemented for grouped FP8 block scaling."); + NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, + "Fused grouped FP8 block-scaling quantize does not support " + "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " + "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); + // dbias is computed in-kernel and reduced per-expert inside group_quantize_blockwise_{1d,2d} + // (mirrors MXFP8); those also handle the two-call workspace sizing protocol. + GroupedTensor *dbias_arg = IS_DBIAS ? dbias_tensor : nullptr; + Tensor *workspace_arg = IS_DBIAS ? workspace_tensor : nullptr; + if (scaling_mode == NVTE_BLOCK_SCALING_1D) { + fp8_blockwise::group_quantize_blockwise_1d(grad_tensor, output_tensor, noop_tensor, + quant_config_cpp.amax_epsilon, stream, dbias_arg, + workspace_arg); + } else { + fp8_blockwise::group_quantize_blockwise_2d(grad_tensor, output_tensor, noop_tensor, + quant_config_cpp.amax_epsilon, stream, dbias_arg, + workspace_arg); + } + break; + } default: NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + "."); } diff --git a/transformer_engine/common/cast/fp8_blockwise/group_dequantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_dequantize_fp8_blockwise.cuh new file mode 100644 index 0000000000..556dcad428 --- /dev/null +++ b/transformer_engine/common/cast/fp8_blockwise/group_dequantize_fp8_blockwise.cuh @@ -0,0 +1,519 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file group_dequantize_fp8_blockwise.cuh + * \brief CUDA kernels to dequantize grouped tensors from FP8 with 1D/2D + * block scaling (rowwise or columnwise) back to BF16 / FP16 / FP32. Mirrors + * the per-expert layouts written by ``group_quantize_fp8_blockwise``. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_FP8_BLOCKWISE_CUH_ +#define TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_FP8_BLOCKWISE_CUH_ + +#include +#include +#include + +#include "../../common.h" +#include "../../utils.cuh" +#include "../core/common.cuh" +#include "group_quantize_fp8_blockwise.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8_blockwise { + +namespace group_dequantize_kernel { + +// Resolve which expert a tile (blocks_X x total_row_blocks grid) belongs to and its row range. +// tensor_M (the expert's first-dim length) addresses the per-expert (K, M_t) transposed block +// in the columnwise modes. +struct TileExpertInfo { + size_t tensor_id; + size_t tensor_block_y_base; + size_t tensor_row_blocks; + size_t tensor_row_base; + size_t tensor_M; + bool in_bounds; +}; + +template +__device__ __forceinline__ TileExpertInfo resolve_tile_expert( + size_t tile_y_global, size_t num_tensors, size_t common_first_dim_blocks, size_t K, + size_t total_row_blocks, const int64_t* __restrict__ tensor_offsets_ptr) { + TileExpertInfo info{}; + info.in_bounds = false; + if (tile_y_global >= total_row_blocks) return info; + const size_t tile_row_stride = static_cast(kTileDim) * K; + info.tensor_id = find_tensor_id_by_block_y( + tile_y_global, num_tensors, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr); + info.tensor_block_y_base = + kSameBothDims + ? (info.tensor_id * common_first_dim_blocks) + : tensor_block_y_base_from_offsets(info.tensor_id, tensor_offsets_ptr, tile_row_stride); + info.tensor_row_blocks = kSameBothDims + ? common_first_dim_blocks + : (tensor_block_y_base_from_offsets( + info.tensor_id + 1, tensor_offsets_ptr, tile_row_stride) - + info.tensor_block_y_base); + if (tile_y_global >= info.tensor_block_y_base + info.tensor_row_blocks) return info; + info.tensor_row_base = info.tensor_block_y_base * kTileDim; + info.tensor_M = info.tensor_row_blocks * kTileDim; + info.in_bounds = true; + return info; +} + +// ===== 1D rowwise ===== +// Per-expert scale layout: (blocks_X, roundup(M_t, 4)) floats. +// scale[expert_off + tile_x * roundup(M_t, 4) + r_local] +template +__global__ void __launch_bounds__(kThreadsPerBlock) + group_dequantize_blockwise_1d_rw_kernel(const IType* __restrict__ input_base, + OType* __restrict__ output_base, + const CType* __restrict__ scale_inv_base, + const int64_t* __restrict__ tensor_offsets_ptr, + const size_t num_tensors, + const size_t common_first_dim_blocks, const size_t K, + const size_t total_row_blocks, const size_t R_total) { +#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 + const size_t tile_x = blockIdx.x; + const size_t tile_y_global = blockIdx.y; + const auto info = resolve_tile_expert( + tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr); + if (!info.in_bounds) return; + + const size_t blocks_X = DIVUP(K, static_cast(kTileDim)); + const size_t tile_row_stride = static_cast(kTileDim) * K; + const size_t expert_offset = expert_scale_offset_1d_rowwise( + info.tensor_id, blocks_X, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr); + const size_t per_expert_stride = DIVUP_TO_MULTIPLE(info.tensor_M, kScaleColAlign); + const CType* const tile_scale_inv_base = + scale_inv_base + expert_offset + tile_x * per_expert_stride; + + const size_t global_row_base = tile_y_global * kTileDim; + const size_t global_col_base = tile_x * kTileDim; + + constexpr int kThreadsPerRow = 8; + constexpr int kVec = 16; + constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow; // 32 + constexpr int kIters = kTileDim / kRowsPerIter; // 4 + + const int tid = threadIdx.x; + const int thr_col = tid % kThreadsPerRow; + const int thr_row = tid / kThreadsPerRow; + const size_t c = global_col_base + static_cast(thr_col) * kVec; + +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int row_local = thr_row + it * kRowsPerIter; + const size_t r_global = global_row_base + row_local; + if (r_global >= R_total) continue; + + const size_t r_local = r_global - info.tensor_row_base; + const CType s_inv = tile_scale_inv_base[r_local]; + + Vec in_vec; + if (c + kVec <= K) { + in_vec.load_from(input_base + r_global * K + c); + } else if (c < K) { + in_vec.load_from_elts(input_base + r_global * K + c, 0, K - c); + } else { + continue; + } + + Vec out_vec; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv); + } + + if (c + kVec <= K) { + out_vec.store_to(output_base + r_global * K + c); + } else if (c < K) { + out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c); + } + } +#endif +} + +// ===== 1D columnwise ===== +// Data layout per-expert: (K, M_t) transposed -- element (r, c) is at +// input[tensor_row_base * K + c * tensor_M + r_local]. +// Scale layout GLOBAL: (total_row_blocks, roundup(K, 4)) floats. +// scale_inv[tile_y_global * scale_t_stride_aligned_K + c] +template +__global__ void __launch_bounds__(kThreadsPerBlock) + group_dequantize_blockwise_1d_cw_kernel(const IType* __restrict__ input_base, + OType* __restrict__ output_base, + const CType* __restrict__ scale_inv_base, + const size_t scale_t_stride_aligned_K, + const int64_t* __restrict__ tensor_offsets_ptr, + const size_t num_tensors, + const size_t common_first_dim_blocks, const size_t K, + const size_t total_row_blocks, const size_t R_total) { +#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 + const size_t tile_x = blockIdx.x; + const size_t tile_y_global = blockIdx.y; + const auto info = resolve_tile_expert( + tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr); + if (!info.in_bounds) return; + + const size_t expert_data_off = info.tensor_row_base * K; + const CType* const tile_scale_base = scale_inv_base + tile_y_global * scale_t_stride_aligned_K; + + const size_t global_row_base = tile_y_global * kTileDim; + const size_t global_col_base = tile_x * kTileDim; + + constexpr int kThreadsPerRow = 8; + constexpr int kVec = 16; + constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow; + constexpr int kIters = kTileDim / kRowsPerIter; + + const int tid = threadIdx.x; + const int thr_col = tid % kThreadsPerRow; + const int thr_row = tid / kThreadsPerRow; + const size_t c = global_col_base + static_cast(thr_col) * kVec; + + // 1D columnwise has one scale per column. Pre-load this thread's 16 columns. + CType s_inv[kVec]; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + s_inv[e] = (c + e < K) ? tile_scale_base[c + e] : static_cast(0.f); + } + +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int row_local = thr_row + it * kRowsPerIter; + const size_t r_global = global_row_base + row_local; + if (r_global >= R_total) continue; + + const size_t r_local = r_global - info.tensor_row_base; + + // Per-expert (K, M_t) transposed input: strided by M_t per column (no vector load). + // K % 128 == 0 so c+e < K always holds; the explicit else just keeps correctness + // independent of that invariant. + Vec in_vec; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + in_vec.data.elt[e] = (c + e < K) + ? input_base[expert_data_off + (c + e) * info.tensor_M + r_local] + : static_cast(0); + } + + Vec out_vec; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv[e]); + } + + if (c + kVec <= K) { + out_vec.store_to(output_base + r_global * K + c); + } else if (c < K) { + out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c); + } + } +#endif +} + +// ===== 2D rowwise ===== +// Data layout: (M, K) flat. One scale per 128x128 tile. +// Scale layout GLOBAL: (total_row_blocks, roundup(blocks_X, 4)) floats. +// scale[tile_y_global * scale_stride_y + tile_x] +template +__global__ void __launch_bounds__(kThreadsPerBlock) + group_dequantize_blockwise_2d_rw_kernel(const IType* __restrict__ input_base, + OType* __restrict__ output_base, + const CType* __restrict__ scale_inv_base, + const size_t scale_stride_y, + const int64_t* __restrict__ tensor_offsets_ptr, + const size_t num_tensors, + const size_t common_first_dim_blocks, const size_t K, + const size_t total_row_blocks, const size_t R_total) { +#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 + const size_t tile_x = blockIdx.x; + const size_t tile_y_global = blockIdx.y; + const auto info = resolve_tile_expert( + tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr); + if (!info.in_bounds) return; + + // 2D: one scale per tile. + const CType s_inv = scale_inv_base[tile_y_global * scale_stride_y + tile_x]; + + const size_t global_row_base = tile_y_global * kTileDim; + const size_t global_col_base = tile_x * kTileDim; + + constexpr int kThreadsPerRow = 8; + constexpr int kVec = 16; + constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow; + constexpr int kIters = kTileDim / kRowsPerIter; + + const int tid = threadIdx.x; + const int thr_col = tid % kThreadsPerRow; + const int thr_row = tid / kThreadsPerRow; + const size_t c = global_col_base + static_cast(thr_col) * kVec; + +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int row_local = thr_row + it * kRowsPerIter; + const size_t r_global = global_row_base + row_local; + if (r_global >= R_total) continue; + + Vec in_vec; + if (c + kVec <= K) { + in_vec.load_from(input_base + r_global * K + c); + } else if (c < K) { + in_vec.load_from_elts(input_base + r_global * K + c, 0, K - c); + } else { + continue; + } + + Vec out_vec; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv); + } + + if (c + kVec <= K) { + out_vec.store_to(output_base + r_global * K + c); + } else if (c < K) { + out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c); + } + } +#endif +} + +// ===== 2D columnwise ===== +// Data layout per-expert: (K, M_t) transposed. +// Scale layout per-expert: (blocks_X, roundup(blocks_y_t, 4)) floats. +// scale[expert_off + tile_x * roundup(blocks_y_t, 4) + local_tile_y] +template +__global__ void __launch_bounds__(kThreadsPerBlock) + group_dequantize_blockwise_2d_cw_kernel(const IType* __restrict__ input_base, + OType* __restrict__ output_base, + const CType* __restrict__ scale_inv_base, + const int64_t* __restrict__ tensor_offsets_ptr, + const size_t num_tensors, + const size_t common_first_dim_blocks, const size_t K, + const size_t total_row_blocks, const size_t R_total) { +#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 + const size_t tile_x = blockIdx.x; + const size_t tile_y_global = blockIdx.y; + const auto info = resolve_tile_expert( + tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr); + if (!info.in_bounds) return; + + // `info.in_bounds` is derived from blockIdx.y and is uniform across the CTA, + // so the early return above never strands a sibling thread inside + // compute_2d_cw_expert_offset's __syncthreads(). + __shared__ size_t warp_offset_partials[kNumWarps]; + + const int tid = threadIdx.x; + const int warp_id = tid / kThreadsPerWarp; + const int lane = tid % kThreadsPerWarp; + + const size_t blocks_X = DIVUP(K, static_cast(kTileDim)); + const size_t tile_row_stride = static_cast(kTileDim) * K; + const size_t expert_offset = compute_2d_cw_expert_offset( + info.tensor_id, blocks_X, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr, + warp_offset_partials, tid, warp_id, lane); + const size_t per_expert_stride_t = DIVUP_TO_MULTIPLE(info.tensor_row_blocks, kScaleColAlign); + const size_t local_tile_y = tile_y_global - info.tensor_block_y_base; + const CType s_inv = scale_inv_base[expert_offset + tile_x * per_expert_stride_t + local_tile_y]; + + const size_t expert_data_off = info.tensor_row_base * K; + const size_t global_row_base = tile_y_global * kTileDim; + const size_t global_col_base = tile_x * kTileDim; + + constexpr int kThreadsPerRow = 8; + constexpr int kVec = 16; + constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow; + constexpr int kIters = kTileDim / kRowsPerIter; + + const int thr_col = tid % kThreadsPerRow; + const int thr_row = tid / kThreadsPerRow; + const size_t c = global_col_base + static_cast(thr_col) * kVec; + +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int row_local = thr_row + it * kRowsPerIter; + const size_t r_global = global_row_base + row_local; + if (r_global >= R_total) continue; + + const size_t r_local = r_global - info.tensor_row_base; + + // Explicit else zeroes out-of-range lanes (c+e < K always holds since K % 128 == 0; + // this keeps correctness independent of that invariant). + Vec in_vec; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + in_vec.data.elt[e] = (c + e < K) + ? input_base[expert_data_off + (c + e) * info.tensor_M + r_local] + : static_cast(0); + } + + Vec out_vec; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv); + } + + if (c + kVec <= K) { + out_vec.store_to(output_base + r_global * K + c); + } else if (c < K) { + out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c); + } + } +#endif +} + +} // namespace group_dequantize_kernel + +// Host-side dispatcher. Supports all four combinations of {1D, 2D} block +// scaling x {rowwise, columnwise} data, matching the layouts written by +// ``group_quantize_fp8_blockwise``. The input GroupedTensor must have exactly +// one of rowwise / columnwise data populated (the dequantize API rejects +// both). +inline void group_dequantize(const GroupedTensor* input, GroupedTensor* output, + cudaStream_t stream) { + using namespace group_dequantize_kernel; + + const int sm = transformer_engine::cuda::sm_arch(); + NVTE_CHECK(sm >= 90 && sm < 100, + "Grouped FP8 block-scaling dequantize is only supported on Hopper (SM90-SM99); " + "got SM", + sm, "."); + NVTE_CHECK( + input->scaling_mode == NVTE_BLOCK_SCALING_1D || input->scaling_mode == NVTE_BLOCK_SCALING_2D, + "Grouped FP8 block-scaling dequantize requires 1D or 2D block scaling " + "(got scaling_mode=", + to_string(input->scaling_mode), ")."); + NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input must have FP8 type."); + NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision."); + NVTE_CHECK(!is_fp4_dtype(output->dtype()), "Output must not be FP4."); + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Number of input and output tensors must match."); + + const bool use_rowwise = input->has_data(); + const bool use_colwise = input->has_columnwise_data(); + NVTE_CHECK(use_rowwise || use_colwise, "Input must have rowwise or columnwise data populated."); + NVTE_CHECK(!(use_rowwise && use_colwise), + "Grouped FP8 block-scaling dequantize accepts exactly one direction at a " + "time (not both rowwise and columnwise simultaneously)."); + NVTE_CHECK(!input->with_gemm_swizzled_scales, + "Grouped FP8 block-scaling dequantize requires compact (un-swizzled) scales."); + + const size_t first_logical_dim = input->logical_shape.data[0]; + const size_t last_logical_dim = input->logical_shape.data[1]; + if (first_logical_dim == 0 || last_logical_dim == 0) return; + + const bool same_both_dims = input->all_same_shape(); + const bool varying_first_dim = (!input->all_same_first_dim()) && input->all_same_last_dim(); + NVTE_CHECK(same_both_dims || varying_first_dim, + "Grouped FP8 block-scaling dequantize supports only SAME_BOTH_DIMS and " + "VARYING_FIRST_DIM shape representations."); + + const size_t num_tensors = input->num_tensors; + const size_t K = last_logical_dim; + NVTE_CHECK(K % kTileDim == 0, + "Last dim must be a multiple of 128 for FP8 block-scaling dequantize (got ", K, ")."); + + size_t common_first_dim_blocks = 0; + if (same_both_dims) { + const size_t common_first_dim = input->get_common_first_dim(); + NVTE_CHECK(common_first_dim % kTileDim == 0, + "SAME_BOTH_DIMS first dim must be multiple of 128 (got ", common_first_dim, ")."); + common_first_dim_blocks = common_first_dim / kTileDim; + } + const size_t total_row_blocks = DIVUP(first_logical_dim, static_cast(kTileDim)); + const size_t blocks_X = K / kTileDim; + + const int64_t* tensor_offsets_ptr = + same_both_dims ? nullptr : reinterpret_cast(input->tensor_offsets.dptr); + if (!same_both_dims) { + NVTE_CHECK(tensor_offsets_ptr != nullptr, + "VARYING_FIRST_DIM requires tensor_offsets to be set on the input."); + } + + const dim3 grid(blocks_X, total_row_blocks); + const dim3 block(kThreadsPerBlock); + const bool is_1d = (input->scaling_mode == NVTE_BLOCK_SCALING_1D); + + // Pick the populated direction's data + scale buffer. The global-layout strides (1D cw, 2D rw) + // are derived from K / blocks_X exactly as the quantize launcher does. + const SimpleTensor& input_data = use_rowwise ? input->data : input->columnwise_data; + const SimpleTensor& input_scale_inv = + use_rowwise ? input->scale_inv : input->columnwise_scale_inv; + const size_t scale_t_stride_aligned_K = DIVUP_TO_MULTIPLE(K, kScaleColAlign); + const size_t scale_stride_y = DIVUP_TO_MULTIPLE(blocks_X, kScaleColAlign); + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + output->dtype(), OType, using CType = float; + const IType* const input_dptr = reinterpret_cast(input_data.dptr); + OType* const output_dptr = reinterpret_cast(output->data.dptr); + const CType* const scale_inv_dptr = reinterpret_cast(input_scale_inv.dptr); + + if (is_1d && use_rowwise) { + if (same_both_dims) { + group_dequantize_blockwise_1d_rw_kernel + <<>>( + input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors, + common_first_dim_blocks, K, total_row_blocks, first_logical_dim); + } else { + group_dequantize_blockwise_1d_rw_kernel + <<>>( + input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors, + common_first_dim_blocks, K, total_row_blocks, first_logical_dim); + } + } else if (is_1d && use_colwise) { + if (same_both_dims) { + group_dequantize_blockwise_1d_cw_kernel + <<>>(input_dptr, output_dptr, scale_inv_dptr, + scale_t_stride_aligned_K, tensor_offsets_ptr, + num_tensors, common_first_dim_blocks, K, + total_row_blocks, first_logical_dim); + } else { + group_dequantize_blockwise_1d_cw_kernel + <<>>(input_dptr, output_dptr, scale_inv_dptr, + scale_t_stride_aligned_K, tensor_offsets_ptr, + num_tensors, common_first_dim_blocks, K, + total_row_blocks, first_logical_dim); + } + } else if (!is_1d && use_rowwise) { + if (same_both_dims) { + group_dequantize_blockwise_2d_rw_kernel + <<>>( + input_dptr, output_dptr, scale_inv_dptr, scale_stride_y, tensor_offsets_ptr, + num_tensors, common_first_dim_blocks, K, total_row_blocks, first_logical_dim); + } else { + group_dequantize_blockwise_2d_rw_kernel + <<>>( + input_dptr, output_dptr, scale_inv_dptr, scale_stride_y, tensor_offsets_ptr, + num_tensors, common_first_dim_blocks, K, total_row_blocks, first_logical_dim); + } + } else { // 2D columnwise + if (same_both_dims) { + group_dequantize_blockwise_2d_cw_kernel + <<>>( + input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors, + common_first_dim_blocks, K, total_row_blocks, first_logical_dim); + } else { + group_dequantize_blockwise_2d_cw_kernel + <<>>( + input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors, + common_first_dim_blocks, K, total_row_blocks, first_logical_dim); + } + }); // NOLINT(*) + ); // NOLINT(*) + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace fp8_blockwise +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_FP8_BLOCKWISE_CUH_ diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh new file mode 100644 index 0000000000..1fd1738f93 --- /dev/null +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -0,0 +1,1022 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file group_quantize_fp8_blockwise.cuh + * \brief CUDA kernels to quantize grouped tensors with FP8 1D and 2D + * block scaling. A single launch walks 128x128 tiles across every tensor + * in the group, with each CTA decoding its owning tensor from the device-side + * GroupedTensor metadata. Supports SAME_BOTH_DIMS and VARYING_FIRST_DIM. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_BLOCKWISE_CUH_ +#define TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_BLOCKWISE_CUH_ + +#include +#include +#include +#include + +#include + +#include "../../common.h" +#include "../../recipe/recipe_common.cuh" +#include "../../transpose/cast_transpose.h" +#include "../../util/cuda_runtime.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "../core/common.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8_blockwise { + +using transformer_engine::detail::FP8BlockwiseColumnwiseOption; +using transformer_engine::detail::FP8BlockwiseRowwiseOption; + +constexpr int kTileDim = 128; +constexpr int kThreadsPerWarp = 32; +constexpr int kThreadsPerBlock = 256; +constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp; + +// ---- Per-expert scale layout helpers -------------------------------------------- +// +// cuBLAS grouped FP8 block-scaling GEMM expects each expert's scales to live +// in a contiguous per-expert sub-block of the global scale buffer: +// +// 1D rowwise CW (op_rowwise=true) per expert: (blocks_X, roundup(M_t, 4)) floats +// 1D columnwise (op_rowwise=false) per expert: (blocks_y_t, roundup(K, 4)) floats +// 2D rowwise (op_rowwise=true) per expert: (blocks_y_t, roundup(blocks_X, 4)) +// 2D columnwise (op_rowwise=false) per expert: (blocks_X, roundup(blocks_y_t, 4)) +// +// The grouped kernel writes the GLOBAL buffer in tile-stride order, but the +// position assigned to each tile must map into the per-expert contiguous +// sub-block. These helpers compute the per-expert cumulative byte/float +// offset and the per-expert local stride. SAME_BOTH_DIMS lets us derive +// without walking offsets; VARYING_FIRST_DIM walks `tensor_offsets_ptr` once +// per block (only the writing thread). + +__device__ __host__ constexpr size_t kScaleColAlign = 4; + +// 2D columnwise per-expert scale offset. Per-expert layout is +// (blocks_X, roundup(blocks_y_t, 4)). Two paths, picked at call sites: +// - SAME_BOTH_DIMS: direct formula (no walk, no reduction). +// - VARYING_FIRST_DIM: CTA-cooperative prefix sum. Each thread accumulates a +// partial over a strided subset of the tensors-before-this-one, a +// warp-shuffle reduces inside each warp, then all threads sum the per-warp +// partials read from a kNumWarps-element smem buffer to obtain the same +// total. The non-linear DIVUP_TO_MULTIPLE on each per-tensor blocks_y_t +// prevents a closed form. +// ALL threads of the CTA must call this in lock-step (the cooperative path +// contains __syncthreads()). +template +__device__ __forceinline__ size_t compute_2d_cw_expert_offset( + const size_t tensor_id, const size_t blocks_X, const size_t common_first_dim_blocks, + const size_t tile_row_stride, const int64_t* __restrict__ tensor_offsets_ptr, + size_t* warp_partials_smem, const int tid, const int warp_id, const int lane) { + if constexpr (kSameBothDims) { + return tensor_id * blocks_X * DIVUP_TO_MULTIPLE(common_first_dim_blocks, kScaleColAlign); + } + size_t my_partial = 0; + for (size_t i = static_cast(tid); i < tensor_id; + i += static_cast(kThreadsPerBlock)) { + const size_t blocks_y_i = + static_cast(tensor_offsets_ptr[i + 1] - tensor_offsets_ptr[i]) / tile_row_stride; + my_partial += DIVUP_TO_MULTIPLE(blocks_y_i, kScaleColAlign); + } + my_partial = warp_allreduce_sum(my_partial); + if (lane == 0) warp_partials_smem[warp_id] = my_partial; + __syncthreads(); + size_t total = 0; +#pragma unroll + for (int w = 0; w < kNumWarps; ++w) { + total += warp_partials_smem[w]; + } + return total * blocks_X; +} + +// 1D rowwise: per-expert layout (blocks_X, roundup(M_t, 4)). +template +__device__ __forceinline__ size_t expert_scale_offset_1d_rowwise( + size_t tensor_id, size_t blocks_X, size_t common_first_dim_blocks, size_t tile_row_stride, + const int64_t* __restrict__ tensor_offsets_ptr) { + if constexpr (kSameBothDims) { + const size_t M = common_first_dim_blocks * kTileDim; + return tensor_id * blocks_X * DIVUP_TO_MULTIPLE(M, kScaleColAlign); + } else { + // Each M_i is enforced to be a multiple of kTileDim (=128), hence a + // multiple of kScaleColAlign (=4), so DIVUP_TO_MULTIPLE(M_i, 4) == M_i and + // sum_{i(tensor_offsets_ptr[tensor_id]) / K; + return blocks_X * total_M_before; + } +} + +// ---- Tensor-lookup helpers ---------------------------------------------------- + +// Map a global tile-row index to its owning tensor. Delegates to the shared +// `common::get_current_tensor_id` helper from `cast/core/common.cuh`. The +// helper is parameterized by total `first_logical_dim` rather than per-tensor +// block count, so we reconstruct it here for the SAME_BOTH_DIMS specialization +// (VARYING_FIRST_DIM ignores it and uses `current_offset` + `offsets_ptr`). +template +__device__ __forceinline__ size_t find_tensor_id_by_block_y( + const size_t block_y_global, const size_t num_tensors, const size_t common_first_dim_blocks, + const size_t tile_row_stride, const int64_t* __restrict__ tensor_offsets_ptr) { + constexpr auto shape_rep = + kSameBothDims ? ShapeRepresentation::SAME_BOTH_DIMS : ShapeRepresentation::VARYING_FIRST_DIM; + const size_t first_logical_dim = num_tensors * common_first_dim_blocks * kTileDim; + const size_t tensor_id = common::get_current_tensor_id( + num_tensors, block_y_global * tile_row_stride, block_y_global, first_logical_dim, + /*last_logical_dim=*/0, tensor_offsets_ptr); + if constexpr (!kSameBothDims) { + // tensor_offsets_ptr carries cumulative element counts; tile_row_stride = + // kTileDim * K, so the per-tensor element span is divisible by + // tile_row_stride iff first_dim is a multiple of kTileDim. + if (tensor_id < num_tensors) { + const size_t span = + static_cast(tensor_offsets_ptr[tensor_id + 1] - tensor_offsets_ptr[tensor_id]); + if (span % tile_row_stride != 0) { + NVTE_DEVICE_ERROR( + "Grouped FP8 block-scaling quantize: each tensor's first dimension must be a " + "multiple of 128 (VARYING_FIRST_DIM)."); + } + } + } + return tensor_id; +} + +// Per-tensor block-y base for VARYING_FIRST_DIM (in 128-row block units). +__device__ __forceinline__ size_t tensor_block_y_base_from_offsets( + const size_t tensor_id, const int64_t* __restrict__ tensor_offsets_ptr, + const size_t tile_row_stride) { + return static_cast(tensor_offsets_ptr[tensor_id]) / tile_row_stride; +} + +// Per-vector amax. Uses bf16x2 `max.xorsign.abs` on sm_89+; FP32 fallback otherwise. +template +__device__ __forceinline__ CType compute_row_amax(const Vec& v) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) + if constexpr (std::is_same_v) { + static_assert(kVec % 2 == 0, "kVec must be even for packed bf16x2 amax"); + const ptx::bf16x2* pairs = reinterpret_cast(&v.data.elt[0]); + ptx::bf16x2 amax_x2{static_cast(0.f), static_cast(0.f)}; +#pragma unroll + for (int p = 0; p < kVec / 2; ++p) { + ptx::abs_max_2x(amax_x2, amax_x2, pairs[p]); + } + return static_cast(__hmax(__habs(amax_x2.x), __habs(amax_x2.y))); + } +#endif + CType amax = 0.f; +#pragma unroll + for (int e = 0; e < kVec; ++e) { + amax = fmaxf(amax, fabsf(static_cast(v.data.elt[e]))); + } + return amax; +} + +// Per-tile column sum of the high-precision input -> one fp32 row at +// dbias_workspace[tile_y_global * K + col]. 2 threads/column sum 64 rows each (combined via +// shfl_xor); grouped_reduce_dbias later sums each expert's row-blocks. Tiles are always full +// (experts are 128-row aligned), so all 128 rows are summed. +template +__device__ __forceinline__ void write_tile_dbias_partial(const IType smem_tile[][kTileDim], + const int tid, + const size_t global_col_base, + const size_t K, const size_t tile_y_global, + float* __restrict__ dbias_workspace) { + constexpr int kThreadsPerColDB = 2; + constexpr int kRowsPerThreadDB = kTileDim / kThreadsPerColDB; // 64 + const int col_local = tid / kThreadsPerColDB; // 0..127 + const int sub = tid % kThreadsPerColDB; + const int row_start = sub * kRowsPerThreadDB; + float partial = 0.f; +#pragma unroll + for (int e = 0; e < kRowsPerThreadDB; ++e) { + partial += static_cast(smem_tile[row_start + e][col_local]); + } + partial += __shfl_xor_sync(0xffffffff, partial, 1); + const size_t c_global = global_col_base + col_local; + if (sub == 0 && c_global < K) { + dbias_workspace[tile_y_global * K + c_global] = partial; + } +} + +// Per-vector multiply-and-quantize via fp32 intermediates. +template +__device__ __forceinline__ void quantize_row_vec(Vec& out, const Vec& in, + CType scale) { +#pragma unroll + for (int e = 0; e < kVec; ++e) { + out.data.elt[e] = static_cast(static_cast(in.data.elt[e]) * scale); + } +} + +// Bank-conflict swizzle delta for the 2D smem_T staging buffer. delta carries +// bits 2..5 only so each 4-byte sub-chunk is preserved. +__device__ __forceinline__ int smem_t_swz_delta(int smem_t_row) { + return ((smem_t_row >> 3) & 0xf) << 2; +} + +// Drain smem_T to gmem for the CW path: 4 cols/warp, 8 lanes/col, each lane +// stores a 16-row chunk so the 8 lanes of a col emit one 128 B gmem line. +// +// Columnwise data is stored per-expert contiguously as a (K, M_t) transposed +// block at element offset K * tensor_row_base, matching cuBLAS grouped GEMM's +// per-expert data pointer (compute_grouped_tensor_offset = sum_i M_i * K). +// `tensor_row_base` is expert t's first global row; `tensor_M` is M_t. +// +// When `kSwizzledStaging` is true, the writer applied smem_t_swz_delta to the +// inner index, so the read must XOR back to recover the logical (row, col). +// Per-byte unswizzled reads are required because the smem_T row stride is not +// 16 B aligned for arbitrary column offsets. +template +__device__ __forceinline__ void drain_smem_t_to_gmem( + OType (&smem_T)[kTileDim][kSMemTRowStride], OType* __restrict__ output_t_base, + const size_t global_col_base, const size_t global_row_base, const size_t tensor_row_base, + const size_t tensor_M, const size_t K, const int tid) { + constexpr int kStorePerChunk = 16; + constexpr int kRowChunksPerCol = kTileDim / kStorePerChunk; // 8 + constexpr int kColsPerIter = kThreadsPerBlock / kRowChunksPerCol; // 32 + constexpr int kColIters = kTileDim / kColsPerIter; // 4 + const int warp_id = tid / kThreadsPerWarp; + const int lane = tid % kThreadsPerWarp; + const int col_in_warp = lane / kRowChunksPerCol; + const int row_chunk = lane % kRowChunksPerCol; + const int out_row_off = row_chunk * kStorePerChunk; + const size_t expert_data_off = tensor_row_base * K; + const size_t m_local_base = global_row_base - tensor_row_base; +#pragma unroll + for (int it = 0; it < kColIters; ++it) { + const int out_col_local = it * kColsPerIter + warp_id * 4 + col_in_warp; + const size_t out_col_global = global_col_base + out_col_local; + if (out_col_global < K) { + // Per-expert (K, M_t): index = expert_off + k * M_t + m_local. + OType* out_ptr = + output_t_base + expert_data_off + out_col_global * tensor_M + m_local_base + out_row_off; + const int swz_delta_r = kSwizzledStaging ? smem_t_swz_delta(out_col_local) : 0; + Vec v; +#pragma unroll + for (int e = 0; e < kStorePerChunk; ++e) { + v.data.elt[e] = smem_T[out_col_local][(out_row_off + e) ^ swz_delta_r]; + } + v.store_to(out_ptr); + } + } +} + +// ----- 2D block scaling kernel with TMA input ------------------------------------------------ +// Pass 1: amax over a 128x128 TMA-loaded tile, with input vectors staged in +// registers. Pass 2: quantize from registers, emit rowwise output and the +// transposed smem_T tile, then drain smem_T to gmem. + +template +__global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma_kernel( + const __grid_constant__ CUtensorMap tensor_map_input, OType* __restrict__ output_base, + OType* __restrict__ output_t_base, CType* __restrict__ scale_inv_base, + CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, + const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, + const size_t total_row_blocks, const size_t blocks_X, const size_t scale_stride_y, + const float epsilon, const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { +#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; + + const size_t tile_x = blockIdx.x; + const size_t tile_y_global = blockIdx.y; + if (tile_y_global >= total_row_blocks) return; + + const size_t tile_row_stride = static_cast(kTileDim) * K; + const size_t tensor_id = find_tensor_id_by_block_y( + tile_y_global, num_tensors, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr); + const size_t tensor_block_y_base = + kSameBothDims + ? (tensor_id * common_first_dim_blocks) + : tensor_block_y_base_from_offsets(tensor_id, tensor_offsets_ptr, tile_row_stride); + const size_t tensor_row_blocks = + kSameBothDims + ? common_first_dim_blocks + : (tensor_block_y_base_from_offsets(tensor_id + 1, tensor_offsets_ptr, tile_row_stride) - + tensor_block_y_base); + if (tile_y_global >= tensor_block_y_base + tensor_row_blocks) return; + + const size_t global_row_base = tile_y_global * kTileDim; + const size_t global_col_base = tile_x * kTileDim; + + // Dynamic smem holds the IType input tile (TMA dest, must be 128 B aligned). + // warp_amaxes and tma_mbar are static smem. + extern __shared__ unsigned char smem_raw_2d_tma[]; + IType(*smem_in)[kTileDim] = reinterpret_cast( + common::align_smem_ptr_per_TMA_requirements(smem_raw_2d_tma)); + + __shared__ CType warp_amaxes[kNumWarps]; + __shared__ size_t warp_offset_partials[kNumWarps]; + __shared__ uint64_t tma_mbar; + + const int tid = threadIdx.x; + const bool leading_thread = (tid == 0); + + // ---- TMA async load of the input tile ---- + if (leading_thread) { + ptx::mbarrier_init(&tma_mbar, 1); + // Fence so the TMA engine (async proxy) and the threads (generic proxy) + // both observe the just-initialized mbarrier consistently. + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + if (leading_thread) { + constexpr uint32_t tx_bytes = kTileDim * kTileDim * sizeof(IType); + ptx::mbarrier_arrive_expect_tx(&tma_mbar, tx_bytes); + ptx::cp_async_bulk_tensor_2d_global_to_shared_cta( + reinterpret_cast(smem_in), reinterpret_cast(&tensor_map_input), + static_cast(global_col_base), static_cast(global_row_base), &tma_mbar); + } + ptx::mbarrier_wait_parity(&tma_mbar, 0); + if (leading_thread) ptx::mbarrier_invalid(&tma_mbar); + __syncthreads(); + + // ---- Optional dbias: per-tile column sum of the high-precision input (smem-resident) ---- + if (dbias_workspace != nullptr) { + write_tile_dbias_partial(smem_in, tid, global_col_base, K, tile_y_global, + dbias_workspace); + } + + // ---- Pass 1: tile amax, staging input vectors in registers for reuse in pass 2 ---- + constexpr int kEltsPerThread = 8; + constexpr int kThreadsPerRow = kTileDim / kEltsPerThread; // 16 + constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow; // 16 + constexpr int kIters = kTileDim / kRowsPerIter; // 8 + + const int thr_col = tid % kThreadsPerRow; + const int thr_row = tid / kThreadsPerRow; + + using IVec = Vec; + using OVec = Vec; + + IVec staged[kIters]; + CType thr_amax = 0.f; +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int r_local = thr_row + it * kRowsPerIter; + staged[it].load_from(&smem_in[r_local][thr_col * kEltsPerThread]); + thr_amax = fmaxf(thr_amax, compute_row_amax(staged[it])); + } + CType warp_amax = warp_reduce_max(thr_amax); + const int warp_id = tid / kThreadsPerWarp; + const int lane = tid % kThreadsPerWarp; + if (lane == 0) warp_amaxes[warp_id] = warp_amax; + __syncthreads(); + + CType block_amax = warp_amaxes[0]; +#pragma unroll + for (int w = 1; w < kNumWarps; ++w) { + block_amax = fmaxf(block_amax, warp_amaxes[w]); + } + const CType scale = + compute_scale_from_types(block_amax, epsilon, /*pow_2_scaling=*/false); + + // The 2D colwise per-expert scale offset requires a CTA-cooperative prefix + // sum in the VARYING_FIRST_DIM case, so compute it across all threads before + // the leading-thread-only store. All threads end up with the same value; + // only thread 0 reads it below. + size_t expert_offset_t = 0; + if constexpr (kReturnColwise) { + expert_offset_t = compute_2d_cw_expert_offset( + tensor_id, blocks_X, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr, + warp_offset_partials, tid, warp_id, lane); + } + + if (leading_thread) { + const CType scale_inv = 1.f / scale; + if constexpr (kReturnRowwise) { + // 2D rowwise: kernel-global stride matches per-expert layout naturally. + // Expert t's rows occupy [tensor_block_y_base, +blocks_y_t) of the buffer, + // each row sized roundup(blocks_X, 4), which equals the dispatcher's + // cumulative per-expert offset. + scale_inv_base[tile_y_global * scale_stride_y + tile_x] = scale_inv; + } + if constexpr (kReturnColwise) { + // 2D colwise: rewrite into per-expert sub-block matching cuBLAS grouped + // GEMM's per-expert layout (blocks_X, roundup(blocks_y_t, 4)). + const size_t local_tile_y = tile_y_global - tensor_block_y_base; + const size_t per_expert_stride_t = DIVUP_TO_MULTIPLE(tensor_row_blocks, kScaleColAlign); + scale_inv_t_base[expert_offset_t + tile_x * per_expert_stride_t + local_tile_y] = scale_inv; + } + } + + // ---- Pass 2: quantize from register-staged inputs, emit rowwise + colwise outputs ---- + // 2D block-scaling uses a per-tile (128x128) scalar scale, so the row-wise and + // column-wise quantized bytes are identical -- only the gmem layout differs. The + // columnwise buffer is physically transposed (cuBLAS FP8 block-scaling GEMM is + // TN-only), so we stage into smem_T and drain to the per-expert (K, M_t) transposed block. + if constexpr (kReturnColwise) { + constexpr int kSMemTRowStride = kTileDim + 4; + __shared__ OType smem_T[kTileDim][kSMemTRowStride]; + // Same delta for all 8 elements this thread writes since (thr_col*8 + e) >> 3 == thr_col. + const int swz_delta_w = smem_t_swz_delta(thr_col * kEltsPerThread); +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int row_local = thr_row + it * kRowsPerIter; + const size_t r = global_row_base + row_local; + const size_t c = global_col_base + thr_col * kEltsPerThread; + OVec qo; + quantize_row_vec(qo, staged[it], scale); + if constexpr (kReturnRowwise) { + if (c < K) { + const size_t count = (c + kEltsPerThread <= K) ? kEltsPerThread : (K - c); + qo.store_to_elts(output_base + r * K + c, 0, count); + } + } + const int c_phys = row_local ^ swz_delta_w; +#pragma unroll + for (int e = 0; e < kEltsPerThread; ++e) { + smem_T[thr_col * kEltsPerThread + e][c_phys] = qo.data.elt[e]; + } + } + __syncthreads(); + + const size_t tensor_row_base = tensor_block_y_base * kTileDim; + const size_t tensor_M = tensor_row_blocks * kTileDim; + drain_smem_t_to_gmem( + smem_T, output_t_base, global_col_base, global_row_base, tensor_row_base, tensor_M, K, tid); + } else if constexpr (kReturnRowwise) { +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int row_local = thr_row + it * kRowsPerIter; + const size_t r = global_row_base + row_local; + const size_t c = global_col_base + thr_col * kEltsPerThread; + OVec qo; + quantize_row_vec(qo, staged[it], scale); + if (c < K) { + const size_t count = (c + kEltsPerThread <= K) ? kEltsPerThread : (K - c); + qo.store_to_elts(output_base + r * K + c, 0, count); + } + } + } +#endif // __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 +} + +// ----- 1D block scaling rowwise-only kernel ---------------------------------------------------- +// No smem cache. Each thread loads 16 cols/row, reduces amax across the 8 +// row-mates with shfl_xor, then quantizes and stores. + +template +__global__ void __launch_bounds__(kThreadsPerBlock) + group_block_scaled_1d_rw_kernel(const IType* __restrict__ input_base, + OType* __restrict__ output_base, + CType* __restrict__ scale_inv_base, + const int64_t* __restrict__ tensor_offsets_ptr, + const size_t num_tensors, const size_t common_first_dim_blocks, + const size_t K, const size_t total_row_blocks, + const size_t R_total, const float epsilon, + const float* __restrict__ noop_ptr) { +#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; + + const size_t tile_x = blockIdx.x; + const size_t tile_y_global = blockIdx.y; + if (tile_y_global >= total_row_blocks) return; + + const size_t tile_row_stride = static_cast(kTileDim) * K; + const size_t tensor_id = find_tensor_id_by_block_y( + tile_y_global, num_tensors, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr); + const size_t tensor_block_y_base = + kSameBothDims + ? (tensor_id * common_first_dim_blocks) + : tensor_block_y_base_from_offsets(tensor_id, tensor_offsets_ptr, tile_row_stride); + const size_t tensor_row_blocks = + kSameBothDims + ? common_first_dim_blocks + : (tensor_block_y_base_from_offsets(tensor_id + 1, tensor_offsets_ptr, tile_row_stride) - + tensor_block_y_base); + if (tile_y_global >= tensor_block_y_base + tensor_row_blocks) return; + + const size_t global_row_base = tile_y_global * kTileDim; + const size_t global_col_base = tile_x * kTileDim; + + // 8 threads per row x 16 cols/thread = one 128 B gmem cache line per row, 4 iters per tile. + constexpr int kThreadsPerRow = 8; + constexpr int kVec = 16; + constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow; // 32 + constexpr int kIters = kTileDim / kRowsPerIter; // 4 + + const int tid = threadIdx.x; + const int thr_col = tid % kThreadsPerRow; // 0..7 + const int thr_row = tid / kThreadsPerRow; // 0..31 (row index within an iter) + const size_t c = global_col_base + static_cast(thr_col) * kVec; + + Vec in_vec[kIters]; + +#pragma unroll + for (int it = 0; it < kIters; ++it) { + const int row_local = thr_row + it * kRowsPerIter; + const size_t r_global = global_row_base + row_local; + + // Load this thread's 16 cols of row `row_local`. + if (c + kVec <= K) { + in_vec[it].load_from(input_base + r_global * K + c); + } else if (c < K) { + in_vec[it].load_from_elts(input_base + r_global * K + c, 0, K - c); + } else { + in_vec[it].clear(); + } + + CType amax = compute_row_amax(in_vec[it]); + amax = subwarp_reduce_max_broadcast(amax); + + const CType scale = + compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale_inv = 1.f / scale; + if (thr_col == 0 && r_global < R_total) { + // Per-expert layout: (blocks_X, roundup(M_t, 4)). Compute expert base + // offset + local stride matching cuBLAS grouped GEMM's per-expert view. + const size_t blocks_X = DIVUP(K, static_cast(kTileDim)); + const size_t expert_offset = expert_scale_offset_1d_rowwise( + tensor_id, blocks_X, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr); + const size_t tensor_M = tensor_row_blocks * kTileDim; + const size_t per_expert_stride = DIVUP_TO_MULTIPLE(tensor_M, kScaleColAlign); + const size_t tensor_row_base = tensor_block_y_base * kTileDim; + const size_t r_local = r_global - tensor_row_base; + scale_inv_base[expert_offset + tile_x * per_expert_stride + r_local] = scale_inv; + } + + if (r_global < R_total) { + Vec out_vec; + quantize_row_vec(out_vec, in_vec[it], scale); + if (c + kVec <= K) { + out_vec.store_to(output_base + r_global * K + c); + } else if (c < K) { + out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c); + } + } + } +#endif // __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 +} + +// ----- 1D block scaling kernel with TMA input ------------------------------------------------ +// CW and BOTH path. TMA fills a 128x128 smem input cache. RW pass reads rows +// and stores quantized output. CW pass stages a column slice in registers, +// computes the per-column amax there, then fills smem_T and drains to gmem. + +template +__global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_kernel( + const __grid_constant__ CUtensorMap tensor_map_input, OType* __restrict__ output_base, + OType* __restrict__ output_t_base, CType* __restrict__ scale_inv_base, + CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, + const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, + const size_t total_row_blocks, const size_t blocks_X, const size_t scale_t_stride_aligned_K, + const size_t R_total, const float epsilon, const float* __restrict__ noop_ptr, + float* __restrict__ dbias_workspace) { +#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; + + const size_t tile_x = blockIdx.x; + const size_t tile_y_global = blockIdx.y; + if (tile_y_global >= total_row_blocks) return; + + const size_t tile_row_stride = static_cast(kTileDim) * K; + const size_t tensor_id = find_tensor_id_by_block_y( + tile_y_global, num_tensors, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr); + const size_t tensor_block_y_base = + kSameBothDims + ? (tensor_id * common_first_dim_blocks) + : tensor_block_y_base_from_offsets(tensor_id, tensor_offsets_ptr, tile_row_stride); + const size_t tensor_row_blocks = + kSameBothDims + ? common_first_dim_blocks + : (tensor_block_y_base_from_offsets(tensor_id + 1, tensor_offsets_ptr, tile_row_stride) - + tensor_block_y_base); + if (tile_y_global >= tensor_block_y_base + tensor_row_blocks) return; + + const size_t global_row_base = tile_y_global * kTileDim; + const size_t global_col_base = tile_x * kTileDim; + + // Dynamic smem: IType[kTileDim][kTileDim], 128 B aligned for TMA. Static smem + // (smem_T when CW, tma_mbar) lives outside the dynamic region. + extern __shared__ unsigned char smem_raw_1d_tma[]; + unsigned char* smem_base = common::align_smem_ptr_per_TMA_requirements(smem_raw_1d_tma); + IType(*smem)[kTileDim] = reinterpret_cast(smem_base); + + __shared__ uint64_t tma_mbar; + const int tid = threadIdx.x; + const bool leading_thread = (tid == 0); + + // ---- TMA async load of the input tile ---- + if (leading_thread) { + ptx::mbarrier_init(&tma_mbar, 1); + // Fence so the TMA engine (async proxy) and the threads (generic proxy) + // both observe the just-initialized mbarrier consistently. + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + if (leading_thread) { + constexpr uint32_t tx_bytes = kTileDim * kTileDim * sizeof(IType); + ptx::mbarrier_arrive_expect_tx(&tma_mbar, tx_bytes); + ptx::cp_async_bulk_tensor_2d_global_to_shared_cta( + reinterpret_cast(smem_base), + reinterpret_cast(&tensor_map_input), + static_cast(global_col_base), static_cast(global_row_base), &tma_mbar); + } + ptx::mbarrier_wait_parity(&tma_mbar, 0); + if (leading_thread) ptx::mbarrier_invalid(&tma_mbar); + __syncthreads(); + + // ---- Optional dbias: per-tile column sum of the high-precision input (smem-resident) ---- + if (dbias_workspace != nullptr) { + write_tile_dbias_partial(smem, tid, global_col_base, K, tile_y_global, dbias_workspace); + } + + // ---- RW pass (1x128 scale per row) ---- + // 8 t/row, vec-16 reads from smem; emits rowwise gmem directly. Only entered + // when CW is also requested (BOTH) -- RW-only requests use the dedicated + // group_block_scaled_1d_rw_kernel which skips the TMA load. + if constexpr (kReturnRowwise) { + constexpr int kThreadsPerRowRW = 8; + constexpr int kVec = 16; + constexpr int kRowsPerIterRW = kThreadsPerBlock / kThreadsPerRowRW; // 32 + constexpr int kRwIters = kTileDim / kRowsPerIterRW; // 4 + + const int rw_thr_col = tid % kThreadsPerRowRW; + const int rw_thr_row = tid / kThreadsPerRowRW; + const int col_local = rw_thr_col * kVec; + +#pragma unroll + for (int it = 0; it < kRwIters; ++it) { + const int row_local = rw_thr_row + it * kRowsPerIterRW; + Vec in_vec; + in_vec.load_from(&smem[row_local][col_local]); + + CType amax = compute_row_amax(in_vec); + amax = subwarp_reduce_max_broadcast(amax); + + const CType scale = + compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale_inv = 1.f / scale; + + const size_t r_global = global_row_base + row_local; + const bool row_in_bounds = (r_global < R_total); + + if (row_in_bounds && rw_thr_col == 0) { + // Per-expert layout matches the 1D RW kernel. + const size_t blocks_X_eff = DIVUP(K, static_cast(kTileDim)); + const size_t expert_offset = expert_scale_offset_1d_rowwise( + tensor_id, blocks_X_eff, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr); + const size_t tensor_M = tensor_row_blocks * kTileDim; + const size_t per_expert_stride = DIVUP_TO_MULTIPLE(tensor_M, kScaleColAlign); + const size_t tensor_row_base = tensor_block_y_base * kTileDim; + const size_t r_local = r_global - tensor_row_base; + scale_inv_base[expert_offset + tile_x * per_expert_stride + r_local] = scale_inv; + } + + if (row_in_bounds) { + const size_t cc = global_col_base + col_local; + Vec out_vec; + quantize_row_vec(out_vec, in_vec, scale); + if (cc + kVec <= K) { + out_vec.store_to(output_base + r_global * K + cc); + } else if (cc < K) { + out_vec.store_to_elts(output_base + r_global * K + cc, 0, K - cc); + } + } + } + } + + if constexpr (kReturnRowwise && kReturnColwise) { + __syncthreads(); + } + + // ---- CW pass (128x1 scale per column) ---- + // CW-only: 2 t/col, single column pass per CTA, 64-row reg_data per thread. + // BOTH: 4 t/col, two column passes per CTA, 32-row reg_data per thread. + // + // In BOTH, the RW pass's per-expert offset arithmetic raises register + // footprint past the 85-reg/thread threshold for 3 CTAs/SM on sm_90, so we + // halve the reg_data stage (and accept the extra XOR-reduce stage plus + // doubled column pass count) to recover that occupancy. CW-only is not + // occupancy-bound, so its 2 t/col path stays — 4 t/col increases the + // bank-conflict footprint on the smem load, which costs more than the extra + // CTA/SM gains. + // + // The columnwise buffer is physically transposed (cuBLAS FP8 block-scaling + // GEMM is TN-only); we stage in smem_T and drain to the per-expert (K, M_t) + // transposed block. + if constexpr (kReturnColwise) { + constexpr int kSMemTRowStride = kTileDim + 4; + __shared__ OType smem_T[kTileDim][kSMemTRowStride]; + + constexpr int kThreadsPerColCW = kReturnRowwise ? 4 : 2; + constexpr int kRowsPerThreadCW = kTileDim / kThreadsPerColCW; // 32 (BOTH) / 64 (CW) + constexpr int kColsPerCWPass = kThreadsPerBlock / kThreadsPerColCW; // 64 (BOTH) / 128 (CW) + constexpr int kCWPasses = kTileDim / kColsPerCWPass; // 2 (BOTH) / 1 (CW) + + const int sub = tid % kThreadsPerColCW; + const int col_in_pass = tid / kThreadsPerColCW; + const int row_start = sub * kRowsPerThreadCW; + +#pragma unroll + for (int pass = 0; pass < kCWPasses; ++pass) { + const int col_local = pass * kColsPerCWPass + col_in_pass; + + CType reg_data[kRowsPerThreadCW]; + CType amax = 0.f; +#pragma unroll + for (int e = 0; e < kRowsPerThreadCW; ++e) { + reg_data[e] = static_cast(smem[row_start + e][col_local]); + amax = fmaxf(amax, fabsf(reg_data[e])); + } + amax = subwarp_reduce_max_broadcast(amax); + + const CType scale = + compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale_inv = 1.f / scale; + + const size_t c_global = global_col_base + col_local; + const bool col_in_bounds = (c_global < K); + + if (col_in_bounds && sub == 0) { + scale_inv_t_base[c_global + tile_y_global * scale_t_stride_aligned_K] = scale_inv; + } + + if (col_in_bounds) { +#pragma unroll + for (int e = 0; e < kRowsPerThreadCW; ++e) { + smem_T[col_local][row_start + e] = static_cast(reg_data[e] * scale); + } + } + } + __syncthreads(); + + const size_t tensor_row_base = tensor_block_y_base * kTileDim; + const size_t tensor_M = tensor_row_blocks * kTileDim; + drain_smem_t_to_gmem( + smem_T, output_t_base, global_col_base, global_row_base, tensor_row_base, tensor_M, K, tid); + } +#endif // __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 +} + +// ----- Host-side dispatchers -------------------------------------------------------------------- + +struct GroupedBlockwiseLaunchInfo { + size_t num_tensors; + size_t K; + size_t R_total; + size_t common_first_dim_blocks; + size_t total_row_blocks; + size_t blocks_X; + bool same_both_dims; + const int64_t* tensor_offsets_d = nullptr; +}; + +inline GroupedBlockwiseLaunchInfo prepare_grouped_blockwise_launch(const GroupedTensor* output) { + GroupedBlockwiseLaunchInfo info{}; + const bool same_both_dims = output->all_same_shape(); + const bool varying_first_dim = (!output->all_same_first_dim()) && output->all_same_last_dim(); + NVTE_CHECK(same_both_dims || varying_first_dim, + "Grouped FP8 block-scaling supports only SAME_BOTH_DIMS and VARYING_FIRST_DIM " + "shape representations."); + + info.same_both_dims = same_both_dims; + info.num_tensors = output->num_tensors; + info.K = output->get_common_last_dim(); + NVTE_CHECK(info.K % TMA_GMEM_ALIGNMENT == 0, + "Last dim must be a multiple of TMA_GMEM_ALIGNMENT (", TMA_GMEM_ALIGNMENT, + ") for FP8 alignment."); + + if (same_both_dims) { + const size_t common_first_dim = output->get_common_first_dim(); + NVTE_CHECK(common_first_dim % kTileDim == 0, + "SAME_BOTH_DIMS first dim must be multiple of 128."); + info.common_first_dim_blocks = common_first_dim / kTileDim; + info.R_total = info.num_tensors * common_first_dim; + } else { + info.common_first_dim_blocks = 0; + info.R_total = output->logical_shape.data[0]; + info.tensor_offsets_d = reinterpret_cast(output->tensor_offsets.dptr); + NVTE_CHECK(info.tensor_offsets_d != nullptr, + "VARYING_FIRST_DIM requires tensor_offsets to be set on the GroupedTensor."); + } + info.total_row_blocks = DIVUP(info.R_total, static_cast(kTileDim)); + info.blocks_X = DIVUP(info.K, static_cast(kTileDim)); + return info; +} + +// Public dispatch — 2D block scaling. +// +// When `dbias` is non-null, also accumulate the bias gradient into `workspace` (reduced per +// expert by grouped_reduce_dbias). Two-call protocol: the sizing pass (workspace unallocated) +// reports the [total_row_blocks, K] fp32 shape and returns without launching. +inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTensor* output, + const Tensor* noop, const float epsilon, + cudaStream_t stream, GroupedTensor* dbias = nullptr, + Tensor* workspace = nullptr) { + const int sm = transformer_engine::cuda::sm_arch(); + NVTE_CHECK(sm >= 90 && sm < 100, + "Grouped FP8 block-scaling quantize is only supported on Hopper (SM90-SM99); " + "use MXFP8 on Blackwell (SM100) or newer. Got SM", + sm, "."); + const bool use_rowwise = output->has_data(); + const bool use_colwise = output->has_columnwise_data(); + NVTE_CHECK(use_rowwise || use_colwise, + "Either rowwise or columnwise output data must be allocated."); + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must be FP8."); + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Input and output must have same num_tensors."); + + auto info = prepare_grouped_blockwise_launch(output); + if (info.R_total == 0 || info.K == 0) return; + + float* dbias_workspace = nullptr; + if (dbias != nullptr) { + NVTE_CHECK(workspace != nullptr, "Workspace required for grouped FP8 block-scaling dbias."); + NVTE_CHECK(dbias->dtype() == input->dtype(), "dbias must have the same dtype as the input."); + if (workspace->data.dptr == nullptr) { + workspace->data.shape = {info.total_row_blocks, info.K}; + workspace->data.dtype = DType::kFloat32; + return; // sizing pass + } + dbias_workspace = reinterpret_cast(workspace->data.dptr); + } + + using CType = float; + const float* noop_ptr = + (noop != nullptr) ? reinterpret_cast(noop->data.dptr) : nullptr; + + const size_t scale_stride_y = DIVUP_TO_MULTIPLE(info.blocks_X, 4); + + dim3 grid(info.blocks_X, info.total_row_blocks, 1); + + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + info.same_both_dims, kSameBothDims, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_rowwise, kRowwise, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_colwise, kColwise, if constexpr (kRowwise || kColwise) { + CUtensorMap tensor_map_input{}; + create_2D_tensor_map(tensor_map_input, input->data, info.R_total, info.K, + kTileDim, kTileDim, info.K, 0, sizeof(IType) * 8); + auto tma_kernel = + group_block_scaled_2d_tma_kernel; + const size_t smem_bytes = + kTileDim * kTileDim * sizeof(IType) + TMA_SHMEM_ALIGNMENT - 1; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + tma_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_bytes))); + tma_kernel<<>>( + tensor_map_input, + kRowwise ? reinterpret_cast(output->data.dptr) : nullptr, + kColwise ? reinterpret_cast(output->columnwise_data.dptr) + : nullptr, + kRowwise ? reinterpret_cast(output->scale_inv.dptr) : nullptr, + kColwise ? reinterpret_cast(output->columnwise_scale_inv.dptr) + : nullptr, + info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, + info.K, info.total_row_blocks, info.blocks_X, scale_stride_y, epsilon, + noop_ptr, dbias_workspace); + if (dbias_workspace != nullptr) { + const ShapeRepresentation shape_rep = + info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS + : ShapeRepresentation::VARYING_FIRST_DIM; + common::grouped_reduce_dbias( + shape_rep, info.num_tensors, info.R_total, info.K, + reinterpret_cast(output->tensor_offsets.dptr), + reinterpret_cast(output->first_dims.dptr), + reinterpret_cast(output->last_dims.dptr), dbias, + dbias_workspace, kTileDim, stream); + } + }))))); + + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// Public dispatch — 1D block scaling. +// +// See group_quantize_blockwise_2d for dbias/workspace semantics. RW-only without dbias uses the +// no-smem fast path; RW-only with dbias routes through the TMA kernel (input in smem) so the +// per-tile column partial can be computed. +inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTensor* output, + const Tensor* noop, const float epsilon, + cudaStream_t stream, GroupedTensor* dbias = nullptr, + Tensor* workspace = nullptr) { + const int sm = transformer_engine::cuda::sm_arch(); + NVTE_CHECK(sm >= 90 && sm < 100, + "Grouped FP8 block-scaling quantize is only supported on Hopper (SM90-SM99); " + "use MXFP8 on Blackwell (SM100) or newer. Got SM", + sm, "."); + const bool use_rowwise = output->has_data(); + const bool use_colwise = output->has_columnwise_data(); + NVTE_CHECK(use_rowwise || use_colwise, + "Either rowwise or columnwise output data must be allocated."); + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must be FP8."); + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Input and output must have same num_tensors."); + + auto info = prepare_grouped_blockwise_launch(output); + if (info.R_total == 0 || info.K == 0) return; + + float* dbias_workspace = nullptr; + if (dbias != nullptr) { + NVTE_CHECK(workspace != nullptr, "Workspace required for grouped FP8 block-scaling dbias."); + NVTE_CHECK(dbias->dtype() == input->dtype(), "dbias must have the same dtype as the input."); + if (workspace->data.dptr == nullptr) { + workspace->data.shape = {info.total_row_blocks, info.K}; + workspace->data.dtype = DType::kFloat32; + return; // sizing pass + } + dbias_workspace = reinterpret_cast(workspace->data.dptr); + } + + using CType = float; + const float* noop_ptr = + (noop != nullptr) ? reinterpret_cast(noop->data.dptr) : nullptr; + + const size_t scale_t_stride_aligned_K = DIVUP_TO_MULTIPLE(info.K, 4); + + dim3 grid(info.blocks_X, info.total_row_blocks, 1); + + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + info.same_both_dims, kSameBothDims, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_rowwise, kRowwise, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_colwise, kColwise, + // RW-only without dbias uses the no-smem fast path; all else uses TMA. + constexpr bool kRwOnly = kRowwise && !kColwise; + const bool use_rw_fast_path = kRwOnly && (dbias_workspace == nullptr); + if (use_rw_fast_path) { + if constexpr (kRwOnly) { + group_block_scaled_1d_rw_kernel + <<>>( + reinterpret_cast(input->data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->scale_inv.dptr), + info.tensor_offsets_d, info.num_tensors, + info.common_first_dim_blocks, info.K, info.total_row_blocks, + info.R_total, epsilon, noop_ptr); + } + } else if constexpr (kRowwise || kColwise) { + // CW-only, BOTH, or RW-only WITH dbias: smem-cached TMA kernel. + const size_t smem_bytes = kTileDim * kTileDim * sizeof(IType); + constexpr size_t kStaticSmemCWBytes = + (kTileDim * (kTileDim + 4)) * sizeof(OType); + const size_t static_smem_bytes = kColwise ? kStaticSmemCWBytes : 0; + const size_t tma_smem_bytes = smem_bytes + TMA_SHMEM_ALIGNMENT - 1; + const size_t total_smem_tma = tma_smem_bytes + static_smem_bytes; + auto tma_kernel = + group_block_scaled_1d_tma_kernel; + if (total_smem_tma >= 48 * 1024) { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + tma_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(tma_smem_bytes))); + } + CUtensorMap tensor_map_input{}; + create_2D_tensor_map(tensor_map_input, input->data, info.R_total, info.K, + kTileDim, kTileDim, info.K, 0, sizeof(IType) * 8); + tma_kernel<<>>( + tensor_map_input, + kRowwise ? reinterpret_cast(output->data.dptr) : nullptr, + kColwise ? reinterpret_cast(output->columnwise_data.dptr) + : nullptr, + kRowwise ? reinterpret_cast(output->scale_inv.dptr) : nullptr, + kColwise ? reinterpret_cast(output->columnwise_scale_inv.dptr) + : nullptr, + info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, + info.K, info.total_row_blocks, info.blocks_X, scale_t_stride_aligned_K, + info.R_total, epsilon, noop_ptr, dbias_workspace); + if (dbias_workspace != nullptr) { + const ShapeRepresentation shape_rep = + info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS + : ShapeRepresentation::VARYING_FIRST_DIM; + common::grouped_reduce_dbias( + shape_rep, info.num_tensors, info.R_total, info.K, + reinterpret_cast(output->tensor_offsets.dptr), + reinterpret_cast(output->first_dims.dptr), + reinterpret_cast(output->last_dims.dptr), dbias, + dbias_workspace, kTileDim, stream); + } + }))))); + + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace fp8_blockwise +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_BLOCKWISE_CUH_ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 88a57fe989..2814aa3490 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -128,22 +128,22 @@ constexpr bool is_supported_arch() { // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-init __device__ __forceinline__ void mbarrier_init(uint64_t *mbar, const uint32_t count) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); asm volatile("mbarrier.init.shared.b64 [%0], %1;" ::"r"(mbar_ptr), "r"(count) : "memory"); #else - NVTE_DEVICE_ERROR("mbarrier_init is only supported on SM 10.0+."); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + NVTE_DEVICE_ERROR("mbarrier_init is only supported on SM 9.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-inval __device__ __forceinline__ void mbarrier_invalid(uint64_t *mbar) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); asm volatile("mbarrier.inval.shared.b64 [%0];" ::"r"(mbar_ptr) : "memory"); #else - NVTE_DEVICE_ERROR("mbarrier_invalid is only supported on SM 10.0+."); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + NVTE_DEVICE_ERROR("mbarrier_invalid is only supported on SM 9.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive @@ -158,13 +158,13 @@ __device__ __forceinline__ void mbarrier_arrive(uint64_t *mbar) { // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive __device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t *mbar, const uint32_t tx_count) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(mbar_ptr), "r"(tx_count) : "memory"); #else - NVTE_DEVICE_ERROR("mbarrier_arrive_expect_tx is only supported on SM 10.0+."); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + NVTE_DEVICE_ERROR("mbarrier_arrive_expect_tx is only supported on SM 9.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } __device__ __forceinline__ void mbarrier_arrive_expect_tx_cta_relaxed_shared_cta( @@ -230,8 +230,26 @@ __device__ __forceinline__ void cp_async_bulk_tensor_2d_global_to_shared( #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +// global -> shared::cta (no cluster; valid on Hopper sm_90+ and Blackwell with +// cluster size 1). Used by the FP8 block-scaling grouped quantize kernels. +__device__ __forceinline__ void cp_async_bulk_tensor_2d_global_to_shared_cta( + uint64_t *dst_shmem, const uint64_t *tensor_map_ptr, const uint32_t offset_x, + const uint32_t offset_y, uint64_t *mbar) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + uint32_t dst_shmem_ptr = __cvta_generic_to_shared(dst_shmem); + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile" + ".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::"r"(dst_shmem_ptr), + "l"(tensor_map_ptr), "r"(offset_x), "r"(offset_y), "r"(mbar_ptr) + : "memory"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_tensor_2d_global_to_shared_cta is only supported on SM 9.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + __device__ __forceinline__ bool mbarrier_try_wait_parity(uint32_t mbar_ptr, const uint32_t parity) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) uint32_t waitComplete; asm volatile( "{\n\t .reg .pred P_OUT; \n\t" @@ -243,19 +261,19 @@ __device__ __forceinline__ bool mbarrier_try_wait_parity(uint32_t mbar_ptr, cons : "memory"); return static_cast(waitComplete); #else - NVTE_DEVICE_ERROR("mbarrier_try_wait_parity is only supported on SM 10.0+."); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + NVTE_DEVICE_ERROR("mbarrier_try_wait_parity is only supported on SM 9.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) return true; } __device__ __forceinline__ void mbarrier_wait_parity(uint64_t *mbar, const uint32_t parity) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); while (!mbarrier_try_wait_parity(mbar_ptr, parity)) { } #else - NVTE_DEVICE_ERROR("mbarrier_wait_parity is only supported on SM 10.0+."); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + NVTE_DEVICE_ERROR("mbarrier_wait_parity is only supported on SM 9.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } __device__ __forceinline__ void mbarrier_wait_parity_acquire_cta_shared_cta(uint64_t *mbar, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index f57bd8797e..8d77a9e349 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -316,6 +316,7 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const FP8_CURRENT_SCALING_GROUPED_QUANTIZE, MXFP8_GROUPED_QUANTIZE, NVFP4_GROUPED_QUANTIZE, + FP8_BLOCKWISE_GROUPED_QUANTIZE, INVALID_FOR_GROUPED_QUANTIZE }; GroupedQuantizationMode grouped_quantization_mode = @@ -326,6 +327,8 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const grouped_quantization_mode = GroupedQuantizationMode::NVFP4_GROUPED_QUANTIZE; } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr())) { grouped_quantization_mode = GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE; + } else if (detail::IsFloat8BlockwiseQuantizers(quantizer.ptr())) { + grouped_quantization_mode = GroupedQuantizationMode::FP8_BLOCKWISE_GROUPED_QUANTIZE; } if (empty_input_buffer) { @@ -371,11 +374,23 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const }); break; } + case GroupedQuantizationMode::FP8_BLOCKWISE_GROUPED_QUANTIZE: { + Float8BlockQuantizer *fp8_block_quantizer_cpp = + static_cast(quantizer_cpp.get()); + QuantizationConfigWrapper quant_config_cpp; + quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); + quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), + quant_config_cpp, at::cuda::getCurrentCUDAStream()); + }); + break; + } case GroupedQuantizationMode::INVALID_FOR_GROUPED_QUANTIZE: default: NVTE_ERROR( - "group_quantize: only supports MXFP8, NVFP4, or " - "Float8CurrentScalingQuantizer."); + "group_quantize: only supports MXFP8, NVFP4, Float8CurrentScalingQuantizer, or " + "Float8Blockwise quantizer."); break; } @@ -479,8 +494,9 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0; - NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()), - "bgrad_group_quantize: only MXFP8 quantizer is supported."); + NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()) || + detail::IsFloat8BlockwiseQuantizers(quantizer.ptr()), + "bgrad_group_quantize: only MXFP8 and FP8 block-scaling quantizers are supported."); auto quantizer_cpp = convert_quantizer(quantizer); @@ -587,12 +603,19 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o // Build input GroupedTensorWrapper. // Data tensors are stored as flat 1D buffers; use the quantizer's dtype // (e.g. kFloat8E4M3) rather than the raw tensor scalar_type (uint8). - auto input_cpp = GroupedTensorWrapper(num_tensors, logical_shape, quantizer->get_scaling_mode()); + const NVTEScalingMode scaling_mode = quantizer->get_scaling_mode(); + const bool is_block_scaling = + (scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D); + const bool is_nvfp4 = (scaling_mode == NVTE_NVFP4_1D_SCALING); + const DType scale_dtype = is_block_scaling ? DType::kFloat32 + : is_nvfp4 ? DType::kFloat8E4M3 + : DType::kFloat8E8M0; + auto input_cpp = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode); if (rowwise_data.has_value()) { input_cpp.set_rowwise_data(rowwise_data->data_ptr(), quantizer->dtype, std::vector{static_cast(rowwise_data->numel())}); if (rowwise_scale_inv.has_value()) { - input_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E8M0, + input_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), scale_dtype, getTensorShape(*rowwise_scale_inv)); } } @@ -601,7 +624,7 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o columnwise_data->data_ptr(), quantizer->dtype, std::vector{static_cast(columnwise_data->numel())}); if (columnwise_scale_inv.has_value()) { - input_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E8M0, + input_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), scale_dtype, getTensorShape(*columnwise_scale_inv)); } } diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index a39d0143a0..d2a3be888f 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1154,6 +1154,14 @@ std::pair Float8BlockQuantizer::create_grouped const size_t logical_last_dim) const { using namespace pybind11::literals; + // The fused grouped FP8 block-scaling path uses unconstrained FP32 scales and does not + // implement power-of-2 scaling. Reject force_pow_2_scales rather than silently ignoring it; + // the unfused per-tensor split-quantize path still honors it. + NVTE_CHECK(!force_pow_2_scales, + "Fused grouped FP8 block-scaling quantize does not support force_pow_2_scales=True. " + "Set force_pow_2_scales=False, or use the unfused split-quantize path " + "(NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0) which supports power-of-2 scales."); + const auto tensor_offsets = resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, logical_first_dim, logical_last_dim); @@ -1169,18 +1177,36 @@ std::pair Float8BlockQuantizer::create_grouped std::optional columnwise_scale_inv; const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; + // cuBLAS FP8 block-scaling grouped GEMM consumes each expert's scales from a + // contiguous per-expert sub-block in the COMPACT (no 128x4 swizzle) layout — + // VEC128_32F / BLK128x128_32F, unlike MXFP8/NVFP4 which require the swizzle. + // Per-tensor first dims are multiples of 128, so the per-expert padded sizes + // sum exactly to `get_scale_shape` of the logical total shape for 1D and + // 2D-rowwise. The only case where the per-expert sum can exceed the + // totals-based size is 2D columnwise (per-expert roundup(blocks_y_t, 4)), + // so reserve a small slack there. Sizing from totals (instead of reading + // first_dims on host) keeps allocation CUDA-graph-safe: no device->host copy. + constexpr size_t kBlockLen = 128; + const size_t blocks_X = ceildiv(logical_last_dim, kBlockLen); + + auto grouped_scale_elems = [&](bool columnwise) -> int64_t { + const auto sh = get_scale_shape(logical_shape_vec, columnwise); + int64_t total = static_cast(product(sh)); + if (columnwise && block_scaling_dim == 2 && num_tensors > 1) { + // sum_t roundup(blocks_y_t, 4) <= roundup(total_blocks_y, 4) + 4*(num_tensors-1). + total += static_cast(blocks_X * 4 * (num_tensors - 1)); + } + return total; + }; + if (rowwise_usage) { rowwise_data = at::empty({total_elements}, uint8_opts); - const auto scale_shape = get_scale_shape(logical_shape_vec, false); - const int64_t total_scale_elements = static_cast(product(scale_shape)); - rowwise_scale_inv = at::empty({total_scale_elements}, float_opts); + rowwise_scale_inv = at::empty({grouped_scale_elems(false)}, float_opts); } if (columnwise_usage) { columnwise_data = at::empty({total_elements}, uint8_opts); - const auto scale_shape = get_scale_shape(logical_shape_vec, true); - const int64_t total_scale_elements = static_cast(product(scale_shape)); - columnwise_scale_inv = at::empty({total_scale_elements}, float_opts); + columnwise_scale_inv = at::empty({grouped_scale_elems(true)}, float_opts); } GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); @@ -1195,6 +1221,8 @@ std::pair Float8BlockQuantizer::create_grouped out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat32, getTensorShape(*columnwise_scale_inv)); } + // FP8 block-scaling grouped GEMM reads compact scales; never swizzle. + out_cpp.set_with_gemm_swizzled_scales(false); if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); } @@ -1431,7 +1459,7 @@ std::vector Float8BlockQuantizer::get_scale_shape(const std::vector Date: Wed, 1 Jul 2026 11:41:44 -0700 Subject: [PATCH 28/42] [JAX] Keep the routing map format alive and EP multiprocess tests in L2 Jax dist (#3159) * Keep the routing map format alive Signed-off-by: Kshitij Lakhani * Fix incorrectly launched multi process EP tests in L2 Jax instead of L2 jax dist Signed-off-by: Kshitij Lakhani --------- Signed-off-by: Kshitij Lakhani --- qa/L2_jax_distributed_unittest/test.sh | 3 +++ qa/L2_jax_unittest/test.sh | 2 +- transformer_engine/jax/cpp_extensions/router.py | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index 04fbdf1643..330b254e7d 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/qa/L2_jax_distributed_unittest/test.sh @@ -13,3 +13,6 @@ mkdir -p "$XML_LOG_DIR" # Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate. XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_* + +# NCCL EP multi-process suite. The launcher skips when fewer than 4 GPUs or no NVLink is detected. +TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh index 38cbc8ad3d..f455ec0df3 100644 --- a/qa/L2_jax_unittest/test.sh +++ b/qa/L2_jax_unittest/test.sh @@ -28,7 +28,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" -NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax -k 'not distributed' || test_fail "tests/jax/*not_distributed_*" +NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax --ignore=$TE_PATH/tests/jax/test_multi_process_ep.py -k 'not distributed' || test_fail "tests/jax/*not_distributed_*" pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements" # Note: mnist intentionally does NOT set --xla_gpu_deterministic_ops because it diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index 3245439689..8cc94fcaaf 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -412,7 +412,7 @@ def partition( arg_infos, result_infos, ): - del result_infos, routing_map_format + del result_infos grad_spec = get_padded_spec(arg_infos[2]) out_sharding = NamedSharding(mesh, PartitionSpec(*grad_spec)) arg_shardings = (arg_infos[0].sharding, arg_infos[1].sharding, arg_infos[2].sharding) From f72111270941633da00d3c721a867619c1a74599 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Thu, 2 Jul 2026 07:08:27 -0700 Subject: [PATCH 29/42] docs: document attention backend selection (#3142) Signed-off-by: Santosh Bhavani --- docs/envvars.rst | 15 ++++++- docs/examples/attention/attention.ipynb | 39 +++++++++---------- .../dot_product_attention.py | 9 ++++- 3 files changed, 40 insertions(+), 23 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 044a7f6a0d..e8f90a5412 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -122,6 +122,19 @@ These environment variables control the behavior of Transformer Engine during ex Attention Backend Selection ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Transformer Engine attention selects a backend in two stages. First, it filters the available +backends by environment variables, GPU architecture, installed ``flash-attn`` and cuDNN versions, +data type and FP8 recipe, training or inference mode, and the provided attention configuration. +Then it applies a performance-based preference order among the remaining eligible backends. + +In PyTorch, the broad preference order is ``FlashAttention > FusedAttention > +UnfusedDotProductAttention`` on supported pre-Hopper GPUs such as Ampere/Ada, and +``FusedAttention > FlashAttention > UnfusedDotProductAttention`` on Hopper and newer GPUs, +including Blackwell. In JAX, Transformer Engine uses cuDNN fused attention when +``NVTE_FUSED_ATTN=1`` and an eligible cuDNN kernel is available; otherwise it falls back to the +JAX-native implementation. See :doc:`examples/attention/attention` for a longer +backend-selection overview. + .. envvar:: NVTE_FLASH_ATTN :Type: ``int`` (0 or 1) @@ -144,7 +157,7 @@ Attention Backend Selection :Type: ``int`` (1 or 2) :Default: Auto-selected - :Description: Force a specific FusedAttention backend. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. + :Description: Request a cuDNN FusedAttention backend when that request is supported by the active fused-attention path. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. BF16/FP16 attention uses sub-backend ``1`` when eligible. FP8 attention uses sub-backend ``2`` when FP8 DPA is enabled and supported by the architecture, cuDNN version, and input configuration. .. envvar:: NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index e7253415d2..c1c8ff38bf 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -110,14 +110,6 @@ " Additional info\n", " \n", " \n", - " 0\n", - " Non-Flash\n", - " BF16/FP16\n", - " ≤512 \n", - " sm80, 90 \n", - " [cuDNN](https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html#fused-attention-fprop) \n", - " \n", - " \n", " 1\n", " Flash\n", " BF16/FP16\n", @@ -208,11 +200,11 @@ "source": [ "## 2. Backend Selection\n", "\n", - "Given the various attention backends, Transformer Engine has a selection logic in place to choose the most appropriate backend for a particular set of user inputs and runtime environment. The selection logic is based on both backend availability and backend performance.\n", + "Given the various attention backends, Transformer Engine first determines which backends are eligible for the provided inputs and runtime environment, then applies a preference order among the eligible backends. Eligibility is affected by user environment variables, GPU architecture, installed `flash-attn` and cuDNN versions, data type and FP8 recipe, QKV layout, training or inference mode, dropout, and other attention features.\n", "\n", - "Backend availability is determined by factors such as model configuration, training hyper-parameters, software versions, and the GPU architecture in question. For example, some considerations are the sequence length, number of attention heads, head size, attention mask type, attention bias type, training or inference mode, self or cross attention, MHA or MQA/GQA, `flash-attn`/cuDNN library versions, and the compute capability of the GPU.\n", + "In PyTorch, the candidates are FlashAttention (`flash-attn` v2, v3, or v4), FusedAttention (cuDNN sub-backends), and UnfusedDotProductAttention. Users can disable whole backend families with `NVTE_FLASH_ATTN`, `NVTE_FUSED_ATTN`, or `NVTE_UNFUSED_ATTN`. In JAX, Transformer Engine checks whether a cuDNN fused-attention kernel is available when `NVTE_FUSED_ATTN=1`; otherwise it falls back to the JAX-native implementation.\n", "\n", - "When there are multiple backends available, Transformer Engine makes backend selection based on performance. In general, there are a few rules being followed in our selection logic (see table below). As we monitor the performance of different backends, the selection logic may change.\n", + "At a high level, the architecture-specific PyTorch selection order is:\n", "\n", "\n", " \n", @@ -220,22 +212,29 @@ " \n", " \n", " \n", - " \n", - " \n", + " \n", + " \n", " \n", " \n", - " \n", + " \n", " \n", " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", " \n", " \n", - "
Selection Order
PyTorchsm90: cuDNN attention > flash-attention > PyTorch-native attentionPyTorchsm8x (Ampere/Ada): flash-attention > cuDNN attention > PyTorch-native attention
sm80: flash-attention > cuDNN attention > PyTorch-native attentionsm90 (Hopper): cuDNN attention > flash-attention > PyTorch-native attention
\n", - " cuDNN attention: sub-backend 1 > sub-backend 0\n", - " sm100/sm120 (Blackwell): cuDNN attention > flash-attention > PyTorch-native attention
cuDNN attention: BF16/FP16 uses sub-backend 1 when eligible; FP8 uses sub-backend 2 when enabled and eligible
JAXcuDNN attention > JAX-native attention
" + "\n", + "\n", + "Within FlashAttention, TE uses the installed implementation that is supported for the architecture and input. FlashAttention 3 is Hopper-only (`sm90`). FlashAttention 4 supports `sm80`, `sm90`, `sm100`, and `sm120`; on Hopper, TE prefers FlashAttention 3 over FlashAttention 4 when both are installed and eligible. On Blackwell, FlashAttention 4 is the Blackwell-specific flash-attention path when installed and eligible, while FlashAttention 2 can still be eligible depending on the installed version and input configuration.\n", + "\n", + "Within cuDNN FusedAttention, TE asks the fused-attention helper which sub-backend is eligible. Sub-backend 1 is the BF16/FP16 flash-based path when available; sub-backend 2 is the FP8 path when FP8 DPA is enabled and the architecture, cuDNN version, and input configuration support it. Hopper supports eligible FP8 DPA through cuDNN sub-backend 2. In the current PyTorch selector, eligible FP8 DPA on Blackwell is an `sm100` path and is disabled on `sm120`.\n", + "\n", + "When all optimized backends are disabled or ineligible, TE falls back to UnfusedDotProductAttention if it is enabled. If no backend is eligible, backend selection returns no backend and the caller raises an error. As we monitor the performance of different backends, the selection logic may change." ] }, { @@ -350,7 +349,7 @@ "**cuDNN attention sub-backends:**\n", "This environment variable allows users to express their preference of cuDNN attention sub-backends. However, the elected sub-backend will only be used *if* it is eligible, i.e. if it has support for the provided inputs and runtime environment.\n", "```\n", - "NVTE_FUSED_ATTN_BACKEND = 0/1/2 # user preference of cuDNN sub-backend\n", + "NVTE_FUSED_ATTN_BACKEND = 1/2 # user preference of cuDNN sub-backend\n", "```\n", "\n", "**Execution paths of cuDNN sub-backend 1:**\n", @@ -369,7 +368,7 @@ "
\n", "Note\n", " \n", - "Environment variables NVTE_FLASH_ATTN, NVTE_FUSED_ATTN, NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT and NVTE_ALLOW_NONDETERMINISTIC_ALGO are only supported in PyTorch, and will be added to JAX in the future.\n", + "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, NVTE_FUSED_ATTN_BACKEND, NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", "
\n", "\n", "### 2.3 Example Tests\n", 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 aa1481a384..d3ee1a2e2c 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 @@ -1059,8 +1059,13 @@ def forward( Users can use environment variables :attr:`NVTE_FLASH_ATTN`, :attr:`NVTE_FUSED_ATTN`, and :attr:`NVTE_FUSED_ATTN_BACKEND` to control which DotProductAttention backend, - and FusedAttention backend if applicable, to use. Transformer Engine prioritizes - FlashAttention over FusedAttention and over UnfusedDotProductAttention. + and FusedAttention backend if applicable, to use. Transformer Engine first filters + backends by support for the runtime environment and input configuration, then applies + a performance-based preference order. On supported pre-Hopper GPUs, FlashAttention is + preferred over FusedAttention and UnfusedDotProductAttention when both optimized + backends are eligible. On Hopper and newer GPUs, including Blackwell, FusedAttention is + preferred over FlashAttention and UnfusedDotProductAttention when both optimized + backends are eligible. If FusedAttention is being used, users can also choose to switch to flash-attn's implementation for backward by setting :attr:`NVTE_FUSED_ATTN_USE_FAv2_BWD=1` (default: 0), because of the performance differences between various versions of From dc5795824762a9df54bbb46b9378afb15a3fb5ff Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Thu, 2 Jul 2026 17:36:00 -0700 Subject: [PATCH 30/42] Skip MXFP8 MFSDP tests on hopper (#3163) * skip tests on hopper Signed-off-by: Varun Thumbe * Update qa/L1_pytorch_mcore_fsdp_integration/test.sh Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- qa/L1_pytorch_mcore_fsdp_integration/test.sh | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/qa/L1_pytorch_mcore_fsdp_integration/test.sh b/qa/L1_pytorch_mcore_fsdp_integration/test.sh index d63c66f2ea..e08cb8bb98 100644 --- a/qa/L1_pytorch_mcore_fsdp_integration/test.sh +++ b/qa/L1_pytorch_mcore_fsdp_integration/test.sh @@ -4,6 +4,15 @@ set -e +# This test uses the MXFP8 recipe (--fp8-recipe mxfp8), which is only supported +# on Blackwell (compute capability 10.0) and newer. +DEVICE_ARCH_RAW=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -n 1) +DEVICE_ARCH=$(echo "${DEVICE_ARCH_RAW}" | sed 's/[^0-9]//g') +if [[ -z "${DEVICE_ARCH}" || ${DEVICE_ARCH} -lt 100 ]]; then + echo "Skipping L1_pytorch_mcore_fsdp_integration: MXFP8 requires compute capability 10.0+ (Blackwell), detected compute_cap=${DEVICE_ARCH_RAW:-unknown}." + exit 0 +fi + # Megatron-LM / Megatron-FSDP commit for main branch on Apr. 10, 2026. # Necessary to support wgrad accumulate fusion and Megatron-FSDP NCCL UBR, # and fixes decoupled_grad <> DistOpt usage in Megatron-LM. From 7cb8b313d55021eb12c5efd4d02f4c8ff7679453 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 2 Jul 2026 23:38:27 -0500 Subject: [PATCH 31/42] [Common] Blackwell skip condition for C++ grouped FP8 block-scaling tests (#3174) fixing Blackwell skip condition for grouped FP8 block-scaling tests in C++ Signed-off-by: Alp Dener --- tests/cpp/operator/test_cast_float8blockwise_grouped.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu index c7d7a73475..bc9f104e17 100644 --- a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu +++ b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu @@ -74,7 +74,8 @@ template void perform_test(ShapeRep shape_rep, BlockDim block_dim, ScalingDir dir, const std::vector& first_dims_h, size_t K, bool force_pow_2_scales, float epsilon) { - if (getDeviceComputeCapability() < hopperComputeCapability) { + if (getDeviceComputeCapability() < hopperComputeCapability || + getDeviceComputeCapability() >= blackwellComputeCapability) { GTEST_SKIP(); } From 6fb6a0dd181dc179d7b68ef407f8c24b0c990389 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 16:07:32 +0200 Subject: [PATCH 32/42] Address review comments: qualname registry, import and comment cleanups - Replace the class-attribute value-opaque flag with a module-level set of class qualnames: a set of class objects is untraceable under fullgraph=True (opaque classes have no equality rule in Dynamo), but the qualname constant-folds to a plain string; also avoids falsely reporting unregistered subclasses. - Register MXFP8Quantizer right after the class like the other quantizers. - Clarify the amax_reduction_group exclusion comment in Float8CurrentScalingQuantizer._value_fields. - Restore import order in test_torch_compile.py, import NVFP4Quantizer from transformer_engine.pytorch, drop unused Float8Quantizer import. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 5 ++--- .../pytorch/dynamo/quantizer_opaque.py | 21 +++++++++++-------- .../pytorch/tensor/float8_tensor.py | 6 ++++-- .../pytorch/tensor/mxfp8_tensor.py | 6 +++--- 4 files changed, 21 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index c7f4ceb71c..07c2863990 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -24,18 +24,17 @@ 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.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer -from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer +from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, - Float8Quantizer, Float8BlockQuantizer, MXFP8Quantizer, + NVFP4Quantizer, ) from utils import recipe_id diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 4ba4761421..4049af2492 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,19 +10,22 @@ from ..constants import DType -# Registration marks the class with this attribute rather than recording it in a -# module-level set. It looks odd but is a deliberate workaround: the check must -# stay traceable when it runs inside a torch.compile graph -- Dynamo can bake a -# ``getattr`` on the opaque quantizer into a constant, but cannot evaluate -# ``type(q) in some_set`` (no equality/hash rules for the opaque class object), -# which would graph-break under ``fullgraph=True``. -_VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" +# Registration records the class *qualname* rather than the class object. The +# check must stay traceable when it runs inside a torch.compile graph, and a +# set of class objects would not be: once a class is registered as opaque, +# Dynamo traces the class itself as ``OpaqueObjectClassVariable``, which +# defines no equality rule, so ``type(q) in some_set`` falls back to an +# iterate-and-compare polyfill that dies on ``is`` between two opaque class +# variables -- a hard ``Unsupported`` error under ``fullgraph=True``. +# ``type(q).__qualname__`` instead constant-folds to a plain string, and +# string-in-set membership is traceable. +_VALUE_OPAQUE_QUALNAMES: set = set() def is_value_opaque_quantizer(quantizer: Any) -> bool: """Whether *quantizer*'s class is registered as a torch.compile value-opaque type.""" - return getattr(quantizer, _VALUE_OPAQUE_FLAG, False) + return type(quantizer).__qualname__ in _VALUE_OPAQUE_QUALNAMES def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: @@ -111,4 +114,4 @@ def register_value_opaque_quantizer(cls: type) -> None: # experimental opaque-object support. return - setattr(cls, _VALUE_OPAQUE_FLAG, True) + _VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 6687d916fd..4a82f5279a 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -389,8 +389,10 @@ def supports_only_rowwise_all_gather(self) -> bool: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group (not a value). If one is actually stored, ``__fx_repr__`` - # raises so it can never be baked into a torch.compile graph. + # process group, not a value. Quantizers that store one are rejected up + # front (``_value_key`` raises before anything is baked into a + # torch.compile graph), and ``_rebuild_quantizer`` restores the field + # as ``None`` on reconstruction. return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 45af51afec..0436ff69ae 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -184,6 +184,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return MXFP8BlockScaling +register_value_opaque_quantizer(MXFP8Quantizer) + + class MXFP8Tensor(MXFP8TensorStorage, QuantizedTensor): """Experimental tensor class with FP8 data @@ -1065,6 +1068,3 @@ def backward( ) return dgrad, None return grad.view(ctx.shape), None - - -register_value_opaque_quantizer(MXFP8Quantizer) From f9c0e183afd628a5ec55dc48f436de104751621b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 16:14:58 +0200 Subject: [PATCH 33/42] Derive quantizer value fields from class annotations Make the class annotations the single source of truth for what defines a quantizer's value, instead of hand-written per-class _value_fields lists: - Quantizer._value_fields is now derived from the annotations across the MRO (subsuming _BASE_VALUE_FIELDS); register_value_opaque_quantizer is the explicit opt-in and validates at import time that no annotated field is a tensor or process group. - Drop the four per-class _value_fields overrides. - Remove the deprecated amax_reduction_group annotation from Float8CurrentScalingQuantizer and NVFP4Quantizer (the attribute is still set for backward compatibility). - NVFP4: rename _with_random_sign_mask to with_random_sign_mask (annotated, matching the constructor argument), stop storing the derived rht_matrix_random_sign_mask_t in the value key and rebuild it together with rht_matrix in _rebuild_derived_state (lru-cached getters), which __init__ now also uses. copy() now propagates with_random_sign_mask. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 36 ++++++++-- .../pytorch/quantized_tensor.py | 48 ++++++++----- .../pytorch/tensor/float8_blockwise_tensor.py | 3 - .../pytorch/tensor/float8_tensor.py | 17 ++--- .../pytorch/tensor/mxfp8_tensor.py | 3 - .../pytorch/tensor/nvfp4_tensor.py | 68 ++++++++----------- 6 files changed, 97 insertions(+), 78 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 4049af2492..70e5ef3155 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -78,18 +78,40 @@ class itself in the FX globals so codegen can resolve them with no global ) +def _check_value_annotations(cls: type) -> None: + """Verify that every annotated field of *cls* can be part of a value key. + + The value fields are derived from the class annotations (see + ``Quantizer._annotated_value_fields``), so derived tensors and process + groups must not be annotated. Checking annotation strings is enough: the + tensor modules use ``from __future__ import annotations``. + """ + for klass in cls.__mro__: + for name, ann in klass.__dict__.get("__annotations__", {}).items(): + ann = str(ann) + if "Tensor" in ann or "dist_group_type" in ann or "ProcessGroup" in ann: + raise TypeError( + f"{cls.__name__} cannot be a torch.compile value quantizer: " + f"annotated field {name!r} ({ann}) is not a plain value. " + "Remove the annotation and rebuild the field in " + "``_rebuild_derived_state`` instead." + ) + + def register_value_opaque_quantizer(cls: type) -> None: """Register a tensorless quantizer class as a torch.compile value opaque type. - Attaches ``__fx_repr__`` and registers the class with + This is the opt-in point for value semantics: it flips + ``cls._is_value_quantizer`` (enabling config-based ``__eq__`` / ``__hash__`` + with fields derived from the class annotations, see + :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`), attaches + ``__fx_repr__`` and registers the class with ``torch._library.opaque_object``. Safe to call on any PyTorch build: on - versions without the opaque-object API it only attaches ``__fx_repr__`` - (harmless), so Transformer Engine keeps importing and running in eager mode. - - The quantizer class must already provide value ``__eq__`` / ``__hash__`` and - a non-``None`` ``_value_fields`` (see - :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). + versions without the opaque-object API the value semantics still apply, + only the torch.compile specialization is skipped. """ + _check_value_annotations(cls) + cls._is_value_quantizer = True # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 6161539b9c..0ffdc2b144 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -409,24 +409,40 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } - #: Attributes shared by every quantizer that take part in value identity. - _BASE_VALUE_FIELDS: Tuple[str, ...] = ( - "rowwise_usage", - "columnwise_usage", - "internal", - "optimize_for_gemm", - ) + # Flipped to True by ``register_value_opaque_quantizer``. Deliberately not + # annotated: value fields are derived from class annotations, and this + # switch is not part of the value. + _is_value_quantizer = False - def _value_fields(self) -> Optional[Tuple[str, ...]]: - """Subclass-specific value-defining attribute names, or ``None``. + @classmethod + def _annotated_value_fields(cls) -> Tuple[str, ...]: + """Value-defining attribute names, derived from class annotations. + + Collects the annotated fields of every ``Quantizer`` class in the MRO + (base first), so the class annotations are the single source of truth + for what defines a quantizer's value. Fields that are not values + (derived tensors, process groups) must not be annotated; + ``register_value_opaque_quantizer`` enforces this at import time. + """ + fields: Dict[str, None] = {} + for klass in reversed(cls.__mro__): + if issubclass(klass, Quantizer): + fields.update(dict.fromkeys(klass.__dict__.get("__annotations__", {}))) + return tuple(fields) - Returning ``None`` (the default) means the quantizer cannot be represented as - a value opaque object and keeps identity-based equality/hashing. - This also means that passing such a quantizer as an argument to a custom op - causes a graph break under torch.compile, since it cannot be baked into the - FX graph as a constant. + def _value_fields(self) -> Optional[Tuple[str, ...]]: + """Value-defining attribute names, or ``None``. + + ``None`` (any class not registered via + ``register_value_opaque_quantizer``) means the quantizer cannot be + represented as a value opaque object and keeps identity-based + equality/hashing. This also means that passing such a quantizer as an + argument to a custom op causes a graph break under torch.compile, + since it cannot be baked into the FX graph as a constant. """ - return None + if not self._is_value_quantizer: + return None + return self._annotated_value_fields() def _check_value_has_no_process_group(self) -> None: # A value quantizer cannot carry live distributed state into the FX @@ -450,7 +466,7 @@ def _value_key(self) -> Tuple[Any, ...]: assert fields is not None, f"{type(self).__name__} is not a value quantizer" self._check_value_has_no_process_group() items = [] - for name in self._BASE_VALUE_FIELDS + tuple(fields): + for name in fields: value = getattr(self, name) if name == "dtype": # ``DType`` is an ``IntEnum``; store the int so the key stays diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 697a5a96df..18975c4c6d 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -70,9 +70,6 @@ def copy(self) -> Float8BlockQuantizer: return quantizer - def _value_fields(self) -> Tuple[str, ...]: - return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim") - def update_quantized( self, src: torch.Tensor, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 4a82f5279a..abf2417f11 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -206,9 +206,14 @@ class Float8CurrentScalingQuantizer(Quantizer): """FP8 datatype""" dtype: DType - """amax reduction options""" + """amax reduction options + + The deprecated ``amax_reduction_group`` attribute is intentionally not + annotated: annotations define the torch.compile value key, and a process + group is not a value (``_value_key`` rejects a stored group; + ``_rebuild_quantizer`` restores the attribute as ``None``). + """ with_amax_reduction: bool - amax_reduction_group: Optional[dist_group_type] """Options about how to quantize the tensor""" force_pow_2_scales: bool amax_epsilon: float @@ -387,14 +392,6 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True - def _value_fields(self) -> Tuple[str, ...]: - # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group, not a value. Quantizers that store one are rejected up - # front (``_value_key`` raises before anything is baked into a - # torch.compile graph), and ``_rebuild_quantizer`` restores the field - # as ``None`` on reconstruction. - return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") - register_value_opaque_quantizer(Float8CurrentScalingQuantizer) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 0436ff69ae..3045662216 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -58,9 +58,6 @@ def copy(self) -> MXFP8Quantizer: return quantizer - def _value_fields(self) -> Tuple[str, ...]: - return ("dtype",) - def update_quantized( self, src: torch.Tensor, diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 25fcfa9d14..0016f1a0b0 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -118,9 +118,14 @@ class NVFP4Quantizer(Quantizer): """Random Hadamard Transform""" with_rht: bool with_post_rht_amax: bool - """amax reduction options""" + """amax reduction options + + The deprecated ``amax_reduction_group`` attribute is intentionally not + annotated: annotations define the torch.compile value key, and a process + group is not a value (``_value_key`` rejects a stored group; + ``_rebuild_quantizer`` restores the attribute as ``None``). + """ with_amax_reduction: bool - amax_reduction_group: Optional[dist_group_type] """2D block scaling, only applicable for weights.""" with_2d_quantization: bool @@ -137,9 +142,14 @@ class NVFP4Quantizer(Quantizer): """NVFP4 4over6 candidate-selection error mode.""" nvfp4_4over6_err_mode: str - """RHT matrix random sign mask""" - rht_matrix_random_sign_mask_t: int - rht_matrix: torch.Tensor + """Whether the RHT sign mask is randomized. + + The derived ``rht_matrix_random_sign_mask_t`` (int) and ``rht_matrix`` + (tensor) attributes are intentionally not annotated: annotations define + the torch.compile value key, and both are rebuilt from this flag by + ``_rebuild_derived_state``. + """ + with_random_sign_mask: bool def __init__( self, @@ -174,11 +184,8 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") - self._with_random_sign_mask = with_random_sign_mask - self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( - with_random_sign_mask, torch.cuda.current_device() - ) - self.rht_matrix = get_rht_matrix(with_random_sign_mask, torch.cuda.current_device()) + self.with_random_sign_mask = with_random_sign_mask + self._rebuild_derived_state() def __getstate__(self): """Exclude unpicklable process group from serialized state.""" @@ -187,14 +194,19 @@ def __getstate__(self): return state def _rebuild_derived_state(self) -> None: - """Restore the derived ``rht_matrix`` after value-key reconstruction. + """Build the derived RHT state (also used after value-key reconstruction). - ``rht_matrix`` is a ``torch.Tensor`` built from ``_with_random_sign_mask`` - and the device, so it cannot be part of the (hashable) value key. - ``_rebuild_quantizer`` calls this hook to rebuild it; the ``lru_cache`` on - :func:`get_rht_matrix` makes an already-seen (flag, device) a cheap hit. + ``rht_matrix`` is a ``torch.Tensor`` and ``rht_matrix_random_sign_mask_t`` + is derived from ``with_random_sign_mask``, so neither is part of the + (hashable) value key. ``__init__`` and ``_rebuild_quantizer`` both call + this hook; the ``lru_cache`` on the getters makes an already-seen + (flag, device) pair a cheap hit. """ - self.rht_matrix = get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device()) + device = torch.cuda.current_device() + self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( + self.with_random_sign_mask, device + ) + self.rht_matrix = get_rht_matrix(self.with_random_sign_mask, device) def update_quantized( self, @@ -242,11 +254,10 @@ def copy(self) -> NVFP4Quantizer: nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, + with_random_sign_mask=self.with_random_sign_mask, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm - quantizer.rht_matrix = self.rht_matrix - quantizer.rht_matrix_random_sign_mask_t = self.rht_matrix_random_sign_mask_t return quantizer @@ -345,27 +356,6 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type: def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling - def _value_fields(self) -> Tuple[str, ...]: - # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group, not a value (``_value_key`` rejects a stored group). - # ``rht_matrix_random_sign_mask_t`` is a device-independent int derived - # from ``_with_random_sign_mask``; kept in the key so the rebuilt - # quantizer carries it without recomputation. - return ( - "dtype", - "with_rht", - "with_post_rht_amax", - "with_2d_quantization", - "stochastic_rounding", - "row_scaled_nvfp4", - "nvfp4_use_4over6", - "nvfp4_e4m3_max", - "nvfp4_4over6_err_mode", - "_with_random_sign_mask", - "rht_matrix_random_sign_mask_t", - "with_amax_reduction", - ) - register_value_opaque_quantizer(NVFP4Quantizer) From dd969562edaf37b5260ab7563f8e51a55e6a90a5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 16:34:56 +0200 Subject: [PATCH 34/42] Drop redundant annotation-exclusion comments Signed-off-by: Pawel Gadzinski --- .../pytorch/tensor/float8_tensor.py | 8 +------- .../pytorch/tensor/nvfp4_tensor.py | 16 ++-------------- 2 files changed, 3 insertions(+), 21 deletions(-) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index abf2417f11..5776200daa 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -206,13 +206,7 @@ class Float8CurrentScalingQuantizer(Quantizer): """FP8 datatype""" dtype: DType - """amax reduction options - - The deprecated ``amax_reduction_group`` attribute is intentionally not - annotated: annotations define the torch.compile value key, and a process - group is not a value (``_value_key`` rejects a stored group; - ``_rebuild_quantizer`` restores the attribute as ``None``). - """ + """amax reduction options""" with_amax_reduction: bool """Options about how to quantize the tensor""" force_pow_2_scales: bool diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 0016f1a0b0..2b435fbdb0 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -118,13 +118,7 @@ class NVFP4Quantizer(Quantizer): """Random Hadamard Transform""" with_rht: bool with_post_rht_amax: bool - """amax reduction options - - The deprecated ``amax_reduction_group`` attribute is intentionally not - annotated: annotations define the torch.compile value key, and a process - group is not a value (``_value_key`` rejects a stored group; - ``_rebuild_quantizer`` restores the attribute as ``None``). - """ + """amax reduction options""" with_amax_reduction: bool """2D block scaling, only applicable for weights.""" @@ -142,13 +136,7 @@ class NVFP4Quantizer(Quantizer): """NVFP4 4over6 candidate-selection error mode.""" nvfp4_4over6_err_mode: str - """Whether the RHT sign mask is randomized. - - The derived ``rht_matrix_random_sign_mask_t`` (int) and ``rht_matrix`` - (tensor) attributes are intentionally not annotated: annotations define - the torch.compile value key, and both are rebuilt from this flag by - ``_rebuild_derived_state``. - """ + """Whether the RHT sign mask is randomized""" with_random_sign_mask: bool def __init__( From 447a4e1cac5b55e81c480f697fe8cba4b6655c09 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 16:46:09 +0200 Subject: [PATCH 35/42] Shorten _is_value_quantizer comment Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/quantized_tensor.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 0ffdc2b144..443acdd015 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -409,9 +409,7 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } - # Flipped to True by ``register_value_opaque_quantizer``. Deliberately not - # annotated: value fields are derived from class annotations, and this - # switch is not part of the value. + # Flipped to True by ``register_value_opaque_quantizer``. _is_value_quantizer = False @classmethod From 02a8fc91e61fab1fc6f9f8ebc8f8111e3aaba66e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 16:48:02 +0200 Subject: [PATCH 36/42] Fold annotation walk into _value_fields, trim comments Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 2 +- .../pytorch/quantized_tensor.py | 36 +++++++------------ 2 files changed, 14 insertions(+), 24 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 70e5ef3155..d25d647b95 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -82,7 +82,7 @@ def _check_value_annotations(cls: type) -> None: """Verify that every annotated field of *cls* can be part of a value key. The value fields are derived from the class annotations (see - ``Quantizer._annotated_value_fields``), so derived tensors and process + ``Quantizer._value_fields``), so derived tensors and process groups must not be annotated. Checking annotation strings is enough: the tensor modules use ``from __future__ import annotations``. """ diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 443acdd015..90682ae842 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -412,35 +412,25 @@ def get_usages(self) -> Dict[str, bool]: # Flipped to True by ``register_value_opaque_quantizer``. _is_value_quantizer = False - @classmethod - def _annotated_value_fields(cls) -> Tuple[str, ...]: - """Value-defining attribute names, derived from class annotations. - - Collects the annotated fields of every ``Quantizer`` class in the MRO - (base first), so the class annotations are the single source of truth - for what defines a quantizer's value. Fields that are not values - (derived tensors, process groups) must not be annotated; - ``register_value_opaque_quantizer`` enforces this at import time. - """ - fields: Dict[str, None] = {} - for klass in reversed(cls.__mro__): - if issubclass(klass, Quantizer): - fields.update(dict.fromkeys(klass.__dict__.get("__annotations__", {}))) - return tuple(fields) - def _value_fields(self) -> Optional[Tuple[str, ...]]: """Value-defining attribute names, or ``None``. - ``None`` (any class not registered via - ``register_value_opaque_quantizer``) means the quantizer cannot be - represented as a value opaque object and keeps identity-based - equality/hashing. This also means that passing such a quantizer as an - argument to a custom op causes a graph break under torch.compile, - since it cannot be baked into the FX graph as a constant. + Derived from the class annotations across the MRO (base first), so the + annotations are the single source of truth for what defines a + quantizer's value; ``register_value_opaque_quantizer`` checks at import + time that no annotated field is a derived tensor or a process group. + ``None`` (any class not registered) keeps identity-based + equality/hashing and graph-breaks under torch.compile when passed to a + custom op, since such a quantizer cannot be baked into the FX graph as + a constant. """ if not self._is_value_quantizer: return None - return self._annotated_value_fields() + fields: Dict[str, None] = {} + for klass in reversed(type(self).__mro__): + if issubclass(klass, Quantizer): + fields.update(dict.fromkeys(klass.__dict__.get("__annotations__", {}))) + return tuple(fields) def _check_value_has_no_process_group(self) -> None: # A value quantizer cannot carry live distributed state into the FX From 5f55a51529bea247830dbf17160ec0212b442255 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 16:49:24 +0200 Subject: [PATCH 37/42] Simplify qualname-registry comment Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index d25d647b95..651b0cef07 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,15 +10,10 @@ from ..constants import DType -# Registration records the class *qualname* rather than the class object. The -# check must stay traceable when it runs inside a torch.compile graph, and a -# set of class objects would not be: once a class is registered as opaque, -# Dynamo traces the class itself as ``OpaqueObjectClassVariable``, which -# defines no equality rule, so ``type(q) in some_set`` falls back to an -# iterate-and-compare polyfill that dies on ``is`` between two opaque class -# variables -- a hard ``Unsupported`` error under ``fullgraph=True``. -# ``type(q).__qualname__`` instead constant-folds to a plain string, and -# string-in-set membership is traceable. +# Registered classes are recorded by qualname because the check may run inside +# a torch.compile'd region: a name-in-set test is traceable there, while a set +# of opaque-registered class objects is not (Dynamo cannot compare opaque +# classes, so ``type(q) in some_set`` graph-breaks under ``fullgraph=True``). _VALUE_OPAQUE_QUALNAMES: set = set() From 02afbdc236121350feb02e36df364c8ba09ce401 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 16:50:38 +0200 Subject: [PATCH 38/42] Reword qualname-registry comment for outside readers Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/quantizer_opaque.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 651b0cef07..a8a9065125 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,10 +10,11 @@ from ..constants import DType -# Registered classes are recorded by qualname because the check may run inside -# a torch.compile'd region: a name-in-set test is traceable there, while a set -# of opaque-registered class objects is not (Dynamo cannot compare opaque -# classes, so ``type(q) in some_set`` graph-breaks under ``fullgraph=True``). +# Qualnames of the registered quantizer classes. The set holds strings rather +# than the classes themselves so that ``is_value_opaque_quantizer`` can be +# called inside a ``torch.compile``'d function without a graph break: Dynamo +# can evaluate ``type(q).__qualname__ in ``, but not set +# membership of a class registered as an opaque type. _VALUE_OPAQUE_QUALNAMES: set = set() From 1590931cc300da67b621f2e3d5f639fd53d78f58 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 17:19:58 +0200 Subject: [PATCH 39/42] Fix review findings: pickle compat, subclass opt-in, cached value fields - NVFP4: stop storing with_random_sign_mask; the annotated value field is the derived (deterministic, device-independent) rht_matrix_random_sign_mask_t and rht_matrix is rebuilt from it in _rebuild_derived_state. Quantizers pickled before this PR already carry the mask in __dict__, so old checkpoints keep working (a boolean back-fill default could lie for quantizers created with with_random_sign_mask=False). - Value semantics no longer leak into unregistered subclasses: register_value_opaque_quantizer stores the field tuple on the class and _value_fields looks it up in the class's own __dict__, so a subclass must register explicitly (it previously inherited value eq/hash that ignored its unannotated fields and skipped the annotation check). - The value-field tuple is computed once at registration instead of an MRO walk per __eq__/__hash__ call (these run per compiled-function invocation via the EQUALS_MATCH guard); registration validation and field derivation now share one annotation walk (Quantizer._annotated_fields). Drop the unreachable other._value_fields() branch and do the cheap type check first in __eq__. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 44 ++++++++----------- .../pytorch/quantized_tensor.py | 33 +++++++------- .../pytorch/tensor/nvfp4_tensor.py | 28 ++++++------ 3 files changed, 50 insertions(+), 55 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index a8a9065125..aa8242f962 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -74,40 +74,34 @@ class itself in the FX globals so codegen can resolve them with no global ) -def _check_value_annotations(cls: type) -> None: - """Verify that every annotated field of *cls* can be part of a value key. - - The value fields are derived from the class annotations (see - ``Quantizer._value_fields``), so derived tensors and process - groups must not be annotated. Checking annotation strings is enough: the - tensor modules use ``from __future__ import annotations``. - """ - for klass in cls.__mro__: - for name, ann in klass.__dict__.get("__annotations__", {}).items(): - ann = str(ann) - if "Tensor" in ann or "dist_group_type" in ann or "ProcessGroup" in ann: - raise TypeError( - f"{cls.__name__} cannot be a torch.compile value quantizer: " - f"annotated field {name!r} ({ann}) is not a plain value. " - "Remove the annotation and rebuild the field in " - "``_rebuild_derived_state`` instead." - ) - - def register_value_opaque_quantizer(cls: type) -> None: """Register a tensorless quantizer class as a torch.compile value opaque type. - This is the opt-in point for value semantics: it flips - ``cls._is_value_quantizer`` (enabling config-based ``__eq__`` / ``__hash__`` - with fields derived from the class annotations, see + This is the opt-in point for value semantics: it derives the value fields + from the class annotations and stores them on the class (enabling + config-based ``__eq__`` / ``__hash__``, see :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`), attaches ``__fx_repr__`` and registers the class with ``torch._library.opaque_object``. Safe to call on any PyTorch build: on versions without the opaque-object API the value semantics still apply, only the torch.compile specialization is skipped. + + Derived tensors and process groups must not be annotated (rebuild them in + ``_rebuild_derived_state`` instead); this is checked here at import time. + Checking annotation strings is enough: the tensor modules use + ``from __future__ import annotations``. """ - _check_value_annotations(cls) - cls._is_value_quantizer = True + fields = cls._annotated_fields() + for name, ann in fields.items(): + ann = str(ann) + if "Tensor" in ann or "dist_group_type" in ann or "ProcessGroup" in ann: + raise TypeError( + f"{cls.__name__} cannot be a torch.compile value quantizer: " + f"annotated field {name!r} ({ann}) is not a plain value. " + "Remove the annotation and rebuild the field in " + "``_rebuild_derived_state`` instead." + ) + cls._value_field_names = tuple(fields) # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 90682ae842..ce820af9c8 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -409,28 +409,31 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } - # Flipped to True by ``register_value_opaque_quantizer``. - _is_value_quantizer = False + @classmethod + def _annotated_fields(cls) -> Dict[str, Any]: + """Annotated fields (name -> annotation) across the ``Quantizer`` MRO, + base first. The class annotations are the single source of truth for + what defines a quantizer's value.""" + fields: Dict[str, Any] = {} + for klass in reversed(cls.__mro__): + if issubclass(klass, Quantizer): + fields.update(klass.__dict__.get("__annotations__", {})) + return fields def _value_fields(self) -> Optional[Tuple[str, ...]]: """Value-defining attribute names, or ``None``. - Derived from the class annotations across the MRO (base first), so the - annotations are the single source of truth for what defines a - quantizer's value; ``register_value_opaque_quantizer`` checks at import - time that no annotated field is a derived tensor or a process group. + Computed from the class annotations and stored on the class by + ``register_value_opaque_quantizer``, which also checks that no + annotated field is a derived tensor or a process group. Looked up in + the class's own ``__dict__`` so a subclass of a registered quantizer + does not silently inherit value semantics without registering itself. ``None`` (any class not registered) keeps identity-based equality/hashing and graph-breaks under torch.compile when passed to a custom op, since such a quantizer cannot be baked into the FX graph as a constant. """ - if not self._is_value_quantizer: - return None - fields: Dict[str, None] = {} - for klass in reversed(type(self).__mro__): - if issubclass(klass, Quantizer): - fields.update(dict.fromkeys(klass.__dict__.get("__annotations__", {}))) - return tuple(fields) + return type(self).__dict__.get("_value_field_names") def _check_value_has_no_process_group(self) -> None: # A value quantizer cannot carry live distributed state into the FX @@ -469,9 +472,7 @@ def __eq__(self, other: object) -> Any: # fall back to identity). ``_value_key`` rejects a stored ProcessGroup. if self is other: return True - if self._value_fields() is None or type(self) is not type(other): - return NotImplemented - if other._value_fields() is None: + if type(self) is not type(other) or self._value_fields() is None: return NotImplemented return self._value_key() == other._value_key() diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 2b435fbdb0..e66d8c1a6d 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -136,8 +136,8 @@ class NVFP4Quantizer(Quantizer): """NVFP4 4over6 candidate-selection error mode.""" nvfp4_4over6_err_mode: str - """Whether the RHT sign mask is randomized""" - with_random_sign_mask: bool + """RHT sign mask (0 when sign randomization is disabled)""" + rht_matrix_random_sign_mask_t: int def __init__( self, @@ -172,7 +172,9 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") - self.with_random_sign_mask = with_random_sign_mask + self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( + with_random_sign_mask, torch.cuda.current_device() + ) self._rebuild_derived_state() def __getstate__(self): @@ -182,19 +184,17 @@ def __getstate__(self): return state def _rebuild_derived_state(self) -> None: - """Build the derived RHT state (also used after value-key reconstruction). + """Build the derived ``rht_matrix`` (also used after value-key reconstruction). - ``rht_matrix`` is a ``torch.Tensor`` and ``rht_matrix_random_sign_mask_t`` - is derived from ``with_random_sign_mask``, so neither is part of the - (hashable) value key. ``__init__`` and ``_rebuild_quantizer`` both call - this hook; the ``lru_cache`` on the getters makes an already-seen - (flag, device) pair a cheap hit. + ``rht_matrix`` is a ``torch.Tensor`` derived from the sign mask, so it + cannot be part of the (hashable) value key. ``__init__`` and + ``_rebuild_quantizer`` both call this hook; the ``lru_cache`` on + :func:`get_rht_matrix` makes an already-seen (flag, device) pair a + cheap hit. """ - device = torch.cuda.current_device() - self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( - self.with_random_sign_mask, device + self.rht_matrix = get_rht_matrix( + self.rht_matrix_random_sign_mask_t != 0, torch.cuda.current_device() ) - self.rht_matrix = get_rht_matrix(self.with_random_sign_mask, device) def update_quantized( self, @@ -242,7 +242,7 @@ def copy(self) -> NVFP4Quantizer: nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, - with_random_sign_mask=self.with_random_sign_mask, + with_random_sign_mask=self.rht_matrix_random_sign_mask_t != 0, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm From 43a4083202f51646e8b717af42e9c4ed8b2de0bb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 17:35:18 +0200 Subject: [PATCH 40/42] Validate value-field annotations by resolved type, not annotation text Replace the substring blocklist with get_type_hints + an allowlist of value types (int/bool/float/str/enum): aliased tensor types no longer slip through and benign types whose name merely contains "Tensor" are no longer rejected. Runs once per class at import time, not in any hot path. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index aa8242f962..692857b4a9 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -5,7 +5,8 @@ """Value-opaque quantizers for torch.compile.""" from __future__ import annotations -from typing import Any, Dict, Tuple +import enum +from typing import Any, Dict, Tuple, get_type_hints from ..constants import DType @@ -86,20 +87,25 @@ def register_value_opaque_quantizer(cls: type) -> None: versions without the opaque-object API the value semantics still apply, only the torch.compile specialization is skipped. - Derived tensors and process groups must not be annotated (rebuild them in - ``_rebuild_derived_state`` instead); this is checked here at import time. - Checking annotation strings is enough: the tensor modules use - ``from __future__ import annotations``. + Only plain value types (``int``/``bool``/``float``/``str`` and enums) may + be annotated: anything else (derived tensors, process groups, containers) + cannot be hashed into the value key or rebuilt from its repr, so it must + be left unannotated and rebuilt in ``_rebuild_derived_state`` instead. + This runs once per class at import time, not in any hot path, so resolving + the annotation strings to real types is affordable. """ fields = cls._annotated_fields() - for name, ann in fields.items(): - ann = str(ann) - if "Tensor" in ann or "dist_group_type" in ann or "ProcessGroup" in ann: + resolved = get_type_hints(cls) + for name in fields: + typ = resolved[name] + if typ not in (int, bool, float, str) and not ( + isinstance(typ, type) and issubclass(typ, enum.Enum) + ): raise TypeError( f"{cls.__name__} cannot be a torch.compile value quantizer: " - f"annotated field {name!r} ({ann}) is not a plain value. " - "Remove the annotation and rebuild the field in " - "``_rebuild_derived_state`` instead." + f"annotated field {name!r} ({typ!r}) is not a plain value type " + "(int/bool/float/str/enum). Remove the annotation and rebuild " + "the field in ``_rebuild_derived_state`` instead." ) cls._value_field_names = tuple(fields) # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the From 10994b71e4d6ad281c307343212fe8567830c84b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 17:53:13 +0200 Subject: [PATCH 41/42] Drop the amax_reduction_group fixup from the generic rebuilder _rebuild_quantizer no longer back-fills the deprecated amax_reduction_group (and loses the field_names set that existed only for that check): a rebuilt quantizer deliberately lacks the attribute, so anything that genuinely needs it fails loudly instead of silently getting None. The only unconditional readers were the two copy() methods, which now tolerate the absent field; _canonicalized_amax_reduction_group (used by the kernel only when with_amax_reduction is set) still raises AttributeError on a rebuilt quantizer, which is the intended behavior. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/quantizer_opaque.py | 11 +++++------ transformer_engine/pytorch/tensor/float8_tensor.py | 3 ++- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 3 ++- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 692857b4a9..a342dd1e6c 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -30,21 +30,20 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the generated FX code calls this to materialize the quantizer constant. + + Only the value fields (plus derived state via ``_rebuild_derived_state``) + are restored. Non-value attributes such as the deprecated + ``amax_reduction_group`` are deliberately absent on the rebuilt quantizer, + so accessing them fails loudly unless set explicitly. """ # Bypass ``__init__`` and restore the value attributes directly: the value # items already capture every value-defining field (including derived ones), # and the constructors have heterogeneous signatures / side effects. obj = cls.__new__(cls) - field_names = set() for name, value in items: if name == "dtype": value = DType.cast(value) object.__setattr__(obj, name, value) - field_names.add(name) - # The deprecated amax-reduction group is not a value field; initialize it to - # None so attribute access keeps working on the rebuilt quantizer. - if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): - object.__setattr__(obj, "amax_reduction_group", None) # Restore non-value derived state that ``__init__`` would normally build but # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor). finalize = getattr(obj, "_rebuild_derived_state", None) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 5776200daa..2e0491d49a 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -257,7 +257,8 @@ def copy(self) -> Float8CurrentScalingQuantizer: rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, with_amax_reduction=self.with_amax_reduction, - amax_reduction_group=self.amax_reduction_group, + # Absent on quantizers rebuilt from a value key (deprecated field). + amax_reduction_group=getattr(self, "amax_reduction_group", None), force_pow_2_scales=self.force_pow_2_scales, amax_epsilon=self.amax_epsilon, ) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index e66d8c1a6d..ccf06ac166 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -233,7 +233,8 @@ def copy(self) -> NVFP4Quantizer: rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, with_amax_reduction=self.with_amax_reduction, - amax_reduction_group=self.amax_reduction_group, + # Absent on quantizers rebuilt from a value key (deprecated field). + amax_reduction_group=getattr(self, "amax_reduction_group", None), with_rht=self.with_rht, with_post_rht_amax=self.with_post_rht_amax, with_2d_quantization=self.with_2d_quantization, From ba520cad11247e00fa4a3b7d297965c09f6ce8a0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 6 Jul 2026 19:56:34 +0200 Subject: [PATCH 42/42] Enable post-RHT amax in the NVFP4 value-object test factory The quantize kernel rejects with_rht=True without with_post_rht_amax=True (pre-RHT amax unsupported); mirror the recipe, which always sets both together. Unnoticed locally because the NVFP4 round-trip is skipped on non-NVFP4 hardware. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 07c2863990..8ebce563ce 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -418,12 +418,14 @@ def _current_scaling(amax_epsilon=0.0): def _nvfp4(with_rht=True): # Default with_rht=True so the quantize round-trip below exercises the # derived ``rht_matrix`` tensor (the field most likely to be dropped on - # value-key reconstruction). + # value-key reconstruction). Post-RHT amax is required by the kernel + # whenever RHT is on (pre-RHT amax is unsupported). return NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, rowwise=True, columnwise=True, with_rht=with_rht, + with_post_rht_amax=with_rht, )