diff --git a/benchmarks/benchmark_group_quantize_current_scaling.py b/benchmarks/benchmark_group_quantize_current_scaling.py new file mode 100644 index 0000000000..6d3efd3933 --- /dev/null +++ b/benchmarks/benchmark_group_quantize_current_scaling.py @@ -0,0 +1,732 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""End-to-end benchmark for grouped FP8 *current scaling* quantization. + +Times the full Float8CurrentScalingQuantizer.group_quantize path +(amax kernel + scale-from-amax kernel + cast/transpose kernels). + +Default shape mirrors the production use case: + sum(first_dims) = 98304, hidden = 2880 + +By default the benchmark sweeps both num_groups=16 and num_groups=64 so that +we can see how kernel performance scales with group count for the same total +work. Override with --num-groups 16 (or 64, or 16 64 ...) to restrict. +""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import dataclass +from typing import List, Optional + +# IMPORTANT: import transformer_engine before torch to avoid cublasLt symbol-resolution +# issues caused by torch's bundled CUDA libs. +import transformer_engine.pytorch # noqa: F401 - registers extension +from transformer_engine.pytorch import Float8CurrentScalingQuantizer +import transformer_engine_torch as tex +import torch + + +MODES = ("rowwise", "columnwise", "both") +SHAPE_CASES = ( + "same-shape", + "varying-first", + "varying-first-overalloc", + "varying-first-mild", + "varying-first-zipf", + "varying-first-heavy", + "varying-first-overalloc-mild", + "varying-first-overalloc-zipf", + "varying-first-overalloc-heavy", +) + + +def _parse_shape(spec: str, default_hidden: int) -> tuple: + """Parse a ``rows[xX,:]hidden`` shape spec into ``(actual_rows, hidden)``. + + Accepts separators ``x``, ``X``, ``,`` or ``:``. ``rows`` and ``hidden`` may + contain ``*`` so you can write the work directly, e.g. ``4096*16x4096``. If + no hidden is given, ``default_hidden`` is used (e.g. ``98304``). + """ + + def _eval_int(token: str) -> int: + token = token.strip() + # Only allow simple integer multiplication like "4096*16". + value = 1 + for part in token.split("*"): + value *= int(part) + return value + + sep = next((c for c in ("x", "X", ",", ":") if c in spec), None) + if sep is None: + return _eval_int(spec), default_hidden + rows_str, hidden_str = spec.split(sep, 1) + return _eval_int(rows_str), _eval_int(hidden_str) + + +@dataclass +class CaseResult: + shape_case: str + mode: str + actual_rows: int + allocated_rows: int + hidden: int + num_groups: int + iters: int + elapsed_ms_total: float + per_iter_us: float + relevant_bytes: int + bw_actual_TBps: float + bw_physical_TBps: float + loop: str = "eager" + # Per-kernel profiling metrics (populated only when --profile is set). Each + # is a dict with keys: per_iter_us, per_launch_us, launches_per_iter, + # bytes_per_launch, bw_TBps. + amax_profile: Optional[dict] = None + cast_profile: Optional[dict] = None + + +def _make_quantizer(mode: str) -> Float8CurrentScalingQuantizer: + q = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + force_pow_2_scales=False, + amax_epsilon=0.0, + ) + q.set_usage(rowwise=mode in ("rowwise", "both"), columnwise=mode in ("columnwise", "both")) + return q + + +def _generate_first_dims(shape_case: str, actual_rows: int, num_groups: int) -> List[int]: + if "mild" in shape_case: + # Mild imbalance: linear variation +/- 20% around average + if num_groups > 1: + weights = [0.8 + 0.4 * i / (num_groups - 1) for i in range(num_groups)] + else: + weights = [1.0] + elif "zipf" in shape_case: + # Zipf-based imbalance: w_i = 1 / (i + 1)^s, using s = 0.7 + weights = [1.0 / ((i + 1) ** 0.7) for i in range(num_groups)] + elif "heavy" in shape_case: + # Heavy imbalance: w_i = 1 / (i + 1)^1.5 + weights = [1.0 / ((i + 1) ** 1.5) for i in range(num_groups)] + else: + # Uniform + assert actual_rows % num_groups == 0, "actual_rows must divide num_groups" + return [actual_rows // num_groups] * num_groups + + # Normalize weights so they sum to actual_rows exactly + sum_weights = sum(weights) + first_dims = [int(round(w * actual_rows / sum_weights)) for w in weights] + + # Ensure all elements are >= 1 + for i in range(num_groups): + if first_dims[i] < 1: + first_dims[i] = 1 + + current_sum = sum(first_dims) + diff = actual_rows - current_sum + if diff > 0: + # Add 1 to the largest elements to preserve distribution shape + sorted_indices = sorted(range(num_groups), key=lambda idx: first_dims[idx], reverse=True) + for i in range(diff): + first_dims[sorted_indices[i % num_groups]] += 1 + elif diff < 0: + # Subtract 1 from the largest elements (as long as they remain > 1) + sorted_indices = sorted(range(num_groups), key=lambda idx: first_dims[idx], reverse=True) + idx = 0 + while diff < 0: + target_idx = sorted_indices[idx % num_groups] + if first_dims[target_idx] > 1: + first_dims[target_idx] -= 1 + diff += 1 + idx += 1 + + return first_dims + + +def _make_inputs( + shape_case: str, + actual_rows: int, + allocated_rows: int, + hidden: int, + num_groups: int, + num_buffers: int, +): + """Returns (inputs, first_dims_or_None, info).""" + dtype = torch.bfloat16 + info = { + "actual_rows": actual_rows, + "allocated_rows": allocated_rows, + "hidden": hidden, + "num_groups": num_groups, + } + if shape_case == "same-shape": + inputs = [ + torch.randn(actual_rows, hidden, dtype=dtype, device="cuda") for _ in range(num_buffers) + ] + return inputs, None, info + if shape_case in ( + "varying-first", + "varying-first-mild", + "varying-first-zipf", + "varying-first-heavy", + ): + first_dims_list = _generate_first_dims(shape_case, actual_rows, num_groups) + first_dims = torch.tensor(first_dims_list, dtype=torch.int64, device="cuda") + inputs = [ + torch.randn(actual_rows, hidden, dtype=dtype, device="cuda") for _ in range(num_buffers) + ] + info["first_dims"] = first_dims_list + return inputs, first_dims, info + if shape_case in ( + "varying-first-overalloc", + "varying-first-overalloc-mild", + "varying-first-overalloc-zipf", + "varying-first-overalloc-heavy", + ): + first_dims_list = _generate_first_dims(shape_case, actual_rows, num_groups) + first_dims = torch.tensor(first_dims_list, dtype=torch.int64, device="cuda") + inputs = [] + for _ in range(num_buffers): + t = torch.randn(allocated_rows, hidden, dtype=dtype, device="cuda") + # Tail rows must NOT influence amax (they shouldn't be read). + # Fill with poison values: a kernel that reads them will produce + # amax that disagrees with current_scaling reference. + t[actual_rows:].fill_(1e30) + inputs.append(t) + info["first_dims"] = first_dims_list + return inputs, first_dims, info + raise ValueError(shape_case) + + +def _verify_no_tail_read( + quantizer, inp: torch.Tensor, first_dims, num_groups: int, actual_rows: int +): + """Sanity check: amax must not be polluted by the poisoned tail rows.""" + out = tex.group_quantize(inp, quantizer, num_groups, first_dims) + amax_max = float(out.amax.max().item()) + # If tail rows (1e30) were read, amax >= 1e30. Real bf16 N(0,1) data: amax ~ 5. + if amax_max > 1e10: + raise RuntimeError( + f"group_quantize is reading the over-allocated tail rows: amax_max={amax_max}. " + "This means kernels are scanning unused memory (a perf and correctness bug)." + ) + return out + + +def _timed_eager(quantizer, inputs, first_dims, num_groups: int, iters: int) -> float: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for it in range(iters): + i = it % len(inputs) + tex.group_quantize(inputs[i], quantizer, num_groups, first_dims) + end.record() + end.synchronize() + return start.elapsed_time(end) # ms + + +def _timed_cuda_graph( + quantizer, inputs, first_dims, num_groups: int, iters: int, calls_per_replay: int = 32 +) -> float: + """Capture `calls_per_replay` group_quantize calls into one CUDA graph and + replay it `iters / calls_per_replay` times. Effectively eliminates Python / + cudaLaunchKernel overhead. The output tensor is allocated inside each call + by ``group_quantize``; under CUDA-graph capture this routes through the + graph allocator pool.""" + static_input = inputs[0] + + # Warmup must happen on a side stream before capture (per torch docs) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + tex.group_quantize(static_input, quantizer, num_groups, first_dims) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for _ in range(calls_per_replay): + tex.group_quantize(static_input, quantizer, num_groups, first_dims) + + replays = max(1, iters // calls_per_replay) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(replays): + g.replay() + end.record() + end.synchronize() + elapsed = start.elapsed_time(end) + actual_iters = replays * calls_per_replay + return elapsed, actual_iters + + +def _produced_out_copies(quantizer, inp: torch.Tensor, num_groups: int, first_dims) -> int: + """Number of FP8 output buffers the kernel actually materializes. + + Don't infer this from ``mode``: on Blackwell, FP8 *current scaling* reuses the + rowwise FP8 data as the columnwise GEMM operand (cuBLAS transposes on the fly), + so a ``both``-usage quantizer writes only the rowwise buffer -- not two. Detect + it empirically from the produced GroupedTensor so the byte/BW accounting matches + what the kernels actually move on this device. + """ + out = tex.group_quantize(inp, quantizer, num_groups, first_dims) + copies = 0 + for attr in ("rowwise_data", "columnwise_data"): + buf = getattr(out, attr, None) + if buf is not None and buf.numel() > 0: + copies += 1 + return copies + + +def _bytes_per_call(actual_rows: int, hidden: int, out_copies: int) -> int: + in_bytes = actual_rows * hidden * 2 # bf16 read + out_bytes = actual_rows * hidden * out_copies # fp8 write(s) + return in_bytes + out_bytes + + +def _physical_bytes_per_call(actual_rows: int, hidden: int, out_copies: int) -> int: + # Under current scaling, the input is read twice: + # Pass 1: amax kernel reads input (2 bytes per element) + # Pass 2: cast kernel reads input (2 bytes per element) and writes output (1 byte per element per copy) + in_bytes = actual_rows * hidden * 2 * 2 # bf16 read twice + out_bytes = actual_rows * hidden * out_copies # fp8 write(s) + return in_bytes + out_bytes + + +def _kernel_bytes_per_call( + bucket: str, actual_rows: int, hidden: int, out_copies: int, num_tensors: int +) -> int: + """Bytes that one launch of the kernel(s) in ``bucket`` actually moves. + + Counts both reads and writes. Tiny kernels (amax_zero, compute_scale) move + only ``num_tensors``-sized metadata buffers so their reported BW will look + low; that's expected -- they're launch-overhead dominated, not bandwidth + bound. + """ + input_bytes = actual_rows * hidden * 2 # bf16 input + if bucket == "amax_zero": + # Writes ``num_tensors`` floats to zero the amax buffer. + return num_tensors * 4 + if bucket == "amax": + # Reads the full input over the active region (the kernel uses tensor_offsets + # to skip rows past sum(first_dims)) and atomicMaxes into the per-group amax + # slots. Output traffic is negligible vs the input scan. + return input_bytes + num_tensors * 4 + if bucket == "compute_scale": + # Reads ``num_tensors`` amax floats, writes ``num_tensors`` of (scale, + # scale_inv) -- ~3 floats per group. + return num_tensors * 4 * 3 + if bucket == "splits_to_offsets": + # Reads ``num_tensors`` elements, writes ``num_tensors + 1`` elements. + return (num_tensors * 2 + 1) * 8 + if bucket == "cast": + # Reads bf16 input, writes fp8 rowwise and/or columnwise (1 byte per + # element per materialized direction). Mirrors the eager BW_actual byte count. + return _bytes_per_call(actual_rows, hidden, out_copies) + return 0 + + +# Substring patterns used to bucket CUDA kernels in --profile mode. Order matters: +# the first matching bucket wins. Anything that doesn't match falls into "cast". +# Patterns are matched against the kernel name AFTER lowercasing -- both the +# kernel name and the pattern strings are lowercased before substring match, +# so patterns here must already be lowercase. Camelcased identifiers like +# ``ComputeScaleAndScaleInvFunctor`` flatten to ``computescaleandscaleinvfunctor``. +_KERNEL_BUCKETS = ( + # Split amax into the trivial zero-init kernel and the real reduction kernel, + # otherwise the bucket average misleadingly suggests the zero kernel is slow. + ("amax_zero", ("grouped_amax_zero",)), + ("amax", ("grouped_amax",)), + # The compute-scale kernel: the grouped current-scaling path launches + # ``grouped_compute_scale_kernel`` directly, while the legacy per-tensor path + # launches it via the multi_tensor_apply framework as + # ``multi_tensor_apply_kernel<...ComputeScaleAndScaleInvFunctor>``. Match both + # (the underscored kernel name and the flattened functor name) after + # lowercasing. Without the underscored pattern the grouped kernel would fall + # through to the ``cast`` bucket and inflate its reported BW. + ("compute_scale", ("compute_scale", "computescale")), + # Splits to offsets helper kernel + ("splits_to_offsets", ("splits_to_offsets",)), +) + + +def _bucket_for_kernel(name: str) -> str: + lname = name.lower() + for bucket, patterns in _KERNEL_BUCKETS: + if any(p in lname for p in patterns): + return bucket + return "cast" + + +def _profile_breakdown(quantizer, inputs, first_dims, num_groups: int, iters: int) -> dict: + """Run ``iters`` group_quantize calls under torch.profiler and aggregate CUDA + kernel time into {amax, compute_scale, cast} buckets. + + Returns dict: bucket -> {"us_total": float, "calls": int}, plus a "_total_us" + key holding the summed GPU time across all kernels (not wall-clock). + """ + from torch.profiler import profile, ProfilerActivity + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for it in range(iters): + i = it % len(inputs) + tex.group_quantize(inputs[i], quantizer, num_groups, first_dims) + torch.cuda.synchronize() + + agg = {b: {"us_total": 0.0, "calls": 0} for b, _ in _KERNEL_BUCKETS} + agg["cast"] = {"us_total": 0.0, "calls": 0} + total_us = 0.0 + for evt in prof.key_averages(): + # CPU events have cuda_time_total == 0; kernel events have device time. + if evt.device_type.name != "CUDA": + continue + us = float(evt.self_device_time_total) + if us <= 0: + continue + bucket = _bucket_for_kernel(evt.key) + agg[bucket]["us_total"] += us + agg[bucket]["calls"] += evt.count + total_us += us + agg["_total_us"] = total_us + return agg + + +def run_case( + shape_case: str, + mode: str, + *, + actual_rows: int, + allocated_rows: int, + hidden: int, + num_groups: int, + num_buffers: int, + warmup: int, + iters: int, + mode_loop: str = "eager", + verbose: bool = True, + profile_breakdown: bool = False, + print_breakdown: bool = True, +) -> CaseResult: + inputs, first_dims, _ = _make_inputs( + shape_case, actual_rows, allocated_rows, hidden, num_groups, num_buffers + ) + quantizer = _make_quantizer(mode) + + # Run once for correctness sanity check (no tail-read leak). + if shape_case == "varying-first-overalloc": + _verify_no_tail_read(quantizer, inputs[0], first_dims, num_groups, actual_rows) + + # Number of FP8 output buffers actually written (e.g. on Blackwell, "both" + # materializes only the rowwise buffer), measured from a real quantize so the + # BW accounting reflects the bytes the kernels truly move. + out_copies = _produced_out_copies(quantizer, inputs[0], num_groups, first_dims) + + # Warmup so GPU clocks are at boost. + for it in range(warmup): + i = it % len(inputs) + tex.group_quantize(inputs[i], quantizer, num_groups, first_dims) + torch.cuda.synchronize() + + if mode_loop == "graph": + elapsed_ms, actual_iters = _timed_cuda_graph( + quantizer, inputs, first_dims, num_groups, iters + ) + else: + elapsed_ms = _timed_eager(quantizer, inputs, first_dims, num_groups, iters) + actual_iters = iters + + relevant = _bytes_per_call(actual_rows, hidden, out_copies) + physical_relevant = _physical_bytes_per_call(actual_rows, hidden, out_copies) + total_bytes_actual = relevant * actual_iters + total_bytes_physical = physical_relevant * actual_iters + elapsed_s = elapsed_ms / 1000.0 + bw_actual = total_bytes_actual / elapsed_s / 1.0e12 + bw_physical = total_bytes_physical / elapsed_s / 1.0e12 + res = CaseResult( + shape_case=shape_case, + mode=mode, + actual_rows=actual_rows, + allocated_rows=allocated_rows, + hidden=hidden, + num_groups=num_groups, + iters=actual_iters, + elapsed_ms_total=elapsed_ms, + per_iter_us=elapsed_ms * 1000.0 / actual_iters, + relevant_bytes=relevant, + bw_actual_TBps=bw_actual, + bw_physical_TBps=bw_physical, + loop=mode_loop, + ) + if verbose: + print( + f" {shape_case:30s} groups={num_groups:3d} mode={mode:10s} " + f"loop={mode_loop:5s} " + f"per_iter={res.per_iter_us:7.2f}us " + f"BW_algo={res.bw_actual_TBps:5.2f} TB/s " + f"BW_phys={res.bw_physical_TBps:5.2f} TB/s" + ) + + if profile_breakdown: + profile_iters = min(iters, 50) + agg = _profile_breakdown(quantizer, inputs, first_dims, num_groups, iters=profile_iters) + total = agg["_total_us"] if agg["_total_us"] > 0 else 1.0 + + def _bucket_metrics(bucket: str) -> dict: + entry = agg[bucket] + per_iter_us = entry["us_total"] / profile_iters + launches_per_iter = entry["calls"] / profile_iters + per_launch_us = entry["us_total"] / entry["calls"] if entry["calls"] else 0.0 + bytes_per_launch = _kernel_bytes_per_call( + bucket, actual_rows, hidden, out_copies, num_groups + ) + bw_tbps = ( + bytes_per_launch / (per_launch_us * 1.0e-6) / 1.0e12 if per_launch_us > 0 else 0.0 + ) + return { + "per_iter_us": per_iter_us, + "per_launch_us": per_launch_us, + "launches_per_iter": launches_per_iter, + "bytes_per_launch": bytes_per_launch, + "bw_TBps": bw_tbps, + "pct": 100.0 * entry["us_total"] / total, + } + + # Always capture the two kernels of interest so the consolidated table + # can be printed at the end. + res.amax_profile = _bucket_metrics("amax") + res.cast_profile = _bucket_metrics("cast") + + if print_breakdown: + per_iter_total = total / profile_iters + print( + f" kernel breakdown over {profile_iters} iters " + f"(per-iter sum={per_iter_total:6.2f}us):" + ) + print( + f" {'bucket':14s} {'per_iter_us':>12s} {'(%)':>7s} " + f"{'launches/iter':>14s} {'per_launch_us':>15s} " + f"{'bytes/launch':>14s} {'BW_TBps':>9s}" + ) + for bucket in ("amax_zero", "amax", "compute_scale", "splits_to_offsets", "cast"): + m = _bucket_metrics(bucket) + print( + f" {bucket:14s} {m['per_iter_us']:12.2f} " + f"{m['pct']:7.1f} {m['launches_per_iter']:14.2f} " + f"{m['per_launch_us']:15.2f} " + f"{m['bytes_per_launch']:14d} {m['bw_TBps']:9.2f}" + ) + return res + + +def print_kernel_summary(results: List[CaseResult]) -> None: + """Print one clean consolidated table containing only the amax and cast + kernel metrics for every profiled use-case.""" + rows = [r for r in results if r.amax_profile is not None and r.cast_profile is not None] + if not rows: + return + + print() + header = ( + f"{'rows x hidden':16s} {'shape_case':30s} {'groups':>6s} " + f"{'mode':10s} {'loop':6s} " + f"{'amax_us':>9s} {'amax_BW':>9s} {'cast_us':>9s} {'cast_BW':>9s}" + ) + print("=" * len(header)) + print("amax + cast kernel profile (per-iter device time)") + print("=" * len(header)) + print(header) + print("-" * len(header)) + for r in rows: + a = r.amax_profile + c = r.cast_profile + shape = f"{r.actual_rows}x{r.hidden}" + print( + f"{shape:16s} {r.shape_case:30s} {r.num_groups:6d} " + f"{r.mode:10s} {r.loop:6s} " + f"{a['per_iter_us']:9.2f} {a['bw_TBps']:9.2f} " + f"{c['per_iter_us']:9.2f} {c['bw_TBps']:9.2f}" + ) + print("-" * len(header)) + print("amax_us/cast_us = per-iter device time (us); BW in TB/s.") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--actual-rows", type=int, default=98304) + parser.add_argument( + "--allocated-rows", + type=int, + default=None, + help=( + "For varying-first-overalloc; default overalloc * " + "actual-rows. Ignored when --shapes is given." + ), + ) + parser.add_argument("--hidden", type=int, default=2880) + parser.add_argument( + "--shapes", + nargs="+", + default=None, + help=( + "Sweep multiple shapes. Each entry is " + "'rows[xX,:]hidden' (hidden optional, falls back to " + "--hidden). '*' is allowed, e.g. " + "'4096*16x4096 98304x2880'. Overrides " + "--actual-rows/--hidden." + ), + ) + parser.add_argument( + "--overalloc", + type=int, + default=2, + help=( + "Over-allocation factor: allocated_rows = " + "overalloc * actual_rows (default 2). Used for the " + "varying-first-overalloc* shape cases." + ), + ) + parser.add_argument( + "--num-groups", + type=int, + nargs="+", + default=[16, 64], + help="Sweep one or more group counts (default: 16 64).", + ) + parser.add_argument("--num-buffers", type=int, default=4) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--shape-cases", nargs="+", default=list(SHAPE_CASES)) + parser.add_argument("--modes", nargs="+", default=list(MODES)) + parser.add_argument( + "--loop", + choices=("eager", "graph", "both"), + default="both", + help="eager Python loop, captured CUDA graph, or both", + ) + parser.add_argument( + "--profile", + action="store_true", + help=( + "After each timed case, also run a torch.profiler pass " + "and break down CUDA kernel time into " + "{amax, compute_scale, cast} buckets. Adds overhead." + ), + ) + parser.add_argument( + "--profile-table-only", + action="store_true", + help=( + "Implies --profile. Suppress the verbose per-case " + "kernel breakdown and only print the final " + "consolidated amax + cast table." + ), + ) + parser.add_argument("--json-out", default=None) + args = parser.parse_args() + + if args.profile_table_only: + args.profile = True + + # Build the list of (actual_rows, hidden, allocated_rows) shapes to sweep. + if args.shapes: + parsed = [_parse_shape(s, args.hidden) for s in args.shapes] + else: + parsed = [(args.actual_rows, args.hidden)] + shapes = [] + for actual_rows, hidden in parsed: + if args.shapes is None and args.allocated_rows is not None: + allocated_rows = args.allocated_rows + else: + allocated_rows = args.overalloc * actual_rows + shapes.append((actual_rows, hidden, allocated_rows)) + + # Validate that every actual_rows is divisible by every requested group count + # so the equal-first-dims helper produces well-defined splits. + for actual_rows, hidden, _ in shapes: + for ng in args.num_groups: + if actual_rows % ng != 0: + raise SystemExit(f"actual_rows={actual_rows} not divisible by num_groups={ng}") + + print(f"GPU: {torch.cuda.get_device_name(0)}") + shapes_desc = ", ".join(f"{r}x{h}(alloc={a})" for r, h, a in shapes) + print( + f"Config: shapes=[{shapes_desc}], overalloc={args.overalloc}, " + f"num_groups_sweep={args.num_groups}, " + f"iters={args.iters}, warmup={args.warmup}" + ) + print() + + loop_modes = ("eager", "graph") if args.loop == "both" else (args.loop,) + # In table-only mode only the eager loop is profiled, so skip the graph loop + # entirely (it would produce no table rows) and silence the chatty per-case + # output so the final consolidated table stands alone. + if args.profile_table_only: + loop_modes = ("eager",) + quiet = args.profile_table_only + + results: List[CaseResult] = [] + for actual_rows, hidden, allocated_rows in shapes: + if not quiet: + print( + f"################ shape: actual_rows={actual_rows} " + f"hidden={hidden} allocated_rows={allocated_rows} ################" + ) + for num_groups in args.num_groups: + if not quiet: + print(f"#### num_groups={num_groups} ####") + for shape_case in args.shape_cases: + if shape_case not in SHAPE_CASES: + raise SystemExit(f"unknown shape_case={shape_case}") + if not quiet: + print(f"== {shape_case} ==") + for mode in args.modes: + if mode not in MODES: + raise SystemExit(f"unknown mode={mode}") + for loop in loop_modes: + # Only profile-breakdown the eager loop -- profiling a + # captured CUDA graph aggregates everything into the graph + # launch event, which hides per-kernel time. Eager mode + # gives true per-kernel device time. + do_profile = args.profile and loop == "eager" + results.append( + run_case( + shape_case, + mode, + actual_rows=actual_rows, + allocated_rows=allocated_rows, + hidden=hidden, + num_groups=num_groups, + num_buffers=args.num_buffers, + warmup=args.warmup, + iters=args.iters, + mode_loop=loop, + verbose=not quiet, + profile_breakdown=do_profile, + print_breakdown=not args.profile_table_only, + ) + ) + if not quiet: + print() + + if args.profile: + print_kernel_summary(results) + + if args.json_out: + with open(args.json_out, "w") as f: + json.dump([r.__dict__ for r in results], f, indent=2) + print(f"Wrote {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 9b67c09f34..1c4d86a3a8 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -5,6 +5,7 @@ add_executable(test_operator test_cast.cu test_cast_current_scaling.cu + test_cast_current_scaling_grouped.cu test_cast_dbias.cu test_cast_dbias_dgelu.cu test_cast_gated_swiglu.cu diff --git a/tests/cpp/operator/test_cast_current_scaling_grouped.cu b/tests/cpp/operator/test_cast_current_scaling_grouped.cu new file mode 100644 index 0000000000..3076a9e0bc --- /dev/null +++ b/tests/cpp/operator/test_cast_current_scaling_grouped.cu @@ -0,0 +1,468 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, +}; + +enum ScalingDirection { ROWWISE = 0, COLWISE = 1, BOTH = 2 }; + +// Reference per-group amax / scale / scale_inv for FP8 current (tensor) scaling. +// Mirrors compute_scale_from_amax on device (epsilon floor, __fdiv_rn, FLT_MAX +// for non-representable scales) with force_pow_2_scales disabled. +template +void compute_amax_scale_ref(const InputType *data, const size_t size, float *amax_ptr, + float *scale_ptr, float *scale_inv_ptr, const float max_fp8, + const float epsilon) { + float current_max = 0.0f; + for (size_t i = 0; i < size; ++i) { + const float current = static_cast(data[i]); + current_max = fmaxf(current_max, fabsf(current)); + } + *amax_ptr = current_max; + + float clamp_amax = current_max; + if (clamp_amax < epsilon) { + clamp_amax = epsilon; + } + + float scale = 1.0f; + float scale_inv = 1.0f; + if (std::isinf(clamp_amax) || clamp_amax == 0.0f || std::isnan(clamp_amax)) { + *scale_ptr = scale; + *scale_inv_ptr = scale_inv; + return; + } + + scale = max_fp8 / clamp_amax; + if (std::isinf(scale)) { + scale = std::numeric_limits::max(); + } + if (std::isnan(scale)) { + scale = 1.0f; + } + scale_inv = 1.0f / scale; + + *scale_ptr = scale; + *scale_inv_ptr = scale_inv; +} + +// Element-wise comparison that tolerates round-to-nearest ambiguity (a quantized +// value that lands exactly between two representable FP8 values may round either +// way depending on tiny scale differences between host and device). +template +void compare_scaled_elts(const std::string &name, const T *ref_data, const T *test_data, + const size_t numel, const size_t tolerable_mismatches_limit = 0) { + size_t mismatches_num = 0; + int64_t first_mismatch_idx = -1; + for (size_t i = 0; i < numel; ++i) { + const double t = static_cast(test_data[i]); + const double r = static_cast(ref_data[i]); + if (t == r) { + continue; + } + // Allow a single-ULP round-to-nearest disagreement. + const double mean = (t + r) / 2; + const double mean_p = mean >= 0 ? mean * (1 + 1e-6) : mean * (1 - 1e-6); + const double mean_m = mean >= 0 ? mean * (1 - 1e-6) : mean * (1 + 1e-6); + const double cast_mean_p = static_cast(static_cast(mean_p)); + const double cast_mean_m = static_cast(static_cast(mean_m)); + const bool round_ambiguity = + (cast_mean_m == std::min(t, r) && cast_mean_p == std::max(t, r)); + if (round_ambiguity) { + continue; + } + mismatches_num++; + if (first_mismatch_idx == -1) { + first_mismatch_idx = static_cast(i); + } + if (mismatches_num > tolerable_mismatches_limit) { + GTEST_FAIL() << mismatches_num << " mismatch(es) in tensor " << name + << " (limit " << tolerable_mismatches_limit << "). First at " + << first_mismatch_idx << ": " + << static_cast(test_data[first_mismatch_idx]) << " vs " + << static_cast(ref_data[first_mismatch_idx]); + } + } +} + +template +void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors, + const std::vector &logical_shape_vec, + const std::vector &first_dims_h, const std::vector &last_dims_h, + const std::vector &offsets_h, const bool rowwise, const bool colwise) { + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + const size_t cols = logical_shape_vec[1]; + + float max_fp8 = Quantized_Limits::max(); + + size_t elts_num = 0; + for (size_t t = 0; t < num_tensors; ++t) { + elts_num += first_dims_h[t] * last_dims_h[t]; + } + + // Host inputs + std::vector in_data(elts_num); + std::mt19937 gen(12345); + std::uniform_real_distribution<> dis(-2.0, 1.0); + for (size_t i = 0; i < elts_num; ++i) { + in_data[i] = static_cast(dis(gen)); + } + + // Reference amax / scale / scale_inv (one per group) and output data. + std::vector ref_amax(num_tensors, 0.0f); + std::vector ref_scale(num_tensors, 1.0f); + std::vector ref_scale_inv(num_tensors, 1.0f); + std::vector ref_out_rowwise(rowwise ? elts_num : 0, static_cast(0.0f)); + std::vector ref_out_colwise(colwise ? elts_num : 0, static_cast(0.0f)); + + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t K = last_dims_h[t]; + const size_t base = offsets_h[t]; + const size_t group_elts = M * K; + + compute_amax_scale_ref(in_data.data() + base, group_elts, &ref_amax[t], + &ref_scale[t], &ref_scale_inv[t], max_fp8, /*epsilon=*/0.0f); + + const float scale = ref_scale[t]; + for (size_t i = 0; i < M; ++i) { + for (size_t j = 0; j < K; ++j) { + const float val = static_cast(in_data[base + i * K + j]) * scale; + if (rowwise) { + ref_out_rowwise[base + i * K + j] = static_cast(val); + } + if (colwise) { + // Columnwise output of one group is stored transposed: [col * M + row]. + ref_out_colwise[base + j * M + i] = static_cast(val); + } + } + } + } + + // Device allocations + const size_t in_data_size = elts_num * sizeof(InputType); + const size_t out_data_size = elts_num * sizeof(OutputType); + const size_t per_group_f32_size = num_tensors * sizeof(float); + const size_t int64_arr_size = num_tensors * sizeof(int64_t); + const size_t offsets_size = (num_tensors + 1) * sizeof(int64_t); + + // first_dims/last_dims/offsets must be int64 on device. + std::vector first_dims_i64(first_dims_h.begin(), first_dims_h.end()); + std::vector last_dims_i64(last_dims_h.begin(), last_dims_h.end()); + std::vector offsets_i64(offsets_h.begin(), offsets_h.end()); + + InputType *in_data_d = nullptr; + float *amax_d = nullptr; + float *scale_d = nullptr; + float *scale_inv_d = nullptr; + OutputType *out_rowwise_d = nullptr; + OutputType *out_colwise_d = nullptr; + int64_t *first_dims_d = nullptr; + int64_t *last_dims_d = nullptr; + int64_t *offsets_d = nullptr; + + NVTE_CHECK_CUDA(cudaMalloc((void **)&in_data_d, in_data_size)); + NVTE_CHECK_CUDA(cudaMalloc((void **)&amax_d, per_group_f32_size)); + NVTE_CHECK_CUDA(cudaMalloc((void **)&scale_d, per_group_f32_size)); + NVTE_CHECK_CUDA(cudaMalloc((void **)&scale_inv_d, per_group_f32_size)); + NVTE_CHECK_CUDA(cudaMalloc((void **)&first_dims_d, int64_arr_size)); + NVTE_CHECK_CUDA(cudaMalloc((void **)&last_dims_d, int64_arr_size)); + NVTE_CHECK_CUDA(cudaMalloc((void **)&offsets_d, offsets_size)); + + NVTE_CHECK_CUDA(cudaMemcpy(in_data_d, in_data.data(), in_data_size, cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemset(amax_d, 0, per_group_f32_size)); + NVTE_CHECK_CUDA(cudaMemset(scale_d, 0, per_group_f32_size)); + NVTE_CHECK_CUDA(cudaMemset(scale_inv_d, 0, per_group_f32_size)); + NVTE_CHECK_CUDA( + cudaMemcpy(first_dims_d, first_dims_i64.data(), int64_arr_size, cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA( + cudaMemcpy(last_dims_d, last_dims_i64.data(), int64_arr_size, cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(offsets_d, offsets_i64.data(), offsets_size, cudaMemcpyHostToDevice)); + + if (rowwise) { + NVTE_CHECK_CUDA(cudaMalloc((void **)&out_rowwise_d, out_data_size)); + NVTE_CHECK_CUDA(cudaMemset(out_rowwise_d, 0, out_data_size)); + } + if (colwise) { + NVTE_CHECK_CUDA(cudaMalloc((void **)&out_colwise_d, out_data_size)); + NVTE_CHECK_CUDA(cudaMemset(out_colwise_d, 0, out_data_size)); + } + + // Shapes + NVTEShape logical_shape = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size()); + std::vector data_1d_shape = {elts_num}; + NVTEShape data_shape = nvte_make_shape(data_1d_shape.data(), data_1d_shape.size()); + std::vector per_group_shape_vec = {num_tensors}; + NVTEShape per_group_shape = nvte_make_shape(per_group_shape_vec.data(), per_group_shape_vec.size()); + + NVTEShape first_dims_shape; + NVTEShape last_dims_shape; + NVTEShape offsets_shape; + first_dims_shape.ndim = 1; + last_dims_shape.ndim = 1; + offsets_shape.ndim = 1; + first_dims_shape.data[0] = num_tensors; + last_dims_shape.data[0] = num_tensors; + offsets_shape.data[0] = num_tensors + 1; + + // Input grouped tensor (high precision) + NVTEGroupedTensor in_group_tensor = + nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape); + NVTEBasicTensor in_data_tensor = {in_data_d, static_cast(itype), data_shape}; + nvte_set_grouped_tensor_param(in_group_tensor, kNVTEGroupedRowwiseData, &in_data_tensor, + sizeof(in_data_tensor)); + + // Output grouped tensor (FP8 tensor scaling) + NVTEGroupedTensor out_group_tensor = + nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape); + + NVTEBasicTensor amax_tensor = {amax_d, kNVTEFloat32, per_group_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, kNVTEGroupedAmax, &amax_tensor, + sizeof(amax_tensor)); + NVTEBasicTensor scale_tensor = {scale_d, kNVTEFloat32, per_group_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, kNVTEGroupedScale, &scale_tensor, + sizeof(scale_tensor)); + + if (rowwise) { + NVTEBasicTensor out_data_tensor = {out_rowwise_d, static_cast(otype), data_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, kNVTEGroupedRowwiseData, &out_data_tensor, + sizeof(out_data_tensor)); + NVTEBasicTensor scale_inv_tensor = {scale_inv_d, kNVTEFloat32, per_group_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, kNVTEGroupedRowwiseScaleInv, &scale_inv_tensor, + sizeof(scale_inv_tensor)); + } + if (colwise) { + NVTEBasicTensor out_data_tensor = {out_colwise_d, static_cast(otype), data_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, kNVTEGroupedColumnwiseData, &out_data_tensor, + sizeof(out_data_tensor)); + // Current scaling reuses a single scale_inv per group across directions; alias + // the same buffer (this matches the production quantizer layout). + NVTEBasicTensor scale_inv_tensor = {scale_inv_d, kNVTEFloat32, per_group_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, kNVTEGroupedColumnwiseScaleInv, + &scale_inv_tensor, sizeof(scale_inv_tensor)); + } + + // Shape metadata on both grouped tensors. + auto set_shape_metadata = [&](NVTEGroupedTensor gt) { + if (shape_rep == VARYING_FIRST_DIM) { + NVTEBasicTensor t = {first_dims_d, kNVTEInt64, first_dims_shape}; + nvte_set_grouped_tensor_param(gt, kNVTEGroupedFirstDims, &t, sizeof(t)); + } + if (shape_rep == VARYING_LAST_DIM) { + NVTEBasicTensor t = {last_dims_d, kNVTEInt64, last_dims_shape}; + nvte_set_grouped_tensor_param(gt, kNVTEGroupedLastDims, &t, sizeof(t)); + } + if (shape_rep != SAME_BOTH_DIMS) { + NVTEBasicTensor t = {offsets_d, kNVTEInt64, offsets_shape}; + nvte_set_grouped_tensor_param(gt, kNVTEGroupedTensorOffsets, &t, sizeof(t)); + } + }; + set_shape_metadata(in_group_tensor); + set_shape_metadata(out_group_tensor); + + // FP8 current-scaling grouped quantize: + // 1) per-group amax, 2) per-group scale/scale_inv, 3) cast/transpose. + QuantizationConfigWrapper quant_config; + nvte_group_compute_amax_with_config(in_group_tensor, out_group_tensor, quant_config, 0); + nvte_group_compute_scale_from_amax(out_group_tensor, quant_config, 0); + nvte_group_quantize(in_group_tensor, out_group_tensor, quant_config, 0); + + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + // Copy results back. + std::vector amax_h(num_tensors); + std::vector scale_h(num_tensors); + std::vector scale_inv_h(num_tensors); + NVTE_CHECK_CUDA(cudaMemcpy(amax_h.data(), amax_d, per_group_f32_size, cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(scale_h.data(), scale_d, per_group_f32_size, cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA( + cudaMemcpy(scale_inv_h.data(), scale_inv_d, per_group_f32_size, cudaMemcpyDeviceToHost)); + + auto [atol_f32, rtol_f32] = getTolerances(DType::kFloat32); + for (size_t t = 0; t < num_tensors; ++t) { + EXPECT_NEAR(amax_h[t], ref_amax[t], atol_f32 + rtol_f32 * std::fabs(ref_amax[t])) + << "amax mismatch at group " << t; + EXPECT_NEAR(scale_h[t], ref_scale[t], atol_f32 + rtol_f32 * std::fabs(ref_scale[t])) + << "scale mismatch at group " << t; + EXPECT_NEAR(scale_inv_h[t], ref_scale_inv[t], atol_f32 + rtol_f32 * std::fabs(ref_scale_inv[t])) + << "scale_inv mismatch at group " << t; + } + + if (rowwise) { + std::vector out_h(elts_num); + NVTE_CHECK_CUDA( + cudaMemcpy(out_h.data(), out_rowwise_d, out_data_size, cudaMemcpyDeviceToHost)); + compare_scaled_elts("rowwise_output", ref_out_rowwise.data(), out_h.data(), + elts_num); + } + if (colwise) { + std::vector out_h(elts_num); + NVTE_CHECK_CUDA( + cudaMemcpy(out_h.data(), out_colwise_d, out_data_size, cudaMemcpyDeviceToHost)); + compare_scaled_elts("colwise_output", ref_out_colwise.data(), out_h.data(), + elts_num); + } + + nvte_destroy_grouped_tensor(in_group_tensor); + nvte_destroy_grouped_tensor(out_group_tensor); + cudaFree(in_data_d); + cudaFree(amax_d); + cudaFree(scale_d); + cudaFree(scale_inv_d); + cudaFree(first_dims_d); + cudaFree(last_dims_d); + cudaFree(offsets_d); + if (rowwise) cudaFree(out_rowwise_d); + if (colwise) cudaFree(out_colwise_d); +} + +// {shape_representation, num_tensors, [logical_shape_M, logical_shape_K], [M_i] or [K_i]}. +// Constant dimensions must be a small-vector multiple of +// (16 for bf16/fp16, 8 for fp32), so the configs below intentionally use sizes +// that are multiples of 16. +std::vector> input_configs = { + {SAME_BOTH_DIMS, 1, 128, 128}, + {SAME_BOTH_DIMS, 2, 256, 144}, + {SAME_BOTH_DIMS, 3, 384, 80}, + {VARYING_FIRST_DIM, 2, 512, 128, 128, 384}, + {VARYING_FIRST_DIM, 3, 1024, 144, 128, 384, 512}, + {VARYING_FIRST_DIM, 4, 1536, 160, 128, 384, 512, 512}, + {VARYING_FIRST_DIM, 3, 800, 96, 100, 300, 400}, + // Empty tensor in the middle must not break the per-group loop. + {VARYING_FIRST_DIM, 4, 512, 160, 128, 0, 128, 256}, + {VARYING_LAST_DIM, 3, 256, 896, 128, 256, 512}, + {VARYING_LAST_DIM, 2, 160, 384, 128, 256}, + {VARYING_LAST_DIM, 3, 80, 512, 128, 128, 256}, +}; + +std::vector scaling_directions = { + ScalingDirection::ROWWISE, + ScalingDirection::COLWISE, + ScalingDirection::BOTH, +}; + +} // namespace + +class GroupedCastCurrentScalingTestSuite + : public ::testing::TestWithParam, // Config + transformer_engine::DType, // InputType + transformer_engine::DType // OutputType + >> {}; + +TEST_P(GroupedCastCurrentScalingTestSuite, Test) { + if (getDeviceComputeCapability() < hopperComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const ScalingDirection scaling_direction = std::get<0>(GetParam()); + const std::vector config = std::get<1>(GetParam()); + const DType input_type = std::get<2>(GetParam()); + const DType output_type = std::get<3>(GetParam()); + + const ShapeRepresentation shape_rep = static_cast(config[0]); + const size_t num_tensors = config[1]; + const std::vector logical_shape = {config[2], config[3]}; + + std::vector first_dims(num_tensors); + std::vector last_dims(num_tensors); + std::vector offsets(num_tensors + 1, 0); + for (size_t t = 0; t < num_tensors; ++t) { + switch (shape_rep) { + case SAME_BOTH_DIMS: + first_dims[t] = logical_shape[0] / num_tensors; + last_dims[t] = logical_shape[1]; + break; + case VARYING_FIRST_DIM: + first_dims[t] = config[t + 4]; + last_dims[t] = logical_shape[1]; + break; + case VARYING_LAST_DIM: + first_dims[t] = logical_shape[0]; + last_dims[t] = config[t + 4]; + break; + } + offsets[t + 1] = offsets[t] + first_dims[t] * last_dims[t]; + } + + bool rowwise = false; + bool colwise = false; + switch (scaling_direction) { + case ScalingDirection::ROWWISE: rowwise = true; break; + case ScalingDirection::COLWISE: colwise = true; break; + case ScalingDirection::BOTH: rowwise = true; colwise = true; break; + } + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( + input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY( + output_type, OutputType, + performTest(shape_rep, num_tensors, logical_shape, first_dims, + last_dims, offsets, rowwise, colwise););); +} + +std::string MakeGroupedCastCurrentScalingTestName( + const testing::TestParamInfo &info) { + std::string name; + switch (std::get<0>(info.param)) { + case ScalingDirection::ROWWISE: name += "ROWWISE_"; break; + case ScalingDirection::COLWISE: name += "COLWISE_"; break; + case ScalingDirection::BOTH: name += "BIDIMENSIONAL_"; break; + } + + const std::vector input = std::get<1>(info.param); + switch (static_cast(input[0])) { + case ShapeRepresentation::SAME_BOTH_DIMS: name += "SAME_BOTH_DIMS"; break; + case ShapeRepresentation::VARYING_FIRST_DIM: name += "VARYING_FIRST_DIM"; break; + case ShapeRepresentation::VARYING_LAST_DIM: name += "VARYING_LAST_DIM"; break; + } + name += "_N_" + std::to_string(input[1]); + name += "_SHAPE_" + std::to_string(input[2]) + "X" + std::to_string(input[3]); + name += "_" + test::typeName(std::get<2>(info.param)); + name += "_" + test::typeName(std::get<3>(info.param)); + return name; +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, GroupedCastCurrentScalingTestSuite, + ::testing::Combine(::testing::ValuesIn(scaling_directions), ::testing::ValuesIn(input_configs), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)), + MakeGroupedCastCurrentScalingTestName); diff --git a/tests/cpp/operator/test_splits_to_offsets.cu b/tests/cpp/operator/test_splits_to_offsets.cu index faac4b7b6f..f0b72d9a6b 100644 --- a/tests/cpp/operator/test_splits_to_offsets.cu +++ b/tests/cpp/operator/test_splits_to_offsets.cu @@ -78,3 +78,102 @@ INSTANTIATE_TEST_SUITE_P( std::to_string(std::get<1>(info.param)); return name; }); + +namespace { + +// Allocate a device buffer holding `host` stored as `dtype` (int32 or int64). +void *copy_to_device(const std::vector &host, transformer_engine::DType dtype) { + using namespace transformer_engine; + NVTE_CHECK(dtype == DType::kInt32 || dtype == DType::kInt64, + "splits_to_offsets test only supports int32/int64."); + void *dptr = nullptr; + if (dtype == DType::kInt32) { + std::vector tmp(host.begin(), host.end()); + NVTE_CHECK_CUDA(cudaMalloc(&dptr, sizeof(int32_t) * tmp.size())); + NVTE_CHECK_CUDA( + cudaMemcpy(dptr, tmp.data(), sizeof(int32_t) * tmp.size(), cudaMemcpyHostToDevice)); + } else { + NVTE_CHECK_CUDA(cudaMalloc(&dptr, sizeof(int64_t) * host.size())); + NVTE_CHECK_CUDA( + cudaMemcpy(dptr, host.data(), sizeof(int64_t) * host.size(), cudaMemcpyHostToDevice)); + } + return dptr; +} + +// Copy a device buffer of `n` `dtype` (int32 or int64) elements back to host as int64. +std::vector copy_to_host(const void *dptr, size_t n, transformer_engine::DType dtype) { + using namespace transformer_engine; + NVTE_CHECK(dtype == DType::kInt32 || dtype == DType::kInt64, + "splits_to_offsets test only supports int32/int64."); + std::vector out(n); + if (dtype == DType::kInt32) { + std::vector tmp(n); + NVTE_CHECK_CUDA(cudaMemcpy(tmp.data(), dptr, sizeof(int32_t) * n, cudaMemcpyDeviceToHost)); + out.assign(tmp.begin(), tmp.end()); + } else { + NVTE_CHECK_CUDA(cudaMemcpy(out.data(), dptr, sizeof(int64_t) * n, cudaMemcpyDeviceToHost)); + } + return out; +} + +} // namespace + +class SplitsToOffsets2DTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(SplitsToOffsets2DTestSuite, TestSplitsToOffsets2D) { + using namespace transformer_engine; + + const size_t num_tensors = std::get<0>(GetParam()); + const DType dtype = std::get<1>(GetParam()); + + // Generate per-tensor first/last dims. Vary both dimensions so the test + // exercises the 2D prefix sum (offset[i+1] = sum_{j<=i} first_dims[j] * last_dims[j]). + std::vector h_first_dims(num_tensors); + std::vector h_last_dims(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + h_first_dims[i] = static_cast((i % 17) + 1); + h_last_dims[i] = static_cast((i % 5) + 1) * 16; + } + + std::vector h_expected(num_tensors + 1, 0); + for (size_t i = 0; i < num_tensors; ++i) { + h_expected[i + 1] = h_expected[i] + h_first_dims[i] * h_last_dims[i]; + } + + void *d_first_dims = copy_to_device(h_first_dims, dtype); + void *d_last_dims = copy_to_device(h_last_dims, dtype); + + std::vector h_output_init(num_tensors + 1, -1); + void *d_output = copy_to_device(h_output_init, dtype); + + TensorWrapper first_dims_w(d_first_dims, std::vector{num_tensors}, dtype); + TensorWrapper last_dims_w(d_last_dims, std::vector{num_tensors}, dtype); + TensorWrapper output_w(d_output, std::vector{num_tensors + 1}, dtype); + + nvte_splits_to_offsets_2d(first_dims_w.data(), last_dims_w.data(), output_w.data(), + 0 /* stream */); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + + std::vector h_output = copy_to_host(d_output, num_tensors + 1, dtype); + + NVTE_CHECK_CUDA(cudaFree(d_first_dims)); + NVTE_CHECK_CUDA(cudaFree(d_last_dims)); + NVTE_CHECK_CUDA(cudaFree(d_output)); + + for (size_t i = 0; i < h_output.size(); ++i) { + EXPECT_EQ(h_output[i], h_expected[i]) + << "Mismatch at index " << i << ": expected " << h_expected[i] << ", got " << h_output[i]; + } +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, SplitsToOffsets2DTestSuite, + ::testing::Combine(::testing::ValuesIn(splits_to_offsets_num_tensors), + ::testing::Values(transformer_engine::DType::kInt32, + transformer_engine::DType::kInt64)), + [](const testing::TestParamInfo &info) { + std::string name = + std::to_string(std::get<0>(info.param)) + "X" + test::typeName(std::get<1>(info.param)); + return name; + }); diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 235825fb37..e1e15e6875 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -4,7 +4,8 @@ """Tests for GroupedTensor class""" -from typing import List, Tuple +from typing import List, Optional, Tuple +import os import pytest import torch import transformer_engine.pytorch as te @@ -18,6 +19,7 @@ NVFP4Quantizer, ) from transformer_engine.pytorch.constants import TE_DType_To_Torch +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported import transformer_engine_torch as tex # Import test utilities @@ -443,46 +445,81 @@ def test_quantize_varying_shapes(self, quantization: str) -> None: cumulative_numel += tensor_shape[0] * tensor_shape[1] @pytest.mark.parametrize( - "shape", - [[(256, 512), (512, 512), (768, 512)], [(512, 512), (512, 512), (512, 512)]], + "shape_case", + ["varying_first", "varying_last", "varying_both"], ) @pytest.mark.parametrize("output_dbias", [False, True]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) - def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]], output_dbias: bool) -> None: - """Test grouped quantization for MXFP8 against per-tensor quantization.""" - # Test wont pass until the grouped quantization PR from Oleg is merged. - num_tensors = 2 - shape = [(512, 1024) for _ in range(num_tensors)] - - # Create BF16 input tensors and pack into a 2D tensor - input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] - grouped_input = torch.cat(input_tensors, dim=0) + def test_quantize_grouped_mxfp8(self, shape_case: str, output_dbias: bool) -> None: + """Test grouped MXFP8 quantization against per-tensor quantization for + varying first/last/both dimensions.""" + if output_dbias and shape_case != "varying_first": + pytest.skip("bgrad_group_quantize requires constant last dimension") + + # Per-tensor shapes are chosen to satisfy MXFP8 alignment requirements: + # - rowwise scale block size 32 -> last dim must be a multiple of 32 + # - kernel chunk size 128 along the first dim -> first dim must be a + # multiple of 128 (per-tensor for VARYING_BOTH_DIMS, otherwise the + # logical first dim). + if shape_case == "varying_first": + per_tensor_shapes = [(128, 512), (256, 512), (384, 512)] + elif shape_case == "varying_last": + per_tensor_shapes = [(512, 128), (512, 256), (512, 384)] + else: # varying_both + per_tensor_shapes = [(128, 256), (256, 512), (384, 384)] + + num_tensors = len(per_tensor_shapes) + + # Each tensor occupies a contiguous chunk of a flat buffer; the kernel + # locates each chunk via tensor_offsets, so the 2D view below only needs + # to encode the correct total number of elements. + input_tensors = [ + torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in per_tensor_shapes + ] + flat_buffer = torch.cat([t.reshape(-1) for t in input_tensors]) + + first_dims_host: Optional[List[int]] + last_dims_host: Optional[List[int]] + if shape_case == "varying_first": + first_dims_host = [s[0] for s in per_tensor_shapes] + last_dims_host = None + common_last = per_tensor_shapes[0][1] + grouped_input = flat_buffer.view(sum(first_dims_host), common_last) + elif shape_case == "varying_last": + first_dims_host = None + last_dims_host = [s[1] for s in per_tensor_shapes] + common_first = per_tensor_shapes[0][0] + grouped_input = flat_buffer.view(common_first, sum(last_dims_host)) + else: # varying_both + first_dims_host = [s[0] for s in per_tensor_shapes] + last_dims_host = [s[1] for s in per_tensor_shapes] + grouped_input = flat_buffer.view(1, -1) + + first_dims = ( + torch.tensor(first_dims_host, dtype=torch.int64, device="cuda") + if first_dims_host is not None + else None + ) + last_dims = ( + torch.tensor(last_dims_host, dtype=torch.int64, device="cuda") + if last_dims_host is not None + else None + ) - # Create MXFP8 output grouped tensor (rowwise only for easier validation) - quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) quantizer.set_usage(rowwise=True, columnwise=False) - first_dims = torch.tensor( - [shape[0][0] for _ in range(num_tensors)], - dtype=torch.int64, - device="cuda", - ) - # Quantize using grouped API if output_dbias: grouped_output, dbias = tex.bgrad_group_quantize( - grouped_input, - quantizer, - num_tensors, - first_dims, + grouped_input, quantizer, num_tensors, first_dims, last_dims ) else: grouped_output = tex.group_quantize( - grouped_input, - quantizer, - num_tensors, - first_dims, + grouped_input, quantizer, num_tensors, first_dims, last_dims ) - # Build expected output by quantizing each tensor independently + + # Reference: quantize each tensor independently and concatenate the + # rowwise data / scale_inv buffers in tensor order. expected_data = [] expected_scale_inv = [] for tensor in input_tensors: @@ -579,75 +616,304 @@ def test_bgrad_group_quantize_zero_size_tensor(self) -> None: assert dbias.shape == (num_tensors, last_dim) assert torch.all(dbias == 0) + @pytest.mark.parametrize( + "quantization", + [ + pytest.param( + "fp8_current_scaling", + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), + ), + ], + ) @pytest.mark.parametrize("output_dbias", [False, True]) - @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) - def test_group_quantize_cudagraph_capturable(self, output_dbias: bool) -> None: + @pytest.mark.parametrize("shape_case", ["varying_first", "varying_last"]) + def test_group_quantize_cudagraph_capturable( + self, quantization: str, output_dbias: bool, shape_case: str + ) -> None: """Ensure group_quantize is CUDA graph capturable.""" - num_tensors = 2 - shape = [(512, 1024) for _ in range(num_tensors)] - input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] - grouped_input = torch.cat(input_tensors, dim=0) - - quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + if output_dbias and quantization != "mxfp8": + pytest.skip("bgrad_group_quantize only supports MXFP8") + if output_dbias and shape_case == "varying_last": + pytest.skip("bgrad_group_quantize does not accept last_dims") + + if shape_case == "varying_last": + rows = 128 + last_dims_host = [256, 128, 384] + num_tensors = len(last_dims_host) + total_cols = sum(last_dims_host) + flat_input = torch.empty(rows * total_cols, dtype=torch.bfloat16, device="cuda") + offset = 0 + for cols in last_dims_host: + member = torch.randn(rows, cols, dtype=torch.bfloat16, device="cuda") + flat_input[offset : offset + member.numel()].copy_(member.reshape(-1)) + offset += member.numel() + grouped_input = flat_input.view(rows, total_cols) + first_dims = None + last_dims = torch.tensor(last_dims_host, dtype=torch.int64, device="cuda") + else: + first_dims_host = [256, 128, 384] + num_tensors = len(first_dims_host) + hidden = 1024 + shape = [(r, hidden) for r in first_dims_host] + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + grouped_input = torch.cat(input_tensors, dim=0) + first_dims = torch.tensor(first_dims_host, dtype=torch.int64, device="cuda") + last_dims = None + + if quantization == "mxfp8": + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + else: + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + force_pow_2_scales=False, + amax_epsilon=0.0, + ) quantizer.set_usage(rowwise=True, columnwise=False) - first_dims = torch.tensor( - [shape[0][0] for _ in range(num_tensors)], - dtype=torch.int64, - device="cuda", - ) torch.cuda.synchronize() static_input = grouped_input.clone() - static_first_dims = first_dims.clone() + static_first_dims = first_dims.clone() if first_dims is not None else None + static_last_dims = last_dims.clone() if last_dims is not None else None + + def _run_group_quantize(input_tensor): + """Return (output, dbias) where dbias is None when output_dbias is False.""" + if output_dbias: + out, dbias = tex.bgrad_group_quantize( + input_tensor, quantizer, num_tensors, static_first_dims + ) + return out, dbias + out = tex.group_quantize( + input_tensor, + quantizer, + num_tensors, + static_first_dims, + last_dims=static_last_dims, + ) + return out, None # Warmup to initialize kernels and allocator state - if output_dbias: - _ = tex.bgrad_group_quantize(static_input, quantizer, num_tensors, static_first_dims) - else: - _ = tex.group_quantize(static_input, quantizer, num_tensors, static_first_dims) + _ = _run_group_quantize(static_input) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - if output_dbias: - static_output, static_dbias = tex.bgrad_group_quantize( - static_input, - quantizer, - num_tensors, - static_first_dims, - ) - else: - static_output = tex.group_quantize( - static_input, - quantizer, - num_tensors, - static_first_dims, - ) + static_output, static_dbias = _run_group_quantize(static_input) - fresh_input = torch.cat( - [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape], - dim=0, - ) + fresh_input = torch.randn_like(grouped_input) static_input.copy_(fresh_input) graph.replay() torch.cuda.synchronize() - if output_dbias: - expected_out, expected_dbias = tex.bgrad_group_quantize( - static_input, - quantizer, - num_tensors, - static_first_dims, - ) - else: - expected_out = tex.group_quantize( - static_input, quantizer, num_tensors, static_first_dims - ) + expected_out, expected_dbias = _run_group_quantize(static_input) assert torch.equal(static_output.rowwise_data, expected_out.rowwise_data) assert torch.equal(static_output.scale_inv, expected_out.scale_inv) if output_dbias: assert torch.allclose(static_dbias, expected_dbias) + @pytest.mark.parametrize("mode", ["rowwise", "columnwise", "both"]) + @pytest.mark.parametrize( + "shape_case", + [ + "uniform", + "varying_first", + "empty_split", + "varying_last", + ], + ) + @pytest.mark.parametrize("overallocated", [False, True]) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_group_quantize_fp8_current_scaling( + self, + mode: str, + shape_case: str, + overallocated: bool, + ) -> None: + """Test grouped FP8 current scaling matches per-tensor current scaling + across shape topologies: + + - ``uniform``: no ``first_dims``/``last_dims`` (kernel partitions implicitly). + - ``varying_first``: ``first_dims`` set, values vary. + - ``empty_split``: ``first_dims`` set with one zero entry. + - ``varying_last``: ``last_dims`` set, values vary. + + When ``overallocated`` is True the backing buffer is twice the active size + in the test case, so the kernel sees an active region followed by an unused tail. + The unused tail elements are poisoned with a large sentinel value (1e4); + since the per-group amax assertion compares against amax computed over the + active input tensors only, any tail read by the kernel would explode the + per-group amax and the assertion would fail. Overallocation is skipped for + ``uniform`` because it partitions the buffer implicitly (no metadata-defined + active region to over-allocate against). + """ + if overallocated and shape_case == "uniform": + pytest.skip( + "Overallocation is not meaningful for ``uniform`` " + "(the kernel partitions the buffer implicitly, so there is no " + "metadata-defined active region to over-allocate against)." + ) + + # Per-tensor shapes for each shape_case. + if shape_case == "uniform": + per_tensor_shapes = [(128, 256)] * 3 + first_dims_host = None + last_dims_host = None + elif shape_case == "varying_first": + # Only common dimension needs to be 16B aligned. + first_dims_host = [65, 128, 97] + last_dims_host = None + per_tensor_shapes = [(r, 256) for r in first_dims_host] + elif shape_case == "empty_split": + first_dims_host = [128, 0, 96] + last_dims_host = None + per_tensor_shapes = [(r, 256) for r in first_dims_host] + elif shape_case == "varying_last": + first_dims_host = None + last_dims_host = [513, 1027, 259] + per_tensor_shapes = [(256, c) for c in last_dims_host] + else: + raise ValueError(f"Unknown shape_case: {shape_case}") + + num_tensors = len(per_tensor_shapes) + first_dims = ( + torch.tensor(first_dims_host, dtype=torch.int64, device="cuda") + if first_dims_host is not None + else None + ) + last_dims = ( + torch.tensor(last_dims_host, dtype=torch.int64, device="cuda") + if last_dims_host is not None + else None + ) + + # Per-tensor data + flat buffer (tensor i occupies a contiguous chunk). + input_tensors = [ + torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in per_tensor_shapes + ] + actual_numel = sum(t.numel() for t in input_tensors) + allocated_numel = actual_numel * 2 if overallocated else actual_numel + flat_buffer = torch.empty(allocated_numel, dtype=torch.bfloat16, device="cuda") + offset = 0 + for t in input_tensors: + flat_buffer[offset : offset + t.numel()].copy_(t.reshape(-1)) + offset += t.numel() + if overallocated: + flat_buffer[actual_numel:].fill_(10000.0) + + # View flat buffer as the 2D shape expected by group_quantize. + if shape_case in ("varying_last",): + common_first = per_tensor_shapes[0][0] + allocated_last = allocated_numel // common_first + grouped_input = flat_buffer.view(common_first, allocated_last) + else: + common_last = per_tensor_shapes[0][1] + allocated_first = allocated_numel // common_last + grouped_input = flat_buffer.view(allocated_first, common_last) + + requested_rowwise = mode in ("rowwise", "both") + requested_columnwise = mode in ("columnwise", "both") + rowwise = requested_rowwise or (requested_columnwise and is_non_tn_fp8_gemm_supported()) + columnwise = requested_columnwise and not is_non_tn_fp8_gemm_supported() + + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + force_pow_2_scales=False, + amax_epsilon=0.0, + ) + quantizer.set_usage(rowwise=requested_rowwise, columnwise=requested_columnwise) + + grouped_output = tex.group_quantize( + grouped_input, quantizer, num_tensors, first_dims, last_dims=last_dims + ) + + # Metadata validation for the varying-last code path. + if shape_case in ("varying_last",): + assert grouped_output.first_dims is None + assert torch.equal(grouped_output.last_dims, last_dims) + + # When ``overallocated`` is True, the input has poisoned rows past + # sum(first_dims) (filled with 1e4). If the kernel were to read those + # tail rows, per-group amax would explode well above what bf16 N(0,1) + # produces. We catch that implicitly via the per-group amax assertion in + # ``_assert_fp8_cs_group_quantize_matches_reference`` below. + self._assert_fp8_cs_group_quantize_matches_reference( + grouped_output=grouped_output, + input_tensors=input_tensors, + rowwise=rowwise, + columnwise=columnwise, + ) + + @staticmethod + def _assert_fp8_cs_group_quantize_matches_reference( + *, + grouped_output, + input_tensors: List[torch.Tensor], + rowwise: bool, + columnwise: bool, + ) -> None: + """Validate amax/scale/scale_inv/per-tensor data of an FP8 current-scaling + grouped output against per-tensor Float8CurrentScalingQuantizer references.""" + expected_amax = torch.stack( + [ + ( + tensor.abs().max().float() + if tensor.numel() > 0 + else torch.zeros((), dtype=torch.float32, device="cuda") + ) + for tensor in input_tensors + ] + ) + torch.testing.assert_close(grouped_output.amax, expected_amax, rtol=0.0, atol=0.0) + + scale_inv = grouped_output.scale_inv + if scale_inv is None: + scale_inv = grouped_output.columnwise_scale_inv + torch.testing.assert_close( + scale_inv, + torch.reciprocal(grouped_output.scale), + rtol=1e-6, + atol=1e-6, + ) + if rowwise and columnwise: + torch.testing.assert_close( + grouped_output.columnwise_scale_inv, + grouped_output.scale_inv, + rtol=0.0, + atol=0.0, + ) + + expected_rowwise = [] + expected_columnwise = [] + for tensor in input_tensors: + if tensor.numel() == 0: + continue + ref_quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + force_pow_2_scales=False, + amax_epsilon=0.0, + ) + ref = ref_quantizer(tensor) + ref_rowwise_data = ref._data.reshape(tensor.shape) + if rowwise: + expected_rowwise.append(ref_rowwise_data.reshape(-1)) + if columnwise: + expected_columnwise.append(ref_rowwise_data.T.contiguous().reshape(-1)) + + if rowwise and expected_rowwise: + expected = torch.cat(expected_rowwise) + assert torch.equal(grouped_output.rowwise_data[: expected.numel()], expected) + if columnwise and expected_columnwise: + expected = torch.cat(expected_columnwise) + assert torch.equal(grouped_output.columnwise_data[: expected.numel()], expected) + @pytest.mark.parametrize( "shape", [[(512, 1024), (512, 1024)], [(256, 512), (512, 512), (768, 512)]], @@ -758,7 +1024,8 @@ def test_grouped_linear_load_state_dict_multi_to_single_param(self, tmp_path) -> in_features = 64 out_features = 32 dtype = torch.float32 - + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0": + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") src = te.GroupedLinear( num_gemms=num_gemms, in_features=in_features, @@ -809,6 +1076,8 @@ def test_grouped_linear_load_state_dict_multi_to_single_param(self, tmp_path) -> def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) -> None: """Load grouped-parameter checkpoint from disk into per-GEMM parameter format.""" + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0": + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") num_gemms = 3 in_features = 64 out_features = 32 diff --git a/transformer_engine/common/cast/cast_grouped.cu b/transformer_engine/common/cast/cast_grouped.cu index 853634c811..c48d94666e 100644 --- a/transformer_engine/common/cast/cast_grouped.cu +++ b/transformer_engine/common/cast/cast_grouped.cu @@ -8,10 +8,12 @@ #include #include #include +#include #include "../common.h" #include "dispatch/dequantize.cuh" #include "dispatch/quantize.cuh" +#include "fp8/group_amax_fp8.cuh" void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, const NVTEQuantizationConfig quant_config, cudaStream_t stream) { @@ -30,6 +32,54 @@ void nvte_group_dequantize(const NVTEGroupedTensor input, NVTEGroupedTensor outp convertNVTEGroupedTensorCheck(output), stream); } +void nvte_group_compute_scale_from_amax(NVTEGroupedTensor output, + const NVTEQuantizationConfig config, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_compute_scale_from_amax); + using namespace transformer_engine; + + NVTE_CHECK(output != nullptr, "Invalid grouped output tensor (got NULL)"); + auto &out = *convertNVTEGroupedTensorCheck(output); + const size_t num_tensors = out.num_tensors; + if (num_tensors == 0) { + return; + } + NVTE_CHECK( + is_fp8_dtype(out.dtype()), + "Grouped scale update requires an FP8 output tensor, but got dtype=", to_string(out.dtype())); + + // FP8 current scaling keeps a single amax/scale/scale_inv per group; the + // scale_inv buffer is aliased by both directions, so pick whichever is set. + float *amax_ptr = reinterpret_cast(out.amax.dptr != nullptr ? out.amax.dptr + : out.columnwise_amax.dptr); + NVTE_CHECK(amax_ptr != nullptr, "Grouped scale update requires an amax buffer."); + float *scale_ptr = reinterpret_cast(out.scale.dptr); + NVTE_CHECK(scale_ptr != nullptr, "Grouped scale update requires a scale buffer."); + float *scale_inv_ptr = reinterpret_cast( + out.scale_inv.dptr != nullptr ? out.scale_inv.dptr : out.columnwise_scale_inv.dptr); + NVTE_CHECK(scale_inv_ptr != nullptr, "Grouped scale update requires a scale_inv buffer."); + + float max_fp8 = 0.f; + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + out.dtype(), DType, + max_fp8 = Quantized_Limits::max_norm;); // NOLINT(*) + + bool force_pow_2_scales = false; + float epsilon = 0.f; + float *noop_ptr = nullptr; + if (config != nullptr) { + const auto *config_cpp = reinterpret_cast(config); + force_pow_2_scales = config_cpp->force_pow_2_scales; + epsilon = config_cpp->amax_epsilon; + const NVTETensor noop = config_cpp->noop_tensor; + noop_ptr = reinterpret_cast(noop != nullptr ? convertNVTETensorCheck(noop)->data.dptr + : nullptr); + } + + dispatch::fp8::launch_grouped_compute_scale_kernel(amax_ptr, scale_ptr, scale_inv_ptr, + num_tensors, max_fp8, force_pow_2_scales, + epsilon, noop_ptr, stream); +} + // Group quantize assumes contiguous inputs and outputs in memory allocation. // Note: this API assumes knowing split sections from the host. If split information // comes from D2H copy, it will break cuda graph capture. diff --git a/transformer_engine/common/cast/core/common.cuh b/transformer_engine/common/cast/core/common.cuh index 3e6eb55b73..84749b0d0e 100644 --- a/transformer_engine/common/cast/core/common.cuh +++ b/transformer_engine/common/cast/core/common.cuh @@ -6,577 +6,22 @@ /*! \file common.cuh * \brief Common functions in quantize. + * + * Umbrella header. The contents are split into: + * - grouped_layout.cuh: architecture-neutral work-decomposition helpers + * (offset/tensor-id lookup, job/block descriptors, dbias reductions). + * - grouped_tma.cuh: architecture-specific TMA descriptor management and + * bulk-copy staging (pulls in arch-specific PTX via ptx.cuh). + * + * Sources that only need the arch-neutral helpers should include + * grouped_layout.cuh directly so they are not forced into arch-specific + * (smXXXa/smXXXf) compilation. */ #ifndef TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ #define TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ -#include -#include -#include -#include - -#include "../../common.h" -#include "../../util/ptx.cuh" -#include "../../utils.cuh" - -namespace transformer_engine { -namespace dispatch { -namespace common { - -constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; - -struct alignas(128) TensorMapStorage { - alignas(128) CUtensorMap input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; - alignas(128) CUtensorMap act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; - alignas(128) CUtensorMap output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; - alignas(128) CUtensorMap output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; -}; - -// Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs. -static __device__ TensorMapStorage g_tensor_maps; - -inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_block) { - const size_t N = product(t->data.shape); - const bool isFullTile = (N % elems_per_block == 0); - return isFullTile; -} - -inline bool dimensions_supported_by_TMA(const Tensor *const t) { - const size_t cols = t->flat_last_dim(); - constexpr size_t TMA_bytes = 16; - const size_t alignment_requirement = (TMA_bytes * 8) / typeToNumBits(t->dtype()); - return cols % alignment_requirement == 0; -} - -__device__ __forceinline__ unsigned char *align_smem_ptr_per_TMA_requirements(unsigned char *p) { - size_t addr = reinterpret_cast(p); - addr = (addr + TMA_SHMEM_ALIGNMENT - 1) & ~(TMA_SHMEM_ALIGNMENT - 1); - return reinterpret_cast(addr); -} - -namespace kernel { - -constexpr size_t THREADS_PER_BLOCK = 256; -template -__global__ void __launch_bounds__(THREADS_PER_BLOCK) - reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial, - const size_t rows, const size_t cols) { - using ComputeVec = Vec; - using OutputVec = Vec; - - const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; - - if (thread_id * nvec >= cols) { - return; - } - - const float *const thread_in_base = dbias_partial + thread_id * nvec; - OType *const thread_out_base = dbias_output + thread_id * nvec; - - ComputeVec ldg_vec; - ComputeVec acc_vec; - acc_vec.clear(); - for (int i = 0; i < rows; ++i) { - ldg_vec.load_from(thread_in_base + i * cols); -#pragma unroll - for (int e = 0; e < nvec; ++e) { - acc_vec.data.elt[e] += ldg_vec.data.elt[e]; - } - } - - OutputVec stg_vec; -#pragma unroll - for (int e = 0; e < nvec; ++e) { - stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); - } - stg_vec.store_to(thread_out_base); -} - -template -__global__ void __launch_bounds__(THREADS_PER_BLOCK) - group_reduce_dbias_kernel(const ShapeRepresentation shape_rep, const size_t num_tensors, - const size_t first_logical_dim, const size_t last_logical_dim, - const int64_t *const offsets_ptr, const int64_t *const first_dims_ptr, - const int64_t *const last_dims_ptr, OType *const dbias_output, - const float *dbias_partial, const size_t chunk_dim_Y) { - using ComputeVec = Vec; - using OutputVec = Vec; - - const size_t tensor_id = blockIdx.y; - const size_t tensor_rows = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) - ? (first_logical_dim / num_tensors) - : static_cast(first_dims_ptr[tensor_id]); - - const size_t rows = tensor_rows / chunk_dim_Y; - const size_t cols = last_logical_dim; - - const size_t dbias_in_offset_Y = - (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) - ? (tensor_id * (tensor_rows / chunk_dim_Y)) - : (static_cast(offsets_ptr[tensor_id]) / cols / chunk_dim_Y); - - const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; - - if (thread_id * nvec >= cols) { - return; - } - - const float *const thread_in_base = dbias_partial + dbias_in_offset_Y * cols + thread_id * nvec; - OType *const thread_out_base = dbias_output + tensor_id * cols + thread_id * nvec; - - ComputeVec ldg_vec; - ComputeVec acc_vec; - acc_vec.clear(); - for (int i = 0; i < rows; ++i) { - ldg_vec.load_from(thread_in_base + i * cols); -#pragma unroll - for (int e = 0; e < nvec; ++e) { - acc_vec.data.elt[e] += ldg_vec.data.elt[e]; - } - } - - OutputVec stg_vec; -#pragma unroll - for (int e = 0; e < nvec; ++e) { - stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); - } - stg_vec.store_to(thread_out_base); -} -} // namespace kernel - -template -void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, const size_t cols, - cudaStream_t stream) { - using namespace kernel; - constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 - constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); - - NVTE_CHECK(cols % reduce_dbias_nvec == 0, "Unsupported shape."); - const size_t reduce_dbias_num_blocks = DIVUP(cols, THREADS_PER_BLOCK * reduce_dbias_nvec); - - reduce_dbias_kernel - <<>>( - reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); - NVTE_CHECK_CUDA(cudaGetLastError()); -} - -template -void grouped_reduce_dbias(const ShapeRepresentation shape_rep, const size_t num_tensors, - const size_t first_logical_dim, const size_t last_logical_dim, - const int64_t *const data_tensor_offsets_ptr, - const int64_t *const data_tensor_first_dims_ptr, - const int64_t *const data_tensor_last_dims_ptr, GroupedTensor *dbias, - const float *workspace_ptr, const size_t chunk_dim_Y, - cudaStream_t stream) { - using namespace kernel; - constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 - constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); - - NVTE_CHECK(last_logical_dim % reduce_dbias_nvec == 0, "Unsupported shape."); - - const size_t blocks_X = DIVUP(last_logical_dim, THREADS_PER_BLOCK * reduce_dbias_nvec); - const size_t blocks_Y = num_tensors; - const dim3 grid(blocks_X, blocks_Y); - - group_reduce_dbias_kernel<<>>( - shape_rep, num_tensors, first_logical_dim, last_logical_dim, data_tensor_offsets_ptr, - data_tensor_first_dims_ptr, data_tensor_last_dims_ptr, - reinterpret_cast(dbias->data.dptr), workspace_ptr, chunk_dim_Y); - - NVTE_CHECK_CUDA(cudaGetLastError()); -} - -template -__device__ __forceinline__ size_t -get_current_tensor_id(const size_t num_tensors, const size_t current_offset, const size_t block_Y, - const size_t first_logical_dim, const size_t last_logical_dim, - const int64_t *const __restrict__ offsets_ptr) { - if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { - const size_t current_row = block_Y * CHUNK_DIM_Y; - const size_t rows_per_tensor = first_logical_dim / num_tensors; - return current_row / rows_per_tensor; - } else { - size_t low = 1; - size_t hi = num_tensors; // [low, hi] - - while (low < hi) { - const size_t mid = low + (hi - low) / 2; - const size_t mid_offset = static_cast(offsets_ptr[mid]); - - if (mid_offset <= current_offset) { - low = mid + 1; - } else { - hi = mid; - } - } - return low - 1; - } -} - -template -__device__ __forceinline__ size_t -get_tensor_rows_num(const size_t tensor_id, const size_t first_logical_dim, - const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { - size_t rows_num = 0; - if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || - SHAPE_REP == ShapeRepresentation::VARYING_LAST_DIM) { - rows_num = first_logical_dim; - } else { - rows_num = static_cast(first_dims_ptr[tensor_id]); - } - if (rows_num % 128 != 0) { - NVTE_DEVICE_ERROR("First dimension of each tensor in a group must be divisible by 128."); - } - return rows_num; -} - -__device__ __forceinline__ size_t get_tensor_rows_num( - const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t first_logical_dim, - const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { - switch (shape_rep) { - case ShapeRepresentation::SAME_BOTH_DIMS: - return get_tensor_rows_num(tensor_id, first_logical_dim, - first_dims_ptr, num_tensors); - case ShapeRepresentation::VARYING_FIRST_DIM: - return get_tensor_rows_num( - tensor_id, first_logical_dim, first_dims_ptr, num_tensors); - case ShapeRepresentation::VARYING_LAST_DIM: - return get_tensor_rows_num( - tensor_id, first_logical_dim, first_dims_ptr, num_tensors); - case ShapeRepresentation::VARYING_BOTH_DIMS: - return get_tensor_rows_num( - tensor_id, first_logical_dim, first_dims_ptr, num_tensors); - } - return 0; -} - -template -__device__ __forceinline__ size_t -get_tensor_cols_num(const size_t tensor_id, const size_t last_logical_dim, - const int64_t *const __restrict__ last_dims_ptr) { - size_t cols_num = 0; - if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || - SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { - cols_num = last_logical_dim; - } else { - cols_num = static_cast(last_dims_ptr[tensor_id]); - if (cols_num % 128 != 0) { - NVTE_DEVICE_ERROR( - "For varying last dimensions support, the last dimension of each tensor in a group " - "must be divisible by 128."); - } - } - return cols_num; -} - -__device__ __forceinline__ size_t get_tensor_cols_num( - const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t last_logical_dim, - const int64_t *const __restrict__ last_dims_ptr) { - switch (shape_rep) { - case ShapeRepresentation::SAME_BOTH_DIMS: - return get_tensor_cols_num(tensor_id, last_logical_dim, - last_dims_ptr); - case ShapeRepresentation::VARYING_FIRST_DIM: - return get_tensor_cols_num( - tensor_id, last_logical_dim, last_dims_ptr); - case ShapeRepresentation::VARYING_LAST_DIM: - return get_tensor_cols_num(tensor_id, last_logical_dim, - last_dims_ptr); - case ShapeRepresentation::VARYING_BOTH_DIMS: - return get_tensor_cols_num( - tensor_id, last_logical_dim, last_dims_ptr); - } - return 0; -} - -// Logical work-item decoded from CTA coordinates. -struct JobDescriptor { - size_t block_id = 0; - size_t block_global_offset = 0; - size_t tensor_id = 0; - size_t rows = 0; - size_t cols = 0; - - __host__ __device__ __forceinline__ constexpr JobDescriptor() = default; - - __host__ __device__ __forceinline__ constexpr JobDescriptor(const size_t block_id_, - const size_t block_global_offset_, - const size_t tensor_id_, - const size_t rows_, - const size_t cols_) - : block_id(block_id_), - block_global_offset(block_global_offset_), - tensor_id(tensor_id_), - rows(rows_), - cols(cols_) {} -}; - -// Tensor-local coordinates for a work-item. -struct BlockDescriptor { - size_t tensor_base = 0; - size_t block_id_in_current_tensor = 0; - size_t block_id_Y = 0; - size_t block_id_X = 0; - size_t block_offset_Y = 0; - size_t block_offset_X = 0; - - __host__ __device__ __forceinline__ constexpr BlockDescriptor() = default; - - __host__ __device__ __forceinline__ constexpr BlockDescriptor( - const size_t tensor_base_, const size_t block_id_in_current_tensor_, const size_t block_id_Y_, - const size_t block_id_X_, const size_t block_offset_Y_, const size_t block_offset_X_) - : tensor_base(tensor_base_), - block_id_in_current_tensor(block_id_in_current_tensor_), - block_id_Y(block_id_Y_), - block_id_X(block_id_X_), - block_offset_Y(block_offset_Y_), - block_offset_X(block_offset_X_) {} -}; - -template -__device__ __forceinline__ JobDescriptor decode_job( - const size_t num_tensors, const size_t first_logical_dim, const size_t last_logical_dim, - const size_t work_blocks_X, const int32_t ctaid_X, const int32_t ctaid_Y, - const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, - const int64_t *const __restrict__ last_dims_ptr) { - constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; - constexpr bool is_single_tensor = (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || - SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM); - const size_t block_id = ctaid_Y * work_blocks_X + ctaid_X; - const size_t block_global_offset = - is_single_tensor ? (ctaid_Y * CHUNK_DIM_Y * last_logical_dim + ctaid_X * CHUNK_DIM_X) - : (block_id * ELTS_PER_CHUNK); - const size_t tensor_id = get_current_tensor_id( - num_tensors, block_global_offset, ctaid_Y, first_logical_dim, last_logical_dim, offsets_ptr); - const size_t rows = - get_tensor_rows_num(tensor_id, first_logical_dim, first_dims_ptr, num_tensors); - const size_t cols = get_tensor_cols_num(tensor_id, last_logical_dim, last_dims_ptr); - return JobDescriptor(block_id, block_global_offset, tensor_id, rows, cols); -} - -template -__device__ __forceinline__ bool is_job_valid(const JobDescriptor &job, - const size_t total_work_blocks, - const int64_t *const __restrict__ offsets_ptr) { - const bool is_valid = (job.block_id < total_work_blocks); - if (!is_valid) { - return false; - } - if (job.rows == 0 || job.cols == 0) { - return true; - } - if constexpr (SHAPE_REP == SAME_BOTH_DIMS) { - return true; - } - - const size_t tensor_start_offset = static_cast(offsets_ptr[job.tensor_id]); - const size_t tensor_end_offset = static_cast(offsets_ptr[job.tensor_id + 1]); - if (job.block_global_offset >= tensor_end_offset) { - return false; - } - - const size_t tensor_offset_from_start = job.block_global_offset - tensor_start_offset; - const size_t block_offset_Y_in_tensor = tensor_offset_from_start / job.cols; - if (block_offset_Y_in_tensor >= job.rows) { - return false; - } - - return true; -} - -__device__ __forceinline__ bool job_has_work(const JobDescriptor &job) { - return job.rows != 0 && job.cols != 0; -} - -__device__ __forceinline__ void advance_to_next_job(bool &job_finished, int32_t &ctaid_X, - int32_t &ctaid_Y, size_t &static_next_block_id, - const size_t static_block_stride, - const size_t total_work_blocks, - const size_t work_blocks_X) { - if (static_next_block_id < total_work_blocks) { - ctaid_X = static_cast(static_next_block_id % work_blocks_X); - ctaid_Y = static_cast(static_next_block_id / work_blocks_X); - static_next_block_id += static_block_stride; - } else { - job_finished = true; - } -} - -template -__device__ __forceinline__ BlockDescriptor -decode_block(const JobDescriptor &job, const int64_t *const __restrict__ offsets_ptr) { - constexpr bool is_single_tensor = (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || - SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM); - constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; - const size_t blocks_X_num_in_current_tensor = DIVUP(job.cols, CHUNK_DIM_X); - const size_t tensor_base = is_single_tensor ? 0 : static_cast(offsets_ptr[job.tensor_id]); - const size_t block_id_in_current_tensor = - is_single_tensor ? job.block_id : (job.block_id - tensor_base / ELTS_PER_CHUNK); - const size_t block_id_Y = block_id_in_current_tensor / blocks_X_num_in_current_tensor; - const size_t block_id_X = block_id_in_current_tensor % blocks_X_num_in_current_tensor; - const size_t block_offset_Y = block_id_Y * CHUNK_DIM_Y; - const size_t block_offset_X = block_id_X * CHUNK_DIM_X; - return BlockDescriptor(tensor_base, block_id_in_current_tensor, block_id_Y, block_id_X, - block_offset_Y, block_offset_X); -} - -// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index -__device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map, - CUtensorMap *global_tensor_map, - const uintptr_t global_data_ptr, - const size_t global_dim_Y, - const size_t global_dim_X, - const size_t data_type_size_bytes) { - __shared__ CUtensorMap shared_tensor_map; - shared_tensor_map = base_tensor_map; // Copy the base tensor map into shmem - constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; - if constexpr (is_blackwell) { - const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; - if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { - NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned."); - } - if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { - NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); - } - - asm volatile( - "{\n\t" - ".reg.b64 tensor_map_ptr; \n\t" - "mov.b64 tensor_map_ptr, %0; \n\t" - "tensormap.replace.tile.global_address.b1024.b64 [tensor_map_ptr], %1; \n\t" - "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 1, %2; \n\t" // DIM Y - "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 0, %3; \n\t" // DIM X - "tensormap.replace.tile.global_stride.b1024.b64 [tensor_map_ptr], 0, %4; \n" - "}\n" ::"l"(reinterpret_cast(&shared_tensor_map)), - "l"(global_data_ptr), "r"(static_cast(global_dim_Y)), - "r"(static_cast(global_dim_X)), "l"(static_cast(global_stride_bytes)) - : "memory"); - *global_tensor_map = shared_tensor_map; - } else { - NVTE_DEVICE_ERROR("tensormap.replace is architecture-specific. "); - } -} - -template -__global__ void __launch_bounds__(1) - update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input, - const __grid_constant__ CUtensorMap base_tensor_map_act_input, - const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise, - const __grid_constant__ CUtensorMap base_tensor_map_output_colwise, - const IType *const __restrict__ input_data_ptr, - const IType *const __restrict__ act_input_data_ptr, - const OType *const __restrict__ output_rowwise_data_ptr, - const OType *const __restrict__ output_colwise_data_ptr, - const ShapeRepresentation shape_rep, const size_t num_tensors, - const size_t first_logical_dim, const size_t last_logical_dim, - const int64_t *const __restrict__ offsets_ptr, - const int64_t *const __restrict__ first_dims_ptr, - const int64_t *const __restrict__ last_dims_ptr, const bool rowwise, - const bool colwise, const bool compute_dactivations) { - const size_t tensor_id = blockIdx.x; - const size_t rows = - get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); - const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); - - const size_t offset_elts = offsets_ptr[tensor_id]; - - // Zero-sized groups: skip TMA descriptor update. The main kernel already returns - // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension - // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. - if (rows == 0 || cols == 0) { - return; - } - - if (tensor_id < num_tensors) { - { - CUtensorMap *modified_tensor_map_input = &g_tensor_maps.input[tensor_id]; - const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_input, modified_tensor_map_input, global_data_ptr, - rows, cols, sizeof(IType)); - } - if (compute_dactivations) { - CUtensorMap *modified_tensor_map_act_input = &g_tensor_maps.act_input[tensor_id]; - const uintptr_t global_data_ptr = - reinterpret_cast(act_input_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_act_input, modified_tensor_map_act_input, - global_data_ptr, rows, cols, sizeof(IType)); - } - if (rowwise) { - CUtensorMap *modified_tensor_map_output_rowwise = &g_tensor_maps.output_rowwise[tensor_id]; - const uintptr_t global_data_ptr = - reinterpret_cast(output_rowwise_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_output_rowwise, modified_tensor_map_output_rowwise, - global_data_ptr, rows, cols, sizeof(OType)); - } - if (colwise) { - CUtensorMap *modified_tensor_map_output_colwise = &g_tensor_maps.output_colwise[tensor_id]; - const uintptr_t global_data_ptr = - reinterpret_cast(output_colwise_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_output_colwise, modified_tensor_map_output_colwise, - global_data_ptr, rows, cols, sizeof(OType)); - } - } -} - -__device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map)); -#else - NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+."); -#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) -} - -// Issue TMA global->shared transfer for one stage of input (and optional activation input). -template -__device__ __forceinline__ void prefetch_input_stage( - IType *in_sh, IType *act_in_sh, const CUtensorMap &tensor_map_input, - const CUtensorMap &tensor_map_act_input, const size_t global_offset_X, - const size_t global_offset_Y, const size_t buff_offset, const size_t shmem_buff_size, - uint64_t *barrier, const bool leading_thread) { - if (leading_thread) { - ptx::mbarrier_arrive_expect_tx(barrier, shmem_buff_size); - ptx::cp_async_bulk_tensor_2d_global_to_shared( - reinterpret_cast(&in_sh[buff_offset]), - reinterpret_cast(&tensor_map_input), global_offset_X, global_offset_Y, - barrier); - if constexpr (IS_DACT) { - ptx::cp_async_bulk_tensor_2d_global_to_shared( - reinterpret_cast(&act_in_sh[buff_offset]), - reinterpret_cast(&tensor_map_act_input), global_offset_X, - global_offset_Y, barrier); - } - } -} - -// Issue TMA shared->global transfer for one stage of outputs. -template -__device__ __forceinline__ void store_output_stage( - OType *out_rowwise_data_sh, OType *out_colwise_data_sh, - const CUtensorMap &tensor_map_output_rowwise, const CUtensorMap &tensor_map_output_colwise, - const size_t global_offset_X, const size_t global_offset_Y, const size_t buff_offset, - const bool leading_thread) { - if (!leading_thread) { - return; - } - - if constexpr (ROWWISE_SCALING) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output_rowwise), global_offset_X, - global_offset_Y, reinterpret_cast(&out_rowwise_data_sh[buff_offset])); - } - if constexpr (COLWISE_SCALING) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output_colwise), global_offset_X, - global_offset_Y, reinterpret_cast(&out_colwise_data_sh[buff_offset])); - } - if constexpr (ROWWISE_SCALING || COLWISE_SCALING) { - ptx::cp_async_bulk_commit_group(); - } -} - -} // namespace common -} // namespace dispatch -} // namespace transformer_engine +#include "grouped_layout.cuh" +#include "grouped_tma.cuh" #endif // TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ diff --git a/transformer_engine/common/cast/core/grouped_layout.cuh b/transformer_engine/common/cast/core/grouped_layout.cuh new file mode 100644 index 0000000000..8337e44815 --- /dev/null +++ b/transformer_engine/common/cast/core/grouped_layout.cuh @@ -0,0 +1,413 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file grouped_layout.cuh + * \brief Architecture-neutral work-decomposition helpers for grouped quantization. + * + * Maps CTA coordinates to (tensor, tile) work-items for grouped tensors of + * varying shapes, plus the dbias reductions. Contains no arch-specific PTX or + * TMA machinery, so it can be included by sources that compile for the generic + * (non-suffixed) Blackwell architectures. The TMA descriptor management and + * bulk-copy staging live in grouped_tma.cuh. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_CORE_GROUPED_LAYOUT_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_CORE_GROUPED_LAYOUT_CUH_ + +#include +#include + +#include "../../common.h" +#include "../../utils.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace common { + +inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_block) { + const size_t N = product(t->data.shape); + const bool isFullTile = (N % elems_per_block == 0); + return isFullTile; +} + +namespace kernel { + +constexpr size_t THREADS_PER_BLOCK = 256; +template +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial, + const size_t rows, const size_t cols) { + using ComputeVec = Vec; + using OutputVec = Vec; + + const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; + + if (thread_id * nvec >= cols) { + return; + } + + const float *const thread_in_base = dbias_partial + thread_id * nvec; + OType *const thread_out_base = dbias_output + thread_id * nvec; + + ComputeVec ldg_vec; + ComputeVec acc_vec; + acc_vec.clear(); + for (int i = 0; i < rows; ++i) { + ldg_vec.load_from(thread_in_base + i * cols); +#pragma unroll + for (int e = 0; e < nvec; ++e) { + acc_vec.data.elt[e] += ldg_vec.data.elt[e]; + } + } + + OutputVec stg_vec; +#pragma unroll + for (int e = 0; e < nvec; ++e) { + stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); + } + stg_vec.store_to(thread_out_base); +} + +template +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + group_reduce_dbias_kernel(const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const offsets_ptr, const int64_t *const first_dims_ptr, + const int64_t *const last_dims_ptr, OType *const dbias_output, + const float *dbias_partial, const size_t chunk_dim_Y) { + using ComputeVec = Vec; + using OutputVec = Vec; + + const size_t tensor_id = blockIdx.y; + const size_t tensor_rows = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) + ? (first_logical_dim / num_tensors) + : static_cast(first_dims_ptr[tensor_id]); + + const size_t rows = tensor_rows / chunk_dim_Y; + const size_t cols = last_logical_dim; + + const size_t dbias_in_offset_Y = + (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) + ? (tensor_id * (tensor_rows / chunk_dim_Y)) + : (static_cast(offsets_ptr[tensor_id]) / cols / chunk_dim_Y); + + const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; + + if (thread_id * nvec >= cols) { + return; + } + + const float *const thread_in_base = dbias_partial + dbias_in_offset_Y * cols + thread_id * nvec; + OType *const thread_out_base = dbias_output + tensor_id * cols + thread_id * nvec; + + ComputeVec ldg_vec; + ComputeVec acc_vec; + acc_vec.clear(); + for (int i = 0; i < rows; ++i) { + ldg_vec.load_from(thread_in_base + i * cols); +#pragma unroll + for (int e = 0; e < nvec; ++e) { + acc_vec.data.elt[e] += ldg_vec.data.elt[e]; + } + } + + OutputVec stg_vec; +#pragma unroll + for (int e = 0; e < nvec; ++e) { + stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); + } + stg_vec.store_to(thread_out_base); +} +} // namespace kernel + +template +void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, const size_t cols, + cudaStream_t stream) { + using namespace kernel; + constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 + constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); + + NVTE_CHECK(cols % reduce_dbias_nvec == 0, "Unsupported shape."); + const size_t reduce_dbias_num_blocks = DIVUP(cols, THREADS_PER_BLOCK * reduce_dbias_nvec); + + reduce_dbias_kernel + <<>>( + reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void grouped_reduce_dbias(const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const data_tensor_offsets_ptr, + const int64_t *const data_tensor_first_dims_ptr, + const int64_t *const data_tensor_last_dims_ptr, GroupedTensor *dbias, + const float *workspace_ptr, const size_t chunk_dim_Y, + cudaStream_t stream) { + using namespace kernel; + constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 + constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); + + NVTE_CHECK(last_logical_dim % reduce_dbias_nvec == 0, "Unsupported shape."); + + const size_t blocks_X = DIVUP(last_logical_dim, THREADS_PER_BLOCK * reduce_dbias_nvec); + const size_t blocks_Y = num_tensors; + const dim3 grid(blocks_X, blocks_Y); + + group_reduce_dbias_kernel<<>>( + shape_rep, num_tensors, first_logical_dim, last_logical_dim, data_tensor_offsets_ptr, + data_tensor_first_dims_ptr, data_tensor_last_dims_ptr, + reinterpret_cast(dbias->data.dptr), workspace_ptr, chunk_dim_Y); + + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +__device__ __forceinline__ size_t +find_tensor_from_offsets(const OffsetType *const __restrict__ offsets_ptr, const size_t num_tensors, + const size_t offset) { + size_t low = 1; + size_t hi = num_tensors; // [low, hi] + + while (low < hi) { + const size_t mid = low + (hi - low) / 2; + const size_t mid_offset = static_cast(offsets_ptr[mid]); + + if (mid_offset <= offset) { + low = mid + 1; + } else { + hi = mid; + } + } + return low - 1; +} + +template +__device__ __forceinline__ size_t +get_current_tensor_id(const size_t num_tensors, const size_t current_offset, const size_t block_Y, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr) { + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + const size_t current_row = block_Y * CHUNK_DIM_Y; + const size_t rows_per_tensor = first_logical_dim / num_tensors; + return current_row / rows_per_tensor; + } else { + return find_tensor_from_offsets(offsets_ptr, num_tensors, current_offset); + } +} + +template +__device__ __forceinline__ size_t +get_tensor_rows_num(const size_t tensor_id, const size_t first_logical_dim, + const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { + size_t rows_num = 0; + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_LAST_DIM) { + rows_num = first_logical_dim; + } else { + rows_num = static_cast(first_dims_ptr[tensor_id]); + } + if (rows_num % 128 != 0) { + NVTE_DEVICE_ERROR("First dimension of each tensor in a group must be divisible by 128."); + } + return rows_num; +} + +__device__ __forceinline__ size_t get_tensor_rows_num( + const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t first_logical_dim, + const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + return get_tensor_rows_num(tensor_id, first_logical_dim, + first_dims_ptr, num_tensors); + case ShapeRepresentation::VARYING_FIRST_DIM: + return get_tensor_rows_num( + tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + case ShapeRepresentation::VARYING_LAST_DIM: + return get_tensor_rows_num( + tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + case ShapeRepresentation::VARYING_BOTH_DIMS: + return get_tensor_rows_num( + tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + } + return 0; +} + +template +__device__ __forceinline__ size_t +get_tensor_cols_num(const size_t tensor_id, const size_t last_logical_dim, + const int64_t *const __restrict__ last_dims_ptr) { + size_t cols_num = 0; + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + cols_num = last_logical_dim; + } else { + cols_num = static_cast(last_dims_ptr[tensor_id]); + if (cols_num % 128 != 0) { + NVTE_DEVICE_ERROR( + "For varying last dimensions support, the last dimension of each tensor in a group " + "must be divisible by 128."); + } + } + return cols_num; +} + +__device__ __forceinline__ size_t get_tensor_cols_num( + const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t last_logical_dim, + const int64_t *const __restrict__ last_dims_ptr) { + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + return get_tensor_cols_num(tensor_id, last_logical_dim, + last_dims_ptr); + case ShapeRepresentation::VARYING_FIRST_DIM: + return get_tensor_cols_num( + tensor_id, last_logical_dim, last_dims_ptr); + case ShapeRepresentation::VARYING_LAST_DIM: + return get_tensor_cols_num(tensor_id, last_logical_dim, + last_dims_ptr); + case ShapeRepresentation::VARYING_BOTH_DIMS: + return get_tensor_cols_num( + tensor_id, last_logical_dim, last_dims_ptr); + } + return 0; +} + +// Logical work-item decoded from CTA coordinates. +struct JobDescriptor { + size_t block_id = 0; + size_t block_global_offset = 0; + size_t tensor_id = 0; + size_t rows = 0; + size_t cols = 0; + + __host__ __device__ __forceinline__ constexpr JobDescriptor() = default; + + __host__ __device__ __forceinline__ constexpr JobDescriptor(const size_t block_id_, + const size_t block_global_offset_, + const size_t tensor_id_, + const size_t rows_, + const size_t cols_) + : block_id(block_id_), + block_global_offset(block_global_offset_), + tensor_id(tensor_id_), + rows(rows_), + cols(cols_) {} +}; + +// Tensor-local coordinates for a work-item. +struct BlockDescriptor { + size_t tensor_base = 0; + size_t block_id_in_current_tensor = 0; + size_t block_id_Y = 0; + size_t block_id_X = 0; + size_t block_offset_Y = 0; + size_t block_offset_X = 0; + + __host__ __device__ __forceinline__ constexpr BlockDescriptor() = default; + + __host__ __device__ __forceinline__ constexpr BlockDescriptor( + const size_t tensor_base_, const size_t block_id_in_current_tensor_, const size_t block_id_Y_, + const size_t block_id_X_, const size_t block_offset_Y_, const size_t block_offset_X_) + : tensor_base(tensor_base_), + block_id_in_current_tensor(block_id_in_current_tensor_), + block_id_Y(block_id_Y_), + block_id_X(block_id_X_), + block_offset_Y(block_offset_Y_), + block_offset_X(block_offset_X_) {} +}; + +template +__device__ __forceinline__ JobDescriptor decode_job( + const size_t num_tensors, const size_t first_logical_dim, const size_t last_logical_dim, + const size_t work_blocks_X, const int32_t ctaid_X, const int32_t ctaid_Y, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr) { + constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; + constexpr bool is_single_tensor = (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM); + const size_t block_id = ctaid_Y * work_blocks_X + ctaid_X; + const size_t block_global_offset = + is_single_tensor ? (ctaid_Y * CHUNK_DIM_Y * last_logical_dim + ctaid_X * CHUNK_DIM_X) + : (block_id * ELTS_PER_CHUNK); + const size_t tensor_id = get_current_tensor_id( + num_tensors, block_global_offset, ctaid_Y, first_logical_dim, last_logical_dim, offsets_ptr); + const size_t rows = + get_tensor_rows_num(tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, last_logical_dim, last_dims_ptr); + return JobDescriptor(block_id, block_global_offset, tensor_id, rows, cols); +} + +template +__device__ __forceinline__ bool is_job_valid(const JobDescriptor &job, + const size_t total_work_blocks, + const int64_t *const __restrict__ offsets_ptr) { + const bool is_valid = (job.block_id < total_work_blocks); + if (!is_valid) { + return false; + } + if (job.rows == 0 || job.cols == 0) { + return true; + } + if constexpr (SHAPE_REP == SAME_BOTH_DIMS) { + return true; + } + + const size_t tensor_start_offset = static_cast(offsets_ptr[job.tensor_id]); + const size_t tensor_end_offset = static_cast(offsets_ptr[job.tensor_id + 1]); + if (job.block_global_offset >= tensor_end_offset) { + return false; + } + + const size_t tensor_offset_from_start = job.block_global_offset - tensor_start_offset; + const size_t block_offset_Y_in_tensor = tensor_offset_from_start / job.cols; + if (block_offset_Y_in_tensor >= job.rows) { + return false; + } + + return true; +} + +__device__ __forceinline__ bool job_has_work(const JobDescriptor &job) { + return job.rows != 0 && job.cols != 0; +} + +__device__ __forceinline__ void advance_to_next_job(bool &job_finished, int32_t &ctaid_X, + int32_t &ctaid_Y, size_t &static_next_block_id, + const size_t static_block_stride, + const size_t total_work_blocks, + const size_t work_blocks_X) { + if (static_next_block_id < total_work_blocks) { + ctaid_X = static_cast(static_next_block_id % work_blocks_X); + ctaid_Y = static_cast(static_next_block_id / work_blocks_X); + static_next_block_id += static_block_stride; + } else { + job_finished = true; + } +} + +template +__device__ __forceinline__ BlockDescriptor +decode_block(const JobDescriptor &job, const int64_t *const __restrict__ offsets_ptr) { + constexpr bool is_single_tensor = (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM); + constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; + const size_t blocks_X_num_in_current_tensor = DIVUP(job.cols, CHUNK_DIM_X); + const size_t tensor_base = is_single_tensor ? 0 : static_cast(offsets_ptr[job.tensor_id]); + const size_t block_id_in_current_tensor = + is_single_tensor ? job.block_id : (job.block_id - tensor_base / ELTS_PER_CHUNK); + const size_t block_id_Y = block_id_in_current_tensor / blocks_X_num_in_current_tensor; + const size_t block_id_X = block_id_in_current_tensor % blocks_X_num_in_current_tensor; + const size_t block_offset_Y = block_id_Y * CHUNK_DIM_Y; + const size_t block_offset_X = block_id_X * CHUNK_DIM_X; + return BlockDescriptor(tensor_base, block_id_in_current_tensor, block_id_Y, block_id_X, + block_offset_Y, block_offset_X); +} + +} // namespace common +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_CORE_GROUPED_LAYOUT_CUH_ diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh new file mode 100644 index 0000000000..7ff501bb49 --- /dev/null +++ b/transformer_engine/common/cast/core/grouped_tma.cuh @@ -0,0 +1,217 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file grouped_tma.cuh + * \brief Architecture-specific TMA helpers for grouped quantization. + * + * TMA descriptor storage/update and bulk-copy staging built on Blackwell + * family/arch-specific PTX (via ptx.cuh). Translation units that include this + * header must be compiled for the smXXXa/smXXXf targets (i.e. listed in + * transformer_engine_cuda_arch_specific_sources in CMakeLists.txt). The + * arch-neutral work-decomposition helpers live in grouped_layout.cuh. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_CORE_GROUPED_TMA_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_CORE_GROUPED_TMA_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "grouped_layout.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace common { + +constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; + +struct alignas(128) TensorMapStorage { + alignas(128) CUtensorMap input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + alignas(128) CUtensorMap act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + alignas(128) CUtensorMap output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + alignas(128) CUtensorMap output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; +}; + +// Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs. +static __device__ TensorMapStorage g_tensor_maps; + +inline bool dimensions_supported_by_TMA(const Tensor *const t) { + const size_t cols = t->flat_last_dim(); + constexpr size_t TMA_bytes = 16; + const size_t alignment_requirement = (TMA_bytes * 8) / typeToNumBits(t->dtype()); + return cols % alignment_requirement == 0; +} + +__device__ __forceinline__ unsigned char *align_smem_ptr_per_TMA_requirements(unsigned char *p) { + size_t addr = reinterpret_cast(p); + addr = (addr + TMA_SHMEM_ALIGNMENT - 1) & ~(TMA_SHMEM_ALIGNMENT - 1); + return reinterpret_cast(addr); +} + +// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index +__device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map, + CUtensorMap *global_tensor_map, + const uintptr_t global_data_ptr, + const size_t global_dim_Y, + const size_t global_dim_X, + const size_t data_type_size_bytes) { + __shared__ CUtensorMap shared_tensor_map; + shared_tensor_map = base_tensor_map; // Copy the base tensor map into shmem + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; + if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned."); + } + if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); + } + + asm volatile( + "{\n\t" + ".reg.b64 tensor_map_ptr; \n\t" + "mov.b64 tensor_map_ptr, %0; \n\t" + "tensormap.replace.tile.global_address.b1024.b64 [tensor_map_ptr], %1; \n\t" + "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 1, %2; \n\t" // DIM Y + "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 0, %3; \n\t" // DIM X + "tensormap.replace.tile.global_stride.b1024.b64 [tensor_map_ptr], 0, %4; \n" + "}\n" ::"l"(reinterpret_cast(&shared_tensor_map)), + "l"(global_data_ptr), "r"(static_cast(global_dim_Y)), + "r"(static_cast(global_dim_X)), "l"(static_cast(global_stride_bytes)) + : "memory"); + *global_tensor_map = shared_tensor_map; + } else { + NVTE_DEVICE_ERROR("tensormap.replace is architecture-specific. "); + } +} + +template +__global__ void __launch_bounds__(1) + update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input, + const __grid_constant__ CUtensorMap base_tensor_map_act_input, + const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise, + const __grid_constant__ CUtensorMap base_tensor_map_output_colwise, + const IType *const __restrict__ input_data_ptr, + const IType *const __restrict__ act_input_data_ptr, + const OType *const __restrict__ output_rowwise_data_ptr, + const OType *const __restrict__ output_colwise_data_ptr, + const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr, + const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, const bool rowwise, + const bool colwise, const bool compute_dactivations) { + const size_t tensor_id = blockIdx.x; + const size_t rows = + get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + + const size_t offset_elts = offsets_ptr[tensor_id]; + + // Zero-sized groups: skip TMA descriptor update. The main kernel already returns + // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension + // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. + if (rows == 0 || cols == 0) { + return; + } + + if (tensor_id < num_tensors) { + { + CUtensorMap *modified_tensor_map_input = &g_tensor_maps.input[tensor_id]; + const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_input, modified_tensor_map_input, global_data_ptr, + rows, cols, sizeof(IType)); + } + if (compute_dactivations) { + CUtensorMap *modified_tensor_map_act_input = &g_tensor_maps.act_input[tensor_id]; + const uintptr_t global_data_ptr = + reinterpret_cast(act_input_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_act_input, modified_tensor_map_act_input, + global_data_ptr, rows, cols, sizeof(IType)); + } + if (rowwise) { + CUtensorMap *modified_tensor_map_output_rowwise = &g_tensor_maps.output_rowwise[tensor_id]; + const uintptr_t global_data_ptr = + reinterpret_cast(output_rowwise_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_output_rowwise, modified_tensor_map_output_rowwise, + global_data_ptr, rows, cols, sizeof(OType)); + } + if (colwise) { + CUtensorMap *modified_tensor_map_output_colwise = &g_tensor_maps.output_colwise[tensor_id]; + const uintptr_t global_data_ptr = + reinterpret_cast(output_colwise_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_output_colwise, modified_tensor_map_output_colwise, + global_data_ptr, rows, cols, sizeof(OType)); + } + } +} + +__device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map)); +#else + NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// Issue TMA global->shared transfer for one stage of input (and optional activation input). +template +__device__ __forceinline__ void prefetch_input_stage( + IType *in_sh, IType *act_in_sh, const CUtensorMap &tensor_map_input, + const CUtensorMap &tensor_map_act_input, const size_t global_offset_X, + const size_t global_offset_Y, const size_t buff_offset, const size_t shmem_buff_size, + uint64_t *barrier, const bool leading_thread) { + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(barrier, shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&in_sh[buff_offset]), + reinterpret_cast(&tensor_map_input), global_offset_X, global_offset_Y, + barrier); + if constexpr (IS_DACT) { + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&act_in_sh[buff_offset]), + reinterpret_cast(&tensor_map_act_input), global_offset_X, + global_offset_Y, barrier); + } + } +} + +// Issue TMA shared->global transfer for one stage of outputs. +template +__device__ __forceinline__ void store_output_stage( + OType *out_rowwise_data_sh, OType *out_colwise_data_sh, + const CUtensorMap &tensor_map_output_rowwise, const CUtensorMap &tensor_map_output_colwise, + const size_t global_offset_X, const size_t global_offset_Y, const size_t buff_offset, + const bool leading_thread) { + if (!leading_thread) { + return; + } + + if constexpr (ROWWISE_SCALING) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_rowwise), global_offset_X, + global_offset_Y, reinterpret_cast(&out_rowwise_data_sh[buff_offset])); + } + if constexpr (COLWISE_SCALING) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_colwise), global_offset_X, + global_offset_Y, reinterpret_cast(&out_colwise_data_sh[buff_offset])); + } + if constexpr (ROWWISE_SCALING || COLWISE_SCALING) { + ptx::cp_async_bulk_commit_group(); + } +} + +} // namespace common +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_CORE_GROUPED_TMA_CUH_ diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 6c71285cd4..97dd27aec6 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -17,6 +17,7 @@ #include "../../transpose/cast_transpose.h" #include "../../util/vectorized_pointwise.h" #include "../core/common.cuh" +#include "../fp8/group_quantize_fp8.cuh" #include "../fp8/quantize_fp8.cuh" #include "../mxfp8/group_quantize_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" @@ -460,6 +461,11 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor // Dispatch to quantization kernel depending on data format switch (scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: { + fp8::group_quantize(input_tensor, noop_tensor, output_tensor, + &quant_config_cpp, stream); + break; + } case NVTE_MXFP8_1D_SCALING: { mxfp8::group_quantize( input_tensor, activations_tensor, noop_tensor, output_tensor, dbias_tensor, diff --git a/transformer_engine/common/cast/fp8/group_amax_fp8.cuh b/transformer_engine/common/cast/fp8/group_amax_fp8.cuh new file mode 100644 index 0000000000..5040f15381 --- /dev/null +++ b/transformer_engine/common/cast/fp8/group_amax_fp8.cuh @@ -0,0 +1,280 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file group_amax_fp8.cuh + * \brief CUDA kernels to compute per-group amax values for grouped FP8 + * tensor scaling. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_AMAX_FP8_CUH_ +#define TRANSFORMER_ENGINE_GROUP_AMAX_FP8_CUH_ + +#include +#include + +#include +#include + +#include "../../common.h" +#include "../../recipe/recipe_common.cuh" +#include "../../util/cuda_runtime.h" +#include "../../utils.cuh" +#include "../core/grouped_layout.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8 { +namespace group_amax_kernel { + +constexpr int GROUPED_AMAX_KERNEL_THREADS = 256; +// Target block-per-tensor multiplier so that, even with few groups, we cover +// enough SMs on H100/GB200-class GPUs (132 SMs) to be HBM-bandwidth bound. +// blocks_per_tensor is bounded by both work and a hard cap to avoid launching +// far more blocks than there is work for. +constexpr int GROUPED_AMAX_BLOCKS_PER_TENSOR_CAP = 64; +constexpr size_t GROUPED_AMAX_MIN_ELTS_PER_BLOCK = 8 * 1024; // ~16KB of bf16 + +// Zero per-tensor amax buffer so the main kernel can use atomicMax updates. +// ``static`` gives the kernel internal linkage so this header can be included +// from multiple translation units without violating the ODR. +static __launch_bounds__(GROUPED_AMAX_KERNEL_THREADS) __global__ + void grouped_amax_zero_kernel(float *amax_ptr, const size_t num_tensors, + const float *noop_ptr) { + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { + return; + } + const size_t tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid < num_tensors) { + amax_ptr[tid] = 0.0f; + } +} + +// Flat vectorized amax kernel for imbalanced/varying shapes and uniform shapes alike. +// +// Grid: (num_blocks). +// Each block scans a flat chunk of elements within a specific tensor, or spanning across multiple tensors. +// Fully load-balanced! +template +__launch_bounds__(GROUPED_AMAX_KERNEL_THREADS) __global__ + void grouped_amax_flat_kernel(const InputType *__restrict__ input, float *__restrict__ amax, + const size_t num_tensors, const size_t first_logical_dim, + const size_t last_logical_dim, + const int64_t *__restrict__ offsets_ptr, + const float *__restrict__ noop_ptr) { + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { + return; + } + + // Dynamic shared memory layout: + // s_offsets: int64_t[num_tensors + 1] + extern __shared__ char s_mem[]; + int64_t *s_offsets = reinterpret_cast(s_mem); + + const size_t tid = threadIdx.x; + + // 1. Copy or compute offsets to shared memory + if (offsets_ptr != nullptr) { + for (size_t i = tid; i <= num_tensors; i += blockDim.x) { + s_offsets[i] = offsets_ptr[i]; + } + } else { + const size_t rows = first_logical_dim / num_tensors; + const size_t numel_per_tensor = rows * last_logical_dim; + for (size_t i = tid; i <= num_tensors; i += blockDim.x) { + s_offsets[i] = i * numel_per_tensor; + } + } + __syncthreads(); + + const size_t total_active_elements = s_offsets[num_tensors]; + if (total_active_elements == 0) { + return; + } + + // Chunk based on the active elements and the grid size + size_t block_chunk_size = (total_active_elements + gridDim.x - 1) / gridDim.x; + // Align block_chunk_size to a multiple of NVEC to ensure perfect vector alignment + block_chunk_size = DIVUP(block_chunk_size, static_cast(NVEC)) * NVEC; + // + if (block_chunk_size < GROUPED_AMAX_MIN_ELTS_PER_BLOCK) { + block_chunk_size = GROUPED_AMAX_MIN_ELTS_PER_BLOCK; + } + + // Calculate this block's local work range + const size_t start_elt = blockIdx.x * block_chunk_size; + const size_t end_elt = start_elt + block_chunk_size < total_active_elements + ? start_elt + block_chunk_size + : total_active_elements; + + if (start_elt >= end_elt) { + return; + } + + constexpr int kWarps = GROUPED_AMAX_KERNEL_THREADS / THREADS_PER_WARP; + + using IVecT = Vec; + + size_t curr_elt = start_elt; + size_t tensor_id = transformer_engine::dispatch::common::find_tensor_from_offsets( + s_offsets, num_tensors, curr_elt); + + while (curr_elt < end_elt) { + const size_t t_next_offset = s_offsets[tensor_id + 1]; + const size_t t_end = t_next_offset < end_elt ? t_next_offset : end_elt; + + const size_t numel = t_end - curr_elt; + if (numel > 0) { + const InputType *base = input + curr_elt; + const size_t total_vecs = numel / NVEC; + const size_t tail_start = total_vecs * NVEC; + + InputType thread_amax_val = 0.f; + const bool aligned = (reinterpret_cast(base) % IVecT::BYTES) == 0; + + if (aligned) { + InputType acc0 = 0.f; + InputType acc1 = 0.f; + InputType acc2 = 0.f; + InputType acc3 = 0.f; + size_t v = tid; + // 4-way vectorized load and reduce + for (; v + 3 * blockDim.x < total_vecs; v += 4 * blockDim.x) { + IVecT vec0, vec1, vec2, vec3; + vec0.load_from(base + v * NVEC); + vec1.load_from(base + (v + blockDim.x) * NVEC); + vec2.load_from(base + (v + 2 * blockDim.x) * NVEC); + vec3.load_from(base + (v + 3 * blockDim.x) * NVEC); +#pragma unroll + for (int i = 0; i < NVEC; ++i) { + acc0 = max_val(acc0, abs_val(vec0.data.elt[i])); + acc1 = max_val(acc1, abs_val(vec1.data.elt[i])); + acc2 = max_val(acc2, abs_val(vec2.data.elt[i])); + acc3 = max_val(acc3, abs_val(vec3.data.elt[i])); + } + } + thread_amax_val = max_val(max_val(acc0, acc1), max_val(acc2, acc3)); + // 1-way vectorized load and reduce for the tail block + for (; v < total_vecs; v += blockDim.x) { + IVecT vec; + vec.load_from(base + v * NVEC); +#pragma unroll + for (int i = 0; i < NVEC; ++i) { + thread_amax_val = max_val(thread_amax_val, abs_val(vec.data.elt[i])); + } + } + } else { + for (size_t v = tid; v < total_vecs; v += blockDim.x) { + const InputType *p = base + v * NVEC; +#pragma unroll + for (int i = 0; i < NVEC; ++i) { + thread_amax_val = max_val(thread_amax_val, abs_val(p[i])); + } + } + } + + // Tail elements + for (size_t i = tail_start + tid; i < numel; i += blockDim.x) { + thread_amax_val = max_val(thread_amax_val, abs_val(base[i])); + } + + // Reduce the per-thread amax over the block and update the global amax. + const int warp_id = tid / THREADS_PER_WARP; + const float block_amax = reduce_max(static_cast(thread_amax_val), warp_id); + if (tid == 0) { + atomicMaxFloat(&amax[tensor_id], block_amax); + } + // Guard reduce_max's shared staging against reuse by the next iteration. + __syncthreads(); + curr_elt = t_end; + } + + if (curr_elt >= t_next_offset) { + tensor_id++; + } + } +} + +// Per-group scale update: derive scale = max_fp8 / amax and scale_inv = 1/scale, +// one thread per group. This is the grouped counterpart of compute_scale_from_amax in +// the per-tensor current-scaling path. ``static`` gives the kernel internal +// linkage so this header can be included from multiple translation units. +static __launch_bounds__(GROUPED_AMAX_KERNEL_THREADS) __global__ + void grouped_compute_scale_kernel(const float *__restrict__ amax_ptr, + float *__restrict__ scale_ptr, + float *__restrict__ scale_inv_ptr, const size_t num_tensors, + const float max_fp8, const bool force_pow_2_scales, + const float epsilon, const float *__restrict__ noop_ptr) { + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { + return; + } + const size_t tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid < num_tensors) { + const float scale = compute_scale_from_amax(amax_ptr[tid], max_fp8, force_pow_2_scales, epsilon, + std::numeric_limits::max()); + scale_ptr[tid] = scale; + reciprocal(scale_inv_ptr + tid, scale); + } +} + +} // namespace group_amax_kernel + +template +void launch_grouped_amax_kernel(const InputType *input, float *amax, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const ShapeRepresentation shape_rep, const int64_t *offsets_ptr, + const int64_t *first_dims_ptr, const int64_t *last_dims_ptr, + const float *noop_ptr, cudaStream_t stream) { + using namespace group_amax_kernel; + if (num_tensors == 0) { + return; + } + (void)shape_rep; + (void)first_dims_ptr; + (void)last_dims_ptr; + + // Zero out the per-tensor amax buffer so atomicMaxFloat works. + const size_t zero_blocks = DIVUP(num_tensors, static_cast(GROUPED_AMAX_KERNEL_THREADS)); + grouped_amax_zero_kernel<<>>( + amax, num_tensors, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + + constexpr int kVecBytes = 16; + constexpr int kNvec = kVecBytes / sizeof(InputType); + static_assert(kNvec >= 1, "Vector width must be at least 1"); + + const int num_sms = ::transformer_engine::cuda::sm_count(); + const size_t grid_size = 8 * num_sms; + const size_t shared_mem_bytes = (num_tensors + 1) * sizeof(int64_t); + + grouped_amax_flat_kernel + <<>>( + input, amax, num_tensors, first_logical_dim, last_logical_dim, offsets_ptr, noop_ptr); + + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// Launch the per-group scale/scale_inv update. The amax, scale and scale_inv +// buffers each hold one FP32 entry per group. ``noop_ptr`` is the optional +// graph-safe no-op flag (skip when 1.0f); pass nullptr to always update. +inline void launch_grouped_compute_scale_kernel(const float *amax, float *scale, float *scale_inv, + const size_t num_tensors, const float max_fp8, + const bool force_pow_2_scales, const float epsilon, + const float *noop_ptr, cudaStream_t stream) { + using namespace group_amax_kernel; + if (num_tensors == 0) { + return; + } + const size_t blocks = DIVUP(num_tensors, static_cast(GROUPED_AMAX_KERNEL_THREADS)); + grouped_compute_scale_kernel<<>>( + amax, scale, scale_inv, num_tensors, max_fp8, force_pow_2_scales, epsilon, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace fp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_GROUP_AMAX_FP8_CUH_ diff --git a/transformer_engine/common/cast/fp8/group_quantize_fp8.cuh b/transformer_engine/common/cast/fp8/group_quantize_fp8.cuh new file mode 100644 index 0000000000..95fe5eea28 --- /dev/null +++ b/transformer_engine/common/cast/fp8/group_quantize_fp8.cuh @@ -0,0 +1,1226 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file group_quantize_fp8.cuh + * \brief CUDA kernels to quantize grouped tensors to FP8 tensor scaling. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_CUH_ +#define TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_CUH_ + +#include +#include + +#include +#include + +#include "../../common.h" +#include "../../util/cuda_runtime.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "../core/grouped_layout.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8 { +namespace group_quantize_kernel { + +using namespace dispatch::common; + +constexpr size_t WARPS_PER_TILE = 8; +constexpr size_t THREADS_PER_TILE = THREADS_PER_WARP * WARPS_PER_TILE; +constexpr size_t ROWWISE_LOAD_SIZE_BYTES = 16; +constexpr size_t TRANSPOSE_LOAD_SIZE_BYTES = 8; +constexpr size_t ROWWISE_STORE_SIZE_BYTES = 16; +constexpr size_t TRANSPOSE_STORE_SIZE_BYTES = 8; +constexpr size_t ROWWISE_FLAT_LOAD_SIZE_BYTES = 32; +constexpr size_t ROWWISE_FLAT_THREADS = 512; +constexpr size_t VARYING_FIRST_TRANSPOSE_MAX_ROW_TILES_PER_TENSOR = 4; +constexpr size_t TRANSPOSE_SHARED_PAD = 1; + +template +struct supports_fast_scaled_fp8_cvt_4 { + static constexpr bool value = + (std::is_same::value || std::is_same::value || + std::is_same::value) && + (std::is_same::value || std::is_same::value); +}; + +template +__device__ __forceinline__ void fast_scaled_fp8_cvt_4(const IType *input, OType *output, + const float scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using IType4 = ptx::FPx4; + using OType4 = ptx::FPx4; + const ptx::floatx2 scale_2x{scale, scale}; + const IType4 &in4x = *reinterpret_cast(input); + OType4 &out4x = *reinterpret_cast(output); + ptx::mul_cvt_4x(out4x, in4x, scale_2x); +#else + (void)input; + (void)output; + (void)scale; +#endif +} + +template +__device__ __forceinline__ void scaled_fp8_cvt_vec(const Vec &input, + Vec &output, + const size_t valid_cols, const float scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr bool kUseFastCvt = + !IS_ACT && (NUM_ELTS % 4 == 0) && supports_fast_scaled_fp8_cvt_4::value; +#else + constexpr bool kUseFastCvt = false; +#endif + + if constexpr (kUseFastCvt) { + if (valid_cols == NUM_ELTS) { +#pragma unroll + for (size_t base = 0; base < NUM_ELTS; base += 4) { + fast_scaled_fp8_cvt_4(input.data.elt + base, output.data.elt + base, scale); + } + return; + } + } + +#pragma unroll + for (size_t j = 0; j < NUM_ELTS; ++j) { + if (j < valid_cols) { + float elt = static_cast(input.data.elt[j]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + output.data.elt[j] = static_cast(elt * scale); + } + } +} + +template +__device__ __forceinline__ void scaled_fp8_cvt_vec_full(const Vec &input, + Vec &output, + const float scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr bool kUseFastCvt = + !IS_ACT && (NUM_ELTS % 4 == 0) && supports_fast_scaled_fp8_cvt_4::value; +#else + constexpr bool kUseFastCvt = false; +#endif + + if constexpr (kUseFastCvt) { +#pragma unroll + for (size_t base = 0; base < NUM_ELTS; base += 4) { + fast_scaled_fp8_cvt_4(input.data.elt + base, output.data.elt + base, scale); + } + } else { +#pragma unroll + for (size_t j = 0; j < NUM_ELTS; ++j) { + float elt = static_cast(input.data.elt[j]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + output.data.elt[j] = static_cast(elt * scale); + } + } +} + +template +__device__ __forceinline__ void store_fp8_vec_streaming(OType *ptr, + const Vec &vec) { + constexpr size_t bytes = Vec::BYTES; + if constexpr (bytes == 4) { + const uint32_t value = vec.data.vec; + asm volatile("st.global.cs.u32 [%0], %1;" ::"l"(ptr), "r"(value) : "memory"); + } else if constexpr (bytes == 8) { + const uint64_t value = vec.data.vec; + asm volatile("st.global.cs.u64 [%0], %1;" ::"l"(ptr), "l"(value) : "memory"); + } else if constexpr (bytes == 16) { + const uint4 value = vec.data.vec; + asm volatile("st.global.cs.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(value.x), + "r"(value.y), "r"(value.z), "r"(value.w) + : "memory"); + } else { + vec.store_to(ptr); + } +} + +// Copy a warp's contiguous SEG-byte column segment from shared to global using +// the widest store the destination alignment allows. The destination alignment +// is warp-uniform (all lanes of a warp target the same column, so the segment +// base offset is identical), so the width selection below never diverges. The +// shared source is contiguous and 8-byte aligned, so every width up to 8 bytes +// is safe to read. Striping element ``e = k*WARP + lane`` keeps each store +// instruction fully coalesced across the warp regardless of width. +template +__device__ __forceinline__ void copy_column_segment_strided(char *gl, const char *sh, size_t lane) { + static_assert(SEG % (sizeof(T) * THREADS_PER_WARP) == 0, + "segment must tile evenly into warp-wide element stripes"); + constexpr size_t STEPS = SEG / sizeof(T) / THREADS_PER_WARP; + T *const g = reinterpret_cast(gl); + const T *const s = reinterpret_cast(sh); +#pragma unroll + for (size_t k = 0; k < STEPS; ++k) { + const size_t e = k * THREADS_PER_WARP + lane; + g[e] = s[e]; + } +} + +template +__device__ __forceinline__ void store_column_segment(OVecT *shared_segment, char *gl, size_t lane) { + // The segment is the WARP contiguous OVecT slots produced for one column. + constexpr size_t SEG = THREADS_PER_WARP * sizeof(OVecT); + const char *const sh = reinterpret_cast(shared_segment); + const uint64_t a = reinterpret_cast(gl); + if ((a & 7) == 0) { + copy_column_segment_strided(gl, sh, lane); + } else if ((a & 3) == 0) { + copy_column_segment_strided(gl, sh, lane); + } else if ((a & 1) == 0) { + copy_column_segment_strided(gl, sh, lane); + } else { + copy_column_segment_strided(gl, sh, lane); + } +} + +template +__global__ void __launch_bounds__(THREADS_PER_TILE) + group_cast_fp8_kernel(const IType *__restrict__ input, OType *__restrict__ output_rowwise, + OType *__restrict__ output_colwise, const float *__restrict__ scale_ptr, + const float *__restrict__ noop, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *__restrict__ offsets_ptr, + const int64_t *__restrict__ first_dims_ptr, + const int64_t *__restrict__ last_dims_ptr) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + constexpr bool ROWWISE_OUTPUT = + (SCALING_TYPE == ScalingType::ROWWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + constexpr bool COLWISE_OUTPUT = + (SCALING_TYPE == ScalingType::COLWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + constexpr size_t LOAD_SIZE_BYTES = + COLWISE_OUTPUT ? TRANSPOSE_LOAD_SIZE_BYTES : ROWWISE_LOAD_SIZE_BYTES; + constexpr size_t STORE_SIZE_BYTES = + COLWISE_OUTPUT ? TRANSPOSE_STORE_SIZE_BYTES : ROWWISE_STORE_SIZE_BYTES; + constexpr size_t kNVecIn = LOAD_SIZE_BYTES / sizeof(IType); + constexpr size_t kNVecOut = STORE_SIZE_BYTES / sizeof(OType); + constexpr size_t tile_dim_m = THREADS_PER_WARP * kNVecOut; + constexpr size_t tile_dim_n = THREADS_PER_WARP * kNVecIn; + constexpr size_t num_iterations = THREADS_PER_WARP / WARPS_PER_TILE; + + using IVecT = Vec; + using OVecC = Vec; + using OVecT = Vec; + + size_t tile_col = blockIdx.x * tile_dim_n; + size_t tensor_id = 0; + size_t rows = 0; + size_t tensor_base = 0; + size_t tile_row = 0; + size_t cols = last_logical_dim; + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + if (tile_col >= last_logical_dim) { + return; + } + const size_t rows_per_tensor = first_logical_dim / num_tensors; + const size_t tiles_per_tensor = DIVUP(rows_per_tensor, tile_dim_m); + if (tiles_per_tensor == 0) { + return; + } + tensor_id = blockIdx.y / tiles_per_tensor; + if (tensor_id >= num_tensors) { + return; + } + rows = rows_per_tensor; + tensor_base = tensor_id * rows * cols; + tile_row = (blockIdx.y - tensor_id * tiles_per_tensor) * tile_dim_m; + } else if constexpr (SHAPE_REP == ShapeRepresentation::VARYING_LAST_DIM) { + rows = first_logical_dim; + const size_t row_tiles_per_tensor = DIVUP(rows, tile_dim_m); + if (row_tiles_per_tensor == 0) { + return; + } + tensor_id = blockIdx.y / row_tiles_per_tensor; + if (tensor_id >= num_tensors) { + return; + } + cols = static_cast(last_dims_ptr[tensor_id]); + tensor_base = static_cast(offsets_ptr[tensor_id]); + tile_row = (blockIdx.y - tensor_id * row_tiles_per_tensor) * tile_dim_m; + if (tile_col >= cols) { + return; + } + } else { + if (tile_col >= last_logical_dim) { + return; + } + // The grid is sized to the (possibly over-allocated) logical first dim, so + // for capacity-padded buffers up to half the Y blocks fall past the active + // region. Bound the active tile count from offsets_ptr[num_tensors] (the + // active element total, one cached load) and drop those blocks in O(1) + // instead of paying the full O(num_tensors) scan below. This is graph-safe: + // it reads device offsets rather than requiring a host-side active count. + const size_t active_rows = static_cast(offsets_ptr[num_tensors]) / last_logical_dim; + const size_t max_active_tiles = DIVUP(active_rows, tile_dim_m) + num_tensors; + if (blockIdx.y >= max_active_tiles) { + return; + } + size_t local_tile_id = blockIdx.y; + bool found_tensor = false; + for (size_t i = 0; i < num_tensors; ++i) { + const size_t tensor_rows = static_cast(first_dims_ptr[i]); + const size_t tensor_tiles = DIVUP(tensor_rows, tile_dim_m); + if (local_tile_id < tensor_tiles) { + tensor_id = i; + rows = tensor_rows; + tensor_base = static_cast(offsets_ptr[i]); + tile_row = local_tile_id * tile_dim_m; + found_tensor = true; + break; + } + local_tile_id -= tensor_tiles; + } + if (!found_tensor) { + return; + } + } + + if (rows == 0 || cols == 0) { + return; + } + if (tile_row >= rows) { + return; + } + + const float scale = scale_ptr == nullptr ? 1.0f : scale_ptr[tensor_id]; + const size_t tid = threadIdx.x; + const size_t tidx = tid % THREADS_PER_WARP; + const size_t tidy = tid / THREADS_PER_WARP; + + if constexpr (COLWISE_OUTPUT) { + __shared__ OVecT + shared_output_t[kNVecIn][THREADS_PER_WARP][THREADS_PER_WARP + TRANSPOSE_SHARED_PAD]; + + // Interior tiles (every fragment full) can skip *all* per-fragment bounds + // checks. That instruction overhead -- not occupancy -- is what holds the + // generic kernel to ~58% of HBM; the branchless path below is HBM-bound. + // The condition is block-uniform (derived from tile_row/tile_col/rows/cols), + // so there is no warp divergence: only edge tiles (partial trailing columns, + // or the last row tile of a group) fall through to the bounds-checked path. + // + // Alignment is split between the two phases because they have independent + // requirements. Phase 1 (load bf16 + convert + optional rowwise store + + // transpose into shared) is governed by the *column* stride and is the bulk + // of global traffic (a 2-byte read per element). Phase 2 (store the + // transposed fp8 columns) is governed by the per-group *row* count, which + // for varying-first groups is an arbitrary integer and is therefore rarely + // a multiple of kNVecOut. We still take the branchless interior path + // whenever the loads are aligned -- the common case -- and only the final + // store falls back to per-element stores when the row stride is unaligned, + // so imbalanced group shapes keep the branchless phase-1 win. + const bool full_tile = (tile_row + tile_dim_m <= rows) && (tile_col + tile_dim_n <= cols); + const bool aligned_load = (cols % kNVecIn == 0) && (tensor_base % kNVecIn == 0); + if (full_tile && aligned_load) { +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidy + iter * WARPS_PER_TILE; + const size_t j1 = tidx; + const size_t base_row = tile_row + i1 * kNVecOut; + const size_t base_col = tile_col + j1 * kNVecIn; + OVecT local_output_t[kNVecIn]; +#pragma unroll + for (size_t i2 = 0; i2 < kNVecOut; ++i2) { + const size_t row = base_row + i2; + IVecT local_input; + OVecC local_output; + local_input.load_from(input + tensor_base + row * cols + base_col); + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + if constexpr (ROWWISE_OUTPUT) { + OType *const output_ptr = output_rowwise + tensor_base + row * cols + base_col; + if constexpr (SCALING_TYPE == ScalingType::BIDIMENSIONAL) { + store_fp8_vec_streaming(output_ptr, local_output); + } else { + local_output.store_to(output_ptr); + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + local_output_t[j2].data.elt[i2] = local_output.data.elt[j2]; + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + shared_output_t[j2][j1][i1] = local_output_t[j2]; + } + } + __syncthreads(); +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + // All 32 lanes of a warp share j1 (= tidy + iter*WARPS_PER_TILE) and + // therefore the same column, so shared_output_t[j2][j1][0..WARP-1] is a + // contiguous kNVecOut*WARP-byte segment in shared and maps to a + // contiguous segment of one output column. Store it with the widest + // type the (warp-uniform) destination offset allows: row strides that + // are multiples of kNVecOut give a single vectorized store per lane; + // arbitrary varying-first strides degrade gracefully to narrower + // coalesced stores instead of faulting. + const size_t j1 = tidy + iter * WARPS_PER_TILE; + const size_t col = tile_col + j1 * kNVecIn + j2; + char *const gl = + reinterpret_cast(output_colwise + tensor_base + col * rows + tile_row); + store_column_segment(&shared_output_t[j2][j1][0], gl, tidx); + } + } + return; + } + +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidy + iter * WARPS_PER_TILE; + const size_t j1 = tidx; + const size_t base_row = tile_row + i1 * kNVecOut; + const size_t base_col = tile_col + j1 * kNVecIn; + const size_t fragment_cols = + base_col < cols ? ((cols - base_col) < kNVecIn ? (cols - base_col) : kNVecIn) : 0; + if (base_row >= rows || fragment_cols == 0) { + continue; + } + const bool full_fragment = (base_row + kNVecOut <= rows) && (fragment_cols == kNVecIn); + OVecT local_output_t[kNVecIn]; +#pragma unroll + for (size_t i2 = 0; i2 < kNVecOut; ++i2) { + const size_t row = base_row + i2; + if (full_fragment || row < rows) { + IVecT local_input; + OVecC local_output; + if (full_fragment) { + const IType *const input_ptr = input + tensor_base + row * cols + base_col; + if (reinterpret_cast(input_ptr) % IVecT::BYTES == 0) { + local_input.load_from(input_ptr); + } else { + local_input.load_from_elts(input_ptr); + } + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + if constexpr (ROWWISE_OUTPUT) { + OType *const output_ptr = output_rowwise + tensor_base + row * cols + base_col; + if constexpr (SCALING_TYPE == ScalingType::BIDIMENSIONAL) { + if (reinterpret_cast(output_ptr) % OVecC::BYTES == 0) { + store_fp8_vec_streaming(output_ptr, local_output); + } else { + local_output.store_to_elts(output_ptr); + } + } else if (reinterpret_cast(output_ptr) % OVecC::BYTES == 0) { + local_output.store_to(output_ptr); + } else { + local_output.store_to_elts(output_ptr); + } + } + } else { + local_input.load_from_elts(input + tensor_base + row * cols + base_col, 0, + fragment_cols); + scaled_fp8_cvt_vec(local_input, local_output, fragment_cols, + scale); + if constexpr (ROWWISE_OUTPUT) { + local_output.store_to_elts(output_rowwise + tensor_base + row * cols + base_col, 0, + fragment_cols); + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + if (j2 < fragment_cols) { + local_output_t[j2].data.elt[i2] = local_output.data.elt[j2]; + } + } + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + if (j2 < fragment_cols) { + shared_output_t[j2][j1][i1] = local_output_t[j2]; + } + } + } + + __syncthreads(); +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidx; + const size_t j1 = tidy + iter * WARPS_PER_TILE; + const size_t row = tile_row + i1 * kNVecOut; + const size_t col = tile_col + j1 * kNVecIn + j2; + if (col < cols) { + const size_t valid_rows = + row < rows ? ((rows - row) < kNVecOut ? (rows - row) : kNVecOut) : 0; + if (valid_rows > 0) { + OType *const output_ptr = output_colwise + tensor_base + col * rows + row; + const OVecT local_output_t = shared_output_t[j2][j1][i1]; + if (valid_rows == kNVecOut && + reinterpret_cast(output_ptr) % OVecT::BYTES == 0) { + local_output_t.store_to(output_ptr); + } else { + local_output_t.store_to_elts(output_ptr, 0, valid_rows); + } + } + } + } + } + } else { +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidy + iter * WARPS_PER_TILE; + const size_t j1 = tidx; +#pragma unroll + for (size_t i2 = 0; i2 < kNVecOut; ++i2) { + const size_t row = tile_row + i1 * kNVecOut + i2; + const size_t col = tile_col + j1 * kNVecIn; + if (row < rows) { + IVecT local_input; + OVecC local_output; + if (col + kNVecIn <= cols) { + const IType *const input_ptr = input + tensor_base + row * cols + col; + if (reinterpret_cast(input_ptr) % IVecT::BYTES == 0) { + local_input.load_from(input_ptr); + } else { + local_input.load_from_elts(input_ptr); + } + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + OType *const output_ptr = output_rowwise + tensor_base + row * cols + col; + if (reinterpret_cast(output_ptr) % OVecC::BYTES == 0) { + local_output.store_to(output_ptr); + } else { + local_output.store_to_elts(output_ptr); + } + } else { + const size_t valid_cols = + col < cols ? ((cols - col) < kNVecIn ? (cols - col) : kNVecIn) : 0; + if (valid_cols == 0) { + continue; + } + local_input.load_from_elts(input + tensor_base + row * cols + col, 0, valid_cols); + scaled_fp8_cvt_vec(local_input, local_output, valid_cols, scale); + local_output.store_to_elts(output_rowwise + tensor_base + row * cols + col, 0, + valid_cols); + } + } + } + } + } +} + +template +__global__ void __launch_bounds__(ROWWISE_FLAT_THREADS) + group_cast_fp8_rowwise_flat_kernel(const IType *__restrict__ input, + OType *__restrict__ output_rowwise, + const float *__restrict__ scale_ptr, + const float *__restrict__ noop, const size_t rows_per_tensor, + const size_t cols, const size_t vecs_per_tensor) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + constexpr size_t nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); + using IVecT = Vec; + using OVecT = Vec; + + const size_t tensor_id = blockIdx.y; + const size_t tensor_base = tensor_id * rows_per_tensor * cols; + const float scale = scale_ptr == nullptr ? 1.0f : scale_ptr[tensor_id]; + + for (size_t vec_id = blockIdx.x * blockDim.x + threadIdx.x; vec_id < vecs_per_tensor; + vec_id += gridDim.x * blockDim.x) { + const size_t offset = tensor_base + vec_id * nvec; + IVecT local_input; + OVecT local_output; + local_input.load_from(input + offset); + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + local_output.store_to(output_rowwise + offset); + } +} + +// Flat, dynamically load-balanced varying-first rowwise group cast. The block +// -> tensor mapping (1D grid, shared-mem prefix sum, binary search, +// size-proportional chunks) distributes work proportionally to each group's +// size. Requires vector-aligned inputs (row length a multiple of nvec and base +// pointers vector aligned); the caller guarantees this, so every access is a +// full-width vectorized load/store with no per-iteration masking. +template +__global__ void __launch_bounds__(ROWWISE_FLAT_THREADS) + group_cast_fp8_varying_first_rowwise_aligned_flat_kernel( + const IType *__restrict__ input, OType *__restrict__ output_rowwise, + const float *__restrict__ scale_ptr, const float *__restrict__ noop, + const size_t num_tensors, const size_t total_elements, + const int64_t *__restrict__ offsets_ptr, const size_t target_blocks) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + constexpr size_t nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); + using IVecT = Vec; + using OVecT = Vec; + + // Dynamic shared memory layout: + // s_offsets: int64_t[num_tensors + 1] + // s_block_offsets: int[num_tensors + 1] + extern __shared__ char s_mem[]; + int64_t *s_offsets = reinterpret_cast(s_mem); + int *s_block_offsets = reinterpret_cast(s_mem + (num_tensors + 1) * sizeof(int64_t)); + + const size_t tid = threadIdx.x; + + // 1. Copy offsets to shared memory + for (size_t i = tid; i <= num_tensors; i += blockDim.x) { + s_offsets[i] = offsets_ptr[i]; + } + __syncthreads(); + + // Derive the per-block chunk size from the *active* element count + // (offsets_ptr[num_tensors]) rather than the over-allocated buffer, so that + // over-allocation does not inflate the chunk and reduce the number of active + // blocks. Every thread computes the same value from shared memory. + const size_t active_elements = static_cast(s_offsets[num_tensors]); + size_t block_chunk_size = DIVUP(active_elements, target_blocks); + block_chunk_size = DIVUP(block_chunk_size, nvec) * nvec; + if (block_chunk_size < 8192) { + block_chunk_size = 8192; + } + + // 2. Compute block counts and prefix sum. + // Do the expensive per-tensor ceil-division in parallel (one tensor per + // thread); the remaining exclusive scan is then just integer adds. This keeps + // the single-thread work O(num_tensors) cheap adds instead of O(num_tensors) + // divisions running on lane 0 alone -- the latter scales with group count and + // showed up as per-warp divergence / exposed block startup in profiling. + for (size_t i = tid; i < num_tensors; i += blockDim.x) { + const size_t tensor_size = s_offsets[i + 1] - s_offsets[i]; + s_block_offsets[i] = static_cast(DIVUP(tensor_size, block_chunk_size)); + } + __syncthreads(); + if (tid == 0) { + int sum = 0; + for (size_t i = 0; i < num_tensors; ++i) { + const int blocks = s_block_offsets[i]; + s_block_offsets[i] = sum; // exclusive prefix + sum += blocks; + } + s_block_offsets[num_tensors] = sum; + } + __syncthreads(); + + // Guard against blocks that are out of bounds of the active work + if (blockIdx.x >= s_block_offsets[num_tensors]) { + return; + } + + // 3. Binary search to find which tensor this block belongs to + const size_t tensor_id = transformer_engine::dispatch::common::find_tensor_from_offsets( + s_block_offsets, num_tensors, blockIdx.x); + + // 4. Calculate this block's local work range within the tensor + const size_t local_block_id = blockIdx.x - s_block_offsets[tensor_id]; + const size_t tensor_base = s_offsets[tensor_id]; + const size_t tensor_size = s_offsets[tensor_id + 1] - tensor_base; + const size_t start_elt = local_block_id * block_chunk_size; + const size_t end_elt = + start_elt + block_chunk_size < tensor_size ? start_elt + block_chunk_size : tensor_size; + + if (start_elt >= end_elt) { + return; + } + + const IType *base_input = input + tensor_base + start_elt; + OType *base_output = output_rowwise + tensor_base + start_elt; + const size_t numel = end_elt - start_elt; + const float scale = scale_ptr == nullptr ? 1.0f : scale_ptr[tensor_id]; + + // Inputs are guaranteed vector aligned by the caller, so chunk start and + // length are nvec-aligned and every access is a full-width vectorized + // load/store with no per-iteration masking. + const size_t tensor_vecs = numel / nvec; + for (size_t local_vec_id = tid; local_vec_id < tensor_vecs; local_vec_id += blockDim.x) { + const size_t offset = local_vec_id * nvec; + IVecT local_input; + OVecT local_output; + local_input.load_from(base_input + offset); + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + local_output.store_to(base_output + offset); + } +} + +template +__global__ void __launch_bounds__(ROWWISE_FLAT_THREADS) group_cast_fp8_variable_rowwise_flat_kernel( + const IType *__restrict__ input, OType *__restrict__ output_rowwise, + const float *__restrict__ scale_ptr, const float *__restrict__ noop, const size_t num_tensors, + const size_t total_elements, const int64_t *__restrict__ offsets_ptr) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + constexpr size_t nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); + using IVecT = Vec; + using OVecT = Vec; + + const size_t active_elements = static_cast(offsets_ptr[num_tensors]); + const size_t bounded_elements = + active_elements < total_elements ? active_elements : total_elements; + const size_t total_vecs = DIVUP(bounded_elements, nvec); + for (size_t vec_id = blockIdx.x * blockDim.x + threadIdx.x; vec_id < total_vecs; + vec_id += gridDim.x * blockDim.x) { + const size_t offset = vec_id * nvec; + const size_t tensor_id = transformer_engine::dispatch::common::find_tensor_from_offsets( + offsets_ptr, num_tensors, offset); + const size_t tensor_end = static_cast(offsets_ptr[tensor_id + 1]); + const float scale = scale_ptr == nullptr ? 1.0f : scale_ptr[tensor_id]; + IVecT local_input; + OVecT local_output; + if (offset + nvec <= tensor_end && offset + nvec <= bounded_elements && + reinterpret_cast(input + offset) % IVecT::BYTES == 0 && + reinterpret_cast(output_rowwise + offset) % OVecT::BYTES == 0) { + local_input.load_from(input + offset); + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + local_output.store_to(output_rowwise + offset); + } else { + const size_t valid_elts = + offset < bounded_elements + ? ((bounded_elements - offset) < nvec ? (bounded_elements - offset) : nvec) + : 0; + if (valid_elts == 0) { + continue; + } +#pragma unroll + for (size_t i = 0; i < nvec; ++i) { + if (i < valid_elts) { + const size_t elt_offset = offset + i; + const size_t elt_tensor_id = + elt_offset < tensor_end + ? tensor_id + : transformer_engine::dispatch::common::find_tensor_from_offsets( + offsets_ptr, num_tensors, elt_offset); + const float elt_scale = scale_ptr == nullptr ? 1.0f : scale_ptr[elt_tensor_id]; + float elt = static_cast(input[elt_offset]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + local_output.data.elt[i] = static_cast(elt * elt_scale); + } + } + local_output.store_to_elts(output_rowwise + offset, 0, valid_elts); + } + } +} + +template +__global__ void __launch_bounds__(THREADS_PER_TILE) + group_cast_fp8_same_shape_full_tile_kernel(const IType *__restrict__ input, + OType *__restrict__ output_rowwise, + OType *__restrict__ output_colwise, + const float *__restrict__ scale_ptr, + const float *__restrict__ noop, + const size_t rows_per_tensor, const size_t cols) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + constexpr bool ROWWISE_OUTPUT = + (SCALING_TYPE == ScalingType::ROWWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + constexpr bool COLWISE_OUTPUT = + (SCALING_TYPE == ScalingType::COLWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + static_assert(COLWISE_OUTPUT, "The full-tile grouped FP8 kernel is for columnwise outputs."); + + constexpr size_t kNVecIn = TRANSPOSE_LOAD_SIZE_BYTES / sizeof(IType); + constexpr size_t kNVecOut = TRANSPOSE_STORE_SIZE_BYTES / sizeof(OType); + constexpr size_t tile_dim_m = THREADS_PER_WARP * kNVecOut; + constexpr size_t tile_dim_n = THREADS_PER_WARP * kNVecIn; + constexpr size_t num_iterations = THREADS_PER_WARP / WARPS_PER_TILE; + + using IVecT = Vec; + using OVecC = Vec; + using OVecT = Vec; + + const size_t tiles_per_tensor = rows_per_tensor / tile_dim_m; + const size_t tensor_id = blockIdx.y / tiles_per_tensor; + const size_t tensor_tile_id = blockIdx.y - tensor_id * tiles_per_tensor; + const size_t tensor_base = tensor_id * rows_per_tensor * cols; + const size_t tile_row = tensor_tile_id * tile_dim_m; + const size_t tile_col = blockIdx.x * tile_dim_n; + const float scale = scale_ptr == nullptr ? 1.0f : scale_ptr[tensor_id]; + + const size_t tid = threadIdx.x; + const size_t tidx = tid % THREADS_PER_WARP; + const size_t tidy = tid / THREADS_PER_WARP; + + __shared__ OVecT + shared_output_t[kNVecIn][THREADS_PER_WARP][THREADS_PER_WARP + TRANSPOSE_SHARED_PAD]; + +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidy + iter * WARPS_PER_TILE; + const size_t j1 = tidx; + const size_t base_row = tile_row + i1 * kNVecOut; + const size_t base_col = tile_col + j1 * kNVecIn; + OVecT local_output_t[kNVecIn]; +#pragma unroll + for (size_t i2 = 0; i2 < kNVecOut; ++i2) { + const size_t row = base_row + i2; + IVecT local_input; + OVecC local_output; + const IType *const input_ptr = input + tensor_base + row * cols + base_col; + local_input.load_from(input_ptr); + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + if constexpr (ROWWISE_OUTPUT) { + OType *const output_ptr = output_rowwise + tensor_base + row * cols + base_col; + if constexpr (SCALING_TYPE == ScalingType::BIDIMENSIONAL) { + store_fp8_vec_streaming(output_ptr, local_output); + } else { + local_output.store_to(output_ptr); + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + local_output_t[j2].data.elt[i2] = local_output.data.elt[j2]; + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + shared_output_t[j2][j1][i1] = local_output_t[j2]; + } + } + + __syncthreads(); +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidx; + const size_t j1 = tidy + iter * WARPS_PER_TILE; + const size_t row = tile_row + i1 * kNVecOut; + const size_t col = tile_col + j1 * kNVecIn + j2; + OType *const output_ptr = output_colwise + tensor_base + col * rows_per_tensor + row; + const OVecT local_output_t = shared_output_t[j2][j1][i1]; + local_output_t.store_to(output_ptr); + } + } +} + +template +__global__ void __launch_bounds__(THREADS_PER_TILE) group_cast_fp8_varying_first_tile_kernel( + const IType *__restrict__ input, OType *__restrict__ output_rowwise, + OType *__restrict__ output_colwise, const float *__restrict__ scale_ptr, + const float *__restrict__ noop, const size_t num_tensors, const size_t rows_upper_bound, + const size_t cols, const size_t total_elements, const int64_t *__restrict__ offsets_ptr, + const int64_t *__restrict__ first_dims_ptr) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + constexpr bool ROWWISE_OUTPUT = + (SCALING_TYPE == ScalingType::ROWWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + constexpr bool COLWISE_OUTPUT = + (SCALING_TYPE == ScalingType::COLWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + static_assert(COLWISE_OUTPUT, "The varying-first tile kernel is for columnwise outputs."); + + constexpr size_t kNVecIn = TRANSPOSE_LOAD_SIZE_BYTES / sizeof(IType); + constexpr size_t kNVecOut = TRANSPOSE_STORE_SIZE_BYTES / sizeof(OType); + constexpr size_t tile_dim_m = THREADS_PER_WARP * kNVecOut; + constexpr size_t tile_dim_n = THREADS_PER_WARP * kNVecIn; + constexpr size_t num_iterations = THREADS_PER_WARP / WARPS_PER_TILE; + + using IVecT = Vec; + using OVecC = Vec; + using OVecT = Vec; + + const size_t tensor_id = blockIdx.z; + if (tensor_id >= num_tensors || cols == 0) { + return; + } + + const size_t raw_tensor_base = static_cast(offsets_ptr[tensor_id]); + const size_t raw_tensor_end = static_cast(offsets_ptr[tensor_id + 1]); + const size_t tensor_base = raw_tensor_base < total_elements ? raw_tensor_base : total_elements; + const size_t tensor_end = raw_tensor_end < total_elements ? raw_tensor_end : total_elements; + if (tensor_end <= tensor_base) { + return; + } + + const size_t rows_from_offsets = (tensor_end - tensor_base) / cols; + const size_t rows_from_dims = static_cast(first_dims_ptr[tensor_id]); + size_t rows = rows_from_dims < rows_from_offsets ? rows_from_dims : rows_from_offsets; + rows = rows < rows_upper_bound ? rows : rows_upper_bound; + if (rows == 0) { + return; + } + + const size_t tile_col = blockIdx.x * tile_dim_n; + if (tile_col >= cols) { + return; + } + const size_t row_tiles = DIVUP(rows, tile_dim_m); + const float scale = scale_ptr == nullptr ? 1.0f : scale_ptr[tensor_id]; + const size_t tid = threadIdx.x; + const size_t tidx = tid % THREADS_PER_WARP; + const size_t tidy = tid / THREADS_PER_WARP; + + __shared__ OVecT + shared_output_t[kNVecIn][THREADS_PER_WARP][THREADS_PER_WARP + TRANSPOSE_SHARED_PAD]; + + for (size_t tensor_tile_id = blockIdx.y; tensor_tile_id < row_tiles; + tensor_tile_id += gridDim.y) { + const size_t tile_row = tensor_tile_id * tile_dim_m; + +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidy + iter * WARPS_PER_TILE; + const size_t j1 = tidx; + const size_t base_row = tile_row + i1 * kNVecOut; + const size_t base_col = tile_col + j1 * kNVecIn; + const size_t fragment_cols = + base_col < cols ? ((cols - base_col) < kNVecIn ? (cols - base_col) : kNVecIn) : 0; + if (base_row >= rows || fragment_cols == 0) { + continue; + } + const bool full_fragment = (base_row + kNVecOut <= rows) && (fragment_cols == kNVecIn); + OVecT local_output_t[kNVecIn]; +#pragma unroll + for (size_t i2 = 0; i2 < kNVecOut; ++i2) { + const size_t row = base_row + i2; + if (full_fragment || row < rows) { + IVecT local_input; + OVecC local_output; + if (full_fragment) { + const IType *const input_ptr = input + tensor_base + row * cols + base_col; + if (reinterpret_cast(input_ptr) % IVecT::BYTES == 0) { + local_input.load_from(input_ptr); + } else { + local_input.load_from_elts(input_ptr); + } + scaled_fp8_cvt_vec_full(local_input, local_output, scale); + if constexpr (ROWWISE_OUTPUT) { + OType *const output_ptr = output_rowwise + tensor_base + row * cols + base_col; + if constexpr (SCALING_TYPE == ScalingType::BIDIMENSIONAL) { + if (reinterpret_cast(output_ptr) % OVecC::BYTES == 0) { + store_fp8_vec_streaming(output_ptr, local_output); + } else { + local_output.store_to_elts(output_ptr); + } + } else if (reinterpret_cast(output_ptr) % OVecC::BYTES == 0) { + local_output.store_to(output_ptr); + } else { + local_output.store_to_elts(output_ptr); + } + } + } else { + local_input.load_from_elts(input + tensor_base + row * cols + base_col, 0, + fragment_cols); + scaled_fp8_cvt_vec(local_input, local_output, fragment_cols, + scale); + if constexpr (ROWWISE_OUTPUT) { + local_output.store_to_elts(output_rowwise + tensor_base + row * cols + base_col, 0, + fragment_cols); + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + if (j2 < fragment_cols) { + local_output_t[j2].data.elt[i2] = local_output.data.elt[j2]; + } + } + } + } +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { + if (j2 < fragment_cols) { + shared_output_t[j2][j1][i1] = local_output_t[j2]; + } + } + } + + __syncthreads(); +#pragma unroll + for (size_t j2 = 0; j2 < kNVecIn; ++j2) { +#pragma unroll + for (size_t iter = 0; iter < num_iterations; ++iter) { + const size_t i1 = tidx; + const size_t j1 = tidy + iter * WARPS_PER_TILE; + const size_t row = tile_row + i1 * kNVecOut; + const size_t col = tile_col + j1 * kNVecIn + j2; + if (col < cols) { + const size_t valid_rows = + row < rows ? ((rows - row) < kNVecOut ? (rows - row) : kNVecOut) : 0; + if (valid_rows > 0) { + OType *const output_ptr = output_colwise + tensor_base + col * rows + row; + const OVecT local_output_t = shared_output_t[j2][j1][i1]; + if (valid_rows == kNVecOut && + reinterpret_cast(output_ptr) % OVecT::BYTES == 0) { + local_output_t.store_to(output_ptr); + } else { + local_output_t.store_to_elts(output_ptr, 0, valid_rows); + } + } + } + } + } + __syncthreads(); + } +} + +} // namespace group_quantize_kernel + +template +void group_quantize(const GroupedTensor *input, const Tensor *noop, GroupedTensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { + using namespace group_quantize_kernel; + (void)quant_config; + + CheckNoopTensor(*noop, "cast_noop"); + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Number of input and output tensors must be same."); + NVTE_CHECK(input->has_data(), "Cannot quantize tensor without rowwise data."); + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); + NVTE_CHECK(output->scale.has_data(), "Grouped FP8 tensor-scaling output scale must be set."); + NVTE_CHECK(output->scale.dtype == DType::kFloat32, + "Grouped FP8 tensor-scaling scale must be FP32."); + NVTE_CHECK(output->scale.numel() >= output->num_tensors, + "Grouped FP8 tensor-scaling scale must have at least one entry per tensor."); + const bool use_rowwise_output = output->has_data(); + const bool use_colwise_output = output->has_columnwise_data(); + NVTE_CHECK(use_rowwise_output || use_colwise_output, + "Either rowwise or columnwise output data need to be allocated."); + if (use_rowwise_output) { + NVTE_CHECK(output->scale_inv.has_data(), "Rowwise scale_inv must be allocated."); + } + if (use_colwise_output) { + NVTE_CHECK(output->columnwise_scale_inv.has_data(), "Columnwise scale_inv must be allocated."); + } + + ScalingType scaling_type = ScalingType::BIDIMENSIONAL; + if (!use_colwise_output) { + scaling_type = ScalingType::ROWWISE; + } else if (!use_rowwise_output) { + scaling_type = ScalingType::COLWISE; + } + + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (output->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (output->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (output->all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else { + NVTE_ERROR( + "Grouped FP8 tensor-scaling quantization only supports same-shape, varying " + "first-dimension, or varying last-dimension grouped tensors."); + } + + NVTE_CHECK(input->logical_shape.data[1] == output->logical_shape.data[1], + "Grouped FP8 tensor-scaling input and output must have the same last dimension."); + NVTE_CHECK(output->logical_shape.data[0] <= input->logical_shape.data[0], + "Grouped FP8 tensor-scaling output first dimension must not exceed input first " + "dimension."); + + // For varying-dim grouped tensors, logical_shape may be larger than the active region + // (sum of first_dims for varying-first, sum of last_dims for varying-last). The backing + // buffer is sized to logical_shape, but the kernel must only touch the active rows/cols; + // metadata (first_dims/last_dims/tensor_offsets) is consulted on device to skip the unused + // tail. We size the grid from logical_shape so the launch covers every potential payload + // row, and rely on per-block bounds checks to drop blocks past the active region. + const size_t first_logical_dim = output->logical_shape.data[0]; + const size_t last_logical_dim = output->logical_shape.data[1]; + const size_t num_tensors = input->num_tensors; + if (first_logical_dim == 0 || last_logical_dim == 0) { + return; + } + const size_t store_size_bytes = + use_colwise_output ? TRANSPOSE_STORE_SIZE_BYTES : ROWWISE_STORE_SIZE_BYTES; + const size_t row_tile_size = THREADS_PER_WARP * store_size_bytes; + size_t work_blocks_Y = DIVUP(first_logical_dim, row_tile_size); + if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { + NVTE_CHECK(first_logical_dim % num_tensors == 0, + "First logical dimension must be divisible by num_tensors."); + const size_t rows_per_tensor = first_logical_dim / num_tensors; + work_blocks_Y = num_tensors * DIVUP(rows_per_tensor, row_tile_size); + } else if (shape_rep == ShapeRepresentation::VARYING_FIRST_DIM) { + // first_dims live on device; over-allocate the grid enough to cover one partial tile per group. + work_blocks_Y += num_tensors; + } else if (shape_rep == ShapeRepresentation::VARYING_LAST_DIM) { + // last_dims live on device; map Y to per-group row tiles and over-allocate X by total width. + work_blocks_Y = num_tensors * DIVUP(first_logical_dim, row_tile_size); + } + + const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(output->last_dims.dptr); + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + + const dim3 block(THREADS_PER_TILE); + + const float *const scale_ptr = reinterpret_cast(output->scale.dptr); + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_SCALING_TYPE_SWITCH( + scaling_type, SCALING_TYPE, + { + constexpr bool colwise_output = (SCALING_TYPE == ScalingType::COLWISE) || + (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + constexpr size_t load_size_bytes = + colwise_output ? TRANSPOSE_LOAD_SIZE_BYTES : ROWWISE_LOAD_SIZE_BYTES; + constexpr size_t kNVecIn = load_size_bytes / sizeof(IType); + constexpr size_t tile_dim_n = THREADS_PER_WARP * kNVecIn; + TRANSFORMER_ENGINE_GROUP_TENSOR_SHAPE_REPRESENTATION_SWITCH(shape_rep, SHAPE_REP, { + if constexpr (SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + NVTE_CHECK(offsets_ptr != nullptr && first_dims_ptr != nullptr, + "Varying first-dimension grouped FP8 quantization requires " + "first_dims and tensor_offsets."); + } else if constexpr (SHAPE_REP == ShapeRepresentation::VARYING_LAST_DIM) { + NVTE_CHECK(offsets_ptr != nullptr && last_dims_ptr != nullptr, + "Varying last-dimension grouped FP8 quantization requires " + "last_dims and tensor_offsets."); + } + bool launched_fast_path = false; + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + const size_t rows_per_tensor = first_logical_dim / num_tensors; + if constexpr (SCALING_TYPE == ScalingType::ROWWISE) { + constexpr size_t flat_nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); + using IVecT = Vec; + using OVecT = Vec; + const size_t elems_per_tensor = rows_per_tensor * last_logical_dim; + const bool flat_aligned = + reinterpret_cast(input->data.dptr) % IVecT::BYTES == 0 && + reinterpret_cast(output->data.dptr) % OVecT::BYTES == 0; + if (elems_per_tensor % flat_nvec == 0 && flat_aligned) { + const size_t vecs_per_tensor = elems_per_tensor / flat_nvec; + const dim3 flat_block(ROWWISE_FLAT_THREADS); + const dim3 flat_grid(DIVUP(vecs_per_tensor, ROWWISE_FLAT_THREADS), + num_tensors); + group_cast_fp8_rowwise_flat_kernel + <<>>( + reinterpret_cast(input->data.dptr), + reinterpret_cast(output->data.dptr), scale_ptr, noop_ptr, + rows_per_tensor, last_logical_dim, vecs_per_tensor); + launched_fast_path = true; + } + } else if constexpr (colwise_output) { + constexpr size_t store_size_bytes = TRANSPOSE_STORE_SIZE_BYTES; + constexpr size_t kNVecOut = store_size_bytes / sizeof(OType); + constexpr size_t tile_dim_m = THREADS_PER_WARP * kNVecOut; + if (rows_per_tensor % tile_dim_m == 0 && last_logical_dim % tile_dim_n == 0) { + const dim3 full_grid(last_logical_dim / tile_dim_n, + num_tensors * (rows_per_tensor / tile_dim_m)); + group_cast_fp8_same_shape_full_tile_kernel + <<>>( + reinterpret_cast(input->data.dptr), + use_rowwise_output ? reinterpret_cast(output->data.dptr) + : nullptr, + reinterpret_cast(output->columnwise_data.dptr), scale_ptr, + noop_ptr, rows_per_tensor, last_logical_dim); + launched_fast_path = true; + } + } + } + if constexpr (SHAPE_REP == ShapeRepresentation::VARYING_LAST_DIM) { + if constexpr (SCALING_TYPE == ScalingType::ROWWISE) { + const size_t total_elements = first_logical_dim * last_logical_dim; + constexpr size_t flat_nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); + const size_t total_vecs = DIVUP(total_elements, flat_nvec); + const dim3 flat_block(ROWWISE_FLAT_THREADS); + const dim3 flat_grid(DIVUP(total_vecs, ROWWISE_FLAT_THREADS)); + group_cast_fp8_variable_rowwise_flat_kernel + <<>>( + reinterpret_cast(input->data.dptr), + reinterpret_cast(output->data.dptr), scale_ptr, noop_ptr, + num_tensors, total_elements, offsets_ptr); + launched_fast_path = true; + } + } + if constexpr (SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + if constexpr (SCALING_TYPE == ScalingType::ROWWISE) { + const size_t total_elements = first_logical_dim * last_logical_dim; + constexpr size_t flat_nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); + using IVecT = Vec; + using OVecT = Vec; + const bool flat_aligned = + last_logical_dim % flat_nvec == 0 && + reinterpret_cast(input->data.dptr) % IVecT::BYTES == 0 && + reinterpret_cast(output->data.dptr) % OVecT::BYTES == 0; + // Only vector-aligned inputs are supported. A non-multiple + // row length or a misaligned base pointer would force + // uncoalesced scalar access, which we deliberately do not + // support on this path. + NVTE_CHECK( + flat_aligned, + "Varying first-dimension rowwise grouped FP8 quantization requires " + "vector-aligned inputs: the last logical dimension must be a " + "multiple of ", + flat_nvec, " and the input/output base pointers must be vector aligned."); + const dim3 flat_block(ROWWISE_FLAT_THREADS); + const int num_sms = ::transformer_engine::cuda::sm_count(); + // Target ~8 blocks per SM of *active* work. The actual + // per-block chunk size is derived inside the kernel from + // the device-side offsets (the active element count), so + // that an over-allocated backing buffer does not inflate + // the chunk and starve the grid of active blocks. We + // launch enough blocks to cover target_blocks plus up to + // one partial block per tensor (from per-tensor rounding). + const size_t target_blocks = 8 * num_sms; + const size_t grid_size = target_blocks + num_tensors; + const size_t shared_mem_bytes = + (num_tensors + 1) * sizeof(int64_t) + (num_tensors + 1) * sizeof(int); + group_cast_fp8_varying_first_rowwise_aligned_flat_kernel + <<>>( + reinterpret_cast(input->data.dptr), + reinterpret_cast(output->data.dptr), scale_ptr, noop_ptr, + num_tensors, total_elements, offsets_ptr, target_blocks); + launched_fast_path = true; + } + } + // Varying-first columnwise intentionally falls through to the + // generic group_cast_fp8_kernel below. The dedicated + // group_cast_fp8_varying_first_tile_kernel used a grid-stride + // row loop whose high register pressure limited occupancy. The + // generic kernel processes one tile per block and now carries a + // branchless interior-tile fast path, so it is both lighter on + // registers and HBM-bound on the interior of every group. + if (!launched_fast_path) { + const size_t work_blocks_X = DIVUP(last_logical_dim, tile_dim_n); + const dim3 grid(work_blocks_X, work_blocks_Y); + group_cast_fp8_kernel<<>>( + reinterpret_cast(input->data.dptr), + use_rowwise_output ? reinterpret_cast(output->data.dptr) : nullptr, + use_colwise_output ? reinterpret_cast(output->columnwise_data.dptr) + : nullptr, + scale_ptr, noop_ptr, num_tensors, first_logical_dim, last_logical_dim, + offsets_ptr, first_dims_ptr, last_dims_ptr); + } + NVTE_CHECK_CUDA(cudaGetLastError()); + }); + })); // NOLINT(*) + ); // NOLINT(*) +} + +} // namespace fp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_CUH_ diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index cad27a2992..47539a89a1 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -99,6 +99,20 @@ void nvte_compute_amax(const NVTETensor input, NVTETensor output, cudaStream_t s void nvte_compute_amax_with_config(const NVTETensor input, NVTETensor output, const NVTEQuantizationConfig config, cudaStream_t stream); +/*! \brief Compute per-group FP8 amax values for a grouped tensor. + * + * The grouped tensor's shape metadata is read on device, so this API is safe + * to use inside CUDA graph capture. The output grouped tensor supplies the + * per-group shape metadata and receives one amax value per group. + * + * \param[in] input Input grouped tensor. Must be unquantized. + * \param[in,out] output Grouped FP8 tensor with per-tensor scaling. + * \param[in] config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_compute_amax_with_config(const NVTEGroupedTensor input, NVTEGroupedTensor output, + const NVTEQuantizationConfig config, cudaStream_t stream); + /*! \brief Update an FP8 tensor's scale based on its amax. * * This is only supported for FP8 tensors with per-tensor scaling. @@ -111,6 +125,20 @@ void nvte_compute_amax_with_config(const NVTETensor input, NVTETensor output, void nvte_compute_scale_from_amax(NVTETensor output, const NVTEQuantizationConfig config, cudaStream_t stream); +/*! \brief Update a grouped FP8 tensor's per-group scale/scale_inv from its amax. + * + * Grouped counterpart of nvte_compute_scale_from_amax: for each group derives + * scale = max_fp8 / amax and scale_inv = 1/scale. All metadata is read on + * device, so this is safe inside CUDA graph capture.Currently used in grouped FP8 + * current-scaling. + * + * \param[in,out] output Grouped FP8 tensor with per-tensor scaling. + * \param[in] config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_compute_scale_from_amax(NVTEGroupedTensor output, + const NVTEQuantizationConfig config, cudaStream_t stream); + /*! \brief Compute partial amax for FP8 blockwise scaling. * * This function computes the maximum absolute values for each block of the original tensor. diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index c32a561fb7..aa0405e177 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -509,6 +509,24 @@ void nvte_memset(void *ptr, int value, size_t size_in_bytes, cudaStream_t stream void nvte_splits_to_offsets(const int64_t *split_sizes, int64_t *output, size_t num_splits, int64_t stride, cudaStream_t stream); +/*! \brief Compute prefix-sum element offsets for grouped tensors with both dimensions varying. + * + * Given per-tensor "first" and "last" dimensions, compute the cumulative + * element offsets: + * + * output[0] = 0 + * output[i] = sum_{j < i} first_dims[j] * last_dims[j] for i in [1, N] + * + * where N is the number of entries in first_dims (== entries in last_dims). + * + * \param[in] first_dims Per-tensor first dim, int32/int64 1D tensor of shape [N]. + * \param[in] last_dims Per-tensor last dim, int32/int64 1D tensor of shape [N]. + * \param[out] output Int32/int64 1D output tensor of shape [N + 1]. + * \param[in] stream CUDA stream to use for the operation. + */ +void nvte_splits_to_offsets_2d(const NVTETensor first_dims, const NVTETensor last_dims, + NVTETensor output, cudaStream_t stream); + /*! \brief Compute multiple scaled prefix-sum offsets for grouped tensors. * * Computes a prefix-sum over the values in split_sizes, and for each diff --git a/transformer_engine/common/recipe/current_scaling.cu b/transformer_engine/common/recipe/current_scaling.cu index 15ec1621bc..3ba51dba75 100644 --- a/transformer_engine/common/recipe/current_scaling.cu +++ b/transformer_engine/common/recipe/current_scaling.cu @@ -10,6 +10,7 @@ #include #include +#include "../cast/fp8/group_amax_fp8.cuh" #include "../common.h" #include "../util/logging.h" #include "../util/vectorized_pointwise.h" @@ -175,6 +176,72 @@ void compute_amax_impl(const NVTETensor input_, const NVTETensor output_, cudaSt stream);); // NOLINT(*) } +void group_compute_amax_impl(const NVTEGroupedTensor input_, NVTEGroupedTensor output_, + const NVTEQuantizationConfig config_, cudaStream_t stream) { + using namespace transformer_engine; + + NVTE_CHECK(input_ != nullptr, "Invalid grouped input tensor (got NULL)"); + const auto &input = *convertNVTEGroupedTensorCheck(input_); + NVTE_CHECK(output_ != nullptr, "Invalid grouped output tensor (got NULL)"); + auto &output = *convertNVTEGroupedTensorCheck(output_); + NVTE_CHECK(input.num_tensors == output.num_tensors, + "Number of grouped input and output tensors must match."); + NVTE_CHECK(input.has_data(), "Grouped amax input must have rowwise data."); + NVTE_CHECK(!is_fp8_dtype(input.data.dtype), + "Grouped amax input must be unquantized, but got dtype=", to_string(input.data.dtype)); + NVTE_CHECK(output.amax.has_data() || output.columnwise_amax.has_data(), + "Grouped amax output must have an amax buffer."); + + CheckInputGroupedTensor(input, "group_compute_amax_input"); + CheckOutputGroupedTensor(output, "group_compute_amax_output", true); + + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (output.all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (output.all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (output.all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + + const int64_t *offsets_ptr = reinterpret_cast(output.tensor_offsets.dptr); + const int64_t *first_dims_ptr = reinterpret_cast(output.first_dims.dptr); + const int64_t *last_dims_ptr = reinterpret_cast(output.last_dims.dptr); + if (shape_rep != ShapeRepresentation::SAME_BOTH_DIMS) { + NVTE_CHECK(offsets_ptr != nullptr, + "Grouped amax requires tensor_offsets when a grouped dimension varies."); + } + if (shape_rep == ShapeRepresentation::VARYING_FIRST_DIM || + shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { + NVTE_CHECK(first_dims_ptr != nullptr, + "Grouped amax requires first_dims for varying first dimensions."); + } + if (shape_rep == ShapeRepresentation::VARYING_LAST_DIM || + shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { + NVTE_CHECK(last_dims_ptr != nullptr, + "Grouped amax requires last_dims for varying last dimensions."); + } + + float *noop_ptr = nullptr; + if (config_ != nullptr) { + const QuantizationConfig *config_cpp = reinterpret_cast(config_); + const NVTETensor noop = config_cpp ? config_cpp->noop_tensor : nullptr; + noop_ptr = reinterpret_cast( + (noop != nullptr ? convertNVTETensorCheck(noop)->data.dptr : nullptr)); + } + + float *amax_ptr = reinterpret_cast( + output.amax.dptr != nullptr ? output.amax.dptr : output.columnwise_amax.dptr); + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.data.dtype, IType, + dispatch::fp8::launch_grouped_amax_kernel( + reinterpret_cast(input.data.dptr), amax_ptr, output.num_tensors, + output.logical_shape.data[0], output.logical_shape.data[1], shape_rep, offsets_ptr, + first_dims_ptr, last_dims_ptr, noop_ptr, stream);); // NOLINT(*) +} + } // anonymous namespace void nvte_compute_amax(const NVTETensor input_, const NVTETensor output_, cudaStream_t stream) { @@ -188,6 +255,12 @@ void nvte_compute_amax_with_config(const NVTETensor input_, const NVTETensor out compute_amax_impl(input_, output_, stream, config_); } +void nvte_group_compute_amax_with_config(const NVTEGroupedTensor input, NVTEGroupedTensor output, + const NVTEQuantizationConfig config, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_compute_amax_with_config); + group_compute_amax_impl(input, output, config, stream); +} + namespace transformer_engine { namespace { diff --git a/transformer_engine/common/util/splits_to_offsets.cu b/transformer_engine/common/util/splits_to_offsets.cu index 721f88c074..2a06603daf 100644 --- a/transformer_engine/common/util/splits_to_offsets.cu +++ b/transformer_engine/common/util/splits_to_offsets.cu @@ -119,6 +119,50 @@ __global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) { } } +__global__ void __launch_bounds__(kThreadsPerBlock) + splits_to_offsets_2d_kernel(const void *__restrict__ first_dims, DType first_dims_dtype, + const void *__restrict__ last_dims, DType last_dims_dtype, + void *__restrict__ output, DType output_dtype, size_t num_tensors) { + __shared__ int64_t block_scan[kThreadsPerBlock]; + __shared__ int64_t chunk_prefix; + + const size_t tid = threadIdx.x; + if (tid == 0) { + store_output(output, output_dtype, 0, 0); + chunk_prefix = 0; + } + __syncthreads(); + + for (size_t chunk_start = 0; chunk_start < num_tensors; chunk_start += kThreadsPerBlock) { + const size_t idx = chunk_start + tid; + int64_t value = 0; + if (idx < num_tensors) { + value = load_split_size(first_dims, first_dims_dtype, idx) * + load_split_size(last_dims, last_dims_dtype, idx); + } + block_scan[tid] = value; + __syncthreads(); + + // Inclusive scan in shared memory. + for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) { + const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0; + __syncthreads(); + block_scan[tid] += addend; + __syncthreads(); + } + + if (idx < num_tensors) { + store_output(output, output_dtype, idx + 1, chunk_prefix + block_scan[tid]); + } + __syncthreads(); + + if (tid == kThreadsPerBlock - 1) { + chunk_prefix += block_scan[tid]; + } + __syncthreads(); + } +} + } // namespace } // namespace transformer_engine::splits_to_offsets @@ -147,6 +191,42 @@ void nvte_splits_to_offsets(const int64_t *split_sizes, int64_t *output, size_t NVTE_CHECK_CUDA(cudaGetLastError()); } +void nvte_splits_to_offsets_2d(const NVTETensor first_dims, const NVTETensor last_dims, + NVTETensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_splits_to_offsets_2d); + using namespace transformer_engine; + namespace s2o = transformer_engine::splits_to_offsets; + + const auto is_integer_dtype = [](DType dtype) { + return dtype == DType::kInt32 || dtype == DType::kInt64; + }; + + const auto *first_dims_tensor = convertNVTETensorCheck(first_dims); + const auto *last_dims_tensor = convertNVTETensorCheck(last_dims); + auto *output_tensor = convertNVTETensorCheck(output); + + const auto num_tensors = first_dims_tensor->numel(); + NVTE_CHECK(num_tensors > 0 && first_dims_tensor->dim() == 1, + "first_dims must be a non-empty 1D tensor, but got shape=", first_dims_tensor->shape(), + "."); + NVTE_CHECK(last_dims_tensor->shape() == first_dims_tensor->shape(), + "first_dims and last_dims must have the same shape, but got ", + first_dims_tensor->shape(), " and ", last_dims_tensor->shape(), "."); + const Shape expected_output_shape = {num_tensors + 1}; + NVTE_CHECK(output_tensor->shape() == expected_output_shape, + "Expected output to have shape=", expected_output_shape, + ", but got shape=", output_tensor->shape(), "."); + NVTE_CHECK(is_integer_dtype(first_dims_tensor->dtype()) && + is_integer_dtype(last_dims_tensor->dtype()) && + is_integer_dtype(output_tensor->dtype()), + "first_dims, last_dims and output must be int32/int64 tensors."); + + s2o::splits_to_offsets_2d_kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>( + first_dims_tensor->data.dptr, first_dims_tensor->dtype(), last_dims_tensor->data.dptr, + last_dims_tensor->dtype(), output_tensor->data.dptr, output_tensor->dtype(), num_tensors); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, const int64_t *strides, const int *include_leading_zero, size_t num_outputs, cudaStream_t stream) { diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index b322ce8fba..635c7a36d2 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -938,6 +938,36 @@ __device__ __forceinline__ float ordered_uint_to_float(unsigned int u) { return __uint_as_float(u ^ mask); } +template +__device__ __forceinline__ T abs_val(T val) { + if constexpr (std::is_same_v) { +#if __CUDA_ARCH__ >= 800 + return __habs(val); +#else + return static_cast<__nv_bfloat16>(fabsf(static_cast(val))); +#endif + } else if constexpr (std::is_same_v) { + return __habs(val); + } else { + return fabsf(val); + } +} + +template +__device__ __forceinline__ T max_val(T a, T b) { + if constexpr (std::is_same_v) { +#if __CUDA_ARCH__ >= 800 + return __hmax(a, b); +#else + return static_cast<__nv_bfloat16>(fmaxf(static_cast(a), static_cast(b))); +#endif + } else if constexpr (std::is_same_v) { + return __hmax(a, b); + } else { + return fmaxf(a, b); + } +} + template __device__ __forceinline__ T warp_allreduce_sum(T x) { // Butterfly reduction diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 779b145dd9..e4aba728b6 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -119,7 +119,8 @@ class Quantizer { virtual std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, py::object quantizer, const std::optional& first_dims, - const std::optional& tensor_offsets, size_t logical_first_dim, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const = 0; /*! @brief Convert a PyTorch tensor into a Transformer Engine C++ tensor @@ -162,7 +163,8 @@ class NoneQuantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, py::object quantizer, const std::optional& first_dims, - const std::optional& tensor_offsets, size_t logical_first_dim, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct a tensor with pre-initialized data */ @@ -194,7 +196,8 @@ class Float8Quantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, py::object quantizer, const std::optional& first_dims, - const std::optional& tensor_offsets, size_t logical_first_dim, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct a tensor with pre-initialized data */ @@ -230,7 +233,8 @@ class Float8CurrentScalingQuantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, py::object quantizer, const std::optional& first_dims, - const std::optional& tensor_offsets, size_t logical_first_dim, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct an unquantized tensor with a freshly allocated amax buffer. @@ -294,7 +298,8 @@ class Float8BlockQuantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, py::object quantizer, const std::optional& first_dims, - const std::optional& tensor_offsets, size_t logical_first_dim, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; std::pair convert_and_update_tensor(py::object shape) const override; @@ -320,7 +325,8 @@ class MXFP8Quantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, py::object quantizer, const std::optional& first_dims, - const std::optional& tensor_offsets, size_t logical_first_dim, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; std::pair convert_and_update_tensor(py::object shape) const override; @@ -365,7 +371,8 @@ class NVFP4Quantizer : public Quantizer { std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, py::object quantizer, const std::optional& first_dims, - const std::optional& tensor_offsets, size_t logical_first_dim, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, size_t logical_first_dim, size_t logical_last_dim) const override; /*! @brief Construct an unquantized tensor that shares NVFP4 tensor's amax pointer diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 13d872392d..e66698dc2c 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -343,12 +343,14 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz 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 tensor_offsets); + std::optional first_dims, std::optional last_dims, + std::optional tensor_offsets, + std::optional noop_flag); py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, + std::optional last_dims, const at::Tensor &rowwise_amax, const at::Tensor &columnwise_amax, std::optional tensor_offsets); @@ -357,6 +359,7 @@ py::object group_dequantize(const py::handle &input, DType otype); py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, + std::optional last_dims, std::optional tensor_offsets); std::vector multi_tensor_quantize(const std::vector &tensor_list, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index aab5a87b9a..f57bd8797e 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -7,6 +7,7 @@ #include "transformer_engine/cast.h" #include +#include #include #include #include @@ -17,8 +18,11 @@ #include "../extensions.h" #include "common.h" +#include "common/common.h" #include "common/util/system.h" #include "pybind.h" +#include "transformer_engine/multi_tensor.h" +#include "transformer_engine/recipe.h" #include "transformer_engine/transformer_engine.h" namespace transformer_engine { @@ -204,17 +208,86 @@ void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, }); } +float fp8_max_for_dtype(const DType dtype) { + if (!is_fp8_dtype(dtype)) { + NVTE_ERROR("Expected FP8 dtype for grouped current-scaling quantization, got ", + to_string(dtype), "."); + } + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + dtype, T, { return transformer_engine::detail::TypeExtrema::max; }); + return 0.0f; +} + +// Computes per-group amax and scale/scale_inv for grouped FP8 current scaling. +// +// 1) nvte_group_compute_amax_with_config: per-group amax over the input tensor +// (current scaling -> amax is computed from this batch's inputs, never from +// history). +// 2) Optional NCCL allreduce on the amax buffer (when DP/TP amax sync is on). +// 3) nvte_group_compute_scale_from_amax: graph-safe per-group math kernel that +// derives scale = max_fp8 / amax and scale_inv = 1/scale per group in a +// single launch. It honors the config's no-op flag, so a skipped weight +// update (skip_fp8_weight_update) preserves the cached scale instead of +// recomputing it from a stale/uninitialized amax. +// +// After this returns, the per-group `scale` / `scale_inv` buffers are +// populated and the actual cast/transpose is performed by `nvte_group_quantize`. +void compute_grouped_fp8_current_scaling_amax_and_scale( + const GroupedTensorWrapper &grouped_input_tensor, + const GroupedTensorWrapper &grouped_output_tensor, const py::object &grouped_output_py, + Float8CurrentScalingQuantizer *quantizer_cpp, const std::optional &noop_flag) { + QuantizationConfigWrapper quant_config; + quant_config.set_force_pow_2_scales(quantizer_cpp->force_pow_2_scales); + quant_config.set_amax_epsilon(quantizer_cpp->amax_epsilon); + if (noop_flag.has_value()) { + quant_config.set_noop_tensor(noop_flag->data()); + } + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_compute_amax_with_config(grouped_input_tensor.data(), grouped_output_tensor.data(), + quant_config, stream); + }); + + if (quantizer_cpp->with_amax_reduction) { + // NCCL collectives require an at::Tensor; the amax buffer lives on the Python + // object (same allocation as the grouped tensor's amax), so fetch it only on + // this path. + auto amax = grouped_output_py.attr("amax").cast(); + c10d::AllreduceOptions opts; + opts.reduceOp = c10d::ReduceOp::MAX; + std::vector tensors = {amax}; + NVTE_SCOPED_GIL_RELEASE( + { quantizer_cpp->amax_reduction_group->allreduce(tensors, opts)->wait(); }); + } + + // Derive per-group scale/scale_inv from the (possibly reduced) amax. The same + // quant_config carries force_pow_2_scales/amax_epsilon and, crucially, the + // no-op flag, so the update is skipped on device when skip_fp8_weight_update is + // set -- preserving the cached scale instead of recomputing it from a + // stale/uninitialized amax. Buffers and the FP8 max are read from the grouped + // tensor inside the kernel launch. + NVTE_SCOPED_GIL_RELEASE( + { nvte_group_compute_scale_from_amax(grouped_output_tensor.data(), quant_config, stream); }); +} + } // namespace -// NOTE: Only supports varying first dim. py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, - std::optional first_dims, - std::optional tensor_offsets) { + std::optional first_dims, std::optional last_dims, + std::optional tensor_offsets, + std::optional noop_flag) { using namespace transformer_engine::pytorch::detail; init_extension(); NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + // No-op flag for CUDA graph weight caching (skip_fp8_weight_update). + std::optional noop_flag_cpp; + if (noop_flag.has_value()) { + noop_flag_cpp = makeTransformerEngineTensor(*noop_flag); + } + std::vector logical_shape; for (const auto &d : tensor.sizes()) { logical_shape.push_back(d); @@ -228,17 +301,19 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const // Create input GroupedTensor. auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); - grouped_input_tensor.set_rowwise_data( - tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor)); + grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), + 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, tensor_offsets, logical_first_dim, - logical_last_dim); + py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, + logical_first_dim, logical_last_dim); // dispatch to scaling methods enum class GroupedQuantizationMode { + FP8_CURRENT_SCALING_GROUPED_QUANTIZE, MXFP8_GROUPED_QUANTIZE, NVFP4_GROUPED_QUANTIZE, INVALID_FOR_GROUPED_QUANTIZE @@ -249,6 +324,8 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const grouped_quantization_mode = GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE; } else if (detail::IsNVFP4Quantizers(quantizer.ptr())) { grouped_quantization_mode = GroupedQuantizationMode::NVFP4_GROUPED_QUANTIZE; + } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr())) { + grouped_quantization_mode = GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE; } if (empty_input_buffer) { @@ -261,13 +338,33 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const switch (grouped_quantization_mode) { case GroupedQuantizationMode::NVFP4_GROUPED_QUANTIZE: { // NVFP4 grouped quantization + NVTE_CHECK(!last_dims.has_value(), + "group_quantize: varying last dim is not supported with NVFP4."); NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), true); break; } + case GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE: { + auto *fp8_quantizer_cpp = static_cast(quantizer_cpp.get()); + compute_grouped_fp8_current_scaling_amax_and_scale( + grouped_input_tensor, grouped_output_tensor_cpp, grouped_output_py, fp8_quantizer_cpp, + noop_flag_cpp); + QuantizationConfigWrapper quant_config_cpp; + 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()); + }); + break; + } case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: { QuantizationConfigWrapper quant_config_cpp; + 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()); @@ -276,7 +373,9 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const } case GroupedQuantizationMode::INVALID_FOR_GROUPED_QUANTIZE: default: - NVTE_ERROR("group_quantize: only support NVFP4 or MXFP8 quantizer."); + NVTE_ERROR( + "group_quantize: only supports MXFP8, NVFP4, or " + "Float8CurrentScalingQuantizer."); break; } @@ -286,6 +385,7 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, + std::optional last_dims, const at::Tensor &rowwise_amax, const at::Tensor &columnwise_amax, std::optional tensor_offsets) { @@ -293,6 +393,8 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q init_extension(); NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(!last_dims.has_value(), + "nvfp4_group_quantize_with_amax: varying last dim is not supported with NVFP4."); NVTE_CHECK(rowwise_amax.is_cuda() && columnwise_amax.is_cuda(), "Precomputed amax tensors must be CUDA tensors."); NVTE_CHECK( @@ -323,8 +425,8 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), - py::reinterpret_borrow(quantizer), first_dims, tensor_offsets, logical_first_dim, - logical_last_dim); + py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, + logical_first_dim, logical_last_dim); if (grouped_output_tensor_cpp.get_amax().data_ptr != nullptr) { grouped_output_tensor_cpp.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, @@ -358,11 +460,15 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q py::object bgrad_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) { using namespace transformer_engine::pytorch::detail; init_extension(); NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(!last_dims.has_value(), + "bgrad_group_quantize: varying last dim is not supported because the underlying " + "MXFP8 dbias kernel requires a constant last dimension across grouped tensors."); std::vector logical_shape; for (const auto &d : tensor.sizes()) { @@ -379,13 +485,14 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto quantizer_cpp = convert_quantizer(quantizer); auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); - grouped_input_tensor.set_rowwise_data( - tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor)); + grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), + GetTransformerEngineDType(tensor.scalar_type()), + std::vector{static_cast(tensor.numel())}); auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), - py::reinterpret_borrow(quantizer), first_dims, tensor_offsets, logical_first_dim, - logical_last_dim); + py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, + logical_first_dim, logical_last_dim); if (empty_input_buffer) { at::Tensor dbias_torch = @@ -473,7 +580,7 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o NoneQuantizer q{py::none()}; auto [out_cpp, out_py] = q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims, - tensor_offsets, logical_first_dim, logical_last_dim); + last_dims, tensor_offsets, logical_first_dim, logical_last_dim); return py::reinterpret_borrow(out_py); } @@ -512,7 +619,7 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o // Create output GroupedTensor using NoneQuantizer. NoneQuantizer q{py::none()}; auto [out_cpp, out_py] = - q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims, + q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims, last_dims, tensor_offsets, logical_first_dim, logical_last_dim); NVTE_SCOPED_GIL_RELEASE({ diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index d1890872c0..ef4f5d17d8 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -169,7 +169,7 @@ void bind_quantize_with_amax_extensions(py::module_ &m) { py::arg("quantizer"), py::arg("rowwise_amax"), py::arg("columnwise_amax")); m.def("nvfp4_group_quantize_with_amax", nvfp4_group_quantize_with_amax, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), - py::arg("rowwise_amax"), py::arg("columnwise_amax"), + py::arg("last_dims") = py::none(), py::arg("rowwise_amax"), py::arg("columnwise_amax"), py::arg("tensor_offsets") = py::none()); } @@ -203,13 +203,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("dtype"), py::arg("device"), py::arg("pin_memory")); m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), - py::arg("tensor_offsets") = py::none()); + py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(), + py::arg("noop_flag") = py::none()); transformer_engine::pytorch::bind_quantize_with_amax_extensions(m); m.def("group_dequantize", transformer_engine::pytorch::group_dequantize, "Dequantize group tensor", py::arg("input"), py::arg("otype")); m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), - py::arg("tensor_offsets") = py::none()); + py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none()); m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize, "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer")); m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)", diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 5fc50953a1..a39d0143a0 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -4,6 +4,7 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include "common.h" @@ -84,31 +85,81 @@ std::vector convert_shape_for_fp4(const std::vector& shape) { return ret; } +/*! @brief Validate an optional 1D int64 CUDA grouped-tensor metadata tensor + * (first_dims / last_dims / tensor_offsets) of a given expected length. */ +void check_grouped_metadata_tensor(const at::Tensor& metadata_tensor, const char* metadata_name, + const size_t expected_len) { + NVTE_CHECK(metadata_tensor.is_cuda(), metadata_name, " must be on CUDA."); + NVTE_CHECK(metadata_tensor.scalar_type() == at::kLong, metadata_name, " must have dtype int64."); + NVTE_CHECK(metadata_tensor.is_contiguous(), metadata_name, " must be contiguous."); + NVTE_CHECK(static_cast(metadata_tensor.numel()) == expected_len, metadata_name, + " must have length ", expected_len, "."); +} + std::optional build_grouped_tensor_offsets(const size_t num_tensors, const std::optional& first_dims, + const std::optional& last_dims, + const size_t logical_first_dim, const size_t logical_last_dim) { - if (!first_dims.has_value()) { + if (!first_dims.has_value() && !last_dims.has_value()) { return std::nullopt; } - const auto& first_dims_tensor = first_dims.value(); - NVTE_CHECK(first_dims_tensor.is_cuda(), "first_dims must be on CUDA."); - NVTE_CHECK(first_dims_tensor.scalar_type() == at::kLong, "first_dims must have dtype int64."); - NVTE_CHECK(static_cast(first_dims_tensor.numel()) == num_tensors, - "first_dims must have length ", num_tensors, "."); + // Validate dims before the splits-to-offsets kernel reads them. + if (first_dims.has_value()) { + check_grouped_metadata_tensor(*first_dims, "first_dims", num_tensors); + } + if (last_dims.has_value()) { + check_grouped_metadata_tensor(*last_dims, "last_dims", num_tensors); + } - const int64_t logical_last_dim_i64 = static_cast(logical_last_dim); - const auto first_dims_contiguous = first_dims_tensor.contiguous(); - auto tensor_offsets = - at::empty({static_cast(num_tensors) + 1}, first_dims_contiguous.options()); - NVTE_SCOPED_GIL_RELEASE({ - nvte_splits_to_offsets(static_cast(first_dims_contiguous.data_ptr()), - static_cast(tensor_offsets.data_ptr()), num_tensors, - logical_last_dim_i64, at::cuda::getCurrentCUDAStream()); - }); + const at::TensorOptions options = + first_dims.has_value() ? first_dims->options() : last_dims->options(); + auto tensor_offsets = at::empty({static_cast(num_tensors) + 1}, options); + if (first_dims.has_value() && last_dims.has_value()) { + auto first_dims_nvte = makeTransformerEngineTensor(*first_dims); + auto last_dims_nvte = makeTransformerEngineTensor(*last_dims); + auto tensor_offsets_nvte = makeTransformerEngineTensor(tensor_offsets); + NVTE_SCOPED_GIL_RELEASE({ + nvte_splits_to_offsets_2d(first_dims_nvte.data(), last_dims_nvte.data(), + tensor_offsets_nvte.data(), at::cuda::getCurrentCUDAStream()); + }); + } else if (first_dims.has_value()) { + NVTE_SCOPED_GIL_RELEASE({ + nvte_splits_to_offsets(static_cast(first_dims->data_ptr()), + static_cast(tensor_offsets.data_ptr()), num_tensors, + static_cast(logical_last_dim), + at::cuda::getCurrentCUDAStream()); + }); + } else { + NVTE_SCOPED_GIL_RELEASE({ + nvte_splits_to_offsets(static_cast(last_dims->data_ptr()), + static_cast(tensor_offsets.data_ptr()), num_tensors, + static_cast(logical_first_dim), + at::cuda::getCurrentCUDAStream()); + }); + } return tensor_offsets; } +/*! @brief Validate grouped-tensor offset metadata and resolve the final offsets, + * whether they are precomputed or built from first_dims/last_dims. */ +std::optional resolve_grouped_tensor_offsets( + const size_t num_tensors, const std::optional& first_dims, + const std::optional& last_dims, + const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, + const size_t logical_last_dim) { + // Precomputed offsets take priority; otherwise build them from first_dims/last_dims + // (which validates the dims internally). + if (precomputed_tensor_offsets.has_value()) { + // tensor_offsets uses a CSR-style prefix-sum layout with num_tensors+1 entries. + check_grouped_metadata_tensor(*precomputed_tensor_offsets, "tensor_offsets", num_tensors + 1); + return precomputed_tensor_offsets; + } + return build_grouped_tensor_offsets(num_tensors, first_dims, last_dims, logical_first_dim, + logical_last_dim); +} + at::TensorOptions grouped_tensor_data_options(const DType dtype) { return at::TensorOptions().dtype(GetATenDType(dtype)).device(torch::kCUDA); } @@ -174,14 +225,14 @@ std::pair NoneQuantizer::create_tensor(const std::vec std::pair NoneQuantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, + const std::optional& last_dims, const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - precomputed_tensor_offsets.has_value() - ? precomputed_tensor_offsets - : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, + logical_first_dim, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -207,6 +258,9 @@ std::pair NoneQuantizer::create_grouped_tensor if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); } + if (last_dims.has_value()) { + out_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims)); + } if (tensor_offsets.has_value()) { out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, getTensorShape(*tensor_offsets)); @@ -231,7 +285,7 @@ std::pair NoneQuantizer::create_grouped_tensor kwargs["columnwise_amax"] = py::none(); kwargs["scale"] = py::none(); kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); - kwargs["last_dims"] = py::none(); + kwargs["last_dims"] = last_dims.has_value() ? py::cast(*last_dims) : py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); @@ -380,14 +434,14 @@ std::pair Float8Quantizer::create_tensor( std::pair Float8Quantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, + const std::optional& last_dims, const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - precomputed_tensor_offsets.has_value() - ? precomputed_tensor_offsets - : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, + logical_first_dim, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -425,6 +479,9 @@ std::pair Float8Quantizer::create_grouped_tens if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); } + if (last_dims.has_value()) { + out_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims)); + } if (tensor_offsets.has_value()) { out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, getTensorShape(*tensor_offsets)); @@ -449,7 +506,7 @@ std::pair Float8Quantizer::create_grouped_tens kwargs["columnwise_amax"] = py::none(); kwargs["scale"] = py::none(); kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); - kwargs["last_dims"] = py::none(); + kwargs["last_dims"] = last_dims.has_value() ? py::cast(*last_dims) : py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); @@ -693,14 +750,19 @@ std::pair Float8CurrentScalingQuantizer::create_tenso std::pair Float8CurrentScalingQuantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, + const std::optional& last_dims, const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, const size_t logical_last_dim) const { using namespace pybind11::literals; + // Group Quantize is not implemented for varying both dims yet. + NVTE_CHECK(!(first_dims.has_value() && last_dims.has_value()), + "FP8 current-scaling grouped quantization does not support varying both " + "first and last dimensions."); + const auto tensor_offsets = - precomputed_tensor_offsets.has_value() - ? precomputed_tensor_offsets - : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, + logical_first_dim, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -713,23 +775,31 @@ std::pair Float8CurrentScalingQuantizer::creat std::optional columnwise_scale_inv; at::Tensor scale = at::empty({static_cast(num_tensors)}, float_opts); at::Tensor amax = at::empty({static_cast(num_tensors)}, float_opts); + const bool is_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported(); + const bool with_rowwise_data = rowwise_usage || is_non_tn_fp8_gemm_supported; + const bool with_columnwise_data = columnwise_usage && !is_non_tn_fp8_gemm_supported; - if (rowwise_usage) { + // FP8 current scaling has a single per-tensor scale + std::optional scale_inv; + if (with_rowwise_data || with_columnwise_data) { + scale_inv = at::empty({static_cast(num_tensors)}, float_opts); + } + if (with_rowwise_data) { rowwise_data = at::empty({total_elements}, uint8_opts); - rowwise_scale_inv = at::empty({static_cast(num_tensors)}, float_opts); + rowwise_scale_inv = scale_inv; } - if (columnwise_usage) { + if (with_columnwise_data) { columnwise_data = at::empty({total_elements}, uint8_opts); - columnwise_scale_inv = at::empty({static_cast(num_tensors)}, float_opts); + columnwise_scale_inv = scale_inv; } GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); - if (rowwise_usage) { + if (with_rowwise_data) { out_cpp.set_rowwise_data(rowwise_data->data_ptr(), this->dtype, getTensorShape(*rowwise_data)); out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat32, getTensorShape(*rowwise_scale_inv)); } - if (columnwise_usage) { + if (with_columnwise_data) { out_cpp.set_columnwise_data(columnwise_data->data_ptr(), this->dtype, getTensorShape(*columnwise_data)); out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat32, @@ -740,6 +810,9 @@ std::pair Float8CurrentScalingQuantizer::creat if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); } + if (last_dims.has_value()) { + out_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims)); + } if (tensor_offsets.has_value()) { out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, getTensorShape(*tensor_offsets)); @@ -764,8 +837,9 @@ std::pair Float8CurrentScalingQuantizer::creat kwargs["columnwise_amax"] = py::none(); kwargs["scale"] = scale; kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); - kwargs["last_dims"] = py::none(); + kwargs["last_dims"] = last_dims.has_value() ? py::cast(*last_dims) : py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + kwargs["offsets"] = py::none(); kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { @@ -1075,14 +1149,14 @@ std::pair Float8BlockQuantizer::create_tensor( std::pair Float8BlockQuantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, + const std::optional& last_dims, const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - precomputed_tensor_offsets.has_value() - ? precomputed_tensor_offsets - : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, + logical_first_dim, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -1124,6 +1198,9 @@ std::pair Float8BlockQuantizer::create_grouped if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); } + if (last_dims.has_value()) { + out_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims)); + } if (tensor_offsets.has_value()) { out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, getTensorShape(*tensor_offsets)); @@ -1148,7 +1225,7 @@ std::pair Float8BlockQuantizer::create_grouped kwargs["columnwise_amax"] = py::none(); kwargs["scale"] = py::none(); kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); - kwargs["last_dims"] = py::none(); + kwargs["last_dims"] = last_dims.has_value() ? py::cast(*last_dims) : py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); @@ -1496,14 +1573,14 @@ std::pair MXFP8Quantizer::create_tensor( std::pair MXFP8Quantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, + const std::optional& last_dims, const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - precomputed_tensor_offsets.has_value() - ? precomputed_tensor_offsets - : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, + logical_first_dim, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); @@ -1515,17 +1592,34 @@ std::pair MXFP8Quantizer::create_grouped_tenso std::optional columnwise_scale_inv; const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; + // For VARYING_BOTH_DIMS each tensor in the group must have its first and last + // dim be a multiple of 128 (grouped-kernel tile alignment). MXFP8 stores one + // E8M0 scale per MXFP8_BLOCK_SIZE (32) elements, so the 128-alignment + // guarantees the per-group scale counts tile evenly and the total scale count + // is simply scale_inv_numel = data_numel / MXFP8_BLOCK_SIZE (i.e. / 32). + const bool is_varying_both = first_dims.has_value() && last_dims.has_value(); + if (rowwise_usage) { rowwise_data = at::empty({total_elements}, uint8_opts); - const auto scale_shape = get_scale_shape(logical_shape_vec, false); - const int64_t total_scale_elements = static_cast(product(scale_shape)); + int64_t total_scale_elements; + if (is_varying_both) { + total_scale_elements = total_elements / static_cast(MXFP8_BLOCK_SIZE); + } else { + const auto scale_shape = get_scale_shape(logical_shape_vec, false); + total_scale_elements = static_cast(product(scale_shape)); + } rowwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); } if (columnwise_usage) { columnwise_data = at::empty({total_elements}, uint8_opts); - const auto scale_shape = get_scale_shape(logical_shape_vec, true); - const int64_t total_scale_elements = static_cast(product(scale_shape)); + int64_t total_scale_elements; + if (is_varying_both) { + total_scale_elements = total_elements / static_cast(MXFP8_BLOCK_SIZE); + } else { + const auto scale_shape = get_scale_shape(logical_shape_vec, true); + total_scale_elements = static_cast(product(scale_shape)); + } columnwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); } @@ -1544,6 +1638,9 @@ std::pair MXFP8Quantizer::create_grouped_tenso if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); } + if (last_dims.has_value()) { + out_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims)); + } if (tensor_offsets.has_value()) { out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, getTensorShape(*tensor_offsets)); @@ -1570,7 +1667,7 @@ std::pair MXFP8Quantizer::create_grouped_tenso kwargs["columnwise_amax"] = py::none(); kwargs["scale"] = py::none(); kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); - kwargs["last_dims"] = py::none(); + kwargs["last_dims"] = last_dims.has_value() ? py::cast(*last_dims) : py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = this->optimize_for_gemm; PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); @@ -1962,14 +2059,14 @@ std::pair NVFP4Quantizer::create_tensor( std::pair NVFP4Quantizer::create_grouped_tensor( const size_t num_tensors, const std::vector& logical_shape, const DType dtype, py::object quantizer, const std::optional& first_dims, + const std::optional& last_dims, const std::optional& precomputed_tensor_offsets, const size_t logical_first_dim, const size_t logical_last_dim) const { using namespace pybind11::literals; const auto tensor_offsets = - precomputed_tensor_offsets.has_value() - ? precomputed_tensor_offsets - : build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, + logical_first_dim, logical_last_dim); const int64_t total_elements = static_cast(logical_first_dim) * static_cast(logical_last_dim); NVTE_CHECK(total_elements % 2 == 0, "NVFP4 data size must be divisible by 2."); @@ -2031,6 +2128,9 @@ std::pair NVFP4Quantizer::create_grouped_tenso if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); } + if (last_dims.has_value()) { + out_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims)); + } if (tensor_offsets.has_value()) { out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, getTensorShape(*tensor_offsets)); @@ -2057,7 +2157,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso kwargs["columnwise_amax"] = maybe_tensor_to_py(columnwise_amax); kwargs["scale"] = py::none(); kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); - kwargs["last_dims"] = py::none(); + kwargs["last_dims"] = last_dims.has_value() ? py::cast(*last_dims) : py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = this->optimize_for_gemm; kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 39180f098e..da3aa13156 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -202,8 +202,8 @@ def _group_quantize_with_amax_for_grouped_mlp( quantizer, num_groups, split_sizes, - rowwise_amax, - columnwise_amax, + rowwise_amax=rowwise_amax, + columnwise_amax=columnwise_amax, tensor_offsets=tensor_offsets, ) diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index c112634024..7ede75943a 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -18,6 +18,7 @@ from .mxfp8_tensor_storage import MXFP8TensorStorage from .float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from .nvfp4_tensor_storage import NVFP4TensorStorage +from ...utils import is_non_tn_fp8_gemm_supported class GroupedTensorStorage: @@ -464,7 +465,6 @@ def make_grouped_tensor_with_shapes( # First dim first_dim_list = [s[0] for s in shapes] uniform_first_dim = all(first_dim_list[0] == x for x in first_dim_list) - logical_first_dim = sum(first_dim_list) if uniform_first_dim: first_dims = None else: @@ -472,13 +472,27 @@ def make_grouped_tensor_with_shapes( # Last dim last_dim_list = [s[1] for s in shapes] - logical_last_dim = last_dim_list[0] - assert all(logical_last_dim == x for x in last_dim_list), "Last dims should be uniform" + uniform_last_dim = all(last_dim_list[0] == x for x in last_dim_list) + if uniform_last_dim: + last_dims = None + else: + last_dims = torch.tensor([s[1] for s in shapes], dtype=torch.int64, device=device) + if first_dims is not None and last_dims is not None: + total_elements = sum(f * l for f, l in zip(first_dim_list, last_dim_list)) + logical_first_dim = 1 + logical_last_dim = total_elements + else: + logical_first_dim = ( + first_dim_list[0] + if uniform_first_dim and not uniform_last_dim + else sum(first_dim_list) + ) + logical_last_dim = sum(last_dim_list) if not uniform_last_dim else last_dim_list[0] return GroupedTensorStorage.make_grouped_tensor( num_tensors=num_tensors, first_dims=first_dims, - last_dims=None, + last_dims=last_dims, logical_first_dim=logical_first_dim, logical_last_dim=logical_last_dim, quantizer=quantizer, @@ -602,12 +616,12 @@ def copy(self) -> "GroupedTensorStorage": ) @staticmethod - def make_tensor_offsets(first_dims: torch.Tensor, logical_last_dim: int) -> torch.Tensor: - """Calculate GPU offsets from first dim splits.""" + def make_tensor_offsets(split_dims: torch.Tensor, common_dim: int) -> torch.Tensor: + """Calculate GPU element offsets from one varying dimension and one common dimension.""" return torch.cat( [ - torch.zeros(1, device=first_dims.device, dtype=first_dims.dtype), - torch.cumsum(first_dims * logical_last_dim, dim=0), + torch.zeros(1, device=split_dims.device, dtype=split_dims.dtype), + torch.cumsum(split_dims * common_dim, dim=0), ] ) @@ -648,7 +662,6 @@ def make_grouped_tensor( all_same_first = first_dims is None all_same_last = last_dims is None - assert all_same_last, "Last dim must be uniform for GroupedTensor" assert logical_first_dim >= 0, "Logical first dim must be non-negative for GroupedTensor" assert logical_last_dim > 0, "Logical last dim must be positive for GroupedTensor" @@ -661,26 +674,16 @@ def make_grouped_tensor( tensor_offsets = None offsets = None shape = [] - if not all_same_first: + if not all_same_first and not all_same_last: + tensor_offsets = GroupedTensorStorage.make_tensor_offsets(first_dims * last_dims, 1) + elif not all_same_first: # Need explicit offsets for non-uniform shapes # Offsets are based on number of elements and not pointers. # Kernels need to calculate precise pointers based on size of elements. - # TODO(ksivaman): Single kernel + remove the host offset calculation. tensor_offsets = GroupedTensorStorage.make_tensor_offsets(first_dims, logical_last_dim) - if ( - first_dims.device.type == "cuda" - and torch.cuda.is_available() - and torch.cuda.is_current_stream_capturing() - ): - # Avoid host sync during CUDA graph capture. - offsets = None - shape = None - else: - offsets = tensor_offsets.tolist() - first_dims_list = first_dims.tolist() - for i in range(num_tensors): - shape.append((first_dims_list[i], logical_last_dim)) + elif not all_same_last: + tensor_offsets = GroupedTensorStorage.make_tensor_offsets(last_dims, logical_first_dim) else: offsets = [ i * logical_first_dim * logical_last_dim // num_tensors @@ -693,9 +696,27 @@ def make_grouped_tensor( logical_shape = (logical_first_dim, logical_last_dim) no_quantization = quantizer is None - rowwise_usage = quantizer.rowwise_usage if not no_quantization else True columnwise_usage = quantizer.columnwise_usage if not no_quantization else False + compatible_recipe = None if no_quantization else quantizer._get_compatible_recipe() + + if not shape: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise ValueError( + "Varying-dimension grouped tensor construction is not graph-safe: it" + " requires device-to-host copies of per-tensor shapes/offsets." + ) + offsets = tensor_offsets.tolist() + if first_dims is not None and last_dims is not None: + first_dims_list = first_dims.tolist() + last_dims_list = last_dims.tolist() + shape = [(first_dims_list[i], last_dims_list[i]) for i in range(num_tensors)] + elif first_dims is not None: + first_dims_list = first_dims.tolist() + shape = [(rows, logical_last_dim) for rows in first_dims_list] + else: + last_dims_list = last_dims.tolist() + shape = [(logical_first_dim, cols) for cols in last_dims_list] # Calculate total elements across all tensors total_elements = logical_first_dim * logical_last_dim @@ -721,7 +742,7 @@ def make_grouped_tensor( if columnwise_usage: # Allocate columnwise data buffer (1D flattened, uint8) columnwise_data = torch.empty(total_elements, dtype=dtype, device=device) - elif quantizer._get_compatible_recipe().mxfp8(): + elif compatible_recipe.mxfp8(): if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) @@ -750,7 +771,7 @@ def make_grouped_tensor( columnwise_scale_inv = torch.empty( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) - elif quantizer._get_compatible_recipe().delayed(): + elif compatible_recipe.delayed(): if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) @@ -769,7 +790,7 @@ def make_grouped_tensor( # Amax buffer for delayed scaling - one per tensor amax = torch.empty(num_tensors, dtype=torch.float32, device=device) - elif quantizer._get_compatible_recipe().nvfp4(): + elif compatible_recipe.nvfp4(): row_scaled_nvfp4 = quantizer.row_scaled_nvfp4 nvfp4_use_4over6 = quantizer.nvfp4_use_4over6 nvfp4_e4m3_max = quantizer.nvfp4_e4m3_max @@ -816,7 +837,7 @@ def make_grouped_tensor( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device) - elif quantizer._get_compatible_recipe().float8_block_scaling(): + elif compatible_recipe.float8_block_scaling(): if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) @@ -843,21 +864,25 @@ def make_grouped_tensor( columnwise_scale_inv = torch.empty( total_columnwise_scale_elements, dtype=torch.float32, device=device ) - elif quantizer._get_compatible_recipe().float8_current_scaling(): - # Current scaling - per-tensor scaling computed on the fly - if rowwise_usage: + elif compatible_recipe.float8_current_scaling(): + # only rowwise would be be needed for Blackwell FP8 GEMM + non_tn_fp8_gemm_supported = is_non_tn_fp8_gemm_supported() + fp8_rowwise_usage = rowwise_usage or non_tn_fp8_gemm_supported + fp8_columnwise_usage = columnwise_usage and not non_tn_fp8_gemm_supported + shared_scale_inv = None + if fp8_rowwise_usage or fp8_columnwise_usage: + shared_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) + if fp8_rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) - # Scale inverse - one per tensor - scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) + scale_inv = shared_scale_inv # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors scale_inv_offsets = list(range(num_tensors + 1)) - if columnwise_usage: + if fp8_columnwise_usage: # Allocate columnwise data buffer (1D flattened, uint8) columnwise_data = torch.empty(total_elements, dtype=torch.uint8, device=device) - # Columnwise scale inverse - one per tensor - columnwise_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) + columnwise_scale_inv = shared_scale_inv # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors columnwise_scale_inv_offsets = list(range(num_tensors + 1)) @@ -931,11 +956,19 @@ def split_into_quantized_tensors( # if self.tensor_shapes is None, then trigger D2H copy and get the shape (not graph safe) if self.tensor_shapes is None: - first_dims_list = ( - [self.logical_shape[0]] * self.num_tensors - if self.first_dims is None - else self.first_dims.tolist() - ) + # Default per-tensor first dim: depends on whether last_dims varies. + # * same-shape (first/last both None): logical_shape[0] // num_tensors + # * varying-last-dim (first None, last given): logical_shape[0] (every tensor + # shares the full row count) + # When self.first_dims is set we use it directly below, so no default is needed. + if self.first_dims is None: + if self.last_dims is None: + default_first_dim = self.logical_shape[0] // self.num_tensors + else: + default_first_dim = self.logical_shape[0] + first_dims_list = [default_first_dim] * self.num_tensors + else: + first_dims_list = self.first_dims.tolist() last_dims_list = ( [self.logical_shape[1]] * self.num_tensors if self.last_dims is None @@ -1119,6 +1152,8 @@ def split_into_quantized_tensors( scale_inv = None if self.scale_inv is not None: scale_inv = self.scale_inv[i : i + 1] + elif self.columnwise_scale_inv is not None: + scale_inv = self.columnwise_scale_inv[i : i + 1] if quantizer.internal: float8_tensor_class = Float8TensorStorage