diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions index 705ca8eb38..9f47d6eb3b 160000 --- a/3rdparty/nccl-extensions +++ b/3rdparty/nccl-extensions @@ -1 +1 @@ -Subproject commit 705ca8eb38297f3a8c4af6adf4185176a1116cb9 +Subproject commit 9f47d6eb3b60962d8157a579b4caaaa4ae6b19f4 diff --git a/benchmarks/linear/benchmark_graph_safe_grouped_linear.py b/benchmarks/linear/benchmark_graph_safe_grouped_mlp.py similarity index 97% rename from benchmarks/linear/benchmark_graph_safe_grouped_linear.py rename to benchmarks/linear/benchmark_graph_safe_grouped_mlp.py index d8230c38fe..00f7f516c8 100644 --- a/benchmarks/linear/benchmark_graph_safe_grouped_linear.py +++ b/benchmarks/linear/benchmark_graph_safe_grouped_mlp.py @@ -13,21 +13,21 @@ Example: - python benchmarks/linear/benchmark_graph_safe_grouped_linear.py + python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py Forward-only: - python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --fwd-only + python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --fwd-only Nsight Systems: (optionally: unset DEBUGINFOD_URLS) nsys profile \ - --output=./benchmarks/linear/graph_safe_grouped_linear_mxfp8 \ + --output=./benchmarks/linear/graph_safe_grouped_mlp_mxfp8 \ --force-overwrite true \ --trace=cuda,nvtx,cudnn,cublas \ - python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --profile + python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --profile """ # Match the Qwen MXFP8 SFT launch toggles before importing TE. diff --git a/benchmarks/linear/benchmark_grouped_gemm_kernels.py b/benchmarks/linear/benchmark_grouped_gemm_kernels.py new file mode 100644 index 0000000000..88b3f8a466 --- /dev/null +++ b/benchmarks/linear/benchmark_grouped_gemm_kernels.py @@ -0,0 +1,581 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark multi-stream and cuBLASLt grouped GEMM kernels. + +Unlike ``benchmark_grouped_linear_cublas_grouped_gemm.py``, this benchmark does +not construct a GroupedLinear module or invoke autograd. Inputs are allocated +and, for MXFP8, quantized and GEMM-swizzled before timing. The timed region only +calls TE's low-level grouped GEMM wrappers. + +The default shape is Qwen3.5-397B-A17B with sequence length 4096, top-10 +routing, and EP32. Each rank owns 16 experts and processes 2560 rows per expert. +All six expert GEMMs are measured independently: + +============== ====== ====== ====== ====== ================================== +Projection Pass Layout M K N +============== ====== ====== ====== ====== ================================== +FC1 fwd TN 2560 4096 2048 +FC1 dgrad NN 2560 4096 2048 +FC1 wgrad NT 2560 4096 2048 +FC2 fwd TN 2560 1024 4096 +FC2 dgrad NN 2560 1024 4096 +FC2 wgrad NT 2560 1024 4096 +============== ====== ====== ====== ====== ================================== + +The layouts describe TE's operands for each expert: + +* Forward (TN): A=``weight[N,K]``, B=``input[M,K]``; + ``input @ weight.T -> output[M,N]``. +* Dgrad (NN): A=``weight[N,K]``, B=``dy[M,N]``; + ``dy @ weight -> dx[M,K]``. +* Wgrad (NT): A=``input[M,K]``, B=``dy[M,N]``; + ``dy.T @ input -> dweight[N,K]``. + +The multi-stream path receives lists of discrete tensors. The cuBLASLt grouped +path receives packed GroupedTensor activations and discrete weights/wgrads, +matching GroupedLinear with discrete parameters. + +Examples +-------- +Run the full BF16 and MXFP8 matrix: + + python benchmarks/linear/benchmark_grouped_gemm_kernels.py + +Run only MXFP8 FC1 wgrad: + + python benchmarks/linear/benchmark_grouped_gemm_kernels.py \ + --precision mxfp8 --projection fc1 --gemm wgrad +""" + +import argparse +import os +from dataclasses import dataclass +from typing import Any, Optional + +import pandas as pd +import torch +import torch.utils.benchmark as benchmark + +from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe +from transformer_engine.pytorch import MXFP8Quantizer +from transformer_engine.pytorch.cpp_extensions import ( + general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, +) +from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported +from transformer_engine.pytorch.module.base import ( + _2X_ACC_DGRAD, + _2X_ACC_FPROP, + _2X_ACC_WGRAD, +) +from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.tensor import GroupedTensor, GroupedTensorStorage +import transformer_engine_torch as tex + + +QWEN_NUM_EXPERTS = 512 +QWEN_TOP_K = 10 +QWEN_SEQUENCE_LENGTH = 4096 +QWEN_EXPERT_PARALLEL_SIZE = 32 +QWEN_HIDDEN_SIZE = 4096 +QWEN_MOE_INTERMEDIATE_SIZE = 1024 + + +@dataclass(frozen=True) +class GemmSpec: + """One expert GEMM in the Qwen routed MLP.""" + + projection: str + gemm: str + layout: str + k: int + n: int + + +GEMM_SPECS = tuple( + GemmSpec(projection, gemm, layout, k, n) + for projection, k, n in ( + ("fc1", QWEN_HIDDEN_SIZE, 2 * QWEN_MOE_INTERMEDIATE_SIZE), + ("fc2", QWEN_MOE_INTERMEDIATE_SIZE, QWEN_HIDDEN_SIZE), + ) + for gemm, layout in (("fwd", "TN"), ("dgrad", "NN"), ("wgrad", "NT")) +) + + +@dataclass +class PreparedGemm: + """Preallocated operands and outputs for both grouped GEMM paths.""" + + spec: GemmSpec + split_sizes: list[int] + multistream_a: list[Any] + multistream_b: list[Any] + multistream_out: list[torch.Tensor] + grouped_a: Any + grouped_b: GroupedTensorStorage + grouped_out: Any + + +def _make_grouped_bf16( + packed: Optional[torch.Tensor], + split_sizes: list[int], + last_dim: int, +) -> GroupedTensorStorage: + """Create packed GroupedTensor storage and optionally initialize its data.""" + first_dims = torch.tensor(split_sizes, dtype=torch.int64, device="cuda") + grouped = GroupedTensor.make_grouped_tensor( + num_tensors=len(split_sizes), + first_dims=first_dims, + last_dims=None, + logical_first_dim=sum(split_sizes), + logical_last_dim=last_dim, + quantizer=None, + device=torch.device("cuda"), + dtype=torch.bfloat16, + ) + if packed is not None: + grouped.rowwise_data.view(-1).copy_(packed.reshape(-1)) + return grouped + + +def _new_mxfp8_quantizer(*, rowwise: bool, columnwise: bool) -> MXFP8Quantizer: + """Construct a quantizer that emits GEMM-ready scales before timing.""" + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + quantizer.optimize_for_gemm = True + return quantizer + + +def _quantize_discrete_mxfp8( + tensors: list[torch.Tensor], + *, + rowwise: bool, + columnwise: bool, +) -> list[Any]: + """Quantize per-expert tensors outside the timed region.""" + quantizer = _new_mxfp8_quantizer(rowwise=rowwise, columnwise=columnwise) + return [quantizer(tensor) for tensor in tensors] + + +def _quantize_grouped_mxfp8( + packed: torch.Tensor, + split_sizes: list[int], + *, + rowwise: bool, + columnwise: bool, +) -> GroupedTensorStorage: + """Quantize one packed activation into GEMM-ready grouped storage.""" + quantizer = _new_mxfp8_quantizer(rowwise=rowwise, columnwise=columnwise) + first_dims = torch.tensor(split_sizes, dtype=torch.int64, device="cuda") + return tex.group_quantize( + packed, + quantizer, + len(split_sizes), + first_dims, + ) + + +def _make_high_precision_operands( + spec: GemmSpec, + split_sizes: list[int], +) -> tuple[list[torch.Tensor], Optional[torch.Tensor], torch.Tensor]: + """Create canonical BF16 operands shared by both precision paths.""" + num_experts = len(split_sizes) + total_rows = sum(split_sizes) + + if spec.gemm in ("fwd", "dgrad"): + weights = [ + torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + for _ in range(num_experts) + ] + b_width = spec.k if spec.gemm == "fwd" else spec.n + packed_b = torch.randn(total_rows, b_width, dtype=torch.bfloat16, device="cuda") + return weights, None, packed_b + + packed_x = torch.randn(total_rows, spec.k, dtype=torch.bfloat16, device="cuda") + packed_dy = torch.randn(total_rows, spec.n, dtype=torch.bfloat16, device="cuda") + return [], packed_x, packed_dy + + +def _operand_usage(spec: GemmSpec) -> tuple[tuple[bool, bool], tuple[bool, bool]]: + """Return rowwise/columnwise MXFP8 storage used by A and B.""" + if spec.gemm == "fwd": + return (True, False), (True, False) + if spec.gemm == "dgrad": + return (False, True), (True, False) + return (False, True), (False, True) + + +def _allocate_outputs( + spec: GemmSpec, + split_sizes: list[int], +) -> tuple[list[torch.Tensor], Any]: + """Allocate output storage matching each low-level API's native contract.""" + num_experts = len(split_sizes) + if spec.gemm == "wgrad": + multistream_out = [ + torch.empty(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + for _ in range(num_experts) + ] + grouped_out = [ + torch.empty(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + for _ in range(num_experts) + ] + return multistream_out, grouped_out + + out_features = spec.n if spec.gemm == "fwd" else spec.k + multistream_out = [ + torch.empty(sum(split_sizes), out_features, dtype=torch.bfloat16, device="cuda") + ] + grouped_out = _make_grouped_bf16(None, split_sizes, out_features) + return multistream_out, grouped_out + + +def _prepare_gemm( + spec: GemmSpec, + split_sizes: list[int], + precision: str, +) -> PreparedGemm: + """Prepare one GEMM without leaving quantization inside the timed region.""" + hp_a, packed_a, packed_b = _make_high_precision_operands(spec, split_sizes) + a_usage, b_usage = _operand_usage(spec) + + if spec.gemm in ("fwd", "dgrad"): + if precision == "mxfp8": + multistream_a = _quantize_discrete_mxfp8( + hp_a, + rowwise=a_usage[0], + columnwise=a_usage[1], + ) + else: + multistream_a = hp_a + grouped_a = multistream_a + elif precision == "mxfp8": + assert packed_a is not None + grouped_a = _quantize_grouped_mxfp8( + packed_a, + split_sizes, + rowwise=a_usage[0], + columnwise=a_usage[1], + ) + else: + assert packed_a is not None + grouped_a = _make_grouped_bf16(packed_a, split_sizes, spec.k) + if spec.gemm == "wgrad": + multistream_a = list(grouped_a.split_into_quantized_tensors()) + + if precision == "mxfp8": + grouped_b = _quantize_grouped_mxfp8( + packed_b, + split_sizes, + rowwise=b_usage[0], + columnwise=b_usage[1], + ) + else: + b_width = spec.k if spec.gemm == "fwd" else spec.n + grouped_b = _make_grouped_bf16(packed_b, split_sizes, b_width) + # Give multi-stream zero-copy member views into the exact same packed operand used by + # cuBLASLt. This prevents quantization or input-data differences from entering the result. + multistream_b = list(grouped_b.split_into_quantized_tensors()) + + multistream_out, grouped_out = _allocate_outputs(spec, split_sizes) + prepared = PreparedGemm( + spec=spec, + split_sizes=split_sizes, + multistream_a=multistream_a, + multistream_b=multistream_b, + multistream_out=multistream_out, + grouped_a=grouped_a, + grouped_b=grouped_b, + grouped_out=grouped_out, + ) + if precision == "mxfp8": + _assert_mxfp8_operands_are_gemm_ready(prepared) + return prepared + + +def _iter_operands(operand: Any): + """Yield discrete members or one packed operand.""" + if isinstance(operand, (list, tuple)): + yield from operand + else: + yield operand + + +def _assert_mxfp8_operands_are_gemm_ready(prepared: PreparedGemm) -> None: + """Ensure scale swizzling cannot leak into the GEMM timing.""" + operands = ( + prepared.multistream_a, + prepared.multistream_b, + prepared.grouped_a, + prepared.grouped_b, + ) + for operand_group in operands: + for operand in _iter_operands(operand_group): + assert getattr( + operand, "_with_gemm_swizzled_scales", False + ), "MXFP8 operand was not prepared with GEMM-swizzled scales" + + +def _use_split_accumulator(spec: GemmSpec) -> bool: + """Match GroupedLinear's accumulator policy for each training GEMM.""" + if spec.gemm == "fwd": + return _2X_ACC_FPROP + if spec.gemm == "dgrad": + return _2X_ACC_DGRAD + return _2X_ACC_WGRAD + + +def _run_multistream(prepared: PreparedGemm, iterations: int) -> None: + """Launch the multi-stream grouped GEMM repeatedly.""" + spec = prepared.spec + for _ in range(iterations): + general_grouped_gemm( + prepared.multistream_a, + prepared.multistream_b, + prepared.multistream_out, + [None] * len(prepared.split_sizes), + torch.bfloat16, + layout=spec.layout, + m_splits=prepared.split_sizes, + grad=spec.gemm != "fwd", + single_output=spec.gemm != "wgrad", + use_split_accumulator=_use_split_accumulator(spec), + ) + + +def _run_cublas_grouped(prepared: PreparedGemm, iterations: int) -> None: + """Launch the device-described cuBLASLt grouped GEMM repeatedly.""" + for _ in range(iterations): + general_grouped_gemm_for_grouped_tensor( + prepared.grouped_a, + prepared.grouped_b, + prepared.grouped_out, + layout=prepared.spec.layout, + use_split_accumulator=_use_split_accumulator(prepared.spec), + ) + + +def _canonical_output(prepared: PreparedGemm, *, grouped: bool) -> torch.Tensor: + """Return either path's output in one comparable packed tensor.""" + output = prepared.grouped_out if grouped else prepared.multistream_out + if prepared.spec.gemm == "wgrad": + return torch.stack(output, dim=0) + if grouped: + out_features = prepared.spec.n if prepared.spec.gemm == "fwd" else prepared.spec.k + return output.rowwise_data.view(sum(prepared.split_sizes), out_features) + return output[0] + + +def _validate_outputs(prepared: PreparedGemm, precision: str) -> None: + """Check that execution-path selection does not change GEMM numerics unexpectedly.""" + _run_multistream(prepared, 1) + _run_cublas_grouped(prepared, 1) + if precision == "mxfp8": + tolerances = {"rtol": 0.125, "atol": 0.0675} + else: + tolerances = {"rtol": 1e-2, "atol": 1e-2} + torch.testing.assert_close( + _canonical_output(prepared, grouped=True), + _canonical_output(prepared, grouped=False), + **tolerances, + ) + + +def _benchmark_path( + prepared: PreparedGemm, + *, + path: str, + iterations_per_run: int, + warmup_iterations: int, + min_run_time: float, + profile: bool, + label: str, +) -> float: + """Return milliseconds for one grouped GEMM launch.""" + run = _run_multistream if path == "multistream" else _run_cublas_grouped + run(prepared, warmup_iterations) + torch.cuda.synchronize() + + if profile: + torch.cuda.nvtx.range_push(label) + timing = benchmark.Timer( + stmt="run(prepared, iterations_per_run)", + globals={ + "run": run, + "prepared": prepared, + "iterations_per_run": iterations_per_run, + }, + num_threads=1, + ).blocked_autorange(min_run_time=min_run_time) + if profile: + torch.cuda.nvtx.range_pop() + return timing.median * 1000 / iterations_per_run + + +def _gemm_tflops(spec: GemmSpec, total_rows: int, time_ms: float) -> float: + """Compute aggregate GEMM throughput across all local experts.""" + flops = 2 * total_rows * spec.k * spec.n + return flops / time_ms / 1e9 + + +def _shape_description(spec: GemmSpec, rows_per_expert: str) -> str: + """Format the actual per-expert operands passed to TE.""" + if spec.gemm == "fwd": + return ( + f"A[{spec.n},{spec.k}] B[{rows_per_expert},{spec.k}] -> D[{rows_per_expert},{spec.n}]" + ) + if spec.gemm == "dgrad": + return ( + f"A[{spec.n},{spec.k}] B[{rows_per_expert},{spec.n}] -> D[{rows_per_expert},{spec.k}]" + ) + return f"A[{rows_per_expert},{spec.k}] B[{rows_per_expert},{spec.n}] -> D[{spec.n},{spec.k}]" + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--precision", choices=("all", "bf16", "mxfp8"), default="all") + parser.add_argument("--projection", choices=("all", "fc1", "fc2"), default="all") + parser.add_argument("--gemm", choices=("all", "fwd", "dgrad", "wgrad"), default="all") + parser.add_argument("--iterations-per-run", type=int, default=1000) + parser.add_argument("--warmup-iterations", type=int, default=500) + parser.add_argument("--min-run-time", type=float, default=5.0) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--m-splits", + type=str, + default=None, + help="Optional comma-separated per-expert rows; defaults to Qwen3.5-397B EP32.", + ) + parser.add_argument("--skip-correctness", action="store_true") + parser.add_argument("--profile", action="store_true", help="Add per-path NVTX ranges.") + parser.add_argument("--output-csv", type=str, default=None) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + if os.getenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "0") == "1": + raise RuntimeError( + "Unset NVTE_USE_CUTLASS_GROUPED_GEMM: this benchmark requires the " + "multi-stream cuBLAS baseline." + ) + if args.iterations_per_run < 1 or args.warmup_iterations < 1: + raise ValueError("iterations-per-run and warmup-iterations must be positive.") + + if args.m_splits is None: + num_local_experts = QWEN_NUM_EXPERTS // QWEN_EXPERT_PARALLEL_SIZE + rows_per_expert = QWEN_SEQUENCE_LENGTH * QWEN_TOP_K // num_local_experts + split_sizes = [rows_per_expert] * num_local_experts + else: + split_sizes = [int(value) for value in args.m_splits.split(",") if value] + if not split_sizes or any(value <= 0 for value in split_sizes): + raise ValueError("m_splits must contain positive integers.") + num_local_experts = len(split_sizes) + + total_rows = sum(split_sizes) + precisions = ("bf16", "mxfp8") if args.precision == "all" else (args.precision,) + specs = [ + spec + for spec in GEMM_SPECS + if (args.projection == "all" or spec.projection == args.projection) + and (args.gemm == "all" or spec.gemm == args.gemm) + ] + recipes: dict[str, Optional[Recipe]] = { + "bf16": None, + "mxfp8": MXFP8BlockScaling(), + } + mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available() + + uniform_rows = str(split_sizes[0]) if len(set(split_sizes)) == 1 else "variable" + print("Qwen3.5-397B-A17B grouped GEMM kernel benchmark") + print(f" local experts: {num_local_experts}") + print(f" total rows: {total_rows}") + print(f" m_splits: {split_sizes}") + print(" quantization is outside the timed region") + print() + for spec in specs: + print( + f" {spec.projection} {spec.gemm:5s} {spec.layout}: " + f"{_shape_description(spec, uniform_rows)}" + ) + print() + + rows = [] + for precision in precisions: + recipe = recipes[precision] + if precision == "mxfp8" and not mxfp8_available: + print(f"Skipping MXFP8: {reason_for_no_mxfp8}") + continue + if not is_module_grouped_tensor_path_supported(recipe, torch.bfloat16): + print( + f"Skipping {precision}: cuBLASLt grouped GEMM is unsupported on this " + "GPU or cuBLASLt version." + ) + continue + + for spec in specs: + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + prepared = _prepare_gemm(spec, split_sizes, precision) + if not args.skip_correctness: + _validate_outputs(prepared, precision) + + timings = {} + for path in ("multistream", "cublas_grouped_gemm"): + label = f"{precision}_{spec.projection}_{spec.gemm}_{path}" + timings[path] = _benchmark_path( + prepared, + path=path, + iterations_per_run=args.iterations_per_run, + warmup_iterations=args.warmup_iterations, + min_run_time=args.min_run_time, + profile=args.profile, + label=label, + ) + + speedup = timings["multistream"] / timings["cublas_grouped_gemm"] + for path, time_ms in timings.items(): + rows.append( + { + "precision": precision, + "projection": spec.projection, + "gemm": spec.gemm, + "layout": spec.layout, + "execution_path": path, + "num_local_experts": num_local_experts, + "total_rows": total_rows, + "k": spec.k, + "n": spec.n, + "time_ms": time_ms, + "tflops": _gemm_tflops(spec, total_rows, time_ms), + "speedup_vs_multistream": ( + speedup if path == "cublas_grouped_gemm" else 1.0 + ), + } + ) + + print( + f"{precision:6s} {spec.projection} {spec.gemm:5s}: " + f"multistream={timings['multistream']:.3f} ms, " + f"cuBLAS grouped={timings['cublas_grouped_gemm']:.3f} ms, " + f"speedup={speedup:.3f}x" + ) + + results = pd.DataFrame(rows) + print() + print(results.to_string(index=False)) + if args.output_csv is not None: + results.to_csv(args.output_csv, index=False) + print(f"\nWrote {args.output_csv}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py b/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py new file mode 100644 index 0000000000..1b400c0c24 --- /dev/null +++ b/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py @@ -0,0 +1,520 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Compare GroupedLinear's multi-stream and cuBLASLt grouped GEMM paths. + +The default problem is the routed-expert shape from Qwen3.5-397B-A17B with +sequence length 4096, top-10 routing, and expert parallelism 32: + + 512 global experts / EP32 = 16 local experts + 4096 tokens * top-10 / 16 local experts = 2560 rows per expert + +The 2560-row split is already 256-aligned, so both paths execute identical +GEMM work. FC1 and FC2 are benchmarked independently: + + FC1: 16 x (M=2560, K=4096, N=2048) + FC2: 16 x (M=2560, K=1024, N=4096) + +Both paths use discrete parameters initialized with identical values: + +* ``use_grouped_tensor=False`` receives CPU splits and launches the + multi-stream grouped GEMM implementation. +* ``use_grouped_tensor=True`` receives CUDA int64 splits and launches the + cuBLASLt grouped GEMM implementation. + +The MXFP8 cases use BF16 primary parameters. Each timing bundle refreshes the +MXFP8 weight cache on its first microbatch and reuses it on later microbatches. + +Examples +-------- +Run the complete BF16 and MXFP8 comparison: + + python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py + +Run only MXFP8 FC1 forward and backward: + + python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py \ + --precision mxfp8 --projection fc1 --mode fwd_bwd + +Profile one microbatch per timing invocation: + + nsys profile \ + --output=grouped_linear_cublas_grouped_gemm \ + --force-overwrite=true \ + --trace=cuda,nvtx,cublas \ + python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py \ + --precision mxfp8 --projection fc1 --mode fwd_bwd \ + --num-microbatches 1 --profile +""" + +import argparse +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Optional + +import pandas as pd +import torch +import torch.utils.benchmark as benchmark + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe +from transformer_engine.pytorch.module import ( + GroupedLinear, + is_module_grouped_tensor_path_supported, +) +from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + +QWEN_NUM_EXPERTS = 512 +QWEN_TOP_K = 10 +QWEN_SEQUENCE_LENGTH = 4096 +QWEN_EXPERT_PARALLEL_SIZE = 32 +QWEN_HIDDEN_SIZE = 4096 +QWEN_MOE_INTERMEDIATE_SIZE = 1024 + + +@dataclass(frozen=True) +class Projection: + """GroupedLinear dimensions for one Qwen routed-expert projection.""" + + name: str + in_features: int + out_features: int + + +PROJECTIONS = { + "fc1": Projection( + name="fc1", + in_features=QWEN_HIDDEN_SIZE, + out_features=2 * QWEN_MOE_INTERMEDIATE_SIZE, + ), + "fc2": Projection( + name="fc2", + in_features=QWEN_MOE_INTERMEDIATE_SIZE, + out_features=QWEN_HIDDEN_SIZE, + ), +} + + +@dataclass(frozen=True) +class ExecutionPath: + """GroupedLinear path and the corresponding m_splits representation.""" + + name: str + use_grouped_tensor: bool + + +EXECUTION_PATHS = ( + ExecutionPath(name="multistream", use_grouped_tensor=False), + ExecutionPath(name="cublas_grouped_gemm", use_grouped_tensor=True), +) + + +def _quantization_context(recipe: Optional[Recipe]): + """Construct a fresh quantization context for one invocation.""" + if recipe is None: + return nullcontext() + return te.autocast(enabled=True, recipe=recipe) + + +def _make_m_splits( + *, + use_grouped_tensor: bool, + split_sizes: list[int], +) -> torch.Tensor: + """Use the split representation consumed natively by each execution path.""" + device = "cuda" if use_grouped_tensor else "cpu" + return torch.tensor(split_sizes, dtype=torch.int64, device=device) + + +def _build_layer( + *, + projection: Projection, + num_local_experts: int, + use_grouped_tensor: bool, +) -> GroupedLinear: + """Construct a discrete-parameter GroupedLinear module.""" + return GroupedLinear( + num_local_experts, + projection.in_features, + projection.out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + use_grouped_tensor=use_grouped_tensor, + ) + + +def _run_microbatches( + layer: GroupedLinear, + x: torch.Tensor, + m_splits: torch.Tensor, + grad_output: torch.Tensor, + *, + recipe: Optional[Recipe], + mode: str, + num_microbatches: int, +) -> torch.Tensor: + """Run one timed bundle while preserving realistic MXFP8 weight caching.""" + layer.zero_grad(set_to_none=True) + x.grad = None + + if mode == "fwd": + with torch.no_grad(), _quantization_context(recipe): + for microbatch in range(num_microbatches): + output = layer( + x, + m_splits, + is_first_microbatch=(microbatch == 0), + ) + return output + + with _quantization_context(recipe): + for microbatch in range(num_microbatches): + output = layer( + x, + m_splits, + is_first_microbatch=(microbatch == 0), + ) + output.backward(grad_output) + return output + + +def _run_correctness_step( + layer: GroupedLinear, + x: torch.Tensor, + m_splits: torch.Tensor, + grad_output: torch.Tensor, + *, + recipe: Optional[Recipe], + mode: str, +) -> tuple[torch.Tensor, Optional[torch.Tensor], list[Optional[torch.Tensor]]]: + """Run one microbatch and retain outputs and gradients for path parity.""" + output = _run_microbatches( + layer, + x, + m_splits, + grad_output, + recipe=recipe, + mode=mode, + num_microbatches=1, + ) + input_grad = None if x.grad is None else x.grad.detach().clone() + parameter_grads = [ + None if param.grad is None else param.grad.detach().clone() for param in layer.parameters() + ] + return output.detach().clone(), input_grad, parameter_grads + + +def _validate_path_parity( + *, + multistream_layer: GroupedLinear, + grouped_layer: GroupedLinear, + multistream_x: torch.Tensor, + grouped_x: torch.Tensor, + multistream_splits: torch.Tensor, + grouped_splits: torch.Tensor, + grad_output: torch.Tensor, + recipe: Optional[Recipe], + mode: str, +) -> None: + """Require the two execution paths to produce compatible results.""" + reference = _run_correctness_step( + multistream_layer, + multistream_x, + multistream_splits, + grad_output, + recipe=recipe, + mode=mode, + ) + actual = _run_correctness_step( + grouped_layer, + grouped_x, + grouped_splits, + grad_output, + recipe=recipe, + mode=mode, + ) + + tolerances = {"rtol": 1e-2, "atol": 1e-2} + + torch.testing.assert_close(actual[0], reference[0], **tolerances) + if mode == "fwd": + return + + assert actual[1] is not None and reference[1] is not None + torch.testing.assert_close(actual[1], reference[1], **tolerances) + assert len(actual[2]) == len(reference[2]) + for actual_grad, reference_grad in zip(actual[2], reference[2]): + assert actual_grad is not None and reference_grad is not None + torch.testing.assert_close(actual_grad, reference_grad, **tolerances) + + +def _benchmark_path( + *, + layer: GroupedLinear, + x: torch.Tensor, + m_splits: torch.Tensor, + grad_output: torch.Tensor, + recipe: Optional[Recipe], + mode: str, + num_microbatches: int, + warmup_steps: int, + min_run_time: float, + profile: bool, + label: str, +) -> float: + """Benchmark one execution path and return milliseconds per microbatch.""" + _run_microbatches( + layer, + x, + m_splits, + grad_output, + recipe=recipe, + mode=mode, + num_microbatches=warmup_steps, + ) + torch.cuda.synchronize() + + if profile: + torch.cuda.nvtx.range_push(label) + + timing = benchmark.Timer( + stmt=( + "_run_microbatches(layer, x, m_splits, grad_output, recipe=recipe, " + "mode=mode, num_microbatches=num_microbatches)" + ), + globals={ + "_run_microbatches": _run_microbatches, + "layer": layer, + "x": x, + "m_splits": m_splits, + "grad_output": grad_output, + "recipe": recipe, + "mode": mode, + "num_microbatches": num_microbatches, + }, + num_threads=1, + ).blocked_autorange(min_run_time=min_run_time) + + if profile: + torch.cuda.nvtx.range_pop() + + return timing.median * 1000 / num_microbatches + + +def _gemm_tflops( + *, + total_rows: int, + projection: Projection, + mode: str, + time_ms: float, +) -> float: + """Report GEMM FLOP/s while timing also includes quantization and Python overhead.""" + flops = 2 * total_rows * projection.in_features * projection.out_features + if mode == "fwd_bwd": + flops *= 3 + return flops / time_ms / 1e9 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--precision", + choices=("all", "bf16", "mxfp8"), + default="all", + ) + parser.add_argument( + "--projection", + choices=("all", "fc1", "fc2"), + default="all", + ) + parser.add_argument( + "--mode", + choices=("fwd", "fwd_bwd"), + default="fwd_bwd", + ) + parser.add_argument("--num-microbatches", type=int, default=16) + parser.add_argument("--warmup-steps", type=int, default=500) + parser.add_argument("--min-run-time", type=float, default=30.0) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--m-splits", + type=str, + default=None, + help="Optional comma-separated per-expert rows; defaults to Qwen3.5-397B EP32.", + ) + parser.add_argument( + "--skip-correctness", + action="store_true", + help="Skip output, input-gradient, and parameter-gradient parity checks.", + ) + parser.add_argument("--profile", action="store_true", help="Add per-path NVTX ranges.") + parser.add_argument("--output-csv", type=str, default=None) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + if args.num_microbatches < 1 or args.warmup_steps < 1: + raise ValueError("num_microbatches and warmup_steps must both be positive.") + + if args.m_splits is None: + num_local_experts = QWEN_NUM_EXPERTS // QWEN_EXPERT_PARALLEL_SIZE + rows_per_expert = QWEN_SEQUENCE_LENGTH * QWEN_TOP_K // num_local_experts + split_sizes = [rows_per_expert] * num_local_experts + else: + split_sizes = [int(value) for value in args.m_splits.split(",") if value] + if not split_sizes or any(value < 0 for value in split_sizes): + raise ValueError("m_splits must contain non-negative integers.") + num_local_experts = len(split_sizes) + + total_rows = sum(split_sizes) + if total_rows == 0: + raise ValueError("At least one routed token row is required for this benchmark.") + + precision_names = ("bf16", "mxfp8") if args.precision == "all" else (args.precision,) + projection_names = ("fc1", "fc2") if args.projection == "all" else (args.projection,) + + recipes: dict[str, Optional[Recipe]] = { + "bf16": None, + "mxfp8": MXFP8BlockScaling(), + } + mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available() + + print("Qwen3.5-397B-A17B GroupedLinear benchmark") + print(f" local experts: {num_local_experts}") + print(f" total routed rows: {total_rows}") + print(f" m_splits: {split_sizes}") + print(f" mode: {args.mode}") + print(" primary parameter dtype: BF16") + print(f" microbatches per timing invocation: {args.num_microbatches}") + print() + + rows = [] + for precision_name in precision_names: + recipe = recipes[precision_name] + if precision_name == "mxfp8" and not mxfp8_available: + print(f"Skipping MXFP8: {reason_for_no_mxfp8}") + continue + if not is_module_grouped_tensor_path_supported(recipe, torch.bfloat16): + print( + f"Skipping {precision_name}: the cuBLASLt grouped-tensor path is unsupported " + "on this GPU or cuBLASLt version." + ) + continue + + for projection_name in projection_names: + projection = PROJECTIONS[projection_name] + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + layers = { + path.name: _build_layer( + projection=projection, + num_local_experts=num_local_experts, + use_grouped_tensor=path.use_grouped_tensor, + ) + for path in EXECUTION_PATHS + } + layers["cublas_grouped_gemm"].load_state_dict(layers["multistream"].state_dict()) + + base_x = torch.randn( + total_rows, + projection.in_features, + dtype=torch.bfloat16, + device="cuda", + ) + inputs = { + path.name: base_x.detach().clone().requires_grad_(args.mode == "fwd_bwd") + for path in EXECUTION_PATHS + } + grad_output = torch.randn( + total_rows, + projection.out_features, + dtype=torch.bfloat16, + device="cuda", + ) + splits = { + path.name: _make_m_splits( + use_grouped_tensor=path.use_grouped_tensor, + split_sizes=split_sizes, + ) + for path in EXECUTION_PATHS + } + + if not args.skip_correctness: + _validate_path_parity( + multistream_layer=layers["multistream"], + grouped_layer=layers["cublas_grouped_gemm"], + multistream_x=inputs["multistream"], + grouped_x=inputs["cublas_grouped_gemm"], + multistream_splits=splits["multistream"], + grouped_splits=splits["cublas_grouped_gemm"], + grad_output=grad_output, + recipe=recipe, + mode=args.mode, + ) + + timings = {} + for path in EXECUTION_PATHS: + label = f"{precision_name}_{projection_name}_{args.mode}_{path.name}" + timing_ms = _benchmark_path( + layer=layers[path.name], + x=inputs[path.name], + m_splits=splits[path.name], + grad_output=grad_output, + recipe=recipe, + mode=args.mode, + num_microbatches=args.num_microbatches, + warmup_steps=args.warmup_steps, + min_run_time=args.min_run_time, + profile=args.profile, + label=label, + ) + timings[path.name] = timing_ms + + speedup = timings["multistream"] / timings["cublas_grouped_gemm"] + for path in EXECUTION_PATHS: + timing_ms = timings[path.name] + rows.append( + { + "precision": precision_name, + "projection": projection_name, + "mode": args.mode, + "execution_path": path.name, + "num_local_experts": num_local_experts, + "total_rows": total_rows, + "in_features": projection.in_features, + "out_features": projection.out_features, + "time_ms": timing_ms, + "gemm_tflops": _gemm_tflops( + total_rows=total_rows, + projection=projection, + mode=args.mode, + time_ms=timing_ms, + ), + "speedup_vs_multistream": (1.0 if path.name == "multistream" else speedup), + } + ) + + print( + f"{precision_name:6s} {projection_name} {args.mode}: " + f"multistream={timings['multistream']:.3f} ms, " + f"cuBLAS grouped={timings['cublas_grouped_gemm']:.3f} ms, " + f"speedup={speedup:.3f}x" + ) + + results = pd.DataFrame(rows) + print() + print(results.to_string(index=False)) + if args.output_csv is not None: + results.to_csv(args.output_csv, index=False) + print(f"\nWrote {args.output_csv}") + + +if __name__ == "__main__": + main() diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index cbb8838b00..30ea83a55b 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -4,11 +4,12 @@ """Installation script.""" +import copy import os import subprocess import sys import sysconfig -import copy +import tempfile import time from pathlib import Path @@ -111,23 +112,37 @@ def run(self) -> None: for ext in self.extensions: package_path = Path(self.get_ext_fullpath(ext.name)) install_dir = package_path.resolve().parent - if isinstance(ext, CMakeExtension): - print(f"Building CMake extension {ext.name}") - # Set up incremental builds for CMake extensions - build_dir = os.getenv("NVTE_CMAKE_BUILD_DIR") - if build_dir: - build_dir = Path(build_dir).resolve() - else: - root_dir = Path(__file__).resolve().parent.parent - build_dir = root_dir / "build" / "cmake" - - # Ensure the directory exists + if not isinstance(ext, CMakeExtension): + continue + + print(f"Building CMake extension {ext.name}") + configured_build_dir = os.getenv("NVTE_CMAKE_BUILD_DIR") + if not configured_build_dir and self.inplace: + root_dir = Path(__file__).resolve().parent.parent + configured_build_dir = root_dir / "build" / "cmake" + + if configured_build_dir: + # A persistent build directory enables incremental builds. + build_dir = Path(configured_build_dir).resolve() build_dir.mkdir(parents=True, exist_ok=True) - ext._build_cmake( build_dir=build_dir, install_dir=install_dir, ) + continue + + # Isolate CMake state between concurrent and successive builds. + build_temp = Path(self.build_temp) + build_temp.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix=f"cmake-build-{ext.name}-", + dir=build_temp, + ) as build_dir: + print(f"Building CMake extension {ext.name} in temporary directory {build_dir}") + ext._build_cmake( + build_dir=Path(build_dir), + install_dir=install_dir, + ) # Build non-CMake extensions as usual all_extensions = self.extensions diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index 0be19bba67..3e4742d8ae 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -40,7 +40,6 @@ pip3 install -r $TE_PATH/examples/jax/encoder/requirements.txt || error_exit "Fa export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py" # Test without custom calls -export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" NVTE_JAX_CUSTOM_CALLS="false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder_without_custom_call.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py without custom calls" # Exercise the docs/examples/jax tutorials. The multi-GPU tests are diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index 3efa462628..36efe485f5 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -22,6 +22,11 @@ FAILED_CASES="" : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage. +export NVTE_FLASH_ATTN_V2=1 +export NVTE_FLASH_ATTN_V3=0 +export NVTE_FLASH_ATTN_V4=0 + # Config with the dummy feature which prevents nvinspect from being disabled. # Nvinspect will be disabled if no feature is active. : ${NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE:=$TE_PATH/tests/pytorch/debug/test_configs/dummy_feature.yaml} diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 6a41515602..ee3ef89e66 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -22,6 +22,11 @@ set -x : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage. +export NVTE_FLASH_ATTN_V2=1 +export NVTE_FLASH_ATTN_V3=0 +export NVTE_FLASH_ATTN_V4=0 + pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" @@ -60,6 +65,7 @@ NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$X 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_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.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_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 90ff6ba2fd..ec19492ee7 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -24,20 +24,22 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # Run CP tests (deterministic + non-deterministic) first so they can be parallelized. # Each needs 4 GPUs, so >=8 GPUs allows them to run concurrently on disjoint GPU sets. +# Main's CP implementation supports FA2/FA3. Keep FA4 disabled at the suite +# boundary so both the reference and CP halves use the same backend generation. NUM_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())") echo "Detected $NUM_GPUS GPU(s)" if [ "$NUM_GPUS" -ge 8 ]; then echo "Running CP tests in parallel: non-deterministic on GPUs 0-3, deterministic on GPUs 4-7" - CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + CUDA_VISIBLE_DEVICES=0,1,2,3 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & PID_CP_NONDET=$! - CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & PID_CP_DET=$! wait $PID_CP_NONDET || test_fail "test_attention_with_cp.py" wait $PID_CP_DET || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" else echo "Running CP tests sequentially: need >=8 GPUs for parallel execution" - python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" - NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" + NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" + NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" fi python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py" diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 468ef04d76..047358b301 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -31,37 +31,57 @@ export NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 # Iterate over Flash Attention versions sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); print(sm[0]*10+sm[1])"` export FLASH_ATTN_CUDA_ARCHS=$sm_arch -# CP tests are expensive and run only once per arch: -# - sm90 (H100): FA3 (3.0.0b1) - context_parallel.py only supports FA3 on Hopper -# - sm>90 (B200): latest FA4 - FA3 is not built/installed for sm>90 -# Non-CP tests still run for every FA version in the array. +# Run one architecture-owned FlashAttention generation. CP remains FA3-only +# until the production selector and runner support FA4 CP end to end. +CP_FA_VERSION="" if [ $sm_arch -gt 90 ] then - FA_versions=(2.8.3 4.0.0b11) - CP_FA_VERSION="${FA_versions[-1]}" + FA_versions=(4.0.0b11) elif [ $sm_arch -eq 90 ] then - FA_versions=(2.8.3 3.0.0b1 4.0.0b11) + FA_versions=(3.0.0b1) CP_FA_VERSION="3.0.0b1" +else + error_exit "No L3 FlashAttention generation is defined for sm${sm_arch}" fi for fa_version in "${FA_versions[@]}" do + # The FA distributions share the flash_attn namespace. Keep exactly one + # installed so import-time discovery and the iteration label cannot disagree. + pip3 uninstall -y flash-attn flash-attn-3 flash-attn-4 \ + || error_exit "Failed to isolate Flash Attention $fa_version" + export NVTE_FLASH_ATTN_V2=0 + export NVTE_FLASH_ATTN_V3=0 + export NVTE_FLASH_ATTN_V4=0 + # Build Flash Attention if [ "${fa_version}" \< "3.0.0" ] then - pip3 install flash-attn==${fa_version} --no-build-isolation + export NVTE_FLASH_ATTN_V2=1 + pip3 install flash-attn==${fa_version} --no-build-isolation \ + || error_exit "Failed to install Flash Attention $fa_version" elif [[ "${fa_version}" == 4.* ]] then - pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 --no-build-isolation + export NVTE_FLASH_ATTN_V4=1 + # FA4 is intentionally last in every version array. Its b11 test pin needs + # CUTLASS DSL 4.4.2, so replace the image-matched stack only for this final + # iteration; later iterations would otherwise need that stack restored. + pip3 uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ + nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 \ + || error_exit "Failed to isolate CUTLASS DSL for Flash Attention $fa_version" + pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 \ + --no-build-isolation || error_exit "Failed to install Flash Attention $fa_version" else + export NVTE_FLASH_ATTN_V3=1 # FA3 source build (~20 min). Skip if FA3 is already installed. if python3 -c "import flash_attn_3" 2>/dev/null; then echo "FA3 already installed (from base image); skipping source build" else git clone https://github.com/Dao-AILab/flash-attention.git - cd flash-attention/hopper && python setup.py install + cd flash-attention/hopper && python setup.py install \ + || error_exit "Failed to install Flash Attention $fa_version" cd ../../ fi fi @@ -82,7 +102,7 @@ do # test_attention.py reloads its own trusted delayed-scaling FP8 checkpoint, # whose legacy extra state requires an explicit pickle opt-in. - if [ "$fa_version" = "$CP_FA_VERSION" ]; then + if [ -n "$CP_FA_VERSION" ] && [ "$fa_version" = "$CP_FA_VERSION" ]; then echo "Running CP tests with FA $fa_version (CP version for sm$sm_arch)" if [ "$NUM_GPUS" -ge 5 ]; then CP_NUM_GPUS=$(( NUM_GPUS - 1 > 4 ? 4 : NUM_GPUS - 1 )) @@ -107,7 +127,11 @@ do NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_CP $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py (FA $fa_version)" fi else - echo "Skipping CP tests for FA $fa_version (CP only runs with FA $CP_FA_VERSION on sm$sm_arch)" + if [ -n "$CP_FA_VERSION" ]; then + echo "Skipping CP tests for FA $fa_version (CP uses FA $CP_FA_VERSION on sm$sm_arch)" + else + echo "CP tests are not scheduled for the FA generation on sm$sm_arch" + fi NVTE_TORCH_COMPILE=0 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" fi done diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index c7c778ce1e..edeb87ebe1 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -7,10 +7,12 @@ #include #include #include +#include #include #include #include +#include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -735,3 +737,239 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::ValuesIn(input_scenarios)), test_name_generator); + +// ============================================================================ +// Swizzled-scales cast-only tests +// +// Validate the WITH_GEMM_SWIZZLED_SCALES=true code path added by the +// CastTraitsSwizzle port. The specialized kernel dispatches to +// CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> whenever the +// output tensor has set_with_gemm_swizzled_scales(true), producing scales +// directly in GEMM-swizzled layout. +// +// Reference construction: run the well-tested linear-scale path, then apply +// nvte_swizzle_scaling_factors (independently tested by SwizzleTestSuite) to +// transform the linear scales into the swizzled layout. Byte-compare the +// direct swizzled cast output against this reference for both scale tensors +// and the FP8 output data itself. +// +// Pass criteria per test: +// 1. FP8 rowwise data byte-identical between linear and swizzled paths. +// 2. FP8 colwise data byte-identical between linear and swizzled paths. +// 3. Swizzled rowwise scale bytes match linear->swizzle reference. +// 4. Swizzled colwise scale bytes match linear->swizzle reference. +// 5. Rowwise-only FP8 data and swizzled scales match the same reference. +// 6. No CUDA errors from any launch. +// ============================================================================ + +class SwizzledScalesFusedCastMXFP8TestSuite : public ::testing::TestWithParam< + std::tuple, + transformer_engine::DType, + transformer_engine::DType>> {}; + +TEST_P(SwizzledScalesFusedCastMXFP8TestSuite, TestSwizzledCastMXFP8) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const auto shape = std::get<0>(GetParam()); + const DType itype = std::get<1>(GetParam()); + const DType otype = std::get<2>(GetParam()); + + // BIDIMENSIONAL scaling needs a 2D+ input. + if (shape.size() < 2) { + GTEST_SKIP(); + } + + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + const size_t out_bytes = rows * cols; // fp8 = 1 byte/elem + + // Input filled with the same values for all casts. + Tensor input("input", shape, itype); + fillUniform(&input); + + // Target: swizzled-scale cast. Dispatcher routes to + // CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> on kernel #3. + Tensor output_swizzled("output_swizzled", shape, otype, true, true, NVTE_MXFP8_1D_SCALING); + output_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_swizzled.data(), 0); + + // Rowwise-only activations take the pointer-based cast-only kernel. This + // directly exercises its WITH_GEMM_SWIZZLED_SCALES trait specialization + // for shapes whose column count is a multiple of 128. + Tensor output_rowwise_swizzled("output_rowwise_swizzled", shape, otype, + /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + output_rowwise_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_rowwise_swizzled.data(), 0); + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "swizzled-scale nvte_quantize failed"; + + // Reference construction: nvte_swizzle_scaling_factors accepts tensors with + // exactly one scale direction, so we build the rowwise and colwise references + // independently. Each is a single-direction linear cast followed by a + // single-direction swizzle transform, using the same input as the target. + // MXFP8 is a deterministic per-direction operation, so a rowwise-only cast + // produces the same rowwise fp8 bytes and scales as a rowwise+colwise cast. + + // Rowwise reference. + Tensor linear_row("linear_row", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_row.data(), 0); + Tensor ref_row_swz("ref_row_swz", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + ref_row_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_row_swz.rowwise_dptr(), linear_row.rowwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_row.data(), ref_row_swz.data(), 0); + } + + // Colwise reference. + Tensor linear_col("linear_col", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_col.data(), 0); + Tensor ref_col_swz("ref_col_swz", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + ref_col_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_col_swz.columnwise_dptr(), linear_col.columnwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_col.data(), ref_col_swz.data(), 0); + } + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "reference construction failed"; + + // ---- Comparisons ---- + + // (1) & (2): FP8 output data byte-identical to single-direction linear casts. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test_row = reinterpret_cast( + output_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref_row = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_row[i], ref_row[i]) + << "rowwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_row[i]) + << " linear=" << static_cast(ref_row[i]) << ")"; + } + const uint8_t *test_col = reinterpret_cast( + output_swizzled.columnwise_cpu_dptr()); + const uint8_t *ref_col = reinterpret_cast( + linear_col.columnwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_col[i], ref_col[i]) + << "colwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_col[i]) + << " linear=" << static_cast(ref_col[i]) << ")"; + } + } + ); + + // (3): Swizzled rowwise scale bytes — directly-written vs linear-then-transform. + { + const size_t n = product(output_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled rowwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (4): Swizzled colwise scale bytes — exercises the uint16-pair-packed flush + // path added by CACHE_COLWISE_SCALE_IN_SMEM + WITH_SWIZZLED_SCALES. + { + const size_t n = product(output_swizzled.columnwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.columnwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_col_swz.columnwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled colwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (5) & (6): The rowwise-only swizzled path matches the same independently + // constructed linear-then-swizzle reference. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test = reinterpret_cast( + output_rowwise_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only fp8 data mismatch at index " << i; + } + } + ); + { + const size_t n = product(output_rowwise_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = + output_rowwise_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only swizzled scale mismatch at offset " << i; + } + } + +} + +std::string swizzled_test_name_generator( + const testing::TestParamInfo& info) { + std::string name; + const auto &shape = std::get<0>(info.param); + for (size_t i = 0; i < shape.size(); ++i) { + if (i > 0) name += "x"; + name += std::to_string(shape[i]); + } + name += "X" + test::typeName(std::get<1>(info.param)) + + "X" + test::typeName(std::get<2>(info.param)); + return name; +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest_FusedCastMXFP8_SwizzledCastOnly, + SwizzledScalesFusedCastMXFP8TestSuite, + ::testing::Values( + // 1. Aligned, small — sanity. + std::make_tuple(std::vector{128, 128}, DType::kBFloat16, DType::kFloat8E4M3), + // 2. Aligned, small — second dtype pair. + std::make_tuple(std::vector{128, 128}, DType::kFloat32, DType::kFloat8E5M2), + // 3. Aligned, medium. + std::make_tuple(std::vector{256, 384}, DType::kBFloat16, DType::kFloat8E4M3), + // 4. Odd number of scale rows (96/32 = 3) - exercises colwise flush + // odd-row scalar tail. + std::make_tuple(std::vector{96, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 5. cols/32 not a multiple of 4 - exercises rowwise flush scalar tail + // (need cols multiple of 8 for TMA-alignment on 16-bit input). + std::make_tuple(std::vector{256, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 6. Odd colwise scale rows + rowwise tail cols - combined tail paths. + std::make_tuple(std::vector{96, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 7. Small but tile-aligned. + std::make_tuple(std::vector{32, 64}, DType::kBFloat16, DType::kFloat8E4M3), + // 8. Minimum aligned size - single 32x32 tile. + std::make_tuple(std::vector{32, 32}, DType::kBFloat16, DType::kFloat8E4M3), + // 9. Multi-CTA at scale — stresses gmem writes across many CTAs. + std::make_tuple(std::vector{4096, 32768}, DType::kBFloat16, DType::kFloat8E4M3), + // 10. 4D input — matches the existing suite's rank coverage. + std::make_tuple(std::vector{16, 8, 4, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 11. Third input dtype. + std::make_tuple(std::vector{128, 128}, DType::kFloat16, DType::kFloat8E4M3), + // 12. Larger fp32. + std::make_tuple(std::vector{1024, 1024}, DType::kFloat32, DType::kFloat8E4M3) + ), + swizzled_test_name_generator); diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 47732196ed..be6dbed3ab 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -97,6 +97,26 @@ static std::vector expected_recv_values_sorted( return vals; } +// Sorted multiset of per-token identity bytes each recv_rank should receive, +// derived from the deterministic routing map (same id_of stamp as the test). +static std::vector expected_recv_ids_sorted( + int recv_rank, int num_processes, int num_tokens, int top_k, + int num_experts, int num_local_experts) { + int base = recv_rank * num_local_experts; + std::vector ids; + for (int src = 0; src < num_processes; ++src) { + auto idx = routing_balanced(src, num_tokens, top_k, num_experts, num_local_experts); + for (int t = 0; t < num_tokens; ++t) + for (int k = 0; k < top_k; ++k) { + int64_t e = idx[t * top_k + k]; + if (e >= base && e < base + num_local_experts) + ids.push_back(static_cast((src * num_tokens + t + 1) & 0xFF)); + } + } + std::sort(ids.begin(), ids.end()); + return ids; +} + // 2^-5 relative tolerance for BF16 (matches mantissa precision with margin), // plus a small atol floor for near-zero expected values. static constexpr float kBf16Rtol = 1.0f / 32.0f; @@ -368,6 +388,123 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) { NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } +// ============================================================================= +// EPDispatchScaledTest: MXFP8 dispatch routes each token's e8m0 scale row in +// lockstep with its e4m3 data row. +// ============================================================================= + +// Each source token stamps a nonzero identity byte across its whole data row +// and its whole scale row. Dispatch is a pure permutation, so every filled recv +// slot must carry the same id in both buffers; a zero id means scales were not +// routed. +TEST_F(EpOpTestBase, MXFP8DispatchScales) { + if (g_sm_major < 9) GTEST_SKIP() << "EP requires SM_90+"; + // MXFP8 dispatch is supported on all HT EM modes (local_permute, local_dup, + // nvlink_dup); the mode is selected by the NCCL_EP_HT_EM_* env at group init. + // MXFP8 e8m0 scale bytes/token = hidden/32; NCCL EP HT requires 16B alignment. + if (hidden_dim_ % 512 != 0) + GTEST_SKIP() << "MXFP8 dispatch needs hidden_dim % 512 == 0 (16B scale alignment)"; + const int scale_cols = hidden_dim_ / 32; + using Shape = std::vector; + + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + auto h_idx = routing_balanced(g_process_id, num_tokens_, top_k_, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens_ * top_k_, 1.0f / top_k_); + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + + auto id_of = [&](int tok) { + return static_cast((g_process_id * num_tokens_ + tok + 1) & 0xFF); + }; + std::vector h_data(num_tokens_ * hidden_dim_); + std::vector h_scale(num_tokens_ * scale_cols); + for (int tok = 0; tok < num_tokens_; ++tok) { + const uint8_t id = id_of(tok); + std::fill_n(&h_data[tok * hidden_dim_], hidden_dim_, id); + std::fill_n(&h_scale[tok * scale_cols], scale_cols, id); + } + + DevBuf d_data(num_tokens_ * hidden_dim_), d_scale(num_tokens_ * scale_cols); + DevBuf d_recv_data(buf.recv_capacity * hidden_dim_); + DevBuf d_recv_scale(buf.recv_capacity * scale_cols); + NVTE_CHECK_CUDA(cudaMemcpy(d_data.get(), h_data.data(), h_data.size(), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(d_scale.get(), h_scale.data(), h_scale.size(), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemset(d_recv_data.get(), 0, d_recv_data.bytes())); + NVTE_CHECK_CUDA(cudaMemset(d_recv_scale.get(), 0, d_recv_scale.bytes())); + + TensorWrapper tokens(NVTE_MXFP8_1D_SCALING); + tokens.set_rowwise_data(d_data.get(), DType::kFloat8E4M3, + Shape{(size_t)num_tokens_, (size_t)hidden_dim_}); + tokens.set_rowwise_scale_inv(d_scale.get(), DType::kFloat8E8M0, + Shape{(size_t)num_tokens_, (size_t)scale_cols}); + TensorWrapper recv_tokens(NVTE_MXFP8_1D_SCALING); + recv_tokens.set_rowwise_data(d_recv_data.get(), DType::kFloat8E4M3, + Shape{buf.recv_capacity, (size_t)hidden_dim_}); + recv_tokens.set_rowwise_scale_inv(d_recv_scale.get(), DType::kFloat8E8M0, + Shape{buf.recv_capacity, (size_t)scale_cols}); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), + t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), + tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector counts(num_local_experts_); + NVTE_CHECK_CUDA(cudaMemcpy(counts.data(), buf.recv_tokens_per_expert.get(), + num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); + auto exp_counts = expected_recv_tokens_per_expert(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); + int total_recv = 0; + for (int e = 0; e < num_local_experts_; ++e) { + EXPECT_EQ(counts[e], exp_counts[e]) << "local expert " << e; + total_recv += exp_counts[e]; + } + ASSERT_LE(total_recv, static_cast(buf.recv_capacity)); + + std::vector r_data(buf.recv_capacity * hidden_dim_); + std::vector r_scale(buf.recv_capacity * scale_cols); + NVTE_CHECK_CUDA(cudaMemcpy(r_data.data(), d_recv_data.get(), r_data.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(r_scale.data(), d_recv_scale.get(), r_scale.size(), cudaMemcpyDeviceToHost)); + + std::vector got_data_ids, got_scale_ids; + got_data_ids.reserve(total_recv); + got_scale_ids.reserve(total_recv); + size_t slot = 0; + for (int e = 0; e < num_local_experts_; ++e) { + for (int i = 0; i < counts[e]; ++i, ++slot) { + got_data_ids.push_back(r_data[slot * hidden_dim_]); + got_scale_ids.push_back(r_scale[slot * scale_cols]); + } + } + std::sort(got_data_ids.begin(), got_data_ids.end()); + std::sort(got_scale_ids.begin(), got_scale_ids.end()); + + auto exp_ids = expected_recv_ids_sorted(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); + ASSERT_EQ(got_data_ids.size(), exp_ids.size()); + ASSERT_EQ(got_scale_ids.size(), exp_ids.size()); + for (size_t i = 0; i < exp_ids.size(); ++i) { + EXPECT_EQ(got_data_ids[i], exp_ids[i]) << "recv data id mismatch at sorted index " << i; + EXPECT_EQ(got_scale_ids[i], exp_ids[i]) << "recv scale id mismatch at sorted index " << i; + } + + if (g_process_id == 0) + printf(" MXFP8DispatchScales: passed (recv=%d, data + scales match routing map)\n", total_recv); + + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + // ============================================================================= // EPCombineTest: round-trip identity expert -> result == top_k * tokens. // ============================================================================= diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index 7cf6017090..37f324eabb 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -7,7 +7,7 @@ /* * Shared TE EP test infrastructure. Include once per TU; ep_bootstrap() in * each test binary's main() populates process-level globals. - * Defaults: 4 experts/rank, hidden_dim=256, max_tokens_per_rank=64. + * Defaults: 4 experts/rank, hidden_dim=512, max_tokens_per_rank=64. */ #pragma once @@ -48,7 +48,7 @@ static int g_num_processes = -1; static int g_sm_major = -1; // set by ep_bootstrap; -1 until then static int g_ep_size = -1; static int g_num_experts = -1; -static int g_hidden_dim = 256; +static int g_hidden_dim = 512; static int g_max_tokens_per_rank = 64; static NVTEDType g_max_token_dtype = kNVTEFloat32; // staging-buffer sizing static bool g_ep_initialized = false; diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 9678195942..c1cc2b8282 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -539,7 +539,7 @@ def check_has_backend_for_mask(mask_type): "window_size", [ pytest.param((-1, -1), id="window_size(-1, -1)"), - pytest.param((5, 0), id="window_size(8, 0)"), + pytest.param((5, 0), id="window_size(5, 0)"), ], ) @pytest.mark.parametrize( diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 4121e5d1b4..47af0b0c39 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -22,6 +22,7 @@ """ import os +import re import sys import unittest @@ -689,9 +690,9 @@ def run(idx, toks, w): "JAX/XLA lacks the gpu_stream:collective annotation (openxla/xla#39604)", ) def test_z_dispatch_combine_on_collective_stream(self): - """Every EP FFI custom call must carry the collective-stream annotation - so XLA schedules them on the collective stream instead of overlapping - them with other collectives.""" + """Every EP FFI custom call must run on the collective stream. compute_on + puts the annotation on the async wrapper XLA generates, so assert each EP + call is reachable from a wrapper that carries it.""" T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() dp_spec = PartitionSpec(("dp", "ep"), None) ep_spec_3d = PartitionSpec(("dp", "ep"), None, None) @@ -719,18 +720,50 @@ def run(idx, toks, w): hlo = run.lower(topk_idx, tokens, topk_w).compile().as_text() - # Every te_ep_* FFI custom call must carry the collective-stream - # annotation so XLA places it on the collective stream. - ep_lines = [l for l in hlo.splitlines() if 'custom_call_target="te_ep_' in l] - self.assertTrue(ep_lines, f"no te_ep_* custom calls in compiled HLO:\n{hlo}") + # Parse the HLO into computations and follow the call graph: a call + # carrying the collective-stream annotation places its callee (and every + # nested callee) on the collective stream. + comps = {} + cur = None + for line in hlo.splitlines(): + stripped = line.strip() + header = re.match(r"(?:ENTRY\s+)?(%[\w.\-]+)\s*\(", stripped) + if header and stripped.endswith("{"): + cur = header.group(1) + comps[cur] = [] + elif stripped == "}": + cur = None + elif cur is not None: + comps[cur].append(line) + + callees = lambda l: re.findall(r"(?:calls|to_apply)=(%[\w.\-]+)", l) + edges = {c: {x for l in ls for x in callees(l)} for c, ls in comps.items()} + + collective = set() + for ls in comps.values(): + for l in ls: + if '_xla_stream_annotation="collective"' in l.replace(" ", ""): + collective.update(callees(l)) + stack = list(collective) + while stack: + for callee in edges.get(stack.pop(), ()): + if callee not in collective: + collective.add(callee) + stack.append(callee) + + ep_calls = [ + (c, l) for c, ls in comps.items() for l in ls if 'custom_call_target="te_ep_' in l + ] + self.assertTrue(ep_calls, f"no te_ep_* custom calls in compiled HLO:\n{hlo}") missing = [ l.strip()[:200] - for l in ep_lines - if '_xla_stream_annotation="collective"' not in l.replace(" ", "") + for c, l in ep_calls + if c not in collective + and '_xla_stream_annotation="collective"' not in l.replace(" ", "") ] self.assertFalse( missing, - "te_ep_* custom calls missing collective-stream annotation:\n" + "\n".join(missing), + "te_ep_* custom calls not on the collective stream:\n" + "\n".join(missing), ) def test_z_no_unexpected_reshard_in_hlo_bwd(self): diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 55ed415d21..0e9cb6598d 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -140,7 +140,7 @@ def _read_mp_options(): ("exp", EP_AXIS), ("embed", FSDP_AXIS), ("mlp", None), - ("batch", (EP_AXIS, FSDP_AXIS)), + ("batch", (FSDP_AXIS, EP_AXIS)), ) # Small shapes so the parity tests stay tight on bf16. The block still @@ -364,6 +364,7 @@ def _make_block( use_expert_routing_bias=False, score_function="softmax", expert_bias_init=None, + input_axes=("batch", None, None), ): kwargs = dict( num_experts=NUM_EXPERTS, @@ -375,6 +376,7 @@ def _make_block( use_expert_routing_bias=use_expert_routing_bias, score_function=score_function, dtype=DTYPE, + input_axes=input_axes, ) # Custom expert_bias_init lets tests inject a non-zero expert_bias without # poking variables['params'] post-init. diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 82791084d8..c4f852f112 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -341,9 +341,18 @@ def test_dpa_num_splits(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." +fa4_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) and bool( + int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) +) +requires_fa4 = pytest.mark.skipif( + not fa4_enabled + or not FlashAttentionUtils.v4_is_installed + or device_compute_capability < (9, 0), + reason="Enabled Flash-attn v4 and compute capability >= SM90 are required.", ) + + +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_base]) @pytest.mark.parametrize("model", model_configs_fa4_base.keys()) @@ -362,9 +371,7 @@ def test_dpa_fa4_base(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.skipif( device_compute_capability not in ((10, 0), (10, 3)), reason="FA4 head_dim=256 dedicated kernel is SM100/103-only.", @@ -442,9 +449,7 @@ def test_dpa_d256(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mla]) @pytest.mark.parametrize("model", model_configs_fa4_mla.keys()) @@ -466,9 +471,7 @@ def test_dpa_fa4_mla(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_swa]) @pytest.mark.parametrize("model", model_configs_fa4_swa.keys()) @@ -489,9 +492,7 @@ def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_varlen]) @pytest.mark.parametrize("model", model_configs_fa4_varlen.keys()) @@ -514,9 +515,7 @@ def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mask]) @pytest.mark.parametrize("model", model_configs_fa4_mask.keys()) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 2c0a5d9217..d7eb16b862 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -9,7 +9,6 @@ import subprocess import sys import threading -import time import pathlib import logging import copy @@ -29,7 +28,7 @@ from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig, get_available_attention_backends pytest_logging_level = logging.getLevelName(logging.root.level) @@ -300,7 +299,10 @@ def _submit(pool: PoolWorker, **kwargs) -> None: qkv_formats = ["sbhd", "thd"] -@pytest.mark.skipif(not FlashAttentionUtils.v2_plus, reason="Flash-attn 2.0+ is required.") +@pytest.mark.skipif( + not (FlashAttentionUtils.v2_plus or FlashAttentionUtils.v3_is_installed), + reason="Flash-attn v2 or v3 is required.", +) @pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.") @pytest.mark.parametrize("dtype", dtypes) @pytest.mark.parametrize("model", model_configs_flash_attn.keys()) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py new file mode 100644 index 0000000000..2c061607fc --- /dev/null +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -0,0 +1,137 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for FusedMLAQUpProjRopeQuant. + +Run: + pytest tests/pytorch/attention/test_fused_mla_q_uproj.py -v +""" + +import pytest +import torch + +import transformer_engine.pytorch # registers transformer_engine_torch +import transformer_engine_torch as tex +from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor + +# DSv3 671B MLA dims +NUM_HEADS = 128 +HEAD_DIM_NOPE = 128 +HEAD_DIM_ROPE = 64 +HEAD_DIM = HEAD_DIM_NOPE + HEAD_DIM_ROPE # 192 +Q_LORA_RANK = 1536 +PROJ_DIM = NUM_HEADS * HEAD_DIM # 24576 + +SEED = 42 + +fused_supported, reason_not_supported = ( + (True, "") + if FusedMLAQUpProjRopeQuant.is_supported() + else ( + False, + ( + "FusedMLAQUpProjRopeQuant.is_supported() returned False " + "(SM100+, cudnn-frontend >= 1.27.0, and NVTE_FUSED_MLA_Q_UPROJ=1 required)" + ), + ) +) + + +def _dequantize_fused_output(query: MXFP8Tensor, s: int, b: int) -> torch.Tensor: + """Dequantize the rowwise fused output to bf16 [s, b, nh, head_dim]. + + TE's C++ dequantize kernel requires 2D layout, so reshape before calling dequantize(). + """ + tokens = s * b + q_2d = MXFP8Tensor( + shape=(tokens, PROJ_DIM), + dtype=torch.bfloat16, + rowwise_data=query._rowwise_data.view(tokens, PROJ_DIM), + rowwise_scale_inv=query._rowwise_scale_inv.view(tokens, PROJ_DIM // 32), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=query._quantizer, + requires_grad=False, + fp8_dtype=query._fp8_dtype, + with_gemm_swizzled_scales=False, + ) + return q_2d.dequantize().to(torch.bfloat16).view(s, b, NUM_HEADS, HEAD_DIM) + + +def _reference_q_uproj( + x: torch.Tensor, + w_mxfp8: MXFP8Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + s: int, + b: int, +) -> torch.Tensor: + """Unfused bf16 reference: dequantize-then-GEMM + RoPE. Returns [s, b, nh, head_dim] bf16.""" + x_dq = ( + MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)(x) + .dequantize() + .to(torch.bfloat16) + ) + w_dq = w_mxfp8.dequantize().to(torch.bfloat16) + out = (x_dq @ w_dq.t()).view(s, b, NUM_HEADS, HEAD_DIM) + + q_nope = out[..., :HEAD_DIM_NOPE] + q_rope = out[..., HEAD_DIM_NOPE:] + cos_ = cos[:, None, None, :].to(q_rope.dtype) + sin_ = sin[:, None, None, :].to(q_rope.dtype) + half = HEAD_DIM_ROPE // 2 + x1, x2 = q_rope[..., 0::2], q_rope[..., 1::2] + q_rope_out = torch.cat( + [ + x1 * cos_[..., :half] - x2 * sin_[..., :half], + x2 * cos_[..., half:] + x1 * sin_[..., half:], + ], + dim=-1, + ) + return torch.cat([q_nope, q_rope_out], dim=-1) + + +def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / ( + 10000 + ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) / HEAD_DIM_ROPE) + ) + freqs = torch.cat( + [torch.outer(torch.arange(tokens, device=device, dtype=torch.float32), inv_freq)] * 2, + dim=-1, + ) + return freqs.cos().to(torch.bfloat16), freqs.sin().to(torch.bfloat16) + + +@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) +@pytest.mark.parametrize("tokens", [256]) +def test_fused_mla_q_uproj(tokens: int) -> None: + """Forward numerics and x_saved properties for FusedMLAQUpProjRopeQuant.run(). + + Full forward+backward autograd testing (via _FusedMLAQUpProjFunction) lives in + Megatron-Core. + """ + s, b = tokens, 1 + device = torch.device("cuda") + torch.manual_seed(SEED) + torch.cuda.manual_seed(SEED) + + x = torch.randn(tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)( + torch.randn(PROJ_DIM, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + ) + cos, sin = _build_rope_tables(tokens, device) + + query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + + # Forward numerics: FP8 GEMM + output quantize introduce ~10% relative error. + fused_dq = _dequantize_fused_output(query, s, b) + ref_dq = _reference_q_uproj(x, w, cos, sin, s, b) + torch.testing.assert_close(fused_dq, ref_dq, atol=0.5, rtol=0.1) + + # x_saved: must be MXFP8 with only columnwise data retained for wgrad. + assert isinstance(x_saved, MXFP8Tensor) + assert x_saved._columnwise_data is not None, "x_saved must retain columnwise data for wgrad" + assert x_saved._rowwise_data is None, "x_saved rowwise data should be dropped after forward" diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index cdd98d2445..2a857a10dc 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -33,7 +33,8 @@ ) _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ( ModelConfig, reset_rng_states, diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 067f75d725..6e09ed316f 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,7 @@ import torch import torch.distributed as dist +from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -19,12 +20,12 @@ ep_dispatch, ep_combine, symm_mem_alloc, + release_symm_mem_pool, is_symm_backed, _ep_combine_raw, _ep_dispatch_raw, ) - ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" OVERFLOW = os.environ.get("NVTE_EP_OVERFLOW", "0") == "1" @@ -32,11 +33,13 @@ # Must come after the transformer_engine import so libtransformer_engine.so is loaded. import transformer_engine_torch as tex # noqa: F401 - NUM_LOCAL_EXPERTS = 2 -HIDDEN_DIM = 32 +# MXFP8 dispatch needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0. Defaults +# satisfy both so the MXFP8 tests run by default; override via NVTE_EP_HIDDEN_DIM / +# NVTE_EP_TOKENS_PER_RANK. +HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512")) TOP_K = 2 -TOKENS_PER_RANK = 4 +TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32")) def _zero_copy_test_include(fn): @@ -57,6 +60,18 @@ def _overflow_test_include(fn): return fn +# MXFP8 grouped dispatch needs a per-expert alignment of 128, but the EP backend caches a single +# alignment per process, so alignment=128 tests cannot share a process with the alignment=0 tests. +# They run in a dedicated pass (NVTE_EP_MXFP8_PASS=1) instead. +MXFP8_PASS = os.environ.get("NVTE_EP_MXFP8_PASS", "0") == "1" + + +def _mxfp8_align_test(fn): + """Mark a test that dispatches with alignment=128; runs only in the MXFP8 pass.""" + fn._mxfp8_align_test = True + return fn + + class _StageToSymm(torch.autograd.Function): """Identity op that stages ``src`` into a symm-mem buffer; grad passes through. Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine. @@ -121,6 +136,16 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _degroup_mxfp8(recv_grouped, valid_counts=None): + """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order. + With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert + slot; otherwise return every (padded) row.""" + parts = recv_grouped.split_into_quantized_tensors() + if valid_counts is None: + return torch.cat([p.dequantize() for p in parts], dim=0) + return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0) + + class _Cfg: rank: int world_size: int @@ -171,6 +196,16 @@ def setUpClass(cls): ) def setUp(self): + # alignment=128 MXFP8 tests run only in the dedicated MXFP8 pass; everything else skips + # there (and the MXFP8 tests skip outside it) since the backend pins one alignment/process. + is_mxfp8_align = getattr(getattr(self, self._testMethodName), "_mxfp8_align_test", False) + if MXFP8_PASS and not is_mxfp8_align: + self.skipTest("only alignment=128 MXFP8 tests run in the MXFP8 pass") + if not MXFP8_PASS and is_mxfp8_align: + self.skipTest("alignment=128 MXFP8 tests run in the dedicated MXFP8 pass") + # MXFP8 quantization requires Blackwell (SM 10.0) or newer. + if is_mxfp8_align and torch.cuda.get_device_capability() < (10, 0): + self.skipTest("MXFP8 EP tests require Blackwell (SM 10.0) or newer") # Only the zero-copy-capable tests run in the zero-copy pass. if ZERO_COPY and not getattr( getattr(self, self._testMethodName), "_zero_copy_test_include", False @@ -185,7 +220,13 @@ def setUp(self): ): self.skipTest("not exercised in overflow mode") - def _make_buffer(self, alignment=0, top_k=TOP_K): + def _make_buffer( + self, + alignment=0, + top_k=TOP_K, + dispatch_fwd_quant_recipe=None, + combine_bwd_quant_recipe=None, + ): return EpBuffer( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, @@ -193,13 +234,17 @@ def _make_buffer(self, alignment=0, top_k=TOP_K): num_local_experts=NUM_LOCAL_EXPERTS, recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank, alignment=alignment, + dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, + combine_bwd_quant_recipe=combine_bwd_quant_recipe, ) def _expert_out(self, expert_out): """Stage the combine input into symm-mem under zero-copy (combine requires it).""" if not ZERO_COPY: return expert_out - symm_buf = symm_mem_alloc(tuple(expert_out.shape), expert_out.dtype, self.ep_group) + symm_buf = symm_mem_alloc( + tuple(expert_out.shape), expert_out.dtype, self.ep_group, use_pool=True + ) return _StageToSymm.apply(expert_out, symm_buf) def _stage_grad_symm(self, x, symm_buf=None): @@ -372,6 +417,70 @@ def test_dispatch_autograd(self): tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2 ) + # MXFP8 dispatch + + def _mxfp8_quantizer(self): + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + + def _require_mxfp8_shapes(self): + if HIDDEN_DIM % 512 != 0 or TOKENS_PER_RANK % 32 != 0: + self.skipTest( + "MXFP8 needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0 " + "(set NVTE_EP_HIDDEN_DIM / NVTE_EP_TOKENS_PER_RANK)" + ) + + def _assert_mxfp8_matches_bf16(self, recv_mx, tokens, topk_idx, w, tc): + """Dequantized MXFP8 recv matches a bf16 dispatch of the same tokens. Both share the + alignment=128 padded expert-major layout, so compare the full prefix [0:sum(padded)].""" + ref_tokens = self._mxfp8_quantizer().quantize(tokens).dequantize() + ref_recv, _rw, _tc = ep_dispatch(self._make_buffer(alignment=128), ref_tokens, topk_idx, w) + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(recv_mx).float(), ref_recv.float()[:n], atol=1e-2, rtol=1e-2 + ) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_dispatch_mxfp8(self): + """MXFP8 dispatch quantizes bf16 tokens internally; recv (a per-expert GroupedTensor) + dequantized matches a bf16 dispatch of the same tokens. Under zero-copy the recv data and + scales are symm-mem backed.""" + self._require_mxfp8_shapes() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + if ZERO_COPY: + self.assertTrue(is_symm_backed(recv_mx.rowwise_data)) + self.assertTrue(is_symm_backed(recv_mx.scale_inv)) + self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + + @_zero_copy_test_include + @_mxfp8_align_test + def test_caller_provides_dispatch_recv_mxfp8(self): + """One caller-supplied buffer holds the recv data followed by the e8m0 scales; ep_dispatch + slices it and the returned GroupedTensor views the data and scale regions of that buffer.""" + self._require_mxfp8_shapes() + from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE + + rc = self.cfg.recv_capacity_per_rank + cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE + nbytes = rc * (HIDDEN_DIM + cols) # fp8 data + e8m0 scales, one byte per element + if ZERO_COPY: + recv_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group) + else: + recv_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w, recv_tokens=recv_buf) + # the returned GroupedTensor views the caller buffer's data then scale regions + self.assertEqual(recv_mx.rowwise_data.data_ptr(), recv_buf.data_ptr()) + self.assertEqual(recv_mx.scale_inv.data_ptr(), recv_buf.data_ptr() + rc * HIDDEN_DIM) + self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + @_zero_copy_test_include def test_caller_provides_dispatch_recv_tokens(self): """Caller-supplied recv_tokens (symm-mem-backed in zero-copy): ep_dispatch @@ -419,6 +528,85 @@ def test_caller_provides_grad_expert_out(self): # the caller-owned buffer was used as the combine-bwd scatter target self.assertGreater(gbuf.abs().sum().item(), 0.0) + @_zero_copy_test_include + @_mxfp8_align_test + def test_combine_bwd_mxfp8_caller_grad_out(self): + """MXFP8 combine backward into a single caller buffer sliced into data + e8m0 scales: the + returned per-expert GroupedTensor views those regions and, dequantized, matches a bf16 + combine backward reference on the same routing. Under zero-copy the caller buffer and combine + input are symm-mem backed.""" + self._require_mxfp8_shapes() + from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE + + rc = self.cfg.recv_capacity_per_rank + cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + eo_vals = ( + torch.linspace(-0.5, 0.5, rc * HIDDEN_DIM, device=self.cfg.device) + .reshape(rc, HIDDEN_DIM) + .to(torch.bfloat16) + ) + # MXFP8 combine backward writes into one caller buffer (data then e8m0 scales) + buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w) # seeds the routing + nbytes = rc * (HIDDEN_DIM + cols) + if ZERO_COPY: + grad_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group) + else: + grad_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + src_mx = eo_vals.detach().clone().requires_grad_(True) + out_mx = ep_combine(buf_mx, self._expert_out(src_mx), grad_out=grad_buf) + (0.5 * (out_mx.float() ** 2).sum()).backward() + g_mx = src_mx.grad # per-expert GroupedTensor viewing grad_buf + self.assertEqual(g_mx.rowwise_data.data_ptr(), grad_buf.data_ptr()) + self.assertEqual(g_mx.scale_inv.data_ptr(), grad_buf.data_ptr() + rc * HIDDEN_DIM) + # bf16 reference combine backward on the same routing + buf_bf = self._make_buffer(alignment=128) + ep_dispatch(buf_bf, tokens, topk_idx, w) + src_bf = eo_vals.detach().clone().requires_grad_(True) + out_bf = ep_combine(buf_bf, self._expert_out(src_bf)) + (0.5 * (out_bf.float() ** 2).sum()).backward() + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 + ) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_combine_bwd_mxfp8(self): + """MXFP8 combine backward with an internally allocated grad target: the returned per-expert + GroupedTensor, dequantized, matches a bf16 combine backward reference on the same routing. + Under zero-copy the combine input is symm-mem backed. + """ + self._require_mxfp8_shapes() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w) # seeds the routing + # Combine input rows match the recv total (per-step in eager, capacity otherwise). + rows = int(buf_mx.total_recv_tokens.item()) if EAGER else self.cfg.recv_capacity_per_rank + eo_vals = ( + torch.linspace(-0.5, 0.5, rows * HIDDEN_DIM, device=self.cfg.device) + .reshape(rows, HIDDEN_DIM) + .to(torch.bfloat16) + ) + src_mx = eo_vals.detach().clone().requires_grad_(True) + out_mx = ep_combine(buf_mx, self._expert_out(src_mx)) + (0.5 * (out_mx.float() ** 2).sum()).backward() + g_mx = src_mx.grad # per-expert GroupedTensor + # bf16 reference combine backward on the same routing + buf_bf = self._make_buffer(alignment=128) + ep_dispatch(buf_bf, tokens, topk_idx, w) + src_bf = eo_vals.detach().clone().requires_grad_(True) + out_bf = ep_combine(buf_bf, self._expert_out(src_bf)) + (0.5 * (out_bf.float() ** 2).sum()).backward() + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 + ) + @_zero_copy_test_include def test_zero_copy_pool_auto_alloc(self): """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO @@ -527,18 +715,29 @@ def _run_1f1b(self, capture): recv = [None, None, None] # Per-microbatch grad-staging buffers, symm-mem under zero-copy and - # pre-allocated so nothing is allocated/freed mid-interleave. The recv - # outputs are owned by each EpBuffer (symm-mem under zero-copy). + # pre-allocated so nothing is allocated/freed mid-interleave. recv_w = [None, None, None] rc = self.cfg.recv_capacity_per_rank if ZERO_COPY: gbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales] gbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales] + # Persistent symm-mem recv buffers per microbatch: leaving recv None + # pool-allocates, which is not CUDA-graph capturable. + rbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales] + rbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales] else: gbuf_t = gbuf_w = [None, None, None] + rbuf_t = rbuf_w = [None, None, None] def fwd(k): - rt, rw, _ = ep_dispatch(buffers[k], tokens_p[k], idx, w) + rt, rw, _ = ep_dispatch( + buffers[k], + tokens_p[k], + idx, + w, + recv_tokens=rbuf_t[k], + recv_topk_weights=rbuf_w[k], + ) recv[k] = self._stage_grad_symm(rt, gbuf_t[k]) recv_w[k] = self._stage_grad_symm(rw, gbuf_w[k]) @@ -631,5 +830,7 @@ def _init_distributed(): result = runner.run(suite) dist.barrier() ep_finalize() + # Deregister symm-mem windows while the comm is still valid. + release_symm_mem_pool() dist.destroy_process_group() sys.exit(0 if result.wasSuccessful() else 1) diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 10e814eb80..5788c2ac9f 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -47,11 +47,13 @@ run_pass() { local zc="$2" local eager="${3:-0}" local overflow="${4:-0}" + local mxfp8="${5:-0}" local log="stdout_ep_${label}.txt" echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun. - NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" setsid timeout --foreground \ - --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ + NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" \ + NVTE_EP_MXFP8_PASS="${mxfp8}" \ + setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ "${SCRIPT}" 2>&1 | tee "${log}" local rc=${PIPESTATUS[0]} @@ -72,5 +74,11 @@ run_pass "default" 0 run_pass "zero_copy" 1 run_pass "eager" 0 1 run_pass "overflow" 0 0 1 +# MXFP8 grouped dispatch pins the per-expert alignment to 128, which the backend caches +# process-wide, so its tests get their own passes (normal + zero-copy + eager IO). mxfp8 is the +# 5th arg; eager is the 3rd. +run_pass "mxfp8" 0 0 0 1 +run_pass "mxfp8_zero_copy" 1 0 0 1 +run_pass "mxfp8_eager" 0 1 0 1 exit $RET diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py index 9d08a81a8f..d733286093 100644 --- a/tests/pytorch/distributed/test_fusible_ops.py +++ b/tests/pytorch/distributed/test_fusible_ops.py @@ -32,7 +32,8 @@ # Import utility functions _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import dtype_tols, make_recipe, quantization_tols diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index eb43ba7e75..07dffebf5f 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -37,7 +37,8 @@ # Import utility functions _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import dtype_tols, make_recipe, run_distributed, str_to_dtype # Check if FP8 is supported diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py index 2e7a63e0a2..1b8bf5890f 100644 --- a/tests/pytorch/distributed/test_sanity.py +++ b/tests/pytorch/distributed/test_sanity.py @@ -19,7 +19,8 @@ from transformer_engine.common import recipe _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig model_configs = { diff --git a/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py b/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py new file mode 100644 index 0000000000..64bfe6012e --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import MXFP8Quantizer + +recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_dequantize_extreme_e8m0_scale_codes() -> None: + """UE8M0 scale code 0 is 2^-127 and code 255 is NaN, not 0.0 and +Inf.""" + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, columnwise=False) + x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + qx = quantizer(x) + data = qx._rowwise_data.view(torch.uint8) + scales = qx._rowwise_scale_inv.view(torch.uint8) + data[0, :32] = 56 + scales[0, 0] = 0 + data[1, :32] = 56 + scales[1, 0] = 255 + y = qx.dequantize(dtype=torch.float32) + assert (y[0, :32] == 2.0**-127).all() + assert torch.isnan(y[1, :32]).all() 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 d07953ce37..23662eb0c3 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py +++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py @@ -60,6 +60,13 @@ def generate_split_sections(M: int, N: int, edge_cases: str) -> list[int]: split_sections = [avg_split] * (num_chunks - 2) + [0] + [avg_split * 2] elif edge_cases == "random_uneven_split": split_sections = generate_random_multiples_sum(M, num_chunks, least_multiple) + elif edge_cases == "imbalanced_avg_misaligned": + # Three groups whose 128-aligned sizes average to a non-128-aligned value, so any + # buffer sized from num_groups * round_up(M // num_groups, 128) diverges from the + # per-group padded sum. "random_uneven_split" cannot catch that: with 4 chunks the + # average only depends on M, which is always 128 * num_chunks aligned here. + assert M >= 3 * least_multiple, "M too small for the imbalanced case" + split_sections = [least_multiple, least_multiple, M - 2 * least_multiple] else: raise ValueError(f"Invalid edge case: {edge_cases}") @@ -469,3 +476,419 @@ def test_grouped_tensor_mxfp8_with_paged_stashing( valid_M=valid_M, optimize_for_gemm=optimize_for_gemm, ) + + +# --------------------------------------------------------------------------------------------- +# Pre-quantized MXFP8 input (FP8 token dispatch) +# +# tex.group_requantize_inplace takes a grouped tensor that arrives +# ALREADY rowwise-quantized (its high-precision form no longer exists), and makes it GEMM-ready +# in both directions: the rowwise data passes through verbatim with its scales swizzled, and the +# columnwise copy is rebuilt via dequantize + columnwise-only requantize. +# +# These mirror the edge-case matrices of test_grouped_tensor_mxfp8_versus_reference and +# test_grouped_tensor_mxfp8_with_paged_stashing so the same shapes, zero-token placements and +# uneven splits exercise this path. +# --------------------------------------------------------------------------------------------- + + +def make_prequantized_wire_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor): + """Rowwise-only, unswizzled grouped tensor, as FP8 dispatch delivers it.""" + wire_quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False) + # Must stay unswizzled: the requantize path asserts on it and dequantize needs compact scales. + wire_quantizer.optimize_for_gemm = False + wire = fused_grouped_quantize(x, split_section_tensor, wire_quantizer) + assert wire.columnwise_data is None + assert not wire._with_gemm_swizzled_scales + return wire + + +def make_op_quantizer(columnwise: bool): + """The op's input quantizer, configured the way the ops layer configures it. + + ``columnwise`` mirrors ``weight_requires_grad``: it tells the helper whether a wgrad GEMM + will consume a columnwise copy. + """ + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=columnwise) + quantizer.optimize_for_gemm = True + return quantizer + + +def make_gemm_ready_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor): + """Grouped tensor already GEMM-ready in both directions (swizzled scales).""" + quantizer = make_op_quantizer(columnwise=True) + tensor = fused_grouped_quantize(x, split_section_tensor, quantizer) + assert tensor.columnwise_data is not None + assert tensor._with_gemm_swizzled_scales + return tensor + + +def check_prequantized_requantize_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + split_sections: list[int], +) -> None: + """Run the pre-quantized requantize path and check both directions against a reference. + + The reference is derived from dequantize(wire), not from the original high-precision x: that + is the only data a consumer can see after dispatch, and MXFP8 rowwise requantization is + idempotent, so it is an exact reference rather than an approximate one. + """ + device = "cuda" + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + # The buffer is always M rows. Paged stashing is just the case where the groups cover fewer + # than M of them (valid_M < M) and the tail holds garbage the kernels must leave alone; the + # non-paged case is the same code path with sum(split_sections) == M. + x = torch.randn((M, N), dtype=x_dtype, device=device) + split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device=device) + num_groups = len(split_sections) + # Rows the groups actually cover. Beyond this the buffers hold whatever the allocator handed + # out, so nothing past it may be compared. + valid_rows = sum(split_sections) + + wire = make_prequantized_wire_tensor(x, split_section_tensor) + + # Snapshot what must survive verbatim, plus the compact scales the reference swizzles. + rowwise_scale_shape_before = wire.scale_inv.shape + rowwise_data_before = wire.rowwise_data.clone() + wire_splits_before = [ + (t._rowwise_data.view(dtype=torch.uint8).clone(), t._rowwise_scale_inv.clone()) + for t in wire.split_into_quantized_tensors() + ] + + # Reference high-precision input: everything downstream is derived from this. Only the live + # rows are kept -- dequantize allocates M rows but writes only the ones the groups cover. + dequantized_ref = ( + tex.group_dequantize(wire, te.DType.kBFloat16) + .rowwise_data.view(M, N)[:valid_rows, :] + .clone() + ) + + # Reference columnwise copy, quantized per group from the dequantized data. + colwise_quantizers = [ + MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=False, columnwise=True) + for _ in range(num_groups) + ] + _, _, colwise_data_ref, colwise_scale_ref = reference_group_quantize( + dequantized_ref, + colwise_quantizers, + split_sections, + return_rowwise=False, + return_transpose=True, + ) + + # ---- the code under test ---- + # The op's quantizer, configured as the ops layer does: columnwise_usage says a wgrad GEMM + # will consume the columnwise copy, so the helper builds it and switches rowwise off itself. + dequantized = tex.group_requantize_inplace( + wire, + make_op_quantizer(columnwise=True), + num_groups, + split_section_tensor, + te.DType.kBFloat16, + return_dequantized=True, + ) + + assert wire.columnwise_data is not None, "columnwise data must be built" + assert wire.columnwise_scale_inv is not None, "columnwise scales must survive the swizzle" + + # The rowwise swizzle is size-preserving; the per-group content checks below cannot catch a + # wrongly sized buffer because they read through offsets derived from the splits, so the + # capacity must be checked explicitly. + assert ( + wire.scale_inv.shape == rowwise_scale_shape_before + ), "the swizzled rowwise scale buffer must keep the compact buffer's capacity" + + # The returned dequantized tensor is what bias gradients are reduced from. Compare only the + # live rows: both this and the reference allocate M rows but write only the covered ones, and + # their tails are separate uninitialized allocations. + torch.testing.assert_close(dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0) + + # The rowwise DATA must pass through untouched; only its scales are re-laid-out. + torch.testing.assert_close(wire.rowwise_data, rowwise_data_before, atol=0.0, rtol=0.0) + + if valid_rows > 0: + # A tensor whose groups are all empty has no scales to lay out, so the swizzle is a no-op + # and leaves the flag unset; every other case must come back swizzled. + assert wire._with_gemm_swizzled_scales, "rowwise scales must be marked swizzled" + + # Per-group comparison, same structure as check_grouped_tensor_mxfp8_versus_reference. + outputs = wire.split_into_quantized_tensors() + x_splits = torch.split(dequantized_ref, split_sections) + + for i, out in enumerate(outputs): + rows_i = split_sections[i] + scale_before = wire_splits_before[i][1] + colwise_data = out._columnwise_data.view(dtype=torch.uint8) + colwise_scale = out._columnwise_scale_inv + + if rows_i == 0: + # Buffers for empty groups are never written, so only shape and dtype are meaningful. + assert_same_shape_and_dtype(colwise_data, colwise_data_ref[i]) + assert_same_shape_and_dtype(colwise_scale, colwise_scale_ref[i]) + continue + + # Rowwise scales: the swizzled form of the compact scales this group arrived with. The + # rowwise DATA is covered by the whole-buffer identity check above. + torch.testing.assert_close( + out._rowwise_scale_inv, + swizzle_mxfp8_scale(rows_i, N, scale_before, columnwise=False), + atol=0.0, + rtol=0.0, + ) + + # Columnwise: rebuilt from the dequantized data, and swizzled by the quantize kernel + # because the caller sets optimize_for_gemm. + torch.testing.assert_close(colwise_data, colwise_data_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, True) + assert ( + valid_scale_shape == colwise_scale.shape + ), "The columnwise scale shape is not correctly aligned" + torch.testing.assert_close( + colwise_scale, + swizzle_mxfp8_scale(rows_i, N, colwise_scale_ref[i], columnwise=True), + atol=0.0, + rtol=0.0, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # edge case, zero tokens for all + (0, 512), + # full tile cases + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 8192), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + "imbalanced_avg_misaligned", + ], +) +def test_prequantized_requantize_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, +) -> None: + split_sections = generate_split_sections(M, N, edge_cases) + check_prequantized_requantize_versus_reference( + x_dtype=x_dtype, + M=M, + N=N, + split_sections=split_sections, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # M won't be empty in paged stashing + (1024, 256), + (8192, 1024), + (16384, 8192), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + "zero_tokens_all", + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + "imbalanced_avg_misaligned", + ], +) +def test_prequantized_requantize_with_paged_stashing( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, +) -> None: + # Paged stashing: the buffer holds M rows but only valid_M carry live tokens; the rest is + # garbage the kernels must not touch. + valid_M = 0 if edge_cases == "zero_tokens_all" else M // 2 + split_sections = generate_split_sections(valid_M, N, edge_cases) + assert sum(split_sections) == valid_M + + check_prequantized_requantize_versus_reference( + x_dtype=x_dtype, + M=M, + N=N, + split_sections=split_sections, + ) + + +def _requantize_setup(M: int = 1024, N: int = 256): + """Common inputs for the state-dispatch tests below.""" + torch.manual_seed(0) + split_sections = [M // 4] * 4 + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + return x, torch.tensor(split_sections, dtype=torch.int64, device="cuda"), len(split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_passes_through_gemm_ready_input(): + """A tensor already GEMM-ready in both directions is left untouched.""" + x, splits, num_groups = _requantize_setup() + tensor = make_gemm_ready_tensor(x, splits) + rowwise_before = tensor.rowwise_data.clone() + columnwise_before = tensor.columnwise_data.clone() + scale_before = tensor.scale_inv.clone() + + out = tex.group_requantize_inplace( + tensor, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16 + ) + + assert out is None + assert tensor._with_gemm_swizzled_scales + torch.testing.assert_close(tensor.rowwise_data, rowwise_before, atol=0.0, rtol=0.0) + torch.testing.assert_close(tensor.columnwise_data, columnwise_before, atol=0.0, rtol=0.0) + torch.testing.assert_close(tensor.scale_inv, scale_before, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_skips_columnwise_when_not_needed(): + """columnwise_usage=False (frozen weights) swizzles rowwise without building columnwise.""" + x, splits, num_groups = _requantize_setup() + wire = make_prequantized_wire_tensor(x, splits) + + out = tex.group_requantize_inplace( + wire, make_op_quantizer(columnwise=False), num_groups, splits, te.DType.kBFloat16 + ) + + assert out is None + assert wire._with_gemm_swizzled_scales, "the GEMM still needs swizzled rowwise scales" + assert wire.columnwise_data is None, "no wgrad GEMM, so no columnwise copy should be built" + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_rejects_dequantized_from_gemm_ready_input(): + """Bias grads cannot be served from an already-swizzled input, so this must raise.""" + x, splits, num_groups = _requantize_setup() + tensor = make_gemm_ready_tensor(x, splits) + + with pytest.raises(RuntimeError, match="compact format"): + tex.group_requantize_inplace( + tensor, + make_op_quantizer(columnwise=True), + num_groups, + splits, + te.DType.kBFloat16, + return_dequantized=True, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_rejects_swizzled_without_columnwise(): + """Swizzled rowwise scales with no columnwise copy: it can no longer be rebuilt.""" + x, splits, num_groups = _requantize_setup() + wire = make_prequantized_wire_tensor(x, splits) + # Swizzle in place, leaving the tensor rowwise-only. + tex.grouped_swizzle_for_gemm(wire, True, False) + + with pytest.raises(RuntimeError, match="cannot be rebuilt"): + tex.group_requantize_inplace( + wire, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16 + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_rejects_dtype_mismatch(): + """The helper keeps the input's format; it does not convert between formats.""" + x, splits, num_groups = _requantize_setup() + wire = make_prequantized_wire_tensor(x, splits) + mismatched = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E5M2, rowwise=True, columnwise=True) + mismatched.optimize_for_gemm = True + + with pytest.raises(RuntimeError, match="dtype"): + tex.group_requantize_inplace(wire, mismatched, num_groups, splits, te.DType.kBFloat16) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise", "columnwise"]) +@pytest.mark.parametrize( + "split_sections", + [ + # Imbalanced groups whose mean is not 128-aligned: sizing the swizzled output from + # num_tensors * round_up(mean, 128) would give 3 * 512 = 1536 rows instead of the + # input's 1280. + [128, 128, 1024], + # Same property with a zero-token group in the mix. + [0, 256, 128, 1152], + ], + ids=["imbalanced", "imbalanced_with_empty"], +) +def test_grouped_swizzle_variable_shape_preserves_scale_capacity( + columnwise: bool, split_sections: list[int] +): + """Swizzling a variable-shape grouped tensor must keep the scale buffer's exact capacity. + + The swizzle kernel walks input and output with identical per-group padded strides, so the + operation is size-preserving. Sizing the output from the per-tensor average instead breaks + any consumer that derives per-group offsets from the split sizes, because the average of + 128-aligned group sizes is generally not 128-aligned. The content checks alone cannot catch + a wrong allocation (per-group offsets are derived from the splits, not the buffer), so the + shape assertion is the actual regression check. + """ + torch.manual_seed(0) + N = 256 + M = sum(split_sections) + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + splits = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + quantizer = MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise + ) + # Compact (unswizzled) scales, as they arrive over the wire. + quantizer.optimize_for_gemm = False + tensor = fused_grouped_quantize(x, splits, quantizer) + assert not tensor._with_gemm_swizzled_scales + + scale_attr = "columnwise_scale_inv" if columnwise else "scale_inv" + compact_shape = getattr(tensor, scale_attr).shape + compact_groups = [ + (t._columnwise_scale_inv if columnwise else t._rowwise_scale_inv).clone() + for t in tensor.split_into_quantized_tensors() + ] + + tex.grouped_swizzle_for_gemm(tensor, not columnwise, columnwise) + + assert tensor._with_gemm_swizzled_scales + assert ( + getattr(tensor, scale_attr).shape == compact_shape + ), "grouped swizzle must preserve the scale buffer's shape for variable-shape tensors" + + # The swizzled content of each group must match a per-group dense swizzle of the compact + # scales it arrived with. + for rows_i, group_before, out in zip( + split_sections, compact_groups, tensor.split_into_quantized_tensors() + ): + if rows_i == 0: + continue + out_scale = out._columnwise_scale_inv if columnwise else out._rowwise_scale_inv + torch.testing.assert_close( + out_scale, + swizzle_mxfp8_scale(rows_i, N, group_before, columnwise=columnwise), + atol=0.0, + rtol=0.0, + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py index 127b487650..16a2d75de6 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py +++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py @@ -129,3 +129,36 @@ def test_mxfp8_quantize_swizzle_fusion( return_rowwise=return_rowwise, return_transpose=return_transpose, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("M, N", [(96, 160), (4096, 576), (4096, 2112)]) +def test_mxfp8_bidirectional_swizzled_row_scale_padding(M: int, N: int) -> None: + """The specialized bidirectional kernel must not overwrite padded row scales.""" + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + quantizer.optimize_for_gemm = True + scale = quantizer(x)._rowwise_scale_inv.view(torch.uint8) + + scale_rows = torch.arange(M, device=scale.device, dtype=torch.int64).view(-1, 1) + scale_cols = torch.arange(N // 32, device=scale.device, dtype=torch.int64).view(1, -1) + num_tiles_x = math.ceil(N / 128) + scale_indices = ( + ((scale_rows // 128) * num_tiles_x + scale_cols // 4) * (128 * 4) + + (scale_rows % 32) * 16 + + ((scale_rows % 128) // 32) * 4 + + scale_cols % 4 + ) + valid_mask = torch.zeros(scale.numel(), dtype=torch.bool, device=scale.device) + valid_mask[scale_indices.view(-1)] = True + + torch.testing.assert_close( + scale.view(-1)[~valid_mask], + torch.zeros_like(scale.view(-1)[~valid_mask]), + atol=0, + rtol=0, + ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py index b14eeb815b..2ed496c730 100755 --- a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py @@ -431,3 +431,39 @@ def test_group_stochastic_rounding_quantization_versus_reference( num_splits=num_splits, use_tex_split_quantize=use_tex_split_quantize, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("use_2D", [False, True], ids=str) +def test_stochastic_rounding_picks_an_adjacent_fp4_value( + x_dtype: torch.dtype, use_2D: bool +) -> None: + device = "cuda" + torch.manual_seed(seed) + M, N = 512, 1024 + x = torch.randn((M, N), dtype=x_dtype, device=device) + amax = torch.max(torch.abs(x)).float() + + q, s, q_t, s_t = quantize_fp4(x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False) + q_redraw, _, q_t_redraw, _ = quantize_fp4( + x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False + ) + assert not torch.equal(q, q_redraw), "two stochastic rounding draws are identical" + assert not torch.equal(q_t, q_t_redraw), "two stochastic rounding draws are identical" + + def check_adjacent(qx: torch.Tensor, sx: torch.Tensor, reference: torch.Tensor) -> None: + # The kernel rounds reference / block_scale onto the E2M1 grid, so every + # output has to be one of the two grid values that bracket it, i.e. within + # one grid step. An all-zero output fails this everywhere the input is not + # already near zero. + block_scale = sx.repeat_interleave(16, dim=1).view(torch.float8_e4m3fn).to(torch.float32) + block_scale = block_scale[: reference.shape[0], : reference.shape[1]] * (amax / (6.0 * 448)) + exact = (reference.float() / block_scale).clamp(-6.0, 6.0) + magnitude = exact.abs() + step = torch.where(magnitude >= 4.0, 2.0, torch.where(magnitude >= 2.0, 1.0, 0.5)) + gap = (fp4_to_fp32(unpack_fp4(qx)) - exact).abs() + assert torch.all(gap <= step * 1.0001), f"largest gap to the exact value is {gap.max()}" + + check_adjacent(q, s, x) + check_adjacent(q_t, s_t, x.t().contiguous()) diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index 6832ef89dd..50dc8e724c 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -166,6 +166,31 @@ def test_frozen_model(self): torch.testing.assert_close(ref_param, tst_param) + @pytest.mark.parametrize("capturable", [False, True]) + def test_empty_param_group_advances_step(self, capturable): + # An empty param group must advance its step counter like a populated one. + # The same group can be empty on one data-parallel rank and populated on + # another, and "step" is part of the checkpoint, so a group that stops + # counting makes a resumed run apply a stale bias correction. + param = torch.nn.Parameter(torch.rand(4, dtype=torch.float, device="cuda")) + tst_optim = self.fused_optim( + [{"params": [param]}, {"params": []}], capturable=capturable, **self.options + ) + + num_steps = 3 + for _ in range(num_steps): + param.grad = torch.rand_like(param) + tst_optim.step() + + populated_step, empty_step = (int(g["step"]) for g in tst_optim.param_groups) + assert populated_step == num_steps + assert empty_step == populated_step + + # The counter is checkpointed through the param groups, so the empty group + # has to round trip with the same value as the populated one. + checkpoint = tst_optim.state_dict() + assert int(checkpoint["param_groups"][1]["step"]) == num_steps + def test_empty_param_at_end_of_group(self): tensors = [ torch.ones(4, dtype=torch.float, device="cuda"), diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 10d0dbd227..1dd2b16ade 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -14,6 +14,8 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, Float8Quantizer, Fp8Padding, Fp8Unpadding, @@ -30,6 +32,11 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) +from transformer_engine.pytorch.constants import TE_DType +from transformer_engine.pytorch.module.grouped_linear import ( + _GroupedLinear, + is_module_grouped_tensor_path_supported, +) from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, get_align_size_for_quantization, @@ -1154,46 +1161,124 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_b torch.testing.assert_close(o, o_ref, **tols) +@pytest.mark.parametrize("use_bias_scale", [False, True]) +def test_grouped_gemm_grouped_tensor_zero_work_bias(use_bias_scale) -> None: + """A grouped bias operation is a no-op when every group has zero rows. + + Zero-sized CUDA tensors may legally have a null data pointer. Exercise both bias entry points + so neither the ordinary nor scaled path mistakes that pointer for missing output storage. + This BF16 case runs on both Hopper and Blackwell when grouped cuBLASLt GEMM is available. + """ + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor GEMM is unavailable.") + + num_groups = 4 + in_features = 256 + out_features = 256 + m_sizes = [0] * num_groups + dtype = torch.bfloat16 + device = torch.device("cuda") + + weights = [ + torch.randn(out_features, in_features, dtype=dtype, device=device) + for _ in range(num_groups) + ] + biases = [torch.randn(1, out_features, dtype=dtype, device=device) for _ in range(num_groups)] + grouped_weights = _make_grouped_tensor_uniform( + num_groups, out_features, in_features, device, dtype + ) + grouped_input = _make_grouped_tensor_from_splits(m_sizes, in_features, device, dtype) + grouped_output = _make_grouped_tensor_from_splits(m_sizes, out_features, device, dtype) + grouped_bias = _make_grouped_tensor_uniform(num_groups, 1, out_features, device, dtype) + _pack_grouped_tensor(grouped_weights, weights) + _pack_grouped_tensor(grouped_bias, biases) + + bias_scale = torch.empty(0, dtype=torch.float32, device=device) if use_bias_scale else None + general_grouped_gemm_for_grouped_tensor( + grouped_weights, + grouped_input, + grouped_output, + layout="TN", + bias=grouped_bias, + bias_scale=bias_scale, + ) + torch.cuda.synchronize() + + assert grouped_output.rowwise_data.numel() == 0 + + @pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) @pytest.mark.parametrize("accumulate", [False, True]) -@pytest.mark.parametrize("quant_type", ["bf16", "mxfp8"]) +@pytest.mark.parametrize( + "quant_type", ["bf16", "fp8_current_scaling", "mxfp8", "fp8_block_scaling"] +) def test_grouped_gemm_grouped_tensor_zero_work(layout, accumulate, quant_type) -> None: """Grouped GEMM with all-zero split sizes (zero total work). For wgrad (NT layout) the output should be zero when not accumulating, or unchanged when accumulating with beta=1. """ - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") if not is_bf16_available(): pytest.skip("bfloat16 is required for grouped GEMM test.") - if quant_type == "mxfp8" and not mxfp8_available: - pytest.skip(reason_for_no_mxfp8) z = 4 k, n = 256, 256 dtype = torch.bfloat16 device = torch.device("cuda") - use_mxfp8 = quant_type == "mxfp8" + + test_recipe = { + "bf16": None, + "fp8_current_scaling": recipe.Float8CurrentScaling(), + "mxfp8": recipe.MXFP8BlockScaling(), + "fp8_block_scaling": recipe.Float8BlockScaling(), + }[quant_type] + if not is_module_grouped_tensor_path_supported( + test_recipe, + dtype, + ): + pytest.skip(f"{quant_type} grouped-tensor GEMM is unavailable") transa = layout[0] == "T" transb = layout[1] == "T" zero_first_dims = torch.zeros(z, dtype=torch.int64, device=device) + def _make_quantizer(fp8_dtype, rowwise, columnwise): + if quant_type == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device=device) + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + elif quant_type == "mxfp8": + quantizer = MXFP8Quantizer( + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + ) + elif quant_type == "fp8_block_scaling": + quantizer = Float8BlockQuantizer( + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + else: + raise ValueError(f"Unsupported quantized zero-work test type {quant_type}") + quantizer.optimize_for_gemm = True + return quantizer + def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): """Create a GroupedTensor with non-zero logical_shape but zero first_dims.""" buf = torch.randn(0, logical_last_dim, dtype=dtype, device=device) - if use_mxfp8: + if test_recipe is not None: if is_a: rowwise, columnwise = transa, not transa else: rowwise, columnwise = not transb, transb - quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=rowwise, - columnwise=columnwise, - ) - quantizer.optimize_for_gemm = True + fp8_dtype = TE_DType[torch.float8_e4m3fn] + quantizer = _make_quantizer(fp8_dtype, rowwise, columnwise) return tex.group_quantize(buf, quantizer, z, zero_first_dims) return GroupedTensor.make_grouped_tensor( num_tensors=z, @@ -1208,12 +1293,18 @@ def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): if layout in ("TN", "NN"): weight_tensors = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] - if use_mxfp8: - grouped_A = _make_grouped_tensor_quantized_mxfp8( - weight_tensors, + if test_recipe is not None: + grouped_weight = torch.cat(weight_tensors, dim=0) + weight_quantizer = _make_quantizer( + TE_DType[torch.float8_e4m3fn], rowwise=transa, columnwise=not transa, - device=device, + ) + grouped_A = tex.group_quantize( + grouped_weight, + weight_quantizer, + z, + torch.full((z,), n, dtype=torch.int64, device=device), ) else: grouped_A = _make_grouped_tensor_uniform(z, n, k, device, dtype) @@ -1506,16 +1597,638 @@ def test_fp8_grouped_gemm(shape, accumulate): @pytest.fixture(autouse=True) def _reset_fp8_state(monkeypatch): - monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "0") + monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) yield FP8GlobalStateManager.reset() monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) +@pytest.mark.parametrize( + "m_splits,exception", + [([256, 256], ValueError), (torch.tensor([256, 256]), ValueError)], + ids=["python-list", "cpu-tensor"], +) +def test_single_grouped_weight_rejects_host_m_splits(monkeypatch, m_splits, exception): + """A single parent parameter must never fall back to host-split per-expert GEMMs.""" + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("Native GroupedTensor GEMM is unavailable on this system.") + + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + grouped_linear = GroupedLinear( + 2, + 64, + 64, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + x = torch.randn(512, 64, dtype=torch.bfloat16, device="cuda") + with pytest.raises(exception, match="requires.*CUDA"): + grouped_linear(x, m_splits) + + +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + ], +) +def test_single_grouped_weight_matches_discrete_grouped_tensor_path(monkeypatch, fp8_recipe): + """Match single and discrete weights while both use CUDA m_splits and grouped GEMM.""" + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + torch.bfloat16, + ): + pytest.skip("Recipe is not supported with a single grouped weight on this system.") + + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + + num_gemms = 3 + in_features = 256 + out_features = 256 + m_splits = torch.tensor([256, 512, 256], dtype=torch.int64, device="cuda") + total_tokens = int(m_splits.sum()) + weights = torch.randn( + num_gemms, + out_features, + in_features, + dtype=torch.bfloat16, + device="cuda", + ) + + discrete = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + use_grouped_tensor=True, + ) + single = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + with torch.no_grad(): + for idx in range(num_gemms): + getattr(discrete, f"weight{idx}").copy_(weights[idx]) + single.weight.rowwise_data.view_as(weights).copy_(weights) + + x = torch.randn(total_tokens, in_features, dtype=torch.bfloat16, device="cuda") + dy = torch.randn(total_tokens, out_features, dtype=torch.bfloat16, device="cuda") + x_discrete = x.detach().clone().requires_grad_(True) + x_single = x.detach().clone().requires_grad_(True) + + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + y_discrete = discrete(x_discrete, m_splits) + y_single = single(x_single, m_splits) + y_discrete.backward(dy) + y_single.backward(dy) + + tolerances = dict(rtol=1e-2, atol=5e-3) + torch.testing.assert_close(y_single.float(), y_discrete.float(), **tolerances) + torch.testing.assert_close(x_single.grad.float(), x_discrete.grad.float(), **tolerances) + discrete_wgrad = torch.stack( + [getattr(discrete, f"weight{idx}").grad for idx in range(num_gemms)] + ) + torch.testing.assert_close(single.weight.grad.float(), discrete_wgrad.float(), **tolerances) + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +def test_single_grouped_weight_mxfp8_workspace_cache(monkeypatch): + """BF16 primary weights update one persistent MXFP8 grouped workspace per iteration.""" + mxfp8_recipe = recipe.MXFP8BlockScaling() + if not is_module_grouped_tensor_path_supported( + mxfp8_recipe, + torch.bfloat16, + ): + pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + + with autocast(enabled=True, recipe=mxfp8_recipe): + grouped_linear(x, m_splits, is_first_microbatch=True) + workspace = grouped_linear._fp8_workspaces["weight"] + assert isinstance(workspace, GroupedTensor) + pointers = ( + workspace.rowwise_data.data_ptr(), + workspace.columnwise_data.data_ptr(), + workspace.scale_inv.data_ptr(), + workspace.columnwise_scale_inv.data_ptr(), + ) + old_data = workspace.rowwise_data.clone() + + with torch.no_grad(): + grouped_linear.weight.rowwise_data.add_(1) + grouped_linear(x, m_splits, is_first_microbatch=False) + assert torch.equal(workspace.rowwise_data, old_data) + + grouped_linear(x, m_splits, is_first_microbatch=True) + assert pointers == ( + workspace.rowwise_data.data_ptr(), + workspace.columnwise_data.data_ptr(), + workspace.scale_inv.data_ptr(), + workspace.columnwise_scale_inv.data_ptr(), + ) + assert not torch.equal(workspace.rowwise_data, old_data) + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +@pytest.mark.parametrize("fp8_recipe", [recipe.MXFP8BlockScaling()], ids=recipe_id) +def test_single_grouped_weight_with_disabled_weight_preswizzle(monkeypatch, fp8_recipe): + """Grouped weight preparation preserves a disabled preswizzle decision.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + with quantized_model_init(enabled=True, recipe=fp8_recipe): + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + weight_quantizer = grouped_linear.weight.quantizer + assert weight_quantizer is not None + preswizzle = grouped_linear._enable_weight_preswizzle( + weight_quantizer, + grouped_linear.weight, + ) + assert preswizzle is False + weight_quantizer.optimize_for_gemm = preswizzle + grouped_weight, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm( + (grouped_linear.weight,), + [weight_quantizer], + [None], + num_gemms=grouped_linear.num_gemms, + single_grouped_weight=True, + with_quantized_compute=True, + columnwise_usage=True, + activation_dtype=torch.bfloat16, + is_first_microbatch=True, + skip_fp8_weight_update=None, + cache_weight=True, + ) + + assert weight_quantizer.optimize_for_gemm is False + assert len(new_workspaces) == 1 + assert new_workspaces[0] is None + assert grouped_weight is grouped_linear.weight + assert hasattr(grouped_weight, "_with_gemm_swizzled_scales") + assert grouped_weight._with_gemm_swizzled_scales is False + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +def test_single_grouped_primary_mxfp8_bypasses_weight_workspace(monkeypatch): + """An MXFP8 primary grouped parameter is already GEMM-ready and is not requantized.""" + mxfp8_recipe = recipe.MXFP8BlockScaling() + if not is_module_grouped_tensor_path_supported( + mxfp8_recipe, + torch.bfloat16, + ): + pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + with quantized_model_init(enabled=True, recipe=mxfp8_recipe): + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + with torch.no_grad(), autocast(enabled=True, recipe=mxfp8_recipe): + grouped_linear(x, m_splits, is_first_microbatch=True) + assert grouped_linear.weight.quantizer is not None + assert "weight" not in grouped_linear._fp8_workspaces + + def _clone_outputs(outputs): return [None if out is None else out.detach().clone() for out in outputs] +def _grouped_linear_weight_params(module): + if module.single_grouped_weight: + return [module.weight] + return [getattr(module, f"weight{i}") for i in range(module.num_gemms)] + + +def _grouped_linear_bias_params(module): + if not module.use_bias: + return [] + if module.single_grouped_bias: + return [module.bias] + return [getattr(module, f"bias{i}") for i in range(module.num_gemms)] + + +def _run_grouped_parameter_layout( + *, + use_grouped_tensor, + fp8_recipe, + single_grouped_weight, + single_grouped_bias, + use_bias, + delay_wgrad_compute, + fuse_wgrad_accumulation, + x_base, + dy, + weights, + biases, + m_splits, + save_original_input=False, +): + """Run one layout and return all numerically observable forward/backward results.""" + FP8GlobalStateManager.reset() + num_gemms, out_features, in_features = weights.shape + module = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=torch.bfloat16, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + delay_wgrad_compute=delay_wgrad_compute, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + use_grouped_tensor=use_grouped_tensor, + save_original_input=save_original_input, + ) + + with torch.no_grad(): + if single_grouped_weight: + module.weight.rowwise_data.view_as(weights).copy_(weights) + else: + for i in range(num_gemms): + getattr(module, f"weight{i}").copy_(weights[i]) + if use_bias: + if single_grouped_bias: + module.bias.rowwise_data.view_as(biases).copy_(biases) + else: + for i in range(num_gemms): + getattr(module, f"bias{i}").copy_(biases[i]) + + flat_main_grad = None + initial_main_grad = None + weight_params = _grouped_linear_weight_params(module) + if fuse_wgrad_accumulation: + # MCore owns one flat FP32 grad buffer. Discrete parameters receive per-expert + # views, while a single grouped parameter receives one view over the full range. + flat_main_grad = torch.full( + (weights.numel(),), + 0.25, + dtype=torch.float32, + device="cuda", + ) + packed_main_grad = flat_main_grad.view_as(weights) + main_grad_views = ( + [packed_main_grad] + if single_grouped_weight + else [packed_main_grad[i] for i in range(num_gemms)] + ) + for param, main_grad in zip(weight_params, main_grad_views): + param.main_grad = main_grad + param.overwrite_main_grad = False + param.zero_out_wgrad = False + param.grad_added_to_main_grad = False + assert ( + param.main_grad.untyped_storage().data_ptr() + == flat_main_grad.untyped_storage().data_ptr() + ) + initial_main_grad = flat_main_grad.clone() + + x = x_base.detach().clone().requires_grad_(True) + m_splits_arg = ( + torch.tensor(m_splits, dtype=torch.int64, device="cuda") if use_grouped_tensor else m_splits + ) + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + y = module( + x, + m_splits_arg, + is_first_microbatch=False if fuse_wgrad_accumulation else None, + ) + y.backward(dy) + + if fuse_wgrad_accumulation and delay_wgrad_compute: + torch.testing.assert_close(flat_main_grad, initial_main_grad, rtol=0, atol=0) + + # The grouped-tensor path computes dbias during the main backward even when dW is delayed. + if use_bias and use_grouped_tensor: + assert all(param.grad is not None for param in _grouped_linear_bias_params(module)) + + if delay_wgrad_compute: + module.backward_dw() + + if fuse_wgrad_accumulation: + assert not torch.equal(flat_main_grad, initial_main_grad) + for param in weight_params: + assert param.grad_added_to_main_grad + assert ( + param.main_grad.untyped_storage().data_ptr() + == flat_main_grad.untyped_storage().data_ptr() + ) + packed_wgrad = flat_main_grad.view_as(weights) + elif single_grouped_weight: + packed_wgrad = module.weight.grad.view_as(weights) + else: + packed_wgrad = torch.stack([param.grad for param in weight_params]) + + packed_dbias = None + if use_bias: + bias_params = _grouped_linear_bias_params(module) + if single_grouped_bias: + packed_dbias = bias_params[0].grad.view_as(biases) + else: + packed_dbias = torch.stack([param.grad for param in bias_params]) + + return { + "output": y.detach().clone(), + "dgrad": x.grad.detach().clone(), + "wgrad": packed_wgrad.detach().clone(), + "dbias": None if packed_dbias is None else packed_dbias.detach().clone(), + } + + +_GROUPED_PARAMETER_LAYOUTS = [ + pytest.param(False, False, False, id="no-bias-discrete-weight"), + pytest.param(False, True, False, id="no-bias-single-weight"), + pytest.param(True, False, False, id="bias-discrete-weight-discrete-bias"), + pytest.param(True, True, False, id="bias-single-weight-discrete-bias"), + pytest.param(True, False, True, id="bias-discrete-weight-single-bias"), + pytest.param(True, True, True, id="bias-single-weight-single-bias"), +] + + +@pytest.mark.parametrize( + "use_bias,single_grouped_weight,single_grouped_bias", _GROUPED_PARAMETER_LAYOUTS +) +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + ], +) +@pytest.mark.parametrize("delay_wgrad_compute", _ALL_BOOLEAN) +@pytest.mark.parametrize("fuse_wgrad_accumulation", _ALL_BOOLEAN) +def test_grouped_parameter_layout_matches_cpu_m_splits( + monkeypatch, + use_bias, + single_grouped_weight, + single_grouped_bias, + fp8_recipe, + delay_wgrad_compute, + fuse_wgrad_accumulation, +): + """Match CUDA m_splits and all meaningful parameter layouts against the legacy path.""" + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + torch.bfloat16, + ): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + FP8GlobalStateManager.reset() + + torch.manual_seed(1234) + # MXFP8 and the grouped FP8 recipes require aligned expert problems. Use the same + # 256-aligned shapes for every recipe so all precision modes exercise one layout. + num_gemms = 2 + in_features = 256 + out_features = 256 + m_splits = [256, 512] + total_tokens = sum(m_splits) + x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(torch.bfloat16) + dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(torch.bfloat16) + weights = (0.1 * torch.randn(num_gemms, out_features, in_features, device="cuda")).to( + torch.bfloat16 + ) + biases = None + if use_bias: + biases = (0.1 * torch.randn(num_gemms, out_features, device="cuda")).to(torch.bfloat16) + + # The CPU m_splits baseline is explicitly the legacy, discrete-parameter contract. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") + reference = _run_grouped_parameter_layout( + use_grouped_tensor=False, + fp8_recipe=fp8_recipe, + single_grouped_weight=False, + single_grouped_bias=False, + use_bias=use_bias, + delay_wgrad_compute=delay_wgrad_compute, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + ) + + # Enable single parameters only for the CUDA m_splits target. The explicit layout flags + # below still decide whether this particular case uses discrete or grouped parameters. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + result = _run_grouped_parameter_layout( + use_grouped_tensor=True, + fp8_recipe=fp8_recipe, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + use_bias=use_bias, + delay_wgrad_compute=delay_wgrad_compute, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + ) + + tolerances = dict(rtol=1e-2, atol=5e-3) + + for name in ("output", "dgrad", "wgrad", "dbias"): + if reference[name] is None: + assert result[name] is None + else: + torch.testing.assert_close( + result[name].float(), + reference[name].float(), + **tolerances, + msg=f"Mismatch for {name}", + ) + + +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + pytest.param( + recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), + marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), + id="nvfp4", + ), + ], +) +@pytest.mark.parametrize("single_grouped_weight", _ALL_BOOLEAN) +def test_grouped_tensor_save_original_input_matches_saved_grouped_input( + monkeypatch, + fp8_recipe, + single_grouped_weight, +): + """Saving raw input must preserve native grouped forward, dgrad, and wgrad numerics.""" + if not is_module_grouped_tensor_path_supported(fp8_recipe, torch.bfloat16): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + if single_grouped_weight and fp8_recipe is not None and fp8_recipe.nvfp4(): + pytest.skip( + "NVFP4 grouped GEMM with single_grouped_weight is not supported yet; " + "only discrete weights are supported." + ) + + def reject_split_fallback(*_args, **_kwargs): + pytest.fail("save_original_input unexpectedly selected the split-quantize path") + + monkeypatch.setattr( + "transformer_engine.pytorch.module.grouped_linear._split_quantize", + reject_split_fallback, + ) + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + torch.manual_seed(1234) + num_gemms = 2 + in_features = 256 + out_features = 256 + m_splits = [256, 512] + total_tokens = sum(m_splits) + x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(torch.bfloat16) + dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(torch.bfloat16) + weights = (0.1 * torch.randn(num_gemms, out_features, in_features, device="cuda")).to( + torch.bfloat16 + ) + + saved_grouped = _run_grouped_parameter_layout( + use_grouped_tensor=True, + fp8_recipe=fp8_recipe, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=False, + use_bias=False, + delay_wgrad_compute=False, + fuse_wgrad_accumulation=False, + x_base=x_base, + dy=dy, + weights=weights, + biases=None, + m_splits=m_splits, + save_original_input=False, + ) + saved_original = _run_grouped_parameter_layout( + use_grouped_tensor=True, + fp8_recipe=fp8_recipe, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=False, + use_bias=False, + delay_wgrad_compute=False, + fuse_wgrad_accumulation=False, + x_base=x_base, + dy=dy, + weights=weights, + biases=None, + m_splits=m_splits, + save_original_input=True, + ) + + for name in ("output", "dgrad", "wgrad"): + torch.testing.assert_close( + saved_original[name].float(), + saved_grouped[name].float(), + rtol=1e-2, + atol=5e-3, + msg=f"Mismatch for {name}", + ) + + def _run_grouped_linear_path( *, enable_grouped_tensor_path: bool, @@ -1609,33 +2322,16 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( fp8_recipe, bias, fp8_model_params, delay_wgrad_compute, monkeypatch ): use_fp8 = fp8_recipe is not None - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." - ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor - # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. - is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() - if is_block_scaling and not (9, 0) <= device_capability < (10, 0): - pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") - if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): - pytest.skip( - "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500: - pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + dtype = torch.bfloat16 + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + dtype, + ): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") if fp8_model_params and not use_fp8: pytest.skip("fp8_model_params requires FP8") - dtype = torch.bfloat16 num_gemms = 3 in_features = 128 out_features = 128 @@ -1692,8 +2388,11 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monkeypatch): - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("GroupedTensor grouped GEMM path requires SM100+") + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") @@ -1722,6 +2421,102 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() +def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch): + """return_bias preserves the grouped parent and accumulates dbias into it. + + This mirrors how MCore applies a returned MoE bias:: + + x -> GroupedLinear (bias not applied) -> output + + + grouped bias [2, 128] + | + +-> repeat_interleave([256, 256]) + -> per-token bias [512, 128] + | + * routing probabilities + | + +-> loss.backward() -> grouped bias.grad + + Since the loss sums every output feature, each feature of expert ``i`` receives + ``sum(probs_for_expert_i)``. The identity assertion also ensures that TE returns the + registered grouped parent rather than copied or split bias tensors. + """ + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + dtype = torch.bfloat16 + num_gemms = 2 + in_features = 128 + out_features = 128 + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + total_tokens = 512 + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=True, + return_bias=True, + params_dtype=dtype, + device="cuda", + single_grouped_weight=False, + single_grouped_bias=True, + use_grouped_tensor=True, + ) + + x = torch.randn( + total_tokens, + in_features, + dtype=dtype, + device="cuda", + requires_grad=True, + ) + + probs = torch.cat( + ( + torch.full((256,), 0.25, dtype=dtype, device="cuda"), + torch.full((256,), 0.5, dtype=dtype, device="cuda"), + ) + ) + output, returned_bias = grouped_linear(x, m_splits) + + assert returned_bias is grouped_linear.bias + assert returned_bias.shape == (num_gemms, out_features) + + bias_per_token = torch.repeat_interleave( + returned_bias, + m_splits, + dim=0, + output_size=total_tokens, + ) + biased_output = (output + bias_per_token * probs.reshape(-1, 1)).to(output.dtype) + biased_output.sum().backward() + + # For token t and output feature j, the externally applied bias contributes + # bias[expert(t), j] * probs[t] to the summed loss. Therefore every bias feature for + # expert e receives sum(probs[t]) over that expert's token range. m_splits places the + # first 256 tokens on expert 0 and the remaining 256 tokens on expert 1. + expected_dbias = ( + torch.stack( + (probs[:256].float().sum(), probs[256:].float().sum()), + ) + .unsqueeze(-1) + .expand(num_gemms, out_features) + ) + assert grouped_linear.bias.grad is not None + # Megatron applies the packed bias in BF16. Compare its gradient against an independently + # accumulated FP32 reference with a tolerance appropriate for a 256-element BF16 reduction. + torch.testing.assert_close( + grouped_linear.bias.grad.float(), + expected_dbias, + rtol=5e-2, + atol=5e-3, + ) + + @pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"]) @pytest.mark.parametrize("supply", ["out", "dgrad_out", "both"]) def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatch): @@ -1730,18 +2525,11 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc Checks that a supplied buffer is written in place (bit-for-bit vs internal allocation) and, on the fused path with a padded input, that only the valid rows are touched. """ - if use_fused_path: - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell" - " (SM10x and SM110)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + if use_fused_path and not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if use_fused_path else "0") give_out = supply in ("out", "both") @@ -1819,19 +2607,16 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) -def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): - """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct. - - Graph-safe grouped quantization currently requires RHT, so requesting NVFP4 with - ``disable_rht=True`` while the fused grouped-tensor path is enabled falls back to the - legacy path internally. We verify the output and gradients against a reference built from - per-GEMM ``te.Linear`` modules that share the same weights and use the same NVFP4 recipe; - the grouped GEMM should match the loop of single GEMMs. +def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(): + """Non-RHT NVFP4 falls back to split-quantize for discrete parameters. + + Graph-safe grouped quantization currently requires RHT. Discrete parameters have a valid + split-quantize fallback, so enabling the grouped-tensor path is a preference rather than a + hard requirement for this parameter layout. """ if torch.cuda.get_device_capability() < (10, 0): pytest.skip("NVFP4 GroupedTensor grouped GEMM path requires SM100+") - monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") FP8GlobalStateManager.reset() dtype = torch.bfloat16 @@ -1854,7 +2639,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): disable_stochastic_rounding=True, ) - # Grouped path: fused path enabled, but non-RHT NVFP4 falls back to legacy internally. grouped_linear = GroupedLinear( num_gemms, in_features, @@ -1862,6 +2646,7 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): bias=False, params_dtype=dtype, device="cuda", + use_grouped_tensor=True, ) with torch.no_grad(): for i in range(num_gemms): @@ -1872,7 +2657,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): y = grouped_linear(x, m_splits) y.backward(dy) - # Reference: one te.Linear per GEMM sharing the same weights and NVFP4 recipe. ref_linears = torch.nn.ModuleList( [ Linear(in_features, out_features, bias=False, params_dtype=dtype, device="cuda") @@ -1890,7 +2674,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): ) y_ref.backward(dy) - # cuBLAS grouped GEMM should match the loop of single GEMMs bit-for-bit. tols = dict(rtol=0, atol=0) torch.testing.assert_close(y.float(), y_ref.float(), **tols) torch.testing.assert_close(x.grad.float(), x_ref.grad.float(), **tols) @@ -1902,6 +2685,33 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): ) +def test_grouped_linear_delay_wgrad_rejects_implicit_fallback(monkeypatch): + """Delayed wgrad reports when a grouped-tensor request used the legacy path.""" + monkeypatch.setattr( + "transformer_engine.pytorch.module.grouped_linear.is_module_grouped_tensor_path_supported", + lambda *_args, **_kwargs: False, + ) + grouped_linear = GroupedLinear( + 2, + 64, + 64, + bias=True, + params_dtype=torch.bfloat16, + device="cuda", + delay_wgrad_compute=True, + use_grouped_tensor=True, + ) + x = torch.randn(16, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + m_splits = torch.tensor([8, 8], dtype=torch.int64, device="cuda") + + grouped_linear(x, m_splits).sum().backward() + with pytest.raises( + RuntimeError, + match="implicit fallback is unsupported with delay_wgrad_compute=True", + ): + grouped_linear.backward_dw() + + @pytest.mark.parametrize( "fp8_recipe", [ @@ -1931,33 +2741,16 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch): """Fused GroupedTensor GEMM path should be CUDA graph capturable.""" use_fp8 = fp8_recipe is not None - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." - ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor - # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. - is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() - if is_block_scaling and not (9, 0) <= device_capability < (10, 0): - pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") - if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): - pytest.skip( - "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500: - pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + dtype = torch.bfloat16 + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + dtype, + ): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") FP8GlobalStateManager.reset() - dtype = torch.bfloat16 device = "cuda" num_gemms = 3 in_features = 128 diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d195eb2f78..c6d73cba63 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -15,6 +15,7 @@ import torch import transformer_engine.pytorch as te +from transformer_engine.pytorch.constants import TE_DType import transformer_engine.pytorch.ops.fused.grouped_mlp as grouped_mlp_module from transformer_engine.pytorch.ops.fused.grouped_mlp import ( _cudnn_frontend_supports_grouped_gemm_srelu, @@ -23,6 +24,7 @@ from transformer_engine.pytorch.ops.basic.grouped_linear import ( OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, + is_op_fuser_grouped_tensor_path_supported, ) from transformer_engine.pytorch import ( QuantizedTensor, @@ -230,6 +232,23 @@ def make_reference_and_test_tensors( return ref, test +class _InjectGrad(torch.autograd.Function): + """Replace the gradient flowing into ``x`` with ``grad``. + + Mirrors how an FP8 token dispatch delivers a pre-quantized ``GroupedTensor`` + grad output: the downstream op simply returns one as its grad input. + """ + + @staticmethod + def forward(ctx, x, grad): # pylint: disable=arguments-differ + ctx.injected_grad = grad + return x + + @staticmethod + def backward(ctx, grad_output): # pylint: disable=arguments-differ + return ctx.injected_grad, None + + class TestGroupedLinearOp: """Tests for advanced features with grouped linear basic op""" @@ -260,6 +279,26 @@ def test_meta_single_grouped_weight_with_delayed_wgrad(self, monkeypatch) -> Non assert grouped_weight.skip_backward_post_hook + def test_single_grouped_bias_uses_registered_packed_storage(self, monkeypatch) -> None: + """The grouped bias compute view must alias the registered trainable parent.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + op = te.ops.GroupedLinear( + 2, + 128, + 128, + bias=True, + device="cuda", + dtype=torch.bfloat16, + single_grouped_bias=True, + ) + + bias_packed = op._get_packed_bias_tensor(torch.bfloat16) + + assert bias_packed.shape == (op.num_groups, op.out_features) + assert bias_packed.untyped_storage().data_ptr() == op.bias.rowwise_data.data_ptr() + assert op.bias.requires_grad + assert dict(op.named_parameters())["bias"] is op.bias + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) @@ -320,6 +359,16 @@ def test_grouped_linear( if single_grouped_bias and not bias: pytest.skip("single_grouped_bias requires bias=True") + recipe = make_recipe(quantization) + compute_recipe = recipe if quantized_compute else None + if ( + single_grouped_weight or single_grouped_bias + ) and not is_op_fuser_grouped_tensor_path_supported( + compute_recipe, + dtype, + ): + # Single grouped parameters intentionally have no split-quantize fallback. + pytest.skip("Single grouped parameters require the native grouped-tensor path") if single_grouped_weight and quantized_weight and quantization in ("fp8_delayed_scaling"): pytest.skip( "single_grouped_weight does not support FP8 delayed scaling " @@ -375,7 +424,6 @@ def test_grouped_linear( y_ref.backward(dy_ref) # Construct fusible operation - recipe = make_recipe(quantization) with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): op = te.ops.GroupedLinear( group_size, @@ -464,6 +512,189 @@ def test_grouped_linear( else: assert b_test.grad is None + @staticmethod + def _make_rowwise_mxfp8_wire_input( + x_hp: torch.Tensor, + group_size: int, + split_sizes: torch.Tensor, + ) -> "GroupedTensor": + """Rowwise-only MXFP8 GroupedTensor with compact scales (FP8 dispatch wire format).""" + wire_quantizer = MXFP8Quantizer( + fp8_dtype=TE_DType[torch.float8_e4m3fn], rowwise=True, columnwise=False + ) + x_wire = tex.group_quantize( + x_hp, wire_quantizer, group_size, split_sizes.to(dtype=torch.int64) + ) + # Sanity: the wire tensor is rowwise-only with unswizzled scales. + assert x_wire.columnwise_data is None + assert not x_wire._with_gemm_swizzled_scales + return x_wire + + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_linear_prequantized_mxfp8_input( + self, + *, + group_size: int = 4, + weight_shape: tuple[int, int] = (256, 256), + split_alignment: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + weight_requires_grad: bool, + ) -> None: + """Rowwise-only MXFP8 GroupedTensor input (FP8 token dispatch wire format). + + The input arrives already rowwise-quantized with compact scales. The op + must feed the rowwise data to the forward GEMM as-is and manufacture the + columnwise copy for the wgrad GEMM. The reference run consumes the + *dequantized* wire tensor (the only data a layer can see after FP8 + dispatch) on the normal quantize-from-BF16 path; because MXFP8 requant + is idempotent along the rowwise axis and both paths derive the + columnwise copy from the same dequantized data, the two runs must match + bit-for-bit. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype) + + # Split sizes (including an empty group) + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + out_features, in_features = weight_shape + total_tokens = int(split_sizes.sum().item()) + in_shape = (total_tokens, in_features) + + # Wire-format input and its exact dequantization (the reference input). + x_hp = torch.rand(in_shape, dtype=dtype, device=device) - 0.5 + x_wire = self._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes) + x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view(in_shape) + + dy = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5 + recipe = make_recipe("mxfp8") + op = te.ops.GroupedLinear( + group_size, in_features, out_features, bias=False, device=device, dtype=dtype + ) + with torch.no_grad(): + for param in op.parameters(): + param.requires_grad_(requires_grad=weight_requires_grad) + + def _run(x): + with te.autocast(enabled=True, recipe=recipe): + y = op(x, split_sizes) + wgrads = [] + if weight_requires_grad: + y.backward(dy) + for group_idx in range(group_size): + weight = getattr(op, f"weight{group_idx}") + wgrads.append(weight.grad.detach().clone()) + weight.grad = None + return y.detach(), wgrads + + y_ref, wgrads_ref = _run(x_ref) + y_test, wgrads_test = _run(x_wire) + + # Bit-exact match expected (identical quantized inputs and kernels). + torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(wgrads_test, wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + + @pytest.mark.parametrize("bias_mode", ("none", "plain", "scaled")) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_linear_prequantized_mxfp8_grad( + self, + *, + group_size: int = 4, + weight_shape: tuple[int, int] = (256, 256), + split_alignment: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + bias_mode: str, + weight_requires_grad: bool, + ) -> None: + """Rowwise-only MXFP8 GroupedTensor grad output (FP8 token dispatch, backward). + + Mirrors ``test_grouped_linear_prequantized_mxfp8_input`` on the backward + side: the rowwise data feeds the dgrad GEMM and the columnwise copy is + manufactured for wgrad. Covers all three bias gradient sources: none, + ``plain`` (fused into the columnwise stage of the quantize kernel, or + reduced from the dequantized grad when frozen weights leave no columnwise + stage), and ``scaled`` (``scale_bias``, whose dbias/dscales need the + dequantized grad because they depend on the routing probabilities). + + With frozen weights TE also requires frozen biases, so bias gradients are + not observable there; ``dscales`` still is, since the probabilities are an + input rather than a parameter. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype) + + has_bias = bias_mode != "none" + use_scale_bias = bias_mode == "scaled" + + # Split sizes (including an empty group) + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + out_features, in_features = weight_shape + total_tokens = int(split_sizes.sum().item()) + + x = torch.rand((total_tokens, in_features), dtype=dtype, device=device) - 0.5 + dy_hp = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + + # Wire-format grad output and its exact dequantization (the reference grad). + dy_wire = self._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes) + dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, out_features + ) + + recipe = make_recipe("mxfp8") + op = te.ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=has_bias, + device=device, + dtype=dtype, + scale_bias=use_scale_bias, + ) + if not weight_requires_grad: + # TE requires bias.requires_grad to match weight.requires_grad. + for param in op.parameters(): + param.requires_grad_(False) + + def _run(grad): + x_in = x.detach().clone().requires_grad_() + probs_in = probs.detach().clone().requires_grad_(use_scale_bias) + extra_inputs = (split_sizes, probs_in) if use_scale_bias else (split_sizes,) + with te.autocast(enabled=True, recipe=recipe): + y = op(x_in, *extra_inputs) + # Deliver ``grad`` as the op's grad output, as FP8 dispatch would. + _InjectGrad.apply(y, grad).backward(torch.ones_like(y)) + grads = [("dx", x_in.grad)] + if use_scale_bias: + grads.append(("dprobs", probs_in.grad)) + if weight_requires_grad: + for group_idx in range(group_size): + weight = getattr(op, f"weight{group_idx}") + grads.append((f"w{group_idx}", weight.grad.detach().clone())) + weight.grad = None + if has_bias: + bias_param = getattr(op, f"bias{group_idx}") + grads.append((f"b{group_idx}", bias_param.grad.detach().clone())) + bias_param.grad = None + return grads + + grads_ref = _run(dy_ref) + grads_test = _run(dy_wire) + + # Bit-exact match expected (identical quantized grads and kernels). + for (_, grad_test), (_, grad_ref) in zip(grads_test, grads_ref): + torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0) + @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) @pytest.mark.parametrize( "quantization", @@ -506,22 +737,6 @@ def test_grouped_linear_cuda_graph_safe( "single_grouped_weight/single_grouped_bias requires" " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" ) - device_capability = torch.cuda.get_device_capability() - if device_capability < (9, 0): - pytest.skip( - "Grouped GEMM CUDA-graph-safe path requires Hopper (SM90) or Blackwell (SM100+)" - ) - # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path, - # but MXFP8/NVFP4 grouped quantization kernels require Blackwell (SM100+). - requires_blackwell = quantization is not None and quantization != "fp8_current_scaling" - if requires_blackwell and device_capability < (10, 0): - pytest.skip("MXFP8/NVFP4 grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") - # Grouped GEMM on Hopper requires cuBLAS 13.4+; Blackwell requires cuBLAS 13.3+. - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") if quantization is None and quantized_weight: pytest.skip("quantized_weight requires a quantization recipe") if ( @@ -539,6 +754,13 @@ def test_grouped_linear_cuda_graph_safe( "only discrete weights (single_grouped_weight=False) are supported." ) + recipe = make_recipe(quantization) + if not is_op_fuser_grouped_tensor_path_supported( + recipe, + dtype, + ): + pytest.skip("Configuration falls back to the non-CUDA-graph-safe path") + single_grouped_bias = bias and single_grouped_weight # Split sizes (statically pinned for graph capture) @@ -551,7 +773,6 @@ def test_grouped_linear_cuda_graph_safe( in_shape = (num_active_tokens + token_padding, in_features) out_shape = (in_shape[0], out_features) - recipe = make_recipe(quantization) with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): op = te.ops.GroupedLinear( group_size, @@ -755,6 +976,13 @@ def test_grouped_mlp( maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) if dtype == torch.bfloat16 and not is_bf16_available(): pytest.skip("BF16 requires SM 8.0+") + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and ( + single_grouped_weight or single_grouped_bias + ): + pytest.skip( + "single_grouped_weight/single_grouped_bias requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) if single_grouped_weight and quantization != "mxfp8": pytest.skip("single_grouped_weight is only supported for MXFP8 quantization") if single_grouped_bias and not bias: @@ -1026,8 +1254,10 @@ def _make_module(): or ( quantization == "nvfp4_rht" and dtype == torch.bfloat16 - and activation == "scaled_srelu" - and glu_interleave_size is None + and ( + (not activation_is_glu and glu_interleave_size is None) + or (activation_is_glu and glu_interleave_size == 32) + ) ) ) if expected_grouped_mlp_fusion: @@ -1107,6 +1337,208 @@ def _make_module(): assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols) assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_mlp_prequantized_mxfp8_input( + self, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + weight_requires_grad: bool, + ) -> None: + """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor input. + + Production (FP8 token dispatch) path: FC1 receives an already + rowwise-quantized input with compact scales. The fused op must feed the + rowwise data to the forward GEMM and manufacture FC1's columnwise copy + for wgrad. Compared bit-for-bit against a run on the dequantized wire + input (see ``test_grouped_linear_prequantized_mxfp8_input``). + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + maybe_skip_quantization( + "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype + ) + + # Split sizes (including an empty group); sum is a multiple of 128. + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + glu_interleave_size = 32 + + # Wire-format FC1 input and its exact dequantization (reference input). + x_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes) + x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, hidden_size + ) + + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype + ) + fc2 = te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + with torch.no_grad(): + for param in module.parameters(): + param.requires_grad_(requires_grad=weight_requires_grad) + + def _run(x): + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + fc1_wgrads, fc2_wgrads = [], [] + if weight_requires_grad: + y.backward(dy) + for group_idx in range(group_size): + fc1_w = getattr(fc1, f"weight{group_idx}") + fc2_w = getattr(fc2, f"weight{group_idx}") + fc1_wgrads.append(fc1_w.grad.detach().clone()) + fc2_wgrads.append(fc2_w.grad.detach().clone()) + fc1_w.grad = None + fc2_w.grad = None + return y.detach(), fc1_wgrads, fc2_wgrads + + y_ref, fc1_wgrads_ref, fc2_wgrads_ref = _run(x_ref) + y_test, fc1_wgrads_test, fc2_wgrads_test = _run(x_wire) + + # Confirm the CuTeDSL fused op was actually formed (not the fallback). + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Bit-exact match expected (identical quantized inputs and kernels). + torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(fc1_wgrads_test, fc1_wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(fc2_wgrads_test, fc2_wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_mlp_prequantized_mxfp8_grad( + self, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + bias: bool, + weight_requires_grad: bool, + ) -> None: + """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor grad output. + + FC2 receives the pre-quantized grad, as an FP8 token dispatch delivers it + on the backward pass. With ``bias`` FC2 uses ``scale_bias``, whose + dbias/dscales need the dequantized grad rather than the fused dbias. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + if not weight_requires_grad: + # Independent of pre-quantization: the fused forward saves the input + # activations whenever anything requires grad, then asserts they carry + # columnwise data -- which is only built when the weights need grads. + pytest.skip("Fused grouped MLP does not support frozen weights") + maybe_skip_quantization( + "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype + ) + + # Split sizes (including an empty group); sum is a multiple of 128. + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + glu_interleave_size = 32 + + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + # Wire-format grad output and its exact dequantization (the reference grad). + dy_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes) + dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, hidden_size + ) + + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=bias, device=device, dtype=dtype + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=bias, + device=device, + dtype=dtype, + scale_bias=bias, + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + + # Frozen experts (weights) with a still-training bias/router is the case + # where the dequantized grad is needed but is not a byproduct of wgrad. + if not weight_requires_grad: + for fc in (fc1, fc2): + for group_idx in range(group_size): + getattr(fc, f"weight{group_idx}").requires_grad_(False) + + def _run(grad): + x_in = x.detach().clone().requires_grad_() + fc2_extra = (split_sizes, probs) if bias else (split_sizes,) + with te.autocast(enabled=True, recipe=recipe): + y = module(x_in, split_sizes, probs, *fc2_extra) + _InjectGrad.apply(y, grad).backward(torch.ones_like(y)) + grads = [("dx", x_in.grad)] + for name, fc in (("fc1", fc1), ("fc2", fc2)): + for group_idx in range(group_size): + if weight_requires_grad: + weight = getattr(fc, f"weight{group_idx}") + grads.append((f"{name}_w{group_idx}", weight.grad.detach().clone())) + weight.grad = None + if bias: + bias_param = getattr(fc, f"bias{group_idx}") + grads.append((f"{name}_b{group_idx}", bias_param.grad.detach().clone())) + bias_param.grad = None + return grads + + grads_ref = _run(dy_ref) + grads_test = _run(dy_wire) + + # Confirm the CuTeDSL fused op was actually formed (not the fallback). + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Bit-exact match expected (identical quantized grads and kernels), except + # bias gradients: the fused kernels generate them with an accumulation + # that is not reproducible run to run (two runs on identical inputs differ + # by one BF16 ulp). Same tolerances as + # ``test_grouped_mlp_single_weight_numerics``. + bias_tols = {"rtol": 0.05, "atol": 0.015625} + for (name, grad_test), (_, grad_ref) in zip(grads_test, grads_ref): + if "_b" in name: + torch.testing.assert_close(grad_test, grad_ref, **bias_tols) + else: + torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0) + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize( @@ -1147,13 +1579,15 @@ def test_grouped_mlp_single_group_mxfp8( """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" if ( runtime_offsets_supported - and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets() + and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets( + te.ops.ScaledSwiGLU + ) ): pytest.skip("Requires cuDNN frontend >= 1.27.0") monkeypatch.setattr( grouped_mlp_module, "_cudnn_frontend_supports_single_group_runtime_offsets", - lambda: runtime_offsets_supported, + lambda _activation_type: runtime_offsets_supported, ) self.test_grouped_mlp( group_size=1, @@ -1266,6 +1700,8 @@ def test_grouped_mlp_single_weight_numerics( ) -> None: """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP.""" + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0": + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") @@ -1584,6 +2020,8 @@ def test_grouped_mlp_overwrite_main_grad( that read ``.grad`` don't see stale bytes from the cached dummy). """ + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight: + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") @@ -1715,6 +2153,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( ) -> None: """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8).""" + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight: + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") if dtype not in (torch.bfloat16, torch.float16): diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 4dd52ee2bd..ac8c493290 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -577,7 +577,7 @@ def test_quantize_grouped_mxfp8(self, shape_case: str, output_dbias: bool) -> No if output_dbias: expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) - assert torch.allclose(dbias, expected_dbias) + torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3) @pytest.mark.parametrize("output_dbias", [False, True]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @@ -805,6 +805,152 @@ def _run(inp): assert torch.equal(static_output.rowwise_data, expected.rowwise_data) assert torch.equal(static_output.scale_inv, expected.scale_inv) + @pytest.mark.parametrize( + "quantization", + [ + pytest.param( + "fp8_current_scaling", + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + ), + pytest.param( + "fp8_blockwise", + marks=pytest.mark.skipif( + not fp8_block_scaling_grouped_available, + reason=reason_for_no_fp8_block_scaling_grouped, + ), + ), + ], + ) + def test_group_quantize_reuses_destination_and_honors_noop_fp8(self, quantization: str) -> None: + """Check pointer-stable weight-cache updates and no-op CUDA graph replays.""" + num_tensors = 2 + shape = (num_tensors * 256, 256) + if quantization == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + force_pow_2_scales=False, + amax_epsilon=0.0, + ) + quantizer.set_usage(rowwise=True, columnwise=True) + else: + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + output = tex.group_quantize(source, quantizer, num_tensors, None) + output_buffers = tuple( + buffer + for buffer in ( + output.rowwise_data, + output.columnwise_data, + output.scale_inv, + output.columnwise_scale_inv, + output.amax, + output.columnwise_amax, + output.scale, + ) + if buffer is not None + ) + assert output_buffers + pointers = tuple(buffer.data_ptr() for buffer in output_buffers) + + # Re-quantization must update the existing allocations because a captured GEMM keeps + # their raw addresses rather than looking them up again through the Python object. + updated_source = source + 1 + updated = tex.group_quantize( + updated_source, + quantizer, + num_tensors, + None, + output=output, + ) + reference = tex.group_quantize(updated_source, quantizer, num_tensors, None) + assert updated is output + assert pointers == tuple(buffer.data_ptr() for buffer in output_buffers) + reference_buffers = tuple( + buffer + for buffer in ( + reference.rowwise_data, + reference.columnwise_data, + reference.scale_inv, + reference.columnwise_scale_inv, + reference.amax, + reference.columnwise_amax, + reference.scale, + ) + if buffer is not None + ) + assert len(output_buffers) == len(reference_buffers) + for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True): + assert torch.equal(output_buffer, reference_buffer) + + # A nonzero device flag must preserve every cached payload and metadata buffer. + old_buffers = tuple(buffer.clone() for buffer in output_buffers) + noop = torch.ones(1, dtype=torch.float32, device="cuda") + tex.group_quantize( + updated_source * 2, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True): + assert torch.equal(output_buffer, old_buffer) + + # Capture one pointer-stable update. Replays change only source contents and the device + # no-op value, matching TE's CUDA graph microbatch weight-cache protocol. + graph_source = updated_source.clone() + noop.zero_() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + tex.group_quantize( + graph_source, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + + graph_source.copy_(source - 1) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + reference = tex.group_quantize(graph_source, quantizer, num_tensors, None) + reference_buffers = tuple( + buffer + for buffer in ( + reference.rowwise_data, + reference.columnwise_data, + reference.scale_inv, + reference.columnwise_scale_inv, + reference.amax, + reference.columnwise_amax, + reference.scale, + ) + if buffer is not None + ) + assert len(output_buffers) == len(reference_buffers) + for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True): + assert torch.equal(output_buffer, reference_buffer) + + old_buffers = tuple(buffer.clone() for buffer in output_buffers) + graph_source.copy_(source + 2) + noop.fill_(1) + graph.replay() + torch.cuda.synchronize() + for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True): + assert torch.equal(output_buffer, old_buffer) + @pytest.mark.parametrize("mode", ["rowwise", "columnwise", "both"]) @pytest.mark.parametrize( "shape_case", @@ -1121,7 +1267,7 @@ def test_quantize_grouped_fp8_blockwise( if output_dbias: expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) - assert torch.allclose(dbias, expected_dbias) + torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3) @pytest.mark.parametrize( "shape", @@ -1204,6 +1350,147 @@ def test_group_dequantize_cudagraph_capturable(self) -> None: for exp, got in zip(expected_tensors, static_tensors): assert torch.equal(got, exp) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_quantize_reuses_destination_and_honors_noop(self) -> None: + """Verify the in-place and graph-safe contracts of grouped MXFP8 quantization. + + A captured training graph records the addresses of all MXFP8 weight storage, not just + the Python ``GroupedTensor`` object. Re-quantizing an updated BF16 weight must therefore + refresh the existing FP8 payloads and scales without replacing any allocation. During + graph replay, ``noop_flag`` must additionally allow the captured quantization kernel to + execute as a no-op on microbatches that should reuse the cached weight. + """ + num_tensors = 2 + # Treat two contiguous [256, 256] matrices as one grouped BF16 source. These dimensions + # satisfy the MXFP8 block-alignment requirements in both rowwise and columnwise layouts. + shape = (num_tensors * 256, 256) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + # Forward GEMM consumes rowwise weight storage, while backward dgrad consumes columnwise + # storage. Both representations and both scale tensors must be updated and kept stable. + quantizer.set_usage(rowwise=True, columnwise=True) + # Store the quantized output in GEMM-ready (swizzled) scale layout. This makes the test + # cover the exact destination format used by cached MXFP8 model weights. + quantizer.optimize_for_gemm = True + + # The first call owns allocation: it creates the persistent destination that subsequent + # optimizer updates and CUDA graph replays must modify in place. + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + output = tex.group_quantize(source, quantizer, num_tensors, None) + + # CUDA graphs retain these raw addresses. Object identity alone is insufficient because + # replacing one internal payload or scale allocation would leave the graph with a stale + # pointer even if ``output`` remained the same Python object. + pointers = ( + output.rowwise_data.data_ptr(), + output.columnwise_data.data_ptr(), + output.scale_inv.data_ptr(), + output.columnwise_scale_inv.data_ptr(), + ) + + # Eager in-place update: quantize new BF16 values into the previously allocated output. + # An independently allocated quantization provides the numerical reference. + updated_source = source + 1 + updated = tex.group_quantize( + updated_source, + quantizer, + num_tensors, + None, + output=output, + ) + reference = tex.group_quantize(updated_source, quantizer, num_tensors, None) + + # The API must return the caller-provided destination and preserve every captured address. + assert updated is output + assert pointers == ( + output.rowwise_data.data_ptr(), + output.columnwise_data.data_ptr(), + output.scale_inv.data_ptr(), + output.columnwise_scale_inv.data_ptr(), + ) + + # In-place quantization must still produce exactly the same FP8 bytes and scales as a + # normal allocating call, for both GEMM orientations. + assert torch.equal(output.rowwise_data, reference.rowwise_data) + assert torch.equal(output.columnwise_data, reference.columnwise_data) + assert torch.equal(output.scale_inv, reference.scale_inv) + assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv) + + # Eager no-op: a nonzero device flag tells the kernel to preserve the cached destination. + # Clone all four buffers so this checks payloads and metadata independently. + old_buffers = ( + output.rowwise_data.clone(), + output.columnwise_data.clone(), + output.scale_inv.clone(), + output.columnwise_scale_inv.clone(), + ) + noop = torch.ones(1, dtype=torch.float32, device="cuda") + + # The source is deliberately different. If the kernel ignores ``noop_flag``, at least + # one payload or scale tensor below will change and expose the failure. + tex.group_quantize( + updated_source * 2, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + assert torch.equal(output.rowwise_data, old_buffers[0]) + assert torch.equal(output.columnwise_data, old_buffers[1]) + assert torch.equal(output.scale_inv, old_buffers[2]) + assert torch.equal(output.columnwise_scale_inv, old_buffers[3]) + + # Capture one fixed launch. ``graph_source``, ``noop``, and ``output`` keep stable device + # addresses; replay changes only their contents. This is how weight caching works when + # different microbatches replay the same CUDA graph. + graph_source = updated_source.clone() + # zero means "perform quantization"; nonzero means "leave destination untouched". + noop.zero_() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + tex.group_quantize( + graph_source, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + + # Replay with new source contents and noop=0. The captured kernel must refresh the same + # destination allocations, and the result must equal a fresh eager quantization. + graph_source.copy_(source - 1) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + reference = tex.group_quantize(graph_source, quantizer, num_tensors, None) + assert torch.equal(output.rowwise_data, reference.rowwise_data) + assert torch.equal(output.columnwise_data, reference.columnwise_data) + assert torch.equal(output.scale_inv, reference.scale_inv) + assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv) + + # Snapshot the successfully refreshed cache before testing the captured no-op branch. + old_buffers = ( + output.rowwise_data.clone(), + output.columnwise_data.clone(), + output.scale_inv.clone(), + output.columnwise_scale_inv.clone(), + ) + + # Replay the exact same captured graph with different source data but noop=1. The launch + # still occurs, which is required for CUDA graph structural consistency, but the kernel + # must not mutate any cached FP8 payload or scale buffer. + graph_source.copy_(source + 2) + noop.fill_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(output.rowwise_data, old_buffers[0]) + assert torch.equal(output.columnwise_data, old_buffers[1]) + assert torch.equal(output.scale_inv, old_buffers[2]) + assert torch.equal(output.columnwise_scale_inv, old_buffers[3]) + @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"]) @pytest.mark.parametrize("direction", ["rowwise", "columnwise"]) @pytest.mark.skipif( diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 8249c7fedd..e6a83d92bc 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -42,6 +42,10 @@ is_nvfp4_available, ) from transformer_engine.pytorch import checkpoint as te_checkpoint +from transformer_engine.pytorch.distributed import ( + is_fp8_activation_recompute_enabled, + in_fp8_activation_recompute_phase, +) from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.common import recipe from transformer_engine.pytorch import DType @@ -82,6 +86,9 @@ all_boolean = [True, False] +# fp8_meta key written by FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute +_FP8_RECOMPUTE_KEY = "global_fp8_buffer_pos_fwd_recompute" + all_activations = [ "gelu", "geglu", @@ -648,7 +655,15 @@ def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, recipe, fp8_m def _test_e2e_full_recompute( - bs, dtype, config, fp8, recipe, fp8_model_params=False, recompute=False, use_reentrant=True + bs, + dtype, + config, + fp8, + recipe, + fp8_model_params=False, + recompute=False, + use_reentrant=True, + inner_autocast=False, ): reset_rng_states() FP8GlobalStateManager.reset() @@ -685,10 +700,17 @@ def _test_e2e_full_recompute( te_inp_hidden_states.retain_grad() te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q) - with autocast(enabled=fp8, recipe=recipe): + forward = block + if inner_autocast: + + def forward(*args, **kwargs): + with autocast(enabled=fp8, recipe=recipe): + return block(*args, **kwargs) + + with autocast(enabled=fp8 and not inner_autocast, recipe=recipe): if recompute: te_out = te_checkpoint( - block, + forward, te_inp_hidden_states, attention_mask=te_inp_attn_mask, checkpoint_core_attention=False, @@ -697,7 +719,7 @@ def _test_e2e_full_recompute( use_reentrant=use_reentrant, ) else: - te_out = block( + te_out = forward( te_inp_hidden_states, attention_mask=te_inp_attn_mask, checkpoint_core_attention=False, @@ -787,6 +809,152 @@ def test_gpt_full_activation_recompute( ) +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkeypatch): + """Check recompute numerics when FP8 autocast starts inside the checkpointed callable.""" + if not use_reentrant: + # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion. + monkeypatch.setenv("NVTE_BIAS_GELU_NVFUSION", "0") + + dtype = torch.bfloat16 + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + config = model_configs["126m"] + + # Reference also opens the autocast inside the callable, so the only difference + # between the two runs is activation recompute. + outputs, names = _test_e2e_full_recompute( + 1, + dtype, + config, + True, + fp8_recipe, + recompute=False, + use_reentrant=use_reentrant, + inner_autocast=True, + ) + + # Before the fix, phase 1 skipped the stash while the recompute still restored it, which + # surfaced as a KeyError on the recompute buffer lookup. Count both sides to pin that down. + stash_counts, restore_counts = {}, {} + stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute + restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute + + def record_stash(fp8_meta): + stash_fn(fp8_meta) + if _FP8_RECOMPUTE_KEY in fp8_meta: + stash_counts[id(fp8_meta)] = stash_counts.get(id(fp8_meta), 0) + 1 + + def record_restore(fp8_meta): + # The restore site is not gated on delayed scaling, but only delayed scaling stashes. + if not fp8_meta["recipe"].delayed(): + restore_fn(fp8_meta) + return + key = id(fp8_meta) + assert key in stash_counts, "Recompute restored a scale that was never stashed" + restore_counts[key] = restore_counts.get(key, 0) + 1 + restore_fn(fp8_meta) + + monkeypatch.setattr( + FP8GlobalStateManager, + "copy_forward_fp8_meta_tensors_for_recompute", + staticmethod(record_stash), + ) + monkeypatch.setattr( + FP8GlobalStateManager, + "get_old_fp8_meta_tensors_for_recompute", + staticmethod(record_restore), + ) + + outputs_recompute, _ = _test_e2e_full_recompute( + 1, + dtype, + config, + True, + fp8_recipe, + recompute=True, + use_reentrant=use_reentrant, + inner_autocast=True, + ) + + assert stash_counts, "No FP8 module stashed a forward scale for the recompute phase" + assert restore_counts == stash_counts, "Stash and restore of forward scales are unbalanced" + + for name, ref, test in zip(names, outputs, outputs_recompute): + torch.testing.assert_close( + test, + ref, + msg=f"Mismatch in tensor {name}", + rtol=0.125, + atol=0.0675, + ) + + +def _checkpointed_linear_backward(body, use_reentrant, *layers): + """Run a checkpointed callable end to end and check the gradients are finite.""" + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = te_checkpoint(body, inp, use_reentrant=use_reentrant) + loss = out.float().sum() + loss.backward() + torch.cuda.synchronize() + + assert torch.isfinite(loss) + assert inp.grad is not None and torch.isfinite(inp.grad).all() + for layer in layers: + assert layer.weight.grad is not None + assert torch.isfinite(layer.weight.grad).all() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_checkpoint_inner_autocast_is_an_fp8_recompute_region(use_reentrant): + """An FP8 autocast opened inside a checkpointed callable is an FP8 recompute region.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + + observed = [] + + def body(value): + outside = is_fp8_activation_recompute_enabled() + with autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + outside, + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + ) + ) + return layer(value) + + _checkpointed_linear_backward(body, use_reentrant, layer) + + # One entry for the checkpointed forward, one for the recompute during backward. The + # query is only an FP8 recompute region inside the autocast, in both phases. + assert observed == [(False, True, False), (False, True, True)] + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state(use_reentrant): + """Only the inner FP8 region of a mixed checkpoint saves recompute metadata.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + non_fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + + def body(value): + value = non_fp8_layer(value) + with autocast(enabled=True, recipe=fp8_recipe): + return fp8_layer(value) + + _checkpointed_linear_backward(body, use_reentrant, non_fp8_layer, fp8_layer) + + assert _FP8_RECOMPUTE_KEY not in non_fp8_layer.fp8_meta + assert _FP8_RECOMPUTE_KEY in fp8_layer.fp8_meta + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index de62e56e53..c9b620fa1e 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -37,6 +37,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data +from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported from utils import ModelConfig, recipe_id, skip_unsupported_backward_override # Only run FP8 tests on supported devices. @@ -86,11 +87,11 @@ def is_fp8_supported(config: ModelConfig): def nvfp4_vanilla(): - nvfp4_recipe = recipe.NVFP4BlockScaling() - nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() - nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() - nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() - return nvfp4_recipe + return recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ) def nvfp4_row_scaled(): @@ -123,7 +124,7 @@ def nvfp4_4over6(): if mxfp8_available: fp8_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: - fp8_recipes.append(nvfp4_vanilla()) # TODO: fix check for this + fp8_recipes.append(nvfp4_vanilla()) fp8_recipes.append(nvfp4_4over6()) if fp8_block_scaling_available: fp8_recipes.append(recipe.Float8BlockScaling()) @@ -563,6 +564,7 @@ def test_sanity_linear_with_zero_tokens( out = te_linear(inp_hidden_states) loss = out.sum() loss.backward() + torch.cuda.synchronize() assert out.shape == (num_tokens, ffn_hidden_size) @@ -603,17 +605,34 @@ def test_sanity_grouped_linear( if fp8_recipe is not None: fp8_recipe = copy.deepcopy(fp8_recipe) fp8_recipe.backward_override = backward_override + if single_param and not is_module_grouped_tensor_path_supported( + fp8_recipe, + dtype, + ): + pytest.skip("Single grouped parameters require the native grouped-tensor path") + if single_param: + # Single grouped parameters intentionally have no split-quantize fallback, so this + # test must satisfy the native grouped kernels' shape contract. MCore pads each + # expert's token count to 256; TE requires at least 128-row alignment. Weight K must + # be 64-aligned. + tokens_per_nonempty_expert = bs * config.max_seqlen_q + if tokens_per_nonempty_expert % 128 != 0: + pytest.skip("Single grouped parameters require each nonempty m_split to be 128-aligned") + k_alignment = 64 + if config.hidden_size % k_alignment != 0: + pytest.skip(f"Single grouped parameters require GEMM K to be {k_alignment}-aligned") if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") if fp8_recipe.nvfp4(): - if not getattr(fp8_recipe, "row_scaled_activation", False): - pytest.skip("NVFP4 not supported for grouped linear") - if single_param: - pytest.skip("Row-scaled NVFP4 does not support GroupedTensor grouped linear") - if dtype == torch.float16: - pytest.skip("FP16 output for NVFP4 not supported") + if dtype != torch.bfloat16: + pytest.skip("NVFP4 GroupedLinear requires BF16") + if single_param and not fp8_model_params: + pytest.skip( + "NVFP4 single grouped BF16 primary weights require unsupported non-RHT " + "grouped weight quantization; enable quantized model initialization" + ) use_fp8 = fp8_recipe is not None with quantized_model_init(enabled=use_fp8 and fp8_model_params, recipe=fp8_recipe): @@ -625,6 +644,7 @@ def test_sanity_grouped_linear( params_dtype=dtype, single_grouped_weight=single_param, single_grouped_bias=single_param, + use_grouped_tensor=single_param, ).cuda() # Verify grouped linear exposes a single grouped weight parameter(and bias when applicable). @@ -644,11 +664,25 @@ def test_sanity_grouped_linear( m_splits[-1] = 0 elif empty_split == "middle": m_splits[num_gemms // 2] = 0 + if single_param: + m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cuda") + + if NVTE_TEST_NVINSPECT_ENABLED and single_param: + # DebugQuantizer operates on per-GEMM tensors, while single grouped parameters + # intentionally have no split-quantize fallback. + with pytest.raises( + RuntimeError, + match="TE debug features do not support single grouped parameters", + ): + with autocast(enabled=use_fp8, recipe=fp8_recipe): + te_grouped_linear(inp_hidden_states, m_splits) + return with autocast(enabled=use_fp8, recipe=fp8_recipe): out = te_grouped_linear(inp_hidden_states, m_splits) loss = out.sum() loss.backward() + torch.cuda.synchronize() assert out.shape == (num_tokens, ffn_hidden_size) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f61e7b4111..6ebc408dcb 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,6 +4,7 @@ import abc import contextlib +import warnings import pytest import torch @@ -21,6 +22,8 @@ except ImportError: _opaque_available = False +from torch._dynamo.utils import counters + import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.common import recipe @@ -41,7 +44,7 @@ MXFP8Quantizer, NVFP4Quantizer, ) -from utils import recipe_id +from utils import ModelConfig, dtype_tols, get_available_attention_backends, recipe_id from transformer_engine.pytorch.attention.dot_product_attention.backends import ( UnfusedDotProductAttention, ) @@ -547,14 +550,628 @@ def fn(q, k, v, extra): compiled = torch.compile(fn, fullgraph=True, mode="reduce-overhead") for _ in range(3): + # Compared against eager rather than only checked for finiteness: a + # replay reusing stale buffers would pass the latter. q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True) - out = compiled(q, k, v, extra) + inputs = (q, k, v, extra) + replayed = _run_and_capture(compiled, inputs, {}, [q, k, v]) + eager = _run_and_capture(fn, inputs, {}, [q, k, v]) + _assert_run_matches(replayed, eager, "unfused", dtype) + + +# --------------------------------------------------------------------------- +# DotProductAttention under torch.compile +# --------------------------------------------------------------------------- + + +# Model configurations, described with the same ModelConfig the eager attention +# tests use. Which backends can run each of them is not hardcoded here -- +# get_available_attention_backends() answers that, so configurations only one +# backend supports (arbitrary masks, biases, MLA head dims, ...) are covered +# rather than avoided. +def _cfg( + model_config, + qkv_format="bshd", + packed=None, + interleave_dim=-3, + share_cu_seqlens=False, + kv_cache=False, +): + """A model configuration plus what DotProductAttention is handed it as. + + `packed` is "qkv" or "kv" for the declarative packed inputs, interleaved at + `interleave_dim`, or None for separate q/k/v. `kv_cache` runs the call as a + decoding step against an InferenceParams KV cache. + """ + return dict( + model_config=model_config, + qkv_format=qkv_format, + packed=packed, + interleave_dim=interleave_dim, + share_cu_seqlens=share_cu_seqlens, + kv_cache=kv_cache, + ) + + +def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str: + """bshd + 3 @ -3 -> bs3hd; bshd + 2 @ -2 -> bsh2d; as DotProductAttention + derives it from the declaration.""" + position = len(qkv_format) + interleave_dim + 1 + return qkv_format[:position] + str(packed_dim) + qkv_format[position:] + + +_DPA_COMPILE_CONFIGS = { + "self_bshd_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal")), + "self_sbhd_no_mask": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd"), + "self_bshd_swa": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0))), + "gqa_bshd_causal": _cfg(ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal")), + "cross_bshd_no_mask": _cfg( + ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask") + ), + "self_bshd_padding_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal")), + "self_thd_padding_causal": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd" + ), + # Self attention naturally passes one cu_seqlens tensor for both q and kv, + # which flash-attn hands to two inputs of the same autograd.Function. + "self_thd_shared_cu_seqlens": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd", share_cu_seqlens=True + ), + "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv"), + "packed_qkv_bsh3d": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv", interleave_dim=-2 + ), + "packed_kv_bshd_bs2hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="kv"), + "packed_kv_thd_th2d": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), + "thd", + packed="kv", + interleave_dim=-2, + ), + # Configurations below are supported by one backend only, or take a code + # path of their own inside DotProductAttention. + "alibi_bshd_causal": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi") + ), + "post_scale_bias_bshd": _cfg( + ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss") + ), + "arbitrary_mask_bshd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary")), + "mla_bshd_causal": _cfg(ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal")), + "sink_softmax_bshd": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one") + ), + # FlashAttention takes a fused sbhd->bshd split for sbh3d with a head + # dimension of 128 and at least 512 tokens, and a plain transpose otherwise. + "packed_qkv_sbh3d_fused_split": _cfg( + ModelConfig(4, 128, 4, 128, attn_mask_type="no_mask"), + "sbhd", + packed="qkv", + interleave_dim=-2, + ), + # A decoding step against a KV cache. FlashAttention 2 wants the non-paged + # cache length divisible by 256. + "kv_cache_bshd": _cfg( + ModelConfig(2, 8, 4, 64, max_seqlen_kv=256, attn_mask_type="padding_causal_bottom_right"), + kv_cache=True, + ), +} + + +def _qkv_layout(spec: dict) -> str: + qkv_format, packed = spec["qkv_format"], spec["packed"] + if packed is None: + return "_".join([qkv_format] * 3) + if packed == "qkv": + return _packed_layout(qkv_format, 3, spec["interleave_dim"]) + return f"{qkv_format}_{_packed_layout(qkv_format, 2, spec['interleave_dim'])}" + + +def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention: + config = spec["model_config"] + if spec["kv_cache"]: + # KV caching addresses the cache by layer number, and a decoding step + # runs in inference mode. + return ( + te.DotProductAttention( + num_attention_heads=config.num_heads, + kv_channels=config.kv_channels, + num_gqa_groups=config.num_gqa_groups, + qkv_format=spec["qkv_format"], + attn_mask_type=config.attn_mask_type, + layer_number=1, + ) + .to(dtype=dtype, device="cuda") + .eval() + ) + return te.DotProductAttention( + num_attention_heads=config.num_heads, + kv_channels=config.kv_channels, + num_gqa_groups=config.num_gqa_groups, + attention_dropout=config.dropout_p, + qkv_format=spec["qkv_format"], + attn_mask_type=config.attn_mask_type, + window_size=config.window_size, + attention_type=config.attn_type, + softmax_type=config.softmax_type, + ).to(dtype=dtype, device="cuda") + + +def _cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor: + cu = torch.zeros(seqlens.numel() + 1, dtype=torch.int32, device="cuda") + cu[1:] = torch.cumsum(seqlens, dim=0) + return cu + + +def _make_dpa_inputs(spec: dict, dtype: torch.dtype): + """Build the (args, kwargs) that `DotProductAttention.forward` is called + with, plus the list of tensors whose gradients the test compares.""" + config = spec["model_config"] + qkv_format, packed = spec["qkv_format"], spec["packed"] + b = config.batch_size + s_q, s_kv = config.max_seqlen_q, config.max_seqlen_kv + h, g = config.num_heads, config.num_gqa_groups + d_qk, d_v = config.head_dim_qk, config.head_dim_v + padded = "padding" in config.attn_mask_type + + kwargs = {} + if padded: + # Sequences shorter than the maximum, so the padding mask is not + # degenerate: with all sequences full it would be all-False, and any + # difference in how masking is compiled would be invisible. + seqlens_q = torch.randint(1, s_q, [b], dtype=torch.int32, device="cuda") + seqlens_kv = ( + seqlens_q + if config.attn_type == "self" + else (torch.randint(1, s_kv, [b], dtype=torch.int32, device="cuda")) + ) + cu_q = _cu_seqlens(seqlens_q) + cu_kv = cu_q if spec["share_cu_seqlens"] else _cu_seqlens(seqlens_kv) + kwargs.update(cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv, max_seqlen_q=s_q, max_seqlen_kv=s_kv) + t_q, t_kv = int(cu_q[-1]), int(cu_kv[-1]) + else: + t_q, t_kv = b * s_q, b * s_kv + + def _shape(s, t, heads, head_dim): + return { + "bshd": (b, s, heads, head_dim), + "sbhd": (s, b, heads, head_dim), + "thd": (t, heads, head_dim), + }[qkv_format] + + def _randn(shape): + return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + + if spec["kv_cache"]: + from collections import OrderedDict + from transformer_engine.pytorch.attention import InferenceParams + + inference_params = InferenceParams( + max_batch_size=b, + max_sequence_length=s_kv, + num_heads_kv=g, + head_dim_k=d_qk, + dtype=dtype, + qkv_format=qkv_format, + ) + inference_params.allocate_memory(1) + inference_params.pre_step(OrderedDict((i, s_q) for i in range(b))) + qkv = [ + torch.randn(_shape(s_q, t_q, heads, d_qk), dtype=dtype, device="cuda") + for heads in (h, g, g) + ] + # The sequence lengths come from the cache, not from cu_seqlens, and a + # decoding step has no gradients to compare. + return tuple(qkv), {"inference_params": inference_params}, [] + + if packed is not None: + # Declarative packed inputs: q/k/v are derived from one buffer by DPA + # itself, and the layout comes from the declaration -- the only packed + # layout that torch.compile supports. + assert h == g and d_qk == d_v, "packed inputs require uniform heads and head dims" + interleave_dim = spec["interleave_dim"] + + def _packed(seqlen, tokens, packed_dim): + leading = { + "bshd": (b, seqlen), + "sbhd": (seqlen, b), + "thd": (tokens,), + }[qkv_format] + trailing = (packed_dim, h, d_qk) if interleave_dim == -3 else (h, packed_dim, d_qk) + return _randn(leading + trailing) + + kwargs["qkv_interleave_dim"] = interleave_dim + if packed == "qkv": + qkv = _packed(s_q, t_q, 3) + kwargs["qkv_layer"] = qkv + grad_tensors, args = [qkv], () + else: + q = _randn(_shape(s_q, t_q, h, d_qk)) + kv = _packed(s_kv, t_kv, 2) + kwargs["kv_layer"] = kv + grad_tensors, args = [q, kv], (q,) + else: + q = _randn(_shape(s_q, t_q, h, d_qk)) + k = _randn(_shape(s_kv, t_kv, g, d_qk)) + v = _randn(_shape(s_kv, t_kv, g, d_v)) + grad_tensors, args = [q, k, v], (q, k, v) + + if config.attn_mask_type == "arbitrary": + kwargs["attention_mask"] = torch.zeros(b, 1, s_q, s_kv, dtype=torch.bool, device="cuda") + if config.attn_bias_type != "no_bias": + kwargs["core_attention_bias_type"] = config.attn_bias_type + if config.attn_bias_type == "post_scale_bias": + kwargs["core_attention_bias"] = torch.randn(1, h, s_q, s_kv, dtype=dtype, device="cuda") + + return args, kwargs, grad_tensors + + +def _skip_unsupported( + spec: dict, backend: str, dtype, compiled: bool = True, inference_params=None +) -> None: + """Skip what the backend under test cannot run, or -- for a test that + compiles it -- cannot be compiled.""" + if compiled and backend == "fused": + # FusedAttention's forward carries @no_torch_dynamo, so there is nothing + # to compile: it runs as an eager island. Drop this skip once it traces, + # and the tests below cover it as they do the others. + pytest.skip("FusedAttention is an eager island and does not compile") + available, _, _ = get_available_attention_backends( + spec["model_config"], + dtype, + _qkv_layout(spec), + inference_params=inference_params, + is_training=inference_params is None, + ) + flash_supported, fused_supported, unfused_supported = available + supported = { + "flash": flash_supported, + "fused": fused_supported, + "unfused": unfused_supported, + }[backend] + if not supported: + pytest.skip(f"the {backend} backend does not support this configuration") + + +def _force_dpa_backend(monkeypatch, backend: str) -> None: + """Restrict DotProductAttention to a single backend.""" + from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import ( + _attention_backends, + ) + from transformer_engine.pytorch.attention.dot_product_attention.utils import ( + FlashAttentionUtils, + ) + + if backend == "flash" and not FlashAttentionUtils.is_installed: + pytest.skip("flash-attn is not installed") + + for name, var in ( + ("flash", "NVTE_FLASH_ATTN"), + ("fused", "NVTE_FUSED_ATTN"), + ("unfused", "NVTE_UNFUSED_ATTN"), + ): + monkeypatch.setenv(var, "1" if name == backend else "0") + # Backend selection is cached on the attention params only, so the env vars + # above are not enough to invalidate it. + _attention_backends["backend_selection_requires_update"] = True + + +def _assert_dpa_backend(backend: str) -> None: + from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import ( + _attention_backends, + ) + + assert _attention_backends[f"use_{backend}_attention"], ( + f"expected the {backend} backend to run, selected:" + f" flash={_attention_backends['use_flash_attention']}," + f" fused={_attention_backends['use_fused_attention']}," + f" unfused={_attention_backends['use_unfused_attention']}" + ) + + +def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) -> None: + """Assert a compiled result matches the eager one. + + A backend that calls the same kernel either way has to match exactly: + compiling changes what surrounds it, not its arithmetic. The unfused + backend is instead built from PyTorch ops, which inductor fuses and + reassociates; that error scales with the softmax sums rather than with each + output element, hence an absolute tolerance taken from the tensor's scale. + """ + if backend != "unfused": + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + return + tols = dtype_tols(dtype) + torch.testing.assert_close( + actual, + expected, + rtol=tols["rtol"], + atol=max(tols["atol"], 1e-3 * expected.abs().max().item()), + ) + + +def _run_and_capture(fn, args, kwargs, grads): + """Run `fn` and backward through its output, returning the output and a copy + of every input gradient. + + The output is cloned because CUDA graphs hand back tensors owned by their + memory pool, which the next replay overwrites, and the gradients are cleared + so that the same inputs can be reused for the other run. + """ + out = fn(*args, **kwargs).clone() + if grads: out.sum().backward() - torch.cuda.synchronize() - assert torch.isfinite(out).all() - assert q.grad is not None - assert k.grad is not None - assert v.grad is not None + torch.cuda.synchronize() + captured = [] + for tensor in grads: + assert tensor.grad is not None + captured.append(tensor.grad.clone()) + tensor.grad = None + return out, captured + + +def _assert_run_matches(actual, expected, backend: str, dtype: torch.dtype) -> None: + """Compare an (output, gradients) pair produced by `_run_and_capture`.""" + (out, out_grads), (ref, ref_grads) = actual, expected + _assert_matches_eager(out, ref, backend, dtype) + for out_grad, ref_grad in zip(out_grads, ref_grads): + _assert_matches_eager(out_grad, ref_grad, backend, dtype) + + +def _compare_compiled_to_eager( + module, args, kwargs, grads, monkeypatch, backend: str, dtype: torch.dtype, **compile_kwargs +) -> None: + """Run the module eagerly and compiled on the same inputs, and compare both + the output and every input gradient.""" + eager = _run_and_capture(module, args, kwargs, grads) + + torch._dynamo.reset() + # Force backend selection to be re-run (and traced) inside the compiled + # region instead of being served from the cache the eager call populated. + _force_dpa_backend(monkeypatch, backend) + compiled = _run_and_capture(torch.compile(module, **compile_kwargs), args, kwargs, grads) + + _assert_dpa_backend(backend) + _assert_run_matches(compiled, eager, backend, dtype) + + +@pytest.mark.parametrize("backend", ["flash", "fused", "unfused"]) +@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys()) +def test_dpa_torch_compile(monkeypatch, backend, config): + """`DotProductAttention` under `torch.compile(fullgraph=True)` must match + eager in forward and backward, for every backend that supports the + configuration. + + `fullgraph=True` makes the test fail on any graph break, so it covers the + whole module: input unpacking, qkv layout, backend selection and the backend + itself. + """ + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS[config] + module = _make_dpa(spec, dtype) + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + _skip_unsupported(spec, backend, dtype, inference_params=kwargs.get("inference_params")) + _force_dpa_backend(monkeypatch, backend) + + _compare_compiled_to_eager( + module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=True + ) + + +def test_dpa_torch_compile_around_fused(monkeypatch): + """FusedAttention itself is an eager island, but everything around it is + compiled: DotProductAttention traces up to the backend call, breaks the + graph there and resumes afterwards. What crosses that break has to survive + it -- the sub-backend enum did not, and reached cuDNN as the function that + produced it.""" + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"] + _skip_unsupported(spec, "fused", dtype, compiled=False) + _force_dpa_backend(monkeypatch, "fused") + + module = _make_dpa(spec, dtype) + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + # No fullgraph: the graph break at the eager island is the point here. + _compare_compiled_to_eager(module, args, kwargs, grads, monkeypatch, "fused", dtype) + + +@pytest.mark.parametrize("backend", ["flash", "unfused"]) +@pytest.mark.parametrize("config", ["self_bshd_causal", "kv_cache_bshd"]) +def test_dpa_torch_compile_cudagraphs(monkeypatch, backend, config): + """`mode="reduce-overhead"`: forward and backward of DotProductAttention + are captured into CUDA graphs and replayed on subsequent iterations.""" + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS[config] + _force_dpa_backend(monkeypatch, backend) + + module = _make_dpa(spec, dtype) + + torch._dynamo.reset() + counters.clear() + compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead") + + for _ in range(3): + # Fresh inputs every iteration: a replay that reuses stale buffers would + # still produce finite values and non-None gradients, so only comparing + # against eager catches it. + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + replayed = _run_and_capture(compiled, args, kwargs, grads) + eager = _run_and_capture(module, args, kwargs, grads) + _assert_run_matches(replayed, eager, backend, dtype) + _assert_dpa_backend(backend) + # Without this, inductor declining to capture -- a mutated input, a CPU + # scalar -- would leave the test passing while measuring nothing. + assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs" + + +@pytest.mark.parametrize("backend", ["flash", "unfused"]) +@pytest.mark.parametrize("paged", [False, True], ids=["non_paged", "paged"]) +@pytest.mark.parametrize("cuda_graphs", [False, True], ids=["default", "cudagraphs"]) +def test_dpa_torch_compile_kv_cache_decoding(monkeypatch, backend, paged, cuda_graphs): + """Generation against a KV cache, one prefill and three single-token steps. + + The cache carries state from one step to the next, so a step that updates it + wrongly is only visible in the step after -- hence comparing every step, and + a cache of its own for each of the two runs. + """ + from collections import OrderedDict + from transformer_engine.pytorch.attention import InferenceParams + + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS["kv_cache_bshd"] + config = spec["model_config"] + b, ctx_len = config.batch_size, config.max_seqlen_q + h, g, d = config.num_heads, config.num_gqa_groups, config.head_dim_qk + + def make_cache(): + kwargs = dict( + max_batch_size=b, + max_sequence_length=config.max_seqlen_kv, + num_heads_kv=g, + head_dim_k=d, + dtype=dtype, + qkv_format=spec["qkv_format"], + ) + if paged: + # One page per sequence, and FlashAttention 2 wants it divisible by 256. + kwargs.update(is_paged=True, page_size=config.max_seqlen_kv, total_num_pages=b) + inference_params = InferenceParams(**kwargs) + inference_params.allocate_memory(1) + return inference_params + + _skip_unsupported(spec, backend, dtype, inference_params=make_cache()) + + gen = torch.Generator(device="cuda").manual_seed(1234) + steps = [ + [ + torch.randn(b, seqlen, heads, d, device="cuda", dtype=dtype, generator=gen) + for heads in (h, g, g) + ] + for seqlen in (ctx_len, 1, 1, 1) + ] + + _force_dpa_backend(monkeypatch, backend) + module = _make_dpa(spec, dtype) + + def generate(fn): + inference_params = make_cache() + outputs = [] + with torch.no_grad(): + for args in steps: + inference_params.pre_step(OrderedDict((i, args[0].shape[1]) for i in range(b))) + _force_dpa_backend(monkeypatch, backend) + # Cloned because a CUDA graph replay overwrites what it returned. + outputs.append(fn(*args, inference_params=inference_params).clone()) + return outputs + + eager = generate(module) + + torch._dynamo.reset() + counters.clear() + compile_kwargs = {"mode": "reduce-overhead"} if cuda_graphs else {} + compiled = generate(torch.compile(module, fullgraph=True, **compile_kwargs)) + + _assert_dpa_backend(backend) + for out, ref in zip(compiled, eager): + _assert_matches_eager(out, ref, backend, dtype) + if cuda_graphs: + assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("fp8_attention", [False, True], ids=["fp8_gemms_only", "fp8_attention"]) +def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention): + """FP8 attention is not supported on the compiled path and falls back to + eager. FP8 elsewhere in the model with attention in high precision -- the + common training setup -- must stay compiled.""" + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"] + _force_dpa_backend(monkeypatch, "unfused") + + module = _make_dpa(spec, dtype) + args, kwargs, _ = _make_dpa_inputs(spec, dtype) + fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention) + + def fn(*args, **kwargs): + with te.autocast(enabled=True, recipe=fp8_recipe): + return module(*args, **kwargs) + + torch._dynamo.reset() + _force_dpa_backend(monkeypatch, "unfused") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + try: + torch.compile(fn)(*args, **kwargs) + except Exception: # pylint: disable=broad-except + # FP8 attention is not available on every device; what matters here + # is which path was taken, which the warning below reports. + pass + fell_back = any("Falling back to eager" in str(w.message) for w in caught) + + assert fell_back == fp8_attention + + +def _packed_views_inputs(spec, dtype, interleave_dim=-3): + """Packed q/k/v handed over as plain views, without declaring the packing.""" + config = spec["model_config"] + b, s = config.batch_size, config.max_seqlen_q + h, d = config.num_heads, config.head_dim_qk + shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d) + qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + q, k, v = [qkv.select(interleave_dim, i) for i in range(3)] + return (q, k, v), {}, [qkv] + + +def _thd_without_max_seqlen_inputs(spec, dtype): + """thd inputs that leave max_seqlen to be derived from cu_seqlens.""" + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + del kwargs["max_seqlen_q"], kwargs["max_seqlen_kv"] + return args, kwargs, grads + + +_EAGER_FALLBACK_CASES = { + # name: (config name, input builder) + "packed_views_bs3hd": ( + "self_bshd_causal", + lambda spec, dtype: _packed_views_inputs(spec, dtype, -3), + ), + "packed_views_bsh3d": ( + "self_bshd_causal", + lambda spec, dtype: _packed_views_inputs(spec, dtype, -2), + ), + "thd_without_max_seqlen": ("self_thd_padding_causal", _thd_without_max_seqlen_inputs), +} + + +@pytest.mark.parametrize("backend", ["flash", "unfused"]) +@pytest.mark.parametrize("case", _EAGER_FALLBACK_CASES.keys()) +def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case): + """Calls that cannot be traced run as an eager island instead, with a + warning: recognizing packed q/k/v takes the data pointers dynamo cannot + read, and deriving max_seqlen off cu_seqlens is a device synchronization. + Both keep working, and both are avoidable -- by declaring the packing via + qkv_layer/kv_layer, or by passing max_seqlen.""" + dtype = torch.bfloat16 + config_name, make_inputs = _EAGER_FALLBACK_CASES[case] + spec = _DPA_COMPILE_CONFIGS[config_name] + _force_dpa_backend(monkeypatch, backend) + + module = _make_dpa(spec, dtype) + args, kwargs, grads = make_inputs(spec, dtype) + + eager = _run_and_capture(module, args, kwargs, grads) + + torch._dynamo.reset() + _force_dpa_backend(monkeypatch, backend) + with pytest.warns(UserWarning, match="Falling back to eager execution"): + fell_back = _run_and_capture(torch.compile(module), args, kwargs, grads) + _assert_run_matches(fell_back, eager, backend, dtype) + + # The same eager island is an error when the user asked for a full graph. + torch._dynamo.reset() + _force_dpa_backend(monkeypatch, backend) + with pytest.raises(Exception, match="torch.compiler.disable"): + torch.compile(module, fullgraph=True)(*args, **kwargs) # --------------------------------------------------------------------------- diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 9734e24fbe..6ceadc7405 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -371,6 +371,7 @@ target_link_libraries(transformer_engine PUBLIC CUDA::cublas CUDA::cudart CUDNN::cudnn_all) +target_link_libraries(transformer_engine PRIVATE ${CMAKE_DL_LIBS}) target_include_directories(transformer_engine PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 231680321e..6b276d22a8 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -235,24 +235,49 @@ def _get_sys_extension() -> str: raise RuntimeError(f"Unsupported operating system ({system})") +def _cuda_runtime_major(cuda_runtime: ctypes.CDLL) -> Optional[int]: + """Return the major version of the CUDA runtime loaded with Transformer Engine.""" + + runtime_version = ctypes.c_int() + get_runtime_version = cuda_runtime.cudaRuntimeGetVersion + get_runtime_version.argtypes = [ctypes.POINTER(ctypes.c_int)] + get_runtime_version.restype = ctypes.c_int + if get_runtime_version(ctypes.byref(runtime_version)) != 0 or runtime_version.value <= 0: + return None + return runtime_version.value // 1000 + + @functools.lru_cache(maxsize=None) -def _nvidia_cudart_include_dir() -> str: +def _nvidia_cudart_include_dir(cuda_major_version: int) -> str: """Returns the include directory for cuda_runtime.h if exists in python environment.""" + # This is primarily here to support editable installs. cuda_runtime.cpp handles the + # resolution for install via wheel or when using the shared library ABI directly. + try: import nvidia except ModuleNotFoundError: return "" - # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia" - # above doesn't throw. However, they don't set "__file__" attribute. + # NVIDIA packages may use either a regular package or a namespace package spread + # across multiple package roots. if nvidia.__file__ is not None: - nvidia_root = Path(nvidia.__file__).parent + nvidia_roots = (Path(nvidia.__file__).parent,) else: - nvidia_root = Path(nvidia.__path__[0]) # namespace package + nvidia_roots = tuple(Path(path) for path in nvidia.__path__) + + layouts = [f"cu{cuda_major_version}"] + if cuda_major_version == 12: + layouts.append("cuda_runtime") - include_dir = nvidia_root / "cuda_runtime" - return str(include_dir) if include_dir.exists() else "" + for layout in layouts: + for nvidia_root in nvidia_roots: + cuda_root = nvidia_root / layout + if (cuda_root / "cuda_runtime.h").is_file() or ( + cuda_root / "include" / "cuda_runtime.h" + ).is_file(): + return str(cuda_root) + return "" @functools.lru_cache(maxsize=None) @@ -382,5 +407,8 @@ def _load_core_library(): _TE_LIB_CTYPES = _load_core_library() # Needed to find the correct headers for NVRTC kernels. - if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _nvidia_cudart_include_dir(): - os.environ["NVTE_CUDA_INCLUDE_DIR"] = _nvidia_cudart_include_dir() + _cuda_major_version = _cuda_runtime_major(_TE_LIB_CTYPES) + if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _cuda_major_version is not None: + cuda_include_dir = _nvidia_cudart_include_dir(_cuda_major_version) + if cuda_include_dir: + os.environ["NVTE_CUDA_INCLUDE_DIR"] = cuda_include_dir diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh index 522a9add8f..52ab6f49e9 100644 --- a/transformer_engine/common/cast/fp8/gated_fp8.cuh +++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh @@ -68,11 +68,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_elems = SHMEM_DIM_Y * SHMEM_DIM_X; constexpr size_t buff_elems_total = BUFFERS_NUM * buff_elems; diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 83b5a49cae..879c6befc3 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -124,11 +124,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) __shared__ float subamax_colwise_buff[SUBAMAX_BUFF_DIM_Y][CHUNK_DIM_X]; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 14832573d7..bc00c5824f 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -572,9 +572,18 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise); const size_t tensor_base = current_block.tensor_base; - const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1) - ? static_cast(offsets_ptr[tensor_id]) - : tensor_base; + size_t tensor_base_for_scales = tensor_base; + size_t tensor_rows_for_scales = rows; + if constexpr (WITH_GEMM_SWIZZLED_SCALES && SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + // The payload is one tall tensor, but GEMM scales are swizzled independently per member. + // Uniform groups omit first_dims/tensor_offsets, so derive the member geometry directly. + tensor_rows_for_scales = first_logical_dim / num_tensors; + tensor_base_for_scales = tensor_id * tensor_rows_for_scales * cols; + } + if constexpr (WITH_GEMM_SWIZZLED_SCALES && + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + tensor_base_for_scales = static_cast(offsets_ptr[tensor_id]); + } const size_t block_id_Y = current_block.block_id_Y; const size_t block_id_X = current_block.block_id_X; const size_t block_offset_Y = current_block.block_offset_Y; @@ -680,8 +689,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel process_colwise_stage( buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise, - scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr, - sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); + scale_stride_colwise, tensor_base_for_scales, tensor_rows_for_scales, cols, sIn_ptr, + sActIn_ptr, sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); } if constexpr (ROWWISE_SCALING) { diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 8c57f112d2..3900f072bb 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -130,11 +130,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0); extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); @@ -648,6 +646,30 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, float *const amax_ptr = reinterpret_cast(output->amax.dptr); const float *noop_ptr = reinterpret_cast(noop->data.dptr); + // Clear padding before either the generic or specialized kernel writes + // directly into the GEMM-swizzled scale layout. + if (with_gemm_swizzled_scales && (cols % 128 != 0 || rows % 128 != 0)) { + constexpr size_t zero_threads = 256; + if (use_rowwise_scaling) { + const size_t size_bytes = output->scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + if (use_colwise_scaling) { + const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->columnwise_scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + } + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( @@ -669,17 +691,26 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, bidimensional_traits::blockDIM::M) <= max_grid_dim_y; const bool is_full_rowwise_chunk = (cols % 128 == 0); + const bool has_full_bidimensional_chunks = + (rows % bidimensional_traits::colChunkElems == 0) && + (cols % bidimensional_traits::rowChunkElems == 0); + // Both rowwise and bidimensional cast-only kernels select their + // scale layout from WITH_GEMM_SWIZZLED_SCALES. const bool scaling_type_has_specialized_support = (scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk && rowwise_specialized_grid_fits) || - (scaling_type == ScalingType::BIDIMENSIONAL && + (scaling_type == ScalingType::BIDIMENSIONAL && has_full_bidimensional_chunks && bidimensional_specialized_grid_fits); - if (specialized::hasSpec() && - !WITH_GEMM_SWIZZLED_SCALES && scaling_type_has_specialized_support) { + // Specialized cast-only kernels do not consume the device noop flag. + // Preserve cached outputs by keeping noop-aware calls on the generic path. + if (noop_ptr == nullptr && + specialized::hasSpec() && + scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { - using traits = specialized::CastTraits; + using traits = specialized::CastTraits; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -698,7 +729,11 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, break; } case ScalingType::BIDIMENSIONAL: { - using traits = specialized::CastTraits; + using traits = + specialized::CastTraitsSwizzle; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -788,38 +823,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - // Zero out swizzled scales if padding is needed - /// TODO (tmoon) Handle this within the cast kernel - if (with_gemm_swizzled_scales) { - constexpr size_t TILE_DIM_X = 128; // Tile dim in data buffer - constexpr size_t TILE_DIM_Y = 128; - if (cols % TILE_DIM_X != 0 || rows % TILE_DIM_Y != 0) { - // Use a noop-aware zero kernel so that the clear is skipped - // when quantization is a noop (e.g. FP8 weight caching). - constexpr size_t zero_threads = 256; - if (use_rowwise_scaling) { - const size_t size_bytes = output->scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->scale_inv.dptr), size_bytes, - noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - if (use_colwise_scaling) { - const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->columnwise_scale_inv.dptr), - size_bytes, noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - } - } - switch (scaling_type) { case ScalingType::ROWWISE: { auto kernel = quantize_mxfp8_kernel #include "../../../util/ptx.cuh" +#include "../swizzle.cuh" // gemm_swizzled_scale_idx (parent dir, GEMM scale swizzle) #include "state_counter.cuh" -#include "swizzle.cuh" +#include "swizzle.cuh" // specialized/swizzle.cuh (TMA input bank-conflict swizzle) namespace transformer_engine { namespace dispatch { @@ -24,6 +25,10 @@ namespace quantize_kernel { namespace specialized { namespace ptx = transformer_engine::ptx; + +// Bring in the GEMM-swizzled scale index helper (from ../swizzle.cuh). +// Used only when the kernel is instantiated with CastTraits::_with_swizzled_scales=true. +using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; namespace { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -122,19 +127,20 @@ struct Layout { static constexpr int32_t num = M * N; }; -template +template struct CastTraits; // 1x32 -template -struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { +template +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false, _kSwizzled> { static constexpr bool isRowwise = true; static constexpr bool isColwise = false; using IType = _IType; using OType = _OType; static constexpr int32_t chunkElems = 32; - using threadLayout = Layout<1, 32>; + using threadLayout = Layout<1, THREADS_PER_WARP>; static constexpr int32_t numThreadsPerChunk = 1; static constexpr int32_t warpDimM = threadLayout::M; static constexpr int32_t warpDimN = threadLayout::N * chunkElems; @@ -151,14 +157,22 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { using iterLayout = Layout<1, 1>; static constexpr int32_t blockDimM = iterLayout::M * blockIterDimM; static constexpr int32_t blockDimN = iterLayout::N * blockIterDimN; + static constexpr int32_t rowwiseScaleStride = blockDimN / chunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 1; static constexpr int32_t numPrefetch = numStages - 1; static constexpr bool _use_cvt_4x = true; static constexpr bool _cache_rowwise_scale_in_smem = true; + static constexpr bool _with_swizzled_scales = _kSwizzled; - static constexpr int32_t numThreads = warpLayout::num * 32; + static constexpr int32_t numThreads = warpLayout::num * THREADS_PER_WARP; static constexpr size_t smem_rowwise_scale = _cache_rowwise_scale_in_smem ? (blockDimM * (blockDimN / chunkElems) * sizeof(e8m0_t)) : 0ul; @@ -498,13 +512,8 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re block_coords.y = blockIdx.y * CastTraits::blockDimM; block_coords.x = blockIdx.x * CastTraits::blockDimN; - constexpr int32_t stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDimM, rows); @@ -514,7 +523,40 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::chunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + if constexpr (CastTraits::_with_swizzled_scales) { + // Four adjacent rowwise scale columns are contiguous in the GEMM scale + // layout, so write them as one uint32_t whenever possible. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::chunkElems; + const size_t num_tiles_x = DIVUP(cols, static_cast(128)); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / groups_per_row; + const int32_t col = (i % groups_per_row) * cols_per_group; + const uint32_t value = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + *reinterpret_cast(&scales_rowwise[idx]) = value; + } + + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / remaining_per_row; + const int32_t col = remaining_start + (i % remaining_per_row); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + scales_rowwise[idx] = sRowwiseScale[row * stride_in_smem + col]; + } + } + } else if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { using DataType = int32_t; constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; @@ -527,8 +569,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -546,8 +589,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -577,11 +621,12 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr int32_t rowChunkElems = 32; static constexpr int32_t colChunkElems = 32; - using rowThreadLayout = Layout<32, 1>; // 32x1 + using rowThreadLayout = Layout; // 32x1 using colThreadLayout = Layout; // 1x32 static_assert(rowThreadLayout::num == colThreadLayout::num, "rowThreadLayout::num must be equal to colThreadLayout::num"); - static_assert(rowThreadLayout::num == 32, "rowThreadLayout::num must be 32"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); using rowWarpDim = Layout; using colWarpDim = Layout; @@ -603,6 +648,13 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { using iterLayout = Layout<1, 4>; using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 2; @@ -636,7 +688,7 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { "It requires aligned smem pointer"); static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; - static constexpr int32_t numThreads = numWarps * 32; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); @@ -660,7 +712,9 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; // && _colwise_reduce_max != ColwiseReduceMax::Redux; static constexpr size_t smem_colwise_reduce = - _need_smem_for_colwise_reduce ? 32 * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul; + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; static constexpr size_t smem = _reuse_input_out_smem @@ -670,9 +724,138 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { smem_alignment + smem_rowwise_scale + smem_colwise_reduce); }; -__device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) { - return (x + align - 1) & ~((align)-1); -} +// Standalone trait for the non-warp-specialized rowwise+colwise cast_only kernel. +// Exposes numStages, iterM, iterN, and the two colwise-scale features as +// caller-controllable template axes. Both colwise flags default to true, +// giving callers the swizzled + colwise-scale-cached fast path by default. +// +// This trait duck-types the same interface CastTraits<_, _, true, true> exposes, +// so it drops into quantize_mxfp8_kernel_cast_only without any +// changes to the kernel signature. Kernel #1 (rowwise-only) and Kernel #2 +// (warp-specialized row+col) don't accept it: kernel #1 requires isColwise=false, +// kernel #2 requires _use_warp_specialization=true - both are wrong here. +template +struct CastTraitsSwizzle { + static constexpr bool isRowwise = true; + static constexpr bool isColwise = true; + using IType = _IType; + using OType = _OType; + + static constexpr int32_t rowChunkElems = 32; + static constexpr int32_t colChunkElems = 32; + + using rowThreadLayout = Layout; // 32x1 + using colThreadLayout = Layout; // 1x32 + static_assert(rowThreadLayout::num == colThreadLayout::num, + "rowThreadLayout::num must be equal to colThreadLayout::num"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); + + using rowWarpDim = Layout; + using colWarpDim = Layout; + using warpDim = + Layout; + + static constexpr bool _tma_swizzle = true; + using warpLayout = Layout<1, 2>; + static_assert(_tma_swizzle ? (warpLayout::N == 2) : true); + static constexpr CUtensorMapSwizzle input_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + static constexpr CUtensorMapSwizzle output_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + using blockIterDim = Layout; + + using iterLayout = Layout<_IterM, _IterN>; + using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; + + static constexpr int32_t numStages = _NumStages; + + using inputUnitType = uint4; + static constexpr int32_t rowNumElemsPerUnit = sizeof(inputUnitType) / sizeof(IType); + static constexpr int32_t rowNumUnitsPerChunk = rowChunkElems / rowNumElemsPerUnit; + using inputElemSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>; + using inputUnitSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>; + + using colIndexSwz = swz::Swizzle<5, 0, 5>; + + using rowOutputUnitType = uint4; + static constexpr int32_t rowNumOutUnitsPerChunk = + rowChunkElems * sizeof(OType) / sizeof(rowOutputUnitType); + static constexpr int32_t rowOutNumElemsPerUnit = sizeof(rowOutputUnitType) / sizeof(OType); + + using rowOutputChunkSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>; + using colOutputSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>; + + static constexpr bool _use_cvt_4x = true; + static constexpr bool _use_warp_specialization = false; + static constexpr bool _need_wait_group = iterLayout::num > numStages; + static constexpr bool _reuse_input_out_smem = false; + static_assert(_reuse_input_out_smem == false, "Just don't use it"); + static constexpr bool _cache_rowwise_scale_in_smem = true; + + static constexpr bool _colwise_source_coming_from_rowwise = true; + static constexpr ColwiseReduceMax _colwise_reduce_max = ColwiseReduceMax::Redux; + static_assert(_colwise_reduce_max != ColwiseReduceMax::RedAsync, + "It requires aligned smem pointer"); + + // The two colwise-scale features exposed as caller-controllable trait axes. + // Both default to true so callers get the vectorized swizzled path by default. + static constexpr bool _cache_colwise_scale_in_smem = _kCacheColwise; + static constexpr bool _with_swizzled_scales = _kSwizzled; + + static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; + static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); + + static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); + static constexpr size_t smemInputPerBlock = smemInputPerWarp * warpLayout::num; + + static constexpr size_t smemRowwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemRowwiseOutputPerBlock = smemRowwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemColwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemColwiseOutputPerBlock = smemColwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemInput = smemInputPerBlock * numStages; + static constexpr size_t smemRowwiseOutput = smemRowwiseOutputPerBlock * numStages; + static constexpr size_t smemColwiseOutput = smemColwiseOutputPerBlock * numStages; + + static constexpr size_t smem_rowwise_scale = + _cache_rowwise_scale_in_smem ? (blockDIM::M * (blockDIM::N / rowChunkElems) * sizeof(e8m0_t)) + : 0ul; + + // Extra shmem for cached colwise scales - only when the flag is on. + static constexpr size_t smem_colwise_scale = + _cache_colwise_scale_in_smem ? (blockDIM::M / colChunkElems) * blockDIM::N * sizeof(e8m0_t) + : 0ul; + + using ColwiseReduceDataType = float; + static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; + static constexpr size_t smem_colwise_reduce = + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; + + static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; + static constexpr size_t smem = + _reuse_input_out_smem + ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce) + : (smemInput + smemRowwiseOutput + smemColwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce); +}; // 32x32 template ( - align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + char *smemAligned = align_up(smem, CastTraits::smem_alignment); IType *sInput = reinterpret_cast(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(sInput); @@ -728,12 +910,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } // TODO: maybe we can assign a different barrier for each warp @@ -744,8 +926,8 @@ __global__ void quantize_mxfp8_kernel_cast_only( #pragma unroll for (int32_t i = 0; i < CastTraits::numStages; i++) { ptx::mbarrier_init(&ldg_producer[i], 1); - ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * 32); - ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * 32); + ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); + ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); ptx::mbarrier_init(&stg_consumer[i], 1); } ptx::fence_mbarrier_init_release_cluster(); @@ -1085,15 +1267,10 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { - ptx::numbered_barrier_sync(CastTraits::warpLayout::num * 32, 0u); + ptx::numbered_barrier_sync(CastTraits::warpLayout::num * THREADS_PER_WARP, 0u); - constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); @@ -1117,8 +1294,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1137,8 +1315,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1179,8 +1358,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( block_coords.x = blockIdx.x * CastTraits::blockDIM::N; extern __shared__ char smem[]; - char *smemAligned = reinterpret_cast( - align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + char *smemAligned = align_up(smem, CastTraits::smem_alignment); IType *sInput = reinterpret_cast(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(sInput); @@ -1191,15 +1369,20 @@ __global__ void quantize_mxfp8_kernel_cast_only( // colwise output will reuse input buffer OType *sColOutput; e8m0_t *sRowwiseScale = nullptr; + e8m0_t *sColwiseScale = nullptr; ColwiseReduceDataType *sColwiseReduce = nullptr; if constexpr (CastTraits::_reuse_input_out_smem) { sColOutput = reinterpret_cast(sInput); if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sRowOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1211,9 +1394,13 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1223,7 +1410,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( rowOutputUnitType *sColOutputUnit = reinterpret_cast(sColOutput); if constexpr (CastTraits::_need_smem_for_colwise_reduce) { - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } __shared__ uint64_t producer[CastTraits::numStages]; @@ -1243,7 +1430,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_colwise_source_coming_from_rowwise && CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { - ptx::mbarrier_init(colwise_reduce_barrier, 32); + ptx::mbarrier_init(colwise_reduce_barrier, THREADS_PER_WARP); } ptx::fence_mbarrier_init_release_cluster(); @@ -1263,18 +1450,35 @@ __global__ void quantize_mxfp8_kernel_cast_only( (threadIdx.x % CastTraits::rowThreadLayout::N) * (CastTraits::rowChunkElems / CastTraits::rowNumElemsPerUnit); - size_t rowwise_scale_base_offset = - (block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N)) * - static_cast(scale_stride_rowwise) + + // Scale coordinates in absolute (compact) scale-tensor space. Shared by both + // compact-layout offsets and by the CastTraits::_with_swizzled_scales branches below. + const int32_t row_scale_row_base = + block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N); + const int32_t row_scale_col_base = (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::rowThreadLayout::N) * CastTraits::rowChunkElems) / - CastTraits::rowChunkElems; + CastTraits::rowChunkElems; + const int32_t col_scale_row_base = + (block_coords.y + warp_coords.y + + (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / + CastTraits::colChunkElems; + const int32_t col_scale_col_base = + block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N); + + size_t rowwise_scale_base_offset = + static_cast(row_scale_row_base) * static_cast(scale_stride_rowwise) + + row_scale_col_base; size_t colwise_scale_base_offset = - ((block_coords.y + warp_coords.y + - (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / - CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N)); + static_cast(col_scale_row_base) * static_cast(scale_stride_colwise) + + col_scale_col_base; + + // Precomputed swizzle constants (each swizzle tile is 128 rows x 4 cols in scale space). + // Rowwise scale tensor has DIVUP(cols, 128) tiles across; + // colwise scale tensor has DIVUP(rows, 128) tiles across (X/Y axes are transposed). + const size_t row_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(cols, static_cast(128)) : 0; + const size_t col_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(rows, static_cast(128)) : 0; constexpr int32_t rowwise_scale_stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; @@ -1401,6 +1605,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( iter_m * CastTraits::blockIterDim::M * rowwise_scale_stride_in_smem + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); sRowwiseScale[rowwise_scale_offset] = row_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = row_scale_row_base + iter_m * CastTraits::blockIterDim::M; + int32_t abs_col = row_scale_col_base + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = row_biased_exponent; } else { size_t rowwise_scale_offset = rowwise_scale_base_offset + @@ -1416,12 +1626,32 @@ __global__ void quantize_mxfp8_kernel_cast_only( e8m0_t col_biased_exponent = to_e8m0(col_amax); float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); sColwiseReduce[threadIdx.x] = col_scale_inverse; - size_t colwise_scale_offset = - colwise_scale_base_offset + - iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - iter_n * CastTraits::blockIterDim::N; - scales_colwise[colwise_scale_offset] = col_biased_exponent; + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + // Cache in shmem; end-of-kernel flush handles gmem indexing. + int32_t smem_row = (warp_coords.y + (threadIdx.x / CastTraits::colThreadLayout::N) * + CastTraits::colChunkElems) / + CastTraits::colChunkElems + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t smem_col = warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N) + + iter_n * CastTraits::blockIterDim::N; + sColwiseScale[smem_row * CastTraits::blockDIM::N + smem_col] = col_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = col_scale_row_base + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t abs_col = col_scale_col_base + iter_n * CastTraits::blockIterDim::N; + // Colwise scale tensor's X/Y axes are transposed vs rowwise + // (see col_swz_num_tiles_X = DIVUP(rows, 128)), so pass + // (abs_col, abs_row) - abs_col is the swizzle "row" dim. + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + scales_colwise[idx] = col_biased_exponent; + } else { + size_t colwise_scale_offset = + colwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + iter_n * CastTraits::blockIterDim::N; + scales_colwise[colwise_scale_offset] = col_biased_exponent; + } __syncwarp(); } } @@ -1546,58 +1776,210 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); - end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, - scale_stride_rowwise); + if constexpr (CastTraits::_with_swizzled_scales) { + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + DIVUP(cols, static_cast(CastTraits::rowChunkElems))); + } else { + // The compact layout's padded entries are consumed by a later swizzle. + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); + } int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { - using DataType = int32_t; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + if constexpr (CastTraits::_with_swizzled_scales) { + // Swizzled flush: within a 128x4 swizzle tile, 4 consecutive column entries + // for the same row live at 4 consecutive gmem bytes. Group by 4 so each + // thread writes a uint32_t when col%4 == 0, then scalar tail for remainder. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::rowChunkElems; + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + + uint32_t val4 = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); + + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + *reinterpret_cast(&scales_rowwise[idx]) = val4; + } - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + // Tail (valid_coords.x % 4 != 0) + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_per_row; + int32_t col = remaining_start + (i % remaining_per_row); + e8m0_t val = sRowwiseScale[row * stride_in_smem + col]; + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = val; + } } } else { - using DataType = PreferredDataType; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + using PreferredDataType = typename CastTraits::PreferredDataType; - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + using DataType = int32_t; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } else { + using DataType = PreferredDataType; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } + } + } + + // Cached colwise scale flush (swizzled path only). Same barrier semantics as + // the rowwise flush above: the last-iter __syncthreads already ordered every + // in-loop sColwiseScale byte store before this block reads them. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + // In GEMM swizzle, contiguous bytes run over four scale-row indices for a + // fixed logical column. A 64-row CTA owns two such rows; pack them as a + // uint16 store only when the pair stays inside the same 4-row swizzle group. + const int32_t row_pairs = valid_rows / 2; + const int32_t total_pairs = row_pairs * valid_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_pairs; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = (i / valid_cols) * 2; + int32_t col = i % valid_cols; + const int32_t abs_col = block_coords.x + col; + const int32_t abs_row = scale_row_base + row; + e8m0_t val0 = sColwiseScale[row * CastTraits::blockDIM::N + col]; + e8m0_t val1 = sColwiseScale[(row + 1) * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + if (((abs_row & 3) != 3) && + ((reinterpret_cast(&scales_colwise[idx]) & (alignof(uint16_t) - 1)) == 0)) { + uint16_t val2 = static_cast(val0) | (static_cast(val1) << 8); + *reinterpret_cast(&scales_colwise[idx]) = val2; + } else { + scales_colwise[idx] = val0; + size_t idx1 = gemm_swizzled_scale_idx(abs_col, abs_row + 1, col_swz_num_tiles_X); + scales_colwise[idx1] = val1; + } + } + // Odd-row tail. + if ((valid_rows & 1) != 0) { + const int32_t row = valid_rows - 1; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < valid_cols; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t col = i; + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(block_coords.x + col, scale_row_base + row, + col_swz_num_tiles_X); + scales_colwise[idx] = val; + } + } + } + + // Cached colwise scale flush (non-swizzled linear layout). + // Rows in scales_colwise are indexed by (block_coords.y / colChunkElems), + // columns are logical input columns; 4 adjacent columns for the same scale + // row live at 4 consecutive gmem bytes, so pack as uint32 stores. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && !CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_cols / cols_per_group; + const int32_t total_groups = valid_rows * groups_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + uint32_t val4 = + *reinterpret_cast(&sColwiseScale[row * CastTraits::blockDIM::N + col]); + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + *reinterpret_cast(&scales_colwise[idx]) = val4; + } + // Column tail (valid_cols not multiple of 4). + const int32_t remaining_start = groups_per_row * cols_per_group; + if (remaining_start < valid_cols) { + const int32_t remaining_cols = valid_cols - remaining_start; + const int32_t total_remaining = valid_rows * remaining_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_cols; + int32_t col = remaining_start + (i % remaining_cols); + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + scales_colwise[idx] = val; } } } @@ -1605,7 +1987,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( ptx::cp_async_bulk_wait_group_read<0>(); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} +} // NOLINT(readability/fn_size) } // namespace specialized } // namespace quantize_kernel diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index 91c6af26b5..0b980c5086 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -262,11 +262,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a38a620ebe..8fc4cc7c46 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -410,11 +410,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); @@ -952,11 +950,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..8fd7d7480c 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -1153,6 +1153,26 @@ inline bool is_aligned_ptr(const void *ptr, size_t alignment) { return reinterpret_cast(ptr) % alignment == 0; } +/*! \brief Align a shared-memory base pointer up to `align` bytes. + * + * The result is derived from `p` by pointer arithmetic on purpose, without losing its + * identity as a pointer in between, so the address is never rounded through an integer + * -- in which case the compiler would lose the link back to the `extern __shared__` + * object, and ptxas could no longer prove the address lives in the shared window and + * would fall back to generic address-space accesses (`LD.E`/`ST.E`) instead of + * `LDS`/`STS`. + * + * `align` must be a power of two. + */ +__device__ __forceinline__ char *align_up(char *p, uintptr_t align) { + const uintptr_t misalign = reinterpret_cast(p) & (align - 1); + // If p is not aligned, (align - misalign) & (align - 1) is the number of bytes to fill the gap between p and + // the next aligned address. + // If p is aligned, misalign is 0 and (align - misalign) is align itself, so we use & (align - 1) + // to make it 0 and return p itself. + return p + ((align - misalign) & (align - 1)); +} + inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) { return is_aligned_ptr(static_cast(t.data.dptr), alignment); } diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 2725008257..b726784d03 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -47,26 +47,50 @@ ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) { return ncclFloat8e4m3; case kNVTEFloat8E5M2: return ncclFloat8e5m2; + case kNVTEFloat8E8M0: + return ncclUint8; default: NVTE_ERROR("Unsupported NVTEDType for NCCL dtype conversion: ", static_cast(dtype)); } return ncclFloat32; // unreachable } -// shape_out is caller-owned; desc.sizes aliases shape_out.data and must -// outlive the NCCL EP call. +// Which part of a TE tensor an NCCL descriptor points at: the data payload, or +// (for block-scaled tensors) the rowwise scale-inverse that rides alongside it. +enum class DescSource { kData, kScaleInv }; + +// Build an NCCL descriptor for a TE tensor's data (kData) or its rowwise +// scale-inverse (kScaleInv). shape_out is caller-owned; desc.sizes aliases +// shape_out.data and must outlive the NCCL EP call. Uses the matching window +// field (win.window / win.scale_window) when set, else the raw pointer. inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, NVTEShape& shape_out, - const NVTECommWindow& win = {}) { - shape_out = nvte_tensor_shape(t); + const NVTECommWindow& win = {}, + DescSource source = DescSource::kData) { ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT; + void* raw_ptr = nullptr; + ncclWindow_t win_hdl = nullptr; + uint64_t win_offset = 0; + if (source == DescSource::kData) { + shape_out = nvte_tensor_shape(t); + desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); + raw_ptr = nvte_tensor_data(t); + win_hdl = win.window; + win_offset = win.offset; + } else { + const SimpleTensor& si = convertNVTETensorCheck(t)->scale_inv; + shape_out = nvte_make_shape(si.shape.data(), si.shape.size()); + desc.datatype = te_dtype_to_nccl_dtype(static_cast(si.dtype)); + raw_ptr = si.dptr; + win_hdl = win.scale_window; + win_offset = win.scale_offset; + } desc.ndim = shape_out.ndim; desc.sizes = shape_out.data; - desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); - if (win.window != nullptr) { - desc.win_hdl = win.window; - desc.win_offset = win.offset; + if (win_hdl != nullptr) { + desc.win_hdl = win_hdl; + desc.win_offset = win_offset; } else { - desc.data = nvte_tensor_data(t); + desc.data = raw_ptr; NVTE_CHECK(desc.data != nullptr || nvte_tensor_numel(t) == 0, "non-empty tensor data must not be null"); } @@ -412,16 +436,44 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_shape, recv_topk_weights_win); } + // Block-scaled (e.g. MXFP8): route the per-token scale-inverse alongside the + // data. High-precision (bf16/fp16/fp32) and per-tensor FP8 payloads carry the + // default delayed scaling mode and skip this. Keys on is_block_scaling so + // NVFP4 can reuse this path later. + const NVTEScalingMode tokens_scaling_mode = nvte_tensor_scaling_mode(tokens); + const bool is_scaled = is_block_scaling(tokens_scaling_mode); + NVTEShape scales_in_shape, scales_out_shape; + ncclEpTensor_t nccl_scales_in = NCCL_EP_TENSOR_INIT, nccl_scales_out = NCCL_EP_TENSOR_INIT; + if (is_scaled) { + NVTE_CHECK(is_mxfp8_scaling(tokens_scaling_mode), + "EP dispatch supports MXFP8 block scaling only; got scaling mode ", + static_cast(tokens_scaling_mode)); + NVTE_CHECK(nvte_tensor_scaling_mode(recv_tokens) == tokens_scaling_mode, + "recv_tokens scaling mode must match tokens scaling mode"); + nccl_scales_in = + make_nccl_ep_tensor(tokens, scales_in_shape, tokens_win, DescSource::kScaleInv); + nccl_scales_out = + make_nccl_ep_tensor(recv_tokens, scales_out_shape, recv_tokens_win, DescSource::kScaleInv); + } else { + NVTE_CHECK(!is_fp8_dtype(static_cast(tok_dtype)), + "EP dispatch of FP8 tokens requires a block scaling mode (e.g. MXFP8); " + "per-tensor (delayed) FP8 scaling is not supported"); + } + ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT; in_struct.tokens = &nccl_tokens_in; in_struct.topk_weights = is_forward ? &nccl_topk_weights_in : nullptr; + in_struct.scales = is_scaled ? &nccl_scales_in : nullptr; ncclEpDispatchOutputs_t out_struct = NCCL_EP_DISPATCH_OUTPUTS_INIT; out_struct.tokens = &nccl_tokens_out; out_struct.topk_weights = is_forward ? &nccl_topk_weights_out : nullptr; + out_struct.scales = is_scaled ? &nccl_scales_out : nullptr; ncclEpDispatchConfig_t dispatch_cfg = NCCL_EP_DISPATCH_CONFIG_INIT; dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS; + // Block-scaled payloads forward the per-token scale-inverse; select the matching recipe. + dispatch_cfg.quant_recipe = is_scaled ? NCCL_EP_DISP_QUANT_FWD : NCCL_EP_DISP_QUANT_NONE; std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 3997e5249d..3663106d50 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1324,7 +1324,8 @@ __global__ void setup_grouped_gemm_kernel( char *a_base, char *b_base, char *c_base, char *d_base, TensorShapeInfo A_meta, TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_bits_per_elem, size_t b_bits_per_elem, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, - float *beta_ptr, bool use_per_group_alpha_beta, + float *beta_ptr, bool use_per_group_alpha_beta, bool a_is_discrete, bool c_is_discrete, + bool d_is_discrete, // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise, @@ -1339,12 +1340,9 @@ __global__ void setup_grouped_gemm_kernel( if (idx >= num_tensors) return; // Get dimensions for this tensor (from array or uniform value) - const bool has_a_multi_tensor = (a_base == nullptr); - const bool has_c_multi_tensor = (c_base == nullptr); - const bool has_d_multi_tensor = (d_base == nullptr); int64_t a_first = 0; int64_t a_last = 0; - if (!has_a_multi_tensor) { + if (!a_is_discrete) { a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first; a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last; } @@ -1354,24 +1352,25 @@ __global__ void setup_grouped_gemm_kernel( int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last; // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform) - int64_t a_offset = has_a_multi_tensor ? 0 : compute_grouped_tensor_offset(A_meta, idx); + int64_t a_offset = a_is_discrete ? 0 : compute_grouped_tensor_offset(A_meta, idx); int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx); int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx); int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx); // Compute data pointers - A_ptrs[idx] = has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx] - : (a_base + (a_offset * a_bits_per_elem) / 8); - B_ptrs[idx] = b_base + (b_offset * b_bits_per_elem) / 8; - C_ptrs[idx] = - has_c_multi_tensor ? c_multi_tensor_args.data_ptrs[idx] : (c_base + c_offset * c_elem_size); - D_ptrs[idx] = - has_d_multi_tensor ? d_multi_tensor_args.data_ptrs[idx] : (d_base + d_offset * d_elem_size); + A_ptrs[idx] = a_is_discrete + ? a_multi_tensor_args.data_ptrs[idx] + : (a_base == nullptr ? nullptr : a_base + (a_offset * a_bits_per_elem) / 8); + B_ptrs[idx] = b_base == nullptr ? nullptr : b_base + (b_offset * b_bits_per_elem) / 8; + C_ptrs[idx] = c_is_discrete ? c_multi_tensor_args.data_ptrs[idx] + : (c_base == nullptr ? nullptr : c_base + c_offset * c_elem_size); + D_ptrs[idx] = d_is_discrete ? d_multi_tensor_args.data_ptrs[idx] + : (d_base == nullptr ? nullptr : d_base + d_offset * d_elem_size); // Compute storage dimensions for cuBLAS matrix layouts from logical dims. // Rowwise and MXFP8 columnwise storage use logical row-major layout, viewed as // column-major rows=last, cols=first. Transposed columnwise storage reverses this. - if (has_a_multi_tensor) { + if (a_is_discrete) { a_rows[idx] = a_multi_tensor_args.rows[idx]; a_cols[idx] = a_multi_tensor_args.cols[idx]; } else if (a_storage_transposed) { @@ -1388,7 +1387,7 @@ __global__ void setup_grouped_gemm_kernel( b_rows[idx] = static_cast(b_last); b_cols[idx] = static_cast(b_first); } - if (has_d_multi_tensor) { + if (d_is_discrete) { d_rows[idx] = d_multi_tensor_args.rows[idx]; d_cols[idx] = d_multi_tensor_args.cols[idx]; } else { @@ -1403,7 +1402,7 @@ __global__ void setup_grouped_gemm_kernel( if (use_per_group_alpha_beta) { float a_amax_val = 0.0f; bool has_a_amax = false; - if (has_a_multi_tensor) { + if (a_is_discrete) { auto *a_amax_p = static_cast(a_multi_tensor_args.amax_ptrs[idx]); if (a_amax_p != nullptr) { a_amax_val = *a_amax_p; @@ -1471,10 +1470,12 @@ __global__ void setup_grouped_gemm_kernel( } }; - if (a_scale_base) { + if (a_is_discrete) { + a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; + } else if (a_scale_base) { fill_scale_ptr(a_scale_inv_ptrs, a_scale_base, A_meta, a_rowwise, a_scaling_mode); } else { - a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; + a_scale_inv_ptrs[idx] = nullptr; } if (b_scale_base) { fill_scale_ptr(b_scale_inv_ptrs, b_scale_base, B_meta, b_rowwise, b_scaling_mode); @@ -1491,7 +1492,7 @@ inline void launch_grouped_gemm_setup( const transformer_engine::Tensor *beta_tensor, bool use_per_group_alpha_beta, size_t num_tensors, cudaStream_t stream, const MultiTensorGroupGemmInputArgs &a_multi_tensor_args, const NVTETensor *C_list, - const NVTETensor *D_list, char *a_base, transformer_engine::DType c_dtype, + const NVTETensor *D_list, bool a_is_discrete, char *a_base, transformer_engine::DType c_dtype, transformer_engine::DType d_dtype) { // Use logical shape info from selection; storage transposes are tracked separately. TensorShapeInfo A_meta = A_sel.logical_tensor_shape; @@ -1499,31 +1500,31 @@ inline void launch_grouped_gemm_setup( TensorShapeInfo C_meta{}; TensorShapeInfo D_meta{}; - const bool has_d_multi_tensor = (D_list != nullptr); - const bool has_c_multi_tensor = (C_list != nullptr) || has_d_multi_tensor; + const bool d_is_discrete = (D_list != nullptr); + const bool c_is_discrete = (C_list != nullptr) || d_is_discrete; MultiTensorGroupGemmOutputArgs c_multi_tensor_args{}; MultiTensorGroupGemmOutputArgs d_multi_tensor_args{}; - if (has_d_multi_tensor) { + if (d_is_discrete) { d_multi_tensor_args = build_grouped_gemm_multi_out_args(D_list, num_tensors, num_tensors, d_dtype, "D"); } if (C_list != nullptr) { c_multi_tensor_args = build_grouped_gemm_multi_out_args(C_list, num_tensors, num_tensors, d_dtype, "C"); - } else if (has_d_multi_tensor) { + } else if (d_is_discrete) { c_multi_tensor_args = d_multi_tensor_args; } char *c_base = nullptr; char *d_base = nullptr; - if (!has_c_multi_tensor) { + if (!c_is_discrete) { NVTE_CHECK(C != nullptr && D != nullptr, "Grouped GEMM: C/D grouped tensors are required when no C list is provided"); C_meta = TensorShapeInfo::create_shape_info_for_C(C, D); c_base = static_cast(C->data.dptr); } - if (!has_d_multi_tensor) { + if (!d_is_discrete) { NVTE_CHECK(D != nullptr, "Grouped GEMM: D grouped tensor is required when no D list is provided"); D_meta = TensorShapeInfo::from_tensor(D); @@ -1544,8 +1545,8 @@ inline void launch_grouped_gemm_setup( const bool b_rowwise = B_sel.rowwise; // NVFP4 alpha needs A's amax from either A_sel.amax (grouped) or amax_ptrs (discrete). - const bool a_has_amax = (A_sel.amax != nullptr) || - (A_sel.dptr == nullptr && a_multi_tensor_args.amax_ptrs[0] != nullptr); + const bool a_has_amax = + a_is_discrete ? a_multi_tensor_args.amax_ptrs[0] != nullptr : A_sel.amax != nullptr; const bool needs_nvfp4_alpha = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) && a_has_amax && (B_sel.amax != nullptr); @@ -1554,11 +1555,12 @@ inline void launch_grouped_gemm_setup( ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs, A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_bits_per_elem, b_bits_per_elem, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr), - static_cast(beta_tensor->data.dptr), use_per_group_alpha_beta, - reinterpret_cast(A_sel.scale_inv), reinterpret_cast(B_sel.scale_inv), - a_rowwise, b_rowwise, A_sel.storage_transposed, B_sel.storage_transposed, A_sel.scaling_mode, - B_sel.scaling_mode, num_tensors, a_multi_tensor_args, c_multi_tensor_args, - d_multi_tensor_args, A_sel.amax ? static_cast(A_sel.amax) : nullptr, + static_cast(beta_tensor->data.dptr), use_per_group_alpha_beta, a_is_discrete, + c_is_discrete, d_is_discrete, reinterpret_cast(A_sel.scale_inv), + reinterpret_cast(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.storage_transposed, + B_sel.storage_transposed, A_sel.scaling_mode, B_sel.scaling_mode, num_tensors, + a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args, + A_sel.amax ? static_cast(A_sel.amax) : nullptr, B_sel.amax ? static_cast(B_sel.amax) : nullptr, needs_nvfp4_alpha ? ws.nvfp4_computed_alpha : nullptr); @@ -1633,8 +1635,8 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, stream, - a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr, - inputC->dtype(), outputD->dtype()); + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, + /*a_is_discrete=*/false, A_sel.dptr, inputC->dtype(), outputD->dtype()); // Compute average dimensions for heuristics // K dimension: if transa, K is A's last dim; if not, K is A's first dim @@ -1788,8 +1790,8 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, stream, - a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr, - inputC->dtype(), outputD->dtype()); + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, + /*a_is_discrete=*/true, nullptr, inputC->dtype(), outputD->dtype()); GroupedGemmConfig gemm_config; gemm_config.use_split_accumulator = config_.use_split_accumulator; @@ -1873,8 +1875,8 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, /*C=*/nullptr, /*D=*/nullptr, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, - stream, a_multi_tensor_args, C_list, D_list, A_sel.dptr, d_dtype, - d_dtype); + stream, a_multi_tensor_args, C_list, D_list, + /*a_is_discrete=*/false, A_sel.dptr, d_dtype, d_dtype); GroupedGemmConfig gemm_config; gemm_config.use_split_accumulator = config_.use_split_accumulator; @@ -1909,8 +1911,6 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD, NVTE_CHECK(outputD->num_tensors >= 1, api_name, ": number of tensors must be at least 1"); NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, api_name, ": output and bias must have the same number of tensors"); - NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data"); - NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data"); NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), api_name, ": output and bias must have matching dtypes"); NVTE_CHECK(bias_tensor->all_same_first_dim(), api_name, @@ -1921,15 +1921,23 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD, NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), api_name, ": output and bias last dims must match"); + const int num_tensors = static_cast(outputD->num_tensors); + NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups, + " tensors, got ", num_tensors); + const int total_rows = static_cast(outputD->logical_shape.data[0]); + // A valid zero-sized CUDA allocation may have a null data pointer. + if (total_rows == 0) { + return; + } + + NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data"); + NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data"); + const TensorShapeInfo d_meta = TensorShapeInfo::from_tensor(outputD); const DType dtype = outputD->dtype(); constexpr int kThreads = 128; - const int num_tensors = static_cast(outputD->num_tensors); - NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups, - " tensors, got ", num_tensors); - const int total_rows = static_cast(outputD->logical_shape.data[0]); const int n = static_cast(outputD->get_common_last_dim()); const size_t elem_size = typeToSize(dtype); @@ -1997,8 +2005,6 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias); const Tensor *scale_tensor = convertNVTETensorCheck(scale); - NVTE_CHECK(scale_tensor->data.dptr != nullptr, - "Grouped scaled bias add: scale tensor must not be null"); NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32, "Grouped scaled bias add: scale must be float32"); NVTE_CHECK(scale_tensor->data.shape.size() == 1, @@ -2007,6 +2013,10 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou const size_t total_rows = static_cast(outputD->logical_shape.data[0]); NVTE_CHECK(scale_tensor->data.shape[0] == total_rows, "Grouped scaled bias add: scale size (", scale_tensor->data.shape[0], ") must equal total rows (", total_rows, ")"); + if (total_rows > 0) { + NVTE_CHECK(scale_tensor->data.dptr != nullptr, + "Grouped scaled bias add: scale tensor must not be null"); + } const float *scale_ptr = static_cast(scale_tensor->data.dptr); launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream); diff --git a/transformer_engine/common/include/transformer_engine/comm_window.h b/transformer_engine/common/include/transformer_engine/comm_window.h index 424c350bbd..ed67774493 100644 --- a/transformer_engine/common/include/transformer_engine/comm_window.h +++ b/transformer_engine/common/include/transformer_engine/comm_window.h @@ -26,6 +26,9 @@ struct ncclWindow_vidmem; typedef struct { struct ncclWindow_vidmem* window; /*!< NCCL window, or NULL to use the raw data pointer. */ uint64_t offset; /*!< Byte offset of the payload within window. */ + struct ncclWindow_vidmem* + scale_window; /*!< Window for a block-scaled tensor's scale-inverse, or NULL for raw. */ + uint64_t scale_offset; /*!< Byte offset of the scale-inverse within scale_window. */ } NVTECommWindow; #ifdef __cplusplus diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index a4800944b5..31417fbff1 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -150,6 +150,12 @@ void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv * *_win arguments enable zero-copy via symmem windows; pass NVTECommWindow{} * when unused. Requires a prior nvte_ep_prepare on this handle_mem. * + * tokens/recv_tokens may be high-precision (bf16/fp16) or FP8: + * for the latter, set rowwise data and rowwise scale-inverse (unswizzled + * [T, hidden_dim/block]) on the tensor and the scales are routed alongside the + * data. tokens and recv_tokens must share a scaling mode. For now, only MXFP8 + * (NVTE_MXFP8_1D_SCALING, e4m3 data + e8m0 scales) is supported. + * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] topk_idx [T, top_k] int64 sparse routing indices. * \param[in] tokens [T, hidden_dim] input tokens. diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index d5f2fa9a2c..7605798758 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -269,11 +269,11 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_stochastic_ro : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x), "r"(rbits)); return *reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt.rs PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - uint16_t dummy = 0; - return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy); + const float q0 = ptx::stochastic_round_fp4_e2m1(in01.x, rbits); + const float q1 = ptx::stochastic_round_fp4_e2m1(in01.y, rbits >> 8); + const float q2 = ptx::stochastic_round_fp4_e2m1(in23.x, rbits >> 16); + const float q3 = ptx::stochastic_round_fp4_e2m1(in23.y, rbits >> 24); + return __nv_fp4x4_e2m1(make_float4(q0, q1, q2, q3)); } } diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 504d761bb1..5690ef5e3c 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -7,6 +7,7 @@ #include "../util/cuda_runtime.h" #include +#include #include #include @@ -26,6 +27,83 @@ namespace { // String with build-time CUDA include path #include "string_path_cuda_include.h" +// Get the runtime directory of the shared library that contains this code +std::filesystem::path shared_library_directory() { + static const char library_anchor = 0; + Dl_info library_info{}; + if (dladdr(static_cast(&library_anchor), &library_info) == 0 || + library_info.dli_fname == nullptr) { + return {}; + } + + std::filesystem::path library_path = library_info.dli_fname; + if (library_path.is_relative()) { + std::error_code error; + library_path = std::filesystem::absolute(library_path, error); + if (error) { + return {}; + } + } + + return library_path.parent_path(); +} + +std::string runtime_cuda_major_version() { + int runtime_version = 0; + // Header discovery is best-effort, so do not throw if the runtime cannot + // report its version. + if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess || runtime_version <= 0) { + return {}; + } + + return std::to_string(runtime_version / 1000); +} + +std::filesystem::path python_cuda_directory() { + using Path = std::filesystem::path; + + // Find the Python package root from the installed Transformer Engine package. Do not + // assume that the root is named site-packages or dist-packages since valid installs + // may use an arbitrary target directory. + Path te_package_directory = shared_library_directory(); + while (true) { + if (te_package_directory.filename() == "transformer_engine") { + break; + } + + const Path parent = te_package_directory.parent_path(); + if (parent == te_package_directory) { + // Root directory reached + return {}; + } + + te_package_directory = parent; + } + + const Path nvidia_directory = te_package_directory.parent_path() / "nvidia"; + const auto cuda_major_version = runtime_cuda_major_version(); + if (cuda_major_version.empty()) { + return {}; + } + + std::error_code error; + const Path cuda_directory = nvidia_directory / ("cu" + cuda_major_version); + if (std::filesystem::is_directory(cuda_directory, error)) { + return cuda_directory; + } + + // CUDA 12 Python wheels use the older nvidia/cuda_runtime layout. + if (cuda_major_version == "12") { + error.clear(); + const Path legacy_cuda_directory = nvidia_directory / "cuda_runtime"; + if (std::filesystem::is_directory(legacy_cuda_directory, error)) { + return legacy_cuda_directory; + } + } + + return {}; +} + } // namespace int num_devices() { @@ -152,6 +230,7 @@ const std::string &include_directory(bool required) { std::vector> search_paths = {{"NVTE_CUDA_INCLUDE_DIR", ""}, {"CUDA_HOME", ""}, {"CUDA_DIR", ""}, + {"", python_cuda_directory()}, {"", string_path_cuda_include}, {"", "/usr/local/cuda"}}; for (auto &[env, p] : search_paths) { diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 2351ddae71..8dbb8b3044 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -380,6 +380,10 @@ __device__ __forceinline__ bf16 exp2f_rcp(e8m0_t biased_exp) { } __device__ __forceinline__ float exp2f(e8m0_t biased_exp) { + // Handle the special case of NaN. + if (biased_exp == 255) return __int_as_float(0x7fffffff); + // 2^-127 is subnormal, so it cannot be built by shifting into the exponent field. + if (biased_exp == 0) return __int_as_float(0x00400000); return __int_as_float(biased_exp << FP32_MANTISSA_BITS); } @@ -600,6 +604,29 @@ __device__ __forceinline__ void mul_cvt_4x(fp4e2m1x4 &out, const Tx2 &in01, cons out = fp4e2m1x4(make_float4(x0, x1, x2, x3)); } +// Software stochastic rounding onto the e2m1 grid, for architectures without +// cvt.rs. Takes 8 random bits per element, which is what +// cvt.rs.satfinite.e2m1x4.f32 takes from its rbits operand. Follows satfinite +// for the edges: NaN becomes positive MAX_NORM and larger magnitudes clamp to +// MAX_NORM with the sign kept. The result is exactly representable in e2m1, so +// packing it is lossless. +__device__ __forceinline__ float stochastic_round_fp4_e2m1(const float x, const uint32_t rbits8) { + constexpr float max_norm = 6.0f; + const float u = static_cast(rbits8 & 0xFFu) * (1.0f / 256.0f); + const float a = fabsf(x); + // Grid step at |x|: 0.5 below 2, 1 in [2, 4), 2 in [4, 6]. + const float step = (a >= 4.0f) ? 2.0f : ((a >= 2.0f) ? 1.0f : 0.5f); + const float t = fmaf(u, step, a); + // The jitter can carry t across one region boundary, so the rounding step + // derives from t. Flooring in units of that step lands on the e2m1 grid in + // every region, and the saturation also maps a NaN t to max_norm, which the + // sign select keeps positive. + const float step_t = (t >= 4.0f) ? 2.0f : ((t >= 2.0f) ? 1.0f : 0.5f); + const float inv_step_t = (t >= 4.0f) ? 0.5f : ((t >= 2.0f) ? 1.0f : 2.0f); + const float q = fminf(floorf(t * inv_step_t) * step_t, max_norm); + return copysignf(q, (x != x) ? 1.0f : x); +} + __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding( const uint64_t in_4x, const float2 scale, const uint32_t rbits) { uint16_t out_4x = 0; @@ -633,9 +660,14 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_roun : "=h"(out_4x) : "l"(in_4x), "l"(reinterpret_cast(scale)), "r"(rbits)); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + // mul.f32x2 above applies scale.x to the even elements and scale.y to the odd ones. + const bf16 *vals = reinterpret_cast(&in_4x); + const float q0 = stochastic_round_fp4_e2m1(static_cast(vals[0]) * scale.x, rbits); + const float q1 = stochastic_round_fp4_e2m1(static_cast(vals[1]) * scale.y, rbits >> 8); + const float q2 = stochastic_round_fp4_e2m1(static_cast(vals[2]) * scale.x, rbits >> 16); + const float q3 = stochastic_round_fp4_e2m1(static_cast(vals[3]) * scale.y, rbits >> 24); + const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3)); + out_4x = *reinterpret_cast(&packed); } return *reinterpret_cast(&out_4x); } @@ -725,9 +757,12 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_stochastic_roun "l"(reinterpret_cast(in23)), "l"(reinterpret_cast(scale)), "r"(rbits)); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + const float q0 = stochastic_round_fp4_e2m1(in01.x * scale.x, rbits); + const float q1 = stochastic_round_fp4_e2m1(in01.y * scale.y, rbits >> 8); + const float q2 = stochastic_round_fp4_e2m1(in23.x * scale.x, rbits >> 16); + const float q3 = stochastic_round_fp4_e2m1(in23.y * scale.y, rbits >> 24); + const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3)); + out_4x = *reinterpret_cast(&packed); } return *reinterpret_cast(&out_4x); } @@ -950,9 +985,26 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding( NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); } } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + constexpr bool known_coeff = std::is_same::value || + std::is_same::value; + if constexpr (known_coeff) { + const float coeff = static_cast(scaling_coefficient); + const bf16 *vals03 = reinterpret_cast(&in03); + const bf16 *vals47 = reinterpret_cast(&in47); + float q[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + q[i] = stochastic_round_fp4_e2m1(static_cast(vals03[i]) * coeff, rbits03 >> (8 * i)); + q[i + 4] = + stochastic_round_fp4_e2m1(static_cast(vals47[i]) * coeff, rbits47 >> (8 * i)); + } + const fp4e2m1x4 lo(make_float4(q[0], q[1], q[2], q[3])); + const fp4e2m1x4 hi(make_float4(q[4], q[5], q[6], q[7])); + out_8x = static_cast(*reinterpret_cast(&lo)) | + (static_cast(*reinterpret_cast(&hi)) << 16); + } else { + NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); + } } return out_8x; } diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index d0d472c6f2..ca70ea145c 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -14,6 +14,7 @@ compound ``(dp, ep)`` axis on the leading dim. """ +import functools from dataclasses import dataclass import jax @@ -34,7 +35,20 @@ def _on_collective_stream(func): return func from jax.experimental.compute_on import compute_on - return compute_on("gpu_stream:collective")(func) # pylint: disable=not-callable + @functools.wraps(func) + def wrapper(*args, **kwargs): + # compute_on traces its callee and abstract-evals every argument, so it + # cannot take the static EpLayerConfig/PartitionSpec args directly. Wrap + # a nullary thunk that closes over them; the array operands are captured + # as consts and lifted to real operands, outputs stay on device. XLA + # async-wraps the resulting call onto the collective stream. + annotated = compute_on( # pylint: disable=not-callable + compute_type="gpu_stream:collective", + out_memory_spaces=jax.memory.Space.Device, + )(lambda: func(*args, **kwargs)) + return annotated() + + return wrapper __all__ = [ diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index a49d4f80c2..21ef59ba5a 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -751,6 +751,8 @@ def _body(*args): num_local_tokens=(B, S), out_partition_spec=out_partition_spec, ) + # output of MLP should be sharded the same way as the activation input + output = with_sharding_constraint_by_logical_axes(output, input_axes) ( casted_sorted_x_lhs_trans, diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e797e83440..9d73ecd3b5 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -69,9 +69,16 @@ def is_triton_autotuned_alias_safe() -> bool: _COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.1" +@lru_cache(maxsize=None) def is_collective_stream_supported() -> bool: """Return True if the installed JAX supports the gpu_stream:collective annotation.""" - return jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION) + if not jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION): + return False + try: + from jax.experimental.compute_on import compute_on # pylint: disable=unused-import + except ImportError: + return False + return True def is_triton_extension_supported() -> bool: diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 4c6b7fc67a..2b1803bfb2 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -29,6 +29,7 @@ from transformer_engine.pytorch.module import destroy_ub from transformer_engine.pytorch.module import UserBufferQuantizationMode from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant from transformer_engine.pytorch.attention import MultiheadAttention from transformer_engine.pytorch.attention import InferenceParams from transformer_engine.pytorch.attention import RotaryPositionEmbedding diff --git a/transformer_engine/pytorch/attention/__init__.py b/transformer_engine/pytorch/attention/__init__.py index c4c2aa3e72..f6e4f0b37f 100644 --- a/transformer_engine/pytorch/attention/__init__.py +++ b/transformer_engine/pytorch/attention/__init__.py @@ -5,12 +5,14 @@ """Python interface for attention""" from .dot_product_attention import DotProductAttention +from .fused_mla_q_uproj import FusedMLAQUpProjRopeQuant from .multi_head_attention import MultiheadAttention from .inference import InferenceParams from .rope import RotaryPositionEmbedding __all__ = [ "DotProductAttention", + "FusedMLAQUpProjRopeQuant", "MultiheadAttention", "InferenceParams", "RotaryPositionEmbedding", diff --git a/transformer_engine/pytorch/attention/custom_ops.py b/transformer_engine/pytorch/attention/custom_ops.py new file mode 100644 index 0000000000..f364633607 --- /dev/null +++ b/transformer_engine/pytorch/attention/custom_ops.py @@ -0,0 +1,114 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Attention kernels wrapped as custom ops, so they don't graph-break under torch.compile.""" + +import torch +import transformer_engine_torch as tex + +from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat + +# The ops take the format's value rather than the pybind enum: converting the +# enum inside a traced region makes dynamo recurse until it gives up. +QKV_FORMAT_VALUE = {name: int(fmt) for name, fmt in QKVFormat.items()} +_QKV_FORMAT_BY_VALUE = {int(fmt): fmt for fmt in QKVFormat.values()} + + +@torch.library.custom_op( + "te_kv_cache::copy_to_kv_cache", + mutates_args=("k_cache", "v_cache"), + device_types="cuda", +) +def copy_to_kv_cache( + new_k: torch.Tensor, + new_v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cu_new_lens: torch.Tensor, + cu_cached_lens: torch.Tensor, + qkv_format: int, + b: int, + max_ctx_len: int, + max_seq_len: int, + max_pages_per_seq: int, + is_non_paged: bool, +) -> None: + """Copy new key/value tokens into the KV cache.""" + tex.copy_to_kv_cache( + new_k, + new_v, + k_cache, + v_cache, + page_table, + cu_new_lens, + cu_cached_lens, + _QKV_FORMAT_BY_VALUE[qkv_format], + b, + max_ctx_len, + max_seq_len, + max_pages_per_seq, + is_non_paged, + ) + + +@copy_to_kv_cache.register_fake +def _copy_to_kv_cache_fake(*_args, **_kwargs) -> None: + return None + + +@torch.library.custom_op("te_kv_cache::convert_bshd_to_thd", mutates_args=(), device_types="cuda") +def convert_bshd_to_thd(tensor: torch.Tensor, cu_seqlens: torch.Tensor, t: int) -> torch.Tensor: + """Convert a tensor from bshd to thd.""" + return tex.convert_bshd_to_thd(tensor, cu_seqlens, t) + + +@convert_bshd_to_thd.register_fake +def _convert_bshd_to_thd_fake( + tensor: torch.Tensor, cu_seqlens: torch.Tensor, t: int +) -> torch.Tensor: + del cu_seqlens + return tensor.new_empty((t, *tensor.shape[2:])) + + +@torch.library.custom_op("te_kv_cache::convert_thd_to_bshd", mutates_args=(), device_types="cuda") +def convert_thd_to_bshd( + tensor: torch.Tensor, cu_seqlens: torch.Tensor, b: int, max_seq_len: int +) -> torch.Tensor: + """Convert a tensor from thd to bshd.""" + return tex.convert_thd_to_bshd(tensor, cu_seqlens, b, max_seq_len) + + +@convert_thd_to_bshd.register_fake +def _convert_thd_to_bshd_fake( + tensor: torch.Tensor, cu_seqlens: torch.Tensor, b: int, max_seq_len: int +) -> torch.Tensor: + del cu_seqlens + return tensor.new_empty((b, max_seq_len, *tensor.shape[1:])) + + +@torch.library.custom_op("te_attention::fa_prepare_fwd", mutates_args=(), device_types="cuda") +def fa_prepare_fwd(qkvi: torch.Tensor) -> torch.Tensor: + """Split interleaved sbh3d QKV into bshd q/k/v.""" + return tex.fa_prepare_fwd(qkvi) + + +@fa_prepare_fwd.register_fake +def _fa_prepare_fwd_fake(qkvi: torch.Tensor) -> torch.Tensor: + # qkvi is the q view into the packed buffer, and its strides cover all of it. + s, b, n, h = qkvi.shape + return qkvi.new_empty((3, b, s, n, h)) + + +@torch.library.custom_op("te_attention::fa_prepare_bwd", mutates_args=(), device_types="cuda") +def fa_prepare_bwd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Pack bshd gradients back into an interleaved sbh3d buffer.""" + return tex.fa_prepare_bwd(q, k, v) + + +@fa_prepare_bwd.register_fake +def _fa_prepare_bwd_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + del k, v + b, s, n, h = q.shape + return q.new_empty((s, b, n, 3 * h)) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 891a5c661d..86c1f9cbe1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -15,7 +15,6 @@ import torch import torch.nn.functional as F -import transformer_engine_torch as tex from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, @@ -50,6 +49,12 @@ ) from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager from transformer_engine.pytorch.distributed import get_distributed_world_size +from transformer_engine.pytorch.attention.custom_ops import ( + convert_bshd_to_thd, + convert_thd_to_bshd, + fa_prepare_bwd, + fa_prepare_fwd, +) from transformer_engine.pytorch.jit import no_torch_dynamo from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( attn_forward_func_with_cp, @@ -167,14 +172,39 @@ flash_attn_func_v4 = None flash_attn_varlen_func_v4 = None else: - from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module - flash_attn_func as flash_attn_func_v4, - flash_attn_varlen_func as flash_attn_varlen_func_v4, - _validate_head_dims as _fa4_validate_head_dims, - ) + try: + cutlass_dsl_version = PkgVersion(get_pkg_version("nvidia-cutlass-dsl")) + + # FA4 4.0.0b24 requires CUTLASS DSL 4.6.2 or newer. + if fa_utils.fa4_version == PkgVersion("4.0.0b24") and cutlass_dsl_version < PkgVersion( + "4.6.2" + ): + raise ImportError( + "flash-attn-4 4.0.0b24 requires nvidia-cutlass-dsl>=4.6.2; " + f"found {cutlass_dsl_version}" + ) + + from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module + flash_attn_func as _flash_attn_func_v4, + flash_attn_varlen_func as _flash_attn_varlen_func_v4, + _validate_head_dims as _fa4_validate_head_dims, + ) + except ImportError as exc: + flash_attn_func_v4 = None + flash_attn_varlen_func_v4 = None + warnings.warn( + f"FlashAttention 4 is installed but cannot be loaded: {exc}", + RuntimeWarning, + stacklevel=2, + ) + else: + # Unlike versions 2 and 3, FlashAttention 4 registers no custom ops: it builds + # its kernels through the CUTLASS DSL as it runs. Keep it an eager island. + flash_attn_func_v4 = no_torch_dynamo()(_flash_attn_func_v4) + flash_attn_varlen_func_v4 = no_torch_dynamo()(_flash_attn_varlen_func_v4) - fa_utils.v4_validate_head_dims = _fa4_validate_head_dims - fa_utils.set_flash_attention_4_params() + fa_utils.v4_validate_head_dims = _fa4_validate_head_dims + fa_utils.set_flash_attention_4_params() # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" @@ -437,7 +467,7 @@ def _forward( if qkv_format == "thd_2bshd": batch_size = key_layer.shape[0] - query_layer = tex.convert_thd_to_bshd( + query_layer = convert_thd_to_bshd( query_layer, cu_seqlens_q, batch_size, @@ -773,7 +803,7 @@ def forward( # All inputs received are non-contiguous tensors. # The `query_layer` tensor is used to access the # full memory region of the QKV tensor. - qkv = tex.fa_prepare_fwd(query_layer) + qkv = fa_prepare_fwd(query_layer) q, k, v = split_tensor_along_dim(qkv, 0, 3) query_layer = torch.squeeze(q, 0) key_layer = torch.squeeze(k, 0) @@ -788,11 +818,23 @@ def backward( dv: torch.Tensor, ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring - dqkv = tex.fa_prepare_bwd(dq, dk, dv) + dqkv = fa_prepare_bwd(dq, dk, dv) dq, dk, dv = split_tensor_along_dim(dqkv, -1, 3) return dq, dk, dv +def _unalias_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor): + """Copy kv's cumulative sequence lengths when they are q's as well. + + Self attention passes one tensor for both, and flash-attn forwards them to + two inputs of the same autograd.Function -- which dynamo cannot trace. In + eager the duplicate is harmless, so nothing is copied there. + """ + if cu_q is cu_kv and torch.compiler.is_compiling(): + return cu_q, cu_kv.clone() + return cu_q, cu_kv + + class FlashAttention(torch.nn.Module): """Dot product attention, using HazyResearch flash-attn package: https://github.com/Dao-AILab/flash-attention @@ -1017,7 +1059,7 @@ def forward( # convert from bshd to thd_2bshd for flash_attn_varlen_func/_with_kvcache; # kernel assumes tensor is contiguous if isinstance(query_layer, Float8Tensor): - query_layer._data = tex.convert_bshd_to_thd( + query_layer._data = convert_bshd_to_thd( query_layer._data, cu_seqlens_q, batch_size * context_len, @@ -1026,7 +1068,7 @@ def forward( query_layer, data=query_layer._data, shape=query_layer._data.shape ) else: - query_layer = tex.convert_bshd_to_thd( + query_layer = convert_bshd_to_thd( query_layer, cu_seqlens_q, batch_size * context_len, @@ -1134,12 +1176,12 @@ def forward( else: func = flash_attn_with_kvcache_v3 # pylint: disable=possibly-used-before-assignment if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None): - fa_optional_forward_args_thd.append( - cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q - ) - fa_optional_forward_args_thd.append( - cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv + cu_q, cu_kv = _unalias_cu_seqlens( + cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q, + cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv, ) + fa_optional_forward_args_thd.append(cu_q) + fa_optional_forward_args_thd.append(cu_kv) fa_optional_forward_args_thd.append(max_seqlen_q) fa_optional_forward_args_thd.append(max_seqlen_kv) if use_flash_attn_4: @@ -1150,8 +1192,9 @@ def forward( if inference_params is None: fa_4_optional_forward_kwargs["deterministic"] = self.deterministic if func is flash_attn_varlen_func_v4: - fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q - fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_seqlens_kv + cu_q, cu_kv = _unalias_cu_seqlens(cu_seqlens_q, cu_seqlens_kv) + fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q + fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv output = func( @@ -1286,7 +1329,7 @@ def convert_to_torch_float8(tensor, dtype): # all KV caching cases use thd_2bshd for calculation # convert results back to bshd from thd_2bshd if isinstance(query_layer, Float8Tensor): - output._data = tex.convert_thd_to_bshd( + output._data = convert_thd_to_bshd( output._data, cu_seqlens_q, batch_size, @@ -1294,7 +1337,7 @@ def convert_to_torch_float8(tensor, dtype): ) output = Float8Tensor.make_like(output, data=output._data, shape=output._data.shape) else: - output = tex.convert_thd_to_bshd( + output = convert_thd_to_bshd( output, cu_seqlens_q, batch_size, @@ -1363,6 +1406,7 @@ def forward( deterministic, softmax_offset, fp8_output, + bf16_backward, layer_number, return_max_logit, packed_qkv=None, @@ -1427,6 +1471,9 @@ def forward( # fp8_dtype = tex.DType.kFloat8E4M3 if is_input_fp8: q_fp8, k_fp8, v_fp8 = q, k, v + + if fp8_recipe.mxfp8(): + qkv_scale_inv_format = "bhsd" # Same as what combine_and_quantize would give else: q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( qkv_layout, @@ -1602,6 +1649,8 @@ def forward( ctx.is_input_fp8 = is_input_fp8 ctx.is_output_fp8 = is_output_fp8 + # Return dQ/dK/dV in bf16 even if is_input_fp8 + ctx.bf16_backward = bf16_backward tensors_to_save, tensor_objects = prepare_for_saving( *fp8_tensors, @@ -1860,7 +1909,8 @@ def backward(ctx, d_out, *_args): # dq, dk, dv: torch.Tensor; dtype = torch.float16 or torch.bfloat16 dq, dk, dv = dq_, dk_, dv_ is_quantized_tensor = isinstance(dq_, QuantizedTensorStorage) - if is_quantized_tensor and not ctx.is_input_fp8: + + if is_quantized_tensor and (not ctx.is_input_fp8 or ctx.bf16_backward): # return in F16 dq, dk, dv = combine_and_dequantize( ctx.dqkv_layout, @@ -1869,7 +1919,7 @@ def backward(ctx, d_out, *_args): dv_, src_nominal_dtype=dq_.dtype, ) - if not is_quantized_tensor and ctx.is_input_fp8: + if not is_quantized_tensor and ctx.is_input_fp8 and not ctx.bf16_backward: # return in FP8 dq, dk, dv, _, _ = combine_and_quantize( ctx.dqkv_layout, dq_, dk_, dv_, ctx.dQKV_quantizer @@ -1968,6 +2018,7 @@ def backward(ctx, d_out, *_args): None, None, # packed_qkv None, # packed_kv + None, ) @@ -2064,6 +2115,7 @@ def forward( score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, packed_qkv: Optional[torch.Tensor] = None, packed_kv: Optional[torch.Tensor] = None, + bf16_backward: bool = False, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2280,6 +2332,7 @@ def forward( self.deterministic, softmax_offset, fp8_output, + bf16_backward, self.layer_number, self.return_max_logit, packed_qkv, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index b444f034c8..ea89ca97eb 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3087,10 +3087,12 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + or use_flash_attn_3 or fa_utils.v2_3_plus ), ( "cp_comm_type='all_gather' only supports SWA through FusedAttention or FlashAttention" - f" >= 2.3. Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + f" >= 2.3. Found {use_fused_attention=}, {use_flash_attn_3=}, " + f"and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" @@ -4231,10 +4233,11 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + or use_flash_attn_3 or fa_utils.v2_3_plus ), ( "cp_comm_type='a2a' only supports SWA through FusedAttention or FlashAttention >= 2.3." - f" Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + f" Found {use_fused_attention=}, {use_flash_attn_3=}, and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( "cp_comm_type='a2a' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" 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 11d6500ab7..d5adbbcadf 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 @@ -12,6 +12,7 @@ import torch import torch.nn.functional as F +from torch.fx.experimental.symbolic_shapes import guard_scalar from torch.nn.parameter import Parameter from transformer_engine.common.recipe import ( @@ -32,6 +33,7 @@ Float8BlockScalingRecipeState, ) from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage +from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage 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, DType @@ -265,6 +267,46 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1) +def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]: + """Why this DotProductAttention call has to run outside the graph, or None. + + `call` maps `DotProductAttention.forward`'s parameter names to the arguments + this call passed, including `self`. + """ + # FP8 GEMMs with the attention itself in high precision -- the common + # training setup -- stay on the compiled path; only FP8 attention bails out. + qstate = FP8GlobalStateManager.quantization_state + fp8_recipe = qstate.fp8_recipe + if qstate.fp8_enabled and fp8_recipe is not None: + if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha: + return "FP8 attention" + + if call["self"].cp_group is not None: + return "context parallelism" + + if call.get("checkpoint_core_attention", False): + return "activation checkpointing of the attention" + + qkv_format = call.get("qkv_format") or call["self"].qkv_format + if qkv_format == "thd" and ( + call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None + ): + # Deriving it reads the sequence lengths off cu_seqlens, which is a + # device synchronization and a data-dependent value while tracing. + return "deriving max_seqlen from cu_seqlens" + + if call.get("qkv_layer") is None and call.get("kv_layer") is None: + qkv = [call.get(name) for name in ("query_layer", "key_layer", "value_layer")] + if dpa_utils.qkv_layout_needs_detection(*qkv): + return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer" + + if call["self"].rng_states_tracker is not None and call["self"].attention_dropout > 0: + # Forking the tracker swaps the global CUDA generator state, which Dynamo + # refuses to trace. With no dropout there is nothing to fork for. + return "attention dropout under a CUDA RNG states tracker" + return None + + def _unpack_packed_qkv( qkv_layer: Optional[torch.Tensor], kv_layer: Optional[torch.Tensor], @@ -706,7 +748,10 @@ def __init__( else: self.rng_states_tracker = get_rng_state_tracker() set_all_rng_states(self.rng_states_tracker.get_states()) - attention_dropout_ctx = self.rng_states_tracker.fork + # Forking only matters if the dropout actually draws from the generator. + attention_dropout_ctx = ( + self.rng_states_tracker.fork if attention_dropout > 0 else nullcontext + ) if softmax_scale is None: softmax_scale = 1.0 / math.sqrt( @@ -876,6 +921,15 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: """ _original_recipe = self.fp8_meta.get("recipe", None) + # get_fp8_recipe() may build a default Recipe, which asserts it is not tracing. + # With quantization off the base class does all that is needed. + qstate = FP8GlobalStateManager.quantization_state + if torch.compiler.is_compiling() and not ( + qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters + ): + super().init_fp8_metadata(num_gemms=num_gemms) + return + # global recipe set in autocast() fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): @@ -1322,7 +1376,7 @@ def get_quantizer_roles( ] return base[:num_quantizers] - @no_torch_dynamo(recursive=False) + @no_torch_dynamo(when=_needs_eager_dpa) def forward( self, query_layer: Optional[torch.Tensor] = None, @@ -1347,6 +1401,7 @@ def forward( inference_params: Optional[InferenceParams] = None, pad_between_seqs: Optional[bool] = None, fp8_output: Optional[bool] = False, + bf16_backward: Optional[bool] = False, num_splits: Optional[int] = 1, score_mod: Optional[Callable] = None, score_mod_bprop: Optional[Callable] = None, @@ -1719,6 +1774,10 @@ def forward( batch_size = query_layer.shape[0] max_seqlen_q = query_layer.shape[1] if max_seqlen_q is None else max_seqlen_q max_seqlen_kv = key_layer.shape[1] if max_seqlen_kv is None else max_seqlen_kv + # Backend selection bakes these in, which it cannot do symbolically. + if not is_in_onnx_export_mode(): + max_seqlen_q = guard_scalar(max_seqlen_q) + max_seqlen_kv = guard_scalar(max_seqlen_kv) if qkv_format == "thd": assert all( len(x.shape) == 3 for x in (query_layer, key_layer, value_layer) @@ -1817,6 +1876,25 @@ def forward( qkv_format=qkv_format, inference_params=inference_params, ) + elif all( + isinstance(x, MXFP8TensorStorage) for x in [query_layer, key_layer, value_layer] + ): + # Pre-quantized MXFP8 q/k/v: the wrapper has no real storage, so run + # layout detection on the underlying rowwise data (mirrors the Float8 path). + ( + qkv_layout, + query_layer._rowwise_data, + key_layer._rowwise_data, + value_layer._rowwise_data, + q_format, + kv_format, + ) = dpa_utils.get_qkv_layout( + query_layer._rowwise_data, + key_layer._rowwise_data, + value_layer._rowwise_data, + qkv_format=qkv_format, + inference_params=inference_params, + ) else: ( qkv_layout, @@ -2052,18 +2130,25 @@ def forward( _attention_backends["fused_attention_backend"] = fused_attention_backend _attention_backends["use_unfused_attention"] = use_unfused_attention _attention_backends["backend_selection_requires_update"] = False + # logging.Logger methods graph-break under torch.compile, so + # selection is only logged in eager -- as in + # get_attention_backend. Note the arguments below are + # evaluated either way, so they have to stay traceable. + logger = ( + dpa_utils.no_op_logger if torch.compiler.is_compiling() else self.logger + ) if use_flash_attention: - self.logger.info( + logger.info( "Running with FlashAttention backend (version %s)", flash_attention_backend, ) elif use_fused_attention: - self.logger.info( + logger.info( "Running with FusedAttention backend (sub-backend %s)", int(fused_attention_backend), ) elif use_unfused_attention: - self.logger.info("Running with UnfusedDotProductAttention backend") + logger.info("Running with UnfusedDotProductAttention backend") else: use_flash_attention = _attention_backends["use_flash_attention"] flash_attention_backend = _attention_backends["flash_attention_backend"] @@ -2190,6 +2275,7 @@ def forward( fp8_output=fp8_output, packed_qkv=qkv_layer, packed_kv=kv_layer, + bf16_backward=bf16_backward, ) return self.fused_attention( query_layer, @@ -2227,6 +2313,7 @@ def forward( score_mod_bprop_tensors=score_mod_bprop_tensors, packed_qkv=qkv_layer, packed_kv=kv_layer, + bf16_backward=bf16_backward, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index d99ee4c8e8..b5e5633e14 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -346,7 +346,7 @@ def error(self, *args, **kwargs): """No-op.""" -_no_op_logger = _NoOpLogger() +no_op_logger = _NoOpLogger() @torch.compiler.assume_constant_result @@ -361,11 +361,17 @@ def _get_fused_attn_backend( *args, ): """Constant-foldable tex.get_fused_attn_backend: the result depends only on - the attention config, and the python-side enum keeps it traceable by - torch.compile (see the FusedAttnBackend docstring). Layout/bias/mask/softmax - are taken as their string keys and resolved to the pybind enums here, so - that every argument is a python literal or a python enum.""" - return FusedAttnBackend.cast( + the attention config. Layout/bias/mask/softmax are taken as their string + keys and resolved to the pybind enums here, so that every argument is a + python literal or a python enum. + + Returns a plain int rather than a FusedAttnBackend member: dynamo + reconstructs the result of an assume_constant_result call by re-emitting the + call, which is only valid inside the frame that made it. An int survives a + graph break because it is baked into the graph as a literal, while an enum + member comes out of the reconstruction corrupted (see the cast at the call + site, which restores the enum).""" + return int( tex.get_fused_attn_backend( is_training, q_type, @@ -395,8 +401,11 @@ def get_attention_backend( Whether the `FlashAttention` backend has been selected. use_fused_attention : bool Whether the `FusedAttention` backend has been selected. - fused_attention_backend : FusedAttnBackend - If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`. + fused_attention_backend : int + If `use_fused_attention = True`, the integer value of one of `FusedAttention`'s three + sub-backends, else `None`. It is not a `FusedAttnBackend` member because that does not + survive a graph break under `torch.compile`; `FusedAttnBackend.cast` turns it into one, + and comparing it against a member works either way. use_unfused_attention : bool Whether the `UnfusedDotProductAttention` backend has been selected. available_backends : List[bool] @@ -452,7 +461,7 @@ def get_attention_backend( if torch.compiler.is_compiling(): # logging.Logger methods graph-break under torch.compile; backend # selection logs are only emitted in eager mode. - logger = _no_op_logger + logger = no_op_logger else: logger = logging.getLogger("DotProductAttention") logger.setLevel(AttentionLogging._log_level) @@ -563,10 +572,10 @@ def _disable_all_flash_attention() -> None: if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: logger.debug("Disabling FlashAttention 3 for compute capability != sm90") use_flash_attention_3 = False - # FA4 supports SM80, SM90, SM100, SM120 - if device_compute_capability < (8, 0): + # FA4 does not currently support SM8x. + if device_compute_capability < (9, 0): if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 for compute capability < sm80") + logger.debug("Disabling FlashAttention 4 for compute capability < sm90") use_flash_attention_4 = False # On SM90, prefer FA3 over FA4 when FA3 is available. # FA3 is more mature on Hopper; FA4's SM90 backward has limitations @@ -996,6 +1005,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt device_compute_capability[0] * 10 + device_compute_capability[1], ) use_flash_attention_4 = False + # FA4's validator currently accepts symmetric (512, 512) on SM100/SM110, + # but the generic forward kernel exceeds its TMEM allocation for that shape. + # Preserve the supported asymmetric (64, 512) MLA path while D512 support + # is completed upstream. + if ( + use_flash_attention_4 + and (10, 0) <= device_compute_capability < (12, 0) + and head_dim_qk == head_dim_v == 512 + ): + logger.debug( + "Disabling FlashAttention 4 for unsupported symmetric head_dim=512 on SM100/SM110." + ) + use_flash_attention_4 = False # flash-attn-4 4.0.0b11 validates (256, 256) on SM100, but its dedicated # hd256 kernel diverges from the reference for cross-attention/decode-like # shapes such as sq=1, skv=2048. Keep FA4 enabled for the self-attention @@ -1460,11 +1482,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt cuda_graph, deterministic, ) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: + if fused_attention_backend == FusedAttnBackend.No_Backend.value: 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"]: + elif ( + has_score_mod and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value + ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " "F16/BF16 arbitrary-seqlen", @@ -1514,7 +1538,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend["FP8"] + fused_attention_backend == FusedAttnBackend.FP8.value and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1525,7 +1549,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] + fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value and is_training and ( device_compute_capability < (9, 0) @@ -1992,21 +2016,21 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor: tensor of shape [batch_size * max_seqlen, 1, 1] containing the indices for the valid tokens in a batch. """ + # Built with device-side ops only: reading the sequence lengths on the host + # would synchronize the device once per sequence. bs = len(cu_seqlens) - 1 seqlens = cu_seqlens[1:] - cu_seqlens[:-1] - indices = [i * max_seqlen + ii for i, j in enumerate(seqlens) for ii in range(j)] - indices = torch.Tensor(indices).unsqueeze(1).unsqueeze(1).to(dtype=torch.int64, device="cuda") - - num_nonzeros = indices.shape[0] - pad_amount = bs * max_seqlen - num_nonzeros - indices = F.pad( - input=indices, - pad=(0, 0, 0, 0, 0, pad_amount), - mode="constant", - value=float(bs * max_seqlen), + positions = torch.arange(max_seqlen, device=cu_seqlens.device) + valid = (positions.unsqueeze(0) < seqlens.unsqueeze(1)).flatten() + # Sorting the mask is stable, so the valid positions come first in order, + # and the invalid tail is replaced by the out-of-range padding index. + ordered = torch.argsort(~valid, stable=True) + indices = torch.where( + torch.arange(bs * max_seqlen, device=cu_seqlens.device) < valid.sum(), + ordered, + bs * max_seqlen, ) - - return indices + return indices.to(dtype=torch.int64, device="cuda").unsqueeze(1).unsqueeze(1) def get_full_cu_seqlens( @@ -2031,6 +2055,11 @@ def _get_cu_seqlens(batch_size, max_seqlen, device): ) if is_in_onnx_export_mode(): + # A tensor cached by an earlier call would be baked into the exported graph. + return _get_cu_seqlens(batch_size, max_seqlen, device) + if torch.compiler.is_compiling(): + # torch.is_inference_mode_enabled(), part of the cache key, graph-breaks, and the + # cache only saves one arange. return _get_cu_seqlens(batch_size, max_seqlen, device) is_inference = torch.is_inference_mode_enabled() @@ -2370,6 +2399,20 @@ def get_qkv_format( return qkv_format, q_format, kv_format +def qkv_layout_needs_detection(*qkv: Optional[torch.Tensor]) -> bool: + """Whether the layout of these q/k/v can only be told by inspecting memory. + + True for tensors that may be slices of a packed buffer, i.e. strided ones + that do not own their whole storage. + """ + return any( + x is not None + and not x.is_contiguous() + and x.untyped_storage().size() != x.numel() * x.element_size() + for x in qkv + ) + + def get_qkv_layout( q: torch.Tensor, k: torch.Tensor, @@ -2543,10 +2586,25 @@ def run_iteratively(q, k, v): return qkv_layout - if not is_in_onnx_export_mode(): - qkv_layout = run_iteratively(q, k, v) - else: + if is_in_onnx_export_mode(): + # Checked first: the ONNX exporter runs through dynamo, so it also sets + # is_compiling(), and it has its own handling below. qkv_layout = "not_supported" + elif torch.compiler.is_compiling(): + # run_iteratively reads data pointers and storage offsets, which dynamo + # cannot trace; unpacked q/k/v need no detection anyway. + assert not qkv_layout_needs_detection(q, k, v), ( + "q/k/v may be views into a packed buffer, whose layout cannot be detected under" + " torch.compile. Pass the packed buffer explicitly via DotProductAttention's" + " qkv_layer/kv_layer (with qkv_interleave_dim)." + ) + q, k, v = [x if x.is_contiguous() else x.contiguous() for x in (q, k, v)] + if is_same_q_kv_format: + qkv_layout = "_".join([qkv_format] * 3) + else: + qkv_layout = q_format + "_" + kv_format + "_" + kv_format + else: + qkv_layout = run_iteratively(q, k, v) if qkv_layout == "not_supported": # force q,k,v to be contiguous and run get_layout again q, k, v = [x.contiguous() for x in [q, k, v]] @@ -2786,10 +2844,37 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): """ if not tensor_quantizer_pairs: return [], src_format + + fp8_tensors = mxfp8_quantize_only(tensor_quantizer_pairs, src_format) + mxfp8_transpose_swizzle(fp8_tensors, src_format) + return fp8_tensors, "bhsd" + + +def mxfp8_quantize_only(tensor_quantizer_pairs, src_format): + """Phase 1 of mxfp8_quantize_fast_path: quantize only, no BHSD transpose or GEMM swizzle. + + Returns MXFP8Tensors with data and scale_invs reshaped to src_format layout. + Call mxfp8_transpose_swizzle to complete the BHSD permute + swizzle when ready + (e.g. after pre-quantized tensors from fused kernels are also available). + + Parameters + ---------- + tensor_quantizer_pairs : list of (torch.Tensor, MXFP8Quantizer) + Same contract as mxfp8_quantize_fast_path. + src_format : str + ``"bshd"`` or ``"sbhd"``. + + Returns + ------- + fp8_tensors : list of MXFP8Tensor + Data and scale_invs in src_format layout; NOT yet BHSD-permuted or swizzled. + """ + if not tensor_quantizer_pairs: + return [] assert src_format in ( "bshd", "sbhd", - ), f"mxfp8_quantize_fast_path only supports bshd/sbhd, got {src_format!r}." + ), f"mxfp8_quantize_only only supports bshd/sbhd, got {src_format!r}." _s_dim = {"bshd": 1, "sbhd": 0} _d_dim = {"bshd": 3, "sbhd": 3} @@ -2800,45 +2885,74 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): rs_shape[_d_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE cs_shape = list(original_shape) cs_shape[_s_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE - - # view tensor as 2D for quantization - # BSHD -> (B*S, H*D) - # SBHD -> (S, B*H*D) if src_format == "bshd": - tensor = tensor.view(*tensor.shape[:2], -1) + t2d = tensor.view(*tensor.shape[:2], -1) else: - tensor = tensor.view(tensor.shape[0], -1) - - # quantize + t2d = tensor.view(tensor.shape[0], -1) orig_optimize = quantizer.optimize_for_gemm quantizer.optimize_for_gemm = False - fp8_tensor = quantizer(tensor) + fp8_2d = quantizer(t2d) quantizer.optimize_for_gemm = orig_optimize - - # reshape rowwise/columnwise data to original shape - fp8_tensor._rowwise_data = ( - fp8_tensor._rowwise_data.view(original_shape) - if fp8_tensor._rowwise_data is not None - else None - ) - fp8_tensor._columnwise_data = ( - fp8_tensor._columnwise_data.view(original_shape) - if fp8_tensor._columnwise_data is not None - else None - ) - fp8_tensor._rowwise_scale_inv = ( - fp8_tensor._rowwise_scale_inv.view(rs_shape) - if fp8_tensor._rowwise_scale_inv is not None - else None - ) - fp8_tensor._columnwise_scale_inv = ( - fp8_tensor._columnwise_scale_inv.view(cs_shape) - if fp8_tensor._columnwise_scale_inv is not None - else None + # Re-wrap with the original 4D SBHD/BSHD shape so that shape[-1] equals the per-head + # dimension (matching Q's wrapper shape) and fused_attn_bwd produces 4D dkv that + # matches key/value's expected gradient shape in _KFQuantizeKVForAttn.backward. + fp8_t = MXFP8Tensor( + shape=original_shape, + dtype=tensor.dtype, + rowwise_data=( + fp8_2d._rowwise_data.view(original_shape) + if fp8_2d._rowwise_data is not None + else None + ), + rowwise_scale_inv=( + fp8_2d._rowwise_scale_inv.view(rs_shape) + if fp8_2d._rowwise_scale_inv is not None + else None + ), + columnwise_data=( + fp8_2d._columnwise_data.view(original_shape) + if fp8_2d._columnwise_data is not None + else None + ), + columnwise_scale_inv=( + fp8_2d._columnwise_scale_inv.view(cs_shape) + if fp8_2d._columnwise_scale_inv is not None + else None + ), + quantizer=quantizer, + requires_grad=False, + fp8_dtype=fp8_2d._fp8_dtype, + with_gemm_swizzled_scales=False, ) - fp8_tensors.append(fp8_tensor) + fp8_tensors.append(fp8_t) + return fp8_tensors + + +def mxfp8_transpose_swizzle(fp8_tensors, src_format): + """Phase 2 of mxfp8_quantize_fast_path: batched BHSD-transpose + GEMM-swizzle. + + For tensors whose data is already quantized (e.g. from a fused GEMM+quant kernel + or from mxfp8_quantize_only), permutes each tensor's scale_invs from src_format to + BHSD and applies the GEMM swizzle in-place. Complements mxfp8_quantize_only to + allow pre-quantized tensors (like a fused-kernel Q) to be processed in the same + batched operation as freshly quantized K/V. + + Parameters + ---------- + fp8_tensors : list of MXFP8Tensor + Tensors with _rowwise_scale_inv / _columnwise_scale_inv in src_format layout. + Modified in-place: scale_invs are replaced with BHSD-permuted, swizzled versions. + src_format : str + ``"bshd"`` or ``"sbhd"``. + """ + if not fp8_tensors: + return + + assert src_format in ( + "bshd", + "sbhd", + ), f"mxfp8_transpose_swizzle only supports bshd/sbhd, got {src_format!r}." - # ---- Pad + permute + swizzle scale_inv to BHSD ---- rs_list = [t._rowwise_scale_inv for t in fp8_tensors] cs_list = [t._columnwise_scale_inv for t in fp8_tensors] @@ -2872,50 +2986,25 @@ def _build_outputs(scale_list, alignment): buf = torch.empty(total, dtype=torch.uint8, device=device) return [buf[e[0] : e[0] + e[1]].view(e[2]) if e is not None else None for e in entries] - # allocate buffers with padding in mind rs_outs = _build_outputs(rs_list, 4) cs_outs = _build_outputs(cs_list, 128) - # permute scale_invs to BHSD; batched rs_permuted = tex.multi_tensor_transpose_to_bhsd( - rs_list, - original_format=src_format, - outputs=rs_outs, + rs_list, original_format=src_format, outputs=rs_outs ) cs_permuted = tex.multi_tensor_transpose_to_bhsd( - cs_list, - original_format=src_format, - outputs=cs_outs, + cs_list, original_format=src_format, outputs=cs_outs ) - # build output tensors - result = [] for t, rp, cp in zip(fp8_tensors, rs_permuted, cs_permuted): - rp = rp.view(-1, rp.shape[-1]) if rp is not None else None - cp = cp.view(-1, cp.shape[-1]) if cp is not None else None - result.append( - MXFP8Tensor( - shape=t.shape, - dtype=t.dtype, - rowwise_data=t._rowwise_data, - rowwise_scale_inv=rp, - columnwise_data=t._columnwise_data, - columnwise_scale_inv=cp, - quantizer=t._quantizer, - requires_grad=False, - fp8_dtype=t._fp8_dtype, - with_gemm_swizzled_scales=t._with_gemm_swizzled_scales, - ) - ) + t._rowwise_scale_inv = rp.view(-1, rp.shape[-1]) if rp is not None else None + t._columnwise_scale_inv = cp.view(-1, cp.shape[-1]) if cp is not None else None - # swizzle in place; batched - tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, True, False) - tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, False, True) - for t in result: + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, True, False) + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, False, True) + for t in fp8_tensors: t._with_gemm_swizzled_scales = True - return result, "bhsd" - def combine_and_quantize( qkv_layout, diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py new file mode 100644 index 0000000000..c176985254 --- /dev/null +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused MLA Q up-projection + per-head RoPE + MXFP8 quantize.""" + +from __future__ import annotations +import functools +import os +from importlib.metadata import PackageNotFoundError, version as get_pkg_version + +import torch +import transformer_engine_torch as tex +from packaging.version import Version as PkgVersion + +from ..constants import MXFP8_BLOCK_SCALING_SIZE +from ..quantized_tensor import QuantizedTensor +from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +from ..utils import get_device_compute_capability + +_CUDNN_FRONTEND_MIN_VERSION = "1.27.0" + + +def _cudnn_frontend_version_supported() -> bool: + """Check that the installed nvidia-cudnn-frontend meets the minimum version.""" + try: + return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion( + _CUDNN_FRONTEND_MIN_VERSION + ) + except PackageNotFoundError: + return False + + +class FusedMLAQUpProjRopeQuant: + """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel. + + - If w is already a QuantizedTensor (primary FP8 parameter in MXFP8BlockScaling recipe), + this performs an MXFP8 GEMM within the fusion (and quantizes the input if necessary) + - Otherwise (plain BF16 weight), x and w are passed as-is to the BF16 kernel variant. + """ + + @classmethod + @functools.lru_cache(maxsize=None) + def _kernel(cls): + # Import directly from the subpackage to avoid depending on cudnn/__init__.py + # lazy-import registration (which would require overlaying cudnn/__init__.py and + # could revert atomicrmw fixes present in the container's version). + try: + from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 + + return gemm_proj_rope_mxfp8_wrapper_sm100 + except ImportError: + return None + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + """Whether the cuDNN FE fused gemm rope quant wrapper is available""" + if int(os.environ.get("NVTE_FUSED_MLA_Q_UPROJ", "1")) <= 0: + return False + if not _cudnn_frontend_version_supported(): + return False + if get_device_compute_capability()[0] < 10: + return False + if cls._kernel() is None: + return False + return True + + @classmethod + def run( + cls, + x: torch.Tensor, + w, # MXFP8Tensor (primary FP8 param) or bf16 torch.Tensor + cos: torch.Tensor, + sin: torch.Tensor, + s: int, + b: int, + ) -> "tuple[MXFP8Tensor, torch.Tensor]": + """Run the fused kernel; return (Q MXFP8Tensor, activation saved for the wgrad backward). + + The kernel precision is selected by the weight precision. + """ + + from cuda.bindings import driver as cuda + + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + wrapper = cls._kernel() + + if isinstance(w, QuantizedTensor): + assert isinstance(w, MXFP8Tensor), ( + "FusedMLAQUpProjRopeQuant expects an MXFP8Tensor weight (MXFP8BlockScaling" + f" recipe), got {type(w).__name__}. Use the unfused path for other quantization" + " recipes." + ) + # ---- FP8 projection: MXFP8-cast x (both usages) + reuse w's fp8 codes -> mxfp8in ---- + # Quantize x with both rowwise (for the forward GEMM) and columnwise (for the FP8 + # wgrad in backward, matching the unfused path). + x_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ) + x_mxfp8 = x_quantizer(x) + x_code = x_mxfp8._rowwise_data.view(torch.float8_e4m3fn) # [tokens, K] + x_scale = x_mxfp8._rowwise_scale_inv # [tokens, K//32] uint8 + + # Primary FP8 parameter: already quantized; use its rowwise FP8 codes + E8M0 scales. + w.update_usage(rowwise_usage=True, columnwise_usage=None) + w_code = w._rowwise_data.view(torch.float8_e4m3fn) # [N, K] + w_scale = w._rowwise_scale_inv # [N, K//32] uint8 + + out = wrapper( + x_code, + w_code, + cos, + sin, + x_scale=x_scale, + w_scale=w_scale, + w_out_in=True, + stream=stream, + ) + + # Drop rowwise data now. + # Only columnwise x is needed for the FP8 wgrad in backward. + x_mxfp8.update_usage(rowwise_usage=False, columnwise_usage=True) + x_saved = x_mxfp8 + else: + # ---- 16-bit projection: bf16 GEMM inputs -> bf16in (the projection stays bf16) ---- + out = wrapper(x, w, cos, sin, w_out_in=True, stream=stream) + x_saved = x + + nh = out["out_fp8_row"].shape[1] + d = out["out_fp8_row"].shape[2] + query = cls.wrap_mxfp8( + out["out_fp8_row"], + out["out_scales_row"], + out["out_fp8_col"], + out["out_scales_col"], + s, + b, + nh, + d, + ) + # 2nd return is the activation to save for wgrad: MXFP8 (fp8 path) or bf16 (16-bit path). + return query, x_saved + + @classmethod + def wrap_mxfp8( + cls, + fp8_row: torch.Tensor, + scales_row: torch.Tensor, + fp8_col: torch.Tensor, + scales_col: torch.Tensor, + s: int, + b: int, + nh: int, + d: int, + ) -> MXFP8Tensor: + """Wrap raw data and scale tensors into an MXFP8Tensor""" + + blk = MXFP8_BLOCK_SCALING_SIZE + # Both rowwise and columnwise Q are required: + # - Forward QK^T uses rowwise + # - cuDNN backward (fused_attn_fp8_bwd_impl) requires columnwise for dK gradient + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + return MXFP8Tensor( + shape=(s, b, nh, d), + dtype=torch.bfloat16, + rowwise_data=fp8_row.view(s, b, nh, d), + rowwise_scale_inv=scales_row.view(s, b, nh, d // blk), + columnwise_data=fp8_col.view(s, b, nh, d), + columnwise_scale_inv=scales_col.view(s // blk, b, nh, d), + quantizer=quantizer, + requires_grad=False, + fp8_dtype=tex.DType.kFloat8E4M3, + with_gemm_swizzled_scales=False, + ) diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py index 08e50aad8b..f98d1693a6 100644 --- a/transformer_engine/pytorch/attention/inference.py +++ b/transformer_engine/pytorch/attention/inference.py @@ -10,8 +10,10 @@ import torch -import transformer_engine_torch as tex -from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat +from transformer_engine.pytorch.attention.custom_ops import ( + copy_to_kv_cache, + QKV_FORMAT_VALUE, +) __all__ = ["InferenceParams", "KVCacheManager", "NonPagedKVCacheManager", "PagedKVCacheManager"] @@ -459,6 +461,9 @@ def allocate_memory(self, layer_number): dtype=self.dtype, device=torch.cuda.current_device(), ) + # Mutated in place, so CUDA graphs are only captured if the addresses are fixed. + torch._dynamo.mark_static_address(k_cache) + torch._dynamo.mark_static_address(v_cache) self.cache[layer_number] = (k_cache, v_cache) def pre_step( @@ -549,7 +554,7 @@ def step( batch_size = new_k.shape[1] ctx_len = new_k.shape[0] - tex.copy_to_kv_cache( + copy_to_kv_cache( new_k, new_v, k_cache, @@ -557,7 +562,7 @@ def step( self.batch_indices, cu_new_seqlens, cu_cached_seqlens, - QKVFormat[qkv_format], + QKV_FORMAT_VALUE[qkv_format], batch_size, ctx_len, self.max_seqlen, @@ -656,6 +661,9 @@ def allocate_memory(self, layer_number): dtype=self.dtype, device=torch.cuda.current_device(), ) + # Mutated in place, so CUDA graphs are only captured if the addresses are fixed. + torch._dynamo.mark_static_address(k_cache) + torch._dynamo.mark_static_address(v_cache) self.cache[layer_number] = (k_cache, v_cache) def print_cache(self): @@ -807,7 +815,7 @@ def step( batch_size = new_k.shape[1] ctx_len = new_k.shape[0] - tex.copy_to_kv_cache( + copy_to_kv_cache( new_k, new_v, k_cache, @@ -815,7 +823,7 @@ def step( self.page_table, cu_new_seqlens, cu_cached_seqlens, - QKVFormat[qkv_format], + QKV_FORMAT_VALUE[qkv_format], batch_size, ctx_len, self.max_seqlen, diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 046019ee58..9a33df7634 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -107,10 +107,13 @@ class FusedAttnBackend(IntEnum): ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum value-for-value, and instances of the two enums compare equal when they share the same integer value. Unlike the pybind enum, a plain-python - ``IntEnum`` is traceable by ``torch.compile``: comparisons constant-fold - cleanly and instances safely cross the ``assume_constant_result`` boundary - in ``get_attention_backend``. Lookup by name (``FusedAttnBackend["FP8"]``) - works the same way as with the dict this used to be. + ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member + constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works + the same way as with the dict this used to be. + + Members do not survive a graph break, though, so ``get_attention_backend`` + returns the sub-backend as a plain int and ``cast`` turns it back into a + member. """ No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 95f02c64f0..9114b1e453 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -290,6 +290,31 @@ py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float l py::object clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer, float limit, float alpha, float glu_linear_offset); + +/* Scaled activation */ +py::object scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, int64_t glu_interleave_size); + +py::object scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size); + +py::object scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer); + +py::tuple scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + int64_t glu_interleave_size, bool compute_scale_grad); + +py::tuple scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, float limit, + float alpha, float glu_linear_offset, int64_t glu_interleave_size, + bool compute_scale_grad); + +py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + bool compute_scale_grad); + /*************************************************************************************************** * LayerNorm **************************************************************************************************/ @@ -350,7 +375,7 @@ 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 last_dims, std::optional tensor_offsets, - std::optional noop_flag); + std::optional noop_flag, const py::object &output); py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, @@ -367,6 +392,11 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, std::optional last_dims, std::optional tensor_offsets); +py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + DType otype, std::optional tensor_offsets, + bool return_dequantized); + std::vector multi_tensor_quantize(const std::vector &tensor_list, std::vector quantizer_list); @@ -691,14 +721,18 @@ void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_pe at::Tensor total_recv_tokens); void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, - at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights); + at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights, + std::optional tokens_scale_inv = std::nullopt, + std::optional recv_scale_inv = std::nullopt); void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result); void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_topk_weights, at::Tensor grad_tokens, at::Tensor grad_topk_weights); -void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out); +void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out, + std::optional grad_scale_inv = std::nullopt, + std::optional grad_expert_out_scale_inv = std::nullopt); // Registers the EP pybind functions on `m`. Defined under NVTE_WITH_NCCL_EP. void register_ep_bindings(pybind11::module_ &m); diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 58a8f84f85..544ff92c1b 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -342,5 +342,150 @@ py::object clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, py:: glu_linear_offset); } +/* Scaled activation helpers (activation + per-row scale via nvte_scaled_*). */ + +template +at::Tensor scaled_activation_compute(const at::Tensor& input, const at::Tensor& act_scales, + int shape_divisor, Args&&... args) { + init_extension(); + NVTE_CHECK(input.dim() >= 1, "scaled activation input must have at least 1 dimension"); + NVTE_CHECK(shape_divisor > 0 && input.size(-1) % shape_divisor == 0, + "scaled activation input width is not compatible with activation"); + + auto input_tensor = input.contiguous(); + auto scales_tensor = act_scales.contiguous().reshape({-1}); + const int64_t rows = input_tensor.numel() / input_tensor.size(-1); + NVTE_CHECK(scales_tensor.numel() == rows, "scaled activation expects one scale per input row"); + + std::vector output_sizes(input_tensor.sizes().begin(), input_tensor.sizes().end()); + output_sizes.back() /= shape_divisor; + auto output = at::empty(output_sizes, input_tensor.options()); + + const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor); + const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_tensor); + const TensorWrapper& output_nvte = makeTransformerEngineTensor(output); + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + act_func(input_nvte.data(), scales_nvte.data(), output_nvte.data(), std::forward(args)..., + stream); + }); + return output; +} + +template +std::tuple scaled_dactivation_compute(const at::Tensor& grad, + const at::Tensor& input, + const at::Tensor& act_scales, + bool compute_scale_grad, + Args&&... args) { + init_extension(); + NVTE_CHECK(input.dim() >= 1 && grad.dim() >= 1, + "scaled dactivation input and grad must have at least 1 dimension"); + + auto grad_tensor = grad.contiguous(); + auto input_tensor = input.contiguous(); + auto scales_tensor = act_scales.contiguous(); + const int64_t rows = input_tensor.numel() / input_tensor.size(-1); + NVTE_CHECK(scales_tensor.numel() == rows, "scaled dactivation expects one scale per input row"); + + auto scales_flat = scales_tensor.reshape({-1}); + auto grad_input = at::empty_like(input_tensor); + auto grad_scales = compute_scale_grad ? at::empty_like(scales_tensor) : at::Tensor(); + auto grad_scales_flat = compute_scale_grad ? grad_scales.reshape({-1}) : at::Tensor(); + + const TensorWrapper& grad_nvte = makeTransformerEngineTensor(grad_tensor); + const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor); + const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_flat); + const TensorWrapper& grad_input_nvte = makeTransformerEngineTensor(grad_input); + std::optional grad_scales_nvte; + if (compute_scale_grad) { + grad_scales_nvte.emplace(makeTransformerEngineTensor(grad_scales_flat)); + } + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + dact_func(grad_nvte.data(), input_nvte.data(), scales_nvte.data(), grad_input_nvte.data(), + compute_scale_grad ? grad_scales_nvte->data() : nullptr, std::forward(args)..., + stream); + }); + return {grad_input, grad_scales}; +} + +py::object maybe_quantize(const at::Tensor& tensor, py::handle quantizer) { + if (quantizer.is_none()) { + return py::cast(tensor); + } + auto quantizer_cpp = convert_quantizer(quantizer); + const TensorWrapper& tensor_nvte = makeTransformerEngineTensor(tensor); + const auto shape_te = tensor_nvte.shape(); + const std::vector shape(shape_te.data, shape_te.data + shape_te.ndim); + auto fake_dtype = GetTransformerEngineDType(tensor.scalar_type()); + auto [out_nvte, out_py] = quantizer_cpp->create_tensor(shape, fake_dtype); + quantizer_cpp->quantize(tensor_nvte, out_nvte); + return out_py; +} + +template +py::object scaled_activation_helper(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, int shape_divisor, Args&&... args) { + auto output = scaled_activation_compute(input, act_scales, shape_divisor, + std::forward(args)...); + return maybe_quantize(output, quantizer); +} + +template +py::tuple scaled_dactivation_helper(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + bool compute_scale_grad, Args&&... args) { + auto [grad_input, grad_scales] = scaled_dactivation_compute( + grad, input, act_scales, compute_scale_grad, std::forward(args)...); + return py::make_tuple(maybe_quantize(grad_input, quantizer), + compute_scale_grad ? py::cast(grad_scales) : py::none()); +} + +py::object scaled_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, int64_t glu_interleave_size) { + return scaled_activation_helper(input, act_scales, quantizer, + /*shape_divisor=*/2, glu_interleave_size); +} + +py::object scaled_clamped_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size) { + return scaled_activation_helper( + input, act_scales, quantizer, /*shape_divisor=*/2, limit, alpha, glu_linear_offset, + glu_interleave_size); +} + +py::object scaled_srelu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer) { + return scaled_activation_helper(input, act_scales, quantizer, + /*shape_divisor=*/1); +} + +py::tuple scaled_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + int64_t glu_interleave_size, bool compute_scale_grad) { + return scaled_dactivation_helper(grad, input, act_scales, quantizer, + compute_scale_grad, glu_interleave_size); +} + +py::tuple scaled_clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, float limit, + float alpha, float glu_linear_offset, int64_t glu_interleave_size, + bool compute_scale_grad) { + return scaled_dactivation_helper( + grad, input, act_scales, quantizer, compute_scale_grad, limit, alpha, glu_linear_offset, + glu_interleave_size); +} + +py::tuple scaled_dsrelu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + bool compute_scale_grad) { + return scaled_dactivation_helper(grad, input, act_scales, quantizer, + compute_scale_grad); +} + } // namespace pytorch } // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 882481e0f4..89c0418f01 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -277,7 +277,7 @@ void compute_grouped_fp8_current_scaling_amax_and_scale( py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, - std::optional noop_flag) { + std::optional noop_flag, const py::object &output) { using namespace transformer_engine::pytorch::detail; init_extension(); @@ -306,11 +306,34 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const GetTransformerEngineDType(tensor.scalar_type()), std::vector{static_cast(tensor.numel())}); - // 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, last_dims, tensor_offsets, - logical_first_dim, logical_last_dim); + // Create a GroupedTensor or reuse an existing destination. Reusing the destination is + // required for weight caching and CUDA graph replay because captured GEMMs retain the + // original data and scale pointers. + py::object grouped_output_py; + auto grouped_output_tensor_cpp = [&]() -> GroupedTensorWrapper { + if (!output.is_none()) { + NVTE_CHECK(!first_dims.has_value() && !last_dims.has_value() && !tensor_offsets.has_value(), + "group_quantize: output reuse currently requires uniform tensor shapes."); + NVTE_CHECK(output.attr("num_tensors").cast() == num_tensors, + "group_quantize: output has a different number of tensors."); + NVTE_CHECK(output.attr("logical_shape").cast>() == logical_shape, + "group_quantize: output has an incompatible logical shape."); + py::object output_quantizer = output.attr("quantizer"); + NVTE_CHECK(!output_quantizer.is_none(), + "group_quantize: output must have quantized storage."); + NVTE_CHECK(output_quantizer.is(quantizer), + "group_quantize: output must have been created by the same quantizer."); + grouped_output_py = output; + return GroupedTensorFromPyTorchGroupedTensor(output); + } + + auto result = quantizer_cpp->create_grouped_tensor( + num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, + logical_first_dim, logical_last_dim); + grouped_output_py = std::move(result.second); + return std::move(result.first); + }(); // dispatch to scaling methods enum class GroupedQuantizationMode { @@ -381,6 +404,9 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const QuantizationConfigWrapper quant_config_cpp; quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + if (noop_flag_cpp.has_value()) { + quant_config_cpp.set_noop_tensor(noop_flag_cpp->data()); + } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), quant_config_cpp, at::cuda::getCurrentCUDAStream()); @@ -664,6 +690,102 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o return py::reinterpret_borrow(out_py); } +py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + DType otype, std::optional tensor_offsets, + bool return_dequantized) { + init_extension(); + + const bool has_rowwise = + !grouped_x.attr("rowwise_data").is_none() && !grouped_x.attr("scale_inv").is_none(); + const bool has_columnwise = !grouped_x.attr("columnwise_data").is_none() && + !grouped_x.attr("columnwise_scale_inv").is_none(); + const bool swizzled = grouped_x.attr("_with_gemm_swizzled_scales").cast(); + + NVTE_CHECK(has_rowwise, "Grouped input has no rowwise data and scales for the GEMM to consume."); + + // The tensor's own quantization must match what the op expects on every path: even a + // pass-through hands its data straight to the GEMM. This keeps the input's format rather than + // converting between formats. + const auto input_quantizer = grouped_x.attr("quantizer"); + NVTE_CHECK(!input_quantizer.is_none(), "Grouped input has no quantizer."); + NVTE_CHECK(Py_TYPE(input_quantizer.ptr()) == Py_TYPE(quantizer.ptr()), + "Grouped input and the op disagree on quantization format."); + NVTE_CHECK(input_quantizer.attr("dtype").cast() == quantizer.attr("dtype").cast(), + "Grouped input and the quantizer disagree on the FP8 dtype."); + + // The columnwise copy is only worth building when a wgrad GEMM will consume it. Read this + // before the usage is overridden below. + const bool need_columnwise = quantizer.attr("columnwise_usage").cast(); + + if (swizzled) { + // Already GEMM-ready. Nothing can be derived from here, since dequantization requires scales + // in compact format, so the input must already carry everything that will be consumed. No + // quantization kernel runs, which makes this path format-agnostic. + NVTE_CHECK(has_columnwise || !need_columnwise, + "Grouped input has swizzled scales but no columnwise data for the wgrad GEMM. It " + "cannot be rebuilt, because dequantization requires scales in compact format."); + NVTE_CHECK(!return_dequantized, + "Cannot return a dequantized tensor for an already-swizzled grouped input: " + "dequantization requires scales in compact format."); + return py::none(); + } + + NVTE_CHECK(!has_columnwise, + "Grouped input already has columnwise data but unswizzled scales; requantizing " + "from it is not supported."); + + // Everything below runs quantization kernels, so it is MXFP8-only. + NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()), + "Requantizing a grouped input is only supported for MXFP8."); + + const auto logical_shape = grouped_x.attr("logical_shape").cast(); + const auto total_tokens = logical_shape[0].cast(); + const auto hidden_dim = logical_shape[1].cast(); + // Each group's token count must be a multiple of 128 too, so that every group's scales start + // on a swizzle-tile boundary. Those counts live on the device (host reads would break CUDA + // graph capture), so that half is the caller's contract rather than an assertion. + NVTE_CHECK(total_tokens % 128 == 0 && hidden_dim % 128 == 0, + "Requantizing a grouped input requires dims that are multiples of 128, but got (", + total_tokens, ", ", hidden_dim, ")."); + + // Dequantize first: it reads the rowwise scales, which the swizzle below replaces. Left + // undefined when nothing consumes it, which skips the pass entirely. + at::Tensor dequantized; + if (need_columnwise || return_dequantized) { + dequantized = group_dequantize(grouped_x, otype) + .attr("rowwise_data") + .cast() + .view({static_cast(total_tokens), static_cast(hidden_dim)}); + } + + // Swizzle the rowwise scales before attaching any columnwise data: a rowwise-only swizzle + // resets columnwise_scale_inv to None, which would strand the columnwise data below with a + // null scale pointer. + grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false); + // The swizzle hands back a 2D [num_tensors * padded_m, padded_k] scale buffer, but grouped + // tensors carry scales as a flat array indexed by element offsets (scale_inv_offsets), so + // per-group slicing breaks unless it is flattened back. + grouped_x.attr("scale_inv") = grouped_x.attr("scale_inv").attr("reshape")(-1); + + if (need_columnwise) { + // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise + // data because the two directions scale along perpendicular axes. Quantizing rowwise as well + // would redo work we already have, so that direction is switched off; the caller sets + // optimize_for_gemm, which makes the kernel emit swizzled columnwise scales directly. + quantizer.attr("set_usage")(py::arg("rowwise") = false, py::arg("columnwise") = true); + auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt, + tensor_offsets, std::nullopt, py::none()); + grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data"); + grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv"); + } + + if (return_dequantized) { + return py::cast(dequantized); + } + return py::none(); +} + namespace { void multi_tensor_quantize_impl(const std::vector &input_list, diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 9bf821721f..cef489dbbb 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -73,8 +73,14 @@ NVTECommWindow maybe_make_window(const at::Tensor& t) { NVTE_CHECK(nccl_sm != nullptr, "Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to " "\"NCCL\" before allocating EP payload buffers."); - return NVTECommWindow{static_cast(nccl_sm->get_window()), - static_cast(nccl_sm->get_offset())}; + // rendezvous resolves ``t`` by its storage base, so get_offset() is the allocation's offset in + // the NCCL window. Add ``t``'s own storage offset so a slice/view of a symm-mem allocation + // (e.g. the scale region carved from a shared recv buffer) resolves to its true position in the + // window rather than the allocation base. + const uint64_t offset = + static_cast(nccl_sm->get_offset()) + + static_cast(t.storage_offset()) * static_cast(t.element_size()); + return NVTECommWindow{static_cast(nccl_sm->get_window()), offset}; #else (void)t; return kNoWindow; @@ -114,6 +120,34 @@ DType check_topk_idx_dtype(at::Tensor topk_idx) { using Shape = std::vector; +// EP block scaling supports only E4M3 MXFP8 today. A future block-scaled recipe (e.g. NVFP4) +// would also set is_scaled but carry a non-FP8 token dtype, so key the guard on the FP8 dtype. +bool is_mxfp8_scaled(bool is_scaled, const at::Tensor& data) { + return is_scaled && data.scalar_type() == at::kFloat8_e4m3fn; +} + +// Validate the scale-inverse pair of a block-scaled EP op and return the scale column count +// (hidden/block). Both scales must be 2D contiguous with matching cols dividing H and numels +// equal to their row counts times cols; the recv scale must be symm-mem-backed under zero-copy. +size_t check_mxfp8_scale_pair(const at::Tensor& send_scale, const at::Tensor& recv_scale, + size_t send_rows, size_t recv_rows, size_t H, const char* recv_name) { + NVTE_CHECK(send_scale.dim() >= 2 && recv_scale.dim() >= 2, + "scale-inverses must be at least 2D [., H/block]"); + NVTE_CHECK(send_scale.is_contiguous() && recv_scale.is_contiguous(), + "scale-inverses must be contiguous"); + const size_t sc_cols = static_cast(send_scale.size(-1)); + NVTE_CHECK(sc_cols > 0 && H % sc_cols == 0, "scale cols (", sc_cols, + ") must be a non-zero divisor of hidden (", H, ")"); + NVTE_CHECK(static_cast(recv_scale.size(-1)) == sc_cols, + "recv scale cols must match send scale cols"); + NVTE_CHECK(static_cast(send_scale.numel()) == send_rows * sc_cols, + "send scale numel must equal rows * cols"); + NVTE_CHECK(static_cast(recv_scale.numel()) == recv_rows * sc_cols, + "recv scale numel must equal rows * cols"); + check_symm_mem_required(recv_scale, recv_name); + return sc_cols; +} + } // namespace bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_relaxed); } @@ -218,8 +252,13 @@ void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_pe total_recv_tokens_te.data(), &layer_cfg, stream); } +// tokens_scale_inv / recv_scale_inv are set only for block-scaled dispatch (for +// now MXFP8): tokens/recv_tokens carry e4m3 data and the scale tensors carry the +// unswizzled e8m0 scale-inverses [T, H/block]. Both null => bf16/fp16/fp32. void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, - at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights) { + at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights, + std::optional tokens_scale_inv, + std::optional recv_scale_inv) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(tokens.dim() >= 2, "tokens must be at least 2D [..., H]"); NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]"); @@ -250,16 +289,39 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, check_symm_mem_required(recv_tokens, "recv_tokens"); check_symm_mem_required(recv_topk_weights, "recv_topk_weights"); + // Block-scaled dispatch: tokens carry e4m3 data and the scale tensors carry + // unswizzled e8m0 scale-inverses [T, H/block]. Scales ride in the tensor; the + // backend keys on the tensor's scaling mode. is_mxfp8 is split from is_scaled + // so future block-scaled recipes can reuse the scale-routing plumbing while + // building their own TE tensors; only MXFP8 is supported for now. + const bool is_scaled = tokens_scale_inv.has_value(); + const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, tokens); + size_t sc_cols = 0; + if (is_scaled) { + NVTE_CHECK(recv_scale_inv.has_value(), + "recv_scale_inv must be provided together with tokens_scale_inv"); + NVTE_CHECK(is_mxfp8, "EP dispatch currently supports only E4M3 MXFP8 block scaling"); + sc_cols = check_mxfp8_scale_pair(*tokens_scale_inv, *recv_scale_inv, T_flat, recv_pr, H, + "recv_scale_inv"); + } + auto tok_dtype = GetTransformerEngineDType(tokens.scalar_type()); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); auto topk_idx_te = makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, idx_dtype); - auto tokens_te = makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype); + auto tokens_te = + is_mxfp8 ? makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype, + nullptr, nullptr, tokens_scale_inv->data_ptr(), + Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype); auto topk_w_te = makeTransformerEngineTensor(topk_weights.data_ptr(), Shape{T_flat, topk_n}, DType::kFloat32); auto recv_tokens_te = - makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype); + is_mxfp8 ? makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype, + nullptr, nullptr, recv_scale_inv->data_ptr(), + Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype); auto recv_topk_w_te = makeTransformerEngineTensor(recv_topk_weights.data_ptr(), Shape{recv_pr}, DType::kFloat32); @@ -269,6 +331,17 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, NVTECommWindow topk_w_win = maybe_make_window(topk_weights); NVTECommWindow recv_tokens_win = maybe_make_window(recv_tokens); NVTECommWindow recv_topk_w_win = maybe_make_window(recv_topk_weights); + // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window. + // Send scales (tokens_scale_inv) stay staged like the send data; recv scales + // are the one-sided write target and must be symm-mem-backed under zero-copy. + if (is_scaled) { + const NVTECommWindow tsi_win = maybe_make_window(*tokens_scale_inv); + const NVTECommWindow rsi_win = maybe_make_window(*recv_scale_inv); + tokens_win.scale_window = tsi_win.window; + tokens_win.scale_offset = tsi_win.offset; + recv_tokens_win.scale_window = rsi_win.window; + recv_tokens_win.scale_offset = rsi_win.offset; + } nvte_ep_dispatch(handle_mem_te.data(), topk_idx_te.data(), tokens_te.data(), tokens_win, topk_w_te.data(), topk_w_win, recv_tokens_te.data(), recv_tokens_win, recv_topk_w_te.data(), recv_topk_w_win, stream); @@ -344,7 +417,9 @@ void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_t g_recv_w_win, grad_tokens_te.data(), grad_topk_w_te.data(), stream); } -void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out) { +void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out, + std::optional grad_scale_inv, + std::optional grad_expert_out_scale_inv) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(grad.dim() >= 2, "grad must be at least 2D [..., H]"); NVTE_CHECK(grad_expert_out.dim() >= 2, "grad_expert_out must be at least 2D [..., recv_pr, H]"); @@ -363,17 +438,51 @@ void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expe // EpBuffer-owned scatter target and must be symm-mem in zero-copy mode. check_symm_mem_required(grad_expert_out, "grad_expert_out"); + // Block-scaled (MXFP8) backward: grad/grad_expert_out carry e4m3 data and the scale + // tensors carry the unswizzled e8m0 scale-inverses [., H/block]; the reverse-direction + // dispatch forwards them like the forward path. is_mxfp8 is split from is_scaled so + // future block-scaled recipes can reuse the plumbing; only MXFP8 is supported for now. + const bool is_scaled = grad_scale_inv.has_value(); + const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, grad); + size_t sc_cols = 0; + if (is_scaled) { + NVTE_CHECK(grad_expert_out_scale_inv.has_value(), + "grad_expert_out_scale_inv must be provided together with grad_scale_inv"); + NVTE_CHECK(is_mxfp8, "EP combine backward currently supports only E4M3 MXFP8 block scaling"); + sc_cols = check_mxfp8_scale_pair(*grad_scale_inv, *grad_expert_out_scale_inv, T_flat, recv_pr, + H, "grad_expert_out_scale_inv"); + } + auto g_dtype = GetTransformerEngineDType(grad.scalar_type()); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); - auto grad_te = makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype); + auto grad_te = is_mxfp8 + ? makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype, + nullptr, nullptr, grad_scale_inv->data_ptr(), + Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype); auto grad_expert_out_te = - makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype); + is_mxfp8 + ? makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype, + nullptr, nullptr, grad_expert_out_scale_inv->data_ptr(), + Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype); // grad is autograd-allocated (staged); grad_expert_out resolves to a symm-mem // window in zero-copy mode, else kNoWindow for the staged path. NVTECommWindow grad_win = maybe_make_window(grad); NVTECommWindow grad_expert_out_win = maybe_make_window(grad_expert_out); + // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window, + // mirroring the forward dispatch. Send scales stay staged like the send data; + // recv scales are the one-sided write target and must be symm-mem-backed. + if (is_scaled) { + const NVTECommWindow gsi_win = maybe_make_window(*grad_scale_inv); + const NVTECommWindow gesi_win = maybe_make_window(*grad_expert_out_scale_inv); + grad_win.scale_window = gsi_win.window; + grad_win.scale_offset = gsi_win.offset; + grad_expert_out_win.scale_window = gesi_win.window; + grad_expert_out_win.scale_offset = gesi_win.offset; + } nvte_ep_combine_bwd(handle_mem_te.data(), grad_te.data(), grad_win, grad_expert_out_te.data(), grad_expert_out_win, stream); } @@ -396,11 +505,16 @@ void register_ep_bindings(pybind11::module_& m) { py::arg("tokens_per_expert"), py::arg("top_k"), py::arg("dispatch_output_per_expert_alignment"), py::arg("total_recv_tokens"), py::call_guard()); - m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::call_guard()); + m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::arg("handle_mem"), py::arg("topk_idx"), + py::arg("tokens"), py::arg("topk_weights"), py::arg("recv_tokens"), + py::arg("recv_topk_weights"), py::arg("tokens_scale_inv") = std::nullopt, + py::arg("recv_scale_inv") = std::nullopt, py::call_guard()); m.def("ep_combine", &ep_combine, "EP combine", py::call_guard()); m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward", py::call_guard()); - m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", + m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", py::arg("handle_mem"), + py::arg("grad"), py::arg("grad_expert_out"), py::arg("grad_scale_inv") = std::nullopt, + py::arg("grad_expert_out_scale_inv") = std::nullopt, py::call_guard()); } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c6bf7eb516..7c8793b0f6 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -209,13 +209,19 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(), - py::arg("noop_flag") = py::none()); + py::arg("noop_flag") = py::none(), py::arg("output") = 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("last_dims") = py::none(), py::arg("tensor_offsets") = py::none()); + m.def("group_requantize_inplace", transformer_engine::pytorch::group_requantize_inplace, + "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle " + "its rowwise scales for GEMM, in place", + py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("otype"), py::arg("tensor_offsets") = py::none(), + py::arg("return_dequantized") = false); 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)", @@ -285,6 +291,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Backward of SwiGLU used in GPT OSS", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f); + /* Scaled activation */ + m.def("scaled_swiglu", transformer_engine::pytorch::scaled_swiglu, "Scaled SwiGLU activation", + py::arg("input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("glu_interleave_size") = 0); + m.def("scaled_clamped_swiglu", transformer_engine::pytorch::scaled_clamped_swiglu, + "Scaled clamped SwiGLU activation", py::arg("input"), py::arg("act_scales"), + py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, + py::arg("glu_linear_offset") = 1.0f, py::arg("glu_interleave_size") = 0); + m.def("scaled_srelu", transformer_engine::pytorch::scaled_srelu, "Scaled SReLU activation", + py::arg("input"), py::arg("act_scales"), py::arg("quantizer")); + m.def("scaled_dswiglu", transformer_engine::pytorch::scaled_dswiglu, "Scaled SwiGLU backward", + py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true); + m.def("scaled_clamped_dswiglu", transformer_engine::pytorch::scaled_clamped_dswiglu, + "Scaled clamped SwiGLU backward", py::arg("grad"), py::arg("fwd_input"), + py::arg("act_scales"), py::arg("quantizer"), py::arg("limit") = 7.0f, + py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f, + py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true); + m.def("scaled_dsrelu", transformer_engine::pytorch::scaled_dsrelu, "Scaled SReLU backward", + py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("compute_scale_grad") = true); /* DBias + DAct fusions*/ m.def("dbias_dgelu", transformer_engine::pytorch::dbias_dgelu, "DGeLU + DBias + Quantize", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer")); diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index c90a7d6d0d..59e3a63b29 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -389,6 +389,23 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW tensor_offsets.data_ptr, static_cast(tensor_offsets.dtype), tensor_offsets.shape); } + // Varying per-tensor dimensions. Leaving these unset declares the grouped tensor uniform, + // which selects the uniform-shape swizzle kernel. + const auto first_dims = input.get_first_dims(); + if (first_dims.data_ptr != nullptr) { + swizzle_input.set_first_dims(first_dims.data_ptr, static_cast(first_dims.dtype), + first_dims.shape); + swizzle_output.set_first_dims(first_dims.data_ptr, static_cast(first_dims.dtype), + first_dims.shape); + } + const auto last_dims = input.get_last_dims(); + if (last_dims.data_ptr != nullptr) { + swizzle_input.set_last_dims(last_dims.data_ptr, static_cast(last_dims.dtype), + last_dims.shape); + swizzle_output.set_last_dims(last_dims.data_ptr, static_cast(last_dims.dtype), + last_dims.shape); + } + // Per-tensor logical dimensions (uniform-shape grouped tensor). const size_t num_tensors = input.num_tensors(); const auto logical_shape_nvte = input.logical_shape(); @@ -398,11 +415,21 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW const size_t per_tensor_last_dim = logical_shape_nvte.data[logical_shape_nvte.ndim - 1]; constexpr size_t kMxfp8BlockSize = 32; - // Output is always allocated in the per-tensor padded ("swizzle-ready") layout - // so the cuDNN grouped GEMM consumer sees the correct stride between experts. - // The swizzle kernel itself handles converting from the kernel-emitted compact - // layout (per-tensor first dim is the unpadded value) to this padded layout. - auto compute_padded_grouped_scale_shape = [&](bool rowwise) { + const bool variable_shape = first_dims.data_ptr != nullptr || last_dims.data_ptr != nullptr; + + // Output is allocated in the layout the swizzle kernel writes so its consumer sees the + // correct stride between experts. + auto compute_padded_grouped_scale_shape = [&](bool rowwise) -> std::vector { + if (variable_shape) { + // Grouped variable-shape scale storage is a concatenation of per-tensor padded + // regions whose sizes live on the device. The swizzle kernel walks input and output + // with identical per-group strides, so swizzling is size-preserving and the output + // needs exactly the input's shape. The uniform-average formula below would be wrong + // here: the average of per-group sizes is generally not tile-aligned, so its + // rounded-up total matches neither the kernel's walk nor the input's capacity. + const auto &scales = rowwise ? row_scales : col_scales; + return nvte_shape_to_vector(scales.shape); + } const size_t m = rowwise ? per_tensor_first_dim : per_tensor_last_dim; const size_t k = rowwise ? per_tensor_last_dim : per_tensor_first_dim; const size_t padded_m = ceildiv(m, size_t{128}) * 128; diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index ddb85808a5..7a3b284aca 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -219,12 +219,23 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { } auto ret = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode); + auto get_initialized_storage_shape = [](const at::Tensor &data) { + auto shape = getTensorShape(data); + if (data.numel() == 0) { + // PyTorch may use a null pointer for a valid zero-sized allocation. TE Common reserves + // {nullptr, {0}} for uninitialized storage, so use an orientation-neutral 2D empty shape + // to distinguish explicitly provided rowwise or columnwise storage from missing storage. + shape = {0, 0}; + } + return shape; + }; + // Rowwise data if (!tensor.attr("rowwise_data").is_none()) { const auto &data = tensor.attr("rowwise_data").cast(); DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; - ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + ret.set_rowwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data)); } else if (quantizer_dtype != DType::kNumTypes) { ret.set_rowwise_data(nullptr, quantizer_dtype, std::vector{0}); } @@ -234,7 +245,7 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { const auto &data = tensor.attr("columnwise_data").cast(); DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; - ret.set_columnwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + ret.set_columnwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data)); } else if (quantizer_dtype != DType::kNumTypes) { ret.set_columnwise_data(nullptr, quantizer_dtype, std::vector{0}); } diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d1525b53f0..8605a4746b 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Methods needed for distributed training (DP/TP).""" + from __future__ import annotations from collections.abc import Iterable @@ -50,7 +51,6 @@ from .tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..debug.pytorch.debug_quantization import DebugQuantizedTensor - __all__ = ["checkpoint", "CudaRNGStatesTracker"] @@ -62,8 +62,8 @@ _USE_REENTRANT_ACTIVATION_RECOMPUTE = True -_FP8_ACTIVATION_RECOMPUTE_ENABLED = False -_FP8_ACTIVATION_RECOMPUTE_PHASE = False +_IN_ACTIVATION_RECOMPUTE_REGION = False +_ACTIVATION_RECOMPUTE_PHASE = False _ALL_ACTIVE_RNG_STATES = {} @@ -255,11 +255,14 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F self.recompute_phase = recompute_phase def __enter__(self): - global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = ( - self.activation_recompute and FP8GlobalStateManager.is_fp8_enabled() - ) - _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase + global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE + # Track the checkpoint region independently of the FP8 state at entry. + # A checkpointed callable may open its own FP8 autocast context (for + # example, to select precision per layer). Delayed-scaling modules in + # that inner context must still save their scale and amax metadata for + # the recompute forward. + _IN_ACTIVATION_RECOMPUTE_REGION = self.activation_recompute + _ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase qstate = FP8GlobalStateManager.quantization_state if self.activation_recompute and not self.recompute_phase: @@ -268,19 +271,19 @@ def __enter__(self): qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0) def __exit__(self, *exc_details): - global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = False - _FP8_ACTIVATION_RECOMPUTE_PHASE = False + global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE + _IN_ACTIVATION_RECOMPUTE_REGION = False + _ACTIVATION_RECOMPUTE_PHASE = False def is_fp8_activation_recompute_enabled() -> bool: - """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_ENABLED + """Whether we are in an activation recompute region with FP8 currently enabled""" + return _IN_ACTIVATION_RECOMPUTE_REGION and FP8GlobalStateManager.is_fp8_enabled() def in_fp8_activation_recompute_phase() -> bool: """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_PHASE + return _ACTIVATION_RECOMPUTE_PHASE def _get_active_autocast_contexts(): @@ -1882,6 +1885,11 @@ def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_gr _SYMM_MEM_POOL = None _SYMM_MEM_POOL_BACKEND = None +# Device the pool was created for; the torch symm_mem._symm_mem_pools cache is keyed by it. +_SYMM_MEM_POOL_DEVICE = None +# True when the pool was created via torch's get_mem_pool, which caches it in the private +# symm_mem._symm_mem_pools dict; release then has to drop that cached reference. +_SYMM_MEM_POOL_TORCH_CACHED = False def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): @@ -1890,12 +1898,14 @@ def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): backend arg, so the (process-global) backend is always set before the pool is created. The collective rendezvous cost is amortized across allocations (paid per new segment, not per buffer). """ - global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED if _SYMM_MEM_POOL is None: symm_mem.set_backend(backend) _SYMM_MEM_POOL_BACKEND = backend + _SYMM_MEM_POOL_DEVICE = device if hasattr(symm_mem, "get_mem_pool"): _SYMM_MEM_POOL = symm_mem.get_mem_pool(device) + _SYMM_MEM_POOL_TORCH_CACHED = True elif hasattr(torch.cuda, "MemPool") and hasattr(symm_mem, "get_mempool_allocator"): _SYMM_MEM_POOL = torch.cuda.MemPool(symm_mem.get_mempool_allocator(device)) else: @@ -1911,6 +1921,40 @@ def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): return _SYMM_MEM_POOL +def release_symm_mem_pool() -> None: + """Free the process-wide symm-mem pool's segments, deregistering their NCCL windows. + + Call before ``dist.destroy_process_group()``: the pool's windows are registered on + the group's NCCL comm, which becomes invalid once the group is destroyed. No-op if + no pool was created. + """ + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED + if _SYMM_MEM_POOL is None: + return + # The torch symm_mem._symm_mem_pools cache is keyed by the pool's creation device. + device = _SYMM_MEM_POOL_DEVICE + _SYMM_MEM_POOL = None + _SYMM_MEM_POOL_BACKEND = None + _SYMM_MEM_POOL_DEVICE = None + torch_cached = _SYMM_MEM_POOL_TORCH_CACHED + _SYMM_MEM_POOL_TORCH_CACHED = False + # A pool from torch's get_mem_pool is also cached in the private module dict + # symm_mem._symm_mem_pools; drop that reference so the segments' refcount reaches zero + # and their NCCL windows deregister. Fail loudly if this internal has changed shape, + # otherwise the windows would silently leak past destroy_process_group(). + if torch_cached: + pools = getattr(symm_mem, "_symm_mem_pools", None) + if not isinstance(pools, dict) or device not in pools: + raise RuntimeError( + "torch symmetric-memory pool cache (symm_mem._symm_mem_pools) is missing or " + "has changed layout; cannot release the pooled segments and their NCCL windows " + "would leak past destroy_process_group(). This torch version needs an updated " + "release_symm_mem_pool()." + ) + pools.pop(device, None) + torch.cuda.empty_cache() + + def symm_mem_alloc( shape, dtype: torch.dtype, diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 940812b6a4..b0c66df4cc 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -7,7 +7,7 @@ import atexit import warnings -from typing import Optional +from typing import Optional, TYPE_CHECKING import torch import torch.distributed as dist @@ -15,16 +15,23 @@ import transformer_engine_torch as tex from .cpu_offload import mark_not_offload -from .distributed import symm_mem_alloc +from .distributed import symm_mem_alloc, release_symm_mem_pool +from .quantized_tensor import QuantizedTensor +# Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of +# common.recipe (the concrete recipe classes are imported lazily where used). +if TYPE_CHECKING: + from ..common.recipe import Recipe __all__ = [ "EpBuffer", "ep_bootstrap", + "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", "ep_combine", "symm_mem_alloc", + "release_symm_mem_pool", "is_symm_backed", ] @@ -115,7 +122,9 @@ def ep_bootstrap( ``zero_copy`` opts the EP group into the symm-mem zero-copy IO path; pass ``True`` only when payload tensors are allocated via ``symm_mem_alloc``. - Requires ``recv_capacity_per_rank``. + Requires ``recv_capacity_per_rank``. To capture a CUDA graph, supply + persistent recv_tokens / grad_out buffers to dispatch/combine; the pool-based + auto-allocation used when they are omitted is not CUDA-graph capturable. ``num_topk`` is the per-token top-k; it sizes NCCL EP internal buffers. @@ -169,18 +178,27 @@ def ep_bootstrap( _ATEXIT_REGISTERED = True +def is_ep_bootstrapped() -> bool: + """Whether EP has been initialized in this process.""" + return _BOOTSTRAPPED + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent. An atexit handler covers normal interpreter shutdown, so most users do not need to call this. Call it explicitly only before ``dist.destroy_process_group()``, since the borrowed NCCL comm becomes - invalid once the PG is destroyed. + invalid once the PG is destroyed. This also releases the symm-mem pool, so + a caller that used ``symm_mem_alloc(use_pool=True)`` does not need a separate + ``release_symm_mem_pool()`` before destroying the PG. """ global _BOOTSTRAPPED, _EP_GROUP, _EAGER if not _BOOTSTRAPPED: return try: + # Deregister pooled symm-mem windows while the group's comm is still valid. + release_symm_mem_pool() tex.ep_finalize() finally: _BOOTSTRAPPED = False @@ -231,6 +249,8 @@ class EpBuffer: "eager", "total_recv_tokens", "_host_total_recv_tokens", + "dispatch_fwd_quant_recipe", + "combine_bwd_quant_recipe", ) def __init__( @@ -243,6 +263,8 @@ def __init__( alignment: int = 0, payload_dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, + dispatch_fwd_quant_recipe: Optional["Recipe"] = None, + combine_bwd_quant_recipe: Optional["Recipe"] = None, ) -> None: if not _BOOTSTRAPPED: raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") @@ -268,6 +290,8 @@ def __init__( self.payload_dtype = payload_dtype self.device = device self.zero_copy = bool(tex.ep_get_zero_copy()) + self.dispatch_fwd_quant_recipe = dispatch_fwd_quant_recipe + self.combine_bwd_quant_recipe = combine_bwd_quant_recipe size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) @@ -313,7 +337,7 @@ def _(*_args, **_kw): @torch.library.custom_op( f"{_LIB}::dispatch", - mutates_args=("recv_tokens", "recv_topk_weights"), + mutates_args=("recv_tokens", "recv_topk_weights", "recv_scale_inv"), device_types="cuda", ) def _dispatch_op( @@ -323,8 +347,19 @@ def _dispatch_op( topk_weights: torch.Tensor, recv_tokens: torch.Tensor, recv_topk_weights: torch.Tensor, + tokens_scale_inv: Optional[torch.Tensor] = None, + recv_scale_inv: Optional[torch.Tensor] = None, ) -> None: - tex.ep_dispatch(handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights) + tex.ep_dispatch( + handle_mem, + topk_idx, + tokens, + topk_weights, + recv_tokens, + recv_topk_weights, + tokens_scale_inv, + recv_scale_inv, + ) @_dispatch_op.register_fake @@ -372,15 +407,17 @@ def _(*_args, **_kw): @torch.library.custom_op( f"{_LIB}::combine_bwd", - mutates_args=("grad_expert_out",), + mutates_args=("grad_expert_out", "grad_expert_out_scale_inv"), device_types="cuda", ) def _combine_bwd_op( handle_mem: torch.Tensor, grad: torch.Tensor, grad_expert_out: torch.Tensor, + grad_scale_inv: Optional[torch.Tensor] = None, + grad_expert_out_scale_inv: Optional[torch.Tensor] = None, ) -> None: - tex.ep_combine_bwd(handle_mem, grad, grad_expert_out) + tex.ep_combine_bwd(handle_mem, grad, grad_expert_out, grad_scale_inv, grad_expert_out_scale_inv) @_combine_bwd_op.register_fake @@ -442,32 +479,84 @@ class _EpDispatch(torch.autograd.Function): def forward( # type: ignore[override] ctx, handle_mem: torch.Tensor, - recv_tokens: torch.Tensor, - recv_topk_weights: torch.Tensor, + recv_tokens: Optional[torch.Tensor], + recv_topk_weights: Optional[torch.Tensor], topk_idx: torch.Tensor, tokens: torch.Tensor, topk_weights: torch.Tensor, + tokens_scale_inv: Optional[torch.Tensor] = None, + token_counts: Optional[torch.Tensor] = None, + num_recv_tokens: Optional[int] = None, + payload_dtype: torch.dtype = torch.bfloat16, ): - """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand.""" + """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand. When scales are set + (MXFP8 for now), ``tokens`` is the quantized tensor kept as the autograd operand so grad + reaches the pre-quant input. Recv outputs are carved/allocated here: a caller may supply + ``recv_tokens`` / ``recv_topk_weights``, else they are sized to ``num_recv_tokens``.""" + is_scaled = tokens_scale_inv is not None + tokens_data = tokens._rowwise_data if isinstance(tokens, QuantizedTensor) else tokens + assert tokens_data.dim() == 2, "EP dispatch tokens must be 2D [num_tokens, hidden]" + hidden = tokens_data.shape[-1] + device = tokens_data.device + zero_copy = tex.ep_get_zero_copy() + + recv_scale_inv = None + if is_scaled: + if tokens._fp8_dtype != tex.DType.kFloat8E4M3: + raise NotImplementedError("EP dispatch supports only E4M3 MXFP8 tokens for now.") + # recv data + scales share one buffer (data then scales); carve or allocate it here. + recv_tokens, recv_scale_inv = _scale_alloc_io( + recv_tokens, + num_recv_tokens, + hidden, + tokens_scale_inv.shape[-1], + tokens_data.dtype, + tokens_scale_inv.dtype, + device, + zero_copy, + ) + # Reinterpret byte-backed FP8 data as the fp8 dtype so the backend sees a scaled tensor. + dispatch_tokens = tokens_data.view(torch.float8_e4m3fn) + dispatch_recv = recv_tokens.view(torch.float8_e4m3fn) + else: + if recv_tokens is None: + recv_tokens = _alloc_io((num_recv_tokens, hidden), payload_dtype, device, zero_copy) + dispatch_tokens = tokens_data + dispatch_recv = recv_tokens + if recv_topk_weights is None: + recv_topk_weights = _alloc_io((num_recv_tokens,), torch.float32, device, zero_copy) torch.ops.transformer_engine_ep.dispatch( handle_mem, topk_idx, - tokens, + dispatch_tokens, topk_weights, - recv_tokens, + dispatch_recv, recv_topk_weights, + tokens_scale_inv, + recv_scale_inv, ) ctx.save_for_backward(handle_mem) ctx.tokens_shape = tokens.shape - ctx.tokens_dtype = tokens.dtype ctx.topk_weights_shape = topk_weights.shape - ctx.tokens_T_flat = tokens.numel() // tokens.shape[-1] + ctx.num_tokens = tokens_data.shape[0] ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1] ctx.top_k = topk_weights.shape[-1] - ctx.hidden_dim = tokens.shape[-1] + ctx.hidden_dim = hidden # Detach so the long-lived buffers aren't tracked as differentiable outputs; - # autograd re-attaches grad_fn pointing back at this Function. - return recv_tokens.detach(), recv_topk_weights.detach() + # autograd re-attaches grad_fn pointing back at this Function. For scaled inputs + # the expert-major recv data + scales are wrapped into a per-expert GroupedTensor + # so downstream grouped GEMM and autograd see a proper quantized grouped tensor. + if is_scaled: + recv_out = _make_grouped_mxfp8( + recv_tokens.view(tokens._rowwise_data.dtype), + recv_scale_inv, + token_counts, + tokens._fp8_dtype, + tokens.dtype, + ) + else: + recv_out = recv_tokens.detach() + return recv_out, recv_topk_weights.detach() @staticmethod def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] @@ -476,8 +565,10 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] device = handle_mem.device g_recv_tokens = g_recv_tokens.contiguous() g_recv_topk_weights = g_recv_topk_weights.contiguous() + # Dispatch grad follows the recv grad's (high-precision) dtype; the quantizer's STE + # owns the fp8 boundary for scaled inputs. grad_tokens = torch.empty( - ctx.tokens_T_flat, ctx.hidden_dim, dtype=ctx.tokens_dtype, device=device + ctx.num_tokens, ctx.hidden_dim, dtype=g_recv_tokens.dtype, device=device ) grad_topk_weights = torch.empty( ctx.topk_T_flat, ctx.top_k, dtype=torch.float32, device=device @@ -496,21 +587,21 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] None, # topk_idx grad_tokens.view(ctx.tokens_shape), grad_topk_weights.view(ctx.topk_weights_shape), + None, # tokens_scale_inv (scales; non-differentiable) + None, # token_counts (per-expert counts; non-differentiable) + None, # num_recv_tokens (sizing scalar) + None, # payload_dtype (sizing scalar) ) class _EpCombine(torch.autograd.Function): - """Autograd combine. - - bwd scatters the expert_out grad into ``grad_out``. When the caller supplies it - (mcore-managed mode) that buffer is used as-is; otherwise it is allocated on the - fly here — from the symm-mem pool in zero-copy mode (one-sided target), or a plain - tensor in normal mode (keeps allocation torch.compile / CUDA-graph safe and lets - autograd own the grad's lifetime). + """Autograd combine; bwd scatters the expert_out grad into ``grad_out``. When the caller + supplies it that buffer is used as-is; otherwise it is allocated in the backward from the + symm-mem pool in zero-copy mode (one-sided target) or a plain tensor in normal mode (keeps + allocation torch.compile / CUDA-graph safe and lets autograd own the grad's lifetime). - ``grad_out`` is the backward's scatter target (an output it writes, never reads), - so it is stashed as a plain ctx attribute rather than via save_for_backward, which - would version-track a tensor we mutate. + ``grad_out`` is a write-only scatter target, so it is stashed as a plain ctx attribute rather + than via save_for_backward, which would version-track a tensor we mutate. """ @staticmethod @@ -521,55 +612,100 @@ def forward( # type: ignore[override] hidden_dim: int, grad_out: Optional[torch.Tensor], expert_out: torch.Tensor, + bwd_quant_recipe=None, + token_counts: Optional[torch.Tensor] = None, ): - """Combine fwd; stashes the bwd grad target or expert_out shape to size it.""" + """Combine fwd; stashes the bwd grad target or expert_out shape to size it. When + ``bwd_quant_recipe`` is set, the backward sends the result-grad as MXFP8.""" device = expert_out.device result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) ctx.save_for_backward(handle_mem) ctx.grad_out = grad_out - if grad_out is None: - ctx.expert_out_shape = expert_out.shape - ctx.expert_out_dtype = expert_out.dtype - ctx.device = device + ctx.bwd_quant_recipe = bwd_quant_recipe + ctx.token_counts = token_counts + ctx.expert_out_shape = expert_out.shape + ctx.expert_out_dtype = expert_out.dtype + ctx.device = device return result @staticmethod def backward(ctx, g_result): # type: ignore[override] - """Combine bwd; scatters the result grad into the grad target.""" + """Combine bwd; scatters the result-grad to expert positions. High-precision sends the grad + as-is; a quantized recipe (MXFP8 today) quantizes it and returns the expert_out grad as a + per-expert GroupedTensor.""" if not g_result.is_contiguous(): g_result = g_result.contiguous() (handle_mem,) = ctx.saved_tensors - grad_expert_out = ctx.grad_out - if grad_expert_out is None: - grad_expert_out = _alloc_io( - ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() + + if ctx.bwd_quant_recipe is None: + grad_expert_out = ctx.grad_out + if grad_expert_out is None: + grad_expert_out = _alloc_io( + ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() + ) + torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) + else: + mx, g_scale_inv = _quantize_mxfp8(g_result) + g_data = mx._rowwise_data + recv_pr, hidden = ctx.expert_out_shape[0], ctx.expert_out_shape[-1] + ge_data, ge_scale_inv = _scale_alloc_io( + ctx.grad_out, + recv_pr, + hidden, + g_scale_inv.shape[-1], + g_data.dtype, + g_scale_inv.dtype, + ctx.device, + tex.ep_get_zero_copy(), + ) + # The backend keys on the fp8 scaling mode; reinterpret the byte-backed data as fp8. + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + g_data.view(torch.float8_e4m3fn), + ge_data.view(torch.float8_e4m3fn), + g_scale_inv, + ge_scale_inv, ) - torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) + grad_expert_out = _make_grouped_mxfp8( + ge_data, ge_scale_inv, ctx.token_counts, mx._fp8_dtype, ctx.expert_out_dtype + ) + return ( None, # handle_mem None, # num_local_tokens None, # hidden_dim None, # grad_out grad_expert_out, + None, # bwd_quant_recipe + None, # token_counts ) # Public high-level wrappers -# NCCL EP currently only supports bfloat16 payload tensors. +# NCCL EP inputs are bfloat16; MXFP8 is applied internally via the buffer's dispatch_fwd_quant_recipe. def _require_bf16(name: str, t: torch.Tensor) -> None: if t.dtype is not torch.bfloat16: raise NotImplementedError( - f"NCCL EP currently supports only bfloat16 payloads; got {name}.dtype={t.dtype}." + "NCCL EP currently supports only bfloat16 or MXFP8 payloads; got" + f" {name}.dtype={t.dtype}." ) def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tensor: """Allocate a dispatch/combine IO tensor the caller did not supply: from the symm-mem pool in - zero-copy mode (auto-registered segment, lifecycle managed by torch refcount), else plain.""" + zero-copy mode (auto-registered segment, lifecycle managed by torch refcount), else plain. + + The zero-copy pool path is not CUDA-graph capturable; supply persistent recv_tokens / grad_out + buffers to capture a graph.""" if zero_copy: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "EP zero-copy pool allocation is not CUDA-graph capturable; supply persistent " + "recv_tokens / grad_out buffers (allocated once via symm_mem_alloc) before capture." + ) t = symm_mem_alloc(shape, dtype, _EP_GROUP, device=device, use_pool=True) # symm-mem storage is non-resizable; exempt it from CPU activation offloading (which # releases via storage.resize_(0)). Matters for bf16 recv_tokens (the saved activation). @@ -578,6 +714,103 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso return torch.empty(*shape, dtype=dtype, device=device) +def _quantize_mxfp8(x: torch.Tensor): + """Quantize a high-precision tensor to MXFP8 and return ``(quantized_tensor, scale_inv)`` where + ``scale_inv`` is the compact ``[T, H/block]`` scale the EP backend routes. The quantized tensor + is returned so callers can keep it as the autograd operand; its ``_rowwise_data`` is the fp8 + payload and ``scale_inv.shape[-1]`` the scale-column count. EP routes and returns E4M3 data in + both directions, so quantize to E4M3 regardless of pass. Strips the GEMM scale row padding to + the compact ``[T, H/block]`` layout; requires a 16-byte-aligned scale row.""" + from .constants import MXFP8_BLOCK_SCALING_SIZE + from .tensor.mxfp8_tensor import MXFP8Quantizer + + mx = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False).quantize(x) + if mx._with_gemm_swizzled_scales: + raise RuntimeError( + "internal MXFP8 quantization produced swizzled scales; EP dispatch needs compact." + ) + data = mx._rowwise_data + scale_inv = mx._rowwise_scale_inv + if data is None or scale_inv is None: + raise ValueError("MXFP8 tokens must carry rowwise data and scale_inv for EP dispatch.") + t_flat = x.shape[0] + hidden = x.shape[-1] + cols = hidden // MXFP8_BLOCK_SCALING_SIZE + # The backend forwards each token's scale row with a 16-byte-aligned store, so the row + # (cols * dtype bytes) must be a multiple of 16. + scale_row_bytes = cols * scale_inv.element_size() + if scale_row_bytes % 16 != 0: + raise ValueError( + f"MXFP8 dispatch requires a 16-byte-aligned scale row; hidden={hidden} gives " + f"{scale_row_bytes} bytes. Use a hidden size that is a multiple of " + f"{16 * MXFP8_BLOCK_SCALING_SIZE}." + ) + # scale_inv is 2D [round_up(T, 128), cols]; drop the row padding to the logical [T, H/block] + # the backend expects. cols is a multiple of 4 (16-byte row), so no column padding and the + # slice stays contiguous; assert rather than force a copy. + scale_inv = scale_inv[:t_flat, :cols] + if not scale_inv.is_contiguous(): + raise ValueError( + "MXFP8 dispatch requires compact contiguous scales [T, H/block]; got a " + f"non-contiguous [{t_flat}, {cols}] slice." + ) + return mx, scale_inv + + +def _scale_alloc_io(buf, rows, data_cols, scale_cols, data_dtype, scale_dtype, device, zero_copy): + """Block-scaled output data + scale buffers, each ``rows`` tall, laid out back-to-back + (``[rows, data_cols]`` data of ``data_dtype`` then ``[rows, scale_cols]`` scales of + ``scale_dtype``). Carve both from a single caller ``buf`` when it is large enough, so one + symm-mem window backs both views; else allocate them (symm-mem pool under zero-copy, else + plain). Recipe-agnostic: byte sizes come from the element sizes.""" + data_bytes = rows * data_cols * torch.empty((), dtype=data_dtype).element_size() + scale_bytes = rows * scale_cols * torch.empty((), dtype=scale_dtype).element_size() + if buf is not None: + # Reinterpret in place; a non-contiguous buf would force a copy and leave the caller's + # buffer unwritten, so require contiguous and view rather than reshape. + if not buf.is_contiguous(): + raise ValueError("scaled output buffer must be contiguous.") + flat = buf.view(-1).view(torch.uint8) + if flat.numel() < data_bytes + scale_bytes: + raise ValueError( + f"scaled output buffer too small: need {data_bytes + scale_bytes} bytes " + f"(data + scales), got {flat.numel()}." + ) + data = flat[:data_bytes].view(data_dtype).reshape(rows, data_cols) + scale_inv = ( + flat[data_bytes : data_bytes + scale_bytes].view(scale_dtype).reshape(rows, scale_cols) + ) + return data, scale_inv + data = _alloc_io((rows, data_cols), data_dtype, device, zero_copy) + scale_inv = _alloc_io((rows, scale_cols), scale_dtype, device, zero_copy) + return data, scale_inv + + +def _make_grouped_mxfp8(data, scale_inv, token_counts, fp8_dtype, fake_dtype): + """Wrap expert-major MXFP8 recv data + compact e8m0 scales as a per-expert ``GroupedTensor``. + + ``token_counts`` (int64 [num_local_experts]) is the padded per-expert row counts (128-aligned), + used as the group sizes. Grouping is device-side (first_dims/tensor_offsets), so the counts never + sync to host; the outer shape is the static recv capacity, bounded per expert by first_dims. + """ + from .tensor.grouped_tensor import GroupedTensor + from .tensor.mxfp8_tensor import MXFP8Quantizer + + assert data.dim() == 2, "recv data must be 2D [capacity_rows, hidden]" + capacity_rows, hidden = data.shape + quantizer = MXFP8Quantizer(fp8_dtype, rowwise=True, columnwise=False) + return GroupedTensor( + shape=(capacity_rows, hidden), + dtype=fake_dtype, + num_tensors=token_counts.numel(), + quantizer=quantizer, + data=data.reshape(-1).detach(), + scale_inv=scale_inv.reshape(-1).detach(), + first_dims=token_counts, + tensor_offsets=tex.splits_to_offsets(token_counts, hidden), + ) + + def ep_dispatch( buffer: EpBuffer, tokens: torch.Tensor, @@ -587,39 +820,59 @@ def ep_dispatch( recv_tokens: Optional[torch.Tensor] = None, recv_topk_weights: Optional[torch.Tensor] = None, ): - """Prepare + dispatch with autograd. topk_idx must be int32 or int64. - - ``recv_tokens`` / ``recv_topk_weights`` are the dispatch recv outputs: pass caller-owned buffers - (mcore-managed mode; in zero-copy they must be symm-mem-backed) or leave them None to allocate on - the fly (zero-copy: symm-mem pool; normal: plain). In eager mode the recv outputs are sized to - this step's recv-token total and must not be caller-supplied. Returns (recv_tokens, - recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See ``buffer.total_recv_tokens`` for - the per-step recv total. + """Prepare + dispatch with autograd. ``tokens`` is bfloat16; ``topk_idx`` is int32 or int64. + + When the buffer's ``dispatch_fwd_quant_recipe`` is set (``MXFP8BlockScaling`` only for now), tokens + are quantized internally and recv is returned as a per-expert ``GroupedTensor``; otherwise recv + stays bfloat16. A pre-quantized ``tokens`` is not accepted. + + ``recv_tokens`` / ``recv_topk_weights`` are the recv outputs: pass caller-owned buffers + (symm-mem-backed under zero-copy) or leave them None to allocate. For MXFP8 the recv data and + scales share ``recv_tokens`` (data then scales), so size it to at least + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes. Eager mode sizes the recv outputs + per step and forbids caller-supplied buffers. Under zero-copy, leaving them None allocates from + the symm-mem pool, which is not CUDA-graph capturable; pass persistent buffers to capture a graph. + + Returns (recv_tokens, recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See + ``buffer.total_recv_tokens`` for the per-step recv total. """ - _require_bf16("tokens", tokens) if topk_weights.dtype is not torch.float32: raise TypeError( f"topk_weights must be float32; got dtype={topk_weights.dtype}. " "Cast with topk_weights.float() before calling." ) + if isinstance(tokens, QuantizedTensor): + raise NotImplementedError( + "NCCL EP dispatch takes a bfloat16 input and quantizes internally when the buffer's " + "dispatch_fwd_quant_recipe is set; a pre-quantized tensor is not accepted." + ) + _require_bf16("tokens", tokens) if buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): raise ValueError( "eager mode sizes the recv outputs from the per-step recv-token total " "and cannot use caller-supplied recv_tokens / recv_topk_weights" ) + # Prepare (routing AllGather) up front so the recv outputs can be sized; in # eager mode ep_prepare also host-syncs this step's recv-token total. tokens_per_expert = ep_prepare(buffer, topk_idx) - rows = buffer._host_total_recv_tokens if buffer.eager else buffer.recv_capacity_per_rank - if recv_tokens is None: - recv_tokens = _alloc_io( - (rows, buffer.hidden_dim), - buffer.payload_dtype, - buffer.device, - buffer.zero_copy, - ) - if recv_topk_weights is None: - recv_topk_weights = _alloc_io((rows,), torch.float32, buffer.device, buffer.zero_copy) + num_recv_tokens = ( + buffer._host_total_recv_tokens if buffer.eager else buffer.recv_capacity_per_rank + ) + + tokens_scale_inv = None + if buffer.dispatch_fwd_quant_recipe is not None: + from ..common.recipe import MXFP8BlockScaling + + if not isinstance(buffer.dispatch_fwd_quant_recipe, MXFP8BlockScaling): + raise NotImplementedError( + "EP block-scaled dispatch supports MXFP8BlockScaling only; got " + f"{type(buffer.dispatch_fwd_quant_recipe).__name__}." + ) + # Quantize here (not in forward) so the quantized tensor stays the autograd operand and grad + # reaches the pre-quant input; forward then carves the recv buffers and routes. + tokens, tokens_scale_inv = _quantize_mxfp8(tokens) + recv_tokens, recv_topk_weights = _EpDispatch.apply( buffer.handle_mem, recv_tokens, @@ -627,6 +880,10 @@ def ep_dispatch( topk_idx, tokens, topk_weights, + tokens_scale_inv, + tokens_per_expert, + num_recv_tokens, + buffer.payload_dtype, ) return recv_tokens, recv_topk_weights, tokens_per_expert @@ -640,11 +897,14 @@ def ep_combine( ): """Combine with autograd; caller pre-applies topk weighting. - ``expert_out`` is the combine input (always caller-supplied; in zero-copy it must be symm-mem- - backed). ``grad_out`` is the backward's grad target: pass a caller-owned buffer (mcore-managed - mode) or leave it None to allocate on the fly in the backward (zero-copy: symm-mem pool; normal: - plain). In eager mode the grad target is sized per step and must not be caller-supplied. Result - shape is (num_local_tokens, buffer.hidden_dim); defaults to buffer.max_tokens_per_rank rows. + ``expert_out`` is the combine input (symm-mem-backed under zero-copy). ``grad_out`` is the + backward's grad target: pass a caller-owned buffer or leave it None to allocate. For MXFP8 the + grad data and scales share ``grad_out`` (data then scales), so size it to at least + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes (non-zero-copy only). Eager mode sizes + the grad target per step and forbids a caller-supplied buffer. Under zero-copy, leaving it None + allocates from the symm-mem pool, which is not CUDA-graph capturable; pass a persistent buffer to + capture a graph. Result shape is (num_local_tokens, hidden_dim); num_local_tokens defaults to + buffer.max_tokens_per_rank. """ _require_bf16("expert_out", expert_out) if buffer.eager and grad_out is not None: @@ -654,10 +914,24 @@ def ep_combine( ) if num_local_tokens is None: num_local_tokens = buffer.max_tokens_per_rank + # When combine_bwd_quant_recipe is set the combine backward sends the result-grad over the + # wire as MXFP8 and returns the expert_out grad as a GroupedTensor. + bwd_quant_recipe = None + if buffer.combine_bwd_quant_recipe is not None: + from ..common.recipe import MXFP8BlockScaling + + if not isinstance(buffer.combine_bwd_quant_recipe, MXFP8BlockScaling): + raise NotImplementedError( + "EP combine backward supports MXFP8BlockScaling only; got " + f"{type(buffer.combine_bwd_quant_recipe).__name__}." + ) + bwd_quant_recipe = buffer.combine_bwd_quant_recipe return _EpCombine.apply( buffer.handle_mem, num_local_tokens, buffer.hidden_dim, grad_out, expert_out, + bwd_quant_recipe, + buffer.tokens_per_expert, ) diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 1b93b8254c..22e9c9ae3e 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -3,7 +3,9 @@ # See LICENSE for license information. """NVFuser functions and JIT utilities""" +import inspect import os +import warnings from functools import wraps from typing import Callable, Optional, Tuple import torch @@ -26,6 +28,8 @@ def lazy_compile(func): @wraps(func) def wrapper(*args, **kwargs): nonlocal compiled_func + if torch.compiler.is_compiling(): + return func(*args, **kwargs) if compiled_func is None: compiled_func = torch.compile(func) return compiled_func(*args, **kwargs) @@ -49,22 +53,66 @@ def wrapper(*args, **kwargs): if torch.__version__ >= "2": import torch._dynamo - def no_torch_dynamo(recursive=True): - """Decorator to disable Torch Dynamo, except during ONNX export.""" + def no_torch_dynamo(recursive=True, when=None): + """Decorator to disable Torch Dynamo, except during ONNX export. - def decorator(f): + `when` makes it conditional: it is called with the arguments of the call + keyed by parameter name -- an argument the call left out is absent, so + reading it with .get() yields its default as long as that default is + None -- and returns the reason this particular call cannot be traced, or + None to have it traced as usual. The reason is reported once per + distinct message. + """ + + def _disable(f): # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True - disabled_f = ( + return ( torch._dynamo.disable(f, recursive=recursive) if torch.__version__ >= "2.1" else torch._dynamo.disable(f) ) - @wraps(f) - def wrapper(*args, **kwargs): - if is_in_onnx_export_mode(): + def decorator(f): + disabled_f = _disable(f) + if when is None: + + @wraps(f) + def wrapper(*args, **kwargs): + if is_in_onnx_export_mode(): + return f(*args, **kwargs) + return disabled_f(*args, **kwargs) + + else: + parameters = inspect.signature(f).parameters + # Arguments are matched to names by position, which only holds + # without a *args in between. + assert not any( + p.kind is inspect.Parameter.VAR_POSITIONAL for p in parameters.values() + ), f"no_torch_dynamo(when=...) does not support *args, which {f.__name__} takes" + parameter_names = list(parameters) + + # The warning belongs inside the disabled function: warnings.warn + # graph-breaks on its own, which would mask the break that matters. + def report_and_run(reason, *args, **kwargs): + warnings.warn( + f"Falling back to eager execution under torch.compile: {reason} is" + " unsupported on the compiled path (graph-breaks under fullgraph=True).", + stacklevel=3, + ) return f(*args, **kwargs) - return disabled_f(*args, **kwargs) + + disabled_report_and_run = _disable(report_and_run) + + @wraps(f) + def wrapper(*args, **kwargs): # pylint: disable=function-redefined + if is_in_onnx_export_mode() or not torch.compiler.is_compiling(): + return f(*args, **kwargs) + call = dict(zip(parameter_names, args)) + call.update(kwargs) + reason = when(call) + if reason is None: + return f(*args, **kwargs) + return disabled_report_and_run(reason, *args, **kwargs) return wrapper @@ -72,7 +120,7 @@ def wrapper(*args, **kwargs): else: # Fallback for PyTorch < 2.0: no-op decorator - def no_torch_dynamo(recursive=True): # pylint: disable=unused-argument + def no_torch_dynamo(recursive=True, when=None): # pylint: disable=unused-argument """No-op decorator for PyTorch < 2.0.""" return lambda func: func diff --git a/transformer_engine/pytorch/module/__init__.py b/transformer_engine/pytorch/module/__init__.py index 3cf15efc11..98ab745448 100644 --- a/transformer_engine/pytorch/module/__init__.py +++ b/transformer_engine/pytorch/module/__init__.py @@ -5,7 +5,7 @@ """Module level PyTorch APIs""" from .layernorm_linear import LayerNormLinear from .linear import Linear -from .grouped_linear import GroupedLinear +from .grouped_linear import GroupedLinear, is_module_grouped_tensor_path_supported from .layernorm_mlp import LayerNormMLP from .layernorm import LayerNorm from .rmsnorm import RMSNorm diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 12497d0cb0..fe283cf1a3 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -82,7 +82,6 @@ from ...debug.pytorch.debug_quantization import DebugQuantizer from ...debug.pytorch.debug_state import TEDebugState - _NATIVE_SPLIT_QUANTIZER_TYPES = frozenset( { Float8Quantizer, @@ -446,7 +445,73 @@ def _split_quantize_and_bias( return outputs, grad_biases -__all__ = ["GroupedLinear"] +__all__ = ["GroupedLinear", "is_module_grouped_tensor_path_supported"] + + +def is_module_grouped_tensor_path_supported( + recipe: Optional[Recipe], + dtype: torch.dtype, +) -> bool: + """Whether the module grouped-tensor path supports this recipe and dtype. + + The grouped-tensor path dispatches to ``general_grouped_gemm_for_grouped_tensor`` + and does not inspect split values because they may reside in a CUDA tensor. + Inspecting them on the host would add synchronization and break CUDA Graph safety. + + Supported Compute Capability (CC) and precisions: + + * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 + block scaling. + * Blackwell (CC 10.x and 11.0): BF16/FP16, FP8 per-tensor current scaling, + MXFP8, and NVFP4 with RHT. + * Custom recipes are unsupported because they may assign different + quantizers to input, weight, and grad-output roles. This predicate + currently supports only built-in recipes with known uniform layouts. + * FP8 delayed scaling is unsupported because the required grouped + quantization kernels are unavailable. + * FP8 block scaling is unsupported by this path on Blackwell because it + does not implement the legacy path's MXFP8-broadcast emulation. + * Grouped GEMM requires cuBLASLt 13.3+, with 13.4+ required on Hopper, + 13.5+ required for FP8 per-tensor current scaling on Hopper, and 13.6+ + required for FP8 block scaling on Hopper. + * FP32 is unsupported by the cuBLASLt grouped GEMM. + + Runtime-only restrictions such as debug mode, CPU offloading, calibration, + output quantization, and backend selection are checked separately by + ``GroupedLinear``. + """ + if dtype not in (torch.bfloat16, torch.float16): + return False + + device_capability = get_device_compute_capability() + if not (9, 0) <= device_capability <= (11, 0): + return False + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + return False + if device_capability < (10, 0) and cublaslt_version < 130400: + return False + + if recipe is None: + return True + if recipe.custom(): + return False + if recipe.backward_override is not None: + return False + if recipe.float8_current_scaling(): + return device_capability >= (10, 0) or cublaslt_version >= 130500 + if recipe.float8_block_scaling(): + # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8. + return device_capability < (10, 0) and cublaslt_version >= 130600 + if recipe.mxfp8(): + return device_capability >= (10, 0) + if recipe.nvfp4(): + return ( + device_capability >= (10, 0) + and not recipe.disable_rht + and not recipe.row_scaled_activation + ) + return False class _GroupedLinear(torch.autograd.Function): @@ -464,104 +529,13 @@ def _maybe_dequantize( return tensor.dequantize(dtype=dtype) return cast_if_needed(tensor, dtype) - @staticmethod - def _is_grouped_tensor_path_supported( - *, - fp8: bool, - fp8_calibration: bool, - debug: bool, - cpu_offloading: bool, - backward_override: Optional[str], - save_original_input: bool, - activation_dtype: torch.dtype, - input_quantizers: List[Optional[Quantizer]], - output_quantizers: List[Optional[Quantizer]], - ) -> bool: - """Whether to use cuBLASLt grouped GEMM through GroupedTensor metadata. - - There are no checks whether split sizes are supported. Splits - may be in a CUDA tensor, so checking would hurt performance - and be incompatible with CUDA Graphs. - - Supported Compute Capability (CC) and precisions: - * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 - block scaling (1D/2D, including power-of-2 scales). - * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8 - per-tensor current scaling. - FP8 delayed scaling is not supported because the corresponding grouped - quantization kernels are missing. FP8 block scaling on Blackwell (SM100 and - SM110) raises instead of falling back: the fused path is Hopper-only and has - no MXFP8-broadcast emulation. Architectures outside the fused-path window - (e.g. SM120) fall back to the legacy path like every other recipe. - Grouped GEMM requires cuBLAS 13.3+ (13.4+ on Hopper, 13.5+ for FP8 - per-tensor current scaling on Hopper); otherwise the legacy path is used. - Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization - currently requires RHT. - - Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would - trigger a fatal error in the cuBLASLt grouped GEMM check. - """ - # 1. Filter by environment variable - if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): - return False - # 2. Filter out advanced features - if ( - debug - or cpu_offloading - or fp8_calibration - or backward_override is not None - or save_original_input - ): - return False - # 3. Filter by compute capability and cuBLAS version - device_capability = get_device_compute_capability() - if not (9, 0) <= device_capability <= (11, 0): - return False - cublaslt_version = tex.get_cublasLt_version() - if cublaslt_version < 130300: - return False - if device_capability < (10, 0) and cublaslt_version < 130400: - return False - # 4. Output quantization is not supported. - if any(q is not None for q in output_quantizers): - return False - # 5. Filter by quantization recipes. - if fp8: - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+. - if device_capability < (10, 0) and cublaslt_version < 130500: - return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM - # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast - # emulation. On Blackwell (SM100/SM110, the only other arch that reaches - # this branch) fail loudly rather than silently falling back to the - # unfused path the user explicitly opted out of. - if get_device_compute_capability() >= (10, 0): - raise RuntimeError( - "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the" - " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped" - " FP8 block-scaling path is Hopper-only. Unset" - " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused" - " path (emulated via MXFP8 GEMM on Blackwell)." - ) - return True - # MXFP8 and NVFP4 require Blackwell+. - if not (10, 0) <= device_capability <= (11, 0): - return False - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( - isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers - ) - return activation_dtype in (torch.bfloat16, torch.float16) - @staticmethod def _make_grouped_tensor( data: torch.Tensor, *, num_gemms: int, split_sizes: torch.Tensor, - base_split_offsets: torch.Tensor, + tensor_offsets: torch.Tensor, last_dim: int, dtype: torch.dtype, ) -> GroupedTensorStorage: @@ -573,7 +547,7 @@ def _make_grouped_tensor( quantizer=None, data=data.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * last_dim, + tensor_offsets=tensor_offsets, ) @staticmethod @@ -604,14 +578,101 @@ def _prepare_weights_for_grouped_tensor_gemm( weight_quantizers: List[Optional[Quantizer]], weight_workspaces: List[Optional[QuantizedTensorStorage]], *, + num_gemms: int, + single_grouped_weight: bool, with_quantized_compute: bool, columnwise_usage: bool, activation_dtype: torch.dtype, is_first_microbatch: Optional[bool], skip_fp8_weight_update: Optional[torch.Tensor], cache_weight: bool, - ) -> Tuple[List[torch.Tensor], List[Optional[QuantizedTensorStorage]]]: - """Prepare discrete weight tensors for GroupedTensor GEMM.""" + ) -> Tuple[ + Union[GroupedTensorStorage, List[torch.Tensor]], + List[Optional[QuantizedTensorStorage]], + ]: + """Prepare a grouped parameter or discrete weights for GroupedTensor GEMM.""" + if single_grouped_weight: + weight = weights[0] + if not isinstance(weight, GroupedTensorStorage): + raise TypeError( + "single_grouped_weight requires the weight parameter to be a GroupedTensor." + ) + + new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] + if weight.quantizer is not None: + if not with_quantized_compute: + raise RuntimeError( + "Quantized single grouped weights require quantized grouped GEMM compute." + ) + return weight, new_workspaces + + if not with_quantized_compute: + if weight.rowwise_data is None: + raise RuntimeError("Single grouped weight has no rowwise storage.") + if weight.rowwise_data.dtype == activation_dtype: + return weight, new_workspaces + data = weight.rowwise_data.to(dtype=activation_dtype) + return ( + GroupedTensorStorage( + shape=weight.logical_shape, + dtype=activation_dtype, + num_tensors=num_gemms, + shapes=weight.tensor_shapes, + quantizer=None, + data=data, + ), + new_workspaces, + ) + + weight_quantizer = weight_quantizers[0] + if weight_quantizer is None: + raise RuntimeError("Quantized grouped compute requires a weight quantizer.") + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + # forward() already applied _enable_weight_preswizzle(); preserve that decision + # because not every quantizer and weight shape supports fused quantize-swizzle. + + workspace = weight_workspaces[0] if weight_workspaces else None + + if workspace is not None and ( + workspace.quantizer is not weight_quantizer + or (columnwise_usage and workspace.columnwise_data is None) + ): + workspace = None + + if weight.rowwise_data is None: + raise RuntimeError("Single grouped weight has no rowwise storage to quantize.") + source = weight.rowwise_data.view(weight.logical_shape) + update_workspace = is_first_microbatch is None or is_first_microbatch + if workspace is None: + if cache_weight: + # Match quantize_weight(): persistent workspaces must be Tensor subclasses + # so autograd can save them without decomposing their storage metadata. + saved_internal = weight_quantizer.internal + weight_quantizer.internal = False + grouped_weight = tex.group_quantize( + source, + weight_quantizer, + num_gemms, + None, + ) + if cache_weight: + weight_quantizer.internal = saved_internal + elif skip_fp8_weight_update is not None or update_workspace: + grouped_weight = tex.group_quantize( + source, + weight_quantizer, + num_gemms, + None, + noop_flag=skip_fp8_weight_update, + output=workspace, + ) + else: + grouped_weight = workspace + + if cache_weight: + new_workspaces[0] = grouped_weight + return grouped_weight, new_workspaces + weights_for_gemm: List[torch.Tensor] = [] new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * len(weights) if not with_quantized_compute: @@ -651,13 +712,10 @@ def _validate_or_alloc_output( """ if buffer is None: return torch.empty((rows, cols), dtype=dtype, device=device) - if buffer.dim() != 2: - raise ValueError(f"Output buffer must be 2D, got {buffer.dim()}D.") - if buffer.size(0) != rows: - raise ValueError(f"Output buffer rows {buffer.size(0)} must match input rows {rows}.") - if buffer.size(1) != cols: + expected_shape = (rows, cols) + if buffer.shape != expected_shape: raise ValueError( - f"Output buffer last dim {buffer.size(1)} does not match required {cols}." + f"Output buffer shape {tuple(buffer.shape)} must match required {expected_shape}." ) if buffer.dtype != dtype: raise ValueError(f"Output buffer dtype {buffer.dtype} does not match required {dtype}.") @@ -671,6 +729,43 @@ def _validate_or_alloc_output( raise ValueError("Output buffer must not require gradient.") return buffer + @staticmethod + def _prepare_bias_for_grouped_tensor_gemm( + biases: Tuple[torch.Tensor, ...], + *, + single_grouped_bias: bool, + num_gemms: int, + out_features: int, + dtype: torch.dtype, + ) -> GroupedTensorStorage: + """Prepare grouped or discrete bias storage for grouped GEMM.""" + if not single_grouped_bias: + return _GroupedLinear._make_grouped_bias( + biases, + num_gemms=num_gemms, + out_features=out_features, + dtype=dtype, + ) + + bias = biases[0] + if not isinstance(bias, GroupedTensorStorage): + raise TypeError("single_grouped_bias requires a GroupedTensor parameter.") + bias_data = bias.rowwise_data + if bias_data.dtype != dtype: + bias_data = bias_data.to(dtype=dtype) + + # The parameter exposes a packed [num_gemms, out_features] tensor, but its grouped + # members are 1D vectors. The grouped bias-add kernel consumes those same bytes as + # num_gemms row matrices with shape [1, out_features]. + return GroupedTensorStorage( + shape=(num_gemms, out_features), + dtype=dtype, + num_tensors=num_gemms, + shapes=[(1, out_features)] * num_gemms, + quantizer=None, + data=bias_data.reshape(-1), + ) + @staticmethod def _forward_grouped_tensor( ctx, @@ -692,6 +787,9 @@ def _forward_grouped_tensor( weight_workspaces: List[Optional[QuantizedTensorStorage]], cache_weight: bool, skip_fp8_weight_update: Optional[torch.Tensor], + save_original_input: bool, + single_grouped_weight: bool, + single_grouped_bias: bool, weights: Tuple[torch.Tensor, ...], biases: Tuple[torch.Tensor, ...], out: Optional[torch.Tensor] = None, @@ -701,11 +799,22 @@ def _forward_grouped_tensor( num_gemms = len(m_splits) device = inp.device in_features = weights[0].size(-1) - out_features = weights[0].size(0) + out_features = weights[0].size(-2) weight_requires_grad = weights[0].requires_grad + save_original_input = save_original_input and weight_requires_grad - split_sizes = m_splits.to(device=device) - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + split_sizes, ( + base_split_offsets, + input_tensor_offsets, + output_tensor_offsets, + ) = tex.splits_to_offsets_multi( + m_splits, + device, + strides=[1, in_features, out_features], + include_leading_zero=[True, True, True], + dtypes=[torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) inp_view = inp.reshape(-1, in_features) x = cast_if_needed(inp_view, activation_dtype) @@ -713,16 +822,22 @@ def _forward_grouped_tensor( input_quantizer = input_quantizers[0] input_quantizer.set_usage( rowwise=True, - columnwise=is_grad_enabled and weight_requires_grad, + columnwise=(is_grad_enabled and weight_requires_grad and not save_original_input), ) input_quantizer.optimize_for_gemm = True - grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes) + grouped_x = tex.group_quantize( + x, + input_quantizer, + num_gemms, + split_sizes, + tensor_offsets=input_tensor_offsets, + ) else: grouped_x = _GroupedLinear._make_grouped_tensor( x, num_gemms=num_gemms, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=input_tensor_offsets, last_dim=in_features, dtype=activation_dtype, ) @@ -732,6 +847,8 @@ def _forward_grouped_tensor( weights, weight_quantizers, weight_workspaces, + num_gemms=num_gemms, + single_grouped_weight=single_grouped_weight, with_quantized_compute=fp8, columnwise_usage=columnwise_usage, activation_dtype=activation_dtype, @@ -751,15 +868,16 @@ def _forward_grouped_tensor( out, num_gemms=num_gemms, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=output_tensor_offsets, last_dim=out_features, dtype=activation_dtype, ) grouped_bias = None if use_bias: - grouped_bias = _GroupedLinear._make_grouped_bias( + grouped_bias = _GroupedLinear._prepare_bias_for_grouped_tensor_gemm( biases, + single_grouped_bias=single_grouped_bias, num_gemms=num_gemms, out_features=out_features, dtype=activation_dtype, @@ -781,26 +899,35 @@ def _forward_grouped_tensor( ) if is_grad_enabled: + input_to_save = grouped_x if weight_requires_grad: - # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data - # in backward pass) - if fp8 and grouped_x.columnwise_data is not None: + if save_original_input: + # Save the high-precision input and reconstruct the grouped columnwise + # operand in backward instead of retaining a second quantized copy. + input_to_save = inp + elif fp8 and grouped_x.columnwise_data is not None: + # Wgrad only consumes the columnwise representation. grouped_x.rowwise_data = None grouped_x.scale_inv = None else: - grouped_x = None + input_to_save = None + + weights_to_save = [weights_for_gemm] if single_grouped_weight else weights_for_gemm + if not inp.requires_grad: + weights_to_save = [None] * len(weights_to_save) - weights_to_save = weights_for_gemm if inp.requires_grad else [None] * num_gemms tensors_to_save, tensor_objects = prepare_for_saving( - grouped_x, + input_to_save, *weights_to_save, split_sizes, base_split_offsets, + input_tensor_offsets, + output_tensor_offsets, ) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects - ctx.use_grouped_tensor_path = True + ctx.grouped_tensor_supported = True ctx.weight_quantizers = weight_quantizers ctx.weights_shape_0 = out_features ctx.weights_shape_1 = in_features @@ -808,16 +935,18 @@ def _forward_grouped_tensor( ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_weight_quantizers = grad_weight_quantizers ctx.weights_requires_grad = weight_requires_grad + ctx.single_grouped_weight = single_grouped_weight + ctx.single_grouped_bias = single_grouped_bias if fuse_wgrad_accumulation and ctx.weights_requires_grad: ctx.origin_weight_refs = [weakref.ref(w) for w in weights] ctx.origin_weights_overwrite_main_grad = getattr( weights[0], "overwrite_main_grad", False ) if hasattr(weights[0], "__fsdp_param__"): - ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + ctx.main_grad_funcs = [weight.get_main_grad for weight in weights] else: ctx.main_grad_funcs = [ - lambda j=i: weights[j].main_grad for i in range(num_gemms) + lambda j=i: weights[j].main_grad for i in range(len(weights)) ] ctx.device = device ctx.dgrad_out = dgrad_out @@ -841,7 +970,7 @@ def _forward_grouped_tensor( ) ctx.wgrad_store = wgrad_store ctx.debug = False - ctx.save_original_input = False + ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @@ -885,19 +1014,22 @@ def forward( delayed_scaling_input_quantizer, unsafe_requantization_input_quantizer, debug, + single_grouped_weight, + single_grouped_bias, + use_grouped_tensor, ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override - else: - backward_override = None + recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + backward_override = recipe.backward_override if recipe is not None else None if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": save_original_input = False num_gemms = len(m_splits) - weights = weights_and_biases[:num_gemms] - biases = weights_and_biases[num_gemms:] + num_weight_args = 1 if single_grouped_weight else num_gemms + num_bias_args = 1 if single_grouped_bias else num_gemms + weights = weights_and_biases[:num_weight_args] + biases = weights_and_biases[num_weight_args : num_weight_args + num_bias_args] device = inp.device weight_requires_grad = weights[0].requires_grad @@ -975,17 +1107,46 @@ def forward( f"weight tensor (shape={tuple(weights[0].size())})" ) - if _GroupedLinear._is_grouped_tensor_path_supported( - fp8=fp8, - fp8_calibration=fp8_calibration, - debug=debug, - cpu_offloading=cpu_offloading, - backward_override=backward_override, - save_original_input=save_original_input, - activation_dtype=activation_dtype, - input_quantizers=input_quantizers, - output_quantizers=output_quantizers, + grouped_tensor_supported = False + if use_grouped_tensor and not ( + fp8_calibration + or debug + or cpu_offloading + or any(q is not None for q in output_quantizers) ): + if ( + fp8 + and recipe.float8_block_scaling() + and (10, 0) <= get_device_compute_capability() <= (11, 0) + ): + raise RuntimeError( + "use_grouped_tensor=True does not support the FP8 block-scaling recipe on " + "Blackwell GPUs: the native grouped FP8 block-scaling path is Hopper-only. " + "Set use_grouped_tensor=False, or unset " + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM if it enabled this path, to use " + "the MXFP8-emulated path on Blackwell." + ) + grouped_tensor_supported = is_module_grouped_tensor_path_supported( + recipe, + activation_dtype, + ) + if ( + use_grouped_tensor + and not grouped_tensor_supported + and (single_grouped_weight or single_grouped_bias) + ): + raise RuntimeError( + "Single grouped parameters require the native grouped-tensor path, but the active " + "device, cuBLASLt version, quantization recipe, or GroupedLinear feature " + "configuration does not support it. Disable single_grouped_weight and " + "single_grouped_bias to allow the split-quantize fallback." + ) + if grouped_tensor_supported: + if m_splits.device.type != "cuda": + raise ValueError( + "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " + "tensor, or set use_grouped_tensor=False." + ) return _GroupedLinear._forward_grouped_tensor( ctx, inp=inp, @@ -1005,6 +1166,9 @@ def forward( weight_workspaces=weight_workspaces, cache_weight=cache_weight, skip_fp8_weight_update=skip_fp8_weight_update, + save_original_input=save_original_input, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, weights=weights, biases=biases, out=out, @@ -1096,7 +1260,7 @@ def forward( mark_not_offload(*weights_fp8, *weights) if is_grad_enabled: - ctx.use_grouped_tensor_path = False + ctx.grouped_tensor_supported = False ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] @@ -1213,13 +1377,47 @@ def _backward_grouped_tensor( """Backward path paired with ``_forward_grouped_tensor``.""" saved_tensors = restore_from_func_ctx(ctx) N = ctx.num_gemms - grouped_x = saved_tensors[0] - weights = saved_tensors[1 : 1 + N] - split_sizes = saved_tensors[1 + N] - base_split_offsets = saved_tensors[2 + N] + saved_input = saved_tensors[0] + if ctx.single_grouped_weight: + weights_for_gemm = saved_tensors[1] + weight_tensors = [weights_for_gemm] + split_sizes = saved_tensors[2] + base_split_offsets = saved_tensors[3] + input_tensor_offsets = saved_tensors[4] + output_tensor_offsets = saved_tensors[5] + else: + weight_tensors = saved_tensors[1 : 1 + N] + weights_for_gemm = weight_tensors + split_sizes = saved_tensors[1 + N] + base_split_offsets = saved_tensors[2 + N] + input_tensor_offsets = saved_tensors[3 + N] + output_tensor_offsets = saved_tensors[4 + N] + + if ctx.save_original_input: + x = cast_if_needed( + saved_input.reshape(-1, ctx.weights_shape_1), + ctx.activation_dtype, + ) + if ctx.fp8: + input_quantizer = ctx.input_quantizers[0] + input_quantizer.set_usage(rowwise=False, columnwise=True) + input_quantizer.optimize_for_gemm = True + grouped_x = tex.group_quantize(x, input_quantizer, N, split_sizes) + else: + grouped_x = _GroupedLinear._make_grouped_tensor( + x, + num_gemms=N, + split_sizes=split_sizes, + tensor_offsets=input_tensor_offsets, + last_dim=ctx.weights_shape_1, + dtype=ctx.activation_dtype, + ) + else: + grouped_x = saved_input - origin_weights = [None] * N - main_grads = [None] * N + num_weight_args = 1 if ctx.single_grouped_weight else N + origin_weights = [None] * num_weight_args + main_grads = [None] * num_weight_args if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None @@ -1253,6 +1451,7 @@ def _backward_grouped_tensor( grad_output_quantizer, N, split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = tex.group_quantize( @@ -1260,22 +1459,28 @@ def _backward_grouped_tensor( grad_output_quantizer, N, split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = _GroupedLinear._make_grouped_tensor( dy_2d, num_gemms=N, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=output_tensor_offsets, last_dim=ctx.weights_shape_0, dtype=ctx.activation_dtype, ) - grad_biases = [None] * N if ctx.use_bias: if dbias_packed is None: dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N) - grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] + if ctx.single_grouped_bias: + grad_bias_args = [dbias_packed.to(dtype=ctx.activation_dtype)] + else: + grad_bias_args = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] + else: + num_bias_args = 1 if ctx.single_grouped_bias else N + grad_bias_args = [None] * num_bias_args dgrad = None if ctx.requires_dgrad: @@ -1284,7 +1489,7 @@ def _backward_grouped_tensor( recipe = ctx.fp8_recipe if hasattr(recipe, "fp8_gemm_dgrad"): dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - for weight in weights: + for weight in weight_tensors: if isinstance(weight, QuantizedTensorStorage): weight.update_usage(columnwise_usage=True) dgrad = _GroupedLinear._validate_or_alloc_output( @@ -1298,12 +1503,12 @@ def _backward_grouped_tensor( dgrad, num_gemms=N, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=input_tensor_offsets, last_dim=ctx.weights_shape_1, dtype=ctx.activation_dtype, ) general_grouped_gemm_for_grouped_tensor( - weights, + weights_for_gemm, grouped_dy, grouped_dgrad, layout="NN", @@ -1324,16 +1529,42 @@ def _backward_grouped_tensor( if hasattr(recipe, "fp8_gemm_wgrad"): wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator if ctx.fuse_wgrad_accumulation: - wgrad_list = main_grads + if ctx.single_grouped_weight: + main_grad = main_grads[0] + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=N, + tensor_shape=(ctx.weights_shape_0, ctx.weights_shape_1), + rowwise_data=main_grad.view(-1), + dtype=main_grad.dtype, + ) + wgrad_output = grouped_wgrad + wgrad_list = [main_grad] + else: + wgrad_output = main_grads + wgrad_list = main_grads else: - wgrad_packed = torch.empty( - N, - ctx.weights_shape_0, - ctx.weights_shape_1, - dtype=ctx.activation_dtype, - device=ctx.device, - ) - wgrad_list = [wgrad_packed[i] for i in range(N)] + if ctx.single_grouped_weight: + grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=N, + shapes=[(ctx.weights_shape_0, ctx.weights_shape_1)] * N, + quantizer=None, + device=ctx.device, + dtype=ctx.activation_dtype, + ) + wgrad_output = grouped_wgrad + wgrad_list = [ + grouped_wgrad.rowwise_data.view(N, ctx.weights_shape_0, ctx.weights_shape_1) + ] + else: + wgrad_packed = torch.empty( + N, + ctx.weights_shape_0, + ctx.weights_shape_1, + dtype=ctx.activation_dtype, + device=ctx.device, + ) + wgrad_output = [wgrad_packed[i] for i in range(N)] + wgrad_list = wgrad_output accumulate = ( accumulate_wgrad_into_param_main_grad @@ -1353,9 +1584,9 @@ def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights): return None, [None] * N, None if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_list], grouped_gemm_wgrad) + ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], grouped_gemm_wgrad) else: - grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_list) + grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_output) def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if ctx.weights_requires_grad: @@ -1383,10 +1614,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) ] else: - wgrad_list = [None] * N - - if not ctx.use_bias: - grad_biases = [None] * N + wgrad_list = [None] * num_weight_args if ctx.reduce_and_update_bwd_fp8_tensors: FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) @@ -1397,7 +1625,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): None, # out None, # dgrad_out *wgrad_list, - *grad_biases, + *grad_bias_args, ) @staticmethod @@ -1406,7 +1634,7 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): - if ctx.use_grouped_tensor_path: + if ctx.grouped_tensor_supported: return _GroupedLinear._backward_grouped_tensor(ctx, grad_output) saved_tensors = restore_from_func_ctx(ctx) @@ -1705,7 +1933,9 @@ class GroupedLinear(TransformerEngineBaseModule): when set to ``True``, this module will not apply the additive bias itself, but instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when - the bias addition can be fused to subsequent operations. + the bias addition can be fused to subsequent operations. A single grouped + bias is returned as its packed ``GroupedTensor`` parameter; discrete biases + are returned as a list of per-GEMM tensors. params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters @@ -1730,6 +1960,13 @@ class GroupedLinear(TransformerEngineBaseModule): EXPERIMENTAL and subject to change. Gated by the ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var is not set this argument is forced to ``False`` with a warning. + use_grouped_tensor : bool or None, default = None + Prefer the native GroupedTensor grouped GEMM path. Discrete parameters + fall back to split-quantize when the path is unsupported. Single grouped + parameters require the native path and raise instead of falling back. + The native path requires CUDA ``m_splits``. ``None`` preserves the deprecated + ``NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM`` environment-variable + selection for compatibility. New callers should pass a boolean explicitly. Notes ----- @@ -1763,6 +2000,7 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, name: Optional[str] = None, + use_grouped_tensor: Optional[bool] = None, ) -> None: super().__init__(name) @@ -1778,11 +2016,37 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input + if use_grouped_tensor is None: + use_grouped_tensor_env = os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM") + if use_grouped_tensor_env is not None: + warnings.warn( + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM is deprecated and will be " + "removed in a future release. Pass use_grouped_tensor=True or " + "use_grouped_tensor=False to GroupedLinear instead.", + FutureWarning, + stacklevel=2, + ) + else: + use_grouped_tensor_env = "0" + use_grouped_tensor = bool(int(use_grouped_tensor_env)) + if not isinstance(use_grouped_tensor, bool): + raise TypeError( + f"use_grouped_tensor must be a bool or None, got {type(use_grouped_tensor)}." + ) + self.use_grouped_tensor = use_grouped_tensor single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags( single_grouped_weight, single_grouped_bias ) self.single_grouped_weight = single_grouped_weight self.single_grouped_bias = single_grouped_bias + if self.use_bias and self.single_grouped_weight and not self.single_grouped_bias: + warnings.warn( + "GroupedLinear has single_grouped_weight=True and bias=True, but " + "single_grouped_bias=False. This requires packing the per-GEMM biases on every " + "forward; enable single_grouped_bias to keep both parameters grouped.", + UserWarning, + stacklevel=2, + ) if ub_overlap_rs or ub_overlap_ag: raise ValueError("GroupedLinear doesn't support Userbuffer overlap.") self.init_method = init_method @@ -2000,7 +2264,7 @@ def make_grouped_weights(self, defer_init=False) -> None: if weight_quantizers and weight_quantizers[0] is not None else None ) - if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()): + if recipe is not None and recipe.delayed(): self.set_tensor_parallel_attributes(defer_init=defer_init) return @@ -2288,16 +2552,13 @@ def forward( is_grad_enabled = torch.is_grad_enabled() num_gemms = self.num_gemms - if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = ( - FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor - ) - else: - skip_fp8_weight_update = None - if skip_fp8_weight_update is not None: - is_first_microbatch = False - # Make sure splits are in expected format + if (self.single_grouped_weight or self.single_grouped_bias) and not self.use_grouped_tensor: + raise RuntimeError( + "single_grouped_weight and single_grouped_bias require " + "use_grouped_tensor=True; the split-quantize path only supports discrete " + "parameters." + ) if not isinstance(m_splits, torch.Tensor): # Convert list of ints to tensor for backward compatibility m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu") @@ -2326,6 +2587,7 @@ def forward( try: weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() + use_grouped_bias = self.use_bias and self.single_grouped_bias quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() @@ -2333,6 +2595,13 @@ def forward( if self.no_debug_features_active(list(chain(*quantizers))): debug = False quantizers = self._get_quantizers() + if debug and (self.single_grouped_weight or self.single_grouped_bias): + raise RuntimeError( + "TE debug features do not support single grouped parameters. DebugQuantizer " + "uses the split-quantize path, which only supports discrete parameters. " + "Disable single_grouped_weight and single_grouped_bias, or disable TE debug " + "features for this GroupedLinear." + ) ( input_quantizers, @@ -2358,11 +2627,14 @@ def forward( autograd_ctx = [None] cache_weight = is_first_microbatch is not None - weight_workspaces = ( - [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] - if cache_weight - else [None] * num_gemms - ) + if self.single_grouped_weight: + weight_workspaces = [self._fp8_workspaces.get("weight")] if cache_weight else [None] + else: + weight_workspaces = ( + [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] + if cache_weight + else [None] * num_gemms + ) non_tensor_args = ( self.apply_bias, @@ -2388,6 +2660,9 @@ def forward( self._delayed_scaling_input_quantizer, self._unsafe_requantization_input_quantizer, debug, + self.single_grouped_weight, + use_grouped_bias, + self.use_grouped_tensor, ) out, new_workspaces = linear_fn( *autograd_ctx, @@ -2405,12 +2680,15 @@ def forward( if ws is not None: if isinstance(ws, torch.Tensor): ws = ws.detach() - self._fp8_workspaces[f"weight{i}"] = ws + key = "weight" if self.single_grouped_weight else f"weight{i}" + self._fp8_workspaces[key] = ws finally: self.end_forward() if self.return_bias: + if use_grouped_bias: + return out, bias_tensors[0] return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors] return out @@ -2425,31 +2703,34 @@ def backward_dw(self): return with get_nvtx_range_context("_GroupedLinear_wgrad"): (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() - wgrad_list = tensor_list[2] + wgrad_output = tensor_list[2] weight_params = self._get_weight_tensors() if not self.fuse_wgrad_accumulation: - for i in range(self.num_gemms): - weight_params[i].grad = wgrad_list[i].to(weight_params[i].dtype) + if self.single_grouped_weight: + weight_params[0].grad = wgrad_output.rowwise_data.view( + self.num_gemms, self.out_features, self.in_features + ).to(weight_params[0].dtype) + else: + for i in range(self.num_gemms): + weight_params[i].grad = wgrad_output[i].to(weight_params[i].dtype) has_grad_biases = [ grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_ ] if self.use_bias and any(has_grad_biases): - grouped_bias = getattr(self, "bias", None) - if grouped_bias is not None: - if not all(has_grad_biases): - raise RuntimeError("Expected all grouped bias gradients to be present.") - gstack = torch.stack(grad_biases_, dim=0).to(grouped_bias.dtype) - if grouped_bias.grad is None: - grouped_bias.grad = gstack - else: - grouped_bias.grad.add_(gstack) - else: - bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] - for i in range(self.num_gemms): - if has_grad_biases[i] and bias_params[i].grad is None: - bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) + if self.use_grouped_tensor: + raise RuntimeError( + "GroupedLinear(use_grouped_tensor=True) fell back to the split-quantize " + "path, which produced per-expert bias gradients during delayed wgrad. " + "This implicit fallback is unsupported with delay_wgrad_compute=True. " + "Use a configuration supported by the grouped-tensor path, or set " + "use_grouped_tensor=False to select the legacy path explicitly." + ) + bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + for i in range(self.num_gemms): + if has_grad_biases[i] and bias_params[i].grad is None: + bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) del grad_biases_ - del wgrad_list + del wgrad_output del tensor_list self._trigger_wgrad_accumulation_and_reduce_hooks() @@ -2492,10 +2773,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage """Get the weight tensors of the module.""" grouped_weight = getattr(self, "weight", None) if grouped_weight is not None: - weight_tensors = grouped_weight.quantized_tensors - if weight_tensors is None: - # TODO(ksivaman): Remove this after GEMM integration. - weight_tensors = grouped_weight.split_into_quantized_tensors() + weight_tensors = [grouped_weight] else: weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors): @@ -2510,13 +2788,23 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage return weight_tensors def _get_bias_tensors(self) -> List[torch.Tensor]: - """Per-GEMM bias tensors (views into grouped storage when ``single_grouped_bias``).""" + """Get bias parameters in their registered grouped or per-GEMM layout. + + A single grouped bias remains one packed GroupedTensor. When return_bias=True, + an upper-level framework such as MCore must apply that packed bias accordingly; + Discrete bias parameters retain the existing list-of-per-GEMM contract. + + Example with 2 experts, 128 output features: + + single grouped bias: + GroupedTensor shape = [2, 128] -> [grouped_bias] + + discrete biases: + bias0 [128] + bias1 [128] -> [bias0, bias1] + """ grouped_bias = getattr(self, "bias", None) if grouped_bias is not None: - parts = grouped_bias.quantized_tensors - if parts is None: - parts = grouped_bias.split_into_quantized_tensors() - return [p.reshape(-1) for p in parts] + return [grouped_bias] return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] def _get_weight_quantizers(self) -> List[Quantizer]: diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 6def36ffc7..caab9c9c8d 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -24,7 +24,7 @@ from .bias import Bias from .constant_scale import ConstantScale from .dropout import Dropout -from .grouped_linear import GroupedLinear +from .grouped_linear import GroupedLinear, is_op_fuser_grouped_tensor_path_supported from .identity import Identity from .l2normalization import L2Normalization from .layer_norm import LayerNorm diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index f4beffe90c..8137051322 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -348,18 +348,8 @@ def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: return tex.dsrelu(*args, **kwargs) -class ScaledSReLU(BasicOperation): - r"""Squared ReLU with per-row post-scaling. - - If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied - with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. - - Parameters - ---------- - activation_recompute_in_mlp : bool, default = ``False`` - Enable fused grouped MLP kernels to recompute activation outputs - during backward when supported instead of saving them. - """ +class _ScaledUnary(BasicOperation, metaclass=abc.ABCMeta): + """Unary activation with per-row scales (fused grouped MLP middle op).""" num_extra_inputs: int = 1 @@ -367,6 +357,25 @@ def __init__(self, *, activation_recompute_in_mlp: bool = False) -> None: super().__init__() self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp + @abc.abstractmethod + def _scaled_unary_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: + """Apply the scaled unary activation.""" + + @abc.abstractmethod + def _scaled_unary_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Apply the scaled unary activation backward pass.""" + def op_forward(self, *args, **kwargs) -> None: raise RuntimeError( f"{self.__class__.__name__} operation has " @@ -410,7 +419,7 @@ def fuser_forward( x = maybe_dequantize(input_.contiguous(), dtype) scales = maybe_dequantize(extra_input, dtype) - y = tex.srelu(x, None) * scales.unsqueeze(-1) + y = self._scaled_unary_forward(x, scales) ctx = basic_op_ctxs[0] if ctx.requires_grad: @@ -448,21 +457,57 @@ def fuser_backward( scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype) - grad_input = None - if ctx.input_requires_grad: - grad_srelu_out = grad_output * scales.unsqueeze(-1) - grad_input = tex.dsrelu(grad_srelu_out, x, None) - - grad_extra_input = None - if ctx.extra_input_requires_grad: - srelu_out = tex.srelu(x, None) - grad_extra_input = torch.linalg.vecdot(srelu_out, grad_output) + grad_input, grad_extra_input = self._scaled_unary_backward( + grad_output, + x, + scales, + compute_scale_grad=ctx.extra_input_requires_grad, + ) + if not ctx.input_requires_grad: + grad_input = None clear_tensor_data(ctx.saved_tensors[0]) return grad_input, [()], [(grad_extra_input,)] +class ScaledSReLU(_ScaledUnary): + r"""Squared ReLU with per-row post-scaling. + + If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied + with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. + + Parameters + ---------- + activation_recompute_in_mlp : bool, default = ``False`` + Enable fused grouped MLP kernels to recompute activation outputs + during backward when supported instead of saving them. + """ + + def _scaled_unary_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: + return tex.scaled_srelu(input_, scales, None) + + def _scaled_unary_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + return tex.scaled_dsrelu( + grad_output, + input_, + scales, + None, + compute_scale_grad, + ) + + class SReGLU(_ActivationOperation): r"""Squared Rectified Gated Linear Unit diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea..1a932fb44d 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -14,7 +14,8 @@ import torch import transformer_engine_torch as tex -from ...constants import DType +from transformer_engine.common.recipe import Recipe +from ...constants import DType, TE_DType from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore @@ -24,14 +25,12 @@ _2X_ACC_WGRAD, ) from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload -from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe +from ...quantization import FP8GlobalStateManager, QuantizerRole from ...quantized_tensor import QuantizedTensorStorage from ...tensor import ( Float8BlockQuantizer, - Float8CurrentScalingQuantizer, MXFP8Quantizer, MXFP8Tensor, - NVFP4Quantizer, Quantizer, ) from ...utils import ( @@ -72,6 +71,67 @@ GRAD_INPUT_BUFFER_KEY = "grad_input" +def is_op_fuser_grouped_tensor_path_supported( + recipe: Optional[Recipe], + dtype: torch.dtype, +) -> bool: + """Whether the op-fuser grouped-tensor path supports this recipe and dtype. + + * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, + which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common + library. + * MXFP8 and NVFP4 are supported on Blackwell GPUs with Compute Capability + (CC) 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped + quantization currently requires it. + * FP8 per-tensor current scaling uses grouped current-scaling quantization + through ``tex.group_quantize`` and cuBLASLt grouped GEMM with per-batch + scalar FP8 scaling. It is supported on Hopper and Blackwell, with + cuBLASLt 13.5+ required on Hopper. + * FP8 block scaling uses the grouped-tensor path only on Hopper with + cuBLASLt 13.6+. On other architectures or older cuBLAS versions it + falls back to the split-quantize path for discrete parameters. + * Custom recipes are unsupported because they may assign different + quantizers to input, weight, and grad-output roles. This predicate + currently supports only built-in recipes with known uniform layouts. + * Other quantization recipes, including FP8 delayed scaling, fall back to + split quantization because their grouped quantization kernels are missing. + * Unquantized BF16/FP16 compute is supported on Hopper and Blackwell. FP32 + is excluded because cuBLASLt grouped GEMM does not support it. + * Single grouped parameters have no split-quantize fallback, so callers + must reject them when this function returns ``False``. + """ + if dtype not in (torch.bfloat16, torch.float16): + return False + + device_capability = get_device_compute_capability() + if not (9, 0) <= device_capability <= (11, 0): + return False + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + return False + if device_capability < (10, 0) and cublaslt_version < 130400: + return False + + if recipe is None: + return True + if recipe.custom(): + return False + if recipe.float8_current_scaling(): + return device_capability >= (10, 0) or cublaslt_version >= 130500 + if recipe.float8_block_scaling(): + # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8. + return device_capability < (10, 0) and cublaslt_version >= 130600 + if recipe.mxfp8(): + return device_capability >= (10, 0) + if recipe.nvfp4(): + return ( + device_capability >= (10, 0) + and not recipe.disable_rht + and not recipe.row_scaled_activation + ) + return False + + class GroupedLinear(BasicOperation): r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` @@ -316,17 +376,28 @@ def backward_dw(self) -> None: w.grad = grad_weights[group_idx].to(w.dtype) self._trigger_wgrad_accumulation_and_reduce_hooks() - def _get_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: - """Retrieve per-group bias tensors in the given dtype.""" + def _get_discrete_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: + """Retrieve discrete per-group bias parameters in the given dtype.""" if self.single_grouped_bias: - bias_parts = self.bias.quantized_tensors - if bias_parts is None: - bias_parts = self.bias.split_into_quantized_tensors() - return [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts] + raise RuntimeError( + "Discrete bias tensors were requested for a single grouped bias parameter." + ) return [ maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(self.num_groups) ] + def _get_packed_bias_tensor(self, dtype: torch.dtype) -> torch.Tensor: + """Return all per-group biases as one dense tensor with shape [num_groups, out_features]. + + A single grouped bias is already stored in this layout, so return a view of the + registered parent parameter instead of splitting it into members and stacking it again. + Discrete biases require a stack because they are independent parameters. + """ + if self.single_grouped_bias: + bias_data = self.bias.rowwise_data.view(self.num_groups, self.out_features) + return bias_data if bias_data.dtype == dtype else bias_data.to(dtype=dtype) + return torch.stack(self._get_discrete_bias_tensors(dtype), dim=0) + def num_quantizers(self, mode: str) -> int: if mode == "forward": return 2 * self.num_groups @@ -791,68 +862,6 @@ def op_backward(self, *args, **kwargs): "It overrides `fuser_backward` instead of `op_backward`." ) - @staticmethod - def _is_graph_safe_path_supported( - *, - with_quantized_compute: bool, - input_quantizers: Sequence[Optional[Quantizer]], - dtype: torch.dtype, - single_grouped_weight: bool, - ) -> bool: - """Whether the graph-safe grouped-tensor flow can be used. - - * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, - which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common - library. This filter mirrors cuBLASLt grouped GEMM's architecture - requirement without duplicating its cuBLAS version checks. - * Quantized compute supports MXFP8 and NVFP4 on Blackwell GPUs with Compute Capability (CC) - 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped quantization currently - requires RHT. NVFP4 is additionally restricted to discrete weights: with - ``single_grouped_weight=True`` the weight quantizer is non-RHT and cannot use the - graph-safe grouped quantize kernel, so we fall back to the split-quantize flow. - * FP8 per-tensor current scaling is backed by grouped current-scaling quantization - (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling, - which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0). - * FP8 block scaling uses the grouped-tensor path on Hopper (CC 9.0) with cuBLAS 13.4+; - it is Hopper-only (no MXFP8-broadcast emulation), so elsewhere it falls back. - * Every other quantization recipe (fp8 delayed scaling, ...) falls back to the legacy flow - because the corresponding grouped quantization kernels are missing. - * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0) - -- FP32 is excluded because the cuBLASLt grouped GEMM doesn't support it. - * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it - would trigger a fatal error in the cuBLASLt grouped GEMM check. - """ - if not (9, 0) <= get_device_compute_capability() <= (11, 0): - return False - if with_quantized_compute: - # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM - # path; the compute-capability range was already checked above. On Hopper it - # requires cuBLAS 13.5+; fall back to the legacy flow on older cuBLAS. - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - if ( - get_device_compute_capability() < (10, 0) - and tex.get_cublasLt_version() < 130500 - ): - return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block scaling is Hopper-only and needs cuBLAS 13.4+; elsewhere - # fall back to the split-quantize (MXFP8-emulated) flow. - if get_device_compute_capability() >= (10, 0): - return False - return tex.get_cublasLt_version() >= 130400 - # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. - if not (10, 0) <= get_device_compute_capability() <= (11, 0): - return False - if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): - return True - # NVFP4 graph-safe grouped quantization requires RHT and only supports - # discrete weights; otherwise fall back to the split-quantize flow. - if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): - return not single_grouped_weight - return False - return dtype in (torch.bfloat16, torch.float16) - def _get_grouped_weight_for_gemm( self, weight_param: GroupedTensor, @@ -974,16 +983,7 @@ def _get_grouped_bias_for_gemm( return None num_groups = self.num_groups - if self.single_grouped_bias: - # Already a contiguous (num_groups * out_features) buffer. - bias_data = self.bias.rowwise_data - if bias_data.dtype != dtype: - bias_data = bias_data.to(dtype=dtype) - else: - bias_list = [ - maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups) - ] - bias_data = torch.stack(bias_list, dim=0).contiguous() + bias_data = self._get_packed_bias_tensor(dtype) return GroupedTensorStorage( shape=(num_groups, self.out_features), @@ -1052,16 +1052,22 @@ def fuser_forward( out_buffer = basic_op_kwargs[0].get(OUTPUT_BUFFER_KEY) # Dispatch: graph-safe GroupedTensor flow whenever it can be used. - # See ``_is_graph_safe_path_supported`` for the gating rationale -- + # See ``is_op_fuser_grouped_tensor_path_supported`` for the gating rationale -- # in short it requires Hopper (SM90+) plus a supported dtype / # quantization recipe. Otherwise we fall back to the legacy # ``tex.split_quantize`` + ``general_grouped_gemm`` flow. - use_grouped_tensor_path = self._is_graph_safe_path_supported( - with_quantized_compute=with_quantized_compute, - input_quantizers=input_quantizers, - dtype=dtype, - single_grouped_weight=self.single_grouped_weight, + recipe = FP8GlobalStateManager.get_fp8_recipe() if with_quantized_compute else None + use_grouped_tensor_path = is_op_fuser_grouped_tensor_path_supported( + recipe, + dtype, ) + if (self.single_grouped_weight or self.single_grouped_bias) and not use_grouped_tensor_path: + raise RuntimeError( + "Single grouped parameters require the native grouped-tensor GroupedLinear path, " + "which is unavailable for the current device, dtype, or quantization recipe. " + "Disable single_grouped_weight/single_grouped_bias or use a supported grouped-" + "tensor configuration." + ) if use_grouped_tensor_path: out, tensors_to_save = self._fuser_forward_grouped_tensor( @@ -1136,14 +1142,17 @@ def fuser_forward_save_ctx( # temporary workspaces freshly created in each forward pass. if is_cpu_offload_enabled(): saved = tensors_to_save[0] - offset = 4 if self._scale_bias else 3 + # Metadata prefix: + # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, (scales?)] + offset = 6 if self._scale_bias else 5 if use_grouped_tensor_path: - # Layout: [split_sizes, base_split_offsets, split_points, (scales?), grouped_x, *weights] + # Layout: [..., grouped_x, *weights] grouped_x = saved[offset] if grouped_x is not None: mark_activation_offload(grouped_x) else: - # Layout: [split_sizes, None, None, (scales?), *xs, *ws] + # Layout: [..., *xs, *ws] live_xs = [t for t in saved[offset : offset + self.num_groups] if t is not None] if live_xs: mark_activation_offload(*live_xs) @@ -1169,9 +1178,9 @@ def fuser_forward_save_ctx( ctx.weight_quantizers = weight_quantizers ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_input_quantizers = None - # ``split_sizes`` and ``base_split_offsets`` are routed through - # ``save_for_backward`` (see ``_fuser_forward_split_quantize`` and - # ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). + # ``split_sizes``, offset metadata, and related tensors are routed + # through ``save_for_backward`` (see ``_fuser_forward_split_quantize`` + # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). if torch.is_autocast_enabled(): ctx.dtype = torch.get_autocast_dtype("cuda") else: @@ -1201,22 +1210,22 @@ def _fuser_forward_split_quantize( out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow.""" + if isinstance(input_, GroupedTensor): + raise NotImplementedError( + "Pre-quantized GroupedTensor input is only supported on the " + "graph-safe grouped-tensor path." + ) num_groups = self.num_groups has_bias = self.has_bias # Need CPU split sizes for split_quantize / general_grouped_gemm. split_sizes_int = [int(s) for s in split_sizes.tolist()] - # Extract params - if self.single_grouped_weight: - weights = self.weight.quantized_tensors - if weights is None: - weights = self.weight.split_into_quantized_tensors() - else: - weights = self._forward_weight_list() # materialized when distributed + # Single grouped parameters are rejected before entering this legacy path. + weights = self._forward_weight_list() # materialized when distributed bs = None if has_bias: - bs = self._get_bias_tensors(dtype) + bs = self._get_discrete_bias_tensors(dtype) ws = self._get_discrete_weights_for_gemm( weights, @@ -1285,12 +1294,12 @@ def _fuser_forward_split_quantize( # Build the tuple of tensors to save for backward. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if scale_bias), *xs, *ws] - # ``base_split_offsets`` and ``split_points`` are unused on the - # split-quantize backward path but are included as ``None`` so the - # saved-tensor layout matches the graph-safe - # ``_fuser_forward_grouped_tensor`` path (and the fused MLP forward). - saved: list[Optional[torch.Tensor]] = [split_sizes, None, None] + # Offset metadata slots are unused on the split-quantize backward path + # but are included as ``None`` so the saved-tensor layout matches the + # graph-safe ``_fuser_forward_grouped_tensor`` path. + saved: list[Optional[torch.Tensor]] = [split_sizes, None, None, None, None] if self._scale_bias: saved.append(scales) saved.extend(xs) @@ -1312,28 +1321,61 @@ def _fuser_forward_grouped_tensor( device: torch.device, out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: - """Graph-safe GroupedTensor forward path (pure compute). - Returns ``(output, tensors_to_save)``. ``split_sizes``, - ``base_split_offsets`` and ``split_points`` are returned so that - ``fuser_forward_save_ctx`` can call ``save_for_backward`` on them. - """ + """Build graph-safe grouped input storage and run grouped GEMM.""" num_groups = self.num_groups - has_bias = self.has_bias - - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) - split_points = base_split_offsets[1:].to(dtype=torch.int) - - # Flatten to 2D so the first dim is the total token count. + split_sizes, grouped_tensor_offsets = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=[1, 1, self.in_features, self.out_features], + include_leading_zero=[False, True, True, True], + dtypes=[torch.int32, torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) + split_points = grouped_tensor_offsets[0] + base_split_offsets = grouped_tensor_offsets[1] + input_tensor_offsets = grouped_tensor_offsets[2] + output_tensor_offsets = grouped_tensor_offsets[3] original_shape = list(input_.size()) - x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) - total_tokens = x.size(0) - - # Build the input GroupedTensor. + prequantized_input = with_quantized_compute and isinstance(input_, GroupedTensor) if with_quantized_compute: input_quantizer = input_quantizers[0] input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) input_quantizer.optimize_for_gemm = True - grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes) + if prequantized_input: + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, in_features) layout; just validate the shape. + if input_.dim() != 2 or input_.size(-1) != self.in_features: + raise ValueError( + "GroupedTensor input must have shape (total_tokens, " + f"{self.in_features}), but got {tuple(input_.size())}." + ) + total_tokens = input_.size(0) + else: + # Flatten to 2D so the first dim is the total token count. + x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) + total_tokens = x.size(0) + + # Build the input GroupedTensorStorage for input. + if prequantized_input: + # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data + # for the GEMM and let the helper supply whatever else the GEMMs need. + grouped_x = input_.copy() + tex.group_requantize_inplace( + grouped_x, + input_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=input_tensor_offsets, + ) + elif with_quantized_compute: + grouped_x = tex.group_quantize( + x, + input_quantizer, + num_groups, + split_sizes, + tensor_offsets=input_tensor_offsets, + ) else: # No quantize: wrap the contiguous high-precision buffer. grouped_x = GroupedTensorStorage( @@ -1343,8 +1385,9 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=x.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.in_features, + tensor_offsets=input_tensor_offsets, ) + has_bias = self.has_bias if is_cpu_offload_enabled() and grouped_x is not None: start_offload(grouped_x) @@ -1379,7 +1422,7 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=out.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.out_features, + tensor_offsets=output_tensor_offsets, ) # Bias: hand off to the grouped GEMM (graph-safe, fused). Plain bias @@ -1414,14 +1457,23 @@ def _fuser_forward_grouped_tensor( # Build the tuple of tensors to save for backward. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if _scale_bias), grouped_x, *weights] + # ``output_tensor_offsets`` matches the linear output row layout and is + # reused as ``grad_output`` offsets in backward. if grouped_x is not None: # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data # in backward pass) if with_quantized_compute and grouped_x.columnwise_data is not None: grouped_x.rowwise_data = None grouped_x.scale_inv = None - saved: list[Optional[torch.Tensor]] = [split_sizes, base_split_offsets, split_points] + saved: list[Optional[torch.Tensor]] = [ + split_sizes, + base_split_offsets, + split_points, + input_tensor_offsets, + output_tensor_offsets, + ] if self._scale_bias: saved.append(scales) saved.append(grouped_x) @@ -1471,13 +1523,13 @@ def _fuser_backward_split_quantize( # Saved tensors from forward pass. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if _scale_bias), *xs, *ws] - # ``base_split_offsets`` and ``split_points`` are unused on this path - # but are present so the saved-tensor layout matches the graph-safe - # path (and the fused MLP forward). + # Offset metadata beyond ``split_sizes`` is unused on this path but is + # present so the saved-tensor layout matches the graph-safe path. saved_tensors = ctx.saved_tensors split_sizes = saved_tensors[0] - saved_tensors = saved_tensors[3:] + saved_tensors = saved_tensors[5:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1505,7 +1557,7 @@ def _fuser_backward_split_quantize( offsets = torch.zeros(num_groups + 1, dtype=torch.int64, device=device) offsets[1:] = split_sizes.cumsum(0) if self._scale_bias: - bias_packed = torch.stack(self._get_bias_tensors(ctx.dtype)) + bias_packed = self._get_packed_bias_tensor(ctx.dtype) scales_f32 = scales.to(dtype=torch.float32) dbias_packed, grad_scales = compute_grouped_dbias_dscales( dy_2d, @@ -1653,24 +1705,24 @@ def _fuser_backward_grouped_tensor( Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: + """Graph-safe GroupedTensor backward path.""" num_groups = self.num_groups has_bias = self.has_bias weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device dtype = ctx.dtype - with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) - # Saved tensors from forward pass - # Layout: [split_sizes, base_split_offsets, split_points, - # (scales if _scale_bias), grouped_x, *weights] - # ``split_points`` is unused on this path but is present so the - # saved-tensor layout matches the fused MLP forward (which needs it - # for the cuDNN grouped GEMM kernel). + # Saved tensors from forward pass. Layout: + # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, + # (scales if _scale_bias), grouped_x, *weights] saved_tensors = ctx.saved_tensors split_sizes = saved_tensors[0] base_split_offsets = saved_tensors[1] - saved_tensors = saved_tensors[3:] + input_tensor_offsets = saved_tensors[3] + output_tensor_offsets = saved_tensors[4] + saved_tensors = saved_tensors[5:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1682,8 +1734,21 @@ def _fuser_backward_grouped_tensor( # Flatten grad_output to 2D (total_tokens, out_features) # to figure out total tokens. - dy_2d = grad_output.reshape(-1, self.out_features) - total_tokens = dy_2d.size(0) + prequantized_grad = with_quantized_compute and isinstance(grad_output, GroupedTensor) + if prequantized_grad: + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, out_features) layout; just validate the shape. + if grad_output.dim() != 2 or grad_output.size(-1) != self.out_features: + raise ValueError( + "GroupedTensor grad output must have shape (total_tokens, " + f"{self.out_features}), but got {tuple(grad_output.size())}." + ) + dy_2d = None + total_tokens = grad_output.size(0) + else: + dy_2d = grad_output.reshape(-1, self.out_features) + total_tokens = dy_2d.size(0) + grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] # Build the grad_output GroupedTensor. # Optionally get dbias is fusion available with bgrad_group_quantize @@ -1691,26 +1756,47 @@ def _fuser_backward_grouped_tensor( if with_quantized_compute: grad_output_quantizer = ctx.grad_output_quantizers[0] grad_output_quantizer.set_usage( - rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad + rowwise=ctx.input_requires_grad, + columnwise=ctx.weight_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True - # FP8 block scaling computes dbias in the rowwise (dgrad) pass, so only fuse # when dgrad is required. fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad ) - if has_bias and not self._scale_bias and fuse_bgrad: + if prequantized_grad: + # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its + # rowwise data for the dgrad GEMM. Bias grads are reduced from the dequantized + # grad below, which is only kept when there is a bias. + grouped_dy = grad_output.copy() + dy_2d = tex.group_requantize_inplace( + grouped_dy, + grad_output_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=output_tensor_offsets, + return_dequantized=has_bias, + ) + elif has_bias and not self._scale_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( - dy_2d, grad_output_quantizer, num_groups, split_sizes + dy_2d, + grad_output_quantizer, + num_groups, + split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = tex.group_quantize( - dy_2d, grad_output_quantizer, num_groups, split_sizes + dy_2d, + grad_output_quantizer, + num_groups, + split_sizes, + tensor_offsets=output_tensor_offsets, ) else: dy_2d = maybe_dequantize(dy_2d, dtype) - # Wrap BF16/FP16 buffer as a GroupedTensor for grouped gemm grouped_dy = GroupedTensorStorage( shape=(total_tokens, self.out_features), dtype=dtype, @@ -1718,7 +1804,7 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=dy_2d.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.out_features, + tensor_offsets=output_tensor_offsets, ) # Bias Grads compute if not already computed in bgrad_group_quantize @@ -1726,7 +1812,7 @@ def _fuser_backward_grouped_tensor( grad_scales: Optional[torch.Tensor] = None if has_bias: if self._scale_bias: - bias_packed = torch.stack(self._get_bias_tensors(dtype)) + bias_packed = self._get_packed_bias_tensor(dtype) scales_f32 = scales.to(dtype=torch.float32) dbias_packed, grad_scales = compute_grouped_dbias_dscales( dy_2d, @@ -1735,7 +1821,8 @@ def _fuser_backward_grouped_tensor( offsets=base_split_offsets, ) elif dbias_packed is None: - # BF16/FP16 path + # BF16/FP16 and pre-quantized MXFP8 paths, neither of which fuses dbias + # into a quantize kernel. dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups) if self.single_grouped_bias: final_bias_grads = [dbias_packed.to(dtype=dtype)] @@ -1745,7 +1832,6 @@ def _fuser_backward_grouped_tensor( # ---- dgrad GEMM ---------------------------------------------------- grad_input = None if ctx.input_requires_grad: - grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] grad_input = validate_or_alloc_output( getattr(ctx, "dgrad_out", None), grad_input_shape, dtype, device ) @@ -1756,7 +1842,7 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=grad_input.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.in_features, + tensor_offsets=input_tensor_offsets, ) general_grouped_gemm_for_grouped_tensor( dist_dgrad_weights if is_dist_weight else ws, diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 02f330ede3..fb663c0480 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -387,14 +387,21 @@ def __init__( self.glu_interleave_size: Optional[int] = glu_interleave_size self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: + def _scaled_glu_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: raise NotImplementedError - def _glu_backward( + def _scaled_glu_backward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, - ) -> torch.Tensor: + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: raise NotImplementedError def op_forward(self, *args, **kwargs) -> None: @@ -442,22 +449,7 @@ def fuser_forward( # Make sure inputs are in correct dtype input_ = maybe_dequantize(input_, dtype) scales = maybe_dequantize(extra_input, dtype) - - # Remove gate interleaving if needed - swiglu_in = input_ - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - swiglu_out = self._glu_forward(swiglu_in) - out = swiglu_out * scales.unsqueeze(-1) + out = self._scaled_glu_forward(input_, scales) # Save state for backward pass ctx = basic_op_ctxs[0] @@ -469,7 +461,7 @@ def fuser_forward( ctx.dtype = dtype ctx.save_for_backward( input_, - scales if ctx.input_requires_grad else None, + scales if ctx.input_requires_grad or ctx.extra_input_requires_grad else None, ) return out, [()] @@ -498,41 +490,14 @@ def fuser_backward( scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output, ctx.dtype) - # Remove gate interleaving if needed - swiglu_in = input_ - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - # Compute input grad - grad_input = None - if ctx.input_requires_grad: - grad_swiglu_out = grad_output * scales.unsqueeze(-1) - grad_swiglu_in = self._glu_backward(grad_swiglu_out, swiglu_in) - grad_input = grad_swiglu_in - if self.glu_interleave_size is not None: - shape = grad_input.size() - grad_input = grad_input.reshape( - -1, - 2, - shape[-1] // (2 * self.glu_interleave_size), - self.glu_interleave_size, - ) - grad_input = grad_input.transpose(1, 2).contiguous() - grad_input = grad_input.view(shape) - - # Compute scales grad by recomputing GLU - grad_extra_input = None - if ctx.extra_input_requires_grad: - swiglu_out = self._glu_forward(swiglu_in) - grad_extra_input = torch.linalg.vecdot(swiglu_out, grad_output) + grad_input, grad_extra_input = self._scaled_glu_backward( + grad_output, + input_, + scales, + compute_scale_grad=ctx.extra_input_requires_grad, + ) + if not ctx.input_requires_grad: + grad_input = None # Clear input tensor if possible clear_tensor_data(ctx.saved_tensors[0]) # input_ @@ -558,15 +523,34 @@ class ScaledSwiGLU(_ScaledGLU): """ - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: - return tex.swiglu(swiglu_in, None) - - def _glu_backward( + def _scaled_glu_forward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, ) -> torch.Tensor: - return tex.dswiglu(grad_swiglu_out, swiglu_in, None) + return tex.scaled_swiglu( + input_, + scales, + None, + int(self.glu_interleave_size or 0), + ) + + def _scaled_glu_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + return tex.scaled_dswiglu( + grad_output, + input_, + scales, + None, + int(self.glu_interleave_size or 0), + compute_scale_grad, + ) class ScaledClampedQGeGLU(_ScaledGLU): @@ -614,16 +598,39 @@ def __init__( glu_linear_offset=glu_linear_offset, ) - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: - return self._clamped._tex_clamped_swiglu_forward(swiglu_in, None) - - def _glu_backward( + def _scaled_glu_forward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, ) -> torch.Tensor: - return self._clamped._tex_clamped_dswiglu( - grad_swiglu_out, - swiglu_in, + clamped = self._clamped + return tex.scaled_clamped_swiglu( + input_, + scales, + None, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(self.glu_interleave_size or 0), + ) + + def _scaled_glu_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + clamped = self._clamped + return tex.scaled_clamped_dswiglu( + grad_output, + input_, + scales, None, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(self.glu_interleave_size or 0), + compute_scale_grad, ) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 0113833647..3582d35ccd 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -16,7 +16,7 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex -from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE +from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed_weight import ( @@ -31,7 +31,7 @@ from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from ...tensor.storage.grouped_tensor_storage import GroupedTensorStorage -from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales +from ...triton.grouped_dbias_dscales import compute_grouped_dbias, compute_grouped_dbias_dscales from ...utils import ( ceil_div, clear_tensor_data, @@ -101,9 +101,13 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() -def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool: - """Check cuDNN FE min version for single-group runtime offsets.""" - return _cudnn_frontend_version_at_least("1.27.0") +def _cudnn_frontend_supports_single_group_runtime_offsets( + activation_type: type[FusibleOperation], +) -> bool: + """Check cuDNN FE support for single-group runtime offsets.""" + return not issubclass(activation_type, ScaledSReLU) and _cudnn_frontend_version_at_least( + "1.27.0" + ) def _wrap_single_quantized_as_grouped( @@ -998,7 +1002,16 @@ def fuser_forward( # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) - input_ = input_.reshape(-1, fc1_weight_shape[1]) + if isinstance(input_, GroupedTensor): + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, in_features) layout; just validate the shape. + if input_.dim() != 2 or input_.size(-1) != fc1_weight_shape[1]: + raise ValueError( + "GroupedTensor input must have shape (total_tokens, " + f"{fc1_weight_shape[1]}), but got {tuple(input_.size())}." + ) + else: + input_ = input_.reshape(-1, fc1_weight_shape[1]) in_shape = list(input_.size()) if in_shape[0] % 128 != 0: raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).") @@ -1064,7 +1077,7 @@ def fuser_forward( activation_kernel = self.grouped_gemm_activation_kernel() supports_single_group_runtime_offsets = ( - _cudnn_frontend_supports_single_group_runtime_offsets() + _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) ) # Shared experts have one dense group and all optimized kernels derive M @@ -1091,21 +1104,40 @@ def fuser_forward( # shared-expert path never consumes it. base_split_offsets = split_sizes fc1_x_tensor_offsets = None + fc1_out_tensor_offsets = None fc2_x_tensor_offsets = None fc2_out_tensor_offsets = None else: + # Bulk-allocate every grouped-tensor offset the forward and backward + # passes need, so the backward can reuse them from the context + # instead of recomputing offsets per GEMM. split_sizes, ( split_points, base_split_offsets, fc1_x_tensor_offsets, + fc1_out_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], + strides=[ + 1, + 1, + fc1_weight_shape[1], + fc1_weight_shape[0], + fc2_weight_shape[1], + fc2_weight_shape[0], + ], + include_leading_zero=[False, True, True, True, True, True], + dtypes=[ + torch.int32, + torch.int64, + torch.int64, + torch.int64, + torch.int64, + torch.int64, + ], bulk_allocate=True, ) @@ -1199,42 +1231,18 @@ def fuser_forward( ) fc1_input_quantizer.optimize_for_gemm = True fc1_input_quantizer.internal = True - input_quantizer = getattr(input_, "quantizer", None) - if isinstance(input_, GroupedTensor) and ( - isinstance(fc1_input_quantizer, MXFP8Quantizer) - and isinstance(input_quantizer, MXFP8Quantizer) - or isinstance(fc1_input_quantizer, NVFP4Quantizer) - and isinstance(input_quantizer, NVFP4Quantizer) - ): - # GroupedTensor is a torch.Tensor subclass, so the CPU offload - # infrastructure's prepare_for_saving treats it as a plain tensor - # and does not decompose it into its component data tensors. By - # repacking into a GroupedTensorStorage (not a torch.Tensor), we - # ensure the fuser's prepare_for_saving call correctly decomposes - # the activation before save_for_backward. - grouped_fc1_x = GroupedTensorStorage( - shape=input_.logical_shape, - dtype=input_.fake_dtype, - num_tensors=input_.num_tensors, - shapes=input_.tensor_shapes, - quantizer=input_.quantizer, - data=input_.rowwise_data, - columnwise_data=input_.columnwise_data, - scale_inv=input_.scale_inv, - columnwise_scale_inv=input_.columnwise_scale_inv, - amax=input_.amax, - columnwise_amax=input_.columnwise_amax, - scale=input_.scale, - first_dims=input_.first_dims, - last_dims=input_.last_dims, - tensor_offsets=input_.tensor_offsets, - offsets=input_.offsets, - scale_inv_offsets=input_.scale_inv_offsets, - columnwise_scale_inv_offsets=input_.columnwise_scale_inv_offsets, - with_gemm_swizzled_scales=input_._with_gemm_swizzled_scales, - row_scaled_nvfp4=input_.row_scaled_nvfp4, - nvfp4_use_4over6=input_.nvfp4_use_4over6, - nvfp4_e4m3_max=input_.nvfp4_e4m3_max, + if isinstance(input_, GroupedTensor): + # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data + # for the GEMM and let the helper supply whatever else the GEMMs need. An input that + # is already GEMM-ready in both directions passes through untouched. + grouped_fc1_x = input_.copy() + tex.group_requantize_inplace( + grouped_fc1_x, + fc1_input_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=fc1_x_tensor_offsets, ) else: fc1_x = maybe_dequantize(input_, dtype) @@ -1743,6 +1751,10 @@ def fuser_forward( split_sizes, base_split_offsets, split_points, + fc1_x_tensor_offsets, + fc1_out_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, saved_fc1_x, *fc1_weight_tensors, activation_in, @@ -1786,7 +1798,16 @@ def fuser_backward( # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) - grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) + if isinstance(grad_output, GroupedTensor): + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, out_features) layout; just validate the shape. + if grad_output.dim() != 2 or grad_output.size(-1) != fc2_weight_shape[0]: + raise ValueError( + "GroupedTensor grad output must have shape (total_tokens, " + f"{fc2_weight_shape[0]}), but got {tuple(grad_output.size())}." + ) + else: + grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) out_shape = list(grad_output.size()) num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 @@ -1796,12 +1817,22 @@ def fuser_backward( # Saved tensors from the joint forward. # Layout: [split_sizes, base_split_offsets, split_points, + # fc1_x_tensor_offsets, fc1_out_tensor_offsets, + # fc2_x_tensor_offsets, fc2_out_tensor_offsets, # grouped_fc1_x, *fc1_weights, # activation_in, scales, # grouped_fc2_x, *fc2_weights] saved_tensors = fc1_ctx.saved_tensors - split_sizes, base_split_offsets, split_points = saved_tensors[:3] - saved_tensors = saved_tensors[3:] + ( + split_sizes, + base_split_offsets, + split_points, + fc1_x_tensor_offsets, + fc1_out_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) = saved_tensors[:7] + saved_tensors = saved_tensors[7:] grouped_fc1_x, saved_tensors = saved_tensors[0], saved_tensors[1:] if fc1_op.single_grouped_weight: grouped_fc1_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1846,20 +1877,28 @@ def fuser_backward( output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None fc2_dy = None - grad_output_quantizer = getattr(grad_output, "quantizer", None) - fc2_grad_output_quantizer_matches = ( - isinstance(fc2_grad_output_quantizer, MXFP8Quantizer) - and isinstance(grad_output_quantizer, MXFP8Quantizer) - ) or ( - isinstance(fc2_grad_output_quantizer, NVFP4Quantizer) - and isinstance(grad_output_quantizer, NVFP4Quantizer) - ) - if ( - not output_fc2_dbias - and isinstance(grad_output, GroupedTensor) - and fc2_grad_output_quantizer_matches - ): - grouped_fc2_dy = grad_output + if isinstance(grad_output, GroupedTensor): + # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise + # data for the dgrad GEMM. Bias grads are reduced from the dequantized grad, which is + # only materialized when one is needed. A grad that is already GEMM-ready in both + # directions passes through untouched. + grouped_fc2_dy = grad_output.copy() + fc2_dy = tex.group_requantize_inplace( + grouped_fc2_dy, + fc2_grad_output_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=fc2_out_tensor_offsets, + return_dequantized=output_fc2_dbias or scale_bias, + ) + if output_fc2_dbias and not scale_bias: + # This path has no quantize kernel to fuse dbias into, and the consumer below + # has no fallback, so reduce it here. + fc2_dbias_packed = compute_grouped_dbias(fc2_dy, base_split_offsets, num_groups) + # scale_bias is the only later consumer of the dequantized grad; drop it so the + # buffer is freed rather than held until backward ends. + fc2_dy = None else: fc2_dy = maybe_dequantize(grad_output, dtype) if output_fc2_dbias and not scale_bias: @@ -1868,6 +1907,7 @@ def fuser_backward( fc2_grad_output_quantizer, num_groups, split_sizes, + tensor_offsets=fc2_out_tensor_offsets, ) else: grouped_fc2_dy = _group_quantize_for_grouped_mlp( @@ -1875,9 +1915,7 @@ def fuser_backward( fc2_grad_output_quantizer, num_groups, split_sizes, - tensor_offsets=( - None if num_groups == 1 else base_split_offsets * fc2_weight_shape[0] - ), + tensor_offsets=fc2_out_tensor_offsets, ) use_nvfp4 = ( @@ -1998,7 +2036,7 @@ def fuser_backward( "use_dynamic_sched": True, } dactivation_kernel = self.grouped_gemm_dactivation_kernel() - if _cudnn_frontend_supports_single_group_runtime_offsets(): + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor @@ -2150,7 +2188,7 @@ def fuser_backward( fc2_input_quantizer, num_groups, split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[1], + tensor_offsets=fc2_x_tensor_offsets, ) else: sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") @@ -2172,15 +2210,14 @@ def fuser_backward( 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], + tensor_offsets=fc2_x_tensor_offsets, with_gemm_swizzled_scales=True, ) fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc2_bias_grad_packed: Optional[torch.Tensor] = None if scale_bias: - fc2_biases = fc2_op._get_bias_tensors(dtype) - bias_packed = torch.stack(fc2_biases) + bias_packed = fc2_op._get_packed_bias_tensor(dtype) fc2_dbias_packed_result, grad_scales = compute_grouped_dbias_dscales( fc2_dy, scales_f32, @@ -2215,9 +2252,7 @@ def fuser_backward( fc1_bias_grads = [dbias_2d[group_idx] for group_idx in range(num_groups)] # FC1 grad output for dgrad and wgrad GEMMs - fc1_dy_tensor_offsets = ( - None if num_groups == 1 else base_split_offsets * fc1_weight_shape[0] - ) + fc1_dy_tensor_offsets = fc1_out_tensor_offsets fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0] if use_nvfp4: fc1_grad_output_quantizer.set_usage( @@ -2304,7 +2339,6 @@ def fuser_backward( ) elif use_nvfp4: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) - fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] grouped_grad_input = GroupedTensor( shape=(out_shape[0], fc1_weight_shape[1]), dtype=dtype, @@ -2346,7 +2380,7 @@ def fuser_backward( "use_dynamic_sched": True, } fc1_dgrad_kernel = self.grouped_gemm_quant_kernel() - if _cudnn_frontend_supports_single_group_runtime_offsets(): + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if fc1_op.single_grouped_weight: diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index b48953e8e3..15290d7723 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -557,12 +557,11 @@ def step(self, closure=None, grad_scaler=None): loss = closure() for group in self.param_groups: - if len(group["params"]) == 0: - continue - device = group["params"][0].device - bias_correction = 1 if group["bias_correction"] else 0 - beta1, beta2 = group["betas"] - + # Advance the step counter before skipping empty groups. A param group can be + # empty on some data-parallel ranks and populated on others, so incrementing + # only for populated groups desynchronizes "step" across the ranks that share + # an optimizer state shard. A rank then resumes from a checkpoint written by a + # rank where the group was empty and applies a stale bias correction. # assume same step across group now to simplify things # per parameter step can be easily support by making it tensor, or pass list into kernel if "step" in group: @@ -570,10 +569,25 @@ def step(self, closure=None, grad_scaler=None): 1 if not self.capturable else (self._dummy_overflow_buf != 1).to(torch.int) ) else: + # Empty groups have no parameter to take the device from, so fall back to + # the device of the optimizer's own scratch buffer. + step_device = ( + group["params"][0].device + if len(group["params"]) > 0 + else self._dummy_overflow_buf.device + ) group["step"] = ( - 1 if not self.capturable else torch.tensor([1], dtype=torch.int, device=device) + 1 + if not self.capturable + else torch.tensor([1], dtype=torch.int, device=step_device) ) + if len(group["params"]) == 0: + continue + device = group["params"][0].device + bias_correction = 1 if group["bias_correction"] else 0 + beta1, beta2 = group["betas"] + # create lists for multi-tensor apply p_main_of_fp8_model = [] p_main_of_f16_model = []