diff --git a/benchmarks/benchmark_rht_cast.py b/benchmarks/benchmark_rht_cast.py index badab1d1994..2ef63f705a2 100644 --- a/benchmarks/benchmark_rht_cast.py +++ b/benchmarks/benchmark_rht_cast.py @@ -8,7 +8,6 @@ import torch.utils.benchmark as benchmark import transformer_engine.pytorch as te -import transformer_engine_torch as tex import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer @@ -17,7 +16,7 @@ permute_scale = False TORCH_TO_TE_FLOAT_MAP = { - torch.bfloat16: tex.DType.kBFloat16, + torch.bfloat16: te.DType.kBFloat16, } @@ -31,7 +30,7 @@ def run_kernel(shape, stochastic_rounding: bool, input_dtype=torch.bfloat16): # Quantize nvfp4_quantizer = NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=te.DType.kFloat4E2M1, rowwise=True, columnwise=True, with_amax_reduction=False, diff --git a/benchmarks/linear/benchmark_graph_safe_grouped_linear.py b/benchmarks/linear/benchmark_graph_safe_grouped_linear.py new file mode 100644 index 00000000000..d8230c38fe0 --- /dev/null +++ b/benchmarks/linear/benchmark_graph_safe_grouped_linear.py @@ -0,0 +1,380 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark MXFP8 graph-safe grouped MLP. + +This mirrors ``benchmark_grouped_linear.py`` but targets the graph-safe TE ops +path used by grouped MLP: + + GroupedLinear -> ScaledSwiGLU -> GroupedLinear + +The benchmark intentionally uses CUDA-device ``m_splits`` and MXFP8 only. + +Example: + + python benchmarks/linear/benchmark_graph_safe_grouped_linear.py + +Forward-only: + + python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --fwd-only + +Nsight Systems: + + (optionally: unset DEBUGINFOD_URLS) + + nsys profile \ + --output=./benchmarks/linear/graph_safe_grouped_linear_mxfp8 \ + --force-overwrite true \ + --trace=cuda,nvtx,cudnn,cublas \ + python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --profile +""" + +# Match the Qwen MXFP8 SFT launch toggles before importing TE. +import os + +os.environ.setdefault("CUDA_DEVICE_MAX_CONNECTIONS", "1") +os.environ.setdefault("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") +os.environ.setdefault("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "1") +os.environ.setdefault("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") + +import argparse +from contextlib import nullcontext + +import pandas as pd +import torch +import torch.utils.benchmark as benchmark + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + +MXFP8_AVAILABLE, REASON_FOR_NO_MXFP8 = FP8GlobalStateManager.is_mxfp8_available() + + +def parse_int_list(value: str) -> list[int]: + """Parse comma-separated integers.""" + return [int(x) for x in value.split(",") if x] + + +def make_uniform_splits(total_tokens: int, num_groups: int) -> list[int]: + """Split tokens uniformly across groups.""" + if total_tokens % num_groups != 0: + raise ValueError( + "Uniform split requires total_tokens divisible by num_groups, " + f"got total_tokens={total_tokens}, num_groups={num_groups}" + ) + return [total_tokens // num_groups] * num_groups + + +def build_grouped_mlp( + *, + num_groups: int, + hidden_dim: int, + ffn_hidden_dim: int, + dtype: torch.dtype, + single_grouped_weight: bool, + accumulate_into_main_grad: bool, + glu_interleave_size: int, +) -> te_ops.Sequential: + """Build graph-safe grouped MLP ops sequence.""" + recipe = MXFP8BlockScaling() + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te_ops.GroupedLinear( + num_groups, + hidden_dim, + 2 * ffn_hidden_dim, + bias=False, + device="cuda", + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + ) + fc2 = te_ops.GroupedLinear( + num_groups, + ffn_hidden_dim, + hidden_dim, + bias=False, + device="cuda", + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + ) + return te_ops.Sequential( + fc1, + te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), + fc2, + ) + + +def init_main_grads(module: torch.nn.Module, value: float = 0.0) -> None: + """Initialize Megatron-style main_grad buffers for accumulate_into_main_grad.""" + with torch.no_grad(): + for param in module.parameters(): + if getattr(param, "main_grad", None) is None: + param.main_grad = torch.empty( + param.size(), device=param.device, dtype=torch.float32 + ) + param.main_grad.fill_(value) + + +def zero_grads(module: torch.nn.Module, x: torch.Tensor, scales: torch.Tensor) -> None: + """Reset gradients without changing allocated main_grad buffers.""" + module.zero_grad(set_to_none=True) + x.grad = None + scales.grad = None + + +def run_grouped_mlp_steps( + module: torch.nn.Module, + x: torch.Tensor, + split_sizes: torch.Tensor, + scales: torch.Tensor, + grad_output: torch.Tensor, + *, + recipe: MXFP8BlockScaling, + fwd_only: bool, + num_steps: int, + accumulate_into_main_grad: bool, +) -> torch.Tensor: + """Run eager grouped MLP for a number of synthetic microbatches.""" + quantization_context = te.autocast(enabled=True, recipe=recipe) + + if fwd_only: + with torch.no_grad(), quantization_context: + for _ in range(num_steps): + out = module(x, split_sizes, scales, split_sizes) + return out + + zero_grads(module, x, scales) + if accumulate_into_main_grad: + init_main_grads(module) + + with quantization_context: + for step in range(num_steps): + torch.cuda.nvtx.range_push(f"step_{step}") + out = module(x, split_sizes, scales, split_sizes) + out.backward(grad_output) + torch.cuda.nvtx.range_pop() + return out + + +def benchmark_case( + *, + total_tokens: int, + hidden_dim: int, + ffn_hidden_dim: int, + num_groups: int, + dtype: torch.dtype, + fwd_only: bool, + single_grouped_weight: bool, + accumulate_into_main_grad: bool, + glu_interleave_size: int, + num_microbatches: int, + min_run_time: float, + profile: bool, +) -> float: + """Benchmark one grouped MLP shape.""" + split_sizes_list = make_uniform_splits(total_tokens, num_groups) + split_sizes = torch.tensor(split_sizes_list, dtype=torch.int64, device="cuda") + x = torch.randn( + (total_tokens, hidden_dim), + dtype=dtype, + device="cuda", + requires_grad=not fwd_only, + ) + scales = torch.ones( + (total_tokens,), + dtype=dtype, + device="cuda", + requires_grad=not fwd_only, + ) + grad_output = torch.ones((total_tokens, hidden_dim), dtype=dtype, device="cuda") + + module = build_grouped_mlp( + num_groups=num_groups, + hidden_dim=hidden_dim, + ffn_hidden_dim=ffn_hidden_dim, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + glu_interleave_size=glu_interleave_size, + ) + recipe = MXFP8BlockScaling() + + print( + "case:", + f"tokens={total_tokens}", + f"hidden={hidden_dim}", + f"ffn_hidden={ffn_hidden_dim}", + f"num_groups={num_groups}", + f"fwd_only={fwd_only}", + f"single_grouped_weight={single_grouped_weight}", + f"accumulate_into_main_grad={accumulate_into_main_grad}", + f"glu_interleave_size={glu_interleave_size}", + ) + print(f"m_splits: {split_sizes_list}") + + # Warmup also forces the op-fuser to materialize the expected fused ops. + run_grouped_mlp_steps( + module, + x, + split_sizes, + scales, + grad_output, + recipe=recipe, + fwd_only=fwd_only, + num_steps=128, + accumulate_into_main_grad=accumulate_into_main_grad, + ) + torch.cuda.synchronize() + + forward_ops = module._module_groups[0]._forward_ops + print("forward fused op:", type(forward_ops[0][0]).__name__ if forward_ops else "none") + if not fwd_only: + backward_ops = module._module_groups[0]._backward_ops + print("backward fused op:", type(backward_ops[0][0]).__name__ if backward_ops else "none") + + label = "graph_safe_grouped_mlp_mxfp8_swiglu" + timing_context = ( + torch.autograd.profiler.emit_nvtx(record_shapes=True) if profile else nullcontext() + ) + with timing_context: + torch.cuda.nvtx.range_push(label) + timing = benchmark.Timer( + stmt=( + "run_grouped_mlp_steps(" + "module, x, split_sizes, scales, grad_output, " + "recipe=recipe, fwd_only=fwd_only, num_steps=num_microbatches, " + "accumulate_into_main_grad=accumulate_into_main_grad)" + ), + globals={ + "run_grouped_mlp_steps": run_grouped_mlp_steps, + "module": module, + "x": x, + "split_sizes": split_sizes, + "scales": scales, + "grad_output": grad_output, + "recipe": recipe, + "fwd_only": fwd_only, + "num_microbatches": num_microbatches, + "accumulate_into_main_grad": accumulate_into_main_grad, + }, + num_threads=1, + ).blocked_autorange(min_run_time=min_run_time) + torch.cuda.nvtx.range_pop() + + print(f"mxfp8_swiglu: {timing}\n") + return timing.median * 1000 / num_microbatches + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--profile", action="store_true", help="Enable NVTX profiling annotations") + parser.add_argument( + "--fwd-only", + action="store_true", + default=False, + help="Benchmark forward only. Default benchmarks forward + backward.", + ) + parser.add_argument( + "--num-groups", + type=str, + default="8", + help="Comma-separated local grouped GEMM/expert counts.", + ) + parser.add_argument( + "--token-dims", + type=str, + default="65536", + help="Comma-separated total token counts to benchmark.", + ) + parser.add_argument("--hidden-dim", type=int, default=7168) + parser.add_argument("--ffn-hidden-dim", type=int, default=2048) + parser.add_argument("--num-microbatches", type=int, default=32) + parser.add_argument("--min-run-time", type=float, default=10.0) + parser.add_argument("--glu-interleave-size", type=int, default=32) + parser.add_argument( + "--single-grouped-weight", + action="store_true", + default=False, + help="Use one GroupedTensor parameter for each grouped linear.", + ) + args = parser.parse_args() + + if not MXFP8_AVAILABLE: + raise RuntimeError(f"MXFP8 is not available: {REASON_FOR_NO_MXFP8}") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + + dtype = torch.bfloat16 + accumulate_into_main_grad = True + token_dims = parse_int_list(args.token_dims) + num_groups_list = parse_int_list(args.num_groups) + + print("Environment toggles:") + for name in ( + "CUDA_DEVICE_MAX_CONNECTIONS", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", + "NVTE_CUTEDSL_FUSED_GROUPED_MLP", + "CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", + ): + print(f" {name}={os.environ.get(name)}") + print("Recipe: MXFP8BlockScaling") + print("Activation: ScaledSwiGLU") + print(f"Default GLU interleave size: {args.glu_interleave_size}") + print() + + data = [] + for num_groups in num_groups_list: + for total_tokens in token_dims: + timing_ms = benchmark_case( + total_tokens=total_tokens, + hidden_dim=args.hidden_dim, + ffn_hidden_dim=args.ffn_hidden_dim, + num_groups=num_groups, + dtype=dtype, + fwd_only=args.fwd_only, + single_grouped_weight=args.single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + glu_interleave_size=args.glu_interleave_size, + num_microbatches=args.num_microbatches, + min_run_time=args.min_run_time, + profile=args.profile, + ) + data.append( + [ + total_tokens, + args.hidden_dim, + args.ffn_hidden_dim, + num_groups, + args.glu_interleave_size, + args.single_grouped_weight, + accumulate_into_main_grad, + "fwd" if args.fwd_only else "fwd_bwd", + timing_ms, + ] + ) + + timing_col = "time_per_microbatch_ms" + df = pd.DataFrame( + data=data, + columns=[ + "tokens", + "hidden_dim", + "ffn_hidden_dim", + "num_groups", + "glu_interleave_size", + "single_grouped_weight", + "accumulate_into_main_grad", + "mode", + timing_col, + ], + ) + print(df) + + +if __name__ == "__main__": + main() diff --git a/build_tools/jax.py b/build_tools/jax.py index a7b200f915e..35ed62f832f 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -83,6 +83,20 @@ def setup_jax_extension( # Header files include_dirs = get_cuda_include_dirs() + cudnn_frontend_include_dir = None + for base_path in (Path(common_header_files), *Path(common_header_files).parents): + candidate = base_path / "3rdparty" / "cudnn-frontend" / "include" + if candidate.exists(): + cudnn_frontend_include_dir = candidate + break + if cudnn_frontend_include_dir is None: + for base_path in Path(__file__).resolve().parents: + candidate = base_path / "3rdparty" / "cudnn-frontend" / "include" + if candidate.exists(): + cudnn_frontend_include_dir = candidate + break + if cudnn_frontend_include_dir is not None: + include_dirs.append(cudnn_frontend_include_dir) include_dirs.extend( [ common_header_files, @@ -103,6 +117,9 @@ def setup_jax_extension( setup_mpi_flags(include_dirs, cxx_flags) + if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): + cxx_flags.append("-DNVTE_WITH_CUBLASMP") + # Define TE/JAX as a Pybind11Extension from pybind11.setup_helpers import Pybind11Extension diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 533addaf531..e2e6d09c29a 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -6,6 +6,7 @@ import os from pathlib import Path +from importlib import metadata import setuptools @@ -88,6 +89,9 @@ def setup_pytorch_extension( libraries.append("nvshmem_host") cxx_flags.append("-DNVTE_ENABLE_NVSHMEM") + if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): + cxx_flags.append("-DNVTE_WITH_CUBLASMP") + # Construct PyTorch CUDA extension sources = [str(path) for path in sources] include_dirs = [str(path) for path in include_dirs] diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index e9ec854dbaa..ff422e9d2ca 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -44,6 +44,8 @@ if $BUILD_COMMON ; then sed -i "s/Name: transformer-engine/Name: transformer-engine-cu${CUDA_MAJOR}/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer_engine/Name: transformer_engine_cu${CUDA_MAJOR}/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu${CUDA_MAJOR}-${VERSION}.dist-info" + # Set WHEEL Tag to match the py3-none filename written by line 54. + sed -i "s/Tag: cp310-cp310/Tag: py3-none/g" "${WHL_BASE}/transformer_engine_cu${CUDA_MAJOR}-${VERSION}.dist-info/WHEEL" /opt/python/cp310-cp310/bin/wheel pack ${WHL_BASE} # Rename the wheel to make it python version agnostic. diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 99d850d04d1..926724250f1 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -55,6 +55,16 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy +.. autoapifunction:: transformer_engine.pytorch.interleave_glu_tensor + +.. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor + +Data types +---------- + +.. autoapiclass:: transformer_engine.pytorch.DType() + :members: kByte, kInt32, kFloat32, kFloat16, kBFloat16, kFloat8E4M3, kFloat8E5M2, kFloat4E2M1 + Recipe availability ------------------- diff --git a/docs/examples/quickstart_utils.py b/docs/examples/quickstart_utils.py index 9b21807255d..140b6b9af2e 100644 --- a/docs/examples/quickstart_utils.py +++ b/docs/examples/quickstart_utils.py @@ -204,9 +204,8 @@ def share_parameters_with_transformerlayer_te_model(te_model, basic_model): def cast_to_representable(inp, scale=1.0, fp8_format="e4m3"): from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer - import transformer_engine_torch as tex - fp8_type = tex.DType.kFloat8E4M3 if fp8_format == "e4m3" else tex.DType.kFloat8E5M2 + fp8_type = te.DType.kFloat8E4M3 if fp8_format == "e4m3" else te.DType.kFloat8E5M2 scale = torch.ones(1, dtype=torch.float32, device="cuda") * scale amax_history = torch.zeros(1, 1, dtype=torch.float32, device="cuda") quantizer = Float8Quantizer(scale=scale, amax=amax_history, fp8_dtype=fp8_type) diff --git a/examples/jax/collective_gemm/common.py b/examples/jax/collective_gemm/common.py index 68159323959..2bb61f8ab29 100644 --- a/examples/jax/collective_gemm/common.py +++ b/examples/jax/collective_gemm/common.py @@ -4,11 +4,14 @@ """Shared functions for the collective GEMM tests""" import argparse +import glob +import os import jax import jax.numpy as jnp import numpy as np from jax.experimental import mesh_utils +from jax.experimental.multihost_utils import sync_global_devices from transformer_engine.jax.cpp_extensions.gemm import collective_gemm_bootstrap @@ -56,9 +59,9 @@ def assert_allclose(actual, desired, rtol=None, atol=None, dtype=None, **kwargs) tols["atol"] = atol if not isinstance(actual, float): - actual = actual.astype(jnp.float32) + actual = np.asarray(actual, dtype=np.float32) if not isinstance(desired, float): - desired = desired.astype(jnp.float32) + desired = np.asarray(desired, dtype=np.float32) np.testing.assert_allclose(actual, desired, **tols, **kwargs) @@ -118,6 +121,20 @@ def _initialize_distributed(args): devices_per_process = 1 num_total_devices = args.num_processes + # Remove stale NCCL unique ID files from previous (possibly crashed) runs. + # These files are used for one-time coordination during bootstrap; stale files + # cause non-leader processes to read an old unique ID, breaking NCCL init. + # Only process 0 performs the cleanup; a global barrier ensures all processes + # wait for the cleanup to complete before any TP leader writes a fresh file. + nccl_base_path = os.environ.get("NVTE_JAX_NCCL_FILE_PATH", "/tmp") + if args.process_id == 0: + for f in glob.glob(os.path.join(nccl_base_path, "nccl_*_unique_id_*.bin")): + try: + os.remove(f) + except OSError: + pass + sync_global_devices("nccl_id_cleanup") + print( f"Initializing CGEMM communicator with num_total_devices={num_total_devices}," f" devices_per_process={devices_per_process}, process_id={args.process_id}" @@ -128,6 +145,7 @@ def _initialize_distributed(args): num_devices_per_process=devices_per_process, process_id=args.process_id, tensor_parallel_size=args.tensor_parallel_size, + use_cublasmp=args.use_cublasmp, ) @@ -199,6 +217,16 @@ def cgemm_parser(description="Collective GEMM test on multi-GPU with tensor para parser.add_argument("--seq-len", type=int, default=8192, help="Sequence length for testing") parser.add_argument("--hidden-in", type=int, default=4096, help="Input hidden dimension") parser.add_argument("--hidden-out", type=int, default=8192, help="Output hidden dimension") + parser.add_argument( + "--std", + type=float, + default=0.023, + help=( + "Standard deviation for input/weight/bias tensors. Matches TE/PyTorch's" + " run_gemm_with_overlap.py default so both frameworks evaluate FP8 noise" + " on equal footing." + ), + ) parser.add_argument( "--collective-type", type=str, @@ -224,5 +252,11 @@ def cgemm_parser(description="Collective GEMM test on multi-GPU with tensor para parser.add_argument( "--enable-result-check", action="store_true", default=True, help="Enable result checking" ) + parser.add_argument( + "--use-cublasmp", + action="store_true", + default=False, + help="Use the cuBLASMp backend for overlapping collective operations with GEMM computation", + ) return parser diff --git a/examples/jax/collective_gemm/conftest.py b/examples/jax/collective_gemm/conftest.py index 5be5709ba7d..d8ffbc28530 100644 --- a/examples/jax/collective_gemm/conftest.py +++ b/examples/jax/collective_gemm/conftest.py @@ -5,6 +5,9 @@ """config for collective_gemm tests""" import pytest +import transformer_engine.jax # noqa: F401 - must load libtransformer_engine.so before transformer_engine_jax +from transformer_engine_jax import nvte_built_with_cublasmp + def pytest_addoption(parser): """Pytest hook for collective_gemm tests""" @@ -12,12 +15,21 @@ def pytest_addoption(parser): parser.addoption("--num-processes", action="store", default=1) parser.addoption("--process-id", action="store", default=0) parser.addoption("--local-device-ids", action="store", default=None) + parser.addoption("--use-cublasmp", action="store_true", default=False) @pytest.fixture(autouse=True) def distributed_args(request): """Fixture for querying distributed initialization arguments""" if request.cls: + use_cublasmp = request.config.getoption("--use-cublasmp") + if use_cublasmp and not nvte_built_with_cublasmp(): + pytest.skip( + "Collective GEMM cuBLASMp backend tests require Transformer Engine to be built " + "with NVTE_WITH_CUBLASMP=1." + ) + if use_cublasmp and "mxfp8" in request.node.name.lower(): + pytest.skip("MXFP8 is not supported by the cuBLASMp backend wrappers in TE/common.") request.cls.coordinator_address = request.config.getoption("--coordinator-address") request.cls.num_processes = int(request.config.getoption("--num-processes")) request.cls.process_id = int(request.config.getoption("--process-id")) @@ -27,3 +39,4 @@ def distributed_args(request): if request.cls.local_device_ids is None else len(request.cls.local_device_ids.split(",")) ) + request.cls.use_cublasmp = use_cublasmp diff --git a/examples/jax/collective_gemm/run_test_cgemm.sh b/examples/jax/collective_gemm/run_test_cgemm.sh index c0a095d6b5a..6c2f8e9b8a9 100644 --- a/examples/jax/collective_gemm/run_test_cgemm.sh +++ b/examples/jax/collective_gemm/run_test_cgemm.sh @@ -23,6 +23,30 @@ else echo "NVLINK support detected" fi +echo "*** Checking cuBLASMp support in TE build ***" +CUBLASMP_SUPPORT=$(python3 - <<'PY' +try: + import transformer_engine.jax + from transformer_engine_jax import nvte_built_with_cublasmp +except Exception as exc: + print(f"error:{exc}") + raise SystemExit(0) + +print("1" if nvte_built_with_cublasmp() else "0") +PY +) + +if [[ "$CUBLASMP_SUPPORT" == "1" ]]; then + echo "cuBLASMp backend support detected" + BACKENDS=("cublasmp" "userbuffers") +elif [[ "$CUBLASMP_SUPPORT" == "0" ]]; then + echo "cuBLASMp backend support not detected; skipping cuBLASMp backend tests" + BACKENDS=("userbuffers") +else + echo "Failed to query cuBLASMp support from transformer_engine_jax: $CUBLASMP_SUPPORT" + exit 1 +fi + # Define individual test cases to run (file::class::method) # DelayedScalingFP8 and CurrentScalingFP8 use the same GEMM so we don't need to test both cases all # the time. @@ -93,50 +117,62 @@ for TEST_CASE in "${TEST_CASES[@]}"; do # Clear PIDs array for this test case PIDS=() - for i in $(seq 0 $(($NUM_GPUS - 1))); do - # Define output file for logs - LOG_FILE="${TEST_NAME}_gpu_${i}.log" - - if [ $i -eq 0 ]; then - # For process 0: show live output AND save to log file using tee - echo "=== Live output from process 0 ===" - pytest -s -c "$TE_PATH/tests/jax/pytest.ini" \ - -vs --junitxml=$XML_LOG_DIR/collective_gemm_${TEST_NAME}.xml \ - "$TE_PATH/examples/jax/collective_gemm/$TEST_CASE" \ - --num-processes=$NUM_GPUS \ - --process-id=$i 2>&1 | tee "$LOG_FILE" & - PID=$! - PIDS+=($PID) + for BACKEND in "${BACKENDS[@]}"; do + echo "Setting backend to $BACKEND for test $TEST_NAME" + + for i in $(seq 0 $(($NUM_GPUS - 1))); do + # Define output file for logs + LOG_FILE="${TEST_NAME}_gpu_${i}_${BACKEND}.log" + + test_args=( + "--num-processes=$NUM_GPUS" + "--process-id=$i" + ) + if [ "$BACKEND" == "cublasmp" ]; then + test_args+=("--use-cublasmp") + fi + + if [ $i -eq 0 ]; then + # For process 0: show live output AND save to log file using tee + echo "=== Live output from process 0 ===" + pytest -s -c "${TE_PATH}/tests/jax/pytest.ini" -vs \ + "--junitxml=${XML_LOG_DIR}/${TEST_NAME}_gpu_${i}_${BACKEND}.xml" \ + "${TE_PATH}/examples/jax/collective_gemm/${TEST_CASE}" \ + "${test_args[@]}" 2>&1 | tee "$LOG_FILE" & + PID=$! + PIDS+=($PID) + else + # For other processes: redirect to log files only + pytest -s -c "${TE_PATH}/tests/jax/pytest.ini" -vs \ + "${TE_PATH}/examples/jax/collective_gemm/${TEST_CASE}" \ + "${test_args[@]}" > "$LOG_FILE" 2>&1 & + PID=$! + PIDS+=($PID) + fi + done + + # Wait for all processes to finish + wait + + # Check and print the log content from process 0 + if grep -q "SKIPPED" "${TEST_NAME}_gpu_0_${BACKEND}.log"; then + echo "... $TEST_CASE SKIPPED" + elif grep -q "FAILED" "${TEST_NAME}_gpu_0_${BACKEND}.log"; then + echo "... $TEST_CASE FAILED" + HAS_FAILURE=1 + elif grep -q "PASSED" "${TEST_NAME}_gpu_0_${BACKEND}.log"; then + echo "... $TEST_CASE PASSED" else - # For other processes: redirect to log files only - pytest -s -c "$TE_PATH/tests/jax/pytest.ini" \ - -vs "$TE_PATH/examples/jax/collective_gemm/$TEST_CASE" \ - --num-processes=$NUM_GPUS \ - --process-id=$i > "$LOG_FILE" 2>&1 & - PID=$! - PIDS+=($PID) + echo "... $TEST_CASE INVALID" + HAS_FAILURE=1 fi - done - # Wait for all processes to finish - wait - - # Check and print the log content from process 0 - if grep -q "SKIPPED" "${TEST_NAME}_gpu_0.log"; then - echo "... $TEST_CASE SKIPPED" - elif grep -q "FAILED" "${TEST_NAME}_gpu_0.log"; then - echo "... $TEST_CASE FAILED" - HAS_FAILURE=1 - elif grep -q "PASSED" "${TEST_NAME}_gpu_0.log"; then - echo "... $TEST_CASE PASSED" - else - echo "... $TEST_CASE INVALID" - HAS_FAILURE=1 - fi - - # Remove the log files after processing them - wait - rm ${TEST_NAME}_gpu_*.log + + # Remove the log files after processing them + wait + rm ${TEST_NAME}_gpu_*_${BACKEND}.log + + done done wait diff --git a/examples/jax/collective_gemm/test_dense_grad.py b/examples/jax/collective_gemm/test_dense_grad.py index 1d300f8e905..66f2870d491 100644 --- a/examples/jax/collective_gemm/test_dense_grad.py +++ b/examples/jax/collective_gemm/test_dense_grad.py @@ -95,11 +95,14 @@ def run_dense_grad_tests(args, mesh=None): # Create test data rng = jax.random.PRNGKey(0) rng, x_rng, weight_rng, bias_rng = jax.random.split(rng, 4) - x = jax.random.normal( + std = jnp.asarray(args.std, dtype=jnp.bfloat16) + x = std * jax.random.normal( x_rng, (args.batch_size, args.seq_len, args.hidden_in), dtype=jnp.bfloat16 ) - weight = jax.random.normal(weight_rng, (args.hidden_in, args.hidden_out), dtype=jnp.bfloat16) - bias = jax.random.normal(bias_rng, (args.hidden_out,), dtype=jnp.bfloat16) + weight = std * jax.random.normal( + weight_rng, (args.hidden_in, args.hidden_out), dtype=jnp.bfloat16 + ) + bias = std * jax.random.normal(bias_rng, (args.hidden_out,), dtype=jnp.bfloat16) collective_op = ( CollectiveOp.ALL_GATHER @@ -183,6 +186,7 @@ def setUp(self): self.args.process_id = self.process_id self.args.local_device_ids = self.local_device_ids self.args.num_devices_per_process = self.num_devices_per_process + self.args.use_cublasmp = self.use_cublasmp self.args.enable_data_parallel = True self.args.tensor_parallel_size = _get_dp_and_tp_sizes(self.args)[1] _initialize_distributed(self.args) diff --git a/examples/jax/collective_gemm/test_gemm.py b/examples/jax/collective_gemm/test_gemm.py index 8221d7bbfd0..a82ee1c0424 100644 --- a/examples/jax/collective_gemm/test_gemm.py +++ b/examples/jax/collective_gemm/test_gemm.py @@ -16,6 +16,7 @@ import os from functools import partial +import numpy as np import jax import jax.numpy as jnp from jax.sharding import PartitionSpec, NamedSharding @@ -86,11 +87,14 @@ def run_gemm_tests(args, mesh=None): # Create test data rng = jax.random.PRNGKey(0) rng, x_rng, weight_rng, bias_rng = jax.random.split(rng, 4) - x = jax.random.normal( + std = jnp.asarray(args.std, dtype=jnp.bfloat16) + x = std * jax.random.normal( x_rng, (args.batch_size, args.seq_len, args.hidden_in), dtype=jnp.bfloat16 ) - weight = jax.random.normal(weight_rng, (args.hidden_in, args.hidden_out), dtype=jnp.bfloat16) - bias = jax.random.normal(bias_rng, (args.hidden_out,), dtype=jnp.bfloat16) + weight = std * jax.random.normal( + weight_rng, (args.hidden_in, args.hidden_out), dtype=jnp.bfloat16 + ) + bias = std * jax.random.normal(bias_rng, (args.hidden_out,), dtype=jnp.bfloat16) collective_op = ( CollectiveOp.ALL_GATHER if args.collective_type == "all_gather" @@ -151,20 +155,27 @@ def run_gemm_tests(args, mesh=None): jax.block_until_ready(gathered_output) if args.enable_result_check and args.process_id == 0: - # CGEMM + RS + BF16 uses TE's reduce_bf16 kernel (sequential left-to-right in FP32). - # With catastrophic cancellation the output is near zero while the absolute diff can - # reach 1 ULP of the partial GEMM magnitude (~0.0625 for typical transformer - # activations at O(8) scale), which exceeds the previous atol=1e-5. The 2x - # margin (0.125) covers this worst-case 1-ULP absolute difference. - is_cgemm_rs_bf16 = collective_op == CollectiveOp.REDUCE_SCATTER and not use_quantization - rtol = 1e-2 if is_cgemm_rs_bf16 else None - atol = 0.125 if is_cgemm_rs_bf16 else None - assert_allclose( - gathered_ref_output, - gathered_output, - dtype=get_tolerance_dtype(quantizer_set), - rtol=rtol, - atol=atol, + if use_quantization: + # FP8 quantization noise on near-zero outputs can exceed the rtol + # gate; allow a small absolute tolerance. + rtol, atol = 0.125, 0.625 + else: + rtol, atol = 0.02, 0.002 + # Use NumPy (not JAX) for the result check to avoid triggering new XLA compilations + # on process 0 only, which would deadlock in multi-process JAX because XLA compilation + # of distributed arrays requires collective synchronization across all processes. + actual = np.asarray(gathered_output, dtype=np.float32) + desired = np.asarray(gathered_ref_output, dtype=np.float32) + diff = np.abs(actual - desired) + abs_desired = np.abs(desired) + failures = (diff > atol) & (diff > rtol * abs_desired) + num_failures = int(np.sum(failures)) + assert num_failures == 0, ( + f"NUMERICAL CHECK FAILED: {num_failures}/{diff.size} elements " + f"({100 * num_failures / diff.size:.4f}%) exceed tolerances " + f"(rtol={rtol}, atol={atol}). " + f"Max abs error: {float(np.max(diff)):.6f}, " + f"max rel error: {float(np.max(diff / np.maximum(abs_desired, 1e-5))):.6f}" ) @@ -180,6 +191,7 @@ def setUp(self): self.args.process_id = self.process_id self.args.local_device_ids = self.local_device_ids self.args.num_devices_per_process = self.num_devices_per_process + self.args.use_cublasmp = self.use_cublasmp self.args.enable_data_parallel = True self.args.tensor_parallel_size = _get_dp_and_tp_sizes(self.args)[1] _initialize_distributed(self.args) diff --git a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py index be94c68d37a..a5ba370fd23 100644 --- a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py +++ b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py @@ -139,19 +139,26 @@ def run_layernorm_mlp_grad_tests(args, mesh=None): rng, x_rng, weight_1_rng, bias_1_rng, weight_2_rng, bias_2_rng, gamma_rng = jax.random.split( rng, 7 ) - x = jax.random.normal( + std = jnp.asarray(args.std, dtype=jnp.bfloat16) + x = std * jax.random.normal( x_rng, (args.batch_size, args.seq_len, args.hidden_in), dtype=jnp.bfloat16 ) - weight_1 = jax.random.normal( - weight_1_rng, (args.hidden_in, 1, args.hidden_out), dtype=jnp.bfloat16 - ) / jnp.sqrt(args.hidden_in) - bias_1 = jax.random.normal(bias_1_rng, (1, args.hidden_out), dtype=jnp.bfloat16) - weight_2 = jax.random.normal( - weight_2_rng, (args.hidden_out, args.hidden_in), dtype=jnp.bfloat16 - ) / jnp.sqrt(args.hidden_out) - bias_2 = jax.random.normal(bias_2_rng, (args.hidden_in,), dtype=jnp.bfloat16) - gamma = jax.random.normal(gamma_rng, (args.hidden_in,), dtype=jnp.bfloat16) / jnp.sqrt( - args.hidden_in + weight_1 = ( + std + * jax.random.normal(weight_1_rng, (args.hidden_in, 1, args.hidden_out), dtype=jnp.bfloat16) + / jnp.sqrt(args.hidden_in) + ) + bias_1 = std * jax.random.normal(bias_1_rng, (1, args.hidden_out), dtype=jnp.bfloat16) + weight_2 = ( + std + * jax.random.normal(weight_2_rng, (args.hidden_out, args.hidden_in), dtype=jnp.bfloat16) + / jnp.sqrt(args.hidden_out) + ) + bias_2 = std * jax.random.normal(bias_2_rng, (args.hidden_in,), dtype=jnp.bfloat16) + gamma = ( + std + * jax.random.normal(gamma_rng, (args.hidden_in,), dtype=jnp.bfloat16) + / jnp.sqrt(args.hidden_in) ) collective_op_set_1 = CollectiveOpSet.create(forward_collective_op=CollectiveOp.ALL_GATHER) collective_op_set_2 = CollectiveOpSet.create(forward_collective_op=CollectiveOp.REDUCE_SCATTER) @@ -249,6 +256,7 @@ def setUp(self): self.args.process_id = self.process_id self.args.local_device_ids = self.local_device_ids self.args.num_devices_per_process = self.num_devices_per_process + self.args.use_cublasmp = self.use_cublasmp self.args.enable_data_parallel = True self.args.tensor_parallel_size = _get_dp_and_tp_sizes(self.args)[1] _initialize_distributed(self.args) diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index 12bb027b9e1..9c020ddd33f 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -29,6 +29,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" mkdir -p "$XML_LOG_DIR" 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_*" +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_score_mod.xml $TE_PATH/tests/jax/test_fused_attn_score_mod.py || test_fail "tests/jax/test_fused_attn_score_mod.py" NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_with_determinism.xml $TE_PATH/tests/jax/test_fused_attn.py -k "TestFusedAttnWithDeterminism" || test_fail "tests/jax/test_fused_attn.py" pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements" diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 2d3f75f293d..2485e84b5e2 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -50,7 +50,9 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_e python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py" NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 0d7258a81d4..44ad7c73841 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -21,6 +21,8 @@ project(transformer_engine_distributed_tests LANGUAGES CUDA CXX) add_subdirectory(../../3rdparty/googletest ${PROJECT_BINARY_DIR}/googletest) +enable_testing() + include_directories(${gtest_SOURCE_DIR}/include ${gtest_SOURCE_DIR}) if(NOT DEFINED TE_LIB_PATH) @@ -30,28 +32,45 @@ if(NOT DEFINED TE_LIB_PATH) get_filename_component(TE_LIB_PATH ${TE_LIB_FILE} DIRECTORY) endif() -find_library(TE_LIB NAMES transformer_engine PATHS "${TE_LIB_PATH}/.." ${TE_LIB_PATH} ENV TE_LIB_PATH REQUIRED) - +find_library(TE_LIB + NAMES transformer_engine + PATHS "${TE_LIB_PATH}/.." ${TE_LIB_PATH} + ENV TE_LIB_PATH + REQUIRED) message(STATUS "Found transformer_engine library: ${TE_LIB}") -include_directories(../../transformer_engine/common/include) -include_directories(../../transformer_engine/common) -include_directories(../../transformer_engine) -include_directories(${CMAKE_SOURCE_DIR}) - -find_package(CUDAToolkit REQUIRED) add_executable(test_comm_gemm test_comm_gemm.cu ../cpp/test_common.cu) +list(APPEND test_comm_gemm_INCLUDES + ${CMAKE_SOURCE_DIR}/../../transformer_engine/common/include + ${CMAKE_SOURCE_DIR}/../../transformer_engine/common + ${CMAKE_SOURCE_DIR}/../../transformer_engine + ${CMAKE_SOURCE_DIR} + ${MPI_CXX_INCLUDE_PATH} + $ENV{CUBLASMP_HOME}/include) +target_include_directories(test_comm_gemm PRIVATE ${test_comm_gemm_INCLUDES}) + +find_package(CUDAToolkit REQUIRED) find_package(OpenMP REQUIRED) find_package(MPI REQUIRED) find_library(NCCL_LIB NAMES nccl libnccl PATH_SUFFIXES lib REQUIRED) -target_include_directories(test_comm_gemm PRIVATE ${MPI_CXX_INCLUDE_PATH} $ENV{CUBLASMP_HOME}/include) -target_link_libraries(test_comm_gemm PUBLIC CUDA::cuda_driver CUDA::cudart GTest::gtest ${TE_LIB} CUDA::nvrtc MPI::MPI_CXX ${NCCL_LIB} OpenMP::OpenMP_CXX) +list(APPEND test_comm_gemm_LINKER_LIBS + CUDA::cuda_driver + CUDA::cudart + GTest::gtest_main + ${TE_LIB} + CUDA::nvrtc + ${NCCL_LIB} + OpenMP::OpenMP_CXX + MPI::MPI_CXX) +target_link_libraries(test_comm_gemm PUBLIC ${test_comm_gemm_LINKER_LIBS}) + +target_compile_options(test_comm_gemm PRIVATE -O2 -fopenmp) include(GoogleTest) gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) diff --git a/tests/cpp_distributed/test_comm_gemm.cu b/tests/cpp_distributed/test_comm_gemm.cu index 45f6664567e..b1c4a9395a9 100644 --- a/tests/cpp_distributed/test_comm_gemm.cu +++ b/tests/cpp_distributed/test_comm_gemm.cu @@ -4,6 +4,7 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include #include @@ -104,6 +105,96 @@ std::vector CopyMatrix(const std::vector& data, size_t mstart, size_t nsta return ret; } +template +std::vector CastToFloat(const std::vector& in) { + std::vector out(in.size()); + for (size_t i = 0; i < in.size(); ++i) { + out[i] = static_cast(in[i]); + } + return out; +} + +template +ncclDataType_t NcclDataType(); + +template <> +ncclDataType_t NcclDataType() { + return ncclFloat16; +} + +template <> +ncclDataType_t NcclDataType() { + return ncclBfloat16; +} + +template <> +ncclDataType_t NcclDataType() { + return ncclFloat; +} + +template <> +ncclDataType_t NcclDataType() { + return ncclFloat8e4m3; +} + +template <> +ncclDataType_t NcclDataType() { + return ncclFloat8e5m2; +} + + + +template +std::vector AllGatherColsSharded(const std::vector& local, int64_t rows, int64_t local_cols, + int64_t global_cols) { + std::vector cols_per_rank{}; + int nranks{}; + CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks)); + cols_per_rank.resize(nranks); + CHECK_MPI(MPI_Allgather(&local_cols, 1, MPI_INT64_T, cols_per_rank.data(), 1, MPI_INT64_T, + MPI_COMM_WORLD)); + + std::vector counts(nranks); + std::vector displs(nranks, 0); + for (int r = 0; r < nranks; ++r) { + counts[r] = static_cast(cols_per_rank[r] * rows * sizeof(T)); + if (r > 0) displs[r] = displs[r - 1] + counts[r - 1]; + } + + std::vector gathered(rows * global_cols); + CHECK_MPI(MPI_Allgatherv(local.data(), static_cast(local.size() * sizeof(T)), MPI_BYTE, + gathered.data(), counts.data(), displs.data(), MPI_BYTE, + MPI_COMM_WORLD)); + return gathered; +} + +template +std::vector AllGatherRowsSharded(const std::vector& local, int64_t local_rows, int64_t cols, + int64_t global_rows) { + std::vector rows_per_rank{}; + int nranks{}; + CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks)); + rows_per_rank.resize(nranks); + CHECK_MPI(MPI_Allgather(&local_rows, 1, MPI_INT64_T, rows_per_rank.data(), 1, MPI_INT64_T, + MPI_COMM_WORLD)); + + std::vector counts(nranks); + std::vector displs(nranks, 0); + for (int r = 0; r < nranks; ++r) { + counts[r] = static_cast(rows_per_rank[r] * sizeof(T)); + if (r > 0) displs[r] = displs[r - 1] + counts[r - 1]; + } + + std::vector gathered(global_rows * cols); + for (int64_t col = 0; col < cols; ++col) { + const auto* send = reinterpret_cast(local.data() + col * local_rows); + auto* recv = reinterpret_cast(gathered.data() + col * global_rows); + CHECK_MPI(MPI_Allgatherv(send, static_cast(local_rows * sizeof(T)), MPI_BYTE, recv, + counts.data(), displs.data(), MPI_BYTE, MPI_COMM_WORLD)); + } + return gathered; +} + template test::Tensor Make(size_t m, size_t n, float scale) { test::Tensor ret("", std::vector{n, m}, TypeInfo::dtype); @@ -148,6 +239,12 @@ struct Params { class CommGemmFixure : public ::testing::TestWithParam { protected: + enum class OverlapType { + kAllGather, + kReduceScatter, + kAllReduce, + }; + CommGemmFixure() { CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks_)); CHECK_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank_)); @@ -180,6 +277,8 @@ class CommGemmFixure : public ::testing::TestWithParam { virtual PatternDims DistributeTensors(int64_t m, int64_t n, int64_t k) = 0; + virtual OverlapType overlap_type() const = 0; + virtual void CommGemm(int64_t m, int64_t n, int64_t k, const NVTETensor a, const NVTETensor b, const NVTETensor d, const NVTETensor bias, const NVTETensor pre_act_out, bool transa, bool transb, bool grad, bool accumulate, int comm_sm_count, @@ -213,14 +312,6 @@ class CommGemmFixure : public ::testing::TestWithParam { return static_cast(dist(rng) * bias_scale); }); - auto ga = transa ? MakeFromData(adata, 0, 0, k, m, k, a_scale) - : MakeFromData(adata, 0, 0, m, k, m, a_scale); - auto gb = transb ? MakeFromData(bdata, 0, 0, n, k, n, b_scale) - : MakeFromData(bdata, 0, 0, k, n, k, b_scale); - auto gbias = MakeFromData(biasdata, 0, 0, m, 1, m, bias_scale); - auto gd = Make(m, n, d_scale); - auto gaux = Make(m, n, d_scale); - auto dims = DistributeTensors(m, n, k); auto a = transa ? MakeFromData(adata, dims.a_rows_start, dims.a_cols_start, dims.a_rows_num, dims.a_cols_num, k, a_scale) @@ -240,24 +331,88 @@ class CommGemmFixure : public ::testing::TestWithParam { CommGemm(m, n, k, a.data(), b.data(), d.data(), bias.data(), aux.data(), transa, transb, grad, accumulate, 0 /*comm_sm_count*/, stream); auto workspace = Make(1, 32 << 20, 1.0); - nvte_cublas_gemm(ga.data(), gb.data(), gd.data(), gbias.data(), gaux.data(), transa, transb, - grad, workspace.data(), accumulate, true /* use_split_accumulator */, - 0 /* math_sm_count */, stream); + + std::vector out_golden{}; + if (overlap_type() == OverlapType::kAllGather) { + // Build AG reference input by explicitly all-gathering the sharded input operand. + std::vector b_global_data{}; + if (transb) { + auto b_local = CopyMatrix(bdata, dims.b_cols_start, dims.b_rows_start, dims.b_cols_num, + dims.b_rows_num, n); + b_global_data = AllGatherRowsSharded(b_local, dims.b_cols_num, k, n); + } else { + auto b_local = CopyMatrix(bdata, dims.b_rows_start, dims.b_cols_start, dims.b_rows_num, + dims.b_cols_num, k); + b_global_data = AllGatherColsSharded(b_local, k, dims.b_cols_num, n); + } + + auto b_ref = + transb ? MakeFromData(b_global_data, 0, 0, n, k, n, b_scale) + : MakeFromData(b_global_data, 0, 0, k, n, k, b_scale); + auto d_ref = Make(dims.d_rows_num, dims.d_cols_num, d_scale); + auto aux_ref = Make(dims.d_rows_num, dims.d_cols_num, d_scale); + nvte_cublas_gemm(a.data(), b_ref.data(), d_ref.data(), bias.data(), aux_ref.data(), transa, + transb, grad, workspace.data(), accumulate, + true /* use_split_accumulator */, 0 /* math_sm_count */, stream); + + std::vector out_ref(dims.d_rows_num * dims.d_cols_num); + NVTE_CHECK_CUDA(cudaMemcpy(out_ref.data(), d_ref.rowwise_dptr(), + out_ref.size() * sizeof(out_ref[0]), cudaMemcpyDefault)); + out_golden = CastToFloat(out_ref); + } else { + transformer_engine::TensorWrapper empty_bias(nullptr, std::vector{0}, + TypeInfo::dtype); + auto d_partial = Make(m, n, d_scale); + auto aux_partial = Make(m, n, d_scale); + nvte_cublas_gemm(a.data(), b.data(), d_partial.data(), empty_bias.data(), + aux_partial.data(), transa, transb, grad, workspace.data(), accumulate, + true /* use_split_accumulator */, 0 /* math_sm_count */, stream); + + std::vector partial_host(m * n); + NVTE_CHECK_CUDA(cudaMemcpy(partial_host.data(), d_partial.rowwise_dptr(), + partial_host.size() * sizeof(partial_host[0]), + cudaMemcpyDefault)); + std::vector partial_float = CastToFloat(partial_host); + + auto d_partial_float = Make(m, n, 1.0f); + NVTE_CHECK_CUDA(cudaMemcpy(d_partial_float.rowwise_dptr(), partial_float.data(), + partial_float.size() * sizeof(float), cudaMemcpyDefault)); + + std::vector reduced; + if (overlap_type() == OverlapType::kReduceScatter) { + auto d_reduced_float = Make(dims.d_rows_num, dims.d_cols_num, 1.0f); + CHECK_NCCL(ncclReduceScatter(d_partial_float.rowwise_dptr(), d_reduced_float.rowwise_dptr(), + dims.d_rows_num * dims.d_cols_num, ncclFloat, ncclSum, + comm_, stream)); + + reduced.resize(dims.d_rows_num * dims.d_cols_num); + NVTE_CHECK_CUDA(cudaMemcpy(reduced.data(), d_reduced_float.rowwise_dptr(), + reduced.size() * sizeof(float), cudaMemcpyDefault)); + } else { + CHECK_NCCL(ncclAllReduce(d_partial_float.rowwise_dptr(), d_partial_float.rowwise_dptr(), + m * n, ncclFloat, ncclSum, comm_, stream)); + + reduced.resize(m * n); + NVTE_CHECK_CUDA(cudaMemcpy(reduced.data(), d_partial_float.rowwise_dptr(), + reduced.size() * sizeof(float), cudaMemcpyDefault)); + } + + for (size_t col = 0; col < dims.d_cols_num; ++col) { + for (size_t row = 0; row < m; ++row) { + reduced[col * m + row] += static_cast(biasdata[row]); + } + } + out_golden = std::move(reduced); + } + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); std::vector out(dims.d_rows_num * dims.d_cols_num); NVTE_CHECK_CUDA( cudaMemcpy(out.data(), d.rowwise_dptr(), out.size() * sizeof out[0], cudaMemcpyDefault)); - std::vector out_golden_global(m * n); - NVTE_CHECK_CUDA(cudaMemcpy(out_golden_global.data(), gd.rowwise_dptr(), - out_golden_global.size() * sizeof out_golden_global[0], - cudaMemcpyDefault)); - - auto out_golden = CopyMatrix(out_golden_global, dims.d_rows_start, dims.d_cols_start, - dims.d_rows_num, dims.d_cols_num, m); NVTE_CHECK(out.size() == out_golden.size()); for (size_t i = 0; i < out.size(); ++i) { - EXPECT_NEAR(static_cast(out[i]), static_cast(out_golden[i]), tol); + EXPECT_NEAR(static_cast(out[i]), out_golden[i], tol); } } @@ -268,6 +423,8 @@ class CommGemmFixure : public ::testing::TestWithParam { }; struct AgGemm : public CommGemmFixure { + OverlapType overlap_type() const override { return OverlapType::kAllGather; } + PatternDims DistributeTensors(int64_t m, int64_t n, int64_t k) override { auto a_cols_num = nvte_comm_gemm_numroc(ctx_, m); auto b_cols_num = nvte_comm_gemm_numroc(ctx_, n); @@ -303,6 +460,8 @@ struct AgGemm : public CommGemmFixure { }; struct GemmRs : public CommGemmFixure { + OverlapType overlap_type() const override { return OverlapType::kReduceScatter; } + PatternDims DistributeTensors(int64_t m, int64_t n, int64_t k) override { auto rows_num = nvte_comm_gemm_numroc(ctx_, k); auto d_cols_num = nvte_comm_gemm_numroc(ctx_, n); @@ -337,7 +496,10 @@ struct GemmRs : public CommGemmFixure { } }; +#if CUBLASMP_VERSION >= 900 struct GemmAr : public CommGemmFixure { + OverlapType overlap_type() const override { return OverlapType::kAllReduce; } + PatternDims DistributeTensors(int64_t m, int64_t n, int64_t k) override { auto rows_num = nvte_comm_gemm_numroc(ctx_, k); @@ -368,6 +530,7 @@ struct GemmAr : public CommGemmFixure { accumulate, comm_sm_count, stream, kNVTECommGemmAlgoDefault); } }; +#endif // CUBLASMP_VERSION >= 900 TEST_P(AgGemm, Gemm) { auto [a_type, b_type, d_type, transa, transb, m, n, k, tol] = GetParam(); @@ -389,6 +552,7 @@ TEST_P(GemmRs, Gemm) { d_type, DType, Run(transa, transb, m, n, k, tol);))); } +#if CUBLASMP_VERSION >= 900 TEST_P(GemmAr, Gemm) { auto [a_type, b_type, d_type, transa, transb, m, n, k, tol] = GetParam(); TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( @@ -398,6 +562,7 @@ TEST_P(GemmAr, Gemm) { TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( d_type, DType, Run(transa, transb, m, n, k, tol);))); } +#endif // CUBLASMP_VERSION >= 900 std::string ParamSuffix(const testing::TestParamInfo& info) { const auto [a_type, b_type, d_type, transa, transb, m, n, k, _tol] = info.param; @@ -450,12 +615,13 @@ INSTANTIATE_TEST_SUITE_P(GemmRs, GemmRs, DType::kFloat16, true, false, 64, 128, 256, 7e-2}), &ParamSuffix); +#if CUBLASMP_VERSION >= 900 INSTANTIATE_TEST_SUITE_P( GemmAr, GemmAr, testing::Values(Params{DType::kFloat16, DType::kFloat16, DType::kFloat16, true, false, 64, 64 * 4, 64 * 4, 7e-2}, Params{DType::kBFloat16, DType::kBFloat16, DType::kBFloat16, true, false, 64, - 64 * 4, 64 * 4, 1e-3}, + 64 * 4, 64 * 4, 6e-1}, Params{DType::kFloat8E5M2, DType::kFloat8E4M3, DType::kFloat16, true, false, 128, 128 * 4, 128 * 4, 1.5e-1}, Params{DType::kFloat8E4M3, DType::kFloat8E5M2, DType::kFloat16, true, false, @@ -463,3 +629,4 @@ INSTANTIATE_TEST_SUITE_P( Params{DType::kFloat8E4M3, DType::kFloat8E4M3, DType::kFloat16, true, false, 128, 128 * 4, 128 * 4, 1.5e-1}), &ParamSuffix); +#endif // CUBLASMP_VERSION >= 900 diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 50c5de1db77..78fd0ed4871 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -12,7 +12,15 @@ generate_context_parallel_configs_for_attn, generate_collectives_count, ) -from test_fused_attn import FusedAttnRunner, BiasShape, SeqDescFormat +from test_fused_attn import ( + FusedAttnRunner, + BiasShape, + SeqDescFormat, +) +from test_fused_attn_score_mod import ( + ScoreModFusedAttnRunner, + _has_cudnn_frontend_python, +) from utils import pytest_parametrize_wrapper from transformer_engine.jax.attention import ( is_fused_attn_kernel_available, @@ -272,6 +280,54 @@ def test_cross_attn( runner.test_backward() +DISTRIBUTED_SCORE_MOD_DATA_SHAPES = { + "L0": [], + "L1": [(4, 16, 4, 64)], +} + + +@pytest.mark.skipif(not _has_cudnn_frontend_python(), reason="cuDNN Python frontend is required") +class TestDistributedScoreModSelfAttn: + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper("data_shape", DISTRIBUTED_SCORE_MOD_DATA_SHAPES) + @pytest.mark.parametrize("dtype", DTYPES) + def test_softcap_score_mod_with_aux_params_backward( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + dtype, + ): + ScoreModFusedAttnRunner.require_cudnn_frontend() + batch, seqlen, num_heads, head_dim = data_shape + dp_axis = mesh_resource.dp_resource + tp_axis = mesh_resource.tpsp_resource + + if dp_axis is not None: + dp_size = mesh_shape[mesh_axes.index(dp_axis)] + if batch % dp_size != 0: + pytest.skip(f"{batch=} must be divisible by {dp_size=}") + if tp_axis is not None: + tp_size = mesh_shape[mesh_axes.index(tp_axis)] + if num_heads % tp_size != 0: + pytest.skip(f"{num_heads=} must be divisible by {tp_size=}") + + runner = ScoreModFusedAttnRunner.softcap( + batch, + seqlen, + num_heads, + head_dim, + dtype, + number_of_devices=device_count, + mesh_shape=mesh_shape, + mesh_axes=mesh_axes, + mesh_resource=mesh_resource, + ) + runner.test_backward() + + DISTRIBUTED_CONTEXT_SELF_ATTN_LAYOUTS_MASKS = [ pytest.param(QKVLayout.BSHD_BS2HD, AttnMaskType.CAUSAL_MASK, id="BSHD_KVPACKED-CAUSAL"), pytest.param(QKVLayout.BSHD_BSHD_BSHD, AttnMaskType.CAUSAL_MASK, id="BSHD_SEPARATE-CAUSAL"), diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 1fb0108068b..0d1db0b9e1b 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -7,7 +7,7 @@ from dataclasses import dataclass, field from functools import partial from math import sqrt -from typing import Tuple, Optional, Dict +from typing import Any, Callable, Mapping, Tuple, Optional, Dict import random import jax @@ -64,7 +64,7 @@ def init(): yield -@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11)) +@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12)) def general_dot_product_attention( query: ArrayLike, key: ArrayLike, @@ -78,6 +78,7 @@ def general_dot_product_attention( dropout_rate: float, dropout_rng: ArrayLike, dtype: DTypeLike, + score_mod_reference: Optional[Callable[[Array], Array]] = None, ) -> Array: """ Similar to flax.linen.dot_product_attention but with GQA support @@ -106,6 +107,10 @@ def general_dot_product_attention( mask = jnp.expand_dims(mask, axis=-3) logits = jnp.where(mask, jnp.finfo(dtype).min, logits) + if score_mod_reference is not None: + # Kernel tests use NO_MASK; fused_attn rejects mask+score_mod before this reference path. + logits = score_mod_reference(logits.astype(jnp.float32)) + match softmax_type: case AttnSoftmaxType.VANILLA_SOFTMAX: softmax_out = jax.nn.softmax(logits).astype(dtype) @@ -267,6 +272,28 @@ def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs """ JAX native dot product attention implementation """ + score_mod_reference = kwargs.get("score_mod_reference") + if score_mod_reference is not None: + if ( + bias is not None + or softmax_offset is not None + or mask is not None + or dropout_rng is not None + ): + raise ValueError( + "score_mod reference path expects no bias, mask, softmax offset, or dropout." + ) + if kwargs["qkv_layout"] != QKVLayout.BSHD_BSHD_BSHD: + raise ValueError( + "score_mod reference path expects separate BSHD query/key/value tensors." + ) + deterministic = not kwargs["is_training"] + dropout_probability = kwargs["dropout_probability"] + softmax_type = kwargs["softmax_type"] + else: + deterministic = not kwargs["is_training"] + dropout_probability = kwargs["dropout_probability"] + softmax_type = kwargs["softmax_type"] output = general_dot_product_attention( query, key, @@ -274,12 +301,13 @@ def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs softmax_offset, bias, mask, - deterministic=not kwargs["is_training"], + deterministic=deterministic, scale_factor=kwargs["scaling_factor"], - dropout_rate=kwargs["dropout_probability"], - softmax_type=kwargs["softmax_type"], + dropout_rate=dropout_probability, + softmax_type=softmax_type, dropout_rng=dropout_rng, dtype=jnp.float32, + score_mod_reference=score_mod_reference, ) return output.astype(query.dtype) @@ -316,6 +344,32 @@ def customcall_fused_dpa( ).astype(query.dtype) +def test_fused_attn_score_mod_rejects_masks_before_cudnn_frontend(): + """fused_attn rejects score_mod with masks before cuDNN frontend graph building.""" + q = jax.ShapeDtypeStruct((1, 16, 1, 128), jnp.float16) + k = jax.ShapeDtypeStruct((1, 16, 1, 128), jnp.float16) + v = jax.ShapeDtypeStruct((1, 16, 1, 128), jnp.float16) + + def score_mod(_graph, score, _tensors): + return score + + with pytest.raises(ValueError, match="mutually exclusive with attention masks"): + fused_attn( + (q, k, v), + None, + None, + None, + AttnBiasType.NO_BIAS, + AttnMaskType.CAUSAL_MASK, + QKVLayout.BSHD_BSHD_BSHD, + AttnSoftmaxType.VANILLA_SOFTMAX, + 1.0, + 0.0, + True, + score_mod=score_mod, + ) + + class BiasShape(Enum): """ Enum class to represent the different bias shapes used in the fused attention. @@ -372,6 +426,18 @@ class FusedAttnRunner: # dictionary of expected collective comm bytes coll_count_ref: Optional[Dict[str, int]] = None + # Optional score_mod test hooks. + score_mod: Optional[Callable] = None + score_mod_bprop: Optional[Callable] = None + score_mod_tensors: Optional[Mapping[str, Any]] = None + score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None + score_mod_reference: Optional[Callable[[Array], Array]] = None + input_scale: float = 1.0 + doutput_seed: Optional[int] = None + doutput: Optional[Array] = field(init=False, default=None) + rtol: Optional[float] = None + atol: Optional[float] = None + def __post_init__(self): # Reset defaults for num_segments_per_seq if not explicitly passed if self.num_segments_per_seq is None: @@ -379,6 +445,16 @@ def __post_init__(self): self.num_segments_per_seq = 2 else: self.num_segments_per_seq = 1 + if self.doutput_seed is not None: + output_shape = ( + self.batch_size, + self.max_seqlen_q, + self.num_heads_q, + self.head_dim_v, + ) + self.doutput = jax.random.normal( + jax.random.PRNGKey(self.doutput_seed), output_shape, dtype=self.dtype + ) # See https://docs.nvidia.com/deeplearning/cudnn/latest/release-notes.html#cudnn-9-4-0 for known issue # generating zero-length ragged tensors. This setting adjusts the test to avoid the zero-length cases. @@ -513,10 +589,15 @@ def _setup_inputs(self): else: pytest.fail(f"PyTest attempted to use an unrecognized bias_layout = {self.bias_shape}!") - self.q = jax.random.uniform(q_key, q_shape, self.dtype, -1.0) - self.k = jax.random.uniform(k_key, k_shape, self.dtype, -1.0) - self.v = jax.random.uniform(v_key, v_shape, self.dtype, -1.0) - + self.q = (self.input_scale * jax.random.uniform(q_key, q_shape, self.dtype, -1.0)).astype( + self.dtype + ) + self.k = (self.input_scale * jax.random.uniform(k_key, k_shape, self.dtype, -1.0)).astype( + self.dtype + ) + self.v = (self.input_scale * jax.random.uniform(v_key, v_shape, self.dtype, -1.0)).astype( + self.dtype + ) if self.attn_bias_type != AttnBiasType.NO_BIAS: if self.bias_shape == BiasShape._1HSS: self.bias = jax.random.uniform(bias_key, bias_shape, self.dtype, -1.0) @@ -810,7 +891,18 @@ def test_forward(self): """ self._setup_inputs() - args = [self.q, self.k, self.v, self.bias, self.softmax_offset, self.mask, self.dropout_rng] + reference_mask = ( + self.sequence_desciptor if self.score_mod_reference is not None else self.mask + ) + args = [ + self.q, + self.k, + self.v, + self.bias, + self.softmax_offset, + reference_mask, + self.dropout_rng, + ] customcall_args = [ # Put test data onto each GPU for distributed. @@ -837,7 +929,12 @@ def test_forward(self): "context_parallel_strategy": self.cp_strategy, "context_parallel_causal_load_balanced": self.cp_load_balanced, "stripe_size": self.stripe_size, + "score_mod": self.score_mod, + "score_mod_bprop": self.score_mod_bprop, + "score_mod_tensors": self.score_mod_tensors, + "score_mod_bprop_tensors": self.score_mod_bprop_tensors, } + reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference} customcall_fused_dpa_jit = jit( partial(customcall_fused_dpa, **kwargs), @@ -857,7 +954,7 @@ def test_forward(self): primitive_out = customcall_fused_dpa_jit(*customcall_args) primitive_out = self.cp_inverse_reorder_fn(primitive_out) - reference_out = jax_dpa(*args, **kwargs) + reference_out = jax_dpa(*args, **reference_kwargs) if self.is_training and self.dropout_prob > 0.0: return @@ -866,8 +963,20 @@ def test_forward(self): _split_valid_and_invalid(primitive_out, reference_out, self.pad_q) ) - assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype) - assert_allclose(primitive_valid, reference_valid, dtype=self.dtype) + assert_allclose( + primitive_invalid, + jnp.zeros_like(primitive_invalid), + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) + assert_allclose( + primitive_valid, + reference_valid, + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) if self.coll_count_ref is not None: with self.mesh, autocast(mesh_resource=self.mesh_resource): @@ -886,29 +995,48 @@ def test_backward(self): self._setup_inputs() - def grad_func(func, *args, cp_reverse_out=False, **kwargs): + def grad_func( + func, + q, + k, + v, + bias, + softmax_offset, + sequence_descriptor, + dropout_rng, + doutput=None, + cp_reverse_out=False, + **kwargs, + ): # Gradient is small, use a gradient multiplier to amplify the gradient gradient_multiplier = self.max_seqlen_q * self.num_heads_q if self.attn_mask_type.is_causal(): gradient_multiplier /= 10 + output = func(q, k, v, bias, softmax_offset, sequence_descriptor, dropout_rng, **kwargs) + if cp_reverse_out: + output = self.cp_inverse_reorder_fn(output) # Keep only valid result for the gradient - if not cp_reverse_out: - ret_valid = jnp.where( - self.pad_q[..., jnp.newaxis, jnp.newaxis], - 0, - func(*args, **kwargs), - ) - else: - ret_valid = jnp.where( - self.pad_q[..., jnp.newaxis, jnp.newaxis], - 0, - self.cp_inverse_reorder_fn(func(*args, **kwargs)), + ret_valid = jnp.where(self.pad_q[..., jnp.newaxis, jnp.newaxis], 0, output) + if doutput is not None: + return jnp.sum(ret_valid.astype(jnp.float32) * doutput.astype(jnp.float32)).astype( + self.dtype ) return ( jnp.mean(ret_valid.astype(jnp.float32), dtype=jnp.float32) * gradient_multiplier ).astype(self.dtype) - args = [self.q, self.k, self.v, self.bias, self.softmax_offset, self.mask, self.dropout_rng] + reference_mask = ( + self.sequence_desciptor if self.score_mod_reference is not None else self.mask + ) + args = [ + self.q, + self.k, + self.v, + self.bias, + self.softmax_offset, + reference_mask, + self.dropout_rng, + ] customcall_args = [ # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and # THD params once we support those features on CP. @@ -920,6 +1048,19 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): jax.device_put(self.sequence_desciptor, self.seq_desc_sharding), jax.device_put(self.dropout_rng, self.dropout_rng_sharding), ] + input_shardings = [ + self.qkvo_sharding, + self.qkvo_sharding, + self.qkvo_sharding, + self.bias_sharding, + self.softmax_offset_sharding, + self.seq_desc_sharding, + self.dropout_rng_sharding, + ] + if self.doutput is not None: + args.append(self.doutput) + customcall_args.append(jax.device_put(self.doutput, self.qkvo_sharding)) + input_shardings.append(self.qkvo_sharding) kwargs = { "attn_bias_type": self.attn_bias_type, "attn_mask_type": self.attn_mask_type, @@ -933,7 +1074,12 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): "context_parallel_strategy": self.cp_strategy, "context_parallel_causal_load_balanced": self.cp_load_balanced, "stripe_size": self.stripe_size, + "score_mod": self.score_mod, + "score_mod_bprop": self.score_mod_bprop, + "score_mod_tensors": self.score_mod_tensors, + "score_mod_bprop_tensors": self.score_mod_bprop_tensors, } + reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference} # We can compute dBias only for the [1, h, s, s] layout if self.bias_shape == BiasShape._1HSS: @@ -964,21 +1110,13 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): ), arg_nums, ), - in_shardings=( - self.qkvo_sharding, - self.qkvo_sharding, - self.qkvo_sharding, - self.bias_sharding, - self.softmax_offset_sharding, - self.seq_desc_sharding, - self.dropout_rng_sharding, - ), + in_shardings=tuple(input_shardings), out_shardings=(None, grad_shardings), ) jitted_reference = jit( value_and_grad( lambda q, k, v, bias, softmax_offset, *args: grad_func( - jax_dpa, q, k, v, bias, softmax_offset, *args, **kwargs + jax_dpa, q, k, v, bias, softmax_offset, *args, **reference_kwargs ), arg_nums, ) @@ -996,7 +1134,13 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): print_debug_tensor_stats(f"primitive_out", primitive_out) print_debug_tensor_stats(f"reference_grad_valid", reference_out) print_debug_tensor_stats(f"diff_grad", jnp.abs(primitive_out - reference_out)) - assert_allclose(primitive_out, reference_out, dtype=self.dtype) + assert_allclose( + primitive_out, + reference_out, + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) def check_dqkv(primitive, reference, pad, idx): primitive_valid, primitive_invalid, reference_valid, reference_invalid = ( @@ -1009,9 +1153,27 @@ def check_dqkv(primitive, reference, pad, idx): f"diff_grad[{idx}]", jnp.abs(primitive_valid[idx] - reference_valid[idx]) ) - assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype) - assert_allclose(primitive_invalid, reference_invalid, dtype=self.dtype) - assert_allclose(primitive_valid, reference_valid, dtype=self.dtype) + assert_allclose( + primitive_invalid, + jnp.zeros_like(primitive_invalid), + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) + assert_allclose( + primitive_invalid, + reference_invalid, + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) + assert_allclose( + primitive_valid, + reference_valid, + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) primitive_dq, primitive_dk, primitive_dv = primitive_dgrad[:3] reference_dq, reference_dk, reference_dv = reference_dgrad[:3] @@ -1037,6 +1199,8 @@ def check_dqkv(primitive, reference, pad, idx): assert_allclose( jnp.where(bias_mask, primitive_dbias, 0), jnp.zeros_like(primitive_dbias), + rtol=self.rtol, + atol=self.atol, dtype=self.dtype, ) @@ -1044,6 +1208,8 @@ def check_dqkv(primitive, reference, pad, idx): assert_allclose( jnp.where(bias_mask, primitive_dbias, 0), jnp.where(bias_mask, reference_dbias, 0), + rtol=self.rtol, + atol=self.atol, dtype=self.dtype, ) @@ -1051,6 +1217,8 @@ def check_dqkv(primitive, reference, pad, idx): assert_allclose( jnp.where(bias_mask, 0, primitive_dbias), jnp.where(bias_mask, 0, reference_dbias), + rtol=self.rtol, + atol=self.atol, dtype=self.dtype, ) diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py new file mode 100644 index 00000000000..a44d1933350 --- /dev/null +++ b/tests/jax/test_fused_attn_score_mod.py @@ -0,0 +1,755 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Tests for cuDNN frontend score_mod fused attention.""" +from math import sqrt + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import transformer_engine.jax.cpp_extensions.flex_attention as tex_attention +from transformer_engine.jax.attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + QKVLayout, +) +from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config +from transformer_engine.jax.flax import transformer as flax_transformer +from transformer_engine_jax import get_device_compute_capability +from test_fused_attn import FusedAttnRunner, SeqDescFormat + + +_CONFIG_TEST_HEAD_DIM = 128 +_CONFIG_TEST_SCALING_FACTOR = 1.0 / sqrt(_CONFIG_TEST_HEAD_DIM) + + +def _has_cudnn_frontend_python(): + try: + tex_attention._import_cudnn_for_score_mod() + except ImportError: + return False + return True + + +def _score_mod_causal(graph, score, tensors): + cudnn = tex_attention._import_cudnn_for_score_mod() + + row_index = graph.gen_index( + input=score, + axis=2, + compute_data_type=cudnn.data_type.INT32, + ) + row_index.set_data_type(cudnn.data_type.INT32) + col_index = graph.gen_index( + input=score, + axis=3, + compute_data_type=cudnn.data_type.INT32, + ) + col_index.set_data_type(cudnn.data_type.INT32) + keep = graph.cmp_ge( + input=row_index, + comparison=col_index, + compute_data_type=cudnn.data_type.BOOLEAN, + ) + keep.set_data_type(cudnn.data_type.BOOLEAN) + return graph.binary_select(input0=score, input1=tensors["neg_inf"], mask=keep) + + +def _score_mod_causal_bprop(graph, dscore, tensors): + cudnn = tex_attention._import_cudnn_for_score_mod() + + row_index = graph.gen_index( + input=dscore, + axis=2, + compute_data_type=cudnn.data_type.INT32, + ) + row_index.set_data_type(cudnn.data_type.INT32) + col_index = graph.gen_index( + input=dscore, + axis=3, + compute_data_type=cudnn.data_type.INT32, + ) + col_index.set_data_type(cudnn.data_type.INT32) + keep = graph.cmp_ge( + input=row_index, + comparison=col_index, + compute_data_type=cudnn.data_type.BOOLEAN, + ) + keep.set_data_type(cudnn.data_type.BOOLEAN) + return graph.binary_select(input0=dscore, input1=tensors["zero"], mask=keep) + + +def _score_mod_post_scale_bias(graph, score, tensors): + cudnn = tex_attention._import_cudnn_for_score_mod() + + row_index = graph.gen_index( + input=score, + axis=2, + compute_data_type=cudnn.data_type.INT32, + ) + row_index.set_data_type(cudnn.data_type.INT32) + col_index = graph.gen_index( + input=score, + axis=3, + compute_data_type=cudnn.data_type.INT32, + ) + col_index.set_data_type(cudnn.data_type.INT32) + post_scale_bias = graph.sub( + a=row_index, + b=col_index, + compute_data_type=cudnn.data_type.FLOAT, + ) + post_scale_bias.set_data_type(cudnn.data_type.FLOAT) + return graph.add( + a=score, + b=post_scale_bias, + compute_data_type=cudnn.data_type.FLOAT, + ) + + +class _ScoreModSoftcap: + """cuDNN frontend score_mod implementing softcapping.""" + + def __init__(self): + self.before_tanh_activation = None + + def score_mod_graph_cache_key(self): + """Graph topology key for softcap score_mod.""" + return ("softcap",) + + def forward(self, graph, score, tensors): + cudnn = tex_attention._import_cudnn_for_score_mod() + + self.before_tanh_activation = graph.div( + a=score, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + self.before_tanh_activation.set_data_type(cudnn.data_type.FLOAT) + tanh_out = graph.tanh(input=self.before_tanh_activation) + tanh_out.set_data_type(cudnn.data_type.FLOAT) + return graph.mul( + a=tanh_out, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + + def backward(self, graph, dscore, tensors): + cudnn = tex_attention._import_cudnn_for_score_mod() + + d_tanh_out = graph.mul( + a=dscore, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + d_tanh_out.set_data_type(cudnn.data_type.FLOAT) + d_before_tanh_activation = graph.tanh_backward( + loss=d_tanh_out, + input=self.before_tanh_activation, + compute_data_type=cudnn.data_type.FLOAT, + ) + d_before_tanh_activation.set_data_type(cudnn.data_type.FLOAT) + return graph.div( + a=d_before_tanh_activation, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + + +def _reference_score_mod_causal(scores): + q_pos = jnp.arange(scores.shape[-2])[:, None] + kv_pos = jnp.arange(scores.shape[-1])[None, :] + return jnp.where(q_pos >= kv_pos, scores, -1e9) + + +def _reference_score_mod_post_scale_bias(scores): + q_pos = jnp.arange(scores.shape[-2], dtype=jnp.float32)[:, None] + kv_pos = jnp.arange(scores.shape[-1], dtype=jnp.float32)[None, :] + return scores + q_pos - kv_pos + + +def _reference_score_mod_softcap(softcap): + def score_mod(scores): + return softcap * jnp.tanh(scores / softcap) + + return score_mod + + +class ScoreModFusedAttnRunner(FusedAttnRunner): + """FusedAttnRunner configured for score_mod tests.""" + + ATTN_BIAS_TYPE = AttnBiasType.NO_BIAS + ATTN_MASK_TYPE = AttnMaskType.NO_MASK + SOFTMAX_TYPE = AttnSoftmaxType.VANILLA_SOFTMAX + DROPOUT_PROBABILITY = 0.0 + IS_TRAINING = True + QKV_LAYOUT = QKVLayout.BSHD_BSHD_BSHD + BIAS_SHAPE = None + WINDOW_SIZE = None + SEQ_DESC_FORMAT = SeqDescFormat.Mask + CP_LOAD_BALANCED = False + DOUTPUT_SEED = None + RTOL = 5e-2 + ATOL = 5e-2 + + @staticmethod + def require_cudnn_frontend(): + """Skip unless cuDNN Python frontend supports score_mod SDPA.""" + try: + cudnn = tex_attention._import_cudnn_for_score_mod() + except ImportError: + pytest.skip("cuDNN Python frontend is required for score_mod") + version = tuple(int(part) for part in cudnn.backend_version_string().split(".")[:2]) + if version < (9, 6): + pytest.skip("cuDNN score_mod SDPA requires cuDNN frontend 9.6 or newer") + + @staticmethod + def _compute_input_scale(head_dim): + return 1.0 / sqrt(head_dim) + + @classmethod + def generic( + cls, + batch, + seqlen, + num_heads, + head_dim, + dtype, + *, + score_mod, + score_mod_reference, + score_mod_bprop=None, + score_mod_tensors=None, + score_mod_bprop_tensors=None, + doutput_seed=None, + rtol=None, + atol=None, + number_of_devices=1, + mesh_shape=(1, 1, 1), + mesh_axes=("dp", "cp", "tp"), + mesh_resource=None, + ): + """Build a runner for a separate-Q/K/V score_mod fused-attention case.""" + kwargs = {} + if mesh_resource is not None: + kwargs["mesh_resource"] = mesh_resource + if doutput_seed is None: + doutput_seed = cls.DOUTPUT_SEED + if rtol is None: + rtol = cls.RTOL + if atol is None: + atol = cls.ATOL + return cls( + batch, + seqlen, + seqlen, + num_heads, + num_heads, + head_dim, + head_dim, + cls.ATTN_BIAS_TYPE, + cls.ATTN_MASK_TYPE, + cls.SOFTMAX_TYPE, + cls.DROPOUT_PROBABILITY, + dtype, + cls.IS_TRAINING, + cls.QKV_LAYOUT, + cls.BIAS_SHAPE, + cls.WINDOW_SIZE, + cls.SEQ_DESC_FORMAT, + number_of_devices=number_of_devices, + mesh_shape=mesh_shape, + mesh_axes=mesh_axes, + cp_load_balanced=cls.CP_LOAD_BALANCED, + score_mod=score_mod, + score_mod_bprop=score_mod_bprop, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, + score_mod_reference=score_mod_reference, + input_scale=cls._compute_input_scale(head_dim), + doutput_seed=doutput_seed, + rtol=rtol, + atol=atol, + **kwargs, + ) + + @classmethod + def softcap( + cls, + batch, + seqlen, + num_heads, + head_dim, + dtype, + *, + number_of_devices=1, + mesh_shape=(1, 1, 1), + mesh_axes=("dp", "cp", "tp"), + mesh_resource=None, + ): + """Build a runner for softcap score modification with explicit bprop.""" + softcap = 0.8 + softcap_score_mod = _ScoreModSoftcap() + return cls.generic( + batch, + seqlen, + num_heads, + head_dim, + dtype, + number_of_devices=number_of_devices, + mesh_shape=mesh_shape, + mesh_axes=mesh_axes, + mesh_resource=mesh_resource, + score_mod=softcap_score_mod.forward, + score_mod_bprop=softcap_score_mod.backward, + score_mod_tensors={"softcap": softcap}, + score_mod_bprop_tensors={"softcap": softcap}, + score_mod_reference=_reference_score_mod_softcap(softcap), + doutput_seed=2025, + rtol=7e-2, + atol=7e-2, + ) + + @classmethod + def post_scale_bias( + cls, + batch, + seqlen, + num_heads, + head_dim, + dtype, + *, + number_of_devices=1, + mesh_shape=(1, 1, 1), + mesh_axes=("dp", "cp", "tp"), + mesh_resource=None, + ): + """Build a runner for post-scale-bias score modification without explicit bprop.""" + return cls.generic( + batch, + seqlen, + num_heads, + head_dim, + dtype, + number_of_devices=number_of_devices, + mesh_shape=mesh_shape, + mesh_axes=mesh_axes, + mesh_resource=mesh_resource, + score_mod=_score_mod_post_scale_bias, + score_mod_reference=_reference_score_mod_post_scale_bias, + ) + + @classmethod + def causal( + cls, + batch, + seqlen, + num_heads, + head_dim, + dtype, + *, + number_of_devices=1, + mesh_shape=(1, 1, 1), + mesh_axes=("dp", "cp", "tp"), + mesh_resource=None, + ): + """Build a runner for causal score modification with explicit bprop.""" + return cls.generic( + batch, + seqlen, + num_heads, + head_dim, + dtype, + number_of_devices=number_of_devices, + mesh_shape=mesh_shape, + mesh_axes=mesh_axes, + mesh_resource=mesh_resource, + score_mod=_score_mod_causal, + score_mod_bprop=_score_mod_causal_bprop, + score_mod_tensors={"neg_inf": -1e9}, + score_mod_bprop_tensors={"zero": 0.0}, + score_mod_reference=_reference_score_mod_causal, + ) + + +@pytest.fixture(autouse=True, scope="module") +def init(): + """ + WAR for CUDA uninitialize error + """ + # Calling customcalls before jax may cause CUDA uninitialize error + _ = jnp.zeros(0) + yield + + +def _identity_score_mod(_graph, score, _tensors): + return score + + +def _install_fake_flax_fused_attn(monkeypatch, *, kernel_available=True): + captured = {} + + def fake_fused_attn_kernel_check(*args, **kwargs): + captured.setdefault("kernel_checks", []).append((args, kwargs)) + return kernel_available + + def fake_fused_attn( + qkv, + bias, + sequence_descriptor, + seed, + *, + attn_bias_type, + attn_mask_type, + qkv_layout, + softmax_type, + scaling_factor, + dropout_probability, + is_training, + max_segments_per_seq=1, + window_size=None, + context_parallel_strategy=None, + context_parallel_causal_load_balanced=False, + context_parallel_axis="", + context_checkpoint_name="context", + softmax_offset=None, + stripe_size=None, + score_mod=None, + score_mod_bprop=None, + score_mod_tensors=None, + score_mod_bprop_tensors=None, + ): + captured.update( + qkv=qkv, + bias=bias, + sequence_descriptor=sequence_descriptor, + seed=seed, + attn_bias_type=attn_bias_type, + attn_mask_type=attn_mask_type, + qkv_layout=qkv_layout, + softmax_type=softmax_type, + scaling_factor=scaling_factor, + dropout_probability=dropout_probability, + is_training=is_training, + max_segments_per_seq=max_segments_per_seq, + window_size=window_size, + context_parallel_strategy=context_parallel_strategy, + context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, + context_parallel_axis=context_parallel_axis, + context_checkpoint_name=context_checkpoint_name, + softmax_offset=softmax_offset, + stripe_size=stripe_size, + score_mod=score_mod, + score_mod_bprop=score_mod_bprop, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, + ) + return qkv[0] + + monkeypatch.setattr( + flax_transformer, + "is_fused_attn_kernel_available", + fake_fused_attn_kernel_check, + ) + monkeypatch.setattr(flax_transformer, "fused_attn", fake_fused_attn) + return captured + + +def test_dot_product_attention_score_mod_requires_fused_attn(monkeypatch): + """DotProductAttention rejects score_mod when fused attention is disabled.""" + monkeypatch.setenv("NVTE_FUSED_ATTN", "0") + query = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + key = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + value = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + + dpa = flax_transformer.DotProductAttention( + head_dim=16, + num_attention_heads=1, + num_gqa_groups=1, + attn_mask_type="no_mask", + qkv_layout="bshd_bshd_bshd", + transpose_batch_sequence=False, + score_mod=_identity_score_mod, + ) + + with pytest.raises(ValueError, match="score_mod requires fused attention"): + dpa.apply({}, query, key, value, deterministic=True) + + +def test_dot_product_attention_score_mod_requires_available_fused_attn_kernel(monkeypatch): + """DotProductAttention rejects score_mod when no fused attention kernel is available.""" + _install_fake_flax_fused_attn(monkeypatch, kernel_available=False) + query = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + key = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + value = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + + dpa = flax_transformer.DotProductAttention( + head_dim=16, + num_attention_heads=1, + num_gqa_groups=1, + attn_mask_type="no_mask", + qkv_layout="bshd_bshd_bshd", + transpose_batch_sequence=False, + score_mod=_identity_score_mod, + ) + + with pytest.raises(ValueError, match="no fused attention kernel is available"): + dpa.apply({}, query, key, value, deterministic=True) + + +def test_dot_product_attention_plumbs_score_mod_to_fused_attn(monkeypatch): + """DotProductAttention forwards score_mod operands to fused_attn.""" + captured = _install_fake_flax_fused_attn(monkeypatch) + query = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + key = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + value = jnp.ones((1, 8, 1, 16), dtype=jnp.float16) + aux = jnp.ones((1, 1, 1, 1), dtype=jnp.float32) + + dpa = flax_transformer.DotProductAttention( + head_dim=16, + num_attention_heads=1, + num_gqa_groups=1, + attn_mask_type="no_mask", + qkv_layout="bshd_bshd_bshd", + transpose_batch_sequence=False, + score_mod=_identity_score_mod, + ) + out = dpa.apply({}, query, key, value, deterministic=True, score_mod_tensors={"aux": aux}) + + np.testing.assert_array_equal(out, query) + assert captured["score_mod"] is _identity_score_mod + assert captured["score_mod_tensors"]["aux"].shape == aux.shape + assert captured["score_mod_bprop"] is None + assert captured["score_mod_bprop_tensors"] is None + assert captured["attn_mask_type"] is AttnMaskType.NO_MASK + assert captured["attn_bias_type"] is AttnBiasType.NO_BIAS + assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD + assert captured["softmax_type"] is AttnSoftmaxType.VANILLA_SOFTMAX + assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD + + +def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monkeypatch): + """Packed QKV inputs are unpacked because score_mod requires separate Q/K/V.""" + captured = _install_fake_flax_fused_attn(monkeypatch) + qkv = jnp.ones((1, 8, 3, 1, 16), dtype=jnp.float16) + + dpa = flax_transformer.DotProductAttention( + head_dim=16, + num_attention_heads=1, + num_gqa_groups=1, + attn_mask_type="no_mask", + qkv_layout="bs3hd", + transpose_batch_sequence=False, + score_mod=_identity_score_mod, + ) + out = dpa.apply({}, qkv, None, None, deterministic=True) + + assert out.shape == (1, 8, 1, 16) + assert len(captured["qkv"]) == 3 + assert captured["qkv"][0].shape == (1, 8, 1, 16) + assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD + assert captured["score_mod"] is _identity_score_mod + assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD + + +def test_multi_head_attention_plumbs_score_mod_to_dot_product_attention(monkeypatch): + """MultiHeadAttention passes score_mod tensors through its attention stack.""" + captured = _install_fake_flax_fused_attn(monkeypatch) + + class FakeLayerNormDenseGeneral: + def __init__(self, *, features, return_layernorm_output=False, **_kwargs): + self.features = features + self.return_layernorm_output = return_layernorm_output + + def __call__(self, inputs): + features = self.features if isinstance(self.features, tuple) else (self.features,) + output = jnp.ones((*inputs.shape[:-1], *features), dtype=inputs.dtype) + ln_out = inputs if self.return_layernorm_output else None + return output, ln_out + + class FakeDenseGeneral: + def __init__(self, *, features, **_kwargs): + self.features = features + + def __call__(self, inputs): + features = self.features if isinstance(self.features, tuple) else (self.features,) + return jnp.ones((*inputs.shape[:-1], *features), dtype=inputs.dtype) + + monkeypatch.setattr(flax_transformer, "LayerNormDenseGeneral", FakeLayerNormDenseGeneral) + monkeypatch.setattr(flax_transformer, "DenseGeneral", FakeDenseGeneral) + + inputs = jnp.ones((1, 8, 16), dtype=jnp.float16) + aux = jnp.ones((1, 1, 1, 1), dtype=jnp.float32) + + mha = flax_transformer.MultiHeadAttention( + head_dim=16, + num_attention_heads=1, + num_gqa_groups=1, + input_layernorm=False, + attention_dropout=0.0, + attn_mask_type="no_mask", + fuse_qkv_params=True, + transpose_batch_sequence=False, + score_mod=_identity_score_mod, + ) + + variables = mha.init( + jax.random.key(0), + inputs, + inputs, + deterministic=True, + score_mod_tensors={"aux": aux}, + ) + out, ln_out = mha.apply( + variables, + inputs, + inputs, + deterministic=True, + score_mod_tensors={"aux": aux}, + ) + + assert out.shape == inputs.shape + assert ln_out is None + assert len(captured["qkv"]) == 3 + assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD + assert captured["score_mod"] is _identity_score_mod + assert captured["score_mod_tensors"]["aux"].shape == aux.shape + + +def test_fused_attn_score_mod_config_splits_tensors_and_pass_by_value_scalars(): + """Score-mod config separates array operands from pass-by-value scalars.""" + tensor = jnp.ones((1, 1, 1, 1), dtype=jnp.float32) + + config, tensor_operands, bprop_tensor_operands = make_fused_attn_score_mod_config( + _identity_score_mod, + None, + {"tensor": tensor, "neg_inf": -1e9}, + None, + _CONFIG_TEST_SCALING_FACTOR, + True, + ) + + assert config.score_mod_tensor_names == ("tensor",) + assert len(tensor_operands) == 1 + assert tensor_operands[0].shape == (1, 1, 1, 1) + assert len(bprop_tensor_operands) == 0 + assert len(config.score_mod_scalars) == 1 + assert config.score_mod_scalars[0].name == "neg_inf" + assert config.score_mod_scalars[0].dtype == "float32" + assert len(config.score_mod_scalars[0].value) == np.dtype(np.float32).itemsize + + +def test_fused_attn_score_mod_cudnn_frontend_version_check(monkeypatch): + """cuDNN frontend Python and C++ versions must match.""" + + class FakeCudnn: + __version__ = "1.22.0" + + monkeypatch.setattr( + tex_attention.transformer_engine_jax, + "get_cudnn_frontend_version", + lambda: 12200, + ) + assert tex_attention._check_cudnn_frontend_version_match(FakeCudnn) == 12200 + + monkeypatch.setattr( + tex_attention.transformer_engine_jax, + "get_cudnn_frontend_version", + lambda: 12100, + ) + with pytest.raises(RuntimeError, match="Python/C\\+\\+ version mismatch"): + tex_attention._check_cudnn_frontend_version_match(FakeCudnn) + + +def test_fused_attn_score_mod_config_stabilizes_bound_method_cache_keys(): + """Bound score_mod methods with graph keys share stable cache keys.""" + softcap_score_mod = _ScoreModSoftcap() + first_forward = softcap_score_mod.forward + second_forward = softcap_score_mod.forward + first_backward = softcap_score_mod.backward + second_backward = softcap_score_mod.backward + + assert first_forward is not second_forward + assert first_backward is not second_backward + + config_1, _, _ = make_fused_attn_score_mod_config( + first_forward, + first_backward, + {"softcap": 0.8}, + {"softcap": 0.8}, + _CONFIG_TEST_SCALING_FACTOR, + True, + ) + config_2, _, _ = make_fused_attn_score_mod_config( + second_forward, + second_backward, + {"softcap": 0.8}, + {"softcap": 0.8}, + _CONFIG_TEST_SCALING_FACTOR, + True, + ) + other_softcap_score_mod = _ScoreModSoftcap() + config_3, _, _ = make_fused_attn_score_mod_config( + other_softcap_score_mod.forward, + other_softcap_score_mod.backward, + {"softcap": 0.8}, + {"softcap": 0.8}, + _CONFIG_TEST_SCALING_FACTOR, + True, + ) + + assert config_1 == config_2 + assert hash(config_1) == hash(config_2) + assert config_1 == config_3 + + +def test_fused_attn_score_mod_config_leaves_unkeyed_bound_methods_uncached(): + """Unkeyed bound methods stay uncached to avoid object-id reuse collisions.""" + + class UnkeyedScoreMod: + def forward(self, _graph, score, _tensors): + return score + + score_mod = UnkeyedScoreMod() + config_1, _, _ = make_fused_attn_score_mod_config( + score_mod.forward, None, None, None, _CONFIG_TEST_SCALING_FACTOR, True + ) + config_2, _, _ = make_fused_attn_score_mod_config( + score_mod.forward, None, None, None, _CONFIG_TEST_SCALING_FACTOR, True + ) + + assert config_1 != config_2 + assert tex_attention._graph_cache_key("fwd", config_1, ()) is None + + +@pytest.mark.skipif(not _has_cudnn_frontend_python(), reason="cuDNN Python frontend is required") +def test_fused_attn_score_mod_post_scale_bias_optional_bprop(): + """Post-scale-bias score_mod matches the JAX reference without explicit bprop.""" + ScoreModFusedAttnRunner.require_cudnn_frontend() + runner = ScoreModFusedAttnRunner.post_scale_bias(1, 64, 2, 128, jnp.float16) + runner.test_forward() + runner.test_backward() + + +@pytest.mark.skipif(not _has_cudnn_frontend_python(), reason="cuDNN Python frontend is required") +def test_fused_attn_score_mod_causal_with_bprop(): + """Causal score_mod matches the JAX reference with explicit bprop.""" + ScoreModFusedAttnRunner.require_cudnn_frontend() + runner = ScoreModFusedAttnRunner.causal(1, 64, 2, 128, jnp.float16) + runner.test_forward() + runner.test_backward() + + +@pytest.mark.skipif(not _has_cudnn_frontend_python(), reason="cuDNN Python frontend is required") +@pytest.mark.skipif( + get_device_compute_capability(0) < 90, + reason="Softcap score_mod tests require sm90+", +) +def test_fused_attn_score_mod_softcap_with_bprop(): + """Softcap score_mod matches the JAX reference with explicit bprop.""" + ScoreModFusedAttnRunner.require_cudnn_frontend() + runner = ScoreModFusedAttnRunner.softcap(1, 16, 2, 64, jnp.float16) + runner.test_backward() diff --git a/tests/jax/test_fused_router.py b/tests/jax/test_fused_router.py index 89a32f1ce26..ff62bce7724 100644 --- a/tests/jax/test_fused_router.py +++ b/tests/jax/test_fused_router.py @@ -559,3 +559,139 @@ def loss_fused_fn(probs_): assert jnp.allclose( grad_ref, grad_fused, atol=1e-5, rtol=1e-5 ), f"Grad mismatch: max diff = {jnp.abs(grad_ref - grad_fused).max()}" + + +# ============================================================================= +# Test: routing_map BITMAP_U8 vs BYTEMAP parity (fwd + bwd) +# ============================================================================= + + +def _bytemap_to_bitmap_u8(bytemap): + """Reference packer: bool[T, E] -> uint8[T, ceil(E/8)] LSB-first.""" + import numpy as np + + flat = np.asarray(bytemap).astype(np.uint8) + return np.packbits(flat, axis=-1, bitorder="little") + + +@pytest_parametrize_wrapper("dtype", DTYPES) +@pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + TOPK_CASES, +) +@pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS) +@pytest.mark.triton +def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function): + """fused_topk_with_score_function should produce the same probs and an + LSB-packed bitmap routing_map when routing_map_format=BITMAP_U8, and + backward gradients should match the bytemap path exactly.""" + from transformer_engine.jax.router import RoutingMapFormat + + logits = make_logits(num_tokens, num_experts, score_function, dtype) + expert_bias = None + + fwd_byte = jax.jit( + partial( + fused_topk_with_score_function, + topk=topk, + score_function=score_function, + expert_bias=expert_bias, + routing_map_format=RoutingMapFormat.BYTEMAP, + ) + ) + fwd_bit = jax.jit( + partial( + fused_topk_with_score_function, + topk=topk, + score_function=score_function, + expert_bias=expert_bias, + routing_map_format=RoutingMapFormat.BITMAP_U8, + ) + ) + probs_byte, routing_map_byte = fwd_byte(logits) + probs_bit, routing_map_bit = fwd_bit(logits) + + assert probs_byte.dtype == probs_bit.dtype + assert jnp.array_equal(probs_byte, probs_bit), "Probs must be identical across formats" + + packed_expected = _bytemap_to_bitmap_u8(routing_map_byte) + assert routing_map_bit.shape == ( + num_tokens, + (num_experts + 7) // 8, + ), f"Bitmap shape {routing_map_bit.shape} != ({num_tokens}, {(num_experts + 7) // 8})" + assert routing_map_bit.dtype == jnp.uint8 + assert jnp.array_equal( + routing_map_bit, packed_expected + ), "Bitmap routing_map disagrees with np.packbits(bytemap, bitorder='little')" + + # Backward parity: grad of probs.sum() must be identical for both formats. + def loss_byte(logits_): + p, _ = fused_topk_with_score_function( + logits_, + topk, + score_function=score_function, + expert_bias=expert_bias, + routing_map_format=RoutingMapFormat.BYTEMAP, + ) + return p.sum() + + def loss_bit(logits_): + p, _ = fused_topk_with_score_function( + logits_, + topk, + score_function=score_function, + expert_bias=expert_bias, + routing_map_format=RoutingMapFormat.BITMAP_U8, + ) + return p.sum() + + grad_byte = jax.jit(jax.grad(loss_byte))(logits) + grad_bit = jax.jit(jax.grad(loss_bit))(logits) + assert jnp.allclose(grad_byte, grad_bit, atol=0.0, rtol=0.0), ( + "Backward grad must be bit-identical across routing_map_format; " + f"max diff = {jnp.abs(grad_byte - grad_bit).max()}" + ) + + +@pytest_parametrize_wrapper("dtype", DTYPES) +@pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + SCORE_AUX_LOSS_CASES, +) +@pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS) +@pytest.mark.triton +def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function): + """compute_aux_scores=True path: bitmap routing_map must equal LSB-packed + bytemap; scores must be bitwise identical across formats.""" + from transformer_engine.jax.router import RoutingMapFormat + + logits = make_logits(num_tokens, num_experts, score_function, dtype) + + fwd_byte = jax.jit( + partial( + fused_topk_with_score_function, + topk=topk, + score_function=score_function, + compute_aux_scores=True, + routing_map_format=RoutingMapFormat.BYTEMAP, + ) + ) + fwd_bit = jax.jit( + partial( + fused_topk_with_score_function, + topk=topk, + score_function=score_function, + compute_aux_scores=True, + routing_map_format=RoutingMapFormat.BITMAP_U8, + ) + ) + scores_byte, routing_map_byte = fwd_byte(logits) + scores_bit, routing_map_bit = fwd_bit(logits) + + assert jnp.array_equal(scores_byte, scores_bit), "Scores must be identical across formats" + packed_expected = _bytemap_to_bitmap_u8(routing_map_byte) + assert routing_map_bit.shape == (num_tokens, (num_experts + 7) // 8) + assert routing_map_bit.dtype == jnp.uint8 + assert jnp.array_equal( + routing_map_bit, packed_expected + ), "Bitmap routing_map (aux-loss path) disagrees with packed bytemap" diff --git a/tests/pytorch/attention/mla_rope_utils.py b/tests/pytorch/attention/mla_rope_utils.py new file mode 100644 index 00000000000..90eebfc66ae --- /dev/null +++ b/tests/pytorch/attention/mla_rope_utils.py @@ -0,0 +1,668 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MLA RoPE for DSv3 671B - Triton forward and backward kernels. + +Source: Megatron-LM megatron/core/fusions/fused_mla_yarn_rope_apply.py +Falls back to pure PyTorch when Triton is unavailable. + +Note: DSv3 uses YaRN-scaled RoPE for long-context extrapolation. This test +intentionally uses plain RoPE (base=10000) because it only validates MXFP8 +attention path wiring, tensor shapes, forward/backward flow, and relative BF16 +vs MXFP8 behavior. Both reference and MXFP8 paths use the same RoPE tables. +""" + +import torch + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + HAVE_TRITON = False + +HEAD_DIM_ROPE = 64 +HEAD_DIM_NOPE = 128 +HEAD_DIM_V = 128 +ROTARY_BASE = 10000 + + +def build_rope_tables( + seq_len: int, + emb_dim: int = HEAD_DIM_ROPE, + base: int = ROTARY_BASE, + device: torch.device = None, +) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / ( + base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) + ) + t = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + freqs = torch.cat([freqs, freqs], dim=-1) + return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() + + +if HAVE_TRITON: + + # Not used for non-packed batches; kept for THD compatibility. + @triton.jit + def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): + token_idx = -1 + this_seq_len = 0 + seq_idx = 0 + last_cum_seqlen = tl.load(cu_seqlens) // cp_size + while seq_idx < seq_num: + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size + if token_idx == -1 and cur_cum_seqlen > pid_m: + token_idx = pid_m - last_cum_seqlen + this_seq_len = cur_cum_seqlen - last_cum_seqlen + last_cum_seqlen = cur_cum_seqlen + seq_idx += 1 + if cp_size > 1: + if token_idx < this_seq_len // 2: + token_idx = token_idx + cp_rank * this_seq_len // 2 + else: + token_idx = (token_idx - this_seq_len // 2) + ( + 2 * cp_size - cp_rank - 1 + ) * this_seq_len // 2 + return token_idx + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_H": 1}), + triton.Config({"BLOCK_H": 2}), + triton.Config({"BLOCK_H": 4}), + triton.Config({"BLOCK_H": 8}), + triton.Config({"BLOCK_H": 16}), + triton.Config({"BLOCK_H": 32}), + triton.Config({"BLOCK_H": 64}), + triton.Config({"BLOCK_H": 128}), + ], + key=["emb_dim", "head_num"], + restore_value=["Q"], + ) + @triton.jit + def rotary_fwd_q_kernel( + Q, + COS, + SIN, + qk_head_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + stride_x_seq, + stride_x_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + Q = Q + pid_m * stride_x_seq + x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim + mask = head_offsets[:, None] < head_num + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_1 = tl.load(Q + x_1_off, mask=mask) + x_2 = tl.load(Q + x_2_off, mask=mask) + x_left = x_1 * cos_left - x_2 * sin_left + x_right = x_2 * cos_right + x_1 * sin_right + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_H": 1}), + triton.Config({"BLOCK_H": 2}), + triton.Config({"BLOCK_H": 4}), + triton.Config({"BLOCK_H": 8}), + triton.Config({"BLOCK_H": 16}), + triton.Config({"BLOCK_H": 32}), + triton.Config({"BLOCK_H": 64}), + triton.Config({"BLOCK_H": 128}), + ], + key=["emb_dim", "head_num"], + restore_value=["DO"], + ) + @triton.jit + def rotary_bwd_q_kernel( + DO, + COS, + SIN, + qk_head_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + stride_x_seq, + stride_x_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + DO = DO + pid_m * stride_x_seq + x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim + mask = head_offsets[:, None] < head_num + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1 = x_left * cos_left + x_right * sin_right + x_2 = -x_left * sin_left + x_right * cos_right + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + tl.store(DO + x_1_off, x_1, mask=mask) + tl.store(DO + x_2_off, x_2, mask=mask) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_H": 1}), + triton.Config({"BLOCK_H": 2}), + triton.Config({"BLOCK_H": 4}), + triton.Config({"BLOCK_H": 8}), + triton.Config({"BLOCK_H": 16}), + triton.Config({"BLOCK_H": 32}), + triton.Config({"BLOCK_H": 64}), + triton.Config({"BLOCK_H": 128}), + ], + key=["emb_dim", "k_dim", "v_dim", "head_num"], + ) + @triton.jit + def rotary_fwd_kv_kernel( + KV, + K_POS_EMB, + O_KEY, + O_VALUE, + COS, + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_kv, + stride_kv_seq, + stride_kv_nheads, + stride_emb_seq, + stride_k_seq, + stride_k_nheads, + stride_v_seq, + stride_v_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_kv is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + KV_ptr = KV + pid_m * stride_kv_seq + kv_off = head_offsets[:, None] * stride_kv_nheads + mask = head_offsets[:, None] < head_num + k_in_off = kv_off + tl.arange(0, k_dim)[None, :] + v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] + k = tl.load(KV_ptr + k_in_off, mask=mask) + v = tl.load(KV_ptr + v_in_off, mask=mask) + K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads + V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads + k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] + v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] + tl.store(K_ptr + k_out_off, k, mask=mask) + tl.store(V_ptr + v_out_off, v, mask=mask) + EMB = K_POS_EMB + pid_m * stride_emb_seq + x_1 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2) + x_2 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2 + 1) + x_left = x_1 * cos_left - x_2 * sin_left + x_right = x_2 * cos_right + x_1 * sin_right + x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_left_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 + tl.store(K_ptr + x_left_off, x_left, mask=mask) + tl.store(K_ptr + x_right_off, x_right, mask=mask) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_H": 1}), + triton.Config({"BLOCK_H": 2}), + triton.Config({"BLOCK_H": 4}), + triton.Config({"BLOCK_H": 8}), + triton.Config({"BLOCK_H": 16}), + triton.Config({"BLOCK_H": 32}), + triton.Config({"BLOCK_H": 64}), + triton.Config({"BLOCK_H": 128}), + ], + key=["emb_dim", "k_dim", "v_dim", "head_num"], + ) + @triton.jit + def rotary_bwd_kv_kernel( + dK, + dV, + dKV, + dEMB, + COS, + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_kv, + stride_dk_seq, + stride_dk_nheads, + stride_dv_seq, + stride_dv_nheads, + stride_dkv_seq, + stride_dkv_nheads, + stride_demb_seq, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_kv is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + dKV_ptr = dKV + pid_m * stride_dkv_seq + dkv_off = head_offsets[:, None] * stride_dkv_nheads + mask = head_offsets[:, None] < head_num + dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] + dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] + dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads + dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads + dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] + dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] + dk = tl.load(dK_ptr + dk_in_off, mask=mask) + dv = tl.load(dV_ptr + dv_in_off, mask=mask) + tl.store(dKV_ptr + dk_out_off, dk, mask=mask) + tl.store(dKV_ptr + dv_out_off, dv, mask=mask) + if pid_head == 0: + x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): + head_offsets_i = i * BLOCK_H + tl.arange(0, BLOCK_H) + dK_ptr_i = dK + pid_m * stride_dk_seq + x_off = head_offsets_i[:, None] * stride_dk_nheads + k_dim + mask_i = head_offsets_i[:, None] < head_num + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left_accum += tl.load(dK_ptr_i + x_left_off, mask=mask_i) + x_right_accum += tl.load(dK_ptr_i + x_right_off, mask=mask_i) + x_left_accum = tl.sum(x_left_accum, axis=0) + x_right_accum = tl.sum(x_right_accum, axis=0) + x_left_accum = x_left_accum.to(dEMB.dtype.element_ty) + x_right_accum = x_right_accum.to(dEMB.dtype.element_ty) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + x_1 = x_left_accum * cos_left + x_right_accum * sin_right + x_2 = -x_left_accum * sin_left + x_right_accum * cos_right + dEMB_ptr = dEMB + pid_m * stride_demb_seq + tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2, x_1) + tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) + + def _flattened_token_stride(tensor: torch.Tensor) -> int: + if tensor.dim() == 4: + return tensor.stride(1) + return tensor.stride(0) + + class _MLARoPEQTriton(torch.autograd.Function): + @staticmethod + def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): + s, b, nheads, _ = q.shape + total = s * b + + grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_fwd_q_kernel[grid_q]( + q, + cos, + sin, + head_dim_nope, + head_dim_rope, + nheads, + b, + None, + None, + _flattened_token_stride(q), + q.stride(2), + 0, + 1, + ) + + ctx.save_for_backward(cos, sin) + ctx.head_dim_nope = head_dim_nope + ctx.head_dim_rope = head_dim_rope + ctx.nheads = nheads + ctx.s = s + ctx.b = b + return q + + @staticmethod + def backward(ctx, dq): + cos, sin = ctx.saved_tensors + s, b, nheads = ctx.s, ctx.b, ctx.nheads + total = s * b + + grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_bwd_q_kernel[grid_q]( + dq, + cos, + sin, + ctx.head_dim_nope, + ctx.head_dim_rope, + nheads, + b, + None, + None, + _flattened_token_stride(dq), + dq.stride(2), + 0, + 1, + ) + return dq, None, None, None, None + + class _MLARoPEKVTriton(torch.autograd.Function): + @staticmethod + def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): + s, b, nheads, _ = kv.shape + total = s * b + + o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) + o_value = kv.new_empty(s, b, nheads, head_dim_v) + grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_fwd_kv_kernel[grid_kv]( + kv, + k_pos_emb, + o_key, + o_value, + cos, + sin, + head_dim_rope, + head_dim_nope, + head_dim_v, + nheads, + b, + None, + None, + _flattened_token_stride(kv), + kv.stride(2), + _flattened_token_stride(k_pos_emb), + _flattened_token_stride(o_key), + o_key.stride(2), + _flattened_token_stride(o_value), + o_value.stride(2), + 0, + 1, + ) + + ctx.save_for_backward(cos, sin) + ctx.head_dim_nope = head_dim_nope + ctx.head_dim_rope = head_dim_rope + ctx.head_dim_v = head_dim_v + ctx.nheads = nheads + ctx.s = s + ctx.b = b + return o_key, o_value + + @staticmethod + def backward(ctx, dk_out, dv_out): + cos, sin = ctx.saved_tensors + s, b, nheads = ctx.s, ctx.b, ctx.nheads + ndp, ndr, ndv = ctx.head_dim_nope, ctx.head_dim_rope, ctx.head_dim_v + total = s * b + + d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) + d_emb = dk_out.new_empty(s, b, 1, ndr) + grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_bwd_kv_kernel[grid_kv]( + dk_out, + dv_out, + d_kv, + d_emb, + cos, + sin, + ndr, + ndp, + ndv, + nheads, + b, + None, + None, + _flattened_token_stride(dk_out), + dk_out.stride(2), + _flattened_token_stride(dv_out), + dv_out.stride(2), + _flattened_token_stride(d_kv), + d_kv.stride(2), + _flattened_token_stride(d_emb), + 0, + 1, + ) + return d_kv, d_emb, None, None, None, None, None + + +def _apply_mla_rope_q_with_tables( + q: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + head_dim_nope: int = HEAD_DIM_NOPE, + head_dim_rope: int = HEAD_DIM_ROPE, +) -> torch.Tensor: + if HAVE_TRITON: + return _MLARoPEQTriton.apply( + q, + cos_table, + sin_table, + head_dim_nope, + head_dim_rope, + ) + return _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + + +def _apply_mla_rope_kv_with_tables( + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + head_dim_nope: int = HEAD_DIM_NOPE, + head_dim_rope: int = HEAD_DIM_ROPE, + head_dim_v: int = HEAD_DIM_V, +) -> tuple[torch.Tensor, torch.Tensor]: + if HAVE_TRITON: + return _MLARoPEKVTriton.apply( + kv, + k_pos_emb, + cos_table, + sin_table, + head_dim_nope, + head_dim_rope, + head_dim_v, + ) + return _apply_pytorch_kv( + kv, + k_pos_emb, + cos_table, + sin_table, + head_dim_nope, + head_dim_rope, + head_dim_v, + ) + + +def apply_mla_rope_q( + q: torch.Tensor, + head_dim_nope: int = HEAD_DIM_NOPE, + head_dim_rope: int = HEAD_DIM_ROPE, + base: int = ROTARY_BASE, + cos_table: torch.Tensor | None = None, + sin_table: torch.Tensor | None = None, +) -> torch.Tensor: + if cos_table is None or sin_table is None: + s = q.shape[0] + cos_table, sin_table = build_rope_tables( + s, + emb_dim=head_dim_rope, + base=base, + device=q.device, + ) + return _apply_mla_rope_q_with_tables( + q, + cos_table, + sin_table, + head_dim_nope, + head_dim_rope, + ) + + +def apply_mla_rope_kv( + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + head_dim_nope: int = HEAD_DIM_NOPE, + head_dim_rope: int = HEAD_DIM_ROPE, + head_dim_v: int = HEAD_DIM_V, + base: int = ROTARY_BASE, + cos_table: torch.Tensor | None = None, + sin_table: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if cos_table is None or sin_table is None: + s = kv.shape[0] + cos_table, sin_table = build_rope_tables( + s, + emb_dim=head_dim_rope, + base=base, + device=kv.device, + ) + return _apply_mla_rope_kv_with_tables( + kv, + k_pos_emb, + cos_table, + sin_table, + head_dim_nope, + head_dim_rope, + head_dim_v, + ) + + +def apply_mla_rope( + q: torch.Tensor, + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + head_dim_nope: int = HEAD_DIM_NOPE, + head_dim_rope: int = HEAD_DIM_ROPE, + head_dim_v: int = HEAD_DIM_V, + base: int = ROTARY_BASE, + cos_table: torch.Tensor | None = None, + sin_table: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if cos_table is None or sin_table is None: + s = q.shape[0] + cos_table, sin_table = build_rope_tables( + s, + emb_dim=head_dim_rope, + base=base, + device=q.device, + ) + q = _apply_mla_rope_q_with_tables(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + k, v = _apply_mla_rope_kv_with_tables( + kv, + k_pos_emb, + cos_table, + sin_table, + head_dim_nope, + head_dim_rope, + head_dim_v, + ) + return q, k, v + + +def _rotate_interleaved_to_neox( + x: torch.Tensor, cos_table: torch.Tensor, sin_table: torch.Tensor +) -> torch.Tensor: + cos_ = cos_table[:, None, None, :].to(x.dtype) + sin_ = sin_table[:, None, None, :].to(x.dtype) + half_dim = x.shape[-1] // 2 + x_1 = x[..., 0::2] + x_2 = x[..., 1::2] + x_left = x_1 * cos_[..., :half_dim] - x_2 * sin_[..., :half_dim] + x_right = x_2 * cos_[..., half_dim:] + x_1 * sin_[..., half_dim:] + return torch.cat((x_left, x_right), dim=-1) + + +def _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope): + q_nope = q[..., :head_dim_nope] + q_rope = q[..., head_dim_nope : head_dim_nope + head_dim_rope] + q_rope = _rotate_interleaved_to_neox(q_rope, cos_table, sin_table) + return torch.cat((q_nope, q_rope), dim=-1) + + +def _apply_pytorch_kv( + kv, + k_pos_emb, + cos_table, + sin_table, + head_dim_nope, + head_dim_rope, + head_dim_v, +): + k_nope = kv[..., :head_dim_nope] + v = kv[..., head_dim_nope : head_dim_nope + head_dim_v] + k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table).expand( + -1, -1, kv.shape[2], -1 + ) + return torch.cat((k_nope, k_rope), dim=-1), v diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 6fca61d3c07..52c7dc067c1 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -14,6 +14,7 @@ ) from transformer_engine.pytorch.attention.dot_product_attention.utils import combine_and_quantize import transformer_engine_torch as tex +from transformer_engine.pytorch import DType from test_attention_with_cp import model_configs_flash_attn, model_configs_fused_attn from transformer_engine.pytorch import ( autocast, @@ -328,34 +329,34 @@ def run_dpa_with_cp( ).cuda() if scaling_mode == "delayed": qkv_quantizer = Float8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, scale=torch.tensor([1], dtype=torch.float32).cuda(), amax=torch.tensor([0], dtype=torch.float32).cuda(), ) dout_quantizer = Float8Quantizer( - fp8_dtype=tex.DType.kFloat8E5M2, + fp8_dtype=DType.kFloat8E5M2, scale=torch.tensor([1], dtype=torch.float32).cuda(), amax=torch.tensor([0], dtype=torch.float32).cuda(), ) if scaling_mode == "current": qkv_quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, device="cuda", ) dout_quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E5M2, + fp8_dtype=DType.kFloat8E5M2, device="cuda", ) if scaling_mode == "mxfp8": qkv_quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True, ) qkv_quantizer.optimize_for_gemm = True qkv_quantizer.internal = False dout_quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E5M2, + fp8_dtype=DType.kFloat8E5M2, rowwise=True, columnwise=True, ) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 401bd6f01d2..2dbf94fc205 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1926,14 +1926,14 @@ def get_model(dtype, config): # test: ModelConfig(b, sq, hq, dqk) "fp8_9": ModelConfig( 2, - 4096, + 2048, 128, 192, head_dim_v=128, ), "fp8_10": ModelConfig( - 1, - 4096, + 2, + 2048, 128, 192, head_dim_v=128, @@ -1941,21 +1941,23 @@ def get_model(dtype, config): ), "fp8_11": ModelConfig( 2, - 4096, + 2048, 128, 192, head_dim_v=128, attn_mask_type="causal_bottom_right", ), - "fp8_12": ModelConfig(2, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), - "fp8_13": ModelConfig(2, 8192, 32, 128, attn_mask_type="causal", window_size=(128, 0)), - "fp8_14": ModelConfig(2, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal"), - "fp8_15": ModelConfig(2, 8192, 64, 64, attn_mask_type="causal", window_size=(128, 0)), + "fp8_12": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), + "fp8_13": ModelConfig( + 2, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal", window_size=(128, 0) + ), + "fp8_14": ModelConfig(2, 4096, 64, 64, num_gqa_groups=8, attn_mask_type="causal"), + "fp8_15": ModelConfig(1, 8192, 64, 64, attn_mask_type="causal", window_size=(128, 0)), "fp8_16": ModelConfig( - 2, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" + 1, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" ), "fp8_17": ModelConfig( - 2, 8192, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" + 2, 4096, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" ), "fp8_18": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), "fp8_19": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py new file mode 100644 index 00000000000..55b8b3b9aa6 --- /dev/null +++ b/tests/pytorch/attention/test_flex_attention.py @@ -0,0 +1,708 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import os +import sys +import pathlib + +import pytest +import torch + +from transformer_engine.pytorch import DotProductAttention, is_bf16_available +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends +import transformer_engine.pytorch.attention.dot_product_attention.flex_attention as flex_attention +from transformer_engine.pytorch.utils import get_cudnn_version, get_device_compute_capability + +_current_file = pathlib.Path(__file__).resolve() +sys.path = [str(_current_file.parent.parent)] + sys.path +from utils import ( # pylint: disable=wrong-import-position + reset_rng_states, + ModelConfig, + get_available_attention_backends, +) + +param_types = [torch.float16] +if torch.cuda.is_available() and is_bf16_available(): + param_types.append(torch.bfloat16) + + +def _score_mod_causal(score_mod_graph, score_tensor, tensors): + """cuDNN frontend score_mod implementing top-left causal masking.""" + cudnn = flex_attention._import_cudnn_frontend() + + row_index = score_mod_graph.gen_index(input=score_tensor, axis=2) + row_index.set_data_type(cudnn.data_type.INT32) + col_index = score_mod_graph.gen_index(input=score_tensor, axis=3) + col_index.set_data_type(cudnn.data_type.INT32) + keep = score_mod_graph.cmp_ge( + input=row_index, + comparison=col_index, + compute_data_type=cudnn.data_type.BOOLEAN, + ) + keep.set_data_type(cudnn.data_type.BOOLEAN) + return score_mod_graph.binary_select( + input0=score_tensor, + input1=tensors["neg_inf"], + mask=keep, + ) + + +def _score_mod_causal_bprop(score_mod_graph, dP_tensor, tensors): + """cuDNN frontend score_mod_bprop implementing top-left causal masking.""" + cudnn = flex_attention._import_cudnn_frontend() + + row_index = score_mod_graph.gen_index(input=dP_tensor, axis=2) + row_index.set_data_type(cudnn.data_type.INT32) + col_index = score_mod_graph.gen_index(input=dP_tensor, axis=3) + col_index.set_data_type(cudnn.data_type.INT32) + keep = score_mod_graph.cmp_ge( + input=row_index, + comparison=col_index, + compute_data_type=cudnn.data_type.BOOLEAN, + ) + keep.set_data_type(cudnn.data_type.BOOLEAN) + return score_mod_graph.binary_select( + input0=dP_tensor, + input1=tensors["zero"], + mask=keep, + ) + + +def _score_mod_post_scale_bias(score_mod_graph, score_tensor, _tensors): + """cuDNN frontend score_mod adding post-scale bias.""" + cudnn = flex_attention._import_cudnn_frontend() + + row_index = score_mod_graph.gen_index(input=score_tensor, axis=2) + row_index.set_data_type(cudnn.data_type.INT32) + col_index = score_mod_graph.gen_index(input=score_tensor, axis=3) + col_index.set_data_type(cudnn.data_type.INT32) + post_scale_bias = score_mod_graph.sub( + a=row_index, + b=col_index, + compute_data_type=cudnn.data_type.FLOAT, + ) + post_scale_bias.set_data_type(cudnn.data_type.FLOAT) + return score_mod_graph.add( + a=score_tensor, + b=post_scale_bias, + compute_data_type=cudnn.data_type.FLOAT, + ) + + +def _score_mod_identity_bprop(_score_mod_graph, dP_tensor, _tensors): + """cuDNN frontend score_mod_bprop for score_mods with unit score derivative.""" + return dP_tensor + + +class _ScoreModSoftcap: + """cuDNN frontend score_mod implementing softcapping.""" + + def __init__(self): + self.before_tanh_activation = None + + def score_mod_graph_cache_key(self): + """Graph topology key for softcap score_mod.""" + return ("softcap",) + + def forward(self, score_mod_graph, score_tensor, tensors): + """Apply softcap * tanh(score / softcap).""" + cudnn = flex_attention._import_cudnn_frontend() + + self.before_tanh_activation = score_mod_graph.div( + a=score_tensor, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + self.before_tanh_activation.set_data_type(cudnn.data_type.FLOAT) + tanh_out = score_mod_graph.tanh(input=self.before_tanh_activation) + tanh_out.set_data_type(cudnn.data_type.FLOAT) + return score_mod_graph.mul( + a=tanh_out, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + + def backward(self, score_mod_graph, dP_tensor, tensors): + """Apply softcap derivative to dP.""" + cudnn = flex_attention._import_cudnn_frontend() + + d_tanh_out = score_mod_graph.mul( + a=dP_tensor, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + d_tanh_out.set_data_type(cudnn.data_type.FLOAT) + d_before_tanh_activation = score_mod_graph.tanh_backward( + loss=d_tanh_out, + input=self.before_tanh_activation, + compute_data_type=cudnn.data_type.FLOAT, + ) + d_before_tanh_activation.set_data_type(cudnn.data_type.FLOAT) + return score_mod_graph.div( + a=d_before_tanh_activation, + b=tensors["softcap"], + compute_data_type=cudnn.data_type.FLOAT, + ) + + +def _score_mod_cache_cpu_inputs(): + """Small CPU tensors for score_mod cache-key tests.""" + q = torch.empty((2, 4, 3, 8), dtype=torch.float16) + k = torch.empty((2, 4, 3, 8), dtype=torch.float16) + v = torch.empty((2, 4, 3, 8), dtype=torch.float16) + o = torch.empty((2, 4, 3, 8), dtype=torch.float16) + stats = torch.empty((2, 3, 4, 1), dtype=torch.float32) + return q, k, v, o, stats + + +def test_score_mod_cache_bound_method_requires_explicit_key(): + """Unkeyed bound methods should be uncached instead of keyed by object id.""" + + class UnkeyedScoreMod: + def forward(self, _score_mod_graph, score_tensor, _tensors): + return score_tensor + + key = flex_attention._score_mod_callback_cache_key(UnkeyedScoreMod().forward) + + assert key is flex_attention._SCORE_MOD_UNCACHEABLE + + +def test_score_mod_cache_bound_method_explicit_key_stable(): + """Bound method keys should be stable when a structural graph key is provided.""" + softcap = _ScoreModSoftcap() + key_0 = flex_attention._score_mod_callback_cache_key(softcap.forward) + key_1 = flex_attention._score_mod_callback_cache_key(softcap.forward) + other_key = flex_attention._score_mod_callback_cache_key(_ScoreModSoftcap().forward) + + assert key_0 == key_1 + assert key_0 == other_key + + +def test_score_mod_cache_explicit_key_distinguishes_topology(): + """Stateful score_mods can opt into caching with topology-specific keys.""" + + class LayeredScoreMod: + def __init__(self, num_layers): + self.num_layers = num_layers + + def score_mod_graph_cache_key(self): + return {"num_layers": self.num_layers} + + def forward(self, _score_mod_graph, score_tensor, _tensors): + return score_tensor + + key_0 = flex_attention._score_mod_callback_cache_key(LayeredScoreMod(1).forward) + key_1 = flex_attention._score_mod_callback_cache_key(LayeredScoreMod(1).forward) + key_2 = flex_attention._score_mod_callback_cache_key(LayeredScoreMod(2).forward) + + assert key_0 == key_1 + assert key_0 != key_2 + + +def test_score_mod_cache_module_lambda_keys_do_not_collide(): + """Module-level lambdas should not reuse graphs only because qualnames match.""" + score_mod_0 = lambda _graph, score_tensor, _tensors: score_tensor # noqa: E731 + score_mod_1 = lambda _graph, score_tensor, _tensors: score_tensor # noqa: E731 + score_mod_0.__module__ = __name__ + score_mod_1.__module__ = __name__ + score_mod_0.__qualname__ = "" + score_mod_1.__qualname__ = "" + + key_0 = flex_attention._score_mod_callback_cache_key(score_mod_0) + key_1 = flex_attention._score_mod_callback_cache_key(score_mod_1) + + assert key_0 is not flex_attention._SCORE_MOD_UNCACHEABLE + assert key_1 is not flex_attention._SCORE_MOD_UNCACHEABLE + assert key_0 != key_1 + + +def test_score_mod_cache_key_ignores_pass_by_value_values(): + """Scalar CPU tensor values are runtime inputs, not execution-plan metadata.""" + q, k, v, o, stats = _score_mod_cache_cpu_inputs() + key_0 = flex_attention._cudnn_score_mod_fwd_cache_key( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + _score_mod_causal, + {"softcap": torch.tensor(0.8, dtype=torch.float32)}, + o, + stats, + ) + key_1 = flex_attention._cudnn_score_mod_fwd_cache_key( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + _score_mod_causal, + {"softcap": torch.tensor(1.2, dtype=torch.float32)}, + o, + stats, + ) + key_2 = flex_attention._cudnn_score_mod_fwd_cache_key( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + _score_mod_causal, + {"softcap": torch.tensor([0.8], dtype=torch.float32)}, + o, + stats, + ) + + assert key_0 == key_1 + assert key_0 != key_2 + + +def test_score_mod_cache_fwd_reuses_graph_for_pass_by_value_changes(monkeypatch): + """Fprop graph cache should reuse entries when only scalar CPU tensor values change.""" + q, k, v, o, stats = _score_mod_cache_cpu_inputs() + cache = flex_attention._cudnn_score_mod_graph_cache + saved_cache = dict(cache) + build_entries = [] + + def fake_build( + is_training, + query_layer, + key_layer, + value_layer, + q_format, + kv_format, + attn_scale, + score_mod, + score_mod_tensors, + output_layer, + stats, + ): + del ( + is_training, + query_layer, + key_layer, + value_layer, + q_format, + kv_format, + attn_scale, + score_mod, + score_mod_tensors, + output_layer, + stats, + ) + entry = object() + build_entries.append(entry) + return entry + + monkeypatch.setattr(flex_attention, "_build_cudnn_score_mod_fwd_graph", fake_build) + try: + cache.clear() + entry_0 = flex_attention._get_cudnn_score_mod_fwd_graph( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + _score_mod_causal, + {"softcap": torch.tensor(0.8, dtype=torch.float32)}, + o, + stats, + ) + entry_1 = flex_attention._get_cudnn_score_mod_fwd_graph( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + _score_mod_causal, + {"softcap": torch.tensor(1.2, dtype=torch.float32)}, + o, + stats, + ) + entry_2 = flex_attention._get_cudnn_score_mod_fwd_graph( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + _score_mod_causal, + {"softcap": torch.tensor([0.8], dtype=torch.float32)}, + o, + stats, + ) + finally: + cache.clear() + cache.update(saved_cache) + + assert entry_0 is entry_1 + assert entry_2 is not entry_0 + assert len(build_entries) == 2 + + +def test_score_mod_cache_fwd_skips_cache_for_unkeyed_bound_method(monkeypatch): + """Unkeyed bound methods should build fresh graphs instead of using an id-based key.""" + + class UnkeyedScoreMod: + def forward(self, _score_mod_graph, score_tensor, _tensors): + return score_tensor + + q, k, v, o, stats = _score_mod_cache_cpu_inputs() + score_mod = UnkeyedScoreMod() + cache = flex_attention._cudnn_score_mod_graph_cache + saved_cache = dict(cache) + build_entries = [] + + def fake_build( + is_training, + query_layer, + key_layer, + value_layer, + q_format, + kv_format, + attn_scale, + score_mod, + score_mod_tensors, + output_layer, + stats, + ): + del ( + is_training, + query_layer, + key_layer, + value_layer, + q_format, + kv_format, + attn_scale, + score_mod, + score_mod_tensors, + output_layer, + stats, + ) + entry = object() + build_entries.append(entry) + return entry + + monkeypatch.setattr(flex_attention, "_build_cudnn_score_mod_fwd_graph", fake_build) + try: + cache.clear() + entry_0 = flex_attention._get_cudnn_score_mod_fwd_graph( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + score_mod.forward, + None, + o, + stats, + ) + entry_1 = flex_attention._get_cudnn_score_mod_fwd_graph( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + score_mod.forward, + None, + o, + stats, + ) + assert len(cache) == 0 + finally: + cache.clear() + cache.update(saved_cache) + + assert entry_0 is not entry_1 + assert len(build_entries) == 2 + + +def test_score_mod_tensors_are_version_checked_for_backward(monkeypatch): + """In-place score_mod tensor updates before backward should be rejected.""" + + class FakeEntry: + graph = object() + q = object() + k = object() + v = object() + output = object() + stats = object() + score_mod_graph_tensors = {"softcap": object()} + workspace_size = 1 + + def fake_execute(graph, variant_pack, workspace_size, device): + del graph, variant_pack, workspace_size, device + + q, k, v, _, _ = _score_mod_cache_cpu_inputs() + q = q.requires_grad_() + k = k.requires_grad_() + v = v.requires_grad_() + softcap = torch.tensor(0.8, dtype=torch.float32) + + monkeypatch.setattr(flex_attention, "_get_cudnn_score_mod_fwd_graph", lambda *args: FakeEntry()) + monkeypatch.setattr(flex_attention, "_execute_cudnn_graph", fake_execute) + + out = flex_attention.FusedAttentionWithScoreModFunc.apply( + True, + q, + k, + v, + "bshd", + "bshd", + 1.0, + _score_mod_causal, + None, + {"softcap": softcap}, + None, + False, + ) + softcap.add_(1.0) + + with pytest.raises(RuntimeError, match="modified by an inplace operation"): + out.sum().backward() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") +def test_score_mod_bprop_tensors_require_score_mod_bprop(): + """score_mod_bprop_tensors should not be silently ignored.""" + + q, k, v, _, _ = (tensor.cuda() for tensor in _score_mod_cache_cpu_inputs()) + attention = DotProductAttention( + num_attention_heads=3, + kv_channels=8, + qkv_format="bshd", + attn_mask_type="no_mask", + ).cuda() + + with pytest.raises(AssertionError, match="score_mod_bprop_tensors requires score_mod_bprop"): + attention( + q, + k, + v, + qkv_format="bshd", + attn_mask_type="no_mask", + score_mod=_score_mod_causal, + score_mod_bprop_tensors={"zero": torch.zeros((1, 1, 1, 1), device="cuda")}, + ) + + +def _post_scale_bias(config, dtype): + """Materialize score + (q_idx - kv_idx) as post-scale attention bias.""" + q_idx = torch.arange(config.max_seqlen_q, dtype=torch.float32, device="cuda").view(1, 1, -1, 1) + kv_idx = torch.arange(config.max_seqlen_kv, dtype=torch.float32, device="cuda").view( + 1, 1, 1, -1 + ) + return (q_idx - kv_idx).to(dtype).expand(1, config.num_heads, -1, -1).contiguous() + + +def _to_bhsd(tensor, qkv_format): + """Convert SBHD/BSHD test tensors to logical BHSD.""" + if qkv_format == "sbhd": + return tensor.permute(1, 2, 0, 3) + return tensor.permute(0, 2, 1, 3) + + +def _from_bhsd(tensor, qkv_format): + """Convert logical BHSD test tensors to SBHD/BSHD.""" + if qkv_format == "sbhd": + return tensor.permute(2, 0, 1, 3).contiguous() + return tensor.permute(0, 2, 1, 3).contiguous() + + +def _pytorch_softcap_attention(q, k, v, qkv_format, softmax_scale, softcap): + """PyTorch reference for softcapped scaled dot-product attention.""" + q_bhsd = _to_bhsd(q, qkv_format).float() + k_bhsd = _to_bhsd(k, qkv_format).float() + v_bhsd = _to_bhsd(v, qkv_format).float() + scores = torch.matmul(q_bhsd, k_bhsd.transpose(-2, -1)) * softmax_scale + scores = softcap * torch.tanh(scores / softcap) + probs = torch.softmax(scores, dim=-1) + out = _from_bhsd(torch.matmul(probs, v_bhsd), qkv_format).to(v.dtype) + return out.reshape(*out.shape[:-2], out.shape[-2] * out.shape[-1]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") +@pytest.mark.skipif(get_cudnn_version() < (9, 6, 0), reason="cuDNN 9.6.0+ is required.") +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("qkv_format", ["sbhd", "bshd"]) +@pytest.mark.parametrize( + "score_mod_case, scalar_loss", + [ + ("causal", False), + ("causal", True), + ("softcap", False), + ("post_scale_bias", False), + ], +) +def test_dot_product_attention_score_mod(dtype, qkv_format, score_mod_case, scalar_loss): + """Compare score_mod attention against equivalent reference implementations.""" + try: + flex_attention._import_cudnn_frontend() + except ImportError: + pytest.skip("cuDNN frontend Python package is required for score_mod attention.") + if score_mod_case == "softcap" and get_device_compute_capability() < (9, 0): + pytest.skip("Softcap score_mod tests require sm90+.") + + reset_rng_states() + + config = ModelConfig( + 2, + 64 if score_mod_case == "causal" else 16, + 4, + 64, + attn_mask_type="no_mask", + ) + available_backends, _, fused_attn_backends = get_available_attention_backends( + config, + qkv_dtype=dtype, + qkv_layout=f"{qkv_format}_{qkv_format}_{qkv_format}", + score_mod=True, + score_mod_bprop=True, + ) + if not available_backends[1] or not fused_attn_backends: + pytest.skip("FusedAttention is not available for this score_mod configuration.") + + if score_mod_case == "post_scale_bias": + bias_config = ModelConfig( + config.batch_size, + config.max_seqlen_q, + config.num_heads, + config.head_dim_qk, + attn_mask_type="no_mask", + attn_bias_type="post_scale_bias", + bias_shape="1hss", + ) + bias_available_backends, _, bias_fused_attn_backends = get_available_attention_backends( + bias_config, + qkv_dtype=dtype, + qkv_layout=f"{qkv_format}_{qkv_format}_{qkv_format}", + ) + if not bias_available_backends[1] or not bias_fused_attn_backends: + pytest.skip("FusedAttention is not available for post_scale_bias reference.") + + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True + + if qkv_format == "sbhd": + q_shape = (config.max_seqlen_q, config.batch_size, config.num_heads, config.head_dim_qk) + kv_shape = q_shape + else: + q_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) + kv_shape = q_shape + + if score_mod_case == "softcap": + q = torch.randn(q_shape, dtype=dtype, device="cuda").requires_grad_() + k = torch.randn(kv_shape, dtype=dtype, device="cuda").requires_grad_() + v = (0.1 * torch.randn(kv_shape, dtype=dtype, device="cuda")).requires_grad_() + else: + q = (0.1 * torch.randn(q_shape, dtype=dtype, device="cuda")).requires_grad_() + k = (0.1 * torch.randn(kv_shape, dtype=dtype, device="cuda")).requires_grad_() + v = (0.1 * torch.randn(kv_shape, dtype=dtype, device="cuda")).requires_grad_() + q_ref, k_ref, v_ref = [x.detach().clone().requires_grad_() for x in (q, k, v)] + + flex_attn = DotProductAttention( + config.num_heads, + config.head_dim_qk, + qkv_format=qkv_format, + attn_mask_type="no_mask", + layer_number=1, + ).to(dtype=dtype, device="cuda") + + if score_mod_case == "causal": + score_mod_kwargs = { + "score_mod": _score_mod_causal, + "score_mod_bprop": _score_mod_causal_bprop, + "score_mod_tensors": {"neg_inf": torch.full((1, 1, 1, 1), -1e9)}, + "score_mod_bprop_tensors": {"zero": torch.full((1, 1, 1, 1), 0.0)}, + } + ref_attn = DotProductAttention( + config.num_heads, + config.head_dim_qk, + qkv_format=qkv_format, + attn_mask_type="causal", + layer_number=1, + ).to(dtype=dtype, device="cuda") + out_ref = ref_attn(q_ref, k_ref, v_ref, qkv_format=qkv_format, attn_mask_type="causal") + tols = dict(atol=5e-2, rtol=5e-2) + elif score_mod_case == "softcap": + softcap = 0.8 + softcap_tensor = torch.full((1, 1, 1, 1), softcap) + softcap_score_mod = _ScoreModSoftcap() + score_mod_kwargs = { + "score_mod": softcap_score_mod.forward, + "score_mod_bprop": softcap_score_mod.backward, + "score_mod_tensors": {"softcap": softcap_tensor}, + "score_mod_bprop_tensors": {"softcap": softcap_tensor}, + } + out_ref = _pytorch_softcap_attention( + q_ref, + k_ref, + v_ref, + qkv_format, + 1.0 / config.head_dim_qk**0.5, + softcap, + ) + tols = dict(atol=7e-2, rtol=7e-2) + else: + assert score_mod_case == "post_scale_bias" + score_mod_kwargs = { + "score_mod": _score_mod_post_scale_bias, + "score_mod_bprop": _score_mod_identity_bprop, + } + ref_attn = DotProductAttention( + config.num_heads, + config.head_dim_qk, + qkv_format=qkv_format, + attn_mask_type="no_mask", + layer_number=1, + ).to(dtype=dtype, device="cuda") + out_ref = ref_attn( + q_ref, + k_ref, + v_ref, + qkv_format=qkv_format, + attn_mask_type="no_mask", + core_attention_bias_type="post_scale_bias", + core_attention_bias=_post_scale_bias(config, dtype), + ) + tols = dict(atol=5e-2, rtol=5e-2) + + out = flex_attn( + q, + k, + v, + qkv_format=qkv_format, + attn_mask_type="no_mask", + **score_mod_kwargs, + ) + + if scalar_loss: + out.sum().backward() + out_ref.sum().backward() + else: + d_out = torch.randn_like(out) + out.backward(d_out) + out_ref.backward(d_out) + + torch.testing.assert_close(out, out_ref, **tols) + torch.testing.assert_close(q.grad, q_ref.grad, **tols) + torch.testing.assert_close(k.grad, k_ref.grad, **tols) + torch.testing.assert_close(v.grad, v_ref.grad, **tols) diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index c662252f9e5..cdd98d2445a 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -375,14 +375,19 @@ def get_tols(config, module, backend, dtype): torch.bfloat16: (3.5e-2, 3.5e-2), } else: - if backend == "UnfusedAttention": + if backend in ("UnfusedAttention", "FlashAttention"): tols = { torch.half: (1.6e-2, 1.6e-2), torch.bfloat16: (1.2e-1, 1e-1), } else: + # head_dim > 128 in fp16 is the worst case for accumulated rounding, and the + # full-sequence vs incremental-KV-cache paths use different kernels/mask types. + # On sm80 with older cuDNN the agreement grazes 1e-2 on a single element, so the + # fp16 tolerance is widened slightly. Tolerances were originally calibrated on + # Hopper/Blackwell + newer cuDNN. tols = { - torch.half: (1e-2, 1e-2), + torch.half: (1.5e-2, 1.5e-2), torch.bfloat16: (8e-2, 7e-2), } if module == "DotProductAttention": diff --git a/tests/pytorch/attention/test_linear_mxfp8_attention.py b/tests/pytorch/attention/test_linear_mxfp8_attention.py new file mode 100644 index 00000000000..f1bba7bc9a1 --- /dev/null +++ b/tests/pytorch/attention/test_linear_mxfp8_attention.py @@ -0,0 +1,449 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP8 end-to-end attention unit test - DSv3 671B MLA dimensions. + +Path: LayerNormLinear(Q/KV, MXFP8) -> MLA-RoPE (Triton) + -> DotProductAttention(MXFP8) -> Linear(out, MXFP8). +Tensor layout: sbhd (seq-first) throughout. + +Run: + python3 -m pytest tests/pytorch/attention/test_linear_mxfp8_attention.py -v -s + +Expected benchmark output (GB200, b=1, s=4096): + [PERF] b=1 s=4096: + BF16 fprop: 8.582 ms (477274 tok/s) + BF16 bprop: 14.006 ms (292445 tok/s) + MXFP8 fprop: 5.210 ms (786180 tok/s) + MXFP8 bprop: 8.763 ms (467428 tok/s) + Fprop speedup: 1.65x + Bprop speedup: 1.60x +""" + +import os +import pathlib +import sys + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends +from transformer_engine.pytorch.utils import get_cudnn_version + +_current_file = pathlib.Path(__file__).resolve() +sys.path = [str(_current_file.parent.parent)] + sys.path +from utils import ModelConfig, compare_and_assert, get_available_attention_backends +from mla_rope_utils import apply_mla_rope, build_rope_tables + + +try: + from transformer_engine.common.recipe import MXFP8BlockScaling + + mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +except (ImportError, AttributeError): + mxfp8_available = False + reason_for_no_mxfp8 = "MXFP8BlockScaling not available in this build" + +# DSv3 671B MLA dims (micro_batch=1, seq_len=4096) +NUM_HEADS = 128 +HEAD_DIM_ROPE = 64 +HEAD_DIM_NOPE = 128 +HEAD_DIM_QK = HEAD_DIM_NOPE + HEAD_DIM_ROPE # 192 +HEAD_DIM_V = 128 +HIDDEN_SIZE = NUM_HEADS * HEAD_DIM_V # 16384 +Q_SIZE = NUM_HEADS * HEAD_DIM_QK # 24576 +KV_SIZE = NUM_HEADS * (HEAD_DIM_NOPE + HEAD_DIM_V) # 32768 +KV_PROJ_SIZE = KV_SIZE + HEAD_DIM_ROPE # [K_NOPE | V] per head, plus shared K RoPE +SEED = 42 + +WARMUP_ITERS = 10 +TIMED_ITERS = 100 +ATOL = 5e-1 +RTOL = 5e-2 +RMSE_TOL = 0.11 +_DETERMINISTIC = ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() +) + + +@pytest.fixture(autouse=True) +def reset_global_fp8_state(): + yield + FP8GlobalStateManager.reset() + + +def _set_seed(seed: int = SEED) -> None: + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + +def _build_modules(dtype: torch.dtype = torch.bfloat16): + def _make_modules(): + q_proj = te.LayerNormLinear( + HIDDEN_SIZE, + Q_SIZE, + bias=True, + params_dtype=dtype, + device="cuda", + ) + kv_proj = te.LayerNormLinear( + HIDDEN_SIZE, + KV_PROJ_SIZE, + bias=True, + params_dtype=dtype, + device="cuda", + ) + dpa = te.DotProductAttention( + num_attention_heads=NUM_HEADS, + kv_channels=(HEAD_DIM_QK, HEAD_DIM_V), + attention_dropout=0.0, + qkv_format="sbhd", + ).to(device="cuda") + out = te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=True).to(dtype=dtype, device="cuda") + return q_proj, kv_proj, dpa, out + + base = _make_modules() + mxfp8 = _make_modules() + + with torch.no_grad(): + for dst_module, src_module in zip(mxfp8, base): + for p_dst, p_src in zip(dst_module.parameters(), src_module.parameters()): + p_dst.copy_(p_src) + + for m in base + mxfp8: + m.train() + + return base, mxfp8 + + +def _require_attention_backends( + batch_size: int, + seq_len: int, + fp8_recipe, +) -> None: + if get_cudnn_version() < (9, 2, 1): + pytest.skip("cuDNN 9.2.1+ is required for FP8 fused attention.") + + config = ModelConfig( + batch_size, + seq_len, + NUM_HEADS, + HEAD_DIM_QK, + head_dim_v=HEAD_DIM_V, + ) + fp8_meta = {"recipe": fp8_recipe} + fp8_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=torch.float8_e4m3fn, + qkv_layout="sbhd_sbhd_sbhd", + fp8=True, + fp8_meta=fp8_meta, + is_training=True, + deterministic=_DETERMINISTIC, + ) + _, fused_attn_supported_fp8, _ = fp8_backends + if not fused_attn_supported_fp8: + pytest.skip("No fused FP8 attention backend available for DSv3 MLA shape.") + + bf16_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=torch.bfloat16, + qkv_layout="sbhd_sbhd_sbhd", + is_training=True, + deterministic=_DETERMINISTIC, + ) + if sum(bf16_backends) < 1: + pytest.skip("No BF16 attention backend available for DSv3 MLA shape.") + + _attention_backends["backend_selection_requires_update"] = True + + +def _run_projections( + q_proj: te.LayerNormLinear, + kv_proj: te.LayerNormLinear, + x: torch.Tensor, + is_first_microbatch: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Project Q and workload-shaped KV, then expose metadata-only RoPE views.""" + if is_first_microbatch is None: + q_flat = q_proj(x) + kv_flat = kv_proj(x) + else: + q_flat = q_proj(x, is_first_microbatch=is_first_microbatch) + kv_flat = kv_proj(x, is_first_microbatch=is_first_microbatch) + + s, b, _ = x.shape + q = q_flat.view(s, b, NUM_HEADS, HEAD_DIM_QK) + kv = kv_flat[:, :, :KV_SIZE].view(s, b, NUM_HEADS, HEAD_DIM_NOPE + HEAD_DIM_V) + k_pos_emb = kv_flat[:, :, KV_SIZE:].view(s, b, 1, HEAD_DIM_ROPE) + return q_flat, kv_flat, q, kv, k_pos_emb + + +def _run_forward_bf16( + modules: tuple, + x: torch.Tensor, + rope_tables: tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q_proj, kv_proj, dpa, out_linear = modules + _, _, q, kv, k_pos_emb = _run_projections(q_proj, kv_proj, x) + q, k, v = apply_mla_rope(q, kv, k_pos_emb, cos_table=rope_tables[0], sin_table=rope_tables[1]) + attn_out = dpa(q, k, v, qkv_format="sbhd") + return q, k, v, out_linear(attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE)) + + +def _run_forward_mxfp8( + modules: tuple, + x: torch.Tensor, + recipe, + rope_tables: tuple[torch.Tensor, torch.Tensor], + is_first_microbatch: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """is_first_microbatch=True caches quantized weights; False reuses cache; None re-quantizes.""" + q_proj, kv_proj, dpa, out_linear = modules + + with te.fp8_autocast(enabled=True, fp8_recipe=recipe): + _, _, q, kv, k_pos_emb = _run_projections( + q_proj, + kv_proj, + x, + is_first_microbatch, + ) + q, k, v = apply_mla_rope( + q, + kv, + k_pos_emb, + cos_table=rope_tables[0], + sin_table=rope_tables[1], + ) + attn_out = dpa(q, k, v, qkv_format="sbhd") + out = out_linear( + attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE), + is_first_microbatch=is_first_microbatch, + ) + + return q, k, v, out + + +def _compute_errors(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float]: + diff = (a.float() - b.float()).abs() + return diff.max().item(), diff.pow(2).mean().sqrt().item() + + +def _clear_training_step_grads(modules: tuple, x: torch.Tensor) -> None: + x.grad = None + for module in modules: + for param in module.parameters(): + param.grad = None + + +def _benchmark_training_step( + forward_fn, + modules: tuple, + x: torch.Tensor, + *forward_args, + warmup: int = WARMUP_ITERS, + iters: int = TIMED_ITERS, +) -> tuple[float, float]: + for _ in range(warmup): + _clear_training_step_grads(modules, x) + *_, out = forward_fn(modules, x, *forward_args) + out.sum().backward() + torch.cuda.synchronize() + + fprop_ms = 0.0 + bprop_ms = 0.0 + for _ in range(iters): + _clear_training_step_grads(modules, x) + + fprop_start = torch.cuda.Event(enable_timing=True) + fprop_end = torch.cuda.Event(enable_timing=True) + bprop_start = torch.cuda.Event(enable_timing=True) + bprop_end = torch.cuda.Event(enable_timing=True) + + fprop_start.record() + *_, out = forward_fn(modules, x, *forward_args) + fprop_end.record() + + bprop_start.record() + out.sum().backward() + bprop_end.record() + + torch.cuda.synchronize() + fprop_ms += fprop_start.elapsed_time(fprop_end) + bprop_ms += bprop_start.elapsed_time(bprop_end) + + return fprop_ms / iters, bprop_ms / iters + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("batch_size", [1]) +@pytest.mark.parametrize("seq_len", [4096]) +class TestLinearMXFP8Attention: + + def test_accuracy(self, batch_size: int, seq_len: int) -> None: + """Validate MXFP8 against BF16 using loose tolerances.""" + fp8_recipe = MXFP8BlockScaling(fp8_dpa=True) + _require_attention_backends(batch_size, seq_len, fp8_recipe) + _set_seed() + baseline_modules, mxfp8_modules = _build_modules() + x = torch.randn(seq_len, batch_size, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + rope_tables = build_rope_tables(seq_len, device=x.device) + + q_bf16, k_bf16, v_bf16, out_bf16 = _run_forward_bf16(baseline_modules, x, rope_tables) + q_mxfp8, k_mxfp8, v_mxfp8, out_mxfp8 = _run_forward_mxfp8( + mxfp8_modules, + x, + fp8_recipe, + rope_tables, + ) + + for name, tensor in [("Q", q_mxfp8), ("K", k_mxfp8), ("V", v_mxfp8)]: + assert not torch.isnan(tensor).any(), f"MXFP8 {name} contains NaN" + assert not torch.isinf(tensor).any(), f"MXFP8 {name} contains Inf" + assert tensor.float().abs().max() > 0, f"MXFP8 {name} is all zeros" + + assert not torch.isnan(out_mxfp8).any(), "MXFP8 output contains NaN" + assert not torch.isinf(out_mxfp8).any(), "MXFP8 output contains Inf" + assert out_mxfp8.float().abs().max() > 0, "MXFP8 output is all zeros" + + max_abs_q, rms_q = _compute_errors(q_bf16, q_mxfp8) + print(f"\n[Q] b={batch_size} s={seq_len}: max_abs={max_abs_q:.6f} rms={rms_q:.6f}") + compare_and_assert( + q_mxfp8, + q_bf16, + "q_mxfp8", + "q_bf16", + ATOL, + RTOL, + RMSE_TOL, + True, + ) + + max_abs_k, rms_k = _compute_errors(k_bf16, k_mxfp8) + print(f"[K] b={batch_size} s={seq_len}: max_abs={max_abs_k:.6f} rms={rms_k:.6f}") + compare_and_assert( + k_mxfp8, + k_bf16, + "k_mxfp8", + "k_bf16", + ATOL, + RTOL, + RMSE_TOL, + True, + ) + + max_abs_v, rms_v = _compute_errors(v_bf16, v_mxfp8) + print(f"[V] b={batch_size} s={seq_len}: max_abs={max_abs_v:.6f} rms={rms_v:.6f}") + compare_and_assert( + v_mxfp8, + v_bf16, + "v_mxfp8", + "v_bf16", + ATOL, + RTOL, + RMSE_TOL, + True, + ) + + max_abs_out, rms_out = _compute_errors(out_bf16, out_mxfp8) + print(f"[OUT] b={batch_size} s={seq_len}: max_abs={max_abs_out:.6f} rms={rms_out:.6f}") + compare_and_assert( + out_mxfp8, + out_bf16, + "out_mxfp8", + "out_bf16", + ATOL, + RTOL, + RMSE_TOL, + True, + ) + + def test_backward(self, batch_size: int, seq_len: int) -> None: + """Gradients must flow end-to-end without NaN/Inf.""" + fp8_recipe = MXFP8BlockScaling(fp8_dpa=True) + _require_attention_backends(batch_size, seq_len, fp8_recipe) + _set_seed() + _, mxfp8_modules = _build_modules() + + x = torch.randn( + seq_len, + batch_size, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + requires_grad=True, + ) + rope_tables = build_rope_tables(seq_len, device=x.device) + + *_, out_mxfp8 = _run_forward_mxfp8(mxfp8_modules, x, fp8_recipe, rope_tables) + out_mxfp8.sum().backward() + + assert x.grad is not None, "MXFP8 path: input grad is None" + assert not torch.isnan(x.grad).any(), "MXFP8 path: input grad NaN" + assert not torch.isinf(x.grad).any(), "MXFP8 path: input grad Inf" + + q_fp8, kv_fp8, _, out_fp8 = mxfp8_modules + for name, mod in [("q_proj", q_fp8), ("kv_proj", kv_fp8), ("out_linear", out_fp8)]: + for p in mod.parameters(): + if p.grad is not None: + assert not torch.isnan(p.grad).any(), f"MXFP8 {name} param grad NaN" + assert not torch.isinf(p.grad).any(), f"MXFP8 {name} param grad Inf" + + dx_rms = x.grad.float().pow(2).mean().sqrt().item() + print(f"\n[BPROP] b={batch_size} s={seq_len}: dx rms={dx_rms:.6f}") + assert dx_rms > 0.0, "MXFP8 path: input grad is all zeros (no gradient flow)" + + def test_performance(self, batch_size: int, seq_len: int) -> None: + """Benchmark the normal MXFP8 training step against BF16.""" + fp8_recipe = MXFP8BlockScaling(fp8_dpa=True) + _require_attention_backends(batch_size, seq_len, fp8_recipe) + _set_seed() + baseline_modules, mxfp8_modules = _build_modules() + x = torch.randn( + seq_len, + batch_size, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + requires_grad=True, + ) + rope_tables = build_rope_tables(seq_len, device=x.device) + + mxfp8_fprop_ms, mxfp8_bprop_ms = _benchmark_training_step( + _run_forward_mxfp8, mxfp8_modules, x, fp8_recipe, rope_tables + ) + mxfp8_fprop_tok = (batch_size * seq_len) / (mxfp8_fprop_ms / 1000.0) + mxfp8_bprop_tok = (batch_size * seq_len) / (mxfp8_bprop_ms / 1000.0) + + bf16_fprop_ms, bf16_bprop_ms = _benchmark_training_step( + _run_forward_bf16, baseline_modules, x, rope_tables + ) + bf16_fprop_tok = (batch_size * seq_len) / (bf16_fprop_ms / 1000.0) + bf16_bprop_tok = (batch_size * seq_len) / (bf16_bprop_ms / 1000.0) + fprop_speedup = bf16_fprop_ms / mxfp8_fprop_ms + bprop_speedup = bf16_bprop_ms / mxfp8_bprop_ms + print( + f"\n[PERF] b={batch_size} s={seq_len}:" + f"\n BF16 fprop: {bf16_fprop_ms:.3f} ms ({bf16_fprop_tok:.0f} tok/s)" + f"\n BF16 bprop: {bf16_bprop_ms:.3f} ms ({bf16_bprop_tok:.0f} tok/s)" + f"\n MXFP8 fprop: {mxfp8_fprop_ms:.3f} ms ({mxfp8_fprop_tok:.0f} tok/s)" + f"\n MXFP8 bprop: {mxfp8_bprop_ms:.3f} ms ({mxfp8_bprop_tok:.0f} tok/s)" + f"\n Fprop speedup: {fprop_speedup:.2f}x" + f"\n Bprop speedup: {bprop_speedup:.2f}x" + ) + + assert fprop_speedup > 1.0, ( + "MXFP8 fprop should be faster than BF16: " + f"got {mxfp8_fprop_ms:.3f} ms vs BF16 {bf16_fprop_ms:.3f} ms " + f"(speedup={fprop_speedup:.2f}x)" + ) + assert bprop_speedup > 1.0, ( + "MXFP8 bprop should be faster than BF16: " + f"got {mxfp8_bprop_ms:.3f} ms vs BF16 {bf16_bprop_ms:.3f} ms " + f"(speedup={bprop_speedup:.2f}x)" + ) diff --git a/tests/pytorch/debug/run_distributed.py b/tests/pytorch/debug/run_distributed.py index 285ec7ba0c4..e068f0993dc 100644 --- a/tests/pytorch/debug/run_distributed.py +++ b/tests/pytorch/debug/run_distributed.py @@ -13,7 +13,7 @@ import torch import torch.distributed as dist import transformer_engine -import transformer_engine_torch as tex +from transformer_engine.pytorch import DType import nvdlfw_inspect.api as debug_api from transformer_engine.debug import set_weight_tensor_tp_group_reduce from transformer_engine.pytorch import is_fp8_available @@ -683,7 +683,7 @@ def _run_test_with_combinations( ) # test_fake_quant_fp8 - dtype_options = [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2, None] + dtype_options = [DType.kFloat8E4M3, DType.kFloat8E5M2, None] _run_test_with_combinations( test_fake_quant_fp8, dtype_options, diff --git a/tests/pytorch/debug/test_api_features.py b/tests/pytorch/debug/test_api_features.py index 5387634cb38..b1589d978dd 100644 --- a/tests/pytorch/debug/test_api_features.py +++ b/tests/pytorch/debug/test_api_features.py @@ -4,12 +4,12 @@ import torch from transformer_engine.pytorch import Float8Tensor, Float8Quantizer +from transformer_engine.pytorch import DType import nvdlfw_inspect.api as debug_api try: import transformer_engine - import transformer_engine_torch as tex except (ImportError, ModuleNotFoundError): print("Could not find TransformerEngine package.") exit(1) @@ -128,12 +128,12 @@ def test_per_tensor_scaling(configs_dir, feature_dirs): default_quantizer1 = Float8Quantizer( scale=torch.tensor([1]).cuda(), amax=torch.tensor([0]).cuda(), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, ) default_quantizer2 = Float8Quantizer( scale=torch.tensor([1]).cuda(), amax=torch.tensor([0]).cuda(), - fp8_dtype=tex.DType.kFloat8E5M2, + fp8_dtype=DType.kFloat8E5M2, ) output1 = debug_api.transformer_engine.modify_tensor( @@ -145,7 +145,7 @@ def test_per_tensor_scaling(configs_dir, feature_dirs): tensor=tensor, ) assert type(output1) == Float8Tensor - assert output1._fp8_dtype == tex.DType.kFloat8E4M3 + assert output1._fp8_dtype == DType.kFloat8E4M3 output2 = debug_api.transformer_engine.modify_tensor( "decoder.1.mlp.fc1", @@ -156,7 +156,7 @@ def test_per_tensor_scaling(configs_dir, feature_dirs): iteration=0, ) assert type(output2) == Float8Tensor - assert output2._fp8_dtype == tex.DType.kFloat8E5M2 + assert output2._fp8_dtype == DType.kFloat8E5M2 assert not debug_api.transformer_engine.modify_tensor_enabled( "decoder.1.mlp.fc1", @@ -234,7 +234,7 @@ def test_statistics_collection(configs_dir, feature_dirs): quantizer = Float8Quantizer( scale=torch.full([1], 1.0).cuda(), amax=torch.full([1], 1.0).cuda(), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, ) tensor_fp8 = quantizer(tensor) @@ -372,7 +372,7 @@ def log_stats(): quantizer = Float8Quantizer( scale=torch.full([1], 1.0).cuda(), amax=torch.full([1], 1.0).cuda(), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, ) def fp8_tensor(t): diff --git a/tests/pytorch/debug/test_numerics.py b/tests/pytorch/debug/test_numerics.py index ab9a2d054a0..c5ca7900ab7 100644 --- a/tests/pytorch/debug/test_numerics.py +++ b/tests/pytorch/debug/test_numerics.py @@ -15,7 +15,6 @@ import nvdlfw_inspect.api as debug_api import transformer_engine.debug import transformer_engine.pytorch as tepytorch -import transformer_engine_torch as tex from transformer_engine.common.recipe import DelayedScaling, Format from transformer_engine.pytorch.quantization import _default_sf_compute from transformer_engine.pytorch import ( @@ -57,7 +56,7 @@ def _cast_to_fp8(tensor, scale, dtype): def _get_current_scale(tensor, fp8_dtype): - if fp8_dtype == tex.DType.kFloat8E4M3: + if fp8_dtype == tepytorch.DType.kFloat8E4M3: fp8_max = Format.E4M3.value.max_fwd else: fp8_max = Format.E5M2.value.max_fwd @@ -93,19 +92,19 @@ def _emulate_linear( input: torch.Tensor, weight: torch.Tensor, fprop_fp8: bool = False, - fprop_input_fake_quant: tex.DType = None, + fprop_input_fake_quant: tepytorch.DType = None, fprop_input_scale: torch.Tensor = None, - fprop_weight_fake_quant: tex.DType = None, + fprop_weight_fake_quant: tepytorch.DType = None, fprop_weight_scale: torch.Tensor = None, dgrad_fp8: bool = False, - dgrad_gradient_fake_quant: tex.DType = None, + dgrad_gradient_fake_quant: tepytorch.DType = None, dgrad_gradient_scale: torch.Tensor = None, - dgrad_weight_fake_quant: tex.DType = None, + dgrad_weight_fake_quant: tepytorch.DType = None, dgrad_weight_scale: torch.Tensor = None, wgrad_fp8: bool = False, - wgrad_gradient_fake_quant: tex.DType = None, + wgrad_gradient_fake_quant: tepytorch.DType = None, wgrad_gradient_scale: torch.Tensor = None, - wgrad_input_fake_quant: tex.DType = None, + wgrad_input_fake_quant: tepytorch.DType = None, wgrad_input_scale: torch.Tensor = None, loss_multiplier: float = 1.0, activation_sync=None, @@ -116,10 +115,10 @@ def _emulate_linear( activation = _fp8_gemm_kernel( input, _scalar(fprop_input_scale or 1.0), - tex.DType.kFloat8E4M3, + tepytorch.DType.kFloat8E4M3, weight, _scalar(fprop_weight_scale or 1.0), - tex.DType.kFloat8E4M3, + tepytorch.DType.kFloat8E4M3, _2X_ACC_FPROP, ) activation = activation.clone().detach().contiguous().requires_grad_(True) @@ -152,10 +151,10 @@ def _emulate_linear( dgrad = _fp8_gemm_kernel( weight.T, _scalar(dgrad_weight_scale or 1.0), - tex.DType.kFloat8E4M3, + tepytorch.DType.kFloat8E4M3, gradient, _scalar(dgrad_gradient_scale or 1.0), - tex.DType.kFloat8E5M2, + tepytorch.DType.kFloat8E5M2, _2X_ACC_DGRAD, ).T else: @@ -176,10 +175,10 @@ def _emulate_linear( wgrad = _fp8_gemm_kernel( input.T, _scalar(wgrad_input_scale or 1.0), - tex.DType.kFloat8E4M3, + tepytorch.DType.kFloat8E4M3, gradient.T, _scalar(wgrad_gradient_scale or 1.0), - tex.DType.kFloat8E5M2, + tepytorch.DType.kFloat8E5M2, _2X_ACC_WGRAD, ).T else: @@ -470,17 +469,17 @@ def set_scaling_factors(model, input_kwargs, fp8_kwargs): def set_current_scaling_factors(x, weight, y, input_kwargs, fp8_kwargs): # Compute per tensor scaling factor if respective flag in input_kwargs is set. if input_kwargs["fprop_inp"]: - fp8_kwargs["fprop_input_scale"] = tex.DType.kFloat8E4M3 + fp8_kwargs["fprop_input_scale"] = tepytorch.DType.kFloat8E4M3 if input_kwargs["fprop_weight"]: - fp8_kwargs["fprop_weight_scale"] = tex.DType.kFloat8E4M3 + fp8_kwargs["fprop_weight_scale"] = tepytorch.DType.kFloat8E4M3 if input_kwargs["dgrad_grad"]: - fp8_kwargs["dgrad_gradient_scale"] = tex.DType.kFloat8E5M2 + fp8_kwargs["dgrad_gradient_scale"] = tepytorch.DType.kFloat8E5M2 if input_kwargs["dgrad_weight"]: - fp8_kwargs["dgrad_weight_scale"] = tex.DType.kFloat8E4M3 + fp8_kwargs["dgrad_weight_scale"] = tepytorch.DType.kFloat8E4M3 if input_kwargs["wgrad_grad"]: - fp8_kwargs["wgrad_gradient_scale"] = tex.DType.kFloat8E5M2 + fp8_kwargs["wgrad_gradient_scale"] = tepytorch.DType.kFloat8E5M2 if input_kwargs["wgrad_input"]: - fp8_kwargs["wgrad_input_scale"] = tex.DType.kFloat8E4M3 + fp8_kwargs["wgrad_input_scale"] = tepytorch.DType.kFloat8E4M3 @create_config_file @@ -651,7 +650,7 @@ def init_and_warmup(): all_combinations = list( - itertools.product([tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2, None], repeat=6) + itertools.product([tepytorch.DType.kFloat8E4M3, tepytorch.DType.kFloat8E5M2, None], repeat=6) ) subset_combinations = random.sample(all_combinations, 10) @@ -687,7 +686,7 @@ def test_fake_quant_fp8( def fake_quant_fp8_create_config( fprop_inp, fprop_weight, dgrad_weight, dgrad_grad, wgrad_input, wgrad_grad, config_file ): - format_to_str = {tex.DType.kFloat8E4M3: "FP8E4M3", tex.DType.kFloat8E5M2: "FP8E5M2"} + format_to_str = {tepytorch.DType.kFloat8E4M3: "FP8E4M3", tepytorch.DType.kFloat8E5M2: "FP8E5M2"} gemms = "" def _add_tensor(quant_format, tensor): diff --git a/tests/pytorch/distributed/run_gemm_with_overlap.py b/tests/pytorch/distributed/run_gemm_with_overlap.py index 96a7e432316..a68929ab364 100644 --- a/tests/pytorch/distributed/run_gemm_with_overlap.py +++ b/tests/pytorch/distributed/run_gemm_with_overlap.py @@ -151,6 +151,27 @@ def _parse_args(argv=None, namespace=None): parser.add_argument( "--use-cuda-graphs", action="store_true", default=False, help="Use CUDA graphs." ) + parser.add_argument( + "--use-cublasmp", action="store_true", default=False, help="Use cuBLASMp backend." + ) + parser.add_argument( + "--rtol", + type=float, + default=None, + help=( + "Override the relative-error tolerance used in the numerical check. " + "When unset, defaults to 0.125 for FP8/MXFP8 and 0.02 otherwise." + ), + ) + parser.add_argument( + "--atol", + type=float, + default=None, + help=( + "Override the absolute-error tolerance used in the numerical check. " + "When unset, defaults to 0.0625 for FP8/MXFP8 and 0.002 otherwise." + ), + ) parser.add_argument( "-v", "--verbose", action="store_true", default=False, help="Verbose info messages." ) @@ -203,6 +224,7 @@ def _main(opts): capture_output=True, text=True, shell=True, + check=False, ) if result.stdout == "0": # Extra checks for non-MNNVL platforms @@ -306,7 +328,7 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None helper = ( tex.CommOverlapHelper() if tex.ubuf_built_with_mpi() - else tex.CommOverlapHelper(bootstrap_pg) + else tex.CommOverlapHelper(bootstrap_pg, tp_group) ) # Initialize userbuffers with (M, N) buffer @@ -322,50 +344,59 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None and opts.comm_type == tex.CommOverlapType.AG ): buffer_dtype = torch.uint8 - ub_obj = ( - tex.CommOverlapP2P( + if opts.p2p: + ub_obj = tex.CommOverlapP2P( (outer_size, hidden_size), buffer_dtype, helper, tp_size, # Tensor-parallel group size (may be different than LOCAL_SIZE) opts.comm_type, + use_cublasmp=opts.use_cublasmp, + num_comm_sm=3 if opts.use_cublasmp else 1, set_sm_margin=opts.comm_type == tex.CommOverlapType.RS or opts.atomic, atomic_gemm=opts.atomic, aggregate=opts.aggregate, use_ce=not (opts.atomic and bool(int(os.getenv("NVTE_AG_P2P_MULTI_ATOMIC", "0")))), ) - if opts.p2p - else tex.CommOverlap( + else: + ub_obj = tex.CommOverlap( (outer_size, hidden_size), buffer_dtype, helper, tp_size, # Tensor-parallel group size (may be different than LOCAL_SIZE) + use_cublasmp=opts.use_cublasmp, + comm_type=opts.comm_type, + num_comm_sm=16, atomic_gemm=opts.atomic, ) - ) # Numerical check on AG + atomic GEMM requires testing an AG+RS pair ub_obj2 = None if opts.atomic and opts.comm_type == tex.CommOverlapType.AG and opts.check_numerics: - ub_obj2 = ( - tex.CommOverlapP2P( + ub2_buffer_dtype = torch.uint8 if opts.fp8_output else torch.bfloat16 + if opts.atomic_rs_p2p: + ub_obj2 = tex.CommOverlapP2P( (outer_size, hidden_size), - torch.uint8 if opts.fp8_output else torch.bfloat16, + ub2_buffer_dtype, helper, - tp_size, # Tensor-parallel group size (may be different than LOCAL_SIZE) + tp_size, tex.CommOverlapType.RS, + use_cublasmp=opts.use_cublasmp, + num_comm_sm=16 if opts.use_cublasmp else 1, set_sm_margin=True, atomic_gemm=True, ) - if opts.atomic_rs_p2p - else tex.CommOverlap( + else: + ub_obj2 = tex.CommOverlap( (outer_size, hidden_size), - torch.uint8 if opts.fp8_output else torch.bfloat16, + ub2_buffer_dtype, helper, - tp_size, # Tensor-parallel group size (may be different than LOCAL_SIZE) + tp_size, + use_cublasmp=opts.use_cublasmp, + comm_type=tex.CommOverlapType.RS, + num_comm_sm=3 if opts.use_cublasmp else 16, atomic_gemm=True, ) - ) # Figure out problem sizing: # M = sequence * batch @@ -408,7 +439,7 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None mean=0.0, std=opts.std, ) - if ub_obj2 is not None: + if opts.comm_type == tex.CommOverlapType.AG and ub_obj2 is not None: kernel2_t = torch.nn.init.normal_( torch.empty(local_kernel2_t_shape, dtype=torch.bfloat16, device="cuda"), mean=0.0, @@ -426,25 +457,19 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None ) else: if opts.comm_type == tex.CommOverlapType.AG: - # AG Kernel: (K/P, N) -> gather -> (K, N) -> T -> (N, K) - ker_g = torch.transpose( - te.distributed.gather_along_first_dim(kernel_t, tp_group)[0], 0, 1 - ).to(dtype=torch.float32) - # AG Input: (M/P, N) -> gather -> (M, N) - inp_g = te.distributed.gather_along_first_dim(inp, tp_group)[0].to(dtype=torch.float32) + # AG Kernel: Keep local (K/P, N) for per-rank reference comparison + ker_g = kernel_t + # AG Input: (M/P, N) -> gather -> (M, N); full input needed for local GEMM chunk + inp_g = te.distributed.gather_along_first_dim(inp, tp_group)[0] if ub_obj2 is not None: - ker2_g = te.distributed.gather_along_first_dim( - torch.transpose(kernel2_t, 0, 1), tp_group - )[0].to(dtype=torch.float32) + # Keep kernel2 local (N, K/P) for per-rank reference comparison + ker2_g = kernel2_t else: - # RS Kernel: (N, K/P) -> T -> (K/P, N) -> gather -> (K, N) - ker_g = te.distributed.gather_along_first_dim( - torch.transpose(kernel_t, 0, 1), tp_group - )[0].to(dtype=torch.float32) - # RS Input: (M, K/P) -> T -> (K/P, M) -> gather -> (K, M) -> T -> (M, K) - inp_g = torch.transpose( - te.distributed.gather_along_first_dim(torch.transpose(inp, 0, 1), tp_group)[0], 0, 1 - ).to(dtype=torch.float32) + # RS: Compute local GEMM on each rank, will apply reduce-scatter after + # RS Kernel: (N, K/P) - keep local + ker_g = kernel_t + # RS Input: (M, K/P) - keep local + inp_g = inp if opts.bulk_overlap: if opts.comm_type == tex.CommOverlapType.AG: @@ -456,10 +481,46 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None # Sum the list together for final global result ref_g = torch.stack(bulk_inp_list).sum(dim=0) else: - ref_g = torch.matmul(inp_g, ker_g) + # cuBLASMp always uses split-accumulator internally and does not expose a + # control to disable it, so force the reference to match when testing the + # cuBLASMp backend; otherwise honor the framework default. + ref_use_split_accumulator = True if opts.use_cublasmp else te.module.base._2X_ACC_FPROP + # For AG: ker_g=kernel_t=(K/P,N), inp_g=(M,N) -> ref_g=(M,K/P) local chunk + # For RS: ker_g=kernel_t=(N,K/P), inp_g=(M,K/P) -> ref_g=(M,N) partial + ref_g, *_ = tex.general_gemm( + ker_g, + inp_g, + out_dtype=torch.bfloat16, + use_split_accumulator=ref_use_split_accumulator, + ) + if opts.comm_type == tex.CommOverlapType.RS: + # Apply non-overlapped reduce-scatter to local reference GEMM output + # ref_g is currently (M, N) on each rank (partial result from local GEMM) + # All-gather to collect all (M, N) from all ranks + ref_rs_list = [torch.zeros_like(ref_g) for _ in range(tp_size)] + dist.all_gather(ref_rs_list, ref_g, group=tp_group) + # Stack and sum across ranks to get global result + ref_global = torch.stack(ref_rs_list, dim=0).sum(dim=0) # (M, N) + # Scatter: each rank keeps its portion (M/P, N) + start_idx = tp_rank * (outer_size // tp_size) + end_idx = (tp_rank + 1) * (outer_size // tp_size) + ref_g = ref_global[start_idx:end_idx, :] # (M/P, N) if ub_obj2 is not None: inp2_g = torch.nn.functional.gelu(ref_g) # pylint: disable=not-callable - ref2_g = torch.matmul(inp2_g, ker2_g) + # ker2_g=kernel2_t=(N,K/P), inp2_g=(M,K/P) -> partial (M,N) per rank + ref2_partial, *_ = tex.general_gemm( + ker2_g, + inp2_g, + out_dtype=torch.bfloat16, + use_split_accumulator=ref_use_split_accumulator, + ) + # Apply non-overlapped reduce-scatter to partial results + ref2_rs_list = [torch.zeros_like(ref2_partial) for _ in range(tp_size)] + dist.all_gather(ref2_rs_list, ref2_partial, group=tp_group) + ref2_global = torch.stack(ref2_rs_list, dim=0).sum(dim=0) # (M, N) fully reduced + start_idx = tp_rank * (outer_size // tp_size) + end_idx = (tp_rank + 1) * (outer_size // tp_size) + ref2_g = ref2_global[start_idx:end_idx, :] # (M/P, N) # Initialize quantizers with_quantized_compute = opts.quantization != "none" @@ -473,7 +534,7 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None if opts.quantization == "fp8": # Structure to maintain amax and scale/scale_inv information for the kernel and input num_gemms = 6 if ub_obj2 is not None else 3 - fp8_dtype = tex.DType.kFloat8E4M3 + fp8_dtype = te.DType.kFloat8E4M3 fp8_scales = torch.ones(num_gemms, dtype=torch.float, device="cuda") fp8_amaxes = torch.zeros(num_gemms, dtype=torch.float, device="cuda") @@ -516,7 +577,7 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None fp8_scales[5].clone(), fp8_amaxes[5].clone(), fp8_dtype ) elif opts.quantization == "mxfp8": - fp8_dtype = tex.DType.kFloat8E4M3 + fp8_dtype = te.DType.kFloat8E4M3 inp_quantizer = MXFP8Quantizer(fp8_dtype, columnwise=False) ker_quantizer = MXFP8Quantizer(fp8_dtype) if opts.bulk_overlap and opts.comm_type == tex.CommOverlapType.RS: @@ -580,7 +641,7 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None tp_group, ) gemm_inp = inp - else: + elif not opts.use_cublasmp: ag_out, _ = fill_userbuffers_buffer_for_all_gather( ub_obj, inp_fp8 if with_quantized_compute else inp, @@ -588,6 +649,8 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None tp_group, ) gemm_inp = ag_out + else: + gemm_inp = inp_fp8 if with_quantized_compute else inp if ub_obj2 is not None: rs_out2 = torch.empty( (outer_size // tp_size, hidden_size), dtype=torch.bfloat16, device="cuda" @@ -658,13 +721,13 @@ def _gemm(): # Trace the CUDA graph first g = torch.cuda.CUDAGraph() if with_quantized_compute: - if ub_obj is None: + if ub_obj2 is None: with torch.cuda.graph(g): all_outputs = _fp8_gemm() else: with torch.cuda.graph(g): all_outputs = _fp8_gemm() - _ = _fp8_gemm2(all_outputs[0]) + all_outputs2 = _fp8_gemm2(all_outputs[0]) else: with torch.cuda.graph(g): all_outputs = _gemm() @@ -682,7 +745,7 @@ def _gemm(): all_outputs = _fp8_gemm() end_events[i].record() if ub_obj2 is not None: - _fp8_gemm2(all_outputs[0]) + all_outputs2 = _fp8_gemm2(all_outputs[0]) else: start_events[i].record() all_outputs = _gemm() @@ -758,23 +821,26 @@ def _gemm(): else: if opts.comm_type == tex.CommOverlapType.AG: if ub_obj2 is not None: - # AG+RS Output: (M/P, N) -> gather -> (M, N) - output = rs_out2.to(dtype=torch.float32) - test_out = te.distributed.gather_along_first_dim(output, tp_group)[0] + # AG+RS Output: Keep local (M/P, N) for comparison with local reference + output = ( + rs_out2.to(dtype=torch.float32) + if not opts.use_cublasmp + else (all_outputs2[0].dequantize() if opts.fp8_output else all_outputs2[0]) + ) + test_out = output else: - # AG Output: (M, K/P) -> T -> (K/P, M) -> gather -> (K, M) -> T -> (M, K) + # AG Output: Keep local (M, K/P) for comparison with local reference output = all_outputs[0].dequantize() if opts.fp8_output else all_outputs[0] - test_out = torch.transpose( - te.distributed.gather_along_first_dim( - torch.transpose(output, 0, 1), tp_group - )[0], - 0, - 1, - ) + test_out = output else: - # RS Output: (M/P, N) -> gather -> (M, N) - output = rs_out.to(dtype=torch.float32) - test_out = te.distributed.gather_along_first_dim(output, tp_group)[0] + # RS Output: Keep local (M/P, N) for comparison with local reference + output = ( + rs_out.to(dtype=torch.float32) + if not opts.use_cublasmp + else (all_outputs[0].dequantize() if opts.fp8_output else all_outputs[0]) + ) + # Don't gather - keep local for comparison with local reduce-scattered reference + test_out = output ref_out = ref2_g if ub_obj2 is not None else ref_g test_nonzeros = torch.count_nonzero(test_out) @@ -782,7 +848,9 @@ def _gemm(): nonzero_info = ( f"output nonzeros = {test_nonzeros} " + f"| reference count = {ref_nonzeros}" ) - dist_print(nonzero_info, src=0, section=True, group=tp_group) + + # Both AG and RS now compare local outputs across all ranks + dist_print(nonzero_info, section=True, group=tp_group) sizing_info = ( f"input: {list(inp.shape)} " + f"| GEMM1 weights: {list(kernel_t.shape)[::-1]} " @@ -792,15 +860,17 @@ def _gemm(): sizing_info += f"| output: {list(output.shape)}\n" dist_print(sizing_info, section=True, group=tp_group) + # Both AG and RS now compare local outputs; print per-rank sizing_info_g = ( - f"input: {list(inp_g.shape)} " + f"| GEMM1 weights: {list(ker_g.shape)} " + f"input: {list(inp.shape)} | GEMM1 weights: {list(kernel_t.shape)[::-1]} " ) if ub_obj2 is not None: - sizing_info_g += f"| GEMM2 weights: {list(ker2_g.shape)} " + sizing_info_g += f"| GEMM2 weights: {list(kernel2_t.shape)[::-1]} " sizing_info_g += ( - f"| output: {list(test_out.shape)} " + f"| reference: {list(ref_out.shape)}\n" + f"| output (local): {list(test_out.shape)} | reference (local):" + f" {list(ref_out.shape)}\n" ) - dist_print(sizing_info_g, src=0, group=tp_group) + dist_print(sizing_info_g, group=tp_group) torch.cuda.synchronize() dist.barrier(tp_group) @@ -808,8 +878,14 @@ def _gemm(): m = torch.argmax(diff) abs_err = diff[m].item() rel_err = abs_err / max(abs(ref_out.flatten()[m].item()), 1e-5) - rtol = 0.02 if opts.quantization == "none" else 0.125 - atol = 0.001 if opts.quantization == "none" else 0.0625 + rtol = ( + opts.rtol if opts.rtol is not None else (0.02 if opts.quantization == "none" else 0.125) + ) + atol = ( + opts.atol + if opts.atol is not None + else (0.002 if opts.quantization == "none" else 0.0625) + ) if rel_err > rtol and abs_err > atol: numerics_failed = True numerics_info = ( @@ -828,9 +904,7 @@ def _gemm(): if abs_err <= atol: numerics_info += f"abs. error = {abs_err} (tol = {atol})" - dist_print( - numerics_info, src=0, section=True, info=True, error=numerics_failed, group=tp_group - ) + dist_print(numerics_info, section=True, info=True, error=numerics_failed, group=tp_group) dist.barrier(tp_group) if LOCAL_RANK == 0: diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 53c7a5e7cc4..46795415e51 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -23,6 +23,7 @@ DelayedScaling, Float8CurrentScaling, Format, + MMParams, MXFP8BlockScaling, ) @@ -258,6 +259,30 @@ def _parse_args(argv=None, namespace=None): default=0, help="Number of layers at the end to run in bf16.", ) + parser.add_argument( + "--use-cublasmp", + action="store_true", + default=False, + help="Use cuBLASMp backend.", + ) + parser.add_argument( + "--rtol", + type=float, + default=None, + help=( + "Override the relative-error tolerance used in the numerical check. " + "When unset, defaults to 0.125 for FP8 and 0.025 otherwise." + ), + ) + parser.add_argument( + "--atol", + type=float, + default=None, + help=( + "Override the absolute-error tolerance used in the numerical check. " + "When unset, defaults to 0.0625 for FP8 and 0.00125 otherwise." + ), + ) args = parser.parse_args(argv, namespace) if args.use_cuda_graphs and args.layer_type in [te.MultiheadAttention, te.TransformerLayer]: @@ -436,6 +461,7 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False): dtype=torch.bfloat16, bootstrap_backend=opts.bootstrap_backend, ub_cfgs=ub_cfgs if opts.ub_cfg is None else opts.ub_cfg, + with_cublasmp=opts.use_cublasmp, ) with te.quantized_model_init(enabled=opts.fp8_init): @@ -471,6 +497,11 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False): elif opts.quantization == "mxfp8": fp8_recipe = MXFP8BlockScaling() + if opts.fp8: + fp8_recipe.fp8_gemm_fprop = MMParams(use_split_accumulator=True) + fp8_recipe.fp8_gemm_dgrad = MMParams(use_split_accumulator=True) + fp8_recipe.fp8_gemm_wgrad = MMParams(use_split_accumulator=True) + layer_contexts = [ ( partial( @@ -552,8 +583,8 @@ def run_fwd_bwd(model, x): # Now validate accuracy if not bool(numerics_failed.item()): for i, (test_g, ref_g) in enumerate(zip(test_grads, ref_grads)): - rtol = 0.125 if opts.fp8 else 0.025 - atol = 0.0625 if opts.fp8 else 0.00125 + rtol = opts.rtol if opts.rtol is not None else (0.125 if opts.fp8 else 0.025) + atol = opts.atol if opts.atol is not None else (0.0625 if opts.fp8 else 0.00125) grad_failed, grad_info = _compare_tensors(names[i], test_g, ref_g, rtol, atol) dist_print(grad_info, src=WORLD_RANK, error=grad_failed) numerics_failed[0] = int(grad_failed) diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index 8e24e636e89..fe02f990b44 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -15,7 +15,6 @@ import torch from torch import nn import torch.distributed as dist -import transformer_engine_torch as tex from transformer_engine.common.recipe import ( MXFP8BlockScaling, DelayedScaling, @@ -399,7 +398,7 @@ def _test_quantizer(input_dtype, fp8_dtype): Args: input_dtype (torch.dtype): The data type of the input. - fp8_dtype (tex.DType): The data type of the fp8. + fp8_dtype (te.DType): The data type of the fp8. """ M, N = WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE @@ -443,7 +442,7 @@ def test_quantizer(): return input_dtypes = [torch.float32, torch.bfloat16] - fp8_dtypes = [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2] + fp8_dtypes = [te.DType.kFloat8E4M3, te.DType.kFloat8E5M2] for input_dtype in input_dtypes: for fp8_dtype in fp8_dtypes: @@ -514,7 +513,7 @@ def _test_quantized_all_gather(input_dtype, low_precision_dtype, quantizer_cls): Args: input_dtype (torch.dtype): The data type of the input. - low_precision_dtype (tex.DType): The data type of the low precision, can be fp4 or fp8. + low_precision_dtype (te.DType): The data type of the low precision, can be fp4 or fp8. """ M, N = WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE // 2 @@ -623,8 +622,8 @@ def test_quantized_all_gather(): return input_dtypes = [torch.bfloat16] - fp4_dtype = [tex.DType.kFloat4E2M1] - fp8_dtype = [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2] + fp4_dtype = [te.DType.kFloat4E2M1] + fp8_dtype = [te.DType.kFloat8E4M3, te.DType.kFloat8E5M2] quantizer_cls_nvfp4 = [NVFP4Quantizer] # add FP8 quantizers if needed quantizer_cls_fp8 = [] diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 7a81f93bd64..6b1ad870e91 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -50,7 +50,9 @@ torch._dynamo.reset() -def _run_gemm_with_overlap(comm_type, bulk, p2p, atomic, aggregate, quantization): +def _run_gemm_with_overlap( + comm_type, bulk, p2p, atomic, aggregate, quantization, use_cublasmp=False +): test_path = TEST_ROOT / "run_gemm_with_overlap.py" test_cmd = LAUNCH_CMD + [ str(test_path), @@ -79,6 +81,18 @@ def _run_gemm_with_overlap(comm_type, bulk, p2p, atomic, aggregate, quantization test_cmd.append("--atomic") if aggregate: test_cmd.append("--aggregate") + if use_cublasmp: + if not tex.nvte_built_with_cublasmp(): + pytest.skip("Transformer Engine not built with cuBLASMp (NVTE_WITH_CUBLASMP=0).") + if quantization == "mxfp8": + pytest.skip( + "cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling)." + ) + if comm_type == "RS" and not p2p and not tex.device_supports_multicast(): + pytest.skip( + "cuBLASMp non-P2P reduce-scatter requires NVSwitch (multicast support)." + ) + test_cmd.append("--use-cublasmp") result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) if ( @@ -90,7 +104,13 @@ def _run_gemm_with_overlap(comm_type, bulk, p2p, atomic, aggregate, quantization def _run_layer_with_overlap( - layer_type, linear_parallel_mode, overlap_rs_dgrad, fp8, quantization, num_layers=1 + layer_type, + linear_parallel_mode, + overlap_rs_dgrad, + fp8, + quantization, + num_layers=1, + use_cublasmp=False, ): test_path = TEST_ROOT / "run_layer_with_overlap.py" test_cmd = LAUNCH_CMD + [ @@ -117,6 +137,13 @@ def _run_layer_with_overlap( test_cmd.append("--fp8") test_cmd.append(f"--quantization={quantization}") + if use_cublasmp: + if not tex.nvte_built_with_cublasmp(): + pytest.skip("Transformer Engine not built with cuBLASMp (NVTE_WITH_CUBLASMP=0).") + if fp8 and quantization == "mxfp8": + pytest.skip("cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling).") + test_cmd.append("--use-cublasmp") + os.environ["PYTORCH_JIT"] = "0" os.environ["NVTE_TORCH_COMPILE"] = "0" os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" @@ -125,6 +152,10 @@ def _run_layer_with_overlap( # backward when running with Userbuffers on A100s. This does # not show up in more recent GPUs. os.environ["NVTE_FLASH_ATTN"] = "0" + elif fp8: + # Fused attention is causing non-deterministic FP8 failures on H100s even with + # NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, so disable it entirely for this test. + os.environ["NVTE_FUSED_ATTN"] = "0" result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) @@ -132,6 +163,7 @@ def _run_layer_with_overlap( os.unsetenv("NVTE_TORCH_COMPILE") os.unsetenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO") os.unsetenv("NVTE_FLASH_ATTN") + os.unsetenv("NVTE_FUSED_ATTN") if ( result.returncode != 0 @@ -141,24 +173,26 @@ def _run_layer_with_overlap( raise AssertionError(result.stderr.decode()) +@pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize("quantization", ("none", "fp8", "mxfp8")) @pytest.mark.parametrize("aggregate", (False, True)) -def test_split_all_gather_overlaps(quantization, aggregate): +def test_split_all_gather_overlaps(quantization, aggregate, use_cublasmp): """ Test (split GEMM -> all-gather) overlaps with direct calls to te.cpp_extensions.gemm or te.cpp_extensions.fp8_gemm. """ - _run_gemm_with_overlap("AG", False, True, False, aggregate, quantization) + _run_gemm_with_overlap("AG", False, True, False, aggregate, quantization, use_cublasmp) +@pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize("quantization", ("none", "fp8", "mxfp8")) @pytest.mark.parametrize("p2p", (False, True)) -def test_split_reduce_scatter_overlaps(quantization, p2p): +def test_split_reduce_scatter_overlaps(quantization, p2p, use_cublasmp): """ Test (reduce-scatter -> split GEMM) overlaps with direct calls to te.cpp_extensions.gemm or te.cpp_extensions.fp8_gemm. """ - _run_gemm_with_overlap("RS", False, p2p, False, False, quantization) + _run_gemm_with_overlap("RS", False, p2p, False, False, quantization, use_cublasmp) @pytest.mark.parametrize( @@ -197,6 +231,7 @@ def test_bulk_overlaps(comm_type, quantization, connections): _run_gemm_with_overlap(comm_type, True, False, False, False, quantization) +@pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "fp8", (False,), @@ -235,13 +270,18 @@ def test_bulk_overlaps(comm_type, quantization, connections): ) ], ) -def test_layers_with_overlap_bf16(layer_type, linear_parallel_mode, overlap_rs_dgrad, fp8): +def test_layers_with_overlap_bf16( + layer_type, linear_parallel_mode, overlap_rs_dgrad, use_cublasmp, fp8 +): """ Test Transformer Engine layers with comm+GEMM overlap. """ - _run_layer_with_overlap(layer_type, linear_parallel_mode, overlap_rs_dgrad, fp8, None) + _run_layer_with_overlap( + layer_type, linear_parallel_mode, overlap_rs_dgrad, fp8, None, use_cublasmp=use_cublasmp + ) +@pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "quantization", ["fp8_delayed_scaling", "fp8_current_scaling", "mxfp8"], @@ -282,13 +322,22 @@ def test_layers_with_overlap_fp8( linear_parallel_mode, overlap_rs_dgrad, quantization, + use_cublasmp, ): """ Test Transformer Engine layers with comm+GEMM overlap. """ - _run_layer_with_overlap(layer_type, linear_parallel_mode, overlap_rs_dgrad, True, quantization) + _run_layer_with_overlap( + layer_type, + linear_parallel_mode, + overlap_rs_dgrad, + True, + quantization, + use_cublasmp=use_cublasmp, + ) +@pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "fp8", (False,), @@ -321,16 +370,23 @@ def test_layers_with_overlap_fp8( ], ) def test_multi_layer_with_overlap_bf16( - layer_type, linear_parallel_mode, overlap_rs_dgrad, fp8, num_layers + layer_type, linear_parallel_mode, overlap_rs_dgrad, fp8, num_layers, use_cublasmp ): """ Test Transformer Engine layers with comm+GEMM overlap. """ _run_layer_with_overlap( - layer_type, linear_parallel_mode, overlap_rs_dgrad, fp8, None, num_layers + layer_type, + linear_parallel_mode, + overlap_rs_dgrad, + fp8, + None, + num_layers, + use_cublasmp=use_cublasmp, ) +@pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "quantization", ["fp8_delayed_scaling", "fp8_current_scaling", "mxfp8"], @@ -357,11 +413,17 @@ def test_multi_layer_with_overlap_bf16( ], ) def test_multi_layer_with_overlap_fp8( - layer_type, linear_parallel_mode, overlap_rs_dgrad, quantization, num_layers + layer_type, linear_parallel_mode, overlap_rs_dgrad, quantization, num_layers, use_cublasmp ): """ Test Transformer Engine layers with comm+GEMM overlap. """ _run_layer_with_overlap( - layer_type, linear_parallel_mode, overlap_rs_dgrad, True, quantization, num_layers + layer_type, + linear_parallel_mode, + overlap_rs_dgrad, + True, + quantization, + num_layers, + use_cublasmp=use_cublasmp, ) diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py index c4840389382..e986f7d3895 100644 --- a/tests/pytorch/distributed/test_fusible_ops.py +++ b/tests/pytorch/distributed/test_fusible_ops.py @@ -29,7 +29,6 @@ is_bf16_available, ) import transformer_engine.pytorch.ops as te_ops -import transformer_engine_torch as tex # Import utility functions _current_file = pathlib.Path(__file__).resolve() @@ -107,17 +106,17 @@ def make_reference_and_test_tensors( quantizer = Float8Quantizer( scale=torch.ones(1, dtype=torch.float32, device=test_device).squeeze(), amax=torch.zeros(1, dtype=torch.float32, device=test_device), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) test = quantizer(test) elif quantization == "fp8_current_scaling": quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, device=test_device, ) test = quantizer(test) elif quantization == "mxfp8": - test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) + test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) elif quantization == "nvfp4": test = NVFP4Quantizer( with_rht=False, diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index 3dcefd46fd4..a838656581c 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -156,17 +156,17 @@ def make_reference_and_test_tensors( quantizer = Float8Quantizer( scale=torch.ones(1, dtype=torch.float32, device=test_device).squeeze(), amax=torch.zeros(1, dtype=torch.float32, device=test_device), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) test = quantizer(test) elif quantization == "fp8_current_scaling": quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, device=test_device, ) test = quantizer(test) elif quantization == "mxfp8": - test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) + test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) else: raise ValueError(f"Unsupported quantization scheme ({quantization})") if isinstance(test, QuantizedTensor) and not test_is_quantized: @@ -372,8 +372,15 @@ def _test_linear( tols = dtype_tols( model[0].weight._fp8_dtype if isinstance(model[0].weight, Float8Tensor) - else tex.DType.kFloat8E4M3 + else te.DType.kFloat8E4M3 ) + if te.module.base.using_cublasmp_backend() and not quantized_compute: + # cuBLASMp's GEMM+RS kernel runs a slightly different GEMM algo than Userbuffers + # (e.g. split-accumulator is always enabled) so it very narrowly violates the default + # bf16 rtol. This is not a regression, just a quirk of how the algorithms line up at the + # precision floor. So we relax rtol only (atol stays same) to allow for this without + # masking real regressions. + tols = {**tols, "rtol": max(tols.get("rtol", 0.0), 3.0e-2)} # Check results y_test = y_test.to(dtype=torch.float64, device="cpu") @@ -477,6 +484,7 @@ def main() -> None: parser.add_argument("--head-dim", type=int, default=256) parser.add_argument("--dtype", type=str, default="bfloat16") parser.add_argument("--quantization", type=str, default=None) + parser.add_argument("--use-cublasmp", action="store_true") args = parser.parse_args() # Run parallel tests if needed @@ -517,6 +525,7 @@ def main() -> None: dtype=model_config.dtype, bootstrap_backend=bootstrap_backend, ub_cfgs=userbuffer_configs, + with_cublasmp=args.use_cublasmp, ) # Run tests diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py index c2f8e8de127..d07953ce370 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py +++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py @@ -139,7 +139,7 @@ def check_grouped_tensor_mxfp8_versus_reference( optimize_for_gemm: bool = False, ) -> None: - te_dtype = tex.DType.kFloat8E4M3 + te_dtype = te.DType.kFloat8E4M3 split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda") @@ -236,7 +236,7 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( optimize_for_gemm: bool = False, ) -> None: - te_dtype = tex.DType.kFloat8E4M3 + te_dtype = te.DType.kFloat8E4M3 assert valid_M is not None, "valid_M must be provided when with_paged_stashing is True" assert valid_M < M, "valid_M must be less than M when with_paged_stashing is True" diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py index 6f0700809b3..127b4876504 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py +++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py @@ -4,7 +4,6 @@ import transformer_engine.pytorch as te -import transformer_engine_torch as tex from transformer_engine.pytorch import MXFP8Quantizer from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -43,7 +42,7 @@ def check_mxfp8_quantize_swizzle_fusion( return_transpose: bool, ) -> None: - te_dtype = tex.DType.kFloat8E4M3 + te_dtype = te.DType.kFloat8E4M3 # Setup device and random seed device = "cuda" diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index bd4d0297296..63e69beafc1 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -34,7 +34,7 @@ def check_nvfp4_gemm_versus_reference( ): if nvfp4_e4m3_max != 448 and not use_4over6: pytest.skip("E4M3 max 256 is only meaningful for 4over6") - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 # Setup device and random seed device = "cuda" @@ -252,7 +252,7 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( use_4over6: bool = False, nvfp4_4over6_err_mode: str = "MAE", ): - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 device = "cuda" torch.manual_seed(23) torch.cuda.manual_seed(23) @@ -347,7 +347,7 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( use_4over6: bool = False, nvfp4_4over6_err_mode: str = "MAE", ): - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 device = "cuda" torch.manual_seed(37) torch.cuda.manual_seed(37) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 20a91bf6fe5..d1c15954893 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -15,7 +15,6 @@ from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils -from transformer_engine.pytorch.constants import TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling import pytest @@ -45,7 +44,7 @@ def check_group_quantization_nvfp4_versus_reference( with_random_sign_mask: bool = True, ) -> None: - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 # Setup device and random seed device = "cuda" diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py index d46a874695f..b78c66befdf 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -8,7 +8,6 @@ from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils -from transformer_engine.pytorch.constants import TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor @@ -55,7 +54,7 @@ def check_grouped_tensor_nvfp4_versus_reference( optimize_for_gemm: bool = False, ) -> None: - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda") @@ -172,7 +171,7 @@ def check_grouped_tensor_nvfp4_with_paged_stashing( optimize_for_gemm: bool = False, ) -> None: - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 assert valid_M is not None, "valid_M must be provided when with_paged_stashing is True" assert valid_M < M, "valid_M must be less than M when with_paged_stashing is True" diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 5bb92f70dc1..12366b731e0 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -10,7 +10,6 @@ from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.common.recipe import NVFP4BlockScaling -from transformer_engine.pytorch.constants import TE_DType recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -62,7 +61,7 @@ def check_quantization_nvfp4_versus_reference( row_scaled_nvfp4, return_transpose, with_2d_quantization, use_4over6, x_dtype, M, N ) - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 # Setup device and random seed device = "cuda" @@ -259,7 +258,7 @@ def test_nvfp4_quantization_extrema_versus_reference( row_scaled_nvfp4, return_transpose, use_4over6=use_4over6 ) - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 device = "cuda" seed = 0 @@ -380,7 +379,7 @@ def test_nvfp4_quantization_boundary_values( row_scaled_nvfp4, return_transpose, use_4over6=use_4over6 ) - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 device = "cuda" seed = 123 @@ -505,7 +504,7 @@ def test_nvfp4_quantization_noncontiguous_inputs( row_scaled_nvfp4, return_transpose, use_4over6=use_4over6 ) - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 device = "cuda" seed = 17 diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index 2d159dbf6ab..a65010ef025 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -14,7 +14,6 @@ from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils -from transformer_engine.pytorch.constants import TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling import pytest @@ -46,7 +45,7 @@ def check_quantization_nvfp4_versus_reference( ) -> None: assert with_rht and with_post_rht_amax, "RHT and post-RHT amax reduction must be enabled." - te_dtype = tex.DType.kFloat4E2M1 + te_dtype = te.DType.kFloat4E2M1 # Setup device and random seed device = "cuda" diff --git a/tests/pytorch/references/ref_per_tensor_cs.py b/tests/pytorch/references/ref_per_tensor_cs.py index c4a6d73d707..1eebbd13e78 100644 --- a/tests/pytorch/references/ref_per_tensor_cs.py +++ b/tests/pytorch/references/ref_per_tensor_cs.py @@ -3,9 +3,9 @@ # See LICENSE for license information. import torch -import transformer_engine_torch as tex from transformer_engine.pytorch.constants import TE_DType_To_Torch +from transformer_engine.pytorch import DType from references.quantize_scale_calc import scale_from_amax_tensor @@ -40,7 +40,7 @@ def _multi_dim_transpose(tensor): # current scaling reference quantization def ref_per_tensor_cs_cast( tensor: torch.Tensor, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: DType = DType.kFloat8E4M3, return_transpose: bool = False, force_pow_2_scales: bool = False, amax_epsilon: float = 0.0, diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 35cc98a976f..e4e36da38a7 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -20,7 +20,6 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe from utils import ModelConfig, recipe_id, skip_unsupported_backward_override -import transformer_engine_torch as tex # Check supported quantization schemes fp8_available, _ = FP8GlobalStateManager.is_fp8_available() @@ -186,23 +185,23 @@ def create_tensor(recipe: Optional[recipe.Recipe], requires_grad: bool = False) return tensor elif recipe.delayed(): quantizer = te.tensor.float8_tensor.Float8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, scale=torch.tensor([1.0], device="cuda"), amax=torch.tensor([1.0], device="cuda"), ) return quantizer(tensor) elif recipe.float8_current_scaling(): quantizer = te.tensor.float8_tensor.Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, device="cuda" + fp8_dtype=te.DType.kFloat8E4M3, device="cuda" ) return quantizer(tensor) elif recipe.float8_block_scaling(): quantizer = te.tensor.float8_blockwise_tensor.Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=True ) return quantizer(tensor) elif recipe.mxfp8(): - quantizer = te.tensor.mxfp8_tensor.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = te.tensor.mxfp8_tensor.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) return quantizer(tensor) elif recipe.nvfp4(): qparams = recipe.fp4_quant_fwd_inp diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 62a6291797c..3e6fdb816b1 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -100,10 +100,10 @@ def test_custom_recipe_sanity(module_type): # Single factory: map roles to quantizers def quantizer_factory(role): if role is None: - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") if role.tensor_type == "grad_output": - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") custom_recipe = recipe.CustomRecipe(qfactory=quantizer_factory) @@ -137,10 +137,10 @@ def test_custom_recipe_grouped_linear_sanity(): def quantizer_factory(role): if role is None: - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") if role.tensor_type == "grad_output": - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") custom_recipe = recipe.CustomRecipe(qfactory=quantizer_factory) @@ -183,11 +183,11 @@ def test_custom_recipe_matches_current_scaling(): ref_fwd_out = model_ref.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] ref_bwd_go = model_ref.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] ref_bwd_gi = model_ref.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] - assert ref_fwd_in.dtype == tex.DType.kFloat8E4M3 - assert ref_fwd_w.dtype == tex.DType.kFloat8E4M3 - assert ref_fwd_out.dtype == tex.DType.kFloat8E4M3 - assert ref_bwd_go.dtype == tex.DType.kFloat8E5M2 - assert ref_bwd_gi.dtype == tex.DType.kFloat8E5M2 + assert ref_fwd_in.dtype == te.DType.kFloat8E4M3 + assert ref_fwd_w.dtype == te.DType.kFloat8E4M3 + assert ref_fwd_out.dtype == te.DType.kFloat8E4M3 + assert ref_bwd_go.dtype == te.DType.kFloat8E5M2 + assert ref_bwd_gi.dtype == te.DType.kFloat8E5M2 # Stress dynamic range in grad_output scale = torch.ones(out_features, device="cuda", dtype=torch.float32) @@ -199,10 +199,10 @@ def test_custom_recipe_matches_current_scaling(): # Custom: single factory returning quantizers per role to match Float8CurrentScaling def quantizer_factory(role): if role is None: - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") if role.tensor_type == "grad_output": - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") custom_recipe = recipe.CustomRecipe(qfactory=quantizer_factory) @@ -214,11 +214,11 @@ def quantizer_factory(role): cus_fwd_out = model_custom.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] cus_bwd_go = model_custom.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] cus_bwd_gi = model_custom.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] - assert cus_fwd_in.dtype == tex.DType.kFloat8E4M3 - assert cus_fwd_w.dtype == tex.DType.kFloat8E4M3 - assert cus_fwd_out.dtype == tex.DType.kFloat8E4M3 - assert cus_bwd_go.dtype == tex.DType.kFloat8E5M2 - assert cus_bwd_gi.dtype == tex.DType.kFloat8E4M3 # role=None fallback + assert cus_fwd_in.dtype == te.DType.kFloat8E4M3 + assert cus_fwd_w.dtype == te.DType.kFloat8E4M3 + assert cus_fwd_out.dtype == te.DType.kFloat8E4M3 + assert cus_bwd_go.dtype == te.DType.kFloat8E5M2 + assert cus_bwd_gi.dtype == te.DType.kFloat8E4M3 # role=None fallback loss_custom = (out_custom.float() * scale.view(1, -1)).sum() loss_custom.backward() @@ -256,10 +256,10 @@ def test_custom_recipe_ops_linear_2_1_layout(): def quantizer_factory(role): if role is None: - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") if role.tensor_type == "grad_output": - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") custom = recipe.CustomRecipe(qfactory=quantizer_factory) @@ -300,14 +300,14 @@ def test_custom_recipe_factory_invocation_counts_and_cycling(): def quantizer_factory(role): if role is None: counts[None] += 1 - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device=torch.device("cuda")) + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device=torch.device("cuda")) assert isinstance(role, QuantizerRole), f"Expected QuantizerRole, got {type(role)}" assert role.module_type == "linear" if role.tensor_type in counts: counts[role.tensor_type] += 1 if role.tensor_type == "grad_output": - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device=torch.device("cuda")) - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device=torch.device("cuda")) + return Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device=torch.device("cuda")) + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device=torch.device("cuda")) custom = recipe.CustomRecipe(qfactory=quantizer_factory) @@ -336,7 +336,7 @@ def test_factories_return_distinct_instances_and_buffers(): def factory(): scale = torch.ones(1, dtype=torch.float32, device="cuda") amax = torch.zeros(1, dtype=torch.float32, device="cuda") - return Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) + return Float8Quantizer(scale=scale, amax=amax, fp8_dtype=te.DType.kFloat8E4M3) q1 = factory() q2 = factory() @@ -723,7 +723,7 @@ def test_grouped_linear_module_type_dispatch(): def recording_factory(role): recorded_roles.append(role) - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") custom_recipe = recipe.CustomRecipe(qfactory=recording_factory) @@ -825,7 +825,7 @@ def mixed_factory(role): # Only weight gets delayed scaling, rest get current scaling if role is not None and role.tensor_type == "weight": return DelayedScalingRequest(fp8_format=Format.HYBRID) - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") custom_recipe = recipe.CustomRecipe(qfactory=mixed_factory) diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index 3b964a5af99..cac7ccb66bf 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -761,7 +761,7 @@ class TestFP8CurrentScalingNativeVsRef: def _make_quantizers(rowwise=True, columnwise=True): # TE native FP8 current scaling quantizer te_quant = te.Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, device=torch.device("cuda"), rowwise=rowwise, columnwise=columnwise, diff --git a/tests/pytorch/test_float8blockwisetensor.py b/tests/pytorch/test_float8blockwisetensor.py index 7add4ee5abf..1b5e17f4635 100644 --- a/tests/pytorch/test_float8blockwisetensor.py +++ b/tests/pytorch/test_float8blockwisetensor.py @@ -17,16 +17,17 @@ get_device_compute_capability, ) import transformer_engine_torch as tex +from transformer_engine.pytorch import DType # PyTorch tensor dtypes _dtypes: List[torch.dtype] = [torch.float32, torch.float16, torch.bfloat16] # TE FP8 dtypes -_fp8_dtypes: List[tex.DType] = [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2] +_fp8_dtypes: List[DType] = [DType.kFloat8E4M3, DType.kFloat8E5M2] # Numerical tolerances with FP8 types -_tols: Dict[tex.DType, Dict[str, float]] = { - tex.DType.kFloat8E4M3: dict(rtol=0.125, atol=0.08), - tex.DType.kFloat8E5M2: dict(rtol=0.25, atol=0.125), +_tols: Dict[DType, Dict[str, float]] = { + DType.kFloat8E4M3: dict(rtol=0.125, atol=0.08), + DType.kFloat8E5M2: dict(rtol=0.25, atol=0.125), } @@ -59,7 +60,7 @@ def setup_class(cls) -> None: def test_constructor( self, dims: DimsType = 1, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: DType = DType.kFloat8E4M3, dtype: torch.dtype = torch.float32, is_2D_scaled: bool = True, ) -> None: @@ -140,7 +141,7 @@ def _test_quantize_dequantize( @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) def test_quantize_dequantize_dtypes( - self, fp8_dtype: tex.DType, dtype: torch.dtype, block_scaling_dim: int + self, fp8_dtype: DType, dtype: torch.dtype, block_scaling_dim: int ) -> None: atol = _tols[fp8_dtype]["atol"] rtol = _tols[fp8_dtype]["rtol"] @@ -156,7 +157,7 @@ def test_quantize_dequantize_dtypes( @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("block_scaling_dim", [1]) def test_quantize_dequantize_columnwise_only( - self, fp8_dtype: tex.DType, dtype: torch.dtype, block_scaling_dim: int + self, fp8_dtype: DType, dtype: torch.dtype, block_scaling_dim: int ) -> None: atol = _tols[fp8_dtype]["atol"] rtol = _tols[fp8_dtype]["rtol"] @@ -181,10 +182,10 @@ def test_quantize_dequantize_dims( block_scaling_dim: int, dq_columnwise: bool, ) -> None: - atol = _tols[tex.DType.kFloat8E4M3]["atol"] - rtol = _tols[tex.DType.kFloat8E4M3]["rtol"] + atol = _tols[DType.kFloat8E4M3]["atol"] + rtol = _tols[DType.kFloat8E4M3]["rtol"] quantizer = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=dq_columnwise, block_scaling_dim=block_scaling_dim, @@ -206,10 +207,10 @@ def test_quantize_dequantize_dims( def test_quantize_dequantize_compact_format( self, dims: DimsType, block_scaling_dim: int, dq_columnwise: bool ) -> None: - atol = _tols[tex.DType.kFloat8E4M3]["atol"] - rtol = _tols[tex.DType.kFloat8E4M3]["rtol"] + atol = _tols[DType.kFloat8E4M3]["atol"] + rtol = _tols[DType.kFloat8E4M3]["rtol"] quantizer = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=dq_columnwise, block_scaling_dim=block_scaling_dim, @@ -229,7 +230,11 @@ def test_quantize_dequantize_compact_format( @pytest.mark.parametrize("fp8_dtype", _fp8_dtypes) @pytest.mark.parametrize("dq_columnwise", [True, False]) def test_quantize_dequantize_dims_cpp_allocate_output( - self, dims: DimsType, block_scaling_dim: int, fp8_dtype: tex.DType, dq_columnwise: bool + self, + dims: DimsType, + block_scaling_dim: int, + fp8_dtype: DType, + dq_columnwise: bool, ) -> None: atol = _tols[fp8_dtype]["atol"] rtol = _tols[fp8_dtype]["rtol"] @@ -257,7 +262,7 @@ def test_data_accessors(self, dims: DimsType, block_scaling_dim: int) -> None: x_hp = torch.rand(_to_list(dims), dtype=dtype, device=device) y_hp = torch.rand(_to_list(dims), dtype=dtype, device=device) - fp8_dtype = tex.DType.kFloat8E4M3 + fp8_dtype = DType.kFloat8E4M3 quantizer = Float8BlockQuantizer( fp8_dtype=fp8_dtype, rowwise=True, @@ -283,7 +288,7 @@ def test_serialization(self, dims: DimsType, block_scaling_dim: int) -> None: dtype = torch.bfloat16 x_hp = torch.rand(_to_list(dims), dtype=dtype, device=device) quantizer = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E5M2, + fp8_dtype=DType.kFloat8E5M2, rowwise=True, columnwise=True, block_scaling_dim=block_scaling_dim, @@ -316,12 +321,12 @@ def test_serialization(self, dims: DimsType, block_scaling_dim: int) -> None: x_fp8_loaded_dequant = x_fp8_loaded.dequantize() torch.testing.assert_close(x_fp8_loaded_dequant, x_fp8_dequant) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3], ids=str) + @pytest.mark.parametrize("fp8_dtype", [DType.kFloat8E4M3], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("dims", [[256, 512], [250, 500]]) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) def test_inplace_ops( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int + self, fp8_dtype: DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int ) -> None: """Test in-place operations""" device = "cuda" @@ -353,12 +358,12 @@ def test_inplace_ops( x_fp8.mul_(y_fp8) torch.testing.assert_close(x_fp8.dequantize(), x_hp * y_hp, **_tols[fp8_dtype]) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3], ids=str) + @pytest.mark.parametrize("fp8_dtype", [DType.kFloat8E4M3], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("dims", [[256, 512], [250, 500]]) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) def test_out_of_place_ops( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int + self, fp8_dtype: DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int ) -> None: """Test out-of-place operations""" device = "cuda" @@ -389,12 +394,12 @@ def test_out_of_place_ops( with pytest.raises(AssertionError): torch.testing.assert_close(x_fp8 + y_fp8, x_hp - y_hp, **_tols[fp8_dtype]) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3], ids=str) + @pytest.mark.parametrize("fp8_dtype", [DType.kFloat8E4M3], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("dims", [[256, 512], [250, 500]]) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) def test_view_same_shape( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int + self, fp8_dtype: DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int ) -> None: """Test view operations that preserve tensor shape""" device = "cuda" @@ -419,13 +424,13 @@ def test_view_same_shape( with pytest.raises(AssertionError): torch.testing.assert_close(x_view.dequantize(), -x_hp, **_tols[fp8_dtype]) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2], ids=str) + @pytest.mark.parametrize("fp8_dtype", [DType.kFloat8E4M3, DType.kFloat8E5M2], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32], ids=str) @pytest.mark.parametrize( "dims", [[16, 16, 512], [16, 16, 512, 16], [12, 7, 11], [13, 14, 16], [2, 3, 5]] ) def test_view_and_reshape_1D( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: List[int] + self, fp8_dtype: DType, dtype: torch.dtype, dims: List[int] ) -> None: """Test view operations that preserve tensor shape""" device = "cuda" @@ -472,11 +477,11 @@ def is_bitwise_equal(a, b): assert is_bitwise_equal(x_fp8_reshape._rowwise_data, x_fp8._rowwise_data) assert is_bitwise_equal(x_fp8_reshape._rowwise_scale_inv, x_fp8._rowwise_scale_inv) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2], ids=str) + @pytest.mark.parametrize("fp8_dtype", [DType.kFloat8E4M3, DType.kFloat8E5M2], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32], ids=str) @pytest.mark.parametrize("dims", [[16, 16, 512, 16], [2, 512, 512, 128], [3, 13, 14, 16]]) def test_view_and_reshape_2D( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: List[int] + self, fp8_dtype: DType, dtype: torch.dtype, dims: List[int] ) -> None: """Test view operations that preserve tensor shape""" device = "cuda" @@ -523,12 +528,12 @@ def is_bitwise_equal(a, b): assert is_bitwise_equal(x_fp8_reshape._rowwise_data, x_fp8._rowwise_data) assert is_bitwise_equal(x_fp8_reshape._rowwise_scale_inv, x_fp8._rowwise_scale_inv) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3], ids=str) + @pytest.mark.parametrize("fp8_dtype", [DType.kFloat8E4M3], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("dims", [[256, 512], [250, 500]]) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) def test_reshape_same_shape( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int + self, fp8_dtype: DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int ) -> None: """Test reshape operations that preserve tensor shape""" device = "cuda" @@ -559,12 +564,12 @@ def test_reshape_same_shape( with pytest.raises(AssertionError): torch.testing.assert_close(x_reshape.dequantize(), -x_hp, **_tols[fp8_dtype]) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3], ids=str) + @pytest.mark.parametrize("fp8_dtype", [DType.kFloat8E4M3], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("dims", [[256, 512], [250, 500]]) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) def test_clone_detach( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int + self, fp8_dtype: DType, dtype: torch.dtype, dims: DimsType, block_scaling_dim: int ) -> None: """Test clone and detach operations""" device = "cuda" diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index f54d16abe21..974ccee19cb 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -4,6 +4,7 @@ import torch from typing import Optional from transformer_engine.pytorch.router import ( + RoutingMapFormat, fused_topk_with_score_function, fused_compute_score_for_moe_aux_loss, fused_moe_aux_loss, @@ -461,6 +462,138 @@ 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 _bytemap_to_bitmap_u8(bytemap: torch.Tensor) -> torch.Tensor: + """Reference packer: bool[T, E] -> uint8[T, ceil(E/8)] LSB-first. + + Matches numpy.packbits(..., bitorder='little') + """ + flat = bytemap.to(torch.uint8).cpu().numpy() + import numpy as np + + return torch.from_numpy(np.packbits(flat, axis=-1, bitorder="little")).to(bytemap.device) + + +@pytest.mark.parametrize("dtype", [torch.float32]) +@pytest.mark.parametrize( + "num_tokens,num_experts,topk", + [(128, 32, 4), (256, 128, 8), (256, 130, 8), (128, 1024, 16)], +) +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) +def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function): + """fused_topk_with_score_function should produce identical probs and an + LSB-packed bitmap routing_map when routing_map_format=BITMAP_U8, and + backward gradients should match the bytemap path exactly.""" + if topk >= num_experts: + pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") + if score_function in ("sigmoid", "sqrtsoftplus"): + offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 + logits = ( + torch.arange(-num_experts // 2, num_experts // 2, device="cuda", dtype=dtype) * 1e-2 + ) + logits = logits.unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) + else: + logits = ( + torch.arange( + -num_tokens * num_experts // 2, + num_tokens * num_experts // 2, + device="cuda", + dtype=dtype, + ) + * 1e-4 + ) + logits = logits.view(num_tokens, num_experts) + + logits_byte = logits.detach().clone().requires_grad_(True) + logits_bit = logits.detach().clone().requires_grad_(True) + + probs_byte, routing_map_byte = fused_topk_with_score_function( + logits=logits_byte, + topk=topk, + use_pre_softmax=False, + num_groups=None, + group_topk=None, + scaling_factor=None, + score_function=score_function, + expert_bias=None, + routing_map_format=RoutingMapFormat.BYTEMAP, + ) + probs_bit, routing_map_bit = fused_topk_with_score_function( + logits=logits_bit, + topk=topk, + use_pre_softmax=False, + num_groups=None, + group_topk=None, + scaling_factor=None, + score_function=score_function, + expert_bias=None, + routing_map_format=RoutingMapFormat.BITMAP_U8, + ) + + assert probs_byte.dtype == probs_bit.dtype + torch.testing.assert_close(probs_byte, probs_bit, atol=0.0, rtol=0.0) + + expected_shape = (num_tokens, (num_experts + 7) // 8) + assert ( + routing_map_bit.shape == expected_shape + ), f"Bitmap shape {tuple(routing_map_bit.shape)} != {expected_shape}" + assert routing_map_bit.dtype == torch.uint8 + assert routing_map_byte.dtype == torch.bool + + packed_expected = _bytemap_to_bitmap_u8(routing_map_byte) + torch.testing.assert_close(routing_map_bit, packed_expected, atol=0, rtol=0) + + # Backward parity: grad of probs.sum() must be bit-identical across formats. + probs_byte.sum().backward() + probs_bit.sum().backward() + torch.testing.assert_close(logits_byte.grad, logits_bit.grad, atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize("dtype", [torch.float32]) +@pytest.mark.parametrize( + "num_tokens,num_experts,topk", + [(128, 32, 4), (256, 128, 8), (256, 130, 8)], +) +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) +def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function): + """fused_compute_score_for_moe_aux_loss: bitmap routing_map must equal + LSB-packed bytemap; scores must be bit-identical across formats.""" + if topk >= num_experts: + pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") + offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 + logits = torch.arange(-num_experts // 2, num_experts // 2, device="cuda", dtype=dtype) * 1e-2 + logits = logits.unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) + + logits_byte = logits.detach().clone().requires_grad_(True) + logits_bit = logits.detach().clone().requires_grad_(True) + + routing_map_byte, scores_byte = fused_compute_score_for_moe_aux_loss( + logits=logits_byte, + topk=topk, + score_function=score_function, + routing_map_format="bytemap", + ) + routing_map_bit, scores_bit = fused_compute_score_for_moe_aux_loss( + logits=logits_bit, + topk=topk, + score_function=score_function, + routing_map_format="bitmap_u8", + ) + + torch.testing.assert_close(scores_byte, scores_bit, atol=0.0, rtol=0.0) + + expected_shape = (num_tokens, (num_experts + 7) // 8) + assert routing_map_bit.shape == expected_shape + assert routing_map_bit.dtype == torch.uint8 + assert routing_map_byte.dtype == torch.bool + packed_expected = _bytemap_to_bitmap_u8(routing_map_byte) + torch.testing.assert_close(routing_map_bit, packed_expected, atol=0, rtol=0) + + # Backward parity through scores. + scores_byte.sum().backward() + scores_bit.sum().backward() + torch.testing.assert_close(logits_byte.grad, logits_bit.grad, atol=0.0, rtol=0.0) + + def profile_topk_softmax( dtype, num_tokens, diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index df607bd6dc9..14a52249f3b 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -21,6 +21,7 @@ import transformer_engine.pytorch as te import transformer_engine.pytorch.ops as te_ops from transformer_engine.pytorch.ops._common import ( + _cudnn_frontend_supports_grouped_gemm_srelu, _cudnn_frontend_version_supported, is_glu_activation, ) @@ -181,17 +182,17 @@ def make_reference_and_test_tensors( quantizer = Float8Quantizer( scale=torch.ones(1, dtype=torch.float32, device=test_device).squeeze(), amax=torch.zeros(1, dtype=torch.float32, device=test_device), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) test = quantizer(test) elif quantization == "fp8_current_scaling": quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, device=test_device, ) test = quantizer(test) elif quantization == "mxfp8": - test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) + test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): tensor_type = "input" if quantizer_role is not None: @@ -1933,9 +1934,9 @@ def test_clamped_swiglu( # Expected numerical error tols = dtype_tols(dtype) if quantized_compute and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): - tols = dtype_tols(tex.DType.kFloat4E2M1) + tols = dtype_tols(te.DType.kFloat4E2M1) elif quantized_compute: - tols = dtype_tols(tex.DType.kFloat8E4M3) + tols = dtype_tols(te.DType.kFloat8E4M3) # Check results assert_close(y_test, y_ref, **tols) @@ -3756,16 +3757,19 @@ def test_grouped_mlp( pytest.skip("single_grouped_bias requires bias=True") if with_quantization and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") - if not activation_is_glu and quantization != "mxfp8": - pytest.skip("Scaled unary grouped MLP fusion is only supported with MXFP8") + if not activation_is_glu and quantization not in ("mxfp8", "nvfp4", "nvfp4_rht"): + pytest.skip("Scaled unary grouped MLP is only supported with MXFP8 or NVFP4") if not activation_is_glu and glu_interleave_size is not None: pytest.skip("Unary activations do not use GLU interleaving") if quantization == "nvfp4_4over6": pytest.skip("NVFP4 4over6 grouped quantization is not supported") - if quantization == "nvfp4_rht" and ( - activation != "scaled_swiglu" or bias or glu_interleave_size != 32 - ): - pytest.skip("NVFP4 RHT grouped MLP coverage is limited to fused no-bias SwiGLU") + if activation == "scaled_srelu" and quantization in ("nvfp4", "nvfp4_rht") and bias: + pytest.skip("NVFP4 SReLU grouped MLP coverage is limited to no-bias") + if quantization == "nvfp4_rht": + if activation == "scaled_swiglu" and (bias or glu_interleave_size != 32): + pytest.skip("NVFP4 RHT SwiGLU grouped MLP coverage is limited to no-bias") + if activation not in ("scaled_swiglu", "scaled_srelu"): + pytest.skip("NVFP4 RHT grouped MLP coverage is limited to SwiGLU and SReLU") if ( with_quantization and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") @@ -4001,14 +4005,26 @@ def _make_module(): fc2.backward_dw() # Check for expected fusions - expected_grouped_mlp_fusion = ( - quantization == "mxfp8" - and dtype in (torch.bfloat16, torch.float16) - and ( - (not activation_is_glu and glu_interleave_size is None) - or (activation_is_glu and glu_interleave_size == 32) + cudnn_frontend_supports_grouped_mlp = ( + _cudnn_frontend_supports_grouped_gemm_srelu() + if activation == "scaled_srelu" + else _cudnn_frontend_version_supported() + ) + expected_grouped_mlp_fusion = cudnn_frontend_supports_grouped_mlp and ( + ( + quantization == "mxfp8" + and dtype in (torch.bfloat16, torch.float16) + and ( + (not activation_is_glu and glu_interleave_size is None) + or (activation_is_glu and glu_interleave_size == 32) + ) + ) + or ( + quantization == "nvfp4_rht" + and dtype == torch.bfloat16 + and activation == "scaled_srelu" + and glu_interleave_size is None ) - and _cudnn_frontend_version_supported() ) if expected_grouped_mlp_fusion: if activation_is_glu: @@ -5348,7 +5364,7 @@ def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: total_m = num_groups * m split_sizes = torch.full((num_groups,), m, device=device, dtype=torch.int64) - q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + q = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False) q.optimize_for_gemm = False torch.manual_seed(0) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 0dc253c18cf..c1a9e0a4070 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -4,7 +4,7 @@ import os import random -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Sequence import pytest import torch @@ -20,6 +20,7 @@ GroupedLinear, Linear, MXFP8Quantizer, + NVFP4Quantizer, autocast, is_bf16_available, quantized_model_init, @@ -35,13 +36,19 @@ ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor import transformer_engine_torch as tex -from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override +from utils import ( + ModelConfig, + assert_close, + recipe_id, + reset_rng_states, + skip_unsupported_backward_override, +) # Only run FP8 tests on supported devices. fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) fp8_block_scaling_available, _ = te.is_fp8_block_scaling_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) -nvfp4_available, _ = te.is_nvfp4_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) seed = 1234 reset_rng_states() @@ -866,6 +873,53 @@ def test_grouped_gemm(shape, dtype, layout, accumulate, use_cutlass, monkeypatch torch.testing.assert_close(o, o_ref, rtol=1.5e-2, atol=1.5e-2) +@pytest.mark.skipif( + torch.cuda.get_device_capability() != (9, 0), + reason="Only enable CUTLASS grouped gemm on Hopper", +) +@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) +def test_grouped_gemm_cutlass_empty_groups(layout, monkeypatch): + dtype = torch.bfloat16 + z, k, n = 1, 2048, 1536 + m_splits = [0] * z + + if layout == "TN": + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = [torch.empty(0, k, dtype=dtype, device="cuda") for _ in range(z)] # input + out = [torch.empty(0, n, dtype=dtype, device="cuda")] # output + grad = False + single_output = True + elif layout == "NN": + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = [torch.empty(0, n, dtype=dtype, device="cuda") for _ in range(z)] # grad_output + out = [torch.empty(0, k, dtype=dtype, device="cuda")] # dgrad + grad = True + single_output = True + else: # layout == "NT" + A = [torch.empty(0, k, dtype=dtype, device="cuda") for _ in range(z)] # input + B = [torch.empty(0, n, dtype=dtype, device="cuda") for _ in range(z)] # grad_output + out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad + grad = True + single_output = False + + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") + general_grouped_gemm( + A, + B, + out, + [None] * z, + dtype, + m_splits=m_splits, + grad=grad, + layout=layout, + single_output=single_output, + ) + torch.cuda.synchronize() + + for tensor in out: + torch.testing.assert_close(tensor, torch.zeros_like(tensor), rtol=0, atol=0) + + def _pack_grouped_tensor(grouped_tensor: GroupedTensor, tensors: List[torch.Tensor]) -> None: data = grouped_tensor.rowwise_data if data is None: @@ -1740,3 +1794,121 @@ def _train_step(x, dy, out_buf, *, use_graphed): for graph_grad, param in zip(graph_param_grads, grouped_linear.parameters()): assert param.grad is not None torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols) + + +@pytest.mark.parametrize("swizzle_type", ["mxfp8_rowwise", "mxfp8_columnwise", "nvfp4"]) +def test_swizzle_scales_and_pack_ptrs_for_discrete_weights( + swizzle_type: str, + num_tensors: int = 4, + shape: Sequence[int] = (160, 96), +): + """Helper function for preparing discrete weights for cuDNN group GEMM kernel""" + + # Skip unsupported configurations + if not mxfp8_available and swizzle_type in ("mxfp8_rowwise", "mxfp8_columnwise"): + pytest.skip(reason_for_no_mxfp8) + if not nvfp4_available and swizzle_type == "nvfp4": + pytest.skip(reason_for_no_nvfp4) + + # Construct quantizer + quantizer = None + if swizzle_type in ("mxfp8_rowwise", "mxfp8_columnwise"): + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=swizzle_type == "mxfp8_rowwise", + columnwise=swizzle_type == "mxfp8_columnwise", + ) + elif swizzle_type == "nvfp4": + quantizer = NVFP4Quantizer( + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=False, + ) + + # Per-expert tensors: unquantized, quantized with compact scales, + # quantized with swizzled scales + device = torch.device("cuda") + unquantized_tensors = [ + torch.randn(shape, dtype=torch.bfloat16, device=device) for _ in range(num_tensors) + ] + quantizer.optimize_for_gemm = False + tensors_with_compact_scales = [quantizer(t) for t in unquantized_tensors] + quantizer.optimize_for_gemm = True + tensors_with_swizzled_scales = [quantizer(t) for t in unquantized_tensors] + + # Extract data and scale buffers + if swizzle_type in ("mxfp8_rowwise", "nvfp4"): + data_tensors = [qx._rowwise_data for qx in tensors_with_compact_scales] + scale_tensors = [qx._rowwise_scale_inv for qx in tensors_with_compact_scales] + ref_scale_tensors = [qx._rowwise_scale_inv for qx in tensors_with_swizzled_scales] + elif swizzle_type == "mxfp8_columnwise": + data_tensors = [qx._columnwise_data for qx in tensors_with_compact_scales] + scale_tensors = [qx._columnwise_scale_inv for qx in tensors_with_compact_scales] + ref_scale_tensors = [qx._columnwise_scale_inv for qx in tensors_with_swizzled_scales] + else: + raise ValueError("Unrecogized swizzle type") + + # Call the helper function + data_ptrs, scale_ptrs, swizzled_scales_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + data_tensors, + scale_tensors, + swizzle_type, + device, + ) + ) + + # Check data pointer values + expected_data_ptrs = torch.tensor( + [t.data_ptr() for t in data_tensors], + dtype=torch.int64, + device="cpu", + ) + assert_close(data_ptrs, expected_data_ptrs) + + # Check scale pointer values + scale_bytes = scale_tensors[0].numel() * scale_tensors[0].element_size() + expected_scale_ptrs = torch.tensor( + [swizzled_scales_buffer.data_ptr() + i * scale_bytes for i in range(num_tensors)], + dtype=torch.int64, + device="cpu", + ) + assert_close(scale_ptrs, expected_scale_ptrs) + + # Check swizzled scale values + swizzled_scales_buffer = swizzled_scales_buffer.view(torch.uint8) + expected_swizzled_scales_buffer = ( + torch.cat(ref_scale_tensors).view(torch.uint8).view_as(swizzled_scales_buffer) + ) + assert_close( + swizzled_scales_buffer, + expected_swizzled_scales_buffer, + ) + + # Poison the padded compact scales + if swizzle_type == "mxfp8_rowwise": + unpadded_scale_shape = (shape[0], shape[1] // 32) + elif swizzle_type == "mxfp8_columnwise": + unpadded_scale_shape = (shape[0] // 32, shape[1]) + elif swizzle_type == "nvfp4": + unpadded_scale_shape = (shape[0], shape[1] // 16) + for scale in scale_tensors: + scale[unpadded_scale_shape[0] :, :].view(torch.uint8).fill_(-1) + scale[:, unpadded_scale_shape[1] :].view(torch.uint8).fill_(-1) + + # Check that swizzling removes poisoned pad scales + _, _, swizzled_scales_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + data_tensors, + scale_tensors, + swizzle_type, + device, + ) + ) + assert_close( + swizzled_scales_buffer, + expected_swizzled_scales_buffer, + ) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index c54c9758ff8..235825fb37d 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -20,6 +20,9 @@ from transformer_engine.pytorch.constants import TE_DType_To_Torch import transformer_engine_torch as tex +# Import test utilities +from utils import assert_close + # Check available recipes fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( @@ -61,17 +64,17 @@ def make_quantizer(quantization: str, num_tensors: int, shape: List[Tuple[int, i quantizer = Float8Quantizer( scale=torch.ones(1, dtype=torch.float32, device="cuda"), amax=torch.zeros(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) elif quantization == "fp8_current_scaling": quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, device="cuda", ) quantizer.set_usage(rowwise=True, columnwise=False) elif quantization == "fp8_blockwise": quantizer = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False, force_pow_2_scales=True, @@ -79,7 +82,7 @@ def make_quantizer(quantization: str, num_tensors: int, shape: List[Tuple[int, i block_scaling_dim=1, ) elif quantization == "mxfp8": - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) elif quantization == "nvfp4": quantizer = NVFP4Quantizer( with_rht=False, @@ -163,6 +166,93 @@ def test_basic_construction_varying_first_dim(self) -> None: shape[0][1], ) # sum of first dims + @pytest.mark.parametrize( + "split_sizes_list,logical_last_dim", + [ + pytest.param([3, 4, 5, 2], 7, id="all_nonzero"), + pytest.param([3, 0, 5, 2], 7, id="zero_middle"), + pytest.param([0, 3, 5, 0], 11, id="zero_edges"), + pytest.param([1], 17, id="single_group"), + pytest.param([1, 2, 3, 4, 5, 6, 7, 8], 13, id="many_groups"), + # MoE-style group counts. ``split_points`` (an int32[num_groups] + # tensor packed into a shared buffer alongside int64 outputs) used + # to land at an 8-byte-aligned offset for these counts, which + # tripped cuDNN's 16-byte alignment requirement in grouped GEMM. + pytest.param([8192] * 8, 2048, id="num_groups_8_uniform"), + pytest.param([4096] * 16, 4096, id="num_groups_16_uniform"), + pytest.param([2048] * 32, 7168, id="num_groups_32_uniform"), + pytest.param([1024] * 64, 7168, id="num_groups_64_uniform"), + pytest.param([512] * 128, 7168, id="num_groups_128_uniform"), + # Non-uniform with large totals to also exercise tensor_offsets > 2^31. + pytest.param( + [12345, 0, 8192, 1, 65536, 100, 131072, 7], + 7168, + id="non_uniform_large_totals", + ), + ], + ) + @pytest.mark.parametrize("input_dtype", [torch.int32, torch.int64], ids=["int32", "int64"]) + @pytest.mark.parametrize("input_device", ["cuda", "cpu"], ids=["cuda", "cpu"]) + @pytest.mark.parametrize("bulk_allocate", [False, True], ids=["separate", "bulk"]) + def test_splits_to_offsets_multi( + self, + bulk_allocate: bool, + input_device: str, + input_dtype: torch.dtype, + split_sizes_list: List[int], + logical_last_dim: int, + ) -> None: + """Test fused grouped split metadata preparation.""" + device = torch.device("cuda") + split_sizes = torch.tensor(split_sizes_list, dtype=input_dtype, device=input_device) + + # Exercise the grouped-MLP-shaped call: mix of int32 (no leading zero) + # and int64 (with leading zero) outputs, several strides. + strides = [1, 1, logical_last_dim, 0, logical_last_dim + 17] + include_leading_zero = [False, True, True, True, True] + dtypes = [torch.int32, torch.int64, torch.int64, torch.int64, torch.int64] + split_sizes_out, outputs = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=strides, + include_leading_zero=include_leading_zero, + dtypes=dtypes, + bulk_allocate=bulk_allocate, + ) + + # Reference implementation. + expected_split_sizes_i64 = split_sizes.to(device=device, dtype=torch.int64) + expected_base_offsets = torch.cat( + ( + torch.zeros(1, dtype=torch.int64, device=device), + torch.cumsum(expected_split_sizes_i64, dim=0), + ) + ) + + # Check output split_sizes: always int64, always on the target device. + assert split_sizes_out.device.type == "cuda" + assert split_sizes_out.dtype == torch.int64 + assert_close(split_sizes_out, expected_split_sizes_i64) + + # Check output offsets. + assert len(outputs) == len(strides) + for output, stride, with_zero, dtype in zip(outputs, strides, include_leading_zero, dtypes): + assert output.dtype == dtype + assert output.device.type == "cuda" + expected_length = split_sizes.numel() + (1 if with_zero else 0) + assert output.numel() == expected_length + expected = expected_base_offsets * stride + if not with_zero: + expected = expected[1:] + assert_close(output, expected) + + # Check pointer alignment: cuDNN CuTe-DSL grouped GEMM kernels + # require 16-byte-aligned data pointers. + for idx, output in enumerate(outputs): + assert ( + output.data_ptr() % 16 == 0 + ), f"outputs[{idx}] data_ptr is not 16-byte aligned: {output.data_ptr():#x}" + def test_split_into_quantized_tensors_no_quantization(self) -> None: """Test split_into_quantized_tensors for unquantized tensors""" num_tensors = 3 @@ -369,7 +459,7 @@ def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]], output_dbias grouped_input = torch.cat(input_tensors, dim=0) # Create MXFP8 output grouped tensor (rowwise only for easier validation) - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) quantizer.set_usage(rowwise=True, columnwise=False) first_dims = torch.tensor( [shape[0][0] for _ in range(num_tensors)], @@ -410,6 +500,64 @@ def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]], output_dbias expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) assert torch.allclose(dbias, expected_dbias) + @pytest.mark.parametrize("output_dbias", [False, True]) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_quantize_precomputed_offsets(self, output_dbias: bool) -> None: + """Test grouped quantization can reuse caller-provided tensor offsets.""" + num_tensors = 2 + last_dim = 1024 + split_sizes_list = [512, 512] + input_tensors = [ + torch.randn(split, last_dim, dtype=torch.bfloat16, device="cuda") + for split in split_sizes_list + ] + grouped_input = torch.cat(input_tensors, dim=0) + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=False) + split_sizes = torch.tensor(split_sizes_list, dtype=torch.int64, device="cuda") + split_sizes, (tensor_offsets,) = tex.splits_to_offsets_multi( + split_sizes, + torch.device("cuda"), + strides=[last_dim], + include_leading_zero=[True], + dtypes=[torch.int64], + ) + + if output_dbias: + grouped_output, dbias = tex.bgrad_group_quantize( + grouped_input, + quantizer, + num_tensors, + split_sizes, + tensor_offsets=tensor_offsets, + ) + expected_output, expected_dbias = tex.bgrad_group_quantize( + grouped_input, + quantizer, + num_tensors, + split_sizes, + ) + assert torch.allclose(dbias, expected_dbias) + else: + grouped_output = tex.group_quantize( + grouped_input, + quantizer, + num_tensors, + split_sizes, + tensor_offsets=tensor_offsets, + ) + expected_output = tex.group_quantize( + grouped_input, + quantizer, + num_tensors, + split_sizes, + ) + + assert grouped_output.tensor_offsets.data_ptr() == tensor_offsets.data_ptr() + assert torch.equal(grouped_output.rowwise_data, expected_output.rowwise_data) + assert torch.equal(grouped_output.scale_inv, expected_output.scale_inv) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_bgrad_group_quantize_zero_size_tensor(self) -> None: """Test bgrad_group_quantize handles zero-row input without error.""" @@ -417,7 +565,7 @@ def test_bgrad_group_quantize_zero_size_tensor(self) -> None: last_dim = 1024 grouped_input = torch.empty(0, last_dim, dtype=torch.bfloat16, device="cuda") - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) quantizer.set_usage(rowwise=True, columnwise=False) first_dims = torch.zeros(num_tensors, dtype=torch.int64, device="cuda") @@ -440,7 +588,7 @@ def test_group_quantize_cudagraph_capturable(self, output_dbias: bool) -> None: input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] grouped_input = torch.cat(input_tensors, dim=0) - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) quantizer.set_usage(rowwise=True, columnwise=False) first_dims = torch.tensor( [shape[0][0] for _ in range(num_tensors)], @@ -513,7 +661,7 @@ def test_group_dequantize(self, shape: List[Tuple[int, int]]) -> None: input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] grouped_input = torch.cat(input_tensors, dim=0) - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) quantizer.set_usage(rowwise=True, columnwise=False) first_dims = torch.tensor([s[0] for s in shape], dtype=torch.int64, device="cuda") @@ -521,7 +669,7 @@ def test_group_dequantize(self, shape: List[Tuple[int, int]]) -> None: quantized = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims) # Dequantize. - dequantized = tex.group_dequantize(quantized, tex.DType.kBFloat16) + dequantized = tex.group_dequantize(quantized, te.DType.kBFloat16) # Verify output metadata. assert dequantized.num_tensors == num_tensors @@ -543,7 +691,7 @@ def test_group_dequantize_cudagraph_capturable(self) -> None: input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] grouped_input = torch.cat(input_tensors, dim=0) - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) quantizer.set_usage(rowwise=True, columnwise=False) first_dims = torch.tensor( [shape[0][0] for _ in range(num_tensors)], @@ -556,12 +704,12 @@ def test_group_dequantize_cudagraph_capturable(self) -> None: # Warmup dequantize. torch.cuda.synchronize() - _ = tex.group_dequantize(quantized, tex.DType.kBFloat16) + _ = 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, tex.DType.kBFloat16) + static_output = tex.group_dequantize(quantized, te.DType.kBFloat16) # Replay with different input data. fresh_input = torch.cat( @@ -575,7 +723,7 @@ def test_group_dequantize_cudagraph_capturable(self) -> None: graph.replay() torch.cuda.synchronize() - expected = tex.group_dequantize(quantized, tex.DType.kBFloat16) + 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): diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index e087f1e1cd4..739aefd1f3c 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -43,6 +43,7 @@ from transformer_engine.pytorch import checkpoint as te_checkpoint from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.common import recipe +from transformer_engine.pytorch import DType import transformer_engine_torch as tex from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override @@ -2126,17 +2127,19 @@ def test_transformer_layer_hidden_states_format(dtype, bs, model): @pytest.mark.parametrize( "input_quantizer", [ - Float8CurrentScalingQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, device="cuda"), - MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + Float8CurrentScalingQuantizer(fp8_dtype=DType.kFloat8E4M3, device="cuda"), + MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3), ], ) @pytest.mark.parametrize( "out_quantizer", [ - Float8CurrentScalingQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, device="cuda"), - MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + Float8CurrentScalingQuantizer(fp8_dtype=DType.kFloat8E4M3, device="cuda"), + MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3), Float8Quantizer( - torch.ones(1).cuda().squeeze(), torch.ones(1).cuda().squeeze(), tex.DType.kFloat8E4M3 + torch.ones(1).cuda().squeeze(), + torch.ones(1).cuda().squeeze(), + DType.kFloat8E4M3, ), ], ) diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py index 9aea3bc274a..0ba15368820 100644 --- a/tests/pytorch/test_onnx_export.py +++ b/tests/pytorch/test_onnx_export.py @@ -32,7 +32,6 @@ from onnxruntime_extensions import PyCustomOpDef, get_library_path, onnx_op import transformer_engine.pytorch as te from transformer_engine.common import recipe -import transformer_engine_torch as tex from transformer_engine.pytorch.export import is_in_onnx_export_mode, te_translation_table from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.utils import get_default_init_method @@ -88,7 +87,7 @@ def trt_fp8_quantize(t, scale_inv): q = te.tensor.float8_tensor.Float8Quantizer( scale=1 / torch.from_numpy(scale_inv).cuda(), amax=torch.zeros([1]).cuda(), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) return q(x)._data.cpu().numpy() @@ -108,7 +107,7 @@ def trt_fp8_dequantize(t, scale_inv): q = te.tensor.float8_tensor.Float8Quantizer( scale=1 / torch.from_numpy(scale_inv).cuda(), amax=torch.zeros([1]).cuda(), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) quantizer_tensor = q.create_tensor_from_data(x, fake_dtype=torch.float32) return quantizer_tensor.dequantize().cpu().numpy() @@ -125,7 +124,7 @@ def trt_fp8_dequantize(t, scale_inv): def trt_mxfp8_quantize(t): """MXFP8 quantization extension for ONNX Runtime.""" x = torch.from_numpy(t).cuda() - q = te.tensor.mxfp8_tensor.MXFP8Quantizer(tex.DType.kFloat8E4M3) + q = te.tensor.mxfp8_tensor.MXFP8Quantizer(te.DType.kFloat8E4M3) return q(x)._rowwise_data.cpu().numpy(), q(x)._rowwise_scale_inv.cpu().numpy() @@ -142,7 +141,7 @@ def trt_mxfp8_dequantize(t, scale_inv): """MXFP8 dequantization extension for ONNX Runtime.""" x = torch.from_numpy(t).cuda() scale_inv_tensor = torch.from_numpy(scale_inv).cuda() - q = te.tensor.mxfp8_tensor.MXFP8Quantizer(tex.DType.kFloat8E4M3) + q = te.tensor.mxfp8_tensor.MXFP8Quantizer(te.DType.kFloat8E4M3) quantizer_tensor = q.create_tensor_from_data(x, scale_inv_tensor, fake_dtype=torch.float32) return quantizer_tensor.dequantize().cpu().numpy() @@ -382,9 +381,9 @@ def dtype2str(dtype: torch.dtype, fake_bf16_io=False): def as_te_type(dtype: torch.dtype): return { - torch.float32: tex.DType.kFloat32, - torch.float16: tex.DType.kFloat16, - torch.bfloat16: tex.DType.kBFloat16, + torch.float32: te.DType.kFloat32, + torch.float16: te.DType.kFloat16, + torch.bfloat16: te.DType.kBFloat16, }[dtype] diff --git a/tests/pytorch/test_permutation.py b/tests/pytorch/test_permutation.py index 66c685e1394..517d6e15436 100644 --- a/tests/pytorch/test_permutation.py +++ b/tests/pytorch/test_permutation.py @@ -25,7 +25,6 @@ Float8BlockQuantizer, MXFP8Quantizer, ) -import transformer_engine_torch as tex from transformer_engine.pytorch import Fp8Padding, Fp8Unpadding import copy @@ -191,19 +190,19 @@ def pytorch_sort_chunks_by_index( return output -def dtype_tols(te_dtype: tex.DType) -> Dict[str, float]: +def dtype_tols(te_dtype: te.DType) -> Dict[str, float]: """Estimated tolerances for a datatype Based on tolerances for torch.testing.assert_close. """ - if te_dtype == tex.DType.kFloat32: + if te_dtype == te.DType.kFloat32: return dict(rtol=1.0e-6, atol=1.0e-6) - if te_dtype == tex.DType.kFloat16: + if te_dtype == te.DType.kFloat16: return dict(rtol=3.0e-3, atol=1.0e-5) - if te_dtype == tex.DType.kBFloat16: + if te_dtype == te.DType.kBFloat16: return dict(rtol=2.0e-2, atol=1.0e-5) - if te_dtype == tex.DType.kFloat8E5M2 or te_dtype == tex.DType.kFloat8E4M3: + if te_dtype == te.DType.kFloat8E5M2 or te_dtype == te.DType.kFloat8E4M3: return dict(rtol=2.0e-1, atol=1.0e-1) raise ValueError(f"Unsuppored dtype ({te_dtype})") @@ -255,11 +254,11 @@ def _test_permutation_index_map( ) # Convert TE dtypes to PyTorch dtypes - if te_dtype == tex.DType.kFloat32: + if te_dtype == te.DType.kFloat32: dtype = torch.float32 - elif te_dtype == tex.DType.kFloat16: + elif te_dtype == te.DType.kFloat16: dtype = torch.float16 - elif te_dtype == tex.DType.kBFloat16: + elif te_dtype == te.DType.kBFloat16: dtype = torch.bfloat16 else: pytest.skip("Invalid dtype.") @@ -476,11 +475,11 @@ def _test_permutation_mask_map( ) # Convert TE dtypes to PyTorch dtypes - if te_dtype == tex.DType.kFloat32: + if te_dtype == te.DType.kFloat32: dtype = torch.float32 - elif te_dtype == tex.DType.kFloat16: + elif te_dtype == te.DType.kFloat16: dtype = torch.float16 - elif te_dtype == tex.DType.kBFloat16: + elif te_dtype == te.DType.kBFloat16: dtype = torch.bfloat16 else: pytest.skip("Invalid dtype.") @@ -704,11 +703,11 @@ def _test_permutation_and_padding_mask_map( ) # Convert TE dtypes to PyTorch dtypes - if te_dtype == tex.DType.kFloat32: + if te_dtype == te.DType.kFloat32: dtype = torch.float32 - elif te_dtype == tex.DType.kFloat16: + elif te_dtype == te.DType.kFloat16: dtype = torch.float16 - elif te_dtype == tex.DType.kBFloat16: + elif te_dtype == te.DType.kBFloat16: dtype = torch.bfloat16 else: pytest.skip("Invalid dtype.") @@ -1000,11 +999,11 @@ def _test_permutation_and_padding_with_merging_probs( ) # Convert TE dtypes to PyTorch dtypes - if te_dtype == tex.DType.kFloat32: + if te_dtype == te.DType.kFloat32: dtype = torch.float32 - elif te_dtype == tex.DType.kFloat16: + elif te_dtype == te.DType.kFloat16: dtype = torch.float16 - elif te_dtype == tex.DType.kBFloat16: + elif te_dtype == te.DType.kBFloat16: dtype = torch.bfloat16 else: pytest.skip("Invalid dtype.") @@ -1327,11 +1326,11 @@ def _test_moe_chunk_sort( ) # Convert TE dtypes to PyTorch dtypes - if te_dtype == tex.DType.kFloat32: + if te_dtype == te.DType.kFloat32: dtype = torch.float32 - elif te_dtype == tex.DType.kFloat16: + elif te_dtype == te.DType.kFloat16: dtype = torch.float16 - elif te_dtype == tex.DType.kBFloat16: + elif te_dtype == te.DType.kBFloat16: dtype = torch.bfloat16 else: pytest.skip("Invalid dtype.") @@ -1462,11 +1461,11 @@ def _test_permutation_mask_map_alongside_probs( ) # Convert TE dtypes to PyTorch dtypes - if te_dtype == tex.DType.kFloat32: + if te_dtype == te.DType.kFloat32: dtype = torch.float32 - elif te_dtype == tex.DType.kFloat16: + elif te_dtype == te.DType.kFloat16: dtype = torch.float16 - elif te_dtype == tex.DType.kBFloat16: + elif te_dtype == te.DType.kBFloat16: dtype = torch.bfloat16 else: pytest.skip("Invalid dtype.") @@ -1667,9 +1666,9 @@ def perf_test_cuda_kernel(cuda_kernel_fn): # TE tensor dtypes -_te_dtypes: List[tex.DType] = [tex.DType.kFloat32, tex.DType.kFloat16] +_te_dtypes: List[te.DType] = [te.DType.kFloat32, te.DType.kFloat16] if te.is_bf16_available(): - _te_dtypes.append(tex.DType.kBFloat16) + _te_dtypes.append(te.DType.kBFloat16) @pytest.mark.parametrize("te_dtype", _te_dtypes) @@ -1899,7 +1898,7 @@ def test_permutation_mask_map_alongside_probs_empty_input(te_dtype, use_torch_co @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize("te_dtype", [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]) +@pytest.mark.parametrize("te_dtype", [te.DType.kFloat8E4M3, te.DType.kFloat8E5M2]) @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) @@ -2048,9 +2047,9 @@ def test_chunk_permutation_empty_input(te_dtype, use_torch_compile): def test_permutation_single_case(): print("GPU:", torch.cuda.get_device_name(0)) - # te_dtype = tex.DType.kFloat32 - # te_dtype = tex.DType.kFloat16 - te_dtype = tex.DType.kBFloat16 + # te_dtype = te.DType.kFloat32 + # te_dtype = te.DType.kFloat16 + te_dtype = te.DType.kBFloat16 num_tokens = 12 num_expert = 4 @@ -2216,9 +2215,9 @@ def test_benchmark_multiple_cases(): """Benchmark test - skipped by default. Run with: RUN_BENCHMARK_TESTS=1 pytest -k benchmark""" print("GPU:", torch.cuda.get_device_name(0)) - # te_dtype = tex.DType.kFloat32 - # te_dtype = tex.DType.kFloat16 - te_dtype = tex.DType.kBFloat16 + # te_dtype = te.DType.kFloat32 + # te_dtype = te.DType.kFloat16 + te_dtype = te.DType.kBFloat16 ep_size = 64 tp_size = 2 diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index c5161349eff..f1d64696aca 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -25,7 +25,6 @@ ) from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported -import transformer_engine_torch as tex from references.ref_per_tensor_cs import ref_per_tensor_cs_cast from utils import assert_close @@ -33,12 +32,12 @@ # PyTorch tensor dtypes _dtypes: List[torch.dtype] = [torch.float32, torch.float16, torch.bfloat16] # TE FP8 dtypes -_fp8_dtypes: List[tex.DType] = [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2] +_fp8_dtypes: List[te.DType] = [te.DType.kFloat8E4M3, te.DType.kFloat8E5M2] # Numerical tolerances with FP8 types -_tols: Dict[tex.DType, Dict[str, float]] = { - tex.DType.kFloat8E4M3: dict(rtol=0.125, atol=0.0675), # epsilon = 0.0625 - tex.DType.kFloat8E5M2: dict(rtol=0.25, atol=0.125), # epsilon = 0.125 +_tols: Dict[te.DType, Dict[str, float]] = { + te.DType.kFloat8E4M3: dict(rtol=0.125, atol=0.0675), # epsilon = 0.0625 + te.DType.kFloat8E5M2: dict(rtol=0.25, atol=0.125), # epsilon = 0.125 } @@ -76,7 +75,7 @@ def _to_list(x: Union[Iterable, Any]) -> List: # delayed scaling def to_float8( tensor: torch.Tensor, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, scale: float = 1.0, ) -> Float8Tensor: """Cast tensor to FP8""" @@ -91,7 +90,7 @@ def to_float8( # current scaling def to_float8_CS( tensor: torch.Tensor, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, return_transpose: bool = False, force_pow_2_scales: bool = False, amax_epsilon: float = 0.0, @@ -144,18 +143,18 @@ def make_reference_and_test_tensors( quantizer = Float8Quantizer( scale=torch.ones(1, dtype=torch.float32, device=test_device).squeeze(), amax=torch.zeros(1, dtype=torch.float32, device=test_device), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) test = quantizer(test) elif quantization == "fp8_current_scaling": quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, device=test_device, ) test = quantizer(test) elif quantization == "fp8_blockwise": quantizer = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=True, force_pow_2_scales=True, @@ -164,7 +163,7 @@ def make_reference_and_test_tensors( ) test = quantizer(test) elif quantization == "mxfp8": - test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) + test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): row_scaled_nvfp4 = quantization == "nvfp4_row_scaled" test = NVFP4Quantizer( @@ -201,7 +200,7 @@ def setup_class(cls) -> None: def test_constructor( self, dims: DimsType = 1, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, scale_inv: float = 0.375, dtype: torch.dtype = torch.float32, ) -> None: @@ -220,7 +219,7 @@ def test_constructor( def _test_quantize_dequantize( self, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, scale: float = 3.5, dtype: torch.dtype = torch.float32, dims: DimsType = 23, @@ -245,7 +244,7 @@ def _test_quantize_dequantize( @pytest.mark.parametrize("dtype", _dtypes) def test_quantize_dequantize_dtypes( self, - fp8_dtype: tex.DType, + fp8_dtype: te.DType, dtype: torch.dtype, ) -> None: self._test_quantize_dequantize(fp8_dtype=fp8_dtype, dtype=dtype) @@ -262,7 +261,7 @@ def test_quantize_dequantize_dims(self, dims: DimsType) -> None: @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("noop", [True, False]) def test_quantize_dequantize_noop( - self, fp8_dtype: tex.DType, dtype: torch.dtype, noop: bool + self, fp8_dtype: te.DType, dtype: torch.dtype, noop: bool ) -> None: noop_tensor = torch.zeros(1, dtype=torch.float32, device="cuda") if noop: @@ -287,7 +286,7 @@ def test_quantize_dequantize_noop( def test_basic_ops( self, dims: DimsType = 23, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, scale: float = 3.5, dtype: torch.dtype = torch.float32, ) -> None: @@ -323,7 +322,7 @@ def test_basic_ops( def test_chunk_op( self, dims: DimsType, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, scale: float = 3.5, dtype: torch.dtype = torch.float32, ) -> None: @@ -352,7 +351,7 @@ def test_chunk_op( def test_inplace_ops( self, dims: DimsType = 23, - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, scale: float = 3.5, dtype: torch.dtype = torch.float32, ) -> None: @@ -390,7 +389,7 @@ def test_inplace_ops( def test_serialization( self, dims: DimsType = [2, 3, 5], - fp8_dtype: tex.DType = tex.DType.kFloat8E4M3, + fp8_dtype: te.DType = te.DType.kFloat8E4M3, scale: float = 0.5, dtype: torch.dtype = torch.float32, ): @@ -490,7 +489,7 @@ def setup_class(cls) -> None: @pytest.mark.parametrize("amax_epsilon", [0.0, 1e-6], ids=str) def test_quantize( self, - fp8_dtype: tex.DType, + fp8_dtype: te.DType, dtype: torch.dtype, dims: DimsType, return_transpose: bool, @@ -532,11 +531,11 @@ def test_quantize( x_fp8._transpose, x_fp8_t_ref.view(torch.uint8), atol=0.0, rtol=0.0 ) - @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3], ids=str) + @pytest.mark.parametrize("fp8_dtype", [te.DType.kFloat8E4M3], ids=str) @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("dims", [[], 1, 311, [7, 11], [7, 5, 3], [2, 3, 5, 3]]) def test_quantize_dequantize( - self, fp8_dtype: tex.DType, dtype: torch.dtype, dims: DimsType + self, fp8_dtype: te.DType, dtype: torch.dtype, dims: DimsType ) -> None: """Check numerical error when casting to FP8 and back""" @@ -778,11 +777,11 @@ def test_update_nd_tensor( quantizer = Float8Quantizer( scale=torch.ones(1, dtype=torch.float32, device=device).squeeze(), amax=torch.zeros(1, dtype=torch.float32, device=device), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, ) elif quantization == "fp8_blockwise": quantizer = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=True, force_pow_2_scales=True, @@ -790,7 +789,7 @@ def test_update_nd_tensor( block_scaling_dim=1, ) elif quantization == "mxfp8": - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) elif quantization in ("nvfp4", "nvfp4_2d", "nvfp4_row_scaled", "nvfp4_4over6"): row_scaled_nvfp4 = quantization == "nvfp4_row_scaled" quantizer = NVFP4Quantizer( @@ -835,7 +834,7 @@ def setup_class(cls) -> None: @pytest.mark.parametrize("dims", [[128, 128], [256, 256], [128, 256]]) def test_mxfp8_dequantize_columnwise_only( self, - fp8_dtype: tex.DType, + fp8_dtype: te.DType, dtype: torch.dtype, dims: DimsType, ) -> None: @@ -876,7 +875,7 @@ def test_mxfp8_dequantize_columnwise_only( @pytest.mark.parametrize("dims", [[128, 128], [256, 256]]) def test_mxfp8_dequantize_columnwise_only_quantized_separately( self, - fp8_dtype: tex.DType, + fp8_dtype: te.DType, dims: DimsType, ) -> None: """Check dequantization of MXFP8 tensor quantized with columnwise only""" diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 9a14cee7fd6..fd3c5a3463e 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -297,14 +297,16 @@ def check_metas( @pytest.mark.parametrize("amax_case", ["zero", "tiny", "normal", "inf", "nan"]) @pytest.mark.parametrize("fused_update", [True, False], ids=["fused", "non-fused"]) @pytest.mark.parametrize( - "fp8_dtype", [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2], ids=["E4M3", "E5M2"] + "fp8_dtype", + [te.DType.kFloat8E4M3, te.DType.kFloat8E5M2], + ids=["E4M3", "E5M2"], ) def test_scale_update_numeric_scenarios(self, amax_case, fused_update, fp8_dtype): - if fp8_dtype == tex.DType.kFloat8E4M3: + if fp8_dtype == te.DType.kFloat8E4M3: fp8_format = transformer_engine.common.recipe.Format.E4M3 fp8_max = fp8_format.value.max_fwd - elif fp8_dtype == tex.DType.kFloat8E5M2: + elif fp8_dtype == te.DType.kFloat8E5M2: fp8_format = transformer_engine.common.recipe.Format.HYBRID fp8_max = fp8_format.value.max_bwd else: diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 404fae85fde..7c19dc6536d 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -35,7 +35,6 @@ is_bf16_available, ) from transformer_engine.common import recipe -import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data from utils import ModelConfig, recipe_id, skip_unsupported_backward_override @@ -1064,7 +1063,7 @@ def test_sanity_fp8_gemm_with_unalignment(N, datatype): scales = torch.ones(1).cuda().squeeze() amaxes = torch.ones(1).cuda().squeeze() - dtype = tex.DType.kFloat8E4M3 + dtype = te.DType.kFloat8E4M3 fp8_quantizer = Float8Quantizer(scales, amaxes, dtype) outp_type = datatype @@ -1088,7 +1087,7 @@ def test_replace_raw_data_for_float8tensor(): torch.manual_seed(12345) torch.cuda.manual_seed(12345) - fp8_quantizer = Float8CurrentScalingQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, device="cuda") + fp8_quantizer = Float8CurrentScalingQuantizer(fp8_dtype=te.DType.kFloat8E4M3, device="cuda") fp8_tensor = fp8_quantizer.make_empty([128, 128], dtype=torch.bfloat16, device="cuda") random_bf16_data = torch.randn(fp8_tensor.shape, dtype=torch.bfloat16, device="cuda") fp8_quantizer.update_quantized(random_bf16_data, fp8_tensor) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 137e5f5a778..309a5d124ec 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -94,7 +94,7 @@ class ToyQuantizer(Float8CurrentScalingQuantizer, metaclass=_ToyQuantizerMeta): opaque value type so torch.compile can treat it as a baked-in constant.""" def __init__(self, tag: str): - super().__init__(fp8_dtype=tex.DType.kFloat8E4M3, device=torch.device("cuda")) + super().__init__(fp8_dtype=te.DType.kFloat8E4M3, device=torch.device("cuda")) self.tag = tag def __eq__(self, other): diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 84489f30c1e..f068f3581e1 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -17,9 +17,9 @@ import torch import transformer_engine -import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch import InferenceParams, QuantizedTensor +from transformer_engine.pytorch import DType from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.attention.dot_product_attention.utils import ( get_attention_backend, @@ -70,7 +70,7 @@ def str_to_dtype(dtype: str | torch.dtype) -> torch.dtype: return dtype -def dtype_tols(dtype: torch.dtype | tex.DType) -> dict[str, float]: +def dtype_tols(dtype: torch.dtype | DType) -> dict[str, float]: """Estimated numerical error for a datatype Based on tolerances for torch.testing.assert_close. @@ -78,17 +78,17 @@ def dtype_tols(dtype: torch.dtype | tex.DType) -> dict[str, float]: """ # Transformer Engine dtypes - if isinstance(dtype, tex.DType): - if dtype == tex.DType.kFloat4E2M1: + if isinstance(dtype, DType): + if dtype == DType.kFloat4E2M1: return dict(rtol=0.25, atol=0.125) # epsilon = 0.25 dtype = { - tex.DType.kByte: torch.uint8, - tex.DType.kInt32: torch.int32, - tex.DType.kFloat32: torch.float32, - tex.DType.kFloat16: torch.half, - tex.DType.kBFloat16: torch.bfloat16, - tex.DType.kFloat8E4M3: torch.float8_e4m3fn, - tex.DType.kFloat8E5M2: torch.float8_e5m2, + DType.kByte: torch.uint8, + DType.kInt32: torch.int32, + DType.kFloat32: torch.float32, + DType.kFloat16: torch.half, + DType.kBFloat16: torch.bfloat16, + DType.kFloat8E4M3: torch.float8_e4m3fn, + DType.kFloat8E5M2: torch.float8_e5m2, }[dtype] # PyTorch dtypes @@ -117,9 +117,9 @@ def quantization_tols(name: str) -> dict[str, float]: "mxfp8", "mxfp8_block_scaling", ): - return dtype_tols(tex.DType.kFloat8E4M3) + return dtype_tols(DType.kFloat8E4M3) if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): - return dtype_tols(tex.DType.kFloat4E2M1) + return dtype_tols(DType.kFloat4E2M1) raise ValueError(f"Unsupported quantization scheme ({name})") @@ -341,6 +341,8 @@ def get_available_attention_backends( fp8_meta: Optional[Dict[str, Any]] = None, is_training: bool = True, inference_params: Optional[InferenceParams] = None, + score_mod: bool = False, + score_mod_bprop: bool = False, ) -> Tuple[List, List]: """Check for all available attention backends that support a model configuration""" @@ -402,6 +404,8 @@ def test(): inference_params=inference_params, softmax_type=config.softmax_type, return_max_logit=config.return_max_logit, + has_score_mod=score_mod, + has_score_mod_bprop=score_mod_bprop, # allow all backends to pass so they can be used for testing; # check for FA3 availability later num_splits=1, diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 06d85b6d84f..8f96432ed85 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -139,6 +139,40 @@ function(find_nccl_version OUT_VERSION OUT_INCLUDE_DIR) set(${OUT_INCLUDE_DIR} "${_nvte_nccl_include_dir}" PARENT_SCOPE) endfunction() +function(find_cublasmp_version OUT_VERSION OUT_INCLUDE_DIR SEARCH_DIR) + find_path(_nvte_cublasmp_include_dir + NAMES cublasmp.h + HINTS "${SEARCH_DIR}/include" + PATH_SUFFIXES include + REQUIRED) + + file(STRINGS "${_nvte_cublasmp_include_dir}/cublasmp.h" _nvte_cublasmp_major_line + REGEX "^#define CUBLASMP_VER_MAJOR[ \t]+[0-9]+$") + file(STRINGS "${_nvte_cublasmp_include_dir}/cublasmp.h" _nvte_cublasmp_minor_line + REGEX "^#define CUBLASMP_VER_MINOR[ \t]+[0-9]+$") + file(STRINGS "${_nvte_cublasmp_include_dir}/cublasmp.h" _nvte_cublasmp_patch_line + REGEX "^#define CUBLASMP_VER_PATCH[ \t]+[0-9]+$") + + string(REGEX REPLACE "^#define CUBLASMP_VER_MAJOR[ \t]+([0-9]+)$" "\\1" + _nvte_cublasmp_major "${_nvte_cublasmp_major_line}") + string(REGEX REPLACE "^#define CUBLASMP_VER_MINOR[ \t]+([0-9]+)$" "\\1" + _nvte_cublasmp_minor "${_nvte_cublasmp_minor_line}") + string(REGEX REPLACE "^#define CUBLASMP_VER_PATCH[ \t]+([0-9]+)$" "\\1" + _nvte_cublasmp_patch "${_nvte_cublasmp_patch_line}") + + if ("${_nvte_cublasmp_major}" STREQUAL "" + OR "${_nvte_cublasmp_minor}" STREQUAL "" + OR "${_nvte_cublasmp_patch}" STREQUAL "") + message(FATAL_ERROR + "Failed to parse cuBLASMp version from ${_nvte_cublasmp_include_dir}/cublasmp.h") + endif() + + set(${OUT_VERSION} + "${_nvte_cublasmp_major}.${_nvte_cublasmp_minor}.${_nvte_cublasmp_patch}" + PARENT_SCOPE) + set(${OUT_INCLUDE_DIR} "${_nvte_cublasmp_include_dir}" PARENT_SCOPE) +endfunction() + # Configure Transformer Engine library include_directories(${PROJECT_SOURCE_DIR}/..) set(transformer_engine_SOURCES) @@ -194,6 +228,7 @@ list(APPEND transformer_engine_cuda_sources permutation/permutation.cu util/utils.cu util/padding.cu + util/splits_to_offsets.cu util/topk.cu swizzle/swizzle.cu swizzle/swizzle_block_scaling.cu @@ -352,25 +387,36 @@ endif() option(NVTE_WITH_CUBLASMP "Use cuBLASMp for tensor parallel GEMMs" OFF) if (NVTE_WITH_CUBLASMP) + target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUBLASMP) target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include) find_nccl_version(NCCL_VERSION NCCL_INCLUDE_DIR) + find_cublasmp_version(CUBLASMP_VERSION CUBLASMP_INCLUDE_DIR ${CUBLASMP_DIR}) find_library(CUBLASMP_LIB - NAMES cublasmp libcublasmp + NAMES cublasmp libcublasmp.so libcublasmp.so.0 PATHS ${CUBLASMP_DIR} - PATH_SUFFIXES lib + PATH_SUFFIXES lib lib64 lib/aarch64-linux-gnu lib/sbsa-linux-gnu lib/x86_64-linux-gnu REQUIRED) find_library(NCCL_LIB NAMES nccl libnccl PATH_SUFFIXES lib REQUIRED) - if (NCCL_VERSION VERSION_LESS 2.29.0) + # cuBLASMp 0.8 is the first release with CUDA-graph-safe overlap algos, + # and NCCL 2.30 is the first release with graph-safe one-sided RMA + # primitives (ncclPutSignal/ncclWaitSignal) that those algos use. + if (CUBLASMP_VERSION VERSION_LESS 0.8.0) + message(FATAL_ERROR + "NVTE_WITH_CUBLASMP requires cuBLASMp >= 0.8.0, but found cuBLASMp " + "${CUBLASMP_VERSION} in ${CUBLASMP_INCLUDE_DIR}/cublasmp.h") + endif() + if (NCCL_VERSION VERSION_LESS 2.30.0) message(FATAL_ERROR - "NVTE_WITH_CUBLASMP requires NCCL >= 2.29.0, but found NCCL ${NCCL_VERSION} " - "in ${NCCL_INCLUDE_DIR}/nccl.h") + "NVTE_WITH_CUBLASMP requires NCCL >= 2.30.0 (for graph-capture-safe " + "one-sided RMA primitives used by cuBLASMp's overlap algorithms), but " + "found NCCL ${NCCL_VERSION} in ${NCCL_INCLUDE_DIR}/nccl.h") endif() target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB} ${CUBLASMP_LIB}) - message(STATUS "Using cuBLASMp at: ${CUBLASMP_DIR}") + message(STATUS "Using cuBLASMp ${CUBLASMP_VERSION} at: ${CUBLASMP_DIR}") message(STATUS "Using NCCL ${NCCL_VERSION} at: ${NCCL_LIB}") endif() diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index 60632b99d84..3545347da47 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -181,6 +181,7 @@ void GemmRsInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n get_cuda_dtype(a->dtype()), ctx->grid_row_major.get(), ctx->a_desc.get())); } + // B is (K/P, N) local -- K is distributed across ranks, N is fully replicated. if (transb) { NVTE_CHECK(b1 == n, "Unsupported tensor dimension in B: expected ", n, ", got ", b1); NVTE_CHECK_CUBLASMP(cublasMpMatrixDescriptorInit(n, k, block_size(ctx, n), block_size(ctx, k), @@ -264,10 +265,47 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo "Unsupported scaling mode: " + std::to_string(t->scaling_mode)); } + // Mirror cublaslt_gemm.cu's CanonicalizeGemmInput for FP8 tensor scaling: depending on the + // architecture and the quantizer's usage modes, the appropriate data + scale_inv + // may live on the rowwise or columnwise side of the tensor. + // * Hopper (!nvte_is_non_tn_fp8_gemm_supported): only TN FP8 GEMMs are supported, so an + // FP8 input not already in TN orientation must be swapped to its columnwise (transposed) + // view and the transpose flag flipped. + // * Blackwell+ (nvte_is_non_tn_fp8_gemm_supported): any FP8 GEMM layout is supported, but + // the quantizer usage may have only been set to columnwise. In that case, fall back to + // the columnwise view and flip the transpose flag so the GEMM sees the matching data and + // scale_inv pair. + // The original tensor is never modified; a new Tensor view aliases the columnwise pointers. + const bool fp8_needs_tn = !nvte_is_non_tn_fp8_gemm_supported(); + auto canonicalize_fp8_input = [fp8_needs_tn](const Tensor* t, bool current_trans, bool want_trans, + const char* side) -> std::pair { + if (!is_fp8_dtype(t->dtype())) { + return {*t, current_trans}; + } + const bool hopper_tn_swap = fp8_needs_tn && current_trans != want_trans; + const bool blackwell_missing_rowwise = !fp8_needs_tn && !t->has_data(); + if (!hopper_tn_swap && !blackwell_missing_rowwise) { + return {*t, current_trans}; + } + NVTE_CHECK(t->has_columnwise_data() && is_fp8_dtype(t->columnwise_data.dtype), + "cuBLASMp FP8 GEMM input ", side, " is missing column-wise usage"); + Tensor view; + view.scaling_mode = t->scaling_mode; + view.data = t->columnwise_data; + view.scale_inv = t->columnwise_scale_inv; + // Columnwise data is the transposed view of the original — flip the transpose flag. + return {view, !current_trans}; + }; + + auto [a_used, transa_eff] = canonicalize_fp8_input(a, transa, /*want_trans=*/true, "A"); + auto [b_used, transb_eff] = canonicalize_fp8_input(b, transb, /*want_trans=*/false, "B"); + transa = transa_eff; + transb = transb_eff; + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorInit(ctx->matmul_desc.get(), CUBLAS_COMPUTE_32F)); int64_t ldd{}; - init_matrices_fn(ctx, &ldd, m, n, k, a, b, d, transa, transb); + init_matrices_fn(ctx, &ldd, m, n, k, &a_used, &b_used, d, transa, transb); const cublasOperation_t trans_a = transa ? CUBLAS_OP_T : CUBLAS_OP_N; const cublasOperation_t trans_b = transb ? CUBLAS_OP_T : CUBLAS_OP_N; @@ -283,23 +321,23 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo sizeof algo_attr)); const cublasMpMatmulMatrixScale_t scale_mode = CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32; - if (is_fp8_dtype(a->dtype())) { - NVTE_CHECK(a->scale_inv.dptr, "Scaling must be set for FP8 dtype"); + if (is_fp8_dtype(a_used.dtype())) { + NVTE_CHECK(a_used.scale_inv.dptr, "Scaling must be set for FP8 dtype"); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_MODE, &scale_mode, sizeof scale_mode)); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_POINTER, - &a->scale_inv.dptr, sizeof(void*))); + &a_used.scale_inv.dptr, sizeof(void*))); } - if (is_fp8_dtype(b->dtype())) { - NVTE_CHECK(b->scale_inv.dptr, "Scaling must be set for FP8 dtype"); + if (is_fp8_dtype(b_used.dtype())) { + NVTE_CHECK(b_used.scale_inv.dptr, "Scaling must be set for FP8 dtype"); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_MODE, &scale_mode, sizeof scale_mode)); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_POINTER, - &b->scale_inv.dptr, sizeof(void*))); + &b_used.scale_inv.dptr, sizeof(void*))); } if (is_fp8_dtype(d->dtype())) { NVTE_CHECK(d->scale.dptr, "Scaling must be set for FP8 dtype"); @@ -398,11 +436,11 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo n, k, &alpha, - a->data.dptr, + a_used.data.dptr, 1, 1, ctx->a_desc.get(), - b->data.dptr, + b_used.data.dptr, 1, 1, ctx->b_desc.get(), @@ -436,9 +474,6 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo std::apply(cublasMpMatmul, std::tuple_cat(args, std::tuple{ctx->workspace, ctx->workspace_size, workspace_host.data(), workspace_host.size()}))); - - NVTE_CHECK_CUDA(cudaEventRecord(ctx->event.get(), main_stream)); - NVTE_CHECK_CUDA(cudaStreamWaitEvent(ctx->stream.get(), ctx->event.get(), 0)); } } // namespace @@ -472,6 +507,8 @@ NVTECommGemmCtx* nvte_comm_gemm_ctx_create(ncclComm_t comm, int nranks, int rank .b_desc = std::move(b_desc), .d_desc = std::move(d_desc), .matmul_desc = std::move(matmul_desc), + .workspace = nullptr, + .workspace_size = 0, }; } diff --git a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp index a1a0dd9d0bd..2eda743af99 100644 --- a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp @@ -4,6 +4,7 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include #include @@ -23,6 +24,14 @@ using namespace std::placeholders; +bool nvte_built_with_cublasmp() { +#ifdef NVTE_WITH_CUBLASMP + return true; +#else + return false; +#endif +} + namespace transformer_engine { namespace { @@ -33,10 +42,6 @@ std::vector shape_to_vector(const NVTEShape &shape) { } // namespace -/*************************************************************************************************** - * Comm+GEMM Overlap Common Core - **************************************************************************************************/ - bool ubuf_built_with_mpi() { #ifdef NVTE_UB_WITH_MPI return true; @@ -45,6 +50,11 @@ bool ubuf_built_with_mpi() { #endif } +/*************************************************************************************************** + * Comm+GEMM Overlap Common Core + **************************************************************************************************/ + +// Constructor for Userbuffers backend CommOverlapCore::CommOverlapCore(int myrank, int numranks, int mylocal, int numlocal, int mynode, int numnodes, int tp_size, ExtAllgatherOp allgather_handle, ExtBarrierOp barrier_handle, int num_splits, int num_max_streams, @@ -69,6 +79,40 @@ CommOverlapCore::CommOverlapCore(int myrank, int numranks, int mylocal, int numl num_comm_sm, set_sm_margin, use_ce, atomic_gemm); } +// Constructor for cuBLASMp backend +CommOverlapCore::CommOverlapCore(ncclComm_t nccl_comm_ptr, int tp_rank, int tp_size, + int num_comm_sm, bool is_p2p, bool atomic_gemm) { + NVTE_CHECK( + nvte_built_with_cublasmp(), + "Comm+GEMM overlap with cuBLASMp backend requires TE to be built with NVTE_WITH_CUBLASMP=1."); + _with_cublasmp = true; + _cublasmp_ctx = nvte_comm_gemm_ctx_create(nccl_comm_ptr, tp_size, tp_rank); + + _tp_id = tp_rank; + _tp_size = tp_size; + _num_comm_sm = num_comm_sm; + _is_p2p = is_p2p; + _atomic_gemm = atomic_gemm; + if (transformer_engine::getenv("UB_SKIPMC") || + !transformer_engine::cuda::supports_multicast()) { + _algo_type = kNVTECommGemmAlgoDefault; // default algo includes fallback to non-multicast + } else { + if (_is_p2p) { + if (_atomic_gemm) { + _algo_type = kNVTECommGemmAlgoAtomicP2P; + } else { + _algo_type = kNVTECommGemmAlgoSplitP2P; + } + } else { + if (_atomic_gemm) { + _algo_type = kNVTECommGemmAlgoAtomicMulticast; + } else { + _algo_type = kNVTECommGemmAlgoSplitMulticast; + } + } + } +} + void CommOverlapCore::initialize(int tp_size, int num_splits, int num_max_streams, int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool use_ce, @@ -134,39 +178,43 @@ void CommOverlapCore::initialize(int tp_size, int num_splits, int num_max_stream } CommOverlapCore::~CommOverlapCore() { - cudaEventDestroy(_stop_comm); - cudaEventDestroy(_start_comm); - cudaEventDestroy(_stop_compute); - cudaEventDestroy(_start_compute); - if (_comm_launch_event) { - cudaEventDestroy(_comm_launch_event); - } + if (_with_cublasmp) { + nvte_comm_gemm_ctx_destroy(_cublasmp_ctx); + } else { + cudaEventDestroy(_stop_comm); + cudaEventDestroy(_start_comm); + cudaEventDestroy(_stop_compute); + cudaEventDestroy(_start_compute); + if (_comm_launch_event) { + cudaEventDestroy(_comm_launch_event); + } - if (_atomic_gemm) { - cudaFree(_counter.dptr()); - } + if (_atomic_gemm) { + cudaFree(_counter.dptr()); + } - for (size_t i = 0; i < _stream_compute.size(); i++) { - cudaStreamSynchronize(_stream_compute[i]); - cudaStreamDestroy(_stream_compute[i]); - } + for (size_t i = 0; i < _stream_compute.size(); i++) { + cudaStreamSynchronize(_stream_compute[i]); + cudaStreamDestroy(_stream_compute[i]); + } - auto error = cudaGetLastError(); - if (error != cudaSuccess) { - NVTE_WARN("Error detected while destroying communicator: ", cudaGetErrorString(error)); - } + auto error = cudaGetLastError(); + if (error != cudaSuccess) { + NVTE_WARN("Error detected while destroying communicator: ", cudaGetErrorString(error)); + } - if (_comm_created) { - try { + if (_comm_created) { + try { #ifdef NVTE_UB_WITH_MPI - destroy_communicator_mpi(_ub_comm); + destroy_communicator_mpi(_ub_comm); #else - destroy_communicator(_ub_comm); + destroy_communicator(_ub_comm); #endif - } catch (const std::exception &e) { - NVTE_WARN("Error destroying communicator, cleanup may be incomplete:\n", e.what()); + } catch (const std::exception &e) { + NVTE_WARN("Error destroying communicator, cleanup may be incomplete:\n", e.what()); + } + _comm_created = false; } - _comm_created = false; } } @@ -294,10 +342,57 @@ TensorWrapper CommOverlapCore::get_buffer_chunk_like(const TensorWrapper &source return chunk; } +namespace { + +struct CublasMpDims { + int64_t m, n, k; +}; + +// Resolve the global m/n/k for the three cuBLASMp communication patterns +// from flattened operand shapes. +CublasMpDims compute_ag_dims(const TensorWrapper &A, bool transa, const TensorWrapper &B, + bool transb, int tp_size) { + auto A_tensor = convertNVTETensorCheck(A.data()); + auto B_tensor = convertNVTETensorCheck(B.data()); + int64_t A0 = A_tensor->flat_first_dim(); + int64_t A1 = A_tensor->flat_last_dim(); + int64_t B0 = B_tensor->flat_first_dim(); + int64_t B1 = B_tensor->flat_last_dim(); + + // col-major A: (M/P, K) -- tensor-parallel in the non-contracting dimension + int64_t m = (transa ? A0 : A1) * tp_size; + // col-major B: (K, N/P) -- sequence-parallel in the non-contracting dimension + int64_t n = (transb ? B1 : B0) * tp_size; + // contracting dimension not distributed + int64_t k = transa ? A1 : A0; + return {m, n, k}; +} + +CublasMpDims compute_rs_dims(const TensorWrapper &A, bool transa, const TensorWrapper &B, + bool transb, int tp_size) { + auto A_tensor = convertNVTETensorCheck(A.data()); + auto B_tensor = convertNVTETensorCheck(B.data()); + int64_t A0 = A_tensor->flat_first_dim(); + int64_t A1 = A_tensor->flat_last_dim(); + int64_t B0 = B_tensor->flat_first_dim(); + int64_t B1 = B_tensor->flat_last_dim(); + + // col-major A: (M, K/P) -- tensor-parallel in the contracting dimension + int64_t m = transa ? A0 : A1; + // col-major B: (K/P, N) -- tensor-parallel in the contracting dimension + int64_t n = transb ? B1 : B0; + // contracting dimension is distributed + int64_t k = (transa ? A1 : A0) * tp_size; + return {m, n, k}; +} + +} // namespace + /*************************************************************************************************** * Comm+GEMM Overlap Base (Pipelined / Collective) **************************************************************************************************/ +// Constructor for Userbuffers backend CommOverlapBase::CommOverlapBase(const std::vector &buffer_shape, DType buffer_dtype, int myrank, int numranks, int mylocal, int numlocal, int mynode, int numnodes, int tp_size, ExtAllgatherOp allgather_handle, @@ -335,9 +430,11 @@ void CommOverlapBase::initialize(const std::vector &buffer_shape, DType } CommOverlapBase::~CommOverlapBase() { - cudaEventDestroy(_start_d2dcopy); - cudaStreamSynchronize(_stream_comm); - cudaStreamDestroy(_stream_comm); + if (!_with_cublasmp) { + cudaEventDestroy(_start_d2dcopy); + cudaStreamSynchronize(_stream_comm); + cudaStreamDestroy(_stream_comm); + } } /* @@ -350,6 +447,8 @@ void CommOverlapBase::bulk_overlap(const TensorWrapper &A, bool transa, const Te bool accumulate, bool use_split_accumulator, CommOverlapType comm_type, TensorWrapper &rs_output, cudaStream_t stream_main) { + NVTE_CHECK(!_with_cublasmp, "Bulk overlap is not supported with cuBlasMp"); + int ori_sms = _ub_comm->sms; _ub_comm->use_ce = _use_ce; _ub_comm->sms = _num_comm_sm; @@ -407,10 +506,21 @@ void CommOverlapBase::atomic_gemm_overlap_rs(const TensorWrapper &A, bool transa TensorWrapper &workspace, bool grad, bool accumulate, bool use_split_accumulator, TensorWrapper &rs_output, cudaStream_t stream_main) { + if (_with_cublasmp) { + auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size); + // col-major GEMM compute overlapped with reduce-scatter on the output + // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P) + nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(), + pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm, + stream_main, _algo_type); + return; + } + int ori_sms = _ub_comm->sms; _ub_comm->use_ce = _use_ce; _ub_comm->sms = _num_comm_sm; _ub_comm->cga_size = _cga_size; + // Get GEMM dimensions size_t m = transa ? A.size(0) : A.size(1); size_t k = transa ? A.size(1) : A.size(0); @@ -503,6 +613,16 @@ void CommOverlapBase::split_overlap_rs(const TensorWrapper &A, bool transa, cons TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad, bool accumulate, bool use_split_accumulator, TensorWrapper &rs_output, cudaStream_t stream_main) { + if (_with_cublasmp) { + auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size); + // col-major GEMM compute overlapped with reduce-scatter on the output + // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P) + nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(), + pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm, + stream_main, _algo_type); + return; + } + // Get GEMM dimensions int ori_sms = _ub_comm->sms; _ub_comm->use_ce = _use_ce; @@ -666,6 +786,7 @@ void CommOverlapBase::bulk_overlap_external_ag(cudaStream_t send_stream, cudaStr * Comm+GEMM Overlap P2P Base (Ring-Exchange) **************************************************************************************************/ +// Constructor for Userbuffers backend CommOverlapP2PBase::CommOverlapP2PBase(const std::vector &buffer_shape, DType buffer_dtype, int myrank, int numranks, int mylocal, int numlocal, int mynode, int numnodes, int tp_size, @@ -746,16 +867,20 @@ void CommOverlapP2PBase::initialize(const std::vector &buffer_shape, DTy } CommOverlapP2PBase::~CommOverlapP2PBase() { - cudaEventDestroy(_stop_recv); - cudaEventDestroy(_stop_send); - cudaStreamDestroy(_stream_recv); - for (size_t i = 0; i < _stream_send.size(); i++) { - cudaStreamDestroy(_stream_send[i]); + if (!_with_cublasmp) { + cudaEventDestroy(_stop_recv); + cudaEventDestroy(_stop_send); + cudaStreamDestroy(_stream_recv); + for (size_t i = 0; i < _stream_send.size(); i++) { + cudaStreamDestroy(_stream_send[i]); + } } } void CommOverlapP2PBase::copy_into_buffer(cudaStream_t stream, const TensorWrapper &source, bool local_chunk, bool rowwise) { + if (_with_cublasmp) return; // cuBlasMp executes its own copy-into-buffer op + // Check element size const size_t element_size = source.element_size(); NVTE_CHECK(_ubuf.element_size() == element_size, @@ -810,6 +935,16 @@ void CommOverlapP2PBase::atomic_gemm_overlap_ag( const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb, TensorWrapper &D, TensorWrapper &bias, TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad, bool accumulate, bool use_split_accumulator, TensorWrapper &B_copy, cudaStream_t stream_main) { + if (_with_cublasmp) { + auto [m, n, k] = compute_ag_dims(A, transa, B, transb, _tp_size); + // col-major GEMM compute overlapped with all-gather on input B + // (M/P, K) x [(K, N/P) -(AG)-> (K, N)] = (M/P, N) + nvte_all_gather_gemm(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(), + pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm, + stream_main, _algo_type); + return; + } + int ori_sms = _ub_comm->sms; _ub_comm->use_ce = _use_ce; _ub_comm->sms = _num_comm_sm; @@ -912,6 +1047,16 @@ void CommOverlapP2PBase::split_overlap_ag(const TensorWrapper &A, bool transa, TensorWrapper &workspace, bool grad, bool accumulate, bool use_split_accumulator, TensorWrapper &B_copy, cudaStream_t stream_main) { + if (_with_cublasmp) { + auto [m, n, k] = compute_ag_dims(A, transa, B, transb, _tp_size); + // col-major GEMM compute overlapped with all-gather on input B + // (M/P, K) x [(K, N/P) -(AG)-> (K, N)] = (M/P, N) + nvte_all_gather_gemm(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(), + pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm, + stream_main, _algo_type); + return; + } + int ori_sms = _ub_comm->sms; _ub_comm->use_ce = _use_ce; _ub_comm->sms = _num_comm_sm; @@ -1079,6 +1224,16 @@ void CommOverlapP2PBase::atomic_gemm_overlap_rs( TensorWrapper &bias, TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad, bool accumulate, bool use_split_accumulator, TensorWrapper &rs_output, cudaStream_t stream_main) { + if (_with_cublasmp) { + auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size); + // col-major GEMM compute overlapped with reduce-scatter on the output + // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P) + nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(), + pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm, + stream_main, _algo_type); + return; + } + int ori_sms = _ub_comm->sms; _ub_comm->use_ce = _use_ce; _ub_comm->sms = _num_comm_sm; @@ -1143,6 +1298,16 @@ void CommOverlapP2PBase::split_overlap_rs(const TensorWrapper &A, bool transa, TensorWrapper &workspace, bool grad, bool accumulate, bool use_split_accumulator, TensorWrapper &rs_output, cudaStream_t stream_main) { + if (_with_cublasmp) { + auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size); + // col-major GEMM compute overlapped with reduce-scatter on the output + // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P) + nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(), + pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm, + stream_main, _algo_type); + return; + } + int ori_sms = _ub_comm->sms; _ub_comm->use_ce = _use_ce; _ub_comm->sms = _num_comm_sm; diff --git a/transformer_engine/common/common.cu b/transformer_engine/common/common.cu index 1bdd80a369b..ef757bac5fc 100644 --- a/transformer_engine/common/common.cu +++ b/transformer_engine/common/common.cu @@ -61,6 +61,7 @@ void update_tensor_scale_inv(Tensor *t, cudaStream_t stream) { namespace { constexpr size_t kThreadsPerBlock = 256; + template __global__ void __launch_bounds__(kThreadsPerBlock) memset_kernel(void *__restrict__ ptr, int value, size_t size_in_bytes) { @@ -87,48 +88,6 @@ __global__ void __launch_bounds__(kThreadsPerBlock) reinterpret_cast(ptr)[idx] = data.value; } -__global__ void __launch_bounds__(kThreadsPerBlock) - splits_to_offsets_kernel(const int64_t *__restrict__ first_dims, int64_t *__restrict__ output, - size_t num_tensors, int64_t logical_last_dim) { - __shared__ int64_t block_scan[kThreadsPerBlock]; - __shared__ int64_t chunk_prefix; - - const size_t tid = threadIdx.x; - if (tid == 0) { - output[0] = 0; - chunk_prefix = 0; - } - __syncthreads(); - - for (size_t chunk_start = 0; chunk_start < num_tensors; chunk_start += kThreadsPerBlock) { - const size_t idx = chunk_start + tid; - int64_t value = 0; - if (idx < num_tensors) { - value = first_dims[idx] * logical_last_dim; - } - block_scan[tid] = value; - __syncthreads(); - - // Inclusive scan in shared memory. - for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) { - const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0; - __syncthreads(); - block_scan[tid] += addend; - __syncthreads(); - } - - if (idx < num_tensors) { - output[idx + 1] = chunk_prefix + block_scan[tid]; - } - __syncthreads(); - - if (tid == kThreadsPerBlock - 1) { - chunk_prefix += block_scan[tid]; - } - __syncthreads(); - } -} - } // namespace #define MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, vectorizedType, stream) \ @@ -159,18 +118,6 @@ void nvte_memset(void *ptr, int value, size_t size_in_bytes, cudaStream_t stream MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, uint8_t, stream); } -void nvte_splits_to_offsets(const int64_t *first_dims, int64_t *output, size_t num_tensors, - int64_t logical_last_dim, cudaStream_t stream) { - NVTE_API_CALL(nvte_splits_to_offsets); - NVTE_CHECK(output != nullptr, "Output pointer must be allocated."); - NVTE_CHECK(num_tensors > 0, "num_tensors must be greater than 0."); - NVTE_CHECK(first_dims != nullptr, "first_dims pointer must be allocated."); - NVTE_CHECK(logical_last_dim > 0, "logical_last_dim must be greater than 0."); - - splits_to_offsets_kernel<<<1, kThreadsPerBlock, 0, stream>>>(first_dims, output, num_tensors, - logical_last_dim); - NVTE_CHECK_CUDA(cudaGetLastError()); -} } // extern "C" void checkCuDriverContext(CUstream stream) { diff --git a/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu index 675f071abac..c3495e033c8 100644 --- a/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu @@ -16,12 +16,14 @@ namespace transformer_engine { namespace fused_router { -template +template __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logits, int num_tokens, int num_experts, int topk, int score_function, float *scores, - bool *routing_map, + uint8_t *routing_map, CompType *intermediate_output) { + constexpr bool kIsBitmap = (RoutingMapFormat == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8); /*** * Section: Global Variables/Addresses init * - Each warp is responsible for one token, and has own shared memory buffer. @@ -36,10 +38,20 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi CompType *topk_logits_buf = reinterpret_cast(logits_buf + num_experts * num_token_per_block); int *topk_indices_buf = reinterpret_cast(topk_logits_buf + topk * num_token_per_block); + // Per-warp bitmap accumulator (BITMAP_U8 only). uint32 packing is bit-for-bit + // equivalent to uint8 LSB-first on little-endian devices (CUDA is always LE). + const int bitmap_words_per_warp = (num_experts + 31) / 32; + const int bitmap_row_bytes = (num_experts + 7) / 8; + uint32_t *bitmap_words_buf = nullptr; + if constexpr (kIsBitmap) { + bitmap_words_buf = reinterpret_cast(topk_indices_buf + topk * num_token_per_block); + } // The address of buffers on the current warp CompType *local_logits = logits_buf + warp_id * num_experts; CompType *topk_logits = topk_logits_buf + warp_id * topk; int *topk_indices = topk_indices_buf + warp_id * topk; + uint32_t *local_bitmap_words = + (bitmap_words_buf != nullptr) ? bitmap_words_buf + warp_id * bitmap_words_per_warp : nullptr; /*** * Section: Main Loop @@ -58,13 +70,23 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi * - Load the logits to shmem */ int pos_offset = token_offset_cur_warp * num_experts; - // Clear the routing_map (num_experts) - for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - routing_map[pos_offset + i] = false; - if (score_function == 1) { + // Clear the routing_map. In BYTEMAP mode this clears the row in global memory; + // in BITMAP_U8 mode the row is accumulated in shmem and written wholesale at + // the end of the loop, so no global clear is required. + if (score_function == 1) { + for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { intermediate_output[pos_offset + i] = -std::numeric_limits::infinity(); } } + if constexpr (!kIsBitmap) { + for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { + routing_map[pos_offset + i] = 0; + } + } else { + for (int j = lane_id; j < bitmap_words_per_warp; j += kThreadsPerWarp) { + local_bitmap_words[j] = 0u; + } + } // Load the logits to shmem for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { local_logits[i] = static_cast(logits[pos_offset + i]); @@ -127,8 +149,22 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi __syncwarp(); // Write the routing_map to the output tensor - for (int i = lane_id; i < topk; i += kThreadsPerWarp) { - routing_map[pos_offset + topk_indices[i]] = true; + if constexpr (!kIsBitmap) { + for (int i = lane_id; i < topk; i += kThreadsPerWarp) { + routing_map[pos_offset + topk_indices[i]] = 1; + } + } else { + for (int i = lane_id; i < topk; i += kThreadsPerWarp) { + int e = topk_indices[i]; + atomicOr(&local_bitmap_words[e / 32], 1u << (e % 32)); + } + __syncwarp(); + uint8_t *bitmap_row = + routing_map + static_cast(token_offset_cur_warp) * bitmap_row_bytes; + const uint8_t *local_bitmap_bytes = reinterpret_cast(local_bitmap_words); + for (int j = lane_id; j < bitmap_row_bytes; j += kThreadsPerWarp) { + bitmap_row[j] = local_bitmap_bytes[j]; + } } // Write the scores to the output tensor for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { @@ -139,32 +175,39 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi } } -template +template void fused_score_for_moe_aux_loss_forward_kernel_launcher( const DataType *logits, int num_tokens, int num_experts, int topk, int score_function, - float *scores, bool *routing_map, CompType *intermediate_output, cudaStream_t stream) { + float *scores, uint8_t *routing_map, CompType *intermediate_output, cudaStream_t stream) { // Meta data for the kernel size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block; size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // logits + topk * num_token_per_block * sizeof(CompType) // topk_logits + topk * num_token_per_block * sizeof(int); // topk_indices + if constexpr (RoutingMapFormat == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + size_t bitmap_words_per_warp = (num_experts + 31) / 32; + shared_memory_size += + bitmap_words_per_warp * num_token_per_block * sizeof(uint32_t); // bitmap accumulator + } check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float; // switch at K=16 where naive O(K^2*E) starts to dominate if (topk < 16) { - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - fused_score_for_moe_aux_loss_forward_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); - fused_score_for_moe_aux_loss_forward_kernel + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(fused_score_for_moe_aux_loss_forward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); + fused_score_for_moe_aux_loss_forward_kernel <<>>( logits, num_tokens, num_experts, topk, score_function, scores, routing_map, intermediate_output); } else { - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - fused_score_for_moe_aux_loss_forward_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); - fused_score_for_moe_aux_loss_forward_kernel + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(fused_score_for_moe_aux_loss_forward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); + fused_score_for_moe_aux_loss_forward_kernel <<>>( logits, num_tokens, num_experts, topk, score_function, scores, routing_map, intermediate_output); @@ -172,17 +215,55 @@ void fused_score_for_moe_aux_loss_forward_kernel_launcher( NVTE_CHECK_CUDA(cudaGetLastError()); } +// Build the expected routing_map shape for a given NVTERoutingMapFormat. +// BYTEMAP -> [num_tokens, num_experts] +// BITMAP_U8 -> [num_tokens, ceil(num_experts/8)] +static std::vector expected_routing_map_shape(int num_tokens, int num_experts, + NVTERoutingMapFormat format) { + const size_t t = static_cast(num_tokens); + const size_t e = static_cast(num_experts); + if (format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + return {t, (e + 7) / 8}; + } + return {t, e}; +} + void fused_score_for_moe_aux_loss_forward(const Tensor &logits, int num_tokens, int num_experts, int topk, int score_function, Tensor &scores, - Tensor &routing_map, Tensor &intermediate_output, - cudaStream_t stream) { - TE_ROUTER_PROBS_TYPE_SWITCH_ALL( - logits.data.dtype, DataType, - fused_score_for_moe_aux_loss_forward_kernel_launcher( - reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, - score_function, reinterpret_cast(scores.data.dptr), - reinterpret_cast(routing_map.data.dptr), + Tensor &routing_map, + NVTERoutingMapFormat routing_map_format, + Tensor &intermediate_output, cudaStream_t stream) { + NVTE_CHECK(num_tokens > 0 && num_experts > 0, + "num_tokens and num_experts must be positive; got num_tokens=", num_tokens, + ", num_experts=", num_experts); + const std::vector dense_shape{static_cast(num_tokens), + static_cast(num_experts)}; + NVTE_CHECK(logits.data.shape == dense_shape, "logits shape must be [num_tokens, num_experts]=[", + num_tokens, ", ", num_experts, "], got ", logits.data.shape); + NVTE_CHECK(scores.data.shape == dense_shape, "scores shape must be [num_tokens, num_experts]=[", + num_tokens, ", ", num_experts, "], got ", scores.data.shape); + NVTE_CHECK(intermediate_output.data.shape == dense_shape, + "intermediate_output shape must be [num_tokens, num_experts]=[", num_tokens, ", ", + num_experts, "], got ", intermediate_output.data.shape); + const auto routing_map_shape = + expected_routing_map_shape(num_tokens, num_experts, routing_map_format); + NVTE_CHECK(routing_map.data.shape == routing_map_shape, "routing_map shape mismatch for ", + (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 ? "BITMAP_U8" : "BYTEMAP"), + "; expected ", routing_map_shape, ", got ", routing_map.data.shape); +#define AUX_LOSS_FORWARD_DISPATCH(RoutingMapFormatVal) \ + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ + logits.data.dtype, DataType, \ + fused_score_for_moe_aux_loss_forward_kernel_launcher( \ + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, \ + score_function, reinterpret_cast(scores.data.dptr), \ + reinterpret_cast(routing_map.data.dptr), \ reinterpret_cast(intermediate_output.data.dptr), stream);); + if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + AUX_LOSS_FORWARD_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) + } else { + AUX_LOSS_FORWARD_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BYTEMAP) + } +#undef AUX_LOSS_FORWARD_DISPATCH } template @@ -343,17 +424,31 @@ void fused_score_for_moe_aux_loss_backward(const Tensor &intermediate_output, } // namespace fused_router } // namespace transformer_engine +void nvte_fused_score_for_moe_aux_loss_forward_v2(const NVTETensor logits, int num_tokens, + int num_experts, int topk, int score_function, + NVTETensor scores, NVTETensor routing_map, + NVTERoutingMapFormat routing_map_format, + const NVTETensor intermediate_output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_forward_v2); + using namespace transformer_engine; + fused_router::fused_score_for_moe_aux_loss_forward( + *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, score_function, + *convertNVTETensorCheck(scores), *convertNVTETensorCheck(routing_map), routing_map_format, + *convertNVTETensorCheck(intermediate_output), stream); +} + +// Deprecated V1 entry point: forwards to the V2 above with the BYTEMAP layout. +// Kept for ABI compatibility with external C API consumers. void nvte_fused_score_for_moe_aux_loss_forward(const NVTETensor logits, int num_tokens, int num_experts, int topk, int score_function, - NVTETensor scores, const NVTETensor routing_map, + NVTETensor scores, NVTETensor routing_map, const NVTETensor intermediate_output, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_forward); - using namespace transformer_engine; - fused_router::fused_score_for_moe_aux_loss_forward( - *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, score_function, - *convertNVTETensorCheck(scores), *convertNVTETensorCheck(routing_map), - *convertNVTETensorCheck(intermediate_output), stream); + nvte_fused_score_for_moe_aux_loss_forward_v2( + logits, num_tokens, num_experts, topk, score_function, scores, routing_map, + NVTE_ROUTING_MAP_FORMAT_BYTEMAP, intermediate_output, stream); } void nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_output, diff --git a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu index 9b8d8b92998..01e6ab75f66 100644 --- a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu +++ b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu @@ -15,12 +15,14 @@ namespace transformer_engine { namespace fused_router { -template +template __global__ void fused_topk_with_score_function_forward_kernel( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, - const BiasType *expert_bias, DataType *probs, bool *routing_map, + const BiasType *expert_bias, DataType *probs, uint8_t *routing_map, CompType *intermediate_output) { + constexpr bool kIsBitmap = (RoutingMapFormat == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8); /*** * Section: Global Variables/Addresses init * - Each warp is responsible for one token, and has own shared memory buffer. @@ -42,12 +44,22 @@ __global__ void fused_topk_with_score_function_forward_kernel( } else { topk_indices_buf = reinterpret_cast(topk_scores_buf + topk * num_token_per_block); } + // Per-warp bitmap accumulator (BITMAP_U8 only). uint32 packing is bit-for-bit + // equivalent to uint8 LSB-first on little-endian devices (CUDA is always LE). + const int bitmap_words_per_warp = (num_experts + 31) / 32; + const int bitmap_row_bytes = (num_experts + 7) / 8; + uint32_t *bitmap_words_buf = nullptr; + if constexpr (kIsBitmap) { + bitmap_words_buf = reinterpret_cast(topk_indices_buf + topk * num_token_per_block); + } // The address of buffers on the current warp CompType *scores = scores_buf + warp_id * num_experts; CompType *topk_scores = topk_scores_buf + warp_id * topk; CompType *masked_scores = masked_scores_buf + warp_id * num_experts; CompType *group_scores = group_scores_buf + warp_id * num_groups; int *topk_indices = topk_indices_buf + warp_id * topk; + uint32_t *local_bitmap_words = + (bitmap_words_buf != nullptr) ? bitmap_words_buf + warp_id * bitmap_words_per_warp : nullptr; /*** * Section: Main Loop @@ -66,14 +78,23 @@ __global__ void fused_topk_with_score_function_forward_kernel( * - Load the logits to shmem */ int pos_offset = token_offset_cur_warp * num_experts; - // Clear the probs/routing_map (num_experts) + // BITMAP_U8 accumulates the row in shmem and writes it wholesale at the end + // of the loop; only BYTEMAP needs a global clear here. for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { probs[pos_offset + i] = 0.0; - routing_map[pos_offset + i] = false; if (score_function == 1) { intermediate_output[pos_offset + i] = -std::numeric_limits::infinity(); } } + if constexpr (!kIsBitmap) { + for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { + routing_map[pos_offset + i] = 0; + } + } else { + for (int j = lane_id; j < bitmap_words_per_warp; j += kThreadsPerWarp) { + local_bitmap_words[j] = 0u; + } + } // Load the logits to shmem for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { scores[i] = logits[pos_offset + i]; @@ -229,21 +250,37 @@ __global__ void fused_topk_with_score_function_forward_kernel( } // Write the probs/routing_map to the output tensor - for (int i = lane_id; i < topk; i += kThreadsPerWarp) { - routing_map[pos_offset + topk_indices[i]] = true; - probs[pos_offset + topk_indices[i]] = scaling_factor * topk_scores[i]; + if constexpr (!kIsBitmap) { + for (int i = lane_id; i < topk; i += kThreadsPerWarp) { + routing_map[pos_offset + topk_indices[i]] = 1; + probs[pos_offset + topk_indices[i]] = scaling_factor * topk_scores[i]; + } + } else { + // shmem atomicOr handles same-word collisions across the topk lanes. + for (int i = lane_id; i < topk; i += kThreadsPerWarp) { + int e = topk_indices[i]; + atomicOr(&local_bitmap_words[e / 32], 1u << (e % 32)); + probs[pos_offset + e] = scaling_factor * topk_scores[i]; + } + __syncwarp(); + uint8_t *bitmap_row = + routing_map + static_cast(token_offset_cur_warp) * bitmap_row_bytes; + const uint8_t *local_bitmap_bytes = reinterpret_cast(local_bitmap_words); + for (int j = lane_id; j < bitmap_row_bytes; j += kThreadsPerWarp) { + bitmap_row[j] = local_bitmap_bytes[j]; + } } __threadfence_block(); __syncwarp(); } } -template +template void fused_topk_with_score_function_forward_kernel_launcher( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, - const BiasType *expert_bias, DataType *probs, bool *routing_map, CompType *intermediate_output, - cudaStream_t stream) { + const BiasType *expert_bias, DataType *probs, uint8_t *routing_map, + CompType *intermediate_output, cudaStream_t stream) { size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block; size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // scores @@ -253,22 +290,31 @@ void fused_topk_with_score_function_forward_kernel_launcher( shared_memory_size += num_groups * num_token_per_block * sizeof(CompType); // group_scores shared_memory_size += num_experts * num_token_per_block * sizeof(CompType); // maksed_scores } + if constexpr (RoutingMapFormat == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + size_t bitmap_words_per_warp = (num_experts + 31) / 32; + shared_memory_size += + bitmap_words_per_warp * num_token_per_block * sizeof(uint32_t); // bitmap accumulator + } check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float; // switch at K=16 where naive O(K^2*E) starts to dominate if (topk < 16) { NVTE_CHECK_CUDA(cudaFuncSetAttribute( - fused_topk_with_score_function_forward_kernel, + fused_topk_with_score_function_forward_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); - fused_topk_with_score_function_forward_kernel + fused_topk_with_score_function_forward_kernel <<>>( logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output); } else { NVTE_CHECK_CUDA(cudaFuncSetAttribute( - fused_topk_with_score_function_forward_kernel, + fused_topk_with_score_function_forward_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); - fused_topk_with_score_function_forward_kernel + fused_topk_with_score_function_forward_kernel <<>>( logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output); @@ -276,43 +322,89 @@ void fused_topk_with_score_function_forward_kernel_launcher( NVTE_CHECK_CUDA(cudaGetLastError()); } +// Build the expected routing_map shape for a given NVTERoutingMapFormat. +// BYTEMAP -> [num_tokens, num_experts] +// BITMAP_U8 -> [num_tokens, ceil(num_experts/8)] +static std::vector expected_routing_map_shape(int num_tokens, int num_experts, + NVTERoutingMapFormat format) { + const size_t t = static_cast(num_tokens); + const size_t e = static_cast(num_experts); + if (format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + return {t, (e + 7) / 8}; + } + return {t, e}; +} + void fused_topk_with_score_function_forward(const Tensor logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, const Tensor expert_bias, Tensor probs, Tensor routing_map, + NVTERoutingMapFormat routing_map_format, Tensor intermediate_output, cudaStream_t stream) { - TE_ROUTER_PROBS_TYPE_SWITCH_ALL( - logits.data.dtype, DataType, - if (expert_bias.has_data()) { - TE_ROUTER_PROBS_TYPE_SWITCH_ALL( - expert_bias.data.dtype, BiasType, - fused_topk_with_score_function_forward_kernel_launcher( - reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, - use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, - reinterpret_cast(expert_bias.data.dptr), - reinterpret_cast(probs.data.dptr), - reinterpret_cast(routing_map.data.dptr), - reinterpret_cast(intermediate_output.data.dptr), stream);); - } else { - fused_topk_with_score_function_forward_kernel_launcher( - reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, - use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, nullptr, - reinterpret_cast(probs.data.dptr), - reinterpret_cast(routing_map.data.dptr), - reinterpret_cast(intermediate_output.data.dptr), stream); + NVTE_CHECK(num_tokens > 0 && num_experts > 0, + "num_tokens and num_experts must be positive; got num_tokens=", num_tokens, + ", num_experts=", num_experts); + const std::vector dense_shape{static_cast(num_tokens), + static_cast(num_experts)}; + NVTE_CHECK(logits.data.shape == dense_shape, "logits shape must be [num_tokens, num_experts]=[", + num_tokens, ", ", num_experts, "], got ", logits.data.shape); + NVTE_CHECK(probs.data.shape == dense_shape, "probs shape must be [num_tokens, num_experts]=[", + num_tokens, ", ", num_experts, "], got ", probs.data.shape); + NVTE_CHECK(intermediate_output.data.shape == dense_shape, + "intermediate_output shape must be [num_tokens, num_experts]=[", num_tokens, ", ", + num_experts, "], got ", intermediate_output.data.shape); + const auto routing_map_shape = + expected_routing_map_shape(num_tokens, num_experts, routing_map_format); + NVTE_CHECK(routing_map.data.shape == routing_map_shape, "routing_map shape mismatch for ", + (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 ? "BITMAP_U8" : "BYTEMAP"), + "; expected ", routing_map_shape, ", got ", routing_map.data.shape); + if (expert_bias.has_data()) { + NVTE_CHECK(expert_bias.data.shape == std::vector{static_cast(num_experts)}, + "expert_bias shape must be [num_experts]=[", num_experts, "], got ", + expert_bias.data.shape); + } +#define ROUTER_FORWARD_DISPATCH(RoutingMapFormatVal) \ + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ + logits.data.dtype, DataType, \ + if (expert_bias.has_data()) { \ + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ + expert_bias.data.dtype, BiasType, \ + fused_topk_with_score_function_forward_kernel_launcher( \ + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, \ + use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, \ + reinterpret_cast(expert_bias.data.dptr), \ + reinterpret_cast(probs.data.dptr), \ + reinterpret_cast(routing_map.data.dptr), \ + reinterpret_cast(intermediate_output.data.dptr), stream);); \ + } else { \ + fused_topk_with_score_function_forward_kernel_launcher( \ + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, \ + use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, nullptr, \ + reinterpret_cast(probs.data.dptr), \ + reinterpret_cast(routing_map.data.dptr), \ + reinterpret_cast(intermediate_output.data.dptr), stream); \ }); + if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + ROUTER_FORWARD_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) + } else { + ROUTER_FORWARD_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BYTEMAP) + } +#undef ROUTER_FORWARD_DISPATCH } -template +template __global__ void fused_topk_with_score_function_backward_kernel( // Inputs tensor - const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs, + const uint8_t *routing_map, const CompType *intermediate_output, const DataType *grad_probs, // Other parameters int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor, int score_function, // Output tensor DataType *grad_logits) { + constexpr bool kIsBitmap = (RoutingMapFormat == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8); /*** * Section: Global Variables/Addresses init * - Each warp is responsible for one token, and has own shared memory buffer. @@ -352,11 +444,19 @@ __global__ void fused_topk_with_score_function_backward_kernel( * - Load the dgrad/output_from_fwd to shmem */ int pos_offset = token_offset_cur_warp * num_experts; - // Load the dgrad/output_from_fwd to shmem + // Load the dgrad/output_from_fwd to shmem. The routing_map source layout + // depends on the RoutingMapFormat template parameter (see NVTERoutingMapFormat). + const int bitmap_row_bytes = (num_experts + 7) / 8; + const uint8_t *bitmap_row = + routing_map + static_cast(token_offset_cur_warp) * bitmap_row_bytes; for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { local_grad[i] = grad_probs[pos_offset + i]; local_act_from_fwd[i] = intermediate_output[pos_offset + i]; - local_routing_map[i] = routing_map[pos_offset + i]; + if constexpr (!kIsBitmap) { + local_routing_map[i] = routing_map[pos_offset + i] != 0; + } else { + local_routing_map[i] = (bitmap_row[i / 8] >> (i % 8)) & 1u; + } } __threadfence_block(); __syncwarp(); @@ -474,9 +574,9 @@ __global__ void fused_topk_with_score_function_backward_kernel( } } -template +template void fused_topk_with_score_function_backward_kernel_launcher( - const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs, + const uint8_t *routing_map, const CompType *intermediate_output, const DataType *grad_probs, int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor, int score_function, DataType *grad_logits, cudaStream_t stream) { // Meta data for the kernel @@ -488,10 +588,10 @@ void fused_topk_with_score_function_backward_kernel_launcher( + num_experts * num_token_per_block * sizeof(CompType) // comp_buf + num_experts * num_token_per_block * sizeof(bool); // routing_map check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); - NVTE_CHECK_CUDA(cudaFuncSetAttribute(fused_topk_with_score_function_backward_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - shared_memory_size)); - fused_topk_with_score_function_backward_kernel + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + fused_topk_with_score_function_backward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); + fused_topk_with_score_function_backward_kernel <<>>( routing_map, intermediate_output, grad_probs, num_tokens, num_experts, topk, use_pre_softmax, scaling_factor, score_function, grad_logits); @@ -499,38 +599,98 @@ void fused_topk_with_score_function_backward_kernel_launcher( } void fused_topk_with_score_function_backward(const Tensor &routing_map, + NVTERoutingMapFormat routing_map_format, const Tensor &intermediate_output, const Tensor &grad_probs, int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor, int score_function, Tensor &grad_logits, cudaStream_t stream) { - TE_ROUTER_PROBS_TYPE_SWITCH_ALL( - grad_logits.data.dtype, DataType, - fused_topk_with_score_function_backward_kernel_launcher( - reinterpret_cast(routing_map.data.dptr), - reinterpret_cast(intermediate_output.data.dptr), - reinterpret_cast(grad_probs.data.dptr), num_tokens, num_experts, topk, - use_pre_softmax, scaling_factor, score_function, + NVTE_CHECK(num_tokens > 0 && num_experts > 0, + "num_tokens and num_experts must be positive; got num_tokens=", num_tokens, + ", num_experts=", num_experts); + const std::vector dense_shape{static_cast(num_tokens), + static_cast(num_experts)}; + NVTE_CHECK(intermediate_output.data.shape == dense_shape, + "intermediate_output shape must be [num_tokens, num_experts]=[", num_tokens, ", ", + num_experts, "], got ", intermediate_output.data.shape); + NVTE_CHECK(grad_probs.data.shape == dense_shape, + "grad_probs shape must be [num_tokens, num_experts]=[", num_tokens, ", ", num_experts, + "], got ", grad_probs.data.shape); + NVTE_CHECK(grad_logits.data.shape == dense_shape, + "grad_logits shape must be [num_tokens, num_experts]=[", num_tokens, ", ", num_experts, + "], got ", grad_logits.data.shape); + const auto routing_map_shape = + expected_routing_map_shape(num_tokens, num_experts, routing_map_format); + NVTE_CHECK(routing_map.data.shape == routing_map_shape, "routing_map shape mismatch for ", + (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 ? "BITMAP_U8" : "BYTEMAP"), + "; expected ", routing_map_shape, ", got ", routing_map.data.shape); +#define ROUTER_BACKWARD_DISPATCH(RoutingMapFormatVal) \ + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ + grad_logits.data.dtype, DataType, \ + fused_topk_with_score_function_backward_kernel_launcher( \ + reinterpret_cast(routing_map.data.dptr), \ + reinterpret_cast(intermediate_output.data.dptr), \ + reinterpret_cast(grad_probs.data.dptr), num_tokens, num_experts, topk, \ + use_pre_softmax, scaling_factor, score_function, \ reinterpret_cast(grad_logits.data.dptr), stream);); + if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + ROUTER_BACKWARD_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) + } else { + ROUTER_BACKWARD_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BYTEMAP) + } +#undef ROUTER_BACKWARD_DISPATCH } } // namespace fused_router } // namespace transformer_engine -void nvte_fused_topk_with_score_function_forward( +void nvte_fused_topk_with_score_function_forward_v2( const NVTETensor logits, int num_tokens, int num_experts, int topk, int use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, - NVTETensor intermediate_output, cudaStream_t stream) { - NVTE_API_CALL(nvte_fused_topk_with_score_function_forward); + NVTERoutingMapFormat routing_map_format, NVTETensor intermediate_output, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_topk_with_score_function_forward_v2); using namespace transformer_engine; fused_router::fused_topk_with_score_function_forward( *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, static_cast(use_pre_softmax), num_groups, group_topk, scaling_factor, score_function, *convertNVTETensorCheck(expert_bias), *convertNVTETensorCheck(probs), - *convertNVTETensorCheck(routing_map), *convertNVTETensorCheck(intermediate_output), stream); + *convertNVTETensorCheck(routing_map), routing_map_format, + *convertNVTETensorCheck(intermediate_output), stream); +} + +// Deprecated V1 entry point: forwards to the V2 above with the BYTEMAP layout. +// Kept for ABI compatibility with external C API consumers. +void nvte_fused_topk_with_score_function_forward( + const NVTETensor logits, int num_tokens, int num_experts, int topk, int use_pre_softmax, + int num_groups, int group_topk, float scaling_factor, int score_function, + const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, + NVTETensor intermediate_output, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_topk_with_score_function_forward); + nvte_fused_topk_with_score_function_forward_v2( + logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, + scaling_factor, score_function, expert_bias, probs, routing_map, + NVTE_ROUTING_MAP_FORMAT_BYTEMAP, intermediate_output, stream); } +void nvte_fused_topk_with_score_function_backward_v2(const NVTETensor routing_map, + NVTERoutingMapFormat routing_map_format, + const NVTETensor intermediate_output, + const NVTETensor grad_probs, int num_tokens, + int num_experts, int topk, int use_pre_softmax, + float scaling_factor, int score_function, + NVTETensor grad_logits, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_topk_with_score_function_backward_v2); + using namespace transformer_engine; + fused_router::fused_topk_with_score_function_backward( + *convertNVTETensorCheck(routing_map), routing_map_format, + *convertNVTETensorCheck(intermediate_output), *convertNVTETensorCheck(grad_probs), num_tokens, + num_experts, topk, static_cast(use_pre_softmax), scaling_factor, score_function, + *convertNVTETensorCheck(grad_logits), stream); +} + +// Deprecated V1 entry point: forwards to the V2 above with the BYTEMAP layout. +// Kept for ABI compatibility with external C API consumers. void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map, const NVTETensor intermediate_output, const NVTETensor grad_probs, int num_tokens, @@ -538,10 +698,7 @@ void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map, float scaling_factor, int score_function, NVTETensor grad_logits, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_topk_with_score_function_backward); - using namespace transformer_engine; - fused_router::fused_topk_with_score_function_backward( - *convertNVTETensorCheck(routing_map), *convertNVTETensorCheck(intermediate_output), - *convertNVTETensorCheck(grad_probs), num_tokens, num_experts, topk, - static_cast(use_pre_softmax), scaling_factor, score_function, - *convertNVTETensorCheck(grad_logits), stream); + nvte_fused_topk_with_score_function_backward_v2( + routing_map, NVTE_ROUTING_MAP_FORMAT_BYTEMAP, intermediate_output, grad_probs, num_tokens, + num_experts, topk, use_pre_softmax, scaling_factor, score_function, grad_logits, stream); } diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index e59e9c00c97..a0529c80c03 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -1056,6 +1056,10 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor cudaStream_t stream) { NVTE_API_CALL(nvte_multi_tensor_gemm); + if (num_gemms <= 0) { + return; + } + const int current_device = transformer_engine::cuda::current_device(); const bool is_hopper = (transformer_engine::cuda::sm_arch(current_device) == 90); const bool use_cutlass = transformer_engine::getenv("NVTE_USE_CUTLASS_GROUPED_GEMM", false); diff --git a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h index 6307eab14c9..b1a7447f2a6 100644 --- a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h +++ b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h @@ -9,6 +9,8 @@ #include #include +#include +#include #include #include @@ -17,6 +19,12 @@ #define NVTE_COMM_OVERLAP_MAX_STREAMS 3 +/* \brief Check if TE is built with cuBLASMp. + * + * \return True if TE is built with cuBLASMp. + */ +bool nvte_built_with_cublasmp(); + namespace transformer_engine { /* \brief Check if Userbufers bootstraps with direct calls to MPI collectives. @@ -59,6 +67,10 @@ class CommOverlapCore { bool _atomic_gemm{false}; bool _is_p2p{false}; + bool _with_cublasmp{false}; + NVTECommGemmCtx *_cublasmp_ctx{nullptr}; + NVTECommGemmAlgoType _algo_type = kNVTECommGemmAlgoDefault; + TensorWrapper _ubuf; TensorWrapper _counter; float *_ubuf_scale_inv; @@ -75,12 +87,17 @@ class CommOverlapCore { public: CommOverlapCore() {} // dummy constructor for exposing type to Python + // Constructor for Userbuffers backend CommOverlapCore(int myrank, int numranks, int mylocal, int numlocal, int mynode, int numnodes, int tp_size, ExtAllgatherOp allgather_handle, ExtBarrierOp barrier_handle, int num_splits, int num_max_streams, int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool use_ce, bool atomic_gemm); + // Constructor for cuBLASMp backend + CommOverlapCore(ncclComm_t nccl_comm_ptr, int tp_rank, int tp_size, int num_comm_sm, bool is_p2p, + bool atomic_gemm); + virtual ~CommOverlapCore(); void *get_ubuf_dptr() { return _ubuf.dptr(); } @@ -109,6 +126,8 @@ class CommOverlapCore { bool is_fp8_ubuf() { return _ubuf.element_size() == 1; } + bool with_cublasmp() { return _with_cublasmp; } + virtual void bulk_overlap(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb, TensorWrapper &D, TensorWrapper &bias, TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad, @@ -169,6 +188,7 @@ class CommOverlapBase : public CommOverlapCore { public: CommOverlapBase() {} // dummy constructor for exposing type to Python + // Constructor for Userbuffers backend CommOverlapBase(const std::vector &buffer_shape, DType buffer_dtype, int myrank, int numranks, int mylocal, int numlocal, int mynode, int numnodes, int tp_size, ExtAllgatherOp allgather_handle, ExtBarrierOp barrier_handle, int num_splits = 3, @@ -177,6 +197,11 @@ class CommOverlapBase : public CommOverlapCore { bool set_sm_margin = true, bool atomic_gemm = false, bool rs_overlap_first_gemm = false); + // Constructor for cuBLASMp backend + CommOverlapBase(ncclComm_t nccl_comm_ptr, int tp_rank, int tp_size, int num_comm_sm = 16, + bool atomic_gemm = false) + : CommOverlapCore(nccl_comm_ptr, tp_rank, tp_size, num_comm_sm, false, atomic_gemm) {} + virtual ~CommOverlapBase(); /* @@ -249,6 +274,7 @@ class CommOverlapP2PBase : public CommOverlapCore { public: CommOverlapP2PBase() {} // dummy constructor for exposing type to Python + // Constructor for Userbuffers backend CommOverlapP2PBase(const std::vector &buffer_shape, DType buffer_dtype, int myrank, int numranks, int mylocal, int numlocal, int mynode, int numnodes, int tp_size, ExtAllgatherOp allgather_handle, ExtBarrierOp barrier_handle, @@ -257,6 +283,11 @@ class CommOverlapP2PBase : public CommOverlapCore { int num_comm_sm = 1, bool set_sm_margin = false, bool use_ce = true, bool atomic_gemm = false, bool aggregate = false); + // Constructor for cuBLASMp backend + CommOverlapP2PBase(ncclComm_t nccl_comm_ptr, int tp_rank, int tp_size, int num_comm_sm = 1, + bool atomic_gemm = false) + : CommOverlapCore(nccl_comm_ptr, tp_rank, tp_size, num_comm_sm, true, atomic_gemm) {} + virtual ~CommOverlapP2PBase(); void copy_into_buffer(cudaStream_t stream, const TensorWrapper &source, bool local_chunk, diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 794880d3249..6cee10bd391 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -13,7 +13,26 @@ extern "C" { #endif -/*! \brief Apply topk + softmax/sigmoid to the input tensor. Grouped topk is supported. +/*! \brief Output format of the routing_map tensor. + * + * BYTEMAP — bool/uint8 tensor of shape [num_tokens, num_experts]; one byte + * per (token, expert) pair, 0 or 1. + * BITMAP_U8 — uint8 tensor of shape [num_tokens, ceil(num_experts/8)]; bit + * (e % 8) of byte (e / 8) of row t is 1 iff token t routes to + * expert e (little-endian / LSB-first packing along the expert + * axis). + */ +typedef enum { + NVTE_ROUTING_MAP_FORMAT_BYTEMAP = 0, + NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 = 1, +} NVTERoutingMapFormat; + +/*! \brief Apply topk + softmax/sigmoid to the input tensor. Grouped topk is supported (deprecated). + * + * \deprecated This function has been deprecated in favor of + * nvte_fused_topk_with_score_function_forward_v2, which adds support + * for the NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 routing_map layout. This + * entry point assumes NVTE_ROUTING_MAP_FORMAT_BYTEMAP. * * \param[in] logits Logits from the gating GEMM. * \param[in] num_tokens Number of tokens. @@ -26,7 +45,7 @@ extern "C" { * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. * \param[in] expert_bias Expert bias. (Used at the sigmoid/sqrtsoftplus cases) * \param[out] probs Output tensor for probabilities. - * \param[out] routing_map Output tensor for routing map. + * \param[out] routing_map Output tensor for routing map (BYTEMAP layout). * \param[out] intermediate_output Output tensor for intermediate output. (Softmax/sigmoid output) * \param[in] stream CUDA stream used for the operation. */ @@ -36,9 +55,40 @@ void nvte_fused_topk_with_score_function_forward( const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, NVTETensor intermediate_output, cudaStream_t stream); -/*! \brief Backward pass for fused topk + softmax/sigmoid. +/*! \brief Apply topk + softmax/sigmoid to the input tensor. Grouped topk is supported. + * + * \param[in] logits Logits from the gating GEMM. + * \param[in] num_tokens Number of tokens. + * \param[in] num_experts Number of experts. + * \param[in] topk Topk value. + * \param[in] use_pre_softmax Whether to use softmax before topk. + * \param[in] num_groups Number of groups in grouped topk. + * \param[in] group_topk Grouped topk value. + * \param[in] scaling_factor Scaling factor. + * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. + * \param[in] expert_bias Expert bias. (Used at the sigmoid/sqrtsoftplus cases) + * \param[out] probs Output tensor for probabilities. + * \param[out] routing_map Output tensor for routing map. Shape depends on + * routing_map_format (see NVTERoutingMapFormat). + * \param[in] routing_map_format NVTERoutingMapFormat value selecting the routing_map + * output layout. The caller is responsible for + * allocating routing_map with the matching shape. + * \param[out] intermediate_output Output tensor for intermediate output. (Softmax/sigmoid output) + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_topk_with_score_function_forward_v2( + const NVTETensor logits, int num_tokens, int num_experts, int topk, int use_pre_softmax, + int num_groups, int group_topk, float scaling_factor, int score_function, + const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, + NVTERoutingMapFormat routing_map_format, NVTETensor intermediate_output, cudaStream_t stream); + +/*! \brief Backward pass for fused topk + softmax/sigmoid (deprecated). + * + * \deprecated This function has been deprecated in favor of + * nvte_fused_topk_with_score_function_backward_v2. This entry point + * assumes NVTE_ROUTING_MAP_FORMAT_BYTEMAP. * - * \param[in] routing_map Routing map. + * \param[in] routing_map Routing map (BYTEMAP layout). * \param[in] intermediate_output Intermediate output from the forward pass. (Softmax/sigmoid output) * \param[in] grad_probs Gradient of probs. * \param[in] num_tokens Number of tokens. @@ -57,7 +107,34 @@ void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map, float scaling_factor, int score_function, NVTETensor grad_logits, cudaStream_t stream); -/*! \brief Forward pass for computing scores/routing map for auxiliary loss. +/*! \brief Backward pass for fused topk + softmax/sigmoid. + * + * \param[in] routing_map Routing map (same layout as produced by forward). + * \param[in] routing_map_format NVTERoutingMapFormat value matching the layout of routing_map. + * \param[in] intermediate_output Intermediate output from the forward pass. (Softmax/sigmoid output) + * \param[in] grad_probs Gradient of probs. + * \param[in] num_tokens Number of tokens. + * \param[in] num_experts Number of experts. + * \param[in] topk Topk value. + * \param[in] use_pre_softmax Whether to use softmax before topk. + * \param[in] scaling_factor Scaling factor. + * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. + * \param[out] grad_logits Gradient of logits. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_topk_with_score_function_backward_v2(const NVTETensor routing_map, + NVTERoutingMapFormat routing_map_format, + const NVTETensor intermediate_output, + const NVTETensor grad_probs, int num_tokens, + int num_experts, int topk, int use_pre_softmax, + float scaling_factor, int score_function, + NVTETensor grad_logits, cudaStream_t stream); + +/*! \brief Forward pass for computing scores/routing map for auxiliary loss (deprecated). + * + * \deprecated This function has been deprecated in favor of + * nvte_fused_score_for_moe_aux_loss_forward_v2. This entry point + * assumes NVTE_ROUTING_MAP_FORMAT_BYTEMAP. * * \param[in] logits Logits from the gating GEMM. * \param[in] num_tokens Number of tokens. @@ -65,16 +142,38 @@ void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map, * \param[in] topk Topk value. * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. * \param[out] scores Output tensor for scores. - * \param[in] routing_map Routing map. + * \param[out] routing_map Output tensor for routing map (BYTEMAP layout). * \param[in] intermediate_output Intermediate output from the forward pass. (Softmax/sigmoid output) * \param[in] stream CUDA stream used for the operation. */ void nvte_fused_score_for_moe_aux_loss_forward(const NVTETensor logits, int num_tokens, int num_experts, int topk, int score_function, - NVTETensor scores, const NVTETensor routing_map, + NVTETensor scores, NVTETensor routing_map, const NVTETensor intermediate_output, cudaStream_t stream); +/*! \brief Forward pass for computing scores/routing map for auxiliary loss. + * + * \param[in] logits Logits from the gating GEMM. + * \param[in] num_tokens Number of tokens. + * \param[in] num_experts Number of experts. + * \param[in] topk Topk value. + * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. + * \param[out] scores Output tensor for scores. + * \param[out] routing_map Output tensor for routing map. Shape depends on + * routing_map_format (see NVTERoutingMapFormat). + * \param[in] routing_map_format NVTERoutingMapFormat value selecting the routing_map + * output layout. + * \param[in] intermediate_output Intermediate output from the forward pass. (Softmax/sigmoid output) + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_score_for_moe_aux_loss_forward_v2(const NVTETensor logits, int num_tokens, + int num_experts, int topk, int score_function, + NVTETensor scores, NVTETensor routing_map, + NVTERoutingMapFormat routing_map_format, + const NVTETensor intermediate_output, + cudaStream_t stream); + /*! \brief Backward pass for computing scores/routing map for auxiliary loss. * * \param[in] intermediate_output Intermediate output from the forward pass. (Softmax/sigmoid output) diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index f675b2f5352..c32a561fb7d 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -497,17 +497,36 @@ void nvte_memset(void *ptr, int value, size_t size_in_bytes, cudaStream_t stream * * Computes: * output[0] = 0 - * output[i + 1] = sum_{j=0..i}(first_dims[j] * logical_last_dim) - * for i in [0, num_tensors - 1]. + * output[i + 1] = sum_{j=0..i}(split_sizes[j] * stride) + * for i in [0, num_splits - 1]. * - * \param[in] first_dims Pointer to device int64 array of size num_tensors. - * \param[out] output Pointer to device int64 array of size num_tensors + 1. - * \param[in] num_tensors Number of entries in first_dims. - * \param[in] logical_last_dim Scale factor applied to each first_dims entry. + * \param[in] split_sizes Pointer to device int64 array of size num_splits. + * \param[out] output Pointer to device int64 array of size num_splits + 1. + * \param[in] num_splits Number of entries in split_sizes. + * \param[in] stride Scale factor applied to each split_sizes entry. * \param[in] stream CUDA stream to use for the operation. */ -void nvte_splits_to_offsets(const int64_t *first_dims, int64_t *output, size_t num_tensors, - int64_t logical_last_dim, cudaStream_t stream); +void nvte_splits_to_offsets(const int64_t *split_sizes, int64_t *output, size_t num_splits, + int64_t stride, cudaStream_t stream); + +/*! \brief Compute multiple scaled prefix-sum offsets for grouped tensors. + * + * Computes a prefix-sum over the values in split_sizes, and for each + * output multiplies the prefix-sum by a stride. Inputs and outputs + * can be any combination of int32 and int64 tensors. + * + * \param[in] split_sizes Device int32/int64 split sizes with shape [N]. + * \param[out] outputs Array of int32/int64 1D output tensors, one per scan. + * \param[in] strides Per-output scale factor. Length num_outputs. + * \param[in] include_leading_zero Per-output flag: 0 if outputs[i] has length N + * (inclusive scan), nonzero if outputs[i] has length N + 1 (inclusive + * scan prepended with zero). Length num_outputs. + * \param[in] num_outputs Number of output tensors. + * \param[in] stream CUDA stream to use for the operation. + */ +void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, + const int64_t *strides, const int *include_leading_zero, + size_t num_outputs, cudaStream_t stream); /*! \brief TE Grouped Tensor type * diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 51969e10e3d..38b526360e6 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -1049,6 +1049,11 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s " column-wise scaling factors, but got shape=", output->columnwise_scale_inv.shape, "."); } + // Return early if tensor has no entries + if (m == 0 || k == 0) { + return; + } + // Choose swizzle implementation bool rowwise_swizzle{false}, columnwise_swizzle{false}; switch (scaling_mode) { diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index ed48fe4d616..f7ffb5ad8db 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -119,11 +119,15 @@ pybind11::module_local()) \ .def(py::init([]() { return new transformer_engine::CommOverlapCore(); }), \ py::call_guard()) \ + .def("get_tp_size", &transformer_engine::CommOverlapCore::get_tp_size, \ + py::call_guard()) \ .def("is_atomic_gemm", &transformer_engine::CommOverlapCore::is_atomic_gemm, \ py::call_guard()) \ .def("is_p2p_overlap", &transformer_engine::CommOverlapCore::is_p2p_overlap, \ py::call_guard()) \ .def("is_fp8_ubuf", &transformer_engine::CommOverlapCore::is_fp8_ubuf, \ + py::call_guard()) \ + .def("with_cublasmp", &transformer_engine::CommOverlapCore::with_cublasmp, \ py::call_guard()); \ py::class_, \ @@ -147,6 +151,8 @@ }, \ py::call_guard(), py::arg("device_id") = -1); \ m.def("ubuf_built_with_mpi", &transformer_engine::ubuf_built_with_mpi, \ + py::call_guard()); \ + m.def("nvte_built_with_cublasmp", &nvte_built_with_cublasmp, \ py::call_guard()); #endif diff --git a/transformer_engine/common/util/splits_to_offsets.cu b/transformer_engine/common/util/splits_to_offsets.cu new file mode 100644 index 00000000000..721f88c0741 --- /dev/null +++ b/transformer_engine/common/util/splits_to_offsets.cu @@ -0,0 +1,212 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include +#include + +#include "../common.h" +#include "../util/logging.h" +#include "../utils.cuh" + +namespace transformer_engine::splits_to_offsets { + +namespace { + +constexpr size_t kThreadsPerBlock = 256; + +struct KernelArgs { + static constexpr size_t kMaxNumOutputs = 8; + const void *split_sizes = nullptr; + DType split_sizes_dtype = DType::kNumTypes; + size_t num_splits = 0; + void *outputs[kMaxNumOutputs] = {}; + DType outputs_dtype[kMaxNumOutputs] = {}; + int64_t strides[kMaxNumOutputs] = {}; + bool include_leading_zero[kMaxNumOutputs] = {}; + size_t num_outputs = 0; +}; + +__device__ __forceinline__ int64_t load_split_size(const void *__restrict__ split_sizes, + DType dtype, size_t idx) { + switch (dtype) { + case DType::kInt32: + return static_cast(static_cast(split_sizes)[idx]); + case DType::kInt64: + return static_cast(split_sizes)[idx]; + default: + NVTE_DEVICE_ERROR("Unsupported dtype for split_sizes (expected int32 or int64)."); + return 0; + } +} + +__device__ __forceinline__ void store_output(void *__restrict__ output, DType dtype, size_t idx, + int64_t value) { + switch (dtype) { + case DType::kInt32: + static_cast(output)[idx] = static_cast(value); + return; + case DType::kInt64: + static_cast(output)[idx] = value; + return; + default: + NVTE_DEVICE_ERROR("Unsupported dtype for output (expected int32 or int64)."); + } +} + +__global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) { + const size_t tid = threadIdx.x; + + // Fill leading zeros if needed + if (tid == 0) { + for (size_t out_idx = 0; out_idx < args.num_outputs; ++out_idx) { + if (args.include_leading_zero[out_idx]) { + store_output(args.outputs[out_idx], args.outputs_dtype[out_idx], 0, 0); + } + } + } + + // Workspace for prefix sum chunk + __shared__ int64_t block_scan[kThreadsPerBlock]; + + // Sum from previous chunks + __shared__ int64_t chunk_prefix; + if (tid == 0) { + chunk_prefix = 0; + } + __syncthreads(); + + // Perform prefix sum in chunks + for (size_t chunk_start = 0; chunk_start < args.num_splits; chunk_start += kThreadsPerBlock) { + const size_t idx = chunk_start + tid; + + // Load input from global memory into shared memory + if (idx < args.num_splits) { + block_scan[tid] = load_split_size(args.split_sizes, args.split_sizes_dtype, idx); + } else { + block_scan[tid] = 0; + } + __syncthreads(); + + // Prefix sum in shared memory + for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) { + const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0; + __syncthreads(); + block_scan[tid] += addend; + __syncthreads(); + } + + // Compute global prefix sum, apply strides, and store to output + if (idx < args.num_splits) { + const int64_t prefix = chunk_prefix + block_scan[tid]; + for (size_t out_idx = 0; out_idx < args.num_outputs; ++out_idx) { + const size_t write_idx = idx + (args.include_leading_zero[out_idx] ? 1 : 0); + store_output(args.outputs[out_idx], args.outputs_dtype[out_idx], write_idx, + prefix * args.strides[out_idx]); + } + } + + // Update sum for later chunks + __syncthreads(); + if (tid == kThreadsPerBlock - 1) { + chunk_prefix += block_scan[tid]; + } + __syncthreads(); + } +} + +} // namespace + +} // namespace transformer_engine::splits_to_offsets + +void nvte_splits_to_offsets(const int64_t *split_sizes, int64_t *output, size_t num_splits, + int64_t stride, cudaStream_t stream) { + NVTE_API_CALL(nvte_splits_to_offsets); + NVTE_CHECK(output != nullptr, "Output pointer is NULL."); + NVTE_CHECK(num_splits > 0, "num_splits must be greater than 0."); + NVTE_CHECK(split_sizes != nullptr, "split_sizes pointer is NULL."); + NVTE_CHECK(stride > 0, "stride must be greater than 0."); + + using namespace transformer_engine; + namespace s2o = transformer_engine::splits_to_offsets; + + s2o::KernelArgs args = {}; + args.split_sizes = split_sizes; + args.split_sizes_dtype = DType::kInt64; + args.num_splits = num_splits; + args.outputs[0] = output; + args.outputs_dtype[0] = DType::kInt64; + args.strides[0] = stride; + args.include_leading_zero[0] = true; + args.num_outputs = 1; + s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, + const int64_t *strides, const int *include_leading_zero, + size_t num_outputs, cudaStream_t stream) { + NVTE_API_CALL(nvte_splits_to_offsets_multi); + using namespace transformer_engine; + namespace s2o = transformer_engine::splits_to_offsets; + + if (num_outputs == 0) { + return; + } + NVTE_CHECK(outputs != nullptr, "outputs is NULL."); + NVTE_CHECK(strides != nullptr, "strides is NULL."); + NVTE_CHECK(include_leading_zero != nullptr, "include_leading_zero is NULL."); + + // Check if dtype is supported + const auto is_integer_dtype = [](DType dtype) { + return dtype == DType::kInt32 || dtype == DType::kInt64; + }; + + // Check input tensor + const auto *split_sizes_tensor = convertNVTETensorCheck(split_sizes); + const auto split_sizes_dtype = split_sizes_tensor->dtype(); + const auto num_splits = split_sizes_tensor->numel(); + NVTE_CHECK(num_splits > 0 && split_sizes_tensor->dim() == 1, + "split_sizes must be a non-empty 1D tensor, but got shape=", + split_sizes_tensor->shape(), "."); + NVTE_CHECK(is_integer_dtype(split_sizes_dtype), + "split_sizes must be an int32/int64 tensor, but got dtype=", split_sizes_dtype, "."); + + // Check output tensors + std::vector output_tensors(num_outputs); + for (size_t i = 0; i < num_outputs; ++i) { + const auto *out_tensor = convertNVTETensorCheck(outputs[i]); + const auto out_dtype = out_tensor->dtype(); + const bool has_leading_zero = include_leading_zero[i] != 0; + const Shape expected_shape = {num_splits + (has_leading_zero ? 1 : 0)}; + NVTE_CHECK(out_tensor->shape() == expected_shape, "Expected outputs[", i, + "] to have shape=", expected_shape, ", but got shape=", out_tensor->shape(), "."); + NVTE_CHECK(is_integer_dtype(out_dtype), "Expected outputs[", i, + "] to be an int32/int64 tensor, but got dtype=", out_dtype, "."); + output_tensors[i] = out_tensor; + } + + // Chunk outputs to fit in kernel arguments and launch kernels + for (size_t chunk_start = 0; chunk_start < num_outputs; + chunk_start += s2o::KernelArgs::kMaxNumOutputs) { + const size_t chunk_size = std::min(s2o::KernelArgs::kMaxNumOutputs, num_outputs - chunk_start); + s2o::KernelArgs args = {}; + args.split_sizes = split_sizes_tensor->data.dptr; + args.split_sizes_dtype = split_sizes_dtype; + args.num_splits = num_splits; + args.num_outputs = chunk_size; + for (size_t i = 0; i < chunk_size; ++i) { + const size_t out_idx = chunk_start + i; + args.outputs[i] = output_tensors[out_idx]->data.dptr; + args.outputs_dtype[i] = output_tensors[out_idx]->dtype(); + args.strides[i] = strides[out_idx]; + args.include_leading_zero[i] = include_leading_zero[out_idx] != 0; + } + s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args); + NVTE_CHECK_CUDA(cudaGetLastError()); + } +} diff --git a/transformer_engine/debug/features/fake_quant.py b/transformer_engine/debug/features/fake_quant.py index f48b49b7257..48112abe4fe 100644 --- a/transformer_engine/debug/features/fake_quant.py +++ b/transformer_engine/debug/features/fake_quant.py @@ -13,16 +13,16 @@ from nvdlfw_inspect.utils import append_parent_docstring -import transformer_engine_torch as tex from transformer_engine.debug.features.api import TEConfigAPIMapper from transformer_engine.common.recipe import Format +from transformer_engine.pytorch import DType from transformer_engine.pytorch.tensor import Quantizer from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.quantization import _default_sf_compute -def fake_quantize(tensor: torch.Tensor, fp8_format: tex.DType, out=None): +def fake_quantize(tensor: torch.Tensor, fp8_format: str, out=None): """Input tensor is quantized to fp8 and then dequantized.""" assert tensor.dtype in ( @@ -43,10 +43,10 @@ def fake_quantize(tensor: torch.Tensor, fp8_format: tex.DType, out=None): if fp8_format in ["FP8E4M3", "FP8E5M2"]: if fp8_format == "FP8E4M3": fp8_max = Format.E4M3.value.max_fwd - fp8_dtype = tex.DType.kFloat8E4M3 + fp8_dtype = DType.kFloat8E4M3 else: fp8_max = Format.E5M2.value.max_fwd - fp8_dtype = tex.DType.kFloat8E5M2 + fp8_dtype = DType.kFloat8E5M2 amax = tensor.abs().max().float() one = torch.ones(1, device=tensor.device) scale = _default_sf_compute(amax, one, fp8_max, 0) diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index d26f9ef7f65..96f1b644cfb 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -12,7 +12,7 @@ import nvdlfw_inspect.api as debug_api from nvdlfw_inspect.debug_features.log_tensor_stats import LogTensorStats as BaseLogTensorStats from nvdlfw_inspect.registry import Registry, api_method -import transformer_engine_torch as tex +from transformer_engine.pytorch import DType from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter @@ -251,7 +251,7 @@ def update_aux_dict( Yields the aux_dict. Needs to clean after usage, because it possibly change the usage of the quantized tensor. """ - fp8_dtype = tex.DType.kFloat8E4M3 + fp8_dtype = DType.kFloat8E4M3 if recipe_name in ["fp8_delayed_scaling", "fp8_current_scaling", "fp8_block_scaling"]: assert isinstance( quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer, Float8BlockQuantizer) diff --git a/transformer_engine/debug/features/per_tensor_scaling.py b/transformer_engine/debug/features/per_tensor_scaling.py index a4bab4eaf56..e0fc44bd08f 100644 --- a/transformer_engine/debug/features/per_tensor_scaling.py +++ b/transformer_engine/debug/features/per_tensor_scaling.py @@ -11,7 +11,7 @@ import nvdlfw_inspect.api as debug_api from nvdlfw_inspect.registry import Registry, api_method -import transformer_engine_torch as tex +from transformer_engine.pytorch import DType from transformer_engine.pytorch.tensor import Quantizer from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Tensor, @@ -22,7 +22,7 @@ def per_tensor_cast( - tensor: torch.Tensor, fp8_dtype: tex.DType, out: Float8Tensor = None + tensor: torch.Tensor, fp8_dtype: DType, out: Float8Tensor = None ) -> Float8Tensor: """ This function computes the scaling factors based on the tensor amax and then casts it to the fp8 @@ -35,8 +35,8 @@ def per_tensor_cast( ), "[NVTORCH INSPECT ERROR] Unsupported tensor type for per tensor current scaling" assert tensor.is_cuda, "[NVTORCH INSPECT ERROR] Must be a GPU tensor." assert fp8_dtype in { - tex.DType.kFloat8E4M3, - tex.DType.kFloat8E5M2, + DType.kFloat8E4M3, + DType.kFloat8E5M2, }, "[NVTORCH INSPECT ERROR] Only 2 FP8 types: E4M3 and E5M2 are supported in TE." tensor = tensor.contiguous() diff --git a/transformer_engine/debug/features/utils/stats_computation.py b/transformer_engine/debug/features/utils/stats_computation.py index 66684000171..f01037acd13 100644 --- a/transformer_engine/debug/features/utils/stats_computation.py +++ b/transformer_engine/debug/features/utils/stats_computation.py @@ -11,8 +11,8 @@ import torch import torch.nn.functional as F -import transformer_engine_torch as tex from transformer_engine.common.recipe import Format +from transformer_engine.pytorch import DType class BlockwiseDynamicRangeStat( @@ -142,8 +142,8 @@ def compute_fp8_delayed_scaling_overflows_num(tensor, quantized_tensor): # Map each supported FP8 dtype to its corresponding max forward value. dtype_to_max = { - tex.DType.kFloat8E4M3: Format.E4M3.value.max_fwd, - tex.DType.kFloat8E5M2: Format.E5M2.value.max_fwd, + DType.kFloat8E4M3: Format.E4M3.value.max_fwd, + DType.kFloat8E5M2: Format.E5M2.value.max_fwd, } if dtype not in dtype_to_max: diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index f54a043fd21..ecca4a38711 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -5,7 +5,7 @@ from __future__ import annotations from enum import Enum from functools import partial -from typing import Optional, Tuple, Union +from typing import Any, Callable, Mapping, Optional, Tuple, Union import warnings from jax.ad_checkpoint import checkpoint_name @@ -1391,10 +1391,63 @@ def _fused_attn_bwd_rule( _fused_attn.defvjp(_fused_attn_fwd_rule, _fused_attn_bwd_rule) +@partial(jax.custom_vjp, nondiff_argnums=(3, 4)) +def _fused_attn_score_mod( + qkv: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], + score_mod_tensors: Tuple[jnp.ndarray, ...], + score_mod_bprop_tensors: Tuple[jnp.ndarray, ...], + config, + context_checkpoint_name: str, +): + output, _ = _fused_attn_score_mod_fwd_rule( + qkv, + score_mod_tensors, + score_mod_bprop_tensors, + config, + context_checkpoint_name, + ) + return output + + +def _fused_attn_score_mod_fwd_rule( + qkv, + score_mod_tensors, + score_mod_bprop_tensors, + config, + context_checkpoint_name, +): + output, softmax_stats = tex.fused_attn_score_mod_fwd(qkv, score_mod_tensors, config) + output = checkpoint_name(output, context_checkpoint_name) + softmax_stats = checkpoint_name(softmax_stats, context_checkpoint_name) + return output, (qkv, score_mod_tensors, score_mod_bprop_tensors, output, softmax_stats) + + +def _fused_attn_score_mod_bwd_rule(config, context_checkpoint_name, ctx, dz): + del context_checkpoint_name + qkv, score_mod_tensors, score_mod_bprop_tensors, output, softmax_stats = ctx + grad_qkv = tex.fused_attn_score_mod_bwd( + qkv, + output, + dz, + softmax_stats, + score_mod_tensors, + score_mod_bprop_tensors, + config, + ) + return ( + grad_qkv, + tuple(None for _ in score_mod_tensors), + tuple(None for _ in score_mod_bprop_tensors), + ) + + +_fused_attn_score_mod.defvjp(_fused_attn_score_mod_fwd_rule, _fused_attn_score_mod_bwd_rule) + + def fused_attn( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], - sequence_descriptor: SequenceDescriptor, + sequence_descriptor: Optional[SequenceDescriptor], seed: Optional[jnp.ndarray], attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType, @@ -1411,6 +1464,10 @@ def fused_attn( context_checkpoint_name: str = "context", softmax_offset: Optional[jnp.ndarray] = None, stripe_size: int | None = None, + score_mod: Optional[Callable] = None, + score_mod_bprop: Optional[Callable] = None, + score_mod_tensors: Optional[Mapping[str, Any]] = None, + score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, ): """ Perform cuDNN fused attention. @@ -1453,6 +1510,20 @@ def fused_attn( Currently, a stripe_size > 1 is only supported for CP + THD + Striped + AG, whereas a stripe_size=1 is supported for both, CP + THD + Striped + AG and CP + THD + Striped + P2P(Ring) None indicates no striping strategy + score_mod (Optional[Callable]): Experimental cuDNN frontend score modification callback. + The callback is called as `score_mod(graph, score, tensors)` while building a + cuDNN frontend graph. When provided, this path only supports BSHD_BSHD_BSHD + layout and is mutually exclusive with masks, padding, bias, dropout, context + parallelism, sliding windows, and non-vanilla softmax. + score_mod_bprop (Optional[Callable]): Optional score modification backward callback, + called as `score_mod_bprop(graph, dscore, tensors)`. If omitted, cuDNN uses the + default backward behavior for the forward score modification graph. + score_mod_tensors (Optional[Mapping[str, Any]]): Additional tensors or Python/NumPy + scalars made available to `score_mod` through its `tensors` dictionary. Scalars + are represented as cuDNN pass-by-value tensors. Tensor entries are treated as + non-differentiable auxiliary inputs. + score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or + Python/NumPy scalars made available to `score_mod_bprop`. Returns: (jnp.ndarray): The output tensor from the fused attention. @@ -1485,6 +1556,53 @@ def fused_attn( AttnBiasType.NO_BIAS, AttnMaskType.PADDING_CAUSAL_MASK, QKVLayout.T3HD, 0.125, 0, True, 3) """ + if score_mod is None: + score_mod_only_args = [ + name + for name, value in ( + ("score_mod_bprop", score_mod_bprop), + ("score_mod_tensors", score_mod_tensors), + ("score_mod_bprop_tensors", score_mod_bprop_tensors), + ) + if value is not None + ] + if score_mod_only_args: + raise ValueError(f"{', '.join(score_mod_only_args)} require score_mod to be provided.") + else: + tex.validate_fused_attn_score_mod( + qkv, + bias, + sequence_descriptor, + seed, + attn_bias_type, + attn_mask_type, + qkv_layout, + softmax_type, + dropout_probability, + max_segments_per_seq, + window_size, + context_parallel_strategy, + context_parallel_causal_load_balanced, + context_parallel_axis, + softmax_offset, + stripe_size, + ) + config, tensor_operands, bprop_tensor_operands = tex.make_fused_attn_score_mod_config( + score_mod, + score_mod_bprop, + score_mod_tensors, + score_mod_bprop_tensors, + scaling_factor, + is_training, + ) + return _fused_attn_score_mod( + qkv, + tensor_operands, + bprop_tensor_operands, + config, + context_checkpoint_name, + ) + if sequence_descriptor is None or isinstance(sequence_descriptor, jnp.ndarray): warnings.warn( "Pass mask to fused_attn is deprecated, please use SequenceDescriptor instead. " diff --git a/transformer_engine/jax/cpp_extensions/__init__.py b/transformer_engine/jax/cpp_extensions/__init__.py index fe1f93dc7ab..b42c9097408 100644 --- a/transformer_engine/jax/cpp_extensions/__init__.py +++ b/transformer_engine/jax/cpp_extensions/__init__.py @@ -5,6 +5,7 @@ from .activation import * from .amax import * from .attention import * +from .flex_attention import * from .normalization import * from .quantization import * from .softmax import * diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py new file mode 100644 index 00000000000..71f0176ed79 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -0,0 +1,967 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""cuDNN frontend score_mod fused attention helpers.""" +import hashlib +import importlib +import inspect +import os +from pathlib import Path +import sys +from dataclasses import dataclass +from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple + +import jax +import jax.numpy as jnp +import numpy as np +from jax import ffi + +import transformer_engine_jax + + +__all__ = [ + "FusedAttnScoreModHelper", + "make_fused_attn_score_mod_config", + "validate_fused_attn_score_mod", + "fused_attn_score_mod_fwd", + "fused_attn_score_mod_bwd", +] + +_CUDNN_FRONTEND_PYTHON_PATH = ( + Path(__file__).resolve().parents[3] / "3rdparty" / "cudnn-frontend" / "python" +) + + +def _is_non_deterministic_allowed(): + """Check if non-deterministic kernels are allowed.""" + return bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + + +def _enum_name(value: Any) -> str: + return getattr(value, "name", str(value)) + + +def validate_fused_attn_score_mod( + qkv: Tuple[jnp.ndarray, ...], + bias: Optional[jnp.ndarray], + sequence_descriptor: Optional[Any], + seed: Optional[jnp.ndarray], + attn_bias_type: Any, + attn_mask_type: Any, + qkv_layout: Any, + softmax_type: Any, + dropout_probability: float, + max_segments_per_seq: int, + window_size: Optional[Tuple[int, int]], + context_parallel_strategy: Any, + context_parallel_causal_load_balanced: bool, + context_parallel_axis: str, + softmax_offset: Optional[jnp.ndarray], + stripe_size: int | None, +): + """Validate arguments for the cuDNN frontend score_mod path.""" + header = "score_mod fused_attn" + if _enum_name(qkv_layout) != "BSHD_BSHD_BSHD": + raise ValueError(f"{header} currently only supports QKVLayout.BSHD_BSHD_BSHD.") + if len(qkv) != 3: + raise ValueError(f"{header} requires separate query, key and value tensors.") + if any(tensor.ndim != 4 for tensor in qkv): + raise ValueError(f"{header} requires rank-4 BSHD query/key/value tensors.") + q, k, v = qkv + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError(f"{header} requires query, key and value to have the same dtype.") + if q.dtype not in (jnp.float16, jnp.bfloat16): + raise ValueError(f"{header} only supports FP16/BF16 query, key and value tensors.") + if q.shape[0] != k.shape[0] or q.shape[0] != v.shape[0]: + raise ValueError(f"{header} requires matching batch dimensions.") + if k.shape[1] != v.shape[1]: + raise ValueError(f"{header} requires key and value sequence lengths to match.") + if k.shape[2] != v.shape[2]: + raise ValueError(f"{header} requires key and value head counts to match.") + if q.shape[3] != k.shape[3]: + raise ValueError(f"{header} requires query/key head dimensions to match.") + + if bias is not None or _enum_name(attn_bias_type) != "NO_BIAS": + raise ValueError(f"{header} is mutually exclusive with attention bias.") + if sequence_descriptor is not None: + raise ValueError(f"{header} is mutually exclusive with padding/sequence descriptors.") + if seed is not None: + raise ValueError(f"{header} is mutually exclusive with dropout seed.") + if _enum_name(attn_mask_type) != "NO_MASK": + raise ValueError(f"{header} is mutually exclusive with attention masks.") + if _enum_name(softmax_type) != "VANILLA_SOFTMAX" or softmax_offset is not None: + raise ValueError(f"{header} only supports vanilla softmax without softmax_offset.") + if dropout_probability != 0.0: + raise ValueError(f"{header} is mutually exclusive with dropout.") + if max_segments_per_seq != 1: + raise ValueError(f"{header} is mutually exclusive with packed/ragged sequence metadata.") + if window_size not in (None, (-1, -1)): + raise ValueError(f"{header} is mutually exclusive with sliding-window attention.") + if _enum_name(context_parallel_strategy) != "DEFAULT": + raise ValueError(f"{header} is mutually exclusive with context parallelism.") + if context_parallel_causal_load_balanced or context_parallel_axis: + raise ValueError(f"{header} is mutually exclusive with context parallelism.") + if stripe_size is not None: + raise ValueError(f"{header} is mutually exclusive with striped context parallelism.") + + +@dataclass(frozen=True) +class _ScoreModScalarSpec: + """Static pass-by-value scalar used when building a cuDNN frontend graph.""" + + name: str + dtype: str + value: bytes + dim: Tuple[int, ...] = (1, 1, 1, 1) + stride: Tuple[int, ...] = (1, 1, 1, 1) + + +class _UncacheableScoreModKey: + """Unique static key for callbacks that must not share compiled score_mod graphs.""" + + def __hash__(self): + return id(self) + + def __eq__(self, other): + return self is other + + +def _score_mod_key_is_uncacheable(key: Any) -> bool: + return isinstance(key, _UncacheableScoreModKey) + + +def _freeze_score_mod_cache_key(value: Any) -> Any: + """Convert a user-provided score_mod graph key into a hashable structure.""" + if _is_array_operand(value): + raise TypeError( + "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " + "through score_mod_tensors or score_mod_bprop_tensors instead." + ) + if isinstance(value, Mapping): + items = ( + ( + _freeze_score_mod_cache_key(key), + _freeze_score_mod_cache_key(val), + ) + for key, val in value.items() + ) + return tuple(sorted(items, key=repr)) + if isinstance(value, (list, tuple)): + return tuple(_freeze_score_mod_cache_key(item) for item in value) + if isinstance(value, (set, frozenset)): + items = (_freeze_score_mod_cache_key(item) for item in value) + return tuple(sorted(items, key=repr)) + try: + hash(value) + except TypeError as exc: + raise TypeError( + "score_mod_graph_cache_key() must return a hashable value or a nested " + "combination of mapping/list/tuple/set values." + ) from exc + return value + + +def _score_mod_explicit_cache_key(callback_owner: Any) -> Optional[Any]: + """Return a user-provided structural graph key for a score_mod callback.""" + explicit_key = getattr(callback_owner, "score_mod_graph_cache_key", None) + if explicit_key is None: + return None + explicit_key = explicit_key() if callable(explicit_key) else explicit_key + return _freeze_score_mod_cache_key(explicit_key) + + +def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: + """Create a stable graph cache key for a score_mod callable. + + Module-level functions are assumed to have stable topology. Stateful bound methods and + callable instances need an explicit score_mod_graph_cache_key(); otherwise their graphs + are left uncached to avoid reusing stale graphs after Python object address reuse. + """ + if callback is None: + return None + self_obj = getattr(callback, "__self__", None) + func_obj = getattr(callback, "__func__", None) + if self_obj is not None and func_obj is not None: + explicit_key = _score_mod_explicit_cache_key(self_obj) + if explicit_key is None: + return _UncacheableScoreModKey() + return ( + "bound_method", + type(self_obj), + func_obj.__module__, + func_obj.__qualname__, + explicit_key, + ) + + explicit_key = _score_mod_explicit_cache_key(callback) + if explicit_key is not None: + return ( + "callable", + type(callback), + getattr(callback, "__module__", None), + getattr(callback, "__qualname__", None), + explicit_key, + ) + + if ( + inspect.isfunction(callback) + and callback.__closure__ is None + and "" not in callback.__qualname__ + ): + return ("function", callback.__module__, callback.__qualname__) + + return _UncacheableScoreModKey() + + +@dataclass(frozen=True) +class _FusedAttnScoreModConfig: + """Static configuration for cuDNN frontend score_mod SDPA graphs.""" + + score_mod: Callable + score_mod_bprop: Optional[Callable] + score_mod_key: Any + score_mod_bprop_key: Any + score_mod_tensor_names: Tuple[str, ...] + score_mod_bprop_tensor_names: Tuple[str, ...] + score_mod_scalars: Tuple[_ScoreModScalarSpec, ...] + score_mod_bprop_scalars: Tuple[_ScoreModScalarSpec, ...] + scaling_factor: float + is_training: bool + deterministic: bool + + def __hash__(self): + return hash( + ( + self.score_mod_key, + self.score_mod_bprop_key, + self.score_mod_tensor_names, + self.score_mod_bprop_tensor_names, + self.score_mod_scalars, + self.score_mod_bprop_scalars, + self.scaling_factor, + self.is_training, + self.deterministic, + ) + ) + + def __eq__(self, other): + if not isinstance(other, _FusedAttnScoreModConfig): + return False + return ( + self.score_mod_key == other.score_mod_key + and self.score_mod_bprop_key == other.score_mod_bprop_key + and self.score_mod_tensor_names == other.score_mod_tensor_names + and self.score_mod_bprop_tensor_names == other.score_mod_bprop_tensor_names + and self.score_mod_scalars == other.score_mod_scalars + and self.score_mod_bprop_scalars == other.score_mod_bprop_scalars + and self.scaling_factor == other.scaling_factor + and self.is_training == other.is_training + and self.deterministic == other.deterministic + ) + + +@dataclass(frozen=True) +class _SerializedScoreModGraph: + """Serialized cuDNN frontend graph and static metadata for C++ execution.""" + + serialized_graph: bytes + graph_hash: Tuple[int, int] + cudnn_frontend_version: int + workspace_size: int + input_uids: np.ndarray + output_uids: np.ndarray + scalar_uids: np.ndarray + scalar_sizes: np.ndarray + scalar_values: np.ndarray + + +# cuDNN frontend tensor UIDs are arbitrary, but assigning stable values makes serialized +# graphs deterministic and simplifies debugging across the Python graph builder and C++ executor. +_SCORE_MOD_UID_Q = 1 +_SCORE_MOD_UID_K = 2 +_SCORE_MOD_UID_V = 3 +_SCORE_MOD_UID_O = 4 +_SCORE_MOD_UID_STATS = 5 +_SCORE_MOD_UID_DO = 6 +_SCORE_MOD_UID_DQ = 7 +_SCORE_MOD_UID_DK = 8 +_SCORE_MOD_UID_DV = 9 +_SCORE_MOD_FWD_TENSOR_UID_BASE = 1000 +_SCORE_MOD_BPROP_TENSOR_UID_BASE = 2000 +_SCORE_MOD_FWD_SCALAR_UID_BASE = 3000 +_SCORE_MOD_BPROP_SCALAR_UID_BASE = 4000 + +_score_mod_graph_cache: Dict[Tuple[Any, ...], _SerializedScoreModGraph] = {} + + +def _row_major_stride(shape: Sequence[int]) -> Tuple[int, ...]: + stride = [] + running = 1 + for dim in reversed(tuple(shape)): + stride.append(running) + running *= dim + return tuple(reversed(stride)) + + +def _bshd_as_bhsd_dim_stride(shape: Sequence[int]) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: + if len(shape) != 4: + raise ValueError(f"score_mod requires rank-4 BSHD tensors, got shape={shape}.") + batch, seqlen, heads, head_dim = tuple(shape) + return ( + (batch, heads, seqlen, head_dim), + (seqlen * heads * head_dim, head_dim, heads * head_dim, 1), + ) + + +def _dtype_name(dtype) -> str: + return str(jnp.dtype(dtype)) + + +def _is_array_operand(value: Any) -> bool: + return ( + hasattr(value, "shape") + and hasattr(value, "dtype") + and not isinstance(value, (bool, int, float, complex, np.generic)) + ) + + +def _scalar_to_spec(name: str, value: Any) -> _ScoreModScalarSpec: + if isinstance(value, bool): + dtype = np.bool_ + elif isinstance(value, int): + dtype = np.int32 + elif isinstance(value, float): + dtype = np.float32 + elif isinstance(value, np.generic): + dtype = value.dtype + else: + scalar = np.asarray(value) + if scalar.shape != (): + raise ValueError( + f"score_mod tensor '{name}' is neither a JAX array nor a scalar pass-by-value." + ) + dtype = scalar.dtype + + scalar = np.full((1, 1, 1, 1), value, dtype=dtype) + return _ScoreModScalarSpec(name=name, dtype=str(scalar.dtype), value=scalar.tobytes()) + + +def _split_score_mod_tensors( + tensors: Optional[Mapping[str, Any]], *, argument_name: str +) -> Tuple[Tuple[str, ...], Tuple[jnp.ndarray, ...], Tuple[_ScoreModScalarSpec, ...]]: + if tensors is None: + return (), (), () + if not isinstance(tensors, Mapping): + raise TypeError(f"{argument_name} must be a mapping from string names to tensors/scalars.") + + names = [] + operands = [] + scalars = [] + for name, value in tensors.items(): + if not isinstance(name, str): + raise TypeError(f"{argument_name} keys must be strings, got {type(name).__name__}.") + if _is_array_operand(value): + if len(value.shape) == 0: + raise ValueError( + f"{argument_name}['{name}'] is a rank-0 array. Use a Python/NumPy scalar " + "for cuDNN pass-by-value scalars, or reshape it to a tensor." + ) + names.append(name) + operands.append(jnp.asarray(value)) + else: + scalars.append(_scalar_to_spec(name, value)) + return tuple(names), tuple(operands), tuple(scalars) + + +def _make_fused_attn_score_mod_config( + score_mod: Callable, + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Mapping[str, Any]], + score_mod_bprop_tensors: Optional[Mapping[str, Any]], + scaling_factor: float, + is_training: bool, +) -> Tuple[_FusedAttnScoreModConfig, Tuple[jnp.ndarray, ...], Tuple[jnp.ndarray, ...]]: + """Normalize score_mod operands and create a static graph-build config.""" + if not callable(score_mod): + raise TypeError("score_mod must be callable.") + if score_mod_bprop is not None and not callable(score_mod_bprop): + raise TypeError("score_mod_bprop must be callable when provided.") + if score_mod_bprop is None and score_mod_bprop_tensors: + raise ValueError("score_mod_bprop_tensors requires score_mod_bprop to be provided.") + + tensor_names, tensor_operands, scalars = _split_score_mod_tensors( + score_mod_tensors, argument_name="score_mod_tensors" + ) + bprop_tensor_names, bprop_tensor_operands, bprop_scalars = _split_score_mod_tensors( + score_mod_bprop_tensors, argument_name="score_mod_bprop_tensors" + ) + config = _FusedAttnScoreModConfig( + score_mod=score_mod, + score_mod_bprop=score_mod_bprop, + score_mod_key=_score_mod_callback_cache_key(score_mod), + score_mod_bprop_key=_score_mod_callback_cache_key(score_mod_bprop), + score_mod_tensor_names=tensor_names, + score_mod_bprop_tensor_names=bprop_tensor_names, + score_mod_scalars=scalars, + score_mod_bprop_scalars=bprop_scalars, + scaling_factor=float(scaling_factor), + is_training=bool(is_training), + deterministic=not _is_non_deterministic_allowed(), + ) + return config, tensor_operands, bprop_tensor_operands + + +def _cudnn_data_type(cudnn, dtype): + dtype = jnp.dtype(dtype) + if dtype == jnp.float16: + return cudnn.data_type.HALF + if dtype == jnp.bfloat16: + return cudnn.data_type.BFLOAT16 + if dtype == jnp.float32: + return cudnn.data_type.FLOAT + if dtype == jnp.float64: + return cudnn.data_type.DOUBLE + if dtype == jnp.int32: + return cudnn.data_type.INT32 + if dtype == jnp.int64: + return cudnn.data_type.INT64 + if dtype == jnp.uint8: + return cudnn.data_type.UINT8 + if dtype == jnp.bool_: + return cudnn.data_type.BOOLEAN + raise ValueError(f"Unsupported score_mod tensor dtype: {dtype}.") + + +def _cudnn_data_type_from_name(cudnn, dtype_name: str): + if dtype_name == "bfloat16": + return cudnn.data_type.BFLOAT16 + return _cudnn_data_type(cudnn, np.dtype(dtype_name)) + + +def _graph_tensor_from_aval(cudnn, graph, name: str, aval, uid: int): + shape = tuple(int(dim) for dim in aval.shape) + return graph.tensor( + name=name, + dim=shape, + stride=_row_major_stride(shape), + data_type=_cudnn_data_type(cudnn, aval.dtype), + uid=uid, + ) + + +def _score_mod_graph_tensors( + cudnn, + graph, + names: Tuple[str, ...], + avals: Sequence[Any], + scalars: Tuple[_ScoreModScalarSpec, ...], + tensor_uid_base: int, + scalar_uid_base: int, +): + graph_tensors = {} + tensor_uids = [] + for index, (name, aval) in enumerate(zip(names, avals)): + uid = tensor_uid_base + index + graph_tensors[name] = _graph_tensor_from_aval(cudnn, graph, name, aval, uid) + tensor_uids.append(uid) + + scalar_uids = [] + scalar_values = [] + for index, scalar in enumerate(scalars): + uid = scalar_uid_base + index + graph_tensors[scalar.name] = graph.tensor( + name=scalar.name, + dim=scalar.dim, + stride=scalar.stride, + is_pass_by_value=True, + data_type=_cudnn_data_type_from_name(cudnn, scalar.dtype), + uid=uid, + ) + scalar_uids.append(uid) + scalar_values.append(scalar.value) + + return graph_tensors, tuple(tensor_uids), tuple(scalar_uids), tuple(scalar_values) + + +def _encode_cudnn_frontend_version(version: str) -> int: + public_version = version.split("+", 1)[0].split("-", 1)[0] + parts = public_version.split(".") + if len(parts) < 3: + raise RuntimeError(f"Could not parse cuDNN frontend Python version: {version!r}.") + major, minor, patch = (int(part) for part in parts[:3]) + return major * 10000 + minor * 100 + patch + + +def _check_cudnn_frontend_version_match(cudnn) -> int: + python_version_string = getattr(cudnn, "__version__", None) + if python_version_string is None: + raise RuntimeError("cuDNN frontend Python package does not expose __version__.") + python_version = _encode_cudnn_frontend_version(python_version_string) + cpp_version = int(transformer_engine_jax.get_cudnn_frontend_version()) + if python_version != cpp_version: + raise RuntimeError( + "cuDNN frontend Python/C++ version mismatch for score_mod graph serialization: " + f"Python cudnn.__version__={python_version_string!r} encodes to {python_version}, " + f"but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION={cpp_version}. " + "Use matching cuDNN frontend Python package and C++ headers." + ) + return python_version + + +def _score_mod_graph_hash(serialized_graph: bytes) -> Tuple[int, int]: + digest = hashlib.sha256(serialized_graph).digest() + return ( + int.from_bytes(digest[0:8], byteorder="little", signed=True), + int.from_bytes(digest[8:16], byteorder="little", signed=True), + ) + + +def _pack_score_mod_scalar_values( + scalar_values: Sequence[bytes], +) -> Tuple[np.ndarray, np.ndarray]: + scalar_sizes = np.asarray([len(value) for value in scalar_values], dtype=np.int64) + packed_values = np.zeros((len(scalar_values), 16), dtype=np.uint8) + for index, value in enumerate(scalar_values): + if len(value) > 16: + raise ValueError("score_mod pass-by-value scalars must be at most 16 bytes.") + packed_values[index, : len(value)] = np.frombuffer(value, dtype=np.uint8) + return scalar_sizes, packed_values.reshape(-1) + + +def _serialized_score_mod_graph( + *, + serialized_graph: bytes, + cudnn_frontend_version: int, + workspace_size: int, + input_uids: Sequence[int], + output_uids: Sequence[int], + scalar_uids: Sequence[int], + scalar_values: Sequence[bytes], +) -> _SerializedScoreModGraph: + scalar_sizes, packed_scalar_values = _pack_score_mod_scalar_values(scalar_values) + return _SerializedScoreModGraph( + serialized_graph=serialized_graph, + graph_hash=_score_mod_graph_hash(serialized_graph), + cudnn_frontend_version=int(cudnn_frontend_version), + workspace_size=int(workspace_size), + input_uids=np.asarray(input_uids, dtype=np.int64), + output_uids=np.asarray(output_uids, dtype=np.int64), + scalar_uids=np.asarray(scalar_uids, dtype=np.int64), + scalar_sizes=scalar_sizes, + scalar_values=packed_scalar_values, + ) + + +def _wrap_score_mod(score_mod: Optional[Callable], graph_tensors: Dict[str, Any]): + if score_mod is None: + return None + + def wrapped_score_mod(sdpa_graph, score_tensor): + return score_mod(sdpa_graph, score_tensor, graph_tensors) + + return wrapped_score_mod + + +def _finalize_score_mod_graph(cudnn, graph) -> Tuple[int, bytes, int]: + graph.validate() + graph.build_operation_graph() + try: + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + except cudnn.cudnnGraphNotSupportedError as exc: + raise RuntimeError(f"cuDNN score_mod SDPA graph is not supported: {exc}") from exc + graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) + serialized_graph = bytes(graph.serialize()) + return ( + max(int(graph.get_workspace_size()), 1), + serialized_graph, + _check_cudnn_frontend_version_match(cudnn), + ) + + +def _graph_cache_key( + direction: str, + config: _FusedAttnScoreModConfig, + avals: Sequence[Any], +) -> Optional[Tuple[Any, ...]]: + if _score_mod_key_is_uncacheable(config.score_mod_key) or _score_mod_key_is_uncacheable( + config.score_mod_bprop_key + ): + return None + return ( + direction, + config, + tuple((tuple(aval.shape), _dtype_name(aval.dtype)) for aval in avals), + ) + + +def _shape_dtype(value) -> jax.ShapeDtypeStruct: + return jax.ShapeDtypeStruct(tuple(value.shape), value.dtype) + + +def _import_cudnn_for_score_mod(): + cudnn_frontend_path = str(_CUDNN_FRONTEND_PYTHON_PATH) + cudnn_frontend_package = _CUDNN_FRONTEND_PYTHON_PATH / "cudnn" + if ( + any(cudnn_frontend_package.glob("_compiled_module*")) + and cudnn_frontend_path not in sys.path + ): + sys.path.insert(0, cudnn_frontend_path) + try: + cudnn = importlib.import_module("cudnn") + except ImportError as exc: + raise ImportError( + "score_mod fused_attn requires the cuDNN frontend Python package (`cudnn`)." + ) from exc + _check_cudnn_frontend_version_match(cudnn) + return cudnn + + +def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): + cudnn = _import_cudnn_for_score_mod() + + io_data_type = _cudnn_data_type(cudnn, q_aval.dtype) + graph = cudnn.pygraph( + io_data_type=io_data_type, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) + k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) + v_dim, v_stride = _bshd_as_bhsd_dim_stride(v_aval.shape) + q = graph.tensor( + name="q", dim=q_dim, stride=q_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_Q + ) + k = graph.tensor( + name="k", dim=k_dim, stride=k_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_K + ) + v = graph.tensor( + name="v", dim=v_dim, stride=v_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_V + ) + + score_mod_graph_tensors, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( + cudnn, + graph, + config.score_mod_tensor_names, + score_mod_avals, + config.score_mod_scalars, + _SCORE_MOD_FWD_TENSOR_UID_BASE, + _SCORE_MOD_FWD_SCALAR_UID_BASE, + ) + + output, stats = graph.sdpa( + name="te_score_mod_sdpa", + q=q, + k=k, + v=v, + generate_stats=config.is_training, + attn_scale=config.scaling_factor, + use_causal_mask=False, + score_mod=_wrap_score_mod(config.score_mod, score_mod_graph_tensors), + ) + + batch, q_seqlen, q_heads, _ = q_aval.shape + _, _, _, v_head_dim = v_aval.shape + output_dim, output_stride = _bshd_as_bhsd_dim_stride((batch, q_seqlen, q_heads, v_head_dim)) + output.set_output(True).set_uid(_SCORE_MOD_UID_O).set_dim(output_dim).set_stride(output_stride) + output.set_data_type(io_data_type) + + output_uids = [_SCORE_MOD_UID_O] + if config.is_training: + stats_shape = (batch, q_heads, q_seqlen, 1) + stats.set_output(True).set_uid(_SCORE_MOD_UID_STATS).set_dim(stats_shape).set_stride( + _row_major_stride(stats_shape) + ) + stats.set_data_type(cudnn.data_type.FLOAT) + output_uids.append(_SCORE_MOD_UID_STATS) + + workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph(cudnn, graph) + return _serialized_score_mod_graph( + serialized_graph=serialized_graph, + cudnn_frontend_version=frontend_version, + workspace_size=workspace_size, + input_uids=[_SCORE_MOD_UID_Q, _SCORE_MOD_UID_K, _SCORE_MOD_UID_V, *tensor_uids], + output_uids=output_uids, + scalar_uids=scalar_uids, + scalar_values=scalar_values, + ) + + +def _build_score_mod_bwd_graph( + q_aval, + k_aval, + v_aval, + output_aval, + doutput_aval, + stats_aval, + score_mod_avals, + score_mod_bprop_avals, + config, +): + cudnn = _import_cudnn_for_score_mod() + + io_data_type = _cudnn_data_type(cudnn, q_aval.dtype) + graph = cudnn.pygraph( + io_data_type=io_data_type, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) + k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) + v_dim, v_stride = _bshd_as_bhsd_dim_stride(v_aval.shape) + o_dim, o_stride = _bshd_as_bhsd_dim_stride(output_aval.shape) + do_dim, do_stride = _bshd_as_bhsd_dim_stride(doutput_aval.shape) + q = graph.tensor( + name="q", dim=q_dim, stride=q_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_Q + ) + k = graph.tensor( + name="k", dim=k_dim, stride=k_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_K + ) + v = graph.tensor( + name="v", dim=v_dim, stride=v_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_V + ) + output = graph.tensor( + name="o", dim=o_dim, stride=o_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_O + ) + doutput = graph.tensor( + name="dO", dim=do_dim, stride=do_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_DO + ) + stats = graph.tensor( + name="stats", + dim=tuple(int(dim) for dim in stats_aval.shape), + stride=_row_major_stride(stats_aval.shape), + data_type=cudnn.data_type.FLOAT, + uid=_SCORE_MOD_UID_STATS, + ) + + score_mod_graph_tensors, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( + cudnn, + graph, + config.score_mod_tensor_names, + score_mod_avals, + config.score_mod_scalars, + _SCORE_MOD_FWD_TENSOR_UID_BASE, + _SCORE_MOD_FWD_SCALAR_UID_BASE, + ) + ( + score_mod_bprop_graph_tensors, + bprop_tensor_uids, + bprop_scalar_uids, + bprop_scalar_values, + ) = _score_mod_graph_tensors( + cudnn, + graph, + config.score_mod_bprop_tensor_names, + score_mod_bprop_avals, + config.score_mod_bprop_scalars, + _SCORE_MOD_BPROP_TENSOR_UID_BASE, + _SCORE_MOD_BPROP_SCALAR_UID_BASE, + ) + + dq, dk, dv = graph.sdpa_backward( + name="te_score_mod_sdpa_backward", + q=q, + k=k, + v=v, + o=output, + dO=doutput, + stats=stats, + attn_scale=config.scaling_factor, + use_causal_mask=False, + score_mod=_wrap_score_mod(config.score_mod, score_mod_graph_tensors), + score_mod_bprop=_wrap_score_mod(config.score_mod_bprop, score_mod_bprop_graph_tensors), + use_deterministic_algorithm=config.deterministic, + ) + + dq.set_output(True).set_uid(_SCORE_MOD_UID_DQ).set_dim(q_dim).set_stride(q_stride) + dk.set_output(True).set_uid(_SCORE_MOD_UID_DK).set_dim(k_dim).set_stride(k_stride) + dv.set_output(True).set_uid(_SCORE_MOD_UID_DV).set_dim(v_dim).set_stride(v_stride) + + workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph(cudnn, graph) + return _serialized_score_mod_graph( + serialized_graph=serialized_graph, + cudnn_frontend_version=frontend_version, + workspace_size=workspace_size, + input_uids=[ + _SCORE_MOD_UID_Q, + _SCORE_MOD_UID_K, + _SCORE_MOD_UID_V, + _SCORE_MOD_UID_O, + _SCORE_MOD_UID_DO, + _SCORE_MOD_UID_STATS, + *tensor_uids, + *bprop_tensor_uids, + ], + output_uids=[_SCORE_MOD_UID_DQ, _SCORE_MOD_UID_DK, _SCORE_MOD_UID_DV], + scalar_uids=[*scalar_uids, *bprop_scalar_uids], + scalar_values=[*scalar_values, *bprop_scalar_values], + ) + + +def _fused_attn_score_mod_fwd( + qkv: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], + score_mod_tensors: Tuple[jnp.ndarray, ...], + config: _FusedAttnScoreModConfig, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Run cuDNN frontend SDPA forward with a score_mod callback.""" + q, k, v = qkv + q_aval, k_aval, v_aval = map(_shape_dtype, (q, k, v)) + score_mod_avals = tuple(_shape_dtype(arg) for arg in score_mod_tensors) + key = _graph_cache_key("fwd", config, (q_aval, k_aval, v_aval, *score_mod_avals)) + if key is None: + graph = _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config) + else: + if key not in _score_mod_graph_cache: + _score_mod_graph_cache[key] = _build_score_mod_fwd_graph( + q_aval, k_aval, v_aval, score_mod_avals, config + ) + graph = _score_mod_graph_cache[key] + + batch, q_seqlen, q_heads, _ = q.shape + _, _, _, v_head_dim = v.shape + output_shape = jax.ShapeDtypeStruct((batch, q_seqlen, q_heads, v_head_dim), q.dtype) + stats_shape = (batch, q_heads, q_seqlen, 1) if config.is_training else (0,) + stats = jax.ShapeDtypeStruct(stats_shape, jnp.float32) + workspace = jax.ShapeDtypeStruct((graph.workspace_size,), jnp.uint8) + output, softmax_stats, _ = ffi.ffi_call( + "te_fused_attn_score_mod_forward_ffi", + (output_shape, stats, workspace), + )( + q, + k, + v, + *score_mod_tensors, + serialized_graph=graph.serialized_graph, + graph_hash0=graph.graph_hash[0], + graph_hash1=graph.graph_hash[1], + cudnn_frontend_version=graph.cudnn_frontend_version, + input_uids=graph.input_uids, + output_uids=graph.output_uids, + scalar_uids=graph.scalar_uids, + scalar_sizes=graph.scalar_sizes, + scalar_values=graph.scalar_values, + ) + return output, softmax_stats + + +def _fused_attn_score_mod_bwd( + qkv: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], + output: jnp.ndarray, + doutput: jnp.ndarray, + softmax_stats: jnp.ndarray, + score_mod_tensors: Tuple[jnp.ndarray, ...], + score_mod_bprop_tensors: Tuple[jnp.ndarray, ...], + config: _FusedAttnScoreModConfig, +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Run cuDNN frontend SDPA backward with score_mod callbacks.""" + if not config.is_training: + raise RuntimeError("score_mod backward requires fused_attn(..., is_training=True).") + + q, k, v = qkv + all_inputs = (q, k, v, output, doutput, softmax_stats, *score_mod_tensors) + all_inputs = (*all_inputs, *score_mod_bprop_tensors) + avals = tuple(_shape_dtype(arg) for arg in all_inputs) + key = _graph_cache_key("bwd", config, avals) + if key is None: + graph = _build_score_mod_bwd_graph( + *avals[:6], + avals[6 : 6 + len(score_mod_tensors)], + avals[6 + len(score_mod_tensors) :], + config, + ) + else: + if key not in _score_mod_graph_cache: + _score_mod_graph_cache[key] = _build_score_mod_bwd_graph( + *avals[:6], + avals[6 : 6 + len(score_mod_tensors)], + avals[6 + len(score_mod_tensors) :], + config, + ) + graph = _score_mod_graph_cache[key] + + dq = jax.ShapeDtypeStruct(q.shape, q.dtype) + dk = jax.ShapeDtypeStruct(k.shape, k.dtype) + dv = jax.ShapeDtypeStruct(v.shape, v.dtype) + workspace = jax.ShapeDtypeStruct((graph.workspace_size,), jnp.uint8) + dq, dk, dv, _ = ffi.ffi_call( + "te_fused_attn_score_mod_backward_ffi", + (dq, dk, dv, workspace), + )( + q, + k, + v, + output, + doutput, + softmax_stats, + *score_mod_tensors, + *score_mod_bprop_tensors, + serialized_graph=graph.serialized_graph, + graph_hash0=graph.graph_hash[0], + graph_hash1=graph.graph_hash[1], + cudnn_frontend_version=graph.cudnn_frontend_version, + input_uids=graph.input_uids, + output_uids=graph.output_uids, + scalar_uids=graph.scalar_uids, + scalar_sizes=graph.scalar_sizes, + scalar_values=graph.scalar_values, + ) + return dq, dk, dv + + +class FusedAttnScoreModHelper: + """Helper for cuDNN frontend score_mod fused attention graphs.""" + + make_config = staticmethod(_make_fused_attn_score_mod_config) + forward = staticmethod(_fused_attn_score_mod_fwd) + backward = staticmethod(_fused_attn_score_mod_bwd) + + +def make_fused_attn_score_mod_config( + score_mod: Callable, + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Mapping[str, Any]], + score_mod_bprop_tensors: Optional[Mapping[str, Any]], + scaling_factor: float, + is_training: bool, +) -> Tuple[_FusedAttnScoreModConfig, Tuple[jnp.ndarray, ...], Tuple[jnp.ndarray, ...]]: + """Normalize score_mod operands and create a static graph-build config.""" + return FusedAttnScoreModHelper.make_config( + score_mod, + score_mod_bprop, + score_mod_tensors, + score_mod_bprop_tensors, + scaling_factor, + is_training, + ) + + +def fused_attn_score_mod_fwd( + qkv: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], + score_mod_tensors: Tuple[jnp.ndarray, ...], + config: _FusedAttnScoreModConfig, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Run cuDNN frontend SDPA forward with a score_mod callback.""" + return FusedAttnScoreModHelper.forward(qkv, score_mod_tensors, config) + + +def fused_attn_score_mod_bwd( + qkv: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], + output: jnp.ndarray, + doutput: jnp.ndarray, + softmax_stats: jnp.ndarray, + score_mod_tensors: Tuple[jnp.ndarray, ...], + score_mod_bprop_tensors: Tuple[jnp.ndarray, ...], + config: _FusedAttnScoreModConfig, +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Run cuDNN frontend SDPA backward with score_mod callbacks.""" + return FusedAttnScoreModHelper.backward( + qkv, + output, + doutput, + softmax_stats, + score_mod_tensors, + score_mod_bprop_tensors, + config, + ) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 4ff6d079864..212c8083ec0 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -23,7 +23,9 @@ get_num_compute_streams, JAXX_Collective_Op, get_device_compute_capability, + nvte_built_with_cublasmp, initialize_cgemm_communicator, + is_collective_gemm_with_cublasmp, get_cgemm_num_max_streams, get_grouped_gemm_setup_workspace_size, ) @@ -222,6 +224,7 @@ def collective_gemm_bootstrap( num_sm_for_communication=2, use_ce=True, aggregate_all_gather=False, + use_cublasmp=False, ): """Initialize NCCL communicators for Collective GEMM operations. @@ -260,6 +263,9 @@ def collective_gemm_bootstrap( Can improve performance by offloading memory operations. Default: True. aggregate_all_gather (bool, optional): Aggregate multiple small all-gather operations into larger ones for better efficiency. Default: False. + use_cublasmp (bool, optional): Use cuBLASMp backend for Collective GEMM overlap. + Requires Transformer Engine to be compiled with NVTE_WITH_CUBLASMP=1. + Default: False. Raises: AssertionError: If num_total_devices is not divisible by num_devices_per_process, @@ -295,6 +301,24 @@ def collective_gemm_bootstrap( This function must be called after JAX distributed initialization and before any collective GEMM operations. Each process should call this function with its own unique process_id. + + With the cuBLASMp backend, XLA command buffer capture must include + ``COLLECTIVES`` so that the NCCL calls inside cuBLASMp end up in the + same captured buffer as the CollectiveGemm custom call. Otherwise the + capture aborts with ``CUDA_ERROR_STREAM_CAPTURE_INVALIDATED``. Set the + flag before ``jax.distributed.initialize()``: + + import os + os.environ["XLA_FLAGS"] = ( + os.environ.get("XLA_FLAGS", "") + + " --xla_gpu_enable_command_buffer=+COLLECTIVES" + ) + + This is not required for non-overlapped collective GEMM (when + ``collective_op`` is ``CollectiveOp.NONE`` and JAX/XLA handles the + collective via its own graph-level optimization), nor for the + Userbuffers backend, which uses CUDA multicast APIs and async + memcpy on symmetric memory pointers that XLA already captures. """ if not (num_devices_per_process == 1 and jax.local_device_count() == 1): @@ -306,6 +330,12 @@ def collective_gemm_bootstrap( ) if not 0 <= process_id < num_total_devices: raise ValueError(f"Invalid process_id={process_id}") + if use_cublasmp and not nvte_built_with_cublasmp(): + raise RuntimeError( + "Collective GEMM with cuBLASMp backend was requested, but Transformer Engine " + "was not built with cuBLASMp support. Rebuild with NVTE_WITH_CUBLASMP=1 or " + "disable use_cublasmp." + ) initialize_cgemm_communicator( num_total_devices, num_devices_per_process, @@ -317,6 +347,7 @@ def collective_gemm_bootstrap( num_sm_for_communication, use_ce, aggregate_all_gather, + use_cublasmp, ) @@ -606,7 +637,11 @@ def _dims_are_consecutive(dims): if scaling_mode.is_nvfp4_scaling: workspace_size += lhs_scale_inv.size + rhs_scale_inv.size if not collective_op.is_none: - workspace_size *= get_cgemm_num_max_streams() + if is_collective_gemm_with_cublasmp(): + # cuBlasMp manages its own cuBlasLt workspaces per stream + workspace_size = 0 + else: + workspace_size *= get_cgemm_num_max_streams() # cuBLAS workspace ptr must be 256 bytes aligned but JAX buffers are not # necessarily 256 bytes aligned, we add some padding to ensure alignment. workspace_size += 256 @@ -812,10 +847,10 @@ def batcher( contracting_dims, scaling_mode, use_split_accumulator, - collective_op, transpose_batch_sequence, sequence_dim, is_outer, + collective_op, ): del transpose_batch_sequence, sequence_dim, is_outer if GemmPrimitive.outer_primitive is None: @@ -998,9 +1033,9 @@ def _parse_operand_output_specs( lhs_scale_specs = rhs_scale_specs = (None,) if scaling_mode.is_1d_block_scaling(): rhs_scale_specs = rhs_specs - # Set the seq spec to None to trigger AG the scales as TE/Common CGEMM does not handle - # scale collecting yet - if collective_op.is_all_gather: + # Set the seq spec to None to trigger AG the scales as TE/Common CGEMM w/ Userbuffers + # backend does not handle scale collecting yet (cuBLASMp backend does) + if collective_op.is_all_gather and not is_collective_gemm_with_cublasmp(): lhs_scale_specs = tuple( None if i == sequence_dim else s for i, s in enumerate(lhs_specs) ) @@ -1957,7 +1992,27 @@ def gemm( transpose_batch_sequence: bool, default = False Transpose the batch and sequence dimensions of the input tensor. collective_op: CollectiveOp, default = CollectiveOp.NONE - Collective operation type for collective GEMM. + Collective operation type for collective GEMM. When set to + ``CollectiveOp.ALL_GATHER`` or ``CollectiveOp.REDUCE_SCATTER``, the GEMM + is executed with communication overlap via the Userbuffers or cuBLASMp + backend (see :func:`collective_gemm_bootstrap`). + + .. note:: + Collective GEMM with communication overlap is captured into XLA + command buffers as a custom call. When executing with the cuBLASMp + backend, this captured graph spans NCCL collectives that XLA does not + include in command buffers by default, so add ``COLLECTIVES`` to the + enabled kinds before JAX initialization:: + + os.environ["XLA_FLAGS"] = ( + os.environ.get("XLA_FLAGS", "") + + " --xla_gpu_enable_command_buffer=+COLLECTIVES" + ) + + Without this, capture aborts with + ``CUDA_ERROR_STREAM_CAPTURE_INVALIDATED``. Not required when + ``collective_op`` is ``CollectiveOp.NONE`` or when using the Userbuffers + backend instead of cuBLASMp. Returns ------- diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index 0ae267cbf34..32454396892 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -7,13 +7,14 @@ import jax.numpy as jnp from jax import dtypes, ffi from jax.sharding import NamedSharding, PartitionSpec -from transformer_engine_jax import JAXX_Score_Function +from transformer_engine_jax import JAXX_Routing_Map_Format, JAXX_Score_Function from .base import BasePrimitive, register_primitive from .misc import get_padded_spec __all__ = [ "ScoreFunction", + "RoutingMapFormat", "fused_topk_with_score_function_fwd", "fused_topk_with_score_function_bwd", "fused_moe_aux_loss_fwd", @@ -28,6 +29,19 @@ class ScoreFunction(IntEnum): SOFTMAX = int(JAXX_Score_Function.SOFTMAX) +class RoutingMapFormat(IntEnum): + """Routing-map output layout, synced with C++ JAXX_Routing_Map_Format / NVTERoutingMapFormat. + + BYTEMAP — bool/uint8 tensor of shape [num_tokens, num_experts]. + BITMAP_U8 — uint8 tensor of shape [num_tokens, ceil(num_experts/8)]; bit + (e % 8) of byte (e / 8) of row t is 1 iff token t routes to + expert e (LSB-first packing along the expert axis). + """ + + BYTEMAP = int(JAXX_Routing_Map_Format.BYTEMAP) + BITMAP_U8 = int(JAXX_Routing_Map_Format.BITMAP_U8) + + # =========================================== ================================== # Fused Top-K with Score Function - Forward # ============================================================================= @@ -50,7 +64,9 @@ class FusedTopkWithScoreFunctionFwdPrimitive(BasePrimitive): 6, 7, 8, - ) # topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, compute_aux_scores + 9, + ) # topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, + # compute_aux_scores, routing_map_format inner_primitive = None outer_primitive = None @@ -65,6 +81,7 @@ def abstract( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): """Abstract evaluation: describe output shapes and dtypes.""" del expert_bias_aval, topk, use_pre_softmax, num_groups, group_topk @@ -72,7 +89,15 @@ def abstract( i_dtype = dtypes.canonicalize_dtype(logits_aval.dtype) i_shape = logits_aval.shape probs_aval = logits_aval.update(shape=i_shape, dtype=i_dtype) - routing_map_aval = logits_aval.update(shape=i_shape, dtype=jnp.bool_) + # routing_map shape/dtype depends on the format. In BITMAP_U8 mode the + # expert axis is bit-packed LSB-first into uint8 bytes, so the trailing + # dim becomes ceil(num_experts/8). + if int(routing_map_format) == int(RoutingMapFormat.BITMAP_U8): + packed_experts = (i_shape[-1] + 7) // 8 + routing_map_shape = (*i_shape[:-1], packed_experts) + routing_map_aval = logits_aval.update(shape=routing_map_shape, dtype=jnp.uint8) + else: + routing_map_aval = logits_aval.update(shape=i_shape, dtype=jnp.bool_) # The CUDA kernel always uses float32 (CompType) for intermediate # computations (softmax/sigmoid values saved for backward). intermediate_aval = logits_aval.update(shape=i_shape, dtype=jnp.float32) @@ -91,6 +116,7 @@ def lowering( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): return ffi.ffi_lowering(FusedTopkWithScoreFunctionFwdPrimitive.name)( ctx, @@ -103,6 +129,7 @@ def lowering( scaling_factor=scaling_factor, score_function=score_function, compute_aux_scores=compute_aux_scores, + routing_map_format=routing_map_format, ) @staticmethod @@ -116,6 +143,7 @@ def impl( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): if FusedTopkWithScoreFunctionFwdPrimitive.inner_primitive is None: raise RuntimeError( @@ -131,6 +159,7 @@ def impl( scaling_factor=scaling_factor, score_function=score_function, compute_aux_scores=compute_aux_scores, + routing_map_format=routing_map_format, ) @staticmethod @@ -145,6 +174,7 @@ def batcher( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): if FusedTopkWithScoreFunctionFwdPrimitive.outer_primitive is None: raise RuntimeError( @@ -163,6 +193,7 @@ def batcher( scaling_factor=scaling_factor, score_function=score_function, compute_aux_scores=compute_aux_scores, + routing_map_format=routing_map_format, ), (logits_bdim, logits_bdim, logits_bdim), ) @@ -176,6 +207,7 @@ def partition( scaling_factor, score_function, compute_aux_scores, + routing_map_format, mesh, arg_infos, result_infos, @@ -183,7 +215,14 @@ def partition( del result_infos logits_spec = get_padded_spec(arg_infos[0]) out_sharding = NamedSharding(mesh, PartitionSpec(*logits_spec)) - routing_sharding = NamedSharding(mesh, PartitionSpec(*logits_spec)) + # For bitmap mode the trailing dim is ceil(E/8) instead of E. We keep the + # routing_map sharded the same way logits is along all non-trailing dims + # and replicate the (now packed) expert axis to avoid sharding mid-byte. + if int(routing_map_format) == int(RoutingMapFormat.BITMAP_U8): + routing_spec = (*logits_spec[:-1], None) if len(logits_spec) >= 1 else logits_spec + else: + routing_spec = logits_spec + routing_sharding = NamedSharding(mesh, PartitionSpec(*routing_spec)) intermediate_sharding = NamedSharding(mesh, PartitionSpec(*logits_spec)) out_shardings = [out_sharding, routing_sharding, intermediate_sharding] arg_shardings = (arg_infos[0].sharding, arg_infos[1].sharding) @@ -199,13 +238,28 @@ def sharded_impl(logits, expert_bias): scaling_factor, score_function, compute_aux_scores, + routing_map_format, ) return mesh, sharded_impl, out_shardings, arg_shardings @staticmethod - def shardy_sharding_rule(*args): - del args + def shardy_sharding_rule(*args, **kwargs): + # Static args arrive in impl_static_args order: routing_map_format is the + # last (8th) static arg. Be defensive about positional-vs-kwarg passing + # across JAX versions. + routing_map_format = kwargs.get("routing_map_format") + if routing_map_format is None and len(args) >= 8: + routing_map_format = args[7] + # routing_map's expert axis is the same as logits in BYTEMAP mode; in + # BITMAP_U8 mode it's a packed-byte axis distinct from num_experts. + if routing_map_format is not None and int(routing_map_format) == int( + RoutingMapFormat.BITMAP_U8 + ): + return ( + "num_tokens num_experts, bias_dim ->" + " num_tokens num_experts, num_tokens packed_experts, num_tokens num_experts" + ) return ( "num_tokens num_experts, bias_dim -> num_tokens num_experts, num_tokens num_experts," " num_tokens num_experts" @@ -234,7 +288,8 @@ class FusedTopkWithScoreFunctionBwdPrimitive(BasePrimitive): 5, 6, 7, - ) # topk, use_pre_softmax, scaling_factor, score_function, compute_aux_scores + 8, + ) # topk, use_pre_softmax, scaling_factor, score_function, compute_aux_scores, routing_map_format inner_primitive = None outer_primitive = None @@ -248,9 +303,10 @@ def abstract( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): del topk, use_pre_softmax, scaling_factor, score_function - del compute_aux_scores, routing_map_aval + del compute_aux_scores, routing_map_aval, routing_map_format return intermediate_aval.update( shape=intermediate_aval.shape, dtype=dtypes.canonicalize_dtype(grad_probs_aval.dtype), @@ -268,6 +324,7 @@ def lowering( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): return ffi.ffi_lowering(FusedTopkWithScoreFunctionBwdPrimitive.name)( ctx, @@ -279,6 +336,7 @@ def lowering( scaling_factor=scaling_factor, score_function=score_function, compute_aux_scores=compute_aux_scores, + routing_map_format=routing_map_format, ) @staticmethod @@ -291,6 +349,7 @@ def impl( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): if FusedTopkWithScoreFunctionBwdPrimitive.inner_primitive is None: raise RuntimeError( @@ -305,6 +364,7 @@ def impl( scaling_factor=scaling_factor, score_function=score_function, compute_aux_scores=compute_aux_scores, + routing_map_format=routing_map_format, ) @staticmethod @@ -317,6 +377,7 @@ def batcher( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): if FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive is None: raise RuntimeError( @@ -334,6 +395,7 @@ def batcher( scaling_factor=scaling_factor, score_function=score_function, compute_aux_scores=compute_aux_scores, + routing_map_format=routing_map_format, ), grad_probs_bdim, ) @@ -345,11 +407,12 @@ def partition( scaling_factor, score_function, compute_aux_scores, + routing_map_format, mesh, arg_infos, result_infos, ): - del result_infos + del result_infos, routing_map_format 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) @@ -364,13 +427,24 @@ def sharded_impl(routing_map, intermediate, grad_probs): scaling_factor, score_function, compute_aux_scores, + routing_map_format, ) return mesh, sharded_impl, out_sharding, arg_shardings @staticmethod - def shardy_sharding_rule(*args): - del args + def shardy_sharding_rule(*args, **kwargs): + # routing_map_format is the 6th static arg (impl_static_args index 5). + routing_map_format = kwargs.get("routing_map_format") + if routing_map_format is None and len(args) >= 6: + routing_map_format = args[5] + if routing_map_format is not None and int(routing_map_format) == int( + RoutingMapFormat.BITMAP_U8 + ): + return ( + "num_tokens packed_experts, num_tokens num_experts, num_tokens num_experts ->" + " num_tokens num_experts" + ) return ( "num_tokens num_experts, num_tokens num_experts, num_tokens num_experts -> num_tokens" " num_experts" @@ -592,6 +666,7 @@ def fused_topk_with_score_function_fwd( score_function, expert_bias: jnp.ndarray, compute_aux_scores: bool = False, + routing_map_format: int = int(RoutingMapFormat.BYTEMAP), ): """ Fused top-k with score function forward pass. @@ -620,6 +695,9 @@ def fused_topk_with_score_function_fwd( Expert bias (only used with sigmoid). Pass empty array if unused. compute_aux_scores : bool If True, compute clean scores for aux loss instead of full top-k. + routing_map_format : int + RoutingMapFormat.BYTEMAP (default, bool[T, E]) or RoutingMapFormat.BITMAP_U8 + (uint8[T, ceil(E/8)], LSB-first along the expert axis). Returns ------- @@ -635,6 +713,7 @@ def fused_topk_with_score_function_fwd( scaling_factor=float(scaling_factor), score_function=int(score_function), compute_aux_scores=int(compute_aux_scores), + routing_map_format=int(routing_map_format), ) @@ -647,12 +726,17 @@ def fused_topk_with_score_function_bwd( scaling_factor: float, score_function, compute_aux_scores: bool = False, + routing_map_format: int = int(RoutingMapFormat.BYTEMAP), ): """ Fused top-k with score function backward pass. When compute_aux_scores=True, routing_map is ignored and the score-for-aux-loss backward kernel is used instead. + + routing_map_format must match the layout produced by the matching forward + call (BYTEMAP or BITMAP_U8). The CUDA kernel branches per-lane on this flag + when loading bits into shmem. """ return FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive.bind( routing_map, @@ -663,6 +747,7 @@ def fused_topk_with_score_function_bwd( scaling_factor=float(scaling_factor), score_function=int(score_function), compute_aux_scores=int(compute_aux_scores), + routing_map_format=int(routing_map_format), ) diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 416b18ada01..c0fa3acaebe 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -147,6 +147,10 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler); + +XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler); + NVTE_Fused_Attn_Backend GetFusedAttnBackend( bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, @@ -253,6 +257,7 @@ XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( // ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Score_Function); +XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Routing_Map_Format); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Collective_Op); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Quantize_Layout); diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index ed136d7b9eb..3fd6780d6d2 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -4,6 +4,19 @@ * See LICENSE for license information. ************************************************************************/ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + #include "../extensions.h" #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" @@ -689,5 +702,243 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnBackwardHandler, FusedAttnBackwardFFI, .Attrs(), FFI_CudaGraph_Traits); +namespace { + +struct ScoreModScalarStorage { + alignas(16) std::array data{}; + size_t size = 0; +}; + +struct ScoreModGraphCacheKey { + int device_id = 0; + int64_t hash0 = 0; + int64_t hash1 = 0; + int64_t frontend_version = 0; + + bool operator==(const ScoreModGraphCacheKey &other) const { + return device_id == other.device_id && hash0 == other.hash0 && hash1 == other.hash1 && + frontend_version == other.frontend_version; + } +}; + +struct ScoreModGraphCacheKeyHash { + size_t operator()(const ScoreModGraphCacheKey &key) const { + size_t seed = std::hash{}(key.device_id); + auto combine = [&seed](int64_t value) { + // 64-bit golden ratio constant from boost::hash_combine to spread mixed keys. + seed ^= std::hash{}(value) + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(key.hash0); + combine(key.hash1); + combine(key.frontend_version); + return seed; + } +}; + +using ScoreModGraphPtr = std::shared_ptr; + +std::unordered_map & +getScoreModeGraphCache() { + static std::unordered_map + cache; + return cache; +} + +std::mutex &getScoreModGraphCacheMutex() { + static std::mutex mutex; + return mutex; +} + +struct ScoreModCudnnHandleCache { + std::unordered_map handles; + + cudnnHandle_t GetHandle() { + int device_id = 0; + NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); + auto it = handles.find(device_id); + if (it == handles.end()) { + cudnnHandle_t handle = nullptr; + NVTE_CHECK_CUDNN(cudnnCreate(&handle)); + it = handles.emplace(device_id, handle).first; + } + return it->second; + } + + ~ScoreModCudnnHandleCache() { + for (auto &[_, handle] : handles) { + cudnnDestroy(handle); + } + } +}; + +cudnnHandle_t GetScoreModCudnnHandle() { + static thread_local ScoreModCudnnHandleCache cache; + return cache.GetHandle(); +} + +ScoreModGraphCacheKey GetScoreModGraphCacheKey(Dictionary &attrs) { + const int64_t frontend_version = get_attr_value(attrs, "cudnn_frontend_version"); + NVTE_CHECK(frontend_version == CUDNN_FRONTEND_VERSION, + "cuDNN frontend version mismatch for score_mod graph deserialization: graph was " + "serialized with Python cuDNN frontend version ", + frontend_version, + ", but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION ", + CUDNN_FRONTEND_VERSION, "."); + + int device_id = 0; + NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); + return ScoreModGraphCacheKey{ + device_id, + get_attr_value(attrs, "graph_hash0"), + get_attr_value(attrs, "graph_hash1"), + frontend_version, + }; +} + +ScoreModGraphPtr GetScoreModGraph(cudaStream_t stream, Dictionary &attrs) { + const auto key = GetScoreModGraphCacheKey(attrs); + { + std::lock_guard lock(getScoreModGraphCacheMutex()); + auto &cache = getScoreModeGraphCache(); + auto it = cache.find(key); + if (it != cache.end()) { + return it->second; + } + } + + const auto serialized_graph = get_attr_value(attrs, "serialized_graph"); + std::vector serialized_data(serialized_graph.begin(), serialized_graph.end()); + + auto handle = GetScoreModCudnnHandle(); + NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); + + auto graph = std::make_shared(); + auto status = graph->deserialize(handle, serialized_data); + NVTE_CHECK(status.is_good(), + "Failed to deserialize cuDNN score_mod SDPA graph: ", status.get_message()); + + std::lock_guard lock(getScoreModGraphCacheMutex()); + auto &cache = getScoreModeGraphCache(); + auto it = cache.find(key); + if (it != cache.end()) { + return it->second; + } + cache.emplace(key, graph); + return graph; +} + +Error_Type ExecuteScoreModGraph(cudaStream_t stream, Dictionary &attrs, + const std::vector &input_ptrs, + const std::vector &output_ptrs, void *workspace) { + auto graph = GetScoreModGraph(stream, attrs); + auto input_uids = get_attr_value>(attrs, "input_uids"); + auto output_uids = get_attr_value>(attrs, "output_uids"); + auto scalar_uids = get_attr_value>(attrs, "scalar_uids"); + auto scalar_sizes = get_attr_value>(attrs, "scalar_sizes"); + auto scalar_values = get_attr_value>(attrs, "scalar_values"); + + NVTE_CHECK(input_ptrs.size() == input_uids.size(), "cuDNN score_mod graph expected ", + input_uids.size(), " inputs but got ", input_ptrs.size()); + NVTE_CHECK(output_ptrs.size() >= output_uids.size(), "cuDNN score_mod graph expected at least ", + output_uids.size(), " outputs but got ", output_ptrs.size()); + NVTE_CHECK(scalar_uids.size() == scalar_sizes.size(), + "Mismatched score_mod scalar uid/value-size counts."); + NVTE_CHECK(scalar_values.size() == scalar_uids.size() * 16, + "Mismatched score_mod packed scalar value size."); + + std::unordered_map variant_pack; + for (size_t i = 0; i < input_uids.size(); ++i) { + variant_pack.emplace(input_uids[i], input_ptrs[i]); + } + for (size_t i = 0; i < output_uids.size(); ++i) { + variant_pack.emplace(output_uids[i], output_ptrs[i]); + } + + std::vector scalar_storage(scalar_uids.size()); + for (size_t i = 0; i < scalar_uids.size(); ++i) { + NVTE_CHECK(scalar_sizes[i] >= 0 && scalar_sizes[i] <= 16, + "score_mod pass-by-value scalars must be at most 16 bytes."); + scalar_storage[i].size = static_cast(scalar_sizes[i]); + std::copy_n(scalar_values.begin() + i * 16, 16, scalar_storage[i].data.begin()); + variant_pack.emplace(scalar_uids[i], scalar_storage[i].data.data()); + } + + auto handle = GetScoreModCudnnHandle(); + NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); + auto status = graph->execute(handle, variant_pack, workspace); + NVTE_CHECK(status.is_good(), + "cuDNN score_mod SDPA graph execution failed: ", status.get_message()); + return ffi_with_cuda_error_check(); +} + +void AppendRemainingBuffers(Variadic_Buffer_Type args, std::vector *ptrs) { + ptrs->reserve(ptrs->size() + args.size()); + for (size_t i = 0; i < args.size(); ++i) { + auto maybe_buf = args.get(i); + NVTE_CHECK(!maybe_buf.has_error(), "Failed to decode variadic score_mod input buffer."); + ptrs->push_back(maybe_buf.value().untyped_data()); + } +} + +} // namespace + +Error_Type FusedAttnScoreModForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, + Buffer_Type v_buf, Variadic_Buffer_Type score_mod_args, + Result_Type output_buf, Result_Type stats_buf, + Result_Type workspace_buf, Dictionary attrs) { + std::vector input_ptrs = {q_buf.untyped_data(), k_buf.untyped_data(), + v_buf.untyped_data()}; + AppendRemainingBuffers(score_mod_args, &input_ptrs); + + std::vector output_ptrs = {output_buf->untyped_data(), stats_buf->untyped_data()}; + return ExecuteScoreModGraph(stream, attrs, input_ptrs, output_ptrs, + workspace_buf->untyped_data()); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler, FusedAttnScoreModForwardFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // q + .Arg() // k + .Arg() // v + .RemainingArgs() // score_mod tensor operands + .Ret() // output + .Ret() // stats + .Ret() // workspace + .Attrs()); + +Error_Type FusedAttnScoreModBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, + Buffer_Type v_buf, Buffer_Type output_buf, + Buffer_Type doutput_buf, Buffer_Type stats_buf, + Variadic_Buffer_Type score_mod_args, Result_Type dq_buf, + Result_Type dk_buf, Result_Type dv_buf, + Result_Type workspace_buf, Dictionary attrs) { + std::vector input_ptrs = {q_buf.untyped_data(), k_buf.untyped_data(), + v_buf.untyped_data(), output_buf.untyped_data(), + doutput_buf.untyped_data(), stats_buf.untyped_data()}; + AppendRemainingBuffers(score_mod_args, &input_ptrs); + + std::vector output_ptrs = {dq_buf->untyped_data(), dk_buf->untyped_data(), + dv_buf->untyped_data()}; + return ExecuteScoreModGraph(stream, attrs, input_ptrs, output_ptrs, + workspace_buf->untyped_data()); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler, FusedAttnScoreModBackwardFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // q + .Arg() // k + .Arg() // v + .Arg() // output + .Arg() // doutput + .Arg() // stats + .RemainingArgs() // score_mod tensor operands + .Ret() // dq + .Ret() // dk + .Ret() // dv + .Ret() // workspace + .Attrs()); + } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp b/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp index 36a4a068a4d..da25e3676a5 100644 --- a/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp +++ b/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp @@ -18,11 +18,9 @@ ncclUniqueId CommunicatorHandler::coordinate_nccl_unique_id(const std::string &i int tp_domain_id = get_tp_domain_id(); bool is_tp_leader = (get_local_device_id_within_tp_domain() == 0); - pid_t pgid = getpgid(0); - std::string base_path = getenv("NVTE_JAX_NCCL_FILE_PATH", "/tmp"); - std::string id_file = base_path + "/nccl_" + id_type + "_unique_id_pgid_" + std::to_string(pgid) + - "_" + std::to_string(num_total_devices) + "_" + std::to_string(tp_size) + + std::string id_file = base_path + "/nccl_" + id_type + "_unique_id_" + + std::to_string(num_total_devices) + "_" + std::to_string(tp_size) + "_domain_" + std::to_string(tp_domain_id) + ".bin"; if (is_tp_leader) { @@ -136,7 +134,7 @@ void CommunicatorHandler::init(int num_total_devices, int num_devices_per_proces handler._initialize = true; - // Bootstrap UB via creating a dummy CommOverlapP2PBase object + // Bootstrap UB/cuBlasMp via creating a dummy CommOverlapP2PBase object std::vector buffer_shape{1, 1}; [[maybe_unused]] auto _ = CollectiveGemmPlanRegistry::getInstance().get_executor( buffer_shape, DType::kFloat32, JAXX_Collective_Op::ALL_GATHER); @@ -144,14 +142,20 @@ void CommunicatorHandler::init(int num_total_devices, int num_devices_per_proces void InitializeCgemmCommunicator(int num_total_devices, int num_devices_per_process, int process_id, int tp_size, int num_max_streams, int gemm_priority, - int comm_priority, int num_comm_sm, bool use_ce, - bool aggregate_ag) { + int comm_priority, int num_comm_sm, bool use_ce, bool aggregate_ag, + bool use_cublasmp) { auto &config = CgemmConfig::get(false); - config.init(num_max_streams, gemm_priority, comm_priority, num_comm_sm, use_ce, aggregate_ag); + config.init(num_max_streams, gemm_priority, comm_priority, num_comm_sm, use_ce, aggregate_ag, + use_cublasmp); auto &handler = CommunicatorHandler::get(false); handler.init(num_total_devices, num_devices_per_process, process_id, tp_size); } +bool IsCollectiveGemmWithCublasmp() { + auto &config = CgemmConfig::get(); + return config.use_cublasmp; +} + int GetCgemmNumMaxStreams() { auto &config = CgemmConfig::get(); return config.num_max_streams; @@ -168,7 +172,8 @@ CommOverlapCore *CollectiveGemmPlanRegistry::get_executor(std::vector bu hash_combine(plan_id, buffer_shape[0], buffer_shape[1], static_cast(dtype), static_cast(collective_op), comm_handler.tp_size, cgemm_config.num_max_streams, cgemm_config.gemm_priority, cgemm_config.comm_priority, cgemm_config.num_comm_sm, - cgemm_config.use_ce, cgemm_config.aggregate_ag, device_idx); + cgemm_config.use_ce, cgemm_config.aggregate_ag, device_idx, + cgemm_config.use_cublasmp); auto it = plan_map.find(plan_id); if (it != plan_map.end()) { @@ -192,14 +197,22 @@ CommOverlapCore *CollectiveGemmPlanRegistry::get_executor(std::vector bu } std::unique_ptr executor; - executor = std::make_unique( - buffer_shape, dtype, comm_handler.get_global_rank(), comm_handler.num_total_devices, - comm_handler.get_local_device_id_within_tp_domain(), comm_handler.tp_size, - comm_handler.get_tp_domain_id(), comm_handler.get_tp_num_domains(), comm_handler.tp_size, - comm_handler.allgather_func, comm_handler.barrier_func, get_nvte_collective_op(collective_op), - cgemm_config.num_max_streams, 1 /*comm_cga_size*/, cgemm_config.gemm_priority, - cgemm_config.comm_priority, cgemm_config.num_comm_sm, true /*set_sm_margin*/, - cgemm_config.use_ce, false /*atomic_gemm*/, cgemm_config.aggregate_ag); + if (cgemm_config.use_cublasmp) { + executor = std::make_unique( + comm_handler.get_comm_for_current_device(), + comm_handler.get_local_device_id_within_tp_domain(), comm_handler.tp_size, + cgemm_config.num_comm_sm, false /*atomic_gemm*/); + } else { + executor = std::make_unique( + buffer_shape, dtype, comm_handler.get_global_rank(), comm_handler.num_total_devices, + comm_handler.get_local_device_id_within_tp_domain(), comm_handler.tp_size, + comm_handler.get_tp_domain_id(), comm_handler.get_tp_num_domains(), comm_handler.tp_size, + comm_handler.allgather_func, comm_handler.barrier_func, + get_nvte_collective_op(collective_op), cgemm_config.num_max_streams, 1 /*comm_cga_size*/, + cgemm_config.gemm_priority, cgemm_config.comm_priority, cgemm_config.num_comm_sm, + true /*set_sm_margin*/, cgemm_config.use_ce, false /*atomic_gemm*/, + cgemm_config.aggregate_ag); + } CommOverlapCore *executor_ptr = executor.get(); plan_map[plan_id] = std::move(executor); diff --git a/transformer_engine/jax/csrc/extensions/cgemm_helper.h b/transformer_engine/jax/csrc/extensions/cgemm_helper.h index 2b980e7ee49..9bc8c9cf8d4 100644 --- a/transformer_engine/jax/csrc/extensions/cgemm_helper.h +++ b/transformer_engine/jax/csrc/extensions/cgemm_helper.h @@ -35,9 +35,10 @@ class CgemmConfig { int num_comm_sm; bool use_ce; bool aggregate_ag; + bool use_cublasmp; static void init(int _num_max_streams, int _gemm_priority, int _comm_priority, int _num_comm_sm, - bool _use_ce, bool _aggregate_ag) { + bool _use_ce, bool _aggregate_ag, bool _use_cublasmp = false) { auto &config = get(false); config._initialized = true; config.num_max_streams = _num_max_streams; @@ -46,6 +47,7 @@ class CgemmConfig { config.num_comm_sm = _num_comm_sm; config.use_ce = _use_ce; config.aggregate_ag = _aggregate_ag; + config.use_cublasmp = _use_cublasmp; } static CgemmConfig &get(bool is_initialized = true) { @@ -178,8 +180,10 @@ class CollectiveGemmPlanRegistry { // Function declarations void InitializeCgemmCommunicator(int num_total_devices, int num_devices_per_process, int process_id, int tp_size, int num_max_streams, int gemm_priority, - int comm_priority, int num_comm_sm, bool use_ce, - bool aggregate_ag); + int comm_priority, int num_comm_sm, bool use_ce, bool aggregate_ag, + bool use_cublasmp = false); + +bool IsCollectiveGemmWithCublasmp(); int GetCgemmNumMaxStreams(); diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 6ca907032ce..a36bf6cd222 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -99,6 +99,21 @@ std::tuple> xla_buffer_to_nvte_gemm_operand( return std::make_tuple(std::move(input), input_shape); } +// Build a TensorWrapper for the prepare stage. Operand contents are +// uninitialized at this point and no kernels are launched. +static TensorWrapper prepare_operand_tensor(Buffer_Type buf, Buffer_Type scale_inv, + JAXX_Scaling_Mode scaling_mode, + const std::vector &flat_shape) { + auto dtype = convert_ffi_datatype_to_te_dtype(buf.element_type()); + TensorWrapper t(get_nvte_scaling_mode(scaling_mode)); + t.set_rowwise_data(buf.untyped_data(), dtype, flat_shape); + if (scaling_mode != JAXX_Scaling_Mode::NO_SCALING && scale_inv.element_count() > 0) { + auto scale_dtype = convert_ffi_datatype_to_te_dtype(scale_inv.element_type()); + t.set_rowwise_scale_inv(scale_inv.untyped_data(), scale_dtype, std::vector{1}); + } + return t; +} + Error_Type GemmInitV2FFI(Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type rhs, Buffer_Type rhs_scale_inv, Buffer_Type bias, Buffer_Type alpha, Buffer_Type beta, Result_Type output, Result_Type workspace, @@ -128,8 +143,49 @@ Error_Type GemmInitV2FFI(Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type buffer_shape[0] = out_shape[0]; buffer_shape[1] = out_shape[1]; } - [[maybe_unused]] auto _ = CollectiveGemmPlanRegistry::getInstance().get_executor( + auto *executor = CollectiveGemmPlanRegistry::getInstance().get_executor( buffer_shape, buffer_dtype, config.collective_op); + + // Run a dummy cuBLASMp matmul in the prepare stage so its lazy NCCL + // window registration and workspace allocation happen outside any + // CUDA-graph capture. The throwaway D output gets overwritten on the + // next execute-stage call. + if (IsCollectiveGemmWithCublasmp() && executor != nullptr) { + auto lhs_ = prepare_operand_tensor(lhs, lhs_scale_inv, config.scaling_mode, lhs_shape); + auto rhs_ = prepare_operand_tensor(rhs, rhs_scale_inv, config.scaling_mode, rhs_shape); + // Match GemmV2FFI's local D shape: AG gathers along axis 0, RS scatters along axis 0. + std::vector d_shape = out_shape; + if (config.collective_op == JAXX_Collective_Op::ALL_GATHER) { + d_shape[0] *= comm_handler.tp_size; + } else if (config.collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { + d_shape[0] /= comm_handler.tp_size; + } + TensorWrapper d_(get_nvte_scaling_mode(JAXX_Scaling_Mode::NO_SCALING)); + d_.set_rowwise_data(output->untyped_data(), + convert_ffi_datatype_to_te_dtype(output->element_type()), d_shape); + TensorWrapper bias_(get_nvte_scaling_mode(JAXX_Scaling_Mode::NO_SCALING)); + if (bias.element_count() > 0) { + bias_.set_rowwise_data(bias.untyped_data(), + convert_ffi_datatype_to_te_dtype(bias.element_type()), + std::vector{static_cast(bias.element_count())}); + } + TensorWrapper pre_gelu_out_(get_nvte_scaling_mode(JAXX_Scaling_Mode::NO_SCALING)); + TensorWrapper dummy; + // Match GemmV2FFI's operand swap: rhs becomes A, lhs becomes B. + cudaStream_t prepare_stream = cudaStreamPerThread; + if (config.collective_op == JAXX_Collective_Op::ALL_GATHER) { + executor->split_overlap_ag(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, d_, + bias_, pre_gelu_out_, dummy, false /*grad*/, + false /*accumulate*/, false /*use_split_accumulator*/, dummy, + prepare_stream); + } else if (config.collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { + executor->split_overlap_rs(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, d_, + bias_, pre_gelu_out_, dummy, false /*grad*/, + false /*accumulate*/, false /*use_split_accumulator*/, dummy, + prepare_stream); + } + NVTE_CHECK_CUDA(cudaStreamSynchronize(prepare_stream)); + } } return ffi_with_cuda_error_check(); } @@ -191,8 +247,13 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(CollectiveGemmInitHandler, CollectiveGemmInitFFI, .Attr("fuse_gelu") .Attr("grad") .Attr("use_split_accumulator") - .Attr("collective_op"), - FFI_CudaGraph_Traits); + .Attr("collective_op") +#ifndef NVTE_WITH_CUBLASMP + // enable CUDA graphs only when cuBLASMp is NOT enabled + , + FFI_CudaGraph_Traits +#endif +); Error_Type GemmV2FFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type rhs, Buffer_Type rhs_scale_inv, Buffer_Type bias, @@ -300,18 +361,27 @@ Error_Type GemmV2FFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale buffer_shape, buffer_dtype, config.collective_op); auto pre_gelu_ = TensorWrapper(nullptr, std::vector{0}, DType::kByte); if (config.collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { - auto ubuf_out_ = TensorWrapper(executor->get_ubuf_dptr(), buffer_shape, out_dtype); - // Prepare the auxiliary buffer for the reduce-scattered GEMM output auto out_ = TensorWrapper(output->untyped_data(), out_shape, out_dtype); NVTE_CHECK(out_.numel() == output->element_count(), "cuBLAS GEMM output buffer size is incorrect, expected ", out_.numel(), " elements ", to_string_like(out_shape), " but got ", output->element_count(), " elements ", to_string_like(output->dimensions())); - // Launch GEMM+RS - executor->split_overlap_rs(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, - ubuf_out_, bias_, pre_gelu_, workspace_, false /*grad*/, - false /*accumulate*/, config.use_split_accumulator, out_, stream); + if (IsCollectiveGemmWithCublasmp()) { + // cuBLASMp writes the reduce-scattered result directly into D + auto rs_out_ = TensorWrapper(nullptr, std::vector{0}, out_dtype); + executor->split_overlap_rs(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, out_, + bias_, pre_gelu_, workspace_, false /*grad*/, + false /*accumulate*/, config.use_split_accumulator, rs_out_, + stream); + } else { + // Userbuffers writes the full GEMM result into ubuf, then reduce-scatters into rs_output + auto ubuf_out_ = TensorWrapper(executor->get_ubuf_dptr(), buffer_shape, out_dtype); + executor->split_overlap_rs(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, + ubuf_out_, bias_, pre_gelu_, workspace_, false /*grad*/, + false /*accumulate*/, config.use_split_accumulator, out_, + stream); + } } else if (config.collective_op == JAXX_Collective_Op::ALL_GATHER) { auto aux_out_ = TensorWrapper(nullptr, std::vector{0}, out_dtype); // Empty diff --git a/transformer_engine/jax/csrc/extensions/misc.h b/transformer_engine/jax/csrc/extensions/misc.h index c6f6f87cb4a..62f95ad5290 100644 --- a/transformer_engine/jax/csrc/extensions/misc.h +++ b/transformer_engine/jax/csrc/extensions/misc.h @@ -4,6 +4,7 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include @@ -127,6 +128,14 @@ enum class JAXX_Score_Function : int64_t { SOFTMAX = 1, }; +// Mirror of NVTERoutingMapFormat for JAX FFI plumbing. Values are taken +// directly from the C enum in +// transformer_engine/common/include/transformer_engine/fused_router.h. +enum class JAXX_Routing_Map_Format : int64_t { + BYTEMAP = NVTE_ROUTING_MAP_FORMAT_BYTEMAP, + BITMAP_U8 = NVTE_ROUTING_MAP_FORMAT_BITMAP_U8, +}; + enum class JAXX_Collective_Op : int64_t { NONE = 0, ALL_GATHER = 1, diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 70d0403b3e5..2432f650051 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -64,6 +64,12 @@ pybind11::dict Registrations() { dict["te_fused_attn_backward_ffi"] = pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), pybind11::arg("execute") = EncapsulateFFI(FusedAttnBackwardHandler)); + dict["te_fused_attn_score_mod_forward_ffi"] = + pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), + pybind11::arg("execute") = EncapsulateFFI(FusedAttnScoreModForwardHandler)); + dict["te_fused_attn_score_mod_backward_ffi"] = + pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), + pybind11::arg("execute") = EncapsulateFFI(FusedAttnScoreModBackwardHandler)); // GEMM dict["te_gemm_ffi"] = @@ -112,6 +118,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("get_fused_attn_backend", &GetFusedAttnBackend); m.def("get_cuda_version", &GetCudaRuntimeVersion); m.def("get_cudnn_version", &GetCudnnRuntimeVersion); + m.def("get_cudnn_frontend_version", &GetCudnnFrontendVersion); m.def("get_device_compute_capability", &GetDeviceComputeCapability); m.def("get_num_compute_streams", &nvte_get_num_compute_streams); m.def("get_cublasLt_version", &cublasLtGetVersion); @@ -124,7 +131,9 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("get_topk_workspace_sizes", &GetTopkWorkspaceSizes); m.def("nvte_get_qkv_format", &nvte_get_qkv_format); m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported); + m.def("nvte_built_with_cublasmp", &::nvte_built_with_cublasmp); m.def("initialize_cgemm_communicator", &InitializeCgemmCommunicator); + m.def("is_collective_gemm_with_cublasmp", &IsCollectiveGemmWithCublasmp); m.def("get_cgemm_num_max_streams", &GetCgemmNumMaxStreams); m.def("get_grouped_gemm_setup_workspace_size", &nvte_get_grouped_gemm_setup_workspace_size); @@ -217,6 +226,11 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("SOFTMAX", JAXX_Score_Function::SOFTMAX) .export_values(); + pybind11::enum_(m, "JAXX_Routing_Map_Format", pybind11::module_local()) + .value("BYTEMAP", JAXX_Routing_Map_Format::BYTEMAP) + .value("BITMAP_U8", JAXX_Routing_Map_Format::BITMAP_U8) + .export_values(); + pybind11::enum_(m, "JAXX_Collective_Op", pybind11::module_local()) .value("NONE", JAXX_Collective_Op::NONE) .value("ALL_GATHER", JAXX_Collective_Op::ALL_GATHER) diff --git a/transformer_engine/jax/csrc/extensions/router.cpp b/transformer_engine/jax/csrc/extensions/router.cpp index 79daec3f074..a1d0731b6db 100644 --- a/transformer_engine/jax/csrc/extensions/router.cpp +++ b/transformer_engine/jax/csrc/extensions/router.cpp @@ -21,10 +21,11 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( Buffer_Type logits_buf, // [num_tokens, num_experts] Buffer_Type expert_bias_buf, // [num_experts] or empty Result_Type probs_buf, // [num_tokens, num_experts] (or scores when compute_aux_scores) - Result_Type routing_map_buf, // [num_tokens, num_experts] + Result_Type routing_map_buf, // BYTEMAP: [T, E] uint8 / BITMAP_U8: [T, ceil(E/8)] uint8 Result_Type intermediate_buf, // [num_tokens, num_experts] int64_t topk, int64_t use_pre_softmax, int64_t num_groups, int64_t group_topk, - double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores) { + double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores, + JAXX_Routing_Map_Format routing_map_format) { auto dtype = convert_ffi_datatype_to_te_dtype(logits_buf.element_type()); auto dims = logits_buf.dimensions(); auto num_tokens = static_cast(product(dims, 0, dims.size() - 1)); @@ -40,7 +41,15 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( std::vector{static_cast(num_tokens), static_cast(num_experts)}; auto logits_tensor = TensorWrapper(logits, flat_shape, dtype); auto probs_tensor = TensorWrapper(probs, flat_shape, dtype); - auto routing_map_tensor = TensorWrapper(routing_map, flat_shape, DType::kByte); + // 2D shape for the kernel: trailing dim = num_experts for BYTEMAP, + // ceil(num_experts/8) for BITMAP_U8. + auto routing_map_format_nvte = static_cast(routing_map_format); + size_t routing_map_trailing = (routing_map_format_nvte == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) + ? static_cast((num_experts + 7) / 8) + : static_cast(num_experts); + auto routing_map_shape = + std::vector{static_cast(num_tokens), routing_map_trailing}; + auto routing_map_tensor = TensorWrapper(routing_map, routing_map_shape, DType::kByte); // intermediate is always float32 (CompType) regardless of logits dtype. auto intermediate_dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf->element_type()); NVTE_CHECK( @@ -51,10 +60,10 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( auto intermediate_tensor = TensorWrapper(intermediate, flat_shape, DType::kFloat32); if (compute_aux_scores) { - nvte_fused_score_for_moe_aux_loss_forward( + nvte_fused_score_for_moe_aux_loss_forward_v2( logits_tensor.data(), num_tokens, num_experts, static_cast(topk), static_cast(score_function), probs_tensor.data(), routing_map_tensor.data(), - intermediate_tensor.data(), stream); + routing_map_format_nvte, intermediate_tensor.data(), stream); } else { auto bias_dims = expert_bias_buf.dimensions(); auto expert_bias_tensor = @@ -63,12 +72,12 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( convert_ffi_datatype_to_te_dtype(expert_bias_buf.element_type())) : TensorWrapper(); - nvte_fused_topk_with_score_function_forward( + nvte_fused_topk_with_score_function_forward_v2( logits_tensor.data(), num_tokens, num_experts, static_cast(topk), static_cast(use_pre_softmax), static_cast(num_groups), static_cast(group_topk), static_cast(scaling_factor), static_cast(score_function), expert_bias_tensor.data(), probs_tensor.data(), - routing_map_tensor.data(), intermediate_tensor.data(), stream); + routing_map_tensor.data(), routing_map_format_nvte, intermediate_tensor.data(), stream); } return ffi_with_cuda_error_check(); @@ -89,7 +98,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, .Attr("group_topk") .Attr("scaling_factor") .Attr("score_function") - .Attr("compute_aux_scores"), + .Attr("compute_aux_scores") + .Attr("routing_map_format"), FFI_CudaGraph_Traits); // ============================================================================ @@ -98,12 +108,13 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, Error_Type FusedTopkWithScoreFunctionBackwardFFI( cudaStream_t stream, - Buffer_Type routing_map_buf, // [num_tokens, num_experts] (unused when compute_aux_scores) + Buffer_Type routing_map_buf, // bytemap [T,E] uint8 or bitmap [T, ceil(E/8)] uint8 Buffer_Type intermediate_buf, // [num_tokens, num_experts] Buffer_Type grad_probs_buf, // [num_tokens, num_experts] (grad_scores when compute_aux_scores) Result_Type grad_logits_buf, // [num_tokens, num_experts] int64_t topk, int64_t use_pre_softmax, double scaling_factor, - JAXX_Score_Function score_function, int64_t compute_aux_scores) { + JAXX_Score_Function score_function, int64_t compute_aux_scores, + JAXX_Routing_Map_Format routing_map_format) { // intermediate is always float32 (CompType) regardless of logits dtype. auto intermediate_dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf.element_type()); NVTE_CHECK( @@ -130,14 +141,20 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( static_cast(score_function), grad_logits_tensor.data(), stream); } else { + auto routing_map_format_nvte = static_cast(routing_map_format); + size_t routing_map_trailing = (routing_map_format_nvte == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) + ? static_cast((num_experts + 7) / 8) + : static_cast(num_experts); + auto routing_map_shape = + std::vector{static_cast(num_tokens), routing_map_trailing}; auto routing_map_tensor = - TensorWrapper(routing_map_buf.untyped_data(), flat_shape, DType::kByte); + TensorWrapper(routing_map_buf.untyped_data(), routing_map_shape, DType::kByte); - nvte_fused_topk_with_score_function_backward( - routing_map_tensor.data(), intermediate_tensor.data(), grad_probs_tensor.data(), num_tokens, - num_experts, static_cast(topk), static_cast(use_pre_softmax), - static_cast(scaling_factor), static_cast(score_function), - grad_logits_tensor.data(), stream); + nvte_fused_topk_with_score_function_backward_v2( + routing_map_tensor.data(), routing_map_format_nvte, intermediate_tensor.data(), + grad_probs_tensor.data(), num_tokens, num_experts, static_cast(topk), + static_cast(use_pre_softmax), static_cast(scaling_factor), + static_cast(score_function), grad_logits_tensor.data(), stream); } return ffi_with_cuda_error_check(); @@ -155,7 +172,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler, .Attr("use_pre_softmax") .Attr("scaling_factor") .Attr("score_function") - .Attr("compute_aux_scores"), + .Attr("compute_aux_scores") + .Attr("routing_map_format"), FFI_CudaGraph_Traits); // ============================================================================ diff --git a/transformer_engine/jax/csrc/extensions/utils.cpp b/transformer_engine/jax/csrc/extensions/utils.cpp index 52ab2edf0f7..d8fe1762dd8 100644 --- a/transformer_engine/jax/csrc/extensions/utils.cpp +++ b/transformer_engine/jax/csrc/extensions/utils.cpp @@ -6,6 +6,7 @@ #include "utils.h" #include +#include #include @@ -22,6 +23,8 @@ int GetCudaRuntimeVersion() { size_t GetCudnnRuntimeVersion() { return cudnnGetVersion(); } +size_t GetCudnnFrontendVersion() { return CUDNN_FRONTEND_VERSION; } + int GetDeviceComputeCapability(int gpu_id) { return transformer_engine::cuda::sm_arch(gpu_id); } } // namespace jax diff --git a/transformer_engine/jax/csrc/extensions/utils.h b/transformer_engine/jax/csrc/extensions/utils.h index c55c8d86ce7..939b0a881e5 100644 --- a/transformer_engine/jax/csrc/extensions/utils.h +++ b/transformer_engine/jax/csrc/extensions/utils.h @@ -20,6 +20,7 @@ namespace jax { int GetCudaRuntimeVersion(); size_t GetCudnnRuntimeVersion(); +size_t GetCudnnFrontendVersion(); int GetDeviceComputeCapability(int gpu_id); class cudaDevicePropertiesManager { diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index a2e79208433..76922d2b55f 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -8,7 +8,7 @@ from enum import Enum from math import sqrt import os -from typing import Any, Callable, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Union import warnings import jax @@ -305,6 +305,9 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT context_checkpoint_name: str = "context" softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX + score_mod: Optional[Callable] = None + score_mod_bprop: Optional[Callable] = None + score_mod_requested: bool = False @nn.compact def __call__( @@ -317,6 +320,8 @@ def __call__( *, dropout_rng: Optional[PRNGKey] = None, deterministic: bool = False, + score_mod_tensors: Optional[Mapping[str, Any]] = None, + score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, ) -> Array: seed = None @@ -340,6 +345,26 @@ def __call__( jnp.float32, ) + fused_attn_kwargs = { + "attn_mask_type": self.attn_mask_type, + "attn_bias_type": self.attn_bias_type, + "softmax_type": self.softmax_type, + "scaling_factor": scale_factor, + "dropout_probability": self.attention_dropout, + "is_training": not deterministic, + "window_size": self.window_size, + "max_segments_per_seq": self.max_segments_per_seq, + "context_parallel_causal_load_balanced": self.context_parallel_causal_load_balanced, + "context_parallel_axis": self.context_parallel_axis, + "context_parallel_strategy": self.context_parallel_strategy, + "context_checkpoint_name": self.context_checkpoint_name, + "softmax_offset": softmax_offset, + "score_mod": self.score_mod, + "score_mod_bprop": self.score_mod_bprop, + "score_mod_tensors": score_mod_tensors, + "score_mod_bprop_tensors": score_mod_bprop_tensors, + } + if self.qkv_layout.is_qkvpacked(): """qkvpacked format, treat query: qkvpacked tensor, shape = [..., 3, h, d] @@ -349,25 +374,23 @@ def __call__( qkv_packed = query if self.transpose_batch_sequence: qkv_packed = qkv_packed.transpose([1, 0, 2, 3, 4]) + if self.score_mod_requested: + query, key, value = jnp.split(qkv_packed, [1, 2], axis=-3) + query, key, value = map( + functools.partial(jnp.squeeze, axis=-3), [query, key, value] + ) + qkv_layout = self.qkv_layout.to_separate() + qkv_args = (query, key, value) + else: + qkv_layout = self.qkv_layout + qkv_args = (qkv_packed,) x = fused_attn( - (qkv_packed,), + qkv_args, bias, sequence_descriptor, seed, - attn_mask_type=self.attn_mask_type, - attn_bias_type=self.attn_bias_type, - qkv_layout=self.qkv_layout, - softmax_type=self.softmax_type, - scaling_factor=scale_factor, - dropout_probability=self.attention_dropout, - is_training=not deterministic, - window_size=self.window_size, - max_segments_per_seq=self.max_segments_per_seq, - context_parallel_causal_load_balanced=self.context_parallel_causal_load_balanced, - context_parallel_axis=self.context_parallel_axis, - context_parallel_strategy=self.context_parallel_strategy, - context_checkpoint_name=self.context_checkpoint_name, - softmax_offset=softmax_offset, + qkv_layout=qkv_layout, + **fused_attn_kwargs, ) elif self.qkv_layout.is_kvpacked(): """kvpacked format, treat @@ -379,25 +402,21 @@ def __call__( if self.transpose_batch_sequence: query = query.transpose([1, 0, 2, 3]) kv_packed = kv_packed.transpose([1, 0, 2, 3, 4]) + if self.score_mod_requested: + key, value = jnp.split(kv_packed, [1], axis=-3) + key, value = map(functools.partial(jnp.squeeze, axis=-3), [key, value]) + qkv_layout = self.qkv_layout.to_separate() + qkv_args = (query, key, value) + else: + qkv_layout = self.qkv_layout + qkv_args = (query, kv_packed) x = fused_attn( - (query, kv_packed), + qkv_args, bias, sequence_descriptor, seed, - attn_mask_type=self.attn_mask_type, - attn_bias_type=self.attn_bias_type, - qkv_layout=self.qkv_layout, - softmax_type=self.softmax_type, - scaling_factor=scale_factor, - dropout_probability=self.attention_dropout, - is_training=not deterministic, - window_size=self.window_size, - max_segments_per_seq=self.max_segments_per_seq, - context_parallel_causal_load_balanced=self.context_parallel_causal_load_balanced, - context_parallel_axis=self.context_parallel_axis, - context_parallel_strategy=self.context_parallel_strategy, - context_checkpoint_name=self.context_checkpoint_name, - softmax_offset=softmax_offset, + qkv_layout=qkv_layout, + **fused_attn_kwargs, ) elif self.qkv_layout.is_separate(): if self.transpose_batch_sequence: @@ -409,20 +428,8 @@ def __call__( bias, sequence_descriptor, seed, - attn_mask_type=self.attn_mask_type, - attn_bias_type=self.attn_bias_type, qkv_layout=self.qkv_layout, - softmax_type=self.softmax_type, - scaling_factor=scale_factor, - dropout_probability=self.attention_dropout, - is_training=not deterministic, - window_size=self.window_size, - max_segments_per_seq=self.max_segments_per_seq, - context_parallel_causal_load_balanced=self.context_parallel_causal_load_balanced, - context_parallel_axis=self.context_parallel_axis, - context_parallel_strategy=self.context_parallel_strategy, - context_checkpoint_name=self.context_checkpoint_name, - softmax_offset=softmax_offset, + **fused_attn_kwargs, ) else: raise ValueError(f"Unsupported {self.qkv_layout=}.") @@ -602,6 +609,16 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention (``'zero sink'`` and ``'learnable sink'``). + score_mod: Optional[Callable], default = None + Experimental cuDNN frontend score modification callback. When set, fused attention is + required and the same score_mod limitations as :func:`fused_attn` apply. + score_mod_bprop: Optional[Callable], default = None + Optional score modification backward callback. + score_mod_tensors: Optional[Mapping[str, Any]], default = None + Additional tensors or pass-by-value scalars for ``score_mod``. This is a call-time + argument to keep tensor operands as normal JAX inputs. + score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None + Additional tensors or pass-by-value scalars for ``score_mod_bprop``. Optimization parameters ----------------------- @@ -628,6 +645,8 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods context_parallel_strategy: str = "DEFAULT" context_checkpoint_name: str = "context" softmax_type: str = "vanilla" + score_mod: Optional[Callable] = None + score_mod_bprop: Optional[Callable] = None def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -669,6 +688,8 @@ def __call__( *, deterministic: bool = False, mask: Optional[Union[SequenceDescriptor, Array]] = None, + score_mod_tensors: Optional[Mapping[str, Any]] = None, + score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, ) -> Array: """ Parameters @@ -689,6 +710,10 @@ def __call__( Below parameters are keyword only deterministic: bool, default = False Disable dropout layers if set to True. + score_mod_tensors: Optional[Mapping[str, Any]], default = None + Additional tensor/scalar operands passed to ``score_mod``. + score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None + Additional tensor/scalar operands passed to ``score_mod_bprop``. Returns ------- @@ -715,6 +740,12 @@ def __call__( qkv_layout = QKVLayout[self.qkv_layout.upper()] softmax_type = AttnSoftmaxType.from_str(self.softmax_type) del self.attn_bias_type, self.attn_mask_type, self.qkv_layout + score_mod_requested = ( + self.score_mod is not None + or self.score_mod_bprop is not None + or score_mod_tensors is not None + or score_mod_bprop_tensors is not None + ) if attn_bias_type == AttnBiasType.NO_BIAS: assert ( @@ -760,13 +791,19 @@ def __call__( head_dim_qk = self.head_dim head_dim_v = self.head_dim + if score_mod_requested: + if not enable_fused_attn: + raise ValueError("score_mod requires fused attention, but NVTE_FUSED_ATTN=0.") + kernel_qkv_layout = qkv_layout.to_separate() if score_mod_requested else qkv_layout has_fused_attn_kernel = is_fused_attn_kernel_available( - # This needs to be fixed: TE-Jax has historically correlated training mode with deterministic mode. + # This needs to be fixed: TE-Jax has historically correlated training mode + # with deterministic mode. not deterministic, input_dtype, - # self._assert_dtypes enforces Q, K, V, bias to have the same dtype so using input_dtype as kv dtype is sufficient + # self._assert_dtypes enforces Q, K, V, bias to have the same dtype, so + # using input_dtype as kv dtype is sufficient. input_dtype, - qkv_layout, + kernel_qkv_layout, attn_bias_type, attn_mask_type, softmax_type, @@ -779,6 +816,10 @@ def __call__( head_dim_v, self.window_size, ) + if score_mod_requested and not has_fused_attn_kernel: + raise ValueError( + "score_mod requires fused attention, but no fused attention kernel is available." + ) use_fused_attn = enable_fused_attn and has_fused_attn_kernel @@ -872,6 +913,9 @@ def __call__( context_parallel_strategy=context_parallel_strategy, context_checkpoint_name=self.context_checkpoint_name, softmax_type=softmax_type, + score_mod=self.score_mod, + score_mod_bprop=self.score_mod_bprop, + score_mod_requested=score_mod_requested, )( query, key, @@ -880,6 +924,8 @@ def __call__( bias, dropout_rng=dropout_rng, deterministic=deterministic, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, ) assert x.dtype == input_dtype, f"output_dtype={x.dtype}, input_dtype={input_dtype}" return x @@ -1178,6 +1224,16 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention (``'zero sink'`` and ``'learnable sink'``). + score_mod: Optional[Callable], default = None + Experimental cuDNN frontend score modification callback. When set, fused attention is + required and the same score_mod limitations as :func:`fused_attn` apply. + score_mod_bprop: Optional[Callable], default = None + Optional score modification backward callback. + score_mod_tensors: Optional[Mapping[str, Any]], default = None + Additional tensors or pass-by-value scalars for ``score_mod``. This is a call-time + argument to keep tensor operands as normal JAX inputs. + score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None + Additional tensors or pass-by-value scalars for ``score_mod_bprop``. """ head_dim: int @@ -1210,6 +1266,8 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods float32_logits: bool = False window_size: Optional[Tuple[int, int]] = None softmax_type: str = "vanilla" + score_mod: Optional[Callable] = None + score_mod_bprop: Optional[Callable] = None # Deprecated parameters num_heads: Optional[int] = None @@ -1278,6 +1336,8 @@ def __call__( *, decode: bool = False, deterministic: bool = False, + score_mod_tensors: Optional[Mapping[str, Any]] = None, + score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, ) -> Array: """ MultiHeadAttention Layer: @@ -1300,6 +1360,10 @@ def __call__( Indicate whether to prepare and use an autoregressive cache. deterministic: bool, default = False Disable dropout layers if set to True. + score_mod_tensors: Optional[Mapping[str, Any]], default = None + Additional tensor/scalar operands passed to ``score_mod``. + score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None + Additional tensor/scalar operands passed to ``score_mod_bprop``. Returns ------- @@ -1631,7 +1695,16 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): transpose_batch_sequence=self.transpose_batch_sequence, window_size=self.window_size, softmax_type=self.softmax_type, - )(*dpa_args, mask, bias, deterministic=deterministic) + score_mod=self.score_mod, + score_mod_bprop=self.score_mod_bprop, + )( + *dpa_args, + mask, + bias, + deterministic=deterministic, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, + ) x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3])) attn_context_sharding_constraint = (*LEADING_AXES, HIDDEN_TP_AXES) @@ -1939,6 +2012,17 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention (``'zero sink'`` and ``'learnable sink'``). Only supported for fused attention backend. + score_mod: Optional[Callable], default = None + Experimental cuDNN frontend score modification callback for the self-attention block. + When set, fused attention is required and the same score_mod limitations as + :func:`fused_attn` apply. + score_mod_bprop: Optional[Callable], default = None + Optional score modification backward callback for the self-attention block. + score_mod_tensors: Optional[Mapping[str, Any]], default = None + Additional tensors or pass-by-value scalars for ``score_mod``. This is a call-time + argument to keep tensor operands as normal JAX inputs. + score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None + Additional tensors or pass-by-value scalars for ``score_mod_bprop``. Optimization parameters ----------------------- @@ -2005,6 +2089,8 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods scaled_query_init: bool = True window_size: Optional[Tuple[int, int]] = None softmax_type: str = "vanilla" + score_mod: Optional[Callable] = None + score_mod_bprop: Optional[Callable] = None def __post_init__(self): if self.mha_kernel_init is None: @@ -2029,6 +2115,8 @@ def __call__( deterministic: bool = False, decode: bool = False, max_decode_length: bool = None, + score_mod_tensors: Optional[Mapping[str, Any]] = None, + score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, ): """ Transformer Layer: attention block and a feedforward network (MLP) @@ -2057,6 +2145,10 @@ def __call__( The maximum length to generate relative embedding biases when :attr:`layer_type=TransformerLayerType.DECODER` and :attr:`enable_relative_embedding=True`. + score_mod_tensors: Optional[Mapping[str, Any]], default = None + Additional tensor/scalar operands passed to the self-attention ``score_mod``. + score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None + Additional tensor/scalar operands passed to the self-attention ``score_mod_bprop``. Returns ------- @@ -2168,7 +2260,18 @@ def generate_batch_seqlen_logical_axes(is_shared_seq=None): name=mha_name, window_size=self.window_size, softmax_type=self.softmax_type, - )(inputs, inputs, attention_mask, attn_bias, deterministic=deterministic, decode=decode) + score_mod=self.score_mod, + score_mod_bprop=self.score_mod_bprop, + )( + inputs, + inputs, + attention_mask, + attn_bias, + deterministic=deterministic, + decode=decode, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, + ) def hidden_dropout(x, deterministic): assert isinstance( diff --git a/transformer_engine/jax/router.py b/transformer_engine/jax/router.py index 65f2e8a7ffb..80ec42a95f9 100644 --- a/transformer_engine/jax/router.py +++ b/transformer_engine/jax/router.py @@ -27,6 +27,7 @@ from transformer_engine.jax.cpp_extensions.router import ( ScoreFunction, + RoutingMapFormat, fused_topk_with_score_function_fwd, fused_topk_with_score_function_bwd, fused_moe_aux_loss_fwd, @@ -35,11 +36,35 @@ __all__ = [ "ScoreFunction", + "RoutingMapFormat", "fused_topk_with_score_function", "fused_moe_aux_loss", ] +def _validate_routing_map_format( + routing_map_format: Union[str, RoutingMapFormat, int], +) -> RoutingMapFormat: + """Validate and convert routing_map_format to a RoutingMapFormat enum.""" + if isinstance(routing_map_format, RoutingMapFormat): + return routing_map_format + if isinstance(routing_map_format, int): + try: + return RoutingMapFormat(routing_map_format) + except ValueError: + raise ValueError( + "routing_map_format int must match a RoutingMapFormat value; " + f"got {routing_map_format!r}" + ) from None + try: + return RoutingMapFormat[routing_map_format.upper()] + except (KeyError, AttributeError): + raise ValueError( + "routing_map_format must be 'bytemap', 'bitmap_u8', a RoutingMapFormat enum, " + f"or the matching int; got {routing_map_format!r}" + ) from None + + def _validate_score_function(score_function: Union[str, ScoreFunction]) -> ScoreFunction: """Validate and convert score_function to a ScoreFunction enum.""" if isinstance(score_function, ScoreFunction): @@ -68,6 +93,7 @@ def fused_topk_with_score_function( score_function: Union[str, ScoreFunction] = ScoreFunction.SOFTMAX, expert_bias: Optional[jnp.ndarray] = None, compute_aux_scores: bool = False, + routing_map_format: Union[str, RoutingMapFormat, int] = RoutingMapFormat.BYTEMAP, ) -> Tuple[jnp.ndarray, jnp.ndarray]: """ Fused top-k with score function router. @@ -108,6 +134,11 @@ def fused_topk_with_score_function( compute_aux_scores : bool If True, use the clean score-for-aux-loss kernel. Returns dense scores over all experts instead of sparse probs. + routing_map_format : Union[str, RoutingMapFormat, int] + Output layout for routing_map. "bytemap" / RoutingMapFormat.BYTEMAP (default) + returns a bool[T, E] tensor; "bitmap_u8" / RoutingMapFormat.BITMAP_U8 returns a + uint8[T, ceil(E/8)] tensor with bit (e % 8) of byte (e / 8) set when token t + routes to expert e (LSB-first along the expert axis). Returns ------- @@ -117,8 +148,10 @@ def fused_topk_with_score_function( When compute_aux_scores=True: Dense score tensor, shape [num_tokens, num_experts]. All expert positions contain scores. routing_map : jnp.ndarray - Boolean mask, shape [num_tokens, num_experts]. - True at selected expert positions. + Shape/dtype depend on routing_map_format: + - BYTEMAP: bool[num_tokens, num_experts]; True at selected expert positions. + - BITMAP_U8: uint8[num_tokens, ceil(num_experts/8)]; LSB-first bit-packed + along the expert axis. """ if not isinstance(scaling_factor, (int, float)): raise TypeError( @@ -127,6 +160,7 @@ def fused_topk_with_score_function( ) score_function = _validate_score_function(score_function) + routing_map_format = _validate_routing_map_format(routing_map_format) if compute_aux_scores: expert_bias = jnp.empty((0,), dtype=logits.dtype) @@ -153,12 +187,13 @@ def fused_topk_with_score_function( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ) return probs_or_scores, routing_map -@partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5, 6, 7, 8)) +@partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5, 6, 7, 8, 9)) def _fused_topk_with_score_function( logits: jnp.ndarray, expert_bias: jnp.ndarray, @@ -169,6 +204,7 @@ def _fused_topk_with_score_function( scaling_factor: float, score_function: ScoreFunction, compute_aux_scores: bool, + routing_map_format: RoutingMapFormat, ) -> Tuple[jnp.ndarray, jnp.ndarray]: (probs, routing_map), _ = _fused_topk_with_score_function_fwd( logits, @@ -180,6 +216,7 @@ def _fused_topk_with_score_function( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ) return probs, routing_map @@ -194,6 +231,7 @@ def _fused_topk_with_score_function_fwd( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ): probs, routing_map, saved_scores = fused_topk_with_score_function_fwd( logits, @@ -205,6 +243,7 @@ def _fused_topk_with_score_function_fwd( score_function, expert_bias, compute_aux_scores, + routing_map_format, ) residuals = (routing_map, saved_scores) return (probs, routing_map), residuals @@ -218,6 +257,7 @@ def _fused_topk_with_score_function_bwd( scaling_factor, score_function, compute_aux_scores, + routing_map_format, residuals, g, ): @@ -234,6 +274,7 @@ def _fused_topk_with_score_function_bwd( scaling_factor, score_function, compute_aux_scores, + routing_map_format, ) return grad_logits, None diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 95ee370c81b..f94d3539404 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -485,6 +485,7 @@ def _normalize_grid(grid_tuple): # NVTE_JAX_ENFORCE_TRITON_AUTOTUNING=1 to raise an error instead, prompting the # user to upgrade JAX for improved performance. is_autotuned = isinstance(kernel_fn, autotuner.Autotuner) + used_autotuned_launch = False if is_autotuned and not is_triton_autotuned_alias_safe(): val = os.environ.get("NVTE_JAX_ENFORCE_TRITON_AUTOTUNING", "0") try: @@ -574,6 +575,7 @@ def _normalize_grid(grid_tuple): kernel_calls, input_output_aliases_with_sizes, ) + used_autotuned_launch = True else: # Regular kernel: compile single config. @@ -630,7 +632,9 @@ def _normalize_grid(grid_tuple): ffi_operand_output_aliases = None compressed_call_proto = zlib.compress(call_proto) - if jax_version_meet_requirement(TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION): + if not used_autotuned_launch and jax_version_meet_requirement( + TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION + ): rule = jax.ffi.ffi_lowering( "triton_kernel_call_ffi", backend_config={"opaque": ir.StringAttr.get(compressed_call_proto)}, diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 7653d5992e0..06db28ee270 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -16,6 +16,8 @@ assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}." load_framework_extension("torch") +from transformer_engine.pytorch import constants +from transformer_engine.pytorch.constants import DType from transformer_engine.pytorch.module import LayerNormLinear from transformer_engine.pytorch.module import Linear from transformer_engine.pytorch.module import LayerNormMLP @@ -54,6 +56,8 @@ from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.utils import get_device_compute_capability from transformer_engine.pytorch.utils import is_bf16_available +from transformer_engine.pytorch.utils import deinterleave_glu_tensor +from transformer_engine.pytorch.utils import interleave_glu_tensor from transformer_engine.pytorch.graph import make_graphed_callables from transformer_engine.pytorch.distributed import checkpoint from transformer_engine.pytorch.distributed import CudaRNGStatesTracker @@ -132,7 +136,10 @@ MXFP8Quantizer, NVFP4Quantizer, Float8BlockQuantizer, - # pybind11 enum used as Quantizer.dtype + # Python IntEnum used as Quantizer.dtype. + DType, + # pybind11 enum used as Quantizer.dtype. + # Kept for backward compatibility. tex.DType, # __reduce_ex__ reconstructors (module-level functions). _make_float8_tensor_in_reduce_ex, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 9f2dadb680f..8f42983553c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -53,6 +53,9 @@ from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( attn_forward_func_with_cp, ) +from transformer_engine.pytorch.attention.dot_product_attention.flex_attention import ( + FusedAttentionWithScoreModFunc, +) from transformer_engine.pytorch.attention.dot_product_attention.softmax import FusedScaleMaskSoftmax from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.cpu_offload import ( @@ -89,7 +92,6 @@ _flash_attn_bwd = None _flash_attn_varlen_fwd = None _flash_attn_varlen_bwd = None - # Try to import Flash Attention v2 try: fa_utils.version = PkgVersion(PkgVersion(get_pkg_version("flash-attn")).public) @@ -1985,6 +1987,10 @@ def forward( inference_params: Optional[InferenceParams] = None, softmax_offset: torch.Tensor = None, fp8_output: bool = False, + score_mod: Optional[Callable] = None, + score_mod_bprop: Optional[Callable] = None, + score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2153,6 +2159,21 @@ def forward( layer_number=self.layer_number, return_max_logit=self.return_max_logit, ) + elif score_mod is not None: + output = FusedAttentionWithScoreModFunc.apply( + self.training, + query_layer, + key_layer, + value_layer, + q_format, + kv_format, + self.softmax_scale, + score_mod, + score_mod_bprop, + score_mod_tensors, + score_mod_bprop_tensors, + self.deterministic, + ) else: with self.attention_dropout_ctx(): output = FusedAttnFunc.apply( 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 ca848a94801..a244c8f8143 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 @@ -13,7 +13,6 @@ import torch from torch.nn.parameter import Parameter -import transformer_engine_torch as tex from transformer_engine.common.recipe import ( Format, Recipe, @@ -35,11 +34,7 @@ from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.export import is_in_onnx_export_mode -from transformer_engine.pytorch.constants import ( - AttnMaskTypes, - AttnTypes, - dist_group_type, -) +from transformer_engine.pytorch.constants import AttnMaskTypes, AttnTypes, dist_group_type, DType from transformer_engine.pytorch.distributed import ( get_distributed_world_size, checkpoint, @@ -1014,6 +1009,10 @@ def forward( pad_between_seqs: Optional[bool] = None, fp8_output: Optional[bool] = False, num_splits: Optional[int] = 1, + score_mod: Optional[Callable] = None, + score_mod_bprop: Optional[Callable] = None, + score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1103,6 +1102,12 @@ def forward( :attr:`max_seqlen_q`, :attr:`max_seqlen_kv`}, and for cuDNN >= 9.6, {"t" dimension of :attr:`query_layer`, "t" dimension of :attr:`key_layer`}. + .. note:: + + :attr:`score_mod`, :attr:`score_mod_bprop`, :attr:`score_mod_tensors`, and + :attr:`score_mod_bprop_tensors` are experimental cuDNN frontend Flex Attention + APIs. Their callback signatures and supported configurations may change. + Parameters ---------- query_layer : torch.Tensor @@ -1202,6 +1207,22 @@ def forward( Optional split control for FlashAttention-3 only. When set, this value is forwarded to the FA3 backend to control internal kernel splitting behavior for non-context-parallel cases. It is ignored for other backends and when context parallelism is enabled. + score_mod: Optional[Callable], default = None + Experimental cuDNN frontend score modification callback. This is a cuDNN-only path + and is mutually exclusive with masks, bias, ALiBi, sink attention, dropout, FP8, + context parallelism, THD format, KV caching, and return_max_logit. The callback + signature is ``score_mod(graph, score, tensors) -> score``. + score_mod_bprop: Optional[Callable], default = None + Optional cuDNN frontend callback for the backward pass of score_mod. The callback + signature is ``score_mod_bprop(graph, dP, tensors) -> dP``. + score_mod_tensors: Optional[Dict[str, torch.Tensor]], default = None + Runtime tensors exposed to score_mod as cuDNN graph tensors. Keys are + user-defined string names consumed by the callback through ``tensors[name]``; + there is no predefined set of accepted keys. + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], default = None + Runtime tensors exposed to score_mod_bprop as cuDNN graph tensors. Keys are + user-defined string names consumed by the callback through ``tensors[name]``; + there is no predefined set of accepted keys. """ with self.prepare_forward_ctx( @@ -1232,11 +1253,11 @@ def forward( forward_dtype = get_fp8_te_dtype(self.fp8_meta["recipe"], fprop_tensor=True) backward_dtype = get_fp8_te_dtype(self.fp8_meta["recipe"], fprop_tensor=False) assert forward_dtype in [ - tex.DType.kFloat8E4M3, - tex.DType.kFloat8E5M2, + DType.kFloat8E4M3, + DType.kFloat8E5M2, ] and backward_dtype in [ - tex.DType.kFloat8E4M3, - tex.DType.kFloat8E5M2, + DType.kFloat8E4M3, + DType.kFloat8E5M2, ], """DotProductAttention only supports "E4M3" and "E5M2" FP8 data types.""" else: fp8_output = False @@ -1528,6 +1549,40 @@ def forward( else: pad_between_seqs = False + # Validate experimental Flex Attention API inputs that backend selection + # cannot represent. + if score_mod is None: + assert score_mod_bprop is None, "score_mod_bprop requires score_mod!" + assert score_mod_tensors is None, "score_mod_tensors requires score_mod!" + assert ( + score_mod_bprop_tensors is None + ), "score_mod_bprop_tensors requires score_mod!" + else: + assert callable(score_mod), "score_mod must be callable!" + assert score_mod_bprop is None or callable( + score_mod_bprop + ), "score_mod_bprop must be callable when provided!" + assert ( + not is_in_onnx_export_mode() + ), "Flex Attention is not supported with ONNX export!" + if score_mod_tensors is not None: + assert isinstance(score_mod_tensors, dict), "score_mod_tensors must be a dict!" + assert all( + isinstance(k, str) and isinstance(v, torch.Tensor) + for k, v in score_mod_tensors.items() + ), "score_mod_tensors must map string names to torch.Tensor instances!" + if score_mod_bprop_tensors is not None: + assert ( + score_mod_bprop is not None + ), "score_mod_bprop_tensors requires score_mod_bprop!" + assert isinstance( + score_mod_bprop_tensors, dict + ), "score_mod_bprop_tensors must be a dict!" + assert all( + isinstance(k, str) and isinstance(v, torch.Tensor) + for k, v in score_mod_bprop_tensors.items() + ), "score_mod_bprop_tensors must map string names to torch.Tensor instances!" + # gather attention params for get_attention_backend attention_params = dpa_utils.AttentionParams( qkv_type=type(query_layer), @@ -1563,6 +1618,10 @@ def forward( return_max_logit=self.return_max_logit, cuda_graph=is_graph_capturing(), num_splits=num_splits, + fp8_output=fp8_output, + checkpoint_core_attention=checkpoint_core_attention, + has_score_mod=score_mod is not None, + has_score_mod_bprop=score_mod_bprop is not None, ) global _attention_backends if is_in_onnx_export_mode(): @@ -1744,6 +1803,10 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, + score_mod=score_mod, + score_mod_bprop=score_mod_bprop, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py new file mode 100644 index 00000000000..da8205923f5 --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -0,0 +1,773 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN-backed Flex Attention helpers.""" + +from dataclasses import dataclass +import importlib +import importlib.util +import inspect +from pathlib import Path +import sys +from typing import Any, Callable, Dict, Optional, Tuple + +import torch + +_cudnn_score_mod_handles: Dict[torch.device, Any] = {} +_cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} +_SCORE_MOD_UNCACHEABLE = object() +_CUDNN_FRONTEND_PYTHON_PATH = ( + Path(__file__).resolve().parents[4] / "3rdparty" / "cudnn-frontend" / "python" +) + + +def _import_cudnn_frontend(): + """Import the vendored cuDNN frontend if built, otherwise use the installed package.""" + cudnn_frontend_path = str(_CUDNN_FRONTEND_PYTHON_PATH) + cudnn_frontend_package = _CUDNN_FRONTEND_PYTHON_PATH / "cudnn" + if any(cudnn_frontend_package.glob("_compiled_module*")): + if cudnn_frontend_path not in sys.path: + sys.path.insert(0, cudnn_frontend_path) + return importlib.import_module("cudnn") + + if importlib.util.find_spec("cudnn") is not None: + return importlib.import_module("cudnn") + + raise ImportError( + "cuDNN frontend Python package not found. " + "Install it with: pip install nvidia-cudnn-frontend" + ) + + +def _bhsd_dim_stride( + tensor: torch.Tensor, tensor_format: str +) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: + """Describe an SBHD/BSHD tensor as cuDNN frontend's logical BHSD format.""" + if tensor_format == "sbhd": + return ( + (tensor.shape[1], tensor.shape[2], tensor.shape[0], tensor.shape[3]), + (tensor.stride(1), tensor.stride(2), tensor.stride(0), tensor.stride(3)), + ) + if tensor_format == "bshd": + return ( + (tensor.shape[0], tensor.shape[2], tensor.shape[1], tensor.shape[3]), + (tensor.stride(0), tensor.stride(2), tensor.stride(1), tensor.stride(3)), + ) + raise ValueError(f"Flex Attention only supports SBHD/BSHD tensor formats, got {tensor_format}.") + + +def _bhsd_graph_tensor(graph, tensor: torch.Tensor, tensor_format: str): + """Create a cuDNN graph tensor with BHSD dims and TE-layout strides.""" + dim, stride = _bhsd_dim_stride(tensor, tensor_format) + return graph.tensor(dim=dim, stride=stride, data_type=tensor.dtype) + + +# score_mod graph cache helpers. +def _freeze_score_mod_cache_key(value: Any) -> Any: + """Convert a user-provided score_mod graph key into a hashable structure.""" + if isinstance(value, torch.Tensor): + raise TypeError( + "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " + "through score_mod_tensors or score_mod_bprop_tensors instead." + ) + if isinstance(value, dict): + items = ( + ( + _freeze_score_mod_cache_key(key), + _freeze_score_mod_cache_key(val), + ) + for key, val in value.items() + ) + return tuple(sorted(items, key=repr)) + if isinstance(value, (list, tuple)): + return tuple(_freeze_score_mod_cache_key(item) for item in value) + if isinstance(value, (set, frozenset)): + items = (_freeze_score_mod_cache_key(item) for item in value) + return tuple(sorted(items, key=repr)) + try: + hash(value) + except TypeError as exc: + raise TypeError( + "score_mod_graph_cache_key() must return a hashable value or a nested " + "combination of dict/list/tuple/set values." + ) from exc + return value + + +def _score_mod_explicit_cache_key(callback_owner: Any) -> Optional[Any]: + """Return a user-provided structural graph key for a score_mod callback.""" + explicit_key = getattr(callback_owner, "score_mod_graph_cache_key", None) + if explicit_key is None: + return None + explicit_key = explicit_key() if callable(explicit_key) else explicit_key + return _freeze_score_mod_cache_key(explicit_key) + + +def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: + """Create a stable graph cache key for a score_mod callable. + + Module-level named functions are assumed to have stable topology. Anonymous functions + are keyed by code object because lambdas in the same module can share the same + qualname. Stateful bound methods and callable instances need an explicit + score_mod_graph_cache_key(); otherwise their graphs are left uncached to avoid reusing + stale graphs after Python object address reuse. + """ + if callback is None: + return None + self_obj = getattr(callback, "__self__", None) + func_obj = getattr(callback, "__func__", None) + if self_obj is not None and func_obj is not None: + explicit_key = _score_mod_explicit_cache_key(self_obj) + if explicit_key is None: + return _SCORE_MOD_UNCACHEABLE + return ( + "bound_method", + type(self_obj), + func_obj.__module__, + func_obj.__qualname__, + explicit_key, + ) + + explicit_key = _score_mod_explicit_cache_key(callback) + if explicit_key is not None: + return ( + "callable", + type(callback), + getattr(callback, "__module__", None), + getattr(callback, "__qualname__", None), + explicit_key, + ) + + if ( + inspect.isfunction(callback) + and callback.__closure__ is None + and "" not in callback.__qualname__ + ): + if callback.__name__ == "" or not callback.__qualname__: + return ("function", callback.__module__, callback.__code__) + return ("function", callback.__module__, callback.__qualname__) + + return _SCORE_MOD_UNCACHEABLE + + +def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]: + """Normalize a tensor device for graph cache keys.""" + if device.type == "cuda": + index = device.index + if index is None: + index = torch.cuda.current_device() + return (device.type, index) + return (device.type, device.index) + + +def _score_mod_tensor_metadata(tensor: torch.Tensor) -> Tuple[Any, ...]: + """Describe tensor metadata that can affect cuDNN graph construction.""" + return ( + tuple(tensor.size()), + tuple(tensor.stride()), + tensor.dtype, + _score_mod_device_key(tensor.device), + ) + + +def _score_mod_tensor_dict_metadata( + tensors: Optional[Dict[str, torch.Tensor]], +) -> Tuple[Tuple[str, Tuple[Any, ...]], ...]: + """Describe score_mod tensor parameters without including their values.""" + if tensors is None: + return () + return tuple((name, _score_mod_tensor_metadata(tensor)) for name, tensor in tensors.items()) + + +def _score_mod_bhsd_tensor_metadata(tensor: torch.Tensor, tensor_format: str) -> Tuple[Any, ...]: + """Describe an SBHD/BSHD runtime tensor as a cuDNN BHSD graph tensor.""" + dim, stride = _bhsd_dim_stride(tensor, tensor_format) + return (dim, stride, tensor.dtype, _score_mod_device_key(tensor.device)) + + +def _make_cudnn_graph_tensor_dict(graph, tensors: Optional[Dict[str, torch.Tensor]]): + """Create cuDNN graph tensors matching runtime tensors.""" + if tensors is None: + return {} + return {name: graph.tensor_like(tensor) for name, tensor in tensors.items()} + + +# cuDNN frontend score_mod graph helpers. +def _wrap_score_mod(score_mod: Optional[Callable], graph_tensors: Dict[str, Any]): + """Adapt TE's score_mod signature to cuDNN frontend's two-argument callback.""" + if score_mod is None: + return None + + def _wrapped_score_mod(sdpa_graph, score_tensor): + return score_mod(sdpa_graph, score_tensor, graph_tensors) + + return _wrapped_score_mod + + +def _get_cudnn_current_stream_handle(cudnn, device: torch.device): + """Return a cuDNN handle for device, bound to PyTorch's current stream.""" + if device.type != "cuda": + raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + + handle = _cudnn_score_mod_handles.get(device) + with torch.cuda.device(device): + if handle is None: + handle = cudnn.create_handle() + _cudnn_score_mod_handles[device] = handle + + stream = torch.cuda.current_stream(device).cuda_stream + cudnn.set_stream(handle=handle, stream=stream) + return handle + + +def _build_cudnn_pygraph(dtype: torch.dtype, device: torch.device): + """Create a cuDNN frontend Python graph for F16/BF16 SDPA.""" + cudnn = _import_cudnn_frontend() + + if dtype == torch.float16: + io_data_type = cudnn.data_type.HALF + elif dtype == torch.bfloat16: + io_data_type = cudnn.data_type.BFLOAT16 + else: + raise ValueError(f"Flex Attention only supports FP16/BF16 tensors, got {dtype}.") + + graph = cudnn.pygraph( + io_data_type=io_data_type, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + handle=_get_cudnn_current_stream_handle(cudnn, device), + ) + return graph + + +@dataclass +class _CudnnScoreModFwdGraphEntry: + """Cached cuDNN frontend graph and graph tensor handles for score_mod fprop.""" + + graph: Any + q: Any + k: Any + v: Any + output: Any + stats: Optional[Any] + score_mod_graph_tensors: Dict[str, Any] + workspace_size: int + + +@dataclass +class _CudnnScoreModBwdGraphEntry: + """Cached cuDNN frontend graph and graph tensor handles for score_mod bprop.""" + + graph: Any + q: Any + k: Any + v: Any + output: Any + d_output: Any + stats: Any + dq: Any + dk: Any + dv: Any + score_mod_graph_tensors: Dict[str, Any] + score_mod_bprop_graph_tensors: Dict[str, Any] + workspace_size: int + + +def _finalize_cudnn_graph(graph) -> int: + """Build a cuDNN frontend Python graph and return its workspace size.""" + cudnn = _import_cudnn_frontend() + + graph.validate() + graph.build_operation_graph() + try: + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + except cudnn.cudnnGraphNotSupportedError as exc: + raise RuntimeError(f"cuDNN Flex Attention SDPA graph is not supported: {exc}") from exc + graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) + return max(graph.get_workspace_size(), 1) + + +def _execute_cudnn_graph( + graph, + variant_pack: Dict[Any, torch.Tensor], + workspace_size: int, + device: torch.device, +): + """Execute a built cuDNN frontend Python graph.""" + cudnn = _import_cudnn_frontend() + + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + workspace = torch.empty( + workspace_size, + device=device, + dtype=torch.uint8, + ) + graph.execute( + variant_pack, + workspace, + handle=_get_cudnn_current_stream_handle(cudnn, device), + ) + + +def _cudnn_score_mod_fwd_cache_key( + is_training: bool, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + q_format: str, + kv_format: str, + attn_scale: float, + score_mod: Callable, + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + output_layer: torch.Tensor, + stats: Optional[torch.Tensor], +) -> Optional[Tuple[Any, ...]]: + """Pre-build cache key for score_mod fprop execution plans. + + cuDNN exposes graph.key(), but only after graph construction has run the user callback. + This key avoids rebuilding the Python graph on cache hits. + """ + score_mod_key = _score_mod_callback_cache_key(score_mod) + if score_mod_key is _SCORE_MOD_UNCACHEABLE: + return None + return ( + "fwd", + is_training, + q_format, + kv_format, + attn_scale, + score_mod_key, + _score_mod_bhsd_tensor_metadata(query_layer, q_format), + _score_mod_bhsd_tensor_metadata(key_layer, kv_format), + _score_mod_bhsd_tensor_metadata(value_layer, kv_format), + _score_mod_bhsd_tensor_metadata(output_layer, q_format), + _score_mod_tensor_metadata(stats) if stats is not None else None, + _score_mod_tensor_dict_metadata(score_mod_tensors), + ) + + +def _cudnn_score_mod_bwd_cache_key( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + output_layer: torch.Tensor, + d_out: torch.Tensor, + stats: torch.Tensor, + q_format: str, + kv_format: str, + attn_scale: float, + score_mod: Callable, + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], + deterministic: bool, +) -> Optional[Tuple[Any, ...]]: + """Pre-build cache key for score_mod bprop execution plans.""" + score_mod_key = _score_mod_callback_cache_key(score_mod) + score_mod_bprop_key = _score_mod_callback_cache_key(score_mod_bprop) + if score_mod_key is _SCORE_MOD_UNCACHEABLE or score_mod_bprop_key is _SCORE_MOD_UNCACHEABLE: + return None + return ( + "bwd", + q_format, + kv_format, + attn_scale, + deterministic, + score_mod_key, + score_mod_bprop_key, + _score_mod_bhsd_tensor_metadata(query_layer, q_format), + _score_mod_bhsd_tensor_metadata(key_layer, kv_format), + _score_mod_bhsd_tensor_metadata(value_layer, kv_format), + _score_mod_bhsd_tensor_metadata(output_layer, q_format), + _score_mod_bhsd_tensor_metadata(d_out, q_format), + _score_mod_tensor_metadata(stats), + _score_mod_tensor_dict_metadata(score_mod_tensors), + _score_mod_tensor_dict_metadata(score_mod_bprop_tensors), + ) + + +def _build_cudnn_score_mod_fwd_graph( + is_training: bool, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + q_format: str, + kv_format: str, + attn_scale: float, + score_mod: Callable, + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + output_layer: torch.Tensor, + stats: Optional[torch.Tensor], +) -> _CudnnScoreModFwdGraphEntry: + """Build a cached cuDNN frontend graph for score_mod fprop.""" + cudnn = _import_cudnn_frontend() + + graph = _build_cudnn_pygraph(query_layer.dtype, query_layer.device) + q = _bhsd_graph_tensor(graph, query_layer, q_format) + k = _bhsd_graph_tensor(graph, key_layer, kv_format) + v = _bhsd_graph_tensor(graph, value_layer, kv_format) + + score_mod_graph_tensors = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) + wrapped_score_mod = _wrap_score_mod(score_mod, score_mod_graph_tensors) + + output_dim, output_stride = _bhsd_dim_stride(output_layer, q_format) + output, stats_tensor = graph.sdpa( + name="te_score_mod_sdpa", + q=q, + k=k, + v=v, + generate_stats=is_training, + attn_scale=attn_scale, + use_causal_mask=False, + score_mod=wrapped_score_mod, + ) + output.set_output(True).set_dim(output_dim).set_stride(output_stride) + + if is_training: + assert stats is not None + stats_tensor.set_output(True).set_dim(stats.size()).set_stride( + stats.stride() + ).set_data_type(cudnn.data_type.FLOAT) + else: + stats_tensor = None + + workspace_size = _finalize_cudnn_graph(graph) + return _CudnnScoreModFwdGraphEntry( + graph=graph, + q=q, + k=k, + v=v, + output=output, + stats=stats_tensor, + score_mod_graph_tensors=score_mod_graph_tensors, + workspace_size=workspace_size, + ) + + +def _get_cudnn_score_mod_fwd_graph( + is_training: bool, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + q_format: str, + kv_format: str, + attn_scale: float, + score_mod: Callable, + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + output_layer: torch.Tensor, + stats: Optional[torch.Tensor], +) -> _CudnnScoreModFwdGraphEntry: + """Return a cached cuDNN frontend graph for score_mod fprop.""" + build_args = ( + is_training, + query_layer, + key_layer, + value_layer, + q_format, + kv_format, + attn_scale, + score_mod, + score_mod_tensors, + output_layer, + stats, + ) + key = _cudnn_score_mod_fwd_cache_key(*build_args) + if key is None: + return _build_cudnn_score_mod_fwd_graph(*build_args) + entry = _cudnn_score_mod_graph_cache.get(key) + if entry is None: + entry = _build_cudnn_score_mod_fwd_graph(*build_args) + _cudnn_score_mod_graph_cache[key] = entry + return entry + + +def _build_cudnn_score_mod_bwd_graph( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + output_layer: torch.Tensor, + d_out: torch.Tensor, + stats: torch.Tensor, + q_format: str, + kv_format: str, + attn_scale: float, + score_mod: Callable, + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], + deterministic: bool, +) -> _CudnnScoreModBwdGraphEntry: + """Build a cached cuDNN frontend graph for score_mod bprop.""" + graph = _build_cudnn_pygraph(query_layer.dtype, query_layer.device) + q = _bhsd_graph_tensor(graph, query_layer, q_format) + k = _bhsd_graph_tensor(graph, key_layer, kv_format) + v = _bhsd_graph_tensor(graph, value_layer, kv_format) + output = _bhsd_graph_tensor(graph, output_layer, q_format) + d_output = _bhsd_graph_tensor(graph, d_out, q_format) + stats_tensor = graph.tensor_like(stats) + + score_mod_graph_tensors = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) + score_mod_bprop_graph_tensors = ( + _make_cudnn_graph_tensor_dict(graph, score_mod_bprop_tensors) + if score_mod_bprop is not None + else {} + ) + wrapped_score_mod = _wrap_score_mod(score_mod, score_mod_graph_tensors) + wrapped_score_mod_bprop = _wrap_score_mod(score_mod_bprop, score_mod_bprop_graph_tensors) + + dq_layer = torch.empty_like(query_layer) + dk_layer = torch.empty_like(key_layer) + dv_layer = torch.empty_like(value_layer) + dq_dim, dq_stride = _bhsd_dim_stride(dq_layer, q_format) + dk_dim, dk_stride = _bhsd_dim_stride(dk_layer, kv_format) + dv_dim, dv_stride = _bhsd_dim_stride(dv_layer, kv_format) + dq, dk, dv = graph.sdpa_backward( + name="te_score_mod_sdpa_backward", + q=q, + k=k, + v=v, + o=output, + dO=d_output, + stats=stats_tensor, + attn_scale=attn_scale, + use_causal_mask=False, + score_mod=wrapped_score_mod, + score_mod_bprop=wrapped_score_mod_bprop, + use_deterministic_algorithm=deterministic, + ) + dq.set_output(True).set_dim(dq_dim).set_stride(dq_stride) + dk.set_output(True).set_dim(dk_dim).set_stride(dk_stride) + dv.set_output(True).set_dim(dv_dim).set_stride(dv_stride) + + workspace_size = _finalize_cudnn_graph(graph) + return _CudnnScoreModBwdGraphEntry( + graph=graph, + q=q, + k=k, + v=v, + output=output, + d_output=d_output, + stats=stats_tensor, + dq=dq, + dk=dk, + dv=dv, + score_mod_graph_tensors=score_mod_graph_tensors, + score_mod_bprop_graph_tensors=score_mod_bprop_graph_tensors, + workspace_size=workspace_size, + ) + + +def _get_cudnn_score_mod_bwd_graph( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + output_layer: torch.Tensor, + d_out: torch.Tensor, + stats: torch.Tensor, + q_format: str, + kv_format: str, + attn_scale: float, + score_mod: Callable, + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], + deterministic: bool, +) -> _CudnnScoreModBwdGraphEntry: + """Return a cached cuDNN frontend graph for score_mod bprop.""" + build_args = ( + query_layer, + key_layer, + value_layer, + output_layer, + d_out, + stats, + q_format, + kv_format, + attn_scale, + score_mod, + score_mod_bprop, + score_mod_tensors, + score_mod_bprop_tensors, + deterministic, + ) + key = _cudnn_score_mod_bwd_cache_key(*build_args) + if key is None: + return _build_cudnn_score_mod_bwd_graph(*build_args) + entry = _cudnn_score_mod_graph_cache.get(key) + if entry is None: + entry = _build_cudnn_score_mod_bwd_graph(*build_args) + _cudnn_score_mod_graph_cache[key] = entry + return entry + + +class FusedAttentionWithScoreModFunc(torch.autograd.Function): + """cuDNN frontend Python SDPA path with Flex Attention score_mod support.""" + + @staticmethod + def forward( + ctx, + is_training: bool, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + q_format: str, + kv_format: str, + attn_scale: float, + score_mod: Callable, + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], + deterministic: bool, + ) -> torch.Tensor: + # pylint: disable=missing-function-docstring + q_bhsd_dim, _ = _bhsd_dim_stride(query_layer, q_format) + score_mod_tensors = dict(score_mod_tensors or {}) + score_mod_bprop_tensors = dict(score_mod_bprop_tensors or {}) + output_shape = (*query_layer.shape[:-1], value_layer.shape[-1]) + output_layer = torch.empty(output_shape, device=query_layer.device, dtype=query_layer.dtype) + if is_training: + stats = torch.empty( + (*q_bhsd_dim[:-1], 1), + device=query_layer.device, + dtype=torch.float32, + ) + else: + stats = None + + entry = _get_cudnn_score_mod_fwd_graph( + is_training, + query_layer, + key_layer, + value_layer, + q_format, + kv_format, + attn_scale, + score_mod, + score_mod_tensors, + output_layer, + stats, + ) + variant_pack = { + entry.q: query_layer, + entry.k: key_layer, + entry.v: value_layer, + entry.output: output_layer, + } + if is_training: + variant_pack[entry.stats] = stats + for name, graph_tensor in entry.score_mod_graph_tensors.items(): + variant_pack[graph_tensor] = score_mod_tensors[name] + + _execute_cudnn_graph( + entry.graph, + variant_pack, + entry.workspace_size, + query_layer.device, + ) + + ctx.is_training = is_training + ctx.q_format = q_format + ctx.kv_format = kv_format + ctx.attn_scale = attn_scale + ctx.score_mod = score_mod + ctx.score_mod_bprop = score_mod_bprop + ctx.score_mod_tensor_names = tuple(score_mod_tensors.keys()) + ctx.score_mod_bprop_tensor_names = tuple(score_mod_bprop_tensors.keys()) + ctx.deterministic = deterministic + if is_training: + # save_for_backward records version counters without copying tensor data. + # This catches in-place score_mod tensor updates before backward. + ctx.save_for_backward( + query_layer, + key_layer, + value_layer, + output_layer, + stats, + *score_mod_tensors.values(), + *score_mod_bprop_tensors.values(), + ) + else: + ctx.save_for_backward(query_layer, key_layer, value_layer, output_layer) + + return output_layer + + @staticmethod + def backward(ctx, d_out: torch.Tensor): + # pylint: disable=missing-function-docstring + if not ctx.is_training: + raise RuntimeError( + "score_mod backward requires DotProductAttention to be in training mode." + ) + + saved_tensors = ctx.saved_tensors + query_layer, key_layer, value_layer, output_layer, stats = saved_tensors[:5] + score_mod_tensors_end = 5 + len(ctx.score_mod_tensor_names) + score_mod_tensors = dict( + zip(ctx.score_mod_tensor_names, saved_tensors[5:score_mod_tensors_end]) + ) + score_mod_bprop_tensors = dict( + zip(ctx.score_mod_bprop_tensor_names, saved_tensors[score_mod_tensors_end:]) + ) + d_out = d_out.contiguous() + + dq_layer = torch.empty_like(query_layer) + dk_layer = torch.empty_like(key_layer) + dv_layer = torch.empty_like(value_layer) + entry = _get_cudnn_score_mod_bwd_graph( + query_layer, + key_layer, + value_layer, + output_layer, + d_out, + stats, + ctx.q_format, + ctx.kv_format, + ctx.attn_scale, + ctx.score_mod, + ctx.score_mod_bprop, + score_mod_tensors, + score_mod_bprop_tensors, + ctx.deterministic, + ) + variant_pack = { + entry.q: query_layer, + entry.k: key_layer, + entry.v: value_layer, + entry.output: output_layer, + entry.d_output: d_out, + entry.stats: stats, + entry.dq: dq_layer, + entry.dk: dk_layer, + entry.dv: dv_layer, + } + for name, graph_tensor in entry.score_mod_graph_tensors.items(): + variant_pack[graph_tensor] = score_mod_tensors[name] + for name, graph_tensor in entry.score_mod_bprop_graph_tensors.items(): + variant_pack[graph_tensor] = score_mod_bprop_tensors[name] + + _execute_cudnn_graph( + entry.graph, + variant_pack, + entry.workspace_size, + query_layer.device, + ) + + return ( + None, + dq_layer, + dk_layer, + dv_layer, + None, + None, + None, + None, + None, + None, + None, + None, + ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 989b65f190c..16498bcf02b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -259,6 +259,14 @@ class AttentionParams: Whether support for cuda graph capture is needed or not. num_splits : int, default = 1 The number of kernels to split attention to. + fp8_output : bool, default = False + Whether output is requested in FP8. + checkpoint_core_attention : bool, default = False + Whether core attention is recomputed during backward. + has_score_mod : bool, default = False + Whether a score_mod callback was provided. + has_score_mod_bprop : bool, default = False + Whether a score_mod bprop callback was provided. """ qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor @@ -292,6 +300,10 @@ class AttentionParams: return_max_logit: bool = False cuda_graph: bool = False num_splits: int = 1 + fp8_output: bool = False + checkpoint_core_attention: bool = False + has_score_mod: bool = False + has_score_mod_bprop: bool = False def __eq__(self, other): """ @@ -307,7 +319,7 @@ def __eq__(self, other): if fname != "fp8_meta": if sf != of: return False - elif sf.get("recipe", None) != of.get("recipe", None): + elif (sf or {}).get("recipe", None) != (of or {}).get("recipe", None): return False return True @@ -370,6 +382,10 @@ def get_attention_backend( return_max_logit = attention_params.return_max_logit cuda_graph = attention_params.cuda_graph num_splits = attention_params.num_splits + fp8_output = attention_params.fp8_output + checkpoint_core_attention = attention_params.checkpoint_core_attention + has_score_mod = attention_params.has_score_mod + has_score_mod_bprop = attention_params.has_score_mod_bprop # Run config logger = logging.getLogger("DotProductAttention") @@ -435,7 +451,7 @@ def get_attention_backend( # regardless of whether FA2 or FA3 is installed. If FA2 or FA3 is not installed but is # necessary for performance/functionality, a warning will be issued to prompt users to # install an appropriate FA version. - qkv_format, q_format, _ = get_qkv_format(qkv_layout, inference_params) + qkv_format, q_format, kv_format = get_qkv_format(qkv_layout, inference_params) # Filter: Environment variables use_flash_attention = int(os.getenv("NVTE_FLASH_ATTN", "1")) @@ -456,6 +472,18 @@ def get_attention_backend( if not use_unfused_attention: logger.debug("Disabling UnfusedDotProductAttention due to NVTE_UNFUSED_ATTN=0") + def _any_flash_attention_enabled() -> bool: + return bool( + use_flash_attention + or use_flash_attention_2 + or use_flash_attention_3 + or use_flash_attention_4 + ) + + def _disable_all_flash_attention() -> None: + nonlocal use_flash_attention + use_flash_attention = False + # Filter: Compute capability if device_compute_capability < (8, 0): if use_flash_attention_2 and FlashAttentionUtils.is_installed: @@ -650,6 +678,72 @@ def get_attention_backend( use_unfused_attention = False logger.debug("Disabling all backends for max_logit with FP8 attention") + # Filter: score_mod + if has_score_mod_bprop and not has_score_mod: + logger.debug("Disabling all backends because score_mod_bprop requires score_mod") + _disable_all_flash_attention() + use_fused_attention = False + use_unfused_attention = False + if has_score_mod: + if _any_flash_attention_enabled(): + logger.debug("Disabling FlashAttention for score_mod") + _disable_all_flash_attention() + if use_unfused_attention: + logger.debug("Disabling UnfusedDotProductAttention for score_mod") + use_unfused_attention = False + + score_mod_unsupported_reasons = [] + if qkv_dtype not in [torch.float16, torch.bfloat16]: + score_mod_unsupported_reasons.append( + f"unsupported qkv_dtype = {qkv_dtype}; supported: torch.float16, torch.bfloat16" + ) + if qkv_type is not torch.Tensor: + score_mod_unsupported_reasons.append( + f"unsupported qkv_type = {qkv_type}; supported: torch.Tensor" + ) + if fp8: + score_mod_unsupported_reasons.append("FP8 DotProductAttention is enabled") + if fp8_output: + score_mod_unsupported_reasons.append("fp8_output is enabled") + if inference_params is not None: + score_mod_unsupported_reasons.append("KV caching is enabled") + if context_parallel: + score_mod_unsupported_reasons.append("context parallelism is enabled") + if "thd" in (qkv_format, q_format, kv_format): + score_mod_unsupported_reasons.append( + f"unsupported QKV format: q_format = {q_format}, kv_format = {kv_format}" + ) + if pad_between_seqs: + score_mod_unsupported_reasons.append("pad_between_seqs is enabled") + if attn_mask_type != "no_mask": + score_mod_unsupported_reasons.append(f"attn_mask_type = {attn_mask_type}") + if window_size is not None and window_size != (-1, -1): + score_mod_unsupported_reasons.append(f"window_size = {window_size}") + if core_attention_bias_type != "no_bias": + score_mod_unsupported_reasons.append( + f"core_attention_bias_type = {core_attention_bias_type}" + ) + if alibi_slopes_shape is not None: + score_mod_unsupported_reasons.append("ALiBi slopes were provided") + if softmax_type != "vanilla": + score_mod_unsupported_reasons.append(f"softmax_type = {softmax_type}") + if attention_dropout != 0.0: + score_mod_unsupported_reasons.append(f"attention_dropout = {attention_dropout}") + if return_max_logit: + score_mod_unsupported_reasons.append("return_max_logit is enabled") + if checkpoint_core_attention: + score_mod_unsupported_reasons.append("checkpoint_core_attention is enabled") + if cuda_graph: + score_mod_unsupported_reasons.append("CUDA graph capture is enabled") + if num_splits != 1: + score_mod_unsupported_reasons.append(f"num_splits = {num_splits}") + if score_mod_unsupported_reasons and use_fused_attention: + logger.debug( + "Disabling FusedAttention for score_mod because %s", + "; ".join(score_mod_unsupported_reasons), + ) + use_fused_attention = False + # Filter: KV cache # backend | precision | KV cache | architecture | qkv_format | page_size # --------------------------------------------------------------------------------------- @@ -750,19 +844,12 @@ def get_attention_backend( if ( # pylint: disable=too-many-boolean-expressions use_flash_attention_2 and FlashAttentionUtils.is_installed - and ( - head_dim_qk > 256 - or head_dim_qk % 8 != 0 - or ( - head_dim_qk > 192 - and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) - ) - ) + and (head_dim_qk > 256 or head_dim_qk % 8 != 0) ): logger.debug( "Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. " "Supported: head_dim_qk = head_dim_v, head_dim_qk %%8 = 0, " - "head_dim_qk <= 256 (>192 requires sm80/90/100+). " + "head_dim_qk <= 256. " "Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.", head_dim_qk, head_dim_v, @@ -1260,6 +1347,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # Filter: cuDNN support fused_attention_backend = None if use_fused_attention: + # ``DType`` is implicitly convertible to ``transformer_engine::DType`` + # on the C++ side, so pass it straight to the pybind function. q_type = TE_DType[qkv_dtype] kv_type = q_type if fp8 and fp8_meta["recipe"].fp8_dpa: @@ -1290,6 +1379,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FusedAttention as no backend supports the provided input") use_fused_attention = False fused_attention_backend = None + elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]: + logger.debug( + "Disabling FusedAttention for score_mod because sub-backend %s is not " + "F16/BF16 arbitrary-seqlen", + int(fused_attention_backend), + ) + use_fused_attention = False + fused_attention_backend = None # Filter: Determinism # backend | deterministic # --------------------------------------------- diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index 2aff4fd8e8a..3a145bbb5b4 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -3,43 +3,107 @@ # See LICENSE for license information. """Enums for e2e transformer""" +import enum from types import SimpleNamespace +from typing import Union import torch import torch.distributed import transformer_engine_torch as tex -""" -This is a map: torch.dtype -> int -Used for passing dtypes into cuda -extension. Has one to one mapping -with enum in transformer_engine.h -""" +class DType(enum.IntEnum): + """Transformer Engine data types used to tag tensors passed to the + Transformer Engine backend. + This is the canonical dtype enum for ``transformer_engine.pytorch`` and + is used throughout the library (for example, to specify the precision of + quantized tensors and quantizers). Each member corresponds to a data type + supported by the Transformer Engine backend: + + * ``kByte`` -- 8-bit unsigned integer (``torch.uint8``). + * ``kInt32`` -- 32-bit signed integer (``torch.int32``). + * ``kFloat32`` -- 32-bit floating point (``torch.float32``). + * ``kFloat16`` -- 16-bit floating point (``torch.float16``). + * ``kBFloat16`` -- 16-bit brain floating point (``torch.bfloat16``). + * ``kFloat8E4M3`` -- 8-bit floating point with 4 exponent and 3 mantissa + bits (``torch.float8_e4m3fn``). + * ``kFloat8E5M2`` -- 8-bit floating point with 5 exponent and 2 mantissa + bits (``torch.float8_e5m2``). + * ``kFloat4E2M1`` -- 4-bit floating point with 2 exponent and 1 mantissa + bits. + + The enum mirrors the backend ``transformer_engine_torch.DType`` (pybind11) + enum value-for-value, and instances of the two enums compare equal when + they share the same integer value. + """ + + kByte = int(tex.DType.kByte) + kInt32 = int(tex.DType.kInt32) + kFloat32 = int(tex.DType.kFloat32) + kFloat16 = int(tex.DType.kFloat16) + kBFloat16 = int(tex.DType.kBFloat16) + kFloat8E4M3 = int(tex.DType.kFloat8E4M3) + kFloat8E5M2 = int(tex.DType.kFloat8E5M2) + kFloat4E2M1 = int(tex.DType.kFloat4E2M1) + + @classmethod + def cast(cls, dtype: "Union[DType, tex.DType]") -> "DType": + """Normalize a supported dtype value to the canonical ``DType`` ``IntEnum``. + ``DType`` is the canonical dtype tag used internally throughout + ``transformer_engine.pytorch``, and is what this function always outputs. + The pybind ``transformer_engine_torch.DType`` enum is an additional type + accepted as input (for backward compatibility), which this function maps + to the matching ``DType`` member so stored attributes are always ``DType``. + """ + if isinstance(dtype, cls): + return dtype + return cls(int(dtype)) + + def __eq__(self, other: object) -> bool: + # ``DType`` is an ``IntEnum`` while ``tex.DType`` is a pybind11 enum. + # ``int.__eq__`` returns ``NotImplemented`` for a pybind enum, so without + # this override a comparison such as ``quantizer.dtype == tex.DType.kX`` + # would silently be ``False``. Compare by integer value so the two enums + # stay equivalent (the pybind ``DType.__eq__`` handles the reverse order). + if isinstance(other, tex.DType): + return int(self) == int(other) + return int.__eq__(self, other) + + def __ne__(self, other: object) -> bool: + result = self.__eq__(other) + if result is NotImplemented: + return result + return not result + + def __hash__(self) -> int: + return int.__hash__(self) + + +# Fail fast at import time if a new enumerator is added +# on the C++ side without being mirrored above. +assert {m.name for m in DType} == set(tex.DType.__members__), ( + "DType in python is out of sync with transformer_engine_torch.DType; " + "defined in C++ side. Please make sure TE C++ and python are in sync." +) + + +# One-to-one mapping ``torch.dtype -> DType`` (mirrors the enum order in +# ``transformer_engine.h``). Use the bracket syntax ``TE_DType[torch_dtype]`` +# to resolve a ``torch.dtype`` to its matching ``DType`` member. +# Used for passing dtypes into cuda extension. TE_DType = { - torch.uint8: tex.DType.kByte, - torch.float8_e4m3fn: tex.DType.kFloat8E4M3, - torch.float8_e5m2: tex.DType.kFloat8E5M2, - torch.int32: tex.DType.kInt32, - torch.float32: tex.DType.kFloat32, - torch.half: tex.DType.kFloat16, - torch.bfloat16: tex.DType.kBFloat16, + torch.uint8: DType.kByte, + torch.float8_e4m3fn: DType.kFloat8E4M3, + torch.float8_e5m2: DType.kFloat8E5M2, + torch.int32: DType.kInt32, + torch.float32: DType.kFloat32, + torch.half: DType.kFloat16, + torch.bfloat16: DType.kBFloat16, } -""" -This is a map: int -> torch.dtype -Used for resolving cuda extension types to torch. -Has one to one mapping with enum in -transformer_engine.h -""" -TE_DType_To_Torch = { - tex.DType.kByte: torch.uint8, - tex.DType.kFloat8E4M3: torch.float8_e4m3fn, - tex.DType.kFloat8E5M2: torch.float8_e5m2, - tex.DType.kInt32: torch.int32, - tex.DType.kFloat32: torch.float32, - tex.DType.kFloat16: torch.half, - tex.DType.kBFloat16: torch.bfloat16, -} + +# Map ``DType -> torch.dtype`` for resolving cuda extension types to +# torch. +TE_DType_To_Torch = {value: key for key, value in TE_DType.items()} # Cache enum -> int conversions to avoid repeated PyObject lookups. FP8FwdTensorIdx = SimpleNamespace( diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 2ce939430da..62dcaadc961 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -16,7 +16,7 @@ NVTE_Fused_Attn_Backend, ) from ..quantized_tensor import Quantizer -from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx +from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, DType __all__ = [ @@ -26,12 +26,12 @@ TORCH_DType = { - tex.DType.kFloat8E4M3: torch.uint8, - tex.DType.kFloat8E5M2: torch.uint8, - tex.DType.kFloat16: torch.half, - tex.DType.kBFloat16: torch.bfloat16, - tex.DType.kFloat32: torch.float32, - tex.DType.kInt32: torch.int32, + DType.kFloat8E4M3: torch.uint8, + DType.kFloat8E5M2: torch.uint8, + DType.kFloat16: torch.half, + DType.kBFloat16: torch.bfloat16, + DType.kFloat32: torch.float32, + DType.kInt32: torch.int32, } QKVFormat = { @@ -173,7 +173,7 @@ def fused_attn_fwd( input tensor K; shape sbhd, bshd or thd (see `qkv_layout` for details) v : torch.Tensor input tensor V; shape sbhd, bshd or thd (see `qkv_layout` for details) - fake_dtype : tex.DType + fake_dtype : DType data type of Q, K and V - in case of high precision, fake dtype in case of FP8; in torch.dtype fused_attention_backend : tex.NVTE_Fused_Attn_Backend @@ -455,7 +455,7 @@ def fused_attn_bwd( d_o : torch.Tensor input tensor dO (gradient of O); same data type as Q, K and V; same shape as Q - fake_dtype : tex.DType + fake_dtype : DType data type of Q, K and V - in case of high precision, fake dtype in case of FP8; in torch.dtype aux_ctx_tensors : List[torch.Tensor] diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 48751d8c939..f3b066d50ba 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -9,7 +9,7 @@ import functools import torch import transformer_engine_torch as tex -from ..constants import TE_DType +from ..constants import TE_DType, DType from ..utils import get_sm_count, _empty_tensor from ..quantized_tensor import Quantizer @@ -287,7 +287,7 @@ def general_grouped_gemm( bias: Optional[List[torch.Tensor]] = None, use_bias: bool = False, use_split_accumulator: bool = False, - D_dtype: Optional[tex.DType] = None, + D_dtype: Optional[DType] = None, single_output=False, ) -> Tuple[List[torch.Tensor], ...]: """ diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index 66bb2dc40e1..d85dcda1593 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -86,6 +86,31 @@ transformer_engine::DType getTransformerEngineFP8Type(bool e4m3_if_hybrid, return transformer_engine::DType::kFloat8E5M2; } +pybind11::object MakePythonDType(transformer_engine::DType dtype) { + constexpr size_t num_dtypes = static_cast(transformer_engine::DType::kNumTypes); + const size_t idx = static_cast(dtype); + NVTE_CHECK(idx < num_dtypes, "Invalid DType (", idx, ")."); + + // Cache one ``transformer_engine.pytorch.DType`` member per value, built once in a + // thread-safe C++11 "magic static" initializer. + static const std::array cache = [] { + pybind11::object dtype_cls = + pybind11::module_::import("transformer_engine.pytorch").attr("DType"); + std::array members; + for (pybind11::handle member : dtype_cls) { + const size_t value = member.attr("value").cast(); + if (value < num_dtypes) { + members[value] = pybind11::reinterpret_borrow(member); + } + } + return members; + }(); + const pybind11::object& member = cache[idx]; + NVTE_CHECK(static_cast(member), "No transformer_engine.pytorch.DType member for DType (", + idx, "); the Python DType enum is out of sync with the C++ enum."); + return member; +} + TensorWrapper makeTransformerEngineTensor(py::handle tensor, py::handle quantizer) { NVTE_CHECK(!tensor.is_none(), "Tensor is not allocated!"); std::unique_ptr my_quantizer = convert_quantizer(quantizer); diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index b376b3022d9..323f8d7c715 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -55,6 +55,7 @@ #include "c10/util/ArrayRef.h" #include "common/util/logging.h" +#include "extensions/pybind_dtype_caster.h" namespace transformer_engine::pytorch { @@ -106,10 +107,19 @@ class Quantizer { const std::vector& shape, DType dtype, std::optional device = std::nullopt, bool pin_memory = false) const = 0; - /*! @brief Construct a grouped tensor with uninitialized data */ + /*! @brief Construct a grouped tensor with uninitialized data + * + * @param tensor_offsets If provided, the precomputed inclusive scan of + * ``first_dims * logical_last_dim`` with a leading zero, used to locate + * each per-group sub-tensor in the shared backing buffer. If null, the + * offsets are computed from ``first_dims`` on demand. Passing this in lets + * callers that already have the scan (e.g. from + * ``tex.splits_to_offsets_multi``) skip the redundant kernel launch. + */ virtual std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, - py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + py::object quantizer, const std::optional& first_dims, + const std::optional& tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const = 0; /*! @brief Convert a PyTorch tensor into a Transformer Engine C++ tensor @@ -151,7 +161,8 @@ class NoneQuantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, - py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + py::object quantizer, const std::optional& first_dims, + const std::optional& tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct a tensor with pre-initialized data */ @@ -182,7 +193,8 @@ class Float8Quantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, - py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + py::object quantizer, const std::optional& first_dims, + const std::optional& tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct a tensor with pre-initialized data */ @@ -217,7 +229,8 @@ class Float8CurrentScalingQuantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, - py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + py::object quantizer, const std::optional& first_dims, + const std::optional& tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct an unquantized tensor with a freshly allocated amax buffer. @@ -280,7 +293,8 @@ class Float8BlockQuantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, - py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + py::object quantizer, const std::optional& first_dims, + const std::optional& tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; std::pair convert_and_update_tensor(py::object shape) const override; @@ -305,7 +319,8 @@ class MXFP8Quantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, - py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + py::object quantizer, const std::optional& first_dims, + const std::optional& tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; std::pair convert_and_update_tensor(py::object shape) const override; @@ -349,7 +364,8 @@ class NVFP4Quantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, - py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + py::object quantizer, const std::optional& first_dims, + const std::optional& tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct an unquantized tensor that shares NVFP4 tensor's amax pointer @@ -364,6 +380,8 @@ class NVFP4Quantizer : public Quantizer { void quantize(const TensorWrapper& input, TensorWrapper& out, const std::optional& noop_flag = std::nullopt) override; + void quantize_impl(const TensorWrapper& input, TensorWrapper& out, + const std::optional& noop_flag, bool compute_amax); /*! @brief Quantize to NVFP4, skipping local amax computation * @@ -376,8 +394,6 @@ class NVFP4Quantizer : public Quantizer { std::vector get_scale_shape(const std::vector& shape, bool columnwise) const; private: - void quantize_impl(const TensorWrapper& input, TensorWrapper& out, - const std::optional& noop_flag, bool compute_amax); void quantize_with_rht_unfused_helper(const TensorWrapper& input, TensorWrapper& out, TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, @@ -392,6 +408,15 @@ std::vector getTensorShape(const at::Tensor& t); transformer_engine::DType getTransformerEngineFP8Type(bool e4m3_if_hybrid, const std::string& fp8_recipe); +/*! @brief Wrap a C++ ``transformer_engine::DType`` as the canonical Python + * ``transformer_engine.pytorch.DType`` ``IntEnum`` member. + * + * The returned object is cached per enum value (one ``py::object`` per + * ``DType``), py::object corresponds to the python ``DType`` enum member + * defined in transformer_engine.pytorch. + */ +pybind11::object MakePythonDType(transformer_engine::DType dtype); + inline size_t typeToNumBits(transformer_engine::DType t) { switch (t) { case transformer_engine::DType::kInt64: diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1f9974448bc..2b4f899e1d6 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -7,10 +7,14 @@ #ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_ #define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_ +#include + #include +#include #include #include #include +#include #include #include @@ -29,19 +33,20 @@ namespace transformer_engine::pytorch { std::tuple fused_topk_with_score_function_fwd( at::Tensor logits, int topk, bool use_pre_softmax, std::optional num_groups, std::optional group_topk, std::optional scaling_factor, std::string score_function, - std::optional expert_bias); + std::optional expert_bias, + int routing_map_format = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP)); -void fused_topk_with_score_function_bwd(int num_tokens, int num_experts, at::Tensor routing_map, - at::Tensor intermediate_output, at::Tensor grad_probs, - at::Tensor grad_logits, int topk, bool use_pre_softmax, - std::optional scaling_factor, - std::string score_function); +void fused_topk_with_score_function_bwd( + at::Tensor routing_map, at::Tensor intermediate_output, at::Tensor grad_probs, + at::Tensor grad_logits, int topk, bool use_pre_softmax, std::optional scaling_factor, + std::string score_function, + int routing_map_format = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP)); std::tuple fused_score_for_moe_aux_loss_fwd( - at::Tensor logits, int topk, std::string score_function); + at::Tensor logits, int topk, std::string score_function, + int routing_map_format = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP)); -void fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts, - at::Tensor intermediate_output, at::Tensor grad_probs, +void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor grad_scores, at::Tensor grad_logits, int topk, std::string score_function); std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, @@ -330,15 +335,28 @@ py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output, std::optional noop_flag); +py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, + const at::Tensor &rowwise_amax, + const at::Tensor &columnwise_amax); + py::object dequantize(const py::handle &input, DType otype); py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, - std::optional first_dims); + std::optional first_dims, + std::optional tensor_offsets); + +py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, + const size_t num_tensors, + std::optional first_dims, + const at::Tensor &rowwise_amax, + const at::Tensor &columnwise_amax, + std::optional tensor_offsets); py::object group_dequantize(const py::handle &input, DType otype); py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, - const size_t num_tensors, std::optional first_dims); + const size_t num_tensors, std::optional first_dims, + std::optional tensor_offsets); std::vector multi_tensor_quantize(const std::vector &tensor_list, std::vector quantizer_list); @@ -486,13 +504,34 @@ size_t get_cublasLt_version(); size_t get_cudnn_version(); at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_dim); +std::tuple> splits_to_offsets_multi( + const at::Tensor &split_sizes, const c10::Device &device, const std::vector &strides, + const std::vector &include_leading_zero, const std::vector &dtypes, + bool bulk_allocate_outputs); at::Tensor copy_data_ptrs_to_device(const std::vector &tensors, const c10::Device &device); -std::tuple> transform_and_copy_data_ptrs_to_device( - const std::string &transform_type, const std::vector &tensors, - const c10::Device &device); +/*************************************************************************************************** + * Experimental helpers for the fused grouped MLP + * + * These primarily exist to support cuDNN CuTe DSL grouped GEMM + * kernels. Since those are unstable and under active development, + * these helpers should also be considered unstable. + **************************************************************************************************/ + +namespace grouped_mlp_experimental { + +// Prepare discrete weight tensors for the cuDNN CuTe DSL grouped GEMM +// kernel by swizzling scales and copying data and scale pointers to +// device. All tensors must share a uniform shape and `swizzle_type` +// must be one of "mxfp8_rowwise", "mxfp8_columnwise", or "nvfp4". +// Returns {data_ptrs_device, scale_ptrs_device, swizzled_scales_buffer}. +std::tuple swizzle_scales_and_pack_ptrs_for_discrete_weights( + const std::vector &data_tensors, const std::vector &scale_tensors, + const std::string &swizzle_type, const c10::Device &device); + +} // namespace grouped_mlp_experimental /*************************************************************************************************** * Support THD format for Context Parallel @@ -647,10 +686,18 @@ void newton_schulz(int64_t ctx_ptr, int64_t m, int64_t n, at::Tensor x, int64_t **************************************************************************************************/ class CommOverlapHelper : torch::CustomClassHolder { + public: + // Shared ownership of an ncclComm_t. The deleter calls ncclCommDestroy when + // the last reference (held by the helper and/or any CommOverlap consumers) + // is released, so the communicator outlives whichever owner is destroyed + // first. + using NcclCommSharedPtr = std::shared_ptr::type>; + private: bool initialized{false}; bool backend_is_nccl{false}; - std::map pgs; + std::map torch_pgs; + std::map nccl_comms; public: int myrank = -1; @@ -671,17 +718,33 @@ class CommOverlapHelper : torch::CustomClassHolder { ExtComm comm); void ub_barrier(ExtComm comm); + + NcclCommSharedPtr get_nccl_comm(std::string comm_name); }; class CommOverlap : torch::CustomClassHolder, public transformer_engine::CommOverlapBase { + private: + // Keeps the cuBLASMp NCCL communicator alive for the lifetime of this + // instance, independent of the CommOverlapHelper that created it. + CommOverlapHelper::NcclCommSharedPtr _nccl_comm; + public: CommOverlap(const std::vector &buffer_shape, at::ScalarType buffer_dtype, - CommOverlapHelper *helper, int tp_size, int num_splits = 3, + CommOverlapHelper *helper, int tp_size, int num_splits = 4, int num_max_streams = NVTE_COMM_OVERLAP_MAX_STREAMS, int comm_cga_size = 2, int gemm_priority = 0, int comm_priority = 0, int num_comm_sm = 16, bool set_sm_margin = true, bool atomic_gemm = false, bool rs_overlap_first_gemm = false); + // cuBLASMp variant. `comm_type`, `buffer_shape`, and `buffer_dtype` size + // the construction-time warmup matmul that primes cuBLASMp's lazy NCCL + // window registrations and workspace allocation so subsequent matmuls + // (including those captured in CUDA graphs) avoid the unsafe lazy paths. + CommOverlap(CommOverlapHelper *helper, int tp_rank, int tp_size, + transformer_engine::CommOverlapType comm_type, + const std::vector &buffer_shape, at::ScalarType buffer_dtype, + int num_comm_sm = 16, bool atomic_gemm = false); + ~CommOverlap() {} using transformer_engine::CommOverlapCore::copy_into_buffer; @@ -695,15 +758,26 @@ class CommOverlap : torch::CustomClassHolder, public transformer_engine::CommOve }; // CommOverlap class CommOverlapP2P : torch::CustomClassHolder, public transformer_engine::CommOverlapP2PBase { + private: + // Keeps the cuBLASMp NCCL communicator alive for the lifetime of this + // instance, independent of the CommOverlapHelper that created it. + CommOverlapHelper::NcclCommSharedPtr _nccl_comm; + public: CommOverlapP2P(const std::vector &buffer_shape, at::ScalarType buffer_dtype, CommOverlapHelper *helper, int tp_size, transformer_engine::CommOverlapType comm_type, - int num_max_streams = NVTE_COMM_OVERLAP_MAX_STREAMS, int comm_cga_size = 2, - int gemm_priority = 0, int comm_priority = 0, int num_comm_sm = 3, - bool set_sm_margin = true, bool atomic_gemm = false, bool use_ce = true, + int num_max_streams = NVTE_COMM_OVERLAP_MAX_STREAMS, int comm_cga_size = 1, + int gemm_priority = 0, int comm_priority = 0, int num_comm_sm = 1, + bool set_sm_margin = false, bool atomic_gemm = false, bool use_ce = true, bool aggregate = false); + // cuBLASMp variant. See CommOverlap for the `comm_type`/buffer args. + CommOverlapP2P(CommOverlapHelper *helper, int tp_rank, int tp_size, + transformer_engine::CommOverlapType comm_type, + const std::vector &buffer_shape, at::ScalarType buffer_dtype, + int num_comm_sm = 1, bool atomic_gemm = false); + ~CommOverlapP2P() {} using transformer_engine::CommOverlapP2PBase::copy_into_buffer; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index d1a9cd85877..49c0a60915a 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -31,6 +31,18 @@ std::vector get_tensor_shape(const TensorWrapper &tensor) { return std::vector(shape.data, shape.data + shape.ndim); } +void allreduce_nvfp4_amax_tensors(NVFP4Quantizer *nvfp4_quantizer_cpp, + std::vector &&amax_tensors) { + if (!nvfp4_quantizer_cpp->with_amax_reduction || amax_tensors.empty()) { + return; + } + c10d::AllreduceCoalescedOptions opts; + opts.reduceOp = c10d::ReduceOp::MAX; + NVTE_SCOPED_GIL_RELEASE({ + nvfp4_quantizer_cpp->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait(); + }); +} + } // namespace py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output, @@ -62,6 +74,57 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob // Perform quantization quantizer_cpp->quantize(input_cpp, output_cpp, noop_flag_cpp); + // Post-quantize swizzle for quantizers whose kernel does not bake + // the GEMM-swizzled scale layout in directly + if (quantizer_cpp->optimize_for_gemm && !output_cpp.get_with_gemm_swizzled_scales()) { + inplace_swizzle_scale_for_gemm(output_py); + } + + return output_py; +} + +py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, + const at::Tensor &rowwise_amax, + const at::Tensor &columnwise_amax) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(tensor.dim() >= 2, "Tensor must be at least 2D"); + NVTE_CHECK(rowwise_amax.is_cuda() && columnwise_amax.is_cuda(), + "Precomputed amax tensors must be CUDA tensors."); + NVTE_CHECK( + rowwise_amax.scalar_type() == at::kFloat && columnwise_amax.scalar_type() == at::kFloat, + "Precomputed amax tensors must be float32."); + NVTE_CHECK(rowwise_amax.numel() == 1 && columnwise_amax.numel() == 1, + "nvfp4_quantize_with_amax expects scalar rowwise and columnwise amaxes."); + + auto quantizer_cpp = convert_quantizer(quantizer); + NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()), + "nvfp4_quantize_with_amax only supports NVFP4 quantizers."); + NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); + + auto input_contiguous = tensor.contiguous(); + auto input_cpp = makeTransformerEngineTensor(input_contiguous); + + const auto shape = get_tensor_shape(input_cpp); + const auto fake_dtype = input_cpp.dtype(); + auto [output_cpp, output_py] = quantizer_cpp->create_tensor(shape, fake_dtype); + + if (output_cpp.get_amax().data_ptr != nullptr) { + output_cpp.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, getTensorShape(rowwise_amax)); + output_py.attr("_amax_rowwise") = py::cast(rowwise_amax); + } + if (output_cpp.get_columnwise_amax().data_ptr != nullptr) { + output_cpp.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32, + getTensorShape(columnwise_amax)); + output_py.attr("_amax_columnwise") = py::cast(columnwise_amax); + } + + nvfp4_quantizer_cpp->quantize_impl(input_cpp, output_cpp, std::nullopt, false); + if (quantizer_cpp->optimize_for_gemm && !output_cpp.get_with_gemm_swizzled_scales()) { + inplace_swizzle_scale_for_gemm(output_py); + } + return output_py; } @@ -78,7 +141,8 @@ namespace { // helper functions for NVFP4 grouped quantization (cuda graph safe with shapes stored in device without D2H copy) void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, GroupedTensorWrapper &grouped_output_tensor, - NVFP4Quantizer *nvfp4_quantizer_cpp, cudaStream_t stream) { + NVFP4Quantizer *nvfp4_quantizer_cpp, cudaStream_t stream, + bool compute_amax) { size_t num_tensors = grouped_input_tensor.num_tensors(); // assert the 2D scaling case, since 2D scaling grouped quant kernel is not ready yet @@ -86,6 +150,10 @@ void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, "2D scaling grouped quant kernel is not ready yet"); NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, "NVFP4 4over6 quantization is not supported for grouped quantization."); + NVTE_CHECK(nvfp4_quantizer_cpp->with_rht, + "graph safe grouped quant kernel for non-RHT path is not ready yet"); + NVTE_CHECK(nvfp4_quantizer_cpp->with_post_rht_amax, + "grouped NVFP4 RHT quantization expects post-RHT amax buffers."); auto quant_config_cpp = QuantizationConfigWrapper(); @@ -116,44 +184,32 @@ void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, quant_config_cpp.set_use_fast_math(true); } - // so far, only the RHT path has grouped kernel support - // grouped kernels for non-RHT path will be added later - - if (nvfp4_quantizer_cpp->with_rht) { - // post-RHT amax or not - if (nvfp4_quantizer_cpp->with_post_rht_amax) { - NVTE_SCOPED_GIL_RELEASE({ - nvte_group_hadamard_transform_amax_graph_safe( - grouped_input_tensor.data(), grouped_output_tensor.data(), 0, - nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, stream); - }); - } else { - NVTE_ERROR("graph safe grouped quant kernel for non-RHT path is not ready yet"); - } - - // RHT cast fusion - auto tile_scheduler_workspace_torch = - at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32)); - auto nvte_tile_scheduler_workspace = - makeTransformerEngineTensor(tile_scheduler_workspace_torch); - - auto rht_matrix_nvte = makeTransformerEngineTensor(nvfp4_quantizer_cpp->rht_matrix); + if (compute_amax) { NVTE_SCOPED_GIL_RELEASE({ - nvte_group_hadamard_transform_cast_fusion_graph_safe( - grouped_input_tensor.data(), grouped_output_tensor.data(), rht_matrix_nvte.data(), - quant_config_cpp, nvte_tile_scheduler_workspace.data(), stream); + nvte_group_hadamard_transform_amax_graph_safe( + grouped_input_tensor.data(), grouped_output_tensor.data(), 0, + nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, stream); }); - - } else { - NVTE_ERROR("graph safe grouped quant kernel for non-RHT path is not ready yet"); } + + // RHT cast fusion + auto tile_scheduler_workspace_torch = at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32)); + auto nvte_tile_scheduler_workspace = makeTransformerEngineTensor(tile_scheduler_workspace_torch); + + auto rht_matrix_nvte = makeTransformerEngineTensor(nvfp4_quantizer_cpp->rht_matrix); + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_hadamard_transform_cast_fusion_graph_safe( + grouped_input_tensor.data(), grouped_output_tensor.data(), rht_matrix_nvte.data(), + quant_config_cpp, nvte_tile_scheduler_workspace.data(), stream); + }); } } // namespace // NOTE: Only supports varying first dim. py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, - std::optional first_dims) { + std::optional first_dims, + std::optional tensor_offsets) { using namespace transformer_engine::pytorch::detail; init_extension(); @@ -178,7 +234,7 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const // Create output GroupedTensor. auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), - py::reinterpret_borrow(quantizer), first_dims, logical_first_dim, + py::reinterpret_borrow(quantizer), first_dims, tensor_offsets, logical_first_dim, logical_last_dim); // dispatch to scaling methods @@ -207,7 +263,7 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const // NVFP4 grouped quantization NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, - nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream()); + nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), true); break; } case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: { @@ -227,8 +283,82 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const return py::reinterpret_borrow(grouped_output_py); } +py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, + const size_t num_tensors, + std::optional first_dims, + const at::Tensor &rowwise_amax, + const at::Tensor &columnwise_amax, + std::optional tensor_offsets) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(rowwise_amax.is_cuda() && columnwise_amax.is_cuda(), + "Precomputed amax tensors must be CUDA tensors."); + NVTE_CHECK( + rowwise_amax.scalar_type() == at::kFloat && columnwise_amax.scalar_type() == at::kFloat, + "Precomputed amax tensors must be float32."); + NVTE_CHECK(rowwise_amax.numel() == static_cast(num_tensors), + "Rowwise amax must contain one value per group."); + NVTE_CHECK(columnwise_amax.numel() == static_cast(num_tensors), + "Columnwise amax must contain one value per group."); + + std::vector logical_shape; + for (const auto &d : tensor.sizes()) { + logical_shape.push_back(d); + } + const auto logical_first_dim = logical_shape[0]; + const auto logical_last_dim = logical_shape[1]; + + bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0; + + auto quantizer_cpp = convert_quantizer(quantizer); + NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()), + "nvfp4_group_quantize_with_amax only supports NVFP4 quantizers."); + NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); + + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_input_tensor.set_rowwise_data( + tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor)); + + auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( + num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims, tensor_offsets, logical_first_dim, + logical_last_dim); + + if (grouped_output_tensor_cpp.get_amax().data_ptr != nullptr) { + grouped_output_tensor_cpp.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, + getTensorShape(rowwise_amax)); + grouped_output_py.attr("amax") = py::cast(rowwise_amax); + } + if (grouped_output_tensor_cpp.get_columnwise_amax().data_ptr != nullptr) { + grouped_output_tensor_cpp.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32, + getTensorShape(columnwise_amax)); + grouped_output_py.attr("columnwise_amax") = py::cast(columnwise_amax); + } + + std::vector amax_tensors; + if (grouped_output_tensor_cpp.get_amax().data_ptr != nullptr) { + amax_tensors.push_back(rowwise_amax); + } + if (grouped_output_tensor_cpp.get_columnwise_amax().data_ptr != nullptr) { + amax_tensors.push_back(columnwise_amax); + } + allreduce_nvfp4_amax_tensors(nvfp4_quantizer_cpp, std::move(amax_tensors)); + + if (empty_input_buffer) { + return py::reinterpret_borrow(grouped_output_py); + } + + group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, nvfp4_quantizer_cpp, + at::cuda::getCurrentCUDAStream(), false); + + return py::reinterpret_borrow(grouped_output_py); +} + py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, - const size_t num_tensors, std::optional first_dims) { + const size_t num_tensors, std::optional first_dims, + std::optional tensor_offsets) { using namespace transformer_engine::pytorch::detail; init_extension(); @@ -254,7 +384,7 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), - py::reinterpret_borrow(quantizer), first_dims, logical_first_dim, + py::reinterpret_borrow(quantizer), first_dims, tensor_offsets, logical_first_dim, logical_last_dim); if (empty_input_buffer) { @@ -343,7 +473,7 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o NoneQuantizer q{py::none()}; auto [out_cpp, out_py] = q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims, - logical_first_dim, logical_last_dim); + tensor_offsets, logical_first_dim, logical_last_dim); return py::reinterpret_borrow(out_py); } @@ -381,8 +511,9 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o // Create output GroupedTensor using NoneQuantizer. NoneQuantizer q{py::none()}; - auto [out_cpp, out_py] = q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), - first_dims, logical_first_dim, logical_last_dim); + auto [out_cpp, out_py] = + q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims, + tensor_offsets, logical_first_dim, logical_last_dim); NVTE_SCOPED_GIL_RELEASE({ nvte_group_dequantize(input_cpp.data(), out_cpp.data(), at::cuda::getCurrentCUDAStream()); @@ -575,10 +706,10 @@ std::tuple, std::vector> bulk_allocate_fp py::object columnwise_scale = (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none()); - // Construct Python tensor - tensor_py_list.emplace_back( - Float8BlockwiseQTensorClass(rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, - fp8_dtype, quantizer_py_list[i], is_2D_scaled)); + // Construct Python tensor. + tensor_py_list.emplace_back(Float8BlockwiseQTensorClass( + rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, MakePythonDType(fp8_dtype), + quantizer_py_list[i], is_2D_scaled)); // Construct C++ tensor tensor_cpp_list.emplace_back(makeTransformerEngineTensor( @@ -681,10 +812,10 @@ std::tuple, std::vector> bulk_allocate_mx py::object columnwise_scale = (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none()); - // Construct Python tensor + // Construct Python tensor. tensor_py_list.emplace_back(MXFP8TensorClass(rowwise_data, rowwise_scale, columnwise_data, - columnwise_scale, fp8_dtype, quantizer_py_list[i], - with_gemm_swizzled_scales)); + columnwise_scale, MakePythonDType(fp8_dtype), + quantizer_py_list[i], with_gemm_swizzled_scales)); // Construct C++ tensor tensor_cpp_list.emplace_back(makeTransformerEngineTensor( @@ -870,13 +1001,13 @@ std::tuple, std::vector, bool> bulk_alloc py::object amax_rowwise = rowwise_usage ? py::cast(amax_rowwise_list[i]) : py::none(); py::object amax_columnwise = columnwise_usage ? py::cast(amax_columnwise_list[i]) : py::none(); - // Construct Python tensor - tensor_py_list.emplace_back( - NVFP4TensorClass(rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, - amax_rowwise, amax_columnwise, fp4_dtype, quantizer_py_list[i], - with_gemm_swizzled_scales, py::arg("row_scaled_nvfp4") = row_scaled_nvfp4, - py::arg("nvfp4_use_4over6") = nvfp4_use_4over6, - py::arg("nvfp4_e4m3_max") = nvfp4_e4m3_max)); + // Construct Python tensor. + tensor_py_list.emplace_back(NVFP4TensorClass( + rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, amax_rowwise, + amax_columnwise, MakePythonDType(fp4_dtype), quantizer_py_list[i], + with_gemm_swizzled_scales, py::arg("row_scaled_nvfp4") = row_scaled_nvfp4, + py::arg("nvfp4_use_4over6") = nvfp4_use_4over6, + py::arg("nvfp4_e4m3_max") = nvfp4_e4m3_max)); // Construct C++ tensor // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor, diff --git a/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp b/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp index a126ab0d604..33237f0751d 100644 --- a/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp +++ b/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp @@ -3,6 +3,9 @@ * * See LICENSE for license information. ************************************************************************/ +#ifdef NVTE_WITH_CUBLASMP +#include +#endif #include "../extensions.h" #include "transformer_engine/transformer_engine.h" @@ -28,20 +31,20 @@ CommOverlapHelper::CommOverlapHelper() { CommOverlapHelper::CommOverlapHelper(c10d::ProcessGroup *world_group, std::optional intra_domain_group) { #ifndef NVTE_UB_WITH_MPI - pgs.insert({"world", world_group}); - myrank = pgs["world"]->getRank(); - numranks = pgs["world"]->getSize(); - c10d::ProcessGroup::BackendType backend = pgs["world"]->getBackendType(); + torch_pgs.insert({"world", world_group}); + myrank = torch_pgs["world"]->getRank(); + numranks = torch_pgs["world"]->getSize(); + c10d::ProcessGroup::BackendType backend = torch_pgs["world"]->getBackendType(); backend_is_nccl = (backend == c10d::ProcessGroup::BackendType::NCCL); if (intra_domain_group.has_value()) { // Get local rank on node and number of local ranks NVTE_CHECK(intra_domain_group.value()->getBackendType() == backend, "Internal TE error: Intra-node group must be on the same backend (%s) as the world ", - "group!", pgs["world"]->getBackendName()); - pgs.insert({"intra", intra_domain_group.value()}); - mylocal = pgs["intra"]->getRank(); - numlocal = pgs["intra"]->getSize(); + "group!", torch_pgs["world"]->getBackendName()); + torch_pgs.insert({"intra", intra_domain_group.value()}); + mylocal = torch_pgs["intra"]->getRank(); + numlocal = torch_pgs["intra"]->getSize(); if (numlocal == numranks) { // Intra-node group is same as the world group so there can only be 1 node @@ -60,13 +63,68 @@ CommOverlapHelper::CommOverlapHelper(c10d::ProcessGroup *world_group, // Intra-node group is not set so we assume there is only 1 node mylocal = myrank; numlocal = numranks; - pgs.insert({"intra", world_group}); + torch_pgs.insert({"intra", world_group}); mynode = 0; numnodes = 1; } initialized = true; + +#ifdef NVTE_WITH_CUBLASMP + // Initialize world NCCL communicator via ncclCommInitRank (one GPU per process under torchrun) + ncclUniqueId nccl_world_id; + if (myrank == 0) { + NVTE_CHECK_NCCL(ncclGetUniqueId(&nccl_world_id)); + } + auto nccl_world_id_tensor = + torch::from_blob(reinterpret_cast(&nccl_world_id), {sizeof(ncclUniqueId)}, + at::device(torch::kCPU).dtype(torch::kUInt8)); + nccl_world_id_tensor = (backend_is_nccl) ? nccl_world_id_tensor.cuda() : nccl_world_id_tensor; + { + c10d::BroadcastOptions bcast_opts; + bcast_opts.rootRank = 0; + std::vector bcast_tensors = {nccl_world_id_tensor}; + auto work = torch_pgs["world"]->broadcast(bcast_tensors, bcast_opts); + work->wait(); + } + nccl_world_id_tensor = (backend_is_nccl) ? nccl_world_id_tensor.cpu() : nccl_world_id_tensor; + nccl_world_id = *reinterpret_cast(nccl_world_id_tensor.data_ptr()); + + ncclComm_t nccl_world; + NVTE_CHECK_NCCL(ncclCommInitRank(&nccl_world, numranks, nccl_world_id, myrank)); + nccl_comms.insert({"world", NcclCommSharedPtr(nccl_world, ncclCommDestroy)}); + + if (intra_domain_group.has_value()) { + // Generate a separate unique ID for the intra-node communicator + ncclUniqueId nccl_intra_id; + if (mylocal == 0) { + NVTE_CHECK_NCCL(ncclGetUniqueId(&nccl_intra_id)); + } + + // Broadcast the intra-node unique ID from the local root to all local ranks + auto nccl_intra_id_tensor = + torch::from_blob(reinterpret_cast(&nccl_intra_id), {sizeof(ncclUniqueId)}, + at::device(torch::kCPU).dtype(torch::kUInt8)); + nccl_intra_id_tensor = (backend_is_nccl) ? nccl_intra_id_tensor.cuda() : nccl_intra_id_tensor; + { + c10d::BroadcastOptions bcast_opts; + bcast_opts.rootRank = 0; + std::vector bcast_tensors = {nccl_intra_id_tensor}; + auto work = torch_pgs["intra"]->broadcast(bcast_tensors, bcast_opts); + work->wait(); + } + nccl_intra_id_tensor = (backend_is_nccl) ? nccl_intra_id_tensor.cpu() : nccl_intra_id_tensor; + nccl_intra_id = *reinterpret_cast(nccl_intra_id_tensor.data_ptr()); + + // Initialize intra-node communicator + ncclComm_t nccl_intra; + NVTE_CHECK_NCCL(ncclCommInitRank(&nccl_intra, numlocal, nccl_intra_id, mylocal)); + nccl_comms.insert({"intra", NcclCommSharedPtr(nccl_intra, ncclCommDestroy)}); + } else { + nccl_comms.insert({"intra", nccl_comms["world"]}); + } +#endif #else NVTE_ERROR("Internal TE error: CommOverlapHelper cannot be initialized with valid PyTorch ", "distributed process groups when TE is compiled with NVTE_UB_WITH_MPI=1!"); @@ -75,9 +133,18 @@ CommOverlapHelper::CommOverlapHelper(c10d::ProcessGroup *world_group, CommOverlapHelper::~CommOverlapHelper() { #ifndef NVTE_UB_WITH_MPI - for (auto &pg : pgs) pg.second = nullptr; + for (auto &pg : torch_pgs) { + pg.second = nullptr; + } + torch_pgs.clear(); backend_is_nccl = false; initialized = false; +#ifdef NVTE_WITH_CUBLASMP + // Releasing the helper's references is enough: each shared_ptr's deleter + // calls ncclCommDestroy once the last owner (helper or any consuming + // CommOverlap/CommOverlapP2P) drops it. + nccl_comms.clear(); +#endif #endif } @@ -96,9 +163,10 @@ void CommOverlapHelper::ub_allgather(void *globaldata, size_t globalbytes, void at::device(torch::kCPU).dtype(torch::kUInt8)); auto globaltmp = (backend_is_nccl) ? globaltensor.cuda() : globaltensor; - std::vector> globalchunks = {globaltmp.chunk(pgs[group]->getSize())}; + std::vector> globalchunks = { + globaltmp.chunk(torch_pgs[group]->getSize())}; std::vector localchunk = {localtmp}; - auto work = pgs[group]->allgather(globalchunks, localchunk); + auto work = torch_pgs[group]->allgather(globalchunks, localchunk); work->wait(); if (backend_is_nccl) { @@ -116,7 +184,7 @@ void CommOverlapHelper::ub_barrier(ExtComm group) { #ifndef NVTE_UB_WITH_MPI NVTE_CHECK(initialized, "Internal TE error: tex.CommOverlapHelper() is not initialized ", "with valid process groups!"); - auto work = pgs[group]->barrier(); + auto work = torch_pgs[group]->barrier(); work->wait(); #else NVTE_ERROR("Internal TE error: CommOverlapHelper::ub_barrier is a no-op when TE is compiled ", @@ -124,6 +192,23 @@ void CommOverlapHelper::ub_barrier(ExtComm group) { #endif } +CommOverlapHelper::NcclCommSharedPtr CommOverlapHelper::get_nccl_comm(std::string comm_name) { +#ifdef NVTE_WITH_CUBLASMP + NVTE_CHECK(initialized, "Internal TE error: tex.CommOverlapHelper() is not initialized ", + "with valid process groups!"); + auto it = nccl_comms.find(comm_name); + if (it != nccl_comms.end()) { + return it->second; + } else { + NVTE_ERROR("Internal TE error: No NCCL communicator found with name ", comm_name, "!"); + } +#else + NVTE_ERROR( + "Internal TE error: CommOverlapHelper::get_nccl_comm() is an internal API that requires TE " + "to be built with NVTE_WITH_CUBLASMP=1!"); +#endif +} + /*************************************************************************************************** * CommOverlap **************************************************************************************************/ @@ -141,6 +226,101 @@ CommOverlap::CommOverlap(const std::vector &buffer_shape, at::ScalarType num_max_streams, comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm) {} +namespace { + +// Run a dummy cuBLASMp matmul during construction so its lazy NCCL window +// registration and workspace allocation happen outside any CUDA-graph +// capture. The warmup is sized from the comm buffer so the cached +// workspace covers any matmul the caller will later run with the same +// descriptor. BF16 is used unconditionally; its workspace is at least as +// large as the FP8 workspace for the same m/n/k. +void cublasmp_capture_warmup(te::CommOverlapCore *core, int tp_size, te::CommOverlapType comm_type, + const std::vector &buffer_shape) { + NVTE_CHECK(buffer_shape.size() == 2, "cuBLASMp warmup expects a 2-D buffer shape, got rank ", + buffer_shape.size()); + // Treat the matmul as square in the weight dim so workspace is sized + // for the wider of the two cases. + const int64_t N_global = static_cast(buffer_shape[0]); + const int64_t hidden = static_cast(buffer_shape[1]); + auto ceil_div = [](int64_t a, int64_t b) { return (a + b - 1) / b; }; + const int64_t M_local = ceil_div(hidden, tp_size); + const int64_t N_local = ceil_div(N_global, tp_size); + const int64_t K_local = ceil_div(hidden, tp_size); + const int64_t bf16_bytes = 2; + + std::vector a_shape, b_shape, d_shape; + if (comm_type == te::CommOverlapType::AG) { + // A = (M_local, K), B = (N_local, K), D = (N_global, M_local) + a_shape = {static_cast(M_local), static_cast(hidden)}; + b_shape = {static_cast(N_local), static_cast(hidden)}; + d_shape = {static_cast(N_global), static_cast(M_local)}; + } else { // RS (or AR -- same descriptor-level dims) + // A = (M_global, K_local), B = (N_global, K_local), D = (N_local, M_global) + a_shape = {static_cast(hidden), static_cast(K_local)}; + b_shape = {static_cast(N_global), static_cast(K_local)}; + d_shape = {static_cast(N_local), static_cast(hidden)}; + } + + const size_t a_bytes = a_shape[0] * a_shape[1] * bf16_bytes; + const size_t b_bytes = b_shape[0] * b_shape[1] * bf16_bytes; + const size_t d_bytes = d_shape[0] * d_shape[1] * bf16_bytes; + + void *a_ptr = nullptr; + void *b_ptr = nullptr; + void *d_ptr = nullptr; + NVTE_CHECK_CUDA(cudaMalloc(&a_ptr, a_bytes)); + NVTE_CHECK_CUDA(cudaMalloc(&b_ptr, b_bytes)); + NVTE_CHECK_CUDA(cudaMalloc(&d_ptr, d_bytes)); + NVTE_CHECK_CUDA(cudaMemset(a_ptr, 0, a_bytes)); + NVTE_CHECK_CUDA(cudaMemset(b_ptr, 0, b_bytes)); + + te::TensorWrapper A_tw, B_tw, D_tw, bias_tw, pre_gelu_tw, dummy; + A_tw.set_rowwise_data(a_ptr, te::DType::kBFloat16, a_shape); + B_tw.set_rowwise_data(b_ptr, te::DType::kBFloat16, b_shape); + D_tw.set_rowwise_data(d_ptr, te::DType::kBFloat16, d_shape); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + if (comm_type == te::CommOverlapType::AG) { + if (core->is_atomic_gemm()) { + core->atomic_gemm_overlap_ag( + A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy, + /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream); + } else { + core->split_overlap_ag( + A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy, + /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream); + } + } else { + if (core->is_atomic_gemm()) { + core->atomic_gemm_overlap_rs( + A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy, + /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream); + } else { + core->split_overlap_rs( + A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy, + /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream); + } + } + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + cudaFree(a_ptr); + cudaFree(b_ptr); + cudaFree(d_ptr); +} + +} // namespace + +CommOverlap::CommOverlap(CommOverlapHelper *helper, int tp_rank, int tp_size, + te::CommOverlapType comm_type, const std::vector &buffer_shape, + at::ScalarType buffer_dtype, int num_comm_sm, bool atomic_gemm) + : te::CommOverlapBase(helper->get_nccl_comm("intra").get(), tp_rank, tp_size, num_comm_sm, + atomic_gemm), + _nccl_comm(helper->get_nccl_comm("intra")) { + // buffer_dtype is unused on this path (the warmup runs in BF16); kept in + // the signature for API symmetry with the non-cuBLASMp ctor. + (void)buffer_dtype; + cublasmp_capture_warmup(this, tp_size, comm_type, buffer_shape); +} + /* ** Helper function to copy input to _ubuf */ @@ -240,6 +420,18 @@ CommOverlapP2P::CommOverlapP2P(const std::vector &buffer_shape, at::Scal comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin, use_ce, atomic_gemm, aggregate) {} +CommOverlapP2P::CommOverlapP2P(CommOverlapHelper *helper, int tp_rank, int tp_size, + te::CommOverlapType comm_type, + const std::vector &buffer_shape, at::ScalarType buffer_dtype, + int num_comm_sm, bool atomic_gemm) + : te::CommOverlapP2PBase(helper->get_nccl_comm("intra").get(), tp_rank, tp_size, num_comm_sm, + atomic_gemm), + _nccl_comm(helper->get_nccl_comm("intra")) { + // See CommOverlap constructor for the buffer_dtype rationale. + (void)buffer_dtype; + cublasmp_capture_warmup(this, tp_size, comm_type, buffer_shape); +} + /* ** Copy input to _ubufs[0] */ diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 53ea76d83b4..b1e552ec8be 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -39,7 +39,8 @@ bool is_low_precision(const DType type) { } std::vector getGemmOutputShape(const NVTEShape& A_shape, const bool transa, - const NVTEShape& B_shape, const bool transb) { + const NVTEShape& B_shape, const bool transb, + size_t tp_size = 1, size_t tp_dim = 0) { // Flatten outer dims to get 2D matrices const auto [A0, A1] = get_2d_dims(A_shape); const auto [B0, B1] = get_2d_dims(B_shape); @@ -52,21 +53,36 @@ std::vector getGemmOutputShape(const NVTEShape& A_shape, const bool tran std::vector ret; if (transb) { ret.emplace_back(B1); - } else { + } else if (tp_size == 1) { // Unflatten B0 for (size_t i = 0; i < B_shape.ndim - 1; ++i) { ret.emplace_back(B_shape.data[i]); } + } else { + // Keep output tensor in 2D for comm+GEMM overlap + ret.emplace_back(B0); } if (transa) { ret.emplace_back(A0); } else { ret.emplace_back(A1); } + + // Correct output dims for comm+GEMM overlap if needed + if (tp_size > 1) { + if (tp_dim == 0) { + // Outer dim is sharded, comm+GEMM overlap would need to do all-gather + ret[0] *= tp_size; + } else { + // Inner dim is sharded, comm+GEMM overlap would need to do reduce-scatter + ret[0] /= tp_size; + } + } return ret; } -bool checkGemmShape(const std::vector& expected, const NVTEShape& actual) { +bool checkGemmShape(const std::vector& expected, const NVTEShape& actual, + size_t tp_size = 1, size_t tp_dim = 0) { if (expected.size() != actual.ndim) return false; for (size_t i = 0; i < expected.size(); ++i) { if (expected[i] != actual.data[i]) return false; @@ -152,11 +168,18 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans A_tensor.scaling_mode() == NVTE_BLOCK_SCALING_2D || B_tensor.scaling_mode() == NVTE_BLOCK_SCALING_1D || B_tensor.scaling_mode() == NVTE_BLOCK_SCALING_2D; - + // Get TP info for comm+GEMM overlap + size_t tp_size = 1; + size_t tp_dim = 0; + if (comm_overlap && !bulk_overlap && comm_overlap->with_cublasmp()) { + tp_size = comm_overlap->get_tp_size(); + tp_dim = (comm_type.value() == CommOverlapType::AG) ? 0 : 1; + } // Check tensor dimensions const auto& A_shape = A_tensor.shape(); const auto& B_shape = B_tensor.shape(); - const auto& D_shape = detail::getGemmOutputShape(A_shape, transa, B_shape, transb); + const auto& D_shape = + detail::getGemmOutputShape(A_shape, transa, B_shape, transb, tp_size, tp_dim); NVTE_CHECK(A_shape.ndim >= 1, "Tensor A needs to have at least 1 dimension"); NVTE_CHECK(B_shape.ndim >= 1, "Tensor B needs to have at least 1 dimension"); @@ -538,6 +561,10 @@ std::optional> te_general_grouped_gemm( te_pre_gelu_out_wrappers.emplace_back(std::move(te_pre_gelu_out)); } + if (te_A_wrappers.empty()) { + return bias; + } + // Keep the swizzled scaling factor tensors alive during the GEMM. std::vector> swizzled_scale_inverses_list; diff --git a/transformer_engine/pytorch/csrc/extensions/grouped_mlp_experimental.cpp b/transformer_engine/pytorch/csrc/extensions/grouped_mlp_experimental.cpp new file mode 100644 index 00000000000..0ab8bc8d61b --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/grouped_mlp_experimental.cpp @@ -0,0 +1,158 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Experimental helpers for the fused grouped MLP. + +#include + +#include +#include +#include +#include + +#include "common/common.h" +#include "extensions.h" + +namespace transformer_engine { +namespace pytorch { +namespace grouped_mlp_experimental { + +std::tuple swizzle_scales_and_pack_ptrs_for_discrete_weights( + const std::vector &data_tensors, const std::vector &scale_tensors, + const std::string &swizzle_type_str, const c10::Device &device) { + const size_t num_tensors = data_tensors.size(); + NVTE_CHECK(scale_tensors.size() == num_tensors, + "Expected data_tensors and scale_tensors to have matching sizes, but got ", + num_tensors, " and ", scale_tensors.size(), "."); + + // Parse swizzle type + enum class SwizzleType { Invalid, MXFP8Rowwise, MXFP8Columnwise, NVFP4 }; + SwizzleType swizzle_type = SwizzleType::Invalid; + if (swizzle_type_str == "mxfp8_rowwise") { + swizzle_type = SwizzleType::MXFP8Rowwise; + } else if (swizzle_type_str == "mxfp8_columnwise") { + swizzle_type = SwizzleType::MXFP8Columnwise; + } else if (swizzle_type_str == "nvfp4") { + swizzle_type = SwizzleType::NVFP4; + } else { + NVTE_ERROR("Unsupported swizzle type (", swizzle_type_str, + "). Expected one of: mxfp8_rowwise, mxfp8_columnwise, nvfp4."); + } + + // Trivial case: no tensors. Return empty tensors. + if (num_tensors == 0) { + auto empty_ptrs = at::empty({0}, at::TensorOptions().dtype(at::kLong).device(device)); + auto empty_scales = at::empty({0}, at::TensorOptions().dtype(at::kByte).device(device)); + return {empty_ptrs, empty_ptrs.clone(), std::move(empty_scales)}; + } + + // CUDA stream + auto stream = at::cuda::getCurrentCUDAStream(); + + // Tensor properties + NVTEScalingMode scaling_mode; + transformer_engine::DType data_dtype, scale_dtype; + NVTETensorParam data_param_name, scale_param_name; + switch (swizzle_type) { + case SwizzleType::MXFP8Rowwise: + case SwizzleType::MXFP8Columnwise: + scaling_mode = NVTE_MXFP8_1D_SCALING; + data_dtype = transformer_engine::DType::kFloat8E4M3; + scale_dtype = transformer_engine::DType::kFloat8E8M0; + if (swizzle_type == SwizzleType::MXFP8Rowwise) { + data_param_name = kNVTERowwiseData; + scale_param_name = kNVTERowwiseScaleInv; + } else { + data_param_name = kNVTEColumnwiseData; + scale_param_name = kNVTEColumnwiseScaleInv; + } + break; + case SwizzleType::NVFP4: + scaling_mode = NVTE_NVFP4_1D_SCALING; + data_dtype = transformer_engine::DType::kFloat4E2M1; + scale_dtype = transformer_engine::DType::kFloat8E4M3; + data_param_name = kNVTERowwiseData; + scale_param_name = kNVTERowwiseScaleInv; + break; + default: + NVTE_ERROR("Unsupported swizzle type (", static_cast(swizzle_type), ")."); + } + + // Data shape + NVTEShape data_shape = convertTorchShape(data_tensors[0].sizes()); + if (swizzle_type == SwizzleType::NVFP4) { + // NVFP4 packs two 4-bit values per byte + NVTE_CHECK(data_shape.ndim > 0, "Invalid shape for NVFP4 data tensor (", + getTensorShape(data_tensors[0]), ")."); + data_shape.data[data_shape.ndim - 1] *= 2; + } + + // Scale shape + const NVTEShape scale_shape = convertTorchShape(scale_tensors[0].sizes()); + NVTE_CHECK(scale_shape.ndim == 2, + "Expected 2D scale tensor, but got shape=", getTensorShape(scale_tensors[0]), "."); + const size_t scale_numel = scale_shape.data[0] * scale_shape.data[1]; + const size_t scale_dtype_bits = transformer_engine::pytorch::typeToNumBits(scale_dtype); + const size_t scale_bytes = ceildiv(scale_numel * scale_dtype_bits, 8); + + // Allocate single buffer for swizzled scales. Uses a uniform stride since + // all tensors share the same scale shape. + const size_t swizzled_scales_stride = roundup(scale_bytes, 16); // Align to 16 bytes + auto swizzled_scales = at::empty({static_cast(swizzled_scales_stride * num_tensors)}, + at::TensorOptions().dtype(at::kByte).device(device)); + uint8_t *swizzled_scales_dptr = reinterpret_cast(swizzled_scales.data_ptr()); + + // Allocate input/output NVTETensors as a single batch. The first + // num_tensors entries are inputs; the next num_tensors are outputs. + MultiTensorWrapper nvte_tensors(2 * num_tensors, scaling_mode); + NVTETensor *inputs_nvte = nvte_tensors.data(); + NVTETensor *outputs_nvte = nvte_tensors.data() + num_tensors; + + auto set_param = [](NVTETensor t, NVTETensorParam param, void *dptr, + transformer_engine::DType dtype, const NVTEShape &shape) { + NVTEBasicTensor data{dptr, static_cast(dtype), shape}; + nvte_set_tensor_param_v2(t, param, &data, sizeof(data)); + }; + + // Configure NVTETensors + for (size_t i = 0; i < num_tensors; ++i) { + const uint8_t swizzled_flag = 1; + nvte_set_tensor_param_v2(outputs_nvte[i], kNVTEWithGEMMSwizzledScales, &swizzled_flag, + sizeof(swizzled_flag)); + void *in_scale_ptr = scale_tensors[i].data_ptr(); + void *out_scale_ptr = swizzled_scales_dptr + i * swizzled_scales_stride; + set_param(inputs_nvte[i], data_param_name, nullptr, data_dtype, data_shape); + set_param(inputs_nvte[i], scale_param_name, in_scale_ptr, scale_dtype, scale_shape); + set_param(outputs_nvte[i], data_param_name, nullptr, data_dtype, data_shape); + set_param(outputs_nvte[i], scale_param_name, out_scale_ptr, scale_dtype, scale_shape); + } + + // Launch swizzle kernel + nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte, outputs_nvte, num_tensors, stream); + + // Pack data pointers (first half) and swizzled scale pointers (second half) + // into a single host buffer and copy to device with one kernel launch. + std::vector packed_ptrs_host(2 * num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + packed_ptrs_host[i] = reinterpret_cast(data_tensors[i].data_ptr()); + packed_ptrs_host[num_tensors + i] = + reinterpret_cast(swizzled_scales_dptr + i * swizzled_scales_stride); + } + auto packed_ptrs_device = at::empty({static_cast(2 * num_tensors)}, + at::TensorOptions().dtype(at::kLong).device(device)); + nvte_copy_host_to_device_via_kernel(packed_ptrs_host.data(), packed_ptrs_device.data_ptr(), + 2 * num_tensors * sizeof(uint64_t), stream); + + // Return the two pointer arrays as views into the packed device buffer. + auto data_ptrs = packed_ptrs_device.narrow(0, 0, static_cast(num_tensors)); + auto scale_ptrs = packed_ptrs_device.narrow(0, static_cast(num_tensors), + static_cast(num_tensors)); + return {std::move(data_ptrs), std::move(scale_ptrs), std::move(swizzled_scales)}; +} + +} // namespace grouped_mlp_experimental +} // namespace pytorch +} // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/extensions/misc.cpp b/transformer_engine/pytorch/csrc/extensions/misc.cpp index c5707fa53c9..ba4371ffe1e 100644 --- a/transformer_engine/pytorch/csrc/extensions/misc.cpp +++ b/transformer_engine/pytorch/csrc/extensions/misc.cpp @@ -4,7 +4,15 @@ * See LICENSE for license information. ************************************************************************/ +#include + +#include +#include +#include + #include "../extensions.h" +#include "common/common.h" +#include "pybind.h" namespace transformer_engine::pytorch { @@ -30,4 +38,88 @@ at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_ return output; } +std::tuple> splits_to_offsets_multi( + const at::Tensor &split_sizes, const c10::Device &device, const std::vector &strides, + const std::vector &include_leading_zero, const std::vector &dtypes, + bool bulk_allocate_outputs) { + const size_t num_outputs = strides.size(); + const size_t num_splits = static_cast(split_sizes.numel()); + + // Check inputs. + NVTE_CHECK(include_leading_zero.size() == num_outputs && dtypes.size() == num_outputs, + "strides, include_leading_zero, and dtypes must have matching lengths, but got ", + strides.size(), ", ", include_leading_zero.size(), ", and ", dtypes.size(), "."); + NVTE_CHECK(device.is_cuda(), "device must be CUDA, but got ", device.str(), "."); + + // Convert split sizes to int64 GPU tensor. + const at::Tensor split_sizes_i64 = + split_sizes.scalar_type() == at::kLong ? split_sizes : split_sizes.to(at::kLong); + const at::Tensor split_sizes_out = + split_sizes_i64.device() == device ? split_sizes_i64 : split_sizes_i64.to(device); + + // Allocate outputs. + std::vector outputs; + outputs.reserve(num_outputs); + if (bulk_allocate_outputs) { + std::vector> shapes; + shapes.reserve(num_outputs); + for (size_t i = 0; i < num_outputs; ++i) { + const size_t length = num_splits + (include_leading_zero[i] ? 1 : 0); + shapes.emplace_back(std::vector{length}); + } + // cuDNN CuTe DSL grouped GEMM kernels require padded_offsets + // aligned to 16 bytes. + const std::vector alignments(num_outputs, 16); + outputs = bulk_allocate(shapes, dtypes, device, alignments); + } else { + for (size_t i = 0; i < num_outputs; ++i) { + const int64_t length = static_cast(num_splits) + (include_leading_zero[i] ? 1 : 0); + outputs.emplace_back( + at::empty({length}, at::TensorOptions().dtype(dtypes[i]).device(device))); + } + } + + // Construct NVTETensors. + MultiTensorWrapper outputs_nvte(num_outputs); + std::vector include_leading_zero_int(num_outputs); + for (size_t i = 0; i < num_outputs; ++i) { + const size_t length = num_splits + (include_leading_zero[i] ? 1 : 0); + NVTEShape shape = nvte_make_shape(&length, 1); + NVTEBasicTensor data = {outputs[i].data_ptr(), + static_cast(GetTransformerEngineDType(dtypes[i])), shape}; + nvte_set_tensor_param_v2(outputs_nvte[i], kNVTERowwiseData, &data, sizeof(data)); + include_leading_zero_int[i] = include_leading_zero[i] ? 1 : 0; + } + + auto split_sizes_nvte = makeTransformerEngineTensor(split_sizes_out); + NVTE_SCOPED_GIL_RELEASE({ + nvte_splits_to_offsets_multi(split_sizes_nvte.data(), outputs_nvte.data(), strides.data(), + include_leading_zero_int.data(), num_outputs, + at::cuda::getCurrentCUDAStream()); + }); + + return {split_sizes_out, std::move(outputs)}; +} + +at::Tensor copy_data_ptrs_to_device(const std::vector &tensors, + const c10::Device &device) { + // Collect data pointers + std::vector ptrs_host; + ptrs_host.reserve(tensors.size()); + for (const auto &tensor : tensors) { + ptrs_host.push_back(reinterpret_cast(tensor.data_ptr())); + } + + // Allocate device buffer + auto ptrs_device = at::empty({static_cast(tensors.size())}, + at::TensorOptions().dtype(at::kLong).device(device)); + + // Load pointers on device + nvte_copy_host_to_device_via_kernel(ptrs_host.data(), ptrs_device.data_ptr(), + tensors.size() * sizeof(uint64_t), + at::cuda::getCurrentCUDAStream()); + + return ptrs_device; +} + } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c1b38a22757..d6089b1e01e 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -129,12 +129,70 @@ void init_extension() { }); } +// Pybind11 registrations for the fused MoE router kernels. Split out of +// PYBIND11_MODULE() to keep that function under the cpplint readability/fn_size +// limit. +void init_router_bindings(pybind11::module &m) { + pybind11::enum_(m, "NVTERoutingMapFormat", pybind11::module_local()) + .value("BYTEMAP", NVTE_ROUTING_MAP_FORMAT_BYTEMAP) + .value("BITMAP_U8", NVTE_ROUTING_MAP_FORMAT_BITMAP_U8); + m.def("fused_topk_with_score_function_fwd", &fused_topk_with_score_function_fwd, + py::arg("logits"), py::arg("topk"), py::arg("use_pre_softmax"), py::arg("num_groups"), + py::arg("group_topk"), py::arg("scaling_factor"), py::arg("score_function"), + py::arg("expert_bias"), + py::arg("routing_map_format") = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP), + "Fused topk with score function fwd"); + m.def("fused_topk_with_score_function_bwd", &fused_topk_with_score_function_bwd, + py::arg("routing_map"), py::arg("intermediate_output"), py::arg("grad_probs"), + py::arg("grad_logits"), py::arg("topk"), py::arg("use_pre_softmax"), + py::arg("scaling_factor"), py::arg("score_function"), + py::arg("routing_map_format") = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP), + "Fused topk with score function bwd"); + m.def("fused_score_for_moe_aux_loss_fwd", &fused_score_for_moe_aux_loss_fwd, py::arg("logits"), + py::arg("topk"), py::arg("score_function"), + py::arg("routing_map_format") = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP), + "Fused aux loss with score function fwd"); + m.def("fused_score_for_moe_aux_loss_bwd", &fused_score_for_moe_aux_loss_bwd, + py::arg("intermediate_output"), py::arg("grad_scores"), py::arg("grad_logits"), + py::arg("topk"), py::arg("score_function"), "Fused aux loss with score function bwd"); + 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"); + 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"); +} + +void bind_quantize_with_amax_extensions(py::module_ &m) { + m.def("nvfp4_quantize_with_amax", nvfp4_quantize_with_amax, py::arg("tensor"), + py::arg("quantizer"), py::arg("rowwise_amax"), py::arg("columnwise_amax")); + m.def("nvfp4_group_quantize_with_amax", nvfp4_group_quantize_with_amax, py::arg("tensor"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("rowwise_amax"), py::arg("columnwise_amax"), + py::arg("tensor_offsets") = py::none()); +} + } // namespace transformer_engine::pytorch #include "common/util/pybind_helper.h" PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m) + + // Register __eq__/__ne__ on the pybind ``DType`` enum so it compares by integer. + py::object dtype_class = m.attr("DType"); + dtype_class.attr("__eq__") = py::cpp_function( + [](transformer_engine::DType self, py::object other) -> py::object { + return py::cast(static_cast(self) == other.cast()); + }, + py::is_method(dtype_class)); + dtype_class.attr("__ne__") = py::cpp_function( + [](transformer_engine::DType self, py::object other) -> py::object { + return py::cast(static_cast(self) != other.cast()); + }, + py::is_method(dtype_class)); + m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("output") = py::none(), py::arg("noop") = py::none()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), @@ -144,11 +202,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Create an empty quantized tensor", py::arg("quantizer"), py::arg("shape"), py::arg("dtype"), py::arg("device"), py::arg("pin_memory")); m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"), - py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims")); + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets") = py::none()); + transformer_engine::pytorch::bind_quantize_with_amax_extensions(m); m.def("group_dequantize", transformer_engine::pytorch::group_dequantize, "Dequantize group tensor", py::arg("input"), py::arg("otype")); m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize, - py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims")); + py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets") = py::none()); m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize, "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer")); m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)", @@ -454,31 +515,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Fused Apply QKV RoPE BWD", py::call_guard()); // fused router - m.def("fused_topk_with_score_function_fwd", - &transformer_engine::pytorch::fused_topk_with_score_function_fwd, py::arg("logits"), - py::arg("topk"), py::arg("use_pre_softmax"), py::arg("num_groups"), py::arg("group_topk"), - py::arg("scaling_factor"), py::arg("score_function"), py::arg("expert_bias"), - "Fused topk with score function fwd"); - m.def("fused_topk_with_score_function_bwd", - &transformer_engine::pytorch::fused_topk_with_score_function_bwd, py::arg("num_tokens"), - py::arg("num_experts"), py::arg("routing_map"), py::arg("intermediate_output"), - py::arg("grad_probs"), py::arg("grad_logits"), py::arg("topk"), py::arg("use_pre_softmax"), - py::arg("scaling_factor"), py::arg("score_function"), "Fused topk with score function bwd"); - m.def("fused_score_for_moe_aux_loss_fwd", - &transformer_engine::pytorch::fused_score_for_moe_aux_loss_fwd, py::arg("logits"), - py::arg("topk"), py::arg("score_function"), "Fused aux loss with score function fwd"); - m.def("fused_score_for_moe_aux_loss_bwd", - &transformer_engine::pytorch::fused_score_for_moe_aux_loss_bwd, py::arg("num_tokens"), - py::arg("num_experts"), py::arg("intermediate_output"), py::arg("grad_scores"), - py::arg("grad_logits"), py::arg("topk"), py::arg("score_function"), - "Fused aux loss with score function bwd"); - m.def("fused_moe_aux_loss_fwd", &transformer_engine::pytorch::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"); - m.def("fused_moe_aux_loss_bwd", &transformer_engine::pytorch::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"); + transformer_engine::pytorch::init_router_bindings(m); // Dropout m.def("dropout_fwd", transformer_engine::pytorch::dropout_fwd, "Dropout forward with 8-bit RNG", @@ -494,13 +531,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); m.def("copy_data_ptrs_to_device", &transformer_engine::pytorch::copy_data_ptrs_to_device, py::arg("tensors"), py::arg("device"), py::call_guard()); - m.def("transform_and_copy_data_ptrs_to_device", - &transformer_engine::pytorch::transform_and_copy_data_ptrs_to_device, - py::arg("transform_type"), py::arg("tensors"), py::arg("device"), - py::call_guard()); m.def("splits_to_offsets", &transformer_engine::pytorch::splits_to_offsets, "Compute grouped tensor offsets from split sizes", py::arg("first_dims"), py::arg("logical_last_dim"), py::call_guard()); + m.def("splits_to_offsets_multi", &transformer_engine::pytorch::splits_to_offsets_multi, + "Compute multiple scaled inclusive-scan offsets from a split-sizes vector", + py::arg("split_sizes"), py::arg("device"), py::kw_only(), py::arg("strides"), + py::arg("include_leading_zero"), py::arg("dtypes"), py::arg("bulk_allocate") = false); m.def("get_num_cublas_streams", &nvte_get_num_compute_streams, "Get number of compute streams", py::call_guard()); @@ -607,6 +644,17 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard(), py::arg("allgather_communicator"), py::arg("send_stream"), py::arg("recv_stream")); + // Experimental fused grouped MLP + auto grouped_mlp_experimental = m.def_submodule( + "grouped_mlp_experimental", + "Experimental helpers for the fused grouped MLP (unstable, may change or disappear)."); + grouped_mlp_experimental.def("swizzle_scales_and_pack_ptrs_for_discrete_weights", + &transformer_engine::pytorch::grouped_mlp_experimental:: + swizzle_scales_and_pack_ptrs_for_discrete_weights, + py::arg("data_tensors"), py::arg("scale_tensors"), + py::arg("swizzle_type"), py::arg("device"), + py::call_guard()); + // Data structures py::class_(m, "FP8TensorMeta") .def(py::init<>()) @@ -641,11 +689,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::class_, transformer_engine::CommOverlapBase, transformer_engine::CommOverlapCore>(m, "CommOverlap") - .def(py::init &, at::ScalarType, CommOverlapHelper *, int, int, int, - int, int, int, int, bool, bool, bool>(), + .def(py::init([](const std::vector &buffer_shape, at::ScalarType buffer_dtype, + CommOverlapHelper *helper, int tp_size, bool use_cublasmp, + transformer_engine::CommOverlapType comm_type, int num_splits, + int num_max_streams, int comm_cga_size, int gemm_priority, int comm_priority, + int num_comm_sm, bool set_sm_margin, bool atomic_gemm, + bool rs_overlap_first_gemm) { + if (use_cublasmp) { + return std::make_shared(helper, helper->mylocal, tp_size, comm_type, + buffer_shape, buffer_dtype, num_comm_sm, + atomic_gemm); + } + return std::make_shared( + buffer_shape, buffer_dtype, helper, tp_size, num_splits, num_max_streams, + comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin, + atomic_gemm, rs_overlap_first_gemm); + }), py::call_guard(), py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), - py::arg("num_splits") = 3, py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, + py::arg("use_cublasmp") = false, + py::arg("comm_type") = transformer_engine::CommOverlapType::RS, + py::arg("num_splits") = 4, py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, py::arg("comm_cga_size") = 2, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 16, py::arg("set_sm_margin") = true, py::arg("atomic_gemm") = false, py::arg("rs_overlap_first_gemm") = false) @@ -660,15 +724,28 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::class_, transformer_engine::CommOverlapP2PBase, transformer_engine::CommOverlapCore>( m, "CommOverlapP2P") - .def(py::init &, at::ScalarType, CommOverlapHelper *, int, - transformer_engine::CommOverlapType, int, int, int, int, int, bool, bool, bool, - bool>(), + .def(py::init([](const std::vector &buffer_shape, at::ScalarType buffer_dtype, + CommOverlapHelper *helper, int tp_size, + transformer_engine::CommOverlapType comm_type, int num_max_streams, + int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, + bool set_sm_margin, bool atomic_gemm, bool use_ce, bool aggregate, + bool use_cublasmp) { + if (use_cublasmp) { + return std::make_shared(helper, helper->mylocal, tp_size, comm_type, + buffer_shape, buffer_dtype, num_comm_sm, + atomic_gemm); + } + return std::make_shared(buffer_shape, buffer_dtype, helper, tp_size, + comm_type, num_max_streams, comm_cga_size, + gemm_priority, comm_priority, num_comm_sm, + set_sm_margin, atomic_gemm, use_ce, aggregate); + }), py::call_guard(), py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), py::arg("comm_type"), py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, py::arg("comm_cga_size") = 1, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 1, py::arg("set_sm_margin") = false, py::arg("atomic_gemm") = false, - py::arg("use_ce") = true, py::arg("aggregate") = false) + py::arg("use_ce") = true, py::arg("aggregate") = false, py::arg("use_cublasmp") = false) .def("copy_into_buffer", static_cast( &CommOverlapP2P::copy_into_buffer), @@ -676,4 +753,4 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("get_buffer", &CommOverlapP2P::get_buffer, py::arg("local_chunk") = false, py::arg("shape") = std::nullopt) .def("get_communication_stream", &CommOverlapP2P::get_communication_stream); -} +} // NOLINT(readability/fn_size) diff --git a/transformer_engine/pytorch/csrc/extensions/pybind_dtype_caster.h b/transformer_engine/pytorch/csrc/extensions/pybind_dtype_caster.h new file mode 100644 index 00000000000..3323c7e6c4a --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/pybind_dtype_caster.h @@ -0,0 +1,79 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_PYBIND_DTYPE_CASTERS_H_ +#define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_PYBIND_DTYPE_CASTERS_H_ + +#include +#include + +#include +#include + +namespace pybind11 { +namespace detail { + +/*! @brief Custom type caster for ``transformer_engine::DType``. + * + * pybind exposes ``transformer_engine::DType`` as the ``transformer_engine_torch.DType`` + * enum. C++ functions take ``transformer_engine::DType`` arguments, and from Python a + * caller may pass any of the following, all of which carry the same integer dtype tag: + * - a ``transformer_engine.pytorch.DType`` (an ``IntEnum``; the canonical Python type), or + * - a plain ``int`` (the integer dtype tag), or + * - a ``transformer_engine_torch.DType`` (the pybind enum) --> Deprecated and supported for + * backward compatibility only. + * + * NOTE: As a compile-time specialization this must be visible in every translation unit + * that converts ``transformer_engine::DType`` (it is pulled in via the PyTorch extension's + * ``common.h``). Otherwise different TUs would instantiate different casters for the same + * type, which is an ODR violation. + */ +template <> +struct type_caster { + public: + PYBIND11_TYPE_CASTER(transformer_engine::DType, const_name("DType")); + + bool load(handle src, bool convert) { + if (!src) { + return false; + } + + // Intended for ``transformer_engine.pytorch.DType``, + // the canonical Python type. Cast it to the C++ enum value. + if (PyLong_Check(src.ptr())) { + const long tag = PyLong_AsLong(src.ptr()); // NOLINT(runtime/int) + if (tag == -1 && PyErr_Occurred()) { + PyErr_Clear(); + return false; + } + if (tag < 0 || + static_cast(tag) >= static_cast(transformer_engine::DType::kNumTypes)) { + return false; + } + value = static_cast(tag); + return true; + } + + // Actual ``transformer_engine_torch.DType`` enum instances (and anything else): defer to + // the standard caster, which reads the value out of the instance with no Python round-trip. + type_caster_base base; + if (base.load(src, convert)) { + value = *static_cast(base); + return true; + } + + return false; + } + + static handle cast(transformer_engine::DType src, return_value_policy policy, handle parent) { + return type_caster_base::cast(src, policy, parent); + } +}; + +} // namespace detail +} // namespace pybind11 + +#endif // TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_PYBIND_DTYPE_CASTERS_H_ diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 4df64d8e26e..5dc5c7fe86e 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -4,6 +4,8 @@ * See LICENSE for license information. ************************************************************************/ +#include + #include "../extensions.h" #include "common.h" @@ -12,13 +14,30 @@ namespace transformer_engine::pytorch { static std::map score_function_map = { {"sigmoid", 0}, {"softmax", 1}, {"sqrtsoftplus", 2}}; +// Allocate a routing_map output tensor: +// BYTEMAP -> bool [*leading_dims, num_experts] +// BITMAP_U8 -> uint8[*leading_dims, ceil(num_experts/8)], LSB-first +static at::Tensor allocate_routing_map(c10::IntArrayRef leading_dims, int64_t num_experts, + int routing_map_format) { + std::vector shape(leading_dims.begin(), leading_dims.end()); + if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + shape.push_back((num_experts + 7) / 8); + return at::empty(shape, at::dtype(at::kByte).device(at::kCUDA)); + } + shape.push_back(num_experts); + return at::empty(shape, at::dtype(at::kBool).device(at::kCUDA)); +} + std::tuple fused_topk_with_score_function_fwd( at::Tensor logits, int topk, bool use_pre_softmax, std::optional num_groups, std::optional group_topk, std::optional scaling_factor, std::string score_function, - std::optional expert_bias) { - int num_tokens = logits.size(0); - int num_experts = logits.size(1); - // Check if the input is valid + std::optional expert_bias, int routing_map_format) { + TORCH_CHECK(logits.dim() >= 1, "logits must have at least 1 dim"); + TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous"); + auto sizes = logits.sizes(); + int64_t num_experts = sizes.back(); + int64_t num_tokens = + std::accumulate(sizes.begin(), sizes.end() - 1, int64_t{1}, std::multiplies()); TORCH_CHECK(num_tokens > 0 && num_experts > 0, "num_tokens and num_experts must be greater than 0"); // Expert bias only happens at the sigmoid case @@ -41,100 +60,157 @@ std::tuple fused_topk_with_score_function_fw int num_groups_value = num_groups.has_value() ? num_groups.value() : -1; float scaling_factor_value = scaling_factor.has_value() ? scaling_factor.value() : 1.0f; - // Construct the output tensor - at::Tensor probs = - at::empty({num_tokens, num_experts}, at::dtype(logits.scalar_type()).device(at::kCUDA)); + at::Tensor probs = at::empty(sizes, at::dtype(logits.scalar_type()).device(at::kCUDA)); at::Tensor routing_map = - at::empty({num_tokens, num_experts}, at::dtype(at::kBool).device(at::kCUDA)); - // Intermediate output is used to store the output of the softmax/sigmoid function - at::Tensor intermediate_output = - at::empty({num_tokens, num_experts}, at::dtype(at::kFloat).device(at::kCUDA)); + allocate_routing_map(sizes.slice(0, sizes.size() - 1), num_experts, routing_map_format); + at::Tensor intermediate_output = at::empty(sizes, at::dtype(at::kFloat).device(at::kCUDA)); - auto logits_cu = makeTransformerEngineTensor(logits); - auto probs_cu = makeTransformerEngineTensor(probs); - auto routing_map_cu = makeTransformerEngineTensor(routing_map); - auto intermediate_output_cu = makeTransformerEngineTensor(intermediate_output); + // 2D shape for the kernel (common-layer NVTE_CHECKs require {num_tokens, trailing_dim}). + const std::vector shape_2d = {static_cast(num_tokens), + static_cast(num_experts)}; + const std::vector routing_map_shape_2d = { + static_cast(num_tokens), + static_cast(routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 + ? (num_experts + 7) / 8 + : num_experts)}; + auto logits_dtype = GetTransformerEngineDType(logits.scalar_type()); + auto routing_map_dtype = GetTransformerEngineDType(routing_map.scalar_type()); + + auto logits_cu = makeTransformerEngineTensor(logits.data_ptr(), shape_2d, logits_dtype); + auto probs_cu = makeTransformerEngineTensor(probs.data_ptr(), shape_2d, logits_dtype); + auto routing_map_cu = + makeTransformerEngineTensor(routing_map.data_ptr(), routing_map_shape_2d, routing_map_dtype); + auto intermediate_output_cu = + makeTransformerEngineTensor(intermediate_output.data_ptr(), shape_2d, DType::kFloat32); auto expert_bias_cu = TensorWrapper(); // empty expert_bias_cu tensor if (expert_bias.has_value()) { expert_bias_cu = makeTransformerEngineTensor(expert_bias.value()); } - nvte_fused_topk_with_score_function_forward( - logits_cu.data(), num_tokens, num_experts, topk, use_pre_softmax, num_groups_value, - group_topk_value, scaling_factor_value, score_function_map[score_function], - expert_bias_cu.data(), probs_cu.data(), routing_map_cu.data(), intermediate_output_cu.data(), - at::cuda::getCurrentCUDAStream()); + nvte_fused_topk_with_score_function_forward_v2( + logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, + use_pre_softmax, num_groups_value, group_topk_value, scaling_factor_value, + score_function_map[score_function], expert_bias_cu.data(), probs_cu.data(), + routing_map_cu.data(), static_cast(routing_map_format), + intermediate_output_cu.data(), at::cuda::getCurrentCUDAStream()); return std::make_tuple(probs, routing_map, intermediate_output); } -void fused_topk_with_score_function_bwd(int num_tokens, int num_experts, at::Tensor routing_map, - at::Tensor intermediate_output, at::Tensor grad_probs, - at::Tensor grad_logits, int topk, bool use_pre_softmax, - std::optional scaling_factor, - std::string score_function) { - // Get the value of the parameters +void fused_topk_with_score_function_bwd(at::Tensor routing_map, at::Tensor intermediate_output, + at::Tensor grad_probs, at::Tensor grad_logits, int topk, + bool use_pre_softmax, std::optional scaling_factor, + std::string score_function, int routing_map_format) { + TORCH_CHECK(grad_probs.dim() >= 1, "grad_probs must have at least 1 dim"); + TORCH_CHECK(grad_probs.is_contiguous(), "grad_probs must be contiguous"); + TORCH_CHECK(grad_logits.is_contiguous(), "grad_logits must be contiguous"); + auto sizes = grad_probs.sizes(); + int64_t num_experts = sizes.back(); + int64_t num_tokens = + std::accumulate(sizes.begin(), sizes.end() - 1, int64_t{1}, std::multiplies()); + auto scaling_factor_value = scaling_factor.has_value() ? scaling_factor.value() : 1.0f; auto score_function_value = score_function_map[score_function]; - auto routing_map_cu = makeTransformerEngineTensor(routing_map); - auto intermediate_output_cu = makeTransformerEngineTensor(intermediate_output); - auto grad_probs_cu = makeTransformerEngineTensor(grad_probs); - auto grad_logits_cu = makeTransformerEngineTensor(grad_logits); + const std::vector shape_2d = {static_cast(num_tokens), + static_cast(num_experts)}; + const std::vector routing_map_shape_2d = { + static_cast(num_tokens), + static_cast(routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 + ? (num_experts + 7) / 8 + : num_experts)}; + auto grad_dtype = GetTransformerEngineDType(grad_probs.scalar_type()); + auto routing_map_dtype = GetTransformerEngineDType(routing_map.scalar_type()); - nvte_fused_topk_with_score_function_backward( - routing_map_cu.data(), intermediate_output_cu.data(), grad_probs_cu.data(), num_tokens, - num_experts, topk, use_pre_softmax, scaling_factor_value, score_function_value, - grad_logits_cu.data(), at::cuda::getCurrentCUDAStream()); + auto routing_map_cu = + makeTransformerEngineTensor(routing_map.data_ptr(), routing_map_shape_2d, routing_map_dtype); + auto intermediate_output_cu = + makeTransformerEngineTensor(intermediate_output.data_ptr(), shape_2d, DType::kFloat32); + auto grad_probs_cu = makeTransformerEngineTensor(grad_probs.data_ptr(), shape_2d, grad_dtype); + auto grad_logits_cu = makeTransformerEngineTensor(grad_logits.data_ptr(), shape_2d, grad_dtype); + + nvte_fused_topk_with_score_function_backward_v2( + routing_map_cu.data(), static_cast(routing_map_format), + intermediate_output_cu.data(), grad_probs_cu.data(), static_cast(num_tokens), + static_cast(num_experts), topk, use_pre_softmax, scaling_factor_value, + score_function_value, grad_logits_cu.data(), at::cuda::getCurrentCUDAStream()); } std::tuple fused_score_for_moe_aux_loss_fwd( - at::Tensor logits, int topk, std::string score_function) { - int num_tokens = logits.size(0); - int num_experts = logits.size(1); - // Check if the input is valid + at::Tensor logits, int topk, std::string score_function, int routing_map_format) { + TORCH_CHECK(logits.dim() >= 1, "logits must have at least 1 dim"); + TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous"); + auto sizes = logits.sizes(); + int64_t num_experts = sizes.back(); + int64_t num_tokens = + std::accumulate(sizes.begin(), sizes.end() - 1, int64_t{1}, std::multiplies()); TORCH_CHECK(num_tokens > 0 && num_experts > 0, "num_tokens and num_experts must be greater than 0"); TORCH_CHECK(topk > 0, "topk must be greater than 0"); - // Check if the score function is valid TORCH_CHECK(score_function == "softmax" || score_function == "sigmoid" || score_function == "sqrtsoftplus", "score_function must be softmax, sigmoid or sqrtsoftplus for router fusion"); int score_function_value = score_function_map[score_function]; - // Construct the output tensor - at::Tensor scores = at::empty({num_tokens, num_experts}, at::dtype(at::kFloat).device(at::kCUDA)); + at::Tensor scores = at::empty(sizes, at::dtype(at::kFloat).device(at::kCUDA)); at::Tensor routing_map = - at::empty({num_tokens, num_experts}, at::dtype(at::kBool).device(at::kCUDA)); - at::Tensor intermediate_output = - at::empty({num_tokens, num_experts}, at::dtype(at::kFloat).device(at::kCUDA)); + allocate_routing_map(sizes.slice(0, sizes.size() - 1), num_experts, routing_map_format); + at::Tensor intermediate_output = at::empty(sizes, at::dtype(at::kFloat).device(at::kCUDA)); - auto logits_cu = makeTransformerEngineTensor(logits); - auto scores_cu = makeTransformerEngineTensor(scores); - auto routing_map_cu = makeTransformerEngineTensor(routing_map); - auto intermediate_output_cu = makeTransformerEngineTensor(intermediate_output); + const std::vector shape_2d = {static_cast(num_tokens), + static_cast(num_experts)}; + const std::vector routing_map_shape_2d = { + static_cast(num_tokens), + static_cast(routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 + ? (num_experts + 7) / 8 + : num_experts)}; + auto logits_dtype = GetTransformerEngineDType(logits.scalar_type()); + auto routing_map_dtype = GetTransformerEngineDType(routing_map.scalar_type()); - nvte_fused_score_for_moe_aux_loss_forward( - logits_cu.data(), num_tokens, num_experts, topk, score_function_value, scores_cu.data(), - routing_map_cu.data(), intermediate_output_cu.data(), at::cuda::getCurrentCUDAStream()); + auto logits_cu = makeTransformerEngineTensor(logits.data_ptr(), shape_2d, logits_dtype); + auto scores_cu = makeTransformerEngineTensor(scores.data_ptr(), shape_2d, DType::kFloat32); + auto routing_map_cu = + makeTransformerEngineTensor(routing_map.data_ptr(), routing_map_shape_2d, routing_map_dtype); + auto intermediate_output_cu = + makeTransformerEngineTensor(intermediate_output.data_ptr(), shape_2d, DType::kFloat32); + + nvte_fused_score_for_moe_aux_loss_forward_v2( + logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, + score_function_value, scores_cu.data(), routing_map_cu.data(), + static_cast(routing_map_format), intermediate_output_cu.data(), + at::cuda::getCurrentCUDAStream()); return std::make_tuple(scores, routing_map, intermediate_output); } -void fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts, - at::Tensor intermediate_output, at::Tensor grad_scores, +void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor grad_scores, at::Tensor grad_logits, int topk, std::string score_function) { - // Get the value of the parameters + TORCH_CHECK(grad_scores.dim() >= 1, "grad_scores must have at least 1 dim"); + TORCH_CHECK(grad_scores.is_contiguous(), "grad_scores must be contiguous"); + TORCH_CHECK(grad_logits.is_contiguous(), "grad_logits must be contiguous"); + auto sizes = grad_scores.sizes(); + int64_t num_experts = sizes.back(); + int64_t num_tokens = + std::accumulate(sizes.begin(), sizes.end() - 1, int64_t{1}, std::multiplies()); + int score_function_value = score_function_map[score_function]; - auto intermediate_output_cu = makeTransformerEngineTensor(intermediate_output); - auto grad_scores_cu = makeTransformerEngineTensor(grad_scores); - auto grad_logits_cu = makeTransformerEngineTensor(grad_logits); + const std::vector shape_2d = {static_cast(num_tokens), + static_cast(num_experts)}; + auto grad_logits_dtype = GetTransformerEngineDType(grad_logits.scalar_type()); + + auto intermediate_output_cu = + makeTransformerEngineTensor(intermediate_output.data_ptr(), shape_2d, DType::kFloat32); + auto grad_scores_cu = + makeTransformerEngineTensor(grad_scores.data_ptr(), shape_2d, DType::kFloat32); + auto grad_logits_cu = + makeTransformerEngineTensor(grad_logits.data_ptr(), shape_2d, grad_logits_dtype); nvte_fused_score_for_moe_aux_loss_backward( - intermediate_output_cu.data(), grad_scores_cu.data(), num_tokens, num_experts, topk, - score_function_value, grad_logits_cu.data(), at::cuda::getCurrentCUDAStream()); + intermediate_output_cu.data(), grad_scores_cu.data(), static_cast(num_tokens), + static_cast(num_experts), topk, score_function_value, grad_logits_cu.data(), + at::cuda::getCurrentCUDAStream()); } std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, diff --git a/transformer_engine/pytorch/csrc/extensions/utils.cpp b/transformer_engine/pytorch/csrc/extensions/utils.cpp deleted file mode 100644 index 453f238c0d5..00000000000 --- a/transformer_engine/pytorch/csrc/extensions/utils.cpp +++ /dev/null @@ -1,177 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include - -#include -#include -#include -#include -#include - -#include "common/common.h" -#include "extensions.h" - -namespace transformer_engine::pytorch { - -at::Tensor copy_data_ptrs_to_device(const std::vector &tensors, - const c10::Device &device) { - // Collect data pointers - std::vector ptrs_host; - ptrs_host.reserve(tensors.size()); - for (const auto &tensor : tensors) { - ptrs_host.push_back(reinterpret_cast(tensor.data_ptr())); - } - - // Allocate device buffer - auto ptrs_device = at::empty({static_cast(tensors.size())}, - at::TensorOptions().dtype(at::kLong).device(device)); - - // Load pointers on device - nvte_copy_host_to_device_via_kernel(ptrs_host.data(), ptrs_device.data_ptr(), - tensors.size() * sizeof(uint64_t), - at::cuda::getCurrentCUDAStream()); - - return ptrs_device; -} - -std::tuple> transform_and_copy_data_ptrs_to_device( - const std::string &transform_type, const std::vector &tensors, - const c10::Device &device) { - const size_t num_tensors = tensors.size(); - - // Trivial cases - if (transform_type.empty()) { - // No transform, just load pointers on device - return {copy_data_ptrs_to_device(tensors, device), std::nullopt}; - } - if (num_tensors == 0) { - // No input tensors, return tensor with no elements - return {at::empty({int64_t{0}}, at::TensorOptions().dtype(at::kLong).device(device)), - std::nullopt}; - } - - // CUDA stream - auto stream = at::cuda::getCurrentCUDAStream(); - - // Swizzle scales for GEMM, with uniform tensor sizes - const bool uniform_mxfp8_rowwise_swizzle = transform_type == "uniform_mxfp8_rowwise_swizzle"; - const bool uniform_mxfp8_colwise_swizzle = transform_type == "uniform_mxfp8_columnwise_swizzle"; - const bool uniform_nvfp4_swizzle = transform_type == "uniform_nvfp4_swizzle"; - if (uniform_mxfp8_rowwise_swizzle || uniform_mxfp8_colwise_swizzle || uniform_nvfp4_swizzle) { - // Tensor format - NVTEScalingMode scaling_mode = NVTE_INVALID_SCALING; - if (uniform_mxfp8_rowwise_swizzle || uniform_mxfp8_colwise_swizzle) { - scaling_mode = NVTE_MXFP8_1D_SCALING; - } else if (uniform_nvfp4_swizzle) { - scaling_mode = NVTE_NVFP4_1D_SCALING; - } - - // Data types - transformer_engine::DType data_dtype, scale_dtype; - switch (scaling_mode) { - case NVTE_MXFP8_1D_SCALING: - data_dtype = transformer_engine::DType::kFloat8E4M3; - scale_dtype = transformer_engine::DType::kFloat8E8M0; - break; - case NVTE_NVFP4_1D_SCALING: - data_dtype = transformer_engine::DType::kFloat4E2M1; - scale_dtype = transformer_engine::DType::kFloat8E4M3; - break; - default: - NVTE_ERROR("Unsupported case."); - } - - // Scale shape - const NVTEShape scale_shape = convertTorchShape(tensors[0].sizes()); - NVTE_CHECK(scale_shape.ndim == 2, - "Expected 2D scale tensor, but got shape=", getTensorShape(tensors[0]), "."); - const size_t scale_numel = scale_shape.data[0] * scale_shape.data[1]; - const size_t scale_dtype_bits = transformer_engine::pytorch::typeToNumBits(scale_dtype); - const size_t scale_bytes = ceildiv(scale_numel * scale_dtype_bits, 8); - - // Expected data shape - // Note: May not match actual data shape since the scales are padded. - // This is fine since we're not actually touching the data. - NVTEShape data_shape; - data_shape.ndim = 2; - if (uniform_mxfp8_rowwise_swizzle) { - data_shape.data[0] = scale_shape.data[0]; - data_shape.data[1] = scale_shape.data[1] * 32; - } else if (uniform_mxfp8_colwise_swizzle) { - data_shape.data[0] = scale_shape.data[0] * 32; - data_shape.data[1] = scale_shape.data[1]; - } else if (uniform_nvfp4_swizzle) { - data_shape.data[0] = scale_shape.data[0]; - data_shape.data[1] = scale_shape.data[1] * 16; - } else { - NVTE_ERROR("Unsupported case."); - } - - // Allocate single buffer for swizzled scales. - // Uses a uniform stride since all tensors share the same scale shape. - const size_t swizzled_scales_stride = roundup(scale_bytes, 16); // Align to 16 bytes - auto swizzled_scales = at::empty({static_cast(swizzled_scales_stride * num_tensors)}, - at::TensorOptions().dtype(at::kByte).device(device)); - uint8_t *swizzled_scales_dptr = reinterpret_cast(swizzled_scales.data_ptr()); - - // Allocate input/output NVTETensors as a single batch. The first - // num_tensors entries are inputs; the next num_tensors are outputs. - MultiTensorWrapper nvte_tensors(2 * num_tensors, scaling_mode); - NVTETensor *inputs_nvte = nvte_tensors.data(); - NVTETensor *outputs_nvte = nvte_tensors.data() + num_tensors; - - auto set_param = [](NVTETensor t, NVTETensorParam param, void *dptr, - transformer_engine::DType dtype, const NVTEShape &shape) { - NVTEBasicTensor data{dptr, static_cast(dtype), shape}; - nvte_set_tensor_param_v2(t, param, &data, sizeof(data)); - }; - - for (size_t i = 0; i < num_tensors; ++i) { - const uint8_t swizzled_flag = 1; - nvte_set_tensor_param_v2(outputs_nvte[i], kNVTEWithGEMMSwizzledScales, &swizzled_flag, - sizeof(swizzled_flag)); - void *in_scale_ptr = tensors[i].data_ptr(); - void *out_scale_ptr = swizzled_scales_dptr + i * swizzled_scales_stride; - if (uniform_mxfp8_rowwise_swizzle || uniform_nvfp4_swizzle) { - set_param(inputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_shape); - set_param(inputs_nvte[i], kNVTERowwiseScaleInv, in_scale_ptr, scale_dtype, scale_shape); - set_param(outputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_shape); - set_param(outputs_nvte[i], kNVTERowwiseScaleInv, out_scale_ptr, scale_dtype, scale_shape); - } else if (uniform_mxfp8_colwise_swizzle) { - set_param(inputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_shape); - set_param(inputs_nvte[i], kNVTEColumnwiseScaleInv, in_scale_ptr, scale_dtype, scale_shape); - set_param(outputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_shape); - set_param(outputs_nvte[i], kNVTEColumnwiseScaleInv, out_scale_ptr, scale_dtype, - scale_shape); - } - } - - // Launch kernel - nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte, outputs_nvte, num_tensors, stream); - - // Collect data pointers - std::vector ptrs_host; - ptrs_host.reserve(num_tensors); - for (size_t i = 0; i < num_tensors; ++i) { - ptrs_host.push_back( - reinterpret_cast(swizzled_scales_dptr + i * swizzled_scales_stride)); - } - - // Load pointers on device - auto ptrs_device = at::empty({static_cast(num_tensors)}, - at::TensorOptions().dtype(at::kLong).device(device)); - nvte_copy_host_to_device_via_kernel(ptrs_host.data(), ptrs_device.data_ptr(), - num_tensors * sizeof(uint64_t), stream); - - return {std::move(ptrs_device), std::move(swizzled_scales)}; - } - - // Unsupported transform - NVTE_ERROR("Unsupported transform type (", transform_type, ")"); -} - -} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index bc87b54ba84..ec51a983d79 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -174,11 +174,14 @@ std::pair NoneQuantizer::create_tensor(const std::vec std::pair NoneQuantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, - const size_t logical_first_dim, const size_t logical_last_dim) const { + const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, + const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + precomputed_tensor_offsets.has_value() + ? precomputed_tensor_offsets + : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -319,7 +322,7 @@ std::pair Float8Quantizer::create_tensor( py::tuple args(0); kwargs["data"] = data_py; kwargs["fp8_scale_inv"] = scale_inv_py; - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["data_transpose"] = transpose_py; kwargs["quantizer"] = this->quantizer; kwargs["fake_dtype"] = GetATenDType(dtype); @@ -343,7 +346,7 @@ std::pair Float8Quantizer::create_tensor( kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["data"] = data_py; kwargs["fp8_scale_inv"] = scale_inv_py; - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["data_transpose"] = transpose_py; kwargs["quantizer"] = this->quantizer; kwargs["device"] = py::cast(device); @@ -377,11 +380,14 @@ std::pair Float8Quantizer::create_tensor( std::pair Float8Quantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, - const size_t logical_first_dim, const size_t logical_last_dim) const { + const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, + const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + precomputed_tensor_offsets.has_value() + ? precomputed_tensor_offsets + : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -527,7 +533,7 @@ std::pair Float8Quantizer::convert_and_update_tensor( tensor.attr("_transpose_invalid") = !need_transpose; // Coerce other attrs - tensor.attr("_fp8_dtype") = dtype; + tensor.attr("_fp8_dtype") = MakePythonDType(dtype); // Construct C++ FP8 tensor TensorWrapper out_cpp; @@ -628,7 +634,7 @@ std::pair Float8CurrentScalingQuantizer::create_tenso py::dict kwargs; kwargs["data"] = data_py; kwargs["fp8_scale_inv"] = scale_inv_py; - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["data_transpose"] = transpose_py; kwargs["quantizer"] = this->quantizer; kwargs["fake_dtype"] = GetATenDType(dtype); @@ -651,7 +657,7 @@ std::pair Float8CurrentScalingQuantizer::create_tenso kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["data"] = data_py; kwargs["fp8_scale_inv"] = scale_inv_py; - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["data_transpose"] = transpose_py; kwargs["quantizer"] = this->quantizer; kwargs["device"] = py::cast(device); @@ -687,11 +693,14 @@ std::pair Float8CurrentScalingQuantizer::create_tenso std::pair Float8CurrentScalingQuantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, - const size_t logical_first_dim, const size_t logical_last_dim) const { + const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, + const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + precomputed_tensor_offsets.has_value() + ? precomputed_tensor_offsets + : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -854,7 +863,7 @@ std::pair Float8CurrentScalingQuantizer::convert_and_ tensor.attr("_transpose_invalid") = !need_transpose; // Coerce other attrs - tensor.attr("_fp8_dtype") = dtype; + tensor.attr("_fp8_dtype") = MakePythonDType(dtype); // Construct C++ FP8 tensor TensorWrapper out_cpp; @@ -1019,7 +1028,7 @@ std::pair Float8BlockQuantizer::create_tensor( kwargs["columnwise_data"] = py::cast(data_colwise); kwargs["rowwise_scale_inv"] = py::cast(scale_inv_rowwise); kwargs["columnwise_scale_inv"] = py::cast(scale_inv_colwise); - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["is_2D_scaled"] = py::cast(block_scaling_dim == 2); kwargs["fake_dtype"] = GetATenDType(dtype); @@ -1045,7 +1054,7 @@ std::pair Float8BlockQuantizer::create_tensor( kwargs["columnwise_data"] = py::cast(data_colwise); kwargs["rowwise_scale_inv"] = py::cast(scale_inv_rowwise); kwargs["columnwise_scale_inv"] = py::cast(scale_inv_colwise); - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["is_2D_scaled"] = py::cast(block_scaling_dim == 2); kwargs["device"] = py::cast(device); @@ -1066,11 +1075,14 @@ std::pair Float8BlockQuantizer::create_tensor( std::pair Float8BlockQuantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, - const size_t logical_first_dim, const size_t logical_last_dim) const { + const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, + const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + precomputed_tensor_offsets.has_value() + ? precomputed_tensor_offsets + : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -1423,7 +1435,7 @@ std::pair MXFP8Quantizer::create_tensor( kwargs["columnwise_data"] = columnwise_data_py; kwargs["rowwise_scale_inv"] = rowwise_scale_inv_py; kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["fake_dtype"] = GetATenDType(dtype); @@ -1447,7 +1459,7 @@ std::pair MXFP8Quantizer::create_tensor( kwargs["columnwise_data"] = columnwise_data_py; kwargs["rowwise_scale_inv"] = rowwise_scale_inv_py; kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; - kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["fp8_dtype"] = MakePythonDType(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["device"] = py::cast(device); @@ -1484,11 +1496,14 @@ std::pair MXFP8Quantizer::create_tensor( std::pair MXFP8Quantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, - const size_t logical_first_dim, const size_t logical_last_dim) const { + const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, + const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + precomputed_tensor_offsets.has_value() + ? precomputed_tensor_offsets + : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -1657,7 +1672,7 @@ std::pair MXFP8Quantizer::convert_and_update_tensor( } // Coerce other attrs - tensor.attr("_fp8_dtype") = dtype; + tensor.attr("_fp8_dtype") = MakePythonDType(dtype); tensor.attr("_with_gemm_swizzled_scales") = with_gemm_swizzled_scales; // Construct C++ MXFP8 tensor @@ -1781,7 +1796,12 @@ std::pair NVFP4Quantizer::create_tensor( using namespace pybind11::literals; // Scaling factor format - const bool with_gemm_swizzled_scales = false; /// TODO (tmoon) self->optimize_for_gemm + // + // The NVFP4 quantize kernel requires `with_gemm_swizzled_scales=false` and + // emits compact scales. When `optimize_for_gemm` is set, the caller of + // `quantize` (cast.cpp) applies a post-quantize swizzle that flips this + // flag on the returned tensor. + const bool with_gemm_swizzled_scales = false; // Tensor dimensions const std::vector shape_int64(shape.begin(), shape.end()); @@ -1857,7 +1877,7 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; kwargs["amax_rowwise"] = amax_rowwise_py; kwargs["amax_columnwise"] = amax_columnwise_py; - kwargs["fp4_dtype"] = py::cast(this->dtype); + kwargs["fp4_dtype"] = MakePythonDType(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); @@ -1888,7 +1908,7 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; kwargs["amax_rowwise"] = amax_rowwise_py; kwargs["amax_columnwise"] = amax_columnwise_py; - kwargs["fp4_dtype"] = py::cast(this->dtype); + kwargs["fp4_dtype"] = MakePythonDType(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["device"] = py::cast(device); @@ -1937,11 +1957,14 @@ std::pair NVFP4Quantizer::create_tensor( std::pair NVFP4Quantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, - const size_t logical_first_dim, const size_t logical_last_dim) const { + const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, + const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + precomputed_tensor_offsets.has_value() + ? precomputed_tensor_offsets + : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); NVTE_CHECK(total_elements % 2 == 0, "NVFP4 data size must be divisible by 2."); @@ -2077,7 +2100,12 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( NVTE_CHECK(detail::IsNVFP4Tensor(tensor.ptr()), "NVFP4Quantizer must output to IsNVFP4Tensor."); // Scaling factor format - const bool with_gemm_swizzled_scales = false; // TODO (tmoon) Enable with optimize_for_gemm + // + // The NVFP4 quantize kernel requires `with_gemm_swizzled_scales=false` and + // emits compact scales. When `optimize_for_gemm` is set, the caller of + // `quantize` (cast.cpp) applies a post-quantize swizzle that flips this + // flag on the returned tensor. + const bool with_gemm_swizzled_scales = false; // Extract buffers from Python tensor auto get_tensor = [&tensor](const char* name) -> std::optional { @@ -2109,6 +2137,9 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( } const auto [flat_first_dim, flat_last_dim] = get_2d_dims(shape); + tensor.attr("_with_gemm_swizzled_scales") = py::cast(with_gemm_swizzled_scales); + + // Advanced NVFP4 modes const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; @@ -2300,8 +2331,40 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& out, const std::optional& noop_flag, bool compute_amax) { + auto reduce_amaxes = [&]() { + if (!this->with_amax_reduction) { + return; + } + + std::vector amax_tensors; + auto make_amax_tensor = [](void* data_ptr) { + NVTE_CHECK(data_ptr != nullptr, "Could not find amax pointer for NVFP4 amax reduction."); + return at::from_blob( + data_ptr, std::vector{1}, + [](void*) {}, // deleter doing nothing since it doesn't own the data + at::device(at::kCUDA).dtype(torch::kFloat32)); + }; + if (rowwise_usage) { + amax_tensors.push_back(make_amax_tensor(out.get_amax().data_ptr)); + } + if (columnwise_usage) { + amax_tensors.push_back(make_amax_tensor(out.get_columnwise_amax().data_ptr)); + } + if (amax_tensors.empty()) { + return; + } + + c10d::AllreduceCoalescedOptions opts; + opts.reduceOp = c10d::ReduceOp::MAX; + NVTE_SCOPED_GIL_RELEASE( + { this->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait(); }); + }; + // Nothing to be done if input is empty if (input.numel() == 0) { + if (!compute_amax) { + reduce_amaxes(); + } return; } @@ -2400,10 +2463,12 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // We need: // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) - NVTE_SCOPED_GIL_RELEASE({ - nvte_hadamard_transform_amax(input.data(), out.data(), 0, - this->rht_matrix_random_sign_mask_t, stream); - }); + if (compute_amax) { + NVTE_SCOPED_GIL_RELEASE({ + nvte_hadamard_transform_amax(input.data(), out.data(), 0, + this->rht_matrix_random_sign_mask_t, stream); + }); + } } else { // raise error since it's not supported yet NVTE_ERROR( @@ -2436,27 +2501,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou } } - // amax reduction - if (this->with_amax_reduction) { - std::vector amax_tensors; - // push amax tensors inside if they need to be reduced - auto make_amax_tensor = [](void* data_ptr) { - return at::from_blob( - data_ptr, std::vector{1}, - [](void*) {}, // deleter doing nothing since it doesn't own the data - at::device(at::kCUDA).dtype(torch::kFloat32)); - }; - if (rowwise_usage) { - amax_tensors.push_back(make_amax_tensor(out.get_amax().data_ptr)); - } - if (columnwise_usage) { - amax_tensors.push_back(make_amax_tensor(out.get_columnwise_amax().data_ptr)); - } - c10d::AllreduceCoalescedOptions opts; - opts.reduceOp = c10d::ReduceOp::MAX; - NVTE_SCOPED_GIL_RELEASE( - { this->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait(); }); - } + reduce_amaxes(); // Fast math toggle: RHT transform can be accelerated // What math is accelerated? Only the high precision math, so numerical impact is minimal diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py index d660e5a53bc..e88adbe4cca 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py @@ -38,9 +38,8 @@ from typing import Optional -import transformer_engine_torch as tex - from transformer_engine.pytorch.quantization import QuantizerRole +from ..constants import DType def nvfp4_linear_mxfp8_grouped_linear_factory( @@ -70,7 +69,7 @@ def _make_mxfp8_quantizer(): from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer return MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, ) @@ -87,7 +86,7 @@ def _make_nvfp4_quantizer(role: Optional[QuantizerRole]): if is_weight: return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=DType.kFloat4E2M1, with_rht=False, with_post_rht_amax=False, with_2d_quantization=True, @@ -97,7 +96,7 @@ def _make_nvfp4_quantizer(role: Optional[QuantizerRole]): if is_grad: return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=DType.kFloat4E2M1, rowwise=True, columnwise=True, with_rht=True, @@ -108,7 +107,7 @@ def _make_nvfp4_quantizer(role: Optional[QuantizerRole]): ) return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=DType.kFloat4E2M1, rowwise=True, columnwise=True, with_rht=True, @@ -180,7 +179,7 @@ def nvfp4_linear_fp8_dpa_factory( if is_dpa: return Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, device="cuda", ) @@ -193,7 +192,7 @@ def nvfp4_linear_fp8_dpa_factory( ) if is_dpa_boundary: return Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, device="cuda", ) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py b/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py index 22eafaa6650..45febaf413d 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py @@ -28,9 +28,8 @@ from typing import Optional import torch -import transformer_engine_torch as tex - from transformer_engine.pytorch.quantization import QuantizerRole +from ..constants import DType def delayed_scaling_quantizer_factory( @@ -65,7 +64,7 @@ def current_scaling_quantizer_factory( ) is_backward = role is not None and role.tensor_type == "grad_output" - fp8_dtype = tex.DType.kFloat8E5M2 if is_backward else tex.DType.kFloat8E4M3 + fp8_dtype = DType.kFloat8E5M2 if is_backward else DType.kFloat8E4M3 return Float8CurrentScalingQuantizer( fp8_dtype=fp8_dtype, @@ -86,7 +85,7 @@ def mxfp8_quantizer_factory( from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer return MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, ) @@ -111,7 +110,7 @@ def float8_block_scaling_quantizer_factory( block_scaling_dim = 2 if is_weight else 1 return Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True, amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero @@ -146,7 +145,7 @@ def nvfp4_quantizer_factory( if is_weight: return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=DType.kFloat4E2M1, with_rht=False, with_post_rht_amax=False, with_2d_quantization=True, @@ -156,7 +155,7 @@ def nvfp4_quantizer_factory( if is_grad: return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=DType.kFloat4E2M1, rowwise=True, columnwise=True, with_rht=True, @@ -168,7 +167,7 @@ def nvfp4_quantizer_factory( # For input and unknown roles return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=DType.kFloat4E2M1, rowwise=True, columnwise=True, with_rht=True, diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index a1213fe4931..6c7ba8a8ab3 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -62,17 +62,35 @@ from ...debug.pytorch.debug_quantization import DebugQuantizer, DebugQuantizedTensor from ...debug.pytorch.utils import next_iter_when_debug_should_be_run, any_feature_enabled -__all__ = ["initialize_ub", "destroy_ub", "UserBufferQuantizationMode"] +__all__ = [ + "initialize_ub", + "destroy_ub", + "is_ub_initialized", + "using_cublasmp_backend", + "UserBufferQuantizationMode", +] _2X_ACC_FPROP = False _2X_ACC_DGRAD = True _2X_ACC_WGRAD = True _dummy_wgrads = {} _ub_communicators = None +_ub_initialized = False +_ub_with_cublasmp = False _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = None, None layers_atomic_ring_exchange = [] +def is_ub_initialized() -> bool: + """Whether the Userbuffers communicators have been initialized.""" + return _ub_initialized + + +def using_cublasmp_backend() -> bool: + """Whether the active comm+GEMM overlap backend is cuBLASMp.""" + return _ub_initialized and _ub_with_cublasmp + + class UserBufferQuantizationMode(Enum): """ UserBufferQuantizationMode is an enum that represents the quantization mode of the UserBuffer. @@ -107,6 +125,7 @@ def initialize_ub( dtype: torch.dtype = torch.bfloat16, ub_cfgs: Optional[Union[dict, List[dict]]] = None, bootstrap_backend: Union[str, torch.distributed.Backend] = None, + with_cublasmp: bool = False, ) -> None: r""" Initialize the Userbuffers communicator for overlapping tensor-parallel communications with @@ -159,8 +178,15 @@ def initialize_ub( not available. Setting ``NVTE_UB_WITH_MPI=1`` when building TE overrides this option and always initializes Userbuffers with direct MPI calls in C++, which also requires ``MPI_HOME=/path/to/mpi/root`` to be set at compile time. + with_cublasmp : bool = False + Whether to use cuBlasMp for the all-gather and reduce-scatter overlaps. TE must + be compiled with `NVTE_WITH_CUBLASMP=1` for this option to work. + """ - if not tex.device_supports_multicast(): + # The cuBLASMp backend falls back to a non-multicast algorithm in C++ when multicast is + # unsupported or disabled (via UB_SKIPMC), so only enforce the multicast requirement for + # the Userbuffers backend. + if not tex.device_supports_multicast() and not with_cublasmp: if not bool(int(os.getenv("UB_SKIPMC", "0"))): raise RuntimeError( "CUDA device, driver and/or toolkit version does not support comm+GEMM overlap " @@ -199,12 +225,13 @@ def initialize_ub( f"quantization configurations ({len(quantization_modes)})" ) - global _ub_communicators + global _ub_communicators, _ub_with_cublasmp if _ub_communicators is not None: raise RuntimeError("UB communicators are already initialized.") _ub_communicators = {} + _ub_with_cublasmp = with_cublasmp - if tex.ubuf_built_with_mpi(): + if tex.ubuf_built_with_mpi() and not with_cublasmp: # We're bootstrapping with direct calls to MPI in Userbuffers code so we need to force # an MPI_Init() here by creating a new MPI process group... if not torch.distributed.is_mpi_available(): @@ -216,7 +243,10 @@ def initialize_ub( helper = tex.CommOverlapHelper() else: # Bootstrapping with torch.distributed API, so check backend and construct - # intra/inter-node process groups... + # intra/inter-node process groups. The cuBLASMp backend always takes this path + # (even when Userbuffers is built with MPI) because it derives its NCCL + # communicators from the helper's PyTorch process groups, which the dummy + # MPI-bootstrap helper does not provide. if not torch.distributed.is_initialized(): raise RuntimeError("torch.distributed must be initialized before using Userbuffers") if bootstrap_backend is None: @@ -260,7 +290,7 @@ def initialize_ub( local_rank = world_rank tp_domain_ranks = list(range(world_size)) - helper = tex.CommOverlapHelper(world_group) + helper = tex.CommOverlapHelper(world_group, world_group) if world_rank == 0: print(f"!!! [UB] Number of TP domains: {num_domains}\n", end="", flush=True) @@ -284,6 +314,7 @@ def initialize_ub( ] layers_reduce_scatter_overlap = ["proj_fprop", "fc2_fprop", "qkv_wgrad", "fc1_wgrad"] dgrad_reduce_scatter_overlap = ["qkv_dgrad", "fc1_dgrad"] + # Default overlap methods for layers methods = { "ring_exchange": [ @@ -350,6 +381,11 @@ def add_ub( gemm_priority: int = 0, pipeline_rs_overlap_first_gemm: bool = False, ) -> None: + if with_cublasmp and method in ("bulk", "external"): + raise ValueError( + f"At {name}, cuBLASMp does not support `{method}` overlap method. " + "Please select a different method or set with_cublasmp=False." + ) if atomic_gemm: warnings.warn( "Atomic GEMM uses a beta API from cublas and is not tested for all use cases." @@ -361,7 +397,7 @@ def add_ub( ) if method in ("bulk", "external"): warnings.warn( - f"At {name}, atoimic GEMM not is supported for a bulk overlap." + f"At {name}, atomic GEMM not is supported for a bulk overlap." "Defaulting to `atomic_gemm=False`." ) atomic_gemm = 0 @@ -407,13 +443,15 @@ def add_ub( if (quantization_mode == UserBufferQuantizationMode.FP8 and fp8_buf) else dtype ) + comm_type = tex.CommOverlapType.RS if is_reduce_scatter else tex.CommOverlapType.AG if method == "ring_exchange": ub_obj = tex.CommOverlapP2P( shape, # Communication buffer shape buffer_dtype, # Communication buffer data type helper, # Helper for torch.distributed callbacks during bootstrapping - tp_size, # Tensor-parallel group size (may be different than local_size) - tex.CommOverlapType.RS if is_reduce_scatter else tex.CommOverlapType.AG, + tp_size, # Tensor-parallel group size (may differ from local_size) + comm_type, + use_cublasmp=with_cublasmp, num_max_streams=_NUM_MAX_UB_STREAMS, comm_cga_size=cga_size, num_comm_sm=num_sm, @@ -429,7 +467,9 @@ def add_ub( shape, # Communication buffer shape buffer_dtype, # Communication buffer data type helper, # Helper for torch.distributed callbacks during bootstrapping - tp_size, # Tensor-parallel group size (may be different than local_size) + tp_size, # Tensor-parallel group size (may differ from local_size) + use_cublasmp=with_cublasmp, + comm_type=comm_type, num_splits=num_splits, num_max_streams=_NUM_MAX_UB_STREAMS, comm_cga_size=cga_size, @@ -464,9 +504,30 @@ def add_ub( new_method = user_ub_cfg[name]["method"] methods[new_method].append(name) - for name in ( - methods["ring_exchange"] + methods["pipeline"] + methods["bulk"] + methods["external"] - ): + # Adjust defaults to account for the fact that cuBLASMp does not support + # bulk or external overlaps + if with_cublasmp: + warnings.warn( + "cuBLASMp does not support bulk or external overlaps. " + "'qkv_dgrad' and 'fc1_dgrad' GEMMs will be configured with 'ring_exchange'" + "overlap unless user configuration specifies otherwise. Bulk overlaps for the " + "corresponding 'qkv_wgrad' and 'fc1_wgrad' GEMMs will be disabled." + ) + methods["bulk"] = [] + methods["external"] = [] + external_gemm_to_overlap.clear() + + for name in dgrad_reduce_scatter_overlap: + wgrad_name = name.replace("dgrad", "wgrad") + if name not in layers_reduce_scatter_overlap: + layers_reduce_scatter_overlap.append(name) + if wgrad_name in layers_reduce_scatter_overlap: + layers_reduce_scatter_overlap.remove(wgrad_name) + if name not in methods["ring_exchange"] and name not in methods["pipeline"]: + methods["ring_exchange"].append(name) + + configured_methods = ["ring_exchange", "pipeline", "bulk", "external"] + for name in (m for k in configured_methods for m in methods[k]): ub_cfg = get_default_config(name) if user_ub_cfg is not None and name in user_ub_cfg: fp8_buf = (name in layers_all_gather_overlap) or ( @@ -476,6 +537,9 @@ def add_ub( ub_cfg["fp8_buf"] = fp8_buf add_ub(name, quantization_mode, **ub_cfg) + global _ub_initialized + _ub_initialized = True + def get_ub(name: str, use_fp8: bool): """Get userbuffer communicator corresponding to give key.""" @@ -483,7 +547,7 @@ def get_ub(name: str, use_fp8: bool): # So favour simplicity until the correct design becomes clear. # This is mainly an internal API so we don't need to worry about future changes key = (name, UserBufferQuantizationMode.FP8 if use_fp8 else UserBufferQuantizationMode.NONE) - if _ub_communicators is None: + if not _ub_initialized or _ub_communicators is None: raise RuntimeError("UB manager is not initialized.") if key not in _ub_communicators: raise KeyError(f"UB for {name} with use_fp8={use_fp8} is not registered.") @@ -492,8 +556,10 @@ def get_ub(name: str, use_fp8: bool): def destroy_ub(): """Destroy all allocated userbuffer communicators.""" - global _ub_communicators + global _ub_communicators, _ub_with_cublasmp, _ub_initialized _ub_communicators = None + _ub_with_cublasmp = False + _ub_initialized = False global layers_atomic_ring_exchange layers_atomic_ring_exchange = [] @@ -516,6 +582,15 @@ def fill_userbuffers_buffer_for_all_gather( tensor's metadata, e.g. scaling factors. """ + # cuBlasMp handles its own buffer filling and gather, but its fused all-gather + GEMM + # still needs an already-quantized local tensor for FP8 GEMMs. The standard FP8 paths + # pre-quantize before this call and the quantized-norm path produces an already-quantized + # tensor, but custom-recipe paths pass an unquantized tensor and rely on this helper to + # quantize it. The callers configure the quantizer usage (rowwise) expected by cuBLASMp. + if comm.with_cublasmp(): + if quantizer is not None and not isinstance(local_tensor, QuantizedTensorStorage): + local_tensor = quantizer(local_tensor) + return local_tensor, local_tensor # Tensor dimensions local_shape = local_tensor.size() diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 8c88f3ee82e..7fc96d4779f 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -21,6 +21,8 @@ from .base import ( fill_userbuffers_buffer_for_all_gather, get_ub, + is_ub_initialized, + using_cublasmp_backend, quantize_weight, TransformerEngineBaseModule, get_dummy_wgrad, @@ -397,7 +399,10 @@ def forward( clear_tensor_data(ln_out, ln_out_total) ln_out = ln_out_total = None elif with_input_all_gather and not return_layernorm_output_gathered: - clear_tensor_data(ln_out_total) + # ln_out_total aliases ln_out for the cuBLASMp backend; skip the + # deallocation to avoid corrupting the backward-saved tensor. + if ln_out_total is not ln_out: + clear_tensor_data(ln_out_total) ln_out_total = None # ------------------------------------------------------ @@ -406,7 +411,9 @@ def forward( # ------------------------------------------------------ out = None if ub_overlap_rs_fprop: - out = reduce_scatter_out + # cuBLASMp writes the reduce-scattered output directly into the + # GEMM output tensor; Userbuffers writes it into the extra-output buffer. + out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out elif parallel_mode == "row" and tp_size > 1: nvtx_range_push(f"{nvtx_label}.row_parallel_comm") out = gemm_out @@ -838,7 +845,13 @@ def backward( dgrad = None dgrad_work = None if ctx.ub_overlap_rs_dgrad: - dgrad = reduce_scatter_out + # cuBLASMp writes the reduce-scattered dgrad directly into the + # GEMM output tensor; Userbuffers uses the extra-output buffer. + dgrad = ( + gemm_out + if ub_obj_dgrad is not None and ub_obj_dgrad.with_cublasmp() + else reduce_scatter_out + ) elif ctx.ub_bulk_wgrad: dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) elif ctx.parallel_mode == "column" and ctx.tp_size > 1: @@ -860,6 +873,28 @@ def backward( # Grad input tensor has been computed... # -------------------------------------------------- + # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and + # does not preserve it for wgrad. Userbuffers leaves the gathered + # tensor in its persistent buffer; cuBLASMp does not, so we gather + # here. Route through the same FP8-aware all-gather as the + # non-overlap path in + # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing + # the grad_output quantizer. Columnwise data needed for wgrad is + # produced by ``update_usage(columnwise_usage=True)`` further below. + if ( + ctx.requires_wgrad + and ctx.ub_overlap_ag + and ctx.ub_obj_gradout is not None + and ctx.ub_obj_gradout.with_cublasmp() + ): + if ctx.grad_output_quantizer is not None: + ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + grad_output, _ = gather_along_first_dim( + grad_output, + ctx.tp_group, + quantizer=ctx.grad_output_quantizer, + ) + # -------------------------------------------------- # Compute grad weight # -------------------------------------------------- @@ -1303,6 +1338,8 @@ def __init__( self.ub_overlap_rs_dgrad = ( ub_overlap_rs_dgrad and self.sequence_parallel and self.parallel_mode == "column" ) + # Bulk overlaps require the Userbuffers backend; the cuBLASMp backend + # falls back to async NCCL ops via torch.distributed. self.ub_bulk_wgrad = ( ub_bulk_wgrad and self.sequence_parallel @@ -1333,9 +1370,23 @@ def __init__( self.ub_overlap_ag_dgrad, ] ): + assert is_ub_initialized(), "initialize_ub() must be called before layer construction." assert ub_name is not None, "Userbuffer name [string] is not set." self.ub_name = ub_name + if using_cublasmp_backend(): + if self.ub_bulk_dgrad: + warnings.warn( + f"cuBLASMp backend does not support bulk overlaps for '{self.ub_name}_dgrad' " + f"and '{self.ub_name}_wgrad' GEMMs. Falling back on DGRAD+RS overlap for " + f"'{self.ub_name}_dgrad' GEMM with no bulk overlap for '{self.ub_name}_wgrad' " + "GEMM. In order to enable bulk overlaps for these GEMMs, set " + "`with_cublasmp=False` when calling `initialize_ub()`." + ) + self.ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad or self.ub_bulk_dgrad + self.ub_bulk_dgrad = False + self.ub_bulk_wgrad = False + if self.symmetric_ar_type is not None: assert torch_version() >= ( 2, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index c837af9d337..6c6cca74ef5 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -23,6 +23,8 @@ fill_userbuffers_buffer_for_all_gather, _ub_communicators, get_ub, + is_ub_initialized, + using_cublasmp_backend, quantize_weight, TransformerEngineBaseModule, _2X_ACC_FPROP, @@ -595,8 +597,12 @@ def _forward( # first part of if statement means that we only clear ln_out_total if # 1) checkpointing and not recomputing (in the forward stage, not bwd recompute stage) # 2) not checkpointing and grad disabled - if ((checkpoint and not is_recomputation) or not is_grad_enabled) and ( - ln_out_total is not ln_out_return + # The `is not ln_out` guard avoids clearing the bwd-saved tensor when + # ln_out_total aliases ln_out (cuBLASMp AG-fprop path). + if ( + ((checkpoint and not is_recomputation) or not is_grad_enabled) + and ln_out_total is not ln_out_return + and ln_out_total is not ln_out ): clear_tensor_data(ln_out_total) @@ -695,7 +701,13 @@ def _forward( # Note: Perform tensor-parallel communication if needed fc2_out = None if ub_overlap_rs: - fc2_out = reduce_scatter_out + # cuBLASMp writes the reduce-scattered output directly into the + # GEMM output tensor; Userbuffers writes it into the extra-output buffer. + fc2_out = ( + gemm_out + if ub_obj_fc2out is not None and ub_obj_fc2out.with_cublasmp() + else reduce_scatter_out + ) elif set_parallel_mode and sequence_parallel: fc2_out, _ = reduce_scatter_along_first_dim(gemm_out, tp_group) elif set_parallel_mode and tensor_parallel: @@ -1255,6 +1267,29 @@ def backward( # Finished FC2 DGRAD... # -------------------------------------------------- + # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and + # does not preserve it for fc2_wgrad. Userbuffers leaves the + # gathered tensor in its persistent buffer; cuBLASMp does not, so + # we gather here. Route through the same FP8-aware all-gather as + # the non-overlap path in + # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing + # the grad_output quantizer. Columnwise data needed for fc2_wgrad + # is produced by ``update_usage(columnwise_usage=True)`` further + # below. + if ( + ctx.fc2_weight_requires_grad + and ctx.ub_overlap_ag + and ctx.ub_obj_gradout is not None + and ctx.ub_obj_gradout.with_cublasmp() + ): + if ctx.fc2_grad_output_quantizer is not None: + ctx.fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + grad_output, _ = gather_along_first_dim( + grad_output, + ctx.tp_group, + quantizer=ctx.fc2_grad_output_quantizer, + ) + # -------------------------------------------------- # FC2 WGRAD # -------------------------------------------------- @@ -1524,7 +1559,13 @@ def fc2_wgrad_gemm( fc1_dgrad = None fc1_dgrad_work = None if ctx.ub_overlap_rs_dgrad: - fc1_dgrad = reduce_scatter_out + # cuBLASMp writes the reduce-scattered dgrad directly into the + # GEMM output tensor; Userbuffers uses the extra-output buffer. + fc1_dgrad = ( + gemm_out + if ub_obj_fc1_dgrad is not None and ub_obj_fc1_dgrad.with_cublasmp() + else reduce_scatter_out + ) elif ctx.ub_bulk_wgrad: fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True) elif ctx.set_parallel_mode and not ctx.ub_bulk_wgrad: @@ -1985,6 +2026,29 @@ def __init__( ub_bulk_dgrad and self.sequence_parallel and not self.ub_overlap_rs_dgrad ) + if any( + [ + self.ub_overlap_ag, + self.ub_overlap_rs, + self.ub_overlap_rs_dgrad, + self.ub_bulk_dgrad, + self.ub_bulk_wgrad, + ] + ): + assert is_ub_initialized(), "initialize_ub() must be called before layer construction." + + if using_cublasmp_backend(): + if self.ub_bulk_dgrad: + warnings.warn( + "cuBLASMp backend does not support bulk overlaps for 'fc1_dgrad' and " + "'fc1_wgrad' GEMMs. Falling back on DGRAD+RS overlap for 'fc1_dgrad' GEMM with " + "no bulk overlap for 'fc1_wgrad' GEMM. In order to enable bulk overlaps for " + "these GEMMs, set `with_cublasmp=False` when calling `initialize_ub()`." + ) + self.ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad or self.ub_bulk_dgrad + self.ub_bulk_dgrad = False + self.ub_bulk_wgrad = False + if self.symmetric_ar_type is not None: assert torch_version() >= ( 2, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index dcbb9eaf930..6c2d98d160d 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -21,6 +21,8 @@ fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, + is_ub_initialized, + using_cublasmp_backend, quantize_weight, TransformerEngineBaseModule, _2X_ACC_FPROP, @@ -504,7 +506,9 @@ def _linear_forward_impl( # ------------------------------------------------------ out = None if ub_overlap_rs_fprop: - out = reduce_scatter_out + # cuBLASMp writes the reduce-scattered output directly into the GEMM + # output tensor; Userbuffers writes it into the extra-output buffer. + out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out elif parallel_mode == "row" and args.tp_size > 1: nvtx_range_push(f"{nvtx_label}.row_parallel_comm") out = gemm_out @@ -518,6 +522,13 @@ def _linear_forward_impl( nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") else: out = gemm_out + + # Restore the input's logical rank (e.g., (seq, batch, hidden)) on the output. + # This is mainly to correct for cuBLASMp comm+GEMM operators that unconditionally + # return a 2D output buffer that ends up incompatible with downstream consumers + # (e.g. ``bias_dropout_add`` residual connections inside ``TransformerLayer``). + out = out.view(-1, *inp.shape[1:-1], out_features) + # ------------------------------------------------------ # Output tensor is ready to return... # ------------------------------------------------------ @@ -1010,7 +1021,13 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # Prepare grad input tensor # Note: Perform tensor-parallel communication if bwd_args.ub_overlap_rs_dgrad: - dgrad = reduce_scatter_out + # cuBLASMp writes the reduce-scattered dgrad directly into the + # GEMM output tensor; Userbuffers uses the extra-output buffer. + dgrad = ( + gemm_out + if ub_obj_dgrad is not None and ub_obj_dgrad.with_cublasmp() + else reduce_scatter_out + ) elif bwd_args.ub_bulk_wgrad: dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) elif bwd_args.parallel_mode == "column" and bwd_args.tp_size > 1: @@ -1032,6 +1049,27 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # Grad input tensor has been computed... # -------------------------------------------------- + # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and does + # not preserve it for wgrad. Userbuffers leaves the gathered tensor in + # its persistent buffer; cuBLASMp does not, so we gather here. Route + # through the same FP8-aware all-gather as the non-overlap path in + # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing the + # grad_output quantizer. The columnwise data needed for wgrad is then + # produced by ``update_usage(columnwise_usage=True)`` further below. + if ( + bwd_args.requires_wgrad + and bwd_args.ub_overlap_ag + and bwd_args.ub_obj_gradout is not None + and bwd_args.ub_obj_gradout.with_cublasmp() + ): + if grad_output_quantizer is not None: + grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + grad_output, _ = gather_along_first_dim( + grad_output, + bwd_args.tp_group, + quantizer=grad_output_quantizer, + ) + # -------------------------------------------------- # Compute grad weight # -------------------------------------------------- @@ -1494,6 +1532,8 @@ def __init__( self.ub_overlap_rs_dgrad = ( self.parallel_mode == "column" and self.sequence_parallel and ub_overlap_rs_dgrad ) + # Bulk overlaps require the Userbuffers backend; the cuBLASMp backend falls back on + # DGRAD+RS overlap with no bulk overlap for WGRAD self.ub_bulk_dgrad = ( self.parallel_mode == "column" and self.sequence_parallel @@ -1525,9 +1565,23 @@ def __init__( self.ub_bulk_wgrad, ] ): + assert is_ub_initialized(), "initialize_ub() must be called before layer construction." assert ub_name is not None, f"Comm+GEMM overlap layer '{ub_name}' is not initialized." self.ub_name = ub_name + if using_cublasmp_backend(): + if self.ub_bulk_dgrad: + warnings.warn( + f"cuBLASMp backend does not support bulk overlaps for '{self.ub_name}_dgrad' " + f"and '{self.ub_name}_wgrad' GEMMs. Falling back on DGRAD+RS overlap for " + f"'{self.ub_name}_dgrad' GEMM with no bulk overlap for '{self.ub_name}_wgrad' " + "GEMM. In order to enable bulk overlaps for these GEMMs, set " + "`with_cublasmp=False` when calling `initialize_ub()`." + ) + self.ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad or self.ub_bulk_dgrad + self.ub_bulk_dgrad = False + self.ub_bulk_wgrad = False + if self.symmetric_ar_type is not None: assert torch_version() >= ( 2, diff --git a/transformer_engine/pytorch/onnx_extensions.py b/transformer_engine/pytorch/onnx_extensions.py index 4d3b90bf638..6e97cf655bb 100644 --- a/transformer_engine/pytorch/onnx_extensions.py +++ b/transformer_engine/pytorch/onnx_extensions.py @@ -30,7 +30,7 @@ from .tensor.float8_tensor import Float8Quantizer from .tensor.mxfp8_tensor import MXFP8Quantizer -from .constants import MXFP8_BLOCK_SCALING_SIZE +from .constants import MXFP8_BLOCK_SCALING_SIZE, DType from .utils import round_up_to_nearest_multiple from .export import is_in_onnx_export_mode @@ -85,7 +85,7 @@ def onnx_quantize_fp8_op(tensor: torch.Tensor, scale: float) -> torch.Tensor: """Quantize to Float8Tensor used for inference.""" scale_tensor = torch.tensor(scale, dtype=torch.float32, device=tensor.device) amax_tensor = torch.tensor([1], dtype=torch.float32, device=tensor.device) - quantizer = Float8Quantizer(scale_tensor, amax_tensor, tex.DType.kFloat8E4M3) + quantizer = Float8Quantizer(scale_tensor, amax_tensor, DType.kFloat8E4M3) return quantizer.quantize(tensor)._data @@ -130,9 +130,7 @@ def onnx_quantize_fp8_symbolic( @torch.library.custom_op("tex::fp8_dequantize", mutates_args=[]) def onnx_dequantize_fp8_op(tensor: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor: """Dequantize from Float8Tensor used for inference.""" - quantizer = Float8Quantizer( - 1 / scale_inv, torch.zeros(1).to(tensor.device), tex.DType.kFloat8E4M3 - ) + quantizer = Float8Quantizer(1 / scale_inv, torch.zeros(1).to(tensor.device), DType.kFloat8E4M3) quantizer_tensor = quantizer.create_tensor_from_data(tensor, fake_dtype=torch.float32) return quantizer_tensor.dequantize() @@ -212,7 +210,7 @@ def onnx_quantize_fp8_cs_symbolic( @torch.library.custom_op("tex::mxfp8_quantize", mutates_args=[]) def onnx_quantize_mxfp8_op(tensor: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Quantize to MXFP8Tensor used for inference.""" - quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(DType.kFloat8E4M3) quantized_tensor = quantizer(tensor) return quantized_tensor._rowwise_data, quantized_tensor._rowwise_scale_inv @@ -264,7 +262,7 @@ def onnx_quantize_mxfp8_symbolic( @torch.library.custom_op("tex::mxfp8_dequantize", mutates_args=[]) def onnx_dequantize_mxfp8_op(tensor: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor: """Dequantize from MXFP8Tensor used for inference.""" - quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(DType.kFloat8E4M3) quantizer_tensor = quantizer.create_tensor_from_data( tensor, scale_inv, fake_dtype=torch.float32 ) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 87911d76f49..19501223f1e 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -61,25 +61,15 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() -def _group_quantize_for_grouped_mlp( +def _wrap_single_nvfp4_as_grouped( tensor: torch.Tensor, - quantizer: Quantizer, - num_groups: int, + quantized: NVFP4Tensor | NVFP4TensorStorage, + quantizer: NVFP4Quantizer, split_sizes: Optional[torch.Tensor], *, tensor_offsets: Optional[torch.Tensor] = None, ) -> GroupedTensor: - """Quantize into grouped storage.""" - - # Typical case: group-quantize - if num_groups != 1 or not isinstance(quantizer, NVFP4Quantizer): - return tex.group_quantize(tensor, quantizer, num_groups, split_sizes) - - # -------------------------------------------------- - # Special case: single-tensor NVFP4 quantize - # -------------------------------------------------- - - quantized = tex.quantize(tensor, quantizer) + """Wrap a single NVFP4 tensor in GroupedTensor storage.""" with_gemm_swizzled_scales = quantized._with_gemm_swizzled_scales if quantizer.optimize_for_gemm: tex.swizzle_scales_for_gemm_(quantized) @@ -130,6 +120,78 @@ def _group_quantize_for_grouped_mlp( ) +def _group_quantize_for_grouped_mlp( + tensor: torch.Tensor, + quantizer: Quantizer, + num_groups: int, + split_sizes: Optional[torch.Tensor], + *, + tensor_offsets: Optional[torch.Tensor] = None, +) -> GroupedTensor: + """Quantize into grouped storage.""" + + if num_groups != 1 or not isinstance(quantizer, NVFP4Quantizer): + return tex.group_quantize( + tensor, + quantizer, + num_groups, + split_sizes, + tensor_offsets=tensor_offsets, + ) + + quantized = tex.quantize(tensor, quantizer) + return _wrap_single_nvfp4_as_grouped( + tensor, + quantized, + quantizer, + split_sizes, + tensor_offsets=tensor_offsets, + ) + + +def _group_quantize_with_amax_for_grouped_mlp( + tensor: torch.Tensor, + quantizer: Quantizer, + num_groups: int, + split_sizes: Optional[torch.Tensor], + rowwise_amax: torch.Tensor, + columnwise_amax: torch.Tensor, + *, + tensor_offsets: Optional[torch.Tensor] = None, +) -> GroupedTensor: + """Quantize with precomputed NVFP4 amaxes into grouped storage.""" + if not isinstance(quantizer, NVFP4Quantizer): + return _group_quantize_for_grouped_mlp( + tensor, + quantizer, + num_groups, + split_sizes, + tensor_offsets=tensor_offsets, + ) + + if num_groups != 1: + return tex.nvfp4_group_quantize_with_amax( + tensor, + quantizer, + num_groups, + split_sizes, + rowwise_amax, + columnwise_amax, + tensor_offsets=tensor_offsets, + ) + + quantized = tex.nvfp4_quantize_with_amax( + tensor, quantizer, rowwise_amax.view(-1)[:1], columnwise_amax.view(-1)[:1] + ) + return _wrap_single_nvfp4_as_grouped( + tensor, + quantized, + quantizer, + split_sizes, + tensor_offsets=tensor_offsets, + ) + + def _nvfp4_amax( tensors: GroupedTensor | Iterable[NVFP4TensorStorage], *, diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index eacc36b36cc..f4beffe90c3 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -12,6 +12,7 @@ import torch import transformer_engine_torch as tex +from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data @@ -107,7 +108,7 @@ def op_forward( # Quantize input to FP8 before caching if needed if self.cache_quantized_input: - input_quantizer = Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, x.device) + input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) input_quantizer.set_usage(rowwise=True, columnwise=False) x = input_quantizer(x) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index e9787f96b25..521ee59fa01 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -14,6 +14,7 @@ import torch import transformer_engine_torch as tex +from ...constants import DType from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore @@ -537,7 +538,7 @@ def _quantize_weights_mxfp8( weight = MXFP8Tensor( shape=unpacked_shape, dtype=dtype, - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, rowwise_data=rowwise_data[group_idx], rowwise_scale_inv=rowwise_scales[group_idx], columnwise_data=columnwise_data[group_idx], diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 39571fcd85b..02f330ede30 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -11,6 +11,7 @@ import torch import transformer_engine_torch as tex +from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data @@ -118,7 +119,7 @@ def op_forward( # Quantize input to FP8 before caching if needed if self.cache_quantized_input: input_quantizer = Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, + DType.kFloat8E4M3, input_.device, ) input_quantizer.set_usage(rowwise=True, columnwise=False) @@ -304,7 +305,7 @@ def op_forward( # Quantize input to FP8 before caching if needed if self.cache_quantized_input: - input_quantizer = Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, x.device) + input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) input_quantizer.set_usage(rowwise=True, columnwise=False) x = input_quantizer(x) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 792b6d7811f..25029190c11 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -88,6 +88,10 @@ def _cudnn_compute_wgrad( accumulate: bool, wgrad_kernel_fn, single_grouped_weight: bool, + use_nvfp4: bool, + data_dtype: torch.dtype, + scale_view_dtype: torch.dtype, + sf_vec_size: int, current_stream=None, ): """Compute wgrad using the cuDNN CuTe DSL grouped GEMM wgrad kernel. @@ -100,8 +104,6 @@ def _cudnn_compute_wgrad( out_features, in_features = weight_shape total_tokens = grouped_dy.logical_shape[0] - fp8_dtype = torch.float8_e4m3fn - sfa_leading_dim = round_up_to_nearest_multiple(out_features, 128) sfb_leading_dim = round_up_to_nearest_multiple(in_features, 128) @@ -110,68 +112,116 @@ def _cudnn_compute_wgrad( # Even for this case, cuteDSL still requires the same # stride requirements for the input and scale tensors. device = grouped_dy.columnwise_data.device - a_tensor = torch.empty_strided((out_features, 0), (16, 1), dtype=fp8_dtype, device=device) + a_tensor = torch.empty_strided( + (out_features, 0), + (16, 1), + dtype=data_dtype, + device=device, + ) b_tensor = torch.empty_strided( - (0, in_features), (in_features, 1), dtype=fp8_dtype, device=device + (0, in_features), + (in_features, 1), + dtype=data_dtype, + device=device, ) sfa_tensor = torch.empty_strided( (sfa_leading_dim, 0), (16, 1), - dtype=torch.float8_e8m0fnu, + dtype=scale_view_dtype, device=device, ) sfb_tensor = torch.empty_strided( (sfb_leading_dim, 0), (16, 1), - dtype=torch.float8_e8m0fnu, + dtype=scale_view_dtype, device=device, ) + elif use_nvfp4: + # NVFP4 columnwise data is stored expert-major as per-expert + # (logical_K, group_M / 2) chunks. cuDNN consumes that layout directly + # with input_order="tensor_ragged". + a_tensor = grouped_dy.columnwise_data.view(dtype=data_dtype).view( + out_features, + total_tokens // 2, + ) + b_tensor = ( + grouped_x.columnwise_data.view(dtype=data_dtype) + .view( + in_features, + total_tokens // 2, + ) + .T + ) + sfa_tensor = grouped_dy.columnwise_scale_inv.view(sfa_leading_dim, -1).view( + dtype=scale_view_dtype + ) + sfb_tensor = grouped_x.columnwise_scale_inv.view(sfb_leading_dim, -1).view( + dtype=scale_view_dtype + ) else: a_tensor = ( - grouped_dy.columnwise_data.view(dtype=fp8_dtype).view(total_tokens, out_features).T + grouped_dy.columnwise_data.view(dtype=data_dtype).view(total_tokens, out_features).T ) - b_tensor = grouped_x.columnwise_data.view(dtype=fp8_dtype).view(total_tokens, in_features) + b_tensor = grouped_x.columnwise_data.view(dtype=data_dtype).view(total_tokens, in_features) sfa_tensor = grouped_dy.columnwise_scale_inv.view(sfa_leading_dim, -1).view( - dtype=torch.float8_e8m0fnu + dtype=scale_view_dtype ) sfb_tensor = grouped_x.columnwise_scale_inv.view(sfb_leading_dim, -1).view( - dtype=torch.float8_e8m0fnu + dtype=scale_view_dtype ) + common_wgrad_kwargs = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "offsets_tensor": offsets, + "acc_dtype": torch.float32, + "sf_vec_size": sf_vec_size, + "accumulate_on_output": accumulate, + "current_stream": current_stream, + } + if use_nvfp4: + global_scale_denom = 448.0 * 6.0 + if total_tokens == 0: + global_scale_shape = (offsets.shape[0],) + common_wgrad_kwargs["global_scale_a"] = torch.zeros( + global_scale_shape, + dtype=torch.float32, + device=device, + ) + common_wgrad_kwargs["global_scale_b"] = torch.zeros( + global_scale_shape, + dtype=torch.float32, + device=device, + ) + else: + common_wgrad_kwargs["global_scale_a"] = ( + _nvfp4_amax(grouped_dy, columnwise=True).to(torch.float32) / global_scale_denom + ) + common_wgrad_kwargs["global_scale_b"] = ( + _nvfp4_amax(grouped_x, columnwise=True).to(torch.float32) / global_scale_denom + ) + common_wgrad_kwargs["input_order"] = "tensor_ragged" + # Prepare wgrad output if single_grouped_weight: # Dense mode: single (num_groups, out_features, in_features) tensor wgrad_tensor = wgrad_output.rowwise_data.view(offsets.shape[0], out_features, in_features) wgrad_kernel_fn( - a_tensor=a_tensor, - b_tensor=b_tensor, - sfa_tensor=sfa_tensor, - sfb_tensor=sfb_tensor, - offsets_tensor=offsets, + **common_wgrad_kwargs, output_mode="dense", wgrad_tensor=wgrad_tensor, - acc_dtype=torch.float32, wgrad_dtype=wgrad_tensor.dtype, - sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, - accumulate_on_output=accumulate, - current_stream=current_stream, ) else: # Discrete mode: per-expert wgrad device pointers wgrad_ptrs = tex.copy_data_ptrs_to_device(wgrad_output, wgrad_output[0].device) wgrad_kernel_fn( - a_tensor=a_tensor, - b_tensor=b_tensor, - sfa_tensor=sfa_tensor, - sfb_tensor=sfb_tensor, - offsets_tensor=offsets, + **common_wgrad_kwargs, output_mode="discrete", wgrad_ptrs=wgrad_ptrs, - acc_dtype=torch.float32, wgrad_dtype=wgrad_output[0].dtype, - sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, - accumulate_on_output=accumulate, - current_stream=current_stream, ) @@ -189,6 +239,10 @@ def _compute_grad_params( label="", *, cudnn_wgrad_kernel_fn, + use_nvfp4, + data_dtype, + scale_view_dtype, + sf_vec_size, offsets, ): """Compute weight gradients and build grad_params for a GroupedLinear layer. @@ -258,6 +312,10 @@ def _compute_grad_params( accumulate=accumulate_into_main_grad, wgrad_kernel_fn=cudnn_wgrad_kernel_fn, single_grouped_weight=fc_op.single_grouped_weight, + use_nvfp4=use_nvfp4, + data_dtype=data_dtype, + scale_view_dtype=scale_view_dtype, + sf_vec_size=sf_vec_size, current_stream=torch.cuda.current_stream().cuda_stream, ) elif ( @@ -403,7 +461,8 @@ def fuser_backward( ]: # Get basic operations - fc1_op, _, fc2_op = self.basic_ops + fc1_op, activation_op, fc2_op = self.basic_ops + activation_is_srelu = isinstance(activation_op, ScaledSReLU) fc1_ctx, activation_ctx, fc2_ctx = basic_op_ctxs # Tensor properties @@ -576,13 +635,23 @@ def fuser_backward( if use_nvfp4: nvfp4_fp4_max = 6.0 nvfp4_fp8_max = 448.0 - fc2_alpha_tensor = ( - torch.sqrt( - _nvfp4_amax(grouped_fc2_dy, columnwise=False) - * _nvfp4_amax(grouped_fc2_weight, columnwise=True) + nvfp4_global_scale_denom = nvfp4_fp4_max * nvfp4_fp8_max + fc2_dy_amax = _nvfp4_amax(grouped_fc2_dy, columnwise=False) + fc2_weight_col_amax = _nvfp4_amax(grouped_fc2_weight, columnwise=True) + if activation_is_srelu: + # DSReLU applies alpha once, so pass the full product of the + # two operand global scales. + fc2_alpha_tensor = ( + (fc2_dy_amax * fc2_weight_col_amax / (nvfp4_global_scale_denom**2)) + .to(torch.float32) + .expand(num_groups) ) - / (nvfp4_fp8_max * nvfp4_fp4_max) - ).expand(num_groups) + else: + # DGLU applies alpha to both gate branches, so the wrapper + # expects sqrt(product) to recover the same global-scale factor. + fc2_alpha_tensor = ( + torch.sqrt(fc2_dy_amax * fc2_weight_col_amax) / nvfp4_global_scale_denom + ).expand(num_groups) fc2_beta_tensor = get_cached_ones_tensor(num_groups, torch.float32, device) fc2_norm_const_tensor = None else: @@ -649,17 +718,13 @@ def fuser_backward( fc2_dactivation_kwargs["b_tensor"] = fc2_w_data fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs = tex.copy_data_ptrs_to_device( - [w._columnwise_data for w in grouped_fc2_weight], - device, - ) - swizzle_type = ( - "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_columnwise_swizzle" - ) - fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - swizzle_type, - [w._columnwise_scale_inv for w in grouped_fc2_weight], - device, + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._columnwise_data for w in grouped_fc2_weight], + [w._columnwise_scale_inv for w in grouped_fc2_weight], + "nvfp4" if use_nvfp4 else "mxfp8_columnwise", + device, + ) ) fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs @@ -700,28 +765,44 @@ def fuser_backward( "the recomputed FC2 input tensor." ) - sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") - if sfd_col_d_srelu_tensor is None: - raise RuntimeError( - "SReLU recompute is enabled, but the DSReLU kernel did not return " - "the recomputed FC2 input column scale tensor." + if use_nvfp4: + # cuDNN's grouped dSReLU FP4 path recomputes FC2 input in BF16. + # Re-quantize to produce the columnwise FP4 data, scales, and + # dSReLU-specific amax needed by the NVFP4 wgrad GEMM. + fc2_x_bf16 = d_srelu_tensor.view(out_shape[0], fc2_weight_shape[1]).contiguous() + fc2_input_quantizer = fc2_ctx.input_quantizers[0] + fc2_input_quantizer.set_usage(rowwise=False, columnwise=True) + fc2_input_quantizer.optimize_for_gemm = True + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_x_bf16, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=base_split_offsets * fc2_weight_shape[1], ) + else: + sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") + if sfd_col_d_srelu_tensor is None: + raise RuntimeError( + "SReLU recompute is enabled, but the DSReLU kernel did not return " + "the recomputed FC2 input column scale tensor." + ) - fc2_x_col_data = d_srelu_tensor.view(out_shape[0], fc2_weight_shape[1]) - fc2_x_col_scale = sfd_col_d_srelu_tensor.permute(5, 2, 4, 0, 1, 3) - grouped_fc2_x = GroupedTensor( - shape=(out_shape[0], fc2_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc2_ctx.input_quantizers[0], - data=None, - columnwise_data=fc2_x_col_data.reshape(-1), - scale_inv=None, - columnwise_scale_inv=fc2_x_col_scale.reshape(-1), - first_dims=split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[1], - with_gemm_swizzled_scales=True, - ) + fc2_x_col_data = d_srelu_tensor.view(out_shape[0], fc2_weight_shape[1]) + fc2_x_col_scale = sfd_col_d_srelu_tensor.permute(5, 2, 4, 0, 1, 3) + grouped_fc2_x = GroupedTensor( + shape=(out_shape[0], fc2_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=fc2_ctx.input_quantizers[0], + data=None, + columnwise_data=fc2_x_col_data.reshape(-1), + scale_inv=None, + columnwise_scale_inv=fc2_x_col_scale.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * fc2_weight_shape[1], + with_gemm_swizzled_scales=True, + ) fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc2_bias_grad_packed: Optional[torch.Tensor] = None @@ -793,7 +874,7 @@ def fuser_backward( ) # FC2 wgrad GEMM - wgrad_kernel_fn = None if use_nvfp4 else self.grouped_gemm_wgrad_kernel() + wgrad_kernel_fn = self.grouped_gemm_wgrad_kernel() fc2_grad_params = _compute_grad_params( fc_op=fc2_op, ctx=fc2_ctx, @@ -807,6 +888,10 @@ def fuser_backward( bias_grad_packed=fc2_bias_grad_packed, label="FC2", cudnn_wgrad_kernel_fn=wgrad_kernel_fn, + use_nvfp4=use_nvfp4, + data_dtype=data_dtype, + scale_view_dtype=scale_view_dtype, + sf_vec_size=sf_vec_size, offsets=split_points, ) @@ -911,17 +996,13 @@ def fuser_backward( fc1_dgrad_kwargs["b_tensor"] = fc1_w_data fc1_dgrad_kwargs["sfb_tensor"] = fc1_w_scales else: - fc1_b_ptrs = tex.copy_data_ptrs_to_device( - [w._columnwise_data for w in grouped_fc1_weight], - device, - ) - swizzle_type = ( - "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_columnwise_swizzle" - ) - fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - swizzle_type, - [w._columnwise_scale_inv for w in grouped_fc1_weight], - device, + fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._columnwise_data for w in grouped_fc1_weight], + [w._columnwise_scale_inv for w in grouped_fc1_weight], + "nvfp4" if use_nvfp4 else "mxfp8_columnwise", + device, + ) ) fc1_dgrad_kwargs["b_ptrs"] = fc1_b_ptrs @@ -947,6 +1028,10 @@ def fuser_backward( bias_grad_packed=fc1_bias_grad_packed, label="FC1", cudnn_wgrad_kernel_fn=wgrad_kernel_fn, + use_nvfp4=use_nvfp4, + data_dtype=data_dtype, + scale_view_dtype=scale_view_dtype, + sf_vec_size=sf_vec_size, offsets=split_points, ) @@ -1039,8 +1124,6 @@ def fuse_backward_srelu_ops( ) -> list[FusibleOperation]: """Apply GroupedLinear + ScaledSReLU + GroupedLinear fusion for backward pass.""" - if recipe is None or not recipe.mxfp8(): - return ops return fuse_grouped_mlp_ops( ops, recipe=recipe, diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index f4f21085781..ece670a539c 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -33,6 +33,7 @@ _cudnn_frontend_version_supported, _cudnn_frontend_supports_grouped_gemm_srelu, _group_quantize_for_grouped_mlp, + _group_quantize_with_amax_for_grouped_mlp, _nvidia_cudnn_frontend_supports_wgrad, _nvfp4_amax, _nvfp4_single_tensor_from_grouped, @@ -76,6 +77,12 @@ def _grouped_gemm_dsrelu_backward_supported() -> bool: return grouped_gemm_dsrelu_wrapper_sm100 is not None +@functools.lru_cache(maxsize=1) +def _use_tmem_post_rht_amax() -> bool: + """Whether the FC1 GLU+RHT+amax kernel should use TMEM post-RHT amax.""" + return os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP_FC1_GLU_RHT_AMAX_TMEM", "0") == "1" + + class _ForwardGroupedMLP_CuTeGEMMBase(FusedOperation): """Base fused op for block-scaled GroupedLinear + activation + GroupedLinear. @@ -208,11 +215,22 @@ def fuser_forward( split_sizes = fc1_split_sizes if int(split_sizes.numel()) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") - split_sizes = split_sizes.to(dtype=torch.int64, device=device) - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) - split_points = base_split_offsets[1:].to(dtype=torch.int) - fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] - fc2_x_tensor_offsets = base_split_offsets * fc2_weight_shape[1] + + # Prepare split metadata + split_sizes, ( + split_points, + base_split_offsets, + fc1_x_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=[1, 1, fc1_weight_shape[1], fc2_weight_shape[1], fc2_weight_shape[0]], + include_leading_zero=[False, True, True, True, True], + dtypes=[torch.int32, torch.int64, torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) # Extract per-row activation probabilities from the middle op. scales = basic_op_extra_inputs[1][0] @@ -387,21 +405,37 @@ def fuser_forward( if use_nvfp4: nvfp4_fp4_max = 6.0 nvfp4_fp8_max = 448.0 + nvfp4_global_scale_denom = nvfp4_fp4_max * nvfp4_fp8_max + # cuDNN receives NVFP4 block-scaled inputs without TE's per-group + # global scale factors, so alpha supplies the product of the two + # operand global scales. fc1_alpha_tensor = ( _nvfp4_amax(grouped_fc1_x, columnwise=False) * _nvfp4_amax(grouped_fc1_weight, columnwise=False) - / (nvfp4_fp4_max**2 * nvfp4_fp8_max**2) + / (nvfp4_global_scale_denom**2) ).to(torch.float32) else: fc1_alpha_tensor = alpha_tensor + use_tmem_post_rht_amax = _use_tmem_post_rht_amax() + use_fc1_glu_hadamard = False + use_nvfp4_rht_amax = ( + use_nvfp4 + and isinstance(fc2_input_quantizer, NVFP4Quantizer) + and fc2_input_quantizer.with_rht + and fc2_input_quantizer.with_post_rht_amax + ) + if use_nvfp4_rht_amax and self._cudnn_act_func == "swiglu": + kernel_getter = getattr(self, "grouped_gemm_glu_hadamard_kernel", None) + if kernel_getter is not None: + use_fc1_glu_hadamard = kernel_getter() is not None + fc1_activation_kwargs = { "a_tensor": fc1_x_data, "sfa_tensor": fc1_x_scales, "padded_offsets": split_points, "alpha_tensor": fc1_alpha_tensor, "bias_tensor": fc1_bias_packed, - "norm_const_tensor": fc1_norm_const_tensor, "prob_tensor": fc1_prob_tensor, "acc_dtype": torch.float32, "c_dtype": torch.bfloat16, @@ -409,11 +443,15 @@ def fuser_forward( "cd_major": "n", "sf_vec_size": sf_vec_size, "current_stream": current_stream, - "discrete_col_sfd": not use_nvfp4, "use_dynamic_sched": True, } if self._cudnn_act_func is not None: fc1_activation_kwargs["act_func"] = self._cudnn_act_func + if use_fc1_glu_hadamard: + fc1_activation_kwargs["use_tmem_post_rht_amax"] = use_tmem_post_rht_amax + else: + fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor + fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4 if self._pass_geglu_runtime_params: fc1_activation_kwargs.update( linear_offset=self._cudnn_linear_offset, @@ -449,15 +487,13 @@ def fuser_forward( fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales else: # Discrete-weight kernel: per-expert data/scale pointers - fc1_b_ptrs = tex.copy_data_ptrs_to_device( - [w._rowwise_data for w in grouped_fc1_weight], - device, - ) - swizzle_type = "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_rowwise_swizzle" - fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - swizzle_type, - [w._rowwise_scale_inv for w in grouped_fc1_weight], - device, + fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc1_weight], + [w._rowwise_scale_inv for w in grouped_fc1_weight], + "nvfp4" if use_nvfp4 else "mxfp8_rowwise", + device, + ) ) fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs @@ -465,7 +501,10 @@ def fuser_forward( fc1_activation_kwargs["b_dtype"] = data_dtype fc1_activation_kwargs["b_major"] = "k" - fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + if use_fc1_glu_hadamard: + fc1_kernel_out = self.grouped_gemm_glu_hadamard_kernel()(**fc1_activation_kwargs) + else: + fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) # Unpack kernel outputs # Note: Fused kernel outputs tensors with non-contiguous @@ -497,13 +536,24 @@ def fuser_forward( fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) fc2_input_quantizer.optimize_for_gemm = True - grouped_fc2_x = _group_quantize_for_grouped_mlp( - fc2_in, - fc2_input_quantizer, - num_groups, - split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - ) + if use_fc1_glu_hadamard: + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + fc1_kernel_out["amax_tensor"].view(-1), + fc1_kernel_out["post_rht_amax_tensor"].view(-1), + tensor_offsets=fc2_x_tensor_offsets, + ) + else: + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + ) fc2_out_buf = torch.empty(fc2_out_shape, dtype=dtype, device=device) if ( @@ -537,7 +587,6 @@ def fuser_forward( else: fc2_out_buf = fc2_out_buf + token_bias else: - fc2_out_offsets = base_split_offsets * fc2_weight_shape[0] fc2_out_grouped = GroupedTensor( shape=(in_shape[0], fc2_weight_shape[0]), dtype=dtype, @@ -545,7 +594,7 @@ def fuser_forward( quantizer=None, data=fc2_out_buf.view(-1), first_dims=split_sizes, - tensor_offsets=fc2_out_offsets, + tensor_offsets=fc2_out_tensor_offsets, ) general_grouped_gemm_for_grouped_tensor( grouped_fc2_weight, @@ -625,17 +674,13 @@ def fuser_forward( fc2_quant_kwargs["b_tensor"] = fc2_w_data fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs = tex.copy_data_ptrs_to_device( - [w._rowwise_data for w in grouped_fc2_weight], - device, - ) - swizzle_type = ( - "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_rowwise_swizzle" - ) - fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - swizzle_type, - [w._rowwise_scale_inv for w in grouped_fc2_weight], - device, + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc2_weight], + [w._rowwise_scale_inv for w in grouped_fc2_weight], + "nvfp4" if use_nvfp4 else "mxfp8_rowwise", + device, + ) ) fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs @@ -745,6 +790,19 @@ def grouped_gemm_activation_kernel(cls) -> Callable: return grouped_gemm_glu_wrapper_sm100 + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_glu_hadamard_kernel(cls) -> Optional[Callable]: + """Fused grouped GEMM GLU kernel that also emits NVFP4 RHT amaxes.""" + try: + from cudnn import ( + grouped_gemm_glu_hadamard_wrapper_sm100, + ) # pylint: disable=no-name-in-module,import-outside-toplevel + except ImportError: + return None + + return grouped_gemm_glu_hadamard_wrapper_sm100 + class ForwardGroupedMLP_CuTeGEMMUnary(_ForwardGroupedMLP_CuTeGEMMBase): """Fused op for block-scaled GroupedLinear + scaled unary activation + GroupedLinear.""" @@ -801,8 +859,6 @@ def fuse_forward_srelu_ops( ) -> list[FusibleOperation]: """Apply GroupedLinear + ScaledSReLU + GroupedLinear fusion for forward pass.""" - if recipe is None or not recipe.mxfp8(): - return ops return fuse_grouped_mlp_ops( ops, recipe=recipe, diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 7d67815f9ac..c1070e38a62 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -12,12 +12,13 @@ from transformer_engine_torch import CommOverlapType, bulk_overlap_ag_with_external_gemm from ...cpp_extensions import general_gemm -from ...distributed import get_distributed_world_size +from ...distributed import gather_along_first_dim, get_distributed_world_size from ...module.base import ( _2X_ACC_DGRAD, _2X_ACC_WGRAD, fill_userbuffers_buffer_for_all_gather, get_ub, + using_cublasmp_backend, ) from ...quantized_tensor import Quantizer from ...tensor.mxfp8_tensor import MXFP8Quantizer @@ -234,7 +235,10 @@ def _functional_backward( # Note: Communication patterns are (1) overlap dy all-gather # with dgrad GEMM, (2) overlap x all-gather with dgrad GEMM # and dx reduce-scatter with wgrad GEMM, (3) overlap dx - # reduce-scatter with dgrad GEMM + # reduce-scatter with dgrad GEMM. The cuBLASMp backend does + # not support bulk overlaps; we use (3) for column-parallel + # and gather x separately for wgrad. + is_cublasmp = using_cublasmp_backend() ub_comm_dgrad = None ub_comm_wgrad = None ub_type_dgrad = None @@ -249,7 +253,7 @@ def _functional_backward( ub_type_dgrad = CommOverlapType.AG with_dgrad_all_gather_dy = True elif tensor_parallel_mode == "column": - if input_requires_grad and weight_requires_grad: + if input_requires_grad and weight_requires_grad and not is_cublasmp: with_bulk_overlap = True ub_comm_dgrad = get_ub(ub_comm_name + "_dgrad", with_quantized_compute) ub_type_dgrad = CommOverlapType.AG @@ -265,7 +269,9 @@ def _functional_backward( ub_comm_dgrad = get_ub(ub_comm_name + "_dgrad", with_quantized_compute) ub_type_dgrad = CommOverlapType.RS with_dgrad_reduce_scatter_dx = True - if ub_comm_dgrad.is_fp8_ubuf(): + # cuBLASMp manages its own buffers, so the FP8 buffer + # restriction does not apply. + if not is_cublasmp and ub_comm_dgrad.is_fp8_ubuf(): raise RuntimeError( "Userbuffers reduce-scatter is not supported with FP8 buffers" ) @@ -315,17 +321,50 @@ def _functional_backward( # Cast input tensor dtype if needed x_local = None + cublasmp_column_tp = ( + is_cublasmp and tensor_parallel_mode == "column" and weight_requires_grad + ) if weight_requires_grad: if input is None: raise ValueError("Input tensor is required to compute weight grad") x_local = input if with_quantized_compute: if not is_quantized_tensor(x_local): - input_quantizer.set_usage(columnwise=True) + if cublasmp_column_tp: + # The upcoming all-gather needs rowwise FP8 (transposes + # are computed during the gather) or columnwise MXFP8 + # (rowwise->columnwise can't be done after gather). + if isinstance(input_quantizer, MXFP8Quantizer): + input_quantizer.set_usage(rowwise=False, columnwise=True) + else: + input_quantizer.set_usage(rowwise=True, columnwise=False) + else: + input_quantizer.set_usage(columnwise=True) x_local = input_quantizer(x_local) else: x_local = maybe_dequantize(x_local, dtype) + # cuBLASMp does not support bulk overlap so we start an async gather of x_local here + # to overlap with the dgrad GEMM+RS in the column-parallel case. + x_total = None + x_total_work = None + if cublasmp_column_tp: + quantizer = None + if with_quantized_compute: + quantizer = input_quantizer + if isinstance(quantizer, MXFP8Quantizer): + quantizer.set_usage(rowwise=False, columnwise=True) + else: + # FP8 per-tensor: gather rowwise + create the columnwise + # transpose during the gather. + quantizer.set_usage(rowwise=True, columnwise=True) + x_total, x_total_work = gather_along_first_dim( + x_local, + tensor_parallel_group, + async_op=True, + quantizer=quantizer, + ) + # dgrad GEMM dx_local = None dx = None @@ -394,7 +433,10 @@ def _functional_backward( extra_output=dx_local if with_dgrad_reduce_scatter_dx else None, bulk_overlap=with_bulk_overlap, ) - if not (with_dgrad_reduce_scatter_dx or with_wgrad_reduce_scatter_dx): + + # cuBLASMp's GEMM+RS writes the reduce-scattered output directly into the GEMM output + # tensor; the extra-output buffer is only used with Userbuffers. + if is_cublasmp or not (with_dgrad_reduce_scatter_dx or with_wgrad_reduce_scatter_dx): dx_local = dx # wgrad GEMM @@ -402,7 +444,26 @@ def _functional_backward( if weight_requires_grad: # Initialize grad output - if tensor_parallel_mode == "row" and isinstance(grad_output_quantizer, MXFP8Quantizer): + if is_cublasmp and tensor_parallel_mode == "row": + # cuBLASMp's AG+GEMM DGRAD does not preserve the gathered dy ensor. Re-gather + # dy_local for wgrad with the appropriate quantization direction for each recipe. + if grad_output_quantizer is not None: + if isinstance(grad_output_quantizer, MXFP8Quantizer): + # MXFP8 can't convert rowwise to columnwise, so gather + # columnwise data directly. + grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + else: + # FP8 per-tensor: gather rowwise + create the columnwise + # transpose during the gather. + grad_output_quantizer.set_usage(rowwise=True, columnwise=True) + dy, _ = gather_along_first_dim( + dy_local, + tensor_parallel_group, + quantizer=grad_output_quantizer, + ) + elif tensor_parallel_mode == "row" and isinstance( + grad_output_quantizer, MXFP8Quantizer + ): # UB does not support pipelined overlapping grad output # all-gather with wgrad GEMM. Also, we can't # convert row-scaled MXFP8 to column-scaled, so we @@ -445,6 +506,12 @@ def _functional_backward( # Initialize input tensor if tensor_parallel_mode == "row": x = x_local + elif x_total is not None: + # cuBLASMp+column TP: use the async-gathered x_total. + if x_total_work is not None: + x_total_work.wait() + x_total_work = None + x = x_total if x is None: raise RuntimeError( "wgrad GEMM requires input tensor, which has not been initialized" diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 84073be6f8e..3a8ff5438de 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -251,7 +251,9 @@ def _functional_forward( extra_output=reduce_scatter_output, ) if with_ub_reduce_scatter: - y_local = reduce_scatter_output + # cuBLASMp writes the reduce-scattered output directly into the GEMM + # output tensor; Userbuffers writes it into the extra-output buffer. + y_local = gemm_output if ub_comm.with_cublasmp() else reduce_scatter_output else: y_local = gemm_output diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 828c34f5397..b48953e8e3a 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -15,6 +15,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +from ..constants import DType from .multi_tensor_apply import multi_tensor_applier @@ -422,7 +423,7 @@ def _initialize_state( quantizer = Float8Quantizer( scale=torch.ones([1], dtype=torch.float32, device=param.device), amax=torch.zeros([1], dtype=torch.float32, device=param.device), - fp8_dtype=tex.DType.kFloat8E4M3, + fp8_dtype=DType.kFloat8E4M3, ) self.state[param][state_name] = quantizer.make_empty(data.shape) self.state[param][state_name].quantize_(data.float()) @@ -599,7 +600,7 @@ def step(self, closure=None, grad_scaler=None): state_scales = {"exp_avg": [], "exp_avg_sq": [], "master_param": []} # Only used when extra params include fp8 tensors. Otherwise, it doesn't matter what the out_dtype is. - out_dtype = tex.DType.kFloat32 + out_dtype = DType.kFloat32 has_fp16 = False has_bf16 = False diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index e503b4b560e..d0c241288ce 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -26,7 +26,7 @@ NVFP4BlockScaling, CustomRecipe, ) -from .constants import dist_group_type +from .constants import dist_group_type, DType from .utils import get_device_compute_capability from .jit import jit_fuser @@ -277,23 +277,23 @@ def get_fp8_torch_dtype(fp8_recipe: Recipe, fprop_tensor: bool = True) -> torch. return torch.float8_e5m2 -def get_fp8_te_dtype(fp8_recipe: Recipe, fprop_tensor: bool = True) -> tex.DType: +def get_fp8_te_dtype(fp8_recipe: Recipe, fprop_tensor: bool = True) -> DType: """Get fp8 data type according to recipe and tensor""" if fp8_recipe.fp8_format == Format.E4M3 or ( fp8_recipe.fp8_format == Format.HYBRID and fprop_tensor ): - return tex.DType.kFloat8E4M3 - return tex.DType.kFloat8E5M2 + return DType.kFloat8E4M3 + return DType.kFloat8E5M2 -def get_fp4_te_dtype(fp4_recipe: Recipe) -> tex.DType: +def get_fp4_te_dtype(fp4_recipe: Recipe) -> DType: """Get fp4 data type according to recipe and tensor""" if fp4_recipe.fp4_format == Format.E2M1: - return tex.DType.kFloat4E2M1 + return DType.kFloat4E2M1 raise ValueError(f"Unsupported FP4 format: {fp4_recipe.fp4_format}") -def get_fp8_max(fp8_recipe: Recipe, fprop_tensor: bool = True) -> tex.DType: +def get_fp8_max(fp8_recipe: Recipe, fprop_tensor: bool = True) -> DType: """Get max representible FP8 value.""" if fp8_recipe.fp8_format == Format.E4M3 or ( fp8_recipe.fp8_format == Format.HYBRID and fprop_tensor @@ -1382,7 +1382,7 @@ class DelayedScalingRecipeState(RecipeState): recipe: DelayedScaling mode: str - dtype: tex.DType + dtype: DType scale: torch.Tensor amax_history: torch.Tensor @@ -1436,7 +1436,7 @@ class Float8CurrentScalingRecipeState(RecipeState): recipe: Float8CurrentScaling mode: str - dtype: tex.DType + dtype: DType device: torch.device def __init__( @@ -1480,7 +1480,7 @@ class MXFP8BlockScalingRecipeState(RecipeState): recipe: MXFP8BlockScaling mode: str - dtype: tex.DType + dtype: DType def __init__( self, @@ -1518,9 +1518,9 @@ class Float8BlockScalingRecipeState(RecipeState): recipe: Float8BlockScaling mode: str - qx_dtype: tex.DType - qw_dtype: tex.DType - qgrad_dtype: tex.DType + qx_dtype: DType + qw_dtype: DType + qgrad_dtype: DType def __init__( self, @@ -1605,7 +1605,7 @@ class NVFP4BlockScalingRecipeState(RecipeState): recipe: NVFP4BlockScaling mode: str - dtype: tex.DType + dtype: DType def __init__( self, diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index b56b1cd5eb5..667dee6872e 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -13,12 +13,58 @@ - Only cast to low-precision when necessary and the casting only happens in writing to global memory. For example, the gradient is required to have the same dtype as the input. """ -from typing import Optional +from typing import Optional, Union import torch import transformer_engine_torch as tex +# Re-export the C++ enum NVTERoutingMapFormat under a friendlier Python name. +# Members: +# RoutingMapFormat.BYTEMAP — bool[num_tokens, num_experts] +# RoutingMapFormat.BITMAP_U8 — uint8[num_tokens, ceil(num_experts/8)], +# LSB-first / little-endian packing along the +# expert axis. +RoutingMapFormat = tex.NVTERoutingMapFormat + + +_ROUTING_MAP_FORMAT_FROM_STRING = { + "bytemap": int(RoutingMapFormat.BYTEMAP), + "bitmap_u8": int(RoutingMapFormat.BITMAP_U8), +} +_VALID_ROUTING_MAP_FORMAT_INTS = frozenset(_ROUTING_MAP_FORMAT_FROM_STRING.values()) + + +def _validate_routing_map_format(routing_map_format: Union[str, RoutingMapFormat, int]) -> int: + """Coerce user-supplied routing_map_format into a plain int (0 or 1). + + Accepts the enum, an int matching one of the enum's values, or the + canonical lowercase strings ``"bytemap"`` / ``"bitmap_u8"``. + """ + # NVTERoutingMapFormat is a standalone pybind11 enum, not a subclass of int. + if isinstance(routing_map_format, RoutingMapFormat): + return int(routing_map_format) + if isinstance(routing_map_format, int): + if routing_map_format in _VALID_ROUTING_MAP_FORMAT_INTS: + return routing_map_format + raise ValueError( + f"routing_map_format int must be one of {sorted(_VALID_ROUTING_MAP_FORMAT_INTS)}; " + f"got {routing_map_format!r}" + ) + if isinstance(routing_map_format, str): + val = _ROUTING_MAP_FORMAT_FROM_STRING.get(routing_map_format) + if val is not None: + return val + raise ValueError( + "routing_map_format string must be 'bytemap' or 'bitmap_u8' (lowercase); " + f"got {routing_map_format!r}" + ) + raise TypeError( + "routing_map_format must be an int, a RoutingMapFormat enum, or 'bytemap' / " + f"'bitmap_u8'; got {routing_map_format!r}" + ) + + class FusedTopkScoreFunction(torch.autograd.Function): """ Fused Topk with Score Function router. @@ -36,14 +82,9 @@ def forward( scaling_factor: Optional[float], score_function: str, expert_bias: Optional[torch.Tensor], + routing_map_format: int, ): # pylint: disable=missing-function-docstring - # Save the shape of the logits - tensor_shape = logits.shape - logits = logits.view(-1, tensor_shape[-1]) - # Get the metadata of the viewed logits - num_tokens = logits.size(0) - num_experts = logits.size(1) probs, routing_map, intermediate_output = tex.fused_topk_with_score_function_fwd( logits, topk, @@ -53,33 +94,24 @@ def forward( scaling_factor, score_function, expert_bias, + routing_map_format, ) - # Restore the shape - probs = probs.view(tensor_shape) ctx.save_for_backward(routing_map, intermediate_output) - ctx.num_tokens = num_tokens - ctx.num_experts = num_experts ctx.use_pre_softmax = use_pre_softmax ctx.topk = topk ctx.scaling_factor = scaling_factor ctx.score_function = score_function - ctx.logits_dtype = logits.dtype + ctx.routing_map_format = routing_map_format return probs, routing_map @staticmethod def backward(ctx, grad_probs, _): # pylint: disable=missing-function-docstring routing_map, intermediate_output = ctx.saved_tensors - # Save the shape of the grad_probs - tensor_shape = grad_probs.shape - # Adjust the shape of the grad_probs to 2D shape - grad_probs = grad_probs.contiguous().view(-1, tensor_shape[-1]) - grad_logits = torch.empty( - (ctx.num_tokens, ctx.num_experts), dtype=ctx.logits_dtype, device=grad_probs.device - ) + if not grad_probs.is_contiguous(): + grad_probs = grad_probs.contiguous() + grad_logits = torch.empty_like(grad_probs) tex.fused_topk_with_score_function_bwd( - ctx.num_tokens, - ctx.num_experts, routing_map, intermediate_output, grad_probs, @@ -88,10 +120,9 @@ def backward(ctx, grad_probs, _): ctx.use_pre_softmax, ctx.scaling_factor, ctx.score_function, + ctx.routing_map_format, ) - # Restore the shape - grad_logits = grad_logits.view(tensor_shape) - return grad_logits, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None def fused_topk_with_score_function( @@ -103,6 +134,7 @@ def fused_topk_with_score_function( scaling_factor: Optional[float], score_function: str, expert_bias: Optional[torch.Tensor], + routing_map_format: Union[str, RoutingMapFormat, int] = RoutingMapFormat.BYTEMAP, ): """ Fused topk with score function router. @@ -121,14 +153,27 @@ def fused_topk_with_score_function( currently support "softmax", "sigmoid" and "sqrtsoftplus". expert_bias : torch.Tensor, optional could be used with the sigmoid/sqrtsoftplus score functions. + routing_map_format : Union[str, RoutingMapFormat, int], optional + Output layout for routing_map. ``"bytemap"`` / ``RoutingMapFormat.BYTEMAP`` + (default) returns a bool[T, E] tensor; ``"bitmap_u8"`` / + ``RoutingMapFormat.BITMAP_U8`` returns a uint8[T, ceil(E/8)] tensor with + bit ``(e % 8)`` of byte ``(e / 8)`` set when token ``t`` routes to expert + ``e`` (LSB-first / little-endian packing along the expert axis). Returns ------- probs : torch.Tensor in the same dtype as the "logits". - routing_map : torch.Tensor in bool. + Same shape as ``logits``. + routing_map : torch.Tensor + Same leading dims as ``logits``; trailing dim and dtype depend on + routing_map_format: + - BYTEMAP: bool[*logits.shape[:-1], num_experts] + - BITMAP_U8: uint8[*logits.shape[:-1], ceil(num_experts/8)] + LSB-first bit-packed. """ if logits.dtype == torch.float64: raise ValueError("Current TE does not support float64 router type.") + routing_map_format = _validate_routing_map_format(routing_map_format) return FusedTopkScoreFunction.apply( logits, topk, @@ -138,6 +183,7 @@ def fused_topk_with_score_function( scaling_factor, score_function, expert_bias, + routing_map_format, ) @@ -152,24 +198,20 @@ def forward( logits: torch.Tensor, topk: int, score_function: str, + routing_map_format: int, ): # pylint: disable=missing-function-docstring - # Save the shape of the logits - tensor_shape = logits.shape - logits = logits.view(-1, tensor_shape[-1]) - # Get the metadata of the viewed logits - num_tokens = logits.size(0) - num_experts = logits.size(1) scores, routing_map, intermediate_output = tex.fused_score_for_moe_aux_loss_fwd( logits=logits, topk=topk, score_function=score_function, + routing_map_format=routing_map_format, ) ctx.save_for_backward(intermediate_output) ctx.topk = topk ctx.score_function = score_function - ctx.num_tokens = num_tokens - ctx.num_experts = num_experts + # scores is FP32 but logits/grad_logits may be bf16/fp16 — remember the + # input dtype for the backward allocation. ctx.logits_dtype = logits.dtype return routing_map, scores @@ -177,31 +219,26 @@ def forward( def backward(ctx, _, grad_scores): # pylint: disable=missing-function-docstring intermediate_output = ctx.saved_tensors[0] - # Save the shape of the grad_scores - tensor_shape = grad_scores.shape - # Adjust the shape of the grad_scores to 2D shape - grad_scores = grad_scores.contiguous().view(-1, tensor_shape[-1]) + if not grad_scores.is_contiguous(): + grad_scores = grad_scores.contiguous() grad_logits = torch.empty( - (ctx.num_tokens, ctx.num_experts), dtype=ctx.logits_dtype, device=grad_scores.device + grad_scores.shape, dtype=ctx.logits_dtype, device=grad_scores.device ) tex.fused_score_for_moe_aux_loss_bwd( - num_tokens=ctx.num_tokens, - num_experts=ctx.num_experts, intermediate_output=intermediate_output, grad_scores=grad_scores, grad_logits=grad_logits, topk=ctx.topk, score_function=ctx.score_function, ) - # Restore the shape - grad_logits = grad_logits.view(tensor_shape) - return grad_logits, None, None + return grad_logits, None, None, None def fused_compute_score_for_moe_aux_loss( logits: torch.Tensor, topk: int, score_function: str, + routing_map_format: Union[str, RoutingMapFormat, int] = RoutingMapFormat.BYTEMAP, ): """ Fused compute scores for MoE aux loss, subset of the fused_topk_with_score_function. @@ -211,13 +248,20 @@ def fused_compute_score_for_moe_aux_loss( topk : int score_function : str currently support "softmax", "sigmoid" and "sqrtsoftplus". + routing_map_format : Union[str, RoutingMapFormat, int], optional + Output layout for routing_map; see :func:`fused_topk_with_score_function`. Returns ------- - routing_map : torch.Tensor in bool + routing_map : torch.Tensor + Same leading dims as ``logits``; trailing dim and dtype depend on + routing_map_format (bool[..., num_experts] for BYTEMAP, + uint8[..., ceil(num_experts/8)] for BITMAP_U8). scores : torch.Tensor in fp32 + Same shape as ``logits``. """ - return FusedComputeScoresForMoEAuxLoss.apply(logits, topk, score_function) + routing_map_format = _validate_routing_map_format(routing_map_format) + return FusedComputeScoresForMoEAuxLoss.apply(logits, topk, score_function, routing_map_format) class FusedAuxLoss(torch.autograd.Function): diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index e091e27e597..ba46508d74d 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -11,11 +11,11 @@ import torch import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType 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 ..constants import DType from ..utils import devices_match, round_up_to_nearest_multiple aten = torch.ops.aten @@ -30,7 +30,7 @@ class Float8BlockQuantizer(Quantizer): """ - dtype: TE_DType + dtype: DType block_len: int amax_epsilon: float force_pow_2_scales: bool @@ -38,7 +38,7 @@ class Float8BlockQuantizer(Quantizer): def __init__( self, - fp8_dtype: TE_DType, + fp8_dtype: Union[DType, tex.DType], *, rowwise: bool, columnwise: bool, @@ -47,7 +47,7 @@ def __init__( block_scaling_dim: int = 2, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) - self.dtype = fp8_dtype + self.dtype = DType.cast(fp8_dtype) self.block_len = 128 self.force_pow_2_scales = force_pow_2_scales self.amax_epsilon = amax_epsilon @@ -230,8 +230,9 @@ class Float8BlockwiseQTensor(Float8BlockwiseQTensorStorage, QuantizedTensor): columnwise_scale_inv : Optional[torch.Tensor] FP32 dequantization scales in GEMM format for dequantizing columnwise_data. - fp8_dtype : transformer_engine_torch.DType, default = kFloat8E4M3 - FP8 format. + fp8_dtype : transformer_engine.pytorch.DType or transformer_engine_torch.DType, + optional, default = kFloat8E4M3. FP8 format. transformer_engine_torch.DType + is accepted for backward compatibility. quantizer : Quantizer - the Float8BlockQuantizer that quantized this tensor and holds configuration about quantization and dequantization modes. """ @@ -245,7 +246,7 @@ def __new__( rowwise_scale_inv: Optional[torch.Tensor], columnwise_data: Optional[torch.Tensor], columnwise_scale_inv: Optional[torch.Tensor], - fp8_dtype: TE_DType, + fp8_dtype: DType, quantizer: Quantizer, is_2D_scaled: bool, **kwargs, @@ -443,7 +444,7 @@ def _make_in_reduce_ex( rowwise_scale_inv: torch.Tensor, columnwise_data: torch.Tensor, columnwise_scale_inv: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: DType, dtype: torch.dtype, quantizer: Quantizer, is_2D_scaled: bool, @@ -667,7 +668,7 @@ def _make_float8_blockwise_tensor_in_reduce_ex( rowwise_scale_inv: torch.Tensor, columnwise_data: torch.Tensor, columnwise_scale_inv: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: DType, dtype: torch.dtype, quantizer: Quantizer, is_2D_scaled: bool, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 7842ccc127c..4de8d822175 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -9,7 +9,6 @@ import torch from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType from transformer_engine.common.recipe import ( DelayedScaling, @@ -20,7 +19,7 @@ from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer from ._quantization_helpers import _IdentityFunc -from ..constants import dist_group_type +from ..constants import dist_group_type, DType aten = torch.ops.aten @@ -54,13 +53,13 @@ class Float8Quantizer(Quantizer): """Max-abs value from last FP8 cast""" amax: torch.Tensor """FP8 datatype""" - dtype: TE_DType + dtype: DType def __init__( self, scale: torch.Tensor, amax: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: Union[DType, tex.DType], *, rowwise: bool = True, columnwise: bool = True, @@ -68,7 +67,7 @@ def __init__( super().__init__(rowwise=rowwise, columnwise=columnwise) self.scale = scale self.amax = amax - self.dtype = fp8_dtype + self.dtype = DType.cast(fp8_dtype) def copy(self) -> Float8Quantizer: """Create shallow copy""" @@ -205,7 +204,7 @@ class Float8CurrentScalingQuantizer(Quantizer): """ """FP8 datatype""" - dtype: TE_DType + dtype: DType """amax reduction options""" with_amax_reduction: bool amax_reduction_group: Optional[dist_group_type] @@ -215,7 +214,7 @@ class Float8CurrentScalingQuantizer(Quantizer): def __init__( self, - fp8_dtype: TE_DType, + fp8_dtype: Union[DType, tex.DType], device: torch.device, *, rowwise: bool = True, @@ -237,7 +236,7 @@ def __init__( stacklevel=2, ) del device, use_existing_amax, scale, amax # Kept for backward compatibility - self.dtype = fp8_dtype + self.dtype = DType.cast(fp8_dtype) self.with_amax_reduction = with_amax_reduction self.amax_reduction_group = amax_reduction_group self.force_pow_2_scales = force_pow_2_scales @@ -402,8 +401,9 @@ class Float8Tensor(Float8TensorStorage, QuantizedTensor): Reciprocal of the scaling factor applied when casting to FP8, i.e. the scaling factor that must be applied when casting from FP8 to higher precision. - fp8_dtype : transformer_engine_torch.DType - FP8 format. + fp8_dtype : transformer_engine.pytorch.DType or transformer_engine_torch.DType + optional, default = kFloat8E4M3 FP8 format. transformer_engine_torch.DType + is accepted for backward compatibility. data_transpose : torch.Tensor, optional FP8 transpose data in a uint8 tensor quantizer : Float8Quantizer, Float8CurrentScalingQuantizer, optional @@ -922,7 +922,7 @@ def __reduce_ex__(self, protocol: int) -> tuple: def _make_in_reduce_ex( cls, data: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: DType, fp8_scale_inv: torch.Tensor, dtype: torch.dtype, shape: torch.Size, @@ -1005,7 +1005,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: def _make_float8_tensor_in_reduce_ex( data: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: DType, fp8_scale_inv: torch.Tensor, dtype: torch.dtype, shape: torch.Size, diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 2815aaa96e8..d759aaf5c4f 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -12,10 +12,9 @@ import torch from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe -from ..constants import MXFP8_BLOCK_SCALING_SIZE +from ..constants import MXFP8_BLOCK_SCALING_SIZE, DType from ..utils import devices_match, round_up_to_nearest_multiple from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func from ..quantized_tensor import QuantizedTensor, Quantizer @@ -33,17 +32,17 @@ class MXFP8Quantizer(Quantizer): """ - dtype: TE_DType + dtype: DType def __init__( self, - fp8_dtype: TE_DType, + fp8_dtype: Union[DType, tex.DType], *, rowwise: bool = True, columnwise: bool = True, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) - self.dtype = fp8_dtype + self.dtype = DType.cast(fp8_dtype) def copy(self) -> MXFP8Quantizer: """Create shallow copy""" @@ -148,7 +147,7 @@ def create_tensor_from_data( data: torch.Tensor, scale_inv: torch.Tensor, fake_dtype: torch.dtype, - fp8_dtype: TE_DType = tex.DType.kFloat8E4M3, + fp8_dtype: DType = DType.kFloat8E4M3, ) -> MXFP8Tensor: """Create a new MXFP8Tensor from data and scale_inv.""" return MXFP8Tensor( @@ -193,8 +192,9 @@ class MXFP8Tensor(MXFP8TensorStorage, QuantizedTensor): ---------- data : torch.Tensor Raw FP8 data in a uint8 tensor - fp8_dtype : transformer_engine_torch.DType, default = kFloat8E4M3 - FP8 format. + fp8_dtype : transformer_engine.pytorch.DType or transformer_engine_torch.DType, + optional, default = kFloat8E4M3 FP8 format. transformer_engine_torch.DType + is accepted for backward compatibility. fp8_scale_inv : torch.Tensor Reciprocal of the scaling factor applied when casting to FP8, i.e. the scaling factor that must @@ -214,7 +214,7 @@ def __new__( rowwise_scale_inv: Optional[torch.Tensor], columnwise_data: Optional[torch.Tensor], columnwise_scale_inv: Optional[torch.Tensor], - fp8_dtype: TE_DType, + fp8_dtype: DType, quantizer: Optional[Quantizer], with_gemm_swizzled_scales: bool, **kwargs, @@ -727,7 +727,7 @@ def _make_in_reduce_ex( rowwise_scale_inv: torch.Tensor, columnwise_data: torch.Tensor, columnwise_scale_inv: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: DType, dtype: torch.dtype, shape: torch.Size, quantizer: Optional[Quantizer] = None, @@ -849,7 +849,7 @@ def _make_mxfp8_tensor_in_reduce_ex( rowwise_scale_inv: torch.Tensor, columnwise_data: torch.Tensor, columnwise_scale_inv: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: DType, dtype: torch.dtype, shape: torch.Size, quantizer: Optional[Quantizer] = None, diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 24962d67f20..5a2765b9f54 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -12,10 +12,9 @@ import torch import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe -from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type +from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType from ..utils import ( canonicalize_process_group, devices_match, @@ -114,7 +113,7 @@ def get_random_sign_mask_for_rht(with_random_sign_mask: bool, device: int) -> in class NVFP4Quantizer(Quantizer): """Builder class for NVFP4 tensors with NV block scaling""" - dtype: TE_DType + dtype: DType """Random Hadamard Transform""" with_rht: bool with_post_rht_amax: bool @@ -143,7 +142,7 @@ class NVFP4Quantizer(Quantizer): def __init__( self, - fp4_dtype: TE_DType = tex.DType.kFloat4E2M1, + fp4_dtype: Union[DType, tex.DType] = DType.kFloat4E2M1, rowwise: bool = True, columnwise: bool = True, with_amax_reduction: bool = False, @@ -159,7 +158,7 @@ def __init__( with_random_sign_mask: bool = True, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) - self.dtype = fp4_dtype + self.dtype = DType.cast(fp4_dtype) self.with_rht = with_rht self.with_post_rht_amax = with_post_rht_amax self.with_amax_reduction = with_amax_reduction @@ -352,7 +351,7 @@ class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): Rowwise amax tracking tensor. amax_columnwise : torch.Tensor, optional Columnwise amax tracking tensor. - fp4_dtype : TE_DType + fp4_dtype : DType The FP4 data type used for quantization. quantizer : Quantizer The quantizer instance used for this tensor. @@ -371,7 +370,7 @@ def __new__( columnwise_scale_inv: Optional[torch.Tensor], amax_rowwise: Optional[torch.Tensor], amax_columnwise: Optional[torch.Tensor], - fp4_dtype: TE_DType, + fp4_dtype: DType, quantizer: Quantizer, with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, @@ -805,7 +804,7 @@ def _make_in_reduce_ex( columnwise_scale_inv: torch.Tensor, amax_rowwise: torch.Tensor, amax_columnwise: torch.Tensor, - fp4_dtype: TE_DType, + fp4_dtype: DType, dtype: torch.dtype, quantizer: Quantizer, with_gemm_swizzled_scales: bool = False, @@ -936,7 +935,7 @@ def _make_nvfp4_tensor_in_reduce_ex( columnwise_scale_inv: torch.Tensor, amax_rowwise: torch.Tensor, amax_columnwise: torch.Tensor, - fp4_dtype: TE_DType, + fp4_dtype: DType, dtype: torch.dtype, quantizer: Quantizer, with_gemm_swizzled_scales: bool, 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 ca3913762f9..f7a3dae70b6 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -6,15 +6,14 @@ from __future__ import annotations import math -from typing import Optional, Dict, Any, Tuple +from typing import Optional, Dict, Any, Tuple, Union import torch import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from ...constants import TE_DType_To_Torch +from ...constants import TE_DType_To_Torch, DType from ...utils import _empty_tensor @@ -31,7 +30,7 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): _rowwise_data: Optional[torch.Tensor] _columnwise_data: Optional[torch.Tensor] _quantizer: Quantizer - _fp8_dtype: TE_DType + _fp8_dtype: DType _rowwise_scale_inv: Optional[torch.Tensor] _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool @@ -42,7 +41,7 @@ def __new__( rowwise_scale_inv: Optional[torch.Tensor], columnwise_data: Optional[torch.Tensor], columnwise_scale_inv: Optional[torch.Tensor], - fp8_dtype: TE_DType, + fp8_dtype: Union[DType, tex.DType], quantizer: Quantizer, is_2D_scaled: bool, *args, @@ -57,7 +56,7 @@ def __new__( instance._rowwise_data = rowwise_data instance._columnwise_data = columnwise_data instance._quantizer = quantizer.copy() if quantizer is not None else None - instance._fp8_dtype = fp8_dtype + instance._fp8_dtype = DType.cast(fp8_dtype) instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv instance._is_2D_scaled = is_2D_scaled diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 3a72ec5d1a2..a97162f91c0 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -6,15 +6,14 @@ from __future__ import annotations import math -from typing import Any, Dict, Optional, Tuple +from typing import Any, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch +from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType from ...utils import is_non_tn_fp8_gemm_supported, _empty_tensor @@ -69,7 +68,7 @@ class Float8TensorStorage(QuantizedTensorStorage): _data: Optional[torch.Tensor] _quantizer: Optional[Quantizer] - _fp8_dtype: TE_DType + _fp8_dtype: DType _scale_inv: torch.Tensor # FP8 transpose cache @@ -81,7 +80,7 @@ def __new__( *args, data: Optional[torch.Tensor], fp8_scale_inv: torch.Tensor, - fp8_dtype: TE_DType, + fp8_dtype: Union[DType, tex.DType], fake_dtype: Optional[torch.dtype] = None, data_transpose: Optional[torch.Tensor] = None, quantizer: Optional[Quantizer] = None, @@ -94,7 +93,7 @@ def __new__( instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) instance._data = data instance._quantizer = quantizer.copy() if quantizer is not None else None - instance._fp8_dtype = fp8_dtype + instance._fp8_dtype = DType.cast(fp8_dtype) instance._scale_inv = fp8_scale_inv instance._transpose = data_transpose instance._transpose_invalid = instance._transpose is None diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 874555f465a..ea592cd9893 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -5,17 +5,16 @@ """Mixin class holding data specific for MXFP8Tensor""" from __future__ import annotations -from typing import Optional, Dict, Any, Tuple +from typing import Optional, Dict, Any, Tuple, Union from collections.abc import Iterable import math import torch import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from ...constants import TE_DType as torch_to_transformer_engine_dtype +from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType from ...utils import _empty_tensor @@ -78,7 +77,7 @@ class MXFP8TensorStorage(QuantizedTensorStorage): # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] # FP8 data type - _fp8_dtype: TE_DType + _fp8_dtype: DType # Whether scaling factors are in the swizzled format expected by # GEMM _with_gemm_swizzled_scales: bool @@ -89,7 +88,7 @@ def __new__( rowwise_scale_inv: Optional[torch.Tensor], columnwise_data: Optional[torch.Tensor], columnwise_scale_inv: Optional[torch.Tensor], - fp8_dtype: TE_DType, + fp8_dtype: Union[DType, tex.DType], quantizer: Optional[Quantizer], with_gemm_swizzled_scales: bool, *args, @@ -106,7 +105,7 @@ def __new__( instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv instance._quantizer = quantizer.copy() if quantizer is not None else None - instance._fp8_dtype = fp8_dtype + instance._fp8_dtype = DType.cast(fp8_dtype) instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales return instance diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 250fa6bdb29..53bb5e7c11e 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -14,11 +14,10 @@ import torch import transformer_engine_torch as tex -from transformer_engine_torch import DType as TE_DType from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from ...constants import TE_DType as torch_to_transformer_engine_dtype +from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType from ...utils import _empty_tensor @@ -98,7 +97,7 @@ class NVFP4TensorStorage(QuantizedTensorStorage): # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] # FP4 data type - _fp4_dtype: TE_DType + _fp4_dtype: DType # Whether scaling factors are in the swizzled format expected by # GEMM _with_gemm_swizzled_scales: bool @@ -117,7 +116,7 @@ def __new__( columnwise_scale_inv: torch.Tensor, amax_rowwise: torch.Tensor, amax_columnwise: torch.Tensor, - fp4_dtype: TE_DType, + fp4_dtype: Union[DType, tex.DType], quantizer: Optional[Quantizer], with_gemm_swizzled_scales: bool, *args, @@ -135,7 +134,7 @@ def __new__( instance._rowwise_data = rowwise_data instance._columnwise_data = columnwise_data - instance._fp4_dtype = fp4_dtype + instance._fp4_dtype = DType.cast(fp4_dtype) instance._quantizer = quantizer.copy() if quantizer is not None else None instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 8b22097f7e0..1789a0c98e3 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -21,7 +21,7 @@ from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from ..optimizers.multi_tensor_apply import multi_tensor_applier from ..utils import is_non_tn_fp8_gemm_supported -from ..constants import NVFP4_BLOCK_SCALING_SIZE +from ..constants import NVFP4_BLOCK_SCALING_SIZE, DType def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): @@ -367,9 +367,9 @@ def _cast_master_weights_to_fp8_current_scaling( # --------------------------------------------------------------------------------------------- # Step 3: Update scales and scale_invs. # --------------------------------------------------------------------------------------------- - if fp8_dtype == tex.DType.kFloat8E4M3: + if fp8_dtype == DType.kFloat8E4M3: max_fp8 = 448.0 - elif fp8_dtype == tex.DType.kFloat8E5M2: + elif fp8_dtype == DType.kFloat8E5M2: max_fp8 = 57344.0 else: raise ValueError(f"Unsupported FP8 dtype: {fp8_dtype}") @@ -530,9 +530,9 @@ def _cast_master_weights_to_fp8_blockwise_scaling( # --------------------------------------------------------------------------------------------- # Step 3: Update scales and scale_invs. # --------------------------------------------------------------------------------------------- - if fp8_dtype == tex.DType.kFloat8E4M3: + if fp8_dtype == DType.kFloat8E4M3: max_fp8 = 448.0 - elif fp8_dtype == tex.DType.kFloat8E5M2: + elif fp8_dtype == DType.kFloat8E5M2: max_fp8 = 57344.0 else: raise ValueError(f"Unsupported FP8 dtype: {fp8_dtype}") diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index fd8f817b33e..cfb21e7bffe 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -17,7 +17,13 @@ from ..debug.pytorch.debug_quantization import DebugQuantizedTensor -__all__ = ["get_device_compute_capability", "get_cudnn_version", "is_bf16_available"] +__all__ = [ + "get_device_compute_capability", + "get_cudnn_version", + "is_bf16_available", + "deinterleave_glu_tensor", + "interleave_glu_tensor", +] @functools.lru_cache(maxsize=None) @@ -82,6 +88,103 @@ def get_device_compute_capability() -> Tuple[int, int]: return _get_device_compute_capability(torch.cuda.current_device()) +def deinterleave_glu_tensor(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Convert a block-interleaved GLU fc1 tensor to contiguous gate/linear layout. + + Fused GLU kernels (for example :class:`~transformer_engine.pytorch.ops.SwiGLU` + with ``glu_interleave_size`` set) expect fc1 weights in a block-interleaved + layout along dimension 0. Checkpoints and frameworks such as Megatron-LM typically + store the gate (``W``) and linear (``V``) halves as two contiguous blocks + ``[W_all, V_all]``. This helper reorders along dimension 0 without changing + the total shape. + + **Layouts along dimension 0** (``k = interleave_size``): + + * **Block-interleaved (input):** ``[W0:k, V0:k, Wk:2k, Vk:2k, ...]`` + * **Contiguous (output):** ``[W_all, V_all]`` + + The same convention applies to ``linear_fc1.weight`` (dimension 0 plus any + remaining dimensions) and ``linear_fc1.bias`` (dimension 0 only). + + Parameters + ---------- + tensor : torch.Tensor + Tensor in block-interleaved layout. The length of dimension 0 must be + divisible by ``2 * interleave_size``. + interleave_size : int + Number of rows (for weights) or elements (for bias) per gate/linear block. + Fused TE GLU paths commonly use ``32``. + + Returns + ------- + torch.Tensor + A new tensor with the same shape as ``tensor`` and contiguous + ``[W_all, V_all]`` ordering along dimension 0. + + See Also + -------- + :func:`interleave_glu_tensor` : Inverse transformation (contiguous -> block-interleaved). + """ + shape = tensor.shape + x = tensor.reshape( + shape[0] // (2 * interleave_size), + 2, + interleave_size, + *shape[1:], + ) + x = x.transpose(0, 1).contiguous() + return x.reshape(shape) + + +def interleave_glu_tensor(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Convert a contiguous SwiGLU fc1 tensor to block-interleaved layout. + + This is the inverse of :func:`deinterleave_glu_tensor`. Use it when loading + contiguous ``[W_all, V_all]`` checkpoints into a module that uses fused + interleaved SwiGLU (``glu_interleave_size`` on the activation op). + + **Layouts along dimension 0** (``k = interleave_size``): + + * **Contiguous (input):** ``[W_all, V_all]`` + * **Block-interleaved (output):** ``[W0:k, V0:k, Wk:2k, Vk:2k, ...]`` + + Parameters + ---------- + tensor : torch.Tensor + Tensor in contiguous gate/linear layout. The length of dimension 0 must be + divisible by ``2 * interleave_size``. + interleave_size : int + Number of rows (for weights) or elements (for bias) per gate/linear block. + Must match the ``glu_interleave_size`` used by the fused SwiGLU op. + + Returns + ------- + torch.Tensor + A new tensor with the same shape as ``tensor`` and block-interleaved + ordering along dimension 0. + + See Also + -------- + :func:`deinterleave_glu_tensor` : Inverse transformation (block-interleaved -> contiguous). + """ + if interleave_size <= 0: + raise ValueError(f"interleave_size must be a positive integer, got {interleave_size}") + if tensor.shape[0] % (2 * interleave_size) != 0: + raise ValueError( + f"tensor dimension 0 ({tensor.shape[0]}) must be divisible by " + f"2 * interleave_size ({2 * interleave_size})" + ) + shape = tensor.shape + x = tensor.reshape( + 2, + shape[0] // (2 * interleave_size), + interleave_size, + *shape[1:], + ) + x = x.transpose(0, 1).contiguous() + return x.reshape(shape) + + def resolve_grouped_linear_single_param_flags( single_grouped_weight: bool, single_grouped_bias: bool,